From 43b4aaea3e40e20ef53ce2bfecfa1b72060bf2b5 Mon Sep 17 00:00:00 2001 From: Sriraj <168966222+sricursion@users.noreply.github.com> Date: Wed, 7 Oct 2026 10:25:34 +0530 Subject: [PATCH 001/624] [Bugfix][Frontend] Return 4xx for client errors in /cohere/v2/chat (#60309) Signed-off-by: sricursion Co-authored-by: Claude --- tests/entrypoints/cohere/test_api_router.py | 28 ++++++++++++++++----- vllm/entrypoints/cohere/api_router.py | 23 ++++++----------- 2 files changed, 29 insertions(+), 22 deletions(-) diff --git a/tests/entrypoints/cohere/test_api_router.py b/tests/entrypoints/cohere/test_api_router.py index 1bb2151621c0..209ceb4e3ef4 100644 --- a/tests/entrypoints/cohere/test_api_router.py +++ b/tests/entrypoints/cohere/test_api_router.py @@ -50,7 +50,7 @@ from vllm.entrypoints.serve.exception_handling.handlers.vllm_error import ( vllm_error_handler, ) -from vllm.exceptions import VLLMError +from vllm.exceptions import VLLMError, VLLMValidationError from vllm.sampling_params import SamplingParams @@ -78,7 +78,7 @@ class _Handler: * a :class:`CohereChatV2Response` (non-streaming JSON path); * an async generator yielding SSE frames (streaming path); * an :class:`ErrorResponse` (error envelope path); or - * an exception (router-level 500 path). + * an exception (router-level error path). """ def __init__(self, result): @@ -371,6 +371,22 @@ def test_handler_exception_returns_500_envelope(self): body = r.json() assert body == {"message": "kaboom"} + @pytest.mark.parametrize( + "exc", + [ + VLLMValidationError("bad temperature", parameter="temperature"), + ValueError("bad temperature"), + ], + ids=["vllm_validation_error", "value_error"], + ) + def test_handler_client_error_returns_400_envelope(self, exc): + """Client-caused handler errors are 4xx, matching the OpenAI routes.""" + app = _build_app(handler=_Handler(exc)) + with TestClient(app) as client: + r = client.post("/cohere/v2/chat", json=_minimal_request_body()) + assert r.status_code == HTTPStatus.BAD_REQUEST + assert "bad temperature" in r.json()["message"] + def test_non_json_content_type_rejected(self): """The ``validate_json_request`` dependency raises ``RequestValidationError`` (HTTP 422) for non-JSON content @@ -525,10 +541,10 @@ def test_render_exception_returns_500_envelope(self): assert r.status_code == HTTPStatus.INTERNAL_SERVER_ERROR assert r.json() == {"message": "tokenizer exploded"} - def test_conversion_exception_returns_500_envelope(self): + def test_conversion_client_error_returns_400_envelope(self): """Conversion runs inside the same ``try`` as the render call, so - a bad v2 body that slips past Pydantic still yields the Cohere - error envelope rather than an unhandled 500 with vLLM's shape. + a bad v2 body that slips past Pydantic still yields a 4xx in the + Cohere error envelope rather than an unhandled 500 with vLLM's shape. """ class _Boom: @@ -538,7 +554,7 @@ def to_chat_completion_request(self, request): app = _build_render_app(_Boom(), _RenderHandler(_generate_request())) with TestClient(app) as client: r = client.post("/cohere/v2/chat/render", json=_minimal_request_body()) - assert r.status_code == HTTPStatus.INTERNAL_SERVER_ERROR + assert r.status_code == HTTPStatus.BAD_REQUEST assert r.json() == {"message": "unconvertible message"} def test_invalid_body_returns_422(self): diff --git a/vllm/entrypoints/cohere/api_router.py b/vllm/entrypoints/cohere/api_router.py index db6366df7121..5362b1d9e95d 100644 --- a/vllm/entrypoints/cohere/api_router.py +++ b/vllm/entrypoints/cohere/api_router.py @@ -40,6 +40,9 @@ import vllm.envs as envs from vllm.entrypoints.serve.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.exception_handling.error_response import ( + create_error_response, +) from vllm.entrypoints.serve.exception_handling.utils import sanitize_message from vllm.entrypoints.serve.utils.api_utils import ( load_aware_call, @@ -138,15 +141,9 @@ async def chat_v2(request: CohereChatV2Request, raw_request: Request): try: result = await handler.create_chat_v2(request, raw_request) - except Exception as e: # noqa: BLE001 - report as 500 for parity + except Exception as e: logger.exception("Error in /cohere/v2/chat: %s", e) - return JSONResponse( - status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, - content=CohereError( - message=sanitize_message(str(e)), - id=_request_id(raw_request), - ).model_dump(exclude_none=True), - ) + return _error_response(create_error_response(e), raw_request) match result: case ErrorResponse(): @@ -194,15 +191,9 @@ async def chat_v2_render(request: CohereChatV2Request, raw_request: Request): try: chat_request = handler.to_chat_completion_request(request) result = await render_handler.render_chat_request(chat_request) - except Exception as e: # noqa: BLE001 - report as 500 for parity + except Exception as e: logger.exception("Error in /cohere/v2/chat/render: %s", e) - return JSONResponse( - status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, - content=CohereError( - message=sanitize_message(str(e)), - id=_request_id(raw_request), - ).model_dump(exclude_none=True), - ) + return _error_response(create_error_response(e), raw_request) if isinstance(result, ErrorResponse): return _error_response(result, raw_request) From cbee82324c4ac2e780514122604f8f5fdb7c1d21 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Wed, 7 Oct 2026 14:22:10 +0800 Subject: [PATCH 002/624] [Docs][RL] Move pause/resume and batch invariance examples to examples/rl (#59949) Signed-off-by: wang.yuqi Signed-off-by: aoshen02 Co-authored-by: wang.yuqi Co-authored-by: Claude Opus 5.5 --- .github/CODEOWNERS | 1 - docs/usage/reproducibility.md | 2 +- .../batch_invariance/reproducibility_offline.py | 0 .../{features => rl}/pause_resume/data_parallel_pause_resume.py | 0 examples/{features => rl}/pause_resume/pause_resume_offline.py | 0 mkdocs.yaml | 2 ++ 6 files changed, 3 insertions(+), 2 deletions(-) rename examples/{features => rl}/batch_invariance/reproducibility_offline.py (100%) rename examples/{features => rl}/pause_resume/data_parallel_pause_resume.py (100%) rename examples/{features => rl}/pause_resume/pause_resume_offline.py (100%) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index ac64cf1225d4..71d501eabe7c 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -235,7 +235,6 @@ mkdocs.yaml @hmellor /tests/entrypoints/serve/dev/rlhf @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche @AndreasKaratzas @aoshen02 /tests/v1/worker/test_gpu_worker_weight_transfer.py @aoshen02 /examples/rl @aoshen02 -/examples/features/pause_resume @aoshen02 /docs/training/rlhf.md @aoshen02 /docs/training/async_rl.md @aoshen02 /docs/training/weight_transfer @aoshen02 diff --git a/docs/usage/reproducibility.md b/docs/usage/reproducibility.md index 680791bbe24a..ef34855ef4ff 100644 --- a/docs/usage/reproducibility.md +++ b/docs/usage/reproducibility.md @@ -7,7 +7,7 @@ reproducible results: or enable [batch invariance](../features/batch_invariance.md) to make the outputs insensitive to scheduling. - In online mode, you can only enable [batch invariance](../features/batch_invariance.md). -Example: [examples/features/batch_invariance/reproducibility_offline.py](../../examples/features/batch_invariance/reproducibility_offline.py) +Example: [examples/rl/batch_invariance/reproducibility_offline.py](../../examples/rl/batch_invariance/reproducibility_offline.py) !!! warning diff --git a/examples/features/batch_invariance/reproducibility_offline.py b/examples/rl/batch_invariance/reproducibility_offline.py similarity index 100% rename from examples/features/batch_invariance/reproducibility_offline.py rename to examples/rl/batch_invariance/reproducibility_offline.py diff --git a/examples/features/pause_resume/data_parallel_pause_resume.py b/examples/rl/pause_resume/data_parallel_pause_resume.py similarity index 100% rename from examples/features/pause_resume/data_parallel_pause_resume.py rename to examples/rl/pause_resume/data_parallel_pause_resume.py diff --git a/examples/features/pause_resume/pause_resume_offline.py b/examples/rl/pause_resume/pause_resume_offline.py similarity index 100% rename from examples/features/pause_resume/pause_resume_offline.py rename to examples/rl/pause_resume/pause_resume_offline.py diff --git a/mkdocs.yaml b/mkdocs.yaml index 8c712257eecd..c272e6a68754 100644 --- a/mkdocs.yaml +++ b/mkdocs.yaml @@ -102,6 +102,8 @@ plugins: features/quantization/int8.md: features/quantization/llm_compressor/int8_w8a8.md serving/openai_compatible_server.md: serving/online_serving/README.md examples/others/lmcache.md: examples/disaggregated/lmcache.md + examples/features/batch_invariance.md: examples/rl/batch_invariance.md + examples/features/pause_resume.md: examples/rl/pause_resume.md markdown_extensions: - attr_list From e6fc81bc7892f2f58c0e347a701fc060ceef44bb Mon Sep 17 00:00:00 2001 From: Tommy Date: Wed, 7 Oct 2026 03:32:00 -0400 Subject: [PATCH 003/624] [Bugfix][Scheduler] Let pooling chunked prefill use full context (#48039) Signed-off-by: Tommy Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: OpenAI Codex Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- tests/v1/core/test_scheduler.py | 19 +++++++++++++++++++ tests/v1/core/utils.py | 3 +++ vllm/v1/core/sched/scheduler.py | 10 ++++++++-- 3 files changed, 30 insertions(+), 2 deletions(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index f78db1917db0..7a8a1a3e4ea5 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -35,6 +35,7 @@ MultiModalKwargsItem, PlaceholderRange, ) +from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.utils.hashing import sha256 from vllm.v1.core.encoder_cache_manager import EncoderCacheManager @@ -279,6 +280,24 @@ def test_schedule(enable_prefix_caching: bool, prompt_logprobs: int | None): assert scheduler.running[i] == request +def test_pooling_chunked_prefill_can_finish_at_max_model_len(): + scheduler = create_scheduler( + max_num_seqs=1, max_model_len=8, max_num_batched_tokens=4, runner="pooling" + ) + assert scheduler.max_model_len == 8 + request = Request( + request_id="pool", + prompt_token_ids=[1] * 8, + sampling_params=None, + pooling_params=PoolingParams(task="embed"), + ) + scheduler.add_request(request) + + scheduler.schedule() + # The final chunk must reach max_model_len, not stop one token short. + assert scheduler.schedule().num_scheduled_tokens == {"pool": 4} + + def test_scheduler_stats_route_to_existing_output_client(): scheduler = create_scheduler() request = create_requests(num_requests=1)[0] diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 49aaca3d7c6a..b5e1bf89329e 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -18,6 +18,7 @@ SpeculativeConfig, VllmConfig, ) +from vllm.config.model import RunnerOption from vllm.config.scheduler import SchedulerPolicy from vllm.multimodal.inputs import ( MultiModalFeatureSpec, @@ -78,6 +79,7 @@ def create_scheduler( ec_role: str | None = None, use_v2_model_runner: bool | None = None, kv_cache_spec: KVCacheSpec | None = None, + runner: RunnerOption = "auto", per_request_spec_decode_metrics: str = "none", scheduling_policy: SchedulerPolicy = "fcfs", diffusion_canvas_length: int | None = None, @@ -109,6 +111,7 @@ def create_scheduler( # The scheduler reads model_config.max_model_len, not the # SchedulerConfig one, so both must agree. max_model_len=max_model_len, + runner=runner, ) if mm_encoder_only: model_config.multimodal_config.mm_encoder_only = True diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 8bdf91da86da..f908a31a1b20 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -147,9 +147,15 @@ def __init__( self.kv_events_config is not None and self.kv_events_config.enable_kv_cache_events ) - # Diffusion models may not sample any tokens for a denoising step. + # Pooling and diffusion models do not append a sampled token after + # prefill, so they should not reserve one context slot while scheduling. self.num_sampled_tokens_per_step = ( - 1 if not vllm_config.model_config.is_diffusion else 0 + 0 + if ( + vllm_config.model_config.runner_type == "pooling" + or vllm_config.model_config.is_diffusion + ) + else 1 ) # Create KVConnector for the Scheduler. Note that each Worker From 7436a7f119dedd0c4d5f4d745aa6c2305ee82252 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 7 Oct 2026 15:45:15 +0800 Subject: [PATCH 004/624] [Rust Frontend] Backtrack from safe text at a complete marker (#59563) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- rust/src/parser/src/tool/json/mod.rs | 11 ++++++++-- rust/src/parser/src/tool/json/qwen.rs | 17 ++++++++++++++ rust/src/parser/src/unified/hf/parser.rs | 2 +- rust/src/parser/src/utils.rs | 28 ++++++++++++++++++++---- 4 files changed, 51 insertions(+), 7 deletions(-) diff --git a/rust/src/parser/src/tool/json/mod.rs b/rust/src/parser/src/tool/json/mod.rs index 1bc50588736c..487be423aef2 100644 --- a/rust/src/parser/src/tool/json/mod.rs +++ b/rust/src/parser/src/tool/json/mod.rs @@ -378,9 +378,16 @@ fn safe_text_event( input: &mut JsonToolInput<'_>, config: JsonToolCallConfig, ) -> ModalResult { + // Stop only where a tool call can start: with exact whitespace, a bare + // start marker without it is text. + let with_whitespace = |marker: &str| match config.marker_whitespace { + JsonToolCallWhitespace::Optional => marker.to_string(), + JsonToolCallWhitespace::Exact(whitespace) => format!("{marker}{whitespace}"), + }; + let start_marker = with_whitespace(config.start_marker); match config.framed_start_marker { - Some(marker) => safe_text_len_mul(input, &[marker, config.start_marker]), - None => safe_text_len(input, config.start_marker), + Some(marker) => safe_text_len_mul(input, &[with_whitespace(marker), start_marker]), + None => safe_text_len(input, start_marker), } .map(|len| JsonToolCallEvent::Text { len }) } diff --git a/rust/src/parser/src/tool/json/qwen.rs b/rust/src/parser/src/tool/json/qwen.rs index ca7c5919cfcf..c4db8c2fbc8b 100644 --- a/rust/src/parser/src/tool/json/qwen.rs +++ b/rust/src/parser/src/tool/json/qwen.rs @@ -205,6 +205,23 @@ mod tests { assert!(output.calls().is_empty()); } + #[test] + fn qwen_xml_streams_start_marker_without_newline_as_text() { + for chunk in [ + "{\"name\":\"get_weather\"", + "hi\n{\"name\":\"get_weather\"", + ] { + let mut parser = Qwen3XmlToolParser::new(&test_tools()); + let mut output = ToolParserOutput::default(); + + parser.parse_into(chunk, &mut output).unwrap(); + + // Not a tool-call start: released without waiting for the end of the stream. + assert_eq!(output.normal_text(), chunk); + assert!(output.calls().is_empty()); + } + } + #[test] fn qwen_xml_requires_newline_before_tool_call_end() { let mut parser = Qwen3XmlToolParser::new(&test_tools()); diff --git a/rust/src/parser/src/unified/hf/parser.rs b/rust/src/parser/src/unified/hf/parser.rs index f523c0f31a76..5d5065f109b6 100644 --- a/rust/src/parser/src/unified/hf/parser.rs +++ b/rust/src/parser/src/unified/hf/parser.rs @@ -599,7 +599,7 @@ fn parse_next_event( match safe_text_len_mul(input, &markers) { Ok(_) => return Ok(HfEvent::Text), // A marker, complete or partial, starts here (or the input is empty). - Err(ErrMode::Incomplete(_)) => {} + Err(ErrMode::Backtrack(_) | ErrMode::Incomplete(_)) => {} Err(error) => return Err(error), } diff --git a/rust/src/parser/src/utils.rs b/rust/src/parser/src/utils.rs index 1324e297fbe5..cb7bdce284c5 100644 --- a/rust/src/parser/src/utils.rs +++ b/rust/src/parser/src/utils.rs @@ -118,7 +118,10 @@ pub fn safe_text_len<'i, I: MarkerStream<'i>, M: MarkerLike>( /// the end of the input for plain markers, and only a genuine special token's /// incomplete marker for guarded ones. /// -/// Returns the text length in bytes, and advances the input. +/// Returns the text length in bytes, and advances the input. Fails with +/// `Backtrack` when a complete marker starts the input, so another branch can +/// consume it, and with `Incomplete` when the input is empty or starts with a +/// marker that needs more input to decide. pub fn safe_text_len_mul<'i, I: MarkerStream<'i>, M: MarkerLike>( input: &mut I, markers: &[M], @@ -128,12 +131,18 @@ pub fn safe_text_len_mul<'i, I: MarkerStream<'i>, M: MarkerLike>( } let start = input.offset(); - let stop = match scan_markers(input, markers, start)? { - Scan::Found(at) | Scan::Pending(at) | Scan::Clear(at) => at, + let (stop, found) = match scan_markers(input, markers, start)? { + Scan::Found(at) => (at, true), + Scan::Pending(at) | Scan::Clear(at) => (at, false), }; let emit_len = stop - start; if emit_len == 0 { - return incomplete(); + // A complete marker starts here: there is no text, whatever input follows. + return if found { + Err(ErrMode::Backtrack(ContextError::new())) + } else { + incomplete() + }; } input.next_slice(emit_len); @@ -529,6 +538,17 @@ mod tests { assert!(matches!(error, ErrMode::Incomplete(_))); } + #[test] + fn safe_text_len_backtracks_at_complete_marker() { + let mut input = Partial::new("body"); + let checkpoint = input.checkpoint(); + + let error = safe_text_len(&mut input, "").unwrap_err(); + + assert!(matches!(error, ErrMode::Backtrack(_))); + assert_eq!(input.offset_from(&checkpoint), 0); + } + #[test] fn safe_text_len_mul_stops_before_earliest_marker() { let mut input = Partial::new("hello<|tool_call>"); From 90f4936d9ee05f6108a401a3acc5c5631f5835f7 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Wed, 7 Oct 2026 10:21:47 +0200 Subject: [PATCH 005/624] [CI/Build] Widen the ragged prefill scoring tolerance, drop batch invariance (#60192) Signed-off-by: Stefan Koncarevic --- tests/v1/sample/test_logprobs.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index 813bea0bf3aa..a9b73f0ca2a7 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -1408,10 +1408,6 @@ def test_prompt_logprobs_with_chunking_and_preemption(): def test_prompt_logprob_token_ids_with_chunking_and_preemption(monkeypatch): """Per-row scores stay row-aligned across chunked prefill and preemption.""" monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") - # The ragged check below compares a second generate call against the first - # one. Default kernel selection varies between processes, which moves bf16 - # logprob tails by up to 0.2; under this the two runs are bit-identical. - monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1") prompts = [ "The following numbers of the sequence " @@ -1474,7 +1470,10 @@ def make_params(row_ids): assert float(scores.max()) <= 1e-3, "logprobs must be <= 0" # Ragged rows: each row keeps a prefix of the candidates and must - # reproduce that prefix of the full scores, with -inf padding. + # reproduce that prefix of the full scores, with -inf padding. On ROCm + # the second generate call picks its own kernels, which moves bf16 + # tails by up to 0.19; a misaligned row is off by nats either way. + ragged_atol = 0.2 if current_platform.is_rocm() else 1e-2 ragged_params = make_params(lambda j: candidate_ids(j)[: 1 + j % 4]) ragged_outputs = vllm_model.llm.generate(token_prompts, ragged_params) for output, ragged_output in zip(outputs, ragged_outputs): @@ -1483,7 +1482,7 @@ def make_params(row_ids): n = 1 + j % 4 expected[j, :n] = output.prompt_token_id_logprobs[j, :n] np.testing.assert_allclose( - ragged_output.prompt_token_id_logprobs, expected, atol=1e-2 + ragged_output.prompt_token_id_logprobs, expected, atol=ragged_atol ) # The row count is checked against the prompt at admission. From bc8dd4ad275c838d6315389432b2010d9348c900 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 7 Oct 2026 16:28:27 +0800 Subject: [PATCH 006/624] [Rust Frontend] Share `SchemaRoot` between argument coercion and grammars (#59408) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- rust/src/parser/src/lib.rs | 1 + .../parser/src/output_grammar/arguments.rs | 236 +------- rust/src/parser/src/schema.rs | 499 +++++++++++++++++ rust/src/parser/src/tool/parameters.rs | 505 ++++++++++-------- rust/src/parser/src/tool/qwen_coder.rs | 35 +- .../src/unified/kimi_k3/structural_tag.rs | 22 +- 6 files changed, 857 insertions(+), 441 deletions(-) create mode 100644 rust/src/parser/src/schema.rs diff --git a/rust/src/parser/src/lib.rs b/rust/src/parser/src/lib.rs index fd4aa0aa7f71..26256d332103 100644 --- a/rust/src/parser/src/lib.rs +++ b/rust/src/parser/src/lib.rs @@ -5,6 +5,7 @@ pub mod output_grammar; pub mod reasoning; +pub(crate) mod schema; pub mod tool; pub mod unified; pub(crate) mod utils; diff --git a/rust/src/parser/src/output_grammar/arguments.rs b/rust/src/parser/src/output_grammar/arguments.rs index a7d9fb307875..0f621789a3bc 100644 --- a/rust/src/parser/src/output_grammar/arguments.rs +++ b/rust/src/parser/src/output_grammar/arguments.rs @@ -25,14 +25,11 @@ //! - `minProperties` and `maxProperties` are not enforced. //! - `allOf` with more than one schema accepts any value. -use std::collections::HashSet; - -use serde_json::{Map, Value, json}; +use serde_json::{Value, json}; use xgrammar_structural_tag::format::{Format, JsonSchemaFormat}; -/// Bound on `$ref` and combinator nesting within one value, which keeps a long -/// chain of references from overflowing the stack. -const MAX_SCHEMA_DEPTH: usize = 32; +pub use crate::schema::JsonType; +use crate::schema::{self, OptionSource, SchemaOption, SchemaRoot}; /// The schema accepting any value. static ANY_SCHEMA: Value = Value::Bool(true); @@ -77,95 +74,22 @@ impl ParameterKey<'_> { } } -/// JSON type of a value option. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum JsonType { - String, - Number, - Boolean, - Null, - Object, - Array, -} - -impl JsonType { - const ALL: [Self; 6] = [ - Self::String, - Self::Number, - Self::Boolean, - Self::Null, - Self::Object, - Self::Array, - ]; - - /// The type's JSON name. Integers are numbers. - pub fn name(self) -> &'static str { - match self { - Self::String => "string", - Self::Number => "number", - Self::Boolean => "boolean", - Self::Null => "null", - Self::Object => "object", - Self::Array => "array", - } - } - - fn of_schema_type(name: &str) -> Option { - Some(match name { - "string" => Self::String, - "integer" | "number" => Self::Number, - "boolean" => Self::Boolean, - "null" => Self::Null, - "object" => Self::Object, - "array" => Self::Array, - _ => return None, - }) - } - - fn of_value(value: &Value) -> Self { - match value { - Value::String(_) => Self::String, - Value::Number(_) => Self::Number, - Value::Bool(_) => Self::Boolean, - Value::Null => Self::Null, - Value::Object(_) => Self::Object, - Value::Array(_) => Self::Array, - } - } - - /// Schema accepting every JSON value of this type. - fn any_schema(self) -> Value { - match self { - Self::Object => json!({ "type": "object", "additionalProperties": true }), - Self::Array => json!({ "type": "array", "items": true }), - _ => json!({ "type": self.name() }), - } - } -} - /// One option of a parameter value: a single JSON type under a narrowed /// schema, a constant, or any value of the type. pub struct ValueOption<'a> { - /// The option's JSON type. + /// The option's JSON Schema type. pub ty: JsonType, - source: OptionSource, + source: OptionSource<'a>, cx: &'a ArgumentContext<'a>, } -enum OptionSource { - /// A schema narrowed to [`ValueOption::ty`]. - Schema(Map), - /// A `const` or `enum` value. - Literal(Value), - /// Any value of [`ValueOption::ty`]. - Any, -} - impl ValueOption<'_> { /// The option as JSON text. pub fn json(&self) -> Format { match &self.source { - OptionSource::Schema(schema) => self.cx.json(Value::Object(schema.clone())), + OptionSource::Schema(schema) => { + self.cx.json(Value::Object(schema::narrow(schema, self.ty).into_owned())) + } OptionSource::Literal(value) => self.cx.json(json!({ "const": value })), OptionSource::Any => self.cx.json(self.ty.any_schema()), } @@ -216,26 +140,18 @@ pub struct ArgumentOptions { pub excludes: Vec, } -/// Reference targets expanded while collecting one value's options. -#[derive(Default)] -struct Expansion { - expanded: HashSet<*const Value>, - /// Whether a reference was skipped because its target already expanded. - skipped: bool, -} - /// Request state shared by the value options of one call. struct ArgumentContext<'a> { options: &'a ArgumentOptions, - /// The parameters schema, which local `$ref`s resolve against. - root: &'a Value, + root: SchemaRoot<'a>, } impl ArgumentContext<'_> { /// A JSON value under `schema`, which keeps the root's definitions so its /// local `$ref`s still resolve. fn json(&self, mut schema: Value) -> Format { - if let (Some(schema), Some(root)) = (schema.as_object_mut(), self.root.as_object()) { + if let (Some(schema), Some(root)) = (schema.as_object_mut(), self.root.schema().as_object()) + { for key in ["$defs", "definitions"] { if let Some(definitions) = root.get(key) { schema.entry(key).or_insert_with(|| definitions.clone()); @@ -250,122 +166,14 @@ impl ArgumentContext<'_> { }) } - fn resolve_ref(&self, reference: &str) -> Option<&Value> { - self.root.pointer(reference.strip_prefix('#')?) - } - - /// Follow `$ref`s and single-schema `allOf`s at the root. - fn resolve<'s>(&'s self, schema: &'s Value, depth: usize) -> &'s Value { - if depth > MAX_SCHEMA_DEPTH { - return &ANY_SCHEMA; - } - if let Some(reference) = schema.get("$ref").and_then(Value::as_str) { - return match self.resolve_ref(reference) { - Some(target) => self.resolve(target, depth + 1), - None => &ANY_SCHEMA, - }; - } - match schema.get("allOf").and_then(Value::as_array).map(Vec::as_slice) { - Some([schema]) => self.resolve(schema, depth + 1), - _ => schema, - } - } - /// Every option of a value under `schema`. - /// - /// Each reference target expands at most once. A union gains no values - /// from a repeated alternative, and a definition shared by several - /// alternatives would otherwise expand once per path to it, exponentially - /// in the nesting. A schema whose references only lead back into it - /// accepts any value. - fn value_options(&self, schema: &Value) -> Vec> { - let mut expansion = Expansion::default(); - let options = self.value_options_at(schema, 0, &mut expansion); - if options.is_empty() && expansion.skipped { - return self.any_value(); - } - options - } - - fn value_options_at( - &self, - schema: &Value, - depth: usize, - expansion: &mut Expansion, - ) -> Vec> { - let schema = match schema { - _ if depth > MAX_SCHEMA_DEPTH => return self.any_value(), - Value::Bool(false) => return vec![], - Value::Object(schema) => schema, - _ => return self.any_value(), - }; - if let Some(reference) = schema.get("$ref").and_then(Value::as_str) { - let Some(target) = self.resolve_ref(reference) else { - return self.any_value(); - }; - if !expansion.expanded.insert(std::ptr::from_ref(target)) { - expansion.skipped = true; - return vec![]; - } - return self.value_options_at(target, depth + 1, expansion); - } - if let Some(value) = schema.get("const") { - return vec![self.literal(value)]; - } - if let Some(values) = schema.get("enum").and_then(Value::as_array) { - return values.iter().map(|value| self.literal(value)).collect(); - } - if let Some(options) = - schema.get("anyOf").or_else(|| schema.get("oneOf")).and_then(Value::as_array) - { - return options - .iter() - .flat_map(|option| self.value_options_at(option, depth + 1, expansion)) - .collect(); - } - if let Some(schemas) = schema.get("allOf").and_then(Value::as_array) { - return match schemas.as_slice() { - [schema] => self.value_options_at(schema, depth + 1, expansion), - _ => self.any_value(), - }; - } - match schema.get("type") { - Some(Value::String(name)) => self.typed(schema.clone(), name).into_iter().collect(), - Some(Value::Array(names)) => names - .iter() - .filter_map(Value::as_str) - .filter_map(|name| { - let mut schema = schema.clone(); - schema.insert("type".to_string(), Value::String(name.to_string())); - self.typed(schema, name) - }) - .collect(), - _ => self.any_value(), - } - } - - fn typed(&self, schema: Map, type_name: &str) -> Option> { - Some(ValueOption { - ty: JsonType::of_schema_type(type_name)?, - source: OptionSource::Schema(schema), - cx: self, - }) - } - - fn literal(&self, value: &Value) -> ValueOption<'_> { - ValueOption { - ty: JsonType::of_value(value), - source: OptionSource::Literal(value.clone()), - cx: self, - } - } - - fn any_value(&self) -> Vec> { - JsonType::ALL + fn value_options<'s>(&'s self, schema: &'s Value) -> Vec> { + let root: SchemaRoot<'s> = self.root; + root.options(schema) .into_iter() - .map(|ty| ValueOption { + .map(|SchemaOption { ty, source }| ValueOption { ty, - source: OptionSource::Any, + source, cx: self, }) .collect() @@ -377,12 +185,12 @@ impl ArgumentContext<'_> { pub fn arguments(schema: &Value, syntax: &dyn ArgumentSyntax, options: &ArgumentOptions) -> Format { let cx = ArgumentContext { options, - root: schema, + root: SchemaRoot::new(schema), }; let parameter = |key: ParameterKey<'_>, schema: &Value| syntax.parameter(key, &cx.value_options(schema)); - let schema = cx.resolve(schema, 0); + let schema = cx.root.resolved(); let (properties, additional) = match schema { Value::Object(schema) => { let additional = match schema.get("additionalProperties") { @@ -492,6 +300,8 @@ pub fn group_by<'o, 'a, K: PartialEq>( #[cfg(test)] mod tests { + use serde_json::Map; + use super::*; use crate::output_grammar::test_utils::outline; use expect_test::{Expect, expect}; @@ -617,7 +427,7 @@ mod tests { or `celsius` `fahrenheit` - tag `` json(integer(minimum=1)) any_order `` + tag `` json(integer(minimum=1)) any_order `` "#]], ); } @@ -639,13 +449,13 @@ mod tests { sequence optional or - tag `` json(integer) `` + tag `` json(integer) `` tag `` text excluding [``] `` tag `` json(null) `` optional or tag `` `fast` `` - tag `` json(1) `` + tag `` json(1) `` or tag `` /[a-z]+/ `` tag `` json(string[]) `` diff --git a/rust/src/parser/src/schema.rs b/rust/src/parser/src/schema.rs new file mode 100644 index 000000000000..c27427de2e2c --- /dev/null +++ b/rust/src/parser/src/schema.rs @@ -0,0 +1,499 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Views of tool parameter JSON schemas. +//! +//! Tool parsers coerce generated parameter values by their schema, and output +//! grammars constrain the same values by it. Both read a parameter's schema +//! through [`SchemaRoot`], so they agree on which JSON types the value may take. + +use std::borrow::Cow; +use std::collections::HashSet; + +use serde_json::{Map, Value, json}; + +/// Bound on `$ref` and combinator nesting within one value, which keeps a long +/// chain of references from overflowing the stack. +const MAX_SCHEMA_DEPTH: usize = 32; + +/// The schema accepting any value. +static ANY_SCHEMA: Value = Value::Bool(true); + +/// A tool's parameters schema, which local `$ref`s resolve against. +#[derive(Debug, Clone, Copy)] +pub(crate) struct SchemaRoot<'a> { + schema: &'a Value, +} + +/// Reference targets expanded while collecting one value's options. +#[derive(Default)] +struct Expansion { + expanded: HashSet<*const Value>, + /// Whether a reference was skipped because its target already expanded. + skipped: bool, +} + +/// Roots are equal when they are the same schema, which a pointer compare +/// decides without walking it. +impl PartialEq for SchemaRoot<'_> { + fn eq(&self, other: &Self) -> bool { + std::ptr::eq(self.schema, other.schema) + } +} + +impl Eq for SchemaRoot<'_> {} + +/// One option of a value: a single JSON type under a narrowed schema, a +/// constant, or any value of the type. +#[derive(Debug, Clone)] +pub(crate) struct SchemaOption<'a> { + pub(crate) ty: JsonType, + pub(crate) source: OptionSource<'a>, +} + +/// Where a [`SchemaOption`] comes from. +#[derive(Debug, Clone)] +pub(crate) enum OptionSource<'a> { + /// A schema, of which the option takes the values of its type; [`narrow`] + /// sets that type on it. + Schema(&'a Map), + /// A `const` or `enum` value. + Literal(&'a Value), + /// Any value of the option's type. + Any, +} + +/// JSON Schema type of a value option. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum JsonType { + String, + Integer, + Number, + Boolean, + Null, + Object, + Array, +} + +impl<'a> SchemaRoot<'a> { + pub(crate) fn new(schema: &'a Value) -> Self { + Self { schema } + } + + /// The parameters schema itself. + pub(crate) fn schema(&self) -> &'a Value { + self.schema + } + + /// The parameters schema after following `$ref`s and single-schema + /// `allOf`s. + pub(crate) fn resolved(&self) -> &'a Value { + self.resolve(self.schema) + } + + /// Follow `$ref`s and single-schema `allOf`s. A reference that does not + /// resolve, or a chain deeper than [`MAX_SCHEMA_DEPTH`], accepts any value. + fn resolve(&self, schema: &'a Value) -> &'a Value { + let mut schema = schema; + for _ in 0..=MAX_SCHEMA_DEPTH { + if let Some(reference) = schema.get("$ref").and_then(Value::as_str) { + match self.resolve_ref(reference) { + Some(target) => schema = target, + None => return &ANY_SCHEMA, + } + continue; + } + match schema.get("allOf").and_then(Value::as_array).map(Vec::as_slice) { + Some([inner]) => schema = inner, + _ => return schema, + } + } + &ANY_SCHEMA + } + + /// Every option of a value under `schema`. + /// + /// `const` and `enum` values become literal options, and `anyOf`, `oneOf`, + /// and type arrays one option per alternative, in schema order. An object + /// or array schema without `type` is typed by its `properties`, + /// `additionalProperties`, or `items`. `"nullable": true`, the OpenAPI + /// spelling of an optional value, adds a `null` option. Type names accept + /// the aliases that real tool schemas use, such as `int` or `dict`. + /// + /// Each reference target expands at most once. A union gains no values + /// from a repeated alternative, and a definition shared by several + /// alternatives would otherwise expand once per path to it, exponentially + /// in the nesting. A schema that constrains nothing, or whose references + /// only lead back into it, yields [`Self::any`]. + pub(crate) fn options(&self, schema: &'a Value) -> Vec> { + let mut expansion = Expansion::default(); + let mut options = self.options_at(schema, 0, &mut expansion); + if options.is_empty() && expansion.skipped { + options = Self::any(); + } + let nullable = schema.get("nullable").and_then(Value::as_bool) == Some(true); + if nullable && !options.iter().any(|option| option.ty == JsonType::Null) { + options.push(SchemaOption { + ty: JsonType::Null, + source: OptionSource::Any, + }); + } + options + } + + /// One option per JSON value type, each accepting any value of the type. + pub(crate) fn any() -> Vec> { + JsonType::VALUE_TYPES + .into_iter() + .map(|ty| SchemaOption { + ty, + source: OptionSource::Any, + }) + .collect() + } + + fn resolve_ref(&self, reference: &str) -> Option<&'a Value> { + self.schema.pointer(reference.strip_prefix('#')?) + } + + fn options_at( + &self, + schema: &'a Value, + depth: usize, + expansion: &mut Expansion, + ) -> Vec> { + let schema = match schema { + _ if depth > MAX_SCHEMA_DEPTH => return Self::any(), + Value::Bool(false) => return vec![], + Value::Object(schema) => schema, + _ => return Self::any(), + }; + if let Some(reference) = schema.get("$ref").and_then(Value::as_str) { + let Some(target) = self.resolve_ref(reference) else { + return Self::any(); + }; + if !expansion.expanded.insert(std::ptr::from_ref(target)) { + expansion.skipped = true; + return vec![]; + } + return self.options_at(target, depth + 1, expansion); + } + if let Some(value) = schema.get("const") { + return vec![SchemaOption::literal(value)]; + } + if let Some(values) = schema.get("enum").and_then(Value::as_array) { + return values.iter().map(SchemaOption::literal).collect(); + } + if let Some(options) = + schema.get("anyOf").or_else(|| schema.get("oneOf")).and_then(Value::as_array) + { + return options + .iter() + .flat_map(|option| self.options_at(option, depth + 1, expansion)) + .collect(); + } + if let Some(schemas) = schema.get("allOf").and_then(Value::as_array) { + return match schemas.as_slice() { + [schema] => self.options_at(schema, depth + 1, expansion), + _ => Self::any(), + }; + } + match schema.get("type") { + Some(Value::String(name)) => JsonType::parse(name) + .map(|ty| SchemaOption::typed(schema, ty)) + .into_iter() + .collect(), + Some(Value::Array(names)) => names + .iter() + .filter_map(Value::as_str) + .filter_map(JsonType::parse) + .map(|ty| SchemaOption::typed(schema, ty)) + .collect(), + Some(_) => Self::any(), + None if schema.contains_key("properties") + || schema.contains_key("additionalProperties") => + { + vec![SchemaOption::typed(schema, JsonType::Object)] + } + None if schema.contains_key("items") => { + vec![SchemaOption::typed(schema, JsonType::Array)] + } + None => Self::any(), + } + } +} + +impl<'a> SchemaOption<'a> { + /// The schema of a [`OptionSource::Schema`] option. + pub(crate) fn schema(&self) -> Option<&'a Map> { + match self.source { + OptionSource::Schema(schema) => Some(schema), + OptionSource::Literal(_) | OptionSource::Any => None, + } + } + + /// The option of `schema` that takes its values of type `ty`. + fn typed(schema: &'a Map, ty: JsonType) -> Self { + Self { + ty, + source: OptionSource::Schema(schema), + } + } + + fn literal(value: &'a Value) -> Self { + Self { + ty: JsonType::of_value(value), + source: OptionSource::Literal(value), + } + } +} + +/// `schema` with its `type` set to the canonical name of `ty`, copied only +/// when that changes it. +pub(crate) fn narrow(schema: &Map, ty: JsonType) -> Cow<'_, Map> { + if schema.get("type").and_then(Value::as_str) == Some(ty.name()) { + Cow::Borrowed(schema) + } else { + let mut schema = schema.clone(); + schema.insert("type".to_string(), Value::String(ty.name().to_string())); + Cow::Owned(schema) + } +} + +impl JsonType { + /// The types of JSON values, which do not distinguish integers. + const VALUE_TYPES: [Self; 6] = [ + Self::String, + Self::Number, + Self::Boolean, + Self::Null, + Self::Object, + Self::Array, + ]; + + /// The JSON Schema type name. + pub fn name(self) -> &'static str { + match self { + Self::String => "string", + Self::Integer => "integer", + Self::Number => "number", + Self::Boolean => "boolean", + Self::Null => "null", + Self::Object => "object", + Self::Array => "array", + } + } + + /// Schema accepting every JSON value of this type. + pub(crate) fn any_schema(self) -> Value { + match self { + Self::Object => json!({ "type": "object", "additionalProperties": true }), + Self::Array => json!({ "type": "array", "items": true }), + ty => json!({ "type": ty.name() }), + } + } + + /// Parse a schema type name, including aliases used by real tool schemas. + fn parse(name: &str) -> Option { + let name = name.trim().to_ascii_lowercase(); + Some(match name.as_str() { + "string" | "str" | "text" | "varchar" | "char" | "enum" => Self::String, + "integer" | "int" => Self::Integer, + "number" | "float" | "double" => Self::Number, + "boolean" | "bool" | "binary" => Self::Boolean, + "object" | "dict" | "map" => Self::Object, + "array" | "arr" | "list" | "sequence" => Self::Array, + "null" => Self::Null, + _ if name.starts_with("int") + || name.starts_with("uint") + || name.starts_with("long") + || name.starts_with("short") + || name.starts_with("unsigned") => + { + Self::Integer + } + _ if name.starts_with("num") || name.starts_with("float") => Self::Number, + _ if name.starts_with("dict") => Self::Object, + _ if name.starts_with("list") => Self::Array, + _ => return None, + }) + } + + fn of_value(value: &Value) -> Self { + match value { + Value::String(_) => Self::String, + Value::Number(number) if number.is_f64() => Self::Number, + Value::Number(_) => Self::Integer, + Value::Bool(_) => Self::Boolean, + Value::Null => Self::Null, + Value::Object(_) => Self::Object, + Value::Array(_) => Self::Array, + } + } +} + +#[cfg(test)] +mod tests { + use expect_test::{Expect, expect}; + use serde_json::json; + + use super::*; + + /// A Pydantic v2 `model_json_schema()` for a nested model, an enum, and an + /// optional nested model. + fn pydantic_schema() -> Value { + json!({ + "$defs": { + "Place": { + "properties": { "city": { "title": "City", "type": "string" } }, + "required": ["city"], + "title": "Place", + "type": "object" + }, + "Unit": { "enum": ["celsius", "fahrenheit"], "title": "Unit", "type": "string" } + }, + "properties": { + "place": { "$ref": "#/$defs/Place" }, + "unit": { "$ref": "#/$defs/Unit", "default": "celsius" }, + "backup": { "anyOf": [{ "$ref": "#/$defs/Place" }, { "type": "null" }], "default": null } + }, + "required": ["place"], + "title": "Args", + "type": "object" + }) + } + + fn check_options(root: &Value, schema: &Value, expected: Expect) { + let options = SchemaRoot::new(root) + .options(schema) + .into_iter() + .map(|option| match option.source { + OptionSource::Schema(schema) => format!( + "{} {}", + option.ty.name(), + Value::Object(narrow(schema, option.ty).into_owned()) + ), + OptionSource::Literal(value) => format!("{} = {value}", option.ty.name()), + OptionSource::Any => format!("{} any", option.ty.name()), + }) + .collect::>(); + expected.assert_debug_eq(&options); + } + + #[test] + fn references_resolve_to_their_definitions() { + let root = pydantic_schema(); + let properties = &root["properties"]; + check_options( + &root, + &properties["place"], + expect![[r#" + [ + "object {\"properties\":{\"city\":{\"title\":\"City\",\"type\":\"string\"}},\"required\":[\"city\"],\"title\":\"Place\",\"type\":\"object\"}", + ] + "#]], + ); + check_options( + &root, + &properties["unit"], + expect![[r#" + [ + "string = \"celsius\"", + "string = \"fahrenheit\"", + ] + "#]], + ); + check_options( + &root, + &properties["backup"], + expect![[r#" + [ + "object {\"properties\":{\"city\":{\"title\":\"City\",\"type\":\"string\"}},\"required\":[\"city\"],\"title\":\"Place\",\"type\":\"object\"}", + "null {\"type\":\"null\"}", + ] + "#]], + ); + } + + #[test] + fn unresolvable_and_cyclic_references_accept_any_value() { + let root = json!({ + "$defs": { "a": { "$ref": "#/$defs/b" }, "b": { "$ref": "#/$defs/a" } }, + "properties": { + "missing": { "$ref": "#/$defs/missing" }, + "cycle": { "$ref": "#/$defs/a" }, + "remote": { "$ref": "https://example.com/schema.json" } + } + }); + let schema_root = SchemaRoot::new(&root); + for key in ["missing", "cycle", "remote"] { + assert_eq!( + schema_root.resolve(&root["properties"][key]), + &ANY_SCHEMA, + "{key}" + ); + assert!( + schema_root + .options(&root["properties"][key]) + .iter() + .all(|option| matches!(option.source, OptionSource::Any)), + "{key}" + ); + } + } + + #[test] + fn options_normalize_common_schema_spellings() { + let root = json!({}); + check_options( + &root, + &json!({ "type": ["int", "null"], "minimum": 0 }), + expect![[r#" + [ + "integer {\"type\":\"integer\",\"minimum\":0}", + "null {\"type\":\"null\",\"minimum\":0}", + ] + "#]], + ); + check_options( + &root, + &json!({ "type": "string", "nullable": true }), + expect![[r#" + [ + "string {\"type\":\"string\",\"nullable\":true}", + "null any", + ] + "#]], + ); + check_options( + &root, + &json!({ "enum": ["a", 1, 1.5, null] }), + expect![[r#" + [ + "string = \"a\"", + "integer = 1", + "number = 1.5", + "null = null", + ] + "#]], + ); + check_options( + &root, + &json!({ "allOf": [{ "type": "boolean" }] }), + expect![[r#" + [ + "boolean {\"type\":\"boolean\"}", + ] + "#]], + ); + check_options( + &root, + &json!({ "items": { "type": "string" } }), + expect![[r#" + [ + "array {\"items\":{\"type\":\"string\"},\"type\":\"array\"}", + ] + "#]], + ); + } +} diff --git a/rust/src/parser/src/tool/parameters.rs b/rust/src/parser/src/tool/parameters.rs index ef59ce08759b..391ca90f105f 100644 --- a/rust/src/parser/src/tool/parameters.rs +++ b/rust/src/parser/src/tool/parameters.rs @@ -5,6 +5,7 @@ use std::collections::BTreeMap; use serde_json::{Map, Number, Value}; +use crate::schema::{JsonType, OptionSource, SchemaRoot}; use crate::tool::Tool; /// Normalized parameter schemas for all tools in one request. @@ -21,7 +22,9 @@ pub(crate) struct ToolSchemas { /// downstream tool call execution. #[derive(Debug, Clone, Default, PartialEq, Eq)] pub(super) struct ToolSchema { - params: BTreeMap, + /// The tool's parameters schema, read one level at a time as values are + /// converted, so recursive schemas cost only as deep as the value. + parameters: Value, } /// Parameter input for schema-aware conversion. @@ -47,22 +50,38 @@ pub(crate) struct ParamElement { pub value: ParamInput, } -/// Normalized JSON parameter type used for raw string coercion. +/// One level of a parameter's normalized type, used for raw string coercion. +/// +/// Object and array types normalize their children only when a structured +/// value reaches them. #[derive(Debug, Clone, PartialEq, Eq)] -pub(super) enum JsonParamType { +enum JsonParamType<'a> { String, Integer, Number, Boolean, - Object { - properties: BTreeMap, - additional_properties: Option>, - }, - Array { - items: Option>, - }, + Object(ObjectType<'a>), + Array(ArrayType<'a>), Null, - OneOf(Vec), + /// A non-string `const` or `enum` value: text spelling it as JSON decodes + /// to it. + Literal(&'a Value), + OneOf(Vec>), +} + +/// An object type whose field types are normalized on access. +#[derive(Debug, Clone, PartialEq, Eq)] +struct ObjectType<'a> { + root: SchemaRoot<'a>, + properties: Option<&'a Map>, + additional_properties: Option<&'a Value>, +} + +/// An array type whose item type is normalized on access. +#[derive(Debug, Clone, PartialEq, Eq)] +struct ArrayType<'a> { + root: SchemaRoot<'a>, + items: Option<&'a Value>, } impl ToolSchemas { @@ -70,7 +89,12 @@ impl ToolSchemas { pub(crate) fn from_tools(tools: &[Tool]) -> Self { let tools = tools .iter() - .map(|tool| (tool.name.clone(), ToolSchema::from_schema(&tool.parameters))) + .map(|tool| { + ( + tool.name.clone(), + ToolSchema::from_schema(tool.parameters.clone()), + ) + }) .collect(); Self { tools } @@ -117,25 +141,14 @@ impl ToolSchema { /// parameters to be treated as strings. const fn empty() -> &'static Self { static EMPTY: ToolSchema = ToolSchema { - params: BTreeMap::new(), + parameters: Value::Null, }; &EMPTY } - /// Normalize an OpenAI-style tool parameters JSON schema. - fn from_schema(parameters: &Value) -> Self { - let Some(properties) = parameters.get("properties").and_then(Value::as_object) else { - return Self::default(); - }; - - let params = properties - .iter() - .filter_map(|(name, schema)| { - JsonParamType::from_schema(schema).map(|param_type| (name.clone(), param_type)) - }) - .collect(); - - Self { params } + /// Keep an OpenAI-style tool parameters JSON schema. + fn from_schema(parameters: Value) -> Self { + Self { parameters } } /// Convert one parameter value using its normalized schema type. @@ -144,169 +157,86 @@ impl ToolSchema { /// the value fails to convert, this falls back to returning the raw /// string as a JSON string value, or object-like JSON for structured input. fn convert(&self, name: &str, input: ParamInput) -> Value { - convert_with_optional_schema(self.params.get(name), &input) + let root = SchemaRoot::new(&self.parameters); + let param_type = root + .resolved() + .get("properties") + .and_then(|properties| properties.get(name)) + .and_then(|schema| JsonParamType::from_schema(&root, schema)); + convert_with_optional_schema(param_type.as_ref(), &input) } } -impl JsonParamType { - /// Normalize one parameter property schema. - /// - /// `"nullable": true` (the OpenAPI spelling of an optional parameter) adds - /// `null` to the accepted types. - fn from_schema(schema: &Value) -> Option { - let param_type = Self::from_schema_ignoring_nullable(schema)?; - let nullable = schema.get("nullable").and_then(Value::as_bool) == Some(true); - Some(match param_type { - _ if !nullable || param_type.accepts_null() => param_type, - Self::OneOf(mut types) => { - types.push(Self::Null); - Self::OneOf(types) - } - param_type => Self::OneOf(vec![param_type, Self::Null]), - }) - } - - /// Return whether `null` is one of the accepted types. - fn accepts_null(&self) -> bool { - match self { - Self::Null => true, - Self::OneOf(types) => types.iter().any(Self::accepts_null), - _ => false, +impl<'a> JsonParamType<'a> { + /// Normalize one parameter or nested value schema, or `None` when it does + /// not constrain the value's type. + fn from_schema(root: &SchemaRoot<'a>, schema: &'a Value) -> Option { + let options = root.options(schema); + if options.iter().all(|option| matches!(option.source, OptionSource::Any)) { + return None; } - } - - /// Normalize one parameter property schema, without `nullable`. - fn from_schema_ignoring_nullable(schema: &Value) -> Option { - let schema = schema.as_object()?; - - if let Some(type_value) = schema.get("type") { - return Self::from_type_value(type_value, schema); - } - - if let Some(composite) = schema.get("anyOf").or_else(|| schema.get("oneOf")) { - let param_type = composite - .as_array() - .map(|schemas| schemas.iter().filter_map(Self::from_schema).collect::>()) - .filter(|types| !types.is_empty()) - .map(Self::one_of) - .unwrap_or_else(|| Self::object_from_schema(Some(schema))); - return Some(param_type); - } - - // Typically, these types are already handled by checking the "type" field, but - // we can also infer them from their characteristic fields if "type" is missing. - if let Some(values) = schema.get("enum").and_then(Value::as_array) { - // Enum values are treated as strings, except that a `null` member - // makes the parameter nullable (mirrors Python's enum type - // inference), so a literal "null" coerces to JSON null. - if values.iter().any(Value::is_null) { - return Some(Self::one_of(vec![Self::String, Self::Null])); - } - return Some(Self::String); - } - if schema.contains_key("items") { - return Some(Self::array_from_schema(Some(schema))); - } - if schema.contains_key("properties") || schema.contains_key("additionalProperties") { - return Some(Self::object_from_schema(Some(schema))); - } - - None - } - - /// Normalize a JSON schema `type` value. - fn from_type_value(type_value: &Value, schema: &Map) -> Option { - match type_value { - Value::String(kind) => Self::from_type_name(kind, Some(schema)), - Value::Array(kinds) => { - let types = kinds - .iter() - .filter_map(Value::as_str) - .filter_map(|kind| Self::from_type_name(kind, Some(schema))) - .collect::>(); - if types.is_empty() { - None - } else { - Some(Self::one_of(types)) - } - } - _ => None, - } - } - - /// Normalize one JSON schema type name. - fn from_type_name(kind: &str, schema: Option<&Map>) -> Option { - let kind = kind.trim().to_ascii_lowercase(); - match kind.as_str() { - "string" | "str" | "text" | "varchar" | "char" | "enum" => Some(Self::String), - "integer" | "int" => Some(Self::Integer), - "number" | "float" | "double" => Some(Self::Number), - "boolean" | "bool" | "binary" => Some(Self::Boolean), - "object" | "dict" | "map" => Some(Self::object_from_schema(schema)), - "array" | "arr" | "list" | "sequence" => Some(Self::array_from_schema(schema)), - "null" => Some(Self::Null), - _ if kind.starts_with("int") - || kind.starts_with("uint") - || kind.starts_with("long") - || kind.starts_with("short") - || kind.starts_with("unsigned") => + let mut types = Vec::new(); + for option in &options { + // String literals stay with the string type, which keeps text + // verbatim instead of unquoting text that spells one as JSON. + if let OptionSource::Literal(value) = option.source + && !value.is_string() + && !types.contains(&Self::Literal(value)) { - Some(Self::Integer) + types.push(Self::Literal(value)); + } + let schema = option.schema(); + let param_type = match option.ty { + JsonType::String => Self::String, + JsonType::Integer => Self::Integer, + JsonType::Number => Self::Number, + JsonType::Boolean => Self::Boolean, + JsonType::Null => Self::Null, + JsonType::Object => Self::Object(ObjectType { + root: *root, + properties: schema + .and_then(|schema| schema.get("properties")) + .and_then(Value::as_object), + additional_properties: schema + .and_then(|schema| schema.get("additionalProperties")) + .filter(|schema| schema.is_object()), + }), + JsonType::Array => Self::Array(ArrayType { + root: *root, + items: schema.and_then(|schema| schema.get("items")), + }), + }; + if !types.contains(¶m_type) { + types.push(param_type); } - _ if kind.starts_with("num") || kind.starts_with("float") => Some(Self::Number), - _ if kind.starts_with("dict") => Some(Self::object_from_schema(schema)), - _ if kind.starts_with("list") => Some(Self::array_from_schema(schema)), - _ => None, } - } - - /// Normalize object schema fields. - fn object_from_schema(schema: Option<&Map>) -> Self { - let properties = schema - .and_then(|schema| schema.get("properties")) - .and_then(Value::as_object) - .map(|properties| { - properties - .iter() - .filter_map(|(name, schema)| { - Self::from_schema(schema).map(|param_type| (name.clone(), param_type)) - }) - .collect() - }) - .unwrap_or_default(); - - let additional_properties = - schema.and_then(|schema| schema.get("additionalProperties")).and_then(|schema| { - if schema.is_object() { - Self::from_schema(schema).map(Box::new) - } else { - None - } - }); - - Self::Object { - properties, - additional_properties, + // Text spelling a literal decodes to it before any type, a string in + // particular, accepts the same text. + types.sort_by_key(|param_type| !matches!(param_type, Self::Literal(_))); + match types.len() { + 0 => None, + 1 => types.pop(), + _ => Some(Self::OneOf(types)), } } +} - /// Normalize array schema fields. - fn array_from_schema(schema: Option<&Map>) -> Self { - let items = schema - .and_then(|schema| schema.get("items")) - .and_then(Self::from_schema) - .map(Box::new); - - Self::Array { items } +impl<'a> ObjectType<'a> { + /// The type of the field named `name`, from its declared schema or the + /// additional properties' schema. + fn field(&self, name: &str) -> Option> { + let schema = self + .properties + .and_then(|properties| properties.get(name)) + .or(self.additional_properties)?; + JsonParamType::from_schema(&self.root, schema) } +} - /// Collapse a candidate type list into one normalized type. - fn one_of(mut types: Vec) -> Self { - if types.len() == 1 { - types.remove(0) - } else { - Self::OneOf(types) - } +impl<'a> ArrayType<'a> { + /// The type of the array's items. + fn items(&self) -> Option> { + JsonParamType::from_schema(&self.root, self.items?) } } @@ -318,7 +248,10 @@ fn is_null_literal(value: &str) -> bool { } /// Convert one parameter input to a normalized JSON value. -fn convert_with_optional_schema(param_type: Option<&JsonParamType>, input: &ParamInput) -> Value { +fn convert_with_optional_schema( + param_type: Option<&JsonParamType<'_>>, + input: &ParamInput, +) -> Value { // Coerce the literal text `null` to JSON null, except for `string`-typed // params, where it must stay the string "null": a model emitting the literal // text "null" for a string field means the string, not a missing value. @@ -341,13 +274,13 @@ fn convert_with_optional_schema(param_type: Option<&JsonParamType>, input: &Para ParamInput::Text(value) => Value::String(value.clone()), ParamInput::Elements(elements) => { // Convert structured input to object without a schema. - Value::Object(convert_elements_to_object(elements, &BTreeMap::new(), None)) + Value::Object(convert_elements_to_object(elements, None)) } } } /// Convert one parameter input to a normalized JSON type. -fn try_convert_value(param_type: &JsonParamType, input: &ParamInput) -> Option { +fn try_convert_value(param_type: &JsonParamType<'_>, input: &ParamInput) -> Option { match input { ParamInput::Text(value) => try_convert_text_value(param_type, value), ParamInput::Elements(elements) => try_convert_elements_value(param_type, elements), @@ -358,18 +291,21 @@ fn try_convert_value(param_type: &JsonParamType, input: &ParamInput) -> Option Option { +fn try_convert_text_value(param_type: &JsonParamType<'_>, value: &str) -> Option { match param_type { JsonParamType::String => Some(Value::String(value.to_string())), JsonParamType::Integer => convert_integer_text(value), JsonParamType::Number => convert_number_text(value), JsonParamType::Boolean => try_convert_boolean(value), - JsonParamType::Object { .. } if value.trim().is_empty() => Some(Value::Object(Map::new())), - JsonParamType::Array { .. } if value.trim().is_empty() => Some(Value::Array(Vec::new())), + JsonParamType::Object(_) if value.trim().is_empty() => Some(Value::Object(Map::new())), + JsonParamType::Array(_) if value.trim().is_empty() => Some(Value::Array(Vec::new())), // For composite types with string input, interpret the string as JSON of that type. - JsonParamType::Object { .. } => serde_json::from_str(value).ok().filter(Value::is_object), - JsonParamType::Array { .. } => serde_json::from_str(value).ok().filter(Value::is_array), + JsonParamType::Object(_) => serde_json::from_str(value).ok().filter(Value::is_object), + JsonParamType::Array(_) => serde_json::from_str(value).ok().filter(Value::is_array), JsonParamType::Null => is_null_literal(value).then_some(Value::Null), + JsonParamType::Literal(literal) => serde_json::from_str::(value.trim()) + .ok() + .filter(|value| value == *literal), JsonParamType::OneOf(types) => { types.iter().find_map(|param_type| try_convert_text_value(param_type, value)) } @@ -378,25 +314,24 @@ fn try_convert_text_value(param_type: &JsonParamType, value: &str) -> Option, elements: &[ParamElement], ) -> Option { match param_type { - JsonParamType::Object { - properties, - additional_properties, - } => Some(Value::Object(convert_elements_to_object( + JsonParamType::Object(object) => Some(Value::Object(convert_elements_to_object( elements, - properties, - additional_properties.as_deref(), + Some(object), ))), - JsonParamType::Array { items } => Some(Value::Array( - // Collect all child elements into an array, regardless of their names. - elements - .iter() - .map(|element| convert_with_optional_schema(items.as_deref(), &element.value)) - .collect(), - )), + JsonParamType::Array(array) => { + let items = array.items(); + Some(Value::Array( + // Collect all child elements into an array, regardless of their names. + elements + .iter() + .map(|element| convert_with_optional_schema(items.as_ref(), &element.value)) + .collect(), + )) + } JsonParamType::OneOf(types) => types .iter() .find_map(|param_type| try_convert_elements_value(param_type, elements)), @@ -406,20 +341,20 @@ fn try_convert_elements_value( | JsonParamType::Integer | JsonParamType::Number | JsonParamType::Boolean - | JsonParamType::Null => None, + | JsonParamType::Null + | JsonParamType::Literal(_) => None, } } /// Convert structured elements to an object, using field schemas when present. fn convert_elements_to_object( elements: &[ParamElement], - properties: &BTreeMap, - additional_properties: Option<&JsonParamType>, + object_type: Option<&ObjectType<'_>>, ) -> Map { let mut object = Map::with_capacity(elements.len()); for element in elements { - let param_type = properties.get(&element.name).or(additional_properties); - let value = convert_with_optional_schema(param_type, &element.value); + let param_type = object_type.and_then(|object_type| object_type.field(&element.name)); + let value = convert_with_optional_schema(param_type.as_ref(), &element.value); insert_object_value(&mut object, element.name.clone(), value); } object @@ -508,7 +443,7 @@ mod tests { #[test] fn invalid_schema_converts_everything_as_string() { - let params = ToolSchema::from_schema(&json!({ "type": "object" })); + let params = ToolSchema::from_schema(json!({ "type": "object" })); assert_eq!(params.convert("count", text("42")), json!("42")); assert_eq!(params.convert("count", text("null")), json!(null)); @@ -516,7 +451,7 @@ mod tests { #[test] fn skips_unknown_property_schema_and_unknown_type() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "unknown_schema": true, @@ -532,7 +467,7 @@ mod tests { #[test] fn converts_supported_types() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "text": { "type": "string" }, @@ -571,7 +506,7 @@ mod tests { #[test] fn number_conversion_preserves_json_number_spelling_with_legacy_fallback() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "value": { "type": "number" } @@ -618,7 +553,7 @@ mod tests { #[test] fn non_string_values_ignore_surrounding_whitespace() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "text": { "type": "string" }, @@ -641,7 +576,7 @@ mod tests { #[test] fn numbers_accept_digit_separators() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "count": { "type": "integer" }, @@ -658,7 +593,7 @@ mod tests { #[test] fn nullable_schemas_accept_null() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "text": { "type": "string", "nullable": true }, @@ -675,7 +610,7 @@ mod tests { #[test] fn composite_values_must_decode_to_their_type() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "payload": { "type": "object" }, @@ -695,7 +630,7 @@ mod tests { #[test] fn converts_upstream_aliases() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "s": { "type": "varchar" }, @@ -717,13 +652,13 @@ mod tests { #[test] fn preserves_union_type_order() { - let integer_first = ToolSchema::from_schema(&json!({ + let integer_first = ToolSchema::from_schema(json!({ "type": "object", "properties": { "value": { "type": ["integer", "string"] } } })); - let string_first = ToolSchema::from_schema(&json!({ + let string_first = ToolSchema::from_schema(json!({ "type": "object", "properties": { "value": { "type": ["string", "integer"] } @@ -736,7 +671,7 @@ mod tests { #[test] fn converts_composite_schemas() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "choice": { @@ -745,7 +680,7 @@ mod tests { { "type": "string" } ] }, - "fallback_object": { + "unknown_alternatives": { "oneOf": [ { "type": "mystery" } ] @@ -754,15 +689,139 @@ mod tests { })); assert_eq!(params.convert("choice", text("42")), json!(42)); + // Alternatives of unknown types leave the value unconstrained, like an + // unknown `type`. assert_eq!( - params.convert("fallback_object", text(r#"{"x":1}"#)), - json!({ "x": 1 }) + params.convert("unknown_alternatives", text(r#"{"x":1}"#)), + json!(r#"{"x":1}"#) + ); + } + + #[test] + fn resolves_references_in_pydantic_schemas() { + let params = ToolSchema::from_schema(json!({ + "$defs": { + "Place": { + "properties": { "city": { "type": "string" }, "days": { "type": "integer" } }, + "type": "object" + }, + "Count": { "type": "integer" } + }, + "type": "object", + "properties": { + "place": { "$ref": "#/$defs/Place" }, + "backup": { "anyOf": [{ "$ref": "#/$defs/Place" }, { "type": "null" }] }, + "count": { "$ref": "#/$defs/Count" } + } + })); + + assert_eq!( + params.convert("place", text(r#"{"city": "Paris", "days": 3}"#)), + json!({ "city": "Paris", "days": 3 }) ); + assert_eq!( + params.convert("backup", text(r#"{"city": "Paris"}"#)), + json!({ "city": "Paris" }) + ); + assert_eq!(params.convert("backup", text("null")), Value::Null); + assert_eq!(params.convert("count", text("3")), json!(3)); + } + + #[test] + fn recursive_references_convert_without_unbounded_expansion() { + let params = ToolSchema::from_schema(json!({ + "$defs": { + "Node": { + "type": "object", + "properties": { + "id": { "type": "integer" }, + "children": { "type": "array", "items": { "$ref": "#/$defs/Node" } } + } + } + }, + "type": "object", + "properties": { "tree": { "$ref": "#/$defs/Node" } } + })); + + assert_eq!( + params.convert( + "tree", + text(r#"{"id": 1, "children": [{"id": 2, "children": []}]}"#) + ), + json!({ "id": 1, "children": [{ "id": 2, "children": [] }] }) + ); + } + + #[test] + fn branching_recursive_references_convert_structured_values_by_level() { + // Two recursive children per node expand to 2^depth types if the schema + // is walked ahead of the value. + let params = ToolSchema::from_schema(json!({ + "$defs": { + "Tree": { + "type": "object", + "properties": { + "value": { "type": "integer" }, + "left": { "$ref": "#/$defs/Tree" }, + "right": { "$ref": "#/$defs/Tree" } + } + } + }, + "type": "object", + "properties": { "tree": { "$ref": "#/$defs/Tree" } } + })); + + let leaf = |value: &str| ParamInput::Elements(vec![elem("value", text(value))]); + let tree = ParamInput::Elements(vec![ + elem("value", text("1")), + elem( + "left", + ParamInput::Elements(vec![elem("value", text("2")), elem("right", leaf("3"))]), + ), + elem("right", leaf("4")), + ]); + assert_eq!( + params.convert("tree", tree), + json!({ + "value": 1, + "left": { "value": 2, "right": { "value": 3 } }, + "right": { "value": 4 } + }) + ); + } + + #[test] + fn mixed_enums_decode_text_spelling_a_non_string_literal_to_it() { + let params = ToolSchema::from_schema(json!({ + "type": "object", + "properties": { + "mode": { "enum": ["fast", 1] }, + "flag": { "enum": ["auto", true] } + } + })); + + assert_eq!(params.convert("mode", text("1")), json!(1)); + assert_eq!(params.convert("mode", text("fast")), json!("fast")); + // Text matching no literal keeps the schema order, which tries the + // string first. + assert_eq!(params.convert("mode", text("2")), json!("2")); + assert_eq!(params.convert("flag", text("true")), json!(true)); + assert_eq!(params.convert("flag", text("auto")), json!("auto")); + } + + #[test] + fn integer_enums_convert_to_numbers() { + let params = ToolSchema::from_schema(json!({ + "type": "object", + "properties": { "level": { "enum": [1, 2, 3] } } + })); + + assert_eq!(params.convert("level", text("2")), json!(2)); } #[test] fn infers_type_from_schema_shape_without_type() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "choice": { "enum": ["a", "b"] }, @@ -845,7 +904,7 @@ mod tests { // A `string`-typed param whose value is the literal text "null"/"NULL" // must stay a string (the original case is preserved), rather than being // coerced to JSON null. Non-string types keep coercing "null" to null. - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "name": { "type": "string" }, @@ -868,7 +927,7 @@ mod tests { // must coerce to JSON null (matching Python's `extract_types_from_schema`, // which infers `null` from the enum values), while a non-null enum keeps // "null" as a string. - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "mode": { "enum": [null, "auto"] }, diff --git a/rust/src/parser/src/tool/qwen_coder.rs b/rust/src/parser/src/tool/qwen_coder.rs index e233f4720c05..e804ce0bb255 100644 --- a/rust/src/parser/src/tool/qwen_coder.rs +++ b/rust/src/parser/src/tool/qwen_coder.rs @@ -306,7 +306,7 @@ mod tests { use serde_json::{Value, json}; use thiserror_ext::AsReport; - use super::{Qwen3CoderToolParser, ToolParser}; + use super::{Qwen3CoderToolParser, Tool, ToolParser}; use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; use crate::tool::tests::assert_tool_framing_preserves_body_whitespace; use crate::tool::{ToolParserOutput, ToolParserTestExt as _}; @@ -545,6 +545,39 @@ mod tests { ); } + #[test] + fn qwen_coder_parse_complete_resolves_schema_references() { + let tools = [Tool { + name: "plan_trip".to_string(), + description: None, + parameters: json!({ + "$defs": { + "Place": { "type": "object", "properties": { "city": { "type": "string" } } }, + "Days": { "type": "integer" } + }, + "type": "object", + "properties": { + "place": { "$ref": "#/$defs/Place" }, + "days": { "$ref": "#/$defs/Days" } + } + }), + strict: None, + defer_loading: None, + }]; + let mut parser = Qwen3CoderToolParser::new(&tools); + let output = parser + .parse_complete(&build_tool_call( + "plan_trip", + &[("place", "\n{\"city\": \"Paris\"}\n"), ("days", "\n3\n")], + )) + .unwrap(); + + assert_eq!( + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), + json!({ "place": { "city": "Paris" }, "days": 3 }) + ); + } + #[test] fn qwen_coder_streaming_extracts_single_tool_call() { let mut parser = Qwen3CoderToolParser::new(&test_tools()); diff --git a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs index 8f0f098f06d1..4754726c1b5d 100644 --- a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs +++ b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs @@ -46,12 +46,12 @@ impl ArgumentSyntax for XtmlArguments { } ParameterKey::Free => ParameterKey::Free, }; - let arguments = group_by(options, |option| option.ty) + let arguments = group_by(options, |option| xtml_type(option.ty)) .into_iter() - .filter_map(|(ty, options)| { + .filter_map(|(type_name, options)| { let values = options .into_iter() - .filter_map(|option| match ty { + .filter_map(|option| match option.ty { // The parser ends the call at its close marker, so a // string value may contain neither close marker. JsonType::String => option.raw_string(&[ARG_CLOSE, CALL_CLOSE]), @@ -59,7 +59,7 @@ impl ArgumentSyntax for XtmlArguments { }) .collect::>(); (!values.is_empty()).then(|| { - let suffix = format!("\" type=\"{}\"{SEP}", ty.name()); + let suffix = format!("\" type=\"{type_name}\"{SEP}"); key.tag( &format!("{ARG_OPEN} key=\""), &suffix, @@ -73,6 +73,20 @@ impl ArgumentSyntax for XtmlArguments { } } +/// The `type` attribute K3 renders for a value of type `ty`, as +/// `_xtml_type` in the checkpoint's `encoding_k3.py`: integers and floats are +/// both `number`. +fn xtml_type(ty: JsonType) -> &'static str { + match ty { + JsonType::String => "string", + JsonType::Integer | JsonType::Number => "number", + JsonType::Boolean => "boolean", + JsonType::Null => "null", + JsonType::Object => "object", + JsonType::Array => "array", + } +} + /// Kimi K3 XTML structural-tag builder. #[derive(Debug, Clone, Copy, Default)] pub struct KimiK3StructuralTagBuilder; From 42ecc8eebe701faf07e7d6f33dbf5a8218825b5b Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Wed, 7 Oct 2026 16:29:09 +0800 Subject: [PATCH 007/624] [CI][RL] Consolidate RL entrypoint tests under tests/entrypoints/rl (#59948) Signed-off-by: wang.yuqi Signed-off-by: aoshen02 Co-authored-by: wang.yuqi Co-authored-by: Claude Opus 5.5 --- .buildkite/intel_jobs/entrypoints_intel.yaml | 13 ++++++---- .buildkite/test-amd.yaml | 6 +++-- .buildkite/test_areas/entrypoints.yaml | 2 -- .buildkite/test_areas/rl.yaml | 25 +++++++++++++++++++ .buildkite/test_areas/rust_frontend.yaml | 2 -- .github/CODEOWNERS | 3 +-- .../{serve/dev/rlhf => rl}/__init__.py | 0 .../{serve/dev/rlhf => rl}/conftest.py | 0 .../return_routed_experts}/__init__.py | 0 .../return_routed_experts/test_openai_api.py} | 2 +- .../test_scale_out_api.py} | 0 .../rl/state_transitions/__init__.py | 0 .../state_transitions/test_pause_resume.py | 2 +- .../state_transitions}/test_sleep.py | 0 .../state_transitions/test_weight_checker.py | 2 +- .../test_weight_operation_metrics.py | 2 +- .../{ => rl}/weight_transfer/__init__.py | 0 .../weight_transfer/test_offline.py} | 3 +-- .../test_serving_multimodal_tokens.py | 4 +-- tools/pre_commit/mypy.py | 2 +- 20 files changed, 46 insertions(+), 22 deletions(-) create mode 100644 .buildkite/test_areas/rl.yaml rename tests/entrypoints/{serve/dev/rlhf => rl}/__init__.py (100%) rename tests/entrypoints/{serve/dev/rlhf => rl}/conftest.py (100%) rename tests/entrypoints/{serve/dev/rlhf/state_transitions => rl/return_routed_experts}/__init__.py (100%) rename tests/entrypoints/{openai/test_return_routed_experts.py => rl/return_routed_experts/test_openai_api.py} (98%) rename tests/entrypoints/{scale_out/token_in_token_out/test_return_routed_experts.py => rl/return_routed_experts/test_scale_out_api.py} (100%) create mode 100644 tests/entrypoints/rl/state_transitions/__init__.py rename tests/entrypoints/{serve/dev/rlhf => rl}/state_transitions/test_pause_resume.py (98%) rename tests/entrypoints/{serve/dev => rl/state_transitions}/test_sleep.py (100%) rename tests/entrypoints/{serve/dev/rlhf => rl}/state_transitions/test_weight_checker.py (98%) rename tests/entrypoints/{serve/dev/rlhf => rl}/test_weight_operation_metrics.py (97%) rename tests/entrypoints/{ => rl}/weight_transfer/__init__.py (100%) rename tests/entrypoints/{weight_transfer/test_weight_transfer_llm.py => rl/weight_transfer/test_offline.py} (99%) diff --git a/.buildkite/intel_jobs/entrypoints_intel.yaml b/.buildkite/intel_jobs/entrypoints_intel.yaml index 7e823e024563..3cc9c1422d9e 100644 --- a/.buildkite/intel_jobs/entrypoints_intel.yaml +++ b/.buildkite/intel_jobs/entrypoints_intel.yaml @@ -19,14 +19,14 @@ steps: source_file_dependencies: - vllm/entrypoints - tests/entrypoints/unit_tests - - tests/entrypoints/weight_transfer + - tests/entrypoints/rl/weight_transfer - tests/entrypoints/launchers commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pytest -v -s entrypoints/unit_tests && - pytest -v -s entrypoints/weight_transfer' + pytest -v -s entrypoints/rl/weight_transfer' - label: ":intel: (BMG) Entrypoints Integration (LLM)" key: entrypoints-integration-llm-intel @@ -71,14 +71,17 @@ steps: - vllm/ - "!vllm/distributed/kv_transfer/" - tests/entrypoints/serve + - tests/entrypoints/rl commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'export VLLM_WORKER_MULTIPROC_METHOD=spawn && cd tests && - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc - --ignore=entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py - --deselect="tests/entrypoints/serve/dev/rlhf/state_transitions/test_pause_resume.py::TestPauseResume::test_clear_cache_preserves_output_and_controls_prefix_cache[MRV1]" && + pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc && + pytest -v -s entrypoints/rl --ignore=entrypoints/rl/weight_transfer + --ignore=entrypoints/rl/return_routed_experts/test_openai_api.py + --ignore=entrypoints/rl/test_weight_operation_metrics.py + --deselect="tests/entrypoints/rl/state_transitions/test_pause_resume.py::TestPauseResume::test_clear_cache_preserves_output_and_controls_prefix_cache[MRV1]" && PYTHONPATH=/workspace/vllm pytest -v -s entrypoints/serve/dev/rpc' - label: ":intel: (BMG) Entrypoints Integration (API Server Scale Out)" diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 17fe9c0ae58e..8317919ce4b8 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1044,12 +1044,12 @@ steps: source_file_dependencies: - vllm/entrypoints - tests/entrypoints/unit_tests - - tests/entrypoints/weight_transfer + - tests/entrypoints/rl/weight_transfer - tests/entrypoints/launchers - vllm/platforms/rocm.py commands: - pytest -v -s entrypoints/unit_tests - - pytest -v -s entrypoints/weight_transfer + - pytest -v -s entrypoints/rl/weight_transfer - pytest -v -s entrypoints/launchers - label: ":amd: (MI300) Entrypoints Integration (Responses API)" # TBD @@ -3746,11 +3746,13 @@ steps: - vllm/ - tests/entrypoints/serve - tests/entrypoints/scale_out + - tests/entrypoints/rl commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc; fi - pytest -v -s entrypoints/scale_out --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - pytest -v -s entrypoints/rl --ignore=entrypoints/rl/weight_transfer --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - label: ":amd: (MI355 DPX) Entrypoints Integration (OpenAI API completion)" # TBD timeout_in_minutes: 65 diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 5f8e1d1c352c..68c6243d479b 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -31,11 +31,9 @@ steps: source_file_dependencies: - vllm/entrypoints - tests/entrypoints/unit_tests - - tests/entrypoints/weight_transfer - tests/entrypoints/launchers commands: - pytest -v -s entrypoints/unit_tests - - pytest -v -s entrypoints/weight_transfer - pytest -v -s entrypoints/launchers mirror: amd: diff --git a/.buildkite/test_areas/rl.yaml b/.buildkite/test_areas/rl.yaml new file mode 100644 index 000000000000..7141a81bde1a --- /dev/null +++ b/.buildkite/test_areas/rl.yaml @@ -0,0 +1,25 @@ +group: RL +depends_on: + - image-build +steps: +- label: ":nvidia: (H200 MIG 35GB) RL Entrypoints" + device: h200_35gb + key: rl-entrypoints-tests + timeout_in_minutes: 40 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - "!vllm/distributed/kv_transfer/" + - tests/entrypoints/rl + commands: + - pytest -v -s entrypoints/rl/weight_transfer + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/rl --ignore=entrypoints/rl/weight_transfer + mirror: + amd: + label: ":amd: (MI355 DPX) RL Entrypoints" + dind: false + device: mi355_dpx + timeout_in_minutes: 50 + depends_on: + - image-build-amd diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index c5899a9c59bf..59c78525c5e7 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -82,7 +82,6 @@ steps: - tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py - tests/entrypoints/serve/instrumentator/test_basic.py - tests/entrypoints/serve/instrumentator/test_metrics.py - # - tests/entrypoints/serve/dev/test_sleep.py - tests/entrypoints/serve/tokenize/test_tokenization.py commands: - export VLLM_USE_RUST_FRONTEND=1 @@ -93,7 +92,6 @@ steps: # test_generate_logprobs expects Python-style top_logprobs truncation (dedup sampled + cap at max(k, 1)). - pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow" - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist" - # - pytest -v -s entrypoints/serve/dev/test_sleep.py # /tokenizer_info is not implemented in the Rust frontend (the CLI flag is accepted as a no-op). - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info" mirror: diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 71d501eabe7c..083842c932ff 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -231,8 +231,7 @@ mkdocs.yaml @hmellor /vllm/distributed/weight_transfer @aoshen02 /vllm/entrypoints/serve/dev/rlhf @njhill @aoshen02 /tests/distributed/test_weight_transfer.py @aoshen02 -/tests/entrypoints/weight_transfer @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche @AndreasKaratzas @aoshen02 -/tests/entrypoints/serve/dev/rlhf @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche @AndreasKaratzas @aoshen02 +/tests/entrypoints/rl @DarkLight1337 @robertgshaw2-redhat @njhill @aarnphm @NickLucche @AndreasKaratzas @aoshen02 /tests/v1/worker/test_gpu_worker_weight_transfer.py @aoshen02 /examples/rl @aoshen02 /docs/training/rlhf.md @aoshen02 diff --git a/tests/entrypoints/serve/dev/rlhf/__init__.py b/tests/entrypoints/rl/__init__.py similarity index 100% rename from tests/entrypoints/serve/dev/rlhf/__init__.py rename to tests/entrypoints/rl/__init__.py diff --git a/tests/entrypoints/serve/dev/rlhf/conftest.py b/tests/entrypoints/rl/conftest.py similarity index 100% rename from tests/entrypoints/serve/dev/rlhf/conftest.py rename to tests/entrypoints/rl/conftest.py diff --git a/tests/entrypoints/serve/dev/rlhf/state_transitions/__init__.py b/tests/entrypoints/rl/return_routed_experts/__init__.py similarity index 100% rename from tests/entrypoints/serve/dev/rlhf/state_transitions/__init__.py rename to tests/entrypoints/rl/return_routed_experts/__init__.py diff --git a/tests/entrypoints/openai/test_return_routed_experts.py b/tests/entrypoints/rl/return_routed_experts/test_openai_api.py similarity index 98% rename from tests/entrypoints/openai/test_return_routed_experts.py rename to tests/entrypoints/rl/return_routed_experts/test_openai_api.py index 46eeed179531..f4fc4013997a 100644 --- a/tests/entrypoints/openai/test_return_routed_experts.py +++ b/tests/entrypoints/rl/return_routed_experts/test_openai_api.py @@ -7,7 +7,7 @@ import pybase64 as base64 import pytest -from ...utils import RemoteOpenAIServer +from tests.utils import RemoteOpenAIServer MODEL_NAME = "TitanML/tiny-mixtral" diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py b/tests/entrypoints/rl/return_routed_experts/test_scale_out_api.py similarity index 100% rename from tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py rename to tests/entrypoints/rl/return_routed_experts/test_scale_out_api.py diff --git a/tests/entrypoints/rl/state_transitions/__init__.py b/tests/entrypoints/rl/state_transitions/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/entrypoints/serve/dev/rlhf/state_transitions/test_pause_resume.py b/tests/entrypoints/rl/state_transitions/test_pause_resume.py similarity index 98% rename from tests/entrypoints/serve/dev/rlhf/state_transitions/test_pause_resume.py rename to tests/entrypoints/rl/state_transitions/test_pause_resume.py index 3de9e0054963..587b4b73194a 100644 --- a/tests/entrypoints/serve/dev/rlhf/state_transitions/test_pause_resume.py +++ b/tests/entrypoints/rl/state_transitions/test_pause_resume.py @@ -10,7 +10,7 @@ import pytest import requests -from tests.entrypoints.serve.dev.rlhf.conftest import ( +from tests.entrypoints.rl.conftest import ( cached_tokens, completion_with_cache_details, gen, diff --git a/tests/entrypoints/serve/dev/test_sleep.py b/tests/entrypoints/rl/state_transitions/test_sleep.py similarity index 100% rename from tests/entrypoints/serve/dev/test_sleep.py rename to tests/entrypoints/rl/state_transitions/test_sleep.py diff --git a/tests/entrypoints/serve/dev/rlhf/state_transitions/test_weight_checker.py b/tests/entrypoints/rl/state_transitions/test_weight_checker.py similarity index 98% rename from tests/entrypoints/serve/dev/rlhf/state_transitions/test_weight_checker.py rename to tests/entrypoints/rl/state_transitions/test_weight_checker.py index 7cb58df9d902..ec3c8016a870 100644 --- a/tests/entrypoints/serve/dev/rlhf/state_transitions/test_weight_checker.py +++ b/tests/entrypoints/rl/state_transitions/test_weight_checker.py @@ -9,7 +9,7 @@ import torch from transformers import AutoModelForCausalLM -from tests.entrypoints.serve.dev.rlhf.conftest import ( +from tests.entrypoints.rl.conftest import ( MODEL_NAME, gen, health, diff --git a/tests/entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py b/tests/entrypoints/rl/test_weight_operation_metrics.py similarity index 97% rename from tests/entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py rename to tests/entrypoints/rl/test_weight_operation_metrics.py index d7758011f925..108dd77b2f1b 100644 --- a/tests/entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py +++ b/tests/entrypoints/rl/test_weight_operation_metrics.py @@ -7,7 +7,7 @@ from prometheus_client.parser import text_string_to_metric_families from transformers import AutoModelForCausalLM -from tests.entrypoints.serve.dev.rlhf.conftest import MODEL_NAME +from tests.entrypoints.rl.conftest import MODEL_NAME from tests.utils import RemoteOpenAIServer from vllm.distributed.weight_transfer import ( HTTPVLLMWeightSyncClient, diff --git a/tests/entrypoints/weight_transfer/__init__.py b/tests/entrypoints/rl/weight_transfer/__init__.py similarity index 100% rename from tests/entrypoints/weight_transfer/__init__.py rename to tests/entrypoints/rl/weight_transfer/__init__.py diff --git a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py b/tests/entrypoints/rl/weight_transfer/test_offline.py similarity index 99% rename from tests/entrypoints/weight_transfer/test_weight_transfer_llm.py rename to tests/entrypoints/rl/weight_transfer/test_offline.py index e5c3616d6e65..e4a30f17e61d 100644 --- a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py +++ b/tests/entrypoints/rl/weight_transfer/test_offline.py @@ -15,6 +15,7 @@ import pytest import torch +from tests.utils import create_new_process_for_each_test from vllm.config import WeightTransferConfig from vllm.distributed.weight_transfer.base import ( WeightTransferEngine, @@ -24,8 +25,6 @@ WeightTransferUpdateRequest, ) -from ...utils import create_new_process_for_each_test - # Use a tiny model for fast testing MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM" diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py index 1d017bd3cf1b..2e58d70972f9 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py @@ -5,8 +5,8 @@ Mirrors test_serving_tokens.py but exercises the multimodal piping using Qwen/Qwen3-VL-2B-Instruct end-to-end via the server's /render -> /generate -> /detokenize path. Intentionally avoids running the HF -processor in the pytest parent process to keep os.fork() in sibling -tests (e.g. test_weight_transfer_llm.py) deadlock-free. +processor in the pytest parent process to keep os.fork() in other +tests (e.g. entrypoints/rl/weight_transfer/test_offline.py) deadlock-free. """ import httpx diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py index a6c8e4641963..245d0daf5e18 100755 --- a/tools/pre_commit/mypy.py +++ b/tools/pre_commit/mypy.py @@ -38,10 +38,10 @@ "tests/entrypoints/generate", "tests/entrypoints/multimodal", "tests/entrypoints/pooling", + "tests/entrypoints/rl", "tests/entrypoints/serve", "tests/entrypoints/tool_parsers", "tests/entrypoints/unit_tests", - "tests/entrypoints/weight_transfer", "tests/kernels/attention", "tests/kernels/core", "tests/kernels/helion", From df417f780a19b68cebc4379942de6ea0e6f6aceb Mon Sep 17 00:00:00 2001 From: Hexiang Wang <56632993+whx-sjtu@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:35:56 +0800 Subject: [PATCH 008/624] [Bugfix][MoRIIO] Exclude synchronous READ destinations from KV zeroing (#59164) Signed-off-by: Yichao Zhu Signed-off-by: whx-sjtu Signed-off-by: Hexiang Wang <56632993+whx-sjtu@users.noreply.github.com> Co-authored-by: Yichao Zhu Co-authored-by: Codex Co-authored-by: Yifan Qiao --- .../unit/test_moriio_hma_scheduler.py | 70 ++++++++++++------- .../kv_transfer/kv_connector/v1/base.py | 6 +- .../v1/moriio/moriio_connector.py | 18 +++++ 3 files changed, 67 insertions(+), 27 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py index 28bc8bebe261..1de729416e90 100644 --- a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py +++ b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py @@ -7,32 +7,20 @@ GPU or the ``mori`` runtime: per-group block-id splitting, the READ/WRITE ``N-1`` token accounting, P-side prompt truncation, and offset-template cache wiring. - -Like ``test_moriio_kv_layout.py`` the whole module is skipped unless it is -running on ROCm with ``mori`` installed (importing the connector pulls in -``mori``). The authoritative run happens on the MIA recipe image. """ import importlib -import importlib.util +from dataclasses import asdict from types import SimpleNamespace import pytest import torch -from vllm.platforms import current_platform from vllm.v1.core.block_pool import BlockPool +from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.single_type_kv_cache_manager import FullAttentionManager from vllm.v1.kv_cache_interface import FullAttentionSpec, UniformTypeKVCacheSpecs -mori_available = importlib.util.find_spec("mori") is not None - -if not (current_platform.is_rocm() and mori_available): - pytest.skip( - "MoRIIOs are only available on ROCm with mori package installed", - allow_module_level=True, - ) - moriio_connector = importlib.import_module( "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector" ) @@ -525,6 +513,8 @@ def _make_read_scheduler(): request_id_to_transfer_id={}, transfer_id_to_request_id={}, _reqs_need_recv={}, + _reqs_need_save={}, + _reqs_need_send={}, _req_kv_params={}, _max_decode_tail_blocks=1, ) @@ -538,24 +528,52 @@ def _make_read_request(remote_block_ids): "transfer_id": "tx", "remote_engine_id": "prefill", "remote_block_ids": remote_block_ids, + "remote_host": "127.0.0.1", + "remote_handshake_port": 6000, + "remote_notify_port": 6001, }, ) -def test_update_state_drops_decode_recompute_tail_block(): +@pytest.mark.parametrize("external_tokens", [0, 256]) +@pytest.mark.parametrize("has_mamba", [False, True]) +@pytest.mark.parametrize("zero_ids", [None, [], [101], [101, 102, 202, 999]]) +def test_update_state_drops_decode_recompute_tail_block( + external_tokens, has_mamba, zero_ids +): + """Only full hybrid READ destinations skip zeroing; local tails still zero.""" sched = _make_read_scheduler() - request = _make_read_request([[10, 11], [91]]) + sched._has_mamba = has_mamba + remote_blocks = [[10, 11], [91]] if has_mamba else [[10, 11]] + request = _make_read_request(remote_blocks) blocks = _FakeBlocks( - all_groups=([100, 101, 102], [200, 201, 202]), - ) - - sched.update_state_after_alloc(request, blocks, num_external_tokens=256) - - assert sched._reqs_need_recv["req"][1] == [[100, 101], [202]] - assert sched._req_kv_params["req"]["remote_block_ids"] == [ - [10, 11], - [91], - ] + all_groups=([100, 101, 102], [200, 201, 202]) + if has_mamba + else ([100, 101, 102],), + ) + + sched.update_state_after_alloc(request, blocks, num_external_tokens=external_tokens) + + expected_blocks = [[100, 101] if external_tokens else []] + if has_mamba: + expected_blocks.append([202]) + assert sched._reqs_need_recv["req"][1] == expected_blocks + assert sched._req_kv_params["req"]["remote_block_ids"] == remote_blocks + + output = SchedulerOutput.make_empty() + output.new_block_ids_to_zero = None if zero_ids is None else list(zero_ids) + before = asdict(output) + meta = sched.build_connector_meta(output) + assert meta.reqs_to_recv["req"].local_block_ids == expected_blocks + if zero_ids and has_mamba and external_tokens: + before["new_block_ids_to_zero"] = zero_ids[1:] + assert asdict(output) == before + + # A later local allocation of the same pages must still be zeroed. + output.new_block_ids_to_zero = [100, 101, 102] + meta = sched.build_connector_meta(output) + assert not meta.reqs_to_recv + assert output.new_block_ids_to_zero == [100, 101, 102] def test_update_state_pairs_shorter_local_blocks_with_remote_suffix(): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index bb9d4597ce3c..71593f59e55c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -583,7 +583,11 @@ def build_connector_meta( """Build the connector metadata for this step. This function should NOT modify fields in the scheduler_output. - Also, calling this function will reset the state of the connector. + FIXME: one exception: + synchronous READ connectors may remove attention blocks they fully + overwrite this step from new_block_ids_to_zero. They must complete + those loads before the blocks are used and fail the step if a load fails. + Calling this function will reset the state of the connector. Args: scheduler_output (SchedulerOutput): the scheduler output object. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 167c7ef3c299..5a8a9176640e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -1232,6 +1232,24 @@ def build_connector_meta( local_block_ids=block_ids, kv_transfer_params=kv_params, ) + if ( + self.mode == MoRIIOMode.READ + and self._has_mamba + and scheduler_output.new_block_ids_to_zero + ): + # Hybrid models zero recycled attention pages that held Mamba state. + # Host-submitted READs overwrite these pages and can race zeroing. + # Hybrid READ metadata puts the aligned attention pages first. + read_dst_block_ids = { + b + for _, block_ids in self._reqs_need_recv.values() + for b in block_ids[0] + } + scheduler_output.new_block_ids_to_zero = [ + b + for b in scheduler_output.new_block_ids_to_zero + if b not in read_dst_block_ids + ] for req_id, (req, block_ids) in self._reqs_need_save.items(): kv_params = self._req_kv_params.get(req_id, req.kv_transfer_params or {}) From 02d8a94e610f515449d3bc3f44c592e3022e19d8 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Wed, 7 Oct 2026 02:38:31 -0700 Subject: [PATCH 009/624] [CI][ROCm] Skip DSv4.1 decoder replay graph test on ROCm (#60378) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- tests/models/test_deepseek_v41_decoder_replay_layers.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/models/test_deepseek_v41_decoder_replay_layers.py b/tests/models/test_deepseek_v41_decoder_replay_layers.py index 9ba2df5c8f4c..1302e0acda30 100644 --- a/tests/models/test_deepseek_v41_decoder_replay_layers.py +++ b/tests/models/test_deepseek_v41_decoder_replay_layers.py @@ -21,6 +21,7 @@ from vllm.models.deepseek_v41.nvidia.decoder_replay_cudagraph import ( DecoderReplayCudaGraphManager, ) +from vllm.platforms import current_platform from vllm.v1.worker.gpu import cudagraph_utils pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") @@ -125,6 +126,9 @@ def _graph_context(num_tokens): ) +@pytest.mark.skipif( + current_platform.is_rocm(), reason="decoder replay CUDA graphs are NVIDIA-only" +) def test_replay_graph_matches_eager(monkeypatch): """The replay graph of the next captured size pads and runs the rows.""" monkeypatch.setattr(cudagraph_utils, "get_pp_group", MagicMock) From 6b75dfb8a101dcf23c9e2b9e627433c61030cbcf Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Wed, 7 Oct 2026 18:53:17 +0900 Subject: [PATCH 010/624] [Misc] Report a missing DeepSelect extension only when it is requested (#60242) Signed-off-by: lesj0610 --- vllm/model_executor/layers/indexer_topk.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/layers/indexer_topk.py b/vllm/model_executor/layers/indexer_topk.py index f0da0312e127..02c4280ea0a1 100644 --- a/vllm/model_executor/layers/indexer_topk.py +++ b/vllm/model_executor/layers/indexer_topk.py @@ -29,12 +29,10 @@ try: import vllm._deepselect_C # noqa: F401 (registers torch.ops.deep_select) -except ImportError as e: - from vllm.logger import init_logger - init_logger(__name__).warning( - "Failed to import the DeepSelect extension (vllm._deepselect_C): %s", e - ) + _HAS_DEEP_SELECT_C = True +except ImportError: + _HAS_DEEP_SELECT_C = False @functools.lru_cache(maxsize=1) @@ -196,6 +194,8 @@ def resolve_backend( failures.append("requires a CUDA platform") elif not self._has_deep_select: failures.append("requires SM100a/SM103a (10.x device family)") + elif not _HAS_DEEP_SELECT_C: + failures.append("vllm._deepselect_C could not be imported") elif not is_deep_select_supported(logits, topk_tokens): failures.append( f"inputs violate DeepSelect's constraints: dtype={logits.dtype}," From 3ca00a8261c7790ba7487e2001ab9004ad8065f2 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Wed, 7 Oct 2026 12:22:31 +0200 Subject: [PATCH 011/624] [CI] Skip EmbeddingGemma2 pooling tests needing transformers >= 5.19.0 (#60371) Signed-off-by: Stefan Koncarevic --- .../multimodal/pooling/test_embedding_gemma2.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/models/multimodal/pooling/test_embedding_gemma2.py b/tests/models/multimodal/pooling/test_embedding_gemma2.py index 701d0914ddbc..7bb69fa69ad8 100644 --- a/tests/models/multimodal/pooling/test_embedding_gemma2.py +++ b/tests/models/multimodal/pooling/test_embedding_gemma2.py @@ -9,11 +9,20 @@ import pytest import torch +from tests.models.registry import HF_EXAMPLE_MODELS + MODEL = os.environ.get("EG2_MODEL_PATH") needs_model = pytest.mark.skipif(not MODEL, reason="EG2_MODEL_PATH not set") TEXT_COS, MM_COS = 0.999, 0.999 +@pytest.fixture +def needs_hf_embedding_gemma2(): + """Skip when the installed transformers has no `embedding_gemma2` module.""" + model_info = HF_EXAMPLE_MODELS.get_hf_info("EmbeddingGemma2Model") + model_info.check_transformers_version(on_fail="skip") + + def _cos(a, b): return torch.nn.functional.cosine_similarity( torch.tensor(a, dtype=torch.float32), @@ -239,7 +248,7 @@ def nparams(m): ) -def test_encoder_contract(): +def test_encoder_contract(needs_hf_embedding_gemma2): from vllm.model_executor.models.embedding_gemma2 import EmbeddingGemma2Model from vllm.model_executor.models.gemma4_mm import ( Gemma4ForConditionalGeneration as G, @@ -423,7 +432,7 @@ def test_video_loader_duplicate_expansion_e2e(vllm_runner, st_model): assert len(out[0]) == 768 -def test_video_loader_indices(): +def test_video_loader_indices(needs_hf_embedding_gemma2): from transformers.models.embedding_gemma2.video_processing_embedding_gemma2 import ( EmbeddingGemma2VideoProcessor, ) @@ -508,7 +517,7 @@ def test_video_loader_indices(): assert v_4h == h_4h, f"1fps/4h mismatch: {v_4h[-2:]} vs {h_4h[-2:]}" -def test_video_loader_duplicate_expansion(): +def test_video_loader_duplicate_expansion(needs_hf_embedding_gemma2): from transformers.models.embedding_gemma2.video_processing_embedding_gemma2 import ( EmbeddingGemma2VideoProcessor, ) From d9503dc503e75193d2128908b59c0d141d4b5fb1 Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Wed, 7 Oct 2026 21:24:40 +0900 Subject: [PATCH 012/624] [Bugfix][Qwen4Exp] Derive the QSA head_dim from the unsharded head count (#59945) Signed-off-by: lesj0610 Co-authored-by: lesj0610 Co-authored-by: Thien Tran --- vllm/models/qwen4_exp/amd/qsa.py | 2 +- vllm/models/qwen4_exp/nvidia/qsa.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/models/qwen4_exp/amd/qsa.py b/vllm/models/qwen4_exp/amd/qsa.py index 325e0e3d917c..3f03b293b4f0 100644 --- a/vllm/models/qwen4_exp/amd/qsa.py +++ b/vllm/models/qwen4_exp/amd/qsa.py @@ -214,7 +214,7 @@ def __init__( elif tp_size % self.total_num_kv_heads: raise ValueError("TP size must be divisible by replicated QSA KV heads") self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.head_dim = int(config.head_dim or self.hidden_size // self.num_heads) + self.head_dim = int(config.head_dim or self.hidden_size // self.total_num_heads) self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5 diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index a20fcdeec98d..0246db80943d 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -393,7 +393,7 @@ def __init__( elif tp_size % self.total_num_kv_heads: raise ValueError("TP size must be divisible by replicated QSA KV heads") self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.head_dim = int(config.head_dim or self.hidden_size // self.num_heads) + self.head_dim = int(config.head_dim or self.hidden_size // self.total_num_heads) self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5 From 2a54f6b625f28109180b072b704c0b0d372a277d Mon Sep 17 00:00:00 2001 From: Jin Tao Date: Wed, 7 Oct 2026 15:53:40 +0300 Subject: [PATCH 013/624] [ROCm][Perf][GLM-5.3-Flash] Enable AITER fused shared experts on gfx950 (#59221) Signed-off-by: Jin Tao Co-authored-by: Cursor Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> --- .../layers/test_fused_shared_expert.py | 119 +++++++++++++++++- vllm/models/glm5next/common/model.py | 85 ++++++++++++- vllm/models/glm5next/common/mtp.py | 15 ++- 3 files changed, 213 insertions(+), 6 deletions(-) diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py index 6721a4c63b56..d8d5c4fb032a 100644 --- a/tests/model_executor/layers/test_fused_shared_expert.py +++ b/tests/model_executor/layers/test_fused_shared_expert.py @@ -14,7 +14,7 @@ from torch import nn import vllm.config as vllm_config_module -from vllm.config import VllmConfig, set_current_vllm_config +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe import utils as fused_moe_utils from vllm.model_executor.layers.fused_moe.layer import determine_expert_counts from vllm.model_executor.layers.quantization.base_config import QuantizationConfig @@ -145,6 +145,12 @@ def _stub_quant_config() -> QuantizationConfig: return cast(QuantizationConfig, object()) +class _StubGlm5NextAttention(nn.Module): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__() + self.o_proj = SimpleNamespace(reduce_results=True) + + def get_fse_test_model_config( model_type: str, quantization_config: dict[str, Any], @@ -279,6 +285,33 @@ def get_fse_test_model_config( quantization_config=quantization_config, ) return config, Glm4MoeModel + if model_type == "glm5_next": + from transformers import Glm5NextTextConfig + + from vllm.models.glm5next.common.model import Glm5NextModel + + config = Glm5NextTextConfig( + vocab_size=256, + hidden_size=128, + intermediate_size=32, + moe_intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=1, + num_key_value_heads=1, + n_routed_experts=2, + n_shared_experts=1, + num_experts_per_tok=1, + layer_types=["linear_attention"], + mlp_layer_types=["sparse"], + index_topk=1, + index_kpool=1, + pad_token_id=None, + scoring_func="sigmoid", + topk_method="noaux_tc", + mhc=False, + quantization_config=quantization_config, + ) + return config, Glm5NextModel raise ValueError(f"Unsupported FSE test model: {model_type}") @@ -796,7 +829,7 @@ def __init__(self, enabled: bool) -> None: @pytest.mark.parametrize( "model_type", - ["minimax_m3", "deepseek_v4", "qwen3_5", "glm4_moe", "deepseek_v2"], + ["minimax_m3", "deepseek_v4", "qwen3_5", "glm4_moe", "deepseek_v2", "glm5_next"], ) @pytest.mark.parametrize( ("use_fse", "exclude"), @@ -850,6 +883,16 @@ def test_models_fse_init( ) else: vllm_config.quant_config = QuarkConfig(quantization_config) + if model_type == "glm5_next": + from vllm.models.glm5next.common import model as glm5_next_model + + for attention in ("Glm5NextLinearAttention", "Glm5NextMLAAttention"): + monkeypatch.setattr(glm5_next_model, attention, _StubGlm5NextAttention) + # The gate allows fusion only at TP4/TP8 on gfx950, and this test + # runs at TP1; the gate has its own test. + monkeypatch.setattr( + glm5_next_model, "_fused_shared_experts_tuned", lambda _: True + ) import vllm.envs as envs from vllm._aiter_ops import rocm_aiter_ops @@ -902,6 +945,10 @@ def test_models_fse_init( ) mtp = DeepSeekV4MTP(vllm_config=vllm_config) assert model.is_fused_shared_expert_enabled is (fse_enabled and not exclude) + if model_type == "glm5_next": + assert (model.layers[0].mlp.shared_experts is None) is ( + fse_enabled and not exclude + ) # The dummy quant config here uses mixed mxfp4/fp8 for experts/shared_expert # so should just raise a warning. @@ -926,6 +973,74 @@ def test_models_fse_init( rocm_aiter_ops.refresh_env_variables() +@pytest.mark.parametrize( + ("setting", "value", "reason"), + [ + ("tensor_parallel_size", 4, None), + ("tensor_parallel_size", 8, None), + ("tensor_parallel_size", 2, "tensor_parallel_size is 2"), + ("data_parallel_size", 2, "data_parallel_size is 2"), + ("prefill_context_parallel_size", 2, "prefill_context_parallel_size is 2"), + ("enable_expert_parallel", True, "expert parallelism is enabled"), + ("on_gfx950", False, "the GPU is not gfx950"), + ], +) +def test_glm5_next_fuses_shared_experts_only_in_tuned_setups( + monkeypatch: pytest.MonkeyPatch, + setting: str, + value: object, + reason: str | None, +) -> None: + from vllm.models.glm5next.common import model as glm5_next_model + + parallel_config = SimpleNamespace( + tensor_parallel_size=4, + data_parallel_size=1, + prefill_context_parallel_size=1, + enable_expert_parallel=False, + ) + on_gfx950 = value if setting == "on_gfx950" else True + if setting != "on_gfx950": + setattr(parallel_config, setting, value) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: on_gfx950) + + with patch.object(glm5_next_model.logger, "warning_once") as warning: + tuned = glm5_next_model._fused_shared_experts_tuned( + cast(ParallelConfig, parallel_config) + ) + + assert tuned is (reason is None) + if reason is None: + warning.assert_not_called() + else: + warning.assert_called_once() + assert warning.call_args.args[1] == reason + + +def test_glm5_next_fused_shared_expert_loads_into_the_shared_slot() -> None: + from vllm.models.glm5next.common import model as glm5_next_model + + assert glm5_next_model._num_fused_shared_experts(1, True) == 1 + assert glm5_next_model._num_fused_shared_experts(1, False) == 0 + assert glm5_next_model._num_fused_shared_experts(None, True) == 0 + with pytest.raises(NotImplementedError, match="only 1 shared expert"): + glm5_next_model._num_fused_shared_experts(2, True) + + for shared, fused in [ + ( + "model.layers.3.mlp.shared_experts.gate_proj.weight_scale_inv", + "model.layers.3.mlp.experts.288.gate_proj.weight_scale_inv", + ), + ( + "model.layers.45.mtp_block.mlp.shared_experts.down_proj.weight", + "model.layers.45.mtp_block.mlp.experts.288.down_proj.weight", + ), + ]: + assert glm5_next_model._fused_shared_expert_name(shared, 288) == fused + routed = "model.layers.3.mlp.experts.7.down_proj.weight" + assert glm5_next_model._fused_shared_expert_name(routed, 288) == routed + + @pytest.mark.parametrize( ("exclude", "expected"), [([], True), ([r"re:.*\.shared_expert\..*"], False)], diff --git a/vllm/models/glm5next/common/model.py b/vllm/models/glm5next/common/model.py index e15c00ecb88b..778738c13712 100644 --- a/vllm/models/glm5next/common/model.py +++ b/vllm/models/glm5next/common/model.py @@ -23,6 +23,10 @@ GateLinear, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, + resolve_layer_fused_shared_expert, +) from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( @@ -258,7 +262,12 @@ def __init__( ) swiglu_limit = config.swiglu_limit - if config.n_shared_experts is None: + self.is_fused_shared_expert_enabled = False + if config.n_shared_experts is not None: + self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert( + quant_config, prefix + ) and _fused_shared_experts_tuned(parallel_config) + if config.n_shared_experts is None or self.is_fused_shared_expert_enabled: self.shared_experts = None else: intermediate_size = config.moe_intermediate_size * config.n_shared_experts @@ -294,7 +303,10 @@ def __init__( enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, - n_shared_experts=None, + n_shared_experts=config.n_shared_experts + if self.is_fused_shared_expert_enabled + else None, + fuse_shared_experts=self.is_fused_shared_expert_enabled, router_logits_dtype=self.gate.out_dtype, swiglu_limit=swiglu_limit, ) @@ -733,6 +745,9 @@ def get_layer(prefix: str): # The active slice is fixed after construction; cache it so forward # doesn't rebuild the slice (a fresh list) every step. self._active_layers = self.layers[self.start_layer : self.end_layer] + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.layers, Glm5NextMoE, "mlp" + ) if get_pp_group().is_last_rank: self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -834,12 +849,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ), 0, ) + num_fused_shared = _num_fused_shared_experts( + self.config.n_shared_experts, self.is_fused_shared_expert_enabled + ) expert_params_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", - num_experts=self.config.n_routed_experts, + num_experts=self.config.n_routed_experts + num_fused_shared, num_redundant_experts=num_redundant_experts, ) else: @@ -868,6 +886,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Models trained using ColossalAI may include these tensors in # the checkpoint. Skip them. continue + if self.is_fused_shared_expert_enabled: + name = _fused_shared_expert_name(name, self.config.n_routed_experts) # Handle FP8 indexer WK: dequantize to BF16 for fusion with # weights_proj into wk_weights_proj. @@ -1228,6 +1248,65 @@ def get_spec_layer_idx_from_weight_name( return None +def _fused_shared_experts_tuned(parallel_config: ParallelConfig) -> bool: + """AITER has fused-MoE configs tuned for the fused shared-expert shape + (one more expert and one more top-k slot than the routed MoE) only on + gfx950, with every expert on each rank and its weights split by TP4 or + TP8. Data, prefill context and expert parallelism change that split, so + any other GPU or parallel layout would run untuned fallback kernels.""" + from vllm.platforms.rocm import on_gfx950 + + reasons: list[str] = [] + if not on_gfx950(): + reasons.append("the GPU is not gfx950") + if parallel_config.tensor_parallel_size not in (4, 8): + reasons.append( + f"tensor_parallel_size is {parallel_config.tensor_parallel_size}" + ) + if parallel_config.data_parallel_size != 1: + reasons.append(f"data_parallel_size is {parallel_config.data_parallel_size}") + if parallel_config.prefill_context_parallel_size != 1: + reasons.append( + "prefill_context_parallel_size is " + f"{parallel_config.prefill_context_parallel_size}" + ) + if parallel_config.enable_expert_parallel: + reasons.append("expert parallelism is enabled") + + if not reasons: + return True + logger.warning_once( + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is ignored for GLM-5.3-Flash: " + "%s. AITER has tuned configs for its fused shared-expert MoE only on " + "gfx950 at TP4 and TP8, without data, prefill context or expert " + "parallelism. Running the shared experts as a separate MLP.", + "; ".join(reasons), + ) + return False + + +def _num_fused_shared_experts(n_shared_experts: int | None, enabled: bool) -> int: + """Expert slots the fused MoE appends for the shared expert; must match the + ``num_fused_shared_experts`` that ``FusedMoE`` allocates.""" + if not enabled or n_shared_experts is None: + return 0 + if n_shared_experts > 1: + raise NotImplementedError( + "Fused shared-expert loading supports only 1 shared expert per " + f"layer, but config.n_shared_experts is {n_shared_experts}. Set " + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0 to run the shared " + "experts as a separate MLP." + ) + return n_shared_experts + + +def _fused_shared_expert_name(name: str, n_routed_experts: int) -> str: + """Point a checkpoint ``mlp.shared_experts.*`` tensor at the fused MoE's + shared-expert slot, which follows the routed experts; other names are + returned unchanged.""" + return name.replace("mlp.shared_experts.", f"mlp.experts.{n_routed_experts}.", 1) + + def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params): if "indexer.wk." not in name or "wk_weights" in name: return False diff --git a/vllm/models/glm5next/common/mtp.py b/vllm/models/glm5next/common/mtp.py index 816cc46fdc08..2bb8b10d0746 100644 --- a/vllm/models/glm5next/common/mtp.py +++ b/vllm/models/glm5next/common/mtp.py @@ -10,6 +10,9 @@ from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -30,6 +33,8 @@ Glm5NextDecoderLayer, Glm5NextMLAAttention, Glm5NextMoE, + _fused_shared_expert_name, + _num_fused_shared_experts, _try_load_fp8_attn_proj, _try_load_fp8_indexer_wk, get_spec_layer_idx_from_weight_name, @@ -217,6 +222,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) self.set_moe_parameters() + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.model.layers.values(), Glm5NextMoE, "mtp_block.mlp" + ) def set_moe_parameters(self): self.num_moe_layers = self.config.num_nextn_predict_layers @@ -298,12 +306,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ("wk_weights_proj", "wk", 0), ("wk_weights_proj", "weights_proj", 1), ] + num_fused_shared = _num_fused_shared_experts( + self.config.n_shared_experts, self.is_fused_shared_expert_enabled + ) expert_params_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", - num_experts=self.config.n_routed_experts, + num_experts=self.config.n_routed_experts + num_fused_shared, num_redundant_experts=self.num_redundant_experts, ) @@ -328,6 +339,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if spec_layer is None: continue name = self._rewrite_spec_layer_name(spec_layer, name) + if self.is_fused_shared_expert_enabled: + name = _fused_shared_expert_name(name, self.config.n_routed_experts) if _try_load_fp8_indexer_wk( name, From efb8d85f24bd270bf20c143b32b5e7db6df61716 Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Wed, 7 Oct 2026 09:19:22 -0400 Subject: [PATCH 014/624] [ROCm][CI] Add FP32 rounding slack to accuracy test near-tie upper bounds (#60292) Signed-off-by: Divakar Verma Co-authored-by: Andreas Karatzas --- tests/v1/e2e/general/accuracy_utils.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/v1/e2e/general/accuracy_utils.py b/tests/v1/e2e/general/accuracy_utils.py index c8fa770fc5bf..e15ef03ebcfe 100644 --- a/tests/v1/e2e/general/accuracy_utils.py +++ b/tests/v1/e2e/general/accuracy_utils.py @@ -110,8 +110,10 @@ def check_first_divergence( # Log-softmax and reconstructing penalties may round in FP32, but # a near tie cannot excuse selecting the engine's own worse token. assert gap >= -1e-5, f"{detail}: selected token is below its alternative" - assert gap <= tolerance.greedy_atol, ( - f"{detail}: divergence exceeds near-tie bound={tolerance.greedy_atol}" + # Same FP32 rounding slack as above, applied to the upper bound. + assert gap <= tolerance.greedy_atol + 1e-5, ( + f"{detail}: divergence exceeds near-tie bound=" + f"{tolerance.greedy_atol} (+1e-5 slack)" ) @@ -201,9 +203,10 @@ def check_greedy_token( ) best = int(logits.argmax()) gap = float(logits[best] - logits[token]) - assert gap <= tolerance.greedy_atol, ( + # FP32 rounding can push a true near-tie fractionally past the bound. + assert gap <= tolerance.greedy_atol + 1e-5, ( f"{context}: selected token={token}, reference best={best}, " - f"logit gap={gap}, bound={tolerance.greedy_atol}" + f"logit gap={gap}, bound={tolerance.greedy_atol} (+1e-5 slack)" ) return gap From 5281e49908960f490c8faa789eee221d4ab5001d Mon Sep 17 00:00:00 2001 From: 00CC <83797617+tangzzycc@users.noreply.github.com> Date: Wed, 7 Oct 2026 21:21:43 +0800 Subject: [PATCH 015/624] [ROCm][Perf] Enable medium-skinny dispatch in RDNAHybridW4A16LinearKernel (#52619) Signed-off-by: tangzzycc <3081129260@qq.com> --- .../quantization/test_rdna_hybrid_w4a16.py | 81 ++++++++++++++++++- .../mixed_precision/rdna_hybrid_w4a16.py | 16 ++-- 2 files changed, 88 insertions(+), 9 deletions(-) diff --git a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py index d6045a791c1f..da7ef2d437d5 100644 --- a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py +++ b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py @@ -30,6 +30,8 @@ pack_int4_exllama_shuffle = hybrid_module.pack_int4_exllama_shuffle SUPPORTED_GROUP_SIZES = hybrid_module.SUPPORTED_GROUP_SIZES MAX_SKINNY_BATCH_SIZE = hybrid_module.MAX_SKINNY_BATCH_SIZE +LDS_CAPACITY_ELEMENTS = hybrid_module.LDS_CAPACITY_ELEMENTS +MEDIUM_SKINNY_LIMIT_ELEMENTS = hybrid_module.MEDIUM_SKINNY_LIMIT_ELEMENTS # --------------------------------------------------------------------------- @@ -100,8 +102,8 @@ def test_rdna_hybrid_w4a16_apply_matches_reference(dtype, group_size, has_zp, M) """Smoke test the registered custom op for both decode and prefill batches. Verifies the dispatch logic in `_rdna_hybrid_w4a16_apply_impl`: - - M <= MAX_SKINNY_BATCH_SIZE: HIP wvSplitK_int4_g - - M > MAX_SKINNY_BATCH_SIZE: Triton prefill kernel + - supported M and K*M: HIP wvSplitK_int4_g + - otherwise: Triton prefill kernel """ if not torch.cuda.is_available(): pytest.skip("CUDA/HIP device not available") @@ -742,3 +744,78 @@ def test_process_weights_pads_cliff_rows(K, expected_weight_stride, dist_init): assert tuple(layer.weight_packed.view(torch.int32).shape) == (N, K // 8) # Metadata is untouched by this change. assert layer.weight_scale.is_contiguous() + + +@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only") +@pytest.mark.parametrize( + "M,K,expected_path", + [ + (4, 8192, "hip"), + (4, 9728, "hip"), + (4, 9856, "triton"), + (MAX_SKINNY_BATCH_SIZE + 1, 1024, "triton"), + ], + ids=["regular_limit", "medium_range", "above_medium", "batch_too_large"], +) +def test_rdna_hybrid_w4a16_dispatch_boundaries(M, K, expected_path, monkeypatch): + """Route only supported regular and medium skinny shapes to HIP.""" + import vllm._custom_ops as ops + + N, G = 16, 128 + x_mk = torch.empty((M, K), dtype=torch.float16) + w_q = torch.empty((N, K // 2), dtype=torch.int8) + scales = torch.empty((N, K // G), dtype=torch.float16) + called = [] + + def fake_hip(*args, **kwargs): + called.append("hip") + return torch.empty((M, N), dtype=x_mk.dtype) + + def fake_triton(*args, **kwargs): + called.append("triton") + return torch.empty((M, N), dtype=x_mk.dtype) + + monkeypatch.setattr(ops, "wvSplitK_int4_g", fake_hip) + monkeypatch.setattr(hybrid_module, "triton_w4a16_skinny_fmt_gemm", fake_triton) + + output = hybrid_module._rdna_hybrid_w4a16_apply_impl( + x_mk, w_q, scales, None, None, 40, G + ) + + assert output.shape == (M, N) + assert called == [expected_path] + + +@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("has_zp", [False, True]) +def test_rdna_hybrid_w4a16_medium_skinny_matches_reference(dtype, has_zp): + """Validate the HIP medium path against an FP32 dequantization reference.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA/HIP device not available") + + from vllm.utils.platform_utils import num_compute_units + + set_random_seed(0) + M, K, N, G = 4, 9728, 256, 128 + assert LDS_CAPACITY_ELEMENTS < K * M <= MEDIUM_SKINNY_LIMIT_ELEMENTS + + x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype) + w_int4_nk = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32) + w_q = pack_int4_exllama_shuffle(w_int4_nk).contiguous().view(torch.int8) + scales = (0.05 * torch.rand((N, K // G), device=device, dtype=torch.float32)).to( + dtype + ) + zp = ( + torch.randint(0, 16, (N, K // G), device=device, dtype=torch.int32) + if has_zp + else None + ) + w_zp = _pack_zp_rows_for_kernel(zp) if zp is not None else None + + output = torch.ops.vllm.rdna_hybrid_w4a16_apply( + x_mk, w_q, scales, w_zp, None, num_compute_units(), G + ) + reference = _rdna_hybrid_w4a16_reference(x_mk, w_int4_nk, scales, zp, G, bias=None) + + torch.testing.assert_close(output, reference, rtol=2e-2, atol=2e-2) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py index 9d5303469945..bed62e89afce 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py @@ -2,9 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Hybrid W4A16 kernel: Triton for prefill, HIP skinny for decode. -Routes based on batch size M: - M <= MAX_SKINNY_BATCH_SIZE: HIP skinny GEMM (wvSplitK_int4_g) - M > MAX_SKINNY_BATCH_SIZE: Triton W4A16 fused dequant GEMM +Routes based on the activation shape: + M <= MAX_SKINNY_BATCH_SIZE and K*M <= MEDIUM_SKINNY_LIMIT_ELEMENTS: + HIP skinny GEMM (wvSplitK_int4_g) + otherwise: Triton W4A16 fused dequant GEMM Stores the weights ONCE as int8 [N, K//2] (ExLlama shuffle packed). Both paths read this single buffer: the HIP skinny kernel uses it directly, and @@ -56,13 +57,14 @@ def _on_gfx1151() -> bool: return on_gfx1151() -# Maximum batch size M for the HIP skinny kernel path (C++ supports N_in -# up to 5). When M is below this AND K*M fits in LDS, the skinny kernel is -# used; otherwise the Triton prefill path handles the GEMM. +# Maximum activation rows supported by HIP skinny (Python M maps to C++ N_in). +# When K*M exceeds regular LDS capacity, the C++ medium kernel caches the +# prefix in LDS and reads the remaining activation elements from global memory. MAX_SKINNY_BATCH_SIZE = 5 # 64 KiB per-workgroup LDS limit expressed in fp16 elements. # (AMD RDNA has 128 KiB total LDS per CU, but 64 KiB per workgroup.) LDS_CAPACITY_ELEMENTS = 64 * 1024 // 2 # 32768 fp16 elements +MEDIUM_SKINNY_LIMIT_ELEMENTS = int(LDS_CAPACITY_ELEMENTS * 1.2) # --------------------------------------------------------------------------- @@ -464,7 +466,7 @@ def _rdna_hybrid_w4a16_apply_impl( K = x_2d.shape[1] N = w_q.shape[0] - if M <= MAX_SKINNY_BATCH_SIZE and K * M <= LDS_CAPACITY_ELEMENTS: + if M <= MAX_SKINNY_BATCH_SIZE and K * M <= MEDIUM_SKINNY_LIMIT_ELEMENTS: # record_function is not torch.compile-safe; use nullcontext when # compiling to keep the op traceable. ctx = ( From c83935b3008553c72b1bad09f842b4c99b7fa47f Mon Sep 17 00:00:00 2001 From: viktorkou <105151911+viktorkou@users.noreply.github.com> Date: Wed, 7 Oct 2026 17:04:35 +0300 Subject: [PATCH 016/624] [Bugfix][Model Loader] Support reload_weights with runai_streamer load format (#58124) Signed-off-by: Viktor Koukouliev Co-authored-by: Claude Fable 5.1 Co-authored-by: Isotr0py --- .../test_runai_model_streamer_loader.py | 93 +++++++++++++++++++ .../test_runai_model_streamer_s3.py | 71 ++++++++++++++ tests/v1/worker/test_gpu_model_runner.py | 30 ++++++ .../model_loader/runai_streamer_loader.py | 15 ++- vllm/v1/worker/gpu_model_runner.py | 6 +- 5 files changed, 208 insertions(+), 7 deletions(-) diff --git a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py index 9164b8e4bea9..e0a083b3a8bc 100644 --- a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py +++ b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py @@ -6,11 +6,13 @@ from unittest.mock import patch import pytest +import torch from vllm import SamplingParams from vllm.config.load import LoadConfig from vllm.model_executor.model_loader import get_model_loader from vllm.model_executor.model_loader import runai_streamer_loader as rsl +from vllm.utils.mem_constants import GiB_bytes load_format = "runai_streamer" test_model = "openai-community/gpt2" @@ -122,3 +124,94 @@ def test_runai_invalid_extra_config_leaves_environ_untouched(): with pytest.raises(ValueError, match="memory_limit must be an integer >= -1"): _runai_loader({"concurrency": 16, "memory_limit": -5}) assert "RUNAI_STREAMER_CONCURRENCY" not in os.environ + + +@pytest.mark.parametrize( + "model_config,expected_source", + [ + # Object storage: model_weights holds the URI, model the pulled config dir. + ( + types.SimpleNamespace( + model="/tmp/pulled-config-files", + model_weights="s3://bucket/weights", + revision="myrev", + ), + "s3://bucket/weights", + ), + # HF repo or local path: model_weights is empty, model is the source. + ( + types.SimpleNamespace( + model="org/model", model_weights="", revision="myrev" + ), + "org/model", + ), + ], + ids=["object_storage", "hf_repo"], +) +def test_runai_get_all_weights_resolves_source(model_config, expected_source): + fake_self = types.SimpleNamespace( + _get_weights_iterator=lambda path, revision: iter( + [(f"{path}@{revision}", None)] + ) + ) + + weights = list( + rsl.RunaiModelStreamerLoader.get_all_weights( + fake_self, model_config, model=None + ) + ) + + assert weights == [(f"{expected_source}@myrev", None)] + + +@pytest.mark.parametrize( + "base_model,mul_model,add_model", + [ + ( + "Qwen/Qwen3-0.6B", + "inference-optimization/Qwen3-0.6B-debug-multiply", + "inference-optimization/Qwen3-0.6B-debug-add", + ), + ], +) +def test_runai_deep_sleep_reload_weights(base_model, mul_model, add_model, vllm_runner): + free, total = torch.accelerator.get_memory_info() + used_bytes_baseline = total - free + + def sleep_and_reload(llm, path): + # Level 2 discards the parameter memory, so after wake_up the weights + # can only come from the reload: a tensor the streamer fails to + # deliver changes the output instead of keeping a stale valid value. + llm.get_llm().sleep(level=2) + free, total = torch.accelerator.get_memory_info() + assert total - free - used_bytes_baseline < 3 * GiB_bytes + llm.get_llm().wake_up(tags=["weights"]) + llm.collective_rpc("reload_weights", kwargs={"weights_path": path}) + free, total = torch.accelerator.get_memory_info() + assert total - free - used_bytes_baseline < 4 * GiB_bytes + llm.get_llm().wake_up(tags=["kv_cache"]) + + with vllm_runner( + model_name=base_model, + load_format=load_format, + enable_sleep_mode=True, + enable_prefix_caching=False, + max_model_len=16, + max_num_seqs=1, + ) as llm: + base_output = llm.generate_greedy(["3 4 ="], max_tokens=4) + + sleep_and_reload(llm, mul_model) + mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0] + add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0] + assert mul_perp < add_perp + + sleep_and_reload(llm, add_model) + mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0] + add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0] + assert add_perp < mul_perp + + # Round trip to the original weights over discarded memory must + # reproduce the pre-sleep output. + sleep_and_reload(llm, base_model) + assert llm.generate_greedy(["3 4 ="], max_tokens=4) == base_output diff --git a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_s3.py b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_s3.py index d7abc80edbfa..a1440ae5d282 100644 --- a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_s3.py +++ b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_s3.py @@ -3,7 +3,9 @@ from pathlib import Path +import torch from runai_model_streamer.safetensors_streamer.streamer_mock import StreamerPatcher +from safetensors.torch import load_file, save_file from vllm.engine.arg_utils import EngineArgs from vllm.transformers_utils.repo_utils import hf_api @@ -50,3 +52,72 @@ def test_runai_model_loader_download_files_s3_mocked_with_patch( executor = RunaiDummyExecutor(vllm_config) executor.driver_worker.load_model() + + +def test_runai_reload_weights_path_from_mocked_s3(tmp_path: Path, monkeypatch): + # An engine started from object storage keeps the bucket URI in + # model_config.model_weights. A reload with weights_path must stream the + # new checkpoint, not re-stream the one the engine started from. + patcher = StreamerPatcher(str(tmp_path)) + monkeypatch.setattr( + "vllm.transformers_utils.runai_utils.runai_list_safetensors", + patcher.shim_list_safetensors, + ) + monkeypatch.setattr( + "vllm.transformers_utils.runai_utils.runai_pull_files", + patcher.shim_pull_files, + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.weight_utils.SafetensorsStreamer", + patcher.create_mock_streamer, + ) + + base_dir = tmp_path / "gpt2" + hf_api().snapshot_download( + repo_id=test_model, + local_dir=str(base_dir), + allow_patterns=["*.json", "*.txt", "model.safetensors"], + ) + # Second checkpoint under another prefix: the final layer norm scaled, so a + # reload that silently re-streams the original bucket is distinguishable. + tensors = load_file(base_dir / "model.safetensors") + original = tensors["ln_f.weight"].clone() + modified = original * 2 + tensors["ln_f.weight"] = modified + mod_dir = tmp_path / "gpt2-mod" + mod_dir.mkdir() + save_file(tensors, mod_dir / "model.safetensors", metadata={"format": "pt"}) + + vllm_config = EngineArgs( + model="s3://my-mock-bucket/gpt2/", + load_format=load_format, + tensor_parallel_size=1, + ).create_engine_config() + assert vllm_config.model_config.model_weights == "s3://my-mock-bucket/gpt2/" + + executor = RunaiDummyExecutor(vllm_config) + executor.driver_worker.load_model() + model = executor.driver_worker.model_runner.get_model() + (param,) = [p for n, p in model.named_parameters() if n.endswith("ln_f.weight")] + + def current() -> torch.Tensor: + return param.detach().float().cpu() + + torch.testing.assert_close(current(), original, rtol=1e-2, atol=1e-2) + + # No weights_path: re-stream the checkpoint the engine started from, i.e. + # the bucket kept in model_weights. Zero the live tensor first so a reload + # that did nothing is distinguishable from one that streamed the bucket. + param.data.zero_() + executor.driver_worker.reload_weights() + torch.testing.assert_close(current(), original, rtol=1e-2, atol=1e-2) + + # weights_path: the new prefix must win over the stale model_weights. + executor.driver_worker.reload_weights(weights_path="s3://my-mock-bucket/gpt2-mod/") + torch.testing.assert_close(current(), modified, rtol=1e-2, atol=1e-2) + + # After the override, a bare reload means the new checkpoint, not the + # one the engine was started with. + param.data.zero_() + executor.driver_worker.reload_weights() + torch.testing.assert_close(current(), modified, rtol=1e-2, atol=1e-2) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 5cf53899b836..deb4f7b42147 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -1015,6 +1015,36 @@ def test_reload_weights_before_load_model(model_runner): model_runner.reload_weights() +def test_reload_weights_path_replaces_object_storage_source(monkeypatch): + # An engine started from an object-storage URI keeps that URI in + # model_config.model_weights, which the runai_streamer loader prefers over + # model_config.model. A new weights_path must replace it, or the loader + # silently re-streams the original checkpoint. + loader = Mock() + loader.get_all_weights.return_value = iter(()) + monkeypatch.setattr(gpu_model_runner_module, "get_model_loader", lambda _: loader) + monkeypatch.setattr(gpu_model_runner_module, "initialize_layerwise_reload", Mock()) + monkeypatch.setattr(gpu_model_runner_module, "finalize_layerwise_reload", Mock()) + + runner = Mock(lora_config=None) + runner.model_config = SimpleNamespace( + model="/tmp/pulled-config-files", + model_weights="s3://bucket/original", + revision="abc123", + quantization=None, + ) + runner.get_model.return_value.named_parameters.return_value = [] + runner.get_model.return_value.load_weights.return_value = None + + GPUModelRunner.reload_weights(runner, weights_path="org/new-model") + + loader.get_all_weights.assert_called_once_with( + runner.model_config, runner.get_model.return_value + ) + cfg = runner.model_config + assert (cfg.model, cfg.model_weights, cfg.revision) == ("org/new-model", "", None) + + def test_sample_passes_reordered_draft_probs_to_rejection_sampler(): runner = object.__new__(GPUModelRunner) runner.use_async_scheduling = False diff --git a/vllm/model_executor/model_loader/runai_streamer_loader.py b/vllm/model_executor/model_loader/runai_streamer_loader.py index 2f18a2e3aca8..e00116006212 100644 --- a/vllm/model_executor/model_loader/runai_streamer_loader.py +++ b/vllm/model_executor/model_loader/runai_streamer_loader.py @@ -134,11 +134,16 @@ def download_model(self, model_config: ModelConfig) -> None: """Download model if necessary.""" self._prepare_weights(model_config.model, model_config.revision) - def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: - """Load weights into a model.""" + def get_all_weights( + self, + model_config: ModelConfig, + model: nn.Module, + ) -> Generator[tuple[str, torch.Tensor], None, None]: model_weights = model_config.model if model_weights_override := model_config.model_weights: model_weights = model_weights_override - model.load_weights( - self._get_weights_iterator(model_weights, model_config.revision) - ) + yield from self._get_weights_iterator(model_weights, model_config.revision) + + def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: + """Load weights into a model.""" + model.load_weights(self.get_all_weights(model_config, model)) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 76299308f866..0cc3fcc29a8f 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -5534,9 +5534,11 @@ def reload_weights( ) if weights_path is not None: - # The revision belongs to the model we are reloading away from, - # so it must not be carried over to the new path. + # The revision and any object-storage `model_weights` source + # belong to the model we are reloading away from, so they must + # not be carried over to the new path. self.model_config.model = weights_path + self.model_config.model_weights = "" self.model_config.revision = None weights_iterator = model_loader.get_all_weights(self.model_config, model) weights_iterator = cast( From f6e3944e3100fcb526e5cc24248080c31812be0f Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Wed, 7 Oct 2026 09:35:53 -0500 Subject: [PATCH 017/624] [ROCm][DSv4] AITER MegaMoEV2 Integration For DeepSeek V4 (#59685) Signed-off-by: Micah Williamson Co-authored-by: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> --- .buildkite/test-amd.yaml | 2 + .buildkite/test_areas/kernels.yaml | 2 + .../models/test_deepseek_v4_aiter_mega_moe.py | 186 ++++++++++++++++++ tests/models/test_deepseek_v4_vl_rocm.py | 4 +- vllm/config/kernel.py | 4 + .../device_communicators/all2all.py | 5 + vllm/models/deepseek_v4/amd/dspark.py | 14 +- vllm/models/deepseek_v4/amd/mega_moe.py | 154 +++++++++++++++ vllm/models/deepseek_v4/amd/model.py | 81 ++++++++ vllm/models/deepseek_v4/amd/mtp.py | 8 + 10 files changed, 454 insertions(+), 6 deletions(-) create mode 100644 tests/models/test_deepseek_v4_aiter_mega_moe.py create mode 100644 vllm/models/deepseek_v4/amd/mega_moe.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 8317919ce4b8..9e634ea02b4d 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -5858,6 +5858,7 @@ steps: timeout_in_minutes: 20 source_file_dependencies: - tests/models/test_deepseek_v4_vl_rocm.py + - tests/models/test_deepseek_v4_aiter_mega_moe.py - vllm/models/deepseek_v4/amd/ - vllm/models/deepseek_v4/common/ - vllm/model_executor/layers/fused_moe/ @@ -5871,6 +5872,7 @@ steps: - vllm/envs.py commands: - VLLM_ROCM_USE_AITER_MOE=0 pytest -v -s models/test_deepseek_v4_vl_rocm.py + - pytest -v -s models/test_deepseek_v4_aiter_mega_moe.py dind: false working_dir: /vllm-workspace/tests agent_pool: mi355_1 diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 54e7215967cd..cfc628d382c6 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -155,6 +155,7 @@ steps: - image-build-amd source_file_dependencies: - tests/models/test_deepseek_v4_vl_rocm.py + - tests/models/test_deepseek_v4_aiter_mega_moe.py - vllm/models/deepseek_v4/amd/ - vllm/models/deepseek_v4/common/ - vllm/model_executor/layers/fused_moe/ @@ -170,6 +171,7 @@ steps: # Common qnorm/RoPE/cache kernels already run in the H100 AMD mirror. # Pin the tested BF16 Triton routing and shared-expert path before imports. - VLLM_ROCM_USE_AITER_MOE=0 pytest -v -s models/test_deepseek_v4_vl_rocm.py + - pytest -v -s models/test_deepseek_v4_aiter_mega_moe.py # Own job so mHC never shares a process with FlashInfer's import-time # ptxas override. diff --git a/tests/models/test_deepseek_v4_aiter_mega_moe.py b/tests/models/test_deepseek_v4_aiter_mega_moe.py new file mode 100644 index 000000000000..a8f692e07bfb --- /dev/null +++ b/tests/models/test_deepseek_v4_aiter_mega_moe.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.layers.fused_moe import ( + fused_moe_make_expert_params_mapping, +) +from vllm.models.deepseek_v4.amd import mega_moe as mega_moe_module +from vllm.models.deepseek_v4.amd.mega_moe import ( + DeepseekV4AiterMegaMoEExperts, +) +from vllm.models.deepseek_v4.amd.model import DeepseekV4MoE +from vllm.platforms import current_platform + +HIDDEN = 64 +INTER = 64 + + +def _make_experts( + monkeypatch, + *, + ep_rank=0, + ep_size=1, + num_experts=2, + intermediate_size=INTER, + max_num_batched_tokens=3000, + swiglu_limit=None, +): + monkeypatch.setattr( + mega_moe_module, + "get_ep_group", + lambda: SimpleNamespace( + rank_in_group=ep_rank, world_size=ep_size, cpu_group="ep" + ), + ) + vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=max_num_batched_tokens) + ) + return DeepseekV4AiterMegaMoEExperts( + vllm_config, + num_experts=num_experts, + top_k=2, + hidden_size=HIDDEN, + intermediate_size=intermediate_size, + swiglu_limit=swiglu_limit, + ) + + +@pytest.mark.cpu_test +def test_experts_load_into_local_slots(monkeypatch): + model = torch.nn.Module() + model.ffn = torch.nn.Module() + model.ffn.experts = _make_experts(monkeypatch, ep_rank=1, ep_size=2, num_experts=4) + params = dict(model.named_parameters()) + mapping = fused_moe_make_expert_params_mapping( + model, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=4, + routed_experts_prefix="", + ) + + shapes = { + ("w1", "weight"): (INTER, HIDDEN // 2), + ("w3", "weight"): (INTER, HIDDEN // 2), + ("w2", "weight"): (HIDDEN, INTER // 2), + ("w1", "weight_scale"): (INTER, HIDDEN // 32), + ("w3", "weight_scale"): (INTER, HIDDEN // 32), + ("w2", "weight_scale"): (HIDDEN, INTER // 32), + } + shard_value = {"w1": 1, "w3": 2, "w2": 3} + loaded = {} + for expert in range(4): + for (shard, suffix), shape in shapes.items(): + name = f"ffn.experts.{expert}.{shard}.{suffix}" + weight = torch.full( + shape, 10 * expert + shard_value[shard], dtype=torch.uint8 + ) + for param_name, weight_name, expert_id, shard_id in mapping: + if weight_name not in name: + continue + mapped = name.replace(weight_name, param_name) + param = params[mapped] + if param.weight_loader( + param, + weight, + mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ): + loaded[name] = mapped + break + + assert sorted({n.split(".")[2] for n in loaded}) == ["2", "3"] + assert len(loaded) == 2 * len(shapes) + for slot, expert in enumerate((2, 3)): + for suffix in ("weight", "weight_scale"): + w13 = params[f"ffn.experts.w13_{suffix}"][slot] + w2 = params[f"ffn.experts.w2_{suffix}"][slot] + assert torch.all(w13[:INTER] == 10 * expert + 1) + assert torch.all(w13[INTER:] == 10 * expert + 2) + assert torch.all(w2 == 10 * expert + 3) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="Requires ROCm") +@pytest.mark.parametrize("hash_routing", [False, True]) +@pytest.mark.parametrize("num_tokens", [1, 37]) +def test_mega_moe_routing_matches_reference(monkeypatch, hash_routing, num_tokens): + num_experts, top_k, vocab = 16, 4, 64 + scaling = 2.5 + device = "cuda" + torch.manual_seed(0) + + class Gate(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.randn(num_experts, HIDDEN, device=device) * 0.1 + self.e_score_correction_bias = ( + None if hash_routing else torch.randn(num_experts, device=device) + ) + self.tid2eid = ( + torch.stack( + [ + torch.randperm(num_experts, device=device)[:top_k] + for _ in range(vocab) + ] + ).to(torch.int32) + if hash_routing + else None + ) + self.bias_vl = None + + def forward(self, x): + return x.float() @ self.weight.t(), None + + captured = {} + + def experts(x, topk_weights, topk_ids): + captured["weights"], captured["ids"] = topk_weights, topk_ids + return 2 * x + + moe = DeepseekV4MoE.__new__(DeepseekV4MoE) + torch.nn.Module.__init__(moe) + moe.gate = Gate() + moe.experts = experts + moe.shared_experts = lambda x: 3 * x + moe.scoring_func = "sqrtsoftplus" + moe.n_activated_experts = top_k + moe.renormalize = True + moe.routed_scaling_factor = scaling + moe.image_sentinel_lo = 0 + + x = torch.randn(num_tokens, HIDDEN, device=device, dtype=torch.bfloat16) + input_ids = torch.randint(0, vocab, (num_tokens,), device=device) + out = moe._forward_mega_moe(x, input_ids) + + ids, weights = captured["ids"], captured["weights"] + assert ids.dtype == torch.int32 and ids.is_contiguous() + assert weights.dtype == torch.float32 and weights.is_contiguous() + + scores = torch.nn.functional.softplus(x.float() @ moe.gate.weight.t()).sqrt() + if hash_routing: + expected_ids = moe.gate.tid2eid[input_ids].long() + else: + expected_ids = (scores + moe.gate.e_score_correction_bias).topk(top_k).indices + expected_weights = scores.gather(1, expected_ids) + expected_weights *= scaling / expected_weights.sum(dim=-1, keepdim=True) + + order = ids.long().sort(dim=-1).indices + expected_order = expected_ids.sort(dim=-1).indices + torch.testing.assert_close( + ids.long().gather(1, order), expected_ids.gather(1, expected_order) + ) + torch.testing.assert_close( + weights.gather(1, order), + expected_weights.gather(1, expected_order), + rtol=1e-3, + atol=1e-4, + ) + torch.testing.assert_close(out, 5 * x) diff --git a/tests/models/test_deepseek_v4_vl_rocm.py b/tests/models/test_deepseek_v4_vl_rocm.py index 036efbaa4248..762f1cd6c0e0 100644 --- a/tests/models/test_deepseek_v4_vl_rocm.py +++ b/tests/models/test_deepseek_v4_vl_rocm.py @@ -7,7 +7,7 @@ import torch from torch import nn -from vllm.config import ParallelConfig +from vllm.config import KernelConfig, ParallelConfig from vllm.model_executor.models.utils import WeightsMapper from vllm.platforms import current_platform @@ -58,6 +58,7 @@ def test_rocm_moe_routing_and_shared_experts_match_reference( model_config=SimpleNamespace(hf_config=config), quant_config=None, parallel_config=ParallelConfig(), + kernel_config=default_vllm_config.kernel_config, ), prefix="model.layers.0.ffn", ) @@ -230,6 +231,7 @@ def fake_factory(**kwargs): model_config=SimpleNamespace(hf_config=config), quant_config=None, parallel_config=ParallelConfig(), + kernel_config=KernelConfig(), ) hash_moe = rocm_model.DeepseekV4MoE(vllm_config, prefix="model.layers.0.ffn") diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index b8bfca082a4a..fa3ecf5df95d 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -137,6 +137,7 @@ def with_default( "humming", "triton_unfused", "aiter", + "aiter_mega_moe", "aiter_triton_mxfp4_bf16", "flydsl", "hpc", @@ -317,6 +318,9 @@ class KernelConfig: - "humming": Use Humming Mixed Precision kernels - "triton_unfused": Use Triton unfused MoE kernels - "aiter": Use AMD AITer kernels (ROCm only) + - "aiter_mega_moe": Use the AITER MegaMoE fused dispatch/GEMM/combine + kernel (ROCm gfx950, requires DP+EP with a MoRI all2all backend; + DeepSeek-V4 only) - "aiter_triton_mxfp4_bf16": Use the AITER Triton MXFP4 W4A16 (moe_gemm_a16w4) MoE kernel (ROCm gfx942/gfx950/gfx1250) - "flydsl": Use AMD FlyDSL kernels (ROCm only) diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index d8790dbf42cb..ae31c0b1d3cd 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -1008,6 +1008,11 @@ def __init__(self, cpu_group, all2all_backend: str): self.handle_cache = Cache() torch._C._distributed_c10d._register_process_group("mori", cpu_group) + if get_current_vllm_config().kernel_config.moe_backend == "aiter_mega_moe": + # MegaMoEV2 places its dispatch/combine workspaces on the MoRI + # symmetric heap, which defaults to 2 GB. MegaMoEV2 requires > 4GB + heap_size = os.environ.setdefault("MORI_SHMEM_HEAP_SIZE", "8G") + logger.info_once("AITER MegaMoE: MORI_SHMEM_HEAP_SIZE=%s", heap_size) mori.shmem.shmem_torch_process_group_init("mori") def _make_all2all_kwargs( diff --git a/vllm/models/deepseek_v4/amd/dspark.py b/vllm/models/deepseek_v4/amd/dspark.py index e0c69c03287c..e62431d7b62b 100644 --- a/vllm/models/deepseek_v4/amd/dspark.py +++ b/vllm/models/deepseek_v4/amd/dspark.py @@ -11,9 +11,8 @@ tilelang / triton / torch) instead of calling the tilelang kernels directly, and gate the trailing ``mhc_post`` on ``use_fused_mhc`` (True when AITER or TileLang fused MHC is available; False only on the torch fallback); - * drop the mega-MoE weight path (``make_deepseek_v4_expert_params_mapping``, - ``use_native_mega_moe`` and ``finalize_mega_moe_weights`` do not exist in - amd/model.py). + * use the AITER MegaMoE path (``aiter_mega_moe``) instead of the DeepGEMM + mega-MoE path. Everything else — the semi-autoregressive drafting hooks, the Markov head, the sliding-window context-KV insert, and the checkpoint ``mtp.*`` weight remap — is @@ -50,6 +49,7 @@ ) from vllm.model_executor.models.utils import maybe_prefix +from .mega_moe import finalize_mega_moe_weights from .model import ( DeepseekV4DecoderLayer, ) @@ -385,14 +385,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: Non-mtp weights (embed/head/main layers) belong to the target model and are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target. """ - # AMD DeepseekV4MoE has no mega-MoE path; always use the standard - # per-expert fused-MoE mapping (mirrors amd/mtp.py). expert_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="w1", ckpt_down_proj_name="w2", ckpt_up_proj_name="w3", num_experts=self.config.n_routed_experts, + routed_experts_prefix=( + "" if self.model.layers[0].ffn.use_mega_moe else "routed_experts" + ), ) expert_scale_suffix = ( ".weight_scale" @@ -496,6 +497,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: logger.info_once("DSpark draft model loaded: %d params", len(loaded_params)) return loaded_params + def process_weights_after_loading(self) -> None: + finalize_mega_moe_weights(self) + def _remap_dspark_name(self, name: str) -> str | None: """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path. diff --git a/vllm/models/deepseek_v4/amd/mega_moe.py b/vllm/models/deepseek_v4/amd/mega_moe.py new file mode 100644 index 000000000000..2fdff9335145 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mega_moe.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""AITER MegaMoE experts (fused dispatch, MXFP4 GEMMs, combine) for DeepSeek V4.""" + +from typing import Any + +import torch +import torch.nn as nn + +from vllm._aiter_ops import rocm_aiter_ops +from vllm.config import VllmConfig +from vllm.distributed import get_ep_group +from vllm.model_executor.utils import set_weight_attrs +from vllm.utils.math_utils import next_power_of_2 + + +class DeepseekV4AiterMegaMoEExperts(nn.Module): + # MegaMoEV2 owns large symmetric-heap buffers, so every layer shares one + # instance and passes its own weights per call. + _runtime_cache: dict[tuple, Any] = {} + + def __init__( + self, + vllm_config: VllmConfig, + *, + num_experts: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + swiglu_limit: float | None, + prefix: str = "", + ): + super().__init__() + ep_group = get_ep_group() + self.ep_rank = ep_group.rank_in_group + self.ep_size = ep_group.world_size + self.num_experts = num_experts + self.num_local_experts = num_experts // self.ep_size + self.experts_start_idx = self.ep_rank * self.num_local_experts + self.top_k = top_k + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.swiglu_limit = swiglu_limit or 0.0 + self.max_num_tokens = next_power_of_2( + vllm_config.scheduler_config.max_num_batched_tokens + ) + self.prefix = prefix + + def make_param(rows: int, cols: int) -> nn.Parameter: + param = nn.Parameter( + torch.zeros(self.num_local_experts, rows, cols, dtype=torch.uint8), + requires_grad=False, + ) + set_weight_attrs(param, {"weight_loader": self.weight_loader}) + return param + + self.w13_weight = make_param(2 * intermediate_size, hidden_size // 2) + self.w13_weight_scale = make_param(2 * intermediate_size, hidden_size // 32) + self.w2_weight = make_param(hidden_size, intermediate_size // 2) + self.w2_weight_scale = make_param(hidden_size, intermediate_size // 32) + self._runtime: Any = None + + def weight_loader( + self, + param: nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: int, + return_success: bool = False, + ) -> bool | None: + local_expert_id = expert_id - self.experts_start_idx + if not 0 <= local_expert_id < self.num_local_experts: + return False if return_success else None + expert_data = param.data[local_expert_id] + if shard_id in ("w1", "w3"): + shard_offset = 0 if shard_id == "w1" else self.intermediate_size + expert_data = expert_data.narrow(0, shard_offset, self.intermediate_size) + elif shard_id != "w2": + raise ValueError(f"Unsupported expert shard id: {shard_id}") + expert_data.copy_(loaded_weight.view(torch.uint8)) + return True if return_success else None + + def finalize_weights(self) -> None: + if self._runtime is not None: + return + n = self.num_local_experts + self.w13_weight.data = rocm_aiter_ops.shuffle_weight_a16w4( + self.w13_weight.data, 16, True + ) + self.w13_weight_scale.data = rocm_aiter_ops.shuffle_scale_a16w4( + self.w13_weight_scale.data.view(-1, self.hidden_size // 32), n, True + ) + self.w2_weight.data = rocm_aiter_ops.shuffle_weight_a16w4( + self.w2_weight.data, 16, False + ) + self.w2_weight_scale.data = rocm_aiter_ops.shuffle_scale_a16w4( + self.w2_weight_scale.data.view(-1, self.intermediate_size // 32), n, False + ) + key = ( + id(get_ep_group().cpu_group), + torch.accelerator.current_device_index(), + self.num_experts, + self.top_k, + self.hidden_size, + self.intermediate_size, + self.max_num_tokens, + self.swiglu_limit, + ) + runtime = self._runtime_cache.get(key) + if runtime is None: + from aiter.ops.flydsl.kernels.mega_moe import MegaMoEV2 + + runtime = MegaMoEV2( + rank=self.ep_rank, + world_size=self.ep_size, + model_dim=self.hidden_size, + inter_dim=self.intermediate_size, + experts=self.num_experts, + topk=self.top_k, + quant="a8w4", + w1=self.w13_weight, + w1_scale=self.w13_weight_scale, + w2=self.w2_weight, + w2_scale=self.w2_weight_scale, + max_tok_per_rank=self.max_num_tokens, + swiglu_limit=self.swiglu_limit, + ) + self._runtime_cache[key] = runtime + self._runtime = runtime + + def forward( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + ) -> torch.Tensor: + runtime = self._runtime + assert runtime is not None + # MegaMoEV2 reads these per launch; rebind to this layer's experts. + runtime._s1_w1 = self.w13_weight + runtime._s1_w1_scale = self.w13_weight_scale + runtime.w2 = self.w2_weight + runtime.w2_scale = self.w2_weight_scale + return runtime.forward(hidden_states.contiguous(), topk_weights, topk_ids) + + +DeepseekV4AiterMegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] + + +def finalize_mega_moe_weights(model: nn.Module) -> None: + for module in model.modules(): + if isinstance(module, DeepseekV4AiterMegaMoEExperts): + module.finalize_weights() diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index f8306dd12c23..8b26fc0060e7 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -30,6 +30,9 @@ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( rocm_aiter_fused_experts, ) +from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( + fused_topk_bias, +) from vllm.model_executor.layers.fused_moe.utils import ( is_model_fused_shared_expert_compatible, ) @@ -71,6 +74,7 @@ make_layers, maybe_prefix, ) +from vllm.models.deepseek_v4.amd.mega_moe import DeepseekV4AiterMegaMoEExperts from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention from vllm.platforms import current_platform from vllm.platforms.rocm import on_gfx950 @@ -529,6 +533,39 @@ def _fuse_shared_experts_enabled(config, parallel_config: ParallelConfig) -> boo ) +def _validate_aiter_mega_moe_config(vllm_config: VllmConfig) -> None: + parallel_config = vllm_config.parallel_config + if not rocm_aiter_ops.is_fused_moe_enabled(): + raise ValueError( + "--moe-backend aiter_mega_moe requires VLLM_ROCM_USE_AITER=1 " + " and VLLM_ROCM_USE_AITER_MOE=1." + ) + if not on_gfx950(): + raise NotImplementedError("AITER MegaMoE requires gfx950.") + if ( + not parallel_config.enable_expert_parallel + or parallel_config.tensor_parallel_size != 1 + or parallel_config.pipeline_parallel_size != 1 + ): + raise NotImplementedError( + "AITER MegaMoE requires data parallel attention with expert " + "parallel MoE: --data-parallel-size N --enable-expert-parallel " + "with TP=1 and PP=1." + ) + if parallel_config.all2all_backend not in ( + "mori_high_throughput", + "mori_low_latency", + ): + raise ValueError( + "AITER MegaMoE runs on MoRI shmem; add " + "--all2all-backend mori_high_throughput." + ) + if parallel_config.enable_eplb: + raise NotImplementedError("AITER MegaMoE does not support EPLB.") + if getattr(vllm_config.model_config.hf_config, "expert_dtype", "fp4") != "fp4": + raise NotImplementedError("AITER MegaMoE only supports fp4 experts.") + + class DeepseekV4MoE(nn.Module): def __init__( self, @@ -643,6 +680,20 @@ def __init__( self.experts_start_idx = self.tp_rank * self.n_local_experts self.experts_end_idx = self.experts_start_idx + self.n_local_experts + self.use_mega_moe = vllm_config.kernel_config.moe_backend == "aiter_mega_moe" + if self.use_mega_moe: + _validate_aiter_mega_moe_config(vllm_config) + self.experts = DeepseekV4AiterMegaMoEExperts( + vllm_config, + num_experts=config.n_routed_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + swiglu_limit=self.swiglu_limit, + prefix=f"{prefix}.experts", + ) + return + fuse_shared_into_routed = ( self.is_fused_shared_expert_enabled or self.fuse_heterogeneous_shared_expert ) @@ -692,6 +743,8 @@ def forward( raise ValueError("DeepSeek V4 vision MoE routing requires input_ids.") org_shape = hidden_states.shape + if self.use_mega_moe: + return self._forward_mega_moe(hidden_states, input_ids).view(org_shape) final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=hidden_states, @@ -700,6 +753,30 @@ def forward( return final_hidden_states.view(org_shape) + def _forward_mega_moe( + self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None + ) -> torch.Tensor: + router_logits, _ = self.gate(hidden_states) + topk_weights, topk_ids = fused_topk_bias( + hidden_states=hidden_states, + gating_output=router_logits, + scoring_func=self.scoring_func, + e_score_correction_bias=self.gate.e_score_correction_bias, + topk=self.n_activated_experts, + renormalize=self.renormalize, + input_tokens=input_ids, + hash_indices_table=self.gate.tid2eid, + routed_scaling_factor=self.routed_scaling_factor, + bias_vl=self.gate.bias_vl, + image_sentinel_lo=self.image_sentinel_lo, + ) + final_hidden_states = self.experts(hidden_states, topk_weights, topk_ids) + if self.shared_experts is not None: + final_hidden_states = final_hidden_states + self.shared_experts( + hidden_states + ) + return final_hidden_states + # Hidden sizes supported by AITER mhc_pre_big_fuse_rmsnorm. _AITER_MHC_FUSED_RMSNORM_SIZES = frozenset({1280, 2560, 4096, 7168}) @@ -993,6 +1070,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.fuse_heterogeneous_shared_expert = _heterogeneous_shared_expert_enabled( vllm_config ) + self.use_mega_moe = vllm_config.kernel_config.moe_backend == "aiter_mega_moe" self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, lambda prefix: DeepseekV4DecoderLayer( @@ -1315,6 +1393,7 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: ckpt_down_proj_name="w2", ckpt_up_proj_name="w3", num_experts=num_experts, + routed_experts_prefix="" if self.use_mega_moe else "routed_experts", ) @@ -1453,6 +1532,8 @@ def process_weights_after_loading(self) -> None: module.prepare_attn_preshuffle() elif isinstance(module, DeepseekV4MLP): module.prepare_gateup_preshuffle() + elif isinstance(module, DeepseekV4AiterMegaMoEExperts): + module.finalize_weights() if fused_compressor_layers: logger.info( "Fused the C4 compressor GEMMs in %d DeepSeek V4 layers", diff --git a/vllm/models/deepseek_v4/amd/mtp.py b/vllm/models/deepseek_v4/amd/mtp.py index 48ab036563f1..3746d6e38772 100644 --- a/vllm/models/deepseek_v4/amd/mtp.py +++ b/vllm/models/deepseek_v4/amd/mtp.py @@ -46,6 +46,7 @@ from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from .mega_moe import finalize_mega_moe_weights from .model import DeepseekV4DecoderLayer logger = init_logger(__name__) @@ -392,12 +393,16 @@ def _resolve_scale_name(name: str) -> str: head_rank_end = n_local_head * (tp_rank + 1) # Pre-compute expert mapping ONCE. + first_layer = next(iter(self.model.layers.values())) expert_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="w1", ckpt_down_proj_name="w2", ckpt_up_proj_name="w3", num_experts=self.config.n_routed_experts, + routed_experts_prefix=( + "" if first_layer.mtp_block.ffn.use_mega_moe else "routed_experts" + ), ) # FP8 experts register ``..._weight_scale_inv`` (block_quant) while @@ -526,6 +531,9 @@ def _resolve_scale_name(name: str) -> str: logger.info_once("MTP draft model loaded: %d params", len(loaded_params)) return loaded_params + def process_weights_after_loading(self) -> None: + finalize_mega_moe_weights(self) + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: """Rewrite the weight name to match the format of the original model. Add .mtp_block for modules in transformer layer block for spec layer From 548597367e66dcc90e24576c08e0e25f88c9ee6c Mon Sep 17 00:00:00 2001 From: Mikko Tukiainen Date: Wed, 7 Oct 2026 17:44:22 +0300 Subject: [PATCH 018/624] [ROCm][Perf] Add AITER FlyDSL GDN prefill backend (#57560) Signed-off-by: Mikko Tukiainen Co-authored-by: Cursor --- tests/v1/attention/test_gdn_aiter_flydsl.py | 199 ++++++++++++++++++ vllm/_aiter_ops.py | 113 ++++++++++ vllm/engine/arg_utils.py | 6 +- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 99 ++++++++- vllm/v1/attention/backends/gdn_attn.py | 57 ++++- 5 files changed, 462 insertions(+), 12 deletions(-) create mode 100644 tests/v1/attention/test_gdn_aiter_flydsl.py diff --git a/tests/v1/attention/test_gdn_aiter_flydsl.py b/tests/v1/attention/test_gdn_aiter_flydsl.py new file mode 100644 index 000000000000..3de625e9361d --- /dev/null +++ b/tests/v1/attention/test_gdn_aiter_flydsl.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from vllm._aiter_ops import rocm_aiter_ops +from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + ChunkGatedDeltaRule, + _resolve_gdn_prefill_backend, +) +from vllm.platforms import current_platform + + +def _make_config( + *, + backend: str = "aiter_flydsl", + head_k_dim: int | None = 128, + head_v_dim: int | None = 128, + dtype: torch.dtype = torch.bfloat16, +): + return SimpleNamespace( + additional_config={"gdn_prefill_backend": backend}, + model_config=SimpleNamespace( + dtype=dtype, + hf_text_config=SimpleNamespace( + linear_key_head_dim=head_k_dim, + linear_value_head_dim=head_v_dim, + ), + ), + ) + + +def _resolve_on_rocm(config, *, kernels_available: bool = True): + with ( + patch.object( + qwen_gdn_linear_attn.current_platform, "is_rocm", return_value=True + ), + patch.object( + qwen_gdn_linear_attn.rocm_aiter_ops, + "is_gdn_flydsl_prefill_available", + return_value=kernels_available, + ), + ): + return _resolve_gdn_prefill_backend(config) + + +@pytest.mark.parametrize( + "head_k_dim,head_v_dim,dtype,expected", + [ + (128, 128, torch.bfloat16, "aiter_flydsl"), + (64, 128, torch.bfloat16, "triton"), + (128, 64, torch.bfloat16, "triton"), + (128, 128, torch.float16, "triton"), + ], +) +def test_resolve_aiter_flydsl_gdn_prefill_backend( + head_k_dim: int, + head_v_dim: int, + dtype: torch.dtype, + expected: str, +): + config = _make_config(head_k_dim=head_k_dim, head_v_dim=head_v_dim, dtype=dtype) + + requested, active = _resolve_on_rocm(config) + + assert requested == "aiter_flydsl" + assert active == expected + + +def test_kda_style_model_cannot_select_flydsl(): + """Kimi K3 KDA reports head dims through linear_attn_config, not these. + + Its builder overrides _build_chunk_metadata, which the AITER path skips, + so this pins the reason the two never meet: the resolver turns the model + down before the builder is ever constructed. + """ + config = _make_config(head_k_dim=None, head_v_dim=None) + + assert _resolve_on_rocm(config) == ("aiter_flydsl", "triton") + + +@pytest.mark.parametrize("on_rocm", [True, False], ids=["kernels_missing", "not_rocm"]) +def test_explicit_aiter_flydsl_fails_closed(on_rocm: bool): + """An explicit request that cannot be honoured is a configuration error. + + Falling back silently would report the run as healthy while the requested + kernels were never used. + """ + config = _make_config() + if on_rocm: + with pytest.raises(RuntimeError, match="not importable"): + _resolve_on_rocm(config, kernels_available=False) + return + with ( + patch.object( + qwen_gdn_linear_attn.current_platform, "is_rocm", return_value=False + ), + patch.object( + qwen_gdn_linear_attn.current_platform, "is_cuda", return_value=True + ), + pytest.raises(RuntimeError, match="aiter_flydsl"), + ): + _resolve_gdn_prefill_backend(config) + + +def _flydsl_prefill_unavailable() -> bool: + return ( + not current_platform.is_rocm() + or not rocm_aiter_ops.is_gdn_flydsl_prefill_available() + ) + + +@pytest.mark.skipif( + _flydsl_prefill_unavailable(), + reason="needs ROCm with the AITER FlyDSL GDN prefill kernels installed", +) +@pytest.mark.parametrize( + "state_dtype", [torch.float32, torch.bfloat16], ids=["fp32_state", "bf16_state"] +) +def test_aiter_flydsl_prefill_matches_triton_reference(state_dtype: torch.dtype): + """Run both prefill paths on device and compare. + + The reference is the same AITER entry point with the FlyDSL prepare and K5 + kernels switched off, so the only thing that differs between the two calls + is the kernels this backend exists to select. The state dtype follows + --mamba-ssm-cache-dtype, which may be float32 or bfloat16. + """ + from aiter.ops.triton.gated_delta_net import chunk_gated_delta_rule_opt_vk + + torch.manual_seed(0) + device = torch.device("cuda") + seq_lens = [70, 130] + total_tokens = sum(seq_lens) + num_k_heads, num_v_heads, head_dim = 2, 4, 128 + + # use_qk_l2norm_in_kernel=False matches the prefill path, where the model + # has already normalised q and k. Feeding raw normals instead lets the + # delta rule run away to NaN, which says nothing about either kernel. + q = torch.nn.functional.normalize( + torch.randn(1, total_tokens, num_k_heads, head_dim, device=device), dim=-1 + ).to(torch.bfloat16) + k = torch.nn.functional.normalize( + torch.randn(1, total_tokens, num_k_heads, head_dim, device=device), dim=-1 + ).to(torch.bfloat16) + v = torch.randn( + 1, total_tokens, num_v_heads, head_dim, dtype=torch.bfloat16, device=device + ) + # g is a log-domain decay, so keep it negative or the recurrence diverges + # and the comparison stops saying anything about the kernels. + g = -torch.rand(1, total_tokens, num_v_heads, dtype=torch.float32, device=device) + beta = torch.rand(1, total_tokens, num_v_heads, dtype=torch.float32, device=device) + cu_seqlens = torch.tensor([0, 70, 200], dtype=torch.int32, device=device) + initial_state = torch.randn( + len(seq_lens), + num_v_heads, + head_dim, + head_dim, + device=device, + ).to(state_dtype) + + shared = { + "q": q, + "k": k, + "v": v, + "g": g, + "beta": beta, + "output_final_state": True, + "cu_seqlens": cu_seqlens, + "use_qk_l2norm_in_kernel": False, + } + + reference_out, reference_state = chunk_gated_delta_rule_opt_vk( + initial_state=initial_state.clone(), + state_dtype=state_dtype, + use_chunk_flydsl=False, + use_prepare_flydsl=False, + **shared, + ) + flydsl_out, flydsl_state = ChunkGatedDeltaRule.forward_aiter_flydsl( + None, + initial_state=initial_state.clone(), + aiter_prefill_metadata=rocm_aiter_ops.build_gdn_flydsl_prefill_metadata( + seq_lens, cu_seqlens=cu_seqlens + ), + **shared, + ) + + assert flydsl_state.dtype == state_dtype + torch.testing.assert_close( + flydsl_out.float(), reference_out.float(), rtol=2e-2, atol=2e-2 + ) + torch.testing.assert_close( + flydsl_state.float(), reference_state.float(), rtol=2e-2, atol=2e-2 + ) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 941f66e5b71f..44b89aeaf12c 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1997,6 +1997,8 @@ def get_moe_dispatch_policy(cls) -> int: # Check if the env variable is set _AITER_ENABLED = envs.VLLM_ROCM_USE_AITER + # Chunk length the AITER FlyDSL K5 prefill kernels are compiled for. + GDN_FLYDSL_CHUNK_SIZE = 64 _CUSTOM_ALL_REDUCE_ENABLED = envs.VLLM_ROCM_USE_AITER_CUSTOM_AR _LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR _FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE @@ -2467,6 +2469,23 @@ def _gdn_triton_kernels_importable() -> bool: except (ImportError, ModuleNotFoundError): return False + @staticmethod + def _gdn_flydsl_prefill_kernels_importable() -> bool: + try: + from aiter.ops.flydsl.linear_attention_prefill_kernels import ( # noqa: F401 + chunk_gated_delta_rule_fwd_h_flydsl_opt, + gdn_prepare_flydsl_supported, + gdn_prepare_fwd_flydsl, + ) + from aiter.ops.triton.gated_delta_net import ( # noqa: F401 + build_gated_delta_rule_prefill_metadata, + chunk_gated_delta_rule_opt_vk, + ) + + return True + except (ImportError, ModuleNotFoundError): + return False + @classmethod @if_aiter_supported def are_gdn_triton_kernels_available(cls) -> bool: @@ -2478,6 +2497,100 @@ def are_gdn_triton_kernels_available(cls) -> bool: """ return cls._AITER_ENABLED and cls._gdn_triton_kernels_importable() + @classmethod + @functools.cache + def is_gdn_flydsl_prefill_available(cls) -> bool: + """Whether the opt-in AITER FlyDSL GDN prefill path can be used. + + Selecting the backend is an explicit opt-in in itself, but it still + runs AITER kernels, so VLLM_ROCM_USE_AITER remains the one switch that + turns all of them off. + """ + return ( + cls._AITER_ENABLED + and is_aiter_found_and_supported() + and cls._gdn_flydsl_prefill_kernels_importable() + ) + + @classmethod + def build_gdn_flydsl_prefill_metadata( + cls, + seq_lens_cpu: list[int], + cu_seqlens: torch.Tensor, + ) -> object: + """Build the reusable varlen metadata FlyDSL GDN prefill runs against. + + The K5 recurrence is compiled for a fixed 64-token chunk, so the chunk + size is a property of the AITER kernel rather than of FLA, and is not + the caller's to choose. + """ + from aiter.ops.triton.gated_delta_net import ( + build_gated_delta_rule_prefill_metadata, + ) + + return build_gated_delta_rule_prefill_metadata( + seq_lens_cpu, + cu_seqlens=cu_seqlens, + chunk_size=cls.GDN_FLYDSL_CHUNK_SIZE, + ) + + @classmethod + def gdn_flydsl_prefill( + cls, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None, + aiter_prefill_metadata: object | None, + use_qk_l2norm_in_kernel: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run GDN prefill through AITER's VK path with the FlyDSL kernels. + + Every argument AITER takes is named here rather than forwarded as + ``**kwargs``, so a rename or reordering on the AITER side fails at this + call instead of silently binding to the wrong parameter. + """ + from aiter.ops.flydsl.linear_attention_prefill_kernels import ( + gdn_prepare_flydsl_supported, + ) + from aiter.ops.triton.gated_delta_net import chunk_gated_delta_rule_opt_vk + + if cu_seqlens is not None and aiter_prefill_metadata is None: + raise RuntimeError( + "AITER FlyDSL GDN prefill requires reusable varlen prefill metadata." + ) + if not gdn_prepare_flydsl_supported(k, v): + raise RuntimeError( + "AITER FlyDSL GDN prepare does not support the runtime input shape, " + "dtype, or device." + ) + logger.info_once( + "Dispatching AITER FlyDSL GDN prefill (prepare=flydsl, chunk=flydsl)." + ) + return chunk_gated_delta_rule_opt_vk( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + use_chunk_flydsl=True, + use_prepare_flydsl=True, + state_dtype=initial_state.dtype, + prefill_metadata=aiter_prefill_metadata, + # The caller stages the state densely and writes it back, so the + # kernel gets a plain [N, H, V, K] buffer and never an index into + # the pool; there is correspondingly nothing to update in place. + inplace_final_state=False, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + @classmethod def is_rdna_gdn_triton_kernels_available(cls) -> bool: """RDNA4 (gfx12) analog of are_gdn_triton_kernels_available().""" diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 527b8956d033..800d1b2f2ccf 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -839,7 +839,9 @@ class EngineArgs: ) fail_on_environ_validation: bool = False - gdn_prefill_backend: Literal["flashinfer", "triton", "cutedsl"] | None = None + gdn_prefill_backend: ( + Literal["flashinfer", "triton", "cutedsl", "aiter_flydsl"] | None + ) = None kda_prefill_backend: ( Literal["auto", "triton", "flashkda", "flashinfer", "fused"] | None ) = None @@ -1890,7 +1892,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: parser.add_argument( "--gdn-prefill-backend", dest="gdn_prefill_backend", - choices=["flashinfer", "triton", "cutedsl"], + choices=["flashinfer", "triton", "cutedsl", "aiter_flydsl"], default=None, help="Select GDN prefill backend.", ) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 7b164f1dd3df..1e7883864027 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -92,7 +92,7 @@ def _resolve_gdn_prefill_backend( vllm_config: VllmConfig, -) -> tuple[str, Literal["triton", "flashinfer", "cutedsl"]]: +) -> tuple[str, Literal["triton", "flashinfer", "cutedsl", "aiter_flydsl"]]: """Resolve GDN prefill backend. FlashInfer's GDN prefill kernel is chosen when: @@ -106,6 +106,11 @@ def _resolve_gdn_prefill_backend( In-tree CuteDSL GDN prefill kernel is chosen when: * "cutedsl" is requested; (opt-in only) * Blackwell (SM10.x) with ``head_k_dim == 128``; + + AITER FlyDSL GDN prefill kernels are chosen when: + * "aiter_flydsl" is requested; (opt-in only) + * ROCm AITER exposes the optimized VK prefill API and FlyDSL kernels; + * the model uses BF16 activations with ``head_k_dim == head_v_dim == 128``. """ additional_config = vllm_config.additional_config backend_cfg = ( @@ -115,12 +120,50 @@ def _resolve_gdn_prefill_backend( ) backend = str(backend_cfg).strip().lower() - if not current_platform.is_cuda(): - return backend, "triton" - head_k_dim = getattr( vllm_config.model_config.hf_text_config, "linear_key_head_dim", None ) + head_v_dim = getattr( + vllm_config.model_config.hf_text_config, "linear_value_head_dim", None + ) + + if current_platform.is_rocm(): + if backend == "aiter_flydsl": + if not rocm_aiter_ops.is_gdn_flydsl_prefill_available(): + raise RuntimeError( + "GDN prefill backend 'aiter_flydsl' was requested but the " + "AITER FlyDSL prefill kernels are not importable. Check " + "that VLLM_ROCM_USE_AITER=1, that this is a CDNA 3 or " + "newer GPU, and that the installed AITER exports them." + ) + if head_k_dim != 128 or head_v_dim != 128: + logger.warning_once( + "GDN prefill backend 'aiter_flydsl' was requested but " + "linear head dims are K=%s V=%s; FlyDSL requires 128/128. " + "Falling back to Triton/FLA.", + head_k_dim, + head_v_dim, + ) + return backend, "triton" + if vllm_config.model_config.dtype != torch.bfloat16: + logger.warning_once( + "GDN prefill backend 'aiter_flydsl' was requested but " + "model dtype is %s; FlyDSL requires bfloat16. " + "Falling back to Triton/FLA.", + vllm_config.model_config.dtype, + ) + return backend, "triton" + return backend, "aiter_flydsl" + return backend, "triton" + + if backend == "aiter_flydsl": + raise RuntimeError( + "GDN prefill backend 'aiter_flydsl' was requested but it is a ROCm " + f"AITER backend and the current platform is {current_platform.device_name}." + ) + + if not current_platform.is_cuda(): + return backend, "triton" supports_flashinfer = False supports_cutedsl = False @@ -172,6 +215,7 @@ def _log_gdn_backend_decision( chosen = { "flashinfer": "FlashInfer", "cutedsl": "CuteDSL", + "aiter_flydsl": "AITER FlyDSL", "triton": "Triton/FLA", }[active_backend] logger.info_once( @@ -259,6 +303,8 @@ def __init__(self) -> None: self._forward_method = self.forward_cuda elif active_backend == "cutedsl": self._forward_method = self.forward_cutedsl + elif active_backend == "aiter_flydsl": + self._forward_method = self.forward_aiter_flydsl else: self._forward_method = self.forward_native @@ -276,6 +322,7 @@ def forward_cuda( chunk_offsets: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = True, core_attn_out: torch.Tensor | None = None, + aiter_prefill_metadata: object | None = None, ): o, final_state = fi_chunk_gated_delta_rule( q=q, @@ -308,6 +355,7 @@ def forward_native( chunk_offsets: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = True, core_attn_out: torch.Tensor | None = None, + aiter_prefill_metadata: object | None = None, ): return fla_chunk_gated_delta_rule( q=q, @@ -338,6 +386,7 @@ def forward_cutedsl( chunk_offsets: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = True, core_attn_out: torch.Tensor | None = None, + aiter_prefill_metadata: object | None = None, ): from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import ( chunk_gated_delta_rule_cutedsl, @@ -367,6 +416,39 @@ def forward_cutedsl( final_state = None return o, final_state + def forward_aiter_flydsl( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_offsets: torch.Tensor | None = None, + use_qk_l2norm_in_kernel: bool = True, + core_attn_out: torch.Tensor | None = None, + aiter_prefill_metadata: object | None = None, + ): + o, final_state = rocm_aiter_ops.gdn_flydsl_prefill( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + aiter_prefill_metadata=aiter_prefill_metadata, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + if core_attn_out is not None: + o_flat = o.reshape(-1) + core_attn_out.reshape(-1)[: o_flat.numel()].copy_(o_flat) + return o, final_state + @PluggableLayer.register("qwen_gated_delta_net_attention") class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): @@ -1132,12 +1214,18 @@ def _warmup_prefill_kernels(self, qkv_or_qkvz: torch.Tensor, v_dim: int) -> None # CuteDSL kernels require metadata chunk_indices = None chunk_offsets = None + aiter_prefill_metadata = None if self.gdn_prefill_backend == "cutedsl": from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import ( prepare_metadata_cutedsl, ) chunk_indices, chunk_offsets = prepare_metadata_cutedsl(cu_seqlens, T) + elif self.gdn_prefill_backend == "aiter_flydsl": + aiter_prefill_metadata = rocm_aiter_ops.build_gdn_flydsl_prefill_metadata( + [T], + cu_seqlens=cu_seqlens, + ) try: self.chunk_gated_delta_rule( @@ -1152,6 +1240,7 @@ def _warmup_prefill_kernels(self, qkv_or_qkvz: torch.Tensor, v_dim: int) -> None chunk_indices=chunk_indices, chunk_offsets=chunk_offsets, use_qk_l2norm_in_kernel=False, + aiter_prefill_metadata=aiter_prefill_metadata, ) except Exception: logger.warning( @@ -1182,6 +1271,7 @@ def _warmup_prefill_kernels(self, qkv_or_qkvz: torch.Tensor, v_dim: int) -> None cu_seqlens, chunk_indices, chunk_offsets, + aiter_prefill_metadata, ) torch.accelerator.empty_cache() @@ -1523,6 +1613,7 @@ def _forward_core( chunk_indices=attn_metadata.chunk_indices, chunk_offsets=attn_metadata.chunk_offsets, use_qk_l2norm_in_kernel=False, + aiter_prefill_metadata=attn_metadata.aiter_prefill_metadata, ) # Init cache ssm_state[prefill_state_indices] = last_recurrent_state.to(ssm_state.dtype) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 95f6a808a597..08369adcb1a3 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -7,6 +7,7 @@ import torch +from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig from vllm.model_executor.layers.mamba.checkpoint import ( MambaPrefillCheckpointBuilder, @@ -79,6 +80,7 @@ class GDNAttentionMetadata: prefill_query_start_loc: torch.Tensor | None = None prefill_state_indices: torch.Tensor | None = None prefill_has_initial_state: torch.Tensor | None = None + aiter_prefill_metadata: object | None = None # The following attributes are for triton implementation of causal_conv1d nums_dict: dict | None = None @@ -109,8 +111,11 @@ def __init__( _resolve_gdn_prefill_backend, ) - self.gdn_prefill_backend: Literal["triton", "flashinfer", "cutedsl"] + self.gdn_prefill_backend: Literal[ + "triton", "flashinfer", "cutedsl", "aiter_flydsl" + ] _, self.gdn_prefill_backend = _resolve_gdn_prefill_backend(vllm_config) + self._check_chunk_metadata_override(type(self), self.gdn_prefill_backend) if self.speculative_config: assert self.speculative_config.num_speculative_tokens is not None @@ -186,6 +191,32 @@ def __init__( device=device, ) + @staticmethod + def _check_chunk_metadata_override( + builder_cls: type["GDNAttentionMetadataBuilder"], backend: str + ) -> None: + """Reject a subclass whose chunk metadata the AITER path would skip. + + AITER brings its own varlen prefill metadata and never calls + ``_build_chunk_metadata``, so a builder that overrides it would lose + that override without a word. No in-tree subclass can get here -- + ``_resolve_gdn_prefill_backend`` only selects this backend for models + with 128-dim GDN key and value heads, which the KDA builders are not -- + but a future one should be told rather than quietly ignored. + """ + if backend != "aiter_flydsl": + return + if ( + builder_cls._build_chunk_metadata + is GDNAttentionMetadataBuilder._build_chunk_metadata + ): + return + raise RuntimeError( + f"{builder_cls.__name__} builds its own FLA chunk metadata, which " + "the 'aiter_flydsl' GDN prefill backend does not use. Select a " + "different gdn_prefill_backend for this model." + ) + def _build_chunk_metadata( self, prefill_query_start_loc: torch.Tensor, @@ -431,6 +462,7 @@ def build( # type: ignore[override] prefill_query_start_loc: torch.Tensor | None = None prefill_state_indices: torch.Tensor | None = None prefill_has_initial_state: torch.Tensor | None = None + aiter_prefill_metadata: object | None = None if num_prefills > 0: # In a mixed non-spec batch, decodes are peeled off to the recurrent # kernel (decode-first front slice), so build chunk metadata from the @@ -452,11 +484,23 @@ def build( # type: ignore[override] prefill_query_start_loc_cpu = non_spec_query_start_loc_cpu prefill_state_indices = non_spec_state_indices_tensor - chunk_indices, chunk_offsets = self._build_chunk_metadata( - prefill_query_start_loc, - prefill_query_start_loc_cpu, - query_start_loc.device, - ) + if self.gdn_prefill_backend == "aiter_flydsl": + # AITER carries its own reusable varlen metadata and has no use + # for FLA's chunk indices, so it replaces them rather than + # extending what _build_chunk_metadata returns. + assert prefill_query_start_loc_cpu is not None + aiter_prefill_metadata = ( + rocm_aiter_ops.build_gdn_flydsl_prefill_metadata( + torch.diff(prefill_query_start_loc_cpu).tolist(), + cu_seqlens=prefill_query_start_loc, + ) + ) + else: + chunk_indices, chunk_offsets = self._build_chunk_metadata( + prefill_query_start_loc, + prefill_query_start_loc_cpu, + query_start_loc.device, + ) if num_prefills > 0: context_lens_tensor = m.compute_num_computed_tokens() @@ -568,6 +612,7 @@ def build( # type: ignore[override] prefill_query_start_loc=prefill_query_start_loc, prefill_state_indices=prefill_state_indices, prefill_has_initial_state=prefill_has_initial_state, + aiter_prefill_metadata=aiter_prefill_metadata, spec_query_start_loc=spec_query_start_loc, non_spec_query_start_loc=non_spec_query_start_loc, spec_state_indices_tensor=spec_state_indices_tensor, From 14e3902d5136e2fb00ef86c05fdc8d098f736283 Mon Sep 17 00:00:00 2001 From: Alec <35311602+alec-flowers@users.noreply.github.com> Date: Wed, 7 Oct 2026 07:46:44 -0700 Subject: [PATCH 019/624] [Frontend][Rust] Add gRPC forbidden token sequences and cache usage (#59837) Signed-off-by: Alec Flowers --- rust/proto/inference.proto | 4 + rust/src/server/src/grpc/convert.rs | 46 +++++++++++- rust/src/server/src/grpc/tests.rs | 2 + .../routes/openai/chat_completions/convert.rs | 1 + .../src/routes/openai/completions/convert.rs | 1 + rust/src/text/src/lower.rs | 75 ++++++++++++++++--- rust/src/text/src/lower/token_ids.rs | 5 ++ rust/src/text/src/request.rs | 3 + 8 files changed, 125 insertions(+), 12 deletions(-) diff --git a/rust/proto/inference.proto b/rust/proto/inference.proto index 2bc35bcb64a2..8189758212d4 100644 --- a/rust/proto/inference.proto +++ b/rust/proto/inference.proto @@ -102,6 +102,9 @@ message DecodingParameters { bool json_object = 10; string structural_tag = 11; } + + // Token sequences that must not be generated. Each sequence must be nonempty. + repeated TokenIds bad_words_token_ids = 12; } message StoppingCriteria { @@ -210,6 +213,7 @@ message FinishInfo { google.protobuf.Struct kv_transfer_params = 6; //uint64 seed = 7; google.protobuf.Struct ec_transfer_params = 8; + optional uint32 num_cached_tokens = 9; } // Info for candidate tokens other than the input/sampled diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index d5dd1ae64df4..8f2c589f154b 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -203,6 +203,10 @@ fn build_sampling_params( if !d.allowed_token_ids.is_empty() { params.allowed_token_ids = Some(d.allowed_token_ids.clone()); } + if !d.bad_words_token_ids.is_empty() { + params.bad_words_token_ids = + Some(d.bad_words_token_ids.iter().map(|sequence| sequence.ids.clone()).collect()); + } params.structured_outputs = convert_structured_output(d)?; } @@ -403,6 +407,7 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> Result>, /// Words to avoid during generation (tokenized to IDs during lowering). pub bad_words: Option>, + /// Pre-tokenized sequences to prohibit, in addition to `bad_words`. + pub bad_words_token_ids: Option>>, /// Specific token IDs for which log probabilities should be returned at /// each position. /// @@ -173,6 +175,7 @@ impl Default for SamplingParams { logit_bias: None, allowed_token_ids: None, bad_words: None, + bad_words_token_ids: None, logprob_token_ids: None, structured_outputs: None, skip_reading_prefix_cache: None, From 71f9a749203c31a263ccad03395623b1d6fc5538 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Wed, 7 Oct 2026 16:57:52 +0200 Subject: [PATCH 020/624] [CI][ROCm] Re-enable the DSv4.1 decoder replay graph test on ROCm (#60392) Signed-off-by: Stefan Koncarevic --- tests/models/test_deepseek_v41_decoder_replay_layers.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/models/test_deepseek_v41_decoder_replay_layers.py b/tests/models/test_deepseek_v41_decoder_replay_layers.py index 1302e0acda30..f3b124c5ac80 100644 --- a/tests/models/test_deepseek_v41_decoder_replay_layers.py +++ b/tests/models/test_deepseek_v41_decoder_replay_layers.py @@ -21,7 +21,6 @@ from vllm.models.deepseek_v41.nvidia.decoder_replay_cudagraph import ( DecoderReplayCudaGraphManager, ) -from vllm.platforms import current_platform from vllm.v1.worker.gpu import cudagraph_utils pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") @@ -126,9 +125,6 @@ def _graph_context(num_tokens): ) -@pytest.mark.skipif( - current_platform.is_rocm(), reason="decoder replay CUDA graphs are NVIDIA-only" -) def test_replay_graph_matches_eager(monkeypatch): """The replay graph of the next captured size pads and runs the rows.""" monkeypatch.setattr(cudagraph_utils, "get_pp_group", MagicMock) @@ -136,6 +132,8 @@ def test_replay_graph_matches_eager(monkeypatch): monkeypatch.setattr(cudagraph_utils, "graph_capture", lambda device: capture) compilation = CompilationConfig(decoder_replay_cudagraph_capture_sizes=[4]) cfg = MagicMock(compilation_config=compilation, speculative_config=None) + # Mocking it true would capture into the cuMem pool, which breaks on ROCm. + cfg.use_cumem_cudagraph_pool = False cfg.scheduler_config = MagicMock(max_num_seqs=2, max_num_batched_tokens=64) cfg.cache_config.use_kda_recoverssm = False cfg.model_config.hf_config = MagicMock(hc_mult=HC, hidden_size=HIDDEN) From 9367d8b9693c3d2fb6c3f99d8ccb4aa7673c0981 Mon Sep 17 00:00:00 2001 From: xaguilar-amd Date: Wed, 7 Oct 2026 17:22:32 +0200 Subject: [PATCH 021/624] [Bugfix][MLA] Scope DSpark non-causal capability to builder layers (#57992) Signed-off-by: Xavier Aguilar Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/attention/test_mla_noncausal.py | 121 ++++++++++++++++++ .../layers/attention/mla_attention.py | 8 +- 2 files changed, 127 insertions(+), 2 deletions(-) diff --git a/tests/v1/attention/test_mla_noncausal.py b/tests/v1/attention/test_mla_noncausal.py index c06ded5c1392..57d27ecd464f 100644 --- a/tests/v1/attention/test_mla_noncausal.py +++ b/tests/v1/attention/test_mla_noncausal.py @@ -19,6 +19,21 @@ class _NonCausalMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): supports_non_causal_multi_token_decode = True +class _CausalOnlyMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): + pass + + +class _DraftMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): + """Triton-like: non-causal decode is supported, DCP draft decode is not.""" + + supports_non_causal_multi_token_decode = True + + +class _DcpCapableMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): + supports_non_causal_multi_token_decode = True + supports_non_causal_multi_token_dcp = True + + def _metadata( query_start_loc: list[int], num_tokens: int | None = None, @@ -60,6 +75,54 @@ def _builder(marked: bool = True) -> _NonCausalMLAMetadataBuilder: return builder +def _mla_layer(*, non_causal: bool = False) -> SimpleNamespace: + return SimpleNamespace( + non_causal_multi_token_decode=non_causal, + q_lora_rank=None, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + prefill_backend=SimpleNamespace(clone=lambda: None), + ) + + +def _merged_mla_spec() -> MLAAttentionSpec: + return MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + non_causal_multi_token_decode=True, + ) + + +def _dspark_dcp_vllm_config( + static_forward_context: dict[str, SimpleNamespace], +) -> SimpleNamespace: + return SimpleNamespace( + speculative_config=SimpleNamespace( + method="dspark", num_speculative_tokens=None + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=2, + prefill_context_parallel_size=1, + cp_kv_cache_interleave_size=1, + ), + compilation_config=SimpleNamespace( + static_forward_context=static_forward_context + ), + model_config=SimpleNamespace( + dtype=torch.bfloat16, + max_model_len=128, + get_num_attention_heads=lambda _parallel_config: 16, + ), + cache_config=SimpleNamespace(cache_dtype="auto", block_size=64), + scheduler_config=SimpleNamespace(max_num_seqs=1), + attention_config=SimpleNamespace(use_prefill_query_quantization=False), + ) + + def test_group_capability_keeps_runtime_causality_per_step(): builder = _builder() @@ -137,3 +200,61 @@ def test_mla_cache_marker_is_promoted_to_group_capability(): [unmarked, unmarked] ).non_causal_multi_token_decode assert MLAAttentionSpec.merge([unmarked, marked]).non_causal_multi_token_decode + + +def test_builder_scopes_noncausal_capability_to_its_layers(): + merged_spec = _merged_mla_spec() + static_forward_context = { + "target": _mla_layer(non_causal=False), + "draft": _mla_layer(non_causal=True), + # Indexer and compressor caches share an MLA group but predate the flag. + "indexer": SimpleNamespace(), + } + vllm_config = _dspark_dcp_vllm_config(static_forward_context) + device = torch.device("cpu") + + def _flag(layer_names: list[str]) -> bool: + return _DcpCapableMLAMetadataBuilder( + merged_spec, + layer_names, + vllm_config, + device, + supports_dcp_with_varlen=True, + ).non_causal_multi_token_decode + + assert not _flag(["target"]) + assert _flag(["draft"]) + assert _flag(["target", "draft"]) + assert not _flag(["target", "indexer"]) + + +def test_merged_group_spec_does_not_mark_a_target_only_builder(): + # Both builders receive the same merged spec; only the draft layer is + # non-causal. Constructing through __init__ would still fail DCP + # validation if the builder read the group flag. + merged_spec = _merged_mla_spec() + static_forward_context = { + "target": _mla_layer(non_causal=False), + "draft": _mla_layer(non_causal=True), + } + vllm_config = _dspark_dcp_vllm_config(static_forward_context) + device = torch.device("cpu") + assert merged_spec.non_causal_multi_token_decode + + target = _CausalOnlyMLAMetadataBuilder( + merged_spec, + ["target"], + vllm_config, + device, + supports_dcp_with_varlen=True, + ) + assert not target.non_causal_multi_token_decode + + with pytest.raises(ValueError, match="non-causal draft"): + _DraftMLAMetadataBuilder( + merged_spec, + ["draft"], + vllm_config, + device, + supports_dcp_with_varlen=True, + ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 00e70d6fd0ae..f550dc3d90d4 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -2384,8 +2384,12 @@ def __init__( self.vllm_config = vllm_config self.device = device self.use_pcp = parallel_config.prefill_context_parallel_size > 1 - self.non_causal_multi_token_decode = getattr( - kv_cache_spec, "non_causal_multi_token_decode", False + # merge() unions this flag across the KV group, so a shared target/draft + # group would mark the target builder as non-causal. + ctx = self.compilation_config.static_forward_context + self.non_causal_multi_token_decode = any( + getattr(ctx[name], "non_causal_multi_token_decode", False) + for name in layer_names ) self._validate_dspark_dcp_support(supports_dcp_with_varlen) From f7f62d19160010bdc84b31af407bb8b88e5392d6 Mon Sep 17 00:00:00 2001 From: pmanczak Date: Wed, 7 Oct 2026 18:02:13 +0200 Subject: [PATCH 022/624] [XPU][MoE] Tune Triton fused MoE for Intel XPU (#53065) --- benchmarks/kernels/benchmark_moe.py | 47 ++++- ...ame=Intel(R)_Arc(TM)_Pro_B70_Graphics.json | 164 ++++++++++++++++++ ...ame=Intel(R)_Arc(TM)_Pro_B70_Graphics.json | 164 ++++++++++++++++++ ...ame=Intel(R)_Arc(TM)_Pro_B70_Graphics.json | 56 ++++++ 4 files changed, 427 insertions(+), 4 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/configs/E=128,N=384,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json create mode 100644 vllm/model_executor/layers/fused_moe/configs/E=64,N=1024,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json create mode 100644 vllm/model_executor/layers/fused_moe/configs/E=64,N=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index 1fcc7ddde7f8..f746e9aab339 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -303,9 +303,13 @@ def run(): run() torch.accelerator.synchronize() - # Capture 10 invocations with CUDA graph - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): + # Capture 10 invocations via graph(); CUDA needs a side stream + graph = ( + torch.cuda.CUDAGraph() + if current_platform.is_cuda_alike() + else torch.xpu.XPUGraph() + ) + with current_platform.graph(graph): for _ in range(10): run() torch.accelerator.synchronize() @@ -361,11 +365,33 @@ def get_rocm_tuning_space(use_fp16): return param_ranges +def get_xpu_tuning_space(use_fp16): + # BLOCK_SIZE_N=32 and num_warps=32 never won; grf_mode=256 beat 128 on B70. + block_m_range = [16, 32, 64, 128, 256] + block_n_range = [64, 128] + block_k_range = [32, 64] + num_warps_range = [4, 8, 16] + group_m_range = [1, 8, 16] + num_stage_range = [3, 4] + + return { + "BLOCK_SIZE_M": block_m_range, + "BLOCK_SIZE_N": block_n_range, + "BLOCK_SIZE_K": block_k_range, + "GROUP_SIZE_M": group_m_range, + "num_warps": num_warps_range, + "num_stages": num_stage_range, + "grf_mode": ["256"], + } + + def get_configs_compute_bound(use_fp16, block_quant_shape) -> list[dict[str, int]]: configs: list[BenchmarkConfig] = [] if current_platform.is_rocm(): param_ranges = get_rocm_tuning_space(use_fp16) + elif current_platform.is_xpu(): + param_ranges = get_xpu_tuning_space(use_fp16) else: # Reduced search space for faster tuning. # TODO(woosuk): Increase the search space and use a performance model to @@ -408,6 +434,14 @@ def get_configs_compute_bound(use_fp16, block_quant_shape) -> list[dict[str, int return configs +def prune_xpu_search_space(num_tokens, search_space, topk, num_experts): + # Cap the tile by rows routed to one expert, not total rows: wider tiles + # just pad and are the slowest to compile. + rows_per_expert = (num_tokens * topk + num_experts - 1) // num_experts + max_block_m = max(16, triton.next_power_of_2(rows_per_expert)) + return [c for c in search_space if c["BLOCK_SIZE_M"] <= max_block_m] + + def prune_rocm_search_space( num_tokens, shard_intermediate_size, hidden_size, search_space, is_fp16, topk ): @@ -522,7 +556,7 @@ def merge_unique_dicts(list1, list2): @ray.remote(num_gpus=1) class BenchmarkWorker: def __init__(self, seed: int) -> None: - torch.set_default_device("cuda") + torch.set_default_device(current_platform.device_type) set_random_seed(seed) self.seed = seed # Get the device ID to allocate tensors and kernels @@ -619,6 +653,10 @@ def tune( is_fp16, topk, ) + elif current_platform.is_xpu(): + search_space = prune_xpu_search_space( + num_tokens, search_space, topk, num_experts + ) need_device_guard = False if current_platform.is_rocm(): @@ -690,6 +728,7 @@ def sort_config(config: BenchmarkConfig) -> BenchmarkConfig: else {} ), **({"kpack": config["kpack"]} if "kpack" in config else {}), + **({"grf_mode": config["grf_mode"]} if "grf_mode" in config else {}), **({"SPLIT_K": config["SPLIT_K"]} if "SPLIT_K" in config else {}), } diff --git a/vllm/model_executor/layers/fused_moe/configs/E=128,N=384,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json b/vllm/model_executor/layers/fused_moe/configs/E=128,N=384,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json new file mode 100644 index 000000000000..6c1788d496d6 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=128,N=384,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4, + "grf_mode": "256" + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4, + "grf_mode": "256" + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3, + "grf_mode": "256" + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + } +} diff --git a/vllm/model_executor/layers/fused_moe/configs/E=64,N=1024,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json b/vllm/model_executor/layers/fused_moe/configs/E=64,N=1024,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json new file mode 100644 index 000000000000..0cdfc2266482 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=64,N=1024,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 4, + "grf_mode": "256" + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 4, + "grf_mode": "256" + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4, + "grf_mode": "256" + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + } +} diff --git a/vllm/model_executor/layers/fused_moe/configs/E=64,N=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json b/vllm/model_executor/layers/fused_moe/configs/E=64,N=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json new file mode 100644 index 000000000000..46bee68b897d --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=64,N=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json @@ -0,0 +1,56 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + } +} From c741bfca70cfb777e2016f827eae31f6e215fe9f Mon Sep 17 00:00:00 2001 From: Andy Lo Date: Wed, 7 Oct 2026 17:05:29 +0100 Subject: [PATCH 023/624] [Docs] Add review areas for andylolu2 (#60416) Signed-off-by: Andy Lo --- docs/community/reviewers.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/community/reviewers.md b/docs/community/reviewers.md index 60d37f0d3b4d..57840ffd8d8e 100644 --- a/docs/community/reviewers.md +++ b/docs/community/reviewers.md @@ -9,7 +9,7 @@ This page lists the active reviewers and describes the role. See the [committers Community members can reach out to the reviewers below for help with PRs in their areas, for example in the `#pr-reviews` Slack channel. Reviewers will add their areas of expertise to this list. Sorted alphabetically by GitHub handle: -- [@andylolu2](https://github.com/andylolu2) +- [@andylolu2](https://github.com/andylolu2): Mistral models, speculative decoding, quantization, scheduler, MLA - [@bnellnm](https://github.com/bnellnm) MoE, Quantization layers - [@BowenBao](https://github.com/BowenBao) - [@cjackal](https://github.com/cjackal): Multimodality; KV Connector and offload From 240ee8de9eb98d4e9269b82054df121adf522c17 Mon Sep 17 00:00:00 2001 From: MohitAMD Date: Wed, 7 Oct 2026 09:33:30 -0700 Subject: [PATCH 024/624] [Bugfix][KV Connector][ROCm] Fix MoRI-IO KV block-offset for speculative decoding (MTP) (#55053) Signed-off-by: Mohit Deopujari Signed-off-by: MohitAMD Co-authored-by: Cursor --- .../unit/test_moriio_connector.py | 348 +++++++++++++++++- .../unit/test_moriio_kv_layout.py | 24 +- .../v1/moriio/moriio_connector.py | 123 +++++-- .../kv_connector/v1/moriio/moriio_layout.py | 14 +- 4 files changed, 452 insertions(+), 57 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index f95e12c601bb..cb33bfbcf03c 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -3,6 +3,8 @@ import importlib.util import socket import uuid +from types import SimpleNamespace +from typing import Any from unittest.mock import MagicMock, patch import msgspec @@ -20,11 +22,13 @@ set_current_vllm_config, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( + ROLE, MoRIIOAgentMetadata, MoRIIOConnectorMetadata, MoRIIOConstants, MoRIIOMode, resolve_host_ip, + set_role, zmq_ctx, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( @@ -147,6 +151,50 @@ def _write_consumer_scheduler_for_finished_request(tp_size: int = 2): return scheduler +def _write_producer_scheduler(block_size: int = 1) -> Any: + """Bare WRITE-mode PRODUCER scheduler exercising the save path without a + real engine/RDMA stack. Uses a single full-attention KV group (index 0).""" + scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler) + scheduler.mode = MoRIIOMode.WRITE + scheduler.is_producer = True + scheduler.block_size = block_size + scheduler.blocks_per_sw = [0] + scheduler._group_block_sizes = [block_size] + scheduler.transfer_id_to_request_id = {} + scheduler._reqs_need_recv = {} + scheduler._reqs_need_save = {} + scheduler._reqs_need_pending_save = {} + scheduler._req_kv_params = {} + scheduler._reqs_need_send = {} + return scheduler + + +def _spec_kv_params(transfer_id: str = "xfer-spec") -> dict[str, Any]: + # request_id embeds no zmq address, so add_new_req resolves the peer here. + return { + "transfer_id": transfer_id, + "remote_engine_id": "remote-engine", + "remote_block_ids": [], + "remote_host": "127.0.0.1", + "remote_handshake_port": 5000, + "remote_notify_port": 5001, + } + + +def _build_meta( + scheduler: Any, req_ids=None, new_block_ids=None, num_scheduled_tokens=None +) -> Any: + # Final-chunk detection reads num_scheduled_tokens, so expose it. + scheduler_output = SimpleNamespace( + scheduled_cached_reqs=SimpleNamespace( + req_ids=req_ids or [], + new_block_ids=new_block_ids or [], + ), + num_scheduled_tokens=num_scheduled_tokens or {}, + ) + return scheduler.build_connector_meta(scheduler_output) + + class FakeMoRIIOWrapper: # A fake MoRIIOWrapper for testing purposes def __init__(self, *args, **kwargs): @@ -700,6 +748,203 @@ def test_resolve_host_ip_prefers_extra_config(): assert resolve_host_ip({"host_ip": ""}) == fallback +def test_write_mode_excludes_spec_lookahead_blocks(): + # The save path must record only the prompt blocks, clamping off the + # trailing lookahead blocks decode never allocates. + set_role(ROLE.PRODUCER) + scheduler = _write_producer_scheduler(block_size=1) + + req_id = "req-spec" + prompt_blocks = [10, 11, 12, 13, 14, 15, 16, 17] + lookahead_blocks = [18, 19, 20] + local_block_ids = [prompt_blocks + lookahead_blocks] + + # block_size=1 => num_prompt_tokens == prompt block count, all scheduled in + # this single chunk, so final-chunk detection fires. + num_prompt_tokens = len(prompt_blocks) + req = SimpleNamespace( + request_id=req_id, + num_prompt_tokens=num_prompt_tokens, + num_computed_tokens=0, + kv_transfer_params=_spec_kv_params(), + ) + scheduler._reqs_need_save[req_id] = (req, local_block_ids) + scheduler._req_kv_params[req_id] = _spec_kv_params() + + meta = _build_meta(scheduler, num_scheduled_tokens={req_id: num_prompt_tokens}) + + assert isinstance(meta, MoRIIOConnectorMetadata) + assert req_id in meta.reqs_to_save + saved = meta.reqs_to_save[req_id] + assert saved.local_block_ids == [prompt_blocks] + + +def test_write_mode_chunked_prefill_clamps_spec_lookahead_blocks(): + # Non-final chunks are buffered untouched; lookahead blocks are only clamped + # on the final-chunk save. + set_role(ROLE.PRODUCER) + scheduler = _write_producer_scheduler(block_size=1) + + req_id = "req-spec-chunked" + num_prompt_tokens = 8 + first_chunk = [10, 11, 12, 13] + final_chunk = [14, 15, 16, 17, 18, 19, 20] # 4 prompt + 3 lookahead + prompt_blocks = [10, 11, 12, 13, 14, 15, 16, 17] + + req = SimpleNamespace( + request_id=req_id, + num_prompt_tokens=num_prompt_tokens, + num_computed_tokens=0, + kv_transfer_params=_spec_kv_params(), + ) + + # Step 1: first chunk (4 of 8 tokens) is buffered, not saved. + scheduler._reqs_need_save[req_id] = (req, [first_chunk]) + scheduler._req_kv_params[req_id] = _spec_kv_params() + meta_step1 = _build_meta(scheduler, num_scheduled_tokens={req_id: 4}) + assert req_id not in meta_step1.reqs_to_save + assert req_id in scheduler._reqs_need_pending_save + assert scheduler._reqs_need_pending_save[req_id][1] == [first_chunk] + + # Step 2: final chunk (computed 4 + scheduled 4 >= 8) assembles the full + # local set, clamps the lookahead, and records only the prompt blocks. + req.num_computed_tokens = 4 + meta_step2 = _build_meta( + scheduler, + req_ids=[req_id], + new_block_ids=[[final_chunk]], + num_scheduled_tokens={req_id: 4}, + ) + assert req_id in meta_step2.reqs_to_save + assert meta_step2.reqs_to_save[req_id].local_block_ids == [prompt_blocks] + assert req_id not in scheduler._reqs_need_pending_save + + +@pytest.mark.parametrize( + ("block_size", "num_prompt_tokens", "num_prompt_blocks"), + [ + # block_size=4: exact multiple (2 full blocks) and non-multiple (3 blocks, + # last one partial). + pytest.param(4, 8, 2, id="bs4-exact-multiple"), + pytest.param(4, 9, 3, id="bs4-non-multiple"), + # block_size=16: exact multiple (2 full blocks) and non-multiple (3 + # blocks, last one partial). + pytest.param(16, 32, 2, id="bs16-exact-multiple"), + pytest.param(16, 40, 3, id="bs16-non-multiple"), + ], +) +def test_write_mode_clamps_prompt_blocks_with_block_size_gt_1( + block_size, num_prompt_tokens, num_prompt_blocks +): + # With block_size > 1 the clamp must keep exactly + # ceil(num_prompt_tokens / block_size) prompt blocks and drop the tail. + set_role(ROLE.PRODUCER) + scheduler = _write_producer_scheduler(block_size=block_size) + + req_id = f"req-bs{block_size}-{num_prompt_tokens}" + prompt_blocks = list(range(10, 10 + num_prompt_blocks)) + lookahead_blocks = [10 + num_prompt_blocks] + local_block_ids = [prompt_blocks + lookahead_blocks] + + req = SimpleNamespace( + request_id=req_id, + num_prompt_tokens=num_prompt_tokens, + num_computed_tokens=0, + kv_transfer_params=_spec_kv_params(), + ) + scheduler._reqs_need_save[req_id] = (req, local_block_ids) + scheduler._req_kv_params[req_id] = _spec_kv_params() + + meta = _build_meta(scheduler, num_scheduled_tokens={req_id: num_prompt_tokens}) + + assert isinstance(meta, MoRIIOConnectorMetadata) + assert req_id in meta.reqs_to_save + assert len(prompt_blocks) == num_prompt_blocks + assert meta.reqs_to_save[req_id].local_block_ids == [prompt_blocks] + + +def test_write_mode_defers_save_until_prompt_complete_with_lookahead(): + # Under speculative decoding the producer reserves lookahead blocks beyond + # the last-scheduled prompt token, inflating the accumulated block count on + # a non-final chunk. Token-progress detection must keep deferring until the + # prompt is complete, where the lookahead is clamped off. + # + # block_size=1, num_prompt_tokens=10, chunk=4, num_lookahead=3: after 8 + # prompt tokens the producer holds 8 + 3 = 11 blocks, but token progress + # 8 < 10 must keep waiting until the final chunk completes tokens 9-10. + set_role(ROLE.PRODUCER) + scheduler = _write_producer_scheduler(block_size=1) + + req_id = "req-lookahead-final-chunk" + num_prompt_tokens = 10 + + chunk1_blocks = [100, 101, 102, 103] # prompt tokens 1-4 + chunk2_blocks = [104, 105, 106, 107, 108, 109, 110] # 4 prompt + 3 lookahead + prompt_blocks = [100, 101, 102, 103, 104, 105, 106, 107, 108, 109] + + req = SimpleNamespace( + request_id=req_id, + num_prompt_tokens=num_prompt_tokens, + num_computed_tokens=0, + kv_transfer_params=_spec_kv_params(), + ) + scheduler._req_kv_params[req_id] = _spec_kv_params() + + # Chunk 1 (tokens 1-4): computed 0 + scheduled 4 < 10 -> buffered. + scheduler._reqs_need_save[req_id] = (req, [chunk1_blocks]) + meta1 = _build_meta(scheduler, num_scheduled_tokens={req_id: 4}) + assert req_id not in meta1.reqs_to_save, "chunk 1 must not be saved" + assert req_id in scheduler._reqs_need_pending_save + assert scheduler._reqs_need_pending_save[req_id][1] == [chunk1_blocks] + + # Chunk 2 (tokens 5-8) inflates the buffered count to 11 via lookahead, but + # token progress 4 + 4 = 8 < 10 keeps the request deferred. + req.num_computed_tokens = 4 + meta2 = _build_meta( + scheduler, + req_ids=[req_id], + new_block_ids=[[chunk2_blocks]], + num_scheduled_tokens={req_id: 4}, + ) + assert req_id not in meta2.reqs_to_save, ( + "second-to-last chunk saved prematurely despite incomplete prompt" + ) + assert req_id in scheduler._reqs_need_pending_save + assert scheduler._reqs_need_pending_save[req_id][1] == [ + chunk1_blocks + chunk2_blocks + ] + + # Final chunk (tokens 9-10, no new blocks): computed 8 + scheduled 2 >= 10 + # emits the save with the trailing lookahead block clamped off. + req.num_computed_tokens = 8 + meta3 = _build_meta( + scheduler, + req_ids=[req_id], + new_block_ids=[None], + num_scheduled_tokens={req_id: 2}, + ) + assert req_id in meta3.reqs_to_save, "final chunk must emit the WRITE save" + assert meta3.reqs_to_save[req_id].local_block_ids == [prompt_blocks] + assert req_id not in scheduler._reqs_need_pending_save + + +def test_clamp_to_prompt_blocks_clamps_each_group_by_its_own_block_size(): + # Two full-attention groups with different block sizes are each clamped to + # their own prompt-block count; a sliding-window group is left untouched. + scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler) + scheduler.blocks_per_sw = [0, 0, 3] + scheduler._group_block_sizes = [16, 32, 16] + + req = SimpleNamespace(num_prompt_tokens=40) + block_ids = [ + [0, 1, 2, 3, 4], # ceil(40/16)=3 -> clamp to 3 + [10, 11, 12], # ceil(40/32)=2 -> clamp to 2 + [20, 21, 22, 23, 24], # sliding-window -> untouched + ] + clamped = scheduler._clamp_to_prompt_blocks(req, block_ids) + assert clamped == [[0, 1, 2], [10, 11], [20, 21, 22, 23, 24]] + + def _make_hybrid_kv_cache_config() -> KVCacheConfig: """One full-attention group + one sliding-window group (Gemma-like).""" full_spec = FullAttentionSpec( @@ -731,6 +976,35 @@ def _make_hybrid_kv_cache_config() -> KVCacheConfig: ) +def _make_sliding_window_only_kv_cache_config() -> KVCacheConfig: + """A uniform sliding-window model (e.g. Mistral-7B): the single group is a + SlidingWindowSpec, so HMA is not required once hybrid management is + disabled, yet the spec guard still applies.""" + sw_spec = SlidingWindowSpec( + block_size=16, + num_kv_heads=4, + head_size=64, + dtype=torch.float16, + sliding_window=32, + ) + num_blocks = 2 + page = sw_spec.page_size_bytes + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=num_blocks * page, + layers=["sw0"], + layer_stride=num_blocks * page, + block_stride=page, + ) + ], + kv_cache_groups=[ + KVCacheGroupSpec(layer_names=["sw0"], kv_cache_spec=sw_spec), + ], + ) + + def _read_scheduler( kv_cache_config: KVCacheConfig, disable_hma: bool = False ) -> MoRIIOConnectorScheduler: @@ -820,15 +1094,6 @@ def test_non_sliding_window_hybrid_is_rejected(): _read_scheduler(config) -def test_token_count_basis_uses_full_attention_group(): - """Chunked-prefill token counting must use an unclipped full-attention - group (blocks_per_sw == 0), not a clipped sliding-window group.""" - scheduler = _read_scheduler(_make_hybrid_kv_cache_config()) - # Group 0 is the full-attention group - assert scheduler._full_attn_group_idx == 0 - assert scheduler._full_attn_block_size == 16 - - def test_get_exchange_clipped_blocks_clips_only_sw_group(): """get_exchange_clipped_blocks keeps the full attn group intact and clips the sliding-window group to its window tail.""" @@ -897,6 +1162,71 @@ def test_single_group_path_unchanged(): assert scheduler.get_exchange_clipped_blocks(blocks) == blocks +class _FakeSpecConfig(SimpleNamespace): + """Minimal speculative_config stand-in. Any drafter-capability probe + (use_dflash/use_eagle/use_dspark/uses_draft_model/...) defaults to False so + VllmConfig.num_lookahead_tokens resolves cleanly as upstream adds methods.""" + + def __getattr__(self, name: str): + if name.startswith(("use_", "uses_")): + return lambda *a, **k: False + raise AttributeError(name) + + +def _fake_speculative_config(num_speculative_tokens: int = 1) -> SimpleNamespace: + return _FakeSpecConfig(num_speculative_tokens=num_speculative_tokens) + + +def _spec_vllm_config() -> VllmConfig: + vllm_config = create_vllm_config(role="kv_producer", read_mode=False) + vllm_config.speculative_config = _fake_speculative_config() + # HMA off => the WRITE+hybrid guard is skipped, so construction reaches the + # spec-decode guard (the only path that hits it). + vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = True + return vllm_config + + +def test_spec_decode_rejected_with_sliding_window_groups(): + """MoRIIO + speculative decoding fails closed on sliding-window models, + since the prompt-block clamp only covers full-attention groups. The only + reachable path is a uniform SWA model with HMA disabled.""" + vllm_config = _spec_vllm_config() + with ( + set_current_vllm_config(vllm_config), + pytest.raises(NotImplementedError, match="speculative decoding"), + ): + MoRIIOConnector( + vllm_config, + KVConnectorRole.SCHEDULER, + _make_sliding_window_only_kv_cache_config(), + ) + + +def test_spec_decode_allowed_with_full_attention_only(): + """Full-attention-only models (e.g. GLM) do not trip the spec guard.""" + vllm_config = _spec_vllm_config() + with set_current_vllm_config(vllm_config): + connector = MoRIIOConnector( + vllm_config, + KVConnectorRole.SCHEDULER, + _make_test_kv_cache_config(), + ) + assert connector.connector_scheduler is not None + + +def test_spec_decode_allowed_with_sliding_window_in_read_mode(): + """The spec guard is WRITE-mode only; READ mode is unaffected.""" + vllm_config = create_vllm_config(role="kv_consumer", read_mode=True) + vllm_config.speculative_config = _fake_speculative_config() + with set_current_vllm_config(vllm_config): + connector = MoRIIOConnector( + vllm_config, + KVConnectorRole.SCHEDULER, + _make_hybrid_kv_cache_config(), + ) + assert connector.connector_scheduler is not None + + def test_hybrid_write_mode_rejected(): """Hybrid KV cache groups are unsupported in WRITE mode and fail closed.""" vllm_config = create_vllm_config(role="kv_producer", read_mode=False) diff --git a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py index 9860a7b7e1aa..5ec6fb5755d7 100644 --- a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py +++ b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py @@ -661,16 +661,34 @@ def test_moriio_wrapper_rejects_invalid_messages(role, payload, match): wrapper._handle_message(payload) -def test_local_block_ids_longer_than_remote_raises_value_error(): - cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16) +def test_local_block_ids_longer_than_remote_raises(): + # A longer local list is a genuine bug here and must fail loudly; the WRITE + # producer save path trims lookahead blocks before they reach this layer. + cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16) # interleaved worker = _worker({"layer": cache}, {"layer": _full_spec()}) with pytest.raises(ValueError, match="longer than remote_block_ids"): moriio_layout.compute_block_transfer_offsets( - "layer", cache, worker.layer_to_spec, [1, 3], [4], _remote_meta().num_blocks + "layer", + cache, + worker.layer_to_spec, + [1, 3, 6, 7], + [4, 5], + _remote_meta().num_blocks, ) +def test_shorter_local_block_ids_still_transfers_only_what_decode_allocated(): + # READ-mode partial transfer keeps its shorter-local behavior. + cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16) + worker = _worker({"layer": cache}, {"layer": _full_spec()}) + + short = moriio_layout.compute_block_transfer_offsets( + "layer", cache, worker.layer_to_spec, [1], [4, 5], _remote_meta().num_blocks + ) + assert len(short[0]) == 1 + + @pytest.mark.parametrize( ("local_block_ids", "remote_block_ids", "match"), [ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 5a8a9176640e..2d17d7a7ac52 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -581,12 +581,31 @@ def __init__( else (0, self.block_size) for g in kv_cache_config.transfer_groups ] - # add 1 to conservatively account for boundary overlap eg window isn't fully - # aligned with blocks. + # +1 conservatively covers window/block boundary overlap. self.blocks_per_sw = [ cdiv(n_tokens, block_size) + 1 if n_tokens else 0 for n_tokens, block_size in sw_sizes_tokens ] + # Per-group block size over transfer_groups (matching blocks_per_sw and + # the connector's block payload), used to clamp each full-attention group + # to its prompt length on the WRITE save path. + self._group_block_sizes = [ + getattr(g.kv_cache_spec, "block_size", self.block_size) + for g in kv_cache_config.transfer_groups + ] + # The WRITE-mode prompt-block clamp only covers full-attention groups, so + # MoRIIO + speculative decoding is not yet safe on sliding-window / + # hybrid models. TODO: lift once SWA + spec is validated end-to-end. + if ( + self.mode == MoRIIOMode.WRITE + and vllm_config.speculative_config is not None + and any(bps != 0 for bps in self.blocks_per_sw) + ): + raise NotImplementedError( + "MoRIIO with speculative decoding is not yet supported for " + "sliding-window / hybrid models: the prompt-block clamp is only " + "validated for full-attention KV groups." + ) # In WRITE mode and chunked prefill, we perform the write after the last chunk. # Hence we need to track once we are in the last chunk. We do this by computing # len(block_ids[g]) * block_size[g] for a full attn group g, as hybrid groups @@ -1178,6 +1197,45 @@ def update_state_after_alloc( params["do_remote_prefill"] = False + def _clamp_to_prompt_blocks(self, req: "Request", block_ids: BlockIds) -> BlockIds: + """Drop trailing speculative-lookahead blocks the consumer never + allocated, keeping each full-attention group's prompt blocks. + + Args: + req: Request being saved. + block_ids: Per-KV-cache-group local block ids. + + Returns: + ``block_ids`` with every full-attention group truncated to + ``ceil(num_prompt_tokens / block_size)`` leading blocks; + sliding-window groups are left unchanged. + + """ + clamped = list(block_ids) + changed = False + for g, group_blocks in enumerate(block_ids): + if self.blocks_per_sw[g] != 0: + continue + num_prompt_blocks = math.ceil( + req.num_prompt_tokens / self._group_block_sizes[g] + ) + if len(group_blocks) > num_prompt_blocks: + clamped[g] = group_blocks[:num_prompt_blocks] + changed = True + return clamped if changed else block_ids + + def _is_final_prefill_chunk( + self, req: "Request", scheduler_output: SchedulerOutput + ) -> bool: + """Whether this step completes the prompt. + + ``build_connector_meta`` runs before ``_update_after_schedule``, so + ``num_computed_tokens`` excludes the current chunk; add this step's + scheduled tokens. Token progress is lookahead-safe unlike block counts. + """ + num_scheduled = scheduler_output.num_scheduled_tokens.get(req.request_id, 0) + return req.num_computed_tokens + num_scheduled >= req.num_prompt_tokens + def build_connector_meta( self, scheduler_output: SchedulerOutput, @@ -1191,38 +1249,36 @@ def build_connector_meta( # It places the request metadata into the saving queue. for i, req_id in enumerate(scheduler_output.scheduled_cached_reqs.req_ids): + # Non-disagg requests aren't registered here; skip them + # instead of KeyError-ing. + if req_id not in self._reqs_need_pending_save: + continue + req, existing_blocks = self._reqs_need_pending_save[req_id] + # The final chunk may allocate no new blocks (the previous chunk + # reserved speculative-lookahead slack), so new_block_ids is None; + # still fall through to the final-chunk check below. new_block_ids = scheduler_output.scheduled_cached_reqs.new_block_ids[i] - if new_block_ids is not None: - # A non-disagg request (no kv_transfer_params, e.g. smoke - # test) is never registered in _reqs_need_pending_save; - # indexing it unconditionally would KeyError and crash the - # EngineCore. Skip it silently. - if req_id not in self._reqs_need_pending_save: - continue - req, existing_blocks = self._reqs_need_pending_save[req_id] - updated_blocks = [ + existing_blocks = [ existing_blocks[g] + new_block_ids[g] for g in range(len(new_block_ids)) ] - self._reqs_need_pending_save[req_id] = (req, updated_blocks) - saved_tokens = ( - len(updated_blocks[self._full_attn_group_idx]) - * self._full_attn_block_size - ) - if saved_tokens >= req.num_prompt_tokens: - # Final chunk: live kv_transfer_params may be cleared, - # so prefer the snapshot from update_state_after_alloc. - kv_params = self._req_kv_params.pop( - req_id, req.kv_transfer_params or {} - ) - meta.add_new_req( - request_id=req_id, - local_block_ids=self._reqs_need_pending_save[req_id][1], - kv_transfer_params=kv_params, - write_mode=True, - ) - del self._reqs_need_pending_save[req_id] + self._reqs_need_pending_save[req_id] = (req, existing_blocks) + if not self._is_final_prefill_chunk(req, scheduler_output): + continue + # Live kv_transfer_params may be cleared; prefer the + # update_state_after_alloc snapshot. + kv_params = self._req_kv_params.pop( + req_id, req.kv_transfer_params or {} + ) + save_block_ids = self._clamp_to_prompt_blocks(req, existing_blocks) + meta.add_new_req( + request_id=req_id, + local_block_ids=save_block_ids, + kv_transfer_params=kv_params, + write_mode=True, + ) + del self._reqs_need_pending_save[req_id] # Loop through scheduled reqs and convert to ReqMeta. for req_id, (req, block_ids) in self._reqs_need_recv.items(): @@ -1253,16 +1309,13 @@ def build_connector_meta( for req_id, (req, block_ids) in self._reqs_need_save.items(): kv_params = self._req_kv_params.get(req_id, req.kv_transfer_params or {}) - tokens_covered = ( - len(block_ids[self._full_attn_group_idx]) * self._full_attn_block_size - ) - if req.num_prompt_tokens > tokens_covered: - # not last chunk prefill + if not self._is_final_prefill_chunk(req, scheduler_output): self._reqs_need_pending_save[req_id] = (req, block_ids) continue + save_block_ids = self._clamp_to_prompt_blocks(req, block_ids) meta.add_new_req( request_id=req_id, - local_block_ids=block_ids, + local_block_ids=save_block_ids, kv_transfer_params=kv_params, write_mode=True, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py index b00692dc5188..8c55837dcd6a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py @@ -377,10 +377,8 @@ def iter_layer_registration_regions( ) region_len = geometry.num_blocks * geometry.regions_per_block * geometry.block_len if geometry.regions_per_block == 1: - # With padded or interleaved pages the block stride exceeds the - # meaningful block_len; register the strided span. The span ends with - # the last block's meaningful bytes, not a whole stride, so it never - # runs past the backing allocation of a strided layer view. + # Padded/interleaved pages stride past block_len: register the strided + # span, ending at the last block's bytes so it stays in-allocation. block_stride_bytes = geometry.block_stride * kv_cache.element_size() region_len = max( region_len, @@ -431,12 +429,8 @@ def compute_block_transfer_offsets( [list[int], list[int], list[int]], tuple[list[int], list[int], list[int]] ] = merge_contiguous_offsets, ) -> tuple[list[int], list[int], list[int]]: - # A shorter (or empty) local list is the READ-mode "drop the transfer, just - # free the prefill blocks" case (full-prefix-hit / aborted-before-scheduled): - # decode pulls fewer blocks than the prefill holds. The zip loop below pairs - # local[i]<->remote[i] and sizes by len(local), so a short local transfers - # only what decode allocated and an empty local is a no-op. A longer local - # list is a genuine bug and still fails loudly. + # A shorter/empty local list transfers only what decode allocated (READ-mode + # partial or full-prefix-hit); a longer local list is a genuine bug. if len(local_block_ids) > len(remote_block_ids): raise ValueError( "local_block_ids longer than remote_block_ids: " From d64a18832bcb97816f3ee88d7594457c2f50d224 Mon Sep 17 00:00:00 2001 From: Daniel Chan Date: Wed, 7 Oct 2026 12:49:59 -0400 Subject: [PATCH 025/624] [ROCm][Triton] Migrating MiniMax-M3 kernels from make_block_ptr to tensor descriptors (#59285) Signed-off-by: jpvillam Signed-off-by: Daniel Co-authored-by: jpvillam Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Rohan Potdar Co-authored-by: Daniel --- vllm/models/minimax_m3/amd/ops/index_topk.py | 15 +-- vllm/models/minimax_m3/amd/ops/sparse_attn.py | 30 ++--- .../minimax_m3/common/ops/sparse_attn.py | 105 ++++++++++-------- 3 files changed, 80 insertions(+), 70 deletions(-) diff --git a/vllm/models/minimax_m3/amd/ops/index_topk.py b/vllm/models/minimax_m3/amd/ops/index_topk.py index ce27ea122057..d2b4602c56b8 100644 --- a/vllm/models/minimax_m3/amd/ops/index_topk.py +++ b/vllm/models/minimax_m3/amd/ops/index_topk.py @@ -229,15 +229,16 @@ def _index_block_score_kernel( if BLOCK_SIZE_Q * pid_q >= q_len: return - q_ptrs = tl.make_block_ptr( + q_desc = tl.make_tensor_descriptor( base=q_ptr + seq_start * stride_q_n + pid_h * stride_q_h, - shape=(q_len, head_dim), - strides=(stride_q_n, stride_q_d), - offsets=(pid_q * BLOCK_SIZE_Q, 0), - block_shape=(BLOCK_SIZE_Q, head_dim), - order=(1, 0), + shape=[q_len, head_dim], + strides=[stride_q_n, stride_q_d], + block_shape=[BLOCK_SIZE_Q, head_dim], + padding_option="zero", ) - q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero") + + # 2. Perform the load by providing the multi-dimensional offset directly + q = q_desc.load([pid_q * BLOCK_SIZE_Q, 0]) q_start = prefix_len + pid_q * BLOCK_SIZE_Q off_q = tl.arange(0, BLOCK_SIZE_Q) + pid_q * BLOCK_SIZE_Q + prefix_len diff --git a/vllm/models/minimax_m3/amd/ops/sparse_attn.py b/vllm/models/minimax_m3/amd/ops/sparse_attn.py index 0719374496be..b4355d18f606 100644 --- a/vllm/models/minimax_m3/amd/ops/sparse_attn.py +++ b/vllm/models/minimax_m3/amd/ops/sparse_attn.py @@ -139,15 +139,15 @@ def _gqa_sparse_fwd_kernel( off_t = tl.arange(0, BLOCK_SIZE_T) topk_idx = tl.load(t_ptr_j + off_t * stride_tk, mask=off_t < max_topk, other=-1) real_topk = tl.sum((topk_idx >= 0).to(tl.int32), axis=0) - q_ptrs = tl.make_block_ptr( + q_desc = tl.make_tensor_descriptor( base=q_ptr + q_start * stride_qn + pid_h * stride_qh, - shape=(q_len, gqa_group_size, head_dim), - strides=(stride_qn, stride_qh, stride_qd), - offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0), - block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D), - order=(2, 1, 0), + shape=[q_len, gqa_group_size, head_dim], + strides=[stride_qn, stride_qh, stride_qd], + block_shape=[BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D], + padding_option="zero", ) - q = tl.load(q_ptrs, boundary_check=(0, 1, 2), padding_option="zero") + + q = q_desc.load([pid_q_j * BLOCK_SIZE_Q, 0, 0]) m_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32) lse_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32) acc_o = tl.zeros((BLOCK_SIZE_QH, BLOCK_SIZE_D), dtype=tl.float32) @@ -228,15 +228,17 @@ def _gqa_sparse_fwd_kernel( lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) acc_o = acc_o * tl.exp2(m_i - lse_i)[:, None] acc_o = tl.reshape(acc_o, BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D) - o_ptrs = tl.make_block_ptr( + o_desc = tl.make_tensor_descriptor( base=o_ptr + q_start * stride_on + pid_h * stride_oh, - shape=(q_len, gqa_group_size, head_dim), - strides=(stride_on, stride_oh, stride_od), - offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0), - block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D), - order=(2, 1, 0), + shape=[q_len, gqa_group_size, head_dim], + strides=[stride_on, stride_oh, stride_od], + block_shape=[BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D], + ) + + o_desc.store( + offsets=[pid_q_j * BLOCK_SIZE_Q, 0, 0], + value=acc_o.to(o_ptr.dtype.element_ty), ) - tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1, 2)) @torch.no_grad() diff --git a/vllm/models/minimax_m3/common/ops/sparse_attn.py b/vllm/models/minimax_m3/common/ops/sparse_attn.py index af703c086226..8dfd6e0c1996 100644 --- a/vllm/models/minimax_m3/common/ops/sparse_attn.py +++ b/vllm/models/minimax_m3/common/ops/sparse_attn.py @@ -110,7 +110,9 @@ def _gqa_sparse_fwd_kernel( bt_row = block_table_ptr + pid_b * stride_bt_b off_n = tl.arange(0, BLOCK_SIZE_K) off_d = tl.arange(0, BLOCK_SIZE_D) + off_h = tl.arange(0, BLOCK_SIZE_H) d_mask = off_d < head_dim + h_mask = off_h < gqa_group_size for j in range(real_q_loop): pid_q_j = pid_q * num_q_loop + j t_ptr_j = t_ptr + (q_block_start + pid_q_j) * stride_tn + pid_kh * stride_th @@ -118,15 +120,20 @@ def _gqa_sparse_fwd_kernel( q_abs = prefix_len + pid_q_j * BLOCK_SIZE_Q valid_blocks = (q_abs + BLOCK_SIZE_K) // BLOCK_SIZE_K real_topk = tl.minimum(max_topk, valid_blocks) - q_ptrs = tl.make_block_ptr( - base=q_ptr + q_start * stride_qn + pid_h * stride_qh, - shape=(q_len, gqa_group_size, head_dim), - strides=(stride_qn, stride_qh, stride_qd), - offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0), - block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D), - order=(2, 1, 0), + off_qi = pid_q_j * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q) + qo_mask = ( + (off_qi < q_len)[:, None, None] + & h_mask[None, :, None] + & d_mask[None, None, :] + ) + q = tl.load( + q_ptr + + (q_start + off_qi)[:, None, None] * stride_qn + + (pid_h + off_h)[None, :, None] * stride_qh + + off_d[None, None, :] * stride_qd, + mask=qo_mask, + other=0.0, ) - q = tl.load(q_ptrs, boundary_check=(0, 1, 2), padding_option="zero") off_q = ( tl.arange(0, BLOCK_SIZE_Q)[:, None] + pid_q_j * BLOCK_SIZE_Q @@ -203,15 +210,14 @@ def _gqa_sparse_fwd_kernel( lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) acc_o = acc_o * tl.exp2(m_i - lse_i)[:, None] acc_o = tl.reshape(acc_o, BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D) - o_ptrs = tl.make_block_ptr( - base=o_ptr + q_start * stride_on + pid_h * stride_oh, - shape=(q_len, gqa_group_size, head_dim), - strides=(stride_on, stride_oh, stride_od), - offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0), - block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D), - order=(2, 1, 0), + tl.store( + o_ptr + + (q_start + off_qi)[:, None, None] * stride_on + + (pid_h + off_h)[None, :, None] * stride_oh + + off_d[None, None, :] * stride_od, + acc_o.to(o_ptr.dtype.element_ty), + mask=qo_mask, ) - tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1, 2)) # --------------------------------------------------------------------------- @@ -306,21 +312,23 @@ def _gqa_sparse_decode_kernel( off_n = tl.arange(0, BLOCK_SIZE_K) off_d = tl.arange(0, BLOCK_SIZE_D) + off_h = tl.arange(0, BLOCK_SIZE_H) d_mask = off_d < head_dim + h_mask = off_h < gqa_group_size + hd_mask = h_mask[:, None] & d_mask[None, :] bt_row = block_table_ptr + req_id * stride_bt_b m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) - q_ptrs = tl.make_block_ptr( - base=q_ptr + pid_b * stride_qn + pid_h * stride_qh, - shape=(gqa_group_size, head_dim), - strides=(stride_qh, stride_qd), - offsets=(0, 0), - block_shape=(BLOCK_SIZE_H, BLOCK_SIZE_D), - order=(1, 0), + q = tl.load( + q_ptr + + pid_b * stride_qn + + (pid_h + off_h)[:, None] * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, ) - q = tl.load(q_ptrs, boundary_check=(0, 1), padding_option="zero") cur_idx_ptr = idx_base + chunk_start_topk * stride_tk for _ in tl.range(chunk_start_topk, chunk_end_topk): @@ -392,24 +400,23 @@ def _gqa_sparse_decode_kernel( # can hit 0 * NaN. All-empty padded rows may still produce NaNs in merge. scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), tl.zeros_like(lse_i)) acc_o = acc_o * scale[:, None] - o_ptrs = tl.make_block_ptr( - base=o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h, - shape=(gqa_group_size, head_dim), - strides=(stride_o_h, stride_o_d), - offsets=(0, 0), - block_shape=(BLOCK_SIZE_H, BLOCK_SIZE_D), - order=(1, 0), + tl.store( + o_ptr + + pid_c * stride_o_c + + pid_b * stride_o_b + + (pid_h + off_h)[:, None] * stride_o_h + + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, ) - tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1)) - lse_ptrs = tl.make_block_ptr( - base=lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h, - shape=(gqa_group_size,), - strides=(stride_l_h,), - offsets=(0,), - block_shape=(BLOCK_SIZE_H,), - order=(0,), + tl.store( + lse_ptr + + pid_c * stride_l_c + + pid_b * stride_l_b + + (pid_h + off_h) * stride_l_h, + lse_i.to(lse_ptr.dtype.element_ty), + mask=h_mask, ) - tl.store(lse_ptrs, lse_i.to(lse_ptr.dtype.element_ty), boundary_check=(0,)) @triton.heuristics( @@ -444,16 +451,16 @@ def _merge_topk_attn_out_kernel( off_c = tl.arange(0, NUM_TOPK_CHUNKS) off_d = tl.arange(0, BLOCK_SIZE_D) - o_ptrs = tl.make_block_ptr( - base=o_ptr + pid_b * stride_o_b + pid_h * stride_o_h, - shape=(NUM_TOPK_CHUNKS, head_dim), - strides=(stride_o_c, stride_o_d), - offsets=(0, 0), - block_shape=(NUM_TOPK_CHUNKS, BLOCK_SIZE_D), - order=(1, 0), + d_mask = off_d < head_dim + o_ptrs = ( + o_ptr + + pid_b * stride_o_b + + pid_h * stride_o_h + + off_c[:, None] * stride_o_c + + off_d[None, :] * stride_o_d ) lse_ptrs = lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c - o = tl.load(o_ptrs, boundary_check=(0, 1), padding_option="zero") + o = tl.load(o_ptrs, mask=d_mask[None, :], other=0.0) lse = tl.load(lse_ptrs) # empty chunks contribute -inf -> weight 0 lse_max = tl.max(lse, axis=0) weights = tl.exp2(lse - lse_max) @@ -462,7 +469,7 @@ def _merge_topk_attn_out_kernel( out_ptrs = ( out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d ) - tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim) + tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=d_mask) # --------------------------------------------------------------------------- From a8260cbc5c020d250b3964141cde8fd839152714 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 7 Oct 2026 18:01:41 +0100 Subject: [PATCH 026/624] [CI] Bump Transformers version to 5.19.0 (#60381) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- requirements/common.txt | 2 +- requirements/test/cpu.txt | 2 +- requirements/test/cuda.in | 2 +- requirements/test/cuda.txt | 2 +- requirements/test/nightly-torch.txt | 2 +- requirements/test/rocm.in | 2 +- requirements/test/rocm.txt | 2 +- requirements/test/xpu.in | 2 +- requirements/test/xpu.txt | 2 +- .../multimodal/pooling/test_embedding_gemma2.py | 15 +++------------ tests/models/registry.py | 1 - vllm/transformers_utils/processor.py | 6 +++++- 12 files changed, 17 insertions(+), 23 deletions(-) diff --git a/requirements/common.txt b/requirements/common.txt index 433494056d0e..ed62ca42f889 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -7,7 +7,7 @@ requests >= 2.33.0 tqdm blake3 py-cpuinfo -transformers >= 5.16.1, < 5.19.0 +transformers >= 5.16.1, < 5.20.0 huggingface_hub >= 1.31.0 # Required for the `huggingface_hub.utils.httpx` re-export tokenizers >= 0.23.2 # 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 diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index 45725d903bb3..7d1f13970fbc 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -1228,7 +1228,7 @@ traitlets==5.16.1 # via # jupyter-client # jupyter-core -transformers==5.18.0 +transformers==5.19.0 # via # -r requirements/test/../common.txt # -r requirements/test/cuda.in diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index b40d80413eee..4a32471652cf 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -40,7 +40,7 @@ datamodel_code_generator>=0.64.0 # required for minicpm3 test nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.18.0 +transformers==5.19.0 schemathesis>=4.0.0 # Required for openai schema test. # quantization buildkite-test-collector==0.1.9 diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 1f64120e31b6..29553d78c7b7 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -1335,7 +1335,7 @@ traitlets==5.16.1 # via # jupyter-client # jupyter-core -transformers==5.18.0 +transformers==5.19.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/nightly-torch.txt b/requirements/test/nightly-torch.txt index a229116859d3..66704d50d3c6 100644 --- a/requirements/test/nightly-torch.txt +++ b/requirements/test/nightly-torch.txt @@ -29,7 +29,7 @@ opencv-python-headless >= 4.13.0 # required for video test datamodel_code_generator>=0.64.0 # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.18.0 +transformers==5.19.0 schemathesis>=4.0.0 # Required for openai schema test. # quantization buildkite-test-collector==0.1.9 diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index 7c0e4848e944..0ad0b255116a 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -36,7 +36,7 @@ datamodel_code_generator>=0.64.0 # required for minicpm3 test nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.18.0 +transformers==5.19.0 schemathesis>=4.0.0 # Required for openai schema test # quantization bitsandbytes==0.50.0 diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index b2919a6e8361..a0ca0fd83310 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -1282,7 +1282,7 @@ traitlets==5.16.1 # via # jupyter-client # jupyter-core -transformers==5.18.0 +transformers==5.19.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/xpu.in b/requirements/test/xpu.in index 4b3d68ddfc9b..c5d08edde9bc 100644 --- a/requirements/test/xpu.in +++ b/requirements/test/xpu.in @@ -20,7 +20,7 @@ nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs modelscope<1.38 runai-model-streamer[s3,gcs,azure]>=0.15.7 tensorizer==2.10.1 -transformers==5.18.0 +transformers==5.19.0 # --- Audio Processing --- av diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 19cd271017ec..96fda8331f33 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -1134,7 +1134,7 @@ traitlets==5.16.1 # via # jupyter-client # jupyter-core -transformers==5.18.0 +transformers==5.19.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/tests/models/multimodal/pooling/test_embedding_gemma2.py b/tests/models/multimodal/pooling/test_embedding_gemma2.py index 7bb69fa69ad8..701d0914ddbc 100644 --- a/tests/models/multimodal/pooling/test_embedding_gemma2.py +++ b/tests/models/multimodal/pooling/test_embedding_gemma2.py @@ -9,20 +9,11 @@ import pytest import torch -from tests.models.registry import HF_EXAMPLE_MODELS - MODEL = os.environ.get("EG2_MODEL_PATH") needs_model = pytest.mark.skipif(not MODEL, reason="EG2_MODEL_PATH not set") TEXT_COS, MM_COS = 0.999, 0.999 -@pytest.fixture -def needs_hf_embedding_gemma2(): - """Skip when the installed transformers has no `embedding_gemma2` module.""" - model_info = HF_EXAMPLE_MODELS.get_hf_info("EmbeddingGemma2Model") - model_info.check_transformers_version(on_fail="skip") - - def _cos(a, b): return torch.nn.functional.cosine_similarity( torch.tensor(a, dtype=torch.float32), @@ -248,7 +239,7 @@ def nparams(m): ) -def test_encoder_contract(needs_hf_embedding_gemma2): +def test_encoder_contract(): from vllm.model_executor.models.embedding_gemma2 import EmbeddingGemma2Model from vllm.model_executor.models.gemma4_mm import ( Gemma4ForConditionalGeneration as G, @@ -432,7 +423,7 @@ def test_video_loader_duplicate_expansion_e2e(vllm_runner, st_model): assert len(out[0]) == 768 -def test_video_loader_indices(needs_hf_embedding_gemma2): +def test_video_loader_indices(): from transformers.models.embedding_gemma2.video_processing_embedding_gemma2 import ( EmbeddingGemma2VideoProcessor, ) @@ -517,7 +508,7 @@ def test_video_loader_indices(needs_hf_embedding_gemma2): assert v_4h == h_4h, f"1fps/4h mismatch: {v_4h[-2:]} vs {h_4h[-2:]}" -def test_video_loader_duplicate_expansion(needs_hf_embedding_gemma2): +def test_video_loader_duplicate_expansion(): from transformers.models.embedding_gemma2.video_processing_embedding_gemma2 import ( EmbeddingGemma2VideoProcessor, ) diff --git a/tests/models/registry.py b/tests/models/registry.py index 52ab94e6163d..4255d439cbdc 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -583,7 +583,6 @@ def check_available_online( ), "EmbeddingGemma2Model": _HfExamplesInfo( "google/embeddinggemma-2", - min_transformers_version="5.19.0", is_available_online=False, ), "Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2"), diff --git a/vllm/transformers_utils/processor.py b/vllm/transformers_utils/processor.py index 64af034d9ac4..9aa592084990 100644 --- a/vllm/transformers_utils/processor.py +++ b/vllm/transformers_utils/processor.py @@ -181,7 +181,11 @@ def get_video_processor_cls_name_from_config( # given type when torchvision is unavailable; callers then use opencv. model_config = get_hf_file_to_dict("config.json", processor_name, revision=revision) if model_config and "model_type" in model_config: - return VIDEO_PROCESSOR_MAPPING_NAMES.get(model_config["model_type"]) + video_processor = VIDEO_PROCESSOR_MAPPING_NAMES.get(model_config["model_type"]) + # Transformers >= 5.19 maps to {"torchvision": name} + if isinstance(video_processor, dict): + video_processor = video_processor.get("torchvision") + return video_processor return None From 3d6853b32c7c872df00d9e3489f642dc22e974e0 Mon Sep 17 00:00:00 2001 From: pondzikk <142788369+pondzikk@users.noreply.github.com> Date: Wed, 7 Oct 2026 19:09:50 +0200 Subject: [PATCH 027/624] [Bugfix] Seed hybrid mamba state index with mamba_block_size on prefix-cache hits (#55601) Signed-off-by: pondzikk Signed-off-by: zjy0516 Co-authored-by: Claude Fable 5.1 Co-authored-by: zjy0516 Co-authored-by: ptorsten <1652174+ptorsten@users.noreply.github.com> --- tests/v1/worker/test_mamba_hybrid_model_state.py | 16 ++++++++++++++++ vllm/v1/worker/gpu/model_states/mamba_hybrid.py | 7 ++++++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index f87201e48de2..fa54ccfcedc5 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -67,6 +67,22 @@ def test_prepare_attn_forwards_positions(monkeypatch: pytest.MonkeyPatch) -> Non assert build_attn_metadata.call_args.kwargs["positions"] is positions +def test_add_request_seeds_state_idx_in_mamba_blocks() -> None: + """A drafter group can lower cache_config.block_size below the mamba + block size; the seeded column must still be in mamba blocks.""" + state = object.__new__(MambaHybridModelState) + state.cache_config = SimpleNamespace(block_size=16, mamba_block_size=880) + state._align_mode = True + state.rope_state = None + state.prompt_embeds_state = None + state.num_accepted_tokens_gpu = torch.ones(2, dtype=torch.int32) + state._mamba_state_idx_gpu = torch.zeros(2, dtype=torch.int32) + + state.add_request(1, SimpleNamespace(num_computed_tokens=107_360)) + + assert state._mamba_state_idx_gpu[1] == 121 + + def test_padded_prompt_tail_builds_as_spec_decode( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 1f8af07898a1..c45304efa512 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -115,8 +115,13 @@ def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: self.num_accepted_tokens_gpu[req_index].fill_(1) if self._align_mode: # Seed the running state block from the resumed/prefilled position. + # The column is in mamba blocks. cache_config.block_size is the + # min over prefix-cacheable groups, so a drafter group with its own + # smaller block (e.g. DFlash) pulls it below mamba_block_size. + mamba_block_size = self.cache_config.mamba_block_size + assert mamba_block_size is not None self._mamba_state_idx_gpu[req_index].fill_( - (new_req_data.num_computed_tokens - 1) // self.cache_config.block_size + (new_req_data.num_computed_tokens - 1) // mamba_block_size ) def _get_mamba_group_info( From 554340f3d3259e321be4c07282be7a02a5aeef83 Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Thu, 8 Oct 2026 02:24:21 +0900 Subject: [PATCH 028/624] [Attention] Support NVFP4 KV cache on SM8x and SM12x with FlashInfer (#46963) Signed-off-by: lesj0610 Signed-off-by: mgoin Co-authored-by: lesj0610 Co-authored-by: mgoin Co-authored-by: Claude --- tests/kernels/attention/test_flashinfer.py | 33 +++++ vllm/v1/attention/backends/flashinfer.py | 138 +++++++++++++++++---- 2 files changed, 149 insertions(+), 22 deletions(-) diff --git a/tests/kernels/attention/test_flashinfer.py b/tests/kernels/attention/test_flashinfer.py index c6cba8d530e4..0d9aecd31e7d 100644 --- a/tests/kernels/attention/test_flashinfer.py +++ b/tests/kernels/attention/test_flashinfer.py @@ -6,6 +6,7 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed +from vllm.v1.kv_cache_layout import KVCacheLayout try: import flashinfer @@ -702,3 +703,35 @@ def test_flashinfer_decode_with_paged_fp8_kv( torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=1e-2), f"{torch.max(torch.abs(output - ref_output))}", ) + + +@pytest.mark.skipif( + not any(current_platform.is_device_capability_family(f) for f in (80, 120)), + reason="NVFP4 fa2 path", +) +@pytest.mark.parametrize("layout", ["LBNHC", "LBHNC", "BLNHC", "BLHNC"]) +def test_flashinfer_nvfp4_kv_cache_update(default_vllm_config, layout): + from vllm.v1.attention.backends import flashinfer as fi + + torch.set_default_device("cuda") + order = KVCacheLayout[layout].stride_order + # [L, B, 2H, N, 64 + 8] NVFP4 pages for two layers; layer 1 is written. + shape = (2, 6, 4, 16, 64 + 8) + layers = torch.zeros([shape[i] for i in order], dtype=torch.uint8) + layers = layers.permute(*(order.index(i) for i in range(5))) + one = torch.ones(()) + layer = torch.nn.Module() + layer._k_scale = layer._v_scale = one + kv = torch.randn(2, 96, 2, 128, dtype=torch.bfloat16) + impl = fi.FlashInferImpl(4, 128, 1.0, 2, None, None, "nvfp4") + impl.do_kv_cache_update(layer, kv[0], kv[1], layers[1], torch.arange(96)) + + data, scales, kv_layout = fi._nvfp4_kv_views(layers[1], 2) + pages = torch.arange(6, dtype=torch.int32)[None] + seq_lens = torch.tensor([96], dtype=torch.int32) + deq = torch.empty_like(kv)[:, None] + flashinfer.nvfp4_kv_dequantize_paged( + data, scales, pages, seq_lens, one, one, *deq, kv_layout + ) + assert ((deq[:, 0] - kv).norm() / kv.norm()) < 0.2 + assert not layers[0].any() diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 91ff6da3f037..a21720ab26e6 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -16,6 +16,7 @@ BatchPrefillWithRaggedKVCacheWrapper, MultiLevelCascadeAttentionWrapper, get_seq_lens, + nvfp4_quantize_append_paged_kv_cache_with_slot_mapping, ) from flashinfer.decode import fast_decode_plan, trtllm_batch_decode_with_kv_cache from flashinfer.prefill import trtllm_batch_context_with_kv_cache @@ -106,6 +107,38 @@ trtllm_workspace_buffer = None +def _nvfp4_kv_on_fa2() -> bool: + return current_platform.is_device_capability_family( + 80 + ) or current_platform.is_device_capability_family(120) + + +_KVPair = tuple[torch.Tensor, torch.Tensor] + + +def _nvfp4_kv_views( + kv_cache: torch.Tensor, num_kv_heads: int +) -> tuple[_KVPair, _KVPair, str]: + """Views of NVFP4 pages laid out as [K data | K scale | V data | V scale].""" + num_blocks, num_slots, block_size, full_dim = kv_cache.shape + hnd = kv_cache.stride()[1:] == (block_size * full_dim, full_dim, 1) + if not hnd and kv_cache.stride()[1:] != (full_dim, num_slots * full_dim, 1): + raise ValueError(f"NVFP4 KV pages must be dense, got {kv_cache.stride()}") + dims = (num_kv_heads, block_size) if hnd else (block_size, num_kv_heads) + side_size = num_kv_heads * block_size * full_dim + (k_data, k_scale), (v_data, v_scale) = ( + nvfp4_split_data_scale( + kv_cache.as_strided( + (num_blocks, *dims, full_dim), + (kv_cache.stride(0), dims[1] * full_dim, full_dim, 1), + kv_cache.storage_offset() + side * side_size, + ) + ) + for side in range(2) + ) + return (k_data, v_data), (k_scale, v_scale), "HND" if hnd else "NHD" + + def _get_trtllm_workspace_buffer(): global trtllm_workspace_buffer if trtllm_workspace_buffer is None: @@ -498,6 +531,9 @@ def get_dtype_for_flashinfer(kv_cache_dtype: str) -> torch.dtype: @classmethod def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool: if kv_cache_dtype is not None and kv_cache_dtype.startswith("nvfp4"): + if _nvfp4_kv_on_fa2(): + # The slot-mapping writer has no scale search. + return kv_cache_dtype == "nvfp4" return ( current_platform.is_device_capability_family(100) and supports_trtllm_attention(is_prefill=True) @@ -505,6 +541,30 @@ def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool: ) return super().supports_kv_cache_dtype(kv_cache_dtype) + @classmethod + def supports_combination( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: CacheDType | None, + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + use_mm_prefix: bool, + device_capability: DeviceCapability, + ) -> str | None: + if ( + kv_cache_dtype is not None + and kv_cache_dtype.startswith("nvfp4") + and device_capability.major in (8, 12) + ): + if head_size == 64: + return "fa2 prefill misreads an NVFP4 KV cache at head_size 64" + if has_sink: + return "the sink prefill wrapper cannot read an NVFP4 KV cache" + return None + @classmethod def get_supported_head_sizes(cls) -> list[int]: # https://github.com/flashinfer-ai/flashinfer/blob/3d55c71a62052c590c130897d3a3db49b14fcc34/include/flashinfer/utils.cuh#L157 @@ -777,7 +837,16 @@ def __init__( # Cannot use self.kv_cache_spec.dtype here because kv_cache_spec # storage dtype may not be the same as the op dtype (uint8 vs fp8_e4m3) self.is_kvcache_nvfp4 = self.cache_dtype.startswith("nvfp4") - if self.is_kvcache_nvfp4: + # fa2 plans NVFP4 by its uint8 storage dtype. + self.nvfp4_fa2 = self.is_kvcache_nvfp4 and _nvfp4_kv_on_fa2() + self.nvfp4_trtllm = self.is_kvcache_nvfp4 and not self.nvfp4_fa2 + if self.nvfp4_fa2 and self.use_dcp: + raise NotImplementedError( + "The DCP prefill wrapper cannot read NVFP4 block scales." + ) + if self.nvfp4_fa2 and self.head_dim == 64: + raise NotImplementedError("fa2 prefill misreads NVFP4 at head_size 64.") + if self.nvfp4_trtllm: if ( force_use_trtllm_attention() is False or not supports_trtllm_attention(is_prefill=True) @@ -798,6 +867,7 @@ def __init__( else: self.cache_dtype = "auto" self.is_kvcache_nvfp4 = False + self.nvfp4_fa2 = self.nvfp4_trtllm = False assert self.kv_cache_spec.dtype == self.model_config.dtype self.kv_cache_dtype = self.kv_cache_spec.dtype @@ -810,8 +880,11 @@ def __init__( # Prefer TRTLLM/XQA for decoding whenever supported. The decode kernel # must be selected statically for FULL cudagraph capture. - can_use_xqa_or_trtllm_gen_decode = can_use_trtllm_attention( - self.num_qo_heads, self.num_kv_heads, is_prefill=False + can_use_xqa_or_trtllm_gen_decode = ( + can_use_trtllm_attention( + self.num_qo_heads, self.num_kv_heads, is_prefill=False + ) + and not self.nvfp4_fa2 ) # Page sizes >= 128 require the trtllm-gen GQA/MQA path (guaranteed by # get_supported_kernel_block_sizes). @@ -894,7 +967,7 @@ def __init__( self.window_left = self.global_hyperparameters.window_left self.logits_soft_cap = self.global_hyperparameters.logits_soft_cap self.has_sinks = self.global_hyperparameters.has_sinks - if self.has_sinks and not FlashInferBackend.supports_sink(): + if self.has_sinks and (self.nvfp4_fa2 or not FlashInferBackend.supports_sink()): raise NotImplementedError( "FlashInfer backend currently does not support attention " "sinks, please use trtllm on blackwell or flash attention on " @@ -970,6 +1043,8 @@ def get_q_data_type(self, is_prefill: bool) -> torch.dtype: return FlashInferBackend.get_dtype_for_flashinfer(cache_dtype) return self.model_config.dtype if cache_dtype.startswith("nvfp4"): + if self.nvfp4_fa2: + return self.model_config.dtype return FlashInferBackend.get_dtype_for_flashinfer("fp8_e4m3") return self.kv_cache_spec.dtype @@ -1002,7 +1077,13 @@ def get_cudagraph_support( num_qo_heads=num_qo_heads, num_kv_heads=spec.num_kv_heads, is_prefill=False, - ) or (is_xqa_arch and not _is_xqa_head_dim_supported(spec.head_size)): + ) or ( + is_xqa_arch + and ( + spec.kv_quant_mode.is_nvfp4 + or not _is_xqa_head_dim_supported(spec.head_size) + ) + ): has_uniform_batch_support = False break @@ -1230,8 +1311,8 @@ def _get_prefill_wrapper( ) else: # NVFP4 KV cache requires the trtllm-gen backend inside - # the wrapper; fa2/fa3 do not support nvfp4. - backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto" + # the wrapper on SM100; SM8x and SM12x read it with fa2. + backend = "trtllm-gen" if self.nvfp4_trtllm else "auto" self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper( self._get_workspace_buffer(), get_flashinfer_layout_string(self.kv_cache_layout), @@ -1256,8 +1337,8 @@ def _get_decode_wrapper(self, batch_size: int, use_cudagraph: bool = False): paged_kv_indices = None paged_kv_last_page_len = None # NVFP4 KV cache requires the trtllm-gen backend inside - # the wrapper; fa2/fa3 do not support nvfp4. - backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto" + # the wrapper on SM100; SM8x and SM12x read it with fa2. + backend = "trtllm-gen" if self.nvfp4_trtllm else "auto" decode_wrapper = BatchDecodeWithPagedKVCacheWrapper( self._get_workspace_buffer(), get_flashinfer_layout_string(self.kv_cache_layout), @@ -1710,7 +1791,7 @@ def build( # use FP8 o_data_type so the wrapper matches the # FP8 output buffer allocated in forward(). o_dtype = ( - FP8_DTYPE if self.is_kvcache_nvfp4 else self.model_config.dtype + FP8_DTYPE if self.nvfp4_trtllm else self.model_config.dtype ) prefill_wrapper.plan( qo_indptr=qo_indptr_prefill_cpu, @@ -1797,9 +1878,7 @@ def build( # NVFP4 trtllm kernel only supports FP8 output; # use FP8 o_data_type so the wrapper matches the # FP8 output buffer allocated in forward(). - o_dtype = ( - FP8_DTYPE if self.is_kvcache_nvfp4 else self.model_config.dtype - ) + o_dtype = FP8_DTYPE if self.nvfp4_trtllm else self.model_config.dtype paged_kv_indptr_cpu = self.paged_kv_indptr.cpu[: num_input_tokens + 1] paged_kv_last_page_len_cpu = self.paged_kv_last_page_len.cpu[ :num_input_tokens @@ -1891,6 +1970,8 @@ def __init__( ) self.cache_dtype = kv_cache_dtype self.is_kvcache_nvfp4 = kv_cache_dtype.startswith("nvfp4") + self.nvfp4_fa2 = self.is_kvcache_nvfp4 and _nvfp4_kv_on_fa2() + self.nvfp4_trtllm = self.is_kvcache_nvfp4 and not self.nvfp4_fa2 self.kv_cache_dtype = "nvfp4" if self.is_kvcache_nvfp4 else kv_cache_dtype self.fp4_data_dim = head_size // 2 if self.is_kvcache_nvfp4 else 0 self.logits_soft_cap = logits_soft_cap @@ -1937,7 +2018,7 @@ def __init__( self.o_sf_scale: float | None = None # Pre-allocated FP8 output buffer for NVFP4 without fused output quant. - if self.is_kvcache_nvfp4 and vllm_config is not None: + if self.nvfp4_trtllm and vllm_config is not None: max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens self._nvfp4_fp8_out = torch.empty( (max_num_tokens, num_heads, head_size), @@ -2195,10 +2276,9 @@ def forward( canonicalize_singleton_dim_strides(k_cache.permute(*stride_order)), canonicalize_singleton_dim_strides(v_cache.permute(*stride_order)), ) - k_data, k_sf = nvfp4_split_data_scale(kv_cache_tuple[0]) - v_data, v_sf = nvfp4_split_data_scale(kv_cache_tuple[1]) - nvfp4_kv_data = (k_data, v_data) - nvfp4_kv_block_scales = (k_sf, v_sf) + nvfp4_kv_data, nvfp4_kv_block_scales, _ = _nvfp4_kv_views( + kv_cache, self.num_kv_heads + ) else: kv_cache_tuple = kv_cache_permute.split(hs, dim=-1) @@ -2276,7 +2356,7 @@ def forward( # Use a pre-allocated FP8 buffer and dequantize # afterwards. needs_fp8_out_prefill = ( - self.is_kvcache_nvfp4 and output.dtype != FP8_DTYPE + self.nvfp4_trtllm and output.dtype != FP8_DTYPE ) if needs_fp8_out_prefill: out_prefill = self._nvfp4_fp8_out[:num_prefill_tokens] @@ -2299,7 +2379,8 @@ def forward( prefill_wrapper.run( prefill_query, kv_cache_for_fi, - q_scale=layer._q_scale_float, + # The fa2 NVFP4 query stays in the model dtype. + q_scale=1.0 if self.nvfp4_fa2 else layer._q_scale_float, k_scale=layer._k_scale_float, v_scale=layer._v_scale_float, out=out_prefill, @@ -2451,7 +2532,7 @@ def forward( # NVFP4 kernel only supports FP8 output. # Use a pre-allocated FP8 buffer and dequantize afterwards. - needs_fp8_out = self.is_kvcache_nvfp4 and output.dtype != FP8_DTYPE + needs_fp8_out = self.nvfp4_trtllm and output.dtype != FP8_DTYPE if needs_fp8_out: out_decode = self._nvfp4_fp8_out[:num_decode_tokens] else: @@ -2488,7 +2569,7 @@ def forward( decode_wrapper.run( decode_query, kv_cache_for_fi, - q_scale=layer._q_scale_float, + q_scale=1.0 if self.nvfp4_fa2 else layer._q_scale_float, k_scale=layer._k_scale_float, v_scale=layer._v_scale_float, out=out_decode, @@ -2668,6 +2749,19 @@ def do_kv_cache_update( # and value[:num_actual_tokens] because the reshape_and_cache_flash # op uses the slot_mapping's shape to determine the number of # actual tokens. + if self.nvfp4_fa2: + data, scales, kv_layout = _nvfp4_kv_views(kv_cache, self.num_kv_heads) + nvfp4_quantize_append_paged_kv_cache_with_slot_mapping( + key, + value, + slot_mapping, + data, + scales, + layer._k_scale, + layer._v_scale, + kv_layout=kv_layout, + ) + return if self.is_kvcache_nvfp4: # (B, 2*H, N, full_dim) -> ((B, N, H, full_dim), # (B, N, H, full_dim)); From fba462e008f0faa16745d1539c339db692d41f13 Mon Sep 17 00:00:00 2001 From: Eduardo Fonseca Date: Wed, 7 Oct 2026 14:01:21 -0400 Subject: [PATCH 029/624] [Bugfix]Fix FlashInfer warmup crash: autotune(tuning_buckets=...) must round M up, not down (Invalid MXFP8 split-K tactic with spec-decode) (#58165) Signed-off-by: Yongye Zhu Co-authored-by: Yongye Zhu Co-authored-by: Claude Opus 5.5 --- .../test_flashinfer_autotune_warmup.py | 31 +++++++++++++++++-- vllm/model_executor/warmup/kernel_warmup.py | 23 ++++++++++++-- 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py index d65f2a4ad116..fe89b09d09bc 100644 --- a/tests/model_executor/test_flashinfer_autotune_warmup.py +++ b/tests/model_executor/test_flashinfer_autotune_warmup.py @@ -115,8 +115,8 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(skip_attn): assert get_buckets.call_args_list == [call(8192), call(128)] assert autotune.call_args_list == [ - call(tuning_buckets=max_buckets), - call(tuning_buckets=deferred_buckets), + call(tuning_buckets=max_buckets, round_up=True), + call(tuning_buckets=deferred_buckets, round_up=True), ] assert runner._dummy_run.call_args_list == [ call( @@ -136,6 +136,33 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(skip_attn): ] +def test_flashinfer_autotune_buckets_cover_drafter_tokens(): + """The drafter's M can exceed the pass size; buckets must include it.""" + runner = _make_runner([]) + runner.max_num_reqs = 256 + runner.max_num_tokens = 8192 + runner.speculator = SimpleNamespace(num_query_per_req=6) + + with ( + patch( + "vllm.model_executor.warmup.kernel_warmup." + "_flashinfer_autotune_token_counts", + return_value=(8192, 128), + ), + patch( + "vllm.utils.flashinfer.flashinfer_get_hybrid_num_tokens_buckets" + ) as get_buckets, + patch("vllm.utils.flashinfer.autotune"), + ): + _run_flashinfer_autotune_dummy_runs(runner) + + assert get_buckets.call_args_list == [call(8192), call(128 * 6)] + assert [c.kwargs["num_tokens"] for c in runner._dummy_run.call_args_list] == [ + 8192, + 128, + ] + + class _AutotuneGroup: def __init__(self, run, ranks): self.run = run diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 06f04b75d959..d6544333beb5 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -363,13 +363,32 @@ def _run_flashinfer_autotune_dummy_runs( dummy_run_kwargs = {"skip_attn": True} if skip_attn else {} for num_tokens in _flashinfer_autotune_token_counts(runner, include_bf16=False): - tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens) + # The drafter may use more tokens than this pass in the same dummy run. + # Include its max M when building the autotune buckets. + max_tuning_tokens = num_tokens + speculator = getattr(runner, "speculator", None) + if speculator is not None: + num_query_per_req = speculator.num_query_per_req + num_draft_reqs = min( + num_tokens, + runner.max_num_reqs, + runner.max_num_tokens // num_query_per_req, + ) + max_tuning_tokens = max( + max_tuning_tokens, num_draft_reqs * num_query_per_req + ) + + tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets( + max_tuning_tokens + ) logger.info( "Running FlashInfer autotune with %d tokens and token buckets %s.", num_tokens, tuning_buckets, ) - with fi_utils.autotune(tuning_buckets=tuning_buckets): + # Round M up to match serving-time bucket selection for non-bucket M. + # See https://github.com/flashinfer-ai/flashinfer/issues/5450 + with fi_utils.autotune(tuning_buckets=tuning_buckets, round_up=True): runner._dummy_run( num_tokens=num_tokens, skip_eplb=True, From 6ee9be6b3c2dd5557ad72cf714b9d8fe3c2009d4 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Wed, 7 Oct 2026 11:01:59 -0700 Subject: [PATCH 030/624] [Spec Decode] Rename AR speculators to StandaloneAR / TargetDependentAR (#60335) Signed-off-by: Giancarlo Delfin --- .buildkite/hardware_tests/cpu.yaml | 2 +- tests/v1/cudagraph/test_cudagraph_manager.py | 4 ++-- ... => test_gpu_target_dependent_ar_speculator.py} | 10 +++++----- vllm/v1/worker/gpu/model_runner.py | 2 +- vllm/v1/worker/gpu/spec_decode/__init__.py | 6 +++--- vllm/v1/worker/gpu/spec_decode/eagle/speculator.py | 8 ++++---- .../v1/worker/gpu/spec_decode/gemma4/speculator.py | 6 +++--- vllm/v1/worker/gpu/spec_decode/mtp/speculator.py | 8 ++++---- .../gpu/spec_decode/multi_module_mtp/speculator.py | 6 +++--- .../{autoregressive => standalone_ar}/__init__.py | 0 .../{draft_model => standalone_ar}/speculator.py | 14 +++++++------- .../__init__.py | 0 .../cudagraph_utils.py | 0 .../speculator.py | 12 +++++++++--- 14 files changed, 42 insertions(+), 36 deletions(-) rename tests/v1/worker/{test_gpu_autoregressive_speculator.py => test_gpu_target_dependent_ar_speculator.py} (98%) rename vllm/v1/worker/gpu/spec_decode/{autoregressive => standalone_ar}/__init__.py (100%) rename vllm/v1/worker/gpu/spec_decode/{draft_model => standalone_ar}/speculator.py (97%) rename vllm/v1/worker/gpu/spec_decode/{draft_model => target_dependent_ar}/__init__.py (100%) rename vllm/v1/worker/gpu/spec_decode/{autoregressive => target_dependent_ar}/cudagraph_utils.py (100%) rename vllm/v1/worker/gpu/spec_decode/{autoregressive => target_dependent_ar}/speculator.py (99%) diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index fbd43d80ddb2..3fdce975d3f4 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -87,7 +87,7 @@ steps: - | bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " pytest -x -v -s \ - tests/v1/worker/test_gpu_autoregressive_speculator.py \ + tests/v1/worker/test_gpu_target_dependent_ar_speculator.py \ tests/v1/worker/test_gpu_extract_hidden_states_speculator.py \ tests/v1/worker/test_gpu_batch_ordering.py \ tests/v1/worker/test_gpu_rejection_sampler_chunking.py \ diff --git a/tests/v1/cudagraph/test_cudagraph_manager.py b/tests/v1/cudagraph/test_cudagraph_manager.py index 668fbdf61b4a..1e4071bcfbfd 100644 --- a/tests/v1/cudagraph/test_cudagraph_manager.py +++ b/tests/v1/cudagraph/test_cudagraph_manager.py @@ -23,10 +23,10 @@ from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp from vllm.v1.worker.gpu.input_batch import InputBuffers from vllm.v1.worker.gpu.pcp_manager import PCPManager -from vllm.v1.worker.gpu.spec_decode.autoregressive import ( +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar import ( cudagraph_utils as spec_cudagraph_utils, ) -from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import ( +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.cudagraph_utils import ( SpeculatorCudaGraphManager, ) diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py similarity index 98% rename from tests/v1/worker/test_gpu_autoregressive_speculator.py rename to tests/v1/worker/test_gpu_target_dependent_ar_speculator.py index 0a8cfd1a130e..20e86f4a7827 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py @@ -22,17 +22,17 @@ from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor from vllm.v1.worker.gpu.spec_decode import speculator as base_spec_module -from vllm.v1.worker.gpu.spec_decode.autoregressive import speculator as spec_module -from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import ( - AutoRegressiveSpeculator, -) from vllm.v1.worker.gpu.spec_decode.multi_module_mtp.speculator import ( MultiModuleMTPSpeculator, ) from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar import speculator as spec_module +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.speculator import ( + TargetDependentARSpeculator, +) -class _TestSpeculator(AutoRegressiveSpeculator): +class _TestSpeculator(TargetDependentARSpeculator): def load_draft_model(self, target_model, target_attn_layer_names): return self.test_draft_model diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index d0c08707f483..373042a41502 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -820,7 +820,7 @@ def _dummy_run( "skip_attn must only be True for initial memory profiling." ) - # Create a dummy scheduler output. Plain draft-model speculation adds + # Create a dummy scheduler output. Standalone AR speculation adds # one correction slot per request during prefill. The scheduler # accounts for these slots, while dummy runs bypass the scheduler. # Adjust the dummy token count only when the expanded draft batch diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index 902a4f085ca9..fc028941cf0d 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/vllm/v1/worker/gpu/spec_decode/__init__.py @@ -71,11 +71,11 @@ def init_speculator( return EagleSpeculator(vllm_config, device) elif speculative_config.uses_draft_model(): - from vllm.v1.worker.gpu.spec_decode.draft_model.speculator import ( - PlainDraftModelSpeculator, + from vllm.v1.worker.gpu.spec_decode.standalone_ar.speculator import ( + StandaloneARSpeculator, ) - return PlainDraftModelSpeculator(vllm_config, device) + return StandaloneARSpeculator(vllm_config, device) elif speculative_config.use_ngram(): from vllm.v1.worker.gpu.spec_decode.ngram.speculator import ( NgramGPUSpeculator, diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index e878872e622b..0611d42b1ce5 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -3,13 +3,13 @@ import torch.nn as nn -from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import ( - AutoRegressiveSpeculator, -) from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.speculator import ( + TargetDependentARSpeculator, +) -class EagleSpeculator(AutoRegressiveSpeculator): +class EagleSpeculator(TargetDependentARSpeculator): def load_draft_model( self, target_model: nn.Module, diff --git a/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py b/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py index 6f76e5055a82..9d553830ae17 100644 --- a/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py @@ -17,8 +17,8 @@ from vllm.logger import init_logger from vllm.model_executor.model_loader import get_model from vllm.model_executor.model_loader.utils import get_draft_load_config -from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import ( - AutoRegressiveSpeculator, +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.speculator import ( + TargetDependentARSpeculator, ) logger = init_logger(__name__) @@ -45,7 +45,7 @@ def _copy_target_kv_scales(attn: nn.Module, target_attn: nn.Module) -> None: getattr(attn, scale_name).copy_(getattr(target_attn, scale_name)) -class Gemma4Speculator(AutoRegressiveSpeculator): +class Gemma4Speculator(TargetDependentARSpeculator): @property def advance_draft_positions(self) -> bool: # Gemma4 MTP is Q-only and reads K/V from the target's existing cache. diff --git a/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py index d94702612d9b..a0064197f756 100644 --- a/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py @@ -3,13 +3,13 @@ import torch.nn as nn -from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import ( - AutoRegressiveSpeculator, -) from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.speculator import ( + TargetDependentARSpeculator, +) -class MTPSpeculator(AutoRegressiveSpeculator): +class MTPSpeculator(TargetDependentARSpeculator): share_mtp_topk_indices: bool = False def load_draft_model( diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py index 71f7f89416be..919e939c03cc 100644 --- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py @@ -14,11 +14,11 @@ from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers -from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import ( - SpeculatorCudaGraphManager, -) from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.cudagraph_utils import ( + SpeculatorCudaGraphManager, +) from vllm.v1.worker.utils import get_uniform_decode_token_count logger = init_logger(__name__) diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/__init__.py b/vllm/v1/worker/gpu/spec_decode/standalone_ar/__init__.py similarity index 100% rename from vllm/v1/worker/gpu/spec_decode/autoregressive/__init__.py rename to vllm/v1/worker/gpu/spec_decode/standalone_ar/__init__.py diff --git a/vllm/v1/worker/gpu/spec_decode/draft_model/speculator.py b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py similarity index 97% rename from vllm/v1/worker/gpu/spec_decode/draft_model/speculator.py rename to vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py index 5be7b0c86ddd..8bdfb8530be6 100644 --- a/vllm/v1/worker/gpu/spec_decode/draft_model/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py @@ -26,22 +26,22 @@ logger = init_logger(__name__) -class PlainDraftModelSpeculator(DraftModelSpeculator): - """Speculative decoding using a separate smaller draft LM. +class StandaloneARSpeculator(DraftModelSpeculator): + """Standalone autoregressive (AR) drafting with a separate smaller LM. Unlike Eagle, the draft model runs fully independently of the target model. Step 0 builds an expanded buffer (accepted + correction token + rejected slots) via a Triton kernel; steps 1..k-1 are single-token decode steps. """ - # Plain draft model needs one extra slot per request for correction. + # Standalone AR drafting needs one extra slot per request for correction. num_extra_query_per_req = 1 def __init__(self, vllm_config: VllmConfig, device: torch.device): super().__init__(vllm_config, device) # draft_max_seq_len is read by the parent's attention metadata builder. - # Plain draft model doesn't do per-batch adjustment; cap at max. + # Standalone AR drafting doesn't do per-batch adjustment; cap at max. self.draft_max_seq_len = self.max_model_len self.last_token_indices = torch.zeros( @@ -65,10 +65,10 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.supports_mm_inputs = False def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: - pass # CUDA graph not yet supported for plain draft model speculator. + pass # CUDA graph not yet supported for standalone AR speculator. def capture(self) -> None: - pass # CUDA graph not yet supported for plain draft model speculator. + pass # CUDA graph not yet supported for standalone AR speculator. def set_attn( self, @@ -415,7 +415,7 @@ def _prepare_prefill_inputs_kernel( max_model_len, BLOCK_SIZE: tl.constexpr, ): - """Per-request step-0 input preparation for the plain draft-model speculator. + """Per-request step-0 input preparation for the standalone AR speculator. Output layout for request i (out_start = query_start_loc[i] + i): [out_start, out_start + num_valid) accepted tokens diff --git a/vllm/v1/worker/gpu/spec_decode/draft_model/__init__.py b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/__init__.py similarity index 100% rename from vllm/v1/worker/gpu/spec_decode/draft_model/__init__.py rename to vllm/v1/worker/gpu/spec_decode/target_dependent_ar/__init__.py diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/cudagraph_utils.py similarity index 100% rename from vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py rename to vllm/v1/worker/gpu/spec_decode/target_dependent_ar/cudagraph_utils.py diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py similarity index 99% rename from vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py rename to vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py index 46ff82681561..9a8ed7b8b4ca 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py @@ -17,16 +17,22 @@ from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState -from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import ( +from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.cudagraph_utils import ( SpeculatorCudaGraphManager, ) -from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator from vllm.v1.worker.utils import AttentionGroup, get_uniform_decode_token_count logger = init_logger(__name__) -class AutoRegressiveSpeculator(DraftModelSpeculator): +class TargetDependentARSpeculator(DraftModelSpeculator): + """Autoregressive (AR) drafting that depends on target model state. + + Base for Eagle, MTP, and Gemma4 drafters, which consume the target's + hidden states (and, for Gemma4, its KV cache) at each draft step. + """ + def __init__(self, vllm_config: VllmConfig, device: torch.device): super().__init__(vllm_config, device) From 32fbfa15e8bacc64182cb1286831bd63d7e4fc12 Mon Sep 17 00:00:00 2001 From: Jinzhen Lin Date: Thu, 8 Oct 2026 02:04:43 +0800 Subject: [PATCH 031/624] [Bugfix][Quantization] Accept folded NVFP4 scales for Humming A16 (#60337) Signed-off-by: jinzhen.ljz Co-authored-by: Kevin H. Luu --- .../quantization/utils/humming/schema.py | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/quantization/utils/humming/schema.py b/vllm/model_executor/layers/quantization/utils/humming/schema.py index c84e8a370d22..aaf6d4a66c87 100644 --- a/vllm/model_executor/layers/quantization/utils/humming/schema.py +++ b/vllm/model_executor/layers/quantization/utils/humming/schema.py @@ -322,6 +322,22 @@ def check_and_fallback_input_schema( input_group_size = input_schema.input_scale_group_size input_scale_dtype = input_schema.input_scale_dtype + def is_compatible(candidate: "HummingInputSchema") -> bool: + check_weight_schema = weight_schema + if ( + candidate.a_dtype in (None, humming_dtypes.float16, humming_dtypes.bfloat16) + and weight_schema.b_dtype == humming_dtypes.float4e2m1 + and weight_group_size == 16 + and weight_schema.bs_dtype + in (None, humming_dtypes.float16, humming_dtypes.bfloat16) + ): + # Humming rejects folded NVFP4 scales despite A16 software dequant + # supporting them. Relax only this check, preserving kernel metadata. + check_weight_schema = dataclasses.replace( + weight_schema, bs_dtype=humming_dtypes.float8e4m3 + ) + return candidate.is_compatible_with(check_weight_schema, param_dtype) + def is_deprecated(dtype: "humming_dtypes.DataType | None") -> bool: if dtype is None or dtype.num_bits == 16: return False @@ -333,7 +349,7 @@ def is_deprecated(dtype: "humming_dtypes.DataType | None") -> bool: is_fp8_deprecated = dtype.num_bits < 16 and 100 <= sm_version < 120 return is_int4_deprecated or is_int8_deprecated or is_fp8_deprecated - if input_schema.is_compatible_with(weight_schema, param_dtype): + if is_compatible(input_schema): if not allow_fallback: if is_deprecated(a_dtype): msg = f"Humming {a_dtype} activation is slow on SM{sm_version} now." @@ -390,7 +406,7 @@ def is_deprecated(dtype: "humming_dtypes.DataType | None") -> bool: input_scale_dtype=scale_dtype, input_quant_mode=quant_mode, ) - if candidate.is_compatible_with(weight_schema, param_dtype): + if is_compatible(candidate): return candidate raise ValueError( From b9767299db6d5c0226e6879dadb079a477524d3a Mon Sep 17 00:00:00 2001 From: Cheese Cake Date: Thu, 8 Oct 2026 00:38:11 +0530 Subject: [PATCH 032/624] [Bugfix] Respect Inductor deterministic mode in combo-kernel defaults (#59074) Signed-off-by: cheesecake Co-authored-by: OpenAI Codex Co-authored-by: Claude Opus 5.5 --- docs/design/debug_vllm_compile.md | 12 ++++++ tests/compile/test_config.py | 67 +++++++++++++++++++++++++++++++ vllm/config/compilation.py | 8 +++- 3 files changed, 86 insertions(+), 1 deletion(-) diff --git a/docs/design/debug_vllm_compile.md b/docs/design/debug_vllm_compile.md index 7edda6fa6476..af9b3f530d4c 100644 --- a/docs/design/debug_vllm_compile.md +++ b/docs/design/debug_vllm_compile.md @@ -251,6 +251,18 @@ that may call 1+ triton kernels. On rare (but unfortunate) occasions, it may produce an incorrect triton kernel. This may manifest as silent incorrectness, CUDA illegal memory accesses, or loud errors. +### Deterministic compilation + +When vLLM selects combo-kernel defaults, it disables combo-kernel benchmarking +if Inductor's `deterministic` setting is enabled. A `deterministic` value in +`inductor_compile_config` takes precedence over the global Inductor setting. +On torch 2.13, `TORCHINDUCTOR_DETERMINISTIC=1` alone is reset by the first +compiled frame ([pytorch#198563](https://github.com/pytorch/pytorch/issues/198563)), +so set `deterministic` in `inductor_compile_config` instead. +Explicit `combo_kernels` or `benchmark_combo_kernel` settings are preserved. +This avoids conflicting compilation defaults; it does not guarantee reproducible +model outputs. + ### Inductor runtime assertions By default (on torch < 2.12), vLLM disables Inductor's runtime assertions diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index 24cb816b5d9d..5a1ba13a57f1 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -1366,3 +1366,70 @@ def test_inductor_asserts_user_override(monkeypatch): assert config.inductor_compile_config.get("size_asserts") is True if not _is_torch_equal_or_newer(torch.__version__, "2.12.0.dev"): assert config.inductor_compile_config.get("alignment_asserts") is False + + +@pytest.mark.parametrize("deterministic", [False, True]) +@pytest.mark.parametrize("override", [None, False, True]) +def test_combo_kernel_benchmarking_respects_deterministic(deterministic, override): + from torch._inductor import config as inductor_config + + overrides = {} if override is None else {"deterministic": override} + with ( + inductor_config.patch(deterministic=deterministic), + patch("vllm.config.compilation.current_platform.is_cpu", return_value=False), + ): + config = CompilationConfig(inductor_compile_config=overrides) + + assert config.inductor_compile_config["combo_kernels"] is True + effective_deterministic = deterministic if override is None else override + assert config.inductor_compile_config["benchmark_combo_kernel"] is ( + not effective_deterministic + ) + + +@pytest.mark.parametrize("deterministic", [False, True]) +@pytest.mark.parametrize( + "overrides", + [ + {"combo_kernels": False}, + {"combo_kernels": True}, + {"benchmark_combo_kernel": False}, + {"benchmark_combo_kernel": True}, + ], +) +def test_combo_kernel_explicit_settings_preserved(deterministic, overrides): + from torch._inductor import config as inductor_config + + with ( + inductor_config.patch(deterministic=deterministic), + patch("vllm.config.compilation.current_platform.is_cpu", return_value=False), + ): + config = CompilationConfig(inductor_compile_config=overrides.copy()) + + for key in ("combo_kernels", "benchmark_combo_kernel"): + assert config.inductor_compile_config.get(key) == overrides.get(key) + + +@pytest.mark.parametrize("is_cpu,torch_version", [(True, "2.13.0"), (False, "2.8.0")]) +def test_combo_kernel_defaults_require_supported_platform_and_torch( + is_cpu, torch_version +): + with ( + patch("vllm.config.compilation.current_platform.is_cpu", return_value=is_cpu), + patch("torch.__version__", torch_version), + ): + config = CompilationConfig() + + assert "combo_kernels" not in config.inductor_compile_config + assert "benchmark_combo_kernel" not in config.inductor_compile_config + + +def test_combo_kernel_defaults_without_inductor_deterministic_setting(): + with ( + patch("torch._inductor.config", SimpleNamespace()), + patch("vllm.config.compilation.current_platform.is_cpu", return_value=False), + ): + config = CompilationConfig() + + assert config.inductor_compile_config["combo_kernels"] is True + assert config.inductor_compile_config["benchmark_combo_kernel"] is True diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index afe576daa9d3..cbd132ab5b68 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -1002,10 +1002,16 @@ def __post_init__(self) -> None: # (fixme @boyuan) combo kernel does not support cpu yet. and not current_platform.is_cpu() ): + from torch._inductor import config as inductor_config + # use horizontal fusion, which is useful for fusing qk-norm and # qk-rope when query and key have different shapes. self.inductor_compile_config["combo_kernels"] = True - self.inductor_compile_config["benchmark_combo_kernel"] = True + + deterministic = self.inductor_compile_config.get( + "deterministic", getattr(inductor_config, "deterministic", False) + ) + self.inductor_compile_config["benchmark_combo_kernel"] = not deterministic if self.use_inductor_graph_partition and not is_torch_equal_or_newer( "2.9.0.dev" From 17c70c895359f7c941199bdf49b723c602ba75de Mon Sep 17 00:00:00 2001 From: Moein Khazraee <33970824+mkhazraee@users.noreply.github.com> Date: Wed, 7 Oct 2026 12:14:07 -0700 Subject: [PATCH 033/624] [KV Offload] Update KVCR adapter to pool-and-index API (#59899) --- .buildkite/scripts/install-kv-offload.sh | 2 +- tests/v1/kv_offload/tiering/test_kvcr_tier.py | 32 ++++++++------- vllm/v1/kv_offload/tiering/kvcr/manager.py | 40 +++++++++---------- 3 files changed, 36 insertions(+), 38 deletions(-) diff --git a/.buildkite/scripts/install-kv-offload.sh b/.buildkite/scripts/install-kv-offload.sh index 46d756b621b1..eb8bab9e393f 100644 --- a/.buildkite/scripts/install-kv-offload.sh +++ b/.buildkite/scripts/install-kv-offload.sh @@ -10,7 +10,7 @@ fi # TODO: Use the PyPI wheel once it is published. uv pip install --system \ - "kvcr @ git+https://github.com/ai-dynamo/kvcr.git@v0.1.0" + "kvcr @ git+https://github.com/ai-dynamo/kvcr.git@v0.1.1" # Keep only the NIXL wheel matching the CI image's CUDA runtime. NIXL_VERSION=$(uv pip show --system nixl | sed -n 's/^Version: //p') diff --git a/tests/v1/kv_offload/tiering/test_kvcr_tier.py b/tests/v1/kv_offload/tiering/test_kvcr_tier.py index 5cbc85230ffb..fc4421eb06b2 100644 --- a/tests/v1/kv_offload/tiering/test_kvcr_tier.py +++ b/tests/v1/kv_offload/tiering/test_kvcr_tier.py @@ -18,7 +18,7 @@ BlockKey, CacheTier, InventoryEvent, - MemDescriptor, + MemoryRef, OpEntryResult, OpEntryStatus, OpHandle, @@ -78,11 +78,9 @@ def __init__(self) -> None: self.submit_hint_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] self.discard_hint_calls: list[str] = [] self.deliver_calls: list[ - tuple[OpHandle, dict[BlockKey, list[MemDescriptor]], str | None] - ] = [] - self.deposit_calls: list[ - tuple[OpHandle, dict[BlockKey, list[MemDescriptor]]] + tuple[OpHandle, dict[BlockKey, list[MemoryRef]], str | None] ] = [] + self.deposit_calls: list[tuple[OpHandle, dict[BlockKey, list[MemoryRef]]]] = [] self.completed: list[tuple[OpHandle, dict[BlockKey, OpEntryResult]]] = [] self._next_op_handle = 1 @@ -107,7 +105,7 @@ def query( def deliver( self, - blocks: Mapping[BlockKey, list[MemDescriptor]], + blocks: Mapping[BlockKey, list[MemoryRef]], request_id: str | None = None, ) -> OpHandle: op_handle = self._next_op_handle @@ -118,7 +116,7 @@ def deliver( def deposit( self, - blocks: Mapping[BlockKey, list[MemDescriptor]], + blocks: Mapping[BlockKey, list[MemoryRef]], ) -> OpHandle: op_handle = self._next_op_handle self._next_op_handle += 1 @@ -248,7 +246,11 @@ def test_kvcr_tier_configures_service_for_local_dp_rank(monkeypatch): compatibility_digest="Opaque-Digest", ) assert kvcr.config is not None - assert kvcr.config.pool_layouts == [("", tier._primary_row_stride)] + assert kvcr.config.pool_layouts == [("", tier.block_size_bytes)] + (region,) = kvcr.backend_configs.framework_regions + assert region.addr == np.asarray(tier._primary_kv_view).ctypes.data + assert (region.size, region.stride or region.size, region.count) == (16, 16, 4) + assert (region.label, region.mem_type, region.device_Id) == ("", "DRAM", 0) assert kvcr.backend_configs.local_dram is None @@ -306,9 +308,7 @@ def test_kvcr_tier_adapts_request_and_load(monkeypatch): assert request_id == "req" assert list(blocks) == [key] (descriptor,) = blocks[key] - assert descriptor.end_point_name == kvcr.nixl_agent_name - assert descriptor.addr == tier._primary_base_addr + 2 * 16 - assert descriptor.size == 16 + assert descriptor == MemoryRef(end_point_name=kvcr.nixl_agent_name, element_index=2) assert list(tier.get_finished_jobs()) == [JobResult(7, True)] # Here we verify that request cleanup discards the request-scoped hint in @@ -351,7 +351,7 @@ def test_kvcr_tier_serves_primary_pin_request(monkeypatch): tier = _make_tier(monkeypatch, kvcr) keys = (BlockKey(b"k0"), BlockKey(b"k1"), BlockKey(b"k2")) hit_keys = (keys[0], keys[2]) - chunk_ids = {keys[0]: 1, keys[2]: 5} + chunk_ids = {keys[0]: 1, keys[2]: 3} lifecycle: list[str] = [] class Parent: @@ -388,7 +388,7 @@ def on_request_finished(self, req_context): block_descriptors = descriptors[keys[2]] assert block_descriptors is not None (descriptor,) = block_descriptors - assert descriptor.addr == (tier._primary_base_addr + 5 * tier._primary_row_stride) + assert descriptor == MemoryRef(end_point_name=kvcr.nixl_agent_name, element_index=3) assert lifecycle == ["new", "finished"] polls = 0 @@ -482,14 +482,16 @@ def test_kvcr_tier_stores_and_emits_inventory(monkeypatch): local_dram = kvcr.backend_configs.local_dram assert local_dram is not None assert [(name, size) for name, _, size in local_dram.pools] == [ - ("", 2 * tier._primary_row_stride) + ("", 2 * tier.block_size_bytes) ] key = OffloadKey(b"k0") tier.submit_store(_job(11, ReqContext(req_id="req"), key=key, chunk_id=2)) _, blocks = kvcr.deposit_calls[0] - assert blocks[key][0].addr == tier._primary_base_addr + 2 * 16 + assert blocks[key] == [ + MemoryRef(end_point_name=kvcr.nixl_agent_name, element_index=2) + ] assert list(tier.get_finished_jobs()) == [JobResult(11, True)] assert kvcr.inventory_sink is not None diff --git a/vllm/v1/kv_offload/tiering/kvcr/manager.py b/vllm/v1/kv_offload/tiering/kvcr/manager.py index 18da2a6c2764..4b6f81b3b71e 100644 --- a/vllm/v1/kv_offload/tiering/kvcr/manager.py +++ b/vllm/v1/kv_offload/tiering/kvcr/manager.py @@ -22,7 +22,6 @@ KVCRBindings, ) from kvcr.config import ( - FrameworkDramInput, G3Options, KeyAdapter, KVCRBackendConfigs, @@ -43,11 +42,12 @@ BlockKey, CacheTier, InventoryEvent, - MemDescriptor, + MemoryRef, OpHandle, PinRequestId, PinResult, QueryStatus, + RegionDescriptor, ) from typing_extensions import override @@ -177,7 +177,7 @@ def cancel_pin_request(self, request: PinRequestId) -> None: def _resolve_pin( self, parent: ParentManager, keys: Collection[BlockKey] - ) -> tuple[str, dict[BlockKey, list[MemDescriptor] | None]] | None: + ) -> tuple[str, dict[BlockKey, list[MemoryRef] | None]] | None: offload_keys = tuple(OffloadKey(bytes(key)) for key in keys) request_id = f"kvcr-source:{self._next_request_id}" self._next_request_id += 1 @@ -202,7 +202,7 @@ def _resolve_pin( if len(job_keys) != len(chunk_ids) or set(job_keys) != set(hit_keys): return None - descriptors: dict[BlockKey, list[MemDescriptor] | None] = { + descriptors: dict[BlockKey, list[MemoryRef] | None] = { BlockKey(bytes(key)): None for key in offload_keys } descriptors.update( @@ -350,12 +350,6 @@ def __init__( with socket.socket() as _s: _s.bind(("", 0)) _nixl_listen_port = _s.getsockname()[1] - self._primary_base_addr = ctypes.addressof( - ctypes.c_char.from_buffer(primary_kv_view) - ) - if primary_kv_view.strides is None: - raise ValueError("primary KV memoryview must expose strides") - self._primary_row_stride = int(primary_kv_view.strides[0]) if secondary_g2_slots < 0: raise ValueError("secondary_g2_slots must be non-negative") if ( @@ -375,7 +369,7 @@ def __init__( ) else: local_mapping = mmap.mmap( - -1, secondary_g2_slots * self._primary_row_stride + -1, secondary_g2_slots * self.block_size_bytes ) local_dram = LocalDramOptions( pools=[ @@ -408,7 +402,7 @@ def __init__( try: self._kvcr = KVCR( KVCRConfig( - pool_layouts=[("", self._primary_row_stride)], + pool_layouts=[("", self.block_size_bytes)], enable_telemetry=enable_telemetry, operation_timeout_ms=operation_timeout_ms, abandon_timeout_ms=abandon_timeout_ms, @@ -429,9 +423,15 @@ def __init__( policy=selected_policy, ), KVCRBackendConfigs( - framework_dram=FrameworkDramInput( - self._primary_base_addr, primary_kv_view.nbytes - ), + framework_regions=[ + RegionDescriptor( + addr=ctypes.addressof( + ctypes.c_char.from_buffer(primary_kv_view) + ), + size=self.block_size_bytes, + count=len(primary_kv_view), + ) + ], local_dram=local_dram, g3=g3_config, remote_fw_dram=RemoteFWDramOptions( @@ -621,12 +621,8 @@ def _record_inventory(self, event: InventoryEvent) -> None: ) ) - def _make_descriptor(self, chunk_id: int) -> MemDescriptor: - return MemDescriptor( + def _make_descriptor(self, chunk_id: int) -> MemoryRef: + return MemoryRef( end_point_name=self._kvcr.config.nixl_agent_name, - mem_type="DRAM", - addr=self._primary_base_addr + chunk_id * self._primary_row_stride, - size=self._primary_row_stride, - device_Id=0, - info="", + element_index=chunk_id, ) From e572e02d21b2b3559b3ce33adb117dca3f50ccd7 Mon Sep 17 00:00:00 2001 From: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Date: Wed, 7 Oct 2026 15:16:19 -0400 Subject: [PATCH 034/624] [UX] Skip checkpoint shards an MTP head does not load (#60152) Signed-off-by: Robert Shaw Co-authored-by: Robert Shaw Co-authored-by: Claude Opus 5.5 --- .../test_filter_duplicate_safetensors.py | 60 +++++++++++++++++++ .../model_loader/default_loader.py | 12 +++- .../model_loader/weight_utils.py | 36 +++++++++++ vllm/model_executor/models/nemotron_h_mtp.py | 15 +++-- vllm/model_executor/models/qwen3_5_mtp.py | 14 +++-- vllm/model_executor/models/qwen3_next_mtp.py | 11 ++-- vllm/models/deepseek_v32/amd/mtp.py | 8 ++- vllm/models/deepseek_v32/nvidia/mtp.py | 8 ++- vllm/models/deepseek_v4/nvidia/mtp.py | 40 +++++++------ 9 files changed, 165 insertions(+), 39 deletions(-) diff --git a/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py b/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py index babc6c061a9c..da04d0108fdf 100644 --- a/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py +++ b/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py @@ -4,11 +4,17 @@ import json import os import tempfile +from types import SimpleNamespace import pytest +import torch +from safetensors.torch import save_file +from vllm.config.load import LoadConfig +from vllm.model_executor.model_loader import DefaultModelLoader from vllm.model_executor.model_loader.weight_utils import ( filter_duplicate_safetensors_files, + filter_safetensors_files_by_weight_name, ) @@ -77,3 +83,57 @@ def test_filter_duplicate_safetensors_files_all_exist(): if __name__ == "__main__": test_filter_duplicate_safetensors_files_missing_weight() test_filter_duplicate_safetensors_files_all_exist() + + +def _skip_non_mtp(name: str) -> bool: + return not name.startswith("mtp.") + + +@pytest.fixture +def mtp_checkpoint(tmp_path): + """Three shards: target only, target + MTP, MTP only.""" + shards = { + "model-00001-of-00003.safetensors": ["model.layers.0.weight"], + "model-00002-of-00003.safetensors": ["model.layers.1.weight", "mtp.a.weight"], + "model-00003-of-00003.safetensors": ["mtp.b.weight"], + } + for filename, names in shards.items(): + save_file({name: torch.zeros(1) for name in names}, tmp_path / filename) + return tmp_path + + +def test_filter_by_weight_name_drops_shards_without_wanted_weights(mtp_checkpoint): + files = sorted(str(f) for f in mtp_checkpoint.glob("*.safetensors")) + + assert filter_safetensors_files_by_weight_name(files, _skip_non_mtp) == files[1:] + # A rule that rejects everything must not leave the loader without files. + assert filter_safetensors_files_by_weight_name(files, lambda _: True) == files + + +def test_loader_does_not_read_shards_the_model_skips(mtp_checkpoint, monkeypatch): + """Whole-file loaders such as InstantTensor only see the kept shards.""" + opened: list[list[str]] = [] + + def fake_iterator(hf_weights_files, *args, **kwargs): + opened.append(sorted(os.path.basename(f) for f in hf_weights_files)) + return iter(()) + + monkeypatch.setattr( + "vllm.model_executor.model_loader.default_loader." + "instanttensor_weights_iterator", + fake_iterator, + ) + + class Drafter(torch.nn.Module): + def is_unused_checkpoint_weight(self, name: str) -> bool: + return _skip_non_mtp(name) + + loader = DefaultModelLoader(LoadConfig(load_format="instanttensor")) + model_config = SimpleNamespace(model=str(mtp_checkpoint), revision=None) + list(loader.get_all_weights(model_config, Drafter())) + list(loader.get_all_weights(model_config, torch.nn.Module())) + + assert opened == [ + ["model-00002-of-00003.safetensors", "model-00003-of-00003.safetensors"], + sorted(f.name for f in mtp_checkpoint.glob("*.safetensors")), + ] diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index c3e09bccd425..9f1c25dd2e34 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/vllm/model_executor/model_loader/default_loader.py @@ -6,7 +6,7 @@ import glob import os import time -from collections.abc import Generator, Iterable +from collections.abc import Callable, Generator, Iterable from typing import TYPE_CHECKING, cast import torch @@ -28,6 +28,7 @@ filter_duplicate_safetensors_files, filter_files_not_needed_for_inference, filter_mm_encoder_only_safetensors_files, + filter_safetensors_files_by_weight_name, get_quant_config, instanttensor_weights_iterator, maybe_download_from_modelscope, @@ -75,6 +76,10 @@ class Source: allow_patterns_overrides: list[str] | None = None """If defined, weights will load exclusively using these patterns.""" + is_unused_weight: Callable[[str], bool] | None = None + """If defined, safetensors shards holding only weights whose checkpoint + name (before *prefix*) it accepts are not read.""" + counter_before_loading_weights: float = 0.0 counter_after_loading_weights: float = 0.0 @@ -283,6 +288,10 @@ def _get_weights_iterator( f"`{source.model_or_path}`; check language_model prefixes " f"{self._encoder_only_lm_prefixes}" ) + if source.is_unused_weight is not None and use_safetensors: + hf_weights_files = filter_safetensors_files_by_weight_name( + hf_weights_files, source.is_unused_weight + ) if self.load_config.load_format == "npcache": # Currently np_cache only support *.bin checkpoints assert use_safetensors is False @@ -359,6 +368,7 @@ def get_all_weights( prefix="", fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True), allow_patterns_overrides=getattr(model, "allow_patterns_overrides", None), + is_unused_weight=getattr(model, "is_unused_checkpoint_weight", None), ) yield from self._get_weights_iterator(primary_weights) diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index 12d557fa1ca2..c6a72c4c33c6 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -785,6 +785,42 @@ def _is_lm_only(weight_file: str) -> bool: return kept +def filter_safetensors_files_by_weight_name( + hf_weights_files: list[str], + is_unused_weight: Callable[[str], bool], +) -> list[str]: + """Drop safetensors shards in which `is_unused_weight` accepts every tensor. + + Loaders that read whole files (InstantTensor, fastsafetensors, multi-thread, + eager, prefetch) cannot skip single tensors, so dropping shards is their + only way to avoid reading weights the model does not load. + + Args: + hf_weights_files: Safetensors shard paths. + is_unused_weight: Returns True for checkpoint weight names the model does + not load. + + Returns: + The shards holding at least one wanted tensor, or `hf_weights_files` + unchanged when no shard does. + + """ + kept: list[str] = [] + for st_file in hf_weights_files: + with safe_open(st_file, framework="pt") as f: + if not all(is_unused_weight(name) for name in f.keys()): # noqa: SIM118 + kept.append(st_file) + if not kept: + return hf_weights_files + if len(kept) < len(hf_weights_files): + logger.info_once( + "Skipped %d/%d safetensors shard(s) holding no weights the model loads", + len(hf_weights_files) - len(kept), + len(hf_weights_files), + ) + return kept + + # explicitly use pure text format, with a newline at the end # this makes it impossible to see the animation in the progress bar # but will avoid messing up with ray or multiprocessing, which wraps diff --git a/vllm/model_executor/models/nemotron_h_mtp.py b/vllm/model_executor/models/nemotron_h_mtp.py index 519f85a63109..af93cc62ae25 100644 --- a/vllm/model_executor/models/nemotron_h_mtp.py +++ b/vllm/model_executor/models/nemotron_h_mtp.py @@ -426,6 +426,11 @@ def compute_logits( ) return self.logits_processor(self.lm_head, hidden_states) + def is_unused_checkpoint_weight(self, name: str) -> bool: + """Only MTP, embedding and LM head weights are loaded.""" + name = name.removeprefix("language_model.") + return not name.startswith(("mtp.", "lm_head.")) and "embeddings" not in name + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load MTP weights with proper name remapping.""" stacked_params_mapping = [ @@ -456,17 +461,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # MTP weights are nested in "language_model." # in Multimodal Nemotron-H checkpoints. name = name.removeprefix("language_model.") + if self.is_unused_checkpoint_weight(name): + continue is_lm_head_weight = name.startswith("lm_head.") if is_lm_head_weight: self.has_own_lm_head = True - # Only process MTP and LM head weights - - # skip all non-MTP and non-LM head weights - if ( - not name.startswith("mtp.") - and "embeddings" not in name - and not is_lm_head_weight - ): - continue # Skip rotary embeddings (computed, not loaded) if "rotary_emb.inv_freq" in name: continue diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py index c17d17014f8d..d7fb1d6e09ea 100644 --- a/vllm/model_executor/models/qwen3_5_mtp.py +++ b/vllm/model_executor/models/qwen3_5_mtp.py @@ -300,16 +300,20 @@ def compute_logits( ) -> torch.Tensor | None: return self.logits_processor(self.lm_head, hidden_states) + def is_unused_checkpoint_weight(self, name: str) -> bool: + return not name.startswith("mtp.") and not any( + key in name for key in ["embed_tokens", "lm_head"] + ) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def remap_weight_names(weights): for name, weight in weights: + if self.is_unused_checkpoint_weight(name): + continue if name.startswith("mtp."): name = name.replace("mtp.", "model.") - elif any(key in name for key in ["embed_tokens", "lm_head"]): - if "embed_tokens" in name: - name = name.replace("language_model.", "") - else: - continue + elif "embed_tokens" in name: + name = name.replace("language_model.", "") yield name, weight loader = AutoWeightsLoader(self) diff --git a/vllm/model_executor/models/qwen3_next_mtp.py b/vllm/model_executor/models/qwen3_next_mtp.py index e5e306bb861c..ef08878eaa7f 100644 --- a/vllm/model_executor/models/qwen3_next_mtp.py +++ b/vllm/model_executor/models/qwen3_next_mtp.py @@ -227,15 +227,18 @@ def compute_logits( ) -> torch.Tensor | None: return self.logits_processor(self.lm_head, hidden_states) - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - shared_weight_names = ["embed_tokens", "lm_head"] + def is_unused_checkpoint_weight(self, name: str) -> bool: + return not name.startswith("mtp.") and not any( + key in name for key in ["embed_tokens", "lm_head"] + ) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def remap_weight_names(weights): for name, weight in weights: + if self.is_unused_checkpoint_weight(name): + continue if name.startswith("mtp."): name = name.replace("mtp.", "model.") - elif not any(key in name for key in shared_weight_names): - continue yield name, weight loader = AutoWeightsLoader(self) diff --git a/vllm/models/deepseek_v32/amd/mtp.py b/vllm/models/deepseek_v32/amd/mtp.py index 70ddeae374c8..6230348ef44c 100644 --- a/vllm/models/deepseek_v32/amd/mtp.py +++ b/vllm/models/deepseek_v32/amd/mtp.py @@ -251,6 +251,12 @@ def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: name = name.replace(f"model.layers.{spec_layer}.", "model.") return name + def is_unused_checkpoint_weight(self, name: str) -> bool: + return ( + get_spec_layer_idx_from_weight_name(self.config, name) is None + and "embed_tokens" not in name + ) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: stacked_params_mapping = [ ("gate_up_proj", "gate_proj", 0), @@ -278,7 +284,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loaded_params: set[str] = set() _pending_wk_fp8: dict = {} for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: + if "rotary_emb.inv_freq" in name or self.is_unused_checkpoint_weight(name): continue spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is None: diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index bdc6b95c15e4..d4dee7b9d462 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -356,6 +356,12 @@ def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: name = name.replace(f"model.layers.{spec_layer}.", "model.") return name + def is_unused_checkpoint_weight(self, name: str) -> bool: + return ( + get_spec_layer_idx_from_weight_name(self.config, name) is None + and "embed_tokens" not in name + ) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: stacked_params_mapping = [ ("gate_up_proj", "gate_proj", 0), @@ -378,7 +384,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loaded_params: set[str] = set() _pending_wk_fp8: dict = {} for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: + if "rotary_emb.inv_freq" in name or self.is_unused_checkpoint_weight(name): continue spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is None: diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py index 64547af038aa..c0f3471dd5b6 100644 --- a/vllm/models/deepseek_v4/nvidia/mtp.py +++ b/vllm/models/deepseek_v4/nvidia/mtp.py @@ -340,6 +340,26 @@ def compute_logits( ) -> torch.Tensor | None: return self.model.compute_logits(hidden_states, spec_step_idx) + def _to_spec_layer_name(self, name: str) -> str: + """Remap V4 `mtp.{i}.*` names to `model.layers.{num_hidden_layers + i}.*` + so that get_spec_layer_idx_from_weight_name can identify them.""" + mtp_layer_idx = 0 + for subname in name.split("."): + try: + # we use the first encountered integer + mtp_layer_idx = int(subname) + break + except ValueError: + continue + return name.replace( + f"mtp.{mtp_layer_idx}.", + f"model.layers.{self.config.num_hidden_layers + mtp_layer_idx}.", + ) + + def is_unused_checkpoint_weight(self, name: str) -> bool: + name = self._to_spec_layer_name(name) + return get_spec_layer_idx_from_weight_name(self.config, name) is None + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Weight name remapping for checkpoint compatibility. # Maps checkpoint weight paths to model parameter paths. @@ -356,16 +376,6 @@ def _remap_weight_name(name: str) -> str: name = name.replace(old_pattern, new_pattern) return name - def _find_mtp_layer_idx(name: str) -> int: - subnames = name.split(".") - for subname in subnames: - try: - # we return the first encountered integer - return int(subname) - except ValueError: - continue - return 0 - stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("gate_up_proj", "w1", 0), @@ -409,15 +419,7 @@ def _find_mtp_layer_idx(name: str) -> int: ) for name, loaded_weight in weights: - mtp_layer_idx = _find_mtp_layer_idx(name) - # V4 checkpoints store MTP weights as `mtp.{i}.*`; remap to - # `model.layers.{num_hidden_layers + i}.*` so that - # get_spec_layer_idx_from_weight_name can identify them. - name = name.replace( - f"mtp.{mtp_layer_idx}.", - f"model.layers.{self.config.num_hidden_layers + mtp_layer_idx}.", - ) - + name = self._to_spec_layer_name(name) spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is None: continue From 08567505b09324f891797f4a13c1bd73fc75a827 Mon Sep 17 00:00:00 2001 From: Rui <57304024+ai-yang@users.noreply.github.com> Date: Thu, 8 Oct 2026 03:36:57 +0800 Subject: [PATCH 035/624] [Bugfix][KV Offload] Restore KVCR construction through the secondary-tier factory (#58088) --- tests/v1/kv_offload/tiering/test_kvcr_tier.py | 53 +++++++++++++------ vllm/v1/kv_offload/tiering/kvcr/manager.py | 6 ++- 2 files changed, 41 insertions(+), 18 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_kvcr_tier.py b/tests/v1/kv_offload/tiering/test_kvcr_tier.py index fc4421eb06b2..0907293fd883 100644 --- a/tests/v1/kv_offload/tiering/test_kvcr_tier.py +++ b/tests/v1/kv_offload/tiering/test_kvcr_tier.py @@ -36,7 +36,9 @@ ReqContext, make_offload_key, ) +from vllm.v1.kv_offload.tiering.backpressure import EMABackpressureDetector from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob +from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory from vllm.v1.kv_offload.tiering.kvcr import manager as kvcr_manager from vllm.v1.kv_offload.tiering.kvcr.manager import KVCRSecondaryTierManager @@ -180,6 +182,7 @@ def _make_tier( g3: dict[str, object] | None = None, control_ports: list[int] | None = None, data_parallel_rank_local: int | None = None, + backpressure: dict[str, Any] | None = None, ) -> KVCRSecondaryTierManager: def make_control(_bind_host, bind_port, advertise_host): return _StubControlChannel(f"tcp://{advertise_host}:{int(bind_port)}") @@ -196,7 +199,25 @@ def make_kvcr(config, bindings, backend_configs, guard_config): monkeypatch.setattr(kvcr_manager, "KVCR", make_kvcr) monkeypatch.setattr(kvcr_manager, "ZmqPeerControlChannel", make_control) - return KVCRSecondaryTierManager( + tier_config = { + "type": "kvcr", + "router_capabilities": ["router_hint"], + "control_host": "127.0.0.1", + "control_ports": control_ports if control_ports is not None else [7777], + "control_advertise_host": "127.0.0.1", + "enable_telemetry": enable_telemetry, + "secondary_g2_slots": secondary_g2_slots, + "kvcr_service_socket_path": kvcr_service_socket_path, + "compatibility_digest": compatibility_digest, + "policy": policy, + "g3": g3, + "local_dram_backend": "UCX", + "remote_fw_dram_backend": "UCX", + } + if backpressure is not None: + tier_config["backpressure"] = backpressure + tier = SecondaryTierFactory.create_secondary_tier( + tier_config=tier_config, offloading_spec=SimpleNamespace( config=SimpleNamespace( parallel=SimpleNamespace( @@ -209,20 +230,9 @@ def make_kvcr(config, bindings, backend_configs, guard_config): ), ), primary_kv_view=memoryview(np.zeros((4, 16), dtype=np.int8)), - tier_type="kvcr", - router_capabilities=["router_hint"], - control_host="127.0.0.1", - control_ports=control_ports if control_ports is not None else [7777], - control_advertise_host="127.0.0.1", - enable_telemetry=enable_telemetry, - secondary_g2_slots=secondary_g2_slots, - kvcr_service_socket_path=kvcr_service_socket_path, - compatibility_digest=compatibility_digest, - policy=policy, - g3=g3, - local_dram_backend="UCX", - remote_fw_dram_backend="UCX", ) + assert isinstance(tier, KVCRSecondaryTierManager) + return tier def test_kvcr_tier_configures_service_for_local_dp_rank(monkeypatch): @@ -271,10 +281,19 @@ def test_kvcr_tier_converts_g3_paths(monkeypatch, tmp_path): ) -def test_kvcr_tier_adapts_request_and_load(monkeypatch): - """Check hint forwarding, key conversion, load descriptors, and cleanup.""" +@pytest.mark.parametrize( + "backpressure", + [None, {"backpressure_cls": "EMABackpressureDetector"}], + ids=["default", "backpressure"], +) +def test_kvcr_tier_adapts_request_and_load(monkeypatch, backpressure): + """Factory-created tiers preserve configuration and complete hinted loads.""" kvcr = RecordingKVCR() - tier = _make_tier(monkeypatch, kvcr) + tier = _make_tier(monkeypatch, kvcr, backpressure=backpressure) + if backpressure is None: + assert tier.bp_detector is None + else: + assert isinstance(tier.bp_detector, EMABackpressureDetector) kv_hint = KvHintsEnvelope(protocol_version="0.1", message_id="msg", actions=[]) ctx = ReqContext( req_id="req", diff --git a/vllm/v1/kv_offload/tiering/kvcr/manager.py b/vllm/v1/kv_offload/tiering/kvcr/manager.py index 4b6f81b3b71e..4804f193e3e8 100644 --- a/vllm/v1/kv_offload/tiering/kvcr/manager.py +++ b/vllm/v1/kv_offload/tiering/kvcr/manager.py @@ -82,6 +82,7 @@ if TYPE_CHECKING: from vllm.v1.kv_offload.base import OffloadingSpec + from vllm.v1.kv_offload.tiering.backpressure import BackpressureDetector _REQUIRED_ROUTER_CAPABILITIES = ROUTER_HINT_CAPABILITIES @@ -295,8 +296,11 @@ def __init__( g3: dict[str, Any] | None = None, local_dram_backend: str = "UCX", remote_fw_dram_backend: str = "UCX", + backpressure_detector: "BackpressureDetector | None" = None, ) -> None: - super().__init__(offloading_spec, primary_kv_view, tier_type) + super().__init__( + offloading_spec, primary_kv_view, tier_type, backpressure_detector + ) selected_policy = _resolve_policy(policy) if (kvcr_service_socket_path is None) != (compatibility_digest is None): raise ValueError( From a982c81e3ed337f76b116338192c83f776a3225c Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Wed, 7 Oct 2026 16:37:13 -0400 Subject: [PATCH 036/624] [Bugfix] Limit GlmMoeDsaForCausalLM fp8 KV cache default to SM100 (#60281) Signed-off-by: mgoin Co-authored-by: Claude Opus 5.5 --- vllm/model_executor/models/config.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 97b9bfa0d38b..1a1bf9206166 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -46,14 +46,22 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: from vllm.platforms import current_platform cache_config = vllm_config.cache_config - if cache_config.cache_dtype == "auto": - if current_platform.is_xpu(): - cache_config.cache_dtype = "bfloat16" - elif current_platform.is_cuda_alike(): - capability = current_platform.get_device_capability() - cache_config.cache_dtype = ( - "fp8_e4m3" if capability and capability.major >= 10 else "bfloat16" - ) + if ( + cache_config.cache_dtype == "auto" + and current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + ): + cache_config.cache_dtype = "fp8_e4m3" + logger.info_once("Using fp8 kv-cache for GlmMoeDsaForCausalLM on SM10x") + # MTP shares the target's KV cache layout; a separate draft model + # must not inherit a default it was never validated with. + spec_config = vllm_config.speculative_config + if ( + spec_config is not None + and spec_config.method != "mtp" + and spec_config.kv_cache_dtype is None + ): + spec_config.kv_cache_dtype = "auto" # For Glm-Moe-DSA, qrep + a2a is better than the default all-gather + ag-rs # in most cases. From b47ec440bc0f2342848d6f1d99942cc3e909a66a Mon Sep 17 00:00:00 2001 From: Thang Nguyen <69278249+Thangnguyenvn98@users.noreply.github.com> Date: Wed, 7 Oct 2026 17:15:39 -0400 Subject: [PATCH 037/624] [CI] Stamp provenance labels on the torch-nightly image so Initialized Snapshot E2E can run (#59487) Signed-off-by: Thang Nguyen Co-authored-by: Kimi --- .buildkite/image_build/image_build_torch_nightly.sh | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.buildkite/image_build/image_build_torch_nightly.sh b/.buildkite/image_build/image_build_torch_nightly.sh index 56c7116e14ac..4d978b5299c8 100755 --- a/.buildkite/image_build/image_build_torch_nightly.sh +++ b/.buildkite/image_build/image_build_torch_nightly.sh @@ -61,6 +61,8 @@ NIGHTLY_UBUNTU_VERSION="24.04" NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu${NIGHTLY_UBUNTU_VERSION}" echo "--- :docker: Building torch nightly image (CUDA ${NIGHTLY_CUDA_VERSION})" +# Provenance labels match the bake `_labels` target in docker/docker-bake.hcl; +# the initialized-snapshot E2E refuses images whose commit labels differ. docker buildx build --file docker/Dockerfile \ --build-arg max_jobs=16 \ --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ @@ -70,6 +72,10 @@ docker buildx build --file docker/Dockerfile \ --build-arg BUILD_BASE_IMAGE="${NIGHTLY_BUILD_BASE_IMAGE}" \ --build-arg FINAL_BASE_IMAGE="${NIGHTLY_FINAL_BASE_IMAGE}" \ --build-arg torch_cuda_arch_list="8.0 8.9 9.0 10.0 12.0" \ + --label "org.opencontainers.image.revision=${BUILDKITE_COMMIT}" \ + --label "ai.vllm.build.commit=${BUILDKITE_COMMIT}" \ + --label "ai.vllm.build.pipeline=${BUILDKITE_PIPELINE_SLUG:-local}" \ + --label "ai.vllm.build.url=${BUILDKITE_BUILD_URL:-}" \ --tag "$IMAGE_TAG" \ --push \ --target test \ From 947dd62c0700851a79b0a8f32cb823e25a003b31 Mon Sep 17 00:00:00 2001 From: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Date: Wed, 7 Oct 2026 23:18:19 +0200 Subject: [PATCH 038/624] [CI] Fix failing test_deepseek_v4_rocm_dspark* tests (#60462) Signed-off-by: simondanielsson --- .buildkite/test_areas/models_basic.yaml | 1 + tests/models/test_deepseek_v4_dspark_rocm.py | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 04efae00fb17..a25c9a78a86c 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -81,6 +81,7 @@ steps: depends_on: - image-build-amd source_file_dependencies: + - vllm/models/deepseek_v4/ - tests/models/test_deepseek_v4_vl_rocm.py - tests/models/test_hyv4_rocm.py - tests/models/transformers/test_layer_registry_e2e.py diff --git a/tests/models/test_deepseek_v4_dspark_rocm.py b/tests/models/test_deepseek_v4_dspark_rocm.py index efa97e10d7f9..a5f8426e4ea5 100644 --- a/tests/models/test_deepseek_v4_dspark_rocm.py +++ b/tests/models/test_deepseek_v4_dspark_rocm.py @@ -15,7 +15,10 @@ def _make_uninitialized_model(confidence_head): object.__setattr__( model, "model", - SimpleNamespace(confidence_head=confidence_head), + SimpleNamespace( + confidence_head=confidence_head, + layers=[SimpleNamespace(ffn=SimpleNamespace(use_mega_moe=False))], + ), ) return model From b803edbc9fa5adbbf31136bf128d88533f7d39c2 Mon Sep 17 00:00:00 2001 From: Luciano Martins <22145370+lucianommartins@users.noreply.github.com> Date: Wed, 7 Oct 2026 18:56:40 -0300 Subject: [PATCH 039/624] [Model] Streamline EmbeddingGemma 2 config resolution and harmonize Triton prefill attention (#60289) Signed-off-by: Luciano Martins Co-authored-by: Luciano Martins Co-authored-by: Matthew Bonanni --- .../pooling/test_embedding_gemma2.py | 125 +++++++----------- vllm/model_executor/models/config.py | 88 +----------- vllm/transformers_utils/config.py | 2 +- 3 files changed, 54 insertions(+), 161 deletions(-) diff --git a/tests/models/multimodal/pooling/test_embedding_gemma2.py b/tests/models/multimodal/pooling/test_embedding_gemma2.py index 701d0914ddbc..0dcb72406da5 100644 --- a/tests/models/multimodal/pooling/test_embedding_gemma2.py +++ b/tests/models/multimodal/pooling/test_embedding_gemma2.py @@ -573,9 +573,10 @@ def test_video_media_io_embedding_gemma2(): assert meta_conn.get("do_sample_frames") is False -def test_config_capping_and_explicit_limits(caplog): - """Verify EmbeddingGemma2ModelConfig capping and explicit user overrides.""" - import logging +def test_config_capping_and_explicit_limits(): + """Verify EmbeddingGemma2ModelConfig backend selection, default capping, + and preservation of explicit limits. + """ from types import SimpleNamespace from unittest.mock import patch @@ -585,19 +586,13 @@ def test_config_capping_and_explicit_limits(caplog): ) from vllm.v1.attention.backends.registry import AttentionBackendEnum - # (a) Defaults: uncapped max_model_len=262144 -> capped to 8192, - # and derived scheduler limits (>= 262144) lowered to 8192. + # (a) When attention backend is None, default to TRITON_ATTN, + # and default capping: max_model_len > 8192 without explicit limit capped to 8192 cfg_default = SimpleNamespace( attention_config=SimpleNamespace(backend=None), model_config=SimpleNamespace( - model="dummy", max_model_len=262144, original_max_model_len=None - ), - scheduler_config=SimpleNamespace( - max_num_batched_tokens=262144, - max_num_encoder_input_tokens=262144, - encoder_cache_size=262144, - max_num_seqs=256, - verify_max_model_len=lambda x: None, + max_model_len=262144, + original_max_model_len=None, ), ) with patch.object(Gemma4Config, "verify_and_update_config"): @@ -605,83 +600,61 @@ def test_config_capping_and_explicit_limits(caplog): assert cfg_default.attention_config.backend == AttentionBackendEnum.TRITON_ATTN assert cfg_default.model_config.max_model_len == 8192 - assert cfg_default.scheduler_config.max_num_batched_tokens == 8192 - assert cfg_default.scheduler_config.max_num_encoder_input_tokens == 8192 - assert cfg_default.scheduler_config.encoder_cache_size == 8192 - - # (b) Reachable explicit limits: explicit max_model_len=8192 and explicit - # max_num_batched_tokens=32768 stays 32768. - # Note: omitting max_model_len while passing max_num_batched_tokens < 262144 - # is rejected earlier by SchedulerConfig validation (32768 < 262144) before - # this hook runs. - cfg_explicit_batched = SimpleNamespace( + + # (b) Default capping with original_max_model_len == -1 + cfg_default_minus1 = SimpleNamespace( attention_config=SimpleNamespace(backend=None), model_config=SimpleNamespace( - model="dummy", max_model_len=8192, original_max_model_len=8192 - ), - scheduler_config=SimpleNamespace( - max_num_batched_tokens=32768, - max_num_encoder_input_tokens=32768, - encoder_cache_size=32768, - max_num_seqs=256, - verify_max_model_len=lambda x: None, + max_model_len=131072, + original_max_model_len=-1, ), ) with patch.object(Gemma4Config, "verify_and_update_config"): - EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_batched) + EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_default_minus1) - assert cfg_explicit_batched.model_config.max_model_len == 8192 - assert cfg_explicit_batched.scheduler_config.max_num_batched_tokens == 32768 - assert cfg_explicit_batched.scheduler_config.max_num_encoder_input_tokens == 32768 - assert cfg_explicit_batched.scheduler_config.encoder_cache_size == 32768 + assert cfg_default_minus1.model_config.max_model_len == 8192 - # (c) Explicit value > uncapped_len (e.g. 300000) with uncapped model len is lowered - # to 8192 and logs a warning because it exceeds uncapped_len. - cfg_explicit_large = SimpleNamespace( + # (c) Explicit limit 4096 preserved + cfg_explicit_4k = SimpleNamespace( attention_config=SimpleNamespace(backend=None), model_config=SimpleNamespace( - model="dummy", max_model_len=262144, original_max_model_len=None + max_model_len=4096, + original_max_model_len=4096, ), - scheduler_config=SimpleNamespace( - max_num_batched_tokens=300000, - max_num_encoder_input_tokens=300000, - encoder_cache_size=300000, - max_num_seqs=256, - verify_max_model_len=lambda x: None, - ), - ) - with ( - caplog.at_level(logging.WARNING), - patch.object(Gemma4Config, "verify_and_update_config"), - ): - EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_large) - - assert cfg_explicit_large.model_config.max_model_len == 8192 - assert cfg_explicit_large.scheduler_config.max_num_batched_tokens == 8192 - assert cfg_explicit_large.scheduler_config.max_num_encoder_input_tokens == 8192 - assert cfg_explicit_large.scheduler_config.encoder_cache_size == 8192 - assert any( - "lowering explicitly set scheduler_config.max_num_batched_tokens" - " from 300000 to 8192" in rec.message - for rec in caplog.records ) + with patch.object(Gemma4Config, "verify_and_update_config"): + EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_4k) - # (d) Explicit max_model_len=16384 is respected and not capped to 8192 - cfg_explicit_model_len = SimpleNamespace( - attention_config=SimpleNamespace(backend=None), + assert cfg_explicit_4k.model_config.max_model_len == 4096 + + # (d) Explicit limit 16384 preserved and explicit backend preserved + cfg_explicit_16k = SimpleNamespace( + attention_config=SimpleNamespace(backend=AttentionBackendEnum.FLASH_ATTN), model_config=SimpleNamespace( - model="dummy", max_model_len=16384, original_max_model_len=16384 - ), - scheduler_config=SimpleNamespace( - max_num_batched_tokens=16384, - max_num_encoder_input_tokens=16384, - encoder_cache_size=16384, - max_num_seqs=256, - verify_max_model_len=lambda x: None, + max_model_len=16384, + original_max_model_len=16384, ), ) with patch.object(Gemma4Config, "verify_and_update_config"): - EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_model_len) + EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_16k) + + assert cfg_explicit_16k.attention_config.backend == AttentionBackendEnum.FLASH_ATTN + assert cfg_explicit_16k.model_config.max_model_len == 16384 + + +def test_sentence_transformer_tokenizer_config_max_seq_length(tmp_path): + """Verify get_sentence_transformer_tokenizer_config accepts config + with only max_seq_length. + """ + import json + + from vllm.transformers_utils.config import ( + get_sentence_transformer_tokenizer_config, + ) - assert cfg_explicit_model_len.model_config.max_model_len == 16384 - assert cfg_explicit_model_len.scheduler_config.max_num_batched_tokens == 16384 + # 1. Synthetic model directory with only max_seq_length in sentence_bert_config.json + cfg_file = tmp_path / "sentence_bert_config.json" + cfg_file.write_text(json.dumps({"max_seq_length": 8192})) + loaded = get_sentence_transformer_tokenizer_config(str(tmp_path)) + assert loaded is not None + assert loaded["max_seq_length"] == 8192 diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 1a1bf9206166..7f6a3aef40a7 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -309,90 +309,10 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: ) model_config = vllm_config.model_config - orig_len = getattr(model_config, "original_max_model_len", None) - if orig_len in (None, -1): - default_len = 8192 - # Check sentence_bert_config.json if available - try: - import json - from pathlib import Path - - cfg_path = Path(model_config.model) / "sentence_bert_config.json" - if cfg_path.exists(): - data = json.loads(cfg_path.read_text()) - if "max_seq_length" in data and isinstance( - data["max_seq_length"], int - ): - default_len = data["max_seq_length"] - except Exception: - pass - - if model_config.max_model_len > default_len: - uncapped_len = model_config.max_model_len - logger.info( - "EmbeddingGemma2: max_model_len not explicitly set; capping from " - "%d to %d.", - model_config.max_model_len, - default_len, - ) - model_config.max_model_len = default_len - scheduler_config = vllm_config.scheduler_config - if scheduler_config is not None: - new_batched_tokens = max( - model_config.max_model_len, scheduler_config.max_num_seqs - ) - if scheduler_config.max_num_batched_tokens >= uncapped_len: - if scheduler_config.max_num_batched_tokens > uncapped_len: - logger.warning( - "EmbeddingGemma2: lowering explicitly set " - "scheduler_config.max_num_batched_tokens from %d to %d " - "to match capped max_model_len.", - scheduler_config.max_num_batched_tokens, - new_batched_tokens, - ) - else: - logger.info( - "EmbeddingGemma2: lowering " - "scheduler_config.max_num_batched_tokens from %d to %d " - "to match capped max_model_len.", - scheduler_config.max_num_batched_tokens, - new_batched_tokens, - ) - scheduler_config.max_num_batched_tokens = new_batched_tokens - if hasattr(vllm_config, "_set_compile_ranges"): - vllm_config._set_compile_ranges() - if scheduler_config.max_num_encoder_input_tokens >= uncapped_len: - if scheduler_config.max_num_encoder_input_tokens > uncapped_len: - logger.warning( - "EmbeddingGemma2: lowering explicitly set " - "scheduler_config.max_num_encoder_input_tokens " - "from %d to %d to match capped max_model_len.", - scheduler_config.max_num_encoder_input_tokens, - new_batched_tokens, - ) - scheduler_config.max_num_encoder_input_tokens = ( - new_batched_tokens - ) - if scheduler_config.encoder_cache_size >= uncapped_len: - if scheduler_config.encoder_cache_size > uncapped_len: - logger.warning( - "EmbeddingGemma2: lowering explicitly set " - "scheduler_config.encoder_cache_size from %d to %d " - "to match capped max_model_len.", - scheduler_config.encoder_cache_size, - new_batched_tokens, - ) - scheduler_config.encoder_cache_size = new_batched_tokens - if hasattr(scheduler_config, "verify_max_model_len"): - scheduler_config.verify_max_model_len( - model_config.max_model_len - ) - elif model_config.max_model_len > 32768: - logger.warning( - "EmbeddingGemma2: max_model_len=%d is large; " - "consider --max-model-len 8192 to bound encoder memory.", - model_config.max_model_len, - ) + if model_config.max_model_len > 8192 and getattr( + model_config, "original_max_model_len", None + ) in (None, -1): + model_config.max_model_len = 8192 class DiffusionGemmaModelForBlockDiffusionConfig(VerifyAndUpdateConfig): diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 49312047144f..6a339297eec5 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -1139,7 +1139,7 @@ def get_sentence_transformer_tokenizer_config( logger.info("Found sentence-transformers tokenize configuration.") - if all(k in encoder_dict for k in ("max_seq_length", "do_lower_case")): + if "max_seq_length" in encoder_dict: return encoder_dict return None From f38c6778ce4cf62b1ca7a8e091892f5bfdf79f98 Mon Sep 17 00:00:00 2001 From: Xuanteng Huang <44627253+xuantengh@users.noreply.github.com> Date: Thu, 8 Oct 2026 06:01:09 +0800 Subject: [PATCH 040/624] [quantization] Add per-token NVFP4 quantization MoE support for ReLU2 (#56740) Signed-off-by: Xuanteng Huang Co-authored-by: Claude Fable 5.1 Co-authored-by: OpenAI Codex --- ...ron-3-Super-120B-A12B-NVFP4-per-token.yaml | 14 ++ .../gsm8k/configs/models-blackwell-ep.txt | 1 + tests/kernels/moe/test_trtllm_nvfp4_moe.py | 130 +++++++++++++++--- .../model_loader/test_weight_cache.py | 46 +++++++ tests/quantization/test_blackwell_moe.py | 13 ++ tests/quantization/test_modelopt.py | 77 ++++++++++- tests/quantization/test_online.py | 32 ++++- .../test_quantization_config_args.py | 51 +++++-- vllm/config/quantization.py | 2 + .../layers/fused_moe/oracle/nvfp4.py | 2 + .../layers/quantization/modelopt.py | 57 ++++++-- .../layers/quantization/online/nvfp4.py | 17 ++- .../quantization/utils/flashinfer_fp4_moe.py | 7 +- .../quantization/utils/flashinfer_utils.py | 14 ++ .../model_loader/weight_cache/protocol.py | 19 ++- .../model_loader/weight_utils.py | 19 ++- vllm/models/deepseek_v4/quant_config.py | 5 + vllm/models/deepseek_v41/quant_config.py | 5 + 18 files changed, 450 insertions(+), 61 deletions(-) create mode 100644 tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4-per-token.yaml diff --git a/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4-per-token.yaml b/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4-per-token.yaml new file mode 100644 index 000000000000..62587ae10a7d --- /dev/null +++ b/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4-per-token.yaml @@ -0,0 +1,14 @@ +model_name: "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4" +accuracy_threshold: 0.93 +num_questions: 1319 +num_fewshot: 5 +startup_max_wait_seconds: 1200 +server_args: >- + --enforce-eager + --max-model-len 4096 + --tensor-parallel-size 2 + --enable-expert-parallel + --mamba-backend flashinfer + --moe-backend flashinfer_trtllm + --quantization-config.moe.activation nvfp4_per_token + --speculative-config '{"method":"mtp","num_speculative_tokens":5}' diff --git a/tests/evals/gsm8k/configs/models-blackwell-ep.txt b/tests/evals/gsm8k/configs/models-blackwell-ep.txt index c2a19051986f..34090d720512 100644 --- a/tests/evals/gsm8k/configs/models-blackwell-ep.txt +++ b/tests/evals/gsm8k/configs/models-blackwell-ep.txt @@ -1,3 +1,4 @@ Qwen3-Next-80B-A3B-NVFP4-EP2.yaml Qwen3-Next-FP8-EP2.yaml Nemotron-3-Super-120B-A12B-NVFP4.yaml +Nemotron-3-Super-120B-A12B-NVFP4-per-token.yaml diff --git a/tests/kernels/moe/test_trtllm_nvfp4_moe.py b/tests/kernels/moe/test_trtllm_nvfp4_moe.py index 292e4549bb7d..4ee095505196 100644 --- a/tests/kernels/moe/test_trtllm_nvfp4_moe.py +++ b/tests/kernels/moe/test_trtllm_nvfp4_moe.py @@ -16,8 +16,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk from tests.kernels.moe.utils import check_deferred_moe_finalize, make_test_quant_config from tests.kernels.quantization.nvfp4_utils import ( - FLOAT4_E2M1_MAX, - FLOAT8_E4M3_MAX, + convert_swizzled_to_linear, dequantize_nvfp4_to_dtype, ) from tests.kernels.utils import torch_moe @@ -39,6 +38,13 @@ from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( TrtLlmNvFp4ExpertsModular, ) +from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( + prepare_static_weights_for_trtllm_fp4_moe, + reorder_w1w3_to_w3w1, +) +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + NVFP4_PER_TOKEN_BASE_GLOBAL_SCALE, +) from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe from vllm.utils.math_utils import next_power_of_2 @@ -96,10 +102,14 @@ def __init__(self, *, compile_native: bool = True) -> None: def _make_trtllm_fp4_moe_kernel( - moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + per_token_activation: bool = False, ) -> mk.FusedMoEKernel: experts = TrtLlmNvFp4ExpertsModular( - moe_config=moe_config, quant_config=quant_config + moe_config=moe_config, + quant_config=quant_config, + per_token_activation=per_token_activation, ) # Mimic the production weight-loader path so per-expert tensors that # are normally precomputed in process_weights_after_loading (g1_scale_c @@ -126,16 +136,37 @@ def _make_trtllm_fp4_moe_kernel( ACTIVATION_CASES = [ - pytest.param(MoEActivation.SILU, MoEActivation.SILU, None, id="silu"), - pytest.param(MoEActivation.SILU, SILU_WITH_CLAMP, _SWIGLU_LIMIT, id="silu_clamp"), - pytest.param(MoEActivation.SITU, SITU, None, id="situ"), + pytest.param(MoEActivation.SILU, MoEActivation.SILU, None, False, id="silu"), + pytest.param( + MoEActivation.SILU, + MoEActivation.SILU, + None, + True, + id="silu_per_token", + ), + pytest.param( + MoEActivation.SILU, + SILU_WITH_CLAMP, + _SWIGLU_LIMIT, + False, + id="silu_clamp", + ), + pytest.param(MoEActivation.SITU, SITU, None, False, id="situ"), pytest.param( MoEActivation.RELU2_NO_MUL, MoEActivation.RELU2_NO_MUL, None, + False, id="relu2_no_mul", ), - pytest.param(MoEActivation.GELU, MoEActivation.GELU, None, id="gelu"), + pytest.param( + MoEActivation.RELU2_NO_MUL, + MoEActivation.RELU2_NO_MUL, + None, + True, + id="relu2_no_mul_per_token", + ), + pytest.param(MoEActivation.GELU, MoEActivation.GELU, None, False, id="gelu"), ] @@ -143,7 +174,10 @@ def _make_trtllm_fp4_moe_kernel( @pytest.mark.parametrize("e", [128]) @pytest.mark.parametrize("topk", [8]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("activation,torch_activation,swiglu_limit", ACTIVATION_CASES) +@pytest.mark.parametrize( + "activation,torch_activation,swiglu_limit,per_token_activation", + ACTIVATION_CASES, +) @torch.inference_mode() def test_trtllm_fp4_moe_no_graph( m: int, @@ -155,6 +189,7 @@ def test_trtllm_fp4_moe_no_graph( activation: MoEActivation, torch_activation: MoEActivation | type[SiluAndMulWithClamp] | type[SituAndMul], swiglu_limit: float | None, + per_token_activation: bool, workspace_init, ): # FlashInfer's trtllm_batched_gemm_runner has no precompiled tile @@ -202,6 +237,42 @@ def test_trtllm_fp4_moe_no_graph( # clamp/output-scale coupling in the FlashInfer kernel wrapper. quant_config.g1_alphas.fill_(_LARGE_OUTPUT1_SCALE) + w1_q_linear = w1_q + w2_q_linear = w2_q + w1_scale_reference = quant_config.w1_scale + w2_scale_reference = quant_config.w2_scale + assert w1_scale_reference is not None + assert w2_scale_reference is not None + w1_scale_linear = torch.stack( + [ + convert_swizzled_to_linear(scale, w1_q.size(1), k, quant_blocksize) + for scale in w1_scale_reference + ] + ) + w2_scale_linear = torch.stack( + [ + convert_swizzled_to_linear(scale, w2_q.size(1), n, quant_blocksize) + for scale in w2_scale_reference + ] + ) + w1_q_for_kernel = w1_q + if is_gated_act: + w1_q_for_kernel, w1_scale_linear = reorder_w1w3_to_w3w1( + w1_q.clone(), w1_scale_linear + ) + w1_q, w1_scale, w2_q, w2_scale = prepare_static_weights_for_trtllm_fp4_moe( + w1_q_for_kernel, + w2_q, + w1_scale_linear, + w2_scale_linear, + hidden_size=k, + intermediate_size=n, + num_experts=e, + is_gated_activation=is_gated_act, + ) + quant_config._w1.scale = w1_scale + quant_config._w2.scale = w2_scale + score = torch.randn((m, e), device="cuda", dtype=dtype) topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) @@ -224,7 +295,11 @@ def test_trtllm_fp4_moe_no_graph( ), ) - trtllm_experts = _make_trtllm_fp4_moe_kernel(moe_config, quant_config) + trtllm_experts = _make_trtllm_fp4_moe_kernel( + moe_config, + quant_config, + per_token_activation=per_token_activation, + ) if activation == MoEActivation.SITU: torch.testing.assert_close( trtllm_experts.fused_experts.g1_scale_c, quant_config.a2_gscale @@ -245,10 +320,20 @@ def test_trtllm_fp4_moe_no_graph( # Reference: round-trip activations and weights through FP4 # quant/dequant so the comparison isolates kernel/activation behavior # from quantization error. - a_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / a.abs().max()).to( - torch.float32 - ) - a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a, a_global_scale) + if per_token_activation: + from flashinfer import SfLayout, nvfp4_quantize + + a_fp4, a_scale_interleaved, per_token_scale = nvfp4_quantize( + a, + NVFP4_PER_TOKEN_BASE_GLOBAL_SCALE, + sfLayout=SfLayout.layout_linear, + per_token_activation=True, + ) + a_global_scale = 1.0 / per_token_scale.unsqueeze(1) + else: + assert quant_config.a1_gscale is not None + a_global_scale = quant_config.a1_gscale[0] + a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a, a_global_scale) a_in_dtype = dequantize_nvfp4_to_dtype( a_fp4, a_scale_interleaved, @@ -256,6 +341,7 @@ def test_trtllm_fp4_moe_no_graph( dtype=a.dtype, device=a.device, block_size=quant_blocksize, + is_sf_linear_layout=per_token_activation, ) w1_d = torch.empty( @@ -264,16 +350,16 @@ def test_trtllm_fp4_moe_no_graph( w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype) for idx in range(e): w1_d[idx] = dequantize_nvfp4_to_dtype( - w1_q[idx], - quant_config.w1_scale[idx], + w1_q_linear[idx], + w1_scale_reference[idx], (1 / quant_config.g1_alphas[idx]), dtype=dtype, device=w1_q.device, block_size=quant_blocksize, ) w2_d[idx] = dequantize_nvfp4_to_dtype( - w2_q[idx], - quant_config.w2_scale[idx], + w2_q_linear[idx], + w2_scale_reference[idx], (1 / quant_config.g2_alphas[idx]), dtype=dtype, device=w2_q.device, @@ -284,6 +370,13 @@ def test_trtllm_fp4_moe_no_graph( a_in_dtype, w1_d, w2_d, score, topk, activation=torch_activation ) + if swiglu_limit is None: + cosine = torch.nn.functional.cosine_similarity( + torch_output.float().flatten(), + trtllm_output.float().flatten(), + dim=0, + ) + assert cosine > 0.99 torch.testing.assert_close(torch_output, trtllm_output, atol=2e-1, rtol=2e-1) @@ -416,5 +509,6 @@ def test_trtllm_fp4_moe_deferred_finalize( MoEActivation.GELU, MoEActivation.GELU, None, + False, None, ) diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py index 731c73262586..7960a237914b 100644 --- a/tests/model_executor/model_loader/test_weight_cache.py +++ b/tests/model_executor/model_loader/test_weight_cache.py @@ -356,6 +356,52 @@ def test_weight_cache_key_distinguishes_dp_ranks(): assert key.mismatched_fields(replace(key, pp_rank=0)) == ["pp_rank"] +@pytest.mark.parametrize("hf_quant_config", [None, {"quant_algo": "NVFP4"}]) +def test_weight_cache_key_distinguishes_nvfp4_activation_override( + tmp_path, hf_quant_config +): + """Cache keys isolate activation modes and reject ambiguous legacy hashes.""" + import json + from dataclasses import replace + from types import SimpleNamespace + + import torch + + from vllm.config.quantization import QuantizationConfigArgs + from vllm.model_executor.model_loader.weight_cache.protocol import WeightCacheKey + from vllm.utils.hashing import safe_hash + + model_config = SimpleNamespace( + model=str(tmp_path), + dtype=torch.bfloat16, + quantization="modelopt_fp4", + quantization_config=None, + revision=None, + hf_config=SimpleNamespace( + architectures=["NemotronHForCausalLM"], + quantization_config=hf_quant_config, + ), + ) + static_key = WeightCacheKey.from_model_config(model_config, tp_size=1, tp_rank=0) + model_config.quantization_config = QuantizationConfigArgs( + moe={"activation": "nvfp4_per_token"} + ) + per_token_key = WeightCacheKey.from_model_config(model_config, tp_size=1, tp_rank=0) + assert static_key.mismatched_fields(per_token_key) == ["quant_config_hash"] + + legacy_hash = ( + "" + if hf_quant_config is None + else safe_hash( + json.dumps(hf_quant_config, sort_keys=True).encode(), + usedforsecurity=False, + ).hexdigest() + ) + assert static_key.mismatched_fields( + replace(static_key, quant_config_hash=legacy_hash) + ) == ["quant_config_hash"] + + def test_ipc_loader_copy_mode_reports_no_external_weight_memory(): """Copy mode clones the weights into the engine, so nothing is external (vllm_config is never touched).""" diff --git a/tests/quantization/test_blackwell_moe.py b/tests/quantization/test_blackwell_moe.py index da70491bbc27..3a4e2b704b81 100644 --- a/tests/quantization/test_blackwell_moe.py +++ b/tests/quantization/test_blackwell_moe.py @@ -297,3 +297,16 @@ def test_nemotron_fp4_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch): hf_overrides=HF_OVERRIDE_TEXT, extra_args=["--moe-backend=flashinfer_trtllm"], ) + + +def test_nemotron_nvfp4_per_token_moe_flashinfer_trtllm( + monkeypatch: pytest.MonkeyPatch, +): + can_initialize( + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4", + hf_overrides=HF_OVERRIDE_TEXT, + extra_args=[ + "--moe-backend=flashinfer_trtllm", + "--quantization-config.moe.activation=nvfp4_per_token", + ], + ) diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 7b1131a00e39..33c798cba09b 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -6,6 +6,7 @@ """ import os +from contextlib import nullcontext from typing import Any, NoReturn from unittest.mock import MagicMock, Mock, patch @@ -17,7 +18,9 @@ load_model_without_vllm_runner, ) from vllm.config import VllmConfig, set_current_vllm_config +from vllm.config.load import LoadConfig from vllm.config.model import ModelConfig +from vllm.config.quantization import QuantizationConfigArgs from vllm.forward_context import set_forward_context from vllm.model_executor.kernels.linear import ( FlashInferCuteDslNvFp4W4A16LinearKernel, @@ -40,12 +43,14 @@ kMxfp8Dynamic, kMxfp8Static, kNvfp4Dynamic, + kNvfp4DynamicToken, kNvfp4Static, ) from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) +from vllm.model_executor.model_loader.weight_utils import get_quant_config from vllm.platforms import current_platform @@ -113,11 +118,13 @@ def _mixed_precision_config(quantized_layers: dict) -> ModelOptMixedPrecisionCon ) -def test_modelopt_nvfp4_quantizes_parallel_lm_head(): +@pytest.mark.parametrize("moe_activation", [None, "nvfp4_per_token"]) +def test_modelopt_nvfp4_quantizes_parallel_lm_head(moe_activation): config = ModelOptNvFp4Config( is_checkpoint_nvfp4_serialized=True, kv_cache_quant_algo=None, exclude_modules=[], + quantization_args=QuantizationConfigArgs(moe={"activation": moe_activation}), ) method = config.get_quant_method(_mock_lm_head(), prefix="lm_head") @@ -790,6 +797,74 @@ def test_modelopt_nvfp4_moe_dispatches_to_marlin_when_w4a16( assert kwargs["activation_key"] is kNvfp4Dynamic +@pytest.mark.parametrize("quantization", ["modelopt_fp4", "modelopt_mixed"]) +@pytest.mark.parametrize("per_token", [False, True]) +def test_modelopt_nvfp4_moe_activation_override(quantization, per_token): + """Runtime overrides reach NVFP4 configs without changing HF metadata.""" + hf_config = { + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION" + if quantization == "modelopt_mixed" + else "NVFP4", + "quantized_layers": {"experts": {"quant_algo": "NVFP4"}}, + } + args = ( + QuantizationConfigArgs(moe={"activation": "nvfp4_per_token"}) + if per_token + else None + ) + config = get_quant_config( + Mock( + quantization=quantization, + quantization_config=args, + hf_config=Mock(quantization_config=hf_config), + ), + LoadConfig(), + ) + if isinstance(config, ModelOptMixedPrecisionConfig): + config = config.nvfp4_config + assert isinstance(config, ModelOptNvFp4Config) + assert config.moe_activation_override == (kNvfp4DynamicToken if per_token else None) + assert "_online_quantization_args" not in hf_config + + +@pytest.mark.parametrize( + "quant_method,activation,backend,error", + [ + ("NVFP4", kNvfp4DynamicToken, "FLASHINFER_TRTLLM", None), + ("W4A16_NVFP4", kNvfp4DynamicToken, "FLASHINFER_TRTLLM", "W4A4 checkpoint"), + ( + "NVFP4", + kMxfp8Dynamic, + "FLASHINFER_TRTLLM", + "Unsupported.*activation override", + ), + ("NVFP4", kNvfp4DynamicToken, "MARLIN", "requires the FlashInfer TRTLLM"), + ], +) +def test_modelopt_nvfp4_moe_validates_activation_override( + quant_method, activation, backend, error +): + from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend + from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4FusedMoE + + config = ModelOptNvFp4Config( + quant_method=quant_method, + quantization_args=QuantizationConfigArgs(moe={"activation": activation}), + ) + expected = pytest.raises(ValueError, match=error) if error else nullcontext() + with ( + patch( + "vllm.model_executor.layers.quantization.modelopt.select_nvfp4_moe_backend", + return_value=(NvFp4MoeBackend[backend], Mock()), + ) as select_backend, + expected, + ): + method = ModelOptNvFp4FusedMoE(config, Mock()) + assert method.per_token_activation + assert select_backend.call_args.kwargs["activation_key"] == activation + + @pytest.mark.parametrize( "per_layer_algo, expected_weight, expected_activation", [ diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index 3f27757b9017..c6b9894009cf 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -103,6 +103,8 @@ amax_for_tp_weight_quant, kMxfp8Dynamic, kMxfp8Static, + kNvfp4DynamicToken, + kNvfp4Static, weight_amax, ) from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -211,15 +213,26 @@ def test_qianfan_online_fp8_keeps_vision_layers_unquantized( assert args.ignore == [] +@pytest.mark.parametrize( + ("is_act_and_mul", "expected_hidden_alignment"), + [(True, 256), (False, 512)], +) def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed( - monkeypatch, + monkeypatch, is_act_and_mul, expected_hidden_alignment ) -> None: - method = object.__new__(Nvfp4OnlineMoEMethod) - method.moe = SimpleNamespace(is_act_and_mul=True) - method.nvfp4_backend = object() - method.experts_cls = object - method.moe_quant_config = None - method.moe_kernel = None + monkeypatch.setattr(current_platform, "is_device_capability_family", lambda _: True) + select_backend = Mock(return_value=(object(), object)) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.online.nvfp4.select_nvfp4_moe_backend", + select_backend, + ) + moe_config = SimpleNamespace(is_act_and_mul=is_act_and_mul) + method = Nvfp4OnlineMoEMethod(moe=moe_config) + select_backend.assert_called_once_with( + config=moe_config, + weight_key=kNvfp4Static, + activation_key=kNvfp4DynamicToken, + ) layer = Mock() converted_weights = [ @@ -274,7 +287,12 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed( assert method.moe_kernel is kernel assert convert_weights.call_count == 3 + assert all( + call.kwargs["trtllm_hidden_alignment"] == expected_hidden_alignment + for call in convert_weights.call_args_list + ) make_kernel.assert_called_once() + assert make_kernel.call_args.kwargs["per_token_activation"] is True get_quant_config.assert_called_once() assert process_weights.call_count == 3 diff --git a/tests/quantization/test_quantization_config_args.py b/tests/quantization/test_quantization_config_args.py index 1d6ced0d48f4..7e2e21a3acbf 100644 --- a/tests/quantization/test_quantization_config_args.py +++ b/tests/quantization/test_quantization_config_args.py @@ -26,6 +26,8 @@ kInt4Static32, kInt8StaticChannelSym, kMxfp8Dynamic, + kNvfp4DynamicToken, + kNvfp4Static, ) # ---- QuantSpec ------------------------------------------------------------ @@ -68,13 +70,17 @@ def test_args_linear_string_resolves_via_quant_key_names(): assert args.moe is None -def test_args_moe_string_resolves_via_online_shorthand(): +@pytest.mark.parametrize( + "shorthand,weight", + [("fp8_per_block", kFp8Static128BlockSym), ("nvfp4_per_token", kNvfp4Static)], +) +def test_args_moe_string_resolves_via_online_shorthand(shorthand, weight): # An online-shorthand name pulls the matching slot from _ONLINE_SHORTHANDS # (so `linear: "fp8_per_block"` and `moe: "fp8_per_block"` produce the # same per-layer-kind spec the `--quantization fp8_per_block` shorthand # would). - args = quant_config_args(moe="fp8_per_block") - assert args.moe == QuantSpec(weight=kFp8Static128BlockSym) + args = quant_config_args(moe=shorthand) + assert args.moe == QuantSpec(weight=weight) def test_args_string_shorthand_missing_slot_raises(): @@ -84,9 +90,14 @@ def test_args_string_shorthand_missing_slot_raises(): quant_config_args(linear="int8_per_channel_weight_only") -def test_args_accepts_dict_form(): - args = quant_config_args(moe={"activation": "mxfp8"}) - assert args.moe == QuantSpec(weight=None, activation=kMxfp8Dynamic) +@pytest.mark.parametrize( + "activation,expected", + [("mxfp8", kMxfp8Dynamic), ("nvfp4_per_token", kNvfp4DynamicToken)], +) +def test_args_accepts_dict_form(activation, expected): + args = quant_config_args(moe={"activation": activation}) + assert args.moe == QuantSpec(weight=None, activation=expected) + assert args.linear is None def test_targets_reject_non_string_keys(): @@ -110,22 +121,34 @@ def test_resolve_colliding_shorthand_is_deferred(quantization: str): assert resolve_quantization_config(quantization, None) is None -def test_resolve_int8_shorthand_leaves_linear_unset(): - # int8_per_channel_weight_only is MoE-only; linear stays None so that - # OnlineQuantizationConfig leaves Linear layers in full precision. - args = resolve_quantization_config("int8_per_channel_weight_only", None) +@pytest.mark.parametrize( + "shorthand,weight", + [ + ("int8_per_channel_weight_only", kInt8StaticChannelSym), + ("nvfp4_per_token", kNvfp4Static), + ], +) +def test_resolve_moe_only_shorthand_leaves_linear_unset(shorthand, weight): + args = resolve_quantization_config(shorthand, None) assert args is not None assert args.linear is None - assert args.moe == QuantSpec(weight=kInt8StaticChannelSym) + assert args.moe == QuantSpec(weight=weight) -def test_resolve_quantization_config_only(): +@pytest.mark.parametrize("quantization", [None, "modelopt_fp4"]) +@pytest.mark.parametrize( + "activation,expected", + [("mxfp8", kMxfp8Dynamic), ("nvfp4_per_token", kNvfp4DynamicToken)], +) +def test_resolve_activation_override(quantization, activation, expected): # When only `quantization_config` is given (e.g. for an already-quantized # checkpoint that needs an activation override), it's returned as-is. - args = resolve_quantization_config(None, {"moe": {"activation": "mxfp8"}}) + args = resolve_quantization_config( + quantization, {"moe": {"activation": activation}} + ) assert args is not None assert args.linear is None - assert args.moe == QuantSpec(weight=None, activation=kMxfp8Dynamic) + assert args.moe == QuantSpec(weight=None, activation=expected) def test_resolve_merges_explicit_over_shorthand(): diff --git a/vllm/config/quantization.py b/vllm/config/quantization.py index d9500847268b..face2c3a7ba7 100644 --- a/vllm/config/quantization.py +++ b/vllm/config/quantization.py @@ -27,6 +27,7 @@ kMxfp4Dynamic, kMxfp4Static, kMxfp8Dynamic, + kNvfp4DynamicToken, kNvfp4Static, ) @@ -40,6 +41,7 @@ "fp8_per_block_dynamic": kFp8Dynamic128Sym, "mxfp8": kMxfp8Dynamic, "mxfp4": kMxfp4Dynamic, + "nvfp4_per_token": kNvfp4DynamicToken, "int8_per_channel_static": kInt8StaticChannelSym, # Load-time MXFP4-to-int4 for Kimi-K3 on gfx942. "int4_per_group_32": kInt4Static32, diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 330e7b6aa13a..80aa9122f447 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -334,6 +334,7 @@ def convert_to_nvfp4_moe_kernel_format( a2_scale: torch.Tensor | None, is_act_and_mul: bool, use_a16: bool = False, + trtllm_hidden_alignment: int = 256, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -419,6 +420,7 @@ def convert_to_nvfp4_moe_kernel_format( w2_scale_2=w2_scale_2, a2_scale=a2_scale, is_act_and_mul=is_act_and_mul, + trtllm_hidden_alignment=trtllm_hidden_alignment, ) elif nvfp4_backend == NvFp4MoeBackend.HUMMING: from vllm.model_executor.layers.quantization.utils.humming import ( diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index e33226cad3c8..98af93c30fad 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -12,7 +12,7 @@ import vllm.envs as envs from vllm.config import get_current_vllm_config_or_none -from vllm.config.quantization import QuantSpec +from vllm.config.quantization import QuantizationConfigArgs, QuantSpec from vllm.logger import init_logger from vllm.model_executor.kernels.linear import ( init_fp8_linear_kernel, @@ -63,6 +63,9 @@ ) from vllm.model_executor.layers.quantization.fp8 import Fp8Config, Fp8MoEMethod from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + trtllm_nvfp4_hidden_alignment, +) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( FP8_SCALE_SENTINEL, process_fp8_input_tensor_strategy_moe, @@ -90,6 +93,7 @@ kMxfp8Dynamic, kMxfp8Static, kNvfp4Dynamic, + kNvfp4DynamicToken, kNvfp4Static, ) from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( @@ -724,12 +728,18 @@ def __init__( kv_cache_quant_algo: str | None = None, exclude_modules: list[str] | None = None, group_size: int = 16, + *, + quantization_args: QuantizationConfigArgs | None = None, ) -> None: if exclude_modules is None: exclude_modules = [] super().__init__(exclude_modules) self.quant_method = quant_method self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized + moe_spec = quantization_args.moe if quantization_args is not None else None + self.moe_activation_override: QuantKey | None = ( + moe_spec.activation if moe_spec is not None else None + ) if is_checkpoint_nvfp4_serialized: self.group_size = group_size self.kv_cache_quant_algo = kv_cache_quant_algo @@ -815,6 +825,7 @@ def _from_config( kv_cache_quant_method, exclude_modules, group_size, + quantization_args=original_config.get("_online_quantization_args"), ) @@ -835,18 +846,36 @@ def __init__( ) -> None: super().__init__(moe_config) self.quant_config = quant_config - # W4A16 mode fires for W4A16_NVFP4 on-disk checkpoints. With - # activation_key=None every W4A4 backend's _supports_quant_scheme - # rejects itself (they all require (kNvfp4Static, kNvfp4Dynamic) - # exactly); only Marlin survives. Marlin's MoE path drops - # activation scales in convert_to_nvfp4_moe_kernel_format, so no - # other change is needed. self.use_a16 = quant_config.quant_method == "W4A16_NVFP4" + activation_key = None if self.use_a16 else kNvfp4Dynamic + + if quant_config.moe_activation_override is not None: + if self.use_a16: + raise ValueError( + "NVFP4 activation overrides require a W4A4 checkpoint." + ) + activation_key = quant_config.moe_activation_override + if activation_key not in (kNvfp4Dynamic, kNvfp4DynamicToken): + raise ValueError( + "Unsupported ModelOpt NVFP4 MoE activation override: " + f"{activation_key}. Use quantization_config.moe.activation=" + "'nvfp4_per_token' or omit the override." + ) + self.per_token_activation = activation_key == kNvfp4DynamicToken + self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend( config=self.moe, weight_key=kNvfp4Static, - activation_key=None if self.use_a16 else kNvfp4Dynamic, + activation_key=activation_key, ) + if ( + self.per_token_activation + and self.nvfp4_backend != NvFp4MoeBackend.FLASHINFER_TRTLLM + ): + raise ValueError( + "Per-token NVFP4 activation for pre-quantized weights " + "requires the FlashInfer TRTLLM MoE backend." + ) self.use_global_sf = is_global_sf_supported_for_nvfp4_backend( self.nvfp4_backend @@ -984,6 +1013,10 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: self._build_moe_kernel(layer) return + if self.per_token_activation: + layer.w13_input_scale.data.fill_(1.0) + layer.w2_input_scale.data.fill_(1.0) + # Use a single gscale for w13. if self.moe.is_act_and_mul and not torch.allclose( layer.w13_weight_scale_2[:, 0], layer.w13_weight_scale_2[:, 1] @@ -1016,6 +1049,10 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: a2_scale=layer.w2_input_scale, is_act_and_mul=self.moe.is_act_and_mul, use_a16=self.use_a16, + trtllm_hidden_alignment=trtllm_nvfp4_hidden_alignment( + per_token_activation=self.per_token_activation, + is_act_and_mul=self.moe.is_act_and_mul, + ), ) replace_parameter(layer, "w13_weight", w13) @@ -1039,6 +1076,7 @@ def _build_moe_kernel(self, layer: RoutedExperts) -> None: experts_cls=self.experts_cls, backend=self.nvfp4_backend, routing_tables=layer._expert_routing_tables(), + per_token_activation=self.per_token_activation, ) self.moe_kernel.fused_experts.process_weights_after_loading(layer) @@ -1588,6 +1626,7 @@ def _from_config( group_size: int | None, **kwargs: Any, ) -> "ModelOptMixedPrecisionConfig": + quantization_args = original_config.get("_online_quantization_args") if "quantization" in original_config: quantized_layers = original_config["quantization"].get( "quantized_layers", {} @@ -1626,6 +1665,7 @@ def _from_config( kv_cache_quant_algo=kv_cache_quant_method, exclude_modules=[], group_size=group_size, + quantization_args=quantization_args, ) # Sibling config for layers that declare quant_algo: "W4A16_NVFP4". # get_quant_method resolves this sub-config to the (kNvfp4Static, None) @@ -1638,6 +1678,7 @@ def _from_config( kv_cache_quant_algo=kv_cache_quant_method, exclude_modules=[], group_size=group_size, + quantization_args=quantization_args, ) mxfp8_config = ModelOptMxFp8Config( diff --git a/vllm/model_executor/layers/quantization/online/nvfp4.py b/vllm/model_executor/layers/quantization/online/nvfp4.py index efc97705c538..6c780c2c46aa 100644 --- a/vllm/model_executor/layers/quantization/online/nvfp4.py +++ b/vllm/model_executor/layers/quantization/online/nvfp4.py @@ -19,6 +19,9 @@ from vllm.model_executor.layers.quantization.online.moe_base import ( OnlineMoEMethodBase, ) +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + trtllm_nvfp4_hidden_alignment, +) from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( FLOAT4_E2M1_MAX, ) @@ -87,6 +90,8 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase): (SM100) only. """ + activation_quant_key = kNvfp4DynamicToken + def __init__( self, *, @@ -100,9 +105,13 @@ def __init__( self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend( config=self.moe, weight_key=kNvfp4Static, - activation_key=kNvfp4DynamicToken, + activation_key=self.activation_quant_key, ) + @property + def per_token_activation(self) -> bool: + return self.activation_quant_key == kNvfp4DynamicToken + def process_weights_after_loading(self, layer: Module) -> None: if getattr(layer, "_already_called_process_weights_after_loading", False): return @@ -156,6 +165,10 @@ def _setup_kernel(self, layer: RoutedExperts) -> None: w2_scale_2=layer.w2_weight_scale_2, a2_scale=layer.w2_input_scale, is_act_and_mul=self.moe.is_act_and_mul, + trtllm_hidden_alignment=trtllm_nvfp4_hidden_alignment( + per_token_activation=self.per_token_activation, + is_act_and_mul=self.moe.is_act_and_mul, + ), ) replace_parameter(layer, "w13_weight", w13) @@ -176,7 +189,7 @@ def _setup_kernel(self, layer: RoutedExperts) -> None: experts_cls=self.experts_cls, backend=self.nvfp4_backend, routing_tables=layer._expert_routing_tables(), - per_token_activation=True, + per_token_activation=self.per_token_activation, ) else: # Reload creates new scale tensors; derived kernel scales must use diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 24b95ef89aee..c5a299eea6e9 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -344,6 +344,7 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( w2_scale_2: torch.Tensor, a2_scale: torch.Tensor, is_act_and_mul: bool, + trtllm_hidden_alignment: int = 256, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -398,7 +399,11 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( # Shuffle weights and scales for FI TRTLLM NVFP4 MoE kernels. if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: w13, w13_scale, w2, w2_scale, padded_hidden = align_fp4_moe_hidden_dim_for_fi( - w13, w13_scale, w2, w2_scale + w13, + w13_scale, + w2, + w2_scale, + min_alignment=trtllm_hidden_alignment, ) if layer.moe_config.hidden_dim_unpadded is None: layer.moe_config.hidden_dim_unpadded = layer.moe_config.hidden_dim diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index be28de9b66d6..5a6ee764286b 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -280,6 +280,20 @@ def align_fp4_moe_weights_for_fi( return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_intermediate +def trtllm_nvfp4_hidden_alignment( + per_token_activation: bool, is_act_and_mul: bool +) -> int: + """Hidden-dim alignment for the TRTLLM-Gen NVFP4 MoE cubins. + + The per-token (dynamic activation scale) non-gated variant has far fewer + tile configs for 256-aligned K and its heuristic default tactic fails with + "No valid config found" (e.g. Nemotron's 2688 -> 2816), so pad to 512. + """ + if per_token_activation and not is_act_and_mul: + return 512 + return 256 + + def align_fp4_moe_hidden_dim_for_fi( w13: torch.Tensor, w13_scale: torch.Tensor, diff --git a/vllm/model_executor/model_loader/weight_cache/protocol.py b/vllm/model_executor/model_loader/weight_cache/protocol.py index 72053805cd8b..535c4de96205 100644 --- a/vllm/model_executor/model_loader/weight_cache/protocol.py +++ b/vllm/model_executor/model_loader/weight_cache/protocol.py @@ -28,6 +28,7 @@ import vllm.version from vllm.config import ModelConfig +from vllm.config.utils import normalize_value from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase from vllm.model_executor.model_loader.weight_cache.utils import ( format_socket_role_suffix, @@ -227,12 +228,17 @@ def verify_peer_is_owner(conn: socket.socket) -> None: raise PermissionError(f"Rejecting weight cache connection from uid {peer_uid}") -def _hash_quant_config(quant_config: Any) -> str: - if quant_config is None: - return "" +def _hash_quant_config(quant_config: Any, runtime_quant_config: Any = None) -> str: if hasattr(quant_config, "to_dict"): quant_config = quant_config.to_dict() - payload = json.dumps(quant_config, sort_keys=True, default=str) + payload = json.dumps( + { + "checkpoint": quant_config, + "runtime": normalize_value(runtime_quant_config), + }, + sort_keys=True, + default=str, + ) return safe_hash(payload.encode(), usedforsecurity=False).hexdigest() @@ -333,7 +339,10 @@ def from_model_config( pp_rank=pp_rank, dtype=str(model_config.dtype), quantization=model_config.quantization, - quant_config_hash=_hash_quant_config(quant_config), + quant_config_hash=_hash_quant_config( + quant_config=quant_config, + runtime_quant_config=model_config.quantization_config, + ), revision=model_config.revision, vllm_version=vllm.version.__version__, is_draft=is_draft, diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index c6a72c4c33c6..ac1da8d6987b 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -200,6 +200,9 @@ def get_quant_config( QuantizationConfigArgs, resolve_quantization_config, ) + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptQuantConfigBase, + ) from vllm.model_executor.layers.quantization.online.base import ( OnlineQuantizationConfig, ) @@ -220,6 +223,14 @@ def maybe_compose_online_quantization( ) return checkpoint_config + def make_checkpoint_config(config: dict[str, Any]) -> QuantizationConfig: + if issubclass(quant_cls, ModelOptQuantConfigBase): + assert online_args is None or isinstance( + online_args, QuantizationConfigArgs + ) + config = {**config, "_online_quantization_args": online_args} + return maybe_compose_online_quantization(quant_cls.from_config(config)) + # Read the quantization config from the HF model config, if available. hf_quant_config = getattr(model_config.hf_config, "quantization_config", None) # some vision model may keep quantization_config in their text_config @@ -260,9 +271,7 @@ def maybe_compose_online_quantization( ): pass # fall through to file-based loading below else: - return maybe_compose_online_quantization( - quant_cls.from_config(hf_quant_config) - ) + return make_checkpoint_config(hf_quant_config) # if hf_quant_config is None, we will try to get config from # hf_overrides @@ -376,14 +385,14 @@ def maybe_compose_online_quantization( if model_config.quantization in ("modelopt", "modelopt_mixed"): if config.get("producer", {}).get("name") == "modelopt": - return maybe_compose_online_quantization(quant_cls.from_config(config)) + return make_checkpoint_config(config) else: raise ValueError( f"Unsupported quantization config" f" found for {model_config.quantization} in {f}." ) - return maybe_compose_online_quantization(quant_cls.from_config(config)) + return make_checkpoint_config(config) def get_sparse_attention_config( diff --git a/vllm/models/deepseek_v4/quant_config.py b/vllm/models/deepseek_v4/quant_config.py index be8846fc79c5..f3aaf72d9b56 100644 --- a/vllm/models/deepseek_v4/quant_config.py +++ b/vllm/models/deepseek_v4/quant_config.py @@ -110,6 +110,11 @@ def _get_nvfp4_config(self) -> ModelOptNvFp4Config: kv_cache_quant_algo=None, exclude_modules=[], group_size=16, + quantization_args=( + self.online_quantization_config.args + if self.online_quantization_config is not None + else None + ), ) return self._nvfp4_config diff --git a/vllm/models/deepseek_v41/quant_config.py b/vllm/models/deepseek_v41/quant_config.py index bed60ae75874..384e9488793c 100644 --- a/vllm/models/deepseek_v41/quant_config.py +++ b/vllm/models/deepseek_v41/quant_config.py @@ -114,6 +114,11 @@ def _get_nvfp4_config(self) -> ModelOptNvFp4Config: kv_cache_quant_algo=None, exclude_modules=[], group_size=16, + quantization_args=( + self.online_quantization_config.args + if self.online_quantization_config is not None + else None + ), ) return self._nvfp4_config From d169c5a90354233d2acec687aec201971fa135fb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Thu, 8 Oct 2026 00:16:05 +0200 Subject: [PATCH 041/624] [Docs] Add disagg PD flow-example guidance to /pr-checklist skill (#60111) Signed-off-by: NickLucche Co-authored-by: Mistral Vibe --- .agents/skills/pr-checklist/SKILL.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.agents/skills/pr-checklist/SKILL.md b/.agents/skills/pr-checklist/SKILL.md index 4d73be6adf3e..ae3099b6a076 100644 --- a/.agents/skills/pr-checklist/SKILL.md +++ b/.agents/skills/pr-checklist/SKILL.md @@ -181,6 +181,8 @@ Summarize model evaluations and benchmarks in compact tables with the configurat When fixing bugs or issues, provide a clear root-cause analysis. Explain what caused the problem, how it was identified, and why the chosen solution effectively addresses it. A patch which resolves the issue without a proper root-cause analysis may lead to recurring problems and increased maintenance overhead. +For fixes that involve disaggregated prefill/decode (PD) logic, such as the NIXL connector (`vllm/distributed/kv_transfer/kv_connector/v1/nixl/`), make a best effort to include a minimal example of the flow being changed and/or the initial flow that caused the failure. A short numbered sequence (or diagram) of a request's path across the P and D instances, e.g. scheduler admission, prefill execution, KV transfer registration and write, decode-side read, token generation, lets reviewers trace the root cause and the fix without reconstructing the distributed flow from the diff. Keep the example minimal: only the steps and components involved in the change. See the [disaggregated prefill](../../../docs/features/disagg_prefill.md), [NIXL KV push](../../../docs/design/nixl_kv_push_connector.md), and [NIXL KV cache lease](../../../docs/design/nixl_kv_cache_lease.md) docs for the flows involved. + #### 4.4: Implementation Details Provide a detailed explanation of the implementation, including design decisions, trade-offs, and any relevant technical details. Highlight any non-obvious choices, tradeoffs, potential limitations, and areas where maintainer attention is required. From ef63c23d35acccbba8e014fc88da19d541da38a9 Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Wed, 7 Oct 2026 15:28:42 -0700 Subject: [PATCH 042/624] [Bugfix][DSv4.1] Skip SWA bounded replay for KV loads that carry the window (#59197) Signed-off-by: Yifan Qiao Co-authored-by: Codex Co-authored-by: Claude Opus 5.5 --- .../aux_output_connector/test_store.py | 16 +++++ tests/v1/core/test_prefix_caching.py | 4 ++ tests/v1/core/test_prefix_replay.py | 27 ++++++++ .../unit/test_nixl_connector_hma.py | 35 +++++++--- tests/v1/kv_connector/unit/utils.py | 3 + .../aux_output_connector/connector.py | 6 +- .../v1/mooncake/mooncake_connector.py | 19 ++++- .../kv_connector/v1/nixl/base_scheduler.py | 13 +++- vllm/v1/core/sched/scheduler.py | 69 ++++++++++++++----- 9 files changed, 161 insertions(+), 31 deletions(-) diff --git a/tests/distributed/aux_output_connector/test_store.py b/tests/distributed/aux_output_connector/test_store.py index 39277f3e7e0e..24d67c3d60e1 100644 --- a/tests/distributed/aux_output_connector/test_store.py +++ b/tests/distributed/aux_output_connector/test_store.py @@ -96,6 +96,10 @@ class _SchedulerRequest: default_factory=lambda: SimpleNamespace(routed_experts_prompt_start=0) ) + @property + def num_prompt_tokens(self) -> int: + return self.num_tokens - self.num_output_tokens + def is_finished(self) -> bool: return self.finished @@ -1439,6 +1443,18 @@ def test_scheduler_starts_worker_output_at_requested_prompt_token(): assert metadata.requests == {request.request_id: 2} + # A P/D prefiller cuts its prompt short of the last token, which can pass + # the requested start; the rows then start at the end of the cut prompt. + connector = _make_connector() + request = _scheduler_request( + "cut", [b"a" * 32], num_tokens=3, num_output_tokens=0, prompt_start=4 + ) + metadata = connector.build_connector_meta( + _step_output([request.request_id], [0], [3]), + {request.request_id: request}, + ) + assert metadata.requests == {request.request_id: 3} + def test_scheduler_connector_preserves_request_finish_order(): connector = _make_connector() diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 3972fe5230da..427201a7ceef 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -594,8 +594,12 @@ def test_nixl_hisparse_full_block_import_keeps_a_writable_tail(num_tokens): failed_recving_kv_req_ids=set(), finished_recving_kv_req_ids={request.request_id}, prefix_replay_tokens=0, + prefix_replay_group_ids=(), ) scheduler._mark_prefix_replay = MethodType(Scheduler._mark_prefix_replay, scheduler) + scheduler._load_restores_replay_window = MethodType( + Scheduler._load_restores_replay_window, scheduler + ) Scheduler._update_waiting_for_remote_kv(scheduler, request) assert request.num_tokens - request.num_computed_tokens == 1 assert _allocate_scheduled(manager, request, num_new_tokens=1) is not None diff --git a/tests/v1/core/test_prefix_replay.py b/tests/v1/core/test_prefix_replay.py index da5718580848..323413a86dce 100644 --- a/tests/v1/core/test_prefix_replay.py +++ b/tests/v1/core/test_prefix_replay.py @@ -247,6 +247,33 @@ def test_remote_kv_hit_is_taken_in_whole_blocks(): assert swa_blocks[new_req.replay_start // BLOCK_SIZE] is not swa_manager._null_block +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("matched", [WINDOW, 4 * BLOCK_SIZE - 1]) +def test_kv_load_carrying_window_does_not_replay(is_async, matched): + """A load of the request's own blocks (P/D) restores every transfer group, + window included, so the hit is taken as is: no replay, no cut to whole + blocks, and a hit no longer than the window is kept.""" + scheduler = _replay_scheduler( + use_kv_connector=MockKVConfig(matched_tokens=matched, is_async=is_async) + ) + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE + )[0] + scheduler.connector.get_loaded_kv_cache_group_ids = lambda request: (FULL, SWA) + scheduler.add_request(request) + out = scheduler.schedule() + if is_async: + assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS + scheduler.update_from_output( + out, create_model_runner_output([], finished_recving={request.request_id}) + ) + out = scheduler.schedule() + new_req = _new_req_data(out, request) + assert new_req.num_computed_tokens == matched + assert new_req.replay_start == 0 + assert out.num_scheduled_tokens[request.request_id] == NUM_PROMPT_TOKENS - matched + + @pytest.mark.parametrize("via_connector", [False, True]) def test_hit_no_longer_than_window_is_ignored(via_connector): """Such a hit would be recomputed in full anyway. It is not adopted (an diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index f96595572162..b4b5e063445f 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -1559,17 +1559,22 @@ def test_failed_load_rezeroes_unwritten_skipped_blocks(): @pytest.mark.cpu_test @pytest.mark.parametrize( - "has_mamba,is_hma_required,expected_count", + "has_mamba,is_hma_required,bounded_replay,expected_count", [ - (True, True, 9), - (False, False, 10), - (False, True, 10), + (True, True, False, 9), + (False, False, False, 10), + (False, True, False, 10), + (False, True, True, 9), ], - ids=["mamba", "fa_only", "swa_only"], + ids=["mamba", "fa_only", "swa_only", "swa_bounded_replay"], ) -def test_mamba_n1_d_side(has_mamba, is_hma_required, expected_count): - """D-side: Mamba gets N-1 matched tokens, non-Mamba gets N.""" - sched = make_nixl_scheduler(has_mamba=has_mamba, is_hma_required=is_hma_required) +def test_mamba_n1_d_side(has_mamba, is_hma_required, bounded_replay, expected_count): + """D-side: Mamba and SWA bounded replay get N-1 matched tokens, others N.""" + sched = make_nixl_scheduler( + has_mamba=has_mamba, + is_hma_required=is_hma_required, + bounded_replay=bounded_replay, + ) req = create_request(num_tokens=10, do_remote_prefill=True) count, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=0) @@ -1638,6 +1643,20 @@ def test_mamba_n1_p_side_truncation(): fa_sched.on_new_request(fa_req) assert len(fa_req.prompt_token_ids) == fa_original + # SWA bounded replay: the window is not replayed after the load, so the + # prefiller stops short of the last token too. + swa_sched = make_nixl_scheduler(is_hma_required=True, bounded_replay=True) + swa_req = create_request(num_tokens=10, do_remote_decode=True) + swa_sched.on_new_request(swa_req) + assert len(swa_req.prompt_token_ids) == 9 + + # A request that skips reading the prefix cache (prompt logprobs) is not + # cut: the decoder loads nothing and recomputes its whole prompt. + logprobs_req = create_request(num_tokens=10, do_remote_decode=True) + logprobs_req.sampling_params.skip_reading_prefix_cache = True + swa_sched.on_new_request(logprobs_req) + assert len(logprobs_req.prompt_token_ids) == 10 + @pytest.mark.cpu_test @pytest.mark.parametrize( diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 86062fc47848..d738e183952f 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -532,6 +532,7 @@ def make_kv_cache_config( def make_nixl_scheduler( has_mamba: bool = False, is_hma_required: bool = False, + bounded_replay: bool = False, heartbeat: bool = False, kv_lease_duration: int = 30, ): @@ -548,6 +549,7 @@ def make_nixl_scheduler( sched = object.__new__(NixlConnectorScheduler) sched._has_mamba = has_mamba + sched._bounded_replay = bounded_replay sched._is_hma_required = is_hma_required sched.kv_cache_config = make_kv_cache_config( block_size=16, @@ -615,6 +617,7 @@ def make_nixl_push_scheduler( sched.side_channel_port = 5600 sched.is_bidirectional_kv_xfer_enabled = is_bidirectional_kv_xfer_enabled sched._has_mamba = has_mamba + sched._bounded_replay = False sched.kv_cache_config = make_kv_cache_config( block_size=16, mamba_enabled=has_mamba, diff --git a/vllm/distributed/aux_output_connector/connector.py b/vllm/distributed/aux_output_connector/connector.py index bd8cbf4d183d..c27e23ea4be6 100644 --- a/vllm/distributed/aux_output_connector/connector.py +++ b/vllm/distributed/aux_output_connector/connector.py @@ -67,8 +67,12 @@ def build_connector_meta( block_hashes_by_request[request_id] = packed self._sent_hash_counts[request_id] = len(request.block_hashes) assert request.sampling_params is not None + # A P/D prefiller may cut the prompt short of the requested start. scheduled_requests[request_id] = max( - request.sampling_params.routed_experts_prompt_start, + min( + request.sampling_params.routed_experts_prompt_start, + request.num_prompt_tokens, + ), 0 if request.num_output_tokens == 0 else request.num_tokens - 1, ) # A settled token can complete a hash block after the next async schedule diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 5f064259e8b2..8300cd950b0d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -897,6 +897,10 @@ def __init__( # GDN is represented as a MambaSpec in vLLM. This Mooncake MambaSpec # path is currently tested with GDN; Mamba2 is not validated yet. self._has_mamba = kv_cache_config.has_mamba_layers + self._bounded_replay = any( + g.kv_cache_spec.prefix_replay_tokens > 0 + for g in kv_cache_config.kv_cache_groups + ) # Requests that need to start recv/send. # New requests are added by update_state_after_alloc in @@ -941,8 +945,10 @@ def get_sw_clipped_blocks( def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int: """D-side only. Returns N-1 for Mamba models since the decoder - always recomputes the last token and must start from h(N-1).""" - if self._has_mamba and num_prompt_tokens > 1: + always recomputes the last token and must start from h(N-1), and for + DSV41 SWA bounded replay, whose window is not replayed after a P/D load + and must be laid out for that token.""" + if (self._has_mamba or self._bounded_replay) and num_prompt_tokens > 1: return num_prompt_tokens - 1 return num_prompt_tokens @@ -958,6 +964,9 @@ def _truncate_mamba_request_for_prefill(self, request: "Request") -> None: params is not None and not params.get("_p_side_truncated") and request.num_prompt_tokens > 1 + # A request that skips reading the prefix cache (prompt logprobs) + # loads nothing on the decoder, which recomputes its whole prompt. + and not request.get_skip_reading_prefix_cache() ): if request.prompt_token_ids is not None: request.prompt_token_ids.pop() @@ -973,7 +982,11 @@ def _truncate_mamba_request_for_prefill(self, request: "Request") -> None: def on_new_request(self, request: "Request") -> None: params = request.kv_transfer_params - if params is not None and params.get("do_remote_decode") and self._has_mamba: + if ( + params is not None + and params.get("do_remote_decode") + and (self._has_mamba or self._bounded_replay) + ): self._truncate_mamba_request_for_prefill(request) def get_num_new_matched_tokens( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index 92f0cbc58a0b..55ee98d46596 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -101,6 +101,10 @@ def __init__( ) ) self._has_mamba = kv_cache_config.has_mamba_layers + self._bounded_replay = any( + g.kv_cache_spec.prefix_replay_tokens > 0 + for g in kv_cache_config.kv_cache_groups + ) logger.info("Initializing NIXL Scheduler %s", engine_id) if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager: @@ -376,14 +380,16 @@ def _prefill_backoff(self) -> int: """Trailing prompt tokens the prefiller must not compute; the decoder recomputes them locally. - Mamba needs h(N-1) so the decoder can derive h(N) itself. Multi-module + Mamba needs h(N-1) so the decoder can derive h(N) itself. DSV41's sliding + window under bounded replay is not replayed after a P/D load, so it must + be laid out for the token the decoder recomputes. Multi-module MTP needs to keep its whole lookahead window off of the prefiller, which would otherwise embed the unverified drafts in the MTP layer's KV cache. The decoder would never rebuild them, because the update is sized by the rejection count, which is zero for the first decode. """ return max( - 1 if self._has_mamba else 0, + 1 if self._has_mamba or self._bounded_replay else 0, self.vllm_config.num_prefill_lookahead_tokens - 1, ) @@ -412,6 +418,9 @@ def _truncate_request_for_prefill(self, request: "Request") -> None: # Guard against repeated truncation after preemption/reschedule. and not params.get("_p_side_truncated") and request.num_prompt_tokens > backoff + # A request that skips reading the prefix cache (prompt logprobs) + # loads nothing on the decoder, which recomputes its whole prompt. + and not request.get_skip_reading_prefix_cache() ): if request.prompt_token_ids is not None: del request.prompt_token_ids[-backoff:] diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index f908a31a1b20..91ea8e93e42c 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -290,10 +290,14 @@ def __init__( # DSV41 SWA bounded replay: groups that declare a replay window are rebuilt # after a prefix hit by recomputing its trailing tokens. One window # for all such groups, so the rewind matches every group's allocation. - replay_windows = { - group.kv_cache_spec.prefix_replay_tokens - for group in kv_cache_config.kv_cache_groups + self.prefix_replay_group_ids = tuple( + group_id + for group_id, group in enumerate(kv_cache_config.kv_cache_groups) if group.kv_cache_spec.prefix_replay_tokens > 0 + ) + replay_windows = { + kv_cache_config.kv_cache_groups[group_id].kv_cache_spec.prefix_replay_tokens + for group_id in self.prefix_replay_group_ids } assert len(replay_windows) <= 1, ( f"Prefix replay windows should agree: {sorted(replay_windows)}" @@ -973,14 +977,17 @@ def skip_request(from_queue: RequestQueue) -> None: skip_request(request_queue) continue - if self.prefix_replay_tokens: - # SWA bounded replay recomputes the hit's last + if self.prefix_replay_tokens and not ( + ext_tokens and self._load_restores_replay_window(request) + ): + # DSV41 SWA bounded replay recomputes the hit's last # window, from the block holding its first # token; the sliding-window groups retire # whole blocks below the window of the hit's # next token. Hits end on a block boundary, or # a hit ending one token short of one would - # retire the block the replay starts in. + # retire the block the replay starts in. A load + # that carries the window replays nothing. ext_tokens -= ext_tokens % self.block_size load_kv_async = load_kv_async and ext_tokens > 0 @@ -1026,9 +1033,17 @@ def skip_request(from_queue: RequestQueue) -> None: num_new_local_computed_tokens + num_external_computed_tokens ) assert num_computed_tokens <= request.num_tokens - if 0 < num_computed_tokens <= self.prefix_replay_tokens: - # SWA bounded replay: a hit no longer than the replayed - # window would be recomputed in full and save nothing. + # DSV41 SWA bounded replay: a KV load that carries the + # window is taken as is; any other hit no longer than the + # replayed window would be recomputed in full and save + # nothing. + window_loaded = bool( + num_external_computed_tokens + ) and self._load_restores_replay_window(request) + if ( + not window_loaded + and 0 < num_computed_tokens <= self.prefix_replay_tokens + ): new_computed_blocks = ( self.kv_cache_manager.empty_kv_cache_blocks ) @@ -1091,7 +1106,7 @@ def skip_request(from_queue: RequestQueue) -> None: if did_prefix_cache_lookup: # Fresh admission: local and/or sync external hit. num_replay_tokens = self._mark_prefix_replay( - request, num_computed_tokens + request, num_computed_tokens, window_loaded ) num_computed_tokens -= num_replay_tokens # Number of tokens to be scheduled. @@ -3083,16 +3098,31 @@ def _inflight_prefill_reserved_blocks(self) -> int: self._request_remaining_blocks(req) for req in self._inflight_prefills ) - def _mark_prefix_replay(self, request: Request, num_hit_tokens: int) -> int: + def _load_restores_replay_window(self, request: Request) -> bool: + """Whether the KV connector's load for ``request`` restores every DSV41 + SWA bounded-replay group, as a P/D transfer of the request's own blocks + does; a prefix-cache store holds no window.""" + assert self.connector is not None + return bool(self.prefix_replay_group_ids) and set( + self.prefix_replay_group_ids + ).issubset(self.connector.get_loaded_kv_cache_group_ids(request)) + + def _mark_prefix_replay( + self, request: Request, num_hit_tokens: int, window_loaded: bool + ) -> int: """Record where a prefix hit's replay starts and return the number of hit tokens to recompute; the caller rewinds the computed count by it. The replayed tokens are the hit's last window (see ``Request.replay_start``): the worker rebuilds their sliding-window KV - and leaves their cached KV alone. + and leaves their cached KV alone. A hit whose KV load carries the + window replays nothing. """ if not self.prefix_replay_tokens: return 0 + if window_loaded: + request.replay_start = 0 + return 0 num_replay_tokens = min(self.prefix_replay_tokens, num_hit_tokens) request.replay_start = num_hit_tokens - num_replay_tokens return num_replay_tokens @@ -3106,7 +3136,8 @@ def _update_waiting_for_remote_kv(self, request: Request) -> None: """ assert self.connector is not None - if request.request_id in self.failed_recving_kv_req_ids: + load_failed = request.request_id in self.failed_recving_kv_req_ids + if load_failed: # Request had KV load failures; num_computed_tokens was already # updated in _update_requests_with_invalid_blocks if request.num_computed_tokens: @@ -3132,11 +3163,15 @@ def _update_waiting_for_remote_kv(self, request: Request) -> None: # This will cache the blocks iff caching is enabled. self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens) - # SWA bounded replay recomputes the tail of the hit, which covers the - # last token; otherwise a full prompt hit re-computes that token so the - # next one can be sampled. + # DSV41 SWA bounded replay recomputes the tail of a hit loaded without + # its window, or of a failed load, which may have lost it; the replay + # covers the last token. Otherwise a full prompt hit re-computes that + # token so the next one can be sampled. num_replay_tokens = self._mark_prefix_replay( - request, request.num_computed_tokens + request, + request.num_computed_tokens, + window_loaded=not load_failed + and self._load_restores_replay_window(request), ) if num_replay_tokens > 0: request.num_computed_tokens -= num_replay_tokens From 4ac5a36a8e38b79c8084589b0bc4f2a249bf7c2a Mon Sep 17 00:00:00 2001 From: Artem Perevedentsev Date: Thu, 8 Oct 2026 01:44:17 +0300 Subject: [PATCH 043/624] [CI/Build][NVIDIA] Build NIXL with NIXL EP from source in Rubin image (#60090) Signed-off-by: Artem Perevedentsev --- .buildkite/release-pipeline.yaml | 5 + docker/Dockerfile | 42 ++ docker/versions.json | 3 + docs/contributing/dockerfile/dockerfile.md | 5 + .../installation/gpu.cuda.inc.md | 5 + tools/build_nixl_with_torch215_for_rubin.py | 619 ++++++++++++++++++ 6 files changed, 679 insertions(+) create mode 100644 tools/build_nixl_with_torch215_for_rubin.py diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index 5eba9fd2e114..6064f93104f5 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -355,6 +355,10 @@ steps: queue: arm64_cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" + # TODO(arpera): + # Remove BUILD_NIXL=true once NIXL has a wheel built against torch 2.15 + # For details: https://github.com/vllm-project/vllm/pull/60090 + # Tracking issue report: https://github.com/ai-dynamo/nixl/issues/2297 - | DOCKER_BUILDKIT=1 docker buildx build --load \ $(bash .buildkite/scripts/docker-build-metadata-args.sh cu134) \ @@ -366,6 +370,7 @@ steps: --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \ --build-arg INSTALL_KV_CONNECTORS=true \ --build-arg INSTALL_RUBIN_PRERELEASE=true \ + --build-arg BUILD_NIXL=true \ --build-arg TRITON_INSTALL_FROM_SOURCE_REPO=https://github.com/triton-lang/triton.git \ --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \ --build-arg RUN_WHEEL_CHECK=false \ diff --git a/docker/Dockerfile b/docker/Dockerfile index d8a1258a74c5..fb8fa6b264ae 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -92,6 +92,11 @@ ARG UV_KEYRING_PROVIDER=${PIP_KEYRING_PROVIDER} # Flag enables built-in KV-connector dependency libs into docker images ARG INSTALL_KV_CONNECTORS=false +# Flag builds NIXL, including its NIXL EP extension for the installed PyTorch, +# from source and installs it in place of the nixl packages from the KV-connector +# requirements +ARG BUILD_NIXL=false + #################### BASE BUILD IMAGE #################### # prepare basic build environment FROM ${BUILD_BASE_IMAGE} AS base @@ -548,6 +553,24 @@ RUN --mount=type=cache,target=/opt/uv/cache \ (cd "${TRITON_WORKSPACE}/dist" \ && sha256sum build-state > wheel.sha256); \ fi + +# Build NIXL and its NIXL EP extension against the PyTorch installed here when +# requested. The wheels are installed by the final image, which has no compiler. +ARG BUILD_NIXL +COPY tools/build_nixl_with_torch215_for_rubin.py /tmp/build_nixl_for_rubin.py +RUN --mount=type=bind,source=requirements,target=/tmp/vllm-requirements,ro \ + set -eux; \ + NIXL_WORKSPACE=/tmp/nixl_workspace; \ + mkdir -p "${NIXL_WORKSPACE}/dist"; \ + if [ "${BUILD_NIXL}" = "true" ]; then \ + dnf install -y --setopt=install_weak_deps=False autoconf automake libtool; \ + dnf clean all; \ + rm -rf /var/cache/dnf; \ + NIXL_REQUIREMENTS_FILE=/tmp/vllm-requirements/kv_connectors.txt \ + WHEELS_CACHE_HOME="${NIXL_WORKSPACE}/dist" \ + MAX_JOBS="${max_jobs:-}" \ + python3 /tmp/build_nixl_for_rubin.py --skip-system-deps --build-only; \ + fi #################### EXTENSIONS BUILD IMAGE #################### #################### WHEEL BUILD IMAGE #################### @@ -1260,6 +1283,25 @@ RUN --mount=type=cache,target=/opt/uv/cache \ /tmp/install_triton_from_source.sh \ /tmp/triton-source-wheel.sha256 +# Install source-built NIXL last, replacing the nixl packages installed above, so +# the image has a NIXL EP build for its PyTorch version. There is no GPU driver +# at build time, so check for the extension instead of importing it. +ARG BUILD_NIXL +RUN --mount=type=cache,target=/opt/uv/cache \ + --mount=type=bind,from=extensions-build,source=/tmp/nixl_workspace/dist,target=/tmp/nixl-dist,ro \ + set -eu; \ + if [ "${BUILD_NIXL}" = "true" ]; then \ + ls /tmp/nixl-dist/*.whl; \ + uv pip uninstall --system nixl nixl-cu12 nixl-cu13; \ + uv pip install --system --no-deps /tmp/nixl-dist/*.whl; \ + python3 -c "import glob, importlib.util as u, os, torch; \ + spec = u.find_spec('nixl_ep_cu' + os.environ['CUDA_VERSION'].split('.')[0]); \ + mm = ''.join(torch.__version__.split('.')[:2]); \ + ext = glob.glob(os.path.join(spec.submodule_search_locations[0], 'nixl_ep_cpp_torch' + mm + '.*')) if spec else []; \ + assert ext, 'NIXL EP has no build for torch ' + torch.__version__; \ + print('NIXL EP extension:', ext[0])"; \ + fi + ENV VLLM_USAGE_SOURCE production-docker-image ENV VLLM_BUILD_COMMIT=${VLLM_BUILD_COMMIT:-unknown} \ VLLM_BUILD_PIPELINE=${VLLM_BUILD_PIPELINE:-local} \ diff --git a/docker/versions.json b/docker/versions.json index bcec27fd5590..68f426d0c7b5 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -37,6 +37,9 @@ "INSTALL_KV_CONNECTORS": { "default": "false" }, + "BUILD_NIXL": { + "default": "false" + }, "SCCACHE_BUCKET_NAME": { "default": "vllm-build-sccache" }, diff --git a/docs/contributing/dockerfile/dockerfile.md b/docs/contributing/dockerfile/dockerfile.md index db46540f9f2c..279c8f0c6a34 100644 --- a/docs/contributing/dockerfile/dockerfile.md +++ b/docs/contributing/dockerfile/dockerfile.md @@ -24,6 +24,11 @@ The `extensions-build` stage can also produce an optional source-built Triton wheel. `vllm-openai-base` installs that wheel after its other Python dependencies so dependency resolution cannot restore an older Triton version. +With `BUILD_NIXL=true`, `extensions-build` also builds NIXL, including its NIXL +EP extension for the installed PyTorch, with `tools/build_nixl_with_torch215_for_rubin.py`. +`vllm-openai-base` then replaces the NIXL packages installed from the KV-connector +requirements with these wheels. + >
> ![](../../assets/contributing/dockerfile-stages-dependency.png){ align="center" alt="query" width="100%" } >
diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md index 0d0d0c0c85ab..86e38c570788 100644 --- a/docs/getting_started/installation/gpu.cuda.inc.md +++ b/docs/getting_started/installation/gpu.cuda.inc.md @@ -439,6 +439,11 @@ BuildKit does not automatically invalidate cached layers when a mutable Git ref changes. Use `--no-cache-filter extensions-build` to refresh an empty, branch, or tag revision. +Set `BUILD_NIXL=true` to build NIXL from source. NIXL's release wheels do not +include the NIXL EP extension for the PyTorch nightly used by the Rubin build. +The build uses the NIXL version pinned in `requirements/kv_connectors.txt` and +replaces the NIXL packages installed from the KV-connector requirements. + For `FINAL_BASE_IMAGE`, use the public, multi-arch `nvidia/cuda:13.4.1-base-ubuntu24.04` image. Set `NCCL_VERSION` to 2.32.3 or newer, the first NCCL release with Rubin (SM107) support. diff --git a/tools/build_nixl_with_torch215_for_rubin.py b/tools/build_nixl_with_torch215_for_rubin.py new file mode 100644 index 000000000000..ee9fa72ae959 --- /dev/null +++ b/tools/build_nixl_with_torch215_for_rubin.py @@ -0,0 +1,619 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# TODO(arpera): +# This script is temporary. It builds nixl from source code against torch 2.15 +# because currently there is no nixl wheel built against torch 2.15 which +# vLLM docker image for Rubin requires. +# Remove this script from vLLM once there is nixl wheel built against torch 2.15. + +import argparse +import glob +import os +import shutil +import subprocess +import sys +import urllib.request + +# --- Configuration --- +WHEELS_CACHE_HOME = os.environ.get("WHEELS_CACHE_HOME", "/tmp/wheels_cache") +ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) +UCX_DIR = os.path.join("/tmp", "ucx_source") +NIXL_DIR = os.path.join("/tmp", "nixl_source") +UCX_INSTALL_DIR = os.path.join("/tmp", "ucx_install") +NIXL_BUILD_DIR = os.path.join("/tmp", "nixl_build") +BUILD_TOOLS_DIR = os.path.join("/tmp", "nixl_build_tools") +UCX_REPO_URL = "https://github.com/openucx/ucx.git" +NIXL_REPO_URL = "https://github.com/ai-dynamo/nixl.git" +UCX_REF = os.environ.get("UCX_REF", "v1.23.x") +TEMP_WHEEL_DIR = os.path.join("/tmp", "nixl_temp_wheelhouse") +NIXL_REQUIREMENTS_FILE = os.environ.get( + "NIXL_REQUIREMENTS_FILE", + os.path.join(ROOT_DIR, os.pardir, "requirements", "kv_connectors.txt"), +) +MAX_JOBS = int(os.environ.get("MAX_JOBS") or os.cpu_count() or 1) +CUDA_HOME = os.environ.get("CUDA_HOME", "/usr/local/cuda") +# SM targets NIXL EP can run on (it requires sm_90 or newer). +NIXL_EP_CUDA_ARCHS = ("90", "100", "103", "107", "110", "120") + +# Libraries that must come from the host or from torch, not from the wheel. +AUDITWHEEL_EXCLUDES = [ + "libcuda*", + "libcufile*", + "libcuobjclient*", + "libssl*", + "libcrypto*", + "libefa*", + "libhwloc*", + "libfabric*", + "libtorch*", + "libc10*", + "libdoca*", + "libred_client*", + "libred_async*", + "liblz4*", +] + + +# --- Helper Functions --- +def get_latest_nixl_version(): + """Helper function to get latest release version of NIXL. + + Raises: + RuntimeError: If the latest release cannot be determined. + + """ + # /releases/latest redirects to the tag page. Unlike the REST API it is not + # rate limited for anonymous clients, which are shared on compute nodes. + request = urllib.request.Request( + f"{NIXL_REPO_URL.removesuffix('.git')}/releases/latest", method="HEAD" + ) + try: + with urllib.request.urlopen(request, timeout=30) as response: + return response.geturl().rsplit("/", 1)[-1] + except OSError as e: + raise RuntimeError( + "Cannot determine the latest NIXL release, set NIXL_VERSION." + ) from e + + +def get_pinned_nixl_version(): + """Returns the NIXL version pinned in the vLLM requirements, if there is one.""" + try: + with open(NIXL_REQUIREMENTS_FILE) as requirements: + for line in requirements: + requirement = line.split("#")[0].split(";")[0] + name, separator, version = requirement.partition("==") + if separator and name.strip() == "nixl" and version.split(): + return version.split()[0] + except OSError: + pass + return None + + +def resolve_nixl_version(): + """Returns the NIXL version to build and where it comes from.""" + if os.environ.get("NIXL_VERSION"): + return os.environ["NIXL_VERSION"], "the NIXL_VERSION environment variable" + pinned_version = get_pinned_nixl_version() + if pinned_version: + return pinned_version, os.path.normpath(NIXL_REQUIREMENTS_FILE) + return get_latest_nixl_version(), "the latest NIXL release" + + +NIXL_VERSION, NIXL_VERSION_SOURCE = resolve_nixl_version() +# Version as it appears in wheel file names (tags may carry a "v" prefix). +NIXL_WHEEL_VERSION = NIXL_VERSION.removeprefix("v") + + +def run_command(command, cwd=".", env=None): + """Helper function to run a shell command and check for errors.""" + print(f"--> Running command: {' '.join(command)} in '{cwd}'", flush=True) + subprocess.check_call(command, cwd=cwd, env=env) + + +def is_pip_package_installed(package_name): + """Checks if a package is installed via pip without raising an exception.""" + result = subprocess.run( + [sys.executable, "-m", "pip", "show", package_name], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + return result.returncode == 0 + + +def find_nixl_wheel_in_cache(cache_dir): + """Finds the nixl-cuXX wheel file in the specified cache directory.""" + # The repaired wheel will have a 'manylinux' tag, but this glob still works. + search_pattern = os.path.join(cache_dir, f"nixl_*{NIXL_WHEEL_VERSION}*.whl") + wheels = glob.glob(search_pattern) + if wheels: + # Sort to get the most recent/highest version if multiple exist + wheels.sort() + return wheels[-1] + return None + + +def find_nixl_meta_wheel(wheel_dir): + """Finds the `nixl` meta package wheel (the nixl_cuXX dispatcher).""" + wheels = glob.glob(os.path.join(wheel_dir, f"nixl-{NIXL_WHEEL_VERSION}-*.whl")) + return wheels[-1] if wheels else None + + +def get_cuda_version(): + """Returns the (major, minor) CUDA toolkit version, or None without CUDA.""" + nvcc = os.path.join(CUDA_HOME, "bin", "nvcc") + if not os.path.exists(nvcc): + return None + output = subprocess.check_output([nvcc, "--version"], text=True) + _, separator, release = output.partition("release ") + try: + if not separator: + raise ValueError + major, minor = release.split(",")[0].split(".")[:2] + return int(major), int(minor) + except ValueError: + raise RuntimeError( + f"Cannot parse the CUDA version from `{nvcc} --version`." + ) from None + + +def get_default_cuda_arch_list(): + """Returns the SM targets of NIXL EP that the installed nvcc can compile.""" + nvcc = os.path.join(CUDA_HOME, "bin", "nvcc") + supported = subprocess.check_output([nvcc, "--list-gpu-code"], text=True).split() + return ",".join(sm for sm in NIXL_EP_CUDA_ARCHS if f"sm_{sm}" in supported) + + +def is_torch_available(): + """Checks whether torch can be imported by the current interpreter.""" + result = subprocess.run( + [sys.executable, "-c", "import torch"], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + return result.returncode == 0 + + +def is_nixl_installed(build_nixl_ep): + """Checks that the requested NIXL version, and its EP extension for the + installed torch when requested, is already installed.""" + cuda_version = get_cuda_version() + if cuda_version is None: + return is_pip_package_installed("nixl") + checks = [ + "import importlib, importlib.metadata as m", + f"assert m.version('nixl') == '{NIXL_WHEEL_VERSION}'", + f"assert m.version('nixl-cu{cuda_version[0]}') == '{NIXL_WHEEL_VERSION}'", + ] + if build_nixl_ep: + # The package picks its extension by torch version and fails to import + # when none was built for the installed torch. + checks.append(f"importlib.import_module('nixl_ep_cu{cuda_version[0]}')") + result = subprocess.run( + [sys.executable, "-c", "\n".join(checks)], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + return result.returncode == 0 + + +def install_system_dependencies(build_nixl_ep=False): + """Installs required system packages using apt-get if run as root.""" + if shutil.which("apt-get") is None: + print("--> apt-get not found, relying on the build tools in PATH.", flush=True) + return + if os.geteuid() != 0: + print("\n---", flush=True) + print( + "WARNING: Not running as root. \ + Skipping system dependency installation.", + flush=True, + ) + print( + "Please ensure the listed packages are installed on your system:", + flush=True, + ) + print( + " patchelf build-essential git cmake ninja-build \ + autotools-dev automake libtool libtool-bin \ + libibverbs-dev", + flush=True, + ) + print("---\n", flush=True) + return + + print("--- Running as root. Installing system dependencies... ---", flush=True) + apt_packages = [ + "patchelf", # <-- Add patchelf here + "build-essential", + "git", + "cmake", + "ninja-build", + "autotools-dev", + "automake", + "libtool", + "libtool-bin", + "pkg-config", + "libibverbs-dev", + ] + # -base CUDA images lack the headers UCX (NVML) and the torch extension + # (cuSPARSE/cuSOLVER/cuBLAS, included by ATen/cuda/CUDAContext.h) need. + cuda_version = get_cuda_version() + if cuda_version is not None: + major, minor = cuda_version + if not os.path.exists(os.path.join(CUDA_HOME, "include", "nvml.h")): + apt_packages.append(f"cuda-nvml-dev-{major}-{minor}") + if build_nixl_ep: + apt_packages += [ + f"{package}-{major}-{minor}" + for package in ("libcusparse-dev", "libcublas-dev", "libcusolver-dev") + ] + run_command(["apt-get", "update"]) + run_command(["apt-get", "install", "-y"] + apt_packages) + print("--- System dependencies installed successfully. ---\n", flush=True) + + +def install_python_build_dependencies(): + """Installs the Python tools needed to build and repair the NIXL wheel.""" + run_command( + [ + sys.executable, + "-m", + "pip", + "install", + "auditwheel", + "meson-python", + "ninja", + "patchelf", + "pybind11", + "pyyaml", + "setuptools>=80.9.0", + "tomlkit", + ] + ) + # An older meson already on the system (e.g. from the distribution) makes + # meson-python fail to package NIXL and pip would keep it, so use our own. + run_command( + [ + sys.executable, + "-m", + "pip", + "install", + "--upgrade", + f"--target={BUILD_TOOLS_DIR}", + "meson", + ] + ) + + +def build_ucx(ucx_install_path): + """Builds UCX from source, with CUDA support if CUDA is installed.""" + if not os.path.exists(UCX_DIR): + run_command(["git", "clone", UCX_REPO_URL, UCX_DIR]) + ucx_source_path = os.path.abspath(UCX_DIR) + run_command(["git", "checkout", UCX_REF], cwd=ucx_source_path) + run_command(["./autogen.sh"], cwd=ucx_source_path) + configure_command = [ + "./contrib/configure-release-mt", + f"--prefix={ucx_install_path}", + "--enable-shared", + "--disable-static", + "--disable-doxygen-doc", + "--enable-experimental-api", + "--enable-optimizations", + "--without-avx", + "--enable-cma", + "--enable-devel-headers", + *([f"--with-cuda={CUDA_HOME}"] if get_cuda_version() else []), + "--with-verbs", + "--without-gdrcopy", + "--without-dc", + "--without-rdmacm", + "--without-gga", + "--with-ze=no", + ] + run_command(configure_command, cwd=ucx_source_path) + run_command(["make", "-j", str(MAX_JOBS)], cwd=ucx_source_path) + run_command(["make", "install"], cwd=ucx_source_path) + + +def build_nixl_wheel(ucx_path, build_env, wheel_dir, args): + """Builds the (unrepaired) NIXL wheel against the UCX installed in ucx_path.""" + if not os.path.exists(NIXL_DIR): + run_command(["git", "clone", NIXL_REPO_URL, NIXL_DIR]) + else: + run_command(["git", "fetch", "--tags"], cwd=NIXL_DIR) + nixl_source_path = os.path.abspath(NIXL_DIR) + # Tags are "v" for current releases and bare for older ones. + for tag in (f"v{NIXL_WHEEL_VERSION}", NIXL_WHEEL_VERSION): + tag_exists = subprocess.run( + ["git", "rev-parse", "--verify", "--quiet", f"refs/tags/{tag}"], + cwd=nixl_source_path, + stdout=subprocess.DEVNULL, + ) + if tag_exists.returncode == 0: + run_command(["git", "checkout", "--force", tag], cwd=nixl_source_path) + print(f"--> Checked out NIXL version: {tag}", flush=True) + break + else: + raise RuntimeError(f"NIXL has no tag for version {NIXL_WHEEL_VERSION}.") + + cuda_version = get_cuda_version() + if cuda_version is not None: + # The wheel name selects which of nixl-cu12 / nixl-cu13 is built. + run_command( + [ + sys.executable, + "contrib/tomlutil.py", + "--wheel-name", + f"nixl-cu{cuda_version[0]}", + "pyproject.toml", + ], + cwd=nixl_source_path, + ) + + meson_options = { + "build_tests": "false", + "ucx_path": ucx_path, + } + cuda_arch_list = args.cuda_arch_list + if args.build_nixl_ep: + meson_options["build_nixl_ep"] = "true" + meson_options["build_examples"] = "true" + cuda_arch_list = cuda_arch_list or get_default_cuda_arch_list() + if cuda_arch_list: + meson_options["nixl_cuda_arch_list"] = cuda_arch_list + if args.nixl_plugins: + meson_options["enable_plugins"] = args.nixl_plugins + config_settings = [f"-Csetup-args=-D{k}={v}" for k, v in meson_options.items()] + # Keep the build tree: auditwheel needs the subproject libraries it holds. + config_settings.append(f"-Cbuild-dir={NIXL_BUILD_DIR}") + config_settings.append(f"-Ccompile-args=-j{MAX_JOBS}") + + run_command( + [ + sys.executable, + "-m", + "pip", + "wheel", + ".", + "--no-deps", + # Build against the torch already installed, not a fresh isolated one. + "--no-build-isolation", + f"--wheel-dir={wheel_dir}", + *config_settings, + ], + cwd=nixl_source_path, + env=build_env, + ) + + +def repair_nixl_wheel(unrepaired_wheel, repaired_dir, ucx_install_path, build_env): + """Bundles the shared libraries and UCX modules into a self-contained wheel.""" + auditwheel_command = ["auditwheel", "repair"] + for pattern in AUDITWHEEL_EXCLUDES: + auditwheel_command += ["--exclude", pattern] + auditwheel_command += [unrepaired_wheel, f"--wheel-dir={repaired_dir}"] + # The wheel ships prometheus-cpp as libcore.so.1.3.0 but its users ask for the + # soname libcore.so.1.3, which only exists as a symlink in the build tree. + repair_env = build_env.copy() + prometheus_dir = os.path.join(NIXL_BUILD_DIR, "subprojects", "prometheus-cpp") + repair_env["LD_LIBRARY_PATH"] = f"{prometheus_dir}:{build_env['LD_LIBRARY_PATH']}" + run_command(auditwheel_command, env=repair_env) + + repaired_wheel = find_nixl_wheel_in_cache(repaired_dir) + if not repaired_wheel: + raise RuntimeError("Failed to find the repaired NIXL wheel.") + # UCX loads its transports (CUDA, IB, ...) with dlopen, so auditwheel does not + # see them as dependencies; the NIXL plugin itself is already in the wheel. + run_command( + [ + sys.executable, + "contrib/wheel_add_ucx_plugins.py", + "--skip-nixl-plugins", + f"--ucx-plugins-dir={os.path.join(ucx_install_path, 'lib', 'ucx')}", + repaired_wheel, + ], + cwd=os.path.abspath(NIXL_DIR), + env=build_env, + ) + return repaired_wheel + + +def install_nixl_wheels(wheels): + """Installs the NIXL wheels and drops the nixl-cuXX of other CUDA majors.""" + # --force-reinstall: the wheels usually have the version of the release that + # is already installed, which pip would otherwise keep. + # w/o "no-deps", it will install cuda-torch + run_command( + [ + sys.executable, + "-m", + "pip", + "install", + "--force-reinstall", + "--no-deps", + *wheels, + ] + ) + + cuda_version = get_cuda_version() + if cuda_version is None: + return + # The dispatcher only loads nixl-cu, and the meta + # package built here does not depend on the other variant. + for cuda_major in (12, 13): + package = f"nixl-cu{cuda_major}" + if cuda_major != cuda_version[0] and is_pip_package_installed(package): + run_command([sys.executable, "-m", "pip", "uninstall", "-y", package]) + + +def build_and_install_prerequisites(args): + """Builds UCX and NIXL from source, creating a self-contained wheel.""" + print(f"--> NIXL version {NIXL_WHEEL_VERSION} (from {NIXL_VERSION_SOURCE})") + if args.build_nixl_ep is None: + args.build_nixl_ep = get_cuda_version() is not None and is_torch_available() + reuse_existing = not (args.force_reinstall or args.build_only) + if reuse_existing and is_nixl_installed(args.build_nixl_ep): + print("--> NIXL is already installed. Nothing to do.", flush=True) + return + + cached_wheel = find_nixl_wheel_in_cache(WHEELS_CACHE_HOME) + if reuse_existing and cached_wheel: + print( + f"\n--> Found self-contained wheel: \ + {os.path.basename(cached_wheel)}.", + flush=True, + ) + print("--> Installing from cache, skipping all source builds.", flush=True) + cached_wheels = [cached_wheel] + cached_meta_wheel = find_nixl_meta_wheel(WHEELS_CACHE_HOME) + if cached_meta_wheel: + cached_wheels.append(cached_meta_wheel) + install_nixl_wheels(cached_wheels) + print("\n--- Installation from cache complete. ---", flush=True) + return + + print( + "\n--> No installed package or cached wheel found. \ + Starting full build process...", + flush=True, + ) + if not args.skip_system_deps: + install_system_dependencies(args.build_nixl_ep) + print("\n--> Installing Python build dependencies...", flush=True) + install_python_build_dependencies() + print(f"--> Using wheel cache directory: {WHEELS_CACHE_HOME}", flush=True) + os.makedirs(WHEELS_CACHE_HOME, exist_ok=True) + + # -- Step 1: Build UCX from source, unless an existing install is given -- + # With an external UCX (e.g. HPC-X) the wheel keeps depending on it at runtime. + external_ucx = args.ucx_path is not None + if external_ucx: + ucx_install_path = os.path.abspath(args.ucx_path) + print(f"\n[1/3] Using the existing UCX in {ucx_install_path}", flush=True) + else: + ucx_install_path = os.path.abspath(UCX_INSTALL_DIR) + print("\n[1/3] Configuring and building UCX from source...", flush=True) + build_ucx(ucx_install_path) + print("--- UCX build and install complete ---", flush=True) + + # -- Step 2: Build NIXL wheel from source -- + print("\n[2/3] Building NIXL wheel from source...", flush=True) + ucx_lib_path = os.path.join(ucx_install_path, "lib") + ucx_plugin_path = os.path.join(ucx_lib_path, "ucx") + build_env = os.environ.copy() + # pip-installed build tools live next to the running interpreter. + build_env["PATH"] = os.pathsep.join( + [ + os.path.join(BUILD_TOOLS_DIR, "bin"), + os.path.join(CUDA_HOME, "bin"), + os.path.dirname(sys.executable), + build_env["PATH"], + ] + ) + build_env["PYTHONPATH"] = os.pathsep.join( + filter(None, [BUILD_TOOLS_DIR, build_env.get("PYTHONPATH")]) + ) + build_env["PKG_CONFIG_PATH"] = os.path.join(ucx_lib_path, "pkgconfig") + existing_ld_path = os.environ.get("LD_LIBRARY_PATH", "") + build_env["LD_LIBRARY_PATH"] = ( + f"{ucx_lib_path}:{ucx_plugin_path}:{existing_ld_path}".strip(":") + ) + print(f"--> Using LD_LIBRARY_PATH: {build_env['LD_LIBRARY_PATH']}", flush=True) + + build_nixl_wheel(ucx_install_path, build_env, TEMP_WHEEL_DIR, args) + + # -- Step 3: Repair the wheel by copying UCX libraries -- + unrepaired_wheel = find_nixl_wheel_in_cache(TEMP_WHEEL_DIR) + if not unrepaired_wheel: + raise RuntimeError("Failed to find the NIXL wheel after building it.") + if external_ucx: + print("\n[3/3] Skipping the repair step: UCX is provided externally.") + built_wheels = [unrepaired_wheel] + else: + print("\n[3/3] Repairing NIXL wheel to include UCX libraries...", flush=True) + built_wheels = [ + repair_nixl_wheel( + unrepaired_wheel, + os.path.join(TEMP_WHEEL_DIR, "repaired"), + ucx_install_path, + build_env, + ) + ] + + meta_wheel = find_nixl_meta_wheel( + os.path.join(NIXL_BUILD_DIR, "src", "bindings", "python", "nixl-meta") + ) + if meta_wheel: + built_wheels.append(meta_wheel) + else: + print("WARNING: the nixl meta package was not built.", flush=True) + + new_wheels = [] + for wheel in built_wheels: + run_command(["cp", wheel, WHEELS_CACHE_HOME]) + new_wheels.append(os.path.join(WHEELS_CACHE_HOME, os.path.basename(wheel))) + run_command(["rm", "-rf", TEMP_WHEEL_DIR]) + + wheel_names = ", ".join(os.path.basename(w) for w in new_wheels) + if args.build_only: + print(f"--> Built wheels in {WHEELS_CACHE_HOME}: {wheel_names}", flush=True) + return + print( + f"--> Successfully built wheels: {wheel_names}. Now installing...", flush=True + ) + install_nixl_wheels(new_wheels) + print("--- NIXL installation complete ---", flush=True) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Build and install UCX and NIXL dependencies." + ) + parser.add_argument( + "--force-reinstall", + action="store_true", + help="Force rebuild and reinstall of UCX and NIXL \ + even if they are already installed.", + ) + parser.add_argument( + "--build-only", + action="store_true", + help="Build the wheels into WHEELS_CACHE_HOME and stop, without \ + installing them or looking at what is already installed.", + ) + parser.add_argument( + "--skip-system-deps", + action="store_true", + help="Do not install system packages with apt-get; the build tools \ + (git, autotools, cmake, patchelf, ...) must already be on PATH.", + ) + parser.add_argument( + "--build-nixl-ep", + action=argparse.BooleanOptionalAction, + default=None, + help="Build the NIXL EP extension for the installed torch version. \ + Defaults to on when CUDA and torch are available.", + ) + parser.add_argument( + "--cuda-arch-list", + default=os.environ.get("NIXL_CUDA_ARCH_LIST"), + help="Comma-separated CUDA SM targets for NIXL, e.g. 100,103,107. \ + Defaults to the NIXL EP targets the installed nvcc supports.", + ) + parser.add_argument( + "--nixl-plugins", + default=os.environ.get("NIXL_PLUGINS", "UCX"), + help="Comma-separated NIXL plugins to build. Pass an empty string to \ + build every plugin whose dependencies are available.", + ) + parser.add_argument( + "--ucx-path", + default=os.environ.get("UCX_PATH"), + help="Prefix of an existing UCX to build against instead of building one. \ + The resulting wheel is not self-contained.", + ) + args = parser.parse_args() + build_and_install_prerequisites(args) From d6fe5dca686b3f6c5f018d399f4f119fedf842f1 Mon Sep 17 00:00:00 2001 From: Dhairyashil R G Date: Thu, 8 Oct 2026 04:51:51 +0530 Subject: [PATCH 044/624] [Bugfix][Parser] Emit buffered post-reasoning text when no tool parser is configured (#58911) Signed-off-by: Dhairyashil R G <18738735+dhairyashilRG@users.noreply.github.com> Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: sfeng33 <4florafeng@gmail.com> --- tests/parser/test_streaming.py | 72 ++++++++++++++++++++++++++++++++++ vllm/parser/abstract_parser.py | 13 +++++- 2 files changed, 83 insertions(+), 2 deletions(-) diff --git a/tests/parser/test_streaming.py b/tests/parser/test_streaming.py index cad16746be44..d5d43398039f 100644 --- a/tests/parser/test_streaming.py +++ b/tests/parser/test_streaming.py @@ -663,6 +663,78 @@ def test_engine_reasoning_no_tool_batched_content_passthrough(tokenizer, request assert len(tool_calls) == 0 +@pytest.mark.parametrize( + "text_chunks", + [ + # A trailing "<" that the engine holds back (it may start a tag). + ["let me thinka <", "x> b"], + # Whitespace after "" in the same delta as the marker. + ["let me think\n\n", "bold done"], + # The delta that ends reasoning carries only held-back text. + ["let me think", "<", "b>bold done"], + ], +) +def test_engine_reasoning_no_tool_batched_holdback_not_dropped( + tokenizer, request_obj, text_chunks +): + """Text the engine buffers when reasoning ends inside a batched delta is + returned by finish_streaming(); with no tool parser it must still be + emitted as content, not dropped.""" + parser = Qwen3ReasoningNoToolParser(tokenizer) + chunks = [tokenizer.encode(t, add_special_tokens=False) for t in text_chunks] + + results = stream_chunks(parser, tokenizer, chunks, request_obj) + reasoning, content, tool_calls = collect_fields(results) + + assert reasoning == "let me think" + assert content == "".join(text_chunks).split("", 1)[1] + assert len(tool_calls) == 0 + + +@pytest.mark.parametrize( + ("text_chunks", "include_reasoning", "expected_content"), + [ + # The delta that ends reasoning is also the last one: the held-back + # "<" must be emitted exactly once (the end-of-stream flush skips the + # reasoning parser once reasoning has ended). + (["let me thinka <"], True, "a <"), + # Whitespace-only content after the marker at the end of the stream. + (["let me think\n\n"], True, "\n\n"), + # Hiding reasoning must not hide the flushed content. + (["let me think\n\n", "bold"], False, "\n\nbold"), + ], +) +def test_engine_reasoning_no_tool_holdback_finished_and_hidden_reasoning( + tokenizer, text_chunks, include_reasoning, expected_content +): + request = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + include_reasoning=include_reasoning, + ) + parser = Qwen3ReasoningNoToolParser(tokenizer) + chunks = [tokenizer.encode(t, add_special_tokens=False) for t in text_chunks] + + results: list[DeltaMessage | None] = [] + prompt_token_ids: list[int] | None = [] + for i, chunk in enumerate(chunks): + results.append( + parser.parse_delta( + tokenizer.decode(chunk), + chunk, + request, + prompt_token_ids=prompt_token_ids, + finished=i == len(chunks) - 1, + ) + ) + prompt_token_ids = None + reasoning, content, tool_calls = collect_fields(results) + + assert reasoning == ("let me think" if include_reasoning else "") + assert content == expected_content + assert len(tool_calls) == 0 + + def _decode_stream_deltas(tokenizer, groups): """Decode token-ID groups into ``(delta_text, group)`` pairs via the real incremental ``DecodeStream``. diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index c28563fabf65..f75d7ce9028d 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -904,8 +904,17 @@ def parse_delta( (delta_message.content if delta_message else None) or "" ) + ((flush_delta.content if flush_delta else None) or "") if self._engine_based: - if delta_message and self._tool_parser is not None: - delta_message.content = None + if self._tool_parser is not None: + # The tool-call phase below consumes current_text. + if delta_message: + delta_message.content = None + elif current_text: + # No tool phase follows to consume the flushed text, + # and the engine is not fed again after the + # transition, so emit it as content here. + if delta_message is None: + delta_message = DeltaMessage() + delta_message.content = current_text else: delta_text = current_text From 38dc8ee50055622ab0d5fe6bb32a5aee526bcbf5 Mon Sep 17 00:00:00 2001 From: www6v Date: Thu, 8 Oct 2026 07:46:41 +0800 Subject: [PATCH 045/624] [Bugfix] Validate tool-parser/tokenizer compatibility at startup (#59749) Signed-off-by: www6v Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: sfeng33 <4florafeng@gmail.com> --- vllm/parser/parser_manager.py | 58 +++++++++++++++++++++++++++++++++-- 1 file changed, 56 insertions(+), 2 deletions(-) diff --git a/vllm/parser/parser_manager.py b/vllm/parser/parser_manager.py index 719b396130af..cf0be908ed63 100644 --- a/vllm/parser/parser_manager.py +++ b/vllm/parser/parser_manager.py @@ -73,6 +73,29 @@ def get_reasoning_parser( raise TypeError(f"{reasoning_parser_name=} has not been registered") from e return parser + @classmethod + def validate_parser_tokenizer( + cls, + parser_cls: type[Parser], + tokenizer: TokenizerLike, + tool_parser_name: str | None = None, + reasoning_parser_name: str | None = None, + model_name: str | None = None, + ) -> None: + """Fail at startup if the tokenizer lacks what the parsers need.""" + try: + parser_cls(tokenizer) + except Exception as e: + flags = [] + if parser_cls.tool_parser_cls is not None: + flags.append(f"--tool-call-parser {tool_parser_name}") + if parser_cls.reasoning_parser_cls is not None: + flags.append(f"--reasoning-parser {reasoning_parser_name}") + raise TypeError( + f"{' and '.join(flags)} cannot be used with the tokenizer of " + f"{model_name!r}: {e}" + ) from e + @classmethod def get_parser( cls, @@ -97,13 +120,44 @@ def get_parser( If True, HarmonyParser is always returned. tool_strict_level: Server-side floor for tool-call structural tags (``--tool-strict-level``). - tokenizer: Tokenizer whose `response_template` metadata is - validated when the `hf` parser is selected. + tokenizer: Tokenizer the composed parser is validated against + at startup (see :meth:`validate_parser_tokenizer`). ``None`` + skips the check (e.g. ``skip_tokenizer_init``). Returns: A Parser class, or None if neither parser is specified. """ + parser_cls = cls._compose_parser( + tool_parser_name=tool_parser_name, + reasoning_parser_name=reasoning_parser_name, + enable_auto_tools=enable_auto_tools, + model_name=model_name, + is_harmony=is_harmony, + tool_strict_level=tool_strict_level, + tokenizer=tokenizer, + ) + if parser_cls is not None and tokenizer is not None: + cls.validate_parser_tokenizer( + parser_cls, + tokenizer, + tool_parser_name=tool_parser_name, + reasoning_parser_name=reasoning_parser_name, + model_name=model_name, + ) + return parser_cls + + @classmethod + def _compose_parser( + cls, + tool_parser_name: str | None, + reasoning_parser_name: str | None, + enable_auto_tools: bool, + model_name: str | None, + is_harmony: bool, + tool_strict_level: str, + tokenizer: TokenizerLike | None, + ) -> type[Parser] | None: if not tool_parser_name and not reasoning_parser_name: return None From ed825aafc0afb36929e837035dd70760c9bc6953 Mon Sep 17 00:00:00 2001 From: YueshenZ <844936239@qq.com> Date: Thu, 8 Oct 2026 08:07:52 +0800 Subject: [PATCH 046/624] [Bugfix][Frontend] Fix Responses namespace aliases in named tool calling (#59723) Signed-off-by: Shen <312058264+YueshenZ@users.noreply.github.com> Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: sfeng33 <4florafeng@gmail.com> --- docs/features/tool_calling.md | 2 + tests/tool_use/test_tool_choice_required.py | 74 ++++++++++++++++++- vllm/entrypoints/openai/responses/protocol.py | 70 ++++++++++++------ vllm/tool_parsers/utils.py | 2 - 4 files changed, 124 insertions(+), 24 deletions(-) diff --git a/docs/features/tool_calling.md b/docs/features/tool_calling.md index 8dccd3a8b857..3081047dcec4 100644 --- a/docs/features/tool_calling.md +++ b/docs/features/tool_calling.md @@ -94,6 +94,8 @@ For best results, we recommend ensuring that the expected output format / schema To use a named function, you need to define the functions in the `tools` parameter of the chat completion request, and specify the `name` of one of the tools in the `tool_choice` parameter of the chat completion request. +In the Responses API, `tool_choice={"type": "function", "name": "..."}` can select a function inside a namespace by its full `namespace__function` name or by its local name. Names are matched literally first, so a plain function named `other__add` never stands in for `add`. A local name defined in more than one namespace must be selected by its full name. + ## Required Function Calling vLLM supports the `tool_choice='required'` option in the chat completion API. Similar to the named function calling, it also uses structured outputs, so this is enabled by default and will work with any supported model. However, support for alternative decoding backends are on the [roadmap](../usage/v1_guide.md#features) for the V1 engine. diff --git a/tests/tool_use/test_tool_choice_required.py b/tests/tool_use/test_tool_choice_required.py index 26978fbb7c0e..ebdea5aa0b9f 100644 --- a/tests/tool_use/test_tool_choice_required.py +++ b/tests/tool_use/test_tool_choice_required.py @@ -5,7 +5,11 @@ import pytest import regex as re -from openai.types.responses import FunctionTool, ToolChoiceFunction, WebSearchTool +from openai.types.responses import ( + FunctionTool, + ToolChoiceFunction, + WebSearchTool, +) from pydantic import TypeAdapter from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -13,6 +17,7 @@ ChatCompletionNamedToolChoiceParam, ChatCompletionToolsParam, ) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.exceptions import VLLMValidationError from vllm.tool_parsers.streaming import extract_required_tool_call_streaming from vllm.tool_parsers.utils import ( @@ -581,3 +586,70 @@ def test_responses_empty_params_constrains_object(self, params): choice = ToolChoiceFunction(type="function", name="ping") schema = get_json_schema_from_tools(choice, [tool]) assert schema == {"type": "object", "properties": {}} + + +class TestForcedFunctionShortNameAlias: + """Forced choices must use the selected function's parameter schema.""" + + PARAMS = {"type": "object", "properties": {"x": {"type": "integer"}}} + + @pytest.mark.parametrize("plain_first", [True, False]) + @pytest.mark.parametrize( + "namespace,function_name,choice_name", + [ + ("math", "add", "add"), + ("math", "add", "math__add"), + ("math", "vector__add", "vector__add"), + ("math__vector", "vector__add", "vector__add"), + ], + ) + def test_namespace_schema_is_not_shadowed( + self, plain_first, namespace, function_name, choice_name + ): + tools = [ + { + "type": "function", + "name": "other__add", + "parameters": { + "type": "object", + "properties": {"y": {"type": "string"}}, + }, + }, + { + "type": "namespace", + "name": namespace, + "description": "math tools", + "tools": [ + { + "type": "function", + "name": function_name, + "parameters": self.PARAMS, + } + ], + }, + ] + if not plain_first: + tools.reverse() + request = ResponsesRequest.model_validate( + { + "model": "test-model", + "input": "Add the numbers.", + "tools": tools, + "tool_choice": {"type": "function", "name": choice_name}, + } + ) + assert ( + get_json_schema_from_tools(request.tool_choice, request.tools) + == self.PARAMS + ) + + def test_plain_function_double_underscore_name_has_no_phantom_alias(self): + tool = FunctionTool(type="function", name="math__add", parameters=self.PARAMS) + choice = ToolChoiceFunction(type="function", name="add") + with pytest.raises(ValueError, match="has not been passed in `tools`"): + get_json_schema_from_tools(choice, [tool]) + + def test_plain_function_full_name_still_resolves(self): + tool = FunctionTool(type="function", name="math__add", parameters=self.PARAMS) + choice = ToolChoiceFunction(type="function", name="math__add") + assert get_json_schema_from_tools(choice, [tool]) == self.PARAMS diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 5d4a7411ac21..35e6d3aa3dcd 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -159,6 +159,51 @@ def _default_input_image_details(value: Any) -> Any: return {**value, "content": new_content} if changed else value +def _tool_field(tool: Any, field: str) -> Any: + return tool.get(field) if isinstance(tool, dict) else getattr(tool, field, None) + + +def _resolve_named_tool_choice(tool_name: str, tools: list[Any]) -> str: + """Resolve a named tool choice to the flat ``__`` form.""" + from vllm.tool_parsers.utils import flat_namespace_tool_name + + exact_names: set[str] = set() + local_names: dict[str, list[str]] = {} + for tool in tools: + if _tool_field(tool, "type") == "namespace": + namespace = _tool_field(tool, "name") + namespaced_tools = _tool_field(tool, "tools") + if not isinstance(namespaced_tools, list) or not isinstance(namespace, str): + return tool_name + for namespaced_tool in namespaced_tools: + local_name = _tool_field(namespaced_tool, "name") + if not isinstance(local_name, str): + continue + flat_name = flat_namespace_tool_name(namespace, local_name) + exact_names.add(flat_name) + local_names.setdefault(local_name, []).append(flat_name) + else: + name = _tool_field(tool, "name") + if isinstance(name, str): + exact_names.add(name) + if tool_name in exact_names: + return tool_name + candidates = local_names.get(tool_name, []) + if len(candidates) == 1: + return candidates[0] + if len(candidates) > 1: + raise VLLMValidationError( + f"Tool choice '{tool_name}' is ambiguous: it is defined in more " + f"than one namespace ({', '.join(candidates)}). Select it by its " + "full '__' name.", + parameter="tool_choice", + ) + raise VLLMValidationError( + "Tool choice 'function' not found in 'tools' parameter.", + parameter="tool_choice", + ) + + class ResponsesRequest(OpenAIBaseModel): # Ordered by official OpenAI API documentation # https://platform.openai.com/docs/api-reference/responses/create @@ -662,31 +707,14 @@ def check_tool_usage(cls, data): ) elif is_named_tool_choice and tools is not None: tool_name = tool_choice.get("name") - tool_names = set() - for tool in tools: - if isinstance(tool, dict): - if tool.get("type") == "namespace": - namespace = tool.get("name") - namespaced_tools = tool.get("tools") - if not isinstance(namespaced_tools, list): - return data - for namespaced_tool in namespaced_tools: - namespaced_name = ( - namespaced_tool.get("name") - if isinstance(namespaced_tool, dict) - else getattr(namespaced_tool, "name", None) - ) - tool_names.add(namespaced_name) - tool_names.add(f"{namespace}__{namespaced_name}") - else: - tool_names.add(tool.get("name")) - else: - tool_names.add(getattr(tool, "name", None)) - if not tool_name or tool_name not in tool_names: + if not tool_name: raise VLLMValidationError( "Tool choice 'function' not found in 'tools' parameter.", parameter="tool_choice", ) + resolved = _resolve_named_tool_choice(tool_name, tools) + if resolved != tool_name: + data["tool_choice"] = {**tool_choice, "name": resolved} return data diff --git a/vllm/tool_parsers/utils.py b/vllm/tool_parsers/utils.py index 86d68a9b5130..7a93c843b8e3 100644 --- a/vllm/tool_parsers/utils.py +++ b/vllm/tool_parsers/utils.py @@ -432,8 +432,6 @@ def get_json_schema_from_tools( continue for name, params in iter_response_function_tool_info(tool): responses_tool_map[name] = params - if "__" in name: - responses_tool_map.setdefault(name.rsplit("__", 1)[1], params) if tool_name not in responses_tool_map: raise ValueError(f"Tool '{tool_name}' has not been passed in `tools`.") return _params_or_empty_object(responses_tool_map[tool_name]) From 8faac3dd2af39fd631f018277387f01c89ad5782 Mon Sep 17 00:00:00 2001 From: v-shobhit <161510941+v-shobhit@users.noreply.github.com> Date: Wed, 7 Oct 2026 17:10:32 -0700 Subject: [PATCH 047/624] [Bugfix] [Spec Decode] DeepSeek V4/V4.1 NVFP4: keep ignored MTP/DSpark experts on MXFP4 (#59655) Signed-off-by: Shobhit Verma Signed-off-by: v-shobhit <161510941+v-shobhit@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Co-authored-by: Benjamin Chislett --- .buildkite/test_areas/lm_eval.yaml | 2 + .buildkite/test_areas/quantization.yaml | 1 + .../DeepSeek-V4.1-Flash-NVFP4-DSpark-TP4.yaml | 15 ++++ .../gsm8k/configs/models-spec-decode.txt | 1 + tests/quantization/test_fp8.py | 76 +++++++++++++++++++ vllm/models/deepseek_v41/quant_config.py | 32 +++++++- 6 files changed, 125 insertions(+), 2 deletions(-) create mode 100644 tests/evals/gsm8k/configs/DeepSeek-V4.1-Flash-NVFP4-DSpark-TP4.yaml diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index bcb111084843..d1a115a1e1b0 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -322,9 +322,11 @@ steps: optional: true source_file_dependencies: - tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-confidence-TEP4.yaml + - tests/evals/gsm8k/configs/DeepSeek-V4.1-Flash-NVFP4-DSpark-TP4.yaml - tests/evals/gsm8k/configs/Kimi-K3-pruned75-DSpark-TP4.yaml - tests/evals/gsm8k/configs/models-spec-decode.txt - vllm/models/deepseek_v4/ + - vllm/models/deepseek_v41/ - vllm/models/kimi_k3/ - vllm/model_executor/models/qwen3_dspark.py - vllm/v1/spec_decode/metrics.py diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index e73054a4bcaf..b49b4b92c736 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -12,6 +12,7 @@ steps: source_file_dependencies: - csrc/ - vllm/model_executor/layers/quantization + - vllm/models/deepseek_v41/quant_config.py - tests/quantization commands: # Pin torchao to the version compatible with the CI PyTorch and CUDA stack. diff --git a/tests/evals/gsm8k/configs/DeepSeek-V4.1-Flash-NVFP4-DSpark-TP4.yaml b/tests/evals/gsm8k/configs/DeepSeek-V4.1-Flash-NVFP4-DSpark-TP4.yaml new file mode 100644 index 000000000000..78d899569644 --- /dev/null +++ b/tests/evals/gsm8k/configs/DeepSeek-V4.1-Flash-NVFP4-DSpark-TP4.yaml @@ -0,0 +1,15 @@ +model_name: "nvidia/DeepSeek-V4.1-Flash-NVFP4" +accuracy_threshold: 0.88 +tolerance: 0.03 +# The bundled DSpark draft (mtp.*) stays MXFP4 in this checkpoint. Loading it +# as NVFP4 keeps accuracy but drops acceptance length to 1.0 (3.57 when correct). +min_acceptance_length: 2.5 +num_questions: 1319 +num_fewshot: 5 +startup_max_wait_seconds: 1800 +server_args: >- + --tokenizer-mode deepseek_v41 + --trust-remote-code + --max-model-len 8192 + --tensor-parallel-size 4 + --speculative-config '{"method":"dspark","num_speculative_tokens":5}' diff --git a/tests/evals/gsm8k/configs/models-spec-decode.txt b/tests/evals/gsm8k/configs/models-spec-decode.txt index 427c35d912ab..96cead25858c 100644 --- a/tests/evals/gsm8k/configs/models-spec-decode.txt +++ b/tests/evals/gsm8k/configs/models-spec-decode.txt @@ -1,2 +1,3 @@ DeepSeek-V4-Flash-DSpark-confidence-TEP4.yaml Kimi-K3-pruned75-DSpark-TP4.yaml +DeepSeek-V4.1-Flash-NVFP4-DSpark-TP4.yaml diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py index 1cb53f223929..64341843d951 100644 --- a/tests/quantization/test_fp8.py +++ b/tests/quantization/test_fp8.py @@ -318,6 +318,82 @@ def test_deepseek_v41_declines_quark_configs(): ) +_DSV41_NVFP4 = { + "moe_quant_algo": "NVFP4", + "ignore": ["*.attn.*", "*.ffn.shared_experts.*", "head", "mtp.*"], +} +_DSV4_NVFP4_DSPARK = { + "moe_quant_algo": "NVFP4", + "ignore": ["*.attn.*", "*.ffn.shared_experts.*", "head"], +} + + +@pytest.mark.parametrize( + ("ckpt_quant_cfg", "prefix", "expected"), + [ + pytest.param(_DSV41_NVFP4, "model.layers.1.ffn.experts", "nvfp4", id="target"), + pytest.param( + _DSV41_NVFP4, + "language_model.model.layers.0.ffn.experts", + "nvfp4", + id="target-vl-prefix", + ), + pytest.param( + _DSV41_NVFP4, "model.layers.40.ffn.experts", "mxfp4", id="ignored-mtp0" + ), + pytest.param( + _DSV41_NVFP4, "model.layers.42.ffn.experts", "mxfp4", id="ignored-mtp2" + ), + pytest.param( + _DSV4_NVFP4_DSPARK, + "model.layers.41.ffn.experts", + "nvfp4", + id="quantized-mtp1", + ), + pytest.param( + {"moe_quant_algo": "NVFP4"}, + "model.layers.40.ffn.experts", + "nvfp4", + id="no-ignore-list", + ), + pytest.param({}, "model.layers.1.ffn.experts", "mxfp4", id="not-nvfp4"), + ], +) +def test_deepseek_v41_nvfp4_export_keeps_ignored_experts_mxfp4( + monkeypatch, ckpt_quant_cfg, prefix, expected +): + """NVFP4 exports keep the routed experts they ``ignore`` in MXFP4, e.g. the + bundled DSpark draft (``mtp.*``, built as ``layers.{40 + k}``) of + DeepSeek-V4.1-Flash-NVFP4. The ignore list is read from the checkpoint + passed to ``from_config``: the current hf_config below carries none.""" + import vllm.model_executor.layers.quantization.modelopt as modelopt + import vllm.models.deepseek_v41.quant_config as dsv41_qc + from vllm.model_executor.layers.fused_moe import RoutedExperts + + hf_config = SimpleNamespace( + expert_dtype="fp4", + quantization_config={"moe_quant_algo": ckpt_quant_cfg.get("moe_quant_algo")}, + ) + model_config = SimpleNamespace( + hf_config=hf_config, hf_text_config=SimpleNamespace(num_hidden_layers=40) + ) + monkeypatch.setattr( + dsv41_qc, + "get_current_vllm_config", + lambda: SimpleNamespace(model_config=model_config), + ) + monkeypatch.setattr(dsv41_qc, "Mxfp4MoEMethod", lambda moe_config: "mxfp4") + monkeypatch.setattr(modelopt, "ModelOptNvFp4FusedMoE", lambda **kwargs: "nvfp4") + + config = dsv41_qc.DeepseekV4FP8Config.from_config( + {"quant_method": "fp8", "activation_scheme": "dynamic", **ckpt_quant_cfg} + ) + monkeypatch.setattr(config, "_get_nvfp4_config", lambda: None) + layer = RoutedExperts.__new__(RoutedExperts) + layer.moe_config = None + assert config.get_quant_method(layer, prefix) == expected + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="DeepGEMM requires CUDA") @pytest.mark.parametrize("scale_dtype", [torch.uint8, torch.float8_e8m0fnu]) @pytest.mark.parametrize( diff --git a/vllm/models/deepseek_v41/quant_config.py b/vllm/models/deepseek_v41/quant_config.py index 384e9488793c..1d92943db425 100644 --- a/vllm/models/deepseek_v41/quant_config.py +++ b/vllm/models/deepseek_v41/quant_config.py @@ -4,8 +4,11 @@ from __future__ import annotations +import fnmatch from typing import TYPE_CHECKING, cast +import regex as re + from vllm.config import get_current_vllm_config from vllm.config.quantization import QuantSpec from vllm.model_executor.layers.fused_moe import ( @@ -23,6 +26,7 @@ ) _DEEPSEEK_V4_EXPERT_DTYPES = ("fp4", "fp8") +_EXPERTS_LAYER_RE = re.compile(r"(?:^|\.)layers\.(\d+)\.ffn\.experts$") if TYPE_CHECKING: from vllm.model_executor.layers.quantization.modelopt import ( @@ -54,6 +58,7 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._resolved_expert_dtype: str | None = None self._resolved_moe_quant_algo: str | None = None + self._modelopt_ignore: tuple[str, ...] = () self._nvfp4_config: ModelOptNvFp4Config | None = None # ``is_scale_e8m0`` is a property that resolves on first read, # by which time the current vllm_config has been set. @@ -103,6 +108,21 @@ def moe_quant_algo(self) -> str: self._resolve_moe_overrides() return self._resolved_moe_quant_algo or "" + def _is_nvfp4_expert_layer(self, prefix: str) -> bool: + if self.moe_quant_algo != "NVFP4": + return False + match = _EXPERTS_LAYER_RE.search(prefix) + if match is None or not self._modelopt_ignore: + return True + # MTP/draft layers are built as layers.{num_hidden_layers + k} but are + # named mtp.{k} in some checkpoints. + model_config = get_current_vllm_config().model_config + num_layers = model_config.hf_text_config.num_hidden_layers + idx = int(match.group(1)) + layer = f"mtp.{idx - num_layers}" if idx >= num_layers else f"layers.{idx}" + name = f"{layer}.ffn.experts" + return not any(fnmatch.fnmatchcase(name, p) for p in self._modelopt_ignore) + def _get_nvfp4_config(self) -> ModelOptNvFp4Config: if self._nvfp4_config is None: from vllm.model_executor.layers.quantization.modelopt import ( @@ -173,6 +193,9 @@ def override_quantization_method( @classmethod def from_config(cls, config: dict) -> DeepseekV4FP8Config: + # ModelOpt NVFP4 exports keep the experts they ``ignore`` in the native + # MXFP4 format, e.g. the bundled MTP/DSpark layers (``mtp.*``). + ignore = config.get("ignore") # Reroute AMD-Quark fused shared expert MXFP4 checkpoints onto the fp8 # path: the runtime layout matches the DeepSeek-native fp8 checkpoint, # so translate the schema into format Fp8Config.from_config expects. @@ -188,7 +211,12 @@ def from_config(cls, config: dict) -> DeepseekV4FP8Config: name for name in quark_exclude if isinstance(name, str) ], } - return cast("DeepseekV4FP8Config", super().from_config(config)) + quant_config = cast("DeepseekV4FP8Config", super().from_config(config)) + if isinstance(ignore, list): + quant_config._modelopt_ignore = tuple( + p for p in ignore if isinstance(p, str) + ) + return quant_config def get_quant_method(self, layer, prefix): if ( @@ -221,7 +249,7 @@ def get_quant_method(self, layer, prefix): ): return UnquantizedFusedMoEMethod(layer.moe_config) if self.expert_dtype == "fp4": - if self.moe_quant_algo == "NVFP4": + if self._is_nvfp4_expert_layer(prefix): from vllm.model_executor.layers.quantization.modelopt import ( ModelOptNvFp4FusedMoE, ) From a58bdd0d4ec62764a8d6c5d4f9ed6a1dd4bf9b80 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Thu, 8 Oct 2026 02:11:59 +0200 Subject: [PATCH 048/624] [Perf] JIT monitor: hook TileLang JITImpl.compile instead of __call__ (#58685) Signed-off-by: NickLucche --- tests/jit_monitor/test_hooks.py | 5 ++++- vllm/utils/jit_monitor.py | 18 ++++++++---------- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/tests/jit_monitor/test_hooks.py b/tests/jit_monitor/test_hooks.py index 452fe9699504..6d4f25314c08 100644 --- a/tests/jit_monitor/test_hooks.py +++ b/tests/jit_monitor/test_hooks.py @@ -65,10 +65,13 @@ def __call__(self, *args, **kwargs): key, _ = self.func.parse_args(*args, **kwargs) kernel = self._kernel_cache.get(key) if kernel is None: - kernel = "compiled" + kernel = self.compile(*args, **kwargs) self._kernel_cache[key] = kernel return kernel + def compile(self, *args, **kwargs): + return "compiled" + fake_kernel = cast(Any, ModuleType("tilelang.jit.kernel")) fake_kernel.JITKernel = FakeJITKernel diff --git a/vllm/utils/jit_monitor.py b/vllm/utils/jit_monitor.py index ffb63bc493d4..387189c793fb 100644 --- a/vllm/utils/jit_monitor.py +++ b/vllm/utils/jit_monitor.py @@ -499,32 +499,30 @@ def _init_with_monitor(self, *args, **kwargs): jit_kernel_cls.__init__ = _init_with_monitor if jit_impl_cls is not None: - original_call = jit_impl_cls.__call__ + # JITImpl.__call__ only calls compile() on a kernel-cache miss, so + # hooking compile() keeps the per-call hot path untouched. + original_compile = jit_impl_cls.compile - @functools.wraps(original_call) - def _call_with_monitor(self, *args, **kwargs): + @functools.wraps(original_compile) + def _compile_with_monitor(self, *args, **kwargs): global _tilelang_jitimpl_compile_depth - cache_key = _tilelang_cache_miss_key(self, args, kwargs) - if cache_key is None: - return original_call(self, *args, **kwargs) - _tilelang_jitimpl_compile_depth += 1 try: detail = None if _verbose: detail = _format_verbose_tilelang_compile_details( - self, args, kwargs, cache_key + self, args, kwargs, _tilelang_cache_miss_key(self, args, kwargs) ) func = getattr(self, "func", None) orig_func = getattr(func, "orig_func", None) _log_tilelang_jit_compile( _tilelang_kernel_name(orig_func or func), detail ) - return original_call(self, *args, **kwargs) + return original_compile(self, *args, **kwargs) finally: _tilelang_jitimpl_compile_depth -= 1 - jit_impl_cls.__call__ = _call_with_monitor + jit_impl_cls.compile = _compile_with_monitor _tilelang_hook_installed = True From 1283362e7618c79dedb5a7bdd9308d3f0e0d389c Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Wed, 7 Oct 2026 17:17:56 -0700 Subject: [PATCH 049/624] [CI] Add headroom and selection coverage for H200 fast lanes (#56775) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: Sherlock --- .buildkite/test_areas/benchmarks.yaml | 2 +- .buildkite/test_areas/cuda.yaml | 7 ++++++- .buildkite/test_areas/misc.yaml | 4 ++-- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/.buildkite/test_areas/benchmarks.yaml b/.buildkite/test_areas/benchmarks.yaml index 4b681cd589a9..40f8f9796324 100644 --- a/.buildkite/test_areas/benchmarks.yaml +++ b/.buildkite/test_areas/benchmarks.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: ":nvidia: (H200 MIG 18GB) Benchmarks CLI" key: benchmarks-cli-test - timeout_in_minutes: 45 + timeout_in_minutes: 55 device: h200_18gb source_file_dependencies: - vllm/ diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml index 02ac11879fcb..45b3dc768e6d 100644 --- a/.buildkite/test_areas/cuda.yaml +++ b/.buildkite/test_areas/cuda.yaml @@ -33,14 +33,19 @@ steps: - label: ":nvidia: (H200 MIG 18GB) CUDAGraph" device: h200_18gb key: cudagraph - timeout_in_minutes: 30 + timeout_in_minutes: 45 source_file_dependencies: - tests/v1/cudagraph - vllm/v1/cudagraph_dispatcher.py - vllm/config/compilation.py + - vllm/config/vllm.py - vllm/compilation + - vllm/v1/worker/gpu/cudagraph_utils.py + - vllm/v1/worker/gpu/dp_utils.py - vllm/v1/worker/encoder_cudagraph.py - vllm/v1/worker/encoder_cudagraph_defs.py + - vllm/v1/worker/gpu/model_runner.py + - vllm/v1/worker/gpu/ubatch_utils.py commands: - pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py - pytest -v -s v1/cudagraph/test_cudagraph_mode.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index cbcfaeaea7ca..e1490139c052 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: ":nvidia: (H200 MIG 18GB) V1 Sample" key: v1-sample - timeout_in_minutes: 35 + timeout_in_minutes: 45 device: h200_18gb source_file_dependencies: &v1-sample-logits-deps - vllm/config/ @@ -38,7 +38,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) V1 Logits + Oracle" key: v1-logits-oracle - timeout_in_minutes: 35 + timeout_in_minutes: 45 device: h200_18gb source_file_dependencies: *v1-sample-logits-deps commands: From 1995c0fcd02f88d5daa3fdbcd47e0f016e7b89b9 Mon Sep 17 00:00:00 2001 From: Kang Nahoon <49896157+ejrtks1020@users.noreply.github.com> Date: Thu, 8 Oct 2026 09:24:20 +0900 Subject: [PATCH 050/624] [Bugfix][Frontend] Apply output_config and thinking in Anthropic count_tokens (#59180) Signed-off-by: Kang Nahoon Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../test_anthropic_messages_conversion.py | 26 +++++++++++++++++++ vllm/entrypoints/anthropic/protocol.py | 2 ++ vllm/entrypoints/anthropic/serving.py | 8 +++--- 3 files changed, 32 insertions(+), 4 deletions(-) diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index 1bccee46016e..399c898d0b01 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -1987,6 +1987,32 @@ def test_claude_code_payload(self): assert result.include_reasoning is True assert result.thinking_token_budget is None + @pytest.mark.parametrize( + "config", + [ + pytest.param({"output_config": {"effort": "low"}}, id="effort"), + pytest.param({"thinking": {"type": "disabled"}}, id="disabled"), + pytest.param( + {"thinking": {"type": "enabled", "budget_tokens": 2048}}, + id="enabled", + ), + ], + ) + def test_count_tokens_matches_messages(self, config): + """Effort and thinking change the rendered prompt, so count_tokens must + apply them the same way /v1/messages does.""" + messages = [{"role": "user", "content": "Hello"}] + expected = _convert( + AnthropicMessagesRequest( + model="test-model", max_tokens=4096, messages=messages, **config + ) + ) + result = _convert( + AnthropicCountTokensRequest(model="test-model", messages=messages, **config) + ) + assert result.reasoning_effort == expected.reasoning_effort + assert result.thinking_token_budget == expected.thinking_token_budget + class TestProbeDisabledThinkingEffort: """``auto`` falls back to ``low`` when ``none`` cannot turn thinking off.""" diff --git a/vllm/entrypoints/anthropic/protocol.py b/vllm/entrypoints/anthropic/protocol.py index 2a14c6293327..5997d2e3d2e6 100644 --- a/vllm/entrypoints/anthropic/protocol.py +++ b/vllm/entrypoints/anthropic/protocol.py @@ -385,7 +385,9 @@ class AnthropicCountTokensRequest(BaseModel): model: str messages: list[AnthropicMessage] + output_config: AnthropicOutputConfig | None = None system: str | list[AnthropicContentBlock] | None = None + thinking: AnthropicThinkingConfig | None = None tool_choice: AnthropicToolChoice | None = None tools: list[AnthropicTool] | None = None diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index 05c1e24f2a48..f11697fa165e 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -633,8 +633,6 @@ def _handle_thinking( ``display`` is intentionally ignored: suppressing reasoning would mark it ended for structured outputs and drop it from multi-turn history. """ - if isinstance(anthropic_request, AnthropicCountTokensRequest): - return thinking: AnthropicThinkingConfig | None = anthropic_request.thinking if thinking is None: return @@ -655,8 +653,6 @@ def _handle_output_config( anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest, ) -> None: """Handle output configuration such as output format and effort.""" - if isinstance(anthropic_request, AnthropicCountTokensRequest): - return output_config: AnthropicOutputConfig | None = anthropic_request.output_config if output_config and output_config.format and output_config.format.json_schema: req.response_format = ResponseFormat( @@ -1238,9 +1234,13 @@ async def count_tokens( raw_request: Request | None = None, ) -> AnthropicCountTokensResponse | ErrorResponse: """Implements Anthropic's messages.count_tokens endpoint.""" + disabled_thinking_effort: AnthropicDisabledThinkingEffort = "none" + if request.thinking is not None and request.thinking.type == "disabled": + disabled_thinking_effort = await self._get_disabled_thinking_effort() chat_req = self.to_chat_completion_request( request, merge_inline_system=self._merge_inline_system, + disabled_thinking_effort=disabled_thinking_effort, ) result = await self.render_chat_request(chat_req) if isinstance(result, ErrorResponse): From 9e3e37cb3c041fa3cb7dc578bca1fa0a004b715a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Grunwald?= Date: Thu, 8 Oct 2026 02:56:28 +0200 Subject: [PATCH 051/624] [TEST][XPU][CI] disable tests of unsupported quantization (#60383) Signed-off-by: microslaw Co-authored-by: microslaw Co-authored-by: Kunshang Ji --- .buildkite/intel_jobs/test-intel.yaml | 2 +- .buildkite/scripts/hardware_ci/run-intel-ci-test.sh | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 08f8dcdeab89..b1030e372c41 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -153,7 +153,7 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pytest -v -s quantization/test_auto_round.py && - pytest -v -s quantization/test_online.py --deselect=tests/quantization/test_online.py::test_online_quantization_loads_real_weights' + pytest -v -s quantization/test_online.py --deselect=tests/quantization/test_online.py::test_online_quantization_loads_real_weights --deselect="tests/quantization/test_online.py::test_online_quantization[False-requantization_mxfp8_ptcp_fp8]"' - label: "XPU compressed tensors FP8 test" key: xpu-compressed-tensors-fp8-test depends_on: diff --git a/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh b/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh index c679d0aff242..44f0d9944adc 100644 --- a/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh +++ b/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh @@ -57,7 +57,7 @@ case "${test_suite}" in cd tests pytest -v -s quantization/test_auto_round.py - pytest -v -s quantization/test_online.py --deselect=tests/quantization/test_online.py::test_online_quantization_loads_real_weights + pytest -v -s quantization/test_online.py --deselect=tests/quantization/test_online.py::test_online_quantization_loads_real_weights --deselect="tests/quantization/test_online.py::test_online_quantization[False-requantization_mxfp8_ptcp_fp8]" ;; compressed-tensors-fp8) cd tests From 2dd887426e28dd7dfc96465acc45e06e849f3d44 Mon Sep 17 00:00:00 2001 From: dingdangmao <2756823972@qq.com> Date: Thu, 8 Oct 2026 09:03:08 +0800 Subject: [PATCH 052/624] [Bugfix][Frontend] Do not modify caller params in LLM.score() (#59959) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 叮当猫 <2756823972@qq.com> Signed-off-by: apex-mochen <2756823972@qq.com> Co-authored-by: DeepSeek Harness Co-authored-by: Taneem Ibrahim Co-authored-by: Codex --- .../entrypoints/pooling/test_io_processor.py | 37 +++++++++++++++++++ vllm/entrypoints/pooling/offline.py | 2 + 2 files changed, 39 insertions(+) diff --git a/tests/entrypoints/pooling/test_io_processor.py b/tests/entrypoints/pooling/test_io_processor.py index 9d2207fdf15d..79bddc0ffe8f 100644 --- a/tests/entrypoints/pooling/test_io_processor.py +++ b/tests/entrypoints/pooling/test_io_processor.py @@ -7,6 +7,8 @@ from vllm import PoolingParams from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor +from vllm.entrypoints.pooling.offline import PoolingOfflineMixin +from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor from vllm.entrypoints.pooling.typing import OfflineEncodeInputsContext from vllm.exceptions import VLLMValidationError from vllm.renderers import TokenizeParams @@ -221,3 +223,38 @@ def test_qwen3_reranker_no_warning_when_template_provided(monkeypatch): ) mock_logger.warning.assert_not_called() + + +def _make_score_llm(monkeypatch) -> PoolingOfflineMixin: + """A scoring LLM with no model behind it. + + Input validation runs for real; only the rendering, engine and + post-processing boundaries are stubbed out. The pooling task is filled in + before any of them, so the mutation is observable without inference. + """ + proc = CrossEncoderIOProcessor.__new__(CrossEncoderIOProcessor) + proc.is_multimodal_model = False + proc.architecture = "CrossEncoder" + monkeypatch.setattr(proc, "get_request_factory_offline", lambda ctx: (None, 0)) + monkeypatch.setattr(proc, "post_process_offline", lambda ctx: []) + + llm = PoolingOfflineMixin.__new__(PoolingOfflineMixin) + llm.runner_type = "pooling" + llm.pooling_task = "classify" # SCORE_TYPE_MAP -> "cross-encoder" + llm.model_config = SimpleNamespace(hf_config=SimpleNamespace(num_labels=1)) + llm.pooling_io_processors = {"cross-encoder": proc} + monkeypatch.setattr(llm, "_run_tiling_engine", lambda *args, **kwargs: []) + return llm + + +def test_score_leaves_caller_pooling_params_untouched(monkeypatch): + llm = _make_score_llm(monkeypatch) + + params = PoolingParams() + llm.score("query", "doc", pooling_params=params) + + assert params.task is None, ( + "LLM.score() wrote the model's own pooling task into the caller's " + "PoolingParams; reusing that object with another pooling model now " + "fails with 'You cannot overwrite ...' before any inference runs." + ) diff --git a/vllm/entrypoints/pooling/offline.py b/vllm/entrypoints/pooling/offline.py index 0a0304e7a4c5..27388f1ff0a9 100644 --- a/vllm/entrypoints/pooling/offline.py +++ b/vllm/entrypoints/pooling/offline.py @@ -362,6 +362,8 @@ def score( assert isinstance(pooling_params, PoolingParams) pooling_task = io_processor.pooling_task + # Clone to avoid modifying the caller's pooling parameters. + pooling_params = pooling_params.clone() if pooling_params.task is None: pooling_params.task = pooling_task elif pooling_params.task != pooling_task: From fdfc171a7495bed2997609a0758d7e69c69440c0 Mon Sep 17 00:00:00 2001 From: Junyi Yao Date: Wed, 7 Oct 2026 18:07:35 -0700 Subject: [PATCH 053/624] [Bugfix][ToolParser] Stream every tool call when one delta completes several (#60354) Signed-off-by: Junyi Yao Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: sfeng33 <4florafeng@gmail.com> --- tests/tool_use/test_tool_choice_required.py | 152 +++++--- vllm/parser/abstract_parser.py | 15 +- vllm/renderers/online_derenderer.py | 6 - vllm/tool_parsers/streaming.py | 373 ++++++++++++-------- 4 files changed, 341 insertions(+), 205 deletions(-) diff --git a/tests/tool_use/test_tool_choice_required.py b/tests/tool_use/test_tool_choice_required.py index ebdea5aa0b9f..be58ffa83aa9 100644 --- a/tests/tool_use/test_tool_choice_required.py +++ b/tests/tool_use/test_tool_choice_required.py @@ -295,45 +295,61 @@ def test_structured_outputs_json_without_parameters( ) -def _collect_required_tool_streaming_json(output_json: str, delta_len: int) -> str: +def _stream_required_tool_calls( + output_json: str, + deltas: list[str], + tool_call_id_type: str = "random", + tool_call_idx: int | None = None, +) -> list[dict]: + """Feed ``deltas`` (which concatenate to ``output_json``) through the + required-tool streaming helper and rebuild the calls per array index, + checking that every index gets exactly one id/name chunk first.""" + assert "".join(deltas) == output_json previous_text = "" - function_name_returned = False - messages = [] - for i in range(0, len(output_json), delta_len): - delta_text = output_json[i : i + delta_len] + calls: dict[int, dict] = {} + for delta_text in deltas: current_text = previous_text + delta_text - - delta_message, function_name_returned = extract_required_tool_call_streaming( + delta_message, _ = extract_required_tool_call_streaming( previous_text=previous_text, current_text=current_text, - delta_text=delta_text, - function_name_returned=function_name_returned, - tool_call_idx=None, - tool_call_id_type="random", + tool_call_idx=tool_call_idx, + tool_call_id_type=tool_call_id_type, ) - - if delta_message: - messages.append(delta_message) - previous_text = current_text + if delta_message is None: + continue + assert delta_message.tool_calls + for tc in delta_message.tool_calls: + assert tc.function is not None + if tc.id is not None: + assert tc.index not in calls, "id/name sent twice for one index" + assert tc.type == "function" + assert tc.function.name + calls[tc.index] = { + "id": tc.id, + "name": tc.function.name, + "arguments": tc.function.arguments or "", + } + else: + assert tc.index in calls, "arguments before id/name" + assert tc.function.name is None + assert tc.function.arguments + calls[tc.index]["arguments"] += tc.function.arguments + assert list(calls) == list(range(len(calls))) + return [calls[i] for i in range(len(calls))] - assert len(messages) > 0 - combined_messages = "[" - for message in messages: - fn = message.tool_calls[0].function - assert fn is not None - if fn.name: - if len(combined_messages) > 1: - combined_messages += "}," +def _fixed_len_deltas(text: str, delta_len: int) -> list[str]: + return [text[i : i + delta_len] for i in range(0, len(text), delta_len)] - combined_messages += ( - '{"name": "' + fn.name + '", "parameters": ' + (fn.arguments or "") - ) - else: - combined_messages += fn.arguments or "" - combined_messages += "}]" - return combined_messages + +def _assert_streams_to(output: list[dict], deltas: list[str]) -> None: + calls = _stream_required_tool_calls(json.dumps(output), deltas) + assert [ + {"name": c["name"], "parameters": json.loads(c["arguments"])} for c in calls + ] == output + for call, expected in zip(calls, output): + assert call["arguments"] == json.dumps(expected["parameters"]) @pytest.mark.parametrize("output", VALID_TOOLS) @@ -343,11 +359,7 @@ def test_streaming_output_valid(output, empty_params, delta_len): output = deepcopy(output) if empty_params: output = [{"name": o["name"], "parameters": {}} for o in output] - output_json = json.dumps(output) - - combined_messages = _collect_required_tool_streaming_json(output_json, delta_len) - assert json.loads(combined_messages) == output - assert json.dumps(json.loads(combined_messages)) == output_json + _assert_streams_to(output, _fixed_len_deltas(json.dumps(output), delta_len)) @pytest.mark.parametrize( @@ -358,22 +370,78 @@ def test_streaming_output_valid(output, empty_params, delta_len): "a }} b", 'a " } b', r"a \ } b", + # Text that looks like the end of one call and the start of another. + '"}}, {"name": "get_forecast", "parameters": {"city": "x', ], ) @pytest.mark.parametrize("delta_len", [1, 2, 3, 8, 9999]) def test_streaming_output_valid_with_braces_in_string(city, delta_len): - output = [{"name": "get_current_weather", "parameters": {"city": city}}] - output_json = json.dumps(output) - combined_messages = _collect_required_tool_streaming_json(output_json, delta_len) - assert json.loads(combined_messages) == output - assert json.dumps(json.loads(combined_messages)) == output_json + output = [ + {"name": "get_current_weather", "parameters": {"city": city}}, + {"name": "get_forecast", "parameters": {"city": city, "days": 1}}, + ] + _assert_streams_to(output, _fixed_len_deltas(json.dumps(output), delta_len)) def test_streaming_output_valid_with_trailing_extra_data(): output = [{"name": "get_current_weather", "parameters": {"city": "Vienna"}}] output_json = json.dumps(output) + "\nDONE" - combined_messages = _collect_required_tool_streaming_json(output_json, delta_len=3) - assert json.loads(combined_messages) == output + calls = _stream_required_tool_calls(output_json, _fixed_len_deltas(output_json, 3)) + assert [ + {"name": c["name"], "parameters": json.loads(c["arguments"])} for c in calls + ] == output + + +TWO_CALLS = [ + {"name": "get_current_weather", "parameters": {"city": "Dallas"}}, + {"name": "get_forecast", "parameters": {"city": "Dallas", "days": 3}}, +] + + +def _split_at(text: str, cuts: list[int]) -> list[str]: + bounds = [0, *cuts, len(text)] + return [text[a:b] for a, b in zip(bounds, bounds[1:])] + + +def test_streaming_multiple_calls_completed_in_one_delta(): + """Regression test for #60351: a delta that completes more than one + call must stream every call, not only the last one.""" + text = json.dumps(TWO_CALLS) + _assert_streams_to(TWO_CALLS, [text[:1], text[1:]]) + _assert_streams_to(TWO_CALLS, [text]) + + +def test_streaming_delta_crossing_call_boundary(): + """A delta that closes one call and opens the next must attribute the + closing argument text to the first call and still announce the next.""" + text = json.dumps(TWO_CALLS) + boundary = text.index("}, {") + 1 + name_end = text.index('"get_forecast"') + 6 + for cuts in ( + [boundary - 1, name_end], # '}, {"name": "get_f' in one delta + [boundary - 2, boundary + 2], # '"}}, {' then the rest + [boundary + 2], # opening brace of the next call ends a delta + ): + _assert_streams_to(TWO_CALLS, _split_at(text, cuts)) + + +def test_streaming_invalid_name_string_is_never_announced(): + """A closed name that is not a valid JSON string must not be streamed + (the non-streaming path rejects such an array as well).""" + text = '[{"name": "bad\\x", "parameters": {"city": "x"}}]' + assert _stream_required_tool_calls(text, _fixed_len_deltas(text, 4)) == [] + + +def test_streaming_tool_call_idx_increments_per_started_call(): + """kimi_k2 ids number the calls started in one delta consecutively.""" + text = json.dumps(TWO_CALLS) + calls = _stream_required_tool_calls( + text, [text], tool_call_id_type="kimi_k2", tool_call_idx=5 + ) + assert [c["id"] for c in calls] == [ + "functions.get_current_weather:5", + "functions.get_forecast:6", + ] FUNCTION_TOOL = FunctionTool( diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index f75d7ce9028d..4af133c9882c 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -741,8 +741,6 @@ def _extract_tool_calls_streaming( extract_required_tool_call_streaming( previous_text=previous_text, current_text=current_text, - delta_text=delta_text, - function_name_returned=function_name_returned, tool_call_idx=tool_call_idx, tool_call_id_type=tool_call_id_type, ) @@ -952,12 +950,13 @@ def parse_delta( elif not delta_message.reasoning: delta_message.reasoning = reasoning_from_this_batch - if ( - delta_message - and delta_message.tool_calls - and delta_message.tool_calls[0].id is not None - ): - state.history_tool_call_cnt += 1 + if delta_message and delta_message.tool_calls: + # A delta may start several tool calls at once (e.g. the + # "required" path when one delta completes more than one + # call), so count every id rather than only the first. + state.history_tool_call_cnt += sum( + tc.id is not None for tc in delta_message.tool_calls + ) # No phase active: pass through as content. # Skip when reasoning just ended in this delta — the engine already diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index bd5bf88fbe8c..66b73fb4912c 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -640,12 +640,6 @@ def _replay(token_ids: list[int], chunk_lens: list[int]) -> None: if tc.index < len(last_tool_call_ids): # Pin: reuse the ID already recorded for this index # rather than one a from scratch replay regenerated. - # Real trigger not just defensive with - # tool_choice="required", - # extract_required_tool_call_streaming resets - # function_name_returned to False whenever the - # partial JSON transiently fails to parse which - # re-emits id+name for the same index on replay. tc.id = last_tool_call_ids[tc.index] else: last_tool_call_ids.append(tc.id) diff --git a/vllm/tool_parsers/streaming.py b/vllm/tool_parsers/streaming.py index 70da7f3cc6b8..07c1a2e031eb 100644 --- a/vllm/tool_parsers/streaming.py +++ b/vllm/tool_parsers/streaming.py @@ -2,19 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json +from dataclasses import dataclass from typing import TYPE_CHECKING -import partial_json_parser -import regex as re -from partial_json_parser.core.options import Allow - from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, DeltaToolCall, ) -from vllm.tool_parsers.utils import partial_json_loads from vllm.utils.mistral import is_mistral_tokenizer if TYPE_CHECKING: @@ -23,63 +19,6 @@ TokenizerLike = object -def _bracket_level_state( - s: str, opening: str = "{", closing: str = "}" -) -> tuple[int, bool, bool]: - level = 0 - in_string = False - escaped = False - for char in s: - if escaped: - escaped = False - continue - if in_string and char == "\\": - escaped = True - continue - if char == '"': - in_string = not in_string - continue - if not in_string: - if char == opening: - level += 1 - elif char == closing: - level -= 1 - return level, in_string, escaped - - -def filter_delta_text( - delta_text: str, - previous_text: str, -) -> tuple[str, bool]: - """Trim trailing tool-list delimiters from required-tool streaming text.""" - bracket_level, in_string, escaped = _bracket_level_state(previous_text) - updated_delta = "" - passed_zero = False - for char in delta_text: - if escaped: - escaped = False - elif in_string: - if char == "\\": - escaped = True - elif char == '"': - in_string = False - elif char == '"': - in_string = True - elif char == "{": - bracket_level += 1 - passed_zero = bracket_level == 0 - elif char == "}": - bracket_level -= 1 - passed_zero = bracket_level == 0 - - if bracket_level != 0: - updated_delta += char - else: - if not in_string and char == ",": - break - return updated_delta, passed_zero - - def extract_named_tool_call_streaming( *, delta_text: str, @@ -124,100 +63,236 @@ def extract_named_tool_call_streaming( ) +@dataclass +class _RequiredToolCallSpan: + """Where one element of a ``tool_choice="required"`` array sits in the text. + + The model output for ``tool_choice="required"`` is a JSON array of + ``{"name": ..., "parameters": ...}`` objects. ``name`` is only set once + its string value is closed and decodes as JSON; ``name_end`` is the + offset just past that closing quote. ``args_start``/``args_end`` are the + character offsets of the ``parameters`` value (``args_end`` is ``None`` + while the value is still being generated). + + All offsets are prefix-stable: scanning any prefix of the text yields the + same spans truncated to that prefix. That is what lets one scan of the + current text tell what an earlier, shorter text had already revealed. + """ + + name: str | None = None + name_end: int | None = None + args_start: int | None = None + args_end: int | None = None + + def ready_at(self, length: int) -> bool: + """Whether the first ``length`` characters of the text already made + this call streamable: its name was complete and its parameters value + had started.""" + return ( + self.name is not None + and self.name_end is not None + and self.name_end <= length + and self.args_start is not None + and self.args_start < length + ) + + @property + def ready(self) -> bool: + return self.name is not None and self.args_start is not None + + +# Parse states for one element object of the required-tool-call array. +_EXPECT_KEY = 0 +_EXPECT_COLON = 1 +_EXPECT_VALUE = 2 +_IN_STRING_VALUE = 3 +_IN_CONTAINER_VALUE = 4 +_IN_SCALAR_VALUE = 5 +_AFTER_VALUE = 6 + + +def _scan_required_tool_calls(text: str) -> list[_RequiredToolCallSpan]: + """Locate every tool call of a (possibly partial) required-tool-call array. + + A single string- and nesting-aware pass over ``text`` that records, for + each top-level array element, its ``name`` (once complete) and the + character span of its ``parameters`` value. Text after the closing ``]`` + of the array is ignored. Only the element level of the array is + interpreted, so braces, brackets and quotes inside nested values and + strings cannot confuse it. + """ + spans: list[_RequiredToolCallSpan] = [] + depth = 0 # nesting depth of ``{}`` and ``[]``; element objects sit at 2 + in_string = False + escaped = False + string_start = 0 + cur: _RequiredToolCallSpan | None = None + expect = _EXPECT_KEY + key: str | None = None + value_start = 0 + + def end_value(pos: int) -> None: + nonlocal expect + assert cur is not None + if key == "parameters": + cur.args_end = pos + elif key == "name" and expect == _IN_STRING_VALUE: + try: + cur.name = json.loads(text[value_start:pos]) + cur.name_end = pos + except json.JSONDecodeError: + # Not a valid JSON string: the call can never become ready, + # which matches the non-streaming path rejecting the array. + cur.name = None + expect = _AFTER_VALUE + + def start_value(pos: int, state: int) -> None: + nonlocal expect, value_start + assert cur is not None + value_start = pos + if key == "parameters": + cur.args_start = pos + expect = state + + for i, ch in enumerate(text): + if in_string: + if escaped: + escaped = False + elif ch == "\\": + escaped = True + elif ch == '"': + in_string = False + if depth == 2 and cur is not None: + if expect == _EXPECT_KEY: + key = text[string_start + 1 : i] + expect = _EXPECT_COLON + elif expect == _IN_STRING_VALUE: + end_value(i + 1) + continue + + if ch == '"': + in_string = True + string_start = i + if depth == 2 and cur is not None and expect == _EXPECT_VALUE: + start_value(i, _IN_STRING_VALUE) + continue + + if ch in "{[": + if depth == 0: + if ch != "[": + return spans + elif depth == 1: + if ch != "{": + return spans + cur = _RequiredToolCallSpan() + spans.append(cur) + expect = _EXPECT_KEY + key = None + elif depth == 2 and cur is not None and expect == _EXPECT_VALUE: + start_value(i, _IN_CONTAINER_VALUE) + depth += 1 + continue + + if ch in "}]": + if depth == 2 and cur is not None and expect == _IN_SCALAR_VALUE: + end_value(i) + depth -= 1 + if depth == 2 and cur is not None and expect == _IN_CONTAINER_VALUE: + end_value(i + 1) + elif depth == 1: + cur = None + elif depth <= 0: + return spans + continue + + if depth != 2 or cur is None: + continue + + if ch == ":": + if expect == _EXPECT_COLON: + expect = _EXPECT_VALUE + elif ch == ",": + if expect == _IN_SCALAR_VALUE: + end_value(i) + expect = _EXPECT_KEY + key = None + elif ch.isspace(): + if expect == _IN_SCALAR_VALUE: + end_value(i) + elif expect == _EXPECT_VALUE: + start_value(i, _IN_SCALAR_VALUE) + + return spans + + def extract_required_tool_call_streaming( *, previous_text: str, current_text: str | None, - delta_text: str, - function_name_returned: bool, tool_call_idx: int | None, tool_call_id_type: str, ) -> tuple[DeltaMessage | None, bool]: - if current_text is None or current_text == "": - # if the current text is empty, we cannot parse it - return None, function_name_returned - try: - flags = Allow.ALL - obj, _ = partial_json_loads(current_text, flags) - except ( - partial_json_parser.core.exceptions.MalformedJSON, - json.JSONDecodeError, - ): - obj = None - - # check if the current text is a valid array - # containing a partial tool calling object - # if not repeat - if obj is None or not isinstance(obj, list) or not len(obj) > 0: - function_name_returned = False - delta_message = None - else: - _, finishes_previous_tool = filter_delta_text(delta_text, previous_text) - # take the last tool call from the generated list - current_tool_call = obj[-1] - - # once parameters have been generated the name is complete as well - if not finishes_previous_tool and ( - "name" not in current_tool_call or "parameters" not in current_tool_call - ): - function_name_returned = False - delta_message = None - else: - if not function_name_returned: - # get partly generated arguments from the latest tool call - param_match = re.search( - r'.*"parameters":\s*(.*)', current_text, re.DOTALL - ) - if param_match: - arguments = param_match.group(1) - arguments_prefix = current_text[: param_match.start(1)] - arguments, _ = filter_delta_text(arguments, arguments_prefix) - else: - arguments = "" - - # if this iteration finishes a previous tool call but a - # new incomplete tool is already generated, take the - # previous from the list - if finishes_previous_tool and "parameters" not in current_tool_call: - current_tool_call = obj[-2] - - function_name_returned = True - tool_call_id = make_tool_call_id( - id_type=tool_call_id_type, - func_name=current_tool_call["name"], - idx=tool_call_idx, - ) - delta_message = DeltaMessage( - tool_calls=[ - DeltaToolCall( - id=tool_call_id, - function=DeltaFunctionCall( - name=current_tool_call["name"], arguments=arguments - ), - index=len(obj) - 1, - type="function", - ) - ] - ) + """Stream the tool calls of a ``tool_choice="required"`` JSON array. + + ``previous_text`` must be a prefix of ``current_text`` (the accumulated + text before and after this delta). Every array element that has something + not yet sent is emitted, in array order, so a delta may carry several + ``DeltaToolCall`` entries: the first chunk of a call carries its id, name + and the arguments generated so far; later chunks carry only new argument + text. What was already sent is read off the same scan, since the spans + are prefix-stable, so no state is needed beyond the two texts. - else: - delta_text, _ = filter_delta_text(delta_text, previous_text) - - if delta_text != "": - delta_message = DeltaMessage( - tool_calls=[ - DeltaToolCall( - function=DeltaFunctionCall( - # OpenAI API returns None - # instead of name every time - name=None, - arguments=delta_text, - ), - index=len(obj) - 1, - ) - ] + Returns the delta (``None`` when nothing is new) and whether at least one + call has been announced so far. + """ + if not current_text: + return None, False + + spans = _scan_required_tool_calls(current_text) + sent_len = len(previous_text) + + tool_calls: list[DeltaToolCall] = [] + started = 0 + for index, span in enumerate(spans): + if not span.ready: + # Elements are generated in order: nothing later can be ready. + break + assert span.args_start is not None + end = span.args_end if span.args_end is not None else len(current_text) + if span.ready_at(sent_len): + # Already announced: send only the new argument text. + sent_end = min(end, sent_len) + new_arguments = current_text[sent_end:end] + if new_arguments: + tool_calls.append( + DeltaToolCall( + index=index, + function=DeltaFunctionCall( + # OpenAI API returns None instead of name every time + name=None, + arguments=new_arguments, + ), ) - else: - delta_message = None + ) + continue - return delta_message, function_name_returned + idx = None if tool_call_idx is None else tool_call_idx + started + started += 1 + tool_calls.append( + DeltaToolCall( + id=make_tool_call_id( + id_type=tool_call_id_type, func_name=span.name, idx=idx + ), + type="function", + index=index, + function=DeltaFunctionCall( + name=span.name, + arguments=current_text[span.args_start : end], + ), + ) + ) + + function_name_returned = any(span.ready for span in spans) + if not tool_calls: + return None, function_name_returned + return DeltaMessage(tool_calls=tool_calls), function_name_returned From 8733a21d5517d56fb7ef573b379eaeb7c08e1d03 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Thu, 8 Oct 2026 03:22:04 +0200 Subject: [PATCH 054/624] [Core] Bound execute_dummy_batch RPC waits by the execute-model timeout (#60398) Signed-off-by: NickLucche --- tests/v1/executor/test_multiproc_executor.py | 22 +++++++++++++------- vllm/v1/executor/multiproc_executor.py | 6 +++++- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/tests/v1/executor/test_multiproc_executor.py b/tests/v1/executor/test_multiproc_executor.py index 706bfae38189..6a3b3d516ebf 100644 --- a/tests/v1/executor/test_multiproc_executor.py +++ b/tests/v1/executor/test_multiproc_executor.py @@ -70,16 +70,24 @@ def fail(): @pytest.mark.parametrize("stalled", [False, True]) -def test_take_draft_token_ids_uses_execute_model_timeout(monkeypatch, stalled): - """A wedged draft-token readback must not block EngineCore forever.""" +@pytest.mark.parametrize( + "method,result", + [ + ("take_draft_token_ids", DraftTokenIds(["req"], [[1, 2, 3]])), + ("execute_dummy_batch", None), + ], +) +def test_model_rpc_uses_execute_model_timeout(monkeypatch, method, result, stalled): + """A wedged model RPC must not block EngineCore forever. For + execute_dummy_batch this is an idle DP rank stuck in a collective with a + hung peer.""" monkeypatch.setenv("VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS", "7") - draft = DraftTokenIds(["req"], [[1, 2, 3]]) def dequeue(*, timeout=None): assert timeout is not None and 0 < timeout <= 7 if stalled: raise TimeoutError - return WorkerProc.ResponseStatus.SUCCESS, draft + return WorkerProc.ResponseStatus.SUCCESS, result executor: Any = MultiprocExecutor.__new__(MultiprocExecutor) executor.is_failed = False @@ -89,7 +97,7 @@ def dequeue(*, timeout=None): executor.response_mqs = [None, SimpleNamespace(dequeue=dequeue)] if stalled: - with pytest.raises(TimeoutError, match="take_draft_token_ids timed out"): - executor.take_draft_token_ids() + with pytest.raises(TimeoutError, match=f"{method} timed out"): + getattr(executor, method)() else: - assert executor.take_draft_token_ids() is draft + assert getattr(executor, method)() is result diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index 78290d25bafc..cad34368546d 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -364,7 +364,11 @@ def sample_tokens( # type: ignore[override] ) def execute_dummy_batch(self) -> None: - self.collective_rpc("execute_dummy_batch", unique_reply_rank=self.output_rank) + self.collective_rpc( + "execute_dummy_batch", + unique_reply_rank=self.output_rank, + timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS, + ) def take_draft_token_ids(self) -> DraftTokenIds | None: # OPTIMIZATION: Get output only from a single worker (output_rank) From 9dba339d07eaf16c9aa874eb042838f20c3ee0c8 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Wed, 7 Oct 2026 18:38:06 -0700 Subject: [PATCH 055/624] [Bugfix][ROCm] Keep MLA derived weights in place on weight reload (#60389) Signed-off-by: Matthew Wong Co-authored-by: Claude Opus 5.5 (1M context) --- .../model_loader/test_reload.py | 43 ++++++++++++++++++- .../layers/attention/mla_attention.py | 38 ++++++++++------ 2 files changed, 68 insertions(+), 13 deletions(-) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index 85ddea39469a..aaad16778f12 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -1051,6 +1051,37 @@ def loaded(value): assert not hasattr(rnorm, "qnorm_weight") +def _device_tensor_ptrs(worker) -> dict[str, int]: + """Addresses of every device tensor held by a module; graphs capture them.""" + ptrs = {} + for name, module in worker.get_model().named_modules(): + tensors = {**vars(module), **module._parameters, **module._buffers} + for attr, tensor in tensors.items(): + if isinstance(tensor, torch.Tensor) and tensor.device.type != "cpu": + ptrs[f"{name}.{attr}"] = tensor.data_ptr() + return ptrs + + +def _moved_tensors(llm, ptrs: list[dict[str, int]]) -> list[list[str]]: + after = llm.collective_rpc(_device_tensor_ptrs) + return [ + sorted(n for n, p in old.items() if new.get(n) != p) + for old, new in zip(ptrs, after) + ] + + +@pytest.mark.parametrize( + "use_aiter", + [ + False, + pytest.param( + True, + marks=pytest.mark.skipif( + not current_platform.is_rocm(), reason="AITER is ROCm-only" + ), + ), + ], +) @pytest.mark.parametrize( "tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])] ) @@ -1094,13 +1125,19 @@ def loaded(value): ), ], ) -def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner): +def test_reload_weights( + base_model, mul_model, add_model, tp_size, use_aiter, vllm_runner, monkeypatch +): if current_platform.device_count() < tp_size: pytest.skip(reason="Not enough CUDA devices") if "FP8" in base_model and _fp8_reload_unsupported(): pytest.skip(reason="Requires FP8 support") + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + if use_aiter: + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + with vllm_runner( model_name=base_model, tensor_parallel_size=tp_size, @@ -1109,12 +1146,16 @@ def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner): max_model_len=16, max_num_seqs=1, ) as llm: + # Reloading must update tensors in place: captured graphs keep pointers. + ptrs = llm.collective_rpc(_device_tensor_ptrs) llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_model}) + assert _moved_tensors(llm, ptrs) == [[]] * len(ptrs) mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0] add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0] assert mul_perp < add_perp llm.collective_rpc("reload_weights", kwargs={"weights_path": add_model}) + assert _moved_tensors(llm, ptrs) == [[]] * len(ptrs) mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0] add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0] assert add_perp < mul_perp diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index f550dc3d90d4..4286502ea9fe 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1272,23 +1272,34 @@ def process_weights_after_loading(self, act_dtype: torch.dtype): quark_quantize_weight_to_mxfp4, ) - self.W_K, self.W_K_scale = quark_quantize_weight_to_mxfp4(W_UK) - # Convert from (L, N, P) to (N, L, P) - self.W_K = self.W_K.transpose(0, 1) - self.W_K_scale = self.W_K_scale.transpose(0, 1) - - self.W_V, self.W_V_scale = quark_quantize_weight_to_mxfp4( - W_UV.permute(1, 2, 0) - ) + W_K, W_K_scale = quark_quantize_weight_to_mxfp4(W_UK) + W_V, W_V_scale = quark_quantize_weight_to_mxfp4(W_UV.permute(1, 2, 0)) + # Copy in place on reload so captured CUDA graphs stay valid. + # Convert W_K from (L, N, P) to (N, L, P). + for name, tensor in ( + ("W_K", W_K.transpose(0, 1)), + ("W_K_scale", W_K_scale.transpose(0, 1)), + ("W_V", W_V), + ("W_V_scale", W_V_scale), + ): + replace_parameter(self, name, tensor, prefer_copy=True) elif self.is_aiter_triton_fp8_bmm_enabled: W_K = W_UK.transpose(0, 1) # 16 512 128 W_V = W_UV.permute(1, 2, 0) # 16 128 512 - self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant( + W_K, W_K_scale = dynamic_per_batched_tensor_quant( W_K, dtype=current_platform.fp8_dtype() ) - self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant( + W_V, W_V_scale = dynamic_per_batched_tensor_quant( W_V, dtype=current_platform.fp8_dtype() ) + # Copy in place on reload so captured CUDA graphs stay valid. + for name, tensor in ( + ("W_K", W_K), + ("W_K_scale", W_K_scale), + ("W_V", W_V), + ("W_V_scale", W_V_scale), + ): + replace_parameter(self, name, tensor, prefer_copy=True) # The kernel operates on non-padded inputs. Hence, pre-compiling # triton kernel to avoid runtime compilation for unseen batch sizes @@ -1327,8 +1338,11 @@ def process_weights_after_loading(self, act_dtype: torch.dtype): # Convert from (L, N, P) to (N, P, L) replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True) if self.dcp_q_replicate: - self.W_UK_T_dcp_qrep = get_dcp_group().all_gather( - self.W_UK_T.contiguous(), dim=0 + replace_parameter( + self, + "W_UK_T_dcp_qrep", + get_dcp_group().all_gather(self.W_UK_T.contiguous(), dim=0), + prefer_copy=True, ) # If we should not load quant weights, we initialize the scales to 1.0 From a4b14e57b6d5ad25a2215eff93d6f93ba05005df Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Thu, 8 Oct 2026 09:56:19 +0800 Subject: [PATCH 056/624] [Platform] Allow BackendEnum dynamic register (#41091) Signed-off-by: wangxiyuan Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/test_attention_backend_registry.py | 141 ++++++++++++++++++ .../v1/attention/test_mla_prefill_registry.py | 120 +++++++++++---- .../backends/mla/prefill/registry.py | 54 ++++++- vllm/v1/attention/backends/registry.py | 85 ++++++++++- 4 files changed, 364 insertions(+), 36 deletions(-) diff --git a/tests/test_attention_backend_registry.py b/tests/test_attention_backend_registry.py index 5885411ed96a..14fe06699d01 100644 --- a/tests/test_attention_backend_registry.py +++ b/tests/test_attention_backend_registry.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest import torch from vllm.platforms.interface import DeviceCapability @@ -71,6 +72,48 @@ def get_builder_cls(): return None +class StringDecoratorAttentionBackend(AttentionBackend): + """Mock backend defined at module level so get_class() can resolve + the auto-generated class path (local classes are not importable).""" + + @staticmethod + def get_name(): + return "DECORATED" + + @staticmethod + def get_impl_cls(): + return CustomAttentionImpl + + @staticmethod + def get_builder_cls(): + return None + + @staticmethod + def get_required_kv_cache_layout(): + return None + + +class StringDecoratorMambaBackend(AttentionBackend): + """Mock mamba backend defined at module level so get_class() can + resolve the auto-generated class path.""" + + @staticmethod + def get_name(): + return "DECORATED" + + @staticmethod + def get_impl_cls(): + return CustomMambaAttentionImpl + + @staticmethod + def get_builder_cls(): + return None + + @staticmethod + def get_required_kv_cache_layout(): + return None + + def test_custom_attention_backend_requires_dcp_opt_in(): """A new implementation remains usable without implicitly opting into DCP.""" config = AttentionSelectorConfig( @@ -178,3 +221,101 @@ def test_register_custom_mamba_backend_with_class_path(): backend_cls = MambaAttentionBackendEnum.CUSTOM.get_class() assert backend_cls.get_name() == "CUSTOM_MAMBA" assert backend_cls.get_impl_cls() == CustomMambaAttentionImpl + + +@pytest.fixture(params=[AttentionBackendEnum, MambaAttentionBackendEnum]) +def _enum_cls(request): + return request.param + + +def _remove_dynamic_member(enum_cls, member): + value = member.value + name = member.name + enum_cls._member_map_.pop(name, None) + enum_cls._member_names_.remove(name) + delattr(enum_cls, name) + if enum_cls._value2member_map_.get(value) is member: + del enum_cls._value2member_map_[value] + + +def test_register_dynamic_enum_member(_enum_cls): + backend_path = { + AttentionBackendEnum: ( + "tests.test_attention_backend_registry.CustomAttentionBackend" + ), + MambaAttentionBackendEnum: ( + "tests.test_attention_backend_registry.CustomMambaAttentionBackend" + ), + }[_enum_cls] + + member = _enum_cls.register("DYNAMIC_TEST", backend_path) + assert member.name == "DYNAMIC_TEST" + assert member.value == backend_path + assert member is _enum_cls.DYNAMIC_TEST + assert _enum_cls["DYNAMIC_TEST"] is member + _remove_dynamic_member(_enum_cls, member) + + +def test_register_dynamic_enum_member_duplicate_raises(_enum_cls): + _enum_cls.register("DUP_TEST", "some.module.Class") + with pytest.raises(ValueError, match="already exists"): + _enum_cls.register("DUP_TEST", "other.module.OtherClass") + _remove_dynamic_member(_enum_cls, _enum_cls.DUP_TEST) + + +def test_register_dynamic_enum_member_reserved_name_raises(_enum_cls): + with pytest.raises(ValueError, match="Invalid or reserved backend name"): + _enum_cls.register("get_path", "some.module.Class") + + +@pytest.fixture( + params=[ + pytest.param((AttentionBackendEnum, False), id="attention"), + pytest.param((MambaAttentionBackendEnum, True), id="mamba"), + ] +) +def _enum_is_mamba(request): + return request.param + + +def _backend_path(enum_cls): + if enum_cls is AttentionBackendEnum: + return "tests.test_attention_backend_registry.CustomAttentionBackend" + return "tests.test_attention_backend_registry.CustomMambaAttentionBackend" + + +def test_register_backend_with_string_name_direct(_enum_is_mamba): + enum_cls, is_mamba = _enum_is_mamba + path = _backend_path(enum_cls) + register_backend("STRING_DIRECT", path, is_mamba=is_mamba) + + member = enum_cls.STRING_DIRECT + assert member.is_overridden() + assert member.get_path() == path + expected = "CUSTOM" if not is_mamba else "CUSTOM_MAMBA" + assert member.get_class().get_name() == expected + member.clear_override() + _remove_dynamic_member(enum_cls, member) + + +def test_register_backend_with_string_name_decorator(_enum_is_mamba): + enum_cls, is_mamba = _enum_is_mamba + # The decorated class must be module-level: the decorator derives a + # class path from it, and get_class() re-imports that path. A class + # defined inside this test would have a non-importable path. + decorated_cls = ( + StringDecoratorMambaBackend if is_mamba else StringDecoratorAttentionBackend + ) + + result = register_backend("STRING_DECORATOR", is_mamba=is_mamba)(decorated_cls) + assert result is decorated_cls + + member = enum_cls.STRING_DECORATOR + assert member.is_overridden() + assert member.get_path() == ( + f"{decorated_cls.__module__}.{decorated_cls.__qualname__}" + ) + assert member.get_class() is decorated_cls + assert member.get_class().get_name() == "DECORATED" + member.clear_override() + _remove_dynamic_member(enum_cls, member) diff --git a/tests/v1/attention/test_mla_prefill_registry.py b/tests/v1/attention/test_mla_prefill_registry.py index fc5efe3a5be0..0068d53b88cf 100644 --- a/tests/v1/attention/test_mla_prefill_registry.py +++ b/tests/v1/attention/test_mla_prefill_registry.py @@ -54,10 +54,20 @@ def test_prefill_backend_clone_has_isolated_metadata(): def cleanup_overrides(): """Clear any overrides after each test.""" yield - for member in MLAPrefillBackendEnum: + for member in list(MLAPrefillBackendEnum): member.clear_override() +def _remove_dynamic_member(enum_cls, member): + value = member.value + name = member.name + enum_cls._member_map_.pop(name, None) + enum_cls._member_names_.remove(name) + delattr(enum_cls, name) + if enum_cls._value2member_map_.get(value) is member: + del enum_cls._value2member_map_[value] + + def test_custom_is_not_alias_of_any_backend(): all_backends = list(MLAPrefillBackendEnum) @@ -83,21 +93,11 @@ def test_custom_unregistered_raises(): def test_register_custom_backend_with_class_path(): register_mla_prefill_backend( - backend=MLAPrefillBackendEnum.CUSTOM, - class_path=( - "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend" - ), + MLAPrefillBackendEnum.CUSTOM, + "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend", ) - assert MLAPrefillBackendEnum.CUSTOM.is_overridden() - - class_path = MLAPrefillBackendEnum.CUSTOM.get_path() - assert class_path == ( - "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend" - ) - - backend_cls = MLAPrefillBackendEnum.CUSTOM.get_class() - assert backend_cls.get_name() == "CUSTOM" + assert MLAPrefillBackendEnum.CUSTOM.get_class().get_name() == "CUSTOM" def test_register_custom_backend_as_decorator(): @@ -120,33 +120,20 @@ def run_prefill_context_chunk(self, chunk, q, k, v, out=None): def test_override_existing_backend(): - original_path = MLAPrefillBackendEnum.FLASH_ATTN.get_path() - register_mla_prefill_backend( - backend=MLAPrefillBackendEnum.FLASH_ATTN, - class_path=( - "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend" - ), + MLAPrefillBackendEnum.FLASH_ATTN, + "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend", ) - + assert MLAPrefillBackendEnum.FLASH_ATTN.get_class().get_name() == "CUSTOM" assert MLAPrefillBackendEnum.FLASH_ATTN.is_overridden() - assert MLAPrefillBackendEnum.FLASH_ATTN.get_path() != original_path - - backend_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class() - assert backend_cls.get_name() == "CUSTOM" def test_clear_override(): original_path = MLAPrefillBackendEnum.FLASH_ATTN.get_path() - register_mla_prefill_backend( - backend=MLAPrefillBackendEnum.FLASH_ATTN, - class_path=( - "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend" - ), + MLAPrefillBackendEnum.FLASH_ATTN, + "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend", ) - assert MLAPrefillBackendEnum.FLASH_ATTN.is_overridden() - MLAPrefillBackendEnum.FLASH_ATTN.clear_override() assert not MLAPrefillBackendEnum.FLASH_ATTN.is_overridden() assert MLAPrefillBackendEnum.FLASH_ATTN.get_path() == original_path @@ -172,3 +159,72 @@ def test_rocm_aiter_fa_registered(): # The AITER FA path is the fp16/bf16 generic-varlen prefill path. assert backend_cls.supports_dtype(torch.bfloat16) assert backend_cls.supports_dtype(torch.float16) + + +def test_register_dynamic_member(): + p = "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend" + member = MLAPrefillBackendEnum.register("DYNAMIC_TEST", p) + assert member.name == "DYNAMIC_TEST" + assert member is MLAPrefillBackendEnum["DYNAMIC_TEST"] + + _remove_dynamic_member(MLAPrefillBackendEnum, member) + + +def test_register_dynamic_member_duplicate_raises(): + path = "some.module.Class" + MLAPrefillBackendEnum.register("DUP_TEST", path) + with pytest.raises(ValueError, match="already exists"): + MLAPrefillBackendEnum.register("DUP_TEST", "other.module.OtherClass") + _remove_dynamic_member(MLAPrefillBackendEnum, MLAPrefillBackendEnum.DUP_TEST) + + +class StringDecoratorMLAPrefillBackend(MLAPrefillBackend): + """Mock backend defined at module level so get_class() can resolve + the auto-generated class path (local classes are not importable).""" + + supported_dtypes = [torch.bfloat16] + requires_r1_mla_dimensions = False + + @staticmethod + def get_name() -> str: + return "DECORATED" + + def run_prefill_new_tokens(self, q, k, v, return_softmax_lse): + raise NotImplementedError + + def run_prefill_context_chunk(self, chunk_idx, q, k, v): + raise NotImplementedError + + +def test_register_mla_prefill_backend_with_string_name_direct(): + register_mla_prefill_backend( + "STRING_DIRECT", + "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend", + ) + member = MLAPrefillBackendEnum.STRING_DIRECT + assert member.is_overridden() + assert member.get_class().get_name() == "CUSTOM" + + member.clear_override() + _remove_dynamic_member(MLAPrefillBackendEnum, member) + + +def test_register_mla_prefill_backend_with_string_name_decorator(): + # The decorated class must be module-level: the decorator derives a + # class path from it, and get_class() re-imports that path. A class + # defined inside this test would have a non-importable path. + decorated_cls = StringDecoratorMLAPrefillBackend + + result = register_mla_prefill_backend("STRING_DECORATOR")(decorated_cls) + assert result is decorated_cls + + member = MLAPrefillBackendEnum.STRING_DECORATOR + assert member.is_overridden() + assert member.get_path() == ( + f"{decorated_cls.__module__}.{decorated_cls.__qualname__}" + ) + assert member.get_class() is decorated_cls + assert member.get_class().get_name() == "DECORATED" + + member.clear_override() + _remove_dynamic_member(MLAPrefillBackendEnum, member) diff --git a/vllm/v1/attention/backends/mla/prefill/registry.py b/vllm/v1/attention/backends/mla/prefill/registry.py index f7ec31e15361..19d37fbc225b 100644 --- a/vllm/v1/attention/backends/mla/prefill/registry.py +++ b/vllm/v1/attention/backends/mla/prefill/registry.py @@ -10,11 +10,14 @@ from enum import Enum, EnumMeta from typing import TYPE_CHECKING +from vllm.logger import init_logger from vllm.utils.import_utils import resolve_obj_by_qualname if TYPE_CHECKING: from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend +logger = init_logger(__name__) + class _MLAPrefillBackendEnumMeta(EnumMeta): """Metaclass for MLAPrefillBackendEnum to provide better error messages.""" @@ -61,6 +64,37 @@ class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta): # set to None to avoid alias with other backend, whose value is an empty string CUSTOM = None + @classmethod + def register(cls, name: str, value: str) -> "MLAPrefillBackendEnum": + """Dynamically register a new MLA prefill backend enum member. + + Args: + name: The name for the new enum member + value: The fully qualified class path string + + Returns: + The newly created enum member + + """ + if name in cls._member_map_: + raise ValueError( + f"MLA prefill backend {name} already exists in {cls.__name__}. " + f"Use register_mla_prefill_backend(" + f"{cls.__name__}.{name}, '{value}') to override." + ) + if not name.isidentifier() or hasattr(cls, name): + raise ValueError(f"Invalid or reserved backend name: {name}") + + member = object.__new__(cls) + member._name_ = name + member._value_ = value + setattr(cls, name, member) + cls._member_map_[name] = member + cls._value2member_map_[value] = member + cls._member_names_.append(name) + logger.info("Registered new MLA prefill backend: %s -> %s", name, value) + return member + def get_path(self) -> str: """Get the class path for this backend (respects overrides). @@ -107,13 +141,15 @@ def clear_override(self) -> None: def register_mla_prefill_backend( - backend: MLAPrefillBackendEnum, + backend: MLAPrefillBackendEnum | str, class_path: str | None = None, ) -> Callable[[type], type]: """Register or override an MLA prefill backend implementation. Args: - backend: The MLAPrefillBackendEnum member to register. + backend: The MLAPrefillBackendEnum member to register, + or a string name for a new custom backend + (e.g., "CUSTOM_MLA_PREFILL"). class_path: Optional class path. If not provided and used as decorator, will be auto-generated from the class. @@ -131,13 +167,27 @@ class MyCustomFlashAttn(MLAPrefillBackend): class MyCustomPrefillBackend(MLAPrefillBackend): ... + # Register a new custom MLA prefill backend with a dynamic enum name + @register_mla_prefill_backend("CUSTOM_MLA_PREFILL") + class CustomMLAPrefillBackend(MLAPrefillBackend): + ... + # Direct registration register_mla_prefill_backend( MLAPrefillBackendEnum.CUSTOM, "my.module.MyCustomPrefillBackend" ) + # Direct registration with string name + register_mla_prefill_backend( + "CUSTOM_MLA_PREFILL", + "custom.attention.mla_prefill.CustomMLAPrefillBackend" + ) + """ + # Handle dynamic enum creation for string backend names + if isinstance(backend, str): + backend = MLAPrefillBackendEnum.register(backend, class_path or "") def decorator(cls: type) -> type: _MLA_PREFILL_OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}" diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index ad184f091329..207c112a384c 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -157,6 +157,38 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): # set to None to avoid alias with other backend, whose value is an empty string CUSTOM = None + @classmethod + def register(cls, name: str, value: str) -> "AttentionBackendEnum": + """Dynamically register a new backend enum member. + + Args: + name: The name for the new enum member + value: The fully qualified class path string + + Returns: + The newly created enum member + + """ + if name in cls._member_map_: + raise ValueError( + f"Backend {name} already exists in {cls.__name__}. " + f"Use register_backend({cls.__name__}.{name}, '{value}') " + f"to override." + ) + if not name.isidentifier() or hasattr(cls, name): + raise ValueError(f"Invalid or reserved backend name: {name}") + + # Create new enum member dynamically + member = object.__new__(cls) + member._name_ = name + member._value_ = value + setattr(cls, name, member) + cls._member_map_[name] = member + cls._value2member_map_[value] = member + cls._member_names_.append(name) + logger.info("Registered new attention backend: %s -> %s", name, value) + return member + def get_path(self, include_classname: bool = True) -> str: """Get the class path for this backend (respects overrides). @@ -214,6 +246,37 @@ class MambaAttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): backend.get_class() """ + @classmethod + def register(cls, name: str, value: str) -> "MambaAttentionBackendEnum": + """Dynamically register a new mamba backend enum member. + + Args: + name: The name for the new enum member + value: The fully qualified class path string + + Returns: + The newly created enum member + + """ + if name in cls._member_map_: + raise ValueError( + f"Backend {name} already exists in {cls.__name__}. " + f"Use register_backend({cls.__name__}.{name}, '{value}') " + f"to override." + ) + if not name.isidentifier() or hasattr(cls, name): + raise ValueError(f"Invalid or reserved backend name: {name}") + + member = object.__new__(cls) + member._name_ = name + member._value_ = value + setattr(cls, name, member) + cls._member_map_[name] = member + cls._value2member_map_[value] = member + cls._member_names_.append(name) + logger.info("Registered new attention backend: %s -> %s", name, value) + return member + MAMBA1 = "vllm.v1.attention.backends.mamba1_attn.Mamba1AttentionBackend" MAMBA2 = "vllm.v1.attention.backends.mamba2_attn.Mamba2AttentionBackend" SHORT_CONV = "vllm.v1.attention.backends.short_conv_attn.ShortConvAttentionBackend" @@ -275,14 +338,15 @@ def clear_override(self) -> None: def register_backend( - backend: AttentionBackendEnum | MambaAttentionBackendEnum, + backend: AttentionBackendEnum | MambaAttentionBackendEnum | str, class_path: str | None = None, is_mamba: bool = False, ) -> Callable[[type], type]: """Register or override a backend implementation. Args: - backend: The AttentionBackendEnum member to register + backend: The AttentionBackendEnum/MambaAttentionBackendEnum member to register, + or a string name for a new custom backend (e.g., "CUSTOM_MLA"). class_path: Optional class path. If not provided and used as decorator, will be auto-generated from the class. is_mamba: Whether the backend is a Mamba attention backend. @@ -306,13 +370,30 @@ class MyCustomMambaAttn: class MyCustomBackend: ... + # Register a new custom attention backend with a dynamic enum name + @register_backend("CUSTOM_MLA") + class CustomMLABackend: + ... + # Direct registration register_backend( AttentionBackendEnum.CUSTOM, "my.module.MyCustomBackend" ) + # Direct registration with string name + register_backend( + "CUSTOM_MLA", + "custom.attention.mla_v1.CustomMLAABackend" + ) + """ + # Handle dynamic enum creation for string backend names + if isinstance(backend, str): + if is_mamba: + backend = MambaAttentionBackendEnum.register(backend, class_path or "") + else: + backend = AttentionBackendEnum.register(backend, class_path or "") def decorator(cls: type) -> type: if is_mamba: From caff2b44cb4e37c81f2cc09234f557fd7cd96505 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 8 Oct 2026 10:30:35 +0800 Subject: [PATCH 057/624] [BugFix] Pick the DP world-group port at bind time via the coordination store (#51018) Signed-off-by: aoshen02 Signed-off-by: aoshen02 Co-authored-by: aoshen02 Co-authored-by: Claude Co-authored-by: Codex --- .../engine/test_core_engine_actor_manager.py | 69 ++++++++++- tests/v1/engine/test_dp_world_group_port.py | 109 ++++++++++++++++++ vllm/distributed/parallel_state.py | 48 +++++++- vllm/v1/engine/utils.py | 18 +++ 4 files changed, 240 insertions(+), 4 deletions(-) create mode 100644 tests/v1/engine/test_dp_world_group_port.py diff --git a/tests/v1/engine/test_core_engine_actor_manager.py b/tests/v1/engine/test_core_engine_actor_manager.py index f1a67ad9ba28..768f51f4379b 100644 --- a/tests/v1/engine/test_core_engine_actor_manager.py +++ b/tests/v1/engine/test_core_engine_actor_manager.py @@ -13,8 +13,9 @@ import pytest import ray import zmq +from torch.distributed import TCPStore -from vllm.utils.network_utils import make_zmq_socket, split_zmq_path +from vllm.utils.network_utils import get_open_port, make_zmq_socket, split_zmq_path from vllm.v1.engine.core import EngineCoreActorMixin from vllm.v1.engine.core_client import BackgroundResources from vllm.v1.engine.utils import ( @@ -134,6 +135,72 @@ def _make_addresses() -> EngineZmqAddresses: ) +@pytest.mark.parametrize( + ("overrides", "expect_store"), + [ + ({"data_parallel_backend": "mp"}, True), + ({}, True), + ( + { + "data_parallel_backend": "mp", + "data_parallel_rank": 1, + "local_engines_only": True, + }, + False, + ), + ({"data_parallel_backend": "mp", "data_parallel_rank_local": 0}, False), + ({"data_parallel_backend": "mp", "enable_elastic_ep": True}, False), + ({"data_parallel_backend": "mp", "data_parallel_size": 1}, False), + ], + ids=["mp", "ray", "non-master-node", "offline", "elastic-ep", "single-engine"], +) +def test_coordination_store_held_only_by_the_online_dp_master( + monkeypatch: pytest.MonkeyPatch, overrides: dict[str, Any], expect_store: bool +) -> None: + """The DP master node holds a coordination store for both engine backends, + so engines there pick world-group ports at bind time. Every other + deployment keeps the pre-allocated ports and must not hold one.""" + + class FakeManager: + def __init__(self, **kwargs) -> None: + self.kwargs = kwargs + + vllm_config = _make_vllm_config_ray_dp_multinode() + parallel_config = vllm_config.parallel_config + for name, value in overrides.items(): + setattr(parallel_config, name, value) + # A fresh port per case so a lingering handshake socket cannot fail the next. + parallel_config.data_parallel_rpc_port = get_open_port() + parallel_config._coord_store_port = 0 + monkeypatch.setattr("vllm.v1.engine.utils.CoreEngineProcManager", FakeManager) + monkeypatch.setattr("vllm.v1.engine.utils.CoreEngineActorManager", FakeManager) + monkeypatch.setattr( + "vllm.v1.engine.utils.wait_for_engine_startup", lambda *args, **kwargs: None + ) + + with launch_core_engines( + vllm_config, + executor_class=_DummyExecutor, + log_stats=False, + addresses=_make_addresses(), + ) as engine_launch: + assert engine_launch.engine_manager is not None + if not expect_store: + assert not parallel_config._coord_store_port + return + assert parallel_config._coord_store_port + # Engines look the store up while they start, so it must be reachable + # for as long as this frame is alive. + client = TCPStore( + parallel_config.data_parallel_master_ip, + parallel_config._coord_store_port, + is_master=False, + wait_for_workers=False, + ) + client.set("probe", b"1") + assert client.get("probe") == b"1" + + def _make_cpu_placement_group(): pg = ray.util.placement_group( [{"CPU": 0.001}, {"CPU": 1.0}], diff --git a/tests/v1/engine/test_dp_world_group_port.py b/tests/v1/engine/test_dp_world_group_port.py new file mode 100644 index 000000000000..49f5a59ab5c4 --- /dev/null +++ b/tests/v1/engine/test_dp_world_group_port.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The DP world group picks its port when it binds, not ahead of time.""" + +import multiprocessing as mp +import socket +import traceback +from multiprocessing.queues import Queue + +import pytest +import torch + +from vllm.distributed.utils import create_tcp_store +from vllm.utils.network_utils import get_open_port + +RESULT_TIMEOUT_S = 240 + + +def _join_dp_world_group( + dp_rank: int, coord_store_port: int, dp_master_port: int, result: Queue +) -> None: + """Bring up one DP rank of the world group and report what it saw.""" + try: + import torch.distributed as dist + + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.config.parallel import ParallelConfig + from vllm.distributed.parallel_state import init_distributed_environment + + parallel_config = ParallelConfig( + data_parallel_size=2, data_parallel_rank=dp_rank + ) + parallel_config.data_parallel_master_ip = "127.0.0.1" + parallel_config._data_parallel_master_port_list = [dp_master_port] + parallel_config._coord_store_port = coord_store_port + + with set_current_vllm_config(VllmConfig(parallel_config=parallel_config)): + init_distributed_environment( + world_size=1, rank=0, local_rank=0, backend="gloo" + ) + + reduced = torch.tensor([dp_rank + 1]) + dist.all_reduce(reduced) + result.put((dp_rank, int(reduced.item()), None)) + except Exception: + result.put((dp_rank, None, traceback.format_exc())) + + +@pytest.mark.skipif( + not torch.distributed.is_gloo_available(), reason="gloo backend required" +) +@pytest.mark.parametrize("with_coord_store", [True, False]) +def test_dp_world_group_survives_a_taken_preallocated_port(with_coord_store: bool): + """With a coordination store, the two ranks agree on a port the group binds + itself, so a pre-allocated port taken in the meantime cannot stop startup. + Without one, the pre-allocated port is still used as before. + """ + squatter = None + coord_store = None + try: + if with_coord_store: + # Something else grabbed the pre-allocated port after it was picked. + squatter = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + squatter.bind(("127.0.0.1", 0)) + squatter.listen(1) + dp_master_port = squatter.getsockname()[1] + coord_store = create_tcp_store( + "127.0.0.1", 0, is_master=True, world_size=-1, wait_for_workers=False + ) + coord_store_port = coord_store.port + else: + dp_master_port = get_open_port() + coord_store_port = 0 + + ctx = mp.get_context("spawn") + result: Queue = ctx.Queue() + procs = [ + ctx.Process( + target=_join_dp_world_group, + args=(dp_rank, coord_store_port, dp_master_port, result), + ) + for dp_rank in range(2) + ] + for proc in procs: + proc.start() + reports: list[tuple[int, int | None, str | None]] = [] + try: + while len(reports) < len(procs): + reports.append(result.get(timeout=RESULT_TIMEOUT_S)) + if reports[-1][2] is not None: + break # the peer would wait on this rank; do not wait for it + finally: + for proc in procs: + proc.join(timeout=30) + if proc.is_alive(): + proc.kill() + + for dp_rank, reduced, error in sorted(reports): + assert error is None, f"DP rank {dp_rank} failed:\n{error}" + # Both ranks joined the same group: 1 + 2 over the two ranks. + assert reduced == 3, f"DP rank {dp_rank} all-reduced to {reduced}" + + if with_coord_store: + assert coord_store is not None + bound_port = int(coord_store.get("world_pg_port").decode()) + assert bound_port != dp_master_port + finally: + if squatter is not None: + squatter.close() diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index de6680e02745..feb8f25d1eeb 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -42,6 +42,7 @@ import torch.distributed._functional_collectives as funcol import torch.distributed._symmetric_memory from torch.distributed import Backend, ProcessGroup, Store +from torch.distributed.constants import default_pg_timeout import vllm.envs as envs from vllm.distributed.device_communicators.base_device_communicator import ( @@ -49,6 +50,7 @@ ) from vllm.distributed.utils import ( StatelessProcessGroup, + create_tcp_store, get_cached_tcp_store_client, ) from vllm.logger import init_logger @@ -1829,6 +1831,7 @@ def init_distributed_environment( config = get_current_vllm_config_or_none() enable_elastic_ep = config is not None and config.parallel_config.enable_elastic_ep + world_pg_store: Store | None = None if ( config is not None and config.parallel_config.distributed_executor_backend != "external_launcher" @@ -1852,7 +1855,46 @@ def init_distributed_environment( distributed_init_method = get_distributed_init_method(ip, port) else: ip = parallel_config.data_parallel_master_ip - port = parallel_config.get_next_dp_init_port() + if ( + parallel_config._coord_store_port + and not torch.distributed.is_initialized() + ): + # Group rank 0 lets the world-group TCPStore bind port 0 + # and publishes the kernel-assigned port via the + # coordination store. A one-shot client of our own: the + # cached one is inherited across fork and shared with + # sibling workers. + coord_store = create_tcp_store( + ip, + parallel_config._coord_store_port, + is_master=False, + wait_for_workers=False, + ) + if rank == 0: + # Workers can only connect once the port is published. + world_pg_store = create_tcp_store( + ip, + 0, + world_size=world_size, + is_master=True, + wait_for_workers=False, + timeout=timeout or default_pg_timeout, + multi_tenant=True, + ) + port = world_pg_store.port + coord_store.set("world_pg_port", str(port).encode()) + else: + port = int(coord_store.get("world_pg_port").decode()) + world_pg_store = create_tcp_store( + ip, + port, + world_size=world_size, + is_master=False, + timeout=timeout or default_pg_timeout, + multi_tenant=True, + ) + else: + port = parallel_config.get_next_dp_init_port() distributed_init_method = get_distributed_init_method(ip, port) logger.debug( "Adjusting world_size=%d rank=%d distributed_init_method=%s for DP", @@ -1882,8 +1924,8 @@ def init_distributed_environment( "Fallback Gloo backend is not available." ) backend = "gloo" - store = None - if distributed_init_method.startswith("file://"): + store = world_pg_store + if store is None and distributed_init_method.startswith("file://"): store = torch.distributed.FileStore( distributed_init_method.removeprefix("file://"), world_size ) diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index 203d1b407c11..49e6884ac9f1 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -1195,6 +1195,23 @@ def launch_core_engines( else: coordinator = None + # Hold a coordination TCPStore (alive for this frame) so engines pick DP + # master ports at bind time; pre-allocated ports can be taken before use. + coord_store = None + if ( + dp_size > 1 + and not offline_mode + and dp_rank == 0 + and local_engine_count > 0 + and not parallel_config.enable_elastic_ep + ): + from vllm.distributed.utils import create_tcp_store + + coord_store = create_tcp_store( + host, 0, is_master=True, world_size=-1, wait_for_workers=False + ) + parallel_config._coord_store_port = coord_store.port + if parallel_config.data_parallel_backend == "ray": logger.info("Starting ray-based data parallel backend") @@ -1415,6 +1432,7 @@ def wait_for_engine_startup( "data_parallel_master_ip", "data_parallel_master_port", "_data_parallel_master_port_list", + "_coord_store_port", "data_parallel_size", ) } From 3722645539f4cbf11028e69eea0643dee5cdc773 Mon Sep 17 00:00:00 2001 From: bnkaiteam Date: Thu, 8 Oct 2026 11:32:35 +0900 Subject: [PATCH 058/624] [Bugfix][Frontend] Accept successful 2xx responses for batch output uploads (#60507) Signed-off-by: bnkaiteam Co-authored-by: Flora Feng <4florafeng@gmail.com> --- tests/entrypoints/launchers/test_run_batch.py | 29 ++++++++++++++----- vllm/entrypoints/launchers/run_batch.py | 4 +-- 2 files changed, 24 insertions(+), 9 deletions(-) diff --git a/tests/entrypoints/launchers/test_run_batch.py b/tests/entrypoints/launchers/test_run_batch.py index f7a250580727..1376eb0645f3 100644 --- a/tests/entrypoints/launchers/test_run_batch.py +++ b/tests/entrypoints/launchers/test_run_batch.py @@ -1071,17 +1071,32 @@ def _make_aiohttp_put_session(status: int = 200, body_text: str = ""): @pytest.mark.asyncio -async def test_upload_data_uploads_once_on_success(): - """A successful upload must not be retried (regression guard).""" - session = _make_aiohttp_put_session(status=200) - with patch( - "vllm.entrypoints.launchers.run_batch.aiohttp.ClientSession", - return_value=session, +@pytest.mark.parametrize("status", [200, 201, 204]) +@pytest.mark.parametrize("from_file", [False, True]) +async def test_upload_data_uploads_once_on_success(status, from_file, tmp_path): + """Successful PUT responses must not retry either upload path.""" + session = _make_aiohttp_put_session(status=status) + data_or_file = "payload" + if from_file: + path = tmp_path / "output.jsonl" + path.write_text(data_or_file, encoding="utf-8") + data_or_file = str(path) + + with ( + patch( + "vllm.entrypoints.launchers.run_batch.aiohttp.ClientSession", + return_value=session, + ), + patch( + "vllm.entrypoints.launchers.run_batch.asyncio.sleep", + new_callable=AsyncMock, + ) as sleep, ): await upload_data( - "https://example.com/output.jsonl", "payload", from_file=False + "https://example.com/output.jsonl", data_or_file, from_file=from_file ) assert session.put.call_count == 1 + sleep.assert_not_awaited() @pytest.mark.asyncio diff --git a/vllm/entrypoints/launchers/run_batch.py b/vllm/entrypoints/launchers/run_batch.py index f54e87dafefc..266592a731a5 100644 --- a/vllm/entrypoints/launchers/run_batch.py +++ b/vllm/entrypoints/launchers/run_batch.py @@ -385,7 +385,7 @@ async def upload_data(output_url: str, data_or_file: str, from_file: bool) -> No if from_file: with open(data_or_file, "rb") as file: async with session.put(output_url, data=file) as response: - if response.status != 200: + if not 200 <= response.status < 300: error_text = await response.text() raise Exception( f"Failed to upload file.\n" @@ -394,7 +394,7 @@ async def upload_data(output_url: str, data_or_file: str, from_file: bool) -> No ) else: async with session.put(output_url, data=data_or_file) as response: - if response.status != 200: + if not 200 <= response.status < 300: error_text = await response.text() raise Exception( f"Failed to upload data.\n" From 87d9996abe7ae0173fa8c1b6ccc2411041bcac81 Mon Sep 17 00:00:00 2001 From: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Date: Thu, 8 Oct 2026 04:37:51 +0200 Subject: [PATCH 059/624] [ROCm][Perf][GLM-5.3-Flash] BF16 splitk kernel for sparse MLA decode (#58584) Signed-off-by: simondanielsson Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> --- .../benchmark_sparse_mla_decode_bf16.py | 130 +++++++ .../attention/test_rocm_aiter_mla_sink.py | 6 + .../attention/test_rocm_triton_attn_dsv4.py | 96 +++++ .../backends/mla/rocm_aiter_mla_sparse.py | 53 ++- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 355 ++++++++++++++++++ 5 files changed, 627 insertions(+), 13 deletions(-) create mode 100644 benchmarks/kernels/benchmark_sparse_mla_decode_bf16.py diff --git a/benchmarks/kernels/benchmark_sparse_mla_decode_bf16.py b/benchmarks/kernels/benchmark_sparse_mla_decode_bf16.py new file mode 100644 index 000000000000..c32c8519aadb --- /dev/null +++ b/benchmarks/kernels/benchmark_sparse_mla_decode_bf16.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark the ROCm bf16 sparse-MLA decode split-K kernel in isolation. + +Example usage: +python3 benchmarks/kernels/benchmark_sparse_mla_decode_bf16.py \ + --rows 1 4 16 64 --lens 2048 --splits auto 4 16 32 +""" + +import argparse +import itertools +import json + +import torch + +from vllm.triton_utils import triton +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _rocm_sparse_attn_decode_ragged_bf16_triton, + _rocm_sparse_attn_prefill_ragged_triton, + rocm_sparse_decode_bf16_num_splits, +) + +BLOCK_H = 16 + + +def make_inputs(rows: int, live: int, args): + torch.manual_seed(args.seed) + kv = ( + torch.randn(args.capacity, args.head_dim, device="cuda", dtype=torch.bfloat16) + * 0.125 + ) + q = ( + torch.randn( + rows, args.heads, args.head_dim, device="cuda", dtype=torch.bfloat16 + ) + * 0.125 + ) + slots = torch.stack( + [torch.randperm(args.capacity, device="cuda")[:live] for _ in range(rows)] + ) + indices = slots.reshape(-1).to(torch.int32) + indptr = torch.arange(rows + 1, device="cuda", dtype=torch.int32) * live + sink = ( + torch.randn(args.heads, device="cuda", dtype=torch.float32) + if args.sink + else None + ) + return q, kv, indices, indptr, sink + + +def resolve_splits(spec: str, rows: int, live: int, args) -> int: + if spec != "auto": + return int(spec) + return rocm_sparse_decode_bf16_num_splits(rows, args.heads, live) + + +def benchmark(backend: str, rows: int, live: int, spec: str, args) -> dict: + q, kv, indices, indptr, sink = make_inputs(rows, live, args) + nope, rope = args.head_dim, 0 + splits = resolve_splits(spec, rows, live, args) + + def ragged(): + return _rocm_sparse_attn_prefill_ragged_triton( + q, kv, indices, indptr, args.scale, sink, nope, rope + ) + + def splitk(): + return _rocm_sparse_attn_decode_ragged_bf16_triton( + q, kv, indices, indptr, args.scale, sink, nope, rope, splits + ) + + run = ragged if backend == "ragged" else splitk + torch.testing.assert_close(run(), ragged(), atol=2e-2, rtol=2e-2) + ms = triton.testing.do_bench(run, warmup=args.warmup, rep=args.rep) + return dict( + backend=backend, + rows=rows, + live=live, + heads=args.heads, + head_dim=args.head_dim, + splits=splits if backend == "splitk" else 1, + workgroups=rows * triton.cdiv(args.heads, BLOCK_H) * splits, + us=ms * 1000, + kv_gbps=rows * live * args.head_dim * 2 / (ms * 1e6), + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", type=int, nargs="+", default=[1, 2, 4, 8, 16, 64]) + parser.add_argument( + "--lens", + type=int, + nargs="+", + default=[2048], + help="Selected KV slots per row; GLM-5.3-Flash uses index_topk=2048", + ) + parser.add_argument( + "--splits", + nargs="+", + default=["auto"], + help="Split counts to sweep, or auto for the runtime heuristic", + ) + parser.add_argument( + "--backends", + nargs="+", + choices=["ragged", "splitk"], + default=["ragged", "splitk"], + ) + parser.add_argument( + "--heads", type=int, default=8, help="Heads per rank; 64 at TP8 gives 8" + ) + parser.add_argument("--head-dim", type=int, default=512) + parser.add_argument("--capacity", type=int, default=65536) + parser.add_argument("--scale", type=float, default=0.044194173824159216) + parser.add_argument("--sink", action="store_true") + parser.add_argument("--warmup", type=int, default=25) + parser.add_argument("--rep", type=int, default=200) + parser.add_argument("--seed", type=int, default=123) + args = parser.parse_args() + + print(json.dumps(dict(gpu=torch.cuda.get_device_name(0), args=vars(args)))) + for backend, rows, live in itertools.product(args.backends, args.rows, args.lens): + specs = ["1"] if backend == "ragged" else args.splits + for spec in specs: + print(json.dumps(benchmark(backend, rows, live, spec, args)), flush=True) + + +if __name__ == "__main__": + main() diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sink.py b/tests/kernels/attention/test_rocm_aiter_mla_sink.py index 40b888c3622e..0b007dbc875c 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sink.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sink.py @@ -156,6 +156,8 @@ def test_sparse_mla_sink_matches_ragged_reference( num_decodes=batch_size, num_decode_tokens=batch_size, max_query_len=1, + max_seq_len=max(seq_lens), + topk_tokens=max(seq_lens), ) sinks = torch.linspace(-2.0, 6.0, real_heads, device=device) @@ -367,12 +369,15 @@ def test_sparse_mla_sink_matches_dense_attention_with_empty_rows_and_paged_cache for page, offset in rows ] lengths = torch.tensor([0, *(len(rows) for rows in selected_rows)]) + longest_row = max(len(rows) for rows in selected_rows) metadata = SimpleNamespace( block_size=block_size, num_prefills=0, num_decodes=len(selected_rows), num_decode_tokens=len(selected_rows), max_query_len=1, + max_seq_len=longest_row, + topk_tokens=longest_row, qo_indptr=torch.arange( len(selected_rows) + 1, dtype=torch.int32, device="cuda" ), @@ -479,6 +484,7 @@ def test_sparse_mla_sink_forward_mqa_preserves_split_query(dtype): num_decodes=num_tokens, num_decode_tokens=num_tokens, max_query_len=1, + max_seq_len=kv.shape[0], block_size=block_size, topk_tokens=impl.topk_indices_buffer.shape[1], req_id_per_token=torch.zeros(num_tokens, dtype=torch.int32, device="cuda"), diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 3b200be514d5..ca2cf2eb7017 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -596,6 +596,87 @@ def test_sparse_attn_prefill_ragged_kernel() -> None: torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) +NOPE_ONLY_HEAD_DIM = 512 + + +def _nope_only_ragged_inputs( + kv_lens: list[int], num_heads: int, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Build a GLM-like 512+0 NoPE decode batch with ragged KV segments.""" + num_kv = 4096 + kv = ( + torch.randn(num_kv, NOPE_ONLY_HEAD_DIM, dtype=torch.bfloat16, device=device) + * 0.125 + ) + q = ( + torch.randn( + len(kv_lens), + num_heads, + NOPE_ONLY_HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + * 0.125 + ) + rows = [torch.randperm(num_kv)[:length].tolist() for length in kv_lens] + indices, indptr = _ragged_from_rows(rows, device) + return q, kv, indices, indptr + + +# Lengths below the split count leave splits empty; non-multiples force a tail. +_NOPE_ONLY_KV_LENS = [0, 1, 3, 37, 129, 2048] + + +@requires_split_decode_arch +@pytest.mark.parametrize("num_splits", [2, 32]) +@pytest.mark.parametrize("with_sink", [True, False]) +@torch.inference_mode() +def test_sparse_attn_decode_bf16_split_k_matches_ragged( + num_splits: int, with_sink: bool +) -> None: + """Split-K decode must reproduce the single-workgroup ragged kernel.""" + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _rocm_sparse_attn_decode_ragged_bf16_triton, + _rocm_sparse_attn_prefill_ragged_triton, + ) + + device = torch.device("cuda") + set_random_seed(0) + num_heads = 16 + q, kv, indices, indptr = _nope_only_ragged_inputs( + _NOPE_ONLY_KV_LENS, num_heads, device + ) + attn_sink = ( + torch.randn(num_heads, dtype=torch.float32, device=device) + if with_sink + else None + ) + scale = NOPE_ONLY_HEAD_DIM**-0.5 + + expected = _rocm_sparse_attn_prefill_ragged_triton( + q=q, + kv=kv, + indices=indices, + indptr=indptr, + scale=scale, + attn_sink=attn_sink, + nope_head_dim=NOPE_ONLY_HEAD_DIM, + rope_head_dim=0, + ) + actual = _rocm_sparse_attn_decode_ragged_bf16_triton( + q=q, + kv=kv, + indices=indices, + indptr=indptr, + scale=scale, + attn_sink=attn_sink, + nope_head_dim=NOPE_ONLY_HEAD_DIM, + rope_head_dim=0, + num_splits=num_splits, + ) + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + def _sparse_prefill_ragged_inputs(nope_dim: int, rope_dim: int) -> dict: device = torch.device("cuda") set_random_seed(7) @@ -1105,6 +1186,21 @@ def test_decode_num_splits_gfx950(monkeypatch) -> None: assert mod._decode_gfx950_num_splits(1, 1, 128, 8192) == 32 assert mod._decode_gfx950_num_splits(17, 1, 128, 32) == 4 assert mod._decode_gfx950_num_splits(512, 1, 128, 7812) == 1 + assert mod._decode_gfx950_num_splits(4, 1, 2048.0, 0.0, 32) == 32 + + +def test_sparse_decode_bf16_num_splits_floor_and_tile_clamp(monkeypatch) -> None: + """Short rows skip splitting, and splits never exceed the tiles they walk.""" + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + monkeypatch.setattr(mod, "_decode_num_splits", lambda *args: 32) + monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: 32) + + floor = mod._SPARSE_DECODE_BF16_MIN_SPLIT_LEN + block_k = mod._SPARSE_DECODE_BF16_BLOCK_K + assert mod.rocm_sparse_decode_bf16_num_splits(1, 1, floor - 1) == 1 + assert mod.rocm_sparse_decode_bf16_num_splits(1, 1, floor) == floor // block_k + assert mod.rocm_sparse_decode_bf16_num_splits(1, 1, 2048) == 32 @requires_split_decode_arch diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index e473f859480b..e3d982e882bb 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -37,7 +37,9 @@ ) from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + rocm_sparse_attn_decode_bf16, rocm_sparse_attn_prefill, + rocm_sparse_decode_bf16_num_splits, ) from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout from vllm.v1.worker.workspace import current_workspace_manager @@ -904,20 +906,45 @@ def _forward_mla( self.sinks.reshape(1, self.num_heads, 1), q.shape[1], ).reshape(-1) - rocm_sparse_attn_prefill( - q=q, - kv=kv_c_and_k_pe_cache.view(-1, 1, q.shape[-1]), - indices=None, - topk_length=None, - scale=self.scale, - head_dim=q.shape[-1], - nope_head_dim=self.kv_lora_rank, - rope_head_dim=q.shape[-1] - self.kv_lora_rank, - attn_sink=triton_sinks, - output=output, - ragged_indices=attn_metadata.paged_kv_indices, - ragged_indptr=attn_metadata.paged_kv_indptr, + kv = kv_c_and_k_pe_cache.view(-1, 1, q.shape[-1]) + decode_num_splits = ( + rocm_sparse_decode_bf16_num_splits( + num_tokens, + q.shape[1], + min(attn_metadata.max_seq_len, attn_metadata.topk_tokens), + ) + if attn_metadata.num_decode_tokens == num_tokens + else 1 ) + if decode_num_splits > 1: + rocm_sparse_attn_decode_bf16( + q=q, + kv=kv, + scale=self.scale, + head_dim=q.shape[-1], + nope_head_dim=self.kv_lora_rank, + rope_head_dim=q.shape[-1] - self.kv_lora_rank, + attn_sink=triton_sinks, + output=output, + ragged_indices=attn_metadata.paged_kv_indices, + ragged_indptr=attn_metadata.paged_kv_indptr, + num_splits=decode_num_splits, + ) + else: + rocm_sparse_attn_prefill( + q=q, + kv=kv, + indices=None, + topk_length=None, + scale=self.scale, + head_dim=q.shape[-1], + nope_head_dim=self.kv_lora_rank, + rope_head_dim=q.shape[-1] - self.kv_lora_rank, + attn_sink=triton_sinks, + output=output, + ragged_indices=attn_metadata.paged_kv_indices, + ragged_indptr=attn_metadata.paged_kv_indptr, + ) output = AiterMLAHelper.get_mla_unpadded_o(self.num_heads, output) return output, None diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index a3412aa83508..90a73c101c0d 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -2208,6 +2208,119 @@ def _sparse_attn_prefill_ragged_kernel( ) +@triton.jit +def _sparse_attn_decode_ragged_bf16_partial_kernel( + q_ptr, + kv_ptr, + kv_indices_ptr, + kv_indptr_ptr, + part_m_ptr, + part_l_ptr, + part_acc_ptr, + q_stride_t, + q_stride_h, + q_stride_d, + kv_stride_n, + kv_stride_d, + pm_stride0, + pm_stride_s, + pa_stride0, + pa_stride_s, + pa_stride_h, + num_heads, + head_dim, + num_kv, + scale, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_K: tl.constexpr, + NUM_SPLITS: tl.constexpr, +): + query_idx = tl.program_id(0) + split_id = tl.program_id(1) + pid_h = tl.program_id(2) + + head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + dim_offsets = tl.arange(0, BLOCK_D) + head_mask = head_offsets < num_heads + dim_mask = dim_offsets < head_dim + + q = tl.load( + q_ptr + + query_idx * q_stride_t + + head_offsets[:, None] * q_stride_h + + dim_offsets[None, :] * q_stride_d, + mask=head_mask[:, None] & dim_mask[None, :], + other=0.0, + ) + + neg_large = -3.4028234663852886e38 + m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32) + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) + acc = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32) + + kv_start = tl.load(kv_indptr_ptr + query_idx) + kv_end = tl.load(kv_indptr_ptr + query_idx + 1) + kv_len = kv_end - kv_start + chunk = (kv_len + NUM_SPLITS - 1) // NUM_SPLITS + kv_lo = split_id * chunk + kv_hi = tl.minimum(kv_lo + chunk, kv_len) + + k_offsets = tl.arange(0, BLOCK_K) + slot = tl.load( + kv_indices_ptr + kv_start + kv_lo + k_offsets, + mask=kv_lo + k_offsets < kv_hi, + other=-1, + ) + for k_start in tl.range(kv_lo, kv_hi, BLOCK_K): + k_pos = k_start + k_offsets + in_range = k_pos < kv_hi + valid = in_range & (slot >= 0) & (slot < num_kv) + safe_slot = tl.where(valid, slot, 0) + + kv = tl.load( + kv_ptr + + _sparse_kv_row_offset(safe_slot[:, None], kv_stride_n) + + dim_offsets[None, :] * kv_stride_d, + mask=valid[:, None] & dim_mask[None, :], + other=0.0, + ) + + next_k_pos = k_start + BLOCK_K + k_offsets + slot = tl.load( + kv_indices_ptr + kv_start + next_k_pos, mask=next_k_pos < kv_hi, other=-1 + ) + + scores = tl.dot(q, tl.trans(kv)) * scale + scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large) + + m_block = tl.max(scores, axis=1) + m_new = tl.maximum(m_i, m_block) + alpha = tl.exp(m_i - m_new) + p = tl.exp(scores - m_new[:, None]) + p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0) + l_new = l_i * alpha + tl.sum(p, axis=1) + + acc = acc * alpha[:, None] + tl.dot(p.to(kv.dtype), kv) + m_i = m_new + l_i = l_new + + # Every launched slot is written; empty splits store the (neg_large, 0.0) + # sentinel the reduce combines to nothing. + pm_base = query_idx * pm_stride0 + split_id * pm_stride_s + head_offsets + tl.store(part_m_ptr + pm_base, m_i, mask=head_mask) + tl.store(part_l_ptr + pm_base, l_i, mask=head_mask) + tl.store( + part_acc_ptr + + query_idx * pa_stride0 + + split_id * pa_stride_s + + head_offsets[:, None] * pa_stride_h + + dim_offsets[None, :], + acc, + mask=head_mask[:, None] & dim_mask[None, :], + ) + + @triton.jit def _decode_e8m0_scales_triton(encoded_scales): scale_bits = encoded_scales.to(tl.int32) << 23 @@ -3664,6 +3777,182 @@ def _decode_gfx950_num_splits( return num_splits +# Empirically single-pass bf16 decode is faster than split-K below this number of +# selected slots per query, so we do not split. +_SPARSE_DECODE_BF16_MIN_SPLIT_LEN = 512 +_SPARSE_DECODE_BF16_BLOCK_H = 16 +_SPARSE_DECODE_BF16_BLOCK_K = 32 + + +def rocm_sparse_decode_bf16_num_splits( + num_queries: int, num_heads: int, sparse_len: int +) -> int: + """Number or kv splits in splitK for the sparse bf16 decode, or 1 for single-pass. + + Args: + num_queries: Decode rows in the batch. + num_heads: Query heads per row. + sparse_len: Longest selected KV run any decode row can walk. + + Returns: + The split count, or 1 when the caller should use the single-pass kernel. + + """ + if sparse_len < _SPARSE_DECODE_BF16_MIN_SPLIT_LEN: + return 1 + block_k = _SPARSE_DECODE_BF16_BLOCK_K + heads_blocks = triton.cdiv(num_heads, _SPARSE_DECODE_BF16_BLOCK_H) + select = _decode_gfx950_num_splits if _ON_GFX950 else _decode_num_splits + num_splits = select(num_queries, heads_blocks, sparse_len, 0.0, block_k) + # Number of splits cannot exceed the available k tiles + return max(1, min(num_splits, math.ceil(sparse_len / block_k))) + + +def _rocm_sparse_attn_decode_ragged_bf16_triton( + q: torch.Tensor, + kv: torch.Tensor, + indices: torch.Tensor, + indptr: torch.Tensor, + scale: float, + attn_sink: torch.Tensor | None, + nope_head_dim: int, + rope_head_dim: int, + num_splits: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Split-K decode over an bf16 ragged KV cache. + + Partitions each query's selected tokens across different workgroups and combines + the partials through reduction. + + Args: + q: Queries laid out as ``[sq, h, d]``. + kv: Unquantized KV rows laid out as ``[skv, d]``. + indices: Flattened per-query KV slots. + indptr: Segment offsets into ``indices``, ``[sq + 1]``. + scale: Softmax scale. + attn_sink: Optional per-head sink logits. + nope_head_dim: NoPE width of ``d``. + rope_head_dim: RoPE width of ``d``. + num_splits: Number of KV splits per query. + out: Optional destination with ``d`` trailing elements. + + Returns: + The attention output, ``out`` when provided. + + """ + assert q.ndim == 3, f"expected q=[sq,h,d], got {q.shape}" + assert kv.ndim == 2, f"expected kv=[skv,d], got {kv.shape}" + assert indices.ndim == 1, f"expected indices=[nnz], got {indices.shape}" + assert indptr.ndim == 1, f"expected indptr=[sq+1], got {indptr.shape}" + assert not q.is_cpu and not kv.is_cpu and not indices.is_cpu and not indptr.is_cpu + + indices = _as_int32_contiguous_1d(indices) + indptr = _as_int32_contiguous_1d(indptr) + has_attn_sink = attn_sink is not None + if attn_sink is None: + attn_sink = torch.empty(1, device=q.device, dtype=torch.float32) + else: + attn_sink = attn_sink.contiguous() + + num_queries, num_heads, head_dim = q.shape + assert indptr.numel() == num_queries + 1, ( + f"expected indptr shape [{num_queries + 1}], got {indptr.shape}" + ) + _validate_sparse_dims( + head_dim, + nope_head_dim, + rope_head_dim, + "_rocm_sparse_attn_decode_ragged_bf16_triton", + ) + if out is None: + out = torch.empty_like(q) + assert out.shape[-1] == head_dim, ( + f"expected out trailing dim {head_dim}, got {out.shape[-1]}" + ) + + block_h = _SPARSE_DECODE_BF16_BLOCK_H + block_d = triton.next_power_of_2(head_dim) + block_k = _SPARSE_DECODE_BF16_BLOCK_K + heads_blocks = triton.cdiv(num_heads, block_h) + comb_dim = nope_head_dim + rope_head_dim + + part_m = torch.empty( + (num_queries, num_splits, num_heads), dtype=torch.float32, device=q.device + ) + part_l = torch.empty_like(part_m) + part_acc = torch.empty( + (num_queries, num_splits, num_heads, comb_dim), + dtype=torch.float32, + device=q.device, + ) + + _sparse_attn_decode_ragged_bf16_partial_kernel[ + (num_queries, num_splits, heads_blocks) + ]( + q, + kv, + indices, + indptr, + part_m, + part_l, + part_acc, + q.stride(0), + q.stride(1), + q.stride(2), + kv.stride(0), + kv.stride(1), + part_m.stride(0), + part_m.stride(1), + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), + num_heads, + head_dim, + kv.shape[0], + float(scale), + BLOCK_H=block_h, + BLOCK_D=block_d, + BLOCK_K=block_k, + NUM_SPLITS=num_splits, + num_warps=4, + ) + + _sparse_attn_decode_reduce_kernel[(num_queries, num_heads)]( + part_m, + part_l, + part_acc, + attn_sink, + out, + None, + None, + None, + out.stride(0), + out.stride(1), + 0, + head_dim // 32, + part_m.stride(0), + part_m.stride(1), + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), + 0, + num_heads, + HAS_ATTN_SINK=has_attn_sink, + ADAPTIVE_SPLITS=False, + COMB_DIM=comb_dim, + BLOCK_H=1, + NUM_SPLITS=num_splits, + SPLITS_PAD=triton.next_power_of_2(num_splits), + FUSE_INV_ROPE=False, + NOPE=nope_head_dim, + HALF=rope_head_dim // 2, + QUANT_OUT=False, + num_warps=4, + ) + return out + + def _rocm_sparse_attn_decode_ragged_triton( q: torch.Tensor, main_cache: torch.Tensor, @@ -4143,6 +4432,72 @@ def rocm_sparse_attn_prefill( ) +def rocm_sparse_attn_decode_bf16( + q: torch.Tensor, + kv: torch.Tensor, + scale: float, + head_dim: int, + nope_head_dim: int, + rope_head_dim: int, + attn_sink: torch.Tensor | None, + output: torch.Tensor, + ragged_indices: torch.Tensor, + ragged_indptr: torch.Tensor, + num_splits: int, +) -> None: + """Run split-K sparse attention over decode rows using an unquantized KV cache. + + Args: + q: Decode queries laid out as ``[sq, h, d]``. + kv: KV cache laid out as ``[skv, 1, d]``. + scale: Softmax scale. + head_dim: Post-absorption head width. + nope_head_dim: NoPE width of ``head_dim``. + rope_head_dim: RoPE width of ``head_dim``. + attn_sink: Optional per-head sink logits. + output: Destination, written in place. + ragged_indices: Flattened per-query KV slots. + ragged_indptr: Segment offsets into ``ragged_indices``, ``[sq + 1]``. + num_splits: KV splits per query, from + :func:`rocm_sparse_decode_bf16_num_splits`. + + """ + assert kv.ndim == 3 and kv.shape[1] == 1, ( + f"ROCm Triton sparse decode expects kv=[skv,1,d], got {kv.shape}" + ) + _validate_sparse_dims( + head_dim, + nope_head_dim, + rope_head_dim, + "rocm_sparse_attn_decode_bf16", + ) + num_queries, num_heads = q.shape[0], q.shape[1] + direct = output.shape[-1] == head_dim + out = ( + output + if direct + else torch.empty( + (num_queries, num_heads, head_dim), + dtype=output.dtype, + device=output.device, + ) + ) + _rocm_sparse_attn_decode_ragged_bf16_triton( + q=q, + kv=kv.squeeze(1), + indices=ragged_indices, + indptr=ragged_indptr, + scale=scale, + attn_sink=None if attn_sink is None else attn_sink[: q.shape[1]], + nope_head_dim=nope_head_dim, + rope_head_dim=rope_head_dim, + num_splits=num_splits, + out=out, + ) + if not direct: + output.copy_(out[..., : output.shape[-1]]) + + def rocm_sparse_attn_decode( q: torch.Tensor, kv_cache: torch.Tensor | None, From 8352b2427f704652da1910f0d53f3fdcc2fa2b0d Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Thu, 8 Oct 2026 11:38:09 +0800 Subject: [PATCH 060/624] [Multimodal] Optimize encoder cudagraph padding logic (#58925) Signed-off-by: Isotr0py Co-authored-by: Shanshan Shen <467638484@qq.com> --- tests/v1/cudagraph/test_encoder_cudagraph.py | 132 ++++++++++++++++++ vllm/model_executor/models/kimi_k25.py | 34 ++--- .../deepseek_v41/common/vl_cudagraph.py | 29 +--- vllm/utils/math_utils.py | 8 ++ vllm/v1/worker/encoder_cudagraph.py | 51 ++++--- 5 files changed, 188 insertions(+), 66 deletions(-) diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index 749f9aa20dee..348944387886 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/tests/v1/cudagraph/test_encoder_cudagraph.py @@ -245,6 +245,138 @@ def test_num_graphs_to_capture_tracks_budgets(self): assert mgr.get_num_graphs_to_capture() == 3 +# --------------------------------------------------------------------------- +# _execute_local greedy packing with budget-cliff deferral (no GPU required) +# --------------------------------------------------------------------------- + + +def _run_packing( + mgr: EncoderCudaGraphManager, specs: list[EncoderItemSpec] +) -> list[tuple[list[int], str, int]]: + """Run _execute_local with graph replay and postprocess mocked out. + + Returns: + (batch_indices, path, token_budget) for every graph replay. + + """ + mgr.model = MagicMock() + + runs: list[tuple[list[int], str, int]] = [] + + def fake_run(mm_kwargs, token_budget, path="default", axis_keys=()): + runs.append((list(mm_kwargs["indices"]), path, token_budget)) + n_tokens = sum( + specs[i].get_path_output_tokens(path) for i in mm_kwargs["indices"] + ) + return torch.zeros(n_tokens, 32) + + def fake_postprocess( + graph_outputs, + batch_indices, + per_item_out_tokens, + outputs_by_orig_idx, + clone, + batch_mm_kwargs, + ): + for i in batch_indices: + outputs_by_orig_idx[i] = torch.zeros(per_item_out_tokens[i], 32) + + mgr.model.postprocess_encoder_output = fake_postprocess + with ( + patch.object(mgr, "_get_item_specs", lambda mm_kwargs: specs), + patch.object( + mgr, + "_select_items", + lambda mm_kwargs, indices: ({"indices": indices}, ()), + ), + patch.object(mgr, "_run_budget_graph", fake_run), + ): + result = mgr._execute_local({}) + assert len(result) == len(specs) + return runs + + +class TestExecuteLocalPacking: + """Greedy packing defers items that cross a budget cliff.""" + + @staticmethod + def _spec(tokens: int) -> EncoderItemSpec: + return EncoderItemSpec(input_size=tokens, output_tokens=tokens) + + def test_cliff_deferral_splits_batch(self): + # Merging would force 4200 tokens into the 8192 budget; deferring + # the 4000-token item costs 256 + 4096 instead. + mgr = _make_manager_with_budgets([256, 4096, 8192]) + runs = _run_packing(mgr, [self._spec(4000), self._spec(200)]) + assert runs == [([1], "default", 256), ([0], "default", 4096)] + + def test_tie_does_not_defer(self): + # 4096 + 4096 ties the merged 8192 budget; merging wins because it + # saves one replay. + mgr = _make_manager_with_budgets([256, 4096, 8192]) + runs = _run_packing(mgr, [self._spec(4000), self._spec(4000)]) + assert runs == [([0, 1], "default", 8192)] + + def test_deferral_requires_item_outweighing_batch(self): + # The third 1366-token item would save budget if deferred (4096+2048 + # < 8192) but is smaller than the current batch (4096), so it merges: + # deferring mid-size items fragments flood packing into extra replays. + mgr = _make_manager_with_budgets([256, 2048, 4096, 8192]) + runs = _run_packing(mgr, [self._spec(1366)] * 3) + assert runs == [([0, 1, 2], "default", 8192)] + + def test_packing_without_cliff_unchanged(self): + mgr = _make_manager_with_budgets([256, 4096]) + runs = _run_packing(mgr, [self._spec(100), self._spec(100)]) + assert runs == [([0, 1], "default", 256)] + + def test_small_items_split_to_avoid_cliff(self): + # Merging 3+ items of 100 tokens crosses 256 -> 4096; splitting + # into pairs costs 256 + 256 < 4096. + mgr = _make_manager_with_budgets([256, 4096]) + runs = _run_packing(mgr, [self._spec(100)] * 4) + assert runs == [([0, 1], "default", 256), ([2, 3], "default", 256)] + + def test_max_batch_size_split_unchanged(self): + mgr = _make_manager_with_budgets([1024, 4096]) + mgr.max_batch_size = 3 + runs = _run_packing(mgr, [self._spec(100)] * 4) + assert runs == [([0, 1, 2], "default", 1024), ([3], "default", 1024)] + + def test_oversized_item_still_falls_back_to_eager(self): + mgr = _make_manager_with_budgets([256, 4096]) + runs = _run_packing(mgr, [self._spec(100), self._spec(9000)]) + assert runs == [([0], "default", 256)] + assert mgr.graph_misses == 1 + + def test_multi_path_deferral_skips_zero_token_paths(self): + mgr = _make_manager_with_budgets([256, 4096, 8192]) + mgr.path_token_budgets = { + "a": [0, 256, 4096, 8192], + "b": [256, 4096, 8192], + } + specs = [ + EncoderItemSpec( + input_size=300, + output_tokens=300, + path_output_tokens={"a": 200, "b": 100}, + ), + EncoderItemSpec( + input_size=4000, + output_tokens=4000, + path_output_tokens={"a": 0, "b": 4000}, + ), + ] + runs = _run_packing(mgr, specs) + # Path b alone triggers deferral (256 + 4096 < 8192); the deferred + # item has zero tokens on path a, so that path is skipped. + assert runs == [ + ([0], "a", 256), + ([0], "b", 256), + ([1], "b", 4096), + ] + + # --------------------------------------------------------------------------- # get_cumulative_stats # --------------------------------------------------------------------------- diff --git a/vllm/model_executor/models/kimi_k25.py b/vllm/model_executor/models/kimi_k25.py index d66bed3642bc..8cd354b903af 100644 --- a/vllm/model_executor/models/kimi_k25.py +++ b/vllm/model_executor/models/kimi_k25.py @@ -479,30 +479,16 @@ def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]: # Image-only (t == 1). Video chunks (t > 1) fall back to eager. # # ------------------------------------------------------------------ # - @property - def _encoder_cudagraph_pad_totals(self) -> dict[int, int]: - """Row count of each captured buffer set, keyed by cu_seqlens ptr.""" - totals = self.__dict__.get("_encoder_cg_pad_totals") - if totals is None: - totals = {} - self.__dict__["_encoder_cg_pad_totals"] = totals - return totals - def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig": from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig - pad_totals = self._encoder_cudagraph_pad_totals - def pad_cu_seqlens(dst: torch.Tensor, src: torch.Tensor) -> None: - # Varlen attention requires cu_seqlens[-1] to equal the number of - # rows actually passed in. The captured buffers are sized for the - # full token budget, so a smaller real batch has to be completed - # with one trailing padding sequence; declaring fewer rows than the - # buffer holds is undefined behaviour and returns NaN on FlashAttn. - total = pad_totals.get(dst.data_ptr()) + # Zero-length padding (same as Qwen): tail slots repeat the last + # real offset, so FlashAttn schedules no work for padding rows. n = min(src.shape[0], dst.shape[0]) dst[:n].copy_(src[:n]) - dst[n:] = total if total is not None else src[-1] + if n < dst.shape[0]: + dst[n:] = src[-1] if n else 0 return EncoderCudaGraphConfig( modalities=["vision_chunk"], @@ -633,8 +619,8 @@ def prepare_encoder_cudagraph_capture_inputs( # max_seqlen must cover the worst case: one item consuming the full # budget, i.e. token_budget * kh * kw patches. - # max_batch_size + 1 leaves a spare cu_seqlens slot so replay can append - # a padding sequence covering rows the real batch does not fill. + # max_batch_size + 1 leaves a spare cu_seqlens slot so replay can + # pad the tail with zero-length sequences. metadata = self.vision_tower.prepare_encoder_cudagraph_metadata( grid_thw_list, max_batch_size=max_batch_size + 1, @@ -645,10 +631,6 @@ def prepare_encoder_cudagraph_capture_inputs( values: dict[str, torch.Tensor] = {"pixel_values": dummy_pixel_values} values.update({k: v for k, v in metadata.items() if v is not None}) - cu_seqlens = values.get("cu_seqlens") - if cu_seqlens is not None: - self._encoder_cudagraph_pad_totals[cu_seqlens.data_ptr()] = total_patches - return EncoderCudaGraphCaptureInputs(values=values) def prepare_encoder_cudagraph_replay_buffers( @@ -663,8 +645,8 @@ def prepare_encoder_cudagraph_replay_buffers( grid_thw_list = self._get_grid_thw_list(mm_kwargs) pixel_values = mm_kwargs["pixel_values"] - # Unpadded: pad_cu_seqlens completes the tail with the padding sequence - # so cu_seqlens[-1] matches the captured buffer's row count. + # Unpadded: pad_cu_seqlens pads the tail with zero-length sequences + # whose entries repeat the real total. metadata = self.vision_tower.prepare_encoder_cudagraph_metadata( grid_thw_list, max_batch_size=None, diff --git a/vllm/models/deepseek_v41/common/vl_cudagraph.py b/vllm/models/deepseek_v41/common/vl_cudagraph.py index bd8ed5a3fd9a..19baf339eb17 100644 --- a/vllm/models/deepseek_v41/common/vl_cudagraph.py +++ b/vllm/models/deepseek_v41/common/vl_cudagraph.py @@ -94,8 +94,6 @@ def _get_grid_list(mm_kwargs: dict[str, Any], key: str) -> list[list[int]]: grid = grid.tolist() return [[int(v) for v in row] for row in grid] - _encoder_cg_pad_totals: dict[int, int] = {} - # -- SupportsEncoderCudaGraph protocol -- def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig: @@ -129,24 +127,13 @@ def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig: "cudagraph_mm_encoder." ) - pad_totals = self._encoder_cg_pad_totals - def pad_cu_seqlens(dst: torch.Tensor, src: torch.Tensor) -> None: - # Varlen attention requires cu_seqlens[-1] to equal the number of - # rows actually passed in. The captured buffers are sized for the - # full token budget, so a smaller real batch is completed with one - # trailing padding sequence; declaring fewer rows than the buffer - # holds is undefined behaviour and returns NaN on FlashAttn. - total = pad_totals.get(dst.data_ptr()) - if total is None: - raise RuntimeError( - "cu_seqlens replay buffer was not registered at capture " - "time; the manager must replay into the capture-time " - "buffers." - ) + # Zero-length padding (same as Qwen): tail slots repeat the last + # real offset, so FlashAttn schedules no work for padding rows. n = min(src.shape[0], dst.shape[0]) dst[:n].copy_(src[:n]) - dst[n:] = total + if n < dst.shape[0]: + dst[n:] = src[-1] if n else 0 return EncoderCudaGraphConfig( modalities=["image"], @@ -268,8 +255,8 @@ def prepare_encoder_cudagraph_capture_inputs( max_seqlen_override=total_patches, cached=False, ) - # Spare cu_seqlens slots let replay append a padding sequence that - # covers the rows a smaller real batch does not fill. + # Spare cu_seqlens slots: replay pads smaller real batches with + # zero-length sequences (tail slots repeat the last real offset). real_cu = metadata.pop("cu_seqlens") cu_seqlens = torch.full( (max_batch_size + 2,), total_patches, dtype=torch.int32, device=device @@ -278,8 +265,6 @@ def prepare_encoder_cudagraph_capture_inputs( merge = build_packed_merge_metadata(grids, r, device=device, dtype=dtype) - self._encoder_cg_pad_totals[cu_seqlens.data_ptr()] = total_patches - return EncoderCudaGraphCaptureInputs( values={ "patches": patches, @@ -303,7 +288,7 @@ def prepare_encoder_cudagraph_replay_buffers( patches = patches.to(dtype) # Unpadded: the manager zero-pads patches/cos/sin/merge buffers, and - # pad_cu_seqlens appends the padding sequence covering the tail rows. + # pad_cu_seqlens pads the tail with zero-length sequences. metadata = build_packed_vit_metadata( vit_grid, rope_dim=self.vision.rope_dim, diff --git a/vllm/utils/math_utils.py b/vllm/utils/math_utils.py index 8fb0009d60bf..9854249af28b 100644 --- a/vllm/utils/math_utils.py +++ b/vllm/utils/math_utils.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Math utility functions for vLLM.""" +from bisect import bisect_left + # Approximate value of 1/ln(2), used for log/exp base conversion # Best FP32 approximation: 1.4426950216 (hex 0x3FB8AA3B) RCP_LN2 = 1.4426950216 @@ -30,3 +32,9 @@ def round_down(x: int, y: int) -> int: def largest_power_of_2_divisor(n: int) -> int: """Return the largest power-of-2 that divides *n* (isolate lowest set bit).""" return n & (-n) + + +def round_up_to(values: list[int], n: int) -> int | None: + """Round `n` up to the smallest member of sorted `values` >= n, or None.""" + i = bisect_left(values, n) + return values[i] if i < len(values) else None diff --git a/vllm/v1/worker/encoder_cudagraph.py b/vllm/v1/worker/encoder_cudagraph.py index 3b587ad71f80..51c832001ca8 100644 --- a/vllm/v1/worker/encoder_cudagraph.py +++ b/vllm/v1/worker/encoder_cudagraph.py @@ -23,6 +23,7 @@ from vllm.model_executor.models.utils import scatter_output_slices from vllm.model_executor.models.vision import get_load_balance_assignment from vllm.utils.gpu_sync_debug import gpu_sync_allowed +from vllm.utils.math_utils import round_up_to from vllm.utils.torch_utils import current_stream from vllm.v1.worker.encoder_cudagraph_defs import ( ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG, @@ -347,17 +348,17 @@ def _capture_budget_graph( def _find_smallest_fitting_budget_given_tokens( self, total_tokens: int, budgets: list[int] | None = None ) -> int | None: - """Find smallest budget >= total_tokens. - - Returns: - Token budget if found, None if no fitting budget. - - """ + """Find smallest budget >= total_tokens, or None if none fits.""" budgets = budgets if budgets is not None else self.token_budgets - for budget in budgets: - if budget >= total_tokens: - return budget - return None + return round_up_to(budgets, total_tokens) + + def _path_fitting_budget(self, path: str, tokens: int) -> int | None: + """Smallest captured budget on a path; 0 when idle, None if oversized.""" + if tokens == 0: + return 0 + return self._find_smallest_fitting_budget_given_tokens( + tokens, self.path_token_budgets[path] + ) def _get_item_specs(self, mm_kwargs: dict[str, Any]) -> list[EncoderItemSpec]: """Get item specs from the model.""" @@ -479,24 +480,38 @@ def append_current_batch() -> None: if not current_batch: return path_budgets = { - path: ( - 0 - if current_tokens[path] == 0 - else self._find_smallest_fitting_budget_given_tokens( - current_tokens[path], self.path_token_budgets[path] - ) - ) + path: self._path_fitting_budget(path, current_tokens[path]) for path in paths } batches.append((list(current_batch), path_budgets)) + def should_defer_item(item_tokens: dict[str, int]) -> bool: + """Whether closing the batch strictly reduces summed budgets.""" + for path in paths: + item = item_tokens[path] + if item == 0: + continue + cur = current_tokens[path] + b_cur = self._path_fitting_budget(path, cur) + b_alone = self._path_fitting_budget(path, item) + b_merged = self._path_fitting_budget(path, cur + item) + # Only called when the item fits; lookups land on a budget. + assert ( + b_cur is not None and b_alone is not None and b_merged is not None + ) + # Only defer items that alone outweigh the whole current + # batch; splitting smaller ones fragments flood packing. + if b_cur + b_alone < b_merged and b_alone >= b_cur: + return True + return False + for orig_idx in sorted_indices: item_tokens = {path: per_item_path_tokens[path][orig_idx] for path in paths} fits = len(current_batch) < self.max_batch_size and all( current_tokens[path] + item_tokens[path] <= max_path_budgets[path] for path in paths ) - if current_batch and not fits: + if current_batch and (not fits or should_defer_item(item_tokens)): append_current_batch() current_batch = [] current_tokens = dict.fromkeys(paths, 0) From 5dd4a628dde34d5b5737067f131a2499f672f684 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Wed, 7 Oct 2026 21:26:12 -0700 Subject: [PATCH 061/624] [ROCm][Bugfix] Map sleep-mode cuMem memory for the host (#60367) Signed-off-by: Matthew Wong Co-authored-by: Claude Opus 5.5 (1M context) --- csrc/cumem_allocator.cpp | 23 +++++++++++++++---- .../memory/sleep_mode/test_sleep_mode.py | 1 - 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/csrc/cumem_allocator.cpp b/csrc/cumem_allocator.cpp index 720c990450b6..6965d2124a65 100644 --- a/csrc/cumem_allocator.cpp +++ b/csrc/cumem_allocator.cpp @@ -203,12 +203,25 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem, } #endif - CUmemAccessDesc accessDesc = {}; - accessDesc.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - accessDesc.location.id = device; - accessDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; +#if defined(USE_ROCM) && HIP_VERSION >= 70200000 + // Also map for the host, as PyTorch does for its expandable segments: on + // ROCm >= 7.2, Tensor.item() on a large-BAR device dereferences VRAM + // directly from the host, which segfaults on a device-only mapping. + constexpr int num_desc = 2; +#else + constexpr int num_desc = 1; +#endif + CUmemAccessDesc accessDesc[num_desc] = {}; + accessDesc[0].location.type = CU_MEM_LOCATION_TYPE_DEVICE; + accessDesc[0].location.id = device; + accessDesc[0].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; +#if defined(USE_ROCM) && HIP_VERSION >= 70200000 + accessDesc[1].location.type = hipMemLocationTypeHost; + accessDesc[1].location.id = 0; // ignored + accessDesc[1].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; +#endif - CUDA_CHECK(cuMemSetAccess(d_mem, size, &accessDesc, 1)); + CUDA_CHECK(cuMemSetAccess(d_mem, size, accessDesc, num_desc)); if (error_code != 0) { return; } diff --git a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py index 73af6e39103d..9f6180335ab1 100644 --- a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py +++ b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py @@ -323,7 +323,6 @@ def _attn_kv_scales(model) -> dict[str, list[float]]: } -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Segfaults on ROCm.") @requires_fp8 @create_new_process_for_each_test() def test_deep_sleep_compressed_tensors_kv_scales( From 1fbbfa165709c490dae48884eebdfbf6a1aef330 Mon Sep 17 00:00:00 2001 From: Hyukjoon Lee Date: Thu, 8 Oct 2026 13:33:32 +0900 Subject: [PATCH 062/624] [Bugfix][Kimi-K3] Declare max_tp_shards on the DSpark MLA KV cache spec (#59733) Signed-off-by: hyukjlee Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Summer Yang --- tests/models/test_dspark_mla.py | 33 +++++++++++++++++++++++++++++++ vllm/models/kimi_k3/nvidia/mla.py | 1 + 2 files changed, 34 insertions(+) diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index 76af0eb60c10..53a93b940cd5 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -336,3 +336,36 @@ def __call__(self, main_x): assert torch.equal(calls[1][1], stacked_output.view(2, 3, 4)[:, 2] + 2) assert calls[0][3] is slot_mappings[0] assert calls[1][3] is slot_mappings[2] + + +@pytest.mark.cpu_test +def test_k3_dspark_mla_kv_cache_spec_groups_with_target_mla(): + """The draft's MLA layers must share a KV cache group with the target's.""" + from vllm.model_executor.layers.attention.mla_attention import MLAAttention + from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention + from vllm.v1.core.kv_cache_utils import _get_kv_cache_groups_uniform_page_size + + vllm_config = SimpleNamespace( + model_config=None, cache_config=SimpleNamespace(block_size=64) + ) + target_attn = SimpleNamespace( + kv_cache_dtype="fp8", + head_size=576, + sliding_window=None, + indexer=None, + non_causal_multi_token_decode=False, + attn_backend=SimpleNamespace(get_name=lambda: "ROCM_AITER_MLA"), + _uses_flat_kv_cache=lambda: False, + ) + draft_attn = SimpleNamespace( + kv_cache_dtype="fp8", head_size=576, non_causal_multi_token_decode=True + ) + target_spec = MLAAttention.get_kv_cache_spec(target_attn, vllm_config) + draft_spec = MultiHeadLatentAttention.get_kv_cache_spec(draft_attn, vllm_config) + + kv_cache_spec = {f"model.layers.{i}.attn": target_spec for i in range(24)} + kv_cache_spec |= {f"draft.layers.{i}.attn": draft_spec for i in range(5)} + + groups = _get_kv_cache_groups_uniform_page_size(kv_cache_spec) + assert len(groups) == 1 + assert len(groups[0].layer_names) == 29 diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index bf7c53e9e515..18fd93579be7 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -422,6 +422,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: return MLAAttentionSpec( # type: ignore[call-arg] block_size=vllm_config.cache_config.block_size, num_kv_heads=1, + max_tp_shards=1, head_size=self.head_size, dtype=kv_cache_dtype, cache_dtype_str=self.kv_cache_dtype, From 4187ade7b87e4c683ff60e41ad27cc6659f8739d Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Thu, 8 Oct 2026 12:37:11 +0800 Subject: [PATCH 063/624] [Frontend] Consolidate rl entrypoints. (#57849) Signed-off-by: wang.yuqi Signed-off-by: aoshen02 Co-authored-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- .github/CODEOWNERS | 2 +- .../test_route_request_bodies.py} | 4 +- .../rl/weight_transfer_metrics}/__init__.py | 0 .../test_across_api_servers_metrics.py} | 0 .../test_weight_operation_metrics.py} | 20 +++--- vllm/distributed/weight_transfer/clients.py | 2 +- .../launchers/api_server/routers.py | 4 ++ vllm/entrypoints/llm.py | 63 ++---------------- vllm/entrypoints/rl/__init__.py | 0 vllm/entrypoints/rl/factories.py | 11 ++++ vllm/entrypoints/rl/offline.py | 64 +++++++++++++++++++ vllm/entrypoints/rl/online/__init__.py | 0 .../dev/rlhf => rl/online}/api_router.py | 8 +-- .../{serve/dev/rlhf => rl/online}/metrics.py | 0 vllm/entrypoints/serve/__init__.py | 4 -- 15 files changed, 99 insertions(+), 83 deletions(-) rename tests/entrypoints/{unit_tests/test_dev_route_request_bodies.py => rl/return_routed_experts/test_route_request_bodies.py} (96%) rename {vllm/entrypoints/serve/dev/rlhf => tests/entrypoints/rl/weight_transfer_metrics}/__init__.py (100%) rename tests/entrypoints/rl/{test_weight_operation_metrics.py => weight_transfer_metrics/test_across_api_servers_metrics.py} (100%) rename tests/entrypoints/{unit_tests/test_rl_weight_operation_metrics.py => rl/weight_transfer_metrics/test_weight_operation_metrics.py} (90%) create mode 100644 vllm/entrypoints/rl/__init__.py create mode 100644 vllm/entrypoints/rl/factories.py create mode 100644 vllm/entrypoints/rl/offline.py create mode 100644 vllm/entrypoints/rl/online/__init__.py rename vllm/entrypoints/{serve/dev/rlhf => rl/online}/api_router.py (97%) rename vllm/entrypoints/{serve/dev/rlhf => rl/online}/metrics.py (100%) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 083842c932ff..677e37ff8b60 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -229,7 +229,7 @@ mkdocs.yaml @hmellor # Reinforcement learning /vllm/distributed/weight_transfer @aoshen02 -/vllm/entrypoints/serve/dev/rlhf @njhill @aoshen02 +/vllm/entrypoints/rl @njhill @aoshen02 /tests/distributed/test_weight_transfer.py @aoshen02 /tests/entrypoints/rl @DarkLight1337 @robertgshaw2-redhat @njhill @aarnphm @NickLucche @AndreasKaratzas @aoshen02 /tests/v1/worker/test_gpu_worker_weight_transfer.py @aoshen02 diff --git a/tests/entrypoints/unit_tests/test_dev_route_request_bodies.py b/tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py similarity index 96% rename from tests/entrypoints/unit_tests/test_dev_route_request_bodies.py rename to tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py index 78fe5eef5dde..b51d08db0758 100644 --- a/tests/entrypoints/unit_tests/test_dev_route_request_bodies.py +++ b/tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py @@ -9,7 +9,7 @@ from fastapi import FastAPI from fastapi.testclient import TestClient -from vllm.entrypoints.serve.dev.rlhf.api_router import router as rlhf_router +from vllm.entrypoints.rl.online.api_router import router as rl_router from vllm.entrypoints.serve.dev.rpc.api_router import router as rpc_router from vllm.entrypoints.serve.exception_handling.register import init_exception_handler @@ -30,7 +30,7 @@ async def call(*args, **kwargs): @pytest.fixture def client(): app = FastAPI() - app.include_router(rlhf_router) + app.include_router(rl_router) app.include_router(rpc_router) init_exception_handler(app) app.state.engine_client = Engine() diff --git a/vllm/entrypoints/serve/dev/rlhf/__init__.py b/tests/entrypoints/rl/weight_transfer_metrics/__init__.py similarity index 100% rename from vllm/entrypoints/serve/dev/rlhf/__init__.py rename to tests/entrypoints/rl/weight_transfer_metrics/__init__.py diff --git a/tests/entrypoints/rl/test_weight_operation_metrics.py b/tests/entrypoints/rl/weight_transfer_metrics/test_across_api_servers_metrics.py similarity index 100% rename from tests/entrypoints/rl/test_weight_operation_metrics.py rename to tests/entrypoints/rl/weight_transfer_metrics/test_across_api_servers_metrics.py diff --git a/tests/entrypoints/unit_tests/test_rl_weight_operation_metrics.py b/tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py similarity index 90% rename from tests/entrypoints/unit_tests/test_rl_weight_operation_metrics.py rename to tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py index bdddf3d247ba..22ec5e611252 100644 --- a/tests/entrypoints/unit_tests/test_rl_weight_operation_metrics.py +++ b/tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py @@ -13,9 +13,9 @@ from fastapi.testclient import TestClient from prometheus_client import CollectorRegistry -from vllm.entrypoints.serve.dev.rlhf import api_router -from vllm.entrypoints.serve.dev.rlhf import metrics as rlhf_metrics -from vllm.entrypoints.serve.dev.rlhf.metrics import WeightOperationMetrics +from vllm.entrypoints.rl.online import api_router +from vllm.entrypoints.rl.online import metrics as rl_metrics +from vllm.entrypoints.rl.online.metrics import WeightOperationMetrics from vllm.entrypoints.serve.exception_handling.register import init_exception_handler from vllm.v1.metrics.prometheus import get_prometheus_registry @@ -52,13 +52,13 @@ async def call(*args, **kwargs): @pytest.fixture def registry(monkeypatch): registry = CollectorRegistry() - monkeypatch.setattr(rlhf_metrics, "_metrics", WeightOperationMetrics(registry)) + monkeypatch.setattr(rl_metrics, "_metrics", WeightOperationMetrics(registry)) return registry @pytest.mark.parametrize("fail", [False, True]) def test_record_releases_in_flight_and_observes_duration(registry, fail): - metrics = rlhf_metrics.weight_operation_metrics() + metrics = rl_metrics.weight_operation_metrics() with ( pytest.raises(RuntimeError) if fail else nullcontext(), metrics.record("update"), @@ -108,7 +108,7 @@ async def test_cancelled_route_releases_in_flight(registry): HOLD_OPERATION = """ import sys -from vllm.entrypoints.serve.dev.rlhf.metrics import weight_operation_metrics +from vllm.entrypoints.rl.online.metrics import weight_operation_metrics with weight_operation_metrics().record("update"): print("entered", flush=True) sys.stdin.readline() @@ -163,19 +163,19 @@ def test_default_recorder_is_created_lazily_on_the_default_registry(): probe = """ from prometheus_client import REGISTRY -from vllm.entrypoints.serve.dev.rlhf import metrics as rlhf_metrics +from vllm.entrypoints.rl.online import metrics as rl_metrics NAMES = [ "vllm:rl_weight_update_operation_duration_seconds", "vllm:rl_weight_update_operations_in_flight", ] -assert rlhf_metrics._metrics is None +assert rl_metrics._metrics is None assert not set(NAMES) & set(REGISTRY._names_to_collectors) -metrics = rlhf_metrics.weight_operation_metrics() +metrics = rl_metrics.weight_operation_metrics() assert REGISTRY._names_to_collectors[NAMES[0]] is metrics.duration assert REGISTRY._names_to_collectors[NAMES[1]] is metrics.in_flight -assert rlhf_metrics.weight_operation_metrics() is metrics +assert rl_metrics.weight_operation_metrics() is metrics print("lazy registration ok") """ result = subprocess.run( diff --git a/vllm/distributed/weight_transfer/clients.py b/vllm/distributed/weight_transfer/clients.py index 1bd2f98937a9..745e67779e0a 100644 --- a/vllm/distributed/weight_transfer/clients.py +++ b/vllm/distributed/weight_transfer/clients.py @@ -58,7 +58,7 @@ def _json_safe_update_payload( class HTTPVLLMWeightSyncClient: """Talks to a vLLM server over the RLHF HTTP routes. - Mirrors `vllm/entrypoints/serve/dev/rlhf/api_router.py`: + Mirrors `vllm/entrypoints/rl/online/api_router.py`: `/init_weight_transfer_engine`, `/start_weight_update`, `/update_weights`, `/finish_weight_update`. """ diff --git a/vllm/entrypoints/launchers/api_server/routers.py b/vllm/entrypoints/launchers/api_server/routers.py index 39d670c019b7..40d04862afb2 100644 --- a/vllm/entrypoints/launchers/api_server/routers.py +++ b/vllm/entrypoints/launchers/api_server/routers.py @@ -36,6 +36,10 @@ def register_api_routers( register_vllm_dev_api_routers(app) + from vllm.entrypoints.rl.factories import register_rl_api_routers + + register_rl_api_routers(app) + if "generate" in supported_tasks: from vllm.entrypoints.generate.api_router import ( register_generate_api_routers, diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 45b6aede583c..c38789a2e29c 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -28,10 +28,6 @@ TokenizerMode, ) from vllm.config.quantization import QuantizationConfigArgs -from vllm.distributed.weight_transfer.base import ( - WeightTransferInitRequest, - WeightTransferUpdateRequest, -) from vllm.engine.arg_utils import EngineArgs from vllm.entrypoints.chat_utils import ( ChatCompletionMessageParam, @@ -40,6 +36,7 @@ ) from vllm.entrypoints.generate.beam_search.offline import BeamSearchOfflineMixin from vllm.entrypoints.pooling.offline import PoolingOfflineMixin +from vllm.entrypoints.rl.offline import RLOfflineMixin from vllm.entrypoints.serve.utils.api_utils import log_non_default_args from vllm.inputs import PromptType from vllm.logger import configure_logging_if_needed, init_logger @@ -64,7 +61,9 @@ logger = init_logger(__name__) -class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin): +class LLM( + BeamSearchOfflineMixin, PoolingOfflineMixin, RLOfflineMixin, OfflineInferenceMixin +): """An LLM for generating texts from given prompts and sampling parameters. This class includes a tokenizer, a language model (possibly distributed @@ -884,60 +883,6 @@ def get_metrics(self) -> list["Metric"]: """ return self.llm_engine.get_metrics() - def init_weight_transfer_engine( - self, request: WeightTransferInitRequest | dict - ) -> None: - """Initialize weight transfer for RL training. - - Args: - request: Weight transfer initialization request with backend-specific info - - """ - init_info_dict = ( - request["init_info"] if isinstance(request, dict) else request.init_info - ) - - self.llm_engine.collective_rpc( - "init_weight_transfer_engine", kwargs={"init_info": init_info_dict} - ) - - def start_weight_update(self) -> None: - """Start a new weight update.""" - self.llm_engine.collective_rpc("start_weight_update") - - def start_draft_weight_update(self) -> None: - """Start a new weight update targeting the speculative draft model.""" - self.llm_engine.collective_rpc("start_draft_weight_update") - - def update_weights(self, request: WeightTransferUpdateRequest | dict) -> None: - """Update the weights of the model. - - Args: - request: Weight update request with backend-specific update info - - """ - update_info_dict = ( - request["update_info"] if isinstance(request, dict) else request.update_info - ) - - self.llm_engine.collective_rpc( - "update_weights", kwargs={"update_info": update_info_dict} - ) - - def finish_weight_update(self, weight_version: str | None = None) -> None: - """Finish the weight update and set its version if provided.""" - self.llm_engine.collective_rpc("finish_weight_update") - if weight_version is not None: - self.llm_engine.set_weight_version(weight_version) - - def update_weight_version(self, new_version: str) -> None: - """Set the weight version without updating weights.""" - self.llm_engine.set_weight_version(new_version) - - def get_weight_version(self) -> str: - """Return the latest committed weight version.""" - return self.llm_engine.get_weight_version() - def __repr__(self) -> str: """Return a transformers-style hierarchical view of the model.""" # Cache the result to avoid repeated collective_rpc calls diff --git a/vllm/entrypoints/rl/__init__.py b/vllm/entrypoints/rl/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/entrypoints/rl/factories.py b/vllm/entrypoints/rl/factories.py new file mode 100644 index 000000000000..a5d237ec0bf6 --- /dev/null +++ b/vllm/entrypoints/rl/factories.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +from fastapi import FastAPI + + +def register_rl_api_routers(app: FastAPI): + from .online.api_router import router as rl_router + + app.include_router(rl_router) diff --git a/vllm/entrypoints/rl/offline.py b/vllm/entrypoints/rl/offline.py new file mode 100644 index 000000000000..02fc63104d8f --- /dev/null +++ b/vllm/entrypoints/rl/offline.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.distributed.weight_transfer.base import ( + WeightTransferInitRequest, + WeightTransferUpdateRequest, +) +from vllm.entrypoints.offline_utils import OfflineInferenceMixin + + +class RLOfflineMixin(OfflineInferenceMixin): + def init_weight_transfer_engine( + self, request: WeightTransferInitRequest | dict + ) -> None: + """Initialize weight transfer for RL training. + + Args: + request: Weight transfer initialization request with backend-specific info + + """ + init_info_dict = ( + request["init_info"] if isinstance(request, dict) else request.init_info + ) + + self.llm_engine.collective_rpc( + "init_weight_transfer_engine", kwargs={"init_info": init_info_dict} + ) + + def start_weight_update(self) -> None: + """Start a new weight update.""" + self.llm_engine.collective_rpc("start_weight_update") + + def start_draft_weight_update(self) -> None: + """Start a new weight update targeting the speculative draft model.""" + self.llm_engine.collective_rpc("start_draft_weight_update") + + def update_weights(self, request: WeightTransferUpdateRequest | dict) -> None: + """Update the weights of the model. + + Args: + request: Weight update request with backend-specific update info + + """ + update_info_dict = ( + request["update_info"] if isinstance(request, dict) else request.update_info + ) + + self.llm_engine.collective_rpc( + "update_weights", kwargs={"update_info": update_info_dict} + ) + + def finish_weight_update(self, weight_version: str | None = None) -> None: + """Finish the weight update and set its version if provided.""" + self.llm_engine.collective_rpc("finish_weight_update") + if weight_version is not None: + self.llm_engine.set_weight_version(weight_version) + + def update_weight_version(self, new_version: str) -> None: + """Set the weight version without updating weights.""" + self.llm_engine.set_weight_version(new_version) + + def get_weight_version(self) -> str: + """Return the latest committed weight version.""" + return self.llm_engine.get_weight_version() diff --git a/vllm/entrypoints/rl/online/__init__.py b/vllm/entrypoints/rl/online/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/entrypoints/serve/dev/rlhf/api_router.py b/vllm/entrypoints/rl/online/api_router.py similarity index 97% rename from vllm/entrypoints/serve/dev/rlhf/api_router.py rename to vllm/entrypoints/rl/online/api_router.py index 7e52e512170c..ef9269749f3b 100644 --- a/vllm/entrypoints/serve/dev/rlhf/api_router.py +++ b/vllm/entrypoints/rl/online/api_router.py @@ -5,7 +5,7 @@ from http import HTTPStatus from typing import Annotated, Any, Literal -from fastapi import APIRouter, Body, FastAPI, HTTPException, Query, Request +from fastapi import APIRouter, Body, HTTPException, Query, Request from fastapi.responses import JSONResponse from vllm.distributed.weight_transfer.base import ( @@ -14,7 +14,7 @@ WeightTransferUpdateRequest, ) from vllm.engine.protocol import EngineClient -from vllm.entrypoints.serve.dev.rlhf.metrics import weight_operation_metrics +from vllm.entrypoints.rl.online.metrics import weight_operation_metrics from vllm.logger import init_logger from vllm.v1.engine import PauseMode @@ -280,7 +280,3 @@ async def get_world_size( else: world_size = parallel_config.world_size return JSONResponse(content={"world_size": world_size}) - - -def attach_router(app: FastAPI): - app.include_router(router) diff --git a/vllm/entrypoints/serve/dev/rlhf/metrics.py b/vllm/entrypoints/rl/online/metrics.py similarity index 100% rename from vllm/entrypoints/serve/dev/rlhf/metrics.py rename to vllm/entrypoints/rl/online/metrics.py diff --git a/vllm/entrypoints/serve/__init__.py b/vllm/entrypoints/serve/__init__.py index c358b9a6db26..d8254cb33548 100644 --- a/vllm/entrypoints/serve/__init__.py +++ b/vllm/entrypoints/serve/__init__.py @@ -50,10 +50,6 @@ def register_vllm_dev_api_routers(app: FastAPI): attach_cache_router(app) - from .dev.rlhf.api_router import attach_router as attach_rlhf_router - - attach_rlhf_router(app) - from .dev.rpc.api_router import attach_router as attach_rpc_router attach_rpc_router(app) From 775792f2d5cdbdb0dcf433943259744b17452a7e Mon Sep 17 00:00:00 2001 From: aniskumar-nv Date: Wed, 7 Oct 2026 21:40:55 -0700 Subject: [PATCH 064/624] [Bugfix] Reland #53020: tie lm_head.weight for Nemotron Parse when checkpoint omits it (#60272) Signed-off-by: Anish Kumar Co-authored-by: Cyrus Leung Co-authored-by: Claude Fable 5.1 --- .../multimodal/generation/test_nemotron_parse.py | 11 ++++++++++- vllm/model_executor/models/nemotron_parse.py | 9 +++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/models/multimodal/generation/test_nemotron_parse.py b/tests/models/multimodal/generation/test_nemotron_parse.py index 69e41b21e342..57e40430c63c 100644 --- a/tests/models/multimodal/generation/test_nemotron_parse.py +++ b/tests/models/multimodal/generation/test_nemotron_parse.py @@ -104,7 +104,16 @@ def run_test( ) -@pytest.mark.parametrize("model", ["nvidia/NVIDIA-Nemotron-Parse-v1.2"]) +@pytest.mark.parametrize( + "model", + [ + "nvidia/NVIDIA-Nemotron-Parse-v1.2", + # v2.0's checkpoint ties lm_head to the decoder's input embeddings + # instead of shipping a separate lm_head.weight tensor; this + # exercises that path (see tie_weights() in nemotron_parse.py). + "nvidia/NVIDIA-Nemotron-Parse-2.0", + ], +) @pytest.mark.parametrize("dtype", ["bfloat16"]) @pytest.mark.parametrize("num_logprobs", [5]) def test_models( diff --git a/vllm/model_executor/models/nemotron_parse.py b/vllm/model_executor/models/nemotron_parse.py index 9ac38c0021db..3ba731d2e0f2 100644 --- a/vllm/model_executor/models/nemotron_parse.py +++ b/vllm/model_executor/models/nemotron_parse.py @@ -568,6 +568,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.lm_head = ParallelLMHead( config.decoder.vocab_size, config.decoder.d_model, quant_config=quant_config ) + # Some checkpoints (e.g. compact exports) tie the output head to the + # decoder's input embeddings instead of materializing a separate + # lm_head.weight tensor. + tie_word_embeddings = bool( + getattr(config, "tie_word_embeddings", False) + or getattr(config.decoder, "tie_word_embeddings", False) + ) + if tie_word_embeddings: + self.lm_head = self.lm_head.tie_weights(self.decoder.embed_tokens) self.logits_processor = LogitsProcessor( self.vocab_size, config.decoder.vocab_size ) From 4ac0179f8a4c6f459cd73dc6ac2e0bd22a5ad294 Mon Sep 17 00:00:00 2001 From: Mehdi Ghanimifard Date: Thu, 8 Oct 2026 06:49:49 +0200 Subject: [PATCH 065/624] [ROCm][Perf] Use native merge_attn_states kernel on gfx942/gfx950 (#60159) Signed-off-by: mghanimi Co-authored-by: Cursor Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> --- vllm/v1/attention/ops/merge_attn_states.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/vllm/v1/attention/ops/merge_attn_states.py b/vllm/v1/attention/ops/merge_attn_states.py index d28c78bdb9fa..b16786a13b47 100644 --- a/vllm/v1/attention/ops/merge_attn_states.py +++ b/vllm/v1/attention/ops/merge_attn_states.py @@ -139,8 +139,19 @@ def supported_headdim(prefix: torch.Tensor) -> bool: return headdim % 4 == 0 return headdim % 8 == 0 + # CDNA3 and newer (gfx942, gfx950, gfx1250). gfx90a is CDNA2 and + # stays on the Triton merge. + def supported_platform() -> bool: + if current_platform.is_cuda(): + return True + if current_platform.is_rocm(): + from vllm.platforms.rocm import get_cdna_version + + return get_cdna_version() >= 3 + return False + if ( - current_platform.is_cuda() + supported_platform() and supported_dtypes(prefix_output) and supported_headdim(prefix_output) ): From 81198e97ba7eee2a22540caaa756b7fdddcb4d93 Mon Sep 17 00:00:00 2001 From: Hui Shi <70356247+huishi-hs@users.noreply.github.com> Date: Thu, 8 Oct 2026 13:05:39 +0800 Subject: [PATCH 066/624] [Refactor] Abstract DeepGemm MegaMOE backend (#58053) Signed-off-by: Hui Shi Co-authored-by: LuneExplorer Co-authored-by: ArmageddonKnight --- tests/models/test_deepseek_v4_mega_moe.py | 15 +- .../layers/fused_moe/deep_gemm_mega_moe.py | 466 ++++++++++++++++++ vllm/models/deepseek_v4/nvidia/model.py | 306 ++++-------- .../deepseek_v4/nvidia/ops/prepare_megamoe.py | 107 +++- vllm/models/kimi_k3/nvidia/model.py | 53 +- 5 files changed, 681 insertions(+), 266 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 104ad8e89514..76915b94a32d 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -6,6 +6,9 @@ import pytest import torch +from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import ( + DeepGemmSm100MegaMoEBackend, +) from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( bind_routed_experts_capturer, ) @@ -451,7 +454,9 @@ def test_deepseek_v4_mega_moe_preserves_checkpoint_dimensions( ): param.data.random_(1, 128) originals.append(param.data.clone()) - monkeypatch.setattr(experts, "_check_runtime_supported", lambda: None) + monkeypatch.setattr( + experts, "_ensure_backend", lambda: DeepGemmSm100MegaMoEBackend() + ) monkeypatch.setattr( "vllm.utils.deep_gemm._import_deep_gemm", lambda: SimpleNamespace( @@ -547,7 +552,9 @@ def transform_weights_for_mega_moe(cls, l1_weights, l2_weights): intermediate_size=intermediate_size, num_shared_experts=1, ) - experts._check_runtime_supported = lambda: None + monkeypatch.setattr( + experts, "_ensure_backend", lambda: DeepGemmSm100MegaMoEBackend() + ) def fp8_parameter(*shape): return torch.nn.Parameter( @@ -664,6 +671,10 @@ def forward(self, hidden_states): class FakeExperts(torch.nn.Module): has_fused_shared_experts = fused + _backend = DeepGemmSm100MegaMoEBackend() + + def _ensure_backend(self): + return self._backend def forward(self, hidden_states, *args, **kwargs): return torch.ones_like(hidden_states) diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py b/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py new file mode 100644 index 000000000000..633ced108208 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py @@ -0,0 +1,466 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepGEMM MegaMoE backends: fused EP dispatch, expert GEMMs and combine.""" + +from inspect import signature +from typing import Any + +import torch +import torch.nn as nn + +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp8Dynamic, +) +from vllm.platforms import current_platform + +logger = init_logger(__name__) + + +def ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor: + """Reinterpret a uint8 UE8M0 exponent tensor as float32 scale values.""" + return (sf.to(torch.int32) << 23).view(torch.float32) + + +class DeepGemmMegaMoEBackend: + """Abstract MegaMoE backend that hides deep_gemm-specific details. + + A backend advertises its expected hidden-state quantization layout via + ``hidden_quant`` and its expected MMA type via ``mma_type``. It is + responsible for the routed-expert weight transform run at load time, the + optional shared-expert weight transform, and the actual MegaMoE kernel + invocation at inference time. Runtime capability checks live in + :func:`get_deep_gemm_mega_moe_backend`: obtaining a backend already implies + that the current device and shapes are supported. + """ + + # Human-readable MMA type tag used by deep_gemm helpers. SM100 default. + mma_type: str = "fp8xfp4" + hidden_quant: QuantKey = kMxfp8Dynamic + + def transform_weights( + self, + *, + w13_weight: torch.Tensor, + w13_weight_scale: torch.Tensor, + w2_weight: torch.Tensor, + w2_weight_scale: torch.Tensor, + num_local_experts: int, + hidden_size: int, + intermediate_size: int, + activation: str | None = None, + ) -> tuple[ + tuple[torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor], + ]: + raise NotImplementedError + + def supports_shared_experts(self, deep_gemm) -> bool: + return False + + def transform_shared_expert_weights( + self, + *, + shared_experts: nn.Module, + num_shared_experts: int, + hidden_size: int, + intermediate_size: int, + prefix: str, + ) -> ( + tuple[ + tuple[torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor], + ] + | None + ): + """Return ``(l1, l2)`` transformed shared-expert weights, or ``None``. + + ``shared_experts`` must expose ``gate_up_proj`` and ``down_proj`` + linear layers. A ``None`` return means the backend cannot fuse shared + experts for this module; the caller should fall back to a separate + shared MLP. + """ + return None + + def run_mega_moe( + self, + *, + y: torch.Tensor, + l1_weights: tuple[torch.Tensor, torch.Tensor], + l2_weights: tuple[torch.Tensor, torch.Tensor], + symm_buffer, + activation_clamp: float | None, + fast_math: bool, + activation: str | None = None, + activation_alpha: float | None = None, + activation_beta: float | None = None, + shared_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None, + shared_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> None: + raise NotImplementedError + + def supports_bf16_mega_gate(self) -> bool: + """Whether this backend can run the DeepGEMM bf16 MegaMoE gate kernel. + + Defaults to ``False`` (the symbol is not guaranteed to exist in every + deep_gemm build). Backends that target a deep_gemm build exposing the + ``bf16_mega_gate`` symbol override this and :meth:`bf16_mega_gate`. + """ + return False + + def bf16_mega_gate( + self, + *, + x: torch.Tensor, + weight: torch.Tensor, + num_topk: int, + scoring_func: str, + routed_scaling_factor: float, + ep_rank: int, + bias: torch.Tensor | None = None, + image_bias: torch.Tensor | None = None, + image_token_mask: torch.Tensor | None = None, + fix_routing_mask: torch.Tensor | None = None, + unmapped_topk_idx: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + raise NotImplementedError + + +class DeepGemmSm100MegaMoEBackend(DeepGemmMegaMoEBackend): + """Default MegaMoE backend targeting SM100 via deep_gemm fp8_fp4_mega_moe.""" + + mma_type = "fp8xfp4" + hidden_quant = kMxfp8Dynamic + + def transform_weights( + self, + *, + w13_weight: torch.Tensor, + w13_weight_scale: torch.Tensor, + w2_weight: torch.Tensor, + w2_weight_scale: torch.Tensor, + num_local_experts: int, + hidden_size: int, + intermediate_size: int, + activation: str | None = None, + ) -> tuple[ + tuple[torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor], + ]: + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + + w13_scale = deep_gemm.transform_sf_into_required_layout( + ue8m0_uint8_to_float(w13_weight_scale).contiguous(), + 2 * intermediate_size, + hidden_size, + (1, 32), + num_local_experts, + ) + w2_scale = deep_gemm.transform_sf_into_required_layout( + ue8m0_uint8_to_float(w2_weight_scale).contiguous(), + hidden_size, + intermediate_size, + (1, 32), + num_local_experts, + ) + kwargs = {} if activation is None else {"activation": activation} + return deep_gemm.transform_weights_for_mega_moe( + (w13_weight.view(torch.int8).contiguous(), w13_scale), + (w2_weight.view(torch.int8).contiguous(), w2_scale), + **kwargs, + ) + + def supports_shared_experts(self, deep_gemm) -> bool: + """Check the Python API before touching a symmetric-memory group. + + This also gives users of an older precompiled vLLM wheel a safe serial + fallback instead of failing halfway through multi-rank buffer setup. + """ + try: + buffer_params = signature(deep_gemm.get_symm_buffer_for_mega_moe).parameters + kernel_params = signature(deep_gemm.fp8_fp4_mega_moe).parameters + except (TypeError, ValueError): + return False + return ( + hasattr(deep_gemm, "get_block_m_for_mega_moe") + and hasattr(deep_gemm, "transform_weights_for_mega_moe") + and "num_shared_experts" in buffer_params + and "shared_l1_weights" in kernel_params + and "shared_l2_weights" in kernel_params + ) + + def transform_shared_expert_weights( + self, + *, + shared_experts: nn.Module, + num_shared_experts: int, + hidden_size: int, + intermediate_size: int, + prefix: str, + ) -> ( + tuple[ + tuple[torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor], + ] + | None + ): + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + + gate_up = shared_experts.gate_up_proj + down = shared_experts.down_proj + gate_up_weight = gate_up.weight.data + gate_up_scale = ( + gate_up.weight_scale + if hasattr(gate_up, "weight_scale") + else gate_up.weight_scale_inv + ).data + down_weight = down.weight.data + down_scale = ( + down.weight_scale + if hasattr(down, "weight_scale") + else down.weight_scale_inv + ).data + + # MegaMoE's shared FP8 MMA consumes a 1x32 scale for every weight row, + # while the checkpoint uses coarser block-FP8 scales (usually + # 128x128). Build a dedicated, numerically equivalent scale view. + ue8m0_scale_dtypes = (torch.float8_e8m0fnu, torch.uint8, torch.int32) + if ( + gate_up_scale.dtype in ue8m0_scale_dtypes + and down_scale.dtype in ue8m0_scale_dtypes + ): + gate_up_scale = self._prepare_shared_expert_scale( + deep_gemm, + gate_up, + gate_up_scale, + gate_up_weight.shape[0], + gate_up_weight.shape[1], + prefix, + ) + down_scale = self._prepare_shared_expert_scale( + deep_gemm, + down, + down_scale, + down_weight.shape[0], + down_weight.shape[1], + prefix, + ) + + if gate_up_scale is None or down_scale is None: + return None + + shared_intermediate_size = intermediate_size * num_shared_experts + expected_gate_up_shape = ( + 2 * shared_intermediate_size, + hidden_size, + ) + expected_down_shape = (hidden_size, shared_intermediate_size) + if ( + gate_up_weight.dtype != torch.float8_e4m3fn + or down_weight.dtype != torch.float8_e4m3fn + or gate_up_scale.dtype != torch.int32 + or down_scale.dtype != torch.int32 + or tuple(gate_up_weight.shape) != expected_gate_up_shape + or tuple(down_weight.shape) != expected_down_shape + ): + logger.warning( + "Disabling native MegaMoE shared-expert fusion for %s: expected " + "replicated block-FP8 weights with gate_up=%s, down=%s, and " + "DeepGEMM int32 scales; got gate_up=%s/%s/%s and down=%s/%s/%s.", + prefix, + expected_gate_up_shape, + expected_down_shape, + tuple(gate_up_weight.shape), + gate_up_weight.dtype, + gate_up_scale.dtype, + tuple(down_weight.shape), + down_weight.dtype, + down_scale.dtype, + ) + return None + + transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe( + (gate_up_weight, gate_up_scale), + (down_weight, down_scale), + ) + # L1 interleaving allocates a full copy. Re-home the loader Parameter on + # that storage so the original 2*intermediate*hidden FP8 tensor can be + # released instead of adding roughly 0.7 GiB per rank on DSV4-Flash. + # The generic linear post-load hook may still repack the serial scales, + # but this shared MLP is never called after native fusion is enabled. + gate_up.weight.data = transformed_l1[0] + transformed_l1 = (gate_up.weight.data, transformed_l1[1]) + return transformed_l1, transformed_l2 + + @staticmethod + def _prepare_shared_expert_scale( + deep_gemm, + linear: nn.Module, + scale: torch.Tensor, + mn: int, + k: int, + prefix: str, + ) -> torch.Tensor | None: + block_size = getattr(linear, "weight_block_size", None) + if block_size is None or len(block_size) != 2: + logger.warning( + "Disabling native MegaMoE shared-expert fusion for %s: " + "shared FP8 weight block size is unavailable.", + prefix, + ) + return None + + block_m, block_k = block_size + num_k_blocks = (k + block_k - 1) // block_k + # The linear hook's DeepGEMM layout packs four k-block bytes per int32. + packed = scale.dtype == torch.int32 + expected_shape = ( + (mn, (num_k_blocks + 3) // 4) + if packed + else ((mn + block_m - 1) // block_m, num_k_blocks) + ) + if block_k % 32 != 0 or tuple(scale.shape) != expected_shape: + logger.warning( + "Disabling native MegaMoE shared-expert fusion for %s: " + "cannot convert shared scale shape %s with block size %s " + "to MegaMoE's 1x32 layout for weight (%d, %d).", + prefix, + tuple(scale.shape), + tuple(block_size), + mn, + k, + ) + return None + + if packed: + row_scale = scale.flatten().view(torch.uint8).view(mn, -1)[:, :num_k_blocks] + else: + row_scale = scale.view(torch.uint8).repeat_interleave(block_m, dim=0)[:mn] + scale_1x32 = ( + ue8m0_uint8_to_float(row_scale) + .repeat_interleave(block_k // 32, dim=1)[:, : k // 32] + .contiguous() + ) + # The grouped API is used with a singleton dimension to request the + # MN-major, TMA-aligned packed-UE8M0 strides, then squeezed back to the + # 2D layout required for a shared expert. + return deep_gemm.transform_sf_into_required_layout( + scale_1x32.unsqueeze(0), + mn, + k, + (1, 32), + 1, + ).squeeze(0) + + def run_mega_moe( + self, + *, + y: torch.Tensor, + l1_weights: tuple[torch.Tensor, torch.Tensor], + l2_weights: tuple[torch.Tensor, torch.Tensor], + symm_buffer, + activation_clamp: float | None, + fast_math: bool, + activation: str | None = None, + activation_alpha: float | None = None, + activation_beta: float | None = None, + shared_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None, + shared_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> None: + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + kwargs: dict[str, Any] = { + name: value + for name, value in ( + ("activation", activation), + ("activation_alpha", activation_alpha), + ("activation_beta", activation_beta), + ) + if value is not None + } + if shared_l1_weights is not None and shared_l2_weights is not None: + kwargs["shared_l1_weights"] = shared_l1_weights + kwargs["shared_l2_weights"] = shared_l2_weights + deep_gemm.fp8_fp4_mega_moe( + y, + l1_weights, + l2_weights, + symm_buffer, + activation_clamp=activation_clamp, + fast_math=fast_math, + **kwargs, + ) + + def supports_bf16_mega_gate(self) -> bool: + """SM100 deep_gemm build is expected to provide the bf16 gate kernel.""" + return True + + def bf16_mega_gate( + self, + *, + x: torch.Tensor, + weight: torch.Tensor, + num_topk: int, + scoring_func: str, + routed_scaling_factor: float, + ep_rank: int, + bias: torch.Tensor | None = None, + image_bias: torch.Tensor | None = None, + image_token_mask: torch.Tensor | None = None, + fix_routing_mask: torch.Tensor | None = None, + unmapped_topk_idx: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + from vllm.utils.deep_gemm import bf16_mega_gate + + return bf16_mega_gate( + x, + weight, + num_topk, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor, + ep_rank=ep_rank, + bias=bias, + image_bias=image_bias, + image_token_mask=image_token_mask, + fix_routing_mask=fix_routing_mask, + unmapped_topk_idx=unmapped_topk_idx, + ) + + +def get_deep_gemm_mega_moe_backend( + device: torch.device, + hidden_size: int, + intermediate_size: int, +) -> DeepGemmMegaMoEBackend: + """Return the DeepGEMM MegaMoE backend for ``device``. + + A successful return implies the backend can run on this device with the + given ``hidden_size`` / ``intermediate_size``. All runtime capability + checks (arch + shape divisibility) live here; there is no separate check + on the returned backend. + + Currently only the SM100 deep_gemm backend is registered here. SM90 and + other backends are expected to be provided by new backends or plugins + that override this factory. + + Raises: + NotImplementedError: If the device architecture is unsupported. + ValueError: If the shapes are unsupported. + + """ + if not current_platform.is_device_capability_family( + 100, device_id=device.index or 0 + ): + raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.") + if hidden_size % 128 != 0 or intermediate_size % 128 != 0: + raise ValueError( + "DeepGEMM MegaMoE requires hidden and intermediate sizes " + "to be multiples of 128." + ) + return DeepGemmSm100MegaMoEBackend() diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index b9996d007c99..fc10cd07c1cf 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import typing from collections.abc import Callable, Iterable -from inspect import signature from itertools import islice import regex as re @@ -33,6 +32,11 @@ FusedMoEFactory, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import ( + DeepGemmMegaMoEBackend, + get_deep_gemm_mega_moe_backend, + ue8m0_uint8_to_float, +) from vllm.model_executor.layers.fused_moe.router.base_router import ( eplb_map_to_physical_and_record, ) @@ -205,7 +209,9 @@ def make_deepseek_v4_expert_params_mapping( class DeepseekV4MegaMoEExperts(nn.Module): - _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int, int], object] = {} + _symm_buffer_cache: dict[ + tuple[int, int, int, int, int, int, int, int, str], object + ] = {} def __init__( self, @@ -296,6 +302,8 @@ def __init__( tuple[torch.Tensor, torch.Tensor] | None ) = None + self._backend: DeepGemmMegaMoEBackend | None = None + # Register in the static forward context so the custom-op wrapper # can look up this module by name from within a torch.compile graph. compilation_config = vllm_config.compilation_config @@ -357,214 +365,71 @@ def weight_loader( @staticmethod def _ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor: - return (sf.to(torch.int32) << 23).view(torch.float32) + return ue8m0_uint8_to_float(sf) - def _check_runtime_supported(self) -> None: - device = self.w13_weight.device - if torch.cuda.get_device_capability(device)[0] != 10: - raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.") - if self.hidden_size % 128 != 0 or self.intermediate_size % 128 != 0: - raise ValueError( - "DeepGEMM MegaMoE requires hidden and intermediate sizes " - "to be multiples of 128." - ) - - @staticmethod - def _deep_gemm_supports_shared_experts(deep_gemm) -> bool: - """Check the Python API before touching a symmetric-memory group. + def _ensure_backend(self) -> DeepGemmMegaMoEBackend: + """Lazily resolve the MegaMoE backend. - This also gives users of an older precompiled vLLM wheel a safe serial - fallback instead of failing halfway through multi-rank buffer setup. + The backend is looked up via :func:`get_deep_gemm_mega_moe_backend`, which + performs all runtime capability checks. Failure propagates directly + to the caller. This must be called before :meth:`finalize_weights` + releases the loader-side parameters. """ - try: - buffer_params = signature(deep_gemm.get_symm_buffer_for_mega_moe).parameters - kernel_params = signature(deep_gemm.fp8_fp4_mega_moe).parameters - except (TypeError, ValueError): - return False - return ( - hasattr(deep_gemm, "get_block_m_for_mega_moe") - and hasattr(deep_gemm, "transform_weights_for_mega_moe") - and "num_shared_experts" in buffer_params - and "shared_l1_weights" in kernel_params - and "shared_l2_weights" in kernel_params + if self._backend is not None: + return self._backend + assert self.w13_weight is not None, ( + "_ensure_backend() must be called before finalize_weights() " + "releases the loader-side parameters." ) - - def _finalize_shared_expert_weights( - self, deep_gemm, shared_experts: DeepseekV4MLP - ) -> None: - gate_up = shared_experts.gate_up_proj - down = shared_experts.down_proj - gate_up_weight = gate_up.weight.data - gate_up_scale = ( - gate_up.weight_scale - if hasattr(gate_up, "weight_scale") - else gate_up.weight_scale_inv - ).data - down_weight = down.weight.data - down_scale = ( - down.weight_scale - if hasattr(down, "weight_scale") - else down.weight_scale_inv - ).data - - # MegaMoE's shared FP8 MMA consumes a 1x32 scale for every weight row, - # while the checkpoint uses coarser block-FP8 scales (usually - # 128x128). Build a dedicated, numerically equivalent scale view. - ue8m0_scale_dtypes = (torch.float8_e8m0fnu, torch.uint8, torch.int32) - if ( - gate_up_scale.dtype in ue8m0_scale_dtypes - and down_scale.dtype in ue8m0_scale_dtypes - ): - gate_up_scale = self._prepare_shared_expert_scale( - deep_gemm, - gate_up, - gate_up_scale, - gate_up_weight.shape[0], - gate_up_weight.shape[1], - ) - down_scale = self._prepare_shared_expert_scale( - deep_gemm, - down, - down_scale, - down_weight.shape[0], - down_weight.shape[1], - ) - - if gate_up_scale is None or down_scale is None: - self.num_shared_experts = 0 - return - - shared_intermediate_size = self.intermediate_size * self.num_shared_experts - expected_gate_up_shape = ( - 2 * shared_intermediate_size, + self._backend = get_deep_gemm_mega_moe_backend( + self.w13_weight.device, self.hidden_size, + self.intermediate_size, ) - expected_down_shape = (self.hidden_size, shared_intermediate_size) - if ( - gate_up_weight.dtype != torch.float8_e4m3fn - or down_weight.dtype != torch.float8_e4m3fn - or gate_up_scale.dtype != torch.int32 - or down_scale.dtype != torch.int32 - or tuple(gate_up_weight.shape) != expected_gate_up_shape - or tuple(down_weight.shape) != expected_down_shape - ): - logger.warning( - "Disabling native MegaMoE shared-expert fusion for %s: expected " - "replicated block-FP8 weights with gate_up=%s, down=%s, and " - "DeepGEMM int32 scales; got gate_up=%s/%s/%s and down=%s/%s/%s.", - self.prefix, - expected_gate_up_shape, - expected_down_shape, - tuple(gate_up_weight.shape), - gate_up_weight.dtype, - gate_up_scale.dtype, - tuple(down_weight.shape), - down_weight.dtype, - down_scale.dtype, - ) - self.num_shared_experts = 0 - return + return self._backend - transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe( - (gate_up_weight, gate_up_scale), - (down_weight, down_scale), - ) - # L1 interleaving allocates a full copy. Re-home the loader Parameter on - # that storage so the original 2*intermediate*hidden FP8 tensor can be - # released instead of adding roughly 0.7 GiB per rank on DSV4-Flash. - # The generic linear post-load hook may still repack the serial scales, - # but this shared MLP is never called after native fusion is enabled. - gate_up.weight.data = transformed_l1[0] - self._transformed_shared_l1_weights = ( - gate_up.weight.data, - transformed_l1[1], - ) - self._transformed_shared_l2_weights = transformed_l2 + def _check_runtime_supported(self) -> None: + """Preserved for API compatibility; delegates to the backend check.""" + self._ensure_backend() - def _prepare_shared_expert_scale( - self, - deep_gemm, - linear: nn.Module, - scale: torch.Tensor, - mn: int, - k: int, - ) -> torch.Tensor | None: - block_size = getattr(linear, "weight_block_size", None) - if block_size is None or len(block_size) != 2: - logger.warning( - "Disabling native MegaMoE shared-expert fusion for %s: " - "shared FP8 weight block size is unavailable.", - self.prefix, - ) - return None - - block_m, block_k = block_size - num_k_blocks = (k + block_k - 1) // block_k - # The linear hook's DeepGEMM layout packs four k-block bytes per int32. - packed = scale.dtype == torch.int32 - expected_shape = ( - (mn, (num_k_blocks + 3) // 4) - if packed - else ((mn + block_m - 1) // block_m, num_k_blocks) - ) - if block_k % 32 != 0 or tuple(scale.shape) != expected_shape: - logger.warning( - "Disabling native MegaMoE shared-expert fusion for %s: " - "cannot convert shared scale shape %s with block size %s " - "to MegaMoE's 1x32 layout for weight (%d, %d).", - self.prefix, - tuple(scale.shape), - tuple(block_size), - mn, - k, - ) - return None + def _finalize_shared_expert_weights(self, shared_experts: DeepseekV4MLP) -> None: + """Delegate shared-expert weight transform to the backend and update state. - if packed: - row_scale = scale.flatten().view(torch.uint8).view(mn, -1)[:, :num_k_blocks] - else: - row_scale = scale.view(torch.uint8).repeat_interleave(block_m, dim=0)[:mn] - scale_1x32 = ( - self._ue8m0_uint8_to_float(row_scale) - .repeat_interleave(block_k // 32, dim=1)[:, : k // 32] - .contiguous() + On success, populates ``self._transformed_shared_l{1,2}_weights``. If + the backend cannot fuse shared experts (returns ``None``), disables + native fusion by setting ``self.num_shared_experts = 0``. + """ + backend = self._ensure_backend() + result = backend.transform_shared_expert_weights( + shared_experts=shared_experts, + num_shared_experts=self.num_shared_experts, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + prefix=self.prefix, ) - # The grouped API is used with a singleton dimension to request the - # MN-major, TMA-aligned packed-UE8M0 strides, then squeezed back to the - # 2D layout required for a shared expert. - return deep_gemm.transform_sf_into_required_layout( - scale_1x32.unsqueeze(0), - mn, - k, - (1, 32), - 1, - ).squeeze(0) + if result is None: + self.num_shared_experts = 0 + return + ( + self._transformed_shared_l1_weights, + self._transformed_shared_l2_weights, + ) = result def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None: + backend = self._ensure_backend() from vllm.utils.deep_gemm import _import_deep_gemm deep_gemm = _import_deep_gemm() - if self._transformed_l1_weights is None: - self._check_runtime_supported() - w13_scale = deep_gemm.transform_sf_into_required_layout( - self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(), - 2 * self.intermediate_size, - self.hidden_size, - (1, 32), - self.num_local_experts, - ) - w2_scale = deep_gemm.transform_sf_into_required_layout( - self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(), - self.hidden_size, - self.intermediate_size, - (1, 32), - self.num_local_experts, - ) self._transformed_l1_weights, self._transformed_l2_weights = ( - deep_gemm.transform_weights_for_mega_moe( - (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale), - (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale), + backend.transform_weights( + w13_weight=self.w13_weight.data, + w13_weight_scale=self.w13_weight_scale.data, + w2_weight=self.w2_weight.data, + w2_weight_scale=self.w2_weight_scale.data, + num_local_experts=self.num_local_experts, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, ) ) # Drop the original loader-side parameters: the MegaMoE kernels only @@ -583,7 +448,7 @@ def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None: return if self._transformed_shared_l1_weights is not None: return - if not self._deep_gemm_supports_shared_experts(deep_gemm): + if not backend.supports_shared_experts(deep_gemm): logger.warning_once( "Disabling native MegaMoE shared-expert fusion because the " "installed DeepGEMM Python API is older than the vLLM " @@ -591,7 +456,8 @@ def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None: ) self.num_shared_experts = 0 return - self._finalize_shared_expert_weights(deep_gemm, shared_experts) + + self._finalize_shared_expert_weights(shared_experts) @property def has_fused_shared_experts(self) -> bool: @@ -602,6 +468,7 @@ def get_symm_buffer(self): deep_gemm = _import_deep_gemm() + backend = self._ensure_backend() group = get_ep_group().device_group device = torch.accelerator.current_device_index() key = ( @@ -613,6 +480,7 @@ def get_symm_buffer(self): self.hidden_size, self.intermediate_size, self.num_shared_experts if self.has_fused_shared_experts else 0, + backend.mma_type, ) symm_buffer = self._symm_buffer_cache.get(key) if symm_buffer is None: @@ -626,6 +494,7 @@ def get_symm_buffer(self): num_shared_experts=( self.num_shared_experts if self.has_fused_shared_experts else 0 ), + mma_type=backend.mma_type, ) self._symm_buffer_cache[key] = symm_buffer return symm_buffer @@ -700,7 +569,6 @@ def forward( deep_gemm = _import_deep_gemm() - symm_buffer = self.get_symm_buffer() num_tokens = hidden_states.shape[0] is_padding = None if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): @@ -732,6 +600,8 @@ def forward( else None, ) + backend = self._ensure_backend() + symm_buffer = self.get_symm_buffer() shared_x_sf = None shared_block_m = None if self.has_fused_shared_experts: @@ -742,7 +612,7 @@ def forward( symm_buffer.num_max_tokens_per_rank, num_tokens, self.top_k, - "fp8xfp4", + backend.mma_type, ) prepare_megamoe_inputs( @@ -756,27 +626,28 @@ def forward( is_padding=is_padding, shared_x_sf=shared_x_sf, shared_block_m=shared_block_m, + hidden_quant=backend.hidden_quant, ) assert self._transformed_l1_weights is not None assert self._transformed_l2_weights is not None if self.has_fused_shared_experts: - deep_gemm.fp8_fp4_mega_moe( - y, - self._transformed_l1_weights, - self._transformed_l2_weights, - symm_buffer, + backend.run_mega_moe( + y=y, + l1_weights=self._transformed_l1_weights, + l2_weights=self._transformed_l2_weights, + symm_buffer=symm_buffer, shared_l1_weights=self._transformed_shared_l1_weights, shared_l2_weights=self._transformed_shared_l2_weights, activation_clamp=activation_clamp, fast_math=fast_math, ) else: - deep_gemm.fp8_fp4_mega_moe( - y, - self._transformed_l1_weights, - self._transformed_l2_weights, - symm_buffer, + backend.run_mega_moe( + y=y, + l1_weights=self._transformed_l1_weights, + l2_weights=self._transformed_l2_weights, + symm_buffer=symm_buffer, activation_clamp=activation_clamp, fast_math=fast_math, ) @@ -1039,9 +910,18 @@ def forward( return self._forward_fused_moe(hidden_states, input_ids) org_shape = hidden_states.shape + # The MegaMoE experts module owns the backend; resolve and assert it. + self.experts._ensure_backend() + assert self.experts._backend is not None, ( + "MegaMoE backend must be resolved before routing." + ) + mega_backend = self.experts._backend # Small local padded batches favor GateLinear; 128-expert gates cross earlier. gate_threshold = 1 if self.gate.weight.shape[0] == 128 else 16 - if hidden_states.shape[0] <= gate_threshold: + if ( + not mega_backend.supports_bf16_mega_gate() + or hidden_states.shape[0] <= gate_threshold + ): router_logits, _ = self.gate(hidden_states) topk_weights, topk_ids = fused_topk_bias( hidden_states=hidden_states, @@ -1060,8 +940,6 @@ def forward( image_sentinel_lo=self.image_sentinel_lo if bias_vl is not None else 0, ) else: - from vllm.utils.deep_gemm import bf16_mega_gate - if mega_gate_metadata is None: raise RuntimeError( "MegaMoE routing metadata must be prepared once by the model." @@ -1079,10 +957,10 @@ def forward( ] fix_routing_mask = mega_gate_metadata.hash_token_mask - topk_weights, topk_ids = bf16_mega_gate( - hidden_states, - self.gate.weight, - self.n_activated_experts, + topk_weights, topk_ids = mega_backend.bf16_mega_gate( + x=hidden_states, + weight=self.gate.weight, + num_topk=self.n_activated_experts, scoring_func=self.scoring_func, routed_scaling_factor=self.routed_scaling_factor, ep_rank=self.ep_rank, diff --git a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py b/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py index aa9482bed643..2d93b0d43032 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py +++ b/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py @@ -9,8 +9,15 @@ import torch +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + FP8_DTYPE, + QuantKey, + kMxfp8Dynamic, +) from vllm.triton_utils import tl, triton +_E8M0_SCALE_DTYPES = (torch.uint8, torch.float8_e8m0fnu) + @triton.jit(do_not_specialize=["num_tokens"]) def _prepare_megamoe_inputs_kernel( @@ -48,6 +55,7 @@ def _prepare_megamoe_inputs_kernel( GROUP_K: tl.constexpr, BLOCK_TOPK: tl.constexpr, SHARED_BLOCK_M: tl.constexpr, + USE_UE8M0: tl.constexpr, ) -> None: token_id = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) token_mask = token_id < num_tokens @@ -69,15 +77,18 @@ def _prepare_megamoe_inputs_kernel( amax = tl.maximum(amax, 1.0e-4) scale = amax / 448.0 - scale_bits = scale.to(tl.uint32, bitcast=True) - scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to( - tl.uint32 - ) - scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254) - rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True) + if USE_UE8M0: + scale_bits = scale.to(tl.uint32, bitcast=True) + scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to( + tl.uint32 + ) + scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254) + quant_scale = (scale_exp << 23).to(tl.float32, bitcast=True) + else: + quant_scale = scale hidden_groups = tl.reshape(hidden, [BLOCK_M, num_groups, GROUP_K]) - scaled = hidden_groups * (1.0 / rounded_scale)[:, :, None] + scaled = hidden_groups * (1.0 / quant_scale)[:, :, None] scaled = tl.reshape(scaled, [BLOCK_M, BLOCK_K]) fp8 = scaled.to(tl.float8e4nv) tl.store( @@ -86,21 +97,32 @@ def _prepare_megamoe_inputs_kernel( mask=k_mask, ) - scale_offsets = tl.arange(0, num_groups) - packed_scale = tl.sum(scale_exp << (scale_offsets[None, :] * 8), axis=1).to( - tl.int32 - ) - tl.store( - x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k, - packed_scale, - mask=token_mask, - ) + if USE_UE8M0: + scale_offsets = tl.arange(0, num_groups) + packed_scale = tl.sum(scale_exp << (scale_offsets[None, :] * 8), axis=1).to( + tl.int32 + ) + tl.store( + x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k, + packed_scale, + mask=token_mask, + ) + else: + sf_offset = tl.arange(0, num_groups) + tl.store( + x_sf + + token_id[:, None] * x_sf_stride_m + + k_block_id * x_sf_stride_k + + sf_offset[None, :], + quant_scale, + mask=token_mask[:, None], + ) # DeepGEMM's SM100 shared-expert TMA loads require the activation scales # in an MN-major layout whose row permutation depends on the MegaMoE # scheduler's runtime BLOCK_M. Write that view while the packed UE8M0 scale # is already resident, avoiding another kernel and temporary tensor. - if shared_x_sf is not None: + if USE_UE8M0 and shared_x_sf is not None: m_block_id = token_id // SHARED_BLOCK_M m_in_block = token_id % SHARED_BLOCK_M aligned_block_m: tl.constexpr = triton.cdiv(SHARED_BLOCK_M, 128) * 128 @@ -160,6 +182,27 @@ def _prepare_megamoe_inputs_kernel( ) +def _resolve_hidden_quant(hidden_quant: QuantKey, block_k: int) -> tuple[int, bool]: + """Map a hidden-state ``QuantKey`` to the kernel's ``(GROUP_K, USE_UE8M0)``.""" + scale = hidden_quant.scale + group_shape = scale.group_shape + if ( + hidden_quant.dtype != FP8_DTYPE + or not hidden_quant.symmetric + or hidden_quant.scale2 is not None + or scale.static + or not group_shape.is_per_group() + or block_k % group_shape.col != 0 + or scale.dtype not in (*_E8M0_SCALE_DTYPES, torch.float32) + ): + raise ValueError( + "DeepSeek V4 MegaMoE input staging requires symmetric dynamic fp8 " + f"per-group quantization with a group size dividing {block_k} and " + f"E8M0 or fp32 scales, got {hidden_quant}." + ) + return group_shape.col, scale.dtype in _E8M0_SCALE_DTYPES + + def prepare_megamoe_inputs( hidden_states: torch.Tensor, topk_weights: torch.Tensor, @@ -171,7 +214,26 @@ def prepare_megamoe_inputs( is_padding: torch.Tensor | None = None, shared_x_sf: torch.Tensor | None = None, shared_block_m: int | None = None, + hidden_quant: QuantKey = kMxfp8Dynamic, ) -> None: + """Quantize hidden states and repack top-k routing for DeepGEMM MegaMoE. + + Args: + hidden_states: Input activations of shape ``[num_tokens, hidden_dim]``. + topk_weights: Router top-k weights of shape ``[num_tokens, top_k]``. + topk_ids: Router top-k expert ids of shape ``[num_tokens, top_k]``. + x_fp8: Output buffer for the fp8-quantized hidden states. + x_sf: Output buffer for the hidden-state scale factors. + topk_idx_out: Output buffer for the repacked top-k expert ids. + topk_weights_out: Output buffer for the repacked top-k weights. + is_padding: Optional per-token mask; padded tokens are not routed. + shared_x_sf: Optional output buffer for shared-expert scale factors. + Must be given together with ``shared_block_m``. + shared_block_m: Block M used to lay out ``shared_x_sf``. + hidden_quant: Hidden-state quantization scheme. E8M0 scales are packed + four per int32 into ``x_sf``; fp32 scales are stored directly. + + """ num_tokens, hidden_size = hidden_states.shape if num_tokens == 0: return @@ -180,6 +242,8 @@ def prepare_megamoe_inputs( "DeepSeek V4 MegaMoE input staging requires hidden_size to be " "a multiple of 128." ) + block_k = 128 + group_k, use_ue8m0 = _resolve_hidden_quant(hidden_quant, block_k) top_k = topk_ids.shape[1] if topk_weights.shape != topk_ids.shape: raise ValueError( @@ -191,6 +255,11 @@ def prepare_megamoe_inputs( "DeepSeek V4 MegaMoE shared input staging requires both " "shared_x_sf and shared_block_m." ) + if shared_x_sf is not None and not use_ue8m0: + raise ValueError( + "DeepSeek V4 MegaMoE shared input staging currently requires " + "UE8M0-packed hidden scales." + ) if shared_x_sf is not None: assert shared_block_m is not None if shared_block_m <= 0: @@ -209,7 +278,6 @@ def prepare_megamoe_inputs( f"{required_rows}, got {shared_x_sf.shape[0]}." ) - block_k = 128 # On GB200, eight-row tiles win from 64 tokens; keep smaller batches untiled. block_m = 8 if num_tokens >= 64 else 1 grid = (triton.cdiv(num_tokens, block_m), triton.cdiv(hidden_size, block_k)) @@ -247,8 +315,9 @@ def prepare_megamoe_inputs( top_k, BLOCK_M=block_m, BLOCK_K=block_k, - GROUP_K=32, + GROUP_K=group_k, BLOCK_TOPK=block_topk, SHARED_BLOCK_M=shared_block_m or 1, + USE_UE8M0=use_ue8m0, num_warps=4, ) diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index cc1ca609a6ba..84e93b82fbb4 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -365,27 +365,16 @@ def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None: return self._check_runtime_supported() - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - w13_scale = deep_gemm.transform_sf_into_required_layout( - self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(), - 2 * self.intermediate_size, - self.hidden_size, - (1, 32), - self.num_local_experts, - ) - w2_scale = deep_gemm.transform_sf_into_required_layout( - self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(), - self.hidden_size, - self.intermediate_size, - (1, 32), - self.num_local_experts, - ) + backend = self._ensure_backend() self._transformed_l1_weights, self._transformed_l2_weights = ( - deep_gemm.transform_weights_for_mega_moe( - (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale), - (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale), + backend.transform_weights( + w13_weight=self.w13_weight.data, + w13_weight_scale=self.w13_weight_scale.data, + w2_weight=self.w2_weight.data, + w2_weight_scale=self.w2_weight_scale.data, + num_local_experts=self.num_local_experts, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, activation=self.activation, ) ) @@ -407,6 +396,7 @@ def get_symm_buffer(self): from vllm.utils.deep_gemm import _import_deep_gemm deep_gemm = _import_deep_gemm() + backend = self._ensure_backend() group = get_ep_group().device_group device = torch.accelerator.current_device_index() key = ( @@ -418,6 +408,7 @@ def get_symm_buffer(self): self.hidden_size, self.intermediate_size, self.activation, + backend.mma_type, ) symm_buffer = self._kimi_symm_buffer_cache.get(key) if symm_buffer is None: @@ -429,6 +420,7 @@ def get_symm_buffer(self): self.hidden_size, self.intermediate_size, activation=self.activation, + mma_type=backend.mma_type, ) self._kimi_symm_buffer_cache[key] = symm_buffer return symm_buffer @@ -449,10 +441,6 @@ def forward( f"but its symmetric buffer supports {self.max_num_tokens}." ) y = torch.empty_like(hidden_states, dtype=torch.bfloat16) - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - symm_buffer = self.get_symm_buffer() num_tokens = hidden_states.shape[0] is_padding = None if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): @@ -483,6 +471,8 @@ def forward( else None, ) + backend = self._ensure_backend() + symm_buffer = self.get_symm_buffer() prepare_megamoe_inputs( hidden_states, topk_weights, @@ -492,22 +482,23 @@ def forward( symm_buffer.topk_idx[:num_tokens], symm_buffer.topk_weights[:num_tokens], is_padding=is_padding, + hidden_quant=backend.hidden_quant, ) self.finalize_weights() assert self._transformed_l1_weights is not None assert self._transformed_l2_weights is not None - deep_gemm.fp8_fp4_mega_moe( - y, - self._transformed_l1_weights, - self._transformed_l2_weights, - symm_buffer, + backend.run_mega_moe( + y=y, + l1_weights=self._transformed_l1_weights, + l2_weights=self._transformed_l2_weights, + symm_buffer=symm_buffer, activation_clamp=activation_clamp, - activation=self.activation, + fast_math=fast_math, # DeepGEMM names the SiTU gate tanh scale `activation_alpha` and the # linear/up tanh scale `activation_beta`; beta=0 leaves up untouched. + activation=self.activation, activation_alpha=self.activation_beta or 1.0, activation_beta=self.activation_linear_beta or 0.0, - fast_math=fast_math, ) return y From 816173ca9d9256f85c36d09d133a7ed219ccbd8e Mon Sep 17 00:00:00 2001 From: "Jiahan Chang (Cyrus)" <173873397+jiahanc@users.noreply.github.com> Date: Thu, 8 Oct 2026 13:38:07 +0800 Subject: [PATCH 067/624] [Kernel] Remove WAR for LL CuTeDSL BF16 GEMM (#60434) Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com> --- tests/kernels/test_flashinfer_bf16_gemm.py | 2 + .../test_flashinfer_autotune_warmup.py | 9 ++-- vllm/model_executor/layers/utils.py | 5 +- vllm/model_executor/warmup/kernel_warmup.py | 50 +------------------ vllm/utils/flashinfer.py | 43 +++------------- 5 files changed, 16 insertions(+), 93 deletions(-) diff --git a/tests/kernels/test_flashinfer_bf16_gemm.py b/tests/kernels/test_flashinfer_bf16_gemm.py index 63a5a39c2e67..28520a6083e4 100644 --- a/tests/kernels/test_flashinfer_bf16_gemm.py +++ b/tests/kernels/test_flashinfer_bf16_gemm.py @@ -25,6 +25,8 @@ def require_flashinfer_bf16_cutedsl() -> None: (8, True, False), (16, False, True), (32, True, True), + (64, True, True), + (512, False, False), ], ) def test_flashinfer_bf16_cutedsl_correctness( diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py index fe89b09d09bc..f7ea6847f26f 100644 --- a/tests/model_executor/test_flashinfer_autotune_warmup.py +++ b/tests/model_executor/test_flashinfer_autotune_warmup.py @@ -39,12 +39,11 @@ def _make_moe(*, max_deferred_tokens: int = 128, enabled: bool = True): return moe -def _make_runner(modules, *, max_tokens: int = 8192, linear_backend: str = "auto"): +def _make_runner(modules, *, max_tokens: int = 8192): """Create a runner carrying only the state used by FlashInfer warmup.""" return SimpleNamespace( scheduler_config=SimpleNamespace(max_num_batched_tokens=max_tokens), vllm_config=SimpleNamespace( - kernel_config=SimpleNamespace(linear_backend=linear_backend), attention_config=SimpleNamespace(hisparse_config=None), parallel_config=SimpleNamespace(data_parallel_rank=0), ), @@ -63,21 +62,19 @@ def test_flashinfer_autotune_token_counts_include_deferred_moe_limits(): _make_moe(max_deferred_tokens=128), _make_moe(max_deferred_tokens=64, enabled=False), _make_moe(max_deferred_tokens=-1), - ], - linear_backend="flashinfer_cutedsl", + ] ) with patch("vllm.model_executor.layers.fused_moe.MoERunner", _FakeMoERunner): token_counts = _flashinfer_autotune_token_counts(runner) - assert token_counts == (8192, 32, 128) + assert token_counts == (8192, 128) def test_flashinfer_autotune_token_counts_are_bounded_and_deduplicated(): runner = _make_runner( [_make_moe(max_deferred_tokens=4096)], max_tokens=32, - linear_backend="flashinfer_cutedsl", ) with patch("vllm.model_executor.layers.fused_moe.MoERunner", _FakeMoERunner): diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index d4ee153fb251..80eef7c75a12 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -128,16 +128,13 @@ def _can_use_flashinfer_cutedsl_bf16( if ( k <= 0 or n <= 0 + or x.numel() == 0 or weight.shape[1] != k or k % 128 != 0 or x.data_ptr() % 32 != 0 or weight.data_ptr() % 32 != 0 ): return False - - m = x.numel() // k - if not 1 <= m <= 32: - return False return bias is None or ( bias.is_cuda and bias.device == x.device diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index d6544333beb5..e5751730e85d 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -316,9 +316,6 @@ def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None: return None -_FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS = 32 - - def _flashinfer_deferred_moe_token_counts( runner: "GPUModelRunner", ) -> tuple[int, ...]: @@ -339,19 +336,10 @@ def _flashinfer_deferred_moe_token_counts( return tuple(dict.fromkeys(token_counts)) -def _flashinfer_autotune_token_counts( - runner: "GPUModelRunner", *, include_bf16: bool = True -) -> tuple[int, ...]: +def _flashinfer_autotune_token_counts(runner: "GPUModelRunner") -> tuple[int, ...]: max_tokens = runner.scheduler_config.max_num_batched_tokens # Tune the widest bucket set first so bounded passes reuse its configs. token_counts = [max_tokens] - linear_backend = runner.vllm_config.kernel_config.linear_backend - if ( - include_bf16 - and linear_backend == "flashinfer_cutedsl" - and max_tokens > _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS - ): - token_counts.append(_FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS) token_counts.extend(_flashinfer_deferred_moe_token_counts(runner)) return tuple(dict.fromkeys(token_counts)) @@ -362,7 +350,7 @@ def _run_flashinfer_autotune_dummy_runs( import vllm.utils.flashinfer as fi_utils dummy_run_kwargs = {"skip_attn": True} if skip_attn else {} - for num_tokens in _flashinfer_autotune_token_counts(runner, include_bf16=False): + for num_tokens in _flashinfer_autotune_token_counts(runner): # The drafter may use more tokens than this pass in the same dummy run. # Include its max M when building the autotune buckets. max_tuning_tokens = num_tokens @@ -398,36 +386,6 @@ def _run_flashinfer_autotune_dummy_runs( ) -def _run_flashinfer_bf16_autotune_dummy_run( - runner: "GPUModelRunner", - *, - skip_ops: set[str] | None = None, - skip_attn: bool = False, -) -> None: - import vllm.utils.flashinfer as fi_utils - - if ( - runner.vllm_config.kernel_config.linear_backend != "flashinfer_cutedsl" - or runner.scheduler_config.max_num_batched_tokens - <= _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS - ): - return - - num_tokens = _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS - tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens) - logger.info("Running FlashInfer BF16-only autotune with %d tokens.", num_tokens) - # Other ops execute from their existing cache or fallback, without the - # BF16 bucket cap. No full-model autotune context may enclose this pass. - with fi_utils.autotune_bf16_only(tuning_buckets, skip_ops=skip_ops): - runner._dummy_run( - num_tokens=num_tokens, - skip_eplb=True, - is_profile=True, - randomize_inputs=True, - **({"skip_attn": True} if skip_attn else {}), - ) - - def _autotune_cache_fingerprint(path: Path) -> tuple[str, int] | None: """Identify a saved autotune file by its FlashInfer metadata and size.""" try: @@ -534,10 +492,6 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: _run_flashinfer_autotune_dummy_runs(runner, skip_attn=hisparse_enabled) replayssm_autotune_warmup(runner) _autotune_kimi_k3_kda_qkvg(runner.get_model()) - with torch.inference_mode(): - _run_flashinfer_bf16_autotune_dummy_run( - runner, skip_ops=skip_ops, skip_attn=hisparse_enabled - ) finally: set_autotune_process_group(None) diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index d7fa807c33b3..f9d6c4439f56 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -11,8 +11,7 @@ import importlib.util import os import shutil -from collections.abc import Callable, Iterator -from contextvars import ContextVar +from collections.abc import Callable from typing import Any, NoReturn import requests @@ -27,24 +26,6 @@ logger = init_logger(__name__) -_bf16_autotune_buckets: ContextVar[tuple[int, ...] | None] = ContextVar( - "flashinfer_bf16_autotune_buckets", default=None -) - - -@contextlib.contextmanager -def autotune_bf16_only( - tuning_buckets: tuple[int, ...], *, skip_ops: set[str] | None = None -) -> Iterator[None]: - """Tune BF16 calls with bounded buckets, outside full-model autotuning.""" - token = _bf16_autotune_buckets.set(tuning_buckets) - try: - with autotune(tune_mode=False, skip_ops=skip_ops): - yield - finally: - _bf16_autotune_buckets.reset(token) - - def flashinfer_bf16_mm( a: torch.Tensor, b: torch.Tensor, @@ -54,21 +35,14 @@ def flashinfer_bf16_mm( ) -> torch.Tensor: from flashinfer import mm_bf16 - tuning_buckets = _bf16_autotune_buckets.get() - tuning = ( - autotune(tune_mode=True, tuning_buckets=tuning_buckets) - if tuning_buckets is not None - else contextlib.nullcontext() + return mm_bf16( + a, + b, + bias=bias, + pdl=pdl, + out_dtype=torch.bfloat16, + backend=backend, ) - with tuning: - return mm_bf16( - a, - b, - bias=bias, - pdl=pdl, - out_dtype=torch.bfloat16, - backend=backend, - ) # This is the storage path for the cubins, it can be replaced @@ -1295,7 +1269,6 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool: __all__ = [ "has_flashinfer", "flashinfer_bf16_mm", - "autotune_bf16_only", "has_flashinfer_bf16_gemm", "is_flashinfer_bf16_gemm_supported", "is_flashinfer_cutedsl_bf16_gemm_supported", From 95307ac6d4cd8be2c1e38f236e9104c569af9025 Mon Sep 17 00:00:00 2001 From: Shanshan Shen <467638484@qq.com> Date: Thu, 8 Oct 2026 14:00:58 +0800 Subject: [PATCH 068/624] [Doc] Update ViT CUDA graph design doc (#60519) Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com> --- docs/design/cuda_graphs_multimodal.md | 96 ++++++++++++--------------- 1 file changed, 43 insertions(+), 53 deletions(-) diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index dca200e2de03..0a8a831f0164 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -7,27 +7,20 @@ For two-tower vision encoders (e.g., DeepSeek-OCR's SAM + CLIP with dynamic tili !!! note Encoder CUDA Graphs are orthogonal to decoder CUDA Graphs — both can be enabled simultaneously. Encoder graphs capture the vision encoder execution (e.g., ViT in Qwen3-VL), while decoder graphs capture the language model execution as described in the [CUDA Graphs design document](cuda_graphs.md). -## Compatibility Matrix - -!!! note - The symbols used below have the following meanings: - - - ✅ = Full compatibility - - 🟠 = Partial compatibility - - ❌ = No compatibility - - ❔ = Unknown or TBD - -### Model x Feature +## Supported Models | Architecture | Models | CG for Image | CG for Video | Multi-Path Graph | | ------------ | ------ | ------------ | ------------ | --------------- | | `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ | +| `DeepseekV41ForCausalLM` | `DeepSeek-V4.1-Flash` | ✅︎ | ❌︎ | ❌︎ | | `Ernie4_5_VLMoeForConditionalGeneration` | `ERNIE-4.5-VL` | ✅︎ | ❌︎ | ❌︎ | | `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ | | `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ | | `Gemma4ForConditionalGeneration` | `Gemma-4` | ✅︎ | ✅︎ | ❌︎ | | `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ | | `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ | +| `KimiK25ForConditionalGeneration` | `Kimi-K2.5, Kimi-K2.6` | ✅︎ | ❌︎ | ❌︎ | +| `KimiK3ForConditionalGeneration` | `Kimi-K3` | ✅︎ | ❌︎ | ❌︎ | | `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | ❌︎ | ❌︎ | | `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ | ❌︎ | | `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ | ❌︎ | @@ -38,26 +31,6 @@ For two-tower vision encoders (e.g., DeepSeek-OCR's SAM + CLIP with dynamic tili | `MiniCPMV` | `MiniCPMV2.5` | ✅︎ | ❌︎ | ❌︎ | | `MiniCPMV` | `MiniCPMV2.6`, `MiniCPMV4.0` | ✅︎ | ✅︎ | ❌︎ | -### Model x Hardware - -| Architecture | NV Blackwell | NV Ampere | AMD MI300X | AMD MI350X / MI355X | -| ------------ | ---------------- | ------------- | -------------- | --------------------- | -| `DeepseekOCRForCausalLM` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Ernie4_5_VLMoeForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Gemma3ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Glm4vForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Gemma4ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `InternVLChatModel` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `KimiVLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Llama4ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Qwen2VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Qwen2_5_VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Qwen3VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Qwen3_5ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Qwen3_5MoeForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Step3VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `MiniCPMV` | ✅︎ | ✅︎ | ❔ | ❔ | - !!! note Encoder CUDA Graph has currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs. For Qwen2-VL and Qwen2.5-VL only FA2 and FA3 has been tested. @@ -80,7 +53,7 @@ The encoder CUDA Graph system uses a **budget-based capture/replay** strategy, m * [EncoderItemSpec][vllm.v1.worker.encoder_cudagraph_defs.EncoderItemSpec]: describes a single encoder input item (image or video) with its input size and output token count. * [BudgetGraphMetadata][vllm.v1.worker.encoder_cudagraph.BudgetGraphMetadata]: holds the captured CUDA Graph and its associated I/O buffers for a single token budget level. -### Budget-based graph capture +### Budget-Based Graph Capture Multiple CUDA Graphs are pre-captured at different **token budget** levels (e.g., `[2048, 4096, 8192, 13824]`). Each budget defines a fixed token capacity, and all budgets share the same maximum batch size (number of images). The `BudgetGraphMetadata` for each level stores the graph along with pre-allocated input, metadata, and output buffers: @@ -99,20 +72,7 @@ Budgets are auto-generated as power-of-2 levels from a model-provided range via Each entry in `EncoderCudaGraphConfig.paths` defines an independently captured encoder path. A path can provide its own minimum token budget and opt into zero-token batches; the manager generates and stores a separate budget graph set for every configured path. -### Greedy bin-packing at runtime - -When a batch of images arrives, the manager sorts images by output token count (smallest first) and greedily packs as many images as possible into each sub-batch while staying within the **largest** token budget and the maximum batch size. Once a sub-batch is finalized (the next image would overflow either constraint), the manager finds the **smallest** budget that fits the sub-batch's total tokens and replays the corresponding CUDA Graph. This repeats until the batch is exhausted. Images that exceed all budgets fall back to eager execution. - -For multi-path models, the same greedy packing loop constrains every configured path simultaneously (see [Multi-Path graph capture](#multi-path-graph-capture)). - -For each graph replay: - -1. Call `prepare_encoder_cudagraph_replay_buffers()` to compute buffer values (including `pixel_values` and precomputed metadata) from actual batch inputs. -2. Zero the pre-allocated `input_buffers`, then slice-copy the replay values into them. -3. Replay the CUDA Graph. -4. Clone outputs from `output_buffer` (cloning is necessary since the buffer is reused across replays). - -### Multi-Path graph capture +### Multi-Path Graph Capture `EncoderCudaGraphConfig.paths` maps path names to `EncoderCudaGraphPathConfig` capture policies. For example, DeepSeek-OCR configures a **global** image path and a **local** patch path, which are captured independently under `budget_graphs["global"]` and `budget_graphs["local"]`. @@ -127,7 +87,7 @@ Both lists are capped at the same `max_budget`. * Sort images by total output tokens (global + local), smallest first. * Greedily pack images: an image is added to the current sub-batch only if every accumulated path token count is within that path's maximum budget, with the image count ≤ `max_batch_size`. -* Once any path constraint would overflow, finalize the sub-batch and find the smallest fitting budget **independently** for each path. +* Once any path constraint would overflow, or the budget-cliff deferral condition fires on any path, finalize the sub-batch and find the smallest fitting budget **independently** for each path. * Repeat until all images are packed. **Partial graph fallback.** For each non-empty path, the manager replays the smallest fitting graph or runs only that path eagerly when no graph fits. Paths with zero tokens are skipped; a `0`-budget graph is never captured or replayed. @@ -139,11 +99,41 @@ Both lists are capped at the same `max_budget`. !!! note The dual-path design enables partial CUDA graph coverage — one path can hit while the other falls back to eager. This avoids wasted compute on zero-padded patch buffers for untiled images and avoids graph invalidation caused by variable `crop_shape` per image. -### Data-parallel support +### Greedy Bin-Packing at Runtime + +When a batch of images arrives, the manager sorts images by output token count (smallest first) and greedily packs as many images as possible into each sub-batch while staying within the **largest** token budget and the maximum batch size. Once a sub-batch is finalized (the next image would overflow either constraint, or deferring it is cheaper), the manager finds the **smallest** budget that fits the sub-batch's total tokens and replays the corresponding CUDA Graph. This repeats until the batch is exhausted. Images that exceed all budgets fall back to eager execution. + +For multi-path models, the same greedy packing loop constrains every configured path simultaneously. + +For each graph replay: + +1. Call `prepare_encoder_cudagraph_replay_buffers()` to compute buffer values (including `pixel_values` and precomputed metadata) from actual batch inputs. +2. Zero the pre-allocated `input_buffers`, then slice-copy the replay values into them. +3. Replay the CUDA Graph. +4. Clone outputs from `output_buffer` (cloning is necessary since the buffer is reused across replays). + +**Budget-cliff deferral:** + +Greedy packing can be wasteful when a large image arrives after small ones: merging it into the current sub-batch may force the whole batch into a much larger budget graph even though splitting would cost less. Because token budgets are discrete (e.g., `[256, 4096, 8192]`), one extra token can double the required graph level — a "budget cliff". Images that overflow the largest budget still finalize the sub-batch unconditionally; for images that *fit*, the packer additionally applies a deferral test: finalize the current sub-batch (replaying it at its own small budget) and start a new sub-batch with the deferred image. + +For each candidate image, the packer compares three budgets per path: `b_cur` (the smallest budget fitting the current sub-batch), `b_alone` (the smallest budget fitting the image alone), and `b_merged` (the smallest budget fitting both). Merging costs `b_merged`; deferring costs `b_cur + b_alone` plus one extra replay. The image is deferred when, on any path, both conditions hold: + +1. `b_cur + b_alone < b_merged` — splitting strictly saves budget. The strict inequality makes ties merge: with budgets `[256, 4096, 8192]`, two 4000-token images tie (`4096 + 4096 == 8192`), and merging wins because a single replay avoids an extra graph launch. +2. `b_alone >= b_cur` — the image alone occupies at least as large a budget bucket as the entire current sub-batch. + +**Condition 1 is the cost test.** For example, a 200-token image followed by a 4000-token image with budgets `[256, 4096, 8192]`: merging takes `b_merged = 8192` (4200 tokens round up to 8192), while deferring takes `256 + 4096 = 4352` — the image is deferred, saving 3840 tokens of padding compute. + +**Condition 2 separates "whales" from "tag-alongs".** Because images are sorted smallest first, a whale arrives last, and deferring it lets the small sub-batch replay at its own small budget instead of being dragged into a large bucket for one image's sake. Without condition 2, deferring mid-size images fragments a flood of similar images into extra replays: with budgets `[256, 2048, 4096, 8192]` and six equal 1366-token images, condition 1 alone defers every third image and yields three 4096 replays (12288 budget tokens in total), while with condition 2 five images merge into one 8192 replay plus a single 2048 replay — fewer replays *and* fewer budget tokens (10240). + +The conditions compose cleanly: when `b_alone == b_cur`, condition 1 requires `b_merged > 2 * b_cur`, i.e. deferral fires only when merging jumps at least two budget levels, so ties always merge. This also lets small images defer across a huge cliff: with budgets `[256, 4096]` and four 100-token images, adding the third image (300 tokens) would jump from 256 to 4096, so the packer defers it and packs the images as two 256-budget pairs. + +For multi-path items, the conditions are evaluated per path and OR-ed: any path on which the image is a whale defers the whole image, while paths where the image contributes zero tokens are skipped. + +### ViT DP Mode Support When `mm_encoder_tp_mode="data"`, the manager distributes images across TP ranks using load-balanced assignment via `get_load_balance_assignment`, executes locally on each rank, then gathers results back in the original order via `tensor_model_parallel_all_gather`. -### Video inference support +### Video Inference Support Following (ViT full CUDA graph support for image inference), extends the encoder CUDA graph framework to support video inference for Qwen3-VL. Previously, the CUDA graph capture/replay path only handled image inputs (`pixel_values` + `image_grid_thw`). Video inputs use different keys (`pixel_values_videos` + `video_grid_thw`) and require larger `cu_seqlens` buffers because each video item contributes multiple frames (`T` attention sequences). This PR generalizes the protocol and manager to handle both modalities through a single shared graph manager. @@ -152,7 +142,7 @@ Following (ViT full CUDA graph Mixed inputs (image+video) per prompt are also supported now. -## Model integration via `SupportsEncoderCudaGraph` +## Model Integration via `SupportsEncoderCudaGraph` Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGraph][vllm.model_executor.models.interfaces.SupportsEncoderCudaGraph] protocol. This protocol encapsulates all model-specific logic so that the manager remains model-agnostic. The protocol defines the following methods: @@ -180,9 +170,9 @@ Four fields in `CompilationConfig` control encoder CUDA Graphs: Multi-path mode is configured at the model level through `EncoderCudaGraphConfig.paths`. Each `EncoderCudaGraphPathConfig` can set a path-specific minimum budget and whether zero-token batches are allowed. The manager automatically generates separate budget lists and uses the same execution loop for single- and multi-path models. -## Usage guide +## Usage Guide -### Image inference +### Image Inference Enable encoder CUDA Graphs via `compilation_config`: @@ -226,7 +216,7 @@ model = vllm.LLM( The manager tracks hit/miss statistics and logs them periodically. A "hit" means an image was processed via CUDA Graph replay; a "miss" means eager fallback (image exceeded all budgets). -### Video inference +### Video Inference Enable encoder CUDA Graphs via `compilation_config`: From fb2ac82458cf6bff45b21378a85fa659213b51e5 Mon Sep 17 00:00:00 2001 From: JohnQinAMD Date: Wed, 7 Oct 2026 23:15:18 -0700 Subject: [PATCH 069/624] [ROCm][DSv4.1][Perf] Replace the Torch prefill combine with the shared combine_topk_swa_indices kernel (#56638) Signed-off-by: JohnQinAMD Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Ren Yuzhou <54501155+yuzhouo7@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> --- .../attention/test_rocm_triton_attn_dsv4.py | 158 ++++++++++++++++- .../attention/test_deepseek_v4_swa_visible.py | 23 ++- vllm/models/deepseek_v41/amd/rocm.py | 164 +----------------- .../deepseek_v41/common/ops/cache_utils.py | 10 +- 4 files changed, 185 insertions(+), 170 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index ca2cf2eb7017..362ac6b7ce0a 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -41,7 +41,7 @@ def _on_gfx950() -> bool: # architectures take the fallback decode kernel, so its tests are skipped there. requires_split_decode_arch = pytest.mark.skipif( not _on_split_decode_arch(), - reason="split-K decode kernel is only tuned for AMD gfx942/gfx950", + reason="split-K decode and the V4.1 combine tests need an AMD gfx942/gfx950", ) requires_gfx950 = pytest.mark.skipif( not _on_gfx950(), @@ -2407,3 +2407,159 @@ def test_paged_mqa_logits_gate_v41_shuffle_vs_c4a_flat() -> None: torch.testing.assert_close( gated_c4a[:, :seq_len], ref[:, :seq_len], atol=2e-3, rtol=2e-3 ) + + +def _combine_topk_swa_indices_torch( + topk_indices, + query_start_loc, + seq_lens, + gather_lens, + window_size, + compress_ratio, + topk, + M, + N, +): + """The Torch path this kernel replaces, copied from before the switch. + + Kept verbatim so the reference cannot drift toward the kernel. + """ + from vllm.models.deepseek_v41.common.ops.cache_utils import ( + _SPARSE_PREFILL_TOPK_ALIGNMENT, + ) + + topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous() + num_tokens = topk_indices.shape[0] + combined_topk = ( + (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) + // _SPARSE_PREFILL_TOPK_ALIGNMENT + * _SPARSE_PREFILL_TOPK_ALIGNMENT + ) + combined_indices = torch.full( + (num_tokens, combined_topk), + fill_value=-1, + dtype=torch.int32, + device=topk_indices.device, + ) + combined_lens = torch.empty( + num_tokens, dtype=torch.int32, device=topk_indices.device + ) + + # query_start_loc may have a non-zero base for a narrowed mixed batch. + query_lens = query_start_loc[1:] - query_start_loc[:-1] + req_ids = torch.repeat_interleave( + torch.arange(seq_lens.shape[0], device=seq_lens.device), query_lens + ) + query_starts = query_start_loc[:-1] - query_start_loc[0] + token_offsets = torch.arange(num_tokens, device=seq_lens.device) - ( + torch.repeat_interleave(query_starts, query_lens) + ) + positions = seq_lens[req_ids] - query_lens[req_ids] + token_offsets + + logical_topk_width = min(topk, topk_indices.shape[1]) + topk_lens = torch.minimum( + (positions + 1) // compress_ratio, + torch.full_like(positions, logical_topk_width), + ).clamp_min(0) + topk_offsets = torch.arange(logical_topk_width, device=seq_lens.device) + topk_mask = topk_offsets[None, :] < topk_lens[:, None] + topk_values = topk_indices[:, :logical_topk_width].to(torch.int32) + topk_valid = topk_mask & (topk_values >= 0) & (topk_values < N) + combined_indices[:, :logical_topk_width] = torch.where( + topk_valid, + topk_values + (M * req_ids).to(torch.int32)[:, None], + -1, + ) + + swa_lens = torch.minimum( + positions + 1, torch.full_like(positions, window_size) + ).clamp_min(0) + swa_offsets = torch.arange(window_size, device=seq_lens.device) + swa_mask = swa_offsets[None, :] < swa_lens[:, None] + swa_columns = topk_lens[:, None] + swa_offsets[None, :] + gather_starts = seq_lens - gather_lens + swa_values = ( + M * req_ids[:, None] + + N + + swa_offsets[None, :] + + positions[:, None] + - swa_lens[:, None] + + 1 + - gather_starts[req_ids, None] + ).to(torch.int32) + rows = torch.arange(num_tokens, device=seq_lens.device)[:, None].expand_as( + swa_columns + ) + flat_dst = rows[swa_mask] * combined_topk + swa_columns[swa_mask] + combined_indices.view(-1).index_copy_(0, flat_dst, swa_values[swa_mask]) + combined_lens.copy_((topk_lens + swa_lens).to(torch.int32)) + return combined_indices, combined_lens + + +def _v41_combine_case(case, window): + dev = "cuda" + ratio, topk = 2, 512 + if case == "seq_len_below_query_len": + qsl = torch.tensor([0, 64], dtype=torch.int32, device=dev) + sl = torch.tensor([32], dtype=torch.int32, device=dev) + rows = 64 + elif case == "seq_len_zero": + # The warmup batch: every pos is negative, so both lengths clamp to 0. + qsl = torch.tensor([0, 32, 64], dtype=torch.int32, device=dev) + sl = torch.zeros(2, dtype=torch.int32, device=dev) + rows = 64 + elif case == "compress_ratio_zero": + qsl = torch.tensor([0, 4, 8], dtype=torch.int32, device=dev) + sl = torch.tensor([4, 4], dtype=torch.int32, device=dev) + rows, ratio, topk = 8, 0, 0 + elif case == "nonzero_base": + # Prefill rows of a mixed batch: query_start_loc starts past the decodes. + qsl = torch.tensor([32, 64, 128], dtype=torch.int32, device=dev) + sl = torch.tensor([48, 80], dtype=torch.int32, device=dev) + rows = 96 + else: + qsl = torch.tensor([0, 32, 96], dtype=torch.int32, device=dev) + sl = torch.tensor([48, 80], dtype=torch.int32, device=dev) + rows = 96 + # What the SWA builder gathers: the query plus up to window - 1 rows before it. + gl = torch.minimum(sl, qsl[1:] - qsl[:-1] + window - 1) + set_random_seed(0) + ti = torch.randint(0, 4096, (rows, max(topk, 1)), dtype=torch.int32) + if case == "invalid_topk": + # Failed candidates stay -1 instead of pointing into another request. + ti[:, ::7] = -1 + return ti.to(dev), qsl, sl, gl, window, ratio, topk, 100000, 4096 + + +@requires_split_decode_arch +@pytest.mark.parametrize( + "case,window", + [ + ("seq_len_below_query_len", 128), + ("seq_len_zero", 128), + ("compress_ratio_zero", 128), + ("nonzero_base", 128), + ("invalid_topk", 128), + ("plain", 64), + ("plain", 128), + ("plain", 256), + ], +) +def test_v41_combine_topk_swa_indices_matches_torch(case, window) -> None: + from vllm.models.deepseek_v41.amd.rocm import combine_topk_swa_indices + + args = _v41_combine_case(case, window) + got_indices, got_lens = combine_topk_swa_indices(*args) + ref_indices, ref_lens = _combine_topk_swa_indices_torch(*args) + assert torch.equal(got_lens, ref_lens), ( + f"lens differ: ref={ref_lens.tolist()} got={got_lens.tolist()}" + ) + assert torch.equal(got_indices, ref_indices) + + +def test_v41_rocm_combine_uses_shared_kernel() -> None: + """ROCm prefill takes the shared combine, so it inherits the replay bound.""" + from vllm.models.deepseek_v41.amd import rocm + from vllm.models.deepseek_v41.common.ops import cache_utils + + assert rocm.combine_topk_swa_indices is cache_utils.combine_topk_swa_indices diff --git a/tests/v1/attention/test_deepseek_v4_swa_visible.py b/tests/v1/attention/test_deepseek_v4_swa_visible.py index 7a74ce8c81be..8f08929549a8 100644 --- a/tests/v1/attention/test_deepseek_v4_swa_visible.py +++ b/tests/v1/attention/test_deepseek_v4_swa_visible.py @@ -302,8 +302,13 @@ def combine_case( with_image: bool, replay_starts: list[int] | None = None, combine_fn=combine_topk_swa_indices, + in_range_topk: bool = False, ): - """Run combine_topk_swa_indices and return (indices, lens, expected).""" + """Run combine_topk_swa_indices and return (indices, lens, expected). + + ``in_range_topk`` draws top-k indices below N, for combiners that drop + out-of-range ones. + """ device = torch.device("cuda") num_reqs = len(seq_lens) replay_starts = replay_starts or [0] * num_reqs @@ -326,7 +331,11 @@ def combine_case( M = N + int(gather_lens.max()) + 8 gen = torch.Generator(device="cpu").manual_seed(0) topk_indices = torch.randint( - 0, 4096, (num_tokens, max(topk, 1)), generator=gen, dtype=torch.int32 + 0, + N if in_range_topk else 4096, + (num_tokens, max(topk, 1)), + generator=gen, + dtype=torch.int32, ).to(device) topk_indices = topk_indices[:, : max(topk, 1)] @@ -415,24 +424,26 @@ def test_combine_topk_swa_with_image_spans(cfg): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @pytest.mark.parametrize("cfg", COMBINE_CASES) -def test_v41_combine_topk_swa_stops_at_replay_start(cfg): +@pytest.mark.parametrize("query_len", [24, 20]) +def test_v41_combine_topk_swa_stops_at_replay_start(cfg, query_len): """SWA bounded replay: the gathered buffer starts at replay_start, so the window never indexes below it.""" from vllm.models.deepseek_v41.common.ops.cache_utils import ( combine_topk_swa_indices as combine_v41, ) - # Request 0 replays [16, 40): the windows of its first rows would - # otherwise reach below 16. + # Request 0 replays from 16: a 24-token chunk starts there, a 20-token + # chunk leaves only 4 context tokens above it (fewer than the window). indices, lens, rows, exp_lens = combine_case( cfg["compress_ratio"], cfg["topk"], seq_lens=[40, 12], - query_lens=[24, 12], + query_lens=[query_len, 12], spans=[[], []], with_image=False, replay_starts=[16, 0], combine_fn=combine_v41, + in_range_topk=True, ) assert lens.cpu().tolist() == exp_lens assert indices.cpu().tolist() == rows diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 64b4a9d3d1f2..3aa54f9b0e97 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -26,7 +26,10 @@ DeepseekV4IndexerCache, _replace_layer_index, ) -from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache +from vllm.models.deepseek_v41.common.ops import ( + combine_topk_swa_indices, + dequantize_and_gather_k_cache, +) from vllm.models.deepseek_v41.sparse_mla import ( DeepseekV4FlashMLAMetadata, DeepseekV4SparseMLABackend, @@ -129,165 +132,6 @@ def apply_pre_quantized_block_scaled_mm( return out.to(dtype=kernel.config.out_dtype) -# ROCm sparse prefill keeps this dense combine local so AMD-specific SWA changes -# do not touch the shared DeepSeek V4 cache utilities. -_SPARSE_PREFILL_TOPK_ALIGNMENT = 128 - - -@triton.jit -def _combine_topk_swa_indices_kernel( - combined_indices_ptr, - combined_indices_stride, - combined_lens_ptr, - topk_indices_ptr, - topk_indices_stride, - query_start_loc_ptr, - seq_lens_ptr, - gather_lens_ptr, - M, - N, - TOP_K: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, - WINDOW_SIZE: tl.constexpr, - TOPK_WIDTH: tl.constexpr, - PADDED_TOP_K: tl.constexpr, -): - batch_idx = tl.program_id(0) - worker_id = tl.program_id(1) - num_workers = tl.num_programs(1) - - base = tl.load(query_start_loc_ptr) - query_start = tl.load(query_start_loc_ptr + batch_idx) - base - query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base - query_len = query_end - query_start - seq_len = tl.load(seq_lens_ptr + batch_idx) - gather_len = tl.load(gather_lens_ptr + batch_idx) - start_pos = seq_len - query_len - gather_start = seq_len - gather_len - - for token_idx in range(query_start + worker_id, query_end, num_workers): - token_idx_in_query = token_idx - query_start - pos = start_pos + token_idx_in_query - topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) - swa_len = tl.minimum(pos + 1, WINDOW_SIZE) - - topk_offset = tl.arange(0, PADDED_TOP_K) - topk_mask = topk_offset < topk_len - safe_topk_offset = tl.where(topk_offset < TOPK_WIDTH, topk_offset, 0) - topk_indices = tl.load( - topk_indices_ptr + token_idx * topk_indices_stride + safe_topk_offset, - mask=topk_mask, - other=-1, - ) - valid_topk = (topk_indices >= 0) & (topk_indices < N) - topk_indices = tl.where(valid_topk, topk_indices + M * batch_idx, -1) - tl.store( - combined_indices_ptr + token_idx * combined_indices_stride + topk_offset, - topk_indices, - mask=topk_mask, - ) - - swa_offset = tl.arange(0, WINDOW_SIZE) - tl.store( - combined_indices_ptr - + token_idx * combined_indices_stride - + topk_len - + swa_offset, - M * batch_idx + N + swa_offset + pos - swa_len + 1 - gather_start, - mask=swa_offset < swa_len, - ) - - tl.store(combined_lens_ptr + token_idx, topk_len + swa_len) - - -def combine_topk_swa_indices( - topk_indices: torch.Tensor, - query_start_loc: torch.Tensor, - seq_lens: torch.Tensor, - gather_lens: torch.Tensor, - window_size: int, - compress_ratio: int, - topk: int, - M: int, - N: int, -) -> tuple[torch.Tensor, torch.Tensor]: - """Combine compressed-attention and sliding-window indices with Torch. - - The Triton implementation inherited from DeepSeek V4 launches a - two-dimensional grid with 128 workers per request. On gfx950 it can issue - an out-of-bounds access for V4.1's mixed prefill metadata (including the - synthetic mixed-token warmup). This path is prefill-only and the tensors - are small, so use ordinary Torch indexing until a gfx950-safe fused kernel - is available. - """ - topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous() - num_tokens = topk_indices.shape[0] - combined_topk = ( - (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) - // _SPARSE_PREFILL_TOPK_ALIGNMENT - * _SPARSE_PREFILL_TOPK_ALIGNMENT - ) - combined_indices = torch.full( - (num_tokens, combined_topk), - fill_value=-1, - dtype=torch.int32, - device=topk_indices.device, - ) - combined_lens = torch.empty( - num_tokens, dtype=torch.int32, device=topk_indices.device - ) - - # query_start_loc may have a non-zero base for a narrowed mixed batch. - query_lens = query_start_loc[1:] - query_start_loc[:-1] - req_ids = torch.repeat_interleave( - torch.arange(seq_lens.shape[0], device=seq_lens.device), query_lens - ) - query_starts = query_start_loc[:-1] - query_start_loc[0] - token_offsets = torch.arange(num_tokens, device=seq_lens.device) - ( - torch.repeat_interleave(query_starts, query_lens) - ) - positions = seq_lens[req_ids] - query_lens[req_ids] + token_offsets - - logical_topk_width = min(topk, topk_indices.shape[1]) - topk_lens = torch.minimum( - (positions + 1) // compress_ratio, - torch.full_like(positions, logical_topk_width), - ).clamp_min(0) - topk_offsets = torch.arange(logical_topk_width, device=seq_lens.device) - topk_mask = topk_offsets[None, :] < topk_lens[:, None] - topk_values = topk_indices[:, :logical_topk_width].to(torch.int32) - topk_valid = topk_mask & (topk_values >= 0) & (topk_values < N) - combined_indices[:, :logical_topk_width] = torch.where( - topk_valid, - topk_values + (M * req_ids).to(torch.int32)[:, None], - -1, - ) - - swa_lens = torch.minimum( - positions + 1, torch.full_like(positions, window_size) - ).clamp_min(0) - swa_offsets = torch.arange(window_size, device=seq_lens.device) - swa_mask = swa_offsets[None, :] < swa_lens[:, None] - swa_columns = topk_lens[:, None] + swa_offsets[None, :] - gather_starts = seq_lens - gather_lens - swa_values = ( - M * req_ids[:, None] - + N - + swa_offsets[None, :] - + positions[:, None] - - swa_lens[:, None] - + 1 - - gather_starts[req_ids, None] - ).to(torch.int32) - rows = torch.arange(num_tokens, device=seq_lens.device)[:, None].expand_as( - swa_columns - ) - flat_dst = rows[swa_mask] * combined_topk + swa_columns[swa_mask] - combined_indices.view(-1).index_copy_(0, flat_dst, swa_values[swa_mask]) - combined_lens.copy_((topk_lens + swa_lens).to(torch.int32)) - return combined_indices, combined_lens - - @triton.jit def _compute_topk_lens_kernel( topk_lens_ptr, diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py index b7f07d8d84c0..a6139c71f836 100644 --- a/vllm/models/deepseek_v41/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py @@ -1089,7 +1089,8 @@ def kernel( token_idx_in_query = token_idx - query_start pos = start_pos + token_idx_in_query if COMPRESS_RATIO > 0: - topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) + # Warmup batches can put pos below zero; keep the length >= 0. + topk_len = tl.maximum(tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K), 0) else: # SWA-only layers pass TOP_K=0; skip the division entirely # (integer div by constexpr 0 is UB and yields garbage, which @@ -1100,7 +1101,7 @@ def kernel( # start (SWA bounded replay), so the window cannot start before # the gathered buffer does. swa_start = tl.maximum(swa_start, gather_start) - swa_len = pos - swa_start + 1 + swa_len = tl.maximum(pos - swa_start + 1, 0) offset = tl.arange(0, PADDED_TOP_K) mask = offset < topk_len @@ -1108,9 +1109,12 @@ def kernel( topk_indices_ptr + token_idx * topk_indices_stride + offset, mask=mask, ) + # A failed (-1) or out-of-pool candidate stays -1 instead of landing + # in the previous request's rows or this request's SWA rows. + valid = (topk_indices >= 0) & (topk_indices < N) tl.store( combined_indices_ptr + token_idx * combined_indices_stride + offset, - topk_indices + M * batch_idx, + tl.where(valid, topk_indices + M * batch_idx, -1), mask=mask, ) # Index into gathered buffer: N + (position - gather_start) From f226af0e114bb67bf59fe6ee2fc762dcbf956b4c Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Thu, 8 Oct 2026 14:25:27 +0800 Subject: [PATCH 070/624] [Bugfix] Eliminate frontend ZMQ port TOCTOU with inherited listeners (#54113) Signed-off-by: Bugen Zhao Co-authored-by: OpenAI Codex Co-authored-by: Claude Opus 5.5 --- rust/Cargo.lock | 62 +++- rust/Cargo.toml | 2 +- rust/src/cmd/src/cli.rs | 22 +- rust/src/cmd/src/cli/tests.rs | 192 ++++++------ rust/src/engine-core-client/Cargo.toml | 2 + rust/src/engine-core-client/src/client.rs | 25 +- .../engine-core-client/src/tests/client.rs | 38 ++- rust/src/engine-core-client/src/transport.rs | 68 +++- rust/src/server/src/grpc/tests.rs | 10 +- .../api_server/_api_server_spawn_workers.py | 14 - .../test_api_server_process_manager.py | 294 +++++++++--------- tests/utils_/test_network_utils.py | 47 +++ .../engine/test_core_engine_actor_manager.py | 90 ++---- vllm/entrypoints/cli/serve.py | 35 +-- vllm/utils/network_utils.py | 103 +++++- vllm/v1/engine/core_client.py | 57 ++-- vllm/v1/engine/utils.py | 85 ++++- vllm/v1/utils.py | 235 +++++++------- 18 files changed, 793 insertions(+), 588 deletions(-) delete mode 100644 tests/entrypoints/launchers/api_server/_api_server_spawn_workers.py diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 9dcdb0bdbc2c..331ebfa8dfaa 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -593,6 +593,17 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "chacha20" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "rand_core 0.10.1", +] + [[package]] name = "chrono" version = "0.4.44" @@ -1614,6 +1625,7 @@ dependencies = [ "cfg-if", "libc", "r-efi 6.0.0", + "rand_core 0.10.1", "wasip2", "wasip3", ] @@ -3451,6 +3463,17 @@ dependencies = [ "rand_core 0.9.5", ] +[[package]] +name = "rand" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65c9fb96cbc91e3478eaae79a69fcd3f1ae4ad052e471fe6732fff548984b4af" +dependencies = [ + "chacha20", + "getrandom 0.4.2", + "rand_core 0.10.1", +] + [[package]] name = "rand_chacha" version = "0.3.1" @@ -3489,6 +3512,12 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + [[package]] name = "rand_distr" version = "0.5.1" @@ -3936,9 +3965,9 @@ checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" [[package]] name = "saa" -version = "5.5.0" +version = "5.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "16c7f49c9d5caa3bf4b3106900484b447b9253fe99670ceb81cb6cb5027855e1" +checksum = "bd85992a680de788217d9c70b74aa7685b108cf1aba02f728fb27f1b27199ad1" [[package]] name = "same-file" @@ -3960,12 +3989,12 @@ dependencies = [ [[package]] name = "scc" -version = "3.6.9" +version = "3.8.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "45bb5ce9efd4a6e7b0f86c2697fe4c1d78d1f4e6d988c54b752d577cafe22fe8" +checksum = "5d5f8ca67c4f6208caf824037a6a9df0211b4e0937fccc4f543441f294d04511" dependencies = [ "saa", - "sdd 4.7.3", + "sdd 4.8.11", ] [[package]] @@ -4039,9 +4068,12 @@ checksum = "490dcfcbfef26be6800d11870ff2df8774fa6e86d047e3e8c8a76b25655e41ca" [[package]] name = "sdd" -version = "4.7.3" +version = "4.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b21a75f5913ab130e4b369fb8693be25f29b983e2ecad4279df9bfa5dd8aaf3e" +checksum = "a31877eee03686b905845993b5d9809a220c9dd5a10e9d6021d7ca67384c9070" +dependencies = [ + "saa", +] [[package]] name = "secrecy" @@ -4991,7 +5023,6 @@ dependencies = [ "bytes", "libc", "mio", - "parking_lot", "pin-project-lite", "signal-hook-registry", "socket2", @@ -5799,6 +5830,7 @@ dependencies = [ "half", "hex", "itertools 0.14.0", + "libc", "parking_lot", "rmp-serde", "rmpv", @@ -5808,6 +5840,7 @@ dependencies = [ "serde_repr", "serde_tuple", "serde_with", + "socket2", "task-local", "tempfile", "thiserror 2.0.18", @@ -6643,22 +6676,21 @@ checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" [[package]] name = "zeromq" version = "0.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "efb2c254fd8f366755335c9e43b865f8484fe3bd717d65ffe7c3f28852863030" +source = "git+https://github.com/zeromq/zmq.rs?rev=67dcba97559ebe8e8206dab753710d8b9cb07591#67dcba97559ebe8e8206dab753710d8b9cb07591" dependencies = [ "async-trait", "asynchronous-codec", "bytes", "crossbeam-queue", "futures", + "futures-channel", "log", - "num-traits", - "once_cell", "parking_lot", - "rand 0.9.5", + "rand 0.10.3", "regex", - "scc 3.6.9", - "thiserror 1.0.69", + "scc 3.8.8", + "socket2", + "thiserror 2.0.18", "tokio", "tokio-util", "uuid", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 38436467f862..57f680cc147e 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -155,7 +155,7 @@ vllm-tokenizer = { path = "src/tokenizer" } vllm-tracing = { path = "src/tracing" } winnow = { version = "1.0.2", features = ["simd"] } xgrammar-structural-tag = "0.4.1" -zeromq = { version = "0.6.0", default-features = false, features = [ +zeromq = { git = "https://github.com/zeromq/zmq.rs", rev = "67dcba97559ebe8e8206dab753710d8b9cb07591", default-features = false, features = [ "tokio-runtime", "all-transport", ] } diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 1269e0b8afa2..d3e369bbe388 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -501,8 +501,8 @@ impl SharedRuntimeArgs { self, listen_fd: i32, grpc_listen_fd: Option, - input_address: String, - output_address: String, + input_listener_fd: i32, + output_listener_fd: i32, coordinator_address: Option, engine_start_index: u32, engine_count: usize, @@ -518,8 +518,8 @@ impl SharedRuntimeArgs { Config { transport_mode: TransportMode::Bootstrapped { - input_address, - output_address, + input_listener_fd, + output_listener_fd, engine_start_index, engine_count, data_parallel_size, @@ -697,14 +697,12 @@ pub struct FrontendArgs { /// supervisor. When not set, no gRPC server is started. #[arg(long)] pub grpc_listen_fd: Option, - /// Frontend input ROUTER socket address that the Python engines will - /// connect to. + /// Inherited frontend input ROUTER listener file descriptor. #[arg(long)] - pub input_address: String, - /// Frontend output PULL socket address that the Python engines will push - /// responses to. + pub input_listener_fd: i32, + /// Inherited frontend output PULL listener file descriptor. #[arg(long)] - pub output_address: String, + pub output_listener_fd: i32, /// Optional Python-owned frontend-side DP coordinator socket address for /// external coordinator mode in the bootstrapped frontend path, i.e., /// `stats_update_address`. @@ -733,8 +731,8 @@ impl FrontendArgs { self.runtime.into_bootstrapped_config( self.listen_fd, self.grpc_listen_fd, - self.input_address, - self.output_address, + self.input_listener_fd, + self.output_listener_fd, self.coordinator_address, self.engine_start_index, self.engine_count, diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 944febc36c10..aa164bc618f8 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -668,10 +668,10 @@ fn frontend_args_json_passes_tls_into_config() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","ssl_certfile":"/tmp/cert.pem","ssl_keyfile":"/tmp/key.pem"}"#, ]) @@ -693,10 +693,10 @@ fn frontend_args_json_rejects_out_of_range_cert_reqs() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","ssl_certfile":"/tmp/cert.pem","ssl_cert_reqs":5}"#, ]) @@ -718,10 +718,10 @@ fn frontend_args_json_passes_enable_request_id_headers_into_config() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","enable_request_id_headers":true}"#, ]) @@ -807,10 +807,10 @@ fn frontend_args_json_passes_lora_modules_into_config() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","lora_modules":[{"name":"alice","path":"org/alice","base_model_name":null,"is_3d_lora_weight":true}]}"#, ]) @@ -862,10 +862,10 @@ fn frontend_args_json_accepts_api_key_string() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","api_key":"secret"}"#, ]) @@ -885,10 +885,10 @@ fn frontend_args_json_accepts_api_key_list() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","api_key":["secret-a","secret-b"]}"#, ]) @@ -991,10 +991,10 @@ fn frontend_args_accept_json() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--coordinator-address", "tcp://127.0.0.1:7000", "--args-json", @@ -1008,8 +1008,8 @@ fn frontend_args_accept_json() { FrontendArgs { listen_fd: 3, grpc_listen_fd: None, - input_address: "ipc:///tmp/input.sock", - output_address: "ipc:///tmp/output.sock", + input_listener_fd: 999999, + output_listener_fd: 999998, coordinator_address: Some( "tcp://127.0.0.1:7000", ), @@ -1083,10 +1083,10 @@ fn frontend_args_json_applies_defaults() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--engine-count", "4", "--args-json", @@ -1115,10 +1115,10 @@ fn frontend_args_json_ignores_engine_owned_max_model_len() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","max_model_len":-1}"#, ]) @@ -1165,10 +1165,10 @@ fn hf_overrides_preserve_merge_patch_in_cli_and_python_bootstrap() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", &payload.to_string(), ]) @@ -1204,10 +1204,10 @@ fn frontend_args_json_accepts_supported_non_default_fields() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","generation_config":"vllm","revision":"release","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_logprobs":-1,"shutdown_timeout":3}"#, ]) @@ -1286,10 +1286,10 @@ fn frontend_args_json_ignores_unknown_fields() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","uds":"/tmp/vllm.sock","nested_unknown":{"x":1}}"#, ]) @@ -1308,10 +1308,10 @@ fn frontend_args_json_sets_prompt_tokens_details_flag() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","api_server_count":2,"enable_prompt_tokens_details":true}"#, ]) @@ -1371,10 +1371,10 @@ fn frontend_args_json_parses_cors_fields() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","allowed_origins":["http://a.com"],"allow_credentials":true}"#, ]) @@ -1396,10 +1396,10 @@ fn frontend_args_json_rejects_unsupported_fields() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","root_path":"/prefix"}"#, ]) @@ -1423,10 +1423,10 @@ fn frontend_args_json_aggregates_multiple_unsupported_fields() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","response_role":"assistant","root_path":"/prefix"}"#, ]) @@ -1452,10 +1452,10 @@ fn frontend_args_json_rejects_malformed_json() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B""#, ]) @@ -1999,10 +1999,10 @@ fn frontend_args_reject_legacy_handshake_flags() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#, "--handshake-address", @@ -2043,10 +2043,10 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--coordinator-address", "tcp://127.0.0.1:7000", "--engine-start-index", @@ -2068,8 +2068,8 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present expect![[r#" Config { transport_mode: Bootstrapped { - input_address: "ipc:///tmp/input.sock", - output_address: "ipc:///tmp/output.sock", + input_listener_fd: 999999, + output_listener_fd: 999998, engine_start_index: 3, engine_count: 1, data_parallel_size: 4, @@ -2200,10 +2200,10 @@ fn frontend_config_adopts_inherited_grpc_listener() { "3", "--grpc-listen-fd", "4", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","grpc_port":50051}"#, ]) @@ -2227,10 +2227,10 @@ fn frontend_args_json_enables_profiling_when_profiler_config_set() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","profiler_config":{"profiler":"torch","torch_profiler_dir":"/tmp/profile"}}"#, ]) @@ -2251,10 +2251,10 @@ fn frontend_args_json_disables_profiling_when_profiler_config_absent() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#, ]) @@ -2275,10 +2275,10 @@ fn frontend_args_json_disables_profiling_when_profiler_type_is_null() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","profiler_config":{"profiler":null}}"#, ]) diff --git a/rust/src/engine-core-client/Cargo.toml b/rust/src/engine-core-client/Cargo.toml index 14b9fe142349..86ffe4c4cdae 100644 --- a/rust/src/engine-core-client/Cargo.toml +++ b/rust/src/engine-core-client/Cargo.toml @@ -18,6 +18,7 @@ futures.workspace = true half.workspace = true hex.workspace = true itertools.workspace = true +libc.workspace = true parking_lot.workspace = true rmp-serde.workspace = true rmpv.workspace = true @@ -27,6 +28,7 @@ serde_json.workspace = true serde_repr.workspace = true serde_tuple.workspace = true serde_with.workspace = true +socket2.workspace = true task-local.workspace = true tempfile = { workspace = true, optional = true } thiserror.workspace = true diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 36a4fbffa5e1..b3516936fef2 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -53,16 +53,15 @@ pub enum TransportMode { local_output_address: Option, }, - /// The Python supervisor has already chosen the frontend transport - /// addresses, and the Rust process only needs to bind them and wait for - /// engine registration frames. + /// The Python supervisor has already bound the frontend transport + /// listeners. Rust adopts them and waits for engine registration frames. + /// `EngineCoreClient::connect` consumes each descriptor exactly once; + /// copied integer values do not create additional descriptor ownership. Bootstrapped { - /// Input ROUTER socket address that engines will connect to for - /// requests. - input_address: String, - /// Output PULL socket address that engines will connect to for - /// responses. - output_address: String, + /// Raw input ROUTER listener descriptor inherited from the supervisor. + input_listener_fd: i32, + /// Raw output PULL listener descriptor inherited from the supervisor. + output_listener_fd: i32, /// First data-parallel engine rank expected to register on this /// transport. engine_start_index: u32, @@ -336,8 +335,8 @@ impl EngineCoreClient { } TransportMode::Bootstrapped { - input_address, - output_address, + input_listener_fd, + output_listener_fd, engine_start_index, engine_count, ready_timeout, @@ -348,8 +347,8 @@ impl EngineCoreClient { } transport::connect_bootstrapped( - input_address, - output_address, + *input_listener_fd, + *output_listener_fd, *engine_start_index, *engine_count, *ready_timeout, diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 0ac81764a9ed..dbd3fa69ceb2 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -4,6 +4,8 @@ use std::collections::{BTreeMap, BTreeSet}; use std::convert::TryFrom; use std::io::Cursor; +use std::os::fd::IntoRawFd; +use std::os::unix::net::UnixListener; use std::path::PathBuf; use std::process::Command; use std::sync::Once; @@ -311,10 +313,17 @@ fn bootstrapped_test_config( client_index: u32, coordinator_mode: Option, ) -> EngineCoreClientConfig { + fn listener(address: &str) -> i32 { + let path = address.strip_prefix("ipc://").expect("bootstrapped tests use IPC listeners"); + let _ = std::fs::remove_file(path); + let listener = UnixListener::bind(path).expect("bind inherited test listener"); + listener.into_raw_fd() + } + EngineCoreClientConfig { transport_mode: TransportMode::Bootstrapped { - input_address, - output_address, + input_listener_fd: listener(&input_address), + output_listener_fd: listener(&output_address), engine_start_index: 0, engine_count, data_parallel_size: engine_count, @@ -362,15 +371,22 @@ fn bootstrapped_test_config_with_start_index( #[test] fn client_config_validates_bootstrapped_dp_range() { - let mut config = bootstrapped_test_config_with_start_index( - "ipc://unused-input".to_string(), - "ipc://unused-output".to_string(), - 1, - 2, - Duration::from_secs(1), - 0, - None, - ); + // Validation never consumes the descriptors, so placeholders avoid opening + // listeners that nothing would close. + let mut config = EngineCoreClientConfig { + transport_mode: TransportMode::Bootstrapped { + input_listener_fd: -1, + output_listener_fd: -1, + engine_start_index: 1, + engine_count: 2, + data_parallel_size: 3, + ready_timeout: Duration::from_secs(1), + }, + coordinator_mode: None, + model_name: "test-model".to_string(), + client_index: 0, + engine_stats_enabled: true, + }; config.validate().expect("frontend may own a subset of global DP ranks"); let TransportMode::Bootstrapped { diff --git a/rust/src/engine-core-client/src/transport.rs b/rust/src/engine-core-client/src/transport.rs index 7aac8774cac7..e453cafdc8fb 100644 --- a/rust/src/engine-core-client/src/transport.rs +++ b/rust/src/engine-core-client/src/transport.rs @@ -4,6 +4,8 @@ use std::collections::{BTreeMap, BTreeSet}; use std::fmt::Debug; use std::ops::Deref; +use std::os::fd::{FromRawFd, OwnedFd, RawFd}; +use std::os::unix::net::UnixListener as StdUnixListener; use std::time::Duration; use bytes::Bytes; @@ -327,8 +329,8 @@ pub async fn connect_handshake( /// identities are synthesized from contiguous rank order instead of being /// discovered through a Rust-owned handshake. pub async fn connect_bootstrapped( - input_address: &str, - output_address: &str, + input_listener_fd: RawFd, + output_listener_fd: RawFd, engine_start_index: u32, engine_count: usize, ready_timeout: Duration, @@ -354,10 +356,16 @@ pub async fn connect_bootstrapped( } let mut input_socket = RouterSocket::new(); - let input_address = input_socket.bind(input_address).await?.to_string(); + let input_address = input_socket + .bind_listener(adopt_listener(input_listener_fd)?) + .await? + .to_string(); let mut output_socket = PullSocket::new(); - let output_address = output_socket.bind(output_address).await?.to_string(); + let output_address = output_socket + .bind_listener(adopt_listener(output_listener_fd)?) + .await? + .to_string(); let engines = wait_for_input_registrations( &mut input_socket, @@ -385,6 +393,34 @@ pub async fn connect_bootstrapped( }) } +fn adopt_listener(fd: RawFd) -> Result { + if fd < 0 || unsafe { libc::fcntl(fd, libc::F_GETFD) } < 0 { + return Err(Error::InvalidClientConfig { + message: format!("inherited ZMQ listener fd {fd} is not open"), + }); + } + + // SAFETY: Python's Popen(pass_fds) gives the frontend its own descriptor. + // The bootstrapped startup path calls this function once per descriptor, + // transferring ownership to the listener constructed below. + let fd = unsafe { OwnedFd::from_raw_fd(fd) }; + let socket: socket2::Socket = fd.into(); + let domain = socket.domain()?; + + if domain == socket2::Domain::IPV4 || domain == socket2::Domain::IPV6 { + let listener: std::net::TcpListener = socket.into(); + return Ok(listener.into()); + } + if domain == socket2::Domain::UNIX { + let listener: StdUnixListener = socket.into(); + return Ok(listener.into()); + } + + Err(Error::InvalidClientConfig { + message: format!("unsupported inherited listener domain {domain:?}"), + }) +} + /// Bind new input and output sockets. async fn bind_local_sockets( local_host: &str, @@ -603,7 +639,12 @@ pub async fn run_output_loop( #[cfg(test)] mod tests { - use super::bind_local_sockets; + use std::os::fd::IntoRawFd; + + use zeromq::RouterSocket; + use zeromq::prelude::Socket; + + use super::{adopt_listener, bind_local_sockets}; #[tokio::test] async fn bind_local_sockets_resolves_zero_port_bindings() { @@ -614,4 +655,21 @@ mod tests { assert!(output_address.starts_with("tcp://127.0.0.1:")); assert_ne!(input_address, output_address); } + + #[tokio::test] + async fn inherited_tcp_listener_is_adopted_without_rebinding() { + let listener = std::net::TcpListener::bind(("127.0.0.1", 0)).expect("bind listener"); + let expected = listener.local_addr().expect("listener address"); + let mut router = RouterSocket::new(); + + let endpoint = router + .bind_listener( + adopt_listener(listener.into_raw_fd()).expect("adopt inherited listener fd"), + ) + .await + .expect("bind inherited listener"); + + assert_eq!(endpoint.to_string(), format!("tcp://{expected}")); + assert!(router.close().await.is_empty()); + } } diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index d8d0fb4e6cdd..631eb76b36ab 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -5,6 +5,8 @@ use std::collections::BTreeMap; use std::fs; use std::future::Future; use std::io; +use std::os::fd::IntoRawFd; +use std::os::unix::net::UnixListener; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -2230,10 +2232,14 @@ async fn control_aggregates_multi_engine_capacity() { ready_1.world_size = 12; ready_1.data_parallel_rank = start_rank + 1; + let listener_fd = |address: &str| { + let path = address.strip_prefix("ipc://").expect("IPC test endpoint"); + UnixListener::bind(path).expect("bind inherited test listener").into_raw_fd() + }; let client_config = EngineCoreClientConfig { transport_mode: TransportMode::Bootstrapped { - input_address: input_address.clone(), - output_address: output_address.clone(), + input_listener_fd: listener_fd(&input_address), + output_listener_fd: listener_fd(&output_address), engine_start_index: start_rank, engine_count: 2, data_parallel_size: global_size, diff --git a/tests/entrypoints/launchers/api_server/_api_server_spawn_workers.py b/tests/entrypoints/launchers/api_server/_api_server_spawn_workers.py deleted file mode 100644 index c8ff969a8a36..000000000000 --- a/tests/entrypoints/launchers/api_server/_api_server_spawn_workers.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Spawn worker targets kept free of heavy imports. - -``multiprocessing`` with the ``spawn`` start method re-imports the module that -defines a process target in the child. Housing these stubs in a stdlib-only -module keeps child startup fast and deterministic, instead of paying a multi- -second ``import vllm`` before the child can run. -""" - - -def exit_before_report_worker(listen_address, sock, args, client_config=None): - """Exit immediately without touching ``actual_address_pipe``.""" - return diff --git a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py index 6c579df8383c..318f57b86ff7 100644 --- a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py +++ b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import multiprocessing +import os import socket import threading import time @@ -11,14 +12,10 @@ import pytest import zmq -from tests.entrypoints.launchers.api_server._api_server_spawn_workers import ( - exit_before_report_worker, -) -from vllm.utils.network_utils import make_zmq_socket, split_zmq_path +from vllm.utils.network_utils import make_zmq_listener, make_zmq_socket from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.utils import ( APIServerProcessManager, - get_engine_client_zmq_addr, wait_for_completion_or_failure, ) @@ -46,36 +43,30 @@ def update_admission_stats_worker(listen_address, sock, args, client_config): args.results.put(stats.get_num_requests()) -# Module-level stub for the gather_actual_addresses test. Must be -# importable by `multiprocessing.spawn` (no closures, no nesting). -def defer_addresses_stub_worker(listen_address, sock, args, client_config): - """Bind ROUTER/PULL with a kernel-assigned port, report the actual - endpoints back via the pipe, then exit.""" +def _adopt_zmq_listener_worker(listen_address, sock, args, client_config): ctx = zmq.Context() + router = make_zmq_socket( + ctx, + client_config["input_address"], + zmq.ROUTER, + bind=True, + listener=client_config["input_listener"], + ) + pull = make_zmq_socket( + ctx, + client_config["output_address"], + zmq.PULL, + bind=True, + listener=client_config["output_listener"], + ) try: - in_sock = make_zmq_socket( - ctx, client_config["input_address"], zmq.ROUTER, bind=True - ) - out_sock = make_zmq_socket( - ctx, client_config["output_address"], zmq.PULL, bind=True - ) - try: - pipe = client_config["actual_address_pipe"] - try: - pipe.send( - { - "input_address": in_sock.getsockopt(zmq.LAST_ENDPOINT).decode(), - "output_address": out_sock.getsockopt( - zmq.LAST_ENDPOINT - ).decode(), - } - ) - finally: - pipe.close() - finally: - in_sock.close(linger=0) - out_sock.close(linger=0) + identity, payload = router.recv_multipart() + assert payload == b"input" + assert pull.recv() == b"output" + router.send_multipart([identity, b"ok"]) finally: + router.close(linger=0) + pull.close(linger=0) ctx.term() @@ -89,15 +80,11 @@ def api_server_args(): "sock": sock, "args": "test_args", # Simple string to avoid pickling issues "num_servers": 3, - "input_addresses": [ - "tcp://127.0.0.1:5001", - "tcp://127.0.0.1:5002", - "tcp://127.0.0.1:5003", + "input_listeners": [ + make_zmq_listener("tcp://127.0.0.1:0", zmq.ROUTER) for _ in range(3) ], - "output_addresses": [ - "tcp://127.0.0.1:6001", - "tcp://127.0.0.1:6002", - "tcp://127.0.0.1:6003", + "output_listeners": [ + make_zmq_listener("tcp://127.0.0.1:0", zmq.PULL) for _ in range(3) ], "stats_update_address": "tcp://127.0.0.1:7000", } @@ -120,6 +107,13 @@ def test_api_server_process_manager_init(api_server_args, with_stats_update): try: # Verify the manager was initialized correctly assert len(manager.processes) == 3 + assert all( + listener.socket.fileno() == -1 + for listener in ( + *args["input_listeners"], + *args["output_listeners"], + ) + ) # Verify all processes are running for proc in manager.processes: @@ -145,6 +139,70 @@ def test_api_server_process_manager_init(api_server_args, with_stats_update): assert not proc.is_alive() +def _assert_endpoint_reserved(address: str) -> None: + """A competing bind must fail while the endpoint has a listener owner.""" + if address.startswith("ipc://"): + competitor = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + bind_address: str | tuple[str, int] = address.removeprefix("ipc://") + else: + competitor = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + bind_address = ("127.0.0.1", int(address.rsplit(":", 1)[1])) + with competitor, pytest.raises(OSError, match="Address already in use"): + competitor.bind(bind_address) + + +@pytest.mark.parametrize("scheme", ["tcp", "ipc"]) +def test_api_server_child_adopts_inherited_zmq_listeners(tmp_path, scheme): + if scheme == "tcp": + input_path = output_path = "tcp://127.0.0.1:0" + else: + input_path = f"ipc://{tmp_path}/input.sock" + output_path = f"ipc://{tmp_path}/output.sock" + input_listener = make_zmq_listener(input_path, zmq.ROUTER) + output_listener = make_zmq_listener(output_path, zmq.PULL) + input_address = input_listener.address + output_address = output_listener.address + # The supervisor reserves each endpoint as soon as it allocates it. + _assert_endpoint_reserved(input_address) + _assert_endpoint_reserved(output_address) + http_socket = socket.socket() + manager = APIServerProcessManager( + target_server_fn=_adopt_zmq_listener_worker, + listen_address="localhost:8000", + sock=http_socket, + args="test_args", + num_servers=1, + input_listeners=[input_listener], + output_listeners=[output_listener], + ) + + ctx = zmq.Context() + dealer = make_zmq_socket( + ctx, input_address, zmq.DEALER, bind=False, identity=b"test" + ) + push = make_zmq_socket(ctx, output_address, zmq.PUSH, bind=False) + dealer.setsockopt(zmq.RCVTIMEO, 30_000) + try: + assert input_listener.socket.fileno() == -1 + assert output_listener.socket.fileno() == -1 + # After the parent closes its copies, the child's inherited descriptors + # keep both endpoints reserved; closing must not unlink IPC paths. + assert manager.processes[0].is_alive() + _assert_endpoint_reserved(input_address) + _assert_endpoint_reserved(output_address) + dealer.send(b"input") + push.send(b"output") + assert dealer.recv() == b"ok" + manager.processes[0].join(timeout=10) + assert manager.processes[0].exitcode == 0 + finally: + manager.shutdown() + dealer.close(linger=0) + push.close(linger=0) + ctx.term() + http_socket.close() + + @pytest.mark.skip_global_cleanup def test_api_server_processes_share_admission_stats(api_server_args): args = api_server_args.copy() @@ -354,96 +412,7 @@ def run_with_exception_capture(): WORKER_RUNTIME_SECONDS = prev_worker_runtime -@pytest.mark.timeout(60) -def test_gather_actual_addresses_end_to_end(): - """Each child binds ROUTER/PULL with a kernel-picked port and reports - the bound endpoints back via its per-child pipe; the manager surfaces - them via :py:meth:`gather_actual_addresses`.""" - host = "127.0.0.1" - num_servers = 4 - - placeholder_inputs = [ - get_engine_client_zmq_addr(local_only=False, host=host) - for _ in range(num_servers) - ] - placeholder_outputs = [ - get_engine_client_zmq_addr(local_only=False, host=host) - for _ in range(num_servers) - ] - for addr in placeholder_inputs + placeholder_outputs: - assert addr == f"tcp://{host}:0", addr - - sock = socket.socket() - manager = APIServerProcessManager( - listen_address=f"tcp://{host}:0", - sock=sock, - args="test_args", - num_servers=num_servers, - input_addresses=placeholder_inputs, - output_addresses=placeholder_outputs, - target_server_fn=defer_addresses_stub_worker, - ) - - try: - assert len(manager.processes) == num_servers - actual_inputs, actual_outputs = manager.gather_actual_addresses(timeout=15.0) - finally: - manager.shutdown() - time.sleep(0.2) - sock.close() - - assert len(actual_inputs) == num_servers - assert len(actual_outputs) == num_servers - - for addr in actual_inputs + actual_outputs: - scheme, parsed_host, port = split_zmq_path(addr) - assert scheme == "tcp", addr - assert parsed_host == host, addr - assert port and int(port) > 0, addr - - all_addrs = actual_inputs + actual_outputs - assert len(set(all_addrs)) == len(all_addrs), all_addrs - - -@pytest.mark.timeout(30) -def test_gather_actual_addresses_child_crash_before_report(): - """A child that exits before sending its endpoints must surface a - clear ``RuntimeError`` rather than hang or return ``None`` slots.""" - host = "127.0.0.1" - num_servers = 2 - placeholder_inputs = [ - get_engine_client_zmq_addr(local_only=False, host=host) - for _ in range(num_servers) - ] - placeholder_outputs = [ - get_engine_client_zmq_addr(local_only=False, host=host) - for _ in range(num_servers) - ] - - sock = socket.socket() - manager = APIServerProcessManager( - listen_address=f"tcp://{host}:0", - sock=sock, - args="test_args", - num_servers=num_servers, - input_addresses=placeholder_inputs, - output_addresses=placeholder_outputs, - # exit_before_report_worker exits without touching - # ``actual_address_pipe`` — simulates a child that dies before - # reporting its bound addresses. - target_server_fn=exit_before_report_worker, - ) - try: - # Sentinel-first vs pipe-EOF-first both produce "reporting". - with pytest.raises(RuntimeError, match="reporting"): - manager.gather_actual_addresses(timeout=10.0) - finally: - manager.shutdown() - time.sleep(0.2) - sock.close() - - -def test_rust_frontend_launch_log_redacts_credentials(monkeypatch, caplog): +def test_rust_frontend_launch_log_redacts_credentials(monkeypatch): """The Rust frontend command carries every non-default arg as JSON. Credentials must not reach the log line.""" import subprocess as subprocess_mod @@ -469,15 +438,17 @@ def wait(self, timeout=None): return 0 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + input_listener = make_zmq_listener("tcp://127.0.0.1:0", zmq.ROUTER) + output_listener = make_zmq_listener("tcp://127.0.0.1:0", zmq.PULL) try: monkeypatch.setattr(subprocess_mod, "Popen", lambda *a, **kw: _FakeProc()) - with caplog.at_level("INFO", logger="vllm.v1.utils"): + with patch("vllm.v1.utils.logger.info") as log_info: RustFrontendProcessManager( binary_path="/nonexistent/vllm-rs", sock=sock, args=args, - input_address="ipc:///tmp/in", - output_address="ipc:///tmp/out", + input_listener=input_listener, + output_listener=output_listener, engine_start_index=0, engine_count=1, data_parallel_size=1, @@ -485,11 +456,13 @@ def wait(self, timeout=None): finally: sock.close() - message = caplog.text + log_format, *log_args = log_info.call_args.args + message = log_format % tuple(log_args) assert "Launching Rust frontend:" in message assert hf_token not in message assert api_key not in message assert '"hf_token": "***"' in message + assert '"api_key": "***"' in message def test_rust_frontend_uses_config_model_as_model_tag(monkeypatch, caplog, tmp_path): @@ -519,6 +492,8 @@ def wait(self, timeout=None): return 0 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + input_listener = make_zmq_listener("tcp://127.0.0.1:0", zmq.ROUTER) + output_listener = make_zmq_listener("tcp://127.0.0.1:0", zmq.PULL) try: monkeypatch.setattr(subprocess_mod, "Popen", lambda *a, **kw: _FakeProc()) with caplog.at_level("INFO", logger="vllm.v1.utils"): @@ -526,8 +501,8 @@ def wait(self, timeout=None): binary_path="/nonexistent/vllm-rs", sock=sock, args=args, - input_address="ipc:///tmp/in", - output_address="ipc:///tmp/out", + input_listener=input_listener, + output_listener=output_listener, engine_start_index=0, engine_count=1, data_parallel_size=1, @@ -568,14 +543,18 @@ def fake_popen(cmd, **kwargs): sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) grpc_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + input_listener = make_zmq_listener("tcp://127.0.0.1:0", zmq.ROUTER) + output_listener = make_zmq_listener("tcp://127.0.0.1:0", zmq.PULL) + input_fd = input_listener.socket.fileno() + output_fd = output_listener.socket.fileno() try: monkeypatch.setattr(subprocess_mod, "Popen", fake_popen) RustFrontendProcessManager( binary_path="/nonexistent/vllm-rs", sock=sock, args=args, - input_address="ipc:///tmp/in", - output_address="ipc:///tmp/out", + input_listener=input_listener, + output_listener=output_listener, engine_start_index=0, engine_count=1, data_parallel_size=1, @@ -585,7 +564,44 @@ def fake_popen(cmd, **kwargs): grpc_fd = grpc_sock.fileno() flag_index = cmd.index("--grpc-listen-fd") assert cmd[flag_index + 1] == str(grpc_fd) - assert list(kwargs["pass_fds"]) == [sock.fileno(), grpc_fd] + assert list(kwargs["pass_fds"]) == [ + sock.fileno(), + input_fd, + output_fd, + grpc_fd, + ] finally: sock.close() grpc_sock.close() + + +def test_rust_frontend_spawn_failure_cleans_up_listeners(tmp_path): + """No child owns the listeners when Popen fails, so the parent removes them.""" + from vllm.entrypoints.launchers.cli_args import make_arg_parser + from vllm.utils.argparse_utils import FlexibleArgumentParser + from vllm.v1.utils import RustFrontendProcessManager + + args = make_arg_parser(FlexibleArgumentParser()).parse_args( + ["--model", "org/model"] + ) + input_listener = make_zmq_listener(f"ipc://{tmp_path}/input.sock", zmq.ROUTER) + output_listener = make_zmq_listener(f"ipc://{tmp_path}/output.sock", zmq.PULL) + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + with pytest.raises(FileNotFoundError): + RustFrontendProcessManager( + binary_path=str(tmp_path / "missing-vllm-rs"), + sock=sock, + args=args, + input_listener=input_listener, + output_listener=output_listener, + engine_start_index=0, + engine_count=1, + data_parallel_size=1, + ) + finally: + sock.close() + + for listener in (input_listener, output_listener): + assert listener.socket.fileno() == -1 + assert not os.path.exists(listener.address.removeprefix("ipc://")) diff --git a/tests/utils_/test_network_utils.py b/tests/utils_/test_network_utils.py index a6c359e26930..7ba3d6ae3ecc 100644 --- a/tests/utils_/test_network_utils.py +++ b/tests/utils_/test_network_utils.py @@ -13,6 +13,7 @@ get_open_ports_list, get_tcp_uri, join_host_port, + make_zmq_listener, make_zmq_path, make_zmq_socket, split_host_port, @@ -205,6 +206,52 @@ def test_make_zmq_socket_ipv6(): ctx.term() +@pytest.mark.parametrize("scheme", ["tcp", "ipc"]) +def test_make_zmq_socket_adopts_listener(tmp_path, scheme): + path = ( + "tcp://127.0.0.1:0" if scheme == "tcp" else f"ipc://{tmp_path}/inherited.sock" + ) + listener = make_zmq_listener(path, zmq.ROUTER) + ctx = zmq.Context() + router = dealer = None + try: + router = make_zmq_socket( + ctx, + listener.address, + zmq.ROUTER, + bind=True, + listener=listener.socket, + ) + assert listener.socket.fileno() == -1 + dealer = make_zmq_socket( + ctx, listener.address, zmq.DEALER, bind=False, identity=b"client" + ) + dealer.send(b"ping") + assert router.recv_multipart() == [b"client", b"ping"] + finally: + if router is not None: + router.close(linger=0) + if dealer is not None: + dealer.close(linger=0) + ctx.term() + listener.cleanup() + + +def test_make_zmq_listener_failed_ipc_bind_keeps_existing_path(tmp_path): + path = tmp_path / "taken.sock" + owner = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + owner.bind(str(path)) + owner.listen() + try: + with pytest.raises(OSError, match="Address already in use"): + make_zmq_listener(f"ipc://{path}", zmq.ROUTER) + + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client: + client.connect(str(path)) + finally: + owner.close() + + def test_make_zmq_path(): assert make_zmq_path("tcp", "127.0.0.1", 5555) == "tcp://127.0.0.1:5555" assert make_zmq_path("tcp", "::1", 5555) == "tcp://[::1]:5555" diff --git a/tests/v1/engine/test_core_engine_actor_manager.py b/tests/v1/engine/test_core_engine_actor_manager.py index 768f51f4379b..1b23356134b0 100644 --- a/tests/v1/engine/test_core_engine_actor_manager.py +++ b/tests/v1/engine/test_core_engine_actor_manager.py @@ -23,7 +23,7 @@ EngineZmqAddresses, _dp_nodes_master_first, _node_ip_from_resources, - get_engine_zmq_addresses, + bind_engine_zmq_listeners, launch_core_engines, ) from vllm.v1.utils import APIServerProcessManager @@ -67,35 +67,27 @@ def get_addresses(self) -> tuple[list[str], list[str]]: # Module-level stub worker for the Ray-DP regression test. Must be importable -# by ``multiprocessing.spawn`` (no closures, no nesting). Mirrors the worker -# in ``tests/entrypoints/test_api_server_process_manager.py``. -def _bind_and_report_worker(listen_address, sock, args, client_config): - """Bind ROUTER/PULL with a kernel-assigned port, report the actual - endpoints back via ``actual_address_pipe``, then exit.""" +# by ``multiprocessing.spawn`` (no closures, no nesting). +def _adopt_listener_worker(listen_address, sock, args, client_config): + """Adopt the supervisor-bound ROUTER/PULL listeners, then exit.""" ctx = zmq.Context() try: in_sock = make_zmq_socket( - ctx, client_config["input_address"], zmq.ROUTER, bind=True + ctx, + client_config["input_address"], + zmq.ROUTER, + bind=True, + listener=client_config["input_listener"], ) out_sock = make_zmq_socket( - ctx, client_config["output_address"], zmq.PULL, bind=True + ctx, + client_config["output_address"], + zmq.PULL, + bind=True, + listener=client_config["output_listener"], ) - try: - pipe = client_config["actual_address_pipe"] - try: - pipe.send( - { - "input_address": in_sock.getsockopt(zmq.LAST_ENDPOINT).decode(), - "output_address": out_sock.getsockopt( - zmq.LAST_ENDPOINT - ).decode(), - } - ) - finally: - pipe.close() - finally: - in_sock.close(linger=0) - out_sock.close(linger=0) + in_sock.close(linger=0) + out_sock.close(linger=0) finally: ctx.term() @@ -327,17 +319,9 @@ def test_ray_dp_addresses_resolved_before_actor_creation( ``launch_core_engines`` Ray branch pickles ``addresses`` into each engine actor at ``.remote()`` time, and ``EngineCoreActorMixin._perform_handshakes`` is a no-op, so the actor uses that pickled snapshot for the rest of its - life. If ``run_multi_api_server`` allocates ``addresses`` as - ``tcp://host:0`` placeholders (its default), the actors hold placeholders - forever and DEALER-connect to port 0 — ZMQ ``connect`` is async and does - not raise, so the failure mode is a deterministic hang. - - The Ray-DP carve-out in ``run_multi_api_server`` forces - ``defer_api_server_ports=False`` when ``data_parallel_backend == "ray"`` - so addresses are pre-allocated in the driver and Ray pickles real ports - into each actor. This test mirrors that call-site logic and asserts the - actors hold real (non-placeholder) endpoints. If the carve-out is - removed without an alternative fix, the test fails. + life. The supervisor therefore binds the frontend listeners before actor + construction and Ray pickles their resolved endpoints. This test asserts + that each actor holds real (non-placeholder) endpoints. """ created_placement_groups: list[Any] = [] @@ -356,15 +340,10 @@ def create_dp_placement_groups(vllm_config: Any): vllm_config = _make_vllm_config_ray_dp_multinode() - # Mirror run_multi_api_server's address-allocation logic. The Ray DP - # carve-out forces pre-allocation so the addresses pickled into engine - # actors at .remote() time are real, not ``tcp://host:0``. - is_ray_dp = vllm_config.parallel_config.data_parallel_backend == "ray" - addresses = get_engine_zmq_addresses( - vllm_config, - num_api_servers=2, - defer_api_server_ports=not is_ray_dp, - ) + # Mirror run_multi_api_server: bind listeners in the supervisor so the + # addresses pickled into Ray actors contain kernel-assigned ports. + listeners = bind_engine_zmq_listeners(vllm_config, num_api_servers=2) + addresses = listeners.addresses sock = socket.socket() engine_manager: CoreEngineActorManager | None = None @@ -381,26 +360,17 @@ def create_dp_placement_groups(vllm_config: Any): engine_manager = engine_launch.engine_manager assert isinstance(engine_manager, CoreEngineActorManager) - # API-server children bind to the pre-allocated ports. + # API-server children adopt the already-bound listener FDs. api_server_manager = APIServerProcessManager( listen_address="tcp://127.0.0.1:0", sock=sock, args="test_args", num_servers=2, - input_addresses=addresses.inputs, - output_addresses=addresses.outputs, - target_server_fn=_bind_and_report_worker, + input_listeners=listeners.inputs, + output_listeners=listeners.outputs, + target_server_fn=_adopt_listener_worker, ) - # run_multi_api_server skips ``gather_actual_addresses`` for - # Ray DP (addresses are already real). Mirror that. - if not is_ray_dp: - actual_inputs, actual_outputs = ( - api_server_manager.gather_actual_addresses(timeout=15.0) - ) - addresses.inputs = actual_inputs - addresses.outputs = actual_outputs - # Snapshot what each Ray actor actually holds. actors = ( engine_manager.local_engine_actors + engine_manager.remote_engine_actors @@ -427,9 +397,9 @@ def create_dp_placement_groups(vllm_config: Any): assert scheme == "tcp", url assert port and int(port) > 0, ( f"Ray actor was pickled with placeholder address {url!r}; " - "``run_multi_api_server`` must pre-allocate ports for the " - "Ray DP backend so the actors hold real endpoints by the " - "time they DEALER-connect. See PR #42585 / Ray-DP " + "``run_multi_api_server`` must bind frontend listeners before " + "constructing Ray DP actors so they hold real endpoints when " + "they DEALER-connect. See PR #42585 / Ray-DP " "multi-API-server regression." ) diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py index 4c9410488b3f..94e6545c0899 100644 --- a/vllm/entrypoints/cli/serve.py +++ b/vllm/entrypoints/cli/serve.py @@ -328,19 +328,10 @@ def signal_handler(signum, frame): None ) - from vllm.v1.engine.utils import get_engine_zmq_addresses - - # Defer port allocation to the child's bind() to avoid TOCTOU, except - # for Rust front-end and Ray DP, which can't see the post-bind rebind - # (CLI-arg subprocess / pickled-into-actor snapshot respectively) and - # so pre-allocate driver-side -- reintroducing the original race only - # there. - is_ray_dp = parallel_config.data_parallel_backend == "ray" - addresses = get_engine_zmq_addresses( - vllm_config, - num_api_servers, - defer_api_server_ports=not (rust_frontend_path or is_ray_dp), - ) + from vllm.v1.engine.utils import bind_engine_zmq_listeners + + zmq_listeners = bind_engine_zmq_listeners(vllm_config, num_api_servers) + addresses = zmq_listeners.addresses with launch_core_engines( vllm_config, executor_class, log_stats, addresses @@ -365,8 +356,8 @@ def signal_handler(signum, frame): sock=sock, grpc_sock=grpc_sock, args=args, - input_address=addresses.inputs[0], - output_address=addresses.outputs[0], + input_listener=zmq_listeners.inputs[0], + output_listener=zmq_listeners.outputs[0], engine_start_index=expected_engine_start_index, engine_count=expected_engine_count, data_parallel_size=parallel_config.data_parallel_size, @@ -379,22 +370,12 @@ def signal_handler(signum, frame): sock=sock, args=args, num_servers=num_api_servers, - input_addresses=addresses.inputs, - output_addresses=addresses.outputs, + input_listeners=zmq_listeners.inputs, + output_listeners=zmq_listeners.outputs, stats_update_address=stats_update_address, tensor_queue=engine_launch.tensor_queue, ) - if not is_ray_dp: - # Forward each child's bound endpoints to the engine handshake - # (runs on ``with`` exit). Skipped for Ray DP, where addresses - # are pre-allocated above and Ray actors already hold them. - actual_inputs, actual_outputs = ( - api_server_manager.gather_actual_addresses() - ) - addresses.inputs = actual_inputs - addresses.outputs = actual_outputs - # Set frontend processes to watch during engine startup. # If any of these processes exit before the engines are up, the engine startup # will be aborted with an error. diff --git a/vllm/utils/network_utils.py b/vllm/utils/network_utils.py index fa24d4d5f9ab..6bc4f9c9fdf8 100644 --- a/vllm/utils/network_utils.py +++ b/vllm/utils/network_utils.py @@ -11,6 +11,7 @@ Iterator, Sequence, ) +from dataclasses import dataclass from typing import Any from uuid import uuid4 @@ -25,6 +26,23 @@ logger = init_logger(__name__) +@dataclass +class ZmqListener: + """A supervisor-bound listener that a ZMQ socket can adopt.""" + + address: str + socket: socket.socket + + def close(self) -> None: + self.socket.close() + + def cleanup(self) -> None: + self.close() + if self.address.startswith("ipc://"): + with contextlib.suppress(FileNotFoundError): + os.unlink(self.address.removeprefix("ipc://")) + + def close_sockets(sockets: Sequence[zmq.Socket | zmq.asyncio.Socket]): for sock in sockets: if sock is not None: @@ -307,6 +325,27 @@ def make_zmq_path(scheme: str, host: str, port: int | None = None) -> str: # Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L783 # noqa: E501 +# Apply the existing make_zmq_socket memory policy to every socket owner. +_ZMQ_LARGE_BUFFER_SIZE = 512 * 1024**2 +_ZMQ_LARGE_BUFFER_MIN_TOTAL_MEMORY = 32 * 1024**3 +_ZMQ_LARGE_BUFFER_MIN_AVAILABLE_MEMORY = 16 * 1024**3 + + +def _get_zmq_socket_buffer_size() -> int: + """Choose the shared libzmq and inherited-listener buffer policy.""" + mem = psutil.virtual_memory() + # For systems with substantial memory (>32GB total, >16GB available): + # - Set a large 0.5GB buffer to improve throughput + # For systems with less memory: + # - Use system default (-1) to avoid excessive memory consumption + if ( + mem.total > _ZMQ_LARGE_BUFFER_MIN_TOTAL_MEMORY + and mem.available > _ZMQ_LARGE_BUFFER_MIN_AVAILABLE_MEMORY + ): + return _ZMQ_LARGE_BUFFER_SIZE + return -1 + + def make_zmq_socket( ctx: zmq.asyncio.Context | zmq.Context, # type: ignore[name-defined] path: str, @@ -315,22 +354,19 @@ def make_zmq_socket( identity: bytes | None = None, linger: int | None = None, router_handover: bool = False, + listener: socket.socket | None = None, ) -> zmq.Socket | zmq.asyncio.Socket: # type: ignore[name-defined] - """Make a ZMQ socket with the proper bind/connect semantics.""" - mem = psutil.virtual_memory() - socket = ctx.socket(socket_type) + """Make a ZMQ socket with the proper bind/connect semantics. - # Calculate buffer size based on system memory - total_mem = mem.total / 1024**3 - available_mem = mem.available / 1024**3 - # For systems with substantial memory (>32GB total, >16GB available): - # - Set a large 0.5GB buffer to improve throughput - # For systems with less memory: - # - Use system default (-1) to avoid excessive memory consumption - buf_size = int(0.5 * 1024**3) if total_mem > 32 and available_mem > 16 else -1 + When supplied, ``listener`` is detached and its fd moves to libzmq. + """ + socket = ctx.socket(socket_type) + buf_size = _get_zmq_socket_buffer_size() if bind is None: bind = socket_type not in (zmq.PUSH, zmq.SUB, zmq.XSUB) + if listener is not None and not bind: + raise ValueError("An inherited ZMQ listener requires bind=True") if socket_type in (zmq.PULL, zmq.DEALER, zmq.ROUTER): socket.setsockopt(zmq.RCVHWM, 0) @@ -353,6 +389,9 @@ def make_zmq_socket( if socket_type == zmq.XPUB: socket.setsockopt(zmq.XPUB_VERBOSE, True) + if listener is not None: + socket.setsockopt(zmq.USE_FD, listener.detach()) + # Determine if the path is a TCP socket with an IPv6 address. # Enable IPv6 on the zmq socket if so. scheme, host, _ = split_zmq_path(path) @@ -367,6 +406,48 @@ def make_zmq_socket( return socket +def make_zmq_listener(path: str, socket_type: Any) -> ZmqListener: + """Bind and listen on a raw socket for later ZMQ adoption.""" + scheme, host, port = split_zmq_path(path) + if scheme == "tcp": + family = socket.AF_INET6 if is_valid_ipv6_address(host) else socket.AF_INET + bind_address: str | tuple[str, int] = (host, int(port)) + elif scheme == "ipc": + family = socket.AF_UNIX + bind_address = path.removeprefix("ipc://") + os.makedirs(os.path.dirname(bind_address), exist_ok=True) + else: + raise ValueError(f"Cannot inherit a {scheme} ZMQ listener") + + listener = socket.socket(family, socket.SOCK_STREAM) + bound = False + try: + buf_size = _get_zmq_socket_buffer_size() + if buf_size >= 0: + # Accepted stream sockets inherit these settings from the listener. + # ZMQ_USE_FD adopts this pre-created socket, so the policy belongs + # here as well as on the eventual libzmq socket. + if socket_type in (zmq.PULL, zmq.DEALER, zmq.ROUTER): + listener.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, buf_size) + if socket_type in (zmq.PUSH, zmq.DEALER, zmq.ROUTER): + listener.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, buf_size) + + listener.bind(bind_address) + bound = True + listener.listen() + if scheme == "tcp": + path = get_tcp_uri(host, listener.getsockname()[1]) + return ZmqListener(address=path, socket=listener) + except BaseException: + listener.close() + # A failed bind leaves any existing pathname owned by its listener. + if scheme == "ipc" and bound: + assert isinstance(bind_address, str) + with contextlib.suppress(FileNotFoundError): + os.unlink(bind_address) + raise + + @contextlib.contextmanager def zmq_socket_ctx( path: str, diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 88fdc1d1e7f1..74520036a7c8 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -11,7 +11,6 @@ from collections.abc import Awaitable, Callable, Sequence from concurrent.futures import Future from dataclasses import dataclass -from multiprocessing.connection import Connection from multiprocessing.queues import Queue from threading import Thread from typing import Any, TypeAlias, TypeVar @@ -56,7 +55,8 @@ from vllm.v1.engine.utils import ( CoreEngineActorManager, CoreEngineProcManager, - get_engine_zmq_addresses, + EngineZmqAddresses, + get_engine_zmq_bind_addresses, launch_core_engines, ) from vllm.v1.executor import Executor @@ -634,63 +634,50 @@ def __init__( self.stats_update_address = client_addresses.get("stats_update_address") # Tensor queues passed via client_addresses for multi-API-server case tensor_queue = client_addresses.get("tensor_queue") + input_listener = client_addresses.get("input_listener") + output_listener = client_addresses.get("output_listener") self.input_socket = self.resources.input_socket = make_zmq_socket( self.ctx, input_address, zmq.ROUTER, bind=True, router_handover=enable_input_socket_handover, + listener=input_listener, ) self.resources.output_socket = make_zmq_socket( - self.ctx, output_address, zmq.PULL + self.ctx, + output_address, + zmq.PULL, + listener=output_listener, ) - # Report bound endpoints back so the parent can forward - # them to engines (mirrors the DPCoordinator pattern). - actual_address_pipe: Connection | None = client_addresses.get( - "actual_address_pipe" - ) - if actual_address_pipe is not None: - try: - actual_input = self.input_socket.getsockopt( - zmq.LAST_ENDPOINT - ).decode() - actual_output = self.resources.output_socket.getsockopt( - zmq.LAST_ENDPOINT - ).decode() - actual_address_pipe.send( - { - "input_address": actual_input, - "output_address": actual_output, - } - ) - finally: - actual_address_pipe.close() # Engines are managed externally: this process does not fork # them, so there is no fork race; start the MM warmup now. self._start_mm_warmup() else: # Engines are managed by this client. - addresses = get_engine_zmq_addresses(vllm_config) + bind_addresses = get_engine_zmq_bind_addresses(vllm_config) self.input_socket = self.resources.input_socket = make_zmq_socket( self.ctx, - addresses.inputs[0], + bind_addresses.inputs[0], zmq.ROUTER, bind=True, router_handover=enable_input_socket_handover, ) self.resources.output_socket = make_zmq_socket( - self.ctx, addresses.outputs[0], zmq.PULL + self.ctx, bind_addresses.outputs[0], zmq.PULL ) - # Resolve ``tcp://host:0`` placeholders to bound endpoints - # before engines DEALER-connect. No-op for IPC. - addresses.inputs[0] = self.input_socket.getsockopt( - zmq.LAST_ENDPOINT - ).decode() - addresses.outputs[0] = self.resources.output_socket.getsockopt( - zmq.LAST_ENDPOINT - ).decode() + # EngineZmqAddresses contains only endpoints resolved by the + # sockets that own the bind operation. + addresses = EngineZmqAddresses( + inputs=[self.input_socket.getsockopt(zmq.LAST_ENDPOINT).decode()], + outputs=[ + self.resources.output_socket.getsockopt( + zmq.LAST_ENDPOINT + ).decode() + ], + ) with launch_core_engines( vllm_config, executor_class, log_stats, addresses diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index 49e6884ac9f1..4f1be3bacd1c 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -23,9 +23,10 @@ from vllm.ray.ray_env import get_env_vars_to_copy from vllm.utils import numa_utils from vllm.utils.network_utils import ( - get_open_port, + ZmqListener, get_open_zmq_ipc_path, get_tcp_uri, + make_zmq_listener, zmq_socket_ctx, ) from vllm.utils.system_utils import get_mp_context, set_env_var @@ -82,8 +83,18 @@ def __init__(self, index: int = 0, local: bool = True): self.state = CoreEngineState.NEW +@dataclass +class EngineZmqBindAddresses: + """Frontend endpoints before their owning sockets bind.""" + + inputs: list[str] + outputs: list[str] + + @dataclass class EngineZmqAddresses: + """Resolved frontend endpoints advertised to engine processes.""" + # ZMQ input socket addresses for each front-end client (requests) inputs: list[str] # ZMQ output socket addresses for each front-end client (responses) @@ -98,6 +109,27 @@ class EngineZmqAddresses: frontend_stats_publish_address: str | None = None +@dataclass +class EngineZmqListeners: + inputs: list[ZmqListener] + outputs: list[ZmqListener] + + @property + def addresses(self) -> EngineZmqAddresses: + return EngineZmqAddresses( + inputs=[listener.address for listener in self.inputs], + outputs=[listener.address for listener in self.outputs], + ) + + def cleanup(self) -> None: + for listener in (*self.inputs, *self.outputs): + listener.cleanup() + + def __del__(self) -> None: + with contextlib.suppress(Exception): + self.cleanup() + + @dataclass class EngineHandshakeMetadata: """Metadata sent to each engine process during startup handshake, @@ -1079,22 +1111,16 @@ def shutdown(self, timeout: float | None = None) -> None: ray.util.remove_placement_group(pg) -def get_engine_zmq_addresses( +def get_engine_zmq_bind_addresses( vllm_config: VllmConfig, num_api_servers: int = 1, - *, - defer_api_server_ports: bool = True, -) -> EngineZmqAddresses: - """Allocate ZMQ addresses for engine-client communication. - - By default each TCP address is a ``tcp://host:0`` placeholder; the - consumer (API-server child or single-process ``MPClient``) binds, then - recovers the kernel-assigned port via ``getsockopt(zmq.LAST_ENDPOINT)`` - and writes it back into ``addresses`` before the engine handshake. - - Set ``defer_api_server_ports=False`` only when the consumer cannot - report a bound port back (e.g. the Rust front-end). IPC paths are - unaffected.""" +) -> EngineZmqBindAddresses: + """Create ZMQ endpoint templates for engine-client communication. + + Each TCP address is a ``tcp://host:0`` placeholder. A single-process + ``MPClient`` binds its own ZMQ sockets. Cross-process frontends use + ``bind_engine_zmq_listeners`` so the supervisor binds raw listeners before + launching engines and transfers listener ownership to each frontend.""" parallel_config = vllm_config.parallel_config local_engine_count = parallel_config.data_parallel_size_local local_start_index = parallel_config.data_parallel_rank_local @@ -1119,14 +1145,39 @@ def get_engine_zmq_addresses( def _addr() -> str: if client_local_only: return get_open_zmq_ipc_path() - return get_tcp_uri(host, 0 if defer_api_server_ports else get_open_port()) + return get_tcp_uri(host, 0) - return EngineZmqAddresses( + return EngineZmqBindAddresses( inputs=[_addr() for _ in range(num_api_servers)], outputs=[_addr() for _ in range(num_api_servers)], ) +def bind_engine_zmq_listeners( + vllm_config: VllmConfig, + num_api_servers: int = 1, +) -> EngineZmqListeners: + """Bind frontend transport listeners before launching engines.""" + bind_addresses = get_engine_zmq_bind_addresses( + vllm_config, + num_api_servers, + ) + inputs: list[ZmqListener] = [] + outputs: list[ZmqListener] = [] + try: + inputs.extend( + make_zmq_listener(address, zmq.ROUTER) for address in bind_addresses.inputs + ) + outputs.extend( + make_zmq_listener(address, zmq.PULL) for address in bind_addresses.outputs + ) + return EngineZmqListeners(inputs=inputs, outputs=outputs) + except BaseException: + for listener in (*inputs, *outputs): + listener.cleanup() + raise + + FrontendProcess = BaseProcess | _SubprocessWrapper diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index ac877263a5f6..868864f0ca41 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -29,7 +29,7 @@ import vllm.envs as envs from vllm.logger import init_logger from vllm.usage.usage_lib import UsageContext, is_usage_stats_enabled, usage_message -from vllm.utils.network_utils import get_open_zmq_ipc_path, get_tcp_uri +from vllm.utils.network_utils import ZmqListener, get_open_zmq_ipc_path, get_tcp_uri from vllm.utils.system_utils import decorate_logs, kill_process_tree, set_process_title from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.core.sched.output import SchedulerOutput @@ -177,28 +177,22 @@ def __init__( sock: Any, args: argparse.Namespace, num_servers: int, - input_addresses: list[str], - output_addresses: list[str], + input_listeners: list[ZmqListener], + output_listeners: list[ZmqListener], target_server_fn: Callable | None = None, stats_update_address: str | None = None, tensor_queue: Queue | None = None, ): """Initialize and start API server worker processes. - ``input_addresses``/``output_addresses`` may contain - ``tcp://host:0`` placeholders; each child must report the actual - bound endpoint over its ``actual_address_pipe`` in ``client_config`` - and the parent collects them via - :py:meth:`gather_actual_addresses`. - Args: target_server_fn: Override function to call for each API server process listen_address: Address to listen for client connections sock: Socket for client connections args: Command line arguments num_servers: Number of API server processes to start - input_addresses: Input addresses for each API server - output_addresses: Output addresses for each API server + input_listeners: Input listeners for each API server + output_listeners: Output listeners for each API server stats_update_address: Optional stats update address tensor_queue: Optional tensor IPC queue for sharing MM tensors @@ -209,7 +203,6 @@ def __init__( spawn_context = multiprocessing.get_context("spawn") self.processes: list[BaseProcess] = [] - self._address_pipes: list[connection.Connection] = [] admission_counters = None if num_servers > 1 and getattr(args, "max_num_queued_reqs", None) is not None: @@ -220,126 +213,65 @@ def __init__( SharedAdmissionStats.num_counters(num_servers), ) - for i, in_addr, out_addr in zip( - range(num_servers), input_addresses, output_addresses - ): - client_config: dict[str, Any] = { - "input_address": in_addr, - "output_address": out_addr, - "client_count": num_servers, - "client_index": i, - } - if admission_counters is not None: - client_config["mp_admission_counters"] = admission_counters - if stats_update_address is not None: - client_config["stats_update_address"] = stats_update_address - if tensor_queue is not None: - client_config["tensor_queue"] = tensor_queue - - parent_recv, child_send = spawn_context.Pipe(duplex=False) - self._address_pipes.append(parent_recv) - client_config["actual_address_pipe"] = child_send - - proc = spawn_context.Process( - target=target_server_fn or run_api_server_worker_proc, - name=f"ApiServer_{i}", - args=(listen_address, sock, args, client_config), + listeners = [*input_listeners, *output_listeners] + self._listeners = listeners + try: + listener_pairs = zip( + range(num_servers), input_listeners, output_listeners, strict=True ) - self.processes.append(proc) - proc.start() - - # Drop parent's write end so reader sees EOF on child death. - child_send.close() + for i, input_listener, output_listener in listener_pairs: + client_config: dict[str, Any] = { + "input_address": input_listener.address, + "output_address": output_listener.address, + "input_listener": input_listener.socket, + "output_listener": output_listener.socket, + "client_count": num_servers, + "client_index": i, + } + if admission_counters is not None: + client_config["mp_admission_counters"] = admission_counters + if stats_update_address is not None: + client_config["stats_update_address"] = stats_update_address + if tensor_queue is not None: + client_config["tensor_queue"] = tensor_queue + + proc = spawn_context.Process( + target=target_server_fn or run_api_server_worker_proc, + name=f"ApiServer_{i}", + args=(listen_address, sock, args, client_config), + ) + self.processes.append(proc) + proc.start() + input_listener.close() + output_listener.close() + except BaseException: + _shutdown_processes_and_cleanup_listeners(self.processes, listeners) + raise logger.info("Started %d API server processes", len(self.processes)) # Shutdown only the API server processes on garbage collection # The extra processes are managed by their owners - self._finalizer = weakref.finalize(self, shutdown, self.processes) - - def gather_actual_addresses( - self, - timeout: float = envs.VLLM_ENGINE_READY_TIMEOUT_S, - ) -> tuple[list[str], list[str]]: - """Return (inputs, outputs) reported by each child, indexed by - ``client_index``. Raises ``RuntimeError`` on timeout or premature - child exit.""" - n = len(self._address_pipes) - inputs: list[str | None] = [None] * n - outputs: list[str | None] = [None] * n - pending: dict[connection.Connection, int] = { - pipe: i for i, pipe in enumerate(self._address_pipes) - } - sentinel_to_idx: dict[Any, int] = { - proc.sentinel: i for i, proc in enumerate(self.processes) - } - - deadline = time.monotonic() + timeout - try: - while pending: - remaining = deadline - time.monotonic() - if remaining <= 0: - missing = [self.processes[i].name for i in pending.values()] - raise RuntimeError( - f"Timed out after {timeout:.1f}s waiting for " - f"API server(s) to report bound ZMQ addresses: " - f"{missing}" - ) - waitables: list[Any] = list(pending.keys()) + list( - sentinel_to_idx.keys() - ) - ready = connection.wait(waitables, timeout=remaining) - # Drain pipes before checking sentinels: a child that sent - # its message and then exited can surface both events in - # the same poll, and we must record the success first. - for item in ready: - if isinstance(item, connection.Connection) and item in pending: - idx = pending.pop(item) - try: - msg: dict[str, str] = item.recv() - except EOFError as e: - raise RuntimeError( - f"API server {self.processes[idx].name} " - f"closed its address pipe without " - f"reporting its bound ZMQ addresses" - ) from e - inputs[idx] = msg["input_address"] - outputs[idx] = msg["output_address"] - item.close() - for item in ready: - if item in sentinel_to_idx: - idx = sentinel_to_idx.pop(item) - pipe = self._address_pipes[idx] - if pipe in pending: - proc = self.processes[idx] - raise RuntimeError( - f"API server process {proc.name} exited " - f"(code={proc.exitcode}) before reporting " - f"its bound ZMQ addresses" - ) - finally: - for pipe in pending: - with contextlib.suppress(Exception): - pipe.close() - - return inputs, outputs # type: ignore[return-value] + self._finalizer = weakref.finalize( + self, + _shutdown_processes_and_cleanup_listeners, + self.processes, + listeners, + ) def shutdown(self, timeout: float | None = None) -> None: """Shutdown API server processes with configurable timeout""" - for pipe in self._address_pipes: - with contextlib.suppress(Exception): - pipe.close() - self._address_pipes = [] - if self._finalizer.detach() is not None: - shutdown(self.processes, timeout=timeout) + _shutdown_processes_and_cleanup_listeners( + self.processes, self._listeners, timeout=timeout + ) class RustFrontendProcessManager: """Manages a single Rust frontend subprocess. Launches the Rust vllm-rs binary in 'frontend' mode, passing the - listening socket fd and ZMQ transport addresses. Provides the same + HTTP and ZMQ listener file descriptors. Provides the same interface as APIServerProcessManager for process monitoring. """ @@ -348,8 +280,8 @@ def __init__( binary_path: str, sock: Any, args: argparse.Namespace, - input_address: str, - output_address: str, + input_listener: ZmqListener, + output_listener: ZmqListener, engine_start_index: int, engine_count: int, data_parallel_size: int, @@ -359,18 +291,22 @@ def __init__( import os import subprocess - fd = sock.fileno() - os.set_inheritable(fd, True) + listen_fd = sock.fileno() + input_fd = input_listener.socket.fileno() + output_fd = output_listener.socket.fileno() + inherited_fds = [listen_fd, input_fd, output_fd] + for fd in inherited_fds: + os.set_inheritable(fd, True) cmd = [ binary_path, "frontend", "--listen-fd", - str(fd), - "--input-address", - input_address, - "--output-address", - output_address, + str(listen_fd), + "--input-listener-fd", + str(input_fd), + "--output-listener-fd", + str(output_fd), "--engine-start-index", str(engine_start_index), "--engine-count", @@ -378,12 +314,11 @@ def __init__( "--data-parallel-size", str(data_parallel_size), ] - pass_fds = [fd] if grpc_sock is not None: grpc_fd = grpc_sock.fileno() os.set_inheritable(grpc_fd, True) cmd.extend(["--grpc-listen-fd", str(grpc_fd)]) - pass_fds.append(grpc_fd) + inherited_fds.append(grpc_fd) if stats_update_address is not None: cmd.extend(["--coordinator-address", stats_update_address]) from vllm.entrypoints.serve.utils.api_utils import jsonify_non_default_args @@ -425,18 +360,34 @@ def __init__( redacted_json = json.dumps(redact_sensitive_args(args_dict), sort_keys=True) logger.info("Launching Rust frontend: %s", " ".join(cmd[:-1] + [redacted_json])) - self._proc = subprocess.Popen(cmd, pass_fds=pass_fds) + self._listeners = [input_listener, output_listener] + try: + self._proc = subprocess.Popen(cmd, pass_fds=inherited_fds) + except BaseException: + for listener in self._listeners: + listener.cleanup() + raise + else: + input_listener.close() + output_listener.close() # Create a process wrapper with a sentinel fd for monitoring self.processes: list[_SubprocessWrapper] = [ _SubprocessWrapper(self._proc, "RustFrontend") ] - self._finalizer = weakref.finalize(self, _shutdown_subprocesses, self.processes) + self._finalizer = weakref.finalize( + self, + _shutdown_subprocesses_and_cleanup_listeners, + self.processes, + self._listeners, + ) def shutdown(self, timeout: float | None = None) -> None: if self._finalizer.detach() is not None: - _shutdown_subprocesses(self.processes, timeout=timeout) + _shutdown_subprocesses_and_cleanup_listeners( + self.processes, self._listeners, timeout=timeout + ) class _SubprocessWrapper: @@ -533,6 +484,18 @@ def _shutdown_subprocesses( logger.debug_once("[shutdown] Subprocess manager: complete") +def _shutdown_subprocesses_and_cleanup_listeners( + procs: list[_SubprocessWrapper], + listeners: list[ZmqListener], + timeout: float | None = None, +) -> None: + try: + _shutdown_subprocesses(procs, timeout=timeout) + finally: + for listener in listeners: + listener.cleanup() + + def run_api_server_worker_proc( listen_address, sock, args, client_config=None, **uvicorn_kwargs ) -> None: @@ -687,6 +650,18 @@ def shutdown(procs: list[BaseProcess], timeout: float | None = None) -> None: logger.debug_once("[shutdown] Process manager: complete") +def _shutdown_processes_and_cleanup_listeners( + procs: list[BaseProcess], + listeners: list[ZmqListener], + timeout: float | None = None, +) -> None: + try: + shutdown(procs, timeout=timeout) + finally: + for listener in listeners: + listener.cleanup() + + def copy_slice( from_tensor: torch.Tensor, to_tensor: torch.Tensor, length: int ) -> torch.Tensor: From ee406adf0cbcd1e7b073a6b0ca6a3afeec2b9171 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:28:44 +0400 Subject: [PATCH 071/624] [Bugfix] Fix ZeroDivisionError in Qwen3-VL video sampling when the source fps is unknown (#60501) Signed-off-by: Vadim Gimpelson Co-authored-by: Claude --- tests/multimodal/test_video.py | 33 ++++++++++++++++++++++++++++++++- vllm/multimodal/video.py | 10 ++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index d593edba2912..407bffed20a5 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -13,7 +13,7 @@ import numpy.typing as npt import pytest import torch -from transformers import AutoVideoProcessor +from transformers import AutoVideoProcessor, Qwen3VLVideoProcessor from transformers.video_utils import VideoMetadata from vllm.assets.base import get_vllm_public_assets @@ -35,6 +35,7 @@ get_video_loader_backend_for_processor, ) from vllm.multimodal.video_decoders import decode_video, resolve_video_backend_kwargs +from vllm.multimodal.video_decoders.opencv import OpenCVVideoBackendMixin from vllm.multimodal.video_decoders.pynvvideocodec import ( PYNVVIDEOCODEC_DECODER_CACHE_SIZE, PyNvVideoCodecDecoderSlot, @@ -730,6 +731,36 @@ def test_cosmos3_edge_uses_qwen3_vl_video_backend(): assert isinstance(VIDEO_LOADER_REGISTRY.load(backend), Qwen3VLVideoBackend) +def test_qwen3_vl_unknown_source_fps_matches_hf( + dummy_video_path, monkeypatch: pytest.MonkeyPatch +): + """An unknown source fps (reported as 0) samples like HF's fps=None, and + the metadata carries the fps HF assumes, which the frame timestamps use.""" + get_video_metadata = OpenCVVideoBackendMixin.get_video_metadata + + def unknown_fps(cap) -> VideoSourceMetadata: + return get_video_metadata(cap)._replace(original_fps=0.0, duration=0.0) + + monkeypatch.setattr( + OpenCVVideoBackendMixin, "get_video_metadata", staticmethod(unknown_fps) + ) + _, metadata = Qwen3VLVideoBackend.load_bytes( + dummy_video_path.read_bytes(), fps=2, backend="opencv" + ) + + total_frames = metadata["total_num_frames"] + hf_metadata = VideoMetadata(total_num_frames=total_frames, fps=None) + hf_indices = Qwen3VLVideoProcessor().sample_frames(hf_metadata, fps=2).tolist() + assert metadata["frames_indices"] == hf_indices + assert metadata["fps"] == hf_metadata.fps + + source = VideoSourceMetadata(total_frames, original_fps=0.0, duration=0.0) + target = VideoTargetMetadata(num_frames=-1, fps=2, max_duration=300) + assert ( + Qwen3VLVideoBackend.compute_frames_index_to_sample(source, target) == hf_indices + ) + + @pytest.mark.parametrize( "model_repo, expected_loader_cls, hf_sample_kwargs", [ diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index 8f22bffe958c..fc0cdc28d478 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -361,6 +361,15 @@ class Qwen3VLVideoBackend(VideoBackend): _MAX_FRAMES: ClassVar[int] = 768 _MAX_FPS: ClassVar[int] = 30 + @classmethod + def _prepare_source(cls, source: VideoSourceMetadata) -> VideoSourceMetadata: + # vLLM reports an unknown source fps as 0. HF's Qwen3VLVideoProcessor + # assumes a 24 fps source then, and keeps it in the metadata for the + # frame timestamps. + if source.original_fps > 0: + return source + return VideoSourceMetadata(source.total_frames_num, 24.0, source.duration) + @classmethod def compute_frames_index_to_sample( cls, @@ -368,6 +377,7 @@ def compute_frames_index_to_sample( target: VideoTargetMetadata, **kwargs, ) -> list[int]: + source = cls._prepare_source(source) total_frames_num = source.total_frames_num original_fps = source.original_fps fps = min(target.fps, cls._MAX_FPS) From 814301b26800a9d22a7d967a161aa07da4cbb5f2 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Thu, 8 Oct 2026 08:38:41 +0200 Subject: [PATCH 072/624] [Bugfix][Quant] Read safetensors headers of checkpoints with other file names (#60411) Signed-off-by: Stefan Koncarevic Co-authored-by: Andreas Karatzas --- tests/quantization/test_auto_gptq.py | 18 ++++++++- tests/transformers_utils/test_config.py | 29 ++++++++++++++ .../layers/quantization/utils/gptq_utils.py | 5 +++ vllm/transformers_utils/config.py | 40 ++++++++++++++++++- 4 files changed, 90 insertions(+), 2 deletions(-) diff --git a/tests/quantization/test_auto_gptq.py b/tests/quantization/test_auto_gptq.py index 6e7791654c41..7e54599c8098 100644 --- a/tests/quantization/test_auto_gptq.py +++ b/tests/quantization/test_auto_gptq.py @@ -23,7 +23,10 @@ AutoGPTQLinearMethod, AutoGPTQMoEMethod, ) -from vllm.model_executor.layers.quantization.utils.gptq_utils import override_config +from vllm.model_executor.layers.quantization.utils.gptq_utils import ( + is_layer_gptq_quantized, + override_config, +) from vllm.platforms import current_platform PROMPT = "On the surface of Mars, we found" @@ -71,6 +74,19 @@ def test_auto_gptq_config_get_name(): assert AutoGPTQConfig.get_name() == "auto_gptq" +def test_auto_gptq_quantizes_every_layer_when_the_module_list_is_unknown(): + """An empty module list means the safetensors metadata could not be read, + not that the checkpoint holds no quantized layer.""" + assert is_layer_gptq_quantized( + prefix="model.layers.0.mlp.down_proj", + quantized_layers=[], + ) + assert not is_layer_gptq_quantized( + prefix="model.layers.0.mlp.down_proj", + quantized_layers=["self_attn.q_proj"], + ) + + @pytest.mark.parametrize( "dynamic", [ diff --git a/tests/transformers_utils/test_config.py b/tests/transformers_utils/test_config.py index 4fd27ba27879..9b13e725604a 100644 --- a/tests/transformers_utils/test_config.py +++ b/tests/transformers_utils/test_config.py @@ -241,6 +241,7 @@ def test_safetensors_metadata_of_repo_without_safetensors(): ) api = SimpleNamespace( get_safetensors_metadata=get_safetensors_metadata, + list_repo_files=MagicMock(return_value=["pytorch_model.bin"]), snapshot_download=MagicMock(side_effect=LocalEntryNotFoundError("no cache")), ) @@ -250,6 +251,34 @@ def test_safetensors_metadata_of_repo_without_safetensors(): get_safetensors_metadata.assert_called_once() +def test_safetensors_metadata_of_repo_with_a_nonstandard_file_name(): + """`get_safetensors_metadata` only knows `model.safetensors` and its index, + so older checkpoints are read through the file listing instead.""" + from huggingface_hub.errors import NotASafetensorsRepoError + from huggingface_hub.utils import TensorInfo + + weights = "gptq_model-4bit-128g.safetensors" + tensor = TensorInfo(dtype="I32", shape=[5632], data_offsets=(0, 22528)) + parse_safetensors_file_metadata = MagicMock( + return_value=SimpleNamespace(tensors={"layers.0.mlp.down_proj.qweight": tensor}) + ) + api = SimpleNamespace( + get_safetensors_metadata=MagicMock( + side_effect=NotASafetensorsRepoError("not a safetensors repo") + ), + list_repo_files=MagicMock(return_value=["config.json", weights]), + parse_safetensors_file_metadata=parse_safetensors_file_metadata, + ) + + with patch.object(config_module, "hf_api", lambda: api): + metadata = get_safetensors_params_metadata("some/old-gptq-model") + + assert metadata["layers.0.mlp.down_proj.qweight"]["dtype"] == "I32" + parse_safetensors_file_metadata.assert_called_once_with( + "some/old-gptq-model", weights, revision=None + ) + + @pytest.mark.parametrize( ("section_key", "mrope_section", "expected_num_dims"), [ diff --git a/vllm/model_executor/layers/quantization/utils/gptq_utils.py b/vllm/model_executor/layers/quantization/utils/gptq_utils.py index d134abf8dcec..024564e6f34c 100644 --- a/vllm/model_executor/layers/quantization/utils/gptq_utils.py +++ b/vllm/model_executor/layers/quantization/utils/gptq_utils.py @@ -133,6 +133,11 @@ def is_layer_gptq_quantized( quantized_layers = flatten_list(quantized_layers) + # Empty means no metadata could be read at all, not that the checkpoint + # holds no quantized layer, and loading GPTQ weights unquantized cannot work. + if not quantized_layers: + return True + # Fused layers like gate_up_proj or qkv_proj will not be fused # in the safetensors checkpoint. So, we convert the name # from the fused version to unfused + check to make sure that diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 6a339297eec5..c22f6fc3c2ea 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -1355,6 +1355,30 @@ def _read_safetensors_metadata_in_dir(local_dir: Path) -> dict[str, Any]: } +def _read_safetensors_metadata_in_repo( + model: str, + *, + revision: str | None = None, +) -> dict[str, Any]: + api = hf_api() + try: + filenames = [ + filename + for filename in api.list_repo_files(model, revision=revision) + if filename.endswith(".safetensors") + ] + return { + param_name: asdict(info) + for filename in filenames + for param_name, info in api.parse_safetensors_file_metadata( + model, filename, revision=revision + ).tensors.items() + } + except Exception as e: + logger.debug("Could not read safetensors headers for %s: %s", model, e) + return {} + + def get_safetensors_params_metadata( model: str, *, @@ -1372,6 +1396,12 @@ def get_safetensors_params_metadata( for param_name, info in file_mt.tensors.items() } + # `get_safetensors_metadata` only looks for `model.safetensors` and its + # index, so older checkpoints naming their weights differently land here + # even though their headers are perfectly readable. + if metadata := _read_safetensors_metadata_in_repo(model, revision=revision): + return metadata + # Hub fetch failed (e.g. 429, network unreachable). Fall back to the # local HF cache: weights may already be cached from a prior run, and # weight loading itself uses the same cache. @@ -1390,7 +1420,15 @@ def get_safetensors_params_metadata( str(e), ) return {} - return _read_safetensors_metadata_in_dir(Path(local_dir)) + + metadata = _read_safetensors_metadata_in_dir(Path(local_dir)) + if not metadata: + logger.warning_once( + "Could not retrieve safetensors metadata for %s " + "(Hub fetch failed and the cached snapshot holds no weights yet).", + model, + ) + return metadata @cache From 032d39e9b451e858ac9cb18eb0f587ca005361b5 Mon Sep 17 00:00:00 2001 From: xiangdong <40376367+zxd1997066@users.noreply.github.com> Date: Thu, 8 Oct 2026 14:53:37 +0800 Subject: [PATCH 073/624] [XPU][CI] Split Intel jobs to improve B50 utilization (#58817) Signed-off-by: zengxian Signed-off-by: xiangdong <40376367+zxd1997066@users.noreply.github.com> --- .buildkite/intel_jobs/lora_intel.yaml | 37 +++++++-- .buildkite/intel_jobs/misc_intel.yaml | 82 +++++++++++++++++-- .../intel_jobs/model_executor_intel.yaml | 32 +++++++- .../intel_jobs/models_multimodal_intel.yaml | 13 +-- .buildkite/intel_jobs/samplers_intel.yaml | 29 ++++++- .buildkite/intel_jobs/test-intel.yaml | 35 ++++++-- 6 files changed, 201 insertions(+), 27 deletions(-) diff --git a/.buildkite/intel_jobs/lora_intel.yaml b/.buildkite/intel_jobs/lora_intel.yaml index 2cecbf18f773..885e4b546691 100644 --- a/.buildkite/intel_jobs/lora_intel.yaml +++ b/.buildkite/intel_jobs/lora_intel.yaml @@ -9,7 +9,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -33,7 +33,6 @@ steps: pytest -v -s lora/test_peft_helper.py && pytest -v -s lora/test_resolver.py && pytest -v -s lora/test_utils.py && - pytest -v -s lora/test_add_lora.py && pytest -v -s lora/test_worker.py' - label: LoRA Fused/MoE Kernels @@ -120,7 +119,7 @@ steps: agent_tags: label: production gpu: 2+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -138,8 +137,7 @@ steps: pytest -v -s lora/test_quant_model.py --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model0]" --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model1]" --deselect="tests/lora/test_quant_model.py::test_quant_model_tp_equality[model0]" && pytest -v -s lora/test_transformers_model.py && pytest -v -s lora/test_chatglm3_tp.py && - pytest -v -s lora/test_llama_tp.py::test_llama_lora && - pytest -s -v lora/test_minicpmv_tp.py' + pytest -v -s lora/test_llama_tp.py::test_llama_lora' - label: LoRA Multimodal key: lora-multimodal @@ -148,7 +146,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -163,5 +161,30 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && export VLLM_WORKER_MULTIPROC_METHOD=spawn && - pytest -v -s lora/test_default_mm_loras.py && pytest -v -s lora/test_whisper.py' + +- label: LoRA Add + Multimodal + key: lora-add-multimodal + timeout_in_minutes: 45 + device: intel_gpu + agent_tags: + label: production + gpu: 2+ + mem: 24+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/lora + - tests/lora + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + export VLLM_WORKER_MULTIPROC_METHOD=spawn && + pytest -v -s lora/test_add_lora.py && + pytest -v -s lora/test_default_mm_loras.py && + pytest -s -v lora/test_minicpmv_tp.py' diff --git a/.buildkite/intel_jobs/misc_intel.yaml b/.buildkite/intel_jobs/misc_intel.yaml index 5fa4d833d7f0..b7ae9f907846 100644 --- a/.buildkite/intel_jobs/misc_intel.yaml +++ b/.buildkite/intel_jobs/misc_intel.yaml @@ -10,7 +10,7 @@ steps: agent_tags: label: production gpu: 2+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -37,7 +37,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -61,7 +61,40 @@ steps: export VLLM_WORKER_MULTIPROC_METHOD=spawn && cd tests && pytest -v -s v1/executor && - VLLM_BATCH_INVARIANT=1 pytest -v -s -m "not cpu_test" v1/ec_connector/unit' + VLLM_BATCH_INVARIANT=1 pytest -v -s -m "not cpu_test" v1/ec_connector/unit + --ignore=v1/ec_connector/unit/test_ec_cpu_connector.py' + +- label: EC CPU Connector (XPU) + key: ec-cpu-connector-xpu + timeout_in_minutes: 45 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/ + - tests/v1/core + - tests/v1/executor + - tests/v1/kv_offload + - tests/v1/worker + - tests/v1/kv_connector/unit + - tests/v1/ec_connector/unit + - tests/v1/metrics + - tests/entrypoints/openai/correctness/test_lmeval.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'pip install -r requirements/kv_connectors.txt && + export VLLM_WORKER_MULTIPROC_METHOD=spawn && + cd tests && + pytest -v -s v1/ec_connector/unit/test_ec_cpu_connector.py' - label: V1 Sample + Logits key: v1-sample-logits-intel @@ -70,7 +103,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -265,7 +298,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -293,10 +326,47 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - pytest -v -s detokenizer && + pytest -v -s detokenizer --ignore=detokenizer/test_stop_strings.py && pytest -v -s -m "not cpu_test" ./multimodal && pytest -v -s utils_ --ignore=utils_/test_mem_utils.py' +- label: Detokenizer Stop Strings (XPU) + key: detokenizer-stop-strings-xpu + timeout_in_minutes: 45 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/assets/ + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/inputs/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/platforms/ + - vllm/sampling_params.py + - vllm/tokenizers/ + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + - tests/detokenizer + - tests/multimodal + - tests/utils_ + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s detokenizer/test_stop_strings.py' + - label: Fusion Unit Tests key: fusion-unit-tests timeout_in_minutes: 30 diff --git a/.buildkite/intel_jobs/model_executor_intel.yaml b/.buildkite/intel_jobs/model_executor_intel.yaml index 98e4b0ff4511..5da075e235b6 100644 --- a/.buildkite/intel_jobs/model_executor_intel.yaml +++ b/.buildkite/intel_jobs/model_executor_intel.yaml @@ -9,7 +9,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -28,7 +28,7 @@ steps: export VLLM_WORKER_MULTIPROC_METHOD=spawn && export PYTHONFAULTHANDLER=1 && cd tests && - pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload" + pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --ignore="model_executor/model_loader/tensorizer_loader/test_tensorizer.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload" --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormModule::test_matches_reference_above_cudnn_grid_limit" --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype0-out_dtype0]" --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype1-out_dtype1]" @@ -38,3 +38,31 @@ steps: --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype5-out_dtype5]" --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormShapes" --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormInputHandling"' + +- label: Tensorizer Model Loader Tests (XPU) + key: tensorizer-model-loader-tests-xpu + timeout_in_minutes: 45 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/engine/arg_utils.py + - vllm/config/model.py + - vllm/model_executor + - tests/model_executor + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'apt-get update && apt-get install -y curl libsodium23 && + export VLLM_WORKER_MULTIPROC_METHOD=spawn && + export PYTHONFAULTHANDLER=1 && + cd tests && + pytest -v -s model_executor/model_loader/tensorizer_loader/test_tensorizer.py -m "not slow_test"' diff --git a/.buildkite/intel_jobs/models_multimodal_intel.yaml b/.buildkite/intel_jobs/models_multimodal_intel.yaml index eb36e0c02eb8..061be1b136ed 100644 --- a/.buildkite/intel_jobs/models_multimodal_intel.yaml +++ b/.buildkite/intel_jobs/models_multimodal_intel.yaml @@ -9,7 +9,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -33,7 +33,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -72,7 +72,8 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma"' + pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" && + pytest -v -s models/multimodal/generation/test_gemma_ragged_audio.py -m core_model' - label: "Multi-Modal Models (Standard) 4: other + whisper" key: multi-modal-models-standard-4-other-whisper @@ -81,7 +82,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -95,7 +96,7 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing' + pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_gemma_ragged_audio.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing' - label: Multi-Modal Processor # 44min key: multi-modal-processor @@ -131,7 +132,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: diff --git a/.buildkite/intel_jobs/samplers_intel.yaml b/.buildkite/intel_jobs/samplers_intel.yaml index 41ec3c4df77b..027237225d3a 100644 --- a/.buildkite/intel_jobs/samplers_intel.yaml +++ b/.buildkite/intel_jobs/samplers_intel.yaml @@ -6,6 +6,33 @@ steps: key: samplers-test-xpu-kernel-intel timeout_in_minutes: 40 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/model_executor/layers + - vllm/sampling_metadata.py + - tests/samplers + - tests/conftest.py + - vllm/entrypoints/generate/beam_search + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + VLLM_XPU_USE_SAMPLER_KERNEL=1 pytest -v -s samplers + --ignore=samplers/test_beam_search.py' + +- label: Samplers Beam Search Test (XPU Kernel) + key: samplers-beam-search-test-xpu-kernel-intel + timeout_in_minutes: 40 + device: intel_gpu agent_tags: label: production gpu: 1+ @@ -26,4 +53,4 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - VLLM_XPU_USE_SAMPLER_KERNEL=1 pytest -v -s samplers' + VLLM_XPU_USE_SAMPLER_KERNEL=1 pytest -v -s samplers/test_beam_search.py' diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index b1030e372c41..a2a0e14be18c 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -23,7 +23,7 @@ steps: agent_tags: label: production gpu: 2+ - mem: 24+ + mem: 16+ no_plugin: true env: REGISTRY: "public.ecr.aws/q9t5s3a7" @@ -46,7 +46,6 @@ steps: python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 && python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel && python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192 && - VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192 && python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-LLMC -tp 2 --max-model-len 8192 ' - label: "XPU W8A8 FP8 Linear Examples" @@ -58,7 +57,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true env: REGISTRY: "public.ecr.aws/q9t5s3a7" @@ -86,7 +85,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true env: REGISTRY: "public.ecr.aws/q9t5s3a7" @@ -106,8 +105,34 @@ steps: pytest -v -s v1/test_serial_utils.py && pytest -v -s v1/e2e/general/test_correctness_sliding_window.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 --deselect=tests/v1/spec_decode/test_mtp.py::test_glm_mtp_defers_lm_head && - pytest -v -s v1/kv_connector/unit --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 --ignore=v1/kv_connector/unit/test_nixl_connector_hma.py --deselect=tests/v1/kv_connector/unit/test_nixl_connector.py::test_abort_timeout_on_prefiller[ray]' + pytest -v -s v1/kv_connector/unit --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 --ignore=v1/kv_connector/unit/test_nixl_connector_hma.py --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --deselect=tests/v1/kv_connector/unit/test_nixl_connector.py::test_abort_timeout_on_prefiller[ray]' parallelism: 2 + - label: "XPU Qwen3-30B + V1 KV Connector" + key: xpu-qwen3-30b-v1-kv-connector + depends_on: + - image-build-xpu + timeout_in_minutes: 90 + device: intel_gpu + agent_tags: + label: production + gpu: 2+ + mem: 24+ + no_plugin: true + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + source_file_dependencies: + - vllm/ + - .buildkite/intel_jobs/test-intel.yaml + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pip install dpctl && + pytest -v -s v1/kv_connector/unit/test_multi_connector.py && + pytest -v -s v1/kv_connector/unit/test_example_connector.py && + cd .. && + VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192' - label: "XPU server test" key: xpu-server-test depends_on: From e58f8dde7f02631829142f11eaa89099bcd94a98 Mon Sep 17 00:00:00 2001 From: ARAVINDHAN T Date: Thu, 8 Oct 2026 00:04:32 -0700 Subject: [PATCH 074/624] [Bugfix] Use FusedMoEWithLoRA for shared MoE LoRAs on 3D models (#60098) Signed-off-by: ARAVINDHAN T Co-authored-by: linitra24 --- vllm/config/lora.py | 5 +++-- vllm/lora/model_manager.py | 10 ++++------ 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/vllm/config/lora.py b/vllm/config/lora.py index f32b89ae5996..6c994a82f8de 100644 --- a/vllm/config/lora.py +++ b/vllm/config/lora.py @@ -81,8 +81,9 @@ class LoRAConfig: """If True, load MoE expert adapters in the "shared-outer" layout, where the gate/up (`w1`/`w3`) lora_A and the down (`w2`) lora_B are shared across all experts (stored once with expert-dim 1) instead of per-expert. The shared - factors are broadcast to the expert count at kernel time. Only meaningful for - MoE models whose adapters use this layout; ignored otherwise.""" + factors are broadcast to the expert count at kernel time. 3D-weight MoE + models use `FusedMoEWithLoRA` for this layout. Only meaningful for MoE + models whose adapters use this layout; ignored otherwise.""" max_lora_cls_labels: int | None = Field(default=None, ge=1) """Maximum output size for LoRA classification heads. Defaults to the base classification head size.""" diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index 802e5b9233d7..b97a0a412a16 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -144,20 +144,18 @@ def __init__( else set() ) - # When the engine is started with enable_mixed_moe_lora_format=True - # we force the universal 2D wrapper (FusedMoEWithLoRA) regardless of - # the model's 3D flag, so 2D and 3D adapters can coexist. + # Mixed-format and shared-outer adapters both use FusedMoEWithLoRA. + # FusedMoE3DWithLoRA only implements the fused gate/up PEFT pair. self._enable_mixed_moe_lora_format = ( is_moe and lora_config.enable_mixed_moe_lora_format ) + self._enable_moe_shared_loras = is_moe and lora_config.enable_moe_shared_loras self._is_3d_moe_model = ( self._is_moe and self.model.is_3d_moe_weight and not self._enable_mixed_moe_lora_format + and not self._enable_moe_shared_loras ) - # Shared MoE adapters: w13 lora_A / w2 lora_B shared across experts, - # stored as pre-stacked experts.w{1,2,3} tensors (startup opt-in). - self._enable_moe_shared_loras = is_moe and lora_config.enable_moe_shared_loras self.packed_modules_mapping = process_packed_modules_mapping( self.model, force_2d_moe=self._enable_mixed_moe_lora_format, From fd0216d3b598501e28041d84ed0b7dd16453aac6 Mon Sep 17 00:00:00 2001 From: Rehan Khan Date: Thu, 8 Oct 2026 12:45:53 +0530 Subject: [PATCH 075/624] [Bugfix][CPU] Fix fused Gumbel-max wrapping-window bias (#60221) Signed-off-by: Rehan Khan Co-authored-by: Bob --- .buildkite/hardware_tests/cpu.yaml | 2 + .../scripts/hardware_ci/run-cpu-test-arm.sh | 1 + .buildkite/test_areas/kernels.yaml | 1 + csrc/cpu/sampling_kernels.cpp | 151 +++++++++-------- tests/kernels/test_cpu_fused_sampling.py | 153 ++++++++++++++++-- vllm/v1/sample/ops/topk_topp_sampler.py | 8 +- 6 files changed, 236 insertions(+), 80 deletions(-) diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index 3fdce975d3f4..d54688b57491 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -23,6 +23,7 @@ steps: - tests/test_zentorch_version_pin.py - vllm/model_executor/layers/fused_moe/experts/cpu_moe.py - tests/kernels/test_onednn.py + - tests/kernels/test_cpu_fused_sampling.py - tests/kernels/test_awq_int4_to_int8.py - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py - tests/kernels/quantization/test_zen_da8w4.py @@ -49,6 +50,7 @@ steps: tests/test_zentorch_version_pin.py \ tests/kernels/mamba/test_cpu_short_conv.py \ tests/kernels/test_onednn.py \ + tests/kernels/test_cpu_fused_sampling.py \ tests/kernels/test_awq_int4_to_int8.py \ tests/kernels/quantization/test_cpu_fp8_scaled_mm.py \ tests/kernels/quantization/test_zen_da8w4.py \ diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh index 689ca92d6a05..1f598b277bb2 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh @@ -43,6 +43,7 @@ kernel_tests() { docker exec "$CONTAINER_NAME" bash -c " set -e pytest -x -v -s tests/kernels/test_onednn.py + pytest -x -v -s tests/kernels/test_cpu_fused_sampling.py pytest -x -v -s tests/kernels/attention/test_cpu_attn.py pytest -x -v -s tests/kernels/core/test_cpu_activation.py pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index cfc628d382c6..d3c8f657902c 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -256,6 +256,7 @@ steps: --ignore=kernels/test_fused_qk_norm_rope_gate.py --ignore=kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py --ignore=kernels/test_top_k_per_row.py + --ignore=kernels/test_cpu_fused_sampling.py --ignore=kernels/test_bf16x3_router_gemm_cutedsl.py --ignore=kernels/test_bf16x3_router_gemm_rocm.py --ignore=kernels/test_ll_bf16_gemm.py diff --git a/csrc/cpu/sampling_kernels.cpp b/csrc/cpu/sampling_kernels.cpp index 3e2f00a0ba95..1d30b92a8d9a 100644 --- a/csrc/cpu/sampling_kernels.cpp +++ b/csrc/cpu/sampling_kernels.cpp @@ -1,82 +1,73 @@ #include "cpu_types.hpp" -#include +#include #include #include #include -#include -#include namespace { -constexpr int GUMBEL_TABLE_SIZE = 1 << 20; -constexpr int GUMBEL_TABLE_MASK = GUMBEL_TABLE_SIZE - 1; +// Gumbel-max: argmax_i(logit_i + g_i) with g_i ~ Gumbel(0, 1) matches +// softmax. Noise is SplitMix64(seed, i), so tokens are independent and +// consecutive seeds are not a 1-token table shift (vllm#59786). +constexpr uint64_t SPLITMIX_GAMMA = 0x9E3779B97F4A7C15ULL; +constexpr double GUMBEL_MAX = 37.5; // 53-bit u in (0,1) => -log(-log u) < 37.5 -static float g_gumbel_table[GUMBEL_TABLE_SIZE]; -static std::once_flag g_gumbel_init_flag; - -static void init_gumbel_table() { - std::mt19937 rng(42); - std::uniform_real_distribution uniform(1.0e-7f, 1.0f - 1.0e-7f); - for (int i = 0; i < GUMBEL_TABLE_SIZE; ++i) { - float u = uniform(rng); - g_gumbel_table[i] = -std::log(-std::log(u)); - } +static inline uint64_t splitmix64(uint64_t z) { + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL; + return z ^ (z >> 31); } -static inline void ensure_gumbel_table() { - std::call_once(g_gumbel_init_flag, init_gumbel_table); +static void fused_gumbel_argmax_row(int64_t* __restrict__ output, + const float* __restrict__ row, + uint64_t seed, int64_t vocab_size) { + const uint64_t key = splitmix64(seed); + + // Compare in (logit, noise) space: x + g > best_x + best_g is + // (x - best_x) + g > best_g. Adding g into a 1e38 logit loses it in + // double ULP, so ties at the float max would always keep the lower index. + double best_logit = 0.0; + double best_g = 0.0; + int64_t best_idx = 0; + bool have_best = false; + for (int64_t i = 0; i < vocab_size; ++i) { + const double x = row[i]; + if (x == -std::numeric_limits::infinity()) continue; + if (have_best && (x - best_logit) + GUMBEL_MAX <= best_g) continue; + const uint64_t bits = + splitmix64(key + uint64_t(i + 1) * SPLITMIX_GAMMA) >> 11; + const double one_minus_u = + (double((1ULL << 53) - 1 - bits) + 0.5) * 0x1.0p-53; + if (have_best && one_minus_u >= std::exp((x - best_logit) - best_g)) + continue; + const double u = (double(bits) + 0.5) * 0x1.0p-53; + const double g = -std::log(-std::log(u)); + if (!have_best || (x - best_logit) + g > best_g) { + have_best = true; + best_logit = x; + best_g = g; + best_idx = i; + } + } + *output = best_idx; } -// Fused Gumbel-max kernel - static void fused_gumbel_argmax_kernel(int64_t* __restrict__ output, const float* __restrict__ logits, const int64_t* __restrict__ seeds, const int64_t batch_size, const int64_t vocab_size) { - ensure_gumbel_table(); - constexpr int VEC_ELEM_NUM = vec_op::FP32Vec8::VEC_ELEM_NUM; - const int64_t vec_end = vocab_size - (vocab_size % VEC_ELEM_NUM); - + if (batch_size <= 0) return; + if (batch_size == 1) { + fused_gumbel_argmax_row(output, logits, uint64_t(seeds[0]), vocab_size); + return; + } #pragma omp parallel for schedule(static) for (int64_t b = 0; b < batch_size; ++b) { - const float* row = logits + b * vocab_size; - const uint32_t seed = static_cast(seeds[b]); - - float best_score = -std::numeric_limits::infinity(); - int64_t best_idx = 0; - - for (int64_t i = 0; i < vec_end; i += VEC_ELEM_NUM) { - vec_op::FP32Vec8 logit_vec(row + i); - - float gumbel_buf[VEC_ELEM_NUM]; - for (int k = 0; k < VEC_ELEM_NUM; ++k) - gumbel_buf[k] = g_gumbel_table[(seed + i + k) & GUMBEL_TABLE_MASK]; - vec_op::FP32Vec8 gumbel_vec(gumbel_buf); - - vec_op::FP32Vec8 score = logit_vec + gumbel_vec; - - float score_buf[VEC_ELEM_NUM]; - score.save(score_buf); - for (int k = 0; k < VEC_ELEM_NUM; ++k) { - if (score_buf[k] > best_score) { - best_score = score_buf[k]; - best_idx = i + k; - } - } - } - - for (int64_t i = vec_end; i < vocab_size; ++i) { - float s = row[i] + g_gumbel_table[(seed + i) & GUMBEL_TABLE_MASK]; - if (s > best_score) { - best_score = s; - best_idx = i; - } - } - - output[b] = best_idx; + fused_gumbel_argmax_row(output + b, logits + b * vocab_size, + uint64_t(seeds[b]), vocab_size); } } @@ -91,22 +82,41 @@ static void greedy_argmax_kernel(int64_t* __restrict__ output, for (int64_t b = 0; b < batch_size; ++b) { const float* row = logits + b * vocab_size; + // Vector max has no consistent NaN rule, so NaNs are caught via the sum: + // any NaN (or inf + -inf, a harmless false positive) makes it NaN. vec_op::FP32Vec16 vmax(-std::numeric_limits::infinity()); + vec_op::FP32Vec16 vsum(0.0f); for (int64_t i = 0; i < vec_end; i += VEC_ELEM_NUM) { - vmax = vmax.max(vec_op::FP32Vec16(row + i)); + vec_op::FP32Vec16 v(row + i); + vmax = vmax.max(v); + vsum = vsum + v; } float best_val = vmax.reduce_max(); + bool maybe_nan = std::isnan(vsum.reduce_sum()); for (int64_t i = vec_end; i < vocab_size; ++i) { + maybe_nan |= std::isnan(row[i]); if (row[i] > best_val) { best_val = row[i]; } } - int64_t best_idx = 0; - for (int64_t i = 0; i < vocab_size; ++i) { - if (row[i] == best_val) { - best_idx = i; - break; + int64_t best_idx = -1; + if (maybe_nan) { + // torch.argmax treats NaN as the maximum: return the first NaN. + for (int64_t i = 0; i < vocab_size; ++i) { + if (std::isnan(row[i])) { + best_idx = i; + break; + } + } + } + if (best_idx < 0) { + best_idx = 0; + for (int64_t i = 0; i < vocab_size; ++i) { + if (row[i] == best_val) { + best_idx = i; + break; + } } } output[b] = best_idx; @@ -117,21 +127,28 @@ static void greedy_argmax_kernel(int64_t* __restrict__ output, torch::Tensor fused_gumbel_argmax(const torch::Tensor& logits, const torch::Tensor& seeds) { + TORCH_CHECK(logits.device().is_cpu(), "logits must be a CPU tensor"); TORCH_CHECK(logits.dim() == 2, "logits must be 2-D [batch, vocab]"); TORCH_CHECK(logits.scalar_type() == torch::kFloat32, "logits must be float32"); + TORCH_CHECK(logits.size(1) > 0, "vocab_size must be positive"); + TORCH_CHECK(seeds.device().is_cpu(), "seeds must be a CPU tensor"); + TORCH_CHECK(seeds.scalar_type() == torch::kInt64, "seeds must be int64"); TORCH_CHECK(seeds.dim() == 1 && seeds.size(0) == logits.size(0), "seeds must be 1-D with batch_size elements"); auto logits_contig = logits.contiguous(); + auto seeds_contig = seeds.contiguous(); auto output = torch::empty({logits_contig.size(0)}, torch::kInt64); - fused_gumbel_argmax_kernel( - output.data_ptr(), logits_contig.data_ptr(), - seeds.data_ptr(), logits_contig.size(0), logits_contig.size(1)); + fused_gumbel_argmax_kernel(output.data_ptr(), + logits_contig.data_ptr(), + seeds_contig.data_ptr(), + logits_contig.size(0), logits_contig.size(1)); return output; } torch::Tensor greedy_argmax(const torch::Tensor& logits) { + TORCH_CHECK(logits.device().is_cpu(), "logits must be a CPU tensor"); TORCH_CHECK(logits.dim() == 2, "logits must be 2-D [batch, vocab]"); TORCH_CHECK(logits.scalar_type() == torch::kFloat32, "logits must be float32"); diff --git a/tests/kernels/test_cpu_fused_sampling.py b/tests/kernels/test_cpu_fused_sampling.py index 8a3857fe9ee7..cdb10ab7e32e 100644 --- a/tests/kernels/test_cpu_fused_sampling.py +++ b/tests/kernels/test_cpu_fused_sampling.py @@ -43,20 +43,20 @@ def test_negative_logits(self): class TestFusedGumbelArgmax: - @pytest.mark.parametrize("vocab_size", VOCAB_SIZES) - def test_distribution_chi_squared(self, vocab_size: int): + def test_distribution_chi_squared(self): """Verify sampling distribution via chi-squared goodness of fit.""" - small_vocab = min(vocab_size, 100) + # Small support is enough for chi-squared power; large-vocab bias + # is covered by test_noise_not_sliding_window (#59786). + small_vocab = 100 logits = torch.randn(1, small_vocab, dtype=torch.float32) probs = logits.softmax(dim=-1).squeeze(0) n_samples = 100_000 - counts = torch.zeros(small_vocab) - for trial in range(n_samples): - seed = torch.tensor([trial * 7 + 13], dtype=torch.long) - tile = logits.expand(1, -1).contiguous() - idx = torch.ops._C.fused_gumbel_argmax(tile, seed) - counts[idx.item()] += 1 + seeds = torch.arange(n_samples, dtype=torch.long) * 7 + 13 + idx = torch.ops._C.fused_gumbel_argmax( + logits.expand(n_samples, -1).contiguous(), seeds + ) + counts = torch.bincount(idx, minlength=small_vocab).float() expected = probs * n_samples mask = expected > 5 @@ -70,6 +70,39 @@ def test_distribution_chi_squared(self, vocab_size: int): f"dof={dof}, p={p_value:.6f}" ) + def test_noise_not_sliding_window(self): + """Table[(seed+i)&MASK] made consecutive seeds a 1-token shift. + + On uniform logits that yields P(w(s+1) == w(s)-1) ~ 1-2/V. Independent + per-token noise yields ~1/V. See #59786. + """ + vocab, n = 128, 2048 + logits = torch.zeros(n, vocab, dtype=torch.float32) + seeds = torch.arange(n, dtype=torch.long) + winners = torch.ops._C.fused_gumbel_argmax(logits, seeds) + shifted = (winners[:-1] - 1) % vocab + frac = (winners[1:] == shifted).float().mean().item() + assert frac < 0.05, ( + f"consecutive-seed winners follow a table shift (frac={frac:.3f})" + ) + + def test_zipf_head_matches_softmax(self): + """Zipf head probability should match softmax (#59786).""" + vocab, n_samples, batch = 8192, 16384, 256 + logits = -1.1 * torch.arange(1, vocab + 1, dtype=torch.float32).log() + p0 = torch.softmax(logits.double(), dim=-1)[0].item() + row = logits.view(1, -1).expand(batch, vocab).contiguous() + n_top = 0 + for start in range(0, n_samples, batch): + seeds = torch.arange(start, start + batch, dtype=torch.long) + tokens = torch.ops._C.fused_gumbel_argmax(row, seeds) + n_top += int((tokens == 0).sum()) + q0 = n_top / n_samples + se = (p0 * (1.0 - p0) / n_samples) ** 0.5 + assert abs(q0 - p0) < 8.0 * se, ( + f"P(top-1) kernel={q0:.4f} exact={p0:.4f} se={se:.4f}" + ) + def test_deterministic_same_seed(self): """Same seed produces same result.""" logits = torch.randn(4, 32000, dtype=torch.float32) @@ -95,6 +128,108 @@ def test_output_in_range(self, vocab_size: int, batch_size: int): assert result.min() >= 0 assert result.max() < vocab_size + def test_samples_only_from_finite_logits(self): + """top-k/top-p mask with -inf; those IDs must never win.""" + allowed = torch.tensor([3, 17, 99, 1024, 8191], dtype=torch.long) + logits = torch.full((64, 32000), float("-inf"), dtype=torch.float32) + logits[:, allowed] = torch.randn(64, allowed.numel()) + tokens = torch.ops._C.fused_gumbel_argmax( + logits, torch.arange(64, dtype=torch.long) + ) + assert torch.isin(tokens, allowed).all() + + def test_single_unmasked_token(self): + logits = torch.full((8, 4096), float("-inf"), dtype=torch.float32) + logits[:, 123] = 0.0 + seeds = torch.tensor([0, -1, -(2**63), 2**63 - 1, 42, 99, 1000, 7]) + tokens = torch.ops._C.fused_gumbel_argmax(logits, seeds) + assert (tokens == 123).all() + + def test_tokens_beyond_gumbel_max_cannot_win(self): + """A gap larger than GUMBEL_MAX is unreachable.""" + logits = torch.zeros(256, 4, dtype=torch.float32) + logits[:, 1:] = -80.0 + tokens = torch.ops._C.fused_gumbel_argmax( + logits, torch.arange(256, dtype=torch.long) + ) + assert (tokens == 0).all() + + def test_tied_huge_logits_not_index_biased(self): + """Noise must still break ties when logits sit at float32 max.""" + logits = torch.full( + (4096, 2), torch.finfo(torch.float32).max, dtype=torch.float32 + ) + tokens = torch.ops._C.fused_gumbel_argmax( + logits, torch.arange(4096, dtype=torch.long) + ) + frac0 = (tokens == 0).float().mean().item() + assert set(tokens.tolist()) == {0, 1}, tokens.unique().tolist() + assert 0.35 < frac0 < 0.65, f"P(index 0)={frac0:.3f} expected ~0.5" + + def test_serial_batch1_matches_batched(self): + """OpenMP rows must match the batch=1 serial path.""" + logits = torch.randn(8, 2048, dtype=torch.float32) + seeds = torch.arange(8, dtype=torch.long) + batched = torch.ops._C.fused_gumbel_argmax(logits, seeds) + serial = torch.stack( + [ + torch.ops._C.fused_gumbel_argmax(logits[i : i + 1], seeds[i : i + 1])[0] + for i in range(8) + ] + ) + torch.testing.assert_close(batched, serial) + + def test_noncontiguous_matches_contiguous(self): + logits = torch.randn(8, 1024, dtype=torch.float32)[:, ::2] + seeds = torch.arange(16, dtype=torch.long)[::2] + assert not logits.is_contiguous() and not seeds.is_contiguous() + expected = torch.ops._C.fused_gumbel_argmax( + logits.contiguous(), seeds.contiguous() + ) + actual = torch.ops._C.fused_gumbel_argmax(logits, seeds) + torch.testing.assert_close(actual, expected) + + def test_empty_batch(self): + tokens = torch.ops._C.fused_gumbel_argmax( + torch.empty(0, 7), torch.empty(0, dtype=torch.long) + ) + assert tokens.shape == (0,) and tokens.dtype == torch.int64 + + def test_masked_support_matches_softmax(self): + """Finite subset after -inf mask should match softmax.""" + from scipy.stats import chisquare + + logits = torch.tensor( + [0.0, -0.2, -0.7, -1.5, -3.0, float("-inf")], dtype=torch.float32 + ) + n_samples = 20_000 + tokens = torch.ops._C.fused_gumbel_argmax( + logits.expand(n_samples, -1).contiguous(), + torch.arange(n_samples, dtype=torch.long), + ) + counts = torch.bincount(tokens, minlength=logits.numel()).double() + assert counts[-1] == 0 + probs = logits.double().softmax(-1) + _, p_value = chisquare(counts[:-1].numpy(), (probs[:-1] * n_samples).numpy()) + assert p_value > 1e-7, f"masked-support p-value: {p_value}" + + def test_zipf_tail_matches_softmax(self): + """Independent 53-bit Gumbel must not starve the Zipf tail (#59786).""" + vocab, n_samples, batch = 8192, 16384, 256 + logits = -1.1 * torch.arange(1, vocab + 1, dtype=torch.float32).log() + p_tail = torch.softmax(logits.double(), dim=-1)[vocab // 2 :].sum().item() + row = logits.view(1, -1).expand(batch, vocab).contiguous() + n_tail = 0 + for start in range(0, n_samples, batch): + seeds = torch.arange(start, start + batch, dtype=torch.long) + tokens = torch.ops._C.fused_gumbel_argmax(row, seeds) + n_tail += int((tokens >= vocab // 2).sum()) + q_tail = n_tail / n_samples + se = (p_tail * (1.0 - p_tail) / n_samples) ** 0.5 + assert abs(q_tail - p_tail) < 8.0 * se, ( + f"P(tail) kernel={q_tail:.4f} exact={p_tail:.4f} se={se:.4f}" + ) + class TestMemory: def test_fused_no_intermediate_allocs(self): diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index a3084a09c28e..daa94ebc3af2 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -319,11 +319,11 @@ def forward_cpu( k: torch.Tensor | None, p: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Fused Gumbel-max sampling for CPU. + """CPU Gumbel-max sampling with per-token SplitMix64 noise. - Uses a precomputed Gumbel table + single-pass argmax over logits, - skipping softmax and intermediate allocations entirely. - Falls back to the native path when fp64 Gumbel noise is requested. + Independent (seed, token) streams avoid the wrapping-window bias + (#59786). Falls back to native sampling when fp64 Gumbel noise + is requested. """ logits = apply_top_k_top_p(logits, k, p) logits_to_return = None From 29ae81fe83712bb74c00a4ce0515d031f1225c9b Mon Sep 17 00:00:00 2001 From: Alec <35311602+alec-flowers@users.noreply.github.com> Date: Thu, 8 Oct 2026 00:27:47 -0700 Subject: [PATCH 076/624] [Rust Frontend] Preserve explicit zero sampling values in gRPC requests (#60115) Signed-off-by: Alec Flowers --- rust/proto/inference.proto | 6 +-- rust/src/server/src/grpc/convert.rs | 63 +++++++++++++++++++++++------ rust/src/server/src/grpc/tests.rs | 34 ++++++++-------- 3 files changed, 71 insertions(+), 32 deletions(-) diff --git a/rust/proto/inference.proto b/rust/proto/inference.proto index 8189758212d4..79c1d0a147e2 100644 --- a/rust/proto/inference.proto +++ b/rust/proto/inference.proto @@ -75,9 +75,9 @@ message KVHintAction { message RandomSampling { uint32 num_sequences = 1; // "n", default (0) means 1 - uint32 top_k = 2; // 0 means default - float top_p = 3; // 0 means default - float min_p = 4; // 0 means default + optional uint32 top_k = 2; // Omitted means default; 0 disables top-k. + optional float top_p = 3; // Omitted means default; must be in (0, 1]. + optional float min_p = 4; // Omitted means default; 0 disables min-p. optional int64 seed = 5; } diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 8f2c589f154b..a409cc7f79ee 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -161,10 +161,7 @@ fn build_sampling_params( ..SamplingParams::default() }; - // RandomSampling: for every remaining sampling field the protobuf default (`0`) - // is treated as "unset" and leaves the resolved value to the lowering - // stage, which falls back to the model-provided default or a - // neutral/disabled value otherwise. + // Preserve explicit values; omitted fields inherit model defaults during lowering. if let Some(s) = sampling { // num_sequences (n > 1) is not supported yet by the TextLlm layer; the response // path also hardcodes SequenceOutput.index = 0, so accepting >1 would silently @@ -174,15 +171,9 @@ fn build_sampling_params( "num_sequences > 1 is not supported", )); } - if s.top_k != 0 { - params.top_k = Some(s.top_k); - } - if s.top_p != 0.0 { - params.top_p = Some(s.top_p); - } - if s.min_p != 0.0 { - params.min_p = Some(s.min_p); - } + params.top_k = s.top_k; + params.top_p = s.top_p; + params.min_p = s.min_p; params.seed = s.seed; } @@ -790,6 +781,52 @@ mod tests { } } + #[test] + fn grpc_sampling_preserves_explicit_values_and_model_defaults_through_lowering() { + for ((top_k, top_p, min_p), expected) in [ + ((None, None, None), (8, 0.9, 0.2)), + ((Some(0), Some(1.0), Some(0.0)), (0, 1.0, 0.0)), + ((Some(50), Some(0.8), Some(0.1)), (50, 0.8, 0.1)), + ] { + let request = pb::GenerateRequest { + temperature: Some(0.7), + sampling: Some(pb::RandomSampling { + top_k, + top_p, + min_p, + ..Default::default() + }), + ..base_request() + }; + let encoded = request.encode_to_vec(); + for stream in [false, true] { + let decoded = pb::GenerateRequest::decode(encoded.as_slice()).unwrap(); + let text = to_text_request(decoded, stream, &["test-model".to_string()]).unwrap(); + let engine = vllm_text::lower_text_request( + text, + vec![1], + SamplingHints { + default_top_k: Some(8), + default_top_p: Some(0.9), + default_min_p: Some(0.2), + ..Default::default() + }, + SamplingLimits { + max_model_len: 256, + max_logprobs: 20, + model_vocab_size: 512, + tokenizer_vocab_size: 512, + }, + &TestTokenizer::new(), + ) + .unwrap() + .generate_request; + let params = engine.sampling_params; + assert_eq!((params.top_k, params.top_p, params.min_p), expected); + } + } + } + #[test] fn bypass_prefix_cache_maps_to_skip_reading_prefix_cache() { let req = pb::GenerateRequest { diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 631eb76b36ab..0ddbfb096e43 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -1366,22 +1366,24 @@ async fn unary_generate_invalid_sampling_params_returns_invalid_argument() { ) .await; - let status = client - .generate(pb::GenerateRequest { - request_id: "test-invalid-sampling".to_string(), - model: "test-model".to_string(), - prompt: Some(pb::generate_request::Prompt::Text("hi".to_string())), - sampling: Some(pb::RandomSampling { - top_p: 2.0, + for top_p in [0.0, 2.0] { + let status = client + .generate(pb::GenerateRequest { + request_id: "test-invalid-sampling".to_string(), + model: "test-model".to_string(), + prompt: Some(pb::generate_request::Prompt::Text("hi".to_string())), + sampling: Some(pb::RandomSampling { + top_p: Some(top_p), + ..Default::default() + }), ..Default::default() - }), - ..Default::default() - }) - .await - .expect_err("should fail when top_p is out of range"); + }) + .await + .expect_err("should fail when top_p is out of range"); - assert_eq!(status.code(), tonic::Code::InvalidArgument); - assert!(status.message().contains("top_p")); + assert_eq!(status.code(), tonic::Code::InvalidArgument); + assert!(status.message().contains("top_p")); + } server_task.abort(); } @@ -1519,8 +1521,8 @@ async fn unary_generate_with_sampling_params() { prompt: Some(pb::generate_request::Prompt::Text("test".to_string())), temperature: Some(0.7), sampling: Some(pb::RandomSampling { - top_k: 50, - top_p: 0.9, + top_k: Some(50), + top_p: Some(0.9), seed: Some(42), ..Default::default() }), From 0f112d180033f3757cf283723f5593a40b02ed82 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Thu, 8 Oct 2026 09:33:28 +0200 Subject: [PATCH 077/624] [Bugfix][Pooling] Preserve cache_salt for chunked embeddings (#47696) Signed-off-by: Ting Sun --- vllm/entrypoints/pooling/embed/io_processor.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index 3c5b75cee62c..1b23a77ec60d 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -142,13 +142,16 @@ def maybe_pre_process_chunked(self, ctx: PoolingServeContext) -> None: ) prompt_token_ids = cast(list[int], token_ids) + cache_salt = cast(str | None, engine_input["prompts"].get("cache_salt")) for chunk_idx, chunk_tokens in enumerate( chunk_list(prompt_token_ids, max_model_len) ): chunked_engine_inputs.append( PoolingEngineInput( - prompts=tokens_input(prompt_token_ids=chunk_tokens), + prompts=tokens_input( + prompt_token_ids=chunk_tokens, cache_salt=cache_salt + ), params=engine_input["params"], lora_requests=engine_input["lora_requests"], priorities=engine_input["priorities"], From 4cf1307e72cb1648e4e3c3ab4e0bf85d5616fcb3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Ma=C5=84czak?= Date: Thu, 8 Oct 2026 10:36:05 +0200 Subject: [PATCH 078/624] [XPU] Add tuned Mamba SSU configs for 2 more B70 shapes (#57565) Signed-off-by: pmanczak --- ..._Pro_B70_Graphics,cache_dtype=float16.json | 47 +++++++++++++++++++ ..._Pro_B70_Graphics,cache_dtype=float32.json | 39 +++++++++++++++ ..._Pro_B70_Graphics,cache_dtype=float16.json | 47 +++++++++++++++++++ ..._Pro_B70_Graphics,cache_dtype=float32.json | 47 +++++++++++++++++++ 4 files changed, 180 insertions(+) create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json new file mode 100644 index 000000000000..d51b085769b1 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json @@ -0,0 +1,47 @@ +{ + "triton_version": "3.8.0", + "48": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "96": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "192": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "384": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "768": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "1536": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "3072": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "6144": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "12288": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "24576": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "49152": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json new file mode 100644 index 000000000000..666506768c01 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json @@ -0,0 +1,39 @@ +{ + "triton_version": "3.8.0", + "48": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "384": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "768": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "1536": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "3072": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "6144": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "12288": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "24576": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "49152": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json new file mode 100644 index 000000000000..aed0a3c63515 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json @@ -0,0 +1,47 @@ +{ + "triton_version": "3.8.0", + "80": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "640": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "1280": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "2560": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "5120": { + "BLOCK_SIZE_M": 32, + "num_warps": 8 + }, + "10240": { + "BLOCK_SIZE_M": 32, + "num_warps": 8 + }, + "20480": { + "BLOCK_SIZE_M": 32, + "num_warps": 8 + }, + "40960": { + "BLOCK_SIZE_M": 32, + "num_warps": 8 + }, + "81920": { + "BLOCK_SIZE_M": 32, + "num_warps": 8 + }, + "122880": { + "BLOCK_SIZE_M": 32, + "num_warps": 8 + }, + "163840": { + "BLOCK_SIZE_M": 32, + "num_warps": 8 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json new file mode 100644 index 000000000000..7c642c06bc2e --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json @@ -0,0 +1,47 @@ +{ + "triton_version": "3.8.0", + "80": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "640": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "1280": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "2560": { + "BLOCK_SIZE_M": 8, + "num_warps": 4 + }, + "5120": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "10240": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "20480": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "40960": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + }, + "81920": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + }, + "122880": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + }, + "163840": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + } +} \ No newline at end of file From 7d47ac2ddba430510f1503dae424bb298fe5a300 Mon Sep 17 00:00:00 2001 From: Thang Nguyen <69278249+Thangnguyenvn98@users.noreply.github.com> Date: Thu, 8 Oct 2026 04:52:59 -0400 Subject: [PATCH 079/624] [CI] Enroll five steps in automatic sharding (#60492) Signed-off-by: Thang Nguyen Co-authored-by: Claude Opus 5.5 Co-authored-by: khluu --- .buildkite/test_areas/entrypoints.yaml | 1 + .buildkite/test_areas/kernels.yaml | 1 + .buildkite/test_areas/quantization.yaml | 1 + .buildkite/test_areas/spec_decode.yaml | 1 + 4 files changed, 4 insertions(+) diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 68c6243d479b..34a190e810b9 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -231,6 +231,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (Pooling)" device: h200_35gb key: entrypoints-integration-pooling + automatic_shard: true timeout_in_minutes: 75 working_dir: "/vllm-workspace/tests" source_file_dependencies: diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index d3c8f657902c..e0d6c4ce40a6 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -504,6 +504,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Mamba Kernels" device: h200_35gb key: kernels-mamba-test + automatic_shard: true timeout_in_minutes: 60 source_file_dependencies: - csrc/mamba/ diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index b49b4b92c736..429b01f27320 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -127,6 +127,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Quantized Models" device: h200_35gb key: quantized-models-test + automatic_shard: true timeout_in_minutes: 65 env: VLLM_USE_V2_MODEL_RUNNER: "0" diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index 6882c24b3745..d031d4fdea9b 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -5,6 +5,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) V1 Spec Decode" device: h200_35gb key: v1-spec-decode + automatic_shard: true timeout_in_minutes: 40 source_file_dependencies: - vllm/config/ From 70bdce6133d0094f420a3714bb9ee186238c82f6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Thu, 8 Oct 2026 11:05:52 +0200 Subject: [PATCH 080/624] [Model][GLM-5.3-Flash] Warm up the DCP top-k merge kernels for the kpool indexer (#60032) Signed-off-by: NickLucche --- vllm/models/glm5next/nvidia/sparse_indexer.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 806b633c5836..77fa4c409152 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -26,6 +26,7 @@ from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops from vllm.platforms import current_platform from vllm.utils.deep_gemm import has_deep_gemm +from vllm.utils.import_utils import has_cutedsl from vllm.utils.torch_utils import ( LayerNameType, _resolve_layer_name, @@ -717,6 +718,14 @@ def __init__( ) self.dcp_world_size = _parallel.decode_context_parallel_size self.dcp_rank = get_dcp_group().rank_in_group + if current_platform.is_cuda() and has_cutedsl(): + from vllm.model_executor.kernels.attention.dsa.dcp_indexer_cutedsl import ( # noqa: E501 + _PACK_DCP_TOPK_CANDIDATES_KERNEL, + _STABLE_TOPK_FROM_GATHERED_CANDIDATES_KERNEL, + ) + + _PACK_DCP_TOPK_CANDIDATES_KERNEL.register_warmup() + _STABLE_TOPK_FROM_GATHERED_CANDIDATES_KERNEL.register_warmup() @property def cp_kv_cache_interleave_size(self) -> int: From 0004536dd00f8237fd6dd68171fbda9b680c07b2 Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Thu, 8 Oct 2026 17:13:48 +0800 Subject: [PATCH 081/624] [XPU][Graph] Replace boolean-mask assignment with `torch.where` (#60529) Signed-off-by: zhenwei-intel Co-authored-by: Kunshang Ji --- vllm/model_executor/models/mimo_mtp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/models/mimo_mtp.py b/vllm/model_executor/models/mimo_mtp.py index 65566230a562..5e331fa8f002 100644 --- a/vllm/model_executor/models/mimo_mtp.py +++ b/vllm/model_executor/models/mimo_mtp.py @@ -73,7 +73,7 @@ def forward( ) -> torch.Tensor: assert inputs_embeds is not None # masking inputs at position 0, as not needed by MTP - inputs_embeds[positions == 0] = 0 + inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) inputs_embeds = self.token_layernorm(inputs_embeds) previous_hidden_states = self.hidden_layernorm(previous_hidden_states) From 21d97584593de08f7a1a44dad8aca32aa4a01c51 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:30:52 +0100 Subject: [PATCH 082/624] Remove Python 3.10 support now that it is EOL (#60402) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- .pre-commit-config.yaml | 28 +++------- AGENTS.md | 3 -- CMakeLists.txt | 48 ++++++++--------- docker/Dockerfile.cpu | 2 +- docs/contributing/README.md | 10 +--- docs/deployment/frameworks/skypilot.md | 8 +-- docs/getting_started/installation/cpu.md | 2 +- docs/getting_started/installation/gpu.md | 2 +- docs/getting_started/quickstart.md | 2 +- .../structured_outputs/pyproject.toml | 2 +- pyproject.toml | 5 +- requirements/cpu.txt | 1 - setup.py | 14 ++--- .../launchers/test_dp_supervisor.py | 2 +- tests/entrypoints/launchers/test_shutdown.py | 6 +-- .../tensorizer_loader/test_tensorizer.py | 2 +- tests/multimodal/media/test_connector.py | 7 ++- tests/v1/distributed/test_async_llm_dp.py | 2 +- tools/ci/kernel_symbol_map.py | 2 +- tools/pre_commit/check_forbidden_imports.py | 6 ++- tools/pre_commit/mypy.py | 4 +- vllm/benchmarks/serve.py | 4 +- vllm/benchmarks/sweep/plot.py | 4 +- vllm/benchmarks/sweep/serve_workload.py | 3 +- vllm/benchmarks/sweep/server.py | 2 +- vllm/config/engram.py | 3 +- vllm/config/lora.py | 3 +- vllm/config/parallel.py | 3 +- vllm/config/profiler.py | 3 +- vllm/config/scheduler.py | 3 +- vllm/config/speculative.py | 3 +- vllm/config/structured_outputs.py | 3 +- vllm/config/watermarking.py | 3 +- vllm/distributed/utils.py | 12 ++--- vllm/distributed/weight_transfer/base.py | 2 +- .../distributed/weight_transfer/ipc_engine.py | 3 +- .../weight_transfer/m2n_trainer.py | 3 +- .../weight_transfer/nccl_engine.py | 3 +- .../weight_transfer/sharded_rdt_trainer.py | 3 +- .../weight_transfer/sparse_nccl_engine.py | 3 +- vllm/entrypoints/chat_utils.py | 3 +- vllm/entrypoints/launchers/dp_supervisor.py | 2 +- vllm/entrypoints/llm.py | 3 +- vllm/entrypoints/pooling/embed/serving.py | 3 +- vllm/entrypoints/pooling/pooling/serving.py | 3 +- vllm/entrypoints/pooling/scoring/typing.py | 4 +- vllm/entrypoints/serve/tokenize/api_router.py | 2 +- vllm/inputs/engine.py | 4 +- vllm/inputs/llm.py | 4 +- vllm/model_executor/kernels/linear/base.py | 3 +- .../scaled_mm/BlockScaledMMLinearKernel.py | 3 +- vllm/model_executor/model_loader/utils.py | 3 +- vllm/model_executor/models/interfaces.py | 3 +- .../deepseek_v4/common/mm_preprocess.py | 3 +- vllm/multimodal/parse.py | 2 +- vllm/multimodal/processing/processor.py | 3 +- vllm/renderers/hf.py | 2 +- vllm/snapshot/manifest.py | 3 +- vllm/snapshot/runtime.py | 4 +- vllm/tokenizers/registry.py | 4 +- vllm/usage/usage_lib.py | 2 +- vllm/utils/collection_utils.py | 4 +- vllm/utils/import_utils.py | 3 +- vllm/v1/kv_cache_interface.py | 3 +- vllm/v1/metrics/perf.py | 3 +- vllm/v1/structured_output/backend_outlines.py | 53 +++++++------------ vllm/v1/worker/workspace.py | 7 +-- 67 files changed, 141 insertions(+), 219 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3e418e005ec8..92963cf69a22 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -199,29 +199,13 @@ repos: additional_dependencies: [PyYAML, regex] verbose: true # full-scan `note:` lines are advisory, so the hook passes require_serial: true # one full scan, not one per pre-commit file batch - - id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward - name: Run mypy for Python 3.10 - entry: python tools/pre_commit/mypy.py "3.10" - <<: &mypy_common - language: python - types_or: [python, pyi] - require_serial: true - additional_dependencies: ["mypy==1.20.2", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic] - - id: mypy-3.11 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward - name: Run mypy for Python 3.11 + - id: mypy # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward + name: Run mypy for the oldest supported Python version entry: python tools/pre_commit/mypy.py "3.11" - <<: *mypy_common - stages: [manual] # Only run in CI - - id: mypy-3.12 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward - name: Run mypy for Python 3.12 - entry: python tools/pre_commit/mypy.py "3.12" - <<: *mypy_common - stages: [manual] # Only run in CI - - id: mypy-3.13 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward - name: Run mypy for Python 3.13 - entry: python tools/pre_commit/mypy.py "3.13" - <<: *mypy_common - stages: [manual] # Only run in CI + language: python + types_or: [python, pyi] + require_serial: true + additional_dependencies: ["mypy==1.20.2", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic] - id: png-lint name: Lint PNG exports from excalidraw entry: tools/pre_commit/png-lint.sh diff --git a/AGENTS.md b/AGENTS.md index 933fa05cfcab..a31d156b4b63 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -114,9 +114,6 @@ pre-commit run --all-files # Run a specific hook: pre-commit run ruff-check --all-files - -# Run mypy as it is in CI: -pre-commit run mypy-3.12 --all-files --hook-stage manual ``` The line length limit for Python code is 88 characters. If you are not sure, use pre-commit to check. diff --git a/CMakeLists.txt b/CMakeLists.txt index 0a5414b4bf41..4ed5ff2573c1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -46,7 +46,7 @@ install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS) # Supported python versions. These versions will be searched in order, the # first match will be selected. These should be kept in sync with setup.py. # -set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13" "3.14") +set(PYTHON_SUPPORTED_VERSIONS "3.11" "3.12" "3.13" "3.14") # Supported AMD GPU architectures. set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1250;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201") @@ -134,40 +134,34 @@ endif() # # spinloop extension (pure CXX; must stay above the non-CUDA device branch so # CPU builds define the target before the early return) -# This extension requires SABI 3.11 since it relies on Py_buffer support. Loading -# failure is handled gracefully on vLLM side for lower Python versions. -# Skip the target entirely on Python < 3.11 so the build doesn't break. +# This extension requires SABI 3.11 since it relies on Py_buffer support. # -if(Python_VERSION VERSION_GREATER_EQUAL "3.11") - set(VLLM_SPINLOOP_EXT_SRC "csrc/spinloop.cpp") - set(SPINLOOP_COMPILE_FLAGS "") - if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64") - list(APPEND SPINLOOP_COMPILE_FLAGS "-mmwaitx") - endif() - define_extension_target( - spinloop - DESTINATION vllm - LANGUAGE CXX - SOURCES ${VLLM_SPINLOOP_EXT_SRC} - COMPILE_FLAGS ${SPINLOOP_COMPILE_FLAGS} - USE_SABI 3.11 - WITH_SOABI) +set(VLLM_SPINLOOP_EXT_SRC "csrc/spinloop.cpp") +set(SPINLOOP_COMPILE_FLAGS "") +if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64") + list(APPEND SPINLOOP_COMPILE_FLAGS "-mmwaitx") endif() +define_extension_target( + spinloop + DESTINATION vllm + LANGUAGE CXX + SOURCES ${VLLM_SPINLOOP_EXT_SRC} + COMPILE_FLAGS ${SPINLOOP_COMPILE_FLAGS} + USE_SABI 3.11 + WITH_SOABI) # # fs_io extension (pure CXX; must stay above the non-CUDA device branch # so CPU builds define the target before the early return). # GIL-releasing filesystem helpers for FileSystemTierManager. # -if(Python_VERSION VERSION_GREATER_EQUAL "3.11") - define_extension_target( - fs_io_C - DESTINATION vllm - LANGUAGE CXX - SOURCES csrc/fs_io.cpp - USE_SABI 3.11 - WITH_SOABI) -endif() +define_extension_target( + fs_io_C + DESTINATION vllm + LANGUAGE CXX + SOURCES csrc/fs_io.cpp + USE_SABI 3.11 + WITH_SOABI) # # Forward the non-CUDA device extensions to external CMake scripts. diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index 655a956ba7cd..b3943a122564 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -13,7 +13,7 @@ # vllm-dev: used for development # # Build arguments: -# PYTHON_VERSION=3.13|3.12 (default)|3.11|3.10 +# PYTHON_VERSION=3.13|3.12 (default)|3.11 # VLLM_CPU_X86=false (default)|true (for cross-compilation) # VLLM_CPU_ARM_BF16=false (default)|true (for cross-compilation) # diff --git a/docs/contributing/README.md b/docs/contributing/README.md index f8e9cfcbb1a5..e7a6448ca63e 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -74,7 +74,7 @@ For an optimized workflow when iterating on C++/CUDA kernels, see the [Increment For JIT kernel warmup conventions, see [JIT Kernel Warmup](./jit_kernel_warmup.md). !!! tip - vLLM is compatible with Python versions 3.10 to 3.13. However, vLLM's default [Dockerfile](../../docker/Dockerfile) ships with Python 3.12 and tests in CI (except `mypy`) are run with Python 3.12. + vLLM is compatible with Python versions 3.11 to 3.14. However, vLLM's default [Dockerfile](../../docker/Dockerfile) ships with Python 3.12 and tests in CI (except `mypy`) are run with Python 3.12. Therefore, we recommend developing with Python 3.12 to minimise the chance of your local environment clashing with our CI environment. @@ -102,14 +102,6 @@ integration, use the configuration files in `tools/pre_commit/`: pre-commit run -a # runs on all files (short for --all-files) ``` - --- - - Some `pre-commit` hooks only run in CI. If you need to, you can run them locally with: - - ```bash - pre-commit run --hook-stage manual mypy-3.11 - ``` - ### Documentation MkDocs is a fast, simple and downright gorgeous static site generator that's geared towards building project documentation. Documentation source files are written in Markdown, and configured with a single YAML configuration file, [mkdocs.yaml](../../mkdocs.yaml). diff --git a/docs/deployment/frameworks/skypilot.md b/docs/deployment/frameworks/skypilot.md index e032769f15e8..309367068d6b 100644 --- a/docs/deployment/frameworks/skypilot.md +++ b/docs/deployment/frameworks/skypilot.md @@ -37,7 +37,7 @@ See the vLLM SkyPilot YAML for serving, [serving.yaml](https://github.com/skypil HF_TOKEN: # Change to your own huggingface token, or use --env to pass. setup: | - conda create -n vllm python=3.10 -y + conda create -n vllm python=3.12 -y conda activate vllm pip install vllm==0.4.0.post1 @@ -136,7 +136,7 @@ SkyPilot can scale up the service to multiple service replicas with built-in aut HF_TOKEN: # Change to your own huggingface token, or use --env to pass. setup: | - conda create -n vllm python=3.10 -y + conda create -n vllm python=3.12 -y conda activate vllm pip install vllm==0.4.0.post1 @@ -248,7 +248,7 @@ This will scale the service up to when the QPS exceeds 2 for each replica. HF_TOKEN: # Change to your own huggingface token, or use --env to pass. setup: | - conda create -n vllm python=3.10 -y + conda create -n vllm python=3.12 -y conda activate vllm pip install vllm==0.4.0.post1 @@ -293,7 +293,7 @@ It is also possible to access the Llama-3 service with a separate GUI frontend, cpus: 2 setup: | - conda create -n vllm python=3.10 -y + conda create -n vllm python=3.12 -y conda activate vllm # Install Gradio for web UI. diff --git a/docs/getting_started/installation/cpu.md b/docs/getting_started/installation/cpu.md index fbd20f334a6c..d27753413823 100644 --- a/docs/getting_started/installation/cpu.md +++ b/docs/getting_started/installation/cpu.md @@ -30,7 +30,7 @@ When open a Github issue about the CPU backend, please add `[CPU Backend]` in th ## Requirements -- Python: 3.10 -- 3.13 +- Python: 3.11 -- 3.14 === "Intel/AMD x86" diff --git a/docs/getting_started/installation/gpu.md b/docs/getting_started/installation/gpu.md index 91d933dd4e86..fc8d9148c56b 100644 --- a/docs/getting_started/installation/gpu.md +++ b/docs/getting_started/installation/gpu.md @@ -25,7 +25,7 @@ vLLM is a Python library that supports the following GPU variants. Select your G ## Requirements - OS: Linux -- Python: 3.10 -- 3.13 +- Python: 3.11 -- 3.14 !!! note vLLM does not support Windows natively. To run vLLM on Windows, you can use the Windows Subsystem for Linux (WSL) with a compatible Linux distribution, or use some community-maintained forks, e.g. [https://github.com/SystemPanic/vllm-windows](https://github.com/SystemPanic/vllm-windows). diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index 6f17094a9f47..afa9713341b2 100644 --- a/docs/getting_started/quickstart.md +++ b/docs/getting_started/quickstart.md @@ -8,7 +8,7 @@ This guide will help you quickly get started with vLLM to perform: ## Prerequisites - OS: Linux -- Python: 3.10 -- 3.13 +- Python: 3.11 -- 3.14 !!! note vLLM also works on macOS with [vLLM-Metal](https://github.com/vllm-project/vllm-metal) for Apple Silicon GPU acceleration. See the [GPU installation guide](installation/gpu.md) and select the "Apple Silicon" tab. diff --git a/examples/features/structured_outputs/pyproject.toml b/examples/features/structured_outputs/pyproject.toml index 5e366ab0a03d..d5ffe3c8d462 100644 --- a/examples/features/structured_outputs/pyproject.toml +++ b/examples/features/structured_outputs/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "examples-online-structured-outputs" -requires-python = ">=3.10, <3.14" +requires-python = ">=3.11, <3.15" dependencies = ["openai==1.78.1", "pydantic==2.11.4"] version = "0.0.0" diff --git a/pyproject.toml b/pyproject.toml index 53ce6dfbe06d..bc924440fa59 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,6 @@ license-files = ["LICENSE"] readme = "README.md" description = "A high-throughput and memory-efficient inference and serving engine for LLMs" classifiers = [ - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", @@ -32,7 +31,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Scientific/Engineering :: Information Analysis", ] -requires-python = ">=3.10,<3.15" +requires-python = ">=3.11,<3.15" dynamic = [ "version", "dependencies", "optional-dependencies"] [project.urls] @@ -103,6 +102,8 @@ ignore = [ "B007", # f-string format "UP032", + # `StrEnum` changes `str()` and `format()` output of members + "UP042", ] [tool.ruff.lint.pydocstyle] diff --git a/requirements/cpu.txt b/requirements/cpu.txt index 191ed26ae047..0570517e74cb 100644 --- a/requirements/cpu.txt +++ b/requirements/cpu.txt @@ -22,7 +22,6 @@ torchcodec >= 0.14; platform_machine != "riscv64" and platform_machine != "ppc64 intel-openmp==2024.2.1; platform_machine == "x86_64" # triton-cpu, only for x86_64 CPUs -triton @ https://github.com/triton-lang/triton-cpu/releases/download/v3.8.0/triton-3.8.0+cpu-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_machine == "x86_64" and python_version == "3.10" triton @ https://github.com/triton-lang/triton-cpu/releases/download/v3.8.0/triton-3.8.0+cpu-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_machine == "x86_64" and python_version == "3.11" triton @ https://github.com/triton-lang/triton-cpu/releases/download/v3.8.0/triton-3.8.0+cpu-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_machine == "x86_64" and python_version >= "3.12" diff --git a/setup.py b/setup.py index 0ac53a45b5b1..5e49f74e6fcd 100644 --- a/setup.py +++ b/setup.py @@ -268,8 +268,8 @@ def configure(self, ext: CMakeExtension) -> None: cfg = envs.CMAKE_BUILD_TYPE or default_cfg cmake_args = [ - "-DCMAKE_BUILD_TYPE={}".format(cfg), - "-DVLLM_TARGET_DEVICE={}".format(VLLM_TARGET_DEVICE), + f"-DCMAKE_BUILD_TYPE={cfg}", + f"-DVLLM_TARGET_DEVICE={VLLM_TARGET_DEVICE}", ] verbose = envs.VERBOSE @@ -293,7 +293,7 @@ def configure(self, ext: CMakeExtension) -> None: # Pass the python executable to cmake so it can find an exact # match. - cmake_args += ["-DVLLM_PYTHON_EXECUTABLE={}".format(sys.executable)] + cmake_args += [f"-DVLLM_PYTHON_EXECUTABLE={sys.executable}"] # Pass the python path to cmake so it can reuse the build dependencies # on subsequent calls to python. @@ -305,7 +305,7 @@ def configure(self, ext: CMakeExtension) -> None: # To override this, set the FETCHCONTENT_BASE_DIR environment variable. fc_base_dir = os.path.join(ROOT_DIR, ".deps") fc_base_dir = os.environ.get("FETCHCONTENT_BASE_DIR", fc_base_dir) - cmake_args += ["-DFETCHCONTENT_BASE_DIR={}".format(fc_base_dir)] + cmake_args += [f"-DFETCHCONTENT_BASE_DIR={fc_base_dir}"] # # Setup parallelism and build tool @@ -313,13 +313,13 @@ def configure(self, ext: CMakeExtension) -> None: num_jobs, nvcc_threads = self.compute_num_jobs() if nvcc_threads: - cmake_args += ["-DNVCC_THREADS={}".format(nvcc_threads)] + cmake_args += [f"-DNVCC_THREADS={nvcc_threads}"] if is_ninja_available(): build_tool = ["-G", "Ninja"] cmake_args += [ "-DCMAKE_JOB_POOL_COMPILE:STRING=compile", - "-DCMAKE_JOB_POOLS:STRING=compile={}".format(num_jobs), + f"-DCMAKE_JOB_POOLS:STRING=compile={num_jobs}", ] else: # Default build tool to whatever cmake picks. @@ -1374,7 +1374,7 @@ def _read_requirements(filename: str) -> list[str]: # copying the relevant .py files from the source repository. ext_modules.append(CMakeExtension(name="vllm.triton_kernels", optional=True)) -if not _is_xpu() and sys.version_info >= (3, 11): +if not _is_xpu(): ext_modules.append(CMakeExtension(name="vllm.spinloop")) ext_modules.append(CMakeExtension(name="vllm.fs_io_C")) diff --git a/tests/entrypoints/launchers/test_dp_supervisor.py b/tests/entrypoints/launchers/test_dp_supervisor.py index c10f7ee69899..174157f12a5d 100644 --- a/tests/entrypoints/launchers/test_dp_supervisor.py +++ b/tests/entrypoints/launchers/test_dp_supervisor.py @@ -630,7 +630,7 @@ async def _poll_until_api_server_running( timeout=request_timeout, ): return - except (aiohttp.ClientError, asyncio.TimeoutError) as exc: + except (TimeoutError, aiohttp.ClientError) as exc: last_exc = exc print("Test detected not started yet, sleeping for 1s") remaining_s = deadline - time.monotonic() diff --git a/tests/entrypoints/launchers/test_shutdown.py b/tests/entrypoints/launchers/test_shutdown.py index f03afbff09d9..2a78b4e80ed6 100644 --- a/tests/entrypoints/launchers/test_shutdown.py +++ b/tests/entrypoints/launchers/test_shutdown.py @@ -285,7 +285,7 @@ async def test_wait_timeout_completes_requests(): sigterm_sent.set() await asyncio.wait_for(request_task, timeout=_SHUTDOWN_DETECTION_TIMEOUT) - except asyncio.TimeoutError: + except TimeoutError: pass finally: state.stop_requesting = True @@ -463,7 +463,7 @@ async def test_abort_timeout_fails_inflight_requests(): try: await asyncio.wait_for(request_task, timeout=5) - except asyncio.TimeoutError: + except TimeoutError: pass finally: state.stop_requesting = True @@ -592,7 +592,7 @@ async def test_multi_api_server_shutdown(): try: await asyncio.wait_for(request_task, timeout=_SHUTDOWN_DETECTION_TIMEOUT) - except asyncio.TimeoutError: + except TimeoutError: pass finally: state.stop_requesting = True diff --git a/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py b/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py index 070c8329674b..c19765c462c4 100644 --- a/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py +++ b/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py @@ -579,7 +579,7 @@ async def test_serialize_and_serve_entrypoints(tmp_path): try: await asyncio.wait_for(fut, 180) - except asyncio.TimeoutError: + except TimeoutError: pytest.fail("Server did not start successfully") finally: proc.terminate() diff --git a/tests/multimodal/media/test_connector.py b/tests/multimodal/media/test_connector.py index 5949c66b562a..f688b51534c8 100644 --- a/tests/multimodal/media/test_connector.py +++ b/tests/multimodal/media/test_connector.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import asyncio import mimetypes import os import shutil @@ -292,7 +291,7 @@ async def test_fetch_video_http(video_url: str, num_frames: int): try: video_sync, metadata_sync = connector.fetch_video(video_url) video_async, metadata_async = await connector.fetch_video_async(video_url) - except (TimeoutError, asyncio.TimeoutError) as e: + except TimeoutError as e: pytest.skip(f"Timeout fetching video (CI network flakiness): {e}") assert np.array_equal(video_sync, video_async) @@ -324,7 +323,7 @@ async def test_fetch_video_http_with_dynamic_loader( try: video_sync, metadata_sync = connector.fetch_video(video_url) video_async, metadata_async = await connector.fetch_video_async(video_url) - except (TimeoutError, asyncio.TimeoutError) as e: + except TimeoutError as e: pytest.skip(f"Timeout fetching video (CI network flakiness): {e}") assert np.array_equal(video_sync, video_async) @@ -402,7 +401,7 @@ async def test_allowed_media_domains(video_url: str, num_frames: int): try: video_sync, metadata_sync = connector.fetch_video(video_url) video_async, metadata_async = await connector.fetch_video_async(video_url) - except (TimeoutError, asyncio.TimeoutError) as e: + except TimeoutError as e: pytest.skip(f"Timeout fetching video (CI network flakiness): {e}") assert np.array_equal(video_sync, video_async) diff --git a/tests/v1/distributed/test_async_llm_dp.py b/tests/v1/distributed/test_async_llm_dp.py index ecdef1f38b0b..fb17ccb1d0fc 100644 --- a/tests/v1/distributed/test_async_llm_dp.py +++ b/tests/v1/distributed/test_async_llm_dp.py @@ -719,7 +719,7 @@ async def consume_gen(req_id: str) -> None: # Drive the staggered barrier. Old code deadlocks here. try: await asyncio.wait_for(client.call_utility_async("barrier"), timeout=30) - except asyncio.TimeoutError: + except TimeoutError: for t in mid_barrier_tasks: t.cancel() pytest.fail( diff --git a/tools/ci/kernel_symbol_map.py b/tools/ci/kernel_symbol_map.py index 2a1e8964d055..c88f2d9ad551 100755 --- a/tools/ci/kernel_symbol_map.py +++ b/tools/ci/kernel_symbol_map.py @@ -255,7 +255,7 @@ def main() -> int: or os.environ.get("VLLM_BUILD_COMMIT") or "", "cuda": cuda_release(nvcc), - "generated_at": dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds"), + "generated_at": dt.datetime.now(dt.UTC).isoformat(timespec="seconds"), "build_dirs": [], "incomplete": False, "errors": [], diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index 126d20094a2a..531f6fb8a2b4 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -88,7 +88,11 @@ class ForbiddenImport: allowed_pattern=re.compile(r"^\s*import\s+pybase64(\s*|\s+as\s+base64\s*)$"), ), "re": ForbiddenImport( - pattern=r"^\s*(?:import\s+re(?:$|\s|,)|from\s+re\s+import)", + # regex has no equivalent of the private sre modules + pattern=( + r"^\s*(?:import\s+re(?:$|\s|,)" + r"|from\s+re\s+import(?!(?:\s*,?\s*_(?:constants|parser))+\s*(?:#.*)?$))" + ), tip="Replace 'import re' with 'import regex as re' or 'import regex'.", allowed_pattern=re.compile(r"^\s*import\s+regex(\s*|\s+as\s+re\s*)$"), allowed_files={"setup.py"}, diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py index 245d0daf5e18..876e19e1fdb9 100755 --- a/tools/pre_commit/mypy.py +++ b/tools/pre_commit/mypy.py @@ -10,7 +10,7 @@ python tools/pre_commit/mypy.py Args: - python_version: Python version to use (e.g., "3.10") or "local" to use + python_version: Python version to use (e.g., "3.11") or "local" to use the local Python version. changed_files: List of changed files to check. @@ -152,7 +152,7 @@ def mypy( Args: targets: List of files or directories to check. - python_version: Python version to use (e.g., "3.10") or None to use + python_version: Python version to use (e.g., "3.11") or None to use the default mypy version. follow_imports: Value for the --follow-imports option or None to use the default mypy behavior. diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py index 9e3c0b57e45f..895e5c6c1890 100644 --- a/vllm/benchmarks/serve.py +++ b/vllm/benchmarks/serve.py @@ -253,7 +253,7 @@ async def fetch_spec_decode_metrics( num_accepted_tokens=num_accepted_tokens, accepted_per_pos=accepted_per_pos, ) - except (aiohttp.ClientError, asyncio.TimeoutError): + except (TimeoutError, aiohttp.ClientError): return None @@ -315,7 +315,7 @@ async def fetch_diffusion_metrics( num_canvas_positions=num_canvas_positions, num_committed_tokens=num_committed_tokens, ) - except (aiohttp.ClientError, asyncio.TimeoutError): + except (TimeoutError, aiohttp.ClientError): return None diff --git a/vllm/benchmarks/sweep/plot.py b/vllm/benchmarks/sweep/plot.py index 59a90636f0e0..6b36a7acc601 100644 --- a/vllm/benchmarks/sweep/plot.py +++ b/vllm/benchmarks/sweep/plot.py @@ -8,9 +8,9 @@ from functools import partial from pathlib import Path from types import TracebackType -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING, ClassVar, Self -from typing_extensions import Self, override +from typing_extensions import override from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.utils.collection_utils import full_groupby diff --git a/vllm/benchmarks/sweep/serve_workload.py b/vllm/benchmarks/sweep/serve_workload.py index a488837e2e6f..acf0ce5f7f3d 100644 --- a/vllm/benchmarks/sweep/serve_workload.py +++ b/vllm/benchmarks/sweep/serve_workload.py @@ -4,10 +4,9 @@ import math from dataclasses import dataclass, fields from pathlib import Path -from typing import ClassVar, Literal, get_args +from typing import ClassVar, Literal, assert_never, get_args import numpy as np -from typing_extensions import assert_never from vllm.benchmarks.datasets import DEFAULT_NUM_PROMPTS from vllm.utils.argparse_utils import FlexibleArgumentParser diff --git a/vllm/benchmarks/sweep/server.py b/vllm/benchmarks/sweep/server.py index 87d841ac827f..dff7e069242b 100644 --- a/vllm/benchmarks/sweep/server.py +++ b/vllm/benchmarks/sweep/server.py @@ -6,9 +6,9 @@ import subprocess import time from types import TracebackType +from typing import Self import requests -from typing_extensions import Self class ServerProcess: diff --git a/vllm/config/engram.py b/vllm/config/engram.py index fd95bc2acd2b..7a7ed8f277ef 100644 --- a/vllm/config/engram.py +++ b/vllm/config/engram.py @@ -1,10 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Self from pydantic import model_validator -from typing_extensions import Self from vllm.config.utils import config, get_hash_factors, hash_factors diff --git a/vllm/config/lora.py b/vllm/config/lora.py index 6c994a82f8de..a77660f3b130 100644 --- a/vllm/config/lora.py +++ b/vllm/config/lora.py @@ -2,11 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, Self import torch from pydantic import ConfigDict, Field, model_validator -from typing_extensions import Self from vllm import envs from vllm.config.utils import config diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 7df8f4e0a4ed..37affbaf9ebf 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -5,13 +5,12 @@ import socket from collections.abc import Callable from datetime import timedelta -from typing import TYPE_CHECKING, Any, Literal, overload +from typing import TYPE_CHECKING, Any, Literal, Self, overload import regex as re import torch from pydantic import Field, field_validator, model_validator from torch.distributed import ProcessGroup, ReduceOp, Store -from typing_extensions import Self import vllm.envs as envs from vllm.config.fault_tolerance import FaultToleranceConfig diff --git a/vllm/config/profiler.py b/vllm/config/profiler.py index 7547546d140d..0ad920e21eda 100644 --- a/vllm/config/profiler.py +++ b/vllm/config/profiler.py @@ -2,10 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os -from typing import Any, Literal +from typing import Any, Literal, Self from pydantic import Field, model_validator -from typing_extensions import Self from vllm.config.utils import config from vllm.logger import init_logger diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index 1fb94d3efd07..0d10103ed4ec 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -3,10 +3,9 @@ from collections.abc import Callable from dataclasses import InitVar -from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast +from typing import TYPE_CHECKING, Any, ClassVar, Literal, Self, cast from pydantic import Field, field_validator -from typing_extensions import Self from vllm.config.utils import config from vllm.logger import init_logger diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index cd42206f983a..4c61ead82dab 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -4,10 +4,9 @@ import copy import functools from collections.abc import Callable, Mapping -from typing import TYPE_CHECKING, Any, Literal, get_args +from typing import TYPE_CHECKING, Any, Literal, Self, get_args from pydantic import Field, SkipValidation, field_validator, model_validator -from typing_extensions import Self from vllm.config import LoadConfig from vllm.config.cache import CacheDType diff --git a/vllm/config/structured_outputs.py b/vllm/config/structured_outputs.py index 059581e37e65..4a13900b5622 100644 --- a/vllm/config/structured_outputs.py +++ b/vllm/config/structured_outputs.py @@ -1,10 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Any, Literal +from typing import Any, Literal, Self from pydantic import model_validator -from typing_extensions import Self from vllm.config.utils import config from vllm.utils.hashing import safe_hash diff --git a/vllm/config/watermarking.py b/vllm/config/watermarking.py index 8701c7b72e4d..4d559c29ce64 100644 --- a/vllm/config/watermarking.py +++ b/vllm/config/watermarking.py @@ -2,10 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import hashlib -from typing import Literal +from typing import Literal, Self from pydantic import Field, model_validator -from typing_extensions import Self from vllm.config.utils import config from vllm.logger import init_logger diff --git a/vllm/distributed/utils.py b/vllm/distributed/utils.py index c780316bb4e1..3618361f2a66 100644 --- a/vllm/distributed/utils.py +++ b/vllm/distributed/utils.py @@ -38,9 +38,7 @@ # We prefer to use os.sched_yield as it results in tighter polling loops, # measured to be around 3e-7 seconds. However on earlier versions of Python # os.sched_yield() does not release the GIL, so we fall back to time.sleep(0) -USE_SCHED_YIELD = (sys.version_info[:3] >= (3, 11, 1)) or ( - sys.version_info[:2] == (3, 10) and sys.version_info[2] >= 8 -) +USE_SCHED_YIELD = sys.version_info[:3] >= (3, 11, 1) def sched_yield(): @@ -52,8 +50,8 @@ def sched_yield(): def ensure_divisibility(numerator, denominator): """Ensure that numerator is divisible by the denominator.""" - assert numerator % denominator == 0, "{} is not divisible by {}".format( - numerator, denominator + assert numerator % denominator == 0, ( + f"{numerator} is not divisible by {denominator}" ) @@ -146,9 +144,7 @@ def get_pp_indices( try: partitions = [int(layer) for layer in partition_list_str.split(",")] except ValueError as err: - raise ValueError( - "Invalid partition string: {}".format(partition_list_str) - ) from err + raise ValueError(f"Invalid partition string: {partition_list_str}") from err if len(partitions) != pp_size: raise ValueError(f"{len(partitions)=} does not match {pp_size=}.") if sum(partitions) != num_hidden_layers: diff --git a/vllm/distributed/weight_transfer/base.py b/vllm/distributed/weight_transfer/base.py index 2fb60fb71b41..3140b70e9656 100644 --- a/vllm/distributed/weight_transfer/base.py +++ b/vllm/distributed/weight_transfer/base.py @@ -11,12 +11,12 @@ ClassVar, Generic, Protocol, + Self, TypeVar, runtime_checkable, ) import torch -from typing_extensions import Self if TYPE_CHECKING: from vllm.config import VllmConfig diff --git a/vllm/distributed/weight_transfer/ipc_engine.py b/vllm/distributed/weight_transfer/ipc_engine.py index 15864811e96c..13da117b99e2 100644 --- a/vllm/distributed/weight_transfer/ipc_engine.py +++ b/vllm/distributed/weight_transfer/ipc_engine.py @@ -4,12 +4,11 @@ import pickle from dataclasses import asdict, dataclass -from typing import TYPE_CHECKING, Any, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar, Self import pybase64 as base64 import torch from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor -from typing_extensions import Self from vllm import envs from vllm.config.weight_transfer import WeightTransferConfig diff --git a/vllm/distributed/weight_transfer/m2n_trainer.py b/vllm/distributed/weight_transfer/m2n_trainer.py index 68c62f72dbac..2d1d9f47ec02 100644 --- a/vllm/distributed/weight_transfer/m2n_trainer.py +++ b/vllm/distributed/weight_transfer/m2n_trainer.py @@ -15,10 +15,9 @@ from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, ClassVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, Self, cast import torch -from typing_extensions import Self from vllm.distributed.weight_transfer.base import ( TrainerInitInfo, diff --git a/vllm/distributed/weight_transfer/nccl_engine.py b/vllm/distributed/weight_transfer/nccl_engine.py index 5b98845870fd..77d8b608727e 100644 --- a/vllm/distributed/weight_transfer/nccl_engine.py +++ b/vllm/distributed/weight_transfer/nccl_engine.py @@ -5,10 +5,9 @@ from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor from dataclasses import asdict, dataclass -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING, ClassVar, Self import torch -from typing_extensions import Self if TYPE_CHECKING: from vllm.config import VllmConfig diff --git a/vllm/distributed/weight_transfer/sharded_rdt_trainer.py b/vllm/distributed/weight_transfer/sharded_rdt_trainer.py index 6b755e772e18..51ea401ef92d 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_trainer.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_trainer.py @@ -24,7 +24,7 @@ from collections.abc import Callable, Collection from concurrent.futures import ThreadPoolExecutor from dataclasses import asdict, dataclass, field -from typing import Any, ClassVar +from typing import Any, ClassVar, Self import ray import torch @@ -33,7 +33,6 @@ rebuild_cuda_tensor, reduce_tensor, ) -from typing_extensions import Self from vllm.distributed.weight_transfer.base import ( ParamMeta, diff --git a/vllm/distributed/weight_transfer/sparse_nccl_engine.py b/vllm/distributed/weight_transfer/sparse_nccl_engine.py index 9c704ab344ea..01b29529deb7 100644 --- a/vllm/distributed/weight_transfer/sparse_nccl_engine.py +++ b/vllm/distributed/weight_transfer/sparse_nccl_engine.py @@ -11,10 +11,9 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext from dataclasses import asdict, dataclass -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING, ClassVar, Self import torch -from typing_extensions import Self if TYPE_CHECKING: from vllm.config import VllmConfig diff --git a/vllm/entrypoints/chat_utils.py b/vllm/entrypoints/chat_utils.py index 2d875a974bef..f18407e535e9 100644 --- a/vllm/entrypoints/chat_utils.py +++ b/vllm/entrypoints/chat_utils.py @@ -17,6 +17,7 @@ Final, Generic, Literal, + Required, TypeAlias, TypeVar, Union, @@ -48,7 +49,7 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter # pydantic needs the TypedDict from typing_extensions -from typing_extensions import Required, TypedDict, override +from typing_extensions import TypedDict, override from vllm import envs from vllm.config import ModelConfig diff --git a/vllm/entrypoints/launchers/dp_supervisor.py b/vllm/entrypoints/launchers/dp_supervisor.py index 192bbb169814..f2a84c4d9dfd 100644 --- a/vllm/entrypoints/launchers/dp_supervisor.py +++ b/vllm/entrypoints/launchers/dp_supervisor.py @@ -200,7 +200,7 @@ async def _probe_endpoint( # vLLM returns 503 on EngineDeadError, so we should return # immediately if vLLM responds with a non-200 status code. return response.status == HTTPStatus.OK - except (aiohttp.ClientError, asyncio.TimeoutError) as e: + except (TimeoutError, aiohttp.ClientError) as e: # Allow retry of connection errors. logger.debug( "Probe attempt %d/%d failed on port %d: %r", diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index c38789a2e29c..9dfff78ed9ff 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -3,13 +3,12 @@ from collections.abc import Callable, Sequence from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, overload import cloudpickle import torch.nn as nn from pydantic import ValidationError from tqdm.auto import tqdm -from typing_extensions import overload from vllm.config import ( AttentionConfig, diff --git a/vllm/entrypoints/pooling/embed/serving.py b/vllm/entrypoints/pooling/embed/serving.py index a9f6dd310ec7..4fdb73d363df 100644 --- a/vllm/entrypoints/pooling/embed/serving.py +++ b/vllm/entrypoints/pooling/embed/serving.py @@ -1,9 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TypeAlias, cast +from typing import TypeAlias, assert_never, cast from fastapi.responses import JSONResponse, Response, StreamingResponse -from typing_extensions import assert_never from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput diff --git a/vllm/entrypoints/pooling/pooling/serving.py b/vllm/entrypoints/pooling/pooling/serving.py index 56bc7c7865b1..4e4ea1fe1bd5 100644 --- a/vllm/entrypoints/pooling/pooling/serving.py +++ b/vllm/entrypoints/pooling/pooling/serving.py @@ -1,9 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import cast +from typing import assert_never, cast from fastapi.responses import JSONResponse, Response, StreamingResponse -from typing_extensions import assert_never from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger diff --git a/vllm/entrypoints/pooling/scoring/typing.py b/vllm/entrypoints/pooling/scoring/typing.py index a67843e57770..9afaada6d666 100644 --- a/vllm/entrypoints/pooling/scoring/typing.py +++ b/vllm/entrypoints/pooling/scoring/typing.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass -from typing import TypeAlias +from typing import Required, TypeAlias -from typing_extensions import Required, TypedDict +from typing_extensions import TypedDict from vllm.entrypoints.chat_utils import ( ChatCompletionContentPartImageEmbedsParam, diff --git a/vllm/entrypoints/serve/tokenize/api_router.py b/vllm/entrypoints/serve/tokenize/api_router.py index 9cdd82fee3f3..cd515d31d2e8 100644 --- a/vllm/entrypoints/serve/tokenize/api_router.py +++ b/vllm/entrypoints/serve/tokenize/api_router.py @@ -3,10 +3,10 @@ from http import HTTPStatus +from typing import assert_never from fastapi import APIRouter, Depends, FastAPI, Request from fastapi.responses import JSONResponse -from typing_extensions import assert_never from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.entrypoints.serve.tokenize.protocol import ( diff --git a/vllm/inputs/engine.py b/vllm/inputs/engine.py index 26b840f76f3f..aa8c9054335e 100644 --- a/vllm/inputs/engine.py +++ b/vllm/inputs/engine.py @@ -3,9 +3,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Literal, TypeAlias +from typing import TYPE_CHECKING, Literal, NotRequired, TypeAlias, assert_never -from typing_extensions import NotRequired, TypedDict, assert_never +from typing_extensions import TypedDict from vllm.exceptions import VLLMValidationError diff --git a/vllm/inputs/llm.py b/vllm/inputs/llm.py index a95b274c1853..1ef43fd23b17 100644 --- a/vllm/inputs/llm.py +++ b/vllm/inputs/llm.py @@ -3,9 +3,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar, final +from typing import TYPE_CHECKING, Any, NotRequired, TypeAlias, TypeVar, final -from typing_extensions import NotRequired, TypedDict +from typing_extensions import TypedDict if TYPE_CHECKING: import torch diff --git a/vllm/model_executor/kernels/linear/base.py b/vllm/model_executor/kernels/linear/base.py index 0db595a3ca42..af1f597ec60c 100644 --- a/vllm/model_executor/kernels/linear/base.py +++ b/vllm/model_executor/kernels/linear/base.py @@ -3,10 +3,9 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Any, ClassVar, Generic, TypeVar +from typing import Any, ClassVar, Generic, Self, TypeVar import torch -from typing_extensions import Self from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey diff --git a/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py b/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py index 3a053412f3cf..865bebeada54 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py @@ -3,10 +3,9 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import ClassVar +from typing import ClassVar, Self import torch -from typing_extensions import Self from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 from vllm.model_executor.layers.quantization.utils.fp8_utils import ( diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index c12b8e76c5f5..65d1ff2808d7 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -6,11 +6,10 @@ import time import warnings from contextlib import contextmanager, nullcontext -from typing import Any +from typing import Any, assert_never import torch from torch import nn -from typing_extensions import assert_never import vllm.envs as envs from vllm.config import ( diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 433c373c8e35..272b78da50a4 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -21,6 +21,7 @@ ClassVar, Literal, Protocol, + Self, TypeAlias, overload, runtime_checkable, @@ -31,7 +32,7 @@ import torch.nn as nn from torch import Tensor from transformers.models.whisper.tokenization_whisper import LANGUAGES -from typing_extensions import Self, TypeIs +from typing_extensions import TypeIs from vllm.logger import init_logger from vllm.model_executor.layers.quantization import QuantizationConfig diff --git a/vllm/models/deepseek_v4/common/mm_preprocess.py b/vllm/models/deepseek_v4/common/mm_preprocess.py index 62a5c2070f5b..9c1a4664c83d 100644 --- a/vllm/models/deepseek_v4/common/mm_preprocess.py +++ b/vllm/models/deepseek_v4/common/mm_preprocess.py @@ -22,13 +22,12 @@ import math from collections.abc import Mapping, Sequence -from typing import Any +from typing import Any, assert_never import numpy as np import torch from PIL import Image, ImageOps from transformers import BatchFeature -from typing_extensions import assert_never from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict diff --git a/vllm/multimodal/parse.py b/vllm/multimodal/parse.py index 98f8a8b2f3e9..546eb6abfda2 100644 --- a/vllm/multimodal/parse.py +++ b/vllm/multimodal/parse.py @@ -13,11 +13,11 @@ TypeAlias, TypeGuard, TypeVar, + assert_never, ) import numpy as np import torch -from typing_extensions import assert_never from vllm.inputs import ModalityData, MultiModalDataDict, MultiModalUUIDDict from vllm.utils.collection_utils import is_list_of diff --git a/vllm/multimodal/processing/processor.py b/vllm/multimodal/processing/processor.py index 726702035355..abdb58540255 100644 --- a/vllm/multimodal/processing/processor.py +++ b/vllm/multimodal/processing/processor.py @@ -13,10 +13,11 @@ NamedTuple, Protocol, TypeAlias, + assert_never, ) import torch -from typing_extensions import TypeVar, assert_never +from typing_extensions import TypeVar from vllm.config import SchedulerConfig from vllm.inputs import ( diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py index 4301c6917613..4a8bcd0c03d7 100644 --- a/vllm/renderers/hf.py +++ b/vllm/renderers/hf.py @@ -1000,7 +1000,7 @@ async def _render_with_timeout(self, coro): return await coro try: return await asyncio.wait_for(coro, timeout=timeout) - except asyncio.TimeoutError: + except TimeoutError: self._replace_executor() raise TimeoutError( f"Chat template rendering timed out after " diff --git a/vllm/snapshot/manifest.py b/vllm/snapshot/manifest.py index a029c544c52e..c550ee80c5cf 100644 --- a/vllm/snapshot/manifest.py +++ b/vllm/snapshot/manifest.py @@ -6,7 +6,7 @@ import os import stat from pathlib import Path -from typing import Annotated, Any, Literal +from typing import Annotated, Any, Literal, Self from pydantic import ( BaseModel, @@ -19,7 +19,6 @@ field_validator, model_validator, ) -from typing_extensions import Self class SnapshotCompatibilityError(RuntimeError): diff --git a/vllm/snapshot/runtime.py b/vllm/snapshot/runtime.py index fdced5116cdf..83c70e473833 100644 --- a/vllm/snapshot/runtime.py +++ b/vllm/snapshot/runtime.py @@ -24,7 +24,7 @@ from collections.abc import Iterable from contextlib import suppress from dataclasses import dataclass -from datetime import datetime, timezone +from datetime import UTC, datetime from pathlib import Path from typing import NoReturn @@ -872,7 +872,7 @@ def make_manifest( return SnapshotManifest( schema_version=1, boundary="post-engine-init-reloadable-state-released", - created_at=datetime.now(timezone.utc).isoformat(), + created_at=datetime.now(UTC).isoformat(), artifact_bytes=self._artifact_bytes(workdir), **identity.model_dump(), model=source_model, diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index e7064fc252f2..ebf85a11fd25 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -4,10 +4,10 @@ from dataclasses import dataclass, field from functools import lru_cache from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, assert_never import huggingface_hub -from typing_extensions import TypeVar, assert_never +from typing_extensions import TypeVar from vllm import envs from vllm.logger import init_logger diff --git a/vllm/usage/usage_lib.py b/vllm/usage/usage_lib.py index 5197aada59ac..b75a4fa16184 100644 --- a/vllm/usage/usage_lib.py +++ b/vllm/usage/usage_lib.py @@ -69,7 +69,7 @@ def is_usage_stats_enabled(): def _get_current_timestamp_ns() -> int: - return int(datetime.datetime.now(datetime.timezone.utc).timestamp() * 1e9) + return int(datetime.datetime.now(datetime.UTC).timestamp() * 1e9) def _detect_cloud_provider() -> str: diff --git a/vllm/utils/collection_utils.py b/vllm/utils/collection_utils.py index de32ed0c15c2..8a77da86cd56 100644 --- a/vllm/utils/collection_utils.py +++ b/vllm/utils/collection_utils.py @@ -8,9 +8,9 @@ import math from collections import defaultdict from collections.abc import Callable, Generator, Hashable, Iterable, Mapping, Sequence -from typing import Generic, Literal, TypeVar +from typing import Generic, Literal, TypeVar, assert_never, overload -from typing_extensions import TypeIs, assert_never, overload +from typing_extensions import TypeIs T = TypeVar("T") diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index aaacccc7cb5a..04a7ebec2eab 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -12,10 +12,9 @@ import sys from functools import cache from types import ModuleType -from typing import Any +from typing import Any, Never import regex as re -from typing_extensions import Never from vllm.logger import init_logger diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index b6913003be59..cf707611b47d 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -11,10 +11,9 @@ from fractions import Fraction from functools import cached_property from math import prod -from typing import TYPE_CHECKING, TypeVar +from typing import TYPE_CHECKING, Self, TypeVar import torch -from typing_extensions import Self from vllm.logger import init_logger from vllm.utils.math_utils import cdiv, round_up diff --git a/vllm/v1/metrics/perf.py b/vllm/v1/metrics/perf.py index 1cd8fa79f1d3..a2b3f534d88c 100644 --- a/vllm/v1/metrics/perf.py +++ b/vllm/v1/metrics/perf.py @@ -10,12 +10,11 @@ from abc import ABC, abstractmethod from collections.abc import Iterable from dataclasses import asdict, dataclass, field -from typing import Any, Protocol +from typing import Any, Protocol, Self import prometheus_client import torch from pydantic import BaseModel, Field, ValidationError, model_validator -from typing_extensions import Self import vllm.envs as envs from vllm.config import VllmConfig diff --git a/vllm/v1/structured_output/backend_outlines.py b/vllm/v1/structured_output/backend_outlines.py index dab4f2aa0462..655b8b14631f 100644 --- a/vllm/v1/structured_output/backend_outlines.py +++ b/vllm/v1/structured_output/backend_outlines.py @@ -4,10 +4,9 @@ from __future__ import annotations import ast -import importlib import json -import sys from dataclasses import dataclass, field +from re import _constants, _parser # type: ignore[attr-defined] from typing import TYPE_CHECKING import torch @@ -36,19 +35,6 @@ oc = LazyLoader("oc", globals(), "outlines_core") json_schema = LazyLoader("json_schema", globals(), "outlines_core.json_schema") -# Python 3.11+ sre_parse and sre_constants -# are deprecated, so we must import them from re -if sys.version_info >= (3, 11): - # Hack to get around pre-commit regex module rule - # because going through re is the only way to get sre_parse - # and sre_constants in Python 3.11+ - _re = importlib.import_module("re") - sre_parse = _re._parser - sre_constants = _re._constants -else: - import sre_constants - import sre_parse - @dataclass class OutlinesBackend(StructuredOutputBackend): @@ -261,20 +247,20 @@ def subpattern_consumes(parsed) -> bool: tokens = parsed.data if hasattr(parsed, "data") else parsed for ttype, tval in tokens: # literal, character class, or dot always consumes - if ttype in (sre_parse.LITERAL, sre_parse.IN, sre_parse.ANY): + if ttype in (_parser.LITERAL, _parser.IN, _parser.ANY): return True # quantified subpattern: check inner pattern - elif ttype == sre_parse.MAX_REPEAT: + elif ttype == _parser.MAX_REPEAT: _, mx, sub = tval if mx != 0 and subpattern_consumes(sub): return True # alternation: if any branch consumes, the whole does - elif ttype == sre_parse.BRANCH: + elif ttype == _parser.BRANCH: _, branches = tval if any(subpattern_consumes(br) for br in branches): return True # grouped subpattern: recurse into its contents - elif ttype == sre_parse.SUBPATTERN and subpattern_consumes(tval[3]): + elif ttype == _parser.SUBPATTERN and subpattern_consumes(tval[3]): return True # No consumers, return False return False @@ -282,14 +268,11 @@ def subpattern_consumes(parsed) -> bool: tokens = parsed.data if hasattr(parsed, "data") else parsed for ttype, tval in tokens: # Direct anchors or look-around - if ttype == sre_parse.AT or ttype in ( - sre_constants.ASSERT, - sre_constants.ASSERT_NOT, - ): + if ttype == _parser.AT or ttype in (_constants.ASSERT, _constants.ASSERT_NOT): return True # Nested subpattern: check - if ttype == sre_parse.SUBPATTERN: + if ttype == _parser.SUBPATTERN: # tval: (group, add_flags, del_flags, subpattern) if _prefix_needs_context(tval[3]): return True @@ -298,7 +281,7 @@ def subpattern_consumes(parsed) -> bool: # if any branch has a prefix anchor => True, # else if at least one branch consumes => prefix ends => False - elif ttype == sre_parse.BRANCH: + elif ttype == _parser.BRANCH: saw_consumer = False for br in tval[1]: if _prefix_needs_context(br): @@ -309,11 +292,11 @@ def subpattern_consumes(parsed) -> bool: return False # Immediate consumer tokens - elif ttype in (sre_parse.LITERAL, sre_parse.IN, sre_parse.ANY): + elif ttype in (_parser.LITERAL, _parser.IN, _parser.ANY): return False # if subpattern has anchor => True, if it can consume => stop - elif ttype == sre_parse.MAX_REPEAT: + elif ttype == _parser.MAX_REPEAT: if _prefix_needs_context(tval[2]): return True if subpattern_consumes(tval[2]): @@ -327,25 +310,25 @@ def _check_unsupported(parsed) -> None: tokens = parsed.data if hasattr(parsed, "data") else parsed for ttype, tval in tokens: # backreference - if ttype in (sre_parse.GROUPREF, sre_parse.GROUPREF_EXISTS): + if ttype in (_parser.GROUPREF, _parser.GROUPREF_EXISTS): raise ValueError("Backreferences are unsupported.") # look-around assertion - elif ttype in (sre_constants.ASSERT, sre_constants.ASSERT_NOT): + elif ttype in (_constants.ASSERT, _constants.ASSERT_NOT): raise ValueError("Look-Around assertion are unsupported.") # unicode word boundaries - elif ttype == sre_parse.AT: - if tval in (sre_constants.AT_BOUNDARY, sre_constants.AT_NON_BOUNDARY): + elif ttype == _parser.AT: + if tval in (_constants.AT_BOUNDARY, _constants.AT_NON_BOUNDARY): raise ValueError("Unicode word boundaries are unsupported.") - elif ttype == sre_parse.BRANCH: + elif ttype == _parser.BRANCH: # tval is (None, branches) for branch in tval[1]: _check_unsupported(branch) # tval is (min, max, subpattern) - elif ttype == sre_parse.MAX_REPEAT: + elif ttype == _parser.MAX_REPEAT: _check_unsupported(tval[2]) @@ -357,9 +340,9 @@ def validate_regex_is_buildable(pattern: str) -> None: https://docs.rs/regex-automata/latest/regex_automata/dfa/trait.Automaton.html#method.universal_start_state """ try: - parsed = sre_parse.parse(pattern) + parsed = _parser.parse(pattern) - except sre_constants.error as e: + except _constants.error as e: raise VLLMValidationError(f"Error parsing regex: {e}") from e try: diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py index 4bae1f52bec6..e69e55bd8216 100644 --- a/vllm/v1/worker/workspace.py +++ b/vllm/v1/worker/workspace.py @@ -229,12 +229,7 @@ def get_caller_info() -> str: # Walk up the stack skipping WorkspaceManager frames curr_frame = curr_frame.f_back while curr_frame is not None: - # TODO: This only catches instance methods (self), missing - # classmethods and staticmethods. Once Python 3.11+ is the - # minimum supported version, use co_qualname instead: - # qualname = curr_frame.f_code.co_qualname - # if qualname.startswith("WorkspaceManager."): - if isinstance(curr_frame.f_locals.get("self"), WorkspaceManager): + if curr_frame.f_code.co_qualname.startswith("WorkspaceManager."): curr_frame = curr_frame.f_back continue filename = os.path.basename(curr_frame.f_code.co_filename) From c4cd88d9fee92ff080837c6c1cdf1de6fe4013fe Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Thu, 8 Oct 2026 02:31:13 -0700 Subject: [PATCH 083/624] [CI] Make (H100) Distributed DP + EP optional (#60540) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/distributed.yaml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 3710f80bfdb8..78a6af2e0ca4 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -397,6 +397,7 @@ steps: - label: ":nvidia: (H100) Distributed DP + EP" key: distributed-tests-8xh100 timeout_in_minutes: 20 + optional: true device: h100 num_devices: 8 working_dir: "/vllm-workspace/tests" @@ -419,6 +420,7 @@ steps: amd: label: ":amd: (MI355) Distributed DP + EP" dind: false + optional: false device: mi355_8 timeout_in_minutes: 50 depends_on: From 7816afdb259dbb18f9e98d11b05194dd3446d0fa Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Thu, 8 Oct 2026 11:33:41 +0200 Subject: [PATCH 084/624] [Bugfix][Core] Drain async KV loads in pause(mode="wait"), and pause the drain test only once a request is in flight (#60572) Signed-off-by: aoshen02 Signed-off-by: Stefan Koncarevic Co-authored-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- tests/v1/core/test_scheduler.py | 90 +++++++++++++++++++++++ tests/v1/distributed/test_async_llm_dp.py | 7 +- vllm/v1/core/sched/interface.py | 2 +- vllm/v1/core/sched/scheduler.py | 12 ++- 4 files changed, 103 insertions(+), 8 deletions(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 7a8a1a3e4ea5..615dd82a1c7b 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -2573,6 +2573,96 @@ def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]): return initial_ecos +def _runner_output( + req_ids: list[str], finished_recving: set[str] | None = None +) -> ModelRunnerOutput: + return ModelRunnerOutput( + req_ids=req_ids, + req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)}, + sampled_token_ids=[[1] for _ in req_ids], + kv_connector_output=KVConnectorOutput(finished_recving=finished_recving), + ) + + +@pytest.mark.parametrize( + "load_done_before_pause", [False, True], ids=["load-in-flight", "load-done"] +) +def test_pause_new_drains_async_kv_loads(load_done_before_pause: bool): + """A PAUSED_NEW (pause mode="wait") drain must finish requests that hold KV + blocks for an async load, keep blockless queued requests queued, and leave + no blocks behind, so the post-drain cache reset succeeds.""" + block_size = 16 + scheduler = create_scheduler( + enable_prefix_caching=True, + use_kv_connector=mock_kv(matched_tokens=2 * block_size, is_async=True), + block_size=block_size, + ) + loading, queued = create_requests( + num_requests=2, + num_tokens=3 * block_size, + max_tokens=2, + block_size=block_size, + req_ids=["loading", "queued"], + ) + scheduler.add_request(loading) + output = scheduler.schedule() + assert loading.status == RequestStatus.WAITING_FOR_REMOTE_KVS + pending_recv = {"loading"} + if load_done_before_pause: + scheduler.update_from_output(output, _runner_output([], pending_recv)) + pending_recv = None + else: + scheduler.update_from_output(output, _runner_output([])) + + scheduler.set_pause_state(PauseState.PAUSED_NEW) + scheduler.add_request(queued) + + num_steps = 0 + while scheduler.has_requests(): + num_steps += 1 + assert num_steps <= 4, "pause drain did not finish" + output = scheduler.schedule() + assert set(output.num_scheduled_tokens) <= {"loading"} + scheduler.update_from_output( + output, _runner_output(list(output.num_scheduled_tokens), pending_recv) + ) + pending_recv = None + + # What pause runs once drained; it raises if the load still holds blocks. + assert scheduler.reset_prefix_cache(reset_running_requests=True) + assert loading.status == RequestStatus.FINISHED_LENGTH_CAPPED + assert loading.request_id not in scheduler.requests + assert list(scheduler.waiting) == [queued] + assert queued.status == RequestStatus.WAITING + + scheduler.set_pause_state(PauseState.UNPAUSED) + scheduler.schedule() + assert queued.status == RequestStatus.WAITING_FOR_REMOTE_KVS + + +def test_pause_new_ignores_streaming_session_waiting_for_input(): + """A session idle in WAITING_FOR_STREAMING_REQ cannot progress, so it must + not keep a PAUSED_NEW drain from completing.""" + stop_token = 7 + scheduler = create_scheduler() + (session,) = create_requests( + num_requests=1, num_tokens=4, stop_token_ids=[stop_token], req_ids=["s"] + ) + session.resumable = True + scheduler.add_request(session) + output = scheduler.schedule() + scheduler.update_from_output( + output, + ModelRunnerOutput( + req_ids=["s"], req_id_to_index={"s": 0}, sampled_token_ids=[[stop_token]] + ), + ) + assert session in scheduler.kv_holding_waiting + + scheduler.set_pause_state(PauseState.PAUSED_NEW) + assert scheduler.get_num_unfinished_requests() == 0 + + @pytest.mark.parametrize( ("load_modes", "expected_has_sync_loads"), [ diff --git a/tests/v1/distributed/test_async_llm_dp.py b/tests/v1/distributed/test_async_llm_dp.py index fb17ccb1d0fc..f4b4ca2a7902 100644 --- a/tests/v1/distributed/test_async_llm_dp.py +++ b/tests/v1/distributed/test_async_llm_dp.py @@ -747,12 +747,13 @@ async def test_dp_pause_wait_mode_drains_in_flight(): prompt=DP_PAUSE_PROMPT, params=SamplingParams(max_tokens=64), ) + # Wait for the first output, so drain has something to observe: a + # request still queued holds no blocks and is not waited on. + out = await asyncio.wait_for(collector.get(), timeout=30) await engine.pause_generation(mode="wait") assert await engine.is_paused() - while True: + while not out.finished: out = await asyncio.wait_for(collector.get(), timeout=30) - if out.finished: - break await engine.resume_generation() async for out in engine.generate( diff --git a/vllm/v1/core/sched/interface.py b/vllm/v1/core/sched/interface.py index c868d249c8a5..dd88d99c1f79 100644 --- a/vllm/v1/core/sched/interface.py +++ b/vllm/v1/core/sched/interface.py @@ -26,7 +26,7 @@ class PauseState(enum.IntEnum): - UNPAUSED: Normal operation - PAUSE_NEW: No new requests are scheduled, requests already in - running state are scheduled. + running state or holding KV blocks are scheduled. - PAUSE_ALL: No requests are scheduled """ diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 91ea8e93e42c..9739dbb3d516 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -887,7 +887,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: assert len(scheduled_loras) <= self.lora_config.max_loras # Next, schedule the WAITING requests. - if not preempted_reqs and self._pause_state == PauseState.UNPAUSED: + if not preempted_reqs and self._pause_state != PauseState.PAUSED_ALL: step_skipped_waiting: deque[Request] = deque() step_skipped_kv_holding: deque[Request] = deque() @@ -900,8 +900,11 @@ def skip_request(from_queue: RequestQueue) -> None: step_skipped_waiting.appendleft(request_to_skip) while token_budget > 0: - # Requests holding KV blocks are always drained first. - request_queue = self.kv_holding_waiting or self.waiting + # Requests holding KV blocks are always drained first, and are + # the only ones admitted under PAUSED_NEW (e.g. async KV loads). + request_queue = self.kv_holding_waiting + if not request_queue and self._pause_state == PauseState.UNPAUSED: + request_queue = self.waiting if not request_queue or input_budget <= draft_slots: break # Paused streaming sessions (WAITING_FOR_STREAMING_REQ) are not @@ -2802,7 +2805,8 @@ def get_num_unfinished_requests(self) -> int: return 0 num_running, num_waiting = self.get_request_counts() if self._pause_state == PauseState.PAUSED_NEW: - return num_running + # Requests holding KV blocks still drain; queued ones stay queued. + num_waiting = len(self.kv_holding_waiting) num_waiting -= self.num_waiting_for_streaming_input return num_waiting + num_running From e265b77982e5e2724ce0b2841a2b05fe339d62ee Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Thu, 8 Oct 2026 17:39:56 +0800 Subject: [PATCH 085/624] [CI][Mamba] Use well-conditioned prompts for FlashInfer ReplaySSM parity tests (#60515) Signed-off-by: zjy0516 Co-authored-by: Claude Opus 5.5 --- tests/v1/e2e/test_replayssm_decode.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 1a369c781f03..ed3c6b0217c8 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -25,6 +25,12 @@ "The capital of France is", "Once upon a time, in a small village,", ] +# The NVFP4 model is not reproducible across engine instances; these prompts +# avoid near-ties and decode confidently past the replay window. +FLASHINFER_PROMPTS = [ + "The capital of France is", + "It was the best of times, it was the worst of times,", +] @pytest.fixture(autouse=True) @@ -43,6 +49,7 @@ def _check_replayssm_parity( *, tensor_parallel_size=1, mamba_backend: str = "triton", + prompts: list[str] = PROMPTS, name_1: str = "replayssm", expected_v2: bool | None = None, ): @@ -60,7 +67,7 @@ def _check_replayssm_parity( with vllm_runner(model_name, **common) as llm: if expected_v2 is not None: assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is expected_v2 - baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) + baseline = llm.generate_greedy_logprobs(prompts, max_tokens=32, num_logprobs=5) with vllm_runner( model_name, use_replayssm=True, @@ -69,7 +76,7 @@ def _check_replayssm_parity( ) as llm: if expected_v2 is not None: assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is expected_v2 - replay = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) + replay = llm.generate_greedy_logprobs(prompts, max_tokens=32, num_logprobs=5) check_logprobs_close( outputs_0_lst=baseline, @@ -105,6 +112,7 @@ def test_replayssm_flashinfer_decode_matches_baseline( vllm_runner, model_name, mamba_backend="flashinfer", + prompts=FLASHINFER_PROMPTS, name_1="replayssm_flashinfer", expected_v2=use_v2_model_runner, ) @@ -127,7 +135,9 @@ def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_na }, ) with vllm_runner(model_name, mamba_backend="flashinfer", **common) as llm: - baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) + baseline = llm.generate_greedy_logprobs( + FLASHINFER_PROMPTS, max_tokens=32, num_logprobs=5 + ) with vllm_runner( model_name, use_replayssm=True, @@ -135,7 +145,9 @@ def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_na mamba_backend="flashinfer", **common, ) as llm: - replay = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) + replay = llm.generate_greedy_logprobs( + FLASHINFER_PROMPTS, max_tokens=32, num_logprobs=5 + ) check_logprobs_close( outputs_0_lst=baseline, From 24be8c0bca2fd806548ad84c0999885a65f0805d Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Thu, 8 Oct 2026 02:42:07 -0700 Subject: [PATCH 086/624] [CI] Remove dead in-tree TPU (torch_xla) CI scripts (#60591) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- .../scripts/hardware_ci/run-cpu-test.sh | 7 +- .../hardware_ci/run-tpu-v1-test-part2.sh | 166 ----------------- .../scripts/hardware_ci/run-tpu-v1-test.sh | 172 ------------------ .buildkite/scripts/tpu/cleanup_docker.sh | 24 --- .buildkite/scripts/tpu/config_v6e_1.env | 14 -- .buildkite/scripts/tpu/docker_run_bm.sh | 91 --------- .buildkite/scripts/tpu/quantized_v6e_1.env | 14 -- .buildkite/scripts/tpu/run_bm.sh | 93 ---------- .github/mergify.yml | 4 - .../run_tpu_disagg_accuracy_test.sh | 157 ---------------- .../run_tpu_edge_case_test.sh | 125 ------------- 11 files changed, 3 insertions(+), 864 deletions(-) delete mode 100755 .buildkite/scripts/hardware_ci/run-tpu-v1-test-part2.sh delete mode 100755 .buildkite/scripts/hardware_ci/run-tpu-v1-test.sh delete mode 100755 .buildkite/scripts/tpu/cleanup_docker.sh delete mode 100644 .buildkite/scripts/tpu/config_v6e_1.env delete mode 100755 .buildkite/scripts/tpu/docker_run_bm.sh delete mode 100644 .buildkite/scripts/tpu/quantized_v6e_1.env delete mode 100755 .buildkite/scripts/tpu/run_bm.sh delete mode 100644 tests/v1/kv_connector/nixl_integration/run_tpu_disagg_accuracy_test.sh delete mode 100644 tests/v1/kv_connector/nixl_integration/run_tpu_edge_case_test.sh diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh index eecf70ea2bde..5f9355f9e943 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh @@ -19,10 +19,9 @@ DISK_USAGE_THRESHOLD=${DISK_USAGE_THRESHOLD:-80} CACHE_MAX_AGE=${CACHE_MAX_AGE:-24h} # Reclaim disk only when the host is under pressure, aging out anything unused -# for less than CACHE_MAX_AGE so hot layers survive -- same `--filter -# until=h` pattern the TPU CI scripts already rely on. `docker buildx -# prune --max-used-space` is a no-op on this host's `docker` driver (BuildKit -# embedded in dockerd never enforces the size cap), so we don't use it. +# for less than CACHE_MAX_AGE so hot layers survive (`--filter until=h`). +# `docker buildx prune --max-used-space` is a no-op on this host's `docker` +# driver (BuildKit embedded in dockerd never enforces the size cap), so we don't use it. prune_if_disk_pressure() { local docker_root disk_usage docker_root=$(docker info -f '{{.DockerRootDir}}' 2>/dev/null || true) diff --git a/.buildkite/scripts/hardware_ci/run-tpu-v1-test-part2.sh b/.buildkite/scripts/hardware_ci/run-tpu-v1-test-part2.sh deleted file mode 100755 index f300b7f80991..000000000000 --- a/.buildkite/scripts/hardware_ci/run-tpu-v1-test-part2.sh +++ /dev/null @@ -1,166 +0,0 @@ -#!/bin/bash - -set -xu - - -remove_docker_container() { - docker rm -f tpu-test || true; -} - -trap remove_docker_container EXIT - -# Remove the container that might not be cleaned up in the previous run. -remove_docker_container - -# Build the docker image. -docker build -f docker/Dockerfile.tpu -t vllm-tpu . - -# Set up cleanup. -cleanup_docker() { - # Get Docker's root directory - docker_root=$(docker info -f '{{.DockerRootDir}}') - if [ -z "$docker_root" ]; then - echo "Failed to determine Docker root directory." - exit 1 - fi - echo "Docker root directory: $docker_root" - # Check disk usage of the filesystem where Docker's root directory is located - disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//') - # Define the threshold - threshold=70 - if [ "$disk_usage" -gt "$threshold" ]; then - echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..." - # Remove dangling images (those that are not tagged and not used by any container) - docker image prune -f - # Remove unused volumes / force the system prune for old images as well. - docker volume prune -f && docker system prune --force --filter "until=72h" --all - echo "Docker images and volumes cleanup completed." - else - echo "Disk usage is below $threshold%. No cleanup needed." - fi -} -cleanup_docker - -# For HF_TOKEN. -source /etc/environment - -docker run --privileged --net host --shm-size=16G -it \ - -e "HF_TOKEN=$HF_TOKEN" --name tpu-test \ - vllm-tpu /bin/bash -c ' -set -e # Exit immediately if a command exits with a non-zero status. -set -u # Treat unset variables as an error. - -echo "--- Starting script inside Docker container ---" - -# Create results directory -RESULTS_DIR=$(mktemp -d) -# If mktemp fails, set -e will cause the script to exit. -echo "Results will be stored in: $RESULTS_DIR" - -# Install dependencies -echo "--- Installing Python dependencies ---" -python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \ - && python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \ - && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.12" \ - && python3 -m pip install --progress-bar off hf-transfer tblib==3.1.0 -echo "--- Python dependencies installed ---" - -export VLLM_XLA_CHECK_RECOMPILATION=1 -export VLLM_XLA_CACHE_PATH= - -echo "--- Hardware Information ---" -# tpu-info -echo "--- Starting Tests ---" -set +e -overall_script_exit_code=0 - -# --- Test Definitions --- -# If a test fails, this function will print logs and will not cause the main script to exit. -run_test() { - local test_num=$1 - local test_name=$2 - local test_command=$3 - local log_file="$RESULTS_DIR/test_${test_num}.log" - local actual_exit_code - - echo "--- TEST_$test_num: Running $test_name ---" - - # Execute the test command. - eval "$test_command" > >(tee -a "$log_file") 2> >(tee -a "$log_file" >&2) - actual_exit_code=$? - - echo "TEST_${test_num}_COMMAND_EXIT_CODE: $actual_exit_code" # This goes to main log - echo "TEST_${test_num}_COMMAND_EXIT_CODE: $actual_exit_code" >> "$log_file" # Also to per-test log - - if [ "$actual_exit_code" -ne 0 ]; then - echo "TEST_$test_num ($test_name) FAILED with exit code $actual_exit_code." >&2 - echo "--- Log for failed TEST_$test_num ($test_name) ---" >&2 - if [ -f "$log_file" ]; then - cat "$log_file" >&2 - else - echo "Log file $log_file not found for TEST_$test_num ($test_name)." >&2 - fi - echo "--- End of log for TEST_$test_num ($test_name) ---" >&2 - return "$actual_exit_code" # Return the failure code - else - echo "TEST_$test_num ($test_name) PASSED." - return 0 # Return success - fi -} - -# Helper function to call run_test and update the overall script exit code -run_and_track_test() { - local test_num_arg="$1" - local test_name_arg="$2" - local test_command_arg="$3" - - # Run the test - run_test "$test_num_arg" "$test_name_arg" "$test_command_arg" - local test_specific_exit_code=$? - - # If the test failed, set the overall script exit code to 1 - if [ "$test_specific_exit_code" -ne 0 ]; then - # No need for extra echo here, run_test already logged the failure. - overall_script_exit_code=1 - fi -} - -# --- Actual Test Execution --- -run_and_track_test 1 "test_struct_output_generate.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/entrypoints/llm/test_struct_output_generate.py -k \"not test_structured_output_with_reasoning_matrices\"" -run_and_track_test 2 "test_moe_pallas.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/tpu/test_moe_pallas.py" -run_and_track_test 3 "test_lora.py" \ - "VLLM_XLA_CHECK_RECOMPILATION=0 python3 -m pytest -s -v /workspace/vllm/tests/tpu/lora/test_lora.py" -run_and_track_test 4 "test_tpu_qkv_linear.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_tpu_qkv_linear.py" -run_and_track_test 5 "test_spmd_model_weight_loading.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_spmd_model_weight_loading.py" -run_and_track_test 6 "test_kv_cache_update_kernel.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_kv_cache_update_kernel.py" -run_and_track_test 7 "test_tpu_int8.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_tpu_int8.py" - -# After all tests have been attempted, exit with the overall status. -if [ "$overall_script_exit_code" -ne 0 ]; then - echo "--- One or more tests FAILED. Overall script exiting with failure code 1. ---" -else - echo "--- All tests have completed and PASSED. Overall script exiting with success code 0. ---" -fi -exit "$overall_script_exit_code" -' # IMPORTANT: This is the closing single quote for the bash -c "..." command. Ensure it is present and correct. - -# Capture the exit code of the docker run command -DOCKER_RUN_EXIT_CODE=$? - -# The trap will run for cleanup. -# Exit the main script with the Docker run command's exit code. -if [ "$DOCKER_RUN_EXIT_CODE" -ne 0 ]; then - echo "Docker run command failed with exit code $DOCKER_RUN_EXIT_CODE." - exit "$DOCKER_RUN_EXIT_CODE" -else - echo "Docker run command completed successfully." - exit 0 -fi -# TODO: This test fails because it uses RANDOM_SEED sampling -# pytest -v -s /workspace/vllm/tests/tpu/test_custom_dispatcher.py \ diff --git a/.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh b/.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh deleted file mode 100755 index 6f4e0d37db7d..000000000000 --- a/.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh +++ /dev/null @@ -1,172 +0,0 @@ -#!/bin/bash - -set -xu - - -remove_docker_container() { - docker rm -f tpu-test || true; -} - -trap remove_docker_container EXIT - -# Remove the container that might not be cleaned up in the previous run. -remove_docker_container - -# Build the docker image. -docker build -f docker/Dockerfile.tpu -t vllm-tpu . - -# Set up cleanup. -cleanup_docker() { - # Get Docker's root directory - docker_root=$(docker info -f '{{.DockerRootDir}}') - if [ -z "$docker_root" ]; then - echo "Failed to determine Docker root directory." - exit 1 - fi - echo "Docker root directory: $docker_root" - # Check disk usage of the filesystem where Docker's root directory is located - disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//') - # Define the threshold - threshold=70 - if [ "$disk_usage" -gt "$threshold" ]; then - echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..." - # Remove dangling images (those that are not tagged and not used by any container) - docker image prune -f - # Remove unused volumes / force the system prune for old images as well. - docker volume prune -f && docker system prune --force --filter "until=72h" --all - echo "Docker images and volumes cleanup completed." - else - echo "Disk usage is below $threshold%. No cleanup needed." - fi -} -cleanup_docker - -# For HF_TOKEN. -source /etc/environment - -docker run --privileged --net host --shm-size=16G -it \ - -e "HF_TOKEN=$HF_TOKEN" --name tpu-test \ - vllm-tpu /bin/bash -c ' -set -e # Exit immediately if a command exits with a non-zero status. -set -u # Treat unset variables as an error. - -echo "--- Starting script inside Docker container ---" - -# Create results directory -RESULTS_DIR=$(mktemp -d) -# If mktemp fails, set -e will cause the script to exit. -echo "Results will be stored in: $RESULTS_DIR" - -# Install dependencies -echo "--- Installing Python dependencies ---" -python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \ - && python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \ - && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.12" \ - && python3 -m pip install --progress-bar off hf-transfer tblib==3.1.0 -echo "--- Python dependencies installed ---" - -export VLLM_XLA_CHECK_RECOMPILATION=1 -export VLLM_XLA_CACHE_PATH= - -echo "--- Hardware Information ---" -# tpu-info -echo "--- Starting Tests ---" -set +e -overall_script_exit_code=0 - -# --- Test Definitions --- -# If a test fails, this function will print logs and will not cause the main script to exit. -run_test() { - local test_num=$1 - local test_name=$2 - local test_command=$3 - local log_file="$RESULTS_DIR/test_${test_num}.log" - local actual_exit_code - - echo "--- TEST_$test_num: Running $test_name ---" - - # Execute the test command. - eval "$test_command" > >(tee -a "$log_file") 2> >(tee -a "$log_file" >&2) - actual_exit_code=$? - - echo "TEST_${test_num}_COMMAND_EXIT_CODE: $actual_exit_code" # This goes to main log - echo "TEST_${test_num}_COMMAND_EXIT_CODE: $actual_exit_code" >> "$log_file" # Also to per-test log - - if [ "$actual_exit_code" -ne 0 ]; then - echo "TEST_$test_num ($test_name) FAILED with exit code $actual_exit_code." >&2 - echo "--- Log for failed TEST_$test_num ($test_name) ---" >&2 - if [ -f "$log_file" ]; then - cat "$log_file" >&2 - else - echo "Log file $log_file not found for TEST_$test_num ($test_name)." >&2 - fi - echo "--- End of log for TEST_$test_num ($test_name) ---" >&2 - return "$actual_exit_code" # Return the failure code - else - echo "TEST_$test_num ($test_name) PASSED." - return 0 # Return success - fi -} - -# Helper function to call run_test and update the overall script exit code -run_and_track_test() { - local test_num_arg="$1" - local test_name_arg="$2" - local test_command_arg="$3" - - # Run the test - run_test "$test_num_arg" "$test_name_arg" "$test_command_arg" - local test_specific_exit_code=$? - - # If the test failed, set the overall script exit code to 1 - if [ "$test_specific_exit_code" -ne 0 ]; then - # No need for extra echo here, run_test already logged the failure. - overall_script_exit_code=1 - fi -} - -# --- Actual Test Execution --- -run_and_track_test 0 "test_perf.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_perf.py" -run_and_track_test 1 "test_compilation.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/tpu/test_compilation.py" -run_and_track_test 2 "test_basic.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_basic.py" -run_and_track_test 3 "test_accuracy.py::test_lm_eval_accuracy_v1_engine" \ - "python3 -m pytest -s -v /workspace/vllm/tests/entrypoints/llm/test_accuracy.py::test_lm_eval_accuracy_v1_engine" -run_and_track_test 4 "test_quantization_accuracy.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/tpu/test_quantization_accuracy.py" -run_and_track_test 6 "test_tpu_model_runner.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/worker/test_tpu_model_runner.py" -run_and_track_test 7 "test_sampler.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_sampler.py" -run_and_track_test 8 "test_topk_topp_sampler.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_topk_topp_sampler.py" -run_and_track_test 9 "test_multimodal.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_multimodal.py" -run_and_track_test 10 "test_pallas.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_pallas.py" - -# After all tests have been attempted, exit with the overall status. -if [ "$overall_script_exit_code" -ne 0 ]; then - echo "--- One or more tests FAILED. Overall script exiting with failure code 1. ---" -else - echo "--- All tests have completed and PASSED. Overall script exiting with success code 0. ---" -fi -exit "$overall_script_exit_code" -' # IMPORTANT: This is the closing single quote for the bash -c "..." command. Ensure it is present and correct. - -# Capture the exit code of the docker run command -DOCKER_RUN_EXIT_CODE=$? - -# The trap will run for cleanup. -# Exit the main script with the Docker run command's exit code. -if [ "$DOCKER_RUN_EXIT_CODE" -ne 0 ]; then - echo "Docker run command failed with exit code $DOCKER_RUN_EXIT_CODE." - exit "$DOCKER_RUN_EXIT_CODE" -else - echo "Docker run command completed successfully." - exit 0 -fi -# TODO: This test fails because it uses RANDOM_SEED sampling -# pytest -v -s /workspace/vllm/tests/tpu/test_custom_dispatcher.py \ diff --git a/.buildkite/scripts/tpu/cleanup_docker.sh b/.buildkite/scripts/tpu/cleanup_docker.sh deleted file mode 100755 index 740d81fb39bb..000000000000 --- a/.buildkite/scripts/tpu/cleanup_docker.sh +++ /dev/null @@ -1,24 +0,0 @@ -#!/bin/bash - -set -euo pipefail - -docker_root=$(docker info -f '{{.DockerRootDir}}') -if [ -z "$docker_root" ]; then - echo "Failed to determine Docker root directory." - exit 1 -fi -echo "Docker root directory: $docker_root" -# Check disk usage of the filesystem where Docker's root directory is located -disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//') -# Define the threshold -threshold=70 -if [ "$disk_usage" -gt "$threshold" ]; then - echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..." - # Remove dangling images (those that are not tagged and not used by any container) - docker image prune -f - # Remove unused volumes / force the system prune for old images as well. - docker volume prune -f && docker system prune --force --filter "until=24h" --all - echo "Docker images and volumes cleanup completed." -else - echo "Disk usage is below $threshold%. No cleanup needed." -fi diff --git a/.buildkite/scripts/tpu/config_v6e_1.env b/.buildkite/scripts/tpu/config_v6e_1.env deleted file mode 100644 index c9e3c26571e7..000000000000 --- a/.buildkite/scripts/tpu/config_v6e_1.env +++ /dev/null @@ -1,14 +0,0 @@ -# Environment config -TEST_NAME=llama8b -CONTAINER_NAME=tpu-test - -# vllm config -MODEL=meta-llama/Llama-3.1-8B-Instruct -MAX_NUM_SEQS=256 -MAX_NUM_BATCHED_TOKENS=1024 -TENSOR_PARALLEL_SIZE=1 -MAX_MODEL_LEN=2048 -DOWNLOAD_DIR=/mnt/disks/persist -EXPECTED_THROUGHPUT=8.0 -INPUT_LEN=1800 -OUTPUT_LEN=128 diff --git a/.buildkite/scripts/tpu/docker_run_bm.sh b/.buildkite/scripts/tpu/docker_run_bm.sh deleted file mode 100755 index efb632e0a854..000000000000 --- a/.buildkite/scripts/tpu/docker_run_bm.sh +++ /dev/null @@ -1,91 +0,0 @@ -#!/bin/bash - -if [ ! -f "$1" ]; then - echo "Error: The env file '$1' does not exist." - exit 1 # Exit the script with a non-zero status to indicate an error -fi - -ENV_FILE=$1 - -# For testing on local vm, use `set -a` to export all variables -source /etc/environment -# shellcheck source=/dev/null -source "$ENV_FILE" - -remove_docker_container() { - docker rm -f "$CONTAINER_NAME" || true; -} - -trap remove_docker_container EXIT - -# Remove the container that might not be cleaned up in the previous run. -remove_docker_container - -LOG_ROOT=$(mktemp -d) -# If mktemp fails, set -e will cause the script to exit. -echo "Results will be stored in: $LOG_ROOT" - -if [ -z "$HF_TOKEN" ]; then - echo "Error: HF_TOKEN is not set or is empty." - exit 1 -fi - -# Make sure mounted disk or dir exists -if [ ! -d "$DOWNLOAD_DIR" ]; then - echo "Error: Folder $DOWNLOAD_DIR does not exist. This is useually a mounted drive. If no mounted drive, just create a folder." - exit 1 -fi - -echo "Run model $MODEL" -echo - -echo "starting docker...$CONTAINER_NAME" -echo -docker run \ - -v "$DOWNLOAD_DIR":"$DOWNLOAD_DIR" \ - --env-file "$ENV_FILE" \ - -e HF_TOKEN="$HF_TOKEN" \ - -e TARGET_COMMIT="$BUILDKITE_COMMIT" \ - -e MODEL="$MODEL" \ - -e WORKSPACE=/workspace \ - --name "$CONTAINER_NAME" \ - -d \ - --privileged \ - --network host \ - -v /dev/shm:/dev/shm \ - vllm/vllm-tpu-bm tail -f /dev/null - -echo "run script..." -echo -docker exec "$CONTAINER_NAME" /bin/bash -c ".buildkite/scripts/tpu/run_bm.sh" - -echo "copy result back..." -VLLM_LOG="$LOG_ROOT/$TEST_NAME"_vllm_log.txt -BM_LOG="$LOG_ROOT/$TEST_NAME"_bm_log.txt -docker cp "$CONTAINER_NAME:/workspace/vllm_log.txt" "$VLLM_LOG" -docker cp "$CONTAINER_NAME:/workspace/bm_log.txt" "$BM_LOG" - -throughput=$(grep "Request throughput (req/s):" "$BM_LOG" | sed 's/[^0-9.]//g') -echo "throughput for $TEST_NAME at $BUILDKITE_COMMIT: $throughput" - -if [ "$BUILDKITE" = "true" ]; then - echo "Running inside Buildkite" - buildkite-agent artifact upload "$VLLM_LOG" - buildkite-agent artifact upload "$BM_LOG" -else - echo "Not running inside Buildkite" -fi - -# -# compare the throughput with EXPECTED_THROUGHPUT -# and assert meeting the expectation -# -if [[ -z "$throughput" || ! "$throughput" =~ ^[0-9]+([.][0-9]+)?$ ]]; then - echo "Failed to get the throughput" - exit 1 -fi - -if (( $(echo "$throughput < $EXPECTED_THROUGHPUT" | bc -l) )); then - echo "Error: throughput($throughput) is less than expected($EXPECTED_THROUGHPUT)" - exit 1 -fi diff --git a/.buildkite/scripts/tpu/quantized_v6e_1.env b/.buildkite/scripts/tpu/quantized_v6e_1.env deleted file mode 100644 index ecb98d4516bd..000000000000 --- a/.buildkite/scripts/tpu/quantized_v6e_1.env +++ /dev/null @@ -1,14 +0,0 @@ -# Environment config -TEST_NAME=llama8bw8a8 -CONTAINER_NAME=tpu-test - -# vllm config -MODEL=RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8 -MAX_NUM_SEQS=128 -MAX_NUM_BATCHED_TOKENS=1024 -TENSOR_PARALLEL_SIZE=1 -MAX_MODEL_LEN=2048 -DOWNLOAD_DIR=/mnt/disks/persist -EXPECTED_THROUGHPUT=8.7 -INPUT_LEN=1800 -OUTPUT_LEN=128 diff --git a/.buildkite/scripts/tpu/run_bm.sh b/.buildkite/scripts/tpu/run_bm.sh deleted file mode 100755 index b5d001bea0fe..000000000000 --- a/.buildkite/scripts/tpu/run_bm.sh +++ /dev/null @@ -1,93 +0,0 @@ -#!/bin/bash - -set -euo pipefail - -VLLM_LOG="$WORKSPACE/vllm_log.txt" -BM_LOG="$WORKSPACE/bm_log.txt" - -if [ -n "$TARGET_COMMIT" ]; then - head_hash=$(git rev-parse HEAD) - if [ "$TARGET_COMMIT" != "$head_hash" ]; then - echo "Error: target commit $TARGET_COMMIT does not match HEAD: $head_hash" - exit 1 - fi -fi - -echo "model: $MODEL" -echo - -# -# create a log folder -# -mkdir "$WORKSPACE/log" - -# TODO: Move to image building. -pip install pandas -pip install datasets - -# -# create sonnet_4x -# -echo "Create sonnet_4x.txt" -echo "" > benchmarks/sonnet_4x.txt -for _ in {1..4} - do - cat benchmarks/sonnet.txt >> benchmarks/sonnet_4x.txt -done - -# -# start vllm service in backend -# -echo "lanching vllm..." -echo "logging to $VLLM_LOG" -echo - -vllm serve "$MODEL" \ - --seed 42 \ - --max-num-seqs "$MAX_NUM_SEQS" \ - --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \ - --tensor-parallel-size "$TENSOR_PARALLEL_SIZE" \ - --no-enable-prefix-caching \ - --download_dir "$DOWNLOAD_DIR" \ - --max-model-len "$MAX_MODEL_LEN" > "$VLLM_LOG" 2>&1 & - - -echo "wait for 20 minutes.." -echo -# sleep 1200 -# wait for 10 minutes... -for _ in {1..120}; do - # TODO: detect other type of errors. - if grep -Fq "raise RuntimeError" "$VLLM_LOG"; then - echo "Detected RuntimeError, exiting." - exit 1 - elif grep -Fq "Application startup complete" "$VLLM_LOG"; then - echo "Application started" - break - else - echo "wait for 10 seconds..." - sleep 10 - fi -done - -# -# run test -# -echo "run benchmark test..." -echo "logging to $BM_LOG" -echo -vllm bench serve \ - --backend vllm \ - --model "$MODEL" \ - --dataset-name sonnet \ - --dataset-path benchmarks/sonnet_4x.txt \ - --sonnet-input-len "$INPUT_LEN" \ - --sonnet-output-len "$OUTPUT_LEN" \ - --ignore-eos > "$BM_LOG" - -echo "completed..." -echo - -throughput=$(grep "Request throughput (req/s):" "$BM_LOG" | sed 's/[^0-9.]//g') -echo "throughput: $throughput" -echo diff --git a/.github/mergify.yml b/.github/mergify.yml index 0798654d34a0..98c7d0327ab1 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -549,8 +549,6 @@ pull_request_rules: - label != stale - or: - files~=tpu.py - - files~=_tpu - - files~=tpu_ - files~=/tpu/ actions: label: @@ -564,8 +562,6 @@ pull_request_rules: - label != stale - and: - -files~=tpu.py - - -files~=_tpu - - -files~=tpu_ - -files~=/tpu/ actions: label: diff --git a/tests/v1/kv_connector/nixl_integration/run_tpu_disagg_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_tpu_disagg_accuracy_test.sh deleted file mode 100644 index 9274e3c573cb..000000000000 --- a/tests/v1/kv_connector/nixl_integration/run_tpu_disagg_accuracy_test.sh +++ /dev/null @@ -1,157 +0,0 @@ -#!/bin/bash -set -xe - -# Hosts / ports -PREFILL_HOST=${PREFILL_HOST:-"localhost"} -PREFILL_PORT=${PREFILL_PORT:-8100} -PREFILL_NIXL_SIDE_PORT=${PREFILL_NIXL_SIDE_PORT:-5577} -DECODE_HOST=${DECODE_HOST:-"localhost"} -DECODE_PORT=${DECODE_PORT:-8200} -PROXY_HOST=${PROXY_HOST:-"localhost"} -PROXY_PORT=${PROXY_PORT:-8192} -BASELINE_HOST=${BASELINE_HOST:-"localhost"} -BASELINE_PORT=${BASELINE_PORT:-9290} - - -# Model to run. -MODEL_NAME=${MODEL_NAME:-"meta-llama/Llama-3.2-3B-Instruct"} -MAX_MODEL_LEN=${MAX_MODEL_LEN:-1024} -BLOCK_SIZE=${BLOCK_SIZE:-32} - - -# execution env -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" -GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" -EXP_ROOT="${GIT_ROOT}/tests/v1/kv_connector/nixl_integration" -CONDA_PATH=${CONDA_PATH:-"/home/${USER}/anaconda3"} -CONDA_ENV_NAME=${CONDA_ENV_NAME:-"nixl"} - -OUTPUT_FILE=${OUTPUT_FILE:-"${EXP_ROOT}/.tpu_accuracy_test_outputs.txt"} - -# Trap the SIGINT signal (triggered by Ctrl+C) -trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT - - -# Waits for vLLM server to start. -wait_for_server() { - local host=$1 - local port=$2 - timeout 1200 bash -c " - until curl -s ${host}:${port}/v1/completions > /dev/null; do - sleep 1 - done" && return 0 || return 1 -} - -# Cleanup function -cleanup() { - echo "Caught Ctrl+C, cleaning up..." - # Cleanup commands - pgrep python | xargs kill -9 || true - # pkill -f python || true - echo "Cleanup complete. Exiting." -} - -launch_baseline() { - BASELINE_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME}; - VLLM_LOGGING_LEVEL=DEBUG \ - PJRT_DEVICE=TPU \ - VLLM_WORKER_MULTIPROC_METHOD=spawn \ - VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \ - --host ${BASELINE_HOST} \ - --port ${BASELINE_PORT} \ - --max-model-len ${MAX_MODEL_LEN}\ - --seed 42 \ - --block-size ${BLOCK_SIZE} \ - --gpu-memory-utilization 0.5 \ - --enforce-eager" - echo "${BASELINE_BASE_CMD}" - ssh -tt "${BASELINE_HOST}" "${BASELINE_BASE_CMD}" & -} - -launch_pd() { - PREFILL_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME}; - UCX_TLS=tcp \ - VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \ - VLLM_LOGGING_LEVEL=DEBUG \ - VLLM_NIXL_SIDE_CHANNEL_HOST=${PREFILL_HOST} \ - VLLM_NIXL_SIDE_CHANNEL_PORT=${PREFILL_NIXL_SIDE_PORT} \ - PJRT_DEVICE=TPU \ - VLLM_WORKER_MULTIPROC_METHOD=spawn \ - VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \ - --host ${PREFILL_HOST} \ - --port ${PREFILL_PORT} \ - --max-model-len ${MAX_MODEL_LEN}\ - --seed 42 \ - --block-size ${BLOCK_SIZE} \ - --enforce-eager \ - --gpu-memory-utilization 0.5 \ - --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"cpu\"}'" - - - DECODE_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME}; - UCX_TLS=tcp \ - VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \ - VLLM_LOGGING_LEVEL=DEBUG \ - PJRT_DEVICE=TPU \ - VLLM_WORKER_MULTIPROC_METHOD=spawn \ - VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \ - --host ${DECODE_HOST} \ - --port ${DECODE_PORT} \ - --max-model-len ${MAX_MODEL_LEN}\ - --seed 42 \ - --block-size ${BLOCK_SIZE} \ - --enforce-eager \ - --gpu-memory-utilization 0.5 \ - --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"cpu\"}'" - - echo "${PREFILL_BASE_CMD}" - echo "${DECODE_BASE_CMD}" - sleep 2 - - # execute on hosts - ssh -tt "${PREFILL_HOST}" "${PREFILL_BASE_CMD}" & - ssh -tt "${DECODE_HOST}" "${DECODE_BASE_CMD}" & - sleep 1 - wait_for_server "${PREFILL_HOST}" "${PREFILL_PORT}" - sleep 1 - wait_for_server "${DECODE_HOST}" "${DECODE_PORT}" - sleep 1 -} - -launch_pd_proxy(){ - PROXY_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME}; - python3 ${EXP_ROOT}/toy_proxy_server.py \ - --prefiller-host ${PREFILL_HOST} --prefiller-port ${PREFILL_PORT} \ - --decoder-host ${DECODE_HOST} --decoder-port ${DECODE_PORT} \ - --host=${PROXY_HOST} --port ${PROXY_PORT}" - echo "${PROXY_BASE_CMD}" - ssh -tt "${PROXY_HOST}" "${PROXY_BASE_CMD}" & -} - -run_tests(){ - local service_url=$1 - local mode=$2 - python3 "${EXP_ROOT}"/test_disagg_accuracy.py --service_url="${service_url}" --model_name="${MODEL_NAME}" --mode="${mode}" --file_name="${OUTPUT_FILE}" -} - - -# run non-disagg. baseline & save outputs -launch_baseline -sleep 2 -wait_for_server "${BASELINE_HOST}" "${BASELINE_PORT}" -run_tests "http://${BASELINE_HOST}:${BASELINE_PORT}" "baseline" -cleanup -sleep 10 - - -# run disagg. & do exact-match with the outputs from baseline -launch_pd -launch_pd_proxy -sleep 10 -run_tests "http://${PROXY_HOST}:${PROXY_PORT}" "disagg" -echo "-----P/D success----" - -rm "${OUTPUT_FILE}" -cleanup - -exit 0 diff --git a/tests/v1/kv_connector/nixl_integration/run_tpu_edge_case_test.sh b/tests/v1/kv_connector/nixl_integration/run_tpu_edge_case_test.sh deleted file mode 100644 index 5969455025e0..000000000000 --- a/tests/v1/kv_connector/nixl_integration/run_tpu_edge_case_test.sh +++ /dev/null @@ -1,125 +0,0 @@ -#!/bin/bash -set -xe - -# Hosts / ports -PREFILL_HOST=${PREFILL_HOST:-"localhost"} -PREFILL_PORT=${PREFILL_PORT:-8100} -PREFILL_NIXL_SIDE_PORT=${PREFILL_NIXL_SIDE_PORT:-5577} -DECODE_HOST=${DECODE_HOST:-"localhost"} -DECODE_PORT=${DECODE_PORT:-8200} -PROXY_HOST=${PROXY_HOST:-"localhost"} -PROXY_PORT=${PROXY_PORT:-8192} -BASELINE_HOST=${BASELINE_HOST:-"localhost"} -BASELINE_PORT=${BASELINE_PORT:-9290} - - -# Model to run. -MODEL_NAME=${MODEL_NAME:-"meta-llama/Llama-3.2-3B-Instruct"} -MAX_MODEL_LEN=${MAX_MODEL_LEN:-1024} -BLOCK_SIZE=${BLOCK_SIZE:-32} - - -# execution env -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" -GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" -EXP_ROOT="${GIT_ROOT}/tests/v1/kv_connector/nixl_integration" -CONDA_PATH=${CONDA_PATH:-"/home/${USER}/anaconda3"} -CONDA_ENV_NAME=${CONDA_ENV_NAME:-"nixl"} - -OUTPUT_FILE=${OUTPUT_FILE:-"${EXP_ROOT}/.tpu_accuracy_test_outputs.txt"} - -# Trap the SIGINT signal (triggered by Ctrl+C) -trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT - -# Waits for vLLM server to start. -wait_for_server() { - local host=$1 - local port=$2 - timeout 1200 bash -c " - until curl -s ${host}:${port}/v1/completions > /dev/null; do - sleep 1 - done" && return 0 || return 1 -} - -# Cleanup function -cleanup() { - echo "Caught Ctrl+C, cleaning up..." - # Cleanup commands - pgrep python | xargs kill -9 || true - # pkill -f python || true - echo "Cleanup complete. Exiting." -} - - -launch_pd() { - PREFILL_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME}; - UCX_TLS=tcp \ - VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \ - VLLM_LOGGING_LEVEL=DEBUG \ - VLLM_NIXL_SIDE_CHANNEL_HOST=${PREFILL_HOST} \ - VLLM_NIXL_SIDE_CHANNEL_PORT=${PREFILL_NIXL_SIDE_PORT} \ - PJRT_DEVICE=TPU \ - VLLM_WORKER_MULTIPROC_METHOD=spawn \ - VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \ - --host ${PREFILL_HOST} \ - --port ${PREFILL_PORT} \ - --max-model-len ${MAX_MODEL_LEN}\ - --seed 42 \ - --block-size ${BLOCK_SIZE} \ - --enforce-eager \ - --gpu-memory-utilization 0.5 \ - --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"cpu\"}'" - - - DECODE_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME}; - UCX_TLS=tcp \ - VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \ - VLLM_LOGGING_LEVEL=DEBUG \ - PJRT_DEVICE=TPU \ - VLLM_WORKER_MULTIPROC_METHOD=spawn \ - VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \ - --host ${DECODE_HOST} \ - --port ${DECODE_PORT} \ - --max-model-len ${MAX_MODEL_LEN}\ - --seed 42 \ - --block-size ${BLOCK_SIZE} \ - --enforce-eager \ - --gpu-memory-utilization 0.5 \ - --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"cpu\"}'" - - echo "${PREFILL_BASE_CMD}" - echo "${DECODE_BASE_CMD}" - sleep 2 - - # execute on hosts - ssh -tt "${PREFILL_HOST}" "${PREFILL_BASE_CMD}" & - ssh -tt "${DECODE_HOST}" "${DECODE_BASE_CMD}" & - sleep 1 - wait_for_server "${PREFILL_HOST}" "${PREFILL_PORT}" - sleep 1 - wait_for_server "${DECODE_HOST}" "${DECODE_PORT}" - sleep 1 -} - -launch_pd_proxy(){ - PROXY_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME}; - python3 ${EXP_ROOT}/toy_proxy_server.py \ - --prefiller-host ${PREFILL_HOST} --prefiller-port ${PREFILL_PORT} \ - --decoder-host ${DECODE_HOST} --decoder-port ${DECODE_PORT} \ - --host=${PROXY_HOST} --port ${PROXY_PORT}" - echo "${PROXY_BASE_CMD}" - ssh -tt "${PROXY_HOST}" "${PROXY_BASE_CMD}" & -} - - -# run disagg. & do exact-match with the outputs from baseline -launch_pd -launch_pd_proxy -sleep 10 - -PREFILL_HOST=${PREFILL_HOST} \ -PREFILL_PORT=${PREFILL_PORT} \ -DECODE_HOST=${DECODE_HOST} \ -DECODE_PORT=${DECODE_PORT} \ -PROXY_HOST=${PROXY_HOST} \ -PROXY_PORT=${PROXY_PORT} python -m pytest -s -v "${GIT_ROOT}"/tests/v1/kv_connector/nixl_integration/test_edge_cases.py From dfaabf91ac3dfa7e8e6bd0f56c7478026a533b2f Mon Sep 17 00:00:00 2001 From: Chaojun Zhang Date: Thu, 8 Oct 2026 17:48:47 +0800 Subject: [PATCH 087/624] [XPU] [CI] Fix stale ignore path for weight_transfer_metrics tests on Intel CI (#60558) Signed-off-by: Chaojun Zhang --- .buildkite/intel_jobs/entrypoints_intel.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/intel_jobs/entrypoints_intel.yaml b/.buildkite/intel_jobs/entrypoints_intel.yaml index 3cc9c1422d9e..96cbd2002a2f 100644 --- a/.buildkite/intel_jobs/entrypoints_intel.yaml +++ b/.buildkite/intel_jobs/entrypoints_intel.yaml @@ -80,7 +80,7 @@ steps: pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc && pytest -v -s entrypoints/rl --ignore=entrypoints/rl/weight_transfer --ignore=entrypoints/rl/return_routed_experts/test_openai_api.py - --ignore=entrypoints/rl/test_weight_operation_metrics.py + --ignore=entrypoints/rl/weight_transfer_metrics --deselect="tests/entrypoints/rl/state_transitions/test_pause_resume.py::TestPauseResume::test_clear_cache_preserves_output_and_controls_prefix_cache[MRV1]" && PYTHONPATH=/workspace/vllm pytest -v -s entrypoints/serve/dev/rpc' From e73895a5dfc242ed8243bacb394ab75438b591d8 Mon Sep 17 00:00:00 2001 From: Muhammad Usman Mateen <74776275+usmanmateen@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:52:50 +0100 Subject: [PATCH 088/624] [Bugfix] Make `suppress_stdout` redirect fd 1 instead of `sys.stdout.fileno()` (#59336) Signed-off-by: usmanmateen Co-authored-by: Misha Goin --- tests/utils_/test_system_utils.py | 25 ++++++++++++++++++++++++- vllm/utils/system_utils.py | 3 ++- 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/tests/utils_/test_system_utils.py b/tests/utils_/test_system_utils.py index 5ef55877a7c7..55ed35eeb56d 100644 --- a/tests/utils_/test_system_utils.py +++ b/tests/utils_/test_system_utils.py @@ -1,11 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import contextlib +import io import os +import sys import tempfile from pathlib import Path -from vllm.utils.system_utils import _maybe_force_spawn, unique_filepath +from vllm.utils.system_utils import ( + _maybe_force_spawn, + suppress_stdout, + unique_filepath, +) def test_unique_filepath(): @@ -25,3 +32,19 @@ def test_numa_bind_forces_spawn(monkeypatch): monkeypatch.setattr("sys.argv", ["vllm", "serve", "--numa-bind"]) _maybe_force_spawn() assert os.environ["VLLM_WORKER_MULTIPROC_METHOD"] == "spawn" + + +def test_suppress_stdout_when_sys_stdout_has_no_fd(capfd): + with contextlib.redirect_stdout(io.StringIO()), suppress_stdout(): + os.write(1, b"c library output\n") + assert "c library output" not in capfd.readouterr().out + + +def test_suppress_stdout_keeps_stderr_when_sys_stdout_is_stderr(capfd, monkeypatch): + monkeypatch.setattr(sys, "stdout", sys.stderr) + with suppress_stdout(): + os.write(1, b"c library output\n") + os.write(sys.stderr.fileno(), b"error message\n") + out, err = capfd.readouterr() + assert "c library output" not in out + assert "error message" in err diff --git a/vllm/utils/system_utils.py b/vllm/utils/system_utils.py index 5b90438dca49..a15dbf3c6cdf 100644 --- a/vllm/utils/system_utils.py +++ b/vllm/utils/system_utils.py @@ -72,7 +72,8 @@ def suppress_stdout(): yield return - stdout_fd = sys.stdout.fileno() + # C libraries write to fd 1, whatever sys.stdout is bound to. + stdout_fd = 1 stdout_dup = os.dup(stdout_fd) devnull_fd = os.open(os.devnull, os.O_WRONLY) From 459764104d2f8ac1b8a947f9a5a2603bdb07e360 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Thu, 8 Oct 2026 17:56:05 +0800 Subject: [PATCH 089/624] [Frontend] Move offline_utils.py to entrypoints/common. (#58052) Signed-off-by: wang.yuqi --- .github/CODEOWNERS | 2 +- tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py | 4 ++-- tests/entrypoints/unit_tests/test_offline_utils.py | 2 +- vllm/entrypoints/common/__init__.py | 0 vllm/entrypoints/{offline_utils.py => common/offline.py} | 0 vllm/entrypoints/generate/beam_search/offline.py | 2 +- vllm/entrypoints/llm.py | 2 +- vllm/entrypoints/pooling/offline.py | 2 +- vllm/entrypoints/rl/offline.py | 2 +- 9 files changed, 8 insertions(+), 8 deletions(-) create mode 100644 vllm/entrypoints/common/__init__.py rename vllm/entrypoints/{offline_utils.py => common/offline.py} (100%) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 677e37ff8b60..65d9454b283b 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -45,7 +45,7 @@ /vllm/entrypoints/serve @njhill /vllm/entrypoints/*.py @njhill /vllm/entrypoints/chat_utils.py @DarkLight1337 -/vllm/entrypoints/offline_utils.py @DarkLight1337 +/vllm/entrypoints/common @DarkLight1337 /vllm/entrypoints/llm.py @DarkLight1337 # Rust Frontend diff --git a/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py b/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py index c03efa8f10da..55990ed1dda0 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py +++ b/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py @@ -212,7 +212,7 @@ def test_preprocess_cmpl_applies_mm_processor_kwargs_to_renderer( llm.renderer = renderer monkeypatch.setattr( - "vllm.entrypoints.offline_utils.parse_model_prompt", + "vllm.entrypoints.common.offline.parse_model_prompt", lambda _model_config, parsed_prompt: parsed_prompt, ) @@ -246,7 +246,7 @@ def test_preprocess_cmpl_keeps_prompt_mm_processor_kwargs_when_no_override( llm.renderer = renderer monkeypatch.setattr( - "vllm.entrypoints.offline_utils.parse_model_prompt", + "vllm.entrypoints.common.offline.parse_model_prompt", lambda _model_config, parsed_prompt: parsed_prompt, ) diff --git a/tests/entrypoints/unit_tests/test_offline_utils.py b/tests/entrypoints/unit_tests/test_offline_utils.py index 39f77672a41f..a37969c9b366 100644 --- a/tests/entrypoints/unit_tests/test_offline_utils.py +++ b/tests/entrypoints/unit_tests/test_offline_utils.py @@ -4,7 +4,7 @@ import pytest from vllm import SamplingParams -from vllm.entrypoints.offline_utils import OfflineInferenceMixin +from vllm.entrypoints.common.offline import OfflineInferenceMixin from vllm.exceptions import VLLMValidationError from vllm.lora.request import LoRARequest diff --git a/vllm/entrypoints/common/__init__.py b/vllm/entrypoints/common/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/entrypoints/offline_utils.py b/vllm/entrypoints/common/offline.py similarity index 100% rename from vllm/entrypoints/offline_utils.py rename to vllm/entrypoints/common/offline.py diff --git a/vllm/entrypoints/generate/beam_search/offline.py b/vllm/entrypoints/generate/beam_search/offline.py index 1ad451ac13ff..0acd0b43a52f 100644 --- a/vllm/entrypoints/generate/beam_search/offline.py +++ b/vllm/entrypoints/generate/beam_search/offline.py @@ -8,7 +8,7 @@ from tqdm import tqdm from vllm import RequestOutput, TextPrompt, TokensPrompt -from vllm.entrypoints.offline_utils import OfflineInferenceMixin +from vllm.entrypoints.common.offline import OfflineInferenceMixin from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.pooling_params import PoolingParams diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 9dfff78ed9ff..a879b614844e 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -33,6 +33,7 @@ ChatTemplateContentFormatOption, load_chat_template, ) +from vllm.entrypoints.common.offline import _O, _R, OfflineInferenceMixin from vllm.entrypoints.generate.beam_search.offline import BeamSearchOfflineMixin from vllm.entrypoints.pooling.offline import PoolingOfflineMixin from vllm.entrypoints.rl.offline import RLOfflineMixin @@ -52,7 +53,6 @@ from vllm.v1.sample.logits_processor import LogitsProcessor from ..renderers import ChatParams -from .offline_utils import _O, _R, OfflineInferenceMixin if TYPE_CHECKING: from vllm.v1.metrics.reader import Metric diff --git a/vllm/entrypoints/pooling/offline.py b/vllm/entrypoints/pooling/offline.py index 27388f1ff0a9..81d78101c429 100644 --- a/vllm/entrypoints/pooling/offline.py +++ b/vllm/entrypoints/pooling/offline.py @@ -7,7 +7,7 @@ from tqdm.auto import tqdm from vllm.entrypoints.chat_utils import ChatTemplateConfig -from vllm.entrypoints.offline_utils import OfflineInferenceMixin +from vllm.entrypoints.common.offline import OfflineInferenceMixin from vllm.inputs import DataPrompt, PromptType from vllm.logger import init_logger from vllm.lora.request import LoRARequest diff --git a/vllm/entrypoints/rl/offline.py b/vllm/entrypoints/rl/offline.py index 02fc63104d8f..1db6a82533dd 100644 --- a/vllm/entrypoints/rl/offline.py +++ b/vllm/entrypoints/rl/offline.py @@ -5,7 +5,7 @@ WeightTransferInitRequest, WeightTransferUpdateRequest, ) -from vllm.entrypoints.offline_utils import OfflineInferenceMixin +from vllm.entrypoints.common.offline import OfflineInferenceMixin class RLOfflineMixin(OfflineInferenceMixin): From e161a0798c2305252e2b9bae483c537e9d0fc027 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Thu, 8 Oct 2026 12:17:11 +0200 Subject: [PATCH 090/624] [Perf][Core] Reduce free KV cache queue overhead (#60269) Signed-off-by: NickLucche --- tests/v1/core/test_kv_cache_utils.py | 18 ++++++++++ vllm/v1/core/kv_cache_utils.py | 53 +++++++++++++++++----------- 2 files changed, 50 insertions(+), 21 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 7a34d07c07c6..53833ba83711 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -966,6 +966,24 @@ def test_free_kv_cache_block_queue_popleft_n(): assert block.next_free_block is None +@pytest.mark.parametrize("num_blocks", [1, 2, 16]) +def test_free_kv_cache_block_queue_popleft_n_drain_and_refill(num_blocks): + """Draining a batch must leave the sentinels ready for reuse.""" + blocks = [KVCacheBlock(block_id=i) for i in range(num_blocks)] + queue = FreeKVCacheBlockQueue(blocks) + + assert queue.popleft_n(num_blocks) == blocks + assert queue.num_free_blocks == 0 + assert queue.fake_free_list_head.next_free_block is queue.fake_free_list_tail + assert queue.fake_free_list_tail.prev_free_block is queue.fake_free_list_head + assert all(b.prev_free_block is None and b.next_free_block is None for b in blocks) + assert queue.popleft_n(0) == [] + + queue.append_n(blocks) + assert queue.popleft_n(num_blocks) == blocks + assert queue.num_free_blocks == 0 + + def test_free_kv_cache_block_queue_get_all_free_blocks(): # Create a list of KVCacheBlock objects blocks = [KVCacheBlock(block_id=i) for i in range(5)] diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 46934d530b54..79ec2fc5e4ed 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -303,19 +303,17 @@ def popleft(self) -> KVCacheBlock: The first free block. """ - if ( - self.fake_free_list_head.next_free_block is self.fake_free_list_tail - or self.fake_free_list_head.next_free_block is None - ): + head = self.fake_free_list_head + first_block = head.next_free_block + if first_block is self.fake_free_list_tail or first_block is None: assert self.num_free_blocks == 0, ( f"num_free_blocks ({self.num_free_blocks}) is out of sync " "with the free list." ) raise ValueError("No free blocks available") - first_block: KVCacheBlock = self.fake_free_list_head.next_free_block - - if first_block.next_free_block is None: + next_block = first_block.next_free_block + if next_block is None: # This should not happen if the block is from the free list. # It indicates a bug in the caller's logic. raise RuntimeError( @@ -325,8 +323,8 @@ def popleft(self) -> KVCacheBlock: # Connect fake_head and the next block of first_block (i.e. second block # or fake tail). - self.fake_free_list_head.next_free_block = first_block.next_free_block - first_block.next_free_block.prev_free_block = self.fake_free_list_head + head.next_free_block = next_block + next_block.prev_free_block = head # Remove the block from the linked list. first_block.prev_free_block = first_block.next_free_block = None @@ -350,6 +348,16 @@ def popleft_n(self, n: int) -> list[KVCacheBlock]: self.num_free_blocks -= n curr_block = self.fake_free_list_head.next_free_block + if n == 1: + assert curr_block is not None + next_block = curr_block.next_free_block + curr_block.prev_free_block = None + curr_block.next_free_block = None + if next_block is not None: + self.fake_free_list_head.next_free_block = next_block + next_block.prev_free_block = self.fake_free_list_head + return [curr_block] + # Pop n blocks from the head of the list ret = [] for _ in range(n): @@ -375,15 +383,17 @@ def remove(self, block: KVCacheBlock) -> None: block: The block to remove. """ - if block.prev_free_block is None or block.next_free_block is None: + prev_block = block.prev_free_block + next_block = block.next_free_block + if prev_block is None or next_block is None: # This should not happen if the block is from the free list. # It indicates a bug in the caller's logic. raise RuntimeError(f"remove() called on an invalid block: {block}") # Link the previous block to the next block. - block.prev_free_block.next_free_block = block.next_free_block + prev_block.next_free_block = next_block # Link the next block to the previous block. - block.next_free_block.prev_free_block = block.prev_free_block + next_block.prev_free_block = prev_block # Remove the block from the linked list. block.prev_free_block = block.next_free_block = None @@ -397,19 +407,19 @@ def append(self, block: KVCacheBlock) -> None: block: The block to append. """ - if self.fake_free_list_tail.prev_free_block is None: + tail = self.fake_free_list_tail + last_block = tail.prev_free_block + if last_block is None: raise RuntimeError( "prev_free_block of fake_free_list_tail should always exist" ) - last_block: KVCacheBlock = self.fake_free_list_tail.prev_free_block - # Connect the new block after the last block. last_block.next_free_block = block block.prev_free_block = last_block # Connect the fake tail after the new block. - block.next_free_block = self.fake_free_list_tail - self.fake_free_list_tail.prev_free_block = block + block.next_free_block = tail + tail.prev_free_block = block self.num_free_blocks += 1 @@ -441,10 +451,11 @@ def append_n(self, blocks: list[KVCacheBlock]) -> None: blocks: The blocks to append. """ - if len(blocks) == 0: + if not blocks: return - last_block = self.fake_free_list_tail.prev_free_block + tail = self.fake_free_list_tail + last_block = tail.prev_free_block assert last_block is not None, ( "prev_free_block of fake_free_list_tail should always exist" ) @@ -455,8 +466,8 @@ def append_n(self, blocks: list[KVCacheBlock]) -> None: last_block = block # Connect the last block of to the fake tail - last_block.next_free_block = self.fake_free_list_tail - self.fake_free_list_tail.prev_free_block = last_block + last_block.next_free_block = tail + tail.prev_free_block = last_block self.num_free_blocks += len(blocks) From ba77c4c13018aa19244545cdc7ae8e016d66955e Mon Sep 17 00:00:00 2001 From: Andy Lo Date: Thu, 8 Oct 2026 11:35:22 +0100 Subject: [PATCH 091/624] [Bugfix][Attention] Preserve small FP8 softmax weights in Triton attention (#60156) Signed-off-by: Andy Lo --- .../test_triton_unified_attention.py | 47 +++++++++++++++++++ .../attention/ops/triton_unified_attention.py | 10 +++- 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/tests/kernels/attention/test_triton_unified_attention.py b/tests/kernels/attention/test_triton_unified_attention.py index e30d9225eac8..d936e377040c 100644 --- a/tests/kernels/attention/test_triton_unified_attention.py +++ b/tests/kernels/attention/test_triton_unified_attention.py @@ -181,6 +181,53 @@ def ref_paged_attn( return torch.cat(outputs, dim=0) +@torch.inference_mode() +def test_fp8_softmax_preserves_small_probabilities() -> None: + """Keep exp(-8) contributions that underflow when cast directly to E4M3.""" + device = torch.device(DEVICE_TYPE) + num_tokens = block_size = 32 + head_size = 128 + + query = torch.zeros(1, 1, head_size, dtype=FP8_DTYPE, device=device) + query[..., 0] = 1 + key_cache = torch.zeros(1, block_size, 1, head_size, dtype=FP8_DTYPE, device=device) + key_cache[:, 1:, :, 0] = -8 + value_cache = torch.ones_like(key_cache) + value_cache[:, 0] = 0 + output = torch.empty(1, 1, head_size, dtype=torch.bfloat16, device=device) + + cu_seqlens_q = torch.tensor([0, 1], dtype=torch.int32, device=device) + seqused_k = torch.tensor([num_tokens], dtype=torch.int32, device=device) + block_table = torch.tensor([[0]], dtype=torch.int32, device=device) + scale = torch.ones(1, dtype=torch.float32, device=device) + + unified_attention( + q=query, + k=key_cache, + v=value_cache, + out=output, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + seqused_k=seqused_k, + max_seqlen_k=num_tokens, + softmax_scale=1.0, + causal=True, + window_size=(-1, -1), + block_table=block_table, + softcap=0, + q_descale=scale, + k_descale=scale, + v_descale=scale, + kv_quant_mode=KVQuantMode.FP8_PER_TENSOR, + ) + + scores = torch.einsum("qhd,thd->qht", query.float(), key_cache[0].float()) + probabilities = torch.softmax(scores, dim=-1) + expected = torch.einsum("qht,thd->qhd", probabilities, value_cache[0].float()) + + torch.testing.assert_close(output.float(), expected, atol=5e-5, rtol=1e-2) + + def ref_paged_clamped_mm_attn( query: torch.Tensor, key_cache: torch.Tensor, diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index 4d73e49e91b3..63a9cdc7f956 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -586,7 +586,15 @@ def kernel_unified_attention( P_v = (P * v_token_head_scales[None, :]).to(V.dtype) acc += tl.dot(P_v, V) else: - acc += tl.dot(P.to(V.dtype), V) + # Keep softmax weights representable when the P @ V dot uses FP8 + # tensor cores. Undo the power-of-two scale after the dot. + if V.dtype == tl.float8e4nv: + p_scale = 256.0 + elif V.dtype.is_fp8(): + p_scale = 128.0 + else: + p_scale = 1.0 + acc += tl.dot((P * p_scale).to(V.dtype), V) / p_scale # ---- Epilogue --------------------------------------------------------- if IS_3D: From b49c921a27083c25c16f308be1326055dbc87bbf Mon Sep 17 00:00:00 2001 From: devshah-cohere Date: Thu, 8 Oct 2026 07:11:55 -0400 Subject: [PATCH 092/624] [Bugfix][Rust Frontend] Always include token_ids in completion stream choices (#60286) Co-authored-by: Bugen Zhao Signed-off-by: devshah-cohere --- .../server/src/routes/openai/completions.rs | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs index 7074b8c6293c..646a67d1fb8f 100644 --- a/rust/src/server/src/routes/openai/completions.rs +++ b/rust/src/server/src/routes/openai/completions.rs @@ -271,6 +271,13 @@ async fn completion_chunk_stream( macro_rules! yield_chunk { ($chunk:expr) => {{ let mut chunk = $chunk; + if return_token_ids { + for choice in &mut chunk.choices { + if choice.token_ids.is_none() { + choice.token_ids = Some(Vec::new()); + } + } + } if include_continuous_usage { chunk.usage = Some(continuous_usage.to_usage()); } @@ -787,6 +794,70 @@ mod tests { } } + #[tokio::test] + async fn completion_stream_serializes_requested_token_ids_on_every_choice() { + let stream = stream::iter(vec![ + Ok(DecodedTextEvent::Start { + prompt_token_ids: vec![1, 2].into(), + prompt_logprobs: None, + }), + Ok(decoded_delta( + "ok", + vec![7, 8], + None, + Some(Finished { + usage: vllm_llm::TokenUsage { + prompt_token_count: 2, + output_token_count: 2, + cached_token_count: 0, + }, + finish_reason: FinishReason::Length, + kv_transfer_params: None, + ec_transfer_params: None, + sampling_mask: None, + }), + )), + ]); + + let chunks = completion_chunk_stream( + stream, + "cmpl-1".to_string(), + "model".to_string(), + 1, + ApiServerOptions::default(), + ResponseOptions { + return_token_ids: true, + ..Default::default() + }, + ) + .collect::>() + .await; + + let chunks: Vec<_> = chunks.into_iter().try_collect().expect("stream should succeed"); + let payloads: Vec<_> = chunks + .iter() + .map(|chunk| serde_json::to_value(chunk).expect("chunk should serialize")) + .collect(); + + assert_eq!(payloads.len(), 3); + assert_eq!( + payloads[0]["choices"][0]["prompt_token_ids"], + serde_json::json!([1, 2]) + ); + assert_eq!( + payloads[0]["choices"][0]["token_ids"], + serde_json::json!([]) + ); + assert_eq!( + payloads[1]["choices"][0]["token_ids"], + serde_json::json!([7, 8]) + ); + assert_eq!( + payloads[2]["choices"][0]["token_ids"], + serde_json::json!([]) + ); + } + #[tokio::test] async fn collect_completion_hides_internal_prompt_only_token() { let stream = stream::iter(vec![ From 73c742b09c14f0260ccc9e9e7449444bdaf5475a Mon Sep 17 00:00:00 2001 From: Stu Cao Date: Thu, 8 Oct 2026 04:41:14 -0700 Subject: [PATCH 093/624] [Bugfix] Make disable_any_whitespace actually disable whitespace on xgrammar (#58067) Signed-off-by: Stu Cao Signed-off-by: Artem Perevedentsev Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Artem Perevedentsev --- .../test_backend_xgrammar.py | 93 +++++++++++++++++++ vllm/v1/structured_output/backend_xgrammar.py | 27 +++++- 2 files changed, 118 insertions(+), 2 deletions(-) diff --git a/tests/v1/structured_output/test_backend_xgrammar.py b/tests/v1/structured_output/test_backend_xgrammar.py index 339f0de17f03..bc29112077cc 100644 --- a/tests/v1/structured_output/test_backend_xgrammar.py +++ b/tests/v1/structured_output/test_backend_xgrammar.py @@ -2,12 +2,18 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import ClassVar + import pytest +from transformers import AutoTokenizer, PreTrainedTokenizerBase from xgrammar import Grammar from xgrammar.testing import _is_grammar_accept_string +from vllm.config import StructuredOutputsConfig, VllmConfig from vllm.sampling_params import SamplingParams, StructuredOutputsParams +from vllm.v1.structured_output.backend_types import StructuredOutputOptions from vllm.v1.structured_output.backend_xgrammar import ( + XgrammarBackend, has_xgrammar_unsupported_json_features, validate_xgrammar_grammar, ) @@ -462,3 +468,90 @@ def test_validate_xgrammar_preserves_multiline_choice(self): assert _is_grammar_accept_string(grammar, "yes") assert _is_grammar_accept_string(grammar, "no\nplease") assert not _is_grammar_accept_string(grammar, r"no\nplease") + + +# ================================================ +# Test XgrammarBackend functionality +# ================================================ + + +class TestXgrammarBackend: + class TestPR58067Regressions: + """Here we check behavior logic of our flag + disable_any_whitespace when used with xgrammar backend. + """ + + _TOKENIZER = "openai-community/gpt2" + _VOCAB_SIZE = 50257 + tokenizer: ClassVar[PreTrainedTokenizerBase] + + @classmethod + @pytest.fixture(scope="class", autouse=True) + def _shared_tokenizer(cls): + cls.tokenizer = AutoTokenizer.from_pretrained(cls._TOKENIZER) + + @classmethod + def _backend(cls, *, disable_any_whitespace: bool) -> XgrammarBackend: + vllm_config = VllmConfig( + structured_outputs_config=StructuredOutputsConfig( + backend="xgrammar", + disable_any_whitespace=disable_any_whitespace, + ) + ) + return XgrammarBackend( + vllm_config, tokenizer=cls.tokenizer, vocab_size=cls._VOCAB_SIZE + ) + + @classmethod + def _accepts_json(cls, backend: XgrammarBackend, json_text: str) -> bool: + grammar = backend.compile_grammar(StructuredOutputOptions.JSON_OBJECT, "") + return grammar.accept_tokens("req", cls.tokenizer.encode(json_text)) + + def test_disable_any_whitespace_is_false(self): + """Verify expected behavior of disable_any_whitespace=False""" + backend = self._backend(disable_any_whitespace=False) + + # Check that grammar accepts json with NO space after colon + assert self._accepts_json(backend, '{"no_space_after_me":"yes_ofcourse"}') + + # Check that grammar accepts json with space after colon + assert self._accepts_json(backend, '{"yes_space_after_me": true}') + + # Check that grammar accepts json with NO space after comma + assert self._accepts_json( + backend, + '{"never_space_after_comma":"accepted","good_boy":":3"}', + ) + + # Check that grammar accepts json with space after comma + assert self._accepts_json( + backend, + '{"where_space_after_comma":42, "wow_thats_cool":4242}', + ) + + def test_disable_any_whitespace_is_true(self): + """Verify expected behavior of disable_any_whitespace=True""" + backend = self._backend(disable_any_whitespace=True) + + # Check that grammar accepts json with NO space after colon + assert self._accepts_json(backend, '{"no_space_after_me":"yes_ofcourse"}') + + # Check that grammar rejects json with space after colon + assert not self._accepts_json(backend, '{"yes_space_after_me": true}') + + # FIXME(arpera): + # Comma + separators=(",", ":") is wrong on xgrammar <=0.2.8 + # Seehttps://github.com/mlc-ai/xgrammar/issues/945 + # Re-enable both tests once vLLM pins xgrammar with the fix (>=0.2.9). + # + # # Check that grammar accepts json with NO space after comma + # assert self._accepts_json( + # backend, + # '{"never_space_after_comma":"accepted","good_boy":":3"}', + # ) + # + # # Check that grammar rejects json with space after comma + # assert not self._accepts_json( + # backend, + # '{"where_space_after_comma":42, "wow_thats_cool":4242}', + # ) diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 733ed7db092b..b675a9a3f193 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -89,13 +89,36 @@ def compile_grammar( grammar_spec: str, stop_token_ids: set[int] | None = None, ) -> StructuredOutputGrammar: + # Note(arpera): + # Our flag disable_any_whitespace does NOT map directly to + # xgrammar's flag any_whitespace + # To achieve desired behavior of disable_any_whitespace + # we have to set not only any_whitespace + # but also specify a list of separators after which + # xgrammar must not insert spaces. + # This is a requirement of xgrammar's API, so we must comply with it. + # + # FIXME(arpera): + # Currently xgrammar v0.2.8 DOES emit spaces after comma + # even if we specify it in separators list. + # The bug has been reported to xgrammar team: + # https://github.com/mlc-ai/xgrammar/issues/945 + # Please, track that issue, and once it is resolved remove this comment. + # Upd. this bug was fixed in xgrammar main branch on Oct 8, 2026 + # and will be available in next release. + # So, remove this comment once xgrammar updates to v0.2.9 + separators = (",", ":") if self.disable_any_whitespace else None if request_type == StructuredOutputOptions.JSON: ctx = self.compiler.compile_json_schema( - grammar_spec, any_whitespace=not self.disable_any_whitespace + grammar_spec, + any_whitespace=not self.disable_any_whitespace, + separators=separators, ) elif request_type == StructuredOutputOptions.JSON_OBJECT: ctx = self.compiler.compile_json_schema( - '{"type": "object"}', any_whitespace=not self.disable_any_whitespace + '{"type": "object"}', + any_whitespace=not self.disable_any_whitespace, + separators=separators, ) elif request_type == StructuredOutputOptions.GRAMMAR: if grammar_is_likely_lark(grammar_spec): From 458ba2edf85bc9b7ebc0d5141f34ea658520faf5 Mon Sep 17 00:00:00 2001 From: Hexiang Wang <56632993+whx-sjtu@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:01:47 +0800 Subject: [PATCH 094/624] [ROCm][Kimi-K3] Fuse AttnRes output with per-token FP8 input quantization (#59069) Signed-off-by: whx-sjtu Co-authored-by: OpenAI Codex Co-authored-by: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> --- .../benchmark_kimi_k3_amd_fp8_fusions.py | 165 ++++++++++++++++++ tests/models/kimi_k3/test_amd_attn_res.py | 99 +++++++++++ .../kimi_k3/test_amd_mla_direct_return.py | 81 +++++++++ vllm/models/kimi_k3/amd/kda.py | 16 +- vllm/models/kimi_k3/amd/linear.py | 49 +++++- vllm/models/kimi_k3/amd/mla.py | 29 ++- vllm/models/kimi_k3/amd/ops/attn_res.py | 31 +++- 7 files changed, 455 insertions(+), 15 deletions(-) create mode 100644 benchmarks/kernels/benchmark_kimi_k3_amd_fp8_fusions.py diff --git a/benchmarks/kernels/benchmark_kimi_k3_amd_fp8_fusions.py b/benchmarks/kernels/benchmark_kimi_k3_amd_fp8_fusions.py new file mode 100644 index 000000000000..dd210ae55ddc --- /dev/null +++ b/benchmarks/kernels/benchmark_kimi_k3_amd_fp8_fusions.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Eight-GPU AttnRes+FP8 projection comparison with rotated graph timing.""" + +import argparse +import json +import os +import statistics +from pathlib import Path + +import torch +import torch.distributed as dist + +from vllm._aiter_ops import rocm_aiter_ops +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.models.kimi_k3.amd.ops.attn_res import attn_res +from vllm.platforms import current_platform + + +def capture(fn, repeats): + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(5): + fn() + torch.cuda.current_stream().wait_stream(stream) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(repeats): + fn() + return graph + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--tokens", type=int, nargs="+", default=[1, 4, 16, 64, 256, 1024, 4096, 16384] + ) + parser.add_argument("--blocks", type=int, nargs="+", default=[1, 4, 8]) + args = parser.parse_args() + assert int(os.environ["WORLD_SIZE"]) == 8 + assert torch.accelerator.device_count() == 8 + rank = int(os.environ["LOCAL_RANK"]) + torch.accelerator.set_device_index(rank) + torch.set_num_threads(1) + torch.manual_seed(1729 + rank) + dist.init_process_group("gloo") + args.output.mkdir(parents=True, exist_ok=True) + dtype = current_platform.fp8_dtype() + rows = [] + with set_current_vllm_config(VllmConfig()): + weight = torch.randn(2112, 7168, device="cuda", dtype=torch.bfloat16) * 0.02 + weight, ws = rocm_aiter_ops.per_token_quant(weight, dtype) + weight = rocm_aiter_ops.shuffle_weight(weight) + for m in args.tokens: + prefix = torch.randn(m, 7168, device="cuda", dtype=torch.bfloat16) + delta = torch.zeros_like(prefix) + blocks = torch.randn(m, 9, 7168, device="cuda", dtype=torch.bfloat16) + norm = torch.ones(7168, device="cuda", dtype=torch.bfloat16) + score = torch.randn_like(norm) / 7168**0.5 + for b in args.blocks: + + def produce( + fused, + prefix=prefix, + delta=delta, + blocks=blocks, + norm=norm, + score=score, + b=b, + ): + result = attn_res( + prefix, + delta, + blocks, + norm, + score, + norm, + b, + -1, + 1e-5, + 1e-5, + quant_dtype=dtype if fused else None, + ) + return ( + result + if fused + else rocm_aiter_ops.per_token_quant(result, dtype) + ) + + def segment(fused, produce=produce): + x, scale = produce(fused) + return rocm_aiter_ops.preshuffled_per_token_w8a8_gemm( + x, weight, scale, ws, output_dtype=torch.bfloat16 + ) + + expected, actual = segment(False), segment(True) + relative_l2 = ( + (actual.float() - expected.float()).norm() / expected.float().norm() + ).item() + assert relative_l2 < 0.002, relative_l2 + repeats = 20 if m < 1024 else 4 + functions = { + "baseline_producer": lambda produce=produce: produce(False), + "fused_producer": lambda produce=produce: produce(True), + "baseline_segment": lambda segment=segment: segment(False), + "fused_segment": lambda segment=segment: segment(True), + } + graphs = {k: capture(fn, repeats) for k, fn in functions.items()} + samples = {k: [] for k in functions} + names = list(functions) + for sample in range(9): + for name in names[sample % 4 :] + names[: sample % 4]: + dist.barrier() + graphs[name].replay() + torch.accelerator.synchronize() + start, end = [ + torch.cuda.Event(enable_timing=True) for _ in range(2) + ] + start.record() + graphs[name].replay() + end.record() + end.synchronize() + samples[name].append(start.elapsed_time(end) * 1000 / repeats) + row = { + "tokens": m, + "blocks": b, + "relative_l2": relative_l2, + "samples_us": samples, + "median_us": {k: statistics.median(v) for k, v in samples.items()}, + } + rows.append(row) + print( + json.dumps( + { + "rank": rank, + "tokens": m, + "blocks": b, + "median_us": row["median_us"], + } + ), + flush=True, + ) + del graphs + (args.output / f"rank-{rank}.json").write_text( + json.dumps( + { + "rank": rank, + "world_size": 8, + "torch": torch.__version__, + "rows": rows, + "scope": "AttnRes + FP8 quant + TP8 MLA input projection", + }, + indent=2, + ) + + "\n" + ) + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/models/kimi_k3/test_amd_attn_res.py b/tests/models/kimi_k3/test_amd_attn_res.py index 80d5e1936f4f..f4efdbd5b21d 100644 --- a/tests/models/kimi_k3/test_amd_attn_res.py +++ b/tests/models/kimi_k3/test_amd_attn_res.py @@ -191,3 +191,102 @@ def test_amd_attn_res_fused_contract( original_blocks[:, block_write_idx].copy_(expected_prefix) torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0) assert actual.is_contiguous() + + +@pytest.mark.parametrize( + "num_tokens,num_blocks,has_delta,write_block", + [ + (0, 0, False, False), + (1, 0, True, True), + (17, 4, True, True), + (320, 8, False, False), + ], +) +def test_amd_attn_res_fp8_preserves_prefix_and_quantized_output( + num_tokens, + num_blocks, + has_delta, + write_block, + default_vllm_config, +): + from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 + from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape + + torch.manual_seed(41) + hidden_size = 7168 + prefix = _randn_with_row_padding(num_tokens, hidden_size, padding=7) + delta = _randn_with_row_padding(num_tokens, hidden_size) if has_delta else None + blocks = _randn_with_row_padding(num_tokens, 9, hidden_size, padding=13) + norm = 1 + torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16) * 0.1 + score = torch.randn_like(norm) / hidden_size**0.5 + out_norm = 1 + torch.randn_like(norm) * 0.1 + ref_prefix, ref_blocks = prefix.clone(), blocks.clone() + kwargs = dict( + num_blocks=num_blocks, + block_write_idx=num_blocks if write_block else -1, + eps=1e-5, + output_norm_eps=1e-5, + ) + reference = attn_res(ref_prefix, delta, ref_blocks, norm, score, out_norm, **kwargs) + output, scale = attn_res( + prefix, + delta, + blocks, + norm, + score, + out_norm, + quant_dtype=current_platform.fp8_dtype(), + **kwargs, + ) + assert output.dtype == current_platform.fp8_dtype() + assert scale.shape == (num_tokens, 1) + if num_tokens: + # Keep reference rounding independent of Inductor's fused quantization. + quant = QuantFP8( + static=False, group_shape=GroupShape.PER_TOKEN, compile_native=False + ) + ref_output, ref_scale = quant(reference) + torch.testing.assert_close(scale, ref_scale, atol=1e-7, rtol=1e-6) + # Floating-point fusion can move values at FP8 rounding midpoints. + # Bound each change to the adjacent code and bound how often it occurs. + code_delta = ( + output.view(torch.uint8).to(torch.int16) + - ref_output.view(torch.uint8).to(torch.int16) + ).abs() + assert code_delta.max().item() <= 1 + assert (code_delta != 0).float().mean().item() < 1e-4 + torch.testing.assert_close(prefix, ref_prefix, atol=0, rtol=0) + torch.testing.assert_close(blocks, ref_blocks, atol=0, rtol=0) + + +@pytest.mark.parametrize( + "matching,requires_unquantized", [(True, False), (False, False), (True, True)] +) +def test_mla_input_quant_requires_all_consumers(matching, requires_unquantized): + from types import SimpleNamespace + + from torch import nn + + from vllm.model_executor.layers.fusion.quant_activation import ( + expose_input_quant_key, + ) + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8DynamicTokenSym, + ) + from vllm.models.kimi_k3.amd.mla import KimiK3MultiHeadLatentAttentionWrapper + + attention = object.__new__(KimiK3MultiHeadLatentAttentionWrapper) + nn.Module.__init__(attention) + attention.indexer = None + attention.q_lora_rank = 1536 + attention.fused_qkv_a_proj = nn.Identity() + attention.g_proj = nn.Identity() + kernel = SimpleNamespace(input_quant_key=lambda: kFp8DynamicTokenSym) + expose_input_quant_key(attention.fused_qkv_a_proj, kernel) + if matching: + expose_input_quant_key(attention.g_proj, kernel) + attention.g_proj.requires_unquantized_input = requires_unquantized + actual = attention.get_input_quant_key() + assert actual == ( + kFp8DynamicTokenSym if matching and not requires_unquantized else None + ) diff --git a/tests/models/kimi_k3/test_amd_mla_direct_return.py b/tests/models/kimi_k3/test_amd_mla_direct_return.py index a1c120737df8..50a52f4fe148 100644 --- a/tests/models/kimi_k3/test_amd_mla_direct_return.py +++ b/tests/models/kimi_k3/test_amd_mla_direct_return.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import pytest import torch from torch import nn @@ -36,3 +39,81 @@ def test_mla_self_attention_returns_projection_storage_directly() -> None: assert output is projected assert output.data_ptr() == projected.data_ptr() + + +@pytest.mark.parametrize("quantized", [False, True]) +@pytest.mark.parametrize("fused_decode", [False, True]) +def test_mla_forward_preserves_output_shape_with_fp8_input( + monkeypatch, quantized: bool, fused_decode: bool +) -> None: + """Both MLA paths accept producer-quantized input after decode fusion.""" + from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8DynamicTokenSym, + ) + from vllm.models.kimi_k3.amd.mla import KimiK3MultiHeadLatentAttentionWrapper + + hidden_states = torch.randn(4, 8, dtype=torch.bfloat16) + if quantized: + hidden_states = QuantizedActivation( + hidden_states.to(kFp8DynamicTokenSym.dtype), + torch.ones(4, 1), + hidden_states.dtype, + hidden_states.shape, + kFp8DynamicTokenSym, + ) + projected = torch.randn(4, 8, dtype=torch.bfloat16) + paths = [] + + def input_projection(value): + assert value is hidden_states + return torch.zeros(4, 8, dtype=torch.bfloat16), None + + def attention_output(path, output_shape): + assert output_shape == projected.shape + paths.append(path) + return projected + + attention = SimpleNamespace( + q_lora_rank=2, + kv_lora_rank=4, + qk_rope_head_dim=2, + qk_head_dim=4, + num_heads=2, + v_head_dim=4, + fused_qkv_a_proj=input_projection, + q_a_layernorm=nn.Identity(), + _normalize_q_kv=lambda q, kv: (q, kv), + q_b_proj=lambda q: (torch.zeros(4, 8, dtype=q.dtype), None), + rotary_emb=None, + indexer=None, + dcp_q_replicate=False, + _fused_qk_prep=fused_decode, + _fused_decode=lambda q, kv, pe, pos, shape, *args: attention_output( + "fused", shape + ), + g_proj=None, + o_proj=lambda value: (value, None), + ) + if fused_decode: + attention.mla_attn = SimpleNamespace(layer_name="test") + monkeypatch.setattr( + "vllm.models.kimi_k3.amd.mla.get_attention_context", + lambda _: ( + SimpleNamespace(num_actual_tokens=4, num_decode_tokens=4), + None, + torch.zeros(1), + torch.arange(4), + ), + ) + else: + attention.mla_attn = lambda q, kv, pe, **kwargs: attention_output( + "regular", kwargs["output_shape"] + ) + + output = KimiK3MultiHeadLatentAttentionWrapper.forward( + attention, torch.arange(4), hidden_states + ) + + assert output is projected + assert paths == ["fused" if fused_decode else "regular"] diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py index e2dea74b6b9e..409cf6b1bf90 100644 --- a/vllm/models/kimi_k3/amd/kda.py +++ b/vllm/models/kimi_k3/amd/kda.py @@ -13,6 +13,10 @@ from vllm.distributed import divide from vllm.forward_context import get_forward_context from vllm.logger import init_logger +from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + get_input_quant_key, +) from vllm.model_executor.layers.linear import ( ColumnParallelLinear, RowParallelLinear, @@ -37,6 +41,7 @@ causal_conv1d_fn, causal_conv1d_update, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader from vllm.model_executor.utils import set_weight_attrs from vllm.models.kimi_k3.amd.kda_metadata import KimiK3ROCmKDABackend @@ -300,13 +305,16 @@ def rearrange_mixed_qkv( qkv = qkv.permute(1, 0, 2, 3).contiguous().unsqueeze(1) return qkv.unbind(0) + def get_input_quant_key(self) -> QuantKey | None: + return get_input_quant_key(self.in_proj_qkvgfab) + def forward( self, - hidden_states: torch.Tensor, + hidden_states: torch.Tensor | QuantizedActivation, positions: torch.Tensor, ) -> torch.Tensor: - num_tokens = hidden_states.size(0) projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0] + num_tokens = projected_qkvgfab.shape[0] split_sizes = [ 3 * self.local_projection_size, @@ -326,8 +334,8 @@ def forward( core_attn_out = torch.empty( (1, num_tokens, self.local_num_heads, self.head_dim), - dtype=hidden_states.dtype, - device=hidden_states.device, + dtype=projected_qkvgfab.dtype, + device=projected_qkvgfab.device, ) self._forward( diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py index 65dad661ab12..3a1d827abe47 100644 --- a/vllm/models/kimi_k3/amd/linear.py +++ b/vllm/models/kimi_k3/amd/linear.py @@ -2,11 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable -from typing import Any +from typing import Any, cast import torch from torch import nn +from vllm import envs from vllm.config import CacheConfig, VllmConfig from vllm.distributed import ( get_pp_group, @@ -19,6 +20,10 @@ fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear +from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + get_input_quant_key, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -38,6 +43,10 @@ ) from vllm.model_executor.layers.mla import MLAModules from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kFp8DynamicTokenSym, +) from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, @@ -114,6 +123,9 @@ def __init__( "Only silu and situ are supported." ) + def get_input_quant_key(self) -> QuantKey | None: + return get_input_quant_key(self.gate_up_proj) + def forward(self, x): gate_up, _ = self.gate_up_proj(x) x = self.act_fn(gate_up) @@ -148,8 +160,10 @@ def _apply_attn_res( delta: torch.Tensor | None = None, output_norm: RMSNorm | None = None, block_write_idx: int = -1, -) -> torch.Tensor: - return attn_res( + quant_key: QuantKey | None = None, +) -> torch.Tensor | QuantizedActivation: + fuse_quant = quant_key == kFp8DynamicTokenSym and not envs.VLLM_BATCH_INVARIANT + result = attn_res( prefix_sum, delta, block_residual, @@ -160,7 +174,18 @@ def _apply_attn_res( block_write_idx, norm.variance_epsilon, 0.0 if output_norm is None else output_norm.variance_epsilon, + quant_dtype=cast(torch.dtype, kFp8DynamicTokenSym.dtype) + if fuse_quant + else None, ) + if fuse_quant: + assert isinstance(result, tuple) + data, scale = result + return QuantizedActivation( + data, scale, prefix_sum.dtype, prefix_sum.shape, kFp8DynamicTokenSym + ) + assert isinstance(result, torch.Tensor) + return result class KimiMoE(nn.Module): @@ -452,10 +477,13 @@ def __init__( prefix, ) + def get_input_quant_key(self) -> QuantKey | None: + return self.mla_attn.get_input_quant_key() + def forward( self, positions: torch.Tensor, - hidden_states: torch.Tensor, + hidden_states: torch.Tensor | QuantizedActivation, ) -> torch.Tensor: return self.mla_attn(positions, hidden_states) @@ -579,7 +607,7 @@ def __init__( def _run_self_attn( self, positions: torch.Tensor, - hidden_states: torch.Tensor, + hidden_states: torch.Tensor | QuantizedActivation, ) -> torch.Tensor: return self.self_attn( hidden_states=hidden_states, @@ -626,6 +654,11 @@ def forward_attn_residual( prefix_delta: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: prefix_sum = hidden_states + attention_quant_key = ( + self.self_attn.get_input_quant_key() + if isinstance(self.self_attn, (KimiK3DeltaAttention, KimiMLAAttention)) + else None + ) hidden_states = _apply_attn_res( prefix_sum, block_residual, @@ -635,6 +668,7 @@ def forward_attn_residual( delta=prefix_delta, output_norm=self.input_layernorm, block_write_idx=(self.block_write_idx if self.is_block_write_layer else -1), + quant_key=attention_quant_key, ) if self.is_block_write_layer: @@ -659,6 +693,11 @@ def forward_attn_residual( mlp_valid_blocks, delta=prefix_delta, output_norm=self.post_attention_layernorm, + quant_key=( + self.mlp.get_input_quant_key() + if isinstance(self.mlp, KimiMLP) + else None + ), ) hidden_states = self.mlp(hidden_states) diff --git a/vllm/models/kimi_k3/amd/mla.py b/vllm/models/kimi_k3/amd/mla.py index 11df7d035f31..1e4a51e8bf50 100644 --- a/vllm/models/kimi_k3/amd/mla.py +++ b/vllm/models/kimi_k3/amd/mla.py @@ -8,8 +8,13 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.model_executor.layers.attention.attention import get_attention_context +from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + get_input_quant_key, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.mla import MultiHeadLatentAttentionWrapper +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey from vllm.platforms import current_platform _OPT_KV_LORA_RANK = 512 @@ -49,6 +54,21 @@ def _fused_qk_prep_supported(self) -> bool: not attn.is_aiter_triton_fp8_bmm_enabled and attn.W_UK_T is not None ) + def get_input_quant_key(self) -> QuantKey | None: + if self.indexer is not None: + return None + consumers = ( + [self.fused_qkv_a_proj] + if self.q_lora_rank is not None + else [self.kv_a_proj_with_mqa, self.q_proj] + ) + if self.g_proj is not None: + consumers.append(self.g_proj) + if any(layer is None for layer in consumers): + return None + keys = [get_input_quant_key(layer) for layer in consumers if layer is not None] + return keys[0] if all(key == keys[0] for key in keys) else None + def _normalize_q_kv( self, q_c: torch.Tensor, @@ -173,10 +193,15 @@ def _fused_decode( def forward( self, positions: torch.Tensor, - hidden_states: torch.Tensor, + hidden_states: torch.Tensor | QuantizedActivation, llama_4_scaling: torch.Tensor | None = None, ) -> torch.Tensor: q_c = None + num_tokens = ( + hidden_states.orig_shape[0] + if isinstance(hidden_states, QuantizedActivation) + else hidden_states.shape[0] + ) if self.q_lora_rank is not None: assert self.fused_qkv_a_proj is not None, ( @@ -238,7 +263,7 @@ def forward( if self.dcp_q_replicate: q_dcp_replicated, q = q, q_proj_layer._local_view(q) - output_shape = (hidden_states.shape[0], self.num_heads * self.v_head_dim) + output_shape = (num_tokens, self.num_heads * self.v_head_dim) # The fused decode path covers only the MQA slice, so it is taken for # decode-only batches; anything else falls through to MLAAttention. diff --git a/vllm/models/kimi_k3/amd/ops/attn_res.py b/vllm/models/kimi_k3/amd/ops/attn_res.py index 5802e603670a..5c4fe7a0b0b2 100644 --- a/vllm/models/kimi_k3/amd/ops/attn_res.py +++ b/vllm/models/kimi_k3/amd/ops/attn_res.py @@ -22,6 +22,7 @@ def _attn_res_kernel( qk_weight_ptr, output_norm_weight_ptr, output_ptr, + output_scale_ptr, stride_prefix_m: tl.constexpr, stride_delta_m: tl.constexpr, stride_block_m: tl.constexpr, @@ -35,6 +36,7 @@ def _attn_res_kernel( HAS_DELTA: tl.constexpr, WRITE_BLOCK: tl.constexpr, APPLY_OUTPUT_NORM: tl.constexpr, + QUANT_MAX: tl.constexpr, BLOCK_L: tl.constexpr, BLOCK_D: tl.constexpr, ): @@ -129,6 +131,14 @@ def _attn_res_kernel( output_norm_weight_ptr + d_offsets, mask=d_mask, other=0.0 ).to(tl.float32) output = mixed * output_reciprocal_std * output_norm_weight + if QUANT_MAX > 0: + # Preserve the rounding of the original BF16 output before quantizing. + output = output.to(prefix_ptr.dtype.element_ty).to(tl.float32) + amax = tl.max(tl.where(d_mask, tl.abs(output), 0.0), axis=0) + scale = tl.maximum(amax, 1e-10) * (1.0 / QUANT_MAX) + inv_scale = 1.0 / scale + output = tl.minimum(tl.maximum(output * inv_scale, -QUANT_MAX), QUANT_MAX) + tl.store(output_scale_ptr + row_idx, scale) tl.store( output_ptr + row_idx * stride_output_m + d_offsets, output, @@ -147,7 +157,9 @@ def attn_res( block_write_idx: int, eps: float, output_norm_eps: float, -) -> torch.Tensor: + *, + quant_dtype: torch.dtype | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: num_tokens, hidden_size = prefix.shape assert prefix.stride(-1) == 1 assert delta is None or delta.stride(-1) == 1 @@ -155,9 +167,18 @@ def attn_res( assert norm_weight.stride(-1) == 1 assert qk_weight.stride(-1) == 1 assert output_norm_weight is None or output_norm_weight.stride(-1) == 1 - output = prefix.new_empty(prefix.shape) + if quant_dtype is not None: + assert quant_dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz) + output = torch.empty_like( + prefix, dtype=quant_dtype or prefix.dtype, memory_format=torch.contiguous_format + ) + scale = ( + torch.empty((num_tokens, 1), device=prefix.device, dtype=torch.float32) + if quant_dtype is not None + else None + ) if num_tokens == 0: - return output + return output if scale is None else (output, scale) # Source tiling helps decode, while one-source tiles scale better for prefill. if num_tokens >= 256 or num_blocks <= 1: @@ -172,6 +193,7 @@ def attn_res( qk_weight, output_norm_weight, output, + scale, prefix.stride(0), 0 if delta is None else delta.stride(0), blocks.stride(0), @@ -185,9 +207,10 @@ def attn_res( HAS_DELTA=delta is not None, WRITE_BLOCK=block_write_idx >= 0, APPLY_OUTPUT_NORM=output_norm_weight is not None, + QUANT_MAX=0.0 if quant_dtype is None else torch.finfo(quant_dtype).max, BLOCK_L=block_l, BLOCK_D=triton.next_power_of_2(hidden_size), num_warps=num_warps, num_stages=2, ) - return output + return output if scale is None else (output, scale) From 194da61de1769c28bfad2b6642c60cc7eca544dd Mon Sep 17 00:00:00 2001 From: Vineeth Sai Varikuntla Date: Thu, 8 Oct 2026 06:01:25 -0700 Subject: [PATCH 095/624] [Bugfix] Reject negative --device-ids indices instead of selecting the wrong device (#49053) Signed-off-by: Vineeth Sai Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- tests/engine/test_arg_utils.py | 5 +++-- vllm/engine/arg_utils.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 924d2bafec77..fc9f0229c46c 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -857,13 +857,14 @@ def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch): class TestDeviceIds: - def test_device_ids_with_cvd_out_of_range(self, monkeypatch): + @pytest.mark.parametrize("device_ids", [[0, 2], [-1]]) + def test_device_ids_with_cvd_out_of_range(self, monkeypatch, device_ids): """--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]) + args = EngineArgs(model="m", device_ids=device_ids) with pytest.raises(ValueError, match="out of range"): args._resolve_device_ids() diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 800d1b2f2ccf..94dd3e2ef1c8 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -2113,7 +2113,7 @@ def _resolve_device_ids(self) -> list[int] | None: for x in cvd.split(",") ] for i in int_ids: - if i >= len(cvd_ids): + if not 0 <= i < len(cvd_ids): raise ValueError( f"--device-ids index {i} is out of range for " f"{current_platform.device_control_env_var}" From 4ed4c4433b87e45d72f401f27f1b56c932c90071 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Thu, 8 Oct 2026 10:00:49 -0400 Subject: [PATCH 096/624] [Bugfix][HiSparse] Drop the unaccounted, write-only prefill mirror staging buffer (#60429) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- .../v1/attention/test_sparse_mla_backends.py | 12 +--- tests/v1/worker/test_attn_utils.py | 2 +- .../kv_connector/v1/hisparse/worker.py | 1 - .../layers/attention/mla_attention.py | 12 ---- vllm/v1/hisparse/binding.py | 4 -- vllm/v1/hisparse/runtime.py | 55 +------------------ 6 files changed, 5 insertions(+), 81 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 3d32aed5d6e7..4d33075c9303 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -2578,7 +2578,7 @@ def test_hisparse_multi_step_writes_request_major_output(): @requires_hisparse_ops -def test_hisparse_kv_update_writes_resident_and_staging_caches(): +def test_hisparse_kv_update_writes_resident_cache(): device = torch.device(DEVICE_TYPE) block_size = 4 row_width = 8 @@ -2597,18 +2597,11 @@ def test_hisparse_kv_update_writes_resident_and_staging_caches(): block_table=torch.tensor([[1]], dtype=torch.int32, device=device), slot_mapping=resident_slots, ) - cache_handle.mirror_staging_cache = torch.empty( - (1, block_size, row_width), dtype=torch.float32, device=device - ) - cache_handle.mirror_staging_slots = torch.arange( - block_size, dtype=torch.int64, device=device - ) slots = torch.tensor([3, 7, -1], dtype=torch.int64, device=device) kv_c = torch.randn(8, row_width - 2, device=device) k_pe = torch.randn(8, 1, 2, device=device) cache_handle.num_actual_tokens = slots.numel() cache_handle.decode_batch = False - cache_handle.host_mirror_required = True source_cache = torch.zeros_like(cache_handle.view.cache) impl = object.__new__(FlashMLASparseImpl) layer = SimpleNamespace( @@ -2635,9 +2628,6 @@ def test_hisparse_kv_update_writes_resident_and_staging_caches(): expected.to(device), ) torch.testing.assert_close(source_cache, torch.zeros_like(source_cache)) - staged = cache_handle.mirror_staging_cache.view(-1, row_width) - staged_expected = torch.cat([kv_c[:3], k_pe[:3, 0]], dim=-1) - torch.testing.assert_close(staged[:3], staged_expected) @requires_hisparse_ops diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 41148070528a..9c9c26c3c2c0 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -434,7 +434,7 @@ def test_init_hisparse_rolls_back_shared_region(monkeypatch, failure_phase): cache_config=SimpleNamespace( get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC ), - scheduler_config=SimpleNamespace(max_num_seqs=1, max_num_batched_tokens=1), + scheduler_config=SimpleNamespace(max_num_seqs=1), ) def allocate(*args): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py index bf6329b3bfdf..b9dd01f05420 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py @@ -417,7 +417,6 @@ def start_step( self._clear_forward_mirror_state() for handle in self.cache_handles: handle.all_context_pages_resident = metadata.all_context_pages_resident - handle.mirror_from_resident = True self._copy_host_blocks(metadata.host_block_copies, previous_host_write_event) transfers = ( metadata.command.page_transfers if metadata.command is not None else [] diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 4286502ea9fe..2a3048096c98 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -802,18 +802,6 @@ def update_kv_cache( kv_cache_dtype, k_scale, ) - if cache is not None: - mirror_target = cache.mirror_write_target(kv_c_normed.shape[0]) - if mirror_target is not None: - mirror_cache, mirror_slots = mirror_target - self.impl.do_kv_cache_update( # type: ignore[attr-defined] - kv_c_normed, - k_pe, - mirror_cache, - mirror_slots, - kv_cache_dtype, - k_scale, - ) def prepare_kv_cache_update( self, attn_metadata: "MLACommonMetadata | None" diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index d021d8b7be09..b8c0d5dff269 100644 --- a/vllm/v1/hisparse/binding.py +++ b/vllm/v1/hisparse/binding.py @@ -147,7 +147,6 @@ def init_hisparse_kv_cache( initialize_hisparse_runtime_buffers( cache_handles, max_num_reqs=vllm_config.scheduler_config.max_num_seqs, - max_num_batched_tokens=vllm_config.scheduler_config.max_num_batched_tokens, ) return kv_caches except Exception: @@ -197,9 +196,6 @@ def release_hisparse_profiling_cache(forward_context: dict[str, Any]) -> None: ) ) release_pinned_state(list(runtimes.values()), registered_pools, shared_region) - for cache in cache_handles: - cache.mirror_staging_cache = None - cache.mirror_staging_slots = None def bind_hisparse_kv_caches( diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index fc415fa2eccf..a6378ca9141a 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -1120,10 +1120,7 @@ def __init__(self, runtime: HiSparseRuntime) -> None: self.num_decode_tokens = 0 self.req_id_per_token: torch.Tensor | None = None self.host_mirror_required = False - self.mirror_from_resident = False self.mirror_slot_mapping: torch.Tensor | None = None - self.mirror_staging_cache: torch.Tensor | None = None - self.mirror_staging_slots: torch.Tensor | None = None self.submit_layer_mirror: Callable[[], None] | None = None # Speculator layers write their rows after the target forward. self.draft_layer = False @@ -1160,26 +1157,6 @@ def write_target( num_rows = min(num_rows, self.runtime.max_num_reqs) return self.view.cache, self.slot_mapping[:num_rows], num_rows - def mirror_write_target( - self, num_rows: int - ) -> tuple[torch.Tensor, torch.Tensor] | None: - if ( - not self.host_mirror_required - or self.decode_batch - or self.mirror_from_resident - ): - return None - cache = self.mirror_staging_cache - slots = self.mirror_staging_slots - if cache is None or slots is None: - raise RuntimeError("HiSparse prefill mirror staging is not bound.") - if num_rows > slots.numel(): - raise RuntimeError( - "HiSparse prefill mirror exceeds staging capacity: " - f"{num_rows} > {slots.numel()}." - ) - return cache, slots[:num_rows] - def finish_kv_update(self) -> None: if self.dummy_batch: return @@ -1242,9 +1219,8 @@ def initialize_hisparse_runtime_buffers( cache_handles: list[HiSparseCacheHandle], *, max_num_reqs: int, - max_num_batched_tokens: int, ) -> None: - """Allocate shared request state and per-layer prefill staging buffers.""" + """Allocate the request state shared by every HiSparse runtime.""" assert cache_handles resident = cache_handles[0].view assert resident is not None @@ -1252,26 +1228,8 @@ def initialize_hisparse_runtime_buffers( request_state_indices = torch.full( (max_num_reqs,), -1, dtype=torch.int32, device=device ) - staging_blocks = ( - max_num_batched_tokens + resident.block_size - 1 - ) // resident.block_size - mirror_staging_caches = torch.empty( - ( - len(cache_handles), - staging_blocks, - resident.block_size, - resident.cache.shape[-1], - ), - dtype=resident.cache.dtype, - device=device, - ) - mirror_staging_slots = torch.arange( - max_num_batched_tokens, dtype=torch.int64, device=device - ) - for layer_index, cache_handle in enumerate(cache_handles): + for cache_handle in cache_handles: cache_handle.runtime.request_state_indices = request_state_indices - cache_handle.mirror_staging_cache = mirror_staging_caches[layer_index] - cache_handle.mirror_staging_slots = mirror_staging_slots def create_hisparse_cache_handle( @@ -1325,11 +1283,4 @@ def create_hisparse_cache_handle( config.device_buffer_size, max_num_reqs, ) - handle = HiSparseCacheHandle(runtime) - speculative_config = vllm_config.speculative_config - handle.mirror_from_resident = bool( - vllm_config.scheduler_config.async_scheduling - and speculative_config is not None - and speculative_config.uses_draft_kv_cache() - ) - return handle + return HiSparseCacheHandle(runtime) From daa9085143d65896ba3fa77635b0b110aab3e585 Mon Sep 17 00:00:00 2001 From: Pravein Govindan Kannan Date: Thu, 8 Oct 2026 19:55:02 +0530 Subject: [PATCH 097/624] [KV Connector][NIXL] Register HiSparse host pool under single TP Rank (#60161) Signed-off-by: Pravein Govindan Kannan Signed-off-by: Matthew Bonanni Co-authored-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- .../kv_connector/unit/test_nixl_connector.py | 62 ++++++++++++++++ .../unit/test_nixl_connector_hma.py | 1 + .../kv_connector/v1/nixl/base_worker.py | 72 ++++++++++++++----- .../kv_connector/v1/nixl/pull_worker.py | 11 ++- 4 files changed, 128 insertions(+), 18 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 18097d220e79..52ac746391ff 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -2083,6 +2083,7 @@ def test_mixed_memory_local_descriptors_split_by_memory_type(): worker.kv_caches_base_addr = {"local": {0: [100, 200]}} worker._has_mamba = False worker._mixed_mem_types = True + worker._skip_dram_xfer = False worker.region_mem_types = ["DRAM", "VRAM"] worker.region_num_blocks = [2, 2] worker._desc_is_dram_by_block_size = {} @@ -2150,6 +2151,7 @@ def recv_worker(): worker.enable_permute_local_kv = False worker.enable_heterogeneous_attn_post_process = False worker.tp_rank = 0 + worker._skip_dram_xfer = False worker._log_failure = MagicMock() # type: ignore[method-assign] worker.xfer_stats = NixlKVConnectorStats() worker.nixl_wrapper = MagicMock() @@ -2207,6 +2209,66 @@ def test_mixed_memory_read_failure_does_not_notify_producer(recv_worker): worker.nixl_wrapper.send_notif.assert_not_called() +def test_hisparse_dram_rank0_only_read_keeps_device_part(recv_worker): + worker = recv_worker + worker.tp_rank = 1 + worker._skip_dram_xfer = True + worker._desc_is_dram_by_block_size = {16: np.array([True, True, False, False])} + worker._desc_pos_by_block_size = {16: np.array([0, 1, 0, 1])} + worker._dram_src_handles_by_block_size = {16: None} + worker.nixl_wrapper.make_prepped_xfer.side_effect = [102] + worker.nixl_wrapper.check_xfer_state.return_value = "DONE" + + worker._read_blocks_mixed( + request_id="request", + local_block_size_key=16, + local_device_handle=20, + local_dram_handle=None, + remote_xfer_side_handle=30, + local_block_descs_ids=np.array([0, 2]), + remote_block_descs_ids=np.array([5, 7]), + notif_agent="prefill", + notif_id=b"request:1", + ) + + (device_read,) = worker.nixl_wrapper.make_prepped_xfer.call_args_list + assert device_read.args[:2] == ("READ", 20) + np.testing.assert_array_equal(device_read.args[2], [0]) + np.testing.assert_array_equal(device_read.args[4], [7]) + + assert worker.get_finished() == (set(), {"request"}) + worker.nixl_wrapper.send_notif.assert_called_once_with( + "prefill", notif_msg=b"request:1" + ) + + +def test_hisparse_dram_rank0_only_dram_only_read_completes(recv_worker): + """With nothing left to read, the request still completes and notifies.""" + worker = recv_worker + worker.tp_rank = 1 + worker._skip_dram_xfer = True + worker._desc_is_dram_by_block_size = {16: np.array([True, True, False, False])} + worker._desc_pos_by_block_size = {16: np.array([0, 1, 0, 1])} + + worker._read_blocks_mixed( + request_id="request", + local_block_size_key=16, + local_device_handle=20, + local_dram_handle=None, + remote_xfer_side_handle=30, + local_block_descs_ids=np.array([0, 1]), + remote_block_descs_ids=np.array([5, 6]), + notif_agent="prefill", + notif_id=b"request:1", + ) + + worker.nixl_wrapper.make_prepped_xfer.assert_not_called() + assert worker.get_finished() == (set(), {"request"}) + worker.nixl_wrapper.send_notif.assert_called_once_with( + "prefill", notif_msg=b"request:1" + ) + + def element_byte_addrs(view: torch.Tensor) -> list[int]: """Absolute byte addresses of every element of a (possibly strided) view.""" offsets = torch.zeros(view.shape, dtype=torch.int64) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index b4b5e063445f..f016b26475f0 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -267,6 +267,7 @@ def test_region_pull_completion_zeros_only_own_padding(region_pull_worker): worker.region_num_blocks = [4, 4] worker.block_len_per_layer = [8, 8] worker.block_stride_per_layer = [24, 24] + worker._skip_dram_xfer = False worker._recving_metadata = { "request": ReqMeta( local_block_ids=[[1, 2], [0, 1]], diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index fa29d41a825e..2a2f660cf15e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -789,9 +789,12 @@ def __init__( self.region_names: list[str] = [] self.region_num_blocks: list[int] = [] self._mixed_mem_types = False + # True on TP ranks that share a HiSparse host pool with TP rank 0 and + # leave its NIXL registration and DRAM reads to rank 0. + self._skip_dram_xfer = False self._desc_is_dram_by_block_size: dict[int, np.ndarray] = {} self._desc_pos_by_block_size: dict[int, np.ndarray] = {} - self._dram_src_handles_by_block_size: dict[int, int] = {} + self._dram_src_handles_by_block_size: dict[int, int | None] = {} # PP>1 (push mode): this worker holds a contiguous layer slice and # transfers into the matching sub-range of a PP=1 remote's regions. @@ -828,7 +831,7 @@ def __init__( # Populated dynamically during handshake based on remote configuration. # Per-source split handles, keyed by (tp_ratio, remote_block_size). self.src_xfer_handles_by_tp_ratio: dict[tuple[int, int], list[int]] = {} - self._dram_src_handles_by_tp_ratio: dict[tuple[int, int], list[int]] = {} + self._dram_src_handles_by_tp_ratio: dict[tuple[int, int], list[int | None]] = {} # Map of engine_id -> {tp_rank: nixl_prepped_dlist_handle (int)}. self.dst_xfer_side_handles = defaultdict[EngineId, dict[int, int]](dict) @@ -1818,12 +1821,15 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.num_descs = sum(xfer_region_num_blocks) self._mixed_mem_types = len(set(region_mem_types)) > 1 + self._skip_dram_xfer = self._leaves_dram_xfer_to_rank0() if self._mixed_mem_types: assert self.use_mla and not self._has_mamba, ( "Mixed-device KV registration is only supported for MLA " "models without Mamba layers." ) for mem_type in sorted(set(region_mem_types)): + if mem_type == "DRAM" and self._skip_dram_xfer: + continue ranges_for_mem_type = [ (start, end - start, device_id, "") for (_, cache_mem_type), ( @@ -2191,10 +2197,13 @@ def register_local_xfer_handler( blocks_data[desc_is_dram, 2] = 0 dram_blocks = blocks_data[dram_idx] vram_blocks = blocks_data[vram_idx] - dram_descs = self.nixl_wrapper.get_xfer_descs(dram_blocks, "DRAM") - self._dram_src_handles_by_block_size[block_size] = ( - self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", dram_descs) - ) + if self._skip_dram_xfer: + self._dram_src_handles_by_block_size[block_size] = None + else: + dram_descs = self.nixl_wrapper.get_xfer_descs(dram_blocks, "DRAM") + self._dram_src_handles_by_block_size[block_size] = ( + self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", dram_descs) + ) descs = self.nixl_wrapper.get_xfer_descs(vram_blocks, self.nixl_memory_type) return ( self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs), @@ -2423,12 +2432,15 @@ def add_remote_agent( if self._mixed_mem_types: handle_data = np.asarray(handle_data, dtype=np.uint64) desc_is_dram = self._desc_is_dram_by_block_size[remote_block_size] - dram_descs = self.nixl_wrapper.get_xfer_descs( - handle_data[desc_is_dram], "DRAM" - ) - dram_handle = self.nixl_wrapper.prep_xfer_dlist( - "NIXL_INIT_AGENT", dram_descs - ) + if self._skip_dram_xfer: + dram_handle = None + else: + dram_descs = self.nixl_wrapper.get_xfer_descs( + handle_data[desc_is_dram], "DRAM" + ) + dram_handle = self.nixl_wrapper.prep_xfer_dlist( + "NIXL_INIT_AGENT", dram_descs + ) self._dram_src_handles_by_tp_ratio[split_key].append(dram_handle) handle_data = handle_data[~desc_is_dram] descs = self.nixl_wrapper.get_xfer_descs( @@ -2887,6 +2899,27 @@ def post_process_device_kv_on_receive_heterogeneous_attn( indices=indices, ) + def _leaves_dram_xfer_to_rank0(self) -> bool: + """Whether this rank leaves the DRAM regions' NIXL work to TP rank 0. + + A HiSparse host pool shared by the local TP ranks (one mmap per DP + rank) is registered with NIXL and filled by TP rank 0 alone: every + rank would otherwise register the same pages on every NIC and READ the + same host KV into them. Private pools are unaffected. + """ + if not ( + self._mixed_mem_types and self.kv_cache_config.hisparse_shared_host_pool + ): + return False + skip = self.tp_rank != 0 + logger.info( + "HiSparse host pool: NIXL DRAM registration and reads on TP rank 0 " + "only (tp_rank=%d, skip=%s)", + self.tp_rank, + skip, + ) + return skip + def _zero_region_blocks(self, block_ids: BlockIds) -> None: """Clear clipped physical pages in their owning region only.""" if not any(block_ids): @@ -2895,6 +2928,9 @@ def _zero_region_blocks(self, block_ids: BlockIds) -> None: for region, blocks in enumerate(block_ids): if not blocks: continue + if self._skip_dram_xfer and self.region_mem_types[region] == "DRAM": + # The shared pool is rank 0's to write. + continue cache = self.device_kv_caches[self.region_names[region]] storage = cache.untyped_storage() pages = torch.empty(0, dtype=torch.uint8, device=cache.device).set_( @@ -3795,11 +3831,13 @@ def _finish_shutdown(self) -> None: for handles in self.src_xfer_handles_by_tp_ratio.values(): for handle in handles: self.nixl_wrapper.release_dlist_handle(handle) - for handles in self._dram_src_handles_by_tp_ratio.values(): - for handle in handles: - self.nixl_wrapper.release_dlist_handle(handle) - for handle in self._dram_src_handles_by_block_size.values(): - self.nixl_wrapper.release_dlist_handle(handle) + for dram_handles in self._dram_src_handles_by_tp_ratio.values(): + for dram_handle in dram_handles: + if dram_handle is not None: + self.nixl_wrapper.release_dlist_handle(dram_handle) + for dram_handle in self._dram_src_handles_by_block_size.values(): + if dram_handle is not None: + self.nixl_wrapper.release_dlist_handle(dram_handle) except Exception: logger.exception("NIXL dlist-handle release failed at shutdown.") self.src_xfer_handles_by_block_size.clear() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 5bcfdb9560a0..16a25149ea22 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -593,7 +593,16 @@ def _read_blocks_mixed( remote_ids = np.asarray(remote_block_descs_ids) is_dram = desc_is_dram[local_ids] - assert local_dram_handle is not None + if self._skip_dram_xfer: + # TP rank 0 reads the DRAM (shared host pool) part for every + # local rank; keep only the device part. Register the request + # even if nothing is left, so it still completes and notifies. + keep = ~is_dram + local_ids = local_ids[keep] + remote_ids = remote_ids[keep] + is_dram = is_dram[keep] + self._recving_transfers.setdefault(request_id, []) + assert local_dram_handle is not None or not is_dram.any() reads = ( (is_dram, local_dram_handle), (~is_dram, local_device_handle), From 2f7f05c256d671a7bdf04bcf4b4fc6a1566bbfcb Mon Sep 17 00:00:00 2001 From: zzaebok <44357534+zzaebok@users.noreply.github.com> Date: Thu, 8 Oct 2026 22:28:16 +0800 Subject: [PATCH 098/624] [Bugfix][MoRIIO] Prevent remote KV reload after decoder preemption (#53250) Signed-off-by: Jaebok Lee --- tests/v1/kv_connector/unit/test_moriio_connector.py | 11 +++++++++++ .../v1/kv_connector/unit/test_moriio_hma_scheduler.py | 4 ++-- .../kv_connector/v1/moriio/moriio_connector.py | 3 +++ 3 files changed, 16 insertions(+), 2 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index cb33bfbcf03c..73cc5f0583d0 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -195,6 +195,17 @@ def _build_meta( return scheduler.build_connector_meta(scheduler_output) +def test_write_mode_does_not_reload_remote_kv_after_preemption(): + scheduler = _write_consumer_scheduler_for_finished_request() + request = create_request(request_id=1, num_tokens=32, do_remote_prefill=True) + + assert scheduler.get_num_new_matched_tokens(request, 0) == (32, True) + + assert request.kv_transfer_params is not None + request.kv_transfer_params["do_remote_prefill"] = False + assert scheduler.get_num_new_matched_tokens(request, 0) == (0, False) + + class FakeMoRIIOWrapper: # A fake MoRIIOWrapper for testing purposes def __init__(self, *args, **kwargs): diff --git a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py index 1de729416e90..4bd32f786a28 100644 --- a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py +++ b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py @@ -824,7 +824,7 @@ def test_get_num_new_matched_tokens_write_plain_keeps_all_tokens(): req = SimpleNamespace( num_prompt_tokens=10, prompt_token_ids=list(range(10)), - kv_transfer_params=None, + kv_transfer_params={"do_remote_prefill": True}, ) n, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=2) # Pure-attention WRITE: no N-1 drop; full length minus already-computed. @@ -848,7 +848,7 @@ def test_get_num_new_matched_tokens_supports_embeds_only_prompts( num_prompt_tokens=10, prompt_token_ids=None, prompt_embeds=object(), - kv_transfer_params=None, + kv_transfer_params={"do_remote_prefill": True}, ) assert sched.get_num_new_matched_tokens(req, num_computed_tokens) == ( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 2d17d7a7ac52..082643778e7d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -769,6 +769,9 @@ def get_num_new_matched_tokens( num_prompt_tokens = request.num_prompt_tokens num_external_tokens = max(num_prompt_tokens - num_computed_tokens, 0) if self.mode == MoRIIOMode.WRITE: + params = request.kv_transfer_params or {} + if not params.get("do_remote_prefill"): + return 0, False # MoriiO in write mode, no remote prefill. Hybrid models never get # here: register_kv_caches refuses WRITE mode for them, so there is # no recurrent-state accounting to do. From 273e099693a6c92e1bc516cd10ffc38b18197eec Mon Sep 17 00:00:00 2001 From: "yhcheong (vllmellm)" Date: Thu, 8 Oct 2026 22:44:58 +0800 Subject: [PATCH 099/624] [Bugfix][ROCm] Select page-aligned kernel blocks for pooled indexers so the block table addresses storage pages (GLM-5.3-Flash) (#59412) Signed-off-by: vllmellm --- .../test_kpool_indexer_block_sizes.py | 135 ++++++++++++++++++ vllm/v1/worker/utils.py | 67 +++++---- 2 files changed, 178 insertions(+), 24 deletions(-) create mode 100644 tests/v1/attention/test_kpool_indexer_block_sizes.py diff --git a/tests/v1/attention/test_kpool_indexer_block_sizes.py b/tests/v1/attention/test_kpool_indexer_block_sizes.py new file mode 100644 index 000000000000..a737988dde66 --- /dev/null +++ b/tests/v1/attention/test_kpool_indexer_block_sizes.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU tests for storage-block kernel-block selection in prepare_kernel_block_sizes.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.models.glm5next.common.attention import Glm5NextIndexerCache +from vllm.v1.attention.backend import AttentionBackend, MultipleOf +from vllm.v1.attention.backends.mla import indexer as indexer_mod +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerBackend, + KpoolTailBackend, +) +from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( + ROCMAiterMLASparseBackend, +) +from vllm.v1.kv_cache_interface import MLAAttentionSpec +from vllm.v1.worker.utils import prepare_kernel_block_sizes + +pytestmark = pytest.mark.cpu_test + +# ``index_kpool`` of zai-org/GLM-5.3-Flash: one indexer state pools 4 +# compressed states (hence storage_block_size = page * 4). +INDEX_KPOOL = 4 + +# The hybrid KDA/mamba manager floors at TP8..TP1; 640 is the geometry of the +# original report ("the 640-token block table"). +TP_FLOOR_MANAGER_BLOCKS = [640, 1152, 2176, 4352] + + +def _mock_rocm_platform(monkeypatch: pytest.MonkeyPatch) -> None: + # Platform predicates are mutually exclusive in production. Override both + # so these ROCm policy tests do not inherit CUDA capabilities from the CI + # host running them (same approach as test_deepseek_v4_rocm_adaptive). + monkeypatch.setattr(indexer_mod.current_platform, "is_cuda", lambda: False) + monkeypatch.setattr(indexer_mod.current_platform, "is_rocm", lambda: True) + + +def _kpool_storage_spec(manager_block: int) -> MLAAttentionSpec: + """Real Glm5NextIndexerCache spec for the given cache block size.""" + with set_current_vllm_config(VllmConfig()): + cache = Glm5NextIndexerCache( + head_dim=128, + dtype=torch.bfloat16, + prefix="model.layers.0.indexer.k_cache_probe", + cache_config=SimpleNamespace(block_size=manager_block), + index_kpool=INDEX_KPOOL, + ) + spec = cache.get_kv_cache_spec(VllmConfig()) + assert isinstance(spec, MLAAttentionSpec) + return spec + + +def _prepare( + spec: MLAAttentionSpec, backends: list[type[AttentionBackend]] +) -> list[int]: + kv_cache_config = SimpleNamespace( + kv_cache_groups=[SimpleNamespace(kv_cache_spec=spec)] + ) + attn_groups = [[SimpleNamespace(backend=backend) for backend in backends]] + return prepare_kernel_block_sizes(kv_cache_config, attn_groups) + + +@pytest.mark.parametrize("manager_block", TP_FLOOR_MANAGER_BLOCKS) +def test_prepare_uses_storage_block_for_the_kpool_group( + monkeypatch: pytest.MonkeyPatch, manager_block: int +): + """The bug: on ROCm, kpool groups must get 128/256 - not the 640-4352 + manager block - as their kernel block size. + + Manager-granular selection is what made the block table address manager + blocks while the kpool writer and the index-cache gather read pool pages, + aliasing every cache column past the request's row. + """ + _mock_rocm_platform(monkeypatch) + spec = _kpool_storage_spec(manager_block) + expected_storage = 256 if manager_block % 256 == 0 else 128 + assert spec.storage_block_size == expected_storage + + selected = _prepare( + spec, + [ROCMAiterMLASparseBackend, DeepseekV32IndexerBackend, KpoolTailBackend], + ) + assert selected == [expected_storage] + assert selected != [manager_block] + # The hybrid block-table split stays integral. + assert manager_block % selected[0] == 0 + + +def test_prepare_falls_back_when_storage_block_is_not_supported(): + """No storage block, or a storage block the backends reject (e.g. the + CUDA indexer's exact ``[64]`` vs a 128-token storage block): selection is + `select_common_block_size`'s answer, unchanged from no-fix behavior. + """ + + class Fixed64Backend: + @staticmethod + def get_supported_kernel_block_sizes(): + return [64] + + @staticmethod + def get_name() -> str: + return "FIXED64" + + # Storage block the group's backend does not accept -> backend vote (64). + spec = _kpool_storage_spec(640) + assert spec.storage_block_size == 128 + assert _prepare(spec, [Fixed64Backend]) == [64] + + # No storage block at all (any non-kpool MLA model) -> backend vote. + plain_spec = MLAAttentionSpec( + block_size=640, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + ) + assert plain_spec.storage_block_size is None + assert _prepare(plain_spec, [Fixed64Backend]) == [64] + + # ...and with backends that accept the manager block, the manager wins, + # exactly as on main. + class AcceptAllBackend: + @staticmethod + def get_supported_kernel_block_sizes(): + return [MultipleOf(1)] + + @staticmethod + def get_name() -> str: + return "ACCEPTALL" + + assert _prepare(plain_spec, [AcceptAllBackend]) == [640] diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 422ba8885fd7..b61c4cebed18 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -327,6 +327,30 @@ def update_draft_decode_metadata( self.get_metadata_builder().update_draft_decode_metadata(metadata) +def _block_size_is_supported( + backends: list[type[AttentionBackend]], block_size: int +) -> bool: + """Check if the block size is supported by all backends. + + An exact ``int`` declaration must match exactly; a ``MultipleOf`` + declaration accepts any multiple of its base. + """ + for backend in backends: + is_supported = False + for supported_size in backend.get_supported_kernel_block_sizes(): + if isinstance(supported_size, int): + if block_size == supported_size: + is_supported = True + elif isinstance(supported_size, MultipleOf): + if block_size % supported_size.base == 0: + is_supported = True + else: + raise ValueError(f"Unknown supported size: {supported_size}") + if not is_supported: + return False + return True + + def select_common_block_size( kv_manager_block_size: int, backends: list[type[AttentionBackend]], @@ -348,27 +372,7 @@ def select_common_block_size( ValueError: If no valid block size found. """ - - def block_size_is_supported( - backends: list[type[AttentionBackend]], block_size: int - ) -> bool: - """Check if the block size is supported by all backends.""" - for backend in backends: - is_supported = False - for supported_size in backend.get_supported_kernel_block_sizes(): - if isinstance(supported_size, int): - if block_size == supported_size: - is_supported = True - elif isinstance(supported_size, MultipleOf): - if block_size % supported_size.base == 0: - is_supported = True - else: - raise ValueError(f"Unknown supported size: {supported_size}") - if not is_supported: - return False - return True - - if block_size_is_supported(backends, kv_manager_block_size): + if _block_size_is_supported(backends, kv_manager_block_size): return kv_manager_block_size # MultipleOf constraints also accept the manager size if they accept a divisor. @@ -381,7 +385,7 @@ def block_size_is_supported( } for size in sorted(candidates, reverse=True): - if block_size_is_supported(backends, size): + if _block_size_is_supported(backends, size): return size raise ValueError( f"No common block size for {kv_manager_block_size} (" @@ -493,9 +497,24 @@ def prepare_kernel_block_sizes( # This is an attention backend that supports virtual block splitting. kv_manager_block_size = kv_cache_group.kv_cache_spec.block_size group_backends = [g.backend for g in attn_groups[kv_cache_gid]] - selected_kernel_size = select_common_block_size( - kv_manager_block_size, group_backends + storage_block_size = ( + kv_cache_spec.storage_block_size + if isinstance(kv_cache_spec, MLAAttentionSpec) + else None ) + if storage_block_size is not None and _block_size_is_supported( + group_backends, storage_block_size + ): + # Storage-block specs (e.g. the GLM-5.3-Flash kpool indexer + # cache) address the cache in pool pages, and every other + # consumer (cache views, metadata builders, hisparse) already + # uses storage_block_size as the kernel block. Fall back to the + # backend vote when the group's backends do not accept it. + selected_kernel_size = storage_block_size + else: + selected_kernel_size = select_common_block_size( + kv_manager_block_size, group_backends + ) kernel_block_sizes.append(selected_kernel_size) elif isinstance(kv_cache_spec, MambaSpec): # This is likely Mamba or other non-attention cache, no splitting. From d33872dfeb79db89b734ce13e07af4630e8bdce2 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Thu, 8 Oct 2026 08:48:12 -0600 Subject: [PATCH 100/624] [Feat] Structured decisions endpoint (`/v1/systemone`) (#59299) Signed-off-by: Matt Mastracci Signed-off-by: Lucas Wilkinson Signed-off-by: Francisco Javier Arceo <4163062+franciscojavierarceo@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Co-authored-by: Lucas Wilkinson Co-authored-by: Lucas Wilkinson Co-authored-by: Francisco Javier Arceo <4163062+franciscojavierarceo@users.noreply.github.com> Co-authored-by: Chauncey --- docs/serving/online_serving/README.md | 3 + .../online_serving/structured_decisions.md | 110 ++++++++ docs/usage/security.md | 1 + .../generate/structured_decisions/__init__.py | 0 .../test_question_types.py | 105 ++++++++ .../structured_decisions/test_strategies.py | 65 +++++ .../test_structured_decisions_e2e.py | 217 ++++++++++++++++ vllm/entrypoints/generate/api_router.py | 33 +++ .../generate/generative_scoring/serving.py | 75 ++---- vllm/entrypoints/generate/label_reads.py | 76 ++++++ .../generate/structured_decisions/__init__.py | 2 + .../structured_decisions/api_router.py | 53 ++++ .../generate/structured_decisions/protocol.py | 88 +++++++ .../structured_decisions/question_types.py | 160 ++++++++++++ .../generate/structured_decisions/serving.py | 121 +++++++++ .../structured_decisions/strategies.py | 235 ++++++++++++++++++ vllm/entrypoints/launchers/cli_args.py | 2 + 17 files changed, 1285 insertions(+), 61 deletions(-) create mode 100644 docs/serving/online_serving/structured_decisions.md create mode 100644 tests/entrypoints/generate/structured_decisions/__init__.py create mode 100644 tests/entrypoints/generate/structured_decisions/test_question_types.py create mode 100644 tests/entrypoints/generate/structured_decisions/test_strategies.py create mode 100644 tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py create mode 100644 vllm/entrypoints/generate/label_reads.py create mode 100644 vllm/entrypoints/generate/structured_decisions/__init__.py create mode 100644 vllm/entrypoints/generate/structured_decisions/api_router.py create mode 100644 vllm/entrypoints/generate/structured_decisions/protocol.py create mode 100644 vllm/entrypoints/generate/structured_decisions/question_types.py create mode 100644 vllm/entrypoints/generate/structured_decisions/serving.py create mode 100644 vllm/entrypoints/generate/structured_decisions/strategies.py diff --git a/docs/serving/online_serving/README.md b/docs/serving/online_serving/README.md index 8e1a4f3398ec..fc46a482c157 100644 --- a/docs/serving/online_serving/README.md +++ b/docs/serving/online_serving/README.md @@ -76,6 +76,9 @@ For further details on speech to text, please refer to [this page](speech_to_tex - [Generative Scoring API](generative_scoring.md) (`/generative_scoring`) - Applicable to [CausalLM models](../../models/generative_models.md) (task `"generate"`). - Computes next-token probabilities for specified `label_token_ids`. +- [Structured Decisions API](structured_decisions.md) (`/v1/systemone`) + - Applicable to [text generation models](../../models/generative_models.md) with a [chat template](#chat-template). + - Answers typed questions about a state with a probability for every allowed answer. ## Instrumentator APIs diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md new file mode 100644 index 000000000000..dd26f2130164 --- /dev/null +++ b/docs/serving/online_serving/structured_decisions.md @@ -0,0 +1,110 @@ +# Structured Decisions + +The `/v1/systemone` endpoint answers a set of typed questions about a state and +returns a probability for every allowed answer. + +The endpoint is off unless the server starts with +`--enable-structured-decisions`. With the flag, startup fails for a model that +does not support structured decisions, and with `--logprobs-mode raw_logits` or +`processed_logits`, since `label_mass` needs log probabilities. + +## How it works + +1. Each question is one read. The user message is the state, followed by the + question and its options, labeled `A`, `B`, `C`, ... in order, and + "Answer with the letter of one option only." +2. The read is the reply's first token. Each label must be one token there, + or the request returns 400. +3. The read returns the logprobs of that question's label tokens. Softmax over + the labels gives the answer's probabilities. + +Thinking is off unless `chat_template_kwargs` turns it on. + +This endpoint does not serve diffusion models yet. + +Every read of a request starts with the state, so with +`--enable-prefix-caching` the state is prefilled once and each further +question prefills only its own text, then reads one token. + +## Question types + +| type | criteria | answer | +| --- | --- | --- | +| `choice` | map of option name to a description or `null` | `choice`, `probabilities` by option name, `confidence` | + +A choice has at least one option. A question id is a non-empty string. + +## Example + +```bash +curl -s http://localhost:8000/v1/systemone \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Qwen/Qwen3-0.6B", + "state": {"ticket": "My card was charged twice for one order."}, + "questions": { + "team": { + "type": "choice", + "instructions": "Which team should handle this ticket?", + "criteria": { + "billing": "payments and refunds", + "shipping": "deliveries", + "security": "account access" + } + } + }, + "chat_template_kwargs": {"enable_thinking": false} + }' +``` + +Response, with probabilities that depend on the model: + +```json +{ + "id": "decision-...", + "object": "structured_decision", + "created": 1790000000, + "model": "Qwen/Qwen3-0.6B", + "answers": { + "team": { + "type": "choice", + "choice": "billing", + "probabilities": {"billing": 0.97, "shipping": 0.01, "security": 0.02}, + "confidence": 0.90 + } + }, + "usage": {"input_tokens": 142, "output_tokens": 1}, + "diagnostics": { + "team": {"label_mass": 0.93, "argmax_is_label": true} + } +} +``` + +`diagnostics.label_mass` is the probability the model put on the labels across +its whole vocabulary. A low value means most of the probability went to tokens +that are not labels. `confidence` is the chosen option's probability over the +whole vocabulary: its share of the labels times `label_mass`. + +## Labels + +Options are labeled `A` to `Z` in the order the request lists them, so a +`choice` has at most 26 options. + +## Limits + +| limit | value | +| --- | --- | +| questions per request | 64 | +| options per `choice` | 26 | + +## Request fields + +| field | meaning | +| --- | --- | +| `state` | what the questions are about: a string, or JSON that is sent as its JSON text | +| `questions` | question id to `{type, instructions, criteria}`, asked in this order | +| `instructions` | optional context placed ahead of the questions | +| `chat_template_kwargs` | passed to the chat template, for example `{"enable_thinking": false}` | + +A question with any field other than `type`, `instructions` and `criteria` is +rejected with a 400. diff --git a/docs/usage/security.md b/docs/usage/security.md index 3dad19d6cb03..c0d2c79700e6 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -175,6 +175,7 @@ When `--api-key` is configured, the following endpoints require Bearer token aut - `/v1/responses/{response_id}/cancel` - Cancel a response - `/v1/score` - Scoring API - `/v1/rerank` - Reranking API +- `/v1/systemone` - Structured decisions API - `/v1/load_lora_adapter` - Load a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`) - `/v1/unload_lora_adapter` - Unload a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`) - `/inference/v1/generate` - Generate completions (available when `--enable-scale-out` is set, or with `--tokens-only`) diff --git a/tests/entrypoints/generate/structured_decisions/__init__.py b/tests/entrypoints/generate/structured_decisions/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/entrypoints/generate/structured_decisions/test_question_types.py b/tests/entrypoints/generate/structured_decisions/test_question_types.py new file mode 100644 index 000000000000..e5906accf025 --- /dev/null +++ b/tests/entrypoints/generate/structured_decisions/test_question_types.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + +import pytest + +from vllm.entrypoints.generate.structured_decisions.question_types import ( + QUESTION_TYPES, + Option, + Question, + QuestionType, + StructuredDecisionError, + build_question, + register_question_type, +) + + +def choice(qid="bucket", criteria=None, instructions="Which team?", max_options=128): + return build_question( + qid, + "choice", + instructions, + criteria or {"billing": "money", "outage": None, "other": None}, + max_options, + ) + + +def test_choice_labels_follow_option_order(): + q = choice() + assert q.labels == ("A", "B", "C") + assert [a.name for a in q.options] == ["billing", "outage", "other"] + assert q.options[0].description == "money" + wide = choice(criteria={str(i): None for i in range(26)}) + assert wide.labels[-1] == "Z" + + +def test_choice_prompt(): + q = choice() + assert q.type.prompt(q) == ( + "Question: Which team?\n" + "A: billing - money\n" + "B: outage\n" + "C: other\n" + "Answer with the letter of one option only." + ) + + +def test_choice_answer_shape(): + q = choice() + answer = q.type.answer(q, [0.2, 0.7, 0.1], 0.5) + assert answer == { + "type": "choice", + "choice": "outage", + "probabilities": {"billing": 0.2, "outage": 0.7, "other": 0.1}, + "confidence": 0.35, + } + + +@pytest.mark.parametrize( + "qid,type_name,criteria,match", + [ + ("", "choice", {"x": None, "y": None}, "non-empty"), + ("q", "nope", {"x": None, "y": None}, "unknown question type"), + ("q", "choice", ["x", "y"], "must map option names"), + ], +) +def test_build_question_rejects(qid, type_name, criteria, match): + with pytest.raises(StructuredDecisionError, match=match): + build_question(qid, type_name, "", criteria, 128) + + +def test_one_option_choice(): + q = choice(criteria={"only": None}) + assert q.labels == ("A",) + assert q.type.answer(q, [1.0], 0.4)["confidence"] == 0.4 + + +def test_option_limit(): + with pytest.raises(StructuredDecisionError, match="at most 2 options"): + choice(criteria={"x": None, "y": None, "z": None}, max_options=2) + with pytest.raises(StructuredDecisionError, match="at most 26 options"): + choice(criteria={str(i): None for i in range(27)}) + + +def test_registered_type_plugs_in(): + class BinaryQuestion(QuestionType): + name = "test_binary" + + def parse_options(self, qid: str, criteria: Any) -> list[Option]: + return [Option("yes"), Option("no")] + + def answer( + self, question: Question, probs: list[float], label_mass: float + ) -> dict[str, Any]: + return {"type": self.name, "yes": probs[0]} + + register_question_type(BinaryQuestion) + try: + q = build_question("ok", "test_binary", "Is it fine?", None, 128) + assert q.labels == ("A", "B") + assert q.type.answer(q, [0.9, 0.1], 1.0) == {"type": "test_binary", "yes": 0.9} + with pytest.raises(ValueError, match="already registered"): + register_question_type(BinaryQuestion) + finally: + del QUESTION_TYPES["test_binary"] diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py new file mode 100644 index 000000000000..c4f4bce61be1 --- /dev/null +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace +from typing import Any + +import pytest +from fastapi import FastAPI + +from vllm.entrypoints.generate.structured_decisions.api_router import ( + register_structured_decisions_api_router, +) +from vllm.entrypoints.generate.structured_decisions.question_types import LABELS +from vllm.entrypoints.generate.structured_decisions.strategies import ( + NextTokenStrategy, + reply_label_ids, + select_read_strategy, +) + + +def model(architecture: str, logprobs_mode: str = "raw_logprobs") -> Any: + return SimpleNamespace(architecture=architecture, logprobs_mode=logprobs_mode) + + +def test_strategy_selection(): + qwen = "Qwen3ForCausalLM" + assert select_read_strategy(model(qwen)) is NextTokenStrategy + assert select_read_strategy(model(qwen, "processed_logprobs")) is NextTokenStrategy + with pytest.raises(ValueError, match="does not support LlamaForCausalLM"): + select_read_strategy(model("LlamaForCausalLM")) + with pytest.raises(ValueError, match="not raw_logits"): + select_read_strategy(model(qwen, "raw_logits")) + + +def test_route_needs_the_flag(): + for enabled in (False, True): + app = FastAPI() + app.state.args = SimpleNamespace(enable_structured_decisions=enabled) + register_structured_decisions_api_router(app) + paths = {getattr(route, "path", None) for route in app.routes} + assert ("/v1/systemone" in paths) == enabled + + +@pytest.fixture(scope="module") +def qwen(): + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") + prompt_ids = tokenizer.apply_chat_template( + [{"role": "user", "content": "state"}], + add_generation_prompt=True, + enable_thinking=False, + tokenize=True, + return_dict=False, + ) + return tokenizer, prompt_ids + + +def test_labels_start_the_reply(qwen): + tokenizer, prompt_ids = qwen + tail, ids = reply_label_ids(tokenizer, prompt_ids) + assert tokenizer.decode(tail) == "\n\n" + assert [tokenizer.decode([i]) for i in ids] == list(LABELS) + # After a colon, Qwen writes ":A" as one token, so "A" is not one token. + with pytest.raises(ValueError, match="not one distinct token"): + reply_label_ids(tokenizer, tokenizer.encode("team:")) diff --git a/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py b/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py new file mode 100644 index 000000000000..e060b096ec51 --- /dev/null +++ b/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import requests + +from tests.utils import RemoteOpenAIServer + +MODEL_NAME = "Qwen/Qwen3-0.6B" + + +@pytest.fixture(scope="module") +def server(): + args = [ + "--dtype", + "bfloat16", + "--max-model-len", + "1024", + "--enforce-eager", + "--max-num-seqs", + "32", + "--enable-prefix-caching", + "--enable-structured-decisions", + ] + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: + yield remote_server + + +def post(server, body): + return requests.post(server.url_for("v1/systemone"), json=body) + + +def choice(instructions, *names): + return { + "type": "choice", + "instructions": instructions, + "criteria": dict.fromkeys(names), + } + + +COPY = { + "team": choice( + "Which team does the message name?", "billing", "shipping", "security" + ), + "lang": choice("Which language does the message name?", "English", "French"), +} + + +def test_choice_decision(server): + response = post( + server, + { + "model": MODEL_NAME, + "state": {"team": "billing", "language": "French"}, + "questions": COPY, + "chat_template_kwargs": {"enable_thinking": False}, + }, + ) + assert response.status_code == 200, response.text + body = response.json() + assert body["object"] == "structured_decision" + assert list(body["answers"]) == ["team", "lang"] + team = body["answers"]["team"] + assert team["type"] == "choice" + assert set(team["probabilities"]) == {"billing", "shipping", "security"} + assert sum(team["probabilities"].values()) == pytest.approx(1.0, abs=1e-4) + assert team["confidence"] == pytest.approx( + max(team["probabilities"].values()) * body["diagnostics"]["team"]["label_mass"] + ) + for diag in body["diagnostics"].values(): + assert 0.0 < diag["label_mass"] <= 1.0 + 1e-6 + assert body["usage"]["output_tokens"] == 2 + assert body["usage"]["input_tokens"] > 0 + + +def test_one_option_choice(server): + response = post( + server, + { + "model": MODEL_NAME, + "state": "Team: billing.", + "questions": { + "team": choice("Which team does the message name?", "billing") + }, + "chat_template_kwargs": {"enable_thinking": False}, + }, + ) + assert response.status_code == 200, response.text + body = response.json() + team = body["answers"]["team"] + assert team["choice"] == "billing" + assert team["probabilities"] == {"billing": 1.0} + assert team["confidence"] == pytest.approx( + body["diagnostics"]["team"]["label_mass"] + ) + + +def test_one_option_coin_flip(server): + # A fair coin with only "heads" allowed. The answer is forced, so its + # confidence must come from the model's probability of the label. + response = post( + server, + { + "model": MODEL_NAME, + "state": "I flip a fair coin.", + "questions": {"coin": choice("How did the coin land?", "heads")}, + "chat_template_kwargs": {"enable_thinking": False}, + }, + ) + assert response.status_code == 200, response.text + body = response.json() + coin = body["answers"]["coin"] + assert coin["probabilities"] == {"heads": 1.0} + assert coin["confidence"] == pytest.approx( + body["diagnostics"]["coin"]["label_mass"] + ) + + +@pytest.mark.parametrize( + "state,questions,expected", + [ + ( + {"team": "billing", "language": "French"}, + COPY, + {"team": "billing", "lang": "French"}, + ), + ( + "My card was charged twice for one order.", + { + "card": choice("Does the message mention a card?", "yes", "no"), + "mood": choice("How does the customer feel?", "happy", "angry"), + }, + {"card": "yes", "mood": "angry"}, + ), + ( + "Hello.", + { + "sky": choice( + "What color is the sky on a clear day?", "blue", "green", "red" + ) + }, + {"sky": "blue"}, + ), + ], +) +def test_single_read_answers(server, state, questions, expected): + response = post( + server, + { + "model": MODEL_NAME, + "state": state, + "questions": questions, + "chat_template_kwargs": {"enable_thinking": False}, + }, + ) + assert response.status_code == 200, response.text + answers = response.json()["answers"] + assert {qid: a["choice"] for qid, a in answers.items()} == expected + + +def test_option_limit(server): + body = { + "model": MODEL_NAME, + "state": "Pick option 7.", + "chat_template_kwargs": {"enable_thinking": False}, + } + ok = post( + server, {**body, "questions": {"n": choice("Which?", *map(str, range(26)))}} + ) + assert ok.status_code == 200, ok.text + assert len(ok.json()["answers"]["n"]["probabilities"]) == 26 + over = post( + server, {**body, "questions": {"n": choice("Which?", *map(str, range(27)))}} + ) + assert over.status_code == 400 + assert "at most 26 options" in over.json()["error"]["message"] + + +def test_thinking_is_refused(server): + response = post( + server, + { + "model": MODEL_NAME, + "state": "x", + "questions": {"q": choice("Which?", "a", "b")}, + "chat_template_kwargs": {"enable_thinking": True}, + }, + ) + assert response.status_code == 400 + assert "thinking must be off" in response.json()["error"]["message"] + + +@pytest.mark.parametrize( + "questions,match", + [ + ( + { + "q": { + "type": "choice", + "criteria": {"a": None, "b": None}, + "depends_on": ["x"], + } + }, + "unknown field(s)", + ), + ( + { + f"q{i}": {"type": "choice", "criteria": {"a": None, "b": None}} + for i in range(65) + }, + "at most 64", + ), + ], +) +def test_validation_errors(server, questions, match): + response = post(server, {"model": MODEL_NAME, "state": "x", "questions": questions}) + assert response.status_code == 400 + assert match in response.json()["error"]["message"] diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index c5703515d82b..70869b61b5d2 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -53,6 +53,12 @@ def register_generate_api_routers(app: FastAPI): register_generative_scoring_api_router(app) + from .structured_decisions.api_router import ( + register_structured_decisions_api_router, + ) + + register_structured_decisions_api_router(app) + async def init_generate_state( engine_client: "EngineClient", @@ -220,3 +226,30 @@ async def init_generate_state( state.openai_serving_models, request_logger=request_logger, ) + + from .structured_decisions.serving import ServingStructuredDecisions + from .structured_decisions.strategies import ReadContext, select_read_strategy + + strategy_cls = ( + select_read_strategy(engine_client.model_config) + if "generate" in supported_tasks + and getattr(args, "enable_structured_decisions", False) + else None + ) + state.serving_structured_decisions = ( + ServingStructuredDecisions( + state.openai_serving_models, + strategy_cls( + ReadContext( + engine_client=engine_client, + online_renderer=state.online_renderer, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + default_chat_template_kwargs=default_chat_template_kwargs, + ) + ), + request_logger=request_logger, + ) + if strategy_cls is not None + else None + ) diff --git a/vllm/entrypoints/generate/generative_scoring/serving.py b/vllm/entrypoints/generate/generative_scoring/serving.py index 4d5a15c2f54c..bd46158edcab 100644 --- a/vllm/entrypoints/generate/generative_scoring/serving.py +++ b/vllm/entrypoints/generate/generative_scoring/serving.py @@ -11,13 +11,14 @@ import asyncio import math import time -from collections.abc import AsyncGenerator, Mapping +from collections.abc import Mapping from typing import Literal from fastapi import Request from pydantic import Field from vllm.engine.protocol import EngineClient +from vllm.entrypoints.generate.label_reads import next_token_label_reads from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.serve.engine.protocol import ( ErrorResponse, @@ -28,7 +29,6 @@ from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.inputs import EngineInput, tokens_input from vllm.logger import init_logger -from vllm.outputs import RequestOutput from vllm.sampling_params import SamplingParams from vllm.tokenizers import TokenizerLike from vllm.tracing import ( @@ -37,7 +37,6 @@ log_tracing_disabled_warning, ) from vllm.utils import random_uuid -from vllm.utils.async_utils import merge_async_iterators logger = init_logger(__name__) @@ -271,37 +270,28 @@ async def create_generative_scoring( else await self._get_trace_headers(raw_request.headers) ) - # Schedule requests for all inputs - generators: list[AsyncGenerator[RequestOutput, None]] = [] for i, engine_input in enumerate(engine_inputs): - request_id_item = f"{request_id}-{i}" - self._log_inputs( - request_id_item, + f"{request_id}-{i}", engine_input, params=sampling_params, lora_request=lora_request, ) - generator = self.engine_client.generate( - engine_input, - sampling_params, - request_id_item, + try: + reads = await next_token_label_reads( + self.engine_client, + engine_inputs, + [sampling_params] * len(engine_inputs), + request_id, lora_request=lora_request, trace_headers=trace_headers, priority=request.priority, ) - generators.append(generator) - - # Collect results - result_generator = merge_async_iterators(*generators) - results: list[RequestOutput | None] = [None] * len(engine_inputs) - - try: - async for i, res in result_generator: - results[i] = res except asyncio.CancelledError: return self.create_error_response("Client disconnected") + except ValueError as e: + return self.create_error_response(e) except Exception as e: logger.exception("Error during generation") return self.create_error_response(e) @@ -311,45 +301,8 @@ async def create_generative_scoring( total_prompt_tokens = 0 total_completion_tokens = 0 - for i, result in enumerate(results): - if result is None: - return self.create_error_response( - f"Failed to generate result for item {i}" - ) - - # Check for errors - if result.outputs and result.outputs[0].finish_reason == "error": - return self.create_error_response(f"Generation error for item {i}") - - # Get logprobs from the generated token - if not result.outputs or len(result.outputs) == 0: - return self.create_error_response(f"No output generated for item {i}") - - output = result.outputs[0] - if output.logprobs is None or len(output.logprobs) == 0: - return self.create_error_response( - f"No logprobs available for item {i}. " - "This might indicate an issue with logprobs configuration." - ) - - # The logprobs dict maps token_id -> Logprob object - # For logprobs=-1, this contains all vocab tokens - logprobs_dict = output.logprobs[0] - - # Extract logprobs for label tokens - label_logprobs: dict[int, float] = {} - missing_tokens = [] - for token_id in request.label_token_ids: - if token_id in logprobs_dict: - label_logprobs[token_id] = logprobs_dict[token_id].logprob - else: - missing_tokens.append(token_id) - - if missing_tokens: - return self.create_error_response( - f"Token IDs {missing_tokens} not found in logprobs for item {i}. " - "This might indicate the tokens are outside the model's vocabulary." - ) + for i, read in enumerate(reads): + label_logprobs = dict(zip(request.label_token_ids, read.logprobs)) # Compute probabilities based on apply_softmax setting token_probs = self._compute_probabilities( @@ -370,7 +323,7 @@ async def create_generative_scoring( # Update token counts total_prompt_tokens += prompt_token_counts[i] - total_completion_tokens += len(output.token_ids) + total_completion_tokens += len(read.result.outputs[0].token_ids) # Build response model_name = self.models.model_name(lora_request) diff --git a/vllm/entrypoints/generate/label_reads.py b/vllm/entrypoints/generate/label_reads.py new file mode 100644 index 000000000000..eda61b22bc5a --- /dev/null +++ b/vllm/entrypoints/generate/label_reads.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""One-token reads that return the logprobs of chosen label tokens.""" + +from collections.abc import AsyncGenerator, Mapping, Sequence +from dataclasses import dataclass + +from vllm.engine.protocol import EngineClient +from vllm.inputs import EngineInput +from vllm.lora.request import LoRARequest +from vllm.outputs import RequestOutput +from vllm.sampling_params import SamplingParams +from vllm.utils.async_utils import merge_async_iterators + + +@dataclass(frozen=True) +class LabelRead: + result: RequestOutput + logprobs: list[float] + """Full-vocabulary logprob of each label token, in ``logprob_token_ids`` + order.""" + + +async def next_token_label_reads( + engine_client: EngineClient, + engine_inputs: Sequence[EngineInput], + sampling_params: Sequence[SamplingParams], + request_id: str, + *, + lora_request: LoRARequest | None = None, + trace_headers: Mapping[str, str] | None = None, + priority: int = 0, +) -> list[LabelRead]: + """Runs one read per input at once, as request ``{request_id}-{i}``. Each + read's params set ``max_tokens=1`` and the label tokens as + ``logprob_token_ids``. Raises ValueError naming the item when a read has + no output or lacks a label's logprob.""" + generators: list[AsyncGenerator[RequestOutput, None]] = [ + engine_client.generate( + engine_input, + params, + f"{request_id}-{i}", + lora_request=lora_request, + trace_headers=trace_headers, + priority=priority, + ) + for i, (engine_input, params) in enumerate(zip(engine_inputs, sampling_params)) + ] + results: list[RequestOutput | None] = [None] * len(generators) + async for i, res in merge_async_iterators(*generators): + results[i] = res + + reads = [] + for i, (result, params) in enumerate(zip(results, sampling_params)): + if result is None: + raise ValueError(f"Failed to generate result for item {i}") + if not result.outputs: + raise ValueError(f"No output generated for item {i}") + output = result.outputs[0] + if output.finish_reason == "error": + raise ValueError(f"Generation error for item {i}") + if not output.logprobs: + raise ValueError( + f"No logprobs available for item {i}. " + "This might indicate an issue with logprobs configuration." + ) + logprobs = output.logprobs[0] + label_ids = params.logprob_token_ids or [] + missing = [t for t in label_ids if t not in logprobs] + if missing: + raise ValueError( + f"Token IDs {missing} not found in logprobs for item {i}. " + "This might indicate the tokens are outside the model's vocabulary." + ) + reads.append(LabelRead(result, [logprobs[t].logprob for t in label_ids])) + return reads diff --git a/vllm/entrypoints/generate/structured_decisions/__init__.py b/vllm/entrypoints/generate/structured_decisions/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/entrypoints/generate/structured_decisions/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/entrypoints/generate/structured_decisions/api_router.py b/vllm/entrypoints/generate/structured_decisions/api_router.py new file mode 100644 index 000000000000..a0fbaa518e07 --- /dev/null +++ b/vllm/entrypoints/generate/structured_decisions/api_router.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from http import HTTPStatus + +from fastapi import APIRouter, Depends, FastAPI, Request +from fastapi.responses import JSONResponse + +from vllm.entrypoints.serve.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.utils.api_utils import ( + load_aware_call, + validate_json_request, + with_cancellation, +) + +from .protocol import StructuredDecisionRequest +from .serving import ServingStructuredDecisions + +router = APIRouter() + + +def structured_decisions(request: Request) -> ServingStructuredDecisions | None: + return getattr(request.app.state, "serving_structured_decisions", None) + + +@router.post( + "/v1/systemone", + dependencies=[Depends(validate_json_request)], + responses={ + HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, + HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, + HTTPStatus.NOT_IMPLEMENTED.value: {"model": ErrorResponse}, + }, +) +@with_cancellation +@load_aware_call +async def create_structured_decision( + request: StructuredDecisionRequest, raw_request: Request +): + handler = structured_decisions(raw_request) + if handler is None: + raise NotImplementedError( + "The model does not support the structured decision API" + ) + result = await handler.create_decision(request, raw_request) + if isinstance(result, ErrorResponse): + return JSONResponse(content=result.model_dump(), status_code=result.error.code) + return JSONResponse(content=result.model_dump()) + + +def register_structured_decisions_api_router(app: FastAPI): + if getattr(app.state.args, "enable_structured_decisions", False): + app.include_router(router) diff --git a/vllm/entrypoints/generate/structured_decisions/protocol.py b/vllm/entrypoints/generate/structured_decisions/protocol.py new file mode 100644 index 000000000000..3a193e203bad --- /dev/null +++ b/vllm/entrypoints/generate/structured_decisions/protocol.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Request and response models of the structured decisions API (/v1/systemone).""" + +import time +from typing import Any, Literal + +from pydantic import Field + +from vllm.config import ModelConfig +from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption +from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel +from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs +from vllm.utils import random_uuid + + +class QuestionSpec(OpenAIBaseModel): + type: str + instructions: Any = "" + criteria: Any = None + + +class StructuredDecisionRequest(OpenAIBaseModel): + model: str | None = None + state: Any = Field(..., description="What the questions are about: text or JSON.") + questions: dict[str, QuestionSpec] = Field( + ..., description="Question id to question, asked in this order." + ) + instructions: str | None = Field( + default=None, description="Context placed ahead of the questions." + ) + chat_template_kwargs: dict[str, Any] | None = None + priority: int = Field(default=0, ge=-(2**63), le=2**63 - 1) + request_id: str = Field(default_factory=random_uuid) + + +class DecisionUsage(OpenAIBaseModel): + input_tokens: int + output_tokens: int + + +class QuestionDiagnostics(OpenAIBaseModel): + label_mass: float = Field( + description="Probability the model put on the labels, over the full " + "vocabulary. A low value means most of the probability went to tokens " + "that are not labels." + ) + argmax_is_label: bool + + +class StructuredDecisionResponse(OpenAIBaseModel): + id: str + object: Literal["structured_decision"] = "structured_decision" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + answers: dict[str, dict[str, Any]] + usage: DecisionUsage + diagnostics: dict[str, QuestionDiagnostics] + + +class ReadPromptRequest(OpenAIBaseModel): + """Chat options for one read's prompt: the state and the question in the + user turn, ending at the generation prompt so the label is the reply's + first token. Thinking is off unless the request turns it on, or the label + would follow a thought rather than start the reply.""" + + chat_template_kwargs: dict[str, Any] | None = None + + def build_chat_params( + self, + default_template: str | None, + default_template_content_format: ChatTemplateContentFormatOption, + ) -> ChatParams: + return ChatParams( + chat_template=default_template, + chat_template_content_format=default_template_content_format, + chat_template_kwargs=merge_kwargs( + merge_kwargs({"enable_thinking": False}, self.chat_template_kwargs), + dict(add_generation_prompt=True, continue_final_message=False), + ), + ) + + def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: + return TokenizeParams( + max_total_tokens=model_config.max_model_len, + max_output_tokens=1, + add_special_tokens=False, + ) diff --git a/vllm/entrypoints/generate/structured_decisions/question_types.py b/vllm/entrypoints/generate/structured_decisions/question_types.py new file mode 100644 index 000000000000..a9a6514a0373 --- /dev/null +++ b/vllm/entrypoints/generate/structured_decisions/question_types.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Question types for structured decisions. + +A question type turns a request's criteria into options, writes the question +for the model with a label per option, and builds the answer from the label +probabilities. +""" + +import math +import string +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any, ClassVar + + +class StructuredDecisionError(ValueError): + """An invalid request. Returned as a 400.""" + + +#: Options are labeled A to Z in order. A single capital letter is one token +#: at the start of a reply for the tokenizers tested; every read checks it. +LABELS = tuple(string.ascii_uppercase) + + +@dataclass(frozen=True) +class Option: + name: str + description: str | None = None + + +@dataclass(frozen=True) +class Question: + id: str + type: "QuestionType" + instructions: str + options: tuple[Option, ...] + labels: tuple[str, ...] + + +class QuestionType(ABC): + name: ClassVar[str] + + @abstractmethod + def parse_options(self, qid: str, criteria: Any) -> list[Option]: ... + + @abstractmethod + def answer( + self, question: Question, probs: list[float], label_mass: float + ) -> dict[str, Any]: + """The answer for ``question``. ``probs[i]`` is the probability of + ``question.labels[i]`` among the labels, and the list sums to 1. + ``label_mass`` is the labels' total probability over the vocabulary.""" + + def prompt(self, question: Question) -> str: + """The question as the model reads it, after the state.""" + lines = [f"Question: {question.instructions}"] if question.instructions else [] + for label, o in zip(question.labels, question.options): + lines.append( + f"{label}: {o.name} - {o.description}" + if o.description + else f"{label}: {o.name}" + ) + lines.append("Answer with the letter of one option only.") + return "\n".join(lines) + + +QUESTION_TYPES: dict[str, QuestionType] = {} + + +def register_question_type(cls: type[QuestionType]) -> type[QuestionType]: + if cls.name in QUESTION_TYPES: + raise ValueError(f"question type {cls.name!r} is already registered") + QUESTION_TYPES[cls.name] = cls() + return cls + + +def get_question_type(name: str) -> QuestionType: + try: + return QUESTION_TYPES[name] + except KeyError: + raise StructuredDecisionError( + f"unknown question type {name!r}; supported: {sorted(QUESTION_TYPES)}" + ) from None + + +def build_question( + qid: str, + type_name: str, + instructions: Any, + criteria: Any, + max_options: int, +) -> Question: + if not qid: + raise StructuredDecisionError("question ids must be non-empty") + qtype = get_question_type(type_name) + options = qtype.parse_options(qid, criteria) + if not options: + raise StructuredDecisionError(f"question {qid!r}: needs at least one option") + names = [o.name for o in options] + if len(set(names)) != len(names): + raise StructuredDecisionError(f"question {qid!r}: duplicate option names") + limit = min(max_options, len(LABELS)) + if len(options) > limit: + raise StructuredDecisionError( + f"question {qid!r}: at most {limit} options for this model" + ) + if not isinstance(instructions, str): + instructions = "" if instructions is None else str(instructions) + return Question( + id=qid, + type=qtype, + instructions=instructions, + options=tuple(options), + labels=LABELS[: len(options)], + ) + + +def argmax(values: list[float]) -> int: + return max(range(len(values)), key=values.__getitem__) + + +def label_softmax(logprobs: list[float]) -> list[float]: + """Softmax of the label logprobs.""" + top = max(logprobs) + weights = [math.exp(lp - top) for lp in logprobs] + total = sum(weights) + return [w / total for w in weights] + + +@register_question_type +class ChoiceQuestion(QuestionType): + """Pick one option. ``criteria`` maps each option name to a description, + or to null.""" + + name = "choice" + + def parse_options(self, qid: str, criteria: Any) -> list[Option]: + if not isinstance(criteria, dict) or not criteria: + raise StructuredDecisionError( + f"question {qid!r}: choice criteria must map option names to " + "a description or null" + ) + return [ + Option(str(name), None if desc is None else str(desc)) + for name, desc in criteria.items() + ] + + def answer( + self, question: Question, probs: list[float], label_mass: float + ) -> dict[str, Any]: + top = argmax(probs) + return { + "type": self.name, + "choice": question.options[top].name, + "probabilities": {a.name: p for a, p in zip(question.options, probs)}, + # The chosen label's probability over the whole vocabulary, so a + # read that mostly wanted a non-label token reports low confidence. + "confidence": probs[top] * label_mass, + } diff --git a/vllm/entrypoints/generate/structured_decisions/serving.py b/vllm/entrypoints/generate/structured_decisions/serving.py new file mode 100644 index 000000000000..7e2136b8a25e --- /dev/null +++ b/vllm/entrypoints/generate/structured_decisions/serving.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import asyncio +import json +from typing import Any + +from fastapi import Request + +from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.serve.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.engine.serving import BaseServing +from vllm.entrypoints.serve.utils.request_logger import RequestLogger +from vllm.logger import init_logger + +from .protocol import ( + DecisionUsage, + QuestionDiagnostics, + StructuredDecisionRequest, + StructuredDecisionResponse, +) +from .question_types import Question, StructuredDecisionError, build_question +from .strategies import DecisionLimits, ReadStrategy + +logger = init_logger(__name__) + + +def state_text(state: Any) -> str: + return state if isinstance(state, str) else json.dumps(state) + + +def parse_questions( + request: StructuredDecisionRequest, + limits: DecisionLimits, +) -> list[Question]: + if not request.questions: + raise StructuredDecisionError("questions: needs at least one question") + if len(request.questions) > limits.max_questions: + raise StructuredDecisionError( + f"questions: at most {limits.max_questions} for this model" + ) + questions = [] + for qid, spec in request.questions.items(): + if spec.model_extra: + raise StructuredDecisionError( + f"question {qid!r}: unknown field(s) {sorted(spec.model_extra)}" + ) + questions.append( + build_question( + qid, + spec.type, + spec.instructions, + spec.criteria, + limits.max_options, + ) + ) + return questions + + +class ServingStructuredDecisions(BaseServing): + def __init__( + self, + models: OpenAIServingModels, + strategy: ReadStrategy, + *, + request_logger: RequestLogger | None = None, + ) -> None: + model_config = strategy.context.engine_client.model_config + super().__init__( + models=models, model_config=model_config, request_logger=request_logger + ) + self.strategy = strategy + self.limits = strategy.limits() + + async def create_decision( + self, + request: StructuredDecisionRequest, + raw_request: Request | None = None, + ) -> StructuredDecisionResponse | ErrorResponse: + if (error := await self._check_model(request)) is not None: # type: ignore[arg-type] + return error + engine_client = self.strategy.context.engine_client + if engine_client.errored: + raise engine_client.dead_error + + base_id = self._base_request_id(raw_request, default=request.request_id) + request_id = f"decision-{base_id}" + try: + questions = parse_questions(request, self.limits) + lora_request = self._maybe_get_adapters(request) # type: ignore[arg-type] + engine_client.check_admission(len(questions)) + reads = await self.strategy.read( + questions, + request.instructions, + state_text(request.state), + request_id=request_id, + chat_template_kwargs=request.chat_template_kwargs, + lora_request=lora_request, + priority=request.priority, + ) + except StructuredDecisionError as e: + return self.create_error_response(e) + except asyncio.CancelledError: + return self.create_error_response("Client disconnected") + + answers: dict[str, dict[str, Any]] = {} + diagnostics: dict[str, QuestionDiagnostics] = {} + for q, read in zip(questions, reads): + answers[q.id] = q.type.answer(q, read.probs, read.label_mass) + diagnostics[q.id] = QuestionDiagnostics( + label_mass=read.label_mass, argmax_is_label=read.argmax_is_label + ) + return StructuredDecisionResponse( + id=request_id, + model=self.models.model_name(lora_request), + answers=answers, + usage=DecisionUsage( + input_tokens=sum(r.input_tokens for r in reads), + output_tokens=sum(r.output_tokens for r in reads), + ), + diagnostics=diagnostics, + ) diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py new file mode 100644 index 000000000000..22233ff91b0e --- /dev/null +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Read strategies: how a model is asked for label probabilities. + +Only models in NEXT_TOKEN_ARCHITECTURES are currently supported. +""" + +import math +from abc import ABC, abstractmethod +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +from vllm.config import ModelConfig +from vllm.engine.protocol import EngineClient +from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption +from vllm.entrypoints.generate.label_reads import next_token_label_reads +from vllm.inputs import EngineInput +from vllm.lora.request import LoRARequest +from vllm.renderers.inputs.preprocess import extract_prompt_components +from vllm.renderers.online_renderer import OnlineRenderer +from vllm.sampling_params import SamplingParams +from vllm.tokenizers import TokenizerLike + +from .protocol import ReadPromptRequest +from .question_types import LABELS, Question, StructuredDecisionError, label_softmax + + +@dataclass +class ReadContext: + engine_client: EngineClient + online_renderer: OnlineRenderer + chat_template: str | None + chat_template_content_format: ChatTemplateContentFormatOption + default_chat_template_kwargs: dict[str, Any] + + +@dataclass(frozen=True) +class DecisionLimits: + max_questions: int + max_options: int + + +@dataclass +class QuestionRead: + probs: list[float] + label_mass: float + argmax_is_label: bool + input_tokens: int + output_tokens: int + + +class ReadStrategy(ABC): + def __init__(self, context: ReadContext): + self.context = context + + @abstractmethod + def limits(self) -> DecisionLimits: ... + + @abstractmethod + async def read( + self, + questions: list[Question], + instructions: str | None, + state: str, + *, + request_id: str, + chat_template_kwargs: dict[str, Any] | None, + lora_request: LoRARequest | None, + priority: int, + ) -> list[QuestionRead]: + """One read per question, in order.""" + + +def reply_label_ids( + tokenizer: TokenizerLike, prompt_ids: Sequence[int] +) -> tuple[list[int], list[int]]: + """The prompt's ids after its last added token, and the token each of LABELS + adds after them as the first token of the reply. Raises ValueError if a + label is not one distinct token there.""" + added = set(tokenizer.get_added_vocab().values()) + start = max((i + 1 for i, t in enumerate(prompt_ids) if t in added), default=0) + tail = list(prompt_ids[start:]) + text = tokenizer.decode(tail) + ids: list[int] = [] + for label in LABELS: + extended = tokenizer.encode(text + label, add_special_tokens=False) + if extended[:-1] != tail or extended[-1] in ids: + raise ValueError( + f"label {label!r} is not one distinct token after this model's " + "chat prompt" + ) + ids.append(extended[-1]) + return tail, ids + + +class NextTokenStrategy(ReadStrategy): + """Autoregressive models. Each question is one request: the state, then the + question with its labeled options, and the logprobs of the label tokens as + the reply's first token. The requests share the state. With prefix caching, + it is prefilled once.""" + + def __init__(self, context: ReadContext): + super().__init__(context) + tokenizer = context.online_renderer.renderer.get_tokenizer() + probe = tokenizer.apply_chat_template( + [{"role": "user", "content": "x"}], + **{ + **context.default_chat_template_kwargs, + "chat_template": context.chat_template, + "add_generation_prompt": True, + "tokenize": True, + "return_dict": False, + "enable_thinking": False, + }, + ) + if isinstance(probe, str): + probe = tokenizer.encode(probe, add_special_tokens=False) + # Every prompt ends with the same generation prompt, so the labels' + # tokens are the same for every question. + self.tail, self.label_ids = reply_label_ids(tokenizer, probe) + + def limits(self) -> DecisionLimits: + return DecisionLimits(max_questions=64, max_options=len(LABELS)) + + def _read_request( + self, chat_template_kwargs: dict[str, Any] | None + ) -> ReadPromptRequest: + if (chat_template_kwargs or {}).get("enable_thinking"): + raise StructuredDecisionError( + "a read is the reply's first token, so thinking must be off" + ) + return ReadPromptRequest(chat_template_kwargs=chat_template_kwargs) + + async def _render( + self, read_request: ReadPromptRequest, messages: list[dict[str, Any]] + ) -> tuple[EngineInput, list[int]]: + ctx = self.context + _, (engine_input,) = await ctx.online_renderer.preprocess_chat( + read_request, + messages, + default_template=ctx.chat_template, + default_template_content_format=ctx.chat_template_content_format, + default_template_kwargs=ctx.default_chat_template_kwargs, + ) + prompt_ids = extract_prompt_components( + ctx.engine_client.model_config, engine_input + ).token_ids + return engine_input, list(prompt_ids or []) + + async def read( + self, + questions: list[Question], + instructions: str | None, + state: str, + *, + request_id: str, + chat_template_kwargs: dict[str, Any] | None, + lora_request: LoRARequest | None, + priority: int, + ) -> list[QuestionRead]: + ctx = self.context + read_request = self._read_request(chat_template_kwargs) + + slots, engine_inputs = [], [] + for q in questions: + messages = [{"role": "user", "content": f"{state}\n\n{q.type.prompt(q)}"}] + if instructions: + messages.insert(0, {"role": "system", "content": instructions}) + engine_input, prompt_ids = await self._render(read_request, messages) + if prompt_ids[len(prompt_ids) - len(self.tail) :] != self.tail: + raise StructuredDecisionError( + "these chat options end the prompt differently, so the " + "labels' tokens are unknown" + ) + slots.append(self.label_ids[: len(q.labels)]) + engine_inputs.append(engine_input) + + label_reads = await next_token_label_reads( + ctx.engine_client, + engine_inputs, + [ + SamplingParams(max_tokens=1, temperature=0.0, logprob_token_ids=ids) + for ids in slots + ], + request_id, + lora_request=lora_request, + priority=priority, + ) + + reads = [] + for ids, label_read in zip(slots, label_reads): + output = label_read.result.outputs[0] + reads.append( + QuestionRead( + probs=label_softmax(label_read.logprobs), + label_mass=sum(math.exp(lp) for lp in label_read.logprobs), + argmax_is_label=bool(output.token_ids) + and output.token_ids[0] in ids, + input_tokens=len(label_read.result.prompt_token_ids or ()), + output_tokens=len(output.token_ids), + ) + ) + return reads + + +NEXT_TOKEN_ARCHITECTURES = frozenset( + { + "Qwen3ForCausalLM", + "Qwen3_5ForConditionalGeneration", + "Qwen3_5MoeForConditionalGeneration", + } +) + + +# label_mass sums the labels' full-vocabulary probabilities. The logprobs modes +# return those probabilities. The logits modes return raw logits instead. +LOGPROBS_MODES = frozenset({"raw_logprobs", "processed_logprobs"}) + + +def select_read_strategy(model_config: ModelConfig) -> type[ReadStrategy]: + """Raises ValueError when the model cannot serve structured decisions. The + server opted in with --enable-structured-decisions, so startup fails.""" + if model_config.architecture not in NEXT_TOKEN_ARCHITECTURES: + raise ValueError( + "--enable-structured-decisions does not support " + f"{model_config.architecture}. Supported architectures: " + f"{sorted(NEXT_TOKEN_ARCHITECTURES)}" + ) + if model_config.logprobs_mode not in LOGPROBS_MODES: + raise ValueError( + "--enable-structured-decisions needs --logprobs-mode raw_logprobs or " + f"processed_logprobs, not {model_config.logprobs_mode}" + ) + return NextTokenStrategy diff --git a/vllm/entrypoints/launchers/cli_args.py b/vllm/entrypoints/launchers/cli_args.py index 28fe79fd691e..18f8cb00b503 100644 --- a/vllm/entrypoints/launchers/cli_args.py +++ b/vllm/entrypoints/launchers/cli_args.py @@ -163,6 +163,8 @@ class BaseFrontendArgs: enable_tokenizer_info_endpoint: bool = False """Enable the `/tokenizer_info` endpoint. May expose chat templates and other tokenizer configuration.""" + enable_structured_decisions: bool = False + """Enable the `/v1/systemone` structured decisions endpoint.""" enable_log_outputs: bool = False """If set to True, log model outputs (generations). Requires `--enable-log-requests`. Output text and finish reasons are logged at INFO, From d80ddaac73217770af1548f43e0cd7e68d4ec33d Mon Sep 17 00:00:00 2001 From: kliuae <17350011+kliuae@users.noreply.github.com> Date: Thu, 8 Oct 2026 22:55:20 +0800 Subject: [PATCH 101/624] [ROCm][Perf] Kimi-K3 Store only the current rank's shards in latent MoE up-proj (#59591) Signed-off-by: kliuae --- .../kimi_k3/test_amd_latent_moe_runner.py | 96 +++++++++++++++---- vllm/models/kimi_k3/amd/latent_moe_runner.py | 16 +++- vllm/models/kimi_k3/amd/linear.py | 41 ++++++-- 3 files changed, 125 insertions(+), 28 deletions(-) diff --git a/tests/models/kimi_k3/test_amd_latent_moe_runner.py b/tests/models/kimi_k3/test_amd_latent_moe_runner.py index fd5bcfd22b78..e47b5fc08e87 100644 --- a/tests/models/kimi_k3/test_amd_latent_moe_runner.py +++ b/tests/models/kimi_k3/test_amd_latent_moe_runner.py @@ -22,7 +22,7 @@ from vllm.distributed import get_tp_group from vllm.model_executor.layers.fused_moe.runner import moe_runner from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.linear import ColumnParallelLinear, ReplicatedLinear from vllm.models.kimi_k3.amd import latent_moe_runner from vllm.models.kimi_k3.amd.latent_moe_runner import ROCmLatentMoERunner from vllm.models.kimi_k3.amd.linear import KimiRoutedOutputTransform @@ -40,20 +40,42 @@ DTYPE = torch.bfloat16 -def _build_transform(device: torch.device) -> KimiRoutedOutputTransform: +def _build_transform( + device: torch.device, *, row_sharded: bool = False +) -> tuple[KimiRoutedOutputTransform, torch.Tensor]: norm = RMSNorm(LATENT_SIZE, eps=EPS).to(device=device, dtype=DTYPE) - up_proj = ReplicatedLinear( - LATENT_SIZE, - HIDDEN_SIZE, - bias=False, - params_dtype=DTYPE, - prefix="routed_expert_up_proj", - ).to(device=device) torch.manual_seed(0) norm.weight.data.copy_(1 + 0.1 * torch.randn_like(norm.weight)) - up_proj.weight.data.copy_(torch.randn_like(up_proj.weight) / LATENT_SIZE**0.5) - return KimiRoutedOutputTransform(norm, up_proj) + full_weight = ( + torch.randn(HIDDEN_SIZE, LATENT_SIZE, device=device, dtype=DTYPE) + / LATENT_SIZE**0.5 + ) + + up_proj: ReplicatedLinear | ColumnParallelLinear + if row_sharded: + up_proj = ColumnParallelLinear( + LATENT_SIZE, + HIDDEN_SIZE, + bias=False, + gather_output=True, + params_dtype=DTYPE, + prefix="routed_expert_up_proj", + ).to(device=device) + up_proj.weight_loader(up_proj.weight, full_weight) + else: + up_proj = ReplicatedLinear( + LATENT_SIZE, + HIDDEN_SIZE, + bias=False, + params_dtype=DTYPE, + prefix="routed_expert_up_proj", + ).to(device=device) + up_proj.weight.data.copy_(full_weight) + return ( + KimiRoutedOutputTransform(norm, up_proj, row_sharded=row_sharded), + full_weight, + ) def _tail_runner( @@ -72,6 +94,7 @@ def _tail_runner( runner = object.__new__(ROCmLatentMoERunner) attrs = { "routed_output_transform": transform, + "_up_proj_preshard": bool(getattr(transform, "row_sharded", False)), "_up_proj_shard_size": HIDDEN_SIZE // tp_world, "_logged_sharded_tail": False, "moe_config": SimpleNamespace( @@ -101,9 +124,13 @@ def _rank_partials( def _check_matches_replicated( - device: torch.device, tp_world: int, rank: int, use_ep: bool = False + device: torch.device, + tp_world: int, + rank: int, + use_ep: bool = False, + row_sharded: bool = False, ) -> None: - transform = _build_transform(device) + transform, up_weight = _build_transform(device, row_sharded=row_sharded) runner = _tail_runner(transform, tp_world, use_ep) group = get_tp_group().device_group @@ -118,7 +145,7 @@ def _check_matches_replicated( transform.norm.weight, EPS, ), - transform.up_proj.weight, + up_weight, ) expected.add_(_all_reduced(shared_output, group)) @@ -136,12 +163,37 @@ def _check_matches_replicated_under_ep( _check_matches_replicated(device, tp_world, rank, use_ep=True) +def _check_presharded_matches_replicated( + device: torch.device, tp_world: int, rank: int +) -> None: + _check_matches_replicated(device, tp_world, rank, row_sharded=True) + + +def _check_unfused_tail_gathers_full_hidden( + device: torch.device, tp_world: int, rank: int +) -> None: + transform, up_weight = _build_transform(device, row_sharded=True) + assert transform.up_proj.weight.shape[0] == HIDDEN_SIZE // tp_world + + # all-gather reconstructs the full projection when all ranks feed the same + # input, so use one seed for all ranks. + torch.manual_seed(1234) + latent = torch.randn(4, LATENT_SIZE, device=device, dtype=DTYPE).mul_(0.01) + + torch.testing.assert_close( + transform(latent), + F.linear(transform.norm(latent), up_weight), + atol=8e-2, + rtol=3e-2, + ) + + def _check_writes_only_its_own_shard( device: torch.device, tp_world: int, rank: int ) -> None: """Two ranks that swapped both weight rows and write offsets still sum to the right total, so inspect each slice before the final collective.""" - transform = _build_transform(device) + transform, up_weight = _build_transform(device) runner = _tail_runner(transform, tp_world) group = get_tp_group().device_group @@ -169,7 +221,7 @@ def _capture(states, trunc_size, output_is_reduced=None): shard = HIDDEN_SIZE // tp_world start, end = rank * shard, (rank + 1) * shard local = captured["states"] - projected = F.linear(latent, transform.up_proj.weight[start:end]) + projected = F.linear(latent, up_weight[start:end]) torch.testing.assert_close( local[:, start:end], before[:, start:end] + projected, atol=8e-2, rtol=3e-2 @@ -181,6 +233,8 @@ def _capture(states, trunc_size, output_is_reduced=None): _CHECKS = { "matches_replicated": _check_matches_replicated, "matches_replicated_ep": _check_matches_replicated_under_ep, + "presharded_matches_replicated": _check_presharded_matches_replicated, + "unfused_gathers": _check_unfused_tail_gathers_full_hidden, "own_shard_only": _check_writes_only_its_own_shard, } @@ -219,6 +273,16 @@ def test_sharded_tail_tp8_ep8_matches_replicated_projection() -> None: _run_ranks("matches_replicated_ep", 8) +@multi_gpu_test(num_gpus=4) +def test_presharded_tail_tp4_matches_replicated_projection() -> None: + _run_ranks("presharded_matches_replicated", 4) + + +@multi_gpu_test(num_gpus=4) +def test_presharded_unfused_tail_gathers_the_full_hidden_dim() -> None: + _run_ranks("unfused_gathers", 4) + + @multi_gpu_test(num_gpus=4) def test_sharded_tail_tp4_writes_only_its_own_shard() -> None: _run_ranks("own_shard_only", 4) diff --git a/vllm/models/kimi_k3/amd/latent_moe_runner.py b/vllm/models/kimi_k3/amd/latent_moe_runner.py index dbebfd418203..f6508d6b559f 100644 --- a/vllm/models/kimi_k3/amd/latent_moe_runner.py +++ b/vllm/models/kimi_k3/amd/latent_moe_runner.py @@ -36,18 +36,23 @@ def __init__( up_proj = getattr(transform, "up_proj", None) tp_size = get_tensor_model_parallel_world_size() + self._up_proj_preshard = bool(getattr(transform, "row_sharded", False)) self._up_proj_shard_size = 0 self._tail_shardable = ( up_proj is not None and tp_size > 1 - and up_proj.weight.shape[0] % tp_size == 0 + and (self._up_proj_preshard or up_proj.weight.shape[0] % tp_size == 0) and self._shared_experts is not None and not self.moe_config.is_sequence_parallel and self.routed_scaling_factor == 1.0 ) if self._tail_shardable: assert up_proj is not None - self._up_proj_shard_size = up_proj.weight.shape[0] // tp_size + self._up_proj_shard_size = ( + up_proj.weight.shape[0] + if self._up_proj_preshard + else up_proj.weight.shape[0] // tp_size + ) else: logger.warning_once( "K3 latent-MoE tail is not shardable under this config, " @@ -80,7 +85,12 @@ def _shard_up_proj_tail( shard_size = self._up_proj_shard_size shard_start = get_tensor_model_parallel_rank() * shard_size - up_proj_shard = transform.up_proj.weight.narrow(0, shard_start, shard_size) + weight = transform.up_proj.weight + up_proj_shard = ( + weight + if self._up_proj_preshard + else weight.narrow(0, shard_start, shard_size) + ) hidden_shard = shared_output.narrow(-1, shard_start, shard_size) # hidden_shard += latent @ up_proj_shard.T, accumulated in the GEMM's diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py index 3a1d827abe47..682773189352 100644 --- a/vllm/models/kimi_k3/amd/linear.py +++ b/vllm/models/kimi_k3/amd/linear.py @@ -137,11 +137,13 @@ class KimiRoutedOutputTransform(nn.Module): def __init__( self, norm: RMSNorm | None, - up_proj: ReplicatedLinear, + up_proj: ReplicatedLinear | ColumnParallelLinear, + row_sharded: bool = False, ) -> None: super().__init__() self.norm = norm self.up_proj = up_proj + self.row_sharded = row_sharded def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if self.norm is not None: @@ -265,7 +267,7 @@ def __init__( self.routed_expert_down_proj: ReplicatedLinear | None self.routed_expert_norm: RMSNorm | None - self.routed_expert_up_proj: ReplicatedLinear | None + self.routed_expert_up_proj: ReplicatedLinear | ColumnParallelLinear | None self.routed_output_transform: KimiRoutedOutputTransform | None if self.use_latent_moe: self.routed_expert_down_proj = ReplicatedLinear( @@ -280,15 +282,36 @@ def __init__( if self.latent_moe_use_norm else None ) - self.routed_expert_up_proj = ReplicatedLinear( - self.moe_hidden_size, - hidden_size, - bias=False, - quant_config=None, - prefix=f"{prefix}.routed_expert_up_proj", + shard_up_proj = ( + self.tp_size > 1 + and hidden_size % self.tp_size == 0 + and self.num_shared_experts is not None + and self.routed_scaling_factor == 1.0 ) + if shard_up_proj: + # gather_output is a fallback restoring the full hidden dim for + # the unfused tail. The fused tail reads from .weight directly + # so it is not triggered + self.routed_expert_up_proj = ColumnParallelLinear( + self.moe_hidden_size, + hidden_size, + bias=False, + gather_output=True, + quant_config=None, + prefix=f"{prefix}.routed_expert_up_proj", + ) + else: + self.routed_expert_up_proj = ReplicatedLinear( + self.moe_hidden_size, + hidden_size, + bias=False, + quant_config=None, + prefix=f"{prefix}.routed_expert_up_proj", + ) self.routed_output_transform = KimiRoutedOutputTransform( - self.routed_expert_norm, self.routed_expert_up_proj + self.routed_expert_norm, + self.routed_expert_up_proj, + row_sharded=shard_up_proj, ) else: self.routed_expert_down_proj = None From 994b8b77be74e65c869c1ca8979fefb17b162b11 Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 8 Oct 2026 17:06:39 +0200 Subject: [PATCH 102/624] [Bugfix][ROCm] RDNA3 W4A16 MoE: handle the N-first CT weight layout (#60388) Signed-off-by: JartX --- .../quantization/test_rdna3_moe_w4a16.py | 81 +++++++++++++++++++ .../layers/fused_moe/oracle/int_wna16.py | 17 ++-- .../compressed_tensors_moe_wna16.py | 6 +- 3 files changed, 95 insertions(+), 9 deletions(-) diff --git a/tests/kernels/quantization/test_rdna3_moe_w4a16.py b/tests/kernels/quantization/test_rdna3_moe_w4a16.py index 6212e6f08c68..58aa90aa07a9 100644 --- a/tests/kernels/quantization/test_rdna3_moe_w4a16.py +++ b/tests/kernels/quantization/test_rdna3_moe_w4a16.py @@ -18,6 +18,7 @@ import pytest import torch +import torch.nn.functional as F from vllm.platforms import current_platform @@ -32,6 +33,9 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( # noqa: E402 moe_align_block_size, ) +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( # noqa: E402 + _process_weights_rdna3, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402 pack_quantized_values_into_int32, ) @@ -95,6 +99,16 @@ def _make_qzeros(E, groups, N): return qz.unsqueeze(0).expand(E, -1, -1).contiguous() +def _make_ct_weights(E, N, K, group_size, dtype): + """Compressed-tensors canonical N-first layout: packed ``[E, N, K // 8]`` + and scales ``[E, N, K // group_size]``, plus the dequantized weights.""" + q = torch.randint(0, 16, (E, N, K), dtype=torch.int32, device=device) + s = torch.rand(E, N, K // group_size, dtype=dtype, device=device) * 0.1 + 0.01 + packed = pack_quantized_values_into_int32(q, scalar_types.uint4b8, packed_dim=2) + w = (q - 8).float() * s.float().repeat_interleave(group_size, dim=2) + return packed.contiguous(), s, w + + @gfx1100_only @pytest.mark.parametrize("E, K, N_inter, top_k, group_size", MODEL_CONFIGS) @pytest.mark.parametrize("M", NUM_TOKENS) @@ -359,3 +373,70 @@ def test_expert_id_minus_one(): # Output should remain zero (expert skipped) assert torch.equal(out, torch.zeros_like(out)) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_process_weights_from_compressed_tensors_layout(dtype): + """Weights in the layout compressed-tensors hands to the oracle go through + the RDNA3 post-processing and a full MoE forward, against a dequantized + reference. The tests above build the kernel layout directly, so they do + not notice when the layout upstream of the kernel changes (#52798).""" + E, K, N_inter, top_k, group_size, M = 8, 512, 256, 2, 32, 16 + torch.manual_seed(5) + w13, w13_s, w13_ref = _make_ct_weights(E, 2 * N_inter, K, group_size, dtype) + w2, w2_s, w2_ref = _make_ct_weights(E, K, N_inter, group_size, dtype) + w13, w2, w13_s, w2_s, w13_z, w2_z, *_ = _process_weights_rdna3( + w13, w2, w13_s, w2_s, group_size + ) + assert w13_z is not None and w2_z is not None + + x = torch.randn(M, K, dtype=dtype, device=device) + topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32) + topk_w = torch.softmax(torch.randn(M, top_k, device=device), dim=-1).float() + si, ei, ntp = moe_align_block_size(topk_ids, 1, E) + + w1_out = torch.zeros(M * top_k, 2 * N_inter, dtype=dtype, device=device) + ops.moe_gptq_gemm_rdna3( + x, + w1_out, + w13, + w13_s, + w13_z, + torch.empty(0, device=device), + si, + ei, + ntp, + top_k, + 1, + False, + 0, + ) + act_out = torch.empty(M * top_k, N_inter, dtype=dtype, device=device) + apply_moe_activation(MoEActivation.SILU, act_out, w1_out) + out = torch.zeros(M, K, dtype=dtype, device=device) + ops.moe_gptq_gemm_rdna3( + act_out, + out, + w2, + w2_s, + w2_z, + topk_w.view(-1), + si, + ei, + ntp, + 1, + 1, + True, + top_k, + ) + + ref = torch.zeros(M, K, device=device) + for m in range(M): + for k in range(top_k): + e = int(topk_ids[m, k]) + h = w13_ref[e] @ x[m].float() + a = F.silu(h[:N_inter]) * h[N_inter:] + ref[m] += (w2_ref[e] @ a) * topk_w[m, k] + rel_l2 = (torch.norm(out.float() - ref) / torch.norm(ref)).item() + assert rel_l2 < (0.02 if dtype == torch.float16 else 0.05), rel_l2 diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index 3e91fd9e041c..9b4b1633790f 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -926,15 +926,18 @@ def _process_weights_rdna3( ]: """RDNA3 (gfx1100) W4A16 weight post-processing. - Interleaves the packed nibbles per expert (the exllama shuffle the dense - RDNA3 kernel also uses) and synthesizes the symmetric zero points that - ``moe_gptq_gemm_rdna3`` dequantizes with. The packed layout - ``[E, K // 8, N]`` and the ``[E, groups, N]`` scales are already what the - kernel wants, so neither is repacked. + Transposes the canonical N-first inputs to the K-first layout of + ``moe_gptq_gemm_rdna3`` (weights ``[E, K // 8, N]``, scales + ``[E, groups, N]``), interleaves the packed nibbles per expert (the exllama + shuffle the dense RDNA3 kernel also uses) and synthesizes the symmetric + zero points the kernel dequantizes with. """ device = w13.device num_experts = w13.size(0) + w13 = w13.transpose(1, 2).contiguous() + w2 = w2.transpose(1, 2).contiguous() + for e in range(num_experts): w13_e = w13[e].contiguous() ops.gptq_shuffle(w13_e, 4) @@ -950,8 +953,8 @@ def _qzeros(w: torch.Tensor) -> torch.Tensor: return ( w13, w2, - w13_scale.contiguous(), - w2_scale.contiguous(), + w13_scale.transpose(1, 2).contiguous(), + w2_scale.transpose(1, 2).contiguous(), _qzeros(w13), _qzeros(w2), None, # w13_input_global_scale diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index fdd37d120354..1e06f474133a 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -450,10 +450,12 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if w2_bias is not None: replace_parameter(layer, "w2_bias", w2_bias) - # CPU fused_experts_cpu requires zero points even for symmetric quant. + # CPU fused_experts_cpu and the RDNA3 HIP kernel require zero points + # even for symmetric quant (the oracle synthesizes them). # EMULATION bakes ZP into the dequantized bf16 weights — ZP is None. if ( - not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU + not self.symmetric + or self.wna16_backend in (WNA16MoEBackend.CPU, WNA16MoEBackend.RDNA3) ) and self.wna16_backend != WNA16MoEBackend.EMULATION: assert w13_qzeros is not None and w2_qzeros is not None replace_parameter(layer, "w13_weight_zero_point", w13_qzeros) From fd928e5f3cf17bce1ba090040beec48f86c84ff6 Mon Sep 17 00:00:00 2001 From: yifanFengg Date: Thu, 8 Oct 2026 08:19:57 -0700 Subject: [PATCH 103/624] [Bugfix] Check batch-invariance support for every attention backend (#60271) Signed-off-by: yifanFengg Co-authored-by: Claude Fable 5.1 Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- vllm/v1/worker/gpu/attn_utils.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 4fa308bcf471..8aa9ba3cb8b2 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -8,6 +8,7 @@ import numpy as np import torch +import vllm.envs as envs from vllm.config import ( VllmConfig, get_layers_from_vllm_config, @@ -224,6 +225,14 @@ def init_attn_backend( for layer_name in layer_names: attn_backend = attn_layers[layer_name].get_attn_backend() + if ( + envs.VLLM_BATCH_INVARIANT + and not attn_backend.supports_batch_invariance() + ): + raise RuntimeError( + "VLLM batch_invariant mode is not supported for " + f"{attn_backend.get_name()}." + ) if layer_name in fast_prefill_eligible_layers: attn_backend = create_fast_prefill_custom_backend( "FastPrefill", attn_backend From 072d49bb6eb017a1e33bf3d016a979972438a90f Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Thu, 8 Oct 2026 23:25:21 +0800 Subject: [PATCH 104/624] [Platform] Support platform-specific MLA prefill backend selection (#56959) Signed-off-by: wangxiyuan Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../gen_files/generate_attention_backends.py | 96 +++--- .../v1/attention/test_mla_prefill_selector.py | 297 ++++++++++++------ vllm/platforms/cpu.py | 20 ++ vllm/platforms/interface.py | 84 +++++ vllm/platforms/rocm.py | 55 ++++ .../backends/mla/prefill/aiter_flash_attn.py | 5 +- .../v1/attention/backends/mla/prefill/base.py | 18 +- .../backends/mla/prefill/flashinfer.py | 7 +- .../backends/mla/prefill/selector.py | 142 ++------- .../backends/mla/prefill/tokenspeed_mla.py | 7 +- .../backends/mla/prefill/trtllm_ragged.py | 7 +- 11 files changed, 462 insertions(+), 276 deletions(-) diff --git a/docs/mkdocs/gen_files/generate_attention_backends.py b/docs/mkdocs/gen_files/generate_attention_backends.py index 725d7072ca92..135dcf3a35a2 100644 --- a/docs/mkdocs/gen_files/generate_attention_backends.py +++ b/docs/mkdocs/gen_files/generate_attention_backends.py @@ -51,7 +51,7 @@ MLA_PREFILL_DIR = BACKENDS_DIR / "mla" / "prefill" MLA_PREFILL_REGISTRY_FILE = MLA_PREFILL_DIR / "registry.py" -MLA_PREFILL_SELECTOR_FILE = MLA_PREFILL_DIR / "selector.py" +PLATFORM_INTERFACE_FILE = REPO_ROOT / "vllm" / "platforms" / "interface.py" # --------------------------------------------------------------------------- @@ -320,67 +320,83 @@ def parse_mla_prefill_registry() -> dict[str, str]: def parse_mla_prefill_priorities() -> dict[str, list[str]]: - """Parse MLA prefill backend priorities from selector.py. + """Parse MLA prefill backend priorities from platforms/interface.py. Returns: A dict with keys like 'blackwell' and 'default' containing lists of backend enum names in priority order. """ - if not MLA_PREFILL_SELECTOR_FILE.exists(): + if not PLATFORM_INTERFACE_FILE.exists(): return {} try: - tree = ast.parse(MLA_PREFILL_SELECTOR_FILE.read_text()) + tree = ast.parse(PLATFORM_INTERFACE_FILE.read_text()) except Exception: return {} + def extract_backends(stmt: ast.AST) -> list[str] | None: + """Read a backend enum list from a return/assignment statement.""" + if isinstance(stmt, ast.Return) or ( + isinstance(stmt, ast.Assign) + and len(stmt.targets) == 1 + and isinstance(stmt.targets[0], ast.Name) + ): + value = stmt.value + else: + return None + if not isinstance(value, ast.List): + return None + return [elt.attr for elt in value.elts if isinstance(elt, ast.Attribute)] + + def extract_last_list(nodes: list[ast.stmt], recurse: bool) -> list[str] | None: + """Return the last backend list found among the given statements.""" + result = None + for stmt in nodes: + backends = extract_backends(stmt) + if backends is not None: + result = backends + elif recurse and isinstance(stmt, ast.If): + nested = extract_last_list([*stmt.body, *stmt.orelse], True) + if nested is not None: + result = nested + return result + + def is_blackwell_check(stmt: ast.AST) -> bool: + return ( + isinstance(stmt, ast.If) + and isinstance(stmt.test, ast.Compare) + and isinstance(stmt.test.left, ast.Attribute) + and stmt.test.left.attr == "major" + and stmt.test.comparators + and isinstance(stmt.test.comparators[0], ast.Constant) + and stmt.test.comparators[0].value == 10 + ) + priorities: dict[str, list[str]] = {} for node in ast.walk(tree): - if not isinstance(node, ast.FunctionDef): + if not isinstance(node, ast.ClassDef) or node.name != "Platform": continue - if node.name != "_get_mla_prefill_backend_priorities": + + method = find_method(node, "get_mla_prefill_backend_cls") + if method is None: continue # Look for if statements checking device_capability.major - for stmt in ast.walk(node): + for stmt in ast.walk(method): if not isinstance(stmt, ast.If): continue - # Check if it's a capability.major == 10 check (Blackwell) - is_blackwell = ( - isinstance(stmt.test, ast.Compare) - and isinstance(stmt.test.left, ast.Attribute) - and stmt.test.left.attr == "major" - and stmt.test.comparators - and isinstance(stmt.test.comparators[0], ast.Constant) - and stmt.test.comparators[0].value == 10 - ) - - # Extract backends from return statements - for body_stmt in stmt.body: - if isinstance(body_stmt, ast.Return) and isinstance( - body_stmt.value, ast.List - ): - backends = [] - for elt in body_stmt.value.elts: - if isinstance(elt, ast.Attribute): - backends.append(elt.attr) - if is_blackwell: - priorities["blackwell"] = backends - else: - priorities["default"] = backends - - # Extract from else branch - for else_stmt in stmt.orelse: - if isinstance(else_stmt, ast.Return) and isinstance( - else_stmt.value, ast.List - ): - backends = [] - for elt in else_stmt.value.elts: - if isinstance(elt, ast.Attribute): - backends.append(elt.attr) + if is_blackwell_check(stmt): + # The Blackwell lists live in the nested if/else inside the + # major == 10 branch. + backends = extract_last_list(stmt.body, True) + if backends is not None: + priorities["blackwell"] = backends + else: + backends = extract_last_list([*stmt.body, *stmt.orelse], False) + if backends is not None: priorities["default"] = backends return priorities diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index d5b81695aa41..3c6e1a929571 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for MLA prefill backend selector.""" +from contextlib import suppress from unittest.mock import MagicMock, patch import pytest @@ -9,13 +10,12 @@ from vllm.config import AttentionConfig, ModelConfig, VllmConfig from vllm.platforms import current_platform -from vllm.platforms.interface import DeviceCapability +from vllm.platforms.interface import DeviceCapability, Platform from vllm.v1.attention.backends.mla.prefill.base import MLADimensions from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum from vllm.v1.attention.backends.mla.prefill.selector import ( MLAPrefillSelectorConfig, _auto_select_mla_prefill_backend, - _get_mla_prefill_backend_priorities, get_mla_prefill_backend, ) @@ -55,51 +55,170 @@ def _make_vllm_config( return mock_vllm_config -class TestGetMLAPrefillBackend: - """Tests for get_mla_prefill_backend (public API).""" +class TestPlatformMLAPrefillHooks: + """Tests for the Platform.get_mla_prefill_backend_cls hook itself.""" - def test_cpu_uses_sdpa_prefill(self): - vllm_config = _make_vllm_config() - zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class() + def _get_tried_backends(self, capability, mla_dimensions): + """Run the default hook, returning the backend order tried.""" + selector_config = MLAPrefillSelectorConfig( + dtype=torch.bfloat16, + mla_dimensions=mla_dimensions, + ) + tried: list[MLAPrefillBackendEnum] = [] + + def mock_get_class(backend_enum): + tried.append(backend_enum) + cls = MagicMock() + cls.validate_configuration.return_value = ["not available"] + return cls - # A Zen CPU with zentorch selects the Zen backend, so expect whichever - # one this host offers. - expected = ( - zen_cls if zen_cls.is_available() else MLAPrefillBackendEnum.CPU.get_class() + with ( + patch.object(Platform, "get_device_capability", return_value=capability), + patch.object(MLAPrefillBackendEnum, "get_class", mock_get_class), + suppress(ValueError), + ): + Platform.get_mla_prefill_backend_cls(selector_config) + return tried + + def test_priorities_no_capability(self): + tried = self._get_tried_backends( + None, + MLADimensions(qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128), ) + assert tried == [MLAPrefillBackendEnum.FLASH_ATTN] - with patch("vllm.platforms.current_platform") as mock_platform: - mock_platform.is_cpu.return_value = True + def test_priorities_blackwell_glm_dimensions(self): + tried = self._get_tried_backends( + DeviceCapability(major=10, minor=0), + MLADimensions(qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256), + ) + assert tried == [ + MLAPrefillBackendEnum.TRTLLM_RAGGED, + MLAPrefillBackendEnum.FLASH_ATTN, + MLAPrefillBackendEnum.FLASHINFER, + MLAPrefillBackendEnum.TOKENSPEED_MLA, + ] - backend = get_mla_prefill_backend(vllm_config) - assert backend is expected + def test_priorities_blackwell_deepseek_dimensions(self): + tried = self._get_tried_backends( + DeviceCapability(major=10, minor=0), + MLADimensions(qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128), + ) + assert tried == [ + MLAPrefillBackendEnum.FLASH_ATTN, + MLAPrefillBackendEnum.TRTLLM_RAGGED, + MLAPrefillBackendEnum.FLASHINFER, + MLAPrefillBackendEnum.TOKENSPEED_MLA, + ] - def test_zen_cpu_prefers_zentorch_prefill(self): - vllm_config = _make_vllm_config() - zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class() + def test_priorities_hopper(self): + tried = self._get_tried_backends( + DeviceCapability(major=9, minor=0), + MLADimensions(qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128), + ) + assert tried == [MLAPrefillBackendEnum.FLASH_ATTN] + + def test_skips_import_error_backends(self): + selector_config = MLAPrefillSelectorConfig( + dtype=torch.bfloat16, + mla_dimensions=MLADimensions( + qk_nope_head_dim=192, + qk_rope_head_dim=64, + v_head_dim=256, + ), + ) + valid = MagicMock() + valid.validate_configuration.return_value = [] + + def mock_get_class(backend_enum): + if backend_enum is MLAPrefillBackendEnum.FLASH_ATTN: + return valid + raise ImportError("not installed") with ( - patch("vllm.platforms.current_platform") as mock_platform, - patch.object(zen_cls, "is_available", return_value=True), + patch.object( + Platform, + "get_device_capability", + return_value=DeviceCapability(major=10, minor=0), + ), + patch.object(MLAPrefillBackendEnum, "get_class", mock_get_class), ): - mock_platform.is_cpu.return_value = True + selected = Platform.get_mla_prefill_backend_cls(selector_config) + assert selected is valid - backend = get_mla_prefill_backend(vllm_config) - assert backend is zen_cls + def test_all_invalid_raises_error(self): + selector_config = MLAPrefillSelectorConfig(dtype=torch.bfloat16) - def test_zen_cpu_falls_back_when_zentorch_unavailable(self): - vllm_config = _make_vllm_config() - zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class() + invalid = MagicMock() + invalid.validate_configuration.return_value = ["not available"] with ( - patch("vllm.platforms.current_platform") as mock_platform, - patch.object(zen_cls, "is_available", return_value=False), + patch.object( + Platform, + "get_device_capability", + return_value=DeviceCapability(major=10, minor=0), + ), + patch.object(MLAPrefillBackendEnum, "get_class", return_value=invalid), + pytest.raises(ValueError, match="No valid MLA"), ): - mock_platform.is_cpu.return_value = True + Platform.get_mla_prefill_backend_cls(selector_config) + + def test_cpu_platform_returns_fixed_backend(self): + """CPU backends are returned via the fixed-backend hook.""" + try: + cpu_cls = MLAPrefillBackendEnum.CPU.get_class() + zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class() + except ImportError: + pytest.skip("CPU MLA prefill backend not available") + return + + from vllm.platforms.cpu import CpuPlatform + + expected = zen_cls if zen_cls.is_available() else cpu_cls + assert ( + CpuPlatform.get_mla_prefill_backend_cls( + MLAPrefillSelectorConfig(dtype=torch.bfloat16) + ) + is expected + ) + + +class TestGetMLAPrefillBackend: + """Tests for get_mla_prefill_backend (public API).""" + + def test_cpu_uses_sdpa_prefill(self): + vllm_config = _make_vllm_config() + + with patch("vllm.platforms.current_platform") as mock_platform: + mock_platform.get_mla_prefill_backend_cls.return_value = ( + MLAPrefillBackendEnum.CPU.get_class() + ) backend = get_mla_prefill_backend(vllm_config) assert backend is MLAPrefillBackendEnum.CPU.get_class() + def test_zen_cpu_prefers_zentorch_prefill(self): + from vllm.platforms.cpu import CpuPlatform + + zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class() + + with patch.object(zen_cls, "is_available", return_value=True): + backend = CpuPlatform.get_mla_prefill_backend_cls( + MLAPrefillSelectorConfig(dtype=torch.bfloat16) + ) + assert backend is zen_cls + + def test_zen_cpu_falls_back_when_zentorch_unavailable(self): + from vllm.platforms.cpu import CpuPlatform + + zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class() + + with patch.object(zen_cls, "is_available", return_value=False): + backend = CpuPlatform.get_mla_prefill_backend_cls( + MLAPrefillSelectorConfig(dtype=torch.bfloat16) + ) + assert backend is MLAPrefillBackendEnum.CPU.get_class() + def test_no_device_capability_returns_flash_attn(self): vllm_config = _make_vllm_config() @@ -108,16 +227,8 @@ class FlashAttnBackend: def get_name(): return "FLASH_ATTN" - with ( - patch("vllm.platforms.current_platform") as mock_platform, - patch.object( - MLAPrefillBackendEnum.FLASH_ATTN, - "get_class", - return_value=FlashAttnBackend, - ), - ): - mock_platform.get_device_capability.return_value = None - mock_platform.is_cpu.return_value = False + with patch("vllm.platforms.current_platform") as mock_platform: + mock_platform.get_mla_prefill_backend_cls.return_value = FlashAttnBackend backend = get_mla_prefill_backend(vllm_config) assert backend.get_name() == "FLASH_ATTN" @@ -214,7 +325,7 @@ def test_auto_selection_on_hopper(self, qk_nope_head_dim: int, v_head_dim: int): mock_platform.get_device_capability.return_value = DeviceCapability( major=9, minor=0 ) - mock_platform.is_rocm.return_value = False + mock_platform.get_mla_prefill_backend_cls.return_value = flash_attn_cls backend = get_mla_prefill_backend(vllm_config) assert backend.get_name() == "FLASH_ATTN" @@ -224,7 +335,6 @@ class TestAutoSelectMLAPrefillBackend: """Tests for fallback and error paths in auto-selection.""" def test_blackwell_glm_dimensions_use_trtllm(self): - capability = DeviceCapability(major=10, minor=0) selector_config = MLAPrefillSelectorConfig( dtype=torch.bfloat16, mla_dimensions=MLADimensions( @@ -235,7 +345,6 @@ def test_blackwell_glm_dimensions_use_trtllm(self): ) try: - flash_attn_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class() trtllm_cls = MLAPrefillBackendEnum.TRTLLM_RAGGED.get_class() except ImportError: pytest.skip("MLA prefill backend not available") @@ -243,19 +352,13 @@ def test_blackwell_glm_dimensions_use_trtllm(self): with ( patch("vllm.platforms.current_platform") as mock_platform, - patch.object(flash_attn_cls, "is_available", return_value=True), patch.object(trtllm_cls, "validate_configuration", return_value=[]), ): - # Force the non-ROCm priority on the Blackwell. - mock_platform.is_rocm.return_value = False - backend = _auto_select_mla_prefill_backend( - capability, - selector_config, - ) + mock_platform.get_mla_prefill_backend_cls.return_value = trtllm_cls + backend = _auto_select_mla_prefill_backend(selector_config) assert backend.get_name() == "TRTLLM_RAGGED" - def test_all_fail_raises_error(self): - capability = DeviceCapability(major=10, minor=0) + def test_platform_hook_raises_propagates(self): selector_config = MLAPrefillSelectorConfig( dtype=torch.bfloat16, mla_dimensions=MLADimensions( @@ -265,18 +368,14 @@ def test_all_fail_raises_error(self): ), ) - def mock_get_class(backend_enum): # noqa: ARG001 - cls = MagicMock() - cls.validate_configuration.return_value = ["not available"] - return cls - - with patch.object(MLAPrefillBackendEnum, "get_class", mock_get_class): - _auto_select_mla_prefill_backend.cache_clear() - with pytest.raises(ValueError, match="No valid MLA"): - _auto_select_mla_prefill_backend( - capability, - selector_config, - ) + with ( + patch("vllm.platforms.current_platform") as mock_platform, + pytest.raises(ValueError, match="No valid MLA"), + ): + mock_platform.get_mla_prefill_backend_cls.side_effect = ValueError( + "No valid MLA prefill backend found" + ) + _auto_select_mla_prefill_backend(selector_config) class TestBackendValidation: @@ -391,18 +490,29 @@ class TestROCmAiterFAPrefillSelection: def test_rocm_priorities_prefer_aiter_fa(self): """On ROCm, ROCM_AITER_FA is tried first, FLASH_ATTN as fallback.""" - with patch("vllm.platforms.current_platform") as mock_platform: - mock_platform.is_rocm.return_value = True - priorities = _get_mla_prefill_backend_priorities( - DeviceCapability(major=9, minor=5), - MLADimensions( - qk_nope_head_dim=128, - qk_rope_head_dim=64, - v_head_dim=128, - ), - ) + from vllm.platforms.rocm import RocmPlatform - assert priorities == [ + selector_config = MLAPrefillSelectorConfig(dtype=torch.bfloat16) + tried: list[MLAPrefillBackendEnum] = [] + + def mock_get_class(backend_enum): + tried.append(backend_enum) + cls = MagicMock() + cls.validate_configuration.return_value = ["not available"] + return cls + + with ( + patch.object( + RocmPlatform, + "get_device_capability", + return_value=MagicMock(), + ), + patch.object(MLAPrefillBackendEnum, "get_class", mock_get_class), + suppress(ValueError), + ): + RocmPlatform.get_mla_prefill_backend_cls(selector_config) + + assert tried == [ MLAPrefillBackendEnum.ROCM_AITER_FA, MLAPrefillBackendEnum.FLASH_ATTN, ] @@ -459,50 +569,29 @@ def test_auto_select_prefers_aiter_fa_on_rocm(self): AiterFlashAttnPrefillBackend, ) - # gfx gating is simulated via the mocked validate_configuration, - # not the capability. - capability = MagicMock() selector_config = MLAPrefillSelectorConfig(dtype=torch.bfloat16) - with ( - patch("vllm.platforms.current_platform") as mock_platform, - patch.object( - AiterFlashAttnPrefillBackend, - "validate_configuration", - return_value=[], - ), - ): - mock_platform.is_rocm.return_value = True - backend = _auto_select_mla_prefill_backend(capability, selector_config) + with patch("vllm.platforms.current_platform") as mock_platform: + mock_platform.get_mla_prefill_backend_cls.return_value = ( + AiterFlashAttnPrefillBackend + ) + backend = _auto_select_mla_prefill_backend(selector_config) assert backend.get_name() == "ROCM_AITER_FA" def test_auto_select_falls_back_to_flash_attn_when_aiter_invalid(self): - from vllm.v1.attention.backends.mla.prefill.aiter_flash_attn import ( - AiterFlashAttnPrefillBackend, - ) - try: flash_attn_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class() except ImportError: pytest.skip("FLASH_ATTN backend not available") return - # the fallback is forced by the mocked validate_configuration, + # the fallback is forced by the mocked get_mla_prefill_backend_cls, # not the capability. - capability = MagicMock() selector_config = MLAPrefillSelectorConfig(dtype=torch.bfloat16) - with ( - patch("vllm.platforms.current_platform") as mock_platform, - patch.object( - AiterFlashAttnPrefillBackend, - "validate_configuration", - return_value=["compute capability not supported"], - ), - patch.object(flash_attn_cls, "validate_configuration", return_value=[]), - ): - mock_platform.is_rocm.return_value = True - backend = _auto_select_mla_prefill_backend(capability, selector_config) + with patch("vllm.platforms.current_platform") as mock_platform: + mock_platform.get_mla_prefill_backend_cls.return_value = flash_attn_cls + backend = _auto_select_mla_prefill_backend(selector_config) assert backend.get_name() == "FLASH_ATTN" diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 25d28daf937f..e0257acb8007 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -18,6 +18,7 @@ get_visible_memory_node, ) from vllm.utils.mem_constants import GiB_bytes +from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum from vllm.v1.attention.backends.registry import AttentionBackendEnum from .interface import CpuArchEnum, Platform, PlatformEnum @@ -26,6 +27,10 @@ if TYPE_CHECKING: from vllm.config import VllmConfig + from vllm.v1.attention.backends.mla.prefill.base import ( + MLAPrefillBackend, + ) + from vllm.v1.attention.backends.mla.prefill.selector import MLAPrefillSelectorConfig from vllm.v1.attention.selector import AttentionSelectorConfig else: VllmConfig = None @@ -173,6 +178,21 @@ def get_attn_backend_cls( logger.info("Cannot use %s backend on CPU.", selected_backend) return AttentionBackendEnum.CPU_ATTN.get_path() + @classmethod + def get_mla_prefill_backend_cls( + cls, + mla_selector_config: "MLAPrefillSelectorConfig", + ) -> "type[MLAPrefillBackend]": + """Get the MLA prefill backend class of a device.""" + for backend_enum in (MLAPrefillBackendEnum.ZEN_CPU, MLAPrefillBackendEnum.CPU): + try: + cpu_backend_cls = backend_enum.get_class() + except ImportError: + continue + if cpu_backend_cls.is_available(): + return cpu_backend_cls + raise ValueError("No valid CPU MLA prefill backend found.") + @classmethod def get_device_total_memory(cls, device_id: int = 0) -> int: meminfo = get_memory_node_info(device_id) diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index f4583daba26f..dc3200759c82 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -26,6 +26,8 @@ from vllm.sampling_params import SamplingParams from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.v1.attention.backend import AttentionBackend + from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend + from vllm.v1.attention.backends.mla.prefill.selector import MLAPrefillSelectorConfig from vllm.v1.attention.selector import AttentionSelectorConfig else: FlexibleArgumentParser = object @@ -429,6 +431,88 @@ def get_vit_attn_backend( ) return AttentionBackendEnum.TORCH_SDPA + @classmethod + def get_mla_prefill_backend_cls( + cls, + mla_selector_config: "MLAPrefillSelectorConfig", + ) -> "type[MLAPrefillBackend]": + """Get the MLA prefill backend class of a device. + + Platforms that have a single, fixed MLA prefill backend can override + this to return it directly. The default implementation selects the + best backend from a device-capability-based priority list. + + Returns: + The selected backend class. + + Raises: + ValueError: If no MLA prefill backend is valid for the + configuration. + + """ + from vllm.v1.attention.backends.mla.prefill.base import MLADimensions + from vllm.v1.attention.backends.mla.prefill.registry import ( + MLAPrefillBackendEnum, + ) + + device_capability = cls.get_device_capability() + if device_capability is None: + return MLAPrefillBackendEnum.FLASH_ATTN.get_class() + + if device_capability.major == 10: # Blackwell + if mla_selector_config.mla_dimensions == MLADimensions( + qk_nope_head_dim=192, + qk_rope_head_dim=64, + v_head_dim=256, + ): + priorities = [ + MLAPrefillBackendEnum.TRTLLM_RAGGED, + MLAPrefillBackendEnum.FLASH_ATTN, + MLAPrefillBackendEnum.FLASHINFER, + MLAPrefillBackendEnum.TOKENSPEED_MLA, + ] + else: + priorities = [ + MLAPrefillBackendEnum.FLASH_ATTN, + MLAPrefillBackendEnum.TRTLLM_RAGGED, + MLAPrefillBackendEnum.FLASHINFER, + MLAPrefillBackendEnum.TOKENSPEED_MLA, + ] + else: # Hopper (SM90) and older + priorities = [MLAPrefillBackendEnum.FLASH_ATTN] + + all_invalid_reasons: dict[str, list[str]] = {} + + for backend_enum in priorities: + try: + backend_cls = backend_enum.get_class() + invalid_reasons = backend_cls.validate_configuration( + device_capability, mla_selector_config + ) + except ImportError: + invalid_reasons = ["ImportError"] + if not invalid_reasons: + return backend_cls + all_invalid_reasons[backend_enum.name] = invalid_reasons + + reasons_str = ( + "{" + + ", ".join( + f"{name}: [{', '.join(reasons)}]" + for name, reasons in all_invalid_reasons.items() + ) + + "}" + ) + logger.debug_once( + "Some MLA prefill backends are not valid with %s. Reasons: %s.", + repr(mla_selector_config), + reasons_str, + ) + raise ValueError( + f"No valid MLA prefill backend found with {mla_selector_config!r}. " + f"Reasons: {reasons_str}." + ) + @classmethod def get_device_capability( cls, diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index f394e56bf640..5840ffe3b792 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -24,6 +24,10 @@ from vllm.config.kernel import IrOpPriorityConfig from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.v1.attention.backend import AttentionBackend + from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend + from vllm.v1.attention.backends.mla.prefill.selector import ( + MLAPrefillSelectorConfig, + ) from vllm.v1.attention.selector import AttentionSelectorConfig logger = init_logger(__name__) @@ -904,6 +908,57 @@ def get_vit_attn_backend( logger.info_once("Using Torch SDPA backend for ViT model.") return AttentionBackendEnum.TORCH_SDPA + @classmethod + def get_mla_prefill_backend_cls( + cls, + mla_selector_config: "MLAPrefillSelectorConfig", + ) -> "type[MLAPrefillBackend]": + """On ROCm, prefer the AITER FlashAttention backend, falling back to + FlashAttention. + + Raises: + ValueError: If neither backend is valid for the configuration. + + """ + from vllm.v1.attention.backends.mla.prefill.registry import ( + MLAPrefillBackendEnum, + ) + + device_capability = cls.get_device_capability() + all_invalid_reasons: dict[str, list[str]] = {} + for backend_enum in ( + MLAPrefillBackendEnum.ROCM_AITER_FA, + MLAPrefillBackendEnum.FLASH_ATTN, + ): + try: + backend_cls = backend_enum.get_class() + invalid_reasons = backend_cls.validate_configuration( + device_capability, mla_selector_config + ) + except ImportError: + invalid_reasons = ["ImportError"] + if not invalid_reasons: + return backend_cls + all_invalid_reasons[backend_enum.name] = invalid_reasons + + reasons_str = ( + "{" + + ", ".join( + f"{name}: [{', '.join(reasons)}]" + for name, reasons in all_invalid_reasons.items() + ) + + "}" + ) + logger.debug_once( + "Some MLA prefill backends are not valid with %s. Reasons: %s.", + repr(mla_selector_config), + reasons_str, + ) + raise ValueError( + f"No valid MLA prefill backend found with {mla_selector_config!r}. " + f"Reasons: {reasons_str}." + ) + @classmethod def set_device(cls, device: torch.device) -> None: """Set the device for the current platform.""" diff --git a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py index 1d03b9e77392..8b0e337bb221 100644 --- a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py @@ -31,7 +31,10 @@ def get_name() -> str: return "ROCM_AITER_FA" @classmethod - def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: + def supports_compute_capability( + cls, + device_capability: "DeviceCapability | None", + ) -> bool: if not current_platform.is_rocm(): return False from vllm.platforms.rocm import on_mi3xx diff --git a/vllm/v1/attention/backends/mla/prefill/base.py b/vllm/v1/attention/backends/mla/prefill/base.py index 1b589cb295f4..4b4361214f84 100644 --- a/vllm/v1/attention/backends/mla/prefill/base.py +++ b/vllm/v1/attention/backends/mla/prefill/base.py @@ -49,7 +49,10 @@ def get_name() -> str: raise NotImplementedError @classmethod - def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: + def supports_compute_capability( + cls, + device_capability: "DeviceCapability | None", + ) -> bool: return True @classmethod @@ -90,16 +93,19 @@ def supports_out(self) -> bool: @classmethod def validate_configuration( cls, - device_capability: "DeviceCapability", + device_capability: "DeviceCapability | None", selector_config: "MLAPrefillSelectorConfig", ) -> list[str]: invalid_reasons: list[str] = [] if not cls.supports_compute_capability(device_capability): - invalid_reasons.append( - f"compute capability {device_capability.major}." - f"{device_capability.minor} not supported" - ) + if device_capability is None: + invalid_reasons.append("device capability not available") + else: + invalid_reasons.append( + f"compute capability {device_capability.major}." + f"{device_capability.minor} not supported" + ) if not cls.supports_dtype(selector_config.dtype): invalid_reasons.append(f"dtype {selector_config.dtype} not supported") diff --git a/vllm/v1/attention/backends/mla/prefill/flashinfer.py b/vllm/v1/attention/backends/mla/prefill/flashinfer.py index 777dd33756c7..5cf47ea2caeb 100644 --- a/vllm/v1/attention/backends/mla/prefill/flashinfer.py +++ b/vllm/v1/attention/backends/mla/prefill/flashinfer.py @@ -51,8 +51,11 @@ def get_name() -> str: return "FLASHINFER" @classmethod - def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: - return device_capability.major == 10 + def supports_compute_capability( + cls, + device_capability: "DeviceCapability | None", + ) -> bool: + return device_capability is not None and device_capability.major == 10 @classmethod def is_available(cls) -> bool: diff --git a/vllm/v1/attention/backends/mla/prefill/selector.py b/vllm/v1/attention/backends/mla/prefill/selector.py index 0cb29f51b49a..f0f001580964 100644 --- a/vllm/v1/attention/backends/mla/prefill/selector.py +++ b/vllm/v1/attention/backends/mla/prefill/selector.py @@ -12,9 +12,7 @@ import torch from vllm.logger import init_logger -from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backends.mla.prefill.base import MLADimensions -from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum if TYPE_CHECKING: from vllm.config import VllmConfig @@ -45,54 +43,8 @@ def __repr__(self): ) -def _get_mla_prefill_backend_priorities( - device_capability: DeviceCapability, - mla_dimensions: MLADimensions, -) -> list[MLAPrefillBackendEnum]: - """Get MLA prefill backend priorities based on device capability. - - Args: - device_capability: The device's compute capability. - mla_dimensions: The model's MLA head dimensions. - - Returns: - List of backends in priority order (highest priority first). - - """ - from vllm.platforms import current_platform - - if current_platform.is_rocm(): - return [ - MLAPrefillBackendEnum.ROCM_AITER_FA, - MLAPrefillBackendEnum.FLASH_ATTN, - ] - - if device_capability.major == 10: # Blackwell - if mla_dimensions == MLADimensions( - qk_nope_head_dim=192, - qk_rope_head_dim=64, - v_head_dim=256, - ): - return [ - MLAPrefillBackendEnum.TRTLLM_RAGGED, - MLAPrefillBackendEnum.FLASH_ATTN, - MLAPrefillBackendEnum.FLASHINFER, - MLAPrefillBackendEnum.TOKENSPEED_MLA, - ] - return [ - MLAPrefillBackendEnum.FLASH_ATTN, - MLAPrefillBackendEnum.TRTLLM_RAGGED, - MLAPrefillBackendEnum.FLASHINFER, - MLAPrefillBackendEnum.TOKENSPEED_MLA, - ] - else: # Hopper (SM90) and older - return [ - MLAPrefillBackendEnum.FLASH_ATTN, - ] - - def get_mla_prefill_backend( - vllm_config: "VllmConfig", + vllm_config: "VllmConfig | None", ) -> "type[MLAPrefillBackend]": """Select the MLA prefill backend based on configuration and device. @@ -101,7 +53,8 @@ def get_mla_prefill_backend( priority-based selection. Args: - vllm_config: The vLLM configuration. + vllm_config: The vLLM configuration. May be None before model + resolution, in which case defaults are used. Returns: The selected prefill backend class. @@ -109,30 +62,7 @@ def get_mla_prefill_backend( """ from vllm.platforms import current_platform - if current_platform.is_cpu(): - # CPUs have no compute capability, so the capability-driven priority - # path below does not apply. Prefer an accelerator-specific CPU backend - # when its kernels are present, else the generic SDPA one. - for backend_enum in (MLAPrefillBackendEnum.ZEN_CPU, MLAPrefillBackendEnum.CPU): - try: - cpu_backend_cls = backend_enum.get_class() - except ImportError: - continue - if cpu_backend_cls.is_available(): - logger.info_once("Using %s MLA prefill backend.", backend_enum.name) - return cpu_backend_cls - raise ValueError("No valid CPU MLA prefill backend found.") - - device_capability = current_platform.get_device_capability() - if device_capability is None: - logger.info_once( - "Device capability not available, using FlashAttention MLA prefill backend." - ) - return MLAPrefillBackendEnum.FLASH_ATTN.get_class() - - attention_config = vllm_config.attention_config - - model_config = vllm_config.model_config + model_config = vllm_config.model_config if vllm_config is not None else None if model_config is None: selector_config = MLAPrefillSelectorConfig(dtype=torch.get_default_dtype()) else: @@ -146,11 +76,21 @@ def get_mla_prefill_backend( ), ) - if attention_config.mla_prefill_backend is not None: + if current_platform.is_cpu(): + # CPUs have a single fixed MLA prefill backend, so explicit backend + # configuration does not apply; always auto-select via the platform. + return _auto_select_mla_prefill_backend(selector_config) + + attention_config = vllm_config.attention_config if vllm_config is not None else None + if ( + attention_config is not None + and attention_config.mla_prefill_backend is not None + ): selected_backend = attention_config.mla_prefill_backend backend_cls: type[MLAPrefillBackend] | None = None try: backend_cls = selected_backend.get_class() + device_capability = current_platform.get_device_capability() invalid_reasons = backend_cls.validate_configuration( device_capability, selector_config ) @@ -166,63 +106,27 @@ def get_mla_prefill_backend( logger.info_once("Using %s MLA prefill backend.", selected_backend.name) return backend_cls - return _auto_select_mla_prefill_backend( - device_capability, - selector_config, - ) + return _auto_select_mla_prefill_backend(selector_config) @cache def _auto_select_mla_prefill_backend( - device_capability: DeviceCapability, selector_config: MLAPrefillSelectorConfig, ) -> "type[MLAPrefillBackend]": """Auto-select the best available MLA prefill backend. Args: - device_capability: The device's compute capability. selector_config: Hashable configuration for backend selection. Returns: The selected prefill backend class. - """ - priorities = _get_mla_prefill_backend_priorities( - device_capability, - selector_config.mla_dimensions, - ) - all_invalid_reasons: dict[str, list[str]] = {} + Raises: + ValueError: If the platform has no valid MLA prefill backend. - for backend_enum in priorities: - backend_cls: type[MLAPrefillBackend] | None = None - try: - backend_cls = backend_enum.get_class() - invalid_reasons = backend_cls.validate_configuration( - device_capability, selector_config - ) - except ImportError: - invalid_reasons = ["ImportError"] - if not invalid_reasons: - assert backend_cls is not None - logger.info_once("Using %s MLA prefill backend.", backend_enum.name) - return backend_cls - all_invalid_reasons[backend_enum.name] = invalid_reasons - - reasons_str = ( - "{" - + ", ".join( - f"{name}: [{', '.join(reasons)}]" - for name, reasons in all_invalid_reasons.items() - ) - + "}" - ) - config_str = repr(selector_config) - logger.debug_once( - "Some MLA prefill backends are not valid with %s. Reasons: %s.", - config_str, - reasons_str, - ) + """ + from vllm.platforms import current_platform - raise ValueError( - f"No valid MLA prefill backend found with {config_str}. Reasons: {reasons_str}." - ) + selected_backend = current_platform.get_mla_prefill_backend_cls(selector_config) + logger.info_once("Using %s MLA prefill backend.", selected_backend.get_name()) + return selected_backend diff --git a/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py index 6985b0a70c24..d65721b81b5f 100644 --- a/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py @@ -35,8 +35,11 @@ def get_name() -> str: return "TOKENSPEED_MLA" @classmethod - def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: - return device_capability.major == 10 + def supports_compute_capability( + cls, + device_capability: "DeviceCapability | None", + ) -> bool: + return device_capability is not None and device_capability.major == 10 _INSTALL_HINT = ( "tokenspeed_mla package is not installed. " diff --git a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py index 167738b07af4..3cb498be2b16 100644 --- a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py +++ b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py @@ -43,8 +43,11 @@ def get_name() -> str: return "TRTLLM_RAGGED" @classmethod - def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: - return device_capability.major == 10 + def supports_compute_capability( + cls, + device_capability: "DeviceCapability | None", + ) -> bool: + return device_capability is not None and device_capability.major == 10 @classmethod def is_available(cls) -> bool: From 194bd72be7b9eb1c406e090328ddaf09a4e2d6eb Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Thu, 8 Oct 2026 08:28:23 -0700 Subject: [PATCH 105/624] [Compile] Remove torch.compile-based sequence parallelism and AsyncTP (#60517) --- .buildkite/test-amd.yaml | 108 -- .buildkite/test_areas/compile.yaml | 128 --- docs/design/cuda_graphs.md | 4 +- docs/design/fusions.md | 76 +- docs/design/vllm_ir.md | 12 +- docs/features/batch_invariance.md | 5 +- .../compile/correctness_e2e/test_async_tp.py | 270 ----- .../correctness_e2e/test_sequence_parallel.py | 414 ------- tests/compile/fusions_e2e/common.py | 6 - tests/compile/fusions_e2e/conftest.py | 73 +- tests/compile/fusions_e2e/models.py | 19 - .../compile/fusions_e2e/test_tp2_async_tp.py | 339 ------ .../fusions_e2e/test_tp2_rocm_async_tp.py | 58 - .../compile/fusions_e2e/test_tp2_rocm_moe.py | 60 - .../passes/distributed/test_async_tp.py | 394 ------- .../distributed/test_sequence_parallelism.py | 343 ------ tests/compile/rocm/__init__.py | 2 - tests/compile/rocm/test_async_tp.py | 75 -- tests/compile/test_config.py | 184 +-- .../test_sequence_parallelism_threshold.py | 192 ---- .../test_routed_experts_capture.py | 5 +- tests/test_config.py | 12 +- tests/v1/worker/test_gpu_worker.py | 3 - .../passes/fusion/collective_fusion.py | 1013 ----------------- .../passes/fusion/sequence_parallelism.py | 620 ---------- vllm/compilation/passes/pass_manager.py | 7 - vllm/config/compilation.py | 60 +- vllm/config/vllm.py | 159 --- vllm/platforms/xpu.py | 1 - vllm/utils/flashinfer.py | 75 -- vllm/v1/worker/gpu/model_runner.py | 1 - vllm/v1/worker/gpu_model_runner.py | 42 +- vllm/v1/worker/gpu_worker.py | 38 - vllm/v1/worker/utils.py | 31 +- 34 files changed, 32 insertions(+), 4797 deletions(-) delete mode 100644 tests/compile/correctness_e2e/test_async_tp.py delete mode 100644 tests/compile/correctness_e2e/test_sequence_parallel.py delete mode 100644 tests/compile/fusions_e2e/test_tp2_async_tp.py delete mode 100644 tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py delete mode 100644 tests/compile/fusions_e2e/test_tp2_rocm_moe.py delete mode 100644 tests/compile/passes/distributed/test_async_tp.py delete mode 100644 tests/compile/passes/distributed/test_sequence_parallelism.py delete mode 100644 tests/compile/rocm/__init__.py delete mode 100644 tests/compile/rocm/test_async_tp.py delete mode 100644 tests/compile/test_sequence_parallelism_threshold.py delete mode 100644 vllm/compilation/passes/fusion/collective_fusion.py delete mode 100644 vllm/compilation/passes/fusion/sequence_parallelism.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 9e634ea02b4d..4ae92315b4c9 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -693,28 +693,8 @@ steps: - vllm/platforms/rocm.py commands: - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions -k "dynamo_partition" -# - label: ":amd: (MI300) Sequence Parallel Correctness" # TBD -# timeout_in_minutes: 180 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] -# dind: false -# agent_pool: mi300_2 -# num_gpus: 2 -# optional: true -# working_dir: "/vllm-workspace/" -# source_file_dependencies: -# - vllm/model_executor/layers/ -# - vllm/compilation/ -# - vllm/v1/worker/ -# - vllm/v1/cudagraph_dispatcher.py -# - tests/compile/correctness_e2e/test_sequence_parallel.py -# - vllm/platforms/rocm.py -# commands: -# - export VLLM_TEST_CLEAN_GPU_MEMORY=1 -# - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py - - label: ":amd: (MI300) Distributed Compile + RPC" # TBD timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -3180,7 +3160,6 @@ steps: # TODO: FP4 output-quant and trained MXFP4 AR/RMS fusion remain unsupported. - amd-smi - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_CUSTOM_AR=1 pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k 'test_tp2_ar_rms_fp8_fusions and qwen3 and ROCM_AITER_UNIFIED_ATTN and +quant_fp8 and (dynamo_partition or not +rms_norm)' - - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_moe.py - label: ":amd: (MI355) Fusion and Compile FP8" timeout_in_minutes: 45 @@ -3205,64 +3184,6 @@ steps: - pytest -v -s tests/compile/passes/test_fusion_attn.py tests/compile/passes/test_silu_mul_quant_fusion.py tests/compile/passes/test_silu_mul_quant_manual_fusion.py - pytest -v -s tests/compile/correctness_e2e/test_attn_quant.py -- label: ":amd: (MI355) Sequence Parallel Correctness" - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/" - source_file_dependencies: - - .buildkite/test_areas/compile.yaml - - vllm/compilation/ - - vllm/config/ - - vllm/distributed/ - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/layers/ - - vllm/model_executor/models/llama.py - - vllm/v1/attention/ - - vllm/v1/worker/ - - vllm/v1/cudagraph_dispatcher.py - - vllm/platforms/rocm.py - - tests/compile/correctness_e2e/test_sequence_parallel.py - - tests/models/registry.py - - tests/utils.py - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py - -- label: ":amd: (MI355) AsyncTP Correctness" - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/" - source_file_dependencies: - - .buildkite/test_areas/compile.yaml - - vllm/compilation/ - - vllm/config/ - - vllm/distributed/ - - vllm/model_executor/layers/ - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/models/llama.py - - vllm/model_executor/models/qwen3.py - - vllm/v1/attention/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - vllm/transformers_utils/repo_utils.py - - tests/compile/correctness_e2e/test_async_tp.py - - tests/compile/fusions_e2e/common.py - - tests/models/registry.py - - tests/utils.py - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_rocm_async_tp_bf16_output_correctness - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_async_tp_pass_correctness -k RedHatAI - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_async_tp_pass_nvfp4_correctness - - label: ":amd: (MI355 DPX) Fusion E2E FP8 Config Sweep" timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -5916,35 +5837,6 @@ steps: - amdmi355 optional: true -- label: ':amd: (MI355) Fusion E2E TP2 AsyncTP BF16 Config Sweep' - timeout_in_minutes: 30 - working_dir: /vllm-workspace/ - source_file_dependencies: - - vllm/compilation/ - - vllm/model_executor/layers/linear.py - - vllm/model_executor/layers/utils.py - - vllm/distributed/ - - vllm/v1/attention/backends/triton_attn.py - - tests/compile/fusions_e2e/ - - tests/compile/passes/distributed/test_async_tp.py - - tests/compile/rocm/test_async_tp.py - - tests/compile/correctness_e2e/test_async_tp.py - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/models/qwen3.py - commands: - - rocm-smi - - pytest -v -s tests/compile/rocm/test_async_tp.py - - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_rocm_async_tp_bf16_output_correctness - dind: false - agent_pool: mi355_2 - num_gpus: 2 - mirror_hardwares: - - amdexperimental - - amdproduction - - amdgfx950nightly - - amdmi355 - optional: true - label: ":amd: (MI355 DPX) MRV2 Sampler JIT Warmup" timeout_in_minutes: 10 diff --git a/.buildkite/test_areas/compile.yaml b/.buildkite/test_areas/compile.yaml index acc1bafa7f2a..4a5a4f1a3cea 100644 --- a/.buildkite/test_areas/compile.yaml +++ b/.buildkite/test_areas/compile.yaml @@ -2,84 +2,6 @@ group: Compile depends_on: - image-build steps: -- label: ":nvidia: (B200) Sequence Parallel Correctness" - key: sequence-parallel-correctness-tests-2xb200 - timeout_in_minutes: 45 - working_dir: "/vllm-workspace/" - device: b200-k8s - optional: true - num_devices: 2 - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py - mirror: - amd: - label: ":amd: (MI355) Sequence Parallel Correctness" - dind: false - device: mi355_2 - num_devices: 2 - timeout_in_minutes: 180 - depends_on: - - image-build-amd - source_file_dependencies: - - .buildkite/test_areas/compile.yaml - - vllm/compilation/ - - vllm/config/ - - vllm/distributed/ - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/layers/ - - vllm/model_executor/models/llama.py - - vllm/v1/attention/ - - vllm/v1/worker/ - - vllm/v1/cudagraph_dispatcher.py - - vllm/platforms/rocm.py - - tests/compile/correctness_e2e/test_sequence_parallel.py - - tests/models/registry.py - - tests/utils.py - -- label: ":nvidia: (B200) AsyncTP Correctness" - key: asynctp-correctness-tests-b200 - timeout_in_minutes: 30 - working_dir: "/vllm-workspace/" - device: b200-k8s - optional: true - num_devices: 2 - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py - mirror: - amd: - label: ":amd: (MI355) AsyncTP Correctness" - dind: false - device: mi355_2 - num_devices: 2 - timeout_in_minutes: 60 - working_dir: "/vllm-workspace/" - depends_on: - - image-build-amd - source_file_dependencies: - - .buildkite/test_areas/compile.yaml - - vllm/compilation/ - - vllm/config/ - - vllm/distributed/ - - vllm/model_executor/layers/ - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/models/llama.py - - vllm/model_executor/models/qwen3.py - - vllm/v1/attention/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - vllm/transformers_utils/repo_utils.py - - tests/compile/correctness_e2e/test_async_tp.py - - tests/compile/fusions_e2e/common.py - - tests/models/registry.py - - tests/utils.py - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_rocm_async_tp_bf16_output_correctness - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_async_tp_pass_correctness -k RedHatAI - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_async_tp_pass_nvfp4_correctness - - label: ":nvidia: (H100) Distributed Compile" key: distributed-compile-unit-tests-2xh100 timeout_in_minutes: 45 @@ -300,7 +222,6 @@ steps: - nvidia-smi # Run all models and attn backends but only Inductor partition and native custom ops - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))" - - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))" mirror: amd: label: ":amd: (MI355) Fusion E2E TP2 Quick" @@ -388,53 +309,6 @@ steps: - amd-smi - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_CUSTOM_AR=1 pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k 'test_tp2_ar_rms_fp8_fusions and Llama-3 and ROCM_AITER_UNIFIED_ATTN and (dynamo_partition or +quant_fp8 or +rms_norm)' -- label: ":nvidia: (H100) Fusion E2E TP2 AsyncTP Config Sweep" - key: fusion-e2e-tp2-asynctp-config-sweep-h100 - timeout_in_minutes: 40 - working_dir: "/vllm-workspace/" - device: h100 - num_devices: 2 - source_file_dependencies: - - csrc/quantization/ - - vllm/compilation/ - # can affect pattern matching - - vllm/model_executor/layers/layernorm.py - - vllm/model_executor/layers/activation.py - - vllm/model_executor/layers/attention/attention.py - - vllm/model_executor/layers/quantization/input_quant_fp8.py - - tests/compile/fusions_e2e/ - commands: - - nvidia-smi - # Run just llama3 (fp8 & bf16) for all config combinations - - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "llama-3" - mirror: - amd: - label: ':amd: (MI355) Fusion E2E TP2 AsyncTP BF16 Config Sweep' - dind: false - device: mi355_2 - num_devices: 2 - working_dir: /vllm-workspace/ - timeout_in_minutes: 30 - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/compilation/ - - vllm/model_executor/layers/linear.py - - vllm/model_executor/layers/utils.py - - vllm/distributed/ - - vllm/v1/attention/backends/triton_attn.py - - tests/compile/fusions_e2e/ - - tests/compile/passes/distributed/test_async_tp.py - - tests/compile/rocm/test_async_tp.py - - tests/compile/correctness_e2e/test_async_tp.py - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/models/qwen3.py - commands: - - rocm-smi - - pytest -v -s tests/compile/rocm/test_async_tp.py - - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_rocm_async_tp_bf16_output_correctness - - label: ":nvidia: (B200) Fusion E2E TP2" key: fusion-e2e-tp2-b200 timeout_in_minutes: 45 @@ -453,7 +327,6 @@ steps: # include qwen/deepseek with +quant_fp8 as -quant_fp8 rms+quant fusion is not supported # for ar-rms-quant-fp4, also sweep llama3 - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "(FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))) or Llama-3.1-8B-Instruct-FP4" - - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))" mirror: amd: label: ":amd: (MI355) Fusion E2E TP2 FP8 BF16 rewrites" @@ -480,7 +353,6 @@ steps: # TODO: FP4 output-quant and trained MXFP4 AR/RMS fusion remain unsupported. - amd-smi - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_CUSTOM_AR=1 pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k 'test_tp2_ar_rms_fp8_fusions and qwen3 and ROCM_AITER_UNIFIED_ATTN and +quant_fp8 and (dynamo_partition or not +rms_norm)' - - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_moe.py - label: ':amd: (MI355) Fusion E2E TP2 AR-RMS (Inductor Partition)' key: fusion-e2e-tp2-ar-rms-inductor-partition-amd diff --git a/docs/design/cuda_graphs.md b/docs/design/cuda_graphs.md index abebdd54ae46..9d53a98f6585 100644 --- a/docs/design/cuda_graphs.md +++ b/docs/design/cuda_graphs.md @@ -225,9 +225,9 @@ outputs = model.generate( ) ``` -### Piecewise compilation and full graph custom passes (attention fusion, sequence parallelism) +### Piecewise compilation and full graph custom passes (attention fusion) -Unfortunately, some custom compile passes have to see the whole graph to be effective and hence aren't compatible with piecewise compilation. This includes `AttnQuantFusionPass` and `SequenceParallelismPass`. As a short-term solution, we automatically disable piecewise compilation (by setting `splitting_ops=[]`) when attention fusion is enabled. We use CUDA Graph modes `FULL` or `FULL_DECODE_ONLY` (depending on backend support). However, this leads to another optimization incompatibility and confusing performance tradeoffs. +Unfortunately, some custom compile passes have to see the whole graph to be effective and hence aren't compatible with piecewise compilation. This includes `AttnQuantFusionPass`. As a short-term solution, we automatically disable piecewise compilation (by setting `splitting_ops=[]`) when attention fusion is enabled. We use CUDA Graph modes `FULL` or `FULL_DECODE_ONLY` (depending on backend support). However, this leads to another optimization incompatibility and confusing performance tradeoffs. Long term, we've added the ability to partition the graph in Inductor instead of right after Dynamo. It can be enabled with `CompilationConfig.use_inductor_graph_partition=True` but is currently experimental and only available with `torch>=2.9`. This also increases compilation time as it has to compile the whole graph and cannot reuse piecewise compilation artifacts. Once vLLM supports 2.9, we plan to make this the default approach as it will also speed up piecewise cudagraph capture. diff --git a/docs/design/fusions.md b/docs/design/fusions.md index 002984c113a1..4095c41a7e63 100644 --- a/docs/design/fusions.md +++ b/docs/design/fusions.md @@ -25,8 +25,6 @@ or just on the low or high end. | [MLA Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | MLA Attention output → FP8/NVFP4 quant | Off by default | TBD | Yes | Always | | [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache) | `fuse_rope_kvcache` | Rotary embedding → KV cache write | O2 (ROCm/AITER only) | 2-4% | No | Low | | [QK Norm + RoPE](#qk-norm--rope-enable_qk_norm_rope_fusion) | `enable_qk_norm_rope_fusion` | Q/K RMSNorm → rotary embedding | Off by default | 2-3% | No | Low | -| [Sequence Parallelism](#sequence-parallelism-enable_sp) | `enable_sp` | AllReduce → ReduceScatter + AllGather | Off by default | Prereq for AsyncTP | Yes | High | -| [AsyncTP GEMM + collective](#asynctp-gemm--collective-overlap-fuse_gemm_comms) | `fuse_gemm_comms` | GEMM → reduce-scatter / all-gather → GEMM | Off by default | 7-10% | Yes | High | | [RMSNorm + Quant](#rmsnorm--quantization-fuse_norm_quant) | `fuse_norm_quant` | RMSNorm (+residual add) → FP8/FP4 quant | O1 (conditional) | 1-4% | No | Always | | [SiLU+Mul + Quant](#silumul--quantization-fuse_act_quant) | `fuse_act_quant` | SiLU+Mul activation → FP8/FP4 quant | O1 (conditional) | 1-4% | No | Always | | [RMSNorm + Padding](#rmsnorm--padding-fuse_act_padding) | `fuse_act_padding` | Residual add + RMSNorm → padding | O1 (ROCm/AITER only) | TBD | No | Always | @@ -44,9 +42,7 @@ The table below lists the quantization schemes supported by each fusion on each | `fuse_attn_quant`\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static\* | | `fuse_attn_quant` (MLA)\* | FP8 static\*, FP8 per-group\*, NVFP4\* | FP8 static\*, FP8 per-group\* | FP8 static\*, FP8 per-group\* | — | FP8 static\* (untested) | | `fuse_rope_kvcache` | — | — | — | — | FP16/BF16 | -| `enable_qk_norm_rope_fusion` | FP16/BF16 | FP16/BF16 | FP16/BF16† | FP16/BF16† | — | -| `enable_sp` | FP16/BF16, FP8 static† | FP16/BF16, FP8 static | FP16/BF16† | FP16/BF16† | — | -| `fuse_gemm_comms` | FP16/BF16, FP8 static† | FP16/BF16, FP8 static | FP16/BF16† | FP16/BF16† | — | +| `enable_qk_norm_rope_fusion` | FP16/BF16 | FP16/BF16 | FP16/BF16 | FP16/BF16 | — | | `fuse_norm_quant` | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | — | FP8 static, FP8 per-token, FP8 per-group | | `fuse_act_quant` | FP8 static, NVFP4 | FP8 static, FP8 per-group (128/64) | FP8 static, FP8 per-group (128/64) | — | FP8 per-group | | `fuse_act_padding` | — | — | — | — | FP16/BF16 | @@ -56,10 +52,6 @@ The table below lists the quantization schemes supported by each fusion on each fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant) for per-backend details. -† `enable_sp` and `fuse_gemm_comms` are only autoconfigured for SM90 today; -other architectures support requires setting `PassConfig.sp_min_token_num` explicitly. -SM100 support also requires setting `VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel`. - ## Enabling / Disabling Fusions Fusions are exposed through `PassConfig`, which is nested inside `CompilationConfig`: @@ -186,72 +178,6 @@ If these conditions are set, the fusion is enabled automatically for optimizatio - Pass: [`vllm/compilation/passes/fusion/rope_kvcache_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rope_kvcache_fusion.py) -### Sequence Parallelism (`enable_sp`) - -**What it fuses.** Replaces all-reduce collectives with reduce-scatter + local RMSNorm + all-gather, -splitting the sequence dimension across TP ranks. This restructures the graph so the subsequent AsyncTP -pass can fuse the reduce-scatter / all-gather with the surrounding GEMMs. - -Sequence Parallelism itself does not directly improve performance; it is a prerequisite for the -AsyncTP pass (`fuse_gemm_comms`). SP is only applied above a minimum token threshold that is -autoconfigured based on device capability and model `hidden_size`. Currently only active on -H100/SM90 for models with `hidden_size >= 8192`. The threshold is configurable via -`PassConfig.sp_min_token_num`. - -The general transformation: - -```text -Input → AllReduce → RMSNorm → Output -becomes: -Input → ReduceScatter → local RMSNorm → AllGather → Output -``` - -Patterns covered: - -- First block: `AllReduce → RMSNorm` → `ReduceScatter → RMSNorm → AllGather` -- Middle blocks: `AllReduce → fused_add_RMSNorm` → `ReduceScatter → fused_add_RMSNorm → AllGather` -- Both with optional `→ FP8 static quant` suffix - -Requires: `use_inductor_graph_partition=True` **or** piecewise compilation with static sizes -divisible by `tensor_parallel_size`. - -Supported hardware: Only tested on NVIDIA CUDA, possibly works on ROCm. FP8 all-gather requires sm90+. - -**Code locations.** - -- Pass: [`vllm/compilation/passes/fusion/sequence_parallelism.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/sequence_parallelism.py) - -### AsyncTP GEMM + Collective Overlap (`fuse_gemm_comms`) - -!!! info - Requires `enable_sp=True` (enabled automatically). This pass is a no-op if Sequence Parallelism has not been applied. - -**What it fuses.** After Sequence Parallelism transforms the graph, fuses GEMM kernels with the -surrounding reduce-scatter (output projection) and all-gather (input projection) using -`torch.ops.symm_mem` symmetric-memory primitives, overlapping communication and computation. -This overlap is only profitable for large `num_tokens`, so the fusion (and preceding SP) -is only performed in the higher compiled range above `PassConfig.sp_min_token_num`. - -Patterns covered: - -- `GEMM → reduce-scatter` → `fused_matmul_reduce_scatter` -- `all-gather → GEMM` → `all_gather_matmul` -- FP8 scaled variants of both patterns - -Supported hardware: NVIDIA CUDA with symmetric-memory (`torch.distributed._symmetric_memory`) support. - -On B200, pattern-matching fp8 FlashInfer scaled MM is not supported, so it must be disabled -([#27893](https://github.com/vllm-project/vllm/issues/27893)) - -```shell -VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel ... -``` - -**Code locations.** - -- Pass: [`vllm/compilation/passes/fusion/collective_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/collective_fusion.py) -- Sequence parallelism pass: [`vllm/compilation/passes/fusion/sequence_parallelism.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/sequence_parallelism.py) - ### QK Norm + RoPE (`enable_qk_norm_rope_fusion`) !!! info diff --git a/docs/design/vllm_ir.md b/docs/design/vllm_ir.md index 82628f3762fe..e324ad25d03a 100644 --- a/docs/design/vllm_ir.md +++ b/docs/design/vllm_ir.md @@ -233,19 +233,17 @@ storing this information in vLLM's `PassContext` for later use in clone eliminat ### 3. IR Fusion and Transformation Passes After functionalization, custom vLLM passes operate on the functional FX graph containing high-level IR operations. -These passes can perform fusion, distribute operations for sequence parallelism, and other transformations: +These passes can perform fusion and other transformations: ```python -# Example: Sequence Parallelism (see SequenceParallelismPass) -# Before SP pass +# Example: AllReduce + RMSNorm fusion (see AllReduceFusionPass) +# Before all_reduce = torch.ops.vllm.all_reduce(x, "tp:0") rms_norm = torch.ops.vllm_ir.rms_norm(all_reduce, weight, 1e-5) -# after SP pass -reduce_scatter = torch.ops.vllm.reduce_scatter(x, "tp:0") -rms_norm = torch.ops.vllm_ir.rms_norm(all_reduce, weight, 1e-5) -all_gather = torch.ops.vllm.all_gather(x, "tp:0") +# After +rms_norm = torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm(x, weight, ...) ``` Fusion passes benefit from the high-level representation: they don't need to match against low-level PyTorch operations, diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index e82336d24117..b40ff4da8cc6 100644 --- a/docs/features/batch_invariance.md +++ b/docs/features/batch_invariance.md @@ -148,11 +148,10 @@ When batch invariance is enabled, vLLM: 1. Uses deterministic kernel implementations for attention and other operations 2. Ensures consistent numerical behavior across different batch sizes -3. Disables certain optimizations that may introduce non-determinism (such as sequence parallelism / async TP, whose reduce-scatter path is not batch-invariant) +3. Disables certain optimizations that may introduce non-determinism 4. Under tensor parallelism, keeps custom all-reduce on with a fixed reduction order (the 1-stage kernel is pinned, and large inputs are reduced in fixed-size chunks), and disables FlashInfer, AITER and QuickReduce all-reduce 5. On CUDA devices with tuned matmul table entries for the model's bf16 unquantized forward linear layers (Ada, Hopper, Blackwell), - runs without `torch.compile` using breakable CUDA graphs so tile configs follow the runtime batch size; set `VLLM_USE_BREAKABLE_CUDAGRAPH=0` to opt out - (not applied when sequence parallelism / async TP are enabled, since those are `torch.compile` passes). + runs without `torch.compile` using breakable CUDA graphs so tile configs follow the runtime batch size; set `VLLM_USE_BREAKABLE_CUDAGRAPH=0` to opt out. 6. Splits the LoRA shrink reduction into 8 fixed K chunks that are summed in a fixed order instead of with atomics !!! warning diff --git a/tests/compile/correctness_e2e/test_async_tp.py b/tests/compile/correctness_e2e/test_async_tp.py deleted file mode 100644 index 307f39039b39..000000000000 --- a/tests/compile/correctness_e2e/test_async_tp.py +++ /dev/null @@ -1,270 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import json -import logging -from contextlib import nullcontext - -import pytest - -from tests.compile.fusions_e2e.common import FUSION_LOG_PATTERNS -from tests.models.registry import HF_EXAMPLE_MODELS -from tests.utils import ( - compare_two_settings, - create_new_process_for_each_test, - multi_gpu_test, -) -from vllm.config import ( - CompilationMode, -) -from vllm.platforms import current_platform -from vllm.transformers_utils.repo_utils import hf_api -from vllm.utils.flashinfer import has_flashinfer - -NVFP4_MODEL_ID = "nvidia/Llama-3.1-8B-Instruct-NVFP4" -NVFP4_HF_OVERRIDES = { - "num_hidden_layers": 4, - "hidden_size": 512, - "intermediate_size": 800, - "num_attention_heads": 4, - "num_key_value_heads": 1, -} - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm AsyncTP contract") -def test_rocm_async_tp_bf16_output_correctness(caplog_mp_spawn): - """Trained-model outputs must agree with unfused TP, including padded decode.""" - # A snapshot path pins both servers and the comparison helper's tokenizer. - model = hf_api().snapshot_download( - "Qwen/Qwen3-0.6B", revision="c1899de289a04d12100db370d81485cdf75e47ca" - ) - common_args = [ - "--dtype", - "bfloat16", - "--load-format", - "auto", - "--tensor-parallel-size", - "2", - "--distributed-executor-backend", - "mp", - "--max-model-len", - "512", - "--max-num-batched-tokens", - "512", - "--max-num-seqs", - "8", - "--kv-cache-memory-bytes", - str(64 * 1024**2), - "--gpu-memory-utilization", - "0.01", - "--attention-backend", - "TRITON_ATTN", - "--disable-custom-all-reduce", - ] - - def settings(enabled): - return [ - *common_args, - "--compilation-config", - json.dumps( - { - "mode": CompilationMode.VLLM_COMPILE, - "cudagraph_mode": "NONE", - "splitting_ops": [], - "inductor_compile_config": {"force_disable_caches": True}, - "pass_config": { - "enable_sp": enabled, - "fuse_gemm_comms": enabled, - "fuse_allreduce_rms": False, - "sp_min_token_num": 1, - }, - } - ), - ] - - env = {"VLLM_ROCM_USE_AITER": "0", "VLLM_DISABLE_COMPILE_CACHE": "1"} - with caplog_mp_spawn(logging.DEBUG) as logs: - compare_two_settings( - model, - settings(False), - settings(True), - env1=env, - env2=env, - method="generate", - force_v1_runner=True, - ) - matches = [ - int(count) for count in FUSION_LOG_PATTERNS["async_tp"].findall(logs.text) - ] - print(f"AsyncTP rewrite counts: {matches}") - assert len(matches) >= 2 and all(count > 0 for count in matches), ( - f"Expected positive AsyncTP rewrites on both TP workers, found {matches}" - ) - - -@create_new_process_for_each_test() -@pytest.mark.parametrize( - "model_id", - ["meta-llama/Llama-3.2-1B-Instruct", "RedHatAI/Llama-3.2-1B-Instruct-FP8"], -) -@pytest.mark.parametrize("tp_size", [2]) -@pytest.mark.parametrize("async_tp_enabled", [True]) -@pytest.mark.parametrize("distributed_backend", ["mp"]) -def test_async_tp_pass_correctness( - model_id: str, - tp_size: int, - async_tp_enabled: bool, - distributed_backend: str, - num_gpus_available: int, - monkeypatch, -): - # Disable FlashInfer FP8 scaled_mm kernel as it is incompatible with - # async TP patterns. No-op on H100 (kernel requires CC >= 100). - monkeypatch.setenv("VLLM_DISABLED_KERNELS", "FlashInferFP8ScaledMMLinearKernel") - - model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id) - model_info.check_transformers_version(on_fail="skip") - model_info.check_available_online(on_fail="skip") - - pp_size = 1 - if num_gpus_available < tp_size: - pytest.skip(f"Need at least {tp_size} x {pp_size} GPUs") - - common_args = [ - "--dtype", - "bfloat16", - "--max-model-len", - "2048", - "--max-num-seqs", - "8", - ] - - compilation_config = { - "mode": CompilationMode.VLLM_COMPILE, - "compile_sizes": [2, 4, 8], - "splitting_ops": [], - "pass_config": {"fuse_gemm_comms": async_tp_enabled}, - } - - async_tp_args = [ - *common_args, - "--tensor-parallel-size", - str(tp_size), - "--distributed-executor-backend", - distributed_backend, - "--compilation_config", - json.dumps(compilation_config), - ] - - tp_args = [ - *common_args, - "--tensor-parallel-size", - str(tp_size), - "--distributed-executor-backend", - "mp", - ] - - compare_two_settings( - model_id, - async_tp_args, - tp_args, - method="generate", - force_v1_runner=True, - ) - - -@create_new_process_for_each_test() -def test_async_tp_pass_nvfp4_correctness(num_gpus_available: int, caplog_mp_spawn): - is_blackwell = current_platform.is_cuda() and ( - current_platform.is_device_capability_family(100) - ) - is_rocm_cdna = False - if current_platform.is_rocm(): - from vllm.platforms.rocm import on_cdna - - is_rocm_cdna = on_cdna() - if not (is_blackwell or is_rocm_cdna): - pytest.skip("NVFP4 AsyncTP requires Blackwell or ROCm CDNA") - if is_blackwell and not has_flashinfer(): - pytest.skip("FlashInfer is required for the NVFP4 AsyncTP path") - - tp_size = 2 - if num_gpus_available < tp_size: - pytest.skip(f"Need at least {tp_size} GPUs") - - common_args = [ - "--dtype", - "bfloat16", - "--max-model-len", - "2048", - "--max-num-seqs", - "8", - "--load-format", - "dummy", - *(["--linear-backend", "flashinfer_cutlass"] if is_blackwell else []), - "--hf-overrides", - json.dumps(NVFP4_HF_OVERRIDES), - ] - - compilation_config = { - "mode": CompilationMode.VLLM_COMPILE, - "compile_sizes": [2, 4, 8], - "splitting_ops": [], - "pass_config": { - "enable_sp": True, - "fuse_gemm_comms": True, - "fuse_allreduce_rms": False, - "sp_min_token_num": 1, - }, - } - if is_rocm_cdna: - compilation_config["inductor_compile_config"] = {"force_disable_caches": True} - - async_tp_args = [ - *common_args, - "--tensor-parallel-size", - str(tp_size), - "--distributed-executor-backend", - "mp", - "--compilation_config", - json.dumps(compilation_config), - ] - - tp_args = [ - *common_args, - "--tensor-parallel-size", - str(tp_size), - "--distributed-executor-backend", - "mp", - ] - - env = {"VLLM_DISABLE_COMPILE_CACHE": "1"} if is_rocm_cdna else None - log_context = caplog_mp_spawn(logging.DEBUG) if is_rocm_cdna else nullcontext() - with log_context as logs: - compare_two_settings( - NVFP4_MODEL_ID, - async_tp_args, - tp_args, - env1=env, - env2=env, - method="generate", - force_v1_runner=True, - ) - if is_rocm_cdna: - assert logs is not None - matches_by_rank = { - rank: [ - int(count) - for line in logs.text.splitlines() - if line.startswith(f"(Worker_TP{rank} pid=") - for count in FUSION_LOG_PATTERNS["async_tp"].findall(line) - ] - for rank in range(tp_size) - } - assert all( - any(count > 0 for count in matches) for matches in matches_by_rank.values() - ), ( - "Expected positive NVFP4 AsyncTP rewrites on both TP workers, " - f"found {matches_by_rank}" - ) diff --git a/tests/compile/correctness_e2e/test_sequence_parallel.py b/tests/compile/correctness_e2e/test_sequence_parallel.py deleted file mode 100644 index 10d25320c012..000000000000 --- a/tests/compile/correctness_e2e/test_sequence_parallel.py +++ /dev/null @@ -1,414 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""WARNING: This test runs in both single-node (4 GPUs) and multi-node -(2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is -important to set the distributed backend to "mp" to avoid Ray scheduling -all workers in a node other than the head node, which can cause the test -to fail. -""" - -import json -import os -from dataclasses import dataclass -from typing import Literal, NamedTuple - -import pytest - -from vllm.config.compilation import CompilationMode -from vllm.config.model import RunnerOption -from vllm.logger import init_logger -from vllm.platforms import current_platform -from vllm.utils.torch_utils import is_torch_equal_or_newer - -from ...models.registry import HF_EXAMPLE_MODELS, _HfExamplesInfo -from ...utils import compare_all_settings, create_new_process_for_each_test - -logger = init_logger("test_sequence_parallel") - -VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1" -NVFP4_MODEL_ID = "nvidia/Llama-3.1-8B-Instruct-NVFP4" -NVFP4_MODEL_INFO = _HfExamplesInfo(NVFP4_MODEL_ID) - - -class ParallelSetup(NamedTuple): - tp_size: int - pp_size: int - fuse_norm_quant: bool - fuse_act_quant: bool - eager_mode: bool - - -class SPTestOptions(NamedTuple): - multi_node_only: bool - load_format: str | None = None - model_info: _HfExamplesInfo | None = None - - -@dataclass -class SPTestSettings: - parallel_setups: list[ParallelSetup] - distributed_backends: list[str] - runner: RunnerOption - test_options: SPTestOptions - - @staticmethod - def fast( - *, - tp_base: int = 2, - pp_base: int = 1, - runner: RunnerOption = "auto", - multi_node_only: bool = False, - load_format: str | None = None, - ): - parallel_setups = [ - ParallelSetup( - tp_size=tp_base, - pp_size=pp_multiplier * pp_base, - fuse_norm_quant=False, - fuse_act_quant=False, - eager_mode=False, - ) - for pp_multiplier in [1, 2] - ] - return SPTestSettings( - parallel_setups=parallel_setups, - distributed_backends=["mp", "ray"], - runner=runner, - test_options=SPTestOptions( - multi_node_only=multi_node_only, load_format=load_format - ), - ) - - @staticmethod - def fp8_quant( - *, - tp_base: int = 2, - pp_base: int = 1, - runner: RunnerOption = "auto", - multi_node_only: bool = False, - load_format: str | None = None, - ): - parallel_setups = [] - for fusion_val in [False, True]: - parallel_setups.append( - ParallelSetup( - tp_size=tp_base, - pp_size=pp_base, - fuse_norm_quant=fusion_val, - fuse_act_quant=fusion_val, - eager_mode=True, - ) - ) - return SPTestSettings( - parallel_setups=parallel_setups, - distributed_backends=["mp", "ray"], - runner=runner, - test_options=SPTestOptions( - multi_node_only=multi_node_only, load_format=load_format - ), - ) - - def iter_params(self, model_id: str): - opts = self.test_options - - for parallel_setup in self.parallel_setups: - for backend in self.distributed_backends: - yield ( - model_id, - parallel_setup, - backend, - self.runner, - opts, - ) - - -def _build_sp_args( - model_id: str, - parallel_setup: ParallelSetup, - distributed_backend: str, - runner: RunnerOption, - test_options: SPTestOptions, - num_gpus_available: int, - use_inductor_graph_partition: bool, - fuse_gemm_comms: bool, - enable_prompt_embeds: bool, - *, - is_multimodal: bool, -) -> tuple[list[str], list[str]] | None: - ( - tp_size, - pp_size, - fuse_norm_quant, - fuse_act_quant, - eager_mode, - ) = parallel_setup - - multi_node_only = test_options.multi_node_only - load_format = test_options.load_format - - model_info = test_options.model_info or HF_EXAMPLE_MODELS.find_hf_info(model_id) - model_info.check_transformers_version(on_fail="skip") - - trust_remote_code = model_info.trust_remote_code - tokenizer_mode = model_info.tokenizer_mode - hf_overrides = dict(model_info.hf_overrides) - require_embed_inputs = model_info.require_embed_inputs - - if load_format == "dummy": - # Avoid OOM - text_overrides = { - "num_hidden_layers": 4, - "hidden_size": 512, - "intermediate_size": 800, - "num_attention_heads": 4, - "num_key_value_heads": 1, - } - - if is_multimodal: - hf_overrides.update({"text_config": text_overrides}) - else: - hf_overrides.update(text_overrides) - else: - model_info.check_available_online(on_fail="skip") - - if num_gpus_available < tp_size * pp_size: - return None - if VLLM_MULTI_NODE and distributed_backend == "mp": - return None - if multi_node_only and not VLLM_MULTI_NODE: - return None - - common_args = [ - "--dtype", - "bfloat16", - "--max-model-len", - "2048", - "--max-num-seqs", - "8", - ] - if eager_mode: - common_args.append("-cc.cudagraph_mode=none") - if runner != "auto": - common_args.extend(["--runner", runner]) - if trust_remote_code: - common_args.append("--trust-remote-code") - if tokenizer_mode: - common_args.extend(["--tokenizer-mode", tokenizer_mode]) - if load_format: - common_args.extend(["--load-format", load_format]) - if hf_overrides: - common_args.extend(["--hf-overrides", json.dumps(hf_overrides)]) - if require_embed_inputs: - common_args.extend( - [ - "--skip-tokenizer-init", - "--enable-prompt-embeds", - "--enable-mm-embeds", - ] - ) - elif enable_prompt_embeds: - common_args.append("--enable-prompt-embeds") - - compilation_config = { - "mode": CompilationMode.VLLM_COMPILE, - "compile_sizes": [4, 8], - "pass_config": { - "enable_sp": True, - "fuse_gemm_comms": fuse_gemm_comms, - "fuse_norm_quant": fuse_norm_quant, - "fuse_act_quant": fuse_act_quant, - "fuse_allreduce_rms": False, - "eliminate_noops": True, - "sp_min_token_num": 0, - }, - "use_inductor_graph_partition": use_inductor_graph_partition, - } - if not use_inductor_graph_partition: - compilation_config["splitting_ops"] = [] - - tp_sp_args = [ - *common_args, - "--tensor-parallel-size", - str(tp_size), - "--pipeline-parallel-size", - str(pp_size), - "--distributed-executor-backend", - distributed_backend, - "--compilation_config", - json.dumps(compilation_config), - ] - - tp_args = [ - *common_args, - "--tensor-parallel-size", - str(tp_size), - "--distributed-executor-backend", - "mp", - ] - - return tp_args, tp_sp_args - - -def _compare_sp_settings( - model_id: str, - settings: list[tuple[list[str], list[str]]], - *, - method: Literal["generate", "encode"], -) -> None: - if not settings: - pytest.skip("No supported sequence-parallel configurations") - - settings_by_tp_args: dict[tuple[str, ...], list[list[str]]] = {} - for tp_args, tp_sp_args in settings: - settings_by_tp_args.setdefault(tuple(tp_args), []).append(tp_sp_args) - - for tp_args_key, grouped_tp_sp_args in settings_by_tp_args.items(): - all_args = [list(tp_args_key), *grouped_tp_sp_args] - compare_all_settings( - model_id, - all_args, - [None] * len(all_args), - method=method, - force_v1_runner=True, - ) - - -SP_TEXT_GENERATION_MODELS = { - # [Decoder-only] - "hmellor/tiny-random-LlamaForCausalLM": SPTestSettings.fast(), - "RedHatAI/Llama-3.2-1B-Instruct-FP8": SPTestSettings.fp8_quant(), -} - - -@pytest.mark.parametrize( - ("model_id", "settings"), - [ - pytest.param(model_id, settings, id=model_id) - for model_id, settings in SP_TEXT_GENERATION_MODELS.items() - ], -) -@create_new_process_for_each_test() -def test_tp_sp_generation( - model_id: str, - settings: SPTestSettings, - num_gpus_available, -): - # Skip FP8 SP-only test on sm89 (compute capability 8.9). - # An unknown capability (None) is left to fail in the test body rather than - # being silently skipped here. - capability = current_platform.get_device_capability() - if "fp8" in model_id.lower() and capability is not None and capability < (9, 0): - pytest.skip("FP8 reduction support begins with sm90 capable devices.") - - graph_partition_options = [False] - if is_torch_equal_or_newer("2.9.0.dev"): - graph_partition_options.append(True) - - comparisons = [] - for _, parallel_setup, backend, runner, test_options in settings.iter_params( - model_id - ): - for use_inductor_graph_partition in graph_partition_options: - comparison = _build_sp_args( - model_id, - parallel_setup, - backend, - runner, - test_options, - num_gpus_available, - use_inductor_graph_partition, - fuse_gemm_comms=False, # TODO: enable async TP - enable_prompt_embeds=False, - is_multimodal=False, - ) - if comparison is not None: - comparisons.append(comparison) - - _compare_sp_settings(model_id, comparisons, method="generate") - - -# Focused regression test for the SP + prompt_embeds graph-rewrite path. -# Covers pp_size=1 (SP only) and pp_size=2 (SP + PP); kept small on purpose so -# we don't double the matrix of `test_tp_sp_generation` above. -SP_PROMPT_EMBEDS_PARALLEL_SETUPS = [ - ParallelSetup( - tp_size=2, - pp_size=pp_size, - fuse_norm_quant=False, - fuse_act_quant=False, - eager_mode=False, - ) - for pp_size in [1, 2] -] - - -@create_new_process_for_each_test() -def test_tp_sp_generation_prompt_embeds( - num_gpus_available, -): - model_id = "hmellor/tiny-random-LlamaForCausalLM" - graph_partition_options = [False] - if is_torch_equal_or_newer("2.9.0.dev"): - graph_partition_options.append(True) - - comparisons = [] - for parallel_setup in SP_PROMPT_EMBEDS_PARALLEL_SETUPS: - for use_inductor_graph_partition in graph_partition_options: - comparison = _build_sp_args( - model_id, - parallel_setup, - distributed_backend="mp", - runner="auto", - test_options=SPTestOptions(multi_node_only=False, load_format=None), - num_gpus_available=num_gpus_available, - use_inductor_graph_partition=use_inductor_graph_partition, - fuse_gemm_comms=False, - enable_prompt_embeds=True, - is_multimodal=False, - ) - if comparison is not None: - comparisons.append(comparison) - - _compare_sp_settings(model_id, comparisons, method="generate") - - -@create_new_process_for_each_test() -def test_tp_sp_nvfp4_generation(num_gpus_available: int): - is_blackwell = current_platform.is_cuda() and ( - current_platform.is_device_capability_family(100) - ) - is_rocm_cdna = False - if current_platform.is_rocm(): - from vllm.platforms.rocm import on_cdna - - is_rocm_cdna = on_cdna() - if not (is_blackwell or is_rocm_cdna): - pytest.skip("NVFP4 sequence parallelism requires Blackwell or ROCm CDNA") - - comparison = _build_sp_args( - NVFP4_MODEL_ID, - ParallelSetup( - tp_size=2, - pp_size=1, - fuse_norm_quant=True, - fuse_act_quant=True, - eager_mode=True, - ), - "mp", - "auto", - SPTestOptions( - multi_node_only=False, - load_format="dummy", - model_info=NVFP4_MODEL_INFO, - ), - num_gpus_available, - use_inductor_graph_partition=False, - fuse_gemm_comms=False, - enable_prompt_embeds=False, - is_multimodal=False, - ) - _compare_sp_settings( - NVFP4_MODEL_ID, - [] if comparison is None else [comparison], - method="generate", - ) diff --git a/tests/compile/fusions_e2e/common.py b/tests/compile/fusions_e2e/common.py index 1f92ff5a1988..a4d996b48ddd 100644 --- a/tests/compile/fusions_e2e/common.py +++ b/tests/compile/fusions_e2e/common.py @@ -21,8 +21,6 @@ class Matches(NamedTuple): # distributed ar_rms_fusion: int = 0 aiter_ar_rms_fusion: int = 0 - sequence_parallel: int = 0 - async_tp: int = 0 class ModelFusionInfo(NamedTuple): @@ -143,8 +141,4 @@ def has_cuda_graph_wrapper_metadata() -> bool: "aiter_ar_rms_fusion": re.compile( r"RocmAiterAllReduceFusionPass Replaced (\d+) patterns" ), - "sequence_parallel": re.compile( - r"sequence_parallelism.py:\d+] Replaced (\d+) patterns" - ), - "async_tp": re.compile(r"collective_fusion.py:\d+] Replaced (\d+) patterns"), } diff --git a/tests/compile/fusions_e2e/conftest.py b/tests/compile/fusions_e2e/conftest.py index db7ba05a5734..8da7aa2a5a94 100644 --- a/tests/compile/fusions_e2e/conftest.py +++ b/tests/compile/fusions_e2e/conftest.py @@ -170,26 +170,8 @@ def run( for match_name in matches_check: log_matches = list(int(ms) for ms in log_matches_dict[match_name]) - # AR+RMS skips the largest range; SP skips the smallest. - # When both are enabled, AR+RMS activation count is - # model-dependent (hidden_size affects threshold), so derive - # from log data. - if ( - match_name == "ar_rms_fusion" - and "sequence_parallel" in matches_check - and num_compile_ranges >= 2 - ): - assert ( - len(log_matches) >= tp_size and len(log_matches) % tp_size == 0 - ), ( - f"Expected multiple of {tp_size} ar_rms log entries, " - f"found {len(log_matches)}" - ) - num_ranges_activated = len(log_matches) // tp_size - elif ( - match_name in ("ar_rms_fusion", "sequence_parallel") - and num_compile_ranges >= 2 - ): + # AR+RMS skips the largest range. + if match_name == "ar_rms_fusion" and num_compile_ranges >= 2: num_ranges_activated = num_compile_ranges - 1 else: num_ranges_activated = num_compile_ranges @@ -233,44 +215,6 @@ def run( f"Expecting at least {expected_matches - matches.ar_rms_fusion} " f"where ar+rms+quant was activated" ) - elif ( - match_name == "async_tp" - and "sequence_parallel" in matches_check - and num_compile_ranges >= 2 - ): - # AsyncTP only finds patterns on ranges where SP ran. - n_sp_ranges = num_compile_ranges - 1 - assert ( - sum(m == expected_matches for m in log_matches) - == tp_size * n_sp_ranges - ), ( - f"Expecting {expected_matches} async_tp on " - f"{tp_size * n_sp_ranges} SP-range entries, " - f"found: {log_matches}" - ) - assert sum(m == 0 for m in log_matches) == tp_size, ( - f"Expecting 0 async_tp on {tp_size} small-range entries " - f"(no SP), found: {log_matches}" - ) - elif ( - match_name == "ar_rms_fusion" - and "sequence_parallel" in matches_check - and num_compile_ranges >= 2 - ): - # SP consumes allreduce patterns first, so AR+RMS finds - # full matches only on the smallest range (no SP). - assert sum(m == expected_matches for m in log_matches) == tp_size, ( - f"Expecting {expected_matches} ar_rms on " - f"{tp_size} small-range entries, found: {log_matches}" - ) - assert sum(m == 0 for m in log_matches) == tp_size * ( - num_ranges_activated - 1 - ), ( - f"Expecting 0 ar_rms on " - f"{tp_size * (num_ranges_activated - 1)} large-range " - f"entries (SP took precedence), found: {log_matches}" - ) - elif match_name == "act_quant_fusion": actual_match = match_table.get("activation_quant_fusion_pass", 0) assert actual_match == expected_matches * n_expected, ( @@ -305,17 +249,4 @@ def run( f"(found {len(log_matches)}) in:\n {log_holder.text}" ) - if match_name == "sequence_parallel" and num_compile_ranges >= 2: - log_matches = re.findall( - r"pass_manager.py:\d+] Skipping " - r".*SequenceParallelismPass.* with compile range", - log_holder.text, - ) - - n_expected = tp_size * (num_compile_ranges - num_ranges_activated) - assert len(log_matches) == n_expected, ( - f'Could not find {n_expected} "Skipping SequenceParallelismPass" ' - f"(found {len(log_matches)}) in:\n {log_holder.text}" - ) - return run diff --git a/tests/compile/fusions_e2e/models.py b/tests/compile/fusions_e2e/models.py index 082ab0bf5296..d80d171984c3 100644 --- a/tests/compile/fusions_e2e/models.py +++ b/tests/compile/fusions_e2e/models.py @@ -73,8 +73,6 @@ matches=lambda n_layers: Matches( ar_rms_fusion=n_layers * 2 + 1, aiter_ar_rms_fusion=n_layers * 2 + 1, - sequence_parallel=n_layers * 2 + 1, - async_tp=n_layers * 4, ), ) @@ -85,8 +83,6 @@ act_quant_fusion=n_layers, attn_quant_fusion=n_layers, ar_rms_fusion=n_layers * 2 + 1, - sequence_parallel=n_layers * 2 + 1, - async_tp=n_layers * 4, ), ) @@ -96,8 +92,6 @@ act_quant_fusion=n_layers, attn_quant_fusion=n_layers, ar_rms_fusion=n_layers * 2 + 1, - sequence_parallel=n_layers * 2 + 1, - async_tp=n_layers * 4, ), ) @@ -127,8 +121,6 @@ rms_quant_fusion=n_layers, attn_quant_fusion=n_layers, ar_rms_fusion=n_layers * 2, - sequence_parallel=n_layers * 2, - async_tp=n_layers * 2 - 1, ), ) @@ -138,8 +130,6 @@ matches=lambda n_layers: Matches( attn_quant_fusion=n_layers, ar_rms_fusion=n_layers * 2, - sequence_parallel=n_layers * 2, - async_tp=n_layers * 2 - 1, ), ) @@ -149,8 +139,6 @@ norm_rope_fusion=n_layers, ar_rms_fusion=n_layers * 2 + 1, aiter_ar_rms_fusion=n_layers * 2 + 1, - sequence_parallel=n_layers * 2 + 1, - async_tp=n_layers * 2, ), ) @@ -161,8 +149,6 @@ norm_rope_fusion=n_layers, attn_quant_fusion=0, # attn + group quant not supported ar_rms_fusion=n_layers * 2 + 1, - sequence_parallel=n_layers * 2 + 1, - async_tp=n_layers * 2, ), ) @@ -192,9 +178,6 @@ # MLA attn + per-group FP8 quant attn_quant_fusion=n_layers, ar_rms_fusion=n_layers * 2 + 1, - # TODO - # sequence_parallel= n_layers * 2 + 1, - # async_tp=n_layers * 2, ), ) @@ -213,8 +196,6 @@ matches=lambda n_layers: Matches( ar_rms_fusion=n_layers * 2 + 1, aiter_ar_rms_fusion=n_layers + 1, - sequence_parallel=n_layers * 2 + 1, - async_tp=n_layers * 2, ), model_kwargs=( {"quantization_config": {"moe": {"activation": "mxfp8"}}} diff --git a/tests/compile/fusions_e2e/test_tp2_async_tp.py b/tests/compile/fusions_e2e/test_tp2_async_tp.py deleted file mode 100644 index b0d32dad1c17..000000000000 --- a/tests/compile/fusions_e2e/test_tp2_async_tp.py +++ /dev/null @@ -1,339 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable - -import pytest - -from vllm.config import PassConfig -from vllm.platforms import current_platform - -from ...utils import multi_gpu_test -from .common import ( - INDUCTOR_GRAPH_PARTITION, - AttentionBackendCase, - Matches, - custom_ops_combos, - is_blackwell, - nvfp4_kernel_exposes_input_quant_key, -) -from .models import ( - FLASHINFER_ATTN, - TRITON_ATTN, - llama3_8b, - llama3_8b_fp4, - llama3_8b_fp8, - llama4_scout_fp8, - qwen3_a3b, -) - -pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA") - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b_fp8, llama4_scout_fp8], -) -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN]) -@pytest.mark.parametrize("n_layers", [4]) -@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm")) -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) -def test_tp2_async_tp_fp8_fusions( - model_name: str, - matches_fn: Callable[[int], Matches], - model_kwargs: dict, - hf_overrides: Callable[[int], dict], - attn_backend: AttentionBackendCase, - n_layers: int, - custom_ops: str, - inductor_graph_partition: bool, - run_e2e_fusion_test, -): - matches = matches_fn(n_layers) - - # Reduce size of model and skip weight loading time - model_kwargs["hf_overrides"] = hf_overrides(n_layers) - model_kwargs["load_format"] = "dummy" - model_kwargs["max_model_len"] = 1024 - model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} - - compilation_config = dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=custom_ops.split(","), - pass_config=PassConfig( - fuse_norm_quant=True, - fuse_act_quant=False, - fuse_attn_quant=True, - enable_qk_norm_rope_fusion=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=False, - # Override threshold for testing (models have small hidden_size) - sp_min_token_num=512, - ), - ) - - matches_check = [ - "rms_quant_fusion", - "norm_rope_fusion", - "attn_quant_fusion", - "sequence_parallel", - "async_tp", - ] - - run_e2e_fusion_test( - model_name, - matches, - model_kwargs, - attn_backend, - compilation_config, - matches_check, - tp_size=2, - ) - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b_fp4], -) -@pytest.mark.parametrize("attn_backend", [FLASHINFER_ATTN]) -@pytest.mark.parametrize("n_layers", [4]) -@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) -@pytest.mark.skipif(not is_blackwell(), reason="Blackwell required for fp4") -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA") -def test_tp2_async_tp_nvfp4_fusions( - model_name: str, - matches_fn: Callable[[int], Matches], - model_kwargs: dict, - hf_overrides: Callable[[int], dict], - attn_backend: AttentionBackendCase, - n_layers: int, - custom_ops: str, - inductor_graph_partition: bool, - run_e2e_fusion_test, -): - if nvfp4_kernel_exposes_input_quant_key(): - pytest.skip( - "NVFP4 kernel exposes input_quant_key; manual fusion fires " - "instead of compiler pass-based fusion" - ) - - # NVFP4 currently wires the all-gather + GEMM path only. - matches = matches_fn(n_layers)._replace(async_tp=n_layers * 2) - - # Reduce size of model and skip weight loading time - model_kwargs["hf_overrides"] = hf_overrides(n_layers) - model_kwargs["load_format"] = "dummy" - model_kwargs["max_model_len"] = 1024 - model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} - - compilation_config = dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=custom_ops.split(","), - pass_config=PassConfig( - fuse_act_quant=False, - fuse_attn_quant=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=False, - # Override threshold for testing (models have small hidden_size) - sp_min_token_num=512, - ), - ) - - matches_check = [ - "attn_quant_fusion", - "sequence_parallel", - "async_tp", - ] - - run_e2e_fusion_test( - model_name, - matches, - model_kwargs, - attn_backend, - compilation_config, - matches_check, - tp_size=2, - ) - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b, qwen3_a3b], -) -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN]) -@pytest.mark.parametrize("n_layers", [4]) -@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) -def test_tp2_async_tp_fusions( - model_name: str, - matches_fn: Callable[[int], Matches], - model_kwargs: dict, - hf_overrides: Callable[[int], dict], - attn_backend: AttentionBackendCase, - n_layers: int, - custom_ops: str, - inductor_graph_partition: bool, - run_e2e_fusion_test, -): - matches = matches_fn(n_layers) - - # Reduce size of model and skip weight loading time - model_kwargs["hf_overrides"] = hf_overrides(n_layers) - model_kwargs["load_format"] = "dummy" - model_kwargs["max_model_len"] = 1024 - model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} - - compilation_config = dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=custom_ops.split(","), - pass_config=PassConfig( - enable_qk_norm_rope_fusion=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=False, - # Override threshold for testing (models have small hidden_size) - sp_min_token_num=512, - ), - ) - - matches_check = [ - "norm_rope_fusion", - "sequence_parallel", - "async_tp", - ] - - run_e2e_fusion_test( - model_name, - matches, - model_kwargs, - attn_backend, - compilation_config, - matches_check, - tp_size=2, - ) - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b_fp8, llama4_scout_fp8], -) -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN]) -@pytest.mark.parametrize("n_layers", [4]) -@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm")) -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) -def test_tp2_sp_ar_rms_fp8_fusions( - model_name: str, - matches_fn: Callable[[int], Matches], - model_kwargs: dict, - hf_overrides: Callable[[int], dict], - attn_backend: AttentionBackendCase, - n_layers: int, - custom_ops: str, - inductor_graph_partition: bool, - run_e2e_fusion_test, -): - matches = matches_fn(n_layers) - - # Reduce size of model and skip weight loading time - model_kwargs["hf_overrides"] = hf_overrides(n_layers) - model_kwargs["load_format"] = "dummy" - model_kwargs["max_model_len"] = 1024 - - compilation_config = dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=custom_ops.split(","), - pass_config=PassConfig( - fuse_norm_quant=True, - fuse_act_quant=False, - fuse_attn_quant=True, - enable_qk_norm_rope_fusion=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=True, - # Override threshold for testing (models have small hidden_size) - sp_min_token_num=512, - ), - ) - - matches_check = [ - "rms_quant_fusion", - "norm_rope_fusion", - "attn_quant_fusion", - "ar_rms_fusion", - "sequence_parallel", - "async_tp", - ] - - run_e2e_fusion_test( - model_name, - matches, - model_kwargs, - attn_backend, - compilation_config, - matches_check, - tp_size=2, - ) - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b, qwen3_a3b], -) -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN]) -@pytest.mark.parametrize("n_layers", [4]) -@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) -def test_tp2_sp_ar_rms_fusions( - model_name: str, - matches_fn: Callable[[int], Matches], - model_kwargs: dict, - hf_overrides: Callable[[int], dict], - attn_backend: AttentionBackendCase, - n_layers: int, - custom_ops: str, - inductor_graph_partition: bool, - run_e2e_fusion_test, -): - matches = matches_fn(n_layers) - - # Reduce size of model and skip weight loading time - model_kwargs["hf_overrides"] = hf_overrides(n_layers) - model_kwargs["load_format"] = "dummy" - model_kwargs["max_model_len"] = 1024 - - compilation_config = dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=custom_ops.split(","), - pass_config=PassConfig( - enable_qk_norm_rope_fusion=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=True, - # Override threshold for testing (models have small hidden_size) - sp_min_token_num=512, - ), - ) - - matches_check = [ - "norm_rope_fusion", - "ar_rms_fusion", - "sequence_parallel", - "async_tp", - ] - - run_e2e_fusion_test( - model_name, - matches, - model_kwargs, - attn_backend, - compilation_config, - matches_check, - tp_size=2, - ) diff --git a/tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py b/tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py deleted file mode 100644 index 3143143bec86..000000000000 --- a/tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py +++ /dev/null @@ -1,58 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Bounded ROCm model coverage for actual sequence-parallel AsyncTP rewrites.""" - -import pytest - -from tests.utils import multi_gpu_test -from vllm.config import PassConfig -from vllm.platforms import current_platform - -from .common import INDUCTOR_GRAPH_PARTITION, custom_ops_combos -from .models import TRITON_ATTN, llama3_8b - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm AsyncTP contract") -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) -@pytest.mark.parametrize("custom_ops", tuple(custom_ops_combos("rms_norm"))) -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN]) -def test_tp2_rocm_async_tp_bf16_fusions( - inductor_graph_partition, custom_ops, attn_backend, run_e2e_fusion_test -): - n_layers = 4 - model_kwargs = dict( - hf_overrides=dict( - num_hidden_layers=n_layers, - hidden_size=512, - intermediate_size=1024, - num_attention_heads=8, - num_key_value_heads=2, - head_dim=64, - ), - load_format="dummy", - max_model_len=1024, - kv_cache_memory_bytes=64 * 1024**2, - gpu_memory_utilization=0.01, - max_num_batched_tokens=1024, - max_num_seqs=32, - ) - run_e2e_fusion_test( - llama3_8b.model_name, - llama3_8b.matches(n_layers), - model_kwargs, - attn_backend, - dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=custom_ops.split(","), - pass_config=PassConfig( - enable_qk_norm_rope_fusion=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=False, - sp_min_token_num=512, - ), - ), - ["norm_rope_fusion", "sequence_parallel", "async_tp"], - tp_size=2, - ) diff --git a/tests/compile/fusions_e2e/test_tp2_rocm_moe.py b/tests/compile/fusions_e2e/test_tp2_rocm_moe.py deleted file mode 100644 index 9a228ca76c50..000000000000 --- a/tests/compile/fusions_e2e/test_tp2_rocm_moe.py +++ /dev/null @@ -1,60 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""ROCm MoE coverage for the TP2 GEMM/collective fusion contract.""" - -import pytest - -from tests.utils import multi_gpu_test -from vllm.config import PassConfig -from vllm.platforms import current_platform - -from .common import INDUCTOR_GRAPH_PARTITION -from .models import TRITON_ATTN, qwen3_a3b - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm AsyncTP contract") -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION[:1]) -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN]) -def test_tp2_rocm_moe_async_tp_bf16_fusions( - inductor_graph_partition, attn_backend, run_e2e_fusion_test -): - """MoE dense projections retain the Qwen SP/AsyncTP rewrite counts.""" - n_layers = 4 - run_e2e_fusion_test( - qwen3_a3b.model_name, - qwen3_a3b.matches(n_layers), - dict( - hf_overrides=dict( - num_hidden_layers=n_layers, - hidden_size=512, - intermediate_size=1024, - moe_intermediate_size=128, - num_attention_heads=8, - num_key_value_heads=2, - head_dim=64, - num_experts=8, - num_experts_per_tok=2, - ), - load_format="dummy", - max_model_len=1024, - kv_cache_memory_bytes=64 * 1024**2, - gpu_memory_utilization=0.01, - max_num_batched_tokens=1024, - max_num_seqs=32, - ), - attn_backend, - dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=["-rms_norm"], - pass_config=PassConfig( - enable_qk_norm_rope_fusion=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=False, - sp_min_token_num=512, - ), - ), - ["norm_rope_fusion", "sequence_parallel", "async_tp"], - tp_size=2, - ) diff --git a/tests/compile/passes/distributed/test_async_tp.py b/tests/compile/passes/distributed/test_async_tp.py deleted file mode 100644 index c0fd4c00a8f7..000000000000 --- a/tests/compile/passes/distributed/test_async_tp.py +++ /dev/null @@ -1,394 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - - -import pytest -import torch - -import vllm.envs as envs -from tests.compile.backend import TestBackend -from tests.utils import ( - multi_gpu_test, -) -from vllm.compilation.passes.fusion.collective_fusion import AsyncTPPass -from vllm.config import ( - CompilationConfig, - DeviceConfig, - ModelConfig, - PassConfig, - VllmConfig, - set_current_vllm_config, -) -from vllm.config.utils import Range -from vllm.distributed import ( - tensor_model_parallel_all_gather, - tensor_model_parallel_reduce_scatter, -) -from vllm.distributed.parallel_state import ( - init_distributed_environment, - initialize_model_parallel, -) -from vllm.platforms import current_platform -from vllm.utils.network_utils import get_file_store_init_method -from vllm.utils.system_utils import update_environment_variables -from vllm.utils.torch_utils import set_random_seed - -DEVICE_TYPE = current_platform.device_type -FP8_DTYPE = current_platform.fp8_dtype() - -prompts = [ - "Hello, my name is", - "The president of the United States is", - "The capital of France is", - "The future of AI is", -] - - -class TestMMRSModel(torch.nn.Module): - def __init__(self, hidden_size=16, dtype=torch.bfloat16): - super().__init__() - self.hidden_size = hidden_size - self.dtype = dtype - self.gate_proj = torch.nn.Parameter( - torch.empty((self.hidden_size * 2, hidden_size)), requires_grad=False - ) - # Initialize weights - torch.nn.init.normal_(self.gate_proj, std=0.02) - - def forward(self, hidden_states): - """Forward pass implementing the mm + reduce scatter in the FX graph.""" - # Reshape input - view = hidden_states.reshape(-1, self.hidden_size) - - # matrix multiplication - permute = self.gate_proj.permute(1, 0) - mm = torch.mm(view, permute) - reduce_scatter = tensor_model_parallel_reduce_scatter(mm, dim=0) - return reduce_scatter - - def ops_in_model_before(self): - return [torch.ops.vllm.reduce_scatter.default] - - def ops_in_model_after(self): - return [torch.ops.symm_mem.fused_matmul_reduce_scatter.default] - - -class TestAGMMModel(torch.nn.Module): - def __init__(self, hidden_size=16, dtype=torch.bfloat16): - super().__init__() - self.hidden_size = hidden_size - self.dtype = dtype - self.weight = torch.nn.Parameter( - torch.empty((hidden_size, hidden_size)), requires_grad=False - ) - # Initialize weights - torch.nn.init.normal_(self.weight, std=0.02) - - def forward(self, hidden_states): - """Forward pass implementing the mm + all gather in the FX graph.""" - # Reshape input - view = hidden_states.reshape(-1, self.hidden_size) - all_gather = tensor_model_parallel_all_gather(view, dim=0) - permute = self.weight.permute(1, 0) - mm = torch.mm(all_gather, permute) - return mm - - def ops_in_model_before(self): - return [torch.ops.vllm.all_gather.default] - - def ops_in_model_after(self): - return [torch.ops.symm_mem.fused_all_gather_matmul.default] - - -class _BaseScaledMMModel(torch.nn.Module): - def __init__(self, hidden_size=16, dtype=torch.bfloat16): - super().__init__() - self.hidden_size = hidden_size - self.dtype = dtype - self.weight = ( - torch.empty([hidden_size, hidden_size], dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - - # Initialize scale_b for _scaled_mm. - self.scale_b = torch.ones(1, self.hidden_size, dtype=torch.float32) - - -class TestScaledMMRSModel(_BaseScaledMMModel): - def forward(self, input: torch.Tensor): - """Forward pass implementing the scaled_mm + reduce scatter in the FX graph.""" - fp8_input = input.to(FP8_DTYPE) - scale_a = torch.ones(input.shape[0], 1, dtype=torch.float32) - scaled_mm = torch._scaled_mm( - fp8_input, - self.weight, - scale_a=scale_a, - scale_b=self.scale_b, - out_dtype=self.dtype, - ) - reduce_scatter = tensor_model_parallel_reduce_scatter(scaled_mm, dim=0) - return reduce_scatter - - def ops_in_model_before(self): - return [torch.ops.vllm.reduce_scatter.default] - - def ops_in_model_after(self): - return [torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter.default] - - -class TestAGScaledMMModel(_BaseScaledMMModel): - def forward(self, input: torch.Tensor): - """Forward pass implementing the all gather + scaled_mm in the FX graph.""" - # Reshape input - fp8_input = input.to(FP8_DTYPE) - all_gather = tensor_model_parallel_all_gather(fp8_input, dim=0) - - scale_a = torch.ones(all_gather.shape[0], 1, dtype=torch.float32) - scaled_mm = torch._scaled_mm( - all_gather, - self.weight, - scale_a=scale_a, - scale_b=self.scale_b, - out_dtype=self.dtype, - ) - return scaled_mm - - def ops_in_model_before(self): - return [torch.ops.vllm.all_gather.default] - - def ops_in_model_after(self): - return [torch.ops.symm_mem.fused_all_gather_scaled_matmul.default] - - -class TestCutlassScaledMMRSModel(_BaseScaledMMModel): - def forward(self, input: torch.Tensor): - """Forward pass implementing the cutlass_scaled_mm + reduce scatter - in the FX graph - - """ - fp8_input = input.to(FP8_DTYPE) - scale_a = torch.ones(input.shape[0], 1, dtype=torch.float32) - mm_out = torch.empty( - (fp8_input.shape[0], self.weight.shape[1]), - dtype=self.dtype, - device=input.device, - ) - torch.ops._C.cutlass_scaled_mm( - mm_out, fp8_input, self.weight, scale_a, self.scale_b, None - ) - reduce_scatter = tensor_model_parallel_reduce_scatter(mm_out, dim=0) - return reduce_scatter - - def ops_in_model_before(self): - return [torch.ops.vllm.reduce_scatter.default] - - def ops_in_model_after(self): - return [torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter.default] - - -class TestAGCutlassScaledMMModel(_BaseScaledMMModel): - def forward(self, input: torch.Tensor): - """Forward pass implementing the all gather + cutlass_scaled_mm - in the FX graph - """ - # Reshape input - fp8_input = input.to(FP8_DTYPE) - all_gather = tensor_model_parallel_all_gather(fp8_input, dim=0) - - scale_a = torch.ones(all_gather.shape[0], 1, dtype=torch.float32) - - mm_out = torch.empty( - (all_gather.shape[0], self.weight.shape[1]), - dtype=self.dtype, - device=all_gather.device, - ) - torch.ops._C.cutlass_scaled_mm( - mm_out, all_gather, self.weight, scale_a, self.scale_b, None - ) - return mm_out - - def ops_in_model_before(self): - return [torch.ops.vllm.all_gather.default] - - def ops_in_model_after(self): - return [torch.ops.symm_mem.fused_all_gather_scaled_matmul.default] - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "test_model", - [ - TestMMRSModel, - TestAGMMModel, - TestScaledMMRSModel, - TestAGScaledMMModel, - pytest.param( - TestCutlassScaledMMRSModel, - marks=pytest.mark.skipif( - not hasattr(torch.ops._C, "cutlass_scaled_mm"), - reason="Requires cutlass_scaled_mm", - ), - ), - pytest.param( - TestAGCutlassScaledMMModel, - marks=pytest.mark.skipif( - not hasattr(torch.ops._C, "cutlass_scaled_mm"), - reason="Requires cutlass_scaled_mm", - ), - ), - ], -) -@pytest.mark.parametrize("batch_size", [8]) -@pytest.mark.parametrize("seq_len", [16]) -@pytest.mark.parametrize("hidden_size", [16]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("dynamic", [True, False]) -@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA") -def test_async_tp_pass_replace( - test_model: str, - batch_size: int, - seq_len: int, - hidden_size: int, - dtype: torch.dtype, - dynamic: bool, -): - num_processes = 2 - distributed_init_method = get_file_store_init_method() - - def run_torch_spawn(fn, nprocs): - # need to use torch.mp.spawn otherwise will have problems with - # torch.distributed and cuda - torch.multiprocessing.spawn( - fn, - args=( - num_processes, - test_model, - batch_size, - seq_len, - hidden_size, - dtype, - dynamic, - distributed_init_method, - ), - nprocs=nprocs, - ) - - run_torch_spawn(async_tp_pass_on_test_model, num_processes) - - -def test_async_tp_pass_requires_full_graph_compilation(): - vllm_config = VllmConfig() - vllm_config.compilation_config.use_inductor_graph_partition = False - vllm_config.compilation_config.splitting_ops = [ - "vllm::unified_attention_with_output" - ] - - async_tp_pass = object.__new__(AsyncTPPass) - async_tp_pass.compilation_config = vllm_config.compilation_config - - with pytest.raises( - AssertionError, match="AsyncTPPass requires full-graph compilation" - ): - async_tp_pass.is_applicable_for_range(Range(start=8, end=8)) - - -def async_tp_pass_on_test_model( - local_rank: int, - world_size: int, - test_model_cls: torch.nn.Module, - batch_size: int, - seq_len: int, - hidden_size: int, - dtype: torch.dtype, - dynamic: bool, - distributed_init_method: str, - check_numerics: bool = False, -): - set_random_seed(0) - - device = torch.device(f"{DEVICE_TYPE}:{local_rank}") - torch.accelerator.set_device_index(device) - torch.set_default_device(device) - torch.set_default_dtype(dtype) - - update_environment_variables( - { - "RANK": str(local_rank), - "LOCAL_RANK": str(local_rank), - "WORLD_SIZE": str(world_size), - } - ) - - # initialize distributed - init_distributed_environment( - world_size=world_size, - rank=local_rank, - distributed_init_method=distributed_init_method, - local_rank=local_rank, - ) - - # configure vllm config for SequenceParallelismPass - vllm_config = VllmConfig() - vllm_config.compilation_config = CompilationConfig( - pass_config=PassConfig( - fuse_gemm_comms=True, - ), - ) - vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE)) - - # this is a fake model name to construct the model config - # in the vllm_config, it's not really used. - model_name = "RedHatAI/Llama-3.2-1B-Instruct-FP8" - vllm_config.model_config = ModelConfig( - model=model_name, trust_remote_code=True, dtype=dtype, seed=42 - ) - - with set_current_vllm_config(vllm_config): - initialize_model_parallel(tensor_model_parallel_size=world_size) - - async_tp_pass = AsyncTPPass(vllm_config) - backend = TestBackend(async_tp_pass) - - assert ( - async_tp_pass.compilation_config.splitting_ops - == vllm_config.compilation_config.splitting_ops - ) - assert ( - async_tp_pass.compilation_config.use_inductor_graph_partition - == vllm_config.compilation_config.use_inductor_graph_partition - ) - - model = test_model_cls(hidden_size, dtype) # Pass dtype to model constructor - - hidden_states = torch.randn( - (batch_size * seq_len, hidden_size), dtype=dtype, requires_grad=False - ) - if check_numerics: - hidden_states += local_rank * 0.125 - - if dynamic: - torch._dynamo.mark_dynamic(hidden_states, 0) - - compiled_model = torch.compile(model, backend=backend) - actual = compiled_model(hidden_states) - - if check_numerics: - # These are new checks, so no existing tolerance is relaxed. - # Allow BF16 rounding vs. eager, but keep the lifetime check exact. - expected = model(hidden_states) - torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) - snapshot = actual.clone() - changed = compiled_model(-hidden_states) - torch.testing.assert_close(changed, -expected, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(actual, snapshot, rtol=0, atol=0) - - assert async_tp_pass.matched_count == 1 - - # In pre-nodes, all gather or reduce scatter should exist, - # fused_matmul_reduce_scatter or fused_all_gather_matmul should not - backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False) - - # In post-nodes, fused_matmul_reduce_scatter or \ - # fused_all_gather_matmul should exist - backend.check_after_ops(model.ops_in_model_after()) diff --git a/tests/compile/passes/distributed/test_sequence_parallelism.py b/tests/compile/passes/distributed/test_sequence_parallelism.py deleted file mode 100644 index c40d75f6754a..000000000000 --- a/tests/compile/passes/distributed/test_sequence_parallelism.py +++ /dev/null @@ -1,343 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest -import torch - -import vllm.envs as envs -from tests.compile.backend import TestBackend -from tests.utils import TestFP8Layer, multi_gpu_test -from vllm.compilation.passes.fusion.rms_quant_fusion import RMSNormQuantFusionPass -from vllm.compilation.passes.fusion.sequence_parallelism import SequenceParallelismPass -from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass -from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass -from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass -from vllm.config import ( - CompilationConfig, - CUDAGraphMode, - DeviceConfig, - ModelConfig, - PassConfig, - VllmConfig, - get_current_vllm_config, - set_current_vllm_config, -) -from vllm.config.utils import Range -from vllm.distributed import tensor_model_parallel_all_reduce -from vllm.distributed.parallel_state import ( - init_distributed_environment, - initialize_model_parallel, -) -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - kFp8StaticTensorSym, -) -from vllm.platforms import current_platform -from vllm.utils.system_utils import update_environment_variables -from vllm.utils.torch_utils import set_random_seed - -DEVICE_TYPE = current_platform.device_type - -pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA") - -FP8_DTYPE = current_platform.fp8_dtype() -prompts = [ - "Hello, my name is", - "The president of the United States is", - "The capital of France is", - "The future of AI is", -] - - -class TestAllReduceRMSNormModel(torch.nn.Module): - def __init__(self, hidden_size=16, eps=1e-6): - super().__init__() - self.hidden_size = hidden_size - self.eps = eps - self.norm = [RMSNorm(hidden_size, eps) for i in range(4)] - self.w = [torch.rand(hidden_size, hidden_size) for _ in range(3)] - - def forward(self, x): - z = torch.relu(x) - x = resid = tensor_model_parallel_all_reduce(z) - y = self.norm[0](x) - - z2 = torch.mm(y, self.w[0]) - x2 = tensor_model_parallel_all_reduce(z2) - - y2, resid = self.norm[1](x2, resid) - - z3 = torch.mm(y2, self.w[1]) - x3 = tensor_model_parallel_all_reduce(z3) - - y3, resid = self.norm[2](x3, resid) - - z4 = torch.mm(y3, self.w[2]) - x4 = tensor_model_parallel_all_reduce(z4) - - y4, resid = self.norm[3](x4, resid) - return y4 - - def ops_in_model_before(self): - return [torch.ops.vllm.all_reduce.default] - - def ops_in_model_after(self): - return [ - torch.ops.vllm.all_gather.default, - torch.ops.vllm.reduce_scatter.default, - ] - - def ops_in_model(self): - return [ - torch.ops.vllm_ir.rms_norm, - torch.ops.vllm_ir.fused_add_rms_norm, - ] - - -class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module): - quant_key = kFp8StaticTensorSym - - def __init__(self, hidden_size=16, eps=1e-6): - super().__init__() - self.vllm_config = get_current_vllm_config() - self.hidden_size = hidden_size - self.eps = eps - self.norm = [RMSNorm(hidden_size, eps) for i in range(4)] - self.fp8_linear_layers = [ - TestFP8Layer( - weight_shape=(hidden_size, hidden_size), - activation_quant_key=self.quant_key, - weight_quant_key=self.quant_key, - input_dtype=self.vllm_config.model_config.dtype, - ) - for i in range(3) - ] - - def forward(self, hidden_states): - # avoid having graph input be an arg to a pattern directly - z = torch.relu(hidden_states) - x = resid = tensor_model_parallel_all_reduce(z) - y = self.norm[0](x) - - z2 = self.fp8_linear_layers[0](y) - - x2 = tensor_model_parallel_all_reduce(z2) - y2, resid = self.norm[1](x2, resid) - - z3 = self.fp8_linear_layers[1](y2) - - x3 = tensor_model_parallel_all_reduce(z3) - y3, resid = self.norm[2](x3, resid) # use resid here - - z4 = self.fp8_linear_layers[2](y3) - x4 = tensor_model_parallel_all_reduce(z4) - y4, resid = self.norm[3](x4, resid) # use resid here - return y4 - - def ops_in_model_after(self): - return [ - torch.ops.vllm.all_gather.default, - torch.ops.vllm.reduce_scatter.default, - ] - - def ops_in_model_before(self): - return [ - torch.ops.vllm.all_reduce.default, - ] - - def ops_in_model(self): - if self.vllm_config.compilation_config.pass_config.fuse_norm_quant: - return [torch.ops._C.fused_add_rms_norm_static_fp8_quant.default] - else: - quant_ops = ( - [torch.ops._C.static_scaled_fp8_quant.default] - if any(layer.is_quant_fp8_enabled() for layer in self.fp8_linear_layers) - else [torch.ops.aten.reciprocal] - ) - return [ - torch.ops.vllm_ir.rms_norm, - torch.ops.vllm_ir.fused_add_rms_norm, - *quant_ops, - ] - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "test_model_cls, custom_ops", - [ - (TestAllReduceRMSNormModel, "+rms_norm"), - (TestAllReduceRMSNormModel, "-rms_norm"), - (TestAllReduceRMSNormStaticQuantFP8Model, "+rms_norm,+quant_fp8"), - (TestAllReduceRMSNormStaticQuantFP8Model, "+rms_norm,-quant_fp8"), - (TestAllReduceRMSNormStaticQuantFP8Model, "-rms_norm,+quant_fp8"), - (TestAllReduceRMSNormStaticQuantFP8Model, "-rms_norm,-quant_fp8"), - ], -) -@pytest.mark.parametrize("batch_size", [8]) -@pytest.mark.parametrize("seq_len", [16]) -@pytest.mark.parametrize("hidden_size", [16]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("fuse_norm_quant", [True, False]) -@pytest.mark.parametrize("dynamic", [False, True]) -@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA") -def test_sequence_parallelism_pass( - test_model_cls: type[torch.nn.Module], - custom_ops: str, - batch_size: int, - seq_len: int, - hidden_size: int, - dtype: torch.dtype, - fuse_norm_quant: bool, - dynamic: bool, -): - num_processes = 2 - - def run_torch_spawn(fn, nprocs): - # need to use torch.mp.spawn otherwise will have problems with - # torch.distributed and cuda - torch.multiprocessing.spawn( - fn, - args=( - num_processes, - test_model_cls, - custom_ops, - batch_size, - seq_len, - hidden_size, - dtype, - fuse_norm_quant, - dynamic, - ), - nprocs=nprocs, - ) - - run_torch_spawn(sequence_parallelism_pass_on_test_model, num_processes) - - -def test_sequence_parallelism_pass_requires_full_graph_compilation(): - vllm_config = VllmConfig() - vllm_config.compilation_config.use_inductor_graph_partition = False - vllm_config.compilation_config.splitting_ops = [ - "vllm::unified_attention_with_output" - ] - - sequence_parallelism_pass = object.__new__(SequenceParallelismPass) - sequence_parallelism_pass.compilation_config = vllm_config.compilation_config - sequence_parallelism_pass.min_token_num = 1 - - with pytest.raises( - AssertionError, - match="SequenceParallelismPass requires full-graph compilation", - ): - sequence_parallelism_pass.is_applicable_for_range(Range(start=8, end=8)) - - -def sequence_parallelism_pass_on_test_model( - local_rank: int, - world_size: int, - test_model_cls: type[torch.nn.Module], - custom_ops: str, - batch_size: int, - seq_len: int, - hidden_size: int, - dtype: torch.dtype, - fuse_norm_quant: bool, - dynamic: bool, -): - set_random_seed(0) - - device = torch.device(f"{DEVICE_TYPE}:{local_rank}") - torch.accelerator.set_device_index(device) - torch.set_default_device(device) - torch.set_default_dtype(dtype) - - update_environment_variables( - { - "RANK": str(local_rank), - "LOCAL_RANK": str(local_rank), - "WORLD_SIZE": str(world_size), - "MASTER_ADDR": "localhost", - "MASTER_PORT": "12345", - } - ) - - # initialize distributed - init_distributed_environment() - - # configure vllm config for SequenceParallelismPass - custom_ops_list = custom_ops.split(",") if custom_ops else [] - compilation_config = CompilationConfig( - splitting_ops=[], # avoid automatic rms_norm enablement - cudagraph_mode=CUDAGraphMode.NONE, # avoid piecewise warnings - custom_ops=custom_ops_list, - pass_config=PassConfig( - enable_sp=True, - fuse_norm_quant=fuse_norm_quant, - eliminate_noops=True, - ), - ) # NoOp needed for fusion - device_config = DeviceConfig(device=torch.device(DEVICE_TYPE)) - - # this is a fake model name to construct the model config - # in the vllm_config, it's not really used. - model_name = "RedHatAI/Llama-3.2-1B-Instruct-FP8" - model_config = ModelConfig( - model=model_name, trust_remote_code=True, dtype=dtype, seed=42 - ) - - vllm_config = VllmConfig( - model_config=model_config, - device_config=device_config, - compilation_config=compilation_config, - ) - - with set_current_vllm_config(vllm_config): - initialize_model_parallel(tensor_model_parallel_size=world_size) - noop_pass = NoOpEliminationPass(vllm_config) - sequence_parallelism_pass = SequenceParallelismPass(vllm_config) - cleanup_pass = PostCleanupPass(vllm_config) - assert ( - sequence_parallelism_pass.compilation_config.splitting_ops - == vllm_config.compilation_config.splitting_ops - ) - assert ( - sequence_parallelism_pass.compilation_config.use_inductor_graph_partition - == vllm_config.compilation_config.use_inductor_graph_partition - ) - passes_for_backend: list[VllmInductorPass] = [ - noop_pass, - sequence_parallelism_pass, - ] - - if fuse_norm_quant: - fusion_pass = RMSNormQuantFusionPass(vllm_config) - passes_for_backend.append(fusion_pass) - - passes_for_backend.append(cleanup_pass) - - backend = TestBackend(*passes_for_backend) - - model = test_model_cls(hidden_size) - - hidden_states = torch.randn((batch_size * seq_len, hidden_size), dtype=dtype) - - if dynamic: - torch._dynamo.mark_dynamic(hidden_states, 0) - - compiled_model = torch.compile(model, backend=backend) - compiled_model(hidden_states) - - assert sequence_parallelism_pass.matched_count == 4 - - # In pre-nodes, all reduce should be there, - # reduce scatter and all gather should not - for op in model.ops_in_model_before(): - assert backend.op_count(op, before=True) == 4 - - # In post-nodes, reduce scatter and all gather should be there, - # all reduce should not - for op in model.ops_in_model_after(): - assert backend.op_count(op, before=False) == 4 - - for op in model.ops_in_model(): - assert backend.op_count(op, before=False) > 0 diff --git a/tests/compile/rocm/__init__.py b/tests/compile/rocm/__init__.py deleted file mode 100644 index 208f01a7cb5e..000000000000 --- a/tests/compile/rocm/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/compile/rocm/test_async_tp.py b/tests/compile/rocm/test_async_tp.py deleted file mode 100644 index a4462265f8e7..000000000000 --- a/tests/compile/rocm/test_async_tp.py +++ /dev/null @@ -1,75 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""ROCm AsyncTP numerical coverage kept outside the shared distributed suite.""" - -import pytest -import torch - -from tests.compile.passes.distributed.test_async_tp import ( - TestAGMMModel, - TestMMRSModel, - async_tp_pass_on_test_model, -) -from tests.utils import multi_gpu_test -from vllm.distributed import ( - tensor_model_parallel_all_gather, - tensor_model_parallel_reduce_scatter, -) -from vllm.platforms import current_platform -from vllm.utils.network_utils import get_file_store_init_method - - -class TestRocmMMRSModel(TestMMRSModel): - def forward(self, hidden_states): - view = hidden_states.reshape(-1, self.hidden_size) - mm = torch.ops.vllm.rocm_unquantized_gemm(view, self.gate_proj, None) - return tensor_model_parallel_reduce_scatter(mm, dim=0) - - def ops_in_model_before(self): - return super().ops_in_model_before() + [ - torch.ops.vllm.rocm_unquantized_gemm.default - ] - - -class TestRocmAGMMModel(TestAGMMModel): - def __init__(self, hidden_size=16, dtype=torch.bfloat16): - super().__init__(hidden_size, dtype) - self.weight = torch.nn.Parameter( - torch.empty((hidden_size * 2, hidden_size)), requires_grad=False - ) - torch.nn.init.normal_(self.weight, std=0.02) - - def forward(self, hidden_states): - view = hidden_states.reshape(-1, self.hidden_size) - gathered = tensor_model_parallel_all_gather(view, dim=0) - return torch.ops.vllm.rocm_unquantized_gemm(gathered, self.weight, None) - - def ops_in_model_before(self): - return super().ops_in_model_before() + [ - torch.ops.vllm.rocm_unquantized_gemm.default - ] - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend contract") -@pytest.mark.parametrize( - "test_model", [TestMMRSModel, TestAGMMModel, TestRocmMMRSModel, TestRocmAGMMModel] -) -@pytest.mark.parametrize("dynamic", [False, True]) -def test_rocm_async_tp_bf16_rewrite_and_correctness(test_model, dynamic): - """Real symmetric-memory AG/GEMM and GEMM/RS must replace the collectives.""" - torch.multiprocessing.spawn( - async_tp_pass_on_test_model, - args=( - 2, - test_model, - 8, - 16, - 128, - torch.bfloat16, - dynamic, - get_file_store_init_method(), - True, - ), - nprocs=2, - ) diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index 5a1ba13a57f1..4e7a756f8cbe 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -464,49 +464,41 @@ def test_should_split(): ( "cudagraph_capture_sizes", "max_cudagraph_capture_size", - "tp_size", - "enable_sp", "max_num_batched_tokens", "cudagraph_mode", "expected_max_size", ), [ - (None, None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), - ([1, 2, 4], 4, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 4), + (None, None, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), + ([1, 2, 4], 4, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 4), ( [1, 2, 4], 8, - 1, - False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, ValidationError, ), - ([1, 256], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), - ([], None, 1, False, 2048, CUDAGraphMode.NONE, 0), - (None, 0, 1, False, 2048, CUDAGraphMode.NONE, 0), + ([1, 256], None, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), + ([], None, 2048, CUDAGraphMode.NONE, 0), + (None, 0, 2048, CUDAGraphMode.NONE, 0), # truncated to nearest multiple of 8 or 16 - (None, 257, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), + (None, 257, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), # max_num_batched_tokens <= max_cudagraph_capture_size should always be # captured even if not landing on a 16-stride step - (None, 2048, 1, False, 257, CUDAGraphMode.FULL_AND_PIECEWISE, 257), + (None, 2048, 257, CUDAGraphMode.FULL_AND_PIECEWISE, 257), # max from list - ([1, 2, 4, 15], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 15), - # SP forces full-graph compilation, sizes are filtered by TP - ([1, 2, 4, 15], None, 2, True, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 4), + ([1, 2, 4, 15], None, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 15), # limited by the max_tokens - ([1, 2, 4, 15], None, 1, False, 8, CUDAGraphMode.FULL_AND_PIECEWISE, 4), + ([1, 2, 4, 15], None, 8, CUDAGraphMode.FULL_AND_PIECEWISE, 4), # the list should contain at least 1 element when use cudagraph - ([], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, ValidationError), + ([], None, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, ValidationError), # the max capturing size should be >= 1 when use cudagraph - (None, 0, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, ValidationError), + (None, 0, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, ValidationError), ], ) def test_cudagraph_sizes_post_init( cudagraph_capture_sizes, max_cudagraph_capture_size, - tp_size, - enable_sp, max_num_batched_tokens, cudagraph_mode, expected_max_size, @@ -515,10 +507,7 @@ def test_cudagraph_sizes_post_init( if expected_max_size == ValidationError: ctx = pytest.raises(expected_max_size) - with ( - ctx, - patch.object(current_platform, "device_count", return_value=tp_size), - ): + with ctx: kwargs = {} if cudagraph_capture_sizes is not None: kwargs["cudagraph_capture_sizes"] = cudagraph_capture_sizes @@ -526,18 +515,15 @@ def test_cudagraph_sizes_post_init( kwargs["max_cudagraph_capture_size"] = max_cudagraph_capture_size compilation_config = CompilationConfig( pass_config=PassConfig( - enable_sp=enable_sp, fuse_norm_quant=True, fuse_act_quant=True, eliminate_noops=True, - sp_min_token_num=512 if enable_sp else None, ), cudagraph_mode=cudagraph_mode, **kwargs, ) engine_args = EngineArgs( model="facebook/opt-125m", - tensor_parallel_size=tp_size, max_num_seqs=min(max_num_batched_tokens, 128), max_num_batched_tokens=max_num_batched_tokens, compilation_config=compilation_config, @@ -728,35 +714,6 @@ def test_default_cudagraph_capture_sizes_keep_all_sizes_bounded(): ) -def test_cudagraph_capture_sizes_respect_sequence_parallelism(): - """Sequence-parallel capture sizes stay divisible by tensor parallel size.""" - compilation_config = CompilationConfig( - cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE - ) - compilation_config.pass_config.enable_sp = True - config = _mock_config_for_cudagraph_sizes( - max_num_seqs=32, - num_speculative_tokens=16, - max_num_batched_tokens=32768, - compilation_config=compilation_config, - ) - config.parallel_config = SimpleNamespace(tensor_parallel_size=2) - config.update_sizes_for_sequence_parallelism = lambda sizes: ( - VllmConfig.update_sizes_for_sequence_parallelism(config, sizes) - ) - - with patch.object( - current_platform, - "is_device_capability_family", - return_value=False, - ): - VllmConfig._set_cudagraph_sizes(config) - - assert all(size % 2 == 0 for size in compilation_config.cudagraph_capture_sizes) - assert 544 not in compilation_config.cudagraph_capture_sizes - assert compilation_config.max_cudagraph_capture_size == 512 - - @pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA") def test_real_vllm_config_caps_widest_ngram_decode_batch(): """Real config post-init preserves the default capture ceiling.""" @@ -1046,123 +1003,6 @@ def test_blackwell_cudagraph_default(is_blackwell, expected_max_size): ) -@pytest.mark.skipif( - not current_platform.support_static_graph_mode(), - reason="Skip if not cudagraph mode supported", -) -@pytest.mark.parametrize( - ( - "cudagraph_mode", - "use_inductor_graph_partition", - "expected_enable_sp", - "expected_cudagraph_mode", - "expected_piecewise_compile", - "expected_capture_sizes", - "expected_max_size", - ), - [ - (CUDAGraphMode.PIECEWISE, False, True, CUDAGraphMode.FULL, False, [2, 4], 4), - ( - CUDAGraphMode.FULL_DECODE_ONLY, - False, - True, - CUDAGraphMode.FULL_DECODE_ONLY, - False, - [2, 4], - 4, - ), - ( - CUDAGraphMode.FULL_AND_PIECEWISE, - False, - True, - CUDAGraphMode.FULL, - False, - [2, 4], - 4, - ), - ( - CUDAGraphMode.FULL_AND_PIECEWISE, - True, - True, - CUDAGraphMode.FULL_AND_PIECEWISE, - True, - [2, 4], - 4, - ), - ], -) -def test_sequence_parallelism_requires_full_graph_compilation( - cudagraph_mode: CUDAGraphMode, - use_inductor_graph_partition: bool, - expected_enable_sp: bool, - expected_cudagraph_mode: CUDAGraphMode, - expected_piecewise_compile: bool, - expected_capture_sizes: list[int], - expected_max_size: int, -): - with patch.object(current_platform, "device_count", return_value=2): - vllm_config = VllmConfig( - parallel_config=ParallelConfig(tensor_parallel_size=2), - scheduler_config=SchedulerConfig( - max_num_seqs=128, - max_num_batched_tokens=2048, - max_model_len=2048, - is_encoder_decoder=False, - ), - ) - vllm_config.model_config = MagicMock( - dtype=torch.float16, - enforce_eager=False, - is_moe=False, - disable_cascade_attn=False, - get_hidden_size=MagicMock(return_value=4096), - ) - vllm_config.compilation_config = CompilationConfig( - mode=CompilationMode.VLLM_COMPILE, - cudagraph_capture_sizes=[1, 2, 4, 15], - max_cudagraph_capture_size=None, - compile_sizes=["cudagraph_capture_sizes"], - use_inductor_graph_partition=use_inductor_graph_partition, - pass_config=PassConfig( - enable_sp=True, - fuse_gemm_comms=True, - fuse_norm_quant=True, - fuse_act_quant=True, - eliminate_noops=True, - sp_min_token_num=512, - ), - cudagraph_mode=cudagraph_mode, - ) - vllm_config.compilation_config.set_splitting_ops_for_v1( - all2all_backend=vllm_config.parallel_config.all2all_backend, - data_parallel_size=1, - ) - vllm_config._set_compile_ranges() - vllm_config._set_cudagraph_sizes() - - assert ( - vllm_config.compilation_config.use_inductor_graph_partition - == use_inductor_graph_partition - ) - assert ( - bool(vllm_config.compilation_config.splitting_ops) == expected_piecewise_compile - ) - assert vllm_config.compilation_config.pass_config.enable_sp == expected_enable_sp - assert ( - vllm_config.compilation_config.pass_config.fuse_gemm_comms == expected_enable_sp - ) - assert vllm_config.compilation_config.cudagraph_mode == expected_cudagraph_mode - assert ( - vllm_config.compilation_config.cudagraph_capture_sizes == expected_capture_sizes - ) - assert ( - vllm_config.compilation_config.max_cudagraph_capture_size == expected_max_size - ) - assert ( - 511 in vllm_config.compilation_config.compile_ranges_endpoints - ) == expected_enable_sp - - def test_cached_compilation_config(default_vllm_config): import torch from torch._inductor.utils import run_and_get_code diff --git a/tests/compile/test_sequence_parallelism_threshold.py b/tests/compile/test_sequence_parallelism_threshold.py deleted file mode 100644 index 090b77b330ab..000000000000 --- a/tests/compile/test_sequence_parallelism_threshold.py +++ /dev/null @@ -1,192 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest - -from vllm.compilation.passes.fusion.sequence_parallelism import ( - SP_MIN_HIDDEN_SIZE, - SP_MIN_PER_GPU_SIZE_MB, - get_sequence_parallelism_threshold, -) - - -class TestGetSequenceParallelismThreshold: - """Tests for get_sequence_parallelism_threshold function.""" - - def test_non_cuda_returns_none(self, mock_cuda_platform): - """Non-CUDA platforms should return None.""" - with mock_cuda_platform(is_cuda=False): - result = get_sequence_parallelism_threshold( - hidden_size=8192, tp_size=2, element_size=2 - ) - assert result is None - - def test_unsupported_device_capability_returns_none(self, mock_cuda_platform): - """Unsupported device capabilities (e.g., sm80) should return None.""" - with mock_cuda_platform(capability=(8, 0)): - result = get_sequence_parallelism_threshold( - hidden_size=8192, tp_size=2, element_size=2 - ) - assert result is None - - def test_small_hidden_size_returns_none(self, mock_cuda_platform): - """H100 with hidden_size below threshold should return None.""" - with mock_cuda_platform(capability=(9, 0)): - result = get_sequence_parallelism_threshold( - hidden_size=4096, - tp_size=2, - element_size=2, # 4096 < 8192 - ) - assert result is None - - def test_h100_large_model_returns_threshold(self, mock_cuda_platform): - """H100 with large enough hidden_size should return calculated threshold.""" - with mock_cuda_platform(capability=(9, 0)): - hidden_size = 8192 - tp_size = 2 - element_size = 2 # float16/bfloat16 - - result = get_sequence_parallelism_threshold( - hidden_size=hidden_size, - tp_size=tp_size, - element_size=element_size, - ) - - # Verify calculation: (8 * 2 * 1024 * 1024) // (8192 * 2) = 1024 - MiB = 1024 * 1024 - expected = int( - (SP_MIN_PER_GPU_SIZE_MB[90] * tp_size * MiB) - // (hidden_size * element_size) - ) - assert result == expected - assert result == 1024 - - @pytest.mark.parametrize( - "hidden_size,tp_size,element_size,expected", - [ - # Boundary: exactly at min hidden size threshold, tp_size=1 - # (8 * 1 * 1024 * 1024) // (8192 * 2) = 512 - (8192, 1, 2, 512), - # Larger hidden size reduces token threshold - # (8 * 1 * 1024 * 1024) // (16384 * 2) = 256 - (16384, 1, 2, 256), - # Larger tp_size increases token threshold - # (8 * 4 * 1024 * 1024) // (8192 * 2) = 2048 - (8192, 4, 2, 2048), - # Larger element_size (fp32) reduces token threshold - # (8 * 2 * 1024 * 1024) // (8192 * 4) = 512 - (8192, 2, 4, 512), - ], - ) - def test_threshold_calculation_variations( - self, mock_cuda_platform, hidden_size, tp_size, element_size, expected - ): - """Test threshold calculation with various parameter combinations.""" - with mock_cuda_platform(capability=(9, 0)): - result = get_sequence_parallelism_threshold( - hidden_size=hidden_size, - tp_size=tp_size, - element_size=element_size, - ) - assert result == expected - - def test_hidden_size_boundary(self, mock_cuda_platform): - """Test behavior at the exact hidden_size boundary.""" - with mock_cuda_platform(capability=(9, 0)): - # Just below threshold - result = get_sequence_parallelism_threshold( - hidden_size=SP_MIN_HIDDEN_SIZE[90] - 1, - tp_size=2, - element_size=2, - ) - assert result is None - - # Exactly at threshold - result = get_sequence_parallelism_threshold( - hidden_size=SP_MIN_HIDDEN_SIZE[90], - tp_size=2, - element_size=2, - ) - assert result is not None - - -# XPU-specific constants (must match sequence_parallelism.py values) -_XPU_MIN_HIDDEN_SIZE = 4096 -_XPU_MIN_PER_GPU_SIZE_MB = 8.0 - - -class TestGetSequenceParallelismThresholdXPU: - """Tests for get_sequence_parallelism_threshold on XPU platform.""" - - def test_xpu_small_hidden_size_returns_none(self, mock_xpu_platform): - """XPU with hidden_size below threshold should return None.""" - with mock_xpu_platform(): - result = get_sequence_parallelism_threshold( - hidden_size=_XPU_MIN_HIDDEN_SIZE - 1, - tp_size=2, - element_size=2, - ) - assert result is None - - def test_xpu_large_model_returns_threshold(self, mock_xpu_platform): - """XPU with hidden_size >= threshold should return calculated value.""" - with mock_xpu_platform(): - hidden_size = _XPU_MIN_HIDDEN_SIZE - tp_size = 2 - element_size = 2 - result = get_sequence_parallelism_threshold( - hidden_size=hidden_size, - tp_size=tp_size, - element_size=element_size, - ) - # (8 * 2 * 1024 * 1024) // (4096 * 2) = 2048 - MiB = 1024 * 1024 - expected = int( - (_XPU_MIN_PER_GPU_SIZE_MB * tp_size * MiB) // (hidden_size * element_size) - ) - assert result == expected - assert result == 2048 - - @pytest.mark.parametrize( - "hidden_size,tp_size,element_size,expected", - [ - # (8 * 1 * 1024 * 1024) // (4096 * 2) = 1024 - (4096, 1, 2, 1024), - # (8 * 4 * 1024 * 1024) // (4096 * 2) = 4096 - (4096, 4, 2, 4096), - # (8 * 2 * 1024 * 1024) // (8192 * 2) = 1024 - (8192, 2, 2, 1024), - # (8 * 2 * 1024 * 1024) // (4096 * 4) = 1024 - (4096, 2, 4, 1024), - ], - ) - def test_xpu_threshold_calculation_variations( - self, mock_xpu_platform, hidden_size, tp_size, element_size, expected - ): - """Test XPU threshold calculation with various parameter combinations.""" - with mock_xpu_platform(): - result = get_sequence_parallelism_threshold( - hidden_size=hidden_size, - tp_size=tp_size, - element_size=element_size, - ) - assert result == expected - - def test_xpu_hidden_size_boundary(self, mock_xpu_platform): - """Test behavior at the exact XPU hidden_size boundary.""" - with mock_xpu_platform(): - # Just below threshold - result = get_sequence_parallelism_threshold( - hidden_size=_XPU_MIN_HIDDEN_SIZE - 1, - tp_size=2, - element_size=2, - ) - assert result is None - - # Exactly at threshold - result = get_sequence_parallelism_threshold( - hidden_size=_XPU_MIN_HIDDEN_SIZE, - tp_size=2, - element_size=2, - ) - assert result is not None diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index 2d57236f8efa..5416fdc8ea6d 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/tests/model_executor/test_routed_experts_capture.py @@ -582,10 +582,7 @@ def test_v2_model_runner_accepts_routed_experts(monkeypatch): use_ubatching=False, enable_elastic_ep=False, ), - compilation_config=SimpleNamespace( - mode=CompilationMode.NONE, - pass_config=SimpleNamespace(enable_sp=False), - ), + compilation_config=SimpleNamespace(mode=CompilationMode.NONE), cache_config=SimpleNamespace( kv_sharing_fast_prefill=False, mamba_cache_mode="none", diff --git a/tests/test_config.py b/tests/test_config.py index 0098ecdc3dac..ebf76c76a0c6 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -843,7 +843,6 @@ def test_breakable_cudagraph_platform_default( ("bi-off", 2048, (16, 8), 6144, 1, False), ("opt-out", 2048, (16, 8), 6144, 1, False), ("eager", 2048, (16, 8), 6144, 1, False), - ("sp", 2048, (16, 8), 6144, 1, False), ("fp16", 2048, (16, 8), 6144, 1, False), ("quantized", 2048, (16, 8), 6144, 1, False), ("untuned", 4096, (32, 32), 11008, 1, False), @@ -886,11 +885,7 @@ def test_batch_invariant_breakable_cudagraph( get_num_kv_heads=lambda pc: heads[1], ) config.parallel_config = SimpleNamespace(tensor_parallel_size=tp) - config.compilation_config = ( - CompilationConfig(pass_config=PassConfig(enable_sp=True)) - if case == "sp" - else CompilationConfig() - ) + config.compilation_config = CompilationConfig() try: assert config._maybe_enable_breakable_cudagraph() is expected if expected: @@ -1021,7 +1016,6 @@ def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1( "FakeAttentionBackend", uniform_decode_query_len=4, use_v2_model_runner=use_v2_model_runner, - tensor_parallel_size=1, ) assert cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE @@ -1104,7 +1098,6 @@ def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling(): "FakeAttentionBackend", uniform_decode_query_len=1, use_v2_model_runner=True, - tensor_parallel_size=1, kv_cache_config=kv_cache_config, max_num_reqs=256, ) @@ -1117,7 +1110,6 @@ def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling(): "FakeAttentionBackend", uniform_decode_query_len=1, use_v2_model_runner=True, - tensor_parallel_size=1, kv_cache_config=kv_cache_config, max_num_reqs=256, is_profiling=True, @@ -3359,8 +3351,6 @@ def test_vllm_config_explicit_overrides(): take precedence over callable defaults, across different models and optimization levels. """ - from vllm.config.compilation import PassConfig - quantized_model = ModelConfig("RedHatAI/Llama-3.2-1B-FP8") moe_model = ModelConfig("deepseek-ai/DeepSeek-V2-Lite") regular_model = ModelConfig("Qwen/Qwen1.5-7B") diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 90adce187448..6dc20a817794 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -378,9 +378,6 @@ def run_model(scheduler_output, intermediate_tensors): worker = SimpleNamespace( vllm_config=SimpleNamespace( - compilation_config=SimpleNamespace( - pass_config=SimpleNamespace(enable_sp=False) - ), parallel_config=SimpleNamespace( pipeline_parallel_size=2, distributed_executor_backend="mp" ), diff --git a/vllm/compilation/passes/fusion/collective_fusion.py b/vllm/compilation/passes/fusion/collective_fusion.py deleted file mode 100644 index f2aa6d7d2f1d..000000000000 --- a/vllm/compilation/passes/fusion/collective_fusion.py +++ /dev/null @@ -1,1013 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from collections.abc import Callable -from contextlib import suppress - -import torch -import torch._inductor.pattern_matcher as pm -import torch.distributed.distributed_c10d as c10d -import torch.fx as fx -from torch._inductor.pattern_matcher import PatternMatcherPass -from torch.distributed._symmetric_memory import enable_symm_mem_for_group - -from vllm.config import VllmConfig -from vllm.config.utils import Range -from vllm.distributed import get_tp_group -from vllm.distributed.parallel_state import ( - get_tensor_model_parallel_world_size, -) -from vllm.logger import init_logger -from vllm.platforms import current_platform -from vllm.utils.torch_utils import direct_register_custom_op - -from ..inductor_pass import enable_fake_mode -from ..vllm_inductor_pass import ( - VllmFusionPatternMatcherPass, - VllmInductorPass, - VllmPatternMatcherPass, - VllmPatternReplacement, -) - -FP8_DTYPE = current_platform.fp8_dtype() - -logger = init_logger(__name__) - - -def _flashinfer_scaled_mm_out( - A: torch.Tensor, - B: torch.Tensor, - *, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - out: torch.Tensor, - bias: torch.Tensor | None = None, - scale_result: torch.Tensor | None = None, - out_dtype: torch.dtype | None = None, - use_fast_accum: bool = False, -) -> None: - # Import lazily to avoid a circular import during module initialization - # when docs or other tooling import the pass without FlashInfer. - from vllm.utils.flashinfer import flashinfer_scaled_fp8_mm_out - - assert bias is None, "FlashInfer symm_mem adapter does not support bias" - assert scale_result is None, ( - "FlashInfer symm_mem adapter does not support result scaling" - ) - assert not use_fast_accum, ( - "FlashInfer symm_mem adapter does not support use_fast_accum" - ) - assert A.ndim == 2 and B.ndim == 2 and out.ndim == 2, ( - "FlashInfer symm_mem adapter expects 2D inputs and output" - ) - assert scale_a.numel() == 1 and scale_b.numel() == 1, ( - "FlashInfer symm_mem adapter only supports tensor-wise FP8 scales" - ) - - flashinfer_scaled_fp8_mm_out( - A, - B, - scale_a, - scale_b, - out=out, - out_dtype=out_dtype or out.dtype, - ) - - -def _flashinfer_fp4_mm_out( - A: torch.Tensor, - B: torch.Tensor, - *, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - out: torch.Tensor, - alpha: torch.Tensor, - out_dtype: torch.dtype | None = None, - use_8x4_sf_layout: bool = False, - backend: str = "cutlass", -) -> None: - from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm_out - - assert A.ndim == 2 and B.ndim == 2 and out.ndim == 2, ( - "FlashInfer FP4 symm_mem adapter expects 2D inputs and output" - ) - flashinfer_scaled_fp4_mm_out( - A, - B, - scale_a, - scale_b, - alpha, - out=out, - out_dtype=out_dtype or out.dtype, - use_8x4_sf_layout=use_8x4_sf_layout, - backend=backend, - ) - - -def fused_flashinfer_scaled_matmul_reduce_scatter_fake( - A: torch.Tensor, - B: torch.Tensor, - A_scale: torch.Tensor, - B_scale: torch.Tensor, - reduce_op: str, - orig_scatter_dim: int, - scatter_dim_after_maybe_reshape: int, - group_name: str, - output_shape: list[int], - out_dtype: torch.dtype | None = None, -) -> torch.Tensor: - world_size = c10d._resolve_process_group(group_name).size() - result_shape = list(output_shape) - result_shape[orig_scatter_dim] //= world_size - return torch.empty( - result_shape, - dtype=out_dtype or torch.bfloat16, - device=A.device, - ) - - -def fused_flashinfer_scaled_matmul_reduce_scatter( - A: torch.Tensor, - B: torch.Tensor, - A_scale: torch.Tensor, - B_scale: torch.Tensor, - reduce_op: str, - orig_scatter_dim: int, - scatter_dim_after_maybe_reshape: int, - group_name: str, - output_shape: list[int], - out_dtype: torch.dtype | None = None, -) -> torch.Tensor: - assert orig_scatter_dim == 0 and scatter_dim_after_maybe_reshape == 0, ( - "FlashInfer symm_mem adapter currently only supports scatter_dim=0" - ) - world_size = c10d._resolve_process_group(group_name).size() - assert A.ndim == 2 and B.ndim == 2, "FlashInfer symm_mem adapter expects 2D inputs" - assert A.is_contiguous(), "FlashInfer symm_mem adapter expects contiguous A" - assert A_scale.numel() == 1 and B_scale.numel() == 1, ( - "FlashInfer symm_mem adapter only supports tensor-wise FP8 scales" - ) - assert A.shape[0] % world_size == 0, ( - "FlashInfer symm_mem adapter expects M divisible by world size" - ) - - kwargs = { - "scale_b": B_scale, - "bias": None, - "scale_result": None, - "out_dtype": out_dtype, - "use_fast_accum": False, - } - return torch.distributed._symmetric_memory._fused_scaled_matmul_reduce_scatter_impl( - mm_out_op=_flashinfer_scaled_mm_out, - A=A, - B=B, - A_scale=A_scale, - kwargs=kwargs, - out_dtype=out_dtype, - reduce_op=reduce_op, - orig_scatter_dim=orig_scatter_dim, - scatter_dim_after_maybe_reshape=scatter_dim_after_maybe_reshape, - group_name=group_name, - output_shape=output_shape, - ) - - -def fused_all_gather_flashinfer_scaled_matmul_fake( - A_shard: torch.Tensor, - B: torch.Tensor, - A_scale: torch.Tensor, - B_scale: torch.Tensor, - gather_dim: int, - group_name: str, - out_dtype: torch.dtype | None = None, -) -> torch.Tensor: - world_size = c10d._resolve_process_group(group_name).size() - output_shape = list(A_shard.shape) - output_shape[gather_dim] *= world_size - output_shape[-1] = B.shape[1] - return torch.empty( - output_shape, - dtype=out_dtype or torch.bfloat16, - device=A_shard.device, - ) - - -def fused_all_gather_flashinfer_scaled_matmul( - A_shard: torch.Tensor, - B: torch.Tensor, - A_scale: torch.Tensor, - B_scale: torch.Tensor, - gather_dim: int, - group_name: str, - out_dtype: torch.dtype | None = None, -) -> torch.Tensor: - assert gather_dim == 0, ( - "FlashInfer symm_mem adapter currently only supports gather_dim=0" - ) - _, outputs = torch.distributed._symmetric_memory._fused_all_gather_matmul_impl( - mm_out_op=_flashinfer_scaled_mm_out, - A_shard=A_shard, - Bs=[B], - A_scale=A_scale, - kwargs_list=[ - { - "scale_b": B_scale, - "bias": None, - "scale_result": None, - "out_dtype": out_dtype, - "use_fast_accum": False, - } - ], - out_dtypes=[out_dtype], - gather_dim=gather_dim, - group_name=group_name, - return_A=False, - ) - return outputs[0] - - -def fused_all_gather_flashinfer_fp4_matmul_fake( - A_shard: torch.Tensor, - B: torch.Tensor, - A_scale_shard: torch.Tensor, - B_scale: torch.Tensor, - alpha: torch.Tensor, - gather_dim: int, - group_name: str, - out_dtype: torch.dtype | None = None, - view_a_scale_as_fp8: bool = False, - use_8x4_sf_layout: bool = False, - backend: str = "cutlass", -) -> torch.Tensor: - world_size = c10d._resolve_process_group(group_name).size() - output_shape = list(A_shard.shape) - output_shape[gather_dim] *= world_size - output_shape[-1] = B.shape[1] - return torch.empty( - output_shape, - dtype=out_dtype or torch.bfloat16, - device=A_shard.device, - ) - - -def fused_all_gather_flashinfer_fp4_matmul( - A_shard: torch.Tensor, - B: torch.Tensor, - A_scale_shard: torch.Tensor, - B_scale: torch.Tensor, - alpha: torch.Tensor, - gather_dim: int, - group_name: str, - out_dtype: torch.dtype | None = None, - view_a_scale_as_fp8: bool = False, - use_8x4_sf_layout: bool = False, - backend: str = "cutlass", -) -> torch.Tensor: - assert gather_dim == 0, ( - "FlashInfer FP4 symm_mem adapter currently only supports gather_dim=0" - ) - assert A_shard.ndim == 2 and A_scale_shard.ndim == 2 and B.ndim == 2, ( - "FlashInfer FP4 symm_mem adapter expects 2D inputs" - ) - if view_a_scale_as_fp8: - A_scale_shard = A_scale_shard.view(torch.float8_e4m3fn) - - group = c10d._resolve_process_group(group_name) - world_size = group.size() - output = A_shard.new_empty( - A_shard.shape[0] * world_size, - B.shape[1], - dtype=out_dtype or torch.bfloat16, - ) - output_shards = output.chunk(world_size) - - A = A_shard.new_empty(A_shard.shape[0] * world_size, A_shard.shape[1]) - A_scale = A_scale_shard.new_empty( - A_scale_shard.shape[0] * world_size, - A_scale_shard.shape[1], - ) - - def fp4_shard_consumer(shards: list[torch.Tensor], rank: int) -> None: - _flashinfer_fp4_mm_out( - shards[0], - B, - scale_a=shards[1], - scale_b=B_scale, - alpha=alpha, - out=output_shards[rank], - out_dtype=out_dtype, - use_8x4_sf_layout=use_8x4_sf_layout, - backend=backend, - ) - - torch.distributed._symmetric_memory._pipelined_multi_all_gather_and_consume( - [A_shard, A_scale_shard], - fp4_shard_consumer, - [A, A_scale], - group_name, - False, - ) - return output - - -direct_register_custom_op( - op_name="fused_flashinfer_scaled_matmul_reduce_scatter", - op_func=fused_flashinfer_scaled_matmul_reduce_scatter, - fake_impl=fused_flashinfer_scaled_matmul_reduce_scatter_fake, -) - -direct_register_custom_op( - op_name="fused_all_gather_flashinfer_scaled_matmul", - op_func=fused_all_gather_flashinfer_scaled_matmul, - fake_impl=fused_all_gather_flashinfer_scaled_matmul_fake, -) - -direct_register_custom_op( - op_name="fused_all_gather_flashinfer_fp4_matmul", - op_func=fused_all_gather_flashinfer_fp4_matmul, - fake_impl=fused_all_gather_flashinfer_fp4_matmul_fake, -) - - -class BasePattern: - def __init__(self, dtype: torch.dtype, device: str | None) -> None: - self.dtype = dtype - self.device = device - self.tp = get_tp_group() - self.tp_size = get_tensor_model_parallel_world_size() - - -class GEMMReduceScatterPattern(BasePattern): - rocm_gemm = False - - def get_inputs(self) -> list[torch.Tensor]: - mul = torch.empty([16, 4], device=self.device, dtype=self.dtype) - mm_weight = torch.empty([4, 4], device=self.device, dtype=self.dtype) - return [mul, mm_weight] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor: - if self.rocm_gemm: - mm = torch.ops.vllm.rocm_unquantized_gemm(mul, mm_weight, None) - else: - mm = torch.ops.aten.mm.default(mul, mm_weight) - reduce_scatter = torch.ops.vllm.reduce_scatter.default( - mm, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - return reduce_scatter - - def replacement(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor: - gemm_rs = torch.ops.symm_mem.fused_matmul_reduce_scatter( - mul, - mm_weight.t() if self.rocm_gemm else mm_weight, - "sum", - scatter_dim=0, - group_name=self.tp.device_group.group_name, - ) - - return gemm_rs - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class RocmGEMMReduceScatterPattern(GEMMReduceScatterPattern): - rocm_gemm = True - - -class AllGatherGEMMPattern(BasePattern): - rocm_gemm = False - - def get_inputs(self) -> list[torch.Tensor]: - x = torch.empty([4, 4], device=self.device, dtype=self.dtype) - weight = torch.empty([4, 4], device=self.device, dtype=self.dtype) - - return [x, weight] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - x: torch.Tensor, - weight: torch.Tensor, - ) -> torch.Tensor: - all_gather = torch.ops.vllm.all_gather.default( - x, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - - if self.rocm_gemm: - return torch.ops.vllm.rocm_unquantized_gemm(all_gather, weight, None) - return torch.ops.aten.mm.default(all_gather, weight) - - def replacement(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: - ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_matmul( - x, - [weight.t() if self.rocm_gemm else weight], - gather_dim=0, - group_name=self.tp.device_group.group_name, - ) - return mm_outputs - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class RocmAllGatherGEMMPattern(AllGatherGEMMPattern): - rocm_gemm = True - - -class ScaledMMReduceScatterPattern(BasePattern): - def get_inputs(self) -> list[torch.Tensor]: - input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - mm_weight = ( - torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - scale_a = torch.empty([16, 1], device=self.device, dtype=torch.float32) - scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32) - return [input, mm_weight, scale_a, scale_b] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - input: torch.Tensor, - mat2: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - ) -> torch.Tensor: - scaled_mm = torch.ops.aten._scaled_mm.default( - input, - mat2=mat2, - scale_a=scale_a, - scale_b=scale_b, - bias=None, - scale_result=None, - out_dtype=self.dtype, - ) - reduce_scatter = torch.ops.vllm.reduce_scatter.default( - scaled_mm, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - return reduce_scatter - - def replacement( - input: torch.Tensor, - mat2: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - ) -> torch.Tensor: - # Calculate output shape: input @ mat2 with scatter_dim reduced - output_shape = [*input.shape[:-1], mat2.shape[1]] - scatter_dim = 0 - gemm_rs = torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter( - input, - mat2, - scale_a, - scale_b, - "sum", - scatter_dim, # orig_scatter_dim - scatter_dim, # scatter_dim_after_maybe_reshape - self.tp.device_group.group_name, - output_shape, - None, # bias - None, # result_scale - self.dtype, # out_dtype - False, # use_fast_accum - ) - - return gemm_rs - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class AllGatherScaledMMPattern(BasePattern): - def get_inputs(self) -> list[torch.Tensor]: - x = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE) - weight = ( - torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - - s1 = x.shape[0] * self.tp_size - - scale_a = torch.empty([s1, 1], device=self.device, dtype=torch.float32) - scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32) - - return [x, weight, scale_a, scale_b] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - x: torch.Tensor, - weight: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - ) -> torch.Tensor: - all_gather = torch.ops.vllm.all_gather.default( - x, dim=0, world_size=self.tp_size, group_name=self.tp.unique_name - ) - - return torch.ops.aten._scaled_mm.default( - all_gather, - mat2=weight, - scale_a=scale_a, - scale_b=scale_b, - bias=None, - scale_result=None, - out_dtype=self.dtype, - ) - - def replacement( - x: torch.Tensor, - weight: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - ) -> torch.Tensor: - ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_scaled_matmul( # noqa - x, - [weight], - scale_a, - [scale_b], - gather_dim=0, - biases=[None], - result_scales=[None], - out_dtypes=[self.dtype], - use_fast_accum=[False], - group_name=self.tp.device_group.group_name, - ) - return mm_outputs - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class CutlassScaledMMReduceScatterPattern(BasePattern): - def get_inputs(self) -> list[torch.Tensor]: - input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - mm_weight = ( - torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - scale_a = torch.empty([16, 1], device=self.device, dtype=torch.float32) - scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32) - - cutlass_mm_output = torch.empty([16, 16], device=self.device, dtype=self.dtype) - return [input, mm_weight, scale_a, scale_b, cutlass_mm_output] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - input: torch.Tensor, - weight: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - cutlass_mm_output: torch.Tensor, - ) -> torch.Tensor: - cutlass_scaled_mm = torch.ops.higher_order.auto_functionalized( - torch.ops._C.cutlass_scaled_mm.default, - out=cutlass_mm_output, - a=input, - b=weight, - a_scales=scale_a, - b_scales=scale_b, - bias=None, - ) - - reduce_scatter = torch.ops.vllm.reduce_scatter.default( - cutlass_scaled_mm[1], - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - return reduce_scatter - - def replacement( - input: torch.Tensor, - mat2: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - cutlass_mm_output: torch.Tensor, - ) -> torch.Tensor: - # Calculate output shape: input @ mat2 with scatter_dim reduced - output_shape = [*input.shape[:-1], mat2.shape[1]] - scatter_dim = 0 - gemm_rs = torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter( - input, - mat2, - scale_a, - scale_b, - "sum", - scatter_dim, # orig_scatter_dim - scatter_dim, # scatter_dim_after_maybe_reshape - self.tp.device_group.group_name, - output_shape, - None, # bias - None, # result_scale - self.dtype, # out_dtype - False, # use_fast_accum - ) - - return gemm_rs - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class AllGatherCutlassScaledMMPattern(BasePattern): - def get_inputs(self) -> list[torch.Tensor]: - x = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE) - weight = ( - torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - - s1 = x.shape[0] * self.tp_size - - scale_a = torch.empty([s1, 1], device=self.device, dtype=torch.float32) - scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32) - - s2 = weight.shape[1] - output = torch.empty([s1, s2], device=self.device, dtype=self.dtype) - - return [x, weight, scale_a, scale_b, output] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - x: torch.Tensor, - weight: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - output: torch.Tensor, - ) -> torch.Tensor: - all_gather = torch.ops.vllm.all_gather.default( - x, dim=0, world_size=self.tp_size, group_name=self.tp.unique_name - ) - - cutlass_scaled_mm = torch.ops.higher_order.auto_functionalized( - torch.ops._C.cutlass_scaled_mm.default, - out=output, - a=all_gather, - b=weight, - a_scales=scale_a, - b_scales=scale_b, - bias=None, - ) - return cutlass_scaled_mm[1] - - def replacement( - x: torch.Tensor, - weight: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - output: torch.Tensor, - ) -> torch.Tensor: - ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_scaled_matmul( # noqa - x, - [weight], - scale_a, - [scale_b], - gather_dim=0, - biases=[None], - result_scales=[None], - out_dtypes=[self.dtype], - use_fast_accum=[False], - group_name=self.tp.device_group.group_name, - ) - return mm_outputs - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class FlashInferBMMFP8ReduceScatterPattern( - BasePattern, VllmPatternReplacement[..., torch.Tensor] -): - def get_inputs(self) -> list[torch.Tensor]: - a_2d = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - b_2d = ( - torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - a_scale = torch.empty([1], device=self.device, dtype=torch.float32) - b_scale = torch.empty([1], device=self.device, dtype=torch.float32) - return [a_2d, b_2d, a_scale, b_scale] - - @property - def pattern(self) -> Callable[..., torch.Tensor]: - def _pattern( - a_2d: torch.Tensor, - b_2d: torch.Tensor, - a_scale: torch.Tensor, - b_scale: torch.Tensor, - ) -> torch.Tensor: - bmm = torch.ops.vllm.bmm_fp8.default( - torch.ops.aten.unsqueeze.default(a_2d, 0), - torch.ops.aten.unsqueeze.default(b_2d, 0), - a_scale, - b_scale, - self.dtype, - "auto", - ) - output = torch.ops.aten.reshape.default(bmm, list(bmm.shape[1:])) - return torch.ops.vllm.reduce_scatter.default( - output, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - - return _pattern - - @property - def replacement(self) -> Callable[..., torch.Tensor]: - def _replacement( - a_2d: torch.Tensor, - b_2d: torch.Tensor, - a_scale: torch.Tensor, - b_scale: torch.Tensor, - ) -> torch.Tensor: - return torch.ops.vllm.fused_flashinfer_scaled_matmul_reduce_scatter.default( - a_2d, - b_2d, - a_scale, - b_scale, - "sum", - 0, - 0, - self.tp.device_group.group_name, - [a_2d.shape[0], b_2d.shape[1]], - self.dtype, - ) - - return _replacement - - -class FlashInferAllGatherBMMFP8Pattern( - BasePattern, VllmPatternReplacement[..., torch.Tensor] -): - def get_inputs(self) -> list[torch.Tensor]: - a_shard_2d = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE) - b_2d = ( - torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - a_scale = torch.empty([1], device=self.device, dtype=torch.float32) - b_scale = torch.empty([1], device=self.device, dtype=torch.float32) - return [a_shard_2d, b_2d, a_scale, b_scale] - - @property - def pattern(self) -> Callable[..., torch.Tensor]: - def _pattern( - a_shard_2d: torch.Tensor, - b_2d: torch.Tensor, - a_scale: torch.Tensor, - b_scale: torch.Tensor, - ) -> torch.Tensor: - all_gather = torch.ops.vllm.all_gather.default( - a_shard_2d, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - return torch.ops.vllm.bmm_fp8.default( - torch.ops.aten.unsqueeze.default(all_gather, 0), - torch.ops.aten.unsqueeze.default(b_2d, 0), - a_scale, - b_scale, - self.dtype, - "auto", - ) - - return _pattern - - @property - def replacement(self) -> Callable[..., torch.Tensor]: - def _replacement( - a_shard_2d: torch.Tensor, - b_2d: torch.Tensor, - a_scale: torch.Tensor, - b_scale: torch.Tensor, - ) -> torch.Tensor: - fused = torch.ops.vllm.fused_all_gather_flashinfer_scaled_matmul.default( - a_shard_2d, - b_2d, - a_scale, - b_scale, - 0, - self.tp.device_group.group_name, - self.dtype, - ) - return torch.ops.aten.unsqueeze.default(fused, 0) - - return _replacement - - -class FlashInferAllGatherFP4Pattern( - BasePattern, VllmPatternReplacement[..., torch.Tensor] -): - def __init__( - self, - dtype: torch.dtype, - device: str | None, - backend: str, - use_8x4_sf_layout: bool, - a_scale_view: str, - ) -> None: - super().__init__(dtype, device) - self.backend = backend - self.use_8x4_sf_layout = use_8x4_sf_layout - self.a_scale_view = a_scale_view - - def get_inputs(self) -> list[torch.Tensor]: - a_shard_2d = torch.empty([8, 8], device=self.device, dtype=torch.uint8) - b_2d = torch.empty([8, 16], device=self.device, dtype=torch.uint8) - a_scale_shard = torch.empty([128, 4], device=self.device, dtype=torch.int32) - b_scale = torch.empty([4, 128], device=self.device, dtype=torch.uint8) - alpha = torch.empty([], device=self.device, dtype=torch.float32) - return [ - a_shard_2d, - b_2d, - a_scale_shard, - b_scale, - alpha, - ] - - @property - def pattern(self) -> Callable[..., torch.Tensor]: - def _pattern( - a_shard_2d: torch.Tensor, - b_2d: torch.Tensor, - a_scale_shard: torch.Tensor, - b_scale: torch.Tensor, - alpha: torch.Tensor, - ) -> torch.Tensor: - all_gather_a = torch.ops.vllm.all_gather.default( - a_shard_2d, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - all_gather_a_scale = torch.ops.vllm.all_gather.default( - a_scale_shard, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - a_scale = all_gather_a_scale - if self.a_scale_view in ("float8", "float8_uint8"): - a_scale = torch.ops.aten.view.dtype(a_scale, torch.float8_e4m3fn) - if self.a_scale_view in ("uint8", "float8_uint8"): - a_scale = torch.ops.aten.view.dtype(a_scale, torch.uint8) - return torch.ops.vllm.flashinfer_mm_fp4.default( - all_gather_a, - b_2d, - a_scale, - b_scale, - alpha, - self.dtype, - self.use_8x4_sf_layout, - self.backend, - ) - - return _pattern - - @property - def replacement(self) -> Callable[..., torch.Tensor]: - def _replacement( - a_shard_2d: torch.Tensor, - b_2d: torch.Tensor, - a_scale_shard: torch.Tensor, - b_scale: torch.Tensor, - alpha: torch.Tensor, - ) -> torch.Tensor: - return torch.ops.vllm.fused_all_gather_flashinfer_fp4_matmul.default( - a_shard_2d, - b_2d, - a_scale_shard, - b_scale, - alpha, - 0, - self.tp.device_group.group_name, - self.dtype, - self.a_scale_view in ("float8", "float8_uint8"), - self.use_8x4_sf_layout, - self.backend, - ) - - return _replacement - - -class AsyncTPPass(VllmFusionPatternMatcherPass): - @enable_fake_mode - def __init__(self, config: VllmConfig) -> None: - super().__init__(config, pass_name="async_tp_pass") - - enable_symm_mem_for_group(get_tp_group().device_group.group_name) - GEMMReduceScatterPattern(self.model_dtype, self.device).register(self.pm_pass) - - AllGatherGEMMPattern(self.model_dtype, self.device).register(self.pm_pass) - if current_platform.is_rocm() and self.model_dtype == torch.bfloat16: - RocmGEMMReduceScatterPattern(self.model_dtype, self.device).register( - self.pm_pass - ) - RocmAllGatherGEMMPattern(self.model_dtype, self.device).register( - self.pm_pass - ) - - # These fusions are enabled only for bfloat16 models because - # `scaled_mm` or `cutlass_scaled_mm` with per-token (row-wise) scaling - # only supports bfloat16 as the output dtype. - if self.model_dtype == torch.bfloat16: - ScaledMMReduceScatterPattern(self.model_dtype, self.device).register( - self.pm_pass - ) - AllGatherScaledMMPattern(self.model_dtype, self.device).register( - self.pm_pass - ) - if hasattr(torch.ops._C, "cutlass_scaled_mm"): - CutlassScaledMMReduceScatterPattern( - self.model_dtype, self.device - ).register(self.pm_pass) - AllGatherCutlassScaledMMPattern(self.model_dtype, self.device).register( - self.pm_pass - ) - with suppress(ImportError): - import vllm.utils.flashinfer # noqa: F401 - if hasattr(torch.ops.vllm, "bmm_fp8"): - self.register( - FlashInferAllGatherBMMFP8Pattern(self.model_dtype, self.device) - ) - self.register( - FlashInferBMMFP8ReduceScatterPattern(self.model_dtype, self.device) - ) - if hasattr(torch.ops.vllm, "flashinfer_mm_fp4"): - for backend in ("cutlass", "cudnn"): - for a_scale_view in ("float8_uint8", "uint8"): - self.register( - FlashInferAllGatherFP4Pattern( - self.model_dtype, - self.device, - backend, - use_8x4_sf_layout=False, - a_scale_view=a_scale_view, - ) - ) - for use_8x4_sf_layout in (False, True): - for a_scale_view in ("float8",): - self.register( - FlashInferAllGatherFP4Pattern( - self.model_dtype, - self.device, - "trtllm", - use_8x4_sf_layout=use_8x4_sf_layout, - a_scale_view=a_scale_view, - ) - ) - self.register( - FlashInferAllGatherFP4Pattern( - self.model_dtype, - self.device, - "cute-dsl", - use_8x4_sf_layout=False, - a_scale_view="float8", - ) - ) - # NVFP4 reduce-scatter does not need scale communication: FP4 - # scales are consumed by the local GEMM and only BF16 partial - # outputs are reduced. Keep this PR scoped to the all-gather - # path; reduce-scatter needs a dedicated FP4 producer rather - # than the existing FP8-style helper. - - self.dump_patterns(config, self.pm_pass) - - def is_applicable_for_range(self, compile_range: Range) -> bool: - # This pass is applied on top of the sequence parallelism pass, - # which is only supported in fullgraph compilation mode. - assert ( - self.compilation_config.use_inductor_graph_partition - or not self.compilation_config.splitting_ops - ), "AsyncTPPass requires full-graph compilation" - return True - - @VllmInductorPass.time_and_log - def __call__(self, graph: fx.Graph) -> None: - self.matched_count = self.pm_pass.apply(graph) - VllmPatternMatcherPass.match_table[self.pass_name] += self.matched_count - logger.debug("Replaced %s patterns", self.matched_count) diff --git a/vllm/compilation/passes/fusion/sequence_parallelism.py b/vllm/compilation/passes/fusion/sequence_parallelism.py deleted file mode 100644 index b1a39b1e0146..000000000000 --- a/vllm/compilation/passes/fusion/sequence_parallelism.py +++ /dev/null @@ -1,620 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import functools -from collections.abc import Callable, Sequence -from typing import Any - -import torch -import torch._inductor.pattern_matcher as pm -import torch.fx as fx -from torch._higher_order_ops.auto_functionalize import auto_functionalized -from torch._inductor.pattern_matcher import PatternMatcherPass - -import vllm.ir.ops -from vllm.config import VllmConfig -from vllm.config.utils import Range -from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce -from vllm.distributed.parallel_state import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) -from vllm.logger import init_logger -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - kFp8StaticTensorSym, -) - -from ..inductor_pass import enable_fake_mode -from ..utility.noop_elimination import NoOpEliminationPass -from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass -from .matcher_utils import MatcherQuantFP8 - -logger = init_logger(__name__) - -if hasattr(torch.ops._C, "scaled_fp4_quant"): - SCALED_FP4_QUANT_OUT_OVERLOAD = torch.ops._C.scaled_fp4_quant.out - SCALED_FP4_QUANT_DEFAULT_OVERLOAD = torch.ops._C.scaled_fp4_quant.default - -# Min hidden size per device capability for sequence parallelism -# Only apply sequence parallelism for models with hidden_size >= threshold -SP_MIN_HIDDEN_SIZE: dict[int, int] = { - 90: 8192, # H100: only for models with hidden_size >= 8192 - 100: 8192, # Blackwell family: only for models with hidden_size >= 8192 -} - -# Min size per GPU per device capability for sequence parallelism -# Total min size = min_per_gpu_size * tp_size -# This ensures the threshold scales appropriately with tensor parallelism -SP_MIN_PER_GPU_SIZE_MB: dict[int, float] = { - 90: 8, # 8MB per GPU for H100 - # Use a more conservative threshold on Blackwell so TP8 starts later. - 100: 32, -} - - -def get_sequence_parallelism_threshold( - hidden_size: int, - tp_size: int, - element_size: int, -) -> int | None: - """Calculate the minimum token threshold for applying sequence parallelism. - - Returns None if sequence parallelism should not be applied based on model size. - - Branching logic based on device capability: - - Check if hidden_size >= SP_MIN_HIDDEN_SIZE[device_capability] - - If not, returns None (SP disabled for small models on this device) - - If yes, calculates threshold based on per-GPU size - - Formula: min_token_num = (min_per_gpu_size_mb * tp_size * MiB) // - (hidden_size * element_size) - """ - from vllm.platforms import current_platform - - if current_platform.is_xpu(): - min_hidden_size = 4096 - min_per_gpu_size_mb = 8.0 - elif current_platform.is_cuda(): - capability = current_platform.get_device_capability() - if capability is None: - return None - - # Collapse Blackwell variants (sm100/sm103/...) into one policy bucket. - if current_platform.is_device_capability_family(100): - device_capability = 100 - else: - device_capability = capability.to_int() - - # Check if device has configured thresholds - _hidden = SP_MIN_HIDDEN_SIZE.get(device_capability) - _gpu_mb = SP_MIN_PER_GPU_SIZE_MB.get(device_capability) - if _hidden is None or _gpu_mb is None: - return None - min_hidden_size, min_per_gpu_size_mb = _hidden, _gpu_mb - else: - return None - - # Only apply sequence parallelism for models meeting the size threshold - if hidden_size < min_hidden_size: - return None - - MiB = 1024 * 1024 - min_size = min_per_gpu_size_mb * MiB * tp_size - return int(min_size // (hidden_size * element_size)) - - -def get_first_out_wrapper( - fn: Callable[..., Sequence[torch.Tensor]], -) -> Callable[..., torch.Tensor]: - @functools.wraps(fn) - def wrapper(*args: Any) -> torch.Tensor: - return fn(*args)[0] - - return wrapper - - -class _SequenceParallelPatternHelper: - """Helper for sequence parallelism patterns.""" - - def __init__( - self, - epsilon: float, - dtype: torch.dtype, - device: str | None, - ) -> None: - self.epsilon = epsilon - self.dtype = dtype - self.device = device - self.tp_group = get_tp_group() - self.tp_size = get_tensor_model_parallel_world_size() - self.tp_rank = get_tensor_model_parallel_rank() - - def _all_reduce(self, x: torch.Tensor) -> torch.Tensor: - return tensor_model_parallel_all_reduce(x) - - def _reduce_scatter(self, x: torch.Tensor) -> torch.Tensor: - return torch.ops.vllm.reduce_scatter.default( - x, dim=0, world_size=self.tp_size, group_name=self.tp_group.unique_name - ) - - def _all_gather(self, x: torch.Tensor) -> torch.Tensor: - return torch.ops.vllm.all_gather.default( - x, dim=0, world_size=self.tp_size, group_name=self.tp_group.unique_name - ) - - def empty(self, *args: Any, **kwargs: Any) -> torch.Tensor: - return torch.empty(*args, dtype=self.dtype, device=self.device, **kwargs) - - def empty_f32(self, *args: Any, **kwargs: Any) -> torch.Tensor: - return torch.empty(*args, dtype=torch.float32, device=self.device, **kwargs) - - -class FirstAllReduceRMSNormPattern(_SequenceParallelPatternHelper): - def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None: - super().__init__(epsilon, dtype, device) - - def get_inputs(self) -> list[torch.Tensor]: - # input, weight - return [self.empty([1, 8, 4]), self.empty([4])] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - input: torch.Tensor, - weight: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(input) - rmsnorm = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon) - - return rmsnorm, all_reduce - - def replacement( - input: torch.Tensor, - weight: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - reduce_scatter = self._reduce_scatter(input) - - rmsnorm = vllm.ir.ops.rms_norm(reduce_scatter, weight, self.epsilon) - all_gather = self._all_gather(rmsnorm) - return all_gather, reduce_scatter - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper): - def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None: - super().__init__(epsilon, dtype, device) - - def get_inputs(self) -> list[torch.Tensor]: - mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype) - - residual = torch.empty([4, 4], device=self.device, dtype=self.dtype) - rms_norm_weights = torch.empty([4, 4], device=self.device, dtype=self.dtype) - - return [ - residual, - mm_1, - rms_norm_weights, - ] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - residual: torch.Tensor, - mm_1: torch.Tensor, - rms_norm_weights: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(mm_1) - rmsnorm = vllm.ir.ops.fused_add_rms_norm( - all_reduce, residual, rms_norm_weights, self.epsilon - ) - return rmsnorm[0], rmsnorm[1] - - def replacement( - residual: torch.Tensor, - mm_1: torch.Tensor, - rms_norm_weights: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - # The pattern matcher replaces from the end of the graph - # (last layer first). At the time each match is replaced, - # the preceding layer has NOT been replaced yet, so - # `residual` is still full-size and the slice below is - # correct. Once the preceding layer IS replaced, its - # residual output shrinks to [local_len, H], and this - # slice becomes semantically incorrect (e.g. for rank > 0, - # the indices would be out of bounds). However, since the - # symbolic output shape equals the input shape, - # NoOpEliminationPass (called at the end of - # SequenceParallelismPass.__call__) removes these slices - # before the graph is ever executed or compiled. - reduce_scatter = self._reduce_scatter(mm_1) - local_len = reduce_scatter.size(0) - # when the preceding VocabParallelEmbedding is excluded - # from the FX graph (e.g., passing `inputs_embeds` directly in VLMs), - # the FirstAllReduceRMSNorm pattern is never matched. we must - # perform a proper TP-aware slice here. simply using `[0:local_len]` - # would incorrectly cause all ranks to process rank 0's chunk. - residual = residual[ - self.tp_rank * local_len : self.tp_rank * local_len + local_len, ... - ] - rmsnorm = vllm.ir.ops.fused_add_rms_norm( - reduce_scatter, residual, rms_norm_weights, self.epsilon - ) - all_gather = self._all_gather(rmsnorm[0]) - # residual output is now [local_len, H]; the next layer's - # slice on it is semantically incorrect until - # NoOpEliminationPass removes it. - return all_gather, rmsnorm[1] - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - pm.register_replacement( - get_first_out_wrapper(pattern), - get_first_out_wrapper(replacement), - self.get_inputs(), - pm.fwd_only, - pm_pass, - ) - - -class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper): - def __init__( - self, - epsilon: float, - dtype: torch.dtype, - device: str | None, - ) -> None: - super().__init__(epsilon, dtype, device) - self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym) - - def get_inputs(self) -> list[torch.Tensor]: - # input, weight, scale - return [self.empty([1, 8, 4]), self.empty([4]), self.empty_f32([1, 1])] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - input: torch.Tensor, - weight: torch.Tensor, - scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(input) - rms = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon) - quant, _ = self.quant_matcher(rms, scale) - return quant, all_reduce - - def replacement( - input: torch.Tensor, - weight: torch.Tensor, - scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - reduce_scatter = self._reduce_scatter(input) - rms = vllm.ir.ops.rms_norm(reduce_scatter, weight, self.epsilon) - quant, _ = self.quant_matcher(rms, scale) - all_gather = self._all_gather(quant) - - return all_gather, reduce_scatter - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper): - def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None: - super().__init__(epsilon, dtype, device) - self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym) - - def get_inputs(self) -> list[torch.Tensor]: - mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype) - residual = torch.empty([4, 4], device=self.device, dtype=self.dtype) - rms_norm_weights = torch.empty([4, 4], device=self.device, dtype=self.dtype) - scale = torch.empty([1, 1], device=self.device, dtype=torch.float32) - - return [residual, mm_1, rms_norm_weights, scale] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - residual: torch.Tensor, - mm_1: torch.Tensor, - rms_norm_weights: torch.Tensor, - scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(mm_1) - rms, residual_out = vllm.ir.ops.fused_add_rms_norm( - all_reduce, residual, rms_norm_weights, self.epsilon - ) - quant, _ = self.quant_matcher(rms, scale) - return quant, residual_out - - def replacement( - residual: torch.Tensor, - mm_1: torch.Tensor, - rms_norm_weights: torch.Tensor, - scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - # See MiddleAllReduceRMSNormPattern.replacement for a - # detailed explanation of the temporary slice below: - # it is correct when first inserted, becomes semantically - # incorrect after the preceding layer is replaced, and is - # removed by NoOpEliminationPass before the graph is compiled. - reduce_scatter = self._reduce_scatter(mm_1) - local_len = reduce_scatter.size(0) - # when the preceding VocabParallelEmbedding is excluded - # from the FX graph (e.g., passing `inputs_embeds` directly in VLMs), - # the FirstAllReduceRMSNorm pattern is never matched. we must - # perform a proper TP-aware slice here. simply using `[0:local_len]` - # would incorrectly cause all ranks to process rank 0's chunk. - residual = residual[ - self.tp_rank * local_len : self.tp_rank * local_len + local_len, ... - ] - rms, residual_out = vllm.ir.ops.fused_add_rms_norm( - reduce_scatter, residual, rms_norm_weights, self.epsilon - ) - quant, _ = self.quant_matcher(rms, scale) - all_gather = self._all_gather(quant) - # residual output is now [local_len, H]; the next layer's - # slice on it is semantically incorrect until - # NoOpEliminationPass removes it. - return all_gather, residual_out - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - pm.register_replacement( - get_first_out_wrapper(pattern), - get_first_out_wrapper(replacement), - self.get_inputs(), - pm.fwd_only, - pm_pass, - ) - - -class FirstAllReduceRMSNormStaticNVFP4Pattern(_SequenceParallelPatternHelper): - def get_inputs(self) -> list[torch.Tensor]: - input = self.empty([8, 16]) - weight = self.empty([16]) - input_global_scale = self.empty_f32([1, 1]) - quant_output = torch.empty([8, 8], device=self.device, dtype=torch.uint8) - output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32) - return [input, weight, input_global_scale, quant_output, output_scale] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - input: torch.Tensor, - weight: torch.Tensor, - input_global_scale: torch.Tensor, - quant_output: torch.Tensor, - output_scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(input) - rms = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon) - quant = auto_functionalized( - SCALED_FP4_QUANT_OUT_OVERLOAD, - input=rms, - input_scale=input_global_scale, - is_sf_swizzled_layout=True, - output=quant_output, - output_scale=output_scale, - ) - return quant[1], all_reduce, quant[2] - - def replacement( - input: torch.Tensor, - weight: torch.Tensor, - input_global_scale: torch.Tensor, - quant_output: torch.Tensor, - output_scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - reduce_scatter = self._reduce_scatter(input) - rms = vllm.ir.ops.rms_norm(reduce_scatter, weight, self.epsilon) - rms = torch.ops.aten.view.default(rms, [-1, rms.shape[-1]]) - quant = SCALED_FP4_QUANT_DEFAULT_OVERLOAD( - rms, - input_global_scale, - True, - ) - return ( - self._all_gather(quant[0]), - reduce_scatter, - self._all_gather(quant[1]), - ) - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class MiddleAllReduceRMSNormStaticNVFP4Pattern(_SequenceParallelPatternHelper): - def get_inputs(self) -> list[torch.Tensor]: - mm_1 = self.empty([8, 16]) - residual = self.empty([8, 16]) - rms_norm_weights = self.empty([16]) - input_global_scale = self.empty_f32([1, 1]) - quant_output = torch.empty([8, 8], device=self.device, dtype=torch.uint8) - output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32) - return [ - residual, - mm_1, - rms_norm_weights, - input_global_scale, - quant_output, - output_scale, - ] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - residual: torch.Tensor, - mm_1: torch.Tensor, - rms_norm_weights: torch.Tensor, - input_global_scale: torch.Tensor, - quant_output: torch.Tensor, - output_scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(mm_1) - rms, residual_out = vllm.ir.ops.fused_add_rms_norm( - all_reduce, residual, rms_norm_weights, self.epsilon - ) - quant = auto_functionalized( - SCALED_FP4_QUANT_OUT_OVERLOAD, - input=rms, - input_scale=input_global_scale, - is_sf_swizzled_layout=True, - output=quant_output, - output_scale=output_scale, - ) - return quant[1], residual_out, quant[2] - - def replacement( - residual: torch.Tensor, - mm_1: torch.Tensor, - rms_norm_weights: torch.Tensor, - input_global_scale: torch.Tensor, - quant_output: torch.Tensor, - output_scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - # Keep this slice in sync with the non-quantized SP replacement: - # once the previous SP pattern fires, it becomes a no-op. - reduce_scatter = self._reduce_scatter(mm_1) - residual = residual[0 : reduce_scatter.size(0), ...] - rms, residual_out = vllm.ir.ops.fused_add_rms_norm( - reduce_scatter, residual, rms_norm_weights, self.epsilon - ) - rms = torch.ops.aten.view.default(rms, [-1, rms.shape[-1]]) - quant = SCALED_FP4_QUANT_DEFAULT_OVERLOAD( - rms, - input_global_scale, - True, - ) - return self._all_gather(quant[0]), residual_out, self._all_gather(quant[1]) - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class SequenceParallelismPass(VllmPatternMatcherPass): - """This pass enables sequence parallelism for models. - It identifies patterns where an AllReduce operation is followed by - an RMSNorm (or RMSNorm and then Quantization) operation. - These patterns are replaced with a ReduceScatter operation, followed by - a local RMSNorm/Quantization, and then an AllGather operation. - - The general transformation is: - Input -> AllReduce -> RMSNorm -> Output - becomes - Input -> ReduceScatter -> RMSNorm -> AllGather -> Output - - While this pass itself does not directly yield performance improvements, - it lays the groundwork for subsequent fusion passes, such as - GEMM + ReduceScatter and AllGather + GEMM fusions. These fusions can - significantly reduce communication overhead and improve overall model - performance. - - This pass is only supported when compiling the whole graph (fullgraph - mode, i.e. using Inductor graph partition or empty splitting_ops). - Piecewise compilation is not supported because the residual tensor - gets split across TP ranks, causing size mismatches at subgraph - boundaries. - - This pass splits up the residual tensor across TP ranks and hence divides - its size. The pattern matcher starts at the end of the graph (last layer - first), so when each replacement inserts a residual slice, the preceding - layer has not been replaced yet and the slice is correct. Once the - preceding layer IS replaced, its residual output shrinks and the slice - becomes semantically incorrect (out-of-bounds indices for rank > 0). - The graph is never executed in this intermediate state — - NoOpEliminationPass removes these slices based on symbolic shape equality - (input shape == output shape) before the graph is compiled. - """ - - @enable_fake_mode - def __init__(self, config: VllmConfig) -> None: - super().__init__(config) - - # Get min_token_num threshold - # Read min_token_num from config (calculated during config init) - self.min_token_num = None - if config.model_config is not None: - pass_config = config.compilation_config.pass_config - self.min_token_num = pass_config.sp_min_token_num - - if self.min_token_num is not None: - # Take the min to avoid exceeding max_num_batched_tokens - max_batched = config.scheduler_config.max_num_batched_tokens - if max_batched is not None: - self.min_token_num = min(self.min_token_num, max_batched) - logger.debug_once( - f"Sequence parallelism min token threshold: {self.min_token_num}", - scope="global", - ) - - # Used to clean up redundant views created temporarily - # to circumvent residual shape change issues - self.noop_cleanup = NoOpEliminationPass(config) - self.noop_cleanup.pass_name = f"{self.pass_name}.{self.noop_cleanup.pass_name}" - - self.patterns: PatternMatcherPass = PatternMatcherPass( - pass_name="sequence_parallelism_pass" - ) - - for epsilon in [1e-5, 1e-6]: - # RMSNorm + Static FP8 quantization patterns - FirstAllReduceRMSNormStaticFP8Pattern( - epsilon, self.model_dtype, self.device - ).register(self.patterns) - MiddleAllReduceRMSNormStaticFP8Pattern( - epsilon, self.model_dtype, self.device - ).register(self.patterns) - - if "SCALED_FP4_QUANT_OUT_OVERLOAD" in globals(): - FirstAllReduceRMSNormStaticNVFP4Pattern( - epsilon, self.model_dtype, self.device - ).register(self.patterns) - MiddleAllReduceRMSNormStaticNVFP4Pattern( - epsilon, self.model_dtype, self.device - ).register(self.patterns) - - # Normal RMSNorm patterns - FirstAllReduceRMSNormPattern( - epsilon, self.model_dtype, self.device - ).register(self.patterns) - - MiddleAllReduceRMSNormPattern( - epsilon, self.model_dtype, self.device - ).register(self.patterns) - - self.dump_patterns(config, self.patterns) - - def is_applicable_for_range(self, compile_range: Range) -> bool: - """Determines if sequence parallelism should be applied for the given - compile range. - - SP is only beneficial for larger batch sizes where the communication - overhead is amortized. For small batches, the overhead of splitting - and gathering tensors across TP ranks outweighs the benefits. - - Returns False (SP disabled) when: - - min_token_num is None (SP disabled for this device/config) - - The compile range starts below the minimum token threshold - """ - assert ( - self.compilation_config.use_inductor_graph_partition - or not self.compilation_config.splitting_ops - ), "SequenceParallelismPass requires full-graph compilation" - - # min_token_num is None when SP is disabled for this device/config - # (e.g., non-CUDA platform, unsupported GPU, or small hidden_size) - if self.min_token_num is None: - return False - - # Only apply SP when batch size meets the minimum threshold - return compile_range.start >= self.min_token_num - - @VllmInductorPass.time_and_log - def __call__(self, graph: fx.Graph) -> None: - self.matched_count = self.patterns.apply(graph) - logger.debug("Replaced %s patterns", self.matched_count) - # Clean up reshape nodes - self.noop_cleanup(graph) diff --git a/vllm/compilation/passes/pass_manager.py b/vllm/compilation/passes/pass_manager.py index df0e58b21081..197f2208a15c 100644 --- a/vllm/compilation/passes/pass_manager.py +++ b/vllm/compilation/passes/pass_manager.py @@ -35,13 +35,11 @@ RMSNormReshapeFusionPass, ) from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass - from .fusion.sequence_parallelism import SequenceParallelismPass from .utility.split_coalescing import SplitCoalescingPass if current_platform.is_cuda_alike(): from .fusion.act_quant_fusion import ActivationQuantFusionPass from .fusion.attn_quant_fusion import AttnQuantFusionPass - from .fusion.collective_fusion import AsyncTPPass from .fusion.mla_attn_quant_fusion import MLAAttnQuantFusionPass from .fusion.mla_rope_kvcache_cat_fusion import MLARoPEKVCacheCatFusionPass from .fusion.qk_norm_rope_kvcache_fusion import QkNormRopeKvCacheFusionPass @@ -156,11 +154,6 @@ def configure(self, config: VllmConfig) -> None: if self.pass_config.eliminate_noops: self.passes += [NoOpEliminationPass(config)] - if self.pass_config.enable_sp: - self.passes += [SequenceParallelismPass(config)] - if self.pass_config.fuse_gemm_comms: - self.passes += [AsyncTPPass(config)] - if enable_transformers_norm_canonicalization: self.passes += [AddRMSNormFusionPass(config)] diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index cbd132ab5b68..fc371d94b493 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -133,12 +133,6 @@ class PassConfig: """Fuse the custom Attention and MLAAttention + quant ops.""" eliminate_noops: bool = Field(default=True) """Eliminate no-op ops.""" - enable_sp: bool = None # type: ignore[assignment] - """Enable sequence parallelism. Requires TP>1. Automatically disabled - if the model's hidden_size is too small for SP to be beneficial - (threshold is device-capability dependent).""" - fuse_gemm_comms: bool = None # type: ignore[assignment] - """Enable async TP.""" fuse_allreduce_rms: bool = None # type: ignore[assignment] """Enable flashinfer allreduce fusion.""" enable_qk_norm_rope_fusion: bool = None # type: ignore[assignment] @@ -183,12 +177,6 @@ class PassConfig: 8: 1, # 1MB }, }, where key is the device capability""" - sp_min_token_num: int | None = None - """The minimum number of tokens above which vllm should use - sequence parallelism. Specified as an integer token count. - Unspecified will fallback to default values which are compute - capability and world size dependent.""" - # TODO(luka) better pass enabling system. def flashinfer_max_size(self, world_size: int) -> int | None: @@ -231,8 +219,6 @@ def compute_hash(self) -> str: "fuse_norm_quant", "fuse_act_quant", "fuse_attn_quant", - "enable_sp", - "fuse_gemm_comms", "fuse_allreduce_rms", "fuse_act_padding", "fuse_mla_dual_rms_norm", @@ -1244,25 +1230,6 @@ def set_splitting_ops_for_v1( self.cudagraph_mode = CUDAGraphMode.FULL self.splitting_ops = [] - if ( - not self.use_inductor_graph_partition - and (self.pass_config.enable_sp or self.pass_config.fuse_gemm_comms) - and self.splitting_ops - ): - logger.warning_once( - "Sequence parallelism requires full-graph compilation when " - "use_inductor_graph_partition is off. Setting splitting_ops " - "to an empty list to preserve SP and async TP." - ) - self.splitting_ops = [] - if self.cudagraph_mode.has_piecewise_cudagraphs(): - logger.warning_once( - "Sequence parallelism is incompatible with piecewise " - "cudagraph when use_inductor_graph_partition is off. " - "Setting cudagraph_mode to FULL." - ) - self.cudagraph_mode = CUDAGraphMode.FULL - # Disable CUDA graphs for DeepEP high-throughput since its not CG compatible if ( all2all_backend == "deepep_high_throughput" @@ -1403,7 +1370,6 @@ def resolve_cudagraph_mode_and_sizes( min_cg_attn_backend: str | None, uniform_decode_query_len: int = 1, use_v2_model_runner: bool = False, - tensor_parallel_size: int = 1, kv_cache_config: "KVCacheConfig | None" = None, max_num_reqs: int | None = None, is_profiling: bool = False, @@ -1530,10 +1496,7 @@ def resolve_cudagraph_mode_and_sizes( and cudagraph_mode.decode_mode() == CUDAGraphMode.FULL and uniform_decode_query_len > 1 ): - self.adjust_cudagraph_sizes_for_spec_decode( - uniform_decode_query_len, - tensor_parallel_size, - ) + self.adjust_cudagraph_sizes_for_spec_decode(uniform_decode_query_len) # For Mamba models with FULL decode cudagraphs, each decode # sequence needs one Mamba cache block. The decode cudagraph @@ -1562,25 +1525,8 @@ def resolve_cudagraph_mode_and_sizes( self.cudagraph_mode = cudagraph_mode return cudagraph_mode - def adjust_cudagraph_sizes_for_spec_decode( - self, uniform_decode_query_len: int, tensor_parallel_size: int - ): + def adjust_cudagraph_sizes_for_spec_decode(self, uniform_decode_query_len: int): multiple_of = uniform_decode_query_len - if tensor_parallel_size > 1 and self.pass_config.enable_sp: - multiple_of = max(uniform_decode_query_len, tensor_parallel_size) - if ( - multiple_of % uniform_decode_query_len != 0 - or multiple_of % tensor_parallel_size != 0 - ): - raise ValueError( - f"Can't determine cudagraph shapes that are both a " - f"multiple of {uniform_decode_query_len} " - f"(num_speculative_tokens + 1) required by spec-decode " - f"and {tensor_parallel_size} (tensor_parallel_size) " - f"required by sequence parallelism please adjust " - f"num_speculative_tokens or disable sequence parallelism" - ) - if not self.cudagraph_capture_sizes or multiple_of <= 1: return @@ -1600,7 +1546,7 @@ def adjust_cudagraph_sizes_for_spec_decode( if len(rounded_sizes) == 0: raise ValueError( f"No valid cudagraph sizes after rounding to multiple of {multiple_of} " - f"(num_speculative_tokens + 1 or tp if sequence parallelism is enabled)" + f"(num_speculative_tokens + 1)" f" please adjust num_speculative_tokens ({uniform_decode_query_len - 1}" f") or max_cudagraph_capture_size ({self.max_cudagraph_capture_size})" f" or cudagraph_capture_sizes ({self.cudagraph_capture_sizes})" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 8fc89f1591cd..3fee812700b0 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -145,12 +145,10 @@ class OptimizationLevel(IntEnum): PerformanceMode = Literal["balanced", "interactivity", "throughput"] IS_QUANTIZED = False -IS_DENSE = False # The optimizations that depend on these properties currently set to False # in all cases. # if model_config is not None: # IS_QUANTIZED = lambda c: c.model_config.is_quantized() -# IS_DENSE = lambda c: not c.model_config.is_model_moe() # See https://github.com/vllm-project/vllm/issues/25689. @@ -259,8 +257,6 @@ def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool: "fuse_act_quant": False, "fuse_allreduce_rms": False, "fuse_attn_quant": False, - "enable_sp": False, - "fuse_gemm_comms": False, "fuse_act_padding": False, "fuse_mla_dual_rms_norm": False, "fuse_rope_kvcache": False, @@ -282,8 +278,6 @@ def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool: "fuse_act_quant": enable_act_fusion, "fuse_allreduce_rms": False, "fuse_attn_quant": False, - "enable_sp": False, - "fuse_gemm_comms": False, "fuse_act_padding": enable_norm_pad_fusion, "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion, "fuse_rope_kvcache": False, @@ -305,8 +299,6 @@ def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool: "fuse_act_quant": enable_act_fusion, "fuse_allreduce_rms": enable_allreduce_rms_fusion, "fuse_attn_quant": IS_QUANTIZED, - "enable_sp": IS_DENSE, - "fuse_gemm_comms": IS_DENSE, "fuse_act_padding": enable_norm_pad_fusion, "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion, "fuse_rope_kvcache": enable_rope_kvcache_fusion, @@ -328,8 +320,6 @@ def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool: "fuse_act_quant": enable_act_fusion, "fuse_allreduce_rms": enable_allreduce_rms_fusion, "fuse_attn_quant": IS_QUANTIZED, - "enable_sp": IS_DENSE, - "fuse_gemm_comms": IS_DENSE, "fuse_act_padding": enable_norm_pad_fusion, "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion, "fuse_rope_kvcache": enable_rope_kvcache_fusion, @@ -828,9 +818,6 @@ def _uses_breakable_cudagraph_for_batch_invariance(self) -> bool: or model.dtype != torch.bfloat16 or model.quantization is not None or not current_platform.is_cuda() - # Sequence parallelism / async TP are torch.compile passes. - or self.compilation_config.pass_config.enable_sp - or self.compilation_config.pass_config.fuse_gemm_comms ): return False family = bi._get_tuned_matmul_arch_family( @@ -1770,26 +1757,6 @@ def __post_init__(self): ) self.compilation_config.mode = CompilationMode.NONE - # TODO: This is a stopgap that turns sequence parallelism / async TP - # off under batch invariance. Make the sequence-parallel reduce-scatter - # path batch-invariant instead so they can stay enabled (#56370). - # Sequence parallelism / async TP rewrite all_reduce + rms_norm into a - # reduce-scatter whose reduction order depends on the batch, so they - # are not batch-invariant. Decide this before the breakable-CUDA-graph - # default below, which declines to auto-enable when they are on. - pass_config = self.compilation_config.pass_config - if envs.VLLM_BATCH_INVARIANT and ( - pass_config.enable_sp or pass_config.fuse_gemm_comms - ): - logger.warning_once( - "Disabling sequence parallelism and async TP " - "(pass_config.enable_sp / fuse_gemm_comms) when " - "VLLM_BATCH_INVARIANT is enabled: the reduce-scatter path " - "is not batch-invariant (see vllm-project/vllm#56370)." - ) - pass_config.enable_sp = False - pass_config.fuse_gemm_comms = False - breakable_cudagraph_enabled = self._maybe_enable_breakable_cudagraph() if not breakable_cudagraph_enabled and ( @@ -1920,38 +1887,6 @@ def has_blocked_weights(): breakable_cudagraph_enabled=breakable_cudagraph_enabled ) - # async tp is built on top of sequence parallelism and requires it. - pass_config = self.compilation_config.pass_config - if pass_config.fuse_gemm_comms: - pass_config.enable_sp = True - if pass_config.enable_sp: - if self.parallel_config.tensor_parallel_size == 1: - logger.warning_once("Sequence Parallelism requires TP>1, disabling") - pass_config.enable_sp = False - pass_config.fuse_gemm_comms = False - else: - if pass_config.sp_min_token_num is None: - from vllm.compilation.passes.fusion.sequence_parallelism import ( - get_sequence_parallelism_threshold, - ) - - tp_size = self.parallel_config.tensor_parallel_size - hidden_size = self.model_config.get_hidden_size() - assert isinstance(self.model_config.dtype, torch.dtype) - element_size = self.model_config.dtype.itemsize - pass_config.sp_min_token_num = get_sequence_parallelism_threshold( - hidden_size, tp_size, element_size - ) - - if pass_config.sp_min_token_num is None: - logger.warning_once( - "Model hidden_size too small for the SP " - "threshold heuristic, disabling. To force SP, " - "set pass_config.sp_min_token_num manually." - ) - pass_config.enable_sp = False - pass_config.fuse_gemm_comms = False - from vllm.utils.torch_utils import HAS_OPAQUE_TYPE if HAS_OPAQUE_TYPE: @@ -2184,30 +2119,6 @@ def has_blocked_weights(): data_parallel_size=effective_dp_size, ) - if self.compilation_config.pass_config.enable_sp: - # With pipeline parallelism, native rms norm tracing errors due to - # incorrect residual shape. - # Use custom rms norm to unblock. In the future, - # the pass will operate on higher-level IR to avoid the issue. - # TODO: https://github.com/vllm-project/vllm/issues/27894 - if self.compilation_config.mode != CompilationMode.VLLM_COMPILE: - logger.warning_once( - "Sequence parallelism is enabled, but running in wrong " - "vllm compile mode: %s.", - self.compilation_config.mode, - ) - - if self.parallel_config.pipeline_parallel_size > 1: - if "-rms_norm" not in self.compilation_config.custom_ops: - self.compilation_config.custom_ops.append("+rms_norm") - else: - logger.warning_once( - "Sequence parallelism not supported with " - "native rms_norm when using %s, " - "this will likely lead to an error.", - "pipeline parallelism", - ) - # final check of cudagraph mode after all possible updates if current_platform.is_cuda_alike(): if ( @@ -2395,29 +2306,6 @@ def has_blocked_weights(): # Log the custom passes that are enabled self.compilation_config.pass_config.log_enabled_passes() - def update_sizes_for_sequence_parallelism(self, possible_sizes: list) -> list: - # remove the sizes that not multiple of tp_size when - # enable sequence parallelism - removed_sizes = [ - size - for size in possible_sizes - if size % self.parallel_config.tensor_parallel_size != 0 - ] - if removed_sizes: - logger.warning( - "Batch sizes %s are removed because they are not " - "multiple of tp_size %d when " - "sequence parallelism is enabled", - removed_sizes, - self.parallel_config.tensor_parallel_size, - ) - - return [ - size - for size in possible_sizes - if size % self.parallel_config.tensor_parallel_size == 0 - ] - def _set_max_num_scheduled_tokens(self): """In most cases, the scheduler may schedule a batch with as many tokens as the worker is configured to handle. @@ -2665,16 +2553,6 @@ def request_counts(max_reqs: int) -> list[int]: # de-duplicate and sort the sizes cudagraph_capture_sizes = sorted(set(cudagraph_capture_sizes)) - if ( - self.parallel_config.tensor_parallel_size > 1 - and self.compilation_config.pass_config.enable_sp - ): - # Sequence parallelism only captures TP-divisible sizes, so a - # wider non-divisible decode batch cannot be captured under SP. - cudagraph_capture_sizes = self.update_sizes_for_sequence_parallelism( - cudagraph_capture_sizes - ) - # user-specific compilation_config.max_cudagraph_capture_size get # truncated to valid_max_size when they are inconsistent. valid_max_size = ( @@ -2764,37 +2642,6 @@ def _set_compile_ranges(self): "allreduce-rms fusion will be enabled for all num_tokens." ) - # Add the compile ranges for sequence parallelism - if compilation_config.pass_config.enable_sp: - pass_config = compilation_config.pass_config - - # Calculate min_token_num if not explicitly provided - # User override works regardless of hidden_size - if pass_config.sp_min_token_num is None: - from vllm.compilation.passes.fusion.sequence_parallelism import ( - get_sequence_parallelism_threshold, - ) - - tp_size = self.parallel_config.tensor_parallel_size - hidden_size = self.model_config.get_hidden_size() - assert isinstance(self.model_config.dtype, torch.dtype) - element_size = self.model_config.dtype.itemsize - pass_config.sp_min_token_num = get_sequence_parallelism_threshold( - hidden_size, tp_size, element_size - ) - - min_token_num = pass_config.sp_min_token_num - max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens - if min_token_num is not None and ( - max_num_batched_tokens is not None - and min_token_num < max_num_batched_tokens - and min_token_num > 1 - ): - # Add endpoint at min_token_num - 1 to ensure SP applies - # starting from min_token_num - # This creates ranges: [1, min-1] (no SP), [min, max] (SP applies) - computed_compile_ranges_endpoints.append(min_token_num - 1) - if compilation_config.pass_config.fuse_rope_kvcache: max_token_num = ( compilation_config.pass_config.rope_kvcache_fusion_max_token_num @@ -3124,12 +2971,6 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: unsupported: list[str] = [] speculative_config = self.speculative_config - if ( - self.compilation_config.pass_config.enable_sp - and self.parallel_config.tensor_parallel_size > 1 - ): - unsupported.append("sequence parallelism") - # V2 does not implement the external_launcher (torchrun) PP-output # broadcast that V1 uses to keep all ranks in sync (broadcast_pp_output). if ( diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 6b7db6d0dc55..355a15094db0 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -352,7 +352,6 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: pass_config = compilation_config.pass_config fusion_passes_to_disable = { - "fuse_gemm_comms": "Async TP", "fuse_allreduce_rms": "AllReduce + RMSNorm fusion", "fuse_attn_quant": "Attention + quant fusion", "fuse_act_padding": "Activation + padding fusion", diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index f9d6c4439f56..9b615cd2fa4b 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -1065,47 +1065,6 @@ def flashinfer_scaled_fp4_mm( ) -def flashinfer_scaled_fp4_mm_out( - a: torch.Tensor, - b: torch.Tensor, - block_scale_a: torch.Tensor, - block_scale_b: torch.Tensor, - alpha: torch.Tensor, - out: torch.Tensor, - out_dtype: torch.dtype | None, - use_8x4_sf_layout: bool, - backend: str, -) -> torch.Tensor: - assert a.ndim == 2 and b.ndim == 2 and out.ndim == 2 - assert block_scale_a.ndim == 2 and block_scale_b.ndim == 2 - assert a.stride(-1) == 1 - assert a.shape[1] == b.shape[0] - assert out.shape == (a.shape[0], b.shape[1]) - assert out.device.type == "cuda" - - if backend in ("cutlass", "cudnn"): - if block_scale_a.dtype != torch.uint8: - block_scale_a = block_scale_a.view(torch.uint8) - if block_scale_b.dtype != torch.uint8: - block_scale_b = block_scale_b.view(torch.uint8) - - from flashinfer import mm_fp4 as flashinfer_mm_fp4_ - - flashinfer_mm_fp4_( - a, - b, - block_scale_a, - block_scale_b, - alpha, - out_dtype or out.dtype, - out=out, - block_size=16, - use_8x4_sf_layout=use_8x4_sf_layout, - backend=backend, - ) - return out - - def flashinfer_scaled_fp8_mm( a: torch.Tensor, b: torch.Tensor, @@ -1136,38 +1095,6 @@ def flashinfer_scaled_fp8_mm( return output -def flashinfer_scaled_fp8_mm_out( - a: torch.Tensor, - b: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - out: torch.Tensor, - out_dtype: torch.dtype | None = None, -) -> torch.Tensor: - assert a.ndim == 2 and b.ndim == 2 and out.ndim == 2 - assert a.shape[1] == b.shape[0] - assert out.shape == (a.shape[0], b.shape[1]) - assert scale_a.numel() == 1 and scale_b.numel() == 1 - assert a.dtype == torch.float8_e4m3fn and b.dtype == torch.float8_e4m3fn - assert out.device.type == "cuda" - assert a.is_contiguous() - - from flashinfer import bmm_fp8 as bmm_fp8_ - - bmm_fp8_( - a.unsqueeze(0), - # FlashInfer expects the weight in the same column-major view layout - # consumed by flashinfer_scaled_fp8_mm, so keep the transposed view. - b.unsqueeze(0), - scale_a, - scale_b, - out_dtype or out.dtype, - out.unsqueeze(0), - "auto", - ) - return out - - def flashinfer_quant_nvfp4_8x4_sf_layout( a: torch.Tensor, a_global_sf: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: @@ -1310,9 +1237,7 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool: "use_trtllm_attention", "flashinfer_mxfp4_quantize", "flashinfer_scaled_fp4_mm", - "flashinfer_scaled_fp4_mm_out", "flashinfer_scaled_fp8_mm", - "flashinfer_scaled_fp8_mm_out", "flashinfer_quant_nvfp4_8x4_sf_layout", "flashinfer_fp8_blockscale_gemm", "should_use_flashinfer_for_blockscale_fp8_gemm", diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 373042a41502..61d6197ceb9d 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -730,7 +730,6 @@ def initialize_kv_cache( attn_cg_support.min_cg_attn_backend, self.decode_query_len, use_v2_model_runner=True, - tensor_parallel_size=self.parallel_config.tensor_parallel_size, kv_cache_config=self.kv_cache_config, max_num_reqs=self.max_num_reqs, is_profiling=is_profiling, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 0cc3fcc29a8f..21b0f67e2fe8 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -116,7 +116,7 @@ from vllm.utils import length_from_prompt_token_ids_or_embeds from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture from vllm.utils.gpu_sync_debug import gpu_sync_allowed -from vllm.utils.math_utils import cdiv, round_up +from vllm.utils.math_utils import cdiv from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib from vllm.utils.nvtx_pytorch_hooks import PytHooks from vllm.utils.platform_utils import num_compute_units @@ -222,7 +222,6 @@ ) from vllm.v1.worker.utils import ( EncoderTimingStats, - is_residual_scattered_for_sp, raise_if_nan_logits, ) from vllm.v1.worker.workspace import lock_workspace @@ -3340,20 +3339,9 @@ def sync_and_gather_intermediate_tensors( ) -> IntermediateTensors: assert self.intermediate_tensors is not None - tp = self.vllm_config.parallel_config.tensor_parallel_size - is_rs = is_residual_scattered_for_sp(self.vllm_config, num_tokens) - - # When sequence parallelism is enabled, the "residual" tensor is - # sharded across TP ranks. All-gather it here because downstream - # QKV + Attention needs the full residual before the SP split point. if sync_self: assert intermediate_tensors is not None for k, v in intermediate_tensors.items(): - is_scattered = k == "residual" and is_rs - if is_scattered: - local_len = num_tokens // tp - v = get_tp_group().all_gather(v[:local_len], dim=0) - self.intermediate_tensors[k][:num_tokens].copy_( v[:num_tokens], non_blocking=True ) @@ -3448,14 +3436,6 @@ def _pool( async_output_copy_stream=self._get_or_create_async_output_copy_stream(), ) - def _pad_for_sequence_parallelism(self, num_scheduled_tokens: int) -> int: - # Pad tokens to multiple of tensor_parallel_size when - # enabled collective fusion for SP - tp_size = self.vllm_config.parallel_config.tensor_parallel_size - if self.compilation_config.pass_config.enable_sp and tp_size > 1: - return round_up(num_scheduled_tokens, tp_size) - return num_scheduled_tokens - def _prepare_padding_mask( self, num_tokens_unpadded: int, num_tokens_padded: int ) -> torch.Tensor: @@ -3944,8 +3924,6 @@ def _determine_batch_execution_and_padding( ) has_lora = num_active_loras > 0 if force_has_lora is None else force_has_lora - num_tokens_padded = self._pad_for_sequence_parallelism(num_tokens) - def dispatch_cudagraph(num_tokens, disable_full=False, valid_modes=None): return self.cudagraph_dispatcher.dispatch( num_tokens=num_tokens, @@ -3957,18 +3935,9 @@ def dispatch_cudagraph(num_tokens, disable_full=False, valid_modes=None): ) cudagraph_mode, batch_descriptor = dispatch_cudagraph( - num_tokens_padded, disable_full=use_cascade_attn or has_encoder_output + num_tokens, disable_full=use_cascade_attn or has_encoder_output ) num_tokens_padded = batch_descriptor.num_tokens - if self.compilation_config.pass_config.enable_sp: - assert ( - batch_descriptor.num_tokens - % self.vllm_config.parallel_config.tensor_parallel_size - == 0 - ), ( - "Sequence parallelism requires num_tokens to be " - "a multiple of tensor parallel size" - ) # Extra coordination when running data-parallel since we need to coordinate # across ranks @@ -4453,15 +4422,9 @@ def execute_model( sample_hidden_states = hidden_states[logits_indices] if not get_pp_group().is_last_rank: - all_gather_tensors = { - "residual": not is_residual_scattered_for_sp( - self.vllm_config, num_tokens_padded - ) - } get_pp_group().send_tensor_dict( hidden_states.tensors, all_gather_group=get_tp_group(), - all_gather_tensors=all_gather_tensors, ) logits = None else: @@ -7135,7 +7098,6 @@ def _check_and_update_cudagraph_mode( min_cg_attn_backend, self.uniform_decode_query_len, use_v2_model_runner=False, - tensor_parallel_size=self.parallel_config.tensor_parallel_size, kv_cache_config=self.kv_cache_config, max_num_reqs=self.max_num_reqs, is_profiling=is_profiling, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index d4c64c554273..70427d85a103 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -12,7 +12,6 @@ from types import NoneType from typing import TYPE_CHECKING, Any, cast -import numpy as np import torch import torch.nn as nn @@ -98,7 +97,6 @@ maybe_apply_startup_plan, maybe_save_startup_plan, ) -from vllm.v1.worker.utils import is_residual_scattered_for_sp from vllm.v1.worker.worker_base import CompilationTimes, WorkerBase from vllm.v1.worker.workspace import init_workspace_manager @@ -1273,45 +1271,10 @@ def execute_model( intermediate_tensors = None forward_pass = scheduler_output.total_num_scheduled_tokens > 0 - num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens - all_gather_tensors = {} - compilation_config = self.vllm_config.compilation_config - parallel_config = self.vllm_config.parallel_config - - if ( - parallel_config.pipeline_parallel_size > 1 - and compilation_config.pass_config.enable_sp - and forward_pass - ): - # currently only supported by V1 GPUModelRunner - assert not self.use_v2_model_runner - num_scheduled_tokens_np = np.array( - list(scheduler_output.num_scheduled_tokens.values()), - dtype=np.int32, - ) - # TODO(lucas): This is pretty gross; ideally we should only ever call - # `_determine_batch_execution_and_padding` once (will get called again - # in `execute_model`) but this requires a larger refactor of PP. - _, batch_desc, _, _, _ = ( - self.model_runner._determine_batch_execution_and_padding( - num_tokens=num_scheduled_tokens, - num_reqs=len(num_scheduled_tokens_np), - num_scheduled_tokens_np=num_scheduled_tokens_np, - max_num_scheduled_tokens=num_scheduled_tokens_np.max(), - use_cascade_attn=False, # TODO(lucas): Handle cascade attention - ) - ) - all_gather_tensors = { - "residual": not is_residual_scattered_for_sp( - self.vllm_config, batch_desc.num_tokens - ) - } - if forward_pass and not get_pp_group().is_first_rank: tensor_dict, comm_handles, comm_postprocess = ( get_pp_group().irecv_tensor_dict( all_gather_group=get_tp_group(), - all_gather_tensors=all_gather_tensors, ) ) assert tensor_dict is not None @@ -1349,7 +1312,6 @@ def execute_model( handles = get_pp_group().isend_tensor_dict( output.tensors, all_gather_group=get_tp_group(), - all_gather_tensors=all_gather_tensors, ) self._pp_send_work = handles[1:] diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index b61c4cebed18..db0f81b81c36 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -10,7 +10,7 @@ import numpy as np import torch -from vllm.config import CacheConfig, VllmConfig +from vllm.config import CacheConfig from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.mamba.mamba_mixer2 import share_replayssm_ring_trackers @@ -813,35 +813,6 @@ def get_uniform_decode_token_count( return None -def is_residual_scattered_for_sp( - vllm_config: VllmConfig, num_input_tokens: int -) -> bool: - """Check if the residual tensor is scattered for sequence parallelism. - - The residual tensor is scattered across tensor parallel ranks when sequence - parallelism and tensor parallelism is enabled. SP is only supported in - full-graph compilation mode. - """ - if not vllm_config.compilation_config.pass_config.enable_sp: - return False - - tp = vllm_config.parallel_config.tensor_parallel_size - - if tp == 1: - return False - - assert ( - vllm_config.compilation_config.use_inductor_graph_partition - or not vllm_config.compilation_config.splitting_ops - ), "Sequence parallelism requires full-graph compilation" - - # When sequence parallelism is enabled, we always pad num_input_tokens - # to be a multiple of tensor_parallel_size (tp) earlier. - assert num_input_tokens % tp == 0 - - return True - - @dataclass class EncoderTimingStats: """Per-request timing statistics for encoder forward pass.""" From 0751f5272fb323e0a5c62a6157f0a3c1d93e6bc5 Mon Sep 17 00:00:00 2001 From: Haokai Ma Date: Thu, 8 Oct 2026 08:29:04 -0700 Subject: [PATCH 106/624] [Bugfix] Scope deprecated CLI argument warnings to the parser that defines them (#59917) Signed-off-by: Tyler Ma <158780265+TylerMa-debugg@users.noreply.github.com> Co-authored-by: Claude --- tests/utils_/test_argparse_utils.py | 23 +++++++++++++++++++++++ vllm/utils/argparse_utils.py | 13 +++++++++---- 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/tests/utils_/test_argparse_utils.py b/tests/utils_/test_argparse_utils.py index 89e0f9977d0c..98f2fe790c12 100644 --- a/tests/utils_/test_argparse_utils.py +++ b/tests/utils_/test_argparse_utils.py @@ -4,6 +4,7 @@ import json import os +import sys from argparse import BooleanOptionalAction import pytest @@ -252,6 +253,28 @@ def test_duplicate_dict_args(caplog_vllm, parser): assert "--optimization-level" in caplog_vllm.text +@pytest.mark.skipif( + sys.version_info >= (3, 13), reason="argparse handles deprecated natively" +) +def test_deprecated_arg_scoped_to_subcommand(caplog_vllm, disable_log_dedup): + """A deprecated arg of one subcommand must not warn for a sibling + subcommand that defines a non-deprecated arg with the same dest.""" + parser = FlexibleArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + run_batch = subparsers.add_parser("run-batch") + run_batch.add_argument_group("Frontend").add_argument( + "--url", default="0.0.0.0", deprecated=True + ) + chat = subparsers.add_parser("chat") + chat.add_argument("--url", default="http://localhost:8000/v1") + + parser.parse_args(["chat", "--url", "http://localhost:9000/v1"]) + assert "is deprecated" not in caplog_vllm.text + + parser.parse_args(["run-batch", "--url", "1.2.3.4"]) + assert caplog_vllm.text.count("argument 'url' is deprecated") == 1 + + def test_model_specification( parser_with_config, cli_config_file, cli_config_file_with_model ): diff --git a/vllm/utils/argparse_utils.py b/vllm/utils/argparse_utils.py index b91b1f26fd13..9ad03453fd20 100644 --- a/vllm/utils/argparse_utils.py +++ b/vllm/utils/argparse_utils.py @@ -119,7 +119,6 @@ def add_arguments(self, actions): class FlexibleArgumentParser(ArgumentParser): """ArgumentParser that allows both underscore and dash in names.""" - _deprecated: set[Action] = set() _json_tip: str = ( "When passing JSON CLI arguments, the following sets of arguments " "are equivalent:\n" @@ -132,6 +131,8 @@ class FlexibleArgumentParser(ArgumentParser): _search_keyword: str | None = None def __init__(self, *args, **kwargs): + # Must exist before super().__init__(), which already adds arguments. + self._deprecated: set[Action] = set() # Set the default "formatter_class" to SortedHelpFormatter if "formatter_class" not in kwargs: kwargs["formatter_class"] = SortedHelpFormatter @@ -160,7 +161,7 @@ def error(self, message: str) -> NoReturn: def parse_known_args(self, args=None, namespace=None): namespace, args = super().parse_known_args(args, namespace) - for action in FlexibleArgumentParser._deprecated: + for action in self._deprecated: if ( hasattr(namespace, dest := action.dest) and getattr(namespace, dest) != action.default @@ -172,15 +173,19 @@ def add_argument(self, *args, **kwargs): deprecated = kwargs.pop("deprecated", False) action = super().add_argument(*args, **kwargs) if deprecated: - FlexibleArgumentParser._deprecated.add(action) + self._deprecated.add(action) return action class _FlexibleArgumentGroup(_ArgumentGroup): + def __init__(self, container, *args, **kwargs): + super().__init__(container, *args, **kwargs) + self._deprecated: set[Action] = container._deprecated + def add_argument(self, *args, **kwargs): deprecated = kwargs.pop("deprecated", False) action = super().add_argument(*args, **kwargs) if deprecated: - FlexibleArgumentParser._deprecated.add(action) + self._deprecated.add(action) return action def add_argument_group(self, *args, **kwargs): From e41d5d55439808e79326009599d5d83e157bf2d1 Mon Sep 17 00:00:00 2001 From: simon-veitner-redhat Date: Thu, 8 Oct 2026 17:29:20 +0200 Subject: [PATCH 107/624] [Bugfix] Don't warn about missing watermark config for default requests (#60412) Signed-off-by: Simon Veitner Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Lucas Wilkinson --- tests/watermarking/test_compatibility.py | 5 +++-- vllm/config/vllm.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/watermarking/test_compatibility.py b/tests/watermarking/test_compatibility.py index 748223d94b56..48c27db09de2 100644 --- a/tests/watermarking/test_compatibility.py +++ b/tests/watermarking/test_compatibility.py @@ -93,7 +93,7 @@ def test_default_watermarking_resolves_when_configured(): assert enabled.watermarking is None -@pytest.mark.parametrize("watermarking", [None, True]) +@pytest.mark.parametrize("watermarking", [True]) def test_requested_watermarking_without_engine_config_warns_and_disables( watermarking, caplog_vllm, reset_warning_once ): @@ -148,6 +148,7 @@ def test_incompatible_modes_are_allowed_when_watermarking_is_disabled(params): @pytest.mark.parametrize( "params", [ + SamplingParams(), SamplingParams(temperature=0), SamplingParams(trace_decode_token_ids=[1]), ], @@ -158,7 +159,7 @@ def test_requests_without_engine_config_are_unwatermarked( with caplog_vllm.at_level("WARNING"): resolved = _validate(params, server_uses_watermarking=False) - assert "engine has no watermark configuration" in caplog_vllm.text + assert "engine has no watermark configuration" not in caplog_vllm.text assert resolved is False assert params.watermarking is None diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 3fee812700b0..5d38b6765308 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1310,7 +1310,7 @@ def _check_supports_watermarking( ) -> bool: watermark_config = getattr(self, "watermark_config", None) if watermark_config is None: - if config is not None and config.watermarking is not False: + if config is not None and config.watermarking is True: logger.warning_once( "Watermarking is enabled for this request, but the engine has no " "watermark configuration. This and subsequent requests will run " From 0fc0cefb4eecf685c4bbd20912db9c1294d84c1e Mon Sep 17 00:00:00 2001 From: Justin Barlow <7@netlobo.com> Date: Thu, 8 Oct 2026 09:34:55 -0600 Subject: [PATCH 108/624] [Bugfix] Don't int-sort attention-type names in --kv-cache-dtype-skip-layers for packed KV dtypes (#60364) Signed-off-by: lobo235 <7@netlobo.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- vllm/engine/arg_utils.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 94dd3e2ef1c8..5f32ec8deb9f 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -2285,9 +2285,11 @@ def create_engine_config( ) boundary = TurboQuantConfig.get_boundary_skip_layers(model_config) - existing = set(cache_config.kv_cache_dtype_skip_layers) + # Sorted so the order does not depend on set iteration: the list is + # a CacheConfig hash factor. Entries can also be attention type + # names (e.g. "sliding_window"), so sort them as strings. cache_config.kv_cache_dtype_skip_layers = sorted( - existing | set(boundary), key=int + set(cache_config.kv_cache_dtype_skip_layers) | set(boundary) ) ray_runtime_env = None From e2414ec1dbb177b54ac1bbce60e94769ebacbb28 Mon Sep 17 00:00:00 2001 From: Mikko Tukiainen Date: Thu, 8 Oct 2026 18:47:56 +0300 Subject: [PATCH 109/624] [ROCm][Qwen4Exp] Use fused PLE Triton kernels on the AMD backend (#60021) Signed-off-by: Mikko Tukiainen Co-authored-by: Cursor --- tests/models/qwen4_exp/test_ple.py | 57 +- vllm/models/qwen4_exp/amd/model.py | 18 +- vllm/models/qwen4_exp/amd/mtp.py | 4 +- vllm/models/qwen4_exp/amd/ple_layer.py | 784 +++++------------- vllm/models/qwen4_exp/common/ops/__init__.py | 4 + .../qwen4_exp/{nvidia => common}/ops/ple.py | 0 .../qwen4_exp/nvidia/ngram_embedding.py | 2 +- vllm/models/qwen4_exp/nvidia/ple_layer.py | 2 +- 8 files changed, 279 insertions(+), 592 deletions(-) create mode 100644 vllm/models/qwen4_exp/common/ops/__init__.py rename vllm/models/qwen4_exp/{nvidia => common}/ops/ple.py (100%) diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 1ad44e830457..d7ee461d8f0d 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -45,6 +45,7 @@ Qwen4ExpPLEUnquantizedEmbeddingMethod, ) from vllm.models.qwen4_exp.nvidia.ple_layer import Qwen4ExpPLELayer +from vllm.platforms import current_platform from vllm.utils.torch_utils import weak_ref_tensor from vllm.v1.attention.backends.short_conv_attn import ( PleShortConvAttentionMetadata, @@ -1102,7 +1103,7 @@ def test_fused_ngram_ids_correctness( contexts: list[list[int]], first_token_id: int, ) -> None: - from vllm.models.qwen4_exp.nvidia.ops.ple import ple_ngram_ids + from vllm.models.qwen4_exp.common.ops.ple import ple_ngram_ids device = torch.device("cuda") query_start_loc = torch.tensor( @@ -1857,15 +1858,20 @@ def _make_conv_metadata( ), ], ) +@pytest.mark.parametrize( + "layer_cls", [Qwen4ExpPLELayer, Qwen4ExpPLELayerAMD], ids=["nvidia", "amd"] +) def test_fused_conv_correctness( case: _ConvBatchCase, state_layout: str, + layer_cls: type[nn.Module], ) -> None: device = torch.device("cuda") metadata, num_real_tokens = _make_conv_metadata(case, device) if case.state_index_stride > 1: assert metadata.state_indices_tensor.stride(0) == case.state_index_stride - module = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer) + # layer_cls is type[nn.Module], so layer_cls.__new__ is type.__new__. + module = object.__new__(layer_cls) nn.Module.__init__(module) module.conv_state_len = (case.kernel_size - 1) * case.dilation module.short_conv_dilation = case.dilation @@ -1938,7 +1944,7 @@ def test_fused_conv_correctness( def test_fused_gate_correctness(num_tokens: int, strided_kv: bool) -> None: import math - from vllm.models.qwen4_exp.nvidia.ops.ple import ple_gate + from vllm.models.qwen4_exp.common.ops.ple import ple_gate device = torch.device("cuda") hc, h = 4, 2560 @@ -2038,7 +2044,6 @@ def _build_amd_ngram_embedding( embedding_dim=config.ple_embed_dim, ple_dense_layer_id=0, max_total_tokens=4, - max_num_reqs=2, prefix="test.ple_embedding", layer_name="test.ple", quant_config=( @@ -2224,3 +2229,47 @@ def lookup(ids: torch.Tensor, out: torch.Tensor) -> torch.Tensor: expected = loaded_weight[ngram_ids.cpu()].to(device="cuda:0").flatten(-2) torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0) + + +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="only the ROCm backend runs the n-gram embedding under torch.compile", +) +@pytest.mark.parametrize( + ("query_lens", "contexts"), + [([3, 1], [[5, 6], [0, 7]]), ([1, 0, 2], [[1, 2], [3, 4], [0, 0]])], + ids=["two-requests", "empty-request"], +) +def test_amd_ngram_embedding_matches_reference_under_compile( + monkeypatch: pytest.MonkeyPatch, + query_lens: list[int], + contexts: list[list[int]], +) -> None: + """The AMD n-gram ID op must run under Dynamo with a symbolic request count.""" + module, loaded_weight = _build_amd_ngram_embedding( + monkeypatch, device="cuda:0", cpu_offload=False, fp8_checkpoint=False + ) + device = torch.device("cuda:0") + query_start_loc = torch.tensor( + [0, *accumulate(query_lens)], dtype=torch.int32, device=device + ) + input_ids = torch.tensor( + [9, 0, 11, 12][: sum(query_lens)], dtype=torch.int32, device=device + ) + ngram_context = torch.tensor(contexts, dtype=torch.int32, device=device) + + compiled = torch.compile(module, fullgraph=True, dynamic=True) + output = compiled(input_ids, query_start_loc, ngram_context) + + ngram_ids = _reference_ngram_ids( + input_ids, + query_start_loc, + ngram_context, + module.layer_multipliers, + module.ngram_heads_vocab_sizes, + module.ngram_heads_offsets, + module.eos_token_id, + module.heads_per_ngram, + ) + expected = loaded_weight[ngram_ids.cpu()].flatten(-2).to(device) + torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0) diff --git a/vllm/models/qwen4_exp/amd/model.py b/vllm/models/qwen4_exp/amd/model.py index b3a2b7505335..f36f70bee2ac 100644 --- a/vllm/models/qwen4_exp/amd/model.py +++ b/vllm/models/qwen4_exp/amd/model.py @@ -140,9 +140,9 @@ def _remap_qsa_cache_scale_name( "_input_scale", ] -# The checkpoint keeps down and injection projections separate; runtime packs -# them into adjacent logical shards of one MergedColumnParallelLinear. -_HC_WEIGHTS_MAPPER = WeightsMapper( +# The checkpoint stores these projections separately; runtime packs each group +# into adjacent logical shards of a MergedColumnParallelLinear. +_EXTRA_WEIGHTS_MAPPER = WeightsMapper( orig_to_new_stacked={ "hyper_connection.input_mix_weight_down.weight": ( "hyper_connection.input_mix_weight_down_block_inject.weight", @@ -152,6 +152,8 @@ def _remap_qsa_cache_scale_name( "hyper_connection.input_mix_weight_down_block_inject.weight", 1, ), + "ple.key_proj": ("ple.kv_proj", 0), + "ple.value_proj": ("ple.kv_proj", 1), } ) @@ -276,7 +278,7 @@ def forward( if input_ids is None or query_start_loc is None or ngram_context is None: raise RuntimeError("PLE inputs were not prepared") - hidden_states = hidden_states + self.ple( + hidden_states = self.ple( hidden_states, input_ids, query_start_loc, @@ -371,7 +373,7 @@ def update_physical_experts_metadata( } ) class Qwen4ExpModel(nn.Module): - hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _HC_WEIGHTS_MAPPER + hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _EXTRA_WEIGHTS_MAPPER def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() @@ -594,6 +596,7 @@ class Qwen4ExpForCausalLM( packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"], + "kv_proj": ["key_proj", "value_proj"], "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"], "in_proj_ba": ["in_proj_b", "in_proj_a"], "input_mix_weight_down_block_inject": [ @@ -644,8 +647,6 @@ def forward( inputs_embeds: torch.Tensor | None = None, **kwargs: object, ) -> torch.Tensor | IntermediateTensors: - # Forward kwargs unchanged so the runner's _maybe_add_ngram_kwargs - # path (query_start_loc / ngram_context) reaches Qwen4ExpModel. return self.model( input_ids, positions, @@ -825,11 +826,12 @@ class Qwen4ExpForConditionalGeneration( requires_raw_input_tokens = True packed_modules_mapping = Qwen3_5ForConditionalGeneration.packed_modules_mapping | { + "kv_proj": ["key_proj", "value_proj"], "input_mix_weight_down_block_inject": [ "input_mix_weight_down", "block_inject_weight", "_input_mix_padding", - ] + ], } @staticmethod diff --git a/vllm/models/qwen4_exp/amd/mtp.py b/vllm/models/qwen4_exp/amd/mtp.py index d06f6436e77f..b6513e802d73 100644 --- a/vllm/models/qwen4_exp/amd/mtp.py +++ b/vllm/models/qwen4_exp/amd/mtp.py @@ -50,7 +50,7 @@ from .hyperconnection import GatedResidual, HyperConnectionConfig from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm from .model import ( - _HC_WEIGHTS_MAPPER, + _EXTRA_WEIGHTS_MAPPER, _QWEN4_EXP_IGNORED_MISSING_SUFFIXES, Qwen4ExpDecoderLayer, Qwen4ExpMixtureOfExperts, @@ -172,7 +172,7 @@ def _make_draft_vllm_config( } ) class Qwen4ExpMultiTokenPredictor(nn.Module): - hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _HC_WEIGHTS_MAPPER + hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _EXTRA_WEIGHTS_MAPPER def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() diff --git a/vllm/models/qwen4_exp/amd/ple_layer.py b/vllm/models/qwen4_exp/amd/ple_layer.py index bd0607b61338..fbc07b73ddf8 100644 --- a/vllm/models/qwen4_exp/amd/ple_layer.py +++ b/vllm/models/qwen4_exp/amd/ple_layer.py @@ -2,18 +2,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """GPU-resident Qwen4Exp position-learning enhancement layers.""" -import math from collections.abc import Iterable, Sequence import torch -import torch.nn.functional as F from torch import nn from transformers import Qwen4ExpTextConfig from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config from vllm.forward_context import get_forward_context from vllm.logger import init_logger -from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.linear import MergedColumnParallelLinear from vllm.model_executor.layers.mamba.abstract import MambaBase from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateDtypeCalculator, @@ -30,13 +28,13 @@ PleShortConvAttentionBackend, PleShortConvAttentionMetadata, ) -from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from ..common.ngram_embedding import ( Qwen4ExpPLEDeviceEmbedding, Qwen4ExpPLEEmbeddingMethod, Qwen4ExpPLEPinnedHostEmbedding, ) +from ..common.ops.ple import ple_conv, ple_gate, ple_ngram_ids logger = init_logger(__name__) @@ -177,7 +175,6 @@ def __init__( embedding_dim: int, ple_dense_layer_id: int, max_total_tokens: int, - max_num_reqs: int, prefix: str, layer_name: str, *, @@ -270,55 +267,6 @@ def __init__( self.ngram_embedding.weight.device, self.ngram_embedding.weight.is_pinned(), ) - self.register_buffer( - "positions_buffer", - torch.arange(max_total_tokens, dtype=torch.int64), - persistent=False, - ) - self.register_buffer( - "padded_buffer", - torch.full( - (max_num_reqs, max_total_tokens), - self.eos_token_id, - dtype=torch.int64, - ), - persistent=False, - ) - - @staticmethod - def _shift_precompute( - tokens: torch.Tensor, eos_token_id: int - ) -> tuple[torch.Tensor, torch.Tensor]: - if tokens.dim() != 2: - raise ValueError("tokens must be a 2D tensor") - batch_size, seq_len = tokens.shape - positions = torch.arange(seq_len, device=tokens.device, dtype=torch.int64) - eos_positions = torch.where(tokens == eos_token_id, positions, -1) - previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values - previous_eos = torch.cat( - [ - eos_positions.new_full((batch_size, 1), -1), - previous_eos_inclusive[:, :-1], - ], - dim=1, - ) - return positions, positions.unsqueeze(0) - previous_eos - 1 - - @staticmethod - def _shift_apply( - tokens: torch.Tensor, - positions: torch.Tensor, - position_in_segment: torch.Tensor, - shift: int, - eos_token_id: int, - ) -> torch.Tensor: - if shift == 0: - return tokens - source = positions - shift - gather_indices = source.clamp_min(0).unsqueeze(0).expand(tokens.shape[0], -1) - shifted = tokens.gather(1, gather_indices) - valid = (source.unsqueeze(0) >= 0) & (position_in_segment >= shift) - return torch.where(valid, shifted, tokens.new_full((), eos_token_id)) def forward( self, @@ -326,64 +274,17 @@ def forward( query_start_loc: torch.Tensor, ngram_context: torch.Tensor, ) -> torch.Tensor: - input_ids = input_ids.reshape(-1).long() - query_start_loc = query_start_loc.long() - num_reqs = query_start_loc.numel() - 1 - num_tokens = input_ids.shape[0] - if num_tokens > self.positions_buffer.numel(): - raise ValueError( - f"PLE received {num_tokens} tokens, but its workspace supports " - f"at most {self.positions_buffer.numel()}" - ) - if num_reqs > self.padded_buffer.shape[0]: - raise ValueError( - f"PLE received {num_reqs} requests, but its workspace supports " - f"at most {self.padded_buffer.shape[0]}" - ) - - positions = self.positions_buffer[:num_tokens] - packed = self.padded_buffer[:num_reqs] - packed.fill_(self.eos_token_id) - request_indices = torch.searchsorted(query_start_loc, positions, right=True) - 1 - request_indices.clamp_(max=num_reqs - 1) - columns = (positions - query_start_loc[request_indices]).clamp( - 0, packed.shape[1] - 1 - ) - packed[request_indices, columns] = input_ids - ngram_context = ngram_context[:num_reqs].to( - device=input_ids.device, dtype=torch.long + input_ids = input_ids.reshape(-1) + ngram_ids = input_ids.new_empty( + (input_ids.shape[0], self.ngram_heads), dtype=torch.int64 ) - - context = torch.cat([ngram_context, packed], dim=-1) - positions_2d, position_in_segment = self._shift_precompute( - context, self.eos_token_id + torch.ops.vllm.qwen4_exp_amd_ple_ngram_ids( + input_ids, + query_start_loc, + ngram_context, + ngram_ids, + self.layer_name, ) - shifted = [context] - for shift in range(1, self.ngram_size): - shifted.append( - self._shift_apply( - context, - positions_2d, - position_in_segment, - shift, - self.eos_token_id, - ) - ) - adjusted_columns = columns + self.ngram_size - 1 - id_blocks = [] - for ngram in range(2, self.ngram_size + 1): - start = (ngram - 2) * self.heads_per_ngram - end = start + self.heads_per_ngram - mixed = shifted[0] * self.layer_multipliers[0] - for index in range(1, ngram): - mixed = torch.bitwise_xor( - mixed, shifted[index] * self.layer_multipliers[index] - ) - sizes = self.ngram_heads_vocab_sizes[start:end] - offsets = self.ngram_heads_offsets[start:end] - ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets - id_blocks.append(ids[request_indices, adjusted_columns]) - ngram_ids = torch.cat(id_blocks, dim=-1) embedding = self.ngram_embedding if embedding.supports_prefetch: output = ngram_ids.new_empty( @@ -508,26 +409,21 @@ def __init__( int(config.ple_embed_dim), self.ple_dense_layer_id, vllm_config.scheduler_config.max_num_batched_tokens, - vllm_config.scheduler_config.max_num_seqs, f"{prefix}.ple_embedding", prefix, data_parallel_rank=vllm_config.parallel_config.data_parallel_rank, quant_config=quant_config, params_dtype=model_config.dtype, ) - self.key_proj = ReplicatedLinear( + # The PLE cache is TP-replicated, so this merged projection is too. + self.kv_proj = MergedColumnParallelLinear( int(config.ple_embed_dim), - self.hc_hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.key_proj", - ) - self.value_proj = ReplicatedLinear( - int(config.ple_embed_dim), - self.hidden_size, + [self.hc_hidden_size, self.hidden_size], bias=False, + params_dtype=model_config.dtype, quant_config=quant_config, - prefix=f"{prefix}.value_proj", + prefix=f"{prefix}.kv_proj", + disable_tp=True, ) norm_args = ( self.hc_hidden_size, @@ -580,355 +476,15 @@ def get_state_shape(self) -> Sequence[tuple[int, ...]]: num_spec=self.num_spec_tokens, ) - def _apply_norm( - self, norm: Qwen4ExpPLEGroupedNorm, hidden_states: torch.Tensor - ) -> torch.Tensor: - shape = hidden_states.shape - return norm(hidden_states.flatten(-2)).reshape(shape) - - def _short_conv_fallback(self, inputs: torch.Tensor) -> torch.Tensor: - # Profiling / CUDA graph capture only; conv state is not updated. - inputs_t = inputs.transpose(0, 1).unsqueeze(0) - output = self.conv1d(inputs_t)[..., : inputs_t.size(-1)] - return F.silu(output).squeeze(0).transpose(0, 1) - - def _short_conv_dilated_decode_batched( - self, - x_d: torch.Tensor, - conv_state: torch.Tensor, - conv_weights: torch.Tensor, - state_indices_tensor_d: torch.Tensor, - has_initial_states_d: torch.Tensor | None, - ) -> torch.Tensor: - state_indices = state_indices_tensor_d.to( - device=conv_state.device, dtype=torch.int64 - ) - # FULL cudagraph padded decode rows use NULL_BLOCK_ID. Remap them to - # slot 0 for a safe gather, then zero output and skip write-back. - valid_state = state_indices != NULL_BLOCK_ID - state_indices = torch.where( - valid_state, state_indices, torch.zeros_like(state_indices) - ) - if has_initial_states_d is None: - has_initial_state = valid_state - else: - if has_initial_states_d.numel() < state_indices_tensor_d.numel(): - raise ValueError( - "has_initial_states_d size mismatch: " - f"got {has_initial_states_d.numel()}, " - f"need >= {state_indices_tensor_d.numel()}." - ) - has_initial_state = has_initial_states_d[ - : state_indices_tensor_d.numel() - ].to(device=conv_state.device, dtype=torch.bool) - has_initial_state = has_initial_state & valid_state - - cached_state = conv_state.index_select(0, state_indices) - state = cached_state[..., : self.conv_state_len].to(x_d.dtype) - if self.conv_state_len > 0: - initial_state = torch.where( - has_initial_state.view(-1, 1, 1), - state, - torch.zeros_like(state), - ) - history = torch.cat((initial_state, x_d.unsqueeze(-1)), dim=-1) - else: - history = x_d.unsqueeze(-1) - - conv_output = F.conv1d( - history, - conv_weights.unsqueeze(1).contiguous(), - groups=history.size(1), - dilation=self.short_conv_dilation, - ).squeeze(-1) - output = F.silu(conv_output) - output = output * valid_state.view(-1, 1).to(output.dtype) - - if self.conv_state_len > 0: - next_state = history[..., -self.conv_state_len :] - # Padded rows are remapped to the reserved null slot. Preserve its - # existing value while writing the new states for valid rows. - existing_base_state = cached_state[..., : self.conv_state_len] - safe_next_state = torch.where( - valid_state.view(-1, 1, 1), - next_state.to(conv_state.dtype), - existing_base_state, - ) - cached_state[..., : self.conv_state_len] = safe_next_state - conv_state.index_copy_(0, state_indices, cached_state) - - return output - - def _short_conv_dilated_prefill_batched( - self, - x_p: torch.Tensor, - metadata: PleShortConvAttentionMetadata, - conv_state: torch.Tensor, - conv_weights: torch.Tensor, - state_indices_tensor_p: torch.Tensor, - num_prefills: int, - num_decode_tokens: int, - num_prefill_tokens: int, - ) -> torch.Tensor: - query_start_loc_p = metadata.query_start_loc_p - if query_start_loc_p is None: - raise ValueError("query_start_loc is required for prefill short-conv") - # The metadata builder guarantees that the prefill query offsets start - # at 0 and end at num_prefill_tokens. Avoid reading those values here, - # since doing so would force a device-to-host synchronization. - has_initial_states_p = metadata.has_initial_states_p - if has_initial_states_p is None: - raise ValueError("has_initial_states_p is required for prefill short-conv") - - output = torch.empty_like(x_p) - q_starts = query_start_loc_p.to(torch.int64) - if state_indices_tensor_p.numel() < num_prefills: - raise ValueError( - "state_indices_tensor_p size mismatch: " - f"got {state_indices_tensor_p.numel()}, " - f"need >= {num_prefills}." - ) - if has_initial_states_p.numel() < num_prefills: - raise ValueError( - "has_initial_states_p size mismatch: " - f"got {has_initial_states_p.numel()}, " - f"need >= {num_prefills}." - ) - if num_prefills == 0 or x_p.numel() == 0: - return output - lengths = q_starts[1:] - q_starts[:-1] - # Use the CPU-computed packing width from the metadata builder instead - # of synchronizing on lengths.max(). - max_len = metadata.max_prefill_query_len - if max_len <= 0: - return output - - hidden_size = x_p.shape[1] - positions = torch.arange( - num_prefill_tokens, device=x_p.device, dtype=torch.int64 - ) - req_indices = torch.searchsorted(q_starts[1:], positions, right=True) - col_indices = positions - q_starts[req_indices] - - packed_tokens = x_p.new_zeros((num_prefills, max_len, hidden_size)) - packed_tokens[req_indices, col_indices] = x_p - packed_tokens = packed_tokens.transpose(1, 2).contiguous() - - state_indices = state_indices_tensor_p[:num_prefills].to( - device=conv_state.device, dtype=torch.int64 - ) - valid_state = state_indices != NULL_BLOCK_ID - state_indices = torch.where( - valid_state, state_indices, torch.zeros_like(state_indices) - ) - has_initial = has_initial_states_p[:num_prefills].to( - device=conv_state.device, dtype=torch.bool - ) - if self.conv_state_len > 0: - if conv_state.shape[0] == 0: - state = conv_state.new_zeros( - (num_prefills, hidden_size, self.conv_state_len), - dtype=x_p.dtype, - ) - else: - state = conv_state.index_select(0, state_indices)[ - ..., : self.conv_state_len - ].to(x_p.dtype) - use_initial_mask = (valid_state & has_initial).view(num_prefills, 1, 1) - initial_state = torch.where( - use_initial_mask, - state, - torch.zeros_like(state), - ) - history = torch.cat((initial_state, packed_tokens), dim=-1) - else: - history = packed_tokens - - conv_output = F.conv1d( - history, - conv_weights.unsqueeze(1).contiguous(), - groups=history.size(1), - dilation=self.short_conv_dilation, - ) - conv_output = F.silu(conv_output).transpose(1, 2).contiguous() - - token_positions = torch.arange(max_len, device=x_p.device, dtype=torch.int64) - valid_tokens = token_positions.view(1, max_len) < lengths.view(num_prefills, 1) - valid_output_mask = valid_tokens & valid_state.to(device=x_p.device).view( - num_prefills, 1 - ) - conv_output.masked_fill_(~valid_output_mask.unsqueeze(-1), 0) - output.copy_(conv_output[req_indices, col_indices]) - - if self.conv_state_len > 0 and conv_state.shape[0] > 0: - state_starts = lengths.to(device=history.device, dtype=torch.int64).view( - num_prefills, 1, 1 - ) - state_offsets = torch.arange( - self.conv_state_len, device=history.device, dtype=torch.int64 - ).view(1, 1, self.conv_state_len) - next_state = history.gather( - dim=2, - index=(state_starts + state_offsets).expand(-1, history.size(1), -1), - ) - # Write back without a host synchronization. Valid, non-empty rows - # receive their new state; padding and zero-length rows keep the - # current cache value. - existing_state = conv_state.index_select(0, state_indices) - existing_base_state = existing_state[..., : self.conv_state_len] - update_mask = valid_state & (lengths.to(device=conv_state.device) > 0) - safe_next_state = torch.where( - update_mask.view(num_prefills, 1, 1), - next_state.to(conv_state.dtype), - existing_base_state, - ) - existing_state[..., : self.conv_state_len] = safe_next_state - conv_state.index_copy_(0, state_indices, existing_state) - return output - - def _short_conv_dilated_spec_batched( - self, - x_spec: torch.Tensor, - conv_state: torch.Tensor, - conv_weights: torch.Tensor, - spec_state_indices_tensor: torch.Tensor, - spec_query_start_loc: torch.Tensor, - num_accepted_tokens: torch.Tensor, - spec_query_len: int, - ) -> torch.Tensor: - """Dilated short-conv for speculative-decode (MTP) requests. - - Each spec request feeds multiple (draft + 1) query tokens. The conv - outputs are computed causally after rolling back the previous draft - state by ``num_accepted_tokens - 1``. The current candidate inputs stay - in the extended cache for the next forward, matching - ``causal_conv1d_update``. - - ``spec_query_len`` (== num_speculative_tokens + 1) is the maximum query - length and is a Python int, so no host synchronization is needed; this - keeps the path safe for full CUDA-graph capture/replay where the buffers - are padded at the request level. - """ - num_reqs = spec_state_indices_tensor.numel() - hidden_size = x_spec.size(-1) - # Use a fixed packing width instead of synchronizing on lengths.max(). - max_len = spec_query_len - # Full CUDA graphs can pad these buffers. Only the first num_reqs - # accepted-token counts belong to actual speculative requests. - num_accepted_tokens = num_accepted_tokens[:num_reqs] - q_starts = spec_query_start_loc[: num_reqs + 1].to(torch.int64) - # Keep the number of real speculative tokens on the device. - total_real_tokens = q_starts[num_reqs] - - state_indices = spec_state_indices_tensor.to( - device=conv_state.device, dtype=torch.int64 - ) - valid_state = state_indices != NULL_BLOCK_ID - state_indices = torch.where( - valid_state, state_indices, torch.zeros_like(state_indices) - ) - positions = torch.arange( - x_spec.size(0), device=x_spec.device, dtype=torch.int64 - ) - # Route graph-padded token rows to the discarded dummy request so that - # they cannot overwrite real packed data. - req_indices = torch.searchsorted(q_starts[1:], positions, right=True) - valid_tokens = (positions < total_real_tokens) & (req_indices < num_reqs) - clamped_req_indices = req_indices.clamp_max(max(num_reqs - 1, 0)) - col_indices = (positions - q_starts[clamped_req_indices]).clamp_(0, max_len - 1) - pack_req_indices = torch.where( - valid_tokens, - clamped_req_indices, - torch.full_like(req_indices, num_reqs), - ) - pack_col_indices = torch.where( - valid_tokens, col_indices, torch.zeros_like(col_indices) - ) - - # The last request row is the dummy sink for graph padding. - packed = x_spec.new_zeros((num_reqs + 1, max_len, hidden_size)) - packed[pack_req_indices, pack_col_indices] = x_spec - packed = packed.transpose(1, 2).contiguous() - - if self.conv_state_len > 0: - cached_state = conv_state.index_select(0, state_indices) - rollback_offsets = num_accepted_tokens.to( - device=conv_state.device, dtype=torch.int64 - ).sub(1) - rollback_offsets = torch.where( - valid_state, - rollback_offsets.clamp_(0, max_len - 1), - torch.zeros_like(rollback_offsets), - ) - state_offsets = torch.arange( - self.conv_state_len, device=conv_state.device, dtype=torch.int64 - ).view(1, 1, self.conv_state_len) - rollback_indices = rollback_offsets.view(-1, 1, 1) + state_offsets - state = cached_state.gather( - 2, rollback_indices.expand(-1, hidden_size, -1) - ).to(x_spec.dtype) - state = torch.where( - valid_state.view(num_reqs, 1, 1), - state, - torch.zeros_like(state), - ) - # Append a zeroed dummy-row state to match the [num_reqs + 1] pack. - dummy_state = state.new_zeros((1, hidden_size, self.conv_state_len)) - state_full = torch.cat((state, dummy_state), dim=0) - history = torch.cat((state_full, packed), dim=-1) - else: - history = packed - - conv_output = F.conv1d( - history, - conv_weights.unsqueeze(1).contiguous(), - groups=history.size(1), - dilation=self.short_conv_dilation, - ) - conv_output = F.silu(conv_output).transpose(1, 2).contiguous() - - output = conv_output[pack_req_indices, pack_col_indices] - output = output * valid_tokens.view(-1, 1).to(output.dtype) - - # Keep all current candidate inputs in the extended state. On the next - # target forward, ``num_accepted_tokens - 1`` selects the rollback - # window before processing the newly scheduled tokens. - if self.conv_state_len > 0: - state_capacity = self.conv_state_len + max_len - 1 - if conv_state.size(-1) < state_capacity: - raise RuntimeError( - "PLE short-conv cache cannot retain speculative tokens: " - f"got {conv_state.size(-1)}, need {state_capacity}." - ) - candidate_state = history[:num_reqs, :, 1 : state_capacity + 1] - query_lengths = q_starts[1:] - q_starts[:-1] - state_positions = torch.arange( - state_capacity, device=history.device, dtype=torch.int64 - ).view(1, 1, state_capacity) - update_lengths = (self.conv_state_len + query_lengths - 1).view( - num_reqs, 1, 1 - ) - update_mask = valid_state.view(num_reqs, 1, 1) & ( - state_positions < update_lengths - ) - existing_state = cached_state[..., :state_capacity] - next_state = torch.where( - update_mask, - candidate_state.to(conv_state.dtype), - existing_state, - ) - cached_state[..., :state_capacity] = next_state - conv_state.index_copy_(0, state_indices, cached_state) - - return output - def _short_conv_dilated_dispatch( self, inputs: torch.Tensor, + residual: torch.Tensor, + outer_residual: torch.Tensor, metadata: PleShortConvAttentionMetadata, conv_state: torch.Tensor, conv_weights: torch.Tensor, - ) -> torch.Tensor: + ) -> None: num_prefills = metadata.num_prefills num_decodes = metadata.num_decodes num_decode_tokens = metadata.num_decode_tokens @@ -936,110 +492,145 @@ def _short_conv_dilated_dispatch( has_prefill = num_prefills > 0 has_decode = num_decodes > 0 has_spec = metadata.spec_sequence_masks is not None - x = inputs[: metadata.num_actual_tokens] - - # Split spec / non-spec tokens. - if has_spec: - if has_prefill or has_decode: - assert metadata.spec_token_indx is not None - assert metadata.non_spec_token_indx is not None - x_spec = x.index_select(0, metadata.spec_token_indx.long()) - x_non_spec = x.index_select(0, metadata.non_spec_token_indx.long()) - else: - x_spec = x - x_non_spec = None - else: - x_spec = None - x_non_spec = x + has_non_spec = has_prefill or has_decode + inputs = inputs[: metadata.num_actual_tokens] + residual = residual[: metadata.num_actual_tokens] + outer_residual = outer_residual[: metadata.num_actual_tokens] + + spec_token_indices = None + non_spec_token_indices = None + if has_spec and has_non_spec: + assert metadata.spec_token_indx is not None + assert metadata.non_spec_token_indx is not None + spec_token_indices = metadata.spec_token_indx + non_spec_token_indices = metadata.non_spec_token_indx - spec_output = None - # 1. Run the multi-query speculative-decode part. if has_spec: assert metadata.spec_state_indices_tensor is not None - assert metadata.spec_query_start_loc is not None - assert metadata.num_accepted_tokens is not None - spec_output = self._short_conv_dilated_spec_batched( - x_spec=x_spec, + query_start_loc = metadata.spec_query_start_loc + num_accepted_tokens = metadata.num_accepted_tokens + assert query_start_loc is not None + assert num_accepted_tokens is not None + spec_state_indices = metadata.spec_state_indices_tensor[ + : metadata.num_spec_decodes + ] + # Mixed batches stay in their original row order; the kernels map + # logical spec/non-spec rows instead of materializing both groups. + ple_conv( + inputs=inputs, + residual=residual, conv_state=conv_state, conv_weights=conv_weights, - spec_state_indices_tensor=metadata.spec_state_indices_tensor[ - : metadata.num_spec_decodes - ], - spec_query_start_loc=metadata.spec_query_start_loc, - num_accepted_tokens=metadata.num_accepted_tokens, + state_indices=spec_state_indices, + outer_residual=outer_residual, + mode="spec", + dilation=self.short_conv_dilation, + query_start_loc=query_start_loc, + num_accepted_tokens=num_accepted_tokens, spec_query_len=metadata.spec_query_len, + token_indices=spec_token_indices, ) - # 2. Run regular decode and prefill requests. - conv_out_non_spec = None - state_indices_tensor = metadata.state_indices_tensor - if x_non_spec is not None: - assert state_indices_tensor is not None - if has_prefill: - state_indices_tensor_d, state_indices_tensor_p = torch.split( - state_indices_tensor, - [num_decodes, num_prefills], - dim=0, + if not has_non_spec: + return + + state_indices = metadata.state_indices_tensor + assert state_indices is not None + if has_prefill: + state_indices_d, state_indices_p = torch.split( + state_indices, [num_decodes, num_prefills], dim=0 + ) + if non_spec_token_indices is None: + inputs_d, inputs_p = torch.split( + inputs, [num_decode_tokens, num_prefill_tokens], dim=0 ) - x_d, x_p = torch.split( - x_non_spec, + residual_d, residual_p = torch.split( + residual, [num_decode_tokens, num_prefill_tokens], dim=0 + ) + outer_residual_d, outer_residual_p = torch.split( + outer_residual, [num_decode_tokens, num_prefill_tokens], dim=0, ) - non_spec_parts: list[torch.Tensor] = [] - if has_decode: - non_spec_parts.append( - self._short_conv_dilated_decode_batched( - x_d=x_d, - conv_state=conv_state, - conv_weights=conv_weights, - state_indices_tensor_d=state_indices_tensor_d, - has_initial_states_d=metadata.has_initial_states_d, - ) - ) - non_spec_parts.append( - self._short_conv_dilated_prefill_batched( - x_p=x_p, - metadata=metadata, - conv_state=conv_state, - conv_weights=conv_weights, - state_indices_tensor_p=state_indices_tensor_p, - num_prefills=num_prefills, - num_decode_tokens=num_decode_tokens, - num_prefill_tokens=num_prefill_tokens, - ) - ) - conv_out_non_spec = torch.vstack(non_spec_parts) + token_indices_d = None + token_indices_p = None else: - conv_out_non_spec = self._short_conv_dilated_decode_batched( - x_d=x_non_spec, + inputs_d = inputs_p = inputs + residual_d = residual_p = residual + outer_residual_d = outer_residual_p = outer_residual + token_indices_d, token_indices_p = torch.split( + non_spec_token_indices, + [num_decode_tokens, num_prefill_tokens], + dim=0, + ) + + if has_decode: + ple_conv( + inputs=inputs_d, + residual=residual_d, conv_state=conv_state, conv_weights=conv_weights, - state_indices_tensor_d=state_indices_tensor[: x_non_spec.size(0)], - has_initial_states_d=metadata.has_initial_states_d, + state_indices=state_indices_d, + outer_residual=outer_residual_d, + mode="decode", + dilation=self.short_conv_dilation, + has_initial_states=metadata.has_initial_states_d, + token_indices=token_indices_d, ) - # 3. Merge both parts back into the original token order. - if has_spec and conv_out_non_spec is not None: - assert metadata.spec_token_indx is not None - assert metadata.non_spec_token_indx is not None - assert spec_output is not None - output = x.new_empty((metadata.num_actual_tokens, x.size(-1))) - output.index_copy_(0, metadata.spec_token_indx, spec_output) - output.index_copy_(0, metadata.non_spec_token_indx, conv_out_non_spec) - return output - elif has_spec: - assert spec_output is not None - return spec_output - if conv_out_non_spec is None: - return x - return conv_out_non_spec + query_start_loc = metadata.query_start_loc_p + if query_start_loc is None: + raise ValueError("query_start_loc is required for prefill short-conv") + has_initial_states = metadata.has_initial_states_p + if has_initial_states is None: + raise ValueError( + "has_initial_states_p is required for prefill short-conv" + ) + ple_conv( + inputs=inputs_p, + residual=residual_p, + conv_state=conv_state, + conv_weights=conv_weights, + state_indices=state_indices_p, + outer_residual=outer_residual_p, + mode="prefill", + dilation=self.short_conv_dilation, + query_start_loc=query_start_loc, + has_initial_states=has_initial_states, + token_indices=token_indices_p, + ) + else: + num_decode_rows = ( + non_spec_token_indices.numel() + if non_spec_token_indices is not None + else inputs.size(0) + ) + ple_conv( + inputs=inputs, + residual=residual, + conv_state=conv_state, + conv_weights=conv_weights, + state_indices=state_indices[:num_decode_rows], + outer_residual=outer_residual, + mode="decode", + dilation=self.short_conv_dilation, + has_initial_states=metadata.has_initial_states_d, + token_indices=non_spec_token_indices, + ) - def _short_conv(self, inputs: torch.Tensor) -> torch.Tensor: + def _short_conv( + self, + inputs: torch.Tensor, + residual: torch.Tensor, + outer_residual: torch.Tensor, + ) -> None: forward_context = get_forward_context() attn_metadata = forward_context.attn_metadata + # Profiling omits all metadata or this Mamba entry. Short convolution + # is a no-op there, but preserve the outer residual addition. if attn_metadata is None: - return self._short_conv_fallback(inputs) + residual.add_(outer_residual) + return if not isinstance(attn_metadata, dict): raise RuntimeError( @@ -1049,8 +640,8 @@ def _short_conv(self, inputs: torch.Tensor) -> torch.Tensor: layer_attn_metadata = attn_metadata.get(self.prefix) if layer_attn_metadata is None: - # MRV2 omits Mamba-family metadata during profile warmup. - return self._short_conv_fallback(inputs) + residual.add_(outer_residual) + return if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata): raise TypeError( "Expected PleShortConvAttentionMetadata for layer " @@ -1059,6 +650,7 @@ def _short_conv(self, inputs: torch.Tensor) -> torch.Tensor: ) conv_state = self.kv_cache[0] + # Canonicalize both backend cache layouts to [slot, channel, window]. if not is_conv_state_dim_first(): conv_state = conv_state.transpose(-1, -2) conv_weights = self.conv1d.weight.squeeze(1) @@ -1072,11 +664,13 @@ def _short_conv(self, inputs: torch.Tensor) -> torch.Tensor: f"expect at least {state_capacity}." ) conv_state = conv_state[..., -state_capacity:] - return self._short_conv_dilated_dispatch( - inputs, - layer_attn_metadata, - conv_state, - conv_weights.to(dtype=inputs.dtype), + self._short_conv_dilated_dispatch( + inputs=inputs, + residual=residual, + outer_residual=outer_residual, + metadata=layer_attn_metadata, + conv_state=conv_state, + conv_weights=conv_weights.to(dtype=inputs.dtype), ) def forward( @@ -1095,27 +689,58 @@ def forward( ) embeddings = self.ple_embedding(input_ids, query_start_loc, ngram_context) embeddings = self.ple_embedding.ngram_embedding.dequantize( - embeddings, - self.ple_embedding.ngram_embedding.params_dtype, - ) - key, _ = self.key_proj(embeddings) - value, _ = self.value_proj(embeddings) - token_count = hidden_states.shape[0] - key = key.reshape(token_count, self.hc_count, self.hidden_size) - query = hidden_states.reshape(token_count, self.hc_count, self.hidden_size) - key = self._apply_norm(self.norm_key, key) - query = self._apply_norm(self.norm_query, query) - gate = (key * query).sum(dim=-1, keepdim=True) / math.sqrt(self.hidden_size) - gate = torch.sigmoid(gate.sign() * gate.abs().clamp_min(1e-6).sqrt()) - gated_value = gate * value.unsqueeze(-2) - normalized = self._apply_norm(self.norm_conv, gated_value).flatten(-2) - conv_output = torch.zeros_like(normalized) + embeddings, hidden_states.dtype + ) + kv, _ = self.kv_proj(embeddings) + key, value = kv.split(self.kv_proj.output_sizes, dim=-1) + gated_output, conv_input = ple_gate( + key, + value, + hidden_states, + self.norm_key.weight, + self.norm_query.weight, + self.norm_conv.weight, + self.norm_key.eps, + ) + # The short-conv op is a piecewise-graph splitting op: state routing + # reads runtime request metadata. It accumulates the convolution and + # the outer residual into gated_output. torch.ops.vllm.qwen4_exp_ple_short_conv( - normalized, - conv_output, + conv_input, + gated_output, + hidden_states, self.prefix, ) - return gated_value.flatten(-2) + conv_output + return gated_output + + +def qwen4_exp_amd_ple_ngram_ids( + input_ids: torch.Tensor, + query_start_loc: torch.Tensor, + ngram_context: torch.Tensor, + output: torch.Tensor, + layer_name: str, +) -> None: + """Hash the current request layout into n-gram embedding indices. + + The launch sizes its request binary search from the request count, which + is symbolic under Dynamo, so it runs outside Inductor's FX graph. + """ + layer = get_forward_context().no_compile_layers[layer_name] + if not isinstance(layer, Qwen4ExpPLELayer): + raise TypeError(f"{layer_name} is not a Qwen4Exp PLE owner") + embedding = layer.ple_embedding + ple_ngram_ids( + input_ids=input_ids, + query_start_loc=query_start_loc, + ngram_context=ngram_context, + layer_multipliers=embedding.layer_multipliers, + ngram_heads_vocab_sizes=embedding.ngram_heads_vocab_sizes, + ngram_heads_offsets=embedding.ngram_heads_offsets, + eos_token_id=embedding.eos_token_id, + heads_per_ngram=embedding.heads_per_ngram, + output=output, + ) def qwen4_exp_amd_ple_ngram_embedding( @@ -1155,12 +780,19 @@ def qwen4_exp_amd_ple_ngram_embedding_pinned( def qwen4_exp_ple_short_conv( inputs: torch.Tensor, - output: torch.Tensor, + residual: torch.Tensor, + outer_residual: torch.Tensor, layer_name: str, ) -> None: layer = get_forward_context().no_compile_layers[layer_name] - result = layer._short_conv(inputs) - output[: result.shape[0]].copy_(result) + layer._short_conv(inputs, residual, outer_residual) + + +direct_register_custom_op( + op_name="qwen4_exp_amd_ple_ngram_ids", + op_func=qwen4_exp_amd_ple_ngram_ids, + mutates_args=["output"], +) direct_register_custom_op( @@ -1180,7 +812,7 @@ def qwen4_exp_ple_short_conv( direct_register_custom_op( op_name="qwen4_exp_ple_short_conv", op_func=qwen4_exp_ple_short_conv, - mutates_args=["output"], + mutates_args=["residual"], ) diff --git a/vllm/models/qwen4_exp/common/ops/__init__.py b/vllm/models/qwen4_exp/common/ops/__init__.py new file mode 100644 index 000000000000..a87d2ca4bed3 --- /dev/null +++ b/vllm/models/qwen4_exp/common/ops/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Qwen4Exp kernels shared by the NVIDIA and AMD backends; import leaf modules +directly.""" diff --git a/vllm/models/qwen4_exp/nvidia/ops/ple.py b/vllm/models/qwen4_exp/common/ops/ple.py similarity index 100% rename from vllm/models/qwen4_exp/nvidia/ops/ple.py rename to vllm/models/qwen4_exp/common/ops/ple.py diff --git a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py index 3be6ea6013dd..cf03cbb40d47 100644 --- a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py +++ b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py @@ -23,7 +23,7 @@ Qwen4ExpPLEPinnedHostEmbedding, Qwen4ExpPLEUnquantizedEmbeddingMethod, ) -from .ops.ple import ple_ngram_ids +from ..common.ops.ple import ple_ngram_ids logger = init_logger(__name__) diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index ee8a2b47dd76..bdb16b3adef0 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -24,8 +24,8 @@ PleShortConvAttentionMetadata, ) +from ..common.ops.ple import ple_conv, ple_gate from .ngram_embedding import Qwen4ExpNGramEmbedding -from .ops.ple import ple_conv, ple_gate class Qwen4ExpPLEGroupedNorm(nn.Module): From 9f3d70e08157a95911f15809f2f1446424825ea7 Mon Sep 17 00:00:00 2001 From: fxmarty-amd Date: Thu, 8 Oct 2026 18:04:46 +0200 Subject: [PATCH 110/624] [Quantization] Enable shared expert fusion compatibility with online `shared_expert` quantization (showcase: along Quark MXFP4 routed experts) (#55686) Signed-off-by: Felix Marty Signed-off-by: fxmarty-amd Co-authored-by: OpenAI Codex Co-authored-by: Misha Goin --- .../layers/test_fused_shared_expert.py | 89 ++- .../model_loader/test_reload.py | 1 + tests/models/kimi_k3/test_gfx942_int4.py | 10 +- tests/quantization/test_auto_gptq.py | 1 + tests/quantization/test_config_utils.py | 5 +- tests/quantization/test_online.py | 9 +- .../test_online_shared_expert_fusion.py | 637 ++++++++++++++++++ .../test_quantization_config_args.py | 6 +- tests/quantization/utils.py | 29 +- .../model_executor/layers/fused_moe/config.py | 1 + vllm/model_executor/layers/fused_moe/layer.py | 3 + .../layers/fused_moe/routed_experts.py | 56 ++ vllm/model_executor/layers/fused_moe/utils.py | 34 + .../layers/quantization/base_config.py | 2 +- .../layers/quantization/online/base.py | 44 +- .../layers/quantization/online/fp8.py | 23 +- .../layers/quantization/online/mxfp8.py | 3 + .../layers/quantization/utils/config_utils.py | 98 ++- vllm/model_executor/models/AXK1.py | 1 + vllm/model_executor/models/deepseek_v2.py | 1 + vllm/model_executor/models/glm4_moe.py | 1 + vllm/model_executor/models/qwen3_next.py | 1 + vllm/models/glm5next/common/model.py | 1 + vllm/models/minimax_m3/amd/model.py | 1 + 24 files changed, 984 insertions(+), 73 deletions(-) create mode 100644 tests/quantization/test_online_shared_expert_fusion.py diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py index d8d5c4fb032a..78aa563925ff 100644 --- a/tests/model_executor/layers/test_fused_shared_expert.py +++ b/tests/model_executor/layers/test_fused_shared_expert.py @@ -15,10 +15,12 @@ import vllm.config as vllm_config_module from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config +from vllm.config.quantization import QuantizationConfigArgs from vllm.model_executor.layers.fused_moe import utils as fused_moe_utils from vllm.model_executor.layers.fused_moe.layer import determine_expert_counts -from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.quantization.experts_int8 import ExpertsInt8Config from vllm.model_executor.layers.quantization.fp8 import Fp8Config +from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig from vllm.model_executor.layers.quantization.utils.config_utils import ( get_quark_ocp_mx_group_size, @@ -141,10 +143,6 @@ def get_deepseek_v4_quark_config(exclude: list[str]) -> dict[str, Any]: return quantization_config -def _stub_quant_config() -> QuantizationConfig: - return cast(QuantizationConfig, object()) - - class _StubGlm5NextAttention(nn.Module): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__() @@ -340,7 +338,7 @@ def test_resolve_layer_fused_shared_expert_skips_compatibility_when_disabled( ) assert not fused_moe_utils.resolve_layer_fused_shared_expert( - _stub_quant_config(), "model.layers.0.mlp" + ExpertsInt8Config(), "model.layers.0.mlp" ) @@ -355,7 +353,7 @@ def test_resolve_layer_fused_shared_expert_normalizes_unavailable_aiter( assert ( fused_moe_utils.resolve_layer_fused_shared_expert( - _stub_quant_config(), "model.layers.0.mlp" + ExpertsInt8Config(), "model.layers.0.mlp" ) is False ) @@ -364,7 +362,7 @@ def test_resolve_layer_fused_shared_expert_normalizes_unavailable_aiter( def test_resolve_layer_fused_shared_expert_passes_module_prefixes( monkeypatch: pytest.MonkeyPatch, ) -> None: - quant_config = _stub_quant_config() + quant_config = ExpertsInt8Config() monkeypatch.setattr( fused_moe_utils.rocm_aiter_ops, "is_fusion_moe_shared_experts_enabled", @@ -413,7 +411,7 @@ def test_resolve_layer_fused_shared_expert_rejects_incompatible_quantization( ) assert not fused_moe_utils.resolve_layer_fused_shared_expert( - _stub_quant_config(), "model.layers.0.mlp" + ExpertsInt8Config(), "model.layers.0.mlp" ) assert "shared experts are excluded" in caplog.text @@ -423,6 +421,7 @@ def test_deepseek_v4_shared_expert_fse_uses_mtp_quantization_config_prefix( ) -> None: class DeepseekV4Config: expert_dtype = "fp4" + online_quantization_config = None def _is_quark_mxfp4_ocp(self, hf_config: object) -> bool: return True @@ -1070,6 +1069,73 @@ def test_quark_shared_expert_fse_compatibility( ) +@pytest.mark.parametrize("exclude_shared_expert", [False, True]) +def test_online_targets_off_shared_expert_defer_to_checkpoint( + exclude_shared_expert: bool, +) -> None: + """Unrelated online targets preserve the checkpoint's FSE decision.""" + shared_prefix = "model.layers.0.mlp.shared_expert" + quant_config = QuarkConfig( + { + **_QUARK_FSE_CONFIG, + "exclude": ( + [f"{shared_prefix}.down_proj"] if exclude_shared_expert else [] + ), + } + ) + quant_config.online_quantization_config = OnlineQuantizationConfig( + QuantizationConfigArgs(targets={"*linear_attn.out_proj": "mxfp4"}) + ) + + compatible, reason = is_shared_expert_quant_fse_compatible( + quant_config, + "model.layers.0.mlp.experts", + shared_prefix, + ) + + assert compatible is not exclude_shared_expert + assert reason == ( + f"Quark excludes shared experts at {shared_prefix}" + if exclude_shared_expert + else None + ) + + +def test_online_targeting_only_one_shared_projection_rejects_fse() -> None: + shared_prefix = "model.layers.0.mlp.shared_expert" + quant_config = QuarkConfig({**_QUARK_FSE_CONFIG, "exclude": []}) + quant_config.online_quantization_config = OnlineQuantizationConfig( + QuantizationConfigArgs(targets={f"{shared_prefix}.down_proj": "mxfp4"}) + ) + + assert is_shared_expert_quant_fse_compatible( + quant_config, + "model.layers.0.mlp.experts", + shared_prefix, + ) == ( + False, + "online quantization targets only part of the shared expert at " + f"{shared_prefix}", + ) + + +def test_online_fp8_shared_expert_rejects_quark_mxfp4_fse() -> None: + shared_prefix = "model.layers.0.mlp.shared_expert" + quant_config = QuarkConfig({**_QUARK_FSE_CONFIG, "exclude": []}) + quant_config.online_quantization_config = OnlineQuantizationConfig( + QuantizationConfigArgs(targets={f"{shared_prefix}*": "fp8_per_tensor"}) + ) + + assert is_shared_expert_quant_fse_compatible( + quant_config, + "model.layers.0.mlp.experts", + shared_prefix, + ) == ( + False, + "online shared-expert quantization keys do not match the routed expert keys", + ) + + def test_quark_shared_expert_fse_exclude_is_scoped_to_the_layer() -> None: """Excluding one layer's shared experts must not disable FSE elsewhere. @@ -1335,14 +1401,15 @@ def test_quark_packed_layer_config_must_match_global_config() -> None: def test_non_quark_shared_expert_fse_is_incompatible() -> None: compatible, reason = is_shared_expert_quant_fse_compatible( - _stub_quant_config(), + ExpertsInt8Config(), "model.layers.0.mlp.experts", "model.layers.0.mlp.shared_experts", ) assert not compatible assert reason == ( - "shared-expert FSE quantization compatibility is not implemented for object" + "shared-expert FSE quantization compatibility is not implemented for " + "ExpertsInt8Config" ) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index aaad16778f12..66c5b490508a 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -528,6 +528,7 @@ def test_padded_moe_reload_releases_each_layer( layer.quant_method = method layer.expert_map_manager = SimpleNamespace(map_global_to_local=lambda i: i) layer._loaded_expert_biases = set() + layer._fused_shared_expert_quantizer = None method.create_weights( layer, experts, diff --git a/tests/models/kimi_k3/test_gfx942_int4.py b/tests/models/kimi_k3/test_gfx942_int4.py index 20bef41eb89d..e38a8ecff944 100644 --- a/tests/models/kimi_k3/test_gfx942_int4.py +++ b/tests/models/kimi_k3/test_gfx942_int4.py @@ -4,7 +4,7 @@ import sys from collections.abc import Callable from types import ModuleType, SimpleNamespace -from unittest.mock import Mock, patch +from unittest.mock import patch import pytest @@ -91,9 +91,7 @@ def test_int4_per_group_32_is_not_an_online_moe_method(): QuantizationConfigArgs(moe={"weight": "int4_per_group_32"}) ) assert ( - config.resolve_quant_method_cls( - Mock(spec=RoutedExperts), "model.layers.0.mlp.experts" - ) + config.resolve_quant_method_cls(RoutedExperts, "model.layers.0.mlp.experts") is None ) @@ -103,6 +101,4 @@ def test_unknown_online_moe_weight_still_raises(): QuantizationConfigArgs(moe={"weight": "fp8_per_token"}) ) with pytest.raises(ValueError, match="online quantization"): - config.resolve_quant_method_cls( - Mock(spec=RoutedExperts), "model.layers.0.mlp.experts" - ) + config.resolve_quant_method_cls(RoutedExperts, "model.layers.0.mlp.experts") diff --git a/tests/quantization/test_auto_gptq.py b/tests/quantization/test_auto_gptq.py index 7e54599c8098..4f60a974caed 100644 --- a/tests/quantization/test_auto_gptq.py +++ b/tests/quantization/test_auto_gptq.py @@ -162,6 +162,7 @@ def test_routed_experts_loads_per_expert_biases(): class Loader: quant_config = None quant_method = object() + _fused_shared_expert_quantizer = None moe_config = SimpleNamespace( is_act_and_mul=True, tp_rank=0, diff --git a/tests/quantization/test_config_utils.py b/tests/quantization/test_config_utils.py index 6a0c27748f90..a0cb12a9f8ac 100644 --- a/tests/quantization/test_config_utils.py +++ b/tests/quantization/test_config_utils.py @@ -2,11 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests quantization configuration matching utilities.""" -from unittest.mock import Mock - import pytest from vllm.config.quantization import QuantizationConfigArgs +from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization.compressed_tensors.utils import ( should_ignore_layer, ) @@ -58,7 +57,7 @@ def test_online_ignore_supports_fnmatch_patterns(): ) with pytest.raises(ValueError, match="matches both quantization_config.ignore"): - config.resolve_quant_method_cls(Mock(), layer_name) + config.resolve_quant_method_cls(LinearBase, layer_name) @pytest.mark.parametrize( diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index c6b9894009cf..ce13aa9a65d1 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -195,7 +195,6 @@ def test_qianfan_online_fp8_keeps_vision_layers_unquantized( config, ) - layer = Mock(spec=LinearBase) for prefix in ( "vision_model.encoder.layers.0.attn.qkv", "vision_model.encoder.layers.1.mlp.fc2", @@ -203,10 +202,10 @@ def test_qianfan_online_fp8_keeps_vision_layers_unquantized( "mlp1.1", "mlp1.3", ): - assert config.resolve_quant_method_cls(layer, prefix) is None + assert config.resolve_quant_method_cls(LinearBase, prefix) is None resolved = config.resolve_quant_method_cls( - layer, "language_model.model.layers.0.self_attn.qkv_proj" + LinearBase, "language_model.model.layers.0.self_attn.qkv_proj" ) assert resolved is not None and resolved[-1] is method_cls assert other_config.ignored_layers == [] @@ -1320,9 +1319,7 @@ def test_online_quantization_targets_ignore_collision() -> None: ) ) with pytest.raises(ValueError, match="matches both quantization_config.ignore"): - config.resolve_quant_method_cls( - Mock(spec=LinearBase), "model.layers.0.self_attn.o_proj" - ) + config.resolve_quant_method_cls(LinearBase, "model.layers.0.self_attn.o_proj") def test_online_quantization_targets_reject_unsupported_layer() -> None: diff --git a/tests/quantization/test_online_shared_expert_fusion.py b/tests/quantization/test_online_shared_expert_fusion.py new file mode 100644 index 000000000000..2dc3d5da76bf --- /dev/null +++ b/tests/quantization/test_online_shared_expert_fusion.py @@ -0,0 +1,637 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests ROCm AITER fused shared experts with online quantization.""" + +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +from tests.quantization.utils import load_model_without_vllm_runner +from vllm._aiter_ops import rocm_aiter_ops +from vllm.config.model import ModelConfig +from vllm.config.quantization import ( + QuantizationConfigArgs, + resolve_quantization_config, +) +from vllm.model_executor.layers.fused_moe import FusedMoEFactory +from vllm.model_executor.layers.fused_moe.experts import rocm_aiter_moe +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend +from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig +from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig +from vllm.model_executor.layers.quantization.quark.quark_moe import ( + QuarkOCP_MX_MoEMethod, +) +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + mxfp4_quantize, +) +from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader +from vllm.model_executor.model_loader.reload.layerwise import ( + finalize_layerwise_reload, + initialize_layerwise_reload, + record_metadata_for_reloading, +) +from vllm.model_executor.model_loader.utils import get_model_architecture +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), + reason="ROCm AITER shared-expert tests require ROCm.", +) + +_QUARK_MXFP4_CONFIG: dict[str, Any] = { + "global_quant_config": { + "input_tensors": { + "dtype": "fp4", + "is_dynamic": True, + "qscheme": "per_group", + "ch_axis": -1, + "group_size": 32, + "block_size": None, + "symmetric": None, + "round_method": "half_even", + "scale_type": "float", + "scale_format": "e8m0", + "scale_calculation_mode": "even", + "mx_element_dtype": None, + "observer_cls": "PerBlockMXObserver", + "is_scale_quant": False, + "enable_buffer_reuse": False, + "max_input_numel": 4194304, + }, + "output_tensors": None, + "weight": { + "dtype": "fp4", + "is_dynamic": False, + "qscheme": "per_group", + "ch_axis": -1, + "group_size": 32, + "block_size": None, + "symmetric": None, + "round_method": "half_even", + "scale_type": "float", + "scale_format": "e8m0", + "scale_calculation_mode": "even", + "mx_element_dtype": None, + "observer_cls": "PerBlockMXObserver", + "is_scale_quant": False, + "enable_buffer_reuse": False, + "max_input_numel": 4194304, + }, + "bias": None, + "target_device": None, + }, + "algo_config": None, + "softmax_quant_spec": None, + "quant_method": "quark", + "layer_type_quant_config": {}, + "layer_quant_config": {}, + "exclude": ["re:^(?!.*\\.mlp\\.experts(?:\\.|$)).*$"], + "kv_cache_quant_config": {}, + "kv_cache_post_rope": False, + "quant_mode": "eager_mode", + "version": "0.12+9d3d471cdf1", + "export": { + "kv_cache_group": [], + "min_kv_scale": 0.0, + "pack_method": "reorder", + "weight_format": "real_quantized", + "weight_merge_groups": None, + }, +} + +_QUARK_MXFP4_LAYER_CONFIG: dict[str, Any] = _QUARK_MXFP4_CONFIG["global_quant_config"] +_QUARK_MXFP4_CONFIG = { + **_QUARK_MXFP4_CONFIG, + "global_quant_config": { + "input_tensors": {}, + "output_tensors": None, + "weight": {}, + "bias": None, + "target_device": None, + }, + "layer_quant_config": {"*mlp.experts*": _QUARK_MXFP4_LAYER_CONFIG}, +} + + +class _StubAttention(torch.nn.Module): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__() + self.o_proj = SimpleNamespace(reduce_results=True) + + +@pytest.fixture(autouse=True) +def reset_aiter_shared_expert_topk_metadata(monkeypatch: pytest.MonkeyPatch): + """Isolate AITER's process-global shared-expert routing buffer.""" + monkeypatch.setattr(rocm_aiter_moe, "aiter_topK_meta_data", None) + rocm_aiter_moe.init_aiter_topK_meta_data.cache_clear() + yield + rocm_aiter_moe.init_aiter_topK_meta_data.cache_clear() + + +def _write_minimal_moe_config(model_path: Path, architecture: str) -> None: + """Write a compact MoE config with one shared expert.""" + config = { + "architectures": [architecture], + "hidden_size": 256, + "intermediate_size": 512, + "moe_intermediate_size": 256, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "head_dim": 64, + "num_hidden_layers": 1, + "vocab_size": 256, + "max_position_embeddings": 64, + "q_lora_rank": 64, + "kv_lora_rank": 512, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "v_head_dim": 128, + "n_routed_experts": 4, + "num_experts": 4, + "num_experts_per_tok": 1, + "n_group": 1, + "topk_group": 1, + "n_shared_experts": 1, + "first_k_dense_replace": 0, + "moe_layer_freq": 1, + "rms_norm_eps": 1e-6, + "hidden_act": "silu", + "tie_word_embeddings": False, + "rope_theta": 10000, + "quantization_config": _QUARK_MXFP4_CONFIG, + } + if architecture == "AXK1ForCausalLM": + config["model_type"] = "axk1" + elif architecture == "MiniMaxM3SparseForCausalLM": + text_config = { + **config, + "model_type": "minimax_m3_text", + "dense_intermediate_size": 512, + "hidden_act": "swigluoai", + "num_local_experts": 4, + "n_shared_experts": 1, + "moe_layer_freq": [1], + "sparse_attention_config": None, + "quantization_config": { + **_QUARK_MXFP4_CONFIG, + "layer_quant_config": { + "*block_sparse_moe.experts*": _QUARK_MXFP4_LAYER_CONFIG + }, + "exclude": [r"re:^(?!.*\.block_sparse_moe\.experts(?:\.|$)).*$"], + }, + } + config = { + "architectures": [architecture], + "model_type": "minimax_m3_vl", + "text_config": text_config, + } + elif architecture == "Glm4MoeForCausalLM": + config["model_type"] = "glm4_moe" + elif architecture == "Glm4MoeLiteForCausalLM": + config["model_type"] = "glm4_moe_lite" + elif architecture in ("Qwen3NextForCausalLM", "Qwen3_5MoeForCausalLM"): + config.update( + model_type=( + "qwen3_next" + if architecture == "Qwen3NextForCausalLM" + else "qwen3_5_moe_text" + ), + layer_types=[ + "linear_attention" + if architecture == "Qwen3NextForCausalLM" + else "full_attention" + ], + head_dim=64, + linear_key_head_dim=64, + linear_value_head_dim=64, + linear_num_key_heads=4, + linear_num_value_heads=4, + shared_expert_intermediate_size=256, + ) + elif architecture == "Qwen4ExpForCausalLM": + config.update( + model_type="qwen4_exp_text", + layer_types=["full_attention"], + head_dim=64, + linear_key_head_dim=64, + linear_value_head_dim=64, + linear_num_key_heads=4, + linear_num_value_heads=4, + shared_expert_intermediate_size=256, + eos_token_id=1, + hc_count=2, + hc_lowrank=4, + ple_layer_ids=[], + ) + elif architecture == "Glm5NextForCausalLM": + config.update( + model_type="glm5_next", + layer_types=["linear_attention"], + mlp_layer_types=["sparse"], + qk_rope_head_dim=0, + index_topk=1, + index_kpool=1, + pad_token_id=None, + scoring_func="sigmoid", + topk_method="noaux_tc", + mhc=False, + ) + else: + config["model_type"] = ( + "deepseek_v3" + if "V3" in architecture + else "deepseek_v32" + if architecture == "DeepseekV32ForCausalLM" + else "deepseek_v2" + ) + if architecture in ("DeepseekV32ForCausalLM", "GlmMoeDsaForCausalLM"): + config.update( + index_topk=1, + index_kpool=1, + index_n_heads=1, + index_head_dim=32, + index_kv_lora_rank=32, + ) + (model_path / "config.json").write_text(json.dumps(config)) + + +@pytest.mark.parametrize( + "load_path", ["load_weights", "fused_gate_up", "weight_loader"] +) +def test_online_shared_expert_loads_bf16_weights_into_mxfp4_slot( + default_vllm_config, + dist_init, + load_path: str, +) -> None: + """A BF16 shared expert is quantized while routed MXFP4 weights are loaded. + + `weight_loader` covers DeepSeek-style model loaders, which bypass + `RoutedExperts.load_weights`. + """ + default_vllm_config.model_config = ModelConfig() + hidden_size = intermediate_size = 64 + num_routed_experts = 2 + device = current_platform.device_type + online_config = OnlineQuantizationConfig( + QuantizationConfigArgs(targets={"*shared_expert*": "mxfp4"}) + ) + quant_config = QuarkConfig(_QUARK_MXFP4_CONFIG) + quant_config.online_quantization_config = online_config + + with torch.device(device): + runner = FusedMoEFactory( + num_experts=num_routed_experts, + top_k=1, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + n_shared_experts=1, + fuse_shared_experts=True, + shared_expert_prefix="model.layers.0.mlp.shared_expert", + prefix="model.layers.0.mlp.experts", + quant_config=quant_config, + ) + layer = runner.routed_experts + assert isinstance(layer.quant_method, QuarkOCP_MX_MoEMethod) + + routed_gate, _ = mxfp4_quantize( + torch.randn(intermediate_size, hidden_size, dtype=torch.bfloat16) + ) + routed_up, _ = mxfp4_quantize( + torch.randn(intermediate_size, hidden_size, dtype=torch.bfloat16) + ) + routed_down, _ = mxfp4_quantize( + torch.randn(hidden_size, intermediate_size, dtype=torch.bfloat16) + ) + shared_gate = torch.randn(intermediate_size, hidden_size, dtype=torch.bfloat16) + shared_up = torch.randn(intermediate_size, hidden_size, dtype=torch.bfloat16) + shared_down = torch.randn(hidden_size, intermediate_size, dtype=torch.bfloat16) + + weights = [ + ("0.gate_proj.weight", routed_gate), + ("0.up_proj.weight", routed_up), + ("0.down_proj.weight", routed_down), + ] + if load_path == "load_weights": + weights += [ + ("2.gate_proj.weight", shared_gate), + ("2.up_proj.weight", shared_up), + ("2.down_proj.weight", shared_down), + ] + elif load_path == "fused_gate_up": + weights += [ + ("2.gate_up_proj.weight", torch.cat([shared_gate, shared_up])), + ("2.down_proj.weight", shared_down), + ] + list(layer.load_weights(weights)) + + if load_path == "weight_loader": + for stem, shard_id, weight in ( + ("w13", "w1", shared_gate), + ("w13", "w3", shared_up), + ("w2", "w2", shared_down), + ): + param = getattr(layer, f"{stem}_weight") + param.weight_loader( + param, weight, f"{layer.layer_name}.{stem}_weight", shard_id, 2 + ) + + expected_gate, expected_gate_scale = mxfp4_quantize(shared_gate) + expected_up, expected_up_scale = mxfp4_quantize(shared_up) + expected_down, expected_down_scale = mxfp4_quantize(shared_down) + assert torch.equal(layer.w13_weight[2, :intermediate_size], expected_gate) + assert torch.equal(layer.w13_weight[2, intermediate_size:], expected_up) + assert torch.equal(layer.w2_weight[2], expected_down) + assert torch.equal( + layer.w13_weight_scale[2, :intermediate_size], expected_gate_scale + ) + assert torch.equal(layer.w13_weight_scale[2, intermediate_size:], expected_up_scale) + assert torch.equal(layer.w2_weight_scale[2], expected_down_scale) + + +def test_online_shared_expert_reload_compatibility( + default_vllm_config, + dist_init, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Reload a BF16 shared weight into its MXFP4 slot via meta staging.""" + default_vllm_config.model_config = ModelConfig() + online_config = OnlineQuantizationConfig( + QuantizationConfigArgs(targets={"*shared_expert*": "mxfp4"}) + ) + quant_config = QuarkConfig(_QUARK_MXFP4_CONFIG) + quant_config.online_quantization_config = online_config + + with torch.device(current_platform.device_type): + layer = FusedMoEFactory( + num_experts=2, + top_k=1, + hidden_size=64, + intermediate_size=64, + n_shared_experts=1, + fuse_shared_experts=True, + shared_expert_prefix="model.layers.0.mlp.shared_expert", + prefix="model.layers.0.mlp.experts", + quant_config=quant_config, + ).routed_experts + monkeypatch.setattr( + layer.quant_method, "process_weights_after_loading", lambda _: None + ) + + routed_weights = [] + for expert_id in range(2): + for projection in ("gate_proj", "up_proj", "down_proj"): + weight, scale = mxfp4_quantize( + torch.randn(64, 64, dtype=torch.bfloat16) + ) + prefix = f"{expert_id}.{projection}" + routed_weights.extend( + [(f"{prefix}.weight", weight), (f"{prefix}.weight_scale", scale)] + ) + + shared_gate = torch.randn(64, 64, dtype=torch.bfloat16) + shared_up = torch.randn(64, 64, dtype=torch.bfloat16) + shared_down = torch.randn(64, 64, dtype=torch.bfloat16) + + def checkpoint_weights(gate: torch.Tensor): + return [ + *routed_weights, + ("2.gate_proj.weight", gate), + ("2.up_proj.weight", shared_up), + ("2.down_proj.weight", shared_down), + ] + + list(layer.load_weights(checkpoint_weights(shared_gate))) + old_weight = layer.w13_weight[2, :64].clone() + reloaded_gate = shared_gate + 1 + expected_weight, expected_scale = mxfp4_quantize(reloaded_gate) + assert not torch.equal(old_weight, expected_weight) + + record_metadata_for_reloading(layer) + initialize_layerwise_reload(layer) + + assert layer.w13_weight.is_meta + list(layer.load_weights(checkpoint_weights(reloaded_gate))) + finalize_layerwise_reload(layer, default_vllm_config.model_config) + assert not layer.w13_weight.is_meta + assert torch.equal(layer.w13_weight[2, :64], expected_weight) + assert torch.equal(layer.w13_weight_scale[2, :64], expected_scale) + + +# TODO: Add DeepseekV4ForConditionalGeneration once it uses +# resolve_layer_fused_shared_expert. +@pytest.mark.parametrize( + "architecture", + [ + "AXK1ForCausalLM", + "DeepseekForCausalLM", + "DeepseekV2ForCausalLM", + "DeepseekV32ForCausalLM", + "DeepseekV3ForCausalLM", + "Glm4MoeForCausalLM", + "Glm4MoeLiteForCausalLM", + "Glm5NextForCausalLM", + "GlmMoeDsaForCausalLM", + "MiniMaxM3SparseForCausalLM", + "Qwen3NextForCausalLM", + "Qwen3_5MoeForCausalLM", + "Qwen4ExpForCausalLM", + ], +) +def test_online_quantization( + architecture: str, + tmp_path: Path, + monkeypatch, + dist_init, + workspace_init, +) -> None: + """Online MXFP4 fuses each architecture's shared expert into its MoE.""" + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "1") + rocm_aiter_ops.refresh_env_variables() + if architecture == "Glm5NextForCausalLM": + from vllm.models.glm5next.common import model as glm5_next_model + + for attention in ("Glm5NextLinearAttention", "Glm5NextMLAAttention"): + monkeypatch.setattr(glm5_next_model, attention, _StubAttention) + monkeypatch.setattr( + glm5_next_model, "_fused_shared_experts_tuned", lambda _: True + ) + elif architecture == "Glm4MoeLiteForCausalLM": + from vllm.model_executor.models import glm4_moe_lite + + monkeypatch.setattr(glm4_moe_lite, "Glm4MoeLiteAttention", _StubAttention) + monkeypatch.setattr(glm4_moe_lite, "Glm4MoeLiteMLAAttention", _StubAttention) + elif architecture == "DeepseekV32ForCausalLM": + from vllm.models.deepseek_v32.amd import model as deepseek_v32_model + + monkeypatch.setattr( + deepseek_v32_model, "DeepseekV32MLAAttention", _StubAttention + ) + elif architecture == "Qwen3_5MoeForCausalLM": + from vllm.model_executor.models import qwen3_5 + + monkeypatch.setattr(qwen3_5, "Qwen3NextAttention", _StubAttention) + elif architecture == "Qwen4ExpForCausalLM": + from vllm.models.qwen4_exp.amd import model as qwen4_exp_model + + monkeypatch.setattr(qwen4_exp_model, "Qwen3NextAttention", _StubAttention) + _write_minimal_moe_config(tmp_path, architecture) + expected_shared_expert_name = ( + "shared_expert" if architecture.startswith("Qwen") else "shared_experts" + ) + target_pattern = f"*{expected_shared_expert_name}*" + + logged_messages: list[str] = [] + logged_warnings: list[str] = [] + + def record_info(message: str, *args: object, **_kwargs: object) -> None: + if "Quantizing " in message and "of types" in message: + logged_messages.append(message % args) + + def record_warning(message: str, *args: object, **_kwargs: object) -> None: + logged_warnings.append(message % args) + + monkeypatch.setattr( + "vllm.model_executor.model_loader.base_loader.logger.info", record_info + ) + monkeypatch.setattr( + "vllm.model_executor.layers.fused_moe.utils.logger.warning", record_warning + ) + + model, vllm_config = load_model_without_vllm_runner( + str(tmp_path), + dtype="bfloat16", + quantization="quark", + model_config_kwargs={ + "quantization_config": resolve_quantization_config( + "quark", {"targets": {target_pattern: "mxfp4"}} + ) + }, + model_loader_cls=DummyModelLoader, + ) + expected_model_cls, resolved_architecture = get_model_architecture( + vllm_config.model_config + ) + assert resolved_architecture == architecture + assert isinstance(model, expected_model_cls) + + fused_moes = [ + module + for module in model.modules() + if getattr(module, "is_fused_shared_expert_enabled", False) + and hasattr(module, "experts") + ] + assert len(fused_moes) == 1, f"{architecture} did not enable FSE" + moe = fused_moes[0] + shared_expert_name = ( + "shared_expert" if hasattr(moe, "shared_expert") else "shared_experts" + ) + assert shared_expert_name == expected_shared_expert_name + assert getattr(moe, shared_expert_name) is None + + routed_experts = moe.experts.routed_experts + assert isinstance(routed_experts.quant_method, QuarkOCP_MX_MoEMethod) + shared_expert_prefix = routed_experts.moe_config.shared_expert_prefix + assert shared_expert_prefix is not None + assert shared_expert_prefix.endswith(shared_expert_name) + assert routed_experts._fused_shared_expert_quantizer is not None + expert_map_manager = routed_experts.expert_map_manager + assert expert_map_manager.num_fused_shared_experts == 1 + assert expert_map_manager.map_global_to_local( + routed_experts.moe_config.num_experts + ) == (routed_experts.moe_config.num_experts) + expected_weight_dtype = { + Mxfp4MoeBackend.AITER_MXFP4_BF16: torch.float4_e2m1fn_x2, + Mxfp4MoeBackend.AITER_MXFP4_MXFP4: torch.float4_e2m1fn_x2, + Mxfp4MoeBackend.EMULATION: torch.uint8, + }[routed_experts.quant_method.mxfp4_backend] + assert routed_experts.w13_weight.dtype == expected_weight_dtype + assert routed_experts.w13_weight_scale.dtype == torch.uint8 + assert routed_experts.w2_weight.dtype == expected_weight_dtype + assert routed_experts.w2_weight_scale.dtype == torch.uint8 + assert vllm_config.quant_config is not None + assert vllm_config.quant_config.online_quantization_config is not None + assert not any( + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but " + "cannot be enabled" in warning + for warning in logged_warnings + ) + if architecture == "MiniMaxM3SparseForCausalLM": + # TODO: Use resolve_layer_fused_shared_expert in MiniMax M3 and remove + # this branch once its virtual fused projections are registered. + assert logged_messages == ["Quantizing 0 layers of types: "] + else: + assert logged_messages == [ + "Quantizing 2 layers of types: " + f"mlp.{shared_expert_name}.down_proj: 1 " + f"(from targets: {target_pattern}, mxfp4); " + f"mlp.{shared_expert_name}.gate_up_proj: 1 " + f"(from targets: {target_pattern}, mxfp4)" + ] + + +def test_online_gate_up_proj_target_disables_fse( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + dist_init, + workspace_init, +) -> None: + """A partial shared-expert target quantizes while disabling FSE.""" + architecture = "DeepseekForCausalLM" + target_pattern = "*shared_experts.gate_up_proj*" + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "1") + rocm_aiter_ops.refresh_env_variables() + _write_minimal_moe_config(tmp_path, architecture) + + logged_messages: list[str] = [] + logged_warnings: list[str] = [] + + def record_info(message: str, *args: object, **_kwargs: object) -> None: + if "Quantizing " in message and "of types" in message: + logged_messages.append(message % args) + + def record_warning(message: str, *args: object, **_kwargs: object) -> None: + logged_warnings.append(message % args) + + monkeypatch.setattr( + "vllm.model_executor.model_loader.base_loader.logger.info", record_info + ) + monkeypatch.setattr( + "vllm.model_executor.layers.fused_moe.utils.logger.warning", record_warning + ) + + model, _ = load_model_without_vllm_runner( + str(tmp_path), + dtype="bfloat16", + quantization="quark", + model_config_kwargs={ + "quantization_config": resolve_quantization_config( + "quark", {"targets": {target_pattern: "mxfp4"}} + ) + }, + model_loader_cls=DummyModelLoader, + ) + + moe = next( + module + for module in model.modules() + if hasattr(module, "is_fused_shared_expert_enabled") + and hasattr(module, "shared_experts") + ) + assert not moe.is_fused_shared_expert_enabled + assert moe.shared_experts is not None + assert logged_messages == [ + "Quantizing 1 layers of types: " + "mlp.shared_experts.gate_up_proj: 1 " + f"(from targets: {target_pattern}, mxfp4)" + ] + assert logged_warnings == [ + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but cannot be " + "enabled - skipping for this layer: online quantization targets only " + "part of the shared expert at model.layers.0.mlp.shared_experts." + ] diff --git a/tests/quantization/test_quantization_config_args.py b/tests/quantization/test_quantization_config_args.py index 7e2e21a3acbf..f24566f5fa8a 100644 --- a/tests/quantization/test_quantization_config_args.py +++ b/tests/quantization/test_quantization_config_args.py @@ -2,8 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for QuantizationConfigArgs parsing.""" -from unittest.mock import Mock - import pytest from tests.quantization.utils import quant_config_args, quant_spec @@ -244,6 +242,4 @@ def test_targets_reject_moe_only_shorthand_for_linear_layer(): ) with pytest.raises(ValueError, match="does not define a QuantSpec"): - config.resolve_quant_method_cls( - Mock(spec=LinearBase), "model.layers.0.self_attn.o_proj" - ) + config.resolve_quant_method_cls(LinearBase, "model.layers.0.self_attn.o_proj") diff --git a/tests/quantization/utils.py b/tests/quantization/utils.py index ed9776ad1c25..d8446952f0d7 100644 --- a/tests/quantization/utils.py +++ b/tests/quantization/utils.py @@ -15,7 +15,11 @@ VllmConfig, set_current_vllm_config, ) -from vllm.config.quantization import QuantizationConfigArgs, QuantSpec +from vllm.config.quantization import ( + QuantizationConfigArgs, + QuantSpec, + resolve_quantization_config, +) from vllm.model_executor.layers.quantization import get_quantization_config from vllm.model_executor.model_loader.default_loader import DefaultModelLoader from vllm.platforms import current_platform @@ -40,10 +44,17 @@ def _limit_num_hidden_layers( original_load_weights = model.load_weights def should_load_weight(name: str) -> bool: - for prefix in ("model.layers.", "layers."): - if name.startswith(prefix): - layer_idx = int(name.removeprefix(prefix).split(".", 1)[0]) - return layer_idx < num_hidden_layers + layers_prefix = "layers." + if name.startswith(layers_prefix): + layer_name = name.removeprefix(layers_prefix) + elif ".layers." in name: + layer_name = name.split(".layers.", maxsplit=1)[1] + else: + return True + + layer_idx_str = layer_name.split(".", maxsplit=1)[0] + if layer_idx_str.isdigit(): + return int(layer_idx_str) < num_hidden_layers return True def load_weights(weights): @@ -71,13 +82,19 @@ def load_model_without_vllm_runner( quantization=quantization, **(model_config_kwargs or {}), ) + model_config.quantization_config = resolve_quantization_config( + model_config.quantization, model_config.quantization_config + ) vllm_config_args = dict(vllm_config_kwargs or {}) vllm_config_args.setdefault( "compilation_config", CompilationConfig(mode=CompilationMode.NONE) ) vllm_config = VllmConfig(model_config=model_config, **vllm_config_args) hf_overrides = (model_config_kwargs or {}).get("hf_overrides") or {} - num_hidden_layers = hf_overrides.get("num_hidden_layers") + text_config_overrides = hf_overrides.get("text_config") or {} + num_hidden_layers = text_config_overrides.get( + "num_hidden_layers", hf_overrides.get("num_hidden_layers") + ) with set_current_vllm_config(vllm_config): model_loader = model_loader_cls(vllm_config.load_config) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 158bc90769f3..f20af8a6f826 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1293,6 +1293,7 @@ class FusedMoEConfig: has_bias: bool = False is_lora_enabled: bool = False has_hash_routing: bool = False + shared_expert_prefix: str | None = None # When True, the MoE skips its final cross-rank all-reduce (and the separate # shared-expert reduce), returning the partial per-rank sum. The caller is diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 294c538c7b02..6b7c64bdc429 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -122,6 +122,7 @@ def FusedMoEFactory( is_fused_checkpoint_transposed: bool = False, n_shared_experts: int | None = None, fuse_shared_experts: bool = False, + shared_expert_prefix: str | None = None, router_logits_dtype: torch.dtype | None = None, gate: torch.nn.Module | None = None, shared_experts: torch.nn.Module | None = None, @@ -197,6 +198,7 @@ def FusedMoEFactory( n_shared_experts: Number of shared experts to fuse into the routed grouped GEMM (ROCm; requires aiter FSE or the router-append path) fuse_shared_experts: Whether to enable shared-expert fusion. + shared_expert_prefix: Checkpoint prefix for the fused shared expert. router_logits_dtype: Data type for router logits buffers gate: Pre-configured gate module shared_experts: Pre-configured shared experts module @@ -360,6 +362,7 @@ def FusedMoEFactory( elastic_ep_max_dp_size=vllm_config.parallel_config.elastic_ep_max_dp_size, has_bias=has_bias, is_lora_enabled=vllm_config.lora_config is not None, + shared_expert_prefix=shared_expert_prefix, activation=moe_activation, device=vllm_config.device_config.device, routing_method=router.routing_method_type, # Not ideal diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index adc8a2701007..ae5829a311d9 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -26,6 +26,7 @@ from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( UnquantizedFusedMoEMethod, ) +from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, resolve_quant_method, @@ -153,6 +154,7 @@ def __init__( self.apply_router_weight_on_input = apply_router_weight_on_input # End random parameters self._loaded_expert_biases: set[str] = set() + self._fused_shared_expert_quantizer = self._get_fused_shared_expert_quantizer() self.quant_method = self._get_quant_method( self.layer_name, @@ -630,6 +632,31 @@ def _load_single_value( # _to_scalar's reshape(()) would reject the size-2 weight_shape. param_data[expert_id] = loaded_weight + def _get_fused_shared_expert_quantizer( + self, + ) -> Callable[[torch.Tensor], tuple[torch.Tensor, torch.Tensor]] | None: + """Return the quantizer for an online-quantized fused shared expert.""" + online_config = getattr(self.quant_config, "online_quantization_config", None) + prefix = self.moe_config.shared_expert_prefix + if ( + online_config is None + or prefix is None + or self.expert_map_manager.num_fused_shared_experts == 0 + ): + return None + # FSE compatibility checks require every shared-expert projection to + # use the routed experts' weight key. + resolved = online_config.resolve_quant_method_cls( + LinearBase, f"{prefix}.down_proj" + ) + if resolved is None: + return None + from vllm.model_executor.layers.quantization.online.base import ( + ONLINE_SHARED_EXPERT_QUANTIZERS, + ) + + return ONLINE_SHARED_EXPERT_QUANTIZERS[resolved[3].weight] + @overload def weight_loader( self, @@ -687,6 +714,35 @@ def weight_loader( return False if return_success else None # Hereafter, `expert_id` is local physical id + # A full-precision fused shared expert is quantized into the checkpoint + # layout, then loaded like a pre-quantized expert weight and scale. + if ( + self._fused_shared_expert_quantizer is not None + and self.moe_config.num_logical_experts + <= global_expert_id + < self.moe_config.num_logical_experts + + self.expert_map_manager.num_fused_shared_experts + and weight_name.endswith("_weight") + and loaded_weight.is_floating_point() + and loaded_weight.dtype != param.dtype + ): + weight, weight_scale = self._fused_shared_expert_quantizer( + loaded_weight.to(self.moe_config.device) + ) + stem = "w2" if shard_id == "w2" else "w13" + self.weight_loader( + getattr(self, f"{stem}_weight_scale"), + weight_scale, + f"{weight_name}_scale", + shard_id, + global_expert_id, + return_success=True, + ) + loaded = self.weight_loader( + param, weight, weight_name, shard_id, global_expert_id, True + ) + return loaded if return_success else None + # is_transposed: if the dim to shard the weight # should be flipped. Required by GPTQ/AWQ (K-first format). # should be whatever dimension intermediate_size_per_partition is diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index a98e181caeb2..c724a7c7cb21 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -84,6 +84,40 @@ def resolve_layer_fused_shared_expert( else (True, None) ) is_fused_shared_expert_enabled = fse_requested and fse_compatible + + # In online shared_expert quantization is used, register it to + # `online_quantization_config.quantized_layers` ahead of weight loading. + if is_fused_shared_expert_enabled and quant_config is not None: + online_quantization_config = quant_config.online_quantization_config + if online_quantization_config is not None: + from vllm.model_executor.layers.linear import ( + LinearBase, + UnquantizedLinearMethod, + ) + + for projection_name in ("gate_up_proj", "down_proj"): + projection_prefix = f"{prefix}.{shared_expert_name}.{projection_name}" + quant_method_metadata = ( + online_quantization_config.resolve_quant_method_cls( + LinearBase, projection_prefix + ) + ) + + if quant_method_metadata is None: + continue + + source, quant_key_str, target_pattern, _, quant_method_cls = ( + quant_method_metadata + ) + + if quant_method_cls in (None, UnquantizedLinearMethod): + continue + + online_quantization_config.quantized_layers[projection_prefix] = ( + source.value, + quant_key_str, + target_pattern, + ) if fse_requested and not is_fused_shared_expert_enabled: logger.warning( "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but " diff --git a/vllm/model_executor/layers/quantization/base_config.py b/vllm/model_executor/layers/quantization/base_config.py index 6c1f93473906..3ec2e728e16c 100644 --- a/vllm/model_executor/layers/quantization/base_config.py +++ b/vllm/model_executor/layers/quantization/base_config.py @@ -320,7 +320,7 @@ def resolve_quant_method( base_quant_method, (UnquantizedLinearMethod, UnquantizedFusedMoEMethod) ) online_target = quant_config.online_quantization_config.resolve_quant_method_cls( - layer, prefix + type(layer), prefix ) if checkpoint_is_quantized: diff --git a/vllm/model_executor/layers/quantization/online/base.py b/vllm/model_executor/layers/quantization/online/base.py index 3f35b56a31b8..511714837faa 100644 --- a/vllm/model_executor/layers/quantization/online/base.py +++ b/vllm/model_executor/layers/quantization/online/base.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Mapping +from collections.abc import Callable, Mapping from enum import Enum from types import MappingProxyType from typing import Any @@ -60,6 +60,7 @@ find_matching_patterns, get_layer_name_after_index, ) +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mxfp4_quantize from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kFp8Static128BlockSym, @@ -104,6 +105,15 @@ class OnlineQuantizationSource(str, Enum): kNvfp4Static: Nvfp4OnlineMoEMethod, } +# Quantizers for full-precision fused shared-expert weights, keyed by the routed +# experts' weight QuantKey. Each shard is quantized to the checkpoint layout on +# arrival, so only block-local schemes are supported. +ONLINE_SHARED_EXPERT_QUANTIZERS: dict[ + QuantKey, Callable[[torch.Tensor], tuple[torch.Tensor, torch.Tensor]] +] = { + kMxfp4Static: mxfp4_quantize, +} + def _find_matching_targets( prefix: str, @@ -217,14 +227,14 @@ def _get_method_cls( self, spec: QuantSpec | None, table: dict[QuantKey, type], - layer: torch.nn.Module, + layer_cls: type[torch.nn.Module], ) -> type | None: """Resolve the online method class for a layer's quantization spec. Args: spec: Quantization specification to resolve. table: Mapping from weight quantization keys to method classes. - layer: Layer that will use the resolved method. + layer_cls: Type of the layer that will use the resolved method. Returns: The matching method class, or None when ``spec`` has no weight @@ -239,7 +249,7 @@ def _get_method_cls( cls = table.get(spec.weight) if cls is None: raise ValueError( - f"online quantization for {type(layer).__name__} with " + f"online quantization for {layer_cls.__name__} with " f"weight={spec.weight} is not supported; supported weight " f"keys: {sorted(str(k) for k in table)}" ) @@ -254,12 +264,12 @@ def _get_method_cls( return cls def resolve_quant_method_cls( - self, layer: torch.nn.Module, prefix: str + self, layer_cls: type[torch.nn.Module], prefix: str ) -> tuple[OnlineQuantizationSource, str, str | None, QuantSpec, type] | None: """Resolve quantization metadata and method class without instantiating it. Args: - layer: Layer for which to resolve online quantization. + layer_cls: Type of layer for which to resolve online quantization. prefix: Fully qualified layer name. Returns: @@ -271,17 +281,17 @@ def resolve_quant_method_cls( quant_spec: QuantSpec | None if self.args.targets is not None: resolved_pattern = self._resolve_targets_quant_method_metadata( - prefix, layer + prefix, layer_cls ) if resolved_pattern is None: return None source, quant_key_str, target_pattern, quant_spec, table = resolved_pattern else: - if isinstance(layer, LinearBase): + if issubclass(layer_cls, LinearBase): source = OnlineQuantizationSource.linear quant_spec = self.args.linear table = _ONLINE_LINEAR_METHODS - elif isinstance(layer, RoutedExperts): + elif issubclass(layer_cls, RoutedExperts): source = OnlineQuantizationSource.moe quant_spec = self.args.moe table = _ONLINE_MOE_METHODS @@ -298,14 +308,14 @@ def resolve_quant_method_cls( quant_key_str = str(quant_spec) target_pattern = None - quant_method_cls = self._get_method_cls(quant_spec, table, layer) + quant_method_cls = self._get_method_cls(quant_spec, table, layer_cls) if quant_method_cls is None: return None assert quant_spec is not None return source, quant_key_str, target_pattern, quant_spec, quant_method_cls def _resolve_targets_quant_method_metadata( - self, prefix: str, layer: torch.nn.Module + self, prefix: str, layer_cls: type[torch.nn.Module] ) -> ( tuple[OnlineQuantizationSource, str, str, QuantSpec, dict[QuantKey, type]] | None @@ -314,7 +324,7 @@ def _resolve_targets_quant_method_metadata( Args: prefix: Fully qualified layer name. - layer: Layer matched against configured target patterns. + layer_cls: Type of layer matched against configured target patterns. Returns: A tuple of source, quantization key string, target pattern, spec, @@ -349,22 +359,22 @@ def _resolve_targets_quant_method_metadata( target_pattern = matches[0] quant_key_str = self.args.targets[target_pattern] shorthand = _ONLINE_SHORTHANDS[quant_key_str] - if isinstance(layer, LinearBase): + if issubclass(layer_cls, LinearBase): quant_spec = shorthand.linear table = _ONLINE_LINEAR_METHODS - elif isinstance(layer, RoutedExperts): + elif issubclass(layer_cls, RoutedExperts): quant_spec = shorthand.moe table = _ONLINE_MOE_METHODS else: raise ValueError( f"Layer {prefix} was matched by quantization_config.targets " f"({target_pattern}), but online quantization is not supported for " - f"{type(layer).__name__}." + f"{layer_cls.__name__}." ) if quant_spec is None: raise ValueError( f"targets pattern {target_pattern} = {quant_key_str} does " - f"not define a QuantSpec for {type(layer).__name__} layers " + f"not define a QuantSpec for {layer_cls.__name__} layers " f"(matched at {prefix})." ) return ( @@ -379,7 +389,7 @@ def get_quant_method( self, layer: torch.nn.Module, prefix: str ) -> "QuantizeMethodBase | None": # `targets` takes precedence over `moe` and `linear` and is exclusive. - resolved = self.resolve_quant_method_cls(layer, prefix) + resolved = self.resolve_quant_method_cls(type(layer), prefix) if resolved is not None: source, quant_key_str, target_pattern, _, quant_method_cls = resolved self.quantized_layers[prefix] = ( diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index 23a36c8fbc3b..f4662ad5a66e 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -12,8 +12,6 @@ FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend - from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey - import vllm.envs as envs from vllm import _custom_ops as ops from vllm.config import get_current_vllm_config @@ -40,6 +38,7 @@ ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, + QuantKey, amax_for_moe_weight_quant, amax_for_tp_weight_quant, create_fp8_quant_key, @@ -122,6 +121,7 @@ class OnlineLinearBase(LinearMethodBase): weights onto meta device and materializes them just-in-time.""" uses_meta_device: bool = True + activation_quant_key: QuantKey | None def __init__(self): self.out_dtype = torch.get_default_dtype() @@ -203,6 +203,10 @@ class Fp8PerTensorOnlineLinearMethod(OnlineLinearBase): """Online tensorwise FP8 linear quantization. Loads fp16/bf16 weights and quantizes them per-tensor during loading.""" + activation_quant_key = ( + kFp8DynamicTokenSym if cutlass_fp8_supported() else kFp8DynamicTensorSym + ) + def __init__(self): super().__init__() @@ -211,11 +215,6 @@ def __init__(self): self.use_marlin = False self.marlin_input_dtype = None self.weight_quant_key = kFp8StaticTensorSym - # Use per-token quantization for better perf if dynamic and cutlass - if cutlass_fp8_supported(): - self.activation_quant_key = kFp8DynamicTokenSym - else: - self.activation_quant_key = kFp8DynamicTensorSym def create_weights( self, @@ -237,6 +236,8 @@ def create_weights( **extra_weight_attrs, ) + # TODO: init_fp8_linear_kernel should support activation_quant_key=None + assert self.activation_quant_key is not None self.fp8_linear = init_fp8_linear_kernel( activation_quant_key=self.activation_quant_key, weight_quant_key=self.weight_quant_key, @@ -305,13 +306,11 @@ class Fp8PerBlockOnlineLinearMethod(OnlineLinearBase): """Online blockwise FP8 linear quantization. Loads fp16/bf16 weights and quantizes them per-block during loading.""" + activation_quant_key = kFp8Dynamic128Sym + def __init__(self): super().__init__() self.weight_block_size = [128, 128] - self.activation_quant_key = create_fp8_quant_key( - static=False, - group_shape=GroupShape(1, self.weight_block_size[0]), - ) self.weight_quant_key = create_fp8_quant_key( static=True, group_shape=GroupShape(*self.weight_block_size) ) @@ -337,6 +336,8 @@ def create_weights( ) layer.weight_block_size = self.weight_block_size + # TODO: init_fp8_linear_kernel should support activation_quant_key=None + assert self.activation_quant_key is not None self.fp8_linear = init_fp8_linear_kernel( activation_quant_key=self.activation_quant_key, weight_quant_key=self.weight_quant_key, diff --git a/vllm/model_executor/layers/quantization/online/mxfp8.py b/vllm/model_executor/layers/quantization/online/mxfp8.py index 9cddab85c12a..ecc196e9f185 100644 --- a/vllm/model_executor/layers/quantization/online/mxfp8.py +++ b/vllm/model_executor/layers/quantization/online/mxfp8.py @@ -9,6 +9,7 @@ from torch.nn import Module from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey if TYPE_CHECKING: import vllm.model_executor.layers.fused_moe.modular_kernel as mk @@ -42,6 +43,8 @@ class Mxfp8OnlineLinearMethod(OnlineLinearBase): FP8 with block-32 scales) during weight loading. """ + activation_quant_key: QuantKey | None = None + def create_weights( self, layer: torch.nn.Module, diff --git a/vllm/model_executor/layers/quantization/utils/config_utils.py b/vllm/model_executor/layers/quantization/utils/config_utils.py index eb50ed1a23db..6e294adb3a65 100644 --- a/vllm/model_executor/layers/quantization/utils/config_utils.py +++ b/vllm/model_executor/layers/quantization/utils/config_utils.py @@ -8,9 +8,13 @@ import regex as re +from vllm.logger import init_logger + if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +logger = init_logger(__name__) + def get_quark_ocp_mx_group_size( quant_config: "QuantizationConfig | None", @@ -70,19 +74,101 @@ def is_shared_expert_quant_fse_compatible( A compatibility flag and, when incompatible, the reason. """ - if projection_names is None: - projection_names = ["gate_up_proj", "down_proj"] - - if quant_config is None: - return True, None - from vllm.model_executor.layers.quantization.fp8 import Fp8Config + from vllm.model_executor.layers.quantization.online.fp8 import OnlineLinearBase from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig from vllm.model_executor.layers.quantization.utils.quant_utils import ( is_layer_skipped, ) from vllm.models.deepseek_v4.quant_config import DeepseekV4FP8Config + if projection_names is None: + projection_names = ["gate_up_proj", "down_proj"] + + if quant_config is None: + return True, None + + online_quant_config = quant_config.online_quantization_config + shared_expert_targets = None + if online_quant_config is not None: + from vllm.model_executor.layers.linear import LinearBase + + online_quant_config.packed_modules_mapping = quant_config.packed_modules_mapping + shared_expert_targets = [ + online_quant_config.resolve_quant_method_cls( + LinearBase, f"{shared_expert_prefix}.{projection_name}" + ) + for projection_name in projection_names + ] + if all(target is None for target in shared_expert_targets): + shared_expert_targets = None + + if shared_expert_targets is not None: + from vllm.model_executor.layers.fused_moe import RoutedExperts + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + from vllm.model_executor.layers.quantization.online.base import ( + ONLINE_SHARED_EXPERT_QUANTIZERS, + ) + + # NOTE: online shared experts quantization check is only implemented for quark + # quant method at the moment. This can be extended here for other quant methods. + # Requirement is: + # - Retrieve routed weight quant key + # - Retrieve routed activation quant key + # and ensure the online quantization set on shared expert match them. + if isinstance(quant_config, QuarkConfig): + routed_weight_key, routed_activation_key, _ = ( + quant_config.get_quant_method_target(expert_prefix, RoutedExperts) + ) + else: + reason = ( + "online quantization FSE compatibility check is not implemented for " + f"{type(quant_config).__name__}" + ) + logger.warning_once(reason) + return False, reason + + if routed_weight_key not in ONLINE_SHARED_EXPERT_QUANTIZERS: + return ( + False, + "online shared-expert quantization is not supported for routed " + f"expert weight key {routed_weight_key}", + ) + + for shared_expert_target in shared_expert_targets: + if shared_expert_target is None: + return ( + False, + "online quantization targets only part of the shared expert at " + f"{shared_expert_prefix}", + ) + + _, _, _, shared_quant_spec, shared_method_cls = shared_expert_target + if shared_method_cls in (None, UnquantizedLinearMethod): + return ( + False, + "online quantization excludes FSE weights at " + f"{shared_expert_prefix}", + ) + assert issubclass(shared_method_cls, OnlineLinearBase) + + # NOTE: We can not rely on shared_quant_spec.activation_quant_key + # as _ONLINE_SHORTHANDS: dict[str, QuantSpec] in quantization.py + # does not define the default activation quant key. + # `online_quant_config.resolve_quant_method_cls` should return the + # activation quant key as well. Only some online methods define it + # on the class, so compare the weight key first. + if ( + shared_quant_spec.weight != routed_weight_key + or shared_method_cls.activation_quant_key != routed_activation_key + ): + return ( + False, + "online shared-expert quantization keys do not match the " + "routed expert keys", + ) + return True, None + if isinstance(quant_config, DeepseekV4FP8Config): from vllm.config import get_current_vllm_config from vllm.model_executor.models.utils import extract_layer_index diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index da38035db9bd..d8ceb5083c2f 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -199,6 +199,7 @@ def __init__( if self.is_fused_shared_expert_enabled else None, fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_experts", ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 70d87529f663..a1a919fae644 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -396,6 +396,7 @@ def __init__( if self.is_fused_shared_expert_enabled else None, fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_experts", router_logits_dtype=self.gate.out_dtype, ) diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py index 5e5974e209a3..4e541a49413e 100644 --- a/vllm/model_executor/models/glm4_moe.py +++ b/vllm/model_executor/models/glm4_moe.py @@ -204,6 +204,7 @@ def __init__( config.n_shared_experts if self.is_fused_shared_expert_enabled else None ), fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_experts", ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 19fc8c156286..9580894e22ee 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -232,6 +232,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = ""): is_sequence_parallel=self.is_sequence_parallel, n_shared_experts=1 if self.shared_expert is None else None, fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_expert", shared_expert_gate=self.shared_expert_gate if self.shared_expert is None else None, diff --git a/vllm/models/glm5next/common/model.py b/vllm/models/glm5next/common/model.py index 778738c13712..26becc9619e7 100644 --- a/vllm/models/glm5next/common/model.py +++ b/vllm/models/glm5next/common/model.py @@ -307,6 +307,7 @@ def __init__( if self.is_fused_shared_expert_enabled else None, fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_experts", router_logits_dtype=self.gate.out_dtype, swiglu_limit=swiglu_limit, ) diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 9105f97af848..3bf89286034f 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -465,6 +465,7 @@ def __init__( self.n_shared_experts if self.is_fused_shared_expert_enabled else None ), fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_experts", quant_config=quant_config, prefix=f"{prefix}.experts", ) From f4dde3132c94cccce485a1a59acac1d06739a59d Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Thu, 8 Oct 2026 18:18:05 +0200 Subject: [PATCH 111/624] [CI][ROCm] Reuse AITER's prebuilt MoE module in the modular-kernel harness (#60633) Signed-off-by: Stefan Koncarevic --- .../moe/modular_kernel_tools/common.py | 26 ++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 9ae6092c944c..5f6349a32880 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from dataclasses import dataclass +from dataclasses import dataclass, replace from types import SimpleNamespace from typing import Any @@ -702,6 +702,21 @@ def make_modular_kernel( return modular_kernel +def _shuffle_weights_for_aiter(rank_weights: WeightTensors) -> WeightTensors: + """Pre-shuffle weights so AITER selects its prebuilt `preshuffle_on` module. + + Production shuffles in `process_weights_after_loading`; this harness builds its + tensors directly, so without this every call asks for a `preshuffle_off` module + that is not prebuilt, and JIT-compiles a kernel the image already ships. + """ + from vllm._aiter_ops import rocm_aiter_ops + + w1, w2 = rocm_aiter_ops.shuffle_weights(rank_weights.w1, rank_weights.w2) + w1.is_shuffled = True + w2.is_shuffled = True + return replace(rank_weights, w1=w1, w2=w2) + + def _maybe_convert_weights_for_experts( config: Config, rank_weights: WeightTensors, @@ -715,6 +730,15 @@ def _maybe_convert_weights_for_experts( fe_type = config.fused_experts_type fe_name = getattr(fe_type, "__name__", "") + # AITER's prebuilt modules carry a gfx950 instance table even on gfx942 and reject + # intermediate sizes that are not a multiple of 256, but only for the fp8 per-tensor + # and per-token schemes. Those keep JIT-compiling until + # https://github.com/ROCm/aiter/issues/5766 is fixed. + if fe_name == "AiterExperts" and ( + config.quant_dtype is None or config.quant_block_shape is not None + ): + return _shuffle_weights_for_aiter(rank_weights) + backend: Fp8MoeBackend | None = None if fe_name == "TrtLlmFp8ExpertsModular": backend = Fp8MoeBackend.FLASHINFER_TRTLLM From 8674a794c788d4c7dac10c6f7ddfdae677588d18 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Elvir=20Crn=C4=8Devi=C4=87?= Date: Thu, 8 Oct 2026 18:42:15 +0200 Subject: [PATCH 112/624] [Perf] Enable Kimi K3 shared expert sharding with DeepEP v2 (#56875) Co-authored-by: Misha Goin --- .../models/kimi_k3/test_sequence_parallel.py | 21 ++++++++++++++++ vllm/models/kimi_k3/nvidia/model.py | 24 +++++++++++++++---- 2 files changed, 41 insertions(+), 4 deletions(-) diff --git a/tests/models/kimi_k3/test_sequence_parallel.py b/tests/models/kimi_k3/test_sequence_parallel.py index 9e8faea099cb..3057f85c7e44 100644 --- a/tests/models/kimi_k3/test_sequence_parallel.py +++ b/tests/models/kimi_k3/test_sequence_parallel.py @@ -301,6 +301,27 @@ def test_shard_sequence_parallel_mlp_gating( ) +@pytest.mark.parametrize( + ("moe_backend", "all2all_backend", "expected"), + [ + ("deep_gemm_mega_moe", "naive", True), + ("auto", "deepep_v2", True), + ("auto", "deepep_low_latency", False), + ], +) +def test_shared_expert_sharding_backend_gate( + moe_backend: str, + all2all_backend: str, + expected: bool, +): + vllm_config = SimpleNamespace( + kernel_config=SimpleNamespace(moe_backend=moe_backend), + parallel_config=SimpleNamespace(all2all_backend=all2all_backend), + ) + + assert kimi_model.can_shard_sequence_parallel_shared_expert(vllm_config) is expected + + def test_sharded_sequence_parallel_mlp_matches_replicated(default_vllm_config): """Sharded SP MLP must reproduce the replicated result for every token. diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 84e93b82fbb4..22d2c4a1a9a1 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -156,6 +156,20 @@ def shard_sequence_parallel_mlp( ) +def can_shard_sequence_parallel_shared_expert(vllm_config: VllmConfig) -> bool: + """Whether the active MoE backend preserves the SP shared-expert contract. + + MegaMoE invokes the shared expert directly. DeepEP v2 invokes it through + ``MoERunner``, but its EP combine returns the routed output in the local + sequence-sharded layout, so the shared expert can use the same + all-gather/partial-GEMM/reduce-scatter implementation. + """ + return ( + vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + or vllm_config.parallel_config.all2all_backend == "deepep_v2" + ) + + def maybe_init_gemm_rs_ar(vllm_config: VllmConfig, use_sequence_parallel: bool) -> bool: # Both feature flags may be enabled; the worker's static SP topology binds # its singleton to exactly one mode. @@ -617,10 +631,12 @@ def __init__( quant_config=quant_config, reduce_results=False, use_sequence_parallel=use_sequence_parallel, - # Only the MegaMoE path calls the shared experts directly; the - # FusedMoE path below hands them to the runner, which fuses - # their reduction and assumes the replicated layout. - can_shard_sequence_parallel=self.use_mega_moe, + # MegaMoE calls the shared expert directly. DeepEP v2 hands it + # to MoERunner, but preserves the local SP output layout needed + # by the sharded shared-expert path. + can_shard_sequence_parallel=can_shard_sequence_parallel_shared_expert( + vllm_config + ), run_gemm_rs_ar=run_gemm_rs_ar, prefix=f"{prefix}.shared_experts", activation_situ_beta=activation_situ_beta, From 242e4213fc9845ff6fe607af1aee626fd8acc990 Mon Sep 17 00:00:00 2001 From: Doug Smith Date: Thu, 8 Oct 2026 12:43:39 -0400 Subject: [PATCH 113/624] [Doc] Require GCC 13 for CUDA source builds (#60259) Signed-off-by: Doug Smith --- CMakeLists.txt | 12 ++++++++++++ docs/getting_started/installation/gpu.cuda.inc.md | 9 ++------- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 4ed5ff2573c1..b4dd9ada5eb6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1568,6 +1568,18 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") include(cmake/external_projects/triton_kernels.cmake) endif() +# DeepGEMM and DeepSelect host code uses C++20 , which libstdc++ +# provides only from GCC 13 onward. +if (VLLM_GPU_LANG STREQUAL "CUDA" + AND CMAKE_CXX_COMPILER_ID STREQUAL "GNU" + AND CMAKE_CXX_COMPILER_VERSION VERSION_LESS "13") + message(FATAL_ERROR + "GCC >= 13 is required to build vLLM with CUDA (found ${CMAKE_CXX_COMPILER_VERSION}). " + "The DeepGEMM and DeepSelect dependencies use C++20 in host code, " + "which libstdc++ does not provide before GCC 13. " + "See: https://github.com/vllm-project/vllm/issues/58158") +endif() + # For CUDA we also build and ship some external projects. if (VLLM_GPU_LANG STREQUAL "CUDA") include(cmake/external_projects/deepgemm.cmake) diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md index 86e38c570788..6ad581ac97d1 100644 --- a/docs/getting_started/installation/gpu.cuda.inc.md +++ b/docs/getting_started/installation/gpu.cuda.inc.md @@ -140,13 +140,8 @@ You can find more information about vLLM's wheels in [Install the latest code](# #### Full build (with compilation) {#full-build} !!! note "Compiler requirement" - Building from source requires GCC/G++ ≥ 11.3. PyTorch's C++20 headers are - not compatible with GCC 10 or GCC < 11.3. On Ubuntu 22.04: - ```bash - sudo apt-get install -y gcc-11 g++-11 - sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 110 \ - --slave /usr/bin/g++ g++ /usr/bin/g++-11 - ``` + Building from source requires GCC/G++ ≥ 13 ([#58158](https://github.com/vllm-project/vllm/issues/58158)). + Ubuntu 24.04, the base of the default vLLM image, ships GCC 13 by default. If you want to modify C++ or CUDA code, you'll need to build vLLM from source. This can take several minutes: From 705e124c05cffe7c00dd19434922c94e0610d674 Mon Sep 17 00:00:00 2001 From: Francisco Javier Arceo Date: Thu, 8 Oct 2026 14:00:03 -0400 Subject: [PATCH 114/624] [Bugfix] Followup to #59299: preserve Unicode in structured decision state prompts (#60651) Signed-off-by: Francisco Javier Arceo <4163062+franciscojavierarceo@users.noreply.github.com> Co-authored-by: Misha Goin --- .../generate/structured_decisions/test_strategies.py | 6 ++++++ vllm/entrypoints/generate/structured_decisions/serving.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index c4f4bce61be1..017788f7b1d8 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -10,6 +10,7 @@ register_structured_decisions_api_router, ) from vllm.entrypoints.generate.structured_decisions.question_types import LABELS +from vllm.entrypoints.generate.structured_decisions.serving import state_text from vllm.entrypoints.generate.structured_decisions.strategies import ( NextTokenStrategy, reply_label_ids, @@ -40,6 +41,11 @@ def test_route_needs_the_flag(): assert ("/v1/systemone" in paths) == enabled +def test_structured_state_keeps_unicode_in_prompt(): + state = {"message": "Français 日本語"} + assert state_text(state) == '{"message": "Français 日本語"}' + + @pytest.fixture(scope="module") def qwen(): from transformers import AutoTokenizer diff --git a/vllm/entrypoints/generate/structured_decisions/serving.py b/vllm/entrypoints/generate/structured_decisions/serving.py index 7e2136b8a25e..4921bfcf3574 100644 --- a/vllm/entrypoints/generate/structured_decisions/serving.py +++ b/vllm/entrypoints/generate/structured_decisions/serving.py @@ -25,7 +25,7 @@ def state_text(state: Any) -> str: - return state if isinstance(state, str) else json.dumps(state) + return state if isinstance(state, str) else json.dumps(state, ensure_ascii=False) def parse_questions( From 972195f4f589688fd4cfbbcfb55a5b5502f5af35 Mon Sep 17 00:00:00 2001 From: Tialo <65392801+Tialo@users.noreply.github.com> Date: Thu, 8 Oct 2026 21:19:20 +0300 Subject: [PATCH 115/624] [CI][Build] build macos wheels using abi3 (#60312) Co-authored-by: Misha Goin --- .buildkite/scripts/build-macos-wheel.sh | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.buildkite/scripts/build-macos-wheel.sh b/.buildkite/scripts/build-macos-wheel.sh index cb95d92ff14f..d956f71038d3 100755 --- a/.buildkite/scripts/build-macos-wheel.sh +++ b/.buildkite/scripts/build-macos-wheel.sh @@ -25,6 +25,7 @@ export CMAKE_BUILD_PARALLEL_LEVEL="${CMAKE_BUILD_PARALLEL_LEVEL:-4}" uv venv --python 3.12 uv pip install -r requirements/build/cpu.txt --index-strategy unsafe-best-match -uv build --wheel --no-build-isolation -o artifacts/dist +uv build --wheel --no-build-isolation -o artifacts/dist \ + --config-setting=--build-option=--py-limited-api=cp38 ls -l artifacts/dist/*.whl From 1664c4f44416804dc47e3677e50933e6e9b4f45e Mon Sep 17 00:00:00 2001 From: kyleliang-nv Date: Thu, 8 Oct 2026 11:27:57 -0700 Subject: [PATCH 116/624] [Bugfix][MiniMax-M3] Guard sparse decode sentinel indices (#54597) Signed-off-by: Kyle Liang Signed-off-by: Kyle Liang Co-authored-by: Kyle Liang Co-authored-by: OpenAI Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .buildkite/test_areas/kernels.yaml | 17 +++- tests/kernels/attention/test_minimax_m3.py | 97 +++++++++++++++++++ .../minimax_m3/common/ops/sparse_attn.py | 24 +++-- 3 files changed, 128 insertions(+), 10 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index e0d6c4ce40a6..aa0971b4c55e 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -552,6 +552,21 @@ steps: - pytest -v -s kernels/attention/test_deepgemm_attention.py - pytest -v -s quantization/test_cutlass_w4a16.py +- label: ":nvidia: (B200) MiniMax-M3 Attention Kernels" + key: kernels-minimax-m3-attention-b200 + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/" + device: b200-k8s + num_devices: 1 + source_file_dependencies: + - vllm/models/minimax_m3/ + - vllm/config/cache.py + - vllm/v1/kv_cache_interface.py + - tests/v1/attention/utils.py + - tests/kernels/attention/test_minimax_m3.py + commands: + - pytest -v -s tests/kernels/attention/test_minimax_m3.py + - label: ":nvidia: (B200) Kernels Shard %N" key: kernels-b200 timeout_in_minutes: 30 @@ -601,7 +616,6 @@ steps: - tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py - tests/kernels/mamba/test_gdn_prefill_cutedsl.py - tests/kernels/test_bf16x3_router_gemm_cutedsl.py - - tests/kernels/attention/test_minimax_m3.py - tests/kernels/test_ll_bf16_gemm.py - tests/kernels/test_top_k_per_row.py - tests/kernels/quantization/test_silu_mul_nvfp4_quant.py @@ -647,7 +661,6 @@ steps: pytest -v -s tests/kernels/moe/test_flashinfer.py pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py pytest -v -s tests/kernels/mamba/test_gdn_prefill_cutedsl.py - pytest -v -s tests/kernels/attention/test_minimax_m3.py ;; *) echo "unexpected BUILDKITE_PARALLEL_JOB=$$BUILDKITE_PARALLEL_JOB" >&2 diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index 7f0e07309331..7ee900fd3dd1 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -2612,6 +2612,103 @@ def test_decode_sparse_attention_correctness( assert error.max().item() < 1.7e-2 +@pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True) +@pytest.mark.parametrize("capture_graph", [False, True]) +@pytest.mark.parametrize("num_reqs", [2, 17], ids=["split-k", "single-chunk"]) +@pytest.mark.parametrize( + "invalid_block", [-1, 1_000_000], ids=["sentinel", "out-of-range"] +) +def test_decode_sparse_attention_ignores_invalid_topk( + kv_layout: KVCacheLayout, capture_graph: bool, num_reqs: int, invalid_block: int +): + """Unused top-k entries must not be dereferenced as block-table indices.""" + torch.manual_seed(0) + decode_query_len = 4 + seq_lens = torch.full( + (num_reqs,), decode_query_len, device="cuda", dtype=torch.int32 + ) + seq_lens[1] = 4096 + q = torch.randn( + seq_lens.numel() * decode_query_len, + NUM_Q_HEADS, + HEAD_DIM, + device="cuda", + dtype=DTYPE, + ) + kv_cache = _allocate_main_kv_via_contract(1, kv_layout) + + # All logical blocks alias the single valid page. With a contiguous multi-row + # table, row 1's raw pointer offset -1 resolves to the final element of row + # 0. Poison that exact location so an unsafe load deterministically forms an + # out-of-range KV-cache pointer. + clean_block_table = torch.zeros(num_reqs, 32, device="cuda", dtype=torch.int32) + block_table = clean_block_table.clone() + block_table[0, -1] = 1_000_000 + topk_idx = torch.full( + (NUM_KV_HEADS, q.shape[0], TOPK), + -1, + device="cuda", + dtype=torch.int32, + ) + topk_idx[..., 0] = 0 + # Request 1 has invalid entries both before and after its valid block. + # Two requests produce empty split-K chunks; 17 put the full top-k in one + # chunk, exercising neutral updates before and after a nonempty softmax. + request_topk = topk_idx[:, decode_query_len : 2 * decode_query_len] + request_topk.fill_(invalid_block) + request_topk[..., TOPK // 2] = 0 + + # A valid-only run isolates sentinel handling from BF16-vs-FP32 rounding. + # Include block 0 exactly once; replacing sentinels with 0 would duplicate + # its contribution. The general correctness test covers reference accuracy. + valid_topk = torch.zeros( + (NUM_KV_HEADS, q.shape[0], 1), device="cuda", dtype=torch.int32 + ) + expected = torch.empty_like(q) + minimax_m3_sparse_attn_decode( + q, + kv_cache, + valid_topk, + clean_block_table, + seq_lens, + NUM_KV_HEADS, + SM_SCALE, + expected, + decode_query_len, + ) + actual = torch.empty_like(q) + + def run_decode() -> None: + minimax_m3_sparse_attn_decode( + q, + kv_cache, + topk_idx, + block_table, + seq_lens, + NUM_KV_HEADS, + SM_SCALE, + actual, + decode_query_len, + ) + + if capture_graph: + # An unfixed kernel faults during capture and leaves the CUDA context + # unusable. Run negative-control reproductions in a dedicated process. + run_decode() + current_platform.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_decode() + for _ in range(3): + graph.replay() + else: + run_decode() + current_platform.synchronize() + + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + def test_decode_wrong_layout_breaks_parity(): """Negative (AC-3/AC-5): consuming the physical HND buffer as if it were already contiguous-NHD (i.e. skipping the allocator's inverse permute) diff --git a/vllm/models/minimax_m3/common/ops/sparse_attn.py b/vllm/models/minimax_m3/common/ops/sparse_attn.py index 8dfd6e0c1996..6426115ae356 100644 --- a/vllm/models/minimax_m3/common/ops/sparse_attn.py +++ b/vllm/models/minimax_m3/common/ops/sparse_attn.py @@ -304,7 +304,7 @@ def _gqa_sparse_decode_kernel( # attention range instead of letting padded rows produce negative lengths. kv_len = tl.maximum(query_pos + 1, 0) - # Valid block count from seq_len (no sentinel): min(topk, cdiv(kv_len, blk)). + # Upper bound on top-k entries from causal length; entries may be sentinels. idx_base = t_ptr + pid_kh * stride_th + pid_b * stride_tn num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K real_topk = tl.minimum(max_topk, num_blocks) @@ -334,10 +334,15 @@ def _gqa_sparse_decode_kernel( for _ in tl.range(chunk_start_topk, chunk_end_topk): blk = tl.load(cur_idx_ptr).to(tl.int32) cur_idx_ptr = cur_idx_ptr + stride_tk - c = blk * BLOCK_SIZE_K - page = tl.load(bt_row + blk).to(tl.int64) + # topk_idx uses -1 for unused entries. Select a safe logical block + # before dereferencing the block table; invalid and future blocks then + # contribute a neutral online-softmax update without control flow. + valid_blk = (blk >= 0) & (blk < num_blocks) + safe_blk = tl.where(valid_blk, blk, 0) + c = safe_blk * BLOCK_SIZE_K + page = tl.load(bt_row + safe_blk).to(tl.int64) pos = c + off_n - pos_mask = pos < kv_len + pos_mask = (pos < kv_len) & valid_blk k = tl.load( kv_cache_ptr + page * stride_kv_blk @@ -364,9 +369,11 @@ def _gqa_sparse_decode_kernel( qk += tl.where(pos_mask[None, :], 0, float("-inf")) qk += tl.dot(q, k) * sm_scale_log2e m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) - p = tl.exp2(qk - m_ij[:, None]) + m_safe = tl.where(valid_blk, m_ij, 0.0) + p = tl.where(pos_mask[None, :], tl.exp2(qk - m_safe[:, None]), 0.0) l_ij = tl.sum(p, axis=1) - acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None] + alpha = tl.where(valid_blk, tl.exp2(m_i - m_safe), 1.0) + acc_o = acc_o * alpha[:, None] v = tl.load( kv_cache_ptr + page * stride_kv_blk @@ -390,8 +397,9 @@ def _gqa_sparse_decode_kernel( ) v = (v * v_scale[:, None]).to(q.dtype) acc_o += tl.dot(p.to(v.dtype), v) - m_i = m_ij - lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + lse_ij = m_safe + tl.log2(tl.exp2(lse_i - m_safe) + l_ij) + m_i = tl.where(valid_blk, m_safe, m_i) + lse_i = tl.where(valid_blk, lse_ij, lse_i) if USE_PDL: tl.extra.cuda.gdc_launch_dependents() From d8fa463e2ba751759a3534dd3ce11af92f1eed83 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Thu, 8 Oct 2026 14:33:48 -0400 Subject: [PATCH 117/624] [Bugfix][Prefix caching] Fix inconsistent hybrid Mamba prefix cache boundaries (#60533) Signed-off-by: Wei Zhao Signed-off-by: zjy0516 Co-authored-by: Wei Zhao Co-authored-by: Codex Co-authored-by: zjy0516 Co-authored-by: Claude Opus 5.5 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../models/kimi_k3/test_amd_kda_checkpoint.py | 7 +- tests/models/kimi_k3/test_kda_metadata.py | 60 ++++- tests/models/kimi_k3/test_prefix_cache.py | 40 +++ .../v1/attention/test_gdn_metadata_builder.py | 8 +- .../test_mamba2_prefill_checkpoints.py | 9 +- .../test_partial_prefix_cache_hits.py | 253 ++++++++++++++---- tests/v1/core/test_prefix_caching.py | 38 +-- vllm/config/cache.py | 11 + .../model_executor/layers/mamba/checkpoint.py | 11 +- vllm/v1/core/block_pool.py | 7 +- vllm/v1/core/kv_cache_coordinator.py | 69 ++--- vllm/v1/core/kv_cache_utils.py | 108 ++++++-- vllm/v1/core/sched/scheduler.py | 35 ++- vllm/v1/core/single_type_kv_cache_manager.py | 60 +++-- vllm/v1/engine/core.py | 14 + vllm/v1/kv_cache_interface.py | 9 +- vllm/v1/worker/gpu_worker.py | 9 +- 17 files changed, 547 insertions(+), 201 deletions(-) diff --git a/tests/models/kimi_k3/test_amd_kda_checkpoint.py b/tests/models/kimi_k3/test_amd_kda_checkpoint.py index 0a622bfa0450..6f3129a224e5 100644 --- a/tests/models/kimi_k3/test_amd_kda_checkpoint.py +++ b/tests/models/kimi_k3/test_amd_kda_checkpoint.py @@ -75,7 +75,7 @@ def _builder( vllm_config.cache_config.prefix_match_unit = prefix_match_unit if num_spec: vllm_config.speculative_config = _StubSpeculativeConfig(num_spec) - return KimiK3ROCmKDAMetadataBuilder( + builder = KimiK3ROCmKDAMetadataBuilder( kv_cache_spec=MambaSpec( block_size=block_size, shapes=((16, 64),), @@ -91,6 +91,11 @@ def _builder( vllm_config=vllm_config, device=torch.device("cuda"), ) + vllm_config.cache_config.hash_block_size = prefix_match_unit or block_size + vllm_config.cache_config.cache_hit_alignment_tokens = ( + prefix_match_unit or block_size + ) + return builder def _build(batch: BatchSpec, block_size: int = BLOCK_SIZE, **kw): diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index 74414279552c..63024634d305 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -100,6 +100,8 @@ def _make_builder( prefix_match_unit: int | None = None, use_eagle: bool = False, disable_eagle_block_drop: bool = False, + cache_hit_alignment_tokens: int | None = None, + hash_block_size: int | None = None, ) -> AttentionMetadataBuilder: vllm_config = create_vllm_config( model_name="Qwen/Qwen3.5-0.8B", @@ -122,6 +124,11 @@ def _make_builder( vllm_config.cache_config.use_replayssm = use_recoverssm vllm_config.cache_config.use_kda_recoverssm = use_recoverssm vllm_config.cache_config.prefix_match_unit = prefix_match_unit + hash_block_size = hash_block_size or prefix_match_unit or mamba_block_size + vllm_config.cache_config.hash_block_size = hash_block_size + vllm_config.cache_config.cache_hit_alignment_tokens = ( + cache_hit_alignment_tokens or hash_block_size + ) builder = builder_cls( kv_cache_spec=MambaSpec( block_size=mamba_block_size, @@ -220,6 +227,8 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch): vllm_config=builder_config, device=DEVICE, ) + builder_config.cache_config.hash_block_size = BLOCK_SIZE + builder_config.cache_config.cache_hit_alignment_tokens = BLOCK_SIZE # An all-prefill speculative batch used to leave an all-false spec mask # alive, then access active_non_spec_mask_cpu before it was initialized. @@ -287,18 +296,45 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_internal_checkpoint_metadata_targets_last_aligned_boundary(): +@pytest.mark.parametrize( + "batch,prefix_match_unit,hit_alignment,expected_offsets,expected_states", + [ + ( + BatchSpec(seq_lens=[50, 32], query_lens=[50, 16]), + None, + 16, + [48, 0], + [2, NULL_BLOCK_ID], + ), + ( + BatchSpec(seq_lens=[50, 32], query_lens=[50, 16]), + 16, + 16, + [48, 0], + [2, NULL_BLOCK_ID], + ), + (BatchSpec(seq_lens=[530], query_lens=[530]), None, 128, [512], [32]), + (BatchSpec(seq_lens=[530], query_lens=[50]), None, 128, [32], [32]), + (BatchSpec(seq_lens=[530], query_lens=[530]), 16, 16, [528], [32]), + (BatchSpec(seq_lens=[530], query_lens=[50]), 16, 16, [48], [32]), + ], +) +def test_internal_checkpoint_metadata_targets_last_aligned_boundary( + batch, prefix_match_unit, hit_alignment, expected_offsets, expected_states +): device = torch.device("cuda") - batch = BatchSpec(seq_lens=[50, 32], query_lens=[50, 16]) common_attn_metadata = create_common_attn_metadata( batch, BLOCK_SIZE, device, arange_block_indices=True - ).replace(is_prefilling=torch.tensor([True, True])) + ).replace(is_prefilling=torch.tensor([True] * len(batch.seq_lens))) builder = _make_builder( KimiK3KDAMetadataBuilder, num_speculative_tokens=0, full_cuda_graph=False, mamba_cache_mode="align", num_prefill_checkpoint_blocks=1, + prefix_match_unit=prefix_match_unit, + cache_hit_alignment_tokens=hit_alignment, + hash_block_size=16, device=device, ) assert isinstance(builder, KimiK3KDAMetadataBuilder) @@ -313,26 +349,33 @@ def test_internal_checkpoint_metadata_targets_last_aligned_boundary(): assert actual.checkpoint is not None torch.testing.assert_close( actual.checkpoint.state_indices, - torch.tensor([2, NULL_BLOCK_ID], dtype=torch.int32, device=device), + torch.tensor(expected_states, dtype=torch.int32, device=device), ) torch.testing.assert_close( actual.checkpoint.checkpoint_offsets, - torch.tensor([48, 0], dtype=torch.int32, device=device), + torch.tensor(expected_offsets, dtype=torch.int32, device=device), ) @pytest.mark.parametrize( - ("disable_eagle_block_drop", "prefix_match_unit", "expected_offset"), - [(False, 16, 80), (True, 16, 96), (False, 8, None)], + ("seq_len", "disable_eagle_block_drop", "prefix_match_unit", "expected_offset"), + [ + (100, False, 16, 80), + (100, True, 16, 96), + (100, False, 8, None), + (96, False, 16, 80), + (96, True, 16, 80), + ], ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_spec_internal_checkpoint_metadata_targets_replay_boundary( + seq_len: int, disable_eagle_block_drop: bool, prefix_match_unit: int, expected_offset: int | None, ) -> None: device = torch.device("cuda") - batch = BatchSpec(seq_lens=[100], query_lens=[100]) + batch = BatchSpec(seq_lens=[seq_len], query_lens=[seq_len]) common_attn_metadata = create_common_attn_metadata( batch, BLOCK_SIZE, device, arange_block_indices=True ) @@ -390,6 +433,7 @@ def test_internal_checkpoint_metadata_skips_unaligned_offset(): full_cuda_graph=False, mamba_cache_mode="align", num_prefill_checkpoint_blocks=1, + prefix_match_unit=16, device=device, ) assert isinstance(builder, KimiK3KDAMetadataBuilder) diff --git a/tests/models/kimi_k3/test_prefix_cache.py b/tests/models/kimi_k3/test_prefix_cache.py index 3525550fe970..20e5d6670e5c 100644 --- a/tests/models/kimi_k3/test_prefix_cache.py +++ b/tests/models/kimi_k3/test_prefix_cache.py @@ -8,6 +8,7 @@ from collections.abc import Iterator from dataclasses import dataclass from math import lcm +from uuid import uuid4 import pytest import requests @@ -353,6 +354,45 @@ def _format(turns: list[Turn]) -> str: return "\n".join(rows) +@pytest.mark.parametrize( + "prompt_len", + [ + # Blocks=6144, PMU=128. P=7040 is PMU-aligned but not block-aligned: + # ensure lookup reads the cached attention proof at 7040 despite + # the final hit limit of 7039, then applies EAGLE's 128-token drop. + pytest.param(7040, id="eagle-proof-scan"), + # Blocks=6144, PMU=128. P=7296 is an exact PMU multiple: + # ensure the Mamba checkpoint is 7168, not 7040 from applying both + # the P - 1 cap and EAGLE's 128-token drop. + pytest.param(7296, id="eagle-double-cap"), + ], +) +def test_dspark_exact_resend_reuses_prompt_checkpoint(prompt_len: int) -> None: + """Resends reuse the Mamba checkpoint and its companion EAGLE attention proof.""" + mode = Mode(128, True) + salt = str(uuid4()) + rng = random.Random(0) + prompt = [rng.randint(1000, 150000) for _ in range(prompt_len)] + deployment = DEPLOYMENTS["plain"] + with _serve(deployment, mode) as servers: + url = servers[0].url_root + cold = _complete(url, prompt, salt, max_tokens=1) + resend = _complete(url, prompt, salt, max_tokens=1) + recompute = _complete(url, prompt, f"{salt}-recompute", max_tokens=1) + + # EAGLE's dropped hash unit already leaves tokens to recompute on resend. + expected = (prompt_len // 128 - 1) * 128 + actual = _cached(resend) + assert _cached(cold) == _cached(recompute) == 0 + assert actual == expected, f"{prompt_len=}: {actual=}, {expected=}" + check_logprobs_close( + outputs_0_lst=[_tokens_text_logprobs(recompute)], + outputs_1_lst=[_tokens_text_logprobs(resend)], + name_0="recompute", + name_1="resend", + ) + + @pytest.mark.parametrize( "name, mode", [_case(name, mode) for name in DEPLOYMENTS for mode in MODES] ) diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index cd01fba49e71..4b52eb9c0950 100644 --- a/tests/v1/attention/test_gdn_metadata_builder.py +++ b/tests/v1/attention/test_gdn_metadata_builder.py @@ -385,7 +385,7 @@ def test_cudagraph_capture_batch_stays_decode_only(): def _create_checkpoint_builder_and_batch( - num_spec: int, prefix_match_unit: int = 16 + num_spec: int, prefix_match_unit: int | None = 16 ) -> tuple[GDNAttentionMetadataBuilder, CommonAttentionMetadata, dict]: vllm_config = create_vllm_config( model_name="Qwen/Qwen3.5-0.8B", @@ -413,6 +413,8 @@ def _create_checkpoint_builder_and_batch( vllm_config=vllm_config, device=DEVICE, ) + vllm_config.cache_config.hash_block_size = prefix_match_unit or 64 + vllm_config.cache_config.cache_hit_alignment_tokens = prefix_match_unit or 64 batch = BatchSpec( seq_lens=[65, 100, 100], query_lens=[3 if num_spec else 1, 100, 99] ) @@ -431,7 +433,9 @@ def _create_checkpoint_builder_and_batch( @pytest.mark.parametrize("num_spec", [0, 3]) -@pytest.mark.parametrize("prefix_match_unit, expected_offset", [(16, 96), (8, 96)]) +@pytest.mark.parametrize( + "prefix_match_unit, expected_offset", [(None, 64), (16, 96), (8, 96)] +) def test_checkpoint_metadata_preserves_non_spec_order( num_spec, prefix_match_unit, expected_offset ): diff --git a/tests/v1/attention/test_mamba2_prefill_checkpoints.py b/tests/v1/attention/test_mamba2_prefill_checkpoints.py index e5276786e23c..a3968b5a0bc8 100644 --- a/tests/v1/attention/test_mamba2_prefill_checkpoints.py +++ b/tests/v1/attention/test_mamba2_prefill_checkpoints.py @@ -77,7 +77,9 @@ def _create_mamba2_builder( num_prefill_checkpoint_blocks: int = 1, prefill_checkpoint_alignment: int | None = 1, ) -> Mamba2AttentionMetadataBuilder: - vllm_config = create_vllm_config(block_size=MAMBA_BLOCK_SIZE) + vllm_config = create_vllm_config( + model_name="Qwen/Qwen3.5-0.8B", block_size=MAMBA_BLOCK_SIZE + ) vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode # get_mamba_chunk_size() reads `mamba_chunk_size` then `chunk_size`. vllm_config.model_config.hf_text_config.mamba_chunk_size = CHUNK_SIZE @@ -89,12 +91,15 @@ def _create_mamba2_builder( num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks, prefill_checkpoint_alignment=prefill_checkpoint_alignment, ) - return Mamba2AttentionMetadataBuilder( + builder = Mamba2AttentionMetadataBuilder( kv_cache_spec=spec, layer_names=["layer.0"], vllm_config=vllm_config, device=DEVICE, ) + vllm_config.cache_config.hash_block_size = MAMBA_BLOCK_SIZE + vllm_config.cache_config.cache_hit_alignment_tokens = MAMBA_BLOCK_SIZE + return builder def _build(builder, seq_lens, query_lens): diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 76beffce452a..51317deb8d01 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -22,6 +22,7 @@ get_block_hash, get_group_id, init_none_hash, + resolve_cache_hit_alignment_tokens, ) from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.kv_cache_interface import ( @@ -148,29 +149,84 @@ def make_full_mamba_manager( ) -def test_dcp_fine_hit_retention_uses_hash_alignment_without_eagle(): - """DCP must not discard a reusable Mamba state at hash alignment.""" - hash_block_size = 2 +@pytest.mark.parametrize( + "dcp_world_size,hash_block_size,block_size,prompt_tokens,checkpoint_blocks", + [ + (4, 2, 2, 7, 0), + (8, 16, 16, 512, 1), + (8, 16, 16, 513, 1), + *[(8, 128, 896, p, 1) for p in (127, 128, 129, 7040, 7168, 7296, 7449)], + ], +) +@pytest.mark.parametrize("use_eagle", [False, True]) +@pytest.mark.parametrize("extension_tokens", [0, 32]) +def test_dcp_fine_hit_retention_uses_hash_alignment( + dcp_world_size, + hash_block_size, + block_size, + prompt_tokens, + checkpoint_blocks, + use_eagle, + extension_tokens, +): + """Resends and extensions reuse the retained hash-aligned checkpoint.""" manager = make_full_mamba_manager( - dcp_world_size=4, + dcp_world_size=dcp_world_size, hash_block_size=hash_block_size, - full_block_size=hash_block_size, - mamba_block_size=hash_block_size, - num_blocks=64, + full_block_size=block_size, + mamba_block_size=block_size, + num_blocks=128, + use_eagle=use_eagle, + num_speculative_blocks=4 if use_eagle else 0, + num_prefill_checkpoint_blocks=checkpoint_blocks, ) manager.coordinator.retention_interval = 0 - - token_ids = list(range(7)) - producer = make_request("producer", token_ids, hash_block_size, sha256) - computed_blocks, num_computed, _ = manager.get_computed_blocks(producer) - assert num_computed == 0 - assert manager.allocate_slots(producer, 6, 0, computed_blocks) is not None + assert manager.coordinator.enable_partial_hash_hits + scheduler = SimpleNamespace( + block_size=block_size * dcp_world_size, + cache_config=SimpleNamespace(block_size=block_size), + hash_block_size=hash_block_size, + max_num_scheduled_tokens=8192, + scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), + use_eagle_block_drop=use_eagle, + mamba_partial_cache_hit=True, + mamba_shared_prefix_checkpoint=False, + mamba_has_prefill_checkpoint_blocks=bool(checkpoint_blocks), + mamba_prefill_checkpoint_alignment=16, + ) + producer = make_request( + "producer", list(range(prompt_tokens)), hash_block_size, sha256 + ) + while producer.num_computed_tokens < prompt_tokens: + scheduled = Scheduler._mamba_block_aligned_split( + scheduler, producer, prompt_tokens - producer.num_computed_tokens + ) + assert scheduled > 0 + manager.new_step_starts() + assert manager.allocate_slots(producer, scheduled) is not None + producer.num_computed_tokens += scheduled manager.free(producer) manager.new_step_starts() - consumer = make_request("consumer", token_ids, hash_block_size, sha256) - _, num_computed, _ = manager.get_computed_blocks(consumer) - assert num_computed == 6 + replay = make_request( + "replay", list(range(prompt_tokens + extension_tokens)), hash_block_size, sha256 + ) + _, hit, _ = manager.get_computed_blocks(replay) + expected = ( + max((prompt_tokens // hash_block_size - 1) * hash_block_size, 0) + if use_eagle + else ((prompt_tokens - 1) // hash_block_size * hash_block_size) + ) + assert hit == expected + cached_states = { + end + for end in range(hash_block_size, prompt_tokens + 1, hash_block_size) + if manager.block_pool.get_cached_block( + producer.block_hashes[end // hash_block_size - 1], [1] + ) + is not None + } + assert cached_states == ({expected} if expected else set()) @pytest.mark.parametrize("dcp_world_size", [1, 4]) @@ -339,7 +395,7 @@ def test_hybrid_mamba_align_partial_hash_hit(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert num_computed == 0 blocks = manager.allocate_slots(req0, 6, num_computed, computed_blocks) @@ -488,7 +544,7 @@ def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert num_computed == 0 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -538,7 +594,7 @@ def test_partial_hit_then_internal_checkpoint_uses_distinct_mamba_blocks(): num_prefill_checkpoint_blocks=1, ) - owner = make_request("owner", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + owner = make_request("owner", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(owner) assert manager.allocate_slots(owner, 6, num_computed, computed_blocks) is not None manager.free(owner) @@ -619,44 +675,70 @@ def test_internal_checkpoint_uses_partial_hash_lifecycle(): assert num_computed == 112 -@pytest.mark.parametrize( - "retention_interval,transient_published", - [(None, True), (0, False)], -) +@pytest.mark.parametrize("retention_interval", [None, 0, 32, 64, 256]) +@pytest.mark.parametrize("hash_block_size", [16, 32]) +@pytest.mark.parametrize("use_eagle", [False, True]) def test_internal_checkpoint_publication_respects_retention( - retention_interval, transient_published + retention_interval, hash_block_size, use_eagle ): - """With retention_interval=0, a mid-prompt internal checkpoint is - request-local and must not enter the prefix cache; the prompt-end - checkpoint keeps the existing publication behavior. Dense retention - (None) is unchanged. - """ - hash_block_size = 16 + """Publish only retained block boundaries and the final prompt checkpoint.""" manager = make_full_mamba_manager( dcp_world_size=1, hash_block_size=hash_block_size, full_block_size=hash_block_size, mamba_block_size=32, num_prefill_checkpoint_blocks=1, + use_eagle=use_eagle, ) manager.coordinator.retention_interval = retention_interval request = make_request("producer", list(range(240)), hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(request) - # Mid-prompt chunk: exports a transient checkpoint at 112. + # The EAGLE adjustment can make a chunk checkpoint block-aligned or partial. assert manager.allocate_slots(request, 128, num_computed, computed_blocks) - transient_hash = request.block_hashes[112 // hash_block_size - 1] + checkpoint = (128 if use_eagle else 127) // hash_block_size * hash_block_size - ( + hash_block_size if use_eagle else 0 + ) + retained_checkpoint = checkpoint % 32 == 0 and ( + retention_interval is None + or (retention_interval > 0 and checkpoint % retention_interval == 0) + ) + transient_hash = request.block_hashes[checkpoint // hash_block_size - 1] transient_hit = manager.block_pool.get_cached_block(transient_hash, [1]) - assert (transient_hit is not None) == transient_published + assert (transient_hit is not None) == retained_checkpoint + offloads = drain_boundary_state_offloads(manager).get(request.request_id, []) + assert any(position == checkpoint for _, _, position in offloads) == ( + retained_checkpoint + ) + assert all( + position % 32 == 0 + and ( + retention_interval is None + or (retention_interval > 0 and position % retention_interval == 0) + ) + for _, _, position in offloads + ) + if retention_interval in (None, 32, 64): + retained_hash = request.block_hashes[128 // hash_block_size - 1] + assert manager.block_pool.get_cached_block(retained_hash, [1]) is not None request.num_computed_tokens = 128 manager.new_step_starts() - # Final chunk: the prompt-end checkpoint at 224 stays published. + # The final checkpoint is reusable regardless of the retention interval. assert manager.allocate_slots(request, 112) is not None - end_hash = request.block_hashes[224 // hash_block_size - 1] + prompt_checkpoint = ( + 240 if use_eagle else 239 + ) // hash_block_size * hash_block_size - (hash_block_size if use_eagle else 0) + end_hash = request.block_hashes[prompt_checkpoint // hash_block_size - 1] assert manager.block_pool.get_cached_block(end_hash, [1]) is not None + assert any( + position == prompt_checkpoint + for _, _, position in drain_boundary_state_offloads(manager).get( + request.request_id, [] + ) + ) def test_transient_checkpoint_evicts_retained_boundary_hash(): @@ -802,9 +884,9 @@ def test_hash_aligned_query_end_uses_regular_partial_tail(): mamba_block_size=mamba_block_size, num_prefill_checkpoint_blocks=1, ) - request = make_request("producer", list(range(14)), hash_block_size, sha256) + request = make_request("producer", list(range(15)), hash_block_size, sha256) - new_blocks = manager.allocate_slots(request, request.num_tokens) + new_blocks = manager.allocate_slots(request, 14) assert new_blocks is not None mamba_manager = manager.coordinator.single_type_managers[1] @@ -928,7 +1010,7 @@ def test_boundary_state_offloads_returns_cow_target(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -1004,7 +1086,7 @@ def test_finished_partial_tail_uses_table_source_once(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None source_block_id = manager.get_blocks("0").get_block_ids()[1][1] @@ -1024,7 +1106,7 @@ def test_finished_partial_tail_uses_table_source_once(): enable_caching=True, hash_block_size=hash_block_size, ) - req1 = make_request("1", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req1 = make_request("1", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = advanced_manager.get_computed_blocks(req1) assert ( advanced_manager.allocate_slots(req1, 6, num_computed, computed_blocks) @@ -1212,7 +1294,7 @@ def test_truncate_computed_blocks_preserves_sparse_prefix_positions(): enable_caching=True, hash_block_size=hash_block_size, ) - producer = make_request("producer", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + producer = make_request("producer", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) blocks, num_computed, _ = manager.get_computed_blocks(producer) assert manager.allocate_slots(producer, 6, num_computed, blocks) is not None manager.free(producer) @@ -1267,7 +1349,7 @@ def test_truncate_computed_blocks_allows_short_mamba_group_only(): enable_caching=True, hash_block_size=hash_block_size, ) - producer = make_request("producer", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + producer = make_request("producer", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) blocks, num_computed, _ = manager.get_computed_blocks(producer) assert manager.allocate_slots(producer, 6, num_computed, blocks) is not None manager.free(producer) @@ -1331,7 +1413,7 @@ def test_hybrid_mamba_partial_tail_owner_continue_preserves_later_hit(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert num_computed == 0 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -1411,7 +1493,7 @@ def test_hybrid_mamba_moved_partial_entry_defers_same_step_hit(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert num_computed == 0 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -1470,7 +1552,7 @@ def test_hybrid_full_attention_partial_hash_hit_uses_cow(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert num_computed == 0 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -1544,7 +1626,7 @@ def test_hybrid_partial_hit_cow_target_starts_uncached(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert num_computed == 0 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -1700,7 +1782,7 @@ def test_cow_retained_blocks_returned_for_release(): enable_caching=True, hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -2133,6 +2215,72 @@ def test_hybrid_sliding_window_group_disables_partial_hash_hits(): assert len(computed_blocks.blocks[0]) * hash_block_size == num_computed +@pytest.mark.parametrize("with_sliding_window", [False, True]) +def test_resolved_cache_hit_alignment_matches_coordinator(with_sliding_window): + """Workers place Mamba checkpoints on the engine-resolved alignment, so it + must be where the coordinator lands hits, including when a group without + fine-grained lookups disables partial hits.""" + hash_block_size = 2 + mamba_block_size = 4 * hash_block_size + kv_cache_groups = [ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=mamba_block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ] + if with_sliding_window: + kv_cache_groups.append( + KVCacheGroupSpec( + ["swa"], + SlidingWindowSpec( + block_size=2 * hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=2 * hash_block_size, + ), + ) + ) + kv_cache_config = KVCacheConfig( + num_blocks=64, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups + ) + coordinator = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ).coordinator + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(enable_prefix_caching=True), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + ) + + assert coordinator.enable_partial_hash_hits != with_sliding_window + assert ( + resolve_cache_hit_alignment_tokens( + kv_cache_config, + vllm_config, + coordinator.scheduler_block_size, + hash_block_size, + ) + == coordinator._cache_hit_alignment_tokens + ) + + def test_opted_out_scratch_group_keeps_partial_hash_hits(): hash_block_size = 2 mamba_block_size = 2 * hash_block_size @@ -2179,7 +2327,7 @@ def test_opted_out_scratch_group_keeps_partial_hash_hits(): assert manager.coordinator.enable_partial_hash_hits - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None manager.free(req0) @@ -2240,7 +2388,7 @@ def test_kpool_tail_supports_128_token_partial_hash_hits(): assert manager.coordinator.enable_partial_hash_hits shared_prefix = [10] * (12 * hash_block_size) - req0 = make_request("0", shared_prefix, hash_block_size, sha256) + req0 = make_request("0", shared_prefix + [11], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert ( manager.allocate_slots(req0, len(shared_prefix), num_computed, computed_blocks) @@ -2266,7 +2414,7 @@ def test_hybrid_partial_hash_hit_uses_cow_under_dcp(dcp_world_size: int): ) assert manager.coordinator.enable_partial_hash_hits - req0 = make_request("dcp-owner", [0, 0, 1, 1, 2, 2], 2, sha256) + req0 = make_request("dcp-owner", [0, 0, 1, 1, 2, 2, 3], 2, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None manager.free(req0) @@ -2318,7 +2466,7 @@ def test_dcp_partial_hit_resumes_on_replicated_mamba_snapshot( assert manager.coordinator.single_type_managers[1].block_size == block_size prefix = list(range(12)) - req0 = make_request("snapshot-owner", prefix, block_size, sha256) + req0 = make_request("snapshot-owner", prefix + [12], block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 12, num_computed, computed_blocks) is not None manager.free(req0) @@ -2355,7 +2503,7 @@ def test_dcp_joint_hit_is_bounded_by_replicated_mamba_snapshots(): mamba_block_size=block_size, ) prefix = list(range(12)) - req0 = make_request("joint-owner", prefix, block_size, sha256) + req0 = make_request("joint-owner", prefix + [12], block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 8, num_computed, computed_blocks) is not None manager.new_step_starts() @@ -2702,6 +2850,7 @@ def test_checkpoint_reservation_matches_worker( mamba_block_size=mamba_block_size, checkpoint_alignment=16, drop_eagle_block=use_eagle, + cache_hit_alignment_tokens=mamba_manager.cache_hit_alignment_tokens, ) reserved = mamba_manager._checkpoints.get(request.request_id) if reserved is not None: diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 427201a7ceef..a89fd085f1fd 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -5047,19 +5047,10 @@ def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary(): def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt(): - """An identical resend and a longer sibling resume at DIFFERENT positions. - - How far a lookup matches depends on who is asking. A resend of the same - prompt caps its lookup at ``num_tokens - 1`` (the last token is recomputed - for logits), while a sibling whose prompt merely starts with this one caps - above the prompt. The two coincide unless the prompt length is an exact - multiple of the alignment -- there they differ by one alignment unit, and - under the EAGLE drop BOTH are reachable. - - Retaining only the higher one leaves the resend with every retained state - above every candidate its lookup can produce, and the reconciled hit - collapses to 0 -- the same zero-hit failure sparse retention already fixes - at unaligned prompt lengths. + """EAGLE resend and extension reuse the same retained checkpoint. + + Attention may inspect the prompt's final block before dropping it, so an + aligned resend must not lose another block to the P - 1 lookup limit. """ block_size = 32 num_spec = 3 @@ -5097,16 +5088,12 @@ def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt(): use_eagle=True, ) - # 128 tokens, an exact multiple of the 32-token alignment. A longer sibling - # matches 128 and drops to 96; this prompt's own resend caps at 127, matches - # 96 and drops to 64. Both states must survive retention. + # Both requests prove 128 attention tokens, drop one block, and reuse 96. token_ids = [i for i in range(4) for _ in range(block_size)] req0 = make_request("0", token_ids, block_size, sha256) computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0) assert num_computed_tokens == 0 - # Prefill in block-aligned chunks the way the align-mode scheduler does: a - # state only materializes as a chunk's running-state block, so a - # single-shot prefill could not retain the lower one. + # Materialize the checkpoint through normal block-aligned prefill chunks. for chunk_end in (32, 64, 96, 128): blocks = manager.allocate_slots( req0, @@ -5118,8 +5105,7 @@ def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt(): assert blocks is not None req0.num_computed_tokens = chunk_end - # Block ``i`` ends at token ``(i + 1) * 32``, so positions 64 and 96 are - # mamba blocks 1 and 2. + # Both materialized replay checkpoints survive; resend and extension use 96. pool = manager.block_pool expected_mamba_cached = {1, 2} for i in range(4): @@ -5130,15 +5116,13 @@ def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt(): assert cached is None, f"mamba hash {i} should not be cached" manager.free(req0) - # The identical resend: full attention matches blocks 0-2 (96 tokens, capped - # by num_tokens - 1) and the EAGLE drop caps the candidate at 64. Without - # the lower state retained the reconciled hit would be 0. + # The identical resend must reuse the retained checkpoint, not miss to zero. req1 = make_request("1", token_ids, block_size, sha256) computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1) - assert num_computed_tokens == 2 * block_size - assert [len(blocks) for blocks in computed_blocks.blocks] == [2, 2] + assert num_computed_tokens == 3 * block_size + assert [len(blocks) for blocks in computed_blocks.blocks] == [3, 3] - # The longer sibling resumes one alignment unit higher, off the same prompt. + # The longer sibling reuses that same checkpoint. longer = make_request("2", token_ids + [9] * block_size, block_size, sha256) computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(longer) assert num_computed_tokens == 3 * block_size diff --git a/vllm/config/cache.py b/vllm/config/cache.py index e6436f98f9cb..5eb26726c718 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -230,6 +230,15 @@ def device_memory_utilization(self, value: float) -> None: for hybrid models where requests occupy multiple KV cache groups.""" kv_cache_max_concurrency: float | None = field(default=None, init=False) """Per-DP-engine maximum concurrency at max_model_len tokens.""" + # The two differ when hashes are finer than the scheduler block but partial + # hash hits are off (e.g. a sliding-window group cannot look up sub-block + # hits), so hits only land on the scheduler block. + # TODO(zjy0516): Merge into one value + hash_block_size: int | None = field(default=None, init=False) + """Tokens per prefix-cache block hash.""" + cache_hit_alignment_tokens: int | None = field(default=None, init=False) + """Token granularity at which prefix-cache hits land: `hash_block_size` + with partial hash hits, otherwise the scheduler block size.""" kv_sharing_fast_prefill: bool = False """In some KV sharing setups, e.g. YOCO (https://arxiv.org/abs/2405.05254), @@ -298,6 +307,8 @@ def compute_hash(self) -> str: "effective_attention_block_size", "kv_cache_size_tokens", "kv_cache_max_concurrency", + "hash_block_size", + "cache_hit_alignment_tokens", # Feature toggles not impacting compiled graph shape "kv_sharing_fast_prefill", "swa_bounded_replay", diff --git a/vllm/model_executor/layers/mamba/checkpoint.py b/vllm/model_executor/layers/mamba/checkpoint.py index 45888113e152..2f3122f87058 100644 --- a/vllm/model_executor/layers/mamba/checkpoint.py +++ b/vllm/model_executor/layers/mamba/checkpoint.py @@ -25,6 +25,7 @@ def compute_mamba_prefill_checkpoints( mamba_block_size: int, checkpoint_alignment: int | None, drop_eagle_block: bool, + cache_hit_alignment_tokens: int, ) -> tuple[list[int], list[int]]: """Per-row internal prefill checkpoint offsets and cache block columns. @@ -43,7 +44,7 @@ def compute_mamba_prefill_checkpoints( for seq_len, query_len in zip(seq_lens, query_lens): query_start = seq_len - query_len position = get_mamba_prefill_checkpoint_position( - seq_len, hash_block_size, drop_eagle_block=drop_eagle_block + seq_len, cache_hit_alignment_tokens, drop_eagle_block=drop_eagle_block ) valid = is_mamba_prefill_checkpoint_valid( query_start=query_start, @@ -106,12 +107,17 @@ def build( return None if self.kv_cache_spec.num_prefill_checkpoint_blocks == 0: return None + cache_config = self.vllm_config.cache_config + hash_block_size = cache_config.hash_block_size + cache_hit_alignment_tokens = cache_config.cache_hit_alignment_tokens + if hash_block_size is None or cache_hit_alignment_tokens is None: + # Unresolved only for the minimal KV cache built for profiling. + return None assert m.seq_lens_cpu_upper_bound is not None all_query_lens = m.query_start_loc_cpu.diff().tolist() query_lens = [all_query_lens[row] for row in request_rows] seq_lens = m.seq_lens_cpu_upper_bound.tolist() block_size = self.kv_cache_spec.block_size - hash_block_size = self.vllm_config.cache_config.prefix_match_unit or block_size speculative_config = self.vllm_config.speculative_config drop_eagle_block = ( speculative_config is not None and speculative_config.use_eagle_block_drop() @@ -123,6 +129,7 @@ def build( mamba_block_size=block_size, checkpoint_alignment=self.kv_cache_spec.prefill_checkpoint_alignment, drop_eagle_block=drop_eagle_block, + cache_hit_alignment_tokens=cache_hit_alignment_tokens, ) if not any(checkpoint_offsets): return None diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index 75bdb40c24be..3b0102fe84aa 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -479,8 +479,8 @@ def cache_partial_block( kv_cache_group_id: KV cache group that owns the partial entry. block_size: Cache block size for the owning group. The partial entry hash itself is always the prefix-chain hash at - ``num_tokens``; ``block_size`` is used to assert that the - entry is partial within the owning cache block. + ``num_tokens``; ``block_size`` must be a multiple of the hash + block size. replace_existing_hashes: Whether the block contents were replaced and all existing cache entries must be removed before the new entry is registered. @@ -494,9 +494,6 @@ def cache_partial_block( return None assert block_size % self.hash_block_size == 0 - assert replace_existing_hashes or ( - block_size > self.hash_block_size and num_tokens % block_size != 0 - ) block_hash = self._get_partial_block_hash(request, num_tokens) num_hash_blocks = num_tokens // self.hash_block_size block_hash_with_group_id = make_block_hash_with_group_id( diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index deb7f89967b2..c32173a63554 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -4,13 +4,14 @@ from collections.abc import Sequence from typing import NamedTuple -from vllm.logger import init_logger from vllm.utils.math_utils import cdiv, round_down from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.core.kv_cache_utils import ( BlockHash, KVCacheBlock, + eagle_proof_margin, + partial_hash_hits_enabled, ) from vllm.v1.core.single_type_kv_cache_manager import ( CrossAttentionManager, @@ -28,8 +29,6 @@ ) from vllm.v1.request import Request -logger = init_logger(__name__) - def _validate_prefix_cache_retention_interval( retention_interval: int | None, @@ -695,37 +694,15 @@ def __init__( # managers in every group. TP needs hashing finer than the Mamba block; # DCP accepts equality because it scales the effective full-attention # block instead. - has_partial_mamba_group = any( - isinstance(g.kv_cache_spec, MambaSpec) - and g.kv_cache_spec.mamba_cache_mode == "align" - and ( - (dcp_world_size == 1 and g.kv_cache_spec.block_size > hash_block_size) - or ( - dcp_world_size > 1 and g.kv_cache_spec.block_size >= hash_block_size - ) - ) - for g in kv_cache_config.kv_cache_groups - ) self.enable_partial_hash_hits = ( - allow_partial_hash_hits and has_partial_mamba_group + allow_partial_hash_hits + and partial_hash_hits_enabled( + kv_cache_config.kv_cache_groups, + hash_block_size, + dcp_world_size, + [type(manager) for manager in self.single_type_managers], + ) ) - if self.enable_partial_hash_hits: - unsupported_partial_hit_managers = { - type(manager).__name__ - for manager, group in zip( - self.single_type_managers, kv_cache_config.kv_cache_groups - ) - if group.kv_cache_spec.prefix_cacheable - and not manager.supports_fine_grained_hash_lookup - and manager.block_size != hash_block_size - } - if unsupported_partial_hit_managers: - self.enable_partial_hash_hits = False - logger.warning_once( - "Disabling fine-grained prefix-cache hits because these KV " - "cache managers require block-aligned lookups: %s.", - ", ".join(sorted(unsupported_partial_hit_managers)), - ) cache_hit_alignment_tokens = self._cache_hit_alignment_tokens for manager in self.single_type_managers: manager.cache_hit_alignment_tokens = cache_hit_alignment_tokens @@ -899,15 +876,15 @@ def find_longest_cache_hit( # mamba: its finder never drops (draft models have no mamba # layers), so the hit would grow past the candidate. if drop_eagle_block and not isinstance(spec, MambaSpec): - eagle_margin = ( - self.hash_block_size - if self.enable_partial_hash_hits - and manager_cls.supports_fine_grained_hash_lookup - and group_block_size > self.hash_block_size - else group_block_size + eagle_margin = eagle_proof_margin( + group_block_size, + self.hash_block_size, + self.enable_partial_hash_hits + and manager_cls.supports_fine_grained_hash_lookup, ) _max_length = min( - curr_hit_length + eagle_margin, max_cache_hit_length + curr_hit_length + eagle_margin, + len(block_hashes) * self.hash_block_size, ) hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit( block_hashes=block_hashes, @@ -983,9 +960,21 @@ def find_longest_cache_hit_per_group( for spec, group_ids, manager_cls, use_eagle in self.attention_groups: manager = self.single_type_managers[group_ids[0]] + lookup_length = max_cache_hit_length + if use_eagle and not isinstance(spec, MambaSpec): + eagle_margin = eagle_proof_margin( + manager.block_size, + self.hash_block_size, + self.enable_partial_hash_hits + and manager_cls.supports_fine_grained_hash_lookup, + ) + lookup_length = min( + max_cache_hit_length + eagle_margin, + len(block_hashes) * self.hash_block_size, + ) blocks, group_hit = manager_cls.find_longest_cache_hit( block_hashes=block_hashes, - max_length=max_cache_hit_length, + max_length=lookup_length, kv_cache_group_ids=group_ids, block_pool=manager.block_pool, kv_cache_spec=spec, diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 79ec2fc5e4ed..53a01a5cabe2 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -53,6 +53,7 @@ if TYPE_CHECKING: from vllm.v1.core.block_pool import BlockPool + from vllm.v1.core.single_type_kv_cache_manager import SingleTypeKVCacheManager # BlockHash represents the hash of a single KV-cache block used for @@ -747,6 +748,60 @@ def resolve_dcp_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> KVCache return replace(spec, block_size=block_size) +def partial_hash_hits_enabled( + kv_cache_groups: Sequence[KVCacheGroupSpec], + hash_block_size: int, + dcp_world_size: int = 1, + manager_classes: Sequence[type["SingleTypeKVCacheManager"]] | None = None, +) -> bool: + """Whether aligned Mamba states support sub-block prefix-cache hits. + + If ``manager_classes`` (one per group) is given, also require every other + prefix-cacheable group to support fine-grained lookups. + """ + if not any( + isinstance(spec, MambaSpec) + and spec.mamba_cache_mode == "align" + and ( + (dcp_world_size == 1 and spec.block_size > hash_block_size) + or (dcp_world_size > 1 and spec.block_size >= hash_block_size) + ) + for group in kv_cache_groups + for spec in iter_layer_specs(group.kv_cache_spec) + ): + return False + if manager_classes is None: + return True + unsupported_managers = { + manager_cls.__name__ + for group, manager_cls in zip(kv_cache_groups, manager_classes) + if group.kv_cache_spec.prefix_cacheable + and not manager_cls.supports_fine_grained_hash_lookup + and resolve_dcp_kv_block_size(group.kv_cache_spec, dcp_world_size) + != hash_block_size + } + if unsupported_managers: + logger.warning_once( + "Disabling fine-grained prefix-cache hits because these KV " + "cache managers require block-aligned lookups: %s.", + ", ".join(sorted(unsupported_managers)), + ) + return False + return True + + +def eagle_proof_margin( + block_size: int, hash_block_size: int, fine_grained_lookup: bool +) -> int: + """Tokens an EAGLE group matches past a cache hit before dropping them. + + Fine-grained lookups drop one hash unit; others drop one cache block. + """ + if fine_grained_lookup and block_size > hash_block_size: + return hash_block_size + return block_size + + def resolve_kv_cache_block_sizes( kv_cache_config: KVCacheConfig, vllm_config: VllmConfig, @@ -775,8 +830,7 @@ def resolve_kv_cache_block_sizes( if groups and not groups[0].kv_cache_spec.dcp_sharded: dcp = 1 bs = cache_config.block_size * dcp - # The Mamba prefill checkpoint builder reads prefix_match_unit directly, - # so a value dropped here puts its checkpoints off the scheduler's grid. + # Reject a prefix_match_unit that would otherwise be silently ignored. if ( cache_config.prefix_match_unit not in (None, bs) and len(groups) == 1 @@ -797,8 +851,8 @@ def resolve_kv_cache_block_sizes( ] scheduler_block_size = math.lcm(*group_block_sizes) - logger.info("kv cache group sizes %s", group_block_sizes) - logger.info("kv lcm block sizes %s", scheduler_block_size) + logger.info_once("kv cache group sizes %s", tuple(group_block_sizes)) + logger.info_once("kv lcm block sizes %s", scheduler_block_size) # Block hashes are only consumed by prefix caching and KV connectors # (P/D, offloading); when neither is active, keep hash_block_size equal @@ -839,18 +893,10 @@ def resolve_kv_cache_block_sizes( and isinstance(spec.tokens_per_state, int) and spec.tokens_per_state > 1 } - has_partial_mamba_group = any( - isinstance(spec, MambaSpec) - and spec.mamba_cache_mode == "align" - and ( - (dcp == 1 and block_size > hash_block_size) - or (dcp > 1 and block_size >= hash_block_size) - ) - for group, block_size in zip(groups, group_block_sizes) - for spec in iter_layer_specs(group.kv_cache_spec) - ) cache_hit_alignment = ( - hash_block_size if has_partial_mamba_group else scheduler_block_size + hash_block_size + if partial_hash_hits_enabled(groups, hash_block_size, dcp) + else scheduler_block_size ) if any(cache_hit_alignment % alignment for alignment in prefix_alignments): raise ValueError( @@ -861,6 +907,38 @@ def resolve_kv_cache_block_sizes( return scheduler_block_size, hash_block_size +def resolve_cache_hit_alignment_tokens( + kv_cache_config: KVCacheConfig, + vllm_config: VllmConfig, + scheduler_block_size: int, + hash_block_size: int, +) -> int: + """Token granularity at which prefix-cache hits land. + + Mirrors the alignment the hybrid KV cache coordinator pushes to its + managers: fine-grained hits land on ``hash_block_size``, all others on + ``scheduler_block_size``. + """ + groups = kv_cache_config.kv_cache_groups + if not vllm_config.cache_config.enable_prefix_caching or len(groups) <= 1: + return scheduler_block_size + manager_classes = [] + for group in groups: + manager_cls = KVCacheSpecRegistry.get_manager_class( + group.kv_cache_spec, group.role + ) + assert manager_cls is not None + manager_classes.append(manager_cls) + if partial_hash_hits_enabled( + groups, + hash_block_size, + vllm_config.parallel_config.decode_context_parallel_size, + manager_classes, + ): + return hash_block_size + return scheduler_block_size + + def get_request_block_hasher( hash_block_size: int, caching_hash_fn: Callable[[Any], bytes], diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 9739dbb3d516..0758c6122565 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -381,9 +381,8 @@ def __init__( or group.kv_cache_spec.num_prefill_checkpoint_blocks > 0 for group in kv_cache_config.kv_cache_groups ) - # A finer prefix_match_unit is configured: a mamba partial tail entry - # can only be registered by a step ending exactly at the prompt's last - # hash boundary, so the split adds that stop. + # Fine-grained hits reuse the prompt's hash-aligned Mamba checkpoint. + # Split there only if the backend cannot save it within the final chunk. self.mamba_partial_cache_hit = ( self.need_mamba_block_aligned_split and self.hash_block_size < self.block_size @@ -460,14 +459,19 @@ def _mamba_block_aligned_split( # The last block-aligned position whose state can be cached. With # Eagle, FullAttn prunes the last matching block, so back off one # block to avoid a Mamba cache miss. - last_cache_position = request.num_tokens - request.num_tokens % block_size - if self.use_eagle_block_drop: - last_cache_position = max(last_cache_position - block_size, 0) + ckpt_block_size = ( + self.hash_block_size if self.mamba_partial_cache_hit else self.block_size + ) + last_cache_position = get_mamba_prefill_checkpoint_position( + request.num_tokens, + block_size, + self.use_eagle_block_drop, + ) end = start + num_new_tokens checkpoint_position = get_mamba_prefill_checkpoint_position( prefill_end, - self.hash_block_size, + ckpt_block_size, drop_eagle_block=self.use_eagle_block_drop, ) use_internal_checkpoint = ( @@ -500,18 +504,14 @@ def _mamba_block_aligned_split( next_block_boundary = (start // block_size + 1) * block_size tail_boundary = ( - request.num_prompt_tokens // self.hash_block_size * self.hash_block_size + get_mamba_prefill_checkpoint_position( + request.num_prompt_tokens, + self.hash_block_size, + self.use_eagle_block_drop, + ) if self.mamba_partial_cache_hit and not use_internal_checkpoint else 0 ) - if tail_boundary and self.use_eagle_block_drop: - # Eagle matches one hash unit past the candidate and drops it, so - # nothing proves the prompt's own last hash boundary. Materialize - # the state one unit lower, where the hit can actually land. Keyed on - # the block-drop bit, not plain use_eagle: this shift exists only to - # compensate for the drop, and the Mamba manager's matching gate - # reads the same bit (the coordinator is handed use_eagle_block_drop). - tail_boundary = max(tail_boundary - self.hash_block_size, 0) junction = request.shared_prefix_boundary # Block-floored: a sub-block junction's state is not separately cacheable. block_floored = start + (junction - start) // block_size * block_size @@ -532,8 +532,7 @@ def _mamba_block_aligned_split( else 0, # Never run past the last cacheable block boundary mid-chunk. last_cache_position, - # Fine-grained hits: the prompt's partial-tail entry can only be - # registered by a chunk ending exactly at its last hash boundary. + # Stop at the prompt checkpoint unless the backend saves it internally. tail_boundary if last_cache_position < tail_boundary < request.num_prompt_tokens else 0, diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 5940a1cd0a3f..3e7ae1b7063a 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -876,7 +876,14 @@ def _cache_partial_tail_block( block are intentionally skipped. """ hash_block_size = self.block_pool.hash_block_size - boundary_tokens = request.num_prompt_tokens // hash_block_size * hash_block_size + # A resend matches at most `prompt_len - 1` tokens, but EAGLE lookups + # read the whole prompt before dropping a hash unit. + token_limit = ( + request.num_prompt_tokens + if self.use_eagle + else request.num_prompt_tokens - 1 + ) + boundary_tokens = token_limit // hash_block_size * hash_block_size if boundary_tokens == 0 or boundary_tokens > num_tokens: return if boundary_tokens % self.block_size == 0: @@ -1788,7 +1795,7 @@ def get_num_blocks_to_allocate( # is kept, so `prefill_end` skips reserving blocks for the others. checkpoint_position = get_mamba_prefill_checkpoint_position( num_tokens, - self.block_pool.hash_block_size, + self.cache_hit_alignment_tokens, self.drop_eagle_checkpoint_block, ) if not self._needs_internal_checkpoint( @@ -2073,20 +2080,29 @@ def _cache_partial_tail_block( checkpoint_position == get_mamba_prefill_checkpoint_position( request.num_prompt_tokens, - hash_block_size, + self.cache_hit_alignment_tokens, self.drop_eagle_checkpoint_block, ) ) + is_retained_checkpoint = False + if checkpoint_position % self.block_size == 0: + block_idx = checkpoint_position // self.block_size - 1 + retained = self.reachable_block_mask( + block_idx, + block_idx + 1, + self.cache_hit_alignment_tokens, + self.kv_cache_spec, + self.drop_eagle_checkpoint_block, + retention_interval=retention_interval, + ) + is_retained_checkpoint = retained is None or retained[0] if ( - retention_interval == 0 - and num_tokens < request.num_prompt_tokens - and not is_prompt_checkpoint + not is_prompt_checkpoint + and not is_retained_checkpoint and checkpoint_position != request.shared_prefix_boundary ): - # retention_interval == 0 keeps this transient checkpoint - # request-local. The slot may carry a hash from this step's - # full-block pass; that must go too, since the checkpoint - # state is about to overwrite the block. + # The transient state overwrites this slot, so neither it nor + # the full-block hash registered earlier this step is reusable. self.block_pool._maybe_evict_cached_block(checkpoint_block) return None if checkpoint_block.block_hash_num_tokens == checkpoint_position: @@ -2099,21 +2115,13 @@ def _cache_partial_tail_block( block_size=self.block_size, replace_existing_hashes=True, ) - if self.block_size == hash_block_size: + if num_tokens <= 0 or num_tokens % hash_block_size != 0: return None - if num_tokens % self.block_size == 0: - return None - if num_tokens % hash_block_size != 0: - return None - latest_prompt_hash_boundary = ( - request.num_prompt_tokens // hash_block_size - ) * hash_block_size - if self.drop_eagle_checkpoint_block: - # Eagle groups match one hash unit past the candidate and drop it, - # so register the tail one unit lower. - latest_prompt_hash_boundary = max( - latest_prompt_hash_boundary - hash_block_size, 0 - ) + latest_prompt_hash_boundary = get_mamba_prefill_checkpoint_position( + request.num_prompt_tokens, + hash_block_size, + self.drop_eagle_checkpoint_block, + ) # The junction is the other position a sibling resumes at: where one was # observed to stop, and where the scheduler already ends a chunk. Bounded # to the prompt chunk being computed -- during decode the target is the @@ -2126,7 +2134,7 @@ def _cache_partial_tail_block( ): return None - block_idx = num_tokens // self.block_size + block_idx = cdiv(num_tokens, self.block_size) - 1 blocks = self.req_to_blocks[request.request_id] if block_idx >= len(blocks): return None @@ -2141,7 +2149,7 @@ def _cache_partial_tail_block( kv_cache_group_id=self.kv_cache_group_id, block_size=self.block_size, ) - if partial_hash is not None: + if partial_hash is not None and num_tokens % self.block_size: self._partial_hit_reqs[request.request_id] = (block_idx, source_block) self.num_cached_block[request.request_id] = block_idx # Producer of this partial tail: the boundary state currently lives diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index e0a003af6254..679bbbfd2415 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -53,6 +53,7 @@ get_kv_cache_configs, get_request_block_hasher, init_none_hash, + resolve_cache_hit_alignment_tokens, resolve_kv_cache_block_sizes, update_kv_cache_capacity, ) @@ -352,6 +353,19 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: vllm_config.validate_block_size() + scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes( + scheduler_kv_cache_config, vllm_config + ) + cache_hit_alignment_tokens = resolve_cache_hit_alignment_tokens( + scheduler_kv_cache_config, + vllm_config, + scheduler_block_size, + hash_block_size, + ) + for kv_cache_config in kv_cache_configs: + kv_cache_config.hash_block_size = hash_block_size + kv_cache_config.cache_hit_alignment_tokens = cache_hit_alignment_tokens + self.model_executor.initialize_from_config(kv_cache_configs) if not envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: self.model_executor.compile_or_warm_up_model() diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index cf707611b47d..ca3f485917eb 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -1127,7 +1127,10 @@ def get_mamba_prefill_checkpoint_position( drop_eagle_block: bool, ) -> int: """Return the reusable Mamba checkpoint boundary for a prefill.""" - checkpoint_position = (num_tokens - 1) // hash_block_size * hash_block_size + # Without EAGLE, leave one token to recompute on resend. + # EAGLE's block drop already leaves tokens to recompute. + proof_limit = num_tokens if drop_eagle_block else num_tokens - 1 + checkpoint_position = proof_limit // hash_block_size * hash_block_size if drop_eagle_block: checkpoint_position -= hash_block_size return max(checkpoint_position, 0) @@ -1487,6 +1490,10 @@ class KVCacheConfig: """Resolved retention policy for local prefix-cache checkpoints.""" kv_cache_layout: str | None = None """The KV cache layout resolved by the engine core, adopted by all workers.""" + hash_block_size: int | None = None + """Tokens per prefix-cache block hash, resolved by the engine core.""" + cache_hit_alignment_tokens: int | None = None + """Token granularity of prefix-cache hits, resolved by the engine core.""" hisparse_host_num_blocks: int | None = None """Capacity of the dedicated HiSparse host-block manager, when enabled.""" diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 70427d85a103..f2dfc8a689f0 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -821,10 +821,15 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None: # so that it's available to the warmup stage. self.cache_config.num_gpu_blocks = kv_cache_config.num_blocks - # Adopt the engine core's layout; workers spawned after resolution - # (e.g. elastic EP scale-up) only see it through the config. + # Adopt the engine core's layout and prefix-cache granularity; workers + # spawned after resolution (e.g. elastic EP scale-up) only see them + # through the config. if kv_cache_config.kv_cache_layout is not None: record_kv_cache_layout(self.cache_config, kv_cache_config.kv_cache_layout) + self.cache_config.hash_block_size = kv_cache_config.hash_block_size + self.cache_config.cache_hit_alignment_tokens = ( + kv_cache_config.cache_hit_alignment_tokens + ) # Init kv cache connector here, because it requires # `kv_cache_config`. From 30c740589055c86f29859f61b13ae6620449f3fc Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Thu, 8 Oct 2026 14:39:09 -0400 Subject: [PATCH 118/624] [MoE] MXFP4 oracle: add CUTLASS W4A4 backend and BF16 activation fallback (#60079) Signed-off-by: mgoin Co-authored-by: Claude --- tests/kernels/moe/test_cutlass_moe.py | 84 +++++++++++++++++ tests/kernels/moe/test_ocp_mx_moe.py | 52 ++++++++++- .../layers/fused_moe/oracle/mxfp4.py | 92 ++++++++++++++++--- 3 files changed, 214 insertions(+), 14 deletions(-) diff --git a/tests/kernels/moe/test_cutlass_moe.py b/tests/kernels/moe/test_cutlass_moe.py index 191079662326..050e1bc10487 100644 --- a/tests/kernels/moe/test_cutlass_moe.py +++ b/tests/kernels/moe/test_cutlass_moe.py @@ -37,9 +37,12 @@ run_cutlass_moe_fp4, run_cutlass_moe_fp8, run_cutlass_moe_mxfp4, + swizzle_mxfp4_scales, ) +from vllm.model_executor.layers.fused_moe.oracle import mxfp4 as mxfp4_oracle from vllm.model_executor.layers.fused_moe.oracle import nvfp4 as nvfp4_oracle from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -294,6 +297,87 @@ def run_moe( torch.testing.assert_close(actual[1:], torch.zeros_like(actual[1:])) +@torch.inference_mode() +def test_cutlass_mxfp4_oracle_matches_direct_kernel(monkeypatch, workspace_init): + """The MXFP4 oracle selects CUTLASS for W4A4 and swizzles checkpoint + scales into the layout the kernel expects.""" + if not CutlassExpertsMxfp4._supports_current_device(): + pytest.skip("CUTLASS mxfp4 MoE is not supported on this GPU") + monkeypatch.setattr(mxfp4_oracle, "_user_moe_activation_override", lambda: None) + + set_random_seed(7) + e, m, n, k, topk = 8, 33, 256, 512, 2 + device = torch.device("cuda") + w1 = torch.randint(0, 256, (e, 2 * n, k // 2), device=device, dtype=torch.uint8) + w2 = torch.randint(0, 256, (e, k, n // 2), device=device, dtype=torch.uint8) + w1_scale = torch.randint(118, 124, (e, 2 * n, k // 32), device=device).byte() + w2_scale = torch.randint(118, 124, (e, k, n // 32), device=device).byte() + a = torch.randn((m, k), device=device, dtype=torch.bfloat16) + score = torch.randn((m, e), device=device, dtype=torch.bfloat16) + topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) + + def swizzle(scale: torch.Tensor) -> torch.Tensor: + rows, cols = scale.shape[1:] + return torch.stack( + [swizzle_mxfp4_scales(s, rows, cols * 32).view(rows, cols) for s in scale] + ) + + expected = torch.empty_like(a) + run_cutlass_moe_mxfp4( + output=expected, + a=a, + w1_fp4=w1, + w1_blockscale=swizzle(w1_scale), + w2_fp4=w2, + w2_blockscale=swizzle(w2_scale), + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SILU, + workspace13=torch.empty((m * topk, max(2 * n, k)), device=device).to(a), + workspace2=torch.empty((m * topk, n), device=device).to(a), + m=m, + n=n, + k=k, + e=e, + device=device, + apply_router_weight_on_input=False, + ) + + moe_config = make_dummy_moe_config( + num_experts=e, experts_per_token=topk, hidden_dim=k, intermediate_size=n + ) + with set_current_vllm_config(vllm_config): + backend, experts_cls = mxfp4_oracle.select_mxfp4_moe_backend( + moe_config, activation_key=kMxfp4Dynamic + ) + assert backend == mxfp4_oracle.Mxfp4MoeBackend.CUTLASS_MXFP4_MXFP4 + + w13, w2, w13_scale, w2_scale, _, _ = ( + mxfp4_oracle.convert_weight_to_mxfp4_moe_kernel_format( + backend, torch.nn.Module(), w1, w2, w1_scale, w2_scale + ) + ) + kernel = mxfp4_oracle.make_mxfp4_moe_kernel( + mxfp4_oracle.make_mxfp4_moe_quant_config(backend, w13_scale, w2_scale), + moe_config, + experts_cls, + backend, + ) + actual = kernel.apply( + a, + w13, + w2, + topk_weights, + topk_ids, + activation=MoEActivation.SILU, + global_num_experts=e, + expert_map=None, + apply_router_weight_on_input=False, + ) + + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + @dataclasses.dataclass class MOETensors: a: torch.Tensor diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 5a7053a8b1d4..f20927565ebd 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -7,7 +7,11 @@ import pytest import torch -from tests.kernels.moe.utils import check_accuracy, check_deferred_moe_finalize +from tests.kernels.moe.utils import ( + check_accuracy, + check_deferred_moe_finalize, + make_dummy_moe_config, +) from vllm._aiter_ops import ( is_aiter_found, is_aiter_found_and_supported, @@ -2009,6 +2013,52 @@ def is_supported_config( mxfp4_oracle.select_mxfp4_moe_backend(moe_config) +@pytest.mark.parametrize( + "moe_backend,w4a4,supported,expected", + [ + # The checkpoint's activation wins over BF16 backends listed first. + ("auto", True, {"CUTLASS_MXFP4_MXFP4", "MARLIN"}, "CUTLASS_MXFP4_MXFP4"), + # No W4A4 kernel: run at BF16 rather than another quantized activation. + ("auto", True, {"FLASHINFER_TRTLLM_MXFP4_MXFP8", "MARLIN"}, "MARLIN"), + ("marlin", True, {"MARLIN"}, "MARLIN"), + # A BF16 checkpoint never gets its activations quantized. + ("auto", False, {"CUTLASS_MXFP4_MXFP4", "MARLIN"}, "MARLIN"), + ], +) +def test_select_mxfp4_moe_backend_respects_checkpoint_activation( + monkeypatch: pytest.MonkeyPatch, + moe_backend: str, + w4a4: bool, + supported: set[str], + expected: str, +): + """Activations run at the checkpoint's precision when a backend supports + it and fall back to BF16 otherwise, never to a lower precision.""" + import vllm.model_executor.layers.fused_moe.oracle.mxfp4 as mxfp4_oracle + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Dynamic, + ) + + def backend_to_kernel_cls(backend): + class Experts: + @staticmethod + def is_supported_config(*args): + return backend.name in supported, None + + return [Experts] + + monkeypatch.setattr(mxfp4_oracle, "backend_to_kernel_cls", backend_to_kernel_cls) + monkeypatch.setattr(mxfp4_oracle, "_user_moe_activation_override", lambda: None) + moe_config = make_dummy_moe_config() + moe_config.moe_backend = moe_backend + + backend, _ = mxfp4_oracle.select_mxfp4_moe_backend( + moe_config, activation_key=kMxfp4Dynamic if w4a4 else None + ) + + assert backend.name == expected + + @pytest.mark.parametrize("backend_name", ["AITER_TRITON_MXFP4_BF16", "AITER_MXFP4_FP8"]) @pytest.mark.parametrize("use_ep", [False, True]) def test_aiter_mxfp4_monolithic_rejects_expert_parallel( diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 5463853caba1..f5fccc03436f 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -114,6 +114,18 @@ def _pack_deepgemm_mxfp4_scales( ) +def _swizzle_cutlass_mxfp4_scales(scale: torch.Tensor) -> torch.Tensor: + """Swizzle [E, N, K // 32] block scales into the CUTLASS tiled layout.""" + from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( + swizzle_mxfp4_scales, + ) + + n, scale_k = scale.shape[1:] + return torch.stack( + [swizzle_mxfp4_scales(s, n, scale_k * 32).reshape(n, scale_k) for s in scale] + ) + + class Mxfp4MoeBackend(Enum): NONE = "None" # b12x backends @@ -127,6 +139,8 @@ class Mxfp4MoeBackend(Enum): # FlashInfer CUTLASS backends FLASHINFER_CUTLASS_MXFP4_MXFP8 = "FLASHINFER_CUTLASS_MXFP4_MXFP8" FLASHINFER_CUTLASS_MXFP4_BF16 = "FLASHINFER_CUTLASS_MXFP4_BF16" + # vLLM CUTLASS backend + CUTLASS_MXFP4_MXFP4 = "CUTLASS_MXFP4_MXFP4" # Marlin BATCHED_MARLIN = "BATCHED_MARLIN" MARLIN = "MARLIN" @@ -218,6 +232,13 @@ def backend_to_kernel_cls( return [FlashInferExperts] + elif backend == Mxfp4MoeBackend.CUTLASS_MXFP4_MXFP4: + from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( + CutlassExpertsMxfp4, + ) + + return [CutlassExpertsMxfp4] + elif backend == Mxfp4MoeBackend.TRITON: from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( # noqa: E501 OAITritonExperts, @@ -378,6 +399,7 @@ def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]: Mxfp4MoeBackend.TRITON, Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, + Mxfp4MoeBackend.CUTLASS_MXFP4_MXFP4, # TRITON_UNFUSED has bug with MTP support # TODO re-enable after kernel is fixed # TRITON_UNFUSED @@ -440,7 +462,10 @@ def _backend_activation_key(backend: Mxfp4MoeBackend) -> QuantKey | None: return kMxfp8Dynamic if backend == Mxfp4MoeBackend.AITER_MXFP4_FP8: return kFp8StaticTensorSym - if backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4: + if backend in ( + Mxfp4MoeBackend.AITER_MXFP4_MXFP4, + Mxfp4MoeBackend.CUTLASS_MXFP4_MXFP4, + ): return kMxfp4Dynamic return None # BF16 activation @@ -474,6 +499,22 @@ def _make_log_backend(backend: Mxfp4MoeBackend) -> str: return f"Using '{backend.value}' Mxfp4 MoE backend." +def _warn_on_bf16_fallback( + backend: Mxfp4MoeBackend, requested_activation_key: QuantKey | None +) -> None: + if ( + requested_activation_key is not None + and backend != Mxfp4MoeBackend.EMULATION + and _backend_activation_key(backend) is None + ): + logger.warning_once( + "No Mxfp4 MoE backend supports activation=%s for the deployment " + "configuration. Falling back to '%s' with BF16 activations.", + requested_activation_key, + backend.value, + ) + + def _make_log_unsupported(backend: Mxfp4MoeBackend, reason: str | None) -> str: base = ( f"Mxfp4 MoE backend '{backend.value}' does not support the " @@ -505,17 +546,19 @@ def _filter_by_activation( backends: list[Mxfp4MoeBackend], requested_activation_key: QuantKey | None, ) -> list[Mxfp4MoeBackend]: - """Pick variants matching ``requested_activation_key``; without one, - prefer BF16 if the list has any, else keep the list as-is so explicit - non-BF16 picks (e.g. the ``_afp8`` aliases) still land.""" + """Pick variants matching ``requested_activation_key``, followed by BF16 + variants as a higher-precision fallback; without one, prefer BF16 if the + list has any, else keep the list as-is so explicit non-BF16 picks (e.g. + the ``_afp8`` aliases) still land.""" + bf16 = [b for b in backends if _backend_activation_key(b) is None] if requested_activation_key is not None: - return [ + matching = [ b for b in backends if _backend_activation_key(b) == requested_activation_key or b == Mxfp4MoeBackend.EMULATION ] - bf16 = [b for b in backends if _backend_activation_key(b) is None] + return matching + [b for b in bf16 if b not in matching] return bf16 if bf16 else backends @@ -576,8 +619,9 @@ def select_mxfp4_moe_backend( Args: config: MoE configuration - activation_key: Optional activation quantization key. If provided, - overrides the default activation key for backend selection. + activation_key: Optional activation quantization key declared by + the checkpoint. Backends that match it are preferred, with BF16 + activation backends as a fallback. Use kFp8StaticTensorSym for W4A8 scheme. Note: Shape-specific fallbacks may still occur at runtime. @@ -614,7 +658,7 @@ def select_mxfp4_moe_backend( else _backend_activation_key(requested_backend) ) try: - return _return_or_raise( + selected = _return_or_raise( requested_backend, config, kMxfp4Static, @@ -623,6 +667,9 @@ def select_mxfp4_moe_backend( ) except ValueError as e: last_error = e + else: + _warn_on_bf16_fallback(requested_backend, requested_activation_key) + return selected assert last_error is not None raise last_error @@ -648,10 +695,10 @@ def select_mxfp4_moe_backend( unsupported_reasons = [] for backend in AVAILABLE_BACKENDS: - # Use requested_activation_key if provided, otherwise use backend default + # EMULATION takes the requested key, other backends use their own act_key = ( requested_activation_key - if requested_activation_key is not None + if backend == Mxfp4MoeBackend.EMULATION else _backend_activation_key(backend) ) for k_cls in backend_to_kernel_cls(backend): @@ -660,6 +707,7 @@ def select_mxfp4_moe_backend( ) if supported: logger.info_once(_make_log_backend(backend)) + _warn_on_bf16_fallback(backend, requested_activation_key) return backend, k_cls else: logger.debug_once(_make_log_unsupported(backend, reason)) @@ -789,6 +837,10 @@ def mxfp4_round_up_hidden_size_and_intermediate_size( # DeepGEMM requires M/N/K alignment intermediate_size = round_up(intermediate_size, 128) hidden_size = round_up(hidden_size, 128) + elif backend == Mxfp4MoeBackend.CUTLASS_MXFP4_MXFP4: + # CUTLASS scale tiles cover 128 rows by 128 elements + intermediate_size = round_up(intermediate_size, 128) + hidden_size = round_up(hidden_size, 128) elif backend in (Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN): intermediate_size = round_up(intermediate_size, 128) if current_platform.is_xpu(): @@ -1433,7 +1485,8 @@ def convert_weight_to_mxfp4_moe_kernel_format( ]: """Convert loaded weights into backend-specific kernel format. - Supports DeepGEMM, FlashInfer, TRTLLM MXFP8, Triton and Marlin backends. + Supports CUTLASS, DeepGEMM, FlashInfer, TRTLLM MXFP8, Triton and Marlin + backends. """ is_gfx1250 = False if current_platform.is_rocm(): @@ -1451,6 +1504,16 @@ def convert_weight_to_mxfp4_moe_kernel_format( w2_bias, ) + if mxfp4_backend == Mxfp4MoeBackend.CUTLASS_MXFP4_MXFP4: + return ( + w13_weight.data, + w2_weight.data, + _swizzle_cutlass_mxfp4_scales(w13_weight_scale.data), + _swizzle_cutlass_mxfp4_scales(w2_weight_scale.data), + w13_bias, + w2_bias, + ) + if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales( w13_weight, @@ -2005,7 +2068,10 @@ def make_mxfp4_moe_quant_config( block_shape=None, gemm1_clamp_limit=swiglu_limit, ) - elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4: + elif mxfp4_backend in ( + Mxfp4MoeBackend.AITER_MXFP4_MXFP4, + Mxfp4MoeBackend.CUTLASS_MXFP4_MXFP4, + ): return ocp_mx_moe_quant_config( quant_dtype="mxfp4", w1_bias=w1_bias, From 2274e86842d6e59a5f9f9d1fb6cc1dd2d2bf8350 Mon Sep 17 00:00:00 2001 From: Shijin Zhang <75300765+Dovis01@users.noreply.github.com> Date: Fri, 9 Oct 2026 02:41:10 +0800 Subject: [PATCH 119/624] [Spec Decode][Model] Support DFlash2 draft models with GLM-5.3-Flash (#56983) Signed-off-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com> --- vllm/models/glm5next/common/model.py | 51 +++++++++++++++++++++++++--- 1 file changed, 47 insertions(+), 4 deletions(-) diff --git a/vllm/models/glm5next/common/model.py b/vllm/models/glm5next/common/model.py index 26becc9619e7..fd7c0263bdda 100644 --- a/vllm/models/glm5next/common/model.py +++ b/vllm/models/glm5next/common/model.py @@ -66,9 +66,11 @@ Glm4vForConditionalGeneration, ) from vllm.model_executor.models.interfaces import ( + EagleModelMixin, HasInnerState, IsHybrid, MixtureOfExperts, + SupportsEagle3, SupportsPP, ) from vllm.model_executor.models.utils import ( @@ -685,7 +687,7 @@ def hc_fused_post_pre( ) -class Glm5NextModel(nn.Module): +class Glm5NextModel(nn.Module, EagleModelMixin): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -743,6 +745,7 @@ def get_layer(prefix: str): get_layer, prefix=f"{prefix}.layers", ) + self._aux_post_op = MHCPostOp() # The active slice is fixed after construction; cache it so forward # doesn't rebuild the slice (a fresh list) every step. self._active_layers = self.layers[self.start_layer : self.end_layer] @@ -767,6 +770,26 @@ def get_layer(prefix: str): def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) + def _aux_hidden_state( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + post: torch.Tensor | None, + comb: torch.Tensor | None, + ) -> torch.Tensor: + """Completed residual stream entering a layer, as one hidden vector. + + mHC layers defer their final ``hc_post`` into the next layer's fused + pre-op, so ``hidden_states`` holds the raw layer output while the + widened stream is completed here and contracted back to + ``hidden_size``. Non-mHC layers already return the summed stream. + """ + if post is None: + return hidden_states + assert residual is not None and comb is not None + completed = self._aux_post_op(hidden_states, residual, post, comb) + return hc_contract(completed, self.config.mhc_num_residual_streams) + def forward( self, input_ids: torch.Tensor | None, @@ -796,7 +819,15 @@ def forward( if self.is_sequence_parallel: hidden_states = sp_shard(hidden_states) - for layer in self._active_layers: + aux_hidden_states: list[torch.Tensor] = [] + for idx, layer in enumerate(self._active_layers, start=self.start_layer): + if idx in self.aux_hidden_state_layers: + aux_hidden_state = self._aux_hidden_state( + hidden_states, residual, post, comb + ) + if self.is_sequence_parallel: + aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens] + aux_hidden_states.append(aux_hidden_state) hidden_states, residual, post, comb = layer( positions, hidden_states, residual, post, comb ) @@ -811,10 +842,18 @@ def forward( {"hidden_states": hidden_states, "residual": residual} ) + if self.end_layer in self.aux_hidden_state_layers: + final_aux = self._aux_hidden_state(hidden_states, residual, post, comb) + if self.is_sequence_parallel: + final_aux = sp_all_gather(final_aux)[:full_num_tokens] + aux_hidden_states.append(final_aux) + if self.is_sequence_parallel: hidden_states = sp_all_gather(hidden_states)[:full_num_tokens] hidden_states = self.norm(hidden_states) + if aux_hidden_states: + return hidden_states, aux_hidden_states return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -1006,7 +1045,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: class Glm5NextForCausalLM( - nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid + nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid, SupportsEagle3 ): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -1102,7 +1141,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: dummy_inputs=Glm4vDummyInputsBuilder, ) class Glm5NextForConditionalGeneration( - Glm4vForConditionalGeneration, HasInnerState, IsHybrid, MixtureOfExperts + Glm4vForConditionalGeneration, + HasInnerState, + IsHybrid, + MixtureOfExperts, + SupportsEagle3, ): # The text model (KDA + dense-MLA + MoE) is a hybrid mamba model. The # multimodal wrapper must declare the same interfaces so vLLM treats it as From c4a5ff5c18ceba5f06512ef6b17a48338f275763 Mon Sep 17 00:00:00 2001 From: Ning Xie Date: Fri, 9 Oct 2026 02:46:01 +0800 Subject: [PATCH 120/624] [Frontend] shutdown engine core and its sub processes asap when initializing (#52299) Signed-off-by: Andy Xie Co-authored-by: Simon Mo --- .../test_engine_core_shutdown_signal.py | 224 ++++++++++++++++++ vllm/v1/engine/core.py | 68 ++++-- 2 files changed, 275 insertions(+), 17 deletions(-) create mode 100644 tests/v1/engine/test_engine_core_shutdown_signal.py diff --git a/tests/v1/engine/test_engine_core_shutdown_signal.py b/tests/v1/engine/test_engine_core_shutdown_signal.py new file mode 100644 index 000000000000..73bc794088d6 --- /dev/null +++ b/tests/v1/engine/test_engine_core_shutdown_signal.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Regression tests for termination signals that arrive while EngineCoreProc is +# still being constructed. +# +# When this fix matters +# --------------------- +# Engine shutdown is driven from the frontend. If a frontend process dies while +# EngineCore is still starting up, the process manager sends SIGTERM to its +# children and force kills whatever is still alive after a short timeout (see +# shutdown in vllm/v1/utils.py). Construction is slow - it spawns the executor +# processes, loads weights, profiles and captures CUDA graphs - so that +# shutdown often lands inside the construction window. Bigger models and more +# ranks widen the window. This is the case PR 52299 (follow up to 43417) has to +# survive: 43417 stops EngineCore from exiting silently, but the request itself +# still has to outlive the construction window. +# +# run_engine_core installs the signal handler before engine_core exists, so in +# that window the handler runs with engine_core == None. Losing the request +# there is permanent, not a delay: shutdown_state stays RUNNING, and +# wakeup_engine returns early while consuming the one-shot SignalCallback. The +# busy loop then starts RUNNING and blocks on input_queue forever, so no later +# signal can wake it. Graceful shutdown is gone for good and the process only +# stops when something force kills it, which skips executor shutdown and GPU +# release. In a log this shows up as the signal being received, followed by a +# force kill of EngineCore after the process manager timeout. +# +# What these tests pin +# -------------------- +# A request arriving in the window must be remembered somewhere other than +# engine_core, applied before run_busy_loop starts, and must not consume the +# one-shot wakeup callback, so a later signal can still wake an idle loop. + +import queue +import signal as signal_module +from typing import Any + +from vllm.v1.engine import EngineCoreRequestType +from vllm.v1.engine import core as core_module +from vllm.v1.engine.core import EngineCoreProc, EngineShutdownState + +WAKEUP_ITEM = (EngineCoreRequestType.WAKEUP, None) +CONSTRUCTION_SIGNAL = signal_module.SIGTERM + + +class _FakeSignal: + # Minimal stand-in for the signal module. It records the installed handler + # so the test can fire it by hand, and avoids touching the real process + # signal disposition. Signals is reused from the real module so the handler + # can still resolve signal.Signals(signum).name. + SIGTERM = signal_module.SIGTERM + SIGINT = signal_module.SIGINT + SIG_DFL = signal_module.SIG_DFL + Signals = signal_module.Signals + + def __init__(self): + self.handlers = {} + + def signal(self, signum, handler): + self.handlers[signum] = handler + + def fire(self): + self.handlers[CONSTRUCTION_SIGNAL](CONSTRUCTION_SIGNAL, None) + + +class _FakeSignalCallback: + # One-shot stand-in for SignalCallback. The real class runs the callback on + # a dedicated thread and exits after a single trigger, so every later + # trigger is a no-op. Running the callback inline keeps the test + # deterministic instead of depending on thread scheduling. + instances: list["_FakeSignalCallback"] = [] + + def __init__(self, callback): + self.callback = callback + self.trigger_count = 0 + self.consumed = False + self.stopped = False + _FakeSignalCallback.instances.append(self) + + def trigger(self): + self.trigger_count += 1 + if self.stopped or self.consumed: + return + self.consumed = True + callback, self.callback = self.callback, None + callback() + + def stop(self): + self.stopped = True + + +class _StubParallelConfig: + def __init__(self): + self.data_parallel_size = 1 + self.data_parallel_rank_local = 0 + self.data_parallel_index = 0 + self.data_parallel_rank = 0 + self.numa_bind = False + + def reconfigure_for_independent_dp_rank(self): + pass + + +class _StubModelConfig: + is_moe = False + + +class _StubVllmConfig: + def __init__(self): + self.parallel_config = _StubParallelConfig() + self.model_config = _StubModelConfig() + self.logging_config = None + self.kv_transfer_config = None + self.shutdown_timeout = 0 + + +def _run_engine_core( + monkeypatch, + *, + signal_during_construction=False, + signal_while_idle=False, +): + _FakeSignalCallback.instances.clear() + fake_signal = _FakeSignal() + trace = {} + + class StubEngineCoreProc: + def __init__(self, *args, **kwargs): + self.shutdown_state = EngineShutdownState.RUNNING + # Same queue type as EngineCore, so the wakeup path under test is + # the real one. + self.input_queue = queue.Queue[tuple[EngineCoreRequestType, Any]]() + self.vllm_config = kwargs.get("vllm_config") + if signal_during_construction: + # The termination signal lands mid construction, while the + # caller still sees engine_core == None. + fake_signal.fire() + trace["constructed"] = True + + def run_busy_loop(self): + # What the real loop sees on its first _handle_shutdown() check. + trace["loop_entered"] = True + trace["state_at_loop_entry"] = self.shutdown_state + if signal_while_idle: + fake_signal.fire() + try: + trace["wakeup"] = self.input_queue.get_nowait() + except queue.Empty: + trace["wakeup"] = None + + def shutdown(self): + pass + + def _send_engine_dead(self): + pass + + def has_work(self): + return False + + monkeypatch.setattr(core_module, "signal", fake_signal) + monkeypatch.setattr(core_module, "SignalCallback", _FakeSignalCallback) + monkeypatch.setattr(core_module, "EngineCoreProc", StubEngineCoreProc) + monkeypatch.setattr(core_module, "DPEngineCoreProc", StubEngineCoreProc) + monkeypatch.setattr(core_module, "configure_logging", lambda *a, **k: None) + monkeypatch.setattr( + core_module, "maybe_register_config_serialize_by_value", lambda *a, **k: None + ) + monkeypatch.setattr(core_module, "set_process_title", lambda *a, **k: None) + monkeypatch.setattr(core_module, "maybe_init_worker_tracer", lambda *a, **k: None) + monkeypatch.setattr(core_module, "decorate_logs", lambda *a, **k: None) + + EngineCoreProc.run_engine_core(vllm_config=_StubVllmConfig()) + return trace + + +def test_signal_during_construction_is_applied_before_loop(monkeypatch): + trace = _run_engine_core(monkeypatch, signal_during_construction=True) + + assert trace["constructed"] is True + assert trace["loop_entered"] is True + # The request arrived while engine_core was None, so it could not be stored + # there. It must still be applied before run_busy_loop starts, otherwise the + # loop comes up as RUNNING, the shutdown is silently dropped, and force kill + # becomes the only way out. + assert trace["state_at_loop_entry"] == EngineShutdownState.REQUESTED, ( + "a signal received during construction must be applied before the loop " + f"starts, got {trace['state_at_loop_entry']!r}" + ) + + +def test_signal_during_construction_keeps_later_wakeup_working(monkeypatch): + trace = _run_engine_core( + monkeypatch, + signal_during_construction=True, + signal_while_idle=True, + ) + + assert trace["state_at_loop_entry"] == EngineShutdownState.REQUESTED, ( + "the construction time signal must be applied before the loop starts, " + f"got {trace['state_at_loop_entry']!r}" + ) + # SignalCallback is one-shot. If the construction time signal consumed it, + # no wakeup could be delivered afterwards and the idle busy loop would block + # on input_queue forever instead of shutting down, so the process would sit + # there holding GPU memory until something force kills it. + assert trace["wakeup"] == WAKEUP_ITEM, ( + "a signal after construction must still wake the idle loop, " + f"got {trace['wakeup']!r}" + ) + + +def test_signal_while_idle_still_wakes_the_loop(monkeypatch): + trace = _run_engine_core(monkeypatch, signal_while_idle=True) + + assert trace["state_at_loop_entry"] == EngineShutdownState.RUNNING + assert trace["wakeup"] == WAKEUP_ITEM + + +def test_no_signal_leaves_the_loop_running(monkeypatch): + trace = _run_engine_core(monkeypatch) + + assert trace["state_at_loop_entry"] == EngineShutdownState.RUNNING + assert trace["wakeup"] is None diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 679bbbfd2415..bfbaa38d9478 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -1410,40 +1410,74 @@ def run_engine_core(*args, dp_rank: int = 0, local_dp_rank: int = 0, **kwargs): vllm_config.kv_transfer_config.engine_id, ) - parallel_config.data_parallel_index = dp_rank - if data_parallel and vllm_config.model_config.is_moe: - # Set data parallel rank for this engine process. - parallel_config.data_parallel_rank = dp_rank - engine_core = DPEngineCoreProc(*args, **kwargs) - else: - # Non-MoE DP ranks are completely independent, so treat like DP=1. - # Note that parallel_config.data_parallel_index will still reflect - # the original DP rank. - parallel_config.reconfigure_for_independent_dp_rank() - engine_core = EngineCoreProc(*args, engine_index=dp_rank, **kwargs) - - assert engine_core is not None + # declare engine core as None and register signal handler before + # engine_core initialization in case engine core process is terminated + # in the middle of initialization. For example, frontend process exits + # unexpectedly, if signal handler is not registered, engine core process + # will use default signal handler which will terminate the process + # silently. If engine core model executor has been initialized, it will + # not be terminated asap. Register signal handler early and only log a + # message will make engine core process alive. engine core process and + # its subprocesses(model executor processes) will be force killed in + # CoreEngineProcManager.shutdown. + engine_core = None + # A termination signal can arrive while engine_core is still being + # constructed. At that point the request cannot be recorded on + # engine_core, so keep it in its own flag and apply it later. + shutdown_requested = False def wakeup_engine(): # Wakes up idle engine via input_queue when shutdown is requested # Not safe in a signal handler - we may interrupt the main thread # while it is holding the non-reentrant input_queue.mutex + assert engine_core is not None engine_core.input_queue.put_nowait((EngineCoreRequestType.WAKEUP, None)) - signal_callback = SignalCallback(wakeup_engine) - def signal_handler(signum, frame): + nonlocal shutdown_requested signal_name = signal.Signals(signum).name logger.info( "[shutdown] EngineCore: trigger received signal=%s", signal_name, ) - engine_core.shutdown_state = EngineShutdownState.REQUESTED - signal_callback.trigger() + # Record the request first so it survives the construction window. + shutdown_requested = True + if engine_core is not None: + engine_core.shutdown_state = EngineShutdownState.REQUESTED + # signal_callback is armed only after engine_core exists. Don't + # trigger it earlier: it is one-shot, and a wakeup cannot be + # delivered while there is no input_queue to wake. + # For more info: https://github.com/vllm-project/vllm/pull/52299#pullrequestreview-5370255271 + if signal_callback is not None: + signal_callback.trigger() signal.signal(signal.SIGTERM, signal_handler) signal.signal(signal.SIGINT, signal_handler) + parallel_config.data_parallel_index = dp_rank + if data_parallel and vllm_config.model_config.is_moe: + # Set data parallel rank for this engine process. + parallel_config.data_parallel_rank = dp_rank + engine_core = DPEngineCoreProc(*args, **kwargs) + else: + # Non-MoE DP ranks are completely independent, so treat like DP=1. + # Note that parallel_config.data_parallel_index will still reflect + # the original DP rank. + parallel_config.reconfigure_for_independent_dp_rank() + engine_core = EngineCoreProc(*args, engine_index=dp_rank, **kwargs) + + assert engine_core is not None + + # Apply a signal received during construction before entering the + # loop. run_busy_loop calls _handle_shutdown first, so the loop exits + # right away instead of blocking on an idle input_queue. + if shutdown_requested: + engine_core.shutdown_state = EngineShutdownState.REQUESTED + + # Arm the wakeup callback only now that input_queue exists, so a + # later signal can still wake an idle busy loop. + signal_callback = SignalCallback(wakeup_engine) + engine_core.run_busy_loop() except SystemExit as e: From 56e6da2c4717f6c6c2faa0159f8c0fa2cf508d6f Mon Sep 17 00:00:00 2001 From: Yeonwoo Sung Date: Fri, 9 Oct 2026 03:47:30 +0900 Subject: [PATCH 121/624] [Bugfix][Frontend] Honor stop_token_ids on the Responses API (#56592) Signed-off-by: YeonwooSung Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: Flora Feng <4florafeng@gmail.com> --- tests/entrypoints/openai/responses/test_sampling_params.py | 2 ++ vllm/entrypoints/openai/responses/protocol.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/tests/entrypoints/openai/responses/test_sampling_params.py b/tests/entrypoints/openai/responses/test_sampling_params.py index 877e8724e8b2..01fdfa4594ea 100644 --- a/tests/entrypoints/openai/responses/test_sampling_params.py +++ b/tests/entrypoints/openai/responses/test_sampling_params.py @@ -47,6 +47,7 @@ def test_extra_sampling_params(self): repetition_penalty=1.2, seed=42, stop=["END", "STOP"], + stop_token_ids=[42, 43], ignore_eos=True, vllm_xargs={"custom": "value"}, ) @@ -56,6 +57,7 @@ def test_extra_sampling_params(self): assert sampling_params.repetition_penalty == 1.2 assert sampling_params.seed == 42 assert sampling_params.stop == ["END", "STOP"] + assert sampling_params.stop_token_ids == [42, 43] assert sampling_params.ignore_eos is True assert sampling_params.extra_args == {"custom": "value"} diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 35e6d3aa3dcd..12e40e98a882 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -363,6 +363,7 @@ class ResponsesRequest(OpenAIBaseModel): repetition_penalty: float | None = None seed: int | None = Field(None, ge=_INT64_MIN, le=_INT64_MAX) stop: StopParam = [] + stop_token_ids: list[int] | None = [] ignore_eos: bool = False vllm_xargs: dict[str, str | int | float | list[str | int | float]] | None = Field( default=None, @@ -521,6 +522,7 @@ def to_sampling_params( max_tokens=max_tokens, logprobs=self.top_logprobs if self.is_include_output_logprobs() else None, stop=stop, + stop_token_ids=self.stop_token_ids, frequency_penalty=frequency_penalty, presence_penalty=presence_penalty, repetition_penalty=repetition_penalty, From 9bbc58156095f45bdb91ddbdbd244119d3afb12c Mon Sep 17 00:00:00 2001 From: Aarushi Jain Date: Thu, 8 Oct 2026 13:53:00 -0500 Subject: [PATCH 122/624] [Bugfix] Close MessageQueue zmq resources via finalizer to avoid exit (#60472) Signed-off-by: aarushjain29 Co-authored-by: Andreas Karatzas --- tests/distributed/test_shm_broadcast.py | 22 ++++++++++++++++ .../device_communicators/shm_broadcast.py | 25 +++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py index 7debed4cbbff..d28732261a13 100644 --- a/tests/distributed/test_shm_broadcast.py +++ b/tests/distributed/test_shm_broadcast.py @@ -1,11 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import gc import io import pickle import random import threading import time +import weakref from types import SimpleNamespace from unittest import mock @@ -634,6 +636,26 @@ def test_acquire_read_releases_slot_when_reader_raises(): reader.shutdown() +def test_reader_in_reference_cycle_is_collected_without_blocking(): + writer = MessageQueue( + n_reader=1, + n_local_reader=1, + max_chunk_bytes=1024 * 1024, + max_chunks=1, + ) + reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0) + reader.cycle = reader + reader_ref = weakref.ref(reader) + del reader + + collector = threading.Thread(target=gc.collect, daemon=True) + collector.start() + collector.join(timeout=10) + + assert not collector.is_alive() + assert reader_ref() is None + + def test_warning_logs(caplog_vllm): """Test that warning logs are emitted at VLLM_RINGBUFFER_WARNING_INTERVAL intervals when indefinite=False, and are not emitted when indefinite=True. diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index bafc5728b85e..da2288f77b54 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -9,6 +9,7 @@ import sys import threading import time +import weakref from contextlib import contextmanager from dataclasses import dataclass, field from multiprocessing import shared_memory @@ -472,6 +473,12 @@ class Handle: remote_addr_ipv6: bool = False +def _close_zmq(context: zmq.Context, sockets: list[zmq.Socket]) -> None: + for socket in sockets: + socket.close(linger=0) + context.term() + + class MessageQueue: def __init__( self, @@ -561,9 +568,26 @@ def __init__( remote_subscribe_addr=remote_subscribe_addr, remote_addr_ipv6=remote_addr_ipv6, ) + self._close_zmq_on_collect(context) logger.debug("vLLM message queue communication handle: %s", self.handle) + def _close_zmq_on_collect(self, context: zmq.Context) -> None: + # pyzmq's Context.__del__ can block forever in term() when the context + # and its sockets are collected in the same GC cycle, so the finalizer + # keeps them alive and closes the sockets before terminating. + sockets = [self.local_socket, self.remote_socket] + if self._spin_condition is not None: + sockets += [ + self._spin_condition.local_notify_socket, + self._spin_condition.read_cancel_socket, + self._spin_condition.write_cancel_socket, + ] + finalizer = weakref.finalize( + self, _close_zmq, context, [s for s in sockets if s is not None] + ) + finalizer.atexit = False # type: ignore[misc] + def export_handle(self) -> Handle: return self.handle @@ -613,6 +637,7 @@ def create_from_handle(handle: Handle, rank) -> "MessageQueue": self._spin_condition = None # type: ignore self.shutting_down = False + self._close_zmq_on_collect(context) return self def wait_until_ready(self): From 18751bc6a7b4b24fa5e6bbd5299100dd0b19a22b Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Thu, 8 Oct 2026 15:11:45 -0400 Subject: [PATCH 123/624] [ROCm][CI] Add coverage for AiterExperts token padding: inf/nan garbage rows (#59334) Signed-off-by: Divakar Verma Co-authored-by: Andreas Karatzas --- .../moe/test_modular_kernel_combinations.py | 171 ++++++++++++++++++ 1 file changed, 171 insertions(+) diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index 4b769931ea04..f29c22085240 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -17,6 +17,10 @@ from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import AiterExperts +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8DynamicTensorSym, + kFp8StaticTensorSym, +) from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe from vllm.utils.import_utils import has_aiter, has_deep_ep, has_deep_gemm @@ -988,6 +992,173 @@ def test_aiter_moe_padding_matrix_mxfp4(): padding matrix once MI350/gfx950 hardware is available in CI.""" +# --- 4c: HIP-graph token padding -- inf/nan garbage rows (topk_ids=-1) ------ + +# AiterExperts isolates garbage rows from real rows under both AITER sorting +# backends, and reference_moe_impl's dispatch (`mask = topk_ids == i`) is +# likewise row-local, so comparing real rows against it is valid. Exception: +# AiterExperts's a1 quant scale is normally fixed before garbage injection +# (computed in RankTensors.make()), but "fp8_tensor_token" computes it live +# from the (already-injected) hidden_states -- see the skip below. + +# {mode name -> {row index: fill value}}. Covers multiple garbage-row counts +# and positions, and both inf and nan. +_TOKEN_PADDING_GARBAGE_MODES: dict[str, dict[int, float]] = { + "single_inf": {3: float("inf")}, + "single_nan": {5: float("nan")}, + "multiple_mixed": { + 1: float("inf"), + 4: float("-inf"), + 7: float("nan"), + 12: float("inf"), + }, +} + + +def _apply_token_padding_garbage_rows( + rank_tensors: RankTensors, garbage_rows: dict[int, float] +) -> None: + """Simulate the HIP-graph padding-token contract in-place: garbage rows + get inf/nan hidden_states and their routing masked to topk_ids=-1 / + topk_weights=0, mirroring grouped_topk_router.py's `is_padding` masking. + """ + for row, fill in garbage_rows.items(): + rank_tensors.hidden_states[row].fill_(fill) + rank_tensors.topk_ids[row].fill_(-1) + rank_tensors.topk_weights[row].fill_(0.0) + + +def _aiter_token_padding_worker( + pgi: ProcessGroupInfo, + vllm_config: VllmConfig, + cpu_group, + config: Config, + weights: WeightTensors, + verbose: bool, + garbage_rows: dict[int, float], +): + device = torch.device(f"cuda:{pgi.local_rank}") + init_workspace_manager(device) + set_random_seed(pgi.rank) + + weights = copy.deepcopy(weights) + weights.to_current_device() + + rank_tensors = RankTensors.make(config, pgi) + _apply_token_padding_garbage_rows(rank_tensors, garbage_rows) + + with set_current_vllm_config(vllm_config): + mk_out = run_modular_kernel(pgi, vllm_config, config, weights, rank_tensors) + ref_out = reference_moe_impl(config, weights, rank_tensors) + + real_rows = [i for i in range(config.Ms) if i not in garbage_rows] + + assert torch.isfinite(mk_out[real_rows]).all(), ( + "Garbage (padding) token rows leaked NaN/Inf into real rows' output " + "through AiterExperts." + ) + + # ref_out's magnitude here (~1e-3-1e-2) is below atol=3e-2, so a zeroed-out + # mk_out would still pass the checks below -- guard against that. + ref_scale = ref_out[real_rows].abs().mean() + mk_scale = mk_out[real_rows].abs().mean() + assert mk_scale > 0.5 * ref_scale, ( + f"AiterExperts output magnitude (mean |mk_out|={mk_scale:.6f}) looks " + f"degenerate/zeroed vs. reference (mean |ref_out|={ref_scale:.6f})." + ) + + if config.quant_config is not None: + check_accuracy( + ref_out[real_rows], mk_out[real_rows], atol=3e-2, rtol=3e-2, percent=0.9 + ) + else: + torch.testing.assert_close( + ref_out[real_rows], mk_out[real_rows], atol=3e-2, rtol=3e-2 + ) + + +@require_aiter_moe +@pytest.mark.parametrize("quant_config", _PADDING_QUANT_CONFIGS, ids=_PADDING_QUANT_IDS) +@pytest.mark.parametrize("garbage_mode", list(_TOKEN_PADDING_GARBAGE_MODES)) +def test_aiter_moe_token_padding_garbage_rows( + garbage_mode: str, quant_config: TestMoEQuantConfig | None +): + """See https://github.com/vllm-project/vllm/issues/54966 ("Test padding"). + + Simulates HIP-graph token padding: a subset of rows get inf/nan + hidden_states with topk_ids=-1 / topk_weights=0, mirroring the router's + padding-token masking contract (grouped_topk_router.py). Verifies the + real (non-garbage) rows' output through AiterExperts remains finite and + matches reference_moe_impl, across multiple garbage-row counts/positions + and both inf and nan fill values, for unquantized + the 4 fp8 quant + schemes AiterExperts supports. + """ + from vllm.model_executor.layers.fused_moe.prepare_finalize import ( + MoEPrepareAndFinalizeNoDPEPModular, + ) + + garbage_rows = _TOKEN_PADDING_GARBAGE_MODES[garbage_mode] + + # Reuses 4b's already-confirmed 128-aligned ("padded") K/N -- these are + # known to work across all 5 quant schemes here, including AITER's + # block-quant kernel (which 4b's _PADDING_K_UNPADDED/_PADDING_N_UNPADDED + # sizes do not support). Token-padding safety doesn't depend on + # hidden_pad/intermediate_pad, so there's no need for a separate shape. + config = Config( + Ms=_PADDING_M, + K=_PADDING_K_PADDED, + N=_PADDING_N_PADDED, + E=_PADDING_E, + topks=_PADDING_TOPK, + dtype=torch.bfloat16, + quant_config=quant_config, + prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular, + fused_experts_type=AiterExperts, + world_size=1, + ) + assert config.is_valid()[0] + assert config.fe_supports_quant_scheme(), ( + f"AiterExperts does not support quant scheme {quant_config}." + ) + + # This (weight, activation) pair is AiterExperts's real per-tensor + # dynamic-activation-quant scheme -- its batch-wide max-abs scale is + # poisoned by a real inf garbage row. Upstream bug: + # https://github.com/ROCm/aiter/issues/6275. + if ( + quant_config is not None + and config.fp8_quant_key_pair() == (kFp8StaticTensorSym, kFp8DynamicTensorSym) + and any(v in (float("inf"), float("-inf")) for v in garbage_rows.values()) + ): + pytest.skip( + "Batch-wide dynamic per-tensor quant scale is poisoned by a " + "real inf garbage row. See " + "https://github.com/ROCm/aiter/issues/6275." + ) + + weights = WeightTensors.make(config) + vllm_config, env_dict = config.make_env_data() + + parallel_launch_with_config( + config.world_size, + _aiter_token_padding_worker, + vllm_config, + env_dict, + config, + weights, + False, + garbage_rows, + ) + + +@pytest.mark.skip( + reason="MXFP4 AiterExperts token padding requires MI350/gfx950 hardware. " + 'See https://github.com/vllm-project/vllm/issues/54966 ("Test padding").' +) +def test_aiter_moe_token_padding_garbage_rows_mxfp4(): + pass + + if __name__ == "__main__": # Ability to test individual PrepareAndFinalize and FusedExperts combination from .modular_kernel_tools.cli_args import make_config, make_config_arg_parser From fc494822eaca100b6339a4911ce36784926573fc Mon Sep 17 00:00:00 2001 From: Kevin Li <110002169+Kevinli24@users.noreply.github.com> Date: Thu, 8 Oct 2026 14:14:35 -0500 Subject: [PATCH 124/624] [Test][Quantization] Gate CUTLASS W4A8 kernel tests to SM90 (#58878) Signed-off-by: Kevin Li Co-authored-by: Misha Goin --- tests/kernels/quantization/test_cutlass_w4a8.py | 11 +++++------ tests/kernels/quantization/test_cutlass_w4a8_moe.py | 3 +-- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/tests/kernels/quantization/test_cutlass_w4a8.py b/tests/kernels/quantization/test_cutlass_w4a8.py index 968d5f2c74b9..fc34a15bde56 100644 --- a/tests/kernels/quantization/test_cutlass_w4a8.py +++ b/tests/kernels/quantization/test_cutlass_w4a8.py @@ -29,8 +29,11 @@ # `is_quant_method_supported` conflates kernels with quantization methods # an assumption which is breaking down as quantizations methods can have # have kernels and some kernels support multiple quantization methods. -capability = current_platform.get_device_capability() -IS_SUPPORTED_BY_GPU = capability is not None and capability[0] >= 9 +# CUTLASS W4A8 kernels are only built for SM90 (9.0a); the kernel selector +# requires exactly SM90 too, so match it here rather than accepting >= 90. +IS_SUPPORTED_BY_GPU = ( + current_platform.is_cuda() and current_platform.is_device_capability(90) +) MNK_SHAPES = [ (1, 128, 128), @@ -106,10 +109,6 @@ class Tensors: # TODO: in future PR refactor this and `is_quant_method_supported` in the kernel # unit tests to a common utility function. Currently the use of -# `is_quant_method_supported` conflates kernels with quantization methods -# an assumption which is breaking down as quantizations methods can have -# have kernels and some kernels support multiple quantization methods. -IS_SUPPORTED_BY_GPU = current_platform.has_device_capability(90) # For testing quantized linear kernels diff --git a/tests/kernels/quantization/test_cutlass_w4a8_moe.py b/tests/kernels/quantization/test_cutlass_w4a8_moe.py index b8c10d051056..4050314a1100 100644 --- a/tests/kernels/quantization/test_cutlass_w4a8_moe.py +++ b/tests/kernels/quantization/test_cutlass_w4a8_moe.py @@ -17,9 +17,8 @@ from vllm.scalar_type import ScalarType, scalar_types from vllm.utils.torch_utils import set_random_seed -capability = current_platform.get_device_capability() IS_SUPPORTED_BY_GPU = ( - current_platform.is_cuda() and capability is not None and capability[0] >= 9 + current_platform.is_cuda() and current_platform.is_device_capability(90) ) From fff10d20728780c58a9427986174eb7485ab5c12 Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Fri, 9 Oct 2026 04:16:22 +0900 Subject: [PATCH 125/624] [Attention] Allow an NVFP4 KV cache on SM90 with fa2 (#60620) Signed-off-by: lesj0610 Co-authored-by: Misha Goin --- tests/kernels/attention/test_flashinfer.py | 2 +- vllm/v1/attention/backends/flashinfer.py | 10 ++++------ 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/tests/kernels/attention/test_flashinfer.py b/tests/kernels/attention/test_flashinfer.py index 0d9aecd31e7d..c56af3a6e852 100644 --- a/tests/kernels/attention/test_flashinfer.py +++ b/tests/kernels/attention/test_flashinfer.py @@ -706,7 +706,7 @@ def test_flashinfer_decode_with_paged_fp8_kv( @pytest.mark.skipif( - not any(current_platform.is_device_capability_family(f) for f in (80, 120)), + not any(current_platform.is_device_capability_family(f) for f in (80, 90, 120)), reason="NVFP4 fa2 path", ) @pytest.mark.parametrize("layout", ["LBNHC", "LBHNC", "BLNHC", "BLHNC"]) diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index a21720ab26e6..85b7456dd889 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -108,9 +108,7 @@ def _nvfp4_kv_on_fa2() -> bool: - return current_platform.is_device_capability_family( - 80 - ) or current_platform.is_device_capability_family(120) + return any(current_platform.is_device_capability_family(f) for f in (80, 90, 120)) _KVPair = tuple[torch.Tensor, torch.Tensor] @@ -557,7 +555,7 @@ def supports_combination( if ( kv_cache_dtype is not None and kv_cache_dtype.startswith("nvfp4") - and device_capability.major in (8, 12) + and device_capability.major in (8, 9, 12) ): if head_size == 64: return "fa2 prefill misreads an NVFP4 KV cache at head_size 64" @@ -1311,7 +1309,7 @@ def _get_prefill_wrapper( ) else: # NVFP4 KV cache requires the trtllm-gen backend inside - # the wrapper on SM100; SM8x and SM12x read it with fa2. + # the wrapper on SM100; SM8x, SM90 and SM12x read it with fa2. backend = "trtllm-gen" if self.nvfp4_trtllm else "auto" self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper( self._get_workspace_buffer(), @@ -1337,7 +1335,7 @@ def _get_decode_wrapper(self, batch_size: int, use_cudagraph: bool = False): paged_kv_indices = None paged_kv_last_page_len = None # NVFP4 KV cache requires the trtllm-gen backend inside - # the wrapper on SM100; SM8x and SM12x read it with fa2. + # the wrapper on SM100; SM8x, SM90 and SM12x read it with fa2. backend = "trtllm-gen" if self.nvfp4_trtllm else "auto" decode_wrapper = BatchDecodeWithPagedKVCacheWrapper( self._get_workspace_buffer(), From e3a29aa9b33e894ee576d3d092ae7605fb4e7721 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Thu, 8 Oct 2026 15:32:07 -0400 Subject: [PATCH 126/624] [Bugfix][MoE] Size the permute scratch from its inputs (#60447) Signed-off-by: mgoin Co-authored-by: Claude --- .../benchmark_moe_permute_unpermute.py | 8 - tests/kernels/moe/test_moe.py | 30 +--- .../kernels/moe/test_moe_permute_unpermute.py | 68 ++++---- .../layers/fused_moe/experts/cutlass_moe.py | 35 ++-- .../fused_moe/experts/fused_humming_moe.py | 22 +-- .../layers/fused_moe/moe_permute_unpermute.py | 153 +++++++----------- vllm/v1/worker/workspace.py | 3 +- 7 files changed, 128 insertions(+), 191 deletions(-) diff --git a/benchmarks/kernels/benchmark_moe_permute_unpermute.py b/benchmarks/kernels/benchmark_moe_permute_unpermute.py index 6867274ebe28..f5aa21fb4bd5 100644 --- a/benchmarks/kernels/benchmark_moe_permute_unpermute.py +++ b/benchmarks/kernels/benchmark_moe_permute_unpermute.py @@ -56,13 +56,9 @@ def benchmark_permute( qhidden_states, input_gating, topk, False ) scratch = MoEPermuteScratch( - max_num_tokens=num_tokens, - topk=topk, num_experts=num_experts, num_local_experts=num_experts, device=qhidden_states.device, - hidden_size=hidden_size, - hidden_dtype=qhidden_states.dtype, ) def prepare(i: int): @@ -135,13 +131,9 @@ def benchmark_unpermute( qhidden_states, input_gating, topk, False ) scratch = MoEPermuteScratch( - max_num_tokens=num_tokens, - topk=topk, num_experts=num_experts, num_local_experts=num_experts, device=qhidden_states.device, - hidden_size=hidden_size, - hidden_dtype=qhidden_states.dtype, ) def prepare(): diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 8d6deac4f025..0f786deb5498 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -1437,6 +1437,7 @@ def test_humming_global_valid_shape_m( def test_humming_permute_scratch_is_shared_by_config( monkeypatch: pytest.MonkeyPatch, ): + """Layers share one scratch whose key is independent of batch shape.""" from types import SimpleNamespace from unittest.mock import Mock @@ -1447,34 +1448,19 @@ def test_humming_permute_scratch_is_shared_by_config( manager = workspace.WorkspaceManager(torch.device("cpu")) monkeypatch.setattr(workspace, "_manager", manager) - scratch_topk6 = Mock() - scratch_topk1 = Mock() - scratch_type = Mock(side_effect=[scratch_topk6, scratch_topk1]) + scratch = Mock() + scratch_type = Mock(side_effect=[scratch]) monkeypatch.setattr(humming, "moe_permute_unpermute_supported", lambda: True) monkeypatch.setattr(permute, "MoEPermuteScratch", scratch_type) moe_config = make_dummy_moe_config(max_num_tokens=512, experts_per_token=6) - moe_config.moe_parallel_config.dp_size = 2 experts = SimpleNamespace(moe_config=moe_config) other_layer = SimpleNamespace(moe_config=moe_config) - assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6 - assert ( - humming.HummingExpertsBase._get_permute_scratch(other_layer, 6) is scratch_topk6 - ) - assert humming.HummingExpertsBase._get_permute_scratch(experts, 1) is scratch_topk1 - assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6 - - assert scratch_type.call_count == 2 - first_call, second_call = scratch_type.call_args_list - assert first_call.kwargs["max_num_tokens"] == 1024 - assert first_call.kwargs["topk"] == 6 - assert second_call.kwargs["max_num_tokens"] == 6144 - assert second_call.kwargs["topk"] == 1 - + assert humming.HummingExpertsBase._get_permute_scratch(experts) is scratch + assert humming.HummingExpertsBase._get_permute_scratch(other_layer) is scratch manager.lock() - assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6 - with pytest.raises(AssertionError, match="was not allocated during warmup"): - humming.HummingExpertsBase._get_permute_scratch(experts, 2) + assert humming.HummingExpertsBase._get_permute_scratch(experts) is scratch + assert scratch_type.call_count == 1 def test_humming_delegates_to_instance_activation(): @@ -1531,7 +1517,7 @@ def test_humming_grouped_apply_forwards_valid_prefix( num_experts=2, estimate_local_valid_shape_m=lambda _: 6, prepare_buffers=lambda *_: buffers, - _get_permute_scratch=lambda _, *, indices_only=False: object(), + _get_permute_scratch=lambda: object(), process_input=Mock( side_effect=lambda _, **kwargs: ( kwargs["inputs"], diff --git a/tests/kernels/moe/test_moe_permute_unpermute.py b/tests/kernels/moe/test_moe_permute_unpermute.py index 76f4b57cf07f..0ab5760c3590 100644 --- a/tests/kernels/moe/test_moe_permute_unpermute.py +++ b/tests/kernels/moe/test_moe_permute_unpermute.py @@ -169,13 +169,9 @@ def test_moe_permute_unpermute( scratch = None if use_scratch: scratch = get_moe_permute_scratch( - max_num_tokens=n_token, - topk=topk, num_experts=n_expert, num_local_experts=n_local_expert, device=hidden_states.device, - hidden_size=n_hidden, - hidden_dtype=dtype, ) ( @@ -247,13 +243,9 @@ def test_moe_permute_reuses_scratch_buffers( _, topk_ids, _ = fused_topk(hidden_states, gating_output, topk, False) scratch_config = dict( - max_num_tokens=n_token, - topk=topk, num_experts=n_expert, num_local_experts=n_expert, device=hidden_states.device, - hidden_size=n_hidden, - hidden_dtype=hidden_states.dtype, ) scratch = get_moe_permute_scratch(**scratch_config) @@ -300,15 +292,10 @@ def test_moe_permute_reuses_scratch_buffers( permuted_idx_2, ) = second - torch.testing.assert_close(permuted_hidden_states_1, permuted_hidden_states_2) torch.testing.assert_close(expert_first_token_offset_1, expert_first_token_offset_2) torch.testing.assert_close(inv_permuted_idx_1, inv_permuted_idx_2) torch.testing.assert_close(permuted_idx_1, permuted_idx_2) - assert ( - permuted_hidden_states_1.untyped_storage().data_ptr() - == permuted_hidden_states_2.untyped_storage().data_ptr() - ) assert ( expert_first_token_offset_1.untyped_storage().data_ptr() == expert_first_token_offset_2.untyped_storage().data_ptr() @@ -330,14 +317,11 @@ def test_moe_permute_scratch_reused_across_graph_sizes(workspace_init) -> None: device = torch.device("cuda:0") scratch = get_moe_permute_scratch( - max_num_tokens=128, - topk=6, num_experts=16, num_local_experts=16, device=device, - hidden_size=2048, - hidden_dtype=torch.bfloat16, ) + scratch.reserve(128 * 6) current_workspace_manager().lock() runs = [] for n_token in (128, 1, 33): @@ -372,13 +356,9 @@ def test_moe_permute_scratch_isolated_across_execution_slots(monkeypatch) -> Non ubatch = 0 monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: ubatch) scratch_config = dict( - max_num_tokens=64, - topk=6, num_experts=16, num_local_experts=16, device=device, - hidden_size=2048, - hidden_dtype=torch.bfloat16, ) runs = [] for i, n_token in enumerate((1, 7, 33, 64)): @@ -396,7 +376,7 @@ def test_moe_permute_scratch_isolated_across_execution_slots(monkeypatch) -> Non runs.append((stream, graph, hidden, topk_ids, output, scratch)) torch.cuda.current_stream().wait_stream(stream) - assert len({run[-1].permuted_hidden_states.data_ptr() for run in runs}) == 4 + assert len({run[-1].permuted_idx.data_ptr() for run in runs}) == 4 manager.lock() for _ in range(3): for i, (stream, graph, hidden, topk_ids, _, scratch) in enumerate(runs): @@ -423,8 +403,6 @@ def test_moe_permute_scratch_without_manager(monkeypatch) -> None: monkeypatch.setattr(workspace, "_manager", None) config = dict( - max_num_tokens=4, - topk=2, num_experts=4, num_local_experts=4, device=torch.device("cuda"), @@ -432,6 +410,8 @@ def test_moe_permute_scratch_without_manager(monkeypatch) -> None: first = get_moe_permute_scratch(**config) second = get_moe_permute_scratch(**config) + first.reserve(8) + second.reserve(8) assert ( first.token_expert_indices.data_ptr() != second.token_expert_indices.data_ptr() ) @@ -450,13 +430,9 @@ def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None: topk_ids = torch.tensor([[0], [-1], [1], [4], [2]], device="cuda") expert_map = torch.tensor([0, 1, -1, -1], dtype=torch.int32, device="cuda") scratch = MoEPermuteScratch( - max_num_tokens=5, - topk=1, num_experts=4, num_local_experts=2, device=hidden_states.device, - hidden_size=16, - hidden_dtype=hidden_states.dtype, ) permuted, _, expert_offsets, inverse, _ = moe_permute( @@ -498,3 +474,39 @@ def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None: expected_offsets = torch.tensor([0, 5, 5], dtype=torch.int64, device="cuda") torch.testing.assert_close(indices.flatten(), expected_inverse) torch.testing.assert_close(expert_offsets, expected_offsets) + + +def test_moe_permute_scratch_grows_until_workspace_lock(workspace_init) -> None: + """Capacity follows the inputs seen in warmup, whatever their topk or size.""" + if not moe_permute_unpermute_supported(): + pytest.skip("moe_permute_unpermute is not supported on this platform.") + + device = torch.device("cuda:0") + scratch = get_moe_permute_scratch( + num_experts=16, num_local_experts=16, device=device + ) + + def run(n_token: int, topk: int, dtype: torch.dtype) -> None: + hidden = torch.randn((n_token, 64), device=device).to(dtype) + topk_ids = torch.randint(16, (n_token, topk), device=device) + actual = moe_permute(hidden, None, topk_ids, 16, scratch=scratch) + expected = moe_permute(hidden, None, topk_ids, 16) + for a, e in zip(actual, expected): + if a is not None: + torch.testing.assert_close(a.float(), e.float()) + + run(8, 6, torch.bfloat16) + run(64, 6, torch.bfloat16) + run(384, 1, torch.float8_e4m3fn) + assert scratch.max_expanded_rows == 384 + + current_workspace_manager().lock() + run(64, 6, torch.bfloat16) + with pytest.raises(AssertionError, match="Scratch growth is not allowed"): + run(65, 6, torch.bfloat16) + + # Experts rebuilt after the lock (e.g. a weight reload) keep the capacity. + rebuilt = get_moe_permute_scratch( + num_experts=16, num_local_experts=16, device=device + ) + assert rebuilt.max_expanded_rows == 384 diff --git a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py index a03651c0a93a..6f9a62f282ce 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py @@ -20,6 +20,7 @@ ) from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import ( MoEPermuteScratch, + get_moe_permute_scratch, moe_permute, moe_permute_unpermute_supported, moe_unpermute, @@ -304,7 +305,6 @@ def __init__( self.ab_strides2 = ab_strides2 self.c_strides1 = c_strides1 self.c_strides2 = ab_strides1_c_strides2 - self._permute_scratch: MoEPermuteScratch | None = None @staticmethod def _supports_current_device() -> bool: @@ -335,15 +335,13 @@ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: return TopKWeightAndReduceDelegate() def _get_permute_scratch(self) -> MoEPermuteScratch | None: - if self._permute_scratch is None and moe_permute_unpermute_supported(): - self._permute_scratch = MoEPermuteScratch( - max_num_tokens=self.moe_config.max_num_tokens, - topk=self.moe_config.experts_per_token, - num_experts=self.moe_config.num_experts, - num_local_experts=self.moe_config.num_local_experts, - device=torch.device(self.moe_config.device), - ) - return self._permute_scratch + if not moe_permute_unpermute_supported(): + return None + return get_moe_permute_scratch( + num_experts=self.moe_config.num_experts, + num_local_experts=self.moe_config.num_local_experts, + device=torch.device(self.moe_config.device), + ) def apply( self, @@ -1296,7 +1294,6 @@ def __init__( self.s_strides2[:, 0] = k self.group_size = group_size - self._permute_scratch: MoEPermuteScratch | None = None @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: @@ -1357,15 +1354,13 @@ def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype: return self.out_dtype if self.out_dtype is not None else act_dtype def _get_permute_scratch(self) -> MoEPermuteScratch | None: - if self._permute_scratch is None and moe_permute_unpermute_supported(): - self._permute_scratch = MoEPermuteScratch( - max_num_tokens=self.moe_config.max_num_tokens, - topk=self.moe_config.experts_per_token, - num_experts=self.moe_config.num_experts, - num_local_experts=self.moe_config.num_local_experts, - device=torch.device(self.moe_config.device), - ) - return self._permute_scratch + if not moe_permute_unpermute_supported(): + return None + return get_moe_permute_scratch( + num_experts=self.moe_config.num_experts, + num_local_experts=self.moe_config.num_local_experts, + device=torch.device(self.moe_config.device), + ) def workspace_shapes( self, diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py index a6d44c3141f7..66e03229a9e1 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py @@ -258,25 +258,14 @@ def humming_forward( **kwargs, ) - def _get_permute_scratch( - self, topk: int, indices_only: bool = False - ) -> MoEPermuteScratch | None: + def _get_permute_scratch(self) -> MoEPermuteScratch | None: if not moe_permute_unpermute_supported(): return None - max_expanded_rows = ( - self.moe_config.max_num_tokens - * self.moe_config.dp_size - * self.moe_config.experts_per_token - ) return get_moe_permute_scratch( - max_num_tokens=math.ceil(max_expanded_rows / topk), - topk=topk, num_experts=self.moe_config.num_experts, num_local_experts=self.moe_config.num_local_experts, device=torch.device(self.moe_config.device), - hidden_size=None if indices_only else self.moe_config.hidden_dim, - hidden_dtype=None if indices_only else self.moe_config.in_dtype, ) def get_global_valid_shape_m(self, topk_ids: torch.Tensor): @@ -911,14 +900,18 @@ def apply( ) buffers["output"] = output + scratch = self._get_permute_scratch() if a1q_scale is None: - scratch = self._get_permute_scratch(topk_ids.size(1), indices_only=True) assert scratch is not None scatter_metadata = moe_prepare_scatter(topk_ids, expert_map, scratch) expert_offsets, scatter_idx = scatter_metadata inv_perm = scatter_idx.flatten() num_valid_tokens = expert_offsets[-1:] if expert_map is not None else None else: + # Prequantized inputs skip w13 quantization, so its buffer is free + # to hold the permuted hidden states. + permuted_hidden_states = buffers["quanted_gate_up_input"] + assert permuted_hidden_states.dtype == hidden_states.dtype hidden_states, a1q_scale, expert_offsets, inv_perm, _ = moe_permute( hidden_states=hidden_states, a1q_scale=a1q_scale, @@ -926,7 +919,8 @@ def apply( n_expert=global_num_experts, n_local_expert=self.num_experts, expert_map=expert_map, - scratch=self._get_permute_scratch(topk_ids.size(1)), + permuted_hidden_states=permuted_hidden_states, + scratch=scratch, ) scatter_idx = None num_valid_tokens = None diff --git a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py index d83152fc696a..b4bf3c89163d 100644 --- a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py +++ b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py @@ -13,90 +13,70 @@ @dataclass class MoEPermuteScratch: - # Reused metadata buffers for repeated grouped-MoE permutes. - max_num_tokens: int - topk: int + """Reused metadata buffers for repeated grouped-MoE permutes. + + Capacity follows the largest input seen, like the shared workspace: it grows + during warmup and is fixed once the workspace manager is locked. + """ + num_experts: int num_local_experts: int device: torch.device - hidden_size: int | None = None - hidden_dtype: torch.dtype | None = None + max_expanded_rows: int = field(init=False, default=0) token_expert_indices: torch.Tensor = field(init=False) expert_first_token_offset: torch.Tensor = field(init=False) permuted_idx: torch.Tensor = field(init=False) inv_permuted_idx: torch.Tensor = field(init=False) - permuted_hidden_states: torch.Tensor | None = field(init=False, default=None) sort_workspace: torch.Tensor = field(init=False) permuted_experts_id: torch.Tensor = field(init=False) sorted_row_idx: torch.Tensor = field(init=False) topk_ids_int32: torch.Tensor = field(init=False) topk_ids_for_sort: torch.Tensor = field(init=False) - max_expanded_rows: int = field(init=False) def __post_init__(self) -> None: - assert self.max_num_tokens > 0 - assert self.topk > 0 assert self.num_experts > 0 assert self.num_local_experts > 0 - if self.hidden_size is None: - assert self.hidden_dtype is None - else: - assert self.hidden_dtype is not None - - self.max_expanded_rows = self.max_num_tokens * self.topk - self.token_expert_indices = torch.arange( - self.max_expanded_rows, dtype=torch.int32, device=self.device - ) self.expert_first_token_offset = torch.empty( self.num_local_experts + 1, dtype=torch.int64, device=self.device ) - self.permuted_idx = torch.empty( - self.max_expanded_rows, dtype=torch.int32, device=self.device - ) - self.inv_permuted_idx = torch.empty( - self.max_expanded_rows, dtype=torch.int32, device=self.device - ) - if self.hidden_size is not None: - hidden_numel = self.max_expanded_rows * self.hidden_size - self.permuted_hidden_states = torch.empty( - hidden_numel, dtype=self.hidden_dtype, device=self.device + # torch.device("cuda") in config, after initialized, + # will be changed to cuda:{index}, so we need to refresh here. + self.device = self.expert_first_token_offset.device + + def reserve(self, expanded_rows: int) -> None: + """Ensure capacity for ``expanded_rows`` (num_tokens * topk) rows.""" + if expanded_rows <= self.max_expanded_rows: + return + if ( + is_workspace_manager_initialized() + and current_workspace_manager().is_locked() + ): + raise AssertionError( + f"Workspace is locked but MoE permute scratch requires " + f"{expanded_rows} rows, current capacity is " + f"{self.max_expanded_rows} rows. Scratch growth is not allowed " + "after locking." ) - self.permuted_experts_id = torch.empty( - self.max_expanded_rows, dtype=torch.int32, device=self.device - ) - self.sorted_row_idx = torch.empty( - self.max_expanded_rows, dtype=torch.int32, device=self.device - ) - self.topk_ids_int32 = torch.empty( - self.max_expanded_rows, dtype=torch.int32, device=self.device - ) - self.topk_ids_for_sort = torch.empty( - self.max_expanded_rows, dtype=torch.int32, device=self.device + + def empty_rows() -> torch.Tensor: + return torch.empty(expanded_rows, dtype=torch.int32, device=self.device) + + self.token_expert_indices = torch.arange( + expanded_rows, dtype=torch.int32, device=self.device ) + self.permuted_idx = empty_rows() + self.inv_permuted_idx = empty_rows() + self.permuted_experts_id = empty_rows() + self.sorted_row_idx = empty_rows() + self.topk_ids_int32 = empty_rows() + self.topk_ids_for_sort = empty_rows() sorter_size = torch.ops._moe_C.moe_permute_sort_workspace_size( - self.max_expanded_rows, self.num_experts + expanded_rows, self.num_experts ) self.sort_workspace = torch.empty( sorter_size, dtype=torch.int8, device=self.device ) - # torch.device("cuda") in config, after initialized, - # will be changed to cuda:{index}, so we need to refresh here. - self.device = self.token_expert_indices.device - - def validate(self, hidden_states: torch.Tensor, topk_ids: torch.Tensor) -> None: - n_token, n_hidden = hidden_states.shape - assert hidden_states.device == self.device - assert topk_ids.device == self.device - assert n_token <= self.max_num_tokens - assert topk_ids.size(1) == self.topk - assert topk_ids.size(0) == n_token - if self.hidden_size is not None: - assert n_hidden == self.hidden_size - assert hidden_states.dtype == self.hidden_dtype - assert self.permuted_hidden_states is not None - - def token_expert_indices_view(self, n_token: int) -> torch.Tensor: - return self.token_expert_indices[: n_token * self.topk].view(n_token, self.topk) + self.max_expanded_rows = expanded_rows def prepare_topk_ids(self, topk_ids: torch.Tensor) -> torch.Tensor: if topk_ids.dtype == torch.int32: @@ -114,17 +94,16 @@ def moe_prepare_scatter( ) -> tuple[torch.Tensor, torch.Tensor]: """Generate expert offsets and shared scatter/unpermute destination indices.""" assert topk_ids.device == scratch.device - assert topk_ids.size(0) <= scratch.max_num_tokens - assert topk_ids.size(1) == scratch.topk n_token, topk = topk_ids.shape expanded_rows = topk_ids.numel() + scratch.reserve(expanded_rows) inverse = scratch.inv_permuted_idx[:expanded_rows].view(n_token, topk) if expanded_rows == 0: scratch.expert_first_token_offset.zero_() return scratch.expert_first_token_offset, inverse torch.ops._moe_C.moe_prepare_scatter( scratch.prepare_topk_ids(topk_ids), - scratch.token_expert_indices_view(n_token), + scratch.token_expert_indices[:expanded_rows].view(n_token, topk), expert_map, scratch.num_experts, scratch.num_local_experts, @@ -140,13 +119,9 @@ def moe_prepare_scatter( def get_moe_permute_scratch( *, - max_num_tokens: int, - topk: int, num_experts: int, num_local_experts: int, device: torch.device, - hidden_size: int | None = None, - hidden_dtype: torch.dtype | None = None, ) -> MoEPermuteScratch: """Share scratch across sequential layers in the current ubatch and lane. @@ -155,28 +130,15 @@ def get_moe_permute_scratch( def create_scratch() -> MoEPermuteScratch: return MoEPermuteScratch( - max_num_tokens=max_num_tokens, - topk=topk, num_experts=num_experts, num_local_experts=num_local_experts, device=device, - hidden_size=hidden_size, - hidden_dtype=hidden_dtype, ) if not is_workspace_manager_initialized(): return create_scratch() - key = ( - MoEPermuteScratch, - max_num_tokens, - topk, - num_experts, - num_local_experts, - device, - hidden_size, - hidden_dtype, - ) + key = (MoEPermuteScratch, num_experts, num_local_experts, device) return current_workspace_manager().get_persistent_resource(key, create_scratch) @@ -204,9 +166,8 @@ def moe_permute( parallel shard. permuted_hidden_states (Optional[torch.Tensor]): Optional output tensor. If None, the output tensor will be created in this function. - scratch (Optional[MoEPermuteScratch]): Optional preallocated scratch - buffers. Validated against hidden_states and topk_ids when given, - otherwise the buffers are allocated in this function. + scratch (Optional[MoEPermuteScratch]): Optional reusable metadata + buffers. Otherwise they are allocated in this function. Returns: - permuted_hidden_states (torch.Tensor): permuted activation. @@ -227,20 +188,11 @@ def moe_permute( if n_local_expert == -1: n_local_expert = n_expert if permuted_hidden_states is None: - if scratch is None: - permuted_hidden_states = torch.empty( - (permuted_row_size, n_hidden), - dtype=hidden_states.dtype, - device=hidden_states.device, - ) - else: - scratch.validate(hidden_states, topk_ids) - hidden_numel = permuted_row_size * n_hidden - scratch_hidden_states = scratch.permuted_hidden_states - assert scratch_hidden_states is not None - permuted_hidden_states = scratch_hidden_states[:hidden_numel].view( - permuted_row_size, n_hidden - ) + permuted_hidden_states = torch.empty( + (permuted_row_size, n_hidden), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) assert permuted_hidden_states.size() == (permuted_row_size, n_hidden), ( f"Expected permuted hidden states to be {(permuted_row_size, n_hidden)}" f" but got {permuted_hidden_states.size()}" @@ -278,10 +230,15 @@ def moe_permute( permuted_idx, ) else: - scratch.validate(hidden_states, topk_ids) + assert hidden_states.device == scratch.device + assert topk_ids.device == scratch.device + assert topk_ids.size(0) == n_token assert n_expert == scratch.num_experts assert n_local_expert == scratch.num_local_experts - token_expert_indices = scratch.token_expert_indices_view(n_token) + scratch.reserve(permuted_row_size) + token_expert_indices = scratch.token_expert_indices[:permuted_row_size].view( + n_token, topk + ) expert_first_token_offset = scratch.expert_first_token_offset permuted_idx = scratch.permuted_idx[:permuted_row_size] permuted_idx.fill_(permuted_row_size) diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py index e69e55bd8216..a05534281c11 100644 --- a/vllm/v1/worker/workspace.py +++ b/vllm/v1/worker/workspace.py @@ -139,7 +139,8 @@ def get_persistent_resource(self, key: Hashable, factory: Callable[[], T]) -> T: AssertionError. Keys must identify the resource type and configuration. Callers must not - resize or replace tensor storage held by a cached resource. Uses of the + resize or replace tensor storage held by a cached resource once the + manager is locked. Uses of the same resource must not overlap; include the CUDA stream in the key when streams can execute concurrently. From bab0ac61f8dae8020df6725c90474bdd647f1c40 Mon Sep 17 00:00:00 2001 From: dingdangmao <2756823972@qq.com> Date: Fri, 9 Oct 2026 03:32:46 +0800 Subject: [PATCH 127/624] [Bugfix][Frontend] Forward the reasoning wiring to the engine for batch chat completions (#59263) Signed-off-by: apex-mochen <2756823972@qq.com> Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../test_batched_chat_completions.py | 102 ++++++++++++++++++ .../openai/chat_completion/batch_serving.py | 13 ++- .../openai/chat_completion/serving.py | 37 ++++--- 3 files changed, 139 insertions(+), 13 deletions(-) diff --git a/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py b/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py index 946615f97105..41c74bb5cef8 100644 --- a/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py +++ b/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py @@ -10,6 +10,7 @@ from tests.entrypoints.openai.chat_completion.test_serving_chat import ( BASE_MODEL_PATHS, + CHAT_TEMPLATE, MockHFConfig, MockModelConfig, _build_renderer, @@ -24,10 +25,13 @@ ) from vllm.entrypoints.openai.chat_completion.protocol import ( BatchChatCompletionRequest, + ChatCompletionRequest, ) +from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.exceptions import VLLMValidationError from vllm.outputs import CompletionOutput, RequestOutput +from vllm.reasoning import ReasoningParserManager from vllm.renderers.online_renderer import OnlineRenderer from vllm.v1.engine.async_llm import AsyncLLM @@ -454,3 +458,101 @@ def parse(self, model_output, **kwargs): ) assert prefixes == [[1], [2]] + + +REASONING_PARSER = "kimi_k3" + + +@pytest.mark.asyncio +@pytest.mark.skip_global_cleanup +async def test_batched_reasoning_wiring_matches_single_chat() -> None: + """`/v1/chat/completions/batch` must hand the engine the same reasoning + wiring as `/v1/chat/completions`. + + Both `reasoning_ended` and `reasoning_parser_kwargs` are carried on the + request into `EngineCoreRequest`, where + `StructuredOutputManager._get_reasoner()` rebuilds the request-local + reasoning parser from the forwarded `chat_template_kwargs` -- "so the + structured-output gate observes the same template kwargs used by the + frontend". A divergence makes a template-kwarg-driven parser fall back to + its default (`thinking` on) for a request that disabled thinking, so + `is_reasoning_end()` keeps reporting False, the structured-output gate + never opens, and the grammar is silently never applied. + + Both frontends run over one and the same model and renderer, so the only + thing that can differ between the two calls is the serving code itself. + """ + mock_model_config = MockModelConfig() + shared_renderer = _build_renderer(mock_model_config) + + def _build_serving(cls): + mock_engine = MagicMock(spec=AsyncLLM) + mock_engine.model_config = mock_model_config + mock_engine.errored = False + mock_engine.input_processor = MagicMock() + mock_engine.renderer = shared_renderer + mock_engine.generate.side_effect = lambda *args, **kwargs: _generator(0, "ok") + + models = OpenAIServingModels(mock_engine, BASE_MODEL_PATHS) + online_renderer = OnlineRenderer( + model_config=mock_model_config, + renderer=mock_engine.renderer, + request_logger=None, + chat_template=CHAT_TEMPLATE, + chat_template_content_format="auto", + reasoning_parser=REASONING_PARSER, + ) + return mock_engine, cls( + mock_engine, + models, + response_role="assistant", + online_renderer=online_renderer, + chat_template=CHAT_TEMPLATE, + chat_template_content_format="auto", + request_logger=None, + reasoning_parser=REASONING_PARSER, + ) + + conversation = [{"role": "user", "content": "What is the capital of France?"}] + chat_template_kwargs = {"enable_thinking": False} + + single_engine, single_serving = _build_serving(OpenAIServingChat) + await single_serving.create_chat_completion( + ChatCompletionRequest( + model=MOCK_MODEL_NAME, + messages=conversation, + chat_template_kwargs=chat_template_kwargs, + ) + ) + single_call = single_engine.generate.call_args_list[-1] + + batch_engine, batch_serving = _build_serving(OpenAIServingChatBatch) + await batch_serving.create_batch_chat_completion( + BatchChatCompletionRequest( + model=MOCK_MODEL_NAME, + messages=[conversation], + chat_template_kwargs=chat_template_kwargs, + ) + ) + batch_call = batch_engine.generate.call_args_list[-1] + + # Guard against a vacuous comparison: the reference path must be wired up. + assert single_call.kwargs["reasoning_ended"] is not None + assert single_call.kwargs["reasoning_parser_kwargs"] is not None + + assert batch_call.kwargs["reasoning_ended"] == single_call.kwargs["reasoning_ended"] + assert ( + batch_call.kwargs["reasoning_parser_kwargs"] + == single_call.kwargs["reasoning_parser_kwargs"] + ) + + # ...and the forwarded kwargs must actually turn thinking off for the + # engine-side parser, i.e. the structured-output gate opens right away. + engine_chat_template_kwargs = batch_call.kwargs["reasoning_parser_kwargs"][ + "chat_template_kwargs" + ] + reasoner = ReasoningParserManager.get_reasoning_parser(REASONING_PARSER)( + tokenizer=batch_engine.renderer.tokenizer, + chat_template_kwargs=engine_chat_template_kwargs, + ) + assert reasoner.is_reasoning_end([]) is True diff --git a/vllm/entrypoints/openai/chat_completion/batch_serving.py b/vllm/entrypoints/openai/chat_completion/batch_serving.py index e630b97f182a..d3ea1a1ac3de 100644 --- a/vllm/entrypoints/openai/chat_completion/batch_serving.py +++ b/vllm/entrypoints/openai/chat_completion/batch_serving.py @@ -169,6 +169,12 @@ async def create_batch_chat_completion( else await self._get_trace_headers(raw_request.headers) ) session_id = self._get_session_id(single_request, raw_request) + chat_template_kwargs = self._effective_chat_template_kwargs(single_request) + parser = ( + self._make_parser(single_request, tokenizer, chat_template_kwargs) + if self.parser_cls is not None + else None + ) generators.append( self.engine_client.generate( engine_prompt, @@ -179,7 +185,12 @@ async def create_batch_chat_completion( priority=request.priority, data_parallel_rank=data_parallel_rank, session_id=session_id, - reasoning_ended=None, + **self._engine_reasoning_kwargs( + single_request, + parser, + chat_template_kwargs, + self._extract_prompt_components(engine_prompt).token_ids, + ), ) ) diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 9b0737c25d47..69b46ba67f61 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -233,6 +233,28 @@ def _engine_chat_template_kwargs( """ return chat_template_kwargs + def _engine_reasoning_kwargs( + self, + request: ChatCompletionRequest, + parser: Parser | None, + chat_template_kwargs: dict[str, Any], + prompt_token_ids: list[int] | None, + ) -> dict[str, Any]: + """`reasoning_ended` and `reasoning_parser_kwargs` for `generate()`.""" + reasoning_parser_kwargs = None + if parser is not None and parser.reasoning_parser is not None: + reasoning_parser_kwargs = { + "chat_template_kwargs": self._engine_chat_template_kwargs( + chat_template_kwargs + ), + } + return { + "reasoning_ended": request.resolve_reasoning_ended( + parser, prompt_token_ids or [] + ), + "reasoning_parser_kwargs": reasoning_parser_kwargs, + } + async def render_chat_request( self, request: ChatCompletionRequest, @@ -364,10 +386,6 @@ async def _create_chat_completion( session_id=session_id, ) else: - reasoning_ended = request.resolve_reasoning_ended( - parser, prompt_token_ids or [] - ) - generator = self.engine_client.generate( engine_input, sampling_params, @@ -377,14 +395,9 @@ async def _create_chat_completion( priority=self._get_priority(request, raw_request), data_parallel_rank=data_parallel_rank, session_id=session_id, - reasoning_ended=reasoning_ended, - reasoning_parser_kwargs={ - "chat_template_kwargs": self._engine_chat_template_kwargs( - chat_template_kwargs - ), - } - if parser is not None and parser.reasoning_parser is not None - else None, + **self._engine_reasoning_kwargs( + request, parser, chat_template_kwargs, prompt_token_ids + ), ) generators.append(generator) From c23ca06b6d944579ad9d220297441236755b4744 Mon Sep 17 00:00:00 2001 From: simon-veitner-redhat Date: Thu, 8 Oct 2026 21:38:33 +0200 Subject: [PATCH 128/624] [Bugfix][Structured Output] Feed the token that implicitly ends reasoning to the grammar (#59627) Signed-off-by: Simon Veitner Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Misha Goin Co-authored-by: Yifan Zong <263738792+yzong-rh@users.noreply.github.com> --- tests/parser/engine/test_deepseek_v4.py | 4 +- tests/parser/engine/test_parser_engine.py | 45 ++++++++--- .../spec_decode/test_mtp_structured_output.py | 7 +- .../test_structured_output_manager.py | 48 +++++++++++- vllm/parser/engine/adapters.py | 6 +- vllm/parser/engine/parser_engine.py | 16 +++- vllm/v1/structured_output/__init__.py | 78 +++++++++++-------- 7 files changed, 151 insertions(+), 53 deletions(-) diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index 2cf3a8bb85f3..f5b66d9dc1cb 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -643,8 +643,8 @@ def test_reasoning_end_token_ids_drop_multi_token_tool_opener( # token-ID fast path; the implicit tool-call exit still goes through # the text lexer. assert thinking_parser.reasoning_end_token_ids == {_THINK_END_ID} - assert thinking_parser.find_reasoning_end_offset([7, _THINK_END_ID, 8]) == 1 - assert thinking_parser.find_reasoning_end_offset([7, 8]) == 2 + assert thinking_parser.find_reasoning_end([7, _THINK_END_ID, 8]) == (1, False) + assert thinking_parser.find_reasoning_end([7, 8]) == (2, False) def test_streaming_reasoning_implicit_end(self, thinking_parser): chunks = [ diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py index 8fd7c4d2623f..dc91cef74028 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -28,12 +28,13 @@ from vllm.parser.abstract_parser import DelegatingParser from vllm.parser.engine.adapters import make_adapters from vllm.parser.engine.events import EventType, SemanticEvent -from vllm.parser.engine.parser_engine import ParserEngine +from vllm.parser.engine.parser_engine import ParserEngine, ReasoningEnd from vllm.parser.engine.parser_engine_config import ( ParserEngineConfig, ParserState, Transition, ) +from vllm.parser.glm47_moe import glm47_moe_config from vllm.parser.parser_manager import ParserManager # ── Shared test configs ────────────────────────────────────────────── @@ -158,7 +159,10 @@ def test_every_reasoning_end_exit_contributes(self): (EventType.REASONING_END, EventType.TOOL_CALL_START), ), ) - assert _make_engine(cfg).reasoning_end_token_ids == {201, 202} + engine = _make_engine(cfg) + assert engine.reasoning_end_token_ids == {201, 202} + assert engine.find_reasoning_end([5, 201]) == ReasoningEnd(1, False) + assert engine.find_reasoning_end([5, 202]) == ReasoningEnd(1, True) def test_transition_staying_in_reasoning_is_ignored(self): cfg = _with_reasoning_exits(("THINK_START", ParserState.REASONING, ())) @@ -182,17 +186,40 @@ def test_unresolved_think_end_fails_closed(self): def test_config_without_reasoning_has_empty_set(self): assert _make_engine(_hermes_config()).reasoning_end_token_ids == frozenset() - def test_find_reasoning_end_offset_returns_first_match(self): + def test_find_reasoning_end_returns_first_match(self): engine = _make_engine() - assert engine.find_reasoning_end_offset([5, 201, 6, 201]) == 1 - assert engine.find_reasoning_end_offset([5, 6]) == 2 - assert engine.find_reasoning_end_offset([]) == 0 + assert engine.find_reasoning_end([5, 201, 6, 201]) == ReasoningEnd(1, False) + assert engine.find_reasoning_end([5, 6]) == ReasoningEnd(2, False) + assert engine.find_reasoning_end([]) == ReasoningEnd(0, False) # Rejected-draft placeholders never match. - assert engine.find_reasoning_end_offset([-1, 201]) == 1 + assert engine.find_reasoning_end([-1, 201]) == ReasoningEnd(1, False) - def test_find_reasoning_end_offset_with_empty_set_returns_none(self): + def test_find_reasoning_end_with_empty_set_returns_none(self): engine = _make_engine(_hermes_config()) - assert engine.find_reasoning_end_offset([201]) is None + assert engine.find_reasoning_end([201]) is None + + +@pytest.mark.parametrize( + ("config", "implicit_token"), + [ + pytest.param(glm47_moe_config(), "", id="glm47_moe"), + ], +) +def test_implicit_reasoning_end_of_parsers( + config: ParserEngineConfig, implicit_token: str +): + literals = sorted(set(config.token_id_terminals.values())) + vocab = {literal: 1000 + i for i, literal in enumerate(literals)} + engine = _make_engine(config, vocab=vocab) + implicit_ids = set() + for token_id in engine.reasoning_end_token_ids: + end = engine.find_reasoning_end([5, token_id, 6]) + assert end is not None and end.offset == 1 + kept = engine.extract_content_ids([5, token_id, 6]) == [token_id, 6] + assert end.implicit == kept + if end.implicit: + implicit_ids.add(token_id) + assert implicit_ids == {vocab[implicit_token]} # ── TestEventsToDelta ──────────────────────────────────────────────── diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py index 5e4312227d7e..8a72765db150 100644 --- a/tests/v1/spec_decode/test_mtp_structured_output.py +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -11,6 +11,7 @@ from vllm.config.model import ModelConfig from vllm.config.speculative import SpeculativeConfig from vllm.parser.engine.adapters import ParserEngineReasoningAdapter +from vllm.parser.engine.parser_engine import ReasoningEnd from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -113,12 +114,12 @@ def __init__(self, end_token_ids): def reasoning_end_token_ids(self): return self._end_token_ids - def find_reasoning_end_offset(self, token_ids): + def find_reasoning_end(self, token_ids): self.windows.append(list(token_ids)) for offset, token in enumerate(token_ids): if token in self._end_token_ids: - return offset - return len(token_ids) + return ReasoningEnd(offset, False) + return ReasoningEnd(len(token_ids), False) def is_reasoning_end(self, input_ids): return any(token in self._end_token_ids for token in input_ids) diff --git a/tests/v1/structured_output/test_structured_output_manager.py b/tests/v1/structured_output/test_structured_output_manager.py index 6082bcd6a7b1..fbceb3bd07e4 100644 --- a/tests/v1/structured_output/test_structured_output_manager.py +++ b/tests/v1/structured_output/test_structured_output_manager.py @@ -11,6 +11,8 @@ from vllm.config.model import ModelConfig from vllm.config.speculative import SpeculativeConfig from vllm.config.structured_outputs import StructuredOutputsBackend +from vllm.parser.engine.adapters import ParserEngineReasoningAdapter +from vllm.parser.engine.parser_engine import ReasoningEnd from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -18,6 +20,7 @@ TOKENIZER = "gpt2" THINK_END = "\n" # reasoning-end marker (single GPT-2 token) +IMPLICIT_THINK_END = "z" # accepted in strings, rejected at the object start EOS = "<|eos|>" # resolved to tokenizer.eos_token_id JSON_SCHEMA = '{"type": "object"}' BACKENDS = ("xgrammar", "guidance") @@ -43,6 +46,7 @@ class FlowCase: prefix: str = "" expected_validated: tuple[str, ...] | None = None reasoning_ended: bool | None = False + engine_reasoner: bool = False xfail_guidance: str | None = None @@ -59,6 +63,23 @@ def is_reasoning_end_streaming(self, input_ids, delta_ids): return self.is_reasoning_end(delta_ids) +class MockEngineReasoner(ParserEngineReasoningAdapter): + """Ends reasoning implicitly on ``marker``, which is also content.""" + + def __init__(self, tokenizer, marker: int): + self.marker = marker + + @property + def reasoning_end_token_ids(self): + return frozenset({self.marker}) + + def find_reasoning_end(self, token_ids): + ids = list(token_ids) + if self.marker in ids: + return ReasoningEnd(ids.index(self.marker), True) + return ReasoningEnd(len(ids), False) + + def _single_token(tokenizer, text: str) -> int: token_ids = tokenizer.encode(text) assert len(token_ids) == 1, (text, token_ids) @@ -331,6 +352,28 @@ def _run_real_flow( ), id="becomes_active_terminates", ), + pytest.param( + FlowCase( + prefix='{"a": "', + raw_drafts=(IMPLICIT_THINK_END,), + expected_row_pattern="UC", + expected_reasoning=True, + expect_terminated=False, + engine_reasoner=True, + ), + id="implicit_end_accepted", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", IMPLICIT_THINK_END, "{", "}"), + expected_validated=(" ",), + expected_row_pattern="UUUUU", + expected_reasoning=False, + expect_terminated=False, + engine_reasoner=True, + ), + id="implicit_end_rejected", + ), ] @@ -344,7 +387,8 @@ def test_real_flow( if backend == "guidance" and case.xfail_guidance: pytest.xfail(case.xfail_guidance) - reasoner_kwargs = {"marker": _single_token(tokenizer, THINK_END)} + marker = IMPLICIT_THINK_END if case.engine_reasoner else THINK_END + reasoner_kwargs = {"marker": _single_token(tokenizer, marker)} manager, request = _build_harness( tokenizer, backend, @@ -352,6 +396,8 @@ def test_real_flow( reasoning_ended=case.reasoning_ended, reasoning_parser_kwargs=reasoner_kwargs, ) + if case.engine_reasoner: + manager.reasoner_cls = MockEngineReasoner raw_drafts = _to_token_ids(tokenizer, case.raw_drafts) expected_texts = ( diff --git a/vllm/parser/engine/adapters.py b/vllm/parser/engine/adapters.py index 4f4aa9f5a1e0..cd7aaa729785 100644 --- a/vllm/parser/engine/adapters.py +++ b/vllm/parser/engine/adapters.py @@ -27,7 +27,7 @@ ChatCompletionRequest, ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest - from vllm.parser.engine.parser_engine import ParserEngine + from vllm.parser.engine.parser_engine import ParserEngine, ReasoningEnd from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.utils import Tool @@ -65,8 +65,8 @@ def _skip_tool_parsing(self) -> Iterator[None]: def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: return self._parser_engine.is_reasoning_end(list(input_ids)) - def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None: - return self._parser_engine.find_reasoning_end_offset(token_ids) + def find_reasoning_end(self, token_ids: Sequence[int]) -> ReasoningEnd | None: + return self._parser_engine.find_reasoning_end(token_ids) def adjust_initial_state_from_prompt(self, prompt_token_ids: Sequence[int]) -> None: self._parser_engine.adjust_initial_state_from_prompt(prompt_token_ids) diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index 527c310c1485..b04fdaa7944d 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -9,7 +9,7 @@ import json from collections.abc import Sequence from functools import cached_property -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, NamedTuple import regex as re @@ -45,6 +45,13 @@ logger = init_logger(__name__) +class ReasoningEnd(NamedTuple): + """Where reasoning ends in a token sequence.""" + + offset: int # index of the first end token, or the sequence length + implicit: bool # the end token is also content, e.g. a tool-call start + + class ToolCallSlot: __slots__ = ( "id", @@ -638,14 +645,15 @@ def _derive_reasoning_end_token_ids( def reasoning_end_token_ids(self) -> frozenset[int]: return self._reasoning_end_token_ids - def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None: + def find_reasoning_end(self, token_ids: Sequence[int]) -> ReasoningEnd | None: + """First reasoning end in `token_ids`, and if its token is content.""" end_ids = self._reasoning_end_token_ids if not end_ids: return None for offset, token_id in enumerate(token_ids): if token_id in end_ids: - return offset - return len(token_ids) + return ReasoningEnd(offset, token_id != self._reasoning_end_token_id) + return ReasoningEnd(len(token_ids), False) def is_reasoning_end(self, input_ids: list[int]) -> bool: config = self.parser_engine_config diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 3b35d39d1447..2163a800a563 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -3,7 +3,7 @@ import multiprocessing from collections.abc import Iterable, Sequence from concurrent.futures import Future, ThreadPoolExecutor -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, NamedTuple from vllm.config import VllmConfig from vllm.logger import init_logger @@ -33,6 +33,17 @@ logger = init_logger(__name__) +class ConstraintBounds(NamedTuple): + """Where tokens start being fed to the grammar and being masked.""" + + grammar_start: int + constraint_start: int + + +_CONSTRAINED_FROM_START = ConstraintBounds(0, 0) +_NO_SPEC_TOKENS = ConstraintBounds(1, 1) + + class StructuredOutputManager: """Engine-level manager for structured output requests.""" @@ -217,32 +228,34 @@ def _async_submit_fill_bitmask( ) -> Future: return self.executor_for_fillmask.submit(self._fill_bitmasks, batch) - def _get_constraint_start( + def _get_constraint_bounds( self, request: "Request", spec_tokens: Sequence[int], spec_tokens_committed: bool = False, - ) -> int: + ) -> ConstraintBounds: """Return the index into `spec_tokens` where tokens should start being constrained by the grammar. Assumes `spec_tokens` are stripped of -1 padding tokens. Returns `len(spec_tokens) + 1` if no token should be constrained after accepting all `spec_tokens`. + This is `constraint_start`; `grammar_start` is one less if reasoning ends + implicitly on a content token (e.g. a tool-call start). `spec_tokens_committed` is True if spec_tokens is already in `request.all_token_ids`. """ if self.enable_in_reasoning: - return 0 + return _CONSTRAINED_FROM_START structured_req = request.structured_output_request assert structured_req is not None if structured_req.reasoning_ended: - return 0 + return _CONSTRAINED_FROM_START reasoner = self._get_reasoner(request) if reasoner is None: - return 0 + return _CONSTRAINED_FROM_START if structured_req.reasoning_ended is None: # This should be removed here, but since `openai_gptoss` @@ -250,30 +263,33 @@ def _get_constraint_start( # After unifying the `openai_gptoss` and non-`openai_gptoss` styles, # it can be removed. if reasoner.is_reasoning_end(request.prompt_token_ids or []): - return 0 + return _CONSTRAINED_FROM_START structured_req.reasoning_ended = False num_spec_tokens = len(spec_tokens) if num_spec_tokens <= 0: - return num_spec_tokens + 1 + return _NO_SPEC_TOKENS - # Use `find_reasoning_end_offset` to find constraint start if supported + # Engine-based parsers also report whether the end token is content. if ( isinstance(reasoner, ParserEngineReasoningAdapter) and reasoner.reasoning_end_token_ids ): - offset = reasoner.find_reasoning_end_offset(spec_tokens) - if offset is not None: - return offset + 1 - - # Fallback to `find_reasoning_end_offset` + end = reasoner.find_reasoning_end(spec_tokens) + if end is not None: + constraint_start = end.offset + 1 + if end.implicit: + return ConstraintBounds(end.offset, constraint_start) + return ConstraintBounds(constraint_start, constraint_start) + + # Fallback: the grammar starts after the end token. # TODO: Build a read-only Sequence view over all_token_ids # instead of copying the entire all_token_ids into input_ids. if spec_tokens_committed: if not spec_tokens or not reasoner.is_reasoning_end_streaming( request.all_token_ids, spec_tokens ): - return num_spec_tokens + 1 + return ConstraintBounds(num_spec_tokens + 1, num_spec_tokens + 1) input_ids = request.all_token_ids.copy() delta_ids = list(spec_tokens) @@ -281,15 +297,15 @@ def _get_constraint_start( input_ids = request.all_token_ids.copy() input_ids.extend(spec_tokens) if not reasoner.is_reasoning_end_streaming(input_ids, spec_tokens): - return num_spec_tokens + 1 + return ConstraintBounds(num_spec_tokens + 1, num_spec_tokens + 1) delta_ids = list(spec_tokens) for i in range(num_spec_tokens - 1, 0, -1): input_ids.pop() delta_ids.pop() if not reasoner.is_reasoning_end_streaming(input_ids, delta_ids): - return i + 1 - return 1 + return ConstraintBounds(i + 1, i + 1) + return ConstraintBounds(1, 1) def validate_tokens(self, request: "Request", spec_tokens: list[int]) -> list[int]: """Return the longest unconstrained or grammar-valid prefix of `spec_tokens`.""" @@ -297,8 +313,8 @@ def validate_tokens(self, request: "Request", spec_tokens: list[int]) -> list[in return spec_tokens spec_tokens = strip_speculative_padding(spec_tokens) - constraint_start = self._get_constraint_start(request, spec_tokens) - if constraint_start >= len(spec_tokens): + grammar_start, _ = self._get_constraint_bounds(request, spec_tokens) + if grammar_start >= len(spec_tokens): return spec_tokens structured_req = request.structured_output_request @@ -307,8 +323,8 @@ def validate_tokens(self, request: "Request", spec_tokens: list[int]) -> list[in grammar = structured_req.grammar if TYPE_CHECKING: assert isinstance(grammar, StructuredOutputGrammar) - prefix = spec_tokens[:constraint_start] - validated = grammar.validate_tokens(spec_tokens[constraint_start:]) + prefix = spec_tokens[:grammar_start] + validated = grammar.validate_tokens(spec_tokens[grammar_start:]) return prefix + validated def grammar_bitmask( @@ -358,7 +374,9 @@ def grammar_bitmask( if TYPE_CHECKING: assert isinstance(grammar, StructuredOutputGrammar) - apply_bitmask = self._get_constraint_start(request, ()) == 0 + apply_bitmask = ( + self._get_constraint_bounds(request, ()).constraint_start == 0 + ) batch.append((grammar, cumulative_index, apply_bitmask)) if len(batch) == self.fill_bitmask_parallel_batch_size: promises.append(self._async_submit_fill_bitmask(batch)) @@ -384,7 +402,7 @@ def grammar_bitmask( assert isinstance(grammar, StructuredOutputGrammar) req_tokens = scheduled_spec_decode_tokens.get(req_id, list()) - constraint_start = self._get_constraint_start( + grammar_start, constraint_start = self._get_constraint_bounds( request, strip_speculative_padding(req_tokens) ) state_advancements = 0 @@ -402,7 +420,7 @@ def grammar_bitmask( self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),)) if token == -1: seen_padding = True - elif apply_bitmask: + elif not seen_padding and i >= grammar_start: if not grammar.is_terminated() and grammar.accept_tokens( req_id, [token] ): @@ -453,17 +471,15 @@ def accept_tokens(self, request: "Request", new_token_ids: list[int]) -> bool: if TYPE_CHECKING: assert isinstance(grammar, StructuredOutputGrammar) - constraint_start = self._get_constraint_start( + grammar_start, _ = self._get_constraint_bounds( request, new_token_ids, spec_tokens_committed=True ) - # Early return only when the constraint hasn't started. + # Early return only when no reasoning-end token is present. # Otherwise, latch structured_req.reasoning_ended. - if constraint_start > len(new_token_ids): + if grammar_start > len(new_token_ids): return True structured_req.reasoning_ended = True - return grammar.accept_tokens( - request.request_id, new_token_ids[constraint_start:] - ) + return grammar.accept_tokens(request.request_id, new_token_ids[grammar_start:]) def clear_backend(self) -> None: if self.backend is not None: From 1afb029839a7c84f14dae30bb08f2fea79184431 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Thu, 8 Oct 2026 15:45:44 -0400 Subject: [PATCH 129/624] [Bugfix][HiSparse] Route by swap-row capacity, not the index-group workspace (#60682) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- tests/v1/attention/test_sparse_mla_backends.py | 12 ++++++++++-- vllm/v1/attention/backends/mla/index_group.py | 2 +- vllm/v1/hisparse/runtime.py | 9 ++++++--- 3 files changed, 17 insertions(+), 6 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 4d33075c9303..973a9e03e88b 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -3296,7 +3296,8 @@ def run_kernel(self, *, q, **kwargs): # noqa: ARG001 cache = SimpleNamespace( runtime=SimpleNamespace( - hot=SimpleNamespace(attention_cache=torch.empty(1, device=device)) + hot=SimpleNamespace(attention_cache=torch.empty(1, device=device)), + max_swap_rows=num_tokens, ), source_block_table=torch.empty( num_decodes, 1, dtype=torch.int32, device=device @@ -3604,15 +3605,22 @@ def test_flashinfer_sm120_hisparse_decode_uses_index_group(): def test_hisparse_resident_prefill_uses_attention_block_stride(): + """A batch beyond the swap-row capacity must bypass swap_in. + + The index-group workspace has one more row than the HiSparse swap rows, so + routing on its size sent a 2-token batch with max_num_seqs=1 (a piecewise + CUDA-graph capture) into swap_in, which overflowed its capacity of 1. + """ expected = torch.tensor([[19]], dtype=torch.int32) cache_handle = SimpleNamespace( all_context_pages_resident=True, view=SimpleNamespace(block_size=64, attention_block_stride=832), block_table=torch.tensor([[3]], dtype=torch.int32), + runtime=SimpleNamespace(max_swap_rows=1), ) index_group = object.__new__(HiSparseMLAIndexGroup) index_group.caches = [cache_handle] - index_group.physical_topk_indices = torch.empty((1, 1), dtype=torch.int32) + index_group.physical_topk_indices = torch.empty((2, 1), dtype=torch.int32) index_group._convert_once = MagicMock(return_value=expected) topk = torch.zeros((2, 1), dtype=torch.int32) metadata = SimpleNamespace( diff --git a/vllm/v1/attention/backends/mla/index_group.py b/vllm/v1/attention/backends/mla/index_group.py index 5a18332e41e1..39f205b91d97 100644 --- a/vllm/v1/attention/backends/mla/index_group.py +++ b/vllm/v1/attention/backends/mla/index_group.py @@ -233,7 +233,7 @@ def convert_logical_to_physical_topk( cache = self.cache(layer_index) num_tokens = logical_topk_indices.shape[0] req_id_per_token = attn_metadata.req_id_per_token[:num_tokens] - if num_tokens > self.physical_topk_indices.shape[0]: + if num_tokens > cache.runtime.max_swap_rows: # Prefill-sized batches do not fit the decode residency workspace. # Non-resident prefills are staged before reaching this path. assert cache.all_context_pages_resident diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index a6378ca9141a..1f7a5ad5ca3c 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -894,14 +894,17 @@ def invalidate_written_slots( def begin_forward(self) -> None: self._swap_step = 0 + @property + def max_swap_rows(self) -> int: + return self.index_group.shared_topk.physical_topk_indices.shape[0] + def _step_rows(self, num_tokens: int) -> slice: start = self._swap_step * num_tokens stop = start + num_tokens - if stop > self.index_group.shared_topk.physical_topk_indices.shape[0]: + if stop > self.max_swap_rows: raise ValueError( "HiSparse swap rows exceed the configured speculative decode " - f"capacity: stop={stop}, capacity=" - f"{self.index_group.shared_topk.physical_topk_indices.shape[0]}." + f"capacity: stop={stop}, capacity={self.max_swap_rows}." ) self._swap_step += 1 return slice(start, stop) From 72d59adc2c76de8f77aad640c7a04d663a970735 Mon Sep 17 00:00:00 2001 From: Alberto Massidda Date: Thu, 8 Oct 2026 21:58:13 +0200 Subject: [PATCH 130/624] [Bugfix][KV Connector][NIXL] Transfer the whole PLE short-conv page in disaggregated P/D (#59997) Signed-off-by: Alberto Massidda Co-authored-by: Alberto Massidda Co-authored-by: Claude Fable 5.1 Co-authored-by: Artem Perevedentsev --- .../unit/test_nixl_desc_geometry.py | 60 +++++++++++++++---- .../kv_connector/v1/nixl/base_worker.py | 28 +++++++-- .../kv_connector/v1/nixl/metadata.py | 5 +- 3 files changed, 77 insertions(+), 16 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index 029ffafd3ceb..9f6bc56deeb5 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -1245,8 +1245,12 @@ def test_mismatched_mla_kernel_page_rejected_for_mla_hybrid(): worker.add_remote_agent(meta_r, remote_tp_rank=0, remote_tp_size=2) -def _make_csa_linear_ple_worker(scratch_aliases: str = "compressed"): - """``scratch_aliases`` selects which pages the compressor ring overlays.""" +def _make_csa_linear_ple_worker( + scratch_aliases: str = "compressed", ple_page_size: int = 256 +): + """``scratch_aliases`` selects which pages the compressor ring overlays; + ``ple_page_size`` widens the PLE page beyond the page it shares a region + with (block-outer layout: both start at byte 0 of the block).""" from unittest.mock import MagicMock from vllm.config import set_current_vllm_config @@ -1305,7 +1309,7 @@ def _make_csa_linear_ple_worker(scratch_aliases: str = "compressed"): block_size=1, shapes=((6, 3),), dtypes=(torch.float16,), - page_size_padded=256, + page_size_padded=ple_page_size, mamba_type=MambaAttentionBackendEnum.SHORT_CONV, tp_replicated=True, ) @@ -1339,17 +1343,22 @@ def group(specs): ("compressed.0",), ("compressed.1",), ) - region_size = 512 page_size = 256 + # Region 0 hosts the PLE page: its blocks are as wide as the widest page. + block_strides = [max(page_size, ple_page_size)] + [page_size] * ( + len(tensor_regions) - 1 + ) + region_sizes = [2 * stride for stride in block_strides] + region_offsets = [sum(region_sizes[:index]) for index in range(len(region_sizes))] kv_cache_config = KVCacheConfig( num_blocks=2, kv_cache_tensors=[ KVCacheTensor( - size=len(tensor_regions) * region_size, + size=sum(region_sizes), layers=[layer_name], - layer_stride=region_size, - block_stride=page_size, - offset=region_index * region_size, + layer_stride=region_sizes[region_index], + block_stride=block_strides[region_index], + offset=region_offsets[region_index], ) for region_index, layer_names in enumerate(tensor_regions) for layer_name in layer_names @@ -1386,10 +1395,14 @@ def group(specs): set_current_vllm_config(vllm_config), ): worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config) - tensors = [torch.zeros((2, 256), dtype=torch.uint8) for _ in range(4)] + tensors = [ + torch.zeros((2, stride), dtype=torch.uint8) for stride in block_strides + ] worker.register_kv_caches( { - layer_name: tensors[region_index] + layer_name: tensors[region_index][ + :, : (ple_page_size if layer_name == "mamba.ple" else page_size) + ] for region_index, layer_names in enumerate(tensor_regions) for layer_name in layer_names } @@ -1523,6 +1536,7 @@ def test_csa_linear_remote_ple_is_copied_whole(): kv_cache_layout="HND", block_size=4, ssm_sizes=(24, 32), + ple_block_len=256, attn_backend_name="test", physical_blocks_per_logical_kv_block=1, ) @@ -1535,7 +1549,7 @@ def test_csa_linear_remote_ple_is_copied_whole(): assert descriptors[-2:, 0].tolist() == [0x10000, 0x10100] assert descriptors[-2:, 1].tolist() == [256, 256] - metadata.block_lens[0] = 128 + metadata.ple_block_len = 128 with pytest.raises(ValueError, match="PLE pages require identical"): worker._build_mamba_remote( metadata, @@ -1544,6 +1558,30 @@ def test_csa_linear_remote_ple_is_copied_whole(): ) +@pytest.mark.cpu_test +def test_csa_linear_ple_page_wider_than_the_shared_region_page(): + """The PLE page is discovered in the region of the main KV page registered + first (both start at byte 0 of the block); its descriptors must still cover + the whole PLE page, not the region's block_len.""" + worker = _make_csa_linear_ple_worker(ple_page_size=384) + + assert worker._ple_region_index == 0 + assert worker.block_len_per_layer[0] == 256 + assert worker.block_stride_per_layer[0] == 384 + assert worker._ple_block_len == 384 + + bases = [0x10000, 0x20000, 0x30000, 0x40000] + mamba = worker._build_mamba_local(bases) + assert mamba[-2:, 0].tolist() == [0x10000, 0x10000 + 384] + assert mamba[-2:, 1].tolist() == [384, 384] + + metadata = msgspec.msgpack.decode( + worker.xfer_handshake_metadata.agent_metadata_bytes, type=NixlAgentMetadata + ) + assert metadata.ple_block_len == 384 + assert metadata.block_lens[0] == 256 + + def _make_ring_worker(): """Paged MLA group plus a per-layer ring group, no Mamba: the shape of a model whose sliding-window KV lives in per-request rings.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 2a2f660cf15e..7ce6e93ebb46 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -935,6 +935,9 @@ def __init__( # page it overlays, so its regions cannot be recovered from spec type. self._scratch_region_indices = list[int]() self._ple_region_index: int | None = None + # Page length of the PLE short-conv state; its region may be shared with + # pages of a different (smaller) size, so it cannot use the region's. + self._ple_block_len: int | None = None # Enable different block lengths for different layers *only* when MLA is used. # This is not used for SSM layers, which use the counterpart `mamba_ssm_size`. @@ -1464,6 +1467,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): seen_base_addresses: list[int] = [] self._ssm_region_indices = [] self._scratch_region_indices = [] + self._ple_block_len = None self._ple_region_index = None packed_storage = _share_storage_and_block_stride(list(xfer_buffers.values())) @@ -1751,6 +1755,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): assert self._is_csa_linear assert self._ple_region_index in (None, region_index) self._ple_region_index = region_index + self._ple_block_len = block_len elif region_index not in self._ssm_region_indices: self._ssm_region_indices.append(region_index) elif ( @@ -1885,6 +1890,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): else self.host_buffer_kv_cache_layout, block_size=self.block_size, ssm_sizes=self._mamba_ssm_size, + ple_block_len=self._ple_block_len, attn_backend_name=self.backend_name, physical_blocks_per_logical_kv_block=( self._physical_blocks_per_logical_kv_block @@ -1906,6 +1912,19 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): agent_metadata_bytes=encoder.encode(agent_metadata), ) + def _ple_page_len(self) -> int: + """Transfer length of one PLE state page. + + The PLE page shares its region with the first page of other cache + groups (block-outer layout), whose block_len may be smaller; the + descriptors must cover the whole PLE page or part of the conv state + never reaches the decode side.""" + assert self._ple_region_index is not None + assert self._ple_block_len is not None, ( + "PLE region registered without its page length" + ) + return self._ple_block_len + def _build_mamba_local(self, base_addresses: list[int]) -> np.ndarray: """Build desc regions (conv sub-projections + ssm) per layer for local mamba blocks with DS conv layout, as an Nx3 uint64 array. @@ -1960,7 +1979,7 @@ def _build_mamba_local(self, base_addresses: list[int]) -> np.ndarray: parts.append(self._stack_descs(blk_addrs + conv_size, ssm_size, device_id)) if (region_index := self._ple_region_index) is not None: - block_len = self.block_len_per_layer[region_index] * physical_per_logical + block_len = self._ple_page_len() * physical_per_logical block_stride = ( self.block_stride_per_layer[region_index] * physical_per_logical ) @@ -2019,11 +2038,12 @@ def _build_mamba_remote( if (region_index := self._ple_region_index) is not None: local_block_len = ( - self.block_len_per_layer[region_index] - * self._physical_blocks_per_logical_kv_block + self._ple_page_len() * self._physical_blocks_per_logical_kv_block ) + # Same NIXL_CONNECTOR_VERSION on both sides, so the field is set. + assert nixl_agent_meta.ple_block_len is not None remote_block_len = ( - nixl_agent_meta.block_lens[region_index] * remote_physical_per_logical + nixl_agent_meta.ple_block_len * remote_physical_per_logical ) if local_block_len != remote_block_len: raise ValueError( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index 8d0aad763358..a33f620a0a0d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -49,8 +49,9 @@ # 11: Add per-region transfer geometry and memory types to NixlAgentMetadata # 12: Add per-region member names for PP push # 13: Add packed-member layouts and order-independent packed-push backend hashes +# 14: Add Qwen3.8-Flash-Next PLE page length - #59997 # -NIXL_CONNECTOR_VERSION: int = 13 +NIXL_CONNECTOR_VERSION: int = 14 @dataclass @@ -73,6 +74,8 @@ class NixlAgentMetadata: region_mem_types: list[str] | None = None dcp_size: int = 1 pcp_size: int = 1 + # Page length of the PLE short-conv state (None when the model has none). + ple_block_len: int | None = None # Layer names sharing each advertised region, in region order. region_members: list[list[str]] = field(default_factory=list) # Packed member -> (byte offset in its region's block, bytes per page). From ae9adec9059b6c6d831cfa46ce966a3d629a4526 Mon Sep 17 00:00:00 2001 From: Francisco Javier Arceo Date: Thu, 8 Oct 2026 16:21:05 -0400 Subject: [PATCH 131/624] [Frontend] Expose /v1/systemone without an opt-in flag (#60677) Signed-off-by: Francisco Javier Arceo <4163062+franciscojavierarceo@users.noreply.github.com> Co-authored-by: Codex Co-authored-by: Cyrus Leung --- .../online_serving/structured_decisions.md | 8 ++--- .../structured_decisions/test_strategies.py | 35 +++++++++++++++---- .../test_structured_decisions_e2e.py | 1 - .../entrypoints/openai/test_openai_schema.py | 8 ++++- vllm/entrypoints/generate/api_router.py | 25 ++++++------- .../structured_decisions/api_router.py | 3 +- .../structured_decisions/strategies.py | 7 ++-- vllm/entrypoints/launchers/cli_args.py | 2 -- 8 files changed, 56 insertions(+), 33 deletions(-) diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md index dd26f2130164..7de935a62f2b 100644 --- a/docs/serving/online_serving/structured_decisions.md +++ b/docs/serving/online_serving/structured_decisions.md @@ -3,10 +3,10 @@ The `/v1/systemone` endpoint answers a set of typed questions about a state and returns a probability for every allowed answer. -The endpoint is off unless the server starts with -`--enable-structured-decisions`. With the flag, startup fails for a model that -does not support structured decisions, and with `--logprobs-mode raw_logits` or -`processed_logits`, since `label_mass` needs log probabilities. +The endpoint is available by default on generation models. Models without a +supported read strategy return 501. Supported architectures require +`--logprobs-mode raw_logprobs` or `processed_logprobs`, since `label_mass` needs +log probabilities. ## How it works diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index 017788f7b1d8..51bff19c8ca2 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -5,6 +5,7 @@ import pytest from fastapi import FastAPI +from fastapi.testclient import TestClient from vllm.entrypoints.generate.structured_decisions.api_router import ( register_structured_decisions_api_router, @@ -16,6 +17,9 @@ reply_label_ids, select_read_strategy, ) +from vllm.entrypoints.serve.exception_handling.register import ( + init_exception_handler, +) def model(architecture: str, logprobs_mode: str = "raw_logprobs") -> Any: @@ -32,13 +36,30 @@ def test_strategy_selection(): select_read_strategy(model(qwen, "raw_logits")) -def test_route_needs_the_flag(): - for enabled in (False, True): - app = FastAPI() - app.state.args = SimpleNamespace(enable_structured_decisions=enabled) - register_structured_decisions_api_router(app) - paths = {getattr(route, "path", None) for route in app.routes} - assert ("/v1/systemone" in paths) == enabled +def test_route_is_registered_by_default(): + app = FastAPI() + register_structured_decisions_api_router(app) + paths = {getattr(route, "path", None) for route in app.routes} + assert "/v1/systemone" in paths + + +def test_unsupported_model_returns_501(): + app = FastAPI() + app.state.args = SimpleNamespace(log_error_stack=False) + app.state.serving_structured_decisions = None + init_exception_handler(app) + register_structured_decisions_api_router(app) + + with TestClient(app) as client: + response = client.post( + "/v1/systemone", + json={ + "model": "unsupported", + "state": "x", + "questions": {"answer": {"type": "choice", "criteria": {"yes": None}}}, + }, + ) + assert response.status_code == 501 def test_structured_state_keeps_unicode_in_prompt(): diff --git a/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py b/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py index e060b096ec51..3e0f9bc0fb66 100644 --- a/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py +++ b/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py @@ -19,7 +19,6 @@ def server(): "--max-num-seqs", "32", "--enable-prefix-caching", - "--enable-structured-decisions", ] with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: yield remote_server diff --git a/tests/entrypoints/openai/test_openai_schema.py b/tests/entrypoints/openai/test_openai_schema.py index 18b01b85a8a3..a405550c4cf4 100644 --- a/tests/entrypoints/openai/test_openai_schema.py +++ b/tests/entrypoints/openai/test_openai_schema.py @@ -7,6 +7,7 @@ import schemathesis from hypothesis import HealthCheck, settings from schemathesis import GenerationMode +from schemathesis.checks import not_a_server_error from schemathesis.config import ( ChecksConfig, CoveragePhaseConfig, @@ -161,8 +162,13 @@ def test_openapi_stateless(case: schemathesis.Case): }.get(key, DEFAULT_TIMEOUT_SECONDS) # No need to verify SSL certificate for localhost - case.call_and_validate( + systemone = case.operation.path == "/v1/systemone" + response = case.call_and_validate( verify=False, timeout=timeout, headers={"Content-Type": "application/json"}, + excluded_checks=[not_a_server_error] if systemone else None, ) + if systemone: + # SmolVLM does not support structured decisions. + assert response.status_code < 500 or response.status_code == 501 diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 70869b61b5d2..67a256192443 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -230,16 +230,11 @@ async def init_generate_state( from .structured_decisions.serving import ServingStructuredDecisions from .structured_decisions.strategies import ReadContext, select_read_strategy - strategy_cls = ( - select_read_strategy(engine_client.model_config) - if "generate" in supported_tasks - and getattr(args, "enable_structured_decisions", False) - else None - ) - state.serving_structured_decisions = ( - ServingStructuredDecisions( - state.openai_serving_models, - strategy_cls( + strategy = None + if "generate" in supported_tasks: + try: + strategy_cls = select_read_strategy(engine_client.model_config) + strategy = strategy_cls( ReadContext( engine_client=engine_client, online_renderer=state.online_renderer, @@ -247,9 +242,15 @@ async def init_generate_state( chat_template_content_format=args.chat_template_content_format, default_chat_template_kwargs=default_chat_template_kwargs, ) - ), + ) + except ValueError: + strategy = None + state.serving_structured_decisions = ( + ServingStructuredDecisions( + state.openai_serving_models, + strategy, request_logger=request_logger, ) - if strategy_cls is not None + if strategy is not None else None ) diff --git a/vllm/entrypoints/generate/structured_decisions/api_router.py b/vllm/entrypoints/generate/structured_decisions/api_router.py index a0fbaa518e07..5729b58b2e50 100644 --- a/vllm/entrypoints/generate/structured_decisions/api_router.py +++ b/vllm/entrypoints/generate/structured_decisions/api_router.py @@ -49,5 +49,4 @@ async def create_structured_decision( def register_structured_decisions_api_router(app: FastAPI): - if getattr(app.state.args, "enable_structured_decisions", False): - app.include_router(router) + app.include_router(router) diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index 22233ff91b0e..864443cbb0b9 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -219,17 +219,16 @@ async def read( def select_read_strategy(model_config: ModelConfig) -> type[ReadStrategy]: - """Raises ValueError when the model cannot serve structured decisions. The - server opted in with --enable-structured-decisions, so startup fails.""" + """Raise ValueError when the model cannot serve structured decisions.""" if model_config.architecture not in NEXT_TOKEN_ARCHITECTURES: raise ValueError( - "--enable-structured-decisions does not support " + "The structured decisions API does not support " f"{model_config.architecture}. Supported architectures: " f"{sorted(NEXT_TOKEN_ARCHITECTURES)}" ) if model_config.logprobs_mode not in LOGPROBS_MODES: raise ValueError( - "--enable-structured-decisions needs --logprobs-mode raw_logprobs or " + "Structured decisions need --logprobs-mode raw_logprobs or " f"processed_logprobs, not {model_config.logprobs_mode}" ) return NextTokenStrategy diff --git a/vllm/entrypoints/launchers/cli_args.py b/vllm/entrypoints/launchers/cli_args.py index 18f8cb00b503..28fe79fd691e 100644 --- a/vllm/entrypoints/launchers/cli_args.py +++ b/vllm/entrypoints/launchers/cli_args.py @@ -163,8 +163,6 @@ class BaseFrontendArgs: enable_tokenizer_info_endpoint: bool = False """Enable the `/tokenizer_info` endpoint. May expose chat templates and other tokenizer configuration.""" - enable_structured_decisions: bool = False - """Enable the `/v1/systemone` structured decisions endpoint.""" enable_log_outputs: bool = False """If set to True, log model outputs (generations). Requires `--enable-log-requests`. Output text and finish reasons are logged at INFO, From c5c116138267ac738bc262495326d28ffff834a2 Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Thu, 8 Oct 2026 16:28:38 -0400 Subject: [PATCH 132/624] [ROCm][Build] Fail fast when `setup.py develop` deps aren't pip-preinstalled (#59447) Signed-off-by: Divakar Verma --- setup.py | 70 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/setup.py b/setup.py index 5e49f74e6fcd..16a3d35ecad6 100644 --- a/setup.py +++ b/setup.py @@ -1311,6 +1311,74 @@ def get_vllm_version() -> str: return version +def _check_requirements_preinstalled(requirements: list[str]) -> None: + """`python setup.py develop` satisfies missing `install_requires` via + setuptools' legacy easy_install path instead of pip's wheel-aware + resolver. That path has repeatedly tried (and failed) to compile old + sdists of fast-moving C-extension packages such as aiohttp from source + against the running Python's headers, e.g. + https://github.com/vllm-project/vllm/issues/34073. Fail fast with an + actionable message instead of a cryptic compiler error. + + Skipped under `--no-deps`, which makes setuptools skip `install_requires` + processing entirely, so the easy_install path this guards against is + unreachable anyway. + """ + from importlib.metadata import PackageNotFoundError + from importlib.metadata import requires as installed_requires + from importlib.metadata import version as installed_version + + from packaging.requirements import Requirement + + missing = [] + checked = set() + + def _check(req_str: str, extra: str | None = None) -> None: + # Strip inline "# comment" suffixes (as pip does when reading + # requirements files); packaging.Requirement cannot parse them. + req_str = re.split(r"\s+#", req_str, maxsplit=1)[0].strip() + try: + req = Requirement(req_str) + except Exception: + return + # `extra` is only set when recursing into an extra's own deps + # below; it selects which `; extra == "..."` markers apply. + if req.marker is not None and not req.marker.evaluate( + {"extra": extra} if extra else None + ): + return + if (req.name, extra) in checked: + return + checked.add((req.name, extra)) + try: + installed = installed_version(req.name) + except PackageNotFoundError: + missing.append(req_str) + return + if req.specifier and not req.specifier.contains(installed, prereleases=True): + missing.append(f"{req_str} (found {req.name}=={installed})") + return + # A requested extra (e.g. "fastapi[standard]") pulls in packages + # that easy_install resolves too, so check those as well. + for extra_name in req.extras: + for dep in installed_requires(req.name) or []: + _check(dep, extra=extra_name) + + for req_str in requirements: + _check(req_str) + + if missing: + raise RuntimeError( + "The following dependencies are missing or outdated:\n " + + "\n ".join(missing) + + "\n\n`python setup.py develop` cannot reliably install these " + "itself (see https://github.com/vllm-project/vllm/issues/34073). " + "Install them with pip first, e.g.:\n" + " pip install -r requirements/rocm.txt\n" + "then re-run `python setup.py develop`." + ) + + def get_requirements() -> list[str]: """Get Python package dependencies from requirements.txt.""" requirements_dir = ROOT_DIR / "requirements" @@ -1355,6 +1423,8 @@ def _read_requirements(filename: str) -> list[str]: requirements = modified_requirements elif _is_hip(): requirements = _read_requirements("rocm.txt") + if "develop" in sys.argv[1:] and "--no-deps" not in sys.argv[1:]: + _check_requirements_preinstalled(requirements) elif _is_tpu(): requirements = _read_requirements("tpu.txt") elif _is_cpu(): From 2ddde19e8e0ba1b9d06a3f9b6201f7fc915d27ab Mon Sep 17 00:00:00 2001 From: sungbin1015 Date: Fri, 9 Oct 2026 05:38:17 +0900 Subject: [PATCH 133/624] [Frontend] Support min_p in the Responses API (#48084) Signed-off-by: sungbin1015 Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../online_serving/openai_compatible_server.md | 6 ++++++ tests/benchmarks/test_responses_request_func.py | 4 +--- .../openai/responses/test_sampling_params.py | 13 +++++++++++++ vllm/benchmarks/serve.py | 5 ----- vllm/entrypoints/openai/responses/protocol.py | 7 +++++++ 5 files changed, 27 insertions(+), 8 deletions(-) diff --git a/docs/serving/online_serving/openai_compatible_server.md b/docs/serving/online_serving/openai_compatible_server.md index c50910040d32..7b4e84a17fee 100644 --- a/docs/serving/online_serving/openai_compatible_server.md +++ b/docs/serving/online_serving/openai_compatible_server.md @@ -204,6 +204,12 @@ To get prompt token IDs before generation, use a self-contained Responses request without inference and returns the prompt token IDs, so a separate `/tokenize` call is not needed. +#### Sampling parameters + +In addition to the OpenAI request fields, the Responses API forwards vLLM +[sampling parameters](../../api/README.md#inference-parameters) such as `top_k`, +`min_p`, `repetition_penalty`, `seed`, `stop`, and `ignore_eos` to the engine. + #### Extra parameters The following extra parameters in the request object are supported: diff --git a/tests/benchmarks/test_responses_request_func.py b/tests/benchmarks/test_responses_request_func.py index 38f1bdc50781..1ecf6115fff0 100644 --- a/tests/benchmarks/test_responses_request_func.py +++ b/tests/benchmarks/test_responses_request_func.py @@ -183,6 +183,7 @@ def test_request_payload_is_accepted_by_the_server_schema(): "temperature": 0.7, "top_p": 0.9, "top_k": 20, + "min_p": 0.05, "frequency_penalty": 0.2, "presence_penalty": 0.1, "repetition_penalty": 1.05, @@ -217,9 +218,6 @@ def test_request_payload_is_accepted_by_the_server_schema(): assert set(seen["payload"]) <= set(ResponsesRequest.model_fields) ResponsesRequest.model_validate(seen["payload"]) - # `vllm bench serve` rejects --min-p for this backend for the same reason. - assert "min_p" not in ResponsesRequest.model_fields - def test_latency_accounting_matches_the_other_endpoints(): """`latency - ttft` must equal `sum(itl)`. diff --git a/tests/entrypoints/openai/responses/test_sampling_params.py b/tests/entrypoints/openai/responses/test_sampling_params.py index 01fdfa4594ea..24d54b8960de 100644 --- a/tests/entrypoints/openai/responses/test_sampling_params.py +++ b/tests/entrypoints/openai/responses/test_sampling_params.py @@ -29,6 +29,7 @@ def test_basic_sampling_params(self): temperature=0.8, top_p=0.95, top_k=50, + min_p=0.1, max_output_tokens=100, ) @@ -37,8 +38,20 @@ def test_basic_sampling_params(self): assert sampling_params.temperature == 0.8 assert sampling_params.top_p == 0.95 assert sampling_params.top_k == 50 + assert sampling_params.min_p == 0.1 assert sampling_params.max_tokens == 100 + def test_min_p_default(self): + """Test that min_p defaults to 0.0 when unset (parity with top_k).""" + request = ResponsesRequest( + model="test-model", + input="test input", + ) + + sampling_params = request.to_sampling_params(default_max_tokens=1000) + + assert sampling_params.min_p == 0.0 + def test_extra_sampling_params(self): """Test extra sampling parameters are correctly mapped.""" request = ResponsesRequest( diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py index 895e5c6c1890..01a0f62aeee2 100644 --- a/vllm/benchmarks/serve.py +++ b/vllm/benchmarks/serve.py @@ -2196,11 +2196,6 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]: "Sampling parameters are only supported by openai-compatible backends." ) - # The Responses API accepts every sampling parameter above except - # min_p, which it would silently drop as an unknown field. - if args.backend == "openai-responses" and "min_p" in sampling_params: - raise ValueError("--min-p is not supported by the Responses API.") - if "temperature" not in sampling_params: print( "WARNING: vllm bench serve no longer sets temperature==0 (greedy) " diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 12e40e98a882..d1a4ea4f6c21 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -259,6 +259,7 @@ class ResponsesRequest(OpenAIBaseModel): top_logprobs: TopLogprobsParam = 0 top_p: float | None = None top_k: int | None = None + min_p: float | None = None truncation: Literal["auto", "disabled"] | None = "disabled" user: str | None = None skip_special_tokens: bool = True @@ -444,6 +445,7 @@ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: "temperature": 1.0, "top_p": 1.0, "top_k": 0, + "min_p": 0.0, } def extract_structured_outputs(self) -> StructuredOutputsParams | None: @@ -494,6 +496,10 @@ def to_sampling_params( top_k = default_sampling_params.get( "top_k", self._DEFAULT_SAMPLING_PARAMS["top_k"] ) + if (min_p := self.min_p) is None: + min_p = default_sampling_params.get( + "min_p", self._DEFAULT_SAMPLING_PARAMS["min_p"] + ) if (repetition_penalty := self.repetition_penalty) is None: repetition_penalty = default_sampling_params.get("repetition_penalty", 1.0) @@ -519,6 +525,7 @@ def to_sampling_params( watermarking=self.watermarking, top_p=top_p, top_k=top_k, + min_p=min_p, max_tokens=max_tokens, logprobs=self.top_logprobs if self.is_include_output_logprobs() else None, stop=stop, From 240785b82c299c6c08bc5239e997cf62d67cf7bb Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Thu, 8 Oct 2026 13:44:04 -0700 Subject: [PATCH 134/624] [Bugfix] Keep the GLM-5.3 kpool tail and Qwen4 QSA ring out of the null block (#59528) Signed-off-by: Yifan Qiao Co-authored-by: Claude Opus 5.5 --- .../qwen4_exp/test_qsa_pre_indexer_amd.py | 5 ++-- tests/models/qwen4_exp/test_qsa_reference.py | 18 ++++++----- .../attention/test_kpool_tail_slot_mapping.py | 30 ++++++++++++------- vllm/models/qwen4_exp/common/qsa_cache.py | 7 +++-- vllm/v1/attention/backends/mla/indexer.py | 9 ++++-- 5 files changed, 45 insertions(+), 24 deletions(-) diff --git a/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py b/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py index 8a39504e0067..7698198923ce 100644 --- a/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py +++ b/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py @@ -55,9 +55,10 @@ def _reference_gemma_rmsnorm( def _make_block_table(block_counts): - num_blocks = sum(block_counts) + # Block 0 is the null block, which no request owns. + num_blocks = sum(block_counts) + 1 table = torch.full((len(block_counts), max(block_counts)), -1, dtype=torch.int32) - physical_blocks = torch.randperm(num_blocks) + physical_blocks = torch.randperm(sum(block_counts)) + 1 offset = 0 for request, count in enumerate(block_counts): table[request, :count] = physical_blocks[offset : offset + count] diff --git a/tests/models/qwen4_exp/test_qsa_reference.py b/tests/models/qwen4_exp/test_qsa_reference.py index bdcc4e342370..a55cc028f9db 100644 --- a/tests/models/qwen4_exp/test_qsa_reference.py +++ b/tests/models/qwen4_exp/test_qsa_reference.py @@ -315,7 +315,7 @@ def test_qsa_circular_buffer_metadata_keeps_only_each_requests_suffix() -> None: builder.k_work_metadata_buffer = torch.empty(0, 2, dtype=torch.int32, device=device) query_start_loc = torch.tensor([0, 7, 13, 13], dtype=torch.int32, device=device) token_to_req = torch.tensor([0] * 7 + [1] * 6 + [0] * 3, device=device) - block_table = torch.tensor([[1], [0], [2]], dtype=torch.int32, device=device) + block_table = torch.tensor([[1], [3], [2]], dtype=torch.int32, device=device) common = SimpleNamespace( num_actual_tokens=16, num_reqs=3, @@ -340,10 +340,10 @@ def test_qsa_circular_buffer_metadata_keeps_only_each_requests_suffix() -> None: 4, -1, -1, - 3, - 0, - 1, - 2, + 15, + 12, + 13, + 14, -1, -1, -1, @@ -351,6 +351,10 @@ def test_qsa_circular_buffer_metadata_keeps_only_each_requests_suffix() -> None: assert metadata.slot_mapping.tolist() == expected + # A dummy batch puts every request on the null block, which owns no ring. + block_table.zero_() + assert builder.build(0, common).slot_mapping.tolist() == [-1] * 16 + @pytest.mark.parametrize("chunk_start", list(range(8))) def test_qsa_circular_buffer_survives_one_speculative_step(chunk_start: int) -> None: @@ -367,7 +371,7 @@ def test_qsa_circular_buffer_survives_one_speculative_step(chunk_start: int) -> query_len = num_spec + 1 slots = qsa_cache.circular_qsa_slot_mapping( - torch.tensor([[0]], dtype=torch.int32), + torch.tensor([[1]], dtype=torch.int32), torch.zeros(query_len, dtype=torch.int32), torch.arange(chunk_start, chunk_start + query_len), capacity, @@ -375,7 +379,7 @@ def test_qsa_circular_buffer_survives_one_speculative_step(chunk_start: int) -> ) committed = torch.arange(chunk_start - chunk_start % compress_ratio, chunk_start) - assert set(slots.tolist()).isdisjoint((committed % capacity).tolist()) + assert set((slots % capacity).tolist()).isdisjoint((committed % capacity).tolist()) def _qsa_key_cache( diff --git a/tests/v1/attention/test_kpool_tail_slot_mapping.py b/tests/v1/attention/test_kpool_tail_slot_mapping.py index da25eaefb23b..183c63ec2ffd 100644 --- a/tests/v1/attention/test_kpool_tail_slot_mapping.py +++ b/tests/v1/attention/test_kpool_tail_slot_mapping.py @@ -227,8 +227,10 @@ def test_circular_mapping_matches_generic_for_short_requests(prompt_len): def test_circular_mapping_preserves_padding_and_empty_batch(): - own_blocks = [5, 9] - per_req = [list(range(10)), list(range(12))] + # The last request is a padding request on the null block, which owns no + # tail block and must never be written. + own_blocks = [5, 9, 0] + per_req = [list(range(10)), list(range(12)), [0, 1]] padded_len = sum(len(p) for p in per_req) + 8 positions, qsl, slot_mapping, num_actual, num_reqs = make_batch( per_req, padded_len=padded_len @@ -237,7 +239,8 @@ def test_circular_mapping_preserves_padding_and_empty_batch(): out = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs) assert out.shape == slot_mapping.shape - assert torch.equal(out[num_actual:], torch.full_like(out[num_actual:], -1)) + pad_start = int(qsl[-2]) + assert torch.equal(out[pad_start:], torch.full_like(out[pad_start:], -1)) empty = circular_tail_slots(slot_mapping, bt, qsl, positions[:0], 0, num_reqs) assert torch.equal(empty, slot_mapping) @@ -430,25 +433,27 @@ def test_interleaved_decode_pollution_legacy_vs_circular(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") @pytest.mark.parametrize( - "per_req,num_actual,padded_len", + "per_req,own_blocks,num_actual,padded_len", [ - ([list(range(10)), list(range(12))], 22, 22), - ([list(range(10)), list(range(12))], 22, 30), - ([[3, 4], [0], [7, 8, 9]], 6, 8), - ([[5]], 1, 1), + ([list(range(10)), list(range(12))], [5, 6], 22, 22), + ([list(range(10)), list(range(12))], [5, 6], 22, 30), + ([[3, 4], [0], [7, 8, 9]], [5, 0, 7], 6, 8), + ([[5]], [5], 1, 1), + ([[0], [0]], [0, 0], 2, 4), # dummy run: every row on the null block ], ) -def test_triton_mapping_matches_cpu(per_req, num_actual, padded_len): +def test_triton_mapping_matches_cpu(per_req, own_blocks, num_actual, padded_len): """The CUDA (Triton) path must match the CPU torch reference, including tokens between the last request boundary and num_actual_tokens (mapped to - the last request) and untouched padding beyond num_actual.""" + the last request), untouched padding beyond num_actual, and PAD for + requests on the null block.""" positions, qsl, slot_mapping, _, num_reqs = make_batch( per_req, padded_len=padded_len ) # Replace the all--1 placeholder slots with sentinel values to check the # padding range is copied through untouched. slot_mapping = torch.arange(padded_len, dtype=torch.int64) + 1000 - bt = make_tail_block_table(list(range(5, 5 + num_reqs))) + bt = make_tail_block_table(own_blocks) ref = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs) got = circular_tail_slots( @@ -460,6 +465,9 @@ def test_triton_mapping_matches_cpu(per_req, num_actual, padded_len): num_reqs, ) torch.testing.assert_close(got.cpu(), ref) + for req, blk in enumerate(own_blocks): + if blk == 0: + assert (got[int(qsl[req]) : int(qsl[req + 1])] == -1).all() @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") diff --git a/vllm/models/qwen4_exp/common/qsa_cache.py b/vllm/models/qwen4_exp/common/qsa_cache.py index 13228185ec3c..f960dea07194 100644 --- a/vllm/models/qwen4_exp/common/qsa_cache.py +++ b/vllm/models/qwen4_exp/common/qsa_cache.py @@ -127,7 +127,9 @@ def circular_qsa_slot_mapping( valid = (requests >= 0) & (requests < block_table.shape[0]) & (positions >= 0) safe_requests = requests.clamp(0, block_table.shape[0] - 1) physical_blocks = block_table[safe_requests, 0].long() - valid &= physical_blocks >= 0 + # Dummy and padding requests own no ring and sit on the null block, + # which must never be written. + valid &= physical_blocks > 0 slots = physical_blocks * compressor_state_size + positions.remainder( compressor_state_size ) @@ -288,7 +290,8 @@ def _build_qsa_metadata_kernel( mask=valid, other=-1, ) - valid &= physical_block >= 0 + # Dummy and padding requests own no ring and sit on the null block. + valid &= physical_block > 0 slot = physical_block * circular_buffer_size + ( logical_position % circular_buffer_size ) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 8b186502f2a9..1005b8e4246a 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -702,11 +702,15 @@ def _kpool_tail_slot_mapping_kernel( end = tl.load(query_start_loc_ptr + pid + 1) end = tl.where(pid == num_reqs - 1, num_actual_tokens, end) own_block = tl.load(block_table_ptr + pid * block_table_stride).to(tl.int64) + # Every scheduled request owns a tail block; the runners give dummy and + # padding requests the null block, which must never be written. + has_block = own_block != 0 for i in range(start, end, BLOCK): offs = i + tl.arange(0, BLOCK) mask = offs < end pos = tl.load(positions_ptr + offs, mask=mask, other=0).to(tl.int64) - tl.store(out_ptr + offs, own_block * kpool + pos % kpool, mask=mask) + slot = tl.where(has_block, own_block * kpool + pos % kpool, -1) + tl.store(out_ptr + offs, slot, mask=mask) else: offs = num_actual_tokens + (pid - num_reqs) * BLOCK + tl.arange(0, BLOCK) mask = offs < num_tokens @@ -759,7 +763,8 @@ def compute_kpool_tail_slot_mapping( req = req.clamp_(min=0, max=num_reqs - 1) own_block = block_table[:num_reqs, 0].index_select(0, req).to(torch.int64) pos = positions[:num_actual_tokens].to(torch.int64) - out[:num_actual_tokens] = own_block * kpool + torch.remainder(pos, kpool) + slots = own_block * kpool + torch.remainder(pos, kpool) + out[:num_actual_tokens] = torch.where(own_block != 0, slots, -1) return out From 8d10c38b7304aa42432f4117f0d197ff46258ecb Mon Sep 17 00:00:00 2001 From: Kamran <79936503+spa5k@users.noreply.github.com> Date: Fri, 9 Oct 2026 03:12:56 +0530 Subject: [PATCH 135/624] [Perf] Reduce fallback decode CPU by 69% at 32K (#60497) Signed-off-by: spa5k <79936503+spa5k@users.noreply.github.com> Co-authored-by: Codex --- tests/tokenizers_/test_detokenize.py | 48 +++++++++++++++++++++++++++- vllm/tokenizers/detokenizer_utils.py | 20 ++++++++---- 2 files changed, 61 insertions(+), 7 deletions(-) diff --git a/tests/tokenizers_/test_detokenize.py b/tests/tokenizers_/test_detokenize.py index d0117e4711c0..1dd2a000a6fa 100644 --- a/tests/tokenizers_/test_detokenize.py +++ b/tests/tokenizers_/test_detokenize.py @@ -8,7 +8,11 @@ from transformers import AutoTokenizer, PythonBackend, TokenizersBackend from vllm.sampling_params import SamplingParams -from vllm.tokenizers.detokenizer_utils import convert_ids_list_to_tokens +from vllm.tokenizers.detokenizer_utils import ( + convert_ids_list_to_tokens, + convert_prompt_ids_to_tokens, + detokenize_incrementally, +) from vllm.tokenizers.mistral import MistralTokenizer from vllm.tokenizers.protocol import TokenizerLike from vllm.v1.engine import EngineCoreRequest @@ -98,6 +102,48 @@ def tokenizer(tokenizer_name): ) +@pytest.mark.parametrize( + "tokenizer_name", ["openai-community/gpt2", "mistralai/Pixtral-12B-2409"] +) +def test_slow_decode_long_history_preserves_text_and_token_ids(tokenizer): + truth = "A map. 我很感谢你的热情 🙂🌶️\n" * 32 + ids = tokenizer.encode(truth, add_special_tokens=False) + text, output_ids = _run_incremental_decode( + tokenizer, ids, skip_special_tokens=True, starting_index=0, fast=False + ) + assert text == tokenizer.decode(ids, skip_special_tokens=True) + assert output_ids == ids + + +@pytest.mark.parametrize( + "tokenizer_name", ["openai-community/gpt2", "mistralai/Pixtral-12B-2409"] +) +def test_first_incremental_decode_returns_the_prompt_window(tokenizer): + ids = tokenizer.encode("A map in a library. " * 8, add_special_tokens=False) + prompt_tokens, _, _ = convert_prompt_ids_to_tokens( + tokenizer, ids[:-1], skip_special_tokens=True + ) + new_token = tokenizer.convert_ids_to_tokens([ids[-1]], skip_special_tokens=True) + tokens, _, _, _ = detokenize_incrementally( + tokenizer, ids, None, 0, 0, skip_special_tokens=True + ) + assert tokens == prompt_tokens + new_token + + +@pytest.mark.parametrize( + "tokenizer_name", ["openai-community/gpt2", "mistralai/Pixtral-12B-2409"] +) +def test_incremental_decode_past_the_window_emits_no_text(tokenizer): + ids = tokenizer.encode("A map in a library.", add_special_tokens=False) + tokens = tokenizer.convert_ids_to_tokens(ids[:-1], skip_special_tokens=True) + offset = len(tokens) + 1 + _, text, prefix_offset, read_offset = detokenize_incrementally( + tokenizer, ids, tokens, offset, offset, skip_special_tokens=True + ) + assert text == "" + assert (prefix_offset, read_offset) == (offset, offset) + + @pytest.mark.parametrize("tokenizer_name", ["mistralai/Pixtral-12B-2409"]) @pytest.mark.parametrize( "truth", diff --git a/vllm/tokenizers/detokenizer_utils.py b/vllm/tokenizers/detokenizer_utils.py index ddaf512442c2..fc2ce4d44c82 100644 --- a/vllm/tokenizers/detokenizer_utils.py +++ b/vllm/tokenizers/detokenizer_utils.py @@ -230,7 +230,13 @@ def detokenize_incrementally( _replace_none_with_empty(new_tokens) # type: ignore[arg-type] else: new_tokens = [""] - output_tokens = prev_tokens + new_tokens + output_token_count = len(prev_tokens) + len(new_tokens) + window_start = ( + 0 if is_first_iter else min(prefix_offset, read_offset, len(prev_tokens)) + ) + output_tokens = ( + prev_tokens[window_start:] if window_start else prev_tokens + ) + new_tokens # If this is the first iteration, return all tokens. if is_first_iter: @@ -241,19 +247,21 @@ def detokenize_incrementally( # surrounding ids. if tokenizer.is_fast or not tokenizer.get_added_vocab(): prefix_text = tokenizer.convert_tokens_to_string( - output_tokens[prefix_offset:read_offset] + output_tokens[prefix_offset - window_start : read_offset - window_start] + ) + new_text = tokenizer.convert_tokens_to_string( + output_tokens[prefix_offset - window_start :] ) - new_text = tokenizer.convert_tokens_to_string(output_tokens[prefix_offset:]) else: prefix_text = _convert_tokens_to_string_with_added_encoders( tokenizer, - output_tokens[prefix_offset:read_offset], + output_tokens[prefix_offset - window_start : read_offset - window_start], skip_special_tokens=skip_special_tokens, spaces_between_special_tokens=spaces_between_special_tokens, ) new_text = _convert_tokens_to_string_with_added_encoders( tokenizer, - output_tokens[prefix_offset:], + output_tokens[prefix_offset - window_start :], skip_special_tokens=skip_special_tokens, spaces_between_special_tokens=spaces_between_special_tokens, ) @@ -266,4 +274,4 @@ def detokenize_incrementally( return new_tokens, "", prefix_offset, read_offset new_text = new_text[len(prefix_text) :] - return new_tokens, new_text, read_offset, len(output_tokens) + return new_tokens, new_text, read_offset, output_token_count From ab905a885dfbfc60a2c02286cc9c608c93884de3 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Thu, 8 Oct 2026 18:22:20 -0400 Subject: [PATCH 136/624] [CI] Run GH200 test on the shared arm64 CI image (#60253) Signed-off-by: mgoin Co-authored-by: Claude Opus 5.5 Co-authored-by: Kevin H. Luu --- .buildkite/hardware_tests/gh200.yaml | 13 ++++++--- .../scripts/hardware_ci/run-gh200-test.sh | 29 ------------------- 2 files changed, 9 insertions(+), 33 deletions(-) delete mode 100644 .buildkite/scripts/hardware_ci/run-gh200-test.sh diff --git a/.buildkite/hardware_tests/gh200.yaml b/.buildkite/hardware_tests/gh200.yaml index 2d42aa45416c..54f2e276327e 100644 --- a/.buildkite/hardware_tests/gh200.yaml +++ b/.buildkite/hardware_tests/gh200.yaml @@ -2,10 +2,15 @@ group: Hardware steps: - label: "GH200 Test" key: gh200-test + timeout_in_minutes: 30 soft_fail: true device: gh200 - no_plugin: true optional: true - commands: - - nvidia-smi - - bash .buildkite/scripts/hardware_ci/run-gh200-test.sh + working_dir: "/vllm-workspace" + depends_on: + - arm64-image-build + env: + VLLM_WORKER_MULTIPROC_METHOD: "spawn" + commands: + - nvidia-smi + - python3 examples/basic/offline_inference/generate.py --model meta-llama/Llama-3.2-1B diff --git a/.buildkite/scripts/hardware_ci/run-gh200-test.sh b/.buildkite/scripts/hardware_ci/run-gh200-test.sh deleted file mode 100644 index 20295fc25e90..000000000000 --- a/.buildkite/scripts/hardware_ci/run-gh200-test.sh +++ /dev/null @@ -1,29 +0,0 @@ -#!/bin/bash - -# This script build the GH200 docker image and run the offline inference inside the container. -# It serves a sanity check for compilation and basic model usage. -set -ex - -# Skip the new torch installation during build since we are using the specified version for arm64 in the Dockerfile -python3 tools/use_existing_torch.py - -# Try building the docker image -DOCKER_BUILDKIT=1 docker build . \ - --file docker/Dockerfile \ - --target vllm-openai \ - --platform "linux/arm64" \ - -t gh200-test \ - --build-arg max_jobs=66 \ - --build-arg nvcc_threads=2 \ - --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0-b8b5f17a7d9ccfc25bbc5cf17b3fcea12964a042 \ - --build-arg torch_cuda_arch_list="9.0+PTX" - -# Setup cleanup -remove_docker_container() { docker rm -f gh200-test || true; } -trap remove_docker_container EXIT -remove_docker_container - -# Run the image and test offline inference -docker run -e HF_TOKEN -e VLLM_WORKER_MULTIPROC_METHOD=spawn -v /root/.cache/huggingface:/root/.cache/huggingface --name gh200-test --gpus=all --entrypoint="" gh200-test bash -c ' - python3 examples/basic/offline_inference/generate.py --model meta-llama/Llama-3.2-1B -' From 4274ae956c7df16b07f6d68ba400e6321de1f47e Mon Sep 17 00:00:00 2001 From: DGlan <162426117+DGlan@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:38:16 -0400 Subject: [PATCH 137/624] [Bugfix][EC Connector] Release shared memory on NIXL setup failure (#59875) Signed-off-by: DGlan <162426117+DGlan@users.noreply.github.com> --- .../unit/test_scheduler_nixl_ctor.py | 68 +++++++++++++++++++ .../ec_connector/cpu/scheduler/__init__.py | 53 +++++++-------- 2 files changed, 92 insertions(+), 29 deletions(-) diff --git a/tests/v1/ec_connector/unit/test_scheduler_nixl_ctor.py b/tests/v1/ec_connector/unit/test_scheduler_nixl_ctor.py index 7d2622c44ebc..8cad9ec65c4d 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_ctor.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_ctor.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import uuid +from pathlib import Path import pytest import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod +import vllm.distributed.nixl_utils as nixl_utils from tests.v1.ec_connector.unit.utils import create_ec_vllm_config from vllm.config.ec_transfer import ECTransferConfig from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import ( @@ -57,6 +59,72 @@ def test_string_false_in_extra_config_leaves_nixl_off(monkeypatch): s.shutdown() +@pytest.mark.cpu_test +def test_missing_nixl_releases_shared_region(monkeypatch): + """A failed constructor leaves no scheduler for the caller to shut down.""" + # Import before patching to avoid caching the missing-NIXL values. + from vllm.distributed.ec_transfer.ec_connector.cpu.data import nixl # noqa: F401 + + region = _region() + monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region) + monkeypatch.setattr(nixl_utils, "NixlWrapper", None) + monkeypatch.setattr(nixl_utils, "nixl_agent_config", None) + cfg = create_ec_vllm_config(ec_role="ec_both") + cfg.ec_transfer_config.ec_connector_extra_config["ec_enable_nixl"] = True + + try: + with pytest.raises(RuntimeError, match="ec_enable_nixl requires NIXL"): + ECCPUScheduler(cfg) + + assert not Path(region._mmap_path).exists() + assert region._fd is None + assert region._mmap_obj is None + finally: + region.cleanup() + + +@pytest.mark.cpu_test +def test_control_transport_failure_releases_nixl_and_shared_region(monkeypatch): + import vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq as zmq_mod + import vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl as data_mod + + class RegisteredData: + registered = True + + def deregister(self): + self.registered = False + + data = RegisteredData() + error = RuntimeError("control transport failed") + + def fail_control_transport(): + raise error + + region = _region() + monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region) + monkeypatch.setattr(nixl_utils, "NixlWrapper", object()) + monkeypatch.setattr(nixl_utils, "nixl_agent_config", object()) + monkeypatch.setattr(data_mod, "NixlDataTransport", lambda **kwargs: data) + monkeypatch.setattr(zmq_mod, "ZmqClientTransport", fail_control_transport) + cfg = create_ec_vllm_config(ec_role="ec_consumer") + cfg.model_config.hf_config = None + cfg.model_config.get_inputs_embeds_size.return_value = 8 + cfg.model_config.model = "test-model" + cfg.ec_transfer_config.ec_connector_extra_config["ec_enable_nixl"] = True + + try: + with pytest.raises(RuntimeError, match="control transport failed") as exc: + ECCPUScheduler(cfg) + + assert exc.value is error + assert not data.registered + assert not Path(region._mmap_path).exists() + assert region._fd is None + assert region._mmap_obj is None + finally: + region.cleanup() + + @pytest.mark.skipif(NixlWrapper is None, reason="Requires NIXL package") def test_gate_on_wires_data_transport_and_producer_session(monkeypatch): # Port 0 lets the OS pick an ephemeral port so the real ZMQ ROUTER bind diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py index f1d22a5714a6..65403980b9d0 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py @@ -115,7 +115,11 @@ def __init__(self, vllm_config: "VllmConfig") -> None: self._element_size: int = 0 self._ack_timeout_s: float = 0.0 if self._nixl_enabled: - self._setup_nixl(vllm_config) + try: + self._setup_nixl(vllm_config) + except Exception: + self.shutdown() + raise def _setup_nixl(self, vllm_config: "VllmConfig") -> None: # Lazy imports keep nixl/zmq off the gate-off path. @@ -165,35 +169,26 @@ def _setup_nixl(self, vllm_config: "VllmConfig") -> None: self._peer_host = envs.VLLM_EC_SIDE_CHANNEL_HOST self._peer_port = envs.VLLM_EC_SIDE_CHANNEL_PORT - # Registering the region with NIXL and binding the control sockets are - # the first side effects here. __init__ propagates a failure, so the - # caller never receives a scheduler it could shut down: unwind through - # the same teardown shutdown() uses. - try: - self._data = NixlDataTransport( - agent_name=engine_id, - base_ptr=self._region.blocks.data_ptr(), - num_blocks=self._region.num_blocks, - block_size_bytes=self._region.block_size_bytes, - total_size_bytes=self._region.num_blocks - * self._region.block_size_bytes, + self._data = NixlDataTransport( + agent_name=engine_id, + base_ptr=self._region.blocks.data_ptr(), + num_blocks=self._region.num_blocks, + block_size_bytes=self._region.block_size_bytes, + total_size_bytes=self._region.num_blocks * self._region.block_size_bytes, + ) + if self._is_producer: + assert self._peer_host is not None + assert self._peer_port is not None + self._producer_session = ProducerSession( + transport=ZmqServerTransport( + host=self._peer_host, port=self._peer_port + ), + data=self._data, + cache=self._cache, + compat_hash=self._compat_hash, ) - if self._is_producer: - assert self._peer_host is not None - assert self._peer_port is not None - self._producer_session = ProducerSession( - transport=ZmqServerTransport( - host=self._peer_host, port=self._peer_port - ), - data=self._data, - cache=self._cache, - compat_hash=self._compat_hash, - ) - if self._is_consumer: - self._transport = ZmqClientTransport() - except Exception: - self._teardown_nixl() - raise + if self._is_consumer: + self._transport = ZmqClientTransport() def has_cache_item(self, identifier: str) -> bool: if not self._is_consumer: From 581692cc0fd564546ca1f58475e795e2a154cb3f Mon Sep 17 00:00:00 2001 From: Nishant Revur Date: Thu, 8 Oct 2026 18:44:50 -0700 Subject: [PATCH 138/624] [Bugfix][Spec Decode] Fix multimodal embedding merge in Llama 4 and Mistral Large 3 EAGLE drafters (#57633) Signed-off-by: Nishant Revur Co-authored-by: Claude --- ...test_gpu_target_dependent_ar_speculator.py | 24 +++++++++++++++ vllm/model_executor/models/llama4_eagle.py | 23 +++++++++++++-- .../models/mistral_large_3_eagle.py | 29 +++++++++++++++++-- 3 files changed, 71 insertions(+), 5 deletions(-) diff --git a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py index 20e86f4a7827..b63bd7a2ff05 100644 --- a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py +++ b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py @@ -337,6 +337,30 @@ def test_draft_model_multimodal_embedding_capability(model_cls, expected): assert supports_multimodal_embeddings(model_cls) is expected +@pytest.mark.parametrize( + "model_cls", [EagleLlama4ForCausalLM, EagleMistralLarge3ForCausalLM] +) +def test_mm_capable_drafter_merges_multimodal_embeddings(model_cls): + """Drafters advertising multimodal embeddings must embed text ids and + scatter the multimodal embeddings themselves; they do not inherit + SupportsMultiModal's helpers.""" + table = torch.arange(8, dtype=torch.float32).repeat(3, 1).T # [8, 3] + drafter = object.__new__(model_cls) + drafter.model = SimpleNamespace(embed_input_ids=lambda ids: table[ids]) + input_ids = torch.tensor([1, 5, 6, 2]) + + text_only = model_cls.embed_input_ids(drafter, input_ids) + assert torch.equal(text_only, table[input_ids]) + + mm_embeds = [torch.full((2, 3), -1.0)] + is_multimodal = torch.tensor([False, True, True, False]) + merged = model_cls.embed_input_ids( + drafter, input_ids, mm_embeds, is_multimodal=is_multimodal + ) + assert torch.equal(merged[[0, 3]], table[[1, 2]]) + assert torch.equal(merged[[1, 2]], mm_embeds[0]) + + def test_run_model_unpacks_tuple_return_for_mtp(monkeypatch): logits_hidden = torch.full((4, 3), 1.0) feedback_hidden = torch.full((4, 3), 2.0) diff --git a/vllm/model_executor/models/llama4_eagle.py b/vllm/model_executor/models/llama4_eagle.py index 975ee28dc22e..8634f113d063 100644 --- a/vllm/model_executor/models/llama4_eagle.py +++ b/vllm/model_executor/models/llama4_eagle.py @@ -36,10 +36,15 @@ from vllm.model_executor.models.llama4 import Llama4DecoderLayer, Llama4ForCausalLM from vllm.model_executor.models.utils import extract_layer_index -from .interfaces import SupportsMultiModal, SupportsMultiModalEmbeddings +from .interfaces import ( + MultiModalEmbeddings, + SupportsMultiModalEmbeddings, + _require_is_multimodal, +) from .utils import ( AutoWeightsLoader, WeightsMapper, + _merge_multimodal_embeddings, maybe_prefix, process_eagle_weight, ) @@ -204,7 +209,21 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): def get_language_model(self) -> torch.nn.Module: return self.model - embed_input_ids = SupportsMultiModal.embed_input_ids # type: ignore + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + inputs_embeds = self.model.embed_input_ids(input_ids) + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + return _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=_require_is_multimodal(is_multimodal), + ) def forward( self, diff --git a/vllm/model_executor/models/mistral_large_3_eagle.py b/vllm/model_executor/models/mistral_large_3_eagle.py index d2e263f3e7d1..1398fd19fc91 100644 --- a/vllm/model_executor/models/mistral_large_3_eagle.py +++ b/vllm/model_executor/models/mistral_large_3_eagle.py @@ -23,8 +23,17 @@ ) from vllm.model_executor.models.mistral_large_3 import MistralLarge3ForCausalLM -from .interfaces import SupportsMultiModal, SupportsMultiModalEmbeddings -from .utils import WeightsMapper, make_empty_intermediate_tensors_factory, maybe_prefix +from .interfaces import ( + MultiModalEmbeddings, + SupportsMultiModalEmbeddings, + _require_is_multimodal, +) +from .utils import ( + WeightsMapper, + _merge_multimodal_embeddings, + make_empty_intermediate_tensors_factory, + maybe_prefix, +) logger = init_logger(__name__) @@ -148,7 +157,21 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): def get_language_model(self) -> torch.nn.Module: return self.model - embed_input_ids = SupportsMultiModal.embed_input_ids # type: ignore + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + inputs_embeds = self.model.embed_input_ids(input_ids) + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + return _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=_require_is_multimodal(is_multimodal), + ) def forward( self, From 201e537d98aa2ec75f12134e78f728fb0b80d542 Mon Sep 17 00:00:00 2001 From: "Brian Weston, PhD" Date: Thu, 8 Oct 2026 18:49:13 -0700 Subject: [PATCH 139/624] [Bugfix][LoRA] Validate incompatible MoE adapter formats (#59924) Signed-off-by: Brian Weston Co-authored-by: Codex Co-authored-by: Jee Jee Li Co-authored-by: linitra24 --- docs/features/lora.md | 13 ++++--- tests/lora/test_lora_manager.py | 64 +++++++++++++++++++++++++++++++++ vllm/lora/model_manager.py | 22 ++++++++++++ vllm/lora/worker_manager.py | 1 + 4 files changed, 96 insertions(+), 4 deletions(-) diff --git a/docs/features/lora.md b/docs/features/lora.md index 4c95bde8cd9f..c05bc8c0fa7a 100644 --- a/docs/features/lora.md +++ b/docs/features/lora.md @@ -280,10 +280,10 @@ curl -X POST http://localhost:8000/v1/load_lora_adapter \ !!! warning "You must know your adapter's layout" Under `--enable-mixed-moe-lora-format`, vLLM trusts whatever - `is_3d_lora_weight` the caller declares — it does **not** inspect the - checkpoint to verify. A wrong declaration will load weights into the - wrong stacked buffers and silently produce garbage outputs, with no - error at load time. Confirm the layout before serving: + `is_3d_lora_weight` the caller declares. Loading checks reject fused 3D + expert weights routed to a 2D wrapper without the required flags, but + do not validate every layout mismatch. A wrong declaration can still + cause loading failures or incorrect outputs. Confirm the layout before serving: - **2D (per-expert, megatron-style)** → set `is_3d_lora_weight: false`. Adapter keys look like `...experts.{idx}.gate_proj.lora_A.weight`, @@ -299,6 +299,11 @@ is ignored: vLLM picks the wrapper from the base model's `is_3d_moe_weight` and the adapter is required to match. The field is also ignored for non-MoE models. +For a model using a 2D MoE wrapper (such as Qwen3-MoE), loading fused 3D +expert adapter weights requires both `enable_mixed_moe_lora_format=True` +on the engine and `is_3d_lora_weight=True` on the request. Otherwise, adapter +loading raises a `ValueError` naming these settings before activation. + ## LoRA model lineage in model card The new format of `--lora-modules` is mainly to support the display of parent model information in the model card. Here's an explanation of how your current response supports this: diff --git a/tests/lora/test_lora_manager.py b/tests/lora/test_lora_manager.py index 43fa1975bd79..ffb064ebfbdc 100644 --- a/tests/lora/test_lora_manager.py +++ b/tests/lora/test_lora_manager.py @@ -16,10 +16,12 @@ ReplicatedLinearWithLoRA, RowParallelLinearWithLoRA, ) +from vllm.lora.layers.fused_moe import FusedMoE3DWithLoRA, FusedMoEWithLoRA from vllm.lora.lora_model import LoRAModel from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights from vllm.lora.model_manager import ( DEFAULT_LANGUAGE_WRAPPER_KEY, + AdapterLRUCache, LoRAMapping, LoRAModelManager, LRUCacheLoRAModelManager, @@ -47,6 +49,68 @@ DEFAULT_DTYPE = torch.get_default_dtype() +@pytest.mark.parametrize( + "mixed,declared_3d", [(False, False), (False, True), (True, False)] +) +def test_reject_fused_moe_lora_before_registration(mixed, declared_3d): + """Missing format flags should produce guidance without caching the adapter.""" + module_name = "model.layers.0.mlp.experts" + manager = LoRAModelManager.__new__(LoRAModelManager) + manager.lora_config = LoRAConfig(max_cpu_loras=1) + manager._registered_adapters = AdapterLRUCache(1, manager.deactivate_adapter) + manager._enable_mixed_moe_lora_format = mixed + manager.is_pooling_model = False + manager.modules = {module_name: FusedMoEWithLoRA.__new__(FusedMoEWithLoRA)} + adapter = LoRAModel( + 1, + 8, + { + module_name: LoRALayerWeights( + module_name, 8, 16, torch.zeros(16, 4), torch.zeros(4, 16) + ) + }, + ) + adapter.is_3d_lora_weight = declared_3d + + with pytest.raises(ValueError, match="enable_mixed_moe_lora_format=True") as exc: + manager.add_adapter(adapter) + assert "is_3d_lora_weight=True" in str(exc.value) + assert module_name in str(exc.value) + assert manager.list_adapters() == {} + + +@pytest.mark.parametrize( + "wrapper,mixed,declared_3d,weight_name", + [ + (FusedMoE3DWithLoRA, False, False, "experts"), + (FusedMoEWithLoRA, True, True, "experts"), + (FusedMoEWithLoRA, False, False, "experts.0.down_proj"), + (FusedMoEWithLoRA, True, False, "experts.0.down_proj"), + (FusedMoEWithLoRA, False, False, "q_proj"), + ], +) +def test_moe_format_validation_allows_supported_layouts( + wrapper, mixed, declared_3d, weight_name +): + """Preserve native 3D, mixed-format, per-expert and attention-only adapters.""" + manager = LoRAModelManager.__new__(LoRAModelManager) + manager._enable_mixed_moe_lora_format = mixed + manager.is_pooling_model = False + manager.modules = {"experts": wrapper.__new__(wrapper)} + adapter = LoRAModel( + 1, + 8, + { + weight_name: LoRALayerWeights( + weight_name, 8, 16, torch.zeros(8, 4), torch.zeros(4, 8) + ) + }, + ) + adapter.is_3d_lora_weight = declared_3d + + manager._validate_moe_lora_format(adapter) + + @pytest.mark.parametrize("device", DEVICES) def test_from_lora_tensors(qwen3_lora_files, device): tensors = load_file(os.path.join(qwen3_lora_files, "adapter_model.safetensors")) diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index b97a0a412a16..8dc7b5605600 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -414,6 +414,7 @@ def _deactivate_adapter(self, lora_id: int): pass def _add_adapter(self, lora: LoRAModel): + self._validate_moe_lora_format(lora) self._create_merged_loras_inplace(lora) self._registered_adapters[lora.id] = lora @@ -1252,6 +1253,27 @@ def _get_module_to_save_weights( ) return weights + def _validate_moe_lora_format(self, lora_model: LoRAModel) -> None: + if self._enable_mixed_moe_lora_format and getattr( + lora_model, "is_3d_lora_weight", False + ): + return + for module_name, module in self.modules.items(): + if not isinstance(module, FusedMoEWithLoRA) or isinstance( + module, FusedMoE3DWithLoRA + ): + continue + module_lora = self._get_lora_layer_weights(lora_model, module_name) + if module_lora is not None and torch.is_tensor(module_lora.lora_a): + raise ValueError( + f"LoRA adapter {lora_model.id} contains fused 3D MoE weights " + f"for {module_name!r}, but the model uses a 2D MoE LoRA " + "wrapper. Start the engine with " + "enable_mixed_moe_lora_format=True " + "(--enable-mixed-moe-lora-format) and set " + "is_3d_lora_weight=True on the LoRA request." + ) + def _validate_modules_to_save(self, lora_model: LoRAModel) -> None: if not lora_model.modules_to_save: return diff --git a/vllm/lora/worker_manager.py b/vllm/lora/worker_manager.py index fe266879a859..9427b3e1341f 100644 --- a/vllm/lora/worker_manager.py +++ b/vllm/lora/worker_manager.py @@ -153,6 +153,7 @@ def _load_adapter(self, lora_request: LoRARequest) -> LoRAModel: # adapter manager can route 3D-format checkpoints through the # 3D->2D conversion when running under the universal 2D wrapper. lora.is_3d_lora_weight = lora_request.is_3d_lora_weight + self._adapter_manager._validate_moe_lora_format(lora) # Validate classification-head weights. self._adapter_manager._validate_modules_to_save(lora) From 4611c2c3d20e6f55f65c96d87f921b1be3005be9 Mon Sep 17 00:00:00 2001 From: Mia Ojeda Date: Thu, 8 Oct 2026 20:16:51 -0600 Subject: [PATCH 140/624] [BugFix][DiffusionGemma][CPU] Copy async output snapshots instead of aliasing reused buffers (#59829) Signed-off-by: ojeda-e Co-authored-by: Lucas Wilkinson --- tests/distributed/aux_output_connector/test_store.py | 3 --- tests/v1/worker/test_gpu_model_runner_v2.py | 10 ++++++++++ vllm/v1/worker/gpu/async_utils.py | 2 +- 3 files changed, 11 insertions(+), 4 deletions(-) diff --git a/tests/distributed/aux_output_connector/test_store.py b/tests/distributed/aux_output_connector/test_store.py index 24d67c3d60e1..62d4472811fc 100644 --- a/tests/distributed/aux_output_connector/test_store.py +++ b/tests/distributed/aux_output_connector/test_store.py @@ -251,9 +251,6 @@ def test_next_step_does_not_consume_pending_output(monkeypatch): event = Mock() monkeypatch.setattr(torch.cuda, "Event", lambda **kwargs: event) monkeypatch.setattr(async_utils, "stream", lambda *args: nullcontext()) - monkeypatch.setattr( - async_utils, "async_copy_to_np", lambda tensor: tensor.numpy().copy() - ) worker = _make_worker(1) worker.begin_step( _metadata(0, [_request_metadata("request", 0, 1, 0, [])], {}).metadata diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index a10f54a97c75..dec382725bc1 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -18,6 +18,7 @@ MambaSpec, UniformTypeKVCacheSpecs, ) +from vllm.v1.worker.gpu.async_utils import async_copy_to_np from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.model_runner import ExecuteModelState, GPUModelRunner @@ -357,3 +358,12 @@ def test_get_drafter_hidden_states_tolerates_missing_target_buffer(target_buffer assert torch.equal(out, buffer[:4]) else: assert out is hidden_states + + +def test_async_copy_to_np_does_not_alias_reused_buffer(): + buffer = torch.zeros(4, dtype=torch.int64) + + snapshot = async_copy_to_np(buffer) + buffer.fill_(1) + + assert snapshot.tolist() == [0, 0, 0, 0] diff --git a/vllm/v1/worker/gpu/async_utils.py b/vllm/v1/worker/gpu/async_utils.py index 647b0c297ccc..0d6481585e21 100644 --- a/vllm/v1/worker/gpu/async_utils.py +++ b/vllm/v1/worker/gpu/async_utils.py @@ -262,7 +262,7 @@ def get_output(self) -> ModelRunnerOutput: def async_copy_to_np(x: torch.Tensor) -> np.ndarray: - return x.to("cpu", non_blocking=True).numpy() + return x.to("cpu", non_blocking=True, copy=True).numpy() @contextlib.contextmanager From 295c062db8a48ea73bbd05d1174cb37d24bf86f9 Mon Sep 17 00:00:00 2001 From: Tiezhen WANG <38108242+xianbaoqian@users.noreply.github.com> Date: Thu, 8 Oct 2026 19:24:56 -0700 Subject: [PATCH 141/624] Remove sudo from dockerfile command (#60488) Signed-off-by: Tiezhen WANG <38108242+xianbaoqian@users.noreply.github.com> Co-authored-by: Kevin H. Luu --- docker/Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index fb8fa6b264ae..d524d0fddd47 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -153,7 +153,7 @@ RUN if [ "$USE_SCCACHE" = "1" ]; then \ && export SCCACHE_DOWNLOAD_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/v0.8.1/sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ && curl -L -o sccache.tar.gz ${SCCACHE_DOWNLOAD_URL} \ && tar -xzf sccache.tar.gz \ - && sudo mv sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ + && mv sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ && rm -rf sccache.tar.gz sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl; \ fi From 702ce313f4f85f5f2e76b7796aa5cdc80f3595a1 Mon Sep 17 00:00:00 2001 From: zzaebok <44357534+zzaebok@users.noreply.github.com> Date: Fri, 9 Oct 2026 10:55:36 +0800 Subject: [PATCH 142/624] [Perf][Multimodal] Sort multimodal encoder inputs by modality (#60196) Signed-off-by: Jaebok Lee Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/worker/test_encoder_runner.py | 17 ++++++- tests/v1/worker/test_gpu_model_runner.py | 64 ++++++++++++++++++++++++ vllm/v1/worker/gpu/mm/encoder_runner.py | 7 ++- vllm/v1/worker/gpu/mm/lora.py | 35 +++++++------ vllm/v1/worker/gpu_model_runner.py | 16 +++--- 5 files changed, 112 insertions(+), 27 deletions(-) diff --git a/tests/v1/worker/test_encoder_runner.py b/tests/v1/worker/test_encoder_runner.py index 98aaa568f0b6..69afc6ea93e7 100644 --- a/tests/v1/worker/test_encoder_runner.py +++ b/tests/v1/worker/test_encoder_runner.py @@ -78,7 +78,7 @@ def test_encoder_graph_unsupported_modality_uses_eager_output(): "vllm.v1.worker.gpu.mm.encoder_runner.group_and_batch_mm_kwargs", return_value=[("audio", 1, {"audio_values": torch.zeros(2)})], ): - result = runner.execute_mm_encoder([]) + result = runner.execute_mm_encoder([("audio", MagicMock())]) assert len(result) == 1 assert result[0] is output manager.execute.assert_not_called() @@ -290,6 +290,21 @@ def test_execute_mm_encoder_caches_outputs_without_gathering(): state.encoder_runner.gather_mm_embeddings.assert_not_called() +def test_execute_mm_encoder_sorts_by_modality(): + """Items are encoded sorted by modality; outputs keep the input order.""" + mm_kwargs = [("video", 0), ("image", 1), ("video", 2), ("image", 3)] + runner = _make_runner([], []) + runner.model = MagicMock() + runner.model.embed_multimodal.side_effect = lambda x: [torch.full((1, 1), x)] + with patch( + "vllm.v1.worker.gpu.mm.encoder_runner.group_and_batch_mm_kwargs", + side_effect=lambda items, **_: [(m, 1, {"x": x}) for m, x in items], + ) as group: + outputs = runner.execute_mm_encoder(mm_kwargs) + assert group.call_args.args[0] == [mm_kwargs[i] for i in (1, 3, 0, 2)] + assert [int(o) for o in outputs] == [0, 1, 2, 3] + + def test_execute_mm_encoder_is_a_noop_without_scheduled_items(): """A step that schedules no encoder input must not touch the encoder.""" cache = EncoderCache() diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index deb4f7b42147..3851d5dad6bd 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -624,6 +624,70 @@ def test_set_active_mm_loras_builds_tower_and_connector_mappings(): assert connector_mapping.index_mapping == ((7,) * 14 + (7,) * 13 + (0,) * 12) +def test_set_active_mm_loras_follows_modality_sorted_encoder_order(): + model = Mock() + model.get_mm_lora_token_counts.side_effect = ( + lambda *, modality, mm_kwargs, num_mm_embeds: (num_mm_embeds, None) + ) + lora_manager = Mock() + lora_manager.supports_tower_connector_lora.return_value = True + + encoder_cache = EncoderCache() + for req_id, modality, length in [("req-a", "video", 1), ("req-b", "image", 2)]: + encoder_cache.mm_features[req_id] = [ + MultiModalFeatureSpec( + data=None, + modality=modality, + identifier=f"{req_id}-0", + mm_position=PlaceholderRange(offset=0, length=length), + ) + ] + lora_state = LoraState(max_num_reqs=2) + lora_request = LoRARequest("vision-lora", 7, "/tmp/vision-lora") + lora_state.add_request("req-a", 0, lora_request) + lora_state.add_request("req-b", 1, None) + + set_active_mm_loras( + model=model, + lora_manager=lora_manager, + encoder_cache=encoder_cache, + req_id_to_index={"req-a": 0, "req-b": 1}, + lora_state=lora_state, + scheduled_encoder_inputs={"req-a": [0], "req-b": [0]}, + ) + + # The image is encoded before the video, so its tokens come first. + _, tower_mapping = lora_manager.set_active_adapters.call_args_list[0].args + assert tower_mapping.index_mapping == (0, 0, 7) + + +def test_batch_mm_inputs_from_scheduler_sorts_by_modality(): + features = { + f"req{i}": MultiModalFeatureSpec( + data=Mock(), + modality=m, + identifier=f"hash{i}", + mm_position=PlaceholderRange(offset=i, length=1), + ) + for i, m in enumerate(["video", "image", "video"]) + } + runner = SimpleNamespace( + requests={r: SimpleNamespace(mm_features=[f]) for r, f in features.items()} + ) + scheduler_output = SimpleNamespace( + scheduled_encoder_inputs={r: [0] for r in features} + ) + + mm_hashes, mm_kwargs, mm_lora_refs = GPUModelRunner._batch_mm_inputs_from_scheduler( + runner, scheduler_output + ) + + order = ["req1", "req0", "req2"] + assert mm_hashes == [features[r].identifier for r in order] + assert mm_kwargs == [(features[r].modality, features[r].data) for r in order] + assert mm_lora_refs == [(r, features[r].mm_position) for r in order] + + def test_update_states_new_request(model_runner, dist_init): req_id = "req_0" diff --git a/vllm/v1/worker/gpu/mm/encoder_runner.py b/vllm/v1/worker/gpu/mm/encoder_runner.py index 16e42d140af2..4d7ab9201394 100644 --- a/vllm/v1/worker/gpu/mm/encoder_runner.py +++ b/vllm/v1/worker/gpu/mm/encoder_runner.py @@ -150,9 +150,12 @@ def profile_encoder_cache( def execute_mm_encoder( self, mm_kwargs: list[tuple[str, MultiModalKwargsItem]] ) -> list[torch.Tensor]: + # Stable-sort items by modality so each modality is encoded in as few + # batches as possible. The returned outputs are restored to input order. + order = sorted(range(len(mm_kwargs)), key=lambda i: mm_kwargs[i][0]) encoder_outputs: list[torch.Tensor] = [] for modality, num_items, mm_kwargs_batch in group_and_batch_mm_kwargs( - mm_kwargs, device=self.device, pin_memory=PIN_MEMORY + [mm_kwargs[i] for i in order], device=self.device, pin_memory=PIN_MEMORY ): cg_manager = self.cudagraph_manager cudagraph_output = ( @@ -169,7 +172,7 @@ def execute_mm_encoder( ) sanity_check_mm_encoder_outputs(batch_outputs, expected_num_items=num_items) encoder_outputs.extend(batch_outputs) - return encoder_outputs + return [encoder_outputs[i] for i in np.argsort(order)] @contextmanager def timed_encoder_operation(self, request_ids: Collection[str]): diff --git a/vllm/v1/worker/gpu/mm/lora.py b/vllm/v1/worker/gpu/mm/lora.py index 819f3b1cb36a..e73e028f30f0 100644 --- a/vllm/v1/worker/gpu/mm/lora.py +++ b/vllm/v1/worker/gpu/mm/lora.py @@ -6,6 +6,7 @@ from vllm.lora.layers import LoRAMapping, LoRAMappingType from vllm.lora.worker_manager import WorkerLoRAManager +from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.v1.worker.gpu.lora_utils import LoraState from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache @@ -31,7 +32,7 @@ def set_active_mm_loras( encoder_token_counts: list[int] = [] connector_token_counts: list[int] = [] - # iterate through images + mm_items: list[tuple[int, MultiModalFeatureSpec]] = [] for req_id, encoder_input_ids in scheduled_encoder_inputs.items(): req_idx = req_id_to_index.get(req_id) if req_idx is None: @@ -39,27 +40,29 @@ def set_active_mm_loras( lora_id = int(lora_state.lora_ids[req_idx]) mm_features = encoder_cache.mm_features[req_id] - - # iterate through visual tokens - for mm_input_id in encoder_input_ids: - pos_info = mm_features[mm_input_id].mm_position - - tower_tokens, connector_tokens = model.get_mm_lora_token_counts( - modality=mm_features[mm_input_id].modality, - mm_kwargs=mm_features[mm_input_id].data, - num_mm_embeds=pos_info.get_num_embeds(), - ) - - prompt_lora_mapping.append(lora_id) - token_lora_mapping.extend([lora_id] * tower_tokens) - encoder_token_counts.append(tower_tokens) - connector_token_counts.append(connector_tokens) + mm_items.extend((lora_id, mm_features[i]) for i in encoder_input_ids) if lora_id > 0: lora_request = lora_state.lora_requests.get(req_id) if lora_request is not None: lora_requests.add(lora_request) + # Runs items sorted by modality. + mm_items.sort(key=lambda x: x[1].modality) + for lora_id, mm_feature in mm_items: + pos_info = mm_feature.mm_position + + tower_tokens, connector_tokens = model.get_mm_lora_token_counts( + modality=mm_feature.modality, + mm_kwargs=mm_feature.data, + num_mm_embeds=pos_info.get_num_embeds(), + ) + + prompt_lora_mapping.append(lora_id) + token_lora_mapping.extend([lora_id] * tower_tokens) + encoder_token_counts.append(tower_tokens) + connector_token_counts.append(connector_tokens) + if not prompt_lora_mapping: return diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 21b0f67e2fe8..c02bdc74e22b 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -2929,7 +2929,14 @@ def _batch_mm_inputs_from_scheduler( mm_kwargs.append((mm_feature.modality, mm_feature.data)) mm_lora_refs.append((req_id, mm_feature.mm_position)) - return mm_hashes, mm_kwargs, mm_lora_refs + # Stable-sort by modality so each modality is encoded in as few batches + # as possible. Encoder outputs are cached by mm_hash, so order is free. + order = sorted(range(len(mm_kwargs)), key=lambda i: mm_kwargs[i][0]) + return ( + [mm_hashes[i] for i in order], + [mm_kwargs[i] for i in order], + [mm_lora_refs[i] for i in order], + ) def _cache_encoder_output( self, @@ -2990,13 +2997,6 @@ def _execute_mm_encoder( and scheduler_output.scheduled_encoder_inputs ) - # Batch mm inputs as much as we can: if a request in the batch has - # multiple modalities or a different modality than the previous one, - # we process it separately to preserve item order. - # FIXME(ywang96): This is a hacky way to deal with multiple modalities - # in the same batch while still being able to benefit from batching - # multimodal inputs. The proper solution should be reordering the - # encoder outputs. model = cast(SupportsMultiModal, self.model) if self.lora_config and self.lora_manager.supports_tower_connector_lora(): From a1cd0ec650852bf2ce124b3dea930a38f0d243da Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Delhaye?= Date: Fri, 9 Oct 2026 05:01:51 +0200 Subject: [PATCH 143/624] [Model] Add zerank-2 reranker (#60415) Signed-off-by: rdlhio --- docs/models/pooling_models/scoring.md | 9 +- examples/pooling/score/template/zerank2.jinja | 13 +++ .../pooling/score/using_template_offline.py | 2 + .../pooling/score/using_template_online.py | 1 + .../language/pooling_mteb_test/test_zerank.py | 95 +++++++++++++++++++ 5 files changed, 119 insertions(+), 1 deletion(-) create mode 100644 examples/pooling/score/template/zerank2.jinja create mode 100644 tests/models/language/pooling_mteb_test/test_zerank.py diff --git a/docs/models/pooling_models/scoring.md b/docs/models/pooling_models/scoring.md index 1c33c6852d7e..0198810cbf68 100644 --- a/docs/models/pooling_models/scoring.md +++ b/docs/models/pooling_models/scoring.md @@ -49,7 +49,7 @@ The three supported scoring functions are as illustrated in the figure below. | `LlamaBidirectionalForSequenceClassification`C | Llama-based with bidirectional attention | `nvidia/llama-nemotron-rerank-1b-v2`, etc. | [nemotron-rerank.jinja](../../../examples/pooling/score/template/nemotron-rerank.jinja) | ✅︎ | ✅︎ | | `ModernBertForSequenceClassification` | ModernBERT-based | `Alibaba-NLP/gte-reranker-modernbert-base`, etc. | N/A | | | | `Qwen2ForSequenceClassification`C | Qwen2-based | `mixedbread-ai/mxbai-rerank-base-v2`(see note), etc. | [mxbai_rerank_v2.jinja](../../../examples/pooling/score/template/mxbai_rerank_v2.jinja) | ✅︎ | ✅︎ | -| `Qwen3ForSequenceClassification`C | Qwen3-based | `tomaarsen/Qwen3-Reranker-0.6B-seq-cls`, `Qwen/Qwen3-Reranker-0.6B`(see note), etc. | [qwen3_reranker.jinja](../../../examples/pooling/score/template/qwen3_reranker.jinja) | ✅︎ | ✅︎ | +| `Qwen3ForSequenceClassification`C | Qwen3-based | `tomaarsen/Qwen3-Reranker-0.6B-seq-cls`, `Qwen/Qwen3-Reranker-0.6B`(see note), `polaria-tech/zerank-2-reranker-vllm`(see note), etc. | [qwen3_reranker.jinja](../../../examples/pooling/score/template/qwen3_reranker.jinja), [zerank2.jinja](../../../examples/pooling/score/template/zerank2.jinja) | ✅︎ | ✅︎ | | `RobertaForSequenceClassification` | RoBERTa-based | `cross-encoder/quora-roberta-base`, etc. | N/A | ✅︎ | | | `XLMRobertaForSequenceClassification` | XLM-RoBERTa-based | `BAAI/bge-reranker-v2-m3`, etc. | N/A | ✅︎ | | | `*Model`C, `*ForCausalLM`C, etc. | Generative models | N/A | N/A | \* | \* | @@ -88,6 +88,13 @@ The three supported scoring functions are as illustrated in the figure below. vllm serve Qwen/Qwen3-Reranker-0.6B --hf_overrides '{"architectures": ["Qwen3ForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}' ``` +!!! note + `polaria-tech/zerank-2-reranker-vllm` is [zeroentropy/zerank-2-reranker](https://huggingface.co/zeroentropy/zerank-2-reranker) converted to `Qwen3ForSequenceClassification`; it returns zerank-2's calibrated score `sigmoid(logit / 5)`. Serve it with [zerank2.jinja](../../../examples/pooling/score/template/zerank2.jinja), which also supports the `instruction` request field. + + ```bash + vllm serve polaria-tech/zerank-2-reranker-vllm --chat-template examples/pooling/score/template/zerank2.jinja + ``` + #### Multimodal Models !!! note diff --git a/examples/pooling/score/template/zerank2.jinja b/examples/pooling/score/template/zerank2.jinja new file mode 100644 index 000000000000..8bd16e430121 --- /dev/null +++ b/examples/pooling/score/template/zerank2.jinja @@ -0,0 +1,13 @@ +{#- zerank-2 (zeroentropy/zerank-2-reranker): query in the system turn, document in the user turn. + The assistant suffix is emitted unconditionally, since score templates are rendered + without add_generation_prompt. An optional `instruction` is wrapped as recommended by + ZeroEntropy; without it, the prompt is the official zerank-2 prompt. -#} +{%- set query = messages | selectattr("role", "eq", "query") | map(attribute="content") | list -%} +{%- set document = messages | selectattr("role", "eq", "document") | map(attribute="content") | list -%} +{%- set query = query[0] if query else '' -%} +{%- if instruction is defined and instruction is string and instruction | trim -%} +{%- set query = '' + query + '\n' + instruction + '' -%} +{%- endif -%} +{{- '<|im_start|>system\n' + query + '<|im_end|>\n' -}} +{{- '<|im_start|>user\n' + (document[0] if document else '') + '<|im_end|>\n' -}} +{{- '<|im_start|>assistant\n' -}} diff --git a/examples/pooling/score/using_template_offline.py b/examples/pooling/score/using_template_offline.py index abe8157d4c2f..1b7fe9f7c0ad 100644 --- a/examples/pooling/score/using_template_offline.py +++ b/examples/pooling/score/using_template_offline.py @@ -53,6 +53,7 @@ def get_chat_template(model: str) -> str: "mixedbread-ai/mxbai-rerank-base-v2": "mxbai_rerank_v2.jinja", "mixedbread-ai/mxbai-rerank-large-v2": "mxbai_rerank_v2.jinja", "nvidia/llama-nemotron-rerank-1b-v2": "nemotron-rerank.jinja", + "polaria-tech/zerank-2-reranker-vllm": "zerank2.jinja", } # Get the template filename for the specified model @@ -110,6 +111,7 @@ def get_hf_overrides(model: str) -> dict[str, Any]: "method": "from_2_way_softmax", }, "nvidia/llama-nemotron-rerank-1b-v2": {}, + "polaria-tech/zerank-2-reranker-vllm": {}, } hf_overrides = model_name_to_hf_overrides_map.get(model) diff --git a/examples/pooling/score/using_template_online.py b/examples/pooling/score/using_template_online.py index f7dc109dd81f..9cf0703ca475 100644 --- a/examples/pooling/score/using_template_online.py +++ b/examples/pooling/score/using_template_online.py @@ -19,6 +19,7 @@ vllm serve tomaarsen/Qwen3-Reranker-0.6B-seq-cls --chat-template examples/pooling/score/template/qwen3_reranker.jinja vllm serve mixedbread-ai/mxbai-rerank-base-v2 --hf_overrides '{"architectures": ["Qwen2ForSequenceClassification"],"classifier_from_token": ["0", "1"], "method": "from_2_way_softmax"}' --chat-template examples/pooling/score/template/mxbai_rerank_v2.jinja vllm serve nvidia/llama-nemotron-rerank-1b-v2 --runner pooling --trust-remote-code --chat-template examples/pooling/score/template/nemotron-rerank.jinja + vllm serve polaria-tech/zerank-2-reranker-vllm --runner pooling --chat-template examples/pooling/score/template/zerank2.jinja vllm serve Qwen/Qwen3-Reranker-0.6B --runner pooling --hf_overrides '{"architectures": ["Qwen3ForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}' --chat-template examples/pooling/score/template/qwen3_reranker.jinja """ diff --git a/tests/models/language/pooling_mteb_test/test_zerank.py b/tests/models/language/pooling_mteb_test/test_zerank.py new file mode 100644 index 000000000000..c90c80bc4d40 --- /dev/null +++ b/tests/models/language/pooling_mteb_test/test_zerank.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + +import mteb +import numpy as np +import pytest +import torch +from torch.utils.data import DataLoader + +from tests.conftest import HfRunner +from tests.models.utils import RerankModelInfo + +from .mteb_score_utils import MtebCrossEncoderMixin, mteb_test_rerank_models + +RERANK_MODELS = [ + RerankModelInfo( + "polaria-tech/zerank-2-reranker-vllm", + architecture="Qwen3ForSequenceClassification", + chat_template_name="zerank2.jinja", + seq_pooling_type="LAST", + attn_type="decoder", + is_prefix_caching_supported=True, + is_chunked_prefill_supported=True, + mteb_score=0.34305, + mteb_tol=1e-2, + enable_test=True, + ), +] + + +class ZerankHfRunner(MtebCrossEncoderMixin, HfRunner): + # The converted checkpoint has no lm_head, so the reference is the original + # causal LM: the logit of the "Yes" token at the last position. The converted + # checkpoint returns sigmoid(logit / 5), which ranks documents identically. + original_model = "zeroentropy/zerank-2-reranker" + + def __init__( + self, model_name: str, dtype: str = "auto", *args: Any, **kwargs: Any + ) -> None: + from transformers import AutoModelForCausalLM, AutoTokenizer + + HfRunner.__init__( + self, + model_name=self.original_model, + auto_cls=AutoModelForCausalLM, + dtype=dtype, + **kwargs, + ) + + self.tokenizer = AutoTokenizer.from_pretrained( + self.original_model, padding_side="left" + ) + self.token_true_id = self.tokenizer.convert_tokens_to_ids("Yes") + + @torch.no_grad + def predict( + self, + inputs1: DataLoader[mteb.types.BatchedInput], + inputs2: DataLoader[mteb.types.BatchedInput], + *args, + **kwargs, + ) -> np.ndarray: + queries = [text for batch in inputs1 for text in batch["text"]] + corpus = [text for batch in inputs2 for text in batch["text"]] + + tokenizer = self.tokenizer + prompts = [] + for query, document in zip(queries, corpus): + conversation = [ + {"role": "query", "content": query}, + {"role": "document", "content": document}, + ] + + prompt = tokenizer.apply_chat_template( + conversation=conversation, + tools=None, + chat_template=self.chat_template, + tokenize=False, + ) + prompts.append(prompt) + + scores = [] + for prompt in prompts: + inputs = tokenizer([prompt], return_tensors="pt") + inputs = self.wrap_device(inputs) + logits = self.model(**inputs).logits[:, -1, :] + scores.append(logits[0, self.token_true_id].item()) + return torch.Tensor(scores) + + +@pytest.mark.flaky(reruns=2) +@pytest.mark.parametrize("model_info", RERANK_MODELS) +def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None: + mteb_test_rerank_models(vllm_runner, model_info, hf_runner=ZerankHfRunner) From 51ef2397685b49ebd90b39bb4b4806a509df19b5 Mon Sep 17 00:00:00 2001 From: Alec <35311602+alec-flowers@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:11:19 -0700 Subject: [PATCH 144/624] [Rust Frontend] Carry reasoning controls over gRPC (#59836) Signed-off-by: Alec Flowers Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- rust/Cargo.lock | 2 +- rust/proto/Cargo.toml | 2 +- rust/proto/README.md | 2 +- rust/proto/inference.proto | 22 ++++++++ rust/src/server/src/grpc/convert.rs | 85 ++++++++++++++++++++++++++++- 5 files changed, 108 insertions(+), 5 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 331ebfa8dfaa..c7b3d66e4c1b 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5947,7 +5947,7 @@ dependencies = [ [[package]] name = "vllm-proto" -version = "0.4.0" +version = "0.5.0" dependencies = [ "prost", "prost-types", diff --git a/rust/proto/Cargo.toml b/rust/proto/Cargo.toml index 492cb5a06d67..c2e32d43a716 100644 --- a/rust/proto/Cargo.toml +++ b/rust/proto/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "vllm-proto" -version = "0.4.0" +version = "0.5.0" edition.workspace = true license.workspace = true description = "Protobuf messages and gRPC bindings for the vLLM Rust frontend" diff --git a/rust/proto/README.md b/rust/proto/README.md index d312e27acac6..3a60c05dfdf0 100644 --- a/rust/proto/README.md +++ b/rust/proto/README.md @@ -16,7 +16,7 @@ the schemas on pull requests, including on forks. ```toml [dependencies] -vllm-proto = "0.4" +vllm-proto = "0.5" ``` For example, import `vllm_proto::inference_client::InferenceClient` or diff --git a/rust/proto/inference.proto b/rust/proto/inference.proto index 79c1d0a147e2..2daacbfc35ba 100644 --- a/rust/proto/inference.proto +++ b/rust/proto/inference.proto @@ -58,6 +58,25 @@ message GenerateRequest { optional bool watermarking = 16; KVHintsEnvelope kv_hints = 17; + + // Transitional, to be deprecated: only consulted when the engine runs its own + // reasoning parser and the structured-output constraint covers only the + // output after reasoning. Prefer a constraint that covers the full output + // from the first generated token. + EngineReasoningGate engine_reasoning_gate = 18; +} + +// Inputs to the engine's reasoning gate, which keeps a structured-output +// constraint inactive until the engine's reasoning parser detects the end of +// reasoning. +message EngineReasoningGate { + // Omitted: the engine infers the initial state from the prompt. + // true: apply the constraint from the first generated token. + // false: generation starts inside reasoning. + optional bool reasoning_ended = 1; + // Chat template kwargs used to render the prompt, so the engine's reasoning + // parser observes the same thinking mode. + google.protobuf.Struct chat_template_kwargs = 2; } message KVHintsEnvelope { @@ -118,6 +137,9 @@ message StoppingCriteria { bool include_stop_strings = 5; bool ignore_eos = 6; + + // Omitted or -1 disables the budget; zero immediately closes reasoning. + optional int64 thinking_token_budget = 7; } message ResponseOptions { diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index a409cc7f79ee..d495fea485c7 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -9,6 +9,7 @@ use tonic::Status; use uuid::Uuid; use vllm_engine_core_client::protocol::kv_hints::{KvHintAction, KvHintsEnvelope}; use vllm_engine_core_client::protocol::output::StopReason; +use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams; use vllm_text::{ DecodedLogprobs, DecodedPromptLogprobs, FinishReason, Finished, Prompt, PromptTruncation, @@ -84,6 +85,7 @@ pub fn to_text_request( }; let session_id = req.session_id.filter(|s| !s.is_empty()); let kv_hints = req.kv_hints.map(kv_hints_from_proto); + let reasoning_gate = req.engine_reasoning_gate.unwrap_or_default(); let sampling = req.sampling.as_ref(); let decoding = req.decoding.as_ref(); @@ -137,8 +139,15 @@ pub fn to_text_request( data_parallel_rank: None, session_id, kv_hints, - reasoning_parser_kwargs: Default::default(), - reasoning_ended: None, + reasoning_parser_kwargs: ReasoningParserKwargs { + chat_template_kwargs: reasoning_gate + .chat_template_kwargs + .map(|kwargs| { + kwargs.fields.iter().map(|(k, v)| (k.clone(), proto_value_to_json(v))).collect() + }) + .unwrap_or_default(), + }, + reasoning_ended: reasoning_gate.reasoning_ended, lora_request: None, arrival_time: None, }) @@ -213,6 +222,7 @@ fn build_sampling_params( params.stop_token_ids = Some(s.stop_token_ids.clone()); } params.ignore_eos = s.ignore_eos; + params.thinking_token_budget = s.thinking_token_budget; } // ResponseOptions → logprobs @@ -611,6 +621,77 @@ mod tests { } } + #[test] + fn grpc_reasoning_controls_reach_engine_request_with_presence_intact() { + let chat_template_kwargs = + serde_json::json!({"enable_thinking": false, "reasoning_effort": "high"}); + for (gate, budget) in [ + (None, None), + ( + Some(pb::EngineReasoningGate { + reasoning_ended: Some(false), + chat_template_kwargs: json_to_proto_struct(&chat_template_kwargs), + }), + Some(0), + ), + ( + Some(pb::EngineReasoningGate { + reasoning_ended: Some(true), + chat_template_kwargs: None, + }), + Some(128), + ), + (Some(pb::EngineReasoningGate::default()), Some(-1)), + ] { + let expected = ( + serde_json::json!({ + "chat_template_kwargs": gate + .as_ref() + .and_then(|gate| gate.chat_template_kwargs.as_ref()) + .map_or_else(|| serde_json::json!({}), |_| chat_template_kwargs.clone()), + }), + gate.as_ref().and_then(|gate| gate.reasoning_ended), + budget.filter(|v| *v >= 0).map(|v| v as u64), + ); + let request = pb::GenerateRequest { + engine_reasoning_gate: gate, + stopping: Some(pb::StoppingCriteria { + thinking_token_budget: budget, + ..Default::default() + }), + ..base_request() + }; + let encoded = request.encode_to_vec(); + for stream in [false, true] { + let request = pb::GenerateRequest::decode(encoded.as_slice()).unwrap(); + let text = to_text_request(request, stream, &["test-model".to_string()]) + .expect("reasoning controls accepted"); + let engine = vllm_text::lower_text_request( + text, + vec![1], + SamplingHints::default(), + SamplingLimits { + max_model_len: 256, + max_logprobs: 20, + model_vocab_size: 512, + tokenizer_vocab_size: 512, + }, + &TestTokenizer::new(), + ) + .expect("reasoning controls lower to engine") + .generate_request; + assert_eq!( + ( + serde_json::to_value(engine.reasoning_parser_kwargs).unwrap(), + engine.reasoning_ended, + engine.sampling_params.thinking_token_budget, + ), + expected, + ); + } + } + } + #[test] fn watermarking_defaults_and_opt_out_survive_protobuf_conversion() { for watermarking in [None, Some(true), Some(false)] { From 112a2829157cc4267afae8c5577df20a36136c22 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Varikuntla Date: Thu, 8 Oct 2026 21:20:43 -0700 Subject: [PATCH 145/624] [Bugfix][Frontend] Keep streamed completion text_offset in step with tokens (#60668) Signed-off-by: Vineeth Sai Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../completion/test_completion_text_offset.py | 97 +++++++++++++++++++ vllm/entrypoints/openai/completion/serving.py | 11 ++- 2 files changed, 105 insertions(+), 3 deletions(-) create mode 100644 tests/entrypoints/openai/completion/test_completion_text_offset.py diff --git a/tests/entrypoints/openai/completion/test_completion_text_offset.py b/tests/entrypoints/openai/completion/test_completion_text_offset.py new file mode 100644 index 000000000000..c8ba7df62f8a --- /dev/null +++ b/tests/entrypoints/openai/completion/test_completion_text_offset.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json + +import pytest + +from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata +from vllm.entrypoints.openai.completion.protocol import CompletionRequest +from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion +from vllm.logprobs import Logprob +from vllm.outputs import CompletionOutput, RequestOutput + +TOKENS = [(1, "abc"), (2, " de"), (3, "!")] +# Text deltas as the detokenizer streams them with stop=["\n\n"]: one char is +# held back until the request finishes, so the text lags the tokens. +TEXT_DELTAS = ["ab", "c d", "e!"] + + +def _logprobs(i: int) -> dict[int, Logprob]: + token_id, token = TOKENS[i] + return {token_id: Logprob(logprob=-1.0, rank=1, decoded_token=token)} + + +def _output(text: str, idxs: list[int], finished: bool) -> RequestOutput: + return RequestOutput( + request_id="r", + prompt="Hi", + prompt_token_ids=[9], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text=text, + token_ids=[TOKENS[i][0] for i in idxs], + cumulative_logprob=None, + logprobs=[_logprobs(i) for i in idxs], + finish_reason="stop" if finished else None, + ) + ], + finished=finished, + ) + + +def _serving() -> OpenAIServingCompletion: + serving = OpenAIServingCompletion.__new__(OpenAIServingCompletion) + serving.enable_prompt_tokens_details = False + serving.enable_force_include_usage = False + serving.enable_per_request_metrics = False + serving.system_fingerprint = None + serving.return_tokens_as_token_ids = False + return serving + + +def _request(stream: bool) -> CompletionRequest: + return CompletionRequest( + model="m", prompt="Hi", max_tokens=8, logprobs=0, stop=["\n\n"], stream=stream + ) + + +@pytest.mark.asyncio +async def test_stream_text_offset_matches_non_stream_with_stop_holdback(): + async def results(): + for i, text in enumerate(TEXT_DELTAS): + yield 0, _output(text, [i], finished=i == len(TEXT_DELTAS) - 1) + + serving = _serving() + stream_offsets = [] + async for line in serving.completion_stream_generator( + _request(stream=True), + [{"prompt": "Hi"}], + results(), + "r", + 0, + "m", + num_prompts=1, + tokenizer=None, + request_metadata=RequestResponseMetadata(request_id="r"), + ): + data = line.removeprefix("data: ").strip() + if data == "[DONE]": + continue + for choice in json.loads(data)["choices"]: + stream_offsets += choice["logprobs"]["text_offset"] + + response = serving.request_output_to_completion_response( + [_output("".join(TEXT_DELTAS), [0, 1, 2], finished=True)], + _request(stream=False), + "r", + 0, + "m", + None, + RequestResponseMetadata(request_id="r"), + ) + assert response.choices[0].logprobs is not None + assert response.choices[0].logprobs.text_offset == [0, 3, 6] + assert stream_offsets == [0, 3, 6] diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index 81dd00f0b3bd..7f7bdac0757b 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -289,7 +289,7 @@ async def completion_stream_generator( request_metadata: RequestResponseMetadata, ) -> AsyncGenerator[str, None]: num_choices = 1 if request.n is None else request.n - previous_text_lens = [0] * num_choices * num_prompts + next_text_offsets = [0] * num_choices * num_prompts previous_num_tokens = [0] * num_choices * num_prompts has_echoed = [False] * num_choices * num_prompts num_prompt_tokens = [0] * num_prompts @@ -386,13 +386,18 @@ async def completion_stream_generator( num_output_top_logprobs=request.logprobs, tokenizer=tokenizer, logprob_token_ids=request.logprob_token_ids, - initial_text_offset=previous_text_lens[i], + initial_text_offset=next_text_offsets[i], return_as_token_id=request.return_tokens_as_token_ids, ) + # Advance by the tokens just reported, not by delta_text: + # text held back for stop strings lags the tokens. + if logprobs.tokens: + next_text_offsets[i] = logprobs.text_offset[-1] + len( + logprobs.tokens[-1] + ) else: logprobs = None - previous_text_lens[i] += len(delta_text) previous_num_tokens[i] += len(output.token_ids) stop_reason = output.stop_reason From 9824159d94082deda7e112bfdbadbd441eb8084a Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Fri, 9 Oct 2026 08:20:49 +0400 Subject: [PATCH 146/624] [Bugfix] Fix InternLM2 tool parser dropping characters from streamed arguments (#60709) Signed-off-by: Vadim Gimpelson Co-authored-by: Claude Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../test_internlm2_tool_parser.py | 28 +++++++++++++++++++ vllm/tool_parsers/internlm2_tool_parser.py | 17 +++++++---- 2 files changed, 40 insertions(+), 5 deletions(-) diff --git a/tests/tool_parsers/test_internlm2_tool_parser.py b/tests/tool_parsers/test_internlm2_tool_parser.py index 881933890dfb..28465654e4fa 100644 --- a/tests/tool_parsers/test_internlm2_tool_parser.py +++ b/tests/tool_parsers/test_internlm2_tool_parser.py @@ -10,6 +10,7 @@ ToolParserTestConfig, ToolParserTests, ) +from tests.tool_parsers.utils import run_tool_extraction_streaming from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.internlm2_tool_parser import Internlm2ToolParser @@ -162,3 +163,30 @@ def test_streaming_arguments_in_single_delta(default_tokenizer: TokenizerLike) - streamed += arguments assert json.loads(streamed) == {"city": "Dallas", "state": "TX"} + + +@pytest.mark.parametrize( + "arguments_deltas", + [ + ['{"location": "Beijing', '", "days": 3}}<|action_end|>'], + [*'{"location": "Beijing", "days": 3}}', "<|action_end|>"], + ], + ids=["split_inside_string", "char_by_char"], +) +def test_streaming_arguments_match_final_arguments( + default_tokenizer: TokenizerLike, arguments_deltas: list[str] +) -> None: + """Streamed argument deltas must add up to the final arguments, wherever + the chunk boundaries fall.""" + parser = Internlm2ToolParser(default_tokenizer) + deltas = [ + '<|action_start|><|plugin|>{"name": "get_weather", "parameters": ', + *arguments_deltas, + ] + + reconstructor = run_tool_extraction_streaming(parser, deltas) + + assert len(reconstructor.tool_calls) == 1 + streamed = reconstructor.tool_calls[0].function.arguments + streamed += parser.get_remaining_unstreamed_args() + assert streamed == '{"location": "Beijing", "days": 3}' diff --git a/vllm/tool_parsers/internlm2_tool_parser.py b/vllm/tool_parsers/internlm2_tool_parser.py index 5a99be6afd0e..709fd6c8f2a8 100644 --- a/vllm/tool_parsers/internlm2_tool_parser.py +++ b/vllm/tool_parsers/internlm2_tool_parser.py @@ -26,7 +26,7 @@ Tool, ToolParser, ) -from vllm.tool_parsers.utils import extract_intermediate_diff, is_complete_json +from vllm.tool_parsers.utils import find_common_prefix, is_complete_json logger = init_logger(__name__) @@ -172,10 +172,17 @@ def extract_tool_calls_streaming( elif cur_arguments and prev_arguments: cur_args_json = json.dumps(cur_arguments, ensure_ascii=False) prev_args_json = json.dumps(prev_arguments, ensure_ascii=False) - - argument_diff = extract_intermediate_diff( - cur_args_json, prev_args_json - ) + sent = len(self.streamed_args_for_tool[self.current_tool_id]) + + # Partial JSON is serialized with synthetic closers, so + # only the part shared with the previous parse is final. + if is_complete_json(parsable_arr): + argument_diff = cur_args_json[sent:] + elif cur_args_json != prev_args_json: + prefix = find_common_prefix(prev_args_json, cur_args_json) + argument_diff = prefix[sent:] + else: + argument_diff = "" delta = DeltaMessage( tool_calls=[ From e31b823bf59bdd96d09fdcf310c5fb02e934d3bf Mon Sep 17 00:00:00 2001 From: Kai Hoang <56827800+hpdkhoa@users.noreply.github.com> Date: Fri, 9 Oct 2026 11:20:54 +0700 Subject: [PATCH 147/624] [Bugfix][Tool Parser] Use LongCat tags in the LongCat structural tag (#60747) Signed-off-by: hpdkhoa Co-authored-by: Claude Opus 5.5 Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../test_structural_tag_registry.py | 39 ++++++++++++ vllm/tool_parsers/longcat_tool_parser.py | 2 + vllm/tool_parsers/structural_tag_registry.py | 60 +++++++++++++------ 3 files changed, 84 insertions(+), 17 deletions(-) diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py index 70f1921e459a..817c48a4401d 100644 --- a/tests/tool_parsers/test_structural_tag_registry.py +++ b/tests/tool_parsers/test_structural_tag_registry.py @@ -48,6 +48,7 @@ from vllm.tool_parsers.kimi_k2_tool_parser import KimiK2ToolParser from vllm.tool_parsers.kimi_k3_tool_parser import KimiK3ToolParser from vllm.tool_parsers.llama_tool_parser import Llama3JsonToolParser +from vllm.tool_parsers.longcat_tool_parser import LongcatFlashToolParser from vllm.tool_parsers.mimo_tool_parser import MiMoToolParser from vllm.tool_parsers.minimax_m2_tool_parser import MinimaxM2ToolParser from vllm.tool_parsers.plamo3_engine_tool_parser import Plamo3EngineToolParser @@ -850,6 +851,44 @@ class StepParser(DelegatingParser): ] +@pytest.mark.parametrize( + "tool_choice", + ["required", {"type": "function", "function": {"name": "get_weather"}}], +) +def test_longcat_forced_tool_choice_round_trips_native_tags( + sample_tools: list[ChatCompletionToolsParam], tool_choice +): + """Forced LongCat calls must use , not Hermes tags.""" + + class LongcatParser(DelegatingParser): + tool_parser_cls = LongcatFlashToolParser + + request = ChatCompletionRequest( + messages=[], + model="m", + tools=[tool.model_dump(exclude_none=True) for tool in sample_tools], + tool_choice=tool_choice, + ) + parser = LongcatParser(MagicMock(), tools=sample_tools) + request = parser.adjust_request(request) + + grammar = Grammar.from_structural_tag(request.structured_outputs.structural_tag) + # Rendered exactly as the LongCat-Flash chat template renders tool calls. + output = ( + '\n{"name": "get_weather", ' + '"arguments": {"city": "Dallas"}}\n' + ) + assert _is_grammar_accept_string(grammar, output) + assert not _is_grammar_accept_string(grammar, output.replace("longcat_", "")) + + _, content, tool_calls = parser.parse(output, request, enable_auto_tools=True) + + assert not content + assert [(call.name, json.loads(call.arguments)) for call in tool_calls] == [ + ("get_weather", {"city": "Dallas"}) + ] + + def test_tool_parsers_without_structural_tag_support_required_and_named(): class NonStructuralTagToolParser(ToolParser): pass diff --git a/vllm/tool_parsers/longcat_tool_parser.py b/vllm/tool_parsers/longcat_tool_parser.py index ccfa94167256..919f80143737 100644 --- a/vllm/tool_parsers/longcat_tool_parser.py +++ b/vllm/tool_parsers/longcat_tool_parser.py @@ -9,6 +9,8 @@ class LongcatFlashToolParser(Hermes2ProToolParser): + structural_tag_model = "longcat" + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py index 1d3a7c4a4c7f..f3148a827eb0 100644 --- a/vllm/tool_parsers/structural_tag_registry.py +++ b/vllm/tool_parsers/structural_tag_registry.py @@ -79,7 +79,7 @@ } ) VLLM_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset( - {"glm_4_7", "hermes", "hy_v4", "kimi_k3", "plamo3"} + {"glm_4_7", "hermes", "hy_v4", "kimi_k3", "longcat", "plamo3"} ) SUPPORTED_STRUCTURAL_TAG_MODELS = ( XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS | VLLM_BUILTIN_STRUCTURAL_TAG_MODELS @@ -346,15 +346,19 @@ def get_function_parameters(function) -> dict[str, Any] | bool: return function.parameters if function.parameters is not None else True -def _hermes_tool_tags(tools: list[FunctionToolParam]) -> list[TagFormat]: +def _hermes_tool_tags( + tools: list[FunctionToolParam], + start_tag: str, + end_tag: str, +) -> list[TagFormat]: arguments_field_prefix = '", "arguments": ' formats = [ # # {"name": "t1", "arguments": {"q": "v"}} # - ('\n{"name": "', "}\n"), + (start_tag + '\n{"name": "', "}\n" + end_tag), # {"name": "t1", "arguments": {"q": "v"}} - ('{"name": "', "}"), + (start_tag + '{"name": "', "}" + end_tag), ] return [ @@ -370,35 +374,29 @@ def _hermes_tool_tags(tools: list[FunctionToolParam]) -> list[TagFormat]: ] -@register_vllm_structural_tag("hermes") -def get_hermes_structural_tag( +def _hermes_style_structural_tag( tools: list[FunctionToolParam], - builtin_tools: list[BuiltinToolParam], tool_choice: SimplifiedToolChoice, - reasoning: bool, - token_suffix: str = "", + start_tag: str, + end_tag: str, ) -> StructuralTag: - del builtin_tools, reasoning, token_suffix - - tool_call_trigger = "" - + tags = _hermes_tool_tags(tools, start_tag, end_tag) if tool_choice == "auto": - tags = _hermes_tool_tags(tools) suffix_tag = ( - TriggeredTagsFormat(triggers=[tool_call_trigger], tags=tags) + TriggeredTagsFormat(triggers=[start_tag], tags=tags) if tags else AnyTextFormat() ) elif tool_choice == "forced": suffix_tag = TagsWithSeparatorFormat( - tags=_hermes_tool_tags(tools), + tags=tags, separator="", at_least_one=True, stop_after_first=True, ) else: suffix_tag = TagsWithSeparatorFormat( - tags=_hermes_tool_tags(tools), + tags=tags, separator="", at_least_one=True, ) @@ -406,6 +404,34 @@ def get_hermes_structural_tag( return StructuralTag(format=suffix_tag) +@register_vllm_structural_tag("hermes") +def get_hermes_structural_tag( + tools: list[FunctionToolParam], + builtin_tools: list[BuiltinToolParam], + tool_choice: SimplifiedToolChoice, + reasoning: bool, + token_suffix: str = "", +) -> StructuralTag: + del builtin_tools, reasoning, token_suffix + return _hermes_style_structural_tag( + tools, tool_choice, "", "" + ) + + +@register_vllm_structural_tag("longcat") +def get_longcat_structural_tag( + tools: list[FunctionToolParam], + builtin_tools: list[BuiltinToolParam], + tool_choice: SimplifiedToolChoice, + reasoning: bool, + token_suffix: str = "", +) -> StructuralTag: + del builtin_tools, reasoning, token_suffix + return _hermes_style_structural_tag( + tools, tool_choice, "", "" + ) + + def _minimax_tool_tags(tools: list[FunctionToolParam]) -> list[TagFormat]: return [ TagFormat( From 81d9c45e4b87b8abc82fd69f81be306f5ceae49a Mon Sep 17 00:00:00 2001 From: Rajeev Warrier Date: Fri, 9 Oct 2026 09:53:07 +0530 Subject: [PATCH 148/624] =?UTF-8?q?[Docs]=20[Multimodal]=20Note=20that=20p?= =?UTF-8?q?er-request=20media=5Fio=5Fkwargs=20needs=20--trust-request-mm?= =?UTF-8?q?=E2=80=A6=20(#60598)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Rajeev Warrier --- docs/features/multimodal_inputs.md | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/docs/features/multimodal_inputs.md b/docs/features/multimodal_inputs.md index 40ceb8158b52..cb091c1c79d6 100644 --- a/docs/features/multimodal_inputs.md +++ b/docs/features/multimodal_inputs.md @@ -1020,6 +1020,13 @@ vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ When you extract video frames on the client side and send them as `video/jpeg` (base64-concatenated JPEG frames), you can preserve the original video metadata by using `media_io_kwargs` in your request. This enables more accurate video understanding by preserving temporal information that would otherwise be lost during client-side frame extraction. +!!! warning + Per-request `media_io_kwargs` (and `mm_processor_kwargs`) are rejected by + default because they let a client change media loading and preprocessing + resource usage. Start the server with `--trust-request-mm-kwargs` only when + your API clients are trusted; see + [Per-request multimodal arguments](../usage/security.md#per-request-multimodal-arguments). + **Supported Parameters:** | Parameter | Type | Description | @@ -1042,7 +1049,7 @@ When you extract video frames on the client side and send them as `video/jpeg` ( frames_b64 = ",".join([encode_image(f) for f in frames]) video_url = f"data:video/jpeg;base64,{frames_b64}" - # Pass video metadata via media_io_kwargs + # Pass video metadata via media_io_kwargs (server started with --trust-request-mm-kwargs) response = client.chat.completions.create( model="your-multimodal-model", messages=[{ From fa312090cb1d972aa96f6b0ea174b7a32b4b0548 Mon Sep 17 00:00:00 2001 From: Aaron Hao Date: Thu, 8 Oct 2026 21:40:13 -0700 Subject: [PATCH 149/624] [Frontend] Accept each connection on the least loaded API worker (#60723) Signed-off-by: hao-aaron Co-authored-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- .../launchers/test_balanced_accepts.py | 218 ++++++++++++++++++ .../entrypoints/launchers/api_server/entry.py | 11 +- vllm/entrypoints/launchers/launcher.py | 130 ++++++++++- vllm/v1/utils.py | 10 + 4 files changed, 365 insertions(+), 4 deletions(-) create mode 100644 tests/entrypoints/launchers/test_balanced_accepts.py diff --git a/tests/entrypoints/launchers/test_balanced_accepts.py b/tests/entrypoints/launchers/test_balanced_accepts.py new file mode 100644 index 000000000000..3ac1025ca853 --- /dev/null +++ b/tests/entrypoints/launchers/test_balanced_accepts.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NoSignalServer(peer_loads=...): API workers sharing one listening socket +accept each connection on the worker with the fewest open connections.""" + +import asyncio +import datetime +import multiprocessing +import os +import socket +import ssl +from collections import Counter + +import pytest +import uvicorn +import uvloop + +from vllm.entrypoints.launchers.launcher import NoSignalServer + +TIMEOUT_S = 30 + + +def _bound_socket() -> socket.socket: + # Like setup_server: bound, not yet listening. + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(("127.0.0.1", 0)) + return sock + + +def _app(hold_s: float = 0.0): + pid = str(os.getpid()).encode() + + async def app(scope, receive, send): + if scope["type"] == "lifespan": + while (await receive())["type"] != "lifespan.shutdown": + await send({"type": "lifespan.startup.complete"}) + await send({"type": "lifespan.shutdown.complete"}) + return + await asyncio.sleep(hold_s) + headers = [(b"content-length", str(len(pid)).encode())] + await send({"type": "http.response.start", "status": 200, "headers": headers}) + await send({"type": "http.response.body", "body": pid}) + + return app + + +async def _get(port: int, requests: int = 1, ssl_ctx=None) -> list[bytes]: + """``requests`` requests on one keep-alive connection; returns the bodies.""" + reader, writer = await asyncio.open_connection("127.0.0.1", port, ssl=ssl_ctx) + bodies = [] + for _ in range(requests): + writer.write(b"GET / HTTP/1.1\r\nHost: t\r\n\r\n") + head = await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), TIMEOUT_S) + assert head.startswith(b"HTTP/1.1 200") + length = int(head.split(b"content-length: ")[1].split(b"\r\n")[0]) + bodies.append(await asyncio.wait_for(reader.readexactly(length), TIMEOUT_S)) + writer.close() + return bodies + + +async def _serve(sock, requests, loads=None, **config_kwargs): + config = uvicorn.Config(_app(), lifespan="on", log_level="warning", **config_kwargs) + loads = loads if loads is not None else [0] + server = NoSignalServer(config, peer_loads=(loads, 0)) + task = asyncio.create_task(server.serve(sockets=[sock])) + while not server.started: + assert not task.done(), task.result() + await asyncio.sleep(0.01) + try: + return await requests(sock.getsockname()[1], server) + finally: + server.should_exit = True + await asyncio.wait_for(task, TIMEOUT_S) + # The listener is closed, nothing is left accepting, and peers no + # longer see this worker as a target. + assert sock.fileno() == -1 + assert not server._accept_tasks + assert loads[0] == 2**31 - 1 + + +def _run(coro, loop_impl): + return uvloop.run(coro) if loop_impl == "uvloop" else asyncio.run(coro) + + +@pytest.mark.parametrize("loop_impl", ["asyncio", "uvloop"]) +def test_requests_and_keep_alive_served(loop_impl): + async def requests(port, server): + return await _get(port), await _get(port, requests=3) + + one, three = _run(_serve(_bound_socket(), requests), loop_impl) + assert one == [str(os.getpid()).encode()] + assert len(three) == 3 + + +def _self_signed(tmp_path): + x509 = pytest.importorskip("cryptography.x509") + from cryptography.hazmat.primitives import hashes, serialization + from cryptography.hazmat.primitives.asymmetric import ec + from cryptography.x509.oid import NameOID + + key = ec.generate_private_key(ec.SECP256R1()) + name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "127.0.0.1")]) + now = datetime.datetime.now(datetime.UTC) + cert = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - datetime.timedelta(minutes=1)) + .not_valid_after(now + datetime.timedelta(hours=1)) + .sign(key, hashes.SHA256()) + ) + keyfile, certfile = tmp_path / "key.pem", tmp_path / "cert.pem" + keyfile.write_bytes( + key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + ) + certfile.write_bytes(cert.public_bytes(serialization.Encoding.PEM)) + return str(keyfile), str(certfile) + + +def test_tls_handshake_failure_keeps_serving(tmp_path): + keyfile, certfile = _self_signed(tmp_path) + client_ctx = ssl.create_default_context() + client_ctx.check_hostname = False + client_ctx.verify_mode = ssl.CERT_NONE + + async def requests(port, server): + # A plain-text client fails the handshake; the next TLS client is served. + reader, writer = await asyncio.open_connection("127.0.0.1", port) + writer.write(b"GET / HTTP/1.1\r\nHost: t\r\n\r\n") + await asyncio.wait_for(reader.read(), TIMEOUT_S) + writer.close() + return await _get(port, ssl_ctx=client_ctx) + + bodies = uvloop.run( + _serve(_bound_socket(), requests, ssl_keyfile=keyfile, ssl_certfile=certfile) + ) + assert bodies == [str(os.getpid()).encode()] + + +def test_busier_worker_still_serves(): + # This worker is index 0; peer 1 claims no connections but never accepts. + loads = [0, 0] + + async def requests(port, server): + reader, writer = await asyncio.open_connection("127.0.0.1", port) + bodies = await _get(port) + assert loads[0] >= 1 # the held connection was published + writer.close() + return bodies + + assert uvloop.run(_serve(_bound_socket(), requests, loads)) == [ + str(os.getpid()).encode() + ] + + +def _worker(sock, loads, index, ready, stop): + async def main(): + config = uvicorn.Config(_app(hold_s=1), lifespan="on", log_level="warning") + server = NoSignalServer(config, peer_loads=(loads, index)) + task = asyncio.create_task(server.serve(sockets=[sock])) + while not server.started: + await asyncio.sleep(0.01) + ready.release() + await asyncio.get_running_loop().run_in_executor(None, stop.wait) + server.should_exit = True + await task + + # API workers run on uvloop, whose own server accepts a whole burst at once. + uvloop.run(main()) + + +@pytest.mark.skipif((os.cpu_count() or 1) < 8, reason="needs idle cores per worker") +@pytest.mark.parametrize("arrival_s", [0, 1]) +def test_connections_balanced_over_workers(arrival_s): + """Both a burst and a steady stream of connections end up balanced; with + the event loop's own server one worker takes most of either.""" + workers, connections = 4, 400 + sock = _bound_socket() + ctx = multiprocessing.get_context("spawn") + ready, stop = ctx.Semaphore(0), ctx.Event() + loads = ctx.RawArray("i", workers) + procs = [ + ctx.Process(target=_worker, args=(sock, loads, i, ready, stop)) + for i in range(workers) + ] + for p in procs: + p.start() + try: + for _ in range(workers): + assert ready.acquire(timeout=120) + + async def get_after(delay): + await asyncio.sleep(delay) + return await _get(sock.getsockname()[1]) + + async def arrive(): + delays = [arrival_s * i / connections for i in range(connections)] + results = await asyncio.gather(*(get_after(d) for d in delays)) + return Counter(body for (body,) in results) + + counts = uvloop.run(arrive()) + assert sum(counts.values()) == connections + # Fair share is 100. + assert max(counts.values()) <= 5 * connections // (4 * workers), counts + finally: + stop.set() + for p in procs: + p.join(30) + if p.is_alive(): + p.kill() + sock.close() diff --git a/vllm/entrypoints/launchers/api_server/entry.py b/vllm/entrypoints/launchers/api_server/entry.py index 28bb696f2c4e..e5ba536e1128 100644 --- a/vllm/entrypoints/launchers/api_server/entry.py +++ b/vllm/entrypoints/launchers/api_server/entry.py @@ -8,7 +8,7 @@ import socket import tempfile from argparse import Namespace -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, MutableSequence from contextlib import asynccontextmanager from typing import Any @@ -118,6 +118,7 @@ async def build_and_serve( listen_address: str, sock: socket.socket, args: Namespace, + peer_loads: tuple[MutableSequence[int], int] | None = None, **uvicorn_kwargs, ) -> asyncio.Task: """Build FastAPI app, initialize state, and start serving. @@ -142,6 +143,7 @@ async def build_and_serve( app, sock=sock, enable_ssl_refresh=args.enable_ssl_refresh, + peer_loads=peer_loads, host=args.host, port=args.port, log_level=args.uvicorn_log_level, @@ -185,12 +187,17 @@ async def run_server_worker( if args.reasoning_parser_plugin and len(args.reasoning_parser_plugin) > 3: ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin) + peer_loads = None + if client_config and "api_server_loads" in client_config: + loads = client_config.pop("api_server_loads") + peer_loads = (loads, client_config["client_index"]) + async with build_async_engine_client( args, client_config=client_config, ) as engine_client: shutdown_task = await build_and_serve( - engine_client, listen_address, sock, args, **uvicorn_kwargs + engine_client, listen_address, sock, args, peer_loads, **uvicorn_kwargs ) # NB: Await server shutdown only after the backend context is exited try: diff --git a/vllm/entrypoints/launchers/launcher.py b/vllm/entrypoints/launchers/launcher.py index 9fabfa3277af..44161439b690 100644 --- a/vllm/entrypoints/launchers/launcher.py +++ b/vllm/entrypoints/launchers/launcher.py @@ -5,7 +5,7 @@ import contextlib import signal import socket -from collections.abc import Generator +from collections.abc import Generator, MutableSequence from functools import partial from typing import Any @@ -34,23 +34,149 @@ logger = init_logger(__name__) +# A worker with more connections than another leaves a pending connection to +# the less loaded one, rechecking every 1 ms, for up to 5 ms (it may be busy). +_ACCEPT_DEFER_CHECKS = 5 +# The load of a worker that is not accepting (also the initial value of the +# shared array), so it is never the least loaded. +_NOT_ACCEPTING = 2**31 - 1 + class NoSignalServer(uvicorn.Server): """Uvicorn server that never installs its own SIGINT/SIGTERM handlers. Callers register their own handlers on the event loop for graceful shutdown; uvicorn's would race with and override them (see #49668). + + With ``peer_loads`` (several API workers on one shared socket: a shared + array with one slot per worker, and this worker's index), each worker + publishes its number of open connections, and a pending connection is + accepted by the least loaded worker. The event loop's own server would + accept every queued connection in whichever worker wakes first. """ + def __init__( + self, + config: uvicorn.Config, + peer_loads: tuple[MutableSequence[int], int] | None = None, + ): + super().__init__(config) + self.peer_loads = peer_loads + self._accept_tasks: set[asyncio.Task] = set() + self._connecting = 0 + @contextlib.contextmanager def capture_signals(self) -> Generator[None, None, None]: yield + def _publish_load(self) -> None: + assert self.peer_loads is not None + loads, index = self.peer_loads + loads[index] = len(self.server_state.connections) + self._connecting + + def _least_loaded(self) -> bool: + assert self.peer_loads is not None + loads, index = self.peer_loads + return loads[index] <= min(loads) + + async def startup(self, sockets: list[socket.socket] | None = None) -> None: + if not (self.peer_loads and sockets): + return await super().startup(sockets) + # With no sockets uvicorn starts no asyncio server; _accept serves + # them with uvicorn's protocol instead. + await super().startup(sockets=[]) + if not self.started: # lifespan startup failed (uvicorn < 0.50) + return + self._publish_load() + for sock in sockets: + sock.listen(self.config.backlog) + sock.setblocking(False) + self._track(self._accept(sock)) + + def _track(self, coro) -> None: + task = asyncio.create_task(coro) + self._accept_tasks.add(task) + task.add_done_callback(self._accept_tasks.discard) + + async def _accept(self, sock: socket.socket) -> None: + loop = asyncio.get_running_loop() + while True: + await _readable(loop, sock) + # Connections closed since the last accept; peers refresh likewise. + self._publish_load() + for _ in range(_ACCEPT_DEFER_CHECKS): + if self._least_loaded(): + break + await asyncio.sleep(0.001) + try: + conn, _ = sock.accept() + except (BlockingIOError, InterruptedError, ConnectionAbortedError): + continue # another worker took it + except OSError: + # e.g. EMFILE; asyncio's own server also retries after 1 s. + logger.exception("Error accepting a connection") + await asyncio.sleep(1) + continue + self._connecting += 1 + self._publish_load() + # Its own task, so a TLS handshake does not hold up the next accept. + self._track(self._connect(loop, conn)) + + async def _connect(self, loop: asyncio.AbstractEventLoop, conn: socket.socket): + config = self.config + try: + await loop.connect_accepted_socket( + lambda: config.http_protocol_class( # type: ignore[call-arg] + config=config, + server_state=self.server_state, + app_state=self.lifespan.state, + _loop=loop, + ), + conn, + ssl=config.ssl, + ) + except asyncio.CancelledError: + conn.close() + raise + except Exception: + # e.g. a failed TLS handshake + logger.debug("Accepted connection failed to start", exc_info=True) + conn.close() + finally: + self._connecting -= 1 + self._publish_load() + + async def shutdown(self, sockets: list[socket.socket] | None = None) -> None: + # Before super() closes the listening sockets. + tasks = list(self._accept_tasks) + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + if self.peer_loads is not None: + loads, index = self.peer_loads + loads[index] = _NOT_ACCEPTING + await super().shutdown(sockets) + + +async def _readable(loop: asyncio.AbstractEventLoop, sock: socket.socket) -> None: + ready: asyncio.Future[None] = loop.create_future() + + def wake() -> None: + if not ready.done(): + ready.set_result(None) + + loop.add_reader(sock, wake) + try: + await ready + finally: + loop.remove_reader(sock) + async def serve_http( app: FastAPI, sock: socket.socket | None, enable_ssl_refresh: bool = False, + peer_loads: tuple[MutableSequence[int], int] | None = None, **uvicorn_kwargs: Any, ): """Start a FastAPI app using Uvicorn, with support for custom Uvicorn config @@ -96,7 +222,7 @@ async def serve_http( config.h11_max_incomplete_event_size = h11_max_incomplete_event_size config.h11_max_header_count = h11_max_header_count config.load() - server = NoSignalServer(config) + server = NoSignalServer(config, peer_loads=peer_loads) app.state.server = server loop = asyncio.get_running_loop() diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index 868864f0ca41..b8d98ff1cf74 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -213,6 +213,14 @@ def __init__( SharedAdmissionStats.num_counters(num_servers), ) + # Open connections per API server, to accept on the least loaded one; + # a server that has not started yet is never the least loaded. + api_server_loads = ( + spawn_context.RawArray("i", [2**31 - 1] * num_servers) + if num_servers > 1 + else None + ) + listeners = [*input_listeners, *output_listeners] self._listeners = listeners try: @@ -228,6 +236,8 @@ def __init__( "client_count": num_servers, "client_index": i, } + if api_server_loads is not None: + client_config["api_server_loads"] = api_server_loads if admission_counters is not None: client_config["mp_admission_counters"] = admission_counters if stats_update_address is not None: From 88a419649e36f3412d2650c7162135a3b12bc5a1 Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Thu, 8 Oct 2026 23:45:13 -0500 Subject: [PATCH 150/624] [ROCm][Kimi-K3] Use vLLM's FP8 max in the fused AttnRes quantization (#60717) Signed-off-by: Rohan138 Co-authored-by: Claude Opus 5.5 --- tests/models/kimi_k3/test_amd_attn_res.py | 12 ++++++++---- vllm/models/kimi_k3/amd/ops/attn_res.py | 3 ++- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/tests/models/kimi_k3/test_amd_attn_res.py b/tests/models/kimi_k3/test_amd_attn_res.py index f4efdbd5b21d..871d53faae12 100644 --- a/tests/models/kimi_k3/test_amd_attn_res.py +++ b/tests/models/kimi_k3/test_amd_attn_res.py @@ -247,12 +247,16 @@ def test_amd_attn_res_fp8_preserves_prefix_and_quantized_output( ) ref_output, ref_scale = quant(reference) torch.testing.assert_close(scale, ref_scale, atol=1e-7, rtol=1e-6) + # Floating-point fusion can move values at FP8 rounding midpoints. # Bound each change to the adjacent code and bound how often it occurs. - code_delta = ( - output.view(torch.uint8).to(torch.int16) - - ref_output.view(torch.uint8).to(torch.int16) - ).abs() + # Codes are sign-magnitude, so order them by value to keep a step + # across zero adjacent. + def ordinal(x: torch.Tensor) -> torch.Tensor: + code = x.view(torch.uint8).to(torch.int16) + return torch.where(code >= 0x80, 0x80 - code, code) + + code_delta = (ordinal(output) - ordinal(ref_output)).abs() assert code_delta.max().item() <= 1 assert (code_delta != 0).float().mean().item() < 1e-4 torch.testing.assert_close(prefix, ref_prefix, atol=0, rtol=0) diff --git a/vllm/models/kimi_k3/amd/ops/attn_res.py b/vllm/models/kimi_k3/amd/ops/attn_res.py index 5c4fe7a0b0b2..95f748e72d1f 100644 --- a/vllm/models/kimi_k3/amd/ops/attn_res.py +++ b/vllm/models/kimi_k3/amd/ops/attn_res.py @@ -10,6 +10,7 @@ import torch +from vllm.model_executor.layers.quantization.utils.quant_utils import get_fp8_min_max from vllm.triton_utils import tl, triton @@ -207,7 +208,7 @@ def attn_res( HAS_DELTA=delta is not None, WRITE_BLOCK=block_write_idx >= 0, APPLY_OUTPUT_NORM=output_norm_weight is not None, - QUANT_MAX=0.0 if quant_dtype is None else torch.finfo(quant_dtype).max, + QUANT_MAX=0.0 if quant_dtype is None else get_fp8_min_max()[1], BLOCK_L=block_l, BLOCK_D=triton.next_power_of_2(hidden_size), num_warps=num_warps, From 8cbd5d03006c33185f402249ff2b448efd594986 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Thu, 8 Oct 2026 22:48:16 -0600 Subject: [PATCH 151/624] [Frontend] structured decisions: noul and score question types (#60696) Signed-off-by: Matt Mastracci Co-authored-by: DeepSeek Co-authored-by: Claude Opus 5.5 Co-authored-by: Misha Goin --- .../online_serving/structured_decisions.md | 44 ++++++-- .../test_question_types.py | 74 +++++++++++++ .../structured_decisions/test_strategies.py | 3 + .../test_structured_decisions_e2e.py | 44 ++++++++ .../structured_decisions/question_types.py | 100 ++++++++++++++++-- .../structured_decisions/strategies.py | 32 ++++-- 6 files changed, 270 insertions(+), 27 deletions(-) diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md index 7de935a62f2b..2194b54e47b4 100644 --- a/docs/serving/online_serving/structured_decisions.md +++ b/docs/serving/online_serving/structured_decisions.md @@ -11,8 +11,8 @@ log probabilities. ## How it works 1. Each question is one read. The user message is the state, followed by the - question and its options, labeled `A`, `B`, `C`, ... in order, and - "Answer with the letter of one option only." + question and its options with the type's labels, and an instruction to + answer with one label. 2. The read is the reply's first token. Each label must be one token there, or the request returns 400. 3. The read returns the logprobs of that question's label tokens. Softmax over @@ -31,8 +31,11 @@ question prefills only its own text, then reads one token. | type | criteria | answer | | --- | --- | --- | | `choice` | map of option name to a description or `null` | `choice`, `probabilities` by option name, `confidence` | +| `noul` | `null`, or an object with optional `true` and `false` descriptions | `noul` (probability of yes), `probabilities`, `confidence` | +| `score` | ordered list of level names | `score` (expected 0-indexed level), `legend`, `probabilities` by level index, `confidence` | -A choice has at least one option. A question id is a non-empty string. +A choice has at least one option, a score at least two levels. A question id +is a non-empty string. ## Example @@ -51,6 +54,16 @@ curl -s http://localhost:8000/v1/systemone \ "shipping": "deliveries", "security": "account access" } + }, + "urgent": { + "type": "noul", + "instructions": "Does the customer need a reply within the hour?", + "criteria": {"true": "the account is locked or money is moving"} + }, + "tone": { + "type": "score", + "instructions": "How upset is the customer?", + "criteria": ["calm", "annoyed", "furious"] } }, "chat_template_kwargs": {"enable_thinking": false} @@ -71,11 +84,26 @@ Response, with probabilities that depend on the model: "choice": "billing", "probabilities": {"billing": 0.97, "shipping": 0.01, "security": 0.02}, "confidence": 0.90 + }, + "urgent": { + "type": "noul", + "noul": 0.02, + "probabilities": {"yes": 0.02, "no": 0.98}, + "confidence": 0.95 + }, + "tone": { + "type": "score", + "score": 0.6, + "legend": {"0": "calm", "1": "annoyed", "2": "furious"}, + "probabilities": {"0": 0.7, "1": 0.2, "2": 0.1}, + "confidence": 0.65 } }, - "usage": {"input_tokens": 142, "output_tokens": 1}, + "usage": {"input_tokens": 142, "output_tokens": 3}, "diagnostics": { - "team": {"label_mass": 0.93, "argmax_is_label": true} + "team": {"label_mass": 0.93, "argmax_is_label": true}, + "urgent": {"label_mass": 0.98, "argmax_is_label": true}, + "tone": {"label_mass": 0.91, "argmax_is_label": true} } } ``` @@ -87,8 +115,9 @@ whole vocabulary: its share of the labels times `label_mass`. ## Labels -Options are labeled `A` to `Z` in the order the request lists them, so a -`choice` has at most 26 options. +Labels follow the type: a `choice` labels its options `A` to `Z` in the order +the request lists them, a `noul` labels its options `yes` and `no`, and a +`score` labels its levels `0` to `9`, so a level's label is its score. ## Limits @@ -96,6 +125,7 @@ Options are labeled `A` to `Z` in the order the request lists them, so a | --- | --- | | questions per request | 64 | | options per `choice` | 26 | +| levels per `score` | 10 | ## Request fields diff --git a/tests/entrypoints/generate/structured_decisions/test_question_types.py b/tests/entrypoints/generate/structured_decisions/test_question_types.py index e5906accf025..3b9ad7ba60a9 100644 --- a/tests/entrypoints/generate/structured_decisions/test_question_types.py +++ b/tests/entrypoints/generate/structured_decisions/test_question_types.py @@ -25,6 +25,20 @@ def choice(qid="bucket", criteria=None, instructions="Which team?", max_options= ) +def noul(qid="urgent", criteria=None, instructions="Reply within the hour?"): + return build_question(qid, "noul", instructions, criteria, 128) + + +def score(qid="tone", criteria=None, instructions="How angry?", max_options=128): + return build_question( + qid, + "score", + instructions, + criteria or ["calm", "annoyed", "furious"], + max_options, + ) + + def test_choice_labels_follow_option_order(): q = choice() assert q.labels == ("A", "B", "C") @@ -62,6 +76,10 @@ def test_choice_answer_shape(): ("", "choice", {"x": None, "y": None}, "non-empty"), ("q", "nope", {"x": None, "y": None}, "unknown question type"), ("q", "choice", ["x", "y"], "must map option names"), + ("q", "noul", ["yes", "no"], "must be an object"), + ("q", "noul", {"maybe": "same day"}, "must be an object"), + ("q", "score", {"0": "calm", "1": "furious"}, "ordered list of levels"), + ("q", "score", ["calm"], "ordered list of levels"), ], ) def test_build_question_rejects(qid, type_name, criteria, match): @@ -82,6 +100,62 @@ def test_option_limit(): choice(criteria={str(i): None for i in range(27)}) +def test_noul_labels_and_prompt(): + q = noul() + assert q.labels == ("yes", "no") + assert [a.name for a in q.options] == ["yes", "no"] + assert q.type.prompt(q) == ( + "Question: Reply within the hour?\nAnswer with yes or no only." + ) + described = noul(criteria={"true": "same day", "false": "later"}) + assert described.type.prompt(described) == ( + "Question: Reply within the hour?\n" + "yes: same day\n" + "no: later\n" + "Answer with yes or no only." + ) + + +def test_noul_answer_shape(): + q = noul() + answer = q.type.answer(q, [0.8, 0.2], 0.5) + assert answer == { + "type": "noul", + "noul": 0.8, + "probabilities": {"yes": 0.8, "no": 0.2}, + "confidence": 0.4, + } + + +def test_score_labels_are_levels(): + assert score().labels == ("0", "1", "2") + with pytest.raises(StructuredDecisionError, match="at most 10"): + score(criteria=[str(i) for i in range(11)]) + + +def test_score_prompt(): + q = score() + assert q.type.prompt(q) == ( + "Question: How angry?\n" + "0: calm\n" + "1: annoyed\n" + "2: furious\n" + "Answer with the number of one level only." + ) + + +def test_score_answer_shape(): + q = score() + answer = q.type.answer(q, [0.5, 0.25, 0.25], 0.4) + assert answer == { + "type": "score", + "score": 0.75, + "legend": {"0": "calm", "1": "annoyed", "2": "furious"}, + "probabilities": {"0": 0.5, "1": 0.25, "2": 0.25}, + "confidence": 0.2, + } + + def test_registered_type_plugs_in(): class BinaryQuestion(QuestionType): name = "test_binary" diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index 51bff19c8ca2..3d9538e4dc8b 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -87,6 +87,9 @@ def test_labels_start_the_reply(qwen): tail, ids = reply_label_ids(tokenizer, prompt_ids) assert tokenizer.decode(tail) == "\n\n" assert [tokenizer.decode([i]) for i in ids] == list(LABELS) + # A noul's and a score's labels are one token here too. + _, ids = reply_label_ids(tokenizer, prompt_ids, ("yes", "no", *"0123456789")) + assert len(set(ids)) == 12 # After a colon, Qwen writes ":A" as one token, so "A" is not one token. with pytest.raises(ValueError, match="not one distinct token"): reply_label_ids(tokenizer, tokenizer.encode("team:")) diff --git a/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py b/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py index 3e0f9bc0fb66..a206c56ccc06 100644 --- a/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py +++ b/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py @@ -188,6 +188,48 @@ def test_thinking_is_refused(server): assert "thinking must be off" in response.json()["error"]["message"] +def test_noul_and_score(server): + response = post( + server, + { + "model": MODEL_NAME, + "state": {"ticket": "My card was charged twice for one order."}, + "questions": { + "urgent": {"type": "noul", "instructions": "Is money moving?"}, + "tone": { + "type": "score", + "instructions": "How upset is the customer?", + "criteria": ["calm", "annoyed", "furious"], + }, + }, + "chat_template_kwargs": {"enable_thinking": False}, + }, + ) + assert response.status_code == 200, response.text + body = response.json() + urgent = body["answers"]["urgent"] + assert urgent["type"] == "noul" + assert set(urgent["probabilities"]) == {"yes", "no"} + assert urgent["probabilities"]["yes"] + urgent["probabilities"]["no"] == ( + pytest.approx(1.0, abs=1e-4) + ) + assert urgent["noul"] == urgent["probabilities"]["yes"] + assert urgent["confidence"] == pytest.approx( + max(urgent["probabilities"].values()) + * body["diagnostics"]["urgent"]["label_mass"] + ) + tone = body["answers"]["tone"] + assert tone["type"] == "score" + assert tone["legend"] == {"0": "calm", "1": "annoyed", "2": "furious"} + assert set(tone["probabilities"]) == {"0", "1", "2"} + assert sum(tone["probabilities"].values()) == pytest.approx(1.0, abs=1e-4) + assert tone["score"] == pytest.approx( + sum(i * p for i, p in enumerate(tone["probabilities"].values())) + ) + assert 0 <= tone["score"] <= 2 + assert body["usage"]["output_tokens"] == 2 + + @pytest.mark.parametrize( "questions,match", [ @@ -208,6 +250,8 @@ def test_thinking_is_refused(server): }, "at most 64", ), + ({"q": {"type": "noul", "criteria": ["yes", "no"]}}, "must be an object"), + ({"q": {"type": "score", "criteria": ["calm"]}}, "ordered list of levels"), ], ) def test_validation_errors(server, questions, match): diff --git a/vllm/entrypoints/generate/structured_decisions/question_types.py b/vllm/entrypoints/generate/structured_decisions/question_types.py index a9a6514a0373..0e80694b14eb 100644 --- a/vllm/entrypoints/generate/structured_decisions/question_types.py +++ b/vllm/entrypoints/generate/structured_decisions/question_types.py @@ -18,8 +18,8 @@ class StructuredDecisionError(ValueError): """An invalid request. Returned as a 400.""" -#: Options are labeled A to Z in order. A single capital letter is one token -#: at the start of a reply for the tokenizers tested; every read checks it. +#: A choice labels its options A to Z in order. A single capital letter is one +#: token at the start of a reply for the tokenizers tested. Every read checks it. LABELS = tuple(string.ascii_uppercase) @@ -40,6 +40,9 @@ class Question: class QuestionType(ABC): name: ClassVar[str] + #: The labels of a question's options, in option order. + label_set: ClassVar[tuple[str, ...]] = LABELS + reply_instruction: ClassVar[str] = "Answer with the letter of one option only." @abstractmethod def parse_options(self, qid: str, criteria: Any) -> list[Option]: ... @@ -52,16 +55,22 @@ def answer( ``question.labels[i]`` among the labels, and the list sums to 1. ``label_mass`` is the labels' total probability over the vocabulary.""" + def reply_line(self, label: str, option: Option) -> str | None: + """The option's line in the prompt, or None to leave it out.""" + return ( + f"{label}: {option.name} - {option.description}" + if option.description + else f"{label}: {option.name}" + ) + def prompt(self, question: Question) -> str: """The question as the model reads it, after the state.""" lines = [f"Question: {question.instructions}"] if question.instructions else [] for label, o in zip(question.labels, question.options): - lines.append( - f"{label}: {o.name} - {o.description}" - if o.description - else f"{label}: {o.name}" - ) - lines.append("Answer with the letter of one option only.") + line = self.reply_line(label, o) + if line is not None: + lines.append(line) + lines.append(self.reply_instruction) return "\n".join(lines) @@ -100,7 +109,7 @@ def build_question( names = [o.name for o in options] if len(set(names)) != len(names): raise StructuredDecisionError(f"question {qid!r}: duplicate option names") - limit = min(max_options, len(LABELS)) + limit = min(max_options, len(qtype.label_set)) if len(options) > limit: raise StructuredDecisionError( f"question {qid!r}: at most {limit} options for this model" @@ -112,7 +121,7 @@ def build_question( type=qtype, instructions=instructions, options=tuple(options), - labels=LABELS[: len(options)], + labels=qtype.label_set[: len(options)], ) @@ -158,3 +167,74 @@ def answer( # read that mostly wanted a non-label token reports low confidence. "confidence": probs[top] * label_mass, } + + +@register_question_type +class NoulQuestion(QuestionType): + """Yes or no. ``criteria`` may describe what true and false mean.""" + + name = "noul" + label_set = ("yes", "no") + reply_instruction = "Answer with yes or no only." + + def parse_options(self, qid: str, criteria: Any) -> list[Option]: + if criteria is not None and ( + not isinstance(criteria, dict) or not criteria.keys() <= {"true", "false"} + ): + raise StructuredDecisionError( + f"question {qid!r}: noul criteria must be an object with true and false" + ) + criteria = criteria or {} + true = criteria.get("true") + false = criteria.get("false") + return [ + Option("yes", None if true is None else str(true)), + Option("no", None if false is None else str(false)), + ] + + def reply_line(self, label: str, option: Option) -> str | None: + return f"{label}: {option.description}" if option.description else None + + def answer( + self, question: Question, probs: list[float], label_mass: float + ) -> dict[str, Any]: + return { + "type": self.name, + "noul": probs[0], + "probabilities": {"yes": probs[0], "no": probs[1]}, + # The chosen label's probability over the whole vocabulary, like a + # choice's confidence. + "confidence": max(probs) * label_mass, + } + + +@register_question_type +class ScoreQuestion(QuestionType): + """Rate the state on an ordered scale. ``criteria`` is the ordered list of + level names, from the first to the last level of the scale.""" + + name = "score" + # 0-indexed like the answer, so a level's label is its score. + label_set = tuple(string.digits) + reply_instruction = "Answer with the number of one level only." + + def parse_options(self, qid: str, criteria: Any) -> list[Option]: + if not isinstance(criteria, list) or len(criteria) < 2: + raise StructuredDecisionError( + f"question {qid!r}: score criteria must be an ordered list of levels" + ) + return [Option(str(level)) for level in criteria] + + def answer( + self, question: Question, probs: list[float], label_mass: float + ) -> dict[str, Any]: + top = argmax(probs) + return { + "type": self.name, + # The expected level, 0-indexed, so a level's index in the request's + # criteria is its value. + "score": sum(i * p for i, p in enumerate(probs)), + "legend": {str(i): o.name for i, o in enumerate(question.options)}, + "probabilities": {str(i): p for i, p in enumerate(probs)}, + "confidence": probs[top] * label_mass, + } diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index 864443cbb0b9..af588e13a289 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -23,7 +23,13 @@ from vllm.tokenizers import TokenizerLike from .protocol import ReadPromptRequest -from .question_types import LABELS, Question, StructuredDecisionError, label_softmax +from .question_types import ( + LABELS, + QUESTION_TYPES, + Question, + StructuredDecisionError, + label_softmax, +) @dataclass @@ -73,17 +79,19 @@ async def read( def reply_label_ids( - tokenizer: TokenizerLike, prompt_ids: Sequence[int] + tokenizer: TokenizerLike, + prompt_ids: Sequence[int], + labels: Sequence[str] = LABELS, ) -> tuple[list[int], list[int]]: - """The prompt's ids after its last added token, and the token each of LABELS - adds after them as the first token of the reply. Raises ValueError if a - label is not one distinct token there.""" + """The prompt's ids after its last added token, and the token each of + ``labels`` adds after them as the first token of the reply. Raises + ValueError if a label is not one distinct token there.""" added = set(tokenizer.get_added_vocab().values()) start = max((i + 1 for i, t in enumerate(prompt_ids) if t in added), default=0) tail = list(prompt_ids[start:]) text = tokenizer.decode(tail) ids: list[int] = [] - for label in LABELS: + for label in labels: extended = tokenizer.encode(text + label, add_special_tokens=False) if extended[:-1] != tail or extended[-1] in ids: raise ValueError( @@ -116,9 +124,13 @@ def __init__(self, context: ReadContext): ) if isinstance(probe, str): probe = tokenizer.encode(probe, add_special_tokens=False) - # Every prompt ends with the same generation prompt, so the labels' - # tokens are the same for every question. - self.tail, self.label_ids = reply_label_ids(tokenizer, probe) + # Every prompt ends with the same generation prompt, so a label's token + # is the same for every question. + self.tail, _ = reply_label_ids(tokenizer, probe, ()) + self.label_ids: dict[str, int] = {} + for qtype in QUESTION_TYPES.values(): + _, ids = reply_label_ids(tokenizer, probe, qtype.label_set) + self.label_ids.update(zip(qtype.label_set, ids)) def limits(self) -> DecisionLimits: return DecisionLimits(max_questions=64, max_options=len(LABELS)) @@ -173,7 +185,7 @@ async def read( "these chat options end the prompt differently, so the " "labels' tokens are unknown" ) - slots.append(self.label_ids[: len(q.labels)]) + slots.append([self.label_ids[label] for label in q.labels]) engine_inputs.append(engine_input) label_reads = await next_token_label_reads( From ad73a4740dd780c5620099261738a30b979b262c Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Fri, 9 Oct 2026 14:07:49 +0800 Subject: [PATCH 152/624] [Perf][DSv4.1] Trim the last KV-source layer in decoder replay (#59894) Signed-off-by: Canlin Co-authored-by: Claude Opus 5.5 Co-authored-by: Juntian Liu --- ...test_deepseek_v41_decoder_replay_layers.py | 37 +++++++--- .../models/test_deepseek_v41_replay_batch.py | 45 +++++++++-- vllm/models/deepseek_v41/attention.py | 28 +++++++ .../deepseek_v41/decoder_replay_layers.py | 48 ++++++++---- vllm/models/deepseek_v41/nvidia/model.py | 62 ++++++++++++---- .../models/deepseek_v41/nvidia/model_state.py | 74 +++++++++++-------- 6 files changed, 219 insertions(+), 75 deletions(-) diff --git a/tests/models/test_deepseek_v41_decoder_replay_layers.py b/tests/models/test_deepseek_v41_decoder_replay_layers.py index f3b124c5ac80..614d903ca10d 100644 --- a/tests/models/test_deepseek_v41_decoder_replay_layers.py +++ b/tests/models/test_deepseek_v41_decoder_replay_layers.py @@ -67,10 +67,8 @@ def _states(num_tokens: int): ) -def test_run_gathers_states_and_realigns_shared_indexer_buffers(): - topk = torch.arange(NUM_TOKENS * 4, device=DEVICE).view(NUM_TOKENS, 4).int() - candidates = torch.arange(NUM_TOKENS * 3, device=DEVICE).view(NUM_TOKENS, 3).int() - topk_before, candidates_before = topk.clone(), candidates.clone() +def test_run_gathers_states(): + """Every row's KV is written first; the layers then run on the replay rows.""" seen = {} def run_layers(hidden_states, positions, input_ids, pre_mix, post, res, residual): @@ -80,7 +78,8 @@ def run_layers(hidden_states, positions, input_ids, pre_mix, post, res, residual seen["is_padding"] = replay_context.is_padding return residual, pre_mix - layers = DecoderReplayLayers(WINDOW, run_layers, [topk, candidates], set()) + write_batch_kv = MagicMock() + layers = DecoderReplayLayers(WINDOW, run_layers, write_batch_kv, set(), "swa") states = _states(NUM_TOKENS) hidden, residual = states[0], states[-1] full, sub = object(), object() @@ -90,11 +89,10 @@ def run_layers(hidden_states, positions, input_ids, pre_mix, post, res, residual outputs = layers(*states) assert get_forward_context() is context + write_batch_kv.assert_called_once_with(*states) rows = torch.tensor(REPLAY_ROWS, device=DEVICE) assert torch.equal(seen["hidden_states"], hidden[rows]) assert seen["attn_metadata"] is sub and seen["is_padding"] is None - assert torch.equal(topk[: len(REPLAY_ROWS)], topk_before[rows]) - assert torch.equal(candidates[: len(REPLAY_ROWS)], candidates_before[rows]) assert len(outputs) == 2 and outputs[0].shape == residual.shape assert torch.equal(outputs[0][rows], residual[rows]) assert outputs[0][1:173].abs().sum() == 0 @@ -107,12 +105,14 @@ def run_layers(hidden_states, *rest): seen["attn_metadata"] = get_forward_context().attn_metadata return (hidden_states,) - layers = DecoderReplayLayers(WINDOW, run_layers, [], set()) + write_batch_kv = MagicMock() + layers = DecoderReplayLayers(WINDOW, run_layers, write_batch_kv, set(), "swa") states = _states(NUM_TOKENS) full = object() with override_forward_context(_context(full)): (output,) = layers(*states) assert output is states[0] and seen["attn_metadata"] is full + write_batch_kv.assert_not_called() def _graph_context(num_tokens): @@ -125,6 +125,25 @@ def _graph_context(num_tokens): ) +def test_piecewise_graphs_break_out_by_size(): + """A PIECEWISE graph's KV write and break are fixed at capture, when no replay + batch is set, so both follow the graph's size rather than the replay batch.""" + write_batch_kv = MagicMock() + layers = DecoderReplayLayers( + WINDOW, lambda *states: ("whole",), write_batch_kv, set(), "swa" + ) + layers._run_in_graph_break = lambda *states: ("break",) # type: ignore[method-assign] + states = _states(NUM_TOKENS) + with override_forward_context(_graph_context(1024)): + assert layers(*states) == ("whole",) # no threshold: no break + layers.trim_threshold = 768 + assert layers(*states) == ("break",) + _set_replay_batch(layers, REPLAY_ROWS, object()) + with override_forward_context(_graph_context(576)): + assert not layers.uses_replay_batch() # a replay batch alone does not break out + write_batch_kv.assert_called_once_with(*states) + + def test_replay_graph_matches_eager(monkeypatch): """The replay graph of the next captured size pads and runs the rows.""" monkeypatch.setattr(cudagraph_utils, "get_pp_group", MagicMock) @@ -142,7 +161,7 @@ def test_replay_graph_matches_eager(monkeypatch): def run_layers(hidden, positions, ids, pre, post, mix, residual): return residual + hidden[:, None] + positions[:, None, None], pre + ids[:, None] - layers = DecoderReplayLayers(4, run_layers, [], set()) + layers = DecoderReplayLayers(4, run_layers, MagicMock(), set(), "swa") manager = DecoderReplayCudaGraphManager(cfg, DEVICE, layers) def prepare(desc): diff --git a/tests/models/test_deepseek_v41_replay_batch.py b/tests/models/test_deepseek_v41_replay_batch.py index 810a1dbd9a02..bcd4db64b25c 100644 --- a/tests/models/test_deepseek_v41_replay_batch.py +++ b/tests/models/test_deepseek_v41_replay_batch.py @@ -11,6 +11,12 @@ import torch from vllm.config import CUDAGraphMode +from vllm.forward_context import ( + BatchDescriptor, + ForwardContext, + override_forward_context, +) +from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers @@ -59,7 +65,7 @@ def state(monkeypatch): cfg.parallel_config.data_parallel_size = 1 cfg.compilation_config.fast_moe_cold_start = False cfg.compilation_config.cudagraph_mode = CUDAGraphMode.NONE - layers = SimpleNamespace(window=WINDOW, replay_batch=None, metadata_prefixes=set()) + layers = DecoderReplayLayers(WINDOW, MagicMock(), MagicMock(), set(), "swa_first") model = SimpleNamespace(token_lookback_depth=0, decoder_replay_layers=layers) builds: list = [] @@ -72,7 +78,7 @@ def prepare_attn(self, input_batch, cg_mode, block_tables, slot_mappings, *a, ** attn_metadata=object(), ) builds.append(build) - return {"swa": build.attn_metadata} + return {"swa": build.attn_metadata, "swa_first": build.attn_metadata} monkeypatch.setattr(DefaultModelState, "prepare_attn", prepare_attn) state = DeepseekV41ModelState(cfg, model, None, DEVICE) @@ -120,7 +126,8 @@ def _slot_mappings(num_tokens: int) -> torch.Tensor: def _prepare(state, batch, cg_mode): - """Runs prepare_attn; returns the replay batch and the replay build's inputs.""" + """Runs prepare_attn; returns the replay batch and the replay build's inputs. + The first replay layer's sliding-window build follows the replay build.""" state.prepare_attn( batch, cg_mode, @@ -130,7 +137,7 @@ def _prepare(state, batch, cg_mode): GROUPS, ) replay = state.decoder_replay_layers.replay_batch - return (None, None) if replay is None else (replay, state.builds[-1]) + return (None, None) if replay is None else (replay, state.builds[-2]) def test_replay_batch_keeps_each_request_window(state): @@ -151,8 +158,16 @@ def test_replay_batch_keeps_each_request_window(state): assert build.replay_start.tolist() == [0, 300 - WINDOW, 50] assert torch.equal(sub.positions, batch.positions[rows]) assert torch.equal(build.slot_mappings, _slot_mappings(401)[:, rows]) + # The first replay layer keeps the encoder-side window starts: it wrote + # every row's window KV. + first = state.builds[-1] + assert first.batch is sub and first.cg_mode == CUDAGraphMode.NONE + assert first.replay_start.tolist() == [0, 0, 50] context = replay.forward_context - assert context.attn_metadata == {"swa": build.attn_metadata} + assert context.attn_metadata == { + "swa": build.attn_metadata, + "swa_first": first.attn_metadata, + } assert torch.equal(context.slot_mapping["mla"], _slot_mappings(401)[1, rows]) assert context.dp_metadata is None and context.is_padding is None @@ -190,7 +205,8 @@ def test_replay_batch_keeps_adaptive_verification_query_bound(state): def test_graph_steps_trim_only_piecewise_at_threshold(state): """A CUDA graph keeps the layers on its whole batch, unless it is a PIECEWISE - one at the trim threshold or above; an eager step trims a batch to its rows.""" + one at the trim threshold or above, as its capture decided; an eager step + trims a batch to its rows.""" batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING, num_tokens_after_padding=512) for cg_mode in (CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL): assert _prepare(state, batch, cg_mode) == (None, None) @@ -199,8 +215,21 @@ def test_graph_steps_trim_only_piecewise_at_threshold(state): assert replay.rows.tolist() == REPLAY_ROWS assert build.cg_mode == CUDAGraphMode.NONE assert build.batch.num_tokens_after_padding == len(REPLAY_ROWS) - state._trim_threshold = 256 - replay, _ = _prepare(state, batch, CUDAGraphMode.PIECEWISE) + layers = state.decoder_replay_layers + graph = ForwardContext( + no_compile_layers={}, + attn_metadata={}, + slot_mapping={}, + cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE, + batch_descriptor=BatchDescriptor(num_tokens=512), + ) + for threshold in (513, 512): + layers.trim_threshold = threshold + replay, _ = _prepare(state, batch, CUDAGraphMode.PIECEWISE) + with override_forward_context(graph): # the capture's decision + assert ( + layers.uses_replay_batch() == (replay is not None) == (threshold == 512) + ) assert replay is not None and replay.rows.tolist() == REPLAY_ROWS assert _prepare(state, batch, CUDAGraphMode.FULL) == (None, None) diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 73d9c0c7233d..559d874cc483 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -686,6 +686,34 @@ def forward( ) return self._o_proj(attn_out, positions) + def forward_kv(self, positions: torch.Tensor, hidden_states: torch.Tensor) -> None: + """Write only the KV ``forward`` writes (SWA KV, compressed KV, index K), + skipping wq_b, the indexer's top-k and attention.""" + attn_metadata = get_forward_context().attn_metadata + qr_kv, kv_score, _ = self._run_parallel_input_projections(hidden_states) + _, _, kv = self._split_qkv_and_norm(qr_kv) + if isinstance(attn_metadata, dict): + # The KV half of _fused_qnorm_rope_kv_insert, bit for bit, without + # the unused query; same as dspark's _insert_context_kv. + cache = self.swa_cache_layer.kv_cache + torch.ops._C.fused_deepseek_v4_kv_rope_insert( + kv, + cache, + cast(Any, attn_metadata[self.swa_cache_layer.prefix]).slot_mapping, + positions, + self.rotary_emb.cos_sin_cache, + self.swa_cache_layer.block_size, + self._flashinfer_fp8_kv_scale + if cache.dtype == torch.float8_e4m3fn + else None, + self.kv_mxfp8, + ) + if self.compressor is not None: + latent = self.compressor(kv_score, positions) + self.compressor.insert_cache(latent, positions, self.rotary_emb) + if self.indexer is not None and self.indexer.owns_k: + self.indexer._produce_k(latent, positions, self.indexer_rotary_emb) + def bind_gemm_rs(self) -> None: """Fuse ``wo_b`` with the sequence-parallel TP reduce-scatter. diff --git a/vllm/models/deepseek_v41/decoder_replay_layers.py b/vllm/models/deepseek_v41/decoder_replay_layers.py index d2d7ffd15364..28aa5e399f63 100644 --- a/vllm/models/deepseek_v41/decoder_replay_layers.py +++ b/vllm/models/deepseek_v41/decoder_replay_layers.py @@ -2,8 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Decoder-side SWA bounded replay: running the replay layers on their batch. -The layers past the last KV-source layer own nothing but sliding-window KV, so -in eager prefill steps they run on each request's last ``window`` rows only. +Once the last KV-source layer has written every row's KV, it and the layers past +it (which own only SWA KV) run on each request's last ``window`` rows only. ``DeepseekV41ModelState`` prepares those rows as a sub-batch with attention metadata and a forward context of its own, like a microbatch; ``DecoderReplayLayers`` gathers the layer inputs by its rows, runs the layers @@ -21,6 +21,7 @@ from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.forward_context import ( ForwardContext, + get_forward_context, in_piecewise_cudagraph, override_forward_context, ) @@ -39,35 +40,57 @@ class ReplayBatch: class DecoderReplayLayers: """Runs the replay layers on the step's replay batch. - ``run_layers`` takes a batch's layer inputs and returns its per-row - outputs. ``row_buffers`` hold per-row results the source layer's indexer - publishes for the layers after it; they are compacted to the replay rows - in place. ``metadata_prefixes`` are the attention metadata keys the layers - read, so the replay batch builds only those. + ``run_layers`` takes a batch's layer inputs and returns its per-row outputs; + ``write_batch_kv`` takes them too and writes the first replay layer's KV for + every row. ``metadata_prefixes`` are the attention metadata keys the layers + read, so the replay batch builds only those. ``first_swa_prefix`` keys the + first replay layer's SWA metadata, built with the batch's window starts. """ def __init__( self, window: int, run_layers: Callable[..., tuple[torch.Tensor, ...]], - row_buffers: list[torch.Tensor], + write_batch_kv: Callable[..., None], metadata_prefixes: set[str], + first_swa_prefix: str, ) -> None: self.window = window self.run_layers = run_layers - self.row_buffers = row_buffers + self.write_batch_kv = write_batch_kv self.metadata_prefixes = metadata_prefixes + self.first_swa_prefix = first_swa_prefix # Set by the model state every step; None runs the layers on the batch. self.replay_batch: ReplayBatch | None = None + # Set by the model state with replay graphs: PIECEWISE model graphs of + # this many tokens or more break out to the replay batch. + self.trim_threshold: int | None = None self._break_outputs: list[torch.Tensor] | None = None def __call__( self, hidden_states: torch.Tensor | MoEOutput, *states: torch.Tensor | None ) -> tuple[torch.Tensor, ...]: - if self.replay_batch is not None and in_piecewise_cudagraph(): - return self._run_in_graph_break(hidden_states, *states) + if self.uses_replay_batch(): + # Outside the break, so a PIECEWISE graph captures it. + self.write_batch_kv(hidden_states, *states) + if in_piecewise_cudagraph(): + return self._run_in_graph_break(hidden_states, *states) return self._run(hidden_states, *states) + def uses_replay_batch(self) -> bool: + """Whether this forward uses the replay batch: in a PIECEWISE graph by its + size, fixed at capture (no replay batch is set then); else if one is set.""" + if in_piecewise_cudagraph(): + batch_descriptor = get_forward_context().batch_descriptor + assert batch_descriptor is not None + return self.graph_uses_replay_batch(batch_descriptor.num_tokens) + return self.replay_batch is not None + + def graph_uses_replay_batch(self, num_tokens: int) -> bool: + """Whether a PIECEWISE graph of ``num_tokens`` breaks out to the replay + batch; capture and the model state's runtime replay batch share it.""" + return self.trim_threshold is not None and num_tokens >= self.trim_threshold + @eager_break_during_capture def _run_in_graph_break( self, hidden_states: torch.Tensor | MoEOutput, *states: torch.Tensor | None @@ -91,9 +114,6 @@ def _run( # Replay batch steps hold more than `window` tokens: no deferred MoE finalize. assert isinstance(hidden_states, torch.Tensor) rows = replay_batch.rows - num_rows = rows.shape[0] - for buf in self.row_buffers: - buf[:num_rows].copy_(buf.index_select(0, rows)) inputs = (hidden_states, *states) with override_forward_context(replay_batch.forward_context): if replay_batch.run_graph is not None: diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 94fa0b4b8b67..18f8f27843ab 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -623,6 +623,40 @@ def forward( torch.cuda.current_stream().wait_stream(mhc_stream) return x, residual, post_mix, res_mix, ffn_pre, previous_aux + def write_kv( + self, + x: torch.Tensor | MoEOutput, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + pre_mix: torch.Tensor, + post_mix: torch.Tensor, + res_mix: torch.Tensor, + residual: torch.Tensor, + ) -> None: + """Write the KV ``forward`` would for the replay layers' inputs, skipping + the rest of the layer; the replay batch then reruns its rows.""" + # The replay batch rules out Engram and sequence parallel here. + assert self.engram is None and not self.use_sequence_parallel + *_, x, _, _ = mhc_shifted_post_pre( + x, + residual, + post_mix, + res_mix, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + pre_mix=pre_mix, + norm_weight=self.attn_norm.weight, + norm_eps=self.attn_norm.variance_epsilon, + reduce_results=self.fuse_mhc_all_reduce, + ) + self.attn.forward_kv(positions, x) + class DeepseekV4Model(nn.Module, EagleModelMixin): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -765,10 +799,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): and self._decoder_replay_supported(vllm_config, cut) and self.layers[cut].attn.swa_cache_layer.bounded_replay ): - self.decoder_replay_start = cut + 1 + assert self.layers[cut].attn.compress_ratio <= 1 + # The last KV source writes every row's KV, then replays. + self.decoder_replay_start = cut # The attention metadata keys the replay layers read. metadata_prefixes = set() - for layer in islice(self.layers, cut + 1, self.end_layer): + for layer in islice(self.layers, cut, self.end_layer): attn = typing.cast(DeepseekV4DecoderLayer, layer).attn metadata_prefixes.add(attn.swa_cache_layer.prefix) if attn.compressed_cache_prefix is not None: @@ -778,16 +814,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.decoder_replay_layers = DecoderReplayLayers( config.sliding_window, self._run_replay_layers, - [ - buf - for buf in (self.topk_indices_buffer, self.candidate_block_buffer) - if buf is not None - ], + self.layers[cut].write_kv, metadata_prefixes, + self.layers[cut].attn.swa_cache_layer.prefix, ) logger.info_once( - "Decoder SWA bounded replay: in eager prefill steps, layers " + "Decoder SWA bounded replay: in eager prefill steps, layer %d " + "writes its KV for every token, and the rest of it and layers " "%d-%d run on each request's last %d tokens only.", + cut, cut + 1, self.end_layer - 1, config.sliding_window, @@ -1120,9 +1155,8 @@ def _run_replay_layers( res_mix: torch.Tensor, residual: torch.Tensor, ) -> tuple[torch.Tensor, ...]: - """The layers past the last KV source, on whatever rows they are given; - returns their output, the last FFN's pre-mix and the aux hidden states - they capture.""" + """Layers ``decoder_replay_start``.. on the given rows; returns their output, + the last FFN's pre-mix and the aux hidden states they capture.""" aux_hidden_by_layer: dict[int, torch.Tensor] = {} hidden_states, residual, post_mix, res_mix, pre_mix = self._run_layers( range(self.decoder_replay_start, self.end_layer), @@ -1154,7 +1188,7 @@ def _run_replay_layers( ) def _decoder_replay_supported(self, vllm_config: VllmConfig, cut: int) -> bool: - """Whether this rank may trim the layers after ``cut``; warns when not.""" + """Whether this rank may replay from layer ``cut`` on; warns when not.""" parallel_config = vllm_config.parallel_config spec_config = vllm_config.speculative_config draft_config = spec_config.draft_model_config if spec_config else None @@ -1176,8 +1210,8 @@ def _decoder_replay_supported(self, vllm_config: VllmConfig, cut: int) -> bool: "the replay-layer batch shrinks per rank, which sequence and " "prefill-context parallelism and microbatching cannot follow" ) - elif any(i > cut for i in getattr(self.config, "engram_layer_ids", ())): - reason = "an Engram layer sits after the last KV source layer" + elif any(i >= cut for i in getattr(self.config, "engram_layer_ids", ())): + reason = "an Engram layer sits at or after the last KV source layer" elif draft_config is not None and ( draft_window is None or draft_window > window diff --git a/vllm/models/deepseek_v41/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py index c09e4582ea34..cc3ade80b8ad 100644 --- a/vllm/models/deepseek_v41/nvidia/model_state.py +++ b/vllm/models/deepseek_v41/nvidia/model_state.py @@ -238,12 +238,13 @@ def __init__( self._kept_kv_start = torch.zeros( self.max_num_reqs, dtype=torch.int32, device=device ) - self._replay_attn_groups: list[list[AttentionGroup]] | None = None + # The replay layers' and the first one's SWA attention groups. + self._replay_attn_groups: ( + tuple[list[list[AttentionGroup]], list[list[AttentionGroup]]] | None + ) = None layers = self.decoder_replay_layers compilation_config = vllm_config.compilation_config self.replay_cudagraphs: DecoderReplayCudaGraphManager | None = None - # PIECEWISE graphs of this many tokens or more break out to the replay batch. - self._trim_threshold: int | None = None if ( layers is not None and compilation_config.cudagraph_decoder_replay @@ -263,7 +264,7 @@ def __init__( self.replay_cudagraphs = DecoderReplayCudaGraphManager( vllm_config, device, layers ) - self._trim_threshold = threshold + layers.trim_threshold = threshold def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: super().add_request(req_index, new_req_data) @@ -454,8 +455,7 @@ def _prepare_replay_batch( layers.replay_batch = None force_replay_batch = capture_desc is not None or ( cudagraph_mode == CUDAGraphMode.PIECEWISE - and self._trim_threshold is not None - and input_batch.num_tokens_after_padding >= self._trim_threshold + and layers.graph_uses_replay_batch(input_batch.num_tokens_after_padding) ) if cudagraph_mode != CUDAGraphMode.NONE and not force_replay_batch: return @@ -509,17 +509,35 @@ def _prepare_replay_batch( kept_batch, kept_slot_mappings = self._kept_input_batch( input_batch, slot_mappings, replay_start, kept_lens, num_padded ) + if self._replay_attn_groups is None: + self._replay_attn_groups = ( + self._copy_attn_groups(attn_groups, layers.metadata_prefixes), + self._copy_attn_groups(attn_groups, {layers.first_swa_prefix}), + ) + replay_groups, first_swa_groups = self._replay_attn_groups attn_metadata = super().prepare_attn( kept_batch, CUDAGraphMode.NONE, block_tables, kept_slot_mappings, - self._replay_groups(attn_groups), + replay_groups, kv_cache_config, model_specific_attn_metadata=ReplayAttnMetadata( self._kept_kv_start[:num_reqs] ), ) + # The first replay layer wrote every row's window KV: its window starts + # are the batch's. + first_swa = layers.first_swa_prefix + attn_metadata[first_swa] = super().prepare_attn( + kept_batch, + CUDAGraphMode.NONE, + block_tables, + kept_slot_mappings, + first_swa_groups, + kv_cache_config, + model_specific_attn_metadata=ReplayAttnMetadata(replay_start), + )[first_swa] layers.replay_batch = ReplayBatch( self._kept_rows[:num_tokens], create_forward_context( @@ -636,27 +654,23 @@ def _agree_across_dp( return False, None return True, agreed[1] - def _replay_groups( - self, attn_groups: list[list[AttentionGroup]] + def _copy_attn_groups( + self, attn_groups: list[list[AttentionGroup]], prefixes: set[str] ) -> list[list[AttentionGroup]]: - """The attention groups with metadata builders of their own, like each - microbatch's: a builder keeps the metadata it built, and the runner's - hold the batch's. Only the groups the replay layers read are built.""" - if self._replay_attn_groups is None: - assert self.decoder_replay_layers is not None - prefixes = self.decoder_replay_layers.metadata_prefixes - self._replay_attn_groups = [] - for groups in attn_groups: - replay_groups = [] - for group in groups: - if prefixes.isdisjoint(group.layer_names): - continue - replay_group = replace(group, metadata_builders=[]) - replay_group.create_metadata_builders( - self.vllm_config, - self.device, - group.metadata_builders[0].kernel_block_size, - ) - replay_groups.append(replay_group) - self._replay_attn_groups.append(replay_groups) - return self._replay_attn_groups + """Copies of the groups holding any of ``prefixes``, with builders of their + own like a microbatch's: a builder keeps the metadata it built.""" + own_groups = [] + for groups in attn_groups: + copies = [] + for group in groups: + if prefixes.isdisjoint(group.layer_names): + continue + copy = replace(group, metadata_builders=[]) + copy.create_metadata_builders( + self.vllm_config, + self.device, + group.metadata_builders[0].kernel_block_size, + ) + copies.append(copy) + own_groups.append(copies) + return own_groups From f0a5f111f205b5c73bb34fbd41f8f0d9936b543f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Ma=C5=84czak?= Date: Fri, 9 Oct 2026 08:29:29 +0200 Subject: [PATCH 153/624] [XPU] Add tuned Mamba SSU configs for Intel Arc Pro B60 (#56765) Signed-off-by: pmanczak Co-authored-by: Kunshang Ji --- ..._Pro_B60_Graphics,cache_dtype=float16.json | 51 +++++++++++++++ ..._Pro_B60_Graphics,cache_dtype=float32.json | 51 +++++++++++++++ ..._Pro_B60_Graphics,cache_dtype=float16.json | 63 +++++++++++++++++++ ..._Pro_B60_Graphics,cache_dtype=float32.json | 63 +++++++++++++++++++ ..._Pro_B60_Graphics,cache_dtype=float16.json | 63 +++++++++++++++++++ ..._Pro_B60_Graphics,cache_dtype=float32.json | 63 +++++++++++++++++++ ..._Pro_B60_Graphics,cache_dtype=float16.json | 63 +++++++++++++++++++ ..._Pro_B60_Graphics,cache_dtype=float32.json | 63 +++++++++++++++++++ 8 files changed, 480 insertions(+) create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json new file mode 100644 index 000000000000..64087ce74c28 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json @@ -0,0 +1,51 @@ +{ + "triton_version": "3.8.0", + "8": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "512": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + } +} diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json new file mode 100644 index 000000000000..0a378592c522 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json @@ -0,0 +1,51 @@ +{ + "triton_version": "3.8.0", + "8": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "16": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "32": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "256": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "512": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "1024": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + } +} diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json new file mode 100644 index 000000000000..91c6e28fa794 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json @@ -0,0 +1,63 @@ +{ + "triton_version": "3.8.0", + "8": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "16": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "32": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "64": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "128": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "256": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "512": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "8192": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "16384": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "32768": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "65536": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "131072": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + } +} diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json new file mode 100644 index 000000000000..c30784ce7762 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json @@ -0,0 +1,63 @@ +{ + "triton_version": "3.8.0", + "8": { + "BLOCK_SIZE_M": 8, + "num_warps": 8 + }, + "16": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "512": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "32768": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "65536": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "131072": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + } +} diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json new file mode 100644 index 000000000000..5f29bcd76c22 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json @@ -0,0 +1,63 @@ +{ + "triton_version": "3.8.0", + "8": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "512": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "32768": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "65536": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "131072": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + } +} diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json new file mode 100644 index 000000000000..8f815a1b337d --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json @@ -0,0 +1,63 @@ +{ + "triton_version": "3.8.0", + "8": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "16": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "32": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "64": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "256": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "512": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "1024": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "8192": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "16384": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "32768": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "65536": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "131072": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + } +} diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json new file mode 100644 index 000000000000..2df252554195 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json @@ -0,0 +1,63 @@ +{ + "triton_version": "3.8.0", + "8": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "128": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "256": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "512": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "32768": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "65536": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "131072": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + } +} diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json new file mode 100644 index 000000000000..f943b523c441 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json @@ -0,0 +1,63 @@ +{ + "triton_version": "3.8.0", + "8": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "512": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "32768": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "65536": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "131072": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + } +} From dd30786fef8b15bd83205833c7b3348b0ba0ebc7 Mon Sep 17 00:00:00 2001 From: Chinmay-Kulkarni-AMD Date: Fri, 9 Oct 2026 12:10:21 +0530 Subject: [PATCH 154/624] [CI] Enable Zen CPU image build on AMD CI (#60016) Signed-off-by: Chinmay Kulkarni Co-authored-by: Cursor Co-authored-by: Micah Williamson --- .buildkite/hardware_tests/cpu.yaml | 54 +++++++++++++ .buildkite/image_build/image_build_zen_cpu.sh | 75 +++++++++++++++++++ 2 files changed, 129 insertions(+) create mode 100755 .buildkite/image_build/image_build_zen_cpu.sh diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index d54688b57491..a29b44ed3c71 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -232,3 +232,57 @@ steps: "$$BUILDKITE_PARALLEL_JOB" "$REGISTRY/$REPO:$BUILDKITE_COMMIT-arm64-cpu" parallelism: 3 + +- label: ":docker: Build Zen CPU image" + key: zen-cpu-image-build + # Reuse the shared `:-cpu` image that image-build-cpu publishes + # to public ECR, so Dockerfile.zen just layers on top instead of recompiling + # vLLM from source. If the pull fails, the build script still falls back to + # building the CPU base from source. + depends_on: + - image-build-cpu + soft_fail: true + agents: + queue: amd-cpu + source_file_dependencies: + - .buildkite/hardware_tests/cpu.yaml + - .buildkite/image_build/image_build_zen_cpu.sh + - docker/Dockerfile.zen + - docker/Dockerfile.cpu + - requirements/cpu.txt + - requirements/build/cpu.txt + - requirements/test/cpu.txt + - cmake/cpu_extension.cmake + - vllm/model_executor/layers/utils.py + - vllm/model_executor/kernels/linear/zentorch_utils.py + - vllm/model_executor/kernels/linear/scaled_mm/zentorch.py + - vllm/model_executor/kernels/linear/mixed_precision/zentorch.py + - vllm/platforms/cpu.py + - vllm/platforms/zen_cpu.py + - vllm/platforms/__init__.py + - tests/model_executor/test_cpu_unquantized_gemm_dispatch.py + - tests/test_zen_cpu_platform_detection.py + - tests/kernels/moe/test_zen_cpu_fused_moe.py + - tests/kernels/moe/test_zen_cpu_int8_moe.py + - vllm/v1/attention/backends/zentorch_sdpa.py + - tests/test_zentorch_version_pin.py + - vllm/v1/attention/backends/mla/prefill/zen_cpu_sdpa.py + - vllm/model_executor/layers/fused_moe/experts/zentorch_moe.py + - tests/kernels/quantization/test_zen_da8w4.py + commands: + - .buildkite/image_build/image_build_zen_cpu.sh "$REGISTRY" "$REPO" "$BUILDKITE_COMMIT" + env: + DOCKER_BUILDKIT: "1" + # Install zentorch directly (if-branch in Dockerfile.zen) instead of + # `uv pip install "vllm[zen]"`, which forces uv to re-resolve vLLM's full + # dependency closure and chokes on its transitive triton-cpu URL dependency. + # Keep this pin in sync with the `zen` extra in setup.py. + ZENTORCH_VERSION: "2.13.0.1" + retry: + automatic: + - exit_status: 1 # Transient Docker/BuildKit failure + limit: 1 + - exit_status: -1 # Agent was lost + limit: 2 + - exit_status: -10 # Agent was lost + limit: 2 diff --git a/.buildkite/image_build/image_build_zen_cpu.sh b/.buildkite/image_build/image_build_zen_cpu.sh new file mode 100755 index 000000000000..3e42e38a8d75 --- /dev/null +++ b/.buildkite/image_build/image_build_zen_cpu.sh @@ -0,0 +1,75 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Build the AMD Zen CPU image (vLLM + zentorch) by layering docker/Dockerfile.zen +# on top of a CPU base image that already has vLLM installed. +# +# To avoid recompiling vLLM from source, prefer reusing the shared CPU image that +# the image-build-cpu step already builds and publishes for this commit +# (`:-cpu`). It lives in public ECR, which allows anonymous pulls, +# so no registry credentials are required. If the pull fails (e.g. the image +# isn't published yet), fall back to building the base from docker/Dockerfile.cpu. +# +# The zen image is not pushed to a registry. +# +# See docker/Dockerfile.zen for the build workflow this mirrors. +set -euo pipefail + +if [[ $# -lt 3 ]]; then + echo "Usage: $0 " + exit 1 +fi + +REGISTRY=$1 +REPO=$2 +BUILDKITE_COMMIT=$3 + +image_repo() { + if [[ -n "${REGISTRY}" ]]; then + printf '%s/%s' "${REGISTRY}" "${REPO}" + else + printf '%s' "${REPO}" + fi +} + +IMAGE_REPO="$(image_repo)" + +# The shared CPU image published by image-build-cpu (public ECR, anonymous pull). +SHARED_CPU_IMAGE="" +if [[ -n "${REGISTRY}" ]]; then + SHARED_CPU_IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu" +fi +FALLBACK_BASE_IMAGE="zen-cpu-base:$BUILDKITE_COMMIT" +IMAGE="$IMAGE_REPO:$BUILDKITE_COMMIT-zen-cpu" + +# ZENTORCH_VERSION is optional; when unset the Dockerfile falls back to +# installing zentorch via `vllm[zen]`. +ZENTORCH_VERSION=${ZENTORCH_VERSION:-} + +# Step 1: obtain the CPU base image that Dockerfile.zen layers on. Prefer pulling +# the published `-cpu` image; fall back to building it from source if unavailable. +if [[ -n "${SHARED_CPU_IMAGE}" ]] && docker pull "$SHARED_CPU_IMAGE"; then + echo "--- :docker: Using published CPU image as base: $SHARED_CPU_IMAGE" + BASE_IMAGE="$SHARED_CPU_IMAGE" +else + echo "--- :docker: Published CPU image unavailable; building base from source" + docker build --file docker/Dockerfile.cpu \ + --platform linux/amd64 \ + --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ + --build-arg VLLM_CPU_X86=true \ + --tag "$FALLBACK_BASE_IMAGE" \ + --target vllm-openai \ + --progress plain . + BASE_IMAGE="$FALLBACK_BASE_IMAGE" +fi + +# Step 2: build the zen test image on top of the CPU base. +echo "--- :docker: Building Zen test image" +docker build --file docker/Dockerfile.zen \ + --platform linux/amd64 \ + --build-arg BASE_IMAGE="$BASE_IMAGE" \ + ${ZENTORCH_VERSION:+--build-arg ZENTORCH_VERSION="$ZENTORCH_VERSION"} \ + --tag "$IMAGE" \ + --target vllm-zen-test \ + --progress plain . From 19c0bed584bde3fe785296d2923247ab9d678483 Mon Sep 17 00:00:00 2001 From: Jakub Zakrzewski Date: Fri, 9 Oct 2026 09:12:03 +0200 Subject: [PATCH 155/624] [Fusion] Support 3D activation quantization fusion and enable it for Pixtral (#59612) Signed-off-by: Jakub Zakrzewski Co-authored-by: Codex Co-authored-by: Misha Goin --- .../fusion/test_quant_activation_contract.py | 2 + .../test_flashinfer_nvfp4_scaled_mm.py | 52 +++++++++++++++++++ tests/kernels/test_fused_quant_activation.py | 32 +++++++----- .../kernels/linear/nvfp4/flashinfer.py | 31 +++++++---- .../layers/fusion/fused_act_quant.py | 8 +-- vllm/model_executor/models/pixtral.py | 5 +- 6 files changed, 101 insertions(+), 29 deletions(-) diff --git a/tests/fusion/test_quant_activation_contract.py b/tests/fusion/test_quant_activation_contract.py index 15b9835d6cf5..da507e5caf50 100644 --- a/tests/fusion/test_quant_activation_contract.py +++ b/tests/fusion/test_quant_activation_contract.py @@ -21,6 +21,7 @@ NvFp4LinearLayerConfig, ) from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( + FlashInferCuteDslNvFp4LinearKernel, FlashInferCutlassNvFp4LinearKernel, FlashInferTrtllmNvFp4LinearKernel, ) @@ -59,6 +60,7 @@ SUPPORTING = { CutlassFP8ScaledMMLinearKernel, FlashInferFP8ScaledMMLinearKernel, + FlashInferCuteDslNvFp4LinearKernel, FlashInferCutlassNvFp4LinearKernel, PerTensorTorchFP8ScaledMMLinearKernel, AiterHipbMMPerTokenFp8ScaledMMLinearKernel, diff --git a/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py b/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py index 54dcaa0b97ea..47c25d046d2c 100644 --- a/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py +++ b/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py @@ -13,11 +13,15 @@ from vllm import _custom_ops as ops from vllm.model_executor.kernels.linear.nvfp4 import NvFp4LinearLayerConfig from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( + FlashInferCuteDslNvFp4LinearKernel, FlashInferCuteDslNvFp4W4A16LinearKernel, + FlashInferCutlassNvFp4LinearKernel, ) +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( dequantize_to_dtype, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic from vllm.platforms import current_platform from vllm.utils.flashinfer import ( flashinfer_scaled_fp4_mm, @@ -47,6 +51,54 @@ CUDA_DEVICES = ["cuda:0"] +@pytest.mark.parametrize("token_shape", [(17,), (1, 17), (2, 17)]) +@pytest.mark.parametrize( + "kernel_cls", + [FlashInferCutlassNvFp4LinearKernel, FlashInferCuteDslNvFp4LinearKernel], +) +@torch.inference_mode() +def test_flashinfer_nvfp4_preserves_quantized_input_shape( + token_shape: tuple[int, ...], + kernel_cls: type[FlashInferCutlassNvFp4LinearKernel] + | type[FlashInferCuteDslNvFp4LinearKernel], +) -> None: + """Pre-quantized inputs must match the tensor path and retain token dimensions.""" + supported, reason = kernel_cls.is_supported() + if not supported: + pytest.skip(reason) + + set_random_seed(42) + hidden_size, output_size = 128, 128 + x = torch.randn(*token_shape, hidden_size, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(output_size, hidden_size, dtype=x.dtype, device=x.device) + input_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / x.abs().max()).float() + weight_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / weight.abs().max()).float() + x_fp4, x_blockscale = ops.scaled_fp4_quant(x.reshape(-1, hidden_size), input_scale) + weight_fp4, weight_blockscale = ops.scaled_fp4_quant(weight, weight_scale) + quantized_input = QuantizedActivation( + data=x_fp4.reshape(*token_shape, hidden_size // 2), + scale=x_blockscale, + orig_dtype=x.dtype, + orig_shape=x.shape, + quant_key=kNvfp4Dynamic, + ) + layer = torch.nn.Module() + layer.input_size_per_partition = hidden_size + layer.output_size_per_partition = output_size + layer.weight = weight_fp4 + layer.weight_scale = weight_blockscale + layer.input_global_scale_inv = input_scale + layer.alpha = (input_scale * weight_scale).reciprocal() + kernel = kernel_cls(NvFp4LinearLayerConfig()) + + expected = kernel.apply_weights(layer, x.reshape(-1, hidden_size)) + actual = kernel.apply_weights(layer, quantized_input) + + assert actual.shape == (*token_shape, output_size) + assert actual.dtype == x.dtype + torch.testing.assert_close(actual.reshape(expected.shape), expected) + + def get_ref_results( a_fp4, b_fp4, diff --git a/tests/kernels/test_fused_quant_activation.py b/tests/kernels/test_fused_quant_activation.py index 321cb6978e3d..90309dbd1d9b 100644 --- a/tests/kernels/test_fused_quant_activation.py +++ b/tests/kernels/test_fused_quant_activation.py @@ -115,13 +115,13 @@ class MockLinearRequiresUnquantized(MockLinearFp8Static): requires_unquantized_input = True -@pytest.mark.parametrize("num_tokens", [1, 16, 128]) +@pytest.mark.parametrize("token_shape", [(1,), (16,), (128,), (1, 17), (2, 17)]) @pytest.mark.parametrize("hidden_size", [128, 512, 1024]) @pytest.mark.parametrize("dtype", DTYPES) @torch.inference_mode() def test_maybe_fused_act_quant_fp8_static( default_vllm_config, - num_tokens: int, + token_shape: tuple[int, ...], hidden_size: int, dtype: torch.dtype, ) -> None: @@ -133,26 +133,29 @@ def test_maybe_fused_act_quant_fp8_static( scale = torch.tensor([0.5], device=device, dtype=torch.float32) linear = MockLinearFp8Static(scale) - x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device) + x = torch.randn(*token_shape, hidden_size * 2, dtype=dtype, device=device) result = maybe_fused_act_quant(act_fn, x, linear) assert isinstance(result, QuantizedActivation) assert result.quant_key == kFp8StaticTensorSym assert result.data.dtype == current_platform.fp8_dtype() assert result.orig_dtype == dtype - assert result.orig_shape == (num_tokens, hidden_size) + assert result.orig_shape == (*token_shape, hidden_size) + assert result.data.shape == result.orig_shape - ref_out = ref_impl(act_fn, x, scale) + ref_out = ref_impl(act_fn, x.reshape(-1, hidden_size * 2), scale).reshape( + result.orig_shape + ) torch.testing.assert_close(result.data.to(torch.float32), ref_out.to(torch.float32)) -@pytest.mark.parametrize("num_tokens", [1, 16, 128]) +@pytest.mark.parametrize("token_shape", [(1,), (16,), (128,), (1, 17), (2, 17)]) @pytest.mark.parametrize("hidden_size", [128, 512, 1024]) @pytest.mark.parametrize("dtype", DTYPES) @torch.inference_mode() def test_maybe_fused_act_quant_fp8_dynamic_block( default_vllm_config, - num_tokens: int, + token_shape: tuple[int, ...], hidden_size: int, dtype: torch.dtype, ) -> None: @@ -166,31 +169,34 @@ def test_maybe_fused_act_quant_fp8_dynamic_block( linear = MockLinearFp8Dynamic128() scale = 1 / hidden_size - x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device) * scale + x = torch.randn(*token_shape, hidden_size * 2, dtype=dtype, device=device) * scale result = maybe_fused_act_quant(act_fn, x, linear) assert isinstance(result, QuantizedActivation) assert result.quant_key == kFp8Dynamic128Sym assert result.data.dtype == current_platform.fp8_dtype() assert result.orig_dtype == dtype - assert result.orig_shape == (num_tokens, hidden_size) + assert result.orig_shape == (*token_shape, hidden_size) + assert result.data.shape == result.orig_shape num_groups = hidden_size // group_size - assert result.scale.shape == (num_tokens, num_groups) + assert result.scale.shape == (*token_shape, num_groups) gate, up = x.split(hidden_size, dim=-1) silu_out = F.silu(gate) * up ref_out, ref_scales = per_token_group_quant_fp8( - silu_out, group_size=group_size, use_ue8m0=False + silu_out.reshape(-1, hidden_size), group_size=group_size, use_ue8m0=False ) + ref_out = ref_out.reshape(result.orig_shape) + ref_scales = ref_scales.reshape(*token_shape, num_groups) torch.testing.assert_close(result.scale, ref_scales, rtol=1e-5, atol=1e-5) ref_deq = ref_out.to(torch.float32) * ref_scales.repeat_interleave( - group_size, dim=1 + group_size, dim=-1 ) result_deq = result.data.to(torch.float32) * result.scale.repeat_interleave( - group_size, dim=1 + group_size, dim=-1 ) torch.testing.assert_close(ref_deq, result_deq, atol=5e-2, rtol=5e-2) diff --git a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py index a555ca17bf89..2a69a3be2100 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py +++ b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py @@ -112,6 +112,9 @@ def apply_weights( class FlashInferCuteDslNvFp4LinearKernel(NvFp4LinearKernel): """NVFP4 GEMM via FlashInfer's cutedsl backend.""" + def input_quant_key(self) -> QuantKey | None: + return kNvfp4Dynamic + @classmethod def is_supported( cls, compute_capability: int | None = None @@ -137,19 +140,26 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: def apply_weights( self, layer: torch.nn.Module, - x: torch.Tensor, + x: torch.Tensor | QuantizedActivation, bias: torch.Tensor | None = None, ) -> torch.Tensor: output_size = layer.output_size_per_partition - output_dtype = x.dtype - output_shape = [*x.shape[:-1], output_size] - - x_fp4, x_blockscale = scaled_fp4_quant( - x, - layer.input_global_scale_inv, - is_sf_swizzled_layout=True, - backend="flashinfer-cutedsl", - ) + qa = as_quantized_activation(x, self.input_quant_key()) + if qa is not None: + x_fp4, x_blockscale = qa.data, qa.scale + x_fp4 = x_fp4.reshape(-1, x_fp4.shape[-1]) + output_dtype = qa.orig_dtype + output_shape = [*qa.orig_shape[:-1], output_size] + else: + assert isinstance(x, torch.Tensor) + output_dtype = x.dtype + output_shape = [*x.shape[:-1], output_size] + x_fp4, x_blockscale = scaled_fp4_quant( + x, + layer.input_global_scale_inv, + is_sf_swizzled_layout=True, + backend="flashinfer-cutedsl", + ) x_fp4 = pad_nvfp4_activation_for_cutlass( x_fp4, nvfp4_weight_padding_bytes(layer) @@ -219,6 +229,7 @@ def apply_weights( qa = as_quantized_activation(x, self.input_quant_key()) if qa is not None: x_fp4, x_blockscale = qa.data, qa.scale + x_fp4 = x_fp4.reshape(-1, x_fp4.shape[-1]) x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_bytes) output_dtype = qa.orig_dtype output_shape = [*qa.orig_shape[:-1], output_size] diff --git a/vllm/model_executor/layers/fusion/fused_act_quant.py b/vllm/model_executor/layers/fusion/fused_act_quant.py index a9b14c617000..6261a1bdd019 100644 --- a/vllm/model_executor/layers/fusion/fused_act_quant.py +++ b/vllm/model_executor/layers/fusion/fused_act_quant.py @@ -63,10 +63,10 @@ def _silu_and_mul_fp8_dynamic_block( x: torch.Tensor, linear: LinearBase, group_size: int, quant_key: QuantKey ) -> QuantizedActivation: """SiluAndMul + FP8 dynamic per-block quantization.""" - assert x.ndim == 2, f"Input must be 2D [batch, hidden*2], got {x.shape}" - + assert x.ndim in (2, 3), f"Input must be 2D or 3D, got {x.shape}" d = x.shape[-1] // 2 out_shape = x.shape[:-1] + (d,) + x = x.reshape(-1, x.shape[-1]) num_tokens = x.shape[0] num_groups = d // group_size @@ -102,10 +102,10 @@ def _silu_and_mul_nvfp4_dynamic( x: torch.Tensor, linear: LinearBase ) -> QuantizedActivation: """SiluAndMul + NVFP4 dynamic quantization.""" - assert x.ndim == 2, f"Input must be 2D [batch, hidden*2], got {x.shape}" - + assert x.ndim in (2, 3), f"Input must be 2D or 3D, got {x.shape}" d = x.shape[-1] // 2 out_shape = x.shape[:-1] + (d,) + x = x.reshape(-1, x.shape[-1]) num_tokens = x.shape[0] # NVFP4 packs 2 values into 1 byte diff --git a/vllm/model_executor/models/pixtral.py b/vllm/model_executor/models/pixtral.py index fb10026ad19a..4f96c0acfcc0 100644 --- a/vllm/model_executor/models/pixtral.py +++ b/vllm/model_executor/models/pixtral.py @@ -32,6 +32,7 @@ from vllm.model_executor.layers.activation import SiluAndMul, get_act_and_mul_fn from vllm.model_executor.layers.attention import MMEncoderAttention from vllm.model_executor.layers.conv import Conv2dLayer +from vllm.model_executor.layers.fusion.fused_act_quant import maybe_fused_act_quant from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -758,7 +759,7 @@ def __init__( def forward(self, x): x, _ = self.gate_up_proj(x) - x = self.act_fn(x) + x = maybe_fused_act_quant(self.act_fn, x, self.down_proj) x, _ = self.down_proj(x) return x @@ -1259,7 +1260,7 @@ def __init__( def forward(self, x: torch.Tensor) -> torch.Tensor: gate_up, _ = self.gate_up_proj(x) - x = self.act_and_mul(gate_up) + x = maybe_fused_act_quant(self.act_and_mul, gate_up, self.down_proj) x, _ = self.down_proj(x) return x From 82e89e99440661f3404c58c9c60d712657305cac Mon Sep 17 00:00:00 2001 From: Zheng Cai Date: Fri, 9 Oct 2026 15:18:50 +0800 Subject: [PATCH 156/624] [Perf][Qwen4Exp] Retune H200 M=4 merged QSA LL-GEMM plan (#60776) Signed-off-by: zigzagcai Co-authored-by: zigzagcai Co-authored-by: Claude Code --- vllm/models/qwen4_exp/nvidia/low_latency_gemm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py index acc5b3e6bd98..2ee1c7cb5a54 100644 --- a/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py +++ b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py @@ -79,7 +79,7 @@ (4224, 2560): { 1: SkinnyGemmConfig(1, 128, 2, k_unroll=6, vector_width=2, static_k=2560), 2: SkinnyGemmConfig(2, 128, 2, k_unroll=3, vector_width=2, static_k=2560), - 4: SkinnyGemmConfig(4, 64, 4, k_unroll=2, vector_width=4), + 4: SkinnyGemmConfig(4, 64, 4, k_unroll=1, vector_width=8), }, # GDN and QSA output projections, TP=4. (2560, 1536): { From efaca9547dbdfd2d5269f274aaf3798c1f730435 Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Fri, 9 Oct 2026 00:30:18 -0700 Subject: [PATCH 157/624] [Perf][DSv4.1] Exchange packed FP8 Engram rows across DP and feed wkv MXFP8 (#59923) Signed-off-by: Juntian Liu Co-authored-by: Claude Opus 5.5 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/distributed/test_engram_dp_shard.py | 73 ++++++ tests/kernels/test_engram.py | 13 +- vllm/models/deepseek_v41/common/engram.py | 223 +++++++++++++++++- .../deepseek_v41/common/ops/query_quant.py | 24 +- vllm/models/deepseek_v41/nvidia/model.py | 30 ++- .../deepseek_v41/nvidia/ops/engram_peer.py | 173 ++++++++++++++ 6 files changed, 505 insertions(+), 31 deletions(-) create mode 100644 vllm/models/deepseek_v41/nvidia/ops/engram_peer.py diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py index 637d179a5d21..8b7d3e8cec77 100644 --- a/tests/distributed/test_engram_dp_shard.py +++ b/tests/distributed/test_engram_dp_shard.py @@ -39,9 +39,16 @@ initialize_model_parallel, ) from vllm.forward_context import set_forward_context +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + dequant_mxfp8_to_bf16, +) +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( + convert_swizzled_to_linear, +) from vllm.models.deepseek_v41.common import engram as engram_ops from vllm.models.deepseek_v41.common.engram import ( Engram, + EngramBatch, EngramLayout, ParallelEngramEmbedding, engram_head_shard_rank, @@ -207,10 +214,75 @@ def _worker(rank: int, tp_size: int, port: int) -> None: n_heads, sequence_parallel, ) + if tp_size == 1: + _check_fp8_batch(vllm_config) finally: cleanup_dist_env_and_memory() +def _check_fp8_batch(vllm_config): + """MXFP8 staging matches the tables across peer/all-to-all slots, idle + replicas and graph replays.""" + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture + + dp_rank = vllm_config.parallel_config.data_parallel_rank + layers = [] + for index in range(2): + weight, scale = _full_table(sum(HEAD_SIZES)) + weight = (weight.float() * (1 - 2 * index)).to(torch.float8_e4m3fn) + table = ParallelEngramEmbedding( + sum(HEAD_SIZES), DIM, HEAD_SIZES, cpu_offload=True + ) + table.weight.weight_loader(table.weight, weight) + table.weight_scale_inv.weight_loader(table.weight_scale_inv, scale) + staged = torch.empty(0, device="cuda") + reference = (weight.cuda(), scale.cuda()) + layers.append(SimpleNamespace(embed_tokens=table, staged_rows=staged)) + layers[-1].layer_hash_index, layers[-1].reference = index, reference + batch = EngramBatch(layers, 257) + + def check(ids): + tokens = ids.shape[0] + for layer in layers: + x = batch.wkv_input(layer, tokens) + sf = convert_swizzled_to_linear(x.scale, tokens, x.data.shape[1], BLOCK) + actual = dequant_mxfp8_to_bf16(x.data, sf).unflatten(-1, (-1, DIM)) + weight, scale = layer.reference + expected = _reference(weight, scale, ids[:, layer.layer_hash_index]) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + for counts in ((7, 0, 5, 1), (127,) * 4, (128,) * 4, (257, 0, 129, 1)): + tokens = counts[dp_rank] + seed = 700 + 10 * dp_rank + ids = torch.stack([_make_ids(HEAD_SIZES, tokens, seed + i) for i in (0, 1)], 1) + with set_forward_context( + None, + vllm_config, + num_tokens=tokens, + num_tokens_across_dp=torch.tensor(counts, dtype=torch.int32), + ): + batch.prepare_embeddings(ids) + check(ids) + for capture in ("full", "breakable") if len(set(counts)) == 1 else (): + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + if capture == "full": + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=warmup): + batch.prepare_embeddings(ids) + else: + graph = BreakableCUDAGraphCapture() + with torch.cuda.stream(warmup), graph: + batch.prepare_embeddings(ids) + torch.cuda.current_stream().wait_stream(warmup) + for step in (2, 3): + replay_ids = _make_ids(HEAD_SIZES, tokens, seed + step) + ids.copy_(torch.stack([replay_ids] * 2, 1)) + graph.replay() + check(ids) + del graph + + def _check_shared_storage_lifetime_and_failures(failure): """One peer's failure must roll back every mapping/registration, without hangs.""" import gc @@ -516,6 +588,7 @@ def __call__(self, hidden, positions, input_ids, *args, **kwargs): model.use_native_mega_moe = False model.use_sequence_parallel = False model.fuse_mhc_all_reduce = False + model.engram_batch = None model.engram_hash = state model.engram_swa_prefix = "swa" model.engram_dp_shared_memory = dp_shared_memory diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py index 15bdc9162cf6..3663108d9d1a 100644 --- a/tests/kernels/test_engram.py +++ b/tests/kernels/test_engram.py @@ -7,6 +7,9 @@ import pytest import torch +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + dequant_mxfp8_to_bf16, +) from vllm.models.deepseek_v41.common import engram as engram_ops from vllm.models.deepseek_v41.common.engram import ( Engram, @@ -714,8 +717,9 @@ def test_engram_lookup_reuses_jit_across_token_shapes(): ) @pytest.mark.parametrize("cpu_offload", [False, True]) @pytest.mark.parametrize("background", [False, True]) +@pytest.mark.parametrize("packed", [False, True]) @pytest.mark.parametrize("num_tokens", [1, 7, 256]) -def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens): +def test_engram_lookup_matches_torch(cpu_offload, background, packed, num_tokens): """The fused gather must be bit-exact with the torch dequant path it replaces, from HBM and from pinned host memory alike, and must contribute zeros for rows another TP rank owns.""" @@ -729,8 +733,13 @@ def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens): layer.vocab_start_idx, layer.vocab_end_idx, ) - out = torch.empty(num_tokens, cols, layer.dim, dtype=torch.bfloat16, device="cuda") + width = layer.dim + layer.dim // 32 if packed else layer.dim + dtype = torch.uint8 if packed else torch.bfloat16 + out = torch.empty(num_tokens, cols, width, dtype=dtype, device="cuda") layer.lookup(ids, out, background=background) + if packed: + values = out[..., : layer.dim].view(torch.float8_e4m3fn) + out = dequant_mxfp8_to_bf16(values, out[..., layer.dim :]) assert torch.equal(out, expected) diff --git a/vllm/models/deepseek_v41/common/engram.py b/vllm/models/deepseek_v41/common/engram.py index d4810ee4ea9e..0e160c8b06cf 100644 --- a/vllm/models/deepseek_v41/common/engram.py +++ b/vllm/models/deepseek_v41/common/engram.py @@ -62,9 +62,16 @@ from vllm.distributed.parallel_state import GroupCoordinator from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic from vllm.model_executor.utils import set_weight_attrs +from vllm.models.deepseek_v41.common.ops.query_quant import ( + can_fuse_query_quant, + mxfp8_scale_bytes, + mxfp8_scale_offset, +) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.platform_utils import is_uva_available @@ -72,6 +79,10 @@ logger = init_logger(__name__) +# Per-rank token slot from which the Engram DP exchange switches from peer +# reads to all-to-all. +ENGRAM_A2A_MIN_SLOT = 128 + # Cache value for tokens that take no part in an n-gram (image spans). DEAD_ID = -1 @@ -644,8 +655,10 @@ def _engram_lookup_kernel( BLOCK_R: tl.constexpr, GRID, SORTED: tl.constexpr if _SORTED_IS_CONSTEXPR else None, # type: ignore[valid-type] + PACKED: tl.constexpr, + OUT_STRIDE: tl.constexpr, ): - """Gather fp8 rows, apply their ue8m0 block scales, write bf16. + """Gather FP8 rows as packed values/scales or dequantized BF16. Only this rank's heads are read; padded heads write zeros for all-gather. `weight`/`scales` may address pinned host memory through UVA. If SORTED, @@ -684,13 +697,28 @@ def _engram_lookup_kernel( mask=owned[:, None], other=0, ) - # ue8m0 is a power of two, so its byte *is* the fp32 exponent field. - scale = (scale.to(tl.int32) << 23).to(tl.float32, bitcast=True) - tl.store( - out + dst[:, None] * DIM + cols[None, :], - (values.to(tl.float32) * scale).to(tl.bfloat16), - mask=valid[:, None], - ) + if PACKED: + width: tl.constexpr = DIM + DIM // QUANT_BLOCK + dst = dst.to(tl.int64) + dst = dst // LOCAL_HEADS * OUT_STRIDE + dst % LOCAL_HEADS * width + tl.store( + out + dst[:, None] + cols[None, :], + values.to(tl.uint8, bitcast=True), + mask=valid[:, None], + ) + tl.store( + out + dst[:, None] + DIM + scale_cols[None, :], + scale, + mask=valid[:, None] & (cols[None, :] % QUANT_BLOCK == 0), + ) + else: + # ue8m0 is a power of two, so its byte *is* the fp32 exponent field. + scale = (scale.to(tl.int32) << 23).to(tl.float32, bitcast=True) + tl.store( + out + dst[:, None] * DIM + cols[None, :], + (values.to(tl.float32) * scale).to(tl.bfloat16), + mask=valid[:, None], + ) def engram_head_shard_rank() -> int: @@ -1082,7 +1110,7 @@ def _storage(self) -> tuple[torch.Tensor, torch.Tensor]: def lookup( self, indices: torch.Tensor, out: torch.Tensor, background: bool = False ) -> None: - """Look up local heads of [T, heads] into [T, local_heads, dim] bf16. + """Look up local heads into BF16, or packed FP8 bytes followed by scales. `background` limits the grid to leave SMs for concurrent work. """ @@ -1131,6 +1159,8 @@ def lookup( BLOCK_R=16, GRID=grid, SORTED=sort, + PACKED=out.dtype == torch.uint8, + OUT_STRIDE=out.stride(0), ) def forward(self, indices: torch.Tensor) -> torch.Tensor: @@ -1311,6 +1341,175 @@ def _gather_engram_rows(staged: torch.Tensor, num_tokens: int) -> torch.Tensor: return rows +@triton.jit(do_not_specialize=["num_tokens"]) +def _engram_unpack_fp8_kernel( + packed, + values, + scales, + num_tokens, + values_stride, + scales_stride, + HEADS: tl.constexpr, + LOCAL_HEADS: tl.constexpr, + LAYERS: tl.constexpr, + RANK: tl.constexpr, + DIM: tl.constexpr, + BLOCK: tl.constexpr, + PEER: tl.constexpr, +): + """Split `[rank][token][layer][local head]` packed rows into MXFP8 values and + F8_128x4 scales, one layer per grid column; `packed` holds per-rank + pointers if PEER. Scales past `num_tokens` zero-pad the 128-row tile. + """ + row = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + layer = tl.program_id(1).to(tl.int64) + token, head = row // HEADS, row % HEADS + if PEER: + src = tl.load(packed + head // LOCAL_HEADS) + src = tl.multiple_of(src, DIM // 32).to(tl.pointer_type(tl.uint8)) + source = RANK * num_tokens + token + else: + src = packed + source = head // LOCAL_HEADS * num_tokens + token + source = (source * LAYERS + layer) * LOCAL_HEADS + head % LOCAL_HEADS + # Packed rows are only scale-width aligned. + src += tl.multiple_of(source * (DIM + DIM // 32), DIM // 32) + valid = token < num_tokens + cols = tl.arange(0, DIM) + data = tl.load(src[:, None] + cols[None, :], mask=valid[:, None], other=0) + values += layer * values_stride + row * DIM + tl.store(values[:, None] + cols[None, :], data, mask=valid[:, None]) + scale_cols = tl.arange(0, DIM // 32) + scale = tl.load( + src[:, None] + DIM + scale_cols[None, :], mask=valid[:, None], other=0 + ) + group = head[:, None] * (DIM // 32) + scale_cols[None, :] + offset = mxfp8_scale_offset(token[:, None], group, HEADS * DIM // 32) + tl.store(scales + layer * scales_stride + offset, scale) + + +class EngramBatch(nn.Module): + """Stage every Engram layer's WKV input as MXFP8 with one exchange. + + A module so the model owns the peer exchange's pointer buffers. + """ + + def __init__(self, layers: list["Engram"], max_tokens: int) -> None: + super().__init__() + first = layers[0].layer_hash_index + hash_indices = [layer.layer_hash_index for layer in layers] + assert hash_indices == list(range(first, first + len(layers))) + self.layers = layers + table = layers[0].embed_tokens + device = layers[0].staged_rows.device + self.packed_rows = torch.empty( + max_tokens * table.dp_size, + len(layers), + table.part_n_hash_cols, + table.dim + table.dim // 32, + dtype=torch.uint8, + device=device, + ) + width = table.n_hash_cols * table.dim + self.fp8_values = torch.empty( + len(layers), max_tokens, width, dtype=torch.uint8, device=device + ) + self.fp8_scales = torch.empty( + len(layers), + mxfp8_scale_bytes(max_tokens, width), + dtype=torch.uint8, + device=device, + ) + max_peer_slot = min(max_tokens, ENGRAM_A2A_MIN_SLOT - 1) + self.peer_exchange = None + if table.cpu_offload: + from vllm.models.deepseek_v41.nvidia.ops.engram_peer import ( + EngramPeerExchange, + ) + + self.peer_exchange = EngramPeerExchange.create(layers, max_peer_slot) + for layer in layers: + # Owned by the model; keep it out of each layer's submodules. + object.__setattr__(layer, "_batch", self) + # The layers read this batch's staging instead of their BF16 rows. + layer.staged_rows = layer.staged_rows.new_empty( + (0, *layer.staged_rows.shape[1:]) + ) + + @classmethod + def create(cls, layers: list["Engram"], max_tokens: int) -> "EngramBatch | None": + """Batch DP-sharded tables whose WKV takes FlashInfer MXFP8 input.""" + if not layers: + return None + table = layers[0].embed_tokens + if ( + table.dp_size == 1 + or table.tp_size > 1 + or table.dim % 128 + or get_current_vllm_config().lora_config is not None + or not can_fuse_query_quant([layer.wkv for layer in layers]) + ): + return None + return cls(layers, max_tokens) + + def prepare_embeddings(self, hashes: torch.Tensor) -> None: + """Stage every layer's WKV input before the decoder layers. + + DP-sharded tables decide on the group's slot, so every rank agrees. + """ + slot = engram_gathered_num_tokens() + peer = self.peer_exchange is not None and slot < ENGRAM_A2A_MIN_SLOT + self._stage(hashes if peer else gather_engram_hashes(hashes), slot, peer) + + @eager_break_during_capture + def _stage(self, hashes: torch.Tensor, slot: int, peer: bool) -> None: + if slot == 0: + return + if peer: + assert self.peer_exchange is not None + rows = self.peer_exchange.exchange(hashes, slot) + rank = self.peer_exchange.dp_rank + else: + rows, rank = self.packed_rows[: hashes.shape[0]], 0 + for index, layer in enumerate(self.layers): + layer.embed_tokens.lookup( + hashes[:, layer.layer_hash_index], rows[:, index] + ) + group = get_engram_dp_group() + assert group is not None + sent, rows = rows, torch.empty_like(rows) + torch.distributed.all_to_all_single(rows, sent, group=group.device_group) + table = self.layers[0].embed_tokens + heads = table.n_hash_cols + grid = (triton.cdiv(slot, 128) * 128 * heads // 16, len(self.layers)) + _engram_unpack_fp8_kernel[grid]( + rows, + self.fp8_values, + self.fp8_scales, + slot, + self.fp8_values.stride(0), + self.fp8_scales.stride(0), + HEADS=heads, + LOCAL_HEADS=table.part_n_hash_cols, + LAYERS=len(self.layers), + RANK=rank, + DIM=table.dim, + BLOCK=16, + PEER=peer, + ) + + def wkv_input(self, layer: "Engram", num_tokens: int) -> QuantizedActivation: + index = layer.layer_hash_index - self.layers[0].layer_hash_index + values = self.fp8_values[index, :num_tokens].view(torch.float8_e4m3fn) + return QuantizedActivation( + values, + self.fp8_scales[index, : mxfp8_scale_bytes(num_tokens, values.shape[1])], + torch.bfloat16, + values.shape, + kMxfp8Dynamic, + ) + + class Engram(nn.Module): """Writes an n-gram lookup into the residual stream, gated by how well it matches that stream. @@ -1323,6 +1522,7 @@ class Engram(nn.Module): _prefetch_stream: torch.cuda.Stream | None = None _prefetch_done: torch.cuda.Event | None = None + _batch: EngramBatch | None = None def __init__( self, @@ -1482,7 +1682,10 @@ def forward( """hidden_states: [T, hc_mult, dim]; hash_ids: [T, n_hash_cols] (all tokens, pre sequence-parallel shard); token_mask: [T], False shuts the gate so those positions pass through untouched.""" - kv = self.wkv(self.embed(hash_ids).flatten(-2)) + if self._batch is not None: + kv = self.wkv(self._batch.wkv_input(self, hash_ids.shape[0])) + else: + kv = self.wkv(self.embed(hash_ids).flatten(-2)) num_kv_tokens = hash_ids.shape[0] assert token_mask is None or token_mask.shape == (num_kv_tokens,) if self.use_sequence_parallel: diff --git a/vllm/models/deepseek_v41/common/ops/query_quant.py b/vllm/models/deepseek_v41/common/ops/query_quant.py index 6094d1980cea..97714e29f151 100644 --- a/vllm/models/deepseek_v41/common/ops/query_quant.py +++ b/vllm/models/deepseek_v41/common/ops/query_quant.py @@ -21,6 +21,18 @@ def _norm_row(x, weight, row, stride, eps, SIZE: tl.constexpr, BLOCK: tl.constex return values * rrms * w +@triton.jit +def mxfp8_scale_offset(row, group, GROUPS: tl.constexpr): + """F8_128x4: [row/128, group/4, row%32, row%128/32, group%4].""" + offset = (row // 128 * triton.cdiv(GROUPS, 4) + group // 4) * 512 + return offset + row % 32 * 16 + row % 128 // 32 * 4 + group % 4 + + +def mxfp8_scale_bytes(tokens: int, width: int) -> int: + """Bytes of F8_128x4 scales for `tokens` rows of `width` MXFP8 values.""" + return triton.cdiv(tokens, 128) * 128 * triton.cdiv(width // 32, 4) * 4 + + @triton.jit(do_not_specialize=["num_tokens"]) def _q_kv_norm_quant_kernel( q, @@ -68,14 +80,7 @@ def _q_kv_norm_quant_kernel( sf = tl.full((BLOCK // 32,), 0, tl.uint32) padded_groups: tl.constexpr = triton.cdiv(Q_SIZE // 32, 4) * 4 sf = tl.where(groups < Q_SIZE // 32, sf, 0) - # F8_128x4: [row/128, group/4, row%32, row%128/32, group%4]. - offsets = ( - row // 128 * (128 * padded_groups) - + groups // 4 * 512 - + row % 32 * 16 - + row % 128 // 32 * 4 - + groups % 4 - ) + offsets = mxfp8_scale_offset(row, groups, Q_SIZE // 32) tl.store(scales + offsets, sf, groups < padded_groups) elif row < num_tokens: y = _norm_row(kv, kvw, row, kv_stride, eps, KV_SIZE, BLOCK) @@ -99,9 +104,8 @@ def fused_q_kv_rmsnorm_quant( qo = torch.empty(qr.shape, dtype=torch.float8_e4m3fn, device=qr.device) kvo = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device) padded_tokens = triton.cdiv(tokens, 128) * 128 - padded_groups = triton.cdiv(q_size // 32, 4) * 4 scales = torch.empty( - padded_tokens * padded_groups, dtype=torch.uint8, device=qr.device + mxfp8_scale_bytes(tokens, q_size), dtype=torch.uint8, device=qr.device ) if tokens: block = triton.next_power_of_2(max(q_size, kv_size)) diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 18f8f27843ab..8ec0155f736a 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -100,6 +100,7 @@ from ..common.engram import ( Engram, + EngramBatch, EngramLayout, NgramHashState, can_share_engram_tables, @@ -777,6 +778,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ), prefix=f"{prefix}.layers", ) + self.engram_batch = EngramBatch.create( + [ + layer.engram + for layer in islice(self.layers, self.start_layer, self.end_layer) + if getattr(layer, "engram", None) is not None + ], + vllm_config.scheduler_config.max_num_batched_tokens, + ) if self.fuse_mhc_all_reduce: # A MoE's top-k finalize folds into the next layer's first mHC # boundary, which the last local layer lacks and an engram layer @@ -971,15 +980,18 @@ def forward( if engram_hashes is not None: # Gather all Engram rows before entering the decoder layers. # One gather feeds every layer sharing the DP-split table. - gathered_hashes = gather_engram_hashes( - engram_hashes, dp_shared_memory=self.engram_dp_shared_memory - ) - for layer in islice(self.layers, self.start_layer, self.end_layer): - engram = getattr(layer, "engram", None) - if engram is not None: - engram.prepare_embeddings( - gathered_hashes[:, engram.layer_hash_index] - ) + if self.engram_batch is not None: + self.engram_batch.prepare_embeddings(engram_hashes) + else: + gathered_hashes = gather_engram_hashes( + engram_hashes, dp_shared_memory=self.engram_dp_shared_memory + ) + for layer in islice(self.layers, self.start_layer, self.end_layer): + engram = getattr(layer, "engram", None) + if engram is not None: + engram.prepare_embeddings( + gathered_hashes[:, engram.layer_hash_index] + ) full_num_tokens = positions.shape[0] if self.use_sequence_parallel: diff --git a/vllm/models/deepseek_v41/nvidia/ops/engram_peer.py b/vllm/models/deepseek_v41/nvidia/ops/engram_peer.py new file mode 100644 index 000000000000..4f9423bb1e8b --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/ops/engram_peer.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Symmetric-memory Engram exchange for slots below ENGRAM_A2A_MIN_SLOT.""" + +from typing import TYPE_CHECKING + +import torch +import torch.distributed as dist +from torch import nn + +from vllm.distributed import get_engram_dp_group +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + +if TYPE_CHECKING: + from vllm.models.deepseek_v41.common.engram import Engram + + +@triton.jit(do_not_specialize=["num_tokens"]) +def _peer_lookup_kernel( + metadata, + id_ptrs, + output, + num_tokens, + HEADS: tl.constexpr, + LAYERS: tl.constexpr, + DP: tl.constexpr, + RANK: tl.constexpr, + DIM: tl.constexpr, + BLOCK: tl.constexpr, +): + """Look up this rank's heads for every peer's tokens into packed rows.""" + row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + local_heads: tl.constexpr = HEADS // DP + valid = row < num_tokens * LAYERS * HEADS + token = row // (LAYERS * local_heads) + layer, head = row // local_heads % LAYERS, row % local_heads + ids = tl.load(id_ptrs + token // num_tokens, valid, other=0).to( + tl.pointer_type(tl.int32) + ) + index = tl.load( + ids + (token % num_tokens * LAYERS + layer) * HEADS + RANK * local_heads + head, + valid, + other=-1, + ).to(tl.int64) + # Per-layer metadata: weight pointer, scale pointer, vocab start, vocab end. + table = metadata + layer * 4 + start, end = tl.load(table + 2), tl.load(table + 3) + owned = valid & (index >= start) & (index < end) + local = tl.where(owned, index - start, 0) + weight = tl.multiple_of(tl.load(table), DIM).to(tl.pointer_type(tl.uint8)) + scales = tl.load(table + 1) + scales = tl.multiple_of(scales, DIM // 32).to(tl.pointer_type(tl.uint8)) + col = tl.arange(0, DIM) + value = tl.load( + weight[:, None] + local[:, None] * DIM + col[None, :], owned[:, None], other=0 + ) + dst = row.to(tl.int64) * (DIM + DIM // 32) + tl.store(output + dst[:, None] + col[None, :], value, valid[:, None]) + scale_cols = tl.arange(0, DIM // 32) + scale = tl.load( + scales[:, None] + local[:, None] * (DIM // 32) + scale_cols[None, :], + owned[:, None], + other=0, + ) + tl.store(output + dst[:, None] + DIM + scale_cols[None, :], scale, valid[:, None]) + + +class EngramPeerExchange(nn.Module): + """Two phase barriers fence reuse of each other's buffers across forwards. + + Pointer tables are buffers so level-2 sleep restores them on wake-up. + """ + + def __init__(self, layers: list["Engram"], max_tokens: int): + import torch.distributed._symmetric_memory as symm_mem + + super().__init__() + group = get_engram_dp_group() + assert group is not None + self.dp_size, self.dp_rank = group.world_size, group.rank_in_group + device = layers[0].staged_rows.device + self.num_layers = len(layers) + self.hash_start = layers[0].layer_hash_index + table = layers[0].embed_tokens + self.heads, self.dim = table.n_hash_cols, table.dim + self.ids = symm_mem.empty( + (max_tokens, self.num_layers, self.heads), + dtype=torch.int32, + device=device, + ) + self.rows = symm_mem.empty( + ( + max_tokens * self.dp_size, + self.num_layers, + self.heads // self.dp_size, + self.dim + self.dim // 32, + ), + dtype=torch.uint8, + device=device, + ) + self.id_handle = symm_mem.rendezvous(self.ids, group.device_group) + self.row_handle = symm_mem.rendezvous(self.rows, group.device_group) + metadata = [] + for layer in layers: + table = layer.embed_tokens + weight, scale = table._storage() + metadata.append( + ( + weight.data_ptr(), + scale.data_ptr(), + table.vocab_start_idx, + table.vocab_end_idx, + ) + ) + for name, data in ( + ("metadata", metadata), + ("id_ptrs", self.id_handle.buffer_ptrs), + ("row_ptrs", self.row_handle.buffer_ptrs), + ): + self.register_buffer( + name, + torch.tensor(data, dtype=torch.int64, device=device), + persistent=False, + ) + + @classmethod + def create( + cls, layers: list["Engram"], max_tokens: int + ) -> "EngramPeerExchange | None": + # The Engram DP group is node-local and every rank agrees on these. + group = get_engram_dp_group() + assert group is not None + if not ( + current_platform.is_device_capability_family(100) + and group.world_size in (2, 4, 8) + and layers[0].embed_tokens.n_hash_cols % group.world_size == 0 + ): + return None + physical_id = current_platform.visible_device_id_to_physical_device_id( + torch.accelerator.current_device_index() + ) + physical_ids = [None] * group.world_size + dist.all_gather_object(physical_ids, physical_id, group=group.cpu_group) + if not current_platform.is_fully_connected(physical_ids): + return None + return cls(layers, max_tokens) + + def exchange(self, hashes: torch.Tensor, slot: int) -> torch.Tensor: + """Look up this rank's heads for every peer; return per-rank row pointers.""" + tokens = hashes.shape[0] + assert 0 < slot <= self.ids.shape[0] and tokens <= slot + hash_end = self.hash_start + self.num_layers + self.ids[:tokens].copy_(hashes[:, self.hash_start : hash_end]) + if tokens < slot: + self.ids[tokens:slot].fill_(-1) + # The preceding rows barrier completed every peer's hash reads. This + # barrier completes their previous output reads before rows are reused. + self.id_handle.barrier() + _peer_lookup_kernel[(triton.cdiv(slot * self.num_layers * self.heads, 16),)]( + self.metadata, + self.id_ptrs, + self.rows, + slot, + HEADS=self.heads, + LAYERS=self.num_layers, + DP=self.dp_size, + RANK=self.dp_rank, + DIM=self.dim, + BLOCK=16, + ) + self.row_handle.barrier() + return self.row_ptrs From f43e4c39e30120abf6250a209a42e5cd4849ca23 Mon Sep 17 00:00:00 2001 From: Denis Ermakov Date: Fri, 9 Oct 2026 09:49:47 +0200 Subject: [PATCH 158/624] [Structured Output] Allow propertyNames with additionalProperties on xgrammar (#57594) Signed-off-by: errmakov Signed-off-by: Artem Perevedentsev Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Artem Perevedentsev --- .../test_backend_xgrammar.py | 136 ++++++++++++------ vllm/v1/structured_output/backend_xgrammar.py | 101 +++++++------ 2 files changed, 152 insertions(+), 85 deletions(-) diff --git a/tests/v1/structured_output/test_backend_xgrammar.py b/tests/v1/structured_output/test_backend_xgrammar.py index bc29112077cc..9024b614dd77 100644 --- a/tests/v1/structured_output/test_backend_xgrammar.py +++ b/tests/v1/structured_output/test_backend_xgrammar.py @@ -31,21 +31,6 @@ def grammar_accepts(schema: dict, text: str) -> bool: # ================================================ -@pytest.fixture -def unsupported_string_schemas(): - return [ - {"type": "string", "format": "non_existing_format"}, - # TODO(arpera): - # pattern/format is compiled but length bounds are silently dropped, - # so the combination must be rejected instead of producing quietly - # wrong output - # https://github.com/mlc-ai/xgrammar/issues/749 - {"type": "string", "pattern": "^a+$", "maxLength": 2}, - {"type": "string", "pattern": "^a+$", "minLength": 3}, - {"type": "string", "format": "email", "maxLength": 10}, - ] - - @pytest.fixture def unsupported_multipleOf_schemas(): return [ @@ -59,6 +44,8 @@ def unsupported_multipleOf_schemas(): @pytest.fixture def unsupported_array_schemas(): return [ + # array + some constraints is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/968 {"type": "array", "uniqueItems": True}, {"type": "array", "contains": {"type": "string"}}, {"type": "array", "minContains": 1}, @@ -67,74 +54,113 @@ def unsupported_array_schemas(): @pytest.fixture -def unsupported_property_names_combinations(): +def unsupported_string_schemas(): + return [ + # ======================================================== + # string + format is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/967 + {"type": "string", "format": "non_existing_format"}, + # ======================================================== + # + # ======================================================== + # string + format/pattern + length constraint is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/966 + # + # pattern/format is compiled but length bounds are silently dropped, + # so the combination must be rejected instead of producing quietly + # wrong output + {"type": "string", "pattern": "^a+$", "maxLength": 2}, + {"type": "string", "pattern": "^a+$", "minLength": 3}, + {"type": "string", "format": "email", "maxLength": 10}, + # ======================================================== + ] + + +@pytest.fixture +def unsupported_propertyNames_combinations(): return [ - # TODO(arpera): - # this case is not covered by xgrammar, so we should report this bug to xgrammar + # ======================================================== + # propertyNames + patternProperties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/959 + # # xgrammar drops propertyNames whenever patternProperties is present, # so "grade_12" matches the pattern and escapes the name constraint { "type": "object", - "patternProperties": {"^grade_[0-9]+$": {"type": "integer"}}, "propertyNames": {"pattern": "^grade_[0-9]$"}, # does NOT match "grade_12" + "patternProperties": {"^grade_[0-9]+$": {"type": "integer"}}, }, - # propertyNames makes xgrammar discard the sibling additionalProperties - # value schema: https://github.com/mlc-ai/xgrammar/issues/826 - { - "type": "object", - "propertyNames": {"pattern": "^[a-z]+$"}, - "additionalProperties": {"type": "integer"}, - }, - # propertyNames is a string schema that conventionally omits "type", so - # it escapes the string check while xgrammar drops its length bound all - # the same: https://github.com/mlc-ai/xgrammar/issues/749 + # ======================================================== + # + # ======================================================== + # propertyNames + maxLength is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/960 { "type": "object", "propertyNames": {"pattern": "^a+$", "maxLength": 2}, }, - # xgrammar emits named `properties` separately from `propertyNames` and - # only applies `propertyNames` in its additional-properties branch, so a - # key declared in `properties` escapes the name constraint regardless of - # what `propertyNames` contains. + # ======================================================== + # + # ======================================================== + # propertyNames + properties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/958 { "type": "object", + "propertyNames": {"pattern": "^[a-z]+$"}, "properties": {"Bad": {"type": "integer"}}, "required": ["Bad"], - "propertyNames": {"pattern": "^[a-z]+$"}, }, { "type": "object", - "properties": {"Bad": {"type": "integer"}}, "propertyNames": {"enum": ["good"]}, + "properties": {"Bad": {"type": "integer"}}, }, + # ======================================================== + # + # ======================================================== + # propertyNames + unevaluatedProperties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/961 + # # xgrammar's propertyNames-only branch falls back to an unconstrained # value type, so it drops unevaluatedProperties whether it restricts # values (a schema) or forbids them (false). { "type": "object", "propertyNames": {"pattern": "^[a-z]+$"}, - "unevaluatedProperties": {"type": "integer"}, + "unevaluatedProperties": False, }, + # also keep in mind that unevaluatedProperties can be a schema { "type": "object", "propertyNames": {"pattern": "^[a-z]+$"}, - "unevaluatedProperties": False, + "unevaluatedProperties": {"type": "integer"}, }, + # ======================================================== ] @pytest.fixture -def unsupported_pattern_properties_combinations(): +def unsupported_patternProperties_combinations(): return [ + # ======================================================== + # patternProperties + properties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/964 + # # JSON Schema requires a property matched by both `properties` and # `patternProperties` to satisfy both (conjunction), but xgrammar # compiles them as alternative branches, so satisfying either one is # enough. { "type": "object", - "properties": {"x": {"type": "string"}}, "patternProperties": {"^x$": {"type": "integer"}}, + "properties": {"x": {"type": "string"}}, }, + # ======================================================== + # + # ======================================================== + # patternProperties + patternProperties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/965 + # # The same alternative-branches problem applies to overlapping # patternProperties patterns: a key matching both patterns only has to # satisfy one of them. @@ -145,6 +171,7 @@ def unsupported_pattern_properties_combinations(): "^[a-z]*z$": {"type": "integer"}, }, }, + # ======================================================== ] @@ -273,6 +300,24 @@ def property_names_schema(): return {"type": "object", "propertyNames": {"pattern": "^[a-z_]+$"}} +@pytest.fixture +def property_names_with_additional_properties_schema(): + """This combination propertyNames + additionalProperties + was unsupported some time ago by xgrammar. + So, this test case was in unsupported category before. + But this xgrammar's PR + https://github.com/mlc-ai/xgrammar/pull/836 + added support of propertyNames + additionalProperties + This fix was released in xgrammar v0.2.7 + which vLLM pins in PR #57272 + """ + return { + "type": "object", + "propertyNames": {"pattern": "^[a-z_]+$"}, + "additionalProperties": {"type": "integer"}, + } + + @pytest.fixture def supported_allof_anyof_and_oneof(): """Xgrammar DO support some variants of anyOf, allOf, oneOf @@ -334,8 +379,8 @@ class TestHasXGrammarUnsupportedJsonFeatures: "unsupported_string_schemas", "unsupported_multipleOf_schemas", "unsupported_array_schemas", - "unsupported_property_names_combinations", - "unsupported_pattern_properties_combinations", + "unsupported_propertyNames_combinations", + "unsupported_patternProperties_combinations", # vLLM issue #56556 "unsupported_multibranch_allof", "unsupported_vllm_issue_56556_schema", @@ -354,6 +399,7 @@ def test_unsupported_json_features(self, schema_type, request): "supported_frankenstein_schema", "pattern_properties_schema", "property_names_schema", + "property_names_with_additional_properties_schema", # Additional test cases implemented during work on # vLLM issue #56556 "supported_allof_anyof_and_oneof", @@ -421,6 +467,14 @@ def test_property_names_constrains_keys(self, property_names_schema): assert not grammar_accepts(property_names_schema, '{"Score": 5}') assert not grammar_accepts(property_names_schema, '{"score_1": 5}') + def test_property_names_keeps_additional_properties_value_schema( + self, property_names_with_additional_properties_schema + ): + schema = property_names_with_additional_properties_schema + assert grammar_accepts(schema, '{"score": 5}') + assert not grammar_accepts(schema, '{"Score": 5}') + assert not grammar_accepts(schema, '{"score": "seven"}') + def test_pattern_properties_constrains_keys_and_values( self, pattern_properties_schema ): diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index b675a9a3f193..0199c5f75d6f 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -296,18 +296,29 @@ def check_object(obj: dict[str, Any]) -> bool: schema_types = _schema_types(obj) - # Check for numeric ranges + # integer/number + multipleOf is unsupported by xgrammar + # This is known behavior and xgrammar emits warning in logs: + # [21:18:08] /project/cpp/json_schema_converter.cc:1053: + # Warning: multipleOf is not supported for type:number; ignoring multipleOf + # This warning was added in PR + # https://github.com/mlc-ai/xgrammar/pull/670 + # So, no need to track progress on this if (schema_types & {"integer", "number"}) and ("multipleOf" in obj): return True - # Check for array unsupported keywords + # array + some constraints is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/968 if "array" in schema_types and any( key in obj for key in ("uniqueItems", "contains", "minContains", "maxContains") ): return True - # Unsupported keywords for strings + # string + format with unsupported keywords + # is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/967 + # See tests on this in test_backend_xgrammar.py + # unsupported_string_schemas if ( "string" in schema_types and "format" in obj @@ -315,48 +326,51 @@ def check_object(obj: dict[str, Any]) -> bool: ): return True - # A string mixing a generative constraint (pattern or format) with - # explicit length bounds. xgrammar compiles the pattern/format side - # and silently drops minLength/maxLength from the grammar, so output - # can violate the bound without any error surfacing. Verified against - # the compiled EBNF: pattern/format grammars come out byte-identical - # with and without the length keywords, while maxLength alone lowers - # to {0, N} correctly. + # string + format/pattern + length constraint is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/966 + # See tests on this in test_backend_xgrammar.py + # unsupported_string_schemas if "string" in schema_types and _has_pattern_and_length_bounds(obj): return True - # propertyNames validates names, so it is a string schema even when it - # omits "type", which is the form that escapes the check above. - if ( - "object" in schema_types - and isinstance(obj.get("propertyNames"), dict) - and _has_pattern_and_length_bounds(obj["propertyNames"]) - ): - return True - - # FIXME: propertyNames conflicts with properties/patternProperties/ - # additionalProperties/unevaluatedProperties under xgrammar. - # https://github.com/mlc-ai/xgrammar/issues/826 - if ( - "object" in schema_types - and "propertyNames" in obj - and ( - "properties" in obj - or "patternProperties" in obj - or isinstance(obj.get("additionalProperties"), dict) - or obj.get("unevaluatedProperties", True) is not True - ) - ): - return True - - # FIXME: multiple patternProperties, or patternProperties alongside - # properties, conflict under xgrammar. - if ( - "object" in schema_types - and isinstance(obj.get("patternProperties"), dict) - and ("properties" in obj or len(obj["patternProperties"]) > 1) - ): - return True + # propertyNames is not supported in pair with some constraints + # in xgrammar + # See tests on this in test_backend_xgrammar.py + # unsupported_propertyNames_combinations + if "object" in schema_types and "propertyNames" in obj: + # propertyNames + maxLength is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/960 + property_names = obj.get("propertyNames") + if isinstance(property_names, dict) and _has_pattern_and_length_bounds( + property_names + ): + return True + # propertyNames + patternProperties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/959 + if "patternProperties" in obj: + return True + # propertyNames + properties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/958 + if "properties" in obj: + return True + # propertyNames + unevaluatedProperties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/961 + if obj.get("unevaluatedProperties", True) is not True: + return True + + # patternProperties is not supported in pair with some constraints + # in xgrammar + # See tests on this in test_backend_xgrammar.py + # unsupported_patternProperties_combinations + if "object" in schema_types and isinstance(obj.get("patternProperties"), dict): + # patternProperties + properties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/964 + if "properties" in obj: + return True + # patternProperties + patternProperties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/965 + if len(obj["patternProperties"]) > 1: + return True # Note(arpera): # Xgrammar lacks support of multi-branch allOf @@ -370,8 +384,7 @@ def check_object(obj: dict[str, Any]) -> bool: # would accept any kind of json, such as # "maybe", "", 42, {}, [], {"a": 1}, etc. # which is NOT what is expected. - # Reported this issue to xgrammar team to track progress on resolving: - # https://github.com/mlc-ai/xgrammar/issues/937 + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/937 allof = obj.get("allOf") if isinstance(allof, list) and len(allof) >= 2: return True From 8a03ccf7b3a26dc0c6ccf2a3ba12113475824db3 Mon Sep 17 00:00:00 2001 From: Maroon Ayoub Date: Fri, 9 Oct 2026 11:14:13 +0300 Subject: [PATCH 159/624] [Frontend] Add return_mm_kwargs to render requests (#60728) Signed-off-by: Maroon Ayoub Co-authored-by: Claude --- docs/serving/online_serving/renderer.md | 7 +++++++ .../token_in_token_out/test_mm_serde.py | 14 ++++++++++++++ vllm/entrypoints/anthropic/protocol.py | 11 +++++++++++ .../openai/chat_completion/protocol.py | 12 ++++++++++++ vllm/entrypoints/openai/completion/protocol.py | 12 ++++++++++++ vllm/entrypoints/openai/responses/protocol.py | 11 +++++++++++ vllm/entrypoints/scale_out/render/serving.py | 16 +++++++++++++--- .../scale_out/token_in_token_out/mm_features.py | 16 ++++++++++------ 8 files changed, 90 insertions(+), 9 deletions(-) diff --git a/docs/serving/online_serving/renderer.md b/docs/serving/online_serving/renderer.md index 79d0ade85b4f..93ecf34f4f36 100644 --- a/docs/serving/online_serving/renderer.md +++ b/docs/serving/online_serving/renderer.md @@ -153,6 +153,13 @@ should split those fields: - Legacy clients that ignore `mm_metadata` and keep sending `kwargs_data` continue to work. +Callers that need only the token layout and item hashes, such as cache-aware +routers, can set `"return_mm_kwargs": false` on the render request. The response +then keeps `mm_hashes` and `mm_placeholders` and sets `kwargs_data` and +`mm_metadata` to null, so the processed tensors are neither serialized nor sent. +Do not forward such a response to `/inference/v1/generate`, which reads a null +`kwargs_data` as every item being cached. + ## Example The example below shows how a disaggregated encode / prefill coordinator can diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py b/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py index 5cd60a959fb3..62223f1d7cf2 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py @@ -168,6 +168,20 @@ def test_render_extracts_metadata_fields_separately(): assert torch.equal(full_data["pixel_values"].data, pixel_values.data) +def test_extract_without_kwargs_keeps_layout(): + """Layout-only callers get hashes and placeholders, no serialized tensors.""" + engine_input, _pixel_values, _image_grid_thw = _image_engine_input() + + full = extract_mm_features(engine_input) + layout = extract_mm_features(engine_input, include_mm_kwargs=False) + + assert full is not None and layout is not None + assert layout.kwargs_data is None + assert layout.mm_metadata is None + assert layout.mm_hashes == full.mm_hashes + assert layout.mm_placeholders == full.mm_placeholders + + def test_extract_includes_declared_placeholder_metadata_fields(): """EC placeholder metadata is kept even when keep_on_cpu is unset.""" engine_input, _pixel_values, image_grid_thw = _image_engine_input( diff --git a/vllm/entrypoints/anthropic/protocol.py b/vllm/entrypoints/anthropic/protocol.py index 5997d2e3d2e6..b2d29e0b34fe 100644 --- a/vllm/entrypoints/anthropic/protocol.py +++ b/vllm/entrypoints/anthropic/protocol.py @@ -229,6 +229,17 @@ class AnthropicMessagesRequest(BaseModel): top_p: float | None = None # vLLM-specific fields that are not in Anthropic spec + return_mm_kwargs: bool = Field( + default=True, + description=( + "If false, the render response's `features` set `kwargs_data` " + "and `mm_metadata` to null, for callers that need only the token " + "layout and item hashes, such as cache-aware routers. Do not send " + "such a response to `/inference/v1/generate`, which reads a null " + "`kwargs_data` as every item being cached. Only supported on the " + "render endpoints; ignored on regular generation endpoints." + ), + ) cache_salt: str | None = Field( default=None, min_length=1, diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 15905b8a3da6..dee8f283cf68 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -458,6 +458,18 @@ class ChatCompletionRequest(OpenAIBaseModel): ), ) + return_mm_kwargs: bool = Field( + default=True, + description=( + "If false, the render response's `features` set `kwargs_data` " + "and `mm_metadata` to null, for callers that need only the token " + "layout and item hashes, such as cache-aware routers. Do not send " + "such a response to `/inference/v1/generate`, which reads a null " + "`kwargs_data` as every item being cached. Only supported on the " + "render endpoints; ignored on regular generation endpoints." + ), + ) + cache_salt: str | None = Field( default=None, min_length=1, diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 78722b45ec7b..f26dea2e1f9e 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -197,6 +197,18 @@ class CompletionRequest(OpenAIBaseModel): ), ) + return_mm_kwargs: bool = Field( + default=True, + description=( + "If false, the render response's `features` set `kwargs_data` " + "and `mm_metadata` to null, for callers that need only the token " + "layout and item hashes, such as cache-aware routers. Do not send " + "such a response to `/inference/v1/generate`, which reads a null " + "`kwargs_data` as every item being cached. Only supported on the " + "render endpoints; ignored on regular generation endpoints." + ), + ) + cache_salt: str | None = Field( default=None, min_length=1, diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index d1a4ea4f6c21..406ab2d0818f 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -330,6 +330,17 @@ class ResponsesRequest(OpenAIBaseModel): "if the served model does not use priority scheduling." ), ) + return_mm_kwargs: bool = Field( + default=True, + description=( + "If false, the render response's `features` set `kwargs_data` " + "and `mm_metadata` to null, for callers that need only the token " + "layout and item hashes, such as cache-aware routers. Do not send " + "such a response to `/inference/v1/generate`, which reads a null " + "`kwargs_data` as every item being cached. Only supported on the " + "render endpoints; ignored on regular generation endpoints." + ), + ) cache_salt: str | None = Field( default=None, min_length=1, diff --git a/vllm/entrypoints/scale_out/render/serving.py b/vllm/entrypoints/scale_out/render/serving.py index bddbe90f286f..84771c5f7390 100644 --- a/vllm/entrypoints/scale_out/render/serving.py +++ b/vllm/entrypoints/scale_out/render/serving.py @@ -144,7 +144,9 @@ async def render_chat_request( return GenerateRequest( request_id=request_id, token_ids=token_ids, - features=self._extract_mm_features(engine_input), + features=self._extract_mm_features( + engine_input, include_mm_kwargs=request.return_mm_kwargs + ), sampling_params=params, model=request.model, reasoning_ended=reasoning_ended, @@ -169,6 +171,7 @@ async def render_messages_request( chat_req = AnthropicServingMessages.to_chat_completion_request( request, merge_inline_system=self._merge_inline_system ) + chat_req.return_mm_kwargs = request.return_mm_kwargs return await self.render_chat_request(chat_req) async def render_completion_request( @@ -217,7 +220,9 @@ async def render_completion_request( GenerateRequest( request_id=request_id, token_ids=token_ids, - features=self._extract_mm_features(engine_input), + features=self._extract_mm_features( + engine_input, include_mm_kwargs=request.return_mm_kwargs + ), sampling_params=params, model=request.model, stream=bool(request.stream), @@ -277,7 +282,9 @@ async def render_responses_request( return GenerateRequest( request_id=request.request_id, token_ids=list(token_ids), - features=self._extract_mm_features(engine_input), + features=self._extract_mm_features( + engine_input, include_mm_kwargs=request.return_mm_kwargs + ), sampling_params=params, model=request.model, reasoning_parser_kwargs=self._reasoning_parser_kwargs(request), @@ -328,8 +335,11 @@ def _placeholder_metadata_fields(self, modality: str) -> set[str]: def _extract_mm_features( self, engine_input: EngineInput, + *, + include_mm_kwargs: bool = True, ) -> MultiModalFeatures | None: return extract_mm_features( engine_input, metadata_fields_for=self._placeholder_metadata_fields, + include_mm_kwargs=include_mm_kwargs, ) diff --git a/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py b/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py index 62d006783a2e..a0cb0d4b5eb6 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py @@ -139,13 +139,15 @@ def extract_mm_features( engine_input: EngineInput, *, metadata_fields_for: Callable[[str], Collection[str]] | None = None, + include_mm_kwargs: bool = True, ) -> MultiModalFeatures | None: """Extract multimodal features from a rendered engine prompt. Returns ``None`` for text-only prompts. ``mm_metadata`` keeps the intersection of processed kwargs that prefill needs after EC transfer: fields declared as embedding metadata, plus fields marked - ``keep_on_cpu`` (for example M-RoPE grid dims). + ``keep_on_cpu`` (for example M-RoPE grid dims). With + ``include_mm_kwargs=False`` neither is serialized. """ if engine_input.get("type") != "multimodal": return None @@ -157,11 +159,13 @@ def extract_mm_features( kwargs_data: dict[str, list[str | None]] | None = None mm_metadata: dict[str, list[str | None]] | None = None - if raw_mm_kwargs := mm_engine_input.get("mm_kwargs"): - kwargs_data, mm_metadata = _encode_mm_kwargs_with_metadata( - raw_mm_kwargs, - metadata_fields_for=metadata_fields_for, - ) + if include_mm_kwargs: + raw_mm_kwargs = mm_engine_input.get("mm_kwargs") + if raw_mm_kwargs: + kwargs_data, mm_metadata = _encode_mm_kwargs_with_metadata( + raw_mm_kwargs, + metadata_fields_for=metadata_fields_for, + ) return MultiModalFeatures( mm_hashes=mm_hashes, From 560fae64cdaf84468590fd362916969f18bb4d9b Mon Sep 17 00:00:00 2001 From: Sriram Kumar Date: Fri, 9 Oct 2026 11:16:20 +0300 Subject: [PATCH 160/624] [ROCm][Bugfix] Pass per-sequence context lengths to the AITER paged MQA-logits kernels (#59420) Signed-off-by: Sriram Kumar --- .../test_rocm_paged_mqa_context_lens.py | 32 +++++++++++++++++++ .../v1/attention/ops/rocm_aiter_mla_sparse.py | 16 +++++++--- 2 files changed, 44 insertions(+), 4 deletions(-) create mode 100644 tests/v1/attention/test_rocm_paged_mqa_context_lens.py diff --git a/tests/v1/attention/test_rocm_paged_mqa_context_lens.py b/tests/v1/attention/test_rocm_paged_mqa_context_lens.py new file mode 100644 index 000000000000..3122cc915a93 --- /dev/null +++ b/tests/v1/attention/test_rocm_paged_mqa_context_lens.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm._aiter_ops import is_aiter_found_and_supported +from vllm.platforms import current_platform + +if not current_platform.is_rocm() or not is_aiter_found_and_supported(): + pytest.skip("Requires ROCm with AITER.", allow_module_level=True) + +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _per_sequence_context_lens + + +def test_per_row_table_uses_last_column(): + seq_lens = torch.tensor([8192, 1500, 3000, 700], dtype=torch.int32) + next_n = 3 + offsets = torch.arange(next_n, dtype=torch.int32) - next_n + 1 + per_row = seq_lens.unsqueeze(1) + offsets + + out = _per_sequence_context_lens(per_row) + + assert out.shape == (4,) + assert out.is_contiguous() + assert torch.equal(out, seq_lens) + + +@pytest.mark.parametrize("shape", [(4,), (4, 1)]) +def test_per_sequence_input_is_unchanged(shape): + context_lens = torch.arange(1, 5, dtype=torch.int32).reshape(shape) + assert _per_sequence_context_lens(context_lens) is context_lens diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 90a73c101c0d..6a93d2151127 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -707,6 +707,14 @@ def paged_mqa_logits_module(): return None +def _per_sequence_context_lens(context_lens: torch.Tensor) -> torch.Tensor: + # AITER takes one length per sequence: the last column of a (B, next_n) table. + # TODO: remove once AITER accepts (B, next_n) context_lens (ROCm/aiter#6153). + if context_lens.dim() == 2 and context_lens.shape[1] > 1: + return context_lens[:, -1].contiguous() + return context_lens + + def rocm_fp8_paged_mqa_logits( q_fp8: torch.Tensor, kv_cache_fp8: torch.Tensor, @@ -726,8 +734,8 @@ def rocm_fp8_paged_mqa_logits( kv_cache_fp8: Paged KV-cache in packed FP8+scale layout with shape [num_blocks, block_size, 1, D+4], dtype `torch.uint8`. weights: Tensor of shape [B * next_n, H], dtype `torch.float32`. - context_lens: Tensor of shape [B], dtype int32; effective context length - for each batch element. + context_lens: Tensor of shape [B] or [B, next_n], dtype int32; effective + context length for each batch element, or for each of its Q rows. block_tables: Tensor of shape [B, max_blocks], dtype int32; maps logical block indices to physical blocks in the paged cache. schedule_metadata: Returned by `get_paged_mqa_logits_metadata`; @@ -779,7 +787,7 @@ def rocm_fp8_paged_mqa_logits( kv_cache_fp8, weights, out_logits, - context_lens, + _per_sequence_context_lens(context_lens), block_tables, max_model_len, ChunkK=256, @@ -801,7 +809,7 @@ def rocm_fp8_paged_mqa_logits( kv_cache_fp8, weights, out_qk, - context_lens, + _per_sequence_context_lens(context_lens), block_tables, max_model_len, ChunkQ=heads, From bacbbe187885db62859f5a4a1443f8ec3006e987 Mon Sep 17 00:00:00 2001 From: The Anh Nguyen Date: Fri, 9 Oct 2026 15:22:18 +0700 Subject: [PATCH 161/624] [Bugfix][NIXL] Guard record_transfer against None telemetry (#57734) Signed-off-by: The Anh Nguyen --- vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py index 26c4275efbe1..645525bc6e2b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py @@ -44,8 +44,10 @@ def reset(self): "num_notifications_after_expiry": [], } - def record_transfer(self, res: "nixlXferTelemetry"): + def record_transfer(self, res: "nixlXferTelemetry | None"): # Keep metrics units consistent with rest of the code: time us->s + if res is None: + return self.data["transfer_duration"].append(res.xferDuration / 1e6) self.data["post_duration"].append(res.postDuration / 1e6) self.data["bytes_transferred"].append(res.totalBytes) From a6f43e57d6bdcbe1a42cc07fd5298ae33052180d Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Fri, 9 Oct 2026 16:29:55 +0800 Subject: [PATCH 162/624] [ROCm] Enable the cuMem CUDA-graph pool offload on ROCm (#59523) Signed-off-by: aoshen02 Signed-off-by: indianspeedster Co-authored-by: Claude Fable 5.1 Co-authored-by: indianspeedster --- csrc/cumem_allocator.cpp | 22 ------- docs/features/sleep_mode.md | 2 +- .../memory/cumem/test_cumem.py | 11 ++-- .../memory/sleep_mode/test_sleep_mode.py | 62 +++++++++++++++++-- vllm/_aiter_ops.py | 6 +- vllm/config/vllm.py | 2 +- vllm/device_allocator/cumem.py | 8 ++- .../aiter_custom_all_reduce.py | 7 ++- .../device_communicators/cuda_communicator.py | 2 + 9 files changed, 84 insertions(+), 38 deletions(-) diff --git a/csrc/cumem_allocator.cpp b/csrc/cumem_allocator.cpp index 6965d2124a65..efb1f879d2ae 100644 --- a/csrc/cumem_allocator.cpp +++ b/csrc/cumem_allocator.cpp @@ -226,28 +226,6 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem, return; } -#ifdef USE_ROCM - // ROCm-only: zero the freshly (re)mapped region. - // - // On some amdgpu driver versions the physical VRAM backing a buffer is not - // scrubbed on release, so the pages handed back here (e.g. on wake_up after a - // sleep that discarded a tag) can still contain stale data from a previously - // freed allocation -- possibly from another process. Callers of the device - // allocator expect newly mapped memory to be well defined; in particular - // hybrid Mamba / gated-delta-net (GDN) models keep a persisted conv/recurrent - // state cache, and a new sequence that reads an unzeroed slot will propagate - // garbage/NaN, collapsing generation to token id 0 (decoded as "!"). - // - // This is scoped to the ROCm path on purpose: the CUDA path is unchanged. - // Use CUDA_CHECK so a failure sets the global error_code; the Python caller - // (python_create_and_map) only inspects error_code, so returning without - // setting it would make a failed memset look successful and still hand back - // unzeroed memory. - CUDA_CHECK((CUresult)hipMemset(reinterpret_cast(d_mem), 0, size)); - if (error_code != 0) { - return; - } -#endif // std::cout << "create_and_map: device=" << device << ", size=" << size << ", // d_mem=" << d_mem << ", p_memHandle=" << p_memHandle << std::endl; } diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index a8aa15af7b52..0db82de9d3c5 100644 --- a/docs/features/sleep_mode.md +++ b/docs/features/sleep_mode.md @@ -107,7 +107,7 @@ By default, CUDA graph memory stays on the GPU while asleep. With `sleep_mode_offload_cudagraph=True` (off by default), CUDA graphs are captured into a cuMem pool that sleep backs up to CPU memory at both levels and any wake restores in place, so graphs are replayed, not recaptured. It needs the default -`cumem` backend, CUDA and CUDA graphs; otherwise it has no effect. +`cumem` backend, CUDA or ROCm, and CUDA graphs; otherwise it has no effect. ```python llm = LLM("Qwen/Qwen3-8B", enable_sleep_mode=True, sleep_mode_offload_cudagraph=True) diff --git a/tests/basic_correctness/memory/cumem/test_cumem.py b/tests/basic_correctness/memory/cumem/test_cumem.py index df8c6048824c..c7ebfd2609d6 100644 --- a/tests/basic_correctness/memory/cumem/test_cumem.py +++ b/tests/basic_correctness/memory/cumem/test_cumem.py @@ -359,7 +359,7 @@ def test_runtime_state_survives_sleep(level): @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") @pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator") def test_level2_discards_ordinary_tensor_with_weights_tag(): - """Reproduce the level-2 variant for an ordinary tensor in weights.""" + """Discarded weights-tag memory is remapped; ROCm zeroes it over stale pages.""" allocator = get_mem_allocator_instance() with allocator.use_memory_pool("weights"): @@ -374,8 +374,9 @@ def test_level2_discards_ordinary_tensor_with_weights_tag(): torch.accelerator.synchronize() assert (fake_weight.data_ptr(), ordinary_tensor.data_ptr()) == pointers - assert torch.all(fake_weight == 0xA5) - assert torch.all(ordinary_tensor == 0xA5) + expected = 0 if current_platform.is_rocm() else 0xA5 + assert torch.all(fake_weight == expected) + assert torch.all(ordinary_tensor == expected) @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") @@ -427,7 +428,9 @@ def model(x): @pytest.mark.parametrize("level", [1, 2], ids=["sleep-1", "sleep-2"]) @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") -@pytest.mark.skipif(not current_platform.is_cuda(), reason="cuMem CUDA graph pool") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="cuMem CUDA graph pool" +) def test_cudagraph_pool_sleep(level): """Routing, and the graph pool backed up at both sleep levels and restored in place by any wake.""" diff --git a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py index 9f6180335ab1..d897d710015b 100644 --- a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py +++ b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py @@ -209,8 +209,20 @@ def _cudagraph_bytes(worker) -> int: def _custom_ar_active(worker) -> bool: from vllm.distributed.parallel_state import get_tp_group - ca_comm = get_tp_group().device_communicator.ca_comm - return ca_comm is not None and not ca_comm.disabled + comm = get_tp_group().device_communicator + return any( + c is not None and not c.disabled for c in (comm.ca_comm, comm.aiter_ar_comm) + ) + + +def _capture_registers_buffers(worker) -> bool: + from vllm.distributed.parallel_state import get_tp_group + + comm = get_tp_group().device_communicator + ca, aiter = comm.ca_comm, comm.aiter_ar_comm + return (ca is not None and ca._capture_registered) or ( + aiter is not None and aiter.aiter_ca.enable_register_for_capturing + ) @pytest.mark.parametrize( @@ -227,7 +239,9 @@ def _custom_ar_active(worker) -> bool: ids=["full", "piecewise", "breakable", "custom-ar-tp2", "off-by-default"], ) @create_new_process_for_each_test() -@pytest.mark.skipif(not current_platform.is_cuda(), reason="cuMem CUDA graph pool") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="cuMem CUDA graph pool" +) def test_sleep_cudagraph_pool(monkeypatch, mode, breakable, tp, offload): """Each capture site uses the pool, exact across sleeps; TP=2 needs unregistered custom AR; off changes nothing.""" @@ -255,8 +269,10 @@ def test_sleep_cudagraph_pool(monkeypatch, mode, breakable, tp, offload): assert all(graph_bytes) if offload else not any(graph_bytes) if not offload: return - if tp > 1 and not all(llm.collective_rpc(_custom_ar_active)): - pytest.skip("Custom allreduce is unavailable on these GPUs") + if tp > 1: + if not all(llm.collective_rpc(_custom_ar_active)): + pytest.skip("Custom allreduce is unavailable on these GPUs") + assert not any(llm.collective_rpc(_capture_registers_buffers)) prompt, params = "How are you?", SamplingParams(temperature=0, max_tokens=10) expected = llm.generate(prompt, params)[0].outputs[0].text for level in (1, 2): @@ -268,6 +284,42 @@ def test_sleep_cudagraph_pool(monkeypatch, mode, breakable, tp, offload): assert llm.generate(prompt, params)[0].outputs[0].text == expected +@multi_gpu_test(num_gpus=2) +@create_new_process_for_each_test() +@pytest.mark.skipif(not current_platform.is_rocm(), reason="AITER fused allreduce") +def test_sleep_cudagraph_pool_aiter_fused_ar_group_quant(monkeypatch): + """The fused AITER AR + RMSNorm + per-group FP8 quant captures into the + cuMem pool without IPC-registering its input, across sleeps.""" + for name, value in [ + ("VLLM_USE_V2_MODEL_RUNNER", "1"), + ("VLLM_ROCM_USE_AITER", "1"), + ("VLLM_ALLOW_INSECURE_SERIALIZATION", "1"), + ]: + monkeypatch.setenv(name, value) + monkeypatch.delenv("NCCL_GRAPH_REGISTER", raising=False) + llm = LLM( + "Qwen/Qwen3-0.6B-FP8", + enable_sleep_mode=True, + sleep_mode_offload_cudagraph=True, + tensor_parallel_size=2, + max_model_len=1024, + compilation_config={"pass_config": {"fuse_allreduce_rms": True}}, + ) + if not all(llm.collective_rpc(_custom_ar_active)): + pytest.skip("AITER custom allreduce is unavailable on these GPUs") + assert not any(llm.collective_rpc(_capture_registers_buffers)) + assert all(llm.collective_rpc(_cudagraph_bytes)) + prompt, params = "How are you?", SamplingParams(temperature=0, max_tokens=10) + assert llm.generate(prompt, params)[0].outputs[0].text + for level in (1, 2): + llm.sleep(level=level) + llm.wake_up(tags=["weights"]) + if level == 2: + llm.collective_rpc("reload_weights") + llm.wake_up(tags=["kv_cache"]) + assert llm.generate(prompt, params)[0].outputs[0].text + + @create_new_process_for_each_test() def test_deep_sleep_async(gpu_memory_cleared): async def test(): diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 44b89aeaf12c..0ed84906185c 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1167,7 +1167,8 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl( w=weight, eps=epsilon, group_size=group_size, - registered=torch.cuda.is_current_stream_capturing(), + registered=torch.cuda.is_current_stream_capturing() + and ca.enable_register_for_capturing, use_1stage=use_1stage, ) assert result is not None @@ -1219,7 +1220,8 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl( w=weight, eps=epsilon, group_size=group_size, - registered=torch.cuda.is_current_stream_capturing(), + registered=torch.cuda.is_current_stream_capturing() + and ca.enable_register_for_capturing, use_1stage=use_1stage, emit_bf16=True, ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 5d38b6765308..bbda00af40bb 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -702,7 +702,7 @@ def use_cumem_cudagraph_pool(self) -> bool: and model_config.enable_sleep_mode and model_config.sleep_mode_backend == "cumem" and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE - and current_platform.is_cuda() + and current_platform.is_cuda_alike() ) @property diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index b85be08d28e3..a2e582f85b07 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -98,7 +98,7 @@ class CuMemAllocator: offloaded to CPU memory, and the rest of the tensors will be discarded. When we call `wake_up`, all tensors that are previously offloaded will be loaded back to GPU memory, and the rest of the tensors will - have empty memory. + have empty memory (zeroed on ROCm). Why it needs to be a singleton? When allocated tensors are garbage collected, PyTorch will call @@ -331,7 +331,7 @@ def discard(self, tags: tuple[str, ...] | str) -> None: def wake_up(self, tags: list[str] | None = None) -> None: """Wake up the allocator from sleep mode. All data that is previously offloaded will be loaded back to GPU - memory, and the rest of the data will have empty memory. + memory, and the rest will have empty memory (zeroed on ROCm). Args: tags: The deferrable tags (weights, kv_cache) to load back to GPU @@ -358,6 +358,10 @@ def wake_up(self, tags: list[str] | None = None) -> None: cpu_ptr = cpu_backup_tensor.data_ptr() libcudart.cudaMemcpy(ptr, cpu_ptr, size_in_bytes) data.cpu_backup_tensor = None + elif current_platform.is_rocm(): + # amdgpu <6.14 may return stale VRAM (#44972); drop with its support + libcudart.cudaMemset(ptr, 0, handle[1]) + torch.accelerator.synchronize() @contextmanager def use_memory_pool(self, tag: str | None = None): diff --git a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py index 6fbffcaaf73d..9a5970c565d7 100644 --- a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py +++ b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py @@ -30,6 +30,8 @@ def __init__( group: ProcessGroup, device: int | str | torch.device, max_size: int | None = None, + *, + register_graph_buffers: bool = True, ): from aiter.dist.device_communicators.custom_all_reduce import ( CustomAllreduce as _AiterCustomAllreduce, @@ -38,7 +40,10 @@ def __init__( if max_size is None: max_size = self.MAX_SIZE - self._impl = _AiterCustomAllreduce(group, device, max_size=max_size) + kwargs = {} + if not register_graph_buffers: + kwargs["enable_register_for_capturing"] = False + self._impl = _AiterCustomAllreduce(group, device, max_size=max_size, **kwargs) @property def aiter_ca(self): diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index d4d588bf450d..992d07e15dad 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -151,6 +151,7 @@ def __init__( self.aiter_ar_comm = AiterCustomAllreduce( group=self.cpu_group, device=self.device, + register_graph_buffers=register, ) if use_custom_allreduce and self.aiter_ar_comm is None and self.world_size > 1: @@ -174,6 +175,7 @@ def __init__( self.aiter_ar_comm = AiterCustomAllreduce( group=self.cpu_group, device=self.device, + register_graph_buffers=register, ) if self.aiter_ar_comm.disabled: self.aiter_ar_comm = None From 18c19eaf239f09bc2543637c73f3b2d05f9bf5c5 Mon Sep 17 00:00:00 2001 From: Taewoong Kim <116135174+ultramancode@users.noreply.github.com> Date: Fri, 9 Oct 2026 17:53:50 +0900 Subject: [PATCH 163/624] [Frontend] Add cache salt and tracing to structured decisions (#60786) Signed-off-by: Taewoong Kim Co-authored-by: Chauncey --- .../online_serving/structured_decisions.md | 5 + docs/usage/security.md | 2 +- .../structured_decisions/test_strategies.py | 135 +++++++++++++++++- .../test_non_object_body_validation.py | 15 ++ .../generate/structured_decisions/protocol.py | 22 ++- .../generate/structured_decisions/serving.py | 23 +++ .../structured_decisions/strategies.py | 15 +- 7 files changed, 207 insertions(+), 10 deletions(-) diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md index 2194b54e47b4..a1f58a3769b9 100644 --- a/docs/serving/online_serving/structured_decisions.md +++ b/docs/serving/online_serving/structured_decisions.md @@ -135,6 +135,11 @@ the request lists them, a `noul` labels its options `yes` and `no`, and a | `questions` | question id to `{type, instructions, criteria}`, asked in this order | | `instructions` | optional context placed ahead of the questions | | `chat_template_kwargs` | passed to the chat template, for example `{"enable_thinking": false}` | +| `cache_salt` | optional [prefix-cache salt](../../usage/security.md#cache-salting) applied to every question | A question with any field other than `type`, `instructions` and `criteria` is rejected with a 400. + +When [tracing](../../../examples/observability/opentelemetry/README.md) is enabled +with `--otlp-traces-endpoint`, the HTTP `traceparent` and `tracestate` headers are +forwarded to every question's generation request. diff --git a/docs/usage/security.md b/docs/usage/security.md index c0d2c79700e6..a74ed1d6ebe9 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -572,7 +572,7 @@ An attacker sharing the same backend can measure differences in Time to First To vLLM accepts an optional `cache_salt` parameter on requests. The salt is mixed into the hash of the first KV cache block, so only requests carrying the same salt can share cached prefix blocks. See [Automatic Prefix Caching](../design/prefix_caching.md) for the implementation details. -`cache_salt` is accepted by the OpenAI-compatible chat completions, completions, responses, and pooling (embeddings, classification, scoring) endpoints, and by the Anthropic `/v1/messages` endpoint. +`cache_salt` is accepted by the OpenAI-compatible chat completions, completions, responses, and pooling (embeddings, classification, scoring) endpoints, the Anthropic `/v1/messages` endpoint, and the structured decisions `/v1/systemone` endpoint. #### Usage with the OpenAI Python client diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index 3d9538e4dc8b..5dcb8edb3061 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -2,24 +2,46 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace from typing import Any +from unittest.mock import MagicMock import pytest -from fastapi import FastAPI +from fastapi import FastAPI, Request from fastapi.testclient import TestClient +from vllm.config import ModelConfig +from vllm.engine.protocol import EngineClient from vllm.entrypoints.generate.structured_decisions.api_router import ( register_structured_decisions_api_router, ) +from vllm.entrypoints.generate.structured_decisions.protocol import ( + StructuredDecisionRequest, + StructuredDecisionResponse, +) from vllm.entrypoints.generate.structured_decisions.question_types import LABELS -from vllm.entrypoints.generate.structured_decisions.serving import state_text +from vllm.entrypoints.generate.structured_decisions.serving import ( + ServingStructuredDecisions, + state_text, +) from vllm.entrypoints.generate.structured_decisions.strategies import ( NextTokenStrategy, + ReadContext, reply_label_ids, select_read_strategy, ) +from vllm.entrypoints.openai.models.protocol import BaseModelPath +from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.serve.exception_handling.register import ( init_exception_handler, ) +from vllm.logprobs import Logprob +from vllm.outputs import CompletionOutput, RequestOutput +from vllm.renderers.hf import HfRenderer +from vllm.renderers.online_renderer import OnlineRenderer + +TRACE_HEADERS = { + "traceparent": "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01", + "tracestate": "vendor=value", +} def model(architecture: str, logprobs_mode: str = "raw_logprobs") -> Any: @@ -69,9 +91,9 @@ def test_structured_state_keeps_unicode_in_prompt(): @pytest.fixture(scope="module") def qwen(): - from transformers import AutoTokenizer + from vllm.tokenizers import get_tokenizer - tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") + tokenizer = get_tokenizer("Qwen/Qwen3-0.6B") prompt_ids = tokenizer.apply_chat_template( [{"role": "user", "content": "state"}], add_generation_prompt=True, @@ -93,3 +115,108 @@ def test_labels_start_the_reply(qwen): # After a colon, Qwen writes ":A" as one token, so "A" is not one token. with pytest.raises(ValueError, match="not one distinct token"): reply_label_ids(tokenizer, tokenizer.encode("team:")) + + +@pytest.fixture +def serving_and_engine(qwen): + tokenizer, _ = qwen + model_config = ModelConfig(model="Qwen/Qwen3-0.6B", dtype="float32") + renderer = HfRenderer( + SimpleNamespace( + model_config=model_config, + parallel_config=SimpleNamespace(_api_process_rank=0), + ), + tokenizer, + ) + engine = MagicMock(spec=EngineClient) + engine.errored = False + engine.model_config = model_config + engine.renderer = renderer + engine.input_processor = MagicMock() + online_renderer = OnlineRenderer( + model_config, + renderer, + request_logger=None, + chat_template=None, + chat_template_content_format="auto", + ) + + async def generate(prompt, params, request_id, **kwargs): + label_ids = params.logprob_token_ids + yield RequestOutput( + request_id=request_id, + prompt=None, + prompt_token_ids=prompt["prompt_token_ids"], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text="A", + token_ids=[label_ids[0]], + cumulative_logprob=None, + logprobs=[{token_id: Logprob(-1.0) for token_id in label_ids}], + finish_reason="length", + ) + ], + finished=True, + ) + + engine.generate.side_effect = generate + strategy = NextTokenStrategy(ReadContext(engine, online_renderer, None, "auto", {})) + models = OpenAIServingModels( + engine, [BaseModelPath(name=model_config.model, model_path=model_config.model)] + ) + yield ServingStructuredDecisions(models, strategy), engine + renderer.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "cache_salt,tracing_enabled,request_headers", + [ + ("tenant-a", True, {"authorization": "Bearer test", **TRACE_HEADERS}), + ("tenant-b", False, {"authorization": "Bearer test", **TRACE_HEADERS}), + (None, True, {"authorization": "Bearer test", **TRACE_HEADERS}), + (None, True, {}), + ("tenant-direct", True, None), + ], +) +async def test_request_metadata_reaches_every_question( + serving_and_engine, cache_salt, tracing_enabled, request_headers +): + serving, engine = serving_and_engine + engine.is_tracing_enabled.return_value = tracing_enabled + raw_request = None + if request_headers is not None: + raw_request = Request( + { + "type": "http", + "headers": [ + (k.encode(), v.encode()) for k, v in request_headers.items() + ], + } + ) + request = StructuredDecisionRequest( + state={"team": "billing", "language": "English"}, + questions={ + "team": { + "type": "choice", + "instructions": "Which team does the message name?", + "criteria": {"billing": "payments", "shipping": "deliveries"}, + }, + "lang": { + "type": "choice", + "instructions": "Which language does the message name?", + "criteria": {"English": "English text", "French": "French text"}, + }, + }, + cache_salt=cache_salt, + ) + response = await serving.create_decision(request, raw_request) + assert isinstance(response, StructuredDecisionResponse) + assert engine.generate.call_count == 2 + for call in engine.generate.call_args_list: + assert call.args[0].get("cache_salt") == cache_salt + assert call.kwargs["trace_headers"] == ( + TRACE_HEADERS if tracing_enabled and request_headers else None + ) diff --git a/tests/entrypoints/unit_tests/test_non_object_body_validation.py b/tests/entrypoints/unit_tests/test_non_object_body_validation.py index 5bdc1d57873b..dd0e3eafd0a4 100644 --- a/tests/entrypoints/unit_tests/test_non_object_body_validation.py +++ b/tests/entrypoints/unit_tests/test_non_object_body_validation.py @@ -13,6 +13,9 @@ import pytest from pydantic import ValidationError +from vllm.entrypoints.generate.structured_decisions.protocol import ( + StructuredDecisionRequest, +) from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.completion.protocol import CompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest @@ -74,6 +77,18 @@ def test_tokenize_chat_request_still_validates_dict_bodies(): CACHE_SALT_MODELS = [ + ( + StructuredDecisionRequest, + { + "state": "hello", + "questions": { + "lang": { + "type": "choice", + "criteria": {"English": "Text written in English."}, + } + }, + }, + ), (CompletionRequest, {"prompt": "hello"}), (ChatCompletionRequest, {"messages": [{"role": "user", "content": "hello"}]}), (ResponsesRequest, {"input": "hello"}), diff --git a/vllm/entrypoints/generate/structured_decisions/protocol.py b/vllm/entrypoints/generate/structured_decisions/protocol.py index 3a193e203bad..b15fee171fca 100644 --- a/vllm/entrypoints/generate/structured_decisions/protocol.py +++ b/vllm/entrypoints/generate/structured_decisions/protocol.py @@ -5,10 +5,11 @@ import time from typing import Any, Literal -from pydantic import Field +from pydantic import Field, model_validator from vllm.config import ModelConfig from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption +from vllm.entrypoints.generate.base.protocol import validate_cache_salt from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs from vllm.utils import random_uuid @@ -30,9 +31,27 @@ class StructuredDecisionRequest(OpenAIBaseModel): default=None, description="Context placed ahead of the questions." ) chat_template_kwargs: dict[str, Any] | None = None + cache_salt: str | None = Field( + default=None, + description=( + "If specified, the prefix cache will be salted with the provided " + "string to prevent an attacker to guess prompts in multi-user " + "environments. The salt should be random, protected from " + "access by 3rd parties, and long enough to be " + "unpredictable (e.g., 43 characters base64-encoded, corresponding " + "to 256 bit)." + ), + ) priority: int = Field(default=0, ge=-(2**63), le=2**63 - 1) request_id: str = Field(default_factory=random_uuid) + @model_validator(mode="before") + @classmethod + def check_cache_salt_support(cls, data: Any) -> Any: + if isinstance(data, dict): + validate_cache_salt(data.get("cache_salt")) + return data + class DecisionUsage(OpenAIBaseModel): input_tokens: int @@ -65,6 +84,7 @@ class ReadPromptRequest(OpenAIBaseModel): would follow a thought rather than start the reply.""" chat_template_kwargs: dict[str, Any] | None = None + cache_salt: str | None = None def build_chat_params( self, diff --git a/vllm/entrypoints/generate/structured_decisions/serving.py b/vllm/entrypoints/generate/structured_decisions/serving.py index 4921bfcf3574..e6e990bd46df 100644 --- a/vllm/entrypoints/generate/structured_decisions/serving.py +++ b/vllm/entrypoints/generate/structured_decisions/serving.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio import json +from collections.abc import Mapping from typing import Any from fastapi import Request @@ -11,6 +12,11 @@ from vllm.entrypoints.serve.engine.serving import BaseServing from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.logger import init_logger +from vllm.tracing import ( + contains_trace_headers, + extract_trace_headers, + log_tracing_disabled_warning, +) from .protocol import ( DecisionUsage, @@ -71,6 +77,16 @@ def __init__( self.strategy = strategy self.limits = strategy.limits() + async def _get_trace_headers( + self, headers: Mapping[str, str] + ) -> Mapping[str, str] | None: + if not contains_trace_headers(headers): + return None + if not await self.strategy.context.engine_client.is_tracing_enabled(): + log_tracing_disabled_warning() + return None + return extract_trace_headers(headers) + async def create_decision( self, request: StructuredDecisionRequest, @@ -88,6 +104,11 @@ async def create_decision( questions = parse_questions(request, self.limits) lora_request = self._maybe_get_adapters(request) # type: ignore[arg-type] engine_client.check_admission(len(questions)) + trace_headers = ( + None + if raw_request is None + else await self._get_trace_headers(raw_request.headers) + ) reads = await self.strategy.read( questions, request.instructions, @@ -96,6 +117,8 @@ async def create_decision( chat_template_kwargs=request.chat_template_kwargs, lora_request=lora_request, priority=request.priority, + cache_salt=request.cache_salt, + trace_headers=trace_headers, ) except StructuredDecisionError as e: return self.create_error_response(e) diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index af588e13a289..65b5a9dd7b05 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -7,7 +7,7 @@ import math from abc import ABC, abstractmethod -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import Any @@ -74,6 +74,8 @@ async def read( chat_template_kwargs: dict[str, Any] | None, lora_request: LoRARequest | None, priority: int, + cache_salt: str | None = None, + trace_headers: Mapping[str, str] | None = None, ) -> list[QuestionRead]: """One read per question, in order.""" @@ -136,13 +138,15 @@ def limits(self) -> DecisionLimits: return DecisionLimits(max_questions=64, max_options=len(LABELS)) def _read_request( - self, chat_template_kwargs: dict[str, Any] | None + self, chat_template_kwargs: dict[str, Any] | None, cache_salt: str | None = None ) -> ReadPromptRequest: if (chat_template_kwargs or {}).get("enable_thinking"): raise StructuredDecisionError( "a read is the reply's first token, so thinking must be off" ) - return ReadPromptRequest(chat_template_kwargs=chat_template_kwargs) + return ReadPromptRequest( + chat_template_kwargs=chat_template_kwargs, cache_salt=cache_salt + ) async def _render( self, read_request: ReadPromptRequest, messages: list[dict[str, Any]] @@ -170,9 +174,11 @@ async def read( chat_template_kwargs: dict[str, Any] | None, lora_request: LoRARequest | None, priority: int, + cache_salt: str | None = None, + trace_headers: Mapping[str, str] | None = None, ) -> list[QuestionRead]: ctx = self.context - read_request = self._read_request(chat_template_kwargs) + read_request = self._read_request(chat_template_kwargs, cache_salt) slots, engine_inputs = [], [] for q in questions: @@ -197,6 +203,7 @@ async def read( ], request_id, lora_request=lora_request, + trace_headers=trace_headers, priority=priority, ) From bc713cf8cc861c0d0ccfe46c9ce322ab81993ace Mon Sep 17 00:00:00 2001 From: Chao Lei Date: Fri, 9 Oct 2026 17:19:33 +0800 Subject: [PATCH 164/624] [Bugfix]MooncakeConnect add query timeout and retry (#55923) --- docs/features/mooncake_connector_usage.md | 5 ++-- .../v1/mooncake/mooncake_connector.py | 28 +++++++++++++++++-- vllm/models/deepseek_v41/common/engram.py | 2 +- 3 files changed, 29 insertions(+), 6 deletions(-) diff --git a/docs/features/mooncake_connector_usage.md b/docs/features/mooncake_connector_usage.md index 866d6d6bd48c..00bfdeabc14c 100644 --- a/docs/features/mooncake_connector_usage.md +++ b/docs/features/mooncake_connector_usage.md @@ -49,12 +49,13 @@ cache), see [ECMooncakeConnector Usage Guide](mooncake_ec_connector_usage.md). - For headless instances, must be the same as the master instance - Each instance needs a unique port on its host; using the same port number across different hosts is fine -- `VLLM_MOONCAKE_CONNECTOR_TIMEOUT`: Per-request timeout (in seconds) for MooncakeConnector HTTP calls to the bootstrap server, including prefiller worker registration at startup. (Optional) +- `VLLM_MOONCAKE_CONNECTOR_TIMEOUT`: Per-request timeout (in seconds) for MooncakeConnector HTTP calls to the bootstrap server, including prefiller worker registration at startup and metadata queries. (Optional) - Default: 30.0 - Global rank 0 hosts the bootstrap server in the same process that mounts the Mooncake transfer engine's host segment, so a large `MOONCAKE_GLOBAL_SEGMENT_SIZE` can delay responses for several seconds during startup - Raise this if bootstrap server timeout warnings appear on hosts with very large host-memory segments - Bootstrap registration uses up to 3 attempts internally, with exponential backoff capped at 10 seconds; connection errors and timeouts are both retried - - Once the attempts are exhausted the caller raises instead of blocking, so failures surface rather than hanging + - Once registration attempts are exhausted the caller raises instead of blocking, so failures surface rather than hanging + - Metadata queries reuse the same timeout and 3-attempt limit, with a fixed 100ms delay between attempts. If all attempts fail, the error is logged and waiting queries are notified - `WITH_NVIDIA_PEERMEM`: Selects how mooncake registers GPU memory for RDMA. Read by mooncake, not vLLM. - Default: 1, which uses `ibv_reg_mr()` and requires the `nvidia-peermem` kernel module to be loaded diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 8300cd950b0d..9d5194f33e27 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -2438,9 +2438,30 @@ def process_pulling_result( async def _connect_to_prefiller_bootstrap(self, remote_bootstrap_addr: str): url = remote_bootstrap_addr + "/query" + max_attempts = _BOOTSTRAP_MAX_ATTEMPTS try: - async with httpx.AsyncClient() as client: - response = await client.get(url) + async with httpx.AsyncClient( + timeout=envs.VLLM_MOONCAKE_CONNECTOR_TIMEOUT + ) as client: + retry_delay = 0.1 + for attempt in range(1, max_attempts + 1): + try: + response = await client.get(url) + break + except httpx.RequestError as e: + if attempt == max_attempts: + raise + logger.warning( + "Bootstrap query to %s failed on attempt %d/%d " + "(%s: %s); retrying in %.1f seconds", + remote_bootstrap_addr, + attempt, + max_attempts, + type(e).__name__, + e, + retry_delay, + ) + await asyncio.sleep(retry_delay) response.raise_for_status() data: dict = response.json() for _, dp_entry in data.items(): @@ -2455,8 +2476,9 @@ async def _connect_to_prefiller_bootstrap(self, remote_bootstrap_addr: str): self._tp_size[remote_engine_id] = len(dp_entry["worker_addr"]) except Exception as e: logger.error( - "Failed to connect to bootstrap server %s: %s", + "Failed to connect to bootstrap server %s (%s): %s", remote_bootstrap_addr, + type(e).__name__, e, ) diff --git a/vllm/models/deepseek_v41/common/engram.py b/vllm/models/deepseek_v41/common/engram.py index 0e160c8b06cf..236235dfa82f 100644 --- a/vllm/models/deepseek_v41/common/engram.py +++ b/vllm/models/deepseek_v41/common/engram.py @@ -1305,7 +1305,7 @@ def _engram_select_rows( ) -> None: """Copy one token window out of a rank-major gathered buffer. - Both gathers land rank-major ([rank][token][local width]); this walks the + Both gathers land rank-major (`[rank][token][local width]`); this walks the window the rank keeps and lays its ranks out side by side as width. """ if output.numel() == 0: From 2c99ee9333821030daf84718b2a76ea28d94ba3c Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Fri, 9 Oct 2026 10:29:03 +0100 Subject: [PATCH 165/624] [Chore] Remove vendored code made redundant by Transformers 5.16.1 (#60623) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- .agents/skills/transformers-upgrade/SKILL.md | 178 +++++++++++ .../transformers-upgrade/agents/openai.yaml | 4 + .../scripts/diff_dumps.py | 51 ++++ .../scripts/dump_configs.py | 30 ++ .../scripts/dump_model.py | 114 +++++++ .../scripts/find_upstreamed.py | 147 +++++++++ .buildkite/test-amd.yaml | 10 - .buildkite/test_areas/lm_eval.yaml | 7 - .buildkite/test_areas/misc.yaml | 1 - .buildkite/test_areas/models_language.yaml | 2 - .claude/skills/transformers-upgrade | 1 + tests/config/test_model_arch_config.py | 16 +- .../token_in_token_out/test_serving_tokens.py | 10 +- .../layers/test_fused_shared_expert.py | 2 +- .../multimodal/generation/test_common.py | 5 +- .../processing/test_transformers_image.py | 6 +- tests/models/registry.py | 12 +- .../bitsandbytes/test_bitsandbytes.py | 5 +- tests/tokenizers_/test_registry.py | 3 +- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 2 +- vllm/model_executor/models/AXK1.py | 2 +- vllm/model_executor/models/aimv2.py | 20 +- vllm/model_executor/models/cohere2_moe.py | 14 - vllm/model_executor/models/cosmos3.py | 19 +- vllm/model_executor/models/cosmos3_edge.py | 3 +- vllm/model_executor/models/diffusion_gemma.py | 27 +- .../model_executor/models/embedding_gemma2.py | 7 +- vllm/model_executor/models/gemma3.py | 15 +- vllm/model_executor/models/gemma3n.py | 12 +- vllm/model_executor/models/gemma4.py | 7 +- vllm/model_executor/models/gemma4_dspark.py | 3 +- vllm/model_executor/models/gemma4_mm.py | 16 +- vllm/model_executor/models/gemma4_mtp.py | 3 +- vllm/model_executor/models/granite4_vision.py | 96 +++--- vllm/model_executor/models/granite_speech.py | 7 +- .../model_executor/models/granitemoehybrid.py | 4 - vllm/model_executor/models/hy_v3.py | 3 +- vllm/model_executor/models/hyperclovax.py | 17 +- vllm/model_executor/models/kimi_k25.py | 17 +- vllm/model_executor/models/kimi_k25_vit.py | 44 +-- vllm/model_executor/models/lfm2_moe.py | 2 +- vllm/model_executor/models/modernbert.py | 28 +- vllm/model_executor/models/nemotron.py | 2 +- vllm/model_executor/models/nemotron_h.py | 2 +- vllm/model_executor/models/nemotron_h_mtp.py | 2 +- vllm/model_executor/models/ovis.py | 2 +- vllm/model_executor/models/qwen3_5.py | 11 +- vllm/model_executor/models/qwen3_5_mtp.py | 3 +- vllm/model_executor/models/qwen3_next.py | 2 +- vllm/model_executor/models/qwen3_next_mtp.py | 2 +- vllm/models/kimi_k3/amd/model.py | 1 + vllm/models/kimi_k3/nvidia/model.py | 1 + vllm/transformers_utils/config.py | 47 +-- vllm/transformers_utils/configs/AXK1.py | 216 ------------- vllm/transformers_utils/configs/__init__.py | 40 +-- vllm/transformers_utils/configs/afmoe.py | 87 ------ vllm/transformers_utils/configs/cosmos3.py | 7 - .../configs/cosmos3_edge.py | 11 +- .../configs/deepseek_vl2.py | 5 - .../configs/diffusion_gemma.py | 44 --- vllm/transformers_utils/configs/gemma4.py | 34 --- .../configs/granite4_vision.py | 94 ------ vllm/transformers_utils/configs/hy_v3.py | 185 ----------- .../transformers_utils/configs/hyperclovax.py | 277 ----------------- vllm/transformers_utils/configs/kimi_k25.py | 131 -------- vllm/transformers_utils/configs/lfm2_moe.py | 164 ---------- vllm/transformers_utils/configs/nemotron.py | 219 ------------- vllm/transformers_utils/configs/nemotron_h.py | 288 ------------------ .../transformers_utils/configs/olmo_hybrid.py | 285 ----------------- vllm/transformers_utils/configs/ovis.py | 72 +---- vllm/transformers_utils/configs/qwen3_5.py | 193 ------------ .../transformers_utils/configs/qwen3_5_moe.py | 205 ------------- vllm/transformers_utils/configs/qwen3_next.py | 277 ----------------- .../model_arch_config_convertor.py | 24 +- .../transformers_utils/processors/__init__.py | 4 - .../processors/granite4_vision.py | 78 ----- 76 files changed, 719 insertions(+), 3268 deletions(-) create mode 100644 .agents/skills/transformers-upgrade/SKILL.md create mode 100644 .agents/skills/transformers-upgrade/agents/openai.yaml create mode 100644 .agents/skills/transformers-upgrade/scripts/diff_dumps.py create mode 100644 .agents/skills/transformers-upgrade/scripts/dump_configs.py create mode 100644 .agents/skills/transformers-upgrade/scripts/dump_model.py create mode 100644 .agents/skills/transformers-upgrade/scripts/find_upstreamed.py create mode 120000 .claude/skills/transformers-upgrade delete mode 100644 vllm/transformers_utils/configs/AXK1.py delete mode 100644 vllm/transformers_utils/configs/afmoe.py delete mode 100644 vllm/transformers_utils/configs/cosmos3.py delete mode 100644 vllm/transformers_utils/configs/diffusion_gemma.py delete mode 100644 vllm/transformers_utils/configs/gemma4.py delete mode 100644 vllm/transformers_utils/configs/granite4_vision.py delete mode 100644 vllm/transformers_utils/configs/hy_v3.py delete mode 100644 vllm/transformers_utils/configs/hyperclovax.py delete mode 100644 vllm/transformers_utils/configs/kimi_k25.py delete mode 100644 vllm/transformers_utils/configs/lfm2_moe.py delete mode 100644 vllm/transformers_utils/configs/nemotron.py delete mode 100644 vllm/transformers_utils/configs/nemotron_h.py delete mode 100644 vllm/transformers_utils/configs/olmo_hybrid.py delete mode 100644 vllm/transformers_utils/configs/qwen3_5.py delete mode 100644 vllm/transformers_utils/configs/qwen3_5_moe.py delete mode 100644 vllm/transformers_utils/configs/qwen3_next.py delete mode 100644 vllm/transformers_utils/processors/granite4_vision.py diff --git a/.agents/skills/transformers-upgrade/SKILL.md b/.agents/skills/transformers-upgrade/SKILL.md new file mode 100644 index 000000000000..de9065b6fa81 --- /dev/null +++ b/.agents/skills/transformers-upgrade/SKILL.md @@ -0,0 +1,178 @@ +--- +name: transformers-upgrade +description: "Playbook for changing vLLM's Transformers version bounds. Use when bumping the pinned or maximum Transformers version, raising the minimum Transformers version, or deleting vendored configs, processors, or compat shims that upstream Transformers now provides." +--- + +# Upgrading Transformers + +The supported range lives in `requirements/common.txt`. CI pins one exact +version in `requirements/test/*.in` and the compiled `requirements/test/*.txt`. +A version change is one of two jobs, and they usually come as a pair: + +- **Bump the pin**: a new Transformers release becomes the CI version and + the upper bound. vLLM grows compat code for it. +- **Raise the floor**: the minimum moves up. That compat code, and every + vendored config or processor upstream now ships, becomes dead weight. + Every floor raise ends with an **unvendoring pass**. + +## Bump the pin + +1. Edit the pin line in each `requirements/test/*.in` and `*.txt` and the + upper bound in `requirements/common.txt`. Past bumps edit only the pin + line in the compiled files; see `git log -- requirements/common.txt`. +2. Read the release notes for config validation changes, renamed auto-mapping + formats, and removed attributes; these are what break vLLM. +3. Search `tests/` for `min_transformers_version` and skip markers naming the + new version, and un-gate them. +4. Run the suites in [Testing](#testing) and fix what breaks. Mark every + shim with a `# Transformers >= X` comment so the next floor raise finds it. + +Done when the [Testing](#testing) suites pass on the new pin and every test +waiting on this version is un-gated. + +## Raise the floor (unvendoring pass) + +### 1. Inventory + +Inventory candidates with the floor version installed (no vLLM install needed): + +```bash +uv run --no-project --with transformers== \ + python .agents/skills/transformers-upgrade/scripts/find_upstreamed.py +``` + +It lists: + +- vendored configs whose `model_type` upstream now owns; +- `_CONFIG_REGISTRY` overrides of upstream types; +- registered processors that collide with upstream class names or models; +- version-gated code at or below the floor; +- code that reads the installed Transformers version, whatever its threshold; +- patches citing a Transformers issue or PR (with whether a referenced PR is + in the floor release) or a future minimum version. + +It finds candidates by name only. Every hit still needs a verdict. + +### 2. Port each candidate to upstream + +Every hit is deleted by default. A vendored copy is a liability: it drifts +from upstream, hides upstream fixes, and shadows the native class. Port the +vLLM model code to the upstream schema instead, renaming attribute reads and +deriving values from upstream fields. + +Fan candidates out to parallel subagents. Each one compares the vendored and +upstream classes, then greps every attribute the vLLM model code reads. +[Upstreaming hazards](#upstreaming-hazards) lists the usual fixes. + +Keep a vendored class only when upstream cannot load a checkpoint vLLM must +keep serving, and record why in the PR description. A different upstream +schema, a model-code rewrite, or a needed eval is work for the PR, not a +reason to keep. + +### 3. Prove it on real checkpoints + +For every checkpoint in `tests/models/registry.py` that uses a ported +config, dump what vLLM builds before and after the change, then diff the two. +Run both dumps; each catches what the other misses: + +- `dump_configs.py`: the resolved config, attribute by attribute. +- `dump_model.py`: the model built on the meta device with no GPU or weights. + It records ModelConfig sizes, every parameter shape, and every module's + scalar attributes (`head_dim`, `num_kv_heads`, `scaling`, ...). Missing + fields that are read with defaults only show up here. + +```bash +git worktree add --detach /vllm-base upstream/main +S=.agents/skills/transformers-upgrade/scripts +for dump in dump_configs dump_model; do + PYTHONPATH=/vllm-base .venv/bin/python $S/$dump.py before.jsonl + PYTHONPATH=$PWD .venv/bin/python $S/$dump.py after.jsonl + .venv/bin/python $S/diff_dumps.py before.jsonl after.jsonl +done +``` + +Add `--trust-remote-code` for checkpoints whose `config.json` has an +`auto_map`, and dump both ways. + +How to read the diffs: + +- **A `` config field the model reads:** this crashes or silently + falls back to a default. Port the model code to the upstream field. +- **New upstream-only defaults** (`attention_bias=False`, + `pad_token_id=None`): usually benign. A model diff of 0 confirms it. +- **Any model diff:** this is a numerics change, and the PR needs a model + eval. For example, A.X-K1 gained the checkpoint's YaRN rope scaling, which + the vendored class had been dropping; it shows up as + `layers.*.self_attn.scaling (x61)`. + +### 4. Delete + +For each deleted config: + +- Delete the file. +- Remove its entries from `_CLASS_TO_MODULE` and `__all__` in + `vllm/transformers_utils/configs/__init__.py`. +- Remove its entries from `_CONFIG_REGISTRY` in + `vllm/transformers_utils/config.py`. +- Switch every importer, including tests, to `from transformers import ...`. +- Remove the `source_file_dependencies` lines naming the file from + `.buildkite/`. + +Deleted processors follow the same steps in +`vllm/transformers_utils/processors/__init__.py`. + +Done when every line of the inventory output is deleted, or kept with a +reason that meets the bar in step 2. The PR also carries the config diff +results and evals for any changed values. + +## Upstreaming hazards + +Upstream configs are `@strict` dataclasses. Checkpoint keys upstream does not +declare survive as plain attributes. Defaults that only vLLM sets, and values +vLLM derives, vanish. + +| Hazard | Signal | Fix | +| --- | --- | --- | +| Upstream consumes legacy keys | `__post_init__` pops keys the model reads (`deepseek_v4` `compress_ratios`) | Read the upstream replacement (`layer_types`, `mlp_layer_types`, per-layer `rope_parameters`) | +| Renamed or dropped fields | Diff shows `` for a field the model reads | Read the upstream name; derive dropped flags from what implies them (`enable_moe_block` from `num_experts`) | +| Type check on the config | `ctx.get_hf_config(SomeConfig)` relies on the vendored base class | Accept the upstream class too, e.g. a `Cosmos3ProcessingInfo` overriding `get_hf_config` | +| Remote code takes over | Checkpoint `config.json` has `auto_map`; with `--trust-remote-code`, the hub class replaces the native one once vLLM stops registering its own | Check the remote class loads; if not, register the upstream class for that `model_type` | +| Read-only upstream fields | Properties or derived fields break subclasses and `hf_overrides` (`nemotron_h` `hybrid_override_pattern`) | Override the source field (`layers_block_type`) instead | +| Behavioural rewrite | Upstream derives fields differently (`diffusion_gemma` `sliding_window`) | Upstream matches the reference implementation; verify with an eval | +| Original checkpoint format | Upstream targets `-hf` conversions; vLLM serves the originals (`qwen3_asr`, `qianfan_ocr`) | Keep only if upstream cannot load the original at all | +| Processor output contract | Model code consumes keys upstream does not emit | Port the model to upstream's keys | +| Name collision | A registered vendored processor shares an upstream class name and hijacks `get_processor` (`InternVLProcessor`) | Drop the registry entry | + +Version gates from the inventory need a decision: + +- **Code** only reachable below the floor is deleted. This covers: + - comparisons against `TRANSFORMERS_VERSION`; + - branches labelled `# Transformers < X`; + - fallbacks for fields that upstream configs now always set, e.g. + `mlp_layer_types`, `per_layer_config` or legacy `layer_types` names. + + A `skipif` whose condition is now always true becomes an unconditional + `skip` with the same reason. +- **`min_transformers_version`** at or below the floor is removed. +- **`max_transformers_version`** below the floor means CI never tests that + model; report it rather than deleting it silently. +- **Patches citing a Transformers issue or PR:** + - Delete the patch when its fix is in the floor release. + - For an issue, find the PR that closed it and check that PR. A closed + issue is not a fixed one: #43329 was closed while the bug was still in + 5.16.1. + - Re-run the patched behaviour before deleting. A fix in the release does + not always cover vLLM's path: #47924 made Emu3's processor add BOS, but + vLLM tokenizes without the processor. + - Links that only explain a design choice (e.g. a float32 RoPE note) are + not patches; leave them. + +## Testing + +```bash +.venv/bin/python -m pytest tests/transformers_utils tests/config tests/tokenizers_ +.venv/bin/python -m pytest tests/models/multimodal/processing/test_common.py -k "" +``` + +Add `tests/models/test_initialization.py -k ""` where a GPU is +available, and a model eval for any model whose resolved config changed. diff --git a/.agents/skills/transformers-upgrade/agents/openai.yaml b/.agents/skills/transformers-upgrade/agents/openai.yaml new file mode 100644 index 000000000000..b881c5d7ea27 --- /dev/null +++ b/.agents/skills/transformers-upgrade/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "Transformers Upgrade" + short_description: "Bump vLLM's Transformers version and delete vendored code upstream now provides." + default_prompt: "Use $transformers-upgrade to raise the minimum Transformers version and remove vendored configs, processors, and shims it makes redundant." diff --git a/.agents/skills/transformers-upgrade/scripts/diff_dumps.py b/.agents/skills/transformers-upgrade/scripts/diff_dumps.py new file mode 100644 index 000000000000..76fb2038a151 --- /dev/null +++ b/.agents/skills/transformers-upgrade/scripts/diff_dumps.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Diff two dump_configs.py or dump_model.py outputs, per repo. + +python diff_dumps.py before.jsonl after.jsonl +""" + +import json +import sys +from collections import Counter + +import regex as re + +IGNORED = {"transformers_version", "_name_or_path", "architectures", "auto_map"} + + +def load(path: str) -> dict[str, dict]: + with open(path) as f: + return {record["repo"]: record for record in map(json.loads, f)} + + +def diff(before: dict, after: dict, prefix: str = ""): + for key in sorted(set(before) | set(after)): + if key in IGNORED: + continue + old, new = before.get(key, ""), after.get(key, "") + if isinstance(old, dict) and isinstance(new, dict) and key != "rope_parameters": + yield from diff(old, new, f"{prefix}{key}.") + elif old != new: + yield f"{prefix}{key}", json.dumps(old), json.dumps(new) + + +def collapse(diffs) -> Counter: + """Merge diffs that differ only by layer index, e.g. `layers.*.scaling`.""" + return Counter( + (re.sub(r"\.\d+\.", ".*.", path), old, new) for path, old, new in diffs + ) + + +before, after = load(sys.argv[1]), load(sys.argv[2]) +for repo, old in before.items(): + new = after.get(repo, {"error": "missing from after"}) + if "error" in old or "error" in new: + print(f"== {repo}\n before: {old.get('error', 'ok')}") + print(f" after: {new.get('error', 'ok')}") + continue + diffs = collapse(diff(old["cfg"], new["cfg"])) + print(f"== {repo}: {old['cls']} -> {new['cls']} ({sum(diffs.values())} diffs)") + for (path, old_value, new_value), count in diffs.items(): + times = f" (x{count})" if count > 1 else "" + print(f" {path}{times}: before={old_value} after={new_value}") diff --git a/.agents/skills/transformers-upgrade/scripts/dump_configs.py b/.agents/skills/transformers-upgrade/scripts/dump_configs.py new file mode 100644 index 000000000000..6d6913fb5178 --- /dev/null +++ b/.agents/skills/transformers-upgrade/scripts/dump_configs.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Dump the config vLLM resolves for each checkpoint, one JSON line per repo. + +Run once per vLLM tree (before and after a change), then compare the two +outputs with diff_dumps.py: + + PYTHONPATH= .venv/bin/python dump_configs.py out.jsonl ... + +Pass --trust-remote-code to resolve checkpoints the way users who set it do. +""" + +import json +import sys + +from vllm.transformers_utils.config import get_config + +out, *args = sys.argv[1:] +trust_remote_code = "--trust-remote-code" in args +with open(out, "w") as f: + for repo in (arg for arg in args if arg != "--trust-remote-code"): + try: + config = get_config(repo, trust_remote_code=trust_remote_code) + cls = f"{type(config).__module__}.{type(config).__name__}" + record = {"repo": repo, "cls": cls, "cfg": config.to_dict()} + except Exception as e: + error = f"{type(e).__name__}: {e}".splitlines()[0] + record = {"repo": repo, "error": error} + f.write(json.dumps(record, default=str, sort_keys=True) + "\n") + f.flush() diff --git a/.agents/skills/transformers-upgrade/scripts/dump_model.py b/.agents/skills/transformers-upgrade/scripts/dump_model.py new file mode 100644 index 000000000000..47efbe873f20 --- /dev/null +++ b/.agents/skills/transformers-upgrade/scripts/dump_model.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Build each model on the meta device and dump what vLLM constructs. + +Records ModelConfig-derived sizes, every parameter shape, and scalar +attributes of every module (head_dim, num_kv_heads, sliding_window, ...). +Needs no GPU and no weights. Run once per vLLM tree, then compare the two +outputs with diff_dumps.py: + + PYTHONPATH= .venv/bin/python dump_model.py out.jsonl ... + +Pass --trust-remote-code to resolve checkpoints the way users who set it do. +""" + +import json +import socket +import sys + +import torch + +from vllm.config import set_current_vllm_config +from vllm.config.vllm import VllmConfig +from vllm.distributed import ( + init_distributed_environment, + initialize_model_parallel, + model_parallel_is_initialized, +) +from vllm.engine.arg_utils import EngineArgs +from vllm.model_executor.model_loader.utils import initialize_model +from vllm.utils.torch_utils import set_default_torch_dtype + +# Construction only: skip model runner checks that need a GPU or Triton +VllmConfig._get_v1_model_runner_unsupported_features = lambda self: [] # type: ignore[method-assign] + +# Model code that moves submodules to a device must leave them on meta +_module_to = torch.nn.Module.to + + +def _meta_safe_to(self, *args, **kwargs): + kwargs.pop("device", None) + args = tuple(a for a in args if not isinstance(a, (torch.device, str, int))) + return _module_to(self, *args, **kwargs) + + +torch.nn.Module.to = _meta_safe_to # type: ignore[method-assign] + +SCALARS = (bool, int, float, str) + + +def module_attrs(model: torch.nn.Module) -> dict[str, dict]: + attrs = {} + for name, module in model.named_modules(): + scalars = { + key: value + for key, value in vars(module).items() + if not key.startswith("_") and isinstance(value, SCALARS) + } + if scalars: + attrs[name or ""] = scalars + return attrs + + +def dump(repo: str, trust_remote_code: bool) -> dict: + args = EngineArgs( + model=repo, + load_format="dummy", + enforce_eager=True, + max_model_len=4096, + trust_remote_code=trust_remote_code, + ) + vllm_config = args.create_engine_config() + model_config = vllm_config.model_config + with set_current_vllm_config(vllm_config): + if not model_parallel_is_initialized(): + with socket.socket() as sock: + sock.bind(("", 0)) + port = sock.getsockname()[1] + init_distributed_environment(1, 0, f"tcp://127.0.0.1:{port}", 0, "gloo") + initialize_model_parallel(1, 1) + with set_default_torch_dtype(model_config.dtype), torch.device("meta"): + model = initialize_model(vllm_config=vllm_config) + return { + "repo": repo, + "cls": type(model).__name__, + "cfg": { + "model_config": { + "head_size": model_config.get_head_size(), + "total_num_kv_heads": model_config.get_total_num_kv_heads(), + "sliding_window": model_config.get_sliding_window(), + "num_layers": model_config.get_num_layers(vllm_config.parallel_config), + }, + "params": {n: list(p.shape) for n, p in model.named_parameters()}, + "modules": module_attrs(model), + }, + } + + +def main() -> None: + out, *args = sys.argv[1:] + trust_remote_code = "--trust-remote-code" in args + repos = [arg for arg in args if arg != "--trust-remote-code"] + with open(out, "w") as f: + for repo in repos: + try: + record = dump(repo, trust_remote_code) + except Exception as e: + error = f"{type(e).__name__}: {e}".splitlines()[0] + record = {"repo": repo, "error": error} + f.write(json.dumps(record, default=str, sort_keys=True) + "\n") + f.flush() + + +if __name__ == "__main__": + main() diff --git a/.agents/skills/transformers-upgrade/scripts/find_upstreamed.py b/.agents/skills/transformers-upgrade/scripts/find_upstreamed.py new file mode 100644 index 000000000000..4cd73c29cbac --- /dev/null +++ b/.agents/skills/transformers-upgrade/scripts/find_upstreamed.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""List vendored vLLM code that the installed Transformers may make redundant. + +Run with the Transformers version under test, from the vLLM repo root: + + uv run --no-project --with transformers== \ + python .agents/skills/transformers-upgrade/scripts/find_upstreamed.py + +Imports nothing from vLLM, so no vLLM install is needed. Uses the `gh` CLI, if +available, to check whether referenced Transformers PRs are in this release. +""" + +import json +import subprocess +import sys +from pathlib import Path + +import regex as re +import transformers +from packaging.version import Version +from transformers.models.auto.configuration_auto import CONFIG_MAPPING_NAMES +from transformers.models.auto.processing_auto import PROCESSOR_MAPPING_NAMES + +ROOT = Path(sys.argv[1] if len(sys.argv) > 1 else ".") +FLOOR = Version(transformers.__version__) +UPSTREAM_NAMES = set(dir(transformers)) + +MODEL_TYPE = re.compile(r'model_type\s*(?::\s*str\s*)?=\s*"([^"]+)"') +REGISTRY_KEY = re.compile(r'^\s+(?:\*\*\{)?"?([\w-]+)"?(?:=|": )"(\w+)"', re.M) +CLASS_KEY = re.compile(r'^\s+"(\w+)": ', re.M) +UPSTREAM_REF = re.compile( + r"huggingface/transformers(?:/(?:issues|pull)/|#)(\d+)" + r"|(?:once|until|when) [Tt]ransformers [\d.]+ is the minimum" +) +VERSIONED = re.compile( + r"(?:min_transformers_version|max_transformers_version|check_version|Version)" + r'\(?\s*=?\s*"([45]\.\d+(?:\.\d+)?)[^"]*"' + r"|[Tt]ransformers\s*(?:v|[<>]=?\s*)([45]\.\d+(?:\.\d+)?|4)\b" +) +READS_VERSION = re.compile( + r"transformers import __version__|transformers\.__version__" + r"|version\([\"']transformers[\"']\)|TRANSFORMERS_VERSION|check_version\(" +) + + +def section(title: str) -> None: + print(f"\n## {title}") + + +def vendored_configs() -> None: + section(f"Vendored configs whose model_type is upstream in {FLOOR}") + files = sorted(ROOT.glob("vllm/transformers_utils/configs/*.py")) + files += sorted(ROOT.glob("vllm/models/*/config*.py")) + files += sorted(ROOT.glob("vllm/models/*/configs/*.py")) + for path in files: + types = sorted(set(MODEL_TYPE.findall(path.read_text()))) + hits = [t for t in types if t in CONFIG_MAPPING_NAMES] + if hits: + upstream = ", ".join(f"{t} -> {CONFIG_MAPPING_NAMES[t]}" for t in hits) + print(f"- {path.relative_to(ROOT)}: {upstream}") + + registry = (ROOT / "vllm/transformers_utils/config.py").read_text() + start = registry.index("_CONFIG_REGISTRY") + body = registry[start : registry.index("\n)\n", start)] + section("_CONFIG_REGISTRY entries that override an upstream model_type") + for model_type, cls in REGISTRY_KEY.findall(body): + if model_type in CONFIG_MAPPING_NAMES: + print(f"- {model_type}: vLLM {cls} over {CONFIG_MAPPING_NAMES[model_type]}") + + +def vendored_processors() -> None: + init = (ROOT / "vllm/transformers_utils/processors/__init__.py").read_text() + names = sorted(set(CLASS_KEY.findall(init))) + section("Registered processor names that collide with upstream classes") + for name in names: + if name in UPSTREAM_NAMES: + print(f"- {name}") + section("Processor files named after a model_type with an upstream processor") + for path in sorted(ROOT.glob("vllm/transformers_utils/processors/*.py")): + if path.stem in PROCESSOR_MAPPING_NAMES: + upstream = PROCESSOR_MAPPING_NAMES[path.stem] + print(f"- {path.relative_to(ROOT)} -> {upstream}") + + +def version_gates() -> None: + files = sorted([*ROOT.glob("vllm/**/*.py"), *ROOT.glob("tests/**/*.py")]) + section(f"Version-gated code at or below {FLOOR} (review each)") + for path in files: + for lineno, line in enumerate(path.read_text().splitlines(), 1): + for match in VERSIONED.finditer(line): + version = Version(match.group(1) or match.group(2)) + if version <= FLOOR: + print(f"- {path.relative_to(ROOT)}:{lineno}: {line.strip()}") + break + section("Code that reads the installed Transformers version (review each)") + for path in files: + for lineno, line in enumerate(path.read_text().splitlines(), 1): + if READS_VERSION.search(line): + print(f"- {path.relative_to(ROOT)}:{lineno}: {line.strip()}") + + +def upstream_status(number: str) -> str: + """Whether Transformers issue/PR `number` is resolved in this release.""" + + def gh(path: str) -> dict | None: + try: + out = subprocess.run(["gh", "api", path], capture_output=True, text=True) + except FileNotFoundError: + return None + return json.loads(out.stdout) if out.returncode == 0 else None + + repo = "repos/huggingface/transformers" + if (issue := gh(f"{repo}/issues/{number}")) is None: + return "status unknown" + if "pull_request" not in issue: + return f"issue {issue['state']}, check which PR fixed it" + pr = gh(f"{repo}/pulls/{number}") + if pr is None or not pr["merged_at"]: + return "PR not merged" + compare = gh(f"{repo}/compare/v{FLOOR}...{pr['merge_commit_sha']}") + included = compare and compare["status"] in ("behind", "identical") + return f"PR merged, {'in' if included else 'not in'} v{FLOOR}" + + +def upstream_refs() -> None: + section("Patches referencing Transformers issues/PRs or a future minimum") + files = sorted([*ROOT.glob("vllm/**/*.py"), *ROOT.glob("tests/**/*.py")]) + statuses: dict[str, str] = {} + for path in files: + for lineno, line in enumerate(path.read_text().splitlines(), 1): + if (match := UPSTREAM_REF.search(line)) is None: + continue + note = "" + if number := match.group(1): + if number not in statuses: + statuses[number] = upstream_status(number) + note = f" [#{number}: {statuses[number]}]" + print(f"- {path.relative_to(ROOT)}:{lineno}:{note} {line.strip()}") + + +if __name__ == "__main__": + print(f"# Transformers {FLOOR}") + vendored_configs() + vendored_processors() + version_gates() + upstream_refs() diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 4ae92315b4c9..a3f7dbd83874 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -237,8 +237,6 @@ steps: - tests/models/language/generation_ppl_test - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - - vllm/transformers_utils/configs/qwen3_5.py - - vllm/transformers_utils/configs/qwen3_5_moe.py - vllm/model_executor/models/qwen2.py - vllm/model_executor/models/qwen3.py - vllm/model_executor/models/qwen3_next.py @@ -2605,7 +2603,6 @@ steps: - vllm/model_executor/models/qwen3_next.py - vllm/model_executor/models/registry.py - vllm/transformers_utils/configs/extract_hidden_states.py - - vllm/transformers_utils/configs/qwen3_5.py - vllm/v1/attention/backends/ - vllm/v1/attention/selector.py - vllm/v1/kv_cache_interface.py @@ -3914,8 +3911,6 @@ steps: source_file_dependencies: - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - - vllm/transformers_utils/configs/qwen3_5.py - - vllm/transformers_utils/configs/qwen3_5_moe.py - vllm/model_executor/models/qwen2.py - vllm/model_executor/models/qwen3.py - vllm/model_executor/models/qwen3_next.py @@ -3957,8 +3952,6 @@ steps: - vllm/model_executor/models/qwen3_next_mtp.py - vllm/model_executor/models/nemotron_h.py - vllm/model_executor/models/nemotron_h_mtp.py - - vllm/transformers_utils/configs/qwen3_next.py - - vllm/transformers_utils/configs/nemotron_h.py - vllm/model_executor/model_loader/ - vllm/model_executor/layers/quantization/ - vllm/model_executor/layers/fused_moe/ @@ -4156,7 +4149,6 @@ steps: - tests/evals/gsm8k/ - tests/utils.py - vllm/distributed/ - - vllm/transformers_utils/configs/diffusion_gemma.py - vllm/v1/attention/ops/ - tests/evals/gsm8k/configs/models-small-tp-mi355.txt commands: @@ -5048,8 +5040,6 @@ steps: source_file_dependencies: - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - - vllm/transformers_utils/configs/qwen3_5.py - - vllm/transformers_utils/configs/qwen3_5_moe.py - vllm/model_executor/models/qwen2.py - vllm/model_executor/models/qwen3.py - vllm/model_executor/models/qwen3_next.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index d1a115a1e1b0..49eea0b78f7c 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -208,7 +208,6 @@ steps: - tests/evals/gsm8k/ - tests/utils.py - vllm/distributed/ - - vllm/transformers_utils/configs/diffusion_gemma.py - vllm/v1/attention/ops/ commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt @@ -417,8 +416,6 @@ steps: - vllm/model_executor/models/qwen3_next_mtp.py - vllm/model_executor/models/nemotron_h.py - vllm/model_executor/models/nemotron_h_mtp.py - - vllm/transformers_utils/configs/qwen3_next.py - - vllm/transformers_utils/configs/nemotron_h.py - vllm/model_executor/model_loader/ - vllm/model_executor/layers/quantization/ - vllm/model_executor/layers/fused_moe/ @@ -445,8 +442,6 @@ steps: source_file_dependencies: - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - - vllm/transformers_utils/configs/qwen3_5.py - - vllm/transformers_utils/configs/qwen3_5_moe.py - vllm/model_executor/models/qwen3_next.py - vllm/model_executor/models/qwen3_next_mtp.py - vllm/third_party/flash_linear_attention/ops/ @@ -463,8 +458,6 @@ steps: source_file_dependencies: - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - - vllm/transformers_utils/configs/qwen3_5.py - - vllm/transformers_utils/configs/qwen3_5_moe.py - vllm/model_executor/models/qwen2.py - vllm/model_executor/models/qwen3.py - vllm/model_executor/models/qwen3_next.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index e1490139c052..52ff1d65b3ee 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -296,7 +296,6 @@ steps: - vllm/model_executor/models/qwen3_next.py - vllm/model_executor/models/registry.py - vllm/transformers_utils/configs/extract_hidden_states.py - - vllm/transformers_utils/configs/qwen3_5.py - vllm/v1/attention/backends/ - vllm/v1/attention/selector.py - vllm/v1/kv_cache_interface.py diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index fa8013ec54e8..f123f2c57b13 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -184,8 +184,6 @@ steps: source_file_dependencies: - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - - vllm/transformers_utils/configs/qwen3_5.py - - vllm/transformers_utils/configs/qwen3_5_moe.py - vllm/model_executor/models/qwen2.py - vllm/model_executor/models/qwen3.py - vllm/model_executor/models/qwen3_next.py diff --git a/.claude/skills/transformers-upgrade b/.claude/skills/transformers-upgrade new file mode 120000 index 000000000000..d71bb1472413 --- /dev/null +++ b/.claude/skills/transformers-upgrade @@ -0,0 +1 @@ +../../.agents/skills/transformers-upgrade \ No newline at end of file diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index ae4c7220f24c..8ec16ee7f9fb 100644 --- a/tests/config/test_model_arch_config.py +++ b/tests/config/test_model_arch_config.py @@ -10,13 +10,11 @@ from unittest.mock import Mock import pytest -from transformers import PreTrainedConfig +from transformers import DiffusionGemmaConfig, PreTrainedConfig from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig from vllm.config.model_arch import ModelArchitectureConfig -from vllm.transformers_utils.configs.diffusion_gemma import DiffusionGemmaConfig -from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config from vllm.transformers_utils.model_arch_config_convertor import ( MODEL_ARCH_CONFIG_CONVERTORS, Gemma4ModelArchConfigConvertor, @@ -318,11 +316,7 @@ def _gemma4_text_config(**overrides) -> Gemma4TextConfig: def test_gemma4_head_dims_vary_by_layer_type(): - """Gemma4's full attention layers are wider than its sliding ones. - - Transformers >= 5.15.0 says so in the config; this exercises the convertor - building the same per-layer view from the flat attributes used before that. - """ + """Gemma4's full attention layers are wider than its sliding ones.""" text_config = _gemma4_text_config( num_global_key_value_heads=8, attention_k_eq_v=True ) @@ -332,9 +326,9 @@ def test_gemma4_head_dims_vary_by_layer_type(): assert (arch.head_size, arch.total_num_kv_heads) == (32, 8) assert [arch[i].head_size for i in range(6)] == [16] * 5 + [32] assert [arch[i].total_num_kv_heads for i in range(6)] == [4] * 5 + [8] - # The model files resolve each layer through the same helper, so the KV cache - # vLLM allocates and the projections the model builds cannot disagree. - assert [gemma4_layer_config(text_config, i).head_dim for i in range(6)] == [ + # The model files read the same per-layer configs, so the KV cache vLLM + # allocates and the projections the model builds cannot disagree. + assert [text_config.per_layer_config[i].head_dim for i in range(6)] == [ arch[i].head_size for i in range(6) ] diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py index 0919431abac6..48b9dd509977 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py @@ -14,7 +14,6 @@ import vllm.envs as envs from tests.utils import RemoteOpenAIServer from vllm.config import ModelConfig -from vllm.config.utils import getattr_iter from vllm.v1.engine.detokenizer import check_stop_strings MODEL_NAME = "Qwen/Qwen3-0.6B" @@ -402,14 +401,7 @@ async def test_same_response_as_chat_completions(client, tokenizer, messages): # Post-EOS generation is undefined and may differ eos_tokens = { tokenizer.eos_token_id, - *getattr_iter( - tokenizer, - [ - "extra_special_tokens_ids", # Transformers v5 - "additional_special_tokens_ids", # Transformers v4 - ], - [], - ), + *tokenizer.extra_special_tokens_ids, } # Find first EOS in generated tokens eos_pos = None diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py index 78aa563925ff..0fcc2540e252 100644 --- a/tests/model_executor/layers/test_fused_shared_expert.py +++ b/tests/model_executor/layers/test_fused_shared_expert.py @@ -12,6 +12,7 @@ import pytest import torch from torch import nn +from transformers import Qwen3_5MoeTextConfig import vllm.config as vllm_config_module from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config @@ -34,7 +35,6 @@ from vllm.models.minimax_m3.amd import model as minimax_m3_model from vllm.platforms import current_platform from vllm.transformers_utils.configs.minimax_m3 import MiniMaxM3TextConfig -from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeTextConfig pytestmark = pytest.mark.skipif( current_platform.is_xpu(), diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index 32a9adda8822..9ea2cc43cd6e 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -9,14 +9,12 @@ from pathlib import PosixPath import pytest -from packaging.version import Version from transformers import ( AutoModel, AutoModelForCausalLM, AutoModelForImageTextToText, AutoModelForTextToWaveform, ) -from transformers import __version__ as TRANSFORMERS_VERSION from vllm.platforms import current_platform from vllm.utils.func_utils import identity @@ -818,8 +816,7 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model): patch_hf_runner=model_utils.paddleocr_vl_patch_hf_runner, image_size_factors=[(0.25,)], marks=[ - pytest.mark.skipif( - Version(TRANSFORMERS_VERSION) >= Version("5.0.0"), + pytest.mark.skip( reason="Model's custom code uses ROPE_INIT_FUNCTIONS" "['default'] which was removed in transformers v5", ), diff --git a/tests/models/multimodal/processing/test_transformers_image.py b/tests/models/multimodal/processing/test_transformers_image.py index 26a3a592ca1c..d45912a4b141 100644 --- a/tests/models/multimodal/processing/test_transformers_image.py +++ b/tests/models/multimodal/processing/test_transformers_image.py @@ -249,9 +249,9 @@ def process(): "BAAI/Emu3-Chat-hf", " and more text", marks=pytest.mark.xfail( - reason="Emu3Processor prepends its BOS token only when images are " - "passed, so the unexpanded prompt vLLM tokenizes never " - "gets one. Fixed by huggingface/transformers#47924, unreleased.", + reason="Emu3Processor prepends the BOS token itself because the " + "tokenizer doesn't, so the unexpanded prompt vLLM tokenizes " + "never gets one.", strict=False, ), ), diff --git a/tests/models/registry.py b/tests/models/registry.py index 4255d439cbdc..31c3c5433aec 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -775,10 +775,7 @@ def check_available_online( # [Decoder-only] "AriaForConditionalGeneration": _HfExamplesInfo("rhymes-ai/Aria"), "AudioFlamingo3ForConditionalGeneration": _HfExamplesInfo( - "nvidia/audio-flamingo-3-hf", - transformers_version_reason={ - "vllm": "Needs https://github.com/huggingface/transformers/pull/43538" - }, + "nvidia/audio-flamingo-3-hf" ), "BagelForConditionalGeneration": _HfExamplesInfo("ByteDance-Seed/BAGEL-7B-MoT"), "BailingMoeV3VLForConditionalGeneration": _HfExamplesInfo( @@ -817,12 +814,7 @@ def check_available_online( "nvidia/Cosmos3-Edge", max_model_len=4096, use_original_num_layers=True, - hf_overrides={ - "text_config": { - "num_hidden_layers": 2, - "hybrid_override_pattern": "*-", - } - }, + hf_overrides={"text_config": {"layers_block_type": ["full_attention", "mlp"]}}, ), "DeepseekVLV2ForCausalLM": _HfExamplesInfo( "deepseek-ai/deepseek-vl2-tiny", diff --git a/tests/plugins_tests/bitsandbytes/test_bitsandbytes.py b/tests/plugins_tests/bitsandbytes/test_bitsandbytes.py index f5845efb0466..92803d93151d 100644 --- a/tests/plugins_tests/bitsandbytes/test_bitsandbytes.py +++ b/tests/plugins_tests/bitsandbytes/test_bitsandbytes.py @@ -6,9 +6,7 @@ from unittest.mock import MagicMock, patch import pytest -from packaging.version import Version from transformers import BitsAndBytesConfig -from transformers import __version__ as TRANSFORMERS_VERSION from vllm_bnb_plugin import bitsandbytes_loader as bnb from vllm.platforms import current_platform @@ -174,8 +172,7 @@ def test_load_pp_4bit_bnb_model( ) -@pytest.mark.skipif( - Version(TRANSFORMERS_VERSION) >= Version("5.0.0"), +@pytest.mark.skip( reason="Need to add support for quantizing MoE experts with bnb in " "transformers v5. See https://github.com/bitsandbytes-foundation/" "bitsandbytes/issues/1849", diff --git a/tests/tokenizers_/test_registry.py b/tests/tokenizers_/test_registry.py index 312b0c03a921..b9d9c732f333 100644 --- a/tests/tokenizers_/test_registry.py +++ b/tests/tokenizers_/test_registry.py @@ -6,7 +6,7 @@ from unittest.mock import patch import pytest -from transformers import AutoConfig +from transformers import AutoConfig, Qwen3_5MoeConfig from transformers.models.auto.configuration_auto import CONFIG_MAPPING from vllm.tokenizers import TokenizerLike @@ -18,7 +18,6 @@ get_tokenizer, resolve_tokenizer_args, ) -from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeConfig class TestTokenizer(TokenizerLike): diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 1e7883864027..516ef132365b 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -8,6 +8,7 @@ import torch from einops import rearrange from torch import nn +from transformers import Qwen3NextConfig from vllm import _custom_ops as ops from vllm import envs @@ -58,7 +59,6 @@ ) from vllm.third_party.flash_linear_attention.ops.chunk import l2norm_fwd from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE -from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig from vllm.utils.torch_utils import ( LayerNameType, _encode_layer_name, diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index d8ceb5083c2f..83859946a255 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -29,6 +29,7 @@ import torch from torch import nn +from transformers import AXK1Config from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.decorators import support_torch_compile @@ -78,7 +79,6 @@ from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.AXK1 import AXK1Config from .interfaces import MixtureOfExperts, SupportsEagle, SupportsLoRA, SupportsPP from .utils import ( diff --git a/vllm/model_executor/models/aimv2.py b/vllm/model_executor/models/aimv2.py index 4e31b99944f1..c9e5dcfdcdd1 100644 --- a/vllm/model_executor/models/aimv2.py +++ b/vllm/model_executor/models/aimv2.py @@ -7,6 +7,7 @@ import torch import torch.nn as nn +from transformers import Aimv2VisionConfig from vllm.distributed import get_tensor_model_parallel_world_size from vllm.distributed.utils import divide @@ -21,20 +22,19 @@ ) from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper -from vllm.transformers_utils.configs.ovis import AIMv2Config class AIMv2SwiGLUFFN(nn.Module): def __init__( self, - config: AIMv2Config, + config: Aimv2VisionConfig, quant_config: QuantizationConfig | None, prefix: str, ): super().__init__() hidden_features = config.intermediate_size in_features = config.hidden_size - bias = config.use_bias + bias = config.mlp_bias self.fc13 = MergedColumnParallelLinear( in_features, @@ -60,7 +60,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class AIMv2PatchEmbed(nn.Module): - def __init__(self, config: AIMv2Config): + def __init__(self, config: Aimv2VisionConfig): super().__init__() self.proj = Conv2dLayer( config.num_channels, @@ -77,7 +77,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class AIMv2ViTPreprocessor(nn.Module): - def __init__(self, config: AIMv2Config): + def __init__(self, config: Aimv2VisionConfig): super().__init__() num_patches = (config.image_size // config.patch_size) ** 2 @@ -95,7 +95,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class AIMv2Attention(nn.Module): def __init__( self, - config: AIMv2Config, + config: Aimv2VisionConfig, quant_config: QuantizationConfig | None, prefix: str, ): @@ -124,7 +124,7 @@ def __init__( self.proj = RowParallelLinear( input_size=self.embed_dim, output_size=self.embed_dim, - bias=config.use_bias, + bias=config.mlp_bias, quant_config=quant_config, prefix=f"{prefix}.proj", ) @@ -151,7 +151,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class AIMv2Block(nn.Module): def __init__( self, - config: AIMv2Config, + config: Aimv2VisionConfig, quant_config: QuantizationConfig | None, prefix: str, ): @@ -174,7 +174,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class AIMv2Transformer(nn.Module): def __init__( self, - config: AIMv2Config, + config: Aimv2VisionConfig, quant_config: QuantizationConfig | None, *, require_post_norm: bool | None = None, @@ -214,7 +214,7 @@ class AIMv2Model(torch.nn.Module): def __init__( self, - config: AIMv2Config, + config: Aimv2VisionConfig, quant_config: QuantizationConfig | None, *, require_post_norm: bool | None = None, diff --git a/vllm/model_executor/models/cohere2_moe.py b/vllm/model_executor/models/cohere2_moe.py index 7c17edc2a1db..3a4d8ef039f7 100644 --- a/vllm/model_executor/models/cohere2_moe.py +++ b/vllm/model_executor/models/cohere2_moe.py @@ -403,20 +403,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config.vocab_size, config.hidden_size ) - # Decoder layers read per-layer MLP layout from config.mlp_layer_types - # (dense MLP vs MoE) and use it for weight loading. Transformers >=5.10 - # populates this field; older versions only expose first_k_dense_replace. - # Normalize here so layer construction below sees a consistent layout. - if getattr(config, "mlp_layer_types", None) is None: - first_k_dense_replace = getattr(config, "first_k_dense_replace", None) - n = config.num_hidden_layers - if first_k_dense_replace is not None: - config.mlp_layer_types = ["dense"] * first_k_dense_replace + [ - "sparse" - ] * (n - first_k_dense_replace) - else: - config.mlp_layer_types = ["sparse"] * n - self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, lambda prefix: Cohere2MoeDecoderLayer( diff --git a/vllm/model_executor/models/cosmos3.py b/vllm/model_executor/models/cosmos3.py index 7e4e82745772..5023266351e5 100644 --- a/vllm/model_executor/models/cosmos3.py +++ b/vllm/model_executor/models/cosmos3.py @@ -2,11 +2,28 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import regex +from transformers import Cosmos3OmniConfig -from vllm.model_executor.models.qwen3_vl import Qwen3VLForConditionalGeneration +from vllm.model_executor.models.qwen3_vl import ( + Qwen3VLDummyInputsBuilder, + Qwen3VLForConditionalGeneration, + Qwen3VLMultiModalProcessor, + Qwen3VLProcessingInfo, +) from vllm.model_executor.models.utils import WeightsMapper +from vllm.multimodal import MULTIMODAL_REGISTRY +class Cosmos3ProcessingInfo(Qwen3VLProcessingInfo): + def get_hf_config(self) -> Cosmos3OmniConfig: + return self.ctx.get_hf_config(Cosmos3OmniConfig) + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen3VLMultiModalProcessor, + info=Cosmos3ProcessingInfo, + dummy_inputs=Qwen3VLDummyInputsBuilder, +) class Cosmos3ForConditionalGeneration(Qwen3VLForConditionalGeneration): # Cosmos3 unified checkpoints store a Qwen3-VL understanding tower # alongside a generation tower in a flat key layout. This mapper drops diff --git a/vllm/model_executor/models/cosmos3_edge.py b/vllm/model_executor/models/cosmos3_edge.py index ae65e25c901d..c10fafd0b192 100644 --- a/vllm/model_executor/models/cosmos3_edge.py +++ b/vllm/model_executor/models/cosmos3_edge.py @@ -5,7 +5,7 @@ import torch import torch.nn as nn -from transformers import ProcessorMixin +from transformers import NemotronHConfig, ProcessorMixin from vllm.config import VllmConfig from vllm.model_executor.layers.layernorm import RMSNorm @@ -20,7 +20,6 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.nemotron_h import NemotronHConfig from vllm.utils.torch_utils import async_tensor_h2d from .interfaces import ( diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 1b27df68abb0..7c64980b57bb 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -16,7 +16,7 @@ from __future__ import annotations -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Iterable, Sequence from types import SimpleNamespace from typing import Any @@ -25,7 +25,7 @@ import torch._dynamo from torch import nn from torch.nn import functional as F -from transformers import AutoModel +from transformers import AutoModel, DiffusionGemmaConfig from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode @@ -113,26 +113,13 @@ def forward( class DiffusionGemmaProcessingInfo(Gemma4ProcessingInfo): - """Processing info for DiffusionGemma. + """Processing info for DiffusionGemma, which has no audio tower.""" - Overrides ``get_hf_config`` to accept ``DiffusionGemmaConfig`` - (which inherits from ``PreTrainedConfig``, not ``Gemma4Config``). - Supports image and video modalities. - """ - - def get_hf_config(self): - # DiffusionGemmaConfig doesn't inherit from Gemma4Config, so we - # accept any PreTrainedConfig here. - return self.ctx.get_hf_config() - - def get_supported_mm_limits(self) -> Mapping[str, int | None]: - # DiffusionGemma supports image and video inputs. - return {"image": None, "video": None} + def get_hf_config(self) -> DiffusionGemmaConfig: + return self.ctx.get_hf_config(DiffusionGemmaConfig) - def get_mm_max_tokens_per_item( - self, seq_len: int, mm_counts: Mapping[str, int] - ) -> Mapping[str, int] | None: - return super().get_mm_max_tokens_per_item(seq_len, mm_counts) + def has_audio(self) -> bool: + return False @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) diff --git a/vllm/model_executor/models/embedding_gemma2.py b/vllm/model_executor/models/embedding_gemma2.py index 150128505d5e..0f92aeaa3716 100644 --- a/vllm/model_executor/models/embedding_gemma2.py +++ b/vllm/model_executor/models/embedding_gemma2.py @@ -36,10 +36,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) -from vllm.model_executor.models.gemma4 import ( - Gemma4MLP, - gemma4_layer_config, -) +from vllm.model_executor.models.gemma4 import Gemma4MLP from vllm.model_executor.models.gemma4_mm import ( _SUPPORTED_SOFT_TOKENS, Gemma4DummyInputsBuilder, @@ -244,7 +241,7 @@ def __init__( prefix: str, ): super().__init__() - lc = gemma4_layer_config(config, idx) + lc = config.per_layer_config[idx] self.self_attn = EmbeddingGemma2Attention( config, lc, diff --git a/vllm/model_executor/models/gemma3.py b/vllm/model_executor/models/gemma3.py index e886c2ccdd07..4646ebeca462 100644 --- a/vllm/model_executor/models/gemma3.py +++ b/vllm/model_executor/models/gemma3.py @@ -154,20 +154,7 @@ def __init__( self.is_sliding = layer_type == "sliding_attention" sliding_window = config.sliding_window if self.is_sliding else None - # Initialize the rotary embedding. - if layer_type in config.rope_parameters: - # Transformers v5 rope config. - rope_parameters = config.rope_parameters[layer_type] - else: - # Transformers v4 rope config. - # Global attention. Use the values in config.json. - rope_parameters = config.rope_parameters - # Local attention. Override the values in config.json. - if self.is_sliding: - rope_parameters = dict( - rope_type="default", rope_theta=config.rope_local_base_freq - ) - + rope_parameters = config.rope_parameters[layer_type] self.rotary_emb = get_rope( self.head_dim, max_position=max_position_embeddings, diff --git a/vllm/model_executor/models/gemma3n.py b/vllm/model_executor/models/gemma3n.py index 3c9df39e674c..9e7c26bf59b9 100644 --- a/vllm/model_executor/models/gemma3n.py +++ b/vllm/model_executor/models/gemma3n.py @@ -350,17 +350,7 @@ def __init__( is_sliding = layer_type == "sliding_attention" self.sliding_window = config.sliding_window if is_sliding else None - # Initialize the rotary embedding. - if layer_type in config.rope_parameters: - # Transformers v5 rope config. - rope_parameters = config.rope_parameters[layer_type] - else: - # Transformers v4 rope config. - # Global attention. Use the values in config.json. - rope_parameters = config.rope_parameters.copy() - # Local attention. Override the values in config.json. - if is_sliding: - rope_parameters["rope_theta"] = config.rope_local_base_freq + rope_parameters = config.rope_parameters[layer_type] kv_sharing_target_layer_name = None if self.is_kv_shared_layer: diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index 17b8251e9a45..d1498f16ddb4 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -59,7 +59,6 @@ ) from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import KVSharingFastPrefillMetadata @@ -576,7 +575,7 @@ def __init__( self.layer_idx = layer_idx # Gemma4 uses different head dimensions for sliding vs full attention - layer_config = gemma4_layer_config(config, layer_idx) + layer_config = config.per_layer_config[layer_idx] head_dim = layer_config.head_dim num_kv_heads = layer_config.num_key_value_heads @@ -628,9 +627,7 @@ def __init__( ) # MoE (Mixture of Experts) — router + expert block parallel to MLP - self.enable_moe_block = getattr(config, "enable_moe_block", False) or getattr( - config, "use_second_mlp_block", False - ) + self.enable_moe_block = getattr(config, "num_experts", None) is not None self.router: Gemma4Router | None self.experts: MoERunner | None self.post_feedforward_layernorm_1: RMSNorm | None diff --git a/vllm/model_executor/models/gemma4_dspark.py b/vllm/model_executor/models/gemma4_dspark.py index 33b1a1aada23..13c868e59b38 100644 --- a/vllm/model_executor/models/gemma4_dspark.py +++ b/vllm/model_executor/models/gemma4_dspark.py @@ -20,7 +20,6 @@ VocabParallelEmbedding, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config from .gemma4_mtp import Gemma4MTPAttention, Gemma4MTPDecoderLayer from .qwen3_dflash import DFlashQwen3Model, _dflash_layer_causal @@ -48,7 +47,7 @@ def __init__( config, "attention_k_eq_v", False ) - layer_config = gemma4_layer_config(config, layer_idx) + layer_config = config.per_layer_config[layer_idx] head_dim = layer_config.head_dim num_kv_heads = layer_config.num_key_value_heads super().__init__( diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index 796dfbb6e8fa..204db6ce8d51 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -222,12 +222,11 @@ def get_hf_processor(self, **kwargs: object) -> Gemma4Processor: **kwargs, ) + def has_audio(self) -> bool: + return self.get_hf_config().audio_config is not None + def validate_num_items(self, modality: str, num_items: int) -> None: - if ( - modality == "audio" - and num_items > 0 - and self.get_hf_config().audio_config is None - ): + if modality == "audio" and num_items > 0 and not self.has_audio(): model_config = self.ctx.model_config model = get_served_model_name( model_config.model, model_config.served_model_name @@ -243,7 +242,7 @@ def validate_num_items(self, modality: str, num_items: int) -> None: def get_supported_mm_limits(self) -> Mapping[str, int | None]: limits: dict[str, int | None] = {"image": None} - if self.get_hf_config().audio_config is not None: + if self.has_audio(): limits["audio"] = None limits["video"] = None return limits @@ -261,7 +260,7 @@ def get_mm_max_tokens_per_item( if isinstance(val, int) and val in _SUPPORTED_SOFT_TOKENS: tokens_per_image = val tokens: dict[str, int] = {"image": tokens_per_image} - if config.audio_config is not None: + if self.has_audio(): # Audio max tokens from the processor's audio_seq_length. processor = self.get_hf_processor() tokens["audio"] = processor.audio_seq_length @@ -275,9 +274,8 @@ def get_mm_max_tokens_per_item( return tokens def get_data_parser(self) -> MultiModalDataParser: - config = self.get_hf_config() kwargs: dict[str, Any] = {"video_needs_metadata": True} - if getattr(config, "audio_config", None) is not None: + if self.has_audio(): processor = self.get_hf_processor() kwargs["target_sr"] = processor.feature_extractor.sampling_rate return MultiModalDataParser(**kwargs) diff --git a/vllm/model_executor/models/gemma4_mtp.py b/vllm/model_executor/models/gemma4_mtp.py index 01f226ff07a5..f9931ddafc0a 100644 --- a/vllm/model_executor/models/gemma4_mtp.py +++ b/vllm/model_executor/models/gemma4_mtp.py @@ -45,7 +45,6 @@ VocabParallelEmbedding, ) from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config from .gemma4 import Gemma4MLP, _get_text_config from .utils import ( @@ -265,7 +264,7 @@ def __init__( self.hidden_size = config.hidden_size layer_idx = extract_layer_index(prefix) - layer_config = gemma4_layer_config(config, layer_idx) + layer_config = config.per_layer_config[layer_idx] head_dim = layer_config.head_dim num_kv_heads = layer_config.num_key_value_heads diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py index b62a86543414..d738c665c2ca 100644 --- a/vllm/model_executor/models/granite4_vision.py +++ b/vllm/model_executor/models/granite4_vision.py @@ -493,13 +493,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: prefix=maybe_prefix(prefix, "vision_tower"), ) - # image_newline parameter - if config.use_image_newline_parameter: - self.image_newline = nn.Parameter( - torch.empty(config.text_config.hidden_size) - ) - else: - self.image_newline = None + self.image_newline = nn.Parameter( + torch.empty(config.text_config.hidden_size) + ) cache_config = vllm_config.cache_config @@ -517,20 +513,18 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ) # Spatial projectors: 4 offset groups - self.spatial_projectors = None - if config.use_spatial_sampling: - self.spatial_projectors = nn.ModuleList( - [ - WindowQFormerDownsampler( - config, - quant_config=quant_config, - cache_config=cache_config, - spatial_offset=i, - prefix=maybe_prefix(prefix, f"spatial_projectors.{i}"), - ) - for i in range(4) - ] - ) + self.spatial_projectors = nn.ModuleList( + [ + WindowQFormerDownsampler( + config, + quant_config=quant_config, + cache_config=cache_config, + spatial_offset=i, + prefix=maybe_prefix(prefix, f"spatial_projectors.{i}"), + ) + for i in range(4) + ] + ) # ----- Language model (marked as LM) ----- with self._mark_language_model(vllm_config): @@ -545,12 +539,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: # Store config values we need self._deepstack_layer_map = config.deepstack_layer_map # [[-19, 9], ...] - self._use_spatial_sampling = getattr(config, "use_spatial_sampling", False) - self._spatial_vision_layer = getattr(config, "spatial_vision_layer", -1) - self._spatial_target_layers = getattr(config, "spatial_target_layers", []) - self._vision_feature_select_strategy = getattr( - config, "vision_feature_select_strategy", "full" - ) + self._spatial_vision_layer = config.spatial_vision_layer + self._spatial_target_layers = config.spatial_target_layers + self._vision_feature_select_strategy = config.vision_feature_select_strategy self._downsample_rate = Fraction(config.downsample_rate) # Ordered list of LLM layer indices for each deepstack level. @@ -558,7 +549,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: # (before any embed_multimodal call). self._ds_layer_indices: list[int] = [ llm_layer for _, llm_layer in config.deepstack_layer_map - ] + list(getattr(config, "spatial_target_layers", [])) + ] + list(config.spatial_target_layers) # Share ds_layer_indices with the LLM causal model so # make_empty_intermediate_tensors includes the correct keys @@ -571,7 +562,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: # the updated values written just before each prefill. # Shape: (max_num_batched_tokens, lm_hidden_size) per level. n_layerwise = len(config.deepstack_layer_map) - n_spatial = len(getattr(config, "spatial_target_layers", [])) + n_spatial = len(config.spatial_target_layers) num_ds_levels = n_layerwise + n_spatial lm_hidden = config.text_config.hidden_size max_tokens = vllm_config.scheduler_config.max_num_batched_tokens @@ -644,26 +635,24 @@ def _pack_and_unpad_image_features( ) image_feature = unpad_image(image_feature, image_sizes[image_idx]) - if self.image_newline is not None: - image_feature = torch.cat( - ( - image_feature, - self.image_newline[:, None, None] - .expand(*image_feature.shape[:-1], 1) - .to(image_feature.device, image_feature.dtype), - ), - dim=-1, - ) + image_feature = torch.cat( + ( + image_feature, + self.image_newline[:, None, None] + .expand(*image_feature.shape[:-1], 1) + .to(image_feature.device, image_feature.dtype), + ), + dim=-1, + ) image_feature = image_feature.flatten(1, 2).transpose(0, 1) image_feature = torch.cat((base_image_feature, image_feature), dim=0) else: image_feature = image_feature[0] - if self.image_newline is not None: - image_feature = torch.cat( - (image_feature, self.image_newline[None].to(image_feature)), - dim=0, - ) + image_feature = torch.cat( + (image_feature, self.image_newline[None].to(image_feature)), + dim=0, + ) new_image_features.append(image_feature) @@ -718,16 +707,15 @@ def _get_all_layer_features( ) levels.append((llm_layer, per_image)) - if self._use_spatial_sampling and self.spatial_projectors is not None: - spatial_hidden = all_hidden_states[self._spatial_vision_layer] - if select_strategy == "default": - spatial_hidden = spatial_hidden[:, 1:] - for group_idx, llm_layer in enumerate(self._spatial_target_layers): - projected = self.spatial_projectors[group_idx](spatial_hidden) - per_image = self._pack_and_unpad_image_features( - torch.split(projected, image_num_patches, dim=0), image_sizes - ) - levels.append((llm_layer, per_image)) + spatial_hidden = all_hidden_states[self._spatial_vision_layer] + if select_strategy == "default": + spatial_hidden = spatial_hidden[:, 1:] + for group_idx, llm_layer in enumerate(self._spatial_target_layers): + projected = self.spatial_projectors[group_idx](spatial_hidden) + per_image = self._pack_and_unpad_image_features( + torch.split(projected, image_num_patches, dim=0), image_sizes + ) + levels.append((llm_layer, per_image)) llm_layer_indices = [llm_layer for llm_layer, _ in levels] num_images = len(image_sizes) diff --git a/vllm/model_executor/models/granite_speech.py b/vllm/model_executor/models/granite_speech.py index 3ae85cfaac69..2380bba10a65 100644 --- a/vllm/model_executor/models/granite_speech.py +++ b/vllm/model_executor/models/granite_speech.py @@ -158,9 +158,7 @@ def _get_prompt_updates( feature_extractor = processor.audio_processor vocab = tokenizer.get_vocab() - # Use getattr with default to be compatible with transformers<4.48 - audio_token = getattr(processor, "audio_token", "<|audio|>") - audio_token_id = vocab[audio_token] + audio_token_id = vocab[processor.audio_token] def get_replacement(item_idx: int): audios = mm_items.get_items("audio", AudioProcessorItems) @@ -221,8 +219,7 @@ def get_dummy_mm_data( def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: num_audios = mm_counts.get("audio", 0) hf_processor = self.info.get_hf_processor() - audio_token = getattr(hf_processor, "audio_token", "<|audio|>") - return audio_token * num_audios + return hf_processor.audio_token * num_audios ### QFormer Projector diff --git a/vllm/model_executor/models/granitemoehybrid.py b/vllm/model_executor/models/granitemoehybrid.py index 84a981e6bba2..979d272c13d1 100644 --- a/vllm/model_executor/models/granitemoehybrid.py +++ b/vllm/model_executor/models/granitemoehybrid.py @@ -320,10 +320,6 @@ def forward( ALL_DECODER_LAYER_TYPES = { - # Transformers < 5.13.0 - "attention": GraniteMoeHybridAttentionDecoderLayer, - "mamba": GraniteMoeHybridMambaDecoderLayer, - # Transformers >= 5.13.0 "full_attention": GraniteMoeHybridAttentionDecoderLayer, "linear_attention": GraniteMoeHybridMambaDecoderLayer, } diff --git a/vllm/model_executor/models/hy_v3.py b/vllm/model_executor/models/hy_v3.py index 2e24cf65e8f3..80fc16e23202 100644 --- a/vllm/model_executor/models/hy_v3.py +++ b/vllm/model_executor/models/hy_v3.py @@ -31,7 +31,7 @@ import torch from torch import nn -from transformers import PreTrainedConfig +from transformers import HYV3Config, PreTrainedConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config @@ -67,7 +67,6 @@ maybe_remap_kv_scale_name, ) from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.hy_v3 import HYV3Config from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP from .utils import ( diff --git a/vllm/model_executor/models/hyperclovax.py b/vllm/model_executor/models/hyperclovax.py index 66c3b4a270a3..e86260bca96f 100644 --- a/vllm/model_executor/models/hyperclovax.py +++ b/vllm/model_executor/models/hyperclovax.py @@ -31,6 +31,7 @@ import torch from torch import nn +from transformers import HyperCLOVAXConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig @@ -51,7 +52,6 @@ VocabParallelEmbedding, ) from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.hyperclovax import HyperCLOVAXConfig from .interfaces import SupportsLoRA, SupportsPP from .utils import ( @@ -137,9 +137,7 @@ def __init__( # the KV heads across multiple tensor parallel GPUs. assert tp_size % self.total_num_kv_heads == 0 self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.head_dim = getattr( - config, "head_dim", self.hidden_size // self.total_num_heads - ) + self.head_dim = config.head_dim self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = config.attention_multiplier @@ -166,7 +164,7 @@ def __init__( self.head_dim, max_position=max_position_embeddings, is_neox_style=True, - rope_parameters=getattr(config, "rope_parameters", None), + rope_parameters=config.rope_parameters, dual_chunk_attention_config=dual_chunk_attention_config, ) @@ -217,18 +215,15 @@ def __init__( "dual_chunk_attention_config", None, ) - attention_bias = getattr(config, "attention_bias", False) self.self_attn = HyperCLOVAXAttention( config=config, hidden_size=self.hidden_size, num_heads=config.num_attention_heads, - num_kv_heads=getattr( - config, "num_key_value_heads", config.num_attention_heads - ), + num_kv_heads=config.num_key_value_heads, max_position_embeddings=max_position_embeddings, quant_config=quant_config, - bias=attention_bias, + bias=config.attention_bias, cache_config=cache_config, prefix=f"{prefix}.self_attn", dual_chunk_attention_config=dual_chunk_attention_config, @@ -238,7 +233,7 @@ def __init__( intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, - bias=getattr(config, "mlp_bias", False), + bias=config.mlp_bias, prefix=f"{prefix}.mlp", ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) diff --git a/vllm/model_executor/models/kimi_k25.py b/vllm/model_executor/models/kimi_k25.py index 8cd354b903af..e555b6368d4a 100644 --- a/vllm/model_executor/models/kimi_k25.py +++ b/vllm/model_executor/models/kimi_k25.py @@ -11,7 +11,7 @@ import torch from torch import nn -from transformers import BatchFeature +from transformers import BatchFeature, Kimi_K25Config from vllm.config import VllmConfig from vllm.config.multimodal import MultiModalDummyOptions @@ -63,7 +63,6 @@ ) from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.kimi_k25 import KimiK25Config from vllm.transformers_utils.processor import cached_get_image_processor from vllm.transformers_utils.processors.kimi_k25 import KimiK25Processor from vllm.transformers_utils.processors.kimi_k25_vision_fused import ( @@ -135,7 +134,7 @@ def __init__(self, ctx: InputProcessingContext) -> None: # Resolve token ID from the tokenizer because transformers v5 # may remap token IDs vs config.json. - config_token_id = hf_config.media_placeholder_token_id + config_token_id = hf_config.image_token_id resolved_token_id = tokenizer.convert_tokens_to_ids("<|media_pad|>") unk_token_id = getattr(tokenizer, "unk_token_id", None) is_valid_resolved = isinstance(resolved_token_id, int) and ( @@ -143,7 +142,7 @@ def __init__(self, ctx: InputProcessingContext) -> None: ) if is_valid_resolved and resolved_token_id != config_token_id: logger.warning_once( - "Kimi-K2.5 config.media_placeholder_token_id (%d) disagrees " + "Kimi-K2.5 config.image_token_id (%d) disagrees " "with tokenizer mapping for <|media_pad|> (%d). " "Using tokenizer value.", config_token_id, @@ -151,7 +150,7 @@ def __init__(self, ctx: InputProcessingContext) -> None: ) media_token_id = resolved_token_id # Patch config so downstream code also sees the correct ID. - hf_config.media_placeholder_token_id = resolved_token_id + hf_config.image_token_id = resolved_token_id else: media_token_id = config_token_id @@ -170,7 +169,7 @@ def get_hf_processor(self, **kwargs: object): return self.hf_processor def get_hf_config(self): - return self.ctx.get_hf_config(KimiK25Config) + return self.ctx.get_hf_config(Kimi_K25Config) def get_supported_mm_limits(self) -> Mapping[str, int | None]: # None means unlimited @@ -338,7 +337,7 @@ def __init__( ) -> None: super().__init__() model_config = vllm_config.model_config - config: KimiK25Config = model_config.hf_config + config: Kimi_K25Config = model_config.hf_config self.config = config quant_config = vllm_config.quant_config @@ -347,7 +346,6 @@ def __init__( ) self.hidden_size = config.text_config.hidden_size self.device = current_platform.current_device() - # Build vision tower directly with KimiK25VisionConfig with self._mark_tower_model(vllm_config, "vision_chunk"): self.vision_tower = MoonViT3dPretrainedModel( config.vision_config, @@ -364,6 +362,7 @@ def __init__( self.mm_projector = KimiK25MultiModalProjector( config=config.vision_config, + out_hidden_size=config.text_config.hidden_size, use_data_parallel=self.use_data_parallel, quant_config=self._maybe_ignore_quant_config(quant_config), prefix=maybe_prefix(prefix, "mm_projector"), @@ -383,7 +382,7 @@ def __init__( self.make_empty_intermediate_tensors = ( self.language_model.make_empty_intermediate_tensors ) - self.media_placeholder: int = self.config.media_placeholder_token_id + self.media_placeholder: int = self.config.image_token_id def _maybe_ignore_quant_config(self, quant_config: QuantizationConfig | None): if isinstance(quant_config, compressed_tensors.CompressedTensorsConfig): diff --git a/vllm/model_executor/models/kimi_k25_vit.py b/vllm/model_executor/models/kimi_k25_vit.py index 44e8498d212b..bb1be1ba32f8 100644 --- a/vllm/model_executor/models/kimi_k25_vit.py +++ b/vllm/model_executor/models/kimi_k25_vit.py @@ -15,6 +15,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from transformers import Kimi_K25VisionConfig from transformers.activations import GELUActivation from vllm.distributed import divide, get_tensor_model_parallel_world_size @@ -40,13 +41,21 @@ run_dp_sharded_mrope_vision_model, ) from vllm.platforms import current_platform -from vllm.transformers_utils.configs.kimi_k25 import KimiK25VisionConfig +from vllm.transformers_utils.configs.kimi_k3 import KimiK3VisionConfig from vllm.triton_utils import HAS_TRITON from vllm.utils.torch_utils import async_tensor_h2d logger = init_logger(__name__) +def _get_pos_emb_size( + config: Kimi_K25VisionConfig | KimiK3VisionConfig, dim: str +) -> int: + # Legacy checkpoints and Kimi-K3 use `init_pos_emb_*` + legacy = getattr(config, f"init_pos_emb_{dim}", None) + return legacy if legacy is not None else getattr(config, f"pos_emb_{dim}") + + def _apply_rope_input_validation(x, freqs_cis): assert x.ndim == freqs_cis.ndim + 1, (x.shape, freqs_cis.shape) assert x.shape[:-2] == freqs_cis.shape[:-1], (x.shape, freqs_cis.shape) @@ -698,14 +707,11 @@ def tpool_patch_merger_packed( class MoonViT3dPretrainedModel(nn.Module): - """Main vision tower model. - - Uses KimiK25VisionConfig directly from transformers_utils/configs/kimi_k25.py. - """ + """Main vision tower model.""" def __init__( self, - config: KimiK25VisionConfig, + config: Kimi_K25VisionConfig | KimiK3VisionConfig, quant_config: QuantizationConfig | None = None, input_norm: nn.Module | None = None, prefix: str = "", @@ -713,17 +719,18 @@ def __init__( super().__init__() config = deepcopy(config) self.config = config # Required for run_dp_sharded_mrope_vision_model - self.merge_kernel_size = config.merge_kernel_size + merge_h, merge_w = config.merge_kernel_size + self.merge_kernel_size = (merge_h, merge_w) self.patch_size = config.patch_size - self.merge_type = config.merge_type + self.merge_type = getattr(config, "merge_type", "sd2_tpool") self.patch_embed = MoonVision3dPatchEmbed( out_dim=config.hidden_size, patch_size=config.patch_size, - pos_emb_height=config.init_pos_emb_height, - pos_emb_width=config.init_pos_emb_width, - pos_emb_time=config.init_pos_emb_time, - pos_emb_type=config.pos_emb_type, + pos_emb_height=_get_pos_emb_size(config, "height"), + pos_emb_width=_get_pos_emb_size(config, "width"), + pos_emb_time=_get_pos_emb_size(config, "time"), + pos_emb_type=getattr(config, "pos_emb_type", "divided_fixed"), patch_embed_proj_bias=getattr(config, "patch_embed_proj_bias", True), pos_emb_interpolation_mode=getattr( config, "pos_emb_interpolation_mode", "bicubic" @@ -740,14 +747,14 @@ def __init__( "qkv_hidden_size": getattr(config, "qkv_hidden_size", None), "mlp_dim": config.intermediate_size, "activation": get_act_fn( - getattr(config, "activation_func", "gelu_pytorch_tanh") + getattr(config, "activation_func", None) or config.hidden_act ), "attn_bias": getattr(config, "attn_bias", True), "norm_type": getattr(config, "norm_type", "layernorm"), "mlp_type": getattr(config, "mlp_type", "mlp2"), "linear_bias": getattr(config, "linear_bias", True), }, - video_attn_type=config.video_attn_type, + video_attn_type=getattr(config, "video_attn_type", "spatial_temporal"), quant_config=quant_config, prefix=maybe_prefix(prefix, "encoder"), ) @@ -886,7 +893,8 @@ class KimiK25MultiModalProjector(nn.Module): def __init__( self, - config: KimiK25VisionConfig, + config: Kimi_K25VisionConfig | KimiK3VisionConfig, + out_hidden_size: int, use_data_parallel: bool = False, quant_config: QuantizationConfig | None = None, prefix: str = "", @@ -909,13 +917,13 @@ def __init__( ) self.linear_2 = ReplicatedLinear( self.hidden_size, - getattr(config, "text_hidden_size", config.mm_hidden_size), + out_hidden_size, bias=False, quant_config=quant_config, prefix=f"{prefix}.linear_2", ) self.post_norm = torch.nn.RMSNorm( - getattr(config, "text_hidden_size", config.mm_hidden_size), + out_hidden_size, eps=config.projector_ln_eps, ) self.act = GELUActivation() @@ -931,7 +939,7 @@ def __init__( ) self.linear_2 = ReplicatedLinear( self.hidden_size, - config.mm_hidden_size, + out_hidden_size, bias=True, quant_config=quant_config, prefix=f"{prefix}.linear_2", diff --git a/vllm/model_executor/models/lfm2_moe.py b/vllm/model_executor/models/lfm2_moe.py index f6e819c2639e..66c9520bc935 100644 --- a/vllm/model_executor/models/lfm2_moe.py +++ b/vllm/model_executor/models/lfm2_moe.py @@ -5,6 +5,7 @@ import torch import torch.nn as nn +from transformers import Lfm2MoeConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config @@ -40,7 +41,6 @@ VocabParallelEmbedding, ) from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.lfm2_moe import Lfm2MoeConfig from .interfaces import ( HasInnerState, diff --git a/vllm/model_executor/models/modernbert.py b/vllm/model_executor/models/modernbert.py index 3a8abee67ec3..0c1909445b48 100644 --- a/vllm/model_executor/models/modernbert.py +++ b/vllm/model_executor/models/modernbert.py @@ -94,28 +94,12 @@ def __init__( prefix=f"{prefix}.Wqkv", ) - if layer_types := getattr(config, "layer_types", None): - # Transformers v5 - layer_type = layer_types[layer_id] - rope_parameters = config.rope_parameters[layer_type] - sliding_window: int | None = None - if layer_type == "sliding_attention": - # Treats the local attention boundary as inclusive - sliding_window = config.sliding_window + 1 - else: - # Transformers v4 - sliding_window = None - if layer_id % config.global_attn_every_n_layers != 0: - # ModernBertConfig does not expose sliding_window - sliding_window = config.local_attention // 2 - rope_theta = ( - config.local_rope_theta - if config.local_rope_theta is not None - else config.global_rope_theta - ) - else: - rope_theta = config.global_rope_theta - rope_parameters = {"rope_type": "default", "rope_theta": rope_theta} + layer_type = config.layer_types[layer_id] + rope_parameters = config.rope_parameters[layer_type] + sliding_window: int | None = None + if layer_type == "sliding_attention": + # Treats the local attention boundary as inclusive + sliding_window = config.sliding_window + 1 self.rotary_emb = get_rope( head_size=self.head_dim, diff --git a/vllm/model_executor/models/nemotron.py b/vllm/model_executor/models/nemotron.py index a271432813a3..55218576969d 100644 --- a/vllm/model_executor/models/nemotron.py +++ b/vllm/model_executor/models/nemotron.py @@ -29,6 +29,7 @@ import torch from torch import nn +from transformers import NemotronConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig @@ -49,7 +50,6 @@ VocabParallelEmbedding, ) from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.nemotron import NemotronConfig from .interfaces import SupportsLoRA, SupportsPP from .utils import ( diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index e477c65302b6..ae653950ac48 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -23,6 +23,7 @@ import torch from torch import nn +from transformers import NemotronHConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import ( @@ -87,7 +88,6 @@ sequence_parallel_chunk, ) from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.nemotron_h import NemotronHConfig class NemotronHMLP(nn.Module): diff --git a/vllm/model_executor/models/nemotron_h_mtp.py b/vllm/model_executor/models/nemotron_h_mtp.py index af93cc62ae25..1e393296bc97 100644 --- a/vllm/model_executor/models/nemotron_h_mtp.py +++ b/vllm/model_executor/models/nemotron_h_mtp.py @@ -7,6 +7,7 @@ import torch import torch.nn as nn +from transformers import NemotronHConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, ModelConfig, VllmConfig @@ -33,7 +34,6 @@ maybe_prefix, ) from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.nemotron_h import NemotronHConfig from .interfaces import SupportsPP, SupportsQuant from .nemotron_h import ( diff --git a/vllm/model_executor/models/ovis.py b/vllm/model_executor/models/ovis.py index 7c1668ddabcd..9e7778da1589 100644 --- a/vllm/model_executor/models/ovis.py +++ b/vllm/model_executor/models/ovis.py @@ -125,7 +125,7 @@ def _init_backbone( prefix: str = "", ) -> nn.Module: model_type = config.backbone_config.model_type - if model_type == "aimv2": + if model_type == "aimv2_vision_model": # No post rms_norm in Ovis2's AIMv2 ViT. return AIMv2Model( config=config.backbone_config, diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index 32660be5150f..0d110dd58be1 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -29,6 +29,12 @@ import torch from torch import nn +from transformers import ( + Qwen3_5Config, + Qwen3_5MoeConfig, + Qwen3_5MoeTextConfig, + Qwen3_5TextConfig, +) from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig @@ -59,11 +65,6 @@ from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.sequence import IntermediateTensors from vllm.tokenizers.registry import cached_tokenizer_from_config -from vllm.transformers_utils.configs.qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig -from vllm.transformers_utils.configs.qwen3_5_moe import ( - Qwen3_5MoeConfig, - Qwen3_5MoeTextConfig, -) from .interfaces import ( HasInnerState, diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py index d7fb1d6e09ea..2acac6b5b7c5 100644 --- a/vllm/model_executor/models/qwen3_5_mtp.py +++ b/vllm/model_executor/models/qwen3_5_mtp.py @@ -6,6 +6,7 @@ import torch from torch import nn +from transformers import Qwen3_5MoeTextConfig, Qwen3_5TextConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig @@ -32,8 +33,6 @@ ) from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.qwen3_5 import Qwen3_5TextConfig -from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeTextConfig from .interfaces import ( MultiModalEmbeddings, diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 9580894e22ee..b0bf0fecd19c 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -7,6 +7,7 @@ import torch from torch import nn +from transformers import Qwen3NextConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, ModelConfig, VllmConfig @@ -58,7 +59,6 @@ from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig from vllm.v1.attention.backend import AttentionType from .interfaces import ( diff --git a/vllm/model_executor/models/qwen3_next_mtp.py b/vllm/model_executor/models/qwen3_next_mtp.py index ef08878eaa7f..4ae8f21840a1 100644 --- a/vllm/model_executor/models/qwen3_next_mtp.py +++ b/vllm/model_executor/models/qwen3_next_mtp.py @@ -6,6 +6,7 @@ import torch from torch import nn +from transformers import Qwen3NextConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig @@ -29,7 +30,6 @@ ) from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig from .utils import ( AutoWeightsLoader, diff --git a/vllm/models/kimi_k3/amd/model.py b/vllm/models/kimi_k3/amd/model.py index 227f6dba2487..124fca3e337a 100644 --- a/vllm/models/kimi_k3/amd/model.py +++ b/vllm/models/kimi_k3/amd/model.py @@ -117,6 +117,7 @@ def __init__( self.mm_projector = KimiK25MultiModalProjector( config=config.vision_config, + out_hidden_size=config.text_config.hidden_size, use_data_parallel=self.use_data_parallel, quant_config=self._maybe_ignore_quant_config(quant_config), prefix=maybe_prefix(prefix, "mm_projector"), diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 22d2c4a1a9a1..eb2ac88aa436 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -1825,6 +1825,7 @@ def __init__( self.mm_projector = KimiK25MultiModalProjector( config=config.vision_config, + out_hidden_size=config.text_config.hidden_size, use_data_parallel=self.use_data_parallel, quant_config=self._maybe_ignore_quant_config(quant_config), prefix=maybe_prefix(prefix, "mm_projector"), diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index c22f6fc3c2ea..e046c301d23a 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -79,8 +79,6 @@ def __getitem__(self, key): _CONFIG_REGISTRY: dict[str, type[PreTrainedConfig]] = LazyConfigDict( - afmoe="AfmoeConfig", - axk1="AXK1Config", bagel="BagelConfig", bailing_moe_v3_vl="BailingMoeV3VLConfig", chatglm="ChatGLMConfig", @@ -89,24 +87,20 @@ def __getitem__(self, key): colqwen3="ColQwen3Config", ops_colqwen3="OpsColQwen3Config", qwen3_vl_nemotron_embed="Qwen3VLNemotronEmbedConfig", - cosmos3_omni="Cosmos3Config", cosmos3_edge="Cosmos3EdgeConfig", - diffusion_gemma="DiffusionGemmaConfig", deepseek_vl_v2="DeepseekVLV2Config", deepseek_v32="DeepseekV3Config", deepseek_v4="DeepseekV4Config", dots3_note="Dots3NoteConfig", k3_dspark="K3DSparkConfig", funaudiochat="FunAudioChatConfig", - granite4_vision="Granite4VisionConfig", - hyperclovax="HyperCLOVAXConfig", - hy_v3="HYV3Config", hy_v4="HYV4Config", + hyperclovax="HyperCLOVAXConfig", # Upstream class, hub remote code is broken isaac="IsaacConfig", kimi_k2="DeepseekV3Config", # Kimi K2 uses same architecture as DeepSeek V3 kimi_linear="KimiLinearConfig", kimi_vl="KimiVLConfig", - kimi_k25="KimiK25Config", + kimi_k25="Kimi_K25Config", # Upstream class, hub remote code uses old schema muse_glimmer="MuseGlimmerConfig", muse_glimmer_text="MuseGlimmerTextConfig", muse_glimmer_vision="MuseGlimmerVisionConfig", @@ -123,8 +117,6 @@ def __getitem__(self, key): moss_transcribe_diarize="MossTranscribeDiarizeConfig", eagle="EAGLEConfig", speculators="SpeculatorsConfig", - nemotron="NemotronConfig", - olmo_hybrid="OlmoHybridConfig", openvla="OpenVLAConfig", ovis="OvisConfig", ultravox="UltravoxConfig", @@ -133,12 +125,6 @@ def __getitem__(self, key): step3p5="Step3p5Config", qianfan_ocr="QianfanOCRConfig", qwen3_asr="Qwen3ASRConfig", - qwen3_next="Qwen3NextConfig", - qwen3_5="Qwen3_5Config", - qwen3_5_text="Qwen3_5TextConfig", - qwen3_5_moe="Qwen3_5MoeConfig", - qwen3_5_moe_text="Qwen3_5MoeTextConfig", - lfm2_moe="Lfm2MoeConfig", **{"unlimited-ocr": "UnlimitedOCRConfig"}, **{"deepseek_v41": "DeepseekV41Config"}, inkling_mm_model="InklingMMConfig", @@ -162,14 +148,6 @@ def __getitem__(self, key): # `gemma4_text` ships them in the checkpoint itself. _PATCH_HF_NESTED_ROPE_VALIDATION: set[str] = {"laguna", "gemma4_text"} -# Model types whose checkpoints declare `layer_types` entries that upstream -# transformers has not added to `ALLOWED_LAYER_TYPES` yet, so its strict config -# validation rejects them (e.g. GLM-5.2 `glm_moe_dsa` use -# `deepseek_sparse_attention`). Extend the allowed set for these model types. -_PATCH_HF_ALLOWED_LAYER_TYPES: dict[str, tuple[str, ...]] = { - "glm_moe_dsa": ("deepseek_sparse_attention",), -} - _CONFIG_ATTRS_MAPPING: dict[str, str] = { "llm_config": "text_config", } @@ -314,24 +292,6 @@ def patched_validate_rope(self, *args, **kwargs): ) -def _patch_hf_transformers_allowed_layer_types( - extra_layer_types: tuple[str, ...], -) -> None: - """Extend transformers' ``ALLOWED_LAYER_TYPES`` so its strict config - validation accepts layer types (e.g. ``deepseek_sparse_attention``) that a - checkpoint declares but upstream transformers has not registered yet. - """ - import transformers.configuration_utils as hf_configuration_utils - - missing = tuple( - layer_type - for layer_type in extra_layer_types - if layer_type not in hf_configuration_utils.ALLOWED_LAYER_TYPES - ) - if missing: - hf_configuration_utils.ALLOWED_LAYER_TYPES += missing - - class HFConfigParser(ConfigParserBase): def parse( self, @@ -379,9 +339,6 @@ def parse( if model_type in _PATCH_HF_NESTED_ROPE_VALIDATION: _patch_hf_transformers_nested_rope_validation() - if extra_layer_types := _PATCH_HF_ALLOWED_LAYER_TYPES.get(model_type): - _patch_hf_transformers_allowed_layer_types(extra_layer_types) - if model_type == "vlm": # HyperCLOVAX remote code registers this alias in a bare # `try/except` that fails silently since transformers 5.18 diff --git a/vllm/transformers_utils/configs/AXK1.py b/vllm/transformers_utils/configs/AXK1.py deleted file mode 100644 index 50bc891573c8..000000000000 --- a/vllm/transformers_utils/configs/AXK1.py +++ /dev/null @@ -1,216 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Any - -from transformers import PreTrainedConfig - - -class AXK1Config(PreTrainedConfig): - r"""This is the configuration class to store the configuration of a [`AXK1Model`]. - It is used to instantiate an A.X model according to the specified arguments, - defining the model architecture. Instantiating a configuration with the defaults - will yield a similar configuration to that of the A.X K1. - Configuration objects inherit from [`PreTrainedConfig`] and can be used to control - the model outputs. Read the documentation from [`PreTrainedConfig`] for more - information. - - Args: - vocab_size (`int`, *optional*, defaults to 163840): - Vocabulary size of the A.X K1 model. Defines the number of different - tokens that can be represented by the `inputs_ids` passed when calling - [`AXK1Model`] - hidden_size (`int`, *optional*, defaults to 7168): - Dimension of the hidden representations. - intermediate_size (`int`, *optional*, defaults to 18432): - Dimension of the MLP representations. - moe_intermediate_size (`int`, *optional*, defaults to 2048): - Dimension of the MoE representations. - num_hidden_layers (`int`, *optional*, defaults to 61): - Number of hidden layers in the Transformer decoder. - num_nextn_predict_layers (`int`, *optional*, defaults to 1): - Number of nextn predict layers in the AXK1 Model. - num_attention_heads (`int`, *optional*, defaults to 64): - Number of attention heads for each attention layer in the Transformer - decoder. - n_shared_experts (`int`, *optional*, defaults to 1): - Number of shared experts, None means dense model. - n_routed_experts (`int`, *optional*, defaults to 192): - Number of routed experts, None means dense model. - routed_scaling_factor (`float`, *optional*, defaults to 2.5): - Scaling factor or routed experts. - topk_method (`str`, *optional*, defaults to `noaux_tc`): - Topk method used in routed gate. - n_group (`int`, *optional*, defaults to 8): - Number of groups for routed experts. - topk_group (`int`, *optional*, defaults to 4): - Number of selected groups for each token(for each token, ensuring the - selected experts is only within `topk_group` groups). - num_experts_per_tok (`int`, *optional*, defaults to 8): - Number of selected experts, None means dense model. - moe_layer_freq (`int`, *optional*, defaults to 1): - The frequency of the MoE layer: one expert layer for every - `moe_layer_freq - 1` dense layers. - first_k_dense_replace (`int`, *optional*, defaults to 1): - Number of dense layers in shallow layers - (embed->dense->dense->...->dense->moe->moe...->lm_head). - \--k dense layers--/ - norm_topk_prob (`bool`, *optional*, defaults to True): - Whether to normalize the weights of the routed experts. - scoring_func (`str`, *optional*, defaults to 'sigmoid'): - Method of computing expert weights. - aux_loss_alpha (`float`, *optional*, defaults to 0.0001): - Auxiliary loss weight coefficient. - seq_aux = (`bool`, *optional*, defaults to True): - Whether to compute the auxiliary loss for each individual sample. - num_key_value_heads (`int`, *optional*): - This is the number of key_value heads that should be used to implement - Grouped Query Attention. If `num_key_value_heads=num_attention_heads`, - the model will use Multi Head Attention (MHA), if `num_key_value_heads=1 - the model will use Multi Query Attention (MQA) otherwise GQA is used. When - converting a multi-head checkpoint to a GQA checkpoint, each group key and - value head should be constructed by meanpooling all the original heads - within that group. For more details checkout - [this paper](https://arxiv.org/pdf/2305.13245.pdf). - If it is not specified, will default to `num_attention_heads`. - hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): - The non-linear activation function (function or string) in the decoder. - max_position_embeddings (`int`, *optional*, defaults to 131072): - The maximum sequence length that this model might ever be used with. - initializer_range (`float`, *optional*, defaults to 0.02): - The standard deviation of the truncated_normal_initializer for - initializing all weight matrices. - rms_norm_eps (`float`, *optional*, defaults to 1e-06): - The epsilon used by the rms normalization layers. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values attentions - (not used by all models). Only relevant if `config.is_decoder=True`. - pad_token_id (`int`, *optional*): - Padding token id. - bos_token_id (`int`, *optional*, defaults to 163691): - Beginning of stream token id. - eos_token_id (`int`, *optional*, defaults to 163691): - End of stream token id. - pretraining_tp (`int`, *optional*, defaults to 1): - Experimental feature. Tensor parallelism rank used during pretraining. - Please refer to - [this document](https://huggingface.co/docs/transformers/parallelism) - to understand more about it. This value is necessary to ensure exact - reproducibility of the pretraining results. Please refer to - [this issue](https://github.com/pytorch/pytorch/issues/76232). - tie_word_embeddings (`bool`, *optional*, defaults to `False`): - Whether to tie weight embeddings - rope_theta (`float`, *optional*, defaults to 10000.0): - The base period of the RoPE embeddings. - rope_scaling (`Dict`, *optional*): - Dictionary containing the scaling configuration for the RoPE embeddings. - Currently supports two scaling strategies: linear and dynamic. - Their scaling factor must be a float greater than 1. The expected format - is `{"type": strategy name, "factor": scaling factor}`. When using this - flag, don't update `max_position_embeddings` to the expected new maximum. - attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`): - Whether to use a bias in the query, key, value and output projection - layers during self-attention. - attention_dropout (`float`, *optional*, defaults to 0.0): - The dropout ratio for the attention probabilities. - - """ - - model_type = "axk1" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size: int = 163840, - hidden_size: int = 7168, - intermediate_size: int = 18432, - moe_intermediate_size: int = 2048, - num_hidden_layers: int = 61, - num_nextn_predict_layers: int | None = 1, - num_attention_heads: int = 64, - num_key_value_heads: int = 64, - n_shared_experts: int | None = 1, - n_routed_experts: int | None = 192, - ep_size: int | None = 8, ## Ignored - Expert parallel size - routed_scaling_factor: float | None = 2.5, - kv_lora_rank: int | None = 512, - q_lora_rank: int | None = 1536, - qk_rope_head_dim: int | None = 64, - v_head_dim: int | None = 128, - qk_nope_head_dim: int | None = 128, - topk_method: str | None = "noaux_tc", - n_group: int | None = 8, - topk_group: int | None = 4, - num_experts_per_tok: int | None = 8, - moe_layer_freq: int | None = 1, - first_k_dense_replace: int = 1, - norm_topk_prob: bool = True, - scoring_func: str | None = "sigmoid", - aux_loss_alpha: float | None = 0.0001, - seq_aux: float | None = True, - hidden_act: str | None = "silu", - max_position_embeddings: int | None = 131072, - initializer_range: float | None = 0.02, - rms_norm_eps: float = 1e-6, - use_cache: bool | None = True, - pad_token_id: int | None = None, - bos_token_id: int | None = 163691, - eos_token_id: int | None = 163691, - pretraining_tp: int | None = 1, - tie_word_embeddings: bool | None = False, - rope_theta: float | None = 10000.0, - rope_scaling: dict[str, Any] | None = None, - rope_parameters: dict[str, Any] | None = None, - attention_bias: bool | None = False, - attention_dropout: float | None = 0.0, - **kwargs, - ): - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.moe_intermediate_size = moe_intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_nextn_predict_layers = num_nextn_predict_layers - self.num_attention_heads = num_attention_heads - self.n_shared_experts = n_shared_experts - self.n_routed_experts = n_routed_experts - self.ep_size = ep_size - self.routed_scaling_factor = routed_scaling_factor - self.kv_lora_rank = kv_lora_rank - self.q_lora_rank = q_lora_rank - self.qk_rope_head_dim = qk_rope_head_dim - self.v_head_dim = v_head_dim - self.qk_nope_head_dim = qk_nope_head_dim - self.topk_method = topk_method - self.n_group = n_group - self.topk_group = topk_group - self.num_experts_per_tok = num_experts_per_tok - self.moe_layer_freq = moe_layer_freq - self.first_k_dense_replace = first_k_dense_replace - self.norm_topk_prob = norm_topk_prob - self.scoring_func = scoring_func - self.aux_loss_alpha = aux_loss_alpha - self.seq_aux = seq_aux - # for backward compatibility - if num_key_value_heads is None: - num_key_value_heads = num_attention_heads - - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.pretraining_tp = pretraining_tp - self.use_cache = use_cache - self.rope_theta = rope_theta - self.rope_scaling = rope_scaling - self.rope_parameters = rope_parameters - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - - super().__init__( - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 5bd47fb9482c..336b96c1a316 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -14,8 +14,6 @@ import importlib _CLASS_TO_MODULE: dict[str, str] = { - "AfmoeConfig": "vllm.transformers_utils.configs.afmoe", - "AXK1Config": "vllm.transformers_utils.configs.AXK1", "BagelConfig": "vllm.transformers_utils.configs.bagel", "BailingMoeV3TextConfig": "vllm.transformers_utils.configs.bailing_moe_v3_vl", "BailingMoeV3VisionConfig": "vllm.transformers_utils.configs.bailing_moe_v3_vl", @@ -26,13 +24,10 @@ "ColQwen3Config": "vllm.transformers_utils.configs.colqwen3", "OpsColQwen3Config": "vllm.transformers_utils.configs.colqwen3", "Qwen3VLNemotronEmbedConfig": "vllm.transformers_utils.configs.colqwen3", - "Cosmos3Config": "vllm.transformers_utils.configs.cosmos3", "Cosmos3EdgeConfig": "vllm.transformers_utils.configs.cosmos3_edge", "Cosmos3EdgeProjectorConfig": "vllm.transformers_utils.configs.cosmos3_edge", "Cosmos3EdgeTextConfig": "vllm.transformers_utils.configs.cosmos3_edge", "Cosmos3EdgeVisionConfig": "vllm.transformers_utils.configs.cosmos3_edge", - "DiffusionGemmaConfig": "vllm.transformers_utils.configs.diffusion_gemma", - "DiffusionGemmaTextConfig": "vllm.transformers_utils.configs.diffusion_gemma", "DeepseekVLV2Config": "vllm.transformers_utils.configs.deepseek_vl2", "DeepseekV4Config": "vllm.transformers_utils.configs.deepseek_v4", "DeepseekV41Config": "vllm.transformers_utils.configs.deepseek_v41", @@ -42,16 +37,12 @@ "EAGLEConfig": "vllm.transformers_utils.configs.eagle", "FunAudioChatConfig": "vllm.transformers_utils.configs.funaudiochat", "FunAudioChatAudioEncoderConfig": "vllm.transformers_utils.configs.funaudiochat", - "Granite4VisionConfig": "vllm.transformers_utils.configs.granite4_vision", - "HYV3Config": "vllm.transformers_utils.configs.hy_v3", "HYV4Config": "vllm.transformers_utils.configs.hy_v4", - "HyperCLOVAXConfig": "vllm.transformers_utils.configs.hyperclovax", "IsaacConfig": "vllm.transformers_utils.configs.isaac", # RWConfig is for the original tiiuae/falcon-40b(-instruct) and # tiiuae/falcon-7b(-instruct) models. Newer Falcon models will use the # `FalconConfig` class from the official HuggingFace transformers library. "RWConfig": "vllm.transformers_utils.configs.falcon", - "Lfm2MoeConfig": "vllm.transformers_utils.configs.lfm2_moe", "MedusaConfig": "vllm.transformers_utils.configs.medusa", "MiDashengLMConfig": "vllm.transformers_utils.configs.midashenglm", "MiniMaxM3Config": "vllm.transformers_utils.configs.minimax_m3", @@ -67,16 +58,12 @@ "MoonViTConfig": "vllm.transformers_utils.configs.moonvit", "KimiLinearConfig": "vllm.transformers_utils.configs.kimi_linear", "KimiVLConfig": "vllm.transformers_utils.configs.kimi_vl", - "KimiK25Config": "vllm.transformers_utils.configs.kimi_k25", "MuseGlimmerConfig": "vllm.transformers_utils.configs.muse_glimmer", "MuseGlimmerTextConfig": "vllm.transformers_utils.configs.muse_glimmer", "MuseGlimmerVisionConfig": "vllm.transformers_utils.configs.muse_glimmer", "MuseGlimmerAssistantConfig": "vllm.transformers_utils.configs.muse_glimmer", "KimiK3Config": "vllm.transformers_utils.configs.kimi_k3", "KimiK3VisionConfig": "vllm.transformers_utils.configs.kimi_k3", - "NemotronConfig": "vllm.transformers_utils.configs.nemotron", - "NemotronHConfig": "vllm.transformers_utils.configs.nemotron_h", - "OlmoHybridConfig": "vllm.transformers_utils.configs.olmo_hybrid", "OpenVLAConfig": "vllm.transformers_utils.configs.openvla", "OvisConfig": "vllm.transformers_utils.configs.ovis", "PixelShuffleSiglip2VisionConfig": "vllm.transformers_utils.configs.isaac", @@ -91,22 +78,17 @@ "QianfanOCRConfig": "vllm.transformers_utils.configs.qianfan_ocr", "QianfanOCRVisionConfig": "vllm.transformers_utils.configs.qianfan_ocr", "Qwen3ASRConfig": "vllm.transformers_utils.configs.qwen3_asr", - "Qwen3NextConfig": "vllm.transformers_utils.configs.qwen3_next", - "Qwen3_5Config": "vllm.transformers_utils.configs.qwen3_5", - "Qwen3_5TextConfig": "vllm.transformers_utils.configs.qwen3_5", - "Qwen3_5MoeConfig": "vllm.transformers_utils.configs.qwen3_5_moe", - "Qwen3_5MoeTextConfig": "vllm.transformers_utils.configs.qwen3_5_moe", "InklingModelConfig": "vllm.models.inkling.configs", "InklingAudioConfig": "vllm.models.inkling.configs", "InklingVisionConfig": "vllm.models.inkling.configs", "InklingMMConfig": "vllm.models.inkling.configs", - # Special case: DeepseekV3Config is from HuggingFace Transformers + # Upstream Transformers classes registered in _CONFIG_REGISTRY "DeepseekV3Config": "transformers", + "HyperCLOVAXConfig": "transformers", + "Kimi_K25Config": "transformers", } __all__ = [ - "AfmoeConfig", - "AXK1Config", "BagelConfig", "BailingMoeV3TextConfig", "BailingMoeV3VisionConfig", @@ -117,13 +99,10 @@ "ColQwen3Config", "OpsColQwen3Config", "Qwen3VLNemotronEmbedConfig", - "Cosmos3Config", "Cosmos3EdgeConfig", "Cosmos3EdgeProjectorConfig", "Cosmos3EdgeTextConfig", "Cosmos3EdgeVisionConfig", - "DiffusionGemmaConfig", - "DiffusionGemmaTextConfig", "DeepseekVLV2Config", "DeepseekV3Config", "DeepseekV4Config", @@ -133,13 +112,10 @@ "EAGLEConfig", "FunAudioChatConfig", "FunAudioChatAudioEncoderConfig", - "Granite4VisionConfig", - "HYV3Config", "HYV4Config", "HyperCLOVAXConfig", "IsaacConfig", "RWConfig", - "Lfm2MoeConfig", "MedusaConfig", "MiDashengLMConfig", "MiniMaxM3Config", @@ -153,16 +129,13 @@ "MoonViTConfig", "KimiLinearConfig", "KimiVLConfig", - "KimiK25Config", + "Kimi_K25Config", "MuseGlimmerConfig", "MuseGlimmerTextConfig", "MuseGlimmerVisionConfig", "MuseGlimmerAssistantConfig", "KimiK3Config", "KimiK3VisionConfig", - "NemotronConfig", - "NemotronHConfig", - "OlmoHybridConfig", "OpenVLAConfig", "OvisConfig", "PixelShuffleSiglip2VisionConfig", @@ -177,11 +150,6 @@ "QianfanOCRConfig", "QianfanOCRVisionConfig", "Qwen3ASRConfig", - "Qwen3NextConfig", - "Qwen3_5Config", - "Qwen3_5TextConfig", - "Qwen3_5MoeConfig", - "Qwen3_5MoeTextConfig", "InklingModelConfig", "InklingAudioConfig", "InklingVisionConfig", diff --git a/vllm/transformers_utils/configs/afmoe.py b/vllm/transformers_utils/configs/afmoe.py deleted file mode 100644 index b3d10dc5f274..000000000000 --- a/vllm/transformers_utils/configs/afmoe.py +++ /dev/null @@ -1,87 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from transformers.configuration_utils import PreTrainedConfig - - -class AfmoeConfig(PreTrainedConfig): - model_type = "afmoe" - - def __init__( - self, - vocab_size: int = 200_192, - hidden_size: int = 2048, - intermediate_size: int = 6144, - moe_intermediate_size: int = 1408, - num_hidden_layers: int = 32, - num_dense_layers: int = 1, - num_attention_heads: int = 16, - num_key_value_heads: int | None = None, - head_dim: int = 128, - hidden_act: str = "silu", - max_position_embeddings: int = 131072, - initializer_range: float = 0.02, - rms_norm_eps: float = 1e-5, - use_cache: bool = True, - tie_word_embeddings: bool = False, - rope_parameters: dict | None = None, - rope_scaling: dict | None = None, - num_experts: int = 64, - num_experts_per_tok: int = 6, - num_shared_experts: int = 2, - num_expert_groups: int = 1, - num_limited_groups: int = 1, - score_func: str = "sigmoid", - route_norm: bool = True, - route_scale: float = 1.0, - global_attn_every_n_layers: int = 4, - sliding_window: int = 2048, - layer_types: list[str] | None = None, - attention_dropout: float = 0.0, - mup_enabled: bool = False, - n_group: int = 1, - topk_group: int = 1, - **kwargs, - ): - self.vocab_size = vocab_size - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_dense_layers = num_dense_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads or num_attention_heads - self.head_dim = head_dim - self.hidden_act = hidden_act - self.max_position_embeddings = max_position_embeddings - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.use_cache = use_cache - rope_theta = kwargs.pop("rope_theta", 10000.0) - if rope_parameters is None: - rope_parameters = {"rope_type": "default", "rope_theta": rope_theta} - self.rope_parameters = rope_parameters - self.rope_scaling = rope_scaling - - self.moe_intermediate_size = moe_intermediate_size - self.num_experts = num_experts - self.num_experts_per_tok = num_experts_per_tok - self.num_shared_experts = num_shared_experts - self.num_expert_groups = num_expert_groups - self.num_limited_groups = num_limited_groups - self.score_func = score_func - self.route_norm = route_norm - self.route_scale = route_scale - - self.global_attn_every_n_layers = global_attn_every_n_layers - self.sliding_window = sliding_window - self.layer_types = layer_types - self.attention_dropout = attention_dropout - - self.mup_enabled = mup_enabled - self.n_group = n_group - self.topk_group = topk_group - - super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) - - -__all__ = ["AfmoeConfig"] diff --git a/vllm/transformers_utils/configs/cosmos3.py b/vllm/transformers_utils/configs/cosmos3.py deleted file mode 100644 index 965801888a4b..000000000000 --- a/vllm/transformers_utils/configs/cosmos3.py +++ /dev/null @@ -1,7 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig - - -class Cosmos3Config(Qwen3VLConfig): - model_type = "cosmos3_omni" diff --git a/vllm/transformers_utils/configs/cosmos3_edge.py b/vllm/transformers_utils/configs/cosmos3_edge.py index 447b64cc99ed..ddf17933fd50 100644 --- a/vllm/transformers_utils/configs/cosmos3_edge.py +++ b/vllm/transformers_utils/configs/cosmos3_edge.py @@ -1,11 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from transformers import NemotronHConfig from transformers.configuration_utils import PreTrainedConfig from transformers.models.siglip2.configuration_siglip2 import Siglip2VisionConfig -from .nemotron_h import NemotronHConfig - def _normalize_mrope_parameters(config: NemotronHConfig) -> None: """Normalize the checkpoint's mRoPE configuration for vLLM.""" @@ -31,11 +30,14 @@ def __init__( num_hidden_layers: int = 28, hidden_act: str = "relu2", rms_norm_eps: float = 1e-5, + layers_block_type: list[str] | None = None, **kwargs, ) -> None: + # Each decoder block is an attention layer followed by an MLP layer + if layers_block_type is None: + layers_block_type = ["full_attention", "mlp"] * num_hidden_layers super().__init__( - num_hidden_layers=2 * num_hidden_layers, - hybrid_override_pattern="*-" * num_hidden_layers, + layers_block_type=layers_block_type, mlp_hidden_act=hidden_act, layer_norm_epsilon=rms_norm_eps, **kwargs, @@ -46,7 +48,6 @@ def to_dict(self) -> dict: config_dict["num_hidden_layers"] = self.num_hidden_layers // 2 config_dict["hidden_act"] = config_dict.pop("mlp_hidden_act") config_dict["rms_norm_eps"] = config_dict.pop("layer_norm_epsilon") - config_dict.pop("hybrid_override_pattern", None) return config_dict diff --git a/vllm/transformers_utils/configs/deepseek_vl2.py b/vllm/transformers_utils/configs/deepseek_vl2.py index 3543db833a4d..a819f12f3e4e 100644 --- a/vllm/transformers_utils/configs/deepseek_vl2.py +++ b/vllm/transformers_utils/configs/deepseek_vl2.py @@ -90,11 +90,6 @@ def __init__( @strict class DeepseekVLV2TextConfig(DeepseekV2Config): - # DeepSeek-VL2 checkpoints (e.g. deepseek-vl2-small) omit ``vocab_size`` and - # rely on the original DeepSeek-VL2 default of 102400; the generic - # Transformers ``DeepseekV2Config`` incorrectly defaults it to 32000. - # Upstream fix: https://github.com/huggingface/transformers/pull/48159 - vocab_size: int = 102400 kv_lora_rank: int | None = DeepseekV2Config.kv_lora_rank diff --git a/vllm/transformers_utils/configs/diffusion_gemma.py b/vllm/transformers_utils/configs/diffusion_gemma.py deleted file mode 100644 index 884d176cd0b4..000000000000 --- a/vllm/transformers_utils/configs/diffusion_gemma.py +++ /dev/null @@ -1,44 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Any - -from transformers import PreTrainedConfig -from transformers.models.gemma4.configuration_gemma4 import Gemma4VisionConfig - - -def _init_text_config(self: PreTrainedConfig, **kwargs: Any) -> None: - PreTrainedConfig.__init__(self, **kwargs) - # DiffusionGemma always uses MoE and K=V sharing for full_attention - # layers. The HF reference removed these config fields entirely. - if getattr(self, "num_experts", None): - self.enable_moe_block = True - self.attention_k_eq_v = True - - -class DiffusionGemmaTextConfig(PreTrainedConfig): - model_type = "diffusion_gemma_text" - - def __init__(self, **kwargs: Any): - _init_text_config(self, **kwargs) - - -class DiffusionGemmaConfig(PreTrainedConfig): - model_type = "diffusion_gemma" - - def __init__( - self, - text_config: dict[str, Any] | None = None, - canvas_length: int = 256, - self_conditioning_size: int | None = None, - **kwargs: Any, - ): - self.text_config = DiffusionGemmaTextConfig(**(text_config or {})) - self.canvas_length = canvas_length - self.self_conditioning_size = self_conditioning_size - vision_config = kwargs.pop("vision_config", None) - if isinstance(vision_config, dict): - self.vision_config = Gemma4VisionConfig(**vision_config) - else: - self.vision_config = vision_config - self.audio_config = None - PreTrainedConfig.__init__(self, **kwargs) diff --git a/vllm/transformers_utils/configs/gemma4.py b/vllm/transformers_utils/configs/gemma4.py deleted file mode 100644 index 69cbc06d0482..000000000000 --- a/vllm/transformers_utils/configs/gemma4.py +++ /dev/null @@ -1,34 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Per-layer config resolution for Gemma4 and its variants.""" - -from copy import copy - -from transformers import PreTrainedConfig - - -def gemma4_layer_config( - text_config: PreTrainedConfig, layer_idx: int -) -> PreTrainedConfig: - """The Gemma4 text config as it applies to one layer. - - Gemma4 uses a larger head dimension on its full attention layers than on its - sliding ones, and with `attention_k_eq_v` it uses more KV heads there too. - Transformers >= 5.15.0 says so in the config itself; before that the values - are flat attributes that `layer_types` picks between, so resolve them here. - The result is homogeneous either way, so callers read `head_dim` and - `num_key_value_heads` off it without caring which version they are on. - """ - if getattr(text_config, "is_heterogeneous", False): - return text_config.per_layer_config[layer_idx] - - layer = copy(text_config) - if text_config.layer_types[layer_idx] == "full_attention": - global_head_dim = getattr(text_config, "global_head_dim", None) - layer.head_dim = global_head_dim or text_config.head_dim - if getattr(text_config, "attention_k_eq_v", False): - global_kv_heads = getattr(text_config, "num_global_key_value_heads", None) - layer.num_key_value_heads = ( - global_kv_heads or text_config.num_key_value_heads - ) - return layer diff --git a/vllm/transformers_utils/configs/granite4_vision.py b/vllm/transformers_utils/configs/granite4_vision.py deleted file mode 100644 index 8522bdd47a3e..000000000000 --- a/vllm/transformers_utils/configs/granite4_vision.py +++ /dev/null @@ -1,94 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from typing import Any - -import transformers - - -class Granite4VisionConfig(transformers.PreTrainedConfig): - """Configuration for Granite 4 Vision model. - - This config is needed because the granite4_vision model type is not yet - in the transformers version pinned by vLLM. Once transformers adds native - support, this file can be removed and the _CONFIG_REGISTRY entry dropped. - """ - - model_type = "granite4_vision" - is_composition = False - - def __init__( - self, - vision_config: dict[str, Any] | None = None, - text_config: dict[str, Any] | None = None, - image_token_index: int = 100352, - image_seq_length: int = 576, - image_grid_pinpoints: list[list[int]] | None = None, - vision_feature_select_strategy: str = "full", - vision_feature_layer: int | list[int] = -2, - projector_hidden_act: str = "gelu", - projector_dropout: float = 0.1, - downsample_rate: str | None = None, - use_image_newline_parameter: bool = True, - deepstack_layer_map: list[list[int]] | None = None, - use_spatial_sampling: bool = False, - spatial_stride: int = 2, - spatial_vision_layer: int = -1, - spatial_target_layers: list[int] | None = None, - # Hub aliases — base model config uses different field names - vision_layer_to_llm_layer: list[list[int]] | None = None, - use_checkerboard_sampling: bool | None = None, - checkerboard_stride: int | None = None, - checkerboard_vision_layer: int | None = None, - checkerboard_llm_layers: list[int] | None = None, - **kwargs: Any, - ): - self.image_token_index = image_token_index - self.image_seq_length = image_seq_length - self.image_grid_pinpoints = image_grid_pinpoints or [] - self.vision_feature_select_strategy = vision_feature_select_strategy - self.vision_feature_layer = vision_feature_layer - self.projector_hidden_act = projector_hidden_act - self.projector_dropout = projector_dropout - self.downsample_rate = downsample_rate - self.use_image_newline_parameter = use_image_newline_parameter - self.deepstack_layer_map = deepstack_layer_map or vision_layer_to_llm_layer - self.use_spatial_sampling = ( - use_spatial_sampling - if use_checkerboard_sampling is None - else use_checkerboard_sampling - ) - self.spatial_stride = ( - spatial_stride if checkerboard_stride is None else checkerboard_stride - ) - self.spatial_vision_layer = ( - spatial_vision_layer - if checkerboard_vision_layer is None - else checkerboard_vision_layer - ) - self.spatial_target_layers = ( - spatial_target_layers or checkerboard_llm_layers or [0, 10, 20, 30] - ) - - if vision_config is None: - vision_config = {} - if text_config is None: - text_config = {} - - vision_model_type = vision_config.get("model_type", "siglip_vision_model") - if vision_model_type in transformers.CONFIG_MAPPING: - self.vision_config = transformers.CONFIG_MAPPING[vision_model_type]( - **vision_config - ) - else: - self.vision_config = transformers.PreTrainedConfig(**vision_config) - - text_model_type = text_config.get("model_type", "granite") - if text_model_type in transformers.CONFIG_MAPPING: - self.text_config = transformers.CONFIG_MAPPING[text_model_type]( - **text_config - ) - else: - self.text_config = transformers.PreTrainedConfig(**text_config) - - super().__init__(**kwargs) diff --git a/vllm/transformers_utils/configs/hy_v3.py b/vllm/transformers_utils/configs/hy_v3.py deleted file mode 100644 index 939dd5d78ceb..000000000000 --- a/vllm/transformers_utils/configs/hy_v3.py +++ /dev/null @@ -1,185 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Any - -from transformers.configuration_utils import PreTrainedConfig - - -class HYV3Config(PreTrainedConfig): - r"""This is the configuration class to store the configuration of a [`HYV3Model`]. - It is used to instantiate a HYV3 model (HY V3 MoE language model) according to - the specified arguments. - - Configuration objects inherit from [`PreTrainedConfig`] and can be used to - control the model outputs. Read the documentation from [`PreTrainedConfig`] - for more information. - - Args: - vocab_size (`int`, *optional*, defaults to 120832): - Vocabulary size of the model. - hidden_size (`int`, *optional*, defaults to 4096): - Dimension of the hidden representations. - intermediate_size (`int`, *optional*, defaults to 13312): - Dimension of the dense FFN intermediate representations. - num_hidden_layers (`int`, *optional*, defaults to 80): - Number of hidden layers in the Transformer decoder. - num_attention_heads (`int`, *optional*, defaults to 64): - Number of attention heads for each attention layer. - num_key_value_heads (`int`, *optional*, defaults to 8): - Number of key-value heads for grouped-query attention. - head_dim (`int`, *optional*, defaults to 128): - Dimension per attention head. - hidden_act (`str`, *optional*, defaults to `"silu"`): - Activation function used in FFN layers. - max_position_embeddings (`int`, *optional*, defaults to 131072): - Maximum sequence length supported by the model. - initializer_range (`float`, *optional*, defaults to 0.006): - Standard deviation of the truncated normal initializer for weight - initialization. - rms_norm_eps (`float`, *optional*, defaults to 1e-5): - Epsilon for RMS normalization layers. - use_cache (`bool`, *optional*, defaults to `True`): - Whether to use KV cache for decoding. - pad_token_id (`int`, *optional*): - Padding token id. - bos_token_id (`int`, *optional*): - Beginning-of-sequence token id. - eos_token_id (`int` or `List[int]`, *optional*): - End-of-sequence token id(s). - rope_parameters (`dict`, *optional*): - The parameters of the RoPE embeddings. - qk_norm (`bool`, *optional*, defaults to `True`): - Whether to apply RMSNorm to query and key states before attention. - tie_word_embeddings (`bool`, *optional*, defaults to `False`): - Whether to tie input and output embedding weights. - enable_attention_fp32_softmax (`bool`, *optional*, defaults to `False`): - Whether to upcast attention softmax to float32. Note: the eager attention - path always computes softmax in float32 regardless of this setting; this - flag is reserved for future use with custom attention backends. - enable_lm_head_fp32 (`bool`, *optional*, defaults to `True`): - Whether to upcast the LM head computation to float32. - num_experts (`int`, *optional*, defaults to 192): - Total number of MoE experts. - num_experts_per_tok (`int`, *optional*, defaults to 8): - Number of experts selected per token (top-k routing). - num_shared_experts (`int`, *optional*, defaults to 1): - Number of always-active shared experts combined into a single MLP. - expert_hidden_dim (`int`, *optional*, defaults to 1536): - Intermediate dimension of each individual MoE expert. - moe_router_enable_expert_bias (`bool`, *optional*, defaults to `True`): - Whether to use per-expert load-balancing bias in the router. - moe_router_use_sigmoid (`bool`, *optional*, defaults to `True`): - Whether to use sigmoid (instead of softmax) for router scoring. - route_norm (`bool`, *optional*, defaults to `True`): - Whether to normalize routing scores when using sigmoid routing. - router_scaling_factor (`float`, *optional*): - Optional multiplicative scaling factor applied to routing scores. - use_grouped_mm (`bool`, *optional*, defaults to `False`): - Whether to use grouped GEMM for expert computation (not yet implemented). - enable_moe_fp32_combine (`bool`, *optional*, defaults to `False`): - Whether to accumulate expert outputs in float32. - first_k_dense_replace (`int`, *optional*, defaults to 1): - Number of initial decoder layers that use a dense FFN instead of MoE. - output_router_logits (`bool`, *optional*, defaults to `False`): - Whether to output router logits from each MoE layer. Useful for computing - auxiliary load-balancing loss during training. Disabled by default to avoid - the memory overhead of storing per-layer router tensors during inference. - - Example: - ```python - >>> from transformers import HYV3Config, HYV3Model - - >>> config = HYV3Config() - >>> model = HYV3Model(config) - ``` - - """ - - model_type = "hy_v3" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size=120832, - hidden_size=4096, - intermediate_size=13312, - num_hidden_layers=80, - num_attention_heads=64, - num_key_value_heads=8, - head_dim=128, - hidden_act="silu", - max_position_embeddings=131072, - initializer_range=0.006, - rms_norm_eps=1e-5, - use_cache=True, - pad_token_id=None, - bos_token_id=None, - eos_token_id=None, - rope_parameters: dict[str, Any] | None = None, - qk_norm=True, - tie_word_embeddings=False, - enable_attention_fp32_softmax=False, - enable_lm_head_fp32=True, - # MoE specific - num_experts=192, - num_experts_per_tok=8, - num_shared_experts=1, - expert_hidden_dim=1536, - moe_router_enable_expert_bias=True, - moe_router_use_sigmoid=True, - route_norm=True, - router_scaling_factor=None, - use_grouped_mm=False, - enable_moe_fp32_combine=False, - # Dense/MoE layer control - first_k_dense_replace=1, - output_router_logits=False, - **kwargs, - ): - self.vocab_size = vocab_size - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.head_dim = head_dim - self.hidden_act = hidden_act - self.max_position_embeddings = max_position_embeddings - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.use_cache = use_cache - rope_theta = kwargs.pop("rope_theta", 11158840.0) - if rope_parameters is None: - rope_parameters = {"rope_type": "default", "rope_theta": rope_theta} - self.rope_parameters = rope_parameters - self.qk_norm = qk_norm - self.tie_word_embeddings = tie_word_embeddings - self.enable_lm_head_fp32 = enable_lm_head_fp32 - self.enable_attention_fp32_softmax = enable_attention_fp32_softmax - - # MoE specific - self.num_experts = num_experts - self.num_experts_per_tok = num_experts_per_tok - self.num_shared_experts = num_shared_experts - self.expert_hidden_dim = expert_hidden_dim - self.moe_router_enable_expert_bias = moe_router_enable_expert_bias - self.moe_router_use_sigmoid = moe_router_use_sigmoid - self.route_norm = route_norm - self.use_grouped_mm = use_grouped_mm - self.router_scaling_factor = router_scaling_factor - self.enable_moe_fp32_combine = enable_moe_fp32_combine - - # Dense/MoE layer control - self.first_k_dense_replace = first_k_dense_replace - self.output_router_logits = output_router_logits - - if eos_token_id is not None and isinstance(eos_token_id, int): - eos_token_id = [eos_token_id] - - super().__init__( - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) diff --git a/vllm/transformers_utils/configs/hyperclovax.py b/vllm/transformers_utils/configs/hyperclovax.py deleted file mode 100644 index 5470408b5803..000000000000 --- a/vllm/transformers_utils/configs/hyperclovax.py +++ /dev/null @@ -1,277 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# SPDX-FileCopyrightText: Copyright 2025 NAVER Cloud HyperCLOVA team -# -# Copyright 2025 NAVER Cloud HyperCLOVA team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""HyperCLOVA X model configuration.""" - -from transformers.configuration_utils import PreTrainedConfig - - -class HyperCLOVAXConfig(PreTrainedConfig): - r"""This is the configuration class to store the configuration of a - [`HyperCLOVAXModel`]. It is used to instantiate a HyperCLOVAX model - according to the specified arguments, defining the model architecture. - Configuration objects inherit from [`PreTrainedConfig`] and can be used - to control the model outputs. Read the documentation from - [`PreTrainedConfig`] for more information. - - Args: - vocab_size (`int`, *optional*, defaults to 32000): - Vocabulary size of the HyperCLOVAX model. Defines the number of - different tokens that can be represented by the `input_ids` - passed when calling [`HyperCLOVAXModel`] - hidden_size (`int`, *optional*, defaults to 4096): - Dimension of the hidden representations. - intermediate_size (`int`, *optional*, defaults to 11008): - Dimension of the MLP representations. - num_hidden_layers (`int`, *optional*, defaults to 32): - Number of hidden layers in the Transformer decoder. - num_attention_heads (`int`, *optional*, defaults to 32): - Number of attention heads for each attention layer in the - Transformer decoder. - num_key_value_heads (`int`, *optional*): - This is the number of key_value heads that should be used to - implement Grouped Query Attention. If - `num_key_value_heads=num_attention_heads`, the model will use - Multi Head Attention (MHA), if `num_key_value_heads=1` the model - will use Multi Query Attention (MQA) otherwise GQA is used. When - converting a multi-head checkpoint to a GQA checkpoint, each - group key and value head should be constructed by meanpooling all - the original heads within that group. For more details checkout - [this paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not - specified, will default to `num_attention_heads`. - hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): - The non-linear activation function (function or string) in the - decoder. - max_position_embeddings (`int`, *optional*, defaults to 2048): - The maximum sequence length that this model might ever be used - with. - initializer_range (`float`, *optional*, defaults to 0.02): - The standard deviation of the truncated_normal_initializer for - initializing all weight matrices. - rms_norm_eps (`float`, *optional*, defaults to 1e-06): - The epsilon used by the rms normalization layers. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values - attentions (not used by all models). Only relevant if - `config.is_decoder=True`. - pad_token_id (`int`, *optional*): - Padding token id. - bos_token_id (`int`, *optional*, defaults to 1): - Beginning of stream token id. - eos_token_id (`int`, *optional*, defaults to 2): - End of stream token id. - pretraining_tp (`int`, *optional*, defaults to 1): - Experimental feature. Tensor parallelism rank used during - pretraining. Please refer to [this document](https://huggingface. - co/docs/transformers/main/perf_train_gpu_many#tensor-parallelism) - to understand more about it. This value is necessary to ensure - exact reproducibility of the pretraining results. Please refer to - [this issue](https://github.com/pytorch/pytorch/issues/76232). - tie_word_embeddings (`bool`, *optional*, defaults to `False`): - Whether to tie weight embeddings - rope_theta (`float`, *optional*, defaults to 10000.0): - The base period of the RoPE embeddings. - rope_scaling (`Dict`, *optional*): - Dictionary containing the scaling configuration for the RoPE - embeddings. NOTE: if you apply new rope type and you expect the - model to work on longer `max_position_embeddings`, we recommend - you to update this value accordingly. - Expected contents: - `rope_type` (`str`): - The sub-variant of RoPE to use. Can be one of ['default', - 'linear', 'dynamic', 'yarn', 'longrope', 'llama3'], with - 'default' being the original RoPE implementation. - `factor` (`float`, *optional*): - Used with all rope types except 'default'. The scaling - factor to apply to the RoPE embeddings. In most scaling - types, a `factor` of x will enable the model to handle - sequences of length x * original maximum pre-trained - length. - `original_max_position_embeddings` (`int`, *optional*): - Used with 'dynamic', 'longrope' and 'llama3'. The - original max position embeddings used during pretraining. - `attention_factor` (`float`, *optional*): - Used with 'yarn' and 'longrope'. The scaling factor to be - applied on the attention computation. If unspecified, it - defaults to value recommended by the implementation, using - the `factor` field to infer the suggested value. - `beta_fast` (`float`, *optional*): - Only used with 'yarn'. Parameter to set the boundary for - extrapolation (only) in the linear ramp function. If - unspecified, it defaults to 32. - `beta_slow` (`float`, *optional*): - Only used with 'yarn'. Parameter to set the boundary for - interpolation (only) in the linear ramp function. If - unspecified, it defaults to 1. - `short_factor` (`List[float]`, *optional*): - Only used with 'longrope'. The scaling factor to be - applied to short contexts (< - `original_max_position_embeddings`). Must be a list of - numbers with the same length as the hidden size divided - by the number of attention heads divided by 2 - `long_factor` (`List[float]`, *optional*): - Only used with 'longrope'. The scaling factor to be - applied to long contexts (< - `original_max_position_embeddings`). Must be a list of - numbers with the same length as the hidden size divided - by the number of attention heads divided by 2 - `low_freq_factor` (`float`, *optional*): - Only used with 'llama3'. Scaling factor applied to low - frequency components of the RoPE - `high_freq_factor` (`float`, *optional*): - Only used with 'llama3'. Scaling factor applied to high - frequency components of the RoPE - attention_bias (`bool`, *optional*, defaults to `False`): - Whether to use a bias in the query, key, value and output - projection layers during self-attention. - attention_dropout (`float`, *optional*, defaults to 0.0): - The dropout ratio for the attention probabilities. - mlp_bias (`bool`, *optional*, defaults to `False`): - Whether to use a bias in up_proj, down_proj and gate_proj layers - in the MLP layers. - head_dim (`int`, *optional*): - The attention head dimension. If None, it will default to - hidden_size // num_heads - embedding_multiplier (`float`, *optional*, defaults to `None`): - Multiplier applied to the embedding weights. If `None`, it is - equivalent to `1.0`. - logits_scaling (`float`, *optional*, defaults to `None`): - Scaling factor for logits. If `None`, it is equivalent to `1.0`. - attention_multiplier (`float`, *optional*, defaults to `None`): - Multiplier applied to the attention weights. If `None`, it is - equivalent to `self.head_dim ** -0.5`. - residual_multiplier (`float`, *optional*, defaults to `None`): - Scaling factor for residual connections. If `None`, it is - equivalent to `1.0`. - use_post_norm (`bool`, *optional*, defaults to `True`): - Determines whether to apply Peri-Layer Normalization. Set to - False to disable this feature. - rope_parameters (`dict`, *optional*): - Dictionary containing the RoPE parameters used by vLLM's - `get_rope`. When provided, takes precedence over `rope_theta` - and `rope_scaling`. If `None`, it is derived from `rope_theta` - and `rope_scaling` automatically. - - """ - - model_type = "hyperclovax" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size=32000, - hidden_size=4096, - intermediate_size=11008, - num_hidden_layers=32, - num_attention_heads=32, - num_key_value_heads=None, - hidden_act="silu", - max_position_embeddings=2048, - initializer_range=0.02, - rms_norm_eps=1e-6, - use_cache=True, - pad_token_id=None, - bos_token_id=1, - eos_token_id=2, - pretraining_tp=1, - tie_word_embeddings=False, - rope_theta=10000.0, - rope_scaling=None, - attention_bias=False, - attention_dropout=0.0, - mlp_bias=False, - head_dim=None, - embedding_multiplier=None, # mup - logits_scaling=None, # mup - attention_multiplier=None, # mup - residual_multiplier=None, # mup - use_post_norm=True, # post-norm(peri-LN) - rope_parameters=None, - auto_map=None, - **kwargs, - ): - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - - # for backward compatibility - if num_key_value_heads is None: - num_key_value_heads = num_attention_heads - - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.pretraining_tp = pretraining_tp - self.use_cache = use_cache - self.rope_theta = rope_theta - self.rope_scaling = rope_scaling - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.mlp_bias = mlp_bias - self.head_dim = ( - head_dim - if head_dim is not None - else self.hidden_size // self.num_attention_heads - ) - # Derive rope_parameters for vLLM's get_rope() from rope_theta / - # rope_scaling, unless the caller already provided rope_parameters. - if rope_parameters is None: - if rope_scaling is not None: - # Shallow-copy to avoid mutating the caller's dict. - rope_parameters = dict(rope_scaling) - # BC: 'type' field -> 'rope_type', remove stale key. - if "type" in rope_parameters: - rope_parameters.setdefault("rope_type", rope_parameters.pop("type")) - else: - rope_parameters = {"rope_type": "default"} - if "rope_theta" not in rope_parameters: - rope_parameters["rope_theta"] = rope_theta - self.rope_parameters = rope_parameters - - # BC: keep self.rope_scaling consistent for HF serialization. - if self.rope_scaling is not None and "type" in self.rope_scaling: - self.rope_scaling["rope_type"] = self.rope_scaling["type"] - - # mup - self.embedding_multiplier = ( - embedding_multiplier if embedding_multiplier is not None else 1.0 - ) - self.logits_scaling = logits_scaling if logits_scaling is not None else 1.0 - self.attention_multiplier = ( - attention_multiplier - if attention_multiplier is not None - else self.head_dim**-0.5 - ) - self.residual_multiplier = ( - residual_multiplier if residual_multiplier is not None else 1.0 - ) - - # post-norm (Peri-LN) - self.use_post_norm = use_post_norm - - super().__init__( - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - auto_map=auto_map, - **kwargs, - ) diff --git a/vllm/transformers_utils/configs/kimi_k25.py b/vllm/transformers_utils/configs/kimi_k25.py deleted file mode 100644 index ed3d05d7cc21..000000000000 --- a/vllm/transformers_utils/configs/kimi_k25.py +++ /dev/null @@ -1,131 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Kimi-K2.5 Model Configuration. - -This configuration supports video-chunk as an internal modality type. -A video-chunk is the smallest independently processable unit of video. -""" - -from transformers import DeepseekV3Config -from transformers.configuration_utils import PreTrainedConfig - - -class KimiK25VisionConfig(PreTrainedConfig): - model_type = "kimi_k25_vision" - - def __init__( - self, - # Vision Tower - patch_size: int = 14, - init_pos_emb_height: int = 64, - init_pos_emb_width: int = 64, - init_pos_emb_time: int = 4, - pos_emb_type: str = "divided_fixed", - num_attention_heads: int = 16, - num_hidden_layers: int = 27, - hidden_size: int = 1152, - intermediate_size: int = 4304, - merge_kernel_size: tuple[int, int] = (2, 2), - video_attn_type: str = "spatial_temporal", - merge_type: str = "sd2_tpool", - # MM Projector - mm_projector_type: str = "patchmerger", - mm_hidden_size: int | None = None, - projector_hidden_act: str = "gelu", - projector_ln_eps: float = 1e-5, - **kwargs, - ): - super().__init__(**kwargs) - # Vision Tower - self.patch_size = patch_size - self.init_pos_emb_height = init_pos_emb_height - self.init_pos_emb_width = init_pos_emb_width - self.init_pos_emb_time = init_pos_emb_time - self.pos_emb_type = pos_emb_type - self.num_attention_heads = num_attention_heads - self.num_hidden_layers = num_hidden_layers - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.merge_kernel_size = merge_kernel_size - self.video_attn_type = video_attn_type - self.merge_type = merge_type - # MM Projector - self.mm_projector_type = mm_projector_type - if mm_hidden_size is not None: - self.mm_hidden_size = mm_hidden_size - else: - self.mm_hidden_size = hidden_size - self.projector_hidden_act = projector_hidden_act - self.projector_ln_eps = projector_ln_eps - - -class KimiK25Config(PreTrainedConfig): - """Kimi-K2.5 model configuration. - - Kimi-K2.5 extends Kimi-K2 with vision support using video-chunks. - A video-chunk consists of multiple consecutive frames - that are processed together with temporal pooling. - - Args: - vision_config: Configuration for the vision tower and projector. - text_config: Configuration for the text model (DeepseekV3). - ignore_index: The ignore index for the loss function. - media_placeholder_token_id: The token ID for media placeholders. - pad_token_id: The token ID for padding. - - """ - - model_type = "kimi_k25" - - def __init__( - self, - vision_config: dict | KimiK25VisionConfig | None = None, - text_config: dict | DeepseekV3Config | None = None, - ignore_index: int = -100, - media_placeholder_token_id: int = 163605, - pad_token_id: int = 0, - use_unified_vision_chunk: bool = False, - video_placeholder: str = "<|kimi_k25_video_placeholder|>", - **kwargs, - ): - # Vision config - if vision_config is None: - self.vision_config = KimiK25VisionConfig() - elif isinstance(vision_config, dict): - self.vision_config = KimiK25VisionConfig(**vision_config) - else: - self.vision_config = vision_config - - # Text config - if text_config is None: - self.text_config = DeepseekV3Config() - elif isinstance(text_config, dict): - self.text_config = DeepseekV3Config(**text_config) - else: - self.text_config = text_config - - # Set mm_hidden_size to text hidden size if not explicitly set - if self.vision_config.mm_hidden_size == self.vision_config.hidden_size: - self.vision_config.mm_hidden_size = self.text_config.hidden_size - - # Other config - self.ignore_index = ignore_index - self.media_placeholder_token_id = media_placeholder_token_id - self.use_unified_vision_chunk = use_unified_vision_chunk - self.video_placeholder = video_placeholder - - # Propagate quantization config from text model - if getattr(self.text_config, "quantization_config", None) is not None: - self.quantization_config = self.text_config.quantization_config - - super().__init__(pad_token_id=pad_token_id, **kwargs) - - @property - def hidden_size(self) -> int: - """Get hidden size from text config for compatibility.""" - return self.text_config.hidden_size - - @property - def vocab_size(self) -> int: - """Get vocab size from text config for compatibility.""" - return self.text_config.vocab_size diff --git a/vllm/transformers_utils/configs/lfm2_moe.py b/vllm/transformers_utils/configs/lfm2_moe.py deleted file mode 100644 index 903ae9d54700..000000000000 --- a/vllm/transformers_utils/configs/lfm2_moe.py +++ /dev/null @@ -1,164 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Any - -from transformers.configuration_utils import PreTrainedConfig - - -class Lfm2MoeConfig(PreTrainedConfig): - r"""This is the configuration class to store the configuration of a [`Lfm2MoeModel`]. It is used to instantiate a LFM2 Moe - model according to the specified arguments, defining the model architecture. Instantiating a configuration with the - defaults will yield a similar configuration to that of the LFM2-8B-A1B model. - e.g. [LiquidAI/LFM2-8B-A1B](https://huggingface.co/LiquidAI/LFM2-8B-A1B) - - Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the - documentation from [`PreTrainedConfig`] for more information. - - - Args: - vocab_size (`int`, *optional*, defaults to 65536): - Vocabulary size of the LLaMA model. Defines the number of different tokens that can be represented by the - `inputs_ids` passed when calling [`Lfm2Model`] - hidden_size (`int`, *optional*, defaults to 2048): - Dimension of the hidden representations. - intermediate_size (`int`, *optional*, defaults to 7168): - Dimension of the MLP representations. - moe_intermediate_size (`int`, *optional*, defaults to 1792): - Intermediate size of the routed expert. - num_hidden_layers (`int`, *optional*, defaults to 32): - Number of hidden layers in the Transformer decoder. - pad_token_id (`int`, *optional*, defaults to 0): - Padding token id. - bos_token_id (`int`, *optional*, defaults to 1): - Beginning of stream token id. - eos_token_id (`int`, *optional*, defaults to 2): - End of stream token id. - tie_word_embeddings (`bool`, *optional*, defaults to `True`): - Whether to tie weight embeddings - rope_parameters (`dict`, *optional*): - The parameters of the RoPE embeddings. - max_position_embeddings (`int`, *optional*, defaults to 128000): - The maximum sequence length that this model might ever be used with. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values attentions (not used by all models). Only - relevant if `config.is_decoder=True`. - norm_eps (`float`, *optional*, defaults to 1e-05): - The epsilon used by the rms normalization layers. - num_attention_heads (`int`, *optional*, defaults to 32): - Number of attention heads for each attention layer in the Transformer decoder. - num_key_value_heads (`int`, *optional*, defaults to 8): - This is the number of key_value heads that should be used to implement Grouped Query Attention. If - `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if - `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When - converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed - by meanpooling all the original heads within that group. For more details, check out [this - paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to - `num_attention_heads`. - conv_bias (`bool`, *optional*, defaults to `False`): - Whether to use bias in the conv layers. - conv_L_cache (`int`, *optional*, defaults to 3): - L_cache dim in the conv layers. - num_dense_layers (`int`, *optional*, defaults to 2): - Number of dense Lfm2MoeMLP layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head). - num_experts_per_tok (`int`, *optional*, defaults to 4): - Number of selected experts. - num_experts (`int`, *optional*, defaults to 32): - Number of routed experts. - use_expert_bias (`bool`, *optional*, defaults to `True`): - Whether to use the expert bias on the routing weights. - routed_scaling_factor (`float`, *optional*, defaults to 1.0): - Scaling factor for routed experts in MoE models. - norm_topk_prob (`bool`, *optional*, defaults to `True`): - Whether to normalize the topk probabilities. - layer_types (`Optional`, *optional*): - Type of each layers. - - ```python - >>> from transformers import Lfm2MoeModel, Lfm2MoeConfig - - >>> # Initializing a LFM2 Moe model - >>> configuration = Lfm2MoeConfig() - - >>> # Initializing a model from the LFM2-8B-A1B style configuration - >>> model = Lfm2MoeModel(configuration) - - >>> # Accessing the model configuration - >>> configuration = model.config - ``` - - """ # noqa: E501 - - model_type = "lfm2_moe" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size: int = 65536, - hidden_size: int = 2048, - intermediate_size: int = 7168, - moe_intermediate_size: int = 1792, - num_hidden_layers: int = 32, - pad_token_id: int = 0, - bos_token_id: int = 1, - eos_token_id: int = 2, - tie_word_embeddings: bool = True, - rope_parameters: dict[str, Any] | None = None, - max_position_embeddings: int = 128_000, - use_cache: bool = True, - norm_eps: float = 0.00001, - num_attention_heads: int = 32, - num_key_value_heads: int = 8, - conv_bias: bool = False, - conv_L_cache: int = 3, - num_dense_layers: int = 2, - num_experts_per_tok: int = 4, - num_experts: int = 32, - use_expert_bias: bool = True, - routed_scaling_factor: float = 1.0, - norm_topk_prob: bool = True, - layer_types: list[str] | None = None, - **kwargs, - ): - self.vocab_size = vocab_size - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - rope_theta = kwargs.pop("rope_theta", 1000000.0) - if rope_parameters is None: - rope_parameters = {"rope_type": "default", "rope_theta": rope_theta} - self.rope_parameters = rope_parameters - self.max_position_embeddings = max_position_embeddings - self.use_cache = use_cache - self.norm_eps = norm_eps - - # attn operator config - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - - # custom operator config - self.conv_bias = conv_bias - self.conv_L_cache = conv_L_cache - - # moe config - self.num_dense_layers = num_dense_layers - self.moe_intermediate_size = moe_intermediate_size - self.num_experts_per_tok = num_experts_per_tok - self.num_experts = num_experts - self.use_expert_bias = use_expert_bias - self.routed_scaling_factor = routed_scaling_factor - self.norm_topk_prob = norm_topk_prob - self.layer_types = layer_types - - tie_word_embeddings = kwargs.get( - "tie_embedding", tie_word_embeddings - ) # to fit original config keys - super().__init__( - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) - - -__all__ = ["Lfm2MoeConfig"] diff --git a/vllm/transformers_utils/configs/nemotron.py b/vllm/transformers_utils/configs/nemotron.py deleted file mode 100644 index 135a0e1596dc..000000000000 --- a/vllm/transformers_utils/configs/nemotron.py +++ /dev/null @@ -1,219 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Copyright 2024 HuggingFace Inc. team. All rights reserved. -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Nemotron model configuration.""" - -from transformers import PreTrainedConfig -from transformers.utils import logging - -logger = logging.get_logger(__name__) - - -class NemotronConfig(PreTrainedConfig): - r"""This is the configuration class to store the configuration of a - [`NemotronModel`]. It is used to instantiate a Nemotron model - according to the specified arguments, defining the model architecture. - Instantiating a configuration with the defaults will yield a similar - configuration to that of the Nemotron-8B. - - Configuration objects inherit from [`PreTrainedConfig`] and can be - used to control the model outputs. Read the documentation from - [`PreTrainedConfig`] for more information. - - - Args: - vocab_size (`int`, *optional*, defaults to 256000): - Vocabulary size of the Nemotron model. Defines the number of - different tokens that can be represented by the - `inputs_ids` passed when calling [`NemotronModel`] - hidden_size (`int`, *optional*, defaults to 6144): - Dimension of the hidden representations. - intermediate_size (`int`, *optional*, defaults to 24576): - Dimension of the MLP representations. - num_hidden_layers (`int`, *optional*, defaults to 32): - Number of hidden layers in the Transformer decoder. - num_attention_heads (`int`, *optional*, defaults to 48): - Number of attention heads for each attention layer in the - Transformer decoder. - head_dim (`int`, *optional*): - Projection weights dimension in multi-head attention. Set to - hidden_size // num_attention_heads if None - num_key_value_heads (`int`, *optional*): - This is the number of key_value heads that should be used to - implement Grouped Query Attention. If - `num_key_value_heads=num_attention_heads`, the model will use - Multi Head Attention (MHA), if - `num_key_value_heads=1 the model will use Multi Query Attention - (MQA) otherwise GQA is used. When converting a multi-head - checkpoint to a GQA checkpoint, each group key and value - head should be constructed by meanpooling all the original - heads within that group. For more details checkout - [this paper](https://arxiv.org/pdf/2305.13245.pdf). If it - is not specified, will default to `num_attention_heads`. - hidden_act (`str` or `function`, *optional*, defaults to `"relu2"`): - The non-linear activation function (function or string) in the - decoder. - max_position_embeddings (`int`, *optional*, defaults to 4096): - The maximum sequence length that this model might ever be used - with. - initializer_range (`float`, *optional*, defaults to 0.0134): - The standard deviation of the truncated_normal_initializer for - initializing all weight matrices. - norm_eps (`float`, *optional*, defaults to 1e-05): - The epsilon used by the normalization layers. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values - attentions (not used by all models). Only relevant if - `config.is_decoder=True`. - pad_token_id (`int`, *optional*): - Padding token id. - bos_token_id (`int`, *optional*, defaults to 2): - Beginning of stream token id. - eos_token_id (`int`, *optional*, defaults to 3): - End of stream token id. - tie_word_embeddings (`bool`, *optional*, defaults to `False`): - Whether to tie weight embeddings - rope_parameters (`dict`, *optional*): - The parameters of the RoPE embeddings. Expected contents: - `rope_theta` (`float`): The base period of the RoPE embeddings. - `rope_type` (`str`): - The sub-variant of RoPE to use. Can be one of ['default', 'linear', - 'dynamic', 'yarn', 'longrope', 'llama3'], with 'default' being the - original RoPE implementation. - `partial_rotary_factor` (`float`, *optional*, defaults to 0.5): - Percentage of the query and keys which will have rotary embedding. - attention_bias (`bool`, *optional*, defaults to `False`): - Whether to use a bias in the query, key, value and output - projection layers during self-attention. - attention_dropout (`float`, *optional*, defaults to 0.0): - The dropout ratio for the attention probabilities. - mlp_bias (`bool`, *optional*, defaults to `False`): - Whether to use a bias in up_proj and down_proj layers in the MLP - layers. - - ```python - >>> from transformers import NemotronModel, NemotronConfig - >>> # Initializing a Nemotron nemotron-15b style configuration - >>> configuration = NemotronConfig() - >>> # Initializing a model from the nemotron-15b style configuration - >>> model = NemotronModel(configuration) - >>> # Accessing the model configuration - >>> configuration = model.config - ``` - - """ - - model_type = "nemotron" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size=256000, - hidden_size=6144, - intermediate_size=24576, - num_hidden_layers=32, - num_attention_heads=48, - head_dim=None, - num_key_value_heads=None, - hidden_act="relu2", - max_position_embeddings=4096, - initializer_range=0.0134, - norm_eps=1e-5, - use_cache=True, - pad_token_id=None, - bos_token_id=2, - eos_token_id=3, - tie_word_embeddings=False, - rope_parameters=None, - attention_bias=False, - attention_dropout=0.0, - mlp_bias=False, - **kwargs, - ): - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - head_dim = head_dim or kwargs.get("kv_channels") - self.head_dim = ( - head_dim if head_dim is not None else (hidden_size // num_attention_heads) - ) - - # for backward compatibility - if num_key_value_heads is None: - num_key_value_heads = num_attention_heads - - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.norm_eps = norm_eps - self.use_cache = use_cache - # Try to set `rope_scaling` if available, otherwise use `rope_parameters` - rope_scaling = kwargs.pop("rope_scaling", None) - rope_parameters = rope_scaling or rope_parameters or {"rope_type": "default"} - rope_theta = kwargs.pop("rope_theta", 10000.0) - if "rope_theta" not in rope_parameters: - rope_parameters["rope_theta"] = rope_theta - # for backward compatibility - partial_rotary_factor = ( - kwargs.get("rope_percent") - or kwargs.get("rope_percentage") - or kwargs.get("partial_rotary_factor") - or 0.5 - ) - if "partial_rotary_factor" not in rope_parameters: - rope_parameters["partial_rotary_factor"] = partial_rotary_factor - self.rope_parameters = rope_parameters - self._rope_parameters_validation() - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.mlp_bias = mlp_bias - - super().__init__( - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) - - def _rope_parameters_validation(self): - """Validate the `rope_parameters` configuration.""" - if self.rope_parameters is None: - return - - rope_type: str | None = self.rope_parameters.get("rope_type", None) - factor: float | None = self.rope_parameters.get("factor", None) - - if rope_type not in {"default", "linear", "dynamic"}: - raise ValueError( - "`rope_type` must be one of ['default', 'linear', 'dynamic'], " - f"got {rope_type}" - ) - if rope_type != "default": - if factor is None: - raise ValueError( - "If `rope_type` is not 'default', `rope_parameters` " - "must include a `factor` field. Got `None`." - ) - if not isinstance(factor, float) or factor <= 1.0: - raise ValueError( - "`rope_parameters`'s factor field must be a float > 1, got " - f"{factor}" - ) diff --git a/vllm/transformers_utils/configs/nemotron_h.py b/vllm/transformers_utils/configs/nemotron_h.py deleted file mode 100644 index c25a5d54e3d7..000000000000 --- a/vllm/transformers_utils/configs/nemotron_h.py +++ /dev/null @@ -1,288 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Copyright 2024 HuggingFace Inc. team. All rights reserved. -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""NemotronH model configuration.""" - -import regex as re -from transformers.configuration_utils import PreTrainedConfig -from transformers.utils import logging - -logger = logging.get_logger(__name__) - - -class NemotronHConfig(PreTrainedConfig): - r"""This is the configuration class to store the configuration of a - [`NemotronHModel`]. It is used to instantiate a NemotronH model according - to the specified arguments, defining the model architecture. Instantiating - a configuration with the defaults will yield a similar configuration to - that of the NemotronH-v0.1 model. - - Args: - vocab_size (`int`, *optional*, defaults to 131072): - Vocabulary size of the NemotronH model. Defines the number of - different tokens that can be represented by the `inputs_ids` - passed when calling [`NemotronHModel`] - tie_word_embeddings (`bool`, *optional*, defaults to `False`): - Whether the model's input and output word embeddings should be - tied. Note that this is only relevant if the model has an output - word embedding layer. - hidden_size (`int`, *optional*, defaults to 4096): - Dimension of the hidden representations. - intermediate_size (`int`, *optional*, defaults to 21504): - Dimension of the MLP representations. - num_hidden_layers (`int`, *optional*, defaults to 52): - Number of hidden layers in the Transformer encoder. - hybrid_override_pattern (`str`, *optional*, defaults to - `"M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-"`): - The pattern of the hybrid model. The pattern is a string of - characters where each character represents - M: Mamba2, *: Attention, -: MLP - mtp_hybrid_override_pattern (`str`, *optional*, defaults to `"*E"`): - The pattern of the MTP layers. - num_attention_heads (`int`, *optional*, defaults to 32): - Number of attention heads for each attention layer in the - Transformer encoder. - attention_head_dim (`int`, *optional*, defaults to 128): - Dimension of each attention head. - num_key_value_heads (`int`, *optional*, defaults to 8): - This is the number of key_value heads that should be used to - implement Grouped Query Attention. If - `num_key_value_heads=num_attention_heads`, the model will use - Multi Head Attention (MHA), if `num_key_value_heads=1` the model - will use Multi Query Attention (MQA) otherwise GQA is used. - mlp_hidden_act (`str`, *optional*, defaults to "relu2"): - The non-linear activation function in the MLP layers. - attention_bias (`bool`, *optional*, defaults to `False`): - Whether to use bias in attention layers. - mlp_bias (`bool`, *optional*, defaults to `False`): - Whether to use bias in MLP layers. - use_bias (`bool`, *optional*, defaults to `False`): - Whether to use bias in the model. - initializer_range (`float`, *optional*, defaults to 0.02): - The standard deviation of the truncated_normal_initializer for - initializing all weight matrices. - layer_norm_epsilon (`float`, *optional*, defaults to 1e-5): - The epsilon used by the layer normalization layers. - residual_in_fp32 (`bool`, *optional*, defaults to `False`): - Whether or not residuals should be in `float32`. If set to `False` - residuals will keep the same `dtype` as the rest of the model. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values - attentions (not used by all models). Only relevant if - `config.is_decoder=True`. - num_logits_to_keep (`int` or `None`, *optional*, defaults to 1): - Number of prompt logits to calculate during generation. If `None`, - all logits will be calculated. If an integer value, only last - `num_logits_to_keep` logits will be calculated. - pad_token_id (`int`, *optional*, defaults to 0): - The id of the padding token. - bos_token_id (`int`, *optional*, defaults to 1): - The id of the "beginning-of-sequence" token. - eos_token_id (`int`, *optional*, defaults to 2): - The id of the "end-of-sequence" token. - sliding_window (`int`, *optional*, defaults to None): - Sliding window attention window size. - max_position_embeddings (`int`, *optional*, defaults to 4096): - The maximum sequence length that this model might ever be used - with. - attention_dropout (`float`, *optional*, defaults to 0.0): - The dropout ratio for the attention probabilities. - hidden_dropout (`float`, *optional*, defaults to 0.0): - The dropout ratio for the hidden states. - use_mamba_kernels (`bool`, *optional*, defaults to `True`): - Flag indicating whether or not to use the fast mamba kernels. - These are available only if `mamba-ssm` and `causal-conv1d` - are installed, and the mamba modules are running on a CUDA device. - ssm_state_size (`int`, *optional*, defaults to 128): - The dimension of the mamba state space latents. - mamba_num_heads (`int`, *optional*, defaults to 128): - Number of heads in Mamba layers. - mamba_n_groups (`int`, *optional*, defaults to 8): - Number of groups in Mamba layers. - mamba_head_dim (`int`, *optional*, defaults to 64): - Dimension of each Mamba head. - mamba_d_conv (`int`, *optional*, defaults to 4): - The size of the mamba convolution kernel. - mamba_expand (`int`, *optional*, defaults to 2): - Expanding factor used to determine the mamba intermediate size. - mamba_hidden_act (`str`, *optional*, defaults to "silu"): - The non-linear activation function in the Mamba layers. - mamba_dt_min (`float`, *optional*, defaults to 0.001): - Minimum value for the time step in Mamba. - mamba_dt_max (`float`, *optional*, defaults to 0.1): - Maximum value for the time step in Mamba. - mamba_dt_limit (`tuple`, *optional*, defaults to (0.0, float("inf"))): - Limits for the time step in Mamba. - mamba_dt_init_floor (`float`, *optional*, defaults to 1e-4): - Floor value for time step initialization in Mamba. - mamba_conv_bias (`bool`, *optional*, defaults to `True`): - Whether to use bias in the convolution layer of the mamba mixer - block. - mamba_proj_bias (`bool`, *optional*, defaults to `False`): - Whether to use bias in the input and output projections of the - mamba mixer block. - mamba_chunk_size (`int`, *optional*, defaults to 256): - Size of chunks for Mamba processing. - rescale_prenorm_residual (`bool`, *optional*, defaults to `True`): - Whether to rescale the pre-normalization residual connections. - - """ - - model_type = "nemotron_h" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size=131072, - tie_word_embeddings=False, - hidden_size=4096, - intermediate_size=21504, - num_hidden_layers=52, - hybrid_override_pattern="M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-", - mtp_hybrid_override_pattern="*E", - num_attention_heads=32, - head_dim=128, - num_key_value_heads=8, # nemo: num_query_groups - mlp_hidden_act="relu2", - attention_bias=False, - mlp_bias=False, - use_bias=False, - initializer_range=0.02, # nemo: init_method_std - layer_norm_epsilon=1e-5, # nemo: layernorm_epsilon - residual_in_fp32=False, # Megatron Core default value - use_cache=True, - num_logits_to_keep=1, - pad_token_id=0, - bos_token_id=1, - eos_token_id=2, - sliding_window=None, - max_position_embeddings=4096, - attention_dropout=0.0, - hidden_dropout=0.0, # * ADDED - use_mamba_kernels=True, - ssm_state_size=128, # mamba_state_size - mamba_num_heads=128, - mamba_n_groups=8, # nemo: mamba_ssm_ngroups = num_heads - mamba_head_dim=64, - mamba_d_conv=4, - mamba_expand=2, - mamba_hidden_act="silu", - mamba_dt_min=0.001, - mamba_dt_max=0.1, - mamba_dt_limit=(0.0, float("inf")), - mamba_dt_init_floor=1e-4, - mamba_conv_bias=True, - mamba_proj_bias=False, - mamba_chunk_size=256, - rescale_prenorm_residual=True, - n_routed_experts=8, - n_shared_experts=1, - moe_intermediate_size=7688, - moe_shared_expert_intermediate_size=7688, - moe_latent_size=None, - num_experts_per_tok=2, - routed_scaling_factor=1.0, - n_group=1, - topk_group=1, - norm_topk_prob=True, - **kwargs, - ): - self.vocab_size = vocab_size - self.tie_word_embeddings = tie_word_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.hybrid_override_pattern = hybrid_override_pattern - self.mtp_hybrid_override_pattern = mtp_hybrid_override_pattern - self.num_attention_heads = num_attention_heads - self.head_dim = head_dim - self.sliding_window = sliding_window - self.max_position_embeddings = max_position_embeddings - self.attention_dropout = attention_dropout - self.hidden_dropout = hidden_dropout - - # Validate hybrid_override_pattern - # M: Mamba2, *: Attention, -: MLP - assert len(self.hybrid_override_pattern) == self.num_hidden_layers, ( - "hybrid_override_pattern must have same length as num_hidden_layers" - ) - assert re.match(r"^[*-ME]+$", self.hybrid_override_pattern), ( - "hybrid_override_pattern must only contain characters 'M', '*', '-', or 'E'" - ) - # for backward compatibility - if num_key_value_heads is None: - num_key_value_heads = num_attention_heads - - self.num_key_value_heads = num_key_value_heads - self.mlp_hidden_act = mlp_hidden_act - self.attention_bias = attention_bias - self.mlp_bias = mlp_bias - self.use_bias = use_bias - self.initializer_range = initializer_range - self.layer_norm_epsilon = layer_norm_epsilon - self.residual_in_fp32 = residual_in_fp32 - - self.use_cache = use_cache - self.num_logits_to_keep = num_logits_to_keep - - self.use_mamba_kernels = use_mamba_kernels - self.n_groups = mamba_n_groups - self.mamba_head_dim = mamba_head_dim - self.ssm_state_size = ssm_state_size - self.mamba_num_heads = mamba_num_heads - self.conv_kernel = mamba_d_conv - self.expand = mamba_expand - self.mamba_hidden_act = mamba_hidden_act - self.time_step_min = mamba_dt_min - self.time_step_max = mamba_dt_max - self.time_step_limit = mamba_dt_limit - self.time_step_floor = mamba_dt_init_floor - self.use_conv_bias = mamba_conv_bias - self.mamba_proj_bias = mamba_proj_bias - self.chunk_size = mamba_chunk_size - self.rescale_prenorm_residual = rescale_prenorm_residual - self.n_routed_experts = n_routed_experts - self.n_shared_experts = n_shared_experts - self.moe_intermediate_size = moe_intermediate_size - self.moe_shared_expert_intermediate_size = moe_shared_expert_intermediate_size # noqa: E501 - self.moe_latent_size = moe_latent_size - self.num_experts_per_tok = num_experts_per_tok - self.routed_scaling_factor = routed_scaling_factor - self.n_group = n_group - self.topk_group = topk_group - self.norm_topk_prob = norm_topk_prob - - super().__init__( - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) - - @property - def layers_block_type(self): - return [ - "mamba" - if self.hybrid_override_pattern[i] == "M" - else "attention" - if self.hybrid_override_pattern[i] == "*" - else "mlp" - if self.hybrid_override_pattern[i] == "-" - else "moe" - for i in range(self.num_hidden_layers) - ] diff --git a/vllm/transformers_utils/configs/olmo_hybrid.py b/vllm/transformers_utils/configs/olmo_hybrid.py deleted file mode 100644 index 62d3c4e0d5d1..000000000000 --- a/vllm/transformers_utils/configs/olmo_hybrid.py +++ /dev/null @@ -1,285 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - - -from transformers.configuration_utils import PreTrainedConfig - - -class OlmoHybridConfig(PreTrainedConfig): - r"""Configuration class for [`OlmoHybridModel`]. It is used to - instantiate an OLMo Hybrid model according to the specified - arguments, defining the model architecture. Instantiating a - configuration with the defaults will yield a similar - configuration to that of the - [allenai/Olmo-Hybrid-7B](https://huggingface.co/allenai/Olmo-Hybrid-7B) - model. - - Configuration objects inherit from [`PreTrainedConfig`] and - can be used to control the model outputs. Read the - documentation from [`PreTrainedConfig`] for more information. - - Args: - vocab_size (`int`, *optional*, defaults to 100352): - Vocabulary size of the OlmoHybrid model. Defines - the number of different tokens that can be - represented by the `inputs_ids` passed when - calling [`OlmoHybridModel`]. - hidden_size (`int`, *optional*, defaults to 3840): - Dimension of the hidden representations. - intermediate_size (`int`, *optional*, - defaults to 11008): - Dimension of the MLP representations. - num_hidden_layers (`int`, *optional*, - defaults to 32): - Number of hidden layers in the Transformer - decoder. - num_attention_heads (`int`, *optional*, - defaults to 30): - Number of attention heads for each attention - layer in the Transformer decoder. - num_key_value_heads (`int`, *optional*): - This is the number of key_value heads that - should be used to implement Grouped Query - Attention. If - `num_key_value_heads=num_attention_heads`, - the model will use Multi Head Attention (MHA), - if `num_key_value_heads=1` the model will use - Multi Query Attention (MQA) otherwise GQA is - used. When converting a multi-head checkpoint - to a GQA checkpoint, each group key and value - head should be constructed by meanpooling all - the original heads within that group. For more - details, check out - [this paper](https://huggingface.co/papers/2305.13245). - If it is not specified, will default to - `num_attention_heads`. - hidden_act (`str` or `function`, *optional*, - defaults to `"silu"`): - The non-linear activation function (function - or string) in the decoder. - max_position_embeddings (`int`, *optional*, - defaults to 65536): - The maximum sequence length that this model - might ever be used with. - initializer_range (`float`, *optional*, - defaults to 0.02): - The standard deviation of the - truncated_normal_initializer for initializing - all weight matrices. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last - key/values attentions (not used by all models). - Only relevant if `config.is_decoder=True`. - pad_token_id (`int`, *optional*, - defaults to 100277): - Padding token id. - bos_token_id (`int`, *optional*): - Beginning of stream token id. - eos_token_id (`int`, *optional*, - defaults to 100257): - End of stream token id. - tie_word_embeddings (`bool`, *optional*, - defaults to `False`): - Whether to tie weight embeddings. - rope_parameters (`RopeParameters`, *optional*): - Dictionary containing the configuration - parameters for the RoPE embeddings. Can be - `None` to disable RoPE. - attention_bias (`bool`, *optional*, - defaults to `False`): - Whether to use a bias in the query, key, value - and output projection layers during - self-attention. - attention_dropout (`float`, *optional*, - defaults to 0.0): - The dropout ratio for the attention - probabilities. - rms_norm_eps (`float`, *optional*, - defaults to 1e-06): - The epsilon used by the rms normalization - layers. - layer_types (`list`, *optional*): - Attention pattern for each layer. Can contain - `"full_attention"` or `"linear_attention"`. - Defaults to linear attention for most layers - with full attention for every 4th layer. - linear_num_key_heads (`int`, *optional*): - Number of key heads for the linear attention - layers. Defaults to `num_attention_heads`. - linear_num_value_heads (`int`, *optional*): - Number of value heads for the linear attention - layers. Defaults to `num_attention_heads`. - linear_key_head_dim (`int`, *optional*): - Dimension of each key head in linear attention - layers. Defaults to - `0.75 * hidden_size / linear_num_key_heads`. - linear_value_head_dim (`int`, *optional*): - Dimension of each value head in linear - attention layers. Defaults to - `2 * linear_key_head_dim`. - linear_a_log_min (`float`, *optional*, - defaults to 0.0): - Minimum value for uniform initialization of - A_log in GatedDeltaNet layers. - linear_a_log_max (`float`, *optional*, - defaults to 16.0): - Maximum value for uniform initialization of - A_log in GatedDeltaNet layers. - linear_dt_min (`float`, *optional*, - defaults to 0.001): - Minimum value for dt initialization in - GatedDeltaNet layers. - linear_dt_max (`float`, *optional*, - defaults to 0.1): - Maximum value for dt initialization in - GatedDeltaNet layers. - linear_dt_init_floor (`float`, *optional*, - defaults to 0.0001): - Floor value for clamping dt during - initialization in GatedDeltaNet layers. - linear_conv_kernel_dim (`int`, *optional*, - defaults to 4): - Kernel size for the short convolution applied - to queries, keys, and values in linear - attention layers. - linear_allow_neg_eigval (`bool`, *optional*, - defaults to `True`): - Whether to allow negative eigenvalues in the - GatedDeltaNet recurrence. When `True`, the - beta parameter is scaled by 2.0 to allow - values in range [0, 2] instead of [0, 1]. - ```python - >>> from transformers import ( - ... OlmoHybridModel, - ... OlmoHybridConfig, - ... ) - - >>> configuration = OlmoHybridConfig() - >>> model = OlmoHybridModel(configuration) - >>> configuration = model.config - ``` - - """ - - model_type = "olmo_hybrid" - keys_to_ignore_at_inference = ["past_key_values"] - base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", - "layers.*.self_attn.k_proj": "colwise_gather_output", - "layers.*.self_attn.v_proj": "colwise_gather_output", - "layers.*.self_attn.o_proj": "rowwise_split_input", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", - } - base_model_pp_plan = { - "embed_tokens": (["input_ids"], ["inputs_embeds"]), - "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), - "norm": (["hidden_states"], ["hidden_states"]), - } - - def __init__( - self, - vocab_size: int | None = 100352, - hidden_size: int | None = 3840, - intermediate_size: int | None = 11008, - num_hidden_layers: int | None = 32, - num_attention_heads: int | None = 30, - num_key_value_heads: int | None = None, - hidden_act: str | None = "silu", - max_position_embeddings: int | None = 65536, - initializer_range: float | None = 0.02, - use_cache: bool | None = True, - pad_token_id: int | None = 100277, - bos_token_id: int | None = None, - eos_token_id: int | None = 100257, - tie_word_embeddings: bool | None = False, - rope_parameters=None, - attention_bias: bool | None = False, - attention_dropout: float | None = 0.0, - rms_norm_eps: float | None = 1e-06, - layer_types: list[str] | None = None, - linear_num_key_heads: int | None = None, - linear_num_value_heads: int | None = None, - linear_key_head_dim: int | None = None, - linear_value_head_dim: int | None = None, - linear_a_log_min: float = 0.0, - linear_a_log_max: float = 16.0, - linear_dt_min: float = 0.001, - linear_dt_max: float = 0.1, - linear_dt_init_floor: float = 1e-4, - linear_conv_kernel_dim: int = 4, - linear_allow_neg_eigval: bool = True, - **kwargs, - ): - super().__init__(**kwargs) - - assert num_hidden_layers is not None - assert hidden_size is not None - assert num_attention_heads is not None - - if layer_types is None: - # Default: linear attention for most layers, full attention every 4th layer - layer_types = ["linear_attention"] * int(num_hidden_layers) - for i in range(int(num_hidden_layers)): - if i % 4 == 3: - layer_types[i] = "full_attention" - # Ensure at least one full attention layer for small num_hidden_layers - if "full_attention" not in layer_types: - layer_types[-1] = "full_attention" - - self.layer_types = layer_types - self.validate_layer_type() - if "linear_attention" not in layer_types: - raise ValueError( - "OLMoHybrid expects at least one 'linear_attention' layer." - ) - if all(t == "linear_attention" for t in layer_types): - raise ValueError("OLMoHybrid expects at least one attention layer.") - - self.layer_types = layer_types - - if linear_num_key_heads is None: - linear_num_key_heads = num_attention_heads - if linear_num_value_heads is None: - linear_num_value_heads = num_attention_heads - if linear_key_head_dim is None: - linear_key_head_dim = int(0.75 * hidden_size / linear_num_key_heads) - if linear_value_head_dim is None: - linear_value_head_dim = 2 * linear_key_head_dim - - self.linear_num_key_heads = linear_num_key_heads - self.linear_num_value_heads = linear_num_value_heads - self.linear_key_head_dim = linear_key_head_dim - self.linear_value_head_dim = linear_value_head_dim - self.linear_a_log_min = linear_a_log_min - self.linear_a_log_max = linear_a_log_max - self.linear_dt_min = linear_dt_min - self.linear_dt_max = linear_dt_max - self.linear_dt_init_floor = linear_dt_init_floor - self.linear_conv_kernel_dim = linear_conv_kernel_dim - self.linear_allow_neg_eigval = linear_allow_neg_eigval - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - - # for backward compatibility - if num_key_value_heads is None: - num_key_value_heads = num_attention_heads - - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.use_cache = use_cache - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.rope_parameters = rope_parameters - - self.tie_word_embeddings = tie_word_embeddings - self.pad_token_id = pad_token_id - self.bos_token_id = bos_token_id - self.eos_token_id = eos_token_id diff --git a/vllm/transformers_utils/configs/ovis.py b/vllm/transformers_utils/configs/ovis.py index bc84bfe812fc..d5662d7ef542 100644 --- a/vllm/transformers_utils/configs/ovis.py +++ b/vllm/transformers_utils/configs/ovis.py @@ -2,69 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -# adapted from https://huggingface.co/AIDC-AI/Ovis2-1B/blob/main/configuration_aimv2.py -# and https://huggingface.co/AIDC-AI/Ovis2-1B/blob/main/configuration_ovis.py -# Ovis Config with AimV2 config registration removed for Transformers compatibility +# adapted from https://huggingface.co/AIDC-AI/Ovis2-1B/blob/main/configuration_ovis.py from typing import Any -from transformers import AutoConfig, PreTrainedConfig - - -class AIMv2Config(PreTrainedConfig): - """This is the configuration class to store the configuration of an [`AIMv2Model`]. - Instantiating a configuration with the defaults will yield a similar configuration - to that of the [apple/aimv2-large-patch14-224](https://huggingface.co/apple/aimv2-large-patch14-224). - - Args: - hidden_size: Dimension of the hidden representations. - intermediate_size: Dimension of the SwiGLU representations. - num_hidden_layers: Number of hidden layers in the Transformer. - num_attention_heads: Number of attention heads for each attention layer - in the Transformer. - num_channels: Number of input channels. - image_size: Image size. - patch_size: Patch size. - rms_norm_eps: Epsilon value used for the RMS normalization layer. - attention_dropout: Dropout ratio for attention probabilities. - projection_dropout: Dropout ratio for the projection layer after the attention. - qkv_bias: Whether to add a bias to the queries, keys and values. - use_bias: Whether to add a bias in the feed-forward and projection layers. - kwargs: Keyword arguments for the [`PreTrainedConfig`]. - - """ - - model_type: str = "aimv2" - - def __init__( - self, - hidden_size: int = 1024, - intermediate_size: int = 2816, - num_hidden_layers: int = 24, - num_attention_heads: int = 8, - num_channels: int = 3, - image_size: int = 224, - patch_size: int = 14, - rms_norm_eps: float = 1e-5, - attention_dropout: float = 0.0, - projection_dropout: float = 0.0, - qkv_bias: bool = False, - use_bias: bool = False, - **kwargs: Any, - ): - super().__init__(**kwargs) - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_channels = num_channels - self.patch_size = patch_size - self.image_size = image_size - self.attention_dropout = attention_dropout - self.rms_norm_eps = rms_norm_eps - - self.projection_dropout = projection_dropout - self.qkv_bias = qkv_bias - self.use_bias = use_bias +from transformers import Aimv2VisionConfig, AutoConfig, PreTrainedConfig # ---------------------------------------------------------------------- @@ -96,14 +37,13 @@ def __init__( f"expect `backbone_config` to be instance of PreTrainedConfig or dict, but got {type(backbone_config)} type" ) if not isinstance(backbone_config, PreTrainedConfig): - model_type = backbone_config["model_type"] - if model_type != "aimv2": - backbone_config.pop("model_type") + model_type = backbone_config.pop("model_type") + if model_type == "aimv2": + backbone_config = Aimv2VisionConfig(**backbone_config) + else: backbone_config = AutoConfig.for_model( model_type, **backbone_config ) - else: - backbone_config = AIMv2Config(**backbone_config) self.backbone_config = backbone_config self.hidden_stride = hidden_stride diff --git a/vllm/transformers_utils/configs/qwen3_5.py b/vllm/transformers_utils/configs/qwen3_5.py deleted file mode 100644 index 57f64f0f7814..000000000000 --- a/vllm/transformers_utils/configs/qwen3_5.py +++ /dev/null @@ -1,193 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. -# All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Qwen3.5 model configuration.""" - -from transformers.configuration_utils import PreTrainedConfig - - -class Qwen3_5TextConfig(PreTrainedConfig): - model_type = "qwen3_5_text" - keys_to_ignore_at_inference = ["past_key_values"] - - base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", - } - base_model_pp_plan = { - "embed_tokens": (["input_ids"], ["inputs_embeds"]), - "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), - "norm": (["hidden_states"], ["hidden_states"]), - } - base_config_key = "text_config" - - def __init__( - self, - vocab_size=248320, - hidden_size=4096, - intermediate_size=12288, - num_hidden_layers=32, - num_attention_heads=16, - num_key_value_heads=4, - hidden_act="silu", - max_position_embeddings=32768, - initializer_range=0.02, - rms_norm_eps=1e-6, - use_cache=True, - tie_word_embeddings=False, - rope_parameters=None, - attention_bias=False, - attention_dropout=0.0, - head_dim=256, - linear_conv_kernel_dim=4, - linear_key_head_dim=128, - linear_value_head_dim=128, - linear_num_key_heads=16, - linear_num_value_heads=32, - layer_types=None, - pad_token_id=None, - bos_token_id=None, - eos_token_id=None, - **kwargs, - ): - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.use_cache = use_cache - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.head_dim = head_dim - self.rope_parameters = rope_parameters - kwargs.setdefault("partial_rotary_factor", 0.25) - - self.layer_types = layer_types - if self.layer_types is None: - interval_pattern = kwargs.get("full_attention_interval", 4) - self.layer_types = [ - "linear_attention" - if bool((i + 1) % interval_pattern) - else "full_attention" - for i in range(self.num_hidden_layers) - ] - kwargs["ignore_keys_at_rope_validation"] = { - "mrope_section", - "mrope_interleaved", - } - self.validate_layer_type() - - # linear attention part - self.linear_conv_kernel_dim = linear_conv_kernel_dim - self.linear_key_head_dim = linear_key_head_dim - self.linear_value_head_dim = linear_value_head_dim - self.linear_num_key_heads = linear_num_key_heads - self.linear_num_value_heads = linear_num_value_heads - super().__init__(**kwargs) - # Set these AFTER super().__init__() because transformers v4's - # PreTrainedConfig.__init__ has these as explicit params with different - # defaults (e.g. tie_word_embeddings=True) that would overwrite our values. - self.pad_token_id = pad_token_id - self.bos_token_id = bos_token_id - self.eos_token_id = eos_token_id - self.tie_word_embeddings = tie_word_embeddings - - -class Qwen3_5VisionConfig(PreTrainedConfig): - model_type = "qwen3_5" - base_config_key = "vision_config" - - def __init__( - self, - depth=27, - hidden_size=1152, - hidden_act="gelu_pytorch_tanh", - intermediate_size=4304, - num_heads=16, - in_channels=3, - patch_size=16, - spatial_merge_size=2, - temporal_patch_size=2, - out_hidden_size=3584, - num_position_embeddings=2304, - initializer_range=0.02, - **kwargs, - ): - super().__init__(**kwargs) - - self.depth = depth - self.hidden_size = hidden_size - self.hidden_act = hidden_act - self.intermediate_size = intermediate_size - self.num_heads = num_heads - self.in_channels = in_channels - self.patch_size = patch_size - self.spatial_merge_size = spatial_merge_size - self.temporal_patch_size = temporal_patch_size - self.out_hidden_size = out_hidden_size - self.num_position_embeddings = num_position_embeddings - self.initializer_range = initializer_range - - -class Qwen3_5Config(PreTrainedConfig): - model_type = "qwen3_5" - sub_configs = { - "vision_config": Qwen3_5VisionConfig, - "text_config": Qwen3_5TextConfig, - } - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - text_config=None, - vision_config=None, - image_token_id=248056, - video_token_id=248057, - vision_start_token_id=248053, - vision_end_token_id=248054, - tie_word_embeddings=False, - **kwargs, - ): - if isinstance(vision_config, dict): - self.vision_config = self.sub_configs["vision_config"](**vision_config) - elif vision_config is None: - self.vision_config = self.sub_configs["vision_config"]() - - if isinstance(text_config, dict): - self.text_config = self.sub_configs["text_config"](**text_config) - elif text_config is None: - self.text_config = self.sub_configs["text_config"]() - - self.image_token_id = image_token_id - self.video_token_id = video_token_id - self.vision_start_token_id = vision_start_token_id - self.vision_end_token_id = vision_end_token_id - super().__init__(**kwargs) - # Set after super().__init__() to avoid v4 PreTrainedConfig overwrite - self.tie_word_embeddings = tie_word_embeddings - - -__all__ = ["Qwen3_5Config", "Qwen3_5TextConfig"] diff --git a/vllm/transformers_utils/configs/qwen3_5_moe.py b/vllm/transformers_utils/configs/qwen3_5_moe.py deleted file mode 100644 index 34e66c825ac5..000000000000 --- a/vllm/transformers_utils/configs/qwen3_5_moe.py +++ /dev/null @@ -1,205 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. -# All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Qwen3.5-MoE model configuration.""" - -from transformers.configuration_utils import PreTrainedConfig - - -class Qwen3_5MoeTextConfig(PreTrainedConfig): - model_type = "qwen3_5_moe_text" - keys_to_ignore_at_inference = ["past_key_values"] - - base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.shared_expert.gate_proj": "colwise", - "layers.*.mlp.shared_expert.up_proj": "colwise", - "layers.*.mlp.shared_expert.down_proj": "rowwise", - } - base_model_pp_plan = { - "embed_tokens": (["input_ids"], ["inputs_embeds"]), - "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), - "norm": (["hidden_states"], ["hidden_states"]), - } - base_config_key = "text_config" - - def __init__( - self, - vocab_size=248320, - hidden_size=2048, - num_hidden_layers=40, - num_attention_heads=16, - num_key_value_heads=2, - hidden_act="silu", - max_position_embeddings=32768, - initializer_range=0.02, - rms_norm_eps=1e-6, - use_cache=True, - tie_word_embeddings=False, - rope_parameters=None, - attention_bias=False, - attention_dropout=0.0, - head_dim=256, - linear_conv_kernel_dim=4, - linear_key_head_dim=128, - linear_value_head_dim=128, - linear_num_key_heads=16, - linear_num_value_heads=32, - moe_intermediate_size=512, - shared_expert_intermediate_size=512, - num_experts_per_tok=8, - num_experts=256, - output_router_logits=False, - router_aux_loss_coef=0.001, - layer_types=None, - pad_token_id=None, - bos_token_id=None, - eos_token_id=None, - **kwargs, - ): - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.use_cache = use_cache - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.head_dim = head_dim - self.rope_parameters = rope_parameters - kwargs.setdefault("partial_rotary_factor", 0.25) - - self.layer_types = layer_types - if self.layer_types is None: - interval_pattern = kwargs.get("full_attention_interval", 4) - self.layer_types = [ - "linear_attention" - if bool((i + 1) % interval_pattern) - else "full_attention" - for i in range(self.num_hidden_layers) - ] - kwargs["ignore_keys_at_rope_validation"] = { - "mrope_section", - "mrope_interleaved", - } - self.validate_layer_type() - - # linear attention part - self.linear_conv_kernel_dim = linear_conv_kernel_dim - self.linear_key_head_dim = linear_key_head_dim - self.linear_value_head_dim = linear_value_head_dim - self.linear_num_key_heads = linear_num_key_heads - self.linear_num_value_heads = linear_num_value_heads - self.moe_intermediate_size = moe_intermediate_size - self.shared_expert_intermediate_size = shared_expert_intermediate_size - self.num_experts_per_tok = num_experts_per_tok - self.num_experts = num_experts - self.output_router_logits = output_router_logits - self.router_aux_loss_coef = router_aux_loss_coef - super().__init__(**kwargs) - # Set these AFTER super().__init__() because transformers v4's - # PreTrainedConfig.__init__ has these as explicit params with different - # defaults (e.g. tie_word_embeddings=True) that would overwrite our values. - self.pad_token_id = pad_token_id - self.bos_token_id = bos_token_id - self.eos_token_id = eos_token_id - self.tie_word_embeddings = tie_word_embeddings - - -class Qwen3_5MoeVisionConfig(PreTrainedConfig): - model_type = "qwen3_5_moe" - base_config_key = "vision_config" - - def __init__( - self, - depth=27, - hidden_size=1152, - hidden_act="gelu_pytorch_tanh", - intermediate_size=4304, - num_heads=16, - in_channels=3, - patch_size=16, - spatial_merge_size=2, - temporal_patch_size=2, - out_hidden_size=3584, - num_position_embeddings=2304, - initializer_range=0.02, - **kwargs, - ): - super().__init__(**kwargs) - - self.depth = depth - self.hidden_size = hidden_size - self.hidden_act = hidden_act - self.intermediate_size = intermediate_size - self.num_heads = num_heads - self.in_channels = in_channels - self.patch_size = patch_size - self.spatial_merge_size = spatial_merge_size - self.temporal_patch_size = temporal_patch_size - self.out_hidden_size = out_hidden_size - self.num_position_embeddings = num_position_embeddings - self.initializer_range = initializer_range - - -class Qwen3_5MoeConfig(PreTrainedConfig): - model_type = "qwen3_5_moe" - sub_configs = { - "vision_config": Qwen3_5MoeVisionConfig, - "text_config": Qwen3_5MoeTextConfig, - } - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - text_config=None, - vision_config=None, - image_token_id=248056, - video_token_id=248057, - vision_start_token_id=248053, - vision_end_token_id=248054, - tie_word_embeddings=False, - **kwargs, - ): - if isinstance(vision_config, dict): - self.vision_config = self.sub_configs["vision_config"](**vision_config) - elif vision_config is None: - self.vision_config = self.sub_configs["vision_config"]() - - if isinstance(text_config, dict): - self.text_config = self.sub_configs["text_config"](**text_config) - elif text_config is None: - self.text_config = self.sub_configs["text_config"]() - - self.image_token_id = image_token_id - self.video_token_id = video_token_id - self.vision_start_token_id = vision_start_token_id - self.vision_end_token_id = vision_end_token_id - super().__init__(**kwargs) - # Set after super().__init__() to avoid v4 PreTrainedConfig overwrite - self.tie_word_embeddings = tie_word_embeddings - - -__all__ = ["Qwen3_5MoeConfig", "Qwen3_5MoeTextConfig"] diff --git a/vllm/transformers_utils/configs/qwen3_next.py b/vllm/transformers_utils/configs/qwen3_next.py deleted file mode 100644 index 8b083ff9633d..000000000000 --- a/vllm/transformers_utils/configs/qwen3_next.py +++ /dev/null @@ -1,277 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team. -# All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Qwen3-Next model configuration.""" - -from transformers.configuration_utils import PreTrainedConfig -from transformers.utils import logging - -logger = logging.get_logger(__name__) - - -class Qwen3NextConfig(PreTrainedConfig): - r"""This is the configuration class to store the configuration of a [`Qwen3NextModel`]. It is used to instantiate a - Qwen3-Next model according to the specified arguments, defining the model architecture. - Instantiating a configuration with the defaults will yield a similar configuration to that of - Qwen3-Next-80B-A3B-Instruct [Qwen/Qwen3-Next-80B-A3B-Instruct](https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct). - - Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the - documentation from [`PreTrainedConfig`] for more information. - - - Args: - vocab_size (`int`, *optional*, defaults to 151936): - Vocabulary size of the model. Defines the number of different tokens that can be represented by the - `inputs_ids`. - hidden_size (`int`, *optional*, defaults to 2048): - Dimension of the hidden representations. - intermediate_size (`int`, *optional*, defaults to 5632): - Dimension of the MLP representations. - num_hidden_layers (`int`, *optional*, defaults to 48): - Number of hidden layers in the Transformer encoder. - num_attention_heads (`int`, *optional*, defaults to 16): - Number of attention heads for each attention layer in the Transformer encoder. - num_key_value_heads (`int`, *optional*, defaults to 2): - This is the number of key_value heads that should be used to implement Grouped Query Attention. If - `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if - `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When - converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed - by meanpooling all the original heads within that group. For more details checkout [this - paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`. - hidden_act (`str`, *optional*, defaults to `"silu"`): - The non-linear activation function in the decoder. - max_position_embeddings (`int`, *optional*, defaults to 32768): - The maximum sequence length that this model might ever be used with. - initializer_range (`float`, *optional*, defaults to 0.02): - The standard deviation of the truncated_normal_initializer for initializing all weight matrices. - rms_norm_eps (`float`, *optional*, defaults to 1e-06): - The epsilon used by the rms normalization layers. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values attentions (not used by all models). Only - relevant if `config.is_decoder=True`. - tie_word_embeddings (`bool`, *optional*, defaults to `False`): - Whether the model's input and output word embeddings should be tied. - rope_parameters (`dict`, *optional*): - Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type - and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value - accordingly. - Expected contents: - `rope_theta` (`float`): The base period of the RoPE embeddings. - `rope_type` (`str`): - The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope', - 'llama3'], with 'default' being the original RoPE implementation. - `factor` (`float`, *optional*): - Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In - most scaling types, a `factor` of x will enable the model to handle sequences of length x * - original maximum pre-trained length. - `original_max_position_embeddings` (`int`, *optional*): - Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during - pretraining. - `attention_factor` (`float`, *optional*): - Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention - computation. If unspecified, it defaults to value recommended by the implementation, using the - `factor` field to infer the suggested value. - `beta_fast` (`float`, *optional*): - Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear - ramp function. If unspecified, it defaults to 32. - `beta_slow` (`float`, *optional*): - Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear - ramp function. If unspecified, it defaults to 1. - `short_factor` (`List[float]`, *optional*): - Only used with 'longrope'. The scaling factor to be applied to short contexts (< - `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden - size divided by the number of attention heads divided by 2 - `long_factor` (`List[float]`, *optional*): - Only used with 'longrope'. The scaling factor to be applied to long contexts (< - `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden - size divided by the number of attention heads divided by 2 - `low_freq_factor` (`float`, *optional*): - Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE - `high_freq_factor` (`float`, *optional*): - Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE - `partial_rotary_factor` (`float`, *optional*, defaults to 0.25): - Percentage of the query and keys which will have rotary embedding. - attention_bias (`bool`, *optional*, defaults to `False`): - Whether to use a bias in the query, key, value and output projection layers during self-attention. - attention_dropout (`float`, *optional*, defaults to 0.0): - The dropout ratio for the attention probabilities. - head_dim (`int`, *optional*, defaults to 256): - Projection weights dimension in multi-head attention. - linear_conv_kernel_dim (`int`, *optional*, defaults to 4): - Kernel size of the convolution used in linear attention layers. - linear_key_head_dim (`int`, *optional*, defaults to 128): - Dimension of each key head in linear attention. - linear_value_head_dim (`int`, *optional*, defaults to 128): - Dimension of each value head in linear attention. - linear_num_key_heads (`int`, *optional*, defaults to 16): - Number of key heads used in linear attention layers. - linear_num_value_heads (`int`, *optional*, defaults to 32): - Number of value heads used in linear attention layers. - decoder_sparse_step (`int`, *optional*, defaults to 1): - The frequency of the MoE layer. - moe_intermediate_size (`int`, *optional*, defaults to 512): - Intermediate size of the routed expert. - shared_expert_intermediate_size (`int`, *optional*, defaults to 512): - Intermediate size of the shared expert. - num_experts_per_tok (`int`, *optional*, defaults to 10): - Number of selected experts. - num_experts (`int`, *optional*, defaults to 512): - Number of routed experts. - norm_topk_prob (`bool`, *optional*, defaults to `True`): - Whether to normalize the topk probabilities. - output_router_logits (`bool`, *optional*, defaults to `False`): - Whether or not the router logits should be returned by the model. Enabling this will also - allow the model to output the auxiliary loss, including load balancing loss and router z-loss. - router_aux_loss_coef (`float`, *optional*, defaults to 0.001): - The aux loss factor for the total loss. - mlp_only_layers (`list[int]`, *optional*, defaults to `[]`): - Indicate which layers use Qwen3NextMLP rather than Qwen3NextSparseMoeBlock - The list contains layer index, from 0 to num_layers-1 if we have num_layers layers - If `mlp_only_layers` is empty, `decoder_sparse_step` is used to determine the sparsity. - layer_types (`list[str]`, *optional*): - Types of each layer (attention or linear). - - ```python - >>> from transformers import Qwen3NextModel, Qwen3NextConfig - - >>> # Initializing a Qwen3Next style configuration - >>> configuration = Qwen3NextConfig() - - >>> # Initializing a model from the Qwen3-Next-80B-A3B style configuration - >>> model = Qwen3NextModel(configuration) - - >>> # Accessing the model configuration - >>> configuration = model.config - ``` - - """ # noqa: E501 - - model_type = "qwen3_next" - keys_to_ignore_at_inference = ["past_key_values"] - - base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.*.gate_proj": "colwise", - "layers.*.mlp.experts.*.up_proj": "colwise", - "layers.*.mlp.experts.*.down_proj": "rowwise", - "layers.*.mlp.shared_experts.gate_proj": "colwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", - } - base_model_pp_plan = { - "embed_tokens": (["input_ids"], ["inputs_embeds"]), - "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), - "norm": (["hidden_states"], ["hidden_states"]), - } - - def __init__( - self, - vocab_size=151936, - hidden_size=2048, - intermediate_size=5632, - num_hidden_layers=48, - num_attention_heads=16, - num_key_value_heads=2, - hidden_act="silu", - max_position_embeddings=32768, - initializer_range=0.02, - rms_norm_eps=1e-6, - use_cache=True, - tie_word_embeddings=False, - rope_parameters=None, - attention_bias=False, - attention_dropout=0.0, - head_dim=256, - linear_conv_kernel_dim=4, - linear_key_head_dim=128, - linear_value_head_dim=128, - linear_num_key_heads=16, - linear_num_value_heads=32, - decoder_sparse_step=1, - moe_intermediate_size=512, - shared_expert_intermediate_size=512, - num_experts_per_tok=10, - num_experts=512, - norm_topk_prob=True, - output_router_logits=False, - router_aux_loss_coef=0.001, - mlp_only_layers=None, - layer_types=None, - **kwargs, - ): - if mlp_only_layers is None: - mlp_only_layers = [] - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.use_cache = use_cache - # Try to set `rope_scaling` if available, otherwise use `rope_parameters` - rope_scaling = kwargs.pop("rope_scaling", None) - rope_parameters = rope_scaling or rope_parameters or {"rope_type": "default"} - rope_theta = kwargs.pop("rope_theta", 10000.0) - if "rope_theta" not in rope_parameters: - rope_parameters["rope_theta"] = rope_theta - partial_rotary_factor = kwargs.pop("partial_rotary_factor", 0.25) - if "partial_rotary_factor" not in rope_parameters: - rope_parameters["partial_rotary_factor"] = partial_rotary_factor - self.rope_parameters = rope_parameters - self.partial_rotary_factor = partial_rotary_factor - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.head_dim = head_dim - - self.layer_types = layer_types - if self.layer_types is None: - self.layer_types = [ - "linear_attention" if bool((i + 1) % 4) else "full_attention" - for i in range(self.num_hidden_layers) - ] - self.validate_layer_type() - - # linear attention part - self.linear_conv_kernel_dim = linear_conv_kernel_dim - self.linear_key_head_dim = linear_key_head_dim - self.linear_value_head_dim = linear_value_head_dim - self.linear_num_key_heads = linear_num_key_heads - self.linear_num_value_heads = linear_num_value_heads - - # MoE arguments - self.decoder_sparse_step = decoder_sparse_step - self.moe_intermediate_size = moe_intermediate_size - self.shared_expert_intermediate_size = shared_expert_intermediate_size - self.num_experts_per_tok = num_experts_per_tok - self.num_experts = num_experts - self.norm_topk_prob = norm_topk_prob - self.output_router_logits = output_router_logits - self.router_aux_loss_coef = router_aux_loss_coef - self.mlp_only_layers = mlp_only_layers - super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) - - -__all__ = ["Qwen3NextConfig"] diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 9dee4af924ce..d79c1275b6e7 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -17,7 +17,6 @@ ConfigFormat, get_safetensors_params_metadata, ) -from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config from vllm.utils.torch_utils import common_broadcastable_dtype logger = init_logger(__name__) @@ -41,7 +40,7 @@ def get_per_layer_hf_configs( This is the only place that decides whether a checkpoint is heterogeneous and how many layers it has, so a convertor can declare variation that - Transformers does not express (see `Gemma4ModelArchConfigConvertor`). + Transformers does not express. Each pair is converted independently and the results are diffed, so a convertor that overrides this must not also collapse a varying field in @@ -745,27 +744,6 @@ def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool: == "vision" ) - def get_per_layer_hf_configs( - self, - ) -> list[tuple[PreTrainedConfig, PreTrainedConfig]] | None: - # Gemma4 uses a larger head dimension, and sometimes more KV heads, on its - # full attention layers than on its sliding ones. Transformers >= 5.15.0 - # says so in the config; before that the values are flat attributes picked - # apart by `layer_types`, so build the per-layer configs here instead. Both - # then reach the base convertor as ordinary homogeneous configs. - if getattr(self.hf_text_config, "is_heterogeneous", False): - return super().get_per_layer_hf_configs() - text_config = self.hf_text_config - if not getattr(text_config, "layer_types", None): - return None - # `num_hidden_layers` is what the stack is built from, and it is what - # Transformers sizes `per_layer_config` by. `layer_types` can be longer, - # e.g. when a test truncates the layer count to build a small model. - return [ - (self.hf_config, gemma4_layer_config(text_config, layer_idx)) - for layer_idx in range(self.get_num_hidden_layers()) - ] - class MossAudioModelArchConfigConvertor(ModelArchConfigConvertorBase): def _language_config(self) -> PreTrainedConfig: diff --git a/vllm/transformers_utils/processors/__init__.py b/vllm/transformers_utils/processors/__init__.py index b913545377d0..f2af4901e2e9 100644 --- a/vllm/transformers_utils/processors/__init__.py +++ b/vllm/transformers_utils/processors/__init__.py @@ -17,10 +17,8 @@ "FireRedASR2Processor", "FunASRProcessor", "GLM4VProcessor", - "Granite4VisionProcessor", "H2OVLProcessor", "Moondream3Processor", - "InternVLProcessor", "IsaacProcessor", "KimiAudioProcessor", "KimiK25Processor", @@ -55,9 +53,7 @@ "FireRedASR2Processor": "vllm.transformers_utils.processors.fireredasr2", "FunASRProcessor": "vllm.transformers_utils.processors.funasr", "GLM4VProcessor": "vllm.transformers_utils.processors.glm4v", - "Granite4VisionProcessor": "vllm.transformers_utils.processors.granite4_vision", "H2OVLProcessor": "vllm.transformers_utils.processors.h2ovl", - "InternVLProcessor": "vllm.transformers_utils.processors.internvl", "IsaacProcessor": "vllm.transformers_utils.processors.isaac", "KimiAudioProcessor": "vllm.transformers_utils.processors.kimi_audio", "KimiK25Processor": "vllm.transformers_utils.processors.kimi_k25", diff --git a/vllm/transformers_utils/processors/granite4_vision.py b/vllm/transformers_utils/processors/granite4_vision.py deleted file mode 100644 index 7972015adf9d..000000000000 --- a/vllm/transformers_utils/processors/granite4_vision.py +++ /dev/null @@ -1,78 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from fractions import Fraction - -from transformers import LlavaNextProcessor -from transformers.image_processing_utils import select_best_resolution - - -class Granite4VisionProcessor(LlavaNextProcessor): - """Processor for Granite 4 Vision. - - Extends LlavaNextProcessor to account for the Window Q-Former - downsampling when computing the number of image features. - - This processor is needed because the granite4_vision processor type - is not yet in the transformers version pinned by vLLM. - """ - - model_type = "granite4_vision" - - def __init__( - self, - image_processor=None, - tokenizer=None, - patch_size=None, - vision_feature_select_strategy=None, - chat_template=None, - image_token="", - num_additional_image_tokens=0, - downsample_rate=None, - **kwargs, - ): - super().__init__( - image_processor=image_processor, - tokenizer=tokenizer, - patch_size=patch_size, - vision_feature_select_strategy=vision_feature_select_strategy, - chat_template=chat_template, - image_token=image_token, - num_additional_image_tokens=num_additional_image_tokens, - ) - self.downsample_rate = downsample_rate - - def _get_number_of_features( - self, - orig_height: int, - orig_width: int, - height: int, - width: int, - ) -> int: - image_grid_pinpoints = self.image_processor.image_grid_pinpoints - - height_best_resolution, width_best_resolution = select_best_resolution( - [orig_height, orig_width], image_grid_pinpoints - ) - scale_height = height_best_resolution // height - scale_width = width_best_resolution // width - - patches_height = height // self.patch_size - patches_width = width // self.patch_size - if self.downsample_rate is not None: - ds_rate = Fraction(self.downsample_rate) - patches_height = int(patches_height * ds_rate) - patches_width = int(patches_width * ds_rate) - - unpadded_features, newline_features = self._get_unpadded_features( - orig_height, - orig_width, - patches_height, - patches_width, - scale_height, - scale_width, - ) - base_features = ( - patches_height * patches_width + self.num_additional_image_tokens - ) - return unpadded_features + newline_features + base_features From 2186b89a9a02fac11fa0e274f3201b9812dbf65c Mon Sep 17 00:00:00 2001 From: Damien Laine Date: Fri, 9 Oct 2026 11:40:29 +0200 Subject: [PATCH 166/624] [Bugfix][Model] Voxtral realtime: fix boot OOM and engine crash at max_model_len (#45022) Signed-off-by: Damien Laine Co-authored-by: Cyrus Leung --- .../generation/test_voxtral_realtime.py | 84 ++++++++++++++++++- .../speech_to_text/realtime/connection.py | 29 ++++++- .../model_executor/models/voxtral_realtime.py | 43 +++++++++- 3 files changed, 149 insertions(+), 7 deletions(-) diff --git a/tests/models/multimodal/generation/test_voxtral_realtime.py b/tests/models/multimodal/generation/test_voxtral_realtime.py index 09432b4494cb..e05930b0f36b 100644 --- a/tests/models/multimodal/generation/test_voxtral_realtime.py +++ b/tests/models/multimodal/generation/test_voxtral_realtime.py @@ -106,7 +106,8 @@ def tokenizer() -> MistralTokenizer: @pytest_asyncio.fixture -async def async_engine(): +async def async_engine(request): + engine_config = {**ENGINE_CONFIG, **getattr(request, "param", {})} gpu_memory_utilization = ENGINE_CONFIG.get("gpu_memory_utilization", 0.9) from vllm.platforms import current_platform @@ -115,7 +116,7 @@ async def async_engine(): wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) - engine_args = AsyncEngineArgs(**ENGINE_CONFIG) + engine_args = AsyncEngineArgs(**engine_config) llm = AsyncLLM.from_engine_args(engine_args) try: yield llm @@ -275,3 +276,82 @@ async def input_stream(audio_buffer): f" got: {got!r}\n" f" expected: {expected!r}" ) + + +@pytest.mark.asyncio +async def test_voxtral_realtime_buffer_stops_at_max_model_len(tokenizer): + from vllm.model_executor.models.voxtral_realtime import VoxtralRealtimeBuffer + + audio_config = tokenizer.instruct_tokenizer.audio_encoder.audio_config + audio = Audio.from_file(AudioAsset("winning_call").get_local_path(), strict=False) + req = TranscriptionRequest( + streaming=StreamingMode.OFFLINE, + audio=audio.to_base64(audio.format), + language=None, + ) + audio_enc = tokenizer.encode_transcription(req) + + max_model_len = 64 + buffer = VoxtralRealtimeBuffer( + audio_config, audio_enc.tokens, max_model_len=max_model_len + ) + await buffer.append_audio(audio_enc.audios[0].audio_array) + await buffer.append_audio(None) + # Unlimited feedback tokens so only the length cap can stop the stream. + await buffer.append_tokens([0] * max_model_len) + + n_tokens = 0 + async for prompt in buffer.get_input_stream(): + n_tokens += len(prompt["prompt_token_ids"]) + + assert len(audio_enc.tokens) < n_tokens < max_model_len + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_engine", [{"max_model_len": 128}], indirect=True) +async def test_voxtral_realtime_generator_max_model_len( + audio_assets, tokenizer, async_engine +): + # Lazy import to avoid CUDA-reinitialization error + from vllm.model_executor.models.voxtral_realtime import VoxtralRealtimeBuffer + + max_model_len = async_engine.model_config.max_model_len + sampling_params = SamplingParams(temperature=0.0, max_tokens=1) + audio_config = tokenizer.instruct_tokenizer.audio_encoder.audio_config + + async def input_stream(audio_buffer): + async for prompt in audio_buffer.get_input_stream(): + parsed_prompt = parse_model_prompt(async_engine.model_config, prompt) + (engine_input,) = await async_engine.renderer.render_cmpl_async( + [parsed_prompt] + ) + yield StreamingInput(prompt=engine_input) + + # winning_call is ~300 tokens long, well past max_model_len. + audio = Audio.from_file(audio_assets[1].get_local_path(), strict=False) + req = TranscriptionRequest( + streaming=StreamingMode.OFFLINE, + audio=audio.to_base64(audio.format), + language=None, + ) + audio_enc = tokenizer.encode_transcription(req) + buffer = VoxtralRealtimeBuffer( + audio_config, audio_enc.tokens, max_model_len=max_model_len + ) + await buffer.append_audio(audio_enc.audios[0].audio_array) + await buffer.append_audio(None) + + output_tokens = [] + async for resp in async_engine.generate( + prompt=input_stream(buffer), + sampling_params=sampling_params, + request_id="session-max-len", + ): + tokens = resp.outputs[0].token_ids[-1:] + output_tokens.extend(tokens) + await buffer.append_tokens(tokens) + + assert resp.finished + assert 0 < len(output_tokens) < max_model_len + # The engine is still healthy after the session hit the cap. + assert not async_engine.errored diff --git a/vllm/entrypoints/speech_to_text/realtime/connection.py b/vllm/entrypoints/speech_to_text/realtime/connection.py index ed8ec7f862c5..4a5e6cef9a57 100644 --- a/vllm/entrypoints/speech_to_text/realtime/connection.py +++ b/vllm/entrypoints/speech_to_text/realtime/connection.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio +import contextlib import json from collections.abc import AsyncGenerator from http import HTTPStatus @@ -48,6 +49,7 @@ def __init__(self, websocket: WebSocket, serving: OpenAIServingRealtime): self.serving = serving self.audio_queue: asyncio.Queue[np.ndarray | None] = asyncio.Queue() self.generation_task: asyncio.Task | None = None + self.request_id: str | None = None self._is_connected = False self._is_model_validated = False @@ -204,6 +206,7 @@ async def _run_generation( 6. Cleans up the audio queue """ request_id = f"rt-{self.connection_id}-{uuid4()}" + self.request_id = request_id full_text = "" prompt_token_ids_len: int = 0 @@ -263,7 +266,9 @@ async def _run_generation( except Exception as e: logger.exception("Error in generation: %s", e) - await self.send_error(sanitize_message(str(e)), "processing_error") + # The socket may already be closed; nothing more we can do. + with contextlib.suppress(Exception): + await self.send_error(sanitize_message(str(e)), "processing_error") async def send( self, event: SessionCreated | TranscriptionDelta | TranscriptionDone @@ -282,8 +287,28 @@ async def cleanup(self): # Signal audio stream to stop self.audio_queue.put_nowait(None) - # Cancel generation task if running + # Abort the engine request. Cancelling the asyncio task alone does NOT + # abort the underlying engine request: a request parked in the + # VAD-waiting state keeps occupying a slot and generating into a now + # dead websocket (an orphaned request that never frees and starves + # other streams). Aborting also unblocks the generation task, which is + # otherwise stuck awaiting the next engine output. + if self.request_id is not None: + try: + await self.serving.engine_client.abort(self.request_id) + except Exception: + logger.exception("Failed to abort request %s", self.request_id) + self.request_id = None + + # Cancel the generation task if still running, and await it so the + # cancellation (and any generator cleanup) actually completes. if self.generation_task and not self.generation_task.done(): self.generation_task.cancel() + try: + await self.generation_task + except asyncio.CancelledError: + pass + except Exception: + logger.exception("Error awaiting cancelled generation task") logger.debug("Connection cleanup complete: %s", self.connection_id) diff --git a/vllm/model_executor/models/voxtral_realtime.py b/vllm/model_executor/models/voxtral_realtime.py index 8531d7d18f39..72318976f1d4 100644 --- a/vllm/model_executor/models/voxtral_realtime.py +++ b/vllm/model_executor/models/voxtral_realtime.py @@ -120,9 +120,35 @@ def _expand_tensor(input_tensor: torch.Tensor, scaling: int) -> torch.Tensor: return (base.unsqueeze(1) + offsets).view(-1) +class VoxtralRealtimeProcessingInfo(VoxtralProcessingInfo): + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + # Each streaming update carries a single audio frame. + return {"audio": 1} + + def get_max_audio_array_len(self) -> int: + # Profile with the largest streaming chunk, the first one: it covers + # the prompt prefix (BOS + left padding + delay), later chunks are + # ~1 token. get_max_audio_tokens keeps max_model_len so the encoder + # cache stays as large as before: each paused session keeps its last + # chunk cached until it resumes, and a cache sized to one chunk + # deadlocks a few concurrent sessions. + tokenizer = self.get_tokenizer() + audio_encoder = tokenizer.instruct.audio_encoder + prompt_tokens = ( + tokenizer.instruct.start() + audio_encoder.encode_streaming_tokens() + ) + return len(prompt_tokens) * audio_encoder.audio_config.raw_audio_length_per_tok + + class VoxtralRealtimeBuffer: - def __init__(self, config: AudioConfig, prompt_tokens: list[int]) -> None: + def __init__( + self, + config: AudioConfig, + prompt_tokens: list[int], + max_model_len: int | None = None, + ) -> None: self._config = config + self._max_model_len = max_model_len _look_ahead_in_ms = self._config.streaming_look_ahead_ms _look_back_in_ms = self._config.streaming_look_back_ms @@ -165,7 +191,16 @@ async def append_tokens(self, tokens: Iterable[int]) -> None: await self._token_queue.put(token) async def get_input_stream(self) -> AsyncGenerator[TokensPrompt]: + n_tokens = 0 for frame_size, num_tokens in self._generate_frame_size_and_num_tokens(): + n_tokens += num_tokens + if self._max_model_len is not None and n_tokens >= self._max_model_len: + logger.warning( + "Stopping realtime audio stream: reached max_model_len (%d).", + self._max_model_len, + ) + return + next_tokens = [await self._token_queue.get() for _ in range(num_tokens)] audio_arrays: list[np.ndarray] = ( @@ -198,7 +233,7 @@ async def get_input_stream(self) -> AsyncGenerator[TokensPrompt]: @MULTIMODAL_REGISTRY.register_processor( VoxtralRealtimeMultiModalProcessor, - info=VoxtralProcessingInfo, + info=VoxtralRealtimeProcessingInfo, dummy_inputs=VoxtralDummyInputsBuilder, ) @support_torch_compile @@ -248,7 +283,9 @@ async def buffer_realtime_audio( # Get left/right padding audio left_pad, right_pad = audio_encoder.get_padding_audio() - buffer = VoxtralRealtimeBuffer(config, prompt_tokens) + buffer = VoxtralRealtimeBuffer( + config, prompt_tokens, max_model_len=model_config.max_model_len + ) # Feed audio with padding into buffer in background async def feed_audio(): From 59ffd73e8597d860730f4a8c4f778f27c75a0dfa Mon Sep 17 00:00:00 2001 From: Warren Deng <35615489+warrendeng@users.noreply.github.com> Date: Fri, 9 Oct 2026 02:46:41 -0700 Subject: [PATCH 167/624] [Bugfix][Kernel] Restore FP fusion for indexer Q quantization (#60716) Signed-off-by: Warren Deng --- vllm/models/deepseek_v4/common/ops/fused_indexer_q.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py index 669186edbb1b..54b1b9be7e89 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py +++ b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py @@ -305,6 +305,7 @@ def __call__( FP8_MAX=fp8_max, USE_FNUZ=use_fnuz, USE_EXPLICIT_FMA=current_platform.is_rocm(), + enable_fp_fusion=True, num_warps=1, ) @@ -540,6 +541,7 @@ def __call__( MXFP4_BLOCK=MXFP4_BLOCK_SIZE, index_weights_stride=index_weights.stride(0), index_weights_out_stride=index_weights_out.stride(0), + enable_fp_fusion=True, num_warps=1, ) From b8bd3357548ac07b1877053b6c5ba415763bad36 Mon Sep 17 00:00:00 2001 From: Thang Nguyen <69278249+Thangnguyenvn98@users.noreply.github.com> Date: Fri, 9 Oct 2026 06:26:27 -0400 Subject: [PATCH 168/624] [CI] Enroll five more steps in automatic sharding (#60662) Signed-off-by: Thang Nguyen Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/engine.yaml | 1 + .buildkite/test_areas/entrypoints.yaml | 3 +++ .buildkite/test_areas/pytorch.yaml | 1 + 3 files changed, 5 insertions(+) diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 53499d4a986a..29aec98b07ee 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -108,6 +108,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) E2E Core" device: h200_18gb key: e2e-core-1-gpu + automatic_shard: true timeout_in_minutes: 65 source_file_dependencies: - vllm/v1/ diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 34a190e810b9..d31c66a8cbd8 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -26,6 +26,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Entrypoints Unit" device: h200_35gb key: entrypoints-unit-tests + automatic_shard: true timeout_in_minutes: 40 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -168,6 +169,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (Responses API)" device: h200_35gb key: entrypoints-integration-responses-api + automatic_shard: true timeout_in_minutes: 50 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -188,6 +190,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (Speech to Text)" device: h200_35gb key: entrypoints-integration-speech_to_text + automatic_shard: true timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" source_file_dependencies: diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index 0ba608e72580..60061af41488 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -123,6 +123,7 @@ steps: - label: ":nvidia: (L4) PyTorch Compilation Passes" key: pytorch-compilation-passes-unit-tests + automatic_shard: true device: l4 num_devices: 1 timeout_in_minutes: 45 From efd141047d8e879536806a86e3f93c02e6164add Mon Sep 17 00:00:00 2001 From: Thang Nguyen <69278249+Thangnguyenvn98@users.noreply.github.com> Date: Fri, 9 Oct 2026 06:26:47 -0400 Subject: [PATCH 169/624] [CI] Enroll a third batch of five steps in automatic sharding (#60757) Signed-off-by: Thang Nguyen Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/cuda.yaml | 1 + .buildkite/test_areas/engine.yaml | 1 + .buildkite/test_areas/lm_eval.yaml | 1 + .buildkite/test_areas/misc.yaml | 1 + .buildkite/test_areas/quantization.yaml | 1 + 5 files changed, 5 insertions(+) diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml index 45b3dc768e6d..b3bfa9f09708 100644 --- a/.buildkite/test_areas/cuda.yaml +++ b/.buildkite/test_areas/cuda.yaml @@ -33,6 +33,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) CUDAGraph" device: h200_18gb key: cudagraph + automatic_shard: true timeout_in_minutes: 45 source_file_dependencies: - tests/v1/cudagraph diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 29aec98b07ee..26485e0a172d 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -38,6 +38,7 @@ steps: - label: ":nvidia: (L4) V1 Engine" key: engine-1-gpu + automatic_shard: true device: l4 num_devices: 1 timeout_in_minutes: 45 diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 49eea0b78f7c..06607951057b 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -5,6 +5,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) LM Eval Small Models" device: h200_35gb key: lm-eval-small-models + automatic_shard: true timeout_in_minutes: 45 source_file_dependencies: - csrc/ diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 52ff1d65b3ee..9763b2ce7de4 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: ":nvidia: (H200 MIG 18GB) V1 Sample" key: v1-sample + automatic_shard: true timeout_in_minutes: 45 device: h200_18gb source_file_dependencies: &v1-sample-logits-deps diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index 429b01f27320..55cef9fbd33a 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -80,6 +80,7 @@ steps: - label: ":nvidia: (B200) Quantized MoE" key: quantized-moe-test-b200 + automatic_shard: true timeout_in_minutes: 120 working_dir: "/vllm-workspace/" device: b200-k8s From 4e318ac56319476be1bc81f44acaf99516891afb Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Fri, 9 Oct 2026 18:38:01 +0800 Subject: [PATCH 170/624] [XPU] fix online fp8_per_channel quantization (#56027) Signed-off-by: Yan Ma --- .buildkite/intel_jobs/quantization.yaml | 4 +++- tests/quantization/test_fp8_per_channel.py | 4 +++- vllm/platforms/xpu.py | 1 + 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/.buildkite/intel_jobs/quantization.yaml b/.buildkite/intel_jobs/quantization.yaml index 8f1dfebe5001..af1b23b8a004 100644 --- a/.buildkite/intel_jobs/quantization.yaml +++ b/.buildkite/intel_jobs/quantization.yaml @@ -25,5 +25,7 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_per_token_kv_cache.py --deselect="tests/quantization/test_per_token_kv_cache.py::test_triton_unified_attention_per_token_head_scale[int4-16-128-num_heads0-seq_lens1]" - && VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_auto_awq.py' + && VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_auto_awq.py + && VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_fp8_per_channel.py + && VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/models/quantization/test_fp8_per_channel.py --deselect="tests/models/quantization/test_fp8_per_channel.py::test_fp8_per_channel_logprobs[moe]"' diff --git a/tests/quantization/test_fp8_per_channel.py b/tests/quantization/test_fp8_per_channel.py index 2c3fddcb0fe4..d4f2fe264e6c 100644 --- a/tests/quantization/test_fp8_per_channel.py +++ b/tests/quantization/test_fp8_per_channel.py @@ -61,7 +61,9 @@ def test_scaled_fp8_quant_per_channel_shape() -> None: weight to `ops.scaled_fp8_quant` with `use_per_token_if_dynamic=True` yields one scale per output row -- a [N, 1] fp32 tensor. """ - x = (torch.randn(size=(96, 256), device="cuda") * 13).to(torch.bfloat16) + x = (torch.randn(size=(96, 256), device=current_platform.device_type) * 13).to( + torch.bfloat16 + ) y, s = ops.scaled_fp8_quant(x, scale=None, use_per_token_if_dynamic=True) assert y.shape == (96, 256) assert y.dtype == current_platform.fp8_dtype() diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 355a15094db0..992bfae0534c 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -123,6 +123,7 @@ class XPUPlatform(Platform): "mxfp8", "fp8_per_tensor", "fp8_per_block", + "fp8_per_channel", "online", "gpt_oss_mxfp4", "modelopt", From 5cadabfafde31dcccf4e1f773d00cc6e8a134cfb Mon Sep 17 00:00:00 2001 From: yuvalluria <131433771+yuvalluria@users.noreply.github.com> Date: Fri, 9 Oct 2026 13:41:29 +0300 Subject: [PATCH 171/624] [Docs] Add Qwen2.5-Coder-7B-Instruct to batch invariance tested models (#46368) Signed-off-by: Yuval Luria Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- docs/features/batch_invariance.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index b40ff4da8cc6..03d07896cce4 100644 --- a/docs/features/batch_invariance.md +++ b/docs/features/batch_invariance.md @@ -126,7 +126,7 @@ Batch invariance has been tested and verified on the following models: - **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B`, `Qwen/Qwen3-4B-AWQ`, `Qwen/Qwen3-8B-AWQ` - **Qwen3-VL (Vision-Language)**: `Qwen/Qwen3-VL-2B-Instruct`, `Qwen/Qwen3-VL-4B-Instruct` (single image and video inputs) - **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct`, `Qwen/Qwen3-30B-A3B-Thinking-2507-FP8` -- **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct` +- **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct`, `Qwen/Qwen2.5-Coder-7B-Instruct` - **Llama 3**: Llama3.1 and 3.2 series, `meta-llama/Llama-3.2-1B-Instruct`, `meta-llama/Llama-3.2-3B-Instruct` for example - **GPT-OSS**: `openai/gpt-oss-20b`, `openai/gpt-oss-120b` - **Mistral**: `mistralai/Mistral-7B-v0.3` From b027ac831253e83ff605a2ff5d4aa7c00619802c Mon Sep 17 00:00:00 2001 From: Chaojun Zhang Date: Fri, 9 Oct 2026 18:58:07 +0800 Subject: [PATCH 172/624] [XPU] Support register KV offload mmap region as pinned host memory on XPU (#51956) Signed-off-by: Chaojun Zhang --- .../cpu/test_shared_offload_region.py | 81 +++++++--------- tests/v1/kv_offload/test_factory.py | 3 + vllm/v1/kv_offload/cpu/gpu_worker.py | 45 ++------- vllm/v1/kv_offload/cpu/host_register.py | 97 +++++++++++++++++++ .../kv_offload/cpu/shared_offload_region.py | 21 ++-- vllm/v1/kv_offload/cpu/spec.py | 4 +- 6 files changed, 151 insertions(+), 100 deletions(-) create mode 100644 vllm/v1/kv_offload/cpu/host_register.py diff --git a/tests/v1/kv_offload/cpu/test_shared_offload_region.py b/tests/v1/kv_offload/cpu/test_shared_offload_region.py index f84e1bd7296c..b7259e6fd621 100644 --- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py +++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py @@ -772,17 +772,14 @@ def test_cleanup_idempotent(iid): def test_cleanup_unregisters_every_pinned_chunk(iid, monkeypatch): """cleanup() must release every chunk from a split host registration.""" r = _make_region(iid) - cudart = MagicMock() - cudart.cudaHostUnregister.return_value = MagicMock(value=0) - monkeypatch.setattr(region_module, "current_platform", MagicMock()) - region_module.current_platform.is_cuda_alike.return_value = True - monkeypatch.setattr(region_module.torch.cuda, "cudart", lambda: cudart) + unregister = MagicMock() + monkeypatch.setattr(region_module, "host_unregister", unregister) r.pinned_addresses = [0x100000, 0x200000, 0x300000] r.is_pinned = True r.cleanup() - assert cudart.cudaHostUnregister.call_args_list == [ + assert unregister.call_args_list == [ call(0x300000), call(0x200000), call(0x100000), @@ -791,40 +788,35 @@ def test_cleanup_unregisters_every_pinned_chunk(iid, monkeypatch): @pytest.fixture def host_register(monkeypatch): - """Pretend to run on CUDA with host registrations capped at seven pages. - - Registration goes to a mock CudaRTLibrary and the region's cleanup to a - separate mock torch runtime. Both stay installed through the cleanup that - ``_region`` runs on exit. + """Mock the cross-platform host_register/host_unregister helpers with + registrations capped at seven pages. The same unregister mock backs both + pin_mmap_region's rollback path and the region's cleanup() on exit from + ``_region``. """ - cudart = MagicMock() - cudart.cudaHostRegister.return_value = 0 - cudart.cudaHostUnregister.return_value = 0 - cudart.cudaGetLastError.return_value = 2 - torch_cudart = MagicMock() - torch_cudart.cudaHostUnregister.return_value = MagicMock(value=0) - monkeypatch.setattr(gpu_worker.current_platform, "is_cuda_alike", lambda: True) + register = MagicMock(return_value=True) + unregister = MagicMock() monkeypatch.setattr(gpu_worker, "MAX_HOST_REGISTER_CHUNK_BYTES", 7 * PAGE_SIZE) - monkeypatch.setattr(gpu_worker, "CudaRTLibrary", lambda: cudart) - monkeypatch.setattr(region_module.torch.cuda, "cudart", lambda: torch_cudart) - return cudart, torch_cudart + monkeypatch.setattr(gpu_worker, "host_register", register) + monkeypatch.setattr(gpu_worker, "host_unregister", unregister) + monkeypatch.setattr(region_module, "host_unregister", unregister) + return register, unregister def test_pin_mmap_region_registers_row_aligned_chunks(iid, host_register): """Chunks end on row boundaries: 3-page rows under a 7-page cap register as 6 + 6 + 3 pages, where a raw byte cap would give 7 + 7 + 1.""" - cudart, torch_cudart = host_register + register, unregister = host_register with _region(iid, num_chunks=5, cpu_page_size=3 * PAGE_SIZE) as region: gpu_worker.pin_mmap_region(region) base = region._base.data_ptr() - assert cudart.mock_calls == [ - call.cudaHostRegister(base, 6 * PAGE_SIZE), - call.cudaHostRegister(base + 6 * PAGE_SIZE, 6 * PAGE_SIZE), - call.cudaHostRegister(base + 12 * PAGE_SIZE, 3 * PAGE_SIZE), + assert register.call_args_list == [ + call(base, 6 * PAGE_SIZE), + call(base + 6 * PAGE_SIZE, 6 * PAGE_SIZE), + call(base + 12 * PAGE_SIZE, 3 * PAGE_SIZE), ] assert region.is_pinned - assert torch_cudart.cudaHostUnregister.call_args_list == [ + assert unregister.call_args_list == [ call(base + 12 * PAGE_SIZE), call(base + 6 * PAGE_SIZE), call(base), @@ -833,10 +825,10 @@ def test_pin_mmap_region_registers_row_aligned_chunks(iid, host_register): @pytest.mark.parametrize("fail_at", [0, 1, 2]) def test_pin_mmap_region_failure_leaves_region_pageable(iid, host_register, fail_at): - """A failed chunk drains its pending error and unregisters the chunks - before it on the same handle, so the region is pinned whole or not at all.""" - cudart, torch_cudart = host_register - cudart.cudaHostRegister.side_effect = [0] * fail_at + [2] + """A failed chunk unregisters the chunks before it on the same range, so + the region is pinned whole or not at all.""" + register, unregister = host_register + register.side_effect = [True] * fail_at + [False] with _region(iid, num_chunks=5, cpu_page_size=3 * PAGE_SIZE) as region: gpu_worker.pin_mmap_region(region) base = region._base.data_ptr() @@ -845,25 +837,22 @@ def test_pin_mmap_region_failure_leaves_region_pageable(iid, host_register, fail registered = [base + i * 6 * PAGE_SIZE for i in range(fail_at)] failed_size = min(6, 15 - 6 * fail_at) * PAGE_SIZE - assert cudart.mock_calls == [ - *(call.cudaHostRegister(address, 6 * PAGE_SIZE) for address in registered), - call.cudaHostRegister(base + fail_at * 6 * PAGE_SIZE, failed_size), - call.cudaGetLastError(), - *(call.cudaHostUnregister(address) for address in reversed(registered)), + assert register.call_args_list == [ + *(call(address, 6 * PAGE_SIZE) for address in registered), + call(base + fail_at * 6 * PAGE_SIZE, failed_size), + ] + assert unregister.call_args_list == [ + call(address) for address in reversed(registered) ] - torch_cudart.cudaHostUnregister.assert_not_called() -def test_pin_mmap_region_without_cuda_runtime_stays_pageable( - iid, host_register, monkeypatch +def test_pin_mmap_region_when_host_register_unsupported_stays_pageable( + iid, monkeypatch ): - """A CUDA runtime that cannot be loaded leaves the region pageable instead - of failing worker startup.""" - - def missing_runtime(): - raise AssertionError("libcudart is not loaded in the current process") - - monkeypatch.setattr(gpu_worker, "CudaRTLibrary", missing_runtime) + """A platform without host registration support (host_register() + returning False) leaves the region pageable instead of failing worker + startup.""" + monkeypatch.setattr(gpu_worker, "host_register", lambda ptr, num_bytes: False) with _region(iid) as region: gpu_worker.pin_mmap_region(region) assert not region.is_pinned diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index 01ef756cae53..84b9da36e06b 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -307,6 +307,7 @@ def test_cpu_spec_replicated_disabled_without_shared_region( # those private buffers empty and corrupting subsequent loads. import vllm.v1.kv_offload.cpu.spec as cpu_spec_module + monkeypatch.setattr(cpu_spec_module.current_platform, "is_xpu", lambda: False) monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: rocm) monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: rocm) worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT @@ -341,6 +342,7 @@ def test_cpu_spec_replicated_layout_truth_matrix( # actually allocates on the shared region. import vllm.v1.kv_offload.cpu.spec as cpu_spec_module + monkeypatch.setattr(cpu_spec_module.current_platform, "is_xpu", lambda: False) monkeypatch.setattr( cpu_spec_module.current_platform, "is_cuda_alike", lambda: cuda_alike ) @@ -417,6 +419,7 @@ def fake_worker_ctor(**kwargs): worker_calls.append(kwargs) return MagicMock() + monkeypatch.setattr(cpu_spec_module.current_platform, "is_xpu", lambda: False) monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: rocm) monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: rocm) monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor) diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py index a56febf702ed..f6053480a7ab 100644 --- a/vllm/v1/kv_offload/cpu/gpu_worker.py +++ b/vllm/v1/kv_offload/cpu/gpu_worker.py @@ -11,7 +11,6 @@ import torch from vllm import _custom_ops as ops -from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON, triton @@ -27,6 +26,7 @@ OffloadingWorker, TransferResult, ) +from vllm.v1.kv_offload.cpu.host_register import host_register, host_unregister from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion from vllm.v1.kv_offload.cpu.swap_blocks_triton import ( THRESHOLD_BYTES, @@ -200,26 +200,7 @@ def _canonical_block_sizes( def pin_mmap_region(region: SharedOffloadRegion) -> None: """Register row-aligned chunks, rolling back on failure.""" - if not current_platform.is_cuda_alike(): - logger.info( - "Skipping mmap host registration on %s; cudaHostRegister is only " - "available on CUDA/ROCm.", - current_platform.device_name, - ) - return - rank = region.rank - try: - cudart = CudaRTLibrary() - except (AssertionError, AttributeError, OSError): - logger.warning( - "Could not load the CUDA runtime for host registration on rank=%d; " - "the offload region stays pageable", - rank, - exc_info=True, - ) - return - base_ptr = region._base.data_ptr() total_size = region.total_size_bytes # Chunks end on block-row boundaries, which are page aligned, so neither the @@ -228,42 +209,30 @@ def pin_mmap_region(region: SharedOffloadRegion) -> None: rows_per_chunk = max(MAX_HOST_REGISTER_CHUNK_BYTES // region._row_stride, 1) chunk_size = rows_per_chunk * region._row_stride - # Register, drain and roll back through the same runtime handle, so a - # failed chunk leaves neither a pending error nor a partly pinned region. + # Register, drain and roll back through the same helper, so a failed + # chunk leaves neither a pending error nor a partly pinned region. addresses: list[int] = [] for offset in range(0, total_size, chunk_size): address = base_ptr + offset size = min(chunk_size, total_size - offset) - result = cudart.cudaHostRegister(address, size) - if result == 0: + if host_register(address, size): addresses.append(address) continue - cudart.cudaGetLastError() logger.warning( - "cudaHostRegister failed for rank=%d at %.2f of %.2f GB (code=%d); " + "host_register failed for rank=%d at %.2f of %.2f GB; " "the offload region stays pageable", rank, offset / 1e9, total_size / 1e9, - result, ) for registered in reversed(addresses): - unregister_result = cudart.cudaHostUnregister(registered) - if unregister_result != 0: - cudart.cudaGetLastError() - logger.warning( - "cudaHostUnregister failed for rank=%d at %#x (code=%d); " - "that chunk stays registered until the process exits", - rank, - registered, - unregister_result, - ) + host_unregister(registered) return region.pinned_addresses.extend(addresses) region.is_pinned = True logger.debug( - "cudaHostRegister rank=%d %.2f GB in %d chunk(s)", + "Host-registered mmap region rank=%d %.2f GB in %d chunk(s)", rank, total_size / 1e9, len(addresses), diff --git a/vllm/v1/kv_offload/cpu/host_register.py b/vllm/v1/kv_offload/cpu/host_register.py new file mode 100644 index 000000000000..37448857409f --- /dev/null +++ b/vllm/v1/kv_offload/cpu/host_register.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Registration of existing host memory for pinned DMA transfers. + +Host allocations that the offloading workers share across processes (the +`SharedOffloadRegion` mmap) cannot be allocated through the accelerator's +pinned allocator, so they must be registered after the fact. Pinning is a +pure bandwidth optimization: when registration is unavailable the transfers +still run correctly, just as pageable (slower) copies. +""" + +import torch + +from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + + +def _load_cudart() -> CudaRTLibrary | None: + try: + return CudaRTLibrary() + except (AssertionError, AttributeError, OSError): + logger.warning( + "Could not load the CUDA runtime for host registration; " + "the offload region stays pageable", + exc_info=True, + ) + return None + + +def host_register(ptr: int, num_bytes: int) -> bool: + """Register a host range as device-accessible pinned memory. + + Args: + ptr: Base address of the host range. + num_bytes: Size of the host range in bytes. + + Returns: + True if the range was registered and later needs `host_unregister`. + + """ + if not ptr or num_bytes <= 0: + return False + + if current_platform.is_cuda_alike(): + cudart = _load_cudart() + if cudart is None: + return False + result = cudart.cudaHostRegister(ptr, num_bytes) + if result != 0: + cudart.cudaGetLastError() + logger.warning( + "cudaHostRegister failed (code=%d) — transfers will still work " + "but may be slower (unpinned DMA)", + result, + ) + return False + return True + + if current_platform.is_xpu(): + if not torch.ops._C.xpu_host_register(ptr, num_bytes): + logger.warning( + "xpu_host_register failed — transfers will still work " + "but may be slower (unpinned DMA)" + ) + return False + return True + + logger.info_once( + "Host registration is not implemented on %s; KV offload transfers will " + "use unpinned host memory.", + current_platform.device_name, + ) + return False + + +def host_unregister(ptr: int) -> None: + """Release a host range previously registered by `host_register`.""" + if not ptr: + return + + if current_platform.is_cuda_alike(): + cudart = _load_cudart() + if cudart is None: + return + result = cudart.cudaHostUnregister(ptr) + if result != 0: + cudart.cudaGetLastError() + logger.warning("cudaHostUnregister failed (code=%d)", result) + return + + if current_platform.is_xpu(): + if not torch.ops._C.xpu_host_unregister(ptr): + logger.warning("xpu_host_unregister failed") + return diff --git a/vllm/v1/kv_offload/cpu/shared_offload_region.py b/vllm/v1/kv_offload/cpu/shared_offload_region.py index af183fbc4f06..7e0e97aa1219 100644 --- a/vllm/v1/kv_offload/cpu/shared_offload_region.py +++ b/vllm/v1/kv_offload/cpu/shared_offload_region.py @@ -14,7 +14,7 @@ check_shm_free_space, ) from vllm.logger import init_logger -from vllm.platforms import current_platform +from vllm.v1.kv_offload.cpu.host_register import host_unregister logger = init_logger(__name__) @@ -343,20 +343,11 @@ def create_kv_memoryview(self) -> memoryview: def cleanup(self) -> None: if self.is_pinned and self._base is not None: - if current_platform.is_cuda_alike(): - base_ptr = self._base.data_ptr() - addresses = self.pinned_addresses or [base_ptr] - for address in reversed(addresses): - result = torch.cuda.cudart().cudaHostUnregister(address) - if result.value != 0: - logger.warning( - "cudaHostUnregister failed for rank=%d, " - "address=%#x (code=%d)", - self.rank, - address, - result.value, - ) - self.pinned_addresses.clear() + base_ptr = self._base.data_ptr() + addresses = self.pinned_addresses or [base_ptr] + for address in reversed(addresses): + host_unregister(address) + self.pinned_addresses.clear() self.is_pinned = False # Release views before _base: each view holds a _base reference and a # direct StorageImpl reference. Freeing views first lets both refcounts diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index fddf56608590..bf289d37a65b 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -162,7 +162,9 @@ def get_manager(self) -> OffloadingManager: def _uses_shared_region(self) -> bool: """Whether the worker CPU buffer is the shared mmap region (vs a private per-rank tensor); replicated-layout dedup is gated on this being True.""" - return current_platform.is_cuda_alike() and not current_platform.is_rocm() + return ( + current_platform.is_cuda_alike() and not current_platform.is_rocm() + ) or current_platform.is_xpu() def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: mmap_region: SharedOffloadRegion | None = None From 28d76f3bad3df9270019b850a5c29f241b3d7bbf Mon Sep 17 00:00:00 2001 From: Shreyansh Jain Date: Fri, 9 Oct 2026 16:46:13 +0530 Subject: [PATCH 173/624] [Bugfix][Structured Output] Reject unsupported regex inside groups (#60359) Signed-off-by: Shreyansh Jain Signed-off-by: Artem Perevedentsev Co-authored-by: Artem Perevedentsev --- .../test_backend_outlines.py | 73 +++++++++++++++++++ vllm/v1/structured_output/backend_outlines.py | 8 +- 2 files changed, 78 insertions(+), 3 deletions(-) create mode 100644 tests/v1/structured_output/test_backend_outlines.py diff --git a/tests/v1/structured_output/test_backend_outlines.py b/tests/v1/structured_output/test_backend_outlines.py new file mode 100644 index 000000000000..abff1e7f51b6 --- /dev/null +++ b/tests/v1/structured_output/test_backend_outlines.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.exceptions import VLLMValidationError +from vllm.v1.structured_output.backend_outlines import validate_regex_is_buildable + +pytestmark = pytest.mark.cpu_test + + +# ================================================ +# Test validate_regex_is_buildable functionality +# ================================================ + + +class TestValidateRegexIsBuildable: + class TestPR60359Regressions: + """This class implements regression tests for PR #60359. + Initial bug report issue #60344. + validate_regex_is_buildable misses \b and backref inside groups + So, we check here that these features are rejected by us. + """ + + @pytest.mark.parametrize( + "regex", + [ + # top-level word boundary \b + r"a\b", + # \b inside capturing group + r"(a\b)", + # word boundary \b in nested group + r"((a\b))", + # top-level backref + r"(a)\1", + # backref in nested group + r"((a)\2)", + # positive lookahead + r"(a(?=b))", + # inline flag group + word boundary + r"(?i:a\b)", + # counted repeat over a group containing \b + r"(?:(a\b)){2}", + # lazy repeat + word boundary + r"(?:a\b)*?", + # possessive repeat + word boundary + r"(?:a\b)*+", + # atomic group + word boundary + r"(?>a\b)", + ], + ) + def test_unsupported_regex(self, regex): + with pytest.raises(VLLMValidationError): + validate_regex_is_buildable(regex) + + @pytest.mark.parametrize( + "regex", + [ + # [\b] is character class, not word-boundary \b + r"[\b]", + r"(a[\b]c)", + # escaped backslash + "b", not word-boundary \b + r"a\\b", + r"(a\\bc)", + # no \b or backref here, must be ok + r"(abc)", + # BRANCH and MAX_REPEAT must be ok + r"(a|b)+c", + r"(?:ab)*?c", + ], + ) + def test_supported_regex(self, regex): + validate_regex_is_buildable(regex) diff --git a/vllm/v1/structured_output/backend_outlines.py b/vllm/v1/structured_output/backend_outlines.py index 655b8b14631f..1cfc6c8f5d1a 100644 --- a/vllm/v1/structured_output/backend_outlines.py +++ b/vllm/v1/structured_output/backend_outlines.py @@ -327,9 +327,11 @@ def _check_unsupported(parsed) -> None: for branch in tval[1]: _check_unsupported(branch) - # tval is (min, max, subpattern) - elif ttype == _parser.MAX_REPEAT: - _check_unsupported(tval[2]) + else: + nested = tval if isinstance(tval, tuple) else (tval,) + for sub in nested: + if isinstance(sub, _parser.SubPattern): + _check_unsupported(sub) def validate_regex_is_buildable(pattern: str) -> None: From 83ed0c7c0944e9d124589d8fa3f82a019343a726 Mon Sep 17 00:00:00 2001 From: Luca Motz Date: Fri, 9 Oct 2026 13:30:11 +0200 Subject: [PATCH 174/624] [Bugfix][NIXL] Fix receive post-process for different P/D block sizes (#58860) Signed-off-by: Luca Motz Co-authored-by: Claude Opus 5.5 --- .../kv_connector/unit/test_nixl_connector.py | 16 ++++- .../unit/test_nixl_connector_hma.py | 66 ++++++++++++------- .../kv_transfer/kv_connector/utils.py | 44 ++----------- .../kv_connector/v1/nixl/base_worker.py | 54 ++++++--------- 4 files changed, 83 insertions(+), 97 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 52ac746391ff..d918c38fbf8c 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -1128,17 +1128,24 @@ def test_concurrent_load_kv( return raise TimeoutError("Took too long to complete async handshake.") + @pytest.mark.parametrize( + ("kv_cache_layout", "enable_permute_local_kv"), + # enable_permute_local_kv only covers a local LBNHC cache. + [("LBHNC", False), ("BLHNC", True)], + ) @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_handshake_fails_on_kv_cache_layout_mismatch( - self, default_vllm_config, dist_init + self, default_vllm_config, dist_init, kv_cache_layout, enable_permute_local_kv ): """Verify that adding a remote agent fails if kv_cache_layout differs. This test is only relevant for heterogeneous TP. """ - vllm_config = create_vllm_config() + vllm_config = create_vllm_config( + enable_permute_local_kv=enable_permute_local_kv + ) # Mock TP world size to 2 to force heterogeneous TP when # remote_tp_size=1 @@ -1151,7 +1158,10 @@ def test_handshake_fails_on_kv_cache_layout_mismatch( vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) ) connector.connector_worker = FakeNixlConnectorWorker( - vllm_config, connector.engine_id, hand_shake_latency=0 + vllm_config, + connector.engine_id, + hand_shake_latency=0, + kv_cache_layout=kv_cache_layout, ) worker = connector.connector_worker diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index f016b26475f0..c665ce0c5b88 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -1404,46 +1404,66 @@ def test_map_block_ids_for_block_size_ratio_hybrid(): @pytest.mark.cpu_test -def test_post_process_zeroes_untransferred_tail(): - """The untransferred sub-blocks of the last local block are zeroed on - receive; mamba state caches are untouched by the attention permute.""" +@pytest.mark.parametrize( + ("kv_cache_layout", "enable_permute_local_kv"), + [("LBHNC", False), ("LBNHC", False), ("LBNHC", True)], +) +def test_post_process_zeroes_untransferred_tail( + kv_cache_layout, enable_permute_local_kv +): + """Received remote sub-blocks are regrouped per head and the untransferred + sub-blocks of the last local block are zeroed on receive, once per block + although two attention groups alias the tensor.""" from unittest.mock import MagicMock from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( NixlConnectorWorker, ) - from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec ratio = 4 block_tokens = 8 # 2 tokens per remote sub-block + num_kv_heads = 2 # fewer than ratio worker = MagicMock(spec=NixlConnectorWorker) - worker._group_spec_types = (FullAttentionSpec, MambaSpec) worker.transfer_topo = MagicMock() worker.device_type = "cpu" - worker.enable_permute_local_kv = False - attn_cache = torch.ones(6, block_tokens, 2, 4) - mamba_cache = torch.ones(6, 16) - worker.device_kv_caches = {"attn.0": attn_cache, "mamba.0": mamba_cache} + worker.enable_permute_local_kv = enable_permute_local_kv + worker.kv_cache_layout = kv_cache_layout + # Attention caches are [B, H, N, C]; distinct values per head and token. + expected = torch.arange(6 * num_kv_heads * block_tokens * 4).view( + 6, num_kv_heads, block_tokens, 4 + ) + # Blocks [2, 3] as received: `ratio` head-major remote sub-blocks each. + received = ( + expected[2:4] + .unflatten(2, (ratio, -1)) + .transpose(1, 2) + .reshape(2, num_kv_heads, block_tokens, 4) + ) + if kv_cache_layout == "LBNHC": + attn_cache = expected.transpose(1, 2).contiguous().transpose(1, 2) + if enable_permute_local_kv: + # The remote is LBHNC: its bytes land in token-major memory. + attn_cache.transpose(1, 2)[2:4] = received.view(2, block_tokens, -1, 4) + # Otherwise token-major blocks receive the sub-blocks in token order. + else: + attn_cache = expected.clone() + attn_cache[2:4] = received + worker.device_kv_caches = {"attn.0": attn_cache, "swa.0": attn_cache} fa_group = MagicMock(layer_names=["attn.0"]) - ssm_group = MagicMock(layer_names=["mamba.0"]) - worker.kv_cache_config = MagicMock(transfer_groups=[fa_group, ssm_group]) - # The cached property filters mamba layers out of the permuted caches. - attn_caches = NixlConnectorWorker._attention_kv_caches.func(worker) - assert len(attn_caches) == 1 and attn_caches[0] is attn_cache - worker._attention_kv_caches = attn_caches + swa_group = MagicMock(layer_names=["swa.0"]) + worker.kv_cache_config = MagicMock(transfer_groups=[fa_group, swa_group]) _bind_worker_method(worker, "post_process_device_kv_on_receive") - # Request occupies blocks [2, 3]; only 6 of 8 sub-blocks were received. - worker.post_process_device_kv_on_receive(ratio, [([2, 3], 6)]) + # Group 0 request in blocks [2, 3]; only 6 of 8 sub-blocks were received. + worker.post_process_device_kv_on_receive(ratio, [(0, [2, 3], 6)]) # Block 2 fully covered; block 3 covered for 2 sub-blocks (4 tokens). - assert torch.all(attn_cache[2] == 1) - assert torch.all(attn_cache[3, :4] == 1) - assert torch.all(attn_cache[3, 4:] == 0) - # Untouched blocks and the mamba cache keep their content. - assert torch.all(attn_cache[4] == 1) - assert torch.all(mamba_cache == 1) + assert torch.equal(attn_cache[2], expected[2]) + assert torch.equal(attn_cache[3, :, :4], expected[3, :, :4]) + assert torch.all(attn_cache[3, :, 4:] == 0) + # Untouched blocks keep their content. + assert torch.equal(attn_cache[4], expected[4]) @pytest.mark.cpu_test diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index a0ead27c2ff1..d314d6107ad7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -256,8 +256,6 @@ def kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio): """ blocks_to_update = cache.index_select(0, indices) - # use physical order - blocks_to_update = blocks_to_update.permute(0, 2, 1, 3) n_kv_heads, block_size, head_size = blocks_to_update.shape[1:] remote_block_size = block_size // block_size_ratio n_blocks = block_size_ratio @@ -267,56 +265,28 @@ def kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio): .permute(0, 2, 1, 3, 4) .flatten(2, 3) ) - permuted_blocks = permuted_blocks.permute(0, 2, 1, 3) - cache.index_copy_(0, indices, permuted_blocks) - - -def kv_postprocess_layout_on_receive(cache, indices): - """Transforms the layout of received KV cache blocks to the local format. - - This method corrects layout mismatches from direct memory copies by - permuting the tensor dimensions. - - 4D cache: - - **Source Layout:** `[num_blocks, n_kv_head, block_size, head_dim]` - - **Target Layout:** `[num_blocks, block_size, n_kv_head, head_dim]` - 5D cache: - - **Source Layout:** `[num_blocks, kv_dim, n_kv_head, block_size, head_dim]` - - **Target Layout:** `[num_blocks, kv_dim, block_size, n_kv_head, head_dim]` - - Implementation: - - x = blocks_to_update.reshape(src_shape) # view local kv with sender layout - - permuted_blocks = x.permute(*inv_order) # transpose n_kv_heads, block_size - - cache.index_copy_(0, indices, permuted_blocks) # copy permuted kv back - - """ - blocks_to_update = cache.index_select(0, indices) - target_shape = list(blocks_to_update.shape) - target_shape[0] = -1 - inv_order = [0, 1, 3, 2, 4] if blocks_to_update.ndim == 5 else [0, 2, 1, 3] - src_shape = tuple(target_shape[i] for i in inv_order) - blocks_to_update = cache.index_select(0, indices) - permuted_blocks = blocks_to_update.reshape(src_shape).permute(*inv_order) cache.index_copy_(0, indices, permuted_blocks) def kv_postprocess_blksize_and_layout_on_receive(cache, indices, block_size_ratio): """Transforms the layout of received KV cache to the local block_size and LBHNC. - (Only works for local blocksize > remote blocksize) + (Only works for local blocksize >= remote blocksize) prefill is LBHNC, smaller block_size decode(local) is LBNHC, larger block_size """ blocks_to_update = cache.index_select(0, indices) - block_size, n_kv_heads, head_size = blocks_to_update.shape[1:] + n_kv_heads, block_size, head_size = blocks_to_update.shape[1:] remote_block_size = block_size // block_size_ratio n_blocks = block_size_ratio + # View the received bytes in memory order as the remote head-major blocks. permuted_blocks = ( - blocks_to_update.reshape(-1, n_blocks, n_kv_heads, remote_block_size, head_size) - .permute(0, 1, 3, 2, 4) - .flatten(1, 2) + blocks_to_update.transpose(1, 2) + .reshape(-1, n_blocks, n_kv_heads, remote_block_size, head_size) + .permute(0, 2, 1, 3, 4) + .flatten(2, 3) ) cache.index_copy_(0, indices, permuted_blocks) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 7ce6e93ebb46..38a656954eee 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -14,7 +14,6 @@ from collections import defaultdict from collections.abc import Iterator from concurrent.futures import Future, ThreadPoolExecutor -from functools import cached_property from typing import TYPE_CHECKING, Any, cast import msgspec @@ -31,7 +30,6 @@ get_current_attn_backends, kv_postprocess_blksize_and_layout_on_receive, kv_postprocess_blksize_on_receive, - kv_postprocess_layout_on_receive, ) from vllm.distributed.kv_transfer.kv_connector.v1.base import ( CopyBlocksOp, @@ -2596,6 +2594,7 @@ def _validate_remote_agent_handshake( if ( self.kv_transfer_config.enable_permute_local_kv and nixl_agent_meta.kv_cache_layout == "LBHNC" + and kv_cache_layout == "LBNHC" ): logger.info( "Remote is LBHNC and local is LBNHC, enabled additional permute " @@ -2607,9 +2606,11 @@ def _validate_remote_agent_handshake( self.enable_permute_local_kv = True else: raise RuntimeError( - "Heterogeneous TP expects same kv_cache_layout. " - "Or enable experimental feature to use HND to NHD support by " - "setting 'enable_permute_local_kv'=True in --kv-transfer-config." + "Heterogeneous TP expects same kv_cache_layout (local " + f"{kv_cache_layout}, remote {nixl_agent_meta.kv_cache_layout}). " + "Or enable experimental feature to receive LBHNC into a local " + "LBNHC cache by setting 'enable_permute_local_kv'=True " + "in --kv-transfer-config." ) # if remote_agent used attn is not same as local, # hint heterogenuous attn post process @@ -2795,26 +2796,10 @@ def save_kv_to_host(self, metadata: NixlConnectorMetadata): "d2h", ) - @cached_property - def _attention_kv_caches(self) -> list[torch.Tensor]: - """Device KV caches of attention layers (mamba states excluded), - as consumed by the receive post-process.""" - assert self.device_kv_caches, ( - "_attention_kv_caches accessed before register_kv_caches" - ) - mamba_layers = { - name - for g, group in enumerate(self.kv_cache_config.transfer_groups) - if _is_ssm_spec(self._group_spec_types[g]) - for name in group.layer_names - } - kv_caches = self.device_kv_caches - return [cache for name, cache in kv_caches.items() if name not in mamba_layers] - def post_process_device_kv_on_receive( self, block_size_ratio: int, - block_ids_list: list[tuple[list[int], int]], + block_ids_list: list[tuple[int, list[int], int]], convert: bool = True, ): """Post process device kv cache after receiving from remote. @@ -2863,9 +2848,11 @@ def post_process_device_kv_on_receive( block_size_ratio, ) - attn_caches = self._attention_kv_caches - device = attn_caches[0].device - for block_ids, covered_sub_blocks in block_ids_list: + for group, block_ids, covered_sub_blocks in block_ids_list: + attn_caches = [ + self.device_kv_caches[name] + for name in self.kv_cache_config.transfer_groups[group].layer_names + ] # Blocks the transfer didn't write: the token tail of the last # partially covered block, then everything beyond it. covered_blocks, sub_blocks_in_last = divmod( @@ -2875,17 +2862,16 @@ def post_process_device_kv_on_receive( has_stale = first_stale < len(block_ids) indices = None if convert or has_stale: - indices = async_tensor_h2d(block_ids, device, torch.long) + indices = async_tensor_h2d(block_ids, attn_caches[0].device, torch.long) if convert: for cache in attn_caches: - if self.enable_permute_local_kv and block_size_ratio > 1: + if self.enable_permute_local_kv: kv_postprocess_blksize_and_layout_on_receive( cache, indices, block_size_ratio ) - elif self.enable_permute_local_kv: - kv_postprocess_layout_on_receive(cache, indices) - else: + elif KVCacheLayout[self.kv_cache_layout].is_block_contiguous: + # Token-major blocks get the sub-blocks in token order. kv_postprocess_blksize_on_receive( cache, indices, block_size_ratio ) @@ -2893,10 +2879,10 @@ def post_process_device_kv_on_receive( if sub_blocks_in_last: last_block_id = block_ids[covered_blocks] for cache in attn_caches: - # Both post-processed layouts leave tokens on dim 1. - sub_block_tokens = cache.shape[1] // block_size_ratio + # Attention caches are [B, H, N, C]: tokens on dim 2. + sub_block_tokens = cache.shape[2] // block_size_ratio zero_from = sub_blocks_in_last * sub_block_tokens - cache[last_block_id, zero_from:].zero_() + cache[last_block_id, :, zero_from:].zero_() if has_stale: assert indices is not None stale_ids = indices[first_stale:] @@ -3068,7 +3054,7 @@ def get_transfer_results(self) -> KVConnectorTransferResults: len(meta.remote.block_ids[g]), ) block_ids_for_blocksize_post_process[block_size_ratio].append( - (local_group, covered_sub_blocks) + (g, local_group, covered_sub_blocks) ) # post processing for heterogeneous attention if self.enable_heterogeneous_attn_post_process: From 7368fadebd6c3f0c470d0e64b62cf6eea3b0987d Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Fri, 9 Oct 2026 19:34:38 +0800 Subject: [PATCH 175/624] [CI Failure] Move MTEB tests to h200_35gb (#60817) Signed-off-by: wang.yuqi --- .buildkite/test_areas/models_language.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index f123f2c57b13..002fad8ac6c3 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -220,10 +220,10 @@ steps: depends_on: - image-build-amd -- label: ":nvidia: (H200 MIG 18GB) Language Models (MTEB)" +- label: ":nvidia: (H200 MIG 35GB) Language Models (MTEB)" key: language-models-test-mteb timeout_in_minutes: 68 - device: h200_18gb + device: h200_35gb optional: true source_file_dependencies: - vllm/ From d095e61c0b389753ae86ccd7b2b7d1f2e1f5cd58 Mon Sep 17 00:00:00 2001 From: Chen Cheng <85005591+ischencheng@users.noreply.github.com> Date: Fri, 9 Oct 2026 07:35:48 -0400 Subject: [PATCH 176/624] [Bugfix][Spec Decode] Fix GLM-5.3-Flash DFlash2/EAGLE3 aux hidden-state capture with the upstream config (#60763) Signed-off-by: Chen Cheng <85005591+ischencheng@users.noreply.github.com> Signed-off-by: Isotr0py Co-authored-by: Isotr0py --- vllm/models/glm5next/common/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/models/glm5next/common/model.py b/vllm/models/glm5next/common/model.py index fd7c0263bdda..e35887ca7c54 100644 --- a/vllm/models/glm5next/common/model.py +++ b/vllm/models/glm5next/common/model.py @@ -788,7 +788,7 @@ def _aux_hidden_state( return hidden_states assert residual is not None and comb is not None completed = self._aux_post_op(hidden_states, residual, post, comb) - return hc_contract(completed, self.config.mhc_num_residual_streams) + return hc_contract(completed, self.config.hc_mult) def forward( self, From 1c47ebfc57d9e2415486e6080027ba319c73362a Mon Sep 17 00:00:00 2001 From: Praneeth Reddy A <72988597+0xViking@users.noreply.github.com> Date: Fri, 9 Oct 2026 17:13:55 +0530 Subject: [PATCH 177/624] [Docs] Fix VLLM_CPU_KVCACHE_SPACE default and units (#60053) Signed-off-by: Praneeth Co-authored-by: Li, Jiang --- docs/configuration/conserving_memory.md | 2 +- docs/getting_started/installation/cpu.md | 2 +- vllm/platforms/cpu.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/configuration/conserving_memory.md b/docs/configuration/conserving_memory.md index 96a903bc31d3..b397283a623f 100644 --- a/docs/configuration/conserving_memory.md +++ b/docs/configuration/conserving_memory.md @@ -80,7 +80,7 @@ llm = LLM(model="meta-llama/Llama-3.1-8B-Instruct", enforce_eager=True) If you run out of CPU RAM, try the following options: - (Multi-modal models only) you can set the size of multi-modal cache by setting `mm_processor_cache_gb` engine argument (default 4 GiB). -- (CPU backend only) you can set the size of KV cache using `VLLM_CPU_KVCACHE_SPACE` environment variable (default 4 GiB). +- (CPU backend only) you can set the size of KV cache using `VLLM_CPU_KVCACHE_SPACE` environment variable, in GiB (default 0, meaning vLLM doesn't override the KV cache size through this variable). ## Multi-modal input limits diff --git a/docs/getting_started/installation/cpu.md b/docs/getting_started/installation/cpu.md index d27753413823..1e58a68fcf89 100644 --- a/docs/getting_started/installation/cpu.md +++ b/docs/getting_started/installation/cpu.md @@ -293,7 +293,7 @@ See [AMD Zen optimizations](#amd-zen-optimizations) for detection rules, support ### How to decide `VLLM_CPU_KVCACHE_SPACE`? -This value is 4GB by default. Larger space can support more concurrent requests, longer context length. However, users should take care of memory capacity of each NUMA node. The memory usage of each TP rank is the sum of `weight shard size` and `VLLM_CPU_KVCACHE_SPACE`, if it exceeds the capacity of a single NUMA node, the TP worker will be killed with `exitcode 9` due to out-of-memory. +This value is 0 by default, which means vLLM doesn't override the KV cache size through this variable. When set, the value is in GiB (e.g, `VLLM_CPU_KVCACHE_SPACE`=40 means 40 GiB). Larger space can support more concurrent requests, longer context length. However, users should take care of memory capacity of each NUMA node. The memory usage of each TP rank is the sum of `weight shard size` and `VLLM_CPU_KVCACHE_SPACE`, if it exceeds the capacity of a single NUMA node, the TP worker will be killed with `exitcode 9` due to out-of-memory. This is a legacy setting; use `--kv-cache-memory-bytes` instead. ### How to do performance tuning for vLLM CPU? diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index e0257acb8007..43d472f22e67 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -323,7 +323,7 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: mamba_backend, ) - # Lagecy setting + # Legacy setting env_key = "VLLM_CPU_KVCACHE_SPACE" if env_key in os.environ and os.environ[env_key] != "": kv_cache_space = int(os.environ[env_key]) From 196483e25b43511e13aa179ed52f50574ecd9fee Mon Sep 17 00:00:00 2001 From: Louie Tsai Date: Fri, 9 Oct 2026 04:44:09 -0700 Subject: [PATCH 178/624] [Tools] Recipes sweep workflow (#57307) Signed-off-by: louie-tsai --- tools/recipes/README.md | 89 +- tools/recipes/RUNTIME_TUNING.md | 82 +- tools/recipes/SWEEP_TUNING.md | 662 ++++++- tools/recipes/VISUALIZATION.md | 73 + tools/recipes/hardware_detection.py | 85 +- tools/recipes/recipe_json_to_vllm_config.py | 260 ++- tools/recipes/requirements.txt | 4 + tools/recipes/sweep/__init__.py | 4 + tools/recipes/sweep/adaptive_concurrency.py | 261 +++ .../sweep/concurrency_recommendation.py | 330 ++++ tools/recipes/{ => sweep}/runtime_tuning.py | 77 +- .../recipes/sweep/scheduler_sweep_prepare.py | 225 +++ tools/recipes/sweep/sweep_generation.py | 1531 +++++++++++++++++ tools/recipes/sweep/sweep_html_report.py | 756 ++++++++ .../{ => sweep}/sweep_recommendation.py | 261 ++- tools/recipes/sweep/sweep_visualization.py | 226 +++ tools/recipes/sweep_generation.py | 358 ---- 17 files changed, 4720 insertions(+), 564 deletions(-) create mode 100644 tools/recipes/VISUALIZATION.md create mode 100644 tools/recipes/requirements.txt create mode 100644 tools/recipes/sweep/__init__.py create mode 100755 tools/recipes/sweep/adaptive_concurrency.py create mode 100644 tools/recipes/sweep/concurrency_recommendation.py rename tools/recipes/{ => sweep}/runtime_tuning.py (81%) create mode 100644 tools/recipes/sweep/scheduler_sweep_prepare.py create mode 100644 tools/recipes/sweep/sweep_generation.py create mode 100644 tools/recipes/sweep/sweep_html_report.py rename tools/recipes/{ => sweep}/sweep_recommendation.py (62%) create mode 100644 tools/recipes/sweep/sweep_visualization.py delete mode 100644 tools/recipes/sweep_generation.py diff --git a/tools/recipes/README.md b/tools/recipes/README.md index 90ad33b3eb10..37486ae5487e 100644 --- a/tools/recipes/README.md +++ b/tools/recipes/README.md @@ -19,7 +19,7 @@ flowchart LR ``` The recipe is the baseline. Hardware and workload information can optionally -refine the initial configuration. Sweep tuning is an optional validation step. +refine the initial configuration. Sweep tuning is optional. ## Getting Started @@ -54,7 +54,9 @@ For `xeon6`, the script enables hardware detection automatically. Generated `config.yml` and `env.sh` files remain in the directory where the script was invoked. -## 1. vLLM Recipes Only +## Detailed Deployment Workflow + +### 1. vLLM Recipes Only Use the converter directly when the recipe already contains the deployment settings you need. This path requires only PyYAML; the vLLM Python package is @@ -64,15 +66,7 @@ not required unless optional runtime tuning or sweep generation is requested. pip install pyyaml ``` -Choose whichever recipe-selection method fits the workflow: - -**Interactive discovery** — search models, then choose hardware and strategy: - -```bash -python3 tools/recipes/recipe_json_to_vllm_config.py -``` - -**Non-interactive discovery** — provide model and hardware and use the +For non-interactive discovery, provide the model and hardware and use the Recipes-recommended strategy: ```bash @@ -81,7 +75,17 @@ python3 tools/recipes/recipe_json_to_vllm_config.py \ --hardware xeon6 ``` -**Direct JSON input** — use a Recipes JSON URL or a local JSON file: +#### 1.1 Interactive Discovery + +Search models, then choose hardware and strategy interactively: + +```bash +python3 tools/recipes/recipe_json_to_vllm_config.py +``` + +#### 1.2 Direct JSON Input + +Use a Recipes JSON URL or a local JSON file: ```bash python3 tools/recipes/recipe_json_to_vllm_config.py \ @@ -90,15 +94,27 @@ python3 tools/recipes/recipe_json_to_vllm_config.py \ python3 tools/recipes/recipe_json_to_vllm_config.py recipe.json ``` -All paths generate `config.yml` and `env.sh`. See -[REFERENCE.md](REFERENCE.md) for recipe discovery, strategy selection, direct -JSON input, custom output files, and deployment scope. +#### 1.3 Test a Preview Recipe API -## 2. Hardware Information (Optional) +A Recipes pull request can expose the same JSON API through its Vercel preview. +Use `--api-base` to validate it before the recipe is available in production: -Add `--detect-hardware` when the target host's effective CPU/NUMA/memory -resources should refine deployment-sensitive values such as -`tensor-parallel-size` and `gpu-memory-utilization`. +```bash +PREVIEW=https://vllm-recipes-git-fork-intel-ai-tce-dockerin-f4c148-inferact-inc.vercel.app + +python3 tools/recipes/recipe_json_to_vllm_config.py \ + --api-base "$PREVIEW" \ + --model meta-llama/Llama-3.2-1B-Instruct \ + --hardware xeon6 +``` + +See [REFERENCE.md](REFERENCE.md) for recipe discovery, strategy selection, +custom output files, and deployment scope. + +### 2. Hardware Information (Optional) + +Add `--detect-hardware` when using the converter directly and the target host's +effective CPU/NUMA/memory resources should refine deployment-sensitive values: ```bash python3 tools/recipes/recipe_json_to_vllm_config.py \ @@ -107,14 +123,13 @@ python3 tools/recipes/recipe_json_to_vllm_config.py \ --detect-hardware ``` -Hardware detection is optional and uses vLLM CPU resource utilities only when -requested. See [RUNTIME_TUNING.md](RUNTIME_TUNING.md#hardware-information). +`serve_with_recipe.sh` enables this automatically for `xeon6`. See +[RUNTIME_TUNING.md](RUNTIME_TUNING.md#hardware-information) for details. -## 3. Workload Information (Optional) +### 3. Workload Information (Optional) -Workload hints can refine scheduler settings for one initial deployment -suggestion. Inputs include token lengths, concurrency, and optional latency or -capacity objectives. +Add workload hints when token lengths, concurrency, or latency objectives should +seed optional benchmark tuning: ```bash python3 tools/recipes/recipe_json_to_vllm_config.py \ @@ -127,17 +142,13 @@ python3 tools/recipes/recipe_json_to_vllm_config.py \ --tpot-sla-ms 100 ``` -Hardware detection and workload information are independent optional inputs; -they can also be supplied together. See -[RUNTIME_TUNING.md](RUNTIME_TUNING.md#workload-information) for the supported -inputs and how runtime parameters are calculated. +See [RUNTIME_TUNING.md](RUNTIME_TUNING.md#workload-information) for supported +inputs and runtime calculations. -## 4. Sweep Tuning (Optional) +### 4. Sweep Tuning (Optional) -Use `--generate-sweep` when the initial scheduler suggestion should be validated -with `vllm bench sweep serve`. The sweep benchmarks nearby scheduler values and -`recommend.py` produces one measured `recommended-config.yml` plus the -selection evidence in `recommendation.json`. +Use `--generate-full-sweep` for benchmark-backed tuning of TP/DP, concurrency, +and scheduler settings: ```bash python3 tools/recipes/recipe_json_to_vllm_config.py \ @@ -147,13 +158,15 @@ python3 tools/recipes/recipe_json_to_vllm_config.py \ --input-tokens 128 \ --output-tokens 128 \ --concurrency 32 \ - --generate-sweep + --generate-full-sweep ``` -See [SWEEP_TUNING.md](SWEEP_TUNING.md) for the benchmark, recommendation, and -vLLM CPU Docker-shell workflow. +See [SWEEP_TUNING.md](SWEEP_TUNING.md) for targeted sweep modes, NUMA binding, +failure handling, recommendations, reporting, and visualization. + +### 5. Start vLLM -## Start vLLM +When using the converter directly: ```bash source env.sh diff --git a/tools/recipes/RUNTIME_TUNING.md b/tools/recipes/RUNTIME_TUNING.md index d0b8ee929974..cf442a34d605 100644 --- a/tools/recipes/RUNTIME_TUNING.md +++ b/tools/recipes/RUNTIME_TUNING.md @@ -64,12 +64,10 @@ python3 tools/recipes/recipe_json_to_vllm_config.py \ --tpot-sla-ms 100 ``` -`--output-tokens` is used to estimate steady-state request turnover for -`max-num-batched-tokens`. When both `--target-qps` and `--tpot-sla-ms` are -available, they are also used to estimate prompt arrival pressure per scheduler -step. `--ttft-sla-ms` is collected but is not yet converted directly into a -batch-size formula because the relationship between TTFT and scheduler budget is -model- and hardware-dependent. +Scheduler calculation is opt-in with `--tune-scheduler`; workload hints do +not rewrite `max-num-seqs` or `max-num-batched-tokens` by default. The detailed +scheduler and staged-sweep behavior is documented in +[SWEEP_TUNING.md](SWEEP_TUNING.md). ## Deployment-Time Parameters @@ -81,11 +79,30 @@ on a validation environment that differs from the user's target deployment. | --- | --- | --- | --- | | `tensor-parallel-size` | Available topology and model head counts constrain TP. | Hardware topology + model configuration | Use the largest head-compatible power-of-two TP within the effective NUMA-node count, for DP=1 and PP=1. | | `gpu-memory-utilization` | Available memory can differ by machine size, container limits, and other memory use. The vLLM option name is also used by the CPU backend. | Hardware memory + recipe baseline | Calculate a conservative fraction from the most constrained NUMA node. | -| `max-num-seqs` | The useful scheduler concurrency depends on the number of requests expected to be active at the same time. | Workload concurrency | Set `max-num-seqs` to `--concurrency` when supplied. | -| `max-num-batched-tokens` | Each scheduler iteration must share its token budget between active decodes and incoming prefills. | Input/output token shape, concurrency, and optional QPS/TPOT | Calculate decode budget + expected prefill demand, with vLLM scheduler constraints as floors. | +| `max-num-seqs` | The useful scheduler concurrency depends on the number of requests expected to be active on each DP replica. | Workload concurrency + DP | Keep the recipe/vLLM value in `config.yml`; use `ceil(--concurrency / data-parallel-size)` as an explicit sweep seed. | +| `max-num-batched-tokens` | Each scheduler iteration must share its token budget between active decodes and incoming prefills. | Input/output token shape, concurrency, and optional QPS/TPOT | Keep the recipe/vLLM value in `config.yml`; use decode budget + expected prefill demand as an explicit sweep seed. | | `data-parallel-size` | The required replica count depends on the requested throughput and the capacity of one replica. | Capacity target | Keep the recipe value today because per-replica SLA capacity is not known. | -## How Each Runtime Parameter Is Calculated +For `max-num-seqs` and `max-num-batched-tokens`, workload formulas now create +benchmark-only seeds instead of overriding the directly deployable +`config.yml`. The optional benchmark sweep measures those explicit seeds +against vLLM's normal serving defaults. Default-reference runs leave the +selected scheduler option unset at engine-configuration time so vLLM can +resolve its platform-, world-size-, model-, and usage-context-aware default +rather than assuming a fixed numeric value. See +[SWEEP_TUNING.md](SWEEP_TUNING.md). + +This serving-default comparison is distinct from the `SchedulerConfig` fallback +constants used inside the initial tuning heuristics below. + +## How Initial Runtime Tuning Seeds Are Calculated + +The hardware formulas below produce directly deployable starting values. The +scheduler formulas produce explicit benchmark seeds and are not claimed to be +universal performance optima. This separation avoids baking workload guesses +into `config.yml` before measurements exist. The Tune All workflow benchmarks +TP/DP first, then concurrency, and finally scheduler settings before writing a +final recommendation. See [SWEEP_TUNING.md](SWEEP_TUNING.md). ### `tensor-parallel-size` @@ -147,14 +164,20 @@ candidate = min(recipe value or 0.80, safe_fraction) ### `max-num-seqs` -When `--concurrency` is supplied, it directly represents the requested maximum -number of simultaneously active requests: +When `--concurrency` is supplied, it represents global client concurrency. Each +DP replica needs capacity for its expected share of those requests: ```text -max-num-seqs = concurrency +per_replica_concurrency = ceil(concurrency / data_parallel_size) +max-num-seqs = per_replica_concurrency ``` -If concurrency is not supplied, the converter does not override this parameter. +For DP=1 this reduces to `max-num-seqs = concurrency`. This value seeds the +scheduler benchmark; it is not added to `config.yml`. If concurrency is not +supplied, no seed is calculated. Benchmark results show +that setting `max-num-seqs` below per-replica concurrency can create request +queueing and severe TTFT degradation; scheduler sweeps therefore do not use +below-concurrency candidates. ### `max-num-batched-tokens` @@ -165,7 +188,7 @@ First, the policy determines the active sequence count: ```text active_sequences = - --concurrency + ceil(--concurrency / data_parallel_size) or recipe max-num-seqs or vLLM default max_num_seqs (128) ``` @@ -181,17 +204,19 @@ The policy then estimates how many new prompts need prefill work per scheduler step: ```text +per_replica_qps = target_qps / data_parallel_size + prefills_per_step = max( 1, active_sequences / output_tokens, - target_qps * tpot_sla_ms / 1000 + per_replica_qps * tpot_sla_ms / 1000 ) prefills_per_step = min(active_sequences, prefills_per_step) prefill_budget = ceil(input_tokens * prefills_per_step) ``` -The final scheduler budget is: +The final explicit scheduler sweep seed is: ```text max-num-batched-tokens = max( @@ -201,7 +226,7 @@ max-num-batched-tokens = max( ) ``` -For example, with 128 input tokens, 128 output tokens, and concurrency 32: +For example, with 128 input tokens, 128 output tokens, concurrency 32, and DP=1: ```text decode_budget = 32 @@ -210,6 +235,12 @@ prefill_budget = 128 candidate = max(2048, 32, 32 + 128) = 2048 ``` +The calculated token budget is a scheduler sweep seed and is not added to +`config.yml`. If P99 TTFT still +misses its objective while P99 TPOT has substantial headroom, Tune All should +reduce workload concurrency or select a different TP/DP layout rather than +reducing `max-num-seqs` below per-replica concurrency. + If chunked prefill is explicitly disabled and `max-model-len` is available, the policy also ensures: @@ -226,8 +257,10 @@ also needs measured per-replica throughput that still satisfies TTFT/TPOT: DP = ceil(target_qps / qps_per_replica_at_SLO) ``` -Because the converter does not have that measured capacity yet, it intentionally -keeps the recipe DP value rather than guessing. +Because the converter does not have that measured capacity yet, its initial +configuration intentionally keeps the recipe DP value rather than guessing. +Tune All subsequently benchmarks supported TP/DP layouts and replaces this seed +with a measured layout before concurrency and scheduler tuning. ## Precedence @@ -235,8 +268,11 @@ keeps the recipe DP value rather than guessing. vLLM defaults -> vLLM Recipes baseline -> hardware refinement (optional) - -> workload / SLO refinement (optional) - -> config.yml + env.sh + -> config.yml + env.sh + +workload / SLO hints (optional) + -> explicit scheduler sweep seed + -> benchmark-backed recommended-config.yml ``` ## Runtime-Tuning Hardware Scope @@ -312,10 +348,10 @@ with [SWEEP_TUNING.md](SWEEP_TUNING.md). ## Implementation - `hardware_detection.py` collects effective CPU/NUMA/memory information. -- `runtime_tuning.py` owns hardware-policy selection and independent parameter +- `sweep/runtime_tuning.py` owns hardware-policy selection and independent parameter tuning functions. - `recipe_json_to_vllm_config.py` resolves the recipe, collects optional inputs, applies the selected policy, and generates `config.yml` and `env.sh`. -For benchmark-backed validation of the initial scheduler suggestion, see +For benchmark-backed validation of the explicit scheduler seed, see [SWEEP_TUNING.md](SWEEP_TUNING.md). diff --git a/tools/recipes/SWEEP_TUNING.md b/tools/recipes/SWEEP_TUNING.md index e7282e122d78..e39444aef983 100644 --- a/tools/recipes/SWEEP_TUNING.md +++ b/tools/recipes/SWEEP_TUNING.md @@ -1,13 +1,133 @@ # Sweep Tuning -Sweep tuning is optional. The converter always creates one initial `config.yml` -first. Users can deploy it immediately or benchmark nearby scheduler settings -and generate one measured recommendation. +Sweep tuning is optional. The converter always creates an initial `config.yml` +that can be deployed directly. Scheduler calculation is disabled by default: +TP/DP and concurrency stages remove explicit `max-num-seqs` and +`max-num-batched-tokens` so vLLM resolves its normal scheduler defaults. -## Generate +Use `--tune-scheduler` only when you explicitly want workload-derived scheduler +values before the dedicated scheduler stage. For benchmark-backed tuning, the +recommended workflow is **Tune All**, which measures the serving stack in +dependency order: + +```text +TP/DP -> concurrency -> scheduler +``` + +The examples below assume tuning is run inside the target vLLM CPU container so +hardware detection sees the same CPU, NUMA, memory, and cgroup limits as the +deployment. + +## 1. Start the vLLM CPU Docker Shell + +From the vLLM source tree: ```bash -python3 tools/recipes/recipe_json_to_vllm_config.py \ +mkdir -p recipe-output + +docker run --rm -it \ + --entrypoint bash \ + --security-opt seccomp=unconfined \ + --cap-add SYS_NICE \ + --shm-size=4g \ + -p 8000:8000 \ + -v "$PWD/tools/recipes:/recipes:ro" \ + -v "$PWD/recipe-output:/output" \ + -v "$HOME/.cache/huggingface:/root/.cache/huggingface" \ + -w /output \ + vllm/vllm-openai-cpu:latest-x86_64 +``` + +`/recipes` is mounted read-only from the source tree and `/output` is writable. +Run the converter, generated sweep scripts, and hardware detection inside this +container. + +For additional container details, see +[RUNTIME_TUNING.md](RUNTIME_TUNING.md#vllm-cpu-docker-shell). + +## 2. Recommended: Tune All + +Use `--generate-full-sweep` by default when you want benchmark-backed runtime +tuning. It measures TP/DP first, then the highest useful concurrency, then +scheduler parameters using the selected layout and concurrency. + +```mermaid +flowchart TD + I[["Input
Recipe + detected hardware
fixed input/output workload shape"]] + + A1("1. Tune TP/DP
tensor-parallel-size + data-parallel-size") + O1[/"parallel-layout-config.yml
selected TP/DP"/] + + A2("2. Tune concurrency
max_concurrency") + O2[/"concurrency-recommendation.json
recommended max_concurrency"/] + + A3("3. Tune scheduler
max-num-seqs + max-num-batched-tokens") + O3[/"recommended-config.yml
final vLLM server config"/] + + I --> A1 + A1 --> O1 + O1 --> A2 + A2 --> O2 + O2 --> A3 + A3 --> O3 +``` + +### Workflow options + +The table is ordered by the tuning sequence. **Tune All** is the recommended +default; the other modes are useful when only one part of the serving stack +needs to be re-measured. + +| Goal | Option | Tuning sequence | +| --- | --- | --- | +| **Tune all (recommended)** | `--generate-full-sweep` | **1. TP/DP -> 2. concurrency -> 3. scheduler** | +| Tune TP/DP only | `--generate-parallel-layout-sweep` | 1. TP/DP | +| Tune concurrency only | `--generate-concurrency-sweep` | 2. concurrency | +| Tune scheduler only | `--generate-scheduler-sweep` (`--generate-sweep` alias) | 3. scheduler | + +All workflows keep one fixed input/output workload shape and require +`--input-tokens`, `--output-tokens`, and `--concurrency`. + +For **Tune All**, the supplied `--concurrency` is the representative load used +for the initial TP/DP comparison and to size the benchmark request set. The +concurrency stage then measures the final SLA-feasible `max_concurrency`. + +### Scheduler defaults and `--tune-scheduler` + +By default, `--concurrency`, token lengths, and SLA values are workload hints; +they do not cause the converter to calculate `max-num-seqs` or +`max-num-batched-tokens`. This keeps scheduler capacity independent from the +seed concurrency while TP/DP and maximum concurrency are being measured. + +The default Tune All policy is: + +```text +Stage 1: TP/DP -> vLLM scheduler defaults +Stage 2: concurrency -> vLLM scheduler defaults +Stage 3: scheduler -> calculate and benchmark scheduler candidates +``` + +For a directly deployable heuristic configuration, or to intentionally carry +workload-derived scheduler values into the earlier sweep stages, opt in with: + +```bash +--tune-scheduler +``` + +With that flag, the existing workload formulas are enabled. In particular, +`max-num-seqs` uses the per-replica workload concurrency and +`max-num-batched-tokens` uses the estimated decode and prefill pressure. + +A scheduler-only sweep is also an explicit opt-in to scheduler calculation: it +uses the workload formulas as benchmark seeds and compares them with vLLM +defaults. + +### Generate the Tune All package + +Inside the container: + +```bash +python3 /recipes/recipe_json_to_vllm_config.py \ --model meta-llama/Llama-3.1-8B-Instruct \ --hardware xeon6 \ --detect-hardware \ @@ -16,80 +136,311 @@ python3 tools/recipes/recipe_json_to_vllm_config.py \ --concurrency 32 \ --ttft-sla-ms 3000 \ --tpot-sla-ms 100 \ - --generate-sweep + --config-out /output/config.yml \ + --env-out /output/env.sh \ + --generate-full-sweep \ + --sweep-out-dir /output/sweep +``` + +The initial configuration remains directly deployable: + +```bash +source /output/env.sh +vllm serve --config /output/config.yml +``` + +Before measurements, `config.yml` retains the recipe/vLLM scheduler defaults. +The sweep package separately carries explicit workload-derived values for +`max-num-seqs` and `max-num-batched-tokens`, so the benchmark can compare both +policies without making the workload heuristic the deployment default. + +Stop a manually started server before running the sweep because the generated +sweep scripts start and stop their own vLLM servers. + +Run all tuning stages: + +```bash +/output/sweep/run_full_sweep.sh ``` -The optional package is: +The end-to-end flow produces intermediate recommendations for TP/DP and +concurrency and finishes with: ```text -sweep/ -├── serve_params.json -├── bench_params.json -├── run_sweep.sh -├── recommend.py -└── SWEEP.md +/output/sweep/recommended-config.yml +/output/sweep/recommendation.json +/output/sweep/sweep-report.html ``` -The sweep varies only `max-num-seqs` and `max-num-batched-tokens`. It uses an -eight-point directed design: a batch-budget curve at the initial sequence count -plus lower/higher batch interactions at three-quarters and one-half of that -count. This gives broader coverage than a one-parameter-at-a-time sweep without -the cost of a full Cartesian grid. +Inspect the final recommendation: -## Run and Recommend +```bash +cat /output/sweep/recommendation.json +``` + +## 3. Tune All Stage Details + +### Stage 1: Tune TP/DP + +The parallel-layout stage measures NUMA-aware combinations of: + +- `tensor-parallel-size` +- `data-parallel-size` + +Primary candidates use all effective NUMA nodes: + +```text +tensor-parallel-size * data-parallel-size = effective NUMA nodes +``` + +The effective NUMA-node count comes from hardware detection; it is not fixed by +the sweep generator. The table below shows representative detected topologies. + +TP is restricted to the supported values `1`, `2`, `4`, and `8`. The sweep also +includes the largest supported TP size that does not exceed the effective +NUMA-node count, even if that layout leaves some NUMA nodes idle. + +| Effective NUMA nodes | Generated layouts | +| ---: | --- | +| 2 | `TP=2, DP=1`; `TP=1, DP=2` | +| 4 | `TP=4, DP=1`; `TP=2, DP=2`; `TP=1, DP=4` | +| 6 | `TP=4, DP=1` (4 of 6 nodes); `TP=2, DP=3`; `TP=1, DP=6` | +| 8 | `TP=8, DP=1`; `TP=4, DP=2`; `TP=2, DP=4`; `TP=1, DP=8` | + +The generator then applies model- and runtime-aware safety policy before any +server is started: + +1. For an MoE model, if the CPU communicator does not override the variable-size + `all_gatherv` and `reduce_scatterv` collectives used by MoE data parallelism, + all generated candidates use `DP=1`. Every supported TP value is retained, + so a 6-NUMA-node host tests `TP=4`, `TP=2`, and `TP=1` rather than only `TP=4`. +2. When MoE dimensions are known and expert parallelism is disabled, the fused + MoE path flattens TP, DP, and PCP. A candidate is valid only when + `moe_intermediate_size % (TP * DP * PCP) == 0`. +3. Failures not covered by a preflight rule remain isolated by + `--continue-on-error` when the installed CLI supports it, with bounded + `--resume` retries as the compatibility fallback. + +Skipped candidates and reasons are saved in +`sweep/parallel_layout_skips.json`. If model metadata cannot be loaded, the +generator keeps the generic layouts and relies on the third policy. + +By default, parallel-layout candidates change only TP and DP. The generated +benchmark config removes explicit `max-num-seqs` and +`max-num-batched-tokens`, so every layout uses vLLM-resolved scheduler +defaults. This prevents a scheduler value derived from the seed concurrency +from becoming part of the TP/DP selection. + +With `--tune-scheduler`, the workload-derived per-replica scheduler baseline is +restored for this stage. + +The stage writes: + +```text +sweep/parallel-layout-config.yml +sweep/parallel-layout-recommendation.json +``` + +The selected TP/DP layout becomes the fixed server layout for Stage 2. + +### Temporary TP/DP NUMA-binding workaround + +Xeon TP/DP and Tune All sweeps currently enable a temporary explicit CPU-binding +workaround by default: + +```text +--tp-dp-numa-bind-workaround +``` + +Hardware detection builds one `VLLM_CPU_OMP_THREADS_BIND` CPU list per effective +NUMA node. On x86 it selects one logical CPU per physical core, matching vLLM's +auto-binding SMT policy, and excludes one physical core from each NUMA node for +non-OMP work. The resulting value is written to `env.sh`, for example: ```bash -sweep/run_sweep.sh --dry-run -sweep/run_sweep.sh -sweep/recommend.py +export VLLM_CPU_OMP_THREADS_BIND='48-62|64-78|80-94|96-110' ``` -Resume an interrupted sweep with: +The exact CPU IDs depend on the effective container/cgroup cpuset and NUMA +topology. Discontinuous CPU IDs are preserved as comma-separated ranges. + +`VLLM_CPU_NUM_OF_RESERVED_CPU` is not used by this workaround because vLLM only +applies its reserved-CPU logic in automatic binding mode. With an explicit +`VLLM_CPU_OMP_THREADS_BIND`, reserved cores must already be omitted from the +generated lists. + +After the vLLM CPU DP NUMA-binding issue is fixed, disable the workaround +without removing the implementation: ```bash -sweep/run_sweep.sh --resume +--no-tp-dp-numa-bind-workaround +``` + +An explicit non-`auto` `VLLM_CPU_OMP_THREADS_BIND` supplied by the recipe is +preserved. + +### Stage 2: Tune Concurrency + +The concurrency stage keeps the selected TP/DP layout fixed while finding the +largest useful client concurrency. By default the generated benchmark config +removes explicit `max-num-seqs` and `max-num-batched-tokens`, so the seed +concurrency cannot become an artificial scheduler admission limit. + +With `--tune-scheduler`, explicit scheduler values are preserved as an opt-in +behavior. + +When TTFT and/or TPOT objectives are supplied, Recipes uses an **adaptive SLA +boundary search** instead of uniformly sampling the entire workload range. + +The supplied `--concurrency` is the seed. The controller grows upward until it +finds the first SLA failure (or shrinks downward when the seed fails), then +binary-searches the integer PASS/FAIL bracket. One vLLM server remains running +throughout this refinement, so midpoint probes do not reload the model. + +For example: + +```text +32 PASS, 64 FAIL + -> 48 + -> 40 or 56 + -> ... + -> adjacent PASS / FAIL boundary +``` + +A point is SLA-feasible only when all of the following hold: + +1. Every measured run completes without failed requests. +2. Median P99 TTFT is at or below the supplied TTFT objective. +3. Median P99 TPOT is at or below the supplied TPOT objective. +4. Duration-weighted combined request compliance is at least the configured + minimum (`0.99`, or 99%, by default). + +The combined-compliance check is intentionally stricter than looking at the +TTFT and TPOT P99 columns independently. `vllm bench serve --goodput` counts +requests that satisfy all supplied latency objectives, and the recommender +estimates the combined compliant-request fraction across repeated runs as: + +```text +sum(request_goodput * duration) / sum(completed_requests) +``` + +Therefore a candidate can have P99 TTFT below 3000 ms and P99 TPOT below +100 ms but still be SLA-ineligible if fewer than 99% of requests satisfy both +objectives together. For example, a reported combined compliance of `98.96%` +fails the default `99.00%` requirement even when both displayed P99 metrics +pass. The HTML report shows the combined compliance, required threshold, and +the specific SLA failure reason for each candidate. + +Adaptive concurrency uses a `1800` second server-readiness timeout by default, +which is intentionally longer than the generic benchmark helper default for +CPU model loading and initialization. `run_full_sweep.sh +--server-ready-timeout VALUE` also exports the selected timeout explicitly to +the generated concurrency runner, in addition to forwarding the CLI option. +This avoids falling back to a shorter adaptive-helper timeout when the +concurrency stage is invoked through the full staged workflow. + +The default growth factor is `2` and the default hard cap is `1000`. +Generated runners accept `--growth-factor VALUE` and +`--max-concurrency-cap VALUE`. + +Without TTFT/TPOT objectives there is no SLA boundary. In that mode the +generated runner keeps vLLM Workload Explorer: + +```text +vllm bench sweep serve_workload --workload-var max_concurrency +``` + +The stage writes: + +```text +sweep/concurrency-recommendation.json +``` + +`max_concurrency` is benchmark/deployment metadata. It is **not** written as a +`vllm serve` configuration parameter. + +### Stage 3: Tune Scheduler + +Before the scheduler sweep, `prepare_scheduler.py` recalculates the scheduler +baseline from: + +```text +selected TP/DP + + +selected max_concurrency + + +fixed input/output token shape +``` + +The scheduler stage keeps an explicit `max-num-seqs` at the per-replica +concurrency lower bound while tuning `max-num-batched-tokens`. It also compares +the explicit sequence limit with the vLLM default. + +The per-replica sequence baseline is recalculated as: + +```text +max-num-seqs = ceil(selected max_concurrency / selected data-parallel-size) ``` -`recommend.py` writes: +The directed scheduler sweep does not test `max-num-seqs` below this value; +doing so limits immediately active requests and can turn scheduler queueing into +severe TTFT degradation. It keeps the existing batch-budget curve around the +baseline and also measures three vLLM-default reference candidates: + +| Reference candidate | `max-num-seqs` | `max-num-batched-tokens` | +| --- | --- | --- | +| `vllm_default_max_num_seqs` | vLLM default | generated value | +| `vllm_default_max_num_batched_tokens` | generated value | vLLM default | +| `vllm_defaults` | vLLM default | vLLM default | + +The default references do not hard-code values such as 128 or 2048. The sweep +omits the selected scheduler argument and lets the installed vLLM resolve its +normal platform-, world-size-, model-, and usage-context-aware default. + +If a default-reference candidate wins, the corresponding scheduler key is +omitted from `recommended-config.yml`. + +## 4. Benchmark Size, Warmup, and Failure Handling + +Measured request count scales with concurrency using: ```text -sweep/recommended-config.yml -sweep/recommendation.json +num_prompts = min(1000, max(100, concurrency * 10)) ``` -With TTFT/TPOT objectives, the benchmark uses vLLM `--goodput`. The recommender -calculates duration-weighted combined compliance across repeated runs and -requires both median P99 latency objectives and the minimum compliance ratio -(default `0.99`). Among eligible configurations it selects the highest mean -output-token throughput. Use `recommend.py --minimum-compliance VALUE` to -change the compliance requirement. +This targets about ten concurrency turnovers when neither bound applies. The +100-prompt floor avoids very small samples for P99/compliance measurements, and +the 1000-prompt cap bounds sweep runtime. + +Before a measured parameter combination, generated scripts use one unmeasured +warmup containing one full concurrency window. Warmup output is retained as +`warmup.json` for auditing but is excluded from recommendation statistics. -If no configuration is eligible, `recommend.py` records the highest-goodput -configuration as `best_effort` in `recommendation.json`, does not write a -deployable `recommended-config.yml`, and exits with status 2. Without latency -objectives it selects highest mean output-token throughput. Failed benchmark -configurations are excluded in either mode. +Generated scripts also handle sweep-run failures without discarding completed +work: -When an eligible configuration exists, `recommended-config.yml` changes only -the two swept scheduler parameters. `recommendation.json` records mean, median, -and worst-run P99 values, combined compliance, and the measured evidence for -every candidate. +- If the installed vLLM CLI supports `--continue-on-error`, the generated script + enables it. +- If the installed vLLM does not expose that option, the script omits it and + automatically retries an interrupted sweep with `--resume`. +- The default is two automatic resume retries. Override it with + `VLLM_RECIPE_SWEEP_RETRIES`. -Deploy: +For example: ```bash -source env.sh -vllm serve --config sweep/recommended-config.yml +export VLLM_RECIPE_SWEEP_RETRIES=3 +/output/sweep/run_full_sweep.sh ``` -## vLLM CPU Docker Shell +## 5. Targeted Tuning Workflows + +Use these modes when a full re-tune is unnecessary. -Use the common CPU container setup in -[RUNTIME_TUNING.md](RUNTIME_TUNING.md#vllm-cpu-docker-shell). It ensures -hardware detection runs against the same effective CPU, NUMA, memory, and -cgroup limits used by the deployment. +### Tune TP/DP Only -Inside that container, generate the initial suggestion plus sweep package: +Use this when only the parallel layout needs to be measured. This mode stops +after selecting `tensor-parallel-size` and `data-parallel-size`: ```bash python3 /recipes/recipe_json_to_vllm_config.py \ @@ -103,35 +454,224 @@ python3 /recipes/recipe_json_to_vllm_config.py \ --tpot-sla-ms 100 \ --config-out /output/config.yml \ --env-out /output/env.sh \ - --generate-sweep \ + --generate-parallel-layout-sweep \ --sweep-out-dir /output/sweep ``` -The initial configuration can be deployed directly: +Run: ```bash -source /output/env.sh -vllm serve --config /output/config.yml +/output/sweep/run_parallel_layout_sweep.sh +/output/sweep/recommend_parallel_layout.py +``` + +The result is `/output/sweep/parallel-layout-config.yml`. Use Tune All when +concurrency and scheduler tuning should follow the selected layout. + +### Tune Concurrency Only + +Use this when TP/DP and scheduler configuration are already fixed and only +concurrent serving capacity needs to be measured: + +```bash +python3 /recipes/recipe_json_to_vllm_config.py \ + --model meta-llama/Llama-3.1-8B-Instruct \ + --hardware xeon6 \ + --detect-hardware \ + --input-tokens 128 \ + --output-tokens 128 \ + --concurrency 32 \ + --ttft-sla-ms 3000 \ + --tpot-sla-ms 100 \ + --config-out /output/config.yml \ + --env-out /output/env.sh \ + --generate-concurrency-sweep \ + --sweep-out-dir /output/sweep +``` + +Run: + +```bash +/output/sweep/run_concurrency_sweep.sh +/output/sweep/recommend_concurrency.py +``` + +### Tune Scheduler Only + +Use this when TP/DP and target concurrency are already known: + +```bash +python3 /recipes/recipe_json_to_vllm_config.py \ + --model meta-llama/Llama-3.1-8B-Instruct \ + --hardware xeon6 \ + --detect-hardware \ + --input-tokens 128 \ + --output-tokens 128 \ + --concurrency 32 \ + --ttft-sla-ms 3000 \ + --tpot-sla-ms 100 \ + --config-out /output/config.yml \ + --env-out /output/env.sh \ + --generate-scheduler-sweep \ + --sweep-out-dir /output/sweep ``` -Stop that manually started server before running the sweep, because -`run_sweep.sh` starts and stops its own vLLM servers. +`--generate-sweep` remains a backward-compatible alias for +`--generate-scheduler-sweep`. + +Run: ```bash -/output/sweep/run_sweep.sh --dry-run /output/sweep/run_sweep.sh /output/sweep/recommend.py ``` -Inspect: +## 6. Recommendation Policy + +With TTFT/TPOT objectives, benchmark commands use vLLM `--goodput`. + +For TP/DP and scheduler selection, the recommender: + +1. Excludes invalid configurations and failed benchmark measurements. +2. Calculates duration-weighted combined request compliance across successful + runs from request goodput and run duration. +3. Requires the median P99 TTFT and P99 TPOT objectives independently. +4. Also requires the combined request-compliance ratio to meet + `--minimum-compliance`, which defaults to `0.99`. +5. Establishes the highest mean output-token throughput among eligible + configurations, then applies the selection policy below. + +The P99 checks and combined-compliance check answer different questions. P99 +verifies the tail of each latency metric independently; combined compliance +verifies that the required fraction of requests satisfy all configured +objectives together. A candidate must pass both forms of validation to be +marked SLA eligible. + +For a scheduler-only comparison with one fixed TP/DP layout, candidates within +1% of the highest eligible output-token throughput are treated as practically +equivalent. Within that set, the recommender prefers the candidate with fewer +explicit scheduler overrides, followed by compliance, goodput, and throughput. +This avoids hard-coding a scheduler value for a difference that is likely within +normal run-to-run variation. Change or disable the equivalence range with: ```bash -cat /output/sweep/recommendation.json +/output/sweep/recommend.py --throughput-equivalence-percent VALUE +``` + +Use `0` to restore exact highest-throughput selection. TP/DP comparisons still +use exact highest-throughput selection because changing a parallel layout is a +material deployment decision. + +For concurrency selection, the policy instead selects the **highest +SLA-feasible `max_concurrency`**, with throughput/compliance metrics used to +break ties. + +Change the compliance threshold with: + +```bash +/output/sweep/recommend.py --minimum-compliance VALUE +``` + +If no configuration is SLA-feasible, the recommendation JSON records a +best-effort candidate. A deployable final configuration is produced only when +the required recommendation stages succeed. + +## 7. Post-Benchmark Reporting + +After Tune All finishes, `run_full_sweep.sh` automatically runs `report.py` and +writes a self-contained summary: + +```text +/output/sweep/sweep-report.html +``` + +The HTML report includes the selected configuration, a stage summary, SLA +metrics, candidate tables, and automatic coverage observations. It uses only +Python's standard library. A report-generation error is printed as a warning +and does not change the success of the completed sweep. + +The TTFT and TPOT objectives supplied to `recipe_json_to_vllm_config.py` are +embedded in every generated recommender and in `report.py`. +`run_full_sweep.sh` also passes those values explicitly to every recommendation +stage and the report. Report generation stops with an error if a stale or +manually generated recommendation JSON contains different SLA values. + +Regenerate the report manually when recommendation files change: + +```bash +cd /output/sweep +./report.py +``` + +### Optional figures + +Every generated sweep package includes a standalone analysis helper. It is not +called by `run_full_sweep.sh` or any stage runner, so figure generation cannot +start, resume, or alter a sweep. + +After one or more benchmark stages finish, install the plotting dependencies: + +```bash +cd /output/sweep +python3 -m pip install -r requirements.txt ``` -Deploy the measured result: +Preview the planned figures, then generate them from the completed +`summary.json` files: + +```bash +./visualize.py --dry-run +./visualize.py +``` + +Use `--stage parallel-layout`, `--stage concurrency-tuning`, or +`--stage runtime-tuning` to analyze selected stages. The helper writes PNG files +under each stage's `figures/` directory. See the generated `VISUALIZATION.md` +for the figure set and additional examples. + +## 8. Deploy the Tuned Configuration + +After Tune All completes: ```bash source /output/env.sh vllm serve --config /output/sweep/recommended-config.yml ``` + +## Select sweep stages + +Sweep implementation and reporting helpers live in `tools/recipes/sweep/`. +The converter remains `tools/recipes/recipe_json_to_vllm_config.py`. + +Use `--sweep-stage concurrency` to generate only a concurrency sweep against +the supplied recipe's fixed TP/DP layout, for example: + +```bash +python3 tools/recipes/recipe_json_to_vllm_config.py \ + --model meta-llama/Llama-3.1-8B-Instruct --hardware xeon6 \ + --input-tokens 128 --output-tokens 128 --concurrency 32 \ + --ttft-sla-ms 3000 --tpot-sla-ms 100 \ + --sweep-stage concurrency --sweep-out-dir sweep +./sweep/run_concurrency_sweep.sh +./sweep/recommend_concurrency.py +``` + +`--sweep-stage` accepts `parallel-layout`, `concurrency`, `scheduler`, or `all`. +The existing `--generate-*-sweep` flags remain supported; do not combine them +with `--sweep-stage`. `parallel-layout` and `all` require `--detect-hardware`. + +In a package generated with `--sweep-stage all` or `--generate-full-sweep`, +select the stages to execute: + +```bash +./sweep/run_full_sweep.sh --stages concurrency +./sweep/run_full_sweep.sh --stages concurrency,scheduler +./sweep/run_full_sweep.sh --stages scheduler --resume +``` + +The default is all three stages. Selected stages run in dependency order. +Skipping parallel-layout reuses `parallel-layout-config.yml` if present; +otherwise the runner copies the generated `initial-config.yml` as the fixed +layout. Scheduler-only execution requires `concurrency-recommendation.json` +from an earlier concurrency run. Remaining arguments go to the selected sweep +runners. Use separate output directories for unrelated experiments. diff --git a/tools/recipes/VISUALIZATION.md b/tools/recipes/VISUALIZATION.md new file mode 100644 index 000000000000..4c109102e7ab --- /dev/null +++ b/tools/recipes/VISUALIZATION.md @@ -0,0 +1,73 @@ +# Post-Benchmark Reporting + +## HTML summary + +`run_full_sweep.sh` automatically generates a self-contained HTML report after +the final recommendation: + +```text +sweep-report.html +``` + +The report uses Python's standard library and does not require plotting +dependencies. It summarizes selected settings, stage results, SLA metrics, and +all measured candidates. Report-generation failures are warnings and do not +change sweep success. + +The generated `report.py` uses the same TTFT and TPOT objectives supplied when +the sweep package was created. It displays both objectives and validates them +against every available recommendation JSON before writing the report. + +Regenerate it manually or select a different output path: + +```bash +./report.py +./report.py --output results/my-sweep-report.html +``` + +## Optional figures + +Visualization is an independent analysis step. It reads completed sweep +results and does not start a server, run a benchmark, or change a tuning +recommendation. + +Install the plotting dependencies in the same environment as vLLM: + +```bash +python3 -m pip install -r requirements.txt +``` + +Preview the figures that would be generated: + +```bash +./visualize.py --dry-run +``` + +Generate figures for every completed stage: + +```bash +./visualize.py +``` + +Select one or more stages when needed: + +```bash +./visualize.py --stage concurrency-tuning +./visualize.py --stage parallel-layout --stage runtime-tuning +``` + +Figures are written beside the source results: + +```text +results/parallel-layout/figures/ +results/concurrency-tuning/figures/ +results/runtime-tuning/figures/ +``` + +The analysis produces output-throughput, P99 TTFT, and P99 TPOT figures. The +concurrency stage also produces a throughput/TTFT tradeoff figure and a goodput +figure when the benchmark results contain that metric. + +The native vLLM Pareto command is intentionally not used because its efficiency +axis is expressed as tokens per second per GPU, which is not an accurate label +for Xeon CPU layouts. diff --git a/tools/recipes/hardware_detection.py b/tools/recipes/hardware_detection.py index 2b90dffac492..f5b428dd020d 100644 --- a/tools/recipes/hardware_detection.py +++ b/tools/recipes/hardware_detection.py @@ -4,7 +4,7 @@ """CPU hardware discovery for optional vLLM Recipes runtime tuning. Keep this module focused on facts exposed to the current process/container. -Tuning decisions belong in runtime_tuning.py. +Tuning decisions belong in sweep/runtime_tuning.py. """ from __future__ import annotations @@ -46,6 +46,89 @@ def to_dict(self) -> dict[str, object]: return asdict(self) +def _compress_cpu_ids(cpu_ids: list[int] | tuple[int, ...]) -> str: + """Return a compact Linux CPU-list string while preserving CPU IDs.""" + values = sorted(set(cpu_ids)) + if not values: + return "" + + ranges: list[str] = [] + start = previous = values[0] + for cpu_id in values[1:]: + if cpu_id == previous + 1: + previous = cpu_id + continue + ranges.append(str(start) if start == previous else f"{start}-{previous}") + start = previous = cpu_id + ranges.append(str(start) if start == previous else f"{start}-{previous}") + return ",".join(ranges) + + +def build_numa_omp_threads_bind( + hardware: HardwareInfo, + *, + reserved_cores_per_numa: int = 1, +) -> str: + """Build explicit x86 OMP CPU lists for the effective NUMA nodes. + + Temporary TP/DP sweep workaround: + - select one logical CPU per physical core, matching vLLM x86 auto-binding; + - reserve cores by omitting them from each explicit NUMA CPU list; + - emit one pipe-separated list per effective NUMA node. + + Explicit VLLM_CPU_OMP_THREADS_BIND lists bypass vLLM's auto-binding reserve + logic, so the reserved cores must be excluded here. + """ + if reserved_cores_per_numa < 0: + raise ValueError("reserved_cores_per_numa must be at least 0") + + effective_nodes = {node.node_id for node in hardware.numa_nodes} + allowed_cpus = get_allowed_cpu_list() + + groups: list[str] = [] + for node in hardware.numa_nodes: + node_cpus = [ + cpu + for cpu in allowed_cpus + if cpu.numa_node == node.node_id and cpu.numa_node in effective_nodes + ] + + core_to_logical_ids: dict[int, list[int]] = {} + for cpu in node_cpus: + if cpu.physical_core < 0: + continue + core_to_logical_ids.setdefault(cpu.physical_core, []).append(cpu.id) + + if not core_to_logical_ids: + raise RuntimeError( + f"NUMA node {node.node_id} has no physical-core topology " + "available for explicit OMP binding." + ) + + # Match vLLM's x86 auto-binding selector: for SMT siblings on a physical + # core, use the highest logical CPU ID. + selected_cpu_ids = sorted( + max(logical_ids) for _, logical_ids in sorted(core_to_logical_ids.items()) + ) + + if len(selected_cpu_ids) <= reserved_cores_per_numa: + raise RuntimeError( + f"NUMA node {node.node_id} has only {len(selected_cpu_ids)} " + "usable physical cores, which is not enough after reserving " + f"{reserved_cores_per_numa} core(s)." + ) + + if reserved_cores_per_numa: + selected_cpu_ids = selected_cpu_ids[:-reserved_cores_per_numa] + + groups.append(_compress_cpu_ids(selected_cpu_ids)) + + if not groups: + raise RuntimeError("No effective NUMA CPU lists are available for OMP binding.") + + return "|".join(groups) + + def _read_socket_id(cpu_id: int) -> int | None: path = Path(f"/sys/devices/system/cpu/cpu{cpu_id}/topology/physical_package_id") try: diff --git a/tools/recipes/recipe_json_to_vllm_config.py b/tools/recipes/recipe_json_to_vllm_config.py index bd947b07451c..c5089fbe9bbf 100644 --- a/tools/recipes/recipe_json_to_vllm_config.py +++ b/tools/recipes/recipe_json_to_vllm_config.py @@ -189,17 +189,67 @@ def parse_args() -> argparse.Namespace: type=float, help="Optional capacity target for future DP/capacity tuning.", ) + tuning.add_argument( + "--tune-scheduler", + action="store_true", + help=( + "Opt in to workload-derived max-num-seqs and " + "max-num-batched-tokens. By default scheduler parameters remain " + "at recipe/vLLM defaults." + ), + ) sweep = p.add_argument_group( "optional performance sweep", - ("Generate benchmark files after creating one initial runtime suggestion."), + ("Generate benchmark files alongside the directly deployable runtime config."), + ) + sweep.add_argument( + "--sweep-stage", + choices=("parallel-layout", "concurrency", "scheduler", "all"), + help="Generate only the selected stage, or all stages in dependency order.", ) sweep.add_argument( "--generate-sweep", action="store_true", + help="Backward-compatible alias for --generate-scheduler-sweep.", + ) + sweep.add_argument( + "--generate-scheduler-sweep", + action="store_true", + help=("Generate the max-num-seqs/max-num-batched-tokens scheduler sweep."), + ) + sweep.add_argument( + "--generate-parallel-layout-sweep", + action="store_true", + help=( + "Generate a standalone NUMA-aware TP/DP sweep. Requires --detect-hardware." + ), + ) + sweep.add_argument( + "--generate-concurrency-sweep", + action="store_true", + help=( + "Generate a max_concurrency workload sweep using " + "vllm bench sweep serve_workload." + ), + ) + sweep.add_argument( + "--generate-full-sweep", + action="store_true", + help=( + "Generate the end-to-end TP/DP -> max_concurrency -> scheduler " + "tuning pipeline. Requires --detect-hardware." + ), + ) + sweep.add_argument( + "--tp-dp-numa-bind-workaround", + action=argparse.BooleanOptionalAction, + default=True, help=( - "Generate an optional vllm bench sweep package around the single " - "initial runtime suggestion." + "Temporary Xeon TP/DP NUMA-binding workaround. For parallel-layout " + "and full sweeps, generate an explicit VLLM_CPU_OMP_THREADS_BIND " + "from detected NUMA topology. Enabled by default; disable with " + "--no-tp-dp-numa-bind-workaround after the vLLM DP binding fix." ), ) sweep.add_argument( @@ -766,10 +816,22 @@ def write_config( Path(path).write_text("\n".join(metadata) + "\n" + body, encoding="utf-8") -def write_env(path: str, source: str, recipe: dict[str, Any]) -> None: - env = recipe.get("env") or {} - if not isinstance(env, dict): - raise ValueError(f"Recipe `env` must be an object, got {type(env).__name__}") +def write_env( + path: str, + source: str, + recipe: dict[str, Any], + *, + env_overrides: dict[str, object] | None = None, +) -> None: + recipe_env = recipe.get("env") or {} + if not isinstance(recipe_env, dict): + raise ValueError( + f"Recipe `env` must be an object, got {type(recipe_env).__name__}" + ) + + env = dict(recipe_env) + overrides = env_overrides or {} + env.update(overrides) lines = [ "#!/usr/bin/env bash", @@ -778,6 +840,9 @@ def write_env(path: str, source: str, recipe: dict[str, Any]) -> None: "", ] + if overrides: + lines.append("# Temporary runtime overrides derived from detected hardware.") + if env: for key, value in env.items(): lines.append(f"export {key}={shlex.quote(str(value))}") @@ -793,6 +858,34 @@ def main() -> int: args = parse_args() try: + sweep_modes = { + "--generate-sweep": args.generate_sweep, + "--generate-scheduler-sweep": args.generate_scheduler_sweep, + "--generate-parallel-layout-sweep": args.generate_parallel_layout_sweep, + "--generate-concurrency-sweep": args.generate_concurrency_sweep, + "--generate-full-sweep": args.generate_full_sweep, + } + selected_sweep_modes = [ + name for name, enabled in sweep_modes.items() if enabled + ] + if len(selected_sweep_modes) > 1: + raise ValueError( + "Choose exactly one sweep-generation mode: " + + ", ".join(selected_sweep_modes) + ) + if args.sweep_stage is not None: + if selected_sweep_modes: + raise ValueError( + "Do not combine --sweep-stage with --generate-*-sweep options." + ) + stage_flags = { + "parallel-layout": "generate_parallel_layout_sweep", + "concurrency": "generate_concurrency_sweep", + "scheduler": "generate_scheduler_sweep", + "all": "generate_full_sweep", + } + setattr(args, stage_flags[args.sweep_stage], True) + selected_sweep_modes = ["--sweep-stage=" + args.sweep_stage] source = args.source if source is None: source = discover_recipe_source( @@ -813,7 +906,8 @@ def main() -> int: tuning_requested = ( args.detect_hardware - or args.generate_sweep + or args.tune_scheduler + or bool(selected_sweep_modes) or any( value is not None for value in ( @@ -828,12 +922,15 @@ def main() -> int: ) tuning = None + sweep_config = None workload = None sweep_writer = None + hardware = None + env_overrides: dict[str, object] = {} if tuning_requested: # Keep plain Recipes conversion lightweight. vLLM-specific modules # are imported only for optional runtime tuning or sweep generation. - from runtime_tuning import ( + from sweep.runtime_tuning import ( WorkloadHints, finetune_runtime_config, get_runtime_tuning_policies, @@ -848,24 +945,60 @@ def main() -> int: target_qps=args.target_qps, ) - if args.generate_sweep: - from sweep_generation import ( - validate_sweep_workload, - write_sweep_files, - ) + if selected_sweep_modes: + from sweep.sweep_generation import validate_sweep_workload validate_sweep_workload(workload) + + if args.generate_sweep or args.generate_scheduler_sweep: + from sweep.sweep_generation import write_sweep_files + sweep_writer = write_sweep_files recipe_hardware = recipe.get("hardware") - policies = get_runtime_tuning_policies(recipe_hardware) + scheduler_tuning_requested = bool( + args.tune_scheduler + or args.generate_sweep + or args.generate_scheduler_sweep + ) + policies = get_runtime_tuning_policies( + recipe_hardware, + tune_scheduler=scheduler_tuning_requested, + ) - hardware = None if args.detect_hardware: - from hardware_detection import detect_hardware + from hardware_detection import ( + build_numa_omp_threads_bind, + detect_hardware, + ) hardware = detect_hardware() + if ( + args.tp_dp_numa_bind_workaround + and str(recipe_hardware).lower() == "xeon6" + and ( + args.generate_parallel_layout_sweep or args.generate_full_sweep + ) + ): + recipe_env = recipe.get("env") or {} + current_binding = ( + recipe_env.get("VLLM_CPU_OMP_THREADS_BIND") + if isinstance(recipe_env, dict) + else None + ) + + # Preserve an explicit recipe-provided manual binding. + # Replace an unset or "auto" value with the temporary + # topology-derived binding. + if current_binding in (None, "auto"): + env_overrides["VLLM_CPU_OMP_THREADS_BIND"] = ( + build_numa_omp_threads_bind( + hardware, + reserved_cores_per_numa=1, + ) + ) + tuning = finetune_runtime_config( config, hardware=hardware, @@ -873,19 +1006,73 @@ def main() -> int: policies=policies, ) config.update(tuning.overrides) + if args.tune_scheduler: + config.update(tuning.sweep_overrides) + sweep_config = dict(config) + sweep_config.update(tuning.sweep_overrides) write_config(args.config_out, source, recipe, config) - write_env(args.env_out, source, recipe) + write_env( + args.env_out, + source, + recipe, + env_overrides=env_overrides, + ) sweep_files: list[Path] = [] - if args.generate_sweep: + if args.generate_full_sweep: + if not args.detect_hardware or hardware is None: + raise ValueError("--generate-full-sweep requires --detect-hardware.") + assert workload is not None + from sweep.sweep_generation import write_full_sweep_files + + sweep_files = write_full_sweep_files( + args.sweep_out_dir, + config_path=args.config_out, + env_path=args.env_out, + config=config, + workload=workload, + numa_node_count=hardware.numa_node_count, + tune_scheduler=args.tune_scheduler, + ) + elif args.generate_parallel_layout_sweep: + if not args.detect_hardware or hardware is None: + raise ValueError( + "--generate-parallel-layout-sweep requires --detect-hardware." + ) + assert workload is not None + from sweep.sweep_generation import write_parallel_layout_sweep_files + + sweep_files = write_parallel_layout_sweep_files( + args.sweep_out_dir, + config_path=args.config_out, + env_path=args.env_out, + config=config, + workload=workload, + numa_node_count=hardware.numa_node_count, + tune_scheduler=args.tune_scheduler, + ) + elif args.generate_concurrency_sweep: + assert workload is not None + from sweep.sweep_generation import write_concurrency_sweep_files + + sweep_files = write_concurrency_sweep_files( + args.sweep_out_dir, + config_path=args.config_out, + env_path=args.env_out, + config=config, + workload=workload, + tune_scheduler=args.tune_scheduler, + ) + elif args.generate_sweep or args.generate_scheduler_sweep: assert workload is not None assert sweep_writer is not None + assert sweep_config is not None sweep_files = sweep_writer( args.sweep_out_dir, config_path=args.config_out, env_path=args.env_out, - config=config, + config=sweep_config, workload=workload, ) @@ -894,6 +1081,13 @@ def main() -> int: print("Initial runtime suggestion:") for key, value in tuning.overrides.items(): print(f" {key}: {value}") + if tuning.sweep_overrides: + if args.tune_scheduler: + print("Workload-derived scheduler tuning:") + elif selected_sweep_modes: + print("Explicit scheduler sweep seed:") + for key, value in tuning.sweep_overrides.items(): + print(f" {key}: {value}") for note in tuning.notes: print(f" tuning: {note}") except Exception as exc: @@ -910,15 +1104,31 @@ def main() -> int: print(f" vllm serve --config {shlex.quote(args.config_out)}") if sweep_files: sweep_dir = Path(args.sweep_out_dir) + run_parallel = sweep_dir / "run_parallel_layout_sweep.sh" + recommend_parallel = sweep_dir / "recommend_parallel_layout.py" + run_concurrency = sweep_dir / "run_concurrency_sweep.sh" + recommend_concurrency = sweep_dir / "recommend_concurrency.py" run_sweep = sweep_dir / "run_sweep.sh" recommend = sweep_dir / "recommend.py" + run_full = sweep_dir / "run_full_sweep.sh" print() print("Optional performance sweep:") - print(f" {shlex.quote(str(run_sweep))} --dry-run") - print(f" {shlex.quote(str(run_sweep))}") - print() - print("After the sweep:") - print(f" {shlex.quote(str(recommend))}") + if args.generate_full_sweep: + print(f" {shlex.quote(str(run_full))}") + elif args.generate_concurrency_sweep: + print(f" {shlex.quote(str(run_concurrency))} --dry-run") + print(f" {shlex.quote(str(run_concurrency))}") + print(f" {shlex.quote(str(recommend_concurrency))}") + elif args.generate_parallel_layout_sweep: + print(f" {shlex.quote(str(run_parallel))} --dry-run") + print(f" {shlex.quote(str(run_parallel))}") + print(f" {shlex.quote(str(recommend_parallel))}") + else: + print(f" {shlex.quote(str(run_sweep))} --dry-run") + print(f" {shlex.quote(str(run_sweep))}") + print() + print("After the sweep:") + print(f" {shlex.quote(str(recommend))}") return 0 diff --git a/tools/recipes/requirements.txt b/tools/recipes/requirements.txt new file mode 100644 index 000000000000..36a96c66b4e5 --- /dev/null +++ b/tools/recipes/requirements.txt @@ -0,0 +1,4 @@ +# Dependencies used by `vllm bench sweep plot`. +matplotlib +pandas +seaborn diff --git a/tools/recipes/sweep/__init__.py b/tools/recipes/sweep/__init__.py new file mode 100644 index 000000000000..c6a8d5b25b68 --- /dev/null +++ b/tools/recipes/sweep/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Runtime tuning and benchmark sweep generation helpers.""" diff --git a/tools/recipes/sweep/adaptive_concurrency.py b/tools/recipes/sweep/adaptive_concurrency.py new file mode 100755 index 000000000000..a0fdf36a24a5 --- /dev/null +++ b/tools/recipes/sweep/adaptive_concurrency.py @@ -0,0 +1,261 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Find the maximum SLA-feasible serving concurrency adaptively.""" + +from __future__ import annotations + +import argparse +import shlex +from pathlib import Path +from statistics import median +from typing import Any + +from vllm.benchmarks.sweep.param_sweep import ParameterSweepItem +from vllm.benchmarks.sweep.serve import run_server +from vllm.benchmarks.sweep.serve_workload import run_comb_workload + +DEFAULT_MINIMUM_COMPLIANCE = 0.99 +DEFAULT_MAX_CONCURRENCY_CAP = 1000 +DEFAULT_GROWTH_FACTOR = 2.0 + + +def _number(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return float(value) + + +def _num_prompts(concurrency: int) -> int: + return min(1000, max(100, concurrency * 10)) + + +def _sla_eligible( + rows: list[dict[str, Any]], + *, + expected_runs: int, + ttft_sla_ms: float | None, + tpot_sla_ms: float | None, + minimum_compliance: float, +) -> bool: + if len(rows) < expected_runs: + return False + if any(int(_number(row.get("failed")) or 0) for row in rows): + return False + + completed = sum(int(_number(row.get("completed")) or 0) for row in rows) + if completed <= 0: + return False + + compliant = sum( + (_number(row.get("request_goodput")) or 0.0) + * (_number(row.get("duration")) or 0.0) + for row in rows + ) + if compliant / completed < minimum_compliance: + return False + + if ttft_sla_ms is not None: + values = [ + value + for row in rows + if (value := _number(row.get("p99_ttft_ms"))) is not None + ] + if not values or median(values) > ttft_sla_ms: + return False + + if tpot_sla_ms is not None: + values = [ + value + for row in rows + if (value := _number(row.get("p99_tpot_ms"))) is not None + ] + if not values or median(values) > tpot_sla_ms: + return False + + return True + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Find the highest SLA-feasible max_concurrency using exponential " + "bracketing followed by integer binary search." + ) + ) + parser.add_argument("--serve-cmd", required=True) + parser.add_argument("--bench-cmd", required=True) + parser.add_argument("--results-dir", required=True) + parser.add_argument("--seed-concurrency", type=int, required=True) + parser.add_argument("--input-tokens", type=int, required=True) + parser.add_argument("--output-tokens", type=int, required=True) + parser.add_argument("--ttft-sla-ms", type=float) + parser.add_argument("--tpot-sla-ms", type=float) + parser.add_argument( + "--minimum-compliance", + type=float, + default=DEFAULT_MINIMUM_COMPLIANCE, + ) + parser.add_argument( + "--max-concurrency-cap", + type=int, + default=DEFAULT_MAX_CONCURRENCY_CAP, + ) + parser.add_argument( + "--growth-factor", + type=float, + default=DEFAULT_GROWTH_FACTOR, + ) + parser.add_argument("--num-runs", type=int, default=3) + parser.add_argument("--server-ready-timeout", type=int, default=1800) + parser.add_argument("--show-stdout", action="store_true") + parser.add_argument("--continue-on-error", action="store_true") + parser.add_argument("--resume", action="store_true") + parser.add_argument("--dry-run", action="store_true") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + + if args.seed_concurrency <= 0: + raise ValueError("--seed-concurrency must be positive.") + if not 0.0 < args.minimum_compliance <= 1.0: + raise ValueError("--minimum-compliance must be greater than 0 and at most 1.") + if args.growth_factor <= 1.0: + raise ValueError("--growth-factor must be greater than 1.") + if args.max_concurrency_cap <= 0: + raise ValueError("--max-concurrency-cap must be positive.") + if args.num_runs <= 0: + raise ValueError("--num-runs must be positive.") + if args.ttft_sla_ms is None and args.tpot_sla_ms is None: + raise ValueError("Adaptive search requires TTFT and/or TPOT SLA.") + + seed = args.seed_concurrency + cap = max(seed, args.max_concurrency_cap) + results_dir = Path(args.results_dir) + results_dir.mkdir(parents=True, exist_ok=True) + + if args.dry_run: + first_upper = min( + cap, + max(seed + 1, int(round(seed * args.growth_factor))), + ) + print("Adaptive concurrency search plan:") + print(f" seed: {seed}") + print(f" first upper probe: {first_upper}") + print(f" hard cap: {cap}") + print(" strategy: bracket SLA boundary, then integer binary search") + return 0 + + serve_cmd = shlex.split(args.serve_cmd) + bench_cmd = shlex.split(args.bench_cmd) + serve_comb = ParameterSweepItem() + + def run_point(server: object, concurrency: int) -> bool: + bench_comb = ParameterSweepItem( + { + "_benchmark_name": "user_workload", + "random_input_len": args.input_tokens, + "random_output_len": args.output_tokens, + "num_prompts": _num_prompts(concurrency), + } + ) + print(f"[ADAPTIVE CONCURRENCY] testing max_concurrency={concurrency}") + rows = ( + run_comb_workload( + server, + bench_cmd, + serve_comb=serve_comb, + bench_comb=bench_comb, + link_vars=[], + experiment_dir=results_dir, + num_runs=args.num_runs, + dry_run=False, + warmup_num_prompts=min(concurrency, 1000), + continue_on_error=args.continue_on_error, + workload_var="max_concurrency", + workload_value=concurrency, + ) + or [] + ) + + eligible = _sla_eligible( + rows, + expected_runs=args.num_runs, + ttft_sla_ms=args.ttft_sla_ms, + tpot_sla_ms=args.tpot_sla_ms, + minimum_compliance=args.minimum_compliance, + ) + print( + "[ADAPTIVE CONCURRENCY] " + f"max_concurrency={concurrency}: {'PASS' if eligible else 'FAIL'}" + ) + return eligible + + with run_server( + serve_cmd, + [], + show_stdout=args.show_stdout, + serve_overrides=serve_comb, + dry_run=False, + server_ready_timeout=args.server_ready_timeout, + ) as server: + seed_passes = run_point(server, seed) + + if seed_passes: + low = seed + high: int | None = None + while low < cap: + probe = min( + cap, + max(low + 1, int(round(low * args.growth_factor))), + ) + if run_point(server, probe): + low = probe + if low == cap: + print( + "[ADAPTIVE CONCURRENCY] reached cap without an SLA failure" + ) + return 0 + else: + high = probe + break + if high is None: + return 0 + else: + high = seed + low: int | None = None + probe = seed + while probe > 1: + probe = max(1, probe // 2) + if run_point(server, probe): + low = probe + break + high = probe + if low is None: + print( + "[ADAPTIVE CONCURRENCY] no SLA-feasible concurrency found, " + "including max_concurrency=1" + ) + return 0 + + assert low is not None + assert high is not None + print(f"[ADAPTIVE CONCURRENCY] refining bracket PASS={low}, FAIL={high}") + + while high - low > 1: + midpoint = (low + high) // 2 + if run_point(server, midpoint): + low = midpoint + else: + high = midpoint + + print(f"[ADAPTIVE CONCURRENCY] boundary resolved: PASS={low}, FAIL={high}") + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/recipes/sweep/concurrency_recommendation.py b/tools/recipes/sweep/concurrency_recommendation.py new file mode 100644 index 000000000000..908f8ce1da8b --- /dev/null +++ b/tools/recipes/sweep/concurrency_recommendation.py @@ -0,0 +1,330 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Select one max-concurrency recommendation from a workload sweep.""" + +from __future__ import annotations + +import argparse +import json +from collections import defaultdict +from pathlib import Path +from statistics import mean, median +from typing import Any + +DEFAULT_TTFT_SLA_MS: float | None = None +DEFAULT_TPOT_SLA_MS: float | None = None +DEFAULT_SEED_CONCURRENCY: int | None = None +DEFAULT_MINIMUM_COMPLIANCE: float = 0.99 +DEFAULT_RESULTS_DIR = "results/concurrency-tuning" +DEFAULT_OUTPUT_JSON = "concurrency-recommendation.json" + + +def _number(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return float(value) + + +def _mean(rows: list[dict[str, Any]], key: str) -> float | None: + values = [_number(row.get(key)) for row in rows] + numbers = [value for value in values if value is not None] + return mean(numbers) if numbers else None + + +def _percentile_summary( + rows: list[dict[str, Any]], key: str +) -> tuple[float | None, float | None, float | None]: + values = [_number(row.get(key)) for row in rows] + numbers = [value for value in values if value is not None] + if not numbers: + return None, None, None + return mean(numbers), median(numbers), max(numbers) + + +def _load_runs(results_dir: Path) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for path in sorted(results_dir.rglob("summary.json")): + data = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(data, list): + continue + for row in data: + if isinstance(row, dict): + copied = dict(row) + copied["_summary_path"] = str(path) + rows.append(copied) + return rows + + +def _aggregate_candidates( + rows: list[dict[str, Any]], + *, + use_goodput: bool, + ttft_sla_ms: float | None, + tpot_sla_ms: float | None, + minimum_compliance: float, +) -> list[dict[str, Any]]: + grouped: dict[int, list[dict[str, Any]]] = defaultdict(list) + for row in rows: + concurrency = row.get("max_concurrency") + if ( + isinstance(concurrency, bool) + or not isinstance(concurrency, int) + or concurrency <= 0 + ): + continue + grouped[concurrency].append(row) + + candidates: list[dict[str, Any]] = [] + for concurrency, runs in sorted(grouped.items()): + failed_requests = sum(int(_number(run.get("failed")) or 0) for run in runs) + output_throughput = _mean(runs, "output_throughput") + request_throughput = _mean(runs, "request_throughput") + request_goodput = _mean(runs, "request_goodput") + mean_p99_ttft_ms, median_p99_ttft_ms, worst_p99_ttft_ms = _percentile_summary( + runs, "p99_ttft_ms" + ) + mean_p99_tpot_ms, median_p99_tpot_ms, worst_p99_tpot_ms = _percentile_summary( + runs, "p99_tpot_ms" + ) + + completed_requests = sum( + int(_number(run.get("completed")) or 0) for run in runs + ) + estimated_compliant_requests = sum( + (_number(run.get("request_goodput")) or 0.0) + * (_number(run.get("duration")) or 0.0) + for run in runs + ) + combined_compliance_ratio = ( + estimated_compliant_requests / completed_requests + if completed_requests > 0 + else None + ) + + valid = failed_requests == 0 and output_throughput is not None + reason = None + if failed_requests: + valid = False + reason = f"{failed_requests} failed request(s) across successful runs" + elif use_goodput and request_goodput is None: + valid = False + reason = "request_goodput is missing; rerun with SLO goodput enabled" + + p99_sla_eligible = valid + if ttft_sla_ms is not None: + p99_sla_eligible = bool( + p99_sla_eligible + and median_p99_ttft_ms is not None + and median_p99_ttft_ms <= ttft_sla_ms + ) + if tpot_sla_ms is not None: + p99_sla_eligible = bool( + p99_sla_eligible + and median_p99_tpot_ms is not None + and median_p99_tpot_ms <= tpot_sla_ms + ) + compliance_eligible = bool( + valid + and combined_compliance_ratio is not None + and combined_compliance_ratio >= minimum_compliance + ) + + candidates.append( + { + "max_concurrency": concurrency, + "run_count": len(runs), + "failed_requests": failed_requests, + "mean_request_goodput": request_goodput, + "mean_request_throughput": request_throughput, + "mean_output_throughput": output_throughput, + "mean_p99_ttft_ms": mean_p99_ttft_ms, + "median_p99_ttft_ms": median_p99_ttft_ms, + "worst_p99_ttft_ms": worst_p99_ttft_ms, + "mean_p99_tpot_ms": mean_p99_tpot_ms, + "median_p99_tpot_ms": median_p99_tpot_ms, + "worst_p99_tpot_ms": worst_p99_tpot_ms, + "estimated_compliant_requests": estimated_compliant_requests, + "completed_requests": completed_requests, + "combined_compliance_ratio": combined_compliance_ratio, + "combined_compliance_percent": ( + combined_compliance_ratio * 100 + if combined_compliance_ratio is not None + else None + ), + "benchmark_valid": valid, + "p99_sla_eligible": p99_sla_eligible, + "compliance_eligible": compliance_eligible, + "sla_eligible": p99_sla_eligible and compliance_eligible, + "valid": valid, + "invalid_reason": reason, + "summary_files": sorted({str(run["_summary_path"]) for run in runs}), + } + ) + + return candidates + + +def _select_candidate( + candidates: list[dict[str, Any]], *, use_goodput: bool +) -> tuple[dict[str, Any] | None, dict[str, Any], str]: + valid = [candidate for candidate in candidates if candidate["valid"]] + if not valid: + raise ValueError( + "No valid workload-sweep point was found. Inspect failed runs and metrics." + ) + + if use_goodput: + objective = "highest_sla_feasible_max_concurrency" + eligible = [candidate for candidate in valid if candidate["sla_eligible"]] + winner = ( + max( + eligible, + key=lambda candidate: ( + candidate["max_concurrency"], + candidate["mean_output_throughput"], + candidate["combined_compliance_ratio"] or 0.0, + candidate["mean_request_goodput"] or 0.0, + ), + ) + if eligible + else None + ) + best_effort = max( + valid, + key=lambda candidate: ( + candidate["mean_request_goodput"] or 0.0, + candidate["combined_compliance_ratio"] or 0.0, + candidate["mean_output_throughput"], + candidate["max_concurrency"], + ), + ) + else: + objective = "highest_mean_output_throughput" + winner = max( + valid, + key=lambda candidate: ( + candidate["mean_output_throughput"], + candidate["max_concurrency"], + ), + ) + best_effort = winner + + return winner, best_effort, objective + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Select max_concurrency from vLLM Workload Explorer results." + ) + parser.add_argument("--results-dir", default=DEFAULT_RESULTS_DIR) + parser.add_argument("--ttft-sla-ms", type=float, default=DEFAULT_TTFT_SLA_MS) + parser.add_argument("--tpot-sla-ms", type=float, default=DEFAULT_TPOT_SLA_MS) + parser.add_argument( + "--seed-concurrency", type=int, default=DEFAULT_SEED_CONCURRENCY + ) + parser.add_argument( + "--minimum-compliance", + type=float, + default=DEFAULT_MINIMUM_COMPLIANCE, + ) + parser.add_argument("--output-json", default=DEFAULT_OUTPUT_JSON) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + if not 0.0 < args.minimum_compliance <= 1.0: + raise ValueError("--minimum-compliance must be greater than 0 and at most 1.") + + script_dir = Path(__file__).resolve().parent + results_dir = Path(args.results_dir) + if not results_dir.is_absolute(): + results_dir = script_dir / results_dir + output_json = Path(args.output_json) + if not output_json.is_absolute(): + output_json = script_dir / output_json + + if not results_dir.exists(): + raise ValueError(f"Workload-sweep results were not found at {results_dir}.") + + rows = _load_runs(results_dir) + if not rows: + raise ValueError(f"No summary.json results found under {results_dir}.") + + use_goodput = args.ttft_sla_ms is not None or args.tpot_sla_ms is not None + candidates = _aggregate_candidates( + rows, + use_goodput=use_goodput, + ttft_sla_ms=args.ttft_sla_ms, + tpot_sla_ms=args.tpot_sla_ms, + minimum_compliance=args.minimum_compliance, + ) + winner, best_effort, objective = _select_candidate( + candidates, use_goodput=use_goodput + ) + + recommendation = { + "status": ( + "sla_feasible" if winner is not None else "no_sla_feasible_concurrency" + ), + "selection_objective": objective, + "seed_concurrency": args.seed_concurrency, + "minimum_compliance_ratio": args.minimum_compliance, + "slo": { + "ttft_ms": args.ttft_sla_ms, + "tpot_ms": args.tpot_sla_ms, + }, + "recommended": ( + {"max_concurrency": winner["max_concurrency"]} + if winner is not None + else None + ), + "measured": ( + { + key: winner[key] + for key in ( + "run_count", + "mean_request_throughput", + "mean_request_goodput", + "mean_output_throughput", + "combined_compliance_ratio", + "combined_compliance_percent", + "mean_p99_ttft_ms", + "median_p99_ttft_ms", + "worst_p99_ttft_ms", + "mean_p99_tpot_ms", + "median_p99_tpot_ms", + "worst_p99_tpot_ms", + ) + } + if winner is not None + else None + ), + "best_effort": { + "max_concurrency": best_effort["max_concurrency"], + "mean_request_goodput": best_effort["mean_request_goodput"], + "mean_output_throughput": best_effort["mean_output_throughput"], + "combined_compliance_ratio": best_effort["combined_compliance_ratio"], + "p99_sla_eligible": best_effort["p99_sla_eligible"], + }, + "candidates": candidates, + } + output_json.write_text( + json.dumps(recommendation, indent=2) + "\n", encoding="utf-8" + ) + + if winner is None: + print("No SLA-feasible max_concurrency was found.") + print(f"Wrote diagnostic evidence to {output_json}") + return 2 + + print(f"Recommended max_concurrency: {winner['max_concurrency']}") + print(f"Selection objective: {objective}") + print(f"Wrote {output_json}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/recipes/runtime_tuning.py b/tools/recipes/sweep/runtime_tuning.py similarity index 81% rename from tools/recipes/runtime_tuning.py rename to tools/recipes/sweep/runtime_tuning.py index b6b2d298d24d..f23e6f76ad6b 100644 --- a/tools/recipes/runtime_tuning.py +++ b/tools/recipes/sweep/runtime_tuning.py @@ -47,7 +47,10 @@ def validate(self) -> None: @dataclass class TuningResult: + """Runtime overrides and benchmark-only seeds produced by tuning policies.""" + overrides: dict[str, Any] = field(default_factory=dict) + sweep_overrides: dict[str, Any] = field(default_factory=dict) notes: list[str] = field(default_factory=list) @@ -67,13 +70,24 @@ def _record_override( result.notes.append(f"{key}={value!r}: {reason}") +def _record_sweep_override( + result: TuningResult, + key: str, + value: Any, + reason: str, +) -> None: + """Record an explicit benchmark seed without changing deployment config.""" + result.sweep_overrides[key] = value + result.notes.append(f"sweep {key}={value!r}: {reason}") + + def _get_active_sequence_count( config: dict[str, Any], workload: WorkloadHints, ) -> int: - """Return the sequence count used to size one scheduler iteration.""" + """Return the per-replica sequence count for one scheduler iteration.""" if workload.concurrency is not None: - return workload.concurrency + return ceil(workload.concurrency / _get_data_parallel_size(config)) configured = config.get("max-num-seqs") if ( @@ -86,9 +100,22 @@ def _get_active_sequence_count( return SchedulerConfig.DEFAULT_MAX_NUM_SEQS +def _get_data_parallel_size(config: dict[str, Any]) -> int: + """Return a valid configured DP size, falling back to one replica.""" + configured = config.get("data-parallel-size", 1) + if ( + isinstance(configured, int) + and not isinstance(configured, bool) + and configured > 0 + ): + return configured + return 1 + + def _estimate_prefills_per_step( workload: WorkloadHints, active_sequences: int, + data_parallel_size: int, ) -> float: """Estimate steady-state prompt arrivals per decode scheduler step.""" # Reserve enough prefill budget for at least one representative prompt. @@ -105,9 +132,10 @@ def _estimate_prefills_per_step( # With target QPS + TPOT, approximate how many prompts arrive during one # decode scheduler-step interval. if workload.target_qps is not None and workload.tpot_sla_ms is not None: + per_replica_qps = workload.target_qps / data_parallel_size prefills_per_step = max( prefills_per_step, - workload.target_qps * workload.tpot_sla_ms / 1000.0, + per_replica_qps * workload.tpot_sla_ms / 1000.0, ) # A step cannot replace more requests than are active. @@ -213,15 +241,22 @@ def _resolve_max_num_seqs( workload: WorkloadHints, result: TuningResult, ) -> None: - del config, hardware + del hardware if workload.concurrency is None: return - _record_override( + data_parallel_size = _get_data_parallel_size(config) + per_replica_concurrency = ceil(workload.concurrency / data_parallel_size) + + _record_sweep_override( result, "max-num-seqs", - workload.concurrency, - "derived from the optional target concurrency", + per_replica_concurrency, + ( + "benchmark seed derived from the optional global target concurrency " + f"and data_parallel_size={data_parallel_size}; config.yml keeps the " + "recipe/vLLM scheduler default" + ), ) @@ -235,12 +270,17 @@ def _resolve_max_num_batched_tokens( if workload.input_tokens is None: return + data_parallel_size = _get_data_parallel_size(config) active_sequences = _get_active_sequence_count(config, workload) # Decode requests consume roughly one token per active sequence in one # scheduler iteration. Prefills share the remaining scheduler token budget. decode_budget = active_sequences - prefills_per_step = _estimate_prefills_per_step(workload, active_sequences) + prefills_per_step = _estimate_prefills_per_step( + workload, + active_sequences, + data_parallel_size, + ) prefill_budget = ceil(workload.input_tokens * prefills_per_step) # Use vLLM's scheduler defaults/constraints as floors rather than local @@ -263,16 +303,18 @@ def _resolve_max_num_batched_tokens( ): candidate = max(candidate, max_model_len) - _record_override( + _record_sweep_override( result, "max-num-batched-tokens", candidate, ( - "workload-derived scheduler budget " - f"(active_sequences={active_sequences}, " + "workload-derived benchmark seed " + f"(per_replica_active_sequences={active_sequences}, " + f"data_parallel_size={data_parallel_size}, " f"prefills_per_step={prefills_per_step:.2f}, " f"decode_budget={decode_budget}, " - f"prefill_budget={prefill_budget})" + f"prefill_budget={prefill_budget}); config.yml keeps the " + "recipe/vLLM scheduler default" ), ) @@ -296,9 +338,12 @@ def _resolve_data_parallel_size( DEFAULT_POLICIES: tuple[Policy, ...] = ( _resolve_tensor_parallel_size, _resolve_gpu_memory_utilization, + _resolve_data_parallel_size, +) + +SCHEDULER_TUNING_POLICIES: tuple[Policy, ...] = ( _resolve_max_num_seqs, _resolve_max_num_batched_tokens, - _resolve_data_parallel_size, ) # Select tuning behavior from the hardware requested by the Recipes rendering, @@ -312,6 +357,8 @@ def _resolve_data_parallel_size( def get_runtime_tuning_policies( recipe_hardware: object, + *, + tune_scheduler: bool = False, ) -> tuple[Policy, ...]: if not isinstance(recipe_hardware, str) or not recipe_hardware: raise ValueError( @@ -326,6 +373,8 @@ def get_runtime_tuning_policies( "Runtime tuning is not supported for recipe hardware " f"{recipe_hardware!r}. Currently supported: {supported}." ) + if tune_scheduler: + policies = (*policies, *SCHEDULER_TUNING_POLICIES) return policies @@ -336,7 +385,7 @@ def finetune_runtime_config( workload: WorkloadHints | None = None, policies: tuple[Policy, ...] = DEFAULT_POLICIES, ) -> TuningResult: - """Return overrides; never mutate the recipe-derived config in place.""" + """Return deployment overrides and sweep seeds without mutating config.""" hints = workload or WorkloadHints() hints.validate() diff --git a/tools/recipes/sweep/scheduler_sweep_prepare.py b/tools/recipes/sweep/scheduler_sweep_prepare.py new file mode 100644 index 000000000000..59709a0fc56d --- /dev/null +++ b/tools/recipes/sweep/scheduler_sweep_prepare.py @@ -0,0 +1,225 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Prepare scheduler sweep inputs after TP/DP and concurrency selection.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any + +try: + import yaml +except ImportError as exc: + raise SystemExit("PyYAML is required. Install it with: pip install pyyaml") from exc + +PROMPTS_PER_CONCURRENCY = 10 +MIN_NUM_PROMPTS = 100 +MAX_NUM_PROMPTS = 1000 + + +def _positive_int(config: dict[str, Any], key: str) -> int: + value = config.get(key) + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"Expected positive integer {key!r}, got {value!r}.") + return value + + +def _num_prompts_for_concurrency(concurrency: int) -> int: + return min( + MAX_NUM_PROMPTS, + max(MIN_NUM_PROMPTS, concurrency * PROMPTS_PER_CONCURRENCY), + ) + + +def _strict_lower_power_of_two(value: int) -> int: + if value <= 1: + return 1 + return 1 << ((value - 1).bit_length() - 1) + + +def _strict_upper_power_of_two(value: int) -> int: + return 1 << value.bit_length() + + +def _scheduler_baseline( + config: dict[str, Any], + *, + concurrency: int, + input_tokens: int, + output_tokens: int, + tpot_sla_ms: float | None, + target_qps: float | None, +) -> tuple[int, int]: + data_parallel_size = int(config.get("data-parallel-size", 1)) + if data_parallel_size <= 0: + raise ValueError("data-parallel-size must be positive.") + + per_replica_concurrency = math.ceil(concurrency / data_parallel_size) + prefills_per_step = max(1.0, per_replica_concurrency / max(output_tokens, 1)) + if target_qps is not None and tpot_sla_ms is not None: + per_replica_qps = target_qps / data_parallel_size + prefills_per_step = max( + prefills_per_step, + per_replica_qps * tpot_sla_ms / 1000.0, + ) + prefills_per_step = min(float(per_replica_concurrency), prefills_per_step) + + batch = max( + 2048, + per_replica_concurrency, + per_replica_concurrency + math.ceil(input_tokens * prefills_per_step), + ) + if config.get("enable-chunked-prefill") is False: + max_model_len = config.get("max-model-len") + if isinstance(max_model_len, int) and not isinstance(max_model_len, bool): + batch = max(batch, max_model_len) + + return per_replica_concurrency, batch + + +def _build_serve_params(config: dict[str, Any]) -> list[dict[str, Any]]: + initial_seqs = _positive_int(config, "max-num-seqs") + initial_batch = _positive_int(config, "max-num-batched-tokens") + + minimum_batch = initial_seqs + if config.get("enable-chunked-prefill") is False: + max_model_len = config.get("max-model-len") + if ( + isinstance(max_model_len, int) + and not isinstance(max_model_len, bool) + and max_model_len > 0 + ): + minimum_batch = max(minimum_batch, max_model_len) + + lower_batch = max(minimum_batch, _strict_lower_power_of_two(initial_batch)) + smaller_batch = max(minimum_batch, _strict_lower_power_of_two(lower_batch)) + higher_batch = max(minimum_batch, _strict_upper_power_of_two(initial_batch)) + + candidates: list[dict[str, Any]] = [] + seen: set[tuple[int | None, int | None]] = set() + + def add( + name: str, + max_num_seqs: int | None, + max_num_batched_tokens: int | None, + ) -> None: + if max_num_seqs is not None and max_num_batched_tokens is not None: + max_num_batched_tokens = max(max_num_batched_tokens, max_num_seqs) + signature = (max_num_seqs, max_num_batched_tokens) + if signature in seen: + return + seen.add(signature) + + candidate: dict[str, Any] = {"_benchmark_name": name} + if max_num_seqs is not None: + candidate["max_num_seqs"] = max_num_seqs + if max_num_batched_tokens is not None: + candidate["max_num_batched_tokens"] = max_num_batched_tokens + candidates.append(candidate) + + add("initial", initial_seqs, initial_batch) + add("smaller_batch_budget", initial_seqs, smaller_batch) + add("lower_batch_budget", initial_seqs, lower_batch) + add("higher_batch_budget", initial_seqs, higher_batch) + add("vllm_default_max_num_seqs", None, initial_batch) + add("vllm_default_max_num_batched_tokens", initial_seqs, None) + add("vllm_defaults", None, None) + return candidates + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--config", required=True) + parser.add_argument("--concurrency-recommendation", required=True) + parser.add_argument("--input-tokens", type=int, required=True) + parser.add_argument("--output-tokens", type=int, required=True) + parser.add_argument("--tpot-sla-ms", type=float) + parser.add_argument("--target-qps", type=float) + parser.add_argument("--output-config", required=True) + parser.add_argument("--serve-params", required=True) + parser.add_argument("--bench-params", required=True) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + config_path = Path(args.config) + recommendation_path = Path(args.concurrency_recommendation) + + config = yaml.safe_load(config_path.read_text(encoding="utf-8")) + if not isinstance(config, dict): + raise ValueError(f"{config_path} does not contain a YAML object.") + + recommendation = json.loads(recommendation_path.read_text(encoding="utf-8")) + selected = recommendation.get("recommended") + if not isinstance(selected, dict): + raise ValueError( + f"{recommendation_path} has no SLA-feasible concurrency recommendation." + ) + concurrency = selected.get("max_concurrency") + if ( + isinstance(concurrency, bool) + or not isinstance(concurrency, int) + or concurrency <= 0 + ): + raise ValueError("Recommended max_concurrency is missing or invalid.") + + seqs, batch = _scheduler_baseline( + config, + concurrency=concurrency, + input_tokens=args.input_tokens, + output_tokens=args.output_tokens, + tpot_sla_ms=args.tpot_sla_ms, + target_qps=args.target_qps, + ) + + scheduler_seed = dict(config) + scheduler_seed["max-num-seqs"] = seqs + scheduler_seed["max-num-batched-tokens"] = batch + + scheduler_config = dict(config) + scheduler_config.pop("max-num-seqs", None) + scheduler_config.pop("max-num-batched-tokens", None) + Path(args.output_config).write_text( + yaml.safe_dump( + scheduler_config, + sort_keys=False, + default_flow_style=False, + allow_unicode=True, + ), + encoding="utf-8", + ) + + Path(args.serve_params).write_text( + json.dumps(_build_serve_params(scheduler_seed), indent=2) + "\n", + encoding="utf-8", + ) + Path(args.bench_params).write_text( + json.dumps( + [ + { + "random_input_len": args.input_tokens, + "random_output_len": args.output_tokens, + "max_concurrency": concurrency, + "num_prompts": _num_prompts_for_concurrency(concurrency), + } + ], + indent=2, + ) + + "\n", + encoding="utf-8", + ) + + print(f"Prepared scheduler sweep for max_concurrency={concurrency}") + print(f"Scheduler seed max-num-seqs={seqs}") + print(f"Scheduler seed max-num-batched-tokens={batch}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/recipes/sweep/sweep_generation.py b/tools/recipes/sweep/sweep_generation.py new file mode 100644 index 000000000000..cd3a131f8951 --- /dev/null +++ b/tools/recipes/sweep/sweep_generation.py @@ -0,0 +1,1531 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Generate optional vLLM benchmark sweep files from one initial suggestion.""" + +from __future__ import annotations + +import json +import math +import os +import shlex +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import yaml + +from .runtime_tuning import WorkloadHints + +SUPPORTED_TENSOR_PARALLEL_SIZES = frozenset({1, 2, 4, 8}) + +MOE_EXPERT_COUNT_KEYS = ( + "num_experts", + "num_local_experts", + "num_sparse_experts", + "num_total_experts", + "n_experts", + "n_routed_experts", + "num_routed_experts", +) +MOE_INTERMEDIATE_SIZE_KEYS = ( + "moe_intermediate_size", + "expert_intermediate_size", + "moe_ffn_hidden_size", +) + +PROMPTS_PER_CONCURRENCY = 10 +MIN_NUM_PROMPTS = 100 +MAX_NUM_PROMPTS = 1000 + + +@dataclass(frozen=True) +class ModelParallelMetadata: + """Model dimensions needed to reject invalid MoE layouts before startup.""" + + is_moe: bool | None + num_experts: int | None = None + moe_intermediate_size: int | None = None + + +UNKNOWN_MODEL_PARALLEL_METADATA = ModelParallelMetadata(is_moe=None) + + +def _num_prompts_for_concurrency(concurrency: int) -> int: + return min( + MAX_NUM_PROMPTS, + max(MIN_NUM_PROMPTS, concurrency * PROMPTS_PER_CONCURRENCY), + ) + + +def _positive_int(config: dict[str, Any], key: str) -> int: + value = config.get(key) + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"Sweep generation requires a positive initial {key!r} value.") + return value + + +def _strict_lower_power_of_two(value: int) -> int: + if value <= 1: + return 1 + return 1 << ((value - 1).bit_length() - 1) + + +def _strict_upper_power_of_two(value: int) -> int: + return 1 << value.bit_length() + + +def validate_sweep_workload(workload: WorkloadHints) -> None: + required = { + "--input-tokens": workload.input_tokens, + "--output-tokens": workload.output_tokens, + "--concurrency": workload.concurrency, + } + missing = [name for name, value in required.items() if value is None] + if missing: + raise ValueError("--generate-sweep requires " + ", ".join(missing) + ".") + + +def build_serve_params(config: dict[str, Any]) -> list[dict[str, Any]]: + """Build a bounded batch-budget sweep at the sequence-count lower bound. + + Explicit candidates keep the initial per-replica sequence count. Three + reference points compare those tuned values with vLLM-resolved defaults. + """ + initial_seqs = _positive_int(config, "max-num-seqs") + initial_batch = _positive_int(config, "max-num-batched-tokens") + + minimum_batch = initial_seqs + if config.get("enable-chunked-prefill") is False: + max_model_len = config.get("max-model-len") + if ( + isinstance(max_model_len, int) + and not isinstance(max_model_len, bool) + and max_model_len > 0 + ): + minimum_batch = max(minimum_batch, max_model_len) + + lower_batch = max( + minimum_batch, + _strict_lower_power_of_two(initial_batch), + ) + smaller_batch = max( + minimum_batch, + _strict_lower_power_of_two(lower_batch), + ) + higher_batch = max( + minimum_batch, + _strict_upper_power_of_two(initial_batch), + ) + + candidates: list[dict[str, Any]] = [] + seen: set[tuple[int | None, int | None]] = set() + + def add( + name: str, + max_num_seqs: int | None, + max_num_batched_tokens: int | None, + ) -> None: + if max_num_seqs is not None and max_num_batched_tokens is not None: + max_num_batched_tokens = max(max_num_batched_tokens, max_num_seqs) + signature = (max_num_seqs, max_num_batched_tokens) + if signature in seen: + return + seen.add(signature) + + candidate: dict[str, Any] = {"_benchmark_name": name} + if max_num_seqs is not None: + candidate["max_num_seqs"] = max_num_seqs + if max_num_batched_tokens is not None: + candidate["max_num_batched_tokens"] = max_num_batched_tokens + candidates.append(candidate) + + # Keep the exact initial suggestion as the measured baseline. Additional + # values exist only in the optional sweep package. + add("initial", initial_seqs, initial_batch) + add("smaller_batch_budget", initial_seqs, smaller_batch) + add("lower_batch_budget", initial_seqs, lower_batch) + add("higher_batch_budget", initial_seqs, higher_batch) + # Default-reference candidates intentionally omit one or both scheduler + # keys. The sweep server starts from sweep_config.yml, where both keys are + # removed, so omission lets vLLM resolve its normal runtime default. + add("vllm_default_max_num_seqs", None, initial_batch) + add("vllm_default_max_num_batched_tokens", initial_seqs, None) + add("vllm_defaults", None, None) + + return candidates + + +def _scheduler_baseline_for_dp( + config: dict[str, Any], workload: WorkloadHints, data_parallel_size: int +) -> tuple[int, int]: + """Return per-replica scheduler settings for a DP layout.""" + assert workload.concurrency is not None + assert workload.input_tokens is not None + per_replica_concurrency = math.ceil(workload.concurrency / data_parallel_size) + output_tokens = workload.output_tokens or 1 + prefills_per_step = max(1.0, per_replica_concurrency / output_tokens) + if workload.target_qps is not None and workload.tpot_sla_ms is not None: + per_replica_qps = workload.target_qps / data_parallel_size + prefills_per_step = max( + prefills_per_step, + per_replica_qps * workload.tpot_sla_ms / 1000.0, + ) + prefills_per_step = min(float(per_replica_concurrency), prefills_per_step) + batch = max( + 2048, + per_replica_concurrency, + per_replica_concurrency + math.ceil(workload.input_tokens * prefills_per_step), + ) + if config.get("enable-chunked-prefill") is False: + max_model_len = config.get("max-model-len") + if isinstance(max_model_len, int) and not isinstance(max_model_len, bool): + batch = max(batch, max_model_len) + return per_replica_concurrency, batch + + +def _model_config_value(view: Any, key: str) -> Any: + if isinstance(view, dict): + return view.get(key) + return getattr(view, key, None) + + +def _config_views(hf_config: Any) -> list[Any]: + """Return common nested model config views without duplicates.""" + views = [hf_config] + for view in views: + for key in ( + "text_config", + "language_config", + "llm_config", + "ffn_config", + "moe_config", + ): + value = _model_config_value(view, key) + if value is not None and all(value is not item for item in views): + views.append(value) + return views + + +def _first_positive_model_int(views: list[Any], keys: tuple[str, ...]) -> int | None: + for view in views: + for key in keys: + value = _model_config_value(view, key) + if isinstance(value, int) and not isinstance(value, bool) and value > 0: + return value + return None + + +def _config_identifies_moe(views: list[Any]) -> bool: + for view in views: + model_type = _model_config_value(view, "model_type") + if isinstance(model_type, str) and "moe" in model_type.lower(): + return True + architectures = _model_config_value(view, "architectures") + if isinstance(architectures, list) and any( + isinstance(name, str) and "moe" in name.lower() for name in architectures + ): + return True + return False + + +def model_parallel_metadata_from_config(hf_config: Any) -> ModelParallelMetadata: + """Extract MoE markers from a loaded Transformers-compatible config.""" + views = _config_views(hf_config) + num_experts = _first_positive_model_int(views, MOE_EXPERT_COUNT_KEYS) + moe_intermediate_size = _first_positive_model_int(views, MOE_INTERMEDIATE_SIZE_KEYS) + is_moe = ( + num_experts is not None + or moe_intermediate_size is not None + or _config_identifies_moe(views) + ) + if is_moe and moe_intermediate_size is None: + # Some MoE configs use the generic name for each expert's FFN width. + moe_intermediate_size = _first_positive_model_int(views, ("intermediate_size",)) + return ModelParallelMetadata( + is_moe=is_moe, + num_experts=num_experts, + moe_intermediate_size=moe_intermediate_size, + ) + + +def inspect_model_parallel_metadata(config: dict[str, Any]) -> ModelParallelMetadata: + """Load the model config and extract portable MoE topology markers.""" + model = config.get("model") + if not isinstance(model, str) or not model: + raise ValueError("Parallel-layout sweep requires a model in config.yml.") + + from vllm.transformers_utils.config import get_config + + hf_config = get_config( + model, + trust_remote_code=config.get("trust-remote-code") is True, + revision=config.get("revision"), + code_revision=config.get("code-revision"), + config_format=config.get("config-format", "auto"), + ) + return model_parallel_metadata_from_config(hf_config) + + +def cpu_supports_moe_dp_collectives() -> bool: + """Whether CPU overrides the variable-size collectives used by MoE DP.""" + try: + from vllm.distributed.device_communicators.base_device_communicator import ( + DeviceCommunicatorBase, + ) + from vllm.distributed.device_communicators.cpu_communicator import ( + CpuCommunicator, + ) + except Exception: + return False + + return all( + getattr(CpuCommunicator, name) is not getattr(DeviceCommunicatorBase, name) + for name in ("all_gatherv", "reduce_scatterv") + ) + + +def _parallel_layout_name( + tensor_parallel_size: int, data_parallel_size: int, numa_node_count: int +) -> str: + name = f"tp{tensor_parallel_size}_dp{data_parallel_size}" + used_nodes = tensor_parallel_size * data_parallel_size + if used_nodes != numa_node_count: + name += f"_numa{used_nodes}of{numa_node_count}" + return name + + +def _default_parallel_layout_sizes(numa_node_count: int) -> list[tuple[int, int]]: + max_supported_tp = max( + size for size in SUPPORTED_TENSOR_PARALLEL_SIZES if size <= numa_node_count + ) + layouts = [] + for tensor_parallel_size in range(numa_node_count, 0, -1): + if tensor_parallel_size not in SUPPORTED_TENSOR_PARALLEL_SIZES: + continue + uses_all_numa_nodes = numa_node_count % tensor_parallel_size == 0 + if not uses_all_numa_nodes and tensor_parallel_size != max_supported_tp: + continue + layouts.append( + (tensor_parallel_size, max(1, numa_node_count // tensor_parallel_size)) + ) + return layouts + + +def build_parallel_layout_plan( + config: dict[str, Any], + workload: WorkloadHints, + numa_node_count: int, + *, + model_metadata: ModelParallelMetadata = UNKNOWN_MODEL_PARALLEL_METADATA, + cpu_moe_dp_supported: bool = True, + tune_scheduler: bool = False, +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + """Build candidates and diagnostics after applying parallel safety policy.""" + validate_sweep_workload(workload) + if numa_node_count <= 1: + raise ValueError("Parallel-layout sweep requires at least two NUMA nodes.") + + layouts = _default_parallel_layout_sizes(numa_node_count) + skipped: list[dict[str, Any]] = [] + if model_metadata.is_moe and not cpu_moe_dp_supported: + for tensor_parallel_size, data_parallel_size in layouts: + if data_parallel_size <= 1: + continue + skipped.append( + { + "_benchmark_name": _parallel_layout_name( + tensor_parallel_size, data_parallel_size, numa_node_count + ), + "policy": "cpu_moe_dp_collectives", + "tensor_parallel_size": tensor_parallel_size, + "data_parallel_size": data_parallel_size, + "reason": ( + "CPU communicator does not implement the variable-size " + "collectives required by MoE data parallelism" + ), + } + ) + # Keep TP coverage instead of collapsing the sweep to whichever DP=1 + # layout happened to exist in the full-NUMA candidate set. + layouts = [ + (tensor_parallel_size, 1) + for tensor_parallel_size in sorted( + SUPPORTED_TENSOR_PARALLEL_SIZES, reverse=True + ) + if tensor_parallel_size <= numa_node_count + ] + + pcp_size = config.get("prefill-context-parallel-size", 1) + if not isinstance(pcp_size, int) or isinstance(pcp_size, bool) or pcp_size <= 0: + raise ValueError("prefill-context-parallel-size must be a positive integer.") + uses_expert_parallel = config.get("enable-expert-parallel") is True + + candidates = [] + for tensor_parallel_size, data_parallel_size in layouts: + name = _parallel_layout_name( + tensor_parallel_size, data_parallel_size, numa_node_count + ) + effective_parallel_size = tensor_parallel_size * data_parallel_size * pcp_size + moe_width = model_metadata.moe_intermediate_size + if ( + model_metadata.is_moe + and not uses_expert_parallel + and moe_width is not None + and moe_width % effective_parallel_size != 0 + ): + skipped.append( + { + "_benchmark_name": name, + "policy": "moe_effective_parallel_size", + "tensor_parallel_size": tensor_parallel_size, + "data_parallel_size": data_parallel_size, + "prefill_context_parallel_size": pcp_size, + "effective_parallel_size": effective_parallel_size, + "moe_intermediate_size": moe_width, + "reason": ( + "moe_intermediate_size is not divisible by the effective " + "TP x DP x PCP size while expert parallelism is disabled" + ), + } + ) + continue + + candidate = { + "_benchmark_name": name, + "tensor_parallel_size": tensor_parallel_size, + "data_parallel_size": data_parallel_size, + } + if tune_scheduler: + max_num_seqs, max_num_batched_tokens = _scheduler_baseline_for_dp( + config, workload, data_parallel_size + ) + candidate["max_num_seqs"] = max_num_seqs + candidate["max_num_batched_tokens"] = max_num_batched_tokens + candidates.append(candidate) + + if not candidates: + raise ValueError( + "No valid parallel-layout candidates remain after applying model and " + "runtime capability checks. See parallel_layout_skips.json." + ) + return candidates, skipped + + +def build_parallel_layout_params( + config: dict[str, Any], + workload: WorkloadHints, + numa_node_count: int, + *, + model_metadata: ModelParallelMetadata = UNKNOWN_MODEL_PARALLEL_METADATA, + cpu_moe_dp_supported: bool = True, + tune_scheduler: bool = False, +) -> list[dict[str, Any]]: + """Build full-NUMA layouts plus the largest supported TP layout.""" + candidates, _ = build_parallel_layout_plan( + config, + workload, + numa_node_count, + model_metadata=model_metadata, + cpu_moe_dp_supported=cpu_moe_dp_supported, + tune_scheduler=tune_scheduler, + ) + return candidates + + +def build_bench_params(workload: WorkloadHints) -> list[dict[str, Any]]: + validate_sweep_workload(workload) + assert workload.input_tokens is not None + assert workload.output_tokens is not None + assert workload.concurrency is not None + + return [ + { + "random_input_len": workload.input_tokens, + "random_output_len": workload.output_tokens, + "max_concurrency": workload.concurrency, + "num_prompts": _num_prompts_for_concurrency(workload.concurrency), + } + ] + + +def _benchmark_models(config: dict[str, Any]) -> tuple[str, str]: + model = config.get("model") + if not isinstance(model, str) or not model: + raise ValueError("Sweep generation requires a model in config.yml.") + + served_model_name = config.get("served-model-name") + if isinstance(served_model_name, str) and served_model_name: + return served_model_name, model + if ( + isinstance(served_model_name, list) + and served_model_name + and isinstance(served_model_name[0], str) + ): + return served_model_name[0], model + + return model, model + + +def _relative_to(directory: Path, target: str) -> str: + return os.path.relpath(Path(target).resolve(), directory.resolve()) + + +def _write_json(path: Path, value: object) -> None: + path.write_text( + json.dumps(value, indent=2) + "\n", + encoding="utf-8", + ) + + +def _write_sweep_config(path: Path, config: dict[str, Any]) -> None: + sweep_config = dict(config) + sweep_config.pop("max-num-seqs", None) + sweep_config.pop("max-num-batched-tokens", None) + + body = yaml.safe_dump( + sweep_config, + sort_keys=False, + default_flow_style=False, + allow_unicode=True, + ) + header = ( + "# Generated for runtime sweep only.\n" + "# Scheduler values are supplied per candidate; omitted values use " + "vLLM defaults.\n" + ) + path.write_text(header + body, encoding="utf-8") + + +def _write_run_script( + path: Path, + *, + config_rel: str, + env_rel: str, + request_model: str, + tokenizer: str, + workload: WorkloadHints, + serve_params_name: str = "serve_params.json", + experiment_name: str = "runtime-tuning", + prepare_config_rel: str | None = None, +) -> None: + bench_parts = [ + "vllm bench serve", + "--backend vllm", + f"--model {shlex.quote(request_model)}", + f"--tokenizer {shlex.quote(tokenizer)}", + "--dataset-name random", + "--request-rate inf", + "--ignore-eos", + "--metric-percentiles 99", + ] + + goodput_pairs: list[str] = [] + if workload.ttft_sla_ms is not None: + goodput_pairs.append(f"ttft:{workload.ttft_sla_ms:g}") + if workload.tpot_sla_ms is not None: + goodput_pairs.append(f"tpot:{workload.tpot_sla_ms:g}") + if goodput_pairs: + bench_parts.append( + "--goodput " + " ".join(shlex.quote(pair) for pair in goodput_pairs) + ) + + bench_cmd = " ".join(bench_parts) + prepare_config = "" + if prepare_config_rel is not None: + prepare_config = f""" +SOURCE_CONFIG_PATH="${{SCRIPT_DIR}}/{prepare_config_rel}" +python3 - "${{SOURCE_CONFIG_PATH}}" "${{CONFIG_PATH}}" <<'PY' +import sys +from pathlib import Path + +import yaml + +source = Path(sys.argv[1]) +target = Path(sys.argv[2]) +config = yaml.safe_load(source.read_text(encoding="utf-8")) +if not isinstance(config, dict): + raise SystemExit(str(source) + " does not contain a YAML configuration object.") +config.pop("max-num-seqs", None) +config.pop("max-num-batched-tokens", None) +target.write_text( + yaml.safe_dump( + config, + sort_keys=False, + default_flow_style=False, + allow_unicode=True, + ), + encoding="utf-8", +) +PY + +""" + + script = f"""#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${{BASH_SOURCE[0]}}")" && pwd)" +CONFIG_PATH="${{SCRIPT_DIR}}/{config_rel}" +ENV_PATH="${{SCRIPT_DIR}}/{env_rel}" + +source "${{ENV_PATH}}" + +# Released/nightly vLLM images keep recipe-relative assets such as +# examples/*.jinja under /vllm-workspace. Launch sweep server subprocesses +# from there when available so relative recipe assets resolve like they do +# in serve_with_recipe.sh. +if [[ -d /vllm-workspace/examples ]]; then + cd /vllm-workspace +fi + +# recipe tools may be newer than the installed vLLM package. Only pass the +# resilience flag when this vLLM CLI actually supports it. +CONTINUE_ON_ERROR_ARG="" +if vllm bench sweep serve --help 2>&1 | grep -q -- "--continue-on-error"; then + CONTINUE_ON_ERROR_ARG="--continue-on-error" +fi + +# On older vLLM versions, a failed benchmark aborts the sweep process. Retry the +# interrupted sweep with --resume so one transient run failure does not discard +# completed work or force a manual restart. +SWEEP_RETRY_COUNT="${{VLLM_RECIPE_SWEEP_RETRY_COUNT:-0}}" +MAX_SWEEP_RETRIES="${{VLLM_RECIPE_SWEEP_RETRIES:-2}}" + +{prepare_config}vllm bench sweep serve \ + --serve-cmd "vllm serve --config '${{CONFIG_PATH}}'" \ + --bench-cmd "{bench_cmd}" \ + --serve-params "${{SCRIPT_DIR}}/{serve_params_name}" \ + --bench-params "${{SCRIPT_DIR}}/bench_params.json" \ + --output-dir "${{SCRIPT_DIR}}/results" \ + --experiment-name {experiment_name} \ + --warmup-num-prompts {workload.concurrency} \ + ${{CONTINUE_ON_ERROR_ARG}} \ + "$@" || {{ + if [[ "${{SWEEP_RETRY_COUNT}}" -ge "${{MAX_SWEEP_RETRIES}}" ]]; then + echo "Sweep failed after ${{MAX_SWEEP_RETRIES}} automatic retries." >&2 + exit 1 + fi + + NEXT_RETRY=$((SWEEP_RETRY_COUNT + 1)) + RETRY_ARGS=() + for arg in "$@"; do + [[ "$arg" == "--resume" ]] && continue + RETRY_ARGS+=("$arg") + done + + EXPERIMENT_DIR="${{SCRIPT_DIR}}/results/{experiment_name}" + if [[ -d "${{EXPERIMENT_DIR}}" ]]; then + echo "Sweep command failed;" \ + "retry ${{NEXT_RETRY}}/${{MAX_SWEEP_RETRIES}} with --resume" + RETRY_ARGS=(--resume "${{RETRY_ARGS[@]}}") + else + echo "Sweep command failed before resumable state was created;" \ + "retry ${{NEXT_RETRY}}/${{MAX_SWEEP_RETRIES}} from the start" + fi + + VLLM_RECIPE_SWEEP_RETRY_COUNT="${{NEXT_RETRY}}" exec "$0" "${{RETRY_ARGS[@]}}" + }} +""" + path.write_text(script, encoding="utf-8") + path.chmod(path.stat().st_mode | 0o111) + + +def _write_recommend_script( + path: Path, + *, + config_rel: str, + env_rel: str, + workload: WorkloadHints, + results_dir: str = "results/runtime-tuning", + output_config: str = "recommended-config.yml", + output_json: str = "recommendation.json", +) -> None: + template_path = Path(__file__).with_name("sweep_recommendation.py") + source = template_path.read_text(encoding="utf-8") + + replacements = { + "DEFAULT_CONFIG_PATH: str | None = None": ( + f"DEFAULT_CONFIG_PATH: str | None = {config_rel!r}" + ), + "DEFAULT_ENV_PATH: str | None = None": ( + f"DEFAULT_ENV_PATH: str | None = {env_rel!r}" + ), + "DEFAULT_TTFT_SLA_MS: float | None = None": ( + f"DEFAULT_TTFT_SLA_MS: float | None = {workload.ttft_sla_ms!r}" + ), + "DEFAULT_TPOT_SLA_MS: float | None = None": ( + f"DEFAULT_TPOT_SLA_MS: float | None = {workload.tpot_sla_ms!r}" + ), + 'DEFAULT_RESULTS_DIR = "results/runtime-tuning"': ( + f"DEFAULT_RESULTS_DIR = {results_dir!r}" + ), + 'DEFAULT_OUTPUT_CONFIG = "recommended-config.yml"': ( + f"DEFAULT_OUTPUT_CONFIG = {output_config!r}" + ), + 'DEFAULT_OUTPUT_JSON = "recommendation.json"': ( + f"DEFAULT_OUTPUT_JSON = {output_json!r}" + ), + } + for marker, replacement in replacements.items(): + if marker not in source: + raise ValueError(f"Recommender template marker not found: {marker}") + source = source.replace(marker, replacement, 1) + + path.write_text(source, encoding="utf-8") + path.chmod(path.stat().st_mode | 0o111) + + +def _write_guide(path: Path, workload: WorkloadHints) -> None: + sla_lines: list[str] = [] + if workload.ttft_sla_ms is not None: + sla_lines.append(f"- TTFT objective: `{workload.ttft_sla_ms:g} ms`") + if workload.tpot_sla_ms is not None: + sla_lines.append(f"- TPOT objective: `{workload.tpot_sla_ms:g} ms`") + + sla_text = "\n".join(sla_lines) + if sla_text: + sla_text = ( + "\n## Supplied latency objectives\n\n" + + sla_text + + "\n\nThe generated benchmark uses these values with vLLM " + "`--goodput`.\n" + ) + + content = f"""# Optional Runtime Tuning Sweep + +The generated `config.yml` is the **single initial suggestion** and can be +deployed directly. The scheduler sweep keeps explicit `max-num-seqs` at the +per-replica concurrency lower bound while benchmarking nearby values for: + +- `max-num-batched-tokens` + +It does not test explicit `max-num-seqs` values below per-replica concurrency, +because those values can introduce scheduler queueing and severe TTFT +degradation. The vLLM-default references still test whether an explicit +`max-num-seqs` override is needed. + +In addition to nearby tuned values, the sweep includes three vLLM-default +references: + +- vLLM default `max-num-seqs` with the initial batch-token budget +- initial `max-num-seqs` with vLLM default `max-num-batched-tokens` +- vLLM defaults for both parameters + +The sweep runs every server from a generated `sweep_config.yml` that copies the +initial `config.yml` but removes these two scheduler keys. Tuned candidates add +explicit CLI values; a default-reference candidate simply omits the selected +key, allowing `vllm serve` to resolve its normal runtime default. This avoids +passing the literal string `None` and keeps defaults platform-, world-size-, +model-, and usage-context-aware. + +## Run + +```bash +./run_sweep.sh --dry-run +./run_sweep.sh +./recommend.py +``` + +For a quick one-run experiment: + +```bash +rm -rf results/runtime-tuning +./run_sweep.sh --num-runs 1 +./recommend.py +``` + +Resume an interrupted sweep: + +```bash +./run_sweep.sh --resume +``` + +Before measurement, the generated script runs one unmeasured warmup containing +one full concurrency window ({workload.concurrency} prompts). It is saved as +`warmup.json` for auditing but excluded from recommendation statistics. By +default, vLLM then benchmarks each parameter combination three times. + +The benchmark request count scales with the supplied workload concurrency using +the same model as vLLM's `PROMPTS_PER_CONCURRENCY` performance-benchmark +control: + +```text +num_prompts = min(1000, max(100, concurrency * 10)) +``` + +This targets about ten concurrency turnovers when neither bound applies, keeps +at least 100 requests for percentile/compliance measurements, and caps each run +at 1000 requests. All scheduler candidates use the same request count. +{sla_text} +## Outputs + +`recommend.py` writes: + +```text +recommended-config.yml +recommendation.json +``` + +With TTFT/TPOT objectives it requires duration-weighted combined compliance of +at least 99% plus median P99 TTFT/TPOT compliance. Scheduler candidates within +1% of the highest eligible output-token throughput are treated as practically +equivalent, and the recommender prefers fewer explicit scheduler overrides +within that set. Use `--throughput-equivalence-percent 0` for exact +highest-throughput selection. If no candidate qualifies, it records a +best-effort candidate but does not write `recommended-config.yml`. Without +latency objectives it applies the same scheduler equivalence policy. +Configurations with failed requests are excluded. TP/DP comparisons continue +to use exact highest-throughput selection. + +`recommended-config.yml` copies the initial configuration and changes only +`max-num-seqs` and `max-num-batched-tokens`. If a vLLM-default reference wins, +the corresponding key is removed from `recommended-config.yml`, allowing +`vllm serve` to resolve that parameter from its normal runtime default policy. +""" + path.write_text(content, encoding="utf-8") + + +def _write_report_script(path: Path, *, workload: WorkloadHints) -> None: + template_path = Path(__file__).with_name("sweep_html_report.py") + source = template_path.read_text(encoding="utf-8") + replacements = { + "DEFAULT_TTFT_SLA_MS: float | None = None": ( + f"DEFAULT_TTFT_SLA_MS: float | None = {workload.ttft_sla_ms!r}" + ), + "DEFAULT_TPOT_SLA_MS: float | None = None": ( + f"DEFAULT_TPOT_SLA_MS: float | None = {workload.tpot_sla_ms!r}" + ), + } + for marker, replacement in replacements.items(): + if marker not in source: + raise ValueError(f"HTML report template marker not found: {marker}") + source = source.replace(marker, replacement, 1) + path.write_text(source, encoding="utf-8") + path.chmod(path.stat().st_mode | 0o111) + + +def _write_post_benchmark_analysis_files( + directory: Path, *, workload: WorkloadHints +) -> list[Path]: + """Copy standalone reporting helpers into a generated sweep package.""" + source_dir = Path(__file__).resolve().parent + generated: list[Path] = [] + for source_name, target_name, executable in ( + ("sweep_visualization.py", "visualize.py", True), + ("requirements.txt", "requirements.txt", False), + ("VISUALIZATION.md", "VISUALIZATION.md", False), + ): + source = ( + source_dir / source_name + if source_name.endswith(".py") + else source_dir.parent / source_name + ) + target = directory / target_name + target.write_text(source.read_text(encoding="utf-8"), encoding="utf-8") + if executable: + target.chmod(target.stat().st_mode | 0o111) + generated.append(target) + report = directory / "report.py" + _write_report_script(report, workload=workload) + generated.append(report) + return generated + + +def write_sweep_files( + output_dir: str, + *, + config_path: str, + env_path: str, + config: dict[str, Any], + workload: WorkloadHints, +) -> list[Path]: + """Write an optional sweep package around the initial config suggestion.""" + validate_sweep_workload(workload) + + directory = Path(output_dir) + directory.mkdir(parents=True, exist_ok=True) + + sweep_config = directory / "sweep_config.yml" + serve_params = directory / "serve_params.json" + bench_params = directory / "bench_params.json" + run_script = directory / "run_sweep.sh" + recommend_script = directory / "recommend.py" + guide = directory / "SWEEP.md" + + config_rel = _relative_to(directory, config_path) + sweep_config_rel = _relative_to(directory, str(sweep_config)) + env_rel = _relative_to(directory, env_path) + + _write_sweep_config(sweep_config, config) + _write_json(serve_params, build_serve_params(config)) + _write_json(bench_params, build_bench_params(workload)) + request_model, tokenizer = _benchmark_models(config) + _write_run_script( + run_script, + config_rel=sweep_config_rel, + env_rel=env_rel, + request_model=request_model, + tokenizer=tokenizer, + workload=workload, + ) + _write_recommend_script( + recommend_script, + config_rel=config_rel, + env_rel=env_rel, + workload=workload, + ) + _write_guide(guide, workload) + analysis_files = _write_post_benchmark_analysis_files(directory, workload=workload) + + return [ + sweep_config, + serve_params, + bench_params, + run_script, + recommend_script, + guide, + *analysis_files, + ] + + +def write_parallel_layout_sweep_files( + output_dir: str, + *, + config_path: str, + env_path: str, + config: dict[str, Any], + workload: WorkloadHints, + numa_node_count: int, + tune_scheduler: bool = False, +) -> list[Path]: + """Write a standalone NUMA-aware TP/DP sweep package.""" + validate_sweep_workload(workload) + directory = Path(output_dir) + directory.mkdir(parents=True, exist_ok=True) + + parallel_params = directory / "parallel_layout_serve_params.json" + parallel_skips = directory / "parallel_layout_skips.json" + parallel_config = directory / "parallel-layout-benchmark-config.yml" + bench_params = directory / "bench_params.json" + run_parallel = directory / "run_parallel_layout_sweep.sh" + recommend_parallel = directory / "recommend_parallel_layout.py" + guide = directory / "SWEEP.md" + + initial_config_rel = _relative_to(directory, config_path) + parallel_config_rel = _relative_to(directory, str(parallel_config)) + selected_config_rel = "parallel-layout-config.yml" + env_rel = _relative_to(directory, env_path) + request_model, tokenizer = _benchmark_models(config) + + _write_sweep_config(parallel_config, config) + benchmark_config_rel = initial_config_rel if tune_scheduler else parallel_config_rel + + try: + model_metadata = inspect_model_parallel_metadata(config) + except Exception as exc: + print( + "Warning: could not inspect model parallel metadata; retaining the " + f"generic layouts and relying on sweep error isolation: {exc}", + file=sys.stderr, + ) + model_metadata = UNKNOWN_MODEL_PARALLEL_METADATA + cpu_moe_dp_supported = ( + cpu_supports_moe_dp_collectives() if model_metadata.is_moe else True + ) + parallel_candidates, skipped_candidates = build_parallel_layout_plan( + config, + workload, + numa_node_count, + model_metadata=model_metadata, + cpu_moe_dp_supported=cpu_moe_dp_supported, + tune_scheduler=tune_scheduler, + ) + _write_json(parallel_params, parallel_candidates) + _write_json(parallel_skips, skipped_candidates) + _write_json(bench_params, build_bench_params(workload)) + _write_run_script( + run_parallel, + config_rel=benchmark_config_rel, + env_rel=env_rel, + request_model=request_model, + tokenizer=tokenizer, + workload=workload, + serve_params_name=parallel_params.name, + experiment_name="parallel-layout", + ) + _write_recommend_script( + recommend_parallel, + config_rel=initial_config_rel, + env_rel=env_rel, + workload=workload, + results_dir="results/parallel-layout", + output_config=selected_config_rel, + output_json="parallel-layout-recommendation.json", + ) + + guide.write_text( + f"""# Parallel-Layout Sweep + +This package tunes only the NUMA-aware TP/DP layout. + +Primary candidates use every effective NUMA node: + +```text +tensor_parallel_size * data_parallel_size = {numa_node_count} +``` + +Tensor parallelism is restricted to `1`, `2`, `4`, or `8`. The largest +supported TP not exceeding the NUMA-node count is also included, even when it +leaves some NUMA nodes idle. + +Before benchmarking, the generator applies two MoE safety checks: + +- When the CPU communicator lacks variable-size MoE collectives, MoE candidates + use `DP=1`; every supported TP size is retained as a replacement candidate. +- With expert parallelism disabled and a known MoE intermediate size, candidates + must satisfy `moe_intermediate_size % (TP * DP * PCP) == 0`. + +Rejected layouts and their reasons are written to +`parallel_layout_skips.json`. If model metadata cannot be inspected, generic +layouts are retained and the runner's per-combination error isolation applies. + +Run: + +```bash +./run_parallel_layout_sweep.sh --dry-run +./run_parallel_layout_sweep.sh +./recommend_parallel_layout.py +``` + +The recommender writes: + +```text +parallel-layout-config.yml +parallel-layout-recommendation.json +``` + +This mode stops after TP/DP selection. Use `--generate-full-sweep` when +concurrency and scheduler tuning should follow automatically. +""", + encoding="utf-8", + ) + analysis_files = _write_post_benchmark_analysis_files(directory, workload=workload) + + return [ + parallel_params, + parallel_skips, + bench_params, + run_parallel, + recommend_parallel, + guide, + *analysis_files, + ] + + +def build_concurrency_bench_params(workload: WorkloadHints) -> list[dict[str, Any]]: + """Build the fixed workload shape for concurrency tuning.""" + validate_sweep_workload(workload) + assert workload.input_tokens is not None + assert workload.output_tokens is not None + assert workload.concurrency is not None + return [ + { + "_benchmark_name": "user_workload", + "random_input_len": workload.input_tokens, + "random_output_len": workload.output_tokens, + "num_prompts": _num_prompts_for_concurrency(workload.concurrency), + } + ] + + +def _write_adaptive_concurrency_script(path: Path) -> None: + template_path = Path(__file__).with_name("adaptive_concurrency.py") + path.write_text(template_path.read_text(encoding="utf-8"), encoding="utf-8") + path.chmod(path.stat().st_mode | 0o111) + + +def _write_concurrency_run_script( + path: Path, + *, + config_rel: str, + env_rel: str, + request_model: str, + tokenizer: str, + workload: WorkloadHints, +) -> None: + bench_parts = [ + "vllm bench serve", + "--backend vllm", + f"--model {shlex.quote(request_model)}", + f"--tokenizer {shlex.quote(tokenizer)}", + "--dataset-name random", + "--request-rate inf", + "--ignore-eos", + "--metric-percentiles 99", + ] + goodput_pairs: list[str] = [] + if workload.ttft_sla_ms is not None: + goodput_pairs.append(f"ttft:{workload.ttft_sla_ms:g}") + if workload.tpot_sla_ms is not None: + goodput_pairs.append(f"tpot:{workload.tpot_sla_ms:g}") + if goodput_pairs: + bench_parts.append( + "--goodput " + " ".join(shlex.quote(pair) for pair in goodput_pairs) + ) + + bench_cmd = " ".join(bench_parts) + use_adaptive_sla_search = bool( + workload.ttft_sla_ms is not None or workload.tpot_sla_ms is not None + ) + adaptive_args = [ + f"--seed-concurrency {workload.concurrency}", + f"--input-tokens {workload.input_tokens}", + f"--output-tokens {workload.output_tokens}", + ] + if workload.ttft_sla_ms is not None: + adaptive_args.append(f"--ttft-sla-ms {workload.ttft_sla_ms:g}") + if workload.tpot_sla_ms is not None: + adaptive_args.append(f"--tpot-sla-ms {workload.tpot_sla_ms:g}") + adaptive_args_text = " \\\n ".join(adaptive_args) + + script = f"""#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${{BASH_SOURCE[0]}}")" && pwd)" +CONFIG_PATH="${{SCRIPT_DIR}}/{config_rel}" +ENV_PATH="${{SCRIPT_DIR}}/{env_rel}" + +source "${{ENV_PATH}}" + +# Released/nightly vLLM images keep recipe-relative assets such as +# examples/*.jinja under /vllm-workspace. Launch sweep server subprocesses +# from there when available so relative recipe assets resolve like they do +# in serve_with_recipe.sh. +if [[ -d /vllm-workspace/examples ]]; then + cd /vllm-workspace +fi + +# recipe tools may be newer than the installed vLLM package. Only pass the +# resilience flag when this vLLM CLI actually supports it. +CONTINUE_ON_ERROR_ARG="" +if [[ "{int(use_adaptive_sla_search)}" == "1" ]] || \ + vllm bench sweep serve_workload --help 2>&1 | grep -q -- "--continue-on-error"; then + CONTINUE_ON_ERROR_ARG="--continue-on-error" +fi + +# On older vLLM versions, a failed benchmark aborts the sweep process. Retry the +# interrupted sweep with --resume so one transient run failure does not discard +# completed work or force a manual restart. +SWEEP_RETRY_COUNT="${{VLLM_RECIPE_SWEEP_RETRY_COUNT:-0}}" +MAX_SWEEP_RETRIES="${{VLLM_RECIPE_SWEEP_RETRIES:-2}}" +SERVER_READY_TIMEOUT="${{VLLM_RECIPE_SERVER_READY_TIMEOUT:-1800}}" + +run_concurrency_sweep() {{ + if [[ "{int(use_adaptive_sla_search)}" == "1" ]]; then + python3 "${{SCRIPT_DIR}}/adaptive_concurrency.py" \ + --serve-cmd "vllm serve --config '${{CONFIG_PATH}}'" \ + --bench-cmd "{bench_cmd}" \ + --results-dir "${{SCRIPT_DIR}}/results/concurrency-tuning" \ + {adaptive_args_text} \ + --server-ready-timeout "${{SERVER_READY_TIMEOUT}}" \ + ${{CONTINUE_ON_ERROR_ARG}} \ + "$@" + else + vllm bench sweep serve_workload \ + --serve-cmd "vllm serve --config '${{CONFIG_PATH}}'" \ + --bench-cmd "{bench_cmd}" \ + --bench-params "${{SCRIPT_DIR}}/concurrency_bench_params.json" \ + --output-dir "${{SCRIPT_DIR}}/results" \ + --experiment-name concurrency-tuning \ + --workload-var max_concurrency \ + --workload-iters 10 \ + --warmup-num-prompts {workload.concurrency} \ + ${{CONTINUE_ON_ERROR_ARG}} \ + "$@" + fi +}} + +run_concurrency_sweep || {{ + if [[ "${{SWEEP_RETRY_COUNT}}" -ge "${{MAX_SWEEP_RETRIES}}" ]]; then + echo "Sweep failed after ${{MAX_SWEEP_RETRIES}} automatic retries." >&2 + exit 1 + fi + + NEXT_RETRY=$((SWEEP_RETRY_COUNT + 1)) + RETRY_ARGS=() + for arg in "$@"; do + [[ "$arg" == "--resume" ]] && continue + RETRY_ARGS+=("$arg") + done + + EXPERIMENT_DIR="${{SCRIPT_DIR}}/results/concurrency-tuning" + if [[ -d "${{EXPERIMENT_DIR}}" ]]; then + echo "Sweep command failed;" \ + "retry ${{NEXT_RETRY}}/${{MAX_SWEEP_RETRIES}} with --resume" + RETRY_ARGS=(--resume "${{RETRY_ARGS[@]}}") + else + echo "Sweep command failed before resumable state was created;" \ + "retry ${{NEXT_RETRY}}/${{MAX_SWEEP_RETRIES}} from the start" + fi + + VLLM_RECIPE_SWEEP_RETRY_COUNT="${{NEXT_RETRY}}" exec "$0" "${{RETRY_ARGS[@]}}" + }} +""" + path.write_text(script, encoding="utf-8") + path.chmod(path.stat().st_mode | 0o111) + + +def _write_concurrency_recommend_script(path: Path, *, workload: WorkloadHints) -> None: + template_path = Path(__file__).with_name("concurrency_recommendation.py") + source = template_path.read_text(encoding="utf-8") + replacements = { + "DEFAULT_TTFT_SLA_MS: float | None = None": ( + f"DEFAULT_TTFT_SLA_MS: float | None = {workload.ttft_sla_ms!r}" + ), + "DEFAULT_TPOT_SLA_MS: float | None = None": ( + f"DEFAULT_TPOT_SLA_MS: float | None = {workload.tpot_sla_ms!r}" + ), + "DEFAULT_SEED_CONCURRENCY: int | None = None": ( + f"DEFAULT_SEED_CONCURRENCY: int | None = {workload.concurrency!r}" + ), + } + for marker, replacement in replacements.items(): + if marker not in source: + raise ValueError(f"Concurrency recommender marker not found: {marker}") + source = source.replace(marker, replacement, 1) + path.write_text(source, encoding="utf-8") + path.chmod(path.stat().st_mode | 0o111) + + +def _write_scheduler_prepare_script(path: Path) -> None: + template_path = Path(__file__).with_name("scheduler_sweep_prepare.py") + path.write_text(template_path.read_text(encoding="utf-8"), encoding="utf-8") + path.chmod(path.stat().st_mode | 0o111) + + +def _write_scheduler_wrapper(path: Path, *, workload: WorkloadHints) -> None: + assert workload.input_tokens is not None + assert workload.output_tokens is not None + args = [ + 'python3 "${SCRIPT_DIR}/prepare_scheduler.py"', + '--config "${SCRIPT_DIR}/parallel-layout-config.yml"', + '--concurrency-recommendation "${SCRIPT_DIR}/concurrency-recommendation.json"', + f"--input-tokens {workload.input_tokens}", + f"--output-tokens {workload.output_tokens}", + ] + if workload.tpot_sla_ms is not None: + args.append(f"--tpot-sla-ms {workload.tpot_sla_ms:g}") + if workload.target_qps is not None: + args.append(f"--target-qps {workload.target_qps:g}") + args.extend( + [ + '--output-config "${SCRIPT_DIR}/parallel-layout-sweep-config.yml"', + '--serve-params "${SCRIPT_DIR}/serve_params.json"', + '--bench-params "${SCRIPT_DIR}/bench_params.json"', + ] + ) + prepare_cmd = " \\\n ".join(args) + script = f"""#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${{BASH_SOURCE[0]}}")" && pwd)" + +{prepare_cmd} + +exec "${{SCRIPT_DIR}}/run_sweep.sh" "$@" +""" + path.write_text(script, encoding="utf-8") + path.chmod(path.stat().st_mode | 0o111) + + +def _recommend_sla_args(workload: WorkloadHints) -> list[str]: + args = [] + if workload.ttft_sla_ms is not None: + args.append(f"--ttft-sla-ms {workload.ttft_sla_ms:g}") + if workload.tpot_sla_ms is not None: + args.append(f"--tpot-sla-ms {workload.tpot_sla_ms:g}") + return args + + +def _write_full_run_script(path: Path, *, workload: WorkloadHints) -> None: + recommend_sla_args = " ".join(_recommend_sla_args(workload)) + script = f"""#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${{BASH_SOURCE[0]}}")" && pwd)" +STAGES=all +RUN_ARGS=() +SERVER_READY_TIMEOUT="${{VLLM_RECIPE_SERVER_READY_TIMEOUT:-1800}}" +RECOMMEND_SLA_ARGS=({recommend_sla_args}) +while (($#)); do + case "$1" in + --stages) + if (($# < 2)); then + echo "--stages requires a comma-separated stage list." >&2 + exit 2 + fi + STAGES="$2" + shift 2 + ;; + --stages=*) STAGES="${{1#*=}}"; shift ;; + --server-ready-timeout) + if (($# < 2)); then + echo "--server-ready-timeout requires a value." >&2 + exit 2 + fi + SERVER_READY_TIMEOUT="$2" + RUN_ARGS+=("$1" "$2") + shift 2 + ;; + --server-ready-timeout=*) + SERVER_READY_TIMEOUT="${{1#*=}}" + RUN_ARGS+=("$1") + shift + ;; + --help|-h) + echo "Usage: $0 [--stages all|parallel-layout,concurrency,scheduler] [sweep args]" + echo "Selected stages always run in dependency order. Default: all." + exit 0 + ;; + *) RUN_ARGS+=("$1"); shift ;; + esac +done +case "$STAGES" in + all) STAGES=parallel-layout,concurrency,scheduler ;; + ""|,*|*,|*,,*) echo "Invalid stage list: $STAGES" >&2; exit 2 ;; +esac +IFS=, read -r -a SELECTED_STAGES <<< "$STAGES" +for stage in "${{SELECTED_STAGES[@]}}"; do + case "$stage" in + parallel-layout|concurrency|scheduler) ;; + *) echo "Unknown sweep stage: $stage" >&2; exit 2 ;; + esac +done +has_stage() {{ [[ ",$STAGES," == *",$1,"* ]]; }} + +# Keep the concurrency runner's server-start timeout explicit even when sweep +# arguments pass through multiple generated wrapper scripts. +export VLLM_RECIPE_SERVER_READY_TIMEOUT="${{SERVER_READY_TIMEOUT}}" + +# Check dependencies before starting any benchmarks. +if has_stage scheduler && ! has_stage concurrency && + [[ ! -f "${{SCRIPT_DIR}}/concurrency-recommendation.json" ]]; then + echo "Scheduler requires concurrency-recommendation.json; run concurrency first." >&2 + exit 2 +fi +if ! has_stage parallel-layout && + [[ ! -f "${{SCRIPT_DIR}}/parallel-layout-config.yml" ]]; then + cp "${{SCRIPT_DIR}}/initial-config.yml" "${{SCRIPT_DIR}}/parallel-layout-config.yml" +fi + +if has_stage parallel-layout; then + "${{SCRIPT_DIR}}/run_parallel_layout_sweep.sh" "${{RUN_ARGS[@]}}" + "${{SCRIPT_DIR}}/recommend_parallel_layout.py" "${{RECOMMEND_SLA_ARGS[@]}}" +fi +if has_stage concurrency; then + "${{SCRIPT_DIR}}/run_concurrency_sweep.sh" "${{RUN_ARGS[@]}}" + "${{SCRIPT_DIR}}/recommend_concurrency.py" "${{RECOMMEND_SLA_ARGS[@]}}" +fi +if has_stage scheduler; then + "${{SCRIPT_DIR}}/run_scheduler_sweep.sh" "${{RUN_ARGS[@]}}" + "${{SCRIPT_DIR}}/recommend.py" "${{RECOMMEND_SLA_ARGS[@]}}" +fi + +if ! "${{SCRIPT_DIR}}/report.py" "${{RECOMMEND_SLA_ARGS[@]}}"; then + echo "Warning: sweep completed, but HTML report generation failed." >&2 +fi +""" + path.write_text(script, encoding="utf-8") + path.chmod(path.stat().st_mode | 0o111) + + +def _write_full_guide(path: Path, workload: WorkloadHints) -> None: + content = f"""# End-to-End Runtime Tuning + +This package keeps one fixed workload shape and tunes in dependency order: + +```text +TP/DP -> max_concurrency -> scheduler +``` + +The supplied `--concurrency={workload.concurrency}` is the representative load +used for TP/DP selection and to size the Workload Explorer dataset. Stage 2 +selects the measured SLA-feasible `max_concurrency`. Stage 3 recalculates its +scheduler baseline from that concurrency and the selected DP layout. + +Run all stages: + +```bash +./run_full_sweep.sh +``` + +After the final recommendation, the full runner writes `sweep-report.html`. +Report-generation failures are warnings and do not change sweep success. + +Run selected stages (always executed in dependency order): + +```bash +./run_full_sweep.sh --stages concurrency +./run_full_sweep.sh --stages concurrency,scheduler +./run_full_sweep.sh --stages scheduler --resume +``` + +When parallel-layout tuning is skipped, an existing `parallel-layout-config.yml` +is reused; otherwise `initial-config.yml` supplies the fixed TP/DP layout. +Scheduler-only runs require an existing `concurrency-recommendation.json`. +Other arguments are forwarded to the selected benchmark runners. + +Or run each stage explicitly: + +```bash +./run_parallel_layout_sweep.sh +./recommend_parallel_layout.py +./run_concurrency_sweep.sh +./recommend_concurrency.py +./run_scheduler_sweep.sh +./recommend.py +``` + +Generated recipe sweeps use `--continue-on-error`. An isolated warmup or +benchmark invocation failure is recorded and skipped. A combination with no +successful measured runs remains retryable with `--resume`. +""" + path.write_text(content, encoding="utf-8") + + +def write_concurrency_sweep_files( + output_dir: str, + *, + config_path: str, + env_path: str, + config: dict[str, Any], + workload: WorkloadHints, + tune_scheduler: bool = False, +) -> list[Path]: + """Write a standalone max_concurrency tuning package.""" + validate_sweep_workload(workload) + directory = Path(output_dir) + directory.mkdir(parents=True, exist_ok=True) + + bench_params = directory / "concurrency_bench_params.json" + concurrency_config = directory / "concurrency-benchmark-config.yml" + adaptive_script = directory / "adaptive_concurrency.py" + run_script = directory / "run_concurrency_sweep.sh" + recommend_script = directory / "recommend_concurrency.py" + + config_rel = _relative_to(directory, config_path) + concurrency_config_rel = _relative_to(directory, str(concurrency_config)) + env_rel = _relative_to(directory, env_path) + request_model, tokenizer = _benchmark_models(config) + + _write_sweep_config(concurrency_config, config) + benchmark_config_rel = config_rel if tune_scheduler else concurrency_config_rel + + _write_json(bench_params, build_concurrency_bench_params(workload)) + _write_adaptive_concurrency_script(adaptive_script) + _write_concurrency_run_script( + run_script, + config_rel=benchmark_config_rel, + env_rel=env_rel, + request_model=request_model, + tokenizer=tokenizer, + workload=workload, + ) + _write_concurrency_recommend_script(recommend_script, workload=workload) + analysis_files = _write_post_benchmark_analysis_files(directory, workload=workload) + return [ + bench_params, + concurrency_config, + adaptive_script, + run_script, + recommend_script, + *analysis_files, + ] + + +def write_full_sweep_files( + output_dir: str, + *, + config_path: str, + env_path: str, + config: dict[str, Any], + workload: WorkloadHints, + numa_node_count: int, + tune_scheduler: bool = False, +) -> list[Path]: + """Write TP/DP -> concurrency -> scheduler end-to-end tuning artifacts.""" + files = write_parallel_layout_sweep_files( + output_dir, + config_path=config_path, + env_path=env_path, + config=config, + workload=workload, + numa_node_count=numa_node_count, + tune_scheduler=tune_scheduler, + ) + + directory = Path(output_dir) + concurrency_bench = directory / "concurrency_bench_params.json" + adaptive_concurrency = directory / "adaptive_concurrency.py" + run_concurrency = directory / "run_concurrency_sweep.sh" + recommend_concurrency = directory / "recommend_concurrency.py" + prepare_scheduler = directory / "prepare_scheduler.py" + scheduler_runner = directory / "run_sweep.sh" + recommend_scheduler = directory / "recommend.py" + run_scheduler = directory / "run_scheduler_sweep.sh" + run_full = directory / "run_full_sweep.sh" + guide = directory / "SWEEP.md" + + env_rel = _relative_to(directory, env_path) + request_model, tokenizer = _benchmark_models(config) + + _write_json(concurrency_bench, build_concurrency_bench_params(workload)) + _write_adaptive_concurrency_script(adaptive_concurrency) + _write_concurrency_run_script( + run_concurrency, + config_rel="parallel-layout-config.yml", + env_rel=env_rel, + request_model=request_model, + tokenizer=tokenizer, + workload=workload, + ) + _write_concurrency_recommend_script(recommend_concurrency, workload=workload) + + # Full tuning owns the scheduler stage. The standalone parallel-layout mode + # intentionally stops after TP/DP selection. + _write_scheduler_prepare_script(prepare_scheduler) + _write_run_script( + scheduler_runner, + config_rel="parallel-layout-sweep-config.yml", + env_rel=env_rel, + request_model=request_model, + tokenizer=tokenizer, + workload=workload, + ) + _write_recommend_script( + recommend_scheduler, + config_rel="parallel-layout-config.yml", + env_rel=env_rel, + workload=workload, + ) + _write_scheduler_wrapper(run_scheduler, workload=workload) + + initial_config = directory / "initial-config.yml" + initial_config.write_text(yaml.safe_dump(config, sort_keys=False), encoding="utf-8") + _write_full_run_script(run_full, workload=workload) + _write_full_guide(guide, workload) + + return [ + *files, + concurrency_bench, + adaptive_concurrency, + run_concurrency, + recommend_concurrency, + prepare_scheduler, + scheduler_runner, + recommend_scheduler, + run_scheduler, + run_full, + initial_config, + ] diff --git a/tools/recipes/sweep/sweep_html_report.py b/tools/recipes/sweep/sweep_html_report.py new file mode 100644 index 000000000000..6d2271360862 --- /dev/null +++ b/tools/recipes/sweep/sweep_html_report.py @@ -0,0 +1,756 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Generate a self-contained HTML summary from completed recipe sweeps.""" + +from __future__ import annotations + +import argparse +import json +from datetime import datetime, timezone +from html import escape +from pathlib import Path +from typing import Any + +STAGES = ( + ( + "parallel-layout", + "Parallel layout", + "parallel-layout-recommendation.json", + ), + ( + "concurrency-tuning", + "Concurrency", + "concurrency-recommendation.json", + ), + ("runtime-tuning", "Scheduler", "recommendation.json"), +) + +DEFAULT_TTFT_SLA_MS: float | None = None +DEFAULT_TPOT_SLA_MS: float | None = None + + +def _read_object(path: Path) -> dict[str, Any] | None: + if not path.is_file(): + return None + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise ValueError(f"{path} does not contain a JSON object.") + return value + + +def _summary_runs(stage_dir: Path) -> list[dict[str, Any]]: + runs: list[dict[str, Any]] = [] + for path in sorted(stage_dir.glob("**/summary.json")): + value = json.loads(path.read_text(encoding="utf-8")) + if isinstance(value, list): + runs.extend(item for item in value if isinstance(item, dict)) + return runs + + +def _number(value: Any) -> float | None: + if isinstance(value, bool) or not isinstance(value, int | float): + return None + return float(value) + + +def _fmt(value: Any, digits: int = 1) -> str: + number = _number(value) + if number is None: + return "-" + return f"{number:.{digits}f}" + + +def _fmt_int(value: Any) -> str: + number = _number(value) + return "-" if number is None else f"{number:.0f}" + + +def _fmt_percent(value: Any) -> str: + number = _number(value) + return "-" if number is None else f"{number:.1f}%" + + +def _text(value: Any) -> str: + if value is None: + return "-" + return escape(str(value)) + + +def _resolve_sla( + reports: dict[str, dict[str, Any] | None], + key: str, + expected: float | None, +) -> float | None: + values: dict[str, float | None] = {} + for stage, label, _ in STAGES: + report = reports[stage] + if report is not None: + values[label] = _number((report.get("slo") or {}).get(key)) + + non_null = {value for value in values.values() if value is not None} + if expected is not None: + mismatched = { + label: value for label, value in values.items() if value != expected + } + if mismatched: + details = ", ".join( + f"{label}={value}" for label, value in mismatched.items() + ) + raise ValueError( + f"Recommendation {key} SLA does not match the command-line " + f"value {expected}: {details}" + ) + return expected + + if len(non_null) > 1: + details = ", ".join(f"{label}={value}" for label, value in values.items()) + raise ValueError(f"Recommendation {key} SLA values disagree: {details}") + return next(iter(non_null), None) + + +def _resolve_minimum_compliance( + reports: dict[str, dict[str, Any] | None], +) -> float | None: + values = { + _number(report.get("minimum_compliance_ratio")) + for report in reports.values() + if report is not None + } + non_null = {value for value in values if value is not None} + if len(non_null) > 1: + raise ValueError( + "Recommendation minimum-compliance values disagree: " + + ", ".join(str(value) for value in sorted(non_null)) + ) + return next(iter(non_null), None) + + +def _table(headers: list[str], rows: list[list[Any]]) -> str: + if not rows: + return '

No completed candidates were found.

' + head = "".join(f"{escape(header)}" for header in headers) + body = [] + for row in rows: + cells = "".join(f"{_text(cell)}" for cell in row) + body.append("" + cells + "") + return ( + '
' + + head + + "" + + "".join(body) + + "
" + ) + + +def _status_badge(value: Any) -> str: + status = str(value or "unknown").replace("_", " ") + css_class = "pass" if value == "sla_feasible" else "warn" + return f'{escape(status)}' + + +def _sla_label(candidate: dict[str, Any]) -> str: + if candidate.get("sla_eligible") is True: + return "PASS" + if candidate.get("sla_eligible") is False: + return "FAIL" + return "-" + + +def _sla_reason( + candidate: dict[str, Any], + *, + ttft_sla_ms: float | None, + tpot_sla_ms: float | None, + minimum_compliance: float | None, +) -> str: + if candidate.get("valid") is False or candidate.get("benchmark_valid") is False: + return str(candidate.get("invalid_reason") or "benchmark invalid") + + reasons: list[str] = [] + ttft = _number(candidate.get("median_p99_ttft_ms")) + tpot = _number(candidate.get("median_p99_tpot_ms")) + compliance = _number(candidate.get("combined_compliance_ratio")) + + if ttft_sla_ms is not None and (ttft is None or ttft > ttft_sla_ms): + reasons.append( + "TTFT " + + ("missing" if ttft is None else f"{ttft:.1f} ms > {ttft_sla_ms:.1f} ms") + ) + if tpot_sla_ms is not None and (tpot is None or tpot > tpot_sla_ms): + reasons.append( + "TPOT " + + ("missing" if tpot is None else f"{tpot:.1f} ms > {tpot_sla_ms:.1f} ms") + ) + if minimum_compliance is not None and ( + compliance is None or compliance < minimum_compliance + ): + reasons.append( + "combined compliance " + + ( + "missing" + if compliance is None + else f"{compliance * 100:.2f}% < {minimum_compliance * 100:.2f}%" + ) + ) + + return "PASS" if not reasons else "; ".join(reasons) + + +def _scheduler_value(candidate: dict[str, Any], key: str) -> str: + value = candidate.get(key) + return "vLLM default" if value is None else str(value) + + +def _selected_settings( + parallel: dict[str, Any] | None, + concurrency: dict[str, Any] | None, + scheduler: dict[str, Any] | None, +) -> list[list[Any]]: + parallel_rec = (parallel or {}).get("recommended") or {} + concurrency_rec = (concurrency or {}).get("recommended") or {} + scheduler_rec = (scheduler or {}).get("recommended") or {} + max_num_seqs = ( + _scheduler_value(scheduler_rec, "max_num_seqs") + if scheduler is not None + else "-" + ) + max_num_batched_tokens = ( + _scheduler_value(scheduler_rec, "max_num_batched_tokens") + if scheduler is not None + else "-" + ) + return [ + ["tensor-parallel-size", parallel_rec.get("tensor_parallel_size")], + ["data-parallel-size", parallel_rec.get("data_parallel_size")], + ["max concurrency", concurrency_rec.get("max_concurrency")], + ["max-num-seqs", max_num_seqs], + ["max-num-batched-tokens", max_num_batched_tokens], + ] + + +def _stage_summary( + reports: dict[str, dict[str, Any] | None], +) -> list[list[Any]]: + rows = [] + for stage, label, _ in STAGES: + report = reports[stage] + if report is None: + rows.append([label, "not available", "-", "-", "-", "-"]) + continue + measured = report.get("measured") or {} + rows.append( + [ + label, + str(report.get("status", "unknown")).replace("_", " "), + _fmt(measured.get("mean_output_throughput")), + _fmt(measured.get("mean_request_goodput"), 3), + _fmt_int(measured.get("mean_p99_ttft_ms")), + _fmt(measured.get("mean_p99_tpot_ms")), + ] + ) + return rows + + +def _parallel_rows( + report: dict[str, Any] | None, + *, + ttft_sla_ms: float | None, + tpot_sla_ms: float | None, + minimum_compliance: float | None, +) -> list[list[Any]]: + rows = [] + for candidate in (report or {}).get("candidates", []): + rows.append( + [ + candidate.get("tensor_parallel_size"), + candidate.get("data_parallel_size"), + candidate.get("max_num_seqs"), + candidate.get("max_num_batched_tokens"), + _fmt(candidate.get("mean_output_throughput")), + _fmt_int(candidate.get("mean_p99_ttft_ms")), + _fmt(candidate.get("mean_p99_tpot_ms")), + _fmt_percent(candidate.get("combined_compliance_percent")), + _sla_label(candidate), + _sla_reason( + candidate, + ttft_sla_ms=ttft_sla_ms, + tpot_sla_ms=tpot_sla_ms, + minimum_compliance=minimum_compliance, + ), + ] + ) + return rows + + +def _concurrency_rows( + report: dict[str, Any] | None, + *, + ttft_sla_ms: float | None, + tpot_sla_ms: float | None, + minimum_compliance: float | None, +) -> list[list[Any]]: + candidates = (report or {}).get("candidates", []) + candidates = sorted(candidates, key=lambda item: item.get("max_concurrency", 0)) + rows = [] + for candidate in candidates: + rows.append( + [ + candidate.get("max_concurrency"), + _fmt(candidate.get("mean_output_throughput")), + _fmt(candidate.get("mean_request_throughput"), 3), + _fmt(candidate.get("mean_request_goodput"), 3), + _fmt_percent(candidate.get("combined_compliance_percent")), + _fmt_int(candidate.get("mean_p99_ttft_ms")), + _fmt(candidate.get("mean_p99_tpot_ms")), + _sla_label(candidate), + _sla_reason( + candidate, + ttft_sla_ms=ttft_sla_ms, + tpot_sla_ms=tpot_sla_ms, + minimum_compliance=minimum_compliance, + ), + ] + ) + return rows + + +def _scheduler_rows( + report: dict[str, Any] | None, + *, + ttft_sla_ms: float | None, + tpot_sla_ms: float | None, + minimum_compliance: float | None, +) -> list[list[Any]]: + rows = [] + for candidate in (report or {}).get("candidates", []): + rows.append( + [ + _scheduler_value(candidate, "max_num_seqs"), + _scheduler_value(candidate, "max_num_batched_tokens"), + _fmt(candidate.get("mean_output_throughput")), + _fmt(candidate.get("mean_request_goodput"), 3), + _fmt_int(candidate.get("mean_p99_ttft_ms")), + _fmt(candidate.get("mean_p99_tpot_ms")), + _fmt_percent(candidate.get("combined_compliance_percent")), + _sla_label(candidate), + _sla_reason( + candidate, + ttft_sla_ms=ttft_sla_ms, + tpot_sla_ms=tpot_sla_ms, + minimum_compliance=minimum_compliance, + ), + ] + ) + return rows + + +def _metadata(runs: dict[str, list[dict[str, Any]]]) -> dict[str, Any]: + for stage, _, _ in STAGES: + if runs[stage]: + first = runs[stage][0] + return { + "model": first.get("model_id", "unknown"), + "input_tokens": first.get("random_input_len"), + "output_tokens": first.get("random_output_len"), + "request_rate": first.get("request_rate"), + } + return { + "model": "unknown", + "input_tokens": None, + "output_tokens": None, + "request_rate": None, + } + + +def _observations( + reports: dict[str, dict[str, Any] | None], + runs: dict[str, list[dict[str, Any]]], +) -> list[tuple[str, str]]: + observations: list[tuple[str, str]] = [] + parallel = reports["parallel-layout"] + concurrency = reports["concurrency-tuning"] + scheduler = reports["runtime-tuning"] + + parallel_candidates = (parallel or {}).get("candidates", []) + if len(parallel_candidates) == 1: + observations.append( + ( + "warning", + ( + "Only one parallel-layout candidate completed. The result " + "validates that layout but does not establish an optimum." + ), + ) + ) + + if parallel and concurrency and runs["parallel-layout"]: + representative = runs["parallel-layout"][0].get("max_concurrency") + concurrency_values = { + item.get("max_concurrency") for item in concurrency.get("candidates", []) + } + parallel_passed = parallel_candidates and parallel_candidates[0].get( + "p99_sla_eligible" + ) + if representative not in concurrency_values and parallel_passed: + observations.append( + ( + "warning", + ( + f"Concurrency {representative} passed in the " + "parallel-layout stage but was not included in the " + "concurrency grid. Treat the automated concurrency " + "choice as conservative." + ), + ) + ) + + concurrency_candidates = (concurrency or {}).get("candidates", []) + first_failure = next( + ( + item + for item in sorted( + concurrency_candidates, + key=lambda candidate: candidate.get("max_concurrency", 0), + ) + if item.get("p99_sla_eligible") is False + ), + None, + ) + if first_failure and concurrency: + slo = concurrency.get("slo") or {} + ttft = _number(first_failure.get("mean_p99_ttft_ms")) + tpot = _number(first_failure.get("mean_p99_tpot_ms")) + ttft_slo = _number(slo.get("ttft_ms")) + tpot_slo = _number(slo.get("tpot_ms")) + if ( + ttft + and ttft_slo + and ttft > ttft_slo + and tpot is not None + and (tpot_slo is None or tpot <= tpot_slo) + ): + observations.append( + ( + "info", + ( + "TTFT is the first limiting SLA; TPOT remains within " + "its objective at the first failing concurrency." + ), + ) + ) + + scheduler_candidates = (scheduler or {}).get("candidates", []) + scheduler_throughputs = [ + number + for number in ( + _number(item.get("mean_output_throughput")) for item in scheduler_candidates + ) + if number is not None + ] + if len(scheduler_throughputs) > 1 and min(scheduler_throughputs) > 0: + spread = (max(scheduler_throughputs) / min(scheduler_throughputs) - 1) * 100 + if spread <= 3: + observations.append( + ( + "info", + ( + "Scheduler output throughput is tightly grouped: " + f"{spread:.2f}% total spread across measured candidates." + ), + ) + ) + + if not observations: + observations.append( + ("info", "No automatic coverage or SLA warnings were detected.") + ) + return observations + + +def _observation_html(observations: list[tuple[str, str]]) -> str: + return "".join( + f'
{escape(message)}
' + for level, message in observations + ) + + +def generate_report( + root: Path, + output: Path, + title: str, + ttft_sla_ms: float | None = None, + tpot_sla_ms: float | None = None, +) -> None: + reports = {stage: _read_object(root / filename) for stage, _, filename in STAGES} + if not any(reports.values()): + names = ", ".join(filename for _, _, filename in STAGES) + raise ValueError(f"No recommendation files found under {root}: {names}") + + ttft_sla_ms = _resolve_sla(reports, "ttft_ms", ttft_sla_ms) + tpot_sla_ms = _resolve_sla(reports, "tpot_ms", tpot_sla_ms) + minimum_compliance = _resolve_minimum_compliance(reports) + + runs = {stage: _summary_runs(root / "results" / stage) for stage, _, _ in STAGES} + meta = _metadata(runs) + observations = _observations(reports, runs) + concurrency = reports["concurrency-tuning"] or {} + scheduler = reports["runtime-tuning"] or {} + measured = scheduler.get("measured") or concurrency.get("measured") or {} + recommended_concurrency = (concurrency.get("recommended") or {}).get( + "max_concurrency" + ) + generated = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC") + + styles = """ +:root { color-scheme: light; --ink: #17324d; --blue: #146c94; + --teal: #2a9d8f; --green: #2e7d32; --amber: #d97706; + --red: #c0392b; --line: #d8dee4; --soft: #f3f5f7; } +* { box-sizing: border-box; } +body { margin: 0; background: #eef2f5; color: var(--ink); + font: 15px/1.45 Arial, sans-serif; } +main { max-width: 1180px; margin: 28px auto; padding: 0 20px 40px; } +header, section { background: white; border: 1px solid var(--line); + border-radius: 12px; margin-bottom: 18px; padding: 24px; } +h1 { margin: 0 0 6px; font-size: 30px; } +h2 { margin: 0 0 14px; font-size: 21px; } +h3 { margin: 20px 0 10px; font-size: 16px; color: var(--blue); } +.sub, .muted { color: #607486; } +.cards { display: grid; grid-template-columns: repeat(4, minmax(0, 1fr)); + gap: 12px; margin-top: 20px; } +.card { background: var(--soft); border-radius: 9px; padding: 14px; } +.card strong { display: block; font-size: 22px; } +.card span { color: #607486; font-size: 12px; } +.badge { border-radius: 99px; display: inline-block; font-size: 12px; + font-weight: 700; padding: 4px 9px; text-transform: uppercase; } +.badge.pass { background: #eaf5ec; color: var(--green); } +.badge.warn { background: #fff4e5; color: var(--amber); } +.notice { border-left: 5px solid var(--blue); margin: 9px 0; + padding: 11px 14px; background: #eaf4f8; } +.notice.warning { border-color: var(--amber); background: #fff4e5; } +.table-wrap { overflow-x: auto; } +table { width: 100%; border-collapse: collapse; font-size: 13px; } +th { background: var(--ink); color: white; text-align: left; } +th, td { border: 1px solid var(--line); padding: 8px 9px; white-space: nowrap; } +tbody tr:nth-child(even) { background: var(--soft); } +footer { color: #607486; font-size: 12px; text-align: center; } +@media (max-width: 760px) { .cards { grid-template-columns: 1fr 1fr; } } +@media print { body { background: white; } main { margin: 0; max-width: none; } + header, section { border: 0; break-inside: avoid; padding: 12px 0; } + .table-wrap { overflow: visible; } } +""" + + body = f""" + + + + +{escape(title)} + + +
+
+
vLLM recipe sweep report
+

{escape(title)}

+
{_text(meta["model"])} | {_text(meta["input_tokens"])} + input / {_text(meta["output_tokens"])} output tokens | + request rate {_text(meta["request_rate"])}
+
+
{_text(recommended_concurrency)} + recommended concurrency
+
{_fmt(measured.get("mean_output_throughput"))} + output tokens/s
+
{_fmt_int(measured.get("mean_p99_ttft_ms"))} ms + mean p99 TTFT
+
{_fmt(measured.get("mean_p99_tpot_ms"))} ms + mean p99 TPOT
+
{_fmt_int(ttft_sla_ms)} ms + TTFT SLA
+
{_fmt(tpot_sla_ms)} ms + TPOT SLA
+
{ + _fmt_percent( + minimum_compliance * 100 if minimum_compliance is not None else None + ) + } + required combined compliance
+
+
+
+

Recommended configuration

+ { + _table( + ["Setting", "Value"], + _selected_settings( + reports["parallel-layout"], + reports["concurrency-tuning"], + reports["runtime-tuning"], + ), + ) + } +
+
+

Stage summary

+ { + _table( + [ + "Stage", + "Status", + "Output tok/s", + "Goodput req/s", + "p99 TTFT ms", + "p99 TPOT ms", + ], + _stage_summary(reports), + ) + } +

Automatic observations

+ {_observation_html(observations)} +
+
+

Parallel-layout candidates { + _status_badge((reports["parallel-layout"] or {}).get("status")) + }

+ { + _table( + [ + "TP", + "DP", + "Max seqs", + "Batch tokens", + "Output tok/s", + "p99 TTFT ms", + "p99 TPOT ms", + "Compliance", + "SLA", + "SLA reason", + ], + _parallel_rows( + reports["parallel-layout"], + ttft_sla_ms=ttft_sla_ms, + tpot_sla_ms=tpot_sla_ms, + minimum_compliance=minimum_compliance, + ), + ) + } +
+
+

Concurrency candidates { + _status_badge((reports["concurrency-tuning"] or {}).get("status")) + }

+ { + _table( + [ + "Concurrency", + "Output tok/s", + "Req/s", + "Goodput req/s", + "Compliance", + "p99 TTFT ms", + "p99 TPOT ms", + "SLA", + "SLA reason", + ], + _concurrency_rows( + reports["concurrency-tuning"], + ttft_sla_ms=ttft_sla_ms, + tpot_sla_ms=tpot_sla_ms, + minimum_compliance=minimum_compliance, + ), + ) + } +
+
+

Scheduler candidates { + _status_badge((reports["runtime-tuning"] or {}).get("status")) + }

+ { + _table( + [ + "Max seqs", + "Batch tokens", + "Output tok/s", + "Goodput req/s", + "p99 TTFT ms", + "p99 TPOT ms", + "Compliance", + "SLA", + "SLA reason", + ], + _scheduler_rows( + reports["runtime-tuning"], + ttft_sla_ms=ttft_sla_ms, + tpot_sla_ms=tpot_sla_ms, + minimum_compliance=minimum_compliance, + ), + ) + } +
+
Generated {escape(generated)} from completed recommendation and +summary JSON files. Warmup results are excluded by the recommenders.
+
+ +""" + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(body, encoding="utf-8") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Generate a self-contained HTML report from recipe sweeps." + ) + parser.add_argument( + "--sweep-dir", + default=None, + help="Sweep package directory (default: beside this script).", + ) + parser.add_argument( + "--output", + default=None, + help="Output HTML path (default: /sweep-report.html).", + ) + parser.add_argument( + "--title", + default="vLLM Recipe Sweep", + help="Report title.", + ) + parser.add_argument( + "--ttft-sla-ms", + type=float, + default=DEFAULT_TTFT_SLA_MS, + help="Expected TTFT objective; must match every recommendation file.", + ) + parser.add_argument( + "--tpot-sla-ms", + type=float, + default=DEFAULT_TPOT_SLA_MS, + help="Expected TPOT objective; must match every recommendation file.", + ) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + for name in ("ttft_sla_ms", "tpot_sla_ms"): + value = getattr(args, name) + if value is not None and value <= 0: + option = "--" + name.replace("_", "-") + raise ValueError(f"{option} must be greater than zero.") + + script_dir = Path(__file__).absolute().parent + root = Path(args.sweep_dir).resolve() if args.sweep_dir else script_dir + output = Path(args.output).resolve() if args.output else root / "sweep-report.html" + generate_report( + root, + output, + args.title, + ttft_sla_ms=args.ttft_sla_ms, + tpot_sla_ms=args.tpot_sla_ms, + ) + print(f"Wrote HTML sweep report: {output}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/recipes/sweep_recommendation.py b/tools/recipes/sweep/sweep_recommendation.py similarity index 62% rename from tools/recipes/sweep_recommendation.py rename to tools/recipes/sweep/sweep_recommendation.py index 70efaad62523..1eebf754c33b 100755 --- a/tools/recipes/sweep_recommendation.py +++ b/tools/recipes/sweep/sweep_recommendation.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Select one measured scheduler recommendation from vLLM sweep results.""" +"""Select one measured parallel-layout or scheduler sweep recommendation.""" from __future__ import annotations @@ -24,6 +24,10 @@ DEFAULT_TTFT_SLA_MS: float | None = None DEFAULT_TPOT_SLA_MS: float | None = None DEFAULT_MINIMUM_COMPLIANCE: float = 0.99 +DEFAULT_THROUGHPUT_EQUIVALENCE_PERCENT: float = 1.0 +DEFAULT_RESULTS_DIR = "results/runtime-tuning" +DEFAULT_OUTPUT_CONFIG = "recommended-config.yml" +DEFAULT_OUTPUT_JSON = "recommendation.json" def _number(value: object) -> float | None: @@ -48,11 +52,75 @@ def _percentile_summary( return mean(numbers), median(numbers), max(numbers) -def _positive_int(row: dict[str, Any], key: str) -> int | None: - value = row.get(key) +def _scheduler_value(row: dict[str, Any], key: str) -> tuple[bool, int | None]: + if key not in row: + # Default-reference sweep candidates omit scheduler keys entirely. + return True, None + value = row[key] + if value is None: + return True, None if isinstance(value, bool) or not isinstance(value, int) or value <= 0: - return None - return value + return False, None + return True, value + + +def _parallel_value(row: dict[str, Any], key: str, default: int) -> tuple[bool, int]: + if key not in row: + return True, default + value = row[key] + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + return False, default + return True, value + + +def _scheduler_sort_value(value: int | None) -> int: + return -1 if value is None else value + + +def _scheduler_preference_key(value: object) -> tuple[int, int]: + """Prefer vLLM default, then the smaller explicit value on exact ties.""" + if value is None: + return (1, 0) + if isinstance(value, int) and not isinstance(value, bool): + return (0, -value) + return (0, 0) + + +def _has_fixed_parallel_layout(candidates: list[dict[str, Any]]) -> bool: + """Return whether candidates belong to one scheduler-only comparison.""" + layouts = { + (candidate["tensor_parallel_size"], candidate["data_parallel_size"]) + for candidate in candidates + } + return len(layouts) == 1 + + +def _select_near_equivalent_scheduler_candidate( + candidates: list[dict[str, Any]], + *, + throughput_equivalence_percent: float, +) -> dict[str, Any]: + """Prefer fewer scheduler overrides among near-equal configurations.""" + best_throughput = max( + candidate["mean_output_throughput"] for candidate in candidates + ) + throughput_floor = best_throughput * (1.0 - throughput_equivalence_percent / 100.0) + near_equivalent = [ + candidate + for candidate in candidates + if candidate["mean_output_throughput"] >= throughput_floor + ] + return max( + near_equivalent, + key=lambda candidate: ( + len(candidate["vllm_default_parameters"]), + candidate["combined_compliance_ratio"] or 0.0, + candidate["mean_request_goodput"] or 0.0, + candidate["mean_output_throughput"], + _scheduler_preference_key(candidate["max_num_batched_tokens"]), + _scheduler_preference_key(candidate["max_num_seqs"]), + ), + ) def _resolve(script_dir: Path, value: str | None) -> Path | None: @@ -80,21 +148,40 @@ def _aggregate_candidates( rows: list[dict[str, Any]], *, use_goodput: bool, + default_tensor_parallel_size: int = 1, + default_data_parallel_size: int = 1, ttft_sla_ms: float | None = None, tpot_sla_ms: float | None = None, minimum_compliance: float = DEFAULT_MINIMUM_COMPLIANCE, ) -> list[dict[str, Any]]: - grouped: dict[tuple[int, int], list[dict[str, Any]]] = defaultdict(list) + grouped: dict[tuple[int, int, int | None, int | None], list[dict[str, Any]]] = ( + defaultdict(list) + ) for row in rows: - seqs = _positive_int(row, "max_num_seqs") - batch = _positive_int(row, "max_num_batched_tokens") - if seqs is None or batch is None: + seqs_valid, seqs = _scheduler_value(row, "max_num_seqs") + batch_valid, batch = _scheduler_value(row, "max_num_batched_tokens") + tp_valid, tp = _parallel_value( + row, "tensor_parallel_size", default_tensor_parallel_size + ) + dp_valid, dp = _parallel_value( + row, "data_parallel_size", default_data_parallel_size + ) + if not seqs_valid or not batch_valid or not tp_valid or not dp_valid: continue - grouped[(seqs, batch)].append(row) + grouped[(tp, dp, seqs, batch)].append(row) candidates: list[dict[str, Any]] = [] - for (seqs, batch), runs in sorted(grouped.items()): + sorted_groups = sorted( + grouped.items(), + key=lambda item: ( + item[0][0], + item[0][1], + _scheduler_sort_value(item[0][2]), + _scheduler_sort_value(item[0][3]), + ), + ) + for (tp, dp, seqs, batch), runs in sorted_groups: failed_requests = sum(int(_number(run.get("failed")) or 0) for run in runs) output_throughput = _mean(runs, "output_throughput") request_throughput = _mean(runs, "request_throughput") @@ -150,9 +237,19 @@ def _aggregate_candidates( candidates.append( { + "tensor_parallel_size": tp, + "data_parallel_size": dp, "max_num_seqs": seqs, "max_num_batched_tokens": batch, "run_count": len(runs), + "vllm_default_parameters": [ + name + for name, value in ( + ("max_num_seqs", seqs), + ("max_num_batched_tokens", batch), + ) + if value is None + ], "failed_requests": failed_requests, "mean_request_goodput": request_goodput, "mean_request_throughput": request_throughput, @@ -188,6 +285,7 @@ def _select_candidate( candidates: list[dict[str, Any]], *, use_goodput: bool, + throughput_equivalence_percent: float = (DEFAULT_THROUGHPUT_EQUIVALENCE_PERCENT), ) -> tuple[dict[str, Any] | None, dict[str, Any], str]: valid = [candidate for candidate in candidates if candidate["valid"]] if not valid: @@ -204,8 +302,10 @@ def _select_candidate( candidate["mean_request_goodput"], candidate["combined_compliance_ratio"] or 0.0, candidate["mean_output_throughput"], - -candidate["max_num_batched_tokens"], - -candidate["max_num_seqs"], + -candidate["data_parallel_size"], + -candidate["tensor_parallel_size"], + _scheduler_preference_key(candidate["max_num_batched_tokens"]), + _scheduler_preference_key(candidate["max_num_seqs"]), ), ) if best_effort["mean_request_goodput"] <= 0: @@ -214,30 +314,46 @@ def _select_candidate( "the supplied TTFT/TPOT objectives." ) eligible = [candidate for candidate in valid if candidate["sla_eligible"]] - winner = ( - max( + if eligible and _has_fixed_parallel_layout(valid): + winner = _select_near_equivalent_scheduler_candidate( eligible, + throughput_equivalence_percent=throughput_equivalence_percent, + ) + else: + winner = ( + max( + eligible, + key=lambda candidate: ( + candidate["mean_output_throughput"], + candidate["combined_compliance_ratio"], + candidate["mean_request_goodput"], + -candidate["data_parallel_size"], + -candidate["tensor_parallel_size"], + _scheduler_preference_key(candidate["max_num_batched_tokens"]), + _scheduler_preference_key(candidate["max_num_seqs"]), + ), + ) + if eligible + else None + ) + else: + objective = "highest_mean_output_throughput" + if _has_fixed_parallel_layout(valid): + winner = _select_near_equivalent_scheduler_candidate( + valid, + throughput_equivalence_percent=throughput_equivalence_percent, + ) + else: + winner = max( + valid, key=lambda candidate: ( candidate["mean_output_throughput"], - candidate["combined_compliance_ratio"], - candidate["mean_request_goodput"], - -candidate["max_num_batched_tokens"], - -candidate["max_num_seqs"], + -candidate["data_parallel_size"], + -candidate["tensor_parallel_size"], + _scheduler_preference_key(candidate["max_num_batched_tokens"]), + _scheduler_preference_key(candidate["max_num_seqs"]), ), ) - if eligible - else None - ) - else: - objective = "highest_mean_output_throughput" - winner = max( - valid, - key=lambda candidate: ( - candidate["mean_output_throughput"], - -candidate["max_num_batched_tokens"], - -candidate["max_num_seqs"], - ), - ) best_effort = winner return winner, best_effort, objective @@ -250,6 +366,31 @@ def _load_config(path: Path) -> dict[str, Any]: return data +def _build_recommended_config( + initial_config: dict[str, Any], + winner: dict[str, Any], +) -> dict[str, Any]: + recommended_config = dict(initial_config) + recommended_config["tensor-parallel-size"] = winner["tensor_parallel_size"] + recommended_config["data-parallel-size"] = winner["data_parallel_size"] + for config_key, result_key in ( + ("max-num-seqs", "max_num_seqs"), + ("max-num-batched-tokens", "max_num_batched_tokens"), + ): + value = winner[result_key] + if value is None: + recommended_config.pop(config_key, None) + else: + recommended_config[config_key] = value + return recommended_config + + +def _format_scheduler_value(value: object) -> str: + if value is None: + return "vLLM default" + return str(value) + + def _write_config( path: Path, *, @@ -274,13 +415,13 @@ def _write_config( def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description=( - "Select one recommended max-num-seqs/max-num-batched-tokens pair " - "from generated vLLM sweep results." + "Select one recommended TP/DP layout and scheduler setting from " + "generated vLLM sweep results." ) ) parser.add_argument( "--results-dir", - default="results/runtime-tuning", + default=DEFAULT_RESULTS_DIR, help="Sweep experiment directory (default: results/runtime-tuning).", ) parser.add_argument( @@ -314,14 +455,24 @@ def parse_args() -> argparse.Namespace: "meet all supplied objectives (default: 0.99)." ), ) + parser.add_argument( + "--throughput-equivalence-percent", + type=float, + default=DEFAULT_THROUGHPUT_EQUIVALENCE_PERCENT, + help=( + "For a scheduler-only comparison, prefer fewer explicit scheduler " + "overrides when throughput is within this percentage of the best " + "eligible result (default: 1.0)." + ), + ) parser.add_argument( "--output-config", - default="recommended-config.yml", + default=DEFAULT_OUTPUT_CONFIG, help="Recommended config output path.", ) parser.add_argument( "--output-json", - default="recommendation.json", + default=DEFAULT_OUTPUT_JSON, help="Recommendation evidence output path.", ) return parser.parse_args() @@ -331,6 +482,10 @@ def main() -> int: args = parse_args() if not 0.0 < args.minimum_compliance <= 1.0: raise ValueError("--minimum-compliance must be greater than 0 and at most 1.") + if not 0.0 <= args.throughput_equivalence_percent < 100.0: + raise ValueError( + "--throughput-equivalence-percent must be at least 0 and less than 100." + ) script_dir = Path(__file__).resolve().parent config_path = _resolve(script_dir, args.config) @@ -350,6 +505,7 @@ def main() -> int: f"Sweep results were not found at {results_dir}. Run the sweep first." ) + initial_config = _load_config(config_path) rows = _load_runs(results_dir) if not rows: raise ValueError(f"No summary.json sweep results found under {results_dir}.") @@ -358,19 +514,20 @@ def main() -> int: candidates = _aggregate_candidates( rows, use_goodput=use_goodput, + default_tensor_parallel_size=int(initial_config.get("tensor-parallel-size", 1)), + default_data_parallel_size=int(initial_config.get("data-parallel-size", 1)), ttft_sla_ms=args.ttft_sla_ms, tpot_sla_ms=args.tpot_sla_ms, minimum_compliance=args.minimum_compliance, ) winner, best_effort, objective = _select_candidate( - candidates, use_goodput=use_goodput + candidates, + use_goodput=use_goodput, + throughput_equivalence_percent=args.throughput_equivalence_percent, ) - initial_config = _load_config(config_path) if winner is not None: - recommended_config = dict(initial_config) - recommended_config["max-num-seqs"] = winner["max_num_seqs"] - recommended_config["max-num-batched-tokens"] = winner["max_num_batched_tokens"] + recommended_config = _build_recommended_config(initial_config, winner) _write_config( output_config, source_path=config_path, @@ -384,8 +541,11 @@ def main() -> int: recommended = None if winner is not None: recommended = { + "tensor_parallel_size": winner["tensor_parallel_size"], + "data_parallel_size": winner["data_parallel_size"], "max_num_seqs": winner["max_num_seqs"], "max_num_batched_tokens": winner["max_num_batched_tokens"], + "vllm_default_parameters": winner["vllm_default_parameters"], } measured = { key: winner[key] @@ -413,19 +573,25 @@ def main() -> int: ), "selection_objective": objective, "minimum_compliance_ratio": args.minimum_compliance, + "throughput_equivalence_percent": args.throughput_equivalence_percent, "slo": { "ttft_ms": args.ttft_sla_ms, "tpot_ms": args.tpot_sla_ms, }, "initial": { + "tensor_parallel_size": initial_config.get("tensor-parallel-size"), + "data_parallel_size": initial_config.get("data-parallel-size", 1), "max_num_seqs": initial_config.get("max-num-seqs"), "max_num_batched_tokens": initial_config.get("max-num-batched-tokens"), }, "recommended": recommended, "measured": measured, "best_effort": { + "tensor_parallel_size": best_effort["tensor_parallel_size"], + "data_parallel_size": best_effort["data_parallel_size"], "max_num_seqs": best_effort["max_num_seqs"], "max_num_batched_tokens": best_effort["max_num_batched_tokens"], + "vllm_default_parameters": best_effort["vllm_default_parameters"], "mean_request_goodput": best_effort["mean_request_goodput"], "combined_compliance_ratio": best_effort["combined_compliance_ratio"], "p99_sla_eligible": best_effort["p99_sla_eligible"], @@ -445,8 +611,15 @@ def main() -> int: print("Recommended runtime configuration") print() - print(f" max-num-seqs: {winner['max_num_seqs']}") - print(f" max-num-batched-tokens: {winner['max_num_batched_tokens']}") + print(f" tensor-parallel-size: {winner['tensor_parallel_size']}") + print(f" data-parallel-size: {winner['data_parallel_size']}") + print( + " max-num-seqs: " + _format_scheduler_value(winner["max_num_seqs"]) + ) + print( + " max-num-batched-tokens: " + + _format_scheduler_value(winner["max_num_batched_tokens"]) + ) print() print(f"Selection objective: {objective}") if use_goodput: diff --git a/tools/recipes/sweep/sweep_visualization.py b/tools/recipes/sweep/sweep_visualization.py new file mode 100644 index 000000000000..4681ae5f7f37 --- /dev/null +++ b/tools/recipes/sweep/sweep_visualization.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Generate post-benchmark figures from completed recipe sweep results.""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import subprocess +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +@dataclass(frozen=True) +class PlotSpec: + var_x: str + var_y: str + fig_name: str + + +STAGE_PLOTS: dict[str, tuple[PlotSpec, ...]] = { + "parallel-layout": ( + PlotSpec("recipe_candidate", "output_throughput", "parallel_throughput"), + PlotSpec("recipe_candidate", "p99_ttft_ms", "parallel_p99_ttft"), + PlotSpec("recipe_candidate", "p99_tpot_ms", "parallel_p99_tpot"), + ), + "concurrency-tuning": ( + PlotSpec("max_concurrency", "output_throughput", "concurrency_throughput"), + PlotSpec("max_concurrency", "p99_ttft_ms", "concurrency_p99_ttft"), + PlotSpec("max_concurrency", "p99_tpot_ms", "concurrency_p99_tpot"), + PlotSpec("output_throughput", "p99_ttft_ms", "throughput_p99_ttft"), + ), + "runtime-tuning": ( + PlotSpec("recipe_candidate", "output_throughput", "scheduler_throughput"), + PlotSpec("recipe_candidate", "p99_ttft_ms", "scheduler_p99_ttft"), + PlotSpec("recipe_candidate", "p99_tpot_ms", "scheduler_p99_tpot"), + ), +} + + +def _serve_candidate(summary_path: Path, record: dict[str, Any]) -> str: + directory_name = summary_path.parent.name + prefix = "SERVE--" + separator = "-BENCH--" + if directory_name.startswith(prefix) and separator in directory_name: + return directory_name[len(prefix) :].split(separator, 1)[0] + + benchmark_name = record.get("_benchmark_name") + if isinstance(benchmark_name, str) and benchmark_name: + return benchmark_name + return "fixed_configuration" + + +def _load_stage_records(stage_dir: Path) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + for summary_path in sorted(stage_dir.rglob("summary.json")): + data = json.loads(summary_path.read_text(encoding="utf-8")) + if not isinstance(data, list): + continue + for record in data: + if not isinstance(record, dict): + continue + normalized = dict(record) + normalized["recipe_candidate"] = _serve_candidate( + summary_path, + normalized, + ) + records.append(normalized) + return records + + +def _has_numeric_metric(records: list[dict[str, Any]], key: str) -> bool: + return any( + isinstance(record.get(key), (int, float)) + and not isinstance(record.get(key), bool) + for record in records + ) + + +def _run_plot( + normalized_dir: Path, + figure_dir: Path, + spec: PlotSpec, + *, + dry_run: bool, +) -> None: + command = [ + "vllm", + "bench", + "sweep", + "plot", + str(normalized_dir), + "--fig-dir", + str(figure_dir.resolve()), + "--var-x", + spec.var_x, + "--var-y", + spec.var_y, + "--fig-name", + spec.fig_name, + ] + if dry_run: + command.append("--dry-run") + + environment = dict(os.environ) + environment.setdefault("MPLBACKEND", "Agg") + subprocess.run(command, check=True, env=environment) + + +def generate_figures( + results_dir: Path, + *, + selected_stages: list[str] | None = None, + dry_run: bool = False, +) -> list[Path]: + if shutil.which("vllm") is None: + raise RuntimeError( + "The vllm executable was not found. Run this analysis inside the " + "vLLM benchmark environment." + ) + + stages = selected_stages or list(STAGE_PLOTS) + discovered: list[Path] = [] + with tempfile.TemporaryDirectory(prefix="vllm-recipe-plot-") as temp_dir: + temporary_root = Path(temp_dir) + for stage in stages: + stage_dir = results_dir / stage + if not stage_dir.is_dir(): + continue + + records = _load_stage_records(stage_dir) + if not records: + continue + discovered.append(stage_dir) + + normalized_dir = temporary_root / stage + normalized_dir.mkdir(parents=True, exist_ok=True) + (normalized_dir / "summary.json").write_text( + json.dumps(records, indent=2) + "\n", + encoding="utf-8", + ) + + figure_dir = stage_dir / "figures" + plots = list(STAGE_PLOTS[stage]) + if stage == "concurrency-tuning" and _has_numeric_metric( + records, "request_goodput" + ): + plots.append( + PlotSpec( + "max_concurrency", + "request_goodput", + "concurrency_goodput", + ) + ) + + for spec in plots: + _run_plot( + normalized_dir, + figure_dir, + spec, + dry_run=dry_run, + ) + + if not discovered: + requested = ", ".join(stages) + raise ValueError( + f"No completed sweep summaries found under {results_dir} " + f"for stages: {requested}." + ) + return discovered + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Plot completed recipe sweep results without running a sweep." + ) + parser.add_argument( + "--results-dir", + default=None, + help="Results root (default: results beside this script).", + ) + parser.add_argument( + "--stage", + action="append", + choices=tuple(STAGE_PLOTS), + help="Plot only this stage; repeat to select multiple stages.", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Preview vLLM plot operations without writing figures.", + ) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + script_dir = Path(__file__).resolve().parent + results_dir = ( + Path(args.results_dir).resolve() + if args.results_dir is not None + else script_dir / "results" + ) + stages = generate_figures( + results_dir, + selected_stages=args.stage, + dry_run=args.dry_run, + ) + print() + if args.dry_run: + print("Visualization dry run completed for:") + else: + print("Wrote post-benchmark figures under:") + for stage_dir in stages: + print(f" {stage_dir / 'figures'}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/recipes/sweep_generation.py b/tools/recipes/sweep_generation.py deleted file mode 100644 index 33a5244a725e..000000000000 --- a/tools/recipes/sweep_generation.py +++ /dev/null @@ -1,358 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""Generate optional vLLM benchmark sweep files from one initial suggestion.""" - -from __future__ import annotations - -import json -import os -import shlex -from pathlib import Path -from typing import Any - -from runtime_tuning import WorkloadHints - - -def _positive_int(config: dict[str, Any], key: str) -> int: - value = config.get(key) - if not isinstance(value, int) or isinstance(value, bool) or value <= 0: - raise ValueError(f"Sweep generation requires a positive initial {key!r} value.") - return value - - -def _strict_lower_power_of_two(value: int) -> int: - if value <= 1: - return 1 - return 1 << ((value - 1).bit_length() - 1) - - -def _strict_upper_power_of_two(value: int) -> int: - return 1 << value.bit_length() - - -def validate_sweep_workload(workload: WorkloadHints) -> None: - required = { - "--input-tokens": workload.input_tokens, - "--output-tokens": workload.output_tokens, - "--concurrency": workload.concurrency, - } - missing = [name for name, value in required.items() if value is None] - if missing: - raise ValueError("--generate-sweep requires " + ", ".join(missing) + ".") - - -def build_serve_params(config: dict[str, Any]) -> list[dict[str, Any]]: - """Build a diverse, bounded sweep around the initial suggestion. - - The initial five-point sweep changed only one parameter at a time and - skipped the useful middle sequence count. Eight directed points cover the - batch-budget curve at full concurrency plus interactions at 3/4 and 1/2 of - the initial scheduler concurrency without paying for a full Cartesian grid. - """ - initial_seqs = _positive_int(config, "max-num-seqs") - initial_batch = _positive_int(config, "max-num-batched-tokens") - - minimum_batch = initial_seqs - if config.get("enable-chunked-prefill") is False: - max_model_len = config.get("max-model-len") - if ( - isinstance(max_model_len, int) - and not isinstance(max_model_len, bool) - and max_model_len > 0 - ): - minimum_batch = max(minimum_batch, max_model_len) - - lower_seqs = max(1, (initial_seqs + 1) // 2) - middle_seqs = max(lower_seqs, (3 * initial_seqs + 3) // 4) - lower_batch = max( - minimum_batch, - _strict_lower_power_of_two(initial_batch), - ) - smaller_batch = max( - minimum_batch, - _strict_lower_power_of_two(lower_batch), - ) - higher_batch = max( - minimum_batch, - _strict_upper_power_of_two(initial_batch), - ) - - candidates: list[dict[str, Any]] = [] - seen: set[tuple[int, int]] = set() - - def add(name: str, max_num_seqs: int, max_num_batched_tokens: int) -> None: - max_num_batched_tokens = max( - max_num_batched_tokens, - max_num_seqs, - ) - signature = (max_num_seqs, max_num_batched_tokens) - if signature in seen: - return - seen.add(signature) - candidates.append( - { - "_benchmark_name": name, - "max_num_seqs": max_num_seqs, - "max_num_batched_tokens": max_num_batched_tokens, - } - ) - - # Keep the exact initial suggestion as the measured baseline. Additional - # values exist only in the optional sweep package. - add("initial", initial_seqs, initial_batch) - add("smaller_batch_budget", initial_seqs, smaller_batch) - add("lower_batch_budget", initial_seqs, lower_batch) - add("higher_batch_budget", initial_seqs, higher_batch) - add("middle_seqs_lower_batch", middle_seqs, lower_batch) - add("middle_seqs_higher_batch", middle_seqs, higher_batch) - add("lower_seqs_lower_batch", lower_seqs, lower_batch) - add("lower_seqs_higher_batch", lower_seqs, higher_batch) - - return candidates - - -def build_bench_params(workload: WorkloadHints) -> list[dict[str, Any]]: - validate_sweep_workload(workload) - assert workload.input_tokens is not None - assert workload.output_tokens is not None - assert workload.concurrency is not None - - return [ - { - "_benchmark_name": "user_workload", - "random_input_len": workload.input_tokens, - "random_output_len": workload.output_tokens, - "max_concurrency": workload.concurrency, - } - ] - - -def _benchmark_models(config: dict[str, Any]) -> tuple[str, str]: - model = config.get("model") - if not isinstance(model, str) or not model: - raise ValueError("Sweep generation requires a model in config.yml.") - - served_model_name = config.get("served-model-name") - if isinstance(served_model_name, str) and served_model_name: - return served_model_name, model - if ( - isinstance(served_model_name, list) - and served_model_name - and isinstance(served_model_name[0], str) - ): - return served_model_name[0], model - - return model, model - - -def _relative_to(directory: Path, target: str) -> str: - return os.path.relpath(Path(target).resolve(), directory.resolve()) - - -def _write_json(path: Path, value: object) -> None: - path.write_text( - json.dumps(value, indent=2) + "\n", - encoding="utf-8", - ) - - -def _write_run_script( - path: Path, - *, - config_rel: str, - env_rel: str, - request_model: str, - tokenizer: str, - workload: WorkloadHints, -) -> None: - bench_parts = [ - "vllm bench serve", - "--backend vllm", - f"--model {shlex.quote(request_model)}", - f"--tokenizer {shlex.quote(tokenizer)}", - "--dataset-name random", - "--request-rate inf", - "--ignore-eos", - "--metric-percentiles 99", - ] - - goodput_pairs: list[str] = [] - if workload.ttft_sla_ms is not None: - goodput_pairs.append(f"ttft:{workload.ttft_sla_ms:g}") - if workload.tpot_sla_ms is not None: - goodput_pairs.append(f"tpot:{workload.tpot_sla_ms:g}") - if goodput_pairs: - bench_parts.append( - "--goodput " + " ".join(shlex.quote(pair) for pair in goodput_pairs) - ) - - bench_cmd = " ".join(bench_parts) - script = f"""#!/usr/bin/env bash -set -euo pipefail - -SCRIPT_DIR="$(cd -- "$(dirname -- "${{BASH_SOURCE[0]}}")" && pwd)" -CONFIG_PATH="${{SCRIPT_DIR}}/{config_rel}" -ENV_PATH="${{SCRIPT_DIR}}/{env_rel}" - -source "${{ENV_PATH}}" - -vllm bench sweep serve \ - --serve-cmd "vllm serve --config '${{CONFIG_PATH}}'" \ - --bench-cmd "{bench_cmd}" \ - --serve-params "${{SCRIPT_DIR}}/serve_params.json" \ - --bench-params "${{SCRIPT_DIR}}/bench_params.json" \ - --output-dir "${{SCRIPT_DIR}}/results" \ - --experiment-name runtime-tuning \ - "$@" -""" - path.write_text(script, encoding="utf-8") - path.chmod(path.stat().st_mode | 0o111) - - -def _write_recommend_script( - path: Path, - *, - config_rel: str, - env_rel: str, - workload: WorkloadHints, -) -> None: - template_path = Path(__file__).with_name("sweep_recommendation.py") - source = template_path.read_text(encoding="utf-8") - - replacements = { - "DEFAULT_CONFIG_PATH: str | None = None": ( - f"DEFAULT_CONFIG_PATH: str | None = {config_rel!r}" - ), - "DEFAULT_ENV_PATH: str | None = None": ( - f"DEFAULT_ENV_PATH: str | None = {env_rel!r}" - ), - "DEFAULT_TTFT_SLA_MS: float | None = None": ( - f"DEFAULT_TTFT_SLA_MS: float | None = {workload.ttft_sla_ms!r}" - ), - "DEFAULT_TPOT_SLA_MS: float | None = None": ( - f"DEFAULT_TPOT_SLA_MS: float | None = {workload.tpot_sla_ms!r}" - ), - } - for marker, replacement in replacements.items(): - if marker not in source: - raise ValueError(f"Recommender template marker not found: {marker}") - source = source.replace(marker, replacement, 1) - - path.write_text(source, encoding="utf-8") - path.chmod(path.stat().st_mode | 0o111) - - -def _write_guide(path: Path, workload: WorkloadHints) -> None: - sla_lines: list[str] = [] - if workload.ttft_sla_ms is not None: - sla_lines.append(f"- TTFT objective: `{workload.ttft_sla_ms:g} ms`") - if workload.tpot_sla_ms is not None: - sla_lines.append(f"- TPOT objective: `{workload.tpot_sla_ms:g} ms`") - - sla_text = "\n".join(sla_lines) - if sla_text: - sla_text = ( - "\n## Supplied latency objectives\n\n" - + sla_text - + "\n\nThe generated benchmark uses these values with vLLM " - "`--goodput`.\n" - ) - - content = f"""# Optional Runtime Tuning Sweep - -The generated `config.yml` is the **single initial suggestion** and can be -deployed directly. The sweep benchmarks nearby values for: - -- `max-num-seqs` -- `max-num-batched-tokens` - -## Run - -```bash -./run_sweep.sh --dry-run -./run_sweep.sh -./recommend.py -``` - -For a quick one-run experiment: - -```bash -rm -rf results/runtime-tuning -./run_sweep.sh --num-runs 1 -./recommend.py -``` - -Resume an interrupted sweep: - -```bash -./run_sweep.sh --resume -``` - -By default, vLLM benchmarks each parameter combination three times. -{sla_text} -## Outputs - -`recommend.py` writes: - -```text -recommended-config.yml -recommendation.json -``` - -With TTFT/TPOT objectives it requires duration-weighted combined compliance of -at least 99% plus median P99 TTFT/TPOT compliance, then selects highest mean -output-token throughput. If no candidate qualifies, it records a best-effort -candidate but does not write `recommended-config.yml`. Without latency -objectives it selects highest mean output-token throughput. Configurations with -failed requests are excluded. - -`recommended-config.yml` copies the initial configuration and changes only -`max-num-seqs` and `max-num-batched-tokens`. -""" - path.write_text(content, encoding="utf-8") - - -def write_sweep_files( - output_dir: str, - *, - config_path: str, - env_path: str, - config: dict[str, Any], - workload: WorkloadHints, -) -> list[Path]: - """Write an optional sweep package around the initial config suggestion.""" - validate_sweep_workload(workload) - - directory = Path(output_dir) - directory.mkdir(parents=True, exist_ok=True) - - serve_params = directory / "serve_params.json" - bench_params = directory / "bench_params.json" - run_script = directory / "run_sweep.sh" - recommend_script = directory / "recommend.py" - guide = directory / "SWEEP.md" - - config_rel = _relative_to(directory, config_path) - env_rel = _relative_to(directory, env_path) - - _write_json(serve_params, build_serve_params(config)) - _write_json(bench_params, build_bench_params(workload)) - request_model, tokenizer = _benchmark_models(config) - _write_run_script( - run_script, - config_rel=config_rel, - env_rel=env_rel, - request_model=request_model, - tokenizer=tokenizer, - workload=workload, - ) - _write_recommend_script( - recommend_script, - config_rel=config_rel, - env_rel=env_rel, - workload=workload, - ) - _write_guide(guide, workload) - - return [serve_params, bench_params, run_script, recommend_script, guide] From 0cc046087024bdb403b2032617ea0f713209a185 Mon Sep 17 00:00:00 2001 From: YiSheng5 Date: Fri, 9 Oct 2026 20:01:54 +0800 Subject: [PATCH 179/624] [XPU]Fix the accuracy issue when meet topk_ids=-1 on DP+EP scenarios (#57787) Signed-off-by: yisheng Co-authored-by: Kunshang Ji --- vllm/model_executor/layers/fused_moe/experts/xpu_moe.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py index ddd891b1c090..35fb91607830 100644 --- a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py @@ -175,6 +175,11 @@ def apply( gemm1_clamp_limit=self.gemm1_clamp_limit, ) assert self.fused_moe_impl is not None + if self.moe_config.ep_size > 1: + # Avoid topk_ids = -1, remap them to an + # expert not owned by this rank so they are skipped. + sentinel = global_num_experts - 1 if self.moe_config.ep_rank == 0 else 0 + topk_ids = torch.where(topk_ids < 0, sentinel, topk_ids) self.fused_moe_impl.apply( output=output, hidden_states=hidden_states, From c17fadcfdc80c9f8cfa97a7eec02a66d823989c4 Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Fri, 9 Oct 2026 20:16:35 +0800 Subject: [PATCH 180/624] [XPU] fix fused_input_norm for MM (#59865) Signed-off-by: Yan Ma --- .../intel_jobs/model_executor_intel.yaml | 11 +--------- .../intel_jobs/models_multimodal_intel.yaml | 3 ++- .../layers/fusion/mm_input_norm.py | 20 ++++++++----------- 3 files changed, 11 insertions(+), 23 deletions(-) diff --git a/.buildkite/intel_jobs/model_executor_intel.yaml b/.buildkite/intel_jobs/model_executor_intel.yaml index 5da075e235b6..d641c93e5a9e 100644 --- a/.buildkite/intel_jobs/model_executor_intel.yaml +++ b/.buildkite/intel_jobs/model_executor_intel.yaml @@ -28,16 +28,7 @@ steps: export VLLM_WORKER_MULTIPROC_METHOD=spawn && export PYTHONFAULTHANDLER=1 && cd tests && - pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --ignore="model_executor/model_loader/tensorizer_loader/test_tensorizer.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload" - --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormModule::test_matches_reference_above_cudnn_grid_limit" - --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype0-out_dtype0]" - --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype1-out_dtype1]" - --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype2-out_dtype2]" - --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype3-out_dtype3]" - --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype4-out_dtype4]" - --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype5-out_dtype5]" - --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormShapes" - --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormInputHandling"' + pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --ignore="model_executor/model_loader/tensorizer_loader/test_tensorizer.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload"' - label: Tensorizer Model Loader Tests (XPU) key: tensorizer-model-loader-tests-xpu diff --git a/.buildkite/intel_jobs/models_multimodal_intel.yaml b/.buildkite/intel_jobs/models_multimodal_intel.yaml index 061be1b136ed..89654ec3aedf 100644 --- a/.buildkite/intel_jobs/models_multimodal_intel.yaml +++ b/.buildkite/intel_jobs/models_multimodal_intel.yaml @@ -73,7 +73,8 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" && - pytest -v -s models/multimodal/generation/test_gemma_ragged_audio.py -m core_model' + pytest -v -s models/multimodal/generation/test_gemma_ragged_audio.py -m core_model + pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model' - label: "Multi-Modal Models (Standard) 4: other + whisper" key: multi-modal-models-standard-4-other-whisper diff --git a/vllm/model_executor/layers/fusion/mm_input_norm.py b/vllm/model_executor/layers/fusion/mm_input_norm.py index 3dc6974c84c6..5de95b401df5 100644 --- a/vllm/model_executor/layers/fusion/mm_input_norm.py +++ b/vllm/model_executor/layers/fusion/mm_input_norm.py @@ -313,19 +313,15 @@ def forward_cuda( def forward_xpu( self, pixel_values: torch.Tensor, visual_dtype: torch.dtype ) -> torch.Tensor: - """XPU fused custom kernel path. - - On XPU, fuse the whole rescale + normalise into a single custom - kernel. The eager path materializes an fp32 intermediate and then - casts back, which adds device-side compute that cancels the - bandwidth saving of transferring uint8 pixel_values. The fused - kernel reads uint8 directly and writes ``visual_dtype`` in one pass. + """XPU fused custom kernel path for uint8 and Triton kernel path + for others. """ - # The out-of-tree XPU kernel only supports the uint8 input that - # device-side normalisation guarantees; fail loudly otherwise. - assert pixel_values.dtype == torch.uint8, ( - f"xpu_fused_input_norm requires uint8 input, got {pixel_values.dtype}" - ) + if pixel_values.dtype != torch.uint8: + return self.forward_cuda(pixel_values, visual_dtype) + + # For uint8 input, falls to xpu_fused_input_norm from xpu kernels + import vllm._xpu_ops # noqa: F401 + return torch.ops.vllm.xpu_fused_input_norm( pixel_values, self.weight, self.bias, visual_dtype ) From 1ee69dde9cdfc949cfae896016b44357ca864a50 Mon Sep 17 00:00:00 2001 From: Suraj Sharan Date: Fri, 9 Oct 2026 16:39:58 +0400 Subject: [PATCH 181/624] [Bugfix][LoRA] Reject PEFT adapter features vLLM does not implement (#60225) Signed-off-by: suraj Co-authored-by: linitra24 --- .../serve/lora/test_lora_adapters.py | 1 + tests/lora/test_peft_helper.py | 39 ++++++++++++++++++ tests/lora/test_utils.py | 2 + vllm/lora/peft_helper.py | 40 +++++++++++++++++++ vllm/lora/utils.py | 2 +- 5 files changed, 83 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/serve/lora/test_lora_adapters.py b/tests/entrypoints/serve/lora/test_lora_adapters.py index 8a860dc6251a..484fb9c0f1e6 100644 --- a/tests/entrypoints/serve/lora/test_lora_adapters.py +++ b/tests/entrypoints/serve/lora/test_lora_adapters.py @@ -23,6 +23,7 @@ "is greater than max_lora_rank", ), ("test_dora", {"use_dora": True}, "does not yet support DoRA"), + ("test_pissa", {"init_lora_weights": "pissa"}, "init_lora_weights='pissa'"), ( "test_modules_to_save", {"modules_to_save": ["lm_head"]}, diff --git a/tests/lora/test_peft_helper.py b/tests/lora/test_peft_helper.py index 9fb52ea7275c..bb1246144b0d 100644 --- a/tests/lora/test_peft_helper.py +++ b/tests/lora/test_peft_helper.py @@ -24,6 +24,25 @@ ), ("test_rank_zero", {"r": 0}, "must be a positive integer"), ("test_rank_negative", {"r": -8}, "must be a positive integer"), + ("test_lora_bias", {"lora_bias": True}, "does not support LoRA bias"), + ("test_pissa", {"init_lora_weights": "pissa"}, "init_lora_weights='pissa'"), + ( + "test_pissa_niter", + {"init_lora_weights": "pissa_niter_4"}, + "init_lora_weights='pissa_niter_4'", + ), + ("test_olora", {"init_lora_weights": "olora"}, "init_lora_weights='olora'"), + ("test_corda", {"init_lora_weights": "corda"}, "init_lora_weights='corda'"), + ("test_loftq", {"init_lora_weights": "loftq"}, "init_lora_weights='loftq'"), + ( + "test_unknown_init", + {"init_lora_weights": "future_init"}, + "init_lora_weights='future_init'", + ), + ("test_alora", {"alora_invocation_tokens": [1, 2]}, "Activated LoRA"), + ("test_layer_replication", {"layer_replication": [[0, 2]]}, "layer_replication"), + ("test_bdlora", {"use_bdlora": {"nblocks": 2}}, "BD-LoRA"), + ("test_qalora", {"use_qalora": True}, "QALoRA"), ] @@ -114,3 +133,23 @@ def test_peft_helper_invalid_rank_direct(bad_rank: int): """ with pytest.raises(ValueError, match="must be a positive integer"): PEFTHelper(r=bad_rank, lora_alpha=16, target_modules=["q_proj"]) + + +@pytest.mark.parametrize( + "init_lora_weights", + [True, False, "gaussian", "eva", "orthogonal", "mica", "lora_ga"], +) +def test_peft_helper_init_lora_weights_supported(init_lora_weights): + """A saved adapter with these inits loads in PEFT as a plain LoRA.""" + lora_config = LoRAConfig(max_lora_rank=16, max_cpu_loras=3, max_loras=2) + PEFTHelper.from_dict( + { + "r": 8, + "lora_alpha": 16, + "target_modules": ["q_proj"], + "init_lora_weights": init_lora_weights, + "alora_invocation_tokens": None, + "layer_replication": None, + "use_bdlora": None, + } + ).validate_legal(lora_config) diff --git a/tests/lora/test_utils.py b/tests/lora/test_utils.py index 4a9ba66069d7..0b1be49f7e15 100644 --- a/tests/lora/test_utils.py +++ b/tests/lora/test_utils.py @@ -106,6 +106,8 @@ def test_parse_fine_tuned_lora_name_invalid(): fixture = { "base_model.weight", "base_model.model.weight", + "base_model.model.model.layers.9.mlp.down_proj.lora_A.bias", + "base_model.model.model.layers.9.mlp.down_proj.lora_B.bias", } for name in fixture: with pytest.raises(ValueError, match="unsupported LoRA weight"): diff --git a/vllm/lora/peft_helper.py b/vllm/lora/peft_helper.py index d84e56bece21..de05dea1cfc1 100644 --- a/vllm/lora/peft_helper.py +++ b/vllm/lora/peft_helper.py @@ -33,6 +33,14 @@ class PEFTHelper: use_rslora: bool = field(default=False) # True to use Weight-Decomposed Low-Rank Adaptation (DoRA, see: https://arxiv.org/abs/2402.09353) use_dora: bool = field(default=False) + # PEFT features that change the adapter's semantics and that vLLM does not + # implement. They are read only to reject such adapters. + lora_bias: bool = field(default=False) + init_lora_weights: bool | str = field(default=True) + alora_invocation_tokens: list[int] | None = field(default=None) + layer_replication: list[list[int]] | None = field(default=None) + use_bdlora: dict | None = field(default=None) + use_qalora: bool = field(default=False) # Extra vllm field, start with 'vllm_' to avoid conflict vllm_lora_scaling_factor: float = field(default=1.0) vllm_max_position_embeddings: int | None = field(default=False) @@ -52,8 +60,40 @@ def _validate_features(self) -> list[str]: ) if self.use_dora: error_msg.append("vLLM does not yet support DoRA.") + if self.lora_bias: + error_msg.append("vLLM does not support LoRA bias (lora_bias).") + if not self._init_is_supported(): + error_msg.append( + f"vLLM does not support init_lora_weights={self.init_lora_weights!r}." + " Inits such as PiSSA, OLoRA, CorDA and LoftQ modify the base model" + " weights when PEFT loads the adapter. Convert it to a regular LoRA" + " adapter with PEFT's path_initial_model_for_weight_conversion, or" + " set init_lora_weights to true if the served base model already" + " contains the modified weights." + ) + if self.alora_invocation_tokens: + error_msg.append("vLLM does not support Activated LoRA (aLoRA).") + if self.layer_replication: + error_msg.append("vLLM does not support layer_replication.") + if self.use_bdlora: + error_msg.append("vLLM does not support block-diagonal LoRA (BD-LoRA).") + if self.use_qalora: + error_msg.append("vLLM does not support QALoRA.") return error_msg + def _init_is_supported(self) -> bool: + # These inits only set the initial adapter weights, so a saved adapter + # loads in PEFT as a plain LoRA. lora_ga modifies the base weights only + # during training, when its gradients are attached. + init = self.init_lora_weights + return not isinstance(init, str) or init.lower() in ( + "gaussian", + "eva", + "orthogonal", + "mica", + "lora_ga", + ) + def __post_init__(self): if self.r <= 0: raise ValueError(f"LoRA rank `r` must be a positive integer, got {self.r}.") diff --git a/vllm/lora/utils.py b/vllm/lora/utils.py index 3a4009d28b23..1766c6de626d 100644 --- a/vllm/lora/utils.py +++ b/vllm/lora/utils.py @@ -205,7 +205,7 @@ def parse_fine_tuned_lora_name( parts = name.split(".") if (parts[-1] == "weight" or parts[-1] == "bias") and len(parts) >= 2: - if parts[-2] in ["lora_A", "lora_B"]: + if parts[-2] in ["lora_A", "lora_B"] and parts[-1] == "weight": new_name = ".".join(parts[start_index:-2]) return new_name, parts[-2] == "lora_A" # For modules_to_save in classification. From 9adba10f3f4bb3bb756f9de53fbf6820668157db Mon Sep 17 00:00:00 2001 From: "Li, Jiang" Date: Fri, 9 Oct 2026 20:48:32 +0800 Subject: [PATCH 182/624] [Bugfix] Fix pre-commit check (#60840) Signed-off-by: jiang1.li --- tools/recipes/sweep/sweep_html_report.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tools/recipes/sweep/sweep_html_report.py b/tools/recipes/sweep/sweep_html_report.py index 6d2271360862..bac10f9b4a28 100644 --- a/tools/recipes/sweep/sweep_html_report.py +++ b/tools/recipes/sweep/sweep_html_report.py @@ -8,7 +8,7 @@ import argparse import json -from datetime import datetime, timezone +from datetime import UTC, datetime from html import escape from pathlib import Path from typing import Any @@ -501,7 +501,7 @@ def generate_report( recommended_concurrency = (concurrency.get("recommended") or {}).get( "max_concurrency" ) - generated = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC") + generated = datetime.now(UTC).strftime("%Y-%m-%d %H:%M UTC") styles = """ :root { color-scheme: light; --ink: #17324d; --blue: #146c94; From 81c993658e0bbdbf9d30c92cfc52baa73853b242 Mon Sep 17 00:00:00 2001 From: Mikko Tukiainen Date: Fri, 9 Oct 2026 16:14:49 +0300 Subject: [PATCH 183/624] [ROCm][Perf] Fuse main QK-norm/RoPE/gate and KV-cache write into the AMD QSA prepare launch (#59732) Signed-off-by: Mikko Tukiainen Co-authored-by: Cursor --- .buildkite/test_areas/models_basic.yaml | 2 +- tests/models/qwen4_exp/test_qsa_amd.py | 72 ++++++-- ...indexer_amd.py => test_qsa_prepare_amd.py} | 85 ++++++++- vllm/config/compilation.py | 2 + vllm/models/qwen4_exp/amd/indexer_qsa.py | 82 ++++++--- .../{qsa_pre_indexer.py => qsa_prepare.py} | 167 ++++++++++++++++-- vllm/models/qwen4_exp/amd/qsa.py | 114 ++++++++++-- 7 files changed, 455 insertions(+), 69 deletions(-) rename tests/models/qwen4_exp/{test_qsa_pre_indexer_amd.py => test_qsa_prepare_amd.py} (80%) rename vllm/models/qwen4_exp/amd/ops/{qsa_pre_indexer.py => qsa_prepare.py} (75%) diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index a25c9a78a86c..fd21324a7ece 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -216,7 +216,7 @@ steps: depends_on: - image-build-amd commands: - - pytest -v -s models/qwen4_exp/test_qsa_amd.py models/qwen4_exp/test_qsa_pre_indexer_amd.py + - pytest -v -s models/qwen4_exp/test_qsa_amd.py models/qwen4_exp/test_qsa_prepare_amd.py # PLE and QSA import conflicting HC registrations in a shared process. - pytest -v -s models/qwen4_exp/test_ple.py diff --git a/tests/models/qwen4_exp/test_qsa_amd.py b/tests/models/qwen4_exp/test_qsa_amd.py index 15c490223f2d..c8ebf0190432 100644 --- a/tests/models/qwen4_exp/test_qsa_amd.py +++ b/tests/models/qwen4_exp/test_qsa_amd.py @@ -21,8 +21,8 @@ apply_qsa_rope, ) from vllm.models.qwen4_exp.amd.ops import qsa as qsa_ops -from vllm.models.qwen4_exp.amd.ops.qsa_pre_indexer import ( - qsa_pre_indexer, +from vllm.models.qwen4_exp.amd.ops.qsa_prepare import ( + qsa_prepare, supports_fused_pre_indexer, ) from vllm.platforms import current_platform @@ -127,6 +127,18 @@ def test_qsa_rmsnorm_uses_portable_implementation(default_vllm_config) -> None: torch.testing.assert_close(output, norm.forward_native(tensor)) +def _make_attn(*, use_fused: bool) -> SimpleNamespace: + """The main-attention fields the fused QSA prepare reads from its owner.""" + return SimpleNamespace( + use_fused_qsa_prepare=use_fused, + kv_cache=torch.zeros(1, 2, 16, 8), + q_norm=GemmaRMSNorm(4, eps=1e-6), + k_norm=GemmaRMSNorm(4, eps=1e-6), + _k_scale_float=1.0, + _v_scale_float=1.0, + ) + + def _make_indexer(*, use_fused: bool, num_tokens: int = 2) -> QSAIndexer: """A QSAIndexer whose unfused path runs for real, with no projection weights. @@ -184,29 +196,45 @@ def _make_indexer(*, use_fused: bool, num_tokens: int = 2) -> QSAIndexer: @pytest.mark.usefixtures("default_vllm_config") -def test_qsa_fused_pre_indexer_replaces_the_unfused_chain( +def test_qsa_fused_prepare_replaces_the_unfused_chain( monkeypatch: pytest.MonkeyPatch, ) -> None: indexer = _make_indexer(use_fused=True) + attn = _make_attn(use_fused=True) calls = [] - monkeypatch.setattr( - indexer_qsa, - "qsa_pre_indexer", - lambda *args, **kwargs: calls.append((args, kwargs)), - ) + main_outputs = (torch.empty(0), torch.empty(0)) + + def fake_prepare(*args, **kwargs): + calls.append((args, kwargs)) + return main_outputs + + monkeypatch.setattr(indexer_qsa, "qsa_prepare", fake_prepare) monkeypatch.setattr( QSAIndexer, "_update_and_compress", lambda *args, **kwargs: pytest.fail("fused path ran the unfused chain"), ) rows = torch.tensor([[3, 1, -1], [5, 2, 0]], dtype=torch.int32) - - actual = QSAIndexer.forward(indexer, torch.zeros(2, 8), torch.tensor([7, 8]), rows) + qkv = torch.zeros(4, 16) + slot_mapping = torch.arange(4) + gate_out = torch.empty(4, 1, 4) + + actual, actual_main = QSAIndexer.forward( + indexer, + torch.zeros(2, 8), + torch.tensor([7, 8]), + rows, + attn=attn, + qkv=qkv, + slot_mapping=slot_mapping, + gate_out=gate_out, + ) assert actual is rows + assert actual_main is main_outputs assert len(calls) == 1 args, kwargs = calls[0] - passed = inspect.signature(qsa_pre_indexer).bind(*args, **kwargs).arguments + passed = inspect.signature(qsa_prepare).bind(*args, **kwargs).arguments _, compressed_metadata = indexer.test_metadata # The kernel addresses the raw cache at its full width and writes the MRoPE # position tail itself, so it needs kv_cache, not the key_cache view. @@ -215,6 +243,17 @@ def test_qsa_fused_pre_indexer_replaces_the_unfused_chain( assert passed["k_work_metadata"] is compressed_metadata.k_work_metadata assert passed["compress_ratio"] == indexer.compress_ratio assert passed["rope_pos_offset"] == indexer.index_head_dim + # Main attention: only the real tokens, the vLLM cache viewed per token, + # and the gate landing in the caller's buffer. + assert passed["main_qkv"].shape[0] == 2 + assert passed["main_qkv"].data_ptr() == qkv.data_ptr() + assert passed["main_slot_mapping"].shape == (2,) + assert passed["main_kv_cache"].shape == (1, 16, 2, 8) + assert passed["main_kv_cache"].data_ptr() == attn.kv_cache.data_ptr() + assert passed["main_gate_out"].shape[0] == 2 + assert passed["main_gate_out"].data_ptr() == gate_out.data_ptr() + assert passed["main_q_norm_weight"] is attn.q_norm.weight + assert passed["main_k_norm_weight"] is attn.k_norm.weight @pytest.mark.usefixtures("default_vllm_config") @@ -224,7 +263,7 @@ def test_qsa_unsupported_config_keeps_the_unfused_chain( indexer = _make_indexer(use_fused=False) monkeypatch.setattr( indexer_qsa, - "qsa_pre_indexer", + "qsa_prepare", lambda *args, **kwargs: pytest.fail("unfused path ran the fused kernel"), ) updates = [] @@ -235,9 +274,16 @@ def test_qsa_unsupported_config_keeps_the_unfused_chain( ) rows = torch.tensor([[3, 1, -1], [5, 2, 0]], dtype=torch.int32) - actual = QSAIndexer.forward(indexer, torch.zeros(2, 8), torch.tensor([7, 8]), rows) + actual, main_outputs = QSAIndexer.forward( + indexer, + torch.zeros(2, 8), + torch.tensor([7, 8]), + rows, + attn=_make_attn(use_fused=False), + ) assert actual is rows + assert main_outputs is None assert len(updates) == 1 diff --git a/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py b/tests/models/qwen4_exp/test_qsa_prepare_amd.py similarity index 80% rename from tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py rename to tests/models/qwen4_exp/test_qsa_prepare_amd.py index 7698198923ce..718929720955 100644 --- a/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py +++ b/tests/models/qwen4_exp/test_qsa_prepare_amd.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Correctness tests for the AMD fused QSA pre-indexer.""" +"""Correctness tests for the AMD fused QSA prepare kernel.""" import pytest import torch @@ -11,9 +11,7 @@ qsa_compress_groups_with_ratio, qsa_store_cache_rows, ) -from vllm.models.qwen4_exp.amd.ops.qsa_pre_indexer import ( - qsa_pre_indexer, -) +from vllm.models.qwen4_exp.amd.ops.qsa_prepare import qsa_prepare from vllm.models.qwen4_exp.common.qsa_cache import ( canonical_qsa_rope_positions, circular_qsa_slot_mapping, @@ -52,6 +50,7 @@ def _reference_gemma_rmsnorm( RTOL = 1.6e-2 ATOL = 1e-2 MIXED_BATCH = ([260, 259, 138], [1, 1, 37], [8, 8, 8]) +MAIN_HQ, MAIN_HK, MAIN_D, MAIN_PAGE = 6, 2, 256, 16 def _make_block_table(block_counts): @@ -78,6 +77,76 @@ def assert_fp8_within_one_ulp(actual: torch.Tensor, expected: torch.Tensor) -> N assert bool(((code_diff <= 1) | (abs_diff <= 2**-8)).all()) +def _make_main_inputs(num_tokens: int) -> dict: + """Random main-attention arguments with a paged BF16 cache in vLLM's layout.""" + num_blocks = num_tokens // MAIN_PAGE + 2 + cache = torch.zeros( + num_blocks, + MAIN_HK, + MAIN_PAGE, + 2 * MAIN_D, + dtype=torch.bfloat16, + device="cuda", + ).transpose(1, 2) + slots = torch.randperm(num_blocks * MAIN_PAGE, device="cuda")[:num_tokens] + slots[-1] = -1 + norm_weights = torch.randn(2, MAIN_D, dtype=torch.bfloat16, device="cuda") * 0.2 + return dict( + main_qkv=torch.randn( + num_tokens, + 2 * (MAIN_HQ + MAIN_HK) * MAIN_D, + dtype=torch.bfloat16, + device="cuda", + ), + main_q_norm_weight=norm_weights[0], + main_k_norm_weight=norm_weights[1], + main_eps=EPS, + main_kv_cache=cache, + main_slot_mapping=slots, + main_k_scale=1.0, + main_v_scale=1.0, + # The AMD op hands the gate back through a buffer it does not own. + main_gate_out=torch.empty( + num_tokens, MAIN_HQ, MAIN_D, dtype=torch.bfloat16, device="cuda" + ), + ) + + +def _check_main_outputs(main: dict, q_out, gate_out, rope, positions) -> None: + """Compare with the AMD layer's unfused split, QK-norm, RoPE and cache write.""" + from vllm._custom_ops import reshape_and_cache_flash + + num_tokens = main["main_qkv"].shape[0] + q_gate, k, v = main["main_qkv"].split( + [2 * MAIN_HQ * MAIN_D, MAIN_HK * MAIN_D, MAIN_HK * MAIN_D], dim=-1 + ) + q, gate = torch.chunk(q_gate.view(num_tokens, MAIN_HQ, -1), 2, dim=-1) + q = _reference_gemma_rmsnorm( + q.reshape(-1, MAIN_D), main["main_q_norm_weight"], EPS + ).view(num_tokens, -1) + k = _reference_gemma_rmsnorm( + k.reshape(-1, MAIN_D), main["main_k_norm_weight"], EPS + ).view(num_tokens, -1) + q, k = rope(positions, q, k) + kv_cache = main["main_kv_cache"] + cache = torch.zeros_like(kv_cache) + key_cache, value_cache = cache.split(MAIN_D, dim=-1) + scale = torch.tensor(1.0, device="cuda") + reshape_and_cache_flash( + k.view(-1, MAIN_HK, MAIN_D), + v.view(-1, MAIN_HK, MAIN_D), + key_cache, + value_cache, + main["main_slot_mapping"], + "auto", + scale, + scale, + ) + torch.testing.assert_close(q_out.flatten(1), q, rtol=RTOL, atol=ATOL) + assert torch.equal(gate_out, gate) + torch.testing.assert_close(kv_cache, cache, rtol=RTOL, atol=ATOL) + + @requires_qsa_kernels @pytest.mark.usefixtures("default_vllm_config") @pytest.mark.parametrize("indexer_dtype", [torch.bfloat16, torch.float8_e4m3fn]) @@ -103,7 +172,7 @@ def assert_fp8_within_one_ulp(actual: torch.Tensor, expected: torch.Tensor) -> N pytest.param(True, True, True, 4, [4097], [4097], [0], id="tiled"), ], ) -def test_qsa_fused_pre_indexer_matches_unfused( +def test_qsa_fused_prepare_matches_unfused( indexer_dtype, mrope, is_2d_positions, @@ -269,7 +338,8 @@ def test_qsa_fused_pre_indexer_matches_unfused( k_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2 fused_query = torch.empty(num_tokens, HQ, D, dtype=indexer_dtype, device=device) - qsa_pre_indexer( + main = _make_main_inputs(num_tokens) + main_q_out, main_gate_out = qsa_prepare( projected_qk[:, : HQ * D], projected_qk[:, HQ * D :], positions, @@ -289,7 +359,10 @@ def test_qsa_fused_pre_indexer_matches_unfused( compress_ratio=CR, mrope_section=MROPE_SECTION if mrope else None, rope_pos_offset=ROPE_POS_OFFSET if cache_rope_positions else None, + **main, ) + assert main_gate_out is main["main_gate_out"] + _check_main_outputs(main, main_q_out, main_gate_out, rope, positions) unfused_query = projected_qk[:, : HQ * D].reshape(num_tokens, HQ, D) unfused_query = _reference_gemma_rmsnorm( diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index fc371d94b493..7a16e3d6a046 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -774,6 +774,8 @@ class CompilationConfig: # Qwen4Exp's AMD backend still uses these splitting ops. "vllm::qwen4_exp_ple_short_conv", "vllm::qwen4_exp_qsa_with_output", + # The fused prepare calls this op instead of the one above. + "vllm::qwen4_exp_qsa_prepare_with_output", "vllm::linear_attention", "vllm::qwen_gdn_attention_core", "vllm::qwen_gdn_attention_core_fused_norm_packed", diff --git a/vllm/models/qwen4_exp/amd/indexer_qsa.py b/vllm/models/qwen4_exp/amd/indexer_qsa.py index 162914f39e08..09cb39acf919 100644 --- a/vllm/models/qwen4_exp/amd/indexer_qsa.py +++ b/vllm/models/qwen4_exp/amd/indexer_qsa.py @@ -4,7 +4,7 @@ from __future__ import annotations -from typing import cast +from typing import TYPE_CHECKING, cast import torch from torch import nn @@ -23,7 +23,10 @@ QSAKeyStateCache, canonical_qsa_rope_positions, ) -from .ops.qsa_pre_indexer import qsa_pre_indexer, supports_fused_pre_indexer +from .ops.qsa_prepare import qsa_prepare, supports_fused_pre_indexer + +if TYPE_CHECKING: + from .qsa import Qwen4ExpQSAAttention def apply_qsa_rope( @@ -265,26 +268,31 @@ def _update_and_compress( position_rows, ) - def _fused_pre_indexer( + def _fused_prepare( self, projected_q: torch.Tensor, raw_keys: torch.Tensor, positions: torch.Tensor, raw_metadata: QSAForwardMetadata, compressed_metadata: QSAForwardMetadata, - ) -> torch.Tensor: - """Normalize and rotate Q, compress K, and write both caches in one launch. + attn: Qwen4ExpQSAAttention, + qkv: torch.Tensor, + slot_mapping: torch.Tensor, + gate_out: torch.Tensor | None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + """Prepare the indexer and ``attn``'s main attention in one launch. - Replaces ``project_qk``'s norm/RoPE half and the whole of - ``_update_and_compress``; the two must stay numerically interchangeable. + Normalizes and rotates the index Q, compresses K and writes both side + caches, which replaces ``project_qk``'s norm/RoPE half and the whole of + ``_update_and_compress``. The same launch applies ``attn``'s QK-norm and + RoPE, copies its gate and writes its K/V into ``attn.kv_cache``, + replacing ``_project_qkv_gate`` and ``do_kv_cache_update``. Both halves + must stay numerically interchangeable with what they replace. """ + num_tokens = raw_metadata.num_actual_tokens raw_key_cache = self.raw_key_cache - q = projected_q.new_empty( - raw_metadata.num_actual_tokens, - self.index_n_heads, - self.index_head_dim, - ) - qsa_pre_indexer( + q = projected_q.new_empty(num_tokens, self.index_n_heads, self.index_head_dim) + main_outputs = qsa_prepare( projected_q, raw_keys, positions, @@ -308,8 +316,17 @@ def _fused_pre_indexer( if raw_key_cache.cache_rope_positions else None ), + main_qkv=qkv[:num_tokens], + main_q_norm_weight=attn.q_norm.weight, + main_k_norm_weight=attn.k_norm.weight, + main_eps=attn.q_norm.variance_epsilon, + main_kv_cache=attn.kv_cache.transpose(1, 2), + main_slot_mapping=slot_mapping[:num_tokens], + main_k_scale=attn._k_scale_float, + main_v_scale=attn._v_scale_float, + main_gate_out=None if gate_out is None else gate_out[:num_tokens], ) - return q + return q, main_outputs def _select( self, @@ -336,13 +353,25 @@ def forward( hidden_states: torch.Tensor, positions: torch.Tensor, out: torch.Tensor | None = None, - ) -> torch.Tensor: - """Return fixed-width request-relative token indices padded with ``-1``.""" + *, + attn: Qwen4ExpQSAAttention, + qkv: torch.Tensor | None = None, + slot_mapping: torch.Tensor | None = None, + gate_out: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]: + """Return fixed-width request-relative token indices padded with ``-1``. + + With ``attn.use_fused_qsa_prepare``, the same launch also writes + ``attn``'s K/V into ``attn.kv_cache`` at ``slot_mapping`` and prepares + its Q and gate from ``qkv``, returned as the second element (None + otherwise). The gate is written into ``gate_out`` when given. ``qkv``, + ``slot_mapping`` and ``gate_out`` are only read in that mode. + """ metadata = self._metadata() if metadata is None: # Preserve step-0 indices when later MTP steps reuse the buffer. if self.skip_topk and out is not None: - return out + return out, None result = torch.full( (hidden_states.shape[0], self.output_width), -1, @@ -351,20 +380,27 @@ def forward( ) if out is not None: out.copy_(result) - return out - return result + return out, None + return result, None raw_metadata, compressed_metadata = metadata num_tokens = raw_metadata.num_actual_tokens hidden_states = hidden_states[:num_tokens] positions = positions[..., :num_tokens] - if self.use_fused_pre_indexer: + main_outputs: tuple[torch.Tensor, torch.Tensor] | None = None + if attn.use_fused_qsa_prepare: + if qkv is None or slot_mapping is None: + raise ValueError("fused QSA prepare requires qkv and slot_mapping") projected_q, raw_keys = self.project(hidden_states) - q = self._fused_pre_indexer( + q, main_outputs = self._fused_prepare( projected_q, raw_keys, positions, raw_metadata, compressed_metadata, + attn, + qkv, + slot_mapping, + gate_out, ) else: q, token_k = self.project_qk(hidden_states, positions) @@ -377,8 +413,8 @@ def forward( if self.skip_topk: if out is None: raise RuntimeError("QSA top-k reuse requires an output buffer") - return out - return self._select(q, compressed_metadata, out) + return out, main_outputs + return self._select(q, compressed_metadata, out), main_outputs __all__ = ["QSAIndexer", "apply_qsa_rope"] diff --git a/vllm/models/qwen4_exp/amd/ops/qsa_pre_indexer.py b/vllm/models/qwen4_exp/amd/ops/qsa_prepare.py similarity index 75% rename from vllm/models/qwen4_exp/amd/ops/qsa_pre_indexer.py rename to vllm/models/qwen4_exp/amd/ops/qsa_prepare.py index 6a4873882f2c..bc2fee3b3a47 100644 --- a/vllm/models/qwen4_exp/amd/ops/qsa_pre_indexer.py +++ b/vllm/models/qwen4_exp/amd/ops/qsa_prepare.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Fused QSA pre-indexer kernel for the AMD Qwen4Exp path. +"""Fused QSA prepare kernel for the AMD Qwen4Exp path. Started as a copy of the NVIDIA kernel (``nvidia/ops/qsa_prepare.py``) and is kept separate so either vendor can change its own copy without re-validating @@ -119,6 +119,27 @@ def _norm_rope( return tl.reshape(result, (TILE_T, TILE_H, D)) +@triton.jit +def _to_dst_dtype(x, dst, scale): + """Round to BF16 like the unfused path, then scale for an FP8 destination.""" + out_ty = dst.dtype.element_ty + x = x.to(tl.bfloat16) + if out_ty == tl.float8e4nv: + x = x.to(tl.float32) / scale + return x.to(out_ty) + + +@triton.jit +def _store_rotated(dst, y, o1, o2, scale): + """Store a normalized head whose first ``2 * len(o1)`` dims are rotated.""" + HALF: tl.constexpr = o1.shape[0] + dims = tl.arange(0, y.shape[0]) + rot = tl.arange(0, HALF) + tl.store(dst + dims, _to_dst_dtype(y, dst, scale), mask=dims >= 2 * HALF) + tl.store(dst + rot, _to_dst_dtype(o1, dst, scale)) + tl.store(dst + HALF + rot, _to_dst_dtype(o2, dst, scale)) + + @triton.jit( do_not_specialize=[ "num_tokens", @@ -127,7 +148,7 @@ def _norm_rope( "num_k_work", ] ) -def _qsa_pre_indexer_kernel( +def _qsa_prepare_kernel( q_ptr, q_stride_token, k_ptr, @@ -171,8 +192,86 @@ def _qsa_pre_indexer_kernel( CACHE_HAS_ROPE_POS: tl.constexpr, MROPE_H: tl.constexpr, MROPE_W: tl.constexpr, + main_qkv_ptr, + main_q_norm_weight_ptr, + main_k_norm_weight_ptr, + main_eps, + main_q_out_ptr, + main_gate_out_ptr, + main_cache_ptr, + main_cache_stride_block, + main_cache_stride_token, + main_cache_stride_head, + main_slots_ptr, + main_k_scale, + main_v_scale, + MAIN_HQ: tl.constexpr, + MAIN_HK: tl.constexpr, + MAIN_D: tl.constexpr, + MAIN_PAGE_SIZE: tl.constexpr, ): pid = tl.program_id(0) + num_index_work = num_k_work + tl.cdiv(num_tokens, TILE_T_Q) * tl.cdiv(HQ, TILE_H_Q) + if pid >= num_index_work: + # Main attention: one program per (token, Q or KV head) after the + # indexer work. RoPE covers the first D // 2 dims, the width of the + # shared cos/sin table. + main_pid = pid - num_index_work + token = (main_pid // (MAIN_HQ + MAIN_HK)).to(tl.int64) + head = main_pid % (MAIN_HQ + MAIN_HK) + HALF: tl.constexpr = D // 4 + dims = tl.arange(0, MAIN_D) + rot = tl.arange(0, HALF) + row = main_qkv_ptr + token * (2 * (MAIN_HQ + MAIN_HK) * MAIN_D) + is_k = head >= MAIN_HQ + kv_head = head - MAIN_HQ + if is_k: + src = row + (2 * MAIN_HQ + kv_head) * MAIN_D + weight_ptr = main_k_norm_weight_ptr + else: + src = row + 2 * head * MAIN_D + weight_ptr = main_q_norm_weight_ptr + x = tl.load(src + dims).to(tl.float32) + inv_rms = tl.rsqrt(tl.sum(x * x, axis=0) / MAIN_D + main_eps) + w = tl.load(weight_ptr + dims).to(tl.float32) + 1.0 + y = x * inv_rms * w + x1 = tl.load(src + rot).to(tl.float32) + x2 = tl.load(src + HALF + rot).to(tl.float32) + w1 = tl.load(weight_ptr + rot).to(tl.float32) + 1.0 + w2 = tl.load(weight_ptr + HALF + rot).to(tl.float32) + 1.0 + x1 = (x1 * inv_rms * w1).to(tl.bfloat16).to(tl.float32) + x2 = (x2 * inv_rms * w2).to(tl.bfloat16).to(tl.float32) + pos = tl.load(pos_ptr + token * pos_stride_token).to(tl.int64) + if IS_2D_POSITIONS: + pos_h = tl.load(pos_ptr + pos_stride_axis + token * pos_stride_token) + pos_w = tl.load(pos_ptr + 2 * pos_stride_axis + token * pos_stride_token) + is_h = (rot % 3 == 1) & (rot < 3 * MROPE_H) + is_w = (rot % 3 == 2) & (rot < 3 * MROPE_W) + pos = tl.where( + is_h, pos_h.to(tl.int64), tl.where(is_w, pos_w.to(tl.int64), pos) + ) + cos = tl.load(cos_sin_ptr + pos * (D // 2) + rot).to(tl.float32) + sin = tl.load(cos_sin_ptr + pos * (D // 2) + HALF + rot).to(tl.float32) + o1 = x1 * cos - x2 * sin + o2 = x2 * cos + x1 * sin + if is_k: + slot = tl.load(main_slots_ptr + token).to(tl.int64) + if slot >= 0: + dst = ( + main_cache_ptr + + (slot // MAIN_PAGE_SIZE) * main_cache_stride_block + + (slot % MAIN_PAGE_SIZE) * main_cache_stride_token + + kv_head * main_cache_stride_head + ) + _store_rotated(dst, y, o1, o2, main_k_scale) + v = tl.load(src + MAIN_HK * MAIN_D + dims) + tl.store(dst + MAIN_D + dims, _to_dst_dtype(v, dst, main_v_scale)) + else: + out = (token * MAIN_HQ + head) * MAIN_D + _store_rotated(main_q_out_ptr + out, y, o1, o2, None) + gate = tl.load(src + MAIN_D + dims) + tl.store(main_gate_out_ptr + out + dims, gate) + return # K work occupies the first programs; the remaining programs tile Q. This # keeps both paths in one launch while leaving their register shapes # independent. @@ -181,7 +280,7 @@ def _qsa_pre_indexer_kernel( num_head_tiles: tl.constexpr = tl.cdiv(HQ, TILE_H_Q) token_tile = q_pid // num_head_tiles head_tile = q_pid % num_head_tiles - tokens = token_tile * TILE_T_Q + tl.arange(0, TILE_T_Q) + tokens = (token_tile * TILE_T_Q + tl.arange(0, TILE_T_Q)).to(tl.int64) heads = head_tile * TILE_H_Q + tl.arange(0, TILE_H_Q) valid_tokens = tokens < num_tokens valid_heads = heads < HQ @@ -283,7 +382,8 @@ def _qsa_pre_indexer_kernel( & (source_tokens < num_tokens) ) current_base = ( - k_ptr + tl.maximum(source_tokens, 0)[:, None] * k_stride_token + k_ptr + + tl.maximum(source_tokens, 0).to(tl.int64)[:, None] * k_stride_token ) cached_base = ( state_cache_ptr @@ -397,7 +497,7 @@ def _qsa_pre_indexer_kernel( # One CTA per request commits only the suffix retained by the ring. num_state_rows = tl.minimum(query_len, STATE_SIZE) for state_offset in tl.range(0, num_state_rows): - token = query_end - num_state_rows + state_offset + token = (query_end - num_state_rows + state_offset).to(tl.int64) valid_token = ( (token >= query_start) & (token < query_end) & (token < num_tokens) ) @@ -447,7 +547,7 @@ def _qsa_pre_indexer_kernel( tl.store(tail + 2, pos_w.to(tl.int64), mask=valid_slot) -def qsa_pre_indexer( +def qsa_prepare( q: torch.Tensor, k: torch.Tensor, positions: torch.Tensor, @@ -468,11 +568,34 @@ def qsa_pre_indexer( compress_ratio: int, mrope_section: tuple[int, int, int] | None, rope_pos_offset: int | None, -) -> None: - """Normalize Q, compress K, then update the circular raw state.""" + main_qkv: torch.Tensor, + main_q_norm_weight: torch.Tensor, + main_k_norm_weight: torch.Tensor, + main_eps: float, + main_kv_cache: torch.Tensor, + main_slot_mapping: torch.Tensor, + main_k_scale: float, + main_v_scale: float, + main_gate_out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Normalize Q, compress K, then update the circular raw state. + + Also prepares the main attention (QK-norm/RoPE, gate copy, K/V cache write) + and returns its Q and gate. ``main_gate_out`` lets the gate land in a + caller-owned buffer, e.g. one that outlives an opaque custom op. + """ num_tokens = q.shape[0] + main_head_dim = main_kv_cache.shape[-1] // 2 + num_main_kv_heads = main_kv_cache.shape[2] + num_main_q_heads = main_qkv.shape[1] // (2 * main_head_dim) - num_main_kv_heads + main_q_out = main_qkv.new_empty(num_tokens, num_main_q_heads, main_head_dim) + if main_gate_out is None: + main_gate_out = torch.empty_like(main_q_out) + else: + assert main_gate_out.shape == main_q_out.shape + assert main_gate_out.is_contiguous() if num_tokens == 0: - return + return main_q_out, main_gate_out num_q_heads, head_dim = q_out.shape[1:] assert cos_sin_cache.shape[-1] * 2 == head_dim assert q.shape == (num_tokens, num_q_heads * head_dim) @@ -497,6 +620,9 @@ def qsa_pre_indexer( pos_stride_axis, pos_stride_token = 0, positions.stride(0) section = mrope_section if mrope_section is not None else (0, 0, 0) assert len(section) == 3 + qkv_width = 2 * (num_main_q_heads + num_main_kv_heads) * main_head_dim + assert main_qkv.shape == (num_tokens, qkv_width) and main_qkv.is_contiguous() + assert main_slot_mapping.shape == (num_tokens,) # Swept on MI355X (64-lane wavefronts) over TILE_T_Q 1-16, TILE_H_Q 1-4 and # num_warps 1/2/4, from 4 to 16384 tokens: no config beat these by more @@ -507,7 +633,8 @@ def qsa_pre_indexer( TILE_T_Q, TILE_H_Q = 2, 4 num_k_work = k_work_metadata.shape[0] num_q_work = triton.cdiv(num_tokens, TILE_T_Q) * triton.cdiv(num_q_heads, TILE_H_Q) - _qsa_pre_indexer_kernel[(num_k_work + num_q_work,)]( + num_main_work = num_tokens * (num_main_q_heads + num_main_kv_heads) + _qsa_prepare_kernel[(num_k_work + num_q_work + num_main_work,)]( q, q.stride(0), k, @@ -551,8 +678,26 @@ def qsa_pre_indexer( CACHE_HAS_ROPE_POS=cache_has_rope_pos, MROPE_H=section[1], MROPE_W=section[2], + main_qkv_ptr=main_qkv, + main_q_norm_weight_ptr=main_q_norm_weight, + main_k_norm_weight_ptr=main_k_norm_weight, + main_eps=main_eps, + main_q_out_ptr=main_q_out, + main_gate_out_ptr=main_gate_out, + main_cache_ptr=main_kv_cache, + main_cache_stride_block=main_kv_cache.stride(0), + main_cache_stride_token=main_kv_cache.stride(1), + main_cache_stride_head=main_kv_cache.stride(2), + main_slots_ptr=main_slot_mapping, + main_k_scale=main_k_scale, + main_v_scale=main_v_scale, + MAIN_HQ=num_main_q_heads, + MAIN_HK=num_main_kv_heads, + MAIN_D=main_head_dim, + MAIN_PAGE_SIZE=main_kv_cache.shape[1], num_warps=1, ) + return main_q_out, main_gate_out -__all__ = ["qsa_pre_indexer", "supports_fused_pre_indexer"] +__all__ = ["qsa_prepare", "supports_fused_pre_indexer"] diff --git a/vllm/models/qwen4_exp/amd/qsa.py b/vllm/models/qwen4_exp/amd/qsa.py index 3f03b293b4f0..ff01490bb190 100644 --- a/vllm/models/qwen4_exp/amd/qsa.py +++ b/vllm/models/qwen4_exp/amd/qsa.py @@ -116,8 +116,8 @@ def forward_qsa( self, layer: torch.nn.Module, query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, + key: torch.Tensor | None, + value: torch.Tensor | None, kv_cache: torch.Tensor, attn_metadata: FlashAttentionMetadata, output: torch.Tensor, @@ -294,6 +294,11 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.indexer", ) + self.use_fused_qsa_prepare = ( + self.attn_output_gate + and self.indexer.use_fused_pre_indexer + and self.head_dim & (self.head_dim - 1) == 0 + ) max_tokens = vllm_config.scheduler_config.max_num_batched_tokens self.register_buffer( "topk_indices_buffer", @@ -327,16 +332,22 @@ def _run_qsa( self, hidden_states: torch.Tensor, positions: torch.Tensor, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, + query: torch.Tensor | None, + key: torch.Tensor | None, + value: torch.Tensor | None, output: torch.Tensor, + qkv: torch.Tensor | None = None, + gate: torch.Tensor | None = None, ) -> None: + # query/key/value are None when the fused prepare runs inside the + # indexer launch, which then writes the gate into ``gate``. metadata = get_forward_context().attn_metadata if isinstance(metadata, list): metadata = metadata[0] if not isinstance(metadata, dict): output.zero_() + if gate is not None: + gate.zero_() return main_metadata = cast(FlashAttentionMetadata, metadata[self.layer_name]) if self.kv_cache.numel() == 0: @@ -349,10 +360,14 @@ def _run_qsa( ) if side_metadata.num_actual_tokens != num_tokens: raise RuntimeError("QSA main and side metadata token counts disagree") - selected = self.indexer( + selected, main_outputs = self.indexer( hidden_states, positions, self.topk_indices_buffer[:num_tokens], + attn=self, + qkv=qkv, + slot_mapping=main_metadata.slot_mapping, + gate_out=gate, ) if selected.shape != ( num_tokens, @@ -360,13 +375,22 @@ def _run_qsa( ): raise RuntimeError("QSA indexer returned an invalid selection shape") impl = cast(Qwen4ExpQSAFlashAttentionImpl, self.impl) - impl.do_kv_cache_update( - self, - key, - value, - self.kv_cache, - main_metadata.slot_mapping, - ) + if main_outputs is None: + assert key is not None and value is not None + impl.do_kv_cache_update( + self, + key, + value, + self.kv_cache, + main_metadata.slot_mapping, + ) + else: + query, _ = main_outputs + # The gate is applied to every row outside this op; padded rows + # must stay finite so they keep the zero output forward_qsa gives. + assert gate is not None + gate[num_tokens:].zero_() + assert query is not None impl.forward_qsa( self, query, @@ -384,13 +408,40 @@ def forward( hidden_states: torch.Tensor, ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) - q, k, v, gate = self._project_qkv_gate(qkv, positions) num_tokens = hidden_states.shape[0] + encoded_layer_name = _encode_layer_name(self.layer_name) + if self.use_fused_qsa_prepare: + # Norm/RoPE/gate and the K/V cache write happen inside the op. + attn_output = qkv.new_empty(num_tokens, self.num_heads, self.head_dim) + gate = torch.empty_like(attn_output) + if current_platform.opaque_attention_op(): + torch.ops.vllm.qwen4_exp_qsa_prepare_with_output( + hidden_states, + positions, + qkv, + attn_output, + gate, + encoded_layer_name, + ) + else: + qwen4_exp_qsa_prepare_with_output( + hidden_states, + positions, + qkv, + attn_output, + gate, + encoded_layer_name, + ) + flat_output = attn_output.view(num_tokens, -1) + flat_output = flat_output * torch.sigmoid(gate.view(num_tokens, -1)) + output, _ = self.o_proj(flat_output) + return output + + q, k, v, gate = self._project_qkv_gate(qkv, positions) query = q.view(num_tokens, self.num_heads, self.head_dim) key = k.view(num_tokens, self.num_kv_heads, self.head_dim) value = v.view(num_tokens, self.num_kv_heads, self.head_dim) attn_output = torch.empty_like(query) - encoded_layer_name = _encode_layer_name(self.layer_name) if current_platform.opaque_attention_op(): torch.ops.vllm.qwen4_exp_qsa_with_output( hidden_states, @@ -449,10 +500,43 @@ def qwen4_exp_qsa_with_output( ) +def qwen4_exp_qsa_prepare_with_output( + hidden_states: torch.Tensor, + positions: torch.Tensor, + qkv: torch.Tensor, + output: torch.Tensor, + gate: torch.Tensor, + layer_name: LayerNameType, +) -> None: + """Run the QSA transaction with the main QK-norm/RoPE/gate fused in.""" + layer_name = _resolve_layer_name(layer_name) + layer = get_forward_context().no_compile_layers[layer_name] + if not isinstance(layer, Qwen4ExpQSAAttention): + raise TypeError(f"{layer_name} is not a Qwen4Exp QSA owner") + layer._run_qsa( + hidden_states, + positions, + None, + None, + None, + output, + qkv=qkv, + gate=gate, + ) + + +direct_register_custom_op( + op_name="qwen4_exp_qsa_prepare_with_output", + op_func=qwen4_exp_qsa_prepare_with_output, + mutates_args=["output", "gate"], +) + + __all__ = [ "QSAIndexer", "Qwen4ExpQSAAttention", "Qwen4ExpQSAFlashAttentionBackend", "Qwen4ExpQSAFlashAttentionImpl", + "qwen4_exp_qsa_prepare_with_output", "qwen4_exp_qsa_with_output", ] From 3d52578e13454de0bb8352cac988347cdcc89928 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20P=C3=A9rez=20de=20Algaba?= <124347725+jperezdealgaba@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:23:35 +0200 Subject: [PATCH 184/624] [Security] Do not retain Phi-4 MM audio attention matrices (#60614) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Juan Pérez de Algaba Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/model_executor/test_phi4mm_attention.py | 28 +++++++++++++++++++ vllm/model_executor/models/phi4mm_audio.py | 8 +++--- vllm/model_executor/models/phi4mm_utils.py | 7 ++--- 3 files changed, 34 insertions(+), 9 deletions(-) create mode 100644 tests/model_executor/test_phi4mm_attention.py diff --git a/tests/model_executor/test_phi4mm_attention.py b/tests/model_executor/test_phi4mm_attention.py new file mode 100644 index 000000000000..51710f84bc47 --- /dev/null +++ b/tests/model_executor/test_phi4mm_attention.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.models.phi4mm_utils import MultiHeadedAttention + + +def test_eager_attention_does_not_retain_probability_matrix(): + module = MultiHeadedAttention( + n_head=2, + n_feat=16, + dropout_rate=0.0, + use_pt_scaled_dot_product_attention=False, + ) + module.eval() + batch, time, feat = 2, 8, 16 + x = torch.randn(batch, time, feat) + with torch.no_grad(): + out = module(x, x, x, None, None, None) + + assert out.shape == (batch, time, feat) + assert getattr(module, "attn", None) is None + + +def test_multiheaded_attention_defaults_to_scaled_dot_product(): + module = MultiHeadedAttention(n_head=2, n_feat=16, dropout_rate=0.0) + assert module.use_pt_scaled_dot_product_attention is True diff --git a/vllm/model_executor/models/phi4mm_audio.py b/vllm/model_executor/models/phi4mm_audio.py index 2a16b4ca7522..9daedcf38e7b 100644 --- a/vllm/model_executor/models/phi4mm_audio.py +++ b/vllm/model_executor/models/phi4mm_audio.py @@ -129,7 +129,7 @@ class ConformerEncoderLayer(nn.Module): default False. use_pt_scaled_dot_product_attention: bool, optional if set to True, use pytorch's scaled dot product attention - implementation in training. + implementation in training. Default: True. attn_group_sizes: int, optional the number of groups to use for attention, default 1 (Multi-Head Attention), @@ -163,7 +163,7 @@ def __init__( attention_glu_type: str = "swish", activation_checkpointing: str = "", export: bool = False, - use_pt_scaled_dot_product_attention: bool = False, + use_pt_scaled_dot_product_attention: bool = True, attn_group_sizes: int = 1, ) -> None: super().__init__() @@ -808,7 +808,7 @@ class ConformerEncoder(TransformerEncoderBase): default 4 use_pt_scaled_dot_product_attention: whether to use pytorch scaled dot product attention in training. - Default: False + Default: True nemo_conv_settings: dict, optional A dictionary of settings for NeMo Subsampling. default: None @@ -873,7 +873,7 @@ def __init__( # pylint: disable-all activation_checkpointing: str = "", relative_attention_bias_args: dict[str, Any] | None = None, time_reduction: int = 4, - use_pt_scaled_dot_product_attention: bool = False, + use_pt_scaled_dot_product_attention: bool = True, nemo_conv_settings: dict[str, Any] | None = None, conv2d_extra_padding: Literal["feat", "feat_time", "none", True] = "none", replication_pad_for_subsample_embedding: bool = False, diff --git a/vllm/model_executor/models/phi4mm_utils.py b/vllm/model_executor/models/phi4mm_utils.py index af54defd3fb0..db1b83f34388 100644 --- a/vllm/model_executor/models/phi4mm_utils.py +++ b/vllm/model_executor/models/phi4mm_utils.py @@ -1642,7 +1642,7 @@ class MultiHeadedAttention(nn.Module): NOTE: this will NOT be used in ONNX decoding due to a lack of support. In that case, we use the original attention implementation, which shows no regression. - default: False. + default: True. n_value: int, optional if set to values other than -1, use a different dimension for value. With the default value (i.e. -1), it is backward compatible. @@ -1666,7 +1666,7 @@ def __init__( attention_inner_dim: int = -1, glu_type: str = "swish", bias_in_glu: bool = True, - use_pt_scaled_dot_product_attention: bool = False, + use_pt_scaled_dot_product_attention: bool = True, n_value: int = -1, group_size: int = 1, ) -> None: @@ -1690,7 +1690,6 @@ def __init__( self.linear_v = nn.Linear(n_value, attention_inner_dim // group_size) self.linear_out = nn.Linear(attention_inner_dim // group_size, n_value) - self.attn = torch.jit.Attribute(None, Tensor | None) self.dropout = nn.Dropout(p=dropout_rate) self.dropout_rate = dropout_rate self.use_pt_scaled_dot_product_attention = use_pt_scaled_dot_product_attention @@ -1799,8 +1798,6 @@ def forward( attn = masked_softmax(scores, mask) # (batch, head, time1, time2) - self.attn = attn - p_attn = self.dropout(attn) x = torch.matmul(p_attn.to(v.dtype), v) # (batch, head, time1, d_k) if pos_v is not None: From fc39317392714b571868d7a49239db471f2dee0d Mon Sep 17 00:00:00 2001 From: Warren Deng <35615489+warrendeng@users.noreply.github.com> Date: Fri, 9 Oct 2026 07:25:41 -0700 Subject: [PATCH 185/624] [Bugfix][Kernel] Remove invalid NC from Kimi-K3 KDA autotune keys (#60446) Signed-off-by: Warren Deng Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py | 2 +- vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py index 0c600f3ba33d..bc13318d83cd 100644 --- a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py +++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py @@ -349,7 +349,7 @@ def chunk_kda_fwd_kernel_inter_solve_fused( for num_warps in [1, 2, 4, 8] for num_stages in [2, 3, 4] ], - key=["BK", "NC", "BT", "HV"], + key=["BK", "BT", "HV"], ) @triton.jit(do_not_specialize=["B", "T"]) def chunk_kda_fwd_kernel_intra_sub_chunk( diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py index 8b9aa86f9ed9..07b4a388936e 100644 --- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py @@ -451,7 +451,7 @@ def chunk_kda_fwd_kernel_inter_solve_fused( for num_warps in [1, 2, 4, 8] for num_stages in [2, 3, 4] ], - key=["BK", "NC", "BT", "HV"], + key=["BK", "BT", "HV"], ) @triton.jit(do_not_specialize=["B", "T"]) def chunk_kda_fwd_kernel_intra_sub_chunk( From a6b39f02620e23422113823810148dfda9876304 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Fri, 9 Oct 2026 09:32:47 -0500 Subject: [PATCH 186/624] [Bugfix] Clean up async client sockets after event loop closure (#57757) Signed-off-by: Taneem Ibrahim --- vllm/v1/engine/core_client.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 74520036a7c8..dccc8cba5368 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -538,10 +538,10 @@ def close_sockets_and_tasks(): with contextlib.suppress(Exception): task.cancel() - if loop is not None: + if loop is not None and not loop.is_closed(): if in_loop(loop): close_sockets_and_tasks() - elif not loop.is_closed(): + else: loop.call_soon_threadsafe(close_sockets_and_tasks) else: # Loop has been closed, try to clean up directly. From 970dc63478da41409dec009690f462e8e33d22f2 Mon Sep 17 00:00:00 2001 From: Abhijit Roy <50175934+Roy214@users.noreply.github.com> Date: Fri, 9 Oct 2026 20:11:38 +0530 Subject: [PATCH 187/624] [Docs] Add Giga-Embeddings bidirectional Qwen3 usage to embed docs (#60854) Signed-off-by: Abhijit --- docs/models/pooling_models/embed.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/models/pooling_models/embed.md b/docs/models/pooling_models/embed.md index 13a2e0d38147..4c33126977e9 100644 --- a/docs/models/pooling_models/embed.md +++ b/docs/models/pooling_models/embed.md @@ -50,7 +50,7 @@ You can compute pairwise similarity scores to build a similarity matrix using th | `ModernBertModel` | ModernBERT-based | `Alibaba-NLP/gte-modernbert-base`, etc. | ✅︎ | | | `NomicBertModel` | Nomic BERT | `nomic-ai/nomic-embed-text-v1`, `nomic-ai/nomic-embed-text-v2-moe`, `Snowflake/snowflake-arctic-embed-m-long`, etc. | | | | `Qwen2Model`C, `Qwen2ForCausalLM`C | Qwen2-based | `ssmits/Qwen2-7B-Instruct-embed-base` (see note), `Alibaba-NLP/gte-Qwen2-7B-instruct` (see note), etc. | ✅︎ | ✅︎ | -| `Qwen3Model`C, `Qwen3ForCausalLM`C | Qwen3-based | `Qwen/Qwen3-Embedding-0.6B`, etc. | ✅︎ | ✅︎ | +| `Qwen3Model`C, `Qwen3ForCausalLM`C | Qwen3-based | `Qwen/Qwen3-Embedding-0.6B`, `ai-sage/Giga-Embeddings-instruct-480M-0826`^, etc. | ✅︎ | ✅︎ | | `RobertaModel`, `RobertaForMaskedLM` | RoBERTa-based | `sentence-transformers/all-roberta-large-v1`, etc. | | | | `VoyageQwen3BidirectionalEmbedModel`C | Voyage Qwen3-based with bidirectional attention | `voyageai/voyage-4-nano`, etc. | ✅︎ | ✅︎ | | `XLMRobertaModel` | XLMRobertaModel-based | `BAAI/bge-m3` (see note), `intfloat/multilingual-e5-base`, `jinaai/jina-embeddings-v3` (see note), etc. | | | @@ -97,6 +97,7 @@ You can compute pairwise similarity scores to build a similarity matrix using th | `*ForConditionalGeneration`C, `*ForCausalLM`C, etc. | Generative models | \* | N/A | \* | \* | C Automatically converted into an embedding model via `--convert embed`. ([details](./README.md#model-conversion)) +^ You need to set the architecture name via `--hf-overrides` to match the one in vLLM. \* Feature support is the same as that of the original model. If your model is not in the above list, we will try to automatically convert the model using From 90ba2f34b5a29d85b8301672e5735676f78fa00c Mon Sep 17 00:00:00 2001 From: Greg Pereira Date: Fri, 9 Oct 2026 08:05:54 -0700 Subject: [PATCH 188/624] [Build] Find compatible published wheels for precompiled installs (#58871) Signed-off-by: greg pereira Signed-off-by: mgoin Co-authored-by: mgoin Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/misc.yaml | 3 + .../installation/gpu.cuda.inc.md | 17 +- setup.py | 109 +++++++++- .../test_precompiled_wheel.py | 195 ++++++++++++++++++ 4 files changed, 306 insertions(+), 18 deletions(-) create mode 100644 tests/standalone_tests/test_precompiled_wheel.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 9763b2ce7de4..93f59b50e52f 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -556,6 +556,8 @@ steps: - tests/multimodal - tests/renderers - tests/standalone_tests/lazy_imports.py + - tests/standalone_tests/test_precompiled_wheel.py + - setup.py - tests/reasoning - tools/pre_commit/check_test_tethering.py - tools/pre_commit/test_tethering_allowlist.txt @@ -572,6 +574,7 @@ steps: commands: - pip install cmake # test_cmake_utils.py shells out to cmake; the CPU image has none - python3 standalone_tests/lazy_imports.py + - python3 standalone_tests/test_precompiled_wheel.py - pytest -v -s device_allocator - pytest -v -s test_uvicorn_logging.py - pytest -v -s test_attention_backend_registry.py diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md index 6ad581ac97d1..f125a189fcd2 100644 --- a/docs/getting_started/installation/gpu.cuda.inc.md +++ b/docs/getting_started/installation/gpu.cuda.inc.md @@ -117,18 +117,21 @@ If you need to recompile the `vllm-rs` Rust frontend binary, you can rebuild and This will install the required Rust toolchain if needed, build the binary, and place it in `vllm/vllm-rs`. -In case you see an error about wheel not found when running the above command, it might be because the commit you based on in the `main` branch was just merged and its precompiled wheel is not available yet. You can wait around an hour and retry, or set `VLLM_PRECOMPILED_WHEEL_COMMIT=nightly` to automatically select the most recent already-built commit on `main`. +Wheels are published an hour or two after a commit lands on `main`. If the wheel for +your merge-base with upstream `main` is not available yet, the installer uses the +nearest older commit that has one, searching up to 20 commits back. It stops with an +error if compiled sources or build configuration (`csrc/`, `cmake/`, `CMakeLists.txt`, +`pyproject.toml`, `vllm/_custom_ops.py`) changed between that commit and your +merge-base, and warns if `rust/` or `setup.py` did. This is a source-level check, not +a guarantee of binary compatibility. -```bash -export VLLM_PRECOMPILED_WHEEL_COMMIT=nightly -export VLLM_USE_PRECOMPILED=1 -uv pip install --editable . -``` +Setting `VLLM_PRECOMPILED_WHEEL_LOCATION` or a full SHA in +`VLLM_PRECOMPILED_WHEEL_COMMIT` skips this search and check. There are more environment variables to control the behavior of Python-only build: - `VLLM_PRECOMPILED_WHEEL_LOCATION`: specify the exact wheel URL or local file path of a pre-compiled wheel to use. All other logic to find the wheel will be skipped. -- `VLLM_PRECOMPILED_WHEEL_COMMIT`: override the commit hash to download the pre-compiled wheel. It can be `nightly` to use the last **already built** commit on the main branch. +- `VLLM_PRECOMPILED_WHEEL_COMMIT`: override the commit to download the pre-compiled wheel from. Must be a full 40-character SHA. - `VLLM_PRECOMPILED_WHEEL_VARIANT`: specify the variant subdirectory to use on the nightly index, e.g., `cu129`, `cu130`, `cpu`. If not specified, the variant is auto-detected based on your system's CUDA version (from PyTorch or nvidia-smi). You can also set `VLLM_MAIN_CUDA_VERSION` to override auto-detection. You can find more information about vLLM's wheels in [Install the latest code](#install-the-latest-code). diff --git a/setup.py b/setup.py index 16a3d35ecad6..b173b9dca8fd 100644 --- a/setup.py +++ b/setup.py @@ -896,6 +896,93 @@ def determine_wheel_url_rocm() -> tuple[str, str | None]: print(f"Using ROCm precompiled wheel: {wheel_url}") return wheel_url, download_filename + # Publication lags main by roughly 5-10 commits; anything much older is + # unlikely to still match the checkout's compiled sources. + WHEEL_SEARCH_DEPTH = 20 + # Changes here make an older wheel's extensions unusable. + WHEEL_BLOCKING_PATHS = ( + "csrc", + "cmake", + "CMakeLists.txt", + "pyproject.toml", + "vllm/_custom_ops.py", + ) + # Changes here are usually harmless for an older wheel. + WHEEL_WARNING_PATHS = ("rust", "setup.py") + + @staticmethod + def find_published_wheel_commit( + base_commit: str, variant: str | None, arch: str + ) -> str: + """Find the nearest ancestor of base_commit that has a published wheel. + + Raises: + ValueError: If no wheel is published within WHEEL_SEARCH_DEPTH + commits, or compiled sources changed since the nearest one. + + """ + from urllib.error import HTTPError + + utils = precompiled_wheel_utils + candidates = subprocess.check_output( + [ + "git", + "rev-list", + "--first-parent", + f"--max-count={utils.WHEEL_SEARCH_DEPTH}", + base_commit, + ], + text=True, + ).splitlines() + for distance, commit in enumerate(candidates): + try: + wheels, _ = utils.fetch_metadata_for_variant(commit, variant) + except HTTPError as err: + if err.code != 404: + raise + continue + if not any( + wheel.get("package_name") == "vllm" + and arch in wheel.get("platform_tag", "") + for wheel in wheels + ): + continue + if distance == 0: + return commit + + diff = ["git", "diff", "--name-only", commit, base_commit, "--"] + blocking, warning = ( + subprocess.check_output([*diff, *paths], text=True).strip() + for paths in (utils.WHEEL_BLOCKING_PATHS, utils.WHEEL_WARNING_PATHS) + ) + if blocking: + raise ValueError( + f"No precompiled wheel is published yet for {base_commit}, and " + f"the nearest one ({commit}, {distance} commits older) was " + f"built from different compiled sources:\n{blocking}\n" + "Wait for wheel publication, rebase onto a commit with a " + "published wheel, build from source (unset " + "VLLM_USE_PRECOMPILED), or set VLLM_PRECOMPILED_WHEEL_COMMIT " + "to a full commit SHA to skip this check." + ) + if warning: + logger.warning( + "Precompiled wheel %s predates changes to:\n%s\n" + "The Rust frontend and packaged files may be stale.", + commit, + warning, + ) + print( + f"No precompiled wheel is published yet for {base_commit}; using " + f"{commit} ({distance} commits older)" + ) + return commit + raise ValueError( + f"No published precompiled wheel for variant {variant} and " + f"architecture {arch} within {utils.WHEEL_SEARCH_DEPTH} commits of " + f"{base_commit}. Wait for wheel publication or build vLLM from source." + ) + @staticmethod def determine_wheel_url() -> tuple[str, str | None]: """Try to determine the precompiled wheel URL or path to use. @@ -906,8 +993,8 @@ def determine_wheel_url() -> tuple[str, str | None]: or CUDA variant selected from VLLM_MAIN_CUDA_VERSION, torch, or nvidia-smi If downloading from the nightly repo, the commit can be specified via - VLLM_PRECOMPILED_WHEEL_COMMIT; otherwise, the head commit in the main branch - is used. + VLLM_PRECOMPILED_WHEEL_COMMIT; otherwise, local checkouts search backwards + from their merge-base with upstream main for a compatible published wheel. """ wheel_location = os.getenv("VLLM_PRECOMPILED_WHEEL_LOCATION", None) if wheel_location is not None: @@ -933,6 +1020,14 @@ def determine_wheel_url() -> tuple[str, str | None]: ", trying to fetch base commit in main branch" ) commit = precompiled_wheel_utils.get_base_commit_in_main_branch() + if ( + envs.VLLM_USE_PRECOMPILED + and not envs.VLLM_DOCKER_BUILD_CONTEXT + and commit != "nightly" + ): + commit = precompiled_wheel_utils.find_published_wheel_commit( + commit, variant, arch + ) print(f"Using precompiled wheel commit {commit} with variant {variant}") download_filename = None try: @@ -1149,17 +1244,9 @@ def get_base_commit_in_main_branch() -> str: ["git", "fetch", "https://github.com/vllm-project/vllm", "main"] ) - # Then get the commit hash of the current branch that is the same as - # the upstream main commit. - current_branch = ( - subprocess.check_output(["git", "branch", "--show-current"]) - .decode("utf-8") - .strip() - ) - base_commit = ( subprocess.check_output( - ["git", "merge-base", f"{upstream_main_commit}", current_branch] + ["git", "merge-base", upstream_main_commit, "HEAD"] ) .decode("utf-8") .strip() diff --git a/tests/standalone_tests/test_precompiled_wheel.py b/tests/standalone_tests/test_precompiled_wheel.py new file mode 100644 index 000000000000..883a6cea842f --- /dev/null +++ b/tests/standalone_tests/test_precompiled_wheel.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Exercise setup-time wheel selection without importing torch or running setup.""" + +import ast +import logging +import os +import subprocess +import tempfile +import unittest +from email.message import Message +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from unittest.mock import patch +from urllib.error import HTTPError, URLError + +# setup.py executes setup() on import; load only the wheel utility class. +setup_path = Path(__file__).resolve().parents[2] / "setup.py" +node = next( + node + for node in ast.parse(setup_path.read_text()).body + if isinstance(node, ast.ClassDef) and node.name == "precompiled_wheel_utils" +) +namespace: dict[str, Any] = { + "subprocess": subprocess, + "os": os, + "logger": logging.getLogger("setup"), + "envs": SimpleNamespace(VLLM_DOCKER_BUILD_CONTEXT=False, VLLM_USE_PRECOMPILED=True), +} +exec( + compile(ast.Module(body=[node], type_ignores=[]), str(setup_path), "exec"), + namespace, +) +wheel_utils = namespace["precompiled_wheel_utils"] +WHEEL = { + "package_name": "vllm", + "platform_tag": "manylinux_x86_64", + "path": "wheel.whl", + "filename": "wheel.whl", +} +NOT_FOUND = HTTPError("url", 404, "missing", Message(), None) + + +class TestPublishedWheelSearch(unittest.TestCase): + """Walk back from the merge-base in a real Git repo; only the index is mocked.""" + + def setUp(self): + tmp = tempfile.TemporaryDirectory() + self.addCleanup(tmp.cleanup) + self.repo = Path(tmp.name) + self.addCleanup(os.chdir, os.getcwd()) + os.chdir(self.repo) + self.git("init", "-q") + self.published = {self.commit("vllm/a.py")} + + def git(self, *args: str) -> str: + identity = ["-c", "user.name=t", "-c", "user.email=t@t"] + return subprocess.check_output(["git", *identity, *args], text=True).strip() + + def commit(self, path: str) -> str: + file = self.repo / path + file.parent.mkdir(parents=True, exist_ok=True) + file.write_text(file.read_text() + "x" if file.exists() else "x") + self.git("add", "-A") + self.git("commit", "-q", "-m", path) + return self.git("rev-parse", "HEAD") + + def search(self, base: str, arch: str = "x86_64") -> str: + def fetch(commit, variant): + if commit not in self.published: + raise NOT_FOUND + return [WHEEL], "url" + + with patch.object(wheel_utils, "fetch_metadata_for_variant", fetch): + return wheel_utils.find_published_wheel_commit(base, "cu130", arch) + + def test_published_base_is_used_despite_local_native_edits(self): + base = self.commit("csrc/ops.cpp") + self.published.add(base) + self.commit("csrc/ops.cpp") + (self.repo / "pyproject.toml").write_text("dirty") + self.assertEqual(self.search(base), base) + + def test_publication_lag_uses_nearest_published_ancestor(self): + published = self.commit("vllm/b.py") + self.published.add(published) + self.commit("vllm/c.py") + self.assertEqual(self.search(self.commit("vllm/d.py")), published) + + def test_compiled_changes_since_wheel_are_rejected(self): + (published,) = self.published + for path in wheel_utils.WHEEL_BLOCKING_PATHS: + with self.subTest(path=path): + self.git("reset", "-q", "--hard", published) + file = path if "." in path else f"{path}/file" + with self.assertRaisesRegex(ValueError, file): + self.search(self.commit(file)) + + def test_rust_changes_since_wheel_only_warn(self): + (published,) = self.published + base = self.commit("rust/src/main.rs") + with self.assertLogs("setup", "WARNING") as logs: + self.assertEqual(self.search(base), published) + self.assertIn("rust/src/main.rs", logs.output[0]) + + def test_wrong_architecture_is_not_a_match(self): + with self.assertRaisesRegex(ValueError, "No published precompiled wheel"): + self.search("HEAD", arch="aarch64") + + def test_search_stops_at_depth_limit(self): + with patch.object(wheel_utils, "WHEEL_SEARCH_DEPTH", 2): + self.commit("vllm/b.py") + with self.assertRaisesRegex(ValueError, "within 2 commits"): + self.search(self.commit("vllm/c.py")) + + def test_network_and_server_errors_do_not_trigger_walkback(self): + for error in ( + URLError("offline"), + HTTPError("url", 503, "busy", Message(), None), + ): + with ( + self.subTest(error=error), + patch.object( + wheel_utils, "fetch_metadata_for_variant", side_effect=error + ) as fetch, + self.assertRaises(type(error)), + ): + wheel_utils.find_published_wheel_commit("HEAD", "cu130", "x86_64") + fetch.assert_called_once() + + +class TestDetermineWheelUrl(unittest.TestCase): + def determine(self, base: str, **env: str): + with ( + patch.dict( + os.environ, + {"VLLM_PRECOMPILED_WHEEL_VARIANT": "cu130", **env}, + clear=True, + ), + patch("platform.machine", return_value="x86_64"), + patch.object(wheel_utils, "is_rocm_system", return_value=False), + patch.object( + wheel_utils, "get_base_commit_in_main_branch", return_value=base + ), + patch.object( + wheel_utils, + "fetch_metadata_for_variant", + return_value=([WHEEL], "https://wheels.test/"), + ) as fetch, + patch.object( + wheel_utils, "find_published_wheel_commit", return_value="older" + ) as search, + ): + url = wheel_utils.determine_wheel_url() + return url, fetch, search + + def test_automatic_install_downloads_the_searched_commit(self): + url, fetch, search = self.determine("base") + self.assertEqual(url, ("https://wheels.test/wheel.whl", "wheel.whl")) + search.assert_called_once_with("base", "cu130", "x86_64") + fetch.assert_called_once_with("older", "cu130") + + def test_search_is_skipped_when_it_cannot_apply(self): + cases = { + "nightly fallback": ("nightly", {}, {}), + "explicit commit": ( + "base", + {"VLLM_PRECOMPILED_WHEEL_COMMIT": "a" * 40}, + {}, + ), + "rust only": ("base", {}, {"VLLM_USE_PRECOMPILED": False}), + "docker": ("base", {}, {"VLLM_DOCKER_BUILD_CONTEXT": True}), + } + for name, (base, env, flags) in cases.items(): + with self.subTest(name), patch.dict(vars(namespace["envs"]), flags): + _, fetch, search = self.determine(base, **env) + search.assert_not_called() + commit = env.get("VLLM_PRECOMPILED_WHEEL_COMMIT", base) + fetch.assert_called_once_with(commit, "cu130") + + def test_explicit_wheel_location_bypasses_selection(self): + for location in ("/tmp/local.whl", "https://example.com/custom.whl"): + with self.subTest(location=location): + url, fetch, search = self.determine( + "base", VLLM_PRECOMPILED_WHEEL_LOCATION=location + ) + self.assertEqual(url, (location, None)) + fetch.assert_not_called() + search.assert_not_called() + + +if __name__ == "__main__": + unittest.main() From 75e3b17c055c609f02c64cd4ad3f8692b8170d8e Mon Sep 17 00:00:00 2001 From: Taewoong Kim <116135174+ultramancode@users.noreply.github.com> Date: Sat, 10 Oct 2026 01:07:50 +0900 Subject: [PATCH 189/624] [Frontend] Add cache salt support to generative scoring (#60841) Signed-off-by: Taewoong Kim Co-authored-by: Cyrus Leung --- .../online_serving/generative_scoring.md | 7 +++++ docs/usage/security.md | 2 +- .../test_generative_scoring.py | 28 +++++++++++++++---- .../test_non_object_body_validation.py | 12 ++++++++ .../generate/generative_scoring/serving.py | 28 +++++++++++++++++-- 5 files changed, 68 insertions(+), 9 deletions(-) diff --git a/docs/serving/online_serving/generative_scoring.md b/docs/serving/online_serving/generative_scoring.md index d0c578f2105b..eebe1f965ec6 100644 --- a/docs/serving/online_serving/generative_scoring.md +++ b/docs/serving/online_serving/generative_scoring.md @@ -62,3 +62,10 @@ Here, each item is appended to the query to form prompts like `"Is this city the "usage": {"prompt_tokens": 45, "total_tokens": 48, "completion_tokens": 3} } ``` + +## Cache salting + +The optional `cache_salt` request field is applied to every item's prompt. With +prefix caching enabled, requests with different salts do not share cached prefix +blocks. See [Cache Salting](../../usage/security.md#cache-salting) for how to choose +and manage salts. diff --git a/docs/usage/security.md b/docs/usage/security.md index a74ed1d6ebe9..ab5aa96becc5 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -572,7 +572,7 @@ An attacker sharing the same backend can measure differences in Time to First To vLLM accepts an optional `cache_salt` parameter on requests. The salt is mixed into the hash of the first KV cache block, so only requests carrying the same salt can share cached prefix blocks. See [Automatic Prefix Caching](../design/prefix_caching.md) for the implementation details. -`cache_salt` is accepted by the OpenAI-compatible chat completions, completions, responses, and pooling (embeddings, classification, scoring) endpoints, the Anthropic `/v1/messages` endpoint, and the structured decisions `/v1/systemone` endpoint. +`cache_salt` is accepted by the OpenAI-compatible chat completions, completions, responses, and pooling (embeddings, classification, scoring) endpoints, the Anthropic `/v1/messages` endpoint, the structured decisions `/v1/systemone` endpoint, and the generative scoring `/generative_scoring` endpoint. #### Usage with the OpenAI Python client diff --git a/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py b/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py index 69be10b19d70..ea7c81ea2670 100644 --- a/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py +++ b/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py @@ -139,6 +139,7 @@ def test_request_and_response_all_fields(self): assert req_str.apply_softmax is True # default assert req_str.item_first is False # default assert req_str.add_special_tokens is True # default + assert req_str.cache_salt is None # Test request with pre-tokenized inputs and custom options req_tok = GenerativeScoringRequest( @@ -296,16 +297,26 @@ class TestGeneration: """Tests for the full generation flow with mocked engine.""" @pytest.mark.asyncio - async def test_successful_generation(self): - """Test successful score generation returns valid response.""" + @pytest.mark.parametrize("cache_salt", [None, "tenant-a"]) + @pytest.mark.parametrize( + "query,items", + [ + ("Is this city the capital of France? ", ["Paris", "London"]), + ([100, 101], [[200, 201], [300, 301]]), + ], + ids=["text", "token_ids"], + ) + async def test_successful_generation(self, cache_salt, query, items): + """Test scoring succeeds and preserves cache_salt for every item.""" mock_engine = _create_mock_engine() serving = _create_serving(mock_engine) mock_logprobs = {1234: -0.5, 5678: -2.0, 100: -3.0} mock_output = _create_mock_request_output(mock_logprobs) - state = SimpleNamespace(sampling_params=None) + state = SimpleNamespace(sampling_params=None, prompts=[]) async def mock_generate(*args, **kwargs): + state.prompts.append(args[0]) state.sampling_params = args[1] yield mock_output @@ -313,15 +324,22 @@ async def mock_generate(*args, **kwargs): request = GenerativeScoringRequest( model=MODEL_NAME, - query="Is Paris the capital?", - items=["Yes", "No"], + query=query, + items=items, label_token_ids=[1234, 5678], + cache_salt=cache_salt, ) result = await serving.create_generative_scoring(request, None) assert isinstance(result, GenerativeScoringResponse) assert state.sampling_params is not None assert state.sampling_params.watermarking is False + assert len(state.prompts) == len(items) + for prompt in state.prompts: + if cache_salt is None: + assert "cache_salt" not in prompt + else: + assert prompt["cache_salt"] == cache_salt assert len(result.data) == 2 for item_result in result.data: assert 0.0 <= item_result.score <= 1.0 diff --git a/tests/entrypoints/unit_tests/test_non_object_body_validation.py b/tests/entrypoints/unit_tests/test_non_object_body_validation.py index dd0e3eafd0a4..2440c880ce91 100644 --- a/tests/entrypoints/unit_tests/test_non_object_body_validation.py +++ b/tests/entrypoints/unit_tests/test_non_object_body_validation.py @@ -13,6 +13,9 @@ import pytest from pydantic import ValidationError +from vllm.entrypoints.generate.generative_scoring.serving import ( + GenerativeScoringRequest, +) from vllm.entrypoints.generate.structured_decisions.protocol import ( StructuredDecisionRequest, ) @@ -31,6 +34,7 @@ pytestmark = pytest.mark.skip_global_cleanup REQUEST_MODELS = [ + GenerativeScoringRequest, CompletionRequest, ResponsesRequest, EmbeddingChatRequest, @@ -77,6 +81,14 @@ def test_tokenize_chat_request_still_validates_dict_bodies(): CACHE_SALT_MODELS = [ + ( + GenerativeScoringRequest, + { + "query": "Is this city the capital of France?", + "items": ["Paris"], + "label_token_ids": [9454, 2753], + }, + ), ( StructuredDecisionRequest, { diff --git a/vllm/entrypoints/generate/generative_scoring/serving.py b/vllm/entrypoints/generate/generative_scoring/serving.py index bd46158edcab..74dbfd075997 100644 --- a/vllm/entrypoints/generate/generative_scoring/serving.py +++ b/vllm/entrypoints/generate/generative_scoring/serving.py @@ -12,12 +12,13 @@ import math import time from collections.abc import Mapping -from typing import Literal +from typing import Any, Literal from fastapi import Request -from pydantic import Field +from pydantic import Field, model_validator from vllm.engine.protocol import EngineClient +from vllm.entrypoints.generate.base.protocol import validate_cache_salt from vllm.entrypoints.generate.label_reads import next_token_label_reads from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.serve.engine.protocol import ( @@ -59,6 +60,7 @@ class GenerativeScoringRequest(OpenAIBaseModel): the full vocab for those ids (False). item_first: If True, prepend items to query. Otherwise append items to query. add_special_tokens: Whether to add special tokens when tokenizing. + cache_salt: Optional salt for prefix caching. """ @@ -91,6 +93,17 @@ class GenerativeScoringRequest(OpenAIBaseModel): default=True, description="Whether to add special tokens when tokenizing.", ) + cache_salt: str | None = Field( + default=None, + description=( + "If specified, the prefix cache will be salted with the provided " + "string to prevent an attacker to guess prompts in multi-user " + "environments. The salt should be random, protected from " + "access by 3rd parties, and long enough to be " + "unpredictable (e.g., 43 characters base64-encoded, corresponding " + "to 256 bit)." + ), + ) priority: int = Field( default=0, ge=-(2**63), @@ -104,6 +117,13 @@ class GenerativeScoringRequest(OpenAIBaseModel): description="The request_id related to this request.", ) + @model_validator(mode="before") + @classmethod + def check_cache_salt_support(cls, data: Any) -> Any: + if isinstance(data, dict): + validate_cache_salt(data.get("cache_salt")) + return data + class GenerativeScoringItemResult(OpenAIBaseModel): """Result for a single item in the generative scoring response. @@ -395,7 +415,9 @@ async def _build_prompts( if len(prompt_token_ids) > max_prompt_len: prompt_token_ids = prompt_token_ids[:max_prompt_len] - engine_inputs.append(tokens_input(prompt_token_ids)) + engine_inputs.append( + tokens_input(prompt_token_ids, cache_salt=request.cache_salt) + ) prompt_token_counts.append(len(prompt_token_ids)) return engine_inputs, prompt_token_counts From b74233860d4f2419ac5a763f6c83859bfe742ad0 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Fri, 9 Oct 2026 18:48:43 +0200 Subject: [PATCH 190/624] [CI][ROCm] Wait for VRAM to settle before building LLM in test_full_graph (#60872) Signed-off-by: Stefan Koncarevic --- tests/compile/fullgraph/test_full_graph.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/tests/compile/fullgraph/test_full_graph.py b/tests/compile/fullgraph/test_full_graph.py index 4ce7d6767f3f..fceed6346d1e 100644 --- a/tests/compile/fullgraph/test_full_graph.py +++ b/tests/compile/fullgraph/test_full_graph.py @@ -10,11 +10,17 @@ from tests.quantization.utils import is_quant_method_supported from vllm import LLM, SamplingParams -from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, PassConfig +from vllm.config import ( + CacheConfig, + CompilationConfig, + CompilationMode, + CUDAGraphMode, + PassConfig, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import is_torch_equal_or_newer -from ...utils import create_new_process_for_each_test +from ...utils import create_new_process_for_each_test, wait_for_memory_to_settle def models_list(*, all: bool = True, keywords: list[str] | None = None): @@ -173,6 +179,14 @@ def run_model( if compilation_config.cudagraph_mode is None: compilation_config.cudagraph_mode = CUDAGraphMode.NONE + # V1 startup requires free_memory >= total * gpu_memory_utilization, so a + # device that is not free yet fails engine init rather than waiting for it. + # VllmRunner guards this; building LLM directly does not. No-op off ROCm. + wait_for_memory_to_settle( + threshold_ratio=1.0 + - model_kwargs.get("gpu_memory_utilization", CacheConfig.gpu_memory_utilization) + ) + llm = LLM( model=model, compilation_config=compilation_config, From 89ac8d971be7d2f598598c1551a21c646bdf8b0d Mon Sep 17 00:00:00 2001 From: Andrey Talman Date: Fri, 9 Oct 2026 12:59:17 -0400 Subject: [PATCH 191/624] [CI][Test] Allow 10% ngram acceptance drop under preemption on CUDA in test_async_scheduling_accuracy (#60715) Signed-off-by: Andrey Talman Co-authored-by: Andrey Talman --- tests/v1/e2e/general/test_async_scheduling_accuracy.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/v1/e2e/general/test_async_scheduling_accuracy.py b/tests/v1/e2e/general/test_async_scheduling_accuracy.py index aa785b3270ce..3ffcaefdc3f3 100644 --- a/tests/v1/e2e/general/test_async_scheduling_accuracy.py +++ b/tests/v1/e2e/general/test_async_scheduling_accuracy.py @@ -271,9 +271,9 @@ def run_tests( context = f"config=[{config}], params={params}" if spec_config.get("max_model_len") is None: # Keep the original acceptance-quality floor per batch. - relative_drop = ( - 0.10 if current_platform.is_rocm() and test_preemption else 0.05 - ) + # Preemption recomputes requests and loses n-gram drafts, + # so allow a larger drop there. + relative_drop = 0.10 if test_preemption else 0.05 assert actual >= baseline * (1 - relative_drop), ( f"{context}: acceptance={actual}, baseline={baseline}" ) From 1a3a7a8f01a8c1a3184151b8185072efda11daab Mon Sep 17 00:00:00 2001 From: Yuvraj Singh Bhadoria Date: Fri, 9 Oct 2026 22:45:22 +0530 Subject: [PATCH 192/624] [Bugfix] Treat max_tokens as output upper bound, not input-space reservation (#57035) Signed-off-by: YuvrajSinghBhadoria2 <245782664+YuvrajSinghBhadoria2@users.noreply.github.com> Signed-off-by: YuvrajSinghBhadoria2 Signed-off-by: Yuvraj Singh Bhadoria Co-authored-by: YuvrajSinghBhadoria2 <245782664+YuvrajSinghBhadoria2@users.noreply.github.com> Co-authored-by: Cyrus Leung Co-authored-by: Claude Opus 4.8 (1M context) --- .../chat_completion/test_serving_chat.py | 22 +++- .../responses/test_serving_responses.py | 6 +- .../pooling/embed/test_protocol.py | 59 +++++++++ .../pooling/test_truncation_control.py | 2 + tests/renderers/test_completions.py | 118 ++++++++++++++++++ vllm/entrypoints/pooling/base/protocol.py | 6 +- vllm/renderers/params.py | 71 +++++++---- 7 files changed, 252 insertions(+), 32 deletions(-) diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index bfaa2d4ca0b2..d22f581f67ec 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -1684,7 +1684,14 @@ def test_serving_chat_starts_without_tokenizer(): @pytest.mark.asyncio async def test_serving_chat_mistral_token_ids_prompt_is_validated(): """Regression test: when the Mistral tokenizer path returns token IDs - directly, we must still apply input length + max_tokens validation. + directly, we must still apply input length validation and clamp the + generation budget to the context the prompt leaves over. + + `max_tokens` is an upper bound on generation rather than a reservation off + the context window, so a 95-token prompt against `max_model_len=100` is + accepted even though the client asked for 10 output tokens. Genuine + overflow is covered by + test_serving_chat_mistral_token_ids_prompt_too_long_is_rejected. """ mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False @@ -1696,9 +1703,8 @@ async def test_serving_chat_mistral_token_ids_prompt_is_validated(): MockVllmConfig(mock_engine.model_config, parallel_config=MockParallelConfig()), tokenizer=mock_tokenizer, ) - # Force the Mistral chat template renderer to return token IDs. - # Choose a prompt length that is < max_model_len, but large enough that - # adding max_tokens should exceed the model context window. + # Force the Mistral chat template renderer to return token IDs. The prompt + # fits the context on its own; only the naive prompt+output sum overflows. mock_renderer.render_messages_async = AsyncMock( return_value=( [], @@ -1715,8 +1721,12 @@ async def test_serving_chat_mistral_token_ids_prompt_is_validated(): max_tokens=10, ) - with pytest.raises(VLLMValidationError): - await serving_chat.create_chat_completion(req) + # Must not raise: the prompt alone fits within max_model_len. + await serving_chat.create_chat_completion(req) + + # get_max_tokens() clamps the sampling budget to the 5 tokens left over. + sampling_params = mock_engine.generate.call_args.args[1] + assert 1 <= sampling_params.max_tokens <= 5 @pytest.mark.asyncio diff --git a/tests/entrypoints/openai/responses/test_serving_responses.py b/tests/entrypoints/openai/responses/test_serving_responses.py index acab4cb6e2ea..1622d0e8f8cf 100644 --- a/tests/entrypoints/openai/responses/test_serving_responses.py +++ b/tests/entrypoints/openai/responses/test_serving_responses.py @@ -515,7 +515,11 @@ async def test_online_renderer_applies_responses_token_budget_to_harmony_prompt( result = await renderer.render_responses(request) assert not isinstance(result, ErrorResponse) - assert result.engine_input["prompt_token_ids"] == [4, 5, 6] + # `max_tokens` is an output upper bound rather than a reservation off the + # context window, so the prompt budget is the full max_model_len minus the + # one token held back for output: 6 tokens truncate to 4, not to + # max_model_len - max_output_tokens. + assert result.engine_input["prompt_token_ids"] == [3, 4, 5, 6] @pytest.mark.asyncio diff --git a/tests/entrypoints/pooling/embed/test_protocol.py b/tests/entrypoints/pooling/embed/test_protocol.py index 91c853f91098..ba797b4ff111 100644 --- a/tests/entrypoints/pooling/embed/test_protocol.py +++ b/tests/entrypoints/pooling/embed/test_protocol.py @@ -4,12 +4,15 @@ underlying packing helpers, plus Cohere-specific serving helpers.""" import struct +from unittest.mock import Mock import numpy as np import pybase64 as base64 import pytest +from vllm.config import ModelConfig, PoolerConfig from vllm.entrypoints.pooling.embed.protocol import ( + EmbeddingCompletionRequest, build_typed_embeddings, ) @@ -162,3 +165,59 @@ def test_subset_types(self, sample_embeddings: list[list[float]]): def test_unknown_type_ignored(self, sample_embeddings: list[list[float]]): result = build_typed_embeddings(sample_embeddings, ["float", "unknown_type"]) assert result.float is not None + + +def _embedding_model_config( + *, + max_model_len: int = 128, + max_embed_len: int | None = None, + enable_chunked_processing: bool = False, +) -> Mock: + model_config = Mock(spec=ModelConfig) + model_config.max_model_len = max_model_len + model_config.encoder_config = None + model_config.pooler_config = PoolerConfig( + seq_pooling_type="CLS", + max_embed_len=max_embed_len, + enable_chunked_processing=enable_chunked_processing, + ) + return model_config + + +class TestMaxEmbedLenCapsInput: + """``build_tok_params`` used to enforce ``max_embed_len`` indirectly, by + setting ``max_output_tokens = max_model_len - max_embed_len`` and relying on + ``max_input_tokens == max_total_tokens - max_output_tokens``. The cap is now + expressed directly as ``max_total_tokens``.""" + + def test_non_chunked_max_embed_len_caps_input(self): + cfg = _embedding_model_config(max_model_len=128, max_embed_len=64) + + tok_params = EmbeddingCompletionRequest(model="m", input="hi").build_tok_params( + cfg + ) + + assert tok_params.max_input_tokens == 64 + assert tok_params.max_output_tokens == 0 + + def test_chunked_max_embed_len_caps_input(self): + cfg = _embedding_model_config( + max_model_len=128, max_embed_len=64, enable_chunked_processing=True + ) + + tok_params = EmbeddingCompletionRequest(model="m", input="hi").build_tok_params( + cfg + ) + + assert tok_params.max_input_tokens == 64 + assert tok_params.max_output_tokens == 0 + + def test_no_max_embed_len_falls_back_to_model_len(self): + cfg = _embedding_model_config(max_model_len=128) + + tok_params = EmbeddingCompletionRequest(model="m", input="hi").build_tok_params( + cfg + ) + + assert tok_params.max_input_tokens == 128 + assert tok_params.max_output_tokens == 0 diff --git a/tests/models/language/pooling/test_truncation_control.py b/tests/models/language/pooling/test_truncation_control.py index c4485195558a..003213410c28 100644 --- a/tests/models/language/pooling/test_truncation_control.py +++ b/tests/models/language/pooling/test_truncation_control.py @@ -55,6 +55,8 @@ def test_max_truncation_size(vllm_model): prompt_tokens = vllm_output[0].prompt_token_ids + # Pooling requests emit an embedding rather than generated tokens, so -1 + # truncates to the full context length. assert len(prompt_tokens) == max_model_len diff --git a/tests/renderers/test_completions.py b/tests/renderers/test_completions.py index 0dce7ef68412..4fac06380750 100644 --- a/tests/renderers/test_completions.py +++ b/tests/renderers/test_completions.py @@ -247,6 +247,8 @@ def test_neg_truncation(self): ) assert len(results) == 1 + # max_output_tokens defaults to 0, so there is no output budget to + # reserve and -1 truncates to the full max_total_tokens. assert len(results[0]["prompt_token_ids"]) == 100 # max_total_tokens def test_truncation_left(self): @@ -505,6 +507,122 @@ def test_explicit_side_text_pretokenization_guard(self): assert renderer.tokenizer._captured_text_len <= 100 +class TestMaxTokensNotReservation: + """Regression tests for GitHub issue #42474. + + ``max_tokens`` must be treated as an upper bound on generation length, not + as a hard reservation that is subtracted from the available input space. + """ + + def test_large_max_tokens_does_not_eat_prompt_budget(self): + """Prompt fits in context even when max_tokens is large.""" + renderer = _build_renderer(MockModelConfig()) + + # 50-token prompt with a 100-token context window and max_tokens=60. + # Before the fix: max_input_tokens = 100 - 60 = 40 → 50 > 40 → error. + # After the fix: max_input_tokens = 100 → 50 ≤ 100 → OK. + prompts = renderer.render_prompts( + _preprocess_prompt(renderer.model_config, "x" * 50) + ) + results = renderer.tokenize_prompts( + prompts, + TokenizeParams(max_total_tokens=100, max_output_tokens=60), + ) + assert len(results) == 1 + assert len(results[0]["prompt_token_ids"]) == 50 + + def test_max_tokens_zero_still_rejects_oversized_prompt(self): + """Genuine prompt overflow is still caught regardless of max_tokens.""" + renderer = _build_renderer(MockModelConfig()) + + # 150-token prompt with 100-token context → should still fail. + prompts = renderer.render_prompts( + _preprocess_prompt(renderer.model_config, "x" * 150) + ) + with pytest.raises(VLLMValidationError, match="maximum context length"): + renderer.tokenize_prompts( + prompts, + TokenizeParams(max_total_tokens=100, max_output_tokens=0), + ) + + def test_negative_truncation_uses_max_total(self): + """truncate_prompt_tokens=-1 reserves one output token.""" + renderer = _build_renderer(MockModelConfig()) + + prompts = renderer.render_prompts( + _preprocess_prompt(renderer.model_config, "x" * 200) + ) + results = renderer.tokenize_prompts( + prompts, + TokenizeParams( + max_total_tokens=100, + max_output_tokens=40, + truncate_prompt_tokens=-1, + ), + ) + # -1 maps to max_total_tokens (100) rather than the old + # max_total - max_output (60), minus the one token held back so + # truncation cannot leave a zero-token output budget. + assert len(results[0]["prompt_token_ids"]) == 99 + + def test_negative_pad_uses_max_total(self): + """pad_prompt_tokens=-1 pads to max_total_tokens, not max_input.""" + renderer = _build_renderer(MockModelConfig()) + pad_id = renderer.tokenizer.pad_token_id + + prompts = renderer.render_prompts( + _preprocess_prompt(renderer.model_config, "x" * 30) + ) + results = renderer.tokenize_prompts( + prompts, + TokenizeParams( + max_total_tokens=100, + max_output_tokens=40, + pad_prompt_tokens=-1, + ), + ) + # -1 should map to max_total_tokens (100), not max_total - max_output (60). + assert results[0]["prompt_token_ids"] == list(range(30)) + [pad_id] * 70 + + def test_max_tokens_exceeds_context_is_still_clamped_by_api(self): + """Renderer allows the prompt; API layer clamps max_tokens later.""" + renderer = _build_renderer(MockModelConfig()) + + # Simulates the reported scenario: max_tokens > context - prompt_length. + # Renderer should accept the 40-token prompt; get_max_tokens() clamps + # sampling max_tokens to 60 at the entrypoint layer. + prompts = renderer.render_prompts( + _preprocess_prompt(renderer.model_config, "x" * 40) + ) + results = renderer.tokenize_prompts( + prompts, + TokenizeParams(max_total_tokens=100, max_output_tokens=80), + ) + assert len(results) == 1 + assert len(results[0]["prompt_token_ids"]) == 40 + + def test_with_kwargs_preserves_output_budget(self): + """with_kwargs without an explicit max_length keeps the output budget. + + Under the old semantics `with_kwargs()` recomputed + `max_output_tokens = max_total - max_input`, which (after this fix) + would have silently reset a real output budget to zero. + """ + tok_params = TokenizeParams(max_total_tokens=100, max_output_tokens=60) + chained = tok_params.with_kwargs(truncation_side="left") + + assert chained.max_output_tokens == 60 + assert chained.max_input_tokens == 100 + + def test_with_kwargs_explicit_max_length_sets_budget(self): + """An explicit max_length still re-derives the output budget.""" + tok_params = TokenizeParams(max_total_tokens=100, max_output_tokens=60) + chained = tok_params.with_kwargs(max_length=80) + + assert chained.max_output_tokens == 20 + assert chained.max_input_tokens == 100 + + class TestRenderEmbedPrompt: def _create_test_embed_bytes(self, tensor: torch.Tensor) -> bytes: """Helper to create base64-encoded tensor bytes.""" diff --git a/vllm/entrypoints/pooling/base/protocol.py b/vllm/entrypoints/pooling/base/protocol.py index c1366734f7c1..0c646b969ab5 100644 --- a/vllm/entrypoints/pooling/base/protocol.py +++ b/vllm/entrypoints/pooling/base/protocol.py @@ -183,15 +183,15 @@ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: if pooler_config.enable_chunked_processing: max_total_tokens = pooler_config.max_embed_len else: - max_embed_len = pooler_config.max_embed_len or default_max_total_tokens - max_output_tokens = default_max_total_tokens - max_embed_len + max_total_tokens = ( + pooler_config.max_embed_len or default_max_total_tokens + ) tok_params = self._build_pooling_tok_params( model_config, add_special_tokens=self.add_special_tokens, max_total_tokens=max_total_tokens, max_output_tokens=max_output_tokens, - max_output_tokens_param="max_model_len - max_embed_len", ) if ( diff --git a/vllm/renderers/params.py b/vllm/renderers/params.py index 9e62dfbb2d23..456a4d0ac563 100644 --- a/vllm/renderers/params.py +++ b/vllm/renderers/params.py @@ -207,11 +207,38 @@ class TokenizeParams: @property def max_input_tokens(self) -> int | None: - """Maximum allowed number of input tokens.""" + """Maximum allowed number of input tokens. + + ``max_output_tokens`` is *not* subtracted here: the value represents the + actual model context length, not a reservation of output space. + Downstream clamping of the sampling ``max_tokens`` is handled by + ``get_max_tokens()`` in the entrypoint layer. + """ if self.max_total_tokens is None: return None - return self.max_total_tokens - self.max_output_tokens + return self.max_total_tokens + + @property + def max_truncation_tokens(self) -> int | None: + """Maximum prompt length that truncation may reduce a prompt to. + + For generative requests this is `max_input_tokens` less one, holding a + token of context back for output. Truncating to the full + `max_total_tokens` would leave `get_max_tokens()` clamping the sampling + budget to 0, turning an over-long prompt into a silent zero-token + generation. + + Pooling requests set `max_output_tokens = 0` and emit an embedding + rather than generated tokens, so they keep the full input budget. + """ + max_input_tokens = self.max_input_tokens + if max_input_tokens is None: + return None + if self.max_output_tokens == 0: + return max_input_tokens + + return max(0, max_input_tokens - 1) def __post_init__(self) -> None: max_total_tokens = self.max_total_tokens @@ -247,14 +274,15 @@ def __post_init__(self) -> None: ): raise VLLMValidationError( f"{self.truncate_prompt_tokens_param}={truncate_prompt_tokens} " - f"cannot be greater than {self.max_total_tokens_param} - " - f"{self.max_output_tokens_param} = {max_input_tokens}. " + f"cannot be greater than {self.max_total_tokens_param} " + f"= {max_input_tokens}. " f"Please request a smaller truncation size.", parameter=self.truncate_prompt_tokens_param, value=truncate_prompt_tokens, ) def with_kwargs(self, **tokenization_kwargs: Any): + has_max_length = "max_length" in tokenization_kwargs max_length = tokenization_kwargs.pop("max_length", self.max_input_tokens) pad_prompt_tokens = tokenization_kwargs.pop( "pad_prompt_tokens", self.pad_prompt_tokens @@ -301,13 +329,17 @@ def with_kwargs(self, **tokenization_kwargs: Any): max_total_tokens = self.max_total_tokens + if has_max_length and max_total_tokens is not None and max_length is not None: + max_output_tokens = max_total_tokens - max_length + else: + # No explicit ``max_length`` override: preserve the original output + # budget instead of recomputing it from ``max_input_tokens`` (which + # is now the full context length). + max_output_tokens = self.max_output_tokens + return TokenizeParams( max_total_tokens=max_total_tokens, - max_output_tokens=( - 0 - if max_total_tokens is None or max_length is None - else max_total_tokens - max_length - ), + max_output_tokens=max_output_tokens, pad_prompt_tokens=pad_prompt_tokens, truncate_prompt_tokens=truncate_prompt_tokens, truncation_side=truncation_side, @@ -320,7 +352,7 @@ def get_encode_kwargs(self) -> dict[str, Any]: """The arguments to pass to `tokenizer.encode`.""" max_length = self.truncate_prompt_tokens if max_length is not None and max_length < 0: - max_length = self.max_input_tokens + max_length = self.max_truncation_tokens elif max_length is None and self.max_input_tokens is not None: # This prevents tokenization from taking up more resources than necessary # while still failing `self._token_len_check` as expected by users @@ -355,13 +387,11 @@ def _text_len_check(self, tokenizer: TokenizerLike | None, text: str) -> str: if len(text) > max_input_chars: raise VLLMValidationError( f"This model's maximum context length is " - f"{self.max_total_tokens} tokens. However, you requested " - f"{self.max_output_tokens} output tokens and your prompt " + f"{self.max_total_tokens} tokens. However, your prompt " f"contains {len(text)} characters (more than " f"{max_input_chars} characters, which is the upper bound " f"for {max_input_tokens} input tokens). " - f"Please reduce the length of the input prompt or the " - f"number of requested output tokens.", + f"Please reduce the length of the input prompt.", parameter="input_text", value=len(text), ) @@ -459,7 +489,7 @@ def _truncation_slice( """ max_length = self.truncate_prompt_tokens if max_length is not None and max_length < 0: - max_length = self.max_input_tokens + max_length = self.max_truncation_tokens if max_length is None or max_length >= length: return None @@ -494,15 +524,12 @@ def _token_len_check(self, tokenizer: TokenizerLike | None, tokens: _S) -> _S: # max_input_tokens + 1 (see get_encode_kwargs), so the # actual prompt length could be larger. qualifier = "at least " if token_count == max_input_tokens + 1 else "" - total = token_count + self.max_output_tokens raise VLLMValidationError( f"This model's maximum context length is " - f"{self.max_total_tokens} tokens. However, you requested " - f"{self.max_output_tokens} output tokens and your prompt " - f"contains {qualifier}{token_count} input tokens, " - f"for a total of {qualifier}{total} tokens. " - f"Please reduce the length of the input prompt or the " - f"number of requested output tokens.", + f"{self.max_total_tokens} tokens. However, your prompt " + f"contains {qualifier}{token_count} input tokens, which " + f"exceeds this limit even without any output tokens. " + f"Please reduce the length of the input prompt.", parameter="input_tokens", value=token_count, ) From 9a585c62aade0bfd290dfc7a975d54cab8ebfa54 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Sat, 10 Oct 2026 02:01:16 +0800 Subject: [PATCH 193/624] [Rust Frontend] Box rarely-set fields of per-token output types (#60588) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- rust/src/chat/src/event.rs | 6 ++-- rust/src/chat/src/output/mod.rs | 6 ++-- rust/src/chat/src/output/structured.rs | 4 +-- rust/src/chat/src/stream.rs | 4 +-- .../engine-core-client/src/protocol/output.rs | 20 +++++++----- rust/src/llm/src/output.rs | 32 ++++++++++--------- rust/src/llm/src/request_metrics.rs | 4 +-- rust/src/llm/tests/generate.rs | 20 ++++++------ rust/src/mock-engine/src/engine.rs | 2 +- rust/src/mock-engine/src/tests.rs | 4 +-- rust/src/server/src/grpc/convert.rs | 4 +-- rust/src/server/src/grpc/inference.rs | 4 +-- .../server/src/routes/inference/generate.rs | 18 ++++++----- rust/src/server/src/routes/tests.rs | 4 +-- rust/src/text/src/output/decoded.rs | 6 ++-- rust/src/text/src/output/mod.rs | 6 ++-- 16 files changed, 80 insertions(+), 64 deletions(-) diff --git a/rust/src/chat/src/event.rs b/rust/src/chat/src/event.rs index e57b8bbb86ab..5fe25a9969ef 100644 --- a/rust/src/chat/src/event.rs +++ b/rust/src/chat/src/event.rs @@ -237,9 +237,11 @@ pub enum ChatEvent { usage: ChatTokenUsage, finish_reason: FinishReason, /// Connector-specific KV transfer parameters for disaggregated serving. - kv_transfer_params: Option, + /// Boxed, like `ec_transfer_params`, to keep the per-token variants of + /// this enum small. + kv_transfer_params: Option>, /// Connector-specific encoder cache transfer parameters for /// disaggregated serving. - ec_transfer_params: Option, + ec_transfer_params: Option>, }, } diff --git a/rust/src/chat/src/output/mod.rs b/rust/src/chat/src/output/mod.rs index 8e62f3e8918b..8c1dae37f4de 100644 --- a/rust/src/chat/src/output/mod.rs +++ b/rust/src/chat/src/output/mod.rs @@ -54,10 +54,12 @@ pub(crate) enum AssistantEvent { usage: ChatTokenUsage, finish_reason: FinishReason, /// Connector-specific KV transfer parameters for disaggregated serving. - kv_transfer_params: Option, + /// Boxed, like `ec_transfer_params`, to keep the per-token variants of + /// this enum small. + kv_transfer_params: Option>, /// Connector-specific encoder cache transfer parameters for /// disaggregated serving. - ec_transfer_params: Option, + ec_transfer_params: Option>, }, } diff --git a/rust/src/chat/src/output/structured.rs b/rust/src/chat/src/output/structured.rs index 5f5bbf32fa8e..80035db4cb1e 100644 --- a/rust/src/chat/src/output/structured.rs +++ b/rust/src/chat/src/output/structured.rs @@ -150,8 +150,8 @@ impl StructuredEventState { &mut self, usage: ChatTokenUsage, finish_reason: FinishReason, - kv_transfer_params: Option, - ec_transfer_params: Option, + kv_transfer_params: Option>, + ec_transfer_params: Option>, ) -> Result> { let mut events = Vec::new(); self.close_open_text_block(&mut events); diff --git a/rust/src/chat/src/stream.rs b/rust/src/chat/src/stream.rs index 72b913078c29..509628e8e2af 100644 --- a/rust/src/chat/src/stream.rs +++ b/rust/src/chat/src/stream.rs @@ -95,8 +95,8 @@ impl ChatEventStream { token_ids, usage, finish_reason, - kv_transfer_params, - ec_transfer_params, + kv_transfer_params: kv_transfer_params.map(|value| *value), + ec_transfer_params: ec_transfer_params.map(|value| *value), }); } ChatEvent::ToolCallEnd { call, .. } => { diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index c33062eeb3c1..9339c117a889 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -82,6 +82,10 @@ pub struct EngineCoreEvent { /// /// Original Python definition: /// +/// +/// Fields that are set on few outputs (first or final output, or opt-in +/// features) are boxed: outputs are moved by value through the dispatcher and +/// the request stream once per engine step, so their inline size matters. #[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)] pub struct EngineCoreOutput { pub request_id: String, @@ -95,7 +99,7 @@ pub struct EngineCoreOutput { #[serde(default)] pub new_prompt_logprobs_tensors: Option, #[serde(default)] - pub pooling_output: Option, + pub pooling_output: Option>, #[serde(default)] pub finish_reason: Option, #[serde(default)] @@ -103,17 +107,17 @@ pub struct EngineCoreOutput { #[serde(default)] pub events: Option>, #[serde(default)] - pub kv_transfer_params: Option, + pub kv_transfer_params: Option>, #[serde(default)] - pub ec_transfer_params: Option, + pub ec_transfer_params: Option>, #[serde(default)] - pub trace_headers: Option, + pub trace_headers: Option>, /// Breakdown of the scheduled prefill computation, set on the first output /// of a newly scheduled prefill and elided for subsequent decode outputs. #[serde(default)] - pub prefill_stats: Option, + pub prefill_stats: Option>, #[serde(default)] - pub routed_experts: Option, + pub routed_experts: Option>, /// Number of NaNs seen in logits. Values above zero indicate corruption. #[serde(default)] pub num_nans_in_logits: u32, @@ -133,11 +137,11 @@ pub struct EngineCoreOutput { /// Per-request speculative-decoding acceptance metrics, set on the final /// output when `--per-request-spec-decode-metrics` is enabled. #[serde(default)] - pub spec_decode_metrics: Option, + pub spec_decode_metrics: Option>, /// Log probabilities of `SamplingParams.prompt_logprob_token_ids`, set on /// the first output of a request that asked for them. #[serde(default)] - pub prompt_token_id_logprobs: Option, + pub prompt_token_id_logprobs: Option>, } /// Raw per-sequence speculative-decoding accumulator. diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index 4b6d03353991..4dde96f2e20c 100644 --- a/rust/src/llm/src/output.rs +++ b/rust/src/llm/src/output.rs @@ -152,7 +152,7 @@ pub struct GenerateOutput { pub request_id: String, /// One-time prompt metadata emitted only on the first output for this /// request. - pub prompt_info: Option, + pub prompt_info: Option>, /// Newly produced token IDs for this step. pub token_ids: Vec, /// Sample logprobs for the generated positions in this step. @@ -162,14 +162,14 @@ pub struct GenerateOutput { /// Number of prompt tokens served from cache, when reported by prefill stats. pub cached_token_count: usize, /// Connector-specific KV transfer parameters for disaggregated serving. - pub kv_transfer_params: Option, + pub kv_transfer_params: Option>, /// Connector-specific encoder cache transfer parameters for disaggregated /// serving. - pub ec_transfer_params: Option, + pub ec_transfer_params: Option>, /// Sampling support sets aligned one-to-one with `token_ids`. pub sampling_mask: Option, /// Per-request speculative-decoding metrics, present on terminal outputs. - pub spec_decode_metrics: Option, + pub spec_decode_metrics: Option>, } impl GenerateOutput { @@ -207,10 +207,12 @@ impl GenerateOutput { ) -> Self { Self { request_id: String::new(), - prompt_info: prompt_token_ids.map(|ids| GeneratePromptInfo { - prompt_token_ids: ids, - prompt_logprobs: None, - prompt_token_id_logprobs: None, + prompt_info: prompt_token_ids.map(|ids| { + Box::new(GeneratePromptInfo { + prompt_token_ids: ids, + prompt_logprobs: None, + prompt_token_id_logprobs: None, + }) }), token_ids, logprobs: None, @@ -231,7 +233,7 @@ impl GenerateOutput { /// with `finish_reason = Abort` before the stream ends. /// - For errors or unexpected engine-side closes, the stream terminates with an error. pub struct GenerateOutputStream { - pending_prompt_info: Option, + pending_prompt_info: Option>, raw_stream: EngineCoreOutputStream, request_metrics: RequestMetricsTracker, /// Removes this request's external→internal tracking edge on drop. Held for @@ -249,11 +251,11 @@ impl GenerateOutputStream { request_guard: RequestGuard, ) -> Self { Self { - pending_prompt_info: Some(GeneratePromptInfo { + pending_prompt_info: Some(Box::new(GeneratePromptInfo { prompt_token_ids, prompt_logprobs: None, prompt_token_id_logprobs: None, - }), + })), raw_stream, request_metrics, _request_guard: request_guard, @@ -291,7 +293,7 @@ impl Stream for GenerateOutputStream { if let Some(info) = &mut self.pending_prompt_info && info.prompt_token_id_logprobs.is_none() { - info.prompt_token_id_logprobs = raw.prompt_token_id_logprobs; + info.prompt_token_id_logprobs = raw.prompt_token_id_logprobs.map(|logprobs| *logprobs); } let logprobs = raw.new_logprobs.map(|value| value.into_direct().unwrap()); @@ -432,9 +434,9 @@ impl> + Send> T { output_token_count: collected.token_ids.len(), cached_token_count, }; - collected.kv_transfer_params = output.kv_transfer_params; - collected.ec_transfer_params = output.ec_transfer_params; - collected.spec_decode_metrics = output.spec_decode_metrics; + collected.kv_transfer_params = output.kv_transfer_params.map(|value| *value); + collected.ec_transfer_params = output.ec_transfer_params.map(|value| *value); + collected.spec_decode_metrics = output.spec_decode_metrics.map(|value| *value); if let Some(mask) = collected.sampling_mask.as_ref() && mask.rows.len() != collected.token_ids.len() { diff --git a/rust/src/llm/src/request_metrics.rs b/rust/src/llm/src/request_metrics.rs index 98f9cadd1a35..61cea120940d 100644 --- a/rust/src/llm/src/request_metrics.rs +++ b/rust/src/llm/src/request_metrics.rs @@ -404,14 +404,14 @@ mod tests { timestamp: 9.0, }, ]), - prefill_stats: Some(PrefillStats { + prefill_stats: Some(Box::new(PrefillStats { num_prompt_tokens: 64, num_computed_tokens: 60, num_cached_tokens: 4, num_local_cached_tokens: 4, num_external_cached_tokens: 0, ..Default::default() - }), + })), ..Default::default() }, ); diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index c596c89ba3ba..813171555bba 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -92,8 +92,8 @@ fn request_output_with_logprobs_and_kv( new_logprobs: new_logprobs.map(MaybeWireLogprobs::Direct), new_prompt_logprobs_tensors: prompt_logprobs.map(MaybeWireLogprobs::Direct), finish_reason, - kv_transfer_params, - ec_transfer_params, + kv_transfer_params: kv_transfer_params.map(Box::new), + ec_transfer_params: ec_transfer_params.map(Box::new), ..Default::default() } } @@ -280,11 +280,11 @@ async fn generate_streams_outputs() { assert_eq!(first.request_id, internal_id); assert_eq!( first.prompt_info, - Some(GeneratePromptInfo { + Some(Box::new(GeneratePromptInfo { prompt_token_ids: vec![11, 22].into(), prompt_logprobs: Some(prompt_logprobs()), prompt_token_id_logprobs: None, - }) + })) ); assert_eq!(first.token_ids, vec![1, 2]); assert_eq!( @@ -330,12 +330,12 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() { RequestBatchOutputs { outputs: vec![ EngineCoreOutput { - prefill_stats: Some(PrefillStats { + prefill_stats: Some(Box::new(PrefillStats { num_prompt_tokens: 2, num_cached_tokens: 1, num_local_cached_tokens: 1, ..Default::default() - }), + })), new_sampling_mask: Some(MaybeWireSamplingMask::Direct( SamplingMask { rows: vec![vec![1, 33, 99]], @@ -409,11 +409,11 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() { async fn collect_output_rejects_partial_sampling_mask() { let output = vllm_llm::GenerateOutput { request_id: "req-partial-mask".to_string(), - prompt_info: Some(GeneratePromptInfo { + prompt_info: Some(Box::new(GeneratePromptInfo { prompt_token_ids: Arc::from([11_u32, 22]), prompt_logprobs: None, prompt_token_id_logprobs: None, - }), + })), token_ids: vec![33, 44], logprobs: None, finish_reason: Some(FinishReason::Length), @@ -763,11 +763,11 @@ async fn generate_records_request_metrics_in_prometheus_output() { engine_index: 4, timestamp: 10.0, outputs: vec![EngineCoreOutput { - prefill_stats: Some(PrefillStats { + prefill_stats: Some(Box::new(PrefillStats { num_prompt_tokens: 2, num_computed_tokens: 2, ..Default::default() - }), + })), ..request_output_with_events( &request.request_id, vec![1], diff --git a/rust/src/mock-engine/src/engine.rs b/rust/src/mock-engine/src/engine.rs index 4844aba12018..8685357b1896 100644 --- a/rust/src/mock-engine/src/engine.rs +++ b/rust/src/mock-engine/src/engine.rs @@ -201,7 +201,7 @@ impl ActiveRequest { let finished = self.generated >= self.max_tokens; EngineCoreOutput { - prefill_stats, + prefill_stats: prefill_stats.map(Box::new), ..request_output( self.request_id.clone(), new_token_ids, diff --git a/rust/src/mock-engine/src/tests.rs b/rust/src/mock-engine/src/tests.rs index 71b4286033ad..d82ceac083ca 100644 --- a/rust/src/mock-engine/src/tests.rs +++ b/rust/src/mock-engine/src/tests.rs @@ -123,11 +123,11 @@ async fn chunk_size_one_outputs_one_token_per_update() { assert_eq!(first.finish_reason, None); assert_eq!( first.prefill_stats, - Some(PrefillStats { + Some(Box::new(PrefillStats { num_prompt_tokens: 3, num_computed_tokens: 3, ..Default::default() - }) + })) ); let second = stream.next().await.expect("second").expect("second ok"); assert_eq!(second.new_token_ids.len(), 1); diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index d495fea485c7..8e31bb84d2d3 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -406,8 +406,8 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> Result { if let Some(metrics) = output.spec_decode_metrics { - spec_decode_metrics = Some(SpeculativeDecodingMetrics::from(metrics)); + spec_decode_metrics = Some(SpeculativeDecodingMetrics::from(*metrics)); } if prompt_tokens.is_none() && let Some(info) = output.prompt_info.as_ref() @@ -632,11 +632,11 @@ mod tests { }), Ok(GenerateOutput { request_id: String::new(), - prompt_info: Some(GeneratePromptInfo { + prompt_info: Some(Box::new(GeneratePromptInfo { prompt_token_ids: Arc::from([11_u32, 22_u32]), prompt_logprobs: None, prompt_token_id_logprobs: None, - }), + })), token_ids: vec![33], logprobs: None, finish_reason: Some(FinishReason::stop_eos()), @@ -716,10 +716,12 @@ mod tests { ) -> GenerateOutput { GenerateOutput { request_id: String::new(), - prompt_info: prompt_token_ids.map(|ids| GeneratePromptInfo { - prompt_token_ids: Arc::from(ids), - prompt_logprobs: None, - prompt_token_id_logprobs: None, + prompt_info: prompt_token_ids.map(|ids| { + Box::new(GeneratePromptInfo { + prompt_token_ids: Arc::from(ids), + prompt_logprobs: None, + prompt_token_id_logprobs: None, + }) }), token_ids, logprobs: None, @@ -817,7 +819,7 @@ mod tests { ..Default::default() }; let mut output = stream_output(None, Vec::new(), Some(FinishReason::stop_eos())); - output.spec_decode_metrics = Some(metrics.clone()); + output.spec_decode_metrics = Some(Box::new(metrics.clone())); for include_usage in [false, true] { let chunks: Vec<_> = generate_chunk_stream( diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 0ed2ea4d4590..a9c9831bbd37 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -120,8 +120,8 @@ fn request_output_with_logprobs_and_kv( new_prompt_logprobs_tensors: new_prompt_logprobs_tensors.map(MaybeWireLogprobs::Direct), finish_reason, stop_reason, - kv_transfer_params, - ec_transfer_params, + kv_transfer_params: kv_transfer_params.map(Box::new), + ec_transfer_params: ec_transfer_params.map(Box::new), ..Default::default() } } diff --git a/rust/src/text/src/output/decoded.rs b/rust/src/text/src/output/decoded.rs index d3dc42265641..64df0492bfbe 100644 --- a/rust/src/text/src/output/decoded.rs +++ b/rust/src/text/src/output/decoded.rs @@ -47,10 +47,12 @@ pub struct Finished { pub usage: TokenUsage, pub finish_reason: FinishReason, /// Connector-specific KV transfer parameters for disaggregated serving. - pub kv_transfer_params: Option, + /// Kept boxed as received from `vllm-llm`, so the chat events built from + /// this stay small without boxing again. + pub kv_transfer_params: Option>, /// Connector-specific encoder cache transfer parameters for disaggregated /// serving. - pub ec_transfer_params: Option, + pub ec_transfer_params: Option>, /// Sampling support sets aligned with all generated token positions. pub sampling_mask: Option, } diff --git a/rust/src/text/src/output/mod.rs b/rust/src/text/src/output/mod.rs index 1293ac6025ff..8e09f0761870 100644 --- a/rust/src/text/src/output/mod.rs +++ b/rust/src/text/src/output/mod.rs @@ -101,8 +101,10 @@ impl T { let mut collected = collected.unwrap(); collected.finish_reason = finished.finish_reason; collected.usage = finished.usage; - collected.kv_transfer_params = finished.kv_transfer_params; - collected.ec_transfer_params = finished.ec_transfer_params; + collected.kv_transfer_params = + finished.kv_transfer_params.map(|value| *value); + collected.ec_transfer_params = + finished.ec_transfer_params.map(|value| *value); collected.sampling_mask = finished.sampling_mask; return Ok(collected); } From ce4603cad676af27f08f39b696a31a4022eb9ca4 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Sat, 10 Oct 2026 02:04:00 +0800 Subject: [PATCH 194/624] [Bugfix][Rust Frontend] Pass through raw characters in byte-level decode (#59861) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- rust/src/tokenizer/src/byte_level_decode.rs | 54 ++++++++++++++++----- rust/src/tokenizer/src/hf.rs | 44 +++++++++++++++++ 2 files changed, 86 insertions(+), 12 deletions(-) diff --git a/rust/src/tokenizer/src/byte_level_decode.rs b/rust/src/tokenizer/src/byte_level_decode.rs index ad71f6decbbf..ef7a3c1ef674 100644 --- a/rust/src/tokenizer/src/byte_level_decode.rs +++ b/rust/src/tokenizer/src/byte_level_decode.rs @@ -7,14 +7,16 @@ /// Reverse GPT-2 byte-to-unicode mapping: codepoint → original byte. The GPT-2 /// table only emits codepoints in U+0000..U+0143, so a flat array suffices. -const CHAR_TO_BYTE: [u8; 324] = build_char_to_byte(); +/// Codepoints in that range the table never emits, such as raw whitespace in +/// an added token's content, map to `None`. +const CHAR_TO_BYTE: [Option; 324] = build_char_to_byte(); const fn is_nice(b: u8) -> bool { (b >= b'!' && b <= b'~') || (b >= 0xA1 && b <= 0xAC) || b >= 0xAE } -const fn build_char_to_byte() -> [u8; 324] { - let mut table = [0u8; 324]; +const fn build_char_to_byte() -> [Option; 324] { + let mut table = [None; 324]; let mut b: u16 = 0; while b < 256 { let cp = if is_nice(b as u8) { @@ -22,7 +24,7 @@ const fn build_char_to_byte() -> [u8; 324] { } else { 256 + nice_offset(b as u8) }; - table[cp as usize] = b as u8; + table[cp as usize] = Some(b as u8); b += 1; } table @@ -41,20 +43,23 @@ const fn nice_offset(b: u8) -> u32 { } /// Decode byte-level encoded token strings into a single UTF-8 string, -/// matching `fastokens::decoders::ByteLevelDecoder`. +/// matching the HuggingFace `tokenizers` `ByteLevel` decoder. pub fn decode_byte_level<'a, I: IntoIterator>(tokens: I) -> String { let iter = tokens.into_iter(); let (lower, _) = iter.size_hint(); let mut bytes: Vec = Vec::with_capacity(lower.saturating_mul(4)); for token in iter { + let start = bytes.len(); for c in token.chars() { - let cp = c as usize; - if cp < CHAR_TO_BYTE.len() { - bytes.push(CHAR_TO_BYTE[cp]); + if let Some(&Some(b)) = CHAR_TO_BYTE.get(c as usize) { + bytes.push(b); } else { - // Non-GPT2 codepoints (e.g. DeepSeek's U+FF5C, U+2581) pass through. - let mut buf = [0u8; 4]; - bytes.extend_from_slice(c.encode_utf8(&mut buf).as_bytes()); + // As in HF, a token with any non-GPT2 codepoint is raw text and + // passes through whole: DeepSeek's U+FF5C and U+2581, or raw + // characters in added tokens (e.g. `\n`). + bytes.truncate(start); + bytes.extend_from_slice(token.as_bytes()); + break; } } } @@ -87,7 +92,7 @@ mod tests { for b in 0..=255u8 { let cp = byte_to_char[b as usize] as usize; assert!(cp < CHAR_TO_BYTE.len()); - assert_eq!(CHAR_TO_BYTE[cp], b, "mismatch for byte {b:#x}"); + assert_eq!(CHAR_TO_BYTE[cp], Some(b), "mismatch for byte {b:#x}"); } } @@ -119,4 +124,29 @@ mod tests { let tok = "<\u{FF5C}begin\u{2581}of\u{2581}sentence\u{FF5C}>"; assert_eq!(decode_byte_level([tok]), "<｜begin▁of▁sentence｜>"); } + + #[test] + fn decode_preserves_raw_chars_outside_gpt2_table() { + // Added tokens store raw text, which may contain codepoints below + // U+0144 that the GPT-2 table never emits. + let toks = [ + "", + "\n", + "\t\u{7F}\u{A0}\u{AD}", + ]; + assert_eq!( + decode_byte_level(toks), + "\n\t\u{7F}\u{A0}\u{AD}", + ); + } + + #[test] + fn decode_passes_through_whole_token_with_raw_chars() { + // Characters that the GPT-2 table does emit (`Ġ`, `é`) stay verbatim + // in a raw token instead of decoding to bytes 0x20 and 0xE9. + assert_eq!(decode_byte_level(["\u{120}x\n"]), "\u{120}x\n"); + assert_eq!(decode_byte_level(["\ncafé"]), "\ncafé"); + } } diff --git a/rust/src/tokenizer/src/hf.rs b/rust/src/tokenizer/src/hf.rs index 44ddc6cb27cb..3ae1779b42cc 100644 --- a/rust/src/tokenizer/src/hf.rs +++ b/rust/src/tokenizer/src/hf.rs @@ -554,6 +554,50 @@ mod tests { } } + /// Added tokens store raw text, including whitespace that the GPT-2 byte + /// table never emits, so a byte-level decode must pass it through verbatim, + /// along with characters the table does emit in the same token (`é`). + #[test] + fn added_tokens_with_raw_whitespace_decode_verbatim() { + const OPEN: &str = ">() + ); + let dir = tempdir().expect("create temp dir"); + let path = write_tokenizer_json(dir.path(), "tokenizer.json", &value); + let text = format!("{OPEN}city\">\nParis{CLOSE}{MIXED}"); + + let fastokens = HuggingFaceTokenizer::new_fastokens(&path).expect("load fastokens wrapper"); + assert!(matches!( + fastokens.backend, + super::Backend::FastokensByteLevel(_) + )); + let hf = HuggingFaceTokenizer::new_hf(&path).expect("load hf wrapper"); + for wrapper in [fastokens, hf] { + let ids = wrapper.encode(&text, false).expect("encode"); + assert!((256..=258).all(|id| ids.contains(&id)), "ids={ids:?}"); + assert_eq!(wrapper.decode(&ids, true).expect("decode"), text); + } + } + #[test] fn hf_vocab_size_counts_added_tokens() { let mut tokenizer = tiny_bpe_tokenizer(); From 5dc9bcd8cd9656b852edfd6c5b4a8c9c46a34530 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Sat, 10 Oct 2026 02:04:49 +0800 Subject: [PATCH 195/624] [Rust Frontend] Pass reasoning controls through render -> generate (#60546) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- .../src/protocol/request.rs | 1 + .../src/routes/inference/generate/convert.rs | 52 ++++++++++++++++++- .../src/routes/inference/generate/types.rs | 8 +++ rust/src/server/src/routes/render.rs | 7 +++ rust/src/server/src/routes/tests.rs | 37 +++++++++---- 5 files changed, 92 insertions(+), 13 deletions(-) diff --git a/rust/src/engine-core-client/src/protocol/request.rs b/rust/src/engine-core-client/src/protocol/request.rs index 02fc1c30cad6..5390eed16f2c 100644 --- a/rust/src/engine-core-client/src/protocol/request.rs +++ b/rust/src/engine-core-client/src/protocol/request.rs @@ -64,6 +64,7 @@ impl EngineCoreRequestType { #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] pub struct ReasoningParserKwargs { /// Effective kwargs visible to the chat template for this request. + #[serde(default)] pub chat_template_kwargs: HashMap, } diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs index 36a497f6ddf3..313e9fa1f955 100644 --- a/rust/src/server/src/routes/inference/generate/convert.rs +++ b/rust/src/server/src/routes/inference/generate/convert.rs @@ -85,8 +85,8 @@ pub(super) fn prepare_generate_request( data_parallel_rank: ctx.data_parallel_rank, session_id: ctx.session_id, kv_hints: None, - reasoning_parser_kwargs: Default::default(), - reasoning_ended: None, + reasoning_parser_kwargs: request.reasoning_parser_kwargs.unwrap_or_default(), + reasoning_ended: request.reasoning_ended, lora_request: lora_resolution.lora_request.clone(), arrival_time: None, }; @@ -224,6 +224,54 @@ mod tests { ); } + #[test] + fn prepare_generate_request_forwards_reasoning_controls() { + for (controls, expected) in [ + (json!({}), (json!({"chat_template_kwargs": {}}), None)), + ( + json!({ + "reasoning_ended": false, + "reasoning_parser_kwargs": { + "chat_template_kwargs": {"enable_thinking": true} + } + }), + ( + json!({"chat_template_kwargs": {"enable_thinking": true}}), + Some(false), + ), + ), + ( + json!({"reasoning_ended": true, "reasoning_parser_kwargs": {}}), + (json!({"chat_template_kwargs": {}}), Some(true)), + ), + ] { + let mut body = json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "token_ids": [11, 22, 33], + "sampling_params": {} + }); + body.as_object_mut().unwrap().extend(controls.as_object().unwrap().clone()); + let request: GenerateRequest = serde_json::from_value(body).expect("parse request"); + + let prepared = prepare_generate_request( + request, + &served(&["Qwen/Qwen1.5-0.5B-Chat"]), + ResolvedRequestContext::default(), + None, + ) + .expect("prepare"); + + assert_eq!( + ( + serde_json::to_value(&prepared.text_request.reasoning_parser_kwargs).unwrap(), + prepared.text_request.reasoning_ended, + ), + expected, + "controls: {controls}" + ); + } + } + #[test] fn prepare_generate_request_gates_continuous_usage_on_include_usage() { let request: GenerateRequest = serde_json::from_value(json!({ diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs index bd277ef34a0a..0a37640d63a1 100644 --- a/rust/src/server/src/routes/inference/generate/types.rs +++ b/rust/src/server/src/routes/inference/generate/types.rs @@ -8,6 +8,7 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use validator::Validate; use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics; +use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; use vllm_text::SamplingParams; use crate::routes::openai::utils::types::{Normalizable, StreamOptions, Usage}; @@ -48,6 +49,13 @@ pub struct GenerateRequest { /// Raw multimodal input; server resolves media. Mutually exclusive with `features`. pub content_parts: Option>, pub return_token_ids: Option, + /// Whether reasoning has ended before the first generated token, as + /// resolved by `/render`. `true` applies structured outputs from the first + /// token; `None` lets the engine check the prompt with its reasoning parser. + pub reasoning_ended: Option, + /// Set by `/render` so the engine-side reasoning parser agrees with the + /// frontend on flags such as `enable_thinking`. + pub reasoning_parser_kwargs: Option, #[serde(flatten)] pub other: Map, } diff --git a/rust/src/server/src/routes/render.rs b/rust/src/server/src/routes/render.rs index 2a6f70ec4b6e..aceb28f4993f 100644 --- a/rust/src/server/src/routes/render.rs +++ b/rust/src/server/src/routes/render.rs @@ -10,6 +10,7 @@ use axum::extract::State; use axum::http::{HeaderMap, StatusCode}; use axum::routing::{get, post}; use thiserror_ext::AsReport as _; +use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; use vllm_text::TextRequest; use crate::DEFAULT_REQUEST_BODY_LIMIT_BYTES; @@ -106,6 +107,12 @@ fn lower_render_request( ec_transfer_params: None, content_parts: None, return_token_ids: None, + // Carry what the chat route would pass to the engine, so the engine + // gates structured outputs the same way on `/inference/v1/generate`. + reasoning_ended: text_request.reasoning_ended, + reasoning_parser_kwargs: (text_request.reasoning_parser_kwargs + != ReasoningParserKwargs::default()) + .then_some(text_request.reasoning_parser_kwargs), other: Default::default(), }; validate_generate_request(&request, &state.served_model_names)?; diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index a9c9831bbd37..118e412b45b6 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -542,7 +542,7 @@ impl ChatRenderer for FakeChatBackend { Ok(vllm_chat::RenderedPrompt { prompt: Prompt::Text(prompt), media_order: None, - effective_template_kwargs: Default::default(), + effective_template_kwargs: request.chat_options.template_kwargs.clone(), }) } } @@ -615,7 +615,7 @@ fn qwen_multimodal_model_info_with_limits( async fn test_models_with_engine_outputs_and_backend_inner( engine_id: impl Into, output_specs: Vec<(Vec, Option)>, - expected_prompt_token_ids: Option>, + check_request: impl FnOnce(&EngineCoreRequest) + Send + 'static, backend: Arc, ) -> (ChatLlm, MockEngineTask) { let ipc = IpcNamespace::new().expect("create ipc namespace"); @@ -630,12 +630,7 @@ async fn test_models_with_engine_outputs_and_backend_inner( let add = recv_engine_message(dealer).await; let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).expect("decode request"); - if let Some(expected_prompt_token_ids) = expected_prompt_token_ids { - assert_eq!( - request.prompt_token_ids.as_deref(), - Some(expected_prompt_token_ids.as_slice()) - ); - } + check_request(&request); send_outputs( push, engine_outputs_for_request(&request.request_id, output_specs), @@ -667,7 +662,8 @@ async fn test_models_with_engine_outputs_and_backend( output_specs: Vec<(Vec, Option)>, backend: Arc, ) -> (ChatLlm, MockEngineTask) { - test_models_with_engine_outputs_and_backend_inner(engine_id, output_specs, None, backend).await + test_models_with_engine_outputs_and_backend_inner(engine_id, output_specs, |_| {}, backend) + .await } async fn test_chat_with_engine_outputs( @@ -1440,6 +1436,8 @@ async fn render_completion_returns_generate_request_with_body_request_id() { assert!(json[0].get("stream_options").is_none()); assert!(json[0].get("prompt_token_ids").is_none()); assert!(json[0].get("mm_features").is_none()); + assert!(json[0].get("reasoning_ended").is_none()); + assert!(json[0].get("reasoning_parser_kwargs").is_none()); } #[tokio::test] @@ -1479,7 +1477,8 @@ async fn render_chat_response_can_be_submitted_to_generate_unchanged() { "messages": [{"role": "user", "content": "hello"}], "max_completion_tokens": 8, "bad_words": ["blocked"], - "vllm_xargs": {"custom": 1} + "vllm_xargs": {"custom": 1}, + "chat_template_kwargs": {"enable_thinking": false} }) .to_string(), )) @@ -1502,11 +1501,27 @@ async fn render_chat_response_can_be_submitted_to_generate_unchanged() { json!(["blocked"]) ); assert_eq!(render_json["sampling_params"]["vllm_xargs"]["custom"], 1); + assert_eq!( + render_json["reasoning_parser_kwargs"], + json!({"chat_template_kwargs": {"enable_thinking": false}}) + ); let (chat, engine_task) = test_models_with_engine_outputs_and_backend_inner( b"engine-render-generate-round-trip", default_stream_output_specs(), - Some(token_ids), + move |request| { + assert_eq!( + request.prompt_token_ids.as_deref(), + Some(token_ids.as_slice()) + ); + assert_eq!( + request + .reasoning_parser_kwargs + .as_ref() + .map(|kwargs| &kwargs.chat_template_kwargs["enable_thinking"]), + Some(&json!(false)) + ); + }, Arc::new(FakeChatBackend::new()), ) .await; From c811943137dbfecdf9c7af849bc66270917a9d79 Mon Sep 17 00:00:00 2001 From: premsurawut <92110316+surawut-jirasaktavee@users.noreply.github.com> Date: Sat, 10 Oct 2026 01:06:26 +0700 Subject: [PATCH 196/624] [Bugfix][Frontend] Reject beam search with echo and logprobs in /v1/completions (#60876) Signed-off-by: surawut.j Co-authored-by: Claude Opus 5.5 --- .../openai/completion/test_completion_error.py | 14 ++++++++++++++ vllm/entrypoints/openai/completion/protocol.py | 11 +++++++++++ 2 files changed, 25 insertions(+) diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index 4056ea49d6d3..bbbb61f9c9c8 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -852,6 +852,20 @@ def test_logprobs_below_minus_one_rejected(): ) +@pytest.mark.parametrize("max_tokens", [5, 0]) +def test_beam_search_echo_logprobs_rejected(max_tokens): + """Beam search does not compute prompt logprobs, so echo + logprobs is rejected.""" + with pytest.raises(VLLMValidationError, match="not supported with beam search"): + CompletionRequest( + model=MODEL_NAME, + prompt="Test prompt", + max_tokens=max_tokens, + use_beam_search=True, + echo=True, + logprobs=1, + ) + + class TestCompletionPromptListLimit: """Regression tests for CVE: unbounded prompt list fan-out.""" diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index f26dea2e1f9e..e99de1e28625 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -492,6 +492,17 @@ def check_logprobs(cls, data): parameter="logprob_token_ids", ) + if ( + data.get("use_beam_search") + and data.get("echo") + and data.get("logprobs") is not None + ): + raise VLLMValidationError( + "`echo` with `logprobs` is not supported with beam search " + "because beam search does not compute prompt logprobs.", + parameter="echo", + ) + if ( data.get("logprob_token_ids") and data.get("echo") From a0f28d73a83bbf0491428efb0f11c2b8f9dc9067 Mon Sep 17 00:00:00 2001 From: Ning Xie Date: Sat, 10 Oct 2026 02:08:31 +0800 Subject: [PATCH 197/624] [sharded state loader] support pp in sharded state loader (#52110) Signed-off-by: Andy Xie Co-authored-by: Simon Mo --- .../save_sharded_state_offline.py | 7 +- .../model_loader/test_sharded_state_loader.py | 438 ++++++++++++++++++ vllm/distributed/parallel_state.py | 5 + .../model_loader/sharded_state_loader.py | 152 ++++-- 4 files changed, 572 insertions(+), 30 deletions(-) diff --git a/examples/features/sharded_state/save_sharded_state_offline.py b/examples/features/sharded_state/save_sharded_state_offline.py index 7ed0de1f1054..44b68cecf0e7 100644 --- a/examples/features/sharded_state/save_sharded_state_offline.py +++ b/examples/features/sharded_state/save_sharded_state_offline.py @@ -38,8 +38,11 @@ def parse_args(): parser.add_argument( "--file-pattern", type=str, - default=ShardedStateLoader.DEFAULT_PATTERN, - help="string pattern of saved filenames", + help=f"string pattern of saved filenames." + f"{ShardedStateLoader.TP_ONLY_PATTERN!r} for tensor parallel only " + f"(pipeline parallelism is 1). " + f"{ShardedStateLoader.PP_AND_TP_PATTERN!r} for pipeline and " + f"tensor parallel.", ) parser.add_argument( "--max-file-size", diff --git a/tests/model_executor/model_loader/test_sharded_state_loader.py b/tests/model_executor/model_loader/test_sharded_state_loader.py index 06354fe3c5ea..6ec5f1ce43ab 100644 --- a/tests/model_executor/model_loader/test_sharded_state_loader.py +++ b/tests/model_executor/model_loader/test_sharded_state_loader.py @@ -6,6 +6,7 @@ import os import shutil from tempfile import TemporaryDirectory +from types import SimpleNamespace import pytest import torch @@ -49,6 +50,318 @@ def test_filter_subtensors(): assert tensor is state_dict[key] +def _mock_pp_group( + monkeypatch: pytest.MonkeyPatch, + world_size: int, + *, + pp_rank: int = 0, + tp_rank: int = 0, + tp_world_size: int = 1, +) -> None: + """Stub the parallel-state accessors used by the sharded state loader.""" + monkeypatch.setattr( + "vllm.distributed.get_pp_group", + lambda: SimpleNamespace(world_size=world_size, rank_in_group=pp_rank), + ) + monkeypatch.setattr( + "vllm.distributed.get_tp_group", + lambda: SimpleNamespace(world_size=tp_world_size, rank_in_group=tp_rank), + ) + monkeypatch.setattr( + "vllm.distributed.get_tensor_model_parallel_rank", + lambda: tp_rank, + ) + monkeypatch.setattr( + "vllm.distributed.get_pipeline_model_parallel_rank", + lambda: pp_rank, + ) + + +@pytest.mark.parametrize( + "world_size,expected", + [ + (1, ShardedStateLoader.TP_ONLY_PATTERN), + (2, ShardedStateLoader.PP_AND_TP_PATTERN), + ], +) +def test_get_default_pattern(monkeypatch, world_size, expected): + """The default pattern is chosen from the pipeline parallel group size.""" + _mock_pp_group(monkeypatch, world_size) + assert ShardedStateLoader._get_default_pattern() == expected + + +def test_render_default_pattern_default_tp(monkeypatch): + """Without a custom pattern and PP=1, legacy TP-only names are rendered.""" + _mock_pp_group(monkeypatch, 1) + assert ( + ShardedStateLoader._render_default_pattern(part=0) + == "model-rank-0-part-0.safetensors" + ) + + +def test_render_default_pattern_default_pp(monkeypatch): + """Without a custom pattern and PP>1, pp/tp-aware names are rendered.""" + _mock_pp_group(monkeypatch, 2) + assert ( + ShardedStateLoader._render_default_pattern(part=0) + == "model-pp-0-tp-0-part-0.safetensors" + ) + + +def test_render_default_pattern_custom_tp_only(monkeypatch): + """A custom TP-only pattern is honored when pipeline parallelism is off.""" + _mock_pp_group(monkeypatch, 1) + assert ( + ShardedStateLoader._render_default_pattern( + ShardedStateLoader.TP_ONLY_PATTERN, part=0 + ) + == "model-rank-0-part-0.safetensors" + ) + + +def test_render_default_pattern_tp_only_with_pp_fails(monkeypatch): + """TP-only filenames would collide across pipeline stages; fail fast.""" + _mock_pp_group(monkeypatch, 2) + with pytest.raises(ValueError, match="does not embed the pipeline parallel"): + ShardedStateLoader._render_default_pattern( + ShardedStateLoader.TP_ONLY_PATTERN, part=0 + ) + + +def test_render_default_pattern_pp_pattern_without_pp_fails(monkeypatch): + """A pp/tp-aware pattern cannot be rendered when pipeline parallelism is off.""" + _mock_pp_group(monkeypatch, 1) + with pytest.raises(ValueError, match="requires pipeline parallelism"): + ShardedStateLoader._render_default_pattern( + ShardedStateLoader.PP_AND_TP_PATTERN, part=0 + ) + + +def test_render_default_pattern_custom_pp(monkeypatch): + """A custom pp/tp-aware pattern is honored when pipeline parallelism is on.""" + _mock_pp_group(monkeypatch, 2) + assert ( + ShardedStateLoader._render_default_pattern( + ShardedStateLoader.PP_AND_TP_PATTERN, part=0 + ) + == "model-pp-0-tp-0-part-0.safetensors" + ) + + +@pytest.mark.parametrize( + "pattern,match", + [ + # tp field present but pp field missing: pipeline stages with the + # same tp rank would write the same file. + ( + "model-tp-{tp_rank}-part-{part}.safetensors", + "does not embed the pipeline parallel", + ), + # pp field present but tp field missing: TP workers would overwrite + # each other's shards. + ( + "model-pp-{pp_rank}-part-{part}.safetensors", + "does not embed the tensor parallel", + ), + ], +) +def test_render_default_pattern_incomplete_custom_patterns_fail( + monkeypatch, pattern, match +): + """A custom pattern must distinguish every active rank (pp and tp).""" + _mock_pp_group(monkeypatch, 2, pp_rank=1, tp_rank=0, tp_world_size=2) + with pytest.raises(ValueError, match=match): + ShardedStateLoader._render_default_pattern(pattern, part=0) + + +def test_render_default_pattern_no_rank_field_with_tp_fails(monkeypatch): + """Without any rank field, TP workers would collide even when PP is off.""" + _mock_pp_group(monkeypatch, 1, tp_world_size=2) + with pytest.raises(ValueError, match="does not embed the tensor parallel"): + ShardedStateLoader._render_default_pattern( + "model-part-{part}.safetensors", part=0 + ) + + +def test_render_default_pattern_custom_pp_tp1_allows_pp_only(monkeypatch): + """With TP=1 a pp-only custom pattern is safe and must be accepted.""" + _mock_pp_group(monkeypatch, 2, pp_rank=1, tp_world_size=1) + assert ( + ShardedStateLoader._render_default_pattern( + "model-pp-{pp_rank}-part-{part}.safetensors", part=0 + ) + == "model-pp-1-part-0.safetensors" + ) + + +@pytest.mark.parametrize( + "pattern,match", + [ + # {{pp_rank}} renders as literal text: both stages would write the + # same file, so it must not count as a pp field. + ( + "model-{{pp_rank}}-tp-{tp_rank}-part-{part}.safetensors", + "does not embed the pipeline parallel", + ), + # Same for an escaped tp field: TP workers would overwrite each + # other's shards. + ( + "model-pp-{pp_rank}-tp-{{tp_rank}}-part-{part}.safetensors", + "does not embed the tensor parallel", + ), + ], +) +def test_render_default_pattern_escaped_field_not_rank_field( + monkeypatch, pattern, match +): + """Escaped braces are literal text and cannot pose as rank fields.""" + _mock_pp_group(monkeypatch, 2, pp_rank=1, tp_rank=0, tp_world_size=2) + with pytest.raises(ValueError, match=match): + ShardedStateLoader._render_default_pattern(pattern, part=0) + + +@pytest.mark.parametrize( + "pattern,world_size,tp_world_size,pp_rank,tp_rank,expected", + [ + # Mixed tp alias works when pipeline parallelism is on. + ( + "model-pp-{pp_rank}-rank-{rank}-part-{part}.safetensors", + 2, + 2, + 1, + 0, + "model-pp-1-rank-0-part-0.safetensors", + ), + # tp_rank alias works for TP-only checkpoints too. + ( + "model-tp-{tp_rank}-part-{part}.safetensors", + 1, + 2, + 0, + 3, + "model-tp-3-part-0.safetensors", + ), + ], +) +def test_render_default_pattern_tp_aliases_accepted( + monkeypatch, pattern, world_size, tp_world_size, pp_rank, tp_rank, expected +): + """Both tp aliases are supplied at render time, so either works.""" + _mock_pp_group( + monkeypatch, + world_size, + pp_rank=pp_rank, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ) + assert ShardedStateLoader._render_default_pattern(pattern, part=0) == expected + + +def test_render_default_pattern_unsupported_field_fails(monkeypatch): + """Unknown placeholders are rejected instead of KeyError-ing later.""" + _mock_pp_group(monkeypatch, 2, pp_rank=1, tp_world_size=2) + with pytest.raises(ValueError, match="unsupported placeholder"): + ShardedStateLoader._render_default_pattern( + "model-{pp_rank}-{tp_rank}-{epoch}-part-{part}.safetensors", part=0 + ) + + +@pytest.mark.parametrize( + "pp_rank,tp_rank,expected", + [ + (0, 0, "model-pp-0-tp-0-part-0.safetensors"), + (1, 2, "model-pp-1-tp-2-part-0.safetensors"), + (3, 5, "model-pp-3-tp-5-part-0.safetensors"), + ], +) +def test_render_default_pattern_non_zero_ranks(monkeypatch, pp_rank, tp_rank, expected): + """The rendered filename must embed the actual pp/tp ranks, not zeros.""" + _mock_pp_group(monkeypatch, 2, pp_rank=pp_rank, tp_rank=tp_rank) + assert ( + ShardedStateLoader._render_default_pattern( + ShardedStateLoader.PP_AND_TP_PATTERN, part=0 + ) + == expected + ) + + +def test_render_default_pattern_part_wildcard(monkeypatch): + """The glob wildcard '*' passes through as the part token for file lookup.""" + _mock_pp_group(monkeypatch, 1, tp_rank=4) + assert ( + ShardedStateLoader._render_default_pattern(part="*") + == "model-rank-4-part-*.safetensors" + ) + + +def test_save_model_pp_aware_filenames(monkeypatch, tmp_path): + """save_model writes pp/tp-aware filenames for pipeline-parallel models.""" + _mock_pp_group(monkeypatch, 2, pp_rank=1, tp_rank=2) + + saved: list[str] = [] + + def _capture(tensors, path, **kwargs): + saved.append(os.path.basename(path)) + + monkeypatch.setattr("safetensors.torch.save_file", _capture) + + model = SimpleNamespace( + state_dict=lambda: {"w": torch.zeros(4, dtype=torch.float32)} + ) + ShardedStateLoader.save_model(model, str(tmp_path)) + assert saved == ["model-pp-1-tp-2-part-0.safetensors"] + + +def test_save_model_part_split(monkeypatch, tmp_path): + """save_model rolls over to a new part when max_size is exceeded.""" + _mock_pp_group(monkeypatch, 1, tp_rank=0) + + saved: list[str] = [] + + def _capture(tensors, path, **kwargs): + saved.append(os.path.basename(path)) + + monkeypatch.setattr("safetensors.torch.save_file", _capture) + + # Two float32[4] tensors (16 B each); max_size=20 fits one but not two. + model = SimpleNamespace( + state_dict=lambda: { + "a": torch.zeros(4, dtype=torch.float32), + "b": torch.zeros(4, dtype=torch.float32), + } + ) + ShardedStateLoader.save_model(model, str(tmp_path), max_size=20) + assert saved == [ + "model-rank-0-part-0.safetensors", + "model-rank-0-part-1.safetensors", + ] + + +def test_save_model_rejects_bad_pattern_before_overwrite(monkeypatch, tmp_path): + """An unsafe pattern is rejected before any shard file is written.""" + _mock_pp_group(monkeypatch, 2, pp_rank=1, tp_rank=0, tp_world_size=2) + + # With the escaped pp field this pattern renders both stages to the + # same literal name; the validation must fire before any write. + colliding = tmp_path / "model-{pp_rank}-tp-0-part-0.safetensors" + colliding.write_bytes(b"existing shard") + before = sorted(p.name for p in tmp_path.iterdir()) + + model = SimpleNamespace( + state_dict=lambda: {"w": torch.zeros(4, dtype=torch.float32)} + ) + with pytest.raises(ValueError, match="does not embed the pipeline parallel"): + ShardedStateLoader.save_model( + model, + str(tmp_path), + pattern="model-{{pp_rank}}-tp-{tp_rank}-part-{part}.safetensors", + ) + + # The pre-existing shard is untouched and nothing new was written. + assert colliding.read_bytes() == b"existing shard" + assert sorted(p.name for p in tmp_path.iterdir()) == before + + @pytest.fixture(scope="module") def llama_3p2_1b_files(): input_dir = hf_api().snapshot_download( @@ -163,3 +476,128 @@ def test_sharded_state_loader( queue.join_thread() assert out_before == out_after + + +def _run_sharded_state_round_trip( + input_dir: str, + *, + tp_size: int, + pp_size: int, + num_gpus_available: int, +) -> None: + """Save a TP/PP sharded checkpoint, reload it, and compare outputs. + + Args: + input_dir: directory containing the original (un-sharded) model. + tp_size: tensor parallel size used for both save and load. + pp_size: pipeline parallel size used for both save and load. + num_gpus_available: number of GPUs on the current machine. + + """ + if num_gpus_available < tp_size * pp_size: + pytest.skip(f"Not enough GPUs for TP={tp_size}, PP={pp_size}") + + weights_patterns = ("*.safetensors",) + gpu_memory_utilization = 0.8 + ctx = mp.get_context("spawn") + + # Keep batching deterministic: this test compares exact greedy outputs + # across separate engine processes, whose schedulers may otherwise form + # different batches depending on prompt-processing timing. + platform_args = {"max_num_seqs": 1} + parallel_args = dict( + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + ) + + with TemporaryDirectory() as output_dir: + p = ctx.Process( + target=_run_writer, + args=(input_dir, output_dir, weights_patterns), + kwargs=dict( + **parallel_args, + gpu_memory_utilization=gpu_memory_utilization, + enforce_eager=True, + **platform_args, + ), + ) + p.start() + p.join() + assert p.exitcode == 0, f"writer process failed with exit code {p.exitcode}" + + if pp_size > 1: + # PP checkpoints must use the pp/tp-aware naming scheme so that + # different pipeline stages write distinct files. + sharded_files = [ + f for f in os.listdir(output_dir) if f.endswith(".safetensors") + ] + assert sharded_files, "no sharded checkpoint files were written" + assert any( + fnmatch.fnmatch(f, "model-pp-*-tp-*-part-*.safetensors") + for f in sharded_files + ), f"expected PP-aware filenames, got {sharded_files}" + + queue = ctx.Queue() + p = ctx.Process( + target=_run_generate, + args=(input_dir, queue), + kwargs=dict( + enable_lora=False, + gpu_memory_utilization=gpu_memory_utilization, + **parallel_args, + **platform_args, + ), + ) + p.start() + # Call queue.get() before p.join() to prevent deadlock: + # If p.join() is called before queue.get() and the queue is full, + # the child process may block while writing to the queue and never + # terminate, causing the parent to wait indefinitely on p.join(). + # See: https://github.com/vllm-project/vllm/pull/22371#discussion_r2257773814 + out_before = queue.get() + p.join() + assert p.exitcode == 0, f"generate process failed with exit code {p.exitcode}" + queue.close() + queue.join_thread() + + queue = ctx.Queue() + p = ctx.Process( + target=_run_generate, + args=(output_dir, queue), + kwargs=dict( + enable_lora=False, + gpu_memory_utilization=gpu_memory_utilization, + load_format="sharded_state", + **parallel_args, + **platform_args, + ), + ) + p.start() + out_after = queue.get() + p.join() + assert p.exitcode == 0, f"generate process failed with exit code {p.exitcode}" + queue.close() + queue.join_thread() + + assert out_before == out_after + + +def test_sharded_state_loader_pp(num_gpus_available, llama_3p2_1b_files): + """Round-trip sharded state save/load with pipeline parallelism (PP=2).""" + _run_sharded_state_round_trip( + llama_3p2_1b_files, + tp_size=1, + pp_size=2, + num_gpus_available=num_gpus_available, + ) + + +def test_sharded_state_loader_pp_tp(num_gpus_available, llama_3p2_1b_files): + """Round-trip sharded state save/load with pipeline + tensor parallelism + (PP=2, TP=2).""" + _run_sharded_state_round_trip( + llama_3p2_1b_files, + tp_size=2, + pp_size=2, + num_gpus_available=num_gpus_available, + ) diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index feb8f25d1eeb..fea0c4e3c783 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -2444,6 +2444,11 @@ def get_tensor_model_parallel_rank() -> int: return get_tp_group().rank_in_group +def get_pipeline_model_parallel_rank() -> int: + """Return my rank for the pipeline model parallel group.""" + return get_pp_group().rank_in_group + + def get_node_count() -> int: """Return the total number of nodes in the distributed environment.""" assert _NODE_COUNT is not None, "distributed environment is not initialized" diff --git a/vllm/model_executor/model_loader/sharded_state_loader.py b/vllm/model_executor/model_loader/sharded_state_loader.py index d17bfdc6d495..97a31e8e879a 100644 --- a/vllm/model_executor/model_loader/sharded_state_loader.py +++ b/vllm/model_executor/model_loader/sharded_state_loader.py @@ -4,6 +4,7 @@ import collections import glob import os +import string import time from collections.abc import Generator from copy import copy @@ -28,13 +29,14 @@ class ShardedStateLoader(BaseModelLoader): """Model loader that directly loads each worker's model state dict, which - enables a fast load path for large tensor-parallel models where each worker - only needs to read its own shard rather than the entire checkpoint. See - `examples/features/sharded_state/save_sharded_state_offline.py` for creating - a sharded checkpoint. + enables a fast load path for large tensor-parallel/pipeline-parallel models + where each worker only needs to read its own shard rather than the entire + checkpoint. See `examples/features/sharded_state/save_sharded_state_offline.py` + for creating a sharded checkpoint. """ - DEFAULT_PATTERN = "model-rank-{rank}-part-{part}.safetensors" + TP_ONLY_PATTERN = "model-rank-{rank}-part-{part}.safetensors" + PP_AND_TP_PATTERN = "model-pp-{pp_rank}-tp-{tp_rank}-part-{part}.safetensors" def __init__(self, load_config: LoadConfig): super().__init__(load_config) @@ -44,14 +46,94 @@ def __init__(self, load_config: LoadConfig): if load_config.model_loader_extra_config is None else copy(load_config.model_loader_extra_config) ) - self.pattern = extra_config.pop("pattern", self.DEFAULT_PATTERN) + # A custom shard filename pattern can be supplied via + # model_loader_extra_config; when omitted (None), the default pattern + # is selected at render time based on the parallel configuration. + self.pattern = extra_config.pop("pattern", None) if extra_config: raise ValueError( f"Unexpected extra config keys for load format " - f"{load_config.load_format}: " - f"{load_config.model_loader_extra_config.keys()}" + f"{load_config.load_format}: {extra_config.keys()}" ) + @staticmethod + def _get_default_pattern() -> str: + from vllm.distributed import get_pp_group + + return ( + ShardedStateLoader.PP_AND_TP_PATTERN + if get_pp_group().world_size > 1 + else ShardedStateLoader.TP_ONLY_PATTERN + ) + + @staticmethod + def _render_default_pattern(pattern: str | None = None, *, part: int | str) -> str: + from vllm.distributed import ( + get_pipeline_model_parallel_rank, + get_pp_group, + get_tensor_model_parallel_rank, + get_tp_group, + ) + + # _get_default_pattern is called here because _render_default_pattern + # can be invoked from save_model without instantiating a + # ShardedStateLoader instance. + if pattern is None: + pattern = ShardedStateLoader._get_default_pattern() + + # Parse the real format fields instead of matching placeholder + # text: escaped braces like {{pp_rank}} render as literal text + # and must not count as a rank field. + fields = { + field + for _, field, _, _ in string.Formatter().parse(pattern) + if field is not None + } + unsupported = fields - {"pp_rank", "tp_rank", "rank", "part"} + if unsupported: + raise ValueError( + f"Pattern '{pattern}' uses unsupported placeholder(s) " + f"{sorted(unsupported)}; supported fields are pp_rank, " + "tp_rank, rank and part" + ) + tp_rank = get_tensor_model_parallel_rank() + pp_rank = get_pipeline_model_parallel_rank() + pp_world_size = get_pp_group().world_size + tp_world_size = get_tp_group().world_size + # The pattern must carry a rank field for every active + # parallelism dimension, otherwise shards from different ranks + # collide on the same filename and silently overwrite each other. + has_pp_field = "pp_rank" in fields + has_tp_field = "tp_rank" in fields or "rank" in fields + if pp_world_size > 1: + if not has_pp_field: + raise ValueError( + f"Pattern '{pattern}' does not embed the pipeline parallel " + "rank; with pipeline parallelism every stage would write the " + f"same file. Use '{ShardedStateLoader.PP_AND_TP_PATTERN}' " + "instead." + ) + if tp_world_size > 1 and not has_tp_field: + raise ValueError( + f"Pattern '{pattern}' does not embed the tensor parallel " + "rank; TP workers would overwrite each other's shards." + ) + # Supply both tp aliases so mixed-alias patterns render. + return pattern.format( + pp_rank=pp_rank, tp_rank=tp_rank, rank=tp_rank, part=part + ) + if has_pp_field: + raise ValueError( + f"Pattern '{pattern}' requires pipeline parallelism, but the " + "pipeline parallel size is 1" + ) + if tp_world_size > 1 and not has_tp_field: + raise ValueError( + f"Pattern '{pattern}' does not embed the tensor parallel rank; " + "TP workers would overwrite each other's shards." + ) + return pattern.format(tp_rank=tp_rank, rank=tp_rank, part=part) + @staticmethod def _filter_subtensors( tensors: dict[str, torch.Tensor], @@ -106,30 +188,36 @@ def download_model(self, model_config: ModelConfig) -> None: self._prepare_weights(model_config.model, model_config.revision) def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: - from vllm.distributed import get_tensor_model_parallel_rank - model_weights = model_config.model if model_weights_override := model_config.model_weights: model_weights = model_weights_override local_model_path = model_weights - rank = get_tensor_model_parallel_rank() - pattern = os.path.join( - local_model_path, - self.pattern.format(rank=rank, part="*"), + rendered_pattern = ShardedStateLoader._render_default_pattern( + self.pattern, part="*" ) - filepaths = [] if is_s3(local_model_path): - file_pattern = f"*{self.pattern.format(rank=rank, part='*')}" + file_pattern = f"*{rendered_pattern}" filepaths = s3_glob(path=local_model_path, allow_pattern=[file_pattern]) else: - filepaths = glob.glob(pattern) + filepaths = glob.glob( + os.path.join( + local_model_path, + rendered_pattern, + ) + ) if not filepaths: # TODO: support un-sharded checkpoints too raise ValueError( - f"Could not find checkpoint files '{pattern}', only " - f"pre-sharded checkpoints are currently supported!" + f"Could not find checkpoint files " + f"'{ + os.path.join( + local_model_path, + rendered_pattern, + ) + }', " + f"only pre-sharded checkpoints are currently supported!" ) state_dict = self._filter_subtensors(model.state_dict()) counter_before_loading_weights = time.perf_counter() @@ -137,8 +225,17 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: # If loading with LoRA enabled, additional padding may # be added to certain parameters. We only load into a # narrowed view of the parameter data. - param_data = state_dict[key].data - param_shape = state_dict[key].shape + param = state_dict[key] + if param.dtype != tensor.dtype: + logger.warning( + "Attempted to load weight %s with dtype %s into " + "parameter with dtype %s", + key, + tensor.dtype, + param.dtype, + ) + param_data = param.data + param_shape = param.shape for dim, size in enumerate(tensor.shape): if size < param_shape[dim]: param_data = param_data.narrow(dim, 0, size) @@ -182,11 +279,6 @@ def save_model( ) -> None: from safetensors.torch import save_file - from vllm.distributed import get_tensor_model_parallel_rank - - if pattern is None: - pattern = ShardedStateLoader.DEFAULT_PATTERN - rank = get_tensor_model_parallel_rank() part_idx = 0 total_size = 0 state_dict = ShardedStateLoader._filter_subtensors(model.state_dict()) @@ -194,7 +286,9 @@ def save_model( for key, tensor in state_dict.items(): param_size = tensor.nelement() * tensor.element_size() if max_size is not None and total_size + param_size > max_size: - filename = pattern.format(rank=rank, part=part_idx) + filename = ShardedStateLoader._render_default_pattern( + pattern, part=part_idx + ) save_file( state_dict_part, os.path.join(path, filename), @@ -205,7 +299,9 @@ def save_model( state_dict_part[key] = tensor total_size += param_size if len(state_dict_part) > 0: - filename = pattern.format(rank=rank, part=part_idx) + filename = ShardedStateLoader._render_default_pattern( + pattern, part=part_idx + ) save_file( state_dict_part, os.path.join(path, filename), From 1dc60dd36734f4436dbc438265bbd9c24164cbad Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Fri, 9 Oct 2026 14:42:41 -0400 Subject: [PATCH 198/624] [CI] Pin pydantic in the mypy pre-commit hook (#60894) Signed-off-by: mgoin Co-authored-by: Claude --- .pre-commit-config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 92963cf69a22..8e257edccc12 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -205,7 +205,7 @@ repos: language: python types_or: [python, pyi] require_serial: true - additional_dependencies: ["mypy==1.20.2", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic] + additional_dependencies: ["mypy==1.20.2", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, "pydantic==2.13.5"] - id: png-lint name: Lint PNG exports from excalidraw entry: tools/pre_commit/png-lint.sh From ad8a6065f4bcbdfafd38828a4b3f1c6fac4dbf89 Mon Sep 17 00:00:00 2001 From: Willow Lopez <100782273+Oxygen56@users.noreply.github.com> Date: Sat, 10 Oct 2026 02:59:02 +0800 Subject: [PATCH 199/624] [Bugfix] Validate guidance tokenizer compatibility before EngineCore (#52298) Signed-off-by: Oxygen56 Signed-off-by: Artem Perevedentsev Co-authored-by: Artem Perevedentsev --- tests/v1/structured_output/test_validation.py | 97 ++++++++++++++++++- vllm/sampling_params.py | 19 +++- vllm/v1/structured_output/backend_guidance.py | 44 ++++++++- 3 files changed, 152 insertions(+), 8 deletions(-) diff --git a/tests/v1/structured_output/test_validation.py b/tests/v1/structured_output/test_validation.py index cd5526c7b760..276056b501cd 100644 --- a/tests/v1/structured_output/test_validation.py +++ b/tests/v1/structured_output/test_validation.py @@ -5,6 +5,7 @@ import json import pytest +from transformers import MistralCommonBackend, TokenizersBackend from vllm.config import StructuredOutputsConfig from vllm.exceptions import VLLMClientError, VLLMValidationError @@ -14,6 +15,8 @@ SamplingParams, StructuredOutputsParams, ) +from vllm.tokenizers import mistral as mistral_tokenizers +from vllm.v1.structured_output import backend_guidance pytestmark = pytest.mark.cpu_test @@ -278,10 +281,102 @@ def test_unsupported_grammar_is_a_client_error(backend, structured_outputs): def test_auto_backend_falls_back_on_unsupported_schema(schema, expected_backend): """`auto` falls back on rejection, so it must catch what the validators raise.""" params = SamplingParams(structured_outputs=StructuredOutputsParams(json=schema)) + # trick to create tokenizer that guidance backend accepts + tokenizer = object.__new__(TokenizersBackend) params._validate_structured_outputs( _StubModelConfig(is_diffusion=False), StructuredOutputsConfig(backend="auto"), - tokenizer=object(), + tokenizer=tokenizer, ) assert params.structured_outputs is not None assert params.structured_outputs._backend == expected_backend + + +# ================================================ +# Test validate_structured_outputs functionality +# ================================================ + + +class TestValidateStructuredOutputs: + class TestPR52298Regression: + class _StubSlowTokenizer: + is_fast = False + + class _StubVllmMistralTokenizer: + IS_MISTRAL_TOKENIZER = True + + def __init__(self, is_tekken: bool): + self.is_tekken = is_tekken + self.llg_tokenizer = object() + + @staticmethod + def _validate_guidance(tokenizer): + params = SamplingParams( + structured_outputs=StructuredOutputsParams(json=JSON_SCHEMA) + ) + params._validate_structured_outputs( + _StubModelConfig(is_diffusion=False), + StructuredOutputsConfig(backend="guidance"), + tokenizer=tokenizer, + ) + return params + + def _make_tokenizer(self, tokenizer_id: str): + if tokenizer_id == "slow-hf": + return self._StubSlowTokenizer() + if tokenizer_id == "plain-object": + return object() + if tokenizer_id == "raw-mistral-common": + return object.__new__(MistralCommonBackend) + if tokenizer_id == "vllm-mistral-non-tekken": + return self._StubVllmMistralTokenizer(is_tekken=False) + if tokenizer_id == "fast-hf": + return object.__new__(TokenizersBackend) + if tokenizer_id == "vllm-mistral-tekken": + return self._StubVllmMistralTokenizer(is_tekken=True) + raise ValueError(tokenizer_id) + + @pytest.mark.parametrize( + "tokenizer_id, patch_mistral", + [ + pytest.param("slow-hf", False, id="slow-hf"), + pytest.param("plain-object", False, id="plain-object"), + pytest.param("raw-mistral-common", False, id="raw-mistral-common"), + pytest.param( + "vllm-mistral-non-tekken", + True, + id="vllm-mistral-non-tekken", + ), + ], + ) + def test_unsupported_tokenizer(self, monkeypatch, tokenizer_id, patch_mistral): + if patch_mistral: + monkeypatch.setattr( + mistral_tokenizers, + "MistralTokenizer", + self._StubVllmMistralTokenizer, + ) + with pytest.raises(VLLMValidationError): + self._validate_guidance(self._make_tokenizer(tokenizer_id)) + + @pytest.mark.parametrize( + "tokenizer_id, patch_mistral", + [ + pytest.param("fast-hf", False, id="fast-hf"), + pytest.param("vllm-mistral-tekken", True, id="vllm-mistral-tekken"), + ], + ) + def test_supported_tokenizer(self, monkeypatch, tokenizer_id, patch_mistral): + if patch_mistral: + monkeypatch.setattr( + mistral_tokenizers, + "MistralTokenizer", + self._StubVllmMistralTokenizer, + ) + monkeypatch.setattr( + backend_guidance, + "validate_guidance_grammar", + lambda *_args, **_kwargs: None, + ) + params = self._validate_guidance(self._make_tokenizer(tokenizer_id)) + assert params.structured_outputs._backend == "guidance" diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 8a37c5a3fd17..8331a1707b7d 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -1317,6 +1317,7 @@ def _validate_structured_outputs( ) from vllm.v1.structured_output.backend_guidance import ( + GuidanceBackend, has_guidance_unsupported_json_features, validate_guidance_grammar, ) @@ -1340,6 +1341,13 @@ def _validate_structured_outputs( "Mistral model, the ['xgrammar', 'outlines'] " "backends or tokenizer_mode='hf' instead." ) + if not GuidanceBackend.is_tokenizer_supported(tokenizer): + raise VLLMValidationError( + "The 'guidance' structured output backend only supports fast " + "Hugging Face tokenizers and Tekken-based Mistral tokenizers. " + "Please use the 'xgrammar' backend or configure a fast " + "tokenizer instead." + ) # TODO: ideally we would have the LLTokenizer here as Lark syntax # allows <|special_token|> and similar, see # https://github.com/guidance-ai/llguidance/blob/main/docs/syntax.md#special-tokens @@ -1408,12 +1416,19 @@ def _validate_structured_outputs( skip_guidance = has_guidance_unsupported_json_features(schema) if skip_guidance: - # Fall back to outlines if the tokenizer is non-tekken Mistral or - # the schema contains features unsupported by guidance + # Fall back to outlines for non-Tekken Mistral tokenizers or + # schemas containing features unsupported by guidance. validate_structured_output_request_outlines(self) self.structured_outputs._backend = "outlines" else: # Fall back to guidance by default. + if not GuidanceBackend.is_tokenizer_supported(tokenizer): + raise VLLMValidationError( + "No compatible structured output backend was found for " + "this request. XGrammar rejected the constraint, and " + "Guidance does not support this tokenizer. Configure a " + "fast tokenizer or select a compatible backend explicitly." + ) from None validate_guidance_grammar( self, tokenizer=_get_llg_tokenizer(tokenizer), diff --git a/vllm/v1/structured_output/backend_guidance.py b/vllm/v1/structured_output/backend_guidance.py index 6dc2dc91ff76..4af5ca2da7ac 100644 --- a/vllm/v1/structured_output/backend_guidance.py +++ b/vllm/v1/structured_output/backend_guidance.py @@ -8,10 +8,12 @@ from typing import TYPE_CHECKING, Any import torch +from transformers import TokenizersBackend from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.sampling_params import SamplingParams +from vllm.tokenizers import TokenizerLike from vllm.utils.import_utils import LazyLoader from vllm.utils.mistral import is_mistral_tokenizer from vllm.v1.structured_output.backend_types import ( @@ -128,6 +130,37 @@ def process_for_additional_properties( @dataclass class GuidanceBackend(StructuredOutputBackend): + @staticmethod + def is_tokenizer_supported(tokenizer: TokenizerLike) -> bool: + """Check if tokenizer is supported by guidance backend. + + TODO(arpera): + guidance backend does NOT support some types of tokenizers + so, we MUST verify tokenizer before using guidance + The problem is that guidance API does not provide + a way to validate tokenizer before building LLTokenizer + So, we have to do it manually here in vLLM. + This is ofcourse very bad, we duplicate complex logic + that can change arbitrary in future. + Our check here is based on guidance source code + on function `from_tokenizer` implementation: + https://github.com/guidance-ai/llguidance/blob/main/python/llguidance/hf.py#L29-L47 + If you want to contribute to vLLM, please, add to llguidance + API function to implement this check there. + In that case we will remove this check from vLLM. + That kind of work is really appreciated! + """ + if is_mistral_tokenizer(tokenizer): + return tokenizer.is_tekken + # PreTrainedTokenizerFast is the same as TokenizersBackend + return isinstance(tokenizer, TokenizersBackend) + + @staticmethod + def _build_ll_tokenizer(tokenizer: TokenizerLike, vocab_size: int) -> Any: + if is_mistral_tokenizer(tokenizer): + return tokenizer.llg_tokenizer + return llguidance_hf.from_tokenizer(tokenizer, max(vocab_size, len(tokenizer))) + def __post_init__(self): self.disable_any_whitespace = ( self.vllm_config.structured_outputs_config.disable_any_whitespace @@ -136,12 +169,13 @@ def __post_init__(self): self.vllm_config.structured_outputs_config.disable_additional_properties ) - if is_mistral_tokenizer(self.tokenizer): - self.ll_tokenizer = self.tokenizer.llg_tokenizer - else: - self.ll_tokenizer = llguidance_hf.from_tokenizer( - self.tokenizer, max(self.vocab_size, len(self.tokenizer)) + if not self.is_tokenizer_supported(self.tokenizer): + raise ValueError( + "Only fast tokenizers are supported." + "This error should NOT happen," + "if you see this, please report an issue." ) + self.ll_tokenizer = self._build_ll_tokenizer(self.tokenizer, self.vocab_size) def compile_grammar( self, From 04462732a40a8329ad60593d754b576c54c5fe5d Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Fri, 9 Oct 2026 23:15:10 +0400 Subject: [PATCH 200/624] [Bugfix] Reject max_num_scheduled_tokens=0 at config time instead of hanging (#60714) Signed-off-by: Vadim Gimpelson Co-authored-by: Claude Co-authored-by: Hui Li <40770106+untrall@users.noreply.github.com> --- tests/test_config.py | 6 ++++++ vllm/config/scheduler.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/test_config.py b/tests/test_config.py index ebf76c76a0c6..425d55f29346 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -3486,6 +3486,12 @@ def test_scheduler_config_init(): print(SchedulerConfig.default_factory().max_model_len) +def test_scheduler_config_rejects_zero_max_num_scheduled_tokens(): + """A zero token budget never schedules anything, so generation would hang.""" + with pytest.raises(ValidationError, match="max_num_scheduled_tokens"): + SchedulerConfig.default_factory(max_num_scheduled_tokens=0) + + @pytest.mark.parametrize( ( "model_id", diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index 0d10103ed4ec..1dbe4b72eb7c 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -52,7 +52,7 @@ class SchedulerConfig: In real usage, this should be set in `EngineArgs.create_engine_config`. """ - max_num_scheduled_tokens: int | None = Field(default=None, ge=0) + max_num_scheduled_tokens: int | None = Field(default=None, ge=1) """Maximum number of tokens that the scheduler may issue in a single iteration. This is usually equal to max_num_batched_tokens, but can be smaller in cases From 3766a67a32f7ea8130d84383f72883cee5e8b006 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Fri, 9 Oct 2026 15:34:58 -0400 Subject: [PATCH 201/624] [WideEP] Explain how to fix a too-old NCCL for DeepEP v2 (#60449) Signed-off-by: mgoin Co-authored-by: Claude --- tools/ep_kernels/README.md | 28 ++++++++----- tools/ep_kernels/install_python_libraries.sh | 26 ++++++++++-- .../device_communicators/all2all.py | 6 +-- vllm/utils/import_utils.py | 41 ++++++++++++------- 4 files changed, 70 insertions(+), 31 deletions(-) diff --git a/tools/ep_kernels/README.md b/tools/ep_kernels/README.md index 04bc3ef27680..286a6db7bcf9 100644 --- a/tools/ep_kernels/README.md +++ b/tools/ep_kernels/README.md @@ -11,12 +11,16 @@ Step 2 is necessary for multi-node deployment. All scripts accept a positional argument as workspace path for staging the build, defaulting to `$(pwd)/ep_kernels_workspace`. -## NCCL version requirement (CUDA 13+) +## NCCL version requirement for DeepEPv2 -DeepEPv2 uses the NCCL GIN (GPU-Initiated Networking) backend, which requires -NCCL >= 2.30.4 at both compile time and runtime. PyTorch 2.11 pins -`nvidia-nccl-cu13==2.28.9` as a transitive dependency, so you need to -override it. +DeepEPv2 (`--all2all-backend deepep_v2`) uses the NCCL GIN (GPU-Initiated +Networking) backend, which requires NCCL >= 2.30.4 both when DeepEP is built +and at runtime. PyTorch pins an older release as a dependency (PyTorch 2.13 +pins `nvidia-nccl-cu13==2.29.7`), so a plain `pip install` or `uv pip install` +of vLLM is not enough. The vLLM Docker images already override it. + +Upgrade NCCL before running `install_python_libraries.sh`. Use +`nvidia-nccl-cu12` on CUDA 12. **With uv** (recommended): @@ -29,18 +33,22 @@ export UV_OVERRIDE=/tmp/nccl-override.txt uv pip install vllm ``` -**With pip**: +Keep `UV_OVERRIDE` set for later installs into the same environment. Without +it, any `uv pip install` that resolves PyTorch again restores the pinned NCCL. + +**With pip**, or to fix an existing environment: ```bash -pip install vllm pip install "nvidia-nccl-cu13>=2.30.4" --no-deps ``` -The override / reinstall must happen before building DeepEP (for GIN device -headers) and must remain in place at runtime. You can verify with: +`install_python_libraries.sh` warns when the installed NCCL is too old. If +DeepEP was built against an older NCCL, upgrade NCCL and run the script again. +When NCCL is too old at runtime, vLLM logs the reason at startup and +`deepep_v2` fails with the same message. You can check with: ```bash -python -c "from vllm.utils.import_utils import has_deep_ep_v2; print(has_deep_ep_v2())" +python -c "from vllm.utils.import_utils import deep_ep_v2_unavailable_reason; print(deep_ep_v2_unavailable_reason() or 'deepep_v2 is available')" ``` ## Usage diff --git a/tools/ep_kernels/install_python_libraries.sh b/tools/ep_kernels/install_python_libraries.sh index 9ca6a8d6a99e..fdd7aa20e957 100755 --- a/tools/ep_kernels/install_python_libraries.sh +++ b/tools/ep_kernels/install_python_libraries.sh @@ -78,10 +78,25 @@ mkdir -p "$WHEEL_DIR" pushd "$WORKSPACE" # install dependencies if not installed +UV_PIP_TARGET=() if [ -z "$VIRTUAL_ENV" ]; then - uv pip install --system cmake torch ninja -else - uv pip install cmake torch ninja + UV_PIP_TARGET=(--system) +fi +uv pip install "${UV_PIP_TARGET[@]}" cmake torch ninja + +# DeepEPv2 needs NCCL >= 2.30.4 at both build time and runtime, but PyTorch +# pins an older release. +DEEPEP_V2_MIN_NCCL="2.30.4" +NCCL_PACKAGE="nvidia-nccl-cu${CUDA_VERSION_MAJOR}" +NCCL_WARNING="" +NCCL_INSTALLED=$(uv pip show "${UV_PIP_TARGET[@]}" "$NCCL_PACKAGE" 2>/dev/null | sed -n 's/^Version: //p') +if [ -n "$NCCL_INSTALLED" ] && \ + [ "$(printf '%s\n' "$DEEPEP_V2_MIN_NCCL" "$NCCL_INSTALLED" | sort -V | head -n1)" != "$DEEPEP_V2_MIN_NCCL" ]; then + NCCL_WARNING="WARNING: ${NCCL_PACKAGE} ${NCCL_INSTALLED} is installed, but the deepep_v2 backend +requires NCCL >= ${DEEPEP_V2_MIN_NCCL} when DeepEP is built and when it runs. To use deepep_v2, run + uv pip install \"${NCCL_PACKAGE}>=${DEEPEP_V2_MIN_NCCL}\" --no-deps +and then run this script again. See tools/ep_kernels/README.md." + echo "$NCCL_WARNING" >&2 fi # fetch nvshmem @@ -236,3 +251,8 @@ if [ "$MODE" = "wheel" ]; then echo "All wheels written to $WHEEL_DIR" ls -l "$WHEEL_DIR" fi + +# Repeat after the verbose build output, where the first copy is easy to miss. +if [ -n "$NCCL_WARNING" ]; then + echo "$NCCL_WARNING" >&2 +fi diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index ae31c0b1d3cd..cff34139f968 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -23,6 +23,7 @@ from vllm.utils.func_utils import supports_kw from vllm.utils.import_utils import ( check_moonep_system_support, + deep_ep_v2_unavailable_reason, has_deep_ep, has_deep_ep_v2, has_mori, @@ -1104,10 +1105,7 @@ class DeepEPV2All2AllManager(All2AllManagerBase): """ def __init__(self, cpu_group, tcp_store_group=None, device_group=None): - assert has_deep_ep_v2(), ( - "DeepEP v2 (ElasticBuffer) not available. Requires DeepEP >= 2.0 " - "(https://github.com/deepseek-ai/DeepEP) and NCCL >= 2.30.4." - ) + assert has_deep_ep_v2(), deep_ep_v2_unavailable_reason() super().__init__(cpu_group, tcp_store_group) self._device_group = device_group self.handle_cache = Cache() diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index 04a7ebec2eab..277f987d4800 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -469,31 +469,44 @@ def _format_nccl_raw_version(raw: int) -> str: return f"{s[0]}.{s[1:3].lstrip('0') or '0'}.{s[3:].lstrip('0') or '0'}" -def has_deep_ep_v2() -> bool: - """Whether deep_ep with ElasticBuffer (v2 API) is available. +def deep_ep_v2_unavailable_reason() -> str | None: + """Why the deep_ep v2 (ElasticBuffer) API cannot be used, or None if it can. Requires both the ElasticBuffer class in the deep_ep module and NCCL >= 2.30.4 (GIN backend), checked against the runtime library. """ + install_doc = "See tools/ep_kernels/README.md." if not _has_module("deep_ep"): - return False + return f"DeepEP is not installed. {install_doc}" import deep_ep # type: ignore[import-not-found] if not hasattr(deep_ep, "ElasticBuffer"): - return False + return f"The installed DeepEP predates v2 (no ElasticBuffer). {install_doc}" + nccl_ver = _get_runtime_nccl_version() + if nccl_ver is None or nccl_ver < DEEPEP_V2_MIN_NCCL_VERSION_RAW: + import torch + + min_version = _format_nccl_raw_version(DEEPEP_V2_MIN_NCCL_VERSION_RAW) + cuda_major = (torch.version.cuda or "13").split(".")[0] + return ( + f"DeepEP v2 requires NCCL >= {min_version} but found " + f"{_format_nccl_raw_version(nccl_ver) if nccl_ver else 'unknown'}. " + "PyTorch pins an older NCCL; upgrade it with " + f'`pip install "nvidia-nccl-cu{cuda_major}>={min_version}" --no-deps` ' + f"and rebuild DeepEP against it. {install_doc}" + ) + return None + + +def has_deep_ep_v2() -> bool: + """Whether deep_ep with ElasticBuffer (v2 API) is available.""" try: - nccl_ver = _get_runtime_nccl_version() - if nccl_ver is None or nccl_ver < DEEPEP_V2_MIN_NCCL_VERSION_RAW: - logger.info_once( - "DeepEP v2 requires NCCL >= %s but found %s. " - "deepep_v2 backend will not be available.", - _format_nccl_raw_version(DEEPEP_V2_MIN_NCCL_VERSION_RAW), - _format_nccl_raw_version(nccl_ver) if nccl_ver else "unknown", - ) - return False + reason = deep_ep_v2_unavailable_reason() except Exception: return False - return True + if reason is not None and _has_module("deep_ep"): + logger.info_once("deepep_v2 backend will not be available: %s", reason) + return reason is None def has_deep_gemm() -> bool: From 02e55d2af38adab51ec5ffb985063d355b20805b Mon Sep 17 00:00:00 2001 From: Dakai An <77474977+andakai@users.noreply.github.com> Date: Sat, 10 Oct 2026 03:48:37 +0800 Subject: [PATCH 202/624] [KV Cache] GLM-5.3-Flash: use the generic packed KV layout; CircularBufferSpec tail with speculative ring slots (#57169) Signed-off-by: Yifan Qiao Signed-off-by: Dakai An Signed-off-by: Lucas Wilkinson Co-authored-by: Yifan Qiao Co-authored-by: OpenAI Codex Co-authored-by: OpenAI Codex Co-authored-by: Lucas Wilkinson Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Lucas Wilkinson --- .../attention/test_deepgemm_attention.py | 46 ++ .../attention/test_rocm_aiter_mla_sink.py | 4 +- .../test_jit_warmup_triton_launcher.py | 3 - .../test_dspark_noncausal_sparse_mla.py | 2 +- .../test_flashinfer_mla_sparse_sm90.py | 18 +- ...st_flashmla_nope_sm90_backend_selection.py | 97 +++ .../test_glm5next_indexer_block_sizes.py | 40 ++ .../test_indexer_deepseek_v4_slot_mapping.py | 19 +- .../test_kpool_indexer_block_sizes.py | 54 +- .../attention/test_kpool_tail_slot_mapping.py | 45 +- tests/v1/attention/test_mm_prefix.py | 1 - .../v1/attention/test_rocm_glm5next_sparse.py | 6 + .../v1/attention/test_sparse_mla_backends.py | 69 +- .../test_sparse_mla_kv_cache_layout.py | 21 +- .../test_partial_prefix_cache_hits.py | 8 +- tests/v1/core/test_contiguous_kv_packing.py | 42 +- tests/v1/core/test_kv_cache_utils.py | 600 ++---------------- .../core/test_single_type_kv_cache_manager.py | 14 +- tests/v1/kv_connector/unit/conftest.py | 2 +- .../test_mooncake_connector_hybrid_mamba.py | 9 +- .../unit/test_mooncake_store_hma_e2e.py | 4 +- .../unit/test_nixl_desc_geometry.py | 187 +----- tests/v1/simple_kv_offload/test_scheduler.py | 4 +- .../spec_decode/test_dflash_prepare_inputs.py | 2 +- tests/v1/spec_decode/test_eagle.py | 6 + .../v1/spec_decode/test_llm_base_proposer.py | 24 +- .../test_gpu_model_runner_streaming.py | 1 - tests/v1/worker/test_attn_utils.py | 112 +++- tests/v1/worker/test_gpu_block_table.py | 31 +- tests/v1/worker/test_gpu_input_batch.py | 9 - .../test_gpu_kpool_tail_slot_mapping.py | 8 +- tests/v1/worker/test_gpu_model_runner.py | 108 +--- ...odel_runner_prefill_tail_classification.py | 3 - tests/v1/worker/test_gpu_model_runner_v2.py | 1 - tests/v1/worker/test_gpu_pcp_manager.py | 3 +- ...test_gpu_target_dependent_ar_speculator.py | 51 ++ tests/v1/worker/test_jit_warmup_migration.py | 18 +- .../worker/test_kv_cache_allocation_scope.py | 2 + tests/v1/worker/test_mamba_utils.py | 1 - tests/v1/worker/test_utils.py | 37 ++ .../kv_transfer/kv_connector/utils.py | 30 + .../v1/mooncake/mooncake_connector.py | 28 +- .../kv_connector/v1/nixl/base_worker.py | 124 +--- .../layers/attention/mla_attention.py | 8 - .../models/deepseek_v41/nvidia/model_state.py | 2 +- vllm/models/glm5next/common/attention.py | 64 +- .../glm5next/nvidia/ops/kpool_compress.py | 5 - vllm/platforms/cuda.py | 20 - vllm/platforms/interface.py | 39 +- vllm/utils/deep_gemm.py | 9 +- vllm/v1/attention/backend.py | 8 +- vllm/v1/attention/backends/composite.py | 5 - .../backends/mla/compressor_utils.py | 4 +- .../backends/mla/flashattn_mla_sparse.py | 7 +- .../backends/mla/flashinfer_mla_sparse.py | 22 +- .../mla/flashinfer_mla_sparse_sm120.py | 8 +- .../mla/flashinfer_mla_sparse_sm90.py | 25 +- .../attention/backends/mla/flashmla_sparse.py | 51 +- vllm/v1/attention/backends/mla/indexer.py | 66 +- .../backends/mla/rocm_aiter_mla_sparse.py | 25 +- .../v1/attention/backends/mla/sparse_utils.py | 9 +- .../attention/backends/mla/xpu_mla_sparse.py | 5 + vllm/v1/core/kv_cache_coordinator.py | 8 +- vllm/v1/core/kv_cache_utils.py | 354 ++--------- vllm/v1/core/single_type_kv_cache_manager.py | 13 - vllm/v1/engine/core.py | 4 +- vllm/v1/hisparse/binding.py | 8 +- vllm/v1/hisparse/layout.py | 4 +- vllm/v1/kv_cache_interface.py | 56 +- vllm/v1/spec_decode/gemma4.py | 8 +- vllm/v1/spec_decode/llm_base_proposer.py | 17 +- vllm/v1/spec_decode/step3p5.py | 8 +- vllm/v1/worker/block_table.py | 129 +--- vllm/v1/worker/cp_utils.py | 6 +- vllm/v1/worker/gpu/attn_utils.py | 24 +- vllm/v1/worker/gpu/block_table.py | 26 +- vllm/v1/worker/gpu/input_batch.py | 17 +- vllm/v1/worker/gpu/model_runner.py | 8 +- .../gpu/spec_decode/dflash/speculator.py | 2 +- .../target_dependent_ar/speculator.py | 5 + vllm/v1/worker/gpu_input_batch.py | 2 - vllm/v1/worker/gpu_model_runner.py | 41 +- vllm/v1/worker/utils.py | 252 ++++++-- 83 files changed, 1281 insertions(+), 1987 deletions(-) create mode 100644 tests/v1/attention/test_glm5next_indexer_block_sizes.py diff --git a/tests/kernels/attention/test_deepgemm_attention.py b/tests/kernels/attention/test_deepgemm_attention.py index 5d11063c7a2b..f332fb0fea21 100644 --- a/tests/kernels/attention/test_deepgemm_attention.py +++ b/tests/kernels/attention/test_deepgemm_attention.py @@ -313,3 +313,49 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int): ref_logits = ref_logits.masked_fill(~mask, 0) diff = calc_diff(logits, ref_logits) assert diff < 1e-3, f"{diff=}" + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") +@pytest.mark.skipif( + not current_platform.has_device_capability(90), reason="SM90 and SM100 only" +) +def test_deepgemm_paged_mqa_packed_manager_block_stride(): + torch.manual_seed(0) + num_blocks, block_size, packed_rows = 4, 256, 320 + num_heads, head_dim, row_bytes = 64, 128, 132 + page_size, pages_per_block = 64, block_size // 64 + dense_stride, packed_stride = pages_per_block, packed_rows // page_size + compact_pages = kv_cache_cast_to_fp8( + torch.randn( + num_blocks, block_size, 1, head_dim, device="cuda", dtype=torch.bfloat16 + ).view(num_blocks * pages_per_block, page_size, 1, head_dim) + ) + packed_backing = torch.zeros( + num_blocks, + packed_rows, + row_bytes, + device="cuda", + dtype=torch.uint8, + ) + packed_manager = packed_backing[:, :block_size] + packed_manager.copy_(compact_pages.squeeze(2).view(num_blocks, block_size, -1)) + packed_pages = packed_backing.view(-1, page_size, row_bytes).unsqueeze(2) + offsets = torch.arange(pages_per_block, device="cuda", dtype=torch.int32) + tables = [ + (stride + offsets).unsqueeze(0) for stride in (dense_stride, packed_stride) + ] + + q = torch.randn(1, 1, num_heads, head_dim, device="cuda", dtype=torch.bfloat16).to( + torch.float8_e4m3fn + ) + weights = torch.randn(1, num_heads, device="cuda", dtype=torch.float32) + context_lens = torch.full((1, 1), block_size, device="cuda", dtype=torch.int32) + schedule = get_paged_mqa_logits_metadata(context_lens, page_size, get_num_sms()) + logits = [ + fp8_fp4_paged_mqa_logits( + (q, None), cache, weights, context_lens, table, schedule, block_size, False + ) + for cache, table in zip((compact_pages, packed_pages), tables) + ] + torch.testing.assert_close(*logits, rtol=0, atol=0) diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sink.py b/tests/kernels/attention/test_rocm_aiter_mla_sink.py index 0b007dbc875c..f66d408c8495 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sink.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sink.py @@ -463,7 +463,7 @@ def test_sparse_mla_sink_forward_mqa_preserves_split_query(dtype): set_random_seed(412) num_tokens, num_heads, block_size = 2, 8, 16 q = torch.randn(num_tokens, num_heads, Q_HEAD_DIM, device="cuda").to(dtype) - kv = torch.randn(2 * block_size, 1, Q_HEAD_DIM, device="cuda").to(dtype) + kv = torch.randn(2, 2, block_size, Q_HEAD_DIM, device="cuda").to(dtype)[:, 1] selected = torch.tensor([[1, 7, 18], [0, 3, 20]], dtype=torch.int32, device="cuda") sinks = torch.linspace(-3.0, 5.0, num_heads, device="cuda") impl = object.__new__(ROCMAiterMLASparseImpl) @@ -504,7 +504,7 @@ def test_sparse_mla_sink_forward_mqa_preserves_split_query(dtype): SimpleNamespace(_q_scale=None, _k_scale=None), ) references = [ - _sink_reference(q[i], kv[:, 0][selected[i].long()], sinks, SM_SCALE) + _sink_reference(q[i], kv.flatten(0, 1)[selected[i].long()], sinks, SM_SCALE) for i in range(num_tokens) ] expected, _, value_scale, score_error = ( diff --git a/tests/model_executor/test_jit_warmup_triton_launcher.py b/tests/model_executor/test_jit_warmup_triton_launcher.py index a92303dcf288..c19804eae4b4 100644 --- a/tests/model_executor/test_jit_warmup_triton_launcher.py +++ b/tests/model_executor/test_jit_warmup_triton_launcher.py @@ -553,12 +553,10 @@ def test_compute_slot_mapping_uses_named_launcher_inputs(monkeypatch) -> None: owner = ComputeSlotMappingKernel() compile_key = owner.CompileKey( kv_cache_block_size=16, - blocks_per_kv_block=1, total_cp_world_size=2, total_cp_rank=1, cp_kv_cache_interleave_size=1, block_table_stride=128, - block_size=16, ) launches: list[tuple[Any, ...]] = [] @@ -574,7 +572,6 @@ def launch(launch_spec: Any, inputs: Any) -> None: assert inputs["block_table_stride"] == 128 assert kwargs == { "KV_CACHE_BLOCK_SIZE": 16, - "BLOCKS_PER_KV_BLOCK": 1, "TOTAL_CP_WORLD_SIZE": 2, "TOTAL_CP_RANK": 1, "CP_KV_CACHE_INTERLEAVE_SIZE": 1, diff --git a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py index 7cd7863b0659..0942af4f4570 100644 --- a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py +++ b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py @@ -425,7 +425,7 @@ def _skip_if_backend_unavailable(backend_cls, kv_cache_dtype: str, block_size: i and kv_cache_dtype != "fp8_ds_mla" ): pytest.skip("FlashMLA Sparse fp8 only supports fp8_ds_mla kv-cache dtype") - if block_size not in backend_cls.get_supported_kernel_block_sizes(): + if not backend_cls.supports_block_size(block_size): pytest.skip( f"{backend_cls.get_name()} does not support block_size={block_size}" ) diff --git a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py index 735748ae943e..66113aca0917 100644 --- a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py +++ b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py @@ -29,7 +29,7 @@ TOPK = 128 # triton convert requires width % 128 == 0 -def ref_convert(req_id, block_table, token_indices, BLOCK_SIZE=64, **_): +def ref_convert(req_id, block_table, token_indices, BLOCK_SIZE=64, **kw): out = torch.full_like(token_indices, -1) counts = torch.zeros(token_indices.shape[0], dtype=torch.int32) for t in range(token_indices.shape[0]): @@ -41,7 +41,7 @@ def ref_convert(req_id, block_table, token_indices, BLOCK_SIZE=64, **_): blk = int(block_table[int(req_id[t]), pos // BLOCK_SIZE]) if blk < 0: continue - vals.append(blk * BLOCK_SIZE + pos % BLOCK_SIZE) + vals.append(blk * kw["BLOCK_STRIDE_ROWS"] + pos % BLOCK_SIZE) out[t, : len(vals)] = torch.tensor(vals, dtype=out.dtype) counts[t] = len(vals) return out, counts @@ -133,11 +133,8 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype): meta.state = state q_nope = torch.randn(rows, impl.num_heads, HEAD) q_rope = torch.randn(rows, impl.num_heads, qk_rope) - cache = torch.zeros( - 8 * BLOCK_SIZE, - impl.head_size, - dtype=torch.uint8 if impl.use_fp8_kv_cache else torch.bfloat16, - ) + dtype = torch.uint8 if impl.use_fp8_kv_cache else torch.bfloat16 + cache = torch.zeros(8, 2, BLOCK_SIZE, impl.head_size, dtype=dtype)[:, 0] out, lse = impl.forward_mqa( (q_nope, q_rope), cache, meta, SimpleNamespace(_k_scale_float=0.5) @@ -147,7 +144,10 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype): # Reserved buffers carry this step's slots; lengths are NOT refreshed # here (the builder plans them host-side before capture/replay). ref_slots, ref_counts = ref_convert( - meta.req_id_per_token, meta.block_table, impl.topk_indices_buffer + meta.req_id_per_token, + meta.block_table, + impl.topk_indices_buffer, + BLOCK_STRIDE_ROWS=2 * BLOCK_SIZE, ) offset = 0 for t in range(rows): @@ -161,7 +161,7 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype): assert state.wrapper.run_args is not None q_pe, ckv, kpe, kwargs = state.wrapper.run_args[1:] assert q_pe.shape == (rows, impl.num_heads, qk_rope) - assert ckv.shape == (8 * BLOCK_SIZE, 1, HEAD) + assert ckv.shape == ((8 - 1) * 2 * BLOCK_SIZE + BLOCK_SIZE, 1, HEAD) assert kpe.shape[-1] == qk_rope if impl.use_fp8_kv_cache: assert kwargs["ckv_scale"] == 0.5 and kwargs["kpe_scale"] == 1.0 diff --git a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py index 652358876949..8532d2f43aee 100644 --- a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py +++ b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py @@ -24,6 +24,21 @@ from vllm.v1.attention.backends.mla.flashmla_sparse import ( QUANTIZED_DS_MLA_CACHE_FORMATS, FlashMLASparseBackend, + _fp8_kv_pages, +) +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + KVCacheLayout, + KVCacheTensor, + KVQuantMode, + MLAAttentionSpec, +) +from vllm.v1.worker.utils import ( + AttentionGroup, + allocate_kv_cache, + prepare_kernel_block_sizes, + select_common_block_size, ) SM90 = DeviceCapability(major=9, minor=0) @@ -83,6 +98,88 @@ def test_supported_head_sizes_include_512(): assert FlashMLASparseBackend.get_supported_head_sizes() == [576, 512] +def test_flashmla_bf16_nope_accepts_packed_manager_blocks(): + spec = MLAAttentionSpec( + block_size=1152, + num_kv_heads=1, + head_size=512, + dtype=torch.bfloat16, + cache_dtype_str="bfloat16", + block_stride_alignment=1024, + ) + num_blocks = 2 + layers = ["layer.0", "layer.1"] + page_size = spec.page_size_bytes + config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=num_blocks * len(layers) * page_size, + layers=layers, + layer_stride=page_size, + block_stride=len(layers) * page_size, + ) + ], + kv_cache_groups=[KVCacheGroupSpec(layers, spec)], + ) + groups = [[AttentionGroup(FlashMLASparseBackend, layers, spec, 0)]] + + kernel_block_sizes = prepare_kernel_block_sizes(config, groups) + caches = allocate_kv_cache( + config, + torch.device("cpu"), + KVCacheLayout.BLHNC, + kernel_block_sizes, + ) + + assert kernel_block_sizes == [1152] + assert caches["layer.0"].shape == (num_blocks, 1, 1152, 512) + + +@pytest.mark.parametrize("sm100", [False, True]) +def test_flashmla_quantized_cache_kernel_pages(monkeypatch, sm100): + """Quantized blocks are read as 64-token pages; SM100 only row-aligns them + (TMA) and keeps 64-token blocks.""" + monkeypatch.setattr( + current_platform, + "is_device_capability_family", + lambda family, device_id=0: sm100 and family == 100, + ) + spec = MLAAttentionSpec( + block_size=1152, + num_kv_heads=1, + head_size=576, + dtype=torch.float8_e4m3fn, + cache_dtype_str="fp8_ds_mla", + kv_quant_mode=KVQuantMode.FP8_PER_TENSOR, + state_content_bytes=656, + ) + + kernel_block_size = 64 if sm100 else 1152 + assert ( + select_common_block_size(1152, [FlashMLASparseBackend], [spec]) + == kernel_block_size + ) + + +def test_fp8_kv_pages_address_packed_blocks(): + """Token t of block b is row b * block_stride_rows + t of the 64-row pages; + 64-token blocks are passed through.""" + row, block_size, stride_rows = 656, 128, 192 + backing = torch.arange(3 * stride_rows).repeat_interleave(row) + cache = backing.view(3, stride_rows, row)[:, :block_size] + + pages, block_stride_rows = _fp8_kv_pages(cache, block_size) + assert block_stride_rows == stride_rows + assert pages.shape[1:] == (64, row) + slot = 2 * block_stride_rows + 100 + assert torch.equal(pages[slot // 64, slot % 64], cache[2, 100]) + + pages, block_stride_rows = _fp8_kv_pages(cache[:, :64], 64) + assert block_stride_rows is None + assert torch.equal(pages, cache[:, :64]) + + def test_quantized_ds_mla_formats_are_the_envelope_set(): assert frozenset({"fp8_ds_mla", "nvfp4_ds_mla"}) == QUANTIZED_DS_MLA_CACHE_FORMATS diff --git a/tests/v1/attention/test_glm5next_indexer_block_sizes.py b/tests/v1/attention/test_glm5next_indexer_block_sizes.py new file mode 100644 index 000000000000..9f27d2994636 --- /dev/null +++ b/tests/v1/attention/test_glm5next_indexer_block_sizes.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-5.3-Flash indexer kernel block sizes are DeepGEMM pool pages.""" + +from types import SimpleNamespace + +import pytest + +from vllm.platforms import current_platform +from vllm.v1.attention.backends.mla import indexer +from vllm.v1.attention.backends.mla.indexer import Glm5NextIndexerBackend + +KPOOL = 4 + + +@pytest.mark.parametrize( + "rocm,sm120,expected", + [ + (False, False, [KPOOL * 32, KPOOL * 64]), + # SM120 DeepGEMM takes only 64-entry pages for the fp8 cache. + (False, True, [KPOOL * 64]), + # ROCm's AITER paged-MQA logits take the same pool pages. + (True, False, [KPOOL * 32, KPOOL * 64]), + ], +) +def test_glm5next_indexer_kernel_block_sizes(monkeypatch, rocm, sm120, expected): + monkeypatch.setattr(current_platform, "is_rocm", lambda: rocm) + monkeypatch.setattr( + current_platform, + "is_device_capability_family", + lambda family: sm120 and family == 120, + ) + model_config = SimpleNamespace(hf_text_config=SimpleNamespace(index_kpool=KPOOL)) + config = SimpleNamespace(model_config=model_config) + monkeypatch.setattr(indexer, "get_current_vllm_config", lambda: config) + + spec = SimpleNamespace(tokens_per_state=KPOOL) + assert Glm5NextIndexerBackend.get_supported_kernel_block_sizes(spec) == expected + # Platform block-size selection queries without a spec. + assert Glm5NextIndexerBackend.get_supported_kernel_block_sizes() == expected diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 17dea5d03c19..464fce66f212 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -122,7 +122,6 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): builder = object.__new__(DeepseekV32IndexerMetadataBuilder) builder.dcp_world_size = 1 builder.compress_ratio = 16 - builder.kernel_block_size = None builder.num_sms = num_sms builder.arange_buffer = torch.arange(num_reqs + 1, dtype=torch.int32, device=device) builder.kv_cache_spec = SimpleNamespace(num_states=64) @@ -170,30 +169,24 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.parametrize( - "block_size,kernel_block_size,compress_ratio", [(64, 64, 1), (256, 64, 4)] -) -def test_indexer_draft_decode_refresh_follows_new_batch( - monkeypatch, block_size, kernel_block_size, compress_ratio -): +@pytest.mark.parametrize("compress_ratio", [1, 4]) +def test_indexer_draft_decode_refresh_follows_new_batch(monkeypatch, compress_ratio): """Without a prior build(), the refresh must not keep the last batch's state.""" monkeypatch.setattr(indexer_module, "has_deep_gemm", lambda: False) device = torch.device("cuda") num_reqs, width = 3, 8 - factor = block_size // kernel_block_size builder = object.__new__(DeepseekV32IndexerMetadataBuilder) builder.dcp_world_size = 1 builder.compress_ratio = compress_ratio - builder.kernel_block_size = kernel_block_size builder.arange_buffer = torch.arange(num_reqs + 1, dtype=torch.int32, device=device) - builder.kv_cache_spec = SimpleNamespace(block_size=block_size, num_states=64) + builder.kv_cache_spec = SimpleNamespace(block_size=64, num_states=64) def stale(*shape: int) -> torch.Tensor: return torch.full(shape, -7, dtype=torch.int32, device=device) decode = DeepSeekV32IndexerDecodeMetadata( - block_table=stale(num_reqs, width // factor), + block_table=stale(num_reqs, width), seq_lens=stale(num_reqs, 1), decode_lens=stale(num_reqs), requires_padding=False, @@ -219,8 +212,7 @@ def stale(*shape: int) -> torch.Tensor: builder.update_draft_decode_metadata(metadata) - expected_block_table = request_block_table[:, ::factor] // factor - torch.testing.assert_close(decode.block_table, expected_block_table) + torch.testing.assert_close(decode.block_table, request_block_table) torch.testing.assert_close(decode.indices, builder.arange_buffer[:num_reqs]) assert torch.all(decode.per_req_decode_lens == 1) assert torch.all(decode.decode_lens == 1) @@ -471,7 +463,6 @@ def test_zero_token_pcp_rank_participates_in_compressed_mapping_gather(monkeypat builder.pcp_rank = 0 builder.dcp_world_size = 1 builder.cp_kv_cache_interleave_size = 1 - builder.kernel_block_size = None builder.kv_cache_spec = SimpleNamespace(block_size=64, num_states=64) builder.compressed_slot_mapping_buffer = torch.zeros(8, dtype=torch.int64) diff --git a/tests/v1/attention/test_kpool_indexer_block_sizes.py b/tests/v1/attention/test_kpool_indexer_block_sizes.py index a737988dde66..86e196d539b4 100644 --- a/tests/v1/attention/test_kpool_indexer_block_sizes.py +++ b/tests/v1/attention/test_kpool_indexer_block_sizes.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""CPU tests for storage-block kernel-block selection in prepare_kernel_block_sizes.""" +"""CPU tests for the kpool indexer's kernel block selection.""" from types import SimpleNamespace @@ -11,20 +11,14 @@ from vllm.models.glm5next.common.attention import Glm5NextIndexerCache from vllm.v1.attention.backend import AttentionBackend, MultipleOf from vllm.v1.attention.backends.mla import indexer as indexer_mod -from vllm.v1.attention.backends.mla.indexer import ( - DeepseekV32IndexerBackend, - KpoolTailBackend, -) -from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( - ROCMAiterMLASparseBackend, -) +from vllm.v1.attention.backends.mla.indexer import Glm5NextIndexerBackend from vllm.v1.kv_cache_interface import MLAAttentionSpec -from vllm.v1.worker.utils import prepare_kernel_block_sizes +from vllm.v1.worker.utils import prepare_kernel_block_sizes, select_common_block_size pytestmark = pytest.mark.cpu_test # ``index_kpool`` of zai-org/GLM-5.3-Flash: one indexer state pools 4 -# compressed states (hence storage_block_size = page * 4). +# compressed states (hence a kernel block is page * 4 tokens). INDEX_KPOOL = 4 # The hybrid KDA/mamba manager floors at TP8..TP1; 640 is the geometry of the @@ -59,14 +53,17 @@ def _prepare( spec: MLAAttentionSpec, backends: list[type[AttentionBackend]] ) -> list[int]: kv_cache_config = SimpleNamespace( - kv_cache_groups=[SimpleNamespace(kv_cache_spec=spec)] + kv_cache_groups=[SimpleNamespace(kv_cache_spec=spec, layer_names=[])], + kv_cache_tensors=[], ) - attn_groups = [[SimpleNamespace(backend=backend) for backend in backends]] + attn_groups = [ + [SimpleNamespace(backend=backend, kv_cache_spec=spec) for backend in backends] + ] return prepare_kernel_block_sizes(kv_cache_config, attn_groups) @pytest.mark.parametrize("manager_block", TP_FLOOR_MANAGER_BLOCKS) -def test_prepare_uses_storage_block_for_the_kpool_group( +def test_kpool_group_maps_pool_pages( monkeypatch: pytest.MonkeyPatch, manager_block: int ): """The bug: on ROCm, kpool groups must get 128/256 - not the 640-4352 @@ -78,54 +75,43 @@ def test_prepare_uses_storage_block_for_the_kpool_group( """ _mock_rocm_platform(monkeypatch) spec = _kpool_storage_spec(manager_block) - expected_storage = 256 if manager_block % 256 == 0 else 128 - assert spec.storage_block_size == expected_storage + expected = 256 if manager_block % 256 == 0 else 128 - selected = _prepare( - spec, - [ROCMAiterMLASparseBackend, DeepseekV32IndexerBackend, KpoolTailBackend], - ) - assert selected == [expected_storage] + # The kpool group's AttentionGroup selects its kernel block this way. + selected = [ + select_common_block_size(manager_block, [Glm5NextIndexerBackend], [spec]) + ] + assert selected == [expected] assert selected != [manager_block] # The hybrid block-table split stays integral. assert manager_block % selected[0] == 0 -def test_prepare_falls_back_when_storage_block_is_not_supported(): - """No storage block, or a storage block the backends reject (e.g. the - CUDA indexer's exact ``[64]`` vs a 128-token storage block): selection is - `select_common_block_size`'s answer, unchanged from no-fix behavior. - """ +def test_prepare_uses_the_backend_vote(): + """Unpacked groups take `select_common_block_size`'s answer.""" class Fixed64Backend: @staticmethod - def get_supported_kernel_block_sizes(): + def get_supported_kernel_block_sizes(kv_cache_spec=None): return [64] @staticmethod def get_name() -> str: return "FIXED64" - # Storage block the group's backend does not accept -> backend vote (64). - spec = _kpool_storage_spec(640) - assert spec.storage_block_size == 128 - assert _prepare(spec, [Fixed64Backend]) == [64] - - # No storage block at all (any non-kpool MLA model) -> backend vote. plain_spec = MLAAttentionSpec( block_size=640, num_kv_heads=1, head_size=128, dtype=torch.bfloat16, ) - assert plain_spec.storage_block_size is None assert _prepare(plain_spec, [Fixed64Backend]) == [64] # ...and with backends that accept the manager block, the manager wins, # exactly as on main. class AcceptAllBackend: @staticmethod - def get_supported_kernel_block_sizes(): + def get_supported_kernel_block_sizes(kv_cache_spec=None): return [MultipleOf(1)] @staticmethod diff --git a/tests/v1/attention/test_kpool_tail_slot_mapping.py b/tests/v1/attention/test_kpool_tail_slot_mapping.py index 183c63ec2ffd..062fd783e0b3 100644 --- a/tests/v1/attention/test_kpool_tail_slot_mapping.py +++ b/tests/v1/attention/test_kpool_tail_slot_mapping.py @@ -3,7 +3,7 @@ """CPU tests for the kpool tail slot mapping (no GPU required). The kpool tail cache is a 1-block-per-request circular ring addressed by -``pos % kpool`` (``KpoolTailSpec`` / ``KpoolTailManager``: exactly one block +``pos % kpool`` (``CircularBufferSpec`` / ``CircularBufferManager``: exactly one block allocated per request, never grown, so only column 0 of its block table is ever written; the rest stays zero-initialized). @@ -32,27 +32,25 @@ KpoolTailMetadataBuilder, compute_kpool_tail_slot_mapping, ) -from vllm.v1.kv_cache_interface import KpoolTailSpec, compute_layout_strides +from vllm.v1.kv_cache_interface import CircularBufferSpec, compute_layout_strides from vllm.v1.kv_cache_layout import KVCacheLayout -from vllm.v1.worker.block_table import get_block_table_width KPOOL = 4 def test_tail_backend_layout_matches_kernel_pointer_arithmetic(): (layout,) = KpoolTailBackend.supported_kv_cache_layouts() - spec = KpoolTailSpec( + spec = CircularBufferSpec( block_size=KPOOL, num_kv_heads=2, head_size=128, head_size_v=0, dtype=torch.bfloat16, - sliding_window=KPOOL, ) strides = compute_layout_strides(spec, num_blocks=8, num_layers=3, layout=layout) _, _, head_stride, state_stride, content_stride = strides - assert layout is KVCacheLayout.LBHNC + assert layout is KVCacheLayout.BLHNC assert head_stride == KPOOL * 128 * torch.bfloat16.itemsize assert state_stride == 128 * torch.bfloat16.itemsize assert content_stride == 1 @@ -77,37 +75,14 @@ def test_tail_ring_divides_the_attention_block(num_speculative_tokens, ring): SimpleNamespace(num_speculative_tokens=num_speculative_tokens) ) - assert spec.block_size == spec.sliding_window == ring + assert isinstance(spec, CircularBufferSpec) and spec.block_size == ring assert ring >= KPOOL + num_speculative_tokens assert 640 % ring == 0 -def test_tail_spec_opts_out_of_generic_slot_mapping(): - """The tail row is one block wide (padded to the block-table alignment), so - the generic kernel's ``pos // kpool`` column index runs off the end of the - allocation for long prompts. The spec must opt out of it entirely.""" - spec = KpoolTailSpec( - block_size=KPOOL, - num_kv_heads=2, - head_size=128, - head_size_v=0, - dtype=torch.bfloat16, - sliding_window=KPOOL, - ) - max_len = 1 << 20 - width = get_block_table_width( - spec.max_num_blocks_per_req(None, max_len), - spec.block_size, - token_alignment=spec.block_table_token_alignment, - ) - - assert width * KPOOL < max_len - assert spec.uses_slot_mapping is False - - def make_tail_block_table(own_blocks, width=64): """Tail-group block table as BlockTables produces it: column 0 holds the - request's single KpoolTailManager block, the remaining columns are never + request's single CircularBufferManager block, the remaining columns are never written and stay zero.""" bt = torch.zeros(len(own_blocks), width, dtype=torch.int32) bt[:, 0] = torch.tensor(own_blocks, dtype=torch.int32) @@ -447,6 +422,7 @@ def test_triton_mapping_matches_cpu(per_req, own_blocks, num_actual, padded_len) tokens between the last request boundary and num_actual_tokens (mapped to the last request), untouched padding beyond num_actual, and PAD for requests on the null block.""" + ring_size = 2 * KPOOL positions, qsl, slot_mapping, _, num_reqs = make_batch( per_req, padded_len=padded_len ) @@ -455,14 +431,17 @@ def test_triton_mapping_matches_cpu(per_req, own_blocks, num_actual, padded_len) slot_mapping = torch.arange(padded_len, dtype=torch.int64) + 1000 bt = make_tail_block_table(own_blocks) - ref = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs) - got = circular_tail_slots( + ref = compute_kpool_tail_slot_mapping( + slot_mapping, bt, qsl, positions, num_actual, num_reqs, ring_size + ) + got = compute_kpool_tail_slot_mapping( slot_mapping.cuda(), bt.cuda(), qsl.cuda().to(torch.int32), positions.cuda(), num_actual, num_reqs, + ring_size, ) torch.testing.assert_close(got.cpu(), ref) for req, blk in enumerate(own_blocks): diff --git a/tests/v1/attention/test_mm_prefix.py b/tests/v1/attention/test_mm_prefix.py index f878ede4865f..b8085f0bc584 100644 --- a/tests/v1/attention/test_mm_prefix.py +++ b/tests/v1/attention/test_mm_prefix.py @@ -805,7 +805,6 @@ def variant_uses_mm_prefix(variant): ) layer.mm_prefix_clamp_sliding_window = True builder = backend.get_builder_cls()(spec, ["composite_test"], cfg, DEVICE) - builder.set_kernel_block_size(block_size) # Match production, where unwritten page tails are zero-initialized. cache = torch.zeros( num_blocks, diff --git a/tests/v1/attention/test_rocm_glm5next_sparse.py b/tests/v1/attention/test_rocm_glm5next_sparse.py index e69686f5005d..6077b9eb2fb6 100644 --- a/tests/v1/attention/test_rocm_glm5next_sparse.py +++ b/tests/v1/attention/test_rocm_glm5next_sparse.py @@ -11,6 +11,7 @@ from vllm.utils.torch_utils import set_random_seed from vllm.v1.attention.backends.mla import rocm_aiter_mla_sparse as sparse_mod from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( + ROCMAiterMLASparseBackend, _use_rocm_sparse_triton, fit_kpool_indices_to_aiter, ) @@ -19,6 +20,7 @@ _validate_dsv4_sparse_dims, _validate_sparse_dims, ) +from vllm.v1.kv_cache_interface import KVCacheLayout @triton.jit @@ -104,6 +106,10 @@ def test_fit_kpool_indices_rejects_narrow_input(): ) +def test_rocm_sparse_mla_supports_glm_packed_layout(): + assert KVCacheLayout.BLHNC in ROCMAiterMLASparseBackend.supported_kv_cache_layouts() + + @pytest.mark.parametrize( ( "kv_cache_dtype", diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 973a9e03e88b..160b766da077 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -140,7 +140,7 @@ def test_nope_flashinfer_sparse_mla_uses_model_scale(monkeypatch): metadata = SimpleNamespace( req_id_per_token=torch.zeros(1, dtype=torch.int32), block_table=torch.zeros((1, 1), dtype=torch.int32), - block_size=1, + block_size=32, ) recorded_scale = None @@ -178,7 +178,7 @@ def fake_flashinfer(**kwargs): ) impl.forward_mqa( torch.zeros(1, 1, kv_lora_rank), - torch.zeros(1, 1, kv_lora_rank), + torch.zeros(1, 32, kv_lora_rank), metadata, SimpleNamespace(), ) @@ -437,6 +437,7 @@ def test_sparse_backend_decode_correctness( qk_rope_head_dim: int, v_head_dim: int, monkeypatch, + block_stride_rows: int | None = None, ): if qk_rope_head_dim == 0 and ( backend_cls != FlashMLASparseBackend @@ -478,8 +479,7 @@ def test_sparse_backend_decode_correctness( if device_capability is None or device_capability.major != 10: pytest.skip("The NVFP4 DS-MLA kv-cache dtype requires SM 10.x") - supported_block_sizes = backend_cls.get_supported_kernel_block_sizes() - if block_size not in supported_block_sizes: + if not backend_cls.supports_block_size(block_size): pytest.skip( f"{backend_cls.get_name()} does not support block_size={block_size}" ) @@ -719,6 +719,11 @@ def test_sparse_backend_decode_correctness( kv_cache_dtype=kv_cache_dtype, scale=kv_cache_scale, ) + if block_stride_rows is not None: + pad = (0, 0, 0, block_stride_rows - block_size) + kv_cache = torch.nn.functional.pad(kv_cache, pad)[:, :, :block_size] + slot = int(common_attn_metadata.slot_mapping[-1]) + query_row = kv_cache[slot // block_size, 0, slot % block_size].zero_() # The sparse builder clones the layer's dense-MHA prefill backend from # static_forward_context; register a mock layer carrying one. @@ -814,6 +819,12 @@ def test_sparse_backend_decode_correctness( out_buffer, ) + if block_stride_rows is not None and kv_cache_dtype == "auto": + expected_row = torch.cat((kv_c_vllm[-1], k_pe_vllm[-1, 0])) + torch.testing.assert_close(query_row, expected_row, rtol=0, atol=0) + elif block_stride_rows is not None: + assert query_row.any() + assert backend_output.shape == sdpa_reference.shape assert backend_output.dtype == sdpa_reference.dtype assert torch.isfinite(backend_output).all() @@ -828,6 +839,56 @@ def test_sparse_backend_decode_correctness( torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.01, atol=0.01) +@pytest.mark.parametrize( + ("backend_cls", "kv_cache_dtype", "block_size", "block_stride_rows", "head_dims"), + [ + pytest.param( + FlashInferMLASparseTRTLLMBackend, + "auto", + 256, + 320, + (128, 128, 64, 128), + id="FlashInferTRTLLM", + ), + # GLM-5.3 on SM90: the fp8 kernel's 64-token pages in larger blocks. + pytest.param( + FlashMLASparseBackend, + "fp8_ds_mla", + 128, + 192, + (64, 256, 0, 256), + id="FlashMLA-fp8_ds_mla", + ), + ], +) +def test_sparse_mla_packed_stride( + default_vllm_config, + dist_init, + workspace_init, + monkeypatch, + backend_cls, + kv_cache_dtype, + block_size, + block_stride_rows, + head_dims, +): + test_sparse_backend_decode_correctness( + default_vllm_config, + dist_init, + backend_cls, + "mixed_small", + kv_cache_dtype, + 1, + block_size, + workspace_init, + 1.0, + 1.0, + *head_dims, + monkeypatch, + block_stride_rows=block_stride_rows, + ) + + def _triton_convert_reference_impl( req_ids: torch.Tensor, block_table: torch.Tensor, diff --git a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py index cd3c9b218147..d869141df81a 100644 --- a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py +++ b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py @@ -12,6 +12,7 @@ from vllm.model_executor.layers.attention.mla_attention import MLAAttention from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache from vllm.platforms import current_platform +from vllm.v1.attention.backend import MultipleOf from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( FlashInferMLASparseTRTLLMBackend, ) @@ -25,7 +26,7 @@ from vllm.v1.core.kv_cache_utils import get_kv_cache_config_from_groups from vllm.v1.kv_cache_interface import KVCacheGroupSpec, UniformTypeKVCacheSpecs from vllm.v1.kv_cache_layout import KVCacheLayout -from vllm.v1.worker.utils import allocate_kv_cache +from vllm.v1.worker.utils import allocate_kv_cache, customize_attention_spec pytestmark = pytest.mark.skip_global_cleanup @@ -48,10 +49,10 @@ @pytest.mark.parametrize( "backend,cache_dtype,sm100,stride_alignment", [ - (FlashInferMLASparseTRTLLMBackend, "auto", True, 1152), - (FlashInferMLASparseTRTLLMBackend, "fp8", True, 576), + (FlashInferMLASparseTRTLLMBackend, "auto", True, 32 * 1152), + (FlashInferMLASparseTRTLLMBackend, "fp8", True, 32 * 576), (FlashMLASparseBackend, "auto", True, 1152), - (FlashMLASparseBackend, "fp8_ds_mla", False, None), + (FlashMLASparseBackend, "fp8_ds_mla", False, MultipleOf(64)), (FlashMLASparseBackend, "fp8_ds_mla", True, 656), (FlashMLASparseBackend, "nvfp4_ds_mla", True, 352), ], @@ -79,7 +80,9 @@ def test_allocation_and_warmup_follow_addressing_mode( non_causal_multi_token_decode=False, ) layer._uses_flat_kv_cache = MethodType(MLAAttention._uses_flat_kv_cache, layer) - spec = MLAAttention.get_kv_cache_spec(layer, config) + spec = customize_attention_spec( + backend, MLAAttention.get_kv_cache_spec(layer, config) + ) indexer = SimpleNamespace( cache_config=config.cache_config, head_dim=132, dtype=torch.uint8 ) @@ -102,10 +105,8 @@ def test_allocation_and_warmup_follow_addressing_mode( MLAAttention.bind_kv_cache(layer, views["mla"]) DeepseekV32IndexerCache.bind_kv_cache(indexer, views["indexer"]) assert indexer.kv_cache.stride(0) == views["indexer"].stride(0) - if stride_alignment is not None: - assert ( - views["mla"].stride(0) * views["mla"].element_size() % stride_alignment == 0 - ) + alignment = spec.get_block_stride_alignment() + assert views["mla"].stride(0) * views["mla"].element_size() % alignment == 0 if not layer._uses_flat_kv_cache(): register.assert_not_called() else: @@ -113,7 +114,7 @@ def test_allocation_and_warmup_follow_addressing_mode( register.assert_called_once_with(config, block_stride_rows=stride) layer.kv_cache[1, 0].fill_(3) torch.testing.assert_close(rows[stride], layer.kv_cache[1, 0]) - if stride_alignment is None: + if alignment == 1: assert cache.kv_cache_tensors[0].size == cache.num_blocks * sum( item.page_size_bytes for item in specs.values() ) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 51317deb8d01..4bffe91a6534 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -26,8 +26,8 @@ ) from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, FullAttentionSpec, - KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, MambaSpec, @@ -2308,12 +2308,11 @@ def test_opted_out_scratch_group_keeps_partial_hash_hits(): ), KVCacheGroupSpec( ["tail"], - KpoolTailSpec( + CircularBufferSpec( block_size=mamba_block_size, num_kv_heads=1, head_size=1, dtype=torch.float32, - sliding_window=mamba_block_size, ), ), ], @@ -2366,12 +2365,11 @@ def test_kpool_tail_supports_128_token_partial_hash_hits(): ), KVCacheGroupSpec( ["tail"], - KpoolTailSpec( + CircularBufferSpec( block_size=kpool_block_size, num_kv_heads=1, head_size=1, dtype=torch.float32, - sliding_window=kpool_block_size, ), ), ], diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 638deeedc579..6a705b6e5e8e 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -88,7 +88,7 @@ def test_packed_alignment_preserves_hot_pages_and_generic_stride(): hot = HiSparseHotSpec(block_size=16, page_size=5120, blocks_per_request=2) groups = [KVCacheGroupSpec(["attention"], spec), KVCacheGroupSpec(["hot"], hot)] alignment = lcm(1152, hot.page_size_bytes) - stride = _pool_bytes_per_block(groups) + stride = _pool_bytes_per_block(_mock_vllm_config("BLHNC"), groups) assert stride == alignment layout = KVCacheLayout.BLHNC tensors = _build_hisparse_kv_cache_tensors(groups, 3, 3 * stride, layout, stride) @@ -254,7 +254,7 @@ def test_roles_land_in_separate_groups(self): def test_every_group_fits_one_packed_block(self): config = _shared_layout_config() groups = get_kv_cache_groups(config, _make_csa_linear_specs()) - bytes_per_block = _get_kv_cache_bytes_per_block(groups) + bytes_per_block = _get_kv_cache_bytes_per_block(groups, KVCacheLayout.BLHNC) pages = _pages(groups) for group in groups: @@ -418,7 +418,7 @@ def test_mamba_split_measures_the_block_the_other_groups_already_force(self, wid groups = _get_packed_kv_cache_groups(config, specs) gdn = [g for g in groups if g.layer_names[0].startswith("gdn.")] - assert _get_kv_cache_bytes_per_block(groups) == sum( + assert _get_kv_cache_bytes_per_block(groups, KVCacheLayout.BLHNC) == sum( specs[name].page_size_bytes for name in specs if name.startswith("wide.") ) # That block holds every GDN state at once, so the repeat pattern alone @@ -465,7 +465,7 @@ def test_sliding_window_bucket_is_capped_at_the_main_page(self): pages = _pages(groups) main = next(g for g in groups if "layers.2.attn" in g.layer_names) main_bytes = sum(pages[n] for n in main.layer_names) - assert _get_kv_cache_bytes_per_block(groups) == main_bytes + assert _get_kv_cache_bytes_per_block(groups, KVCacheLayout.BLHNC) == main_bytes swa_groups = [g for g in groups if g.layer_names[0].endswith(".swa")] per_group = main_bytes // pages["layers.0.swa"] assert len(swa_groups) == -(-43 // per_group) @@ -479,9 +479,9 @@ def test_sliding_window_bucket_is_capped_at_the_main_page(self): class TestDensePacking: def test_bytes_per_block_is_largest_group(self): groups, g1, g2 = _mixed_page_groups() - assert _get_kv_cache_bytes_per_block(groups) == _expected_bytes_per_block( - groups - ) + assert _get_kv_cache_bytes_per_block( + groups, KVCacheLayout.BLHNC + ) == _expected_bytes_per_block(groups) config = get_kv_cache_config_from_groups( _mock_vllm_config("BLHNC"), groups, MEMORY @@ -554,6 +554,20 @@ def test_single_group_mixed_pages_follows_layout(self, layout): assert idx_tensor.offset == 2 * _mla(512).page_size_bytes assert mla_tensor.block_stride == block_stride + def test_layer_outer_pads_only_unaligned_pages(self): + """Layer-outer block strides are pages, so only an unaligned page pads.""" + mla_page = _mla(512).page_size_bytes + idx = replace(_mla(128), block_stride_alignment=12288) + groups = [_uniform_group({"mla.0": _mla(512), "idx.0": idx})] + config = get_kv_cache_config_from_groups( + _mock_vllm_config("LBNHC"), groups, MEMORY + ) + mla_tensor, idx_tensor = config.kv_cache_tensors + assert config.num_blocks == MEMORY // (mla_page + 12288) + assert (mla_tensor.block_stride, idx_tensor.block_stride) == (mla_page, 12288) + assert idx_tensor.offset == mla_page * config.num_blocks + _bind(config, "LBNHC") + def test_overlaid_groups_alias_and_stay_isolated(self): groups, g1, g2 = _mixed_page_groups() # Overlay models resolve to a block-outer layout at backend selection (the @@ -562,7 +576,9 @@ def test_overlaid_groups_alias_and_stay_isolated(self): _mock_vllm_config("BLNHC"), groups, MEMORY ) assert config.num_blocks == MEMORY // _expected_bytes_per_block(groups) - assert _pool_bytes_per_block(groups) == _expected_bytes_per_block(groups) + assert _pool_bytes_per_block( + _mock_vllm_config("BLNHC"), groups + ) == _expected_bytes_per_block(groups) views = _bind(config, "BLNHC") assert set(views) == set(g1) | set(g2) @@ -667,7 +683,10 @@ def test_ring_beside_paged_groups_keeps_the_kv_block_as_scheduler_block(self): assert len(ring_groups) == 1 assert not ring_groups[0].kv_cache_spec.prefix_cacheable kv_cache_config = get_kv_cache_config_from_groups( - config, groups, available_memory=64 * _get_kv_cache_bytes_per_block(groups) + config, + groups, + available_memory=64 + * _get_kv_cache_bytes_per_block(groups, KVCacheLayout.BLHNC), ) assert resolve_kv_cache_block_sizes(kv_cache_config, config) == (128, 64) manager = KVCacheManager( @@ -716,7 +735,10 @@ def test_non_cacheable_swa_leaves_hits_to_the_paged_group(self): for g in groups ) == [(False, 128)] * 4 + [(True, 0)] kv_cache_config = get_kv_cache_config_from_groups( - config, groups, available_memory=64 * _get_kv_cache_bytes_per_block(groups) + config, + groups, + available_memory=64 + * _get_kv_cache_bytes_per_block(groups, KVCacheLayout.BLHNC), ) manager = KVCacheManager( generate_scheduler_kv_cache_config([kv_cache_config]), diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 53833ba83711..b06479f1caf1 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -10,7 +10,7 @@ from dataclasses import replace from pathlib import Path from types import SimpleNamespace -from typing import Any, cast +from typing import Any import pytest import torch @@ -19,7 +19,6 @@ import vllm.v1.hisparse.runtime as hisparse_runtime_module from vllm.config import ( CacheConfig, - KVTransferConfig, ModelConfig, SchedulerConfig, VllmConfig, @@ -69,7 +68,6 @@ HiddenStateCacheSpec, HiSparseHotSpec, HiSparseResidentSpec, - KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, @@ -1942,13 +1940,12 @@ def test_dcp_replicated_kpool_tail_keeps_block_interior_layout(monkeypatch): config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) config.parallel_config.decode_context_parallel_size = 4 config.cache_config.kv_cache_layout = None - tail = KpoolTailSpec( + tail = CircularBufferSpec( block_size=4, num_kv_heads=2, head_size=128, head_size_v=0, dtype=torch.bfloat16, - sliding_window=4, dcp_sharded=False, ) layout = resolve_kv_cache_layout(config, [["LBHNC"]], [new_mla_spec(), tail]) @@ -2906,564 +2903,77 @@ def test_generate_scheduler_kv_cache_config(): ) -def _glm5_like_kv_cache_spec( - mamba_spec_factory=new_mamba_spec, -) -> tuple[dict[str, KVCacheSpec], list[str]]: - """(mamba, mamba, mamba, MLA + indexer) * 11 plus a trailing mamba.""" +def test_packed_groups_glm5_like_hybrid(): + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=8192)) + vllm_config.cache_config.kv_cache_layout = "BLHNC" + kpool, block_size = 4, 1024 kv_cache_spec: dict[str, KVCacheSpec] = {} - mamba_layers = [] - for i in range(45): - if i % 4 == 3: - kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( - block_size=1024, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - ) - kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec( - block_size=1024, - num_kv_heads=1, - head_size=132, - dtype=torch.uint8, - tokens_per_state=16, - ) - else: - name = f"layers.{i}.linear_attn" - kv_cache_spec[name] = mamba_spec_factory() - mamba_layers.append(name) - return kv_cache_spec, mamba_layers - - -def _glm5_like_kv_cache_spec_with_tail( - mamba_spec_factory=new_mamba_spec, -) -> dict[str, KVCacheSpec]: - """Production kpool=4 proportions: (mamba*3, MLA + indexer + tail) * 11. - - The tail's logical page (2 * kpool * 2*indexer_head_dim * bf16) must fit - inside the indexer page it parasitizes (block_size//kpool * 132 B), which - is why the fixture drops the tokens_per_state=16 shape used above. - """ - kv_cache_spec, _ = _glm5_like_kv_cache_spec(mamba_spec_factory) - for i in range(3, 45, 4): - kpool = 4 - kv_cache_spec[f"layers.{i}.indexer"] = replace( - cast(MLAAttentionSpec, kv_cache_spec[f"layers.{i}.indexer"]), + for i in range(11): + kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + block_stride_alignment=32 * 576 * 2, + ) + kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, tokens_per_state=kpool, + block_stride_alignment=64 * 132, ) - kv_cache_spec[f"layers.{i}.tail"] = KpoolTailSpec( + kv_cache_spec[f"layers.{i}.tail"] = CircularBufferSpec( block_size=kpool, num_kv_heads=2, head_size=128, head_size_v=0, dtype=torch.bfloat16, - sliding_window=kpool, ) - return kv_cache_spec - - -def _tensor_by_layer(kv_cache_config: KVCacheConfig) -> dict[str, KVCacheTensor]: - return { - layer_name: tensor - for tensor in kv_cache_config.kv_cache_tensors - for layer_name in tensor.layers - } - - -def _layer_offset(tensor: KVCacheTensor, layer_name: str) -> int: - return tensor.offset + tensor.layers.index(layer_name) * tensor.layer_stride - - -def test_get_kv_cache_config_balanced_mamba_hybrid(): - """Hybrid slot sharing: mamba layers co-own the MLA slot tensors.""" - model_config = ModelConfig(max_model_len=8192) - vllm_config = VllmConfig(model_config=model_config) - - kv_cache_spec, mamba_layers = _glm5_like_kv_cache_spec() - mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes - idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes - - groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - uniform_groups = [ - group - for group in groups - if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) - ] - mamba_groups = [ - group for group in groups if isinstance(group.kv_cache_spec, MambaSpec) - ] - - # Round-robin into G = ceil(34 / 11) = 4 groups so every mamba layer - # gets an MLA slot. - assert len(groups) == 5 - assert len(uniform_groups) == 1 - assert len(uniform_groups[0].layer_names) == 22 - assert uniform_groups[0].kv_cache_spec.get_max_layers_per_page_size() == 11 - assert [len(group.layer_names) for group in mamba_groups] == [9, 9, 8, 8] - for k, name in enumerate(mamba_layers): - assert name in mamba_groups[k % 4].layer_names - # Mamba pages are padded up to the MLA page. - for group in mamba_groups: - assert group.kv_cache_spec.page_size_padded == mla_page - assert group.kv_cache_spec.page_size_bytes == mla_page - - bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) - assert bytes_per_block == 11 * mla_page + 11 * idx_page - - # Every block id is charged the full per-block sum. - attn_blocks = uniform_groups[0].kv_cache_spec.max_memory_usage_pages(vllm_config) - mamba_blocks_per_group = 1 + new_mamba_spec().num_speculative_blocks - blocks_per_request = attn_blocks + 4 * mamba_blocks_per_group - assert ( - kv_cache_utils._max_memory_usage_bytes_from_groups(vllm_config, groups) - == blocks_per_request * bytes_per_block - ) - - available_memory = bytes_per_block * 100 + 1 - kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, groups, available_memory - ) - assert kv_cache_config.num_blocks == 100 - - # Every logical layer has a view into one backing allocation. Mamba views - # alias their corresponding MLA slot by using the same byte offset. - assert len(kv_cache_config.kv_cache_tensors) == 56 - assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100} - tensors = _tensor_by_layer(kv_cache_config) - mla_layer_names = [f"layers.{i}.attn" for i in range(3, 45, 4)] - for i, mla_name in enumerate(mla_layer_names): - mla_tensor = tensors[mla_name] - assert mla_tensor.block_stride == mla_page - for group in mamba_groups: - if i < len(group.layer_names): - assert tensors[group.layer_names[i]].offset == mla_tensor.offset - for i in range(11): - assert tensors[f"layers.{4 * i + 3}.indexer"].block_stride == idx_page - - total_allocated = next(iter({t.size for t in kv_cache_config.kv_cache_tensors})) - assert total_allocated == bytes_per_block * 100 - assert 0 <= available_memory - total_allocated < bytes_per_block - - assert get_max_concurrency_for_kv_cache_config( - vllm_config, kv_cache_config - ) == pytest.approx(100 / blocks_per_request) - - -def test_get_kv_cache_config_kpool_tail_coowns_indexer_tensor(): - """The kpool tail parasitizes the indexer tensors instead of getting its - own: sibling idx/tail tensors paired by layer order, zero standalone tail - bytes, one shared block per request, and no prefix-caching leakage from - the kpool-sized scratch group.""" - model_config = ModelConfig(max_model_len=8192) - vllm_config = VllmConfig(model_config=model_config) - - kv_cache_spec = _glm5_like_kv_cache_spec_with_tail() - mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes - idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes - tail_logical_page = kv_cache_spec["layers.3.tail"].page_size_bytes - assert tail_logical_page < idx_page - - groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - tail_group = next( - group - for group in groups - if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) - and all( - isinstance(spec, KpoolTailSpec) - for spec in group.kv_cache_spec.kv_cache_specs.values() + for i in range(34): + kv_cache_spec[f"layers.{i}.kda"] = MambaSpec( + shapes=((3, 6144), (16, 128, 128)), + dtypes=(torch.bfloat16, torch.float32), + block_size=block_size, ) - ) - # The tail never prefix-caches; its page is padded up to the indexer page - # so the runner's strided view rides the indexer storage. - assert not tail_group.kv_cache_spec.prefix_cacheable - tail_inner = cast( - KpoolTailSpec, tail_group.kv_cache_spec.kv_cache_specs["layers.3.tail"] - ) - assert tail_inner.page_size_padded == idx_page - - bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) - assert bytes_per_block == 11 * mla_page + 11 * idx_page - - kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, groups, bytes_per_block * 100 + 1 - ) - assert kv_cache_config.num_blocks == 100 - # Tail layers get logical views but no additional storage: each view aliases - # its sibling indexer at the same offset in the shared allocation. - assert len(kv_cache_config.kv_cache_tensors) == 67 - tensors = _tensor_by_layer(kv_cache_config) - for i in range(11): - idx_name = f"layers.{4 * i + 3}.indexer" - tail_name = f"layers.{4 * i + 3}.tail" - assert tensors[idx_name].offset == tensors[tail_name].offset - assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100} - - # The layout detector surfaces the sibling names in layer order for the - # accounting and connector paths. - layout = kv_cache_utils._glm5_next_tensor_layout(kv_cache_config.kv_cache_groups) - assert layout is not None - assert layout[3] == [f"layers.{4 * i + 3}.indexer" for i in range(11)] - assert layout[6] == [f"layers.{4 * i + 3}.tail" for i in range(11)] - - # Accounting charges exactly one extra shared block per request for the - # tail, not a per-sequence allocation. - attn_group = next( - group - for group in groups - if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) - and group is not tail_group - ) - attn_blocks = attn_group.kv_cache_spec.max_memory_usage_pages(vllm_config) - mamba_blocks_per_group = 1 + new_mamba_spec().num_speculative_blocks - blocks_per_request = attn_blocks + 4 * mamba_blocks_per_group + 1 - assert ( - kv_cache_utils._max_memory_usage_bytes_from_groups(vllm_config, groups) - == blocks_per_request * bytes_per_block - ) - - -def test_glm5_kpool_tail_does_not_drag_hash_block_size(): - """The tail's kpool-sized scratch block (4 tokens) must not constrain the - prefix-cache hash granularity: participating groups alone decide it.""" - model_config = ModelConfig(max_model_len=8192) - vllm_config = VllmConfig(model_config=model_config) - - def align_mamba(): - return new_mamba_spec(mamba_cache_mode="align") - - groups = kv_cache_utils.get_kv_cache_groups( - vllm_config, _glm5_like_kv_cache_spec_with_tail(align_mamba) - ) - kv_cache_config = KVCacheConfig( - num_blocks=1, - kv_cache_tensors=[], - kv_cache_groups=groups, - ) - hash_vllm_config = SimpleNamespace( - cache_config=SimpleNamespace( - block_size=16, - enable_prefix_caching=True, - prefix_match_unit=None, - ), - parallel_config=SimpleNamespace(decode_context_parallel_size=1), - kv_transfer_config=object(), - ) - # gcd(attn 1024, mamba 16) with the tail's 4 excluded; scheduler size is - # the lcm including the tail (1024 % 4 == 0, so it coincides). - assert kv_cache_utils.resolve_kv_cache_block_sizes( - kv_cache_config, hash_vllm_config - ) == (1024, 16) - - -def test_get_kv_cache_config_mamba_hybrid_sharing_infeasible(): - """Reject GLM-5.3-Flash layouts whose Mamba page exceeds the MLA page.""" - model_config = ModelConfig(max_model_len=8192) - vllm_config = VllmConfig(model_config=model_config) - - # (512, 1024) fp32 state = 2 MiB per page > 1.125 MiB MLA page. - def big_mamba_spec(): - return new_mamba_spec(shapes=((512, 1024),), dtypes=(torch.float32,)) - - kv_cache_spec, _ = _glm5_like_kv_cache_spec(big_mamba_spec) - mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes - assert big_mamba_spec().page_size_bytes > mla_page - - with pytest.raises(ValueError, match="does not fit the MLA page"): - kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - - -def test_get_kv_cache_config_mamba_hybrid_sharing_infeasible_no_indexer(): - """Use the generic layout error when no kpool indexer is present.""" - model_config = ModelConfig(max_model_len=8192) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBHNC" - - kv_cache_spec: dict[str, KVCacheSpec] = {} - for i in range(27): - if i % 4 == 3 or i == 26: - kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( - block_size=1024, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - ) - else: - kv_cache_spec[f"layers.{i}.linear_attn"] = new_mamba_spec( - shapes=((512, 1024),), dtypes=(torch.float32,) - ) - mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes - assert kv_cache_spec["layers.0.linear_attn"].page_size_bytes > mla_page - - with pytest.raises(NotImplementedError, match="page size"): - kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - - -def test_get_kv_cache_config_mamba_hybrid_sharing_prepadded_mamba(): - """Platform-prepadded mamba pages (mamba_page_size_padded hint) must not - disable slot sharing; the layout re-pads them to the MLA page.""" - model_config = ModelConfig(max_model_len=8192) - vllm_config = VllmConfig(model_config=model_config) - - def prepadded_mamba_spec(): - return new_mamba_spec(page_size_padded=294_912) - - kv_cache_spec, _ = _glm5_like_kv_cache_spec(prepadded_mamba_spec) - mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes - assert prepadded_mamba_spec().page_size_bytes < mla_page groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - mamba_groups = [ - group for group in groups if isinstance(group.kv_cache_spec, MambaSpec) + attn = [ + g + for g in groups + if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) + and isinstance(g.kv_cache_spec.first_spec, MLAAttentionSpec) ] - assert len(groups) == 5 - assert [len(group.layer_names) for group in mamba_groups] == [9, 9, 8, 8] - for group in mamba_groups: - assert group.kv_cache_spec.page_size_bytes == mla_page - - bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) - kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, groups, bytes_per_block * 100 + 1 - ) - assert kv_cache_config.num_blocks == 100 - assert len(kv_cache_config.kv_cache_tensors) == 56 - - -def test_get_kv_cache_config_mamba_hybrid_sharing_pp_balanced_projection(): - """Round-robin mamba grouping keeps practical PP splits balanced: every - stage's largest projected mamba group slice fits its projected MLA - layers, so per-stage slot tensors all have an MLA owner.""" - model_config = ModelConfig(max_model_len=8192) - vllm_config = VllmConfig(model_config=model_config) - - kv_cache_spec, _ = _glm5_like_kv_cache_spec() - global_groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes - idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes - - # PP=2 split at transformer layer 22/23. - # Stage 0 of a PP=2 split at transformer layer 22/23: 5 MLA(+indexer) - # layers, projected mamba groups [5, 5, 4, 4]. - worker_spec = {n: s for n, s in kv_cache_spec.items() if int(n.split(".")[1]) <= 22} - groups = kv_cache_utils._project_kv_cache_groups_to_worker( - global_groups, worker_spec - ) - mamba_groups = [ - group for group in groups if isinstance(group.kv_cache_spec, MambaSpec) + rings = [ + g + for g in groups + if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) + and isinstance(g.kv_cache_spec.first_spec, CircularBufferSpec) ] - assert [len(group.layer_names) for group in mamba_groups] == [5, 5, 4, 4] - - bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) - assert bytes_per_block == 5 * mla_page + 5 * idx_page - - kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, groups, bytes_per_block * 100 + 1 - ) - assert kv_cache_config.num_blocks == 100 - tensors = _tensor_by_layer(kv_cache_config) - # Every projected Mamba slot aliases an MLA owner. - assert tensors["layers.3.attn"].offset == tensors["layers.0.linear_attn"].offset - assert tensors["layers.3.attn"].offset == tensors["layers.1.linear_attn"].offset - assert tensors["layers.3.attn"].offset == tensors["layers.2.linear_attn"].offset - assert tensors["layers.3.attn"].offset == tensors["layers.4.linear_attn"].offset - # The last slot only hosts the two groups with a 5th projected layer. - assert tensors["layers.19.attn"].offset == tensors["layers.21.linear_attn"].offset - assert tensors["layers.19.attn"].offset == tensors["layers.22.linear_attn"].offset - assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100} - - -def test_get_kv_cache_config_mamba_hybrid_sharing_pp_group_count_bump(monkeypatch): - """PP=4's default partition [11,11,12,11] leaves stage 0 with 2 MLA - layers but a round-robin slice of 3 under the minimum 4 mamba groups; - grouping bumps to 5 groups so every stage's projection keeps sharing - on instead of silently falling back on stage 0.""" - from vllm.config import ParallelConfig - from vllm.distributed.utils import get_pp_indices - - monkeypatch.setattr(ModelConfig, "get_total_num_hidden_layers", lambda self: 45) - vllm_config = VllmConfig( - model_config=ModelConfig(max_model_len=8192), - parallel_config=ParallelConfig(pipeline_parallel_size=4), - ) - - kv_cache_spec, mamba_layers = _glm5_like_kv_cache_spec() - groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - mamba_groups = [g for g in groups if isinstance(g.kv_cache_spec, MambaSpec)] - assert [len(g.layer_names) for g in mamba_groups] == [7, 7, 7, 7, 6] - for k, name in enumerate(mamba_layers): - assert name in mamba_groups[k % 5].layer_names - - for rank in range(4): - start, end = get_pp_indices(45, rank, 4) - worker_spec = { - n: s - for n, s in kv_cache_spec.items() - if start <= int(n.split(".")[1]) < end - } - projected = kv_cache_utils._project_kv_cache_groups_to_worker( - groups, worker_spec + mamba = [ + g + for g in groups + if isinstance( + getattr(g.kv_cache_spec, "first_spec", g.kv_cache_spec), MambaSpec ) - assert kv_cache_utils._glm5_next_tensor_layout(projected) is not None - - -def test_get_kv_cache_config_mamba_hybrid_sharing_pp_starved_stage(monkeypatch): - """Reject PP stages whose Mamba layers have no MLA slot to share.""" - from vllm.config import ParallelConfig - - monkeypatch.setattr(ModelConfig, "get_total_num_hidden_layers", lambda self: 45) - monkeypatch.setenv("VLLM_PP_LAYER_PARTITION", "3,42") - vllm_config = VllmConfig( - model_config=ModelConfig(max_model_len=8192), - parallel_config=ParallelConfig(pipeline_parallel_size=2), - ) - - kv_cache_spec, _ = _glm5_like_kv_cache_spec() - with pytest.raises(ValueError, match="VLLM_PP_LAYER_PARTITION"): - kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - - -def test_get_kv_cache_config_mamba_hybrid_sharing_beats_cross_layers_flag(): - """Hybrid slot sharing must take precedence over the experimental - enable_cross_layers_blocks packed layout: generic packing would give the - MLA slots a strided view, breaking the contiguous virtual split.""" - model_config = ModelConfig(max_model_len=8192) - kv_transfer_config = KVTransferConfig( - kv_connector="NixlConnector", - kv_role="kv_both", - kv_connector_extra_config={"enable_cross_layers_blocks": "true"}, - ) - vllm_config = VllmConfig( - model_config=model_config, kv_transfer_config=kv_transfer_config - ) - - kv_cache_spec, _ = _glm5_like_kv_cache_spec() - groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - - mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes - idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes - bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) - assert bytes_per_block == 11 * mla_page + 11 * idx_page - - kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, groups, bytes_per_block * 100 + 1 - ) - assert kv_cache_config.num_blocks == 100 - assert len(kv_cache_config.kv_cache_tensors) == 56 - assert all( - tensor.block_stride in (mla_page, idx_page) - for tensor in kv_cache_config.kv_cache_tensors - ) - assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100} - - -def test_get_kv_cache_config_mamba_hybrid_sharing_no_indexer(): - """Kimi-Linear-like: MLA without indexer layers, idx_stride == 0.""" - model_config = ModelConfig(max_model_len=8192) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - - # 20 mamba + 7 MLA layers -> G = ceil(20 / 7) = 3 groups of [7, 7, 6]. - kv_cache_spec: dict[str, KVCacheSpec] = {} - for i in range(27): - if i % 4 == 3 or i == 26: - kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( - block_size=1024, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - ) - else: - kv_cache_spec[f"layers.{i}.linear_attn"] = new_mamba_spec() - mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes - - groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - mamba_groups = [ - group for group in groups if isinstance(group.kv_cache_spec, MambaSpec) ] - assert len(groups) == 4 - assert [len(group.layer_names) for group in mamba_groups] == [7, 7, 6] - for group in mamba_groups: - assert group.kv_cache_spec.page_size_bytes == mla_page - - bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) - assert bytes_per_block == 7 * mla_page + assert len(attn) == 1 and len(attn[0].layer_names) == 22 + assert len(rings) == 1 and len(rings[0].layer_names) == 11 + # 11 * (1.18 MB MLA + 34 KB indexer) per block fits 12 fp32-SSM KDA states. + assert len(mamba) == 3 + assert sum(len(g.layer_names) for g in mamba) == 34 kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, groups, bytes_per_block * 100 + 1 - ) - assert kv_cache_config.num_blocks == 100 - # Each group has one tensor descriptor and all groups alias the same slots. - assert len(kv_cache_config.kv_cache_tensors) == 4 - tensors = _tensor_by_layer(kv_cache_config) - mla_names = [f"layers.{i}.attn" for i in (*range(3, 27, 4), 26)] - for index, mla_name in enumerate(mla_names): - mla_offset = _layer_offset(tensors[mla_name], mla_name) - for group in mamba_groups: - if index < len(group.layer_names): - mamba_name = group.layer_names[index] - assert _layer_offset(tensors[mamba_name], mamba_name) == mla_offset - assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100} - - -def test_get_kv_cache_capacity_after_scheduler_unwrap(): - """max_concurrency must survive the scheduler-config unwrap. - - Regression for the balanced-mamba hybrid layout: - ``generate_scheduler_kv_cache_config`` flattens the MLA - ``UniformTypeKVCacheSpecs`` group into a single ``MLAAttentionSpec``, so the - scheduler config holds an MLA spec (page size A) next to ``MambaSpec`` - groups (page size B). ``EngineCore._initialize_kv_caches`` calls - ``get_kv_cache_capacity`` on that unwrapped config; the uniform-page-size - path used to assert-fail (``assert len(page_sizes) == 1``) on this topology. - """ - model_config = ModelConfig(max_model_len=8192) - vllm_config = VllmConfig(model_config=model_config) - - kv_cache_spec: dict[str, KVCacheSpec] = {} - for i in range(45): - if i % 4 == 3: - kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( - block_size=1024, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - ) - kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec( - block_size=1024, - num_kv_heads=1, - head_size=132, - dtype=torch.uint8, - tokens_per_state=16, - ) - else: - kv_cache_spec[f"layers.{i}.linear_attn"] = new_mamba_spec() - - groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) - kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, groups, bytes_per_block * 100 + 1 - ) - - scheduler_config = generate_scheduler_kv_cache_config([kv_cache_config]) - # Confirm we are exercising the regression path: the MLA group is no longer - # UniformTypeKVCacheSpecs, so sibling specs have differing page sizes. - assert not any( - isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) - for g in scheduler_config.kv_cache_groups - ) - - unwrapped_groups = scheduler_config.kv_cache_groups - expected_max_mem = kv_cache_utils._max_memory_usage_bytes_from_groups( - vllm_config, unwrapped_groups - ) - expected_pool = kv_cache_utils._pool_bytes_per_block(unwrapped_groups) - expected_blocks_per_request = ( - expected_max_mem + expected_pool - 1 - ) // expected_pool - - _, max_concurrency = get_kv_cache_capacity(vllm_config, scheduler_config) - assert max_concurrency > 0 - assert max_concurrency == pytest.approx( - scheduler_config.num_blocks / expected_blocks_per_request + vllm_config, groups, available_memory=2 * 1024**3 ) + strides = {t.block_stride for t in kv_cache_config.kv_cache_tensors} + assert len(strides) == 1 + (block_stride,) = strides + # Blocks sit a whole number of 32-row MLA pages and 64-state indexer pages + # apart so both caches can be re-paged for their kernels. + assert block_stride % (32 * 576 * 2) == 0 + assert block_stride % (64 * 132) == 0 def new_mla_spec(cache_dtype_str=None, block_size: int = 16): @@ -4483,6 +3993,7 @@ def test_kv_cache_reserves_null_block_for_max_model_len(use_override): block_size = 16 max_model_len = 512 # needs 512 / 16 = 32 blocks vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) + vllm_config.cache_config.kv_cache_layout = "LBNHC" spec = new_kv_cache_spec(block_size=block_size) # 32 blocks -> only 31 usable after the null block: one short -> reject. @@ -4528,6 +4039,7 @@ def test_check_enough_kv_cache_memory_reserves_null_block(): block_size = 16 max_model_len = 512 # needs 512 / 16 = 32 blocks vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) + vllm_config.cache_config.kv_cache_layout = "LBNHC" spec = new_kv_cache_spec(block_size=block_size) # 32 blocks -> only 31 usable after the null block: one short -> reject. diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index 415609aefb8d..f926f496514a 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -20,7 +20,6 @@ ChunkedLocalAttentionManager, CircularBufferManager, FullAttentionManager, - KpoolTailManager, MambaManager, RSWAManager, SlidingWindowManager, @@ -29,7 +28,6 @@ ChunkedLocalAttentionSpec, CircularBufferSpec, FullAttentionSpec, - KpoolTailSpec, MambaSpec, RSWASpec, SlidingWindowSpec, @@ -406,15 +404,15 @@ def test_circular_buffer_allocates_one_block_for_the_request_lifetime(): @pytest.mark.parametrize("record_for_zeroing", [True, False]) def test_external_kpool_tail_zeroing(record_for_zeroing): - """A transferred ring at index zero skips zeroing even with a local hit.""" - spec = KpoolTailSpec( + """The kpool tail ring is never zeroed, transferred or not: a pool only + reads ring slots its own request wrote.""" + spec = CircularBufferSpec( block_size=4, num_kv_heads=2, head_size=8, dtype=torch.bfloat16, - sliding_window=4, ) - manager = KpoolTailManager( + manager = CircularBufferManager( spec, block_pool=BlockPool(10, enable_caching=True, hash_block_size=4), enable_caching=True, @@ -428,9 +426,7 @@ def test_external_kpool_tail_zeroing(record_for_zeroing): ) blocks = manager.req_to_blocks["request"] assert len(blocks) == 1 - assert manager.take_new_block_ids() == ( - [blocks[0].block_id] if record_for_zeroing else [] - ) + assert manager.take_new_block_ids() == [] assert manager.allocate_new_blocks("request", 20, 20) == [] assert manager.take_new_block_ids() == [] diff --git a/tests/v1/kv_connector/unit/conftest.py b/tests/v1/kv_connector/unit/conftest.py index 58c6e632d3d9..85862da685a7 100644 --- a/tests/v1/kv_connector/unit/conftest.py +++ b/tests/v1/kv_connector/unit/conftest.py @@ -23,5 +23,5 @@ def _mock_cpu_backend_block_sizes(monkeypatch): monkeypatch.setattr( CPUAttentionBackend, "get_supported_kernel_block_sizes", - staticmethod(lambda: [MultipleOf(16)]), + staticmethod(lambda kv_cache_spec=None: [MultipleOf(16)]), ) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py b/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py index 23856b4541aa..969bf65f8c54 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py @@ -164,13 +164,12 @@ def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch): fa_spec = kv_cache_config.kv_cache_groups[0].kv_cache_spec gdn_spec = kv_cache_config.kv_cache_groups[1].kv_cache_spec fa_raw = torch.empty(num_blocks * fa_spec.page_size_bytes, dtype=torch.int8) - gdn_raw = torch.empty(num_blocks * gdn_spec.page_size_bytes, dtype=torch.int8) + gdn_block_stride = gdn_spec.page_size_bytes + 64 + gdn_raw = torch.empty(num_blocks, gdn_block_stride, dtype=torch.int8) (fa_cache,) = dense_kv_cache_views( fa_raw, fa_spec, num_blocks, 1, KVCacheLayout.LBHNC ) - (gdn_cache,) = dense_kv_cache_views( - gdn_raw, gdn_spec, num_blocks, 1, KVCacheLayout.LBHNC - ) + gdn_cache = gdn_raw[:, : gdn_spec.page_size_bytes] worker.register_kv_caches( { @@ -184,7 +183,7 @@ def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch): "model.layers.0.self_attn", "model.layers.1.linear_attn", ] - assert worker.block_len_per_layer == [ + assert worker.kv_block_len_per_layer == [ fa_spec.page_size_bytes, gdn_spec.page_size_bytes, ] diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index 6c2cff0a542c..a8354390d71e 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -31,7 +31,6 @@ from vllm.v1.kv_cache_interface import ( CircularBufferSpec, FullAttentionSpec, - KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor, @@ -711,13 +710,12 @@ def test_worker_setup_tolerates_finer_scratch_group(): mamba_cache_mode="align", ) # Scratch block 4 is not divisible by the hash unit 8 below. - scratch = KpoolTailSpec( + scratch = CircularBufferSpec( block_size=4, num_kv_heads=2, head_size=64, head_size_v=0, dtype=torch.bfloat16, - sliding_window=4, ) cfg = KVCacheConfig( num_blocks=4, diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index 9f6bc56deeb5..4e7a7b843143 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -641,164 +641,43 @@ def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blo @pytest.mark.cpu_test -@pytest.mark.parametrize("logical_block_size", [1152, 640]) -@pytest.mark.parametrize("tail_first", [False, True]) -def test_register_compressed_indexer_uses_virtual_transfer_pages( - logical_block_size, tail_first -): - """Compressed indexer rows must split into contiguous NIXL transfer pages.""" - from unittest.mock import MagicMock - - from vllm.config import set_current_vllm_config - from vllm.distributed.kv_transfer.kv_connector.v1.nixl import ( - base_worker as bw, - ) - from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( - NixlConnectorWorker, - ) - from vllm.v1.kv_cache_interface import ( - KpoolTailSpec, - KVCacheConfig, - KVCacheGroupSpec, - KVCacheLayout, - KVCacheTensor, - MLAAttentionSpec, - UniformTypeKVCacheSpecs, - create_kv_cache_views, - ) - - num_logical_blocks = 3 - transfer_block_size = 64 - kernel_block_size = 128 - tokens_per_state = 4 - state_content_bytes = 132 - - indexer_spec = MLAAttentionSpec( - block_size=logical_block_size, - num_kv_heads=1, - head_size=128, - head_size_v=0, - dtype=torch.uint8, - state_content_bytes=state_content_bytes, - tokens_per_state=tokens_per_state, - ) - indexer_page_size = indexer_spec.page_size_bytes - tail_spec = KpoolTailSpec( - block_size=tokens_per_state, - num_kv_heads=2, - head_size=128, - head_size_v=0, - dtype=torch.bfloat16, - page_size_padded=indexer_page_size, - sliding_window=tokens_per_state, - ) - - allocation_size = num_logical_blocks * indexer_page_size - indexer_tensor = KVCacheTensor( - size=allocation_size, - layers=["indexer"], - layer_stride=allocation_size, - block_stride=indexer_page_size, - ) - tail_tensor = KVCacheTensor( - size=allocation_size, - layers=["tail"], - layer_stride=allocation_size, - block_stride=indexer_page_size, - ) - kv_cache_config = KVCacheConfig( - num_blocks=num_logical_blocks, - kv_cache_tensors=[indexer_tensor, tail_tensor], - kv_cache_groups=[ - KVCacheGroupSpec( - ["indexer"], - UniformTypeKVCacheSpecs( - block_size=logical_block_size, - kv_cache_specs={"indexer": indexer_spec}, - ), - ), - KVCacheGroupSpec( - ["tail"], - UniformTypeKVCacheSpecs( - block_size=tokens_per_state, - kv_cache_specs={"tail": tail_spec}, - ), - ), - ], - ) - - raw = torch.zeros(allocation_size, dtype=torch.int8) - (indexer_cache,) = create_kv_cache_views( - raw, - indexer_spec, - num_logical_blocks, - KVCacheLayout.LBHNC, - indexer_tensor, - kernel_block_size=kernel_block_size, +def test_nixl_keeps_packed_sparse_mla_manager_block(): + """Like the worker, NIXL keeps packed sparse MLA + indexer manager blocks.""" + from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase + from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseBackend + from vllm.v1.attention.backends.mla.indexer import Glm5NextIndexerBackend + from vllm.v1.kv_cache_interface import UniformTypeKVCacheSpecs + + specs = { + "mla": MLAAttentionSpec( + 1152, num_kv_heads=1, head_size=512, dtype=torch.bfloat16 + ), + "indexer": MLAAttentionSpec( + 1152, num_kv_heads=1, head_size=128, dtype=torch.uint8, tokens_per_state=4 + ), + } + backends = {"mla": FlashMLASparseBackend, "indexer": Glm5NextIndexerBackend} + layers = { + name: MagicMock(spec=AttentionLayerBase, **{"get_attn_backend.return_value": b}) + for name, b in backends.items() + } + worker = object.__new__(bw.NixlBaseConnectorWorker) + worker.block_size, worker.num_blocks, worker.attn_backends = 1152, 2, [] + worker._physical_blocks_per_logical_kv_block = 1 + worker.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC + ), + compilation_config=SimpleNamespace(static_forward_context=layers), ) - (tail_cache,) = create_kv_cache_views( - raw, - tail_spec, - num_logical_blocks, - KVCacheLayout.LBHNC, - tail_tensor, + worker.kv_cache_config = KVCacheConfig( + 2, [], [KVCacheGroupSpec(list(specs), UniformTypeKVCacheSpecs(1152, specs))] ) - assert indexer_cache.data_ptr() == tail_cache.data_ptr() == raw.data_ptr() - vllm_config = create_vllm_config(block_size=logical_block_size) - vllm_config.cache_config.kv_cache_layout = "LBHNC" - vllm_config.kv_transfer_config.kv_buffer_device = "cuda" - fake_backend = MagicMock() - fake_backend.get_supported_kernel_block_sizes.return_value = [transfer_block_size] - fake_backend.get_name.return_value = "DEEPSEEK_V32_INDEXER" - fake_backend.full_cls_name.return_value = "fake.DEEPSEEK_V32_INDEXER" - fake_platform = MagicMock() - fake_platform.device_type = "cuda" - fake_platform.get_nixl_memory_type.return_value = "VRAM" - - caches = [("indexer", indexer_cache), ("tail", tail_cache)] - if tail_first: - caches.reverse() + worker._sync_block_size_with_kernel() - with ( - patch.object(bw, "NixlWrapper", _RecordingNixl), - patch.object(bw, "get_tensor_model_parallel_rank", return_value=0), - patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1), - patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]), - patch.object(bw, "current_platform", fake_platform), - set_current_vllm_config(vllm_config), - ): - worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config) - worker.use_mla = True - worker.register_kv_caches(dict(caches)) - - transfer_page_size = transfer_block_size // tokens_per_state * state_content_bytes - num_transfer_blocks = num_logical_blocks * ( - logical_block_size // transfer_block_size - ) - expected_descs = np.asarray( - [ - [ - raw.data_ptr() + block_idx * transfer_page_size, - transfer_page_size, - 0, - ] - for block_idx in range(num_transfer_blocks) - ], - dtype=np.uint64, - ) - - assert worker.block_size == transfer_block_size - assert worker.num_regions == 1 - assert worker.block_len_per_layer == [transfer_page_size] - assert worker.block_stride_per_layer == [transfer_page_size] - assert worker._region_is_mla == [True] - assert worker.kv_caches_base_addr[worker.engine_id][0] == [raw.data_ptr()] - assert worker._registered_descs[0] == [(raw.data_ptr(), raw.nbytes, 0, "")] - np.testing.assert_array_equal(worker.src_blocks_data, expected_descs) - assert expected_descs[-1, 0] + expected_descs[-1, 1] == ( - raw.data_ptr() + raw.nbytes - ) + assert worker.block_size == 1152 + assert worker._physical_blocks_per_logical_kv_block == 1 def _make_remote_meta( diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 08759bc1e34b..6c0cc310c729 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -47,7 +47,6 @@ from vllm.v1.kv_cache_interface import ( CircularBufferSpec, FullAttentionSpec, - KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor, @@ -145,13 +144,12 @@ def _make_scratch_kv_cache_config( register_all_kvcache_specs(vllm_config=None) fa_config = _make_kv_cache_config(num_blocks, num_groups=1) scratch_layers = ["layer_scratch"] - scratch_spec = KpoolTailSpec( + scratch_spec = CircularBufferSpec( block_size=scratch_block_size, num_kv_heads=2, head_size=HEAD_SIZE, head_size_v=0, dtype=DTYPE, - sliding_window=scratch_block_size, ) assert not scratch_spec.prefix_cacheable scratch_bytes = scratch_spec.page_size_bytes * num_blocks diff --git a/tests/v1/spec_decode/test_dflash_prepare_inputs.py b/tests/v1/spec_decode/test_dflash_prepare_inputs.py index 1d004c36677a..3fc5caec2e08 100644 --- a/tests/v1/spec_decode/test_dflash_prepare_inputs.py +++ b/tests/v1/spec_decode/test_dflash_prepare_inputs.py @@ -102,7 +102,7 @@ class InputsPrepared(Exception): block_tables=SimpleNamespace( slot_mappings=query_slot_mapping.unsqueeze(0), input_block_tables=[block_table], - kernel_block_sizes=[4], + block_sizes=[4], cp_rank=cp_rank, cp_size=cp_size, cp_interleave=cp_interleave, diff --git a/tests/v1/spec_decode/test_eagle.py b/tests/v1/spec_decode/test_eagle.py index 2ddd0a977481..62a9a1635c1f 100644 --- a/tests/v1/spec_decode/test_eagle.py +++ b/tests/v1/spec_decode/test_eagle.py @@ -963,6 +963,9 @@ def create_deterministic_logits(token_ids): # Mock draft_attn_groups for attention metadata building. mock_attn_group = mock.MagicMock() mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder + mock_attn_group.build_metadata_for_drafting = ( + attn_metadata_builder.build_for_drafting + ) mock_attn_group.layer_names = list(proposer._draft_attn_layer_names) mock_attn_group.kv_cache_spec = attn_metadata_builder.kv_cache_spec proposer.draft_attn_groups = [mock_attn_group] @@ -1066,6 +1069,9 @@ def fake_compute_probs(logits, sampling_metadata, use_fp64_gumbel): ) mock_attn_group = mock.MagicMock() mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder + mock_attn_group.build_metadata_for_drafting = ( + attn_metadata_builder.build_for_drafting + ) mock_attn_group.layer_names = list(proposer._draft_attn_layer_names) mock_attn_group.kv_cache_spec = attn_metadata_builder.kv_cache_spec proposer.draft_attn_groups = [mock_attn_group] diff --git a/tests/v1/spec_decode/test_llm_base_proposer.py b/tests/v1/spec_decode/test_llm_base_proposer.py index fb1ecdd2df12..086fd467acd4 100644 --- a/tests/v1/spec_decode/test_llm_base_proposer.py +++ b/tests/v1/spec_decode/test_llm_base_proposer.py @@ -2,10 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for SpecDecodeBaseProposer.initialize_attn_backend. -Block tables are stored at kernel-block granularity, so the proposer's -``block_size`` (used for slot-mapping math) must be the kernel block size, -not the KV cache manager's block size — the two differ when manager blocks -are split for the attention kernel. The value must also be deterministic: +Block tables are stored at manager-block granularity, so the proposer's +``block_size`` (used for slot-mapping math) must be the KV cache manager's +block size. Attention groups map to kernel blocks for their metadata builders. +The value must also be deterministic: ``_draft_attn_layer_names`` is a set, whose iteration order varies across processes, so anything derived from iteration order must not leak into ``block_size``. @@ -68,9 +68,9 @@ def _make_kv_cache_config(layer_names: set[str]) -> KVCacheConfig: return SimpleNamespace(kv_cache_groups=[group]) # type: ignore[return-value] -def test_block_size_uses_kernel_block_size(monkeypatch: pytest.MonkeyPatch): - """The proposer's slot-mapping math runs against the kernel-granularity - block table, so block_size must come from kernel_block_sizes.""" +def test_block_size_uses_manager_block_size(monkeypatch: pytest.MonkeyPatch): + """The proposer's slot-mapping math runs against the manager-granularity + block table; attention groups map kernel blocks for metadata builders.""" layer_names = {"draft.0.self_attn.attn"} proposer = _make_proposer(monkeypatch, layer_names) @@ -79,12 +79,8 @@ def test_block_size_uses_kernel_block_size(monkeypatch: pytest.MonkeyPatch): kernel_block_sizes=[KERNEL_BLOCK_SIZE], ) - assert proposer.block_size == KERNEL_BLOCK_SIZE - assert proposer.block_size != SCHEDULER_BLOCK_SIZE - # The metadata builder keeps receiving the kernel block size as well. - group = proposer.draft_attn_groups[0] - assert isinstance(group, _FakeAttentionGroup) - assert group.kernel_block_size == KERNEL_BLOCK_SIZE + assert proposer.block_size == SCHEDULER_BLOCK_SIZE + assert proposer.draft_attn_groups[0].kernel_block_size == KERNEL_BLOCK_SIZE def test_block_size_falls_back_to_kv_cache_spec(monkeypatch: pytest.MonkeyPatch): @@ -112,4 +108,4 @@ def test_draft_layer_iteration_is_deterministic(monkeypatch: pytest.MonkeyPatch) ) assert len(proposer.draft_attn_groups) == 1 assert proposer.draft_attn_groups[0].layer_names == expected_order - assert proposer.block_size == KERNEL_BLOCK_SIZE + assert proposer.block_size == SCHEDULER_BLOCK_SIZE diff --git a/tests/v1/streaming_input/test_gpu_model_runner_streaming.py b/tests/v1/streaming_input/test_gpu_model_runner_streaming.py index 9457c66050db..7002176dcc57 100644 --- a/tests/v1/streaming_input/test_gpu_model_runner_streaming.py +++ b/tests/v1/streaming_input/test_gpu_model_runner_streaming.py @@ -38,7 +38,6 @@ def mock_model_runner_with_input_batch(): device="cpu", vocab_size=32000, block_sizes=[16], - kernel_block_sizes=[16], max_num_blocks_per_req=[64], logitsprocs=None, is_pooling_model=False, diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 9c9c26c3c2c0..8d30403f30fc 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -15,7 +15,11 @@ import torch import vllm.v1.hisparse.binding as attn_utils_module -from tests.v1.attention.utils import dense_kv_cache_views +from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + dense_kv_cache_views, +) from vllm.config.compilation import CompilationConfig, CUDAGraphMode from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy @@ -42,6 +46,8 @@ AttentionGroup, allocate_kv_cache, copy_kv_cache_blocks_inplace, + customize_attention_spec, + map_kv_caches_to_kernel_blocks, ) @@ -79,7 +85,7 @@ def test_get_kv_cache_spec_resolves_hisparse_block_size( backend = SimpleNamespace( get_name=lambda name=name: name, customize_spec=AttentionBackend.customize_spec, - get_supported_kernel_block_sizes=lambda sizes=sizes: sizes, + get_supported_kernel_block_sizes=lambda _=None, sizes=sizes: sizes, ) layers[name] = SimpleNamespace( get_kv_cache_spec=lambda _, spec=specs[name]: spec, @@ -99,7 +105,8 @@ def test_get_kv_cache_spec_resolves_hisparse_block_size( resolved = attn_utils.get_kv_cache_spec(config) assert resolved["main"].block_size == resolved["indexer"].block_size == expected - assert resolved["dense"] is specs["dense"] + dense_backend = layers["dense"].get_attn_backend() + assert resolved["dense"] == customize_attention_spec(dense_backend, specs["dense"]) assert all(spec.block_size == block_size for spec in specs.values()) @@ -495,20 +502,17 @@ def test_reshape_padded_kv_cache_strides_by_padded_page(): @pytest.mark.parametrize( ( "kernel_block_sizes", - "storage_block_size", "expected_num_blocks", "expected_num_states", ), [ - (None, None, 4, 64), - ([256], None, 4, 64), - ([64], None, 16, 16), - ([64], 256, 4, 64), + (None, 4, 64), + ([256], 4, 64), + ([64], 16, 16), ], ) def test_allocate_compressed_mla_cache( kernel_block_sizes: list[int] | None, - storage_block_size: int | None, expected_num_blocks: int, expected_num_states: int, ): @@ -518,7 +522,6 @@ def test_allocate_compressed_mla_cache( head_size=128, dtype=torch.bfloat16, tokens_per_state=4, - storage_block_size=storage_block_size, ) num_pages = 4 config = KVCacheConfig( @@ -744,3 +747,92 @@ def host_allocator(size: int) -> torch.Tensor: backing.untyped_storage().data_ptr() == caches["indexer"].untyped_storage().data_ptr() ) + + +class _TableBuilder: + requires_block_table_width = False + supports_update_block_table = False + + def __init__(self, kv_cache_spec, layer_names, vllm_config, device): + self.kv_cache_spec = kv_cache_spec + + def build(self, common_prefix_len, common_attn_metadata, fast_build=False): + return common_attn_metadata.block_table_tensor + + +_VLLM_CONFIG = SimpleNamespace( + model_config=SimpleNamespace(max_model_len=8192), + scheduler_config=SimpleNamespace(max_num_seqs=4), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), +) + + +def _packed_group(tokens_per_state, row_bytes, supported_sizes): + spec = MLAAttentionSpec( + block_size=1024, + num_kv_heads=1, + head_size=128, + dtype=torch.uint8, + state_content_bytes=row_bytes, + tokens_per_state=tokens_per_state, + ) + backend = SimpleNamespace( + get_builder_cls=lambda: _TableBuilder, + get_supported_kernel_block_sizes=lambda _=None: supported_sizes, + get_name=lambda: "FAKE", + customize_spec=lambda spec: spec, + ) + return AttentionGroup(backend, ["layer.0"], spec, 0) + + +def test_packed_compressed_group_gets_kernel_block_table(): + """A packed indexer cache whose kernel needs 64-state pages gets a table and + a bound view in kernel blocks that address the same packed-block rows.""" + row, stride_rows = 132, 5 * 64 + group = _packed_group(4, row, [128, 256]) + group.create_metadata_builders(_VLLM_CONFIG, "cpu", kernel_block_size=1024) + assert group.get_metadata_builder().kv_cache_spec.block_size == 256 + + backing = torch.arange(3 * stride_rows).repeat_interleave(row) + manager = backing.view(3, 1, stride_rows, row)[:, :, :256] + pages = map_kv_caches_to_kernel_blocks({"layer.0": manager}, [group])["layer.0"] + metadata = create_common_attn_metadata(BatchSpec([2048], [1]), 1024, "cpu") + metadata.block_table_tensor = torch.tensor([[0, 2]], dtype=torch.int32) + table = group.build_metadata(metadata) + assert table.tolist() == [[0, 1, 2, 3, 10, 11, 12, 13]] + assert torch.equal(pages[table[0, 4:], 0].flatten(0, 1), manager[2, 0]) + + +def test_split_group_maps_manager_block_table(): + """Groups split in a dense layout keep manager block tables; builders get + b * blocks_per_kv_block + j against the view allocated in kernel blocks.""" + group = _packed_group(4, 132, [128, 256]) + group.create_metadata_builders(_VLLM_CONFIG, "cpu", kernel_block_size=256) + kernel_view = torch.zeros(12, 1, 64, 132) + bound = map_kv_caches_to_kernel_blocks({"layer.0": kernel_view}, [group]) + assert bound["layer.0"] is kernel_view + + metadata = create_common_attn_metadata(BatchSpec([2048], [1]), 1024, "cpu") + metadata.block_table_tensor = torch.tensor([[0, 2]], dtype=torch.int32) + table = group.build_metadata(metadata) + assert table.tolist() == [[0, 1, 2, 3, 8, 9, 10, 11]] + + +@pytest.mark.parametrize( + ("block_size", "page_states"), [(1024, 64), (384, 32), (256, 0)] +) +def test_split_blocks_align_to_kernel_blocks(block_size, page_states): + """Backends taking only fixed kernel blocks align packed blocks to the kernel + block the final block size is split into, and not at all when unsplit.""" + group = _packed_group(4, 132, [128, 256]) + spec = customize_attention_spec(group.backend, group.kv_cache_spec) + spec = spec.copy_with_new_block_size(block_size) + assert spec.get_block_stride_alignment() == max(page_states * 132, 1) + + +def test_packed_token_cache_cannot_be_split(): + group = _packed_group(1, 1152, [64]) + group.create_metadata_builders(_VLLM_CONFIG, "cpu", kernel_block_size=1024) + packed = torch.zeros(3, 1, 2048, 1152, dtype=torch.uint8)[:, :, :1024] + with pytest.raises(ValueError, match="cannot be split"): + map_kv_caches_to_kernel_blocks({"layer.0": packed}, [group]) diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 5c2231a1f267..dac9f12ae69a 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -22,7 +22,6 @@ def test_block_tables_apply_staged_writes_fuses_kv_groups(monkeypatch): max_num_batched_tokens=64, max_num_blocks_per_group=[8, 8, 8], device=device, - kernel_block_sizes=[16, 16, 8], ) def fail_if_apply_write_called(): @@ -48,33 +47,31 @@ def fail_if_apply_write_called(): block_tables.block_tables[0].gpu[0, :2], torch.tensor([1, 2], dtype=torch.int32, device=device), ) - # Group 1 has blocks_per_kv_block == 2, so each KV block expands to two - # kernel block IDs. assert torch.equal( - block_tables.block_tables[1].gpu[0, :4], - torch.tensor([20, 21, 22, 23], dtype=torch.int32, device=device), + block_tables.block_tables[1].gpu[0, :2], + torch.tensor([10, 11], dtype=torch.int32, device=device), ) assert torch.equal( block_tables.block_tables[0].gpu[1, :1], torch.tensor([3], dtype=torch.int32, device=device), ) assert torch.equal( - block_tables.block_tables[1].gpu[1, :2], - torch.tensor([24, 25], dtype=torch.int32, device=device), + block_tables.block_tables[1].gpu[1, :1], + torch.tensor([12], dtype=torch.int32, device=device), ) assert torch.equal( block_tables.block_tables[2].gpu[1, :2], torch.tensor([5, 6], dtype=torch.int32, device=device), ) assert block_tables.num_blocks.np[0, 0] == 2 - assert block_tables.num_blocks.np[1, 0] == 4 + assert block_tables.num_blocks.np[1, 0] == 2 assert block_tables.num_blocks.np[2, 0] == 0 assert block_tables.num_blocks.np[0, 1] == 1 - assert block_tables.num_blocks.np[1, 1] == 2 + assert block_tables.num_blocks.np[1, 1] == 1 assert block_tables.num_blocks.np[2, 1] == 2 assert torch.equal( block_tables.num_blocks.gpu[:, :2], - torch.tensor([[2, 1], [4, 2], [0, 2]], dtype=torch.int32, device=device), + torch.tensor([[2, 1], [2, 1], [0, 2]], dtype=torch.int32, device=device), ) for block_table in block_tables.block_tables: @@ -96,15 +93,15 @@ def fail_if_apply_write_called(): torch.tensor([1, 2, 7], dtype=torch.int32, device=device), ) assert torch.equal( - block_tables.block_tables[1].gpu[0, :6], - torch.tensor([20, 21, 22, 23, 26, 27], dtype=torch.int32, device=device), + block_tables.block_tables[1].gpu[0, :3], + torch.tensor([10, 11, 13], dtype=torch.int32, device=device), ) assert torch.equal( block_tables.block_tables[2].gpu[0, :1], torch.tensor([8], dtype=torch.int32, device=device), ) assert block_tables.num_blocks.np[0, 0] == 3 - assert block_tables.num_blocks.np[1, 0] == 6 + assert block_tables.num_blocks.np[1, 0] == 3 assert block_tables.num_blocks.np[2, 0] == 1 @@ -116,7 +113,6 @@ def test_block_tables_apply_staged_writes_single_group(): max_num_batched_tokens=16, max_num_blocks_per_group=[4], device=device, - kernel_block_sizes=[16], ) block_tables.append_block_ids( @@ -141,7 +137,6 @@ def test_block_tables_skip_custom_slot_mapping_groups(): max_num_batched_tokens=4, max_num_blocks_per_group=[1, 1], device=device, - kernel_block_sizes=[8, 262144], slot_mapping_enabled=[False, True], ) block_tables.append_block_ids( @@ -170,7 +165,7 @@ def test_block_tables_skip_custom_slot_mapping_groups(): @pytest.mark.parametrize("cp_rank", range(4)) -def test_dcp_slot_mapping_with_smaller_kernel_blocks(cp_rank: int): +def test_dcp_slot_mapping_only_shards_sharded_groups(cp_rank: int): """Only sharded groups use DCP interleave in logical-block coordinates.""" device = torch.device("cuda") block_tables = BlockTables( @@ -179,7 +174,6 @@ def test_dcp_slot_mapping_with_smaller_kernel_blocks(cp_rank: int): max_num_batched_tokens=1024, max_num_blocks_per_group=[2, 8], device=device, - kernel_block_sizes=[64, 64], dcp_sharded=[True, False], cp_size=4, cp_rank=cp_rank, @@ -230,7 +224,6 @@ def test_v1_block_table_move_row_clears_vacated_row(): max_num_batched_tokens=64, pin_memory=False, device=torch.device("cuda"), - kernel_block_size=16, cp_kv_cache_interleave_size=1, ) block_table.add_row([7, 8, 9], row_idx=0) @@ -260,7 +253,6 @@ def test_get_dummy_block_tables_returns_zeroed_rows(): max_num_batched_tokens=64, max_num_blocks_per_group=[8], device=device, - kernel_block_sizes=[16], ) # Simulate a real step: stage a request's blocks and gather them into # the persistent input block tables. @@ -292,7 +284,6 @@ def test_dummy_request_slot_mapping_is_pad(): max_num_batched_tokens=16, max_num_blocks_per_group=[4], device=device, - kernel_block_sizes=[4], ) block_tables.append_block_ids(req_index=0, new_block_ids=([7, 8],), overwrite=True) block_tables.apply_staged_writes() diff --git a/tests/v1/worker/test_gpu_input_batch.py b/tests/v1/worker/test_gpu_input_batch.py index f37e824e6c63..21157aae33b2 100644 --- a/tests/v1/worker/test_gpu_input_batch.py +++ b/tests/v1/worker/test_gpu_input_batch.py @@ -233,7 +233,6 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int): device=torch.device(device), vocab_size=1024, block_sizes=[1], - kernel_block_sizes=[1], max_num_blocks_per_req=[1024], ) reqs: list[CachedRequestState] = [] @@ -327,7 +326,6 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis device=torch.device(device), vocab_size=1024, block_sizes=[1], - kernel_block_sizes=[1], max_num_blocks_per_req=[1024], ) ref_input_batch: InputBatch = InputBatch( @@ -337,7 +335,6 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis device=torch.device(device), vocab_size=1024, block_sizes=[1], - kernel_block_sizes=[1], max_num_blocks_per_req=[1024], ) @@ -388,7 +385,6 @@ def test_condense_clears_stale_allowed_token_ids_mask(device: str): device=torch.device(device), vocab_size=VOCAB_SIZE, block_sizes=[1], - kernel_block_sizes=[1], max_num_blocks_per_req=[1024], ) @@ -446,7 +442,6 @@ def test_swap_states_preserves_allowed_token_ids_mask(device: str): device=torch.device(device), vocab_size=VOCAB_SIZE, block_sizes=[1], - kernel_block_sizes=[1], max_num_blocks_per_req=[1024], ) @@ -516,7 +511,6 @@ def test_pooling_prompt_lens_not_aliased(device: str): device=torch.device(device), vocab_size=VOCAB_SIZE, block_sizes=[16], - kernel_block_sizes=[16], max_num_blocks_per_req=[64], is_pooling_model=True, ) @@ -552,7 +546,6 @@ def test_placeholder_spec_token_ids_written_verbatim(): device=torch.device("cpu"), vocab_size=VOCAB_SIZE, block_sizes=[16], - kernel_block_sizes=[16], max_num_blocks_per_req=[1], ) req = CachedRequestState( @@ -599,7 +592,6 @@ def test_pooling_metadata_token_id_buffers( device=torch.device("cpu"), vocab_size=VOCAB_SIZE, block_sizes=[16], - kernel_block_sizes=[16], max_num_blocks_per_req=[64], is_pooling_model=True, ) @@ -630,7 +622,6 @@ def _make_input_batch(is_pooling_model: bool = False) -> InputBatch: device=torch.device("cpu"), vocab_size=VOCAB_SIZE, block_sizes=[16], - kernel_block_sizes=[16], max_num_blocks_per_req=[64], is_pooling_model=is_pooling_model, ) diff --git a/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py b/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py index 8a7445f7d49a..d9d5c4de63af 100644 --- a/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py +++ b/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py @@ -5,7 +5,7 @@ The kpool tail cache holds one circular block per request, so its block table row is orders of magnitude shorter than the sequence it serves. `_compute_slot_mappings_kernel` indexes block tables by -`position // kernel_block_size` with an effectively unmasked load, so +`position // block_size` with an effectively unmasked load, so this leads to a read of neighboring requests' rows (and for the later requests, outside the buffer entirely). """ @@ -15,7 +15,7 @@ from vllm.platforms import current_platform from vllm.v1.attention.backends.utils import PAD_SLOT_ID -from vllm.v1.kv_cache_interface import KpoolTailSpec +from vllm.v1.kv_cache_interface import CircularBufferSpec from vllm.v1.worker.block_table import get_block_table_width from vllm.v1.worker.gpu.block_table import BlockTables @@ -32,13 +32,12 @@ def make_tail_spec(): - return KpoolTailSpec( + return CircularBufferSpec( block_size=KPOOL, num_kv_heads=2, head_size=128, head_size_v=0, dtype=torch.bfloat16, - sliding_window=KPOOL, ) @@ -67,7 +66,6 @@ def test_kpool_tail_group_never_position_indexes_its_block_table(): max_num_batched_tokens=len(positions), max_num_blocks_per_group=[mla_width, tail_width], device=device, - kernel_block_sizes=[MLA_BLOCK_SIZE, spec.block_size], slot_mapping_enabled=[True, spec.uses_slot_mapping], ) mla_blocks = list(range(1, mla_width + 1)) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 3851d5dad6bd..158befc15027 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -111,9 +111,6 @@ def initialize_kv_cache(runner: GPUModelRunner): device=runner.device, vocab_size=runner.model_config.get_vocab_size(), block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size], - kernel_block_sizes=[ - kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size - ], max_num_blocks_per_req=[NUM_BLOCKS], ) runner.initialize_attn_backend(kv_cache_config) @@ -303,7 +300,7 @@ def get_name(): return "MOCK" @staticmethod - def get_supported_kernel_block_sizes(): + def get_supported_kernel_block_sizes(kv_cache_spec=None): return supported_sizes return _MockBackend() @@ -441,7 +438,7 @@ def get_name() -> str: return "MOCK_EXACT" if exact else "MOCK" @staticmethod - def get_supported_kernel_block_sizes(): + def get_supported_kernel_block_sizes(kv_cache_spec=None): return list(supported) if exact: @@ -1625,7 +1622,6 @@ def test_input_batch_reinitialized_after_late_interleave_adjustment(monkeypatch) runner.device = torch.device("cpu") runner.is_pooling_model = False runner._init_block_sizes = [16] - runner._init_kernel_block_sizes = [16] runner._init_max_num_blocks = [4] runner._init_slot_mapping_modes = [ gpu_model_runner_module.SlotMappingMode.TOKEN_TO_KV_SLOT @@ -1653,7 +1649,7 @@ def test_input_batch_reinitialized_after_late_interleave_adjustment(monkeypatch) lambda _: gpu_model_runner_module.KVCacheSpecKind.FULL_ATTENTION, ) - runner.may_reinitialize_input_batch(kv_cache_config, [16]) + runner.may_reinitialize_input_batch(kv_cache_config) assert input_batch_cls.call_count == 1 assert input_batch_cls.call_args.kwargs["cp_kv_cache_interleave_size"] == 16 @@ -1681,68 +1677,10 @@ class StopInitialization(Exception): assert runner.cp_interleave == 16 -def test_hybrid_block_table_initialization(): - """Test hybrid block table with different kernel and kvcache_manager block - sizes.""" - from vllm.v1.worker.block_table import BlockTable - - # Test configuration: kvcache_manager block size = 32, - # kernel block size = 16 - block_size = 32 - kernel_block_sizes = [16] - max_num_reqs = 10 - max_num_blocks_per_req = 20 - max_num_batched_tokens = 512 - cp_kv_cache_interleave_size = 8 - - block_table = BlockTable( - block_size=block_size, - max_num_reqs=max_num_reqs, - max_num_blocks_per_req=max_num_blocks_per_req, - max_num_batched_tokens=max_num_batched_tokens, - pin_memory=False, - device=torch.device(DEVICE_TYPE), - kernel_block_size=kernel_block_sizes[0], - cp_kv_cache_interleave_size=cp_kv_cache_interleave_size, - ) - - # Verify hybrid block configuration - assert block_table.use_hybrid_blocks is True - assert block_table.block_size == kernel_block_sizes[0] - assert block_table.blocks_per_kv_block == ( - block_size // kernel_block_sizes[0] - ) # Changed to use first element - - # Test block table conversion logic - # One kvcache_manager block should map to multiple kernel blocks - kvcache_manager_blocks = [0, 1, 2] - - # Verify that kvcache_manager blocks can be converted to kernel blocks - # and that block table operations work correctly. - req_index = 0 - block_table.append_row(kvcache_manager_blocks, req_index) - # Get expected kernel blocks from the implementation for verification. - expected_kernel_blocks = block_table.map_to_kernel_blocks( - np.array(kvcache_manager_blocks), - block_table.blocks_per_kv_block, - block_table._kernel_block_arange, - ) - # Verify block table state - assert block_table.num_blocks_per_row[req_index] == len(expected_kernel_blocks) - assert np.array_equal( - block_table.block_table.np[req_index, : len(expected_kernel_blocks)], - expected_kernel_blocks, - ) - - def test_get_block_table_width_aligns_to_128_tokens(): assert get_block_table_width(1875, 64) == 1876 -def test_get_block_table_width_splits_virtual_blocks(): - assert get_block_table_width(235, 256, 64) == 940 - - def test_mamba_state_table_width_is_not_aligned(): block_tables = MultiGroupBlockTable( max_num_reqs=1, @@ -1750,7 +1688,6 @@ def test_mamba_state_table_width_is_not_aligned(): pin_memory=False, device=torch.device("cpu"), block_sizes=[39664], - kernel_block_sizes=[39664], max_num_blocks=[1], slot_mapping_modes=[SlotMappingMode.NONE], ) @@ -1758,42 +1695,6 @@ def test_mamba_state_table_width_is_not_aligned(): assert block_tables[0].max_num_blocks_per_req == 1 -def test_input_batch_with_kernel_block_sizes(): - """Test InputBatch initialization with kernel_block_sizes parameter.""" - max_num_reqs = 10 - max_model_len = 512 - max_num_batched_tokens = 512 - device = torch.device(DEVICE_TYPE) - vocab_size = 50272 - - # Test with different kernel block sizes - block_sizes = [32, 64] - kernel_block_sizes = [16, 32] - - input_batch = InputBatch( - max_num_reqs=max_num_reqs, - max_model_len=max_model_len, - max_num_batched_tokens=max_num_batched_tokens, - device=device, - vocab_size=vocab_size, - block_sizes=block_sizes, - kernel_block_sizes=kernel_block_sizes, - max_num_blocks_per_req=[16, 8], - ) - - # Verify that block tables were created with kernel block sizes - assert len(input_batch.block_table.block_tables) == len(block_sizes) - - for i, (kv_size, kernel_size) in enumerate(zip(block_sizes, kernel_block_sizes)): - block_table = input_batch.block_table.block_tables[i] - if kv_size != kernel_size: - assert block_table.use_hybrid_blocks is True - assert block_table.block_size == kernel_size - else: - assert block_table.use_hybrid_blocks is False - assert block_table.block_size == kernel_size - - def test_hybrid_cache_integration(default_vllm_config, dist_init): """Test hybrid cache architecture integration with GPUModelRunner.""" # Create a new model runner with hybrid cache configuration @@ -1843,9 +1744,8 @@ def test_hybrid_cache_integration(default_vllm_config, dist_init): device=runner.device, vocab_size=runner.model_config.get_vocab_size(), block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size], - kernel_block_sizes=[16], max_num_blocks_per_req=[NUM_BLOCKS], - ) # Use kernel block size + ) runner.initialize_attn_backend(kv_cache_config) diff --git a/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py b/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py index 92207aa38f6b..dd9046de612d 100644 --- a/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py +++ b/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py @@ -119,9 +119,6 @@ def _initialize_kv_cache(runner: GPUModelRunner) -> None: device=runner.device, vocab_size=runner.model_config.get_vocab_size(), block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size], - kernel_block_sizes=[ - kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size - ], max_num_blocks_per_req=[NUM_BLOCKS], ) runner.initialize_attn_backend(kv_cache_config) diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index dec382725bc1..47173da08155 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -237,7 +237,6 @@ def stage_write(self, *_args): block_tables = BlockTables.__new__(BlockTables) block_tables.num_kv_cache_groups = 1 - block_tables.blocks_per_kv_block = [1] block_tables.block_tables = [_BlockTable()] block_tables.num_blocks = SimpleNamespace( np=torch.tensor([[0, 3]], dtype=torch.int32) diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index d4b464e01463..fcf171a89ca3 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -47,8 +47,7 @@ def _make_capture_manager(block_table: torch.Tensor): block_tables = SimpleNamespace( input_block_tables=(block_table,), num_kv_cache_groups=1, - kernel_block_sizes=(2,), - blocks_per_kv_block=(1,), + block_sizes=(2,), ) manager = PCPManager( pcp_world_size=2, diff --git a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py index b63bd7a2ff05..a3ec5f4eb7f4 100644 --- a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py +++ b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py @@ -20,6 +20,7 @@ from vllm.v1.attention.backends import flash_attn as flash_attn_module from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata from vllm.v1.attention.backends.utils import split_decodes_and_prefills +from vllm.v1.kv_cache_interface import MLAAttentionSpec from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor from vllm.v1.worker.gpu.spec_decode import speculator as base_spec_module from vllm.v1.worker.gpu.spec_decode.multi_module_mtp.speculator import ( @@ -30,6 +31,7 @@ from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.speculator import ( TargetDependentARSpeculator, ) +from vllm.v1.worker.utils import AttentionGroup class _TestSpeculator(TargetDependentARSpeculator): @@ -102,6 +104,7 @@ def build(common_prefix_len, common_attn_metadata): get_metadata_builder=lambda _: SimpleNamespace( build=build, supports_update_block_table=False ), + build_metadata=lambda cm, ubatch_idx=0, **_: build(0, cm), layer_names=["draft"], ) ] @@ -427,6 +430,8 @@ def test_multi_step_decode_replays_captured_graph_as_expected( query_start_loc=torch.arange(3), ) speculator.idx_mapping = torch.arange(2) + speculator.attn_groups = [] + speculator.block_tables = SimpleNamespace(input_block_tables=[]) generate_draft = Mock() speculator._generate_draft = generate_draft run_fullgraph = Mock() @@ -450,6 +455,52 @@ def test_multi_step_decode_replays_captured_graph_as_expected( assert run_fullgraph.call_count == expected_graph_replays +def test_fused_full_graph_maps_split_group_block_tables(): + """The fused FULL replay skips build(), so split groups' tables are mapped first.""" + spec = MLAAttentionSpec( + block_size=1024, num_kv_heads=1, head_size=128, dtype=torch.uint8 + ) + group = AttentionGroup( + SimpleNamespace(), + ["layer.0"], + spec, + 0, + kernel_block_size=256, + kernel_block_table=torch.zeros(1, 4, 32, dtype=torch.int32), + kernel_block_offsets=torch.arange(4, dtype=torch.int32), + ) + speculator = object.__new__(_TestSpeculator) + speculator.input_buffers = SimpleNamespace( + positions=torch.arange(2), query_start_loc=torch.arange(3) + ) + speculator.idx_mapping = torch.arange(2) + speculator.attn_groups = [[group]] + speculator.block_tables = SimpleNamespace( + input_block_tables=[torch.tensor([[1, 2], [0, 3]], dtype=torch.int32)] + ) + replayed_tables = [] + speculator.decode_cudagraph_manager = SimpleNamespace( + run_fullgraph=lambda _: replayed_tables.append( + group.kernel_block_table[0, :2, :8].tolist() + ) + ) + + speculator._fused_multi_step_decode( + num_reqs=2, + skip_attn=False, + batch_desc=BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.FULL, num_tokens=2, num_reqs=2 + ), + num_tokens_across_dp=None, + seq_lens_cpu_upper_bound=None, + num_speculative_steps=4, + ) + + assert replayed_tables == [ + [[4, 5, 6, 7, 8, 9, 10, 11], [0, 1, 2, 3, 12, 13, 14, 15]] + ] + + def test_update_draft_decode_metadata_updates_fa3_scheduler_metadata( monkeypatch, ): diff --git a/tests/v1/worker/test_jit_warmup_migration.py b/tests/v1/worker/test_jit_warmup_migration.py index ef887c5351f6..1875e605e53a 100644 --- a/tests/v1/worker/test_jit_warmup_migration.py +++ b/tests/v1/worker/test_jit_warmup_migration.py @@ -12,40 +12,24 @@ from vllm.v1.worker.block_table import ComputeSlotMappingKernel -@pytest.mark.parametrize( - ("kv_cache_block_size", "blocks_per_kv_block", "block_size", "block_size_rep"), - [ - (256, 1, 256, 16), - (256, 4, 64, 16), - (64, 1, 64, 16), - (8, 1, 8, 2), - (4, 1, 4, 2), - ], -) +@pytest.mark.parametrize("kv_cache_block_size", [256, 64, 8, 4]) def test_compute_slot_mapping_warmup_matches_runtime_specializations( kv_cache_block_size: int, - blocks_per_kv_block: int, - block_size: int, - block_size_rep: int, ) -> None: kernel = ComputeSlotMappingKernel() kwargs = dict( kv_cache_block_size=kv_cache_block_size, - blocks_per_kv_block=blocks_per_kv_block, total_cp_world_size=1, total_cp_rank=0, cp_kv_cache_interleave_size=1, block_table_stride=32768, - block_size=block_size, ) expected = kernel.CompileKey( kv_cache_block_size=kv_cache_block_size, - blocks_per_kv_block=blocks_per_kv_block, total_cp_world_size=1, total_cp_rank=0, cp_kv_cache_interleave_size=1, block_table_stride=16, - block_size=block_size_rep, ) assert kernel.dispatch(**kwargs) == expected diff --git a/tests/v1/worker/test_kv_cache_allocation_scope.py b/tests/v1/worker/test_kv_cache_allocation_scope.py index 7f4bea52c227..52cf2a374172 100644 --- a/tests/v1/worker/test_kv_cache_allocation_scope.py +++ b/tests/v1/worker/test_kv_cache_allocation_scope.py @@ -112,6 +112,8 @@ def bind(*args, **kwargs): model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type="test")), compilation_config=SimpleNamespace(static_forward_context={}), kv_caches=[], + attn_groups=[], + speculative_config=None, ) result = gpu_model_runner.GPUModelRunner.initialize_kv_cache_tensors( runner, diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 4407d6fc5430..7e588ccde2e5 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -659,7 +659,6 @@ def test_v2_align_ctx_binds_request_slot_tables_for_state_copies(): max_num_batched_tokens=64, max_num_blocks_per_group=[4], device=device, - kernel_block_sizes=[cfg.block_size], ) state = object.__new__(MambaHybridModelState) diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index f6f7227954bf..e17f214871f3 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -443,6 +443,43 @@ def test_copy_cpu_kv_cache_logical_blocks_ignores_storage_padding(): assert (backing[9] == -1).all() +def test_bind_kv_cache_keeps_scheduler_block_views_for_runner(): + """Block copies index scheduler blocks, so the runner keeps the unmapped + views while a packed layer is bound with kernel blocks.""" + + class _Layer: + def bind_kv_cache(self, kv_cache): + self.kv_cache = kv_cache + + num_blocks, row = 4, 64 + storage = torch.zeros(num_blocks * row) + # Two layers packed into one block row: an 8x4 page, then a 4x2 page + # split into two 2-row kernel blocks. + mla = storage.as_strided((num_blocks, 8, 4), (row, 4, 1)) + indexer = storage.as_strided((num_blocks, 1, 4, 2), (row, 8, 2, 1), 40) + # The indexer in 2-row kernel blocks: block b's kernel block j is b * 16 + j. + mapped = indexer.as_strided((50, 1, 2, 2), (4, 8, 2, 1)) + kv_caches = {"layers.0.attn": mla, "layers.1.indexer": indexer} + ctx = {name: _Layer() for name in kv_caches} + runner_kv_caches: list[torch.Tensor] = [] + + bind_kv_cache( + kv_caches, + ctx, + runner_kv_caches, + layer_kv_caches={**kv_caches, "layers.1.indexer": mapped}, + ) + + assert ctx["layers.1.indexer"].kv_cache is mapped + assert runner_kv_caches == [mla, indexer] + copies = [KVCacheBlockCopy(1, 3)] + with pytest.raises(AssertionError, match="not divisible"): + copy_kv_cache_blocks_inplace([mapped], num_blocks, copies) + storage[row : 2 * row] = 7 + copy_kv_cache_blocks_inplace(runner_kv_caches, num_blocks, copies) + assert (mla[3] == 7).all() and (indexer[3] == 7).all() + + @pytest.mark.parametrize("shared,rank", [(False, 0), (True, 0), (True, 1)]) @pytest.mark.parametrize("has_copies", [False, True]) def test_hisparse_host_copy_waits_for_writes_on_writer_only( diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index d314d6107ad7..d09e5da71b0e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -20,6 +20,7 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backend import AttentionBackend +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput if TYPE_CHECKING: @@ -351,6 +352,35 @@ def get_current_attn_backends( ] +def get_current_attn_backends_and_specs( + vllm_config: VllmConfig, + kv_cache_config: KVCacheConfig, + fallback: list[type[AttentionBackend]], +) -> tuple[list[type[AttentionBackend]], list[AttentionSpec | None]]: + """Distinct (backend, spec) pairs of the transfer layers, else ``fallback``. + + Compressed caches are skipped: their kernel pages never split transferred + blocks. + """ + pairs: dict[tuple[type[AttentionBackend], AttentionSpec | None], None] = {} + layer_type = cast(type[Any], AttentionLayerBase) + for group in kv_cache_config.transfer_groups: + specs = getattr(group.kv_cache_spec, "kv_cache_specs", {}) + layers = get_layers_from_vllm_config(vllm_config, layer_type, group.layer_names) + for name, layer in layers.items(): + spec = specs.get(name, group.kv_cache_spec) + if isinstance(spec, AttentionSpec) and spec.tokens_per_state > 1: + layout = vllm_config.cache_config.get_resolved_kv_cache_layout() + assert layout.is_block_outermost, "Compressed caches must be packed" + continue + if not isinstance(spec, AttentionSpec): + spec = None + pairs[layer.get_attn_backend(), spec] = None + if not pairs: + return fallback, [None] * len(fallback) + return [backend for backend, _ in pairs], [spec for _, spec in pairs] + + def get_current_attn_backend( vllm_config: VllmConfig, layer_names: list[str] | None = None ) -> type[AttentionBackend]: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 9d5194f33e27..41c1337ecd85 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -26,6 +26,7 @@ EngineId, TransferTopology, get_current_attn_backends, + get_current_attn_backends_and_specs, ) from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorBase_V1, @@ -59,8 +60,8 @@ from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import ( AttentionSpec, + CircularBufferSpec, FullAttentionSpec, - KpoolTailSpec, KVCacheSpec, MambaSpec, MLAAttentionSpec, @@ -69,7 +70,6 @@ ) from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.request import RequestStatus -from vllm.v1.worker.block_table import BlockTable from vllm.v1.worker.utils import select_common_block_size logger = init_logger(__name__) @@ -1350,7 +1350,10 @@ def _sync_block_size_with_kernel(self) -> None: # physical block sizes. Pick the common (smallest) block size so that # KV-cache registration and transfer work correctly for both models. backends = get_current_attn_backends(self.vllm_config) - kernel_block_size = select_common_block_size(self.block_size, backends) + backends, specs = get_current_attn_backends_and_specs( + self.vllm_config, self.kv_cache_config, backends + ) + kernel_block_size = select_common_block_size(self.block_size, backends, specs) if self.block_size != kernel_block_size: logger.info_once( "User-specified logical block size (%s) does not match" @@ -1727,11 +1730,12 @@ def _logical_to_kernel_block_ids( ) group_specs = self.kv_cache_config.transfer_groups return [ - BlockTable.map_to_kernel_blocks( - np.array(group), - self._physical_blocks_per_logical_kv_block, - block_arange, - ).tolist() + ( + np.array(group)[:, None] * self._physical_blocks_per_logical_kv_block + + block_arange + ) + .reshape(-1) + .tolist() if not isinstance(group_specs[i].kv_cache_spec, MambaSpec) else group for i, group in enumerate(block_ids) @@ -2054,7 +2058,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): # hetero PP can align by name (#56033). A non-contiguous row cannot # name a sub-span, so it is one region anchored at the allocation. use_packed_row = ( - not isinstance(layer_spec, (MambaSpec, KpoolTailSpec)) + not isinstance(layer_spec, (MambaSpec, CircularBufferSpec)) and storage_is_block_major and block_stride > 0 and (is_mla_region or not hnc_contiguous) @@ -2121,9 +2125,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): block_len = region_cache.stride(0) * region_cache.element_size() region_base_addresses.append(base_addr) - if isinstance(layer_spec, KpoolTailSpec): - kv_block_len = layer_spec.unpadded_page_size_bytes // 2 - elif isinstance(layer_spec, AttentionSpec) and block_is_contiguous: + if isinstance(layer_spec, AttentionSpec) and block_is_contiguous: assert ( layer_spec.page_size_bytes % self._physical_blocks_per_logical_kv_block @@ -2133,6 +2135,8 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): layer_spec.page_size_bytes // self._physical_blocks_per_logical_kv_block ) + elif isinstance(layer_spec, MambaSpec) and block_is_contiguous: + kv_block_len = layer_spec.page_size_bytes else: kv_block_len = block_len if kv_block_len > block_len: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 38a656954eee..778eb0f7560d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -28,6 +28,7 @@ EngineTransferInfo, TransferTopology, get_current_attn_backends, + get_current_attn_backends_and_specs, kv_postprocess_blksize_and_layout_on_receive, kv_postprocess_blksize_on_receive, ) @@ -79,7 +80,6 @@ from vllm.v1.kv_cache_interface import ( CircularBufferSpec, FullAttentionSpec, - KpoolTailSpec, KVCacheLayout, KVCacheSpec, MambaSpec, @@ -88,7 +88,6 @@ UniformTypeKVCacheSpecs, iter_layer_specs, ) -from vllm.v1.worker.block_table import BlockTable from vllm.v1.worker.utils import select_common_block_size if TYPE_CHECKING: @@ -120,46 +119,6 @@ def _share_storage_and_block_stride(caches: list[torch.Tensor]) -> bool: return len(block_strides) == len(storage_ptrs) == 1 -def _tensor_byte_span_end(cache: torch.Tensor) -> int: - """Return the exclusive end address touched by a nonnegative-stride view.""" - if cache.numel() == 0: - return cache.data_ptr() - if any(stride < 0 for stride in cache.stride()): - raise ValueError("NIXL cache views must have nonnegative strides") - max_element_offset = sum( - (size - 1) * stride for size, stride in zip(cache.shape, cache.stride()) - ) - return cache.data_ptr() + (max_element_offset + 1) * cache.element_size() - - -def _uses_dense_virtual_transfer_pages( - layer_spec: KVCacheSpec, - cache: torch.Tensor, - physical_page_size: int, - num_blocks: int, -) -> bool: - """Return whether a compressed kernel view can be split into NIXL pages.""" - if not ( - isinstance(layer_spec, MLAAttentionSpec) - and layer_spec.tokens_per_state > 1 - and cache.ndim == 4 - and cache.shape[1] == 1 - and cache.is_contiguous() - and physical_page_size > 0 - and layer_spec.state_content_size_bytes > 0 - ): - return False - - block_stride = cache.stride(0) * cache.element_size() - return ( - block_stride > physical_page_size - and block_stride % physical_page_size == 0 - and physical_page_size % layer_spec.state_content_size_bytes == 0 - and cache.shape[0] * (block_stride // physical_page_size) == num_blocks - and cache.nbytes == num_blocks * physical_page_size - ) - - class NixlBaseConnectorWorker: """Base implementation of Worker side methods shared by pull and push.""" @@ -985,8 +944,10 @@ def _validate_remote_parallel_config( ) def _sync_block_size_with_kernel(self) -> None: - backends = get_current_attn_backends(self.vllm_config) - kernel_block_size = select_common_block_size(self.block_size, backends) + backends, specs = get_current_attn_backends_and_specs( + self.vllm_config, self.kv_cache_config, self.attn_backends + ) + kernel_block_size = select_common_block_size(self.block_size, backends, specs) # Number of blocks not accounting for kernel block mismatches self._logical_num_blocks = self.num_blocks if self.block_size != kernel_block_size: @@ -1425,7 +1386,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): if self._supports_pp_hma and self._has_packed_cache: # PP can reorder the same backends across stages of a packed model. backend_name = ",".join( - sorted(backend.get_name() for backend in self.attn_backends) + sorted({backend.get_name() for backend in self.attn_backends}) ) self.compat_hash = compute_nixl_compatibility_hash( self.vllm_config, @@ -1478,31 +1439,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ): raise NotImplementedError("PP push with packed KV caches requires MLA") - layer_specs: dict[str, KVCacheSpec] = {} - compressed_region_owners: dict[int, torch.Tensor] = {} - for layer_name, cache in xfer_buffers.items(): - layer_spec = self._layer_specs.get(layer_name) - if isinstance(layer_spec, UniformTypeKVCacheSpecs): - layer_spec = layer_spec.kv_cache_specs[layer_name] - if layer_spec is None: - continue - layer_specs[layer_name] = layer_spec - physical_page_size = ( - layer_spec.page_size_bytes - if isinstance(layer_spec, MambaSpec) - else layer_spec.page_size_bytes - // self._physical_blocks_per_logical_kv_block - ) - num_blocks = ( - self._logical_num_blocks - if isinstance(layer_spec, MambaSpec) - else self.num_blocks - ) - if _uses_dense_virtual_transfer_pages( - layer_spec, cache, physical_page_size, num_blocks - ): - compressed_region_owners.setdefault(cache.data_ptr(), cache) - track_region_layers = self._tracks_region_layers() region_layers: list[list[str]] = [] packed_member_layouts: dict[str, tuple[int, int]] = {} @@ -1527,7 +1463,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): # kernel requires a specific block size. This leads to SSM and FA layers # having different num_blocks. # `_physical_blocks_per_logical_kv_block` ratio is used to adjust for this. - layer_spec = layer_specs.get(layer_name) + layer_spec = self._layer_specs.get(layer_name) if layer_spec is None: logger.debug( "Skipping layer %s as no KVCache spec is present. " @@ -1587,32 +1523,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): region_device_id, ) - if isinstance(layer_spec, KpoolTailSpec): - compressed_owner = compressed_region_owners.get(cache.data_ptr()) - if compressed_owner is not None: - owner_storage = compressed_owner.untyped_storage() - owner_end = compressed_owner.data_ptr() + compressed_owner.nbytes - tail_is_covered = ( - compressed_owner.is_contiguous() - and owner_storage.data_ptr() == storage_addr - and _tensor_byte_span_end(cache) <= owner_end - ) - if not tail_is_covered: - raise AssertionError( - "Kpool tail cache is not fully covered by its compressed " - f"indexer region: layer={layer_name}, " - f"tail_shape={tuple(cache.shape)}, " - f"tail_stride={tuple(cache.stride())}, " - f"owner_shape={tuple(compressed_owner.shape)}, " - f"owner_stride={tuple(compressed_owner.stride())}" - ) - logger.debug( - "Skipping layer %s because its compressed indexer region " - "covers the same storage", - layer_name, - ) - continue - if isinstance(layer_spec, MambaSpec): physical_ratio = self._physical_blocks_per_logical_kv_block block_len = physical_page_size // physical_ratio @@ -1647,9 +1557,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): and cache[0].is_contiguous() and cache[0].nbytes == physical_page_size ) - virtual_transfer_pages = _uses_dense_virtual_transfer_pages( - layer_spec, cache, physical_page_size, num_blocks - ) # Push workers address the pages of packed MLA rows by layer. packed_mla_push = ( track_region_layers @@ -1658,12 +1565,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): and storage_is_block_major and is_mla_region ) - if virtual_transfer_pages: - # A compressed kernel row can contain multiple NIXL transfer pages. - region_specs = [ - (cache.data_ptr(), physical_page_size, physical_page_size) - ] - elif packed_mla_push and use_layer_name_routing: + if packed_mla_push and use_layer_name_routing: # PP>1 prefill stage: register only this layer's page of each row. region_specs = [ (cache.data_ptr(), physical_page_size, block_stride) @@ -1671,7 +1573,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): elif packed_mla_push: # PP=1 decode or prefill: register whole rows, and advertise where # this layer's page sits so a PP>1 prefill can write into it. - storage_block_len = registration_len // num_blocks + storage_block_len = block_stride region_specs = [ (registration_base, storage_block_len, storage_block_len) ] @@ -3388,11 +3290,9 @@ def _logical_to_kernel_block_ids(self, block_ids: BlockIds, ratio: int) -> Block physical_block_ids.append(group) else: physical_block_ids.append( - BlockTable.map_to_kernel_blocks( - np.array(group), - ratio, - block_arange, - ).tolist() + (np.array(group)[:, None] * ratio + block_arange) + .reshape(-1) + .tolist() ) return physical_block_ids diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 2a3048096c98..861c894c383e 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1376,14 +1376,6 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: is_index_group_leader=self.indexer is not None, non_causal_multi_token_decode=self.non_causal_multi_token_decode, ) - # SM100 FlashMLA paged kernels also express TMA coordinates in token rows. - uses_tma_rows = ( - self.attn_backend.get_name() == "FLASHMLA_SPARSE" - and self.kv_cache_dtype in ("fp8_ds_mla", "nvfp4_ds_mla") - and current_platform.is_device_capability_family(100) - ) - if self._uses_flat_kv_cache() or uses_tma_rows: - spec = replace(spec, block_stride_alignment=spec.state_content_size_bytes) return spec def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor): diff --git a/vllm/models/deepseek_v41/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py index cc3ade80b8ad..7fc0bd25ff75 100644 --- a/vllm/models/deepseek_v41/nvidia/model_state.py +++ b/vllm/models/deepseek_v41/nvidia/model_state.py @@ -669,7 +669,7 @@ def _copy_attn_groups( copy.create_metadata_builders( self.vllm_config, self.device, - group.metadata_builders[0].kernel_block_size, + group.kernel_block_size, ) copies.append(copy) own_groups.append(copies) diff --git a/vllm/models/glm5next/common/attention.py b/vllm/models/glm5next/common/attention.py index a264717f0d78..b398a614c070 100644 --- a/vllm/models/glm5next/common/attention.py +++ b/vllm/models/glm5next/common/attention.py @@ -33,9 +33,9 @@ from vllm.models.glm5next.nvidia.ops.kpool_compress import fwht128_quant_fp8 from vllm.models.glm5next.sparse_indexer import SparseAttnIndexerKpool from vllm.platforms import current_platform -from vllm.utils.deep_gemm import PAGED_MQA_PAGE_SIZES from vllm.utils.math_utils import cdiv, next_power_of_2 -from vllm.v1.kv_cache_interface import KpoolTailSpec, MLAAttentionSpec +from vllm.v1.attention.backends.mla.indexer import Glm5NextIndexerBackend +from vllm.v1.kv_cache_interface import CircularBufferSpec, MLAAttentionSpec logger = init_logger(__name__) @@ -89,15 +89,9 @@ class Glm5NextIndexerCache(DeepseekV32IndexerCache): The indexer shares one block with the co-located MLA (a single ``MLAAttentionSpec`` / block_table), so ``block_size`` is the model-wide - ``cache_config.block_size``. DeepGEMM's paged-MQA kernel - (``csrc/apis/attention.hpp``) requires ``block_kv`` to be exactly 32 or - 64, so the storage block is virtually split into pool pages of the - largest such size that tiles it (``storage_kernel_block_size``); this - needs ``block_size`` to be a multiple of ``index_kpool * 32`` (512 for - ``index_kpool = 16``). A smaller block (e.g. the default 64) silently - collapses ``storage_block_size`` (64 // 16 = 4) and only fails later at - the opaque C++ assert; ``get_kv_cache_spec`` guards this up front - instead. + ``cache_config.block_size``. DeepGEMM's paged-MQA kernel takes 32- or + 64-state pages (only 64 on SM120), which ``Glm5NextIndexerBackend`` + declares as its kernel block sizes. """ def __init__( @@ -113,12 +107,6 @@ def __init__( head_dim=head_dim, dtype=dtype, prefix=prefix, cache_config=cache_config ) assert index_kpool > 1, "Glm5NextIndexerCache expects index_kpool > 1" - # Keep chunked-prefill boundaries aligned to complete pools. - assert cache_config.block_size % index_kpool == 0, ( - "Glm5NextIndexerCache: cache_config.block_size " - f"({cache_config.block_size}) must be a multiple of index_kpool " - f"({index_kpool}) so chunked-prefill boundaries stay pool-aligned." - ) self._index_kpool = index_kpool def get_kv_cache_spec(self, vllm_config: VllmConfig): @@ -128,33 +116,10 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig): # ``tokens_per_state`` is the KV-spec representation of kpool # compression in the current cache-layout API. assert isinstance(spec, MLAAttentionSpec) - spec = replace(spec, tokens_per_state=self._index_kpool) - - # DeepGEMM paged-MQA takes block_kv in {32, 64}; the storage block - # (= block_size // index_kpool) is virtually split into pool pages of - # the largest such size that tiles it, so it must be a multiple of 32. - storage_block_size = spec.block_size // self._index_kpool - assert ( - spec.block_size % self._index_kpool == 0 and storage_block_size % 32 == 0 - ), ( - "Glm5NextIndexerCache: kpool indexer requires cache block_size to " - f"be a multiple of index_kpool * 32 ({self._index_kpool * 32}) so " - "that DeepGEMM paged-MQA pool pages (32 or 64 entries) tile the " - f"storage block, got block_size={spec.block_size} -> " - f"storage_block_size={storage_block_size}." - ) - max_page_size = max(PAGED_MQA_PAGE_SIZES) - min_page_size = min(PAGED_MQA_PAGE_SIZES) - if storage_block_size <= max_page_size: - page_size = storage_block_size - elif storage_block_size % max_page_size == 0: - page_size = max_page_size - else: - page_size = min_page_size - return replace( - spec, - storage_block_size=page_size * self._index_kpool, - ) + return replace(spec, tokens_per_state=self._index_kpool) + + def get_attn_backend(self): + return Glm5NextIndexerBackend class Glm5NextTailCache(DeepseekV32IndexerCache): @@ -164,9 +129,9 @@ class Glm5NextTailCache(DeepseekV32IndexerCache): ``ring`` slots per request, overwritten in place by ``pos % ring`` as decode/spec-decode advances. Prefill seeds it (instead of discarding the tail raw K+gate); the connector transfers it across PD; decode reads it to - compress the boundary pool correctly. ``KpoolTailSpec`` / - ``KpoolTailManager`` provide the no-prune, 1-block/req allocation that lets - the in-progress pool survive across steps and across transfer. + compress the boundary pool correctly. ``CircularBufferSpec`` provides the + no-prune, 1-block/req allocation that lets the in-progress pool survive + across steps and across transfer. Stores raw bf16 K (``head_dim``) as the "K" half of each block and the bf16 gate score (``head_dim``) as the "V" half -- not the fp8-compressed @@ -202,13 +167,12 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig): f"({self.cache_config.block_size}) must be a multiple of the " f"tail ring ({ring})" ) - return KpoolTailSpec( + return CircularBufferSpec( block_size=ring, num_kv_heads=2, head_size=self.head_dim, head_size_v=0, dtype=torch.bfloat16, - sliding_window=ring, # Tiny per-request ring; every DCP rank keeps the full tail so the # rank owning a pool can compress it. dcp_sharded=False, @@ -303,7 +267,7 @@ def __init__( # Paged tail cache (in-progress pool's raw K + gate score). Written by # prefill (seeds the boundary pool) and decode (per-step stash); read by # the decode kernel to compress the boundary pool. Transferred across PD - # so the decode side sees the prefill tail. See KpoolTailSpec/Manager. + # so the decode side sees the prefill tail. See CircularBufferManager. self.tail_cache = Glm5NextTailCache( head_dim=self.head_dim, dtype=torch.bfloat16, diff --git a/vllm/models/glm5next/nvidia/ops/kpool_compress.py b/vllm/models/glm5next/nvidia/ops/kpool_compress.py index db496d300b46..eb289c65cf68 100644 --- a/vllm/models/glm5next/nvidia/ops/kpool_compress.py +++ b/vllm/models/glm5next/nvidia/ops/kpool_compress.py @@ -396,11 +396,6 @@ def _kpool_tail_seed_kernel( ahead belongs to a different tail block (or is past the batch / padding, slot < 0). ``tslot = block * RING + pos % RING``; the destination is ``tail[block, {0:K, 1:score}, pos % RING, :]``. - - The tail cache aliases the indexer cache with the indexer's (padded) block - stride, so blocks are addressed through ``TAIL_BLOCK_ELEMS`` / - ``KPOOL_HEAD`` (``tail.stride(0)`` / ``tail.stride(1)``), never as a dense - ``[num_blocks, 2, KPOOL, HEAD_DIM]`` array. """ i = tl.program_id(0) t = tl.load(tslot_ptr + i).to(tl.int64) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 68c3aa7851cd..401a98f66feb 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -414,26 +414,6 @@ def get_valid_backends( return valid_backends_priorities, invalid_reasons - @classmethod - def _get_indexer_block_alignment(cls, vllm_config: VllmConfig) -> int | None: - index_kpool = getattr( - vllm_config.model_config.hf_text_config, "index_kpool", None - ) - if not index_kpool or index_kpool <= 1: - return None - from vllm.utils.deep_gemm import PAGED_MQA_PAGE_SIZES - - # kpool paged-MQA indexer: the storage block (block_size / - # index_kpool) is virtually split into pool pages, so block_size - # must be a multiple of index_kpool times a legal pool page. - page = min(PAGED_MQA_PAGE_SIZES) - if cls.is_device_capability_family(120): - # On sm120 the DeepGEMM paged-MQA kernel only accepts block_kv - # 64 for the fp8 indexer cache, so align to the largest pool - # page here to make the page split land on 64 not the min 32. - page = max(PAGED_MQA_PAGE_SIZES) - return index_kpool * page - @classmethod def get_attn_backend_cls( cls, diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index dc3200759c82..3df44a927531 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -737,14 +737,14 @@ def _preferred_block_size_for_backends( for candidate in candidates: if all(b.supports_block_size(candidate) for b in backend_classes): return candidate - raise ValueError( - "The attention backends share no supported KV cache block size (" - + "; ".join( - f"{b.get_name()}: {b.get_supported_kernel_block_sizes()}" - for b in backend_classes + raise ValueError( + "The attention backends share no supported KV cache block size (" + + "; ".join( + f"{b.get_name()}: {b.get_supported_kernel_block_sizes()}" + for b in backend_classes + ) + + ")." ) - + ")." - ) @classmethod def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: @@ -935,17 +935,6 @@ def per_token_page_bytes(dtype: "torch.dtype", cache_dtype: str) -> int: if cache_config.mamba_page_size_padded is not None: cache_config.mamba_page_size_padded = shared_page - @classmethod - def _get_indexer_block_alignment(cls, vllm_config: "VllmConfig") -> int | None: - """Extra ``block_size`` multiple a sparse indexer needs, else ``None``. - - The CUDA kpool paged-MQA indexer virtually splits each storage block - into pool pages, so ``block_size`` must be a multiple of - ``index_kpool * min(PAGED_MQA_PAGE_SIZES)`` — implemented in the CUDA - platform override. Other platforms impose no extra constraint. - """ - return None - @classmethod def _align_hybrid_block_size( cls, @@ -1077,6 +1066,14 @@ def _align_hybrid_block_size( ), cache_config.block_size, ) + # Also a multiple of every backend's smallest kernel block. + backend_block_alignment = lcm( + *( + min(s.base if isinstance(s, MultipleOf) else s for s in sizes) + for b in cls._find_non_ssm_backends(vllm_config) + if (sizes := b.get_supported_kernel_block_sizes()) + ) + ) if model_config.use_mla: # TRTLLM/FlashInfer MLA decode kernels require the physical # number of kernel blocks to be aligned to 128 / kernel_block_size. @@ -1090,9 +1087,9 @@ def _align_hybrid_block_size( mamba_page_size, kernel_block_alignment_size * attn_page_size_1_token, ) - indexer_align = cls._get_indexer_block_alignment(vllm_config) - if indexer_align: - attn_block_size = indexer_align * cdiv(attn_block_size, indexer_align) + attn_block_size = backend_block_alignment * cdiv( + attn_block_size, backend_block_alignment + ) if cache_config.block_size < attn_block_size: cache_config.block_size = attn_block_size diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 0eb979ce05a5..04c4508f0ba9 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -29,10 +29,11 @@ "qwen3_5_moe_text", } -# KV page sizes (in cache entries) supported by the paged-MQA logits kernels -# (fp8_fp4_paged_mqa_logits / get_paged_mqa_logits_metadata). Larger storage -# blocks must be virtually split into one of these page sizes. -PAGED_MQA_PAGE_SIZES = (32, 64) + +def get_paged_mqa_page_sizes() -> tuple[int, ...]: + """KV page sizes (in cache entries) the paged-MQA logits kernels take (only + 64 on SM120); larger storage blocks are split into one of these.""" + return (64,) if current_platform.is_device_capability_family(120) else (32, 64) def should_auto_disable_deep_gemm(model_type: str | None) -> bool: diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 3387a31794c9..68bb0b3bdd63 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -48,12 +48,10 @@ class AttentionType(str, Enum): """Attention between dec. Q and enc. K/V for encoder-decoder.""" +@dataclass(frozen=True) class MultipleOf: base: int - def __init__(self, base: int): - self.base = base - def __repr__(self) -> str: return f"MultipleOf({self.base})" @@ -622,10 +620,6 @@ def __init__( self.layer_names = layer_names self.vllm_config = vllm_config self.device = device - self.kernel_block_size: int | None = None - - def set_kernel_block_size(self, kernel_block_size: int) -> None: - self.kernel_block_size = kernel_block_size @classmethod def get_cudagraph_support( diff --git a/vllm/v1/attention/backends/composite.py b/vllm/v1/attention/backends/composite.py index 0ba40e99f20a..f77e5e6d9561 100644 --- a/vllm/v1/attention/backends/composite.py +++ b/vllm/v1/attention/backends/composite.py @@ -297,11 +297,6 @@ def __init__(self, kv_cache_spec, layer_names, vllm_config, device, **kwargs): ) self.routing = routing_policy(layer_names, vllm_config) - def set_kernel_block_size(self, kernel_block_size): - super().set_kernel_block_size(kernel_block_size) - for builder in self._builders: - builder.set_kernel_block_size(kernel_block_size) - @classmethod def get_cudagraph_support(cls, vllm_config, kv_cache_spec): return min( diff --git a/vllm/v1/attention/backends/mla/compressor_utils.py b/vllm/v1/attention/backends/mla/compressor_utils.py index 41d973f2963b..bf939dc04990 100644 --- a/vllm/v1/attention/backends/mla/compressor_utils.py +++ b/vllm/v1/attention/backends/mla/compressor_utils.py @@ -95,7 +95,9 @@ def kernel( block_table_ptr + batch_idx * block_table_stride + block_ids, mask=mask & is_valid, ) - slot_ids = block_numbers * block_size + pos_after_compress % block_size + slot_ids = block_numbers.to(tl.int64) * block_size + ( + pos_after_compress % block_size + ) # NOTE slot_ids = tl.where(is_valid, slot_ids, PAD_ID) diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index aabd68224093..0016eafedc8a 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py @@ -24,6 +24,7 @@ from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( + align_blocks_to_rows, flat_kv_row_view, triton_convert_req_index_to_global_index, ) @@ -42,7 +43,11 @@ class FlashAttnMLASparseBackend(AttentionBackend): @staticmethod def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: - return [64] + return [MultipleOf(64)] + + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + return align_blocks_to_rows(spec) @staticmethod def get_name() -> str: diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 4ebda50178cb..21eb66aae490 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """FlashInfer sparse MLA attention backend.""" +import math from dataclasses import dataclass from typing import TYPE_CHECKING, ClassVar @@ -30,6 +31,7 @@ ) from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( + align_blocks_to_rows, flat_kv_row_view, prepare_sparse_mla_safe_lengths, triton_convert_req_index_to_global_index, @@ -69,6 +71,13 @@ def is_mla(cls) -> bool: def is_sparse(cls) -> bool: return True + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + # The kernels read the rows as pages of their kernel block size. + (page,) = cls.get_supported_kernel_block_sizes(spec) + assert isinstance(page, MultipleOf) + return align_blocks_to_rows(spec, page.base) + class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase): """FlashInfer sparse MLA backend using the TRTLLM-gen launcher.""" @@ -84,7 +93,7 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase): @staticmethod def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: - return [32, 64] + return [MultipleOf(32)] @staticmethod def get_impl_cls() -> type[MLAAttentionImpl]: @@ -173,7 +182,7 @@ def get_name() -> str: @staticmethod def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: - return [64, 256] + return [MultipleOf(64)] @staticmethod def get_impl_cls() -> type[MLAAttentionImpl]: @@ -575,7 +584,7 @@ def forward_mqa( assert prefill_lse is not None return output, torch.cat((decode_lse, prefill_lse)) - _, block_stride_rows = flat_kv_row_view( + kv_rows, block_stride_rows = flat_kv_row_view( kv_c_and_k_pe_cache, attn_metadata.block_size ) @@ -602,7 +611,12 @@ def forward_mqa( return self._run_mqa_kernel( q, - kv_c_and_k_pe_cache, + # Block sizes and strides have been aligned to 32 rows. + kv_rows.view( + -1, + math.gcd(block_stride_rows, attn_metadata.block_size, 64), + kv_rows.shape[-1], + ), topk_indices_physical, seq_lens, ) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py index 2c918c0e4acd..2bfcee05c400 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py @@ -16,6 +16,7 @@ ) from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( + flat_kv_row_view, triton_convert_req_index_to_global_index, ) @@ -184,6 +185,9 @@ def forward_mqa( output = torch.cat(outputs) if len(outputs) > 1 else outputs[0] return output, None + kv_rows, block_stride_rows = flat_kv_row_view( + kv_c_and_k_pe_cache, attn_metadata.block_size + ) topk_indices_physical = cast( torch.Tensor, triton_convert_req_index_to_global_index( @@ -191,13 +195,15 @@ def forward_mqa( attn_metadata.block_table, topk_indices, BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, NUM_TOPK_TOKENS=topk_indices.shape[1], ), ) return ( self._run_mqa_kernel( q, - kv_c_and_k_pe_cache, + # The SM120 kernels are built for 64-row pages only. + kv_rows.view(-1, 64, kv_rows.shape[-1]), topk_indices_physical, ), None, diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py index 16d7c12e6720..9e4e85761a67 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py @@ -26,7 +26,7 @@ reserved buffers by kernels inside the captured forward, and captured runs read the refreshed plan buffers on replay. -KV cache format: plain contiguous E4M3 ``[num_blocks, block_size, 512]`` +KV cache format: E4M3 ``[num_blocks, block_size, 512]`` (any block stride) (uint8 storage) with a per-tensor ``k_scale``; BF16 caches also work. The per-token x 128-channel-group ``ckv_scale_arr`` layout is supported by the kernel but not wired yet (it needs a group-quantizing cache-write op). @@ -57,9 +57,11 @@ FlashInferMLASparseMetadataBuilder, ) from vllm.v1.attention.backends.mla.sparse_utils import ( + align_blocks_to_rows, + flat_kv_row_view, triton_convert_req_index_to_global_index, ) -from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout +from vllm.v1.kv_cache_interface import AttentionSpec _FP8_KV_DTYPES = ("fp8", "fp8_e4m3") _WORKSPACE_BYTES = 128 * 1024 * 1024 @@ -99,7 +101,12 @@ def _pack_topk_indices( end = tl.load(indptr + row + 1) mask = (cols < end - start) & (cols < WIDTH) values = tl.load(slots + row * WIDTH + cols, mask=mask, other=0) - tl.store(indices + start + cols, tl.maximum(values, 0), mask=mask) + first_slot = tl.load(slots + row * WIDTH) + tl.store( + indices + start + cols, + tl.where(values >= 0, values, tl.maximum(first_slot, 0)), + mask=mask, + ) @triton.jit @@ -149,6 +156,10 @@ class FlashInferMLASparseSM90Backend(AttentionBackend): def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: return [MultipleOf(64)] + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + return align_blocks_to_rows(spec) + @staticmethod def get_name() -> str: return "FLASHINFER_MLA_SPARSE_SM90" @@ -222,10 +233,6 @@ def get_kv_cache_shape( ) -> tuple[int, ...]: return (num_blocks, block_size, head_size) - @classmethod - def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: - return (KVCacheLayout.LBHNC,) - class _SM90State: """Builder-owned wrapper, capture-stable buffers, and plan parameters. @@ -602,6 +609,9 @@ def forward_mqa( assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_tokens] + kv_c_and_k_pe_cache, block_stride_rows = flat_kv_row_view( + kv_c_and_k_pe_cache, attn_metadata.block_size + ) # return_valid_counts=True keeps the compacted-prefix layout: valid # entries at [0, valid_count), -1 past it — exactly the prefix the # planned per-row lengths address. @@ -610,6 +620,7 @@ def forward_mqa( attn_metadata.block_table, topk_indices, BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, NUM_TOPK_TOKENS=topk_indices.shape[1], return_valid_counts=True, ) diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index d441934c022a..0ff34f8ac5e0 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import TYPE_CHECKING, ClassVar import numpy as np @@ -36,6 +36,7 @@ ) from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( + align_blocks_to_rows, flat_kv_row_view, request_row_bounds, triton_convert_req_index_to_global_index, @@ -53,7 +54,7 @@ flash_mla_with_kvcache, get_mla_metadata, ) -from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.kv_cache_interface import AttentionSpec, KVQuantMode from vllm.v1.worker.workspace import current_workspace_manager if TYPE_CHECKING: @@ -132,6 +133,17 @@ FP8_DS_MLA_ROPE_DIM = 64 +def _fp8_kv_pages( + kv_cache: torch.Tensor, block_size: int +) -> tuple[torch.Tensor, int | None]: + """The cache as the fp8 kernel's 64-token pages, and the block stride in + rows when a block spans several pages.""" + if block_size == 64: + return kv_cache, None + rows, block_stride_rows = flat_kv_row_view(kv_cache, block_size) + return rows.view(-1, 64, rows.shape[-1]), block_stride_rows + + class FlashMLASparseBackend(AttentionBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ @@ -144,7 +156,23 @@ class FlashMLASparseBackend(AttentionBackend): @staticmethod def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: - return [64] + if ( + isinstance(kv_cache_spec, AttentionSpec) + and kv_cache_spec.kv_quant_mode != KVQuantMode.NONE + and current_platform.is_device_capability_family(100) + ): + return [64] + return [MultipleOf(64)] + + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + # Off SM100 the quantized kernel reads 64-token pages, only needed + # when blocks span several. + if spec.kv_quant_mode != KVQuantMode.NONE and ( + not current_platform.is_device_capability_family(100) + ): + return replace(spec, block_stride_alignment=MultipleOf(64)) + return align_blocks_to_rows(spec) @staticmethod def get_name() -> str: @@ -998,13 +1026,16 @@ def _forward_fp8_kv_separate_prefill_decode( # and the prefill rows per chunk, against the gathered workspace. lse: torch.Tensor | None = None topk_length: torch.Tensor | None = None + kv_pages, block_stride_rows = _fp8_kv_pages( + kv_c_and_k_pe_cache, attn_metadata.block_size + ) if self.pcp_dcp_kv_gather: pass elif num_prefill_tokens == 0 and not uses_host_cache: topk_indices, topk_length = self._convert_logical_to_physical_topk( topk_indices, attn_metadata, - block_stride_rows=None, + block_stride_rows=block_stride_rows, return_valid_counts=True, ) elif num_prefill_tokens > 0: @@ -1013,6 +1044,7 @@ def _forward_fp8_kv_separate_prefill_decode( attn_metadata.block_table, topk_indices, BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, NUM_TOPK_TOKENS=topk_indices.shape[1], HAS_PREFILL_WORKSPACE=has_prefill_workspace, prefill_workspace_request_ids=prefill_request_ids, @@ -1056,7 +1088,7 @@ def _fp8_decode( topk_indices = topk_indices.view(num_decodes, seq_len, -1) attn_out, _ = self._fp8_flash_mla_kernel( q=q, - kv_c_and_k_pe_cache=kv_c_and_k_pe_cache, + kv_c_and_k_pe_cache=kv_pages, topk_indices=topk_indices, kernel_metadata=fp8_metadata.decode.kernel_metadata, ) @@ -1173,6 +1205,9 @@ def _forward_fp8_kv_mixed_batch( # req_id_per_token covers the whole batch; slice it to the MQA tokens # (q may exclude prefill tokens routed to dense MHA). req_id_per_token = attn_metadata.req_id_per_token[: topk_indices.shape[0]] + kv_pages, block_stride_rows = _fp8_kv_pages( + kv_c_and_k_pe_cache, attn_metadata.block_size + ) if self.dcp_world_size > 1: # The indexer emits global token ids; keep this rank's shard and # convert to local slots. compact_valid_to_front=False keeps the @@ -1187,6 +1222,7 @@ def _forward_fp8_kv_mixed_batch( dcp_rank=self.dcp_rank, cp_kv_cache_interleave_size=attn_metadata.cp_kv_cache_interleave_size, BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, NUM_TOPK_TOKENS=topk_indices.shape[1], compact_valid_to_front=False, ) @@ -1202,7 +1238,7 @@ def _forward_fp8_kv_mixed_batch( topk_indices = self._convert_logical_to_physical_topk( topk_indices, attn_metadata, - block_stride_rows=None, + block_stride_rows=block_stride_rows, return_valid_counts=False, ) else: @@ -1211,12 +1247,13 @@ def _forward_fp8_kv_mixed_batch( block_table, topk_indices, BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, NUM_TOPK_TOKENS=topk_indices.shape[1], ) _attn_out, _lse = self._fp8_flash_mla_kernel( q=q.unsqueeze(0), # unsqueeze to add batch_dim: (T, H, D) -> (1, T, H, D) - kv_c_and_k_pe_cache=kv_c_and_k_pe_cache, + kv_c_and_k_pe_cache=kv_pages, topk_indices=topk_indices.unsqueeze(0), # (T, topk) -> (1, T, topk) kernel_metadata=fp8_metadata, ) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 1005b8e4246a..e776057b05b7 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -7,7 +7,7 @@ import torch import vllm.envs as envs -from vllm.config import VllmConfig +from vllm.config import VllmConfig, get_current_vllm_config from vllm.distributed import get_dcp_group, get_pcp_group from vllm.logger import init_logger from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs @@ -22,6 +22,7 @@ from vllm.triton_utils import tl, triton from vllm.utils.deep_gemm import ( get_paged_mqa_logits_metadata, + get_paged_mqa_page_sizes, has_deep_gemm, native_next_n_supported, ) @@ -234,12 +235,25 @@ def get_builder_cls() -> type["DeepseekV32IndexerMetadataBuilder"]: return DeepseekV32IndexerMetadataBuilder +class Glm5NextIndexerBackend(DeepseekV32IndexerBackend): + @staticmethod + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: + if kv_cache_spec is not None: + # Worker get_kv_cache_spec: runs without the current config. + index_kpool = int(kv_cache_spec.tokens_per_state) + else: + # Platform block-size selection: no spec yet, under the current config. + model_config = get_current_vllm_config().model_config + index_kpool = model_config.hf_text_config.index_kpool + return [index_kpool * n for n in get_paged_mqa_page_sizes()] + + class KpoolTailBackend(DeepseekV32IndexerBackend): """Storage-only backend for the GLM-5.3-Flash kpool tail cache.""" @classmethod def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: - return (KVCacheLayout.LBHNC,) + return (KVCacheLayout.BLHNC,) @staticmethod def get_name() -> str: @@ -1065,8 +1079,6 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: dtype=torch.int32, device=self.device, ) - self.indexer_decode_block_table_buffer: torch.Tensor | None = None - self._max_num_batched_tokens = scheduler_config.max_num_batched_tokens # DCP not supported yet self.supports_draft_decode_metadata_update = self.dcp_world_size == 1 @@ -1353,16 +1365,7 @@ def build( dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens compressed_slot_mapping = slot_mapping - indexer_block_table = block_table if self.compress_ratio > 1: - kernel_block_size = self.kernel_block_size - if ( - kernel_block_size is not None - and self.kv_cache_spec.block_size != kernel_block_size - and self.kv_cache_spec.block_size % kernel_block_size == 0 - ): - factor = self.kv_cache_spec.block_size // kernel_block_size - indexer_block_table = (block_table[:, ::factor] // factor).contiguous() padded_num_tokens = num_tokens local_slot_mapping = slot_mapping if self.use_pcp: @@ -1378,7 +1381,7 @@ def build( local_slot_mapping, query_start_loc, seq_lens, - indexer_block_table, + block_table, self.kv_cache_spec.num_states, self.compress_ratio, out=self.compressed_slot_mapping_buffer, @@ -1497,7 +1500,7 @@ def build( seq_lens, compressed_seq_lens, compressed_seq_lens_cpu, - indexer_block_table, + block_table, self.compress_ratio, query_slice=query_slice, skip_kv_gather=query_slice.start > 0, @@ -1632,27 +1635,6 @@ def build( ) ) - if self.compress_ratio > 1: - kernel_block_size = self.kernel_block_size - if ( - kernel_block_size is not None - and self.kv_cache_spec.block_size != kernel_block_size - and self.kv_cache_spec.block_size % kernel_block_size == 0 - ): - factor = self.kv_cache_spec.block_size // kernel_block_size - compressed = block_table[:, ::factor] // factor - rows, cols = compressed.shape - if self.indexer_decode_block_table_buffer is None: - self.indexer_decode_block_table_buffer = torch.zeros( - (self._max_num_batched_tokens, cols), - dtype=torch.int32, - device=self.device, - ) - self.indexer_decode_block_table_buffer[:rows, :cols].copy_( - compressed - ) - block_table = self.indexer_decode_block_table_buffer[:rows, :cols] - # Flattening always returns a buffer view, including single-token # batches. Keep its address stable across varlen graph replays. seq_lens_is_buffer_view = not use_native or next_n > 1 @@ -1741,17 +1723,7 @@ def update_draft_decode_metadata( # One query token per request, so the per-token decode block table is # the request block table (at the indexer kernel's page size). build() # may copy it into a builder buffer, so re-derive it for this batch. - block_table = metadata.block_table[: metadata.num_decode_tokens] - kernel_block_size = self.kernel_block_size - if ( - self.compress_ratio > 1 - and kernel_block_size is not None - and self.kv_cache_spec.block_size != kernel_block_size - and self.kv_cache_spec.block_size % kernel_block_size == 0 - ): - factor = self.kv_cache_spec.block_size // kernel_block_size - block_table = block_table[:, ::factor] // factor - decode.block_table.copy_(block_table) + decode.block_table.copy_(metadata.block_table[: metadata.num_decode_tokens]) if decode.indices is not None: decode.indices.copy_(self.arange_buffer[: metadata.num_decode_tokens]) if decode.per_req_decode_lens is not None: diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index e3d982e882bb..e30fec715deb 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -35,6 +35,10 @@ from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( AiterMLAHelper, ) +from vllm.v1.attention.backends.mla.sparse_utils import ( + align_blocks_to_rows, + flat_kv_row_view, +) from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( rocm_sparse_attn_decode_bf16, @@ -160,6 +164,7 @@ def _convert_req_index_to_global_index_kernel( # shapes (compile-time where possible) max_num_blocks_per_req: tl.constexpr, BLOCK_SIZE: tl.constexpr, + BLOCK_STRIDE_ROWS: tl.constexpr, BLOCK_N: tl.constexpr, # tile width along columns # strides (in elements) bt_stride0, @@ -201,9 +206,8 @@ def _convert_req_index_to_global_index_kernel( bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1 base = tl.load(bt_ptr, mask=valid_block, other=0) - # # If token == -1 OR block_id OOB, output 0; else base * BLOCK_SIZE + offset out_val = tl.where( - is_invalid_tok | (~valid_block), 0, base * BLOCK_SIZE + inblock_off + is_invalid_tok | (~valid_block), 0, base * BLOCK_STRIDE_ROWS + inblock_off ) out_ptr_ij = out_ptr + seq_start + indice_id out_ptr_ij_mask = (seq_start + indice_id) < seq_end @@ -218,13 +222,14 @@ def triton_convert_req_index_to_global_index( token_indices: torch.Tensor, # int32 [num_tokens, NUM_TOPK_TOKENS] cu_seqlens: torch.Tensor, # int32 [num_tokens + 1] paged_kv_indices: torch.Tensor, # int32 [num_tokens * topk] out_buffer + BLOCK_STRIDE_ROWS: int, BLOCK_SIZE: int = 64, NUM_TOPK_TOKENS: int = 2048, BLOCK_N: int = 128, # tile width along columns ): """out[token_id, indice_id] = block_table[req_id[token_id], - token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE + token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_STRIDE_ROWS + token_indices[token_id, indice_id] % BLOCK_SIZE Only when token_indices[token_id, indice_id] == -1 do we output -1. @@ -264,6 +269,7 @@ def triton_convert_req_index_to_global_index( # shapes / constexprs max_num_blocks_per_req, BLOCK_SIZE, + BLOCK_STRIDE_ROWS, BLOCK_N, # strides bt_stride0, @@ -345,6 +351,10 @@ class ROCMAiterMLASparseBackend(AttentionBackend): def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: return [1, MultipleOf(16)] + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + return align_blocks_to_rows(spec) + @staticmethod def get_name() -> str: return "ROCM_AITER_MLA_SPARSE" @@ -371,8 +381,7 @@ def is_sparse(cls) -> bool: @classmethod def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: - # Global index conversion assumes contiguous pages within each layer. - return (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC) + return (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC, KVCacheLayout.BLHNC) @classmethod def supports_sink(cls) -> bool: @@ -1155,6 +1164,9 @@ def forward_mqa( self.topk_indices_buffer[:num_actual_toks], attn_metadata.topk_tokens ) + kv_c_and_k_pe_cache, block_stride_rows = flat_kv_row_view( + kv_c_and_k_pe_cache, attn_metadata.block_size + ) triton_convert_req_index_to_global_index( attn_metadata.req_id_per_token, attn_metadata.block_table, @@ -1162,6 +1174,7 @@ def forward_mqa( attn_metadata.paged_kv_indptr, attn_metadata.paged_kv_indices, BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, NUM_TOPK_TOKENS=attn_metadata.topk_tokens, ) @@ -1179,5 +1192,5 @@ def forward_mqa( return output, None mla_padded_q = AiterMLAHelper.get_mla_padded_q(self.num_heads, q) return self._forward_mla( - layer, mla_padded_q, kv_c_and_k_pe_cache, attn_metadata + layer, mla_padded_q, kv_c_and_k_pe_cache.unsqueeze(1), attn_metadata ) diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index 2ec10218d99f..2447df5ceb71 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/vllm/v1/attention/backends/mla/sparse_utils.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Utility functions for sparse MLA backends.""" -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Any import numpy as np @@ -22,6 +22,7 @@ ) from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv +from vllm.v1.kv_cache_interface import AttentionSpec from vllm.v1.worker.block_table import get_block_table_width @@ -40,6 +41,12 @@ def request_row_bounds(req_idx: np.ndarray) -> np.ndarray: return bounds +def align_blocks_to_rows(spec: AttentionSpec, rows: int = 1) -> AttentionSpec: + """Sparse kernels read the cache as flat rows, viewed as ``rows``-row pages, + so blocks must be whole pages apart.""" + return replace(spec, block_stride_alignment=rows * spec.state_content_size_bytes) + + def flat_kv_row_view( kv_cache: torch.Tensor, # [num_blocks, block_size, head_dim] block_size: int, diff --git a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py index 036dfe01730b..87800ecd1727 100644 --- a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py @@ -27,6 +27,7 @@ MLAAttentionImpl, ) from vllm.v1.attention.backends.mla.sparse_utils import ( + align_blocks_to_rows, flat_kv_row_view, triton_convert_req_index_to_global_index, ) @@ -50,6 +51,10 @@ class XPUMLASparseBackend(AttentionBackend): def get_name() -> str: return "XPU_MLA_SPARSE" + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + return align_blocks_to_rows(spec) + @staticmethod def get_metadata_cls() -> type["XPUMLASparseMetadata"]: return XPUMLASparseMetadata diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index c32173a63554..019941a1bf03 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -661,8 +661,8 @@ def __init__( self.hash_block_size = hash_block_size self.dcp_world_size = dcp_world_size # Only groups that participate in prefix caching must satisfy the - # divisibility constraint; groups that opt out (e.g. GLM-5.3-Flash kpool - # tail, block_size=kpool) are scratch buffers and excluded. + # divisibility constraint; groups that opt out (e.g. circular buffers + # such as the GLM-5.3-Flash indexer tail) are scratch and excluded. cacheable_block_sizes = [ manager.block_size for manager, group in zip( @@ -724,8 +724,8 @@ def verify_and_split_kv_cache_groups(self) -> None: """ self.attention_groups: list[SpecGroup] = [] for i, g in enumerate(self.kv_cache_config.kv_cache_groups): - # Skip groups that opt out of prefix caching (e.g. GLM-5.3-Flash - # kpool tail): their blocks are per-request scratch, never + # Skip groups that opt out of prefix caching (e.g. circular + # buffers): their blocks are per-request scratch, never # shareable, so they must not participate in hit lookup (their # manager-level hooks already no-op). Their slot in the per-group # hit tuple stays empty. diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 53a01a5cabe2..919a7b4b3f52 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -10,7 +10,7 @@ from collections.abc import Callable, Iterable, Iterator, Sequence from dataclasses import dataclass, field, replace from functools import partial, reduce -from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, cast, overload +from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, overload from vllm import envs from vllm.config import VllmConfig @@ -32,7 +32,6 @@ FullAttentionSpec, HiddenStateCacheSpec, HiSparseHotSpec, - KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheLayout, @@ -1126,7 +1125,7 @@ def check_enough_kv_cache_memory( # of the specs since grouping may unify them in-place. groups = get_kv_cache_groups(vllm_config, dict(kv_cache_spec)) check_memory = ( - available_memory - _pool_bytes_per_block(groups) + available_memory - _pool_bytes_per_block(vllm_config, groups) if groups else available_memory ) @@ -1237,14 +1236,17 @@ def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int: return num_blocks -def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int: +def _pool_bytes_per_block( + vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec] +) -> int: """Bytes consumed by one block in the worker's shared KV cache pool, mirroring the divisor used by `get_kv_cache_config_from_groups` to convert `available_memory` into `num_blocks`. Used to compute the effective KV cache capacity once `num_gpu_blocks_override` is applied. """ return _get_kv_cache_bytes_per_block( - [group for group in kv_cache_groups if not group.host_resident] + [group for group in kv_cache_groups if not group.host_resident], + vllm_config.cache_config.get_resolved_kv_cache_layout(), ) @@ -1287,205 +1289,6 @@ def _get_kv_cache_groups_uniform_type( return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)] -def _pp_balanced_mamba_group_count( - vllm_config: VllmConfig, - mamba_layer_names: list[str], - mla_layer_names: list[str], -) -> int | None: - """Return a Mamba group count whose PP projections fit the MLA slots.""" - num_groups = cdiv(len(mamba_layer_names), len(mla_layer_names)) - pp_size = vllm_config.parallel_config.pipeline_parallel_size - if pp_size == 1: - return num_groups - - from vllm.distributed.utils import get_pp_indices - from vllm.model_executor.models.utils import extract_layer_index - - total_layers = vllm_config.model_config.get_total_num_hidden_layers() - mamba_indices = [extract_layer_index(name) for name in mamba_layer_names] - mla_indices = [extract_layer_index(name) for name in mla_layer_names] - for rank in range(pp_size): - start, end = get_pp_indices(total_layers, rank, pp_size) - num_mamba = sum(start <= index < end for index in mamba_indices) - num_mla = sum(start <= index < end for index in mla_indices) - if not num_mamba: - continue - if not num_mla: - return None - num_groups = max(num_groups, cdiv(num_mamba, num_mla)) - return num_groups - - -def _get_kv_cache_groups_glm5_next( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec] | None: - """Build GLM-5.3-Flash groups with Mamba/MLA and tail/indexer aliasing.""" - mamba_specs = { - name: spec - for name, spec in kv_cache_spec.items() - if isinstance(spec, MambaSpec) - } - tail_specs = { - name: spec - for name, spec in kv_cache_spec.items() - if isinstance(spec, KpoolTailSpec) - } - attn_specs = { - name: spec - for name, spec in kv_cache_spec.items() - if not isinstance(spec, (MambaSpec, KpoolTailSpec)) - } - if not mamba_specs or not all( - type(spec) is MLAAttentionSpec for spec in attn_specs.values() - ): - return None - - mla_specs = cast(dict[str, MLAAttentionSpec], attn_specs) - idx_pages = { - spec.page_size_bytes for spec in mla_specs.values() if spec.tokens_per_state > 1 - } - if not idx_pages: - return None - - assert all(spec.page_size_padded is None for spec in mla_specs.values()) - assert len(idx_pages) == 1 - mla_names = [name for name, spec in mla_specs.items() if spec.tokens_per_state == 1] - mla_pages = {mla_specs[name].page_size_bytes for name in mla_names} - assert len(mla_pages) == 1 - mla_page = mla_pages.pop() - uniform_spec = UniformTypeKVCacheSpecs.from_specs(attn_specs) - assert uniform_spec is not None - - tail_group: KVCacheGroupSpec | None = None - if tail_specs: - idx_page = next(iter(idx_pages)) - padded_tail_specs: dict[str, KVCacheSpec] = { - name: replace(spec, page_size_padded=idx_page) - for name, spec in tail_specs.items() - } - tail_uniform = UniformTypeKVCacheSpecs.from_specs(padded_tail_specs) - assert tail_uniform is not None - tail_group = KVCacheGroupSpec(list(padded_tail_specs), tail_uniform) - - any_mamba = next(iter(mamba_specs.values())) - assert all(spec == any_mamba for spec in mamba_specs.values()) - if any_mamba.real_page_size_bytes > mla_page: - raise ValueError( - f"the mamba state page ({any_mamba.real_page_size_bytes} bytes) " - f"does not fit the MLA page ({mla_page} bytes); increase tensor " - "parallelism or use a wider KV cache dtype" - ) - padded_specs: dict[str, KVCacheSpec] = { - name: replace(any_mamba, page_size_padded=mla_page) for name in mamba_specs - } - num_groups = _pp_balanced_mamba_group_count( - vllm_config, list(mamba_specs), mla_names - ) - if num_groups is None: - raise ValueError( - "a pipeline stage has mamba layers but no MLA layer to share " - "slots with; realign the stage boundaries (VLLM_PP_LAYER_PARTITION)" - ) - mamba_grouped_names: list[list[str]] = [[] for _ in range(num_groups)] - for index, name in enumerate(mamba_specs): - mamba_grouped_names[index % num_groups].append(name) - - return ( - [KVCacheGroupSpec(list(attn_specs), uniform_spec)] - + ([tail_group] if tail_group is not None else []) - + create_kv_cache_group_specs(padded_specs, mamba_grouped_names) - ) - - -def _glm5_next_tensor_layout( - kv_cache_groups: list[KVCacheGroupSpec], -) -> ( - tuple[ - KVCacheGroupSpec, - list[KVCacheGroupSpec], - list[str], - list[str], - int, - int, - list[str], - int, - ] - | None -): - """Recognize the GLM-5.3-Flash grouping after optional PP projection.""" - uniform_groups = [ - group - for group in kv_cache_groups - if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) - ] - mamba_groups = [ - group for group in kv_cache_groups if isinstance(group.kv_cache_spec, MambaSpec) - ] - attn_group: KVCacheGroupSpec | None = None - tail_group: KVCacheGroupSpec | None = None - for group in uniform_groups: - inner = cast(UniformTypeKVCacheSpecs, group.kv_cache_spec).kv_cache_specs - if all(type(spec) is MLAAttentionSpec for spec in inner.values()): - attn_group = group - elif all(isinstance(spec, KpoolTailSpec) for spec in inner.values()): - tail_group = group - if attn_group is None or not mamba_groups: - return None - if len(uniform_groups) + len(mamba_groups) != len(kv_cache_groups): - return None - - attn_uniform = cast(UniformTypeKVCacheSpecs, attn_group.kv_cache_spec) - mla_inner = cast(dict[str, MLAAttentionSpec], attn_uniform.kv_cache_specs) - if not all( - type(spec) is MLAAttentionSpec and spec.page_size_padded is None - for spec in mla_inner.values() - ): - return None - mla_names = [ - name for name in attn_group.layer_names if mla_inner[name].tokens_per_state == 1 - ] - idx_names = [ - name for name in attn_group.layer_names if mla_inner[name].tokens_per_state > 1 - ] - mla_pages = {mla_inner[name].page_size_bytes for name in mla_names} - idx_pages = {mla_inner[name].page_size_bytes for name in idx_names} - if len(mla_pages) != 1 or len(idx_pages) != 1: - return None - mla_page = mla_pages.pop() - idx_page = idx_pages.pop() - if any(group.kv_cache_spec.page_size_bytes != mla_page for group in mamba_groups): - return None - - tail_names: list[str] = [] - tail_page = 0 - if tail_group is not None: - tail_names = list(tail_group.layer_names) - tail_inner = cast( - UniformTypeKVCacheSpecs, tail_group.kv_cache_spec - ).kv_cache_specs - tail_pages = { - cast(KpoolTailSpec, spec).unpadded_page_size_bytes - for spec in tail_inner.values() - } - if len(tail_pages) != 1 or len(tail_names) != len(idx_names): - return None - tail_page = tail_pages.pop() - if tail_page > idx_page: - return None - - return ( - attn_group, - mamba_groups, - mla_names, - idx_names, - mla_page, - idx_page, - tail_names, - tail_page, - ) - - def unify_kv_cache_spec_page_size( kv_cache_spec: dict[str, KVCacheSpec], ) -> dict[str, KVCacheSpec]: @@ -1709,14 +1512,24 @@ def _get_per_layer_spec( return spec +def _get_layer_block_stride(spec: KVCacheSpec) -> int: + """A layer-outer block stride: the page padded to its alignment.""" + return round_up(spec.page_size_bytes, spec.get_block_stride_alignment()) + + def _get_kv_cache_bytes_per_block( - kv_cache_groups: list[KVCacheGroupSpec], + kv_cache_groups: list[KVCacheGroupSpec], layout: KVCacheLayout ) -> int: """Return the largest cache group's bytes per block.""" - if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None: - _, _, mla_names, idx_names, mla_page, idx_page, _, _ = glm5_layout - return len(mla_names) * mla_page + len(idx_names) * idx_page - + if not layout.is_block_outermost: + # Each layer's block stride is its own padded page. + return max( + sum( + _get_layer_block_stride(_get_per_layer_spec(group, layer_name)) + for layer_name in group.layer_names + ) + for group in kv_cache_groups + ) bytes_per_block = max( sum( _get_per_layer_spec(group, layer_name).page_size_bytes @@ -1731,7 +1544,7 @@ def _get_kv_cache_bytes_per_block( if isinstance(group.kv_cache_spec, HiSparseHotSpec) ] alignments.extend( - _get_per_layer_spec(group, layer_name).block_stride_alignment or 1 + _get_per_layer_spec(group, layer_name).get_block_stride_alignment() for group in kv_cache_groups for layer_name in group.layer_names ) @@ -1806,72 +1619,9 @@ def get_kv_cache_config_from_groups( vllm_config, kv_cache_groups, available_memory, host_budget ) - if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None: - ( - attn_group, - mamba_groups, - mla_names, - idx_names, - mla_page, - idx_page, - tail_names, - _, - ) = glm5_layout - bytes_per_block = len(mla_names) * mla_page + len(idx_names) * idx_page - num_blocks = may_override_num_blocks( - vllm_config, available_memory // bytes_per_block - ) - size = bytes_per_block * num_blocks - attn_specs = cast( - UniformTypeKVCacheSpecs, attn_group.kv_cache_spec - ).kv_cache_specs - - kv_cache_tensors: list[KVCacheTensor] = [] - - def add_tensor(layer_name: str, spec: KVCacheSpec, offset: int) -> None: - kv_cache_tensors.append( - KVCacheTensor( - size=size, - layers=[layer_name], - layer_stride=spec.page_size_bytes * num_blocks, - block_stride=spec.page_size_bytes, - offset=offset, - ) - ) - - for index, mla_name in enumerate(mla_names): - offset = index * mla_page * num_blocks - add_tensor(mla_name, attn_specs[mla_name], offset) - for group in mamba_groups: - if index < len(group.layer_names): - add_tensor(group.layer_names[index], group.kv_cache_spec, offset) - - idx_base = len(mla_names) * mla_page * num_blocks - for index, idx_name in enumerate(idx_names): - offset = idx_base + index * idx_page * num_blocks - add_tensor(idx_name, attn_specs[idx_name], offset) - if tail_names: - tail_name = tail_names[index] - tail_group = next( - group for group in kv_cache_groups if tail_name in group.layer_names - ) - tail_specs = cast( - UniformTypeKVCacheSpecs, tail_group.kv_cache_spec - ).kv_cache_specs - add_tensor(tail_name, tail_specs[tail_name], offset) - - return KVCacheConfig( - num_blocks=num_blocks, - kv_cache_tensors=kv_cache_tensors, - kv_cache_groups=kv_cache_groups, - prefix_cache_retention_interval=( - vllm_config.cache_config.prefix_cache_retention_interval - ), - ) - layout = vllm_config.cache_config.get_resolved_kv_cache_layout() validate_kv_cache_layout(layout, kv_cache_groups) - bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups) + bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups, layout) interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None num_blocks = available_memory // bytes_per_block @@ -1901,17 +1651,17 @@ def add_tensor(layer_name: str, spec: KVCacheSpec, offset: int) -> None: byte_offset = 0 for spec, layer_names in layers_by_spec.items(): + block_stride = interleaved_block_stride + padded_page = _get_layer_block_stride(spec) + if block_stride is None and padded_page != spec.page_size_bytes: + assert layout.is_block_compact, f"{layout.name} cannot pad pages" + block_stride = padded_page layer_stride, block_stride, _, _, _ = compute_layout_strides( spec, num_blocks, len(layer_names), layout, - fixed_strides=(None, interleaved_block_stride, None, None, None), - ) - offset = ( - byte_offset - * max(layer_stride, spec.page_size_bytes) - // spec.page_size_bytes + fixed_strides=(None, block_stride, None, None, None), ) kv_cache_tensors.append( KVCacheTensor( @@ -1919,10 +1669,13 @@ def add_tensor(layer_name: str, spec: KVCacheSpec, offset: int) -> None: layers=layer_names, layer_stride=layer_stride, block_stride=block_stride, - offset=offset, + offset=byte_offset, ) ) - byte_offset += len(layer_names) * spec.page_size_bytes + # Block-outer layers share each block; layer-outer ones own regions. + byte_offset += len(layer_names) * ( + spec.page_size_bytes if layout.is_block_outermost else layer_stride + ) return KVCacheConfig( num_blocks=num_blocks, @@ -2438,9 +2191,6 @@ def get_kv_cache_groups( # full attention, or all layers are sliding window attention with the # same window size). Put all layers into one group. return _get_kv_cache_groups_uniform_type(uniform_spec) - elif glm5_groups := _get_kv_cache_groups_glm5_next(vllm_config, kv_cache_spec): - return glm5_groups - # Hidden-state layers use their own block table and must not be absorbed # into a compatible attention bucket. hidden_specs = { @@ -2594,31 +2344,7 @@ def _max_memory_usage_bytes_from_groups( if not kv_cache_groups: return 0 - if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None: - ( - attn_group, - mamba_groups, - mla_names, - idx_names, - mla_page, - idx_page, - tail_names, - _, - ) = glm5_layout - uniform_spec = cast(UniformTypeKVCacheSpecs, attn_group.kv_cache_spec) - total_blocks = uniform_spec.max_memory_usage_pages(vllm_config) - total_blocks += sum( - cdiv( - group.kv_cache_spec.max_memory_usage_bytes(vllm_config), - group.kv_cache_spec.page_size_bytes, - ) - for group in mamba_groups - ) - if tail_names: - total_blocks += 1 - return total_blocks * (len(mla_names) * mla_page + len(idx_names) * idx_page) - - bytes_per_block = _pool_bytes_per_block(kv_cache_groups) + bytes_per_block = _pool_bytes_per_block(vllm_config, kv_cache_groups) total_blocks = 0 for group in kv_cache_groups: spec = group.kv_cache_spec @@ -2874,7 +2600,7 @@ def get_kv_cache_configs( if not groups: adjusted_memory.append(avail_mem) continue - bytes_per_block = _pool_bytes_per_block(groups) + bytes_per_block = _pool_bytes_per_block(vllm_config, groups) logger.info( "Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d", avail_mem // bytes_per_block, @@ -2887,7 +2613,7 @@ def get_kv_cache_configs( # the capacity check both plan against usable blocks. Allocation below # still uses the full memory. check_memory = [ - avail_mem - _pool_bytes_per_block(groups) if groups else avail_mem + avail_mem - _pool_bytes_per_block(vllm_config, groups) if groups else avail_mem for groups, avail_mem in zip(projected_groups_per_worker, available_memory) ] @@ -2931,7 +2657,9 @@ def get_kv_cache_configs( # strides and offsets stay consistent with the shrunken allocation. groups = kv_cache_config.kv_cache_groups kv_cache_configs[i] = get_kv_cache_config_from_groups( - vllm_config, groups, min_num_blocks * _pool_bytes_per_block(groups) + vllm_config, + groups, + min_num_blocks * _pool_bytes_per_block(vllm_config, groups), ) for kv_cache_config in kv_cache_configs: diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 3e7ae1b7063a..e9f5f2f3c231 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, ClassVar from vllm.distributed.kv_events import MEDIUM_CPU -from vllm.logger import init_logger from vllm.utils.math_utils import cdiv from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_utils import ( @@ -27,7 +26,6 @@ HiddenStateCacheSpec, HiSparseHotSpec, HiSparseResidentSpec, - KpoolTailSpec, KVCacheGroupRole, KVCacheSpec, MambaSpec, @@ -45,8 +43,6 @@ if TYPE_CHECKING: from vllm.v1.hisparse.coordinator import HiSparseCoordinator -logger = init_logger(__name__) - class SingleTypeKVCacheManager(ABC): """An abstract base class for a manager that handle the kv cache management @@ -1299,10 +1295,6 @@ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: return 0 -class KpoolTailManager(CircularBufferManager): - """One-block circular scratch manager for ``KpoolTailSpec``.""" - - class ChunkedLocalAttentionManager(SingleTypeKVCacheManager): def __init__(self, kv_cache_spec: ChunkedLocalAttentionSpec, **kwargs) -> None: super().__init__(kv_cache_spec, **kwargs) @@ -2687,11 +2679,6 @@ def register_all_kvcache_specs(vllm_config): SlidingWindowManager, uniform_type_base_spec=SlidingWindowMLASpec, ) - KVCacheSpecRegistry.register( - KpoolTailSpec, - KpoolTailManager, - uniform_type_base_spec=KpoolTailSpec, - ) KVCacheSpecRegistry.register( MambaSpec, MambaManager, uniform_type_base_spec=MambaSpec diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index bfbaa38d9478..d0dadae27e9b 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -335,8 +335,8 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: vllm_config.cache_config.num_gpu_blocks = scheduler_kv_cache_config.num_blocks kv_cache_groups = scheduler_kv_cache_config.kv_cache_groups if kv_cache_groups: - # Exclude groups that opt out of prefix caching (e.g. GLM-5.3-Flash - # kpool tail, a 1-block/req scratch buffer with block_size=kpool): + # Exclude groups that opt out of prefix caching (e.g. a circular + # buffer, a 1-block/req scratch buffer with a small ring block): # their small block_size would otherwise drag the global block_size # below the real allocator block size and desync it from mamba. participating = [ diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index b8c0d5dff269..60c60a3ff504 100644 --- a/vllm/v1/hisparse/binding.py +++ b/vllm/v1/hisparse/binding.py @@ -55,8 +55,9 @@ def resolve_hisparse_specs( if len(block_sizes) != 1: raise ValueError("HiSparse requires one scheduler block size.") backends = [attn_layers[name].get_attn_backend() for name in mla_specs] + specs = list(mla_specs.values()) try: - block_size = select_common_block_size(block_sizes.pop(), backends) + block_size = select_common_block_size(block_sizes.pop(), backends, specs) except ValueError as error: raise ValueError( "HiSparse requires a GPU block size supported by every sparse " @@ -104,16 +105,13 @@ def allocate_hisparse_kv_caches( if isinstance(host_spec, UniformTypeKVCacheSpecs) else host_spec ) - kernel_block_size = kernel_block_sizes[host_group_id] - if isinstance(spec, MLAAttentionSpec) and spec.storage_block_size is not None: - kernel_block_size = spec.storage_block_size views = create_kv_cache_views( backing, spec, kv_cache_config.num_blocks_of(tensor), layout, tensor, - kernel_block_size=kernel_block_size, + kernel_block_size=kernel_block_sizes[host_group_id], ) kv_caches.update(zip(tensor.layers, views)) return kv_caches diff --git a/vllm/v1/hisparse/layout.py b/vllm/v1/hisparse/layout.py index 98636866f8ac..96d93650ad29 100644 --- a/vllm/v1/hisparse/layout.py +++ b/vllm/v1/hisparse/layout.py @@ -333,7 +333,7 @@ def get_hisparse_kv_cache_config( device_groups = hisparse_layout.device_groups layout = vllm_config.cache_config.get_resolved_kv_cache_layout() validate_kv_cache_layout(layout, device_groups) - bytes_per_block = _get_kv_cache_bytes_per_block(device_groups) + bytes_per_block = _get_kv_cache_bytes_per_block(device_groups, layout) num_blocks = may_override_num_blocks( vllm_config, available_memory // bytes_per_block ) @@ -345,7 +345,7 @@ def get_hisparse_kv_cache_config( host_groups = [hisparse_layout.source_group] host_layout = KVCacheLayout.LBNHC validate_kv_cache_layout(host_layout, host_groups) - host_bytes_per_block = _get_kv_cache_bytes_per_block(host_groups) + host_bytes_per_block = _get_kv_cache_bytes_per_block(host_groups, host_layout) host_size = host_bytes_per_block * hisparse_layout.host_num_blocks kv_cache_tensors[:0] = _build_hisparse_kv_cache_tensors( host_groups, diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index ca3f485917eb..79e88ca91d8f 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -4,6 +4,7 @@ from __future__ import annotations import copy +import math from collections import Counter from collections.abc import Collection, Sequence from dataclasses import dataclass, field, fields, replace @@ -24,6 +25,7 @@ if TYPE_CHECKING: from vllm.config import VllmConfig + from vllm.v1.attention.backend import MultipleOf logger = init_logger(__name__) @@ -168,13 +170,28 @@ class KVCacheSpec: dcp_sharded: bool = field(default=False, kw_only=True) """Whether DCP shards this cache's token positions across ranks.""" - block_stride_alignment: int | None = field(default=None, kw_only=True) - """Required byte alignment between physical blocks, including packed layers.""" + block_stride_alignment: int | MultipleOf | None = field(default=None, kw_only=True) + """Required byte alignment between physical blocks, including packed layers. + ``MultipleOf(n)``: whole kernel blocks, when blocks are split into kernel + blocks of up to ``n`` tokens.""" def __post_init__(self): - if self.block_stride_alignment is not None and self.block_stride_alignment <= 0: + alignment = self.block_stride_alignment + if isinstance(alignment, int) and alignment <= 0: raise ValueError("block_stride_alignment must be positive") + def get_block_stride_alignment(self) -> int: + """Bytes the stride between physical blocks must be a multiple of.""" + from vllm.v1.attention.backend import MultipleOf + + alignment = self.block_stride_alignment + if not isinstance(alignment, MultipleOf): + return alignment or 1 + kernel_block_size = math.gcd(alignment.base, self.block_size) + if kernel_block_size == self.block_size: + return 1 + return self.copy_with_new_block_size(kernel_block_size).page_size_bytes + @property def prefix_cacheable(self) -> bool: """Whether this spec's group participates in prefix caching.""" @@ -674,8 +691,6 @@ class MLAAttentionSpec(FullAttentionSpec): model_version: str | None = None cache_role: SparseCacheRole = SparseCacheRole.SPARSE is_index_group_leader: bool = False - storage_block_size: int | None = None - """Token width used to view storage when it differs from the kernel block.""" # Group capability enabled when any member flattens a non-causal query block # into decode rows. Runtime metadata still selects causal vs. non-causal mode. non_causal_multi_token_decode: bool = False @@ -696,7 +711,6 @@ def merge(cls, specs: list[Self]) -> Self: model_version_set = set(spec.model_version for spec in specs) cache_role_set = {spec.cache_role for spec in specs} index_group_leader_set = {spec.is_index_group_leader for spec in specs} - storage_block_size_set = set(spec.storage_block_size for spec in specs) block_stride_alignment_set = {spec.block_stride_alignment for spec in specs} assert ( len(cache_dtype_str_set) == 1 @@ -704,12 +718,11 @@ def merge(cls, specs: list[Self]) -> Self: and len(model_version_set) == 1 and len(cache_role_set) == 1 and len(index_group_leader_set) == 1 - and len(storage_block_size_set) == 1 and len(block_stride_alignment_set) == 1 ), ( "All attention layers in the same KV cache group must use the same " "quantization method, tokens per state, model version, cache role, " - "index-sharing role, storage block size and block stride alignment." + "index-sharing role, and block stride alignment." ) merged_spec = cls( block_size=specs[0].block_size, @@ -727,7 +740,6 @@ def merge(cls, specs: list[Self]) -> Self: model_version=model_version_set.pop(), cache_role=cache_role_set.pop(), is_index_group_leader=index_group_leader_set.pop(), - storage_block_size=storage_block_size_set.pop(), block_stride_alignment=block_stride_alignment_set.pop(), non_causal_multi_token_decode=any( spec.non_causal_multi_token_decode for spec in specs @@ -1019,32 +1031,6 @@ def is_uniform_with_collection( ) -@dataclass(frozen=True, kw_only=True) -class KpoolTailSpec(SlidingWindowSpec): - """One-block circular scratch cache for a kpool indexer's raw tail.""" - - def max_admission_blocks_per_request( - self, max_in_flight_tokens: int, max_model_len: int - ) -> int: - return 1 - - def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: - return 1 - - def is_uniform_with_collection( - self, kv_cache_specs: dict[str, KVCacheSpec] - ) -> bool: - return all(isinstance(spec, KpoolTailSpec) for spec in kv_cache_specs.values()) - - @property - def prefix_cacheable(self) -> bool: - return False - - @property - def uses_slot_mapping(self) -> bool: - return False - - @dataclass(frozen=True) class MambaSpec(KVCacheSpec): shapes: tuple[tuple[int, ...], ...] diff --git a/vllm/v1/spec_decode/gemma4.py b/vllm/v1/spec_decode/gemma4.py index bea677dde85a..af504b0a1108 100644 --- a/vllm/v1/spec_decode/gemma4.py +++ b/vllm/v1/spec_decode/gemma4.py @@ -95,7 +95,7 @@ def build_per_group_and_layer_attn_metadata( cm.block_table_tensor = self._per_group_block_tables[gid][:batch_size] else: cm = common_attn_metadata - attn_metadata = attn_group.get_metadata_builder().build_for_drafting( + attn_metadata = attn_group.build_metadata_for_drafting( common_attn_metadata=cm, draft_index=draft_index ) per_group_attn_metadata.append(attn_metadata) @@ -291,11 +291,7 @@ def initialize_attn_backend( self.draft_attn_groups = list(attention_groups.values()) if self.draft_attn_groups: self.kv_cache_gid = self.draft_attn_groups[0].kv_cache_group_id - self.block_size = ( - self.draft_attn_groups[0] - .get_metadata_builder() - .kv_cache_spec.block_size - ) + self.block_size = self.draft_attn_groups[0].kv_cache_spec.block_size else: self.kv_cache_gid = 0 self.block_size = kv_cache_config.kv_cache_groups[ diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index e4718bb3d7d1..79be2e8a4d86 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -991,7 +991,7 @@ def build_per_group_and_layer_attn_metadata( per_group_attn_metadata: list[object] = [] per_layer_attn_metadata: dict[str, object] = {} for attn_group in self.draft_attn_groups: - attn_metadata = attn_group.get_metadata_builder().build_for_drafting( + attn_metadata = attn_group.build_metadata_for_drafting( common_attn_metadata=common_attn_metadata, draft_index=draft_index ) per_group_attn_metadata.append(attn_metadata) @@ -1765,20 +1765,7 @@ def initialize_attn_backend( attention_groups[backend_key].layer_names.append(layer_name) self.draft_attn_groups = list(attention_groups.values()) - if kernel_block_sizes is not None and 0 <= self.kv_cache_gid < len( - kernel_block_sizes - ): - # Slot mappings are computed against the block table, which is - # stored at kernel-block granularity. Use the kernel block size - # rather than the KV cache manager's block size; the two differ - # when manager blocks are split for the attention kernel. - self.block_size = kernel_block_sizes[self.kv_cache_gid] - else: - self.block_size = ( - self.draft_attn_groups[0] - .get_metadata_builder() - .kv_cache_spec.block_size - ) + self.block_size = self.draft_attn_groups[0].kv_cache_spec.block_size logger.debug("Using block size %d for drafting layers", self.block_size) def _determine_batch_execution_and_padding( diff --git a/vllm/v1/spec_decode/step3p5.py b/vllm/v1/spec_decode/step3p5.py index 7a73473adc2f..0853067f3906 100644 --- a/vllm/v1/spec_decode/step3p5.py +++ b/vllm/v1/spec_decode/step3p5.py @@ -144,7 +144,7 @@ def build_per_group_and_layer_attn_metadata( cm.slot_mapping = sm else: cm = common_attn_metadata - attn_metadata = attn_group.get_metadata_builder().build_for_drafting( + attn_metadata = attn_group.build_metadata_for_drafting( common_attn_metadata=cm, draft_index=draft_index, ) @@ -244,11 +244,7 @@ def initialize_attn_backend( self.draft_attn_groups = list(attention_groups.values()) if self.draft_attn_groups: self.kv_cache_gid = self.draft_attn_groups[0].kv_cache_group_id - self.block_size = ( - self.draft_attn_groups[0] - .get_metadata_builder() - .kv_cache_spec.block_size - ) + self.block_size = self.draft_attn_groups[0].kv_cache_spec.block_size else: self.kv_cache_gid = 0 self.block_size = kv_cache_config.kv_cache_groups[ diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index c63908dad0d0..2beec249bff0 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -29,24 +29,16 @@ def get_block_table_width( max_num_blocks: int, block_size: int, - kernel_block_size: int | None = None, *, token_alignment: int | None = 128, ) -> int: - """Return the width after optional alignment and virtual block splitting.""" - if kernel_block_size is None: - kernel_block_size = block_size - if block_size % kernel_block_size != 0: - raise ValueError( - f"kernel_block_size {kernel_block_size} must divide " - f"block_size {block_size} evenly" - ) + """Return the width after optional alignment.""" if token_alignment is not None: if token_alignment <= 0: raise ValueError("token_alignment must be positive") block_alignment = token_alignment // math.gcd(token_alignment, block_size) max_num_blocks = cdiv(max_num_blocks, block_alignment) * block_alignment - return max_num_blocks * block_size // kernel_block_size + return max_num_blocks class SlotMappingMode(Enum): @@ -63,7 +55,6 @@ def __init__( max_num_batched_tokens: int, pin_memory: bool, device: torch.device, - kernel_block_size: int, cp_kv_cache_interleave_size: int, slot_mapping_mode: SlotMappingMode = SlotMappingMode.TOKEN_TO_KV_SLOT, ): @@ -74,9 +65,6 @@ def __init__( max_num_batched_tokens: Maximum number of tokens in a batch. pin_memory: Whether to pin memory for faster GPU transfers. device: Target device for the block table. - kernel_block_size: The block_size of underlying attention kernel. - Will be the same as `block_size` if `block_size` is supported - by the attention kernel. slot_mapping_mode: How this cache group maps scheduled tokens to cache slots. Mamba-like state caches do not use token slot mappings and should use SlotMappingMode.NONE. @@ -86,30 +74,8 @@ def __init__( self.max_num_batched_tokens = max_num_batched_tokens self.pin_memory = pin_memory self.device = device - self.kv_cache_block_size = block_size - - if kernel_block_size == block_size: - # Standard case: allocation and computation use same block size - # No block splitting needed, direct mapping - self.block_size = block_size - self.blocks_per_kv_block = 1 - self.use_hybrid_blocks = False - else: - # Hybrid case: allocation block size differs from kernel block size - # Memory blocks are subdivided to match kernel requirements - # Example: 32-token memory blocks with 16-token kernel blocks - # → Each memory block corresponds to 2 kernel blocks - if block_size % kernel_block_size != 0: - raise ValueError( - f"kernel_block_size {kernel_block_size} must divide " - f"kv_manager_block_size size {block_size} evenly" - ) - - self.block_size = kernel_block_size - self.blocks_per_kv_block = block_size // kernel_block_size - self.use_hybrid_blocks = True - - self.max_num_blocks_per_req = max_num_blocks_per_req * self.blocks_per_kv_block + self.block_size = block_size + self.max_num_blocks_per_req = max_num_blocks_per_req self.block_table = self._make_buffer( self.max_num_reqs, self.max_num_blocks_per_req, dtype=torch.int32 @@ -120,13 +86,6 @@ def __init__( self.max_num_batched_tokens, dtype=torch.int64 ) - if self.use_hybrid_blocks: - self._kernel_block_arange = np.arange(0, self.blocks_per_kv_block).reshape( - 1, -1 - ) - else: - self._kernel_block_arange = None - try: self.pcp_world_size = get_pcp_group().world_size self.pcp_rank = get_pcp_group().rank_in_group @@ -145,13 +104,11 @@ def __init__( self.slot_mapping_mode = slot_mapping_mode if self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT: _COMPUTE_SLOT_MAPPING_KERNEL.register_warmup( - kv_cache_block_size=self.kv_cache_block_size, - blocks_per_kv_block=self.blocks_per_kv_block, + kv_cache_block_size=self.block_size, total_cp_world_size=self.dcp_world_size, total_cp_rank=self.dcp_rank, cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, block_table_stride=self.block_table.gpu.stride(0), - block_size=self.block_size, ) def append_row( @@ -162,11 +119,6 @@ def append_row( if not block_ids: return - if self.use_hybrid_blocks: - block_ids = self.map_to_kernel_blocks( - np.array(block_ids), self.blocks_per_kv_block, self._kernel_block_arange - ) - num_blocks = len(block_ids) start = self.num_blocks_per_row[row_idx] self.num_blocks_per_row[row_idx] += num_blocks @@ -219,10 +171,8 @@ def compute_slot_mapping( positions, self.block_table.gpu, self.block_table.gpu.stride(0), - self.block_size, self.slot_mapping.gpu, - self.kv_cache_block_size, - self.blocks_per_kv_block, + self.block_size, self.dcp_world_size, self.dcp_rank, self.cp_kv_cache_interleave_size, @@ -235,37 +185,6 @@ def clear(self) -> None: self.block_table.gpu.fill_(0) self.block_table.cpu.fill_(0) - @staticmethod - def map_to_kernel_blocks( - kv_manager_block_ids: np.ndarray, - blocks_per_kv_block: int, - kernel_block_arange: np.ndarray, - ) -> np.ndarray: - """Convert kv_manager_block_id IDs to kernel block IDs. - - Example: - # kv_manager_block_ids: 32 tokens, - # Kernel block size: 16 tokens - # blocks_per_kv_block = 2 - >>> kv_manager_block_ids = np.array([0, 1, 2]) - >>> Result: [0, 1, 2, 3, 4, 5] - - # Each kv_manager_block_id maps to 2 kernel block id: - # kv_manager_block_id 0 → kernel block id [0, 1] - # kv_manager_block_id 1 → kernel block id [2, 3] - # kv_manager_block_id 2 → kernel block id [4, 5] - - """ - if blocks_per_kv_block == 1: - return kv_manager_block_ids - - kernel_block_ids = ( - kv_manager_block_ids.reshape(-1, 1) * blocks_per_kv_block - + kernel_block_arange - ) - - return kernel_block_ids.reshape(-1) - def get_device_tensor(self, num_reqs: int) -> torch.Tensor: """Returns the device tensor of the block table.""" return self.block_table.gpu[:num_reqs] @@ -296,16 +215,10 @@ def __init__( pin_memory: bool, device: torch.device, block_sizes: list[int], - kernel_block_sizes: list[int], max_num_blocks: list[int], cp_kv_cache_interleave_size: int = 1, slot_mapping_modes: list[SlotMappingMode] | None = None, ) -> None: - if len(kernel_block_sizes) != len(block_sizes): - raise ValueError( - f"kernel_block_sizes length ({len(kernel_block_sizes)}) " - f"must match block_sizes length ({len(block_sizes)})" - ) if slot_mapping_modes is None: slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT] * len(block_sizes) if len(slot_mapping_modes) != len(block_sizes): @@ -339,17 +252,11 @@ def __init__( max_num_batched_tokens, pin_memory, device, - kernel_block_size, cp_kv_cache_interleave_size, slot_mapping_mode=slot_mapping_mode, ) - for ( - block_size, - kernel_block_size, - max_num_blocks_per_req, - slot_mapping_mode, - ) in zip( - block_sizes, kernel_block_sizes, max_num_blocks, slot_mapping_modes + for block_size, max_num_blocks_per_req, slot_mapping_mode in zip( + block_sizes, max_num_blocks, slot_mapping_modes ) ] @@ -403,12 +310,10 @@ class ComputeSlotMappingKernel( @dataclass(frozen=True) class CompileKey: kv_cache_block_size: int - blocks_per_kv_block: int total_cp_world_size: int total_cp_rank: int cp_kv_cache_interleave_size: int block_table_stride: int - block_size: int @staticmethod @triton.jit(do_not_specialize=["num_tokens", "max_num_tokens"]) @@ -419,10 +324,8 @@ def kernel( positions_ptr, # [num_tokens], int64 block_table_ptr, # [max_num_reqs, max_num_blocks_per_req], int32 (flat) block_table_stride, # max_num_blocks_per_req - block_size, slot_mapping_ptr, # [max_num_tokens], int64 KV_CACHE_BLOCK_SIZE: tl.constexpr, - BLOCKS_PER_KV_BLOCK: tl.constexpr, TOTAL_CP_WORLD_SIZE: tl.constexpr, TOTAL_CP_RANK: tl.constexpr, CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr, @@ -463,17 +366,12 @@ def kernel( virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE ) - block_indices = ( - virtual_block_indices * BLOCKS_PER_KV_BLOCK - + local_block_offsets // block_size - ) block_numbers = tl.load( - block_table_ptr + row_offset + block_indices, + block_table_ptr + row_offset + virtual_block_indices, mask=mask & is_local, other=0, ).to(tl.int64) - slot_offsets = local_block_offsets % block_size - slot_ids = block_numbers * block_size + slot_offsets + slot_ids = block_numbers * KV_CACHE_BLOCK_SIZE + local_block_offsets slot_ids = tl.where(is_local, slot_ids, PAD_ID) tl.store(slot_mapping_ptr + offsets, slot_ids, mask=mask) @@ -481,13 +379,11 @@ def dispatch( # type: ignore[override] self, *, block_table_stride: int, - block_size: int, **compile_key_fields: int, ) -> CompileKey: return self.CompileKey( **compile_key_fields, block_table_stride=triton_scalar_specialization_rep(block_table_stride), - block_size=triton_scalar_specialization_rep(block_size), ) def get_warmup_keys(self, **dispatch_kwargs: int) -> list[CompileKey]: @@ -507,10 +403,8 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: shape=(1, compile_key.block_table_stride), ), block_table_stride=compile_key.block_table_stride, - block_size=compile_key.block_size, slot_mapping=int64_ptr, kv_cache_block_size=compile_key.kv_cache_block_size, - blocks_per_kv_block=compile_key.blocks_per_kv_block, total_cp_world_size=compile_key.total_cp_world_size, total_cp_rank=compile_key.total_cp_rank, cp_kv_cache_interleave_size=compile_key.cp_kv_cache_interleave_size, @@ -526,17 +420,14 @@ def __call__( positions: torch.Tensor, block_table: torch.Tensor, block_table_stride: int, - block_size: int, slot_mapping: torch.Tensor, kv_cache_block_size: int, - blocks_per_kv_block: int, total_cp_world_size: int, total_cp_rank: int, cp_kv_cache_interleave_size: int, ) -> LaunchSpec: return (num_reqs + 1,), dict( KV_CACHE_BLOCK_SIZE=kv_cache_block_size, - BLOCKS_PER_KV_BLOCK=blocks_per_kv_block, TOTAL_CP_WORLD_SIZE=total_cp_world_size, TOTAL_CP_RANK=total_cp_rank, CP_KV_CACHE_INTERLEAVE_SIZE=cp_kv_cache_interleave_size, diff --git a/vllm/v1/worker/cp_utils.py b/vllm/v1/worker/cp_utils.py index bcf56dc0bf78..49dcccf44318 100644 --- a/vllm/v1/worker/cp_utils.py +++ b/vllm/v1/worker/cp_utils.py @@ -95,11 +95,9 @@ def prepare_dcp_dummy_context_metadata( max_valid_block_id = kv_cache_config.num_blocks - 1 assert max_valid_block_id > 0 for blk_table in input_batch.block_table.block_tables: - max_row_blocks = ( - blk_table.max_num_blocks_per_req // blk_table.blocks_per_kv_block - ) block_ids = [ - (block_idx % max_valid_block_id) + 1 for block_idx in range(max_row_blocks) + (block_idx % max_valid_block_id) + 1 + for block_idx in range(blk_table.max_num_blocks_per_req) ] for req_idx in range(num_reqs): blk_table.add_row(block_ids, req_idx) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 8aa9ba3cb8b2..d1feaef031d9 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass from typing import TYPE_CHECKING, Any, cast @@ -40,6 +40,8 @@ add_kv_sharing_layers_to_kv_cache_groups, allocate_kv_cache, bind_kv_cache_to_layers, + customize_attention_spec, + map_kv_caches_to_kernel_blocks, prepare_kernel_block_sizes, ) @@ -141,7 +143,7 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: # Skip modules that don't need KV cache (eg encoder-only attention) if spec := attn_module.get_kv_cache_spec(vllm_config): if isinstance(spec, AttentionSpec): - spec = attn_module.get_attn_backend().customize_spec(spec) + spec = customize_attention_spec(attn_module.get_attn_backend(), spec) kv_cache_spec[layer_name] = spec if vllm_config.attention_config.hisparse_config is not None: kv_cache_spec = resolve_hisparse_specs(vllm_config, kv_cache_spec, attn_layers) @@ -382,6 +384,7 @@ def init_kv_cache( kv_cache_allocation_context: AbstractContextManager | None = None, *, block_tables: "BlockTables | None" = None, + attn_groups: Iterable[AttentionGroup] = (), ) -> dict[str, Any]: allocation_context = kv_cache_allocation_context or nullcontext() with allocation_context: @@ -409,9 +412,10 @@ def init_kv_cache( num_attn_module = ( 2 if vllm_config.model_config.hf_config.model_type == "longcat_flash" else 1 ) - bindable_caches = { - name: cache for name, cache in kv_caches.items() if name in forward_context - } + bindable_caches = map_kv_caches_to_kernel_blocks( + {name: cache for name, cache in kv_caches.items() if name in forward_context}, + attn_groups, + ) bind_kv_cache_to_layers( bindable_caches, forward_context, @@ -536,8 +540,8 @@ def build_attn_metadata( cached_metadata[reuse_key], block_table, slot_mapping ) elif for_cudagraph_capture: - metadata = attn_metadata_builder.build_for_cudagraph_capture( - common_attn_metadata + metadata = attn_group.build_metadata_for_cudagraph_capture( + common_attn_metadata, ubatch_idx ) else: attn_metadata_extra_kwargs = ( @@ -548,10 +552,8 @@ def build_attn_metadata( if model_specific_attn_metadata is not None else {} ) - metadata = attn_metadata_builder.build( - common_prefix_len=0, - common_attn_metadata=common_attn_metadata, - **attn_metadata_extra_kwargs, + metadata = attn_group.build_metadata( + common_attn_metadata, ubatch_idx, **attn_metadata_extra_kwargs ) if reuse_key is not None and reuse_key not in cached_metadata: cached_metadata[reuse_key] = metadata diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index c7665c46b876..1ab78864c31e 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -25,7 +25,6 @@ def __init__( max_num_batched_tokens: int, max_num_blocks_per_group: list[int], device: torch.device, - kernel_block_sizes: list[int], cp_size: int = 1, cp_rank: int = 0, cp_interleave: int = 1, @@ -33,7 +32,6 @@ def __init__( dcp_sharded: list[bool] | None = None, ): self.block_sizes = block_sizes - self.kernel_block_sizes = kernel_block_sizes self.max_num_reqs = max_num_reqs self.max_num_batched_tokens = max_num_batched_tokens self.device = device @@ -53,14 +51,9 @@ def __init__( assert len(dcp_sharded) == self.num_kv_cache_groups self.dcp_sharded = torch.tensor(dcp_sharded, dtype=torch.bool, device=device) - self.blocks_per_kv_block = [ - bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes) - ] - # num_kv_cache_groups x [max_num_reqs, max_num_blocks] self.block_tables: list[StagedWriteTensor] = [] - for i in range(self.num_kv_cache_groups): - max_num_blocks = max_num_blocks_per_group[i] * self.blocks_per_kv_block[i] + for max_num_blocks in max_num_blocks_per_group: block_table = StagedWriteTensor( (self.max_num_reqs, max_num_blocks), dtype=torch.int32, device=device ) @@ -109,9 +102,6 @@ def init_block_table_layout_tensors(self) -> None: self.block_sizes_tensor = torch.tensor( self.block_sizes, dtype=torch.int32, device=self.device ) - self.kernel_block_sizes_tensor = torch.tensor( - self.kernel_block_sizes, dtype=torch.int32, device=self.device - ) self.slot_mapping_enabled = torch.tensor( self._slot_mapping_enabled, dtype=torch.bool, device=self.device ) @@ -126,9 +116,6 @@ def append_block_ids( for i in range(self.num_kv_cache_groups): start = self.num_blocks.np[i, req_index] if not overwrite else 0 block_ids = new_block_ids[i] - bpk = self.blocks_per_kv_block[i] - if bpk > 1: - block_ids = [b * bpk + k for b in block_ids for k in range(bpk)] if self.redirect_writes_to_null_block: block_ids = [0] * len(block_ids) end = start + len(block_ids) @@ -218,7 +205,6 @@ def compute_slot_mappings( self.block_table_ptrs, self.block_table_strides, self.block_sizes_tensor, - self.kernel_block_sizes_tensor, self.slot_mapping_enabled, self.dcp_sharded, slot_mappings, @@ -292,7 +278,6 @@ def _compute_slot_mappings_kernel( block_table_ptrs, # [num_kv_cache_groups] block_table_strides, # [num_kv_cache_groups] block_sizes, # [num_kv_cache_groups] - kernel_block_sizes, # [num_kv_cache_groups] slot_mapping_enabled, # [num_kv_cache_groups] dcp_sharded, # [num_kv_cache_groups] slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens] @@ -322,7 +307,6 @@ def _compute_slot_mappings_kernel( block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32) block_table_stride = tl.load(block_table_strides + group_id) kv_block_size = tl.load(block_sizes + group_id) - kernel_block_size = tl.load(kernel_block_sizes + group_id) mapping_enabled = tl.load(slot_mapping_enabled + group_id) if CP_SIZE != 1: sharded = tl.load(dcp_sharded + group_id) @@ -354,16 +338,14 @@ def _compute_slot_mappings_kernel( local_positions = tl.where(sharded, local_positions, positions) is_local = ~sharded | is_local - block_indices = tl.where( - mapping_enabled, local_positions // kernel_block_size, 0 - ) - block_offsets = local_positions % kernel_block_size + block_indices = tl.where(mapping_enabled, local_positions // kv_block_size, 0) + block_offsets = local_positions % kv_block_size block_numbers = tl.load( block_table_ptr + req_state_idx * block_table_stride + block_indices, mask=is_local & is_real_req, other=0, ) - slot_ids = block_numbers * kernel_block_size + block_offsets + slot_ids = block_numbers * kv_block_size + block_offsets if CP_SIZE != 1: slot_ids = tl.where(is_local, slot_ids, PAD_ID) diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index e2f68ad59865..ea1c925dadc0 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -252,17 +252,18 @@ def set_dummy_context( input_batch.positions.copy_(torch.from_numpy(local_pos + context_len)) seq_len = context_len + query_len - for block_table, block_size, bpk in zip( - input_block_tables, - block_tables.kernel_block_sizes, - block_tables.blocks_per_kv_block, - ): + for block_table, block_size in zip(input_block_tables, block_tables.block_sizes): num_blocks = min(cdiv(seq_len, block_size), block_table.shape[1]) # Spans are disjoint until the pool runs out, then they wrap and share # blocks: profiling only needs the reads to be realistic, not distinct. - block_ids = torch.arange( - num_reqs * num_blocks, dtype=block_table.dtype, device=block_table.device - ) % (num_kv_blocks * bpk) + block_ids = ( + torch.arange( + num_reqs * num_blocks, + dtype=block_table.dtype, + device=block_table.device, + ) + % num_kv_blocks + ) block_table[:num_reqs, :num_blocks] = block_ids.view(num_reqs, num_blocks) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 61d6197ceb9d..134b4bfb0e69 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -688,7 +688,6 @@ def initialize_kv_cache( max_num_batched_tokens=self.max_num_tokens, max_num_blocks_per_group=max_num_blocks_per_group, device=self.device, - kernel_block_sizes=self.kernel_block_sizes, slot_mapping_enabled=slot_mapping_enabled, dcp_sharded=dcp_sharded, cp_size=self.dcp_size, @@ -773,6 +772,13 @@ def initialize_kv_cache( self.vllm_config, kv_cache_allocation_context=kv_cache_allocation_context, block_tables=self.block_tables, + # Draft groups build their own metadata, so they map the views too. + attn_groups=( + g + for groups in self.attn_groups + + getattr(self.speculator, "attn_groups", []) + for g in groups + ), ) self.kv_caches = [ cache for cache in kv_caches_dict.values() if cache.device == self.device diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index d59d87ae7e64..ec46039533d7 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -432,7 +432,7 @@ def propose( temperature, seeds, self.block_tables.input_block_tables[gid], - self.block_tables.kernel_block_sizes[gid], + self.block_tables.block_sizes[gid], self.block_tables.cp_rank, self.dcp_size, self.block_tables.cp_interleave, diff --git a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py index 9a8ed7b8b4ca..4894ce77e962 100644 --- a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py @@ -577,6 +577,11 @@ def _fused_multi_step_decode( idx_mapping = self.idx_mapping[:num_reqs] if batch_desc.cg_mode == CUDAGraphMode.FULL: + for groups in self.attn_groups: + for group in groups: + group.map_to_kernel_block_table( + self.block_tables.input_block_tables[group.kv_cache_group_id] + ) assert self.decode_cudagraph_manager is not None self.decode_cudagraph_manager.run_fullgraph(batch_desc) return diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 4707e4f2ca95..73a8c0eb9a59 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -96,7 +96,6 @@ def __init__( device: torch.device, vocab_size: int, block_sizes: list[int], # The block_size of each kv cache group - kernel_block_sizes: list[int], max_num_blocks_per_req: list[int], logitsprocs: LogitsProcessors | None = None, logitsprocs_need_output_token_ids: bool = False, @@ -187,7 +186,6 @@ def __init__( pin_memory=PIN_MEMORY, device=device, block_sizes=block_sizes, - kernel_block_sizes=kernel_block_sizes, max_num_blocks=max_num_blocks_per_req, cp_kv_cache_interleave_size=cp_kv_cache_interleave_size, slot_mapping_modes=slot_mapping_modes, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index c02bdc74e22b..22ad662b9018 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -233,6 +233,8 @@ allocate_kv_cache, bind_kv_cache, copy_kv_cache_blocks_inplace, + customize_attention_spec, + map_kv_caches_to_kernel_blocks, prepare_kernel_block_sizes, sanity_check_mm_encoder_outputs, ) @@ -718,7 +720,6 @@ def __init__( max(self.max_model_len, self.max_encoder_len), placeholder_block_size ) self._init_block_sizes = [placeholder_block_size] - self._init_kernel_block_sizes = [placeholder_block_size] self._init_max_num_blocks = [placeholder_max_num_blocks] self._init_slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT] self.cp_kv_cache_interleave_size = ( @@ -734,7 +735,6 @@ def __init__( device=self.device, vocab_size=self.model_config.get_vocab_size(), block_sizes=[placeholder_block_size], - kernel_block_sizes=[placeholder_block_size], max_num_blocks_per_req=[placeholder_max_num_blocks], num_spec_tokens=self.num_spec_tokens, logitsprocs=build_logitsprocs( @@ -2508,8 +2508,8 @@ def _build_attn_group_metadata( ) if for_cudagraph_capture: - attn_metadata_i = builder.build_for_cudagraph_capture( - common_attn_metadata + attn_metadata_i = attn_group.build_metadata_for_cudagraph_capture( + common_attn_metadata, ubid or 0 ) elif ( cache_key in cached_attn_metadata @@ -2517,13 +2517,16 @@ def _build_attn_group_metadata( ): attn_metadata_i = builder.update_block_table( cached_attn_metadata[cache_key], - common_attn_metadata.block_table_tensor, + attn_group.map_to_kernel_block_table( + common_attn_metadata.block_table_tensor, ubid or 0 + ), common_attn_metadata.slot_mapping, ) else: - attn_metadata_i = builder.build( + attn_metadata_i = attn_group.build_metadata( + common_attn_metadata, + ubid or 0, common_prefix_len=cascade_attn_prefix_len, - common_attn_metadata=common_attn_metadata, **extra_attn_metadata_args, ) if builder.supports_update_block_table: @@ -7143,9 +7146,7 @@ def calculate_reorder_batch_threshold(self) -> None: return self.reorder_batch_threshold = reduce(min_none_high, reorder_batch_thresholds) # type: ignore[assignment] - def may_reinitialize_input_batch( - self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int] - ) -> None: + def may_reinitialize_input_batch(self, kv_cache_config: KVCacheConfig) -> None: """Re-initialize the input batch if the block sizes are different from what it was originally created with. This happens when the final block size (determined after model loading) differs from the @@ -7154,7 +7155,6 @@ def may_reinitialize_input_batch( Args: kv_cache_config: The KV cache configuration. - kernel_block_sizes: The kernel block sizes for each KV cache group. """ block_sizes = [] @@ -7179,14 +7179,12 @@ def may_reinitialize_input_batch( if ( block_sizes != self._init_block_sizes - or kernel_block_sizes != self._init_kernel_block_sizes or max_num_blocks != self._init_max_num_blocks or slot_mapping_modes != self._init_slot_mapping_modes or self.cp_kv_cache_interleave_size != self.parallel_config.cp_kv_cache_interleave_size ): self._init_block_sizes = block_sizes - self._init_kernel_block_sizes = kernel_block_sizes self._init_max_num_blocks = max_num_blocks self._init_slot_mapping_modes = slot_mapping_modes self.cp_kv_cache_interleave_size = ( @@ -7201,7 +7199,6 @@ def may_reinitialize_input_batch( device=self.device, vocab_size=self.model_config.get_vocab_size(), block_sizes=block_sizes, - kernel_block_sizes=kernel_block_sizes, max_num_blocks_per_req=max_num_blocks, num_spec_tokens=self.num_spec_tokens, logitsprocs=self.input_batch.logitsprocs, @@ -7217,10 +7214,6 @@ def may_reinitialize_input_batch( f"InputBatch block_sizes {self._init_block_sizes} != " f"kv_cache block_sizes {block_sizes}" ) - assert self._init_kernel_block_sizes == kernel_block_sizes, ( - f"InputBatch kernel_block_sizes {self._init_kernel_block_sizes} " - f"!= kv_cache kernel_block_sizes {kernel_block_sizes}" - ) def _attn_group_iterator(self) -> Iterator[AttentionGroup]: return itertools.chain.from_iterable(self.attn_groups) @@ -7267,12 +7260,16 @@ def initialize_kv_cache_tensors( num_attn_module = ( 2 if self.model_config.hf_config.model_type == "longcat_flash" else 1 ) + attn_groups = [g for groups in self.attn_groups for g in groups] + if self.speculative_config: + attn_groups += getattr(self.drafter, "draft_attn_groups", []) bind_kv_cache( kv_caches, self.compilation_config.static_forward_context, self.kv_caches, num_attn_module, kv_cache_groups=kv_cache_config.kv_cache_groups, + layer_kv_caches=map_kv_caches_to_kernel_blocks(kv_caches, attn_groups), ) return kv_caches @@ -7336,6 +7333,8 @@ def initialize_kv_cache( # backends for that group only supports block_size 64, we will return # kernel_block_size 64 and split the 256-token-block to 4 blocks with 64 # tokens each. + # Block-outermost packed groups are not split; their attention groups + # map kernel blocks themselves. kernel_block_sizes = prepare_kernel_block_sizes( kv_cache_config, self.attn_groups ) @@ -7345,7 +7344,7 @@ def initialize_kv_cache( self.initialize_metadata_builders(kv_cache_config, kernel_block_sizes) # Reinitialize need to after initialize_attn_backend - self.may_reinitialize_input_batch(kv_cache_config, kernel_block_sizes) + self.may_reinitialize_input_batch(kv_cache_config) # Capture warmup providers that depend on allocated KV-cache strides. with self.jit_warmup_registry.activate(): # type: ignore[attr-defined] kv_caches = self.initialize_kv_cache_tensors( @@ -7422,7 +7421,9 @@ def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: # Skip modules that don't need KV cache (eg encoder-only attention) if spec := attn_module.get_kv_cache_spec(self.vllm_config): if isinstance(spec, AttentionSpec): - spec = attn_module.get_attn_backend().customize_spec(spec) + spec = customize_attention_spec( + attn_module.get_attn_backend(), spec + ) kv_cache_spec[layer_name] = spec return kv_cache_spec diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index db0f81b81c36..d10cdcec6453 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -3,7 +3,7 @@ import math from collections import defaultdict from collections.abc import Iterable, Mapping, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from itertools import product as iprod from typing import Any @@ -24,6 +24,7 @@ from vllm.v1.attention.backend import ( AttentionBackend, AttentionMetadataBuilder, + CommonAttentionMetadata, MultipleOf, ) from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy @@ -35,7 +36,6 @@ KVCacheLayout, KVCacheSpec, MambaSpec, - MLAAttentionSpec, UniformTypeKVCacheSpecs, create_kv_cache_views, ) @@ -166,6 +166,9 @@ def __init__( kv = static_forward_context[layer_name].kv_cache if not isinstance(kv, torch.Tensor): continue + if group.kernel_block_stride: + # Bound views address kernel blocks; zero by manager block. + kv = group.map_kv_cache_to_manager_blocks(kv, num_blocks) if kv.device.type != self.device.type: continue dp = kv.data_ptr() @@ -255,12 +258,22 @@ def warmup(self, num_kv_blocks: int) -> None: self.zero_block_ids([0]) -@dataclass +@dataclass(eq=False) class AttentionGroup: backend: type[AttentionBackend] layer_names: list[str] kv_cache_spec: KVCacheSpec kv_cache_group_id: int + # Set when the backend needs blocks smaller than the KV cache manager's: its + # block tables and KV cache views then address kernel blocks, kernel block j + # of block b being b * kernel_block_stride + j. The stride is set when a + # packed KV cache is bound (other layers' pages sit between its blocks); + # otherwise it stays 0 and the kernel blocks per block are used. + kernel_block_size: int | None = None + kernel_block_stride: int = 0 + # Persistent per ubatch: CUDA graphs capture the mapped block table. + kernel_block_table: torch.Tensor | None = None + kernel_block_offsets: torch.Tensor | None = None # When ubatching is enabled we will have a metadata builder for each ubatch # so that if they use internal persistent buffers for cudagraphs, and they # won't have to worry about conflicting with the other ubatches. @@ -275,27 +288,44 @@ def create_metadata_builders( kernel_block_size: int | None = None, num_metadata_builders: int = 1, ): + spec = self.kv_cache_spec + if kernel_block_size is not None and isinstance(spec, AttentionSpec): + # A backend in a packed KV cache group may need smaller blocks than + # the group is split into. + kernel_block_size = select_common_block_size( + kernel_block_size, [self.backend], [spec] + ) + if kernel_block_size != spec.block_size: + self.kernel_block_size = kernel_block_size if kernel_block_size is None: kv_cache_spec_builder = self.kv_cache_spec - elif ( - isinstance(self.kv_cache_spec, MLAAttentionSpec) - and self.kv_cache_spec.storage_block_size is not None - ): - kv_cache_spec_builder = self.kv_cache_spec.copy_with_new_block_size( - self.kv_cache_spec.storage_block_size - ) else: kv_cache_spec_builder = self.kv_cache_spec.copy_with_new_block_size( kernel_block_size ) builder_cls = self.backend.get_builder_cls() builder_kwargs = {} - if builder_cls.requires_block_table_width: + if builder_cls.requires_block_table_width or self.kernel_block_size is not None: max_num_blocks = self.kv_cache_spec.max_num_blocks_per_req( vllm_config, vllm_config.model_config.max_model_len ) - builder_kwargs["block_table_width"] = get_block_table_width( - max_num_blocks, self.kv_cache_spec.block_size, kernel_block_size + width = get_block_table_width(max_num_blocks, self.kv_cache_spec.block_size) + if kernel_block_size is not None: + width *= self.kv_cache_spec.block_size // kernel_block_size + if builder_cls.requires_block_table_width: + builder_kwargs["block_table_width"] = width + if self.kernel_block_size is not None: + self.kernel_block_table = torch.zeros( + num_metadata_builders, + vllm_config.scheduler_config.max_num_seqs, + width, + dtype=torch.int32, + device=device, + ) + self.kernel_block_offsets = torch.arange( + spec.block_size // self.kernel_block_size, + dtype=torch.int32, + device=device, ) self.metadata_builders = [ builder_cls( @@ -307,9 +337,109 @@ def create_metadata_builders( ) for _ in range(num_metadata_builders) ] - if kernel_block_size is not None: - for builder in self.metadata_builders: - builder.set_kernel_block_size(kernel_block_size) + + def build_metadata( + self, + common_attn_metadata: CommonAttentionMetadata, + ubatch_idx: int = 0, + common_prefix_len: int = 0, + **kwargs, + ) -> Any: + return self.get_metadata_builder(ubatch_idx).build( + common_prefix_len, + self._with_kernel_block_table(common_attn_metadata, ubatch_idx), + **kwargs, + ) + + def build_metadata_for_cudagraph_capture( + self, common_attn_metadata: CommonAttentionMetadata, ubatch_idx: int = 0 + ) -> Any: + return self.get_metadata_builder(ubatch_idx).build_for_cudagraph_capture( + self._with_kernel_block_table(common_attn_metadata, ubatch_idx) + ) + + def build_metadata_for_drafting( + self, common_attn_metadata: CommonAttentionMetadata, draft_index: int + ) -> Any: + return self.get_metadata_builder().build_for_drafting( + self._with_kernel_block_table(common_attn_metadata), draft_index + ) + + def _with_kernel_block_table( + self, common_attn_metadata: CommonAttentionMetadata, ubatch_idx: int = 0 + ) -> CommonAttentionMetadata: + if self.kernel_block_size is None: + return common_attn_metadata + return replace( + common_attn_metadata, + block_table_tensor=self.map_to_kernel_block_table( + common_attn_metadata.block_table_tensor, ubatch_idx + ), + ) + + def map_to_kernel_block_table( + self, block_table: torch.Tensor, ubatch_idx: int = 0 + ) -> torch.Tensor: + if self.kernel_block_size is None: + return block_table + assert self.kernel_block_table is not None + blocks_per_kv_block = self.kv_cache_spec.block_size // self.kernel_block_size + rows, cols = block_table.shape + out = self.kernel_block_table[ubatch_idx, :rows, : cols * blocks_per_kv_block] + torch.add( + self.kernel_block_offsets, + block_table.unsqueeze(-1), + alpha=self.kernel_block_stride or blocks_per_kv_block, + out=out.unflatten(1, (cols, -1)), + ) + return out + + def map_kv_cache_to_kernel_blocks(self, kv_cache: torch.Tensor) -> torch.Tensor: + """Views a packed cache in manager blocks as kernel blocks.""" + assert self.kernel_block_size is not None + spec = self.kv_cache_spec + kernel_rows = spec.get_num_kernel_states(self.kernel_block_size) + if kv_cache.shape[-2] == kernel_rows: + return kv_cache + # Token and pooled caches are written through manager-block slot + # mappings, so only compressed caches (written from their block table) + # can be mapped. + if ( + spec.tokens_per_state <= 1 + or kv_cache.shape[1] != 1 + or kv_cache.stride(0) % (kernel_rows * kv_cache.stride(-2)) + ): + raise ValueError( + f"{self.backend.get_name()} needs {self.kernel_block_size}-token " + f"kernel blocks, which packed {spec.block_size}-token KV cache " + f"blocks cannot be split into. Reduce --block-size to " + f"{self.kernel_block_size} or set VLLM_KV_CACHE_LAYOUT to a " + "layer-compact layout (e.g. LBNHC)." + ) + # Kernel block j of block b is b * kernel_block_stride + j. The view + # spans the other layers' pages between blocks (never addressing them), + # which ``view`` cannot express. + page_stride = kernel_rows * kv_cache.stride(-2) + self.kernel_block_stride = kv_cache.stride(0) // page_stride + num_pages = (kv_cache.shape[0] - 1) * self.kernel_block_stride + ( + kv_cache.shape[-2] // kernel_rows + ) + return kv_cache.as_strided( + (num_pages, *kv_cache.shape[1:-2], kernel_rows, kv_cache.shape[-1]), + (page_stride, *kv_cache.stride()[1:]), + ) + + def map_kv_cache_to_manager_blocks( + self, kv_cache: torch.Tensor, num_blocks: int + ) -> torch.Tensor: + """Views a packed cache in kernel blocks as manager blocks.""" + assert self.kernel_block_stride + spec = self.kv_cache_spec + rows = spec.get_num_kernel_states(spec.block_size) + return kv_cache.as_strided( + (num_blocks, *kv_cache.shape[1:-2], rows, kv_cache.shape[-1]), + (kv_cache.stride(0) * self.kernel_block_stride, *kv_cache.stride()[1:]), + ) def get_metadata_builder(self, ubatch_id: int = 0) -> AttentionMetadataBuilder: assert len(self.metadata_builders) > ubatch_id @@ -327,17 +457,48 @@ def update_draft_decode_metadata( self.get_metadata_builder().update_draft_decode_metadata(metadata) +def map_kv_caches_to_kernel_blocks( + kv_caches: dict[str, torch.Tensor], attn_groups: Iterable[AttentionGroup] +) -> dict[str, torch.Tensor]: + """The views to bind to layers. The manager views in ``kv_caches`` stay for + KV connectors.""" + mapped = dict(kv_caches) + for group in attn_groups: + if group.kernel_block_size is not None: + mapped.update( + (name, group.map_kv_cache_to_kernel_blocks(kv_caches[name])) + for name in group.layer_names + if name in kv_caches + ) + return mapped + + +def customize_attention_spec( + backend: type[AttentionBackend], spec: AttentionSpec +) -> AttentionSpec: + """The backend's spec. Blocks larger than every kernel block size a backend + takes are split into kernel blocks, which packed blocks must keep whole.""" + spec = backend.customize_spec(spec) + sizes = backend.get_supported_kernel_block_sizes(spec) + fixed_sizes = [s for s in sizes if isinstance(s, int)] + if spec.block_stride_alignment is None and sizes and fixed_sizes == sizes: + spec = replace(spec, block_stride_alignment=MultipleOf(max(fixed_sizes))) + return spec + + def _block_size_is_supported( - backends: list[type[AttentionBackend]], block_size: int + backends: list[type[AttentionBackend]], + block_size: int, + specs: Sequence[KVCacheSpec | None], ) -> bool: """Check if the block size is supported by all backends. An exact ``int`` declaration must match exactly; a ``MultipleOf`` declaration accepts any multiple of its base. """ - for backend in backends: + for backend, spec in zip(backends, specs, strict=True): is_supported = False - for supported_size in backend.get_supported_kernel_block_sizes(): + for supported_size in backend.get_supported_kernel_block_sizes(spec): if isinstance(supported_size, int): if block_size == supported_size: is_supported = True @@ -354,6 +515,7 @@ def _block_size_is_supported( def select_common_block_size( kv_manager_block_size: int, backends: list[type[AttentionBackend]], + kv_cache_specs: Sequence[KVCacheSpec | None] | None = None, ) -> int: """Select a block size that is supported by all backends and is a factor of kv_manager_block_size. @@ -364,6 +526,7 @@ def select_common_block_size( Args: kv_manager_block_size: Block size of KV cache. backends: List of attention backend classes. + kv_cache_specs: Per-backend cache specs for spec-specific constraints. Returns: The selected block size. @@ -372,25 +535,27 @@ def select_common_block_size( ValueError: If no valid block size found. """ - if _block_size_is_supported(backends, kv_manager_block_size): + specs = kv_cache_specs or [None] * len(backends) + if _block_size_is_supported(backends, kv_manager_block_size, specs): return kv_manager_block_size # MultipleOf constraints also accept the manager size if they accept a divisor. # Any remaining candidate must therefore be an explicit size from a backend. candidates = { size - for backend in backends - for size in backend.get_supported_kernel_block_sizes() + for backend, spec in zip(backends, specs, strict=True) + for size in backend.get_supported_kernel_block_sizes(spec) if isinstance(size, int) and kv_manager_block_size % size == 0 } for size in sorted(candidates, reverse=True): - if _block_size_is_supported(backends, size): + if _block_size_is_supported(backends, size, specs): return size raise ValueError( f"No common block size for {kv_manager_block_size} (" + "; ".join( - f"{b.get_name()}: {b.get_supported_kernel_block_sizes()}" for b in backends + f"{b.get_name()}: {b.get_supported_kernel_block_sizes(s)}" + for b, s in zip(backends, specs, strict=True) ) + ")." ) @@ -450,9 +615,6 @@ def allocate_kv_cache( kernel_block_size = None if kernel_block_sizes is not None and group_id < len(kernel_block_sizes): kernel_block_size = kernel_block_sizes[group_id] - if isinstance(spec, MLAAttentionSpec) and spec.storage_block_size is not None: - kernel_block_size = spec.storage_block_size - views = create_kv_cache_views( buf, spec, @@ -497,24 +659,20 @@ def prepare_kernel_block_sizes( # This is an attention backend that supports virtual block splitting. kv_manager_block_size = kv_cache_group.kv_cache_spec.block_size group_backends = [g.backend for g in attn_groups[kv_cache_gid]] - storage_block_size = ( - kv_cache_spec.storage_block_size - if isinstance(kv_cache_spec, MLAAttentionSpec) - else None - ) - if storage_block_size is not None and _block_size_is_supported( - group_backends, storage_block_size - ): - # Storage-block specs (e.g. the GLM-5.3-Flash kpool indexer - # cache) address the cache in pool pages, and every other - # consumer (cache views, metadata builders, hisparse) already - # uses storage_block_size as the kernel block. Fall back to the - # backend vote when the group's backends do not accept it. - selected_kernel_size = storage_block_size - else: - selected_kernel_size = select_common_block_size( - kv_manager_block_size, group_backends + group_specs = [g.kv_cache_spec for g in attn_groups[kv_cache_gid]] + # Block-outermost packing puts other layers' pages between the + # blocks, so they are never split group-wide. + selected_kernel_size = ( + kv_manager_block_size + if any( + t.block_stride > t.layer_stride + for t in kv_cache_config.kv_cache_tensors + if not set(t.layers).isdisjoint(kv_cache_group.layer_names) ) + else select_common_block_size( + kv_manager_block_size, group_backends, group_specs + ) + ) kernel_block_sizes.append(selected_kernel_size) elif isinstance(kv_cache_spec, MambaSpec): # This is likely Mamba or other non-attention cache, no splitting. @@ -656,6 +814,7 @@ def bind_kv_cache( runner_kv_caches: list[torch.Tensor], num_attn_module: int = 1, kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None, + layer_kv_caches: dict[str, torch.Tensor] | None = None, ) -> None: """Bind the allocated KV cache to both ModelRunner and forward context so that the KV cache can be used in the forward pass. @@ -674,6 +833,9 @@ def bind_kv_cache( runner_kv_caches: The kv_cache declared by ModelRunner. kv_cache_groups: The KV cache groups of the model, used to resolve layers that share a KV cache. + layer_kv_caches: The views to bind to layers when they differ from + kv_caches (e.g. kernel-block views). The runner keeps kv_caches, + which are addressed by scheduler block IDs. """ # Bind kv_caches to ModelRunner @@ -699,7 +861,7 @@ def bind_kv_cache( runner_kv_caches.append(kv_caches[layer_name]) bind_kv_cache_to_layers( - kv_caches, forward_context, num_attn_module, kv_cache_groups + layer_kv_caches or kv_caches, forward_context, num_attn_module, kv_cache_groups ) From e0cb3efce05388e21454556193d4069984b3154a Mon Sep 17 00:00:00 2001 From: Kai Hoang <56827800+hpdkhoa@users.noreply.github.com> Date: Sat, 10 Oct 2026 03:52:08 +0700 Subject: [PATCH 203/624] [Bugfix][Structured Output] Decode choice specs as JSON in outlines and LMFE (#60748) Signed-off-by: hpdkhoa Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../test_structured_output_manager.py | 18 +++++++++++++++++- .../backend_lm_format_enforcer.py | 3 +-- vllm/v1/structured_output/backend_outlines.py | 3 +-- 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/tests/v1/structured_output/test_structured_output_manager.py b/tests/v1/structured_output/test_structured_output_manager.py index fbceb3bd07e4..67c15825580c 100644 --- a/tests/v1/structured_output/test_structured_output_manager.py +++ b/tests/v1/structured_output/test_structured_output_manager.py @@ -16,7 +16,8 @@ from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager -from vllm.v1.structured_output.backend_outlines import OutlinesGrammar +from vllm.v1.structured_output.backend_outlines import OutlinesBackend, OutlinesGrammar +from vllm.v1.structured_output.request import get_structured_output_key TOKENIZER = "gpt2" THINK_END = "\n" # reasoning-end marker (single GPT-2 token) @@ -560,3 +561,18 @@ def test_outlines_termination(tokenizer): assert grammar.validate_tokens([eos]) == [] assert grammar.accept_tokens(request.request_id, [one, eos, one]) assert grammar.is_terminated() + + +def test_outlines_choice_with_non_bmp_characters(tokenizer): + """The choice spec is JSON, which encodes emoji as surrogate pairs.""" + choices = ["😀 yes", "no"] + backend = OutlinesBackend( + VllmConfig(), tokenizer=tokenizer, vocab_size=len(tokenizer) + ) + request_type, grammar_spec = get_structured_output_key( + StructuredOutputsParams(choice=choices) + ) + + grammar = backend.compile_grammar(request_type, grammar_spec) + + assert grammar.accept_tokens("", tokenizer.encode(choices[0])) diff --git a/vllm/v1/structured_output/backend_lm_format_enforcer.py b/vllm/v1/structured_output/backend_lm_format_enforcer.py index 36eb5eb159e1..9a16ddc1fe38 100644 --- a/vllm/v1/structured_output/backend_lm_format_enforcer.py +++ b/vllm/v1/structured_output/backend_lm_format_enforcer.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import ast import json from dataclasses import dataclass, field from functools import lru_cache @@ -117,7 +116,7 @@ def compile_grammar( grammar_spec, ) elif request_type == StructuredOutputOptions.CHOICE: - choices = ast.literal_eval(grammar_spec) + choices = json.loads(grammar_spec) character_level_parser = lmformatenforcer.UnionParser( [lmformatenforcer.StringParser(choice) for choice in choices] ) diff --git a/vllm/v1/structured_output/backend_outlines.py b/vllm/v1/structured_output/backend_outlines.py index 1cfc6c8f5d1a..176d1a46e684 100644 --- a/vllm/v1/structured_output/backend_outlines.py +++ b/vllm/v1/structured_output/backend_outlines.py @@ -3,7 +3,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from __future__ import annotations -import ast import json from dataclasses import dataclass, field from re import _constants, _parser # type: ignore[attr-defined] @@ -68,7 +67,7 @@ def compile_grammar( elif request_type == StructuredOutputOptions.REGEX: regex = grammar_spec elif request_type == StructuredOutputOptions.CHOICE: - choices = ast.literal_eval(grammar_spec) + choices = json.loads(grammar_spec) choices = [regex_escape(c) for c in choices] regex = "(" + "|".join(choices) + ")" else: From ffdde24f546c503c7e47c5264b180c16fd668888 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Fri, 9 Oct 2026 13:54:40 -0700 Subject: [PATCH 204/624] [Docs] Document adaptive verification without a confidence head (#60913) Signed-off-by: Giancarlo Delfin --- .../adaptive_verification.md | 25 ++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/docs/features/speculative_decoding/adaptive_verification.md b/docs/features/speculative_decoding/adaptive_verification.md index e9314c0e09f0..d560b5af3e4f 100644 --- a/docs/features/speculative_decoding/adaptive_verification.md +++ b/docs/features/speculative_decoding/adaptive_verification.md @@ -4,7 +4,7 @@ Speculative decoding buys fewer decode steps with more compute. At batch size 1 That matters because per-position acceptance decays fast. While the GPU is memory-bound that slot is effectively free and worth the gamble; once it saturates the gamble has a real throughput cost. The crossover moves with load and with workload-dependent acceptance rates, so no static `num_speculative_tokens` is right across concurrencies. -Adaptive verification decides per step how much of the draft to verify instead. Every (request, position) draft slot is scored by its *survival probability*, the running product of that request's per-position confidences, and the highest-scoring slots are admitted until a global budget is spent. Slots compete across requests: position 5 of a confident request can outrank position 1 of a doubtful one, so one request keeps its full block while another could be trimmed after a token or two. +Adaptive verification decides per step how much of the draft to verify instead. Every (request, position) draft slot is scored by its *survival probability*, the running product of that request's per-position acceptance estimates, and the highest-scoring slots are admitted until a global budget is spent. Slots compete across requests: position 5 of a confident request can outrank position 1 of a doubtful one, so one request keeps its full block while another could be trimmed after a token or two. The budget itself comes from a cost model profiled at startup. vLLM measures what a step costs at each shape, then picks the token count that maximizes expected accepted tokens per second. @@ -12,11 +12,14 @@ The practical effect is that one configuration holds up across the whole load ra ## Support -Adaptive verification needs per-position acceptance estimates, so today it is only supported for DSpark with a **confidence head**. +Adaptive verification is supported for model-based drafters (EAGLE, MTP, DFlash, DSpark and standalone draft models). The per-position acceptance estimates come from one of two sources: + +- **Confidence head**: DSpark checkpoints that ship a trained confidence head use its predictions directly. +- **Online acceptance estimator**: every other drafter predicts acceptance from its draft logits with a small logistic model, recalibrated during serving against the target's verification results. ## Usage -It is off by default. Enable it in the speculative config: +It is off by default. Enable it in the speculative config with `enable_adaptive_verification`. For example, with DSpark: ```bash vllm serve deepseek-ai/DeepSeek-V4-Flash-DSpark \ @@ -30,6 +33,21 @@ vllm serve deepseek-ai/DeepSeek-V4-Flash-DSpark \ }' ``` +Or with DFlash on MiMo-V2.5-Pro, whose drafter ships in the checkpoint's `dflash/` subdirectory: + +```bash +hf download XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash --local-dir MiMo-V2.5-Pro-FP4-DFlash + +vllm serve MiMo-V2.5-Pro-FP4-DFlash \ + --tensor-parallel-size 4 --trust-remote-code --max-model-len auto \ + --speculative-config '{ + "method": "dflash", + "model": "MiMo-V2.5-Pro-FP4-DFlash/dflash", + "num_speculative_tokens": 8, + "enable_adaptive_verification": true + }' +``` + Set `enable_adaptive_verification: false` to verify the full block for every request. ## Requirements and limitations @@ -37,6 +55,7 @@ Set `enable_adaptive_verification: false` to verify the full block for every req - The attention backend must tolerate device-decided query lengths, since the CPU lengths only bound them from above. Backends that plan off the CPU lengths are excluded by the attention selector, and rejected at startup for models that hard-wire their backend. - Full cudagraphs are required: step costs are profiled from captured graphs, so `--enforce-eager` is rejected at startup. - Not supported with LoRA (the per-token LoRA mapping is built from CPU-side boundaries) or pipeline parallelism (cost curves and confidences exist only on the last rank). +- Without a confidence head, `use_local_argmax_reduction` is rejected. The estimator needs the full draft logits, which local argmax reduction never materializes. ## Tuning the cost profile From 35e920d671539d7066b648dc82058f8ef5348b9f Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Sat, 10 Oct 2026 05:10:42 +0800 Subject: [PATCH 205/624] [Rust Frontend] Support `--stream-interval` and per-request `stream_interval` (#60589) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- rust/src/cmd/src/cli.rs | 14 + rust/src/cmd/src/cli/tests.rs | 34 ++ rust/src/cmd/src/cli/unsupported.rs | 4 - .../examples/external_engine_logprobs.rs | 3 +- rust/src/engine-core-client/src/client.rs | 21 +- rust/src/engine-core-client/src/client/imp.rs | 73 ++-- .../src/client/output_channel.rs | 321 ++++++++++++++++++ .../engine-core-client/src/client/state.rs | 286 ++++++++++++---- .../engine-core-client/src/client/stream.rs | 83 ++++- rust/src/engine-core-client/src/lib.rs | 2 +- .../src/protocol/sampling.rs | 7 + .../engine-core-client/src/tests/client.rs | 43 +-- rust/src/llm/src/lib.rs | 23 +- rust/src/llm/src/output.rs | 79 ++++- rust/src/llm/tests/generate.rs | 106 ++++++ rust/src/mock-engine/src/tests.rs | 10 +- .../examples/external_engine_openai_qwen.rs | 2 + rust/src/server/src/config.rs | 5 + rust/src/server/src/lib.rs | 4 +- .../routes/openai/chat_completions/convert.rs | 13 + .../routes/openai/chat_completions/types.rs | 8 + .../src/routes/openai/completions/convert.rs | 10 +- .../src/routes/openai/completions/types.rs | 7 + rust/src/text/src/lower.rs | 29 ++ rust/src/text/src/request.rs | 7 + 25 files changed, 1040 insertions(+), 154 deletions(-) create mode 100644 rust/src/engine-core-client/src/client/output_channel.rs diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index d3e369bbe388..14b949c87de4 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -12,6 +12,7 @@ mod unsupported; use std::collections::HashMap; use std::ffi::{OsStr, OsString}; +use std::num::NonZeroU32; use std::path::PathBuf; use std::time::Duration; @@ -268,6 +269,13 @@ pub struct SharedRuntimeArgs { #[arg(long, value_parser = clap::value_parser!(i32).range(-1..), allow_negative_numbers = true)] #[serde(default)] pub max_logprobs: Option, + /// The interval (or buffer size) for streaming in terms of token length. + /// A smaller value (1) makes streaming smoother by sending each token + /// immediately, while a larger value (e.g., 10) reduces host overhead and + /// may increase throughput by batching multiple tokens before sending. + #[arg(long, default_value_t = NonZeroU32::MIN)] + #[serde(default = "default_stream_interval")] + pub stream_interval: NonZeroU32, /// Maximum time to wait for active requests to drain during shutdown. #[arg(long, default_value_t = 0)] #[serde(default)] @@ -547,6 +555,7 @@ impl SharedRuntimeArgs { lora_modules: self.lora_modules, chat_template_content_format: self.chat_template_content_format, max_logprobs: self.max_logprobs, + stream_interval: self.stream_interval, api_server_options, cors, tls, @@ -609,6 +618,7 @@ impl SharedRuntimeArgs { lora_modules: self.lora_modules, chat_template_content_format: self.chat_template_content_format, max_logprobs: self.max_logprobs, + stream_interval: self.stream_interval, api_server_options, cors, tls, @@ -646,6 +656,10 @@ fn default_engine_ready_timeout_secs() -> u64 { 600 } +fn default_stream_interval() -> NonZeroU32 { + NonZeroU32::MIN +} + fn default_cors_wildcard() -> JsonStringList { JsonStringList(vec!["*".to_string()]) } diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index aa164bc618f8..57780027deac 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -189,6 +189,7 @@ fn serve_args_forward_python_flags_with_separator() { renderer: Auto, language_model_only: false, max_logprobs: None, + stream_interval: 1, shutdown_timeout: 0, http_timeout_keep_alive: None, chat_template: None, @@ -373,6 +374,34 @@ fn serve_args_forward_profiler_config_to_managed_engine() { ); } +#[test] +fn serve_args_forward_stream_interval_to_frontend_config() { + let cli = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--stream-interval", + "4", + ]) + .unwrap(); + + let Command::Serve(args) = cli.command else { + panic!("expected serve args"); + }; + let frontend_config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string()); + assert_eq!(frontend_config.stream_interval.get(), 4); + + let error = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--stream-interval", + "0", + ]) + .unwrap_err(); + assert_eq!(error.kind(), clap::error::ErrorKind::ValueValidation); +} + #[test] fn serve_args_forward_max_logprobs_to_frontend_and_managed_engine() { let cli = Cli::try_parse_from([ @@ -1030,6 +1059,7 @@ fn frontend_args_accept_json() { renderer: Auto, language_model_only: false, max_logprobs: None, + stream_interval: 1, shutdown_timeout: 0, http_timeout_keep_alive: None, chat_template: None, @@ -1702,6 +1732,7 @@ fn serve_args_accept_handshake_aliases() { renderer: Auto, language_model_only: false, max_logprobs: None, + stream_interval: 1, shutdown_timeout: 0, http_timeout_keep_alive: None, chat_template: None, @@ -1866,6 +1897,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { lora_modules: [], chat_template_content_format: Auto, max_logprobs: None, + stream_interval: 1, api_server_options: ApiServerOptions { enable_log_requests: false, enable_prompt_tokens_details: false, @@ -1961,6 +1993,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { lora_modules: [], chat_template_content_format: Auto, max_logprobs: None, + stream_interval: 1, api_server_options: ApiServerOptions { enable_log_requests: false, enable_prompt_tokens_details: false, @@ -2100,6 +2133,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present lora_modules: [], chat_template_content_format: Auto, max_logprobs: None, + stream_interval: 1, api_server_options: ApiServerOptions { enable_log_requests: false, enable_prompt_tokens_details: false, diff --git a/rust/src/cmd/src/cli/unsupported.rs b/rust/src/cmd/src/cli/unsupported.rs index f65f766e847c..0a14d1b841c4 100644 --- a/rust/src/cmd/src/cli/unsupported.rs +++ b/rust/src/cmd/src/cli/unsupported.rs @@ -359,10 +359,6 @@ pub struct EngineUnsupportedArgs { #[arg(long)] pub collect_detailed_traces: Option, - /// The interval (or buffer size) for streaming in terms of token length. - #[arg(long)] - pub stream_interval: Option, - /// Maximum number of requests that can be in-flight (waiting or running) /// at the same time. When the limit is reached, new requests are rejected /// with HTTP 503 so the client can retry on another instance. diff --git a/rust/src/engine-core-client/examples/external_engine_logprobs.rs b/rust/src/engine-core-client/examples/external_engine_logprobs.rs index d4ec116f173f..e0da15623ca3 100644 --- a/rust/src/engine-core-client/examples/external_engine_logprobs.rs +++ b/rust/src/engine-core-client/examples/external_engine_logprobs.rs @@ -87,8 +87,9 @@ fn build_request( } async fn wait_for_final_output( - mut stream: vllm_engine_core_client::EngineCoreOutputStream, + stream: vllm_engine_core_client::EngineCoreOutputStream, ) -> Result { + let mut stream = stream.into_outputs(); while let Some(output) = stream.next().await { let output = output.context("failed to receive engine-core output")?; if output.finished() { diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index b3516936fef2..d8b436c6466c 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::BTreeMap; +use std::num::NonZeroU32; use std::sync::Arc; use std::time::Duration; @@ -25,10 +26,11 @@ use crate::runtime::{BackgroundShutdownRuntime, build_zmq_runtime}; use crate::transport::{self, ConnectedEngine}; pub(crate) mod imp; +mod output_channel; mod state; mod stream; -pub use stream::{EngineCoreOutputStream, EngineCoreStreamOutput}; +pub use stream::{EngineCoreOutputStream, EngineCoreStreamDelivery, EngineCoreStreamOutput}; /// How the frontend acquires its request/response transport with Python /// `EngineCoreProc`s. @@ -614,6 +616,12 @@ fn validate_lora_capabilities(engines: &[ConnectedEngine]) -> Result<()> { impl EngineCoreClient { /// Add a new request to the engine and return a per-request raw output /// stream. + /// + /// With `sampling_params.stream_interval` above one, outputs after the + /// first are delivered to the stream in batches of at least that many new + /// tokens, and the terminal output is delivered immediately together with + /// any held-back outputs. The stream still yields every raw output in + /// order; batching only reduces how often the consuming task is woken. pub async fn call(&self, mut req: EngineCoreRequest) -> Result { req.client_index = self.config.client_index; req.validate()?; @@ -628,8 +636,15 @@ impl EngineCoreClient { let request_id = req.request_id.clone(); let lora_name = req.lora_request.as_ref().map(|lora| lora.lora_name.clone()); let data_parallel_rank = req.data_parallel_rank; - let (engine_id, rx) = - self.inner.register_request(request_id.clone(), lora_name, data_parallel_rank)?; + let stream_interval = (req.sampling_params.as_ref()) + .and_then(|params| params.stream_interval) + .unwrap_or(NonZeroU32::MIN); + let (engine_id, rx) = self.inner.register_request( + request_id.clone(), + lora_name, + data_parallel_rank, + stream_interval, + )?; // Construct the output stream first before actually sending the request to the engine. // This ensures that cancelling the future will properly clean up the request registration diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index c8abb95bbdf8..0a4af1a5ee5f 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::{BTreeMap, BTreeSet}; +use std::num::NonZeroU32; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; @@ -15,7 +16,9 @@ use tracing::{debug, info, trace, warn}; use vllm_metrics::METRICS; use zeromq::RouterSendHalf; -use crate::client::state::{OutputReceiver, RequestRegistry, UtilityReceiver, UtilityRegistry}; +use crate::client::state::{ + OutputReceiver, OutputSender, RequestRegistry, UtilityReceiver, UtilityRegistry, +}; use crate::client::stream::EngineCoreStreamOutput; use crate::client::{AbortCause, AbortRequest}; use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output}; @@ -111,12 +114,13 @@ impl ClientInner { request_id: String, lora_name: Option, data_parallel_rank: Option, + stream_interval: NonZeroU32, ) -> Result<(EngineId, OutputReceiver)> { let mut registry = self.request_reg.lock(); if registry.is_closed() { return Err(self.closed_error()); } - registry.register(request_id, lora_name, data_parallel_rank) + registry.register(request_id, lora_name, data_parallel_rank, stream_interval) } /// Allocate the next utility `call_id` and register its waiting receiver. @@ -159,13 +163,15 @@ impl ClientInner { Ok(registry.abortable_request_ids(request_ids)) } - /// Obtain stream senders for a whole engine output batch with one registry - /// lock acquisition. - pub fn take_senders_for_outputs<'a>( + /// Send a whole engine output batch to the request streams with one + /// registry lock acquisition. Returns the outputs whose requests are no + /// longer tracked. + pub fn send_outputs( &self, - outputs: impl IntoIterator, - ) -> Vec>>> { - self.request_reg.lock().senders_for_outputs(outputs) + outputs: impl IntoIterator, + ) -> Vec { + let mut registry = self.request_reg.lock(); + outputs.into_iter().filter_map(|output| registry.send_output(output)).collect() } /// Remove a batch of requests that have finished or aborted, returning @@ -173,7 +179,7 @@ impl ClientInner { pub fn finish_requests<'a>( &self, request_ids: impl IntoIterator, - ) -> Vec>> { + ) -> Vec { self.request_reg.lock().finish_many(request_ids) } @@ -214,7 +220,7 @@ impl ClientInner { // Notify all ongoing requests that the client is closed. for sender in request_senders { - let _ = sender.send(Err(Error::Shared(persistent_error.clone()))); + sender.send_error(Error::Shared(persistent_error.clone())); } for sender in utility_senders { let _ = sender.send(Err(Error::Shared(persistent_error.clone()))); @@ -413,6 +419,16 @@ pub(crate) async fn run_abort_loop( } } +/// Number of tokens one output accounts for in `vllm:iteration_tokens_total`. +fn output_iteration_tokens(output: &EngineCoreOutput) -> u64 { + let mut tokens = output.new_token_ids.len() as u64; + if let Some(prefill) = &output.prefill_stats { + // The engine emits prefill_stats once, on the first output. + tokens += u64::from(prefill.num_computed_tokens); + } + tokens +} + /// Background loop that listens for engine-core outputs and dispatches them to /// the corresponding request streams based on their `request_id`. pub(crate) async fn run_output_dispatcher_loop( @@ -435,29 +451,20 @@ pub(crate) async fn run_output_dispatcher_loop( match outputs { EngineCoreOutputs::RequestBatch(batch) => { let has_outputs = !batch.outputs.is_empty(); - let mut iteration_tokens = 0_u64; - let senders = inner.take_senders_for_outputs(&batch.outputs); - for (output, sender) in batch.outputs.into_iter().zip(senders) { - let request_id = output.request_id.clone(); - let Some(sender) = sender else { - debug!(request_id, "dropping output for inactive request"); - continue; - }; - - iteration_tokens += output.new_token_ids.len() as u64; - if let Some(prefill) = &output.prefill_stats { - // The engine emits prefill_stats once, on the first output. - iteration_tokens += u64::from(prefill.num_computed_tokens); - } - - let wrapped_output = EngineCoreStreamOutput { - engine_index: batch.engine_index, - timestamp: batch.timestamp, - output, - }; - if sender.send(Ok(wrapped_output)).is_err() { - debug!(request_id, "request output stream receiver dropped"); - } + let mut iteration_tokens: u64 = + batch.outputs.iter().map(output_iteration_tokens).sum(); + let (engine_index, timestamp) = (batch.engine_index, batch.timestamp); + let outputs = batch.outputs.into_iter().map(|output| EngineCoreStreamOutput { + engine_index, + timestamp, + output, + }); + for output in inner.send_outputs(outputs) { + debug!( + request_id = output.request_id, + "dropping output for inactive request" + ); + iteration_tokens -= output_iteration_tokens(&output); } if has_outputs diff --git a/rust/src/engine-core-client/src/client/output_channel.rs b/rust/src/engine-core-client/src/client/output_channel.rs new file mode 100644 index 000000000000..7b58d3dc2f1f --- /dev/null +++ b/rust/src/engine-core-client/src/client/output_channel.rs @@ -0,0 +1,321 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Per-request output channel between the output dispatcher and +//! [`EngineCoreOutputStream`](crate::EngineCoreOutputStream). +//! +//! With a stream interval above one, the sender holds back the outputs after +//! the first one until they carry at least that many new tokens, then delivers +//! them in one channel message, which the stream yields as one +//! [`EngineCoreStreamDelivery`](crate::EngineCoreStreamDelivery). The request +//! task is then woken once per delivery instead of once per engine step. +//! +//! Outputs are never merged, reordered, or dropped here; the consumer merges a +//! delivery into one output (see `vllm-llm`). Merging here would hide the raw +//! outputs that request metrics observe one by one (per-step timestamps for +//! inter-token latency, prefill stats, lifecycle events), and would have to +//! combine wire-level fields such as undecoded logprobs tensors and opaque +//! values. + +use std::num::NonZeroU32; + +use tokio::sync::mpsc; +use tracing::{debug, warn}; + +use crate::client::stream::{EngineCoreStreamDelivery as Delivery, EngineCoreStreamOutput}; +use crate::error::{Error, Result}; + +/// Create the output channel for one request. +pub(crate) fn output_channel(stream_interval: NonZeroU32) -> (OutputSender, OutputReceiver) { + let (tx, rx) = mpsc::unbounded_channel(); + let sender = OutputSender { + tx, + stream_interval: stream_interval.get() as usize, + delivered_first: false, + buffered: Vec::new(), + buffered_tokens: 0, + }; + (sender, rx) +} + +/// Receiving half of a request's output channel. +pub(crate) type OutputReceiver = mpsc::UnboundedReceiver>; + +/// Sending half of a request's output channel, owned by its registry entry. +#[derive(Debug)] +pub(crate) struct OutputSender { + tx: mpsc::UnboundedSender>, + /// Minimum number of new tokens per delivery after the first one. + stream_interval: usize, + /// Whether the first output has been delivered. + delivered_first: bool, + /// Outputs held back since the last delivery. + buffered: Vec, + /// Number of new tokens in `buffered`. + buffered_tokens: usize, +} + +impl OutputSender { + /// Send one output, or hold it back until the held-back outputs reach the + /// stream interval. The first output and the terminal output are always + /// delivered immediately, the latter together with any held-back outputs. + pub fn send(&mut self, output: EngineCoreStreamOutput) { + let finished = output.finished(); + if self.buffered.is_empty() + && (finished || !self.delivered_first || self.stream_interval == 1) + { + self.delivered_first = true; + self.deliver(Delivery::One(output)); + return; + } + + self.buffered_tokens += output.new_token_ids.len(); + self.buffered.push(output); + if finished || self.buffered_tokens >= self.stream_interval { + self.flush(); + } + } + + /// Deliver the held-back outputs, if any. + pub fn flush(&mut self) { + if self.buffered.is_empty() { + return; + } + self.buffered_tokens = 0; + let outputs = std::mem::take(&mut self.buffered); + self.deliver(Delivery::Many(outputs)); + } + + /// Deliver the held-back outputs, then terminate the stream with `error`. + pub fn send_error(mut self, error: Error) { + self.flush(); + let _ = self.tx.send(Err(error)); + } + + fn deliver(&self, delivery: Delivery) { + if let Err(mpsc::error::SendError(Ok(delivery))) = self.tx.send(Ok(delivery)) { + debug!( + request_id = delivery.first().request_id, + "request output stream receiver dropped" + ); + } + } +} + +impl Drop for OutputSender { + fn drop(&mut self) { + // Every path that ends a live request delivers held-back outputs first + // (terminal output, `flush`, or `send_error`). Outputs left here are lost. + if !self.buffered.is_empty() && !self.tx.is_closed() { + warn!( + request_id = self.buffered[0].request_id, + count = self.buffered.len(), + "dropping request output sender with held-back outputs" + ); + } + } +} + +#[cfg(test)] +mod tests { + use std::num::NonZeroU32; + + use expect_test::expect; + use tokio::sync::mpsc::error::TryRecvError; + + use super::*; + use crate::protocol::output::{EngineCoreFinishReason, EngineCoreOutput}; + + fn output(token_ids: &[u32]) -> EngineCoreStreamOutput { + EngineCoreStreamOutput { + engine_index: 0, + timestamp: 0.0, + output: EngineCoreOutput { + request_id: "req".to_string(), + new_token_ids: token_ids.to_vec(), + ..EngineCoreOutput::default() + }, + } + } + + fn finished(token_ids: &[u32]) -> EngineCoreStreamOutput { + let mut output = output(token_ids); + output.output.finish_reason = Some(EngineCoreFinishReason::Length); + output + } + + fn channel(stream_interval: u32) -> (OutputSender, OutputReceiver) { + output_channel(NonZeroU32::new(stream_interval).unwrap()) + } + + /// Render a delivery as the token ids of its outputs. + fn token_ids(delivery: Delivery) -> Vec> { + delivery.into_iter().map(|output| output.output.new_token_ids).collect() + } + + /// Drain all pending channel messages, rendering each delivery as the + /// token ids of its outputs. + fn deliveries(receiver: &mut OutputReceiver) -> Vec>> { + let mut deliveries = Vec::new(); + loop { + match receiver.try_recv() { + Ok(Ok(delivery)) => deliveries.push(token_ids(delivery)), + Ok(Err(error)) => panic!("unexpected error delivery: {error:?}"), + Err(TryRecvError::Empty | TryRecvError::Disconnected) => return deliveries, + } + } + } + + #[test] + fn stream_interval_one_delivers_every_output() { + let (mut sender, mut receiver) = channel(1); + sender.send(output(&[1])); + sender.send(output(&[2, 3])); + sender.send(finished(&[4])); + + expect![[r#" + [ + [ + [ + 1, + ], + ], + [ + [ + 2, + 3, + ], + ], + [ + [ + 4, + ], + ], + ] + "#]] + .assert_debug_eq(&deliveries(&mut receiver)); + } + + #[test] + fn stream_interval_batches_outputs_after_the_first() { + let (mut sender, mut receiver) = channel(3); + // The first output is delivered immediately; later ones are held back + // until they carry at least 3 tokens, counted across multi-token + // outputs. The terminal output flushes whatever is held back. + for token_ids in [&[1][..], &[2], &[3], &[4, 5], &[6], &[7, 8, 9, 10], &[11]] { + sender.send(output(token_ids)); + } + sender.send(finished(&[12])); + + expect![[r#" + [ + [ + [ + 1, + ], + ], + [ + [ + 2, + ], + [ + 3, + ], + [ + 4, + 5, + ], + ], + [ + [ + 6, + ], + [ + 7, + 8, + 9, + 10, + ], + ], + [ + [ + 11, + ], + [ + 12, + ], + ], + ] + "#]] + .assert_debug_eq(&deliveries(&mut receiver)); + } + + #[test] + fn terminal_output_without_held_back_outputs_is_delivered_alone() { + let (mut sender, mut receiver) = channel(2); + sender.send(output(&[1])); + sender.send(output(&[2, 3])); + sender.send(finished(&[])); + + expect![[r#" + [ + [ + [ + 1, + ], + ], + [ + [ + 2, + 3, + ], + ], + [ + [], + ], + ] + "#]] + .assert_debug_eq(&deliveries(&mut receiver)); + } + + #[test] + fn flush_delivers_held_back_outputs() { + let (mut sender, mut receiver) = channel(4); + sender.send(output(&[1])); + sender.send(output(&[2])); + sender.flush(); + sender.flush(); + + expect![[r#" + [ + [ + [ + 1, + ], + ], + [ + [ + 2, + ], + ], + ] + "#]] + .assert_debug_eq(&deliveries(&mut receiver)); + } + + #[test] + fn error_follows_held_back_outputs() { + let (mut sender, mut receiver) = channel(4); + sender.send(output(&[1])); + sender.send(output(&[2])); + sender.send_error(Error::EngineCoreDead); + + let mut items = Vec::new(); + while let Ok(item) = receiver.try_recv() { + items.push(item.map(|delivery| token_ids(delivery).len()).ok()); + } + // One delivery for the first output, one for the held-back output, + // then the error. + assert_eq!(items, [Some(1), Some(1), None]); + } +} diff --git a/rust/src/engine-core-client/src/client/state.rs b/rust/src/engine-core-client/src/client/state.rs index a8f0fa3c24ed..9c183b87e4e0 100644 --- a/rust/src/engine-core-client/src/client/state.rs +++ b/rust/src/engine-core-client/src/client/state.rs @@ -2,12 +2,15 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::num::NonZeroU32; use std::sync::atomic::{AtomicU64, Ordering}; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::oneshot; use tracing::trace; use crate::EngineId; +use crate::client::output_channel::output_channel; +pub(crate) use crate::client::output_channel::{OutputReceiver, OutputSender}; use crate::client::stream::EngineCoreStreamOutput; use crate::error::{Error, Result}; use crate::protocol::output::{EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput}; @@ -15,8 +18,6 @@ use crate::protocol::stats::SchedulerStats; use crate::protocol::utility::UtilityOutput; use crate::transport::ConnectedEngine; -pub type OutputSender = mpsc::UnboundedSender>; -pub type OutputReceiver = mpsc::UnboundedReceiver>; pub type UtilitySender = oneshot::Sender>; pub type UtilityReceiver = oneshot::Receiver>; @@ -137,18 +138,22 @@ impl RequestRegistry { /// When `data_parallel_rank` is provided, the request is routed directly to /// the engine at that rank index, bypassing load balancing. Otherwise /// the engine with the lowest routing score is chosen. + /// + /// `stream_interval` sets how many new tokens the channel batches into one + /// delivery after the first output; see [`OutputSender::send`]. pub fn register( &mut self, request_id: String, lora_name: Option, data_parallel_rank: Option, + stream_interval: NonZeroU32, ) -> Result<(EngineId, OutputReceiver)> { if self.requests.contains_key(&request_id) { return Err(Error::DuplicateRequestId { request_id }); } let engine_id = self.choose_engine_for_request(data_parallel_rank)?; - let (tx, rx) = mpsc::unbounded_channel(); + let (tx, rx) = output_channel(stream_interval); let lora = lora_name.map(|adapter_name| LoraRequestState { adapter_name, phase: LoraPhase::Waiting, @@ -213,17 +218,30 @@ impl RequestRegistry { by_engine } - /// Obtain the stream sender for one output. If it indicates the request is - /// finished, it will be removed from the registry. - pub fn sender_for_output(&mut self, output: &EngineCoreOutput) -> Option { - self.apply_lora_events(output); + /// Send one output to its request stream. If it indicates the request is + /// finished, the request will be removed from the registry. Returns the + /// output back if its request is no longer tracked. + /// + /// The send happens under the registry lock: the sender is stateful (see + /// [`OutputSender::send`]) and owned by the registry entry, and sending on + /// its unbounded channel never blocks. + pub fn send_output( + &mut self, + output: EngineCoreStreamOutput, + ) -> Option { + self.apply_lora_events(&output.output); if output.finished() { - self.remove(output.request_id.as_str()).map(|tracked| tracked.0) + match self.remove(output.request_id.as_str()) { + Some((mut sender, _)) => sender.send(output), + None => return Some(output), + } } else { - self.requests - .get(output.request_id.as_str()) - .map(|tracked| tracked.sender.clone()) + match self.requests.get_mut(output.request_id.as_str()) { + Some(tracked) => tracked.sender.send(output), + None => return Some(output), + } } + None } /// Advance the request's LoRA scheduling phase from the engine-core events @@ -267,17 +285,8 @@ impl RequestRegistry { (running, waiting) } - /// Obtain stream senders for a whole engine output batch under one - /// registry lock. Finished outputs are removed before returning. - pub fn senders_for_outputs<'a>( - &mut self, - outputs: impl IntoIterator, - ) -> Vec> { - outputs.into_iter().map(|output| self.sender_for_output(output)).collect() - } - /// Remove a batch of requests that have finished or aborted, returning - /// their stream senders. + /// their stream senders after delivering any held-back outputs. pub fn finish_many<'a>( &mut self, request_ids: impl IntoIterator, @@ -285,6 +294,10 @@ impl RequestRegistry { request_ids .into_iter() .filter_map(|request_id| self.remove(request_id.as_str()).map(|tracked| tracked.0)) + .map(|mut sender| { + sender.flush(); + sender + }) .collect() } @@ -325,7 +338,7 @@ impl RequestRegistry { ) -> Vec { let mut aborted = Vec::new(); for request_id in request_ids { - let Some((sender, engine_id)) = self.remove(request_id) else { + let Some((mut sender, engine_id)) = self.remove(request_id) else { continue; }; let output = EngineCoreStreamOutput { @@ -337,7 +350,7 @@ impl RequestRegistry { ..EngineCoreOutput::default() }, }; - let _ = sender.send(Ok(output)); + sender.send(output); aborted.push(request_id.clone()); } aborted @@ -469,11 +482,13 @@ impl UtilityRegistry { #[cfg(test)] mod tests { use std::collections::BTreeSet; + use std::num::NonZeroU32; use crate::EngineId; use crate::client::state::{ - EngineLoadSnapshot, EngineRoutingState, RequestRegistry, UtilityRegistry, + EngineLoadSnapshot, EngineRoutingState, OutputReceiver, RequestRegistry, UtilityRegistry, }; + use crate::client::stream::EngineCoreStreamOutput; use crate::mock_engine::default_ready_response; use crate::protocol::output::{ EngineCoreEvent, EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput, @@ -508,6 +523,14 @@ mod tests { } } + fn stream_output(output: EngineCoreOutput) -> EngineCoreStreamOutput { + EngineCoreStreamOutput { + engine_index: 0, + timestamp: 0.0, + output, + } + } + fn adapter_names(values: &[&str]) -> BTreeSet { values.iter().map(|name| (*name).to_string()).collect() } @@ -515,8 +538,9 @@ mod tests { #[test] fn registry_rejects_duplicate_request_ids() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); - registry.register("req-1".to_string(), None, None).unwrap(); - let error = registry.register("req-1".to_string(), None, None).unwrap_err(); + registry.register("req-1".to_string(), None, None, NonZeroU32::MIN).unwrap(); + let error = + registry.register("req-1".to_string(), None, None, NonZeroU32::MIN).unwrap_err(); assert!(matches!( error, crate::error::Error::DuplicateRequestId { request_id } if request_id == "req-1" @@ -526,25 +550,99 @@ mod tests { #[test] fn registry_removes_finished_request_on_output() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); - registry.register("req-1".to_string(), None, None).unwrap(); + registry.register("req-1".to_string(), None, None, NonZeroU32::MIN).unwrap(); - let sender = registry.sender_for_output(&EngineCoreOutput { + let unrouted = registry.send_output(stream_output(EngineCoreOutput { request_id: "req-1".to_string(), finish_reason: Some(EngineCoreFinishReason::Length), ..Default::default() - }); + })); - assert!(sender.is_some()); + assert!(unrouted.is_none()); assert!(!registry.contains("req-1")); } + /// Drain the outputs that are ready on `receiver` without waiting, + /// rendering each as its token ids and finish reason. + fn ready_outputs( + receiver: &mut OutputReceiver, + ) -> Vec<(Vec, Option)> { + let mut outputs = Vec::new(); + while let Ok(delivery) = receiver.try_recv() { + for output in delivery.unwrap() { + outputs.push((output.output.new_token_ids, output.output.finish_reason)); + } + } + outputs + } + + #[test] + fn registry_abort_delivers_held_back_outputs_before_terminal_output() { + let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); + let (_, mut receiver) = registry + .register("req-1".to_string(), None, None, NonZeroU32::new(4).unwrap()) + .unwrap(); + + for token_id in 1..=3 { + let unrouted = registry.send_output(stream_output(EngineCoreOutput { + request_id: "req-1".to_string(), + new_token_ids: vec![token_id], + ..Default::default() + })); + assert!(unrouted.is_none()); + } + // Only the first output is delivered; the next two are held back + // below the stream interval. + expect_test::expect![[r#" + [ + ( + [ + 1, + ], + None, + ), + ] + "#]] + .assert_debug_eq(&ready_outputs(&mut receiver)); + + registry.abort_many(&["req-1".to_string()], 0.0); + expect_test::expect![[r#" + [ + ( + [ + 2, + ], + None, + ), + ( + [ + 3, + ], + None, + ), + ( + [], + Some( + Abort, + ), + ), + ] + "#]] + .assert_debug_eq(&ready_outputs(&mut receiver)); + } + #[test] fn registry_tracks_lora_phases_from_engine_events() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); registry - .register("req-lora".to_string(), Some("adapter-a".to_string()), None) + .register( + "req-lora".to_string(), + Some("adapter-a".to_string()), + None, + NonZeroU32::MIN, + ) .unwrap(); - registry.register("req-plain".to_string(), None, None).unwrap(); + registry.register("req-plain".to_string(), None, None, NonZeroU32::MIN).unwrap(); // Registered but not yet scheduled: counted as waiting. The non-LoRA // request never shows up. @@ -554,7 +652,7 @@ mod tests { ); // Queued then scheduled in one output: running. - drop(registry.sender_for_output(&output_with_events( + registry.send_output(stream_output(output_with_events( "req-lora", &[EngineCoreEventType::Queued, EngineCoreEventType::Scheduled], None, @@ -565,7 +663,7 @@ mod tests { ); // Preempted: back to waiting. - drop(registry.sender_for_output(&output_with_events( + registry.send_output(stream_output(output_with_events( "req-lora", &[EngineCoreEventType::Preempted], None, @@ -576,7 +674,7 @@ mod tests { ); // Finished: dropped from tracking entirely. - drop(registry.sender_for_output(&output_with_events( + registry.send_output(stream_output(output_with_events( "req-lora", &[EngineCoreEventType::Scheduled], Some(EngineCoreFinishReason::Stop), @@ -591,18 +689,33 @@ mod tests { fn registry_unions_lora_adapters_across_requests() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); registry - .register("req-a1".to_string(), Some("adapter-a".to_string()), None) + .register( + "req-a1".to_string(), + Some("adapter-a".to_string()), + None, + NonZeroU32::MIN, + ) .unwrap(); registry - .register("req-a2".to_string(), Some("adapter-a".to_string()), None) + .register( + "req-a2".to_string(), + Some("adapter-a".to_string()), + None, + NonZeroU32::MIN, + ) .unwrap(); registry - .register("req-b".to_string(), Some("adapter-b".to_string()), None) + .register( + "req-b".to_string(), + Some("adapter-b".to_string()), + None, + NonZeroU32::MIN, + ) .unwrap(); // One of adapter-a's requests starts running while the other waits: // the adapter appears in both sets. - drop(registry.sender_for_output(&output_with_events( + registry.send_output(stream_output(output_with_events( "req-a1", &[EngineCoreEventType::Scheduled], None, @@ -620,7 +733,7 @@ mod tests { fn registry_counts_only_active_lora_requests() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); - registry.register("req-plain".to_string(), None, None).unwrap(); + registry.register("req-plain".to_string(), None, None, NonZeroU32::MIN).unwrap(); assert_eq!(registry.active_lora_requests(), 0); assert_eq!( registry.lora_adapter_states(), @@ -632,6 +745,7 @@ mod tests { "req-lora-a".to_string(), Some("adapter-a".to_string()), None, + NonZeroU32::MIN, ) .unwrap(); registry @@ -639,6 +753,7 @@ mod tests { "req-lora-b".to_string(), Some("adapter-b".to_string()), None, + NonZeroU32::MIN, ) .unwrap(); assert_eq!(registry.active_lora_requests(), 2); @@ -661,7 +776,12 @@ mod tests { fn registry_clears_lora_count_on_close() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); registry - .register("req-lora".to_string(), Some("adapter-a".to_string()), None) + .register( + "req-lora".to_string(), + Some("adapter-a".to_string()), + None, + NonZeroU32::MIN, + ) .unwrap(); assert_eq!(registry.active_lora_requests(), 1); @@ -677,7 +797,12 @@ mod tests { fn registry_drops_lora_tracking_on_abort() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); registry - .register("req-lora".to_string(), Some("adapter-a".to_string()), None) + .register( + "req-lora".to_string(), + Some("adapter-a".to_string()), + None, + NonZeroU32::MIN, + ) .unwrap(); drop(registry.finish_many(&["req-lora".to_string()])); @@ -691,8 +816,8 @@ mod tests { #[test] fn registry_closes_all_requests_on_failure() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); - registry.register("req-1".to_string(), None, None).unwrap(); - registry.register("req-2".to_string(), None, None).unwrap(); + registry.register("req-1".to_string(), None, None, NonZeroU32::MIN).unwrap(); + registry.register("req-2".to_string(), None, None, NonZeroU32::MIN).unwrap(); let senders = registry.close(); @@ -708,9 +833,12 @@ mod tests { connected_engine(engine_0.clone()), connected_engine(engine_1.clone()), ]); - let (chosen_0, _) = registry.register("req-1".to_string(), None, None).unwrap(); - let (chosen_1, _) = registry.register("req-2".to_string(), None, None).unwrap(); - let (chosen_0_again, _) = registry.register("req-3".to_string(), None, None).unwrap(); + let (chosen_0, _) = + registry.register("req-1".to_string(), None, None, NonZeroU32::MIN).unwrap(); + let (chosen_1, _) = + registry.register("req-2".to_string(), None, None, NonZeroU32::MIN).unwrap(); + let (chosen_0_again, _) = + registry.register("req-3".to_string(), None, None, NonZeroU32::MIN).unwrap(); assert_eq!(chosen_0, engine_0); assert_eq!(chosen_1, engine_1); @@ -737,9 +865,12 @@ mod tests { connected_engine(engine_1.clone()), ]); - let (chosen_0, _) = registry.register("req-1".to_string(), None, None).unwrap(); - let (chosen_1, _) = registry.register("req-2".to_string(), None, None).unwrap(); - let (chosen_0_again, _) = registry.register("req-3".to_string(), None, None).unwrap(); + let (chosen_0, _) = + registry.register("req-1".to_string(), None, None, NonZeroU32::MIN).unwrap(); + let (chosen_1, _) = + registry.register("req-2".to_string(), None, None, NonZeroU32::MIN).unwrap(); + let (chosen_0_again, _) = + registry.register("req-3".to_string(), None, None, NonZeroU32::MIN).unwrap(); assert_eq!(chosen_0, engine_0); assert_eq!(chosen_1, engine_1); @@ -809,7 +940,8 @@ mod tests { } )); - let (chosen, _) = registry.register("req-stats".to_string(), None, None).unwrap(); + let (chosen, _) = + registry.register("req-stats".to_string(), None, None, NonZeroU32::MIN).unwrap(); assert_eq!(chosen, engine_1); } @@ -824,7 +956,9 @@ mod tests { for (rank, running) in counts.into_iter().enumerate() { for index in 0..running { let request_id = format!("old-{rank}-{index}"); - registry.register(request_id.clone(), None, Some(rank as u32)).unwrap(); + registry + .register(request_id.clone(), None, Some(rank as u32), NonZeroU32::MIN) + .unwrap(); old_requests.push(request_id); } registry.apply_scheduler_counts( @@ -839,7 +973,8 @@ mod tests { let mut distribution = [0; 2]; for index in 0..14 { - let (engine, _) = registry.register(format!("new-{index}"), None, None).unwrap(); + let (engine, _) = + registry.register(format!("new-{index}"), None, None, NonZeroU32::MIN).unwrap(); distribution[engine.engine_index().unwrap() as usize] += 1; } distributions.push(distribution); @@ -865,7 +1000,8 @@ mod tests { ); } let request_id = format!("req-{index}"); - let (engine, _) = registry.register(request_id.clone(), None, None).unwrap(); + let (engine, _) = + registry.register(request_id.clone(), None, None, NonZeroU32::MIN).unwrap(); chosen.push(engine.engine_index().unwrap()); drop(registry.finish_many(&[request_id])); } @@ -890,12 +1026,15 @@ mod tests { }, ); for index in 0..4 { - registry.register(format!("pinned-{index}"), None, Some(0)).unwrap(); + registry + .register(format!("pinned-{index}"), None, Some(0), NonZeroU32::MIN) + .unwrap(); } // Even unchanged counts replace the optimistic estimate for that engine. registry.apply_scheduler_counts(0, snapshot); - let (after_refresh, _) = registry.register("after-refresh".into(), None, None).unwrap(); + let (after_refresh, _) = + registry.register("after-refresh".into(), None, None, NonZeroU32::MIN).unwrap(); // An older, empty snapshot must not erase the five admitted requests. for rank in [0, 1] { registry.apply_scheduler_counts( @@ -906,7 +1045,8 @@ mod tests { }, ); } - let (after_empty, _) = registry.register("after-empty".into(), None, None).unwrap(); + let (after_empty, _) = + registry.register("after-empty".into(), None, None, NonZeroU32::MIN).unwrap(); expect_test::expect!["(Some(0), Some(1))"].assert_eq(&format!( "{:?}", (after_refresh.engine_index(), after_empty.engine_index()) @@ -927,10 +1067,11 @@ mod tests { }, ); } - registry.register("pinned".into(), None, Some(9)).unwrap(); + registry.register("pinned".into(), None, Some(9), NonZeroU32::MIN).unwrap(); let mut chosen = Vec::new(); for index in 0..3 { - let (engine, _) = registry.register(format!("auto-{index}"), None, None).unwrap(); + let (engine, _) = + registry.register(format!("auto-{index}"), None, None, NonZeroU32::MIN).unwrap(); chosen.push(engine.engine_index().unwrap()); } expect_test::expect!["[2, 5, 2]"].assert_eq(&format!("{chosen:?}")); @@ -948,15 +1089,18 @@ mod tests { ]); // Explicitly target rank 2 (third engine). - let (chosen, _) = registry.register("req-1".to_string(), None, Some(2)).unwrap(); + let (chosen, _) = + registry.register("req-1".to_string(), None, Some(2), NonZeroU32::MIN).unwrap(); assert_eq!(chosen, engine_2); // Explicitly target rank 0 (first engine). - let (chosen, _) = registry.register("req-2".to_string(), None, Some(0)).unwrap(); + let (chosen, _) = + registry.register("req-2".to_string(), None, Some(0), NonZeroU32::MIN).unwrap(); assert_eq!(chosen, engine_0); // Explicitly target rank 1. - let (chosen, _) = registry.register("req-3".to_string(), None, Some(1)).unwrap(); + let (chosen, _) = + registry.register("req-3".to_string(), None, Some(1), NonZeroU32::MIN).unwrap(); assert_eq!(chosen, engine_1); } @@ -970,11 +1114,12 @@ mod tests { ]); // Load-balance: first two go to engine_0 and engine_1. - registry.register("req-lb-0".to_string(), None, None).unwrap(); + registry.register("req-lb-0".to_string(), None, None, NonZeroU32::MIN).unwrap(); // Now engine_0 has 1 in-flight. Without dp_rank, next would go to engine_1. // But with dp_rank=0, it should still go to engine_0. - let (chosen, _) = registry.register("req-dp".to_string(), None, Some(0)).unwrap(); + let (chosen, _) = + registry.register("req-dp".to_string(), None, Some(0), NonZeroU32::MIN).unwrap(); assert_eq!(chosen, engine_0); } @@ -985,7 +1130,9 @@ mod tests { connected_engine(EngineId::from_engine_index(1)), ]); - let error = registry.register("req-1".to_string(), None, Some(2)).unwrap_err(); + let error = registry + .register("req-1".to_string(), None, Some(2), NonZeroU32::MIN) + .unwrap_err(); assert!(matches!( error, crate::error::Error::InvalidDataParallelRank { @@ -1000,10 +1147,13 @@ mod tests { let engine_3 = EngineId::from_engine_index(3); let mut registry = RequestRegistry::new(&[connected_engine(engine_3.clone())]); - let (chosen, _) = registry.register("req-ok".to_string(), None, Some(3)).unwrap(); + let (chosen, _) = + registry.register("req-ok".to_string(), None, Some(3), NonZeroU32::MIN).unwrap(); assert_eq!(chosen, engine_3); - let error = registry.register("req-bad".to_string(), None, Some(0)).unwrap_err(); + let error = registry + .register("req-bad".to_string(), None, Some(0), NonZeroU32::MIN) + .unwrap_err(); assert!(matches!( error, crate::error::Error::InvalidDataParallelRank { diff --git a/rust/src/engine-core-client/src/client/stream.rs b/rust/src/engine-core-client/src/client/stream.rs index ba408f2c8788..6f363f61017d 100644 --- a/rust/src/engine-core-client/src/client/stream.rs +++ b/rust/src/engine-core-client/src/client/stream.rs @@ -5,8 +5,9 @@ use std::ops::Deref; use std::pin::Pin; use std::task::{Context, Poll}; -use futures::Stream; use futures::stream::FusedStream; +use futures::{Stream, StreamExt as _}; +use itertools::Either; use thiserror_ext::AsReport as _; use tokio::sync::mpsc; use tracing::{debug, error, warn}; @@ -40,6 +41,68 @@ impl Deref for EngineCoreStreamOutput { } } +/// Consecutive outputs of one request, delivered to its stream together and in +/// order. Outputs are grouped by the request's `stream_interval` (see +/// [`EngineCoreClient::call`]) but not merged; consumers merge them after +/// observing each raw output. +/// +/// [`EngineCoreClient::call`]: crate::EngineCoreClient::call +#[derive(Debug)] +#[allow( + clippy::large_enum_variant, + reason = "single outputs stay inline, avoiding an allocation per engine step" +)] +pub enum EngineCoreStreamDelivery { + /// A single output, sent without being held back. + One(EngineCoreStreamOutput), + /// Outputs held back by the stream interval, in order. Never empty. + Many(Vec), +} + +impl EngineCoreStreamDelivery { + /// Return the first output. + pub fn first(&self) -> &EngineCoreStreamOutput { + match self { + Self::One(output) => output, + Self::Many(outputs) => outputs.first().expect("deliveries are never empty"), + } + } + + /// Return the last output, the only one that can be terminal. + pub fn last(&self) -> &EngineCoreStreamOutput { + match self { + Self::One(output) => output, + Self::Many(outputs) => outputs.last().expect("deliveries are never empty"), + } + } + + /// Iterate over the outputs in order. + pub fn iter( + &self, + ) -> Either< + std::iter::Once<&EngineCoreStreamOutput>, + std::slice::Iter<'_, EngineCoreStreamOutput>, + > { + match self { + Self::One(output) => Either::Left(std::iter::once(output)), + Self::Many(outputs) => Either::Right(outputs.iter()), + } + } +} + +impl IntoIterator for EngineCoreStreamDelivery { + type Item = EngineCoreStreamOutput; + type IntoIter = + Either, std::vec::IntoIter>; + + fn into_iter(self) -> Self::IntoIter { + match self { + Self::One(output) => Either::Left(std::iter::once(output)), + Self::Many(outputs) => Either::Right(outputs.into_iter()), + } + } +} + /// Stream of raw engine-core outputs for one request. /// /// The stream yields only [`EngineCoreStreamOutput`] values whose embedded @@ -79,10 +142,21 @@ impl EngineCoreOutputStream { pub fn engine_index(&self) -> u32 { self.engine_index } + + /// Flatten the stream into one item per raw output, for consumers that do + /// not care how outputs are grouped into deliveries. + pub fn into_outputs(self) -> impl Stream> + Send + Unpin { + self.flat_map(|item| { + futures::stream::iter(match item { + Ok(delivery) => Either::Left(delivery.into_iter().map(Ok)), + Err(error) => Either::Right(std::iter::once(Err(error))), + }) + }) + } } impl Stream for EngineCoreOutputStream { - type Item = Result; + type Item = Result; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { if self.is_terminated() { @@ -93,10 +167,11 @@ impl Stream for EngineCoreOutputStream { Poll::Pending => Poll::Pending, Poll::Ready(Some(item)) => { match &item { - Ok(output) => { + Ok(delivery) => { // If the output indicates the request is finished, mark the stream as // terminated with cleanly-finished state and expect no more outputs to - // come. + // come. Only the last output of a delivery can be terminal. + let output = delivery.last(); if output.finished() { if output.finish_reason == Some(EngineCoreFinishReason::Error) { error!( diff --git a/rust/src/engine-core-client/src/lib.rs b/rust/src/engine-core-client/src/lib.rs index 7b60206194ce..8f79e227de11 100644 --- a/rust/src/engine-core-client/src/lib.rs +++ b/rust/src/engine-core-client/src/lib.rs @@ -14,7 +14,7 @@ mod transport; pub use client::{ AbortCause, CoordinatorMode, EngineCoreClient, EngineCoreClientConfig, EngineCoreOutputStream, - EngineCoreStreamOutput, TransportMode, + EngineCoreStreamDelivery, EngineCoreStreamOutput, TransportMode, }; pub use error::{Error, Result}; pub use transport::{ENGINE_CORE_DEAD_SENTINEL, EngineId}; diff --git a/rust/src/engine-core-client/src/protocol/sampling.rs b/rust/src/engine-core-client/src/protocol/sampling.rs index 577dcb651f20..51e80d0f66eb 100644 --- a/rust/src/engine-core-client/src/protocol/sampling.rs +++ b/rust/src/engine-core-client/src/protocol/sampling.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::{BTreeSet, HashMap}; +use std::num::NonZeroU32; use serde::{Deserialize, Serialize}; use serde_default::DefaultFromSerde; @@ -161,6 +162,11 @@ pub struct EngineCoreSamplingParams { /// Number of prompt tokens to skip from returned routed-expert data. /// A value of zero returns routing data for the entire prompt. pub routed_experts_prompt_start: u32, + /// Number of newly generated tokens to batch into each streamed output. + /// The first and final outputs are always emitted immediately. Engine-core + /// ignores it; [`EngineCoreClient::call`](crate::EngineCoreClient::call) + /// batches output deliveries by it. + pub stream_interval: Option, } impl EngineCoreSamplingParams { @@ -195,6 +201,7 @@ impl EngineCoreSamplingParams { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } } } diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index dbd3fa69ceb2..0ad58b027154 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -755,8 +755,8 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { ) .await; - let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap(); - let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap(); + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap().into_outputs(); + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap().into_outputs(); let final_1 = timeout(Duration::from_secs(1), stream_1.next()) .await @@ -778,7 +778,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { tokio::time::sleep(Duration::from_millis(100)).await; - let mut stream_3 = client.call(sample_request_with_id("req-3")).await.unwrap(); + let mut stream_3 = client.call(sample_request_with_id("req-3")).await.unwrap().into_outputs(); let final_3 = timeout(Duration::from_secs(1), stream_3.next()) .await .unwrap() @@ -933,7 +933,7 @@ async fn coordinator_accepts_stats_only_outputs() { ) .await; - let mut stream = client.call(sample_request_with_id("req-stats")).await.unwrap(); + let mut stream = client.call(sample_request_with_id("req-stats")).await.unwrap().into_outputs(); let final_output = timeout(Duration::from_secs(1), stream.next()).await.unwrap().unwrap().unwrap(); assert_eq!(final_output.request_id, "req-stats"); @@ -1026,8 +1026,8 @@ async fn client_fail_closes_when_main_output_path_receives_dp_control() { assert!(client.ready_responses()[0].max_model_len > 0); assert_eq!(client.vllm_version(), "test-vllm-version"); - let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap(); - let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap(); + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap().into_outputs(); + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap().into_outputs(); let error_2 = timeout(Duration::from_secs(1), stream_2.next()) .await @@ -1103,7 +1103,7 @@ async fn duplicate_request_ids_are_rejected_without_sending_a_second_add() { ) .await; - let mut stream = client.call(sample_request()).await.unwrap(); + let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); let error = match client.call(sample_request()).await { Ok(_) => panic!("expected duplicate request error"), Err(error) => error, @@ -1170,7 +1170,7 @@ async fn finished_requests_without_final_output_is_treated_as_unexpected_close() ) .await; - let mut stream = client.call(sample_request()).await.unwrap(); + let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); let error = timeout(Duration::from_secs(1), stream.next()) .await .unwrap() @@ -1233,7 +1233,7 @@ async fn dropping_a_live_stream_triggers_abort() { ) .await; - let mut stream = client.call(sample_request()).await.unwrap(); + let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); let first = timeout(Duration::from_secs(1), stream.next()).await.unwrap().unwrap().unwrap(); assert_eq!(first.new_token_ids, vec![99]); drop(stream); @@ -1316,7 +1316,7 @@ async fn dropping_multiple_live_streams_aborts_all_in_a_burst() { // emits outputs, then drain the first token from each stream. let mut streams = Vec::new(); for id in request_ids { - streams.push(client.call(sample_request_with_id(id)).await.unwrap()); + streams.push(client.call(sample_request_with_id(id)).await.unwrap().into_outputs()); } for stream in streams.iter_mut() { let first = timeout(Duration::from_secs(1), stream.next()).await.unwrap().unwrap().unwrap(); @@ -1363,8 +1363,8 @@ async fn dispatcher_failure_propagates_to_streams_and_future_calls() { ) .await; - let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap(); - let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap(); + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap().into_outputs(); + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap().into_outputs(); let error_1 = timeout(Duration::from_secs(1), stream_1.next()) .await @@ -1862,7 +1862,7 @@ async fn client_decodes_multipart_logprob_outputs() { ) .await; - let stream = client.call(sample_request()).await.unwrap(); + let stream = client.call(sample_request()).await.unwrap().into_outputs(); let outputs = stream.collect::>().await; assert_eq!(outputs.len(), 1); @@ -1950,7 +1950,7 @@ async fn client_sends_large_multimodal_tensor_as_aux_frame() { ) .await; - let outputs = client.call(request).await.unwrap().collect::>().await; + let outputs = client.call(request).await.unwrap().into_outputs().collect::>().await; assert_eq!(outputs.len(), 1); assert!(outputs[0].is_ok()); @@ -2079,8 +2079,8 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() { assert_eq!(client.ready_responses().len(), 2); assert_eq!(client.engine_identities()[0], b"engine-0"); - let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap(); - let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap(); + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap().into_outputs(); + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap().into_outputs(); assert_eq!( timeout(Duration::from_secs(1), engine_0_seen_rx.recv()).await.unwrap().unwrap(), "req-1" @@ -2100,7 +2100,7 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() { assert_eq!(final_1.new_token_ids, vec![10]); assert_eq!(final_1.finish_reason, Some(EngineCoreFinishReason::Length)); - let mut stream_3 = client.call(sample_request_with_id("req-3")).await.unwrap(); + let mut stream_3 = client.call(sample_request_with_id("req-3")).await.unwrap().into_outputs(); assert_eq!( timeout(Duration::from_secs(1), engine_0_seen_rx.recv()).await.unwrap().unwrap(), "req-3" @@ -2271,8 +2271,8 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() { assert!(client.is_sleeping().await.unwrap()); - let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap(); - let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap(); + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap().into_outputs(); + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap().into_outputs(); client .abort(&[ @@ -2803,6 +2803,7 @@ fn python_msgpack_fixtures_match_rust_encoding() { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, }, ); @@ -3266,7 +3267,7 @@ async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends .await; tokio::time::sleep(Duration::from_millis(50)).await; - let mut stream = client.call(sample_request()).await.unwrap(); + let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); let wakeup = timeout( Duration::from_secs(1), @@ -3352,7 +3353,7 @@ async fn bootstrapped_external_coordinator_running_state_suppresses_wakeup() { send_external_coordinator_publish(&mut stats_socket, &(Value::Nil, 5_u32, true)).await; tokio::time::sleep(Duration::from_millis(50)).await; - let mut stream = client.call(sample_request()).await.unwrap(); + let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); assert!( timeout( diff --git a/rust/src/llm/src/lib.rs b/rust/src/llm/src/lib.rs index 015624b409c8..7e9b5e0ba53f 100644 --- a/rust/src/llm/src/lib.rs +++ b/rust/src/llm/src/lib.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::num::NonZeroU32; + use tracing::Span; use vllm_engine_core_client::EngineCoreClient; @@ -34,6 +36,7 @@ use crate::request_metrics::RequestMetricsTracker; pub struct Llm { client: EngineCoreClient, randomize_request_id: bool, + stream_interval: NonZeroU32, stats_logger: Option, inflight: InflightRequests, } @@ -45,6 +48,7 @@ impl Llm { Self { client, randomize_request_id: true, + stream_interval: NonZeroU32::MIN, stats_logger: None, inflight: InflightRequests::new(), } @@ -71,6 +75,14 @@ impl Llm { self } + /// Set the frontend-level stream interval: the minimum number of newly + /// generated tokens batched into each streamed output after the first one. + /// A request's own `stream_interval` can only raise it. + pub fn with_stream_interval(mut self, stream_interval: NonZeroU32) -> Self { + self.stream_interval = stream_interval; + self + } + /// Expose the underlying engine-core client for low-level utility/admin /// calls. pub fn engine_core_client(&self) -> &EngineCoreClient { @@ -79,7 +91,16 @@ impl Llm { /// Submit one tokenized generate request and return a per-request output /// stream. - pub async fn generate(&self, req: GenerateRequest) -> Result { + pub async fn generate(&self, mut req: GenerateRequest) -> Result { + // Clamp the request's stream interval to the frontend-level one, like + // Python. Engine-core ignores it; the client batches deliveries by it. + let stream_interval = (req.sampling_params.stream_interval) + .map_or(self.stream_interval, |interval| { + interval.max(self.stream_interval) + }); + req.sampling_params.stream_interval = + (stream_interval > NonZeroU32::MIN).then_some(stream_interval); + let prepared = req.prepare(self.randomize_request_id)?; let prompt_token_ids = prepared.prompt_token_ids().into(); let external_request_id = prepared diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index 4dde96f2e20c..b6ccb7717508 100644 --- a/rust/src/llm/src/output.rs +++ b/rust/src/llm/src/output.rs @@ -16,7 +16,9 @@ use vllm_engine_core_client::protocol::output::{ }; use vllm_engine_core_client::protocol::sampling_mask::SamplingMask; use vllm_engine_core_client::protocol::tensor::WireNdArray; -use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream}; +use vllm_engine_core_client::{ + AbortCause, EngineCoreOutputStream, EngineCoreStreamDelivery, EngineCoreStreamOutput, +}; use crate::error::Result; use crate::inflight::RequestGuard; @@ -195,6 +197,35 @@ impl GenerateOutput { pub fn finished(&self) -> bool { self.finish_reason.is_some() } + + /// Append a later output of the same request, as if both were produced in + /// one step. Per-token data is concatenated; per-output fields take the + /// later output's value, like Python's `RequestState.make_request_output()` + /// does for the output that flushes a stream interval. + fn merge(&mut self, next: GenerateOutput) { + debug_assert!( + next.prompt_info.is_none(), + "only the first output carries prompt info" + ); + self.token_ids.extend(next.token_ids); + if let Some(next_logprobs) = next.logprobs { + match &mut self.logprobs { + Some(logprobs) => logprobs.positions.extend(next_logprobs.positions), + None => self.logprobs = Some(next_logprobs), + } + } + if let Some(next_mask) = next.sampling_mask { + match &mut self.sampling_mask { + Some(mask) => mask.rows.extend(next_mask.rows), + None => self.sampling_mask = Some(next_mask), + } + } + self.cached_token_count = self.cached_token_count.max(next.cached_token_count); + self.finish_reason = next.finish_reason; + self.kv_transfer_params = next.kv_transfer_params; + self.ec_transfer_params = next.ec_transfer_params; + self.spec_decode_metrics = next.spec_decode_metrics; + } } #[cfg(any(test, feature = "test-util"))] @@ -266,18 +297,27 @@ impl GenerateOutputStream { pub fn request_id(&self) -> &str { self.raw_stream.request_id() } -} - -impl Stream for GenerateOutputStream { - type Item = Result; - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let raw = match ready!(Pin::new(&mut self.raw_stream).poll_next(cx)) { - Some(Ok(raw)) => raw, - Some(Err(error)) => return Poll::Ready(Some(Err(error.into()))), - None => return Poll::Ready(None), + /// Merge the outputs of one delivery into one output, like Python + /// `RequestState.make_request_output()` does for outputs held back by + /// `stream_interval`. Request metrics still observe every raw output. + fn merge_delivery(&mut self, delivery: EngineCoreStreamDelivery) -> Result { + // Match rather than iterate, so the common single output is converted + // without first being moved into an iterator. + let mut outputs = match delivery { + EngineCoreStreamDelivery::One(output) => return self.convert_output(output), + EngineCoreStreamDelivery::Many(outputs) => outputs.into_iter(), }; + let first = outputs.next().expect("deliveries are never empty"); + let mut merged = self.convert_output(first)?; + for raw in outputs { + merged.merge(self.convert_output(raw)?); + } + Ok(merged) + } + /// Convert one raw engine-core output. + fn convert_output(&mut self, raw: EngineCoreStreamOutput) -> Result { let received_at = current_unix_timestamp_secs(); self.request_metrics.observe_output(raw.timestamp, received_at, &raw.output); @@ -301,11 +341,11 @@ impl Stream for GenerateOutputStream { if let Some(mask) = sampling_mask.as_ref() && mask.rows.len() != raw.new_token_ids.len() { - return Poll::Ready(Some(Err(crate::Error::SamplingMaskTokenCountMismatch { + return Err(crate::Error::SamplingMaskTokenCountMismatch { request_id: raw.request_id, token_count: raw.new_token_ids.len(), row_count: mask.rows.len(), - }))); + }); } let cached_token_count = raw .prefill_stats @@ -331,7 +371,20 @@ impl Stream for GenerateOutputStream { spec_decode_metrics: raw.spec_decode_metrics, }; - Poll::Ready(Some(Ok(output))) + Ok(output) + } +} + +impl Stream for GenerateOutputStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let delivery = match ready!(Pin::new(&mut self.raw_stream).poll_next(cx)) { + Some(Ok(delivery)) => delivery, + Some(Err(error)) => return Poll::Ready(Some(Err(error.into()))), + None => return Poll::Ready(None), + }; + Poll::Ready(Some(self.merge_delivery(delivery))) } } diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 813171555bba..e1205ede9961 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::BTreeSet; +use std::num::NonZeroU32; use std::sync::{Arc, Once}; use std::time::Duration; @@ -308,6 +309,111 @@ async fn generate_streams_outputs() { llm.shutdown().await.unwrap(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn generate_merges_outputs_by_stream_interval() { + init_tracing(); + let ipc = IpcNamespace::new().unwrap(); + let handshake_address = ipc.handshake_endpoint(); + let engine_id = b"engine-stream-interval".to_vec(); + + // One logprobs position per token, so merged outputs can be checked + // against the logprobs of their concatenated tokens. + fn token_logprobs(token_ids: &[u32]) -> Logprobs { + Logprobs { + positions: token_ids + .iter() + .map(|&token_id| PositionLogprobs { + entries: vec![TokenLogprob { + token_id, + logprob: -0.5, + rank: 1, + }], + }) + .collect(), + } + } + + let (shutdown_tx, engine_task) = spawn_mock_engine_task( + handshake_address.clone(), + engine_id.clone(), + |dealer, push| { + Box::pin(async move { + let add = recv_engine_message(dealer).await; + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); + + // One engine step per batch. + let steps: [(&[u32], Option); 6] = [ + (&[1], None), + (&[2], None), + (&[3], None), + (&[4, 5], None), + (&[6], None), + (&[7], Some(EngineCoreFinishReason::Length)), + ]; + for (token_ids, finish_reason) in steps { + send_outputs( + push, + RequestBatchOutputs { + outputs: vec![request_output_with_logprobs( + &request.request_id, + token_ids.to_vec(), + finish_reason, + Some(token_logprobs(token_ids)), + None, + )], + ..Default::default() + } + .into(), + ) + .await; + } + }) + }, + ); + + // The request's interval raises the frontend-level one. + let llm = connect_async_llm_with_ipc(handshake_address, 7, "test-model", &ipc) + .await + .with_stream_interval(NonZeroU32::new(2).unwrap()); + let mut request = sample_generate_request("req-stream-interval", 7); + request.sampling_params.stream_interval = NonZeroU32::new(3); + let outputs: Vec<_> = llm + .generate(request) + .await + .unwrap() + .map(|output| output.unwrap()) + .collect() + .await; + + let _ = shutdown_tx.send(()); + engine_task.await.unwrap(); + llm.shutdown().await.unwrap(); + + // The first output is yielded immediately, later ones once they carry at + // least 3 tokens, and the terminal one flushes the rest. + let shapes: Vec<_> = outputs + .iter() + .map(|output| { + ( + output.prompt_info.is_some(), + output.token_ids.clone(), + output.finish_reason.clone(), + ) + }) + .collect(); + assert_eq!( + shapes, + [ + (true, vec![1], None), + (false, vec![2, 3, 4, 5], None), + (false, vec![6, 7], Some(FinishReason::Length)), + ] + ); + for output in &outputs { + assert_eq!(output.logprobs, Some(token_logprobs(&output.token_ids))); + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() { init_tracing(); diff --git a/rust/src/mock-engine/src/tests.rs b/rust/src/mock-engine/src/tests.rs index d82ceac083ca..0b053cbbb16a 100644 --- a/rust/src/mock-engine/src/tests.rs +++ b/rust/src/mock-engine/src/tests.rs @@ -116,7 +116,7 @@ async fn chunk_size_one_outputs_one_token_per_update() { let ipc = IpcNamespace::new().expect("ipc namespace"); let handshake_address = ipc.handshake_endpoint(); let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await; - let mut stream = client.call(sample_request("req-1", 3)).await.expect("call"); + let mut stream = client.call(sample_request("req-1", 3)).await.expect("call").into_outputs(); let first = stream.next().await.expect("first").expect("first ok"); assert_eq!(first.new_token_ids.len(), 1); @@ -147,7 +147,7 @@ async fn chunk_size_clips_final_output_to_max_tokens() { let ipc = IpcNamespace::new().expect("ipc namespace"); let handshake_address = ipc.handshake_endpoint(); let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 4).await; - let mut stream = client.call(sample_request("req-clip", 6)).await.expect("call"); + let mut stream = client.call(sample_request("req-clip", 6)).await.expect("call").into_outputs(); let first = stream.next().await.expect("first").expect("first ok"); assert_eq!(first.new_token_ids.len(), 4); @@ -165,7 +165,11 @@ async fn abort_cancels_active_request_and_emits_terminal_output() { let ipc = IpcNamespace::new().expect("ipc namespace"); let handshake_address = ipc.handshake_endpoint(); let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await; - let mut stream = client.call(sample_request("req-abort", 1_000_000)).await.expect("call"); + let mut stream = client + .call(sample_request("req-abort", 1_000_000)) + .await + .expect("call") + .into_outputs(); let first = stream.next().await.expect("first").expect("first ok"); assert_eq!(first.finish_reason, None); diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs index 43192ab78f47..a391f56caeaa 100644 --- a/rust/src/server/examples/external_engine_openai_qwen.rs +++ b/rust/src/server/examples/external_engine_openai_qwen.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::HashMap; +use std::num::NonZeroU32; use std::time::Duration; use anyhow::{Context, Result, bail}; @@ -81,6 +82,7 @@ async fn main() -> Result<()> { lora_modules: Vec::new(), chat_template_content_format: ChatTemplateContentFormatOption::Auto, max_logprobs: None, + stream_interval: NonZeroU32::MIN, api_server_options: ApiServerOptions::default(), cors: CorsConfig::default(), tls: None, diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index 8e51f1440b3e..3f95b5891cb2 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -3,6 +3,7 @@ use std::collections::HashMap; use std::fmt; +use std::num::NonZeroU32; use std::str::FromStr; use std::time::Duration; @@ -258,6 +259,10 @@ pub struct Config { /// Optional maximum number of top log probabilities accepted by the /// frontend. `None` delegates to the text layer default. pub max_logprobs: Option, + /// Minimum number of newly generated tokens batched into each streamed + /// output after the first one. Requests can raise it with their own + /// `stream_interval`. + pub stream_interval: NonZeroU32, /// HTTP/API-server behavior switches. pub api_server_options: ApiServerOptions, /// CORS settings applied to every HTTP response. diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index f54f8ce3e423..d5999cd3e78b 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -123,7 +123,9 @@ async fn build_state(config: &Config) -> Result> { .await .context("failed to connect to engine core")?; - let llm = Llm::new(client).with_log_stats(!config.disable_log_stats); + let llm = Llm::new(client) + .with_log_stats(!config.disable_log_stats) + .with_stream_interval(config.stream_interval); let text = TextLlm::new(llm, text_backend).with_max_logprobs(config.max_logprobs); let chat = ChatLlm::new(text, chat_backend) diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 4f80dda6c7b0..46aaafbc69e2 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -179,6 +179,7 @@ pub(super) fn prepare_chat_request( merge_ec_transfer_params(request.vllm_xargs, request.ec_transfer_params.as_ref()), request.kv_transfer_params.as_ref(), ), + stream_interval: request.stream_interval, }, chat_options: ChatOptions { generation_prompt_mode, @@ -443,6 +444,7 @@ fn convert_tool_choice(tool_choice: &ToolChoice) -> Result(request).is_err()); + } + #[test] fn chat_http_request_defaults_missing_or_null_model() { for model in [None, Some(serde_json::Value::Null)] { @@ -1004,6 +1015,7 @@ mod tests { frequency_penalty: Some(0.3), presence_penalty: Some(0.4), repetition_penalty: Some(1.1), + stream_interval: NonZeroU32::new(4), ..base_request() }; @@ -1019,6 +1031,7 @@ mod tests { frequency_penalty: Some(0.3), presence_penalty: Some(0.4), repetition_penalty: Some(1.1), + stream_interval: NonZeroU32::new(4), ..VllmSamplingParams::default() }; assert_eq!(prepared.chat_request.sampling_params, expected); diff --git a/rust/src/server/src/routes/openai/chat_completions/types.rs b/rust/src/server/src/routes/openai/chat_completions/types.rs index 13c7c7233f14..204f3bab5aab 100644 --- a/rust/src/server/src/routes/openai/chat_completions/types.rs +++ b/rust/src/server/src/routes/openai/chat_completions/types.rs @@ -3,6 +3,7 @@ use std::collections::HashMap; use std::fmt; +use std::num::NonZeroU32; use std::sync::Arc; use serde::{Deserialize, Serialize}; @@ -256,6 +257,12 @@ pub struct ChatCompletionRequest { /// Parameters for detecting repetitive N-gram patterns in output tokens pub repetition_detection: Option, + + /// Number of tokens to batch into each streamed chunk. Raises the server's + /// `--stream-interval` for this request. Values below the server setting + /// are clamped up to it. The first and last chunks are always sent + /// immediately. Ignored for non-streaming requests. + pub stream_interval: Option, } impl Default for ChatCompletionRequest { @@ -323,6 +330,7 @@ impl Default for ChatCompletionRequest { ec_transfer_params: None, vllm_xargs: None, repetition_detection: None, + stream_interval: None, } } } diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index 0b37f84057f5..33fc6bbe0a07 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -144,6 +144,7 @@ pub(super) fn prepare_completion_request( merge_ec_transfer_params(request.vllm_xargs, request.ec_transfer_params.as_ref()), request.kv_transfer_params.as_ref(), ), + stream_interval: request.stream_interval, }, decode_options: TextDecodeOptions { skip_special_tokens: request.skip_special_tokens, @@ -211,6 +212,8 @@ fn completion_echo_text( #[cfg(test)] mod tests { + use std::num::NonZeroU32; + use axum::http::HeaderMap; use serde_json::json; use validator::Validate; @@ -400,7 +403,8 @@ mod tests { "presence_penalty": 0.3, "repetition_penalty": 1.1, "ignore_eos": true, - "skip_special_tokens": false + "skip_special_tokens": false, + "stream_interval": 4 })) .expect("parse request"); @@ -436,6 +440,10 @@ mod tests { ); assert!(prepared.text_request.sampling_params.ignore_eos); assert!(!prepared.text_request.decode_options.skip_special_tokens); + assert_eq!( + prepared.text_request.sampling_params.stream_interval, + NonZeroU32::new(4) + ); } #[test] diff --git a/rust/src/server/src/routes/openai/completions/types.rs b/rust/src/server/src/routes/openai/completions/types.rs index 2a6901370682..b10f8a23d701 100644 --- a/rust/src/server/src/routes/openai/completions/types.rs +++ b/rust/src/server/src/routes/openai/completions/types.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::HashMap; +use std::num::NonZeroU32; use std::sync::Arc; use serde::{Deserialize, Serialize}; @@ -114,6 +115,12 @@ pub struct CompletionRequest { /// Parameters for detecting repetitive N-gram patterns in output tokens pub repetition_detection: Option, + /// Number of tokens to batch into each streamed chunk. Raises the server's + /// `--stream-interval` for this request. Values below the server setting + /// are clamped up to it. The first and last chunks are always sent + /// immediately. Ignored for non-streaming requests. + pub stream_interval: Option, + /// Length penalty for beam search pub length_penalty: Option, diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index 2b09070049df..9fa8c78f4f91 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -124,6 +124,7 @@ pub fn lower_sampling_params( structured_outputs, skip_reading_prefix_cache, vllm_xargs, + stream_interval, } = sampling_params; validate_logprobs( @@ -214,6 +215,7 @@ pub fn lower_sampling_params( skip_reading_prefix_cache, extra_args: vllm_xargs, routed_experts_prompt_start: 0, + stream_interval, }; validate_resolved_sampling_params(¶ms)?; validate_vocab_range(¶ms, &sampling_limits)?; @@ -341,6 +343,7 @@ fn merge_unique_token_ids( #[cfg(test)] mod tests { use std::collections::{BTreeSet, HashMap}; + use std::num::NonZeroU32; use serial_test::file_serial; use vllm_engine_core_client::protocol::kv_hints::{KvHintAction, KvHintsEnvelope}; @@ -693,6 +696,7 @@ mod tests { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } "#]] .assert_debug_eq(¶ms); @@ -746,6 +750,7 @@ mod tests { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } "#]] .assert_debug_eq(¶ms); @@ -921,6 +926,7 @@ mod tests { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } "#]] .assert_debug_eq(¶ms); @@ -992,6 +998,7 @@ mod tests { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } "#]] .assert_debug_eq(¶ms); @@ -1056,6 +1063,7 @@ mod tests { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } "#]] .assert_debug_eq(¶ms); @@ -1401,6 +1409,7 @@ mod tests { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } "#]] .assert_debug_eq(¶ms); @@ -1430,6 +1439,26 @@ mod tests { assert_eq!(prepared.generate_request.request_id, "text-1"); } + #[test] + fn lower_text_request_carries_stream_interval_in_sampling_params() { + let mut request = sample_request(); + request.sampling_params.stream_interval = NonZeroU32::new(4); + + let prepared = lower_text_request( + request, + vec![1, 2, 3], + sample_sampling_hints(), + sample_sampling_limits(), + &stub_tokenizer(), + ) + .unwrap(); + + assert_eq!( + prepared.generate_request.sampling_params.stream_interval, + NonZeroU32::new(4) + ); + } + #[test] fn lower_text_request_passes_kv_hints_through() { let hints = KvHintsEnvelope { diff --git a/rust/src/text/src/request.rs b/rust/src/text/src/request.rs index 08758748ac79..b00d8f3b2f12 100644 --- a/rust/src/text/src/request.rs +++ b/rust/src/text/src/request.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::HashMap; +use std::num::NonZeroU32; use enum_as_inner::EnumAsInner; use serde::{Deserialize, Serialize}; @@ -147,6 +148,11 @@ pub struct SamplingParams { pub skip_reading_prefix_cache: Option, /// Additional request parameters for custom extensions. pub vllm_xargs: Option>, + /// Number of newly generated tokens to batch into each streamed output. + /// Raises the interval above the frontend-level `--stream-interval`; + /// values below it are clamped up to it. The first and final outputs are + /// always emitted immediately. + pub stream_interval: Option, } #[allow(clippy::derivable_impls)] // more explicit @@ -180,6 +186,7 @@ impl Default for SamplingParams { structured_outputs: None, skip_reading_prefix_cache: None, vllm_xargs: None, + stream_interval: None, } } } From 51144dd83ec557afa6fc6032bbf03d7f50a55b89 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Sat, 10 Oct 2026 05:22:31 +0800 Subject: [PATCH 206/624] [Bugfix][Spec Decode] Trim token_ids/logprobs left past a stop string under speculative decoding (#47616) Signed-off-by: Ting Sun Signed-off-by: Nick Hill Co-authored-by: Nick Hill Co-authored-by: Claude Opus 5.5 --- ...stop_string_while_stop_model_terminates.py | 113 +++++++++++++++++- vllm/v1/engine/detokenizer.py | 24 +++- vllm/v1/engine/output_processor.py | 34 +++++- 3 files changed, 158 insertions(+), 13 deletions(-) diff --git a/tests/detokenizer/test_stop_string_while_stop_model_terminates.py b/tests/detokenizer/test_stop_string_while_stop_model_terminates.py index 3b72d612dd1c..38d860969f49 100644 --- a/tests/detokenizer/test_stop_string_while_stop_model_terminates.py +++ b/tests/detokenizer/test_stop_string_while_stop_model_terminates.py @@ -1,11 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import numpy as np import pytest from vllm.sampling_params import SamplingParams -from vllm.v1.engine import EngineCoreRequest +from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason from vllm.v1.engine.detokenizer import BaseIncrementalDetokenizer +from vllm.v1.engine.output_processor import OutputProcessor +from vllm.v1.outputs import LogprobsLists, SamplingMaskLists @pytest.fixture(params=[True, False]) @@ -82,7 +85,7 @@ def test_stop_string_while_stop_token_terminates(include_stop_str_in_output: boo # Simulate that the last token ('Z') is a stop token (stop_terminated=True). result = detok.update(new_token_ids=token_ids, stop_terminated=True) - # The update should not report a stop string + # The update should report the matched stop string. assert result == stop_string # Output text should reflect stop-string handling: @@ -91,9 +94,111 @@ def test_stop_string_while_stop_token_terminates(include_stop_str_in_output: boo expected_text = "abcd" if include_stop_str_in_output else "ab" assert detok.output_text == expected_text - # The skipped final token should still be recorded in token_ids. - assert detok.output_token_ids == token_ids + # Tokens after the stop string, including the stop token, are dropped. + assert detok.output_token_ids == token_ids[:4] # get_next_output_text should return the full text when finished=True. # (Buffering only applies during streaming when finished=False.) assert detok.get_next_output_text(finished=True, delta=False) == expected_text + + +@pytest.mark.parametrize( + "min_tokens,num_kept", + [(0, 2), (1, 2), (2, 4), (3, 4), (4, None)], +) +@pytest.mark.parametrize("num_prev_tokens", [0, 1]) +def test_stop_string_min_tokens_within_step( + include_stop_str_in_output: bool, min_tokens, num_kept, num_prev_tokens +): + """Stop strings completing within min_tokens are ignored, even when the + min_tokens boundary falls in the middle of a multi-token step.""" + token_ids = [ord(c) for c in "cdcdx"] + req = _make_request( + stop=["cd"], + include_stop_str_in_output=include_stop_str_in_output, + min_tokens=min_tokens, + ) + detok = _DummyDetokenizer(req) + + if num_prev_tokens: + assert detok.update(token_ids[:num_prev_tokens], False) is None + result = detok.update(token_ids[num_prev_tokens:], False) + + if num_kept is None: + assert result is None + assert detok.output_text == "cdcdx" + assert detok.output_token_ids == token_ids + return + assert result == "cd" + text_len = num_kept if include_stop_str_in_output else num_kept - 2 + assert detok.output_text == "cdcdx"[:text_len] + assert detok.output_token_ids == token_ids[:num_kept] + + +@pytest.mark.parametrize( + "decoded_tokens,keep", + [ + (["a", "b", "c", "d", "e", "f"], 4), + (["ab", "cdx", "ef"], 2), + (["ab", "", "c", "", "d", "", "ef"], 5), + ], +) +@pytest.mark.parametrize("stop_terminated", [False, True]) +def test_stop_string_trims_speculative_overflow( + include_stop_str_in_output: bool, + decoded_tokens, + keep, + stop_terminated: bool, + monkeypatch, +): + if stop_terminated: + # The engine also stopped on a stop token at the end of the batch. + decoded_tokens = decoded_tokens + [""] + token_ids = list(range(len(decoded_tokens))) + stop_string = "cd" + expected_token_ids = token_ids[:keep] + + req = _make_request( + stop=[stop_string], include_stop_str_in_output=include_stop_str_in_output + ) + req.external_req_id = req.request_id + assert req.sampling_params is not None + req.sampling_params.logprobs = 0 + detok = _DummyDetokenizer(req) + monkeypatch.setattr(detok, "decode_next", decoded_tokens.__getitem__) + processor = OutputProcessor(tokenizer=None, log_stats=False) + processor.add_request(req, prompt=None) + processor.request_states[req.request_id].detokenizer = detok + + outputs = processor.process_outputs( + [ + EngineCoreOutput( + request_id=req.request_id, + new_token_ids=token_ids, + finish_reason=FinishReason.STOP if stop_terminated else None, + new_logprobs=LogprobsLists( + np.array(token_ids).reshape(-1, 1), + np.zeros((len(token_ids), 1)), + np.ones(len(token_ids), dtype=int), + None, + ), + new_sampling_mask=SamplingMaskLists( + np.array(token_ids), np.arange(len(token_ids) + 1) + ), + routed_experts=np.array(token_ids).reshape(-1, 1, 1), + ) + ] + ) + result = outputs.request_outputs[0].outputs[0] + + assert result.stop_reason == stop_string + expected_text = "abcd" if include_stop_str_in_output else "ab" + assert detok.output_text == expected_text + assert detok.output_token_ids == expected_token_ids + assert result.token_ids == expected_token_ids + assert result.logprobs is not None + assert len(result.logprobs) == len(expected_token_ids) + assert result.sampling_mask is not None + assert result.sampling_mask.token_ids == [[token] for token in expected_token_ids] + assert result.routed_experts is not None + assert result.routed_experts.ravel().tolist() == expected_token_ids diff --git a/vllm/v1/engine/detokenizer.py b/vllm/v1/engine/detokenizer.py index c27e01895cac..adc0be31584b 100644 --- a/vllm/v1/engine/detokenizer.py +++ b/vllm/v1/engine/detokenizer.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import sys from abc import ABC, abstractmethod +from bisect import bisect_left import tokenizers.decoders from tokenizers import Tokenizer @@ -112,21 +113,27 @@ def update(self, new_token_ids: list[int], stop_terminated: bool) -> str | None: skipped_stop_token_id = None # 1) Detokenize the new token ids incrementally. - stop_check_offset = len(self.output_text) + num_prev_tokens = self.num_output_tokens() + first_new_token_index = len(self.token_ids) + # Output text offsets at the boundaries of the new tokens. + token_offsets = [len(self.output_text)] if self.stop else None for new_token_id in new_token_ids: self.token_ids.append(new_token_id) self.output_text += self.decode_next(new_token_id) - # Support min_tokens, see https://github.com/vllm-project/vllm/pull/22014 - if self.min_tokens and self.num_output_tokens() <= self.min_tokens: - stop_check_offset = len(self.output_text) + if token_offsets is not None: + token_offsets.append(len(self.output_text)) if skipped_stop_token_id is not None: # Cleanup after skipping detokenization. self.token_ids.append(skipped_stop_token_id) + if token_offsets is not None: + token_offsets.append(len(self.output_text)) # 2) Evaluate stop strings. stop_string = None - if self.stop and self.num_output_tokens() > self.min_tokens: + if token_offsets is not None and self.num_output_tokens() > self.min_tokens: + # Support min_tokens, see https://github.com/vllm-project/vllm/pull/22014 + stop_check_offset = token_offsets[max(self.min_tokens - num_prev_tokens, 0)] stop = check_stop_strings( output_text=self.output_text, new_char_count=len(self.output_text) - stop_check_offset, @@ -135,6 +142,13 @@ def update(self, new_token_ids: list[int], stop_terminated: bool) -> str | None: ) if stop is not None: stop_string, truncate_to = stop + stop_end = len(self.output_text) if truncate_to == -1 else truncate_to + if not self.include_stop_str_in_output: + stop_end += len(stop_string) + # Drop tokens generated after the one completing the stop string + # (possible when multiple tokens are appended in one step). + num_kept = bisect_left(token_offsets, stop_end, lo=1) + del self.token_ids[first_new_token_index + num_kept :] if truncate_to != -1: self.output_text = self.output_text[:truncate_to] diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index c22c5474c7c0..abb45f7239be 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -736,17 +736,22 @@ def process_outputs( if pooling_output is None: assert req_state.detokenizer is not None assert req_state.logprobs_processor is not None - if engine_core_output.new_sampling_mask is not None: - req_state.sampling_mask_chunks.append( - engine_core_output.new_sampling_mask - ) # 2) Detokenize the token ids into text and perform stop checks. + num_prev_tokens = req_state.detokenizer.num_output_tokens() stop_string = req_state.detokenizer.update( new_token_ids, finish_reason == FinishReason.STOP ) if stop_string: finish_reason = FinishReason.STOP stop_reason = stop_string + self._trim_surplus_tokens( + req_state, engine_core_output, num_prev_tokens + ) + + if engine_core_output.new_sampling_mask is not None: + req_state.sampling_mask_chunks.append( + engine_core_output.new_sampling_mask + ) # 3) Compute sample and prompt logprobs for request, # if required. @@ -799,6 +804,27 @@ def process_outputs( reqs_to_abort=reqs_to_abort, ) + @staticmethod + def _trim_surplus_tokens( + req_state: RequestState, output: EngineCoreOutput, num_prev_tokens: int + ) -> None: + """Drop new tokens generated past a stop string matched by the + detokenizer, e.g. when multiple tokens are accepted in one step.""" + assert req_state.detokenizer is not None + num_kept = req_state.detokenizer.num_output_tokens() - num_prev_tokens + num_dropped = len(output.new_token_ids) - num_kept + if num_dropped <= 0: + return + del output.new_token_ids[num_kept:] + if output.new_logprobs is not None: + output.new_logprobs = output.new_logprobs.slice_request(0, num_kept) + if output.new_sampling_mask is not None: + output.new_sampling_mask = output.new_sampling_mask.slice_request( + 0, num_kept + ) + if output.routed_experts is not None: + req_state.routed_experts_chunks[-1] = output.routed_experts[:-num_dropped] + def _finish_request(self, req_state: RequestState) -> None: req_id = req_state.request_id self.request_states.pop(req_id) From ed7ae705a1f366c2008964af562a7543e91d934e Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 9 Oct 2026 14:38:11 -0700 Subject: [PATCH 207/624] [Bugfix][MRV2] Use int64 index mappings to fix row-offset overflow (#60736) Signed-off-by: Nick Hill Co-authored-by: vllmellm Co-authored-by: Claude Opus 5.5 (1M context) --- tests/v1/worker/test_gpu_row_offsets_i64.py | 203 ++++++++++++++++++ vllm/v1/worker/gpu/buffer_utils.py | 2 +- vllm/v1/worker/gpu/input_batch.py | 6 +- vllm/v1/worker/gpu/model_runner.py | 2 +- vllm/v1/worker/gpu/pcp_manager.py | 4 +- vllm/v1/worker/gpu/sample/batch_shard.py | 20 +- vllm/v1/worker/gpu/sample/penalties.py | 2 +- .../gpu/spec_decode/dflash/speculator.py | 2 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 2 +- .../spec_decode/standalone_ar/speculator.py | 2 +- 10 files changed, 221 insertions(+), 24 deletions(-) create mode 100644 tests/v1/worker/test_gpu_row_offsets_i64.py diff --git a/tests/v1/worker/test_gpu_row_offsets_i64.py b/tests/v1/worker/test_gpu_row_offsets_i64.py new file mode 100644 index 000000000000..374332e2dafa --- /dev/null +++ b/tests/v1/worker/test_gpu_row_offsets_i64.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""int64 row offsets in the MRV2 staged-write and idx-mapping kernels. + +With int32 request-slot indices, `slot * stride` wraps once the product +crosses 2**31 (issue #57030): at max_model_len=1M from row 2148, and from +row 716 for the 3-dim mrope positions. Each test runs a control row below +and an overflow row above the threshold in the same kernel launch. +Requires CUDA/ROCm and ~11 GiB free device memory: the tensors must exceed +2**31 elements by construction. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.utils.math_utils import cdiv +from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor +from vllm.v1.worker.gpu.input_batch import post_update +from vllm.v1.worker.gpu.mm.rope import RopeState +from vllm.v1.worker.gpu.sample.penalties import bincount + +STRIDE = 1_000_000 # max_model_len +ROW = 2559 # ROW * STRIDE > 2**31 (first overflowing row at 1M: 2148) +NUM_ROWS = ROW + 1 # (2560, 1M) int32 = 10.24 GB +NEEDED_BYTES = NUM_ROWS * STRIDE * 4 +VOCAB_SIZE = 8192 +PROMPT_LEN = 1024 +PREFILL_LEN = 2048 + +# mrope: flat row = NUM_DIMS * req + j and stride0 = NUM_DIMS * STRIDE, so +# the overflow threshold drops to row 716 of a 1M-row buffer. +NUM_DIMS = 3 +ROPE_ROW = 716 # NUM_DIMS * ROPE_ROW * STRIDE > 2**31 + +pytestmark = [ + pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires CUDA or ROCm" + ), + pytest.mark.skipif( + not current_platform.is_cuda_alike() + or torch.accelerator.get_memory_info(0)[0] + < NEEDED_BYTES + 2**30, # slack for the outputs + reason=f"requires >= {NEEDED_BYTES / 2**30:.0f} GiB free memory", + ), +] + + +def test_staged_write_high_row(): + """Apply-write past the int32 offset limit must land on the target row.""" + device = torch.device("cuda:0") + content = list(range(9000, 9016)) + staged = StagedWriteTensor((NUM_ROWS, STRIDE), torch.int32, device) + staged.stage_write(ROW, 0, content) + staged.apply_write() + torch.accelerator.synchronize() + + assert staged.gpu[ROW, : len(content)].tolist() == content + # A wrapped offset must not clobber row 0 (silent-corruption variant). + assert staged.gpu[0, : len(content)].tolist() == [0] * len(content) + + +def test_bincount_high_row(): + """penalties.bincount past the int32 offset limit must hit the right row.""" + device = torch.device("cuda:0") + tokens = torch.arange(PREFILL_LEN, dtype=torch.int32, device=device) + + all_token_ids = torch.zeros(NUM_ROWS, STRIDE, dtype=torch.int32, device=device) + all_token_ids[0, :PREFILL_LEN] = tokens + all_token_ids[ROW, :PREFILL_LEN] = tokens + + prompt_len = torch.zeros(NUM_ROWS, dtype=torch.int32, device=device) + prefill_len = torch.zeros(NUM_ROWS, dtype=torch.int32, device=device) + prompt_len[[0, ROW]] = PROMPT_LEN + prefill_len[[0, ROW]] = PREFILL_LEN + + idx_mapping = torch.tensor([0, ROW], dtype=torch.int64, device=device) + prompt_bin_mask = torch.zeros( + NUM_ROWS, cdiv(VOCAB_SIZE, 32), dtype=torch.int32, device=device + ) + output_bin_counts = torch.zeros( + NUM_ROWS, VOCAB_SIZE, dtype=torch.int32, device=device + ) + + bincount( + idx_mapping, + all_token_ids, + prompt_len, + prefill_len, + prompt_bin_mask, + output_bin_counts, + max_prefill_len=PREFILL_LEN, + ) + torch.accelerator.synchronize() + + assert torch.equal(prompt_bin_mask[ROW], prompt_bin_mask[0]) + assert torch.equal(output_bin_counts[ROW], output_bin_counts[0]) + # Sanity: the control row got prompt bits and output counts. + assert prompt_bin_mask[0, : cdiv(PROMPT_LEN, 32)].any().item() + assert ( + output_bin_counts[0, PROMPT_LEN:PREFILL_LEN].sum().item() + == PREFILL_LEN - PROMPT_LEN + ) + + +def test_post_update_high_row(): + """post_update (runs every sampled token) must store on the mapped row.""" + device = torch.device("cuda:0") + # Big buffer first so a wrapped (negative) offset cannot land in-bounds. + all_token_ids = torch.zeros(NUM_ROWS, STRIDE, dtype=torch.int32, device=device) + num_computed = torch.zeros(NUM_ROWS, dtype=torch.int32, device=device) + last_sampled = torch.full((NUM_ROWS, 1), -1, dtype=torch.int64, device=device) + total_len = torch.zeros(NUM_ROWS, dtype=torch.int32, device=device) + total_len[[0, ROW]] = torch.tensor([50, 100], dtype=torch.int32, device=device) + + idx_mapping = torch.tensor([0, ROW], dtype=torch.int64, device=device) + sampled_tokens = torch.tensor([[111], [222]], dtype=torch.int64, device=device) + num_sampled = torch.tensor([1, 1], dtype=torch.int32, device=device) + num_rejected = torch.tensor([0, 0], dtype=torch.int32, device=device) + query_start_loc = torch.tensor([0, 1, 2], dtype=torch.int32, device=device) + + post_update( + idx_mapping, + num_computed, + last_sampled, + None, # output_bin_counts: a second product of the same idx load + sampled_tokens, + num_sampled, + num_rejected, + query_start_loc, + all_token_ids, + total_len, + ) + torch.accelerator.synchronize() + + assert int(all_token_ids[0, 50]) == 111 + assert int(all_token_ids[ROW, 100]) == 222 + # Each mapped row holds exactly its own token; unmapped rows stay zero. + assert int(all_token_ids[0].sum()) == 111 + assert int(all_token_ids[ROW].sum()) == 222 + assert int(all_token_ids[1].sum()) == 0 + assert int(last_sampled[0, 0]) == 111 + assert int(last_sampled[ROW, 0]) == 222 + assert int(total_len[0]) == 51 + assert int(total_len[ROW]) == 101 + assert int(num_computed[0]) == 1 + assert int(num_computed[ROW]) == 1 + + +def test_rope_positions_high_row(): + """M-RoPE prepare_positions must read staged rows past the int32 limit. + + Threshold is 3x lower than the other tests (stride0 = num_dims * 1M), so + this is the variant that trips at default max_num_seqs on 1M context. + """ + device = torch.device("cuda:0") + max_num_reqs = NUM_ROWS // NUM_DIMS + base = torch.arange(PREFILL_LEN, dtype=torch.int32, device=device) + # (num_dims * max_num_reqs, max_model_len), logically [req][dim][pos]. + prefill_positions = torch.zeros(NUM_ROWS, STRIDE, dtype=torch.int32, device=device) + # Distinct per-(dim, request) patterns so mixed-up rows are detectable. + for j in range(NUM_DIMS): + prefill_positions[j, :PREFILL_LEN] = 1_000 * (j + 1) + base + target_row = NUM_DIMS * ROPE_ROW + j + prefill_positions[target_row, :PREFILL_LEN] = 7_000 * (j + 1) + base + + # Skeletal RopeState: the real __init__ pins num_dims * max_num_reqs host + # memory, which the stride arithmetic does not depend on. + rope = object.__new__(RopeState) + rope.num_dims = NUM_DIMS + rope.max_num_reqs = max_num_reqs + rope.max_model_len = STRIDE + rope.prefill_positions = SimpleNamespace(gpu=prefill_positions) + rope.positions = torch.zeros( + (NUM_DIMS, 2 * PREFILL_LEN + 1), dtype=torch.int64, device=device + ) + rope.prefill_delta = SimpleNamespace( + gpu=torch.zeros(max_num_reqs, dtype=torch.int32, device=device) + ) + + idx_mapping = torch.tensor([0, ROPE_ROW], dtype=torch.int64, device=device) + query_start_loc = torch.tensor( + [0, PREFILL_LEN, 2 * PREFILL_LEN], dtype=torch.int32, device=device + ) + prefill_len = torch.zeros(max_num_reqs, dtype=torch.int32, device=device) + prefill_len[[0, ROPE_ROW]] = PREFILL_LEN + num_computed = torch.zeros(max_num_reqs, dtype=torch.int32, device=device) + + rope.prepare_positions(idx_mapping, query_start_loc, prefill_len, num_computed) + torch.accelerator.synchronize() + + expect_control = torch.stack([1_000 * (j + 1) + base for j in range(NUM_DIMS)]).to( + torch.int64 + ) + expect_target = torch.stack([7_000 * (j + 1) + base for j in range(NUM_DIMS)]).to( + torch.int64 + ) + assert torch.equal(rope.positions[:, :PREFILL_LEN], expect_control) + assert torch.equal(rope.positions[:, PREFILL_LEN : 2 * PREFILL_LEN], expect_target) + # Padding column stays zero. + assert bool((rope.positions[:, 2 * PREFILL_LEN :] == 0).all()) diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index 32d997358c13..7364156848bf 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/vllm/v1/worker/gpu/buffer_utils.py @@ -310,7 +310,7 @@ def _apply_write_kernel( MULTI_GROUP: tl.constexpr, ): pid = tl.program_id(0) - row_idx = tl.load(write_indices_ptr + pid) + row_idx = tl.load(write_indices_ptr + pid).to(tl.int64) start_idx = tl.load(write_starts_ptr + pid) cu_start = tl.load(write_cu_lens_ptr + pid - 1) if pid > 0 else 0 diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index ea1c925dadc0..bd55a43f6412 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -137,7 +137,7 @@ def make_dummy( req_ids = [f"req_{i}_{random_uuid()}" for i in range(num_reqs)] idx_mapping_np = np.arange(num_reqs, dtype=np.intp) - idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device) + idx_mapping = torch.arange(num_reqs, dtype=torch.int64, device=device) expanded_idx_mapping = idx_mapping expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=device) @@ -715,7 +715,9 @@ def expand_idx_mapping( ) -> tuple[torch.Tensor, torch.Tensor]: num_reqs = idx_mapping.shape[0] expanded_idx_mapping = idx_mapping.new_empty(total_num_logits) - expanded_local_pos = idx_mapping.new_empty(total_num_logits) + expanded_local_pos = torch.empty( + total_num_logits, dtype=torch.int32, device=idx_mapping.device + ) _expand_idx_mapping_kernel[(num_reqs,)]( idx_mapping, expanded_idx_mapping, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 134b4bfb0e69..2e9925f1ad1c 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1361,7 +1361,7 @@ def prepare_inputs( num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens idx_mapping_np = batch_req_state.idx_mapping_np idx_mapping = async_tensor_h2d( - idx_mapping_np, device=self.device, dtype=torch.int32 + idx_mapping_np, device=self.device, dtype=torch.int64 ) num_reqs = len(req_ids) diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index 6aec96bfe14b..aa245fa267c4 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -653,9 +653,7 @@ def partition_batch( idx_mapping=local_to_global_req_idx, idx_mapping_np=local_to_global_req_idx_np, expanded_idx_mapping=local_to_global_req_idx, - expanded_local_pos=torch.zeros( - num_local_reqs, dtype=torch.int32, device=self.device - ), + expanded_local_pos=input_batch.expanded_local_pos.new_zeros(num_local_reqs), num_scheduled_tokens=local_num_scheduled_tokens, num_tokens=num_local_tokens, num_tokens_after_padding=num_local_tokens_padded, diff --git a/vllm/v1/worker/gpu/sample/batch_shard.py b/vllm/v1/worker/gpu/sample/batch_shard.py index a7e369b1886a..695b1d543563 100644 --- a/vllm/v1/worker/gpu/sample/batch_shard.py +++ b/vllm/v1/worker/gpu/sample/batch_shard.py @@ -214,21 +214,15 @@ def shard_sampler_inputs( local_logits_indices = torch.empty( num_local_logits, dtype=torch.int64, device=self.device ) - local_idx_mapping = torch.empty( - num_local_reqs, dtype=torch.int32, device=self.device + local_idx_mapping = input_batch.idx_mapping.new_empty(num_local_reqs) + local_cu_num_logits = input_batch.cu_num_logits.new_empty(num_local_reqs + 1) + local_expanded_idx_mapping = input_batch.expanded_idx_mapping.new_empty( + num_local_logits ) - local_cu_num_logits = torch.empty( - num_local_reqs + 1, dtype=torch.int32, device=self.device - ) - local_expanded_idx_mapping = torch.empty( - num_local_logits, dtype=torch.int32, device=self.device - ) - local_expanded_local_pos = torch.empty( - num_local_logits, dtype=torch.int32, device=self.device - ) - local_seq_lens = torch.empty( - num_local_reqs, dtype=torch.int32, device=self.device + local_expanded_local_pos = input_batch.expanded_local_pos.new_empty( + num_local_logits ) + local_seq_lens = input_batch.seq_lens.new_empty(num_local_reqs) if num_reqs > 0: _build_shard_plan_kernel[(num_reqs,)]( input_batch.idx_mapping, diff --git a/vllm/v1/worker/gpu/sample/penalties.py b/vllm/v1/worker/gpu/sample/penalties.py index eb751e8be355..5df7c60afd56 100644 --- a/vllm/v1/worker/gpu/sample/penalties.py +++ b/vllm/v1/worker/gpu/sample/penalties.py @@ -62,7 +62,7 @@ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool: def apply_staged_writes(self) -> None: if self._new_penalties_reqs: idx_mapping = async_tensor_h2d( - self._new_penalties_reqs, dtype=torch.int32, device=self.device + self._new_penalties_reqs, dtype=torch.int64, device=self.device ) prefill_lens = self.req_states.prefill_len.np[self._new_penalties_reqs] diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index ec46039533d7..9538dcb4d4bf 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -102,7 +102,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # full buffer before a real batch has populated it, so zero would make # every padding row scatter into request slot 0. self.sample_idx_mapping = torch.full( - (max_num_sampled_tokens,), -1, dtype=torch.int32, device=device + (max_num_sampled_tokens,), -1, dtype=torch.int64, device=device ) # [0, 1, ..., N-1, 0, 1, ..., N-1, ...] -> the per-token column index into # draft_logits[req, step, :]. diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index d07c043de915..cbd94ed0542d 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -152,7 +152,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): device=device, ) self.idx_mapping = torch.zeros( - self.max_num_reqs, dtype=torch.int32, device=device + self.max_num_reqs, dtype=torch.int64, device=device ) self.temperature = torch.zeros( self.max_num_reqs, dtype=torch.float32, device=device diff --git a/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py index 8bdfb8530be6..898de1e3f886 100644 --- a/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py @@ -402,7 +402,7 @@ def _prepare_prefill_inputs_kernel( target_input_ids_ptr, # [src_tokens] int32 target_positions_ptr, # [src_tokens] int64 last_sampled_ptr, # [max_num_reqs] int32 - idx_mapping_ptr, # [num_reqs] int32 + idx_mapping_ptr, # [num_reqs] int64 out_input_ids_ptr, # [src_tokens + num_reqs] int32 out_positions_ptr, # [src_tokens + num_reqs] int64 last_token_indices_ptr, # [max_num_reqs] int64 From 783ac12071ec319d7d3e5af29d039034115d8934 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Fri, 9 Oct 2026 15:43:37 -0700 Subject: [PATCH 208/624] [CI] Accept whitespace around /ci commands and reply instead of going silent (#60925) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- .github/workflows/run-ci-command.yml | 29 +++++++------- .github/workflows/scripts/run_ci_command.py | 34 ++++++++++++++++- .../workflows/scripts/test_run_ci_command.py | 38 ++++++++++++++++++- 3 files changed, 82 insertions(+), 19 deletions(-) diff --git a/.github/workflows/run-ci-command.yml b/.github/workflows/run-ci-command.yml index 3435148ea4e3..b0764bfa0683 100644 --- a/.github/workflows/run-ci-command.yml +++ b/.github/workflows/run-ci-command.yml @@ -18,22 +18,8 @@ jobs: if: >- github.repository_owner == 'vllm-project' && github.event.issue.pull_request && - (github.event.comment.body == '/ci run' || - github.event.comment.body == '/ci run all' || - github.event.comment.body == '/ci run nightly' || - github.event.comment.body == '/ci run --allow-stale' || - github.event.comment.body == '/ci run all --allow-stale' || - github.event.comment.body == '/ci run nightly --allow-stale' || - github.event.comment.body == '/ci retry' || - github.event.comment.body == '/ci cancel' || - github.event.comment.body == '/amd-ci run' || - github.event.comment.body == '/amd-ci run all' || - github.event.comment.body == '/amd-ci run nightly' || - github.event.comment.body == '/amd-ci run --allow-stale' || - github.event.comment.body == '/amd-ci run all --allow-stale' || - github.event.comment.body == '/amd-ci run nightly --allow-stale' || - github.event.comment.body == '/amd-ci retry' || - github.event.comment.body == '/amd-ci cancel') + (startsWith(github.event.comment.body, '/ci') || + startsWith(github.event.comment.body, '/amd-ci')) runs-on: [self-hosted, linux, x64, vllm-runners] steps: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 @@ -54,3 +40,14 @@ jobs: BUILDKITE_PIPELINE: ci CI_TRUSTED_USERS: ${{ vars.CI_TRUSTED_USERS }} GH_TOKEN: ${{ github.token }} + - name: Report failure on the PR + if: failure() + uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 + with: + script: | + const runUrl = `${context.serverUrl}/${context.repo.owner}/${context.repo.repo}/actions/runs/${context.runId}`; + await github.rest.issues.createComment({ + ...context.repo, + issue_number: context.issue.number, + body: `❌ @${context.payload.comment.user.login}, your CI command failed with an unexpected error, so CI may not have started. See the [workflow run](${runUrl}) for details.`, + }); diff --git a/.github/workflows/scripts/run_ci_command.py b/.github/workflows/scripts/run_ci_command.py index 530bd697d008..537efa83c4a5 100644 --- a/.github/workflows/scripts/run_ci_command.py +++ b/.github/workflows/scripts/run_ci_command.py @@ -58,6 +58,7 @@ RETRY_COMMANDS = frozenset({COMMAND_RETRY_FAILED, COMMAND_RETRY_AMD_FAILED}) CANCEL_COMMANDS = frozenset({COMMAND_CANCEL_CI, COMMAND_CANCEL_AMD_CI}) ALL_CI_COMMANDS = UPSTREAM_CI_COMMANDS | AMD_CI_COMMANDS +CI_COMMAND_PREFIXES = frozenset({"/ci", "/amd-ci"}) CI_AUTHORIZED_COMMENT_MARKER = "" READY_LABELS = {"ready", "ready-run-all-tests"} TRUSTED_PERMISSIONS = {"admin", "maintain", "write"} @@ -461,11 +462,17 @@ def list_failed_jobs(self, build_number: int) -> list[dict[str, Any]]: def parse_command(body: str) -> str | None: + body = body.strip() if body in ALL_CI_COMMANDS: return body return None +def is_ci_command_attempt(body: str) -> bool: + words = body.split(maxsplit=1) + return bool(words) and words[0].casefold() in CI_COMMAND_PREFIXES + + def pipeline_for_command( command: str, *, @@ -1086,6 +1093,28 @@ def handle_cancel_ci( ) +def reply_unrecognized_command( + event: Mapping[str, Any], + github: GitHubClient, +) -> None: + if "pull_request" not in event["issue"]: + return + issue_number = event["issue"]["number"] + comment_id = event["comment"]["id"] + actor = event["comment"]["user"]["login"] + if is_already_handled(github, issue_number, comment_id): + print(f"Comment {comment_id} was already handled.") + return + add_reaction_safely(github, comment_id, "-1") + commands = ", ".join(f"`{command}`" for command in sorted(ALL_CI_COMMANDS)) + github.add_comment( + issue_number, + f"❌ @{actor}, that is not a recognized CI command, so no CI was " + "started. Post a new comment containing only one of: " + f"{commands}.\n\n{command_comment_marker(comment_id)}", + ) + + def run( event: Mapping[str, Any], github: GitHubClient, @@ -1217,8 +1246,11 @@ def main() -> None: ) return - command = parse_command(event["comment"]["body"]) + body = event["comment"]["body"] + command = parse_command(body) if not command: + if is_ci_command_attempt(body): + reply_unrecognized_command(event, github) return pipeline = pipeline_for_command( diff --git a/.github/workflows/scripts/test_run_ci_command.py b/.github/workflows/scripts/test_run_ci_command.py index 9750cefaa394..623b41247641 100644 --- a/.github/workflows/scripts/test_run_ci_command.py +++ b/.github/workflows/scripts/test_run_ci_command.py @@ -32,10 +32,12 @@ has_trusted_approval, is_active_build, is_build_for_pr, + is_ci_command_attempt, notify_authorized, parse_command, parse_trusted_users, pipeline_for_command, + reply_unrecognized_command, resolve_workflow_run_pr, run, select_latest_build, @@ -340,14 +342,46 @@ def test_only_exact_ci_commands_are_accepted(self) -> None: self.assertIsNone(parse_command("/ci run please")) self.assertIsNone(parse_command("/ci run all please")) self.assertIsNone(parse_command("/ci cancel please")) - self.assertIsNone(parse_command(" /ci run")) self.assertIsNone(parse_command("/amd-ci run please")) - self.assertIsNone(parse_command("/amd-ci retry ")) self.assertIsNone(parse_command("/AMD-CI run")) self.assertIsNone(parse_command("/amdci run")) self.assertIsNone(parse_command("/ci run --allow-stale all")) self.assertIsNone(parse_command("/amd-ci run --allow-stale please")) + def test_surrounding_whitespace_is_ignored(self) -> None: + self.assertEqual(parse_command("/ci run\r\n"), COMMAND_RUN_CI) + self.assertEqual(parse_command(" /ci run "), COMMAND_RUN_CI) + self.assertEqual(parse_command("/amd-ci retry\n"), COMMAND_RETRY_AMD_FAILED) + + def test_ci_command_attempts_are_detected(self) -> None: + for body in ("/ci", "/ci rn", "/CI run", "/amd-ci run please", " /ci\nrun"): + with self.subTest(body=body): + self.assertTrue(is_ci_command_attempt(body)) + for body in ("", "LGTM", "/cifoo", "/amdci run", "please /ci run"): + with self.subTest(body=body): + self.assertFalse(is_ci_command_attempt(body)) + + def test_unrecognized_command_gets_feedback(self) -> None: + github = FakeGitHub() + + reply_unrecognized_command(make_event("/ci rn", actor="author"), github) + + self.assertEqual(github.reactions, ["-1"]) + self.assertEqual(len(github.comments), 1) + self.assertIn( + "@author, that is not a recognized CI command", github.comments[0] + ) + self.assertIn(f"`{COMMAND_RUN_CI}`", github.comments[0]) + self.assertIn("", github.comments[0]) + + def test_unrecognized_command_feedback_is_not_repeated(self) -> None: + github = FakeGitHub(comments=["❌ earlier\n\n"]) + + reply_unrecognized_command(make_event("/ci rn"), github) + + self.assertEqual(len(github.comments), 1) + self.assertEqual(github.reactions, []) + def test_commands_select_only_their_configured_pipeline(self) -> None: cases = ( (COMMAND_RUN_CI, "ci"), From 4c43254d52785d7a69300fdfb6c79987200c277c Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Fri, 9 Oct 2026 18:45:34 -0400 Subject: [PATCH 209/624] [HiSparse] Report host-pool block residency as HiSparse metrics (#59485) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- docs/usage/metrics.md | 5 +- .../unit/test_hisparse_connector.py | 44 ++++++++++++- .../kv_connector/unit/test_hisparse_stats.py | 31 +++++++++- tests/v1/worker/test_utils.py | 3 + .../kv_connector/v1/hisparse/connector.py | 2 + .../kv_connector/v1/hisparse/stats.py | 61 +++++++++++++++++++ vllm/v1/core/single_type_kv_cache_manager.py | 7 ++- 7 files changed, 147 insertions(+), 6 deletions(-) diff --git a/docs/usage/metrics.md b/docs/usage/metrics.md index b87ff1c28fee..ef75ce11c1fc 100644 --- a/docs/usage/metrics.md +++ b/docs/usage/metrics.md @@ -50,7 +50,10 @@ The following metrics are exposed: These metrics are exposed when the `HiSparseConnector` KV connector is configured. The host-pool gauges are reported by the scheduler once per engine step, for the host pool shared by all tensor-parallel ranks; evictable cached -blocks count as free. +blocks count as free. With `--kv-cache-metrics`, the +`vllm:hisparse_host_block_*_seconds` histograms sample host-pool block +residency, like the device `vllm:kv_block_*_seconds` histograms, which cover +only the GPU pool. --8<-- "gen:metrics-hisparse" diff --git a/tests/v1/kv_connector/unit/test_hisparse_connector.py b/tests/v1/kv_connector/unit/test_hisparse_connector.py index f08982a5e924..f9598a24f82e 100644 --- a/tests/v1/kv_connector/unit/test_hisparse_connector.py +++ b/tests/v1/kv_connector/unit/test_hisparse_connector.py @@ -3,7 +3,7 @@ import contextlib import ctypes from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import torch @@ -25,6 +25,7 @@ HiSparseConnectorWorker, ) from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector +from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.hisparse import runtime as runtime_module from vllm.v1.hisparse.coordinator import get_hisparse_coordinator from vllm.v1.hisparse.runtime import HiSparseCacheHandle @@ -33,6 +34,7 @@ SparseKVPageTransfer, SparseKVRowMirror, ) +from vllm.v1.metrics.stats import KVCacheEvictionEvent from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector @@ -100,6 +102,46 @@ def test_scheduler_stats_report_host_pool_usage(): assert stats.data["host_cache_usage_perc"] == [expected] +def test_hisparse_host_residency_reported_as_hisparse_metrics(): + """Host blocks sample into HiSparse's own residency metrics.""" + from tests.v1.core.test_prefix_caching import ( + make_hisparse_kv_cache_config, + make_kv_cache_manager, + ) + from tests.v1.core.utils import create_requests + + collector = KVCacheMetricsCollector(sample_rate=1.0) + manager = make_kv_cache_manager( + make_hisparse_kv_cache_config(2, 2), + max_model_len=128, + enable_caching=True, + hash_block_size=16, + metrics_collector=collector, + ) + coordinator = get_hisparse_coordinator(manager) + assert coordinator.host_manager is not None + scheduler = HiSparseConnectorScheduler(async_speculative=False) + scheduler.bind_coordinator(coordinator) + device = manager.block_pool + host = coordinator.host_manager.block_pool + request = create_requests(1, num_tokens=16)[0] + for pool, birth in ((device, 1), (host, 2)): + with patch("time.monotonic_ns", return_value=birth * 10**9): + block = pool.get_new_blocks(1)[0] + assert block.block_id == 1 + pool.cache_full_blocks(request, [block], 0, 1, 16, 0) + with patch("time.monotonic_ns", return_value=4_000_000_000): + host.free_blocks([host.blocks[1]]) + host.evict_blocks({1}) + with patch("time.monotonic_ns", return_value=6_000_000_000): + device.evict_blocks({1}) + + assert collector.drain_events() == [KVCacheEvictionEvent(5.0, 5.0, ())] + stats = scheduler.get_kv_connector_stats() + assert stats.data["host_block_lifetime_seconds"] == [2.0] + assert stats.data["host_block_idle_before_evict_seconds"] == [2.0] + + def test_no_forward_enqueues_deferred_hisparse_transfers(): """A zero-token step must still enqueue deferred post-forward transfers.""" connector = object.__new__(ActiveKVConnector) diff --git a/tests/v1/kv_connector/unit/test_hisparse_stats.py b/tests/v1/kv_connector/unit/test_hisparse_stats.py index 91d00b1e2260..f0b09ecf365b 100644 --- a/tests/v1/kv_connector/unit/test_hisparse_stats.py +++ b/tests/v1/kv_connector/unit/test_hisparse_stats.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest + from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import ( HiSparseConnector, ) @@ -82,7 +84,8 @@ def test_build_kv_connector_stats_round_trip(): assert rebuilt.reduce() == stats.reduce() -def test_prom_metrics_observe_host_usage_gauges(): +@pytest.mark.parametrize("kv_cache_metrics", [True, False]) +def test_prom_metrics_observe_host_metrics(kv_cache_metrics: bool): from types import SimpleNamespace from typing import Any @@ -91,12 +94,14 @@ def test_prom_metrics_observe_host_usage_gauges(): from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.stats import ( HiSparsePromMetrics, ) + from vllm.v1.metrics.stats import KVCacheEvictionEvent class _FakeMetric: def __init__(self, **kwargs: Any): self.kwargs = kwargs self.increments: list[int | float] = [] self.set_values: list[int | float] = [] + self.observed: list[float] = [] def labels(self, *labelvalues: object) -> "_FakeMetric": return self @@ -107,6 +112,9 @@ def inc(self, value: int | float) -> None: def set(self, value: int | float) -> None: self.set_values.append(value) + def observe(self, value: float) -> None: + self.observed.append(value) + created: dict[str, _FakeMetric] = {} class _NamedFake(_FakeMetric): @@ -115,7 +123,12 @@ def __init__(self, **kwargs: Any): created[kwargs["name"]] = self prom = HiSparsePromMetrics( - vllm_config=SimpleNamespace(kv_transfer_config=None), + vllm_config=SimpleNamespace( + kv_transfer_config=None, + observability_config=SimpleNamespace( + kv_cache_metrics=kv_cache_metrics, custom_histogram_buckets=None + ), + ), metric_types={Gauge: _NamedFake, Counter: _NamedFake, Histogram: _NamedFake}, labelnames=["model_name"], per_engine_labelvalues={0: ["model"]}, @@ -125,9 +138,21 @@ def __init__(self, **kwargs: Any): stats.record_snapshot(hits=3, misses=2, host_to_device_bytes=32) stats.record_host_usage(usage=0.5, pending_page_transfers=1) stats.record_host_usage(usage=0.75, pending_page_transfers=3) - prom.observe(stats.to_dict()) + stats.record_host_evictions([KVCacheEvictionEvent(5.0, 2.0, (1.0, 2.0))]) + later = HiSparseKVConnectorStats() + later.record_host_evictions([KVCacheEvictionEvent(3.0, 3.0, ())]) + prom.observe(stats.aggregate(later).to_dict()) assert created["vllm:hisparse_cache_hits"].increments == [3] assert created["vllm:hisparse_host_to_device_bytes"].increments == [32] assert created["vllm:hisparse_host_cache_usage_perc"].set_values == [0.75] assert created["vllm:hisparse_pending_page_transfers"].set_values == [3] + if not kv_cache_metrics: + assert not any("host_block" in name for name in created) + return + assert created["vllm:hisparse_host_block_lifetime_seconds"].observed == [5.0, 3.0] + assert created["vllm:hisparse_host_block_idle_before_evict_seconds"].observed == [ + 2.0, + 3.0, + ] + assert created["vllm:hisparse_host_block_reuse_gap_seconds"].observed == [1.0, 2.0] diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index e17f214871f3..98395e90c4dd 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -94,6 +94,9 @@ def test_hisparse_worker_get_kv_connector_stats_reads_completed_snapshot(monkeyp "host_to_device_bytes": [64], "host_cache_usage_perc": [], "pending_page_transfers": [], + "host_block_lifetime_seconds": [], + "host_block_idle_before_evict_seconds": [], + "host_block_reuse_gap_seconds": [], } diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py index 7e8960f644a9..863f52e02514 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py @@ -112,6 +112,8 @@ def get_kv_connector_stats(self) -> HiSparseKVConnectorStats | None: stats.record_host_usage( host_pool.get_usage(), len(self.coordinator.pending_spills) ) + if host_pool.metrics_collector is not None: + stats.record_host_evictions(host_pool.metrics_collector.drain_events()) return stats def build_connector_meta( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py index 6f6cfced9ede..9c7e8750976a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py @@ -12,6 +12,8 @@ PromMetric, PromMetricT, ) +from vllm.v1.metrics.buckets import histogram_buckets +from vllm.v1.metrics.stats import KVCacheEvictionEvent from vllm.v1.metrics.utils import create_metric_per_engine _HISPARSE_LEVEL_KEYS = frozenset({"host_cache_usage_perc", "pending_page_transfers"}) @@ -40,6 +42,9 @@ def reset(self): "host_to_device_bytes": [], "host_cache_usage_perc": [], "pending_page_transfers": [], + "host_block_lifetime_seconds": [], + "host_block_idle_before_evict_seconds": [], + "host_block_reuse_gap_seconds": [], } def record_snapshot(self, hits: int, misses: int, host_to_device_bytes: int): @@ -51,6 +56,12 @@ def record_host_usage(self, usage: float, pending_page_transfers: int) -> None: self.data["host_cache_usage_perc"].append(usage) self.data["pending_page_transfers"].append(pending_page_transfers) + def record_host_evictions(self, events: list[KVCacheEvictionEvent]) -> None: + for event in events: + self.data["host_block_lifetime_seconds"].append(event.lifetime_seconds) + self.data["host_block_idle_before_evict_seconds"].append(event.idle_seconds) + self.data["host_block_reuse_gap_seconds"].extend(event.reuse_gaps_seconds) + def aggregate(self, other: KVConnectorStats) -> KVConnectorStats: if not other.is_empty(): for k, v in other.data.items(): @@ -143,6 +154,53 @@ def __init__( ), } + self.hisparse_histograms: dict[str, Any] = {} + if vllm_config.observability_config.kv_cache_metrics: + residency_buckets = histogram_buckets( + "kv_cache_residency", + overrides=vllm_config.observability_config.custom_histogram_buckets, + ) + histogram_host_block_lifetime = self._histogram_cls( + name="vllm:hisparse_host_block_lifetime_seconds", + documentation=( + "Histogram of HiSparse host KV block lifetime from allocation " + "to eviction. Sampled metrics (controlled by " + "--kv-cache-metrics-sample)." + ), + buckets=residency_buckets, + labelnames=labelnames, + ) + histogram_host_block_idle_before_evict = self._histogram_cls( + name="vllm:hisparse_host_block_idle_before_evict_seconds", + documentation=( + "Histogram of HiSparse host KV block idle time before eviction. " + "Sampled metrics (controlled by --kv-cache-metrics-sample)." + ), + buckets=residency_buckets, + labelnames=labelnames, + ) + histogram_host_block_reuse_gap = self._histogram_cls( + name="vllm:hisparse_host_block_reuse_gap_seconds", + documentation=( + "Histogram of time gaps between consecutive HiSparse host KV " + "block accesses. Sampled metrics (controlled by " + "--kv-cache-metrics-sample)." + ), + buckets=residency_buckets, + labelnames=labelnames, + ) + self.hisparse_histograms = { + "host_block_lifetime_seconds": create_metric_per_engine( + histogram_host_block_lifetime, self.per_engine_labelvalues + ), + "host_block_idle_before_evict_seconds": create_metric_per_engine( + histogram_host_block_idle_before_evict, self.per_engine_labelvalues + ), + "host_block_reuse_gap_seconds": create_metric_per_engine( + histogram_host_block_reuse_gap, self.per_engine_labelvalues + ), + } + def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): for name, counter in self.hisparse_counters.items(): for value in transfer_stats_data.get(name, []): @@ -151,3 +209,6 @@ def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): values = transfer_stats_data.get(name, []) if values: gauge[engine_idx].set(values[-1]) + for name, histogram in self.hisparse_histograms.items(): + for value in transfer_stats_data.get(name, []): + histogram[engine_idx].observe(value) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index e9f5f2f3c231..eed15427b2dd 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -9,6 +9,7 @@ from vllm.distributed.kv_events import MEDIUM_CPU from vllm.utils.math_utils import cdiv from vllm.v1.core.block_pool import BlockPool +from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.core.kv_cache_utils import ( BlockHashList, BlockHashListWithBlockSize, @@ -2272,7 +2273,11 @@ def bind_host_pool(self, num_blocks: int) -> None: enable_caching=self.enable_caching and self.kv_cache_spec.prefix_cacheable, hash_block_size=device_pool.hash_block_size, enable_kv_cache_events=device_pool.enable_kv_cache_events, - metrics_collector=None, + metrics_collector=( + KVCacheMetricsCollector(device_pool.metrics_collector.sample_rate) + if device_pool.metrics_collector is not None + else None + ), medium=MEDIUM_CPU, event_owner=device_pool, ) From fa94ad41431d7942abcd1adfffd3262edaed975a Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Fri, 9 Oct 2026 17:54:06 -0500 Subject: [PATCH 210/624] [Test] Allow one bf16 ulp in the fused mHC post/pre residual check (#60724) Signed-off-by: Rohan138 Co-authored-by: Claude Opus 5.5 --- tests/kernels/mhc/test_mhc_kernels.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/kernels/mhc/test_mhc_kernels.py b/tests/kernels/mhc/test_mhc_kernels.py index 00e9834d6215..4afb45f503ef 100644 --- a/tests/kernels/mhc/test_mhc_kernels.py +++ b/tests/kernels/mhc/test_mhc_kernels.py @@ -439,13 +439,16 @@ def test_deepseek_v41_mhc_fused_post_pre_delayed(num_tokens, hidden_size, carrie return # The post mapping and the collapse are what the unfused kernels produce. + # The two kernels may round a near-midpoint residual differently (FMA + # contraction is up to the compiler), so allow one bf16 ulp there and + # check the collapse against the fused kernel's own residual. residual_ref = torch.ops.vllm.mhc_post_tilelang( x, residual, post_layer_mix, comb_res_mix ) layer_input_ref = mhc_pre_delayed_tilelang( - residual_ref, *mix_args, pre_mix=pre_mix, norm_weight=weight, norm_eps=1e-6 + residual_cur, *mix_args, pre_mix=pre_mix, norm_weight=weight, norm_eps=1e-6 )[2] - torch.testing.assert_close(residual_cur, residual_ref, atol=0, rtol=0) + torch.testing.assert_close(residual_cur, residual_ref, atol=0, rtol=2**-7) torch.testing.assert_close(layer_input, layer_input_ref, atol=0, rtol=0) post_ref, comb_ref, _, next_pre_ref = mhc_fused_post_pre_delayed_ref( From 04b02fc8a56fc0140d8559b3a07d84b2fdc46ae8 Mon Sep 17 00:00:00 2001 From: Bolin Sun Date: Fri, 9 Oct 2026 19:13:24 -0400 Subject: [PATCH 211/624] [Core] shm tensor arena for efficient cpu->gpu worker broadcast (#51207) Signed-off-by: Bolin Sun Signed-off-by: Signed-off-by: Nick Hill Co-authored-by: Claude Co-authored-by: Nick Hill --- docs/design/shm_tensor_arena.md | 302 ++++++++++ tests/distributed/test_shm_broadcast.py | 522 +++++++++++++++++- tools/pre_commit/check_forbidden_imports.py | 1 + vllm/config/parallel.py | 10 + .../device_communicators/shm_broadcast.py | 59 +- .../device_communicators/shm_tensor_arena.py | 397 +++++++++++++ vllm/engine/arg_utils.py | 6 + vllm/v1/executor/multiproc_executor.py | 1 + 8 files changed, 1290 insertions(+), 8 deletions(-) create mode 100644 docs/design/shm_tensor_arena.md create mode 100644 vllm/distributed/device_communicators/shm_tensor_arena.py diff --git a/docs/design/shm_tensor_arena.md b/docs/design/shm_tensor_arena.md new file mode 100644 index 000000000000..c37ac89e363a --- /dev/null +++ b/docs/design/shm_tensor_arena.md @@ -0,0 +1,302 @@ +# Zero-Copy Shared-Memory Tensor Arena for `MessageQueue` + +This document describes an optimization to vLLM's engine→worker IPC path for +multimodal serving under tensor parallelism, implemented in +`vllm/distributed/device_communicators/shm_tensor_arena.py` and wired into +`shm_broadcast.py`'s `MessageQueue`. It **builds on the existing out-of-band +tensor channel** (`_reduce_tensor` / `_rebuild_tensor`, still in +`shm_broadcast.py`), which already keeps CPU-tensor bytes out of the pickle +stream. + +> **Status: experimental, opt-in.** This lands as an increment over the +> existing out-of-band channel, not a replacement for it. vLLM also has an +> existing multimodal-tensor shm caching mechanism; reconciling the two is +> intentionally left as follow-up work rather than blocking this PR. + +## TL;DR + +`MessageQueue.enqueue` already avoids copying CPU tensor bytes into the pickle +stream: `_reduce_tensor` emits each tensor as a protocol-5 out-of-band +`PickleBuffer`, and large payloads are published to the local readers over a +ZMQ socket. That removed the dominant cost of the older *in-band* path — the +per-tensor `pickle` serialize (see [History](#history-the-in-band-pickle-era)). + +Two costs remain for a large multimodal `pixel_values` tensor on a TP=N worker: + +1. **N transport copies.** On the overflow (large-payload) path the out-of-band + buffer is published to each of the N node-local readers via `send_multipart`, + so the tensor bytes are moved once **per reader** by the transport — ~N×200 MB + of host data motion for a payload the model needs exactly once. +2. **A pageable H2D.** Each reader zero-copies a view over its received ZMQ + frame, but that frame is ordinary pageable memory, so the reader's subsequent + `cudaMemcpyAsync` to the GPU pays pageable-copy staging plus first-touch + faults. + +This change adds a slotted shared-memory **arena**: the writer does **one** +memcpy of the tensor into a free slot; every reader takes a **zero-copy view** +of the same slot (no per-reader transport copy); and the mapping is +`cudaHostRegister`-pinned so the H2D is a true DMA. Net for a 200 MB image on +TP=4: **1 copy + a pinned DMA** instead of **4 transport copies + a pageable +H2D**. Small tensors, non-contiguous tensors, and arena exhaustion all fall +through to the existing `_reduce_tensor` path unchanged. + +## Scope of the win (and what still needs measuring) + +The arena is an **incremental** improvement layered on the out-of-band channel, +not a from-scratch fix — the out-of-band channel already eliminated the single +largest cost (the serialize copy). The benefit that remains is: + +- **N→1 copies**, which grows with TP degree (no gain at TP=1), and +- the **pinned vs. pageable H2D**, which shows up as reader-side H2D latency. + +> **The benchmarks in [§4](#4-validation) were measured against the *in-band +> pickle* baseline that predates `_reduce_tensor`**, where the dominant cost was +> the serialize copy. They are retained as motivation and as an upper bound. The +> incremental improvement of the arena *over the current out-of-band path on +> `main`* is smaller and has **not yet been re-measured**; that A/B (arena vs. +> oob, same seeded workload) is the outstanding validation item before this is +> merge-ready. + +## 1. Where it bites + +Any multimodal model served with `--tensor-parallel-size > 1` through the +multiproc executor. The `rpc_broadcast_mq` `MessageQueue` +(`v1/executor/multiproc_executor.py`) carries scheduler output — including the +raw preprocessed `pixel_values` — from the EngineCore to every TP worker each +step. Images at native resolution reach tens to hundreds of MB (a ~4k×4k image +is ~200 MB as bf16 patches). The out-of-band channel keeps those bytes out of +the pickle stream, but on the overflow path they are still published **once per +node-local reader**, and each reader's H2D is from **pageable** memory. The +per-reader copy cost scales with TP degree; the pinning cost is a fixed per-image +tax on the reader H2D. Latency-sensitive (interactive) serving makes the residual +visible in TTFT and, via head-of-line blocking, in the end-to-end latency of +*other* requests. + +### History: the in-band-pickle era + +Before `_reduce_tensor`, `MessageQueue.enqueue` pickled the whole payload — +including tensor bytes — because `torch.Tensor.__reduce_ex__` does not emit +protocol-5 out-of-band buffers. A ~200 MB tensor was byte-copied into the pickle +stream on the engine (`THPStorage_writeFileRaw`, ~6.5 ms/MB → ~1.3 s), *inside* +the EngineCore step loop, so all GPUs on the worker went idle for **0.9–1.3 s per +large image** (the vision-encoder compute for the same image is 50–90 ms → +transport dominated compute ~15:1); each reader then paid a full `pickle.loads` +byte-copy. Per 200 MB image on TP=4 that was ~1.8 GB of host data motion. + +The `_reduce_tensor` out-of-band channel removed the serialize copy (and the +per-reader `pickle.loads` copy on the zero-copy ZMQ-frame path). Note that +merely raising `VLLM_MQ_MAX_CHUNK_BYTES_MB` (ring vs. socket path) did **not** +help in the in-band era: the serialize ran before the ring-vs-socket branch, so +it was common to both — a negative result that pinned the real cost on the +byte-copies, not the transport choice. + +## 2. The arena + +The arena itself (`ShmTensorArena`, `_ArenaPickler`, `_rebuild_arena_tensor`) +lives in `shm_tensor_arena.py`; `shm_broadcast.py`'s `MessageQueue` creates +and wires it in but otherwise treats it as an internal implementation detail. + +### 2.1 `ShmTensorArena` + +A second shared-memory region created alongside the existing `ShmRingBuffer`: + +- **N slots × slot_bytes** (default 8 × 256 MB, internal constants -- see + §3), plus per-slot metadata `[written_flag, reader0_done … readerN_done]`. +- Concurrency uses the **same lock-free single-writer/N-reader protocol as + `ShmRingBuffer`** (memory fences, per-reader done flags), so the model is one + the codebase already trusts. +- Created by the queue **writer** in `MessageQueue.__init__` when all readers are + node-local; readers attach via a new `tensor_arena_handle` field on the queue + `Handle`. Queues with remote readers get no arena and keep today's behavior. +- Pages are allocated lazily by the kernel on first write, so arenas on queues + that never carry big tensors (e.g. worker→engine response queues) cost + approximately nothing beyond the (virtual) reservation. + +### 2.2 Writer path — arena diversion, composing with `_reduce_tensor` + +`enqueue` builds one `dispatch_table` that routes CPU tensors through +`_reduce_tensor` (out-of-band `PickleBuffer`). When a node-local arena exists, +the pickler is an `_ArenaPickler` whose `reducer_override` *additionally* diverts +**large contiguous** CPU tensors into the arena: + +```python +class _ArenaPickler(pickle.Pickler): + def reducer_override(self, obj): + rebuild_fn = _ARENA_REBUILD_FNS.get(type(obj)) # {Tensor: ..., Parameter: ...} + if (rebuild_fn is not None and obj.device.type == "cpu" + and obj.layout is torch.strided and obj.is_contiguous() + and not obj.requires_grad + and obj.numel() * obj.element_size() >= MIN_BYTES): + idx = self.arena.write_tensor(obj) # ONE memcpy into a free slot + if idx is not None: + return (rebuild_fn, + (self.arena.shared_memory.name, idx, nbytes, dtype_str, shape)) + return NotImplemented # fall through to dispatch_table → _reduce_tensor +``` + +Note the exact-type dict lookup (not `isinstance`), matching how +`dispatch_table` itself dispatches: an unrecognized `Tensor` subclass must +decline here too, or it would come back as a plain `Tensor`, silently +losing its subclass identity. `torch.nn.Parameter` is the one subclass this +*does* divert, via its own rebuild function (§2.3). + +`reducer_override` is consulted before an object's normal reduction, and +returning `NotImplemented` falls through to the `dispatch_table` — so a diverted +tensor becomes a ~100-byte rebuild stub, and **everything the arena declines** +(too small, non-contiguous, or arena full) is handled by `_reduce_tensor` +exactly as on `main`. The single arena copy — +`torch.frombuffer(slot).copy_(t.view(torch.uint8))` — is a multithreaded memcpy +at memory bandwidth (~50–70 ms for 200 MB), and, unlike the oob path, the *same* +slot is then read zero-copy by all N readers rather than transported to each. + +**The writer never blocks.** If no slot is free or the tensor exceeds the slot +size, `write_tensor` returns `None`; the pickler falls through to `_reduce_tensor` +— the worst case is exactly the current out-of-band behavior, and deadlock is +structurally impossible. + +### 2.3 Reader path — zero copies + +The stub unpickles through a module-level rebuild function: + +```python +def _rebuild_arena_tensor(arena_name, slot_idx, nbytes, dtype_str, shape): + arena = _TENSOR_ARENAS[arena_name] # this process's registry of attached arenas + # get_tensor: a view over the mapped slot — zero bytes copied + return arena.get_tensor(slot_idx, nbytes, getattr(torch, dtype_str), shape) +``` + +`_rebuild_arena_parameter` wraps the same tensor in +`torch.nn.Parameter(t, requires_grad=False)`, which shares its storage. +`_ARENA_REBUILD_FNS = {torch.Tensor: _rebuild_arena_tensor, torch.nn. +Parameter: _rebuild_arena_parameter}` is the full set of types the arena +diverts; anything else (any other `Tensor` subclass) declines in +`reducer_override` and takes the out-of-band path. + +The registry is keyed by arena shm name (`_TENSOR_ARENAS: +weakref.WeakValueDictionary[str, ShmTensorArena]`) rather than a single +per-process slot, since nothing guarantees a process only ever attaches to +one arena-bearing queue. It holds only a *weak* reference — a strong one +would keep the arena (and its pinned mapping) alive for the process +lifetime — and needs no explicit removal on `MessageQueue.shutdown()` +(unsafe anyway, since shutdown can run on a different thread than an +in-flight `dequeue()`); an entry just disappears once nothing else +references that arena. + +The rebuilt tensor *is* the shared memory — no transport copy and no deserialize +on any rank. + +**Slot lifecycle.** The rebuilt tensor is the *source* of an async H2D while +the worker executes that step, so the reader must not release the slot at +unpickle time. Release is tied to **the lifetime of the slot's tensor +storage**, not a fixed "next dequeue" schedule — needed for callers like +`prompt_embeds` that retain the tensor across many `dequeue` calls during +chunked prefill, where a fixed schedule would let the writer reclaim (and +corrupt) the slot while the worker was still reading it. `get_tensor` builds +the tensor via `torch.from_numpy` over a numpy array wrapping the slot; the +tensor's storage holds the only reference to that array, so a +`weakref.finalize` on it fires only once every tensor, view/slice and wrapper +(e.g. `Parameter`) sharing the storage has been freed. Once queued, the +pinned fast path (§2.4) adds a second gate: the H2D is a true async DMA that +can outlive `execute_model`, so "the storage was freed" alone isn't +sufficient. +`flush_releases` records a CUDA event on the compute stream after each +step's H2D and only marks a slot done once that event completes +(non-blocking `event.query()`; not-yet-done just waits one more dequeue). +Unpinned, `cudaMemcpyAsync` from pageable memory already stages +synchronously, so the slot releases immediately. Either way, the writer +requires every reader's done flag before reusing a slot. + +> **Note.** Assumes the multimodal H2D is issued on the worker's +> current/default compute stream (true today: mm inputs are copied eagerly, outside the decode CUDA +> graph). If a future vLLM issues that copy on a dedicated side stream, the event +> must be recorded at the copy site rather than at `flush_releases`. + +### 2.4 Pinning — the zero-copy trap + +A zero-copy view over shared memory the reader process has never touched makes +the subsequent `cudaMemcpyAsync` pay **first-touch page faults on ~50k pages plus +pageable-copy staging** — the same tax the current oob path pays on its pageable +ZMQ frame. Fix: each reader lazily `cudaHostRegister`s the whole arena mapping +once (~1 s, first use), after which every H2D from the arena is a **pinned-memory +DMA** (~10 ms for 192 MB). Processes without a CUDA context skip registration +silently. (Because the arena mapping is stable and reused, this one-time pin +amortizes across all future images — an option the transient per-message ZMQ +frame does not have.) + +## 3. Configuration + +A single CLI flag toggles the arena; the slot count/size and divert threshold +are internal constants (no env vars). + +| Flag (`ParallelConfig` field) | Default | Meaning | +| --- | --- | --- | +| `--enable-shm-tensor-arena` / `--no-enable-shm-tensor-arena` (`enable_shm_tensor_arena`) | **off** | Opt in to route large CPU tensors through the arena (reserves slots in `/dev/shm`). Off = the out-of-band `_reduce_tensor` path only, identical to stock behavior. | + +Internal constants in `shm_tensor_arena.py`: **8 slots × 256 MB**; tensors larger +than a slot, or smaller than the **8 MB** divert threshold, take the out-of-band +`_reduce_tensor` path. + +## 4. Validation + +> **Baseline caveat.** The A/B numbers below were collected against the *in-band +> pickle* path that predates `_reduce_tensor`. They quantify the arena vs. that +> older baseline (an upper bound), **not** the arena vs. the current out-of-band +> path. The current-`main` oob baseline would land between the "stock" and +> "arena + pinning" rows; measuring exactly where is the open item. + +1. **Unit tests** (`tests/distributed/test_shm_broadcast.py`): byte-exact + zero-copy round-trips (incl. bf16/fp8), slot lifecycle (no reuse until every + reader releases; exhaustion / oversize / non-contiguous fall back to the + out-of-band path), event-gated release on the pinned path, `_ArenaPickler` ⊕ + `_reduce_tensor` composition, and an end-to-end `MessageQueue` broadcast + through forked processes. (Historical microbenchmark: enqueue of a 199 MB + payload took **66.7 ms** vs ~1275 ms for the old in-band serialize.) +2. **Same-seed A/B** (baseline = in-band pickle, pre-`_reduce_tensor`), interactive + multimodal workload, uncapped input images, TP=4 ×2 workers on one 8-GPU node, + low qps so individual images decompose cleanly (~1.2k aligned requests): + + | TTFT (ms) | p50 | p90 | p99 | max | >1 s | >1.5 s | + | --- | ---: | ---: | ---: | ---: | ---: | ---: | + | in-band pickle (old baseline) | 93 | 401 | 1321 | 2139 | 30 | 9 | + | + bigger MQ chunk (config lever) | 91 | 368 | 1287 | 2027 | 22 | 5 | + | **arena + pinning** | **89** | **241** | **862** | **1375** | **7** | **0** | + + p50 unchanged → `reducer_override` adds no measurable overhead when no large + tensor is present. +3. **Per-image stall** (CUDA-level all-GPU idle gap around the pixel H2D), vs the + in-band baseline: 192 MB image 1318 ms → 681 ms (arena only) → **351 ms class** + (arena + pinning); vision-encoder compute for the same image is ~80 ms. +4. **Outstanding — arena vs. out-of-band on `main`.** Re-run the same seeded A/B + with the baseline set to current `main` (i.e. `_reduce_tensor` on, arena off + via `--no-enable-shm-tensor-arena`) vs. arena on, to quantify the incremental N→1 + copy + pinning gain at TP=4 (and ideally TP=2/8 to show the copy-count scaling). + +## 5. Limitations and future work + +- **Value is TP- and size-dependent.** The arena helps only when a queue has + ≥2 node-local readers *and* carries tensors ≥ `MIN_MB`; at TP=1, or for small + payloads, it adds nothing and everything takes the out-of-band path. +- **Shared-memory reservation.** The arena reserves `slots × slot_bytes` of + `/dev/shm` (default 2 GB, lazily paged). Like `ShmRingBuffer`, creation is + guarded by a free-space check (`check_shm_free_space`), so an undersized + `/dev/shm` fails fast with a clear error instead of a `SIGBUS` the first + time a tensor is copied into pages beyond tmpfs capacity. +- **Relationship to `_reduce_tensor`.** This is strictly additive: the arena is + an opt-in fast path for the large-multimodal-tensor case; declining it (or + disabling via env) reverts to the merged out-of-band behavior. +- **Relationship to existing mm tensor shm caching.** vLLM already has a + separate shared-memory caching path for multimodal tensors; this arena is + not yet unified with it (see the experimental note in §1). Reconciling the + two is follow-up work. +- **Slot release granularity**: releases are tied to the slot's tensor + storage being freed and, on the pinned path, further gated on a per-slot + CUDA event recorded after the consuming H2D (§2.3), which closes both the + multi-step-retention and the async-DMA reuse windows. Bursts deeper than + the slot count safely fall back to the out-of-band path when the arena is + exhausted. +- **Fallback observability**: arena exhaustion (no free slot) is a + rate-limited log line today; a counter metric would be better. An + oversize tensor (bigger than a slot) falls back silently, with no log at + all -- worth adding if that path turns out to matter in practice. +- **Scope**: the arena activates only when every queue reader is node-local. + Remote readers (multi-node PP/TP) keep the existing socket path. diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py index d28732261a13..4984a2480a61 100644 --- a/tests/distributed/test_shm_broadcast.py +++ b/tests/distributed/test_shm_broadcast.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import dataclasses import gc import io import pickle @@ -17,7 +18,7 @@ import torch import torch.distributed as dist -from vllm.distributed.device_communicators import shm_broadcast +from vllm.distributed.device_communicators import shm_broadcast, shm_tensor_arena from vllm.distributed.device_communicators.shm_broadcast import ( MessageQueue, ShmRingBuffer, @@ -25,6 +26,10 @@ _reduce_tensor, check_shm_free_space, ) +from vllm.distributed.device_communicators.shm_tensor_arena import ( + ShmTensorArena, + _ArenaPickler, +) from vllm.distributed.utils import StatelessProcessGroup from vllm.utils.network_utils import get_open_port from vllm.utils.system_utils import update_environment_variables @@ -38,15 +43,18 @@ def get_arrays(n: int, seed: int = 0) -> list[np.ndarray]: return [np.random.randint(1, 100, i) for i in sizes] -def distributed_run(fn, world_size, timeout=60): +def distributed_run(fn, world_size, timeout=60, mp_context=None): """Run a function in multiple processes with proper error handling. Args: fn: Function to run in each process world_size: Number of processes to spawn timeout: Maximum time in seconds to wait for processes (default: 60) + mp_context: Optional multiprocess context (e.g. ``mp.get_context("spawn")``) + for workers that must not inherit the parent's thread/CUDA state. """ + ctx = mp_context or mp number_of_processes = world_size processes = [] for i in range(number_of_processes): @@ -57,7 +65,7 @@ def distributed_run(fn, world_size, timeout=60): env["LOCAL_WORLD_SIZE"] = str(number_of_processes) env["MASTER_ADDR"] = "localhost" env["MASTER_PORT"] = "12345" - p = mp.Process(target=fn, args=(env,)) + p = ctx.Process(target=fn, args=(env,)) processes.append(p) p.start() @@ -765,6 +773,20 @@ def test_shm_ring_buffer_creation_checks_free_space(): ShmRingBuffer(n_reader=1, max_chunk_bytes=24 * 1024 * 1024, max_chunks=10) +def test_shm_tensor_arena_creation_checks_free_space(): + """Like ShmRingBuffer, arena creation must fail fast with a clear error + on an undersized /dev/shm instead of discovering the shortfall via a + SIGBUS the first time a tensor is copied into pages past tmpfs capacity.""" + with ( + mock.patch.object( + shm_broadcast.shutil, "disk_usage", return_value=_fake_disk_usage(1 << 20) + ), + mock.patch.object(shm_broadcast.os.path, "isdir", return_value=True), + pytest.raises(RuntimeError, match="Insufficient space"), + ): + ShmTensorArena(n_reader=1, slot_bytes=256 << 20, n_slots=8) + + def test_remote_subscribe_addr_unique_concurrent_writers( monkeypatch: pytest.MonkeyPatch, ): @@ -828,3 +850,497 @@ def reader_main(): for q in queues: q.remote_socket.close(linger=0) + + +# --------------- ShmTensorArena (zero-copy tensor arena) tests --------------- + + +def _make_arena(n_reader: int = 2, slot_bytes: int = 4 << 20, n_slots: int = 3): + """Writer arena plus attached per-reader arenas (same process).""" + writer = ShmTensorArena(n_reader, slot_bytes, n_slots) + readers = [ShmTensorArena(*writer.handle(), reader_rank=i) for i in range(n_reader)] + return writer, readers + + +def _get_view(reader: ShmTensorArena, idx: int, ref: torch.Tensor) -> torch.Tensor: + return reader.get_tensor( + idx, ref.numel() * ref.element_size(), ref.dtype, tuple(ref.shape) + ) + + +def _drain(reader: ShmTensorArena) -> None: + """Flush a reader's releases to completion. On the pinned path the first + flush defers the release behind an H2D-completion event; the second flush + retires it once the event has fired.""" + reader.flush_releases() + if reader._deferred_releases: + torch.accelerator.synchronize() + reader.flush_releases() + assert not reader._pending_release + assert not reader._deferred_releases + + +def test_arena_zero_copy_roundtrip(): + writer, (r0, r1) = _make_arena(n_reader=2) + src = torch.randn(400, 512) + idx = writer.write_tensor(src) + assert idx is not None + a = _get_view(r0, idx, src) + b = _get_view(r1, idx, src) + assert torch.equal(a, src) + assert torch.equal(b, src) + assert a.dtype == src.dtype and a.shape == src.shape + # Readers alias one shared mapping: a write through one reader's view is + # visible through the other's (this is what "zero-copy" means here). + a[0, 0] = 12345.0 + assert b[0, 0].item() == 12345.0 + del a, b + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +def test_arena_dtype_roundtrip(dtype): + """Dtypes numpy doesn't recognize traverse the arena's uint8 view path.""" + writer, (reader,) = _make_arena(n_reader=1) + src = torch.randn(256, 256).to(dtype) + idx = writer.write_tensor(src) + assert idx is not None + got = _get_view(reader, idx, src) + assert got.dtype == dtype + assert torch.equal(got.view(torch.uint8), src.view(torch.uint8)) + del got + + +def test_arena_slot_lifecycle(): + """A slot must not be reusable until EVERY reader has released it — + the writer overwriting a slot a reader still consumes would corrupt data. + Release is tied to the returned tensor's storage lifetime (`get_tensor`): + a reader that still holds its view keeps the slot reserved even across + repeated `flush_releases()` calls, not just until the "next" one.""" + writer, readers = _make_arena(n_reader=2, n_slots=3) + t = torch.ones(1000) + first = writer.write_tensor(t) + views = [_get_view(r, first, t) for r in readers] + # Fill the remaining slots; the arena is now exhausted. + assert all(writer.write_tensor(t) is not None for _ in range(2)) + assert writer.write_tensor(t) is None # exhausted -> caller falls back + + # Flushing while both readers still hold their view releases nothing: + # the slot is only queued once the returned tensor's storage is freed. + readers[0].flush_releases() + readers[1].flush_releases() + assert writer.write_tensor(t) is None + + # Dropping one reader's view queues its release; still not enough. + views[0] = None + _drain(readers[0]) + assert writer.write_tensor(t) is None + # Once every reader has dropped (and released) its view, it's reused. + views[1] = None + _drain(readers[1]) + assert writer.write_tensor(t) == first + + +def test_arena_slot_not_released_while_tensor_retained_across_multiple_flushes(): + """A caller retaining the tensor across many `flush_releases()` calls + (e.g. `prompt_embeds` across chunked-prefill steps) must keep its slot + reserved the whole time, not just until the reader's next flush.""" + writer, (reader,) = _make_arena(n_reader=1, n_slots=2) + t = torch.ones(1000) + idx = writer.write_tensor(t) + retained = _get_view(reader, idx, t) + # Several flush cycles pass while the caller still holds `retained`. + for _ in range(5): + reader.flush_releases() + assert reader._pending_release == [] + # The slot is still reserved: only the arena's other slot is available... + assert writer.write_tensor(t) is not None + assert writer.write_tensor(t) is None # ...and now the arena is exhausted. + # ...and the bytes the caller is reading are unchanged throughout. + assert torch.equal(retained, t) + # Only once the caller drops its reference is the slot released. + del retained + _drain(reader) + assert writer.write_tensor(t) == idx + + +def test_arena_slot_not_released_while_view_retained(): + """A view/slice (or wrapper) sharing the returned tensor's storage keeps + the slot reserved even after the original tensor object is dropped.""" + writer, (reader,) = _make_arena(n_reader=1, n_slots=1) + t = torch.arange(1000, dtype=torch.float32) + idx = writer.write_tensor(t) + view = _get_view(reader, idx, t)[10:20].unsqueeze(0) + param = torch.nn.Parameter(view, requires_grad=False) + del view + reader.flush_releases() + assert reader._pending_release == [] + assert writer.write_tensor(t) is None + assert torch.equal(param[0], t[10:20]) + del param + assert reader._pending_release == [idx] + _drain(reader) + assert writer.write_tensor(t) == idx + + +def test_arena_release_survives_intervening_flush_from_another_tensor(): + """`flush_releases()` must clear `_pending_release` in place, never + rebind it: a finalizer registered before a rebind would append into the + now-orphaned old list and permanently leak its slot. Exercises the + interleaving a long-lived `prompt_embeds` tensor sees once a + shorter-lived tensor through the same reader triggers a flush first.""" + writer, (reader,) = _make_arena(n_reader=1, n_slots=2) + t = torch.ones(1000) + # `a` is long-lived: its finalizer is registered first, while + # _pending_release is still the arena's original (empty) list object. + idx_a = writer.write_tensor(t) + a = _get_view(reader, idx_a, t) + # `b` is short-lived and its finalizer is registered against that SAME + # list object (nothing has flushed yet). + idx_b = writer.write_tensor(t) + b = _get_view(reader, idx_b, t) + del b + assert reader._pending_release == [idx_b] + # This flush must not orphan `a`'s still-pending finalizer. `_drain`, + # not a bare `flush_releases()` call: on the pinned fast path (real + # CUDA present) the release is gated behind an event and may need a + # second flush once it fires -- see `flush_releases`. + _drain(reader) + assert writer.write_tensor(t) == idx_b # b's slot correctly reused. + # `a` is dropped well after that intervening flush. + del a + assert reader._pending_release == [idx_a] + _drain(reader) + assert writer.write_tensor(t) == idx_a # a's slot must NOT be leaked. + + +def test_arena_oversize_falls_back(): + writer, _ = _make_arena(n_reader=1, slot_bytes=1 << 20, n_slots=2) + big = torch.empty((1 << 20) + 4096, dtype=torch.uint8) + assert writer.write_tensor(big) is None + + +def _dumps_arena(obj, arena: ShmTensorArena) -> tuple[bytes, list, list[int]]: + """Pickle `obj` the same way `MessageQueue.enqueue` does when an arena is + attached: arena diversion first (reducer_override), then the tensor + dispatch table, with out-of-band buffers >= 1MiB. Also returns the slot + indices the arena actually wrote to, in write order.""" + buffers = [] + + def callback(buf: pickle.PickleBuffer) -> bool: + raw = buf.raw() + if raw.nbytes < 1024 * 1024: + return True + buffers.append(raw) + return False + + written: list[int] = [] + orig_write_tensor = arena.write_tensor + + def _spy_write_tensor(t): + idx = orig_write_tensor(t) + if idx is not None: + written.append(idx) + return idx + + bio = io.BytesIO() + pickler = _ArenaPickler(bio, arena, buffer_callback=callback) + pickler.dispatch_table = {torch.Tensor: _reduce_tensor} + with mock.patch.object(arena, "write_tensor", side_effect=_spy_write_tensor): + pickler.dump(obj) + return bio.getvalue(), buffers, written + + +def test_arena_pickler_composes(monkeypatch): + """Large contiguous tensors are diverted into the arena; everything the + arena declines falls through to `_reduce_tensor` unchanged.""" + writer, (reader,) = _make_arena(n_reader=1) + monkeypatch.setattr(shm_tensor_arena, "_ARENA_MIN_BYTES", 1 << 20) + monkeypatch.setitem( + shm_tensor_arena._TENSOR_ARENAS, reader.shared_memory.name, reader + ) + big = torch.randn(1024, 1024) # 4MiB -> diverted into the arena + small = torch.randn(16, 16) # 1KiB -> falls through to _reduce_tensor + data, buffers, written = _dumps_arena({"big": big, "small": small}, writer) + # The diverted tensor's bytes are in the arena, not the pickle stream. + assert len(data) + sum(b.nbytes for b in buffers) < big.numel() * 4 + out = pickle.loads(data, buffers=buffers) + assert torch.equal(out["big"], big) + assert torch.equal(out["small"], small) + # "big" is a zero-copy view of the reader's slot; "small" is not. + (idx,) = written + nbytes = big.numel() * big.element_size() + slot_ptr = torch.frombuffer(reader._slot(idx, nbytes), dtype=torch.uint8).data_ptr() + assert out["big"].data_ptr() == slot_ptr + assert out["small"].data_ptr() != slot_ptr + # Not yet queued for release: `out["big"]` is still alive (see get_tensor). + assert reader._pending_release == [] + del out + assert reader._pending_release == [idx] + _drain(reader) + + +def test_arena_pickler_noncontig_falls_through(monkeypatch): + writer, (reader,) = _make_arena(n_reader=1) + monkeypatch.setattr(shm_tensor_arena, "_ARENA_MIN_BYTES", 1 << 20) + monkeypatch.setitem( + shm_tensor_arena._TENSOR_ARENAS, reader.shared_memory.name, reader + ) + nc = torch.randn(2048, 1024)[:, ::2] # non-contiguous, above threshold + assert not nc.is_contiguous() + data, buffers, written = _dumps_arena(nc, writer) + out = pickle.loads(data, buffers=buffers) + assert torch.equal(out, nc) + # The arena never saw it: no slot consumed on the reader. + assert written == [] + assert reader._pending_release == [] + del out + + +@pytest.mark.parametrize("case", ["requires_grad", "conj"]) +def test_arena_pickler_excludes_requires_grad_and_conj(monkeypatch, case: str): + """Tensors `_reduce_tensor` excludes (requires_grad, or an exotic + conjugate-bit view that can't alias as uint8) must fall through the + arena the same way, even when every other criterion qualifies.""" + writer, (reader,) = _make_arena(n_reader=1) + monkeypatch.setattr(shm_tensor_arena, "_ARENA_MIN_BYTES", 1 << 20) + monkeypatch.setitem( + shm_tensor_arena._TENSOR_ARENAS, reader.shared_memory.name, reader + ) + if case == "requires_grad": + tensor = torch.randn(1024, 1024, requires_grad=True) # 4MiB, above threshold + else: + tensor = torch.randn(600, 600, dtype=torch.complex64).conj() # ~2.9MiB + assert tensor.is_contiguous() + assert tensor.numel() * tensor.element_size() >= shm_tensor_arena._ARENA_MIN_BYTES + + data, buffers, written = _dumps_arena(tensor, writer) + # Declined by the arena -- no slot consumed -- and handled by + # _reduce_tensor's identical exclusion instead. + assert written == [] + out = pickle.loads(data, buffers=buffers) + assert torch.equal(out, tensor) + assert out.requires_grad == tensor.requires_grad + + +def test_arena_pickler_excludes_tensor_subclasses(monkeypatch): + """`torch.nn.Parameter` is diverted via its own `_rebuild_arena_parameter` + hook, which wraps an intermediate tensor it keeps no Python reference to. + The slot must stay reserved while the Parameter is alive -- the + `_pending_release` assertions below check this even though + `type(out) is torch.nn.Parameter` alone would not.""" + writer, (reader,) = _make_arena(n_reader=1) + monkeypatch.setattr(shm_tensor_arena, "_ARENA_MIN_BYTES", 1 << 20) + monkeypatch.setitem( + shm_tensor_arena._TENSOR_ARENAS, reader.shared_memory.name, reader + ) + # requires_grad=False to isolate the subclass check from the (already + # separately tested) requires_grad exclusion. + param = torch.nn.Parameter(torch.randn(1024, 1024), requires_grad=False) + assert type(param) is not torch.Tensor + assert param.is_contiguous() + assert param.numel() * param.element_size() >= shm_tensor_arena._ARENA_MIN_BYTES + + data, buffers, written = _dumps_arena(param, writer) + # The Parameter itself is diverted directly (_ARENA_REBUILD_FNS routes + # it to _rebuild_arena_parameter) -- Parameter.__reduce_ex__ never runs, + # so there's no separate "inner" tensor diversion to speak of. + assert written != [] + out = pickle.loads(data, buffers=buffers) + # The outer Parameter's type identity survives. + assert type(out) is torch.nn.Parameter + assert torch.equal(out, param) + assert out.requires_grad == param.requires_grad + # Not yet queued for release: `out` still shares the slot's storage. + assert reader._pending_release == [] + del out + assert reader._pending_release == written + _drain(reader) + + +def test_arena_unregisters_pinned_memory_on_del(monkeypatch): + """A reader that cudaHostRegister-ed its mapping must cudaHostUnregister + the same pointer before the mapping is closed; an unpinned arena must not + unregister memory it never registered.""" + writer, (reader,) = _make_arena(n_reader=1) + cudart = mock.MagicMock() + cudart.cudaHostRegister.return_value = 0 + cudart.cudaHostUnregister.return_value = 0 + monkeypatch.setattr( + shm_tensor_arena, + "current_platform", + SimpleNamespace(is_cuda_alike=lambda: True, cudart=lambda: cudart), + ) + reader._ensure_pinned() + assert reader._pinned + ptr = cudart.cudaHostRegister.call_args.args[0] + del reader + cudart.cudaHostUnregister.assert_called_once_with(ptr) + del writer # never pinned -> no unregister + cudart.cudaHostUnregister.assert_called_once() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_arena_event_gated_release(): + """On the pinned path a slot release is gated on an H2D-completion CUDA + event: the writer must not see the reader's done flag until the async DMA + sourced from the slot has been retired. Release is first queued once the + caller drops the CPU-side view -- the realistic pattern: the view is only + needed as the H2D source, while the GPU tensor is what's retained for the + rest of the step.""" + writer, (reader,) = _make_arena(n_reader=1) + reader._ensure_pinned() + if not reader._pinned: + pytest.skip("cudaHostRegister unavailable in this environment") + src = torch.randn(512, 1024) + idx = writer.write_tensor(src) + view = _get_view(reader, idx, src) + dev = view.to("cuda", non_blocking=True) + # Not yet queued: this test still holds `view`. + assert reader._pending_release == [] + del view + assert reader._pending_release == [idx] + reader.flush_releases() + # Deferred behind the event, not applied eagerly. + assert len(reader._deferred_releases) == 1 + with reader._meta(idx) as meta: + assert meta[1] == 0 + torch.accelerator.synchronize() + reader.flush_releases() + assert reader._deferred_releases == [] + with reader._meta(idx) as meta: + assert meta[1] == 1 + assert torch.equal(dev.cpu(), src) + del dev + + +def test_tensor_arena_registry_is_weak_and_survives_shutdown(): + """`_TENSOR_ARENAS` must not keep an arena alive past its owning + MessageQueue, but also must not be cleared by `shutdown()` -- which can + run on a different thread than an in-flight `dequeue()` still + reconstructing a tensor from that arena. A weak-value registry gets + both: the entry survives `shutdown()` and vanishes only once + unreferenced.""" + writer_mq = MessageQueue(n_reader=1, n_local_reader=1, enable_shm_tensor_arena=True) + handle = writer_mq.export_handle() + reader_mq = MessageQueue.create_from_handle(handle, rank=0) + + name = reader_mq.tensor_arena.shared_memory.name + assert shm_tensor_arena._TENSOR_ARENAS[name] is reader_mq.tensor_arena + + # shutdown() must NOT remove the entry: that would race with a + # concurrent dequeue() still reconstructing a tensor from this arena. + reader_mq.shutdown() + assert shm_tensor_arena._TENSOR_ARENAS[name] is reader_mq.tensor_arena + + # Once nothing else references the arena, the registry entry vanishes + # on its own. + reader_mq.tensor_arena = None + assert name not in shm_tensor_arena._TENSOR_ARENAS + + writer_mq.shutdown() + + +def test_create_from_handle_tolerates_arena_attach_failure(): + """If a reader's `ShmTensorArena.__init__` can't attach to the writer's + segment (`FileNotFoundError`, e.g. deserialized on a different node), + `create_from_handle` must not crash -- same deferred-failure convention + as the sibling `ShmRingBuffer` path.""" + writer_mq = MessageQueue(n_reader=1, n_local_reader=1, enable_shm_tensor_arena=True) + handle = writer_mq.export_handle() + n_reader, slot_bytes, n_slots, _real_name = handle.tensor_arena_handle + bogus_handle = dataclasses.replace( + handle, + tensor_arena_handle=( + n_reader, + slot_bytes, + n_slots, + "vllm-arena-does-not-exist", + ), + ) + + reader_mq = MessageQueue.create_from_handle(bogus_handle, rank=0) + + # Attach failed silently: the arena object exists but never got its + # shared_memory attribute, and was never registered. + assert reader_mq.tensor_arena is not None + assert not hasattr(reader_mq.tensor_arena, "shared_memory") + assert "vllm-arena-does-not-exist" not in shm_tensor_arena._TENSOR_ARENAS + + reader_mq.shutdown() + writer_mq.shutdown() + + +@worker_fn_wrapper +def worker_fn_arena_broadcast(): + rank = dist.get_rank() + writer_rank = 0 + if rank == writer_rank: + message_queue = MessageQueue( + 1, + 1, + local_reader_ranks=[1], + max_chunk_bytes=8 * 1024 * 1024, + enable_shm_tensor_arena=True, + ) + handles = [message_queue.export_handle()] + dist.broadcast_object_list(handles, src=writer_rank) + else: + handles = [None] + dist.broadcast_object_list(handles, src=writer_rank) + message_queue = MessageQueue.create_from_handle(handles[0], rank) + message_queue.wait_until_ready() + + torch.manual_seed(42) + payload = { + # 16MiB: above the arena divert threshold (8MiB) -> arena slot. + "huge": torch.randn(2048, 2048), + # 2MiB: declined by the arena -> out-of-band _reduce_tensor path. + "mid": torch.randn(1024, 512), + } + + if rank == writer_rank: + assert message_queue.tensor_arena is not None + with mock.patch( + "vllm.distributed.device_communicators.shm_broadcast._reduce_tensor", + wraps=_reduce_tensor, + ) as wrapped_reduce: + message_queue.enqueue(payload) + # "huge" was diverted into the arena before the dispatch table was + # consulted; only "mid" went through _reduce_tensor. + assert wrapped_reduce.call_count == 1 + else: + received = message_queue.dequeue(timeout=30) + assert torch.equal(received["huge"], payload["huge"]) + assert torch.equal(received["mid"], payload["mid"]) + # The huge tensor is a zero-copy view into the arena's shared memory + # (not a copy received over the transport). Its slot is not yet + # queued for release while `received` still holds the reference. + arena = message_queue.tensor_arena + assert arena is not None + assert arena._pending_release == [] + huge_ptr = received["huge"].data_ptr() + arena_base = torch.frombuffer( + arena.shared_memory.buf, dtype=torch.uint8 + ).data_ptr() + assert arena_base <= huge_ptr < arena_base + arena.total_bytes + # Dropping the last reference to the slot's storage queues it for + # release (see get_tensor). + del received + assert len(arena._pending_release) == 1 + + dist.barrier() + print(f"arena broadcast passed the test! Rank {rank}") + + +def test_arena_broadcast(): + # Spawn (not fork): by the time this test runs, the pytest process may be + # multi-threaded / CUDA-initialized (earlier GPU tests), and a forked + # child running the arena's slot memcpy can deadlock on inherited lock + # state. Spawned workers start clean; the larger timeout absorbs their + # interpreter + import startup. + distributed_run( + worker_fn_arena_broadcast, + 2, + timeout=180, + mp_context=mp.get_context("spawn"), + ) diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index 531f6fb8a2b4..d80a8e1fd26d 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -56,6 +56,7 @@ class ForbiddenImport: "vllm/distributed/device_communicators/all_reduce_utils.py", "vllm/distributed/device_communicators/shm_broadcast.py", "vllm/distributed/device_communicators/shm_object_storage.py", + "vllm/distributed/device_communicators/shm_tensor_arena.py", "vllm/distributed/weight_transfer/ipc_engine.py", "vllm/distributed/weight_transfer/clients.py", "vllm/model_executor/model_loader/weight_cache/protocol.py", diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 37affbaf9ebf..6960661e0c7e 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -235,6 +235,15 @@ class ParallelConfig: disable_custom_all_reduce: bool = False """Disable the custom all-reduce kernel and fall back to NCCL.""" + enable_shm_tensor_arena: bool = False + """Route large CPU tensors (e.g. multimodal ``pixel_values``) in the + engine→worker broadcast through a zero-copy shared-memory arena instead of + sending a copy to every reader. **Experimental**: may be reconciled with + vLLM's existing multimodal shm tensor caching in a future release. Opt-in; + most beneficial for multimodal serving with large images and tensor + parallelism. Active only when all queue readers are node-local; reserves + slots in ``/dev/shm``.""" + enable_elastic_ep: bool = False """Enable elastic expert parallelism with stateless NCCL groups for DP/EP.""" elastic_ep_max_dp_size: int = Field(default=None, ge=1) # type: ignore[assignment] @@ -901,6 +910,7 @@ def compute_hash(self): "nnodes", "max_parallel_loading_workers", "disable_custom_all_reduce", + "enable_shm_tensor_arena", "ray_workers_use_nsight", "ray_runtime_env", "placement_group", diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index da2288f77b54..5656e8aaa153 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -61,6 +61,10 @@ if TYPE_CHECKING: from _typeshed import SizedBuffer + from vllm.distributed.device_communicators.shm_tensor_arena import ( + ShmTensorArena, + ) + VLLM_RINGBUFFER_WARNING_INTERVAL = envs.VLLM_RINGBUFFER_WARNING_INTERVAL # Cap on how long an idle reader parks before re-reading the authoritative SHM # written-flag. Bounds lost-notify recovery latency to ~5s while the periodic @@ -471,6 +475,7 @@ class Handle: local_notify_addr: str | None = None remote_subscribe_addr: str | None = None remote_addr_ipv6: bool = False + tensor_arena_handle: tuple[int, int, int, str] | None = None def _close_zmq(context: zmq.Context, sockets: list[zmq.Socket]) -> None: @@ -490,6 +495,7 @@ def __init__( max_chunk_bytes: int = 1024 * 1024 * 24, max_chunks: int = 10, connect_ip: str | None = None, + enable_shm_tensor_arena: bool = False, ): if local_reader_ranks is None: local_reader_ranks = list(range(n_local_reader)) @@ -501,12 +507,22 @@ def __init__( self.shutting_down = False context = Context() + self.tensor_arena: ShmTensorArena | None = None if n_local_reader > 0: # for local readers, we will: # 1. create a shared memory ring buffer to communicate small data # 2. create a publish-subscribe socket to communicate large data self.buffer = ShmRingBuffer(n_local_reader, max_chunk_bytes, max_chunks) + # Zero-copy arena for large CPU tensors (see ShmTensorArena). + # Local readers only: remote readers receive the pickled bytes + # over a socket and cannot map the arena, so the substitution + # would break them. + if enable_shm_tensor_arena and n_remote_reader == 0: + from vllm.distributed.device_communicators import shm_tensor_arena + + self.tensor_arena = shm_tensor_arena.ShmTensorArena(n_local_reader) + # XPUB is very similar to PUB, # except that it can receive subscription messages # to confirm the number of subscribers @@ -567,6 +583,9 @@ def __init__( local_notify_addr=local_notify_addr, remote_subscribe_addr=remote_subscribe_addr, remote_addr_ipv6=remote_addr_ipv6, + tensor_arena_handle=( + self.tensor_arena.handle() if self.tensor_arena is not None else None + ), ) self._close_zmq_on_collect(context) @@ -599,6 +618,7 @@ def create_from_handle(handle: Handle, rank) -> "MessageQueue": context = Context() + self.tensor_arena = None if rank in handle.local_reader_ranks: assert handle.buffer_handle is not None self.buffer = ShmRingBuffer(*handle.buffer_handle) @@ -607,6 +627,23 @@ def create_from_handle(handle: Handle, rank) -> "MessageQueue": self._is_local_reader = True self._is_remote_reader = False + arena_handle = handle.tensor_arena_handle + if arena_handle is not None: + from vllm.distributed.device_communicators import shm_tensor_arena + + self.tensor_arena = shm_tensor_arena.ShmTensorArena( + *arena_handle, reader_rank=self.local_reader_rank + ) + # `shared_memory` is unset (not just None) if attaching to + # the writer's segment raised FileNotFoundError -- see + # ShmTensorArena.__init__'s reader branch. Registering a + # broken arena is a no-op deferred failure (same convention + # ShmRingBuffer uses above), not a crash here. + if hasattr(self.tensor_arena, "shared_memory"): + shm_tensor_arena._TENSOR_ARENAS[ + self.tensor_arena.shared_memory.name + ] = self.tensor_arena + self.local_socket = context.socket(SUB) self.local_socket.setsockopt_string(SUBSCRIBE, "") socket_addr = handle.local_subscribe_addr @@ -880,12 +917,21 @@ def oob_callback(buf: PickleBuffer) -> bool: # dispatch table would otherwise shadow them. dispatch_table = dict(copyreg.dispatch_table) dispatch_table[torch.Tensor] = _reduce_tensor + arena = self.tensor_arena with io.BytesIO() as bio: - pickler = pickle.Pickler( - bio, - protocol=pickle.HIGHEST_PROTOCOL, - buffer_callback=oob_callback, - ) + pickler: pickle.Pickler + if arena is not None: + from vllm.distributed.device_communicators.shm_tensor_arena import ( + _ArenaPickler, + ) + + pickler = _ArenaPickler(bio, arena, buffer_callback=oob_callback) + else: + pickler = pickle.Pickler( + bio, + protocol=pickle.HIGHEST_PROTOCOL, + buffer_callback=oob_callback, + ) pickler.dispatch_table = dispatch_table pickler.dump(obj) all_buffers[0] = bio.getvalue() @@ -922,6 +968,9 @@ def dequeue( ): """Read from message queue with optional timeout (in seconds)""" if self._is_local_reader: + # See ShmTensorArena.flush_releases for what this retires and why. + if self.tensor_arena is not None: + self.tensor_arena.flush_releases() with self.acquire_read(timeout, indefinite) as buf: overflow = buf[0] == 1 if not overflow: diff --git a/vllm/distributed/device_communicators/shm_tensor_arena.py b/vllm/distributed/device_communicators/shm_tensor_arena.py new file mode 100644 index 000000000000..257cce06b2b4 --- /dev/null +++ b/vllm/distributed/device_communicators/shm_tensor_arena.py @@ -0,0 +1,397 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Zero-copy shared-memory arena for large CPU tensors (multimodal fast path). + +Problem: `MessageQueue.enqueue` (see `shm_broadcast.py`) pickles the whole +payload, and torch's default tensor pickling copies bytes into the pickle +stream on the writer and back out on every local reader. For a large +`pixel_values` tensor (up to hundreds of MB) this blocks the EngineCore step +loop for ~1s with all GPUs idle. + +Fast path: `_ArenaPickler.reducer_override` intercepts qualifying CPU +tensors during pickling, memcpys them ONCE into a free arena slot, and +pickles a tiny `(arena_name, slot, nbytes, dtype, shape)` stub. Readers +rebuild a zero-copy view of the mapped slot (`torch.frombuffer`). + +Slot lifecycle: single writer, n_reader readers, per-slot metadata +[written_flag, reader0_done, ..., readerN_done] (same protocol as +ShmRingBuffer). A slot is not released when rebuilt: the rebuilt tensor is +the SOURCE of an async H2D copy, and callers like chunked-prefill +`prompt_embeds` retain it across many steps. Release is queued only once the +tensor's storage is freed, i.e. once no tensor or view of the slot remains +(see `get_tensor`), and is further gated on a CUDA event on the pinned fast +path, since the H2D can outlive the step that issued it (see +`flush_releases`). + +The writer NEVER blocks: no free slot (or an oversized tensor) falls back to +the default pickle path, so worst case matches today's behavior. + +Physical memory: slots are paged in lazily, so an arena on a queue that +never carries big tensors costs ~0. Creation is guarded by +`check_shm_free_space` (like `ShmRingBuffer`) so an undersized `/dev/shm` +fails fast instead of a `SIGBUS` mid-copy. +""" + +import pickle +import weakref +from contextlib import suppress +from multiprocessing import shared_memory +from unittest.mock import patch + +import numpy as np +import torch + +from vllm.distributed.device_communicators.shm_broadcast import ( + check_shm_free_space, + memory_fence, +) +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + +# Arena sizing — internal constants (kept off the env/CLI surface). +_ARENA_SLOTS = 8 +_ARENA_SLOT_BYTES = 256 * 1024 * 1024 +# Minimum contiguous CPU-tensor size to divert into the arena; smaller tensors +# take the out-of-band pickle-buffer path. +_ARENA_MIN_BYTES = 8 * 1024 * 1024 + +# Reader-side registry: arena shm name -> attached ShmTensorArena, keyed by +# name rather than a single slot since a process isn't guaranteed to attach +# to only one. Populated by `MessageQueue.create_from_handle`; consumed by +# the `_rebuild_arena_*` unpickle hooks. WEAK-value: a strong reference here +# would keep the arena (and its pinned mapping) alive for the process +# lifetime, and would need `MessageQueue.shutdown` to remove it explicitly -- +# unsafe since shutdown() can run on a different thread than an in-flight +# `dequeue`. Instead an entry just disappears once nothing else references +# that arena. +_TENSOR_ARENAS: weakref.WeakValueDictionary[str, "ShmTensorArena"] = ( + weakref.WeakValueDictionary() +) + + +class ShmTensorArena: + """Slotted shared-memory arena: one writer, n_reader zero-copy readers. + + Memory layout: [slot0 | slot1 | ... | slotN-1 | meta0 | meta1 | ... ] + where each meta is (1 + n_reader) bytes: [written_flag, reader_done...]. + Slot states mirror ShmRingBuffer: + written=0 -> free (never written / being written) + written=1, some reader_done=0 -> in use, cannot reuse + written=1, all reader_done=1 -> consumed, can reuse + """ + + def __init__( + self, + n_reader: int, + slot_bytes: int = _ARENA_SLOT_BYTES, + n_slots: int = _ARENA_SLOTS, + name: str | None = None, + reader_rank: int = -1, + ): + self.n_reader = n_reader + self.slot_bytes = slot_bytes + self.n_slots = n_slots + self.reader_rank = reader_rank # -1 for the writer + self.metadata_size = 1 + n_reader + self.metadata_offset = slot_bytes * n_slots + self.total_bytes = (slot_bytes + self.metadata_size) * n_slots + self._next_slot = 0 + self._fallbacks = 0 + self._pin_attempted = False + self._pinned = False + self._pinned_ptr = 0 + # slots queued for release once the storage of the tensor `get_tensor` + # handed out is freed; flushed at the next dequeue. + self._pending_release: list[int] = [] + # slots whose release is gated on an H2D-completion CUDA event: each + # entry is (event, [slot_idx, ...]). Only populated on the pinned fast + # path, where the source-side DMA is asynchronous (see flush_releases). + self._deferred_releases: list[tuple[torch.cuda.Event, list[int]]] = [] + + if name is None: + self.is_creator = True + # Same guard as ShmRingBuffer: SharedMemory/ftruncate can succeed + # on an undersized /dev/shm, so check up front instead of SIGBUS + # on first write past tmpfs capacity. + check_shm_free_space(self.total_bytes, allocation_name="ShmTensorArena") + self.shared_memory = shared_memory.SharedMemory( + create=True, size=self.total_bytes + ) + assert self.shared_memory.buf is not None + with self.shared_memory.buf[self.metadata_offset :] as meta: + torch.frombuffer(meta, dtype=torch.uint8).fill_(0) + else: + self.is_creator = False + # same resource_tracker workaround as ShmRingBuffer + with patch( + "multiprocessing.resource_tracker.register", + lambda *args, **kwargs: None, + ): + try: + self.shared_memory = shared_memory.SharedMemory(name=name) + assert self.shared_memory.size >= self.total_bytes + except FileNotFoundError: + # deserialized on a different node; arena unused there + pass + + def handle(self): + return (self.n_reader, self.slot_bytes, self.n_slots, self.shared_memory.name) + + def _meta(self, idx: int) -> memoryview: + start = self.metadata_offset + idx * self.metadata_size + assert self.shared_memory.buf is not None + return self.shared_memory.buf[start : start + self.metadata_size] + + def _slot(self, idx: int, nbytes: int) -> memoryview: + start = idx * self.slot_bytes + assert self.shared_memory.buf is not None + return self.shared_memory.buf[start : start + nbytes] + + # ---- writer side ---- + + def write_tensor(self, t: torch.Tensor) -> int | None: + """Copy tensor bytes into a free slot; return slot idx or None + (caller must then fall back to the default pickle path).""" + nbytes = t.numel() * t.element_size() + if nbytes > self.slot_bytes: + return None + try: + # Computed before claiming a slot: exotic tensors (e.g. conjugate + # bit set) raise RuntimeError here, same as `_reduce_tensor`'s + # identical view -- fail before claiming, not after. + src = t.detach().reshape(-1).view(torch.uint8) + except RuntimeError: + return None + memory_fence() + for probe in range(self.n_slots): + idx = (self._next_slot + probe) % self.n_slots + with self._meta(idx) as meta: + free = meta[0] == 0 or sum(meta[1:]) == self.n_reader + if not free: + continue + meta[0] = 0 # claim + slot_mv = self._slot(idx, nbytes) + try: + dst = torch.frombuffer(slot_mv, dtype=torch.uint8, count=nbytes) + dst.copy_(src) + finally: + slot_mv.release() + with self._meta(idx) as meta: + for i in range(1, self.n_reader + 1): + meta[i] = 0 + memory_fence() + meta[0] = 1 + memory_fence() + self._next_slot = (idx + 1) % self.n_slots + return idx + self._fallbacks += 1 + if self._fallbacks == 1 or self._fallbacks % 100 == 0: + logger.info( + "ShmTensorArena: no free slot (%d bytes, %d fallbacks so far); " + "falling back to the out-of-band pickle path.", + nbytes, + self._fallbacks, + ) + return None + + # ---- reader side ---- + + def _ensure_pinned(self): + """Pin the whole arena mapping via cudaHostRegister in THIS process + (lazy, once). Without it the HtoD of a zero-copy tensor pays + first-touch page faults on the tmpfs mapping plus pageable staging + (~hundreds of ms for 200MB); registration allocates+pins the pages + once, making every later HtoD a true DMA. Failure (no CUDA in this + process, etc.) is fine — the copy still works, just slower.""" + if self._pin_attempted: + return + self._pin_attempted = True + try: + if not current_platform.is_cuda_alike(): + return + import ctypes + + buf = self.shared_memory.buf + assert buf is not None + ptr = ctypes.addressof(ctypes.c_char.from_buffer(buf)) + ret = current_platform.cudart().cudaHostRegister(ptr, self.total_bytes, 0) + self._pinned = int(ret) == 0 + if self._pinned: + self._pinned_ptr = ptr + logger.info( + "ShmTensorArena: cudaHostRegister(%d MB) -> %s", + self.total_bytes >> 20, + "pinned" if self._pinned else f"error {int(ret)}", + ) + except Exception as e: + logger.info("ShmTensorArena: host-register skipped: %s", e) + + def _unpin(self): + """Undo the pin from _ensure_pinned via cudaHostUnregister; must run + before the mapping is closed. Failures are ignored — at interpreter + shutdown the CUDA context may already be gone, and the registration + dies with the process anyway.""" + if not self._pinned: + return + self._pinned = False + with suppress(Exception): + current_platform.cudart().cudaHostUnregister(self._pinned_ptr) + + def get_tensor( + self, idx: int, nbytes: int, dtype: torch.dtype, shape: tuple[int, ...] + ) -> torch.Tensor: + """Zero-copy view of a slot as a tensor. The slot is queued for + release once the returned tensor's storage is freed, i.e. once the + caller has dropped it along with any views/slices or wrappers (such + as `torch.nn.Parameter`) sharing that storage.""" + self._ensure_pinned() + # The tensor's storage holds the only reference to `holder` and drops + # it when freed, firing the finalizer. (`weakref.finalize` can't + # target the tensor itself, which views and wrappers don't keep alive, + # nor the memoryview, which doesn't support weakrefs.) + holder = np.frombuffer(self._slot(idx, nbytes), dtype=np.uint8) + # Binds to the current list OBJECT: `flush_releases` must mutate it in + # place, never rebind `self._pending_release`. + weakref.finalize(holder, self._pending_release.append, idx) + return torch.from_numpy(holder).view(dtype).view(shape) + + def _mark_released(self, idxs: list[int]): + """Set THIS reader's done flag on the given slots; a slot becomes + reusable once every reader has done so.""" + for idx in idxs: + with self._meta(idx) as meta: + meta[1 + self.reader_rank] = 1 + memory_fence() + + def _record_release_event(self): + """Record a CUDA event on the current (compute) stream, or return None + if this process has no usable CUDA context. Called from flush_releases, + which runs after the previous step's H2D was enqueued on that same + stream, so the event is ordered strictly after that H2D.""" + try: + event = current_platform.Event() + event.record() + return event + except Exception: + return None + + def flush_releases(self): + """Retire this reader's arena slots that are ready for writer reuse. + + Called on every `MessageQueue.dequeue`. `_pending_release` holds + slots whose tensor storage was just freed (see `get_tensor`). On + the pinned fast path that tensor was the source of an async + `non_blocking=True` H2D that can outlive the collection, so instead + of marking released immediately we record a CUDA event (ordered + after that H2D on the same stream) and move the slot to + `_deferred_releases` until `event.query()` says it's done -- a + not-yet-complete slot just waits one more `dequeue`. Unpinned (or no + CUDA context), the H2D already staged synchronously, so the slot is + safe to release right away. + + A slot only becomes free for the writer (`write_tensor`) once every + reader has marked it released (`_mark_released`). + """ + if self._deferred_releases: + still_pending = [] + for deferred_event, idxs in self._deferred_releases: + if deferred_event.query(): + self._mark_released(idxs) + else: + still_pending.append((deferred_event, idxs)) + self._deferred_releases = still_pending + + if not self._pending_release: + return + # Drain in place with atomic pops, never rebind (see get_tensor): a + # finalizer may append from another thread concurrently. + idxs = [] + while self._pending_release: + idxs.append(self._pending_release.pop()) + + event = self._record_release_event() if self._pinned else None + if event is None: + self._mark_released(idxs) + else: + self._deferred_releases.append((event, idxs)) + + def __del__(self): + if hasattr(self, "shared_memory"): + self._unpin() + try: + self.shared_memory.close() + if self.is_creator: + self.shared_memory.unlink() + except BufferError: + # zero-copy tensor views may still hold exported pointers at + # interpreter shutdown; the mapping dies with the process. + pass + + +def _rebuild_arena_tensor(arena_name, slot_idx, nbytes, dtype_str, shape): + """Unpickle hook: rebuild a tensor as a zero-copy view of an arena slot.""" + arena = _TENSOR_ARENAS[arena_name] + return arena.get_tensor(slot_idx, nbytes, getattr(torch, dtype_str), shape) + + +def _rebuild_arena_parameter(arena_name, slot_idx, nbytes, dtype_str, shape): + """Unpickle hook for `torch.nn.Parameter` (shares the slot's storage).""" + t = _rebuild_arena_tensor(arena_name, slot_idx, nbytes, dtype_str, shape) + return torch.nn.Parameter(t, requires_grad=False) + + +# Types `_ArenaPickler.reducer_override` diverts, mapped to the rebuild hook +# that reconstructs them. Any other type -- including other `torch.Tensor` +# subclasses -- falls through to the normal pickle path, since an arbitrary +# subclass can't be faithfully rebuilt from just the slot's bytes. +_ARENA_REBUILD_FNS = { + torch.Tensor: _rebuild_arena_tensor, + torch.nn.Parameter: _rebuild_arena_parameter, +} + + +class _ArenaPickler(pickle.Pickler): + """Pickler that diverts large contiguous CPU tensors into the arena. + + `reducer_override` is consulted before an object's normal reduction. For + a qualifying tensor (exact type in `_ARENA_REBUILD_FNS`, CPU, strided, + contiguous, not `requires_grad`, above the divert-size threshold -- the + same criteria `_reduce_tensor` uses, plus size) it does one memcpy into a + free slot and returns a tiny rebuild stub. Anything that doesn't qualify, + or finds the arena full, returns `NotImplemented` and falls through to + the normal out-of-band pickle path, exactly as if no arena were attached. + """ + + def __init__(self, file, arena: ShmTensorArena, buffer_callback=None): + super().__init__( + file, + protocol=pickle.HIGHEST_PROTOCOL, + buffer_callback=buffer_callback, + ) + self.arena = arena + + def reducer_override(self, obj): + rebuild_fn = _ARENA_REBUILD_FNS.get(type(obj)) + if ( + rebuild_fn is not None + and obj.device.type == "cpu" + and obj.layout == torch.strided + and obj.is_contiguous() + and not obj.requires_grad + and obj.numel() * obj.element_size() >= _ARENA_MIN_BYTES + ): + idx = self.arena.write_tensor(obj) + if idx is not None: + return ( + rebuild_fn, + ( + self.arena.shared_memory.name, + idx, + obj.numel() * obj.element_size(), + str(obj.dtype).removeprefix("torch."), + tuple(obj.shape), + ), + ) + return NotImplemented diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 5f32ec8deb9f..31777e266f8b 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -622,6 +622,7 @@ class EngineArgs: allow_deprecated_quantization: bool = ModelConfig.allow_deprecated_quantization enforce_eager: bool = ModelConfig.enforce_eager disable_custom_all_reduce: bool = ParallelConfig.disable_custom_all_reduce + enable_shm_tensor_arena: bool = ParallelConfig.enable_shm_tensor_arena language_model_only: bool = MultiModalConfig.language_model_only limit_mm_per_prompt: dict[str, int | dict[str, int]] = get_field( MultiModalConfig, "limit_per_prompt" @@ -1274,6 +1275,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--all2all-backend", **parallel_kwargs["all2all_backend"] ) parallel_group.add_argument("--enable-dbo", **parallel_kwargs["enable_dbo"]) + parallel_group.add_argument( + "--enable-shm-tensor-arena", + **parallel_kwargs["enable_shm_tensor_arena"], + ) parallel_group.add_argument( "--ubatch-size", **parallel_kwargs["ubatch_size"], @@ -2534,6 +2539,7 @@ def create_engine_config( enable_elastic_ep=self.enable_elastic_ep, elastic_ep_max_dp_size=self.elastic_ep_max_dp_size, enable_dbo=self.enable_dbo, + enable_shm_tensor_arena=self.enable_shm_tensor_arena, ubatch_size=self.ubatch_size, dbo_decode_token_threshold=self.dbo_decode_token_threshold, dp_sync_interval=self.dp_sync_interval, diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index cad34368546d..749f46b42872 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -166,6 +166,7 @@ def _init_executor(self) -> None: self.local_world_size, max_chunk_bytes=max_chunk_bytes, connect_ip=mq_connect_ip, + enable_shm_tensor_arena=(self.parallel_config.enable_shm_tensor_arena), ) scheduler_output_handle = self.rpc_broadcast_mq.export_handle() # Create workers From 03a7d3776fc2f520a70c8d2da1724eab6df6d0e1 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Fri, 9 Oct 2026 19:33:47 -0400 Subject: [PATCH 212/624] [Perf] Support minimum KV splits configuration for TokenSpeed MLA decode (#60345) Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Signed-off-by: wzhao18 Co-authored-by: Codex --- tests/engine/test_arg_utils.py | 6 ++ tests/v1/attention/test_mla_backends.py | 27 ++++++-- vllm/config/attention.py | 5 ++ .../attention/backends/mla/tokenspeed_mla.py | 67 ++++++++++++++----- 4 files changed, 83 insertions(+), 22 deletions(-) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index fc9f0229c46c..3e3a82f80fbd 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -503,6 +503,12 @@ def test_attention_config(): engine_args = EngineArgs.from_cli_args(args) assert engine_args.attention_config == AttentionConfig() + args = parser.parse_args( + ["--attention-config", '{"tokenspeed_mla_min_split_kv": 8}'] + ) + engine_args = EngineArgs.from_cli_args(args) + assert engine_args.attention_config.tokenspeed_mla_min_split_kv == 8 + # set backend via dot notation args = parser.parse_args(["--attention-config.backend", "FLASH_ATTN"]) assert args is not None diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 9c504f31ef7e..91bd9008f078 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -1098,6 +1098,7 @@ def test_flashinfer_mla_dspark_dcp_supports_target_and_draft(monkeypatch): ], ) @pytest.mark.parametrize("cp_interleave_size", [1, 16, 896]) +@pytest.mark.parametrize("min_split_kv", [1, 8]) def test_tokenspeed_mla_decode_contract( monkeypatch, causal, @@ -1105,6 +1106,7 @@ def test_tokenspeed_mla_decode_contract( dcp_world_size, dcp_rank, cp_interleave_size, + min_split_kv, ): decode_call = None num_decodes = 2 @@ -1148,7 +1150,12 @@ def fake_decode(**kwargs): impl.scale = 1.0 impl.softmax_scale = 1.0 impl.output_scale = 1.0 - impl._workspace_buffer = torch.empty(1, dtype=torch.int8) + impl._min_split_kv = min_split_kv + impl._max_decode_tokens = num_decode_tokens + workspace = torch.empty(1, dtype=torch.int8) + monkeypatch.setattr( + tokenspeed_mla_module, "_get_workspace", lambda *args: workspace + ) metadata = SimpleNamespace( num_decodes=num_decodes, @@ -1183,6 +1190,8 @@ def fake_decode(**kwargs): assert lse.shape == (num_decode_tokens, num_heads) assert decode_call is not None + assert decode_call["workspace_buffer"] is workspace + assert decode_call["min_split_kv"] == min_split_kv assert decode_call["query"].shape == ( num_decodes, tokens_per_decode, @@ -1217,16 +1226,23 @@ def fake_decode(**kwargs): pytest.param(896, 8 * 896 + 17, id="block-interleave-full-cycle"), ], ) -def test_tokenspeed_mla_dcp_matches_unsharded_decode(cp_interleave_size, seq_len): +def test_tokenspeed_mla_dcp_matches_unsharded_decode( + monkeypatch, cp_interleave_size, seq_len +): from tokenspeed_mla import tokenspeed_mla_decode from vllm.v1.attention.ops.dcp import ( _lse_weighted_combine, mask_dcp_empty_shards_, ) + from vllm.v1.worker.workspace import WorkspaceManager torch.manual_seed(7) device = torch.device("cuda") + workspace_manager = WorkspaceManager(device) + monkeypatch.setattr( + tokenspeed_mla_module, "current_workspace_manager", lambda: workspace_manager + ) dcp_world_size = 8 kernel_block_size = 64 num_heads = 128 @@ -1257,7 +1273,9 @@ def make_paged_cache(tokens: torch.Tensor): global_tokens = (torch.randn(seq_len, head_size, device=device) * 0.1).to( torch.float8_e4m3fn ) - workspace = tokenspeed_mla_module._get_workspace(device, num_heads, kv_lora_rank) + workspace = tokenspeed_mla_module._get_workspace( + device, num_heads, kv_lora_rank, 1, query.shape[1] + ) global_cache, global_block_table = make_paged_cache(global_tokens) global_seq_len = torch.tensor([seq_len], dtype=torch.int32, device=device) reference, _ = tokenspeed_mla_decode( @@ -1301,7 +1319,8 @@ def make_paged_cache(tokens: torch.Tensor): impl.scale = head_size**-0.5 impl.softmax_scale = None impl.output_scale = None - impl._workspace_buffer = workspace + impl._min_split_kv = 1 + impl._max_decode_tokens = query.shape[1] metadata = SimpleNamespace( num_decodes=1, num_decode_tokens=query.shape[1], diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 8a9be6336c5e..80548f58f2ca 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -61,6 +61,11 @@ class AttentionConfig: flash_attn_max_num_splits_for_cuda_graph: int = 32 """Flash Attention max number splits for cuda graph decode.""" + tokenspeed_mla_min_split_kv: int = Field(default=1, ge=1, le=256) + """Minimum KV splits for TokenSpeed MLA decode. The default of 1 preserves + automatic scheduling. Larger values request a split floor and may increase + workspace memory. Only applies to the TokenSpeed MLA backend.""" + tq_max_kv_splits_for_cuda_graph: int = 32 """TurboQuant max NUM_KV_SPLITS for cuda graph decode. Fixes the split count so grid dimensions are constant across captures, diff --git a/vllm/v1/attention/backends/mla/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/tokenspeed_mla.py index 96859d7f9884..8d41cd178ed8 100644 --- a/vllm/v1/attention/backends/mla/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/tokenspeed_mla.py @@ -24,6 +24,7 @@ AttentionType, MultipleOf, ) +from vllm.v1.worker.workspace import current_workspace_manager if TYPE_CHECKING: from vllm.config import VllmConfig @@ -38,21 +39,38 @@ # own buffer check. _TOKENSPEED_MAX_Q_LEN = 8 -_g_workspace: dict[torch.device, torch.Tensor] = {} - def _get_workspace( - device: torch.device, num_heads: int, kv_lora_rank: int + device: torch.device, + num_heads: int, + kv_lora_rank: int, + min_split_kv: int, + max_decode_tokens: int, ) -> torch.Tensor: from tokenspeed_mla import get_num_sm + # TokenSpeed reduction scratch: B * H * Q * split_kv * (D + 1) * 4 bytes, + # using the effective query layout and FP32 partial outputs plus LSE. + # Automatic splitting lowers split_kv as the batch grows; split_kv=1 + # needs no reduction scratch, so an SM-based bound covers all batch sizes. needed = ( get_num_sm(device) * num_heads * _TOKENSPEED_MAX_Q_LEN * (kv_lora_rank + 1) * 4 ) - existing = _g_workspace.get(device) - if existing is None or existing.numel() < needed: - _g_workspace[device] = torch.empty(needed, dtype=torch.int8, device=device) - return _g_workspace[device] + if min_split_kv > 1: + # A forced floor breaks that bound: reserve for the maximum batch. + # Packed queries pad head rows to 128; reserve one tile per query token. + needed = max( + needed, + max_decode_tokens + * max(num_heads, 128) + * min_split_kv + * (kv_lora_rank + 1) + * 4, + ) + logger.info_once( + "TokenSpeed MLA decode workspace: %.2f GiB", needed / (1 << 30) + ) + return current_workspace_manager().get_simultaneous(((needed,), torch.int8))[0] class TokenspeedMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): @@ -212,15 +230,25 @@ def __init__( f"got kv_cache_dtype={self.kv_cache_dtype!r}." ) - # Allocate (or fetch the cached) workspace lazily on first forward — - # __init__ runs before the device is necessarily set on the worker; - # we know it for sure at forward time when we see the input tensor. - self._workspace_buffer: torch.Tensor | None = None + config = get_current_vllm_config() + self._min_split_kv = config.attention_config.tokenspeed_mla_min_split_kv + self._max_decode_tokens = 0 + if self._min_split_kv > 1: + spec = config.speculative_config + query_len = 1 + (spec.num_speculative_tokens if spec else 0) + self._max_decode_tokens = max( + min( + config.scheduler_config.max_num_batched_tokens, + config.scheduler_config.max_num_seqs * query_len, + ), + config.compilation_config.max_cudagraph_capture_size or 0, + ) + self.softmax_scale: float | None = None self.output_scale: float | None = None # NIXL resolves interleaving after model construction; retain the config # rather than caching its initial interleave size. - self._parallel_config = get_current_vllm_config().parallel_config + self._parallel_config = config.parallel_config # Pre-JIT BF16 and FP8 prefill kernels here too — decode impl always # runs when tokenspeed is selected, prefill backend may not (user can @@ -290,11 +318,13 @@ def forward_mqa( ) self.output_scale = layer._k_scale_float - if self._workspace_buffer is None: - # Parallelism can change the runtime query head count. - self._workspace_buffer = _get_workspace( - q.device, q.shape[-2], self.kv_lora_rank - ) + workspace = _get_workspace( + q.device, + q.shape[-2], + self.kv_lora_rank, + self._min_split_kv, + self._max_decode_tokens, + ) # vLLM kv_c_and_k_pe_cache is already (num_blocks, block_size, head_size). # tokenspeed_mla_decode wants 3D — pass as-is (no unsqueeze, unlike trtllm). @@ -302,7 +332,7 @@ def forward_mqa( kernel_out = tokenspeed_mla_decode( query=q, kv_cache=kv_c_and_k_pe_cache, - workspace_buffer=self._workspace_buffer, + workspace_buffer=workspace, kv_lora_rank=self.kv_lora_rank, qk_rope_head_dim=self.qk_rope_head_dim, block_tables=block_tables, @@ -317,6 +347,7 @@ def forward_mqa( cp_world=self.dcp_world_size, cp_rank=self.dcp_rank, cp_interleave_size=self._parallel_config.cp_kv_cache_interleave_size, + min_split_kv=self._min_split_kv, ) if return_lse: o, lse = kernel_out From f8e6e365aec519285d99961015ee270a3871d871 Mon Sep 17 00:00:00 2001 From: Dao007forever Date: Fri, 9 Oct 2026 17:08:08 -0700 Subject: [PATCH 213/624] [Bugfix][Mooncake] Save the EAGLE attention proof with hybrid Mamba partial tails (#60695) Signed-off-by: Dao Le Signed-off-by: Dao007forever Signed-off-by: zjy0516 Co-authored-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Co-authored-by: zjy0516 --- .buildkite/test_areas/models_basic.yaml | 2 + tests/models/kimi_k3/test_prefix_cache.py | 127 +++++- .../unit/test_mooncake_store_coordinator.py | 44 +- .../unit/test_mooncake_store_hma_e2e.py | 95 +++-- .../unit/test_mooncake_store_scheduler.py | 316 ++++++++++++-- .../unit/test_mooncake_store_worker.py | 392 ++++++++++++++++-- .../v1/mooncake/store/coordinator.py | 98 +++-- .../kv_connector/v1/mooncake/store/data.py | 51 ++- .../v1/mooncake/store/scheduler.py | 202 ++++++--- .../kv_connector/v1/mooncake/store/worker.py | 190 ++------- vllm/v1/engine/core.py | 5 +- 11 files changed, 1137 insertions(+), 385 deletions(-) diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index fd21324a7ece..1d9a7fe2f3b1 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -165,6 +165,8 @@ steps: - vllm/v1/core/ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/ + - vllm/distributed/mooncake_store.py - vllm/v1/simple_kv_offload/ - vllm/v1/spec_decode/ - vllm/model_executor/layers/mamba/ diff --git a/tests/models/kimi_k3/test_prefix_cache.py b/tests/models/kimi_k3/test_prefix_cache.py index 20e5d6670e5c..d6819a507ab8 100644 --- a/tests/models/kimi_k3/test_prefix_cache.py +++ b/tests/models/kimi_k3/test_prefix_cache.py @@ -4,7 +4,12 @@ import contextlib import json +import os import random +import shutil +import socket +import subprocess +import time from collections.abc import Iterator from dataclasses import dataclass from math import lcm @@ -60,6 +65,16 @@ } NIXL = {"kv_connector": "NixlConnector", "kv_role": "kv_both"} +MOONCAKE = { + "kv_connector": "MooncakeStoreConnector", + "kv_role": "kv_both", + "kv_load_failure_policy": "recompute", + "kv_connector_extra_config": { + "load_async": True, + "lookup_async": True, + "enable_cross_layers_blocks": False, + }, +} # SimpleCPUOffloadConnector as in the Kimi-K3 recipe. OFFLOAD = { "kv_connector": "SimpleCPUOffloadConnector", @@ -243,10 +258,16 @@ def _tokens_text_logprobs(response: dict): def _reset_gpu_prefix_cache(url: str) -> None: - response = requests.post( - f"{url}/reset_prefix_cache", params={"reset_external": "false"}, timeout=60 - ) - response.raise_for_status() + deadline = time.monotonic() + 60 + while True: + response = requests.post( + f"{url}/reset_prefix_cache", params={"reset_external": "false"}, timeout=60 + ) + response.raise_for_status() + if response.json()["success"]: + return + assert time.monotonic() < deadline, "GPU cache reset timed out" + time.sleep(0.1) @dataclass(frozen=True) @@ -354,29 +375,19 @@ def _format(turns: list[Turn]) -> str: return "\n".join(rows) -@pytest.mark.parametrize( - "prompt_len", - [ - # Blocks=6144, PMU=128. P=7040 is PMU-aligned but not block-aligned: - # ensure lookup reads the cached attention proof at 7040 despite - # the final hit limit of 7039, then applies EAGLE's 128-token drop. - pytest.param(7040, id="eagle-proof-scan"), - # Blocks=6144, PMU=128. P=7296 is an exact PMU multiple: - # ensure the Mamba checkpoint is 7168, not 7040 from applying both - # the P - 1 cap and EAGLE's 128-token drop. - pytest.param(7296, id="eagle-double-cap"), - ], -) -def test_dspark_exact_resend_reuses_prompt_checkpoint(prompt_len: int) -> None: +def _check_resend_reuses_prompt_checkpoint( + deployment: Deployment, prompt_len: int +) -> None: """Resends reuse the Mamba checkpoint and its companion EAGLE attention proof.""" mode = Mode(128, True) salt = str(uuid4()) rng = random.Random(0) prompt = [rng.randint(1000, 150000) for _ in range(prompt_len)] - deployment = DEPLOYMENTS["plain"] with _serve(deployment, mode) as servers: url = servers[0].url_root cold = _complete(url, prompt, salt, max_tokens=1) + if deployment.offload: + _reset_gpu_prefix_cache(url) resend = _complete(url, prompt, salt, max_tokens=1) recompute = _complete(url, prompt, f"{salt}-recompute", max_tokens=1) @@ -393,6 +404,84 @@ def test_dspark_exact_resend_reuses_prompt_checkpoint(prompt_len: int) -> None: ) +@pytest.mark.parametrize( + "prompt_len", + [ + # Blocks=6144, PMU=128. P=7040 is PMU-aligned but not block-aligned: + # ensure lookup reads the cached attention proof at 7040 despite + # the final hit limit of 7039, then applies EAGLE's 128-token drop. + pytest.param(7040, id="eagle-proof-scan"), + # Blocks=6144, PMU=128. P=7296 is an exact PMU multiple: + # ensure the Mamba checkpoint is 7168, not 7040 from applying both + # the P - 1 cap and EAGLE's 128-token drop. + pytest.param(7296, id="eagle-double-cap"), + ], +) +def test_dspark_exact_resend_reuses_prompt_checkpoint(prompt_len: int) -> None: + _check_resend_reuses_prompt_checkpoint(DEPLOYMENTS["plain"], prompt_len) + + +@pytest.fixture +def mooncake_store(tmp_path, monkeypatch) -> Iterator[None]: + """Run a Mooncake master and point the servers' store config at it.""" + if shutil.which("mooncake_master") is None: + pytest.skip("mooncake_master is not installed") + port = get_open_port() + log_path = tmp_path / "mooncake_master.log" + with open(log_path, "w") as log: + master = subprocess.Popen( + [ + "mooncake_master", + f"--port={port}", + f"--metrics_port={get_open_port()}", + "--enable_metric_reporting=false", + "--default_kv_lease_ttl=60000", + ], + stdout=log, + stderr=subprocess.STDOUT, + ) + try: + deadline = time.monotonic() + 60 + while True: + assert master.poll() is None, log_path.read_text() + with ( + contextlib.suppress(OSError), + socket.create_connection(("127.0.0.1", port)), + ): + break + assert time.monotonic() < deadline, "Mooncake master did not start" + time.sleep(0.1) + config_path = tmp_path / "mooncake.json" + config_path.write_text( + json.dumps( + { + "metadata_server": "P2PHANDSHAKE", + "master_server_address": f"127.0.0.1:{port}", + "protocol": "rdma", + "device_name": os.getenv("MOONCAKE_DEVICE_NAME", "mlx5_12"), + "global_segment_size": "4GB", + "local_buffer_size": "1GB", + } + ) + ) + monkeypatch.setenv("MOONCAKE_CONFIG_PATH", str(config_path)) + monkeypatch.setenv("MC_MAX_MR_SIZE", str(4 << 30)) + yield + finally: + master.kill() + master.wait() + + +@pytest.mark.usefixtures("mooncake_store") +def test_mooncake_resend_reuses_prompt_checkpoint() -> None: + """Blocks=6144, PMU=128. P=7449 needs Mamba state at 7296 and attention proof + at 7424, both beyond the normal Mooncake save boundary of 6144: ensure they + remain reusable from the store after a GPU cache reset.""" + _check_resend_reuses_prompt_checkpoint( + Deployment(Instance(kv_config=MOONCAKE), offload=True), 7449 + ) + + @pytest.mark.parametrize( "name, mode", [_case(name, mode) for name in DEPLOYMENTS for mode in MODES] ) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py index 8767b1c0d983..ac52ca07ab33 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py @@ -30,7 +30,13 @@ def _mamba_align(block_size=32): ) -def _make_coord(groups, hash_block_size, use_eagle=False, retention_interval=None): +def _make_coord( + groups, + hash_block_size, + use_eagle=False, + retention_interval=None, + enable_partial_hash_hits=False, +): """Construct a coordinator using the natural LCM of group block sizes as the scheduler block size — mirrors ``resolve_kv_cache_block_sizes`` for the test fixtures.""" @@ -42,6 +48,7 @@ def _make_coord(groups, hash_block_size, use_eagle=False, retention_interval=Non hash_block_size=hash_block_size, use_eagle=use_eagle, retention_interval=retention_interval, + enable_partial_hash_hits=enable_partial_hash_hits, ) @@ -219,8 +226,7 @@ def test_coordinator_fine_grained_partial_tail_hit(): KVCacheGroupSpec(["L0"], _full(32)), KVCacheGroupSpec(["L1"], _mamba_align(32)), ] - coord = _make_coord(groups, hash_block_size=16) - assert coord.enable_partial_hash_hits + coord = _make_coord(groups, hash_block_size=16, enable_partial_hash_hits=True) hs = _hashes(4) # 4 hash units of 16 = 64 tokens; block 0 = [0,32), etc. # Both groups: full block 0 (key = last sub-hash hs[1]) + partial boundary # at token 48 (key = hs[2]). No hs[3] -> block 1 is not full. @@ -287,38 +293,6 @@ def test_store_mask_swa_only_window_around_each_lcm_boundary(): assert masks[1] == [False, False, False, True, False, False, False, True] -def test_store_mask_swa_does_not_double_scale_under_dcp(): - """The coordinator receives DCP-resolved specs (worker.py applies - resolve_dcp_kv_cache_spec) and indexes chunks in that already-scaled block - size. So dcp_world_size must NOT further scale reachable_block_mask; the - mask must be identical for any dcp_world_size given the same specs. - - Regression: previously the coordinator forwarded its dcp_world_size, which - double-scaled the SWA block size and produced a wrong mask when dcp > 1. - """ - full = _full(32) - swa = _swa(block_size=8, sliding_window=8) - groups = [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)] - scheduler_block_size = lcm(32, 8) - - def make(dcp_world_size): - return MooncakeStoreCoordinator( - groups, - scheduler_block_size=scheduler_block_size, - hash_block_size=8, - dcp_world_size=dcp_world_size, - ) - - mask_dcp1 = make(1).store_mask(64) - mask_dcp2 = make(2).store_mask(64) - - # dcp must not change the mask; both match the correct single-block-size - # result (chunks ending at 32 and 64 -> blocks 3 and 7). - expected = [False, False, False, True, False, False, False, True] - assert mask_dcp1[1] == expected - assert mask_dcp2[1] == expected - - def test_store_mask_swa_wider_window_covers_more_blocks_per_lcm(): """Same hybrid layout but sliding_window=16 (= 2 SWA blocks). Each lcm boundary should now span two SWA tail blocks.""" diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index a8354390d71e..0682f9f465f1 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -18,11 +18,15 @@ MooncakeStoreCoordinator, ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( + BoundaryPut, ChunkedTokenDatabase, KeyMetadata, LoadSpec, ReqMeta, ) +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import ( # noqa: E501 + _partial_tail_non_mamba_puts, +) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.worker import ( # noqa: E501 KVCacheStoreRecvingThread, KVCacheStoreSendingThread, @@ -378,6 +382,15 @@ def test_chunked_token_database_hash_block_size_smaller_than_block_size(): assert out[1][2].hex() == fine_hashes[6].hex() +def _resolve_partial_tail(thread, req: ReqMeta) -> ReqMeta: + """Add the tail's non-Mamba puts as the scheduler does before the worker.""" + req.boundary_puts = [BoundaryPut(*put) for put in req.boundary_puts or []] + req.boundary_puts[:0] = _partial_tail_non_mamba_puts( + thread.coord, req, [db.block_size for db in thread.token_databases] + ) + return req + + def test_sub_block_partial_tail_offload_reads_cow_block(): """Sub-block prompt (the 900/128/1536 shape, scaled to 12/4/16): the partial tail is offloaded for both groups under the boundary sub-hash. The @@ -394,7 +407,12 @@ def test_sub_block_partial_tail_offload_reads_cow_block(): KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], mamba), ] - coord = MooncakeStoreCoordinator(groups, scheduler_block_size=16, hash_block_size=4) + coord = MooncakeStoreCoordinator( + groups, + scheduler_block_size=16, + hash_block_size=4, + enable_partial_hash_hits=True, + ) assert coord.enable_partial_hash_hits class _RecordingStore(_DictStore): @@ -441,11 +459,13 @@ def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k): block_ids=([1], [2]), block_hashes=hs, can_save=True, - num_prompt_tokens=20, - boundary_state_offloads=[(1, mamba_cow_block, 12)], + num_prompt_tokens=13, + prefill_end_tokens=20, + completed_token_len=20, + boundary_puts=[(1, mamba_cow_block, 12)], ) - send._maybe_offload_boundary_states(req) + send._maybe_offload_boundary_states(_resolve_partial_tail(send, req)) # boundary = 12 // 4 * 4 = 12 -> keyed by hs[12 // 4 - 1] = hs[2]. partial_hash = hs[2] @@ -473,7 +493,12 @@ def test_offload_syncs_event_before_put(): KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], mamba), ] - coord = MooncakeStoreCoordinator(groups, scheduler_block_size=16, hash_block_size=4) + coord = MooncakeStoreCoordinator( + groups, + scheduler_block_size=16, + hash_block_size=4, + enable_partial_hash_hits=True, + ) event = MagicMock() class _FencedStore(_DictStore): @@ -513,12 +538,13 @@ def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k): block_hashes=hs, can_save=True, num_prompt_tokens=12, + completed_token_len=12, store_job_id=1, - boundary_state_offloads=[(1, 7, 12)], + boundary_puts=[(1, 7, 8)], ) req.current_event = event - send.add_request(req) + send.add_request(_resolve_partial_tail(send, req)) send._handle_request(send.request_queue.get()) assert send.request_queue.qsize() == 0 assert store._data @@ -532,7 +558,8 @@ def test_sub_block_partial_tail_offload_covers_smaller_group_blocks(): offload must persist every FA block up to the boundary — the normal save floors to the lcm, so those blocks are otherwise never written and the consumer's per-group lookup would miss. The mamba boundary block still - reads the core-provided CoW target.""" + reads the core-provided CoW target. Without EAGLE, the 12-token prompt's + resend-safe checkpoint is 8, leaving two attention blocks to save.""" full = FullAttentionSpec(block_size=4, num_kv_heads=8, head_size=64, dtype=None) mamba = MambaSpec( block_size=16, @@ -544,7 +571,12 @@ def test_sub_block_partial_tail_offload_covers_smaller_group_blocks(): KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], mamba), ] - coord = MooncakeStoreCoordinator(groups, scheduler_block_size=16, hash_block_size=4) + coord = MooncakeStoreCoordinator( + groups, + scheduler_block_size=16, + hash_block_size=4, + enable_partial_hash_hits=True, + ) assert coord.enable_partial_hash_hits class _RecordingStore(_DictStore): @@ -590,19 +622,19 @@ def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k): block_hashes=hs, can_save=True, num_prompt_tokens=12, - boundary_state_offloads=[(1, mamba_cow_block, 12)], + completed_token_len=12, + boundary_puts=[(1, mamba_cow_block, 8)], ) - send._maybe_offload_boundary_states(req) + send._maybe_offload_boundary_states(_resolve_partial_tail(send, req)) - # FA (block 4): full blocks ending at 4, 8 and 12, keyed by their normal + # FA (block 4): full blocks ending at 4 and 8, keyed by their normal # block-end hashes; mamba (block 16): the partial boundary block under # the boundary sub-hash, read from the CoW target. expected = { token_dbs[0].key_for(hs[0]): [1 * 512], token_dbs[0].key_for(hs[1]): [2 * 512], - token_dbs[0].key_for(hs[2]): [3 * 512], - token_dbs[1].key_for(hs[2]): [10_000 + mamba_cow_block * 512], + token_dbs[1].key_for(hs[1]): [10_000 + mamba_cow_block * 512], } assert store.puts == expected @@ -611,11 +643,10 @@ def test_worker_lookup_hits_sub_block_partial_tail(): """worker.lookup must query sub-block keys when partial hash hits are on. Regression test: ``lookup`` hard-coded ``fine_grained = False``, so the - sub-block keys persisted by ``_sub_block_tail_puts`` were never probed and - partial prefix hits silently returned 0 while the store side kept writing - them. Here the mamba block (16) exceeds the hash unit (4), so - ``enable_partial_hash_hits`` is on and the lookup must find the stored - boundary at 12. + partial tail's sub-block keys were never probed and partial prefix hits + silently returned 0 while the store side kept writing them. Here the + mamba block (16) exceeds the hash unit (4), so ``enable_partial_hash_hits`` + is on and the lookup must find the stored boundary at 8. """ full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) mamba = MambaSpec( @@ -644,6 +675,8 @@ def test_worker_lookup_hits_sub_block_partial_tail(): KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], mamba), ], + # The hit alignment the engine core resolves for this config. + cache_hit_alignment_tokens=4, ) vllm_config = _minimal_vllm_config(cache_block_size=16) # Hash unit 4 < mamba block 16 -> partial hash hits are enabled. @@ -672,7 +705,7 @@ def test_worker_lookup_hits_sub_block_partial_tail(): replicate_config=MagicMock(), ) - # Persist the sub-block partial tail at boundary 12 (keyed by hs[12//4-1]). + # A 12-token prompt leaves a resend-safe sub-block checkpoint at 8. hs = [BlockHash(bytes([i + 1]) * 4) for i in range(5)] req = ReqMeta( req_id="r0", @@ -680,15 +713,16 @@ def test_worker_lookup_hits_sub_block_partial_tail(): block_ids=([1], [2]), block_hashes=hs, can_save=True, - num_prompt_tokens=20, - boundary_state_offloads=[(1, 7, 12)], + num_prompt_tokens=12, + completed_token_len=12, + boundary_puts=[(1, 7, 8)], ) - send_thread._maybe_offload_boundary_states(req) + send_thread._maybe_offload_boundary_states(_resolve_partial_tail(send_thread, req)) worker.store = store - # A 13-token prompt sharing the prefix must hit the stored boundary at 12. - assert worker.lookup(num_tokens=13, block_hashes=hs).hit_length == 12 + # A 13-token prompt sharing the prefix must hit the stored boundary at 8. + assert worker.lookup(num_tokens=13, block_hashes=hs).hit_length == 8 def test_worker_setup_tolerates_finer_scratch_group(): @@ -744,6 +778,8 @@ def test_worker_setup_tolerates_finer_scratch_group(): KVCacheGroupSpec(["L1"], mamba), KVCacheGroupSpec(["L2"], scratch), ], + # The hit alignment the engine core resolves for this config. + cache_hit_alignment_tokens=8, ) vllm_config = _minimal_vllm_config(cache_block_size=16) # Hash unit 8 divides the participating groups (16) but not the scratch @@ -777,7 +813,7 @@ def test_worker_setup_tolerates_finer_scratch_group(): group_participates=[True, True], ) - # Persist the sub-block partial tail at boundary 12 (keyed by hs[12//8-1]). + # Persist the sub-block partial tail at boundary 8 (keyed by hs[8//8-1]). hs = [BlockHash(bytes([i + 1]) * 8) for i in range(3)] req = ReqMeta( req_id="r0", @@ -785,10 +821,11 @@ def test_worker_setup_tolerates_finer_scratch_group(): block_ids=([1], [2]), block_hashes=hs, can_save=True, - num_prompt_tokens=20, - boundary_state_offloads=[(1, 7, 12)], + num_prompt_tokens=12, + completed_token_len=12, + boundary_puts=[(1, 7, 8)], ) - send_thread._maybe_offload_boundary_states(req) + send_thread._maybe_offload_boundary_states(_resolve_partial_tail(send_thread, req)) worker.store = store diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py index 7ced38d65a41..47796eb14707 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py @@ -7,7 +7,11 @@ import pytest import torch +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( + MooncakeStoreCoordinator, +) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( + BoundaryPut, LoadSpec, MooncakeLookupResult, MooncakeStoreConnectorMetadata, @@ -17,6 +21,7 @@ ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import ( MooncakeStoreScheduler, + _partial_tail_non_mamba_puts, ) from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_manager import KVCacheBlocks @@ -62,6 +67,31 @@ def _make_bare_scheduler( scheduler._next_store_job_id = 0 scheduler._pinned_saves = {} scheduler._boundary_state_group_ids = frozenset({1}) + scheduler._store_coord = MooncakeStoreCoordinator( + [ + KVCacheGroupSpec( + ["attention"], + FullAttentionSpec( + block_size=max(16, hash_block_size), + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=max(16, hash_block_size), + shapes=((1,),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + max(16, hash_block_size), + hash_block_size, + enable_partial_hash_hits=enable_partial_hash_hits, + ) return scheduler @@ -121,6 +151,7 @@ def _make_new_scheduler_output() -> SimpleNamespace: req_id="req-0", num_computed_tokens=0, prompt_token_ids=list(range(32)), + num_prompt_tokens=32, prefill_token_ids=None, block_ids=([0, 1],), block_hashes=[b"h0", b"h1"], @@ -199,6 +230,7 @@ def _add_unfinished_request( ) -> None: request = SimpleNamespace( all_token_ids=token_ids, + num_prompt_tokens=prefill_end_tokens, block_hashes=block_hashes, num_output_placeholders=0, ) @@ -268,6 +300,7 @@ def _make_qsa_hybrid_cache_config(): def test_scheduler_projects_nonprefix_groups_and_mamba_ids(): vllm_config = SimpleNamespace( + speculative_config=None, kv_transfer_config=SimpleNamespace( kv_role="kv_both", kv_connector_extra_config={} ), @@ -528,6 +561,7 @@ def _make_pending_load_unfinished_request( ) -> None: request = SimpleNamespace( num_tokens=num_tokens, + num_prompt_tokens=num_tokens, block_hashes=block_hashes, num_output_placeholders=0, ) @@ -594,6 +628,7 @@ def _make_resumed_unfinished_request( ) -> None: request = SimpleNamespace( all_token_ids=token_ids, + num_prompt_tokens=32, block_hashes=block_hashes, num_computed_tokens=num_computed_tokens, num_output_placeholders=0, @@ -796,13 +831,24 @@ def test_from_request_tracker_load_with_can_load_false_still_saves(): assert tracker.num_saved_tokens == 48 -def test_from_request_tracker_no_load_saves_normally(): +@pytest.mark.parametrize("token_len,save_partial_tail", [(48, False), (13, True)]) +def test_from_request_tracker_no_load_saves_normally(token_len, save_partial_tail): tracker = RequestTracker( req_id="req-0", - token_len=48, + token_len=token_len, allocated_block_ids=([0, 1, 2],), num_saved_tokens=0, + prefill_end_tokens=token_len, ) + if save_partial_tail: + tracker.token_len = 8 + assert ( + ReqMeta.from_request_tracker( + tracker, 16, save_partial_tail=True, num_prompt_tokens=token_len + ) + is None + ) + tracker.token_len = token_len req_meta = ReqMeta.from_request_tracker( tracker, @@ -810,12 +856,70 @@ def test_from_request_tracker_no_load_saves_normally(): load_spec=None, skip_save=False, block_hashes=[b"h0", b"h1", b"h2"], + save_partial_tail=save_partial_tail, + num_prompt_tokens=token_len, ) assert req_meta is not None assert req_meta.can_save is True assert req_meta.load_spec is None - assert tracker.num_saved_tokens == 48 + assert req_meta.token_len_chunk == token_len // 16 * 16 + assert req_meta.completed_token_len == token_len + assert req_meta.boundary_puts is None + assert tracker.num_saved_tokens == token_len // 16 * 16 + assert ( + ReqMeta.from_request_tracker( + tracker, + 16, + save_partial_tail=save_partial_tail, + num_prompt_tokens=token_len, + ) + is None + ) + + +def test_partial_tail_is_resolved_once_on_the_prompt_completing_save(): + scheduler = _make_bare_scheduler( + hash_block_size=4, enable_partial_hash_hits=True, save_decode_cache=True + ) + scheduler._store_group_ids = (0, 1) + _add_unfinished_request( + scheduler, + token_ids=list(range(48)), + block_hashes=[bytes([i]) for i in range(12)], + prefill_end_tokens=45, + ) + block_ids = ([0, 1, 2], [50]) + scheduler._request_trackers["req-0"].allocated_block_ids = block_ids + + def step(num_computed_tokens: int, num_scheduled_tokens: int): + return scheduler.build_connector_meta( + SimpleNamespace( + finished_req_ids=set(), + preempted_req_ids=set(), + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace( + req_ids=["req-0"], + new_block_ids=[()], + num_computed_tokens=[num_computed_tokens], + resumed_req_ids=set(), + ), + num_scheduled_tokens={"req-0": num_scheduled_tokens}, + scheduled_spec_decode_tokens={}, + kv_connector_block_state=_make_connector_block_state(block_ids), + ) + ).requests + + # The save that completes the prompt carries the tail for the worker. + (req_meta,) = step(44, 1) + assert req_meta.publish_partial_tail + assert req_meta.boundary_puts == [(0, 2, 44)] + + # A later save does not recompute or resend it. + (req_meta,) = step(45, 3) + assert req_meta.token_len_chunk == 48 + assert not req_meta.publish_partial_tail + assert req_meta.boundary_puts is None class _StubLookupClient: @@ -1018,7 +1122,7 @@ def _add_pending_partial_tail_request( all_token_ids=list(range(num_tokens)), block_hashes=block_hashes, num_output_placeholders=0, - num_prompt_tokens=12, + num_prompt_tokens=13, ) scheduler._unfinished_requests["req-0"] = (request, block_ids) scheduler._request_trackers["req-0"] = RequestTracker( @@ -1057,7 +1161,7 @@ def test_pending_partial_tail_emits_offload_only_reqmeta(): scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) out = _add_pending_partial_tail_request( scheduler, - num_tokens=12, + num_tokens=13, block_hashes=[b"h0", b"h1", b"h2"], block_ids=([0],), ) @@ -1069,8 +1173,8 @@ def test_pending_partial_tail_emits_offload_only_reqmeta(): assert req_meta.req_id == "req-0" assert req_meta.can_save is True assert req_meta.token_len_chunk == 0 - assert req_meta.boundary_state_offloads == [(1, 7, 12)] - assert req_meta.num_prompt_tokens == 12 + assert req_meta.boundary_puts == [(1, 7, 12)] + assert req_meta.num_prompt_tokens == 13 assert req_meta.block_ids == ([0],) store_job_id = req_meta.store_job_id assert scheduler._pinned_saves[store_job_id][0] == [7] @@ -1081,34 +1185,71 @@ def test_pending_partial_tail_emits_offload_only_reqmeta(): assert tracker.has_pending_offload is True -def test_finished_partial_tail_is_pre_pinned_as_store_job(): +@pytest.mark.parametrize("attention_block_size", [4, 8, 16]) +@pytest.mark.parametrize("use_eagle", [False, True]) +@pytest.mark.parametrize("num_workers", [1, 2]) +def test_finished_partial_tail_is_pre_pinned_as_store_job( + attention_block_size, use_eagle, num_workers +): scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) scheduler._store_group_ids = (0, 1) scheduler.client = SimpleNamespace(discard=lambda *_: None) + scheduler._num_workers = num_workers + groups = list(scheduler._store_coord.kv_cache_groups) + groups[0] = KVCacheGroupSpec( + ["attention"], + FullAttentionSpec( + block_size=attention_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + is_eagle_group=use_eagle, + ) + scheduler._store_coord = MooncakeStoreCoordinator( + groups, 16, 4, use_eagle, enable_partial_hash_hits=True + ) + prompt_tokens = 49 if use_eagle else 45 + attention_ids = list( + range(1, (prompt_tokens + attention_block_size - 1) // attention_block_size + 1) + ) + block_ids = (attention_ids, [50]) request = SimpleNamespace( request_id="req-0", - block_hashes=[b"h0", b"h1", b"h2"], + block_hashes=[bytes([i]) for i in range(prompt_tokens // 4)], + num_computed_tokens=prompt_tokens, + num_prompt_tokens=prompt_tokens, ) scheduler._request_trackers["req-0"] = RequestTracker( req_id="req-0", - token_len=12, - allocated_block_ids=([3], [9]), - token_ids=list(range(12)), - prefill_end_tokens=12, + token_len=prompt_tokens, + allocated_block_ids=block_ids, + token_ids=list(range(prompt_tokens)), + prefill_end_tokens=prompt_tokens, ) - block_ids = ([3], [9]) + pool = scheduler._gpu_block_pool + owned_ids = attention_ids + [50] + pool.touch([pool.blocks[i] for i in owned_ids]) delay_free = scheduler.register_finished_partial_tail( request, block_ids, - [(1, 9, 12)], + [(1, 50, 44)], ) # The exact source is pinned immediately, so the request can be freed # before the next connector metadata build. assert delay_free is False - assert scheduler._gpu_block_pool.blocks[9].ref_cnt == 1 - assert scheduler._gpu_block_pool.blocks[3].ref_cnt == 0 + proof_end = 48 if use_eagle else 44 + expected = attention_ids[ + 32 // attention_block_size : (proof_end + attention_block_size - 1) + // attention_block_size + ] + pool.free_blocks([pool.blocks[i] for i in owned_ids]) + assert all( + pool.blocks[i].ref_cnt == (1 if i in expected else 0) for i in attention_ids + ) + assert pool.blocks[50].ref_cnt == 1 out = SimpleNamespace( finished_req_ids={"req-0"}, @@ -1132,13 +1273,116 @@ def test_finished_partial_tail_is_pre_pinned_as_store_job(): assert req_meta.token_len_chunk == 0 assert req_meta.block_ids == block_ids assert req_meta.block_hashes == request.block_hashes - assert req_meta.boundary_state_offloads == [(1, 9, 12)] - assert scheduler._pinned_saves[req_meta.store_job_id][0] == [9] - assert scheduler._gpu_block_pool.blocks[9].ref_cnt == 1 + assert req_meta.boundary_puts == [ + (0, attention_ids[idx], min((idx + 1) * attention_block_size, proof_end)) + for idx in range( + 32 // attention_block_size, + (proof_end + attention_block_size - 1) // attention_block_size, + ) + ] + [(1, 50, 44)] + assert set(scheduler._pinned_saves[req_meta.store_job_id][0]) == {50, *expected} + assert all(pool.blocks[i].ref_cnt == 1 for i in expected + [50]) assert scheduler._finished_partial_tail_metas == {} + for _ in range(num_workers - 1): + scheduler.update_connector_output( + _make_worker_output({req_meta.store_job_id: 1}) + ) + assert all(pool.blocks[i].ref_cnt == 1 for i in expected + [50]) scheduler.update_connector_output(_make_worker_output({req_meta.store_job_id: 1})) - assert scheduler._gpu_block_pool.blocks[9].ref_cnt == 0 + assert all(pool.blocks[i].ref_cnt == 0 for i in owned_ids) + + +def test_partial_tail_boundary_follows_eagle_block_drop_not_group_flags(): + """Core drops the Mamba checkpoint whenever EAGLE block drop is on, even if + only the Mamba group carries the eagle flag.""" + groups = [ + KVCacheGroupSpec( + ["attention"], + FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=16, + shapes=((1,),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + is_eagle_group=True, + ), + ] + coord = MooncakeStoreCoordinator(groups, 16, 4, use_eagle=True) + # Core's checkpoint for a 49-token prompt: 48 minus one EAGLE hash unit. + req_meta = ReqMeta( + req_id="req-0", + token_len_chunk=0, + block_ids=([1, 2, 3, 4], [50]), + block_hashes=[bytes([i]) for i in range(12)], + num_prompt_tokens=49, + completed_token_len=49, + boundary_puts=[BoundaryPut(1, 50, 44)], + ) + + assert _partial_tail_non_mamba_puts(coord, req_meta, [16, 16]) == [ + BoundaryPut(0, 3, 44) + ] + + +@pytest.mark.parametrize( + "completed, mamba_tails, publish, expected", + [ + # The chunk ending at the junction (40) hands off its Mamba state. Its + # attention KV comes from the normal saves, so nothing is added here. + (40, [40], False, []), + # The prompt-completing save: blocks from the LCM floor 64 up to the + # checkpoint 68 plus the EAGLE proof unit, whether or not the junction + # hand-off is in the same save. + (75, [68], True, [BoundaryPut(0, 17, 68), BoundaryPut(0, 18, 72)]), + (75, [40, 68], True, [BoundaryPut(0, 17, 68), BoundaryPut(0, 18, 72)]), + ], +) +def test_shared_prefix_junction_handoff_adds_no_attention_tail( + completed, mamba_tails, publish, expected +): + """``--enable-mamba-shared-prefix-checkpoint`` also hands off a sub-block + Mamba state at the shared-prefix junction, below the prompt checkpoint.""" + groups = [ + KVCacheGroupSpec( + ["attention"], + FullAttentionSpec( + block_size=4, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + is_eagle_group=True, + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=16, + shapes=((1,),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ] + coord = MooncakeStoreCoordinator( + groups, 16, 4, use_eagle=True, enable_partial_hash_hits=True + ) + # 75-token prompt: checkpoint 72 - 4 = 68. + req_meta = ReqMeta( + req_id="req-0", + token_len_chunk=0, + block_ids=(list(range(1, 20)), [50, 51, 52, 53, 54]), + block_hashes=[bytes([i]) for i in range(18)], + num_prompt_tokens=75, + completed_token_len=completed, + boundary_puts=[BoundaryPut(1, 60 + i, n) for i, n in enumerate(mamba_tails)], + publish_partial_tail=publish, + ) + + assert _partial_tail_non_mamba_puts(coord, req_meta, [4, 16]) == expected def test_decode_boundary_state_offload_dropped_unclaimed(): @@ -1235,7 +1479,7 @@ def test_boundary_state_group_ids_are_remapped_to_store_projection(): scheduler._handle_boundary_state_offloads({"req-0": [(2, 7, 800)]}, meta) - assert meta.requests[0].boundary_state_offloads == [(1, 7, 800)] + assert meta.requests[0].boundary_puts == [(1, 7, 800)] def test_resumed_prefill_claims_boundaries_past_prompt_length(): @@ -1245,13 +1489,13 @@ def test_resumed_prefill_claims_boundaries_past_prompt_length(): # missing for boundaries full attention does store and no joint hit could # complete there. scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) - _register_offload_request(scheduler, prefill_end_tokens=20, num_prompt_tokens=12) - out = _make_offload_only_output([(1, 7, 16), (1, 9, 20), (1, 11, 24)]) + _register_offload_request(scheduler, prefill_end_tokens=37, num_prompt_tokens=13) + out = _make_offload_only_output([(1, 7, 16), (1, 9, 32), (1, 11, 48)]) meta = scheduler.build_connector_meta(out) - # 16 and 20 are inside the resumed prefill; 24 is past it. - assert meta.requests[0].boundary_state_offloads == [(1, 7, 16), (1, 9, 20)] + # Aligned states at 16 and 32 are inside the resumed prefill; 48 is past it. + assert meta.requests[0].boundary_puts == [(1, 7, 16), (1, 9, 32)] store_job_id = meta.requests[0].store_job_id assert scheduler._pinned_saves[store_job_id][0] == [7, 9] @@ -1267,7 +1511,7 @@ def test_boundary_state_job_pins_exact_blocks_once(): meta = scheduler.build_connector_meta(out) req_meta = meta.requests[0] - assert req_meta.boundary_state_offloads == [ + assert req_meta.boundary_puts == [ (1, 7, 16), (1, 8, 32), (1, 9, 48), @@ -1294,6 +1538,7 @@ def test_store_job_pins_current_non_null_non_mamba_blocks(): request = SimpleNamespace( all_token_ids=list(range(48)), block_hashes=[bytes([i]) for i in range(12)], + num_prompt_tokens=48, num_output_placeholders=0, ) scheduler._unfinished_requests["req-0"] = (request, ([7, 2, 0], [21, 0, 22])) @@ -1379,7 +1624,7 @@ def test_preemption_and_request_id_reuse_do_not_release_inflight_job(): _register_offload_request(scheduler, prefill_end_tokens=64, num_prompt_tokens=64) second = scheduler.build_connector_meta(_make_offload_only_output([(1, 8, 32)])) second_job_id = second.requests[0].store_job_id - assert second.requests[0].boundary_state_offloads == [(1, 8, 32)] + assert second.requests[0].boundary_puts == [(1, 8, 32)] scheduler.update_connector_output(_make_worker_output({first_job_id: 1})) assert scheduler._gpu_block_pool.blocks[7].ref_cnt == 0 @@ -1412,9 +1657,9 @@ def test_resumed_partial_tail_uses_exact_boundary(): meta = scheduler.build_connector_meta(out) assert len(meta.requests) == 1 - assert meta.requests[0].boundary_state_offloads == [(1, 7, 12)] - # Ordinary metadata retains the full resumed prefill range. - assert meta.requests[0].num_prompt_tokens == 20 + assert meta.requests[0].boundary_puts == [(1, 7, 12)] + assert meta.requests[0].num_prompt_tokens == 13 + assert meta.requests[0].prefill_end_tokens == 20 tracker = scheduler._request_trackers["req-0"] assert tracker.num_saved_tokens == 0 assert tracker.has_pending_offload is True @@ -1427,7 +1672,7 @@ def test_resumed_partial_tail_attached_to_save_keeps_exact_boundary(): all_token_ids=list(range(48)), block_hashes=[b"h0", b"h1", b"h2"], num_output_placeholders=0, - num_prompt_tokens=36, + num_prompt_tokens=37, ) scheduler._unfinished_requests["req-0"] = (request, ([0, 1],)) scheduler._request_trackers["req-0"] = RequestTracker( @@ -1445,8 +1690,9 @@ def test_resumed_partial_tail_attached_to_save_keeps_exact_boundary(): assert len(meta.requests) == 1 assert meta.requests[0].can_save is True - assert meta.requests[0].boundary_state_offloads == [(0, 7, 36)] - assert meta.requests[0].num_prompt_tokens == 48 + assert meta.requests[0].boundary_puts == [(0, 7, 36)] + assert meta.requests[0].num_prompt_tokens == 37 + assert meta.requests[0].prefill_end_tokens == 48 # Ordinary saving still covers the full resumed prefill range. tracker = scheduler._request_trackers["req-0"] assert tracker.num_saved_tokens == 48 @@ -1460,7 +1706,7 @@ def test_partial_tail_cow_block_is_referenced_for_the_job(): scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) out = _add_pending_partial_tail_request( scheduler, - num_tokens=12, + num_tokens=13, block_hashes=[b"h0", b"h1", b"h2"], block_ids=([0],), ) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index cb6bd8d995be..09d27eb947a0 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -30,6 +30,7 @@ ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( BlobBlockHashes, + BoundaryPut, ChunkedTokenDatabase, KeyMetadata, LBHNCStoreLayout, @@ -44,6 +45,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import ( MooncakeStoreConnectorStats, ) +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import ( # noqa: E501 + _partial_tail_non_mamba_puts, +) from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash from vllm.v1.kv_cache_interface import ( @@ -942,6 +946,15 @@ def test_store_sending_thread_retries_skipped_range_after_pressure(): assert store.batch_put_from_multi_buffers.call_args.args[0] == keys +def _resolve_partial_tail(thread, req: ReqMeta) -> ReqMeta: + """Add the tail's non-Mamba puts as the scheduler does before the worker.""" + req.boundary_puts = [BoundaryPut(*put) for put in req.boundary_puts or []] + req.boundary_puts[:0] = _partial_tail_non_mamba_puts( + thread.coord, req, [db.block_size for db in thread.token_databases] + ) + return req + + def _make_partial_tail_send_thread( store, *, @@ -954,6 +967,8 @@ def _make_partial_tail_send_thread( hash_block_size=4, lcm_block_size=16, mamba_group_ids={1}, + use_eagle=False, + eagle_proof_margin_by_group={}, ) db = ChunkedTokenDatabase( KeyMetadata("test-model", 0, 0, 0, 0), @@ -987,8 +1002,63 @@ def _make_partial_tail_req(block_ids: list[int]) -> ReqMeta: block_ids=(block_ids, [1]), block_hashes=[b"a0", b"a1", b"a2"], can_save=True, - boundary_state_offloads=[(1, 7, 12)], + boundary_puts=[(1, 7, 12)], + num_prompt_tokens=13, + completed_token_len=13, + ) + + +@pytest.mark.parametrize("use_eagle", [False, True]) +def test_partial_tail_ignores_handoff_off_prompt_boundary(use_eagle): + """A Mamba hand-off away from the prompt checkpoint (a shared-prefix + junction) gets no attention tail; only its own Mamba put remains.""" + store = MagicMock() + thread = _make_partial_tail_send_thread(store) + thread.coord.use_eagle = use_eagle + thread.coord.eagle_proof_margin_by_group = {0: 4} if use_eagle else {} + req = _make_partial_tail_req([1, 2, 3]) + req.num_prompt_tokens = 17 if use_eagle else 13 + req.boundary_puts = [(1, 7, 8)] + + assert _resolve_partial_tail(thread, req).boundary_puts == [(1, 7, 8)] + + +def test_eagle_attention_proof_published_after_checkpoint_handoff(): + store = MagicMock() + stored = set() + store.batch_is_exist.side_effect = lambda keys: [int(k in stored) for k in keys] + + def put(keys, *args): + stored.update(keys) + return [256] * len(keys) + + store.batch_put_from_multi_buffers.side_effect = put + thread = _make_partial_tail_send_thread(store) + thread.coord.use_eagle = True + thread.coord.eagle_proof_margin_by_group = {0: 4} + metadata = _make_partial_tail_req([1, 2, 3]) + metadata.num_prompt_tokens = 13 + metadata.completed_token_len = 8 + metadata.boundary_puts = [(1, 7, 8)] + + assert thread._maybe_offload_boundary_states( + _resolve_partial_tail(thread, metadata) ) + mamba_key = thread.token_databases[1].key_for(b"a1") + attention_key = thread.token_databases[0].key_for(b"a2") + assert mamba_key in stored + assert attention_key not in stored + + metadata.completed_token_len = 13 + metadata.boundary_puts = None + metadata.publish_partial_tail = True + assert thread._maybe_offload_boundary_states( + _resolve_partial_tail(thread, metadata) + ) + assert attention_key in stored + keys, addrs, *_ = store.batch_put_from_multi_buffers.call_args.args + assert addrs[keys.index(attention_key)] == [0x1000 + 3 * 256] + assert mamba_key not in keys def test_partial_tail_offload_skips_null_source_blocks(): @@ -997,7 +1067,9 @@ def test_partial_tail_offload_skips_null_source_blocks(): store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) thread = _make_partial_tail_send_thread(store) - assert thread._maybe_offload_boundary_states(_make_partial_tail_req([0, 2, 3])) + assert thread._maybe_offload_boundary_states( + _resolve_partial_tail(thread, _make_partial_tail_req([0, 2, 3])) + ) keys, addrs, _sizes, _replicate_config = ( store.batch_put_from_multi_buffers.call_args.args @@ -1131,9 +1203,11 @@ def make_db(group_id: int, block_size: int, base_addr: int): block_hashes=[b"a0", b"a1", b"a2"], can_save=True, # The scheduler accepted group 1 and omitted group 2 at boundary 12. - boundary_state_offloads=[(1, 7, 12)], + boundary_puts=[(1, 7, 12)], + num_prompt_tokens=13, + completed_token_len=13, ) - assert thread._maybe_offload_boundary_states(req) + assert thread._maybe_offload_boundary_states(_resolve_partial_tail(thread, req)) keys, addrs, _sizes, _replicate_config = ( store.batch_put_from_multi_buffers.call_args.args @@ -1156,7 +1230,9 @@ def test_partial_tail_offload_replaces_stale_group_ids_after_filtering(): supports_group_ids=True, ) - assert thread._maybe_offload_boundary_states(_make_partial_tail_req([1, 2, 3])) + assert thread._maybe_offload_boundary_states( + _resolve_partial_tail(thread, _make_partial_tail_req([1, 2, 3])) + ) keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args assert keys == [ @@ -1190,13 +1266,17 @@ def test_partial_tail_put_failure_activates_pressure_gate(): store.batch_put_from_multi_buffers.return_value = [256, -200, 256] thread = _make_partial_tail_send_thread(store) - _run_store_req(thread, _make_partial_tail_req([1, 2, 3])) + _run_store_req( + thread, _resolve_partial_tail(thread, _make_partial_tail_req([1, 2, 3])) + ) assert thread._store_pressure_active is True assert thread._skip_store_requests == {"req-a"} assert thread._saved_offset.get("req-a", 0) == 0 - _run_store_req(thread, _make_partial_tail_req([1, 2, 3])) + _run_store_req( + thread, _resolve_partial_tail(thread, _make_partial_tail_req([1, 2, 3])) + ) assert store.batch_put_from_multi_buffers.call_count == 1 @@ -1277,9 +1357,9 @@ def test_block_aligned_snapshot_offload_uses_provided_block(): block_ids=([1, 2, 3], [5]), block_hashes=hs, can_save=True, - boundary_state_offloads=[(1, 7, 32)], + boundary_puts=[(1, 7, 32)], ) - assert thread._maybe_offload_boundary_states(req) + assert thread._maybe_offload_boundary_states(_resolve_partial_tail(thread, req)) keys, addrs, _sizes, _ = store.batch_put_from_multi_buffers.call_args.args # boundary 32 is block-aligned for the mamba group (block 16): one key, @@ -1289,7 +1369,9 @@ def test_block_aligned_snapshot_offload_uses_provided_block(): assert addrs == [[0x2000 + 7 * 256]] -def test_mixed_snapshot_and_sub_block_offloads(): +@pytest.mark.parametrize("saved_tokens", [0, 32]) +@pytest.mark.parametrize("use_eagle", [False, True]) +def test_mixed_snapshot_and_sub_block_offloads(saved_tokens, use_eagle): """A retention snapshot and the prompt-end sub-block CoW tail can arrive in one hand-off; the sub-block path covers FA gap blocks but reads the mamba boundary only from the CoW block, never positionally.""" @@ -1297,39 +1379,150 @@ def test_mixed_snapshot_and_sub_block_offloads(): store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) thread = _make_partial_tail_send_thread(store) + thread.coord.use_eagle = use_eagle + thread.coord.eagle_proof_margin_by_group = {0: 4} if use_eagle else {} + thread._saved_offset["req-a"] = saved_tokens - hs = [bytes([i + 1]) * 4 for i in range(11)] # 11 hash units = 44 tokens + hs = [bytes([i + 1]) * 4 for i in range(12)] req = ReqMeta( req_id="req-a", token_len_chunk=0, - block_ids=([1, 2, 3], [0, 0, 0]), + block_ids=(list(range(1, 14)), [0, 0, 0]), block_hashes=hs, can_save=True, - boundary_state_offloads=[(1, 9, 32), (1, 7, 44)], + boundary_puts=[(1, 9, 32), (1, 7, 44)], + num_prompt_tokens=49 if use_eagle else 45, + completed_token_len=49 if use_eagle else 45, ) - assert thread._maybe_offload_boundary_states(req) + assert thread._maybe_offload_boundary_states(_resolve_partial_tail(thread, req)) keys, addrs, _sizes, _ = store.batch_put_from_multi_buffers.call_args.args db_full, db_mamba = thread.token_databases + proof_end = 48 if use_eagle else 44 assert keys == [ + # The tail reads only the LCM gap, even when normal saves lag. + *[db_full.key_for(BlockHash(hs[i])) for i in range(8, proof_end // 4)], # Aligned snapshot: boundary 32 from the handed-off block 9. db_mamba.key_for(BlockHash(hs[7])), - # Sub-block boundary 44: FA gap blocks ending at 4, 8, 12. - db_full.key_for(BlockHash(hs[0])), - db_full.key_for(BlockHash(hs[1])), - db_full.key_for(BlockHash(hs[2])), # Mamba boundary block from the CoW hand-off (block 7). db_mamba.key_for(BlockHash(hs[10])), ] assert addrs == [ + *[[0x1000 + (i + 1) * 256] for i in range(8, proof_end // 4)], [0x2000 + 9 * 256], - [0x1000 + 1 * 256], - [0x1000 + 2 * 256], - [0x1000 + 3 * 256], [0x2000 + 7 * 256], ] +@pytest.mark.parametrize("tp_rank", [0, 1]) +def test_partial_tail_with_smaller_mamba_blocks_writes_one_mamba_key(tp_rank): + """Under DCP2 the attention block spans the LCM (16) and Mamba blocks are + half that. The prompt-completing save writes one attention block and one + Mamba key: the boundary state from the CoW block, never the interior + Mamba blocks.""" + store = MagicMock() + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) + store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) + coord = SimpleNamespace( + enable_partial_hash_hits=True, + hash_block_size=4, + lcm_block_size=16, + mamba_group_ids={1}, + use_eagle=False, + eagle_proof_margin_by_group={}, + ) + db_full = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0), block_size=16, hash_block_size=4 + ) + db_full.set_kv_caches_base_addr([0x1000]) + db_full.set_block_len([256]) + db_mamba = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=1), + block_size=8, + hash_block_size=4, + ) + db_mamba.set_kv_caches_base_addr([0x2000]) + db_mamba.set_block_len([256]) + thread = _make_store_sending_thread( + store, coord=coord, token_databases=[db_full, db_mamba], tp_rank=tp_rank + ) + # Mamba is replicated across two ranks, which split its keys by block. + thread.group_put_steps = [1, 2] + + hs = [bytes([i + 1]) * 4 for i in range(11)] + req = ReqMeta( + req_id="req-a", + token_len_chunk=0, + block_ids=([1, 2, 3], list(range(20, 26))), + block_hashes=hs, + can_save=True, + boundary_puts=[(1, 7, 44)], + num_prompt_tokens=45, + completed_token_len=45, + publish_partial_tail=True, + ) + assert thread._maybe_offload_boundary_states(_resolve_partial_tail(thread, req)) + assert req.boundary_puts == [(0, 3, 44), (1, 7, 44)] + + keys, addrs, _sizes, _ = store.batch_put_from_multi_buffers.call_args.args + # Mamba block cdiv(44, 8) - 1 = 5 belongs to put_step_rank 5 % 2 = 1, + # which is tp_rank 0 for group 1. + mamba_puts = [db_mamba.key_for(BlockHash(hs[10]))] if tp_rank == 0 else [] + assert keys == [db_full.key_for(BlockHash(hs[10])), *mamba_puts] + assert addrs == [[0x1000 + 3 * 256], *([[0x2000 + 7 * 256]] * len(mamba_puts))] + + +def test_shared_prefix_junction_handoff_writes_only_its_mamba_key(): + """A junction hand-off (``--enable-mamba-shared-prefix-checkpoint``) below + the prompt checkpoint stores its Mamba state under the junction hash and + no attention tail.""" + store = MagicMock() + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) + store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) + coord = SimpleNamespace( + enable_partial_hash_hits=True, + hash_block_size=4, + lcm_block_size=16, + mamba_group_ids={1}, + use_eagle=True, + eagle_proof_margin_by_group={0: 4}, + ) + db_full = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0), block_size=4, hash_block_size=4 + ) + db_full.set_kv_caches_base_addr([0x1000]) + db_full.set_block_len([256]) + db_mamba = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=1), + block_size=16, + hash_block_size=4, + ) + db_mamba.set_kv_caches_base_addr([0x2000]) + db_mamba.set_block_len([256]) + thread = _make_store_sending_thread( + store, coord=coord, token_databases=[db_full, db_mamba] + ) + + hs = [bytes([i + 1]) * 4 for i in range(18)] + # 75-token prompt (checkpoint 68), chunk cut at the junction 40. + req = ReqMeta( + req_id="req-b", + token_len_chunk=0, + block_ids=(list(range(1, 11)), [20, 21, 22]), + block_hashes=hs, + can_save=True, + boundary_puts=[(1, 7, 40)], + num_prompt_tokens=75, + completed_token_len=40, + ) + assert thread._maybe_offload_boundary_states(_resolve_partial_tail(thread, req)) + assert req.boundary_puts == [(1, 7, 40)] + + keys, addrs, _sizes, _ = store.batch_put_from_multi_buffers.call_args.args + assert keys == [db_mamba.key_for(BlockHash(hs[9]))] + assert addrs == [[0x2000 + 7 * 256]] + + def test_snapshot_offload_skips_null_handoff_block(): """A hand-off the core could not materialize (null block) carries no committed state; persisting it would poison the boundary key.""" @@ -1338,13 +1531,16 @@ def test_snapshot_offload_skips_null_handoff_block(): hs = [bytes([i + 1]) * 4 for i in range(8)] assert thread._maybe_offload_boundary_states( - ReqMeta( - req_id="req-a", - token_len_chunk=0, - block_ids=([1, 2, 3], [5]), - block_hashes=hs, - can_save=True, - boundary_state_offloads=[(1, NULL_BLOCK_ID, 32)], + _resolve_partial_tail( + thread, + ReqMeta( + req_id="req-a", + token_len_chunk=0, + block_ids=([1, 2, 3], [5]), + block_hashes=hs, + can_save=True, + boundary_puts=[(1, NULL_BLOCK_ID, 32)], + ), ) ) store.batch_is_exist.assert_not_called() @@ -1398,6 +1594,7 @@ def test_store_sending_thread_delta_saves_only_new_masked_chunks(): [True, False], ), ) + coord.store_mask = MagicMock(side_effect=coord.store_mask) db_full = ChunkedTokenDatabase( KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), @@ -1427,6 +1624,8 @@ def test_store_sending_thread_delta_saves_only_new_masked_chunks(): block_ids=([0, 1, 2, 3], [0, 1, 2, 3]), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, + num_prompt_tokens=33, + prefill_end_tokens=64, ), ) @@ -1436,6 +1635,7 @@ def test_store_sending_thread_delta_saves_only_new_masked_chunks(): assert full_hashes == [b"a2".hex(), b"a3".hex()] assert masked_hashes == [b"a2".hex()] + coord.store_mask.assert_called_once_with(64, 32, num_prompt_tokens=64) def test_store_sending_thread_prepares_missing_chunks_once_per_group(): @@ -3580,6 +3780,139 @@ def _make_bare_worker( return worker +@pytest.mark.parametrize("use_eagle", [False, True]) +@pytest.mark.parametrize("prefix_match_unit", [None, 8]) +@pytest.mark.parametrize("dcp_world_size,mamba_block_size", [(1, 32), (8, 16)]) +def test_mooncake_lookup_reuses_resolved_hash_checkpoint( + use_eagle, prefix_match_unit, dcp_world_size, mamba_block_size +): + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import ( + coordinator as mooncake_coordinator, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import ( + MooncakeStoreScheduler, + ) + + groups = [ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + is_eagle_group=use_eagle, + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=mamba_block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ] + config = SimpleNamespace( + speculative_config=None, + kv_transfer_config=SimpleNamespace( + kv_role="kv_both", kv_connector_extra_config={} + ), + kv_events_config=None, + cache_config=SimpleNamespace( + block_size=16, + enable_prefix_caching=True, + prefix_match_unit=prefix_match_unit, + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=dcp_world_size, world_size=dcp_world_size + ), + ) + with patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "scheduler.LookupKeyClient" + ): + scheduler = MooncakeStoreScheduler( + config, + KVCacheConfig( + num_blocks=128, + kv_cache_tensors=[], + kv_cache_groups=groups, + # The hit alignment the engine core resolves for this config. + cache_hit_alignment_tokens=prefix_match_unit or 16, + ), + ) + assert scheduler.enable_partial_hash_hits + hash_block_size = prefix_match_unit or 16 + assert scheduler._hash_block_size == hash_block_size + coordinator = mooncake_coordinator.MooncakeStoreCoordinator( + groups, + scheduler_block_size=math.lcm(16 * dcp_world_size, mamba_block_size), + hash_block_size=hash_block_size, + enable_partial_hash_hits=True, + ) + hashes = [BlockHash(bytes([i]) * 16) for i in range(512 // hash_block_size)] + pool = mooncake_coordinator.ExternalCachedBlockPool( + hash_block_size=hash_block_size, + exists={(group, bytes(h)) for group in range(2) for h in hashes}, + ) + _, hit = coordinator.find_longest_cache_hit(hashes, 511, pool) + assert hit == 512 - hash_block_size + + +@pytest.mark.parametrize("alignment, expected", [(8, True), (32, False), (None, False)]) +def test_scheduler_follows_core_hit_alignment(alignment, expected): + """The store takes partial hits from core's resolved alignment, not from its + own groups: these Mamba specs alone would allow them, but core may still + keep hits on the block (e.g. for a sliding-window group it never sees).""" + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import ( + MooncakeStoreScheduler, + ) + + groups = [ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=32, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=32, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ] + config = SimpleNamespace( + speculative_config=None, + kv_transfer_config=SimpleNamespace( + kv_role="kv_both", kv_connector_extra_config={} + ), + kv_events_config=None, + cache_config=SimpleNamespace( + block_size=32, enable_prefix_caching=True, prefix_match_unit=8 + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1, world_size=1), + ) + with patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "scheduler.LookupKeyClient" + ): + scheduler = MooncakeStoreScheduler( + config, + KVCacheConfig( + num_blocks=128, + kv_cache_tensors=[], + kv_cache_groups=groups, + cache_hit_alignment_tokens=alignment, + ), + ) + + assert scheduler.enable_partial_hash_hits is expected + assert scheduler._store_coord.enable_partial_hash_hits is expected + + def test_lookup_key_prefixes_cover_dcp_rank_namespaces(): worker = _make_bare_worker() worker.tp_size = 4 @@ -3873,6 +4206,7 @@ def test_lookup_partial_tail_uses_hash_alignment(): worker._kv_cache_groups, scheduler_block_size=16, hash_block_size=4, + enable_partial_hash_hits=True, ) _refresh_group_tp_replication_factors(worker) worker.store.batch_is_exist.return_value = [0, 0, 1, 0, 0, 1] @@ -3946,6 +4280,7 @@ def test_lookup_plan_resolves_group_tail_keys_from_existing_hashes(): scheduler_block_size=16, hash_block_size=4, use_eagle=True, + enable_partial_hash_hits=True, ) _refresh_group_tp_replication_factors(worker) hashes = [BlockHash(f"h{i}".encode()) for i in range(6)] @@ -4015,6 +4350,7 @@ def test_lookup_plan_recovers_tail_key_after_multi_chunk_convergence(): worker._kv_cache_groups, scheduler_block_size=16, hash_block_size=4, + enable_partial_hash_hits=True, ) _refresh_group_tp_replication_factors(worker) hashes = [BlockHash(f"h{i}".encode()) for i in range(12)] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index d05b553b64b9..48a36bc82f8b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -3,7 +3,8 @@ """External-store cache-hit coordinator for MooncakeStoreConnector.""" from collections.abc import Sequence -from typing import NamedTuple, cast +from dataclasses import replace +from typing import TYPE_CHECKING, NamedTuple, cast from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( chunk_hashes_for_block_size, @@ -13,12 +14,15 @@ from vllm.v1.core.kv_cache_utils import ( BlockHash, KVCacheBlock, + eagle_proof_margin, + resolve_dcp_kv_cache_spec, ) from vllm.v1.core.single_type_kv_cache_manager import ( SingleTypeKVCacheManager, ) from vllm.v1.kv_cache_interface import ( FullAttentionSpec, + KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, MambaSpec, @@ -26,6 +30,9 @@ ) from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry +if TYPE_CHECKING: + from vllm.config import VllmConfig + class StoreSpecGroup(NamedTuple): spec: KVCacheSpec @@ -82,7 +89,7 @@ def __init__( hash_block_size: int, use_eagle: bool = False, retention_interval: int | None = None, - dcp_world_size: int = 1, + enable_partial_hash_hits: bool = False, ) -> None: # Mirrors core's resolve_kv_cache_block_sizes: the hash unit only has # to divide groups that participate in prefix caching. Non-shareable @@ -110,14 +117,51 @@ def __init__( } self.hash_block_size = hash_block_size self.lcm_block_size = scheduler_block_size - self.enable_partial_hash_hits = partial_hash_hits_enabled( - kv_cache_groups, hash_block_size, dcp_world_size - ) + # Core's decision: whether hits land on the hash unit, below the block. + self.enable_partial_hash_hits = enable_partial_hash_hits self.use_eagle = use_eagle # Mirror vLLM core's KVCacheCoordinator.retention_interval. self.retention_interval = retention_interval self._verify_and_split_kv_cache_groups() + @classmethod + def from_kv_cache_config( + cls, + kv_cache_config: KVCacheConfig, + vllm_config: "VllmConfig", + scheduler_block_size: int, + hash_block_size: int, + ) -> "MooncakeStoreCoordinator": + """Build the coordinator the connector's scheduler and workers share. + + Specs are DCP-resolved, so a block spans the tokens of one block-table + entry. Partial hits follow core's hit alignment, which checks every KV + cache group and its manager, not only the groups the store transfers; + unset (a config the engine core did not resolve) keeps hits on the + block. + """ + dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size + spec_config = vllm_config.speculative_config + alignment = kv_cache_config.cache_hit_alignment_tokens + return cls( + [ + replace( + group, + kv_cache_spec=resolve_dcp_kv_cache_spec( + group.kv_cache_spec, dcp_world_size + ), + ) + for group in kv_cache_config.prefix_cacheable_groups + ], + scheduler_block_size, + hash_block_size, + use_eagle=spec_config is not None and spec_config.use_eagle_block_drop(), + retention_interval=kv_cache_config.prefix_cache_retention_interval, + enable_partial_hash_hits=( + alignment is not None and alignment < scheduler_block_size + ), + ) + def align_lookup_length(self, length: int) -> int: alignment = ( self.hash_block_size @@ -167,6 +211,17 @@ def _verify_and_split_kv_cache_groups(self) -> None: self.eagle_group_ids = { gid for g in attention_groups if g.use_eagle for gid in g.group_ids } + self.eagle_proof_margin_by_group = { + gid: eagle_proof_margin( + group.spec.block_size, + self.hash_block_size, + self.enable_partial_hash_hits + and group.manager_cls.supports_fine_grained_hash_lookup, + ) + for group in attention_groups + if group.use_eagle and not isinstance(group.spec, MambaSpec) + for gid in group.group_ids + } def find_longest_cache_hit( self, @@ -394,14 +449,16 @@ def _find_hit_blocks( # never drops a block, so a widened bound would match past the # attention-verified hit and resume from speculative state (#43559). if drop_eagle_block and not isinstance(spec, MambaSpec): - eagle_margin = ( - self.hash_block_size - if self.enable_partial_hash_hits - and manager_cls.supports_fine_grained_hash_lookup - and spec.block_size > self.hash_block_size - else spec.block_size + eagle_margin = eagle_proof_margin( + spec.block_size, + self.hash_block_size, + self.enable_partial_hash_hits + and manager_cls.supports_fine_grained_hash_lookup, + ) + _max_length = min( + curr_hit_length + eagle_margin, + len(block_hashes) * self.hash_block_size, ) - _max_length = min(curr_hit_length + eagle_margin, max_length) hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit( block_hashes=block_hashes, # type: ignore[arg-type] max_length=_max_length, @@ -449,20 +506,3 @@ def _unwrap_spec(spec: KVCacheSpec) -> KVCacheSpec: if isinstance(spec, UniformTypeKVCacheSpecs): return next(iter(spec.kv_cache_specs.values())) return spec - - -def partial_hash_hits_enabled( - kv_cache_groups: Sequence[KVCacheGroupSpec], - hash_block_size: int, - dcp_world_size: int = 1, -) -> bool: - """Match core's DCP-aware Mamba partial-hit condition.""" - return any( - isinstance(spec := _unwrap_spec(g.kv_cache_spec), MambaSpec) - and spec.mamba_cache_mode == "align" - and ( - (dcp_world_size == 1 and spec.block_size > hash_block_size) - or (dcp_world_size > 1 and spec.block_size >= hash_block_size) - ) - for g in kv_cache_groups - ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py index 2510a7ee8e76..4cd0402b7b4f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py @@ -7,7 +7,7 @@ from collections.abc import Iterable, Sequence from dataclasses import dataclass, field -from typing import cast +from typing import NamedTuple, cast import numpy as np import torch @@ -712,6 +712,7 @@ class RequestTracker: # For a fresh request this is len(prompt). For a resumed-from-preemption # request it includes previously-generated tokens, which are re-prefilled. prefill_end_tokens: int = 0 + partial_tail_sent: bool = False def reset(self) -> None: self.token_len = 0 @@ -720,6 +721,7 @@ def reset(self) -> None: self.token_ids = None self.has_pending_offload = False self.prefill_end_tokens = 0 + self.partial_tail_sent = False def update( self, @@ -739,6 +741,14 @@ def update( existing.extend(new) +class BoundaryPut(NamedTuple): + """A block stored under the hash of the prefix ending at ``num_tokens``.""" + + group_id: int + block_id: int + num_tokens: int + + @dataclass class ReqMeta: """Per-request metadata for store put/get operations.""" @@ -756,15 +766,24 @@ class ReqMeta: # Absolute request offset represented by token_ids[0]. token_ids_start: int = 0 num_prompt_tokens: int | None = None + # Total prefill tokens, including generated tokens replayed after preemption. + prefill_end_tokens: int | None = None # Identifies this store job for the engine's lifetime. A request id cannot # serve that purpose: it is reused once a preempted request resumes, so it # would release the wrong job's blocks. store_job_id: int | None = None - # Core-provided (group_id, block_id, boundary_tokens) mamba "align" - # boundary states. A block-aligned entry is a committed boundary snapshot; - # a non-aligned entry is the sub-block CoW tail. The store-job reference - # keeps each exact block alive until every worker rank finishes the job. - boundary_state_offloads: list[tuple[int, int, int]] | None = None + # Blocks stored under an explicit prefix hash rather than by position: + # core-provided mamba "align" states (a block-aligned entry is a committed + # snapshot, a non-aligned one the sub-block CoW prompt tail) and the + # scheduler-resolved non-Mamba blocks of the prompt's partial tail. The + # store-job reference keeps each block alive until every worker rank + # finishes the job. + boundary_puts: list[BoundaryPut] | None = None + # Total computed prefix length at the end of this step. + completed_token_len: int | None = None + # Set on the save that first covers the whole prompt, so the scheduler + # publishes the prompt's partial tail once rather than on every save. + publish_partial_tail: bool = False @staticmethod def from_request_tracker( @@ -773,8 +792,12 @@ def from_request_tracker( load_spec: LoadSpec | None = None, skip_save: bool | None = False, block_hashes: list[BlockHash] | None = None, + save_partial_tail: bool = False, + num_prompt_tokens: int | None = None, ) -> "ReqMeta | None": """Create ReqMeta from a RequestTracker.""" + if save_partial_tail: + assert num_prompt_tokens is not None if block_hashes is None: block_hashes = [] input_token_len = tracker.token_len @@ -783,7 +806,15 @@ def from_request_tracker( chunk_boundary = cdiv(token_ids_start + 1, block_size) * block_size num_tokens_to_save = input_token_len // block_size * block_size - skip_save = skip_save or num_tokens_to_save < chunk_boundary + publish_tail = ( + save_partial_tail + and not tracker.partial_tail_sent + and num_prompt_tokens is not None + and 0 < num_prompt_tokens <= input_token_len + ) + skip_save = skip_save or ( + num_tokens_to_save < chunk_boundary and not publish_tail + ) # A ReqMeta must never carry both a save AND a load. # The save would also be wasted work — the bytes are being looked up # in the store right now. Later cached_reqs steps save new tokens @@ -795,6 +826,7 @@ def from_request_tracker( if not skip_save: tracker.num_saved_tokens = num_tokens_to_save + tracker.partial_tail_sent |= publish_tail token_ids = None if tracker.token_ids and not skip_save: @@ -826,7 +858,10 @@ def from_request_tracker( block_hashes=block_hashes, token_ids=token_ids, token_ids_start=token_ids_start, - num_prompt_tokens=tracker.prefill_end_tokens, + num_prompt_tokens=num_prompt_tokens, + prefill_end_tokens=tracker.prefill_end_tokens, + completed_token_len=input_token_len, + publish_partial_tail=publish_tail and not skip_save, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py index 8f277d7eea37..239dad736490 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py @@ -5,14 +5,17 @@ # (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/). """Scheduler-side logic for MooncakeStoreConnector.""" +from collections.abc import Sequence + from vllm.config import VllmConfig from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorMetadata, ) -from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501 - partial_hash_hits_enabled, +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( + MooncakeStoreCoordinator, ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501 + BoundaryPut, LoadSpec, MooncakeStoreConnectorMetadata, MooncakeStoreWorkerMetadata, @@ -23,12 +26,19 @@ LookupKeyClient, ) from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_manager import KVCacheBlocks -from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes +from vllm.v1.core.kv_cache_utils import ( + resolve_kv_cache_block_sizes, +) from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput -from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + MambaSpec, + get_mamba_prefill_checkpoint_position, +) from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import Request @@ -49,6 +59,79 @@ def _new_req_prefill_tokens(request: NewRequestData) -> list[int]: return request.prompt_token_ids +def _partial_tail_non_mamba_puts( + coord: MooncakeStoreCoordinator, + req_meta: ReqMeta, + block_sizes: Sequence[int], +) -> list[BoundaryPut]: + """Locate the non-Mamba blocks a partial-tail save publishes. + + A later request resumes at the prompt's Mamba checkpoint (``boundary``) only + if both pieces of KV are stored: + + - Mamba groups: the state at ``boundary``, handed off by the core as a + boundary put of its own, so Mamba groups are not listed here. + - Other groups: the KV from the last LCM-aligned normal save up to + ``boundary`` plus the group's EAGLE proof margin (``proof_end``). Each + block is keyed by the hash at its end, clipped to ``proof_end``. + + Groups whose proof is not computed yet are omitted. For example, with LCM + 16, blocks of 4 tokens, and a checkpoint at 44, the full-attention blocks + are 8-10, keyed at 36, 40 and 44. + + A Mamba state handed off at any other position, i.e. a shared-prefix + junction (``--enable-mamba-shared-prefix-checkpoint``), adds no puts here. + + ``block_sizes`` are the per-group block sizes, indexed like + ``req_meta.block_ids``. + """ + mamba_tails = [ + put.num_tokens + for put in req_meta.boundary_puts or [] + if put.group_id in coord.mamba_group_ids + and put.num_tokens % block_sizes[put.group_id] + ] + # The tail is due on the save that first covers the prompt, or when Mamba + # hands off its checkpoint state. + if not mamba_tails and not req_meta.publish_partial_tail: + return [] + prompt_tokens = req_meta.num_prompt_tokens + completed = req_meta.completed_token_len + assert prompt_tokens is not None and completed is not None + hash_block_size = coord.hash_block_size + # Core drops the Mamba checkpoint whenever EAGLE block drop is on, whichever + # groups carry the eagle flag. + boundary = get_mamba_prefill_checkpoint_position( + prompt_tokens, hash_block_size, coord.use_eagle + ) + # Any other sub-block Mamba hand-off is a shared-prefix junction. Its + # attention KV, EAGLE proof included, ends on an attention block end no later + # than this checkpoint's proof, so the normal and prompt-tail saves store it. + if not req_meta.publish_partial_tail and boundary not in mamba_tails: + return [] + num_hashes = len(req_meta.block_hashes) + if boundary == 0 or boundary // hash_block_size > num_hashes: + return [] + start = boundary // coord.lcm_block_size * coord.lcm_block_size + puts: list[BoundaryPut] = [] + for group_id, block_size in enumerate(block_sizes): + if group_id in coord.mamba_group_ids: + continue + proof_end = boundary + coord.eagle_proof_margin_by_group.get(group_id, 0) + if proof_end > completed or proof_end // hash_block_size > num_hashes: + continue + group_blocks = req_meta.block_ids[group_id] + for block_idx in range(start // block_size, cdiv(proof_end, block_size)): + if block_idx >= len(group_blocks): + break + block_id = group_blocks[block_idx] + if block_id == NULL_BLOCK_ID: + continue + num_tokens = min((block_idx + 1) * block_size, proof_end) + puts.append(BoundaryPut(group_id, block_id, num_tokens)) + return puts + + class MooncakeStoreScheduler: """Scheduler-side component for MooncakeStoreConnector.""" @@ -84,11 +167,6 @@ def __init__( self._block_size, self._hash_block_size = resolve_kv_cache_block_sizes( kv_cache_config, vllm_config ) - self.enable_partial_hash_hits = partial_hash_hits_enabled( - store_groups, - self._hash_block_size, - vllm_config.parallel_config.decode_context_parallel_size, - ) mamba_groups = { group_id: group.kv_cache_spec for group_id, group in enumerate(store_groups) @@ -99,6 +177,11 @@ def __init__( ), "MooncakeStoreScheduler requires mamba_cache_mode='align'" self._boundary_state_group_ids = frozenset(mamba_groups) + self._store_coord = MooncakeStoreCoordinator.from_kv_cache_config( + kv_cache_config, vllm_config, self._block_size, self._hash_block_size + ) + self.enable_partial_hash_hits = self._store_coord.enable_partial_hash_hits + self._gpu_block_pool: BlockPool | None = None self._num_workers = vllm_config.parallel_config.world_size self._next_store_job_id = 0 @@ -273,6 +356,8 @@ def build_connector_meta( # producer. Loads are still carried by the same metadata. skip_save=is_consumer, block_hashes=request_real.block_hashes, + num_prompt_tokens=request_real.num_prompt_tokens, + save_partial_tail=self.enable_partial_hash_hits, ) if req_meta is not None: meta.add_request(req_meta) @@ -323,6 +408,8 @@ def build_connector_meta( load_spec=load_spec, skip_save=is_consumer, block_hashes=request_real.block_hashes, + num_prompt_tokens=request_real.num_prompt_tokens, + save_partial_tail=self.enable_partial_hash_hits, ) else: # Decode/chunked request @@ -364,6 +451,8 @@ def build_connector_meta( load_spec=None, skip_save=False, block_hashes=unfinished_req.block_hashes, + num_prompt_tokens=unfinished_req.num_prompt_tokens, + save_partial_tail=self.enable_partial_hash_hits, ) if req_meta is not None: @@ -398,6 +487,7 @@ def build_connector_meta( load_spec=load_spec, skip_save=None, block_hashes=unfinished_req.block_hashes, + num_prompt_tokens=unfinished_req.num_prompt_tokens, ) if req_meta is not None: meta.add_request(req_meta) @@ -455,6 +545,8 @@ def _reference_save_blocks(self, meta: MooncakeStoreConnectorMetadata) -> None: itself is freed, until every rank reports the job done. """ pool = self._gpu_block_pool + coord = self._store_coord + block_sizes = [g.kv_cache_spec.block_size for g in coord.kv_cache_groups] for req_meta in meta.requests: if not req_meta.can_save: continue @@ -464,27 +556,27 @@ def _reference_save_blocks(self, meta: MooncakeStoreConnectorMetadata) -> None: if req_meta.store_job_id is not None: assert req_meta.store_job_id in self._pinned_saves continue + if tail := _partial_tail_non_mamba_puts(coord, req_meta, block_sizes): + req_meta.boundary_puts = [*tail, *(req_meta.boundary_puts or [])] req_meta.store_job_id = store_job_id = self._next_store_job_id self._next_store_job_id += 1 - block_ids: list[int] = [] - if req_meta.boundary_state_offloads: - block_ids.extend( - block_id for _, block_id, _ in req_meta.boundary_state_offloads - ) + block_ids = [put.block_id for put in req_meta.boundary_puts or []] assert NULL_BLOCK_ID not in block_ids, ( - "A null block cannot back a boundary-state offload" - ) - # Every allocated block is referenced, not just the ones covering - # this job's token range: a rank resumes from its own last - # successful offset, which lags the scheduler's whenever a save was - # skipped or failed, so it may read anywhere below the range. - block_ids.extend( - block_id - for group_id, group in enumerate(req_meta.block_ids) - if group_id not in self._boundary_state_group_ids - for block_id in group - if block_id != NULL_BLOCK_ID + "A null block cannot back a boundary put" ) + if req_meta.token_len_chunk: + # Normal prefix save: pin all non-Mamba sources for retries. + # Every allocated block is referenced, not just the ones covering + # this job's token range: a rank resumes from its own last + # successful offset, which lags the scheduler's whenever a save was + # skipped or failed, so it may read anywhere below the range. + block_ids.extend( + block_id + for group_id, group in enumerate(req_meta.block_ids) + if group_id not in self._boundary_state_group_ids + for block_id in group + if block_id != NULL_BLOCK_ID + ) # An aligned boundary block may also be present in the request's # block table. Take and release exactly one reference per block. block_ids = list(dict.fromkeys(block_ids)) @@ -515,17 +607,39 @@ def register_finished_partial_tail( if boundary_tokens > tracker.prefill_end_tokens: return False - pinned_block_ids: list[int] = [] - remapped_offloads: list[tuple[int, int, int]] = [] + mamba_puts: list[BoundaryPut] = [] for group_id, block_id, boundary in partial_tail_offloads: store_group_id = self._store_group_id_by_kv_cache_group_id.get(group_id) if store_group_id not in self._boundary_state_group_ids: return False if block_id == NULL_BLOCK_ID: return False - pinned_block_ids.append(block_id) - remapped_offloads.append((store_group_id, block_id, boundary)) - pinned_block_ids = list(dict.fromkeys(pinned_block_ids)) + mamba_puts.append(BoundaryPut(store_group_id, block_id, boundary)) + req_meta = ReqMeta( + req_id=request.request_id, + token_len_chunk=0, + block_ids=tuple( + block_ids[group_id].copy() for group_id in self._store_group_ids + ), + block_hashes=list(request.block_hashes), + can_save=True, + num_prompt_tokens=request.num_prompt_tokens, + prefill_end_tokens=tracker.prefill_end_tokens, + boundary_puts=mamba_puts, + completed_token_len=request.num_computed_tokens, + ) + coord = self._store_coord + req_meta.boundary_puts = [ + *_partial_tail_non_mamba_puts( + coord, + req_meta, + [g.kv_cache_spec.block_size for g in coord.kv_cache_groups], + ), + *mamba_puts, + ] + pinned_block_ids = list( + dict.fromkeys(put.block_id for put in req_meta.boundary_puts) + ) pool = self._gpu_block_pool assert pool is not None, ( @@ -537,18 +651,8 @@ def register_finished_partial_tail( self._pinned_saves[store_job_id] = (pinned_block_ids, self._num_workers) pool.touch([pool.blocks[block_id] for block_id in pinned_block_ids]) - self._finished_partial_tail_metas[request.request_id] = ReqMeta( - req_id=request.request_id, - token_len_chunk=0, - block_ids=tuple( - block_ids[group_id].copy() for group_id in self._store_group_ids - ), - block_hashes=list(request.block_hashes), - can_save=True, - num_prompt_tokens=tracker.prefill_end_tokens, - store_job_id=store_job_id, - boundary_state_offloads=remapped_offloads, - ) + req_meta.store_job_id = store_job_id + self._finished_partial_tail_metas[request.request_id] = req_meta tracker.has_pending_offload = True # The store job owns exact block refs, so request cleanup need not wait. return False @@ -575,7 +679,7 @@ def _handle_boundary_state_offloads( # going away, so the offload is conservatively dropped. logger.debug("Dropping boundary-state offload for request %s", req_id) continue - accepted: list[tuple[int, int, int]] = [] + accepted: list[BoundaryPut] = [] for group_id, block_id, boundary_tokens in entries: # Every other group stops saving at the end of this prefill, so # a mamba-only key past it can never complete a joint hybrid @@ -589,12 +693,12 @@ def _handle_boundary_state_offloads( store_group_id = self._store_group_id_by_kv_cache_group_id.get(group_id) if store_group_id not in self._boundary_state_group_ids: continue - accepted.append((store_group_id, block_id, boundary_tokens)) + accepted.append(BoundaryPut(store_group_id, block_id, boundary_tokens)) if not accepted: continue tracker.has_pending_offload = True if (req_meta := save_metas.get(req_id)) is not None: - req_meta.boundary_state_offloads = accepted + req_meta.boundary_puts = accepted continue meta.add_request( ReqMeta( @@ -603,8 +707,10 @@ def _handle_boundary_state_offloads( block_ids=tracker.allocated_block_ids, block_hashes=req_tuple[0].block_hashes, can_save=True, - num_prompt_tokens=tracker.prefill_end_tokens, - boundary_state_offloads=accepted, + num_prompt_tokens=req_tuple[0].num_prompt_tokens, + prefill_end_tokens=tracker.prefill_end_tokens, + boundary_puts=accepted, + completed_token_len=tracker.token_len, ) ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 173aa324e725..004b69286ea3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -74,7 +74,6 @@ from vllm.v1.core.kv_cache_utils import ( BlockHash, maybe_convert_block_hash, - resolve_dcp_kv_cache_spec, resolve_kv_cache_block_sizes, ) from vllm.v1.kv_cache_interface import ( @@ -569,40 +568,43 @@ def _clear_store_pressure(self) -> bool: self._skip_store_requests.clear() return True - def _boundary_snapshot_puts( - self, req_meta: ReqMeta, entries: list[tuple[int, int, int]] + def _boundary_puts( + self, req_meta: ReqMeta ) -> list[tuple[str, list[int], list[int], KeyMetadata]]: - """Puts for committed mamba "align" boundary-state snapshots. - - These are block-aligned boundaries, i.e. exactly what the normal save - would key — but ``store_mask`` masks mamba groups out of it entirely, so - this is their *only* writer. The exclusion is not an optimization: the - normal save resolves a chunk's address as - ``req_meta.block_ids[g][start // block_size]``, and ``block_ids`` is the - connector's append-only mirror of the core's per-group table. An - align-mode table is mutated in place (a superseded state block is freed - and nulled; speculative blocks relocate), and the connector is never - told, so a stale mirror entry is indistinguishable from a live one — a - retry of a failed or pressure-skipped chunk would read a block that now - belongs to another request. - - Each entry's handed-off block *is* the boundary state and is pinned by - the core, so it is uploaded under its boundary-end hash key and never - resolved positionally. + """Puts for ``req_meta.boundary_puts``: each block under the hash of + the prefix ending at its ``num_tokens``. + + The scheduler resolves and pins every block, so none is looked up + positionally here. That matters for mamba "align" states: ``store_mask`` + masks mamba groups out of the normal save, which resolves a chunk's + address as ``req_meta.block_ids[g][start // block_size]``. An align-mode + table is mutated in place (a superseded state block is freed and + nulled; speculative blocks relocate) without telling the connector, so + a retried positional read could hit a block that now belongs to another + request. The handed-off block *is* the state. """ hash_block_size = self.coord.hash_block_size puts: list[tuple[str, list[int], list[int], KeyMetadata]] = [] - for group_id, block_id, boundary in entries: - if boundary == 0 or block_id == NULL_BLOCK_ID: + for group_id, block_id, num_tokens in req_meta.boundary_puts or []: + if ( + not self.group_participates[group_id] + or num_tokens == 0 + or block_id == NULL_BLOCK_ID + ): continue - hash_idx = boundary // hash_block_size - 1 + # A negative index would silently key the block by the last hash. + assert num_tokens % hash_block_size == 0, ( + f"Boundary put at {num_tokens} tokens is not a multiple of the " + f"hash block size {hash_block_size}" + ) + hash_idx = num_tokens // hash_block_size - 1 if hash_idx >= len(req_meta.block_hashes): continue db = self.token_databases[group_id] # Distribute across ranks by the same rule as normal chunks. put_step = self.group_put_steps[group_id] put_step_rank = (self.tp_rank + group_id) % put_step - if (boundary // db.block_size - 1) % put_step != put_step_rank: + if (cdiv(num_tokens, db.block_size) - 1) % put_step != put_step_rank: continue addr, size = db.prepare_value_for_block(block_id) puts.append( @@ -610,115 +612,20 @@ def _boundary_snapshot_puts( ) return puts - def _sub_block_tail_puts( - self, req_meta: ReqMeta, entries: list[tuple[int, int, int]] - ) -> list[tuple[str, list[int], list[int], KeyMetadata]]: - """Puts for the request's sub-block partial tail (its last prompt hash - boundary), so a later request can hit the sub-block prefix. - - Covers every group's blocks from the normal save's lcm floor to the - boundary: the normal save floors to ``lcm_block_size``, so a - smaller-block group's full blocks in that gap are never persisted - elsewhere, and the consumer's lookup needs every group at every probed - boundary. Full blocks are keyed by their block-end hash and the partial - boundary block by the boundary sub-hash; a mamba "align" group - contributes only its boundary block, from the core-provided CoW block. - """ - boundaries = {boundary for _, _, boundary in entries} - if len(boundaries) != 1: - raise ValueError( - "Sub-block partial-tail offloads for one request must share a boundary" - ) - boundary = boundaries.pop() - hash_block_size = self.coord.hash_block_size - if boundary == 0 or boundary // hash_block_size - 1 >= len( - req_meta.block_hashes - ): - return [] - - mamba_offloads = {group_id: block_id for group_id, block_id, _ in entries} - saved = self._saved_offset.get(req_meta.req_id, 0) - puts: list[tuple[str, list[int], list[int], KeyMetadata]] = [] - for g_idx, db in enumerate(self.token_databases): - if not self.group_participates[g_idx]: - continue - group_blocks = req_meta.block_ids[g_idx] - # Distribute across ranks by the same rule as normal chunks. - put_step = self.group_put_steps[g_idx] - put_step_rank = (self.tp_rank + g_idx) % put_step - # Always include the boundary block: its sub-hash key is written - # only here, even if normal saves already advanced past it. - last_block = cdiv(boundary, db.block_size) - 1 - for block_idx in range( - min(saved // db.block_size, last_block), last_block + 1 - ): - if block_idx % put_step != put_step_rank: - continue - valid_end = min((block_idx + 1) * db.block_size, boundary) - key_hash = req_meta.block_hashes[valid_end // hash_block_size - 1] - if g_idx in mamba_offloads: - if valid_end != boundary: - # Interior align-mode state positions are null or - # stale (the block table is not append-only) and never - # valid gap content; only the boundary block is - # persisted, from the core-provided hand-off. - continue - block_id = mamba_offloads[g_idx] - elif g_idx in self.coord.mamba_group_ids: - continue - elif block_idx < len(group_blocks): - block_id = group_blocks[block_idx] - else: - continue - if block_id == NULL_BLOCK_ID: - logger.debug( - "Skipping unavailable partial-tail source block " - "(req=%s, group=%d, block=%d)", - req_meta.req_id, - g_idx, - block_idx, - ) - continue - addr, size = db.prepare_value_for_block(block_id) - puts.append((db.key_for(key_hash), addr, size, db.metadata)) - return puts - def _maybe_offload_boundary_states(self, req_meta: ReqMeta) -> bool: - """Persist connector-pinned mamba "align" boundary states handed off - for this request, deduped against the store. - - This is every mamba key the connector writes — ``store_mask`` excludes - mamba groups from the positional normal save, aligned boundaries - included (see :meth:`_boundary_snapshot_puts`). + """Persist this job's boundary puts, deduped against the store. - The two entry kinds are keyed and sourced differently, so they are - prepared separately and put in one batch: - - - block-aligned for its group: a committed boundary-state snapshot, - the handed-off block itself; - - not block-aligned: the sub-block CoW partial tail, which also has to - cover the other groups' blocks in the normal save's lcm gap. + These are every mamba key the connector writes (``store_mask`` excludes + mamba groups from the positional normal save) plus the other groups' + partial-tail blocks in the normal save's lcm gap. Returns: True when no put is needed or every put succeeds, False otherwise. """ - offloads = req_meta.boundary_state_offloads - if not offloads or not req_meta.block_hashes: + if not req_meta.boundary_puts or not req_meta.block_hashes: return True - - snapshots: list[tuple[int, int, int]] = [] - sub_block: list[tuple[int, int, int]] = [] - for group_id, block_id, boundary in offloads: - entry = (group_id, block_id, boundary) - if boundary % self.token_databases[group_id].block_size == 0: - snapshots.append(entry) - else: - sub_block.append(entry) - - puts = self._boundary_snapshot_puts(req_meta, snapshots) - if sub_block and self.coord.enable_partial_hash_hits: - puts.extend(self._sub_block_tail_puts(req_meta, sub_block)) + puts = self._boundary_puts(req_meta) if not puts: return True @@ -854,9 +761,7 @@ def _handle_request(self, req_meta: ReqMeta): # Offload the handed-off mamba boundary states (independent of the # normal positional save, which may be skipped this step). - if req_meta.boundary_state_offloads is not None and not ( - self._maybe_offload_boundary_states(req_meta) - ): + if not self._maybe_offload_boundary_states(req_meta): return if token_len == 0: @@ -870,7 +775,7 @@ def _handle_request(self, req_meta: ReqMeta): store_masks = self.coord.store_mask( token_len, save_start, - num_prompt_tokens=req_meta.num_prompt_tokens, + num_prompt_tokens=req_meta.prefill_end_tokens, ) starts: list[int] = [] @@ -1603,31 +1508,10 @@ def __init__( ) return - self._kv_cache_groups = [ - dataclasses.replace( - group, - kv_cache_spec=resolve_dcp_kv_cache_spec( - group.kv_cache_spec, - self.dcp_size, - ), - ) - for group in kv_cache_config.prefix_cacheable_groups - ] - spec_cfg = getattr(vllm_config, "speculative_config", None) - use_eagle_block_drop = bool( - spec_cfg.use_eagle_block_drop() - if spec_cfg is not None - and callable(getattr(spec_cfg, "use_eagle_block_drop", None)) - else False - ) - self.coord = MooncakeStoreCoordinator( - self._kv_cache_groups, - scheduler_block_size=self.block_size, - hash_block_size=self.hash_block_size, - use_eagle=use_eagle_block_drop, - retention_interval=kv_cache_config.prefix_cache_retention_interval, - dcp_world_size=self.dcp_size, + self.coord = MooncakeStoreCoordinator.from_kv_cache_config( + kv_cache_config, vllm_config, self.block_size, self.hash_block_size ) + self._kv_cache_groups = self.coord.kv_cache_groups self.store_tp_size, store_namespace, store_layout_cls = ( self._select_store_layout(extra_config) ) diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index d0dadae27e9b..b13764abec3a 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -362,7 +362,10 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: scheduler_block_size, hash_block_size, ) - for kv_cache_config in kv_cache_configs: + # The scheduler's config is a deep copy taken before these are resolved, + # so set them there too: the KV connector's scheduler side reads the same + # fields as its workers. + for kv_cache_config in (*kv_cache_configs, scheduler_kv_cache_config): kv_cache_config.hash_block_size = hash_block_size kv_cache_config.cache_hit_alignment_tokens = cache_hit_alignment_tokens From 9f1c78388673f90d5b2be3ef77242143e7f7a3f6 Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Fri, 9 Oct 2026 19:26:53 -0500 Subject: [PATCH 214/624] [Test] Allow one fp8 step on K bytes in the kpool decode fuzz test (#60725) Signed-off-by: Rohan138 Co-authored-by: Claude Opus 5.5 Co-authored-by: Andreas Karatzas --- .../test_kpool_decode_update_batched.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/tests/kernels/test_kpool_decode_update_batched.py b/tests/kernels/test_kpool_decode_update_batched.py index 1f5693bad92a..4accf527f3d1 100644 --- a/tests/kernels/test_kpool_decode_update_batched.py +++ b/tests/kernels/test_kpool_decode_update_batched.py @@ -220,12 +220,19 @@ def _torch_reference( return kv_out, tail_cpu.to(torch.bfloat16).to(device="cuda") -def _assert_eq(r_ref, r_kern): +def _assert_eq(r_ref, r_kern, k_byte_tol=0): kv_ref, tail_ref = r_ref kv_kern, tail_kern = r_kern - assert torch.equal(kv_ref, kv_kern), ( - "kv_cache differs: max diff " - f"{(kv_ref.int() - kv_kern.int()).abs().max().item()}" + # Each page holds the fp8 K bytes followed by the fp32 scales. + k_region = PAGE_SIZE * HEAD_DIM + pages_ref = kv_ref.view(kv_ref.shape[0], -1) + pages_kern = kv_kern.view(kv_kern.shape[0], -1) + k_diff = (pages_ref[:, :k_region].int() - pages_kern[:, :k_region].int()).abs() + assert k_diff.max().item() <= k_byte_tol, ( + f"kv_cache K differs: max diff {k_diff.max().item()}" + ) + assert torch.equal(pages_ref[:, k_region:], pages_kern[:, k_region:]), ( + "kv_cache scales differ" ) assert torch.equal(tail_ref, tail_kern), ( "tail_kv_cache differs: max diff " @@ -652,4 +659,6 @@ def test_batched_matches_reference_fuzz(seed): r_ref = _torch_reference(kv, tail, tail_slot, key, score, ape, slot_map, pos) r_kern = _run_kernel(kv, tail, tail_slot, key, score, ape, slot_map, pos) - _assert_eq(r_ref, r_kern) + # Random inputs can land the pooled value on a bf16 rounding midpoint, + # where the kernel and the CPU reference may round apart by one fp8 step. + _assert_eq(r_ref, r_kern, k_byte_tol=1) From 7d1d8660d620149bc4eadf795cbb80a1f884ddc2 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 9 Oct 2026 17:58:09 -0700 Subject: [PATCH 215/624] [MRV2] Miscellaneous code structure updates (#60754) Signed-off-by: Nick Hill --- tests/v1/spec_decode/test_ngram_gpu.py | 2 +- tests/v1/worker/test_gpu_kv_connector.py | 12 ++-- tests/v1/worker/test_gpu_model_runner_v2.py | 10 +-- .../worker/test_gpu_model_runner_v2_eplb.py | 2 +- tests/v1/worker/test_gpu_ubatch_slicing.py | 10 +-- vllm/compilation/backends.py | 4 +- vllm/v1/worker/gpu/dp_utils.py | 38 +++++------ vllm/v1/worker/gpu/kv_connector.py | 39 +++++++++++- vllm/v1/worker/gpu/model_runner.py | 63 ++++++------------- .../gpu/model_states/encoder_decoder.py | 2 +- .../gpu/spec_decode/dflash/speculator.py | 10 +-- .../gpu/spec_decode/extract_hidden_states.py | 6 +- .../multi_module_mtp/speculator.py | 4 +- .../gpu/spec_decode/ngram/speculator.py | 2 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 10 +-- .../spec_decode/standalone_ar/speculator.py | 4 +- .../target_dependent_ar/speculator.py | 12 ++-- vllm/v1/worker/gpu/spec_decode/utils.py | 19 ++++++ 18 files changed, 140 insertions(+), 109 deletions(-) diff --git a/tests/v1/spec_decode/test_ngram_gpu.py b/tests/v1/spec_decode/test_ngram_gpu.py index 2359f0cdef2e..0170b58ad104 100644 --- a/tests/v1/spec_decode/test_ngram_gpu.py +++ b/tests/v1/spec_decode/test_ngram_gpu.py @@ -147,7 +147,7 @@ def _propose( next_prefill_tokens=torch.zeros(B, dtype=torch.int32, device=DEVICE), temperature=torch.zeros(B, dtype=torch.float32, device=DEVICE), seeds=torch.zeros(B, dtype=torch.int64, device=DEVICE), - dp_sync=None, + dp_sync_state=None, ) return drafts.cpu().tolist() diff --git a/tests/v1/worker/test_gpu_kv_connector.py b/tests/v1/worker/test_gpu_kv_connector.py index c442c6fc59e2..78fe3c185bd9 100644 --- a/tests/v1/worker/test_gpu_kv_connector.py +++ b/tests/v1/worker/test_gpu_kv_connector.py @@ -63,13 +63,11 @@ def test_load_start_phase( request_indices = torch.tensor([3, 1]) request_ids = ["first", "second"] - attn_metadata = {"layer": object()} - connector.pre_forward( - output, - request_state_indices=request_indices, - request_ids=request_ids, - attn_metadata=attn_metadata, + input_batch = SimpleNamespace( + idx_mapping=request_indices, req_ids=request_ids, num_tokens=5 ) + attn_metadata = {"layer": object()} + connector.pre_forward(output, input_batch, attn_metadata=attn_metadata) assert events == ( ["handle", "bind", "start"] if has_sync_kv_loads else ["handle", "bind"] ) @@ -80,6 +78,8 @@ def test_load_start_phase( kwargs = connector.kv_connector.start_load_kv.call_args.kwargs assert kwargs["request_state_indices"] is request_indices assert kwargs["request_ids"] is request_ids + assert kwargs["num_tokens"] == 5 + assert kwargs["scheduler_output"] is output assert kwargs["attn_metadata"] is attn_metadata # A subsequent step without a forward must not reuse the prior batch. diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 47173da08155..85c3ddfce425 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -21,6 +21,7 @@ from vllm.v1.worker.gpu.async_utils import async_copy_to_np from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.model_runner import ExecuteModelState, GPUModelRunner +from vllm.v1.worker.gpu.spec_decode.utils import get_drafter_hidden_states def test_non_last_pp_rank_uses_global_batch_for_sample_feedback(): @@ -43,7 +44,7 @@ def test_non_last_pp_rank_uses_global_batch_for_sample_feedback(): slot_mappings_by_layer=None, hidden_states=None, aux_hidden_states=None, - dp_sync=None, + dp_sync_state=None, finished_req_ids=set(), ec_connector_output=None, cudagraph_stats=None, @@ -342,16 +343,15 @@ def test_capture_model_profile_only_skips_lock(monkeypatch): @pytest.mark.parametrize("target_buffer", ["absent", "none", "tensor"]) def test_get_drafter_hidden_states_tolerates_missing_target_buffer(target_buffer): """Targets allocate the MTP hidden buffer only for hidden-state drafters.""" - runner = GPUModelRunner.__new__(GPUModelRunner) hidden_states = torch.zeros(4, 8) buffer = torch.arange(16 * 8, dtype=torch.float32).view(16, 8) if target_buffer == "absent": - runner.model = SimpleNamespace() + model = SimpleNamespace() else: returned = buffer if target_buffer == "tensor" else None - runner.model = SimpleNamespace(get_mtp_target_hidden_states=lambda: returned) + model = SimpleNamespace(get_mtp_target_hidden_states=lambda: returned) - out = runner._get_drafter_hidden_states(hidden_states) + out = get_drafter_hidden_states(model, hidden_states) if target_buffer == "tensor": assert torch.equal(out, buffer[:4]) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index d8090de2a359..6b38bfef968d 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -185,7 +185,7 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch): slot_mappings_by_layer=None, hidden_states=None, aux_hidden_states=None, - dp_sync=None, + dp_sync_state=None, finished_req_ids=set(), ec_connector_output=None, routed_experts=None, diff --git a/tests/v1/worker/test_gpu_ubatch_slicing.py b/tests/v1/worker/test_gpu_ubatch_slicing.py index ebf334e45d4f..127d289afa44 100644 --- a/tests/v1/worker/test_gpu_ubatch_slicing.py +++ b/tests/v1/worker/test_gpu_ubatch_slicing.py @@ -350,10 +350,10 @@ def test_decode_threshold_only_applies_when_every_rank_is_a_uniform_decode(): def test_microbatching_pads_all_ranks_to_the_largest(): """Ranks must run the same token count so microbatch sizes line up.""" - desc, dp_sync = _sync_dp([200, 256]) + desc, dp_sync_state = _sync_dp([200, 256]) assert desc.num_tokens == 256 - assert dp_sync is not None - assert dp_sync.num_tokens_across_dp.tolist() == [256, 256] + assert dp_sync_state is not None + assert dp_sync_state.num_tokens_across_dp.tolist() == [256, 256] def test_microbatching_survives_a_rank_that_cannot_fill_it(): @@ -959,7 +959,7 @@ def test_microbatched_graph_needs_every_rank_to_reach_the_split(): assert desc.num_ubatches == 2 # Rank 1 pads the group up to a 128-token graph, but rank 0 only reaches 33. - desc, dp_sync = _sync_dp([33, 128], uniform, cudagraph_manager=manager) + desc, dp_sync_state = _sync_dp([33, 128], uniform, cudagraph_manager=manager) assert desc.cg_mode == CUDAGraphMode.NONE assert desc.num_ubatches == 2 - assert dp_sync is not None and dp_sync.eager + assert dp_sync_state is not None and dp_sync_state.eager diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index d985d8de757a..434c1af47ee8 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -1091,12 +1091,12 @@ def __call__(self, graph: fx.GraphModule, example_inputs: Sequence[Any]) -> Any: # TODO(patchy): the V1 torch.compile ngram-gpu kernel causes vllm # torch compile cache errors. The V2 implementation is pure Triton and # does not need the cache disabled. - is_ngram_gpu_enabled = ( + is_v1_ngram_gpu_enabled = ( vllm_config.speculative_config is not None and vllm_config.speculative_config.use_ngram_gpu() and not vllm_config.use_v2_model_runner ) - disable_cache = disable_cache or is_ngram_gpu_enabled + disable_cache = disable_cache or is_v1_ngram_gpu_enabled if disable_cache: logger.info("%svLLM's torch.compile cache is disabled.", self._log_prefix()) diff --git a/vllm/v1/worker/gpu/dp_utils.py b/vllm/v1/worker/gpu/dp_utils.py index 15db7c197c96..5c71ebcea018 100644 --- a/vllm/v1/worker/gpu/dp_utils.py +++ b/vllm/v1/worker/gpu/dp_utils.py @@ -152,7 +152,7 @@ def sync_cudagraph_and_dp_padding( num_reqs=num_reqs, ) - synced_cg_mode = CUDAGraphMode(int(cg_mode_across_dp.min().item())) + synced_cg_mode = CUDAGraphMode(int(cg_mode_across_dp.min())) # If any rank wants to run eager, all ranks run eager if synced_cg_mode == CUDAGraphMode.NONE: @@ -175,13 +175,13 @@ def sync_cudagraph_and_dp_padding( "cudagraph_manager should only be None during profile run, " "where synced_cg_mode must be NONE across all DP ranks" ) - synced_num_tokens = int(num_tokens_across_dp.max().item()) + synced_num_tokens = int(num_tokens_across_dp.max()) # Varlen decode graphs are selected by the query-length bound, so ranks must agree # on it or they pad to different token counts below. synced_max_query_len: int | None = None - if bool(torch.all(max_query_lens_across_dp != -1).item()): - synced_max_query_len = int(max_query_lens_across_dp.max().item()) + if bool(torch.all(max_query_lens_across_dp != -1)): + synced_max_query_len = int(max_query_lens_across_dp.max()) # Dispatch for the final synced values, use num_reqs instead of synced_num_reqs # so we don't perform request padding for PIECEWISE graphs. @@ -223,12 +223,12 @@ def dispatch_cg_and_sync_dp( parallel_config: ParallelConfig | None = None, allow_ubatching: bool = False, uniform_decode: bool = False, - dp_sync: DPSyncState | None = None, + dp_sync_state: DPSyncState | None = None, ) -> tuple[BatchExecutionDescriptor, DPSyncState | None]: """Pick a cudagraph descriptor for this batch, agreeing it across DP ranks. Runs a collective when dp_size > 1 so every rank dispatches to the same - shape. Pass `dp_sync` from a dispatch already made over this same batch (a + shape. Pass `dp_sync_state` from a dispatch already made over this same batch (a drafter's prefill runs the target's batch shape) to reuse that agreement instead, with no collective. @@ -238,7 +238,7 @@ def dispatch_cg_and_sync_dp( num_reqs: Requests in this rank's batch. num_tokens: Tokens in this rank's batch, already padded by the caller. uniform_token_count: Per-request token count if this rank's batch is a - uniform decode, else None. `dp_sync.uniform_token_count` takes its + uniform decode, else None. `dp_sync_state.uniform_token_count` takes its place when a sync is reused, since that one is agreed across ranks. dp_size: Data-parallel world size. 1 skips all cross-rank work. dp_rank: This rank's index in the DP group. @@ -254,7 +254,7 @@ def dispatch_cg_and_sync_dp( microbatches. Agreed across ranks before it takes effect. uniform_decode: Whether this rank's batch is a uniform decode, used to pick the microbatching split. - dp_sync: Agreement from a prior dispatch over this same batch, to reuse. + dp_sync_state: Agreement from a prior dispatch over this same batch, to reuse. Must come from a batch with this same padded `num_tokens` and the same `uniform_token_count`; `num_reqs` may differ, as neither depends on it. Passing a sync from a different batch is a caller @@ -265,7 +265,7 @@ def dispatch_cg_and_sync_dp( dispatch to reuse. It is None when `dp_size` is 1 or no rank has work. """ - reuse_eager = dp_sync is not None and dp_sync.eager + reuse_eager = dp_sync_state is not None and dp_sync_state.eager if need_eager or reuse_eager: batch_desc = BatchExecutionDescriptor( @@ -282,7 +282,7 @@ def dispatch_cg_and_sync_dp( batch_desc = cudagraph_manager.dispatch( num_reqs, num_tokens, - dp_sync.uniform_token_count if dp_sync is not None else uniform_token_count, + dp_sync_state.uniform_token_count if dp_sync_state else uniform_token_count, num_active_loras=num_active_loras, max_query_len=max_query_len, ) @@ -290,24 +290,24 @@ def dispatch_cg_and_sync_dp( if dp_size == 1: return batch_desc, None - if dp_sync is not None: - assert dp_sync.num_tokens_across_dp[dp_rank] == num_tokens, ( + if dp_sync_state is not None: + assert dp_sync_state.num_tokens_across_dp[dp_rank] == num_tokens, ( "reusing a DP sync taken over a different batch" ) assert ( - dp_sync.uniform_token_count is None - or uniform_token_count == dp_sync.uniform_token_count + dp_sync_state.uniform_token_count is None + or uniform_token_count == dp_sync_state.uniform_token_count ), "reusing a DP sync taken over a different batch" - if not dp_sync.eager and batch_desc.num_tokens != num_tokens: + if not dp_sync_state.eager and batch_desc.num_tokens != num_tokens: # Capture sizes can differ between managers, so this one may # pad further. Every rank pads alike, so report what will run. - dp_sync = replace( - dp_sync, + dp_sync_state = replace( + dp_sync_state, num_tokens_across_dp=torch.full_like( - dp_sync.num_tokens_across_dp, batch_desc.num_tokens + dp_sync_state.num_tokens_across_dp, batch_desc.num_tokens ), ) - return batch_desc, dp_sync + return batch_desc, dp_sync_state return sync_cudagraph_and_dp_padding( cudagraph_manager, diff --git a/vllm/v1/worker/gpu/kv_connector.py b/vllm/v1/worker/gpu/kv_connector.py index 1ba75beae271..e9b71dae3e8c 100644 --- a/vllm/v1/worker/gpu/kv_connector.py +++ b/vllm/v1/worker/gpu/kv_connector.py @@ -24,12 +24,18 @@ if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput + from vllm.v1.worker.gpu.input_batch import InputBatch class KVConnector: """KVConnector interface used by GPUModelRunner.""" - def pre_forward(self, scheduler_output: "SchedulerOutput", **kwargs: Any) -> None: + def pre_forward( + self, + scheduler_output: "SchedulerOutput", + input_batch: "InputBatch | None" = None, + attn_metadata: dict[str, Any] | None = None, + ) -> None: pass def finish_forward(self) -> None: @@ -61,7 +67,12 @@ def __init__( self._pending_load_kwargs: dict[str, Any] | None = None self._disabled = False - def pre_forward(self, scheduler_output: "SchedulerOutput", **kwargs: Any) -> None: + def pre_forward( + self, + scheduler_output: "SchedulerOutput", + input_batch: "InputBatch | None" = None, + attn_metadata: dict[str, Any] | None = None, + ) -> None: if self._disabled: return @@ -69,7 +80,9 @@ def pre_forward(self, scheduler_output: "SchedulerOutput", **kwargs: Any) -> Non assert kv_connector_metadata is not None self.kv_connector.handle_preemptions(kv_connector_metadata) self.kv_connector.bind_connector_metadata(kv_connector_metadata) - self._pending_load_kwargs = kwargs + self._pending_load_kwargs = self._build_load_kwargs( + scheduler_output, input_batch, attn_metadata + ) if scheduler_output.has_sync_kv_loads: # Sync loads need to run before this step's forward. @@ -77,6 +90,26 @@ def pre_forward(self, scheduler_output: "SchedulerOutput", **kwargs: Any) -> Non # Otherwise start the async load after forward to keep submission # off the critical path. + @staticmethod + def _build_load_kwargs( + scheduler_output: "SchedulerOutput", + input_batch: "InputBatch | None", + attn_metadata: dict[str, Any] | None, + ) -> dict[str, Any]: + """Keyword args for the connector's start_load_kv().""" + if input_batch is None: + # No forward this step. + return {} + load_kwargs: dict[str, Any] = dict( + scheduler_output=scheduler_output, + request_state_indices=input_batch.idx_mapping, + request_ids=input_batch.req_ids, + num_tokens=input_batch.num_tokens, + ) + if attn_metadata is not None: + load_kwargs["attn_metadata"] = attn_metadata + return load_kwargs + def _start_load_kv(self) -> None: load_kwargs = self._pending_load_kwargs assert load_kwargs is not None diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 2e9925f1ad1c..f39fd0a6e6a7 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -162,7 +162,10 @@ get_max_chunk_logits, ) from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator -from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler +from vllm.v1.worker.gpu.spec_decode.utils import ( + DraftTokensHandler, + get_drafter_hidden_states, +) from vllm.v1.worker.gpu.states import RequestState from vllm.v1.worker.gpu.structured_outputs import ( StructuredOutputsWorker, @@ -578,9 +581,7 @@ def main_stream(self) -> torch.cuda.Stream: def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]: encoder_runner = getattr(self.model_state, "encoder_runner", None) - if encoder_runner is None: - return {} - return encoder_runner.get_encoder_timing_stats() + return encoder_runner.get_encoder_timing_stats() if encoder_runner else {} def get_kv_cache_spec(self): kv_cache_spec = get_kv_cache_spec(self.vllm_config) @@ -905,7 +906,7 @@ def _dummy_run( slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer hidden_states = self.execute_model_state.hidden_states aux_hidden_states = self.execute_model_state.aux_hidden_states - dp_sync = self.execute_model_state.dp_sync + dp_sync_state = self.execute_model_state.dp_sync_state self.execute_model_state = None self.step_timing.forward_end() @@ -921,7 +922,7 @@ def _dummy_run( ) mm_inputs = [], all_false - spec_hidden_states = self._get_drafter_hidden_states(hidden_states) + spec_hidden_states = get_drafter_hidden_states(self.model, hidden_states) if isinstance(self.sampler, GPUWatermarkSampler): self.speculator.prepare_watermarking( self.sampler, input_batch.idx_mapping @@ -943,7 +944,7 @@ def _dummy_run( next_prefill_tokens=self.req_states.next_prefill_tokens, temperature=self.sampler.sampling_states.temperature.gpu, seeds=self.sampler.sampling_states.seeds.gpu, - dp_sync=dp_sync, + dp_sync_state=dp_sync_state, dummy_run=True, skip_attn_for_dummy_run=skip_attn, mm_inputs=mm_inputs, @@ -1771,7 +1772,7 @@ def execute_model( # cross-attention cache with dynamic encoder outputs. skip_compiled = True - batch_desc, dp_sync = dispatch_cg_and_sync_dp( + batch_desc, dp_sync_state = dispatch_cg_and_sync_dp( self.cudagraph_manager, num_reqs, num_toks, @@ -1983,13 +1984,6 @@ def execute_model( ) self.step_timing.forward_start() - connector_kwargs = dict( - scheduler_output=scheduler_output, - request_state_indices=input_batch.idx_mapping, - request_ids=input_batch.req_ids, - num_tokens=input_batch.num_tokens, - ) - # Run model. if batch_desc.cg_mode == CUDAGraphMode.FULL: # Use explicit cudagraph replay for FULL mode. @@ -1997,7 +1991,7 @@ def execute_model( # because they are already copied to the CUDA graph input buffers. assert self.cudagraph_manager is not None self.kv_connector.pre_forward( - **connector_kwargs, attn_metadata=attn_metadata + scheduler_output, input_batch, attn_metadata=attn_metadata ) model_output = self.cudagraph_manager.run_fullgraph(batch_desc) else: @@ -2013,16 +2007,15 @@ def execute_model( self.vllm_config, num_tokens=input_batch.num_tokens_after_padding, cudagraph_runtime_mode=batch_desc.cg_mode, - num_tokens_across_dp=( - dp_sync.num_tokens_across_dp if dp_sync is not None else None - ), + num_tokens_across_dp=dp_sync_state + and dp_sync_state.num_tokens_across_dp, batch_descriptor=batch_descriptor, ubatch_slices=ubatch_slices, slot_mapping=slot_mappings_by_layer, skip_compiled=skip_compiled, is_padding=input_batch.is_padding, ): - self.kv_connector.pre_forward(**connector_kwargs) + self.kv_connector.pre_forward(scheduler_output, input_batch) if ubatch_state is not None: assert self.ubatch_runner is not None model_output = self.ubatch_runner.run( @@ -2064,7 +2057,7 @@ def execute_model( slot_mappings_by_layer=slot_mappings_by_layer, hidden_states=hidden_states, aux_hidden_states=aux_hidden_states, - dp_sync=dp_sync, + dp_sync_state=dp_sync_state, finished_req_ids=finished_req_ids, ec_connector_output=ec_connector_output, cudagraph_stats=cudagraph_stats, @@ -2080,24 +2073,6 @@ def execute_model( ) return None - def _get_drafter_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: - """Hidden states fed to the drafter. - - Targets such as DeepSeek V4 expose the pre-hc_head residual through - get_mtp_target_hidden_states(). The buffer is sized at - max_num_batched_tokens and only allocated for drafters that consume - target hidden states, so None means "use the regular hidden states". - """ - get_target_hidden_states = getattr( - self.model, "get_mtp_target_hidden_states", None - ) - if get_target_hidden_states is None: - return hidden_states - target_hidden_states = get_target_hidden_states() - if target_hidden_states is None: - return hidden_states - return target_hidden_states[: hidden_states.shape[0]] - @torch.inference_mode() @step_eplb_after() def sample_tokens( @@ -2112,7 +2087,7 @@ def sample_tokens( slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer hidden_states = self.execute_model_state.hidden_states aux_hidden_states = self.execute_model_state.aux_hidden_states - dp_sync = self.execute_model_state.dp_sync + dp_sync_state = self.execute_model_state.dp_sync_state finished_req_ids = self.execute_model_state.finished_req_ids ec_connector_output = self.execute_model_state.ec_connector_output cudagraph_stats = self.execute_model_state.cudagraph_stats @@ -2243,7 +2218,9 @@ def sample_tokens( self.speculator.observe_verification( input_batch.idx_mapping, num_sampled, num_rejected ) - spec_hidden_states = self._get_drafter_hidden_states(draft_hidden_states) + spec_hidden_states = get_drafter_hidden_states( + self.model, draft_hidden_states + ) if isinstance(self.sampler, GPUWatermarkSampler): self.speculator.prepare_watermarking( self.sampler, input_batch.idx_mapping @@ -2262,7 +2239,7 @@ def sample_tokens( self.sampler.sampling_states.temperature.gpu, self.sampler.sampling_states.seeds.gpu, num_speculative_tokens=num_spec_tokens, - dp_sync=dp_sync, + dp_sync_state=dp_sync_state, mm_inputs=mm_inputs, ) if num_spec_tokens < self.num_speculative_steps: @@ -2419,7 +2396,7 @@ class ExecuteModelState(NamedTuple): slot_mappings_by_layer: dict[str, torch.Tensor] | None hidden_states: torch.Tensor | None aux_hidden_states: list[torch.Tensor] | None - dp_sync: DPSyncState | None + dp_sync_state: DPSyncState | None finished_req_ids: set[str] ec_connector_output: ECConnectorOutput | None cudagraph_stats: CUDAGraphStat | None diff --git a/vllm/v1/worker/gpu/model_states/encoder_decoder.py b/vllm/v1/worker/gpu/model_states/encoder_decoder.py index bebb9cd41be2..6363dbad14d3 100644 --- a/vllm/v1/worker/gpu/model_states/encoder_decoder.py +++ b/vllm/v1/worker/gpu/model_states/encoder_decoder.py @@ -135,7 +135,7 @@ def prepare_attn( if for_capture: max_seq_len = self.max_model_len else: - max_seq_len = int(seq_lens_cpu_upper_bound[:num_reqs].max().item()) + max_seq_len = int(seq_lens_cpu_upper_bound[:num_reqs].max()) attn_metadata = build_attn_metadata( attn_groups=attn_groups, num_reqs=num_reqs, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 9538dcb4d4bf..884de5a264be 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -355,7 +355,7 @@ def propose( temperature: torch.Tensor, # [max_num_reqs] seeds: torch.Tensor, - dp_sync: DPSyncState | None = None, + dp_sync_state: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, @@ -446,8 +446,10 @@ def propose( ) batch_sync, num_batch_tokens = ( - self._build_uniform_batch_dp_sync(dp_sync, num_reqs, self.num_query_per_req) - if dp_sync is not None + self._build_uniform_batch_dp_sync( + dp_sync_state, num_reqs, self.num_query_per_req + ) + if dp_sync_state is not None else (None, num_query_tokens) ) # Every DFlash step has exactly num_query_per_req tokens, so we can use FULL CGs @@ -459,7 +461,7 @@ def propose( dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, - dp_sync=batch_sync, + dp_sync_state=batch_sync, ) num_tokens_padded = batch_desc.num_tokens num_tokens_across_dp = ( diff --git a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py index 5afbe640b2f2..e02d7ee11d70 100644 --- a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py +++ b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py @@ -98,7 +98,7 @@ def propose( next_prefill_tokens: torch.Tensor, temperature: torch.Tensor, seeds: torch.Tensor, - dp_sync: DPSyncState | None = None, + dp_sync_state: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, @@ -145,9 +145,7 @@ def propose( draft_attn_metadata, self.vllm_config, num_tokens=num_tokens, - num_tokens_across_dp=( - dp_sync.num_tokens_across_dp if dp_sync is not None else None - ), + num_tokens_across_dp=dp_sync_state and dp_sync_state.num_tokens_across_dp, cudagraph_runtime_mode=CUDAGraphMode.NONE, slot_mapping=draft_slot_mappings, is_padding=input_batch.is_padding[:num_tokens], diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py index 919e939c03cc..a453baf6e3fa 100644 --- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py @@ -150,7 +150,7 @@ def propose( temperature: torch.Tensor, # [max_num_reqs] seeds: torch.Tensor, - dp_sync: DPSyncState | None = None, + dp_sync_state: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, @@ -201,7 +201,7 @@ def propose( dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, - dp_sync=dp_sync, + dp_sync_state=dp_sync_state, ) num_tokens_across_dp = ( batch_sync.num_tokens_across_dp if batch_sync is not None else None diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py index 7a61e3581a11..b74a0ec4f984 100644 --- a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -226,7 +226,7 @@ def propose( next_prefill_tokens: torch.Tensor, temperature: torch.Tensor, seeds: torch.Tensor, - dp_sync: DPSyncState | None = None, + dp_sync_state: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index cbd94ed0542d..c3e8dcbff4d1 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -92,7 +92,7 @@ def propose( temperature: torch.Tensor, # [max_num_reqs] seeds: torch.Tensor, - dp_sync: DPSyncState | None = None, + dp_sync_state: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, @@ -525,18 +525,18 @@ def _copy_request_inputs( def _build_uniform_batch_dp_sync( self, - target_dp_sync: DPSyncState, + target_dp_sync_state: DPSyncState, num_reqs: int, num_query_per_req: int = 1, ) -> tuple[DPSyncState, int]: - num_batch_tokens = target_dp_sync.num_reqs * num_query_per_req + num_batch_tokens = target_dp_sync_state.num_reqs * num_query_per_req assert num_reqs * num_query_per_req <= num_batch_tokens, ( "reusing a DP sync that does not cover this batch's requests" ) return replace( - target_dp_sync, + target_dp_sync_state, num_tokens_across_dp=torch.full_like( - target_dp_sync.num_tokens_across_dp, num_batch_tokens + target_dp_sync_state.num_tokens_across_dp, num_batch_tokens ), uniform_token_count=num_query_per_req, eager=False, diff --git a/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py index 898de1e3f886..04e79debb869 100644 --- a/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py @@ -350,7 +350,7 @@ def propose( next_prefill_tokens: torch.Tensor, temperature: torch.Tensor, seeds: torch.Tensor, - dp_sync: DPSyncState | None = None, + dp_sync_state: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, @@ -362,7 +362,7 @@ def propose( num_reqs = input_batch.num_reqs skip_attn = dummy_run and skip_attn_for_dummy_run num_tokens_across_dp = ( - dp_sync.num_tokens_across_dp if dp_sync is not None else None + dp_sync_state.num_tokens_across_dp if dp_sync_state is not None else None ) # Copy per-request temperature/seeds/idx_mapping into pre-allocated diff --git a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py index 4894ce77e962..02e28875b823 100644 --- a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py @@ -219,7 +219,7 @@ def propose( temperature: torch.Tensor, # [max_num_reqs] seeds: torch.Tensor, - dp_sync: DPSyncState | None = None, + dp_sync_state: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, @@ -299,7 +299,7 @@ def propose( dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, - dp_sync=dp_sync, + dp_sync_state=dp_sync_state, ) num_tokens_across_dp = ( prefill_batch_sync.num_tokens_across_dp @@ -352,8 +352,10 @@ def propose( ) decode_batch_sync, num_batch_tokens = ( - self._build_uniform_batch_dp_sync(dp_sync, num_reqs, num_query_per_req=1) - if dp_sync is not None + self._build_uniform_batch_dp_sync( + dp_sync_state, num_reqs, num_query_per_req=1 + ) + if dp_sync_state is not None else (None, num_reqs) ) # Each request produces exactly 1 token per draft generation step, @@ -366,7 +368,7 @@ def propose( dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, - dp_sync=decode_batch_sync, + dp_sync_state=decode_batch_sync, ) if self.decode_cudagraph_manager is not None: decode_batch_desc = self.decode_cudagraph_manager.specialize_spec_tokens( diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index f73acb2bdaed..b005f5c3f3ae 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -13,6 +13,25 @@ logger = init_logger(__name__) +def get_drafter_hidden_states( + model: torch.nn.Module, hidden_states: torch.Tensor +) -> torch.Tensor: + """Hidden states fed to the drafter. + + Targets such as DeepSeek V4 expose the pre-hc_head residual through + get_mtp_target_hidden_states(). The buffer is sized at + max_num_batched_tokens and only allocated for drafters that consume + target hidden states, so None means "use the regular hidden states". + """ + get_target_hidden_states = getattr(model, "get_mtp_target_hidden_states", None) + if get_target_hidden_states is None: + return hidden_states + target_hidden_states = get_target_hidden_states() + if target_hidden_states is None: + return hidden_states + return target_hidden_states[: hidden_states.shape[0]] + + def get_pp_safe_draft_load_config(load_config: LoadConfig) -> LoadConfig: """Avoid collectives that include PP ranks without a draft model.""" if get_pp_group().world_size > 1 and load_config.load_format == "fastsafetensors": From 3c802af3cf6263d0ef6fb356a7884680ecf15f04 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Fri, 9 Oct 2026 19:33:22 -0600 Subject: [PATCH 216/624] [Feat][Frontend] structured decisions: DiffusionGemma read strategy (#60466) Signed-off-by: Matt Mastracci Signed-off-by: Simon Veitner Co-authored-by: Claude Opus 5.5 Co-authored-by: Simon Veitner Co-authored-by: Misha Goin --- .../online_serving/structured_decisions.md | 5 +- .../test_question_types.py | 2 +- .../structured_decisions/test_strategies.py | 37 +++- vllm/entrypoints/generate/api_router.py | 7 +- .../structured_decisions/question_types.py | 13 +- .../structured_decisions/strategies.py | 169 +++++++++++++++--- 6 files changed, 193 insertions(+), 40 deletions(-) diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md index a1f58a3769b9..a9ed5a2ee5aa 100644 --- a/docs/serving/online_serving/structured_decisions.md +++ b/docs/serving/online_serving/structured_decisions.md @@ -20,7 +20,10 @@ log probabilities. Thinking is off unless `chat_template_kwargs` turns it on. -This endpoint does not serve diffusion models yet. +Probabilities can differ between identical requests. They depend on which +requests share a batch and on whether a read hits the prefix cache. Set +`VLLM_BATCH_INVARIANT=1` for reads that repeat exactly, at a cost in +throughput. Every read of a request starts with the state, so with `--enable-prefix-caching` the state is prefilled once and each further diff --git a/tests/entrypoints/generate/structured_decisions/test_question_types.py b/tests/entrypoints/generate/structured_decisions/test_question_types.py index 3b9ad7ba60a9..26047b5d6c07 100644 --- a/tests/entrypoints/generate/structured_decisions/test_question_types.py +++ b/tests/entrypoints/generate/structured_decisions/test_question_types.py @@ -105,7 +105,7 @@ def test_noul_labels_and_prompt(): assert q.labels == ("yes", "no") assert [a.name for a in q.options] == ["yes", "no"] assert q.type.prompt(q) == ( - "Question: Reply within the hour?\nAnswer with yes or no only." + "Question: Reply within the hour?\nyes\nno\nAnswer with yes or no only." ) described = noul(criteria={"true": "same day", "false": "later"}) assert described.type.prompt(described) == ( diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index 5dcb8edb3061..a63e2587521c 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -17,12 +17,16 @@ StructuredDecisionRequest, StructuredDecisionResponse, ) -from vllm.entrypoints.generate.structured_decisions.question_types import LABELS +from vllm.entrypoints.generate.structured_decisions.question_types import ( + LABELS, + StructuredDecisionError, +) from vllm.entrypoints.generate.structured_decisions.serving import ( ServingStructuredDecisions, state_text, ) from vllm.entrypoints.generate.structured_decisions.strategies import ( + DiffusionGemmaCanvasStrategy, NextTokenStrategy, ReadContext, reply_label_ids, @@ -52,6 +56,10 @@ def test_strategy_selection(): qwen = "Qwen3ForCausalLM" assert select_read_strategy(model(qwen)) is NextTokenStrategy assert select_read_strategy(model(qwen, "processed_logprobs")) is NextTokenStrategy + assert ( + select_read_strategy(model("DiffusionGemmaForBlockDiffusion")) + is DiffusionGemmaCanvasStrategy + ) with pytest.raises(ValueError, match="does not support LlamaForCausalLM"): select_read_strategy(model("LlamaForCausalLM")) with pytest.raises(ValueError, match="not raw_logits"): @@ -220,3 +228,30 @@ async def test_request_metadata_reaches_every_question( assert call.kwargs["trace_headers"] == ( TRACE_HEADERS if tracing_enabled and request_headers else None ) + + +def test_canvas_read(): + strategy = DiffusionGemmaCanvasStrategy.__new__(DiffusionGemmaCanvasStrategy) + strategy.thought, strategy.end, strategy.pad = [10, 11, 12, 13], 106, 0 + strategy.width, strategy.vocab_size = 16, 1000 + strategy.max_model_len = 23 + params = strategy._sampling_params([65, 66], prompt_ids=[1, 2, 3]) + assert params.extra_args is not None + canvas = params.extra_args["diffusion_seed_canvas"] + # The label slot as noise, the end of the turn, padding. + assert canvas[1:] == [106] + [0] * 14 + assert params.max_tokens == 2 and params.logprob_token_ids == [65, 66] + again = strategy._sampling_params([65, 66], prompt_ids=[1, 2, 3]) + assert again.extra_args == params.extra_args + read_input = strategy._read_input( + {"type": "token", "prompt_token_ids": [1, 2, 3]}, [1, 2, 3] + ) + assert read_input == { + "type": "token", + "prompt_token_ids": [1, 2, 3, 10, 11, 12, 13], + } + # One token more and the thought and the canvas no longer fit. + with pytest.raises(StructuredDecisionError, match="max_model_len=23"): + strategy._read_input( + {"type": "token", "prompt_token_ids": [1, 2, 3, 4]}, [1, 2, 3, 4] + ) diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 67a256192443..18dfb87a26da 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -5,6 +5,7 @@ from fastapi import FastAPI import vllm.envs as envs +from vllm.logger import init_logger if TYPE_CHECKING: from argparse import Namespace @@ -17,6 +18,8 @@ else: RequestLogger = object +logger = init_logger(__name__) + def register_generate_api_routers(app: FastAPI): from vllm.entrypoints.openai.chat_completion.api_router import ( @@ -243,7 +246,9 @@ async def init_generate_state( default_chat_template_kwargs=default_chat_template_kwargs, ) ) - except ValueError: + except ValueError as e: + # Info, since every model without a read strategy lands here. + logger.info("/v1/systemone is disabled: %s", e) strategy = None state.serving_structured_decisions = ( ServingStructuredDecisions( diff --git a/vllm/entrypoints/generate/structured_decisions/question_types.py b/vllm/entrypoints/generate/structured_decisions/question_types.py index 0e80694b14eb..18e7a6001855 100644 --- a/vllm/entrypoints/generate/structured_decisions/question_types.py +++ b/vllm/entrypoints/generate/structured_decisions/question_types.py @@ -55,8 +55,7 @@ def answer( ``question.labels[i]`` among the labels, and the list sums to 1. ``label_mass`` is the labels' total probability over the vocabulary.""" - def reply_line(self, label: str, option: Option) -> str | None: - """The option's line in the prompt, or None to leave it out.""" + def reply_line(self, label: str, option: Option) -> str: return ( f"{label}: {option.name} - {option.description}" if option.description @@ -67,9 +66,7 @@ def prompt(self, question: Question) -> str: """The question as the model reads it, after the state.""" lines = [f"Question: {question.instructions}"] if question.instructions else [] for label, o in zip(question.labels, question.options): - line = self.reply_line(label, o) - if line is not None: - lines.append(line) + lines.append(self.reply_line(label, o)) lines.append(self.reply_instruction) return "\n".join(lines) @@ -192,8 +189,10 @@ def parse_options(self, qid: str, criteria: Any) -> list[Option]: Option("no", None if false is None else str(false)), ] - def reply_line(self, label: str, option: Option) -> str | None: - return f"{label}: {option.description}" if option.description else None + def reply_line(self, label: str, option: Option) -> str: + # A side without a description still shows its label. Without the + # label in the prompt, models answer "Yes" or "No" instead. + return f"{label}: {option.description}" if option.description else label def answer( self, question: Question, probs: list[float], label_mass: float diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index 65b5a9dd7b05..f79375e42a83 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -2,10 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Read strategies: how a model is asked for label probabilities. -Only models in NEXT_TOKEN_ARCHITECTURES are currently supported. +Only models in READ_STRATEGIES are currently supported. """ +import hashlib +import json import math +import random from abc import ABC, abstractmethod from collections.abc import Mapping, Sequence from dataclasses import dataclass @@ -15,7 +18,7 @@ from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.generate.label_reads import next_token_label_reads -from vllm.inputs import EngineInput +from vllm.inputs import EngineInput, tokens_input from vllm.lora.request import LoRARequest from vllm.renderers.inputs.preprocess import extract_prompt_components from vllm.renderers.online_renderer import OnlineRenderer @@ -80,6 +83,13 @@ async def read( """One read per question, in order.""" +def prompt_tail(tokenizer: TokenizerLike, prompt_ids: Sequence[int]) -> list[int]: + """The prompt's ids after its last added token.""" + added = set(tokenizer.get_added_vocab().values()) + start = max((i + 1 for i, t in enumerate(prompt_ids) if t in added), default=0) + return list(prompt_ids[start:]) + + def reply_label_ids( tokenizer: TokenizerLike, prompt_ids: Sequence[int], @@ -88,9 +98,7 @@ def reply_label_ids( """The prompt's ids after its last added token, and the token each of ``labels`` adds after them as the first token of the reply. Raises ValueError if a label is not one distinct token there.""" - added = set(tokenizer.get_added_vocab().values()) - start = max((i + 1 for i, t in enumerate(prompt_ids) if t in added), default=0) - tail = list(prompt_ids[start:]) + tail = prompt_tail(tokenizer, prompt_ids) text = tokenizer.decode(tail) ids: list[int] = [] for label in labels: @@ -104,6 +112,18 @@ def reply_label_ids( return tail, ids +def type_label_ids( + tokenizer: TokenizerLike, prompt_ids: Sequence[int] +) -> dict[str, int]: + """The reply's first token for each label of every registered question + type. Raises ValueError as reply_label_ids does.""" + label_ids: dict[str, int] = {} + for qtype in QUESTION_TYPES.values(): + _, ids = reply_label_ids(tokenizer, prompt_ids, qtype.label_set) + label_ids.update(zip(qtype.label_set, ids)) + return label_ids + + class NextTokenStrategy(ReadStrategy): """Autoregressive models. Each question is one request: the state, then the question with its labeled options, and the logprobs of the label tokens as @@ -112,7 +132,18 @@ class NextTokenStrategy(ReadStrategy): def __init__(self, context: ReadContext): super().__init__(context) - tokenizer = context.online_renderer.renderer.get_tokenizer() + # Every prompt ends with the same generation prompt, so a label's token + # is the same for every question. + prompt = self._generation_prompt() + self.tail = prompt_tail(self._tokenizer(), prompt) + self.label_ids = type_label_ids(self._tokenizer(), prompt) + + def _tokenizer(self) -> TokenizerLike: + return self.context.online_renderer.renderer.get_tokenizer() + + def _generation_prompt(self) -> list[int]: + """A chat prompt's ids up to the generation prompt, with thinking off.""" + context, tokenizer = self.context, self._tokenizer() probe = tokenizer.apply_chat_template( [{"role": "user", "content": "x"}], **{ @@ -126,13 +157,20 @@ def __init__(self, context: ReadContext): ) if isinstance(probe, str): probe = tokenizer.encode(probe, add_special_tokens=False) - # Every prompt ends with the same generation prompt, so a label's token - # is the same for every question. - self.tail, _ = reply_label_ids(tokenizer, probe, ()) - self.label_ids: dict[str, int] = {} - for qtype in QUESTION_TYPES.values(): - _, ids = reply_label_ids(tokenizer, probe, qtype.label_set) - self.label_ids.update(zip(qtype.label_set, ids)) + return list(probe) + + def _sampling_params( + self, label_ids: list[int], prompt_ids: list[int] + ) -> SamplingParams: + return SamplingParams( + max_tokens=1, temperature=0.0, logprob_token_ids=label_ids + ) + + def _read_input( + self, engine_input: EngineInput, prompt_ids: list[int] + ) -> EngineInput: + """The engine input for a read of the rendered prompt.""" + return engine_input def limits(self) -> DecisionLimits: return DecisionLimits(max_questions=64, max_options=len(LABELS)) @@ -180,7 +218,7 @@ async def read( ctx = self.context read_request = self._read_request(chat_template_kwargs, cache_salt) - slots, engine_inputs = [], [] + slots, engine_inputs, params = [], [], [] for q in questions: messages = [{"role": "user", "content": f"{state}\n\n{q.type.prompt(q)}"}] if instructions: @@ -192,15 +230,13 @@ async def read( "labels' tokens are unknown" ) slots.append([self.label_ids[label] for label in q.labels]) - engine_inputs.append(engine_input) + engine_inputs.append(self._read_input(engine_input, prompt_ids)) + params.append(self._sampling_params(slots[-1], prompt_ids)) label_reads = await next_token_label_reads( ctx.engine_client, engine_inputs, - [ - SamplingParams(max_tokens=1, temperature=0.0, logprob_token_ids=ids) - for ids in slots - ], + params, request_id, lora_request=lora_request, trace_headers=trace_headers, @@ -223,13 +259,87 @@ async def read( return reads -NEXT_TOKEN_ARCHITECTURES = frozenset( - { - "Qwen3ForCausalLM", - "Qwen3_5ForConditionalGeneration", - "Qwen3_5MoeForConditionalGeneration", - } -) +class DiffusionGemmaCanvasStrategy(NextTokenStrategy): + """DiffusionGemma. Each question is one read-only request with one denoising + step.""" + + THOUGHT = "<|channel>thought\n" + END_OF_TURN = "" + CANVAS_STEP = 16 + + def __init__(self, context: ReadContext): + ReadStrategy.__init__(self, context) + tokenizer = self._tokenizer() + self.tail = prompt_tail(tokenizer, self._generation_prompt()) + self.thought = tokenizer.encode(self.THOUGHT, add_special_tokens=False) + self.label_ids = type_label_ids(tokenizer, self.thought) + end = tokenizer.encode(self.END_OF_TURN, add_special_tokens=False) + self.pad = tokenizer.pad_token_id + if self.pad is None or len(end) != 1: + raise ValueError( + "Structured decisions need the model's " + f"{self.END_OF_TURN!r} and pad tokens" + ) + self.end = end[0] + vllm_config = context.engine_client.vllm_config + served = ( + vllm_config.diffusion_config.canvas_length + if vllm_config.diffusion_config + and vllm_config.diffusion_config.canvas_length + else vllm_config.model_config.hf_config.canvas_length + ) + self.width = self.CANVAS_STEP + if self.width > served: + raise ValueError( + f"Structured decisions need a canvas of {self.width}, " + f"and the served canvas is {served}" + ) + self.vocab_size = vllm_config.model_config.get_vocab_size() + self.max_model_len = vllm_config.model_config.max_model_len + + def _sampling_params( + self, label_ids: list[int], prompt_ids: list[int] + ) -> SamplingParams: + # The noise at the label slot is drawn from the prompt, so a repeated + # read gets the same canvas. + rng = random.Random(hashlib.sha256(json.dumps(prompt_ids).encode()).digest()) + canvas = [rng.randrange(self.vocab_size), self.end] + canvas += [self.pad] * (self.width - len(canvas)) + return SamplingParams( + max_tokens=2, + logprob_token_ids=label_ids, + extra_args={ + "diffusion_seed_canvas": canvas, + "diffusion_canvas_length": self.width, + "diffusion_max_steps": 1, + "diffusion_read_only": True, + }, + ) + + def _read_input( + self, engine_input: EngineInput, prompt_ids: list[int] + ) -> EngineInput: + extra = len(self.thought) + self.width + if len(prompt_ids) + extra > self.max_model_len: + raise StructuredDecisionError( + f"the prompt has {len(prompt_ids)} tokens and a canvas read adds " + f"{extra}, more than max_model_len={self.max_model_len}" + ) + salt = engine_input.get("cache_salt") + # The thought block _must_ appear outside of the canvas (far worse results + # otherwise) + return tokens_input( + prompt_ids + self.thought, + cache_salt=salt if isinstance(salt, str) else None, + ) + + +READ_STRATEGIES: dict[str, type[ReadStrategy]] = { + "Qwen3ForCausalLM": NextTokenStrategy, + "Qwen3_5ForConditionalGeneration": NextTokenStrategy, + "Qwen3_5MoeForConditionalGeneration": NextTokenStrategy, + "DiffusionGemmaForBlockDiffusion": DiffusionGemmaCanvasStrategy, +} # label_mass sums the labels' full-vocabulary probabilities. The logprobs modes @@ -239,15 +349,16 @@ async def read( def select_read_strategy(model_config: ModelConfig) -> type[ReadStrategy]: """Raise ValueError when the model cannot serve structured decisions.""" - if model_config.architecture not in NEXT_TOKEN_ARCHITECTURES: + strategy = READ_STRATEGIES.get(model_config.architecture) + if strategy is None: raise ValueError( "The structured decisions API does not support " f"{model_config.architecture}. Supported architectures: " - f"{sorted(NEXT_TOKEN_ARCHITECTURES)}" + f"{sorted(READ_STRATEGIES)}" ) if model_config.logprobs_mode not in LOGPROBS_MODES: raise ValueError( "Structured decisions need --logprobs-mode raw_logprobs or " f"processed_logprobs, not {model_config.logprobs_mode}" ) - return NextTokenStrategy + return strategy From 193922d6ff5f7e213eacadd3cb7dd96612b30ed3 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Fri, 9 Oct 2026 21:53:18 -0400 Subject: [PATCH 217/624] [ROCm][Perf][DSv4.1] Mono decode layer for DeepSeek-V4.1-Flash on MI355X (#60397) Signed-off-by: Fangzhou Ai Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/models_basic.yaml | 20 +- tests/models/test_dsv41_mono_decode.py | 317 ++++ tests/models/test_dsv41_mono_layers.py | 335 ++++ tests/models/test_dsv41_mono_numerics.py | 529 ++++++ vllm/envs.py | 7 + vllm/models/deepseek_v41/amd/model.py | 17 + vllm/models/deepseek_v41/amd/mono/README.md | 94 + vllm/models/deepseek_v41/amd/mono/__init__.py | 5 + .../amd/mono/attention/__init__.py | 4 + .../deepseek_v41/amd/mono/attention/back.py | 884 +++++++++ .../deepseek_v41/amd/mono/attention/device.py | 504 +++++ .../deepseek_v41/amd/mono/attention/front.py | 662 +++++++ .../deepseek_v41/amd/mono/attention/gemv.py | 234 +++ .../deepseek_v41/amd/mono/attention/plan.py | 173 ++ .../deepseek_v41/amd/mono/common/__init__.py | 9 + .../models/deepseek_v41/amd/mono/common/mx.py | 78 + .../deepseek_v41/amd/mono/common/ops.py | 300 +++ .../amd/mono/common/peer_memory.py | 104 ++ .../deepseek_v41/amd/mono/common/plan.py | 135 ++ .../deepseek_v41/amd/mono/common/sync.py | 177 ++ vllm/models/deepseek_v41/amd/mono/layer.py | 940 ++++++++++ vllm/models/deepseek_v41/amd/mono/runner.py | 437 +++++ vllm/models/deepseek_v41/amd/mono/sources.py | 8 + .../deepseek_v41/amd/mono/stages/__init__.py | 4 + .../deepseek_v41/amd/mono/stages/dims.py | 60 + .../deepseek_v41/amd/mono/stages/gemv.py | 294 +++ .../deepseek_v41/amd/mono/stages/moe.py | 1635 +++++++++++++++++ .../deepseek_v41/amd/mono/stages/moe_plan.py | 54 + .../deepseek_v41/amd/mono/stages/moe_shape.py | 229 +++ .../deepseek_v41/amd/mono/stages/router.py | 86 + .../deepseek_v41/amd/mono/stages/seam.py | 406 ++++ vllm/models/deepseek_v41/amd/mono_decode.py | 376 ++++ vllm/models/deepseek_v41/amd/rocm.py | 4 +- 33 files changed, 9119 insertions(+), 2 deletions(-) create mode 100644 tests/models/test_dsv41_mono_decode.py create mode 100644 tests/models/test_dsv41_mono_layers.py create mode 100644 tests/models/test_dsv41_mono_numerics.py create mode 100644 vllm/models/deepseek_v41/amd/mono/README.md create mode 100644 vllm/models/deepseek_v41/amd/mono/__init__.py create mode 100644 vllm/models/deepseek_v41/amd/mono/attention/__init__.py create mode 100644 vllm/models/deepseek_v41/amd/mono/attention/back.py create mode 100644 vllm/models/deepseek_v41/amd/mono/attention/device.py create mode 100644 vllm/models/deepseek_v41/amd/mono/attention/front.py create mode 100644 vllm/models/deepseek_v41/amd/mono/attention/gemv.py create mode 100644 vllm/models/deepseek_v41/amd/mono/attention/plan.py create mode 100644 vllm/models/deepseek_v41/amd/mono/common/__init__.py create mode 100644 vllm/models/deepseek_v41/amd/mono/common/mx.py create mode 100644 vllm/models/deepseek_v41/amd/mono/common/ops.py create mode 100644 vllm/models/deepseek_v41/amd/mono/common/peer_memory.py create mode 100644 vllm/models/deepseek_v41/amd/mono/common/plan.py create mode 100644 vllm/models/deepseek_v41/amd/mono/common/sync.py create mode 100644 vllm/models/deepseek_v41/amd/mono/layer.py create mode 100644 vllm/models/deepseek_v41/amd/mono/runner.py create mode 100644 vllm/models/deepseek_v41/amd/mono/sources.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/__init__.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/dims.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/gemv.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/moe.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/moe_plan.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/moe_shape.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/router.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/seam.py create mode 100644 vllm/models/deepseek_v41/amd/mono_decode.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 1d9a7fe2f3b1..483a9aceac22 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -91,8 +91,9 @@ steps: - tests/models/test_deepseek_v4_dspark_rocm.py - tests/models/test_deepseek_v4_rocm_wo_a.py - tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py + - tests/models/test_dsv41_mono_decode.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.py - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - label: ":nvidia: (B200) Inkling" @@ -259,3 +260,20 @@ steps: device: cpu-small commands: - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ model_executor/layers/test_activation.py models/test_language_model_cache_is_weak.py models/deepseek_v32/test_sequence_parallel.py models/test_deepseek_v4_fi_moe_ep.py models/test_dspark_mla.py models/transformers/test_layer_registry.py + +- label: ':amd: (MI355) DeepSeek-V4.1 Mono Decode Numerics' + key: dsv41-mono-decode-numerics-amd + dind: false + device: mi355_4 + num_devices: 4 + timeout_in_minutes: 40 + working_dir: /vllm-workspace/ + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/models/deepseek_v41/amd/ + - tests/models/test_dsv41_mono_numerics.py + - tests/models/test_dsv41_mono_layers.py + commands: + - pytest -v -s tests/models/test_dsv41_mono_numerics.py + - pytest -v -s tests/models/test_dsv41_mono_layers.py diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_mono_decode.py new file mode 100644 index 000000000000..684e72c00482 --- /dev/null +++ b/tests/models/test_dsv41_mono_decode.py @@ -0,0 +1,317 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The DeepSeek-V4.1 mono decode layer runs only where its kernels are exact: +decode-only steps of at most MAX_ROWS rows with causal SWA windows, eager or in +a FULL graph. A PIECEWISE capture is replayed for mixed batches, so it must +never record the mono launches. Layers outside the whole-layer kernels run the +FFN launch on the same steps, from wo_b's unreduced output.""" + +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import torch + +from vllm.config import CUDAGraphMode +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend +from vllm.models.deepseek_v41.amd import mono_decode as md + +M = 6 + + +class _Runner: + def __init__(self): + self.calls = [] + + def forward(self, *args, **kwargs): + self.calls.append((args, kwargs)) + return ("mono outputs",) + + def ffn(self, *args): + self.calls.append((args, {})) + return ("ffn outputs",) + + +def _layer(comp="comp"): + cache = lambda: torch.zeros(4, 128 * md.RECORD, dtype=torch.uint8) # noqa: E731 + attn = SimpleNamespace( + swa_cache_layer=SimpleNamespace(prefix="swa", kv_cache=cache()), + compressed_cache_prefix=comp, + topk_indices_buffer=torch.zeros(64, 512, dtype=torch.int32), + compress_ratio=1, + _compressed_kv_cache=cache, + ) + return SimpleNamespace(attn=attn) + + +def _swa(**overrides): + fields = dict( + num_prefills=0, + num_decodes=1, + num_decode_tokens=M, + decode_swa_width=md.SWA_WIDTH, + decode_swa_indices=torch.zeros(64, 1, md.SWA_WIDTH, dtype=torch.int32), + decode_swa_lens=torch.zeros(64, dtype=torch.int32), + token_to_req_indices=torch.zeros(64, dtype=torch.int32), + slot_mapping=torch.zeros(64, dtype=torch.int64), + block_size=32, + ) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def _inputs(rows=M): + return dict( + x=torch.zeros(rows, 5120, dtype=torch.bfloat16), + positions=torch.zeros(rows, dtype=torch.int64), + residual=torch.zeros(rows, 4, 5120, dtype=torch.bfloat16), + post_mix=torch.zeros(rows, 4, 1), + res_mix=torch.zeros(rows, 4, 4), + pre_mix=torch.zeros(rows, 4), + ) + + +@pytest.fixture +def run(monkeypatch): + """Call the mono path for one layer under a forward context built from + ``mode`` and ``metadata``; returns (result, runner calls).""" + + def _run( + mode=CUDAGraphMode.FULL, metadata=None, ffn_only=False, entry=None, **inputs + ): + if metadata is None: + metadata = { + "swa": _swa(), + "comp": SimpleNamespace(block_size=128, block_table=None), + } + fc = SimpleNamespace(cudagraph_runtime_mode=mode, attn_metadata=metadata) + runner = _Runner() + monkeypatch.setattr(md, "is_forward_context_available", lambda: True) + monkeypatch.setattr(md, "get_forward_context", lambda: fc) + monkeypatch.setattr(md, "_mono_runner", lambda device: runner) + monkeypatch.setattr(md.MonoDecodeLayer, "weights", lambda self, layer: None) + args = {**_inputs(), **inputs} + mono = md.MonoDecodeLayer(ffn_only) + if (entry or ("ffn" if ffn_only else "layer")) == "ffn": + # x stands for wo_b's unreduced output; the first layer's attention + # has no compressed cache + part, _ = args.pop("x"), args.pop("positions") + return mono.ffn(_layer(comp=None), part, **args), runner.calls + return mono(_layer(), **args), runner.calls + + return _run + + +@pytest.mark.parametrize("mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL]) +def test_decode_step_runs_mono(run, mode): + out, calls = run(mode) + assert out == ("mono outputs",) and len(calls) == 1 + args, kwargs = calls[0] + # the caches reach the kernels as [blocks, block, record] bytes + assert args[8].shape == (4, 32, md.RECORD) and args[8].stride() == ( + 128 * md.RECORD, + md.RECORD, + 1, + ) + assert kwargs["comp_cache"].shape == (4, 128, md.RECORD) + + +def test_padded_decode_step_runs_mono(run): + """A FULL graph pads the step to its capture size: the kernels take it and + skip the pad rows (their slot is -1).""" + swa = _swa(num_decode_tokens=M - 2) + comp = SimpleNamespace(block_size=128, block_table=None) + out, calls = run(metadata={"swa": swa, "comp": comp}) + assert out == ("mono outputs",) and len(calls) == 1 + + +@pytest.mark.parametrize( + "case", + [ + dict(mode=CUDAGraphMode.PIECEWISE), + dict( + metadata={ + "swa": _swa(num_prefills=1), + "comp": SimpleNamespace(block_size=128), + } + ), + dict( + metadata={ + "swa": _swa(decode_swa_width=256), + "comp": SimpleNamespace(block_size=128), + } + ), + dict( + metadata={ + "swa": _swa(num_decode_tokens=M + 1), + "comp": SimpleNamespace(block_size=128), + } + ), + dict(metadata={}), # a profile / dummy run: no attention metadata + dict(**_inputs(md.MAX_ROWS + 6)), + dict(residual=None), # the first layer's seam broadcasts the embedding + ], + ids=[ + "piecewise", + "prefill", + "noncausal-window", + "tokens-past-rows", + "no-metadata", + "rows", + "first-layer", + ], +) +def test_other_steps_take_the_original_path(run, case): + out, calls = run(**case) + assert out is None and not calls + + +def _decoder_layer(engram=None, **attn): + attn = ( + dict( + layer_id=21, + compressor=None, + indexer=None, + compress_ratio=1, + kv_mxfp8=False, + kv_cache_dtype="fp8_ds_mla", + ) + | attn + ) + backend = Mxfp4MoeBackend.AITER_MXFP4_BF16 + quant_method = SimpleNamespace(mxfp4_backend=backend) + ffn = SimpleNamespace( + experts=SimpleNamespace( + routed_experts=SimpleNamespace(quant_method=quant_method) + ), + shared_experts=object(), + gate=SimpleNamespace(tid2eid=None), + n_routed_experts=384, + n_activated_experts=6, + routed_scaling_factor=1.5, + swiglu_limit=10.0, + scoring_func="sqrtsoftplus", + renormalize=True, + ) + return SimpleNamespace( + attn=SimpleNamespace(**attn), + ffn=ffn, + engram=engram, + use_sequence_parallel=False, + fuse_seam_norm=True, + ) + + +@pytest.mark.parametrize( + "layer, path", + [ + (_decoder_layer(), "whole"), + (_decoder_layer(compress_ratio=0), "ffn"), + (_decoder_layer(engram=object()), "ffn"), + (_decoder_layer(compressor=object(), indexer=object()), "ffn"), + (_decoder_layer(indexer=object()), "ffn"), + (_decoder_layer(kv_cache_dtype="auto"), "ffn"), + (_decoder_layer(layer_id=40), None), + ], + ids=["standard", "first", "engram", "kv-source", "index-source", "kv", "draft"], +) +def test_layers_take_whole_or_ffn_launch(monkeypatch, layer, path): + """The whole-layer kernels serve the standard layers; every other backbone + layer keeps vLLM's attention and runs its FFN half as one launch.""" + _deployment(monkeypatch) + mono = md.MonoDecodeLayer.create(layer, _config()) + got = None if mono is None else "ffn" if mono.ffn_only else "whole" + assert got == path + + +@pytest.mark.parametrize("mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL]) +def test_decode_step_runs_ffn_launch(run, mode): + out, calls = run(mode, ffn_only=True, metadata={"swa": _swa()}) + assert out == ("ffn outputs",) and len(calls) == 1 + assert calls[0][0][1].shape == (M, 5120) # wo_b's unreduced output + + +def test_ffn_layer_never_runs_whole(run): + out, calls = run(ffn_only=True, entry="layer") + assert out is None and not calls + + +@pytest.mark.parametrize( + "case", + [ + dict(mode=CUDAGraphMode.PIECEWISE), + dict(metadata={"swa": _swa(num_prefills=1)}), + dict(metadata={"swa": _swa(num_decode_tokens=M + 1)}), + dict(metadata={}), + dict(**_inputs(md.MAX_ROWS + 6)), + ], + ids=["piecewise", "prefill", "tokens-past-rows", "no-metadata", "rows"], +) +def test_other_steps_reduce_wo_b_themselves(run, case): + """None: the layer all-reduces wo_b's partial and runs the FFN as before.""" + case.setdefault("metadata", {"swa": _swa()}) + out, calls = run(ffn_only=True, **case) + assert out is None and not calls + + +def _config(**parallel): + parallel = ( + dict(enable_expert_parallel=False, enable_eplb=False, data_parallel_size=1) + | parallel + ) + return SimpleNamespace( + model_config=SimpleNamespace(hf_config=SimpleNamespace(num_hidden_layers=40)), + parallel_config=SimpleNamespace(**parallel), + ) + + +def _deployment(monkeypatch, cdna=4, cus=256): + """A deployment the kernels take (the runner module stubbed: no FlyDSL).""" + import vllm.platforms.rocm as rocm + + monkeypatch.setenv("VLLM_ROCM_MONO_DECODE", "1") + monkeypatch.delenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", raising=False) + monkeypatch.setattr(rocm, "get_cdna_version", lambda: cdna) + monkeypatch.setattr(md, "get_tensor_model_parallel_world_size", lambda: 2) + monkeypatch.setattr(md, "_compute_units", lambda: cus) + name = "vllm.models.deepseek_v41.amd.mono.runner" + runner = ModuleType(name) + runner.BLOCKS = 256 # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, name, runner) + + +@pytest.mark.parametrize("cdna", [3, 5]) +def test_opt_in_outside_the_kernels_cdna_raises(monkeypatch, cdna): + _deployment(monkeypatch, cdna=cdna) + with pytest.raises(ValueError, match="needs CDNA4"): + md.MonoDecodeLayer.create(_decoder_layer(), _config()) + + +def _other_backend(): + layer = _decoder_layer() + experts = layer.ffn.experts.routed_experts + experts.quant_method.mxfp4_backend = Mxfp4MoeBackend.TRITON + return layer + + +@pytest.mark.parametrize( + "case, match", + [ + (dict(parallel=dict(enable_expert_parallel=True)), "no expert"), + (dict(parallel=dict(data_parallel_size=2)), "no expert"), + (dict(parallel=dict(enable_eplb=True)), "no expert"), + (dict(a4w4=True), "A8W4"), + (dict(layer=_other_backend), "AITER_MXFP4_BF16"), + (dict(cus=128), "256 compute units"), + ], + ids=["expert-parallel", "data-parallel", "eplb", "a4w4", "moe-backend", "cus"], +) +def test_opt_in_outside_the_kernels_moe_raises(monkeypatch, case, match): + """The kernels read all 384 experts on each rank in AITER's A8W4 layout, with + all 256 CTAs resident: any other deployment is an error, not garbage.""" + _deployment(monkeypatch, cus=case.get("cus", 256)) + if case.get("a4w4"): + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", "1") + layer = case.get("layer", _decoder_layer)() + with pytest.raises(ValueError, match=match): + md.MonoDecodeLayer.create(layer, _config(**case.get("parallel", {}))) diff --git a/tests/models/test_dsv41_mono_layers.py b/tests/models/test_dsv41_mono_layers.py new file mode 100644 index 000000000000..580462a82394 --- /dev/null +++ b/tests/models/test_dsv41_mono_layers.py @@ -0,0 +1,335 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The DeepSeek-V4.1 mono decode path through the real model: a five-layer +DeepSeek-V4.1-Flash at TP2 and TP4 on CDNA4, served with ``VLLM_ROCM_MONO_DECODE`` on +and off on the same random weights: every decoder layer's outputs at every +decode step must match. Unlike ``test_dsv41_mono_numerics.py``, which +checks the kernels against vLLM's ops one stage at a time, this runs +``DeepseekV4DecoderLayer.forward`` as the engine does -- the model's hooks, +the per-step path choice, the weights taken off the layer, the metadata and +the indexer's top-k -- so a change to the decoder layer that the kernels do +not follow fails here. + +The five layers hold both kinds the mono path takes: +- 0: sliding window only -> the FFN launch +- 1: ratio-2 KV and index source -> the FFN launch +- 2: ratio-2 consumer -> the whole layer (K1 + K2) +- 3: ratio-1 KV, index and candidate source -> the FFN launch +- 4: ratio-1 consumer -> the whole layer +""" + +import zlib + +import pytest +import torch + +from vllm.model_executor.model_loader import register_model_loader +from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader +from vllm.platforms import current_platform + +MODEL = "deepseek-ai/DeepSeek-V4.1-Flash" +LOAD_FORMAT = "dsv41_mono_random" +TRUNCATED = { + "num_hidden_layers": 5, + # one entry a layer, the three MTP layers' included + "compress_ratios": [0, 2, 2, 1, 1, 0, 0, 0], + "kv_source_layer_ids": [1, 3], + "index_source_layer_ids": [1, 3], + "candidate_source_layer_id": 3, + "engram_layer_ids": [], + "engram_num_embeddings": [], + "dspark_target_layer_ids": [2, 3, 4], +} +WHOLE, FFN = (2, 4), (0, 1, 3) +STEPS = 5 # decode steps after the prefill's token + + +def _on_cdna4() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import get_cdna_version + + return get_cdna_version() == 4 + + +@torch.no_grad() +def _fill(name: str, t: torch.Tensor) -> None: + """A random value for every weight, seeded by its name (alike on every + rank and in every run), in the ranges the checkpoint's formats hold.""" + if t.device.type == "meta": + return # materialized and filled by the dummy loader's own path + g = torch.Generator(device=t.device).manual_seed(zlib.crc32(name.encode())) + + def randn(scale=1.0, shift=0.0): + return torch.randn(t.shape, generator=g, device=t.device) * scale + shift + + if t.dtype == torch.uint8 and "scale" in name: + # E8M0 2^-12 .. 2^-8, alike over each 32 rows as in the checkpoint (the + # loader folds per-row scales into 32 x 32 blocks only then). Larger + # random layers make the stack chaotic: atomics' order alone moves + # vLLM's own logits run to run. + rows = t.shape[-2] if t.dim() > 1 else 1 + shape = (*t.shape[:-2], -(-rows // 32), t.shape[-1]) if t.dim() > 1 else t.shape + e8m0 = torch.randint(115, 120, shape, generator=g, device=t.device) + if t.dim() > 1: + e8m0 = e8m0.repeat_interleave(32, dim=-2)[..., :rows, :] + t.copy_(e8m0) + elif t.dtype in (torch.uint8, torch.float4_e2m1fn_x2): + u8 = t.view(torch.uint8) + u8.copy_(torch.randint(0, 256, u8.shape, generator=g, device=t.device)) + elif t.dtype == torch.float8_e4m3fn: + t.copy_(randn().clamp(-6, 6)) + elif not torch.is_floating_point(t): + return + elif "norm" in name: + t.copy_(randn(0.1, 1.0)) + elif name.endswith(("hc_attn_scale", "hc_ffn_scale")): + t.copy_(randn(0.5).abs() + 0.5) + elif name.endswith(("hc_attn_base", "hc_ffn_base", "attn_sink")): + t.copy_(randn(0.1)) + elif name.endswith("embed_tokens.weight"): + t.copy_(randn()) + else: + t.copy_(randn(0.02)) + + +@register_model_loader(LOAD_FORMAT) +class RandomWeightLoader(DummyModelLoader): + """Dummy loading with weights in the checkpoint's ranges, before post-load + processing, so every processed copy follows (plain dummy weights leave the + E8M0 scales zero, and both paths would agree on zeros).""" + + def load_weights(self, model, model_config) -> None: + from vllm.distributed import get_tensor_model_parallel_world_size + + super().load_weights(model, model_config) + for name, t in model.state_dict().items(): + _fill(name, t) + # the routed experts' intermediate past this rank's share is the + # loader's zero padding (576 -> 640 at TP4): the kernels skip it + inter = model_config.hf_text_config.moe_intermediate_size + inter //= get_tensor_model_parallel_world_size() + for m in model.modules(): + w13, w2 = getattr(m, "w13_weight", None), getattr(m, "w2_weight", None) + if w13 is None or w2 is None or w13.device.type == "meta": + continue + padded = w13.shape[1] // 2 + if padded > inter: + w13, w2 = w13.data.view(torch.uint8), w2.data.view(torch.uint8) + for half in (0, padded): + w13[:, half + inter : half + padded] = 0 + w2[..., inter // 2 :] = 0 + + +class MonoProbe: + """Worker extension (``worker_extension_cls``): it also brings this module, + and so the loader, into the workers.""" + + def _decoder_layers(self) -> dict: + from vllm.models.deepseek_v41.amd.model import DeepseekV4DecoderLayer + + model = self.model_runner.model # type: ignore[attr-defined] + return { + m.attn.layer_id: m + for m in model.modules() + if isinstance(m, DeepseekV4DecoderLayer) + } + + def capture_steps(self, lens: list[int]) -> None: + """Keep each decoder layer's outputs (x, residual) for every decode + token taken in a decode-only step (the steps the mono path takes), + keyed by (layer, prompt length, position): a request's step and row + can differ between engines (requests reach the engine core while it + runs), its prompt length and position cannot. ``lens``: the prompt + lengths, all distinct; a row's request is the longest prompt not past + its position.""" + from vllm.forward_context import get_forward_context + + self._captured: dict = {} + lens = sorted(lens) + + def hook(layer_id): + def keep(module, args, kwargs, out): + md = get_forward_context().attn_metadata + if not isinstance(md, dict): + return + swa = md.get(module.attn.swa_cache_layer.prefix) + if swa is None or swa.num_prefills != 0: + return + pos = kwargs.get("positions", args[1] if len(args) > 1 else None) + for row, p in enumerate(pos.tolist()): + n = max(n for n in lens if n <= p) + kept = [t[row].float().cpu() for t in out[:2]] + self._captured.setdefault((layer_id, n, p), kept) + + return keep + + def gap_hook(layer_id, layer): + def keep(ffn, args, kwargs): + # vLLM's path only (the mono kernels route inside): each row's + # top-6 margin, sqrt(softplus(logits)) + bias as vLLM routes + md = get_forward_context().attn_metadata + if not isinstance(md, dict): + return + swa = md.get(layer.attn.swa_cache_layer.prefix) + if swa is None or swa.num_prefills != 0: + return + h = args[0] if args else kwargs["hidden_states"] + logits = h.float() @ ffn.gate.weight.float().T + score = torch.nn.functional.softplus(logits).sqrt() + score = score + ffn.gate.e_score_correction_bias.float() + top = score.topk(7, dim=-1).values + gap = (top[:, 5] - top[:, 6]).tolist() + pos = self._positions + for row, p in enumerate(pos.tolist()): + n = max(n for n in lens if n <= p) + self._gaps.setdefault((layer_id, n, p), gap[row]) + + return keep + + def pos_hook(module, args, kwargs): + self._positions = kwargs.get( + "positions", args[1] if len(args) > 1 else None + ) + + self._gaps: dict = {} + for i, layer in self._decoder_layers().items(): + layer.register_forward_hook(hook(i), with_kwargs=True) + layer.register_forward_pre_hook(pos_hook, with_kwargs=True) + gap = gap_hook(i, layer) + layer.ffn.register_forward_pre_hook(gap, with_kwargs=True) + + def save_capture(self, path: str) -> None: + rank = self.rank # type: ignore[attr-defined] + torch.save((self._captured, self._gaps), f"{path}.rank{rank}") + + def mono_layers(self) -> dict: + from vllm.models.deepseek_v41.amd import mono_decode + + layers = { + i: None + if m.mono is None + else ["ffn" if m.mono.ffn_only else "whole", m.mono._weights is not None] + for i, m in self._decoder_layers().items() + } + runner = mono_decode._runner + return { + "layers": layers, + "epoch": -1 if runner is None else int(runner.epoch[0]), + } + + +def _decode(vllm_runner, monkeypatch, tp: int, mono: bool, prompts, path: str): + """Greedy decode of ``prompts`` with the mono path on or off: (each + prompt's output, the workers' mono layers); each decoder layer's outputs + for the decode tokens of decode-only steps, and vLLM's routing margins, + saved under ``path``.""" + from vllm import SamplingParams + from vllm.inputs import TokensPrompt + + monkeypatch.setenv("VLLM_ROCM_MONO_DECODE", "1" if mono else "0") + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + with vllm_runner( + MODEL, + tensor_parallel_size=tp, + load_format=LOAD_FORMAT, + hf_overrides=TRUNCATED, + worker_extension_cls=f"{__name__}.MonoProbe", + moe_backend="aiter", + language_model_only=True, + tokenizer_mode="deepseek_v41", + attention_config={"indexer_kv_dtype": "mxfp4", "indexer_sparse_logits": True}, + block_size=128, + enforce_eager=True, + max_model_len=4096, + max_num_seqs=len(prompts), + # a prompt prefills in one step, so its decode tokens never share one + max_num_batched_tokens=16384, + enable_prefix_caching=False, + gpu_memory_utilization=0.5, + ) as runner: + runner.llm.collective_rpc("capture_steps", args=([len(p) for p in prompts],)) + params = SamplingParams(temperature=0.0, max_tokens=STEPS + 1) + outs = runner.llm.generate( + [TokensPrompt(prompt_token_ids=p) for p in prompts], params + ) + runner.llm.collective_rpc("save_capture", args=(path,)) + probe = runner.llm.collective_rpc("mono_layers") + return [o.outputs[0] for o in outs], probe + + +def _rel(a: torch.Tensor, b: torch.Tensor) -> float: + return ((a - b).norm() / b.norm()).item() + + +@pytest.mark.skipif(not _on_cdna4(), reason="the mono kernels target CDNA4") +@pytest.mark.parametrize("tp", [2, 4]) +def test_mono_decode_matches_the_model(vllm_runner, monkeypatch, tmp_path, tp): + if torch.accelerator.device_count() < tp: + pytest.skip(f"needs {tp} GPUs") + pytest.importorskip("flydsl") + g = torch.Generator().manual_seed(0) + # contexts past both ratios' 512 compressed entries, so top-k selects + lens = torch.randperm(1700, generator=g)[:8].add(300).tolist() # distinct + prompts = [torch.randint(1000, 60000, (n,), generator=g).tolist() for n in lens] + + on, probe_on = _decode( + vllm_runner, monkeypatch, tp, True, prompts, f"{tmp_path}/on" + ) + off, probe_off = _decode( + vllm_runner, monkeypatch, tp, False, prompts, f"{tmp_path}/off" + ) + + # both layer kinds took the mono path on every rank, at every decode step + launches = len(WHOLE + FFN) # one launch (pair) a layer a step + for rank in probe_on: + assert rank["layers"] == { + i: ["whole" if i in WHOLE else "ffn", True] for i in WHOLE + FFN + }, rank + assert rank["epoch"] >= launches * STEPS, rank + assert rank["epoch"] % launches == 0, rank + assert all(set(r["layers"].values()) == {None} for r in probe_off) + # each decode token's input matches between the runs while the greedy + # tokens before it do: (prompt length, position) pairs to compare + same: set[tuple[int, int]] = set() + for prompt, a, b in zip(prompts, on, off): + for step in range(STEPS): + if a.token_ids[: step + 1] != b.token_ids[: step + 1]: + break + same.add((len(prompt), len(prompt) + step)) + + # each layer's outputs for those tokens, mono against vLLM's path, token by + # token. Layer 0 sees the same inputs on both paths at a request's first + # decode token: there the bound is test_dsv41_mono_numerics' (about 3x the + # kernels' 0.4%); its later tokens attend over KV each path wrote, 3%. Each + # later layer takes the previous one's slightly different output, which + # random weights amplify (2-4% by layer 4): 8%. A token whose top-6 routing + # is a near-tie on vLLM's path (6th / 7th score within 5e-3: a 0.4% input + # difference moves scores near 0.8 by about 3e-3) may pick another expert; + # once it has, its residual stream differs and the later layers are exempt + # for it. The residual stream: 3% (0.3% seen). A layer the kernels no + # longer follow is off by far more, on every token. + for rank in range(tp): + got, _ = torch.load(f"{tmp_path}/on.rank{rank}") + want, gaps = torch.load(f"{tmp_path}/off.rank{rank}") + keys: list[tuple[int, int]] = sorted( + same & {k[1:] for k in got} & {k[1:] for k in want} + ) + assert len(keys) >= len(prompts), keys + # the bound that sees a 3% error: a request's first decode token + assert any(p == n for n, p in keys), keys + flipped: set = set() + for i in sorted(WHOLE + FFN): + for k in keys: + if k in flipped: + continue + (x, res), (x_ref, res_ref) = got[(i, *k)], want[(i, *k)] + err = _rel(res, res_ref) + assert err < 0.03, (rank, i, k, "residual", err) + bound = (0.015 if k[1] == k[0] else 0.03) if i == 0 else 0.08 + err = _rel(x, x_ref) + if err >= bound: + assert gaps[(i, *k)] < 5e-3, (rank, i, k, "x", err) + flipped.add(k) + assert len(flipped) <= len(keys) // 4, (rank, sorted(flipped)) diff --git a/tests/models/test_dsv41_mono_numerics.py b/tests/models/test_dsv41_mono_numerics.py new file mode 100644 index 000000000000..f7cd877caf47 --- /dev/null +++ b/tests/models/test_dsv41_mono_numerics.py @@ -0,0 +1,529 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Numerics of the DeepSeek-V4.1 mono decode kernels at TP2 on CDNA4 against the +ops vLLM runs for the same decode step: random weights in the checkpoint's +formats, the step's KV caches as vLLM's insert writes them. + +- The whole layer (K1 + K2) against the mHC seams, the ROCm attention op chain, + the TP all-reduce and ``DeepseekV4MoE``. +- The FFN launch against the same FFN half, from wo_b's unreduced output. + +Each stage is compared on its own input, so a failure names the stage; the MoE +is checked on the kernels' own FFN input, where routing cannot differ.""" + +import pytest +import ray +import torch + +from vllm.platforms import current_platform + +from ..utils import ( + ensure_model_parallel_initialized, + init_test_distributed_environment, + multi_process_parallel, +) + +MODEL = "deepseek-ai/DeepSeek-V4.1-Flash" +HIDDEN, Q_RANK, HEAD_DIM, NOPE, ROPE = 5120, 1280, 512, 448, 64 +N_HEADS, O_GROUPS, O_RANK = 64, 8, 1024 +WINDOW, TOPK, SWA_BLOCK, CACHE_BLOCK = 128, 512, 32, 128 +RECORD, DATA, ALIGN = 584, 576, 576 +REPLAYS = 300 +HC = 4 + + +def _on_cdna4() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import get_cdna_version + + return get_cdna_version() == 4 + + +def _cos(a: torch.Tensor, b: torch.Tensor) -> float: + a, b = a.float().reshape(-1), b.float().reshape(-1) + return torch.nn.functional.cosine_similarity(a, b, dim=0).item() + + +def _rel(a: torch.Tensor, b: torch.Tensor) -> float: + """Relative L2 error of ``a`` against ``b``: unlike the cosine, it sees a + wrong scale.""" + a, b = a.float().reshape(-1), b.float().reshape(-1) + return ((a - b).norm() / b.norm()).item() + + +# ---------------------------------------------------------------- caches + + +def _paged_cache(num_blocks: int, block: int, device) -> torch.Tensor: + """[blocks, block, 584] u8 over pages of round_up(block * 584, 576) bytes, + the way vLLM binds a packed V4 cache.""" + page = -(-block * RECORD // ALIGN) * ALIGN + raw = torch.zeros(num_blocks * page, dtype=torch.uint8, device=device) + return torch.as_strided(raw, (num_blocks, block, RECORD), (page, RECORD, 1)) + + +def _rope(kv: torch.Tensor, positions: torch.Tensor, cos_sin: torch.Tensor): + """GPT-J RoPE of the last 64 dims (fp32 math, bf16 out).""" + cs = cos_sin[positions.long()] + cos, sin = cs[:, :32], cs[:, 32:] + x = kv.float() + e, o = x[:, NOPE::2].clone(), x[:, NOPE + 1 :: 2].clone() + x[:, NOPE::2] = e * cos - o * sin + x[:, NOPE + 1 :: 2] = e * sin + o * cos + return x.to(torch.bfloat16) + + +def _write_records(cache, block: int, slots, kv): + """fp8_ds_mla records of roped bf16 kv rows at ``slots``: the 448 nope dims + in e4m3 with a UE8M0 per 64, the rope dims in bf16, as vLLM's insert.""" + nope = kv.float()[:, :NOPE].view(-1, 7, 64) + exp = torch.ceil(torch.log2(nope.abs().amax(-1).clamp_min(1e-4) / 448.0)) + q = (nope * torch.exp2(-exp)[..., None]).clamp(-448, 448).reshape(-1, NOPE) + data = torch.cat( + [ + q.to(torch.float8_e4m3fn).view(torch.uint8), + kv[:, NOPE:].contiguous().view(torch.uint8), + ], + 1, + ) + scales = torch.zeros(kv.shape[0], 8, dtype=torch.uint8, device=kv.device) + scales[:, :7] = (exp + 127).clamp(0, 255).to(torch.uint8) + n = cache.untyped_storage().nbytes() - cache.storage_offset() + flat = torch.as_strided(cache, (n,), (1,)) + slots = slots.to(kv.device).long() + b, o = slots // block, slots % block + d0 = b * cache.stride(0) + o * DATA + s0 = b * cache.stride(0) + block * DATA + o * 8 + ar = torch.arange(DATA, device=kv.device) + flat[(d0[:, None] + ar).reshape(-1)] = data.reshape(-1) + ar = torch.arange(8, device=kv.device) + flat[(s0[:, None] + ar).reshape(-1)] = scales.reshape(-1) + + +def _normed_kv(n: int, kv_norm, g, device): + x = torch.randn(n, HEAD_DIM, generator=g).to(device) + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True)) * kv_norm.float() + return x.to(torch.bfloat16) + + +def _decode_step(ratio: int, R: int, T: int, kv_norm, cos_sin, seed: int): + """R requests x T tokens past random contexts: the step's metadata as + vLLM's builders lay it out, its window rows and (ratio > 0) the compressed + cache filled, a random top-512 a token.""" + g = torch.Generator().manual_seed(seed) + dev = kv_norm.device + ctx = [int(c) for c in torch.randint(2000, 40000, (R,), generator=g)] + M = R * T + spans = [ + (max(0, c - WINDOW - T) // SWA_BLOCK, (c + T - 1) // SWA_BLOCK) for c in ctx + ] + n_swa = 1 + sum(hi - lo + 1 for lo, hi in spans) + perm = (torch.randperm(n_swa - 1, generator=g) + 1).tolist() + bt = torch.zeros(R, max(hi for _, hi in spans) + 1, dtype=torch.int64) + for r, (lo, hi) in enumerate(spans): + for lb in range(lo, hi + 1): + bt[r, lb] = perm.pop() + swa = _paged_cache(n_swa, SWA_BLOCK, dev) + + def slot(r, p): + return int(bt[r, p // SWA_BLOCK]) * SWA_BLOCK + p % SWA_BLOCK + + pos = [ctx[r] + t for r in range(R) for t in range(T)] + t2r = [r for r in range(R) for _ in range(T)] + swa_idx = torch.full((M, WINDOW), -1, dtype=torch.int32) + swa_len = torch.zeros(M, dtype=torch.int32) + for i in range(M): + ps = range(max(pos[i] - WINDOW + 1, 0), pos[i] + 1) + swa_idx[i, : len(ps)] = torch.tensor([slot(t2r[i], p) for p in ps]) + swa_len[i] = len(ps) + hist = [ + (r, p) for r in range(R) for p in range(max(0, ctx[r] - WINDOW - T), ctx[r]) + ] + hp = torch.tensor([p for _, p in hist], device=dev) + kv = _rope(_normed_kv(len(hist), kv_norm, g, dev), hp, cos_sin) + _write_records(swa, SWA_BLOCK, torch.tensor([slot(r, p) for r, p in hist]), kv) + + step = dict( + M=M, + positions=torch.tensor(pos, dtype=torch.int64, device=dev), + slot_mapping=torch.tensor( + [slot(t2r[i], pos[i]) for i in range(M)], dtype=torch.int64, device=dev + ), + token_to_req=torch.tensor(t2r, dtype=torch.int32, device=dev), + swa_indices=swa_idx.to(dev), + swa_lens=swa_len.to(dev), + swa_cache=swa, + topk=None, + comp_cache=None, + comp_block_table=None, + comp_entries=0, + ) + if ratio: + entries = CACHE_BLOCK // ratio + lens = [(c + T) // ratio for c in ctx] + nblk = [-(-n // entries) for n in lens] + perm = (torch.randperm(sum(nblk), generator=g) + 1).tolist() + cbt = torch.zeros(R, max(nblk), dtype=torch.int32) + for r in range(R): + for b in range(nblk[r]): + cbt[r, b] = perm.pop() + comp = _paged_cache(1 + sum(nblk), entries, dev) + for r in range(R): + idx = torch.arange(lens[r]) + slots = cbt[r, idx // entries].long() * entries + idx % entries + kv = _rope( + _normed_kv(lens[r], kv_norm, g, dev), (idx * ratio).to(dev), cos_sin + ) + _write_records(comp, entries, slots, kv) + topk = torch.full((M, TOPK), -1, dtype=torch.int32) + for i in range(M): + cand = (pos[i] + 1) // ratio + n = min(cand, TOPK) + topk[i, :n] = torch.randperm(cand, generator=g)[:n].to(torch.int32) + step.update( + topk=topk.to(dev), + comp_cache=comp, + comp_block_table=cbt.to(dev), + comp_entries=entries, + ) + return step + + +# ---------------------------------------------------------------- weights + + +def _mxfp8(n: int, k: int, g, device): + """A random MXFP8 weight [n, k] e4m3 and its E8M0 32 x 32 block scales.""" + w = torch.randn(n, k, generator=g).clamp(-6, 6).to(torch.float8_e4m3fn) + s = torch.randint(118, 123, (n // 32, k // 32), generator=g, dtype=torch.uint8) + return w.to(device), s.to(device) + + +def _random_moe(moe, gd, gr, inter: int) -> None: + """Random parameters in the loaded layout, before post-load processing: this + rank's shards (``gd``, a device generator) of the MXFP4 routed experts + (random nibbles, E8M0 scales; the intermediate past ``inter`` zero, as the + loader pads it) and the MXFP8 shared expert (one E8M0 a 32 x 32 block); the + replicated router from ``gr``, seeded alike on every rank.""" + u8 = torch.uint8 + for lin in (moe.shared_experts.gate_up_proj, moe.shared_experts.down_proj): + n = lin.weight.shape[0] + w = torch.randn(lin.weight.shape, generator=gd, device=gd.device) + lin.weight.data.copy_(w.clamp(-6, 6).to(lin.weight.dtype)) + sc = lin.weight_scale.data.view(u8) + rows = -(-n // 32) if sc.shape[0] == n else sc.shape[0] + blocks = torch.randint( + 118, 123, (rows, sc.shape[1]), generator=gd, device=gd.device, dtype=u8 + ) + sc.copy_(blocks.repeat_interleave(32, 0)[:n] if sc.shape[0] == n else blocks) + e = moe.experts.routed_experts + for name, lo, hi in (("weight", 0, 256), ("weight_scale", 118, 122)): + for w in ("w13", "w2"): + t = getattr(e, f"{w}_{name}").data.view(u8) + t.copy_( + torch.randint(lo, hi, t.shape, generator=gd, device=gd.device, dtype=u8) + ) + padded = e.w13_weight.shape[1] // 2 + if padded > inter: + w13, w2 = e.w13_weight.data.view(u8), e.w2_weight.data.view(u8) + for half in (0, padded): + w13[:, half + inter : half + padded] = 0 + w2[..., inter // 2 :] = 0 + for t in (moe.gate.weight, moe.gate.e_score_correction_bias): + t.data.copy_(torch.randn(t.shape, generator=gr, device=gr.device) * 0.02) + + +# ---------------------------------------------------------------- the worker + + +@ray.remote(num_gpus=1, max_calls=1) +def _mono_numerics(monkeypatch, tp_size, pp_size, rank, distributed_init_port): + from vllm._aiter_ops import rocm_aiter_ops + from vllm.v1.worker.workspace import init_workspace_manager + + with monkeypatch.context() as m: + m.delenv("CUDA_VISIBLE_DEVICES", raising=False) + # the checkpoint's MXFP4 experts run on AITER's MoE; its flags are read + # at import, so refresh them + m.setenv("VLLM_ROCM_USE_AITER", "1") + m.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + rocm_aiter_ops.refresh_env_variables() + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) + ensure_model_parallel_initialized(tp_size, pp_size) + init_workspace_manager(device) + _check_layers(rank, tp_size, device) + + +def _check_layers(rank: int, tp: int, device) -> None: + from vllm._aiter_ops import rocm_aiter_ops + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce + from vllm.engine.arg_utils import EngineArgs + from vllm.forward_context import set_forward_context + from vllm.model_executor.kernels.linear.mxfp8.rocm_block32_gemm import ( + rocm_mxfp8_block32_gemm, + ) + from vllm.model_executor.layers.mhc import MHCPreDelayedOp + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, + ) + from vllm.model_executor.model_loader.utils import process_weights_after_loading + from vllm.models.common.ops import fused_q_kv_rmsnorm + from vllm.models.deepseek_v41.amd.model import DeepseekV4MoE + from vllm.models.deepseek_v41.amd.mono.layer import scratch_layout + from vllm.models.deepseek_v41.amd.mono.runner import ( + AttnWeights, + DSV41MonoLayer, + MonoLayerWeights, + ) + from vllm.models.deepseek_v41.amd.rocm import ( + compute_global_topk_ragged_indices_and_indptr, + ) + from vllm.models.deepseek_v41.common.rope import build_deepseek_v4_rope + from vllm.utils.torch_utils import set_default_torch_dtype + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + build_ragged_indices_from_dense, + rocm_inverse_rope_mxfp8_rows, + ) + + vc = EngineArgs( + model=MODEL, + tensor_parallel_size=tp, + moe_backend="aiter", + load_format="dummy", + max_model_len=65536, + language_model_only=True, + attention_config={"indexer_kv_dtype": "mxfp4", "indexer_sparse_logits": True}, + ).create_engine_config() + cfg = vc.model_config.hf_config + cfg = getattr(cfg, "text_config", cfg) + # replicated tensors alike on every rank; each rank's own TP shards + g = torch.Generator().manual_seed(1000) + gs = torch.Generator().manual_seed(1001 + rank) + H, G = N_HEADS // tp, O_GROUPS // tp + bf = torch.bfloat16 + + with set_current_vllm_config(vc), set_default_torch_dtype(bf), device: + moe = DeepseekV4MoE(vc, prefix="model.layers.21.ffn") + mhc = MHCPreDelayedOp() + _random_moe( + moe, + torch.Generator(device=device).manual_seed(2001 + rank), + torch.Generator(device=device).manual_seed(2000), + inter=2304 // tp, + ) + process_weights_after_loading(moe, vc.model_config, device) + + def f32(*shape, s=1.0, gen=g): + return (torch.randn(*shape, generator=gen) * s).to(device) + + mix = (2 + HC) * HC + seams = { + sub: dict( + fn=f32(mix, HC * HIDDEN, s=0.01), + scale=f32(3, s=0.5).abs() + 0.5, + base=f32(mix, s=0.1), + norm=(1 + f32(HIDDEN, s=0.1)).to(bf), + ) + for sub in ("attn", "ffn") + } + + def seam(sub, residual, x, post, comb, pre): + s = seams[sub] + return mhc( + residual, s["fn"], s["scale"], s["base"], cfg.rms_norm_eps, + cfg.hc_eps, cfg.hc_eps, 2.0, cfg.hc_sinkhorn_iters, + pre_mix=pre, sublayer_out=x, post_layer_mix=post, comb_res_mix=comb, + norm_weight=s["norm"], norm_eps=cfg.rms_norm_eps, + ) # fmt: skip + + def moe_out(x): + with set_forward_context(None, vc): + ids = torch.zeros(x.shape[0], dtype=torch.int64, device=device) + return moe(x.clone(), ids) + + wqkv, wqkv_s = _mxfp8(1792, HIDDEN, g, device) + wq_b, wq_b_s = _mxfp8(H * HEAD_DIM, Q_RANK, gs, device) + wo_a, wo_a_s = _mxfp8(G * O_RANK, H // G * HEAD_DIM, gs, device) + wo_b, wo_b_s = _mxfp8(HIDDEN, G * O_RANK, gs, device) + q_norm, kv_norm = (1 + f32(Q_RANK, s=0.1)).to(bf), (1 + f32(HEAD_DIM, s=0.1)).to(bf) + sink = f32(H, gen=gs) + e, sh = moe.experts.routed_experts, moe.shared_experts + u8 = lambda t: t.view(torch.uint8) # noqa: E731 + runner = DSV41MonoLayer(tp, rank, get_tp_group().cpu_group, device) + gw, bias = moe.gate.weight.float(), moe.gate.e_score_correction_bias.float() + + def top6(x): + """The top-6 experts as vLLM routes: sqrt(softplus(logits)) + bias.""" + sc = torch.nn.functional.softplus(x.float() @ gw.T).sqrt() + bias + return sc.topk(6, dim=1).indices.sort(1).values + + steps = [] # each case's launches, on fixed buffers, for the replay stress + for ratio, R, T in ((1, 1, 6), (2, 8, 6)): + with set_current_vllm_config(VllmConfig()): + rope = build_deepseek_v4_rope( + cfg, head_dim=HEAD_DIM, rope_head_dim=ROPE, + max_position_embeddings=cfg.max_position_embeddings, + compress_ratio=ratio, + ) # fmt: skip + cos_sin = rope.cos_sin_cache.to(device=device, dtype=torch.float32) + s = _decode_step(ratio, R, T, kv_norm, cos_sin, seed=7 + ratio) + M = s["M"] + mw = MonoLayerWeights( + attn=AttnWeights( + layer_id=21, wqkv=wqkv, wqkv_scale=wqkv_s, q_norm=q_norm, + kv_norm=kv_norm, wq_b=wq_b, wq_b_scale=wq_b_s, wo_a=wo_a, + wo_a_scale=wo_a_s, wo_b=wo_b, wo_b_scale=wo_b_s, attn_sink=sink, + cos_sin=cos_sin, ratio=ratio, + ), + hc_attn_fn=seams["attn"]["fn"], hc_attn_scale=seams["attn"]["scale"], + hc_attn_base=seams["attn"]["base"], attn_norm=seams["attn"]["norm"], + hc_ffn_fn=seams["ffn"]["fn"], hc_ffn_scale=seams["ffn"]["scale"], + hc_ffn_base=seams["ffn"]["base"], ffn_norm=seams["ffn"]["norm"], + gate_w=moe.gate.weight, bias=moe.gate.e_score_correction_bias, + w13=e.w13_weight, w13_s=e.w13_weight_scale, w2=e.w2_weight, + w2_s=e.w2_weight_scale, sgu=sh.gate_up_proj.weight, + sgu_s=u8(sh.gate_up_proj.weight_scale), sw2=sh.down_proj.weight, + sw2_s=u8(sh.down_proj.weight_scale), + ) # fmt: skip + + # the layer's inputs at the attention seam (the same on every rank) + gi = torch.Generator().manual_seed(11) + res0 = torch.randn(M, HC, HIDDEN, generator=gi) + res0 = ( + (res0 * torch.rsqrt(res0.pow(2).mean(-1, keepdim=True))).to(bf).to(device) + ) + x0 = (torch.randn(M, HIDDEN, generator=gi) * 0.5).to(bf).to(device) + comb0 = torch.rand(M, HC, HC, generator=gi) + 0.1 + for _ in range(20): + comb0 = comb0 / comb0.sum(-1, keepdim=True) + comb0 = comb0 / comb0.sum(-2, keepdim=True) + comb0 = comb0.to(device) + post0 = (torch.sigmoid(torch.randn(M, HC, 1, generator=gi)) * 2).to(device) + pre0 = (torch.sigmoid(torch.randn(M, HC, generator=gi)) + 1e-6).to(device) + + # vLLM's decode chain + res, post, comb, xa, attn_pre = seam("attn", res0, x0, post0, comb0, pre0) + swa_ind, swa_ptr = build_ragged_indices_from_dense( + s["swa_indices"].reshape(M, -1), s["swa_lens"] + ) + top_ind = top_ptr = None + if ratio: + valid = torch.ones(M, dtype=torch.bool, device=device) + top_ind, top_ptr, _ = compute_global_topk_ragged_indices_and_indptr( + s["topk"], s["token_to_req"], s["comp_block_table"], + s["comp_entries"], valid, + ) # fmt: skip + x8, xs = mxfp8_e4m3_quantize(xa) + qr_kv = rocm_mxfp8_block32_gemm(x8, xs, wqkv, wqkv_s, bf) + qr, kv = fused_q_kv_rmsnorm( + qr_kv[:, :Q_RANK], qr_kv[:, Q_RANK:], q_norm, kv_norm, cfg.rms_norm_eps + ) + q8, qs = mxfp8_e4m3_quantize(qr) + q = rocm_mxfp8_block32_gemm(q8, qs, wq_b, wq_b_s, bf).view(M, H, HEAD_DIM) + swa2d = s["swa_cache"].view(s["swa_cache"].shape[0], -1) + q_roped = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q, kv, swa2d, s["slot_mapping"], s["positions"], cos_sin, H, + cfg.rms_norm_eps, SWA_BLOCK, False, False, True, False, + ) # fmt: skip + attn = torch.empty(M, H, HEAD_DIM, dtype=bf, device=device) + rocm_aiter_ops.triton_sparse_mla_fwd( + q_roped, s["swa_cache"], attn, HEAD_DIM**-0.5, swa_ptr, swa_ind, + attn_sink=sink, extra_kv_buffer=s["comp_cache"], + extra_kv_indptr=top_ptr, extra_kv_indices=top_ind, + ) # fmt: skip + a8 = torch.empty(M, H * HEAD_DIM, dtype=torch.float8_e4m3fn, device=device) + a_s = torch.empty(M, H * HEAD_DIM // 32, dtype=torch.uint8, device=device) + rocm_inverse_rope_mxfp8_rows(attn, s["positions"], cos_sin, ROPE, a8, a_s) + z = torch.ops.vllm.rocm_dsv41_mxfp8_wo_a_bmm(a8, a_s, wo_a, wo_a_s, G, O_RANK) + z8, zs = mxfp8_e4m3_quantize(z) + part = rocm_mxfp8_block32_gemm(z8, zs, wo_b, wo_b_s, bf).contiguous() + a = tensor_model_parallel_all_reduce(part.clone()) + res2, post2, comb2, xf, ffn_pre = seam("ffn", res, a, post, comb, attn_pre) + out = moe_out(xf) + + # ---- the whole layer: K1 + K2 + o = runner.forward( + mw, x0, res0, post0, comb0, pre0, s["positions"], s["slot_mapping"], + s["swa_cache"], s["swa_indices"], s["swa_lens"], s["token_to_req"], + topk_indices=s["topk"], comp_cache=s["comp_cache"], + comp_block_table=s["comp_block_table"], + ) # fmt: skip + off = scratch_layout(M, tp)["normed"][0] + normed = runner.scratch(M)[off : off + M * HIDDEN * 2].view(bf).view(M, HIDDEN) + tag = f"ratio {ratio}, {R} x {T}" + assert _cos(runner.res_mid[:M], res) > 0.99999, tag + for got, want in ( + (runner.post_a, post), + (runner.comb_a, comb), + (runner.pre_a, attn_pre), + ): + torch.testing.assert_close( + got[:M].view_as(want), want, rtol=1e-4, atol=1e-5 + ) + same_in = moe_out(normed) + # relative errors about 3x what the kernels show (0.3-0.5%) + assert _cos(o[1], res2) > 0.9999 and _rel(o[1], res2) < 0.01, tag + assert _cos(normed, xf) > 0.999 and _rel(normed, xf) < 0.012, tag + assert _cos(o[0], same_in) > 0.9999 and _rel(o[0], same_in) < 0.015, tag + # per token, where both FFN inputs pick the same experts + same = (top6(xf) == top6(normed)).all(1) + assert same.float().mean() >= 0.5, tag + tok = torch.nn.functional.cosine_similarity(o[0].float(), out.float(), dim=1) + assert tok[same].min() > 0.99, tag + + # ---- the FFN launch, from the same unreduced wo_b output + f = runner.ffn(mw, part, res, post, comb, attn_pre) + torch.testing.assert_close(f[1], res2, rtol=2**-8, atol=1e-6) + torch.testing.assert_close(normed, xf, rtol=2**-7, atol=1e-6) + assert _cos(f[0], out) > 0.9999 and _rel(f[0], out) < 0.015, tag + torch.accelerator.synchronize() + + whole, ffn = runner._outs(M, res0), runner._outs(M, res) + args = (s["positions"], s["slot_mapping"], s["swa_cache"], s["swa_indices"]) + meta = dict( + topk_indices=s["topk"], comp_cache=s["comp_cache"], + comp_block_table=s["comp_block_table"], + ) # fmt: skip + + def step(mw=mw, x0=x0, res0=res0, post0=post0, comb0=comb0, pre0=pre0, + args=args, s=s, meta=meta, whole=whole, part=part, res=res, + post=post, comb=comb, attn_pre=attn_pre, ffn=ffn): # fmt: skip + runner.forward( + mw, x0, res0, post0, comb0, pre0, *args, s["swa_lens"], + s["token_to_req"], **meta, outs=whole, + ) # fmt: skip + runner.ffn(mw, part, res, post, comb, attn_pre, outs=ffn) + return [*whole, *ffn] + + steps.append(step) + + # ---- the hand-off protocol under back-to-back launches: both step widths + # alternating in one graph, every replay bit for bit the first's + outs = [t for step in steps for t in step()] + torch.accelerator.synchronize() + ref = [t.clone() for t in outs] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for step in steps: + step() + for i in range(REPLAYS): + graph.replay() + if i % 50 == 49 or i == REPLAYS - 1: + torch.accelerator.synchronize() + bad = [ + j for j, (t, r) in enumerate(zip(outs, ref)) if not torch.equal(t, r) + ] + assert not bad, f"replay {i}: outputs {bad} differ from the first run" + + +@pytest.mark.skipif(not _on_cdna4(), reason="the mono kernels target CDNA4") +@pytest.mark.parametrize("tp", [2, 4]) +def test_mono_decode_layer_matches_vllm(monkeypatch: pytest.MonkeyPatch, tp: int): + if torch.accelerator.device_count() < tp: + pytest.skip(f"needs {tp} GPUs") + pytest.importorskip("flydsl") + multi_process_parallel(monkeypatch, tp, 1, _mono_numerics) diff --git a/vllm/envs.py b/vllm/envs.py index 851741d5253b..02c282fd8ec2 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -155,6 +155,7 @@ VLLM_ROCM_USE_AITER_FP4BMM: bool = True VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = False + VLLM_ROCM_MONO_DECODE: bool = False VLLM_ROCM_USE_AITER_TRITON_GEMM: bool = True VLLM_ROCM_USE_SKINNY_GEMM: bool = True VLLM_ROCM_FP8_PADDING: bool = True @@ -1387,6 +1388,12 @@ def _resolve_rust_cli_path() -> str | None: os.getenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "False").lower() in ("true", "1") ), + # Run the eligible layers of decode steps on the model's mono decode + # kernels: persistent launches a layer, the TP all-reduces in-kernel + # (DeepSeek-V4.1 on CDNA4). By default is disabled. + "VLLM_ROCM_MONO_DECODE": lambda: ( + os.getenv("VLLM_ROCM_MONO_DECODE", "False").lower() in ("true", "1") + ), # Whether to use aiter triton kernels for gemm ops. # By default is enabled. "VLLM_ROCM_USE_AITER_TRITON_GEMM": lambda: ( diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index cff72f12dc22..5de173a8706c 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -14,6 +14,7 @@ get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, ) from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger @@ -58,6 +59,7 @@ from vllm.models.deepseek_v4.amd.model import ( DeepseekV4MoE as DeepseekV4MoEBase, ) +from vllm.models.deepseek_v41.amd.mono_decode import MonoDecodeLayer from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention from vllm.models.deepseek_v41.attention import DeepseekV4Attention from vllm.sequence import IntermediateTensors @@ -255,6 +257,12 @@ def __init__( # attn_norm / ffn_norm into its collapse, so the separate norms are # skipped for the seams it takes. self.fuse_seam_norm = HAS_AITER_MHC_FUSED_POST_PRE_DELAYED_RMS_NORM + # Decode steps of an eligible layer may run as the mono decode layer. + self.mono = MonoDecodeLayer.create(self, vllm_config) + if self.mono is not None and self.mono.ffn_only: + # A decode step hands wo_b's TP partial to the FFN launch; forward + # reduces it on the others. + self.attn.wo_b.reduce_results = False @staticmethod def _hc_collapse(x: torch.Tensor, pre_mix: torch.Tensor) -> torch.Tensor: @@ -273,6 +281,10 @@ def forward( engram_hashes: torch.Tensor | None = None, engram_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if self.mono is not None: + out = self.mono(self, x, positions, residual, post_mix, res_mix, pre_mix) + if out is not None: + return out # Layer 0's attention seam projects the 2-D embedding with the folded # hc_attn_fn_broadcast instead of the 4-stream residual with hc_attn_fn. # The fused kernel only takes the latter, so that seam keeps the @@ -371,6 +383,11 @@ def forward( x = self.attn(positions, x, None) if self.use_sequence_parallel: x = sp_reduce_scatter(x) + elif self.mono is not None and self.mono.ffn_only: + out = self.mono.ffn(self, x, residual, post_mix, res_mix, attn_pre) + if out is not None: + return out + x = tensor_model_parallel_all_reduce(x) residual, post_mix, res_mix, x, ffn_pre = self.mhc_pre_delayed( residual, diff --git a/vllm/models/deepseek_v41/amd/mono/README.md b/vllm/models/deepseek_v41/amd/mono/README.md new file mode 100644 index 000000000000..88e40d73dd0a --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/README.md @@ -0,0 +1,94 @@ +# DeepSeek-V4.1-Flash mono decode layer (gfx950, FlyDSL) + +One fused decoder layer for vLLM's decode steps on MI355X: the attention seam, +the attention, the attention all-reduce, the FFN seam, the MoE and the MoE +all-reduce, in two persistent launches a layer. The other backbone layers keep +vLLM's attention and run the rest -- the attention all-reduce, the FFN seam, +the MoE and its all-reduce -- as one launch (the FFN launch). ATOM's V4.1 mono decode +(ROCm/ATOM `atom/models/deepseek_v41/mono`) is the reference design: the seam +and MoE stages (`stages/`) and the framework under them (`common/`) are adapted +from it, each file naming its origin; the attention stages (`attention/`) are +new, reading vLLM's KV records and weights as loaded. + +## Scope + +- Layers: the whole layer on every backbone layer with the standard seam and no + compressor / indexer: 3-7, 9-13, 15-19 (compress ratio 2) and 21-23, 25-27, + 29-31, 33-35, 37-39 (ratio 1), 30 of 40. The FFN launch on the other 10: + layer 0 (embedding-broadcast seam), layer 1 (Engram at the seam) and the + compressor / indexer layers 2, 8, 14 (also Engram), 20, 24, 28, 32, 36, + whose wo_b leaves its TP reduce to it. The DSpark draft layers keep vLLM's + path. The paths interleave freely (same tensors at the layer boundary). +- Steps: decode only (FULL graph or eager), M <= 48 rows (8 requests x 6 + DSpark tokens), TP2 / TP4, `fp8_ds_mla` KV records (ROCm), text tokens (the + routing bias is `e_score_correction_bias`). +- Contract: a drop-in for `DeepseekV4DecoderLayer.forward`: + in `x` (the previous FFN output, reduced, bf16 [M, 5120]), `residual` + (bf16 [M, 4, 5120]), `post_mix` [M, 4, 1], `res_mix` [M, 4, 4], + `pre_mix` [M, 4] (f32, the previous seam's) + out the MoE output (reduced), the residual after the FFN seam, and that + seam's post / comb / pre mixes. + +## Launches + + K1 seam (stages/seam.py): slice (160 x 32 columns): fold x into the + residual (post x + sum comb r), the collapsed input (pre . R'), partial + mixes (fn . R') and sums of squares -> gate (S): 24 mixes -> pre / post + / comb (Sinkhorn 20) -> norm (S): attn_norm -> vLLM MXFP8 -> X8 + attention front (attention/front.py): wqkv, q / kv norms + KV insert, + wq_b + RoPE -> Q + K2 attention back (attention/back.py): split attention, combine + inverse + RoPE, wo_a, wo_b; the bf16 wo_b partial pushed to every rank + FFN seam: reduce (the ranks' partials summed in rank order, fp32 -> + bf16: the attention output), slice, gate, norm -> the MoE input + MoE (stages/moe.py): xq, router (bf16 GEMV, f32 logits), route + (sqrtsoftplus + bias, top-6, renormalized x 1.5), shared expert (MXFP8 + GEMVs, clamped SiLU), ug (MXFP4 x MXFP8 per expert of the step's + union), down (+ routing weight, top-k order; + shared) -> push, rank-order + sum -> out + FFN K2 without the attention back: push (160 x 32 columns) of vLLM's + unreduced wo_b output, as its bits, to every rank's attention region, + then K2's FFN seam and MoE (the same code, `layer.run_ffn`) + +## Execution model and hand-offs + +- Grid = 256 CTAs (one a CU, all resident), 512 threads a CTA, LDS a per-stage + union. A stage is a task table: CTA `b` runs tasks `b, b + 256, ...`, stages + in order; a consumer only waits on earlier stages, so no deadlock while the + grid is resident. A GEMV task issues its weight loads before it polls its + activations. +- In a launch: tagged pairs (value + tag in one 8 B store) or plain data + + flags, at device scope. The tag is the launch pair's epoch (K1 `2e + 1`, + K2 or the FFN launch `2e + 2`), a device word their CTA 0 moves on once every + CTA has marked the launch done: no clearing between steps, no host step hook, every argument + a fixed device pointer (graph-safe). +- Across ranks: one symmetric uncached buffer a rank (aiter `UncachedIpcHeap`), + every rank holding every peer's address; pushes are system-scope stores of + (value, tag) pairs. Every rank runs the same launches, so the epochs agree; + peer regions alternate by epoch parity, and a rank cannot run two steps ahead + of a peer (each step's reduce waits on every rank's push). + +## Numerics (the vLLM path each stage reproduces) + +- Seams: aiter `mhc_fused_post_pre_delayed_rmsnorm` (post fold order, fn as + bf16 hi + lo, products summed in order, Triton sigmoid / exp / rcp). +- Attention: vLLM's ROCm chain. `mxfp8(x)`: one E8M0 per 32 values, + `code = clamp(ceil(log2(amax / 448)) + 127, 0, 254)`; weights in vLLM's + row-major e4m3 [N, K] with [N/32, K/32] E8M0 scales (one shared copy); KV + records as vLLM's insert writes them (dims 0..447 e4m3 with a UE8M0 per 64, + RoPE dims bf16, 584 B); keys = the sliding window plus the top-k compressed + rows, softmax in base 2 with the attention sink folded in. +- MoE: vLLM's router (bf16 gate GEMV, f32 logits, `topk_hash_softplus_sqrt`) + and AITER a8w4 `fused_moe` (interleaved gate / up, swiglu limit 10, weights + in aiter's (16, 16) shuffle and `shuffle_scale` order); the shared expert as + vLLM's MXFP8 block-32 linears with `silu_and_mul_with_clamp`. +- All-reduces: exact, fp32 in rank order -> bf16. vLLM's decode all-reduces at + M <= 48 are below its INT4 quick-reduce threshold, so vLLM runs them exact too. + +## Validation + +- Per layer against vLLM's own modules in a TP process group (`DeepseekV4MoE`, + the mHC seams, the attention chain, the all-reduce), real checkpoint weights. +- Bench: layers of one type back to back in one HIP graph, cold weights, + us per layer, M = 1..48, TP2 / TP4; then e2e (AgentX, InferenceX#3690 + recipe). diff --git a/vllm/models/deepseek_v41/amd/mono/__init__.py b/vllm/models/deepseek_v41/amd/mono/__init__.py new file mode 100644 index 000000000000..72a6ec2ab196 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4.1-Flash mono decode layer: one decoder layer of a decode step in +two persistent FlyDSL launches, its TP all-reduces in-kernel (README.md). The +host side is ``runner``; vLLM docks it in ``..mono_decode``.""" diff --git a/vllm/models/deepseek_v41/amd/mono/attention/__init__.py b/vllm/models/deepseek_v41/amd/mono/attention/__init__.py new file mode 100644 index 000000000000..da1bf8c269a0 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/attention/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The mono layer's attention stages: ``front`` (K1: wqkv, q / kv norms and the +KV insert, wq_b) and ``back`` (K2: split attention, combine, wo_a, wo_b).""" diff --git a/vllm/models/deepseek_v41/amd/mono/attention/back.py b/vllm/models/deepseek_v41/amd/mono/attention/back.py new file mode 100644 index 000000000000..4526f8eda7ee --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/attention/back.py @@ -0,0 +1,884 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K2 ``back``: one DeepSeek-V4.1 attention layer from the rotated query to the +rank's wo_b partial. + +One launch a layer and rank, ``BLOCKS`` x ``THREADS``, S rows: + + split (S x chunks of CK keys): flash-decoding over one chunk of a token's + keys (K1's key table: top-k compressed rows, then window rows) for + every head of the rank. The token's q in LDS; each wave one + (head tile, 16-key tile): its fp8_ds_mla rows loaded once, + dequantized into its QK operands and kept raw in LDS; QK (bf16 + MFMA), the chunk's max / sum per head, bf16 p in LDS; PV over the + chunk's raw rows (wave w: dims 64 w ..) -> the chunk's partial + (bf16) and its (max, sum) per head, published + combine (S x H): a (token, head)'s chunks merged with the sink folded in, + bf16, inverse GPT-J RoPE of dims 448.., vLLM's MXFP8 of each 32 + -> XO (published) + wo_a (G x 1024 / 32, two 16-row tiles): the group's XO rows times wo_a + -> bf16 -> MXFP8 of the 32 rows -> X8B (published) + wo_b (320, 16 rows): X8B times wo_b -> bf16 -> ``out`` + +Hand-offs are plain device-scope data behind tagged flags (half the bytes of a +tagged pair per value). Adapted from ATOM's V4.1 ``attn_post`` (K2a), which +reads a bf16 pool and ends in a peer all-reduce. +""" + +from dataclasses import dataclass + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T, as_ir_value + +from .device import ( + BLOCKS, + CM_DEV, + THREADS, + WAVES, + Mailbox, + bf16_round, + bf_hi, + bf_lo, + butterfly, + clamp_fp8, + first_task, + fp8_pack4, + fp8x8_bf16, + fresh, + gload, + gstore, + hw_exp2, + lane_gather, + ld_f32, + ld_i32, + mfma_bf16, + mx_code, + mx_mul, + pow2, + rsrc, + spin_until, + stamp, + traced, + xshfl, +) +from .gemv import RED, gemv_loads, gemv_mfmas, lds_row, row_sum, rows_loads, rows_store +from .plan import ( + DATA, + HALF, + HEAD_DIM, + HEAD_TILE, + KEYS, + NOPE, + ROPE, + ROWS, + WOA_ROWS, + Dims, + back_scratch, + cdiv, + front_end, + tile_rows, + token_tiles, +) + +_STREAM = fx.Stream(None) +NEG = -3.4028234663852886e38 +LOG2E = 1.4426950408889634 +REC = DATA + 8 # a raw record row in LDS: 576 data bytes + 8 scale bytes +Q_ROW = ( + HEAD_DIM // 2 + 4 +) # a q row in LDS (i32 words): 512 bf16 + 16 B, no bank aliasing +P_PAD = 2 # i32 words past a P row: 16 head rows read together hit distinct banks +# wo_a's K order: waves 0..3 / 4..7 take the two 16-row tiles, a contiguous K +# range a wave, summed in wave order +WOA_SPLIT_K = WAVES * ROWS // WOA_ROWS +BACK_POINTS = 16 # start, stage ends, split internals, GEMV internals + + +@dataclass(frozen=True) +class BackBuild: + tokens: int + tp: int + ratio: int # 0: window only; 1 / 2: top-k over a compressed cache + timeline: bool = False + + +def plus(t0, v): + return v if isinstance(t0, int) and t0 == 0 else v + t0 + + +def head_groups(d: Dims, s: int) -> int: + """Wave groups a split unit's head tiles are spread over: two (a tile each, + 64-key chunks, twice the units) while the units fit the grid, else one (a + wave's keys loaded and dequantized once for every head tile).""" + return 2 if d.head_tiles == 2 and s * cdiv(KEYS, 64) <= BLOCKS else 1 + + +def chunk_keys(d: Dims, s: int) -> int: + """A split unit's keys: a 16-key tile a wave of a head group.""" + return 16 * (WAVES // head_groups(d, s)) + + +def p_row(d: Dims, s: int) -> int: + """A P row (one head's chunk of bf16 p) in LDS, i32 words.""" + return chunk_keys(d, s) // 2 + P_PAD + + +@traced +def publish(c, region, i): + """This CTA's stores visible device-wide, then the flag pair (region, i).""" + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if c["tid"] == 0: + c["mb"].put(region, i, fx.Int32(1)) + + +@traced +def await_flags(c, region, n, index): + """Thread u < n polls flag ``index(u)``; then the CTA proceeds.""" + tid = c["tid"] + if tid < n: + c["mb"].poll([(region, index(fx.min(tid, n - 1)), 1)]) + gpu.barrier() + + +# ---------------------------------------------------------------- keys + + +def key_rec(c, t, k): + """(data address, scale address, valid) of token t's key k from K1's key + table: slot | 1 << 31 for a compressed row, a window slot otherwise, -1 for + none (it reads a zero record: its scores are masked, its values zeros).""" + a, ratio = c["args"], c["ratio"] + val = ld_i32(a["kt"], t * KEYS + fx.min(fx.max(k, 0), KEYS - 1)) + ok = val != fx.Int32(-1) + slot = ok.select(val & fx.Int32(0x7FFFFFFF), fx.Int32(0)) + if const_expr(ratio > 0): + is_c = ok & (val < 0) + blk = is_c.select(a["comp_block"], a["swa_block"]) + stride = is_c.select(a["comp_stride"], a["swa_stride"]) + base = is_c.select(a["comp"], a["swa"]) + else: + blk, stride, base = a["swa_block"], a["swa_stride"], a["swa"] + b = slot // blk + off = slot % blk + page = base + fx.Int64(b) * stride + z = a["zrec"] + data = ok.select(page + fx.Int64(off * DATA), z) + sc = ok.select(page + fx.Int64(blk * DATA + off * 8), z + fx.Int64(DATA)) + return data, sc, ok + + +def _scale_byte(scw, blk64): + return (scw[blk64 // 4] >> (blk64 % 4 * 8)) & 0xFF + + +def _fp8_pair_f32(word, scale, hi): + v = rocdl.cvt_scalef32_pk_f32_fp8( + fx.Vector.make_type(2, fx.Float32), + as_ir_value(fx.Int32(word)), + as_ir_value(fx.Float32(scale)), + hi, + ) + v = fx.Vector(v) + return [v[0], v[1]] + + +# ---------------------------------------------------------------- split + + +@traced +def stage_split(c, unit): + """Flash-decoding over chunk ``unit % nchunk`` of token ``unit // nchunk``, + every head of the rank -> the chunk's partial and (max, sum) per head. Wave + (head group hg, key tile kt): keys 16 kt .. of the chunk (loaded and + dequantized once) against the group's head tiles; then dims 64 w .. of the + PV over every key and head.""" + tid, lane, wave, d, s = c["tid"], c["lane"], c["wave"], c["d"], c["S"] + a, sl = c["args"], c["sl"] + nchunk = c["nchunk"] + t = unit // nchunk + ch = unit % nchunk + ck = chunk_keys(d, s) + ngrp = head_groups(d, s) + nkt = WAVES // ngrp # 16-key tiles a chunk + nht = d.head_tiles // ngrp # head tiles a wave + hgi = wave // nkt + kt = wave % nkt + g = lane // 16 + r16 = lane % 16 + k0 = ch * ck + # the key count, this lane's key and q: independent loads, out together + # (K1's key table holds -1 past a token's keys) + kv_len = ld_i32(a["klen"], t) + key = k0 + kt * 16 + r16 + data, sc, okv = key_rec(c, t, key) + qrow = d.heads * HEAD_DIM // 8 + qv = [] + for i in range_constexpr(cdiv(qrow, THREADS)): + u = fresh(tid) + THREADS * i + qv.append( + fx.Vector( + bo.buffer_load( + rsrc(a["q"]), + (t * d.q_rows) // 2 + fx.min(u, qrow - 1) * 4, + vec_width=4, + dtype=T.i32, + ) + ) + ) + if k0 < kv_len: + # a key's 448 e4m3 as 7 x 16 B (lane g: dims 64 c + 16 g ..), its 64 + # bf16 rope dims as 2 x 16 B (32 c + 8 g ..), its 7 scale bytes + scw = gload(sc, words=2) + nope = [ + gload(data + fx.Int64(64 * c4 + 16 * g), words=4) + for c4 in range(NOPE // 64) + ] + rope = [ + gload(data + fx.Int64(NOPE + (cc * 32 + 8 * g) * 2), words=4) + for cc in range(ROPE // 32) + ] + # q of every head into LDS (bf16 rows of 512 + pad), 16 B a thread + for i in range_constexpr(cdiv(qrow, THREADS)): + u = fresh(tid) + THREADS * i + if u < qrow: + fx.ptr_store( + qv[i], + sl["q"] + + (u // (HEAD_DIM // 8)) * Q_ROW + + (u % (HEAD_DIM // 8)) * 4, + ) + # the raw rows to LDS for the PV (head group 0's waves) + if hgi == 0: + rawb = sl["raw"] + (kt * 16 + r16) * (REC // 4) + for c4 in range_constexpr(NOPE // 64): + fx.ptr_store(nope[c4], rawb + (64 * c4 + 16 * g) // 4) + for cc in range_constexpr(ROPE // 32): + fx.ptr_store(rope[cc], rawb + (NOPE + (cc * 32 + 8 * g) * 2) // 4) + if g == 0: + fx.ptr_store(scw, rawb + DATA // 4) + gpu.barrier() + stamp(c, 5) + # QK: A = the keys (dequantized once), B = q of a head tile (LDS); the + # contraction's 32-dim chunk 2 c + h of the nope dims is {64 c + 16 g + + # 8 h ..} (lane g's 16 B load), q read in the same order + scale = fx.Int32(a["qk_scale"]).bitcast(fx.Float32) + st = sl["st"] + accs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(nht)] + for c4 in range_constexpr(NOPE // 64): + ksc = pow2(_scale_byte(scw, c4)) + for h2 in range_constexpr(2): + kw = fp8x8_bf16(nope[c4][2 * h2], nope[c4][2 * h2 + 1], ksc).bitcast( + fx.BFloat16 + ) + for ht in range_constexpr(nht): + qb = sl["q"] + ((hgi * nht + ht) * HEAD_TILE + r16) * Q_ROW + qw = fx.Vector( + fx.ptr_load( + qb + 32 * c4 + 8 * g + 4 * h2, + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + accs[ht] = mfma_bf16(kw, qw.bitcast(fx.BFloat16), accs[ht]) + for cc in range_constexpr(ROPE // 32): + kw = rope[cc].bitcast(fx.BFloat16) + for ht in range_constexpr(nht): + qb = sl["q"] + ((hgi * nht + ht) * HEAD_TILE + r16) * Q_ROW + qw = fx.Vector( + fx.ptr_load( + qb + (NOPE + cc * 32 + 8 * g) // 2, + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + accs[ht] = mfma_bf16(kw, qw.bitcast(fx.BFloat16), accs[ht]) + # scores of keys 4 g + i (C rows), head r16 of a tile (C column) + ok = (key < kv_len) & okv + okw = ok.select(fx.Int32(1), fx.Int32(0)) + valid = [lane_gather(okw, 4 * g + i) != 0 for i in range(4)] + svs = [] + for ht in range_constexpr(nht): + sv = [ + fx.Float32(valid[i].select(accs[ht][i] * scale, fx.Float32(NEG))) + for i in range(4) + ] + mw = butterfly( + fx.max(fx.max(sv[0], sv[1]), fx.max(sv[2], sv[3])), (32, 16), fx.max + ) + if g == 0: + fx.ptr_store(mw, st + ((hgi * nht + ht) * nkt + kt) * HEAD_TILE + r16) + svs.append(sv) + gpu.barrier() + lst = st + d.heads * nkt + for ht in range_constexpr(nht): + hti = hgi * nht + ht + m = fx.ptr_load(st + (hti * nkt) * HEAD_TILE + r16) + for j in range_constexpr(1, nkt): + m = fx.max(m, fx.ptr_load(st + (hti * nkt + j) * HEAD_TILE + r16)) + sv = svs[ht] + p = [ + fx.Float32(valid[i].select(hw_exp2(sv[i] - m), fx.Float32(0.0))) + for i in range(4) + ] + lw = (p[0] + p[2]) + (p[1] + p[3]) + lw = lw + xshfl(lw, 32) + lw = lw + xshfl(lw, 16) + pw = ( + fx.Vector.from_elements(p, fx.Float32).to(fx.BFloat16).bitcast(fx.Int32) + ) + pbase = ( + sl["p"] + (hti * HEAD_TILE + r16) * p_row(d, s) + (kt * 16 + 4 * g) // 2 + ) + fx.ptr_store(fx.Vector.from_elements([pw[0], pw[1]], fx.Int32), pbase) + if g == 0: + fx.ptr_store(lw, lst + (hti * nkt + kt) * HEAD_TILE + r16) + gpu.barrier() + # chunk stats: (max, sum) per head + if tid < d.heads: + hts = tid // HEAD_TILE + hr = tid % HEAD_TILE + mm = fx.ptr_load(st + (hts * nkt) * HEAD_TILE + hr) + ll = fx.ptr_load(lst + (hts * nkt) * HEAD_TILE + hr) + for j in range_constexpr(1, nkt): + mm = fx.max(mm, fx.ptr_load(st + (hts * nkt + j) * HEAD_TILE + hr)) + ll = ll + fx.ptr_load(lst + (hts * nkt + j) * HEAD_TILE + hr) + bo.buffer_store( + fx.Vector.from_elements([mm, ll], fx.Float32), + rsrc(c["pstat"]), + ((t * nchunk + ch) * d.heads + tid) * 2, + cache_modifier=CM_DEV, + ) + # PV: wave w dims 64 w + 4 r16 + dt (dt = 0..3: one LDS word holds the + # four of a key); the rope wave (dims 448..) reads bf16 pairs. Lane + # (r16, g) takes keys 4 g + j of a 16-key tile: A[dim][key] + is_rope = wave == WAVES - 1 + dbase = 64 * wave + sh = (wave % 4) * 8 + accs = [ + [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(4)] + for _ in range(d.head_tiles) + ] + for kt2 in range_constexpr(nkt): + rb0 = sl["raw"] + (kt2 * 16 + 4 * g) * (REC // 4) + vv = fx.Vector.filled(16, 0.0, fx.Float32) # [dt][j] + if is_rope: + rv = [] + for j in range_constexpr(4): + rw = fx.Vector( + fx.ptr_load( + rb0 + j * (REC // 4) + (NOPE + 8 * r16) // 4, + result_type=fx.Vector.make_type(2, fx.Int32), + ) + ) + rv.append([bf_lo(rw[0]), bf_hi(rw[0]), bf_lo(rw[1]), bf_hi(rw[1])]) + vv = fx.Vector.from_elements( + [rv[j][dt] for dt in range(4) for j in range(4)], fx.Float32 + ) + else: + nv = [] + for j in range_constexpr(4): + rb = rb0 + j * (REC // 4) + w = fx.Int32(fx.ptr_load(rb + (dbase + 4 * r16) // 4)) + scb = fx.Int32(fx.ptr_load(rb + DATA // 4 + wave // 4)) + scl = pow2((scb >> sh) & 0xFF) + nv.append( + _fp8_pair_f32(w, scl, False) + _fp8_pair_f32(w, scl, True) + ) + vv = fx.Vector.from_elements( + [nv[j][dt] for dt in range(4) for j in range(4)], fx.Float32 + ) + for ht2 in range_constexpr(d.head_tiles): + pb = fx.Vector( + fx.ptr_load( + sl["p"] + + (ht2 * HEAD_TILE + r16) * p_row(d, s) + + (kt2 * 16 + 4 * g) // 2, + result_type=fx.Vector.make_type(2, fx.Int32), + ) + ).bitcast(fx.Int16) + for dt in range_constexpr(4): + av = ( + fx.Vector.from_elements( + [vv[dt * 4 + j] for j in range(4)], fx.Float32 + ) + .to(fx.BFloat16) + .bitcast(fx.Int16) + ) + accs[ht2][dt] = fx.Vector( + rocdl.mfma_f32_16x16x16bf16_1k( + T.vec(4, T.f32), [av, pb, accs[ht2][dt]] + ) + ) + stamp(c, 6) + # partial[t][ch][head][dim] bf16: lane (r16, g) holds C rows 4 g + i -> + # dims 64 w + 16 g + 4 i + dt of head ht2 * 16 + r16: 32 contiguous + # bytes, two 16 B stores (4 lanes fill a 128 B line) + for ht2 in range_constexpr(d.head_tiles): + hh = ht2 * HEAD_TILE + r16 + words = [] + for i in range_constexpr(4): + v4 = fx.Vector.from_elements( + [accs[ht2][dt][i] for dt in range(4)], fx.Float32 + ) + w2 = v4.to(fx.BFloat16).bitcast(fx.Int32) + words += [w2[0], w2[1]] + off = (((t * nchunk + ch) * d.heads + hh) * HEAD_DIM + dbase + 16 * g) // 2 + for hlf in range_constexpr(2): + bo.buffer_store( + fx.Vector.from_elements(words[4 * hlf : 4 * hlf + 4], fx.Int32), + rsrc(c["part"]), + off + 4 * hlf, + cache_modifier=CM_DEV, + ) + publish(c, c["srdy"], t * nchunk + ch) + stamp(c, 7) + + +# ---------------------------------------------------------------- combine + +COMB_HEADS = WAVES # heads a combine unit: a wave each + + +@traced +def stage_combine(c, unit): + """(token, 8 heads) ``unit``, a wave a head, 8 dims a lane: the chunks + merged with the sink folded in, bf16, inverse RoPE, MXFP8 of each 32 -> XO, + published (a flag a (token, head group)).""" + lane, wave, d, s = c["lane"], c["wave"], c["d"], c["S"] + a = c["args"] + hg_n = d.heads // COMB_HEADS + t = unit // hg_n + hg = unit % hg_n + h = hg * COMB_HEADS + wave + nchunk = c["nchunk"] + ck = chunk_keys(d, s) + kv_len = ld_i32(a["klen"], t) + nlive = (kv_len + ck - 1) // ck + pos = ld_i32(a["pos"], 2 * t) + sink = ld_f32(a["sink"], h) * fx.Float32(LOG2E) + d0 = lane * 8 # this lane's 8 dims + rope_lane = d0 >= NOPE + jj = fx.max((d0 - NOPE) // 2, 0) # 4 pairs from here + cs = [ + rope_lane.select(ld_f32(a["cos_sin"], pos * ROPE + jj + i), fx.Float32(1.0)) + for i in range(4) + ] + sn = [ + rope_lane.select( + ld_f32(a["cos_sin"], pos * ROPE + HALF + jj + i), fx.Float32(0.0) + ) + for i in range(4) + ] + await_flags(c, c["srdy"], nlive, lambda u: t * nchunk + u) + m = fx.Float32(NEG) + ms, ls, vs = [], [], [] + for cc in range_constexpr(nchunk): + live = cc < nlive + stv = fx.Vector( + bo.buffer_load( + rsrc(c["pstat"]), + ((t * nchunk + cc) * d.heads + h) * 2, + vec_width=2, + dtype=T.f32, + ) + ) + pv = fx.Vector( + bo.buffer_load( + rsrc(c["part"]), + (((t * nchunk + cc) * d.heads + h) * HEAD_DIM + d0) // 2, + vec_width=4, + dtype=T.i32, + ) + ) + mc = fx.Float32(live.select(stv[0], fx.Float32(NEG))) + ms.append(mc) + ls.append(fx.Float32(live.select(stv[1], fx.Float32(0.0)))) + vals = [] + for e in range_constexpr(4): + vals += [bf_lo(pv[e]), bf_hi(pv[e])] + vs.append([fx.Float32(live.select(v, fx.Float32(0.0))) for v in vals]) + m = fx.max(m, mc) + m_final = fx.max(m, sink) + lsum = hw_exp2(sink - m_final) + tot = [fx.Float32(0.0)] * 8 + for cc in range_constexpr(nchunk): + wgt = hw_exp2(ms[cc] - m_final) + lsum = lsum + ls[cc] * wgt + tot = [tot[e] + vs[cc][e] * wgt for e in range(8)] + inv = fx.Float32(1.0) / lsum + o = [bf16_round(tot[e] * inv) for e in range(8)] + # inverse GPT-J RoPE of each pair: e' = e cos + o sin, o' = o cos - e sin + y = [] + for i in range_constexpr(4): + e, od = o[2 * i], o[2 * i + 1] + y += [e * cs[i] + od * sn[i], od * cs[i] - e * sn[i]] + amax = abs(y[0]) + for v in y[1:]: + amax = fx.max(amax, abs(v)) + amax = butterfly(amax, (1, 2), fx.max) # 32 dims = 4 lanes + code = mx_code(amax) + mul = mx_mul(code) + w0 = fp8_pack4(*[clamp_fp8(y[i] * mul) for i in range(4)]) + w1 = fp8_pack4(*[clamp_fp8(y[4 + i] * mul) for i in range(4)]) + bo.buffer_store( + fx.Vector.from_elements([w0, w1], fx.Int32), + rsrc(c["xo"]), + t * (d.q_rows // 4) + (h * HEAD_DIM + d0) // 4, + cache_modifier=CM_DEV, + ) + if lane % 4 == 0: + bo.buffer_store( + code, + rsrc(c["xos"]), + t * (d.q_rows // 32) + (h * HEAD_DIM + d0) // 32, + cache_modifier=CM_DEV, + ) + publish(c, c["crdy"], t * hg_n + hg) + + +# ---------------------------------------------------------------- wo_a / wo_b + + +@traced +def stage_woa(c, task, ops, t0, n): + """Rows 32 task .. of wo_a -> bf16 -> vLLM's MXFP8 of the 32 -> X8B.""" + tid, red, d, lane = c["tid"], c["red"], c["d"], c["lane"] + split = WOA_SPLIT_K + gemv_mfmas(c, d.group_k, c["xl"], c["xsl"], red, ops, split, tiled=True, rows=n) + gpu.barrier() + mine = tid < WOA_ROWS * n + tl = fx.min(tid // WOA_ROWS, n - 1) + t = plus(t0, tl) + row = tid % WOA_ROWS + tile = row // ROWS + y = bf16_round( + row_sum(red, row % ROWS, tl, [tile * split + w for w in range(split)]) + ) + amax = butterfly(abs(y), (1, 2, 4, 8, 16), fx.max) + code = mx_code(amax) + q = clamp_fp8(y * mx_mul(code)) + w = fp8_pack4( + q, + *[ + lane_gather(q.bitcast(fx.Int32), fx.min(lane + k, 63)).bitcast(fx.Float32) + for k in (1, 2, 3) + ], + ) + if mine & (row % 4 == 0): + bo.buffer_store( + w, + rsrc(c["x8b"]), + t * (d.o_rows // 4) + (task * WOA_ROWS + row) // 4, + cache_modifier=CM_DEV, + ) + if mine & (row == 0): + bo.buffer_store( + code, rsrc(c["x8bs"]), t * (d.o_rows // 32) + task, cache_modifier=CM_DEV + ) + gpu.barrier() + + +def woa_loads(c, task): + a = c["args"] + return gemv_loads( + c, + a["woa"], + a["woa_s"], + c["d"].group_k, + task * WOA_ROWS, + WOA_SPLIT_K, + tiled=True, + ) + + +def wob_loads(c, task, pred=None): + a = c["args"] + return gemv_loads(c, a["wob"], a["wob_s"], c["d"].o_rows, task * ROWS, pred=pred) + + +def _wob_store(a, t, col, v0, v1): + word = ( + fx.Vector.from_elements([v0, v1], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Int32)[0] + ) + bo.buffer_store(word, rsrc(a["out"]), (t * a["out_stride"] + col) // 2) + + +@traced +def stage_wob(c, task, ops, task2, ops2, has2, t0, n, mark=False): + """Rows 16 task .. (and 16 task2 .., ``has2``) of wo_b -> bf16 -> out, or + (``c["wob_out"]``, the mono layer) the caller's hook with each bf16 pair.""" + tid, red = c["tid"], c["red"] + o_rows = c["d"].o_rows + a = c["args"] + gemv_mfmas(c, o_rows, c["xl"], c["xsl"], red, ops, rows=n) + gemv_mfmas(c, o_rows, c["xl"], c["xsl"], red + RED, ops2, rows=n) + gpu.barrier() + if const_expr(mark): + stamp(c, 15) + per = ROWS * n // 2 # bf16 pairs a task + k = tid // per + emit = c.get("wob_out") or (lambda t, col, v0, v1: _wob_store(a, t, col, v0, v1)) + if (tid < 2 * per) & ((k == 0) | has2): + u = tid % per + rp = u % (ROWS // 2) + tl = u // (ROWS // 2) + t = plus(t0, tl) + rk = red + k * RED + col = (k == 0).select(task, task2) * ROWS + 2 * rp + v0 = row_sum(rk, 2 * rp, tl) + v1 = row_sum(rk, 2 * rp + 1, tl) + emit(t, col, v0, v1) + gpu.barrier() + + +# ---------------------------------------------------------------- kernel + + +def back_plan(key): + """(nchunk, n_split, n_comb, COMB0, WOA0, WOB0, early_wob) of a build.""" + s, d = key.tokens, Dims(key.tp) + ck = chunk_keys(d, s) + nchunk = cdiv(KEYS, ck) + n_split = s * nchunk + n_comb = s * (d.heads // COMB_HEADS) + COMB0 = n_split % BLOCKS + assert d.woa_tasks <= BLOCKS + assert BLOCKS < d.wob_tasks <= 2 * BLOCKS + # the GEMVs past the attention units, so CTAs without attention work have + # their weights in flight from the start + WOA0 = (n_split + n_comb) % BLOCKS + WOB0 = (WOA0 + d.woa_tasks) % BLOCKS + # wo_b's weights before the wo_a stage only while the wo_a CTAs have no + # attention work (loads complete in issue order: else they would hold up + # wo_a's activation rows) + early_wob = n_split + n_comb + d.woa_tasks <= BLOCKS + return nchunk, n_split, n_comb, COMB0, WOA0, WOB0, early_wob + + +def back_lds_members(s: int, d: Dims) -> dict: + """K2's attention LDS structs: the split stage's and the GEMVs' (a union's + members: one at a time).""" + rows = tile_rows(s) + ck = chunk_keys(d, s) + + @fx.struct + class SplitLds: + q: fx.Array[fx.Int32, d.heads * Q_ROW, 16] + raw: fx.Array[fx.Int32, ck * REC // 4, 16] + p: fx.Array[fx.Int32, d.heads * p_row(d, s), 16] + st: fx.Array[fx.Float32, 2 * d.heads * (WAVES // head_groups(d, s)), 16] + + @fx.struct + class GemvLds: + red: fx.Array[fx.Float32, 2 * RED, 16] + xl: fx.Array[fx.Int32, rows * lds_row(d.group_k // 4), 16] + xsl: fx.Array[fx.Int32, rows * lds_row(d.group_k // 32), 16] + + return {"split": SplitLds, "gemv": GemvLds} + + +def back_lds(s: int, d: Dims): + """K2's attention LDS: the union of ``back_lds_members``.""" + m = back_lds_members(s, d) + SplitLds, GemvLds = m["split"], m["gemv"] + + @fx.union + class BackLds: + split: SplitLds # type: ignore[valid-type] + gemv: GemvLds # type: ignore[valid-type] + + return BackLds + + +def back_context(key, lds, args: dict, scratch, tag, layout=None) -> dict: + """The attention back's context: LDS views (``lds``: an allocated + ``back_lds``, or a dict of the members' views), mailbox (tag ``tag``), + scratch regions, arguments by name.""" + s, d = key.tokens, Dims(key.tp) + tid = fx.thread_idx.x + if isinstance(lds, dict): + spl, gml = lds["split"], lds["gemv"] + else: + spl, gml = lds.split.peek(), lds.gemv.peek() + c = { + "S": s, + "tid": tid, + "bid": fx.block_idx.x, + "lane": tid % 64, + "wave": tid // 64, + "red": gml.red.ptr, + "xl": gml.xl.ptr, + "xsl": gml.xsl.ptr, + "d": d, + "ratio": key.ratio, + "mb": Mailbox(tag), + "nchunk": back_plan(key)[0], + "sl": {"q": spl.q.ptr, "raw": spl.raw.ptr, "p": spl.p.ptr, "st": spl.st.ptr}, + "args": args, + } + for region, (off, _) in ( + layout or back_scratch(s, d, start=front_end(s, d)) + ).items(): + c[region] = scratch + fx.Int64(off) + return c + + +EPOCH_MARKS = 4 # word offset of the per-CTA marks in the epoch buffer + + +@traced +def epoch_begin(c, epoch, ep): + """This CTA has read the launch pair's epoch: its own mark word (no + contended atomic) := ep + 1, which ``epoch_end`` checks.""" + if c["tid"] == 0: + bo.buffer_store( + ep + 1, rsrc(epoch), EPOCH_MARKS + c["bid"], cache_modifier=CM_DEV + ) + + +@traced +def epoch_end(c, epoch, ep, reset=None): + """CTA 0, once every CTA marked this epoch (thread u checks CTA u's mark): + ``reset()`` (state later pairs must find cleared), then the epoch moved on. + The next pair's hand-offs carry a tag no earlier pair wrote; later launches' + kernels start after this one ends.""" + if c["bid"] == 0: + tid = c["tid"] + if tid < BLOCKS: + r = rsrc(epoch) + + def load_mark(): + return fx.Int32( + bo.buffer_load( + r, + EPOCH_MARKS + tid, + vec_width=1, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ) + + _ = spin_until(load_mark, lambda v: v != ep + 1) + gpu.barrier() + if tid == 0: + if const_expr(reset is not None): + reset() + gstore(epoch, ep + 1, words=1) + + +@traced +def run_back(c, key, bid): + """The attention back's stages on this CTA: split, combine, wo_a, wo_b.""" + s = key.tokens + dd = Dims(key.tp) + # plain ints: a Python object read inside traced control flow is carried + woa_tasks, groups, heads, group_k = dd.woa_tasks, dd.groups, dd.heads, dd.group_k + q_rows, o_rows, wob_tasks = dd.q_rows, dd.o_rows, dd.wob_tasks + a = c["args"] + tiles = token_tiles(s) + _, n_split, n_comb, COMB0, WOA0, WOB0, early_wob = back_plan(key) + stamp(c, 0) + for unit in range(first_task(bid, 0), n_split, BLOCKS): + stage_split(c, unit) + stamp(c, 1) + for unit in range(first_task(bid, COMB0), n_comb, BLOCKS): + stage_combine(c, unit) + gpu.barrier() + stamp(c, 2) + # the GEMV weights, every CTA right after its attention work: its wo_a + # task (if any), then its one or two wo_b tasks + woa_task = first_task(bid, WOA0) + has_woa = woa_task < woa_tasks + woa_t = fx.min(woa_task, woa_tasks - 1) + ops = gemv_loads( + c, + a["woa"], + a["woa_s"], + group_k, + woa_t * WOA_ROWS, + WOA_SPLIT_K, + tiled=True, + pred=has_woa, + ) + wob_task = first_task(bid, WOB0) + wob_second = wob_task + BLOCKS + has_second = wob_second < wob_tasks + + def wob_weights(): + return ( + wob_loads(c, wob_task), + wob_loads(c, fx.min(wob_second, wob_tasks - 1), pred=has_second), + ) + + if const_expr(early_wob): + ops1, ops2 = wob_weights() + if const_expr(c.get("tl") is not None): + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + stamp(c, 8) + # wo_a: every token tile of its group's XO rows, the next tile's in flight + if has_woa: + grp = woa_t // (woa_tasks // groups) + await_flags(c, c["crdy"], s, lambda u: u * (heads // COMB_HEADS) + grp) + stamp(c, 9) + xo = ( + c["xo"], + c["xos"], + q_rows // 4, + grp * group_k // 4, + group_k // 4, + q_rows // 32, + grp * group_k // 32, + group_k // 32, + ) + regs = rows_loads(c, *xo, *tiles[0]) + for ti in range_constexpr(len(tiles)): + t0, n = tiles[ti] + rows_store(c, regs, group_k // 4, group_k // 32, c["xl"], c["xsl"], n) + if const_expr(ti + 1 < len(tiles)): + regs = rows_loads(c, *xo, *tiles[ti + 1]) + stage_woa(c, woa_t, ops, t0, n) + stamp(c, 10) + publish(c, c["ardy"], woa_t) + stamp(c, 3) + # wo_b + if const_expr(not early_wob): + ops1, ops2 = wob_weights() + stamp(c, 11) + await_flags(c, c["ardy"], woa_tasks, lambda u: u) + stamp(c, 12) + x8b = ( + c["x8b"], + c["x8bs"], + o_rows // 4, + 0, + o_rows // 4, + o_rows // 32, + 0, + o_rows // 32, + ) + regs = rows_loads(c, *x8b, *tiles[0]) + for ti in range_constexpr(len(tiles)): + t0, n = tiles[ti] + rows_store(c, regs, o_rows // 4, o_rows // 32, c["xl"], c["xsl"], n) + if const_expr(ti == 0): + stamp(c, 14) + if const_expr(ti + 1 < len(tiles)): + regs = rows_loads(c, *x8b, *tiles[ti + 1]) + stage_wob( + c, + wob_task, + ops1, + fx.min(wob_second, wob_tasks - 1), + ops2, + has_second, + t0, + n, + mark=ti == 0, + ) + if const_expr(ti == 0): + stamp(c, 13) + stamp(c, 4) diff --git a/vllm/models/deepseek_v41/amd/mono/attention/device.py b/vllm/models/deepseek_v41/amd/mono/attention/device.py new file mode 100644 index 000000000000..ec230cd7074a --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/attention/device.py @@ -0,0 +1,504 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Device primitives of the DeepSeek-V4.1 attention mega kernel. + +Wave reductions and shuffles, hardware math, bf16 / fp8 / MX packing, the +scaled MFMA, buffer and 64-bit global accesses, and the tagged-pair mailbox. +Adapted from ATOM's mono decode framework (``atom/mono/device``), itself ported +from the FlyDSL GLM-5 layer kernel. +""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from aiter.ops.flydsl.kernels.dpp_utils import update_dpp_i32 +from aiter.ops.flydsl.kernels.kernels_common import kernel_signature +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm +from flydsl.compiler.ast_rewriter import ASTRewriter +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import T, as_ir_value + +# The execution model: one CTA a CU of an MI355X, all resident. +BLOCKS = 256 +THREADS = 512 +WAVES = THREADS // 64 + +# gfx95x cache policy bits (LLVM CPol): SC0 = 1, SC1 = 16, NT = 2. SC1 alone is +# device scope (past the per-XCD L2s); NT streams data read once. +CM_DEV = 16 +CM_NT = 2 + +FP8_MAX = 448.0 +FP8, FP4 = 0, 4 # f8f6f4 MFMA operand formats +UNIT_SCALE = 127 # E8M0 code of 1.0 +FLT_MIN = 1.1754943508222875e-38 + +POLL_MAX = 12 # mailbox pairs polled a batch + + +def traced(fn): + """FlyDSL's AST rewrite for a helper, so its ``if`` / ``while`` / ``for`` on + traced values lower to scf ops (kernels rewrite only their own body).""" + return ASTRewriter.transform(fn) + + +def kernel_symbol(stem, **params): + """Kernel symbol that profilers demangle to ``aiter::_``.""" + name = f"{stem}_{kernel_signature(**params)}" + return f"_ZN5aiter{len(name)}{name}E" + + +def first_task(bid, base): + """CTA ``bid``'s first task of a stage placed from CTA ``base``, then every + BLOCKS-th.""" + return (bid - base % BLOCKS + BLOCKS) % BLOCKS + + +# ----------------------------------------------------------------------------- memory + + +def rsrc(addr, nbytes=None): + return bo.create_buffer_resource_from_addr(addr, num_records_bytes=nbytes) + + +def ld_i32(ptr, i, cm=0): + return fx.Int32( + bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.i32, cache_modifier=cm) + ) + + +def ld_f32(ptr, i): + return fx.Float32(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.f32)) + + +def ld_bf(ptr, i): + return fx.Float32( + fx.BFloat16(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.bf16)) + ) + + +def ld_i64(ptr, i): + """Element i of an i64 buffer (two dwords).""" + w = fx.Vector(bo.buffer_load(rsrc(ptr), i * 2, vec_width=2, dtype=T.i32)) + return (fx.Int64(w[1]) << 32) | fx.Int64(fx.Uint32(w[0])) + + +def _gptr(addr): + return llvm.IntToPtrOp( + ir.Type.parse("!llvm.ptr<1>"), as_ir_value(fx.Int64(addr)) + ).result + + +def gload(addr, words=4, nt=False): + """``words`` i32 (1, 2 or 4) at a 64-bit global address: a cache past 4 GiB is + reachable only this way (a buffer offset is 32 bits).""" + i32 = ir.IntegerType.get_signless(32) + ty = i32 if words == 1 else ir.VectorType.get([words], i32) + v = llvm.LoadOp(ty, _gptr(addr), alignment=4 * words, nontemporal=nt or None).result + return fx.Int32(v) if words == 1 else fx.Vector(v) + + +def atomic_add_dev(addr, value): + """Device-scope atomic add on the global i32 at ``addr``; the old value. + Relaxed: an acquire / release here would write back and invalidate L2.""" + return fx.Int32( + fx.atomic_add( + _gptr(addr), + fx.Int32(value), + syncscope="agent", + ordering=fx.AtomicOrdering.Monotonic, + ) + ) + + +def gstore(addr, value, words=4): + """Store ``words`` i32 (an Int32 or a Vector) at a 64-bit global address.""" + v = as_ir_value(fx.Int32(value) if words == 1 else value) + llvm.StoreOp(v, _gptr(addr), alignment=4 * words) + + +# ----------------------------------------------------------------------------- math + + +def fresh(v): + """``v`` through an opaque move, so the lane math built on it is computed at + its use instead of being hoisted out of a task loop and held in registers.""" + return fx.Int32( + llvm.inline_asm( + T.i32, + [fx.Int32(v).ir_value()], + "v_mov_b32 $0, $1", + "=v,v", + has_side_effects=True, + ) + ) + + +def tile_map(tid, width): + """(tid // width, tid % width) for 0 <= tid < THREADS without a division: + shifts for a power of two, else a multiply-shift exact on that range.""" + if width & (width - 1) == 0: + sh = width.bit_length() - 1 + return tid >> sh, tid & (width - 1) + m = -(-(1 << 20) // width) + assert all((t * m) >> 20 == t // width for t in range(THREADS)), width + row = (tid * m) >> 20 + return row, tid - row * width + + +def uniform(v): + return fx.Int32(rocdl.readfirstlane(T.i32, fx.Int32(v).ir_value())) + + +def _hw_f32(name, x): + return fx.Float32( + llvm.call_intrinsic(T.f32, name, [fx.Float32(x).ir_value()], [], []) + ) + + +def hw_rsq(x): + return _hw_f32("llvm.amdgcn.rsq.f32", x) + + +def hw_rcp(x): + return _hw_f32("llvm.amdgcn.rcp.f32", x) + + +def hw_exp2(x): + return _hw_f32("llvm.amdgcn.exp2.f32", x) + + +def fma(a, b, c): + return fx.Float32(fmath.fma(fx.Float32(a), fx.Float32(b), fx.Float32(c))) + + +def div_rn(x, d, r): + """``x / d`` rounded to nearest (IEEE), given ``r = RN(1 / d)``.""" + q = x * r + return fma(fma(-q, d, x), r, q) + + +def permlane_swap(off, x, y): + swap = rocdl.permlane32_swap if off == 32 else rocdl.permlane16_swap + pr = swap( + llvm.StructType.get_literal([T.i32, T.i32]), + as_ir_value(fx.Int32(x)), + as_ir_value(fx.Int32(y)), + False, + False, + ) + return tuple(fx.Int32(llvm.extractvalue(T.i32, pr, [j])) for j in range(2)) + + +def lane_gather(v, src_lane): + """Lane ``src_lane``'s value of v (ds_bpermute).""" + return fx.Int32( + rocdl.ds_bpermute(T.i32, fx.Int32(src_lane) * 4, fx.Int32(v).ir_value()) + ) + + +def lane_f32(v, src_lane): + return lane_gather(v.bitcast(fx.Int32), src_lane).bitcast(fx.Float32) + + +def readlane(v, src_lane): + bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) + r = fx.Int32(rocdl.readlane(T.i32, bits, fx.Int32(src_lane))) + return r.bitcast(fx.Float32) if isinstance(v, fx.Float32) else r + + +def xshfl(v, off): + """Value of lane ``lane ^ off`` (DPP in a row, permlane swaps across).""" + assert off in (1, 2, 4, 8, 16, 32), off + if off >= 16: + return v.shuffle_xor(off, 64) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + if off == 8: + y = fx.Int32(update_dpp_i32(x, x, 0x118, 0xF, 0xC, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x108, 0xF, 0x3, False)) + elif off == 4: + y = fx.Int32(update_dpp_i32(x, x, 0x114, 0xF, 0xA, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x104, 0xF, 0x5, False)) + elif off == 2: + y = fx.Int32(update_dpp_i32(x, x, 0x4E, 0xF, 0xF, False)) + else: + y = fx.Int32(update_dpp_i32(x, x, 0xB1, 0xF, 0xF, False)) + return y.bitcast(fx.Float32) if is_f else y + + +def xred(v, off, op): + if off < 16: + return op(v, xshfl(v, off)) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + a, b = permlane_swap(off, x, x) + if is_f: + return op(a.bitcast(fx.Float32), b.bitcast(fx.Float32)) + return op(type(v)(a), type(v)(b)) + + +def _add(a, b): + return a + b + + +def butterfly(v, offsets, op=_add): + """Butterfly reduction over the xor ``offsets``, in that order.""" + for off in offsets: + v = xred(v, off, op) + return v + + +def wave_sum(v): + return butterfly(v, (32, 16, 8, 4, 2, 1)) + + +def wave_max(v): + return butterfly(v, (32, 16, 8, 4, 2, 1), fx.max) + + +def row_shr(v, lane, k): + """Lane - k of the same 16-lane row, else 0 (DPP row_shr).""" + moved = lane_gather(v.bitcast(fx.Int32), fx.max(lane - k, 0)).bitcast(fx.Float32) + return fx.Float32((lane % 16 >= k).select(moved, fx.Float32(0.0))) + + +# ----------------------------------------------------------------------------- packing + + +def bf16_round(a): + return fx.Float32(fx.Float32(a).to(fx.BFloat16)) + + +def bf16_pair(a, b): + """Two f32 -> one f32-typed word holding (bf16(a), bf16(b)).""" + return ( + fx.Vector.from_elements([a, b], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Float32)[0] + ) + + +def bf_lo(w): + return (w << 16).bitcast(fx.Float32) + + +def bf_hi(w): + return (w & fx.Int32(-65536)).bitcast(fx.Float32) + + +def fp8_pack4(a, b, c, d): + """Four f32 (in range) -> one dword of four E4M3 bytes, RNE.""" + lo = fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, a, b, fx.Int32(0), False)) + return fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, c, d, lo, True)) + + +def fp8x8_bf16(w0, w1, scale): + """Eight E4M3 of two dwords times an f32 ``scale`` -> eight bf16 as a + 4-dword Vector (exact for a power-of-two scale in range).""" + out = [] + for w in (w0, w1): + for hi in (False, True): + p = rocdl.cvt_scalef32_pk_bf16_fp8( + T.bf16x2, as_ir_value(fx.Int32(w)), as_ir_value(fx.Float32(scale)), hi + ) + out.append(fx.Vector(p).bitcast(fx.Int32)[0]) + return fx.Vector.from_elements(out, fx.Int32) + + +def clamp_fp8(v): + return fx.max(fx.min(v, fx.Float32(FP8_MAX)), fx.Float32(-FP8_MAX)) + + +def pow2(code): + """The fp32 power of two a biased exponent (1..254) stands for.""" + return (code << 23).bitcast(fx.Float32) + + +def rcp_pow2(code): + """1 / ``pow2(code)``, exact for code in [1, 253].""" + return ((254 - code) << 23).bitcast(fx.Float32) + + +def ceil_exp(v): + """Biased exponent of the smallest power of two >= v (v > 0 finite, normal).""" + bits = v.bitcast(fx.Int32) + return ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).select( + fx.Int32(1), fx.Int32(0) + ) + + +def mx_code(amax): + """VLLM's ROCm MXFP8 scale code of a 32-group: clamp(ceil(log2(amax / 448)) + + 127, 0, 254), amax floored at FLT_MIN.""" + a = fx.max(amax, fx.Float32(FLT_MIN)) + q = div_rn(a, fx.Float32(FP8_MAX), fx.Float32(1.0 / FP8_MAX)) + return fx.min(fx.max(ceil_exp(q), fx.Int32(0)), fx.Int32(254)) + + +def mx_mul(code): + """2^(127 - code) as fp32, for code in [0, 254] (the quant multiplier).""" + return ((254 - code) << 23).bitcast(fx.Float32) + + +# ----------------------------------------------------------------------------- MFMA + + +def mfma_scaled(a, b, c, sa, sb, a_fmt=FP8, b_fmt=FP8): + """16x16x128 scaled MFMA (gfx950 f8f6f4): lane l (row l % 16) holds K bytes + 16 (l // 16) .. +16 and 64 + 16 (l // 16) .. +16 of the 128-step, and the + E8M0 of K block l // 16.""" + return fx.Vector( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + T.vec(4, T.f32), [a, b, c, a_fmt, b_fmt, 0, sa, 0, sb] + ) + ) + + +def mfma_bf16(a, b, c): + """16x16x32 bf16 MFMA: lane l (row / col l % 16) holds K 8 (l // 16) .. +8.""" + return fx.Vector(rocdl.mfma_f32_16x16x32_bf16(T.vec(4, T.f32), [a, b, c])) + + +# ----------------------------------------------------------------------------- mailbox + + +@traced +def _spin(load_all, pending): + v = load_all() + while pending(v): + rocdl.s_sleep(1) + v = load_all() + return v + + +def spin_until(load, pending): + """``load()`` again (s_sleep between) while ``pending(value)``; the value.""" + return _spin(load, pending) + + +class Mailbox: + """Tagged-pair hand-off between the CTAs of one launch (device scope). + + Every 32-bit value is stored next to the launch's tag in one 8-byte store; a + consumer polls until every tag matches, so a hand-off is one round trip with + no separate flag. The tag is ``epoch * 64 + layer + 1`` (``epoch`` a device + counter bumped once a step): a pair from an earlier layer or step never + carries a live tag, and nothing is cleared between launches.""" + + def __init__(self, tag): + self.tag = tag + + def put(self, base, i, v): + bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) + bo.buffer_store( + fx.Vector.from_elements([bits, self.tag], fx.Int32), + rsrc(base), + i * 2, + cache_modifier=CM_DEV, + ) + + def put_bf(self, base, i, vs): + """Elements i .. i + len(vs) (2 or 4, i aligned) as bf16 pairs, one store.""" + words = [] + for j in range_constexpr(len(vs) // 2): + words += [bf16_pair(vs[2 * j], vs[2 * j + 1]).bitcast(fx.Int32), self.tag] + bo.buffer_store( + fx.Vector.from_elements(words, fx.Int32), + rsrc(base), + i, + cache_modifier=CM_DEV, + ) + + def put_words(self, base, i, words): + """Raw 32-bit words (1 or 2) at pairs i.. (fp8 words must come this way).""" + vals = [] + for w in words: + vals += [fx.Int32(w), self.tag] + bo.buffer_store( + fx.Vector.from_elements(vals, fx.Int32), + rsrc(base), + i * 2, + cache_modifier=CM_DEV, + ) + + def poll(self, specs, batch=POLL_MAX): + """Batched poll of ``specs`` = [(base, pair index, npairs in {1, 2})]: all + re-loaded while any tag is stale. Returns the value words.""" + if const_expr(len(specs) == 0): + return [] + if const_expr(len(specs) > batch): + return self.poll(specs[:batch], batch) + self.poll(specs[batch:], batch) + tag = self.tag + + def load_all(): + words = [] + for b, i, n in specs: + w = fx.Vector( + bo.buffer_load( + rsrc(b), + fx.Int32(i) * 2, + vec_width=2 * n, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ) + words += [w[e] for e in range(2 * n)] + return fx.Vector.from_elements(words, fx.Int32) + + nw = sum(2 * n for _, _, n in specs) + + def pending(v): + bad = v[1] != tag + for e in range_constexpr(3, nw, 2): + bad = bad | (v[e] != tag) + return bad + + v = _spin(load_all, pending) + outs, e = [], 0 + for _, _, n in specs: + outs.append([v[e + 2 * q] for q in range(n)]) + e += 2 * n + return outs + + +LOAD_BATCH = 4 + + +@traced +def batched_rounds(tid, n, load, use, batch=LOAD_BATCH): + """A CTA's strided loop over v = tid + THREADS i < n: ``batch`` rounds' loads + issued before their uses; a round past n loads element n - 1, uses nothing.""" + rounds = -(-n // THREADS) + for b0 in range_constexpr(0, rounds, batch): + idx = range(b0, min(b0 + batch, rounds)) + got = [load(fx.min(tid + THREADS * i, n - 1)) for i in idx] + for i, g in zip(idx, got): + v = tid + THREADS * i + if v < n: + use(v, g) + + +def barrier(): + gpu.barrier() + + +# ----------------------------------------------------------------------------- timeline + + +def memrealtime(): + """s_memrealtime: the 100 MHz device clock.""" + return fx.Int64(llvm.call_intrinsic(T.i64, "llvm.amdgcn.s.memrealtime", [], [], [])) + + +@traced +def stamp(c, point): + """A timeline build's stamp: thread 0 of the CTA stores the clock at + ``point`` (``c["tl"]`` [BLOCKS, points] i64; None: not a timeline build).""" + if const_expr(c.get("tl") is None): + return + if c["tid"] == 0: + t = memrealtime() + w = fx.Vector.from_elements([fx.Int32(t), fx.Int32(t >> 32)], fx.Int32) + bo.buffer_store(w, rsrc(c["tl"]), (c["bid"] * c["tl_points"] + point) * 2) diff --git a/vllm/models/deepseek_v41/amd/mono/attention/front.py b/vllm/models/deepseek_v41/amd/mono/attention/front.py new file mode 100644 index 000000000000..3ac9e3fe9b10 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/attention/front.py @@ -0,0 +1,662 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K1 ``front``: one DeepSeek-V4.1 attention layer from its normed input row to +the rotated query and the step's KV rows in the sliding-window cache. + +One launch a layer and rank, ``BLOCKS`` CTAs x ``THREADS`` threads, S rows. +Every CTA issues its weight loads (its wqkv part and its wq_b tasks) at the +kernel start, before it waits on anything; the stages: + + xq (S > 16: S x 2 K halves, on the 32 CTAs without a wqkv part): the + bf16 input's half -> vLLM's MXFP8 -> X8 (published, a flag each) + kt (S): the token's attention keys for K2 -- top-k compressed slots + through the kv source's block table, then window slots -> KT + wqkv_a (112 row tiles x 2 K halves, 16 rows): the input's MXFP8 half -- + quantized in the CTA (S <= 16) or X8's -- a token tile at a time + (the next tile's rows in flight) times wqkv's -> fp32 partials -> + QKVP + qkv (S): the halves summed -> bf16; q_norm -> bf16 -> MXFP8 -> QX8 + (published, a flag a token); kv_norm -> bf16 -> GPT-J RoPE -> bf16 + -> the fp8_ds_mla record at the token's slot + wq_b (H x 512 / 16, 16 rows): QX8 times wq_b -> bf16 -> GPT-J RoPE of + each head's dims 448.. -> q (global, bf16) + +Each elementwise step follows vLLM's ROCm path (``mxfp8_quantize``, +``FusedQKVRMSNorm``, ``fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert``); +the GEMV reductions differ from the Triton GEMMs' only in their order. +Adapted from ATOM's V4.1 ``attn_pre`` (K1), less its mHC seam. +""" + +from dataclasses import dataclass + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T + +from .device import ( + BLOCKS, + CM_DEV, + THREADS, + Mailbox, + bf16_round, + bf_hi, + bf_lo, + butterfly, + ceil_exp, + clamp_fp8, + div_rn, + first_task, + fp8_pack4, + fresh, + gstore, + hw_rsq, + ld_f32, + ld_i32, + mx_code, + mx_mul, + rsrc, + stamp, + tile_map, + traced, + xshfl, +) +from .gemv import RED, gemv_loads, gemv_mfmas, lds_row, row_sum, rows_loads, rows_store +from .plan import ( + DATA, + EPS, + HALF, + HEAD_DIM, + HIDDEN, + KEYS, + KV_DIM, + NOPE, + Q_RANK, + QKV_ROWS, + ROPE, + ROWS, + TILE, + TOPK, + WINDOW, + WQKV_KPARTS, + WQKV_TASKS, + Dims, + front_scratch, + tile_rows, + token_tiles, +) + +_STREAM = fx.Stream(None) +X_WORDS = HIDDEN // 4 +X_GROUPS = HIDDEN // 32 +QX_WORDS = Q_RANK // 4 +KSTEPS = HIDDEN // 128 // WQKV_KPARTS # a wqkv part's 128-steps +KWORDS = KSTEPS * 32 +KGROUPS = KSTEPS * 4 +XH = HIDDEN // WQKV_KPARTS # a K half's values +XH_COLS = XH // 8 // 2 # 16 B items (8 values) a half of a K half +XH_RPR = THREADS // XH_COLS # half-rows a round +XQ_ITEMS = XH // 8 # an xq unit's 16 B items (one a thread) +FRONT_POINTS = 6 # start, kt done, wqkv done, qkv done, wq_b done, (spare) + + +@dataclass(frozen=True) +class FrontBuild: + tokens: int + tp: int + ratio: int # the layer's compress ratio (0: no top-k keys) + timeline: bool = False + + +def plus(t0, v): + return v if isinstance(t0, int) and t0 == 0 else v + t0 + + +@traced +def publish(c, region, i): + """This CTA's stores visible device-wide, then the flag pair (region, i).""" + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if c["tid"] == 0: + c["mb"].put(region, i, fx.Int32(1)) + + +@traced +def await_flags(c, region, n, index): + """Thread u < n polls flag ``index(u)``; then the CTA proceeds.""" + tid = c["tid"] + if tid < n: + c["mb"].poll([(region, index(fx.min(tid, n - 1)), 1)]) + gpu.barrier() + + +# ---------------------------------------------------------------- xq + + +def xq_loads(c, unit, pred): + """(token, K half) ``unit``'s bf16 input: its load only, 8 values a thread + (threads past the half mirror its last item).""" + a = c["args"] + xr = rsrc(a["x"], pred.select(fx.Int64(0xFFFFFFFF), fx.Int64(0))) + col = fx.min(c["tid"], XQ_ITEMS - 1) + e = (unit >> 1) * a["x_stride"] + (unit & 1) * XH + col * 8 + return fx.Vector(bo.buffer_load(xr, e >> 1, vec_width=4, dtype=T.i32)) + + +@traced +def stage_xq(c, unit, w, live): + """``xq_loads``' values -> vLLM's MXFP8 -> X8 words / X8S codes of the + unit (published by the caller).""" + tid, lane = c["tid"], c["lane"] + vals = [] + for d in range_constexpr(4): + vals += [bf_lo(w[d]), bf_hi(w[d])] + amax = abs(vals[0]) + for v in vals[1:]: + amax = fx.max(amax, abs(v)) + amax = butterfly(amax, (1, 2), fx.max) + code = mx_code(amax) + mul = mx_mul(code) + w0 = fp8_pack4(*[clamp_fp8(vals[q] * mul) for q in range(4)]) + w1 = fp8_pack4(*[clamp_fp8(vals[4 + q] * mul) for q in range(4)]) + t = unit >> 1 + part = unit & 1 + if live & (tid < XQ_ITEMS): + bo.buffer_store( + fx.Vector.from_elements([w0, w1], fx.Int32), + rsrc(c["x8"]), + t * X_WORDS + part * KWORDS + tid * 2, + cache_modifier=CM_DEV, + ) + if lane % 4 == 0: + bo.buffer_store( + code, + rsrc(c["x8s"]), + t * X_GROUPS + part * KGROUPS + (tid >> 2), + cache_modifier=CM_DEV, + ) + + +# ---------------------------------------------------------------- kt + + +@traced +def stage_kt(c, t): + """Token t's keys for K2: ``KT[t][k] = slot | 1 << 31`` for its top-k compressed + rows (k < ntopk), the window slot for k in [ntopk, kv_len), -1 for an + invalid or absent key; KLEN[t] = kv_len (0 for a pad row), KLEN[S + t] = + ntopk.""" + tid = c["tid"] + a, ratio = c["args"], c["ratio"] + slot = ld_i32(a["slot"], 2 * t) + live = slot >= 0 + nswa = live.select(ld_i32(a["swa_lens"], t), fx.Int32(0)) + if const_expr(ratio > 0): + pos = ld_i32(a["pos"], 2 * t) + ntopk = live.select(fx.min((pos + 1) // ratio, fx.Int32(TOPK)), fx.Int32(0)) + req = ld_i32(a["t2r"], t) + else: + ntopk = fx.Int32(0) + kv_len = ntopk + nswa + for i in range_constexpr((KEYS + THREADS - 1) // THREADS): + k = fresh(tid) + THREADS * i + sidx = fx.min(fx.max(k - ntopk, 0), WINDOW - 1) + s_slot = ld_i32(a["swa_idx"], t * WINDOW + sidx) + val = (k < kv_len).select(s_slot, fx.Int32(-1)) + if const_expr(ratio > 0): + local = ld_i32(a["topk"], t * TOPK + fx.min(k, TOPK - 1)) + lc = fx.max(local, 0) + e = a["comp_block"] + blk = ld_i32(a["comp_bt"], req * a["bt_stride"] + lc // e) + c_slot = (blk * e + lc % e) | fx.Int32(-2147483648) + c_val = (local >= 0).select(c_slot, fx.Int32(-1)) + val = (k < ntopk).select(c_val, val) + if k < KEYS: + bo.buffer_store(val, rsrc(a["kt"]), t * KEYS + k) + if tid == 0: + bo.buffer_store(kv_len, rsrc(a["klen"]), t) + bo.buffer_store(ntopk, rsrc(a["klen"]), c["S"] + t) + + +# ---------------------------------------------------------------- wqkv_a + + +def x_half_loads(c, part, pred, t0, n): + """Tokens t0 .. t0 + n's bf16 input, K half ``part``: its loads only (16 B a + lane-item, 8 values; a thread one column of [2 n half-rows, 160]), masked + off on CTAs without a wqkv part.""" + a = c["args"] + xr = rsrc(a["x"], pred.select(fx.Int64(0xFFFFFFFF), fx.Int64(0))) + hr0, col = tile_map(c["tid"], XH_COLS) + e0 = part * XH + col * 8 + out = [] + for i in range((2 * n + XH_RPR - 1) // XH_RPR): + hr = fx.min(hr0 + XH_RPR * i, 2 * n - 1) + e = (t0 + (hr >> 1)) * a["x_stride"] + (hr & 1) * (XH // 2) + e0 + out.append(fx.Vector(bo.buffer_load(xr, e >> 1, vec_width=4, dtype=T.i32))) + return out + + +@traced +def quant_x_half(c, regs, n): + """``x_half_loads``' values -> vLLM's MXFP8 rows of the K half in LDS (a + 32-group = 4 lanes).""" + tid, lane = c["tid"], c["lane"] + x_row, s_row = lds_row(KWORDS), lds_row(KGROUPS) + hr0, col = tile_map(tid, XH_COLS) + act = tid < XH_RPR * XH_COLS + for i in range_constexpr(len(regs)): + hr = hr0 + XH_RPR * i + w = regs[i] + vals = [] + for d in range_constexpr(4): + vals += [bf_lo(w[d]), bf_hi(w[d])] + amax = abs(vals[0]) + for v in vals[1:]: + amax = fx.max(amax, abs(v)) + amax = butterfly(amax, (1, 2), fx.max) + code = mx_code(amax) + mul = mx_mul(code) + w0 = fp8_pack4(*[clamp_fp8(vals[q] * mul) for q in range(4)]) + w1 = fp8_pack4(*[clamp_fp8(vals[4 + q] * mul) for q in range(4)]) + tok = hr >> 1 + half = hr & 1 + if act & (hr < 2 * n): + fx.ptr_store( + fx.Vector.from_elements([w0, w1], fx.Int32), + c["xl"] + (tok * x_row + half * (KWORDS // 2) + col * 2), + ) + if lane % 4 == 0: + fx.ptr_store( + code, c["xsl"] + (tok * s_row + half * (KGROUPS // 2) + (col >> 2)) + ) + gpu.barrier() + + +@traced +def stage_wqkv(c, task, part, ops, t0, n, live): + """Rows 16 task .. of wqkv over K half ``part`` against tokens t0 .. t0 + n + -> fp32 partials QKVP[part] (``live``: a CTA with a wqkv part).""" + tid, red = c["tid"], c["red"] + gemv_mfmas(c, HIDDEN, c["xl"], c["xsl"], red, ops, rows=n, nsteps=KSTEPS) + gpu.barrier() + if (tid < ROWS * n) & live: + r = tid % ROWS + tl = tid // ROWS + t = plus(t0, tl) + v = row_sum(red, r, tl) + c["mb"].put(c["qkvp"], (part * c["S"] + t) * QKV_ROWS + task * ROWS + r, v) + gpu.barrier() + + +# ---------------------------------------------------------------- q / kv norms + + +@traced +def stage_qkv(c, t): + """Token t: the K halves summed -> bf16; q_norm -> bf16 -> MXFP8 (QX8); + kv_norm -> bf16 -> RoPE -> bf16 -> the fp8_ds_mla record at + slot_mapping[t]. 256 threads (8 q columns and 2 kv columns each); threads + 256.. mirror thread 255 and store nothing.""" + tid, lane, wave, red = c["tid"], c["lane"], c["wave"], c["red"] + a, s = c["args"], c["S"] + tt = fx.min(tid, 255) + mine = tid < 256 + qc = [fx.min(tt * 8 + j, Q_RANK - 1) for j in range(8)] + # independent loads first: they overlap the partials' wait + qw = [c["ldbf"](a["qn_w"], qc[j]) for j in range(8)] + kw = [c["ldbf"](a["kvn_w"], tt * 2 + j) for j in range(2)] + pos = ld_i32(a["pos"], 2 * t) + slot = ld_i32(a["slot"], 2 * t) + jr = tt - NOPE // 2 + is_rope = jr >= 0 + jj = fx.max(jr, 0) + cs = is_rope.select(ld_f32(a["cos_sin"], pos * ROPE + jj), fx.Float32(1.0)) + sn = is_rope.select(ld_f32(a["cos_sin"], pos * ROPE + HALF + jj), fx.Float32(0.0)) + specs = [] + for p in range_constexpr(2): + specs += [(c["qkvp"], (p * s + t) * QKV_ROWS + qc[j], 1) for j in range(8)] + specs += [ + (c["qkvp"], (p * s + t) * QKV_ROWS + Q_RANK + tt * 2 + j, 1) + for j in range(2) + ] + got = c["mb"].poll(specs) + vals = [ + bf16_round(got[j][0].bitcast(fx.Float32) + got[10 + j][0].bitcast(fx.Float32)) + for j in range(10) + ] + xq = [((tt * 8 + j) < Q_RANK).select(vals[j], fx.Float32(0.0)) for j in range(8)] + xk = vals[8:10] + ssq = xq[0] * xq[0] + for j in range_constexpr(1, 8): + ssq = ssq + xq[j] * xq[j] + ssq = butterfly(ssq, (32, 16, 8, 4, 2, 1)) + ssk = butterfly(xk[0] * xk[0] + xk[1] * xk[1], (32, 16, 8, 4, 2, 1)) + if (lane == 0) & mine: + fx.ptr_store(ssq, red + wave) + fx.ptr_store(ssk, red + 4 + wave) + gpu.barrier() + tq = (fx.ptr_load(red + 0) + fx.ptr_load(red + 1)) + ( + fx.ptr_load(red + 2) + fx.ptr_load(red + 3) + ) + tk = (fx.ptr_load(red + 4) + fx.ptr_load(red + 5)) + ( + fx.ptr_load(red + 6) + fx.ptr_load(red + 7) + ) + rq = hw_rsq( + div_rn(tq, fx.Float32(float(Q_RANK)), fx.Float32(1.0 / Q_RANK)) + + fx.Float32(EPS) + ) + rk = hw_rsq( + div_rn(tk, fx.Float32(float(KV_DIM)), fx.Float32(1.0 / KV_DIM)) + + fx.Float32(EPS) + ) + # q: bf16 normed row, then vLLM's MXFP8 (a 32-group = 4 threads) + ys = [bf16_round((xq[j] * rq) * qw[j]) for j in range(8)] + amax = abs(ys[0]) + for y in ys[1:]: + amax = fx.max(amax, abs(y)) + code = mx_code(butterfly(amax, (1, 2), fx.max)) + mul = mx_mul(code) + if mine & (tt * 8 < Q_RANK): + w0 = fp8_pack4(*[clamp_fp8(ys[i] * mul) for i in range(4)]) + w1 = fp8_pack4(*[clamp_fp8(ys[4 + i] * mul) for i in range(4)]) + bo.buffer_store( + fx.Vector.from_elements([w0, w1], fx.Int32), + rsrc(c["qx8"]), + t * QX_WORDS + 2 * tt, + cache_modifier=CM_DEV, + ) + if tt % 4 == 0: + bo.buffer_store( + code, + rsrc(c["qx8s"]), + t * (Q_RANK // 32) + tt // 4, + cache_modifier=CM_DEV, + ) + # kv: this thread's pair (2 tt, 2 tt + 1) + kn = [bf16_round((xk[j] * rk) * kw[j]) for j in range(2)] + e0 = bf16_round(kn[0] * cs - kn[1] * sn) + e1 = bf16_round(kn[0] * sn + kn[1] * cs) + # record quant of the nope dims: a 64-block = 32 threads (half a wave) + amax = butterfly(fx.max(abs(e0), abs(e1)), (1, 2, 4, 8, 16), fx.max) + rcode = ceil_exp( + div_rn( + fx.max(amax, fx.Float32(1e-4)), fx.Float32(448.0), fx.Float32(1.0 / 448.0) + ) + ) + inv = mx_mul(rcode) + q0 = clamp_fp8(e0 * inv) + q1 = clamp_fp8(e1 * inv) + n0 = xshfl(q0, 1) + n1 = xshfl(q1, 1) + word = fp8_pack4(q0, q1, n0, n1) + rope_word = ( + fx.Vector.from_elements([e0, e1], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Int32)[0] + ) + blk = a["swa_block"] + sblock = slot // blk + soff = slot % blk + base = a["swa"] + fx.Int64(sblock) * fx.Int64(a["swa_stride"]) + rec = base + fx.Int64(soff * DATA) + if mine & (tt % 32 == 0) & (tt < NOPE // 2): + fx.ptr_store(rcode.bitcast(fx.Float32), red + 8 + tt // 32) + gpu.barrier() + if mine & (slot >= 0): + if tt < NOPE // 2: + if tt % 2 == 0: + gstore(rec + fx.Int64(2 * tt), word, words=1) + else: + gstore(rec + fx.Int64(NOPE + 4 * (tt - NOPE // 2)), rope_word, words=1) + if tt == 0: + sb = [fx.ptr_load(red + 8 + b).bitcast(fx.Int32) for b in range(7)] + w0 = sb[0] | (sb[1] << 8) | (sb[2] << 16) | (sb[3] << 24) + w1 = sb[4] | (sb[5] << 8) | (sb[6] << 16) + sc = base + fx.Int64(blk * DATA + soff * 8) + gstore(sc, fx.Vector.from_elements([w0, w1], fx.Int32), words=2) + publish(c, c["qrdy"], t) + + +# ---------------------------------------------------------------- wq_b + + +def wqb_rope(c, task0, pos_t, n): + """(cos, sin) of this thread's pair in ``stage_wqb``'s epilogue over a tile + of n rows (``pos_t``: its token's position; every task of the CTA shares + the pair's dims: tasks BLOCKS apart are whole heads apart).""" + a = c["args"] + rp = c["tid"] % (ROWS // 2) + dim = (task0 * ROWS) % HEAD_DIM + 2 * rp + jj = fx.max((dim - NOPE) >> 1, 0) + return ld_f32(a["cos_sin"], pos_t * ROPE + jj), ld_f32( + a["cos_sin"], pos_t * ROPE + HALF + jj + ) + + +def wqb_pos(c, t0, n): + """This thread's token position for ``stage_wqb``'s epilogue of tile + t0 .. t0 + n.""" + per = ROWS * n // 2 + tl = fx.min((c["tid"] % per) // (ROWS // 2), n - 1) + return ld_i32(c["args"]["pos"], 2 * (t0 + tl)) + + +@traced +def stage_wqb(c, task0, bops, t0, n, rope): + """Rows 16 (task0 + BLOCKS k) .. of the rank's H x 512 query rows, every + task k of the CTA: GEMVs -> bf16 -> GPT-J RoPE of a head's dims 448.. + (fp32) -> bf16 -> q.""" + tid, red, d = c["tid"], c["red"], c["d"] + a = c["args"] + nt = len(bops) + for k in range_constexpr(nt): + gemv_mfmas(c, Q_RANK, c["qxl"], c["qxsl"], red + k * RED, bops[k], rows=n) + gpu.barrier() + per = ROWS * n // 2 # bf16 pairs a task + if tid < nt * per: + k = tid // per + u = tid % per + rp = u % (ROWS // 2) + tl = u // (ROWS // 2) + t = plus(t0, tl) + r0 = (task0 + BLOCKS * k) * ROWS + 2 * rp + rk = red + k * RED + v0 = bf16_round(row_sum(rk, 2 * rp, tl)) + v1 = bf16_round(row_sum(rk, 2 * rp + 1, tl)) + is_rope = (task0 * ROWS) % HEAD_DIM + 2 * rp >= NOPE + cs, sn = rope + e0 = is_rope.select(v0 * cs - v1 * sn, v0) + e1 = is_rope.select(v0 * sn + v1 * cs, v1) + word = ( + fx.Vector.from_elements([e0, e1], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Int32)[0] + ) + bo.buffer_store(word, rsrc(a["q"]), (t * d.q_rows + r0) // 2) + gpu.barrier() + + +# ---------------------------------------------------------------- kernel + + +def front_lds(s: int, d: Dims): + """K1's LDS (the attention front's stages).""" + rows = tile_rows(s) + wqb_per_cta = d.wqb_tasks // BLOCKS + + @fx.struct + class FrontLds: + red: fx.Array[fx.Float32, wqb_per_cta * RED + 16, 16] + qxl: fx.Array[fx.Int32, rows * lds_row(QX_WORDS), 16] + qxsl: fx.Array[fx.Int32, rows * lds_row(Q_RANK // 32), 16] + xl: fx.Array[fx.Int32, rows * lds_row(KWORDS), 16] + xsl: fx.Array[fx.Int32, rows * lds_row(KGROUPS), 16] + + return FrontLds + + +def front_context(key, lds, args: dict, scratch, tag, layout=None) -> dict: + """The front stages' context: LDS views, mailbox (tag ``tag``), scratch + regions, arguments (``args``: the kernel's, by name).""" + s, d = key.tokens, Dims(key.tp) + tid = fx.thread_idx.x + c = { + "S": s, + "tid": tid, + "bid": fx.block_idx.x, + "lane": tid % 64, + "wave": tid // 64, + "red": lds.red.ptr, + "xl": lds.xl.ptr, + "xsl": lds.xsl.ptr, + "qxl": lds.qxl.ptr, + "qxsl": lds.qxsl.ptr, + "mb": Mailbox(tag), + "d": d, + "ratio": key.ratio, + "ldbf": lambda ptr, i: fx.Float32( + fx.BFloat16(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.bf16)) + ), + "args": args, + } + for region, (off, _) in (layout or front_scratch(s, d)).items(): + c[region] = scratch + fx.Int64(off) + return c + + +@traced +def run_front(c, key, bid, x8_given=False, before_wqkv=None, after_wqkv=None): + """The front's stages on this CTA. ``x8_given``: a stage before this one + (the mono layer's seam norm) publishes X8 / XRDY for every token (no input + quantization here). ``before_wqkv()`` runs ahead of the wqkv weights' loads + (a latency-critical stage of the caller's), ``after_wqkv()`` right after.""" + s = key.tokens + wqb_tasks = Dims(key.tp).wqb_tasks + a = c["args"] + tid = c["tid"] + tiles = token_tiles(s) + wqb_per_cta = wqb_tasks // BLOCKS + assert wqb_tasks % BLOCKS == 0 + n_wqkv = WQKV_TASKS * WQKV_KPARTS # 224 + assert n_wqkv <= BLOCKS + n_free = BLOCKS - n_wqkv # CTAs without a wqkv part: xq, kt + QKV0 = 0 # on wqkv CTAs + WQB0 = 0 + in_cta_quant = s <= TILE and not x8_given + xq_stage = s > TILE and not x8_given + xq_units = WQKV_KPARTS * s + xq_per = (xq_units + n_free - 1) // n_free + stamp(c, 0) + + # ---- Loads complete in issue order, so every CTA issues its latency- + # critical loads before its later stages' weights: a wqkv CTA its first + # input tile and its wqkv part, the next tiles while it quantizes, then + # its wq_b tasks; a CTA with a key-table task those after that work. + # The q / kv norms run on wqkv CTAs, their partials' wait behind the + # wq_b weights that are due by then anyway. + wqkv_unit = first_task(bid, 0) + has_wqkv = wqkv_unit < n_wqkv + task = fx.min(wqkv_unit, n_wqkv - 1) % WQKV_TASKS + part = fx.min(wqkv_unit, n_wqkv - 1) // WQKV_TASKS + wqb_task = first_task(bid, WQB0) + free = bid - n_wqkv + is_free = free >= 0 + if const_expr(in_cta_quant): + xregs = [x_half_loads(c, part, has_wqkv, *tiles[0])] + if const_expr(xq_stage): + xunits = [free + n_free * i for i in range(xq_per)] + xlive = [is_free & (u < xq_units) for u in xunits] + xq_regs = [ + xq_loads(c, fx.min(fx.max(u, 0), xq_units - 1), lv) + for u, lv in zip(xunits, xlive) + ] + if const_expr(before_wqkv is not None): + before_wqkv() + kops = gemv_loads( + c, + a["wqkv"], + a["wqkv_s"], + HIDDEN, + task * ROWS, + kbeg=part * KSTEPS, + nsteps=KSTEPS, + pred=has_wqkv, + ) + if const_expr(after_wqkv is not None): + after_wqkv() + if const_expr(xq_stage): # noqa: SIM102 (a build-time and a run-time test) + if is_free: + for i in range_constexpr(xq_per): + stage_xq( + c, fx.min(fx.max(xunits[i], 0), xq_units - 1), xq_regs[i], xlive[i] + ) + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if tid < xq_per: + u = free + n_free * tid + if u < xq_units: + c["mb"].put(c["xrdy"], u, fx.Int32(1)) + for t in range(is_free.select(free, fx.Int32(s)), s, n_free): + stage_kt(c, t) + stamp(c, 1) + poss = [wqb_pos(c, t0, n) for t0, n in tiles] + + # ---- wqkv_a: one (row tile, K half) a CTA, every token tile; the wq_b + # weights issued behind the input rows' loads + def wqb_weights(): + return [ + gemv_loads(c, a["wqb"], a["wqb_s"], Q_RANK, (wqb_task + BLOCKS * i) * ROWS) + for i in range(wqb_per_cta) + ] + + if const_expr(in_cta_quant): + if has_wqkv: + quant_x_half(c, xregs[0], tiles[0][1]) + bops = wqb_weights() + if has_wqkv: + stage_wqkv(c, task, part, kops, *tiles[0], has_wqkv) + else: + await_flags(c, c["xrdy"], s, lambda u: WQKV_KPARTS * u + part) + xr = ( + c["x8"], + c["x8s"], + X_WORDS, + part * KWORDS, + KWORDS, + X_GROUPS, + part * KGROUPS, + KGROUPS, + ) + regs = rows_loads(c, *xr, tiles[0][0], tiles[0][1], pred=has_wqkv) + if const_expr(len(tiles) == 1): + bops = wqb_weights() + for ti in range_constexpr(len(tiles)): + t0, n = tiles[ti] + rows_store(c, regs, KWORDS, KGROUPS, c["xl"], c["xsl"], n) + if const_expr(ti + 1 < len(tiles)): + t1, n1 = tiles[ti + 1] + regs = rows_loads(c, *xr, t1, n1, pred=has_wqkv) + if const_expr(ti + 2 == len(tiles)): + bops = wqb_weights() + stage_wqkv(c, task, part, kops, t0, n, has_wqkv) + stamp(c, 2) + + # ---- q / kv norms, KV insert + for t in range(first_task(bid, QKV0), s, BLOCKS): + stage_qkv(c, t) + stamp(c, 3) + + # ---- wq_b: every task's rows of a token tile at a time, the next + # tile's rows in flight + ropes = [wqb_rope(c, wqb_task, poss[ti], tiles[ti][1]) for ti in range(len(tiles))] + await_flags(c, c["qrdy"], s, lambda u: u) + qx = (c["qx8"], c["qx8s"], QX_WORDS, 0, QX_WORDS, Q_RANK // 32, 0, Q_RANK // 32) + regs = rows_loads(c, *qx, *tiles[0]) + for ti in range_constexpr(len(tiles)): + t0, n = tiles[ti] + rows_store(c, regs, QX_WORDS, Q_RANK // 32, c["qxl"], c["qxsl"], n) + if const_expr(ti + 1 < len(tiles)): + regs = rows_loads(c, *qx, *tiles[ti + 1]) + stage_wqb(c, wqb_task, bops, t0, n, ropes[ti]) + stamp(c, 4) diff --git a/vllm/models/deepseek_v41/amd/mono/attention/gemv.py b/vllm/models/deepseek_v41/amd/mono/attention/gemv.py new file mode 100644 index 000000000000..9471a1e3f4df --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/attention/gemv.py @@ -0,0 +1,234 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MXFP8 GEMV of the mega kernel: 16 rows of a row-major e4m3 weight [N, K] +(vLLM's loaded layout, E8M0 scales in 32x32 blocks [N / 32, K / 32]) against +every token of a tile, its MXFP8 rows in LDS, on the gfx950 scaled MFMA. + +A task may cover a window of K (``kbeg`` .. ``kbeg + nsteps`` 128-steps): a +split-K part, its activations in LDS holding just that window. The weight +operands are loaded first (all in flight, before the task waits on its +activations); wave w takes the window's steps w, w + 8, ... (or, ``split``, a +contiguous range of waves 0 .. split - 1), its 16 x 16 partial in ``red``, +summed in wave order by ``row_sum``. Adapted from ATOM's +``attn_pre.gemv_fp8_*``, which reads aiter's (16, 16)-preshuffled layout. +""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr +from flydsl.expr.typing import T + +from .device import THREADS, WAVES, mfma_scaled, rsrc, tile_map, traced + +# a token's LDS row of MXFP8 words / E8M0 codes, padded by 16 B so 16 rows read +# by one MFMA operand do not share a bank group +LDS_PAD = 4 +RED = WAVES * 64 * 4 # one task's accumulators in ``red`` (f32) + + +def lds_row(words): + return words + LDS_PAD + + +def lds_at(u, row): + """Element u of a span of rows ``row`` wide, at the ``lds_row`` stride.""" + return u + u // row * LDS_PAD + + +def _steps(c, nsteps, split, tiled): + """This wave's steps of the window: [(local step, live)].""" + wave = c["wave"] + if const_expr(split is None): + per = (nsteps + WAVES - 1) // WAVES + ragged = nsteps % WAVES != 0 + return [ + (wave + WAVES * i, (wave + WAVES * i < nsteps) if ragged else None) + for i in range(per) + ] + assert split <= WAVES and nsteps % split == 0 + per = nsteps // split + part = wave % split if tiled else fx.min(wave, split - 1) + return [(part * per + i, None) for i in range(per)] + + +def _step_load(c, r_w, r_ws, k, row0, st, pred=None): + """Global step ``st``'s weight operand of rows row0 .. +16: two 16 B chunks a + lane (K bytes 16 j .. and 64 + 16 j .. of the step, row row0 + lane % 16) + and the E8M0 byte of its K block. ``pred`` false: masked (out of range: no + memory traffic, zeros).""" + lane = c["lane"] + j = lane // 16 + row = row0 + lane % 16 + base = row * k + 128 * st + 16 * j + lo = fx.Vector(bo.buffer_load(r_w, base // 4, vec_width=4, dtype=T.i32)) + hi = fx.Vector(bo.buffer_load(r_w, (base + 64) // 4, vec_width=4, dtype=T.i32)) + kg = k // 32 + wsi = (row0 // 32) * kg + st * 4 + j + word = fx.Int32(bo.buffer_load(r_ws, wsi // 4, vec_width=1, dtype=T.i32)) + return lo, hi, (word >> (wsi % 4 * 8)) & 0xFF + + +def gemv_loads( + c, w, ws, k, row0, split=None, tiled=False, kbeg=0, nsteps=None, pred=None +): + """The weight operands of rows row0 .. +16 (``tiled``: row0 + 16 (wave // + split)) over steps kbeg .. kbeg + nsteps for this wave, all in flight + (``pred`` false: masked off, no memory traffic).""" + nsteps = k // 128 if nsteps is None else nsteps + if const_expr(pred is None): + r_w, r_ws = rsrc(w), rsrc(ws) + else: + # a masked-off CTA reads through a zero-size buffer: zeros, no traffic + nb = pred.select(fx.Int64(0xFFFFFFFF), fx.Int64(0)) + r_w, r_ws = rsrc(w, nb), rsrc(ws, nb) + if const_expr(tiled): + row0 = row0 + 16 * (c["wave"] // split) + return [ + _step_load( + c, + r_w, + r_ws, + k, + row0, + kbeg + (st if live is None else fx.min(st, nsteps - 1)), + pred, + ) + for st, live in _steps(c, nsteps, split, tiled) + ] + + +def _step_mfma(c, nsteps, xl, xsl, st, ops, acc, live, rows): + lane = c["lane"] + lo, hi, sa = ops + x_row, s_row = lds_row(nsteps * 32), lds_row(nsteps * 4) + j = lane // 16 + col = fx.min(lane % 16, rows - 1) + if const_expr(live is not None): + st = fx.min(st, nsteps - 1) + av = fx.Vector.from_elements( + [lo[d] for d in range(4)] + [hi[d] for d in range(4)], fx.Int32 + ) + xlo = fx.Vector( + fx.ptr_load( + xl + (col * x_row + st * 32 + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + xhi = fx.Vector( + fx.ptr_load( + xl + (col * x_row + st * 32 + 16 + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + xv = fx.Vector.from_elements( + [xlo[d] for d in range(4)] + [xhi[d] for d in range(4)], fx.Int32 + ) + sb = fx.ptr_load(xsl + (col * s_row + st * 4 + j)) + if const_expr(live is not None): + # a step past the end adds zero at unit scale + av = fx.Vector.from_elements( + [live.select(av[d], fx.Int32(0)) for d in range(8)], fx.Int32 + ) + sa = live.select(sa, fx.Int32(127)) + sb = live.select(sb, fx.Int32(127)) + return mfma_scaled(av, xv, acc, sa, sb) + + +def gemv_mfmas(c, k, xl, xsl, red, ops, split=None, tiled=False, rows=16, nsteps=None): + """``gemv_loads``' operands against ``rows`` tokens' MXFP8 rows of the + window in LDS (``xl`` words, ``xsl`` codes, rows ``lds_row`` of the window) + -> ``red[wave][lane]`` (a 16 x 16 C: row 4 (lane // 16) + i, token lane % 16).""" + lane, wave = c["lane"], c["wave"] + nsteps = k // 128 if nsteps is None else nsteps + acc = fx.Vector.filled(4, 0.0, fx.Float32) + for i, (st, live) in enumerate(_steps(c, nsteps, split, tiled)): + acc = _step_mfma(c, nsteps, xl, xsl, st, ops[i], acc, live, rows) + if const_expr(split is not None and not tiled): + acc = fx.Vector.from_elements( + [(wave < split).select(acc[e], fx.Float32(0.0)) for e in range(4)], + fx.Float32, + ) + fx.ptr_store(acc, red + (wave * 64 + lane) * 4) + + +def row_sum(red, r, t, waves=None): + """Row r, token t of the waves' accumulators, summed in wave order.""" + tot = fx.Float32(0.0) + for w in waves if waves is not None else range(WAVES): + tot = tot + fx.ptr_load(red + ((w * 64 + 16 * (r // 4) + t) * 4 + r % 4)) + return tot + + +def span_items(width, n): + """A [n, width] span of 16 B items: rows a round of THREADS threads, rounds.""" + wq = width // 4 + rpr = THREADS // wq + assert rpr >= 1, width + return wq, rpr, (n + rpr - 1) // rpr + + +def rows_loads( + c, + words_base, + codes_base, + stride_w, + off_w, + words, + stride_g, + off_g, + groups, + t0, + n, + pred=None, +): + """Tokens t0 .. t0 + n's published MXFP8 row slices -- ``words`` i32 from + ``off_w`` and ``groups`` codes (an i32 each) from ``off_g``, rows + ``stride_w`` / ``stride_g`` i32 apart -- loaded (issued, not waited on), + 16 B a lane-item: a thread one column, rows a round apart. Plain loads: + their flags are seen, and nothing in the launch read these lines before. + ``pred`` false: masked off (zeros, no memory traffic).""" + tid = c["tid"] + nb = None if pred is None else pred.select(fx.Int64(0xFFFFFFFF), fx.Int64(0)) + out = [] + for base, stride, off, width in ( + (words_base, stride_w, off_w, words), + (codes_base, stride_g, off_g, groups), + ): + wq, rpr, rounds = span_items(width, n) + row0, col = tile_map(tid, wq) + a0 = col * 4 + off + t0 * stride + r = rsrc(base, nb) + regs = [] + for i in range(rounds): + row = row0 + rpr * i + if (i + 1) * rpr > n or rpr * wq < THREADS: + row = fx.min(row, n - 1) + regs.append( + fx.Vector( + bo.buffer_load(r, row * stride + a0, vec_width=4, dtype=T.i32) + ) + ) + out.append(regs) + return out + + +@traced +def _store_span(c, regs, width, dst, n): + tid = c["tid"] + wq, rpr, _ = span_items(width, n) + row0, col = tile_map(tid, wq) + lrow = lds_row(width) + base = dst + (row0 * lrow + col * 4) + for i in range_constexpr(len(regs)): + if const_expr((i + 1) * rpr <= n and rpr * wq == THREADS): + fx.ptr_store(regs[i], base + i * rpr * lrow) + else: + if (row0 + rpr * i < n) & (tid < rpr * wq): + fx.ptr_store(regs[i], base + i * rpr * lrow) + + +def rows_store(c, regs, words, groups, xl, xsl, n): + """``rows_loads``' registers into LDS rows 0 .. n (``lds_row`` strides).""" + _store_span(c, regs[0], words, xl, n) + _store_span(c, regs[1], groups, xsl, n) + gpu.barrier() diff --git a/vllm/models/deepseek_v41/amd/mono/attention/plan.py b/vllm/models/deepseek_v41/amd/mono/attention/plan.py new file mode 100644 index 000000000000..58865ac7d01e --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/attention/plan.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Geometry of the DeepSeek-V4.1 attention mega kernel: model widths, a TP +rank's share of them, task counts and the scratch (mailbox) layout.""" + +from dataclasses import dataclass + +HIDDEN = 5120 +Q_RANK = 1280 +KV_DIM = 512 +QKV_ROWS = Q_RANK + KV_DIM # 1792 +HEAD_DIM = 512 +NOPE = 448 +ROPE = 64 +HALF = ROPE // 2 +N_HEADS = 64 +O_GROUPS = 8 +O_RANK = 1024 +WINDOW = 128 +TOPK = 512 +EPS = 1e-20 + +ROWS = 16 # GEMV rows a task +TILE = 16 # tokens a GEMV pass (an MFMA's N) +HEAD_TILE = 16 # heads an attention MFMA takes +WOA_ROWS = 32 # a wo_a task's rows: one MX group of its output +MAX_TOKENS = 48 + +# KV record (V4 fp8_ds_mla): 576 B data rows (448 e4m3 | 64 bf16), then 8 B of +# scales a token (7 UE8M0 + pad) after the block's rows +DATA = 576 +SCALE_BYTES = 8 + +# attention: a split is 16 keys; a token has at most TOPK + WINDOW keys +BK = 16 +SPLITS = 64 # split slots a (token, head tile): 1024 keys >= TOPK + WINDOW +COLS = 32 # PV columns a group (one MX group of the output) +GROUPS = HEAD_DIM // COLS # 16 + +PAIR = 8 # a mailbox (value, tag) pair + + +def cdiv(a, b): + return -(-a // b) + + +def round_up(a, b): + return cdiv(a, b) * b + + +@dataclass(frozen=True) +class Dims: + tp: int + + def __post_init__(self): + assert self.tp in (2, 4), self.tp + + @property + def heads(self) -> int: + return N_HEADS // self.tp + + @property + def head_tiles(self) -> int: + return self.heads // HEAD_TILE + + @property + def groups(self) -> int: + return O_GROUPS // self.tp + + @property + def group_k(self) -> int: + """wo_a's K a group: its eight heads' outputs.""" + return self.heads // self.groups * HEAD_DIM # 4096 + + @property + def o_rows(self) -> int: + return self.groups * O_RANK + + @property + def q_rows(self) -> int: + return self.heads * HEAD_DIM + + @property + def wqb_tasks(self) -> int: + return self.q_rows // ROWS + + @property + def woa_tasks(self) -> int: + return self.o_rows // WOA_ROWS + + @property + def wob_tasks(self) -> int: + return HIDDEN // ROWS # 320 + + +WQKV_TASKS = QKV_ROWS // ROWS # 112 +WQKV_KPARTS = 2 # wqkv's K split: 224 (row tile, K half) parts +KEYS = TOPK + WINDOW # a token's attention keys at most + + +def token_tiles(s): + return [(t0, min(TILE, s - t0)) for t0 in range(0, s, TILE)] + + +def tile_rows(s): + return min(s, TILE) + + +def pair_layout(items, start=0, align=256): + out, off = {}, start + for name, pairs in items: + off = round_up(off, align) + out[name] = (off, pairs * PAIR) + off += pairs * PAIR + return out + + +def byte_layout(items, start=0, align=256): + out, off = {}, start + for name, nbytes in items: + off = round_up(off, align) + out[name] = (off, nbytes) + off += nbytes + return out + + +def front_scratch(s: int, d: Dims) -> dict: + """K1's scratch regions -> (byte offset, bytes): published data (plain + words behind a flag) and mailbox pairs.""" + return byte_layout( + [ + ("x8", s * HIDDEN), # e4m3 words (S > 16: quantized once, shared) + ("x8s", s * HIDDEN // 32 * 4), # an i32 code a group + ("xrdy", WQKV_KPARTS * s * PAIR), # a flag a (token, K half) + ("qkvp", WQKV_KPARTS * s * QKV_ROWS * PAIR), + ("qx8", s * Q_RANK), # e4m3 words + ("qx8s", s * Q_RANK // 32 * 4), # an i32 code a group + ("qrdy", s * PAIR), + ] + ) + + +def back_scratch(s: int, d: Dims, start: int = 0) -> dict: + """K2's scratch regions, laid out past ``start``: published data (plain) + and flag pairs.""" + ck = ( + 64 if d.head_tiles == 2 and s * cdiv(KEYS, 64) <= 256 else 128 + ) # back.chunk_keys + nchunk = cdiv(KEYS, ck) + return byte_layout( + [ + ("part", s * nchunk * d.heads * HEAD_DIM * 2), # bf16 partials + ("pstat", s * nchunk * d.heads * 8), # (max, sum) f32 + ("srdy", s * nchunk * PAIR), + ("xo", s * d.q_rows), # e4m3 words + ("xos", s * d.q_rows // 32 * 4), # an i32 code a group + ("crdy", s * d.heads // 8 * PAIR), + ("x8b", s * d.o_rows), + ("x8bs", s * d.o_rows // 32 * 4), + ("ardy", d.woa_tasks * PAIR), + ], + start=start, + ) + + +def front_end(s: int, d: Dims) -> int: + """Where K2's regions start: past K1's, which share the launch tag.""" + return round_up(max(o + n for o, n in front_scratch(s, d).values()), 256) + + +def scratch_bytes(s: int, d: Dims) -> int: + b = back_scratch(s, d, start=front_end(s, d)) + return round_up(max(o + n for o, n in b.values()), 4096) diff --git a/vllm/models/deepseek_v41/amd/mono/common/__init__.py b/vllm/models/deepseek_v41/amd/mono/common/__init__.py new file mode 100644 index 000000000000..397642c6b441 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/common/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/__init__.py +"""The mono kernels' shared mechanisms: ``ops`` (wave / math primitives, the TP +group's peer buffers), ``sync`` (the tagged mailbox), ``mx`` (FP8 / MX +arithmetic), ``plan`` (execution model, layouts, build keys) and +``peer_memory`` (symmetric peer buffers).""" diff --git a/vllm/models/deepseek_v41/amd/mono/common/mx.py b/vllm/models/deepseek_v41/amd/mono/common/mx.py new file mode 100644 index 000000000000..9c228abdbf27 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/common/mx.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/device/mx.py +"""FP8 / MX arithmetic the mono kernels share: the E4M3 range, E8M0 scale +codes, the byte of a scale in aiter's ``shuffle_scale`` layout, an FP4 weight +tile's operand loads and the scaled MFMA that takes them. + +An E8M0 code is a biased exponent; ``pow2(code)`` is the scale it stands for.""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import rocdl +from flydsl.expr.typing import T + +from vllm.models.deepseek_v41.amd.mono.common.ops import CM_NT + +FP8_MAX = 448.0 +# the f8f6f4 MFMA's operand formats (cbsz / blgp) +FP8, FP4 = 0, 4 +UNIT_SCALE = 127 # the E8M0 code of 1.0 + + +def clamp_fp8(v): + """``v`` clamped to the E4M3 range.""" + return fx.max(fx.min(v, fx.Float32(FP8_MAX)), fx.Float32(-FP8_MAX)) + + +def code_ceil(v): + """Biased exponent of the smallest power of two >= ``v`` (v >= 0, finite), + clamped to [1, 254]: the codes a scale may take.""" + bits = v.bitcast(fx.Int32) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).select( + fx.Int32(1), fx.Int32(0) + ) + return fx.max(fx.min(e, fx.Int32(254)), fx.Int32(1)) + + +def rcp_pow2(code): + """1 / ``pow2(code)``, exact (code in [1, 253]): x times it is x / scale + bit for bit, a multiply where the division is a dozen instructions.""" + return ((254 - code) << 23).bitcast(fx.Float32) + + +def scale_index(row, col, cols): + """Byte of e8m0 scale (row, col) in aiter's ``shuffle_scale`` layout: + ``view(rows/32, 2, 16, cols/8, 2, 4).permute(0, 3, 5, 2, 4, 1)``.""" + r32, a, b = row // 32, (row // 16) % 2, row % 16 + c8, d, e = col // 8, (col // 4) % 2, col % 4 + return ((((r32 * (cols // 8) + c8) * 4 + e) * 16 + b) * 2 + d) * 2 + a + + +def fp4_tile_load(r_w, base, rg, k, st, at): + """A lane's 16 B of an FP4 weight's (16, 16)-shuffled 16-row group ``rg`` + (K = ``k``) at 128-k step ``st``: dword ``base`` + the group's 2 k + 256 a + step + ``at`` (``lane * 4``: 32 fp4, K block lane / 16 of row lane % 16). + Nontemporal: a weight a step reads once.""" + return fx.Vector( + bo.buffer_load( + r_w, + base + rg * (2 * k) + st * 256 + at, + vec_width=4, + dtype=T.i32, + cache_modifier=CM_NT, + ) + ) + + +def mfma_scaled(a, b, c, sa, sb, a_fmt=FP8, b_fmt=FP8, sa_byte=0, sb_byte=0): + """One 16x16x128 scaled MFMA (gfx950 f8f6f4): A ``a`` (``a_fmt``) with E8M0 + ``sa`` (its byte ``sa_byte``), B ``b`` (``b_fmt``) with E8M0 ``sb`` (its + byte ``sb_byte``).""" + return fx.Vector( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + T.vec(4, T.f32), [a, b, c, a_fmt, b_fmt, sa_byte, sa, sb_byte, sb] + ) + ) diff --git a/vllm/models/deepseek_v41/amd/mono/common/ops.py b/vllm/models/deepseek_v41/amd/mono/common/ops.py new file mode 100644 index 000000000000..f9f7fb22ec92 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/common/ops.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/device/ops.py +# atom/mono/device/ranks.py +# atom/mono/device/stamps.py +"""Device-side primitives of the mono kernels: wave reductions and shuffles, the +approximate hardware math, bf16 / fp8 packing, buffer resources, the +execution-model constants and the kernel symbol naming. + +Ported from the FlyDSL GLM-5 layer kernel (ROCm/FlyDSL branch +``glm5-mla-moe-layer-monokernel`` @ 1e28fd8a, +``kernels/mla_moe_layer/shared_reuse_moe_kernel.py``). Only published flydsl APIs +and aiter's vendored ``buffer_ops`` / ``dpp_utils`` are used. + +Every function but ``kernel_symbol`` is traced inside a ``@flyc.kernel`` body. + +Raw boundaries kept, for want of an equivalent on the pinned FlyDSL: +``buffer_ops`` loads / stores (the kernels take raw Int64 addresses, no tensors +for ``make_buffer_tensor``); raw MFMA intrinsics (hand-laid operands: permlane +transposes, fp8 packing, MXFP4 scales); the LLVM intrinsics here (``expect``, +``s_memrealtime``, the approximate rsq / rcp / exp2) and the struct unpacking of +``permlane*_swap``. + +Also the TP group seen from inside a kernel (every rank's peer buffer, an +all-reduce's rank-ordered sum) and a timeline build's phase stamps.""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from aiter.ops.flydsl.kernels.dpp_utils import update_dpp_i32 +from flydsl._mlir.dialects import llvm +from flydsl.compiler.ast_rewriter import ASTRewriter +from flydsl.expr import const_expr, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import T, as_ir_value + +from vllm.models.deepseek_v41.amd.mono.common.plan import THREADS, WAVES + +# ---------------------------------------------------------------- primitives +# gfx95x cache policy bits (LLVM CPol): SC0 = 1, SC1 = 16. SC1 alone is device +# scope (past the per-XCD caches); SC0 | SC1 is system scope (peers over XGMI). +CM_DEV = 16 + + +CM_SYS = 17 + + +# NT = 2: nontemporal. A weight a layer reads once streams past the caches it +# would only evict the step's reused data from (the MoE at 48 rows: 143 -> 133 us) +CM_NT = 2 + + +def rsrc(addr, nbytes=None): + return bo.create_buffer_resource_from_addr(addr, num_records_bytes=nbytes) + + +def permlane_swap(off, x, y): + """v_permlane{32,16}_swap on registers (x, y): off 32 trades x's upper-half + lanes with y's lower-half lanes; off 16 trades x's odd 16-lane rows with y's + even rows (each within its half). -> (x', y').""" + swap = rocdl.permlane32_swap if off == 32 else rocdl.permlane16_swap + pr = swap( + llvm.StructType.get_literal([T.i32, T.i32]), + as_ir_value(fx.Int32(x)), + as_ir_value(fx.Int32(y)), + False, + False, + ) + return tuple(fx.Int32(llvm.extractvalue(T.i32, pr, [j])) for j in range(2)) + + +def lane_gather(v, src_lane): + """Lane ``src_lane``'s value of v, a source lane a lane (ds_bpermute).""" + return fx.Int32( + rocdl.ds_bpermute(T.i32, fx.Int32(src_lane) * 4, fx.Int32(v).ir_value()) + ) + + +def memrealtime(): + """s_memrealtime: the 100 MHz device clock (no FlyDSL wrapper).""" + return fx.Int64(llvm.call_intrinsic(T.i64, "llvm.amdgcn.s.memrealtime", [], [], [])) + + +def fresh(v): + """``v`` (a per-lane Int32) through an opaque move, here: the lane math built + on it is computed at its use, where LLVM would otherwise hoist it out of a task + loop or CSE it across stages into registers held for the whole kernel (spills, + and a kernel with scratch is refused: ``runtime.widths``).""" + return fx.Int32( + llvm.inline_asm( + T.i32, + [fx.Int32(v).ir_value()], + "v_mov_b32 $0, $1", + "=v,v", + has_side_effects=True, + ) + ) + + +def uniform(v): + return fx.Int32(rocdl.readfirstlane(T.i32, fx.Int32(v).ir_value())) + + +def _hw_f32(name, x): + return fx.Float32( + llvm.call_intrinsic(T.f32, name, [fx.Float32(x).ir_value()], [], []) + ) + + +def hw_rsq(x): + """v_rsq_f32 (~1 ulp, no libm range fixups).""" + return _hw_f32("llvm.amdgcn.rsq.f32", x) + + +def hw_rcp(x): + return _hw_f32("llvm.amdgcn.rcp.f32", x) + + +def hw_exp2(x): + return _hw_f32("llvm.amdgcn.exp2.f32", x) + + +def div_rn(x, d, r): + """``x / d`` rounded to nearest, bit-identical to the IEEE division, given + ``r = 1.0 / d`` from one true division per divisor (Markstein: a quotient + within 1 ulp plus its exact fma residual times RN(1/d) rounds correctly).""" + q = x * r + return fx.Float32(fmath.fma(fx.Float32(fmath.fma(-q, d, x)), r, q)) + + +def xshfl(v, off): + """Value of lane ``lane ^ off``: offsets 32 / 16 via v_permlane*_swap, in-row + offsets via DPP (VALU latency, not LDS). Powers of two only: any other + offset would take the lane ^ 1 path.""" + assert off in (1, 2, 4, 8, 16, 32), off + if off >= 16: + return v.shuffle_xor(off, 64) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + if off == 8: + y = fx.Int32(update_dpp_i32(x, x, 0x118, 0xF, 0xC, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x108, 0xF, 0x3, False)) + elif off == 4: + y = fx.Int32(update_dpp_i32(x, x, 0x114, 0xF, 0xA, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x104, 0xF, 0x5, False)) + elif off == 2: + y = fx.Int32(update_dpp_i32(x, x, 0x4E, 0xF, 0xF, False)) + else: + y = fx.Int32(update_dpp_i32(x, x, 0xB1, 0xF, 0xF, False)) + return y.bitcast(fx.Float32) if is_f else y + + +def xred(v, off, op): + """op(v, lane ^ off) for a symmetric op; 32 / 16 take both halves of one swap.""" + if off < 16: + return op(v, xshfl(v, off)) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + a, b = permlane_swap(off, x, x) + if is_f: + return op(a.bitcast(fx.Float32), b.bitcast(fx.Float32)) + return op(type(v)(a), type(v)(b)) + + +def _add(a, b): + return a + b + + +def butterfly(v, offsets, op=_add): + """Butterfly reduction over the given xor offsets, in the given order. + + The order is the numerical contract: the aiter per-head norm reduces 32 + lanes with offsets 16, 8, 4, 2, 1, and matching it keeps rstd bit-exact.""" + for off in offsets: + v = xred(v, off, op) + return v + + +def traced(fn): + """Apply flydsl's AST rewriting to a helper, so its ``while`` / ``if`` over + traced values lower to scf ops. Kernels rewrite only their own body.""" + return ASTRewriter.transform(fn) + + +LOAD_BATCH = 4 # rounds of a strided loop whose loads go out together + + +@traced +def batched_rounds(tid, n, load, use, batch=LOAD_BATCH): + """A CTA's strided loop over elements v = tid + THREADS i < n (n a Python + int): ``load(v)`` of ``batch`` rounds issued before any ``use(v, loaded)``. + A round whose ``use`` stores to global memory keeps the next round's loads + behind it (they may alias), so the rounds' loads would go out in series. A + round past n loads element n - 1 (in range) and uses nothing.""" + rounds = -(-n // THREADS) + for b0 in range_constexpr(0, rounds, batch): + idx = range(b0, min(b0 + batch, rounds)) + got = [load(fx.min(tid + THREADS * i, n - 1)) for i in idx] + for i, g in zip(idx, got): + v = tid + THREADS * i + if v < n: + use(v, g) + + +def mfma_bf16(a, b, c): + return fx.Vector(rocdl.mfma_f32_16x16x32_bf16(T.vec(4, T.f32), [a, b, c])) + + +def row_sum(red, r, t, waves=None): + """Row r, token t of the waves' 16x16 MFMA accumulators (wave w's at ``red`` + + (w * 64 + lane) * 4), summed in wave order from 0 (``waves``: which, in + order; all of them by default).""" + tot = fx.Float32(0.0) + for w in waves if waves is not None else range(WAVES): + tot = tot + fx.ptr_load(red + ((w * 64 + 16 * (r // 4) + t) * 4 + r % 4)) + return tot + + +def ballot(pred): + """The wave's lanes where ``pred`` holds, a bit a lane.""" + return fx.Int64(rocdl.ballot(T.i64, as_ir_value(pred))) + + +def popcount(x): + return fx.Int32(fx.ctpop(x)) + + +def lanes_below(bal, lane): + """How many of the lanes in ballot ``bal`` are below ``lane``.""" + return popcount(bal & ((fx.Int64(1) << fx.Int64(lane)) - 1)) + + +def load_ptr64(table, i): + """Entry ``i`` of an i64 table in memory, wave-uniform (each half through + ``uniform``): a peer's buffer base, a kernel argument block's pointer.""" + pw = fx.Vector(bo.buffer_load(rsrc(table), i * 2, vec_width=2, dtype=T.i32)) + return (fx.Int64(uniform(pw[1])) << 32) | fx.Int64(fx.Uint32(uniform(pw[0]))) + + +def bf16_round(a): + return fx.Float32(fx.Float32(a).to(fx.BFloat16)) + + +def bf16_pair(a, b): + """Two f32 -> one f32-typed word holding (bf16(a), bf16(b)).""" + return ( + fx.Vector.from_elements([a, b], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Float32)[0] + ) + + +def bf_lo(w): + """The f32 of a packed bf16 pair word's low half.""" + return (w << 16).bitcast(fx.Float32) + + +def bf_hi(w): + """The f32 of a packed bf16 pair word's high half.""" + return (w & fx.Int32(-65536)).bitcast(fx.Float32) + + +def bf2_f32(w): + """Packed bf16 pair word -> (f32 low, f32 high).""" + return bf_lo(w), bf_hi(w) + + +def fp8_pack4(a, b, c, d): + """Four f32 (already in range) -> one dword of four E4M3 bytes, RNE.""" + lo = fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, a, b, fx.Int32(0), False)) + return fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, c, d, lo, True)) + + +# ---------------------------------------------------------------- TP ranks +def peer_bases(peers, tp): + """Every rank's peer buffer base from the ``peers`` table (an i64 a rank), + loaded once: a put's address must not wait on a load of its own.""" + return [load_ptr64(peers, p) for p in range(tp)] + + +def sum_partials(poll, own, pair_of, tp): + """The ``tp`` ranks' bf16 pair partials at ``pair_of(src)`` of the mailbox + region ``own``, summed in fp32 in rank order 0 .. tp - 1: one poll batch -> + (lo sum, hi sum). Every rank sums in the same order, so they agree.""" + ws = poll([(own, pair_of(src), 1) for src in range(tp)]) + acc0, acc1 = bf2_f32(ws[0][0]) + for src in range_constexpr(1, tp): + lo, hi = bf2_f32(ws[src][0]) + acc0 = acc0 + lo + acc1 = acc1 + hi + return acc0, acc1 + + +# ---------------------------------------------------------------- timeline stamps +def stamp(on, tls, tid, k): + # compile-time gate outside, traced condition inside: they cannot be one `and` + if const_expr(on): # noqa: SIM102 + if tid == 0: + fx.ptr_store(memrealtime(), tls + k) diff --git a/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py b/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py new file mode 100644 index 000000000000..58c95a567e07 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/runtime/consensus.py +# atom/mono/runtime/peer_memory.py +"""Symmetric peer memory for the in-kernel TP all-reduces, and the agreement every +rank must reach before anything collective (a rank that turned the mono layer +off alone would leave the others spinning on the GPU).""" + +import torch +import torch.distributed as dist +from aiter.ops.flydsl.quick_allreduce_int4_ipc import UncachedIpcHeap + + +# ---------------------------------------------------------------- TP agreement +class MonoUnsupported(Exception): + """The loaded model or runtime configuration is outside what mono serves.""" + + +def tp_agree(ok: bool, group) -> bool: + """True when every rank of ``group`` passed ``ok`` true. Every rank must call + it at the same point; ``group`` is a CPU (gloo) process group, or None for a + single rank.""" + if group is None: + return ok + flag = torch.tensor([int(ok)], dtype=torch.int32) + dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=group) + return bool(flag.item()) + + +# ---------------------------------------------------------------- peer buffers +class _DeviceBytes: + """``__cuda_array_interface__`` of ``nbytes`` bytes at a raw device address, for a + non-owning torch view.""" + + def __init__(self, ptr: int, nbytes: int): + self.__cuda_array_interface__ = { + "shape": (nbytes,), + "typestr": "|u1", + "data": (ptr, False), + "version": 3, + } + + +class PeerBuffer: + """One uncached buffer per rank, every rank holding every peer's address + (``addresses``, int64, indexed by TP rank). What lives where, and why no rank + overwrites a peer's unread data, is the model's layout's to say. + + The constructor is collective over ``group`` (the handle exchange, then an + agreement that every rank mapped every peer): every rank must reach it, so a + runner reaches it only after ``tp_agree``. A failed mapping on any rank raises + ``MonoUnsupported`` on all of them. + """ + + def __init__( + self, + nbytes: int, + group, + rank: int, + npes: int, + device: torch.device, + ): + self.local = UncachedIpcHeap.alloc_uncached(nbytes) + # a view, not an owner: close() frees the memory under it + self.bytes = torch.as_tensor(_DeviceBytes(self.local, nbytes), device=device) + self.rank = rank + self._opened: list[int] = [] + addresses = [self.local] + if npes > 1: + handles = UncachedIpcHeap.gather_object_list_via_broadcast( + group, UncachedIpcHeap.get_mem_handle_bytes(self.local) + ) + addresses, failure = [], None + try: + for peer, handle in enumerate(handles): + if peer == rank: + addresses.append(self.local) + else: + base = UncachedIpcHeap.open_mem_handle(handle) + self._opened.append(base) + addresses.append(base) + except RuntimeError as err: # every rank must still reach the agreement + failure = err + if not tp_agree(failure is None, group): + self.close() + raise MonoUnsupported( + "peer memory handshake failed on " + f"{'this' if failure else 'another'}" + " TP rank" + ) from failure + self.addresses = torch.tensor(addresses, dtype=torch.int64, device=device) + + def close(self) -> None: + """Close the peer mappings and free the local buffer; idempotent. Only once no + kernel of any rank can still touch it (the runner is being dropped).""" + opened, self._opened = self._opened, [] + for base in opened: + UncachedIpcHeap.close_mem_handle(base) + if self.local: + self.bytes = None + UncachedIpcHeap.free_device_mem(self.local) + self.local = 0 diff --git a/vllm/models/deepseek_v41/amd/mono/common/plan.py b/vllm/models/deepseek_v41/amd/mono/common/plan.py new file mode 100644 index 000000000000..2b05d7a073e9 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/common/plan.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/plan/execution.py +# atom/mono/plan/layout.py +# atom/mono/plan/arith.py +# atom/mono/plan/shard.py +# atom/mono/plan/build_key.py +"""Plain-Python planning of the mono kernels: the execution model (one CTA a CU, +all resident), mailbox layouts, arithmetic that runs on ints and traced values +alike, TP shards, and build keys (every build parameter, and this package's +sources, in FlyDSL's JIT cache key).""" + +import enum +import functools +import hashlib +from dataclasses import dataclass, fields +from pathlib import Path + +# ---------------------------------------------------------------- execution model +BLOCKS = 256 +THREADS = 512 +WAVES = THREADS // 64 + + +def first_task(bid, base): + """CTA ``bid``'s first task of a stage placed from CTA ``base`` (any count of + CTAs past 0: placements that add task counts run past the grid), then every + BLOCKS-th.""" + return (bid - base % BLOCKS + BLOCKS) % BLOCKS + + +# ---------------------------------------------------------------- mailbox layouts +PAIR_BYTES = 8 # a (value, tag) pair + + +def pair_layout(items, align: int = PAIR_BYTES, start: int = 0) -> dict: + """``items`` (name, pairs) laid out in order: name -> (byte offset, bytes), + each region from an ``align``-byte boundary at or past ``start``.""" + out, off = {}, start + for name, pairs in items: + off = (off + align - 1) // align * align + out[name] = (off, pairs * PAIR_BYTES) + off += pairs * PAIR_BYTES + return out + + +# ---------------------------------------------------------------- arithmetic +def sel(pred, a, b): + """A if pred else b: a conditional on Python values, ``select`` on traced + ones (rewrapped in the traced operand's type).""" + if isinstance(pred, bool): + return a if pred else b + out = pred.select(a, b) + traced = [x for x in (a, b) if not isinstance(x, int)] + return type(traced[0])(out) if traced else out + + +def imin(a, b): + return sel(a < b, a, b) + + +def i32(x): + """``x``, asserted inside the device's Int32 range when it is a Python int.""" + if isinstance(x, int): + assert -(2**31) <= x < 2**31, x + return x + + +def cdiv(a, b): + """ceil(a / b) for a >= 0, b > 0.""" + return i32((a + b - 1) // b) + + +# ---------------------------------------------------------------- TP shards +class ShardError(ValueError): + """A width the TP size does not divide, or a TP size the model refuses.""" + + +@dataclass(frozen=True) +class Shard: + tp: int + + def __post_init__(self): + if self.tp < 1: + raise ShardError(f"TP {self.tp}") + + def split(self, full: int, what: str) -> int: + """``full`` / tp, refusing a remainder: a rank's share of ``what``.""" + if full % self.tp: + raise ShardError(f"{what} {full} is not divisible by TP {self.tp}") + return full // self.tp + + +def padded(n: int, align: int) -> int: + return cdiv(n, align) * align + + +# ---------------------------------------------------------------- build keys +_SCALARS = (int, float, bool, str, enum.Enum) +_ROOT = Path(__file__).resolve().parents[1] # the mono package + + +@functools.cache +def source_digest(*paths: str) -> str: + """sha256 of every ``.py`` file under ``paths`` (files or directories, + relative to the mono package), each file's path and contents: one model's + kernel sources, so an edit to another model's rebuilds none of these.""" + h = hashlib.sha256() + for rel in paths: + root = _ROOT / rel + files = [root] if root.is_file() else sorted(root.rglob("*.py")) + if not files: + raise FileNotFoundError(f"no kernel sources at {root}") + for f in files: + h.update(str(f.relative_to(_ROOT)).encode()) + h.update(f.read_bytes()) + return h.hexdigest()[:16] + + +def key_tuple(key, sources: str) -> tuple: + """``key``'s (field name, value) pairs, every value a scalar, then + ("sources", ``sources``): the model's ``source_digest``.""" + out = [] + for f in fields(key): + value = getattr(key, f.name) + if not isinstance(value, _SCALARS): + raise TypeError( + f"{type(key).__name__}.{f.name} = {value!r}: a build parameter must be" + " a scalar to reach the JIT cache key" + ) + out.append((f.name, value)) + out.append(("sources", sources)) + return tuple(out) diff --git a/vllm/models/deepseek_v41/amd/mono/common/sync.py b/vllm/models/deepseek_v41/amd/mono/common/sync.py new file mode 100644 index 000000000000..0220ce12c32a --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/common/sync.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/device/sync.py +"""The tagged mailbox: hand-offs between the CTAs of one GPU (device scope) and +between GPUs (system scope), every address a named region (``Addr``).""" + +from dataclasses import dataclass +from enum import Enum + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T + +from vllm.models.deepseek_v41.amd.mono.common.ops import ( + CM_DEV, + CM_SYS, + bf16_pair, + rsrc, + traced, +) + +POLL_MAX = 12 # mailbox specs polled per batch (bounds live registers) + + +@traced +def _spin(load_all, pending): + """Re-issue ``load_all`` until ``pending`` clears; a side-effecting op in the + loop keeps the loads from being hoisted.""" + v = load_all() + while pending(v): + rocdl.s_sleep(1) + v = load_all() + return v + + +class Space(Enum): + SCRATCH = "scratch" # this GPU's memory: device-scope hand-off between CTAs + PEER = "peer" # a peer-symmetric region: system-scope hand-off between GPUs + + +@dataclass(frozen=True) +class Addr: + """A mailbox region's address: the region's name and space, and the traced + address. The space sets the hand-off's scope: device for SCRATCH, system for + PEER.""" + + region: str + space: Space + value: object + + +def sreg(base, offset, region): + """Scratch region ``region`` at byte ``offset`` of ``base``.""" + return Addr(region, Space.SCRATCH, base + fx.Int64(offset)) + + +def preg(base, offset, region): + """Peer-symmetric region ``region`` at byte ``offset`` of a rank's buffer + ``base``.""" + return Addr(region, Space.PEER, base + fx.Int64(offset)) + + +# the store / load cache policy of a hand-off in each space +_SCOPE_CM = {Space.SCRATCH: CM_DEV, Space.PEER: CM_SYS} + + +class Mailbox: + """Tagged-pair hand-off between CTAs (device scope) and GPUs (system scope). + + Every 32-bit value is stored next to this launch's tag; a consumer + polls the payload until all tags match, so a hand-off costs one round trip + (no store drain, no separate flag). The mono layer's launches tag their + hand-offs with their epoch (``layer.py``), so a pair from an earlier step + never carries a live tag. Every address is an ``Addr``: its space picks the + scope. + """ + + def __init__(self, tag): + self.tag = tag + + def put(self, a, i, v): + bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) + bo.buffer_store( + fx.Vector.from_elements([bits, self.tag], fx.Int32), + rsrc(a.value), + i * 2, + cache_modifier=_SCOPE_CM[a.space], + ) + + def put_bf(self, a, i, vs): + """Elements i .. i+len(vs) (2 or 4, i aligned) as packed bf16 pairs, one + store.""" + words = [] + for j in range_constexpr(len(vs) // 2): + words += [bf16_pair(vs[2 * j], vs[2 * j + 1]).bitcast(fx.Int32), self.tag] + bo.buffer_store( + fx.Vector.from_elements(words, fx.Int32), + rsrc(a.value), + i, + cache_modifier=_SCOPE_CM[a.space], + ) + + def put_words(self, a, i, words): + """Raw 32-bit words (1 or 2) at pairs i.. as (word, tag). Integer payloads + (fp8 words) must come this way: through an f32 value a NaN bit pattern is + not kept.""" + vals = [] + for w in words: + vals += [fx.Int32(w), self.tag] + bo.buffer_store( + fx.Vector.from_elements(vals, fx.Int32), + rsrc(a.value), + i * 2, + cache_modifier=_SCOPE_CM[a.space], + ) + + def poll(self, specs, batch=POLL_MAX): + """Batched poll of pairs ``specs`` = [(Addr, pair index, npairs in {1, 2})], + all of one space. + + All pairs are re-loaded together while any tag is stale, so a batch costs + one round trip after its last producer lands. Returns the value words.""" + if const_expr(len(specs) == 0): + return [] + if const_expr(len(specs) > batch): + return self.poll(specs[:batch], batch) + self.poll(specs[batch:], batch) + spaces = {a.space for a, _, _ in specs} + assert len(spaces) == 1, f"one poll batch spans {spaces}" + cm = _SCOPE_CM[spaces.pop()] + bases = [a.value for a, _, _ in specs] + tag = self.tag + + def load_all(): + words = [] + for b, (_, i, n) in zip(bases, specs): + w = fx.Vector( + bo.buffer_load( + rsrc(b), + fx.Int32(i) * 2, + vec_width=2 * n, + dtype=T.i32, + cache_modifier=cm, + ) + ) + words += [w[e] for e in range(2 * n)] + return fx.Vector.from_elements(words, fx.Int32) + + nw = sum(2 * n for _, _, n in specs) + + def pending(v): + bad = v[1] != tag + for e in range_constexpr(3, nw, 2): + bad = bad | (v[e] != tag) + return bad + + v = _spin(load_all, pending) + outs, e = [], 0 + for _, _, n in specs: + outs.append([v[e + 2 * q] for q in range(n)]) + e += 2 * n + return outs + + +@traced +def publish(put, a, i, v, who): + """This CTA's plain stores made readable by every CTA, then (thread ``who``) + the mailbox pair (``a``, i) := v. The stores are drained first: a flag may + not overtake the data it announces, and a workgroup-scope release fence does + not wait for them (LLVM's AMDGPU memory model omits vmcnt outside tgsplit + mode: the workgroup's waves share one CU, the readers here do not).""" + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if who: + put(a, i, v) diff --git a/vllm/models/deepseek_v41/amd/mono/layer.py b/vllm/models/deepseek_v41/amd/mono/layer.py new file mode 100644 index 000000000000..c1be9e3428ae --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/layer.py @@ -0,0 +1,940 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The mono decode layer's two launches (README): K1 = the attention seam +(``stages.seam``: slice, gate, the norm into vLLM's MXFP8) + the attention front +(``attention.front``); K2 = the attention back (``attention.back``, wo_b pushed +to every rank) + the FFN seam (the TP sum folded in) + the MoE (``stages.moe``, +its all-reduce in-kernel too). A layer whose attention vLLM runs takes the FFN +launch instead: K2 with the attention back replaced by a push of vLLM's +unreduced wo_b output (``seam.stage_push``). + +Tags: K1's hand-offs carry 2 e + 1, K2's 2 e + 2 (e: the launch pair's epoch, +moved on by K2's CTA 0 at its end), so K2's seam may reuse K1's seam regions. +Peer regions alternate by e's parity; every rank runs the same launches, so +the epochs agree across ranks. +""" + +from dataclasses import dataclass + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import Int32, Int64, T + +from .attention import back as mb_back +from .attention import front as mb_front +from .attention.device import BLOCKS, CM_DEV, THREADS, gstore, kernel_symbol, rsrc +from .attention.plan import HIDDEN, Dims, back_scratch, front_scratch +from .common.mx import clamp_fp8 +from .common.ops import butterfly, fp8_pack4, peer_bases, traced +from .common.plan import WAVES, first_task, key_tuple +from .common.sync import Mailbox, preg, publish, sreg +from .sources import SOURCES +from .stages import moe, seam +from .stages.dims import Dims as StageDims +from .stages.moe_shape import SORT_NETS, MoeBuild, route_shape + +_STREAM = fx.Stream(None) +MAX_TOKENS = 48 +X_WORDS = HIDDEN // 4 +X_GROUPS = HIDDEN // 32 +SLICES = seam.SLICES # 160 +# K1's MXFP8 x and the MoE's: plain words in both, the same size, so the two +# launches share the MoE's (placed later); any other name is one region, so a +# word is never a tag to one launch and data to another +SHARED_REGIONS = ("x8", "x8s") +COUNTER_WORDS = 2 * 256 # the MoE's ug queue and down counts, a word a slot +COUNTERS = ("ugq", "dq") +# the epoch buffer, in words: the epoch, a mark a CTA (back.EPOCH_MARKS), then on +# lines of their own the MoE's counters +EPOCH_COUNTERS = -(-(mb_back.EPOCH_MARKS + BLOCKS) // 64) * 64 +EPOCH_WORDS = EPOCH_COUNTERS + COUNTER_WORDS + + +@dataclass(frozen=True) +class MonoBuild: + tokens: int + tp: int + ratio: int # the layer's compress ratio (its attention's key sources) + timeline: bool = False + + +def _moe_key(key: MonoBuild) -> MoeBuild: + return MoeBuild(tokens=key.tokens, tp=key.tp) + + +def _front_key(key: MonoBuild): + return mb_front.FrontBuild(key.tokens, key.tp, key.ratio, key.timeline) + + +def _back_key(key: MonoBuild): + return mb_back.BackBuild(key.tokens, key.tp, key.ratio, key.timeline) + + +def scratch_layout(s: int, tp: int) -> dict: + """Every region of step width s's launches -> (byte offset, bytes), disjoint: + K1's front and seam, K2's back, the MoE's regions, the normed rows and their + flags. Each width has a scratch of its own (``DSV41MonoLayer.scratch``): + under another width's layout a mailbox pair's tag word could be plain data, + which a poll may take for a current tag. The MoE's counters, whose slots + outlive a step width, are in the epoch buffer.""" + d = Dims(tp) + out: dict = {} + off = 0 + + def place(regions): + nonlocal off + base = off + for name, (o, n) in regions.items(): + assert name not in out or name in SHARED_REGIONS, name + out[name] = (base + o, n) + off = max(off, base + max(o + n for o, n in regions.values())) + off = -(-off // 256) * 256 + + place(front_scratch(s, d)) + place(seam.scratch_layout(s)) + place(back_scratch(s, d, start=0)) + regions = moe.scratch_layout(_moe_key(MonoBuild(s, tp, 1))).items() + mr = {n: v for n, v in regions if n not in COUNTERS} + lo = min(o for o, _ in mr.values()) + place({n: (o - lo, n_) for n, (o, n_) in mr.items()}) + place({"normed": (0, s * HIDDEN * 2), "xrdy_moe": (s * HIDDEN * 2 + 256, s * 8)}) + return out + + +def scratch_bytes(s: int = MAX_TOKENS, tp: int = 2) -> int: + return max(o + n for o, n in scratch_layout(s, tp).values()) + + +def peer_half_bytes(tp: int) -> int: + """One parity's peer regions: the attention's partials, then the MoE's.""" + return seam.attn_peer_bytes(MAX_TOKENS, tp) + moe.peer_bytes(MAX_TOKENS, tp) + + +def _epoch(epoch): + return fx.Int32(bo.buffer_load(rsrc(epoch), 0, vec_width=1, dtype=T.i32)) + + +# ---------------------------------------------------------------- K1 + + +@traced +def store_x8(c, t, col, ys, live): + """The attention seam's norm output (``seam.stage_norm``'s 8 columns at + ``col``) -> vLLM's MXFP8 of each 32 (the attention's input quant) -> the + front's X8 / X8S, plain at device scope.""" + amax = abs(ys[0]) + for y in ys[1:]: + amax = fx.max(amax, abs(y)) + code = moe.vllm_mx_code(butterfly(amax, (1, 2), fx.max)) + mul = ((254 - code) << 23).bitcast(fx.Float32) + w0 = fp8_pack4(*[clamp_fp8(ys[i] * mul) for i in range(4)]) + w1 = fp8_pack4(*[clamp_fp8(ys[4 + i] * mul) for i in range(4)]) + if live: + bo.buffer_store( + fx.Vector.from_elements([w0, w1], fx.Int32), + rsrc(c["x8"]), + t * X_WORDS + col // 4, + cache_modifier=CM_DEV, + ) + if c["tid"] % 4 == 0: + bo.buffer_store( + code, rsrc(c["x8s"]), t * X_GROUPS + col // 32, cache_modifier=CM_DEV + ) + + +@traced +def publish_x8(c, t): + """Token t's two XRDY flags (one a wqkv K half) once its X8 / X8S landed.""" + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if c["tid"] < 2: + c["mb"].put(c["xrdy"], 2 * t + c["tid"], fx.Int32(1)) + gpu.barrier() + + +def build_mono_k1(key: MonoBuild): + s, tp = key.tokens, key.tp + assert 1 <= s <= MAX_TOKENS + layout = scratch_layout(s, tp) + fkey = _front_key(key) + FrontLds = mb_front.front_lds(s, Dims(tp)) + NORM0 = SLICES # the norms on the CTAs past the slices: the wqkv waits on them + GATE0 = SLICES + s + assert GATE0 + s <= BLOCKS + + @fx.struct + class SeamLds: + red: fx.Array[fx.Float32, WAVES * 2 * 64 * 4, 16] + rl: fx.Array[fx.Float32, s * seam.KT, 16] + fl: fx.Array[fx.Float32, seam.MIX * seam.KT, 16] + + @fx.union + class K1Lds: + seam: SeamLds + front: FrontLds # type: ignore[valid-type] + + name = kernel_symbol("dsv41_mono_k1", s=s, tp=tp, r=key.ratio, tl=key.timeline) + keyed = key_tuple(key, SOURCES) + + @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) + def k1( + res_in: Int64, + pend: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + attn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + wqkv: Int64, + wqkv_s: Int64, + qn_w: Int64, + kvn_w: Int64, + wqb: Int64, + wqb_s: Int64, + cos_sin: Int64, + pos: Int64, + slot: Int64, + swa: Int64, + swa_stride: Int64, + swa_block: Int32, + q: Int64, + swa_idx: Int64, + swa_lens: Int64, + comp_block: Int32, + comp_bt: Int64, + bt_stride: Int32, + topk: Int64, + t2r: Int64, + kt: Int64, + klen: Int64, + scratch: Int64, + epoch: Int64, + tl: Int64, + ): + _ = keyed + bid = fx.block_idx.x + lds = fx.SharedAllocator().allocate(K1Lds) + sl = lds.seam.peek() + ep = _epoch(epoch) + tag = (ep << 1) + 1 + fargs = { + "x": 0, + "x_stride": 0, + "wqkv": wqkv, + "wqkv_s": wqkv_s, + "qn_w": qn_w, + "kvn_w": kvn_w, + "wqb": wqb, + "wqb_s": wqb_s, + "cos_sin": cos_sin, + "pos": pos, + "slot": slot, + "swa": swa, + "swa_stride": swa_stride, + "swa_block": swa_block, + "q": q, + "swa_idx": swa_idx, + "swa_lens": swa_lens, + "comp_block": comp_block, + "comp_bt": comp_bt, + "bt_stride": bt_stride, + "topk": topk, + "t2r": t2r, + "kt": kt, + "klen": klen, + } + flayout = {n: layout[n] for n in front_scratch(s, Dims(tp))} + c = mb_front.front_context(fkey, lds.front.peek(), fargs, scratch, tag, flayout) + amb = Mailbox(tag) + tid = fx.thread_idx.x + ca = { + "S": s, + "tid": tid, + "bid": bid, + "lane": tid % 64, + "wave": tid // 64, + "rl": sl.rl.ptr, + "fl": sl.fl.ptr, + "red": sl.red.ptr, + "put": amb.put, + "put_bf": amb.put_bf, + "put_words": amb.put_words, + "poll": amb.poll, + "args": { + "res_in": res_in, + "pend": pend, + "post_in": post_in, + "comb_in": comb_in, + "pre_in": pre_in, + "hc_fn": hc_fn, + "hc_scale": hc_scale, + "hc_base": hc_base, + "norm_w": attn_w, + "res_out": res_out, + "post_out": post_out, + "comb_out": comb_out, + "pre_out": pre_out, + }, + "d": StageDims(tp), + } + for region in ("lin", "pmix"): + ca[region] = sreg(scratch, layout[region][0], region) + if const_expr(key.timeline): + c["tl"], c["tl_points"] = tl, mb_front.FRONT_POINTS + + def seam_slice(): + for task in range(first_task(bid, 0), SLICES, BLOCKS): + seam.stage_slice(ca, task) + gpu.barrier() + + def seam_rest(): + for t in range(first_task(bid, NORM0), s, BLOCKS): + seam.stage_norm(ca, t, lambda *v: store_x8(c, *v)) + publish_x8(c, t) + for t in range(first_task(bid, GATE0), s, BLOCKS): + seam.stage_gate(ca, t) + gpu.barrier() + + mb_front.run_front( + c, fkey, bid, x8_given=True, before_wqkv=seam_slice, after_wqkv=seam_rest + ) + + @flyc.jit + def launch( + res_in: Int64, + pend: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + attn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + wqkv: Int64, + wqkv_s: Int64, + qn_w: Int64, + kvn_w: Int64, + wqb: Int64, + wqb_s: Int64, + cos_sin: Int64, + pos: Int64, + slot: Int64, + swa: Int64, + swa_stride: Int64, + swa_block: Int32, + q: Int64, + swa_idx: Int64, + swa_lens: Int64, + comp_block: Int32, + comp_bt: Int64, + bt_stride: Int32, + topk: Int64, + t2r: Int64, + kt: Int64, + klen: Int64, + scratch: Int64, + epoch: Int64, + tl: Int64, + stream: fx.Stream = _STREAM, + ): + _ = keyed + k1( + res_in, + pend, + post_in, + comb_in, + pre_in, + hc_fn, + hc_scale, + hc_base, + attn_w, + res_out, + post_out, + comb_out, + pre_out, + wqkv, + wqkv_s, + qn_w, + kvn_w, + wqb, + wqb_s, + cos_sin, + pos, + slot, + swa, + swa_stride, + swa_block, + q, + swa_idx, + swa_lens, + comp_block, + comp_bt, + bt_stride, + topk, + t2r, + kt, + klen, + scratch, + epoch, + tl, + ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) + + return launch + + +# ---------------------------------------------------------------- K2 + + +@traced +def store_normed(normed, t, col, ys, live): + """The FFN seam's norm output -> the MoE input row (bf16, device scope: + this launch reads it).""" + if live: + for j in range_constexpr(8): + bo.buffer_store( + ys[j].to(fx.BFloat16), + rsrc(normed), + t * HIDDEN + col + j, + cache_modifier=CM_DEV, + ) + + +def _seam_lds(s): + @fx.struct + class SeamLds: + red: fx.Array[fx.Float32, WAVES * 2 * 64 * 4, 16] + rl: fx.Array[fx.Float32, s * seam.KT, 16] + fl: fx.Array[fx.Float32, seam.MIX * seam.KT, 16] + pl: fx.Array[fx.Float32, s * seam.COLS, 16] + + return SeamLds + + +@traced +def run_ffn(key: MonoBuild, lds, a: dict, amb, bases, own, rank, ep, part=None): + """K2's FFN half on this CTA: the FFN seam (the attention's TP partials -- + ``part`` pushed to every rank's ATTN region first -- summed and folded into + the residual), the MoE and its all-reduce into ``a["out"]``, then the + launch pair's epoch moved on.""" + s, tp = key.tokens, key.tp + layout = scratch_layout(s, tp) + mkey = _moe_key(key) + rs = route_shape(mkey) + GATE0 = SLICES + NORM0 = SLICES + s + bid = fx.block_idx.x + tid = fx.thread_idx.x + sl = lds.seam.peek() + scratch = a["scratch"] + seam_args = ( + "res_in", + "post_in", + "comb_in", + "pre_in", + "hc_fn", + "hc_scale", + "hc_base", + "res_out", + "post_out", + "comb_out", + "pre_out", + ) + ca = { + "S": s, + "tid": tid, + "bid": bid, + "lane": tid % 64, + "wave": tid // 64, + "rl": sl.rl.ptr, + "fl": sl.fl.ptr, + "red": sl.red.ptr, + "pend_lds": sl.pl.ptr, + "put": amb.put, + "put_bf": amb.put_bf, + "put_words": amb.put_words, + "poll": amb.poll, + "peer_addr": lambda p: bases[p], + "rank": rank, + "sym": own, + "args": {n: a[n] for n in seam_args} | {"norm_w": a["ffn_w"]}, + "d": StageDims(tp), + } + normed = scratch + fx.Int64(layout["normed"][0]) + for region in ("lin", "pmix"): + ca[region] = sreg(scratch, layout[region][0], region) + mlayout = {n: layout[n] for n in moe.scratch_layout(mkey) if n not in COUNTERS} + mlayout["xrdy"] = layout["xrdy_moe"] + margs = moe.moe_args( + normed, + a["gate_w"], + a["bias"], + a["w13"], + a["w13_s"], + a["w2"], + a["w2_s"], + a["sgu"], + a["sgu_s"], + a["sw2"], + a["sw2_s"], + a["out"], + ) + cb = moe.moe_context( + s, rs, lds.moe.peek(), amb, bases, mlayout, scratch, margs, rank, own + ) + cb["x_cm"], cb["x_ready"] = CM_DEV, True + cb["tag"] = ep & 255 + for i, region in enumerate(COUNTERS): + cb[region] = sreg(a["epoch"], 4 * (EPOCH_COUNTERS + i * moe.TAGS), region) + for task in range(first_task(bid, 0), SLICES, BLOCKS): + if const_expr(part is not None): + seam.stage_push(ca, task, part) + seam.stage_reduce(ca, task) + seam.stage_slice(ca, task) + for t in range(first_task(bid, GATE0), s, BLOCKS): + seam.stage_gate(ca, t) + for t in range(first_task(bid, NORM0), s, BLOCKS): + seam.stage_norm(ca, t, lambda *v: store_normed(normed, *v)) + publish(cb["put"], cb["xrdy"], t, 1, tid == 0) + gpu.barrier() + + # ---- the MoE, its all-reduce into ``out`` + moe.run_moe(cb, mkey, bid, 0, 0) + + def reset(): + # the MoE counter slot 128 launch pairs ahead: its last use long + # done, its next use far off + slot = (ep + 128) & 255 + for region in COUNTERS: + gstore(cb[region].value + fx.Int64(slot * 4), fx.Int32(0), words=1) + + mb_back.epoch_end({"bid": bid, "tid": tid}, a["epoch"], ep, reset) + + +def build_mono_k2(key: MonoBuild): + s, tp = key.tokens, key.tp + assert 1 <= s <= MAX_TOKENS + layout = scratch_layout(s, tp) + bkey = _back_key(key) + mkey = _moe_key(key) + rs = route_shape(mkey) + assert rs.experts // 64 in SORT_NETS + BackM = mb_back.back_lds_members(s, Dims(tp)) + SplitLds, GemvLds = BackM["split"], BackM["gemv"] + MoeLds = moe.moe_smem(s, rs, False) + half = peer_half_bytes(tp) + assert SLICES + 2 * s <= BLOCKS + SeamLds = _seam_lds(s) + + @fx.union + class K2Lds: + split: SplitLds # type: ignore[valid-type] + gemv: GemvLds # type: ignore[valid-type] + seam: SeamLds # type: ignore[valid-type] + moe: MoeLds # type: ignore[valid-type] + + name = kernel_symbol("dsv41_mono_k2", s=s, tp=tp, r=key.ratio, tl=key.timeline) + keyed = key_tuple(key, SOURCES) + + @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) + def k2( + q: Int64, + swa: Int64, + swa_stride: Int64, + swa_block: Int32, + comp: Int64, + comp_stride: Int64, + comp_block: Int32, + kt: Int64, + klen: Int64, + pos: Int64, + sink: Int64, + qk_scale: Int32, + cos_sin: Int64, + woa: Int64, + woa_s: Int64, + wob: Int64, + wob_s: Int64, + zrec: Int64, + res_in: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + ffn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + gate_w: Int64, + bias: Int64, + w13: Int64, + w13_s: Int64, + w2: Int64, + w2_s: Int64, + sgu: Int64, + sgu_s: Int64, + sw2: Int64, + sw2_s: Int64, + out: Int64, + scratch: Int64, + sym: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + tl: Int64, + ): + _ = keyed + bid = fx.block_idx.x + lds = fx.SharedAllocator().allocate(K2Lds) + ep = _epoch(epoch) + tag = (ep << 1) + 2 + par = fx.Int64(ep & 1) * fx.Int64(half) + bases = [b + par for b in peer_bases(peers, tp)] + own = sym + par + bargs = { + "q": q, + "swa": swa, + "swa_stride": swa_stride, + "swa_block": swa_block, + "comp": comp, + "comp_stride": comp_stride, + "comp_block": comp_block, + "kt": kt, + "klen": klen, + "pos": pos, + "sink": sink, + "qk_scale": qk_scale, + "cos_sin": cos_sin, + "woa": woa, + "woa_s": woa_s, + "wob": wob, + "wob_s": wob_s, + "out": 0, + "out_stride": 0, + "zrec": zrec, + } + blayout = {n: layout[n] for n in back_scratch(s, Dims(tp), start=0)} + views = {"split": lds.split.peek(), "gemv": lds.gemv.peek()} + c = mb_back.back_context(bkey, views, bargs, scratch, tag, blayout) + amb = Mailbox(tag) + + def push(t, col, v0, v1): + # this rank's bf16 partial pair to every rank's ATTN region + for p in range_constexpr(tp): + dst = preg(bases[p], 0, "attn") + amb.put_bf(dst, 2 * seam.attn_region_pair(rank, t, s, col), [v0, v1]) + + c["wob_out"] = push + if const_expr(key.timeline): + c["tl"], c["tl_points"] = tl, mb_back.BACK_POINTS + mb_back.epoch_begin(c, epoch, ep) + mb_back.run_back(c, bkey, bid) + gpu.barrier() + + ffn = dict( + res_in=res_in, + post_in=post_in, + comb_in=comb_in, + pre_in=pre_in, + hc_fn=hc_fn, + hc_scale=hc_scale, + hc_base=hc_base, + ffn_w=ffn_w, + res_out=res_out, + post_out=post_out, + comb_out=comb_out, + pre_out=pre_out, + gate_w=gate_w, + bias=bias, + w13=w13, + w13_s=w13_s, + w2=w2, + w2_s=w2_s, + sgu=sgu, + sgu_s=sgu_s, + sw2=sw2, + sw2_s=sw2_s, + out=out, + scratch=scratch, + epoch=epoch, + ) + run_ffn(key, lds, ffn, amb, bases, own, rank, ep) + + @flyc.jit + def launch( + q: Int64, + swa: Int64, + swa_stride: Int64, + swa_block: Int32, + comp: Int64, + comp_stride: Int64, + comp_block: Int32, + kt: Int64, + klen: Int64, + pos: Int64, + sink: Int64, + qk_scale: Int32, + cos_sin: Int64, + woa: Int64, + woa_s: Int64, + wob: Int64, + wob_s: Int64, + zrec: Int64, + res_in: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + ffn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + gate_w: Int64, + bias: Int64, + w13: Int64, + w13_s: Int64, + w2: Int64, + w2_s: Int64, + sgu: Int64, + sgu_s: Int64, + sw2: Int64, + sw2_s: Int64, + out: Int64, + scratch: Int64, + sym: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + tl: Int64, + stream: fx.Stream = _STREAM, + ): + _ = keyed + k2( + q, + swa, + swa_stride, + swa_block, + comp, + comp_stride, + comp_block, + kt, + klen, + pos, + sink, + qk_scale, + cos_sin, + woa, + woa_s, + wob, + wob_s, + zrec, + res_in, + post_in, + comb_in, + pre_in, + hc_fn, + hc_scale, + hc_base, + ffn_w, + res_out, + post_out, + comb_out, + pre_out, + gate_w, + bias, + w13, + w13_s, + w2, + w2_s, + sgu, + sgu_s, + sw2, + sw2_s, + out, + scratch, + sym, + peers, + rank, + epoch, + tl, + ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) + + return launch + + +# ---------------------------------------------------------------- the FFN launch + + +def build_mono_ffn(key: MonoBuild): + """The FFN launch of a layer whose attention vLLM ran (``part``: this rank's + unreduced wo_b output, bf16 [S, HIDDEN]): K2's FFN half, the partial pushed + to every rank first. It takes a launch pair's epoch, as K2 does.""" + s, tp = key.tokens, key.tp + assert 1 <= s <= MAX_TOKENS + rs = route_shape(_moe_key(key)) + assert rs.experts // 64 in SORT_NETS + MoeLds = moe.moe_smem(s, rs, False) + half = peer_half_bytes(tp) + assert SLICES + 2 * s <= BLOCKS + SeamLds = _seam_lds(s) + + @fx.union + class FfnLds: + seam: SeamLds # type: ignore[valid-type] + moe: MoeLds # type: ignore[valid-type] + + name = kernel_symbol("dsv41_mono_ffn", s=s, tp=tp) + keyed = key_tuple(key, SOURCES) + + @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) + def kffn( + part: Int64, + res_in: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + ffn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + gate_w: Int64, + bias: Int64, + w13: Int64, + w13_s: Int64, + w2: Int64, + w2_s: Int64, + sgu: Int64, + sgu_s: Int64, + sw2: Int64, + sw2_s: Int64, + out: Int64, + scratch: Int64, + sym: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + ): + _ = keyed + lds = fx.SharedAllocator().allocate(FfnLds) + ep = _epoch(epoch) + par = fx.Int64(ep & 1) * fx.Int64(half) + bases = [b + par for b in peer_bases(peers, tp)] + mb_back.epoch_begin({"bid": fx.block_idx.x, "tid": fx.thread_idx.x}, epoch, ep) + a = dict( + res_in=res_in, + post_in=post_in, + comb_in=comb_in, + pre_in=pre_in, + hc_fn=hc_fn, + hc_scale=hc_scale, + hc_base=hc_base, + ffn_w=ffn_w, + res_out=res_out, + post_out=post_out, + comb_out=comb_out, + pre_out=pre_out, + gate_w=gate_w, + bias=bias, + w13=w13, + w13_s=w13_s, + w2=w2, + w2_s=w2_s, + sgu=sgu, + sgu_s=sgu_s, + sw2=sw2, + sw2_s=sw2_s, + out=out, + scratch=scratch, + epoch=epoch, + ) + amb = Mailbox((ep << 1) + 2) + run_ffn(key, lds, a, amb, bases, sym + par, rank, ep, part=part) + + @flyc.jit + def launch( + part: Int64, + res_in: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + ffn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + gate_w: Int64, + bias: Int64, + w13: Int64, + w13_s: Int64, + w2: Int64, + w2_s: Int64, + sgu: Int64, + sgu_s: Int64, + sw2: Int64, + sw2_s: Int64, + out: Int64, + scratch: Int64, + sym: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + stream: fx.Stream = _STREAM, + ): + _ = keyed + kffn( + part, + res_in, + post_in, + comb_in, + pre_in, + hc_fn, + hc_scale, + hc_base, + ffn_w, + res_out, + post_out, + comb_out, + pre_out, + gate_w, + bias, + w13, + w13_s, + w2, + w2_s, + sgu, + sgu_s, + sw2, + sw2_s, + out, + scratch, + sym, + peers, + rank, + epoch, + ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) + + return launch diff --git a/vllm/models/deepseek_v41/amd/mono/runner.py b/vllm/models/deepseek_v41/amd/mono/runner.py new file mode 100644 index 000000000000..eb379902754f --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/runner.py @@ -0,0 +1,437 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4.1-Flash mono decode layer: host side. + +One ``DSV41MonoLayer`` a TP rank runs every mono layer of a decode step (M <= 48 +rows): two persistent launches a layer (``layer``) from the layer's +inputs at the attention seam to its outputs at the next one -- the MoE's +reduced output, the residual after the FFN seam and that seam's mixes -- with +both TP all-reduces inside the kernels (symmetric peer memory). A layer whose +attention vLLM runs (``ffn``) takes one launch from its unreduced attention +output on. Every argument is a device pointer: the launches can be captured in a +HIP graph. The kernels move their mailbox epoch on themselves; every rank runs +the same launches, so the ranks' epochs agree. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from pathlib import Path + +import torch + +from .attention.plan import HEAD_DIM, HIDDEN, KEYS, Dims +from .common.plan import BLOCKS +from .layer import ( + EPOCH_WORDS, + MAX_TOKENS, + MonoBuild, + build_mono_ffn, + build_mono_k1, + build_mono_k2, + peer_half_bytes, + scratch_bytes, +) +from .stages.dims import Dims as MoeDims +from .stages.moe_shape import EXPERTS + +__all__ = ["BLOCKS", "MAX_TOKENS", "AttnWeights", "DSV41MonoLayer", "MonoLayerWeights"] + +_LOG2E = 1.4426950408889634 +HC = 4 + + +def _ensure_writable_flydsl_cache() -> None: + """Aiter points FlyDSL at its bundled, read-only cache; new kernels need a + writable one.""" + cur = os.environ.get("FLYDSL_RUNTIME_CACHE_DIR") + if cur and os.access(cur, os.W_OK): + return + path = Path.home() / ".flydsl" / "cache" + path.mkdir(parents=True, exist_ok=True) + os.environ["FLYDSL_RUNTIME_CACHE_DIR"] = str(path) + + +def _check_tensors(owner, tag: str, want: dict) -> None: + """Each named tensor of ``owner`` has the (shape, dtype) the kernels read, + contiguous: their buffer loads are unbounded, so a mismatch would read + garbage rather than fault.""" + for name, (shape, dtype) in want.items(): + t = getattr(owner, name) + assert tuple(t.shape) == shape and t.dtype == dtype and t.is_contiguous(), ( + f"{tag} {name}: {tuple(t.shape)} {t.dtype}, want {shape} {dtype}" + ) + + +@dataclass +class AttnWeights: + """One layer's attention tensors on this rank, in vLLM's loaded layout.""" + + layer_id: int + wqkv: torch.Tensor # [1792, 5120] e4m3: [wq_a; wkv] + wqkv_scale: torch.Tensor # [56, 160] uint8 E8M0 (32 x 32 blocks) + q_norm: torch.Tensor # [1280] bf16 + kv_norm: torch.Tensor # [512] bf16 + wq_b: torch.Tensor # [H * 512, 1280] e4m3 + wq_b_scale: torch.Tensor # [H * 16, 40] uint8 + wo_a: torch.Tensor # [G * 1024, 4096] e4m3 + wo_a_scale: torch.Tensor # [G * 32, 128] uint8 + wo_b: torch.Tensor # [5120, G * 1024] e4m3 + wo_b_scale: torch.Tensor # [160, G * 32] uint8 + attn_sink: torch.Tensor # [>= H] f32 + cos_sin: torch.Tensor # [positions, 64] f32 (the layer's rope table) + ratio: int # 0 / 1 / 2: the layer's compress ratio + + def check(self, d: Dims) -> None: + h, g = d.heads, d.groups + want = { + "wqkv": ((1792, HIDDEN), torch.float8_e4m3fn), + "wqkv_scale": ((56, HIDDEN // 32), torch.uint8), + "wq_b": ((h * HEAD_DIM, 1280), torch.float8_e4m3fn), + "wq_b_scale": ((h * HEAD_DIM // 32, 40), torch.uint8), + "wo_a": ((g * 1024, d.group_k), torch.float8_e4m3fn), + "wo_a_scale": ((g * 32, d.group_k // 32), torch.uint8), + "wo_b": ((HIDDEN, g * 1024), torch.float8_e4m3fn), + "wo_b_scale": ((HIDDEN // 32, g * 32), torch.uint8), + } + _check_tensors(self, f"layer {self.layer_id}", want) + assert self.cos_sin.dtype == torch.float32 and self.cos_sin.shape[-1] == 64 + + +@dataclass +class MonoLayerWeights: + """One layer's tensors on this rank, in vLLM's loaded layout (``attn``: + None for a layer whose attention vLLM runs).""" + + attn: AttnWeights | None + hc_attn_fn: torch.Tensor # [24, 4 * 5120] f32 + hc_attn_scale: torch.Tensor # [3] f32 + hc_attn_base: torch.Tensor # [24] f32 + attn_norm: torch.Tensor # [5120] bf16 + hc_ffn_fn: torch.Tensor + hc_ffn_scale: torch.Tensor + hc_ffn_base: torch.Tensor + ffn_norm: torch.Tensor + gate_w: torch.Tensor # [384, 5120] bf16 + bias: torch.Tensor # [384] f32 (e_score_correction_bias) + w13: torch.Tensor # [384, 2 inter, 2560] fp4x2, aiter (16, 16) shuffle + w13_s: torch.Tensor # its scales, aiter shuffle_scale order + w2: torch.Tensor # [384, 5120, inter / 2] fp4x2 + w2_s: torch.Tensor + sgu: torch.Tensor # shared gate_up [2 inter, 5120] e4m3, row-major + sgu_s: torch.Tensor # [2 inter / 32, 160] E8M0 + sw2: torch.Tensor # shared down [5120, inter] e4m3 + sw2_s: torch.Tensor + + def check(self, tp: int) -> None: + """The MoE's tensors: 384 experts, TP-sharded intermediates (the routed + one padded as the loader pads it), AITER's A8W4 layout, 32 x 32 E8M0 + blocks for the shared expert.""" + d, fp4, e4m3, u8 = ( + MoeDims(tp), + torch.float4_e2m1fn_x2, + torch.float8_e4m3fn, + torch.uint8, + ) + want = { + "gate_w": ((EXPERTS, HIDDEN), torch.bfloat16), + "bias": ((EXPERTS,), torch.float32), + "w13": ((EXPERTS, 2 * d.inter, HIDDEN // 2), fp4), + "w13_s": ((EXPERTS * 2 * d.inter, HIDDEN // 32), u8), + "w2": ((EXPERTS, HIDDEN, d.inter // 2), fp4), + "w2_s": ((EXPERTS * HIDDEN, d.down_scale_cols), u8), + "sgu": ((2 * d.sh_inter, HIDDEN), e4m3), + "sgu_s": ((2 * d.sh_inter // 32, HIDDEN // 32), u8), + "sw2": ((HIDDEN, d.sh_inter), e4m3), + "sw2_s": ((HIDDEN // 32, d.sh_inter // 32), u8), + } + _check_tensors(self, "MoE", want) + + +class DSV41MonoLayer: + """The mono decode runner of one TP rank (``group``: its TP group, for the + peer-memory handle exchange; a gloo / CPU group).""" + + def __init__(self, tp: int, rank: int, group, device: torch.device | str = "cuda"): + _ensure_writable_flydsl_cache() + from .common.peer_memory import PeerBuffer + + self.tp, self.rank = tp, rank + self.d = Dims(tp) + dev = self.device = torch.device(device) + self._scratch: dict[int, torch.Tensor] = {} + # [epoch, -, -, -, a mark per CTA, the MoE's counters] + self.epoch = torch.zeros(EPOCH_WORDS, dtype=torch.int32, device=dev) + self.q = torch.zeros( + MAX_TOKENS, self.d.heads, HEAD_DIM, dtype=torch.bfloat16, device=dev + ) + self.kt = torch.zeros(MAX_TOKENS * KEYS, dtype=torch.int32, device=dev) + self.klen = torch.zeros(2 * MAX_TOKENS, dtype=torch.int32, device=dev) + self._dummy = torch.zeros(4, dtype=torch.int32, device=dev) + self._zero_rec = torch.zeros(1024, dtype=torch.uint8, device=dev) + # the attention seam's outputs, K1 -> K2 + self.res_mid = torch.zeros( + MAX_TOKENS, HC, HIDDEN, dtype=torch.bfloat16, device=dev + ) + self.post_a = torch.zeros(MAX_TOKENS, HC, dtype=torch.float32, device=dev) + self.comb_a = torch.zeros(MAX_TOKENS, HC, HC, dtype=torch.float32, device=dev) + self.pre_a = torch.zeros(MAX_TOKENS, HC, dtype=torch.float32, device=dev) + self.peer = PeerBuffer(2 * peer_half_bytes(tp), group, rank, tp, dev) + self.peer.bytes.zero_() + self._qk_scale = ( + torch.tensor([HEAD_DIM**-0.5 * _LOG2E], dtype=torch.float32) + .view(torch.int32) + .item() + ) + self._kernels: dict = {} + + def scratch(self, tokens: int) -> torch.Tensor: + """Step width ``tokens``'s scratch: a width's layout has memory of its own + (``layer.scratch_layout``). Allocated at the width's first step, which is + eager: vLLM runs every graph's batch eagerly before capturing it.""" + buf = self._scratch.get(tokens) + if buf is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"DSv4.1 mono decode: step width {tokens} first reached inside " + "a CUDA graph capture" + ) + buf = torch.zeros( + scratch_bytes(tokens, self.tp), dtype=torch.uint8, device=self.device + ) + self._scratch[tokens] = buf + return buf + + @staticmethod + def supports(tokens: int) -> bool: + return 1 <= tokens <= MAX_TOKENS + + def kernels(self, tokens: int, ratio: int): + key = (tokens, ratio) + if key not in self._kernels: + b = MonoBuild(tokens, self.tp, ratio) + self._kernels[key] = (build_mono_k1(b), build_mono_k2(b)) + return self._kernels[key] + + def ffn_kernel(self, tokens: int): + key = (tokens, "ffn") + if key not in self._kernels: + self._kernels[key] = build_mono_ffn(MonoBuild(tokens, self.tp, 0)) + return self._kernels[key] + + @staticmethod + def _outs(M: int, residual: torch.Tensor) -> tuple: + dev = residual.device + return ( + torch.empty(M, HIDDEN, dtype=torch.bfloat16, device=dev), + torch.empty_like(residual), + torch.empty(M, HC, 1, dtype=torch.float32, device=dev), + torch.empty(M, HC, HC, dtype=torch.float32, device=dev), + torch.empty(M, HC, dtype=torch.float32, device=dev), + ) + + def forward( + self, + w: MonoLayerWeights, + x: torch.Tensor, # [M, 5120] bf16: the previous FFN output (reduced) + residual: torch.Tensor, # [M, 4, 5120] bf16 + post_mix: torch.Tensor, # [M, 4, 1] f32 + res_mix: torch.Tensor, # [M, 4, 4] f32 + pre_mix: torch.Tensor, # [M, 4] f32 + positions: torch.Tensor, # [M] int64 + slot_mapping: torch.Tensor, # [M] int64 + swa_cache: torch.Tensor, # [blocks, block, 584] uint8 + swa_indices: torch.Tensor, # [>= M, 1, 128] int32 + swa_lens: torch.Tensor, # [>= M] int32 + token_to_req: torch.Tensor, # [>= M] int32 + topk_indices: torch.Tensor | None = None, + comp_cache: torch.Tensor | None = None, + comp_block_table: torch.Tensor | None = None, + outs: tuple | None = None, + ): + """-> (out, residual, post_mix, res_mix, pre_mix): the layer's outputs + (``outs``: buffers to write them into).""" + M = x.shape[0] + assert self.supports(M), M + a = w.attn + assert a is not None + ratio = a.ratio + k1, k2 = self.kernels(M, ratio) + st = torch.cuda.current_stream() + if outs is None: + outs = self._outs(M, residual) + out, res_out, post_out, comb_out, pre_out = outs + for t_ in (x, residual, post_mix, res_mix, pre_mix, positions, slot_mapping): + assert t_.is_contiguous() + assert positions.dtype == slot_mapping.dtype == torch.int64 + assert swa_indices.dtype == swa_lens.dtype == token_to_req.dtype == torch.int32 + assert swa_indices.stride(0) == 128 and swa_cache.dtype == torch.uint8 + assert swa_cache.shape[-1] == 584 and swa_cache.stride(1) == 584 + swa_block = swa_cache.shape[1] + if ratio: + assert topk_indices is not None and comp_cache is not None + assert comp_block_table is not None + assert topk_indices.dtype == torch.int32 and topk_indices.stride(0) == 512 + assert comp_cache.dtype == torch.uint8 and comp_cache.stride(1) == 584 + assert comp_block_table.dtype == torch.int32 + comp, comp_stride, comp_block = ( + comp_cache, + comp_cache.stride(0), + comp_cache.shape[1], + ) + bt, bt_stride, topk = ( + comp_block_table, + comp_block_table.stride(0), + topk_indices, + ) + else: + comp, comp_stride, comp_block = self._dummy, 0, 1 + bt, bt_stride, topk = self._dummy, 0, self._dummy + k1( + residual.data_ptr(), + x.data_ptr(), + post_mix.data_ptr(), + res_mix.data_ptr(), + pre_mix.data_ptr(), + w.hc_attn_fn.data_ptr(), + w.hc_attn_scale.data_ptr(), + w.hc_attn_base.data_ptr(), + w.attn_norm.data_ptr(), + self.res_mid.data_ptr(), + self.post_a.data_ptr(), + self.comb_a.data_ptr(), + self.pre_a.data_ptr(), + a.wqkv.data_ptr(), + a.wqkv_scale.data_ptr(), + a.q_norm.data_ptr(), + a.kv_norm.data_ptr(), + a.wq_b.data_ptr(), + a.wq_b_scale.data_ptr(), + a.cos_sin.data_ptr(), + positions.data_ptr(), + slot_mapping.data_ptr(), + swa_cache.data_ptr(), + swa_cache.stride(0), + swa_block, + self.q.data_ptr(), + swa_indices.data_ptr(), + swa_lens.data_ptr(), + comp_block, + bt.data_ptr(), + bt_stride, + topk.data_ptr(), + token_to_req.data_ptr(), + self.kt.data_ptr(), + self.klen.data_ptr(), + self.scratch(M).data_ptr(), + self.epoch.data_ptr(), + 0, + stream=st, + ) + k2( + self.q.data_ptr(), + swa_cache.data_ptr(), + swa_cache.stride(0), + swa_block, + comp.data_ptr(), + comp_stride, + comp_block, + self.kt.data_ptr(), + self.klen.data_ptr(), + positions.data_ptr(), + a.attn_sink.data_ptr(), + self._qk_scale, + a.cos_sin.data_ptr(), + a.wo_a.data_ptr(), + a.wo_a_scale.data_ptr(), + a.wo_b.data_ptr(), + a.wo_b_scale.data_ptr(), + self._zero_rec.data_ptr(), + self.res_mid.data_ptr(), + self.post_a.data_ptr(), + self.comb_a.data_ptr(), + self.pre_a.data_ptr(), + w.hc_ffn_fn.data_ptr(), + w.hc_ffn_scale.data_ptr(), + w.hc_ffn_base.data_ptr(), + w.ffn_norm.data_ptr(), + res_out.data_ptr(), + post_out.data_ptr(), + comb_out.data_ptr(), + pre_out.data_ptr(), + w.gate_w.data_ptr(), + w.bias.data_ptr(), + w.w13.data_ptr(), + w.w13_s.data_ptr(), + w.w2.data_ptr(), + w.w2_s.data_ptr(), + w.sgu.data_ptr(), + w.sgu_s.data_ptr(), + w.sw2.data_ptr(), + w.sw2_s.data_ptr(), + out.data_ptr(), + self.scratch(M).data_ptr(), + self.peer.local, + self.peer.addresses.data_ptr(), + self.rank, + self.epoch.data_ptr(), + 0, + stream=st, + ) + return outs + + def ffn( + self, + w: MonoLayerWeights, + part: torch.Tensor, # [M, 5120] bf16: this rank's unreduced wo_b output + residual: torch.Tensor, # [M, 4, 5120] bf16, after the attention seam + post_mix: torch.Tensor, # [M, 4, 1] f32, the attention seam's + res_mix: torch.Tensor, # [M, 4, 4] f32 + pre_mix: torch.Tensor, # [M, 4] f32 + outs: tuple | None = None, + ): + """-> (out, residual, post_mix, res_mix, pre_mix) of a layer whose + attention vLLM ran: the attention's TP reduce, the FFN seam, the MoE and + its all-reduce in one launch.""" + M = part.shape[0] + assert self.supports(M), M + for t_ in (part, residual, post_mix, res_mix, pre_mix): + assert t_.is_contiguous() + assert part.dtype == residual.dtype == torch.bfloat16 + if outs is None: + outs = self._outs(M, residual) + out, res_out, post_out, comb_out, pre_out = outs + self.ffn_kernel(M)( + part.data_ptr(), + residual.data_ptr(), + post_mix.data_ptr(), + res_mix.data_ptr(), + pre_mix.data_ptr(), + w.hc_ffn_fn.data_ptr(), + w.hc_ffn_scale.data_ptr(), + w.hc_ffn_base.data_ptr(), + w.ffn_norm.data_ptr(), + res_out.data_ptr(), + post_out.data_ptr(), + comb_out.data_ptr(), + pre_out.data_ptr(), + w.gate_w.data_ptr(), + w.bias.data_ptr(), + w.w13.data_ptr(), + w.w13_s.data_ptr(), + w.w2.data_ptr(), + w.w2_s.data_ptr(), + w.sgu.data_ptr(), + w.sgu_s.data_ptr(), + w.sw2.data_ptr(), + w.sw2_s.data_ptr(), + out.data_ptr(), + self.scratch(M).data_ptr(), + self.peer.local, + self.peer.addresses.data_ptr(), + self.rank, + self.epoch.data_ptr(), + stream=torch.cuda.current_stream(), + ) + return outs diff --git a/vllm/models/deepseek_v41/amd/mono/sources.py b/vllm/models/deepseek_v41/amd/mono/sources.py new file mode 100644 index 000000000000..49f34d8fed67 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/sources.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The mono layer's kernel sources, in every build's JIT cache key +(``common.plan.key_tuple``): an edit anywhere in this package rebuilds.""" + +from vllm.models.deepseek_v41.amd.mono.common.plan import source_digest + +SOURCES = source_digest(".") diff --git a/vllm/models/deepseek_v41/amd/mono/stages/__init__.py b/vllm/models/deepseek_v41/amd/mono/stages/__init__.py new file mode 100644 index 000000000000..ebb582c99b8d --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The V4.1 stages the mono layer adapts from ROCm/ATOM: the mHC seams +(``seam``), the MoE (``moe``) and the GEMV / LDS helpers they share (``gemv``).""" diff --git a/vllm/models/deepseek_v41/amd/mono/stages/dims.py b/vllm/models/deepseek_v41/amd/mono/stages/dims.py new file mode 100644 index 000000000000..6fd119c77206 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/dims.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/dims.py +"""A TP rank's share of V4.1's MoE widths: every constant the MoE stages derive +from the TP size, in one place (a build key's ``tp`` picks it).""" + +from dataclasses import dataclass + +from vllm.models.deepseek_v41.amd.mono.common.plan import Shard, padded + +# the model's full widths +MOE_INTER = 2304 # a routed expert's intermediate width +SHARED_INTER = 2304 # the shared expert's +# the loader pads a routed expert's rank share to a multiple of this (``FusedMoE`` +# ``pad_align``) with zero weights +MOE_PAD = 128 +ROWS = 16 # GEMV rows a task +UG_PART = 32 # intermediate columns an ug task part: one FP4 group + + +@dataclass(frozen=True) +class Dims: + tp: int + + def __post_init__(self): + shard = Shard(self.tp) + shard.split(MOE_INTER, "routed intermediate") + shard.split(SHARED_INTER, "shared intermediate") + + @property + def inter_real(self) -> int: + """A routed expert's real intermediate width on this rank.""" + return MOE_INTER // self.tp + + @property + def inter(self) -> int: + """... padded as the loader pads it (the rest zero weights).""" + return padded(self.inter_real, MOE_PAD) + + @property + def sh_inter(self) -> int: + """The shared expert's intermediate width on this rank.""" + return SHARED_INTER // self.tp + + @property + def shared_tasks(self) -> int: + return self.sh_inter // ROWS + + @property + def mid_words(self) -> int: + """MXFP8 words of a pick's routed intermediate.""" + return self.inter // 4 + + @property + def down_scale_cols(self) -> int: + """w2_s's e8m0 columns: inter / 32, padded to a multiple of 8 + (aiter ``shuffle_scale``).""" + return padded(self.inter // 32, 8) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/gemv.py b/vllm/models/deepseek_v41/amd/mono/stages/gemv.py new file mode 100644 index 000000000000..31a19ea4c72f --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/gemv.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/attn_pre.py (its GEMV and LDS helpers). +"""The FP8 GEMV tiles of the seam and MoE stages and the helpers they share: +token tiles, LDS rows of MXFP8 words, scalar loads and the mailbox-to-LDS +copy.""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr.typing import T + +from vllm.models.deepseek_v41.amd.mono.common.mx import ( + UNIT_SCALE, + mfma_scaled, +) +from vllm.models.deepseek_v41.amd.mono.common.ops import ( + CM_NT, + fresh, + rsrc, + traced, +) +from vllm.models.deepseek_v41.amd.mono.common.plan import WAVES +from vllm.models.deepseek_v41.amd.mono.common.sync import POLL_MAX + +ROWS = 16 # GEMV rows a task +TILE = 16 # an MFMA's N: the tokens a GEMV pass takes (``token_tiles``) + + +def ld_f32(ptr, i): + return fx.Float32(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.f32)) + + +def ld_bf(ptr, i): + return fx.Float32( + fx.BFloat16(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.bf16)) + ) + + +def token_tiles(s): + """A step's tokens in GEMV passes: [(first token, tokens)]. A GEMV loads its + weights once and runs every pass on them, an x tile in LDS at a time.""" + return [(t0, min(TILE, s - t0)) for t0 in range(0, s, TILE)] + + +def tile_rows(s): + """The LDS rows an x tile holds.""" + return min(s, TILE) + + +# a token's LDS row of MXFP8 words / E8M0 codes, 16 of which a MFMA B operand +# gathers, is padded by 16 B: at a multiple of 64 words they share a bank group +LDS_PAD = 4 + + +def lds_row(words): + return words + LDS_PAD + + +def lds_at(u, row): + """Element u of a span of rows ``row`` wide, at the ``lds_row`` stride.""" + return u + u // row * LDS_PAD + + +def lds_tail(k): + """Words a K = ``k`` row's last K step reads past its ``lds_row`` (selected + out): the last row's buffer must hold them.""" + return max(0, 32 * -(-k // 128) - lds_row(k // 4)) + + +def plus(t0, v): + """Token ``v`` of the tile from ``t0`` (a first tile adds nothing; ``t0`` + traced: a tile a CTA picks at run time).""" + return v if isinstance(t0, int) and t0 == 0 else v + t0 + + +# ---------------------------------------------------------------- GEMV +def _fp8_step_load(c, r_w, r_ws, k, rg, st, row_major=False): + """K step st's global operands of rows 16 rg ..: the two 64-column weight + chunks (the upper clamped into the row) and the raw scale word. The (16, 16) + shuffle lays a row group out in 32-column blocks of 16 rows, a chunk two of + them (lanes 0-31, 32-63); ``k`` % 64 = 32 (a last quarter step): the lanes + of a block past the row read its last block. ``row_major`` (vLLM's dense + layout, k % 64 == 0): the same registers from a row-major [N, K] weight -- + lane l row 16 rg + l % 16, K bytes 128 st + 16 (l // 16) .. and + 64.""" + lane = c["lane"] + j = lane // 16 + kg = k // 32 + if const_expr(row_major): + assert k % 64 == 0 + row = rg * 16 + lane % 16 + lo_at = row * (k // 4) + st * 32 + j * 4 + hi_at = row * (k // 4) + fx.min(2 * st + 1, k // 64 - 1) * 16 + j * 4 + lo = fx.Vector( + bo.buffer_load(r_w, lo_at, vec_width=4, dtype=T.i32, cache_modifier=CM_NT) + ) + hi = fx.Vector( + bo.buffer_load(r_w, hi_at, vec_width=4, dtype=T.i32, cache_modifier=CM_NT) + ) + wsi = (rg // 2) * kg + fx.min(st * 4 + j, kg - 1) + word = fx.Int32(bo.buffer_load(r_ws, wsi // 4, vec_width=1, dtype=T.i32)) + return lo, hi, (word >> (wsi % 4 * 8)) & 0xFF + if const_expr(k % 64 == 0): + lo_at = (rg * (k // 64) + 2 * st) * 256 + lane * 4 + else: + lo_at = rg * (k * 4) + fx.min(4 * st + lane // 32, kg - 1) * 128 + lane % 32 * 4 + lo = fx.Vector( + bo.buffer_load(r_w, lo_at, vec_width=4, dtype=T.i32, cache_modifier=CM_NT) + ) + if const_expr(k % 64 == 0): + hi_at = (rg * (k // 64) + fx.min(2 * st + 1, k // 64 - 1)) * 256 + lane * 4 + else: + hi_at = rg * (k * 4) + fx.min(2 * st + 1, k // 64 - 1) * 256 + lane * 4 + hi = fx.Vector( + bo.buffer_load(r_w, hi_at, vec_width=4, dtype=T.i32, cache_modifier=CM_NT) + ) + wsi = (rg // 2) * kg + fx.min(st * 4 + j, kg - 1) + word = fx.Int32(bo.buffer_load(r_ws, wsi // 4, vec_width=1, dtype=T.i32)) + return lo, hi, (word >> (wsi % 4 * 8)) & 0xFF + + +def _fp8_step_mfma(c, k, xl, xsl, st, ops, acc, live=None, rows=None): + """``acc`` through K step st's scaled MFMA: ``_fp8_step_load``'s operands + against every token's MXFP8 row in LDS. ``live`` false: a step past the end + (its operands those of the last step), zero at unit scale. ``k`` % 64 = 32: + a lane's lower block past the row is zero on both sides.""" + s, lane = c["S"] if rows is None else rows, c["lane"] + lo, hi, sa = ops + kg = k // 32 + x_row, s_row = lds_row(k // 4), lds_row(kg) + j = lane // 16 + col = fx.min(lane % 16, s - 1) + if const_expr(live is not None): + st = fx.min(st, (k + 127) // 128 - 1) + lo = [live.select(lo[d], fx.Int32(0)) for d in range(4)] + past = (2 * st + 1) * 64 >= k + if const_expr(live is not None): + past = past | ~live + av = fx.Vector.from_elements( + [lo[d] for d in range(4)] + [past.select(fx.Int32(0), hi[d]) for d in range(4)], + fx.Int32, + ) + grp = st * 4 + j + xlo = fx.Vector( + fx.ptr_load( + xl + (col * x_row + st * 32 + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + xhi = fx.Vector( + fx.ptr_load( + xl + (col * x_row + st * 32 + 16 + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + if const_expr(k % 64 != 0): + over = 4 * st + j // 2 >= kg + av = fx.Vector.from_elements( + [over.select(fx.Int32(0), av[d]) for d in range(4)] + + [av[4 + d] for d in range(4)], + fx.Int32, + ) + xlo = [over.select(fx.Int32(0), xlo[d]) for d in range(4)] + xv = fx.Vector.from_elements( + [xlo[d] for d in range(4)] + + [past.select(fx.Int32(0), xhi[d]) for d in range(4)], + fx.Int32, + ) + sb = fx.ptr_load(xsl + (col * s_row + fx.min(grp, kg - 1))) + unit = grp >= kg + if const_expr(live is not None): + unit = unit | ~live + sa = unit.select(fx.Int32(UNIT_SCALE), sa) + sb = unit.select(fx.Int32(UNIT_SCALE), sb) + return mfma_scaled(av, xv, acc, sa, sb) + + +def _gemv_steps(c, k, split, tiled=False): + """This wave's K steps (a step past the end marked, ``split`` waves past + split repeat the last one's; ``tiled``: wave w takes part w % split of tile + w / split): [(step, live)].""" + wave = c["wave"] + steps = (k + 127) // 128 + if const_expr(split is None): + per = (steps + WAVES - 1) // WAVES + ragged = steps % WAVES != 0 + return [ + (wave + WAVES * i, (wave + WAVES * i < steps) if ragged else None) + for i in range(per) + ] + assert split <= WAVES and steps % split == 0 + per = steps // split + part = wave % split if tiled else fx.min(wave, split - 1) + return [(part * per + i, None) for i in range(per)] + + +def gemv_fp8_loads(c, w, ws, k, rg, split=None, tiled=False, row_major=False): + """The weight operands of ``gemv_fp8``'s rows for this wave, every one in + flight at once (a step's latency each would put them in series): loads + only. A step past the end loads the last step. ``tiled`` (split x tiles = + WAVES): row groups rg .., one a group of ``split`` waves.""" + steps = (k + 127) // 128 + r_w, r_ws = rsrc(w), rsrc(ws) + if const_expr(tiled): + assert WAVES % split == 0 + rg = rg + c["wave"] // split + return [ + _fp8_step_load( + c, + r_w, + r_ws, + k, + rg, + st if live is None else fx.min(st, steps - 1), + row_major, + ) + for st, live in _gemv_steps(c, k, split, tiled) + ] + + +def gemv_fp8_mfmas(c, k, xl, xsl, red, ops, split=None, tiled=False, rows=None): + """``gemv_fp8_loads``' operands against the rows in LDS (``rows``: a token + tile's, every token's by default) -> + ``red[wave][lane]``; the scaled 16x16x128 MFMA, chunk kc in dwords 0-3 and kc + + 1 in 4-7. ``k`` a multiple of 32: a last half step has its upper chunk zeroed + on both sides at unit scale (a NaN byte past the row must not reach the + MFMA), its loads clamped into the row: the last row group's would read past + the tensor; a last quarter step its lower chunk's upper block too. A step + past the end adds zero at unit scale. + + The K order is the original kernel's, which ``row_sum`` completes: + wave w takes steps w, w + 8, ... (``split`` None: aiter's preshuffled group32 + GEMM, P0.5 microbenchmark 1), or waves 0 .. split - 1 one contiguous K range + each, the other waves zero (aiter's split-K bmm, whose last split sums the + partials in split order).""" + lane, wave = c["lane"], c["wave"] + acc = fx.Vector.filled(4, 0.0, fx.Float32) + for i, (st, live) in enumerate(_gemv_steps(c, k, split, tiled)): + acc = _fp8_step_mfma(c, k, xl, xsl, st, ops[i], acc, live, rows) + if const_expr(split is not None and not tiled): + acc = fx.Vector.from_elements( + [(wave < split).select(acc[e], fx.Float32(0.0)) for e in range(4)], + fx.Float32, + ) + fx.ptr_store(acc, red + (wave * 64 + lane) * 4) + + +# ---------------------------------------------------------------- mailbox -> LDS +@traced +def poll_copy(c, idx, width, spans, live=None): + """Mailbox spans [(region, pair_of, n, dst, row)] into LDS: element u (pair + ``pair_of(u)``) to dst + u, or ``lds_at(u, row)`` for a span of rows + ``row`` words wide. Thread ``idx`` of ``width`` takes u = idx, idx + + width, ..., POLL_MAX a round trip, each batch stored before the next is + polled (holding them all spills past S = 16); a thread past a span's end + re-reads its last element and stores nothing. + + ``live``: a span's elements the reader needs (traced, <= n), a span each; + a batch past them is not polled and an element past them neither polled + nor written (one past a step's rows is never written).""" + idx = fresh(idx) # its elements' indices not hoisted out of a task loop + if const_expr(live is None): + _poll_store_batches( + c, + [ + (region, pair_of, n, (dst, row), idx + width * i) + for region, pair_of, n, dst, row in spans + for i in range((n + width - 1) // width) + ], + ) + return + for (region, pair_of, n, dst, row), need in zip(spans, live): + chunks = [ + (region, pair_of, need, (dst, row), idx + width * i) + for i in range((n + width - 1) // width) + ] + for b0 in range_constexpr(0, len(chunks), POLL_MAX): + if chunks[b0][4] < need: + _poll_store_batches(c, chunks[b0 : b0 + POLL_MAX]) + + +@traced +def _poll_store_batches(c, chunks): + """``poll_copy``'s chunks [(region, pair_of, n, (dst, row), u)], POLL_MAX a + round trip, each batch stored before the next is polled.""" + for b0 in range_constexpr(0, len(chunks), POLL_MAX): + batch = chunks[b0 : b0 + POLL_MAX] + got = c["poll"]([(r, p(fx.min(u, n - 1)), 1) for r, p, n, _, u in batch]) + for (_, _, n, (dst, row), u), v in zip(batch, got): + at = u if row is None else lds_at(u, row) + if u < n: + fx.ptr_store(v[0], dst + at) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/moe.py b/vllm/models/deepseek_v41/amd/mono/stages/moe.py new file mode 100644 index 000000000000..63f8947166d2 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/moe.py @@ -0,0 +1,1635 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/moe.py +"""The mono layer's MoE stages: one V4.1 layer's FFN, from the FFN seam's normed +rows to the reduced output (its TP all-reduce in-kernel). + +One launch a layer and rank, ``BLOCKS`` x ``THREADS``, S = ``tokens`` rows, +E routed experts and top-k (``MoeBuild``: the target's 384 / 6, the DSpark +draft's 128 / 3). The rank holds every routed expert's 1/TP of the +intermediate width (TP-sharded MoE) and the shared expert's 1/TP (``Dims``). +Stages, with TP4's task counts: + + xq (S): the row as MXFP8 for the routed experts (aiter + ``fused_mx_quant_moe_sort``: amax seeded 1e-10) and for the shared + one (ATOM ``quantize_fp8``: amax floored 1e-4) -> X8R, X8 (plain + words, then the token's flag X8RDY) + router (E / 16 x 4, 16 rows x a quarter of K): bf16 GEMV -> the part's + f32 sum -> LOGIT + route (S): the parts summed in order -> bf16 logits; aiter + ``topk_gating_kernel_opt`` (sqrtsoftplus + bias, top-k, renormalized + x 1.5), one wave a token -> ROUTE + shared (36, 16 gate + 16 up rows): FP8 GEMV -> bf16 -> aiter silu_and_mul + (gate <= 10, up in [-10, 10], ocml exp) -> bf16 -> SMID + ug (|U| x 9, 2 x (32 gate + 32 up) rows of one expert of the union U of + every token's top-k -- the draft's |U| x 18 of 1 x (32 + 32): + ``RouteShape.ug_groups``; w13 gate / up interleaved a 16-row tile): + the MXFP4 x MXFP8 GEMV, a wave the gate and up 16-row tiles of one + 32-row scale block over a half (the draft's: a quarter) of K (each + B operand from LDS read once for both), the slices summed + pairwise; then its epilogue: + clamp, silu x up, bf16, + and the stage-2 input quant to MXFP8 (amax seeded 1e-10, the ceil + code) -> MID (plain words at device scope, then the task's flag, UGF) + down (320, 16 rows; a CTA's two run together): per expert of U the MXFP4 + x MXFP8 GEMV of its intermediate (x its routing weight -> bf16), + summed in each token's top-k order, rounded to bf16 at every add + (aiter's stage 2 adds bf16 atomically, in no fixed order); the + shared expert's FP8 down projection (quantize_fp8 of SMID) -> bf16; + routed + shared -> bf16 -> pushed to every rank's FFN region + reduce (160, 32 columns): the TP partials in rank order 0 .. TP - 1 + (``sum_partials``), fp32 -> bf16 -> ``out`` +""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import T + +from vllm.models.deepseek_v41.amd.mono.common.mx import ( + FP4, + FP8, + FP8_MAX, + UNIT_SCALE, + clamp_fp8, + code_ceil, + fp4_tile_load, + mfma_scaled, + rcp_pow2, + scale_index, +) +from vllm.models.deepseek_v41.amd.mono.common.ops import ( + CM_DEV, + CM_NT, + _hw_f32, + ballot, + bf16_round, + bf_hi, + bf_lo, + butterfly, + div_rn, + fp8_pack4, + fresh, + hw_exp2, + hw_rcp, + lane_gather, + lanes_below, + popcount, + row_sum, + rsrc, + stamp, + sum_partials, + traced, +) +from vllm.models.deepseek_v41.amd.mono.common.plan import ( + BLOCKS, + THREADS, + WAVES, + first_task, + pair_layout, +) +from vllm.models.deepseek_v41.amd.mono.common.sync import preg, publish, sreg +from vllm.models.deepseek_v41.amd.mono.stages import moe_plan as mplan +from vllm.models.deepseek_v41.amd.mono.stages.dims import UG_PART + +# helpers the kernel calls are imported by name: FlyDSL keys a build by the +# source of the same-directory functions it names, never through a module +from vllm.models.deepseek_v41.amd.mono.stages.gemv import ( + ROWS, + TILE, + gemv_fp8_loads, + gemv_fp8_mfmas, + ld_f32, + lds_row, + lds_tail, + plus, + poll_copy, + tile_rows, + token_tiles, +) +from vllm.models.deepseek_v41.amd.mono.stages.moe_shape import ( + MASK_WORDS, + MoeBuild, + route_shape, +) +from vllm.models.deepseek_v41.amd.mono.stages.router import PARTS as ROUTER_PARTS +from vllm.models.deepseek_v41.amd.mono.stages.router import ( + router_tile_mfma, + router_weight_loads, + router_x_loads, +) +from vllm.models.deepseek_v41.amd.mono.stages.seam import HIDDEN, attn_peer_bytes + +# a pick tile is a token tile's MFMA N: ``gather_x8r`` fills one from the other +assert mplan.TILE == TILE +SCALE = 1.5 +LIMIT = 10.0 # swiglu_limit +LOG2E = 1.4426950408889634 +LN2 = 0.6931471805599453 +DOWN_TASKS = HIDDEN // ROWS # 320 +# (task, slot) scale blocks a down wave holds a round (``stage_down``): 3 took +# no less time, 4 spills +DOWN_R_MAX = 2 +# a down slot's K steps on this many independent MFMA accumulators +DOWN_CHAINS = 2 +# a reduce task's columns: at 128 a thread polled 6 pairs at 48 rows, each +# pass a round trip in series (32: 1-2 passes, S=48 -1 us) +RED_COLS = 32 +RED_TASKS = HIDDEN // RED_COLS # 160 +# CTA b polls route copy b % ROUTE_COPIES: every CTA on one small region +# serializes on its lines (S=24: -0.5 us) +ROUTE_COPIES = 8 +X8_WORDS = HIDDEN // 4 +TAGS = 256 # the ug queue's / down count's slots, a word each: ``c["tag"]`` picks one + +XQ0 = 0 +DOWN0 = 0 +RED0 = 0 + + +# ``timeline`` stamps: kernel start, then the end of each stage's task loop +TL_STAGES = ("xq", "router", "route", "shared", "ug", "down", "reduce") +TL_POINTS = 1 + len(TL_STAGES) + + +# the MoE stages' kernel arguments, in launch order +ARGS = ( + "x", + "gate_w", + "bias", + "w13", + "w13_s", + "w2", + "w2_s", + "sgu", + "sgu_s", + "sw2", + "sw2_s", + "out", +) + + +def scratch_layout(key: MoeBuild) -> dict[str, tuple[int, int]]: + tokens, experts, topk = key.tokens, key.experts, key.topk + rs = route_shape(key) + d = rs.dims + pairs = { + # the routed experts' MXFP8 x (aiter's fused_mx_quant_moe_sort) and the + # shared expert's (ATOM quantize_fp8): their amax floors differ + "x8r": tokens * X8_WORDS // 2, + "x8rs": tokens * HIDDEN // 32 // 2, + "x8": tokens * X8_WORDS // 2, + "x8s": tokens * HIDDEN // 32 // 2, + # a token's four MXFP8 regions written (plain words): ``stage_xq`` + "x8rdy": tokens, + "logit": tokens * experts * ROUTER_PARTS, + "route": ROUTE_COPIES * tokens * 2 * topk, + "smid": tokens * d.sh_inter, + # SMID quantize_fp8'd (``smid_quant``): the down stage's MXFP8 rows + "smq": tokens * d.sh_inter // 4, + "smc": tokens * d.sh_inter // 32, + # a pick a row, in the slot order (``RouteShape.order0``) + "mid": rs.picks * d.mid_words, + "mids": rs.picks * d.inter // 32, + # an ug task's done flag (``stage_ug``) + "ugf": (rs.max_tiles if rs.multi else rs.max_u) * rs.ug_parts, + # the ug task queue: a word a mailbox tag (a layer), counting up from the + # step's clear (``take_ug_task``); plain words, 2 a pair + "ugq": TAGS // 2, + # the down stage's arrival count, a word a tag (``take_down_rank``) + "dq": TAGS // 2, + } + return pair_layout(pairs.items()) + + +def peer_bytes(tokens: int, tp: int) -> int: + """FFN: ``[source rank][token][hidden / 2]`` bf16 pairs.""" + return tp * tokens * HIDDEN // 2 * 8 + + +def ffn_offset(tokens: int, tp: int) -> int: + """FFN's byte offset in the peer buffer: past the ATTN region.""" + return attn_peer_bytes(tokens, tp) + + +def region_pair(src, t, s, col): + return (src * s + t) * (HIDDEN // 2) + col // 2 + + +# ---------------------------------------------------------------- xq +@traced +def stage_xq(c, t): + """Token t, a group of 32 a thread.""" + tid = c["tid"] + a = c["args"] + # thread g < 160 owns the 32 values of group g; the rest mirror thread 159 + g = fx.min(tid, HIDDEN // 32 - 1) + mine = tid < HIDDEN // 32 + # 32 bf16 as four 16-byte loads (a load an element made the stage issue bound) + words = [ + fx.Vector( + bo.buffer_load( + rsrc(a["x"]), + (t * HIDDEN + g * 32) // 2 + 4 * q, + vec_width=4, + dtype=T.i32, + cache_modifier=c["x_cm"], + ) + ) + for q in range(4) + ] + v = [(bf_lo if i % 2 == 0 else bf_hi)(words[i // 8][i % 8 // 2]) for i in range(32)] + amax0 = abs(v[0]) + for y in v[1:]: + amax0 = fx.max(amax0, abs(y)) + # the routed experts' (aiter fused_mx_quant_moe_sort): amax floored 1e-10; + # the shared expert's: vLLM's mxfp8_e4m3_quantize (its MXFP8 linear) + if mine: + put_mxfp8_group(c, "x8r", "x8rs", t, g, v, fx.max(amax0, fx.Float32(1e-10))) + put_mxfp8_group_vllm(c, "x8", "x8s", t, g, v, amax0) + + +def put_mxfp8_group(c, words_mb, codes_mb, t, g, v, amax): + """Token t's group g of 32 values ``v`` as MXFP8 (the ceil code of + ``amax`` / 448, x / scale clamped, RNE) -> scratch ``words_mb`` / + ``codes_mb`` (plain words: the token's X8RDY announces them).""" + code = code_ceil(amax * fx.Float32(1.0 / FP8_MAX)) + inv = rcp_pow2(code) + words = [] + for w4 in range_constexpr(4): + ws = [ + fp8_pack4(*[clamp_fp8(v[8 * w4 + 4 * h + i] * inv) for i in range(4)]) + for h in range(2) + ] + words.append(ws) + for h in range_constexpr(2): + st_dev( + fx.Vector.from_elements(words[2 * h] + words[2 * h + 1], fx.Int32), + c[words_mb], + t * X8_WORDS + g * 8 + 4 * h, + ) + st_dev(fx.Int32(code), c[codes_mb], t * (HIDDEN // 32) + g) + + +FLT_MIN = 1.1754943508222875e-38 + + +def vllm_mx_code(amax): + """VLLM's ``mxfp8_e4m3_quantize`` scale code of a 32-group: the smallest + power of two >= RN(max(amax, FLT_MIN) / 448), clamped to [0, 254].""" + a = fx.max(amax, fx.Float32(FLT_MIN)) + q = div_rn(a, fx.Float32(FP8_MAX), fx.Float32(1.0 / FP8_MAX)) + bits = q.bitcast(fx.Int32) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).select( + fx.Int32(1), fx.Int32(0) + ) + return fx.min(fx.max(e, fx.Int32(0)), fx.Int32(254)) + + +def put_mxfp8_group_vllm(c, words_mb, codes_mb, t, g, v, amax): + """``put_mxfp8_group`` with vLLM's rule (``vllm_mx_code``; x * 2^(127 - code), + which code 0 keeps finite).""" + code = vllm_mx_code(amax) + inv = ((254 - code) << 23).bitcast(fx.Float32) + words = [] + for w4 in range_constexpr(4): + ws = [ + fp8_pack4(*[clamp_fp8(v[8 * w4 + 4 * h + i] * inv) for i in range(4)]) + for h in range(2) + ] + words.append(ws) + for h in range_constexpr(2): + st_dev( + fx.Vector.from_elements(words[2 * h] + words[2 * h + 1], fx.Int32), + c[words_mb], + t * X8_WORDS + g * 8 + 4 * h, + ) + st_dev(fx.Int32(code), c[codes_mb], t * (HIDDEN // 32) + g) + + +# ---------------------------------------------------------------- router +@traced +def stage_router(c, task): + """Router task (token group, row tile, K part): K part ``part`` of rows + 16 ``tile`` .. of the bf16 gate weight [experts, 5120] times the group's + tokens (``RouteShape.router_groups``), a token tile at a time -- the + decode router's tile (``router``), the + same tiles in the same order -> the part's f32 sum (``route`` sums the + parts in order, rounds to bf16). Every tile's x loads issue first: a task + over every token's tiles in series read their x serially (48 rows: a router + CTA 9 us, 8 of them x).""" + s, tid, red, lane, wave = c["S"], c["tid"], c["red"], c["lane"], c["wave"] + a = c["args"] + experts = c["rs"].experts + groups, per = c["rs"].router_groups + per_group = experts // ROWS * ROUTER_PARTS + tile, part = task % per_group // ROUTER_PARTS, task % ROUTER_PARTS + g0 = 0 if groups == 1 else task // per_group * per + # the group's token tiles: (first token, tokens), tokens traced past one group + tiles = [ + ( + plus(g0, j0), + min(TILE, s - j0) + if groups == 1 + else fx.min(fx.Int32(min(TILE, per - j0)), s - plus(g0, j0)), + ) + for j0 in range(0, per, TILE) + ] + wvs = router_weight_loads(a["gate_w"], tile, part, lane, wave, CM_NT) + xs = [router_x_loads(a["x"], part, n, lane, wave, c["x_cm"], t0) for t0, n in tiles] + for (t0, n), xvs in zip(tiles, xs): + router_tile_mfma(wvs, xvs, lane, wave, red) + gpu.barrier() + if tid < ROWS * n: + r = tid % ROWS + tl = tid // ROWS + v = row_sum(red, r, tl) + c["put"]( + c["logit"], + (plus(t0, tl) * experts + tile * ROWS + r) * ROUTER_PARTS + part, + v, + ) + gpu.barrier() + + +# ---------------------------------------------------------------- route +def _score(x): + """sqrt(log2(1 + exp2(x log2e)) ln2), x itself past 20 (topk_gating).""" + sp = fx.Float32( + (x > fx.Float32(20.0)).select( + x, + _hw_f32( + "llvm.amdgcn.log.f32", fx.Float32(1.0) + hw_exp2(x * fx.Float32(LOG2E)) + ) + * fx.Float32(LN2), + ) + ) + return fx.Float32(fmath.sqrt(sp)) + + +@traced +def stage_route(c, t): + """VLLM's topk_hash_softplus_sqrt for token t on wave 0 (f32 logits).""" + lane, wave = c["lane"], c["wave"] + experts, topk = c["rs"].experts, c["rs"].topk + a = c["args"] + if wave == 0: + ept = experts // 64 + # each logit's router part sums, added in part order (vLLM's gate: f32 out) + assert ROUTER_PARTS == 4 + got = c["poll"]( + [ + (c["logit"], (t * experts + lane + 64 * i) * ROUTER_PARTS + h, 2) + for i in range(ept) + for h in (0, 2) + ] + ) + logits = [] + for i in range_constexpr(ept): + parts = got[2 * i] + got[2 * i + 1] + tot = fx.Float32(0.0) + for p in range_constexpr(ROUTER_PARTS): + tot = tot + parts[p].bitcast(fx.Float32) + logits.append(tot) + vals, orig, idxs = [], [], [] + for i in range_constexpr(ept): + e = lane + 64 * i + sc = _score(logits[i]) + # aiter's NaN guards: a NaN score weighs 0 and never wins + orig.append(fx.Float32(fx.isnan(sc).select(fx.Float32(0.0), sc))) + vals.append(fx.max(sc + ld_f32(a["bias"], e), fx.Float32(float("-inf")))) + idxs.append(fx.Int32(e)) + for p, q in c["rs"].sort_net: + sw = vals[p] < vals[q] + vals[p], vals[q] = sw.select(vals[q], vals[p]), sw.select(vals[p], vals[q]) + orig[p], orig[q] = sw.select(orig[q], orig[p]), sw.select(orig[p], orig[q]) + idxs[p], idxs[q] = sw.select(idxs[q], idxs[p]), sw.select(idxs[p], idxs[q]) + cursor = fx.Int32(0) + total = fx.Float32(0.0) + my_w, my_id = fx.Float32(0.0), fx.Int32(0) + for k in range_constexpr(topk): + cv, ci, co = fx.Float32(float("-inf")), fx.Int32(0), fx.Float32(0.0) + for i in range_constexpr(ept): + at = cursor == i + cv, ci, co = ( + at.select(vals[i], cv), + at.select(idxs[i], ci), + at.select(orig[i], co), + ) + best = butterfly(cv, (32, 16, 8, 4, 2, 1), fx.max) + ball = ballot(cv == best) + win = fx.Int32( + (ball == fx.Int64(0)).select(fx.Int32(0), fx.Int32(fmath.cttz(ball))) + ) + win_id = fx.Int32(rocdl.readlane(T.i32, ci.ir_value(), win.ir_value())) + won = (cursor < ept) & (ci == win_id) + mine_o = fx.Float32(won.select(co, fx.Float32(0.0))) + w = fx.Int32( + rocdl.readlane( + T.i32, mine_o.bitcast(fx.Int32).ir_value(), (win_id & 63).ir_value() + ) + ).bitcast(fx.Float32) + cursor = cursor + won.select(1, 0) + total = total + w + my_w = (lane == k).select(w, my_w) + my_id = (lane == k).select(win_id, my_id) + scale = fx.Float32(SCALE) / fx.max(total, fx.Float32(1e-20)) + if lane < topk: + for r in range_constexpr(ROUTE_COPIES): + at = (r * c["S"] + t) * 2 * topk + lane + c["put"](c["route"], at, my_id) + c["put"](c["route"], at + topk, my_w * scale) + gpu.barrier() + + +@traced +def load_route(c, tab): + """Every token's (ids, weights) into LDS ``tab``, then the union U of the ids + (ascending): tab[2 S topk ..] = |U|, the experts of U, each (t, k)'s slot in + U, each expert's slot, the picks in slot order (a slot's ascending) and each + slot's first position (``RouteShape``). + + A pick a thread, no scan in series: each pick ORs its token's bit into its + expert's token bitmap, so an expert's pick count is its bitmap's popcount, + a slot's first position the counts of the experts below it (a wave's prefix + sum; slots follow the experts ascending), and a pick's rank in its slot the + bitmap's tokens below its own. The bitmaps are the same whatever order the + ORs land in, so every CTA derives the same pick order.""" + s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] + rs = c["rs"] + topk, max_u, experts = rs.topk, rs.max_u, rs.experts + n = s * topk + assert s <= 32 * MASK_WORDS + # a thread's route words in one poll batch: a pass each was a round trip + # each in series past 256 picks (48 rows: 5 us from the last route to the + # table); a word past 2n re-reads the last and is not stored + idxs = [plus(pas, tid) for pas in range(0, 2 * n, THREADS)] + copy_at = c["bid"] % ROUTE_COPIES * (2 * n) + got = c["poll"]([(c["route"], copy_at + fx.min(idx, 2 * n - 1), 1) for idx in idxs]) + for i in range_constexpr(len(idxs)): + if idxs[i] < 2 * n: + fx.ptr_store(got[i][0], tab + idxs[i]) + for pas in range_constexpr(0, MASK_WORDS * experts, THREADS): + idx = plus(pas, tid) + if idx < MASK_WORDS * experts: + fx.ptr_store(fx.Int32(0), tab + (rs.mask0 + idx)) + gpu.barrier() + if tid < n: + t = tid // topk + eid = fx.ptr_load(tab + (t * 2 * topk + tid % topk)) + fx.llvm.atomic_or( + tab + (rs.mask0 + eid * MASK_WORDS + t // 32), + fx.Int32(1) << (t % 32), + syncscope="workgroup", + ) + gpu.barrier() + base = 2 * n + # a chunk of 64 experts a wave (wave 0 alone through 6 chunks in series + # was 3.3 us of the table's 5 at 48 rows): each chunk's own counts and + # prefixes, its totals through LDS ``red`` (free here), then each chunk + # past the totals of the chunks below it -- the serial walk's values + chunks = experts // 64 + assert chunks <= WAVES + sums = c["red"] + e = lane + 64 * fx.min(wave, chunks - 1) + count = fx.Int32(0) + for w in range_constexpr(MASK_WORDS): + word = fx.ptr_load(tab + (rs.mask0 + e * MASK_WORDS + w)) + count = count + popcount(word) + ball = ballot(count != 0) + # this chunk's inclusive prefixes, lanes ascending: picks, and pick tiles + incl = lane_prefix(count, lane) + tiles = mplan.pick_tiles(count) + tincl = lane_prefix(tiles, lane) if rs.multi else tiles + totals = [popcount(ball), lane_gather(incl, 63), lane_gather(tincl, 63)] + if (wave < chunks) & (lane == 0): + for k in range_constexpr(3): + fx.ptr_store(totals[k].bitcast(fx.Float32), sums + (wave * 3 + k)) + gpu.barrier() + below_c = [fx.Int32(0), fx.Int32(0), fx.Int32(0)] # n_u, first, tile0 + for j in range_constexpr(chunks - 1): + for k in range_constexpr(3): + v = fx.ptr_load(sums + (j * 3 + k)).bitcast(fx.Int32) + below_c[k] = below_c[k] + (j < wave).select(v, fx.Int32(0)) + n_u, first, tile0 = below_c + slot = n_u + lanes_below(ball, lane) + # e afresh: hoisted, its address held a VGPR across the prefix sums (spill) + e_at = fresh(e) + if wave < chunks: + fx.ptr_store(slot, tab + (rs.expert_slot0 + e_at)) + if count != 0: + fx.ptr_store(fx.Int32(e), tab + (base + 1 + slot)) + fx.ptr_store(first + incl - count, tab + (rs.first0 + slot)) + if const_expr(rs.multi): + # the slot's pick tiles: their first index, and each one's slot + tf = tile0 + tincl - tiles + if count != 0: + fx.ptr_store(tf, tab + (rs.tile_first0 + slot)) + for j in range_constexpr(rs.slot_tiles): + if j < tiles: + fx.ptr_store(slot, tab + (rs.tile_slot0 + tf + j)) + if (wave == chunks - 1) & (lane == 0): + n_all = n_u + totals[0] + fx.ptr_store(n_all, tab + base) + fx.ptr_store(fx.Int32(n), tab + (rs.first0 + n_all)) + if const_expr(rs.multi): + fx.ptr_store(tile0 + totals[2], tab + (rs.tile_first0 + n_all)) + gpu.barrier() + if tid < n: + t = tid // topk + eid = fx.ptr_load(tab + (t * 2 * topk + tid % topk)) + slot = fx.ptr_load(tab + (rs.expert_slot0 + eid)) + fx.ptr_store(slot, tab + (base + 1 + max_u + tid)) + # the tokens below t that picked this expert: the picks before tid in it + rank = fx.Int32(0) + for w in range_constexpr(MASK_WORDS): + word = fx.ptr_load(tab + (rs.mask0 + eid * MASK_WORDS + w)) + lo = t - 32 * w + keep = (lo >= 32).select( + fx.Int32(-1), + (lo > 0).select((fx.Int32(1) << fx.max(lo, 0)) - 1, fx.Int32(0)), + ) + rank = rank + popcount(word & keep) + pos = fx.ptr_load(tab + (rs.first0 + slot)) + rank + fx.ptr_store(tid, tab + (rs.order0 + pos)) + gpu.barrier() + + +def lane_prefix(v, lane): + """The wave's inclusive prefix sum of v, lanes ascending.""" + incl = v + for d in (1, 2, 4, 8, 16, 32): + moved = lane_gather(incl, fx.max(lane - d, 0)) + incl = incl + (lane >= d).select(moved, fx.Int32(0)) + return incl + + +# ---------------------------------------------------------------- shared ug +def _silu_mul(g, u): + """VLLM's silu_and_mul_with_clamp (the shared expert's): gate min limit, + up clamped; silu = g / (1 + expf(-g)) rounded to bf16, then times up.""" + g = bf16_round(fx.min(g, fx.Float32(LIMIT))) + u = fx.max(fx.min(u, fx.Float32(LIMIT)), fx.Float32(-LIMIT)) + act = bf16_round(g / (fx.Float32(1.0) + fx.Float32(fmath.exp(-g)))) + return act * u + + +@traced +def stage_shared(c, task, gops, t0, n): + """Gate rows 16 task .. (weights ``gops``), up rows 576 + 16 task .. against + tokens t0 .. t0 + n (their x in LDS; ``n`` traced: a tile a CTA picks at run + time) -> bf16 -> silu x up -> SMID; an odd task then quantizes its group of + 32 (``smid_quant``).""" + tid, red, d = c["tid"], c["red"], c["d"] + a = c["args"] + rows = None if isinstance(n, int) else n + gemv_fp8_mfmas(c, HIDDEN, c["xl"], c["xsl"], red, gops, rows=rows) + gpu.barrier() + gate = fx.Float32(0.0) + if tid < ROWS * n: + gate = row_sum(red, tid % ROWS, tid // ROWS) + gpu.barrier() + uops = gemv_fp8_loads( + c, a["sgu"], a["sgu_s"], HIDDEN, d.shared_tasks + task, row_major=True + ) + gemv_fp8_mfmas(c, HIDDEN, c["xl"], c["xsl"], red, uops, rows=rows) + gpu.barrier() + if tid < ROWS * n: + r = tid % ROWS + t = plus(t0, tid // ROWS) + up = row_sum(red, r, tid // ROWS) + v = fx.Float32(_silu_mul(bf16_round(gate), bf16_round(up)).to(fx.BFloat16)) + c["put"](c["smid"], t * d.sh_inter + task * ROWS + r, v) + fx.ptr_store(v, c["pair"] + tid) + gpu.barrier() + if task % 2 == 1: + smid_quant(c, task, t0, n) + gpu.barrier() + + +@traced +def smid_quant(c, task, t0, n): + """Odd shared task ``task``: its 16 columns (LDS ``pair``) and the even + task's (SMID) are a quantize_fp8 group of 32 (amax floored 1e-4, the ceil + code) -> SMQ / SMC, a thread a token. Once here, not in each of the 256 CTAs + whose down stage reads them: that took a down CTA 14 us at 48 rows.""" + tid, sh = c["tid"], c["d"].sh_inter + assert c["d"].shared_tasks % 2 == 0 # tasks pair into groups + if tid < n: + t = plus(t0, tid) + got = c["poll"]( + [(c["smid"], t * sh + (task - 1) * ROWS + r, 1) for r in range(ROWS)] + ) + vals = [got[r][0].bitcast(fx.Float32) for r in range(ROWS)] + [ + fx.ptr_load(c["pair"] + (tid * ROWS + r)) for r in range(ROWS) + ] + amax = abs(vals[0]) + for y in vals[1:]: + amax = fx.max(amax, abs(y)) + # vLLM's down_proj input quant (mxfp8_e4m3_quantize) + code = vllm_mx_code(amax) + inv = ((254 - code) << 23).bitcast(fx.Float32) + g = task // 2 + for w in range_constexpr(8): + word = fp8_pack4(*[clamp_fp8(vals[4 * w + i] * inv) for i in range(4)]) + c["put_words"](c["smq"], t * (sh // 4) + g * 8 + w, [word]) + c["put_words"](c["smc"], t * (sh // 32) + g, [code]) + + +# ---------------------------------------------------------------- routed ug +@traced +def gemv_fp4_pair(c, w, ws, e, n_rows, k, rg0, ks, xl, xsl, scale_cols, col): + """Tiles rg0 and rg0 + 1 (16 rows each, one 32-row scale block) of expert + e's (16, 16)-preshuffled MXFP4 weight [n_rows, k] times 16 MXFP8 rows in + LDS, over this wave's K slice ``ks`` (``ug_rounds``) -> both tiles' partial + C, a Vector of 8 (rows 4 (lane / 16) .., column lane % 16, whose LDS row is + ``col``). Each B operand read from LDS feeds both tiles' MFMAs: the waves of + a task share x, so a wave a tile over all of K read it once a tile. + """ + lane, wave = c["lane"], c["wave"] + sl = c["wsl"] + wave * wsl_words(c["d"]) + slices = c["rs"].ug_k_slices + rounds = ug_rounds(k, ks, slices) + accs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(2)] + # rounds whose loads are in flight past the one computing: a quarter's + # 2-step rounds two (the draft's 20 rows: -0.5 us) + ahead = slices // 2 + ops = [ + fp4_pair_loads(c, w, ws, e, n_rows, k, rg0, scale_cols, *rounds[q]) + for q in range(min(ahead, len(rounds))) + ] + for r in range_constexpr(len(rounds)): + sc0, st0, n, sc_steps = rounds[r] + sv, avs = ops[r] + fx.ptr_store(sv, sl + lane * (sc_steps // 2)) + sas = [ + [ + lds_u8( + sl, + scale_index( + t * 16 + lane % 16, (st0 + i) * 4 + lane // 16, scale_cols + ) + - 128 * sc0, + ) + for i in range(n) + ] + for t in range(2) + ] + accs = fp4_pair_mfma(c, k, xl, xsl, st0, n, avs, sas, accs, col) + if const_expr(r + ahead < len(rounds)): + ops.append( + fp4_pair_loads( + c, w, ws, e, n_rows, k, rg0, scale_cols, *rounds[r + ahead] + ) + ) + return fx.Vector.from_elements( + [accs[t][q] for t in range(2) for q in range(4)], fx.Float32 + ) + + +def ug_rounds(k, ks, slices): + """K slice ``ks`` of ``slices``' rounds (first step of the scales loaded, + first step, steps, steps of scales loaded); ``ks`` traced: the steps are. + Halves: the 8-step scale blocks 2 r + ks, the odd block left split 4 + 4, + so both run k / 256 steps. Quarters: steps 8 i + 2 ks, + 1 of each block + i, a round loading only their scales (the shuffled layout keeps 2 steps' + together: a dword a lane, where the whole block was 4).""" + blocks = k // 128 // 8 + if slices == 4: + return [(fx.Int32(8 * i) + ks * 2,) * 2 + (2, 2) for i in range(blocks)] + assert slices == 2 + rounds = [ + (fx.Int32(16 * r) + ks * 8, fx.Int32(16 * r) + ks * 8, 8, 8) + for r in range(blocks // 2) + ] + if blocks % 2: + last = (blocks - 1) * 8 + rounds.append((fx.Int32(last), fx.Int32(last) + ks * 4, 4, 8)) + return rounds + + +def fp4_pair_loads(c, w, ws, e, n_rows, k, rg0, scale_cols, sc0, st0, n, sc_steps): + """``gemv_fp4_pair``'s loads of a round: the scales of steps sc0 .. sc0 + + sc_steps (of the tiles' 32-row block: 128 bytes a step, sc_steps / 2 dwords + a lane) and both tiles' weight operands of steps st0 .. st0 + n. Loads + only.""" + lane = c["lane"] + block = (e * n_rows + rg0 * 16) // 32 * (scale_cols * 8) + words = sc_steps // 2 + sv = bo.buffer_load( + rsrc(ws), + block + 32 * sc0 + lane * words, + vec_width=words, + dtype=T.i32, + cache_modifier=CM_NT, + ) + sv = fx.Vector(sv) if words > 1 else fx.Int32(sv) + return sv, [fp4_tile_weights(c, w, e, n_rows, k, rg0 + t, st0, n) for t in range(2)] + + +def scale_block_loads(c, ws, first_row, scale_cols): + """The loads of the 32-row e8m0 block holding ``first_row``: in aiter's + shuffle_scale order a block is contiguous (both 16-row halves, all of K), so + a few 16-byte loads a lane fetch it, where a byte load a K step doubled the + stage's load instructions and bound it by their issue (2026-09-29, ug). A + lane past the block re-reads its last 16 bytes.""" + lane = c["lane"] + words = scale_cols * 8 + block = first_row // 32 * words + return [ + fx.Vector( + bo.buffer_load( + rsrc(ws), + block + fx.min(256 * i + lane * 4, words - 4), + vec_width=4, + dtype=T.i32, + cache_modifier=CM_NT, + ) + ) + for i in range((words + 255) // 256) + ] + + +def down_scale_words(d) -> int: + """A down tile's 32-row scale block in LDS: ``scale_block_store`` writes + whole 256-word rounds.""" + return down_lds_words(d.down_scale_cols * 8) + + +def wsl_words(d) -> int: + """A wave's LDS scale words: an ug batch's 256, or a down round's blocks -- + a (task, slot) each, ``DOWN_R_MAX`` at most (``stage_down``).""" + return max(256, DOWN_R_MAX * down_scale_words(d)) + + +def scale_block_store(c, svs, sl): + """``scale_block_loads``' words -> this wave's LDS block ``sl`` + (``down_scale_words``; a lane past the block writes past it, inside it).""" + lane = c["lane"] + for i in range_constexpr(len(svs)): + fx.ptr_store(svs[i], sl + (256 * i + lane * 4)) + + +def lds_scales(sl, local_row, lane, st0, blk, scale_cols): + """K steps st0 .. st0 + blk's scale bytes of this lane's row from a 32-row + scale block in LDS (``local_row`` its row in the block). An even step and + the next are bytes 2 apart of one word (``scale_index``'s d): a read a + pair, not a step (down: 3 reads a slot, not 5).""" + out, word = [], None + for i in range(blk): + idx = scale_index(local_row, (st0 + i) * 4 + lane // 16, scale_cols) + if (st0 + i) % 2 == 0 or word is None: + word = fx.ptr_load(sl + idx // 4) + out.append((word >> (idx % 4 * 8)) & 0xFF) + return out + + +def lds_u8(ptr, i): + return (fx.ptr_load(ptr + i // 4) >> (i % 4 * 8)) & 0xFF + + +def fp4_tile_weights(c, w, e, n_rows, k, rg, st0, blk, k_real=None): + """The weight operands of expert e's 16-row tile rg, K steps st0 .. st0 + + blk: loads only, nothing waits on them here. ``k_real``: a lane whose 32 + columns lie past it (zero weights) reads its row group's first lane's + address instead (no extra traffic) and holds zero.""" + lane = c["lane"] + base = e * n_rows * (k // 2) // 4 # dwords: expert e + out = [] + for i in range_constexpr(blk): + past = k_real is not None and ((st0 + i) * 128 + 96) >= k_real + if const_expr(past): + real = (st0 + i) * 128 + lane // 16 * 32 < k_real + at = real.select(lane * 4, lane % 16 * 4) + else: + at = lane * 4 + v = fp4_tile_load(rsrc(w), base, rg, k, st0 + i, at) + if const_expr(past): + v = fx.Vector.from_elements( + [real.select(v[q], fx.Int32(0)) for q in range(4)], fx.Int32 + ) + out.append(v) + return out + + +def fp4_pair_mfma(c, k, xl, xsl, st0, n, avs, sas, accs, col): + """``accs`` (two tiles') chained through ``n`` K steps from st0 against MXFP8 + rows in LDS (column lane % 16: row ``col``): a step's B operand loaded once + for both tiles' MFMAs.""" + lane = c["lane"] + x_row = lds_row(k // 4) + s_row = lds_row(k // 32) + j = lane // 16 + accs = list(accs) + for i in range_constexpr(n): + bv = fp8_operand( + [ + fx.ptr_load( + xl + (col * x_row + (st0 + i) * 32 + 16 * h + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + for h in range(2) + ] + ) + sb = fx.ptr_load(xsl + (col * s_row + (st0 + i) * 4 + j)) + for t in range_constexpr(2): + accs[t] = mfma_scaled(avs[t][i], bv, accs[t], sas[t][i], sb, FP4, FP8) + return accs + + +def fp8_operand(halves): + """A lane's MXFP8 MFMA operand of a K step: its two 16-byte halves (K + 16 j .. and 64 + 16 j .. of the step, j = lane / 16).""" + return fx.Vector.from_elements( + [fx.Vector(h).bitcast(fx.Int32)[d] for h in halves for d in range(4)], + fx.Int32, + ) + + +def down_lds_words(n): + """A down wave's LDS block of n words: ``scale_block_store`` writes whole + 256-word rounds.""" + return -(-n // 256) * 256 + + +def st_dev(v, region, i): + """Word i of scratch ``region`` (plain words, not tagged pairs) at device + scope: past the per-XCD L2, so another CTA reads it after the flag.""" + bo.buffer_store(v, rsrc(region.value), i, cache_modifier=CM_DEV) + + +@traced +def take_ug_task(c): + """This CTA's next ug task past the first round, from the launch's queue + (``ugq``, a word a tag, counting from 0): the CTAs that ran a shared, + router or route task come to it late and so take fewer, where a fixed + placement gave every CTA the same count and down waited on theirs. A + task's result is the same whichever CTA runs it.""" + if c["tid"] == 0: + at = c["ugq"].value + fx.Int64(c["tag"] * 4) + ptr = fx.inttoptr(fx.PointerType.get(T.i32, fx.AddressSpace.Global, 4), at) + task = fx.llvm.atomic_add(ptr, fx.Int32(1), syncscope="agent") + fx.ptr_store(fx.Int32(task), c["tab"] + c["rs"].queue0) + gpu.barrier() + return fx.ptr_load(c["tab"] + c["rs"].queue0) + + +@traced +def stage_ug(c, task): + """One token tile: slot task // ug_parts of U (``RouteShape.ug_parts``), the + ``ug_groups`` FP4 groups of intermediate columns from 32 ug_groups (task % + ug_parts) (``ug_wave_pair``); then silu x up, rounded to bf16, and the + MXFP8 requant of each token's 32 values of a group -> MID (the stage-1 + epilogue and the stage-2 input quant). Past one: ``stage_ug_tile``.""" + s, tid, lane, tab = c["S"], c["tid"], c["lane"], c["tab"] + rs = c["rs"] + ug_parts = rs.ug_parts + u = task // ug_parts + if u < c["rs"].n_union(tab, s): + e = c["rs"].union_expert(tab, s, u) + # the slot's picks, a column each (a column past them repeats the last); + # every token's row is in LDS: a column reads its pick's token's + first = c["rs"].slot_first(tab, u) + count = c["rs"].slot_first(tab, u + 1) - first + pick = c["rs"].pick_at(tab, first + fx.min(lane % 16, count - 1)) + ug_pass(c, e, task % ug_parts * rs.ug_groups, first, count, pick // rs.topk) + publish(c["put"], c["ugf"], task, 1, tid == 0) + gpu.barrier() + + +def ug_wave_pair(c, part0): + """This wave's share of an ug task from group ``part0``: group p of its + ``ug_groups`` (wave / (2 slices)), its 16-column half bb = wave % 2 -- w13's + gate tile 2 b and up tile 2 b + 1 (aiter GUGU interleaves them a 16-row + tile; b = 2 (part0 + p) + bb), one 32-row scale block -- over K slice ks = + wave / 2 % slices (``gemv_fp4_pair``): (p, 2 b, ks, whether the group is in + the real width).""" + wave, slices = c["wave"], c["rs"].ug_k_slices + p = wave // (2 * slices) + b = (part0 + p) * 2 + wave % 2 + return p, 2 * b, wave // 2 % slices, part0 + p < c["d"].inter_real // UG_PART + + +def ug_red_slot(rs, p, bb, ks, t): + """RED's 64-lane slot of an ug task's partial C: group p, half bb, K slice + ks, gate (t 0) or up (t 1).""" + return ((p * 2 + bb) * rs.ug_k_slices + ks) * 2 + t + + +def pairwise_sum(vs): + """Vs summed as a balanced tree, neighbours first: (v0 + v1) + (v2 + v3).""" + while len(vs) > 1: + vs = [vs[i] + vs[i + 1] for i in range(0, len(vs), 2)] + return vs[0] + + +def ug_tile_task(c, task): + """Past one token tile, ug task ``task``: its slot's expert, its first + group, and its pick tile's first pick and the picks from it.""" + tab, rs, s = c["tab"], c["rs"], c["S"] + ug_parts = rs.ug_parts + tile = task // ug_parts + u = rs.tile_slot(tab, tile) + first = mplan.tile_pick0(rs.slot_first(tab, u), tile, rs.tile_first(tab, u)) + count = rs.slot_first(tab, u + 1) - first + return rs.union_expert(tab, s, u), task % ug_parts * rs.ug_groups, first, count + + +@traced +def stage_ug_tile(c, task, total): + """Past one token tile: ug task ``task`` of ``total`` (a pick tile of its + slot), then the next task from the queue (``take_ug_task``). No next task's + weights prefetched under this one: held across its epilogue and gather they + spilled the merged K2 (S = 24 / 48), and without them it ran faster.""" + tid, lane = c["tid"], c["lane"] + e, part0, first, count = ug_tile_task(c, task) + gather_x8r(c, first, count) + ug_pass(c, e, part0, first, count, fx.min(lane % 16, count - 1)) + nxt = (total > BLOCKS).select(take_ug_task(c) + BLOCKS, total) + # the epilogue's MID stores done before the flag + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if tid == 0: + c["put"](c["ugf"], task, 1) + return nxt + + +@traced +def ug_pass(c, e, part0, first, count, col): + """An ug task's GEMV and epilogue for the slot's picks first .. first + + min(count, 16) (column lane % 16's x in LDS row ``col``).""" + tid, lane, wave, red = c["tid"], c["lane"], c["wave"], c["red"] + a, d, rs = c["args"], c["d"], c["rs"] + p, rg0, ks, real = ug_wave_pair(c, part0) + # a group past the real width (the padding to inter) has zero weights: + # its GEMV is zero unread, and the epilogue still writes its zero MID + acc = fx.Vector.filled(8, 0.0, fx.Float32) + if real: + acc = gemv_fp4_pair( + c, + a["w13"], + a["w13_s"], + e, + 2 * d.inter, + HIDDEN, + rg0, + ks, + c["xl"], + c["xsl"], + HIDDEN // 32, + col, + ) + for t in range_constexpr(2): + slot = ug_red_slot(rs, p, wave % 2, ks, t) + fx.ptr_store( + fx.Vector.from_elements([acc[4 * t + q] for q in range(4)], fx.Float32), + red + (slot * 64 + lane) * 4, + ) + gpu.barrier() + # thread (p, j, i): group p, the slot's pick j, intermediate column + # 32 part + i, a pass of THREADS at a time; a group's 32 threads stay + # lane-aligned + cols = rs.tile_picks + for pas in range_constexpr(0, rs.ug_groups * cols * UG_PART, THREADS): + ug_epilogue(c, part0, pas + tid, first, count) + + +@traced +def load_x8_rows(c, words_mb, codes_mb, xl, xsl, token_of, live): + """MXFP8 rows ``token_of(j)``, j < ``live`` (<= TILE; traced or int), of + scratch ``words_mb`` / ``codes_mb`` into LDS rows j of ``xl`` / ``xsl``: + the first ``live`` lanes wait on their token's X8RDY, then plain 16-byte + loads (half the bytes of the tagged pairs a poll would read).""" + tid = c["tid"] + if tid < TILE: + c["poll"]([(c["x8rdy"], token_of(fx.min(tid, live - 1)), 1)]) + gpu.barrier() + for mb, width, dst in ((words_mb, X8_WORDS, xl), (codes_mb, HIDDEN // 32, xsl)): + quads = width // 4 + for i in range_constexpr(-(-TILE * quads // THREADS)): + q = tid + THREADS * i + row = q // quads + if (q < TILE * quads) & (row < live): + v = fx.Vector( + bo.buffer_load( + rsrc(c[mb].value), + token_of(row) * width + q % quads * 4, + vec_width=4, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ) + fx.ptr_store(v, dst + (row * lds_row(width) + q % quads * 4)) + gpu.barrier() + + +def gather_x8r(c, first, count): + """The X8R rows and codes of the slot's picks first .. first + min(count, 16) + into LDS rows 0 .. (only those rows: an ug pass reads a pick's row, the + columns past them repeat the last one's).""" + tab, topk = c["tab"], c["rs"].topk + + def token(j): + return c["rs"].pick_at(tab, first + fx.min(j, count - 1)) // topk + + load_x8_rows(c, "x8r", "x8rs", c["x8rl"], c["x8rsl"], token, fx.min(count, TILE)) + + +@traced +def ug_epilogue(c, part0, idx, first, count): + """Thread ``idx``'s value of an ug task: silu x up of its (group, pick j, + column), the group's MXFP8 scale, its byte of MID (the pick's row, ``first`` + + j in the pick order).""" + lane, red, d, rs = c["lane"], c["red"], c["d"], c["rs"] + cols = rs.tile_picks + if idx < rs.ug_groups * cols * UG_PART: + p = idx // (cols * UG_PART) + part = part0 + p + t = idx % (cols * UG_PART) // UG_PART + i = idx % UG_PART + tile = i // 16 + r = i % 16 + at = (16 * (r // 4) + t) * 4 + r % 4 + g, up = [ + pairwise_sum( + [ + fx.ptr_load(red + ug_red_slot(rs, p, tile, ks, t) * 64 * 4 + at) + for ks in range(rs.ug_k_slices) + ] + ) + for t in range(2) + ] + gc = fx.min(g, fx.Float32(LIMIT)) + uc = fx.max(fx.min(up, fx.Float32(LIMIT)), fx.Float32(-LIMIT)) + sig = hw_rcp(fx.Float32(1.0) + hw_exp2(gc * fx.Float32(-LOG2E))) + # the stage-1 kernel stores the intermediate as bf16; the stage-2 input + # quant (fused_mx_quant_moe_sort) makes it MXFP8: amax floored 1e-10, + # the ceil code, x / scale clamped, RNE + v = bf16_round(gc * sig * uc) + amax = fx.max(butterfly(abs(v), (1, 2, 4, 8, 16), fx.max), fx.Float32(1e-10)) + code = code_ceil(amax * fx.Float32(1.0 / FP8_MAX)) + q8 = clamp_fp8(v * rcp_pow2(code)) + nb = [ + lane_gather(q8.bitcast(fx.Int32), fx.min(lane + q, 63)).bitcast(fx.Float32) + for q in range(1, 4) + ] + # 4 threads' bytes -> one word + word = fp8_pack4(q8, *nb) + row = first + t + if (i % 4 == 0) & (t < count): + st_dev(word, c["mid"], row * d.mid_words + part * 8 + i // 4) + if (i == 0) & (t < count): + st_dev(code, c["mids"], mid_code_index(d, row, part)) + + +# ---------------------------------------------------------------- down +@traced +def stage_down(c, tasks): + """Rows 16 task .. of the hidden output for each of ``tasks`` (a CTA's one or + two, run together: a second pass would double the CTA's time): routed + experts of U (a wave each, round robin), the shared expert, their sum -> + pushed to every rank.""" + s, wave, tab, d = c["S"], c["wave"], c["tab"], c["d"] + n = len(tasks) + nu = c["rs"].n_union(tab, s) + sl = c["wsl"] + wave * wsl_words(d) + sw = down_scale_words(d) + # slots of U a round: the loads of a round stay the same + slots = DOWN_R_MAX // n + per_round = WAVES * slots + # the shared expert's down projection first: its input is ready long + # before the routed results, which a CTA mostly waits on here; after the + # rounds it sat on the path past the last ug task + down_shared(c, tasks) + for rnd in range(0, (nu + per_round - 1) // per_round, 1): + # the round's slots of U (wave, wave + WAVES, ...), every load of them + # and of every task in flight before the first wait: the weights and + # their scale blocks (no ug result: before the round's ug flags), then + # the slots' picks' intermediate rows, each lane its pick's straight + # from MID (``mid_operands``). A slot past U loads the last slot again + # (L2 hits) and contributes nothing + us = [wave + WAVES * (slots * rnd + i) for i in range(slots)] + ucs = [fx.min(u, nu - 1) for u in us] + firsts = [c["rs"].slot_first(tab, u) for u in ucs] + svs, avs = down_weights(c, tasks, rnd, slots, nu) + # this round's slots' ug tasks done, not every slot of the wave: the + # round starts while later slots' ug run + c["poll"]([(c["ugf"], f, 1) for f in down_ug_flags(c, ucs)]) + counts = [c["rs"].slot_first(tab, u + 1) - f for u, f in zip(ucs, firsts)] + mids = [mid_operands(c, f, k) for f, k in zip(firsts, counts)] + for q in range_constexpr(len(svs)): + scale_block_store(c, svs[q], sl + q * sw) + for i in range_constexpr(slots): + for j in range_constexpr(n): + down_mfma( + c, + j, + tasks[j], + sl + (j * slots + i) * sw, + mids[i], + firsts[i], + counts[i], + avs[j * slots + i], + us[i] < nu, + ) + # a slot of more than 16 picks: its next pick tile's rows + if const_expr(s > TILE): + for j0 in range_constexpr(TILE, s, TILE): + if j0 < counts[i]: + more = mid_operands(c, firsts[i] + j0, counts[i] - j0) + for j in range_constexpr(n): + down_mfma( + c, + j, + tasks[j], + sl + (j * slots + i) * sw, + more, + firsts[i] + j0, + counts[i] - j0, + avs[j * slots + i], + us[i] < nu, + ) + gpu.barrier() + for t0, nn in token_tiles(s): + for j in range_constexpr(n): + down_combine(c, tasks[j], j, t0, nn) + + +@traced +def take_down_rank(c): + """This CTA's arrival rank at the down stage (DQ, a word a tag): which + CTAs take the extra tiles. A tile's result is the same whichever CTA + computes it.""" + if c["tid"] == 0: + at = c["dq"].value + fx.Int64(c["tag"] * 4) + ptr = fx.inttoptr(fx.PointerType.get(T.i32, fx.AddressSpace.Global, 4), at) + rank = fx.llvm.atomic_add(ptr, fx.Int32(1), syncscope="agent") + fx.ptr_store(fx.Int32(rank), c["tab"] + c["rs"].queue0) + gpu.barrier() + return fx.ptr_load(c["tab"] + c["rs"].queue0) + + +def down_weights(c, tasks, rnd, slots, nu): + """Round ``rnd``'s weights of ``stage_down``: each (task, slot)'s scale block + and weight operands, this wave's slots of U (a slot past U: the last's). + Loads only.""" + s, wave, tab, a, d = c["S"], c["wave"], c["tab"], c["args"], c["d"] + us = [wave + WAVES * (slots * rnd + i) for i in range(slots)] + es = [c["rs"].union_expert(tab, s, fx.min(u, nu - 1)) for u in us] + svs = [ + scale_block_loads(c, a["w2_s"], e * HIDDEN + task * 16, d.down_scale_cols) + for task in tasks + for e in es + ] + avs = [ + fp4_tile_weights( + c, a["w2"], e, HIDDEN, d.inter, task, 0, d.inter // 128, d.inter_real + ) + for task in tasks + for e in es + ] + return svs, avs + + +def mid_operands(c, first, count): + """The MFMA B operands of picks first .. first + min(count, 16) (column + lane % 16 the pick, past them the last): each K step's two 16-byte halves + of the lane's MID row and its group's code, loaded at device scope straight into + registers -- no LDS block a slot, so a round holds more slots.""" + lane, d = c["lane"], c["d"] + j = lane // 16 + row = first + fx.min(lane % 16, count - 1) + steps = d.inter // 128 + # the step's constant after an opaque lane word folds into the immediate + # offset; summed with j * 4 it was a hoisted VGPR a step (spill) + lane_word = fresh(row * d.mid_words + j * 4) + bvs = [] + for st in range_constexpr(steps): + halves = [] + for h in range_constexpr(2): + v = fx.Vector( + bo.buffer_load( + rsrc(c["mid"].value), + lane_word + (st * 32 + 16 * h), + vec_width=4, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ) + if const_expr(st * 128 + 64 * h + 48 >= d.inter_real): + # columns past the rank's real intermediate (the loader pads + # it with zero weights) are never written: stale bytes, which + # a NaN pattern would carry through the zero weights + past = st * 128 + 64 * h + 16 * j >= d.inter_real + v = fx.Vector.from_elements( + [past.select(fx.Int32(0), v[q]) for q in range(4)], fx.Int32 + ) + halves.append(v) + bvs.append(fp8_operand(halves)) + # this lane's codes, one a K step: contiguous in MIDS (``mid_code_index``), + # 4 a load -- a load a step was 5 of a slot's 16 load instructions, and the + # stage is load-issue bound (probe ``harness/down_probe.py``: -8 %) + base = mid_code_index(d, row, j) + sbs = [] + for w0 in range_constexpr(0, steps, 4): + n = min(4, steps - w0) + v = bo.buffer_load( + rsrc(c["mids"].value), + base + w0, + vec_width=n, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + if const_expr(n == 1): + sbs.append(fx.Int32(v)) + else: + v = fx.Vector(v) + sbs += [fx.Int32(v[k]) for k in range(n)] + for st in range_constexpr(steps): + if const_expr((st + 1) * 128 > d.inter_real): + # a group past the real intermediate has no code written either + past = (4 * st + j) * 32 >= d.inter_real + sbs[st] = past.select(fx.Int32(UNIT_SCALE), sbs[st]) + return bvs, sbs + + +def mid_code_index(d, row, group): + """MIDS word of pick row ``row``'s MXFP8 group ``group`` (32 columns): a K + step's 4 groups apart, so the steps of one group % 4 lie together + (``mid_operands``' lane j reads groups j, j + 4, ..).""" + steps = d.inter // 128 + return row * (d.inter // 32) + group % 4 * steps + group // 4 + + +def down_ug_flags(c, ucs): + """This lane's ug flags of a down round's slots ``ucs``: flag f a (slot, + pick tile past one token tile, part), a tile past the slot's repeating its + last; lane l takes flags l, l + 64, .. (TP2's 18 parts a tile, 48 rows: + 108 flags a round).""" + rs = c["rs"] + ug_parts = rs.ug_parts + per_slot = rs.slot_tiles * ug_parts if rs.multi else ug_parts + n = len(ucs) * per_slot + return [ + _down_ug_flag(c, ucs, fx.min(c["lane"] + 64 * j, n - 1), per_slot) + for j in range(-(-n // 64)) + ] + + +def _down_ug_flag(c, ucs, f, per_slot): + rs, tab = c["rs"], c["tab"] + ug_parts = rs.ug_parts + mine = ucs[0] + for i in range_constexpr(1, len(ucs)): + mine = (f // per_slot == i).select(ucs[i], mine) + if const_expr(not rs.multi): + return mplan.ug_flag(mine, f % ug_parts, ug_parts) + tile0 = rs.tile_first(tab, mine) + tiles = rs.tile_first(tab, mine + 1) - tile0 + return mplan.ug_flag( + mplan.polled_tile(tile0, tiles, f % per_slot // ug_parts), + f % ug_parts, + ug_parts, + ) + + +@traced +def down_shared(c, tasks): + """Each of ``tasks``' rows of the shared expert's down projection + (its MXFP8 intermediate, ``load_smid_tile``) for every + token, bf16 -> LDS ``sdl[task][token][16 rows]`` (``down_combine``): + every task's weights in flight first.""" + s, tid, red, d = c["S"], c["tid"], c["red"], c["d"] + a = c["args"] + shops = [ + gemv_fp8_loads(c, a["sw2"], a["sw2_s"], d.sh_inter, task, row_major=True) + for task in tasks + ] + for t0, nn in token_tiles(s): + load_smid_tile(c, t0, nn, d.sh_inter // 32) + for j in range_constexpr(len(tasks)): + gemv_fp8_mfmas( + c, d.sh_inter, c["shared_xl"], c["shared_xsl"], red, shops[j], rows=nn + ) + gpu.barrier() + if tid < ROWS * nn: + v = row_sum(red, tid % ROWS, tid // ROWS) + at = (j * s + plus(t0, tid // ROWS)) * ROWS + tid % ROWS + fx.ptr_store(bf16_round(v).bitcast(fx.Int32), c["sdl"] + at) + gpu.barrier() + + +@traced +def down_combine(c, task, j, t0, n): + """Down task ``task`` (the CTA's j-th) for tokens t0 .. t0 + n: the routed + contributions summed in bf16 in each token's top-k order, plus the shared expert's + (``down_shared``), the sum pushed to every rank.""" + s, tid, d, tab = c["S"], c["tid"], c["d"], c["tab"] + topk = c["rs"].topk + if tid < ROWS * n: + r = tid % ROWS + tl = tid // ROWS + t = plus(t0, tl) + shared = fx.ptr_load(c["sdl"] + ((j * s + t) * ROWS + r)).bitcast(fx.Float32) + # aiter's stage 2 adds each bf16(v * w) into a bf16 buffer: a rounding + # at every add (its order is the atomics'; here top-k order) + routed = fx.Float32(0.0) + for k in range_constexpr(topk): + w = c["rs"].route_w(tab, s, t, k) + v = contrib_load(c, j, (t * topk + k) * 16 + r) + routed = bf16_round(routed + bf16_round(v * w)) + fx.ptr_store(bf16_round(routed + shared), c["pair"] + tid) + gpu.barrier() + if tid < ROWS * n // 2: + rp = tid % (ROWS // 2) + tl = tid // (ROWS // 2) + t = plus(t0, tl) + v0 = fx.ptr_load(c["pair"] + (tl * ROWS + 2 * rp)) + v1 = fx.ptr_load(c["pair"] + (tl * ROWS + 2 * rp + 1)) + col = task * ROWS + 2 * rp + for p in range_constexpr(d.tp): + dst = preg(c["peer_addr"](p), c["ffn_off"], "ffn") + c["put_bf"](dst, 2 * region_pair(c["rank"], t, s, col), [v0, v1]) + gpu.barrier() + + +@traced +def down_mfma(c, j, task, sl, mid, first, count, av, live): + """Down task ``task`` (the CTA's j-th) on a slot of U: the routed MFMA of its + picks first .. first + min(count, 16) (their B operands ``mid``: + ``mid_operands``; the scale block in ``sl``), each pick's contribution.""" + lane, tab, d = c["lane"], c["tab"], c["d"] + sas = lds_scales( + sl, task % 2 * 16 + lane % 16, lane, 0, d.inter // 128, d.down_scale_cols + ) + jj = lane % 16 + col = fx.min(jj, count - 1) + bvs, sbs = mid + # the K steps on DOWN_CHAINS accumulators in turn, summed in order: one + # chain put every FP8-rate MFMA's latency in series + accs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(DOWN_CHAINS)] + for i in range_constexpr(len(av)): + h = i % DOWN_CHAINS + accs[h] = mfma_scaled(av[i], bvs[i], accs[h], sas[i], sbs[i], FP4, FP8) + acc = accs[0] + for h in range_constexpr(1, DOWN_CHAINS): + acc = acc + accs[h] + pick = c["rs"].pick_at(tab, first + col) + for q in range_constexpr(4): + r = 4 * (lane // 16) + q + if live & (jj < count): + contrib_store(c, j, pick * 16 + r, acc[q]) + + +def contrib_store(c, j, i, v): + """Task j's routed contribution word i (``[pick (t, k)][16 rows]`` f32 bits).""" + fx.ptr_store(v.bitcast(fx.Int32), c["contrib"] + (j * c["rs"].picks * 16 + i)) + + +def contrib_load(c, j, i): + return fx.ptr_load(c["contrib"] + (j * c["rs"].picks * 16 + i)).bitcast(fx.Float32) + + +# ---------------------------------------------------------------- loaders + + +@traced +def load_smid_tile(c, t0, n, groups): + """The shared expert's MXFP8 intermediate of tokens t0 .. t0 + n (SMQ / + SMC, ``smid_quant``) into LDS ``shared_xl`` / ``shared_xsl`` rows 0 .., by every + thread.""" + tid, words = c["tid"], c["d"].sh_inter // 4 + poll_copy( + c, + tid, + THREADS, + [ + (c["smq"], lambda u: plus(t0 * words, u), n * words, c["shared_xl"], words), + ( + c["smc"], + lambda u: plus(t0 * groups, u), + n * groups, + c["shared_xsl"], + groups, + ), + ], + ) + gpu.barrier() + + +# ---------------------------------------------------------------- reduce +@traced +def stage_reduce(c, task): + """Columns 128 task ..: the TP ranks' partials, this rank's first, then + rank + 1, + 2, .., in fp32 -> bf16 -> ``out``.""" + s, tid = c["S"], c["tid"] + for pas in range_constexpr(0, s * RED_COLS // 2, THREADS): + reduce_pair(c, task, pas + tid) + + +@traced +def reduce_pair(c, task, idx): + """Thread ``idx``'s column pair of token idx / 64.""" + s, tp = c["S"], c["d"].tp + a = c["args"] + if idx < s * RED_COLS // 2: + t = idx // (RED_COLS // 2) + col = task * RED_COLS + 2 * (idx % (RED_COLS // 2)) + own = preg(c["sym"], c["ffn_off"], "ffn") + acc0, acc1 = sum_partials( + c["poll"], own, lambda src: region_pair(src, t, s, col), tp + ) + bo.buffer_store(acc0.to(fx.BFloat16), rsrc(a["out"]), t * HIDDEN + col) + bo.buffer_store(acc1.to(fx.BFloat16), rsrc(a["out"]), t * HIDDEN + col + 1) + + +def moe_smem(s, rs, timeline): + """The MoE stages' LDS (``timeline``: with its own stamp record).""" + + @fx.struct + class Smem: + # an ug task's partial C: 2 tiles a wave (``ug_red_slot``) + red: fx.Array[fx.Float32, 2 * WAVES * 64 * 4, 16] + tab: fx.Array[fx.Int32, rs.tab_words, 16] + pair: fx.Array[fx.Float32, ROWS * s, 16] + # one stage's buffers at a time (``stage_lds``) + stage: fx.Array[fx.Int32, sum(stage_lds(s, rs)["words"].values()), 16] + tls: fx.Array[fx.Int64, TL_POINTS if timeline else 1, 16] + + return Smem + + +def stage_lds(s, rs): + """Word offsets into ``Smem.stage`` of the buffers only one stage reads: the + shared stage's MXFP8 x tile; the waves' scale blocks (``wsl``, the ug and + down stages'), then the ug stage's routed MXFP8 x tile, or the down stage's + token tile of the shared expert's MXFP8 intermediate (``load_smid_tile``), + each down task's routed contributions (``contrib_store``) and its shared + expert rows (``down_shared``). A barrier ends every stage.""" + d = rs.dims + # the x tile the shared / ug stages hold (every token's up to one tile) + rows = tile_rows(s) + x8 = rows * lds_row(X8_WORDS) + groups = rows * lds_row(HIDDEN // 32) + wsl = WAVES * wsl_words(d) + shared_xl = rows * lds_row(d.sh_inter // 4) + lds_tail(d.sh_inter) + shared_xsl = rows * lds_row(d.sh_inter // 32) + contrib = 2 * rs.picks * 16 # two down tasks' at most + offsets = { + "x8l": 0, + "x8sl": x8, + "wsl": 0, + "x8rl": wsl, + "x8rsl": wsl + x8, + "shared_xl": wsl, + "shared_xsl": wsl + shared_xl, + "contrib": wsl + shared_xl + shared_xsl, + "sdl": wsl + shared_xl + shared_xsl + contrib, + } + words = { + "shared": x8 + groups, + "ug": wsl + x8 + groups, + "down": wsl + shared_xl + shared_xsl + contrib + 2 * s * ROWS, + } + return {"offsets": offsets, "words": {"all": max(words.values())}} + + +def moe_context(s, rs, lds, mb, bases, layout, scratch, args, rank, sym): + """The MoE stages' context: LDS, mailbox, scratch regions, arguments. Its lane + math starts from a ``fresh`` thread id: in K2 it is not the attention and + seam stages', held from theirs on.""" + tid = fresh(fx.thread_idx.x) + c = { + "S": s, + "tid": tid, + "bid": fx.block_idx.x, + "lane": tid % 64, + "wave": tid // 64, + "red": lds.red.ptr, + "tab": lds.tab.ptr, + "pair": lds.pair.ptr, + "put": mb.put, + "put_bf": mb.put_bf, + "put_words": mb.put_words, + "poll": mb.poll, + "tag": mb.tag, + "peer_addr": lambda p: bases[p], + "rank": rank, + "sym": sym, + "ffn_off": ffn_offset(s, rs.tp), + "rs": rs, + "d": rs.dims, + "args": args, + # x (``ffn_normed``) before this launch: plain loads, no readiness wait + "x_cm": 0, + "x_ready": False, + } + for name, off in stage_lds(s, rs)["offsets"].items(): + c[name] = lds.stage.ptr + off + for region, (off, _) in layout.items(): + c[region] = sreg(scratch, off, region) + return c + + +@traced +def run_moe(c, key, bid, tls, tl0): + """The MoE stages, in order; ``timeline`` stamps from point ``tl0`` + 1.""" + s, timeline, tid = key.tokens, key.timeline, c["tid"] + router_tasks = c["rs"].router_tasks + route0, shared0 = c["rs"].route0, c["rs"].shared0 + shared_tasks = c["d"].shared_tasks + ug_tasks = c["rs"].max_u * c["rs"].ug_parts + for t in range(first_task(bid, XQ0), s, BLOCKS): + await_x(c, [t]) + stage_xq(c, t) + publish(c["put"], c["x8rdy"], t, 1, tid == 0) + stamp(timeline, tls, tid, tl0 + 1) + multi = len(token_tiles(s)) > 1 + for task in range(first_task(bid, c["rs"].router0), router_tasks, BLOCKS): + await_x(c, list(range(s))) + stage_router(c, task) + stamp(timeline, tls, tid, tl0 + 2) + for t in range(first_task(bid, route0), s, BLOCKS): + stage_route(c, t) + stamp(timeline, tls, tid, tl0 + 3) + c["xl"], c["xsl"] = c["x8l"], c["x8sl"] + # a CTA's one shared unit (task, token tile) at most: straight-line code, + # whose tile loads LLVM does not hoist into one live set as it would out of + # a task loop (spills). A task's tiles in series on one CTA put three tiles' + # latency on the path down waits on (the weights read again a tile: L2 hits) + assert c["rs"].shared_units <= BLOCKS + unit = first_task(bid, shared0) + if unit < c["rs"].shared_units: + # the gate weights in flight before x's wait + a = c["args"] + gops = gemv_fp8_loads( + c, a["sgu"], a["sgu_s"], HIDDEN, unit % shared_tasks, row_major=True + ) + if const_expr(not multi): + load_x8_rows(c, "x8", "x8s", c["x8l"], c["x8sl"], lambda j: j, s) + stage_shared(c, unit, gops, 0, s) + else: + t0 = unit // shared_tasks * TILE + n = fx.min(fx.Int32(TILE), s - t0) + load_x8_rows(c, "x8", "x8s", c["x8l"], c["x8sl"], lambda j: t0 + j, n) + stage_shared(c, unit % shared_tasks, gops, t0, n) + stamp(timeline, tls, tid, tl0 + 4) + c["xl"], c["xsl"] = c["x8rl"], c["x8rsl"] + if const_expr(multi): + # past one token tile a task gathers its pick tile's rows itself. The + # first round placed (the shared units' CTAs take its last tasks); the + # rest, past a round, from the queue: a CTA a shared or router task + # delayed takes fewer of them + load_route(c, c["tab"]) + n_tiles = c["rs"].tile_first(c["tab"], c["rs"].n_union(c["tab"], s)) + ug_total = mplan.ug_tasks(n_tiles, c["rs"].ug_parts) + task = first_task(bid, c["rs"].ug0) + while task < ug_total: + task = stage_ug_tile(c, task, ug_total) + else: + # one token tile: every task placed (its x in LDS once a CTA) + loaded = fx.Int32(0) + for task in range(first_task(bid, c["rs"].ug0), ug_tasks, BLOCKS): + if loaded == 0: + # X8R (ready long before the routing) first: after the routing + # lands only its own poll remains + load_x8_rows(c, "x8r", "x8rs", c["x8rl"], c["x8rsl"], lambda j: j, s) + load_route(c, c["tab"]) + loaded = fx.Int32(1) + stage_ug(c, task) + stamp(timeline, tls, tid, tl0 + 5) + assert BLOCKS < DOWN_TASKS <= 2 * BLOCKS + # the route table is in LDS already past one token tile (every CTA loads it + # for the ug queue) or when every CTA ran an ug iteration (more ug task + # slots than CTAs), whose first loads it; else load it + if const_expr(not multi and ug_tasks <= BLOCKS): + load_route(c, c["tab"]) + # every CTA its tile; the first DOWN_TASKS - BLOCKS to get here an extra + # one too: the CTAs whose ug ended first (by bid, those ending last held + # two tiles and ended the stage) + down_task = first_task(bid, DOWN0) + rank = take_down_rank(c) + if rank < DOWN_TASKS - BLOCKS: + stage_down(c, [down_task, BLOCKS + rank]) + else: + stage_down(c, [down_task]) + stamp(timeline, tls, tid, tl0 + 6) + for task in range(first_task(bid, RED0), RED_TASKS, BLOCKS): + stage_reduce(c, task) + stamp(timeline, tls, tid, tl0 + 7) + + +def await_x(c, tokens): + """In a launch that also writes x (K2: the FFN seam's norm): tokens' rows + ready.""" + if const_expr(c["x_ready"]): + c["poll"]([(c["xrdy"], t, 1) for t in tokens]) + + +def moe_args(*values) -> dict: + """The MoE stages' arguments by name, in ``ARGS`` order.""" + assert len(values) == len(ARGS) + return dict(zip(ARGS, values)) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/moe_plan.py b/vllm/models/deepseek_v41/amd/mono/stages/moe_plan.py new file mode 100644 index 000000000000..5b5aa6fcdded --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/moe_plan.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/moe_plan.py +"""Every per-step decision of which MoE ug flags (UGF) and intermediate rows +(MID) the ug stage writes and the down stage reads, defined once. + +The route table (``stage_route``) gives each expert with picks a slot of U, its +picks a run from ``slot_first``, and, past one token tile, its picks tiles of +``TILE`` from ``tile_first`` (each tile's slot ``tile_slot``). An ug task covers +a unit -- a slot, or past one token tile a pick tile -- and one of ``ug_parts`` parts +of the intermediate; its flag is ``ug_flag``. A down round polls each of its +slots' flags (``polled_tile`` its tiles) before it reads their MID rows. Called +by the kernel on traced Int32 and by the CPU property test on Python ints. +""" + +from vllm.models.deepseek_v41.amd.mono.common.plan import cdiv, imin + +TILE = 16 # picks a pick tile: an MFMA's N + + +def pick_tiles(count): + """A slot's pick tiles: its ``count`` picks, ``TILE`` a tile.""" + return cdiv(count, TILE) + + +def ug_parts(groups, ug_groups): + """A unit's ug parts: the expert's ``groups`` FP4 groups of its real width, + ``ug_groups`` a part. A part wholly in the padding is not run: MID's padded + columns stay the step's zeros, against zero weights in down.""" + return cdiv(groups, ug_groups) + + +def ug_flag(unit, part, ug_parts): + """The UGF pair of ug task (``unit``, ``part``), ``ug_parts`` parts a unit.""" + return unit * ug_parts + part + + +def ug_tasks(units, ug_parts): + """A step's ug tasks: every part of each of its ``units``.""" + return units * ug_parts + + +def tile_pick0(slot_first, tile, tile_first): + """Pick tile ``tile``'s first pick, its slot's picks from ``slot_first`` and + tiles from ``tile_first``.""" + return slot_first + (tile - tile_first) * TILE + + +def polled_tile(tile_first, tiles, j): + """The pick tile a down round polls as a slot's ``j``-th, the slot's + ``tiles`` tiles from ``tile_first``: the last past them.""" + return tile_first + imin(j, tiles - 1) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/moe_shape.py b/vllm/models/deepseek_v41/amd/mono/stages/moe_shape.py new file mode 100644 index 000000000000..a07e88cfb064 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/moe_shape.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/moe_shape.py +"""The MoE stages' build key and route shape: what they are built for (``MoeBuild``), +and the task counts, placements and LDS route table that follow from its +routing (``RouteShape``), with the table's readers. Plain Python but for the +readers; the kernel's stages are in ``moe``.""" + +from dataclasses import dataclass, field + +import flydsl.expr as fx + +from vllm.models.deepseek_v41.amd.mono.common.plan import BLOCKS, WAVES +from vllm.models.deepseek_v41.amd.mono.stages import moe_plan as mplan +from vllm.models.deepseek_v41.amd.mono.stages.dims import UG_PART, Dims +from vllm.models.deepseek_v41.amd.mono.stages.gemv import ROWS, TILE, token_tiles +from vllm.models.deepseek_v41.amd.mono.stages.router import PARTS as ROUTER_PARTS + +EXPERTS, TOPK = 384, 6 # the target's routing (a build's: ``MoeBuild``) +ROUTER0 = 8 # the router's first task: past the xq tasks (a token each) +MASK_WORDS = 2 # an expert's token bitmap in the route table: tokens <= 64 +# a route lane's experts (EXPERTS / 64) sorted descending, ties kept in index +# order: the network by their count +SORT_NETS = { + 6: [ + (0, 1), + (2, 3), + (4, 5), + (0, 2), + (1, 4), + (3, 5), + (0, 1), + (2, 3), + (4, 5), + (1, 2), + (3, 4), + (2, 3), + ], + 2: [(0, 1)], +} + + +@dataclass(frozen=True) +class MoeBuild: + tokens: int + # routed experts and a token's top-k: the target's, or the DSpark draft's + experts: int = EXPERTS + topk: int = TOPK + # the TP size: every width a rank holds (``Dims``) + tp: int = 4 + timeline: bool = field(default=False, metadata={"sym": "tl"}) + # FP4 groups an ug task (``RouteShape.ug_groups``) + ug_groups: int = 2 + + +@dataclass(frozen=True) +class RouteShape: + """A build's routing: ``experts`` routed experts, ``topk`` a token of + ``tokens``, the task counts and placements that follow, and the route table + ``load_route`` lays out in LDS: every token's (ids, weights), |U|, the + experts of U, each (t, k)'s slot in U, each expert's slot, the picks in slot + order and each slot's first position in it (past one token tile also each + slot's first pick tile and each pick tile's slot).""" + + experts: int + topk: int + tokens: int + tp: int + # FP4 groups (32 intermediate columns) an ug task, its waves splitting K in + # halves (2) or quarters (1). 1 is twice the tasks, so the CTAs' ug ends + # closer together: faster at the draft's widths (40 rows -4 us), slower at + # the target's (48 rows +3 us; the runners pick) + ug_groups: int = 2 + + @property + def dims(self) -> Dims: + return Dims(self.tp) + + @property + def ug_parts(self) -> int: + """An expert's ug tasks (``moe_plan.ug_parts``).""" + return mplan.ug_parts(self.dims.inter_real // UG_PART, self.ug_groups) + + @property + def ug_k_slices(self) -> int: + """The K slices of an ug task's waves: a wave a (group, 16-column half, + slice).""" + return WAVES // (2 * self.ug_groups) + + @property + def picks(self) -> int: + return self.tokens * self.topk + + @property + def max_u(self) -> int: + return min(self.picks, self.experts) + + @property + def router_groups(self) -> tuple[int, int]: + """``stage_router``'s token groups and their tokens: a group a token + tile, as many as one round of CTAs holds (a CTA running two serialized + them: the route waited on the 32 that did at 48 rows).""" + per_group = self.experts // ROWS * ROUTER_PARTS + groups = max(1, min(len(token_tiles(self.tokens)), BLOCKS // per_group)) + tokens = -(-self.tokens // groups) + assert (groups - 1) * tokens < self.tokens + return groups, tokens + + @property + def router_tasks(self) -> int: + # (token group, 16-row tile, K part): ``stage_router`` + return self.router_groups[0] * self.experts // ROWS * ROUTER_PARTS + + @property + def router0(self) -> int: + """Past the xq tasks (a token each from CTA 0).""" + return max(ROUTER0, self.tokens) + + @property + def route0(self) -> int: + return self.router0 + self.router_tasks + + @property + def shared0(self) -> int: + return self.route0 + max(8, self.tokens) + + @property + def shared_units(self) -> int: + """The shared expert's (task, token tile) units (``stage_shared``).""" + return self.dims.shared_tasks * len(token_tiles(self.tokens)) + + @property + def ug0(self) -> int: + """The first ug round's placement: the shared units' CTAs last.""" + return self.shared0 + self.shared_units + + @property + def sort_net(self): + return SORT_NETS[self.experts // 64] + + @property + def tab_words(self) -> int: + return self.tile_slot0 + self.max_tiles if self.multi else self.queue0 + 1 + + @property + def multi(self) -> bool: + """More than one token tile: an ug task is a (tile of a slot's picks, + part) (``stage_ug``).""" + return self.tokens > TILE + + @property + def slot_tiles(self) -> int: + """A slot's pick tiles at most (a token picks an expert once).""" + return len(token_tiles(self.tokens)) + + @property + def max_tiles(self) -> int: + """The pick tiles of U at most: a slot's partial tile each, then the + whole tiles, and never more than a tile a pick.""" + return min(self.picks, self.max_u + self.picks // TILE) + + @property + def tile_first0(self) -> int: + """Past one token tile: each slot's first pick tile, then their count.""" + return self.queue0 + 1 + + @property + def tile_slot0(self) -> int: + """Past one token tile: each pick tile's slot of U.""" + return self.tile_first0 + self.max_u + 1 + + @property + def queue0(self) -> int: + """The ug task this CTA took from the queue (``take_ug_task``).""" + return self.mask0 + MASK_WORDS * self.experts + + @property + def expert_slot0(self) -> int: + """The route table's per-expert slots (a slot where the expert is in U).""" + return 3 * self.picks + 1 + self.max_u + + @property + def order0(self) -> int: + """The picks (t topk + k) in slot order, a slot's ascending.""" + return self.expert_slot0 + self.experts + + @property + def first0(self) -> int: + """Each slot's first position in the pick order, then the pick count.""" + return self.order0 + self.picks + + @property + def mask0(self) -> int: + """Each expert's tokens that picked it, a bit a token (``MASK_WORDS``).""" + return self.first0 + self.max_u + 1 + + @property + def tile_picks(self) -> int: + """A slot's picks at most: a token picks an expert once.""" + return min(self.tokens, 16) + + def route_w(self, tab, s, t, k): + return fx.ptr_load(tab + (t * 2 * self.topk + self.topk + k)).bitcast( + fx.Float32 + ) + + def n_union(self, tab, s): + return fx.ptr_load(tab + 2 * s * self.topk) + + def union_expert(self, tab, s, u): + return fx.ptr_load(tab + (2 * s * self.topk + 1 + u)) + + def slot_first(self, tab, u): + return fx.ptr_load(tab + (self.first0 + u)) + + def pick_at(self, tab, i): + return fx.ptr_load(tab + (self.order0 + i)) + + def tile_first(self, tab, u): + return fx.ptr_load(tab + (self.tile_first0 + u)) + + def tile_slot(self, tab, i): + return fx.ptr_load(tab + (self.tile_slot0 + i)) + + +def route_shape(key: MoeBuild) -> RouteShape: + return RouteShape(key.experts, key.topk, key.tokens, key.tp, key.ug_groups) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/router.py b/vllm/models/deepseek_v41/amd/mono/stages/router.py new file mode 100644 index 000000000000..a7d6f411a816 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/router.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/model_ops/deepseek_v41/router.py +"""V4.1's router GEMV tile for decode: bf16 x [T, 5120] times the gate weight +[384, 5120] -> f32 logits, in one fixed summation order. + +Rows 16 a tile, K in ``PARTS`` contiguous parts, 16x16x32 bf16 MFMA: in part +p wave w takes the part's K steps w, w + 8, ... (32 columns each) in that +order, and the eight wave partials are summed in wave order +(``common.ops.row_sum``); the MoE's route stage sums the parts in part +order. The MoE stages spread a tile's parts over PARTS CTAs, and past 16 tokens +a token tile a CTA too, which the order allows: a CTA alone cannot read a +tile, or every token's x, fast enough. +""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr.typing import T + +from vllm.models.deepseek_v41.amd.mono.common.ops import mfma_bf16, rsrc +from vllm.models.deepseek_v41.amd.mono.common.plan import WAVES + +HIDDEN = 5120 +ROWS = 16 # a tile's rows: the MFMA's M +PARTS = 4 # contiguous K parts of a tile, summed in order + + +def _k0s(part, lane, wave): + """The K offsets of this lane's loads in part ``part`` (8 bf16 each).""" + per_wave = HIDDEN // 32 // WAVES // PARTS + assert per_wave * 32 * WAVES * PARTS == HIDDEN + return [ + (part * per_wave * WAVES + wave + WAVES * i) * 32 + 8 * (lane // 16) + for i in range(per_wave) + ] + + +def router_weight_loads(w, task, part, lane, wave, w_cm=0): + """``router_tile``'s weight operands (rows 16 task .., K part ``part``). + ``w_cm``: their load cache policy (nontemporal where nothing else of the + step reuses them).""" + row = task * ROWS + lane % 16 + return [ + fx.Vector( + bo.buffer_load( + rsrc(w), + (row * HIDDEN + k0) // 2, + vec_width=4, + dtype=T.i32, + cache_modifier=w_cm, + ) + ) + for k0 in _k0s(part, lane, wave) + ] + + +def router_x_loads(x, part, s, lane, wave, x_cm=0, t0=0): + """``router_tile``'s x operands: the s tokens from ``t0`` (column + lane % 16; ``t0`` traced or an int). ``x_cm``: their load cache policy + (device scope when another CTA of the same launch wrote x).""" + tok = fx.min(lane % 16, s - 1) + if const_expr(not isinstance(t0, int) or t0 != 0): + tok = tok + t0 + return [ + fx.Vector( + bo.buffer_load( + rsrc(x), + (tok * HIDDEN + k0) // 2, + vec_width=4, + dtype=T.i32, + cache_modifier=x_cm, + ) + ) + for k0 in _k0s(part, lane, wave) + ] + + +def router_tile_mfma(wvs, xvs, lane, wave, red): + """``router_tile``'s MFMAs, K steps in order -> ``red[wave][lane]``.""" + acc = fx.Vector.filled(4, 0.0, fx.Float32) + for i in range_constexpr(len(wvs)): + acc = mfma_bf16(wvs[i].bitcast(fx.BFloat16), xvs[i].bitcast(fx.BFloat16), acc) + fx.ptr_store(acc, red + (wave * 64 + lane) * 4) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/seam.py b/vllm/models/deepseek_v41/amd/mono/stages/seam.py new file mode 100644 index 000000000000..82abee9f46fa --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/seam.py @@ -0,0 +1,406 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/attn_pre.py (slice, gate, norm) and +# atom/models/deepseek_v41/mono/kernels/attn_post.py (the attention's TP reduce). +"""The mono layer's mHC seam stages, each a task table of a persistent launch +handing off through the tagged mailbox: + + slice (160, 32 hidden columns x 4 streams each): the owed post folded into + the residual (R_new, bf16 out), the collapsed layer input (mailbox + LIN) and the partial mix projection / sum of squares (mailbox PMIX) + gate (S): the 160 partials -> post / comb (Sinkhorn) and the next pre + norm (S): the RMSNorm -> the bf16 row, handed to the caller's store + reduce (160): the FFN seam's attention output, the TP ranks' partials + (pushed to every rank's ATTN region) summed in the all-reduce's + order -> ``pend_lds``, where its slice reads it + push (160): a rank's partial given in memory (an attention vLLM ran) + to every rank's ATTN region, ahead of the slice's reduce + +Every elementwise step follows aiter's ``mhc_fused_post_pre_delayed_rmsnorm`` +rounding; the mix projection and the sums of squares, split-K reductions, +differ only in their split and order. +""" + +import math + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import T + +from vllm.models.deepseek_v41.amd.mono.common.ops import ( + batched_rounds, + bf16_round, + bf_hi, + bf_lo, + butterfly, + hw_exp2, + hw_rcp, + hw_rsq, + lane_gather, + rsrc, + sum_partials, + traced, +) +from vllm.models.deepseek_v41.amd.mono.common.plan import THREADS, WAVES, pair_layout +from vllm.models.deepseek_v41.amd.mono.common.sync import preg +from vllm.models.deepseek_v41.amd.mono.stages.gemv import ( + ld_bf, + ld_f32, + plus, + token_tiles, +) + +HC = 4 +HIDDEN = 5120 +K_MIX = HC * HIDDEN +MIX = HC * (HC + 2) # 24 projections: pre 4 | post 4 | comb 16 +EPS = 1e-20 +HC_EPS = 1e-6 +SINKHORN = 20 +POST_MULT = 2.0 +LOG2E = 1.0 / math.log(2.0) + +COLS = 32 # hidden columns a slice task +SLICES = HIDDEN // COLS # 160 +KT = HC * COLS # projection columns a slice task +PART = MIX + 1 # a slice's partials a token: the mixes, then the sum of squares +GATE_BLOCK = 8 # slices a gate thread sums + + +def scratch_layout(tokens: int) -> dict[str, tuple[int, int]]: + """Mailbox region -> (byte offset, bytes) of a seam; 8 B a pair.""" + pairs = { + "lin": tokens * HIDDEN // 2, # bf16 pairs + "pmix": SLICES * tokens * PART, + } + return pair_layout(pairs.items()) + + +def attn_peer_bytes(tokens: int, tp: int) -> int: + """ATTN: ``[source rank][token][hidden / 2]`` bf16 pairs.""" + return tp * tokens * HIDDEN // 2 * 8 + + +def attn_region_pair(src, t, s, col): + return (src * s + t) * (HIDDEN // 2) + col // 2 + + +def fma(a, b, c): + return fx.Float32(fmath.fma(a, b, c)) + + +def sigmoid(z): + """Triton's ``tl.sigmoid``: 1 / (1 + exp2(-z log2(e))).""" + return fx.Float32(1.0) / (fx.Float32(1.0) + hw_exp2(z * fx.Float32(-LOG2E))) + + +def lane_f32(v, src): + return lane_gather(v.bitcast(fx.Int32), src).bitcast(fx.Float32) + + +def pend_value(c, t, col, cc): + """The owed post's sublayer output at (token t, column col): read from + ``pend``, or where a caller staged it (``pend_lds``, ``[token][32]`` of this + task's columns: a reduction it did first).""" + if const_expr(c.get("pend_lds") is None): + return ld_bf(c["args"]["pend"], t * HIDDEN + col) + return fx.ptr_load(c["pend_lds"] + (t * COLS + cc)) + + +# ---------------------------------------------------------------- slice +@traced +def slice_input(c, task, idx): + """Thread ``idx``'s pair of the layer input (token idx / 16, columns + 32 task + 2 (idx % 16) ..).""" + s, rl = c["S"], c["rl"] + a = c["args"] + c0 = task * COLS + if idx < s * COLS // 2: + t = idx // (COLS // 2) + cc = 2 * (idx % (COLS // 2)) + outs = [] + for d in range_constexpr(2): + acc = fx.Float32(0.0) + for h in range_constexpr(HC): + rb = fx.ptr_load(rl + (t * KT + h * COLS + cc + d)) + acc = acc + ld_f32(a["pre_in"], t * HC + h) * rb + outs.append(acc) + c["put_bf"](c["lin"], t * HIDDEN + c0 + cc, outs) + + +def slice_at(c0, e): + """Element e of a slice task at column c0: (token, stream, hidden column).""" + return e // KT, (e % KT) // COLS, c0 + e % COLS + + +def slice_loads(c, c0, e): + """Element e's loads of ``stage_slice`` (e clamped into the step's).""" + a = c["args"] + t, o, col = slice_at(c0, e) + r = [ld_bf(a["res_in"], (t * HC + h) * HIDDEN + col) for h in range(HC)] + cm = [ld_f32(a["comb_in"], (t * HC + h) * HC + o) for h in range(HC)] + return r, cm, ld_f32(a["post_in"], t * HC + o) + + +def slice_row(c, c0, e, got): + """Element e of ``stage_slice``'s R_new from ``slice_loads``' values.""" + a = c["args"] + t, o, col = slice_at(c0, e) + r, cm, post = got + # aiter's mhc_fused_post_pre_delayed_rmsnorm: ((post x + c0 r0) + # + c1 r1) + c2 r2, each product rounded, then fma(c3, r3, ..) + v = post * pend_value(c, t, col, e % COLS) + (r[0] * cm[0]) + for h in range_constexpr(1, HC - 1): + v = v + cm[h] * r[h] + # bf16 R_new is what the collapse, the sums of squares and the + # mix projection all read + v = bf16_round(fma(cm[HC - 1], r[HC - 1], v)) + bo.buffer_store(v.to(fx.BFloat16), rsrc(a["res_out"]), (t * HC + o) * HIDDEN + col) + fx.ptr_store(v, c["rl"] + (t * KT + o * COLS + e % COLS)) + + +@traced +def stage_slice(c, task): + """R_new, the layer input and the partial projections of hidden columns + 32 task .. +32, all four streams, every token.""" + s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] + rl, fl = c["rl"], c["fl"] + a = c["args"] + c0 = task * COLS + batched_rounds( + tid, + s * KT, + lambda e: slice_loads(c, c0, e), + lambda e, got: slice_row(c, c0, e, got), + ) + for i in range_constexpr((MIX * KT + THREADS - 1) // THREADS): + e = tid + THREADS * i + if e < MIX * KT: + j = e // KT + o = (e % KT) // COLS + # the original's two bf16 MFMAs take fn as hi = bf16(fn) and lo = + # bf16(fn - hi): their sum, exact in fp32, is the weight it applies + f = ld_f32(a["hc_fn"], j * K_MIX + o * HIDDEN + c0 + e % COLS) + hi = bf16_round(f) + fx.ptr_store(hi + bf16_round(f - hi), fl + e) + gpu.barrier() + # the layer input: bf16 R_new of the four streams by the incoming pre-mix, + # products rounded, summed in order (the Triton collapse, no contraction); + # two columns a thread, one bf16 pair a mailbox store + for pas in range_constexpr(0, s * COLS // 2, THREADS): + slice_input(c, task, plus(pas, tid)) + # sum of squares: token t on wave t % WAVES, two columns a lane + for t in range_constexpr(s): + if wave == t % WAVES: + x0 = fx.ptr_load(rl + (t * KT + 2 * lane)) + x1 = fx.ptr_load(rl + (t * KT + 2 * lane + 1)) + sq = butterfly(x0 * x0 + x1 * x1, (32, 16, 8, 4, 2, 1)) + if lane == 0: + c["put"](c["pmix"], (task * s + t) * PART + MIX, sq) + # mixes: fp32 MFMA 16x16x4, A = fn rows (two 16-row tiles), B = R_new^T, + # a token tile at a time + for t0, n in token_tiles(s): + slice_mix(c, task, t0, n) + + +@traced +def slice_mix(c, task, t0, n): + """The mix projection partials of tokens t0 .. t0 + n (``stage_slice``).""" + s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] + rl, fl, red = c["rl"], c["fl"], c["red"] + tok = plus(t0, fx.min(lane % 16, n - 1)) + kw = KT // WAVES + cs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(2)] + for kk in range_constexpr(kw // 4): + colk = wave * kw + lane // 16 + 4 * kk + b = fx.ptr_load(rl + (tok * KT + colk)) + for h in range_constexpr(2): + row = h * 16 + lane % 16 + av = fx.ptr_load(fl + (fx.min(row, MIX - 1) * KT + colk)) + av = (row < MIX).select(av, fx.Float32(0.0)) + ops = rocdl._split_mfma_operands([av, b, cs[h], 0, 0, 0]) + cs[h] = fx.Vector(rocdl.mfma_f32_16x16x4f32(T.vec(4, T.f32), *ops).result) + for h in range_constexpr(2): + fx.ptr_store(cs[h], red + ((wave * 2 + h) * 64 + lane) * 4) + gpu.barrier() + if tid < MIX * n: + tl = tid // MIX + t = plus(t0, tl) + j = tid % MIX + h = j // 16 + ln = 16 * ((j % 16) // 4) + tl + tot = fx.Float32(0.0) + for w in range_constexpr(WAVES): + tot = tot + fx.ptr_load(red + (((w * 2 + h) * 64 + ln) * 4 + j % 4)) + c["put"](c["pmix"], (task * s + t) * PART + j, tot) + gpu.barrier() + + +# ---------------------------------------------------------------- gate +def _row(v, lane, q): + """Comb lanes 8..23 hold ``comb[r][c]`` at lane 8 + 4 r + c: element q of + the row.""" + return lane_f32(v, (lane - 8) // 4 * 4 + 8 + q) + + +def _col(v, lane, q): + return lane_f32(v, 8 + 4 * q + (lane - 8) % 4) + + +@traced +def stage_gate(c, t): + """Token t's 160 slices, in slice order, -> post / comb / next pre.""" + s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] + red = c["red"] + a = c["args"] + groups = SLICES // GATE_BLOCK # 20 + if tid < groups * PART: + g = tid // PART + j = tid % PART + vals = c["poll"]( + [ + (c["pmix"], ((g * GATE_BLOCK + k) * s + t) * PART + j, 1) + for k in range(GATE_BLOCK) + ] + ) + acc = vals[0][0].bitcast(fx.Float32) + for k in range_constexpr(1, GATE_BLOCK): + acc = acc + vals[k][0].bitcast(fx.Float32) + fx.ptr_store(acc, red + (g * PART + j)) + gpu.barrier() + if wave == 0: + j = fx.min(lane, MIX - 1) + mix = fx.ptr_load(red + j) + sq = fx.ptr_load(red + MIX) + for g in range_constexpr(1, groups): + mix = mix + fx.ptr_load(red + (g * PART + j)) + sq = sq + fx.ptr_load(red + (g * PART + MIX)) + sc = [ld_f32(a["hc_scale"], i) for i in range(3)] + base = ld_f32(a["hc_base"], j) + # aiter's mhc_fused_post_pre_delayed_rmsnorm reduce: one rstd = + # rsq(fma(sum, 1/K, eps)), then each gate fma(mix rstd, scale, base) + v = mix * hw_rsq(fma(sq, fx.Float32(1.0 / K_MIX), fx.Float32(EPS))) + if lane < HC: + pre = sigmoid(fma(v, sc[0], base)) + bo.buffer_store(pre + fx.Float32(HC_EPS), rsrc(a["pre_out"]), t * HC + lane) + if (lane >= HC) & (lane < 2 * HC): + post = sigmoid(fma(v, sc[1], base)) + bo.buffer_store( + post * fx.Float32(POST_MULT), rsrc(a["post_out"]), t * HC + lane - HC + ) + cv = fma(v, sc[2], base) + row = [_row(cv, lane, q) for q in range(4)] + # Triton's exp: exp2 of (x - max) log2(e) + m = fx.max(fx.max(row[0], row[1]), fx.max(row[2], row[3])) + cv = hw_exp2((cv - m) * fx.Float32(LOG2E)) + # fast_dividef is x rcp(y); its first "+ eps" contracts into the fma + row = [_row(cv, lane, q) for q in range(4)] + cv = fma(cv, hw_rcp((row[0] + row[1]) + (row[2] + row[3])), fx.Float32(HC_EPS)) + col = [_col(cv, lane, q) for q in range(4)] + cv = cv * hw_rcp(((col[0] + col[1]) + (col[2] + col[3])) + fx.Float32(HC_EPS)) + for _ in range_constexpr(SINKHORN - 1): + row = [_row(cv, lane, q) for q in range(4)] + cv = cv * hw_rcp( + ((row[0] + row[1]) + (row[2] + row[3])) + fx.Float32(HC_EPS) + ) + col = [_col(cv, lane, q) for q in range(4)] + cv = cv * hw_rcp( + ((col[0] + col[1]) + (col[2] + col[3])) + fx.Float32(HC_EPS) + ) + if (lane >= 2 * HC) & (lane < MIX): + bo.buffer_store(cv, rsrc(a["comb_out"]), t * HC * HC + lane - 2 * HC) + gpu.barrier() + + +# ---------------------------------------------------------------- norm +@traced +def stage_norm(c, t, store): + """Token t's RMSNorm (rstd = rsq(fma(sum, 1/H, eps)), bf16((x rstd) w)); + 256 threads, 3 chunks of 8 at 8 t + 2048 c, each chunk's bf16 values handed + to ``store(t, col, ys, live)`` by every thread (``live``: a thread's chunk + in the row; threads 256.. mirror thread 255).""" + tid, lane, wave, red = c["tid"], c["lane"], c["wave"], c["red"] + a = c["args"] + tt = fx.min(tid, 255) + mine = tid < 256 + cols = [tt * 8 + 2048 * ch for ch in range(3)] + xs = [] + for ch in range_constexpr(3): + base = t * HIDDEN + fx.min(cols[ch], HIDDEN - 8) + words = c["poll"]([(c["lin"], base // 2 + k, 1) for k in range(4)]) + v = [] + for k in range_constexpr(4): + v += [bf_lo(words[k][0]), bf_hi(words[k][0])] + xs.append([(cols[ch] < HIDDEN).select(x, fx.Float32(0.0)) for x in v]) + acc = fx.Float32(0.0) + for ch in range_constexpr(3): + for x in xs[ch]: + acc = acc + x * x + acc = butterfly(acc, (1, 2, 4, 8, 16, 32)) + if (lane == 0) & mine: + fx.ptr_store(acc, red + wave) + gpu.barrier() + tot = butterfly(fx.ptr_load(red + lane % 4), (1, 2)) + r = hw_rsq(fma(tot, fx.Float32(1.0 / HIDDEN), fx.Float32(EPS))) + for ch in range_constexpr(3): + col = fx.min(cols[ch], HIDDEN - 8) + ys = [ + bf16_round((xs[ch][j] * r) * ld_bf(a["norm_w"], col + j)) for j in range(8) + ] + store(t, col, ys, mine & (cols[ch] < HIDDEN)) + gpu.barrier() + + +# ---------------------------------------------------------------- reduce +@traced +def stage_reduce(c, task): + """The attention output at this slice's 32 columns: the TP ranks' partials + in the all-reduce's order (``sum_partials``), in fp32 -> bf16 -> + ``pend_lds``.""" + s, tid = c["S"], c["tid"] + for pas in range_constexpr(0, s * COLS // 2, THREADS): + reduce_attn_pair(c, task, plus(pas, tid)) + gpu.barrier() + + +@traced +def reduce_attn_pair(c, task, idx): + """Thread ``idx``'s column pair of token idx / 16 (``stage_reduce``).""" + s = c["S"] + c0 = task * COLS + if idx < s * COLS // 2: + t = idx // (COLS // 2) + col = c0 + 2 * (idx % (COLS // 2)) + own = preg(c["sym"], 0, "attn") + acc0, acc1 = sum_partials( + c["poll"], + own, + lambda src: attn_region_pair(src, t, s, col), + c["d"].tp, + ) + lds = c["pend_lds"] + (t * COLS + col - c0) + fx.ptr_store(bf16_round(acc0), lds) + fx.ptr_store(bf16_round(acc1), lds + 1) + + +@traced +def stage_push(c, task, part): + """This rank's attention partial (``part``: bf16 [S, HIDDEN], wo_b's + unreduced output) at this slice's 32 columns, every token, to every rank's + ATTN region: the pairs ``stage_reduce`` polls, as the bits vLLM wrote.""" + s, tid, tp = c["S"], c["tid"], c["d"].tp + for pas in range_constexpr(0, s * COLS // 2, THREADS): + idx = plus(pas, tid) + if idx < s * COLS // 2: + t = idx // (COLS // 2) + col = task * COLS + 2 * (idx % (COLS // 2)) + w = bo.buffer_load( + rsrc(part), (t * HIDDEN + col) // 2, vec_width=1, dtype=T.i32 + ) + pair = attn_region_pair(c["rank"], t, s, col) + for p in range_constexpr(tp): + c["put_words"](preg(c["peer_addr"](p), 0, "attn"), pair, [w]) diff --git a/vllm/models/deepseek_v41/amd/mono_decode.py b/vllm/models/deepseek_v41/amd/mono_decode.py new file mode 100644 index 000000000000..b404a42de37a --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono_decode.py @@ -0,0 +1,376 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4.1 mono decode layer on ROCm CDNA4 (``VLLM_ROCM_MONO_DECODE=1``). + +A decode step's backbone layer -- the attention seam, the attention, its TP +all-reduce, the FFN seam, the MoE and its all-reduce -- runs as two persistent +FlyDSL launches (``mono``, this directory) on the layer's loaded +weights, the TP reductions inside the kernels. A layer runs whole when it has +the standard seam and no compressor or indexer of its own (compress ratio 1 or +2) over ``fp8_ds_mla`` records. Any other backbone layer (the first, the Engram +layers, the compressor / indexer layers) keeps vLLM's attention seam and +attention and runs the rest as one launch, the FFN launch, from wo_b's +unreduced output. A step takes the kernels when it is decode only, at most +``MAX_ROWS`` rows (eager or a FULL CUDA graph: the model is not torch.compiled); +a whole layer also needs causal SWA windows. Anything else runs the layer as +before; the paths share every tensor at the layer boundary, so they interleave +freely. + +Every TP rank must take the same path at every layer (the kernels wait on each +other's pushes): the conditions read only metadata every rank holds alike. +""" + +from typing import TYPE_CHECKING, cast + +import torch + +import vllm.envs as envs +from vllm.config import CUDAGraphMode +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + get_tp_group, +) +from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.logger import init_logger +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from vllm.models.deepseek_v41.amd.model import DeepseekV4DecoderLayer + from vllm.models.deepseek_v41.amd.mono.runner import MonoLayerWeights + from vllm.models.deepseek_v41.amd.rocm import DeepseekV4ROCMAiterSparseSWAMetadata + from vllm.models.deepseek_v41.sparse_mla import DeepseekV4FlashMLAMetadata + +logger = init_logger(__name__) + +RECORD = 584 # an fp8_ds_mla KV record: 576 data bytes, 8 scale bytes +SWA_WIDTH = 128 # a causal decode token's window slots +MAX_ROWS = 48 # the kernels' step rows at most: 8 requests x (1 + 5 DSpark drafts) +# the kernels' ISA: CDNA4 (gfx950) scaled MFMA and MX formats +CDNA_VERSIONS = (4,) + +_runner = None + + +def _mono_runner(device: torch.device): + """The process's mono runner (one a TP rank, shared by every layer). Its + peer-memory handshake is collective over the TP group: every rank builds it + at the same layer of the same step.""" + global _runner + if _runner is None: + if torch.cuda.is_current_stream_capturing(): + # its peer memory is allocated and exchanged eagerly; vLLM warms + # every graph up eagerly before capturing it + raise RuntimeError( + "VLLM_ROCM_MONO_DECODE was first reached inside a CUDA graph " + "capture, before an eager decode step." + ) + from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer + + assert MAX_TOKENS >= MAX_ROWS, MAX_TOKENS + _runner = DSV41MonoLayer( + get_tensor_model_parallel_world_size(), + get_tensor_model_parallel_rank(), + get_tp_group().cpu_group, + device, + ) + return _runner + + +def _record_view(kv: torch.Tensor, block: int) -> torch.Tensor: + """A layer's fp8_ds_mla cache as [blocks, block, RECORD] bytes (a block's + data rows, then its scale words: the kernels index both from it).""" + kv = kv if kv.dtype == torch.uint8 else kv.view(torch.uint8) + assert kv.stride(0) >= block * RECORD, (tuple(kv.shape), kv.stride(), block) + return torch.as_strided(kv, (kv.shape[0], block, RECORD), (kv.stride(0), RECORD, 1)) + + +def _scale_bytes(scale: torch.Tensor) -> torch.Tensor: + return scale if scale.dtype == torch.uint8 else scale.view(torch.uint8) + + +def _compute_units() -> int: + return current_platform.num_compute_units(torch.accelerator.current_device_index()) + + +def _unsupported_moe(layer: "DeepseekV4DecoderLayer", vllm_config) -> str | None: + """Why the kernels cannot run this deployment's MoE, or None. They read every + one of the 384 experts on each rank, TP-sharded, in AITER's A8W4 MXFP4 + layout.""" + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend + + pc = vllm_config.parallel_config + if pc.enable_expert_parallel or pc.enable_eplb or pc.data_parallel_size > 1: + return "needs tensor parallelism only: no expert or data parallelism" + if envs.VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4: + return "needs AITER's A8W4 MoE (VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4 unset)" + experts = getattr(getattr(layer.ffn, "experts", None), "routed_experts", None) + if experts is None: + return None # not a routed MoE: the layer checks decline it + backend = getattr(experts.quant_method, "mxfp4_backend", None) + if backend != Mxfp4MoeBackend.AITER_MXFP4_BF16: + return f"needs the AITER_MXFP4_BF16 MoE backend, not {backend}" + return None + + +class MonoDecodeLayer: + """One decoder layer's mono path (``create``: None when the layer is not + eligible): the whole layer (``__call__``) or, with ``ffn_only``, the FFN + launch after vLLM's attention (``ffn``). Its weights are taken from the + layer at the first call, after loading.""" + + def __init__(self, ffn_only: bool = False) -> None: + self.ffn_only = ffn_only + self._weights: MonoLayerWeights | None = None + + @staticmethod + def create( + layer: "DeepseekV4DecoderLayer", vllm_config + ) -> "MonoDecodeLayer | None": + if not envs.VLLM_ROCM_MONO_DECODE: + return None + from vllm.platforms.rocm import get_cdna_version + + # the deployment: an explicit opt-in that cannot run is an error + why = None + if get_cdna_version() not in CDNA_VERSIONS: + why = f"needs CDNA{'/'.join(map(str, CDNA_VERSIONS))}" + elif get_tensor_model_parallel_world_size() not in (2, 4): + why = "needs tensor parallel size 2 or 4" + else: + why = _unsupported_moe(layer, vllm_config) + if why is None: + try: + from vllm.models.deepseek_v41.amd.mono.runner import BLOCKS + except ImportError as err: + why = f"needs FlyDSL and AITER's FlyDSL helpers ({err})" + else: + # every CTA of a launch stays resident: a partitioned GPU deadlocks + cus = _compute_units() + if cus < BLOCKS: + why = f"needs {BLOCKS} compute units, the GPU has {cus}" + if why is not None: + raise ValueError(f"VLLM_ROCM_MONO_DECODE {why}.") + # the layer: the kernels serve the backbone's seams and MoE + config = vllm_config.model_config.hf_config + attn, ffn = layer.attn, layer.ffn + if layer.use_sequence_parallel: + why = "sequence-parallel layer" + elif not layer.fuse_seam_norm: + why = "needs aiter's fused mHC seam" + elif attn.layer_id >= config.num_hidden_layers: + why = "draft layer" + elif ( + ffn.shared_experts is None + or ffn.gate.tid2eid is not None + or ffn.n_routed_experts != 384 + or ffn.n_activated_experts != 6 + or ffn.routed_scaling_factor != 1.5 + or ffn.swiglu_limit != 10.0 + or ffn.scoring_func != "sqrtsoftplus" + or not ffn.renormalize + ): + why = "MoE routing outside the kernels' (384 / 6, sqrtsoftplus, x1.5)" + if why is not None: + logger.info_once("DSv4.1 mono decode off for some layers: %s", why) + return None + # the whole layer: the kernels' attention takes the standard layers only + if layer.engram is not None: + why = "Engram" + elif attn.compressor is not None or attn.indexer is not None: + why = "a compressor / indexer" + elif attn.compress_ratio not in (1, 2): + why = f"compress ratio {attn.compress_ratio}" + elif attn.kv_mxfp8 or attn.kv_cache_dtype != "fp8_ds_mla": + why = f"a {attn.kv_cache_dtype} KV cache" + if why is not None: + logger.info_once( + "DSv4.1 mono decode: the FFN launch for layers with %s", why + ) + return MonoDecodeLayer(ffn_only=True) + logger.info_once("DSv4.1 mono decode: the whole standard backbone layers") + return MonoDecodeLayer() + + def weights(self, layer: "DeepseekV4DecoderLayer"): + if self._weights is not None: + return self._weights + from vllm.models.deepseek_v41.amd.mono.runner import MonoLayerWeights + + a, f = layer.attn, layer.ffn + e, sh = f.experts.routed_experts, f.shared_experts + assert sh is not None # ``create`` took only layers with one + self._weights = MonoLayerWeights( + attn=None if self.ffn_only else self._attn_weights(a), + hc_attn_fn=layer.hc_attn_fn, + hc_attn_scale=layer.hc_attn_scale, + hc_attn_base=layer.hc_attn_base, + attn_norm=layer.attn_norm.weight, + hc_ffn_fn=layer.hc_ffn_fn, + hc_ffn_scale=layer.hc_ffn_scale, + hc_ffn_base=layer.hc_ffn_base, + ffn_norm=layer.ffn_norm.weight, + gate_w=f.gate.weight, + bias=f.gate.e_score_correction_bias, + w13=e.w13_weight, + w13_s=e.w13_weight_scale, + w2=e.w2_weight, + w2_s=e.w2_weight_scale, + sgu=sh.gate_up_proj.weight, + sgu_s=_scale_bytes(sh.gate_up_proj.weight_scale), + sw2=sh.down_proj.weight, + sw2_s=_scale_bytes(sh.down_proj.weight_scale), + ) + self._weights.check(get_tensor_model_parallel_world_size()) + return self._weights + + @staticmethod + def _attn_weights(a): + from vllm.models.deepseek_v41.amd.mono.attention.plan import Dims + from vllm.models.deepseek_v41.amd.mono.runner import AttnWeights + + attn = AttnWeights( + layer_id=a.layer_id, + wqkv=a.fused_wqa_wkv.weight, + wqkv_scale=_scale_bytes(a.fused_wqa_wkv.weight_scale), + q_norm=a.q_norm.weight, + kv_norm=a.kv_norm.weight, + wq_b=a.wq_b.weight, + wq_b_scale=_scale_bytes(a.wq_b.weight_scale), + wo_a=a.wo_a.weight, + wo_a_scale=_scale_bytes(a.wo_a.weight_scale), + wo_b=a.wo_b.weight, + wo_b_scale=_scale_bytes(a.wo_b.weight_scale), + attn_sink=a.attn_sink, + cos_sin=a.rotary_emb.cos_sin_cache, + ratio=a.compress_ratio, + ) + attn.check(Dims(get_tensor_model_parallel_world_size())) + return attn + + @staticmethod + def _skip(why: str) -> None: + logger.debug_once("DSv4.1 mono decode skipped a step: %s", why) + return None + + @staticmethod + def _decode_step(layer: "DeepseekV4DecoderLayer", rows: int): + """(why this step takes the original path -- None for a decode-only step + the kernels take -- the layer's SWA metadata, its compressed cache's).""" + if not 1 <= rows <= MAX_ROWS: + return f"{rows} rows", None, None + if not is_forward_context_available(): + return "no forward context", None, None + fc = get_forward_context() + md = fc.attn_metadata + attn = layer.attn + swa, comp = None, None + if isinstance(md, dict): + swa = cast( + "DeepseekV4ROCMAiterSparseSWAMetadata | None", + md.get(attn.swa_cache_layer.prefix), + ) + if attn.compressed_cache_prefix is not None: + comp = cast( + "DeepseekV4FlashMLAMetadata | None", + md.get(attn.compressed_cache_prefix), + ) + why = None + if fc.cudagraph_runtime_mode not in (CUDAGraphMode.NONE, CUDAGraphMode.FULL): + # a PIECEWISE (breakable) capture is replayed for mixed batches too + why = f"{fc.cudagraph_runtime_mode} step" + elif not isinstance(md, dict): # a profile / dummy run, or DBO ubatches + why = "no attention metadata" + elif swa is None: + why = "layer metadata missing" + elif swa.num_prefills != 0 or swa.num_decodes == 0: + why = "not a decode-only step" + elif swa.num_decode_tokens > rows: + why = "more decode tokens than rows" + return why, swa, comp + + def __call__( + self, + layer: "DeepseekV4DecoderLayer", + x: torch.Tensor, + positions: torch.Tensor, + residual: torch.Tensor | None, + post_mix: torch.Tensor | None, + res_mix: torch.Tensor | None, + pre_mix: torch.Tensor | None, + ) -> tuple[torch.Tensor, ...] | None: + """The layer's outputs (x, residual, post_mix, res_mix, pre_mix) by the + mono kernels, or None when this step takes the original path (always for + an ``ffn_only`` layer).""" + if self.ffn_only: + return None + if residual is None or x.dim() != 2: + return self._skip("first layer") + why, swa, comp = self._decode_step(layer, x.shape[0]) + attn = layer.attn + tensors = (x, residual, post_mix, res_mix, pre_mix) + if why is None: + # the kernels' attention: causal SWA windows, the compressed cache + if comp is None: + why = "layer metadata missing" + elif swa.decode_swa_width != SWA_WIDTH: + why = f"SWA width {swa.decode_swa_width}" + elif ( + swa.decode_swa_indices is None + or swa.decode_swa_lens is None + or swa.token_to_req_indices is None + ): + why = "SWA decode metadata missing" + elif not all(t is not None and t.is_contiguous() for t in tensors): + why = "non-contiguous layer inputs" + if why is not None: + return self._skip(why) + assert swa is not None and comp is not None + runner = _mono_runner(x.device) + logger.info_once( + "DSv4.1 mono decode layer active (decode steps of <= %d rows)", MAX_ROWS + ) + out = runner.forward( + self.weights(layer), + x, + residual, + post_mix, + res_mix, + pre_mix, + positions, + swa.slot_mapping, + _record_view(attn.swa_cache_layer.kv_cache, swa.block_size), + swa.decode_swa_indices, + swa.decode_swa_lens, + swa.token_to_req_indices, + topk_indices=attn.topk_indices_buffer, + comp_cache=_record_view( + attn._compressed_kv_cache(), comp.block_size // attn.compress_ratio + ), + comp_block_table=comp.block_table, + ) + return out + + def ffn( + self, + layer: "DeepseekV4DecoderLayer", + part: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor, + res_mix: torch.Tensor, + pre_mix: torch.Tensor, + ) -> tuple[torch.Tensor, ...] | None: + """The layer's outputs by the FFN launch from ``part``, wo_b's unreduced + output, and the attention seam's outputs; None when this step takes the + original path (the caller then reduces ``part`` itself).""" + why, _, _ = self._decode_step(layer, part.shape[0]) + tensors = (part, residual, post_mix, res_mix, pre_mix) + if why is None and not all(t.is_contiguous() for t in tensors): + why = "non-contiguous layer inputs" + if why is not None: + return self._skip(why) + runner = _mono_runner(part.device) + logger.info_once( + "DSv4.1 mono decode FFN launch active (decode steps of <= %d rows)", + MAX_ROWS, + ) + return runner.ffn(self.weights(layer), *tensors) diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 3aa54f9b0e97..3ba6eebcbd17 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -929,7 +929,9 @@ def _o_proj( def _wo_b_after_wo_a(self, zf: torch.Tensor) -> torch.Tensor: if self._wo_b_scale is not None and zf.dim() == 2: - return self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True) + return self._bpre_attn_gemm( + self.wo_b.weight, self._wo_b_scale, zf, self.wo_b.reduce_results + ) return self.wo_b(zf) def forward_mqa( From 1dfe9fbdae734fc2ee8c0494c543c24d25b563d6 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sat, 10 Oct 2026 10:26:40 +0800 Subject: [PATCH 218/624] [Core] Release FlashInfer allreduce workspace with --enable-nccl-comm-suspend (#59360) Signed-off-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- tests/distributed/test_comm_ops.py | 71 +++++++++++++++++++ .../device_communicators/cuda_communicator.py | 7 ++ .../flashinfer_all_reduce.py | 24 +++++-- 3 files changed, 98 insertions(+), 4 deletions(-) diff --git a/tests/distributed/test_comm_ops.py b/tests/distributed/test_comm_ops.py index 253ec8a188af..78aee1efd18d 100644 --- a/tests/distributed/test_comm_ops.py +++ b/tests/distributed/test_comm_ops.py @@ -325,6 +325,77 @@ def test_cuda_communicator_checkpoints_flashinfer_workspaces( workspace.checkpoint_restore.assert_called_once_with(group) +def _fi_ar_communicator( + monkeypatch: pytest.MonkeyPatch, group: object, normal: Mock, quant: Mock +) -> CudaCommunicator: + monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_workspace", normal) + monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_quant_workspace", quant) + monkeypatch.setattr( + flashinfer_all_reduce, + "_fi_ar_workspace_groups", + {id(normal): group, id(quant): group}, + ) + monkeypatch.setattr( + flashinfer_all_reduce, "TorchDistBackend", lambda group: group, raising=False + ) + communicator = CudaCommunicator.__new__(CudaCommunicator) + communicator.cpu_group = group + communicator.pynccl_comm = Mock() + communicator.all2all_manager = None + return communicator + + +def _symm_mem_workspace() -> Mock: + workspace = Mock() + workspace.checkpoint_prepare.side_effect = NotImplementedError + workspace.checkpoint_restore.side_effect = NotImplementedError + return workspace + + +@pytest.mark.parametrize("symm_mem", [False, True], ids=["mnnvl", "symm-mem+mnnvl"]) +def test_cuda_communicator_suspend_detaches_flashinfer_workspaces( + monkeypatch: pytest.MonkeyPatch, + symm_mem: bool, +) -> None: + group = object() + workspaces = (_symm_mem_workspace() if symm_mem else Mock(), Mock()) + communicator = _fi_ar_communicator(monkeypatch, group, *workspaces) + communicator.suspend() + communicator.resume() + + for workspace in workspaces: + workspace.checkpoint_prepare.assert_called_once_with() + workspace.checkpoint_restore.assert_called_once_with(group) + communicator.pynccl_comm.suspend.assert_called_once_with() + communicator.pynccl_comm.resume.assert_called_once_with() + + +def test_cuda_communicator_checkpoint_rejects_symm_mem_workspace( + monkeypatch: pytest.MonkeyPatch, +) -> None: + communicator = _fi_ar_communicator( + monkeypatch, object(), _symm_mem_workspace(), Mock() + ) + with pytest.raises(NotImplementedError): + communicator.checkpoint_prepare() + with pytest.raises(NotImplementedError): + communicator.checkpoint_restore() + + +def test_cuda_communicator_resume_retries_failed_restore( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workspace = Mock() + workspace.checkpoint_restore.side_effect = [RuntimeError("cuMemCreate"), None] + communicator = _fi_ar_communicator(monkeypatch, object(), workspace, Mock()) + communicator.suspend() + with pytest.raises(RuntimeError): + communicator.resume() + communicator.resume() + + assert workspace.checkpoint_restore.call_count == 2 + + @pytest.mark.parametrize( ("backend", "capability", "world_size", "nodes", "expected"), [ diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 992d07e15dad..29d5bae89d2c 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -814,12 +814,19 @@ def _can_use_aiter_ag_rs(self, sizes: list[int] | None) -> bool: return bd is not None and bd.uniform def suspend(self) -> None: + from .flashinfer_all_reduce import checkpoint_prepare_fi_ar_workspaces + + # FlashInfer AR syncs over the gloo cpu_group, so order vs. NCCL is free. + checkpoint_prepare_fi_ar_workspaces(self.cpu_group, skip_unsupported=True) if self.pynccl_comm is not None: self.pynccl_comm.suspend() def resume(self) -> None: + from .flashinfer_all_reduce import checkpoint_restore_fi_ar_workspaces + if self.pynccl_comm is not None: self.pynccl_comm.resume() + checkpoint_restore_fi_ar_workspaces(self.cpu_group, skip_unsupported=True) def checkpoint_prepare(self) -> None: # Only FlashInfer all-reduce and FlashInfer all2all are supported for now. diff --git a/vllm/distributed/device_communicators/flashinfer_all_reduce.py b/vllm/distributed/device_communicators/flashinfer_all_reduce.py index d51e9e0c0e2c..22bcbb775a65 100644 --- a/vllm/distributed/device_communicators/flashinfer_all_reduce.py +++ b/vllm/distributed/device_communicators/flashinfer_all_reduce.py @@ -360,14 +360,30 @@ def _fi_ar_workspaces_for_group(group: ProcessGroup) -> list[Any]: return group_workspaces -def checkpoint_prepare_fi_ar_workspaces(group: ProcessGroup) -> None: +def checkpoint_prepare_fi_ar_workspaces( + group: ProcessGroup, skip_unsupported: bool = False +) -> None: for workspace in _fi_ar_workspaces_for_group(group): - workspace.checkpoint_prepare() + try: + workspace.checkpoint_prepare() + except NotImplementedError: + if not skip_unsupported: + raise + logger.info_once( + "FlashInfer all-reduce workspace backed by torch symmetric memory " + "stays resident during suspend." + ) -def checkpoint_restore_fi_ar_workspaces(group: ProcessGroup) -> None: +def checkpoint_restore_fi_ar_workspaces( + group: ProcessGroup, skip_unsupported: bool = False +) -> None: for workspace in _fi_ar_workspaces_for_group(group): - workspace.checkpoint_restore(TorchDistBackend(group=group)) + try: + workspace.checkpoint_restore(TorchDistBackend(group=group)) + except NotImplementedError: + if not skip_unsupported: + raise atexit.register(destroy_fi_ar_workspace) From 02aa10539f15c8c73aa7fd4a5ee0f691cbf88dcb Mon Sep 17 00:00:00 2001 From: siliangchen-amd Date: Sat, 10 Oct 2026 10:44:06 +0800 Subject: [PATCH 219/624] [Bugfix][ModelLoader] Make enable_weights_track check quantized layers (#57108) Signed-off-by: siliangchen-amd Co-authored-by: Cursor Co-authored-by: Shanshan Shen <467638484@qq.com> --- .../model_loader/test_registry.py | 82 +++++++++++++++++++ .../layers/fused_moe/runner/moe_runner.py | 5 +- .../model_loader/default_loader.py | 36 +++++++- 3 files changed, 119 insertions(+), 4 deletions(-) diff --git a/tests/model_executor/model_loader/test_registry.py b/tests/model_executor/model_loader/test_registry.py index a795497a11b6..7e4b6c537a52 100644 --- a/tests/model_executor/model_loader/test_registry.py +++ b/tests/model_executor/model_loader/test_registry.py @@ -2,13 +2,18 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +import torch from torch import nn from vllm.config import ModelConfig from vllm.config.load import LoadConfig +from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner +from vllm.model_executor.layers.linear import UnquantizedLinearMethod +from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod from vllm.model_executor.model_loader import get_model_loader, register_model_loader from vllm.model_executor.model_loader.base_loader import BaseModelLoader from vllm.model_executor.model_loader.default_loader import DefaultModelLoader +from vllm.model_executor.models.utils import AutoWeightsLoader @register_model_loader("custom_load_format") @@ -88,3 +93,80 @@ def test_default_loader_hf_still_falls_back_to_pt(tmp_path): ) assert use_safetensors is False assert any(f.endswith("model.pt") for f in files) + + +class _QuantMethod: + uses_meta_device = False + + def process_weights_after_loading(self, layer: nn.Module) -> None: + pass + + +def _layer(quant_method=None, **shapes) -> nn.Module: + layer = nn.Module() + for name, shape in shapes.items(): + param = nn.Parameter(torch.zeros(shape), requires_grad=False) + layer.register_parameter(name, param) + layer.quant_method = quant_method + return layer + + +def _track(model: nn.Module, loaded: set[str], quantized: bool) -> str: + loader = DefaultModelLoader(LoadConfig()) + try: + loader.track_weights_loading(model, set(loaded), quantized=quantized) + except ValueError as e: + return str(e) + return "" + + +def test_weights_track_reports_missing_quantized_weights(): + # Quantized checkpoints may only omit KV-cache quantization params, empty + # placeholders, fbgemm_fp8's input_scale_ub and online-quant params. + online = _QuantMethod() + online.uses_meta_device = True + model = nn.Module() + model.proj = _layer(_QuantMethod(), weight=(4, 4), weight_scale_inv=(1, 1)) + model.unquantized = _layer(UnquantizedLinearMethod(), weight=(4, 4), bias=(4,)) + model.gptq = _layer(_QuantMethod(), qweight=(4, 4), qzeros=(0,)) + model.fbgemm = _layer(_QuantMethod(), weight=(4, 4), input_scale_ub=()) + model.attn = _layer(BaseKVCacheMethod(None), k_scale=(), k_zero_point=()) + model.online = _layer(online, weight=(4, 4)) + loaded = {"proj.weight", "unquantized.weight", "gptq.qweight", "fbgemm.weight"} + error = _track(model, loaded, quantized=True) + assert "proj.weight_scale_inv" in error and "unquantized.bias" in error + assert not any(n in error for n in ("gptq", "fbgemm", "attn", "online")) + + +def test_weights_track_keeps_unquantized_models_unchanged(): + # The check is on by default here, so unquantized layers stay exempt. + model = nn.Module() + model.score = _layer(_QuantMethod(), weight=(2, 4)) + model.norm = _layer(weight=(4,)) + assert _track(model, {"norm.weight"}, quantized=False) == "" + assert "norm.weight" in _track(model, set(), quantized=False) + + +def test_moe_runner_reports_registered_expert_names(): + # AutoWeightsLoader prefixes these names, so they must include routed_experts. + class Experts(nn.Module): + def __init__(self): + super().__init__() + self.w13_weight = nn.Parameter(torch.zeros(2, 8, 4), requires_grad=False) + + def load_weights(self, weights): + for _ in weights: + yield "w13_weight" + + class Runner(nn.Module): + load_weights = MoERunner.load_weights + + def __init__(self): + super().__init__() + self.routed_experts = Experts() + + model = nn.Module() + model.experts = Runner() + weights = [("experts.w13_weight", torch.zeros(2, 8, 4))] + loaded = AutoWeightsLoader(model).load_weights(weights) + assert loaded == {"experts.routed_experts.w13_weight"} diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 9cb7cfbb0f3b..d220f8da36bb 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -301,7 +301,10 @@ def __init__( def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]] ) -> Iterable[str]: - return self.routed_experts.load_weights(weights) + # routed_experts reports names relative to itself, but callers such as + # AutoWeightsLoader qualify what this returns with the runner's prefix. + for name in self.routed_experts.load_weights(weights): + yield f"routed_experts.{name}" def _select_forward(self) -> Callable: if current_platform.is_tpu(): diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index 9f1c25dd2e34..b162f60b0577 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/vllm/model_executor/model_loader/default_loader.py @@ -16,6 +16,7 @@ from vllm.config import ModelConfig from vllm.config.load import LoadConfig from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod from vllm.model_executor.layers.quantization.torchao import torchao_version_at_least from vllm.model_executor.model_loader.base_loader import BaseModelLoader from vllm.model_executor.model_loader.ep_weight_filter import ( @@ -47,6 +48,9 @@ logger = init_logger(__name__) +# fbgemm_fp8 builds input_scale_ub from its config; checkpoints do not carry it. +_CONFIG_BUILT_PARAMS = ("input_scale_ub",) + class DefaultModelLoader(BaseModelLoader): """Model loader that can load different file types from disk.""" @@ -506,10 +510,15 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: else default_enable_weights_track ) if enable_weights_track: - self.track_weights_loading(model, loaded_weights) + self.track_weights_loading( + model, loaded_weights, quantized=model_config.quantization is not None + ) def track_weights_loading( - self, model: nn.Module, loaded_weights: set[str] | None + self, + model: nn.Module, + loaded_weights: set[str] | None, + quantized: bool = False, ) -> None: weights_to_load = {name for name, _ in model.named_parameters()} if loaded_weights is not None: @@ -523,7 +532,28 @@ def track_weights_loading( # ignore kv_cache scale and online quant scale, # which can be missing in checkpoints if has_online_quant or has_postprocess_quant: - for param_name, _ in module.named_parameters(): + for param_name, param in module.named_parameters(): + # On a quantized model, every parameter has to come from + # the checkpoint except those a checkpoint cannot carry: + # - online quantization (uses_meta_device) quantizes the + # weights while loading, so its parameters do not map + # to checkpoint tensors; + # - KV-cache quantization parameters (BaseKVCacheMethod) + # fall back to defaults when a checkpoint has none; + # - empty placeholders hold no data, e.g. the qzeros + # moe_wna16 registers for symmetric GPTQ; + # - _CONFIG_BUILT_PARAMS come from the quantization config. + # Unquantized models, where this check runs by default, + # keep the module-wide exemption. + if ( + quantized + and not has_online_quant + and not isinstance(quant_method, BaseKVCacheMethod) + and param.numel() > 0 + and param_name.rsplit(".", 1)[-1] + not in _CONFIG_BUILT_PARAMS + ): + continue full_name = f"{name}.{param_name}" if name else param_name loaded_weights.add(full_name) weights_not_loaded = weights_to_load - loaded_weights From 10cc2f6ae2c9ba7cc5841ece95e27d0562aef1bf Mon Sep 17 00:00:00 2001 From: DimensionSTP <65501090+DimensionSTP@users.noreply.github.com> Date: Sat, 10 Oct 2026 12:00:16 +0900 Subject: [PATCH 220/624] [Bugfix] Bound GDN speculative decode query widths (#55504) Signed-off-by: DimensionSTP <65501090+DimensionSTP@users.noreply.github.com> Co-authored-by: OpenAI Codex --- tests/v1/worker/test_mamba_hybrid_model_state.py | 11 +++++++++++ vllm/v1/worker/gpu/model_states/mamba_hybrid.py | 9 ++++++++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index fa54ccfcedc5..0dcc4725d779 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -314,6 +314,17 @@ def test_padded_rows_keep_the_sentinel(monkeypatch): assert metadata.num_decode_draft_tokens_cpu.tolist() == [0, 0, -1, -1] +def test_rows_wider_than_the_speculative_state_window_keep_the_sentinel( + monkeypatch: pytest.MonkeyPatch, +) -> None: + state = _mamba_hybrid_state(num_speculative_tokens=7) + input_batch = _input_batch([12, 8, 1], None, [False, False, False]) + + metadata = _prepare_attn_metadata(state, monkeypatch, input_batch) + + assert metadata.num_decode_draft_tokens_cpu.tolist() == [-1, 0, 0] + + def test_chunked_prefill_tail_of_two_or_three_tokens_keeps_the_sentinel( monkeypatch, ): diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index c45304efa512..94c57c1968c3 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -288,7 +288,14 @@ def prepare_attn( # Test request state, not num_scheduled_tokens == draft_count+1: # adaptive rewrites num_scheduled_tokens to an even split, so that # equality rarely holds and would demote every verify row to decode. - is_decode = ~is_prefilling_np & (input_batch.num_scheduled_tokens > 0) + is_decode = ( + ~is_prefilling_np + & (input_batch.num_scheduled_tokens > 0) + & ( + input_batch.num_scheduled_tokens + <= self.vllm_config.num_speculative_tokens + 1 + ) + ) num_decode_draft_tokens_np[: input_batch.num_reqs] = np.where( is_decode, num_draft_tokens_per_req, -1 ) From 2d69083c4a09dfeb17d761ae96d2e571a628309c Mon Sep 17 00:00:00 2001 From: xiaofei-zheng Date: Sat, 10 Oct 2026 11:14:01 +0800 Subject: [PATCH 221/624] [ROCm][MXFP4] Align the AITER W4A4 MoE intermediate size to 128 (#55104) Signed-off-by: xiaofei.zheng Co-authored-by: Cursor Co-authored-by: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> --- tests/kernels/moe/test_ocp_mx_moe.py | 76 +++++++++++++++++++ .../layers/fused_moe/oracle/mxfp4.py | 13 +++- 2 files changed, 85 insertions(+), 4 deletions(-) diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index f20927565ebd..13f29087d951 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -2255,3 +2255,79 @@ def run(quant_config): "w_mxfp4_a_mxfp6_e3m2 output is bit-identical to weight-only w_mxfp4:" " the emulation never fake-quantized the activations" ) + + +# (backend, activation, hidden_in, inter_in, hidden_out, inter_out) +# +# The AITER_MXFP4_MXFP4 (W4A4) rows are the ones this table exists for: the +# per-partition intermediate is aligned to 128, not 256, so a model whose native +# shard is already 128-aligned is allocated at its real width instead of being +# padded up. `hidden_size` deliberately stays on 256 -- the gfx950 stage-1 mxfp4 +# kernels are generated with tile_k 256 only, and stage-1 reduces over model_dim. +# +# The remaining rows pin the neighbouring branches so a future edit to this +# function cannot silently move a backend between them. +ROCM_MXFP4_ROUNDUP_CASES = [ + # W4A4: 3072 / TP8 = 384 stays 384. This is the case the alignment change is + # for; the generic 256 round-up would allocate 512. + ("AITER_MXFP4_MXFP4", "SILU", 6144, 384, 6144, 384), + # W4A4, already 256-aligned (e.g. moe_intermediate 2048 at TP8): identity in + # both the old and the new form, so such models are provably untouched. + ("AITER_MXFP4_MXFP4", "SILU", 6144, 256, 6144, 256), + # W4A4, 1024 / TP8 = 128: the most aggressive case, halved rather than trimmed. + ("AITER_MXFP4_MXFP4", "SILU", 4096, 128, 4096, 128), + # W4A4 ignores the activation -- the branch does not read it. + ("AITER_MXFP4_MXFP4", None, 6144, 384, 6144, 384), + ("AITER_MXFP4_MXFP4", "SWIGLUOAI", 6144, 384, 6144, 384), + # W4A4 hidden_size is still rounded to 256, unlike the intermediate. + ("AITER_MXFP4_MXFP4", "SILU", 384, 384, 512, 384), + # W4A16 CK + SiTU/SILU keeps its own 128/128 carve-out. + ("AITER_MXFP4_BF16", "SITU", 384, 384, 384, 384), + ("AITER_MXFP4_BF16", "SILU", 384, 384, 384, 384), + # ... but only for those activations; anything else takes the generic ROCm arm. + ("AITER_MXFP4_BF16", "SWIGLUOAI", 384, 384, 512, 512), + # Backends that stay on the generic ROCm 256/256 round-up. + ("AITER_MXFP4_FP8", "SILU", 384, 384, 512, 512), + ("AITER_TRITON_MXFP4_BF16", "SILU", 384, 384, 512, 512), + ("TRITON", "SWIGLUOAI", 384, 384, 512, 512), +] + + +@pytest.mark.parametrize( + "backend_name,activation_name,hidden_in,inter_in,hidden_out,inter_out", + ROCM_MXFP4_ROUNDUP_CASES, +) +@pytest.mark.skipif(not ROCM_AVAILABLE, reason="ROCm is required for this test") +def test_rocm_mxfp4_round_up_hidden_and_intermediate( + backend_name: str, + activation_name: str | None, + hidden_in: int, + inter_in: int, + hidden_out: int, + inter_out: int, +): + """Pin the ROCm arm of mxfp4_round_up_hidden_size_and_intermediate_size. + + The allocation this function returns is what every MXFP4 MoE weight buffer, + scale buffer and aiter tuned-config lookup is sized from, so a change here is + invisible locally and expensive everywhere else. Asserting the table keeps + each backend on the alignment its kernels actually require. + """ + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + Mxfp4MoeBackend, + mxfp4_round_up_hidden_size_and_intermediate_size, + ) + + backend = Mxfp4MoeBackend[backend_name] + activation = MoEActivation[activation_name] if activation_name else None + + hidden, intermediate = mxfp4_round_up_hidden_size_and_intermediate_size( + backend, hidden_in, inter_in, activation + ) + + assert (hidden, intermediate) == (hidden_out, inter_out), ( + f"{backend_name} / {activation_name}: " + f"hidden {hidden_in}->{hidden} (expected {hidden_out}), " + f"intermediate {inter_in}->{intermediate} (expected {inter_out})" + ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index f5fccc03436f..7667dfc6aef0 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -875,13 +875,18 @@ def mxfp4_round_up_hidden_size_and_intermediate_size( triton_uses_128 = ( backend == Mxfp4MoeBackend.TRITON_UNFUSED and get_cdna_version() != 4 ) + aiter_w4a4 = backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4 - alignment = ( - 128 if is_situ_or_silu and (aiter_uses_128 or triton_uses_128) else 256 + intermediate_alignment = ( + 128 + if aiter_w4a4 or (is_situ_or_silu and (aiter_uses_128 or triton_uses_128)) + else 256 ) + # W4A4 stage 1 still requires a 256-aligned hidden dimension. + hidden_alignment = 256 if aiter_w4a4 else intermediate_alignment - intermediate_size = round_up(intermediate_size, alignment) - hidden_size = round_up(hidden_size, alignment) + intermediate_size = round_up(intermediate_size, intermediate_alignment) + hidden_size = round_up(hidden_size, hidden_alignment) elif backend == Mxfp4MoeBackend.CPU: # CPU AMX kernel uses BLOCK_N=32, align to 32 intermediate_size = round_up(intermediate_size, 32) From f1a3c60558df01b0c1c61af113663c33bb5a84ed Mon Sep 17 00:00:00 2001 From: Giulio De Pasquale Date: Sat, 10 Oct 2026 04:55:45 +0100 Subject: [PATCH 222/624] [Feature][Attention] Cover a whole spec-verify in one DiffKV split-KV program (#59085) Signed-off-by: Giulio De Pasquale --- .../test_triton_unified_attention_diffkv.py | 95 +++++++++++++++++++ .../ops/triton_unified_attention_diffkv.py | 45 ++++++++- 2 files changed, 136 insertions(+), 4 deletions(-) diff --git a/tests/kernels/attention/test_triton_unified_attention_diffkv.py b/tests/kernels/attention/test_triton_unified_attention_diffkv.py index 3b6cf524e2fc..0addd7ccd332 100644 --- a/tests/kernels/attention/test_triton_unified_attention_diffkv.py +++ b/tests/kernels/attention/test_triton_unified_attention_diffkv.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for the Triton DiffKV unified-attention kernel.""" +from unittest.mock import patch + import pytest import torch @@ -16,6 +18,7 @@ get_flash_attn_version, is_flash_attn_varlen_func_available, ) +from vllm.v1.attention.ops import triton_unified_attention_diffkv as diffkv_module from vllm.v1.attention.ops.triton_unified_attention_diffkv import ( unified_attention_diffkv, ) @@ -258,3 +261,95 @@ def test_triton_unified_attn_diffkv_fp8_query(seq_lens, seq_threshold_3D): seq_threshold_3D=seq_threshold_3D, quant_query=True, ) + + +def _capture_diffkv_launch( + query_lens: list[int], max_seqlen_q: int, segm_first_dim: int +): + """Run the launcher with both Triton kernels mocked; return grid/kwargs.""" + num_query_heads, num_kv_heads = 8, 2 + head_size_qk, head_size_v = 128, 128 + num_seqs = len(query_lens) + num_tokens = sum(query_lens) + + query = torch.zeros( + num_tokens, num_query_heads, head_size_qk, dtype=torch.bfloat16, device="cpu" + ) + kv_cache = torch.zeros( + NUM_BLOCKS, + BLOCK_SIZES[0], + num_kv_heads, + head_size_qk + head_size_v, + dtype=torch.bfloat16, + device="cpu", + ) + key_cache = kv_cache[..., :head_size_qk] + value_cache = kv_cache[..., head_size_qk:] + out = torch.zeros( + num_tokens, num_query_heads, head_size_v, dtype=torch.bfloat16, device="cpu" + ) + cu_seqlens_q = torch.tensor( + [0] + query_lens, dtype=torch.int32, device="cpu" + ).cumsum(dim=0, dtype=torch.int32) + seqused_k = torch.full((num_seqs,), 512, dtype=torch.int32, device="cpu") + block_tables = torch.zeros(num_seqs, 32, dtype=torch.int32, device="cpu") + segm_output, segm_max, segm_expsum = _alloc_segm_buffers( + segm_first_dim, num_query_heads, head_size_v + ) + + with ( + patch.object(diffkv_module, "kernel_unified_attention_diffkv") as mock_kernel, + patch.object(diffkv_module, "kernel_reduce_segments_diffkv") as mock_reduce, + ): + unified_attention_diffkv( + q=query, + k=key_cache, + v=value_cache, + out=out, + cu_seqlens_q=cu_seqlens_q, + seqused_k=seqused_k, + softmax_scale=head_size_qk**-0.5, + causal=True, + window_size=(-1, -1), + block_table=block_tables, + softcap=0, + max_seqlen_q=max_seqlen_q, + seq_threshold_3D=8, + num_par_softmax_segments=NUM_PAR_SOFTMAX_SEGMENTS, + softmax_segm_output=segm_output, + softmax_segm_max=segm_max, + softmax_segm_expsum=segm_expsum, + ) + + grid = mock_kernel.__getitem__.call_args.args[0] + kwargs = mock_kernel.__getitem__.return_value.call_args.kwargs + return grid, kwargs, mock_reduce.__getitem__.called + + +def test_spec_verify_block_m_grouping_and_launch_kwargs() -> None: + """Pin program granularity for verify/decode without launching Triton.""" + verify_grid, verify_kwargs, verify_reduce = _capture_diffkv_launch( + query_lens=[8, 8], max_seqlen_q=8, segm_first_dim=16 + ) + num_queries_per_kv = 4 + expected_block_m = min(128, next_power_of_2(8 * num_queries_per_kv)) + assert verify_grid == (4, 2, NUM_PAR_SOFTMAX_SEGMENTS) + assert verify_kwargs["IS_3D"] is True + assert verify_kwargs["BLOCK_M"] == expected_block_m + assert verify_kwargs["BLOCK_Q"] == expected_block_m // num_queries_per_kv + assert verify_kwargs["TILE_SIZE"] == 16 + assert verify_kwargs["num_warps"] == 4 + assert "num_stages" not in verify_kwargs + assert verify_reduce + + decode_grid, decode_kwargs, decode_reduce = _capture_diffkv_launch( + query_lens=[1], max_seqlen_q=1, segm_first_dim=8 + ) + assert decode_grid == (1, 2, NUM_PAR_SOFTMAX_SEGMENTS) + assert decode_kwargs["IS_3D"] is True + assert decode_kwargs["BLOCK_M"] == 16 + assert decode_kwargs["BLOCK_Q"] == 4 + assert decode_kwargs["TILE_SIZE"] == 16 + assert "num_warps" not in decode_kwargs + assert "num_stages" not in decode_kwargs + assert decode_reduce diff --git a/vllm/v1/attention/ops/triton_unified_attention_diffkv.py b/vllm/v1/attention/ops/triton_unified_attention_diffkv.py index 8480754256b2..27b089be0e13 100644 --- a/vllm/v1/attention/ops/triton_unified_attention_diffkv.py +++ b/vllm/v1/attention/ops/triton_unified_attention_diffkv.py @@ -46,6 +46,13 @@ is_batch_invariant = envs.VLLM_BATCH_INVARIANT +# Whole-verify split-KV grouping: one program covers all q tokens of a +# spec-decode verify request (BLOCK_M = q_len * GQA group, capped at 128). +_SPEC_3D_MAX_Q = 16 +_SPEC_3D_BLOCK_M = 128 +_SPEC_3D_NUM_WARPS = 8 +_SPEC_3D_TILE = 16 + @triton.jit def kernel_unified_attention_diffkv( @@ -290,6 +297,9 @@ def kernel_unified_attention_diffkv( acc, mask=dim_mask_v[None, :] & query_mask_0[:, None] & query_mask_1[:, None], ) + # A row whose keys are all masked in this segment ends at M=0, L=0; + # -inf makes reduce_segments treat that segment as absent. + M = tl.where(L > 0.0, M, float("-inf")) store_segm_reduce_scalars( segm_max_ptr, segm_expsum_ptr, @@ -365,6 +375,8 @@ def kernel_reduce_segments_diffkv( ) segm_max = tl.load(segm_max_ptr + segm_offset, mask=segm_mask, other=float("-inf")) overall_max = tl.max(segm_max) + # Graph-capture dummy lengths can leave a query with no causal keys. + overall_max = tl.where(overall_max > float("-inf"), overall_max, 0.0) segm_expsum = tl.load(segm_expsum_ptr + segm_offset, mask=segm_mask, other=0.0) segm_expsum = segm_expsum * tl.exp(segm_max - overall_max) @@ -439,29 +451,53 @@ def unified_attention_diffkv( BLOCK_M = ( 16 if num_queries_per_kv <= 16 else triton.next_power_of_2(num_queries_per_kv) ) - BLOCK_Q = BLOCK_M // num_queries_per_kv - - total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs + launch_kw: dict[str, int] = {} sliding_window_val = 1 + window_size[0] if window_size[0] >= 0 else 0 # Decide between 2D and 3D launch. Mirrors the standard launcher: # 3D requires preallocated softmax buffers, decode-only batches, and # a small number of sequences (otherwise 2D already saturates the SM). + # Spec-decode verify batches (q_len <= _SPEC_3D_MAX_Q, full attention) + # also take 3D: their 2D grid is only q_blocks x kv_heads programs, each + # walking the whole context. Sliding-window layers stay 2D: their loop + # is already window-bounded. + spec_3d = ( + 1 < max_seqlen_q <= _SPEC_3D_MAX_Q + and sliding_window_val == 0 + and softmax_segm_output is not None + and q.shape[0] <= softmax_segm_output.shape[0] + ) use_3d = not ( seq_threshold_3D is None or num_par_softmax_segments is None or softmax_segm_output is None or softmax_segm_max is None or softmax_segm_expsum is None - or max_seqlen_q > 1 + or (max_seqlen_q > 1 and not spec_3d) or num_seqs > seq_threshold_3D or is_batch_invariant ) + spec_tile: int | None = None + if use_3d and spec_3d and _SPEC_3D_BLOCK_M > BLOCK_M: + spec_bm = min( + _SPEC_3D_BLOCK_M, + triton.next_power_of_2(max_seqlen_q * num_queries_per_kv), + ) + if spec_bm > BLOCK_M and spec_bm % num_queries_per_kv == 0: + BLOCK_M = spec_bm + launch_kw["num_warps"] = _SPEC_3D_NUM_WARPS if BLOCK_M >= 128 else 4 + spec_tile = _SPEC_3D_TILE + BLOCK_Q = BLOCK_M // num_queries_per_kv + + total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs + # Tile size: 32 for prefill-class kernels. Decode (small Q) prefers # smaller tiles to expose more parallelism along the KV dim. tile_size = 32 if not use_3d else (16 if q.element_size() >= 2 else 32) + if spec_tile is not None: + tile_size = spec_tile grid: tuple[Any, ...] if use_3d: @@ -530,6 +566,7 @@ def unified_attention_diffkv( BLOCK_M=BLOCK_M, NUM_SEGMENTS_PER_SEQ=num_segments, IS_3D=use_3d, + **launch_kw, ) if use_3d: From be0823357f2c1c7833805be480470df8995371ec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Grunwald?= Date: Sat, 10 Oct 2026 06:09:13 +0200 Subject: [PATCH 223/624] [TEST][XPU][CI] Fix ci uva kernel test instability 2 (#60404) Signed-off-by: microslaw Co-authored-by: Kunshang Ji --- tests/kernels/core/test_uva.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/kernels/core/test_uva.py b/tests/kernels/core/test_uva.py index 9d88eeced3ba..391a6b763cb1 100644 --- a/tests/kernels/core/test_uva.py +++ b/tests/kernels/core/test_uva.py @@ -21,6 +21,9 @@ @pytest.mark.parametrize("device", DEVICES) def test_cpu_write(device): torch.set_default_device(device) + # Initializes the device; otherwise torch.accelerator.synchronize() is a no-op, + # as the UVA view's ops never trigger lazy init. + torch.accelerator.set_device_index(device) cpu_tensor = torch.zeros(10, 10, device="cpu", pin_memory=True, dtype=torch.int32) gpu_view = get_accelerator_view_from_cpu_tensor(cpu_tensor) assert gpu_view.device.type == DEVICE_TYPE @@ -44,6 +47,9 @@ def test_cpu_write(device): @pytest.mark.parametrize("device", DEVICES) def test_gpu_write(device): torch.set_default_device(device) + # Initializes the device; otherwise torch.accelerator.synchronize() is a no-op, + # as the UVA view's ops never trigger lazy init. + torch.accelerator.set_device_index(device) cpu_tensor = torch.zeros(10, 10, device="cpu", pin_memory=True, dtype=torch.int32) gpu_view = get_accelerator_view_from_cpu_tensor(cpu_tensor) assert gpu_view.device.type == DEVICE_TYPE From 29e4d2d42e1333da6bded6f28dd5bff03573af7c Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Sat, 10 Oct 2026 00:22:35 -0400 Subject: [PATCH 224/624] [Frontend][Build] Add `vllm download-kernels` for FlashInfer precompiled kernels (#58765) Signed-off-by: mgoin Signed-off-by: Misha Goin Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/entrypoints.yaml | 1 + README.md | 1 + docker/Dockerfile | 12 +- docker/versions.json | 3 - docs/cli/README.md | 17 ++- docs/contributing/README.md | 1 + .../installation/gpu.cuda.inc.md | 17 ++- docs/getting_started/quickstart.md | 7 +- docs/mkdocs/gen_files/generate_argparse.py | 7 ++ requirements/cuda.txt | 5 - .../unit_tests/test_download_kernels_cli.py | 113 ++++++++++++++++++ vllm/entrypoints/cli/download_kernels.py | 87 ++++++++++++++ vllm/entrypoints/cli/main.py | 2 + vllm/utils/flashinfer.py | 54 ++++++++- vllm/v1/worker/gpu_worker.py | 2 + 15 files changed, 309 insertions(+), 20 deletions(-) create mode 100644 tests/entrypoints/unit_tests/test_download_kernels_cli.py create mode 100644 vllm/entrypoints/cli/download_kernels.py diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index d31c66a8cbd8..bc7fa701b235 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -31,6 +31,7 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/entrypoints + - vllm/utils/flashinfer.py - tests/entrypoints/unit_tests - tests/entrypoints/launchers commands: diff --git a/README.md b/README.md index ca31a8cce94d..ff92fface442 100644 --- a/README.md +++ b/README.md @@ -67,6 +67,7 @@ Install vLLM with [`uv`](https://docs.astral.sh/uv/) (recommended) or `pip`: ```bash uv pip install vllm +vllm download-kernels # precompiled kernels for faster startup ``` Or [build from source](https://docs.vllm.ai/en/latest/getting_started/installation/gpu/index.html#build-wheel-from-source) for development. diff --git a/docker/Dockerfile b/docker/Dockerfile index d524d0fddd47..e29142dd963a 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -936,13 +936,11 @@ RUN --mount=type=cache,target=/opt/uv/cache \ fi && \ rm /tmp/requirements-cuda.txt /tmp/common.txt -# 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.7.0.post1 -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 '.') +# Install FlashInfer's precompiled kernels for all GPU architectures (the default). +# `vllm download-kernels` wraps this command but vLLM is only installed in later +# stages, and running it there would rebuild this large layer for every wheel. +RUN --mount=type=cache,target=/opt/pip/cache \ + PIP_CACHE_DIR=/opt/pip/cache flashinfer download-kernels # ============================================================ # OPENAI API SERVER DEPENDENCIES diff --git a/docker/versions.json b/docker/versions.json index 68f426d0c7b5..12ec728ec8f3 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -73,9 +73,6 @@ "RUN_WHEEL_CHECK": { "default": "true" }, - "FLASHINFER_VERSION": { - "default": "0.7.0.post1" - }, "GDRCOPY_CUDA_VERSION": { "default": "12.8" }, diff --git a/docs/cli/README.md b/docs/cli/README.md index 27425cd92c78..a1eff4c8ae64 100644 --- a/docs/cli/README.md +++ b/docs/cli/README.md @@ -9,7 +9,7 @@ vllm --help Available Commands: ```bash -vllm {chat,complete,serve,launch,bench,collect-env,run-batch,preload} +vllm {chat,complete,serve,launch,bench,collect-env,download-kernels,run-batch,preload} ``` ## serve @@ -188,6 +188,21 @@ Start collecting environment information. vllm collect-env ``` +## download-kernels + +Install the precompiled kernels that vLLM's dependencies would otherwise download +or compile at startup. Run it after installing or upgrading vLLM. It currently +installs FlashInfer's kernels and does nothing when FlashInfer is not installed. + +```bash +vllm download-kernels + +# Print the install commands without running them +vllm download-kernels --dry-run +``` + +See [vllm download-kernels](./download-kernels.md) for the full reference of all available arguments. + ## run-batch Run batch prompts and write results to file. diff --git a/docs/contributing/README.md b/docs/contributing/README.md index e7a6448ca63e..5acb8d980940 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -41,6 +41,7 @@ If you are only developing vLLM's Python code, install vLLM using: ```bash VLLM_USE_PRECOMPILED=1 uv pip install -e . +vllm download-kernels ``` To rebuild only the Rust frontend binary: diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md index f125a189fcd2..6f18681713ef 100644 --- a/docs/getting_started/installation/gpu.cuda.inc.md +++ b/docs/getting_started/installation/gpu.cuda.inc.md @@ -23,14 +23,28 @@ Therefore, it is recommended to install vLLM with a **fresh new** environment. I ```bash uv pip install vllm --torch-backend=auto +vllm download-kernels ``` ??? console "pip" ```bash # Install vLLM with CUDA 12.9. pip install vllm --extra-index-url https://download.pytorch.org/whl/cu129 + vllm download-kernels ``` +`vllm download-kernels` installs [FlashInfer](https://docs.flashinfer.ai/)'s +precompiled kernels for the installed FlashInfer and CUDA versions. Without them, +vLLM downloads and compiles kernels at startup, which can add several minutes on +Hopper and newer GPUs, **especially Blackwell**, and logs a warning. Run it in the +same Python environment after every vLLM install or upgrade, and in custom +container images. The kernels are tied to the FlashInfer version, and FlashInfer +refuses to import when they do not match, so an upgrade needs them refreshed. +With a PyTorch built for CUDA older than 12.9, only `flashinfer-cubin` is +installed, because FlashInfer does not publish `flashinfer-jit-cache` for it. +For CUDA overrides and nightly kernels, use the +[FlashInfer CLI](https://docs.flashinfer.ai/cli.html#download-kernels) directly. + We recommend leveraging `uv` to [automatically select the appropriate PyTorch index at runtime](https://docs.astral.sh/uv/guides/integration/pytorch/#automatic-backend-selection) by inspecting the installed CUDA driver version via `--torch-backend=auto` (or `UV_TORCH_BACKEND=auto`). To select a specific backend (e.g., `cu130`), set `--torch-backend=cu130` (or `UV_TORCH_BACKEND=cu130`). If this doesn't work, try running `uv self update` to update `uv` first. !!! note @@ -94,9 +108,10 @@ If you only need to change Python code, you can build and install vLLM without c git clone https://github.com/vllm-project/vllm.git cd vllm VLLM_USE_PRECOMPILED=1 uv pip install --editable . --torch-backend=auto +vllm download-kernels ``` -This command will do the following: +The install command will do the following: 1. Look for the current branch in your vLLM clone. 1. Identify the corresponding base commit in the main branch. diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index afa9713341b2..a0939a11c4fd 100644 --- a/docs/getting_started/quickstart.md +++ b/docs/getting_started/quickstart.md @@ -25,8 +25,11 @@ This guide will help you quickly get started with vLLM to perform: uv venv --python 3.12 --seed source .venv/bin/activate uv pip install vllm --torch-backend=auto + vllm download-kernels ``` + `vllm download-kernels` installs precompiled kernels that vLLM otherwise downloads and compiles at startup, which can add several minutes on Hopper and newer GPUs. See the [installation guide](installation/gpu.md) for details. + `uv` can [automatically select the appropriate PyTorch index at runtime](https://docs.astral.sh/uv/guides/integration/pytorch/#automatic-backend-selection) by inspecting the installed CUDA driver version via `--torch-backend=auto` (or `UV_TORCH_BACKEND=auto`). To select a specific backend (e.g., `cu126`), set `--torch-backend=cu126` (or `UV_TORCH_BACKEND=cu126`). Another delightful way is to use `uv run` with `--with [dependency]` option, which allows you to run commands such as `vllm serve` without creating any permanent environment: @@ -329,5 +332,5 @@ Some of the available backend options include: - On AMD ROCm: `TRITON_ATTN`, `ROCM_ATTN`, `ROCM_AITER_FA`, `ROCM_AITER_UNIFIED_ATTN`, `TRITON_MLA`, `ROCM_AITER_MLA` or `ROCM_AITER_TRITON_MLA`. - On Intel XPU: `FLASH_ATTN`, `TRITON_ATTN`, `TRITON_MLA`, `XPU_MLA_SPARSE`, `TORCH_SDPA` or `TURBOQUANT`. -!!! warning - There are no pre-built vllm wheels containing Flash Infer, so you must install it in your environment first. Refer to the [Flash Infer official docs](https://docs.flashinfer.ai/) or see [docker/Dockerfile](../../docker/Dockerfile) for instructions on how to install it. +!!! tip + FlashInfer is installed with vLLM, but its precompiled kernels are not. Run `vllm download-kernels` after installing or upgrading vLLM so they are not downloaded and compiled at startup. diff --git a/docs/mkdocs/gen_files/generate_argparse.py b/docs/mkdocs/gen_files/generate_argparse.py index a45de12f2e4a..ec9511efadf0 100644 --- a/docs/mkdocs/gen_files/generate_argparse.py +++ b/docs/mkdocs/gen_files/generate_argparse.py @@ -178,6 +178,9 @@ def auto_mock(module_name: str, attr: str, max_mocks: int = 100): launch_description = auto_mock("vllm.entrypoints.cli.launch", "DESCRIPTION") RenderSubcommand = auto_mock("vllm.entrypoints.cli.launch", "RenderSubcommand") PreloadSubcommand = auto_mock("vllm.entrypoints.cli.preload", "PreloadSubcommand") +DownloadKernelsSubcommand = auto_mock( + "vllm.entrypoints.cli.download_kernels", "DownloadKernelsSubcommand" +) sweep_subcommands = auto_mock("vllm.benchmarks.sweep.cli", "SUBCOMMANDS") openai_cli_args = auto_mock("vllm.entrypoints.launchers", "cli_args") run_batch = auto_mock("vllm.entrypoints.launchers", "run_batch") @@ -332,6 +335,10 @@ def linkify_docs_urls(text: str) -> str: "cli/complete.md": (create_parser(CompleteCommand.add_cli_args), False), "cli/run-batch.md": (create_parser(run_batch.make_arg_parser), True), "cli/preload.md": (create_parser(PreloadSubcommand.add_cli_args), True), + "cli/download-kernels.md": ( + create_parser(DownloadKernelsSubcommand.add_cli_args), + False, + ), "cli/launch/render.md": (create_parser(RenderSubcommand.add_cli_args), True), "cli/bench/latency.md": (create_parser(bench_latency.add_cli_args), True), # URL kept as `mm_processor` for back-compat; command name is `mm-processor` diff --git a/requirements/cuda.txt b/requirements/cuda.txt index e0db38baadd3..1f5a9642b28a 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -10,12 +10,7 @@ torchaudio==2.11.0 torchvision==0.28.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version torchcodec >= 0.14 PyNvVideoCodec==2.0.4 -# FlashInfer should be updated together with the Dockerfile -# flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from -# install_requires so the published wheel does not carry an unresolvable pin ---extra-index-url https://flashinfer.ai/whl/ flashinfer-python==0.7.0.post1 -flashinfer-cubin==0.7.0.post1 apache-tvm-ffi==0.1.11 tilelang==0.1.12 nvidia-cudnn-frontend>=1.19.1 diff --git a/tests/entrypoints/unit_tests/test_download_kernels_cli.py b/tests/entrypoints/unit_tests/test_download_kernels_cli.py new file mode 100644 index 000000000000..75af038134b0 --- /dev/null +++ b/tests/entrypoints/unit_tests/test_download_kernels_cli.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import importlib.metadata +import os +import subprocess +from importlib.util import find_spec +from pathlib import Path +from types import SimpleNamespace + +import pytest +import regex as re + +import vllm.utils.flashinfer as fi +from vllm.entrypoints.cli import download_kernels + + +@pytest.mark.skipif(find_spec("flashinfer") is None, reason="requires FlashInfer") +def test_download_kernels_resolves_kernel_warning( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + capfd: pytest.CaptureFixture[str], + caplog_vllm: pytest.LogCaptureFixture, + disable_log_dedup, +): + """Runs the installed FlashInfer CLI, so a FlashInfer release that renames + its kernel packages, CLI or version check fails here instead of leaving + users with a warning that its own fix cannot clear.""" + # A vLLM upgrade leaves kernels from the previous FlashInfer behind, which + # FlashInfer rejects on import. The command must work in that state. + stale_cubin = tmp_path / "flashinfer_cubin" + stale_cubin.mkdir() + (stale_cubin / "__init__.py").write_text( + f'__version__ = "0"\ndef get_cubin_dir(): return "{stale_cubin}"\n' + ) + monkeypatch.setenv( + "PYTHONPATH", os.pathsep.join([str(tmp_path), os.getenv("PYTHONPATH", "")]) + ) + + assert download_kernels._download_flashinfer_kernels(dry_run=True) == 0 + kernels = dict(re.findall(r"Requirement: ([\w-]+)==(\S+)", capfd.readouterr().out)) + assert kernels + + installed = {"flashinfer-python": importlib.metadata.version("flashinfer-python")} + + def version(name: str) -> str: + if name not in installed: + raise importlib.metadata.PackageNotFoundError(name) + return installed[name] + + def warns() -> bool: + caplog_vllm.clear() + fi.warn_if_flashinfer_kernels_missing() + return "Run `vllm download-kernels`" in caplog_vllm.text + + monkeypatch.delenv("VLLM_HAS_FLASHINFER_CUBIN", raising=False) + monkeypatch.setattr(importlib.metadata, "version", version) + monkeypatch.setattr(fi, "has_flashinfer", lambda: True) + monkeypatch.setattr( + fi, + "current_platform", + SimpleNamespace(is_cuda=lambda: True, has_device_capability=lambda _: True), + ) + + assert warns() + installed.update(dict.fromkeys(kernels, "0")) + assert warns() + installed.update(kernels) + assert not warns() + + +def test_cuda_without_jit_cache_wheels_needs_only_cubin( + monkeypatch: pytest.MonkeyPatch, + caplog_vllm: pytest.LogCaptureFixture, + disable_log_dedup, +): + """FlashInfer publishes no flashinfer-jit-cache below CUDA 12.9, so the + command must not fail on it and the warning must not ask for it.""" + commands = [] + installed = {"flashinfer-python": "1.0", "flashinfer-cubin": "1.0"} + + def version(name: str) -> str: + if name not in installed: + raise importlib.metadata.PackageNotFoundError(name) + return installed[name] + + def run(cmd: list[str], **kwargs) -> SimpleNamespace: + commands.append(cmd) + return SimpleNamespace(returncode=0) + + monkeypatch.delenv("VLLM_HAS_FLASHINFER_CUBIN", raising=False) + monkeypatch.setattr(fi.torch.version, "cuda", "12.8") + monkeypatch.setattr(download_kernels, "find_spec", lambda name: object()) + monkeypatch.setattr(subprocess, "run", run) + monkeypatch.setattr(importlib.metadata, "version", version) + monkeypatch.setattr(fi, "has_flashinfer", lambda: True) + monkeypatch.setattr( + fi, + "current_platform", + SimpleNamespace(is_cuda=lambda: True, has_device_capability=lambda _: True), + ) + + assert download_kernels._download_flashinfer_kernels(dry_run=False) == 0 + assert commands[0][-1] == "install-cubin-wheel" + fi.warn_if_flashinfer_kernels_missing() + assert "vllm download-kernels" not in caplog_vllm.text + + +def test_download_kernels_skips_without_flashinfer(monkeypatch: pytest.MonkeyPatch): + """Builds without FlashInfer (CPU, ROCm, XPU) have nothing to download.""" + monkeypatch.setattr(download_kernels, "find_spec", lambda name: None) + monkeypatch.setattr(subprocess, "run", pytest.fail) + + assert download_kernels._download_flashinfer_kernels(dry_run=False) == 0 diff --git a/vllm/entrypoints/cli/download_kernels.py b/vllm/entrypoints/cli/download_kernels.py new file mode 100644 index 000000000000..a5d8d35afb75 --- /dev/null +++ b/vllm/entrypoints/cli/download_kernels.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The `download-kernels` subcommand for the vLLM CLI. + +Installs the precompiled kernels that vLLM's dependencies would otherwise +download or compile at startup. Run it after installing or upgrading vLLM: + + vllm download-kernels +""" + +import argparse +import os +import subprocess +import sys +import typing +from importlib.util import find_spec + +from vllm.entrypoints.cli.types import CLISubcommand +from vllm.logger import init_logger + +if typing.TYPE_CHECKING: + from vllm.utils.argparse_utils import FlexibleArgumentParser +else: + FlexibleArgumentParser = argparse.ArgumentParser + +logger = init_logger(__name__) + + +def _download_flashinfer_kernels(dry_run: bool) -> int: + if find_spec("flashinfer") is None: + logger.info("FlashInfer is not installed; skipping its kernels.") + return 0 + from vllm.utils.flashinfer import has_flashinfer_jit_cache_wheels + + cmd = [sys.executable, "-m", "flashinfer"] + if has_flashinfer_jit_cache_wheels(): + cmd.append("download-kernels") + else: + logger.info( + "FlashInfer does not publish flashinfer-jit-cache for this CUDA " + "version; installing flashinfer-cubin only." + ) + cmd.append("install-cubin-wheel") + if dry_run: + cmd.append("--dry-run") + # FlashInfer refuses to import while kernels from another FlashInfer version + # are installed, which would stop its own CLI from replacing them. + env = {**os.environ, "FLASHINFER_DISABLE_VERSION_CHECK": "1"} + return subprocess.run(cmd, env=env, check=False).returncode + + +class DownloadKernelsSubcommand(CLISubcommand): + """The `download-kernels` subcommand for the vLLM CLI.""" + + name = "download-kernels" + + @staticmethod + def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: + parser.add_argument( + "--dry-run", + action="store_true", + help="Print the install commands without running them.", + ) + return parser + + @staticmethod + def cmd(args: argparse.Namespace) -> None: + sys.exit(_download_flashinfer_kernels(args.dry_run)) + + def subparser_init( + self, subparsers: argparse._SubParsersAction + ) -> FlexibleArgumentParser: + parser = subparsers.add_parser( + self.name, + help="Install precompiled kernels for faster startup.", + description=( + "Install the precompiled kernels that vLLM's dependencies would " + "otherwise download or compile at startup. Run it after " + "installing or upgrading vLLM." + ), + usage="vllm download-kernels [--dry-run]", + ) + return DownloadKernelsSubcommand.add_cli_args(parser) + + +def cmd_init() -> list[CLISubcommand]: + return [DownloadKernelsSubcommand()] diff --git a/vllm/entrypoints/cli/main.py b/vllm/entrypoints/cli/main.py index 4e1ae68b008c..c85bec4f1809 100644 --- a/vllm/entrypoints/cli/main.py +++ b/vllm/entrypoints/cli/main.py @@ -42,6 +42,7 @@ def main(): import vllm.entrypoints.cli.benchmark.main import vllm.entrypoints.cli.collect_env + import vllm.entrypoints.cli.download_kernels import vllm.entrypoints.cli.launch import vllm.entrypoints.cli.openai import vllm.entrypoints.cli.preload @@ -60,6 +61,7 @@ def main(): vllm.entrypoints.cli.launch, vllm.entrypoints.cli.benchmark.main, vllm.entrypoints.cli.collect_env, + vllm.entrypoints.cli.download_kernels, vllm.entrypoints.cli.preload, vllm.entrypoints.cli.run_batch, vllm.entrypoints.cli.snapshot, diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 9b615cd2fa4b..69401febea96 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -8,6 +8,7 @@ import contextlib import functools import importlib +import importlib.metadata import importlib.util import os import shutil @@ -16,6 +17,7 @@ import requests import torch +from packaging.version import Version import vllm.envs as envs from vllm.logger import init_logger @@ -98,12 +100,60 @@ def has_flashinfer() -> bool: "FlashInfer kernels are disabled: flashinfer-cubin is not installed " "and nvcc (CUDA_HOME, CUDA_PATH, PATH or /usr/local/cuda) or ninja " "(PATH) is missing. Set CUDA_HOME to a CUDA toolkit and put ninja on " - "PATH, or run `flashinfer download-kernels`." + "PATH, or run `vllm download-kernels`." ) return False return True +def _installed_version(distribution: str) -> str | None: + try: + return Version(importlib.metadata.version(distribution)).public + except importlib.metadata.PackageNotFoundError: + return None + + +def has_flashinfer_jit_cache_wheels() -> bool: + """Return whether FlashInfer publishes flashinfer-jit-cache for the installed + PyTorch's CUDA version, which it does from CUDA 12.9.""" + cuda_version = torch.version.cuda + return cuda_version is not None and Version(cuda_version) >= Version("12.9") + + +def warn_if_flashinfer_kernels_missing() -> None: + """Warn on Hopper and newer GPUs when FlashInfer's precompiled kernels are + missing or were installed for another FlashInfer version.""" + if not ( + current_platform.is_cuda() + and current_platform.has_device_capability(90) + and has_flashinfer() + ): + return + flashinfer_version = _installed_version("flashinfer-python") + packages = [] + if not envs.VLLM_HAS_FLASHINFER_CUBIN: + packages.append("flashinfer-cubin") + if has_flashinfer_jit_cache_wheels(): + packages.append("flashinfer-jit-cache") + installed = {name: _installed_version(name) for name in packages} + if flashinfer_version is None or all( + version == flashinfer_version for version in installed.values() + ): + return + logger.warning_once( + "FlashInfer's precompiled kernels are missing or do not match " + "flashinfer-python %s (%s). FlashInfer then downloads and compiles " + "kernels at startup, which can take several minutes, or fails to import " + "mismatched ones. Run `vllm download-kernels` in this Python environment " + "to install them.", + flashinfer_version, + ", ".join( + f"{name} {version or 'not installed'}" + for name, version in installed.items() + ), + ) + + @functools.cache def has_flashinfer_bf16_gemm() -> bool: """Return whether FlashInfer exposes the BF16 dense GEMM API.""" @@ -1243,4 +1293,6 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool: "should_use_flashinfer_for_blockscale_fp8_gemm", "is_flashinfer_fp8_blockscale_gemm_supported", "is_flashinfer_cudnn_fp8_prefill_attn_supported", + "has_flashinfer_jit_cache_wheels", + "warn_if_flashinfer_kernels_missing", ] diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index f2dfc8a689f0..a01e5dd5f7a3 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -73,6 +73,7 @@ from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask from vllm.tracing import instrument +from vllm.utils.flashinfer import warn_if_flashinfer_kernels_missing from vllm.utils.gc_utils import freeze_gc_heap, maybe_attach_gc_debug_callback from vllm.utils.gpu_sync_debug import enable_gpu_sync_check, with_gpu_sync_check from vllm.utils.mem_constants import GiB_bytes @@ -497,6 +498,7 @@ def init_device(self): if self.use_v2_model_runner: logger.info_once("Using V2 Model Runner") + warn_if_flashinfer_kernels_missing() # Set random seed. set_random_seed(self.model_config.seed) From 9bd074a99631a87095465590362f63dda5997ba1 Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Sat, 10 Oct 2026 13:24:00 +0900 Subject: [PATCH 225/624] [Attention] Decode an NVFP4 KV cache with XQA on SM12x (#60452) Signed-off-by: lesj0610 Co-authored-by: Misha Goin --- .buildkite/test_areas/attention.yaml | 14 +++ tests/v1/attention/test_attention_backends.py | 96 ++++++++++++++++++- vllm/v1/attention/backends/flashinfer.py | 21 ++-- 3 files changed, 121 insertions(+), 10 deletions(-) diff --git a/.buildkite/test_areas/attention.yaml b/.buildkite/test_areas/attention.yaml index b9f430a20ee0..83ba90e2cb1b 100644 --- a/.buildkite/test_areas/attention.yaml +++ b/.buildkite/test_areas/attention.yaml @@ -62,3 +62,17 @@ steps: commands: - pytest -v -s v1/attention --ignore=v1/attention/test_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + +- label: ":nvidia: (DGX) Spark FlashInfer NVFP4 XQA" + key: flashinfer-nvfp4-xqa-dgx-spark + timeout_in_minutes: 30 + device: dgx-spark + optional: true + num_devices: 1 + depends_on: + - arm64-image-build + source_file_dependencies: + - vllm/v1/attention/backends/flashinfer.py + - tests/v1/attention/test_attention_backends.py + commands: + - pytest -v -s v1/attention/test_attention_backends.py -k test_flashinfer_xqa_nvfp4_decode_correctness diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index 77366c02aebf..add611d92b6a 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -31,6 +31,7 @@ STR_DTYPE_TO_TORCH_DTYPE, is_quantized_kv_cache, is_torch_equal_or_newer, + nvfp4_kv_cache_full_dim, set_random_seed, ) from vllm.v1.attention.backend import ( @@ -41,7 +42,7 @@ max_decode_query_len, ) from vllm.v1.attention.backends.registry import AttentionBackendEnum -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout, KVQuantMode BACKENDS_TO_TEST = [ AttentionBackendEnum.FLASH_ATTN, @@ -1698,6 +1699,99 @@ def causal_mask_mod( ) +@pytest.mark.skipif( + AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST + or not current_platform.is_device_capability_family(120), + reason="NVFP4 XQA decode requires SM12x.", +) +@pytest.mark.parametrize("use_cuda_graph", [False, True]) +@pytest.mark.parametrize("q_len", [1, 4]) +def test_flashinfer_xqa_nvfp4_decode_correctness( + default_vllm_config, q_len, use_cuda_graph +): + """Single-token and spec-decode queries over an NVFP4 cache decode with + XQA, eager and replayed, and match the fa2 prefill that wrote the cache.""" + import unittest.mock + + from vllm.v1.attention.backends import flashinfer as fi + + seq_lens, bs, d = [40, 300, 1027], 16, 256 + config = create_vllm_config(max_model_len=max(seq_lens), block_size=bs) + # GQA 6 at head_dim 256, as Qwen3.8-27B at TP=2. + config.model_config.model_arch_config = replace( + config.model_config.model_arch_config, + total_num_attention_heads=24, + total_num_kv_heads=4, + head_size=d, + ) + config.cache_config.cache_dtype = "nvfp4" + if q_len > 1: + config.speculative_config = SpeculativeConfig( + method="ngram", num_speculative_tokens=q_len - 1 + ) + spec = FullAttentionSpec( + block_size=bs, + num_kv_heads=4, + head_size=d, + dtype=torch.uint8, + kv_quant_mode=KVQuantMode.NVFP4, + ) + cg_support = fi.FlashInferMetadataBuilder.get_cudagraph_support(config, spec) + assert cg_support == AttentionCGSupport.UNIFORM_BATCH + + device = torch.device(f"{DEVICE_TYPE}:0") + # Request i owns consecutive pages from starts[i]; page 0 stays null. + pages = [0] + [cdiv(s, bs) for s in seq_lens[:-1]] + starts = 1 + torch.tensor(pages).cumsum(0) + max_pages = cdiv(max(seq_lens), bs) + block_table = (starts[:, None] + torch.arange(max_pages)).int().to(device) + full_dim = nvfp4_kv_cache_full_dim(d) + kv_cache = torch.zeros( + int(starts[-1]) + max_pages, bs, 8, full_dim, dtype=torch.uint8, device=device + ).transpose(1, 2) + q, k, v = ( + torch.randn(sum(seq_lens), h, d).bfloat16().to(device) for h in (24, 4, 4) + ) + + def run(query_lens, use_cuda_graph=False): + req = torch.repeat_interleave(torch.arange(3), torch.tensor(query_lens)) + pos = torch.cat([torch.arange(s - n, s) for s, n in zip(seq_lens, query_lens)]) + rows = pos + torch.tensor([0, *seq_lens]).cumsum(0)[req] + metadata = replace( + create_common_attn_metadata(BatchSpec(seq_lens, query_lens), bs, device), + block_table_tensor=block_table, + slot_mapping=(starts[req] * bs + pos).to(device), + ) + with set_current_vllm_config(config): + return rows, run_attention_backend( + AttentionBackendEnum.FLASHINFER, + spec, + ["placeholder"], + config, + device, + metadata, + q[rows], + k[rows], + v[rows], + kv_cache, + kv_cache_dtype="nvfp4", + use_cuda_graph=use_cuda_graph, + layer_k_scale=0.02, + layer_v_scale=0.03, + ) + + # The full prefill writes every token through the NVFP4 writer; its fa2 + # output for the last q_len tokens of each request is the reference. + _, reference = run(seq_lens) + xqa_fn = fi.flashinfer_xqa_batch_decode_with_kv_cache + with unittest.mock.patch.object( + fi, "flashinfer_xqa_batch_decode_with_kv_cache", wraps=xqa_fn + ) as xqa: + rows, output = run([q_len] * len(seq_lens), use_cuda_graph) + assert xqa.called + torch.testing.assert_close(output, reference[rows], atol=2e-2, rtol=2e-2) + + if current_platform.is_rocm(): # FLASH_ATTN is not supported on ROCm SLIDING_WINDOW_BACKENDS_TO_TEST = [ diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 85b7456dd889..1a4145844705 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -878,12 +878,10 @@ def __init__( # Prefer TRTLLM/XQA for decoding whenever supported. The decode kernel # must be selected statically for FULL cudagraph capture. - can_use_xqa_or_trtllm_gen_decode = ( - can_use_trtllm_attention( - self.num_qo_heads, self.num_kv_heads, is_prefill=False - ) - and not self.nvfp4_fa2 - ) + # XQA reads an NVFP4 cache only on SM12x; SM90 decodes it with fa2. + can_use_xqa_or_trtllm_gen_decode = can_use_trtllm_attention( + self.num_qo_heads, self.num_kv_heads, is_prefill=False + ) and (not self.nvfp4_fa2 or current_platform.is_device_capability_family(120)) # Page sizes >= 128 require the trtllm-gen GQA/MQA path (guaranteed by # get_supported_kernel_block_sizes). assert self.page_size <= 64 or ( @@ -1078,7 +1076,10 @@ def get_cudagraph_support( ) or ( is_xqa_arch and ( - spec.kv_quant_mode.is_nvfp4 + ( + spec.kv_quant_mode.is_nvfp4 + and not current_platform.is_device_capability_family(120) + ) or not _is_xqa_head_dim_supported(spec.head_size) ) ): @@ -2283,7 +2284,6 @@ def forward( use_dcp = self.dcp_world_size > 1 if decode_with_xqa: assert not use_dcp - assert not self.is_kvcache_nvfp4 assert self.o_sf_scale is None assert output.dtype != FP4_DTYPE @@ -2628,7 +2628,10 @@ def forward( flashinfer_xqa_batch_decode_with_kv_cache( query=decode_query, - kv_cache=kv_cache_tuple, + kv_cache=( + nvfp4_kv_data if self.is_kvcache_nvfp4 else kv_cache_tuple + ), + kv_cache_sf=nvfp4_kv_block_scales, workspace_buffer=workspace_buffer, block_tables=block_tables_decode, seq_lens=seq_lens_decode, From d5f356dae093d2e47b97af7c9cd8d2fc1aedc56f Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Sat, 10 Oct 2026 00:25:44 -0400 Subject: [PATCH 226/624] [HiSparse] Own GPU residency instead of editing core block-table rows (#60647) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- csrc/libtorch_stable/hisparse_kernels.cu | 24 ++-- .../v1/attention/test_sparse_mla_backends.py | 63 ++++++++-- tests/v1/core/test_prefix_caching.py | 94 +++++++++++---- .../unit/test_hisparse_connector.py | 52 ++++++-- tests/v1/worker/test_utils.py | 3 + .../kv_connector/v1/hisparse/connector.py | 15 ++- .../kv_connector/v1/hisparse/worker.py | 19 ++- vllm/v1/attention/backends/mla/index_group.py | 12 +- vllm/v1/core/sched/output.py | 3 - vllm/v1/core/single_type_kv_cache_manager.py | 42 ++++++- vllm/v1/hisparse/binding.py | 28 ++++- vllm/v1/hisparse/coordinator.py | 111 +++++++++++------- vllm/v1/hisparse/runtime.py | 76 ++++++++++-- vllm/v1/hisparse/types.py | 9 ++ vllm/v1/worker/gpu/model_runner.py | 6 +- 15 files changed, 429 insertions(+), 128 deletions(-) diff --git a/csrc/libtorch_stable/hisparse_kernels.cu b/csrc/libtorch_stable/hisparse_kernels.cu index 069eddf3b0ca..5c7621c2f598 100644 --- a/csrc/libtorch_stable/hisparse_kernels.cu +++ b/csrc/libtorch_stable/hisparse_kernels.cu @@ -194,7 +194,7 @@ constexpr int kHashLoadInverse = 4; // the page is not resident) and returns its host row (-1 when not host-backed). __device__ __forceinline__ int32_t translate_topk_entry( const int32_t token_index, const int32_t request_row, - const int32_t* __restrict__ source_block_table, + const int32_t state_row, const int32_t* __restrict__ source_block_table, const int32_t* __restrict__ resident_block_table, const int64_t host_rows, const int64_t source_bt_stride, const int32_t source_num_reqs, const int32_t source_num_blocks, const int32_t source_block_size, @@ -221,10 +221,10 @@ __device__ __forceinline__ int32_t translate_topk_entry( if (resident_block_table != nullptr) { const int32_t resident_block = token_index >= 0 ? token_index / resident_block_size : -1; - if (request_row >= 0 && request_row < resident_num_reqs && + if (state_row >= 0 && state_row < resident_num_reqs && resident_block >= 0 && resident_block < resident_num_blocks) { const int32_t physical_block = - resident_block_table[static_cast(request_row) * + resident_block_table[static_cast(state_row) * resident_bt_stride + resident_block]; if (physical_block != resident_null_block && physical_block >= 0) { @@ -424,11 +424,11 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel( for (int i = tid; i < top_k; i += blockDim.x) { int32_t resident_row; const int32_t g = translate_topk_entry( - row_topk[i], request_row, source_block_table, resident_block_table, - host_rows, source_bt_stride, source_num_reqs, source_num_blocks, - source_block_size, resident_bt_stride, resident_num_reqs, - resident_num_blocks, resident_block_size, resident_null_block, - resident_row); + row_topk[i], request_row, state_row, source_block_table, + resident_block_table, host_rows, source_bt_stride, source_num_reqs, + source_num_blocks, source_block_size, resident_bt_stride, + resident_num_reqs, resident_num_blocks, resident_block_size, + resident_null_block, resident_row); if (resolved_global_indices != nullptr) { resolved_global_indices[static_cast(row) * top_k + i] = g; } @@ -602,8 +602,8 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel( g = translate_topk_entry( global_indices[static_cast(row) * input_row_stride + position % top_k], - request_row, source_block_table, resident_block_table, host_rows, - source_bt_stride, source_num_reqs, source_num_blocks, + request_row, state_row, source_block_table, resident_block_table, + host_rows, source_bt_stride, source_num_reqs, source_num_blocks, source_block_size, resident_bt_stride, resident_num_reqs, resident_num_blocks, resident_block_size, resident_null_block, resident_row); @@ -708,8 +708,8 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel( const int32_t g = translate_topk_entry( global_indices[static_cast(row) * input_row_stride + position % top_k], - request_row, source_block_table, resident_block_table, host_rows, - source_bt_stride, source_num_reqs, source_num_blocks, + request_row, state_row, source_block_table, resident_block_table, + host_rows, source_bt_stride, source_num_reqs, source_num_blocks, source_block_size, resident_bt_stride, resident_num_reqs, resident_num_blocks, resident_block_size, resident_null_block, resident_row); diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 160b766da077..6eb089ce3ed6 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -3241,7 +3241,8 @@ def test_hisparse_prefill_staging_plan_resolves_resident_sources(): [[11, 12, 0, 13, 0, 0], [21, 0, 22, 23, 0, 0]], dtype=torch.int32 ) - plan.ensure_gpu_sources(resident_table, resident_block_size) + state_indices = torch.arange(2, dtype=torch.int32) + plan.ensure_gpu_sources(resident_table, state_indices, resident_block_size) assert plan.gpu_row_ids is not None unique_hosts = (plan.row_ids[0].view(-1, block_size)[:, 0] // block_size).tolist() @@ -3271,12 +3272,51 @@ def test_hisparse_prefill_staging_plan_resolves_resident_sources(): check_dtype=False, ) - plan.ensure_gpu_sources(torch.zeros_like(resident_table), resident_block_size) + plan.ensure_gpu_sources( + torch.zeros_like(resident_table), state_indices, resident_block_size + ) assert plan.gpu_row_ids is not None assert (plan.gpu_row_ids == -1).all() +def test_hisparse_prefill_resident_sources_follow_each_resident_group(): + """Layers of different resident groups share one batch staging plan. + + Each group's layers must stage from that group's resident blocks. Keying + the plan's cached resident sources on a per-call batch gather let a later + group's gather reuse the freed address of an earlier one, so its layers + read the earlier group's block ids. + """ + block_size = 4 + plan = build_hisparse_prefill_staging_plan( + torch.tensor([[5, 2]], dtype=torch.int32), + torch.tensor([8], dtype=torch.int32), + block_size, + staging_block_capacity=2, + ) + # The prefill runs in request state row 1; one table holds both groups. + # Rows are wide so that a cache keyed on a freed per-call gather of a row + # would see the next gather reuse its address. + state_rows = torch.zeros((4, 2, 1024), dtype=torch.int32) + state_rows[1, 0, :2] = torch.tensor([11, 12]) + state_rows[1, 1, :2] = torch.tensor([21, 22]) + state_indices = torch.tensor([1], dtype=torch.int32) + host_ids = (plan.row_ids[0].view(-1, block_size)[:, 0] // block_size).tolist() + + for group in (0, 1, 0, 1): + plan.ensure_gpu_sources(state_rows[:, group], state_indices, block_size) + assert plan.gpu_row_ids is not None + gpu_blocks = plan.gpu_row_ids[0].view(-1, block_size)[:, 0] // block_size + staged = { + host_id: gpu_block + for host_id, gpu_block in zip(host_ids, gpu_blocks.tolist()) + if host_id > 0 + } + first_page, second_page = state_rows[1, group, :2].tolist() + assert staged == {5: first_page, 2: second_page} + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_hisparse_gather_prefill_cache_prefers_resident_rows(): """Staged rows come from the resident cache when a shadow page exists.""" @@ -3294,7 +3334,11 @@ def test_hisparse_gather_prefill_cache_prefers_resident_rows(): dtype=torch.int32, device=device, ) - plan.ensure_gpu_sources(resident_table, resident_block_size) + plan.ensure_gpu_sources( + resident_table, + torch.arange(2, dtype=torch.int32, device=device), + resident_block_size, + ) num_host_blocks, num_res_blocks = 10, 24 host_cache = ( @@ -3677,6 +3721,7 @@ def test_hisparse_resident_prefill_uses_attention_block_stride(): all_context_pages_resident=True, view=SimpleNamespace(block_size=64, attention_block_stride=832), block_table=torch.tensor([[3]], dtype=torch.int32), + batch_block_table=lambda: torch.tensor([[3]], dtype=torch.int32), runtime=SimpleNamespace(max_swap_rows=1), ) index_group = object.__new__(HiSparseMLAIndexGroup) @@ -4091,6 +4136,7 @@ def gather(kv_cache, staging_plan, resident_cache=None): cache = SimpleNamespace( runtime=SimpleNamespace( gather_prefill_cache=gather, + request_state_indices=torch.tensor([0], dtype=torch.int32), ), view=SimpleNamespace(cache=resident_cache, block_size=1), block_table=resident_block_table, @@ -4113,10 +4159,6 @@ def gather(kv_cache, staging_plan, resident_cache=None): assert result is staged assert block_table is plan.block_table torch.testing.assert_close(request_ids, metadata.req_id_per_token) - plan.ensure_gpu_sources.assert_called_once() - args = plan.ensure_gpu_sources.call_args.args - torch.testing.assert_close(args[0], resident_block_table) - assert args[1] == 1 assert calls == [(source, plan, resident_cache)] @@ -4134,6 +4176,9 @@ def test_hisparse_fp8_prefill_gather_uses_dedicated_stream(monkeypatch): cache = SimpleNamespace( view=SimpleNamespace(cache=resident_cache, block_size=4), block_table=torch.tensor([[2, 3]], dtype=torch.int32), + runtime=SimpleNamespace( + request_state_indices=torch.tensor([0], dtype=torch.int32) + ), ) index_group = object.__new__(HiSparseMLAIndexGroup) index_group.caches = [cache] @@ -4178,10 +4223,6 @@ def test_hisparse_fp8_prefill_gather_uses_dedicated_stream(monkeypatch): assert args.kwargs["host_cache"].data_ptr() == source.data_ptr() assert args.kwargs["host_row_ids"] is plan.row_ids assert args.kwargs["device_row_ids"] is plan.gpu_row_ids - plan.ensure_gpu_sources.assert_called_once() - ensure_args = plan.ensure_gpu_sources.call_args.args - torch.testing.assert_close(ensure_args[0], cache.block_table) - assert ensure_args[1] == 4 def test_sparse_impl_observes_repointed_indexer_buffer(): diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index a89fd085f1fd..836286d127fa 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -279,6 +279,12 @@ def _allocate_scheduled( return blocks +def _resident_block_ids(manager: KVCacheManager, request_id: str) -> list[int]: + """The GPU block HiSparse reads each resident page from; 0 if host-only.""" + resident = get_hisparse_coordinator(manager).resident_managers[0] + return [block.block_id for block in resident.get_residency_row(request_id)] + + def test_hisparse_does_not_write_back_reprefillable_tokens(): """Multi-module MTP re-prefills the last tokens, so they are not final.""" manager = make_hisparse_kv_cache_manager(32, 16, num_prefill_lookahead=3) @@ -319,7 +325,7 @@ def test_hisparse_builds_dma_row_mirrors_across_pages(): def test_hisparse_async_speculation_mirrors_uncertain_position_range(): """Unresolved drafts must not leave gaps in the eager host mirror.""" coordinator = MagicMock(host_group_id=0) - coordinator.take_block_table_updates.return_value = {} + coordinator.take_residency_updates.return_value = {} coordinator.build_offload_command.return_value = None coordinator.build_row_mirrors.return_value = () scheduler = HiSparseConnectorScheduler( @@ -332,7 +338,6 @@ def test_hisparse_async_speculation_mirrors_uncertain_position_range(): connector.connector_scheduler = scheduler connector.update_state_after_alloc(request, None, 0) scheduler_output = SimpleNamespace( - block_table_updates=None, kv_cache_block_copies=None, scheduled_new_reqs=[], scheduled_cached_reqs=SimpleNamespace( @@ -370,11 +375,66 @@ def test_hisparse_reports_when_context_is_fully_resident(): hot_manager.require_hot(request.request_id) assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) is not None _publish_hisparse_pages(manager) + coordinator.take_residency_updates([request.request_id]) + core_row = manager.get_block_ids(request.request_id)[2] pool = manager.block_pool pool.get_new_blocks(pool.get_num_free_blocks()) assert not coordinator.all_context_pages_resident(scheduled) - assert coordinator.take_block_table_updates().keys() == {request.request_id} + # The block-table row stays append-only; the worker learns of the lost + # pages from the residency update instead. + assert manager.get_block_ids(request.request_id)[2] == core_row + resident_ids = _resident_block_ids(manager, request.request_id) + lost = [page for page, block_id in enumerate(resident_ids) if block_id == 0] + assert lost + update = coordinator.take_residency_updates([request.request_id]) + assert update[request.request_id].pages == lost + assert update[request.request_id].block_ids[0] == [0] * len(lost) + + +def test_hisparse_residency_update_sends_only_changed_pages(): + """Losing one early page must not resend the pages after it: per-step + residency updates would otherwise grow with the context length.""" + manager = make_hisparse_kv_cache_manager(32, 16) + tokens = list(range(6 * HISPARSE_BLOCK_SIZE)) + request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) + coordinator = get_hisparse_coordinator(manager) + for hot_manager in coordinator.hot_managers: + hot_manager.require_hot(request.request_id) + assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) + coordinator.take_residency_updates([request.request_id]) + pool = manager.block_pool + num_never_used = pool.get_num_free_blocks() + # Pages unpin as their host copies complete, so page 0 is reused first. + _publish_hisparse_pages(manager) + pool.get_new_blocks(num_never_used + 1) + + update = coordinator.take_residency_updates([request.request_id]) + + assert update[request.request_id].pages == [0] + assert update[request.request_id].block_ids == ([0],) + + +def test_hisparse_readmitted_request_resends_its_full_residency(): + """A preempted request may resume in a worker state row holding another + request's residency, so its first update must cover every page again.""" + manager = make_hisparse_kv_cache_manager(32, 16) + coordinator = get_hisparse_coordinator(manager) + tokens = list(range(2 * HISPARSE_BLOCK_SIZE)) + request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert _allocate_scheduled(manager, request, len(tokens)) is not None + assert request.request_id in coordinator.take_residency_updates( + [request.request_id] + ) + manager.free(request) + + assert _allocate_scheduled(manager, request, len(tokens)) is not None + update = coordinator.take_residency_updates([request.request_id]) + assert update[request.request_id].pages == [0, 1] + assert update[request.request_id].block_ids[0] == _resident_block_ids( + manager, request.request_id + ) def test_hisparse_host_prefix_can_be_completed_by_indexer_offload(): @@ -416,9 +476,7 @@ def test_hisparse_host_prefix_can_be_completed_by_indexer_offload(): source, indexer, resident, hot = manager.get_blocks(resumed.request_id).blocks assert len(source) == len(indexer) == len(resident) == 4 assert len(hot) == 2 - assert not any(block.is_null for block in resident[:2]) - assert not resident[2].is_null - assert not resident[3].is_null + assert 0 not in _resident_block_ids(manager, resumed.request_id) coordinator = get_hisparse_coordinator(manager) assert not coordinator.build_offload_command().page_transfers coordinator.finish_host_import(resumed.request_id, failed=False) @@ -666,9 +724,10 @@ def test_hisparse_keeps_resident_pages_until_hot_buffer_is_allocated(enable_cach pool.free_blocks(held) assert _allocate_scheduled(manager, first, num_new_tokens=1) is not None assert all(m.has_hot("first") for m in coordinator.hot_managers) + coordinator.take_residency_updates(["first"]) pool.get_new_blocks(pool.get_num_free_blocks()) assert not coordinator.all_context_pages_resident(scheduled) - assert "first" in coordinator.take_block_table_updates() + assert "first" in coordinator.take_residency_updates(["first"]) def test_hisparse_full_pool_keeps_pages_pinned_until_preemption(): @@ -1204,11 +1263,7 @@ def test_hisparse_prefix_hit_adopts_gpu_shadow_pages(): num_new_computed_tokens=num_computed, new_computed_blocks=computed, ) - resident_blocks = manager.get_blocks("resumed").blocks[2] - assert [block.block_id for block in resident_blocks[:3]] == ( - original_resident_ids[:3] - ) - assert not any(block.is_null for block in resident_blocks[:3]) + assert _resident_block_ids(manager, "resumed")[:3] == original_resident_ids[:3] assert get_hisparse_coordinator(manager).all_context_pages_resident( ((resumed.request_id, num_computed, len(tokens) - num_computed),) ) @@ -1249,17 +1304,13 @@ def test_hisparse_prefix_hit_under_pressure_adopts_surviving_copies( new_computed_blocks=computed, ) - resumed_blocks = manager.get_blocks("resumed").blocks - assert [block.block_id for block in resumed_blocks[2][:3]] == [ + resident_ids = _resident_block_ids(manager, "resumed") + assert resident_ids[:3] == [ copy_ids[page] if page in adopted_pages else pool.null_block.block_id for page in range(3) ] - gpu_ids = [ - block.block_id - for group in resumed_blocks[1:] - for block in group - if not block.is_null - ] + _, indexer, _, hot = manager.get_block_ids("resumed") + gpu_ids = [block_id for block_id in indexer + resident_ids + hot if block_id] assert len(gpu_ids) == len(set(gpu_ids)) @@ -1488,8 +1539,7 @@ def test_hisparse_reused_copy_is_not_adopted(): num_new_computed_tokens=num_computed, new_computed_blocks=computed, ) - resident_blocks = manager.get_blocks("resumed").blocks[2] - assert all(block.is_null for block in resident_blocks[:3]) + assert _resident_block_ids(manager, "resumed")[:3] == [0, 0, 0] def test_hisparse_external_import_uses_hard_gpu_footprint(): diff --git a/tests/v1/kv_connector/unit/test_hisparse_connector.py b/tests/v1/kv_connector/unit/test_hisparse_connector.py index f9598a24f82e..9ae0d3bc8eb1 100644 --- a/tests/v1/kv_connector/unit/test_hisparse_connector.py +++ b/tests/v1/kv_connector/unit/test_hisparse_connector.py @@ -28,10 +28,14 @@ from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.hisparse import runtime as runtime_module from vllm.v1.hisparse.coordinator import get_hisparse_coordinator -from vllm.v1.hisparse.runtime import HiSparseCacheHandle +from vllm.v1.hisparse.runtime import ( + HiSparseCacheHandle, + update_hisparse_residency, +) from vllm.v1.hisparse.types import ( SparseKVOffloadCommand, SparseKVPageTransfer, + SparseKVResidencyUpdate, SparseKVRowMirror, ) from vllm.v1.metrics.stats import KVCacheEvictionEvent @@ -183,7 +187,7 @@ def test_full_graph_step_prepares_host_mirror_outside_model(): connector = object.__new__(HiSparseConnector) connector.connector_worker = worker connector._get_connector_metadata = MagicMock( - return_value=HiSparseConnectorMetadata(None, (), (), {}, True) + return_value=HiSparseConnectorMetadata(None, (), (), {}, True, {}) ) req_id_per_token = torch.tensor([0, 1], dtype=torch.int32) attn_metadata = SimpleNamespace( @@ -321,6 +325,7 @@ def swap_blocks_batch(src, dst, sizes): ) }, True, + {}, ) ) # A verification step: four query tokens of one request fill the page. @@ -354,7 +359,7 @@ def swap_blocks_batch(src, dst, sizes): # The next step's start mirrors the drafter's rows, then hands the page over. compute_stream.wait_event.reset_mock() connector._get_connector_metadata.return_value = HiSparseConnectorMetadata( - None, (), (), {}, True + None, (), (), {}, True, {} ) connector.start_load_kv( SimpleNamespace(), @@ -397,7 +402,8 @@ def swap_blocks_batch(src, dst, sizes): def test_scheduled_prefix_hit_publishes_adopted_copies(): - """Copies adopted after scheduling must reach the worker's block table.""" + """Copies adopted after scheduling reach the worker as a residency update, + leaving the block-table row the scheduler output carries untouched.""" from tests.v1.core.test_prefix_caching import ( HISPARSE_BLOCK_SIZE, _allocate_scheduled, @@ -443,7 +449,39 @@ def test_scheduled_prefix_hit_publishes_adopted_copies(): num_scheduled_tokens={resumed.request_id: num_new_tokens}, ) - scheduler.build_connector_meta(scheduler_output) + core_row = list(scheduler_output.scheduled_new_reqs[0].block_ids[2]) + + metadata = scheduler.build_connector_meta(scheduler_output) + + update = metadata.residency_updates[resumed.request_id] + assert update.pages == [0, 1, 2, 3] + assert update.block_ids[0][:3] == copy_ids[:3] + assert manager.get_block_ids(resumed.request_id)[2] == core_row + assert core_row[:3] == [0, 0, 0] + + +def test_residency_updates_persist_by_state_row(): + """Updates name requests by batch row but are stored by state row. + + An update for a lost page and an appended page must leave the other pages + and the other request's row intact across steps that reorder the batch. + """ + table = torch.zeros((4, 2, 4), dtype=torch.int32) + update_hisparse_residency( + table, + { + "a": SparseKVResidencyUpdate([0, 1, 2], ([1, 2, 3], [11, 12, 13])), + "b": SparseKVResidencyUpdate([0, 1], ([4, 5], [14, 15])), + }, + ["a", "b"], + torch.tensor([2, 0], dtype=torch.int32), + ) + update_hisparse_residency( + table, + {"a": SparseKVResidencyUpdate([1, 3], ([0, 6], [0, 16]))}, + ["b", "a"], + torch.tensor([0, 2], dtype=torch.int32), + ) - resident_ids = scheduler_output.block_table_updates[resumed.request_id][2] - assert resident_ids[:3] == copy_ids[:3] + assert table[[0, 2], 0].tolist() == [[4, 5, 0, 0], [1, 0, 3, 6]] + assert table[[0, 2], 1].tolist() == [[14, 15, 0, 0], [11, 0, 13, 16]] diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index 98395e90c4dd..84dd944e1d22 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -1068,6 +1068,7 @@ def test_hisparse_step_waits_for_previous_host_write(monkeypatch, is_host_writer row_mirrors={}, all_context_pages_resident=True, row_mirrors_from_resident=False, + residency_updates={}, ), None, ) @@ -1079,6 +1080,7 @@ def test_hisparse_step_waits_for_previous_host_write(monkeypatch, is_host_writer row_mirrors={}, all_context_pages_resident=True, row_mirrors_from_resident=False, + residency_updates={}, ), None, ) @@ -1152,6 +1154,7 @@ def test_hisparse_empty_step_does_not_replay_stale_host_mirror(monkeypatch): row_mirrors={}, all_context_pages_resident=True, row_mirrors_from_resident=False, + residency_updates={}, ), None, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py index 863f52e02514..4f17f1e737e0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py @@ -36,7 +36,11 @@ from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.hisparse.coordinator import get_hisparse_coordinator from vllm.v1.hisparse.layout import get_hisparse_steady_state_concurrency -from vllm.v1.hisparse.types import SparseKVOffloadCommand, SparseKVRowMirror +from vllm.v1.hisparse.types import ( + SparseKVOffloadCommand, + SparseKVResidencyUpdate, + SparseKVRowMirror, +) from vllm.v1.outputs import KVConnectorOutput if TYPE_CHECKING: @@ -56,6 +60,7 @@ class HiSparseConnectorMetadata(KVConnectorMetadata): source_block_ids: tuple[int, ...] row_mirrors: dict[str, tuple[SparseKVRowMirror, ...]] all_context_pages_resident: bool + residency_updates: dict[str, SparseKVResidencyUpdate] @dataclass @@ -143,7 +148,7 @@ def build_connector_meta( scheduler_output.num_scheduled_tokens.items() ) ) - # Runs first: it plans this step's transfers and table updates. + # Runs first: it plans this step's transfers and residency changes. self.coordinator.advance_scheduled( ( request_id, @@ -151,9 +156,6 @@ def build_connector_meta( ) for request_id, start, count in scheduled_requests ) - scheduler_output.block_table_updates = ( - self.coordinator.take_block_table_updates() or None - ) command = self.coordinator.build_offload_command() host_block_copies = self.coordinator.take_host_block_copies() source_group_id = self.coordinator.host_group_id @@ -192,6 +194,9 @@ def build_connector_meta( tuple(source_block_ids), row_mirrors, self.coordinator.all_context_pages_resident(scheduled_requests), + self.coordinator.take_residency_updates( + scheduler_output.num_scheduled_tokens + ), ) def update_connector_output(self, connector_output: KVConnectorOutput) -> None: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py index b9dd01f05420..bbccecd65784 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py @@ -26,7 +26,11 @@ from vllm.utils.torch_utils import current_stream from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.hisparse.layout import HISPARSE_HOT_SUFFIX -from vllm.v1.hisparse.runtime import HiSparseCacheHandle, release_pinned_state +from vllm.v1.hisparse.runtime import ( + HiSparseCacheHandle, + release_pinned_state, + update_hisparse_residency, +) from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror from vllm.v1.kv_cache_interface import ( HiSparseHotSpec, @@ -334,6 +338,10 @@ def initialize( "HiSparse request-state mapping does not match max_num_seqs." ) self.hot_backing = hot_backing + self.residency: torch.Tensor | None = next( + (cache.residency for cache in cache_handles if cache.residency is not None), + None, + ) self._pending_invalid_block_ids: list[int] = [] # Destination block ids of host copies this worker has run. self._completed_host_copy_dst_ids: list[int] = [] @@ -437,6 +445,15 @@ def start_step( self._pending_invalid_block_ids.extend(metadata.source_block_ids) if request_state_indices is not None: self.set_request_state_indices(request_state_indices) + if metadata.residency_updates: + assert self.residency is not None + assert request_state_indices is not None and request_ids is not None + update_hisparse_residency( + self.residency, + metadata.residency_updates, + request_ids, + request_state_indices, + ) def _clear_forward_mirror_state(self) -> None: self._per_layer_mirrored.clear() diff --git a/vllm/v1/attention/backends/mla/index_group.py b/vllm/v1/attention/backends/mla/index_group.py index 39f205b91d97..0dfda0ded9f0 100644 --- a/vllm/v1/attention/backends/mla/index_group.py +++ b/vllm/v1/attention/backends/mla/index_group.py @@ -243,7 +243,7 @@ def convert_logical_to_physical_topk( layer_index, logical_topk_indices, req_id_per_token, - leader.block_table, + leader.batch_block_table(), leader.view.block_size, block_stride_rows=leader.view.attention_block_stride, return_valid_counts=return_valid_counts, @@ -268,8 +268,11 @@ def stage_prefill_rows( assert staging_plan is not None resident_cache = None if cache.view is not None and cache.block_table is not None: + state_indices = cache.runtime.request_state_indices + assert state_indices is not None staging_plan.ensure_gpu_sources( - cache.block_table[attn_metadata.num_decodes :], + cache.block_table, + state_indices[attn_metadata.num_decodes :], cache.view.block_size, ) resident_cache = cache.view.cache @@ -291,8 +294,11 @@ def _prefill_gather_plan( plan = prefill.host_staging_plan if prefill is not None else None assert plan is not None assert cache.view is not None and cache.block_table is not None + state_indices = cache.runtime.request_state_indices + assert state_indices is not None plan.ensure_gpu_sources( - cache.block_table[attn_metadata.num_decodes :], + cache.block_table, + state_indices[attn_metadata.num_decodes :], cache.view.block_size, ) assert plan.gpu_row_ids is not None diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index 9d16d2f710a3..a5af5a3d5755 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -304,9 +304,6 @@ class SchedulerOutput: # CoW copies to apply after zeroing new blocks and before forward. kv_cache_block_copies: list[KVCacheBlockCopy] | None = None - # Complete block-table rows that replace incrementally appended block IDs. - block_table_updates: dict[str, tuple[list[int], ...]] | None = None - # Scheduler-local; always None by the time this reaches a worker. kv_connector_block_state: KVConnectorBlockState | None = None diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index eed15427b2dd..77c6a4a5f750 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -2503,6 +2503,10 @@ class HiSparseResidentManager(_HiSparseAuxiliaryManager): max_admission_blocks_per_request: int + def __init__(self, kv_cache_spec: HiSparseResidentSpec, **kwargs) -> None: + super().__init__(kv_cache_spec, **kwargs) + self.residency_changes: dict[str, dict[int, KVCacheBlock]] = {} + def get_num_blocks_to_allocate( self, request_id: str, @@ -2546,7 +2550,7 @@ def allocate_external_computed_blocks( tail_page = (num_tokens - 1) // self.block_size if len(blocks) <= tail_page: blocks.extend([self._null_block] * (tail_page + 1 - len(blocks))) - if blocks[tail_page].is_null: + if self.get_resident_page(request_id, tail_page) is None: blocks[tail_page] = self.block_pool.get_new_blocks(1)[0] def add_local_computed_blocks( @@ -2580,25 +2584,55 @@ def allocate_new_blocks( def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: assert self.coordinator is not None self.coordinator.free(request_id) - return super().pop_blocks_for_free(request_id) + blocks = self.get_residency_row(request_id) + self.residency_changes.pop(request_id, None) + super().pop_blocks_for_free(request_id) + return blocks def adopt_resident_page( self, request_id: str, block_idx: int, block: KVCacheBlock ) -> bool: """Point a null prefix page at a pinned GPU copy of its contents.""" blocks = self.req_to_blocks.get(request_id) - if blocks is None or block_idx >= len(blocks) or not blocks[block_idx].is_null: + if blocks is None or block_idx >= len(blocks): + return False + if self.get_resident_page(request_id, block_idx) is not None: return False - blocks[block_idx] = block + self.residency_changes.setdefault(request_id, {})[block_idx] = block return True + def drop_resident_page( + self, request_id: str, block_idx: int + ) -> KVCacheBlock | None: + """Stop reading a page from the GPU; return the block it was read from.""" + block = self.get_resident_page(request_id, block_idx) + if block is not None: + changes = self.residency_changes.setdefault(request_id, {}) + changes[block_idx] = self._null_block + return block + def get_resident_page(self, request_id: str, block_idx: int) -> KVCacheBlock | None: blocks = self.req_to_blocks.get(request_id) if blocks is None or block_idx >= len(blocks): return None block = blocks[block_idx] + changes = self.residency_changes.get(request_id) + if changes: + block = changes.get(block_idx, block) return None if block.is_null else block + def get_residency_row(self, request_id: str, start: int = 0) -> list[KVCacheBlock]: + """The request's row from ``start`` with residency changes applied. Unlike + ``get_resident_page``, host-only pages stay as the null block.""" + blocks = self.req_to_blocks.get(request_id, [])[start:] + changes = self.residency_changes.get(request_id) + if not changes: + return blocks + return [ + changes.get(block_idx, block) + for block_idx, block in enumerate(blocks, start) + ] + def get_manager_for_kv_cache_spec( kv_cache_spec: KVCacheSpec, diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index 60c60a3ff504..1420a45396cd 100644 --- a/vllm/v1/hisparse/binding.py +++ b/vllm/v1/hisparse/binding.py @@ -214,10 +214,26 @@ def bind_hisparse_kv_caches( for tensor_config in kv_cache_config.kv_cache_tensors for layer_index, name in enumerate(tensor_config.layers) } - resident_source_index = 0 - for group_id, group in enumerate(kv_cache_config.kv_cache_groups): - if not isinstance(group.kv_cache_spec, HiSparseResidentSpec): - continue + resident_groups = [ + (group_id, group) + for group_id, group in enumerate(kv_cache_config.kv_cache_groups) + if isinstance(group.kv_cache_spec, HiSparseResidentSpec) + ] + residency: torch.Tensor | None = None + if resident_groups: + table_shapes = { + block_tables.input_block_tables[group_id].shape + for group_id, _ in resident_groups + } + assert len(table_shapes) == 1 + max_num_reqs, max_num_pages = table_shapes.pop() + residency = torch.zeros( + (max_num_reqs, len(resident_groups), max_num_pages), + dtype=torch.int32, + device=block_tables.input_block_tables[resident_groups[0][0]].device, + ) + for resident_source_index, (group_id, group) in enumerate(resident_groups): + assert residency is not None for cache_name in group.layer_names: assert cache_name.endswith(HISPARSE_RESIDENT_SUFFIX) layer_name = cache_name[: -len(HISPARSE_RESIDENT_SUFFIX)] @@ -230,13 +246,13 @@ def bind_hisparse_kv_caches( block_stride=tensor_config.block_stride, num_blocks=kv_cache_config.num_blocks, block_size=group.kv_cache_spec.block_size, - block_table=block_tables.input_block_tables[group_id], + block_table=residency[:, resident_source_index], slot_mapping=block_tables.slot_mappings[group_id], ) + cache_handle.residency = residency assert cache_handle.view is not None kv_caches[cache_name] = cache_handle.view.cache cache_handle.runtime.resident_source_index = resident_source_index - resident_source_index += 1 hot_backing: torch.Tensor | None = None cache_handles: list[HiSparseCacheHandle] = [] diff --git a/vllm/v1/hisparse/coordinator.py b/vllm/v1/hisparse/coordinator.py index 390e0d27ced2..374f100bc889 100644 --- a/vllm/v1/hisparse/coordinator.py +++ b/vllm/v1/hisparse/coordinator.py @@ -22,6 +22,7 @@ ACTIVE_TAIL_PAGES, SparseKVOffloadCommand, SparseKVPageTransfer, + SparseKVResidencyUpdate, SparseKVRowMirror, ) from vllm.v1.kv_cache_interface import ( @@ -65,6 +66,10 @@ class _HiSparseRequestState: unpinned_pages: set[int] = field(default_factory=set) # Prefix pages whose GPU copies are adopted after the admitting allocation. pages_to_adopt: int = 0 + # The worker holds this request's residency for every page below + # synced_pages, except the stale pages changed since it was sent. + synced_pages: int = 0 + stale_pages: set[int] = field(default_factory=set) @dataclass @@ -86,9 +91,10 @@ class HiSparseCoordinator: GPU-resident pages are a write-back cache of the host tier. Once a page's host copy is durable and its owner can read from host (it has, or has asked for, a hot region), the page's allocation reference is released with - ``BlockPool.unpin_blocks``: the block stays in the block table and keeps - being read, but the pool counts it as free and may hand it out, at which - point the owner's page is nulled and its block table republished. + ``BlockPool.unpin_blocks``: the block keeps being read, but the pool counts + it as free and may hand it out, at which point the owner stops reading it. + Residency changes reach the worker through ``take_residency_updates``, not + the block table, which stays append-only. """ def __init__( @@ -167,7 +173,6 @@ def __init__( self._retained_copies: dict[int, tuple[KVCacheBlock, KVCacheBlock]] = {} # request -> prefix length an external load is still filling in. self._pending_imports: dict[str, int] = {} - self.block_table_updates: set[str] = set() # Requests whose resident pages were all non-null when last scanned. # Pages only become null at admission or import (before the first # scan) and when the pool reuses an unpinned page (_lose_page). @@ -269,7 +274,8 @@ def _adopt_copies( if manager.adopt_resident_page(request_id, host_idx, block): manager.block_pool.touch([block]) state.pinned_clean.add(host_idx) - self.block_table_updates.add(request_id) + if host_idx < state.synced_pages: + state.stale_pages.add(host_idx) def _record_copies(self, request_id: str, num_computed_tokens: int) -> None: """Index the GPU copies of just-published host blocks for later hits.""" @@ -342,14 +348,11 @@ def _resident_page_blocks( ) -> list[KVCacheBlock] | None: blocks: list[KVCacheBlock] = [] for manager in self.resident_managers: - req_blocks = manager.req_to_blocks.get(request_id) - if ( - req_blocks is None - or page_idx >= len(req_blocks) - ACTIVE_TAIL_PAGES - or req_blocks[page_idx].is_null - ): + num_pages = len(manager.req_to_blocks.get(request_id, ())) + block = manager.get_resident_page(request_id, page_idx) + if page_idx >= num_pages - ACTIVE_TAIL_PAGES or block is None: return None - blocks.append(req_blocks[page_idx]) + blocks.append(block) return blocks def _unpin_page( @@ -392,16 +395,15 @@ def _lose_page(self, request_id: str, page_idx: int) -> None: if state is None: return for manager in self.resident_managers: - blocks = manager.req_to_blocks.get(request_id) - if blocks is None or page_idx >= len(blocks) or blocks[page_idx].is_null: + block = manager.drop_resident_page(request_id, page_idx) + if block is None: continue - block = blocks[page_idx] - blocks[page_idx] = manager._null_block owners = self._owners.get(block.block_id) if owners is not None: owners.discard((request_id, page_idx)) state.unpinned_pages.discard(page_idx) - self.block_table_updates.add(request_id) + if page_idx < state.synced_pages: + state.stale_pages.add(page_idx) for hot_manager in self.hot_managers: hot_manager.require_hot(request_id) @@ -627,11 +629,10 @@ def build_row_mirrors( mirrors: list[SparseKVRowMirror] = [] for request_id, num_computed_tokens, num_scheduled_tokens in requests: host_blocks = self.host_manager.req_to_blocks.get(request_id) - resident_blocks = [ - manager.req_to_blocks.get(request_id) + if host_blocks is None or any( + request_id not in manager.req_to_blocks for manager in self.resident_managers - ] - if host_blocks is None or any(blocks is None for blocks in resident_blocks): + ): continue token_position = num_computed_tokens end_position = token_position + num_scheduled_tokens @@ -644,14 +645,12 @@ def build_row_mirrors( # window only ever moves forward, so dropping them here would # leave their host rows permanently stale. source_starts = [] - for blocks in resident_blocks: - assert blocks is not None - if page_idx >= len(blocks) or blocks[page_idx].is_null: + for manager in self.resident_managers: + block = manager.get_resident_page(request_id, page_idx) + if block is None: break - source_starts.append( - blocks[page_idx].block_id * block_size + row_offset - ) - if len(source_starts) != len(resident_blocks): + source_starts.append(block.block_id * block_size + row_offset) + if len(source_starts) != len(self.resident_managers): token_position += num_rows continue host_block_idx = page_idx @@ -693,21 +692,53 @@ def all_context_pages_resident( if any( block.is_null for manager in self.resident_managers - for block in manager.req_to_blocks.get(request_id, ()) + for block in manager.get_residency_row(request_id) ): return False self._resident_requests.add(request_id) return True - def take_block_table_updates(self) -> dict[str, tuple[list[int], ...]]: - updates = { - request_id: tuple( - [block.block_id for block in manager.req_to_blocks.get(request_id, [])] - for manager in self.managers + def take_residency_updates( + self, request_ids: Iterable[str] + ) -> dict[str, SparseKVResidencyUpdate]: + """Resident block ids the worker lacks for the scheduled requests: + stale pages plus pages appended since the last update. + + Requests that are not scheduled keep their changes until they are: + only a scheduled request's residency is read. + """ + updates: dict[str, SparseKVResidencyUpdate] = {} + if not self.resident_managers: + return updates + first = self.resident_managers[0] + for request_id in request_ids: + state = self._get_request_state(request_id) + num_pages = len(first.req_to_blocks.get(request_id, ())) + start = state.synced_pages + if not state.stale_pages and start >= num_pages: + continue + stale_pages = sorted(state.stale_pages) + block_ids: list[list[int]] = [] + for manager in self.resident_managers: + null_block_id = manager.block_pool.null_block.block_id + group_block_ids = [ + null_block_id if block is None else block.block_id + for block in ( + manager.get_resident_page(request_id, page_idx) + for page_idx in stale_pages + ) + ] + group_block_ids.extend( + block.block_id + for block in manager.get_residency_row(request_id, start) + ) + block_ids.append(group_block_ids) + updates[request_id] = SparseKVResidencyUpdate( + pages=[*stale_pages, *range(start, num_pages)], + block_ids=tuple(block_ids), ) - for request_id in self.block_table_updates - } - self.block_table_updates.clear() + state.stale_pages.clear() + state.synced_pages = num_pages return updates def build_offload_command(self) -> SparseKVOffloadCommand | None: @@ -846,9 +877,7 @@ def free(self, request_id: str) -> None: # planned, so its copies do not prove that the KV was computed. state.publication = None for manager in self.resident_managers: - blocks = manager.req_to_blocks.get(request_id) - if not blocks: - continue + blocks = manager.get_residency_row(request_id) # Tail first, so a prefix loses its tail pages before its head. for page_idx in reversed(range(len(blocks))): block = blocks[page_idx] @@ -865,7 +894,7 @@ def free(self, request_id: str) -> None: manager.block_pool.unpin_blocks([block], self._on_block_reused) else: continue - blocks[page_idx] = manager._null_block + manager.drop_resident_page(request_id, page_idx) def get_hisparse_coordinator( diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index 1f7a5ad5ca3c..f1b0fa688eda 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -7,10 +7,11 @@ import math import mmap import time -from collections.abc import Callable +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, TypeAlias +import numpy as np import psutil import torch @@ -21,8 +22,9 @@ from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv, round_up from vllm.utils.mem_utils import get_max_shared_memory_bytes -from vllm.utils.torch_utils import current_stream +from vllm.utils.torch_utils import async_tensor_h2d, current_stream from vllm.v1.attention.backend import max_decode_query_len +from vllm.v1.hisparse.types import SparseKVResidencyUpdate from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor @@ -499,16 +501,19 @@ class HiSparsePrefillStagingPlan: def ensure_gpu_sources( self, - resident_block_table: torch.Tensor, + resident_state_rows: torch.Tensor, + state_indices: torch.Tensor, resident_block_size: int, ) -> None: """Resolve which staged rows can be served from the resident cache. - Computed once per plan (the resident block table is shared by every - layer in the group); non-null resident pages become miss_mask=0 rows - gathered device-to-device by ``gather_prefill_cache``. + ``resident_state_rows`` is a resident group's persistent table by + request state row and ``state_indices`` holds each staged request's + state row. Computed once per plan and resident group (every layer in + the group shares the table); non-null resident pages become + miss_mask=0 rows gathered device-to-device by ``gather_prefill_cache``. """ - source_key = (resident_block_table.data_ptr(), resident_block_size) + source_key = (resident_state_rows.data_ptr(), resident_block_size) if self.gpu_source_key == source_key: return block_size = self.block_size @@ -519,8 +524,11 @@ def ensure_gpu_sources( host_ids = self.row_ids[0].view(num_unique, block_size)[:, 0] // block_size new_bt = self.block_table.to(torch.int64) num_rows, num_cols = new_bt.shape - if num_rows == 0 or resident_block_table.shape[0] < num_rows: + if num_rows == 0 or state_indices.shape[0] < num_rows: return + resident_block_table = resident_state_rows.index_select( + 0, state_indices[:num_rows].clamp(min=0) + ) # One representative (row, col) per unique host block: any request # referencing the block holds an equivalent (refcounted) resident view. flat_pos = torch.arange(num_rows * num_cols, device=device) @@ -1107,11 +1115,50 @@ def swap_in( return physical_topk_indices +def update_hisparse_residency( + residency: torch.Tensor, + updates: Mapping[str, SparseKVResidencyUpdate], + request_ids: Sequence[str], + request_state_indices: torch.Tensor, +) -> None: + """Apply scheduled requests' residency changes. + + ``residency`` is ``[state rows, resident groups, pages]``: the GPU block of + each page, or block 0 when the page is read from the host. + """ + if not updates: + return + batch_rows = {request_id: row for row, request_id in enumerate(request_ids)} + num_groups = residency.shape[1] + rows: list[int] = [] + pages: list[int] = [] + block_ids: list[list[int]] = [[] for _ in range(num_groups)] + for request_id, update in updates.items(): + rows.append(batch_rows[request_id]) + pages.extend(update.pages) + for group_block_ids, update_block_ids in zip( + block_ids, update.block_ids, strict=True + ): + group_block_ids.extend(update_block_ids) + host_values = np.empty((len(pages), 2 + num_groups), dtype=np.int32) + host_values[:, 0] = np.repeat( + rows, [len(update.pages) for update in updates.values()] + ) + host_values[:, 1] = pages + host_values[:, 2:] = np.asarray(block_ids, dtype=np.int32).T + values = async_tensor_h2d( + host_values, device=request_state_indices.device, dtype=torch.int32 + ) + state_rows = request_state_indices[values[:, 0]] + residency[state_rows, :, values[:, 1]] = values[:, 2:] + + class HiSparseCacheHandle: """Attention-facing handle for resident KV and sparse offload state.""" def __init__(self, runtime: HiSparseRuntime) -> None: self.view: PagedCacheView | None = None + self.residency: torch.Tensor | None = None self.block_table: torch.Tensor | None = None self.source_block_table: torch.Tensor | None = None self.slot_mapping: torch.Tensor | None = None @@ -1128,6 +1175,7 @@ def __init__(self, runtime: HiSparseRuntime) -> None: # Speculator layers write their rows after the target forward. self.draft_layer = False self.index_group_caches: list[HiSparseCacheHandle] = [self] + self._batch_block_table: torch.Tensor | None = None def prepare_group_for_batch(self, attn_metadata: Any | None) -> None: assert self.runtime.is_group_leader @@ -1136,6 +1184,7 @@ def prepare_group_for_batch(self, attn_metadata: Any | None) -> None: def _prepare_for_batch(self, attn_metadata: Any | None) -> None: self.dummy_batch = attn_metadata is None + self._batch_block_table = None self.runtime.begin_forward() self.num_actual_tokens = ( attn_metadata.num_actual_tokens if attn_metadata is not None else 0 @@ -1151,6 +1200,17 @@ def _prepare_for_batch(self, attn_metadata: Any | None) -> None: not self.decode_batch or self.runtime.eager_host_mirror ) + def batch_block_table(self) -> torch.Tensor: + """Resident rows by this step's batch row, for the prefill paths.""" + if self._batch_block_table is None: + assert self.block_table is not None + indices = self.runtime.request_state_indices + assert indices is not None + self._batch_block_table = self.block_table.index_select( + 0, indices.clamp(min=0) + ) + return self._batch_block_table + def write_target( self, num_input_rows: int, num_slot_rows: int ) -> tuple[torch.Tensor, torch.Tensor, int]: diff --git a/vllm/v1/hisparse/types.py b/vllm/v1/hisparse/types.py index db0ee06e9399..37b25003df25 100644 --- a/vllm/v1/hisparse/types.py +++ b/vllm/v1/hisparse/types.py @@ -28,6 +28,15 @@ class SparseKVRowMirror: num_rows: int +@dataclass(frozen=True) +class SparseKVResidencyUpdate: + """GPU block ids of some of a request's resident pages, per resident group. + A null block id means the page is read from the host.""" + + pages: list[int] + block_ids: tuple[list[int], ...] + + @dataclass class SparseKVOffloadCommand: """Opaque scheduler-to-worker command for one model step.""" diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index f39fd0a6e6a7..674cf23527d1 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1224,17 +1224,13 @@ def add_requests(self, scheduler_output: SchedulerOutput) -> None: def update_requests(self, scheduler_output: SchedulerOutput) -> None: # Add new blocks and update num_computed_tokens for the existing requests. reqs = scheduler_output.scheduled_cached_reqs - table_updates = scheduler_output.block_table_updates or {} - for req_id, block_ids in table_updates.items(): - if (idx := self.req_states.req_id_to_index.get(req_id)) is not None: - self.block_tables.append_block_ids(idx, block_ids, overwrite=True) num_computed_tokens_np = self.req_states.num_computed_tokens_np for req_id, num_computed_tokens, req_new_block_ids in zip( reqs.req_ids, reqs.num_computed_tokens, reqs.new_block_ids ): req_index = self.req_states.req_id_to_index[req_id] num_computed_tokens_np[req_index] = num_computed_tokens - if req_new_block_ids is not None and req_id not in table_updates: + if req_new_block_ids is not None: self.block_tables.append_block_ids( req_index, req_new_block_ids, overwrite=False ) From acec1b1d84c428ced19d53935979107eeab922a2 Mon Sep 17 00:00:00 2001 From: Guancheng Fu <110874468+gc-fu@users.noreply.github.com> Date: Sat, 10 Oct 2026 12:27:35 +0800 Subject: [PATCH 227/624] [Bugfix][XPU] Remove stale W4A8 weight after repacking (#57511) Signed-off-by: Guancheng Fu Co-authored-by: Kunshang Ji --- vllm/model_executor/kernels/linear/mixed_precision/xpu.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/xpu.py b/vllm/model_executor/kernels/linear/mixed_precision/xpu.py index c6545f312a91..c82e862e5870 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/xpu.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/xpu.py @@ -186,7 +186,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Free the original unpacked int8 weight (still registered as "weight") # to avoid double-storing both int8 [N, K] and int32 [N, K/8] in memory. - layer.register_parameter("weight", None) + delattr(layer, "weight") def apply_weights( self, From 5ca26880ad87fe5f55628cd0d09d8b8bfc0ba59b Mon Sep 17 00:00:00 2001 From: Puneet Sharma Date: Sat, 10 Oct 2026 10:09:58 +0530 Subject: [PATCH 228/624] Optimize ppc64le vLLM runtime image size (#48074) Signed-off-by: puneetsharma21 --- docker/Dockerfile.ppc64le | 29 +-- docker/build_vllm_ppc64le.sh | 434 +++++++++++++++++++++++++++-------- 2 files changed, 346 insertions(+), 117 deletions(-) diff --git a/docker/Dockerfile.ppc64le b/docker/Dockerfile.ppc64le index 68e6ab956377..00d24abbbb5c 100644 --- a/docker/Dockerfile.ppc64le +++ b/docker/Dockerfile.ppc64le @@ -1,5 +1,5 @@ # Base UBI image -ARG BASE_UBI_IMAGE_TAG=9.6-1754584681 +ARG BASE_UBI_IMAGE_TAG=9.8 ############################################################### # BUILDER STAGE # @@ -7,7 +7,7 @@ ARG BASE_UBI_IMAGE_TAG=9.6-1754584681 FROM registry.access.redhat.com/ubi9/ubi-minimal:${BASE_UBI_IMAGE_TAG} AS builder-base -ARG VLLM_VERSION="0.22.1" +ARG VLLM_VERSION="0.0.0.dev" ARG PYTHON_VERSION=3.12 ARG VLLM_TARGET_DEVICE=cpu @@ -45,10 +45,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ # copy vllm source code to build cache COPY . . -RUN --mount=type=cache,target=/root/.cache/uv \ - source /opt/rh/gcc-toolset-14/enable && \ - pip install -U uv - # build & install vLLM so that all transitive dependencies are build/downloaded into the uv cache RUN --mount=type=cache,target=/root/.cache/uv \ source /opt/rh/gcc-toolset-14/enable && \ @@ -63,19 +59,24 @@ RUN --mount=type=cache,target=/root/.cache/uv \ export DNNL_ARCH_OPT_FLAGS="-mcpu=power10 -mtune=power10" && \ export C_INCLUDE_PATH=/usr/local/include:$C_INCLUDE_PATH && \ export CPLUS_INCLUDE_PATH=/usr/local/include:$CPLUS_INCLUDE_PATH && \ - uv pip install 'setuptools>=78.1.1' && \ export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/OpenBLAS/lib/:/usr/local/lib64:/usr/local/lib && \ export LIBGOMP=/opt/rh/gcc-toolset-14/root/usr/lib/gcc/ppc64le-redhat-linux/14/libgomp.so && \ - export CMAKE_LIBRARY_PATH=$(dirname $LIBGOMP):${CMAKE_LIBRARY_PATH} && \ export LIBRARY_PATH=$(dirname $LIBGOMP):${LIBRARY_PATH} && \ export LD_LIBRARY_PATH=$(dirname $LIBGOMP):${LD_LIBRARY_PATH} && \ - echo "LIBGOMP=${LIBGOMP}" && \ find /root/.cache/uv -name "*.whl" && \ SETUPTOOLS_SCM_PRETEND_VERSION="$VLLM_VERSION" uv build \ --wheel --out-dir ${WHEEL_DIR} --no-build-isolation && \ - uv pip install "$(echo ${WHEEL_DIR}/vllm*.whl)[tensorizer]" --refresh + uv pip install "$(echo ${WHEEL_DIR}/vllm*.whl)[tensorizer]" --refresh && \ + uv pip install "setuptools==80.10.2" && \ + uv pip uninstall -y pip cmake meson-python scikit-build-core maturin build wheel auditwheel pybind11 nanobind cython setuptools-rust && \ + # Remove packaged test suites + find ${VIRTUAL_ENV}/lib64/python${PYTHON_VERSION}/site-packages \ + -type d \( -name tests -o -name test \) -exec rm -rf {} + && \ + find ${VIRTUAL_ENV} -name "__pycache__" -type d -exec rm -rf {} + && \ + find ${VIRTUAL_ENV} -name "*.pyc" -delete && \ + rm -f ${VIRTUAL_ENV}/lib64/python${PYTHON_VERSION}/site-packages/xgrammar/lib/libxgrammar.a ############################################################### # FINAL VLLM IMAGE STAGE # @@ -94,7 +95,6 @@ ENV PKG_CONFIG_PATH=${PCP_DIR}/usr/lib64/pkgconfig:/usr/local/lib/pkgconfig/ ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH" ENV LD_LIBRARY_PATH=${PCP_DIR}/usr/lib64:${PCP_DIR}/usr/lib:${VIRTUAL_ENV}/lib64/python${PYTHON_VERSION}/site-packages/torch/lib:/usr/local/lib:$LD_LIBRARY_PATH:/usr/local/lib64:/usr/lib64:/usr/lib ENV UV_LINK_MODE=copy -ARG VLLM_VERSION="0.22.1" ARG UV_EXTRA_INDEX_URL="https://wheels.developerfirst.ibm.com/ppc64le/linux/+simple/" ENV UV_EXTRA_INDEX_URL=${UV_EXTRA_INDEX_URL} ENV UV_INDEX_STRATEGY=first-match @@ -103,7 +103,7 @@ ENV UV_INDEX_STRATEGY=first-match RUN --mount=type=cache,target=/root/.cache/uv \ rpm -ivh https://dl.fedoraproject.org/pub/epel/epel-release-latest-9.noarch.rpm && \ microdnf install --nodocs -y \ - libomp libicu tar autoconf automake libtool findutils openssl numactl numactl-devel \ + libomp libicu tar autoconf automake libtool findutils openssl openblas-devel numactl numactl-devel \ pkgconfig xsimd gcc-toolset-14 libsndfile \ libtiff libjpeg openjpeg2 zlib zeromq \ freetype lcms2 libwebp tcl tk utf8proc \ @@ -113,10 +113,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ && microdnf update -y \ && microdnf clean all -# The `lscpu` command was added as a requirement in part of https://github.com/vllm-project/vllm/pull/21032, so installing it. -RUN microdnf install --nodocs -y util-linux && \ - microdnf clean all - COPY --from=builder-base /usr/lib64/libprotobuf.so.25 /usr/lib64/ COPY --from=builder-base /usr/lib64/libprotobuf.so.25.0.0 /usr/lib64/ @@ -154,4 +150,3 @@ USER 2000 ENTRYPOINT ["vllm", "serve"] - diff --git a/docker/build_vllm_ppc64le.sh b/docker/build_vllm_ppc64le.sh index f2fbad8e2199..e49edbd7c687 100644 --- a/docker/build_vllm_ppc64le.sh +++ b/docker/build_vllm_ppc64le.sh @@ -13,7 +13,6 @@ cd "$REPO_ROOT" ######################################## IBM_DEVPI_URL=${IBM_DEVPI_URL:-"https://wheels.developerfirst.ibm.com/ppc64le/linux/+simple/"} -RHOAI_INDEX_URL=${RHOAI_INDEX_URL:-"https://console.redhat.com/api/pypi/public-rhai/rhoai/3.4/cpu-ubi9/simple/"} ######################################## # wheel dir @@ -39,16 +38,65 @@ try_install_from_devpi() { ######################################## cd "$REPO_ROOT" TORCH_VERSION=${TORCH_VERSION:-$(grep -E '^torch==.+==\s*"ppc64le"' requirements/cpu.txt | grep -Eo '\b[0-9\.]+\b' || true)} -TORCH_VERSION=${TORCH_VERSION:-2.11.0} - -TORCHVISION_VERSION=${TORCHVISION_VERSION:-0.26.0} -TORCHAUDIO_VERSION=${TORCHAUDIO_VERSION:-${TORCH_VERSION}} +TORCH_VERSION=${TORCH_VERSION:-2.13.0} + +# Parse version parts +TORCH_MAJOR=$(echo "${TORCH_VERSION}" | cut -d. -f1) +TORCH_MINOR=$(echo "${TORCH_VERSION}" | cut -d. -f2) +TORCH_PATCH=$(echo "${TORCH_VERSION}" | cut -d. -f3 | cut -d+ -f1) +TORCH_PATCH=${TORCH_PATCH:-0} + +# 1. Calculate torchvision (always minor + 15) +VISION_MINOR=$((TORCH_MINOR + 15)) +TORCHVISION_VERSION="0.${VISION_MINOR}.${TORCH_PATCH}" + +# 2. Calculate torchaudio (minor - 2 for >= 2.13, otherwise same as torch) +if [ "${TORCH_MAJOR}" -eq 2 ] && [ "${TORCH_MINOR}" -ge 13 ]; then + AUDIO_MINOR=$((TORCH_MINOR - 2)) + TORCHAUDIO_VERSION="${TORCH_MAJOR}.${AUDIO_MINOR}.${TORCH_PATCH}" +else + TORCHAUDIO_VERSION="${TORCH_VERSION}" +fi export TORCH_VERSION export TORCHVISION_VERSION export TORCHAUDIO_VERSION -export OPENCV_VERSION=${OPENCV_VERSION:-4.13.0.92} -export XGRAMMAR_VERSION=${XGRAMMAR_VERSION:-0.2.1} + +echo "Resolved PyTorch versions:" +echo " torch: ${TORCH_VERSION}" +echo " torchvision: ${TORCHVISION_VERSION}" +echo " torchaudio: ${TORCHAUDIO_VERSION}" + +OPENCV_REQUIREMENT=$(grep -E '^opencv-python-headless[[:space:]]*(==|>=|<=|~=)' \ + requirements/common.txt | + sed -E 's/^opencv-python-headless[[:space:]]*//' | + sed -E 's/[[:space:]]*#.*$//' | + tr -d ' ' | + head -n1) + +export OPENCV_REQUIREMENT +echo "OpenCV requirement: ${OPENCV_REQUIREMENT}" + + +XGRAMMAR_VERSION=$( + sed -nE 's/^xgrammar[[:space:]]*==[[:space:]]*([^,;[:space:]]+).*/\1/p' requirements/common.txt | + head -n1 +) + +if [ -z "${XGRAMMAR_VERSION}" ]; then + XGRAMMAR_VERSION=$( + sed -nE 's/^xgrammar[[:space:]]*>=[[:space:]]*([^,;[:space:]]+).*/\1/p' requirements/common.txt | + head -n1 + ) +fi + +export XGRAMMAR_VERSION + +if [ -z "${XGRAMMAR_VERSION}" ]; then + echo "xgrammar version not found in requirements/common.txt" + exit 1 +fi +echo "XGRAMMAR_VERSION=${XGRAMMAR_VERSION}" ######################################## # install system dependencies @@ -56,15 +104,20 @@ export XGRAMMAR_VERSION=${XGRAMMAR_VERSION:-0.2.1} rpm -ivh https://dl.fedoraproject.org/pub/epel/epel-release-latest-9.noarch.rpm || true +# Dynamically match openssl-devel version with installed openssl-libs +OPENSSL_DEVEL="openssl-devel-$(rpm -q --queryformat '%{VERSION}-%{RELEASE}' openssl-libs 2>/dev/null || echo 'openssl-devel')" + microdnf install -y \ + "${OPENSSL_DEVEL}" \ python3.12 python3.12-devel python3.12-pip gcc \ git jq gcc-toolset-14 gcc-toolset-14-libatomic-devel \ - automake libtool clang-devel openssl-devel \ + automake libtool clang-devel \ harfbuzz-devel kmod lcms2-devel libimagequant-devel libjpeg-turbo-devel \ llvm15-devel libraqm-devel libtiff-devel libwebp-devel libxcb-devel \ ninja-build openjpeg2-devel pkgconfig \ - tcl-devel tk-devel xsimd-devel zeromq-devel zlib-devel patchelf file openblas openblas-devel protobuf numactl numactl-devel openmpi openmpi-devel - + tcl-devel tk-devel xsimd-devel zeromq-devel zlib-devel patchelf file \ + openblas openblas-devel protobuf numactl numactl-devel openmpi openmpi-devel + rpm -ivh --nodeps \ https://mirror.stream.centos.org/9-stream/CRB/ppc64le/os/Packages/protobuf-lite-devel-3.14.0-17.el9.ppc64le.rpm @@ -88,11 +141,9 @@ python --version ######################################## # install build tools (stable uv) ######################################## - -pip install -U pip setuptools-rust pip install uv -pip install "setuptools<70" build wheel cmake auditwheel -uv pip install "setuptools<70" cython meson-python pybind11 "sympy>=1.13.3" --no-build-isolation +pip install build cmake +uv pip install setuptools==78.1.1 cython meson-python pybind11 "sympy>=1.13.3" --no-build-isolation ######################################## # Rust @@ -115,127 +166,310 @@ export MAX_JOBS=${MAX_JOBS:-$(nproc)} export GRPC_PYTHON_BUILD_SYSTEM_OPENSSL=1 ######################################## -# Install Packages From Devpi +# Install Torch packages from DevPI or build from source ######################################## -uv pip install numpy==2.3.5 pillow==12.2.0 --extra-index-url "$IBM_DEVPI_URL" -try_install_from_devpi "opencv-python-headless==${OPENCV_VERSION}" -try_install_from_devpi "torch==${TORCH_VERSION}" -try_install_from_devpi "torchvision==${TORCHVISION_VERSION}" + +is_available_on_devpi() { + local pkg=$1 + local version=$2 + + echo "Checking DevPI for ${pkg}==${version}..." + + if pip index versions "${pkg}" \ + --index-url "${IBM_DEVPI_URL}" 2>/dev/null | + grep -F "${version}" >/dev/null; then + + echo "${pkg}==${version} is available on DevPI" + return 0 + fi + + echo "${pkg}==${version} is NOT available on DevPI" + return 1 +} + +######################################## +# Common packages from DevPI +######################################## +uv pip install numpy==2.3.5 pillow --extra-index-url "$IBM_DEVPI_URL" ######################################## -# torch audio +# OpenCV ######################################## -TEMP_BUILD_DIR=$(mktemp -d) -cd "${TEMP_BUILD_DIR}" -export BUILD_SOX=1 BUILD_KALDI=1 BUILD_RNNT=1 USE_FFMPEG=0 USE_ROCM=0 USE_CUDA=0 -export TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_FFMPEG=1 -git clone --recursive https://github.com/pytorch/audio.git -b "v${TORCHAUDIO_VERSION}" -cd audio -#patching -sed -i ' -s|_CSRC_DIR / "_torchaudio.cpp"|str(_CSRC_DIR / "_torchaudio.cpp")|; -s|_CSRC_DIR / "utils.cpp"|str(_CSRC_DIR / "utils.cpp")|; -s|sources=\[_CSRC_DIR / s for s in sources\]|sources=[str(_CSRC_DIR / s) for s in sources]|; -' tools/setup_helpers/extension.py -MAX_JOBS=${MAX_JOBS:-$(nproc)} \ -BUILD_VERSION=${TORCHAUDIO_VERSION} \ -uv build --wheel --out-dir "${WHEEL_DIR}" --no-build-isolation -uv pip install "${WHEEL_DIR}"/torchaudio*.whl -cd "${REPO_ROOT}" -rm -rf "${TEMP_BUILD_DIR}" +if [[ "${OPENCV_REQUIREMENT}" == ">=4.13.0" ]]; then + echo "OpenCV requirement is ${OPENCV_REQUIREMENT}" + echo "Installing opencv-python-headless from DevPI" + + OPENCV_VERSION="==4.13.0.92+ppc64le1" +else + echo "Using OpenCV requirement: ${OPENCV_REQUIREMENT}" + + OPENCV_VERSION="${OPENCV_REQUIREMENT}" +fi + +try_install_from_devpi "opencv-python-headless${OPENCV_VERSION}" ######################################## -# Xgrammar +# Check Torch / Torchvision / Torchaudio ######################################## -uv pip install \ - "scikit-build-core==0.11.6" \ - "pyproject-metadata<0.8" \ - pathspec \ - packaging \ - distro \ - "setuptools<70" \ - setuptools_scm \ - cmake \ - ninja \ - pybind11 \ - nanobind -uv pip install apache-tvm-ffi==0.1.12 \ - --no-build-isolation \ - --no-cache - -TEMP_BUILD_DIR=$(mktemp -d) - -pushd "${TEMP_BUILD_DIR}" - -export CFLAGS="-fno-lto -mcpu=power9" -export CXXFLAGS="-fno-lto -mcpu=power9" -export LDFLAGS="-fno-lto" -export PATH=/opt/vllm/bin:$PATH -export Python_EXECUTABLE=/opt/vllm/bin/python3 -export Python3_EXECUTABLE=/opt/vllm/bin/python3 -export PYTHON_EXECUTABLE=/opt/vllm/bin/python3 +TORCH_FROM_DEVPI=false +TORCHVISION_FROM_DEVPI=false +TORCHAUDIO_FROM_DEVPI=false -export Python_ROOT_DIR=/opt/vllm -export Python3_ROOT_DIR=/opt/vllm +if is_available_on_devpi "torch" "${TORCH_VERSION}"; then + TORCH_FROM_DEVPI=true +fi -git clone \ - --recursive \ - https://github.com/mlc-ai/xgrammar \ - -b "v${XGRAMMAR_VERSION}" +if is_available_on_devpi "torchvision" "${TORCHVISION_VERSION}"; then + TORCHVISION_FROM_DEVPI=true +fi -cd xgrammar +if is_available_on_devpi "torchaudio" "${TORCHAUDIO_VERSION}"; then + TORCHAUDIO_FROM_DEVPI=true +fi -cp cmake/config.cmake . -export PYTHONPATH=/opt/vllm/lib64/python3.12/site-packages:/opt/vllm/lib/python3.12/site-packages:${PYTHONPATH:-} +######################################## +# Torch +######################################## -uv build \ - --wheel \ - --out-dir "${WHEEL_DIR}" \ - --no-build-isolation +if [[ "${TORCH_FROM_DEVPI}" == "true" ]]; then + echo "Installing torch==${TORCH_VERSION} from DevPI" + + uv pip install \ + --extra-index-url "${IBM_DEVPI_URL}" \ + --index-strategy unsafe-best-match \ + --only-binary=:all: \ + --no-build-isolation \ + "torch==${TORCH_VERSION}" +else + echo "Torch wheel not available on DevPI. Building torch==${TORCH_VERSION} from source." + + TEMP_BUILD_DIR=$(mktemp -d) + cd "${TEMP_BUILD_DIR}" -uv pip install "${WHEEL_DIR}"/xgrammar*.whl -v + git clone \ + --recursive \ + https://github.com/pytorch/pytorch.git \ + -b "v${TORCH_VERSION}" -popd + cd pytorch -rm -rf "${TEMP_BUILD_DIR}" -cd "${REPO_ROOT}" + export BLAS=OpenBLAS + export USE_OPENMP=1 + export USE_MKLDNN=OFF + export USE_MKLDNN_CBLAS=OFF + export _GLIBCXX_USE_CXX11_ABI=1 + uv pip install maturin + pip install -r requirements.txt --extra-index-url "$IBM_DEVPI_URL" + python setup.py develop + rm -f dist/torch*+git*whl + MAX_JOBS=${MAX_JOBS:-$(nproc)} \ + PYTORCH_BUILD_VERSION=${TORCH_VERSION} PYTORCH_BUILD_NUMBER=1 uv build --wheel --out-dir "${WHEEL_DIR}" + uv pip install "${WHEEL_DIR}"/torch*.whl + cd "${REPO_ROOT}" + rm -rf "${TEMP_BUILD_DIR}" +fi ######################################## -# RHOAI Binary Downloads +# Torchvision ######################################## -pip download \ - --index-url "${RHOAI_INDEX_URL}" \ - --only-binary=:all: \ - --no-deps \ - llvmlite==0.47.0 \ - -d "${WHEEL_DIR}" -pip download \ - --index-url "${RHOAI_INDEX_URL}" \ - --only-binary=:all: \ - --no-deps \ - Numba==0.65.0 \ - -d "${WHEEL_DIR}" +if [[ "${TORCHVISION_FROM_DEVPI}" == "true" ]]; then + echo "Installing torchvision==${TORCHVISION_VERSION} from DevPI" + + uv pip install \ + --extra-index-url "${IBM_DEVPI_URL}" \ + --index-strategy unsafe-best-match \ + --only-binary=:all: \ + --no-build-isolation \ + "torchvision==${TORCHVISION_VERSION}" +else + echo "Torchvision wheel not available on DevPI. Building torchvision==${TORCHVISION_VERSION} from source." + + TEMP_BUILD_DIR=$(mktemp -d) + cd "${TEMP_BUILD_DIR}" + + export TORCHVISION_USE_NVJPEG=0 TORCHVISION_USE_FFMPEG=0 + git clone \ + --recursive \ + https://github.com/pytorch/vision.git \ + -b "v${TORCHVISION_VERSION}" + + cd vision + uv pip install standard-pkg-resources --no-build-isolation + MAX_JOBS=${MAX_JOBS:-$(nproc)} \ + BUILD_VERSION=${TORCHVISION_VERSION} \ + uv build --wheel --out-dir "${WHEEL_DIR}" --no-build-isolation + + export BUILD_VERSION="${TORCHVISION_VERSION}" + + uv build \ + --wheel \ + --out-dir "${WHEEL_DIR}" \ + --no-build-isolation + + uv pip install "${WHEEL_DIR}"/torchvision*.whl + + cd "${REPO_ROOT}" + rm -rf "${TEMP_BUILD_DIR}" +fi + +######################################## +# Torchaudio +######################################## + +if [[ "${TORCHAUDIO_FROM_DEVPI}" == "true" ]]; then + echo "Installing torchaudio==${TORCHAUDIO_VERSION} from DevPI" + + uv pip install \ + --extra-index-url "${IBM_DEVPI_URL}" \ + --index-strategy unsafe-best-match \ + --only-binary=:all: \ + --no-build-isolation \ + "torchaudio==${TORCHAUDIO_VERSION}" +else + echo "Torchaudio wheel not available on DevPI. Building torchaudio==${TORCHAUDIO_VERSION} from source." + + TEMP_BUILD_DIR=$(mktemp -d) + cd "${TEMP_BUILD_DIR}" + + export BUILD_SOX=1 + export BUILD_KALDI=1 + export BUILD_RNNT=1 + export USE_FFMPEG=0 + export USE_ROCM=0 + export USE_CUDA=0 + export TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_FFMPEG=1 + + git clone \ + --recursive \ + https://github.com/pytorch/audio.git \ + -b "v${TORCHAUDIO_VERSION}" + + cd audio + + # Patches required for newer Python versions + sed -i ' + s|_CSRC_DIR / "_torchaudio.cpp"|str(_CSRC_DIR / "_torchaudio.cpp")|; + s|_CSRC_DIR / "utils.cpp"|str(_CSRC_DIR / "utils.cpp")|; + s|sources=\[_CSRC_DIR / s for s in sources\]|sources=[str(_CSRC_DIR / s) for s in sources]|; + ' tools/setup_helpers/extension.py + + MAX_JOBS="${MAX_JOBS}" \ + BUILD_VERSION="${TORCHAUDIO_VERSION}" \ + uv build \ + --wheel \ + --out-dir "${WHEEL_DIR}" \ + --no-build-isolation + + uv pip install "${WHEEL_DIR}"/torchaudio*.whl + + cd "${REPO_ROOT}" + rm -rf "${TEMP_BUILD_DIR}" +fi + +######################################## +# Xgrammar +######################################## +XGRAMMAR_FROM_DEVPI=false + +if is_available_on_devpi "xgrammar" "${XGRAMMAR_VERSION}"; then + XGRAMMAR_FROM_DEVPI=true +fi + +if [[ "${XGRAMMAR_FROM_DEVPI}" == "true" ]]; then + echo "Installing xgrammar==${XGRAMMAR_VERSION} from DevPI" + + uv pip install \ + --extra-index-url "${IBM_DEVPI_URL}" \ + --index-strategy unsafe-best-match \ + --only-binary=:all: \ + --no-build-isolation \ + "xgrammar==${XGRAMMAR_VERSION}" +else + echo "Xgrammar wheel not available on DevPI. Building xgrammar==${XGRAMMAR_VERSION} from source." + # Install xgrammar build dependencies only when building from source. + uv pip install \ + "scikit-build-core==0.11.6" \ + "pyproject-metadata<0.8" \ + pathspec \ + packaging \ + distro \ + setuptools_scm \ + cmake \ + ninja \ + pybind11 \ + nanobind + uv pip install apache-tvm-ffi==0.1.12 \ + --no-build-isolation \ + --no-cache + + TEMP_BUILD_DIR=$(mktemp -d) + + pushd "${TEMP_BUILD_DIR}" + + export CFLAGS="-fno-lto -mcpu=power9" + export CXXFLAGS="-fno-lto -mcpu=power9" + export LDFLAGS="-fno-lto" + export PATH=/opt/vllm/bin:$PATH + + export Python_EXECUTABLE=/opt/vllm/bin/python3 + export Python3_EXECUTABLE=/opt/vllm/bin/python3 + export PYTHON_EXECUTABLE=/opt/vllm/bin/python3 + + export Python_ROOT_DIR=/opt/vllm + export Python3_ROOT_DIR=/opt/vllm + + git clone \ + --recursive \ + https://github.com/mlc-ai/xgrammar \ + -b "v${XGRAMMAR_VERSION}" + + cd xgrammar + + cp cmake/config.cmake . + export PYTHONPATH=/opt/vllm/lib64/python3.12/site-packages:/opt/vllm/lib/python3.12/site-packages:${PYTHONPATH:-} + + uv build \ + --wheel \ + --out-dir "${WHEEL_DIR}" \ + --no-build-isolation + + uv pip install "${WHEEL_DIR}"/xgrammar*.whl -v + + popd + + rm -rf "${TEMP_BUILD_DIR}" + cd "${REPO_ROOT}" +fi + +NUMBA_REQUIREMENT=$(grep -E '^numba[[:space:]]*(==|>=|<=|~=)' \ + requirements/cpu.txt | + sed -E 's/^numba[[:space:]]*//' | + sed -E 's/[[:space:]]*#.*$//' | + sed -E 's/;.*$//' | + tr -d ' ' | + head -n1) + +echo "Numba requirement: ${NUMBA_REQUIREMENT}" +try_install_from_devpi "numba${NUMBA_REQUIREMENT}" ######################################## # install built wheels ######################################## uv pip install setuptools_scm maturin setuptools-rust ninja scikit-build-core pybind11 nanobind \ --no-build-isolation -uv pip install "${WHEEL_DIR}"/*.whl ######################################## # install remaining deps ######################################## sed -i.bak -e 's/.*torch.*//g' pyproject.toml requirements/*.txt - -uv pip install "setuptools>=78.1.1" --no-build-isolation - export PKG_CONFIG_PATH=/usr/local/lib/pkgconfig:/usr/local/lib64/pkgconfig:/usr/lib64/pkgconfig uv pip install -r requirements/common.txt \ -r requirements/cpu.txt \ -r requirements/build/cpu.txt --index-strategy unsafe-best-match + + From 364bc81a9ecc72d44d8d1486971d8059d3f8932d Mon Sep 17 00:00:00 2001 From: linitra24 Date: Sat, 10 Oct 2026 12:52:51 +0800 Subject: [PATCH 229/624] [Bugfix][LoRA] Fix GPT-OSS gate/up interleaving with mixed MoE LoRA formats (#60812) Signed-off-by: linitra24 --- tests/lora/test_gptoss_tp.py | 38 ++++++++++++++++++- .../fused_moe/experts/lora_experts_mixin.py | 23 +++++++++-- 2 files changed, 56 insertions(+), 5 deletions(-) diff --git a/tests/lora/test_gptoss_tp.py b/tests/lora/test_gptoss_tp.py index 70b94ea2eb2a..baabaedded63 100644 --- a/tests/lora/test_gptoss_tp.py +++ b/tests/lora/test_gptoss_tp.py @@ -64,7 +64,13 @@ def _enable_deterministic_lora_shrink(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") -def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None: +def generate_and_test( + llm: vllm.LLM, + lora_path: str, + lora_id: int, + *, + is_3d_lora_weight: bool = False, +) -> None: prompts = [ PROMPT_TEMPLATE.format( context="Give the average number of working horses on farms with more than 5000 total horses." # noqa: E501 @@ -80,7 +86,11 @@ def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None: outputs = llm.generate( prompts, sampling_params, - lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None, + lora_request=LoRARequest( + str(lora_id), lora_id, lora_path, is_3d_lora_weight=is_3d_lora_weight + ) + if lora_id + else None, ) # Print the outputs. generated_texts: list[str] = [] @@ -187,3 +197,27 @@ def test_gpt_oss_lora_tp2( generate_and_test(llm, gptoss20b_lora_files, lora_id=1) generate_and_test(llm, gptoss20b_lora_files, lora_id=2) + + +def test_gpt_oss_lora_mixed_moe_format( + gptoss20b_lora_files, + monkeypatch: pytest.MonkeyPatch, +): + _enable_deterministic_lora_shrink(monkeypatch) + + llm = vllm.LLM( + MODEL_PATH, + max_model_len=1024, + enable_lora=True, + enable_mixed_moe_lora_format=True, + max_loras=2, + max_lora_rank=8, + max_num_seqs=2, + max_num_batched_tokens=2048, + compilation_config=vllm.config.CompilationConfig( + cudagraph_specialize_lora=False, + ), + ) + + generate_and_test(llm, gptoss20b_lora_files, lora_id=1, is_3d_lora_weight=True) + generate_and_test(llm, gptoss20b_lora_files, lora_id=2, is_3d_lora_weight=True) diff --git a/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py b/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py index ac75e7bc798b..5d552a88944d 100644 --- a/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py +++ b/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py @@ -3,6 +3,8 @@ import torch +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig from vllm.model_executor.layers.fused_moe.experts.lora_context import MoELoRAContext @@ -22,6 +24,7 @@ class LoRAExpertsMixin: state is on lora_context or passed as arguments. """ + moe_config: FusedMoEConfig _lora_context: MoELoRAContext | None = None def set_lora_context(self, ctx: MoELoRAContext) -> None: @@ -70,8 +73,13 @@ def apply_w13_lora( # concat swap afterwards. w13_lora_a_stacked = w13_lora_a_stacked[::-1] w13_lora_b_stacked = w13_lora_b_stacked[::-1] - return lora_context.punica_wrapper.add_lora_w13( - y, + interleave_w13 = ( + lora_context.w13_num_slices == 2 + and self.moe_config.activation == MoEActivation.SWIGLUOAI + ) + lora_output = torch.zeros_like(y) if interleave_w13 else y + result = lora_context.punica_wrapper.add_lora_w13( + lora_output, x, w13_lora_a_stacked, w13_lora_b_stacked, @@ -89,9 +97,18 @@ def apply_w13_lora( lora_context.w13_num_slices, lora_context.fully_sharded, lora_context.use_tuned_config, - add_inputs=add_inputs, + add_inputs=add_inputs or interleave_w13, token_lora_mapping=lora_context.local_token_lora_mapping, ) + if interleave_w13: + # Map gate/up slices onto even/odd output channels without a copy. + gate_up_delta = lora_output.unflatten(-1, (2, -1)) + gate_up_output = y.unflatten(-1, (-1, 2)).transpose(-1, -2) + if add_inputs: + gate_up_output.add_(gate_up_delta) + else: + gate_up_output.copy_(gate_up_delta) + return result def apply_w2_lora( self, From 8ea646b82713067406344e315a3c4ae211b0f799 Mon Sep 17 00:00:00 2001 From: Alec <35311602+alec-flowers@users.noreply.github.com> Date: Fri, 9 Oct 2026 21:59:35 -0700 Subject: [PATCH 230/624] [Rust Frontend][gRPC] Remove `GenerateStream` producer task so disconnects abort requests (#60922) Co-authored-by: Codex Co-authored-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 Signed-off-by: Alec Flowers Signed-off-by: Bugen Zhao --- rust/Cargo.lock | 1 - rust/src/server/Cargo.toml | 1 - rust/src/server/src/grpc/inference.rs | 130 ++++++++++++++------------ rust/src/server/src/grpc/tests.rs | 64 +++++++++++++ 4 files changed, 132 insertions(+), 64 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index c7b3d66e4c1b..419365580535 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5997,7 +5997,6 @@ dependencies = [ "tls-listener", "tokio", "tokio-openssl", - "tokio-stream", "tokio-util", "tonic", "tonic-health", diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml index 8b22b6344a0e..b7b00398ced1 100644 --- a/rust/src/server/Cargo.toml +++ b/rust/src/server/Cargo.toml @@ -35,7 +35,6 @@ thiserror-ext.workspace = true tls-listener.workspace = true tokio.workspace = true tokio-openssl.workspace = true -tokio-stream.workspace = true tokio-util.workspace = true tonic.workspace = true tonic-health.workspace = true diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs index a1e631b2e4e8..1484a9f9e375 100644 --- a/rust/src/server/src/grpc/inference.rs +++ b/rust/src/server/src/grpc/inference.rs @@ -5,16 +5,17 @@ use std::pin::Pin; use std::sync::Arc; use std::time::Instant; -use futures::{Stream, StreamExt as _}; +use asynk_strim_attr::{TryYielder, try_stream}; +use futures::{Stream, StreamExt as _, pin_mut}; use thiserror_ext::AsReport as _; -use tokio::sync::mpsc; -use tokio_stream::wrappers::ReceiverStream; use tonic::{Request, Response, Status}; use tracing::{Span, info, info_span, warn}; use tracing_futures::Instrument as _; use uuid::Uuid; use vllm_llm::current_unix_timestamp_secs; -use vllm_text::{DecodedTextEvent, Prompt, SampledDelta, TextOutputStreamExt as _, TextRequest}; +use vllm_text::{ + DecodedTextEvent, Prompt, SampledDelta, TextOutputStream, TextOutputStreamExt as _, TextRequest, +}; use super::convert::{self, ResponseOpts}; use super::{InferenceServer, pb}; @@ -304,67 +305,72 @@ impl pb::inference_server::Inference for InferenceServiceImpl { .await .map_err(|error| log_text_error(&request_span, started_at, "submission", error))?; - let (tx, rx) = mpsc::channel(32); - - let task_span = request_span.clone(); - tokio::spawn( - async move { - futures::pin_mut!(stream); - while let Some(event) = stream.next().await { - let response = match event { - Err(error) => Err(log_text_error(&task_span, started_at, "stream", error)), - Ok(DecodedTextEvent::Start { - prompt_token_ids, - prompt_logprobs, - }) => { - let prompt_info = convert::to_prompt_info( - &prompt_token_ids, - prompt_logprobs.as_ref(), - &response_opts, - ); - Ok(pb::GenerateResponse { - prompt_info: Some(prompt_info), - outputs: None, - }) - } - Ok(DecodedTextEvent::TextDelta { - decoded, - sampled: - SampledDelta { - token_ids, - logprobs, - }, - finished, - }) => convert::to_sequence_output( - &decoded.text, - &token_ids, - logprobs.as_ref(), - finished.as_deref(), - &response_opts, - ) - .map(|outputs| pb::GenerateResponse { - prompt_info: None, - outputs: Some(outputs), - }), - }; - - let failed = response.is_err(); - if tx.send(response).await.is_err() || failed { - break; - } - } - info!( - parent: &task_span, - elapsed_ms = started_at.elapsed().as_millis() as u64, - "gRPC inference stream closed" + let response_stream = + generate_stream_responses(stream, response_opts, request_span.clone(), started_at) + .instrument(request_span); + Ok(Response::new(Box::pin(response_stream))) + } +} + +/// Convert one decoded text stream into gRPC streaming responses. +/// +/// The response body owns the generation stream, so a client disconnect drops +/// it and aborts the request in the engine, even before any output arrives. +#[try_stream] +async fn generate_stream_responses( + stream: impl TextOutputStream, + response_opts: ResponseOpts, + request_span: Span, + started_at: Instant, + mut y: TryYielder, +) -> Result<(), Status> { + pin_mut!(stream); + + while let Some(event) = stream.next().await { + let response = match event { + Err(error) => Err(log_text_error(&request_span, started_at, "stream", error)), + Ok(DecodedTextEvent::Start { + prompt_token_ids, + prompt_logprobs, + }) => { + let prompt_info = convert::to_prompt_info( + &prompt_token_ids, + prompt_logprobs.as_ref(), + &response_opts, ); + Ok(pb::GenerateResponse { + prompt_info: Some(prompt_info), + outputs: None, + }) } - .instrument(request_span), - ); - - let response_stream = ReceiverStream::new(rx); - Ok(Response::new(Box::pin(response_stream))) + Ok(DecodedTextEvent::TextDelta { + decoded, + sampled: + SampledDelta { + token_ids, + logprobs, + }, + finished, + }) => convert::to_sequence_output( + &decoded.text, + &token_ids, + logprobs.as_ref(), + finished.as_deref(), + &response_opts, + ) + .map(|outputs| pb::GenerateResponse { + prompt_info: None, + outputs: Some(outputs), + }), + }; + y.yield_ok(response?).await; } + info!( + parent: &request_span, + elapsed_ms = started_at.elapsed().as_millis() as u64, + "gRPC inference stream closed" + ); + Ok(()) } fn text_error_to_status(error: vllm_text::Error) -> Status { diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 0ddbfb096e43..f46c658b0188 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -1827,6 +1827,70 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() { server_task.abort(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn dropping_generate_stream_aborts_before_engine_output() { + let (submitted_tx, submitted_rx) = tokio::sync::oneshot::channel(); + let (inference_service, control_service, engine_health, engine_task) = + setup_grpc_service_with_engine_script( + b"engine-grpc-cancel-pending-output".to_vec(), + default_ready_response(), + Arc::new(FakeTextBackend), + None, + move |dealer, _push| { + boxed_test_future(async move { + let add = recv_engine_message(dealer).await; + assert_eq!(add[0].as_ref(), &[0x00]); + let request: EngineCoreRequest = + rmp_serde::from_slice(&add[1]).expect("decode request"); + submitted_tx.send(()).expect("signal submitted request"); + + // A queued request produces no output before cancellation. + let abort = recv_engine_message(dealer).await; + assert_eq!(abort[0].as_ref(), &[0x01]); + let request_ids: Vec = + rmp_serde::from_slice(&abort[1]).expect("decode abort request IDs"); + assert_eq!(request_ids, vec![request.request_id]); + }) + }, + ) + .await; + let (channel, server_task) = start_grpc_test_server( + inference_service, + control_service, + engine_health, + tokio_util::sync::CancellationToken::new(), + ) + .await; + let mut client = InferenceClient::new(channel); + let stream = client + .generate_stream(pb::GenerateRequest { + request_id: "test-cancel-pending-output".to_string(), + model: "test-model".to_string(), + prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())), + stopping: Some(pb::StoppingCriteria { + max_new_tokens: 10, + ..Default::default() + }), + ..Default::default() + }) + .await + .expect("start generation") + .into_inner(); + tokio::time::timeout(Duration::from_secs(5), submitted_rx) + .await + .expect("timed out waiting for engine submission") + .expect("engine submission signal"); + + drop(stream); + + let result = tokio::time::timeout(Duration::from_secs(5), engine_task).await; + server_task.abort(); + result + .expect("disconnect must abort without waiting for engine output") + .expect("mock engine task"); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn control_abort_resolves_external_id_and_empty_is_noop() { From 7d0b4e57aac4c4323225b0eee9ae960ce75ffffe Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 9 Oct 2026 22:09:19 -0700 Subject: [PATCH 231/624] [ROCm] Remove MRV1 model exceptions (#56360) Signed-off-by: Nick Hill --- tests/test_config.py | 67 ++------------------------------------------ vllm/config/vllm.py | 30 -------------------- 2 files changed, 2 insertions(+), 95 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 425d55f29346..734253edc9f8 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -651,21 +651,14 @@ def test_no_hisparse_connector_keeps_attention_config_unset(monkeypatch): def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): """ROCm keeps the DSA models (DeepSeek V3.2/V4, GLM-5.2) on their compiled - MRV1 paths and off breakable cudagraphs by default.""" - from vllm.config.vllm import ( - ROCM_DEFAULT_MRV1_ARCHITECTURES, - default_breakable_cudagraph_architectures, - ) + paths and off breakable cudagraphs by default.""" + from vllm.config.vllm import default_breakable_cudagraph_architectures from vllm.platforms import current_platform monkeypatch.setattr(current_platform, "is_rocm", lambda: True) # The lookup is lru_cached against a fixed platform. default_breakable_cudagraph_architectures.cache_clear() try: - assert "DeepseekV32ForCausalLM" in ROCM_DEFAULT_MRV1_ARCHITECTURES - assert "DeepseekV4ForCausalLM" in ROCM_DEFAULT_MRV1_ARCHITECTURES - assert "GlmMoeDsaForCausalLM" in ROCM_DEFAULT_MRV1_ARCHITECTURES - breakable_architectures = default_breakable_cudagraph_architectures() assert "DeepseekV32ForCausalLM" not in breakable_architectures assert "DeepseekV32MTPModel" not in breakable_architectures @@ -674,66 +667,10 @@ def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): # supports uniform-batch CUDA graphs, so it must opt in or the # default FULL_AND_PIECEWISE serve path cannot start. assert "DeepseekV41ForCausalLM" in breakable_architectures - - # The carve-out takes effect via the runner-selection property - # (warning_once args must be hashable for its lru_cache). - monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False) - config = SimpleNamespace( - model_config=SimpleNamespace(architectures=["DeepseekV32ForCausalLM"]), - attention_config=AttentionConfig(), - ) - config._get_v1_model_runner_unsupported_features = lambda: [] - assert VllmConfig.use_v2_model_runner.fget(config) is False finally: default_breakable_cudagraph_architectures.cache_clear() -def test_rocm_mrv1_default_yields_to_v1_unsupported_config(monkeypatch): - """The ROCm V1 default is a speed preference, not a capability claim. - - DSpark runs only on V2, so pinning DeepSeek V4 to V1 would fail config - validation instead of serving it. With nothing V1 refuses, it still holds. - """ - from vllm.platforms import current_platform - - monkeypatch.setattr(current_platform, "is_rocm", lambda: True) - monkeypatch.setattr(vllm_config_module, "HAS_TRITON", True) - monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False) - - config = SimpleNamespace( - model_config=SimpleNamespace( - architectures=["DeepseekV4ForCausalLM"], is_diffusion=False - ), - attention_config=AttentionConfig(), - parallel_config=SimpleNamespace( - prefill_context_parallel_size=1, - pipeline_parallel_size=1, - enable_batch_sharded_sampling=False, - ), - scheduler_config=SimpleNamespace(async_scheduling=False), - speculative_config=None, - ) - config._dflash_needs_multi_kv_group = lambda: False - config._is_dflash_candidate_draft = lambda: False - config._get_v2_model_runner_unsupported_features = lambda: [] - # The real predicate, so the test also pins where dspark lands in it. - config._get_v1_model_runner_unsupported_features = lambda: ( - VllmConfig._get_v1_model_runner_unsupported_features(config) - ) - - assert VllmConfig.use_v2_model_runner.fget(config) is False - - config.speculative_config = SimpleNamespace( - method="dspark", enable_adaptive_verification=False - ) - assert VllmConfig.use_v2_model_runner.fget(config) is True - - # Yielding is not the same as selecting V2: the later checks still run, so - # a config neither runner can serve lands on V1 and fails validation there. - config._get_v2_model_runner_unsupported_features = lambda: ["sequence parallelism"] - assert VllmConfig.use_v2_model_runner.fget(config) is False - - @pytest.mark.parametrize( ("model", "architecture"), [ diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index bbda00af40bb..b4298cd145b8 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -71,12 +71,6 @@ logger = init_logger(__name__) -# TODO(rocm): These models are either unsupported by MRV2 or slower with -# MRV2 on AMD GPUs. -ROCM_DEFAULT_MRV1_ARCHITECTURES = frozenset( - {"DeepseekV32ForCausalLM", "DeepseekV4ForCausalLM", "GlmMoeDsaForCausalLM"} -) - DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset( { "DeepseekV32MTPModel", @@ -727,30 +721,6 @@ def use_v2_model_runner(self) -> bool: if use_v2_model_runner is not None: return use_v2_model_runner - from vllm.platforms import current_platform - - model_config = self.model_config - if model_config is not None and current_platform.is_rocm(): - architectures = getattr(model_config, "architectures", ()) - if any(arch in ROCM_DEFAULT_MRV1_ARCHITECTURES for arch in architectures): - # This default is a speed preference, not a claim that V1 can - # serve the config, so it yields where V1 cannot. It yields by - # falling through to the checks below, not by selecting V2. - v1_unsupported = self._get_v1_model_runner_unsupported_features() - if not v1_unsupported: - logger.warning_once( - "Defaulting to V1 model runner on ROCm for model " - "architectures: %s", - ", ".join(architectures), - ) - return False - logger.warning_once( - "Skipping the ROCm V1 model runner default for %s: V1 does " - "not support %s.", - ", ".join(architectures), - ", ".join(v1_unsupported), - ) - if not HAS_TRITON: logger.warning_once( "Model Runner V2 requires Triton; using the V1 model runner instead." From fa5cfd48d561989d2f5e7fd42788e2f6cfb479c6 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Sat, 10 Oct 2026 13:27:58 +0800 Subject: [PATCH 232/624] [Refactor] Deduplicate trace header handling in BaseServing (#60828) Signed-off-by: chaunceyjiang --- .../structured_decisions/test_strategies.py | 3 ++- vllm/entrypoints/generate/base/serving.py | 22 +------------------ .../generate/generative_scoring/serving.py | 20 ----------------- .../generate/structured_decisions/serving.py | 17 +------------- vllm/entrypoints/pooling/base/serving.py | 22 +------------------ vllm/entrypoints/serve/engine/serving.py | 9 ++++++++ vllm/v1/engine/async_llm.py | 6 ++++- 7 files changed, 19 insertions(+), 80 deletions(-) diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index a63e2587521c..41b9a8f01271 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -222,11 +222,12 @@ async def test_request_metadata_reaches_every_question( ) response = await serving.create_decision(request, raw_request) assert isinstance(response, StructuredDecisionResponse) + engine.is_tracing_enabled.assert_not_awaited() assert engine.generate.call_count == 2 for call in engine.generate.call_args_list: assert call.args[0].get("cache_salt") == cache_salt assert call.kwargs["trace_headers"] == ( - TRACE_HEADERS if tracing_enabled and request_headers else None + TRACE_HEADERS if request_headers else None ) diff --git a/vllm/entrypoints/generate/base/serving.py b/vllm/entrypoints/generate/base/serving.py index d1f3e94b006f..a0eda6d05ff4 100644 --- a/vllm/entrypoints/generate/base/serving.py +++ b/vllm/entrypoints/generate/base/serving.py @@ -2,14 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json import time -from collections.abc import Awaitable, Mapping +from collections.abc import Awaitable from dataclasses import dataclass, field from http import HTTPStatus from typing import ClassVar, Generic, TypeVar from fastapi import Request from pydantic import ConfigDict -from starlette.datastructures import Headers from vllm import RequestOutput from vllm.engine.protocol import EngineClient @@ -33,11 +32,6 @@ from vllm.lora.request import LoRARequest from vllm.tokenizers import TokenizerLike from vllm.tokenizers.detokenizer_utils import convert_ids_list_to_tokens -from vllm.tracing import ( - contains_trace_headers, - extract_trace_headers, - log_tracing_disabled_warning, -) from vllm.v1.metrics.stats import RequestStateStats logger = init_logger(__name__) @@ -222,20 +216,6 @@ def _convert_generation_error_to_streaming_response( status_code=e.status_code, ) - async def _get_trace_headers( - self, - headers: Headers, - ) -> Mapping[str, str] | None: - is_tracing_enabled = await self.engine_client.is_tracing_enabled() - - if is_tracing_enabled: - return extract_trace_headers(headers) - - if contains_trace_headers(headers): - log_tracing_disabled_warning() - - return None - @staticmethod def _get_data_parallel_rank(raw_request: Request | None) -> int | None: """Pulls the data parallel rank from a header, if provided.""" diff --git a/vllm/entrypoints/generate/generative_scoring/serving.py b/vllm/entrypoints/generate/generative_scoring/serving.py index 74dbfd075997..92c084771643 100644 --- a/vllm/entrypoints/generate/generative_scoring/serving.py +++ b/vllm/entrypoints/generate/generative_scoring/serving.py @@ -11,7 +11,6 @@ import asyncio import math import time -from collections.abc import Mapping from typing import Any, Literal from fastapi import Request @@ -32,11 +31,6 @@ from vllm.logger import init_logger from vllm.sampling_params import SamplingParams from vllm.tokenizers import TokenizerLike -from vllm.tracing import ( - contains_trace_headers, - extract_trace_headers, - log_tracing_disabled_warning, -) from vllm.utils import random_uuid logger = init_logger(__name__) @@ -462,17 +456,3 @@ def _compute_probabilities( token_id: math.exp(logprob) for token_id, logprob in label_logprobs.items() } - - async def _get_trace_headers( - self, - headers: Mapping[str, str], - ) -> Mapping[str, str] | None: - """Extract trace headers from request headers.""" - if not contains_trace_headers(headers): - return None - - if not await self.engine_client.is_tracing_enabled(): - log_tracing_disabled_warning() - return None - - return extract_trace_headers(headers) diff --git a/vllm/entrypoints/generate/structured_decisions/serving.py b/vllm/entrypoints/generate/structured_decisions/serving.py index e6e990bd46df..2633647a0b65 100644 --- a/vllm/entrypoints/generate/structured_decisions/serving.py +++ b/vllm/entrypoints/generate/structured_decisions/serving.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio import json -from collections.abc import Mapping from typing import Any from fastapi import Request @@ -12,11 +11,6 @@ from vllm.entrypoints.serve.engine.serving import BaseServing from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.logger import init_logger -from vllm.tracing import ( - contains_trace_headers, - extract_trace_headers, - log_tracing_disabled_warning, -) from .protocol import ( DecisionUsage, @@ -74,19 +68,10 @@ def __init__( super().__init__( models=models, model_config=model_config, request_logger=request_logger ) + self.engine_client = strategy.context.engine_client self.strategy = strategy self.limits = strategy.limits() - async def _get_trace_headers( - self, headers: Mapping[str, str] - ) -> Mapping[str, str] | None: - if not contains_trace_headers(headers): - return None - if not await self.strategy.context.engine_client.is_tracing_enabled(): - log_tracing_disabled_warning() - return None - return extract_trace_headers(headers) - async def create_decision( self, request: StructuredDecisionRequest, diff --git a/vllm/entrypoints/pooling/base/serving.py b/vllm/entrypoints/pooling/base/serving.py index 92cdfaec252f..2fb0b8e1bc72 100644 --- a/vllm/entrypoints/pooling/base/serving.py +++ b/vllm/entrypoints/pooling/base/serving.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio from abc import ABC, abstractmethod -from collections.abc import AsyncGenerator, Mapping +from collections.abc import AsyncGenerator from concurrent.futures import Executor from http import HTTPStatus from typing import ClassVar @@ -10,7 +10,6 @@ import torch from fastapi import Request from fastapi.responses import Response -from starlette.datastructures import Headers from vllm import PoolingRequestOutput, envs from vllm.config import VllmConfig @@ -24,11 +23,6 @@ from vllm.exceptions import GenerationError, RetryableRequestError from vllm.lora.request import LoRARequest from vllm.renderers.base import BaseRenderer -from vllm.tracing import ( - contains_trace_headers, - extract_trace_headers, - log_tracing_disabled_warning, -) from vllm.utils.async_utils import make_async, merge_async_iterators from ...serve.engine.protocol import ErrorResponse @@ -262,20 +256,6 @@ def _validate_request(self, ctx: PoolingServeContext) -> None: return None - async def _get_trace_headers( - self, - headers: Headers, - ) -> Mapping[str, str] | None: - is_tracing_enabled = await self.engine_client.is_tracing_enabled() - - if is_tracing_enabled: - return extract_trace_headers(headers) - - if contains_trace_headers(headers): - log_tracing_disabled_warning() - - return None - class PoolingServing(PoolingBaseServing, ABC): def __init__(self, *args, **kwargs): diff --git a/vllm/entrypoints/serve/engine/serving.py b/vllm/entrypoints/serve/engine/serving.py index 8b4244955295..15f6adc37744 100644 --- a/vllm/entrypoints/serve/engine/serving.py +++ b/vllm/entrypoints/serve/engine/serving.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Mapping from http import HTTPStatus from fastapi import Request @@ -23,6 +24,7 @@ extract_prompt_len, ) from vllm.sampling_params import BeamSearchParams +from vllm.tracing import contains_trace_headers, extract_trace_headers from vllm.utils import random_uuid @@ -37,6 +39,13 @@ def __init__( self.model_config = model_config self.request_logger = request_logger + async def _get_trace_headers( + self, headers: Mapping[str, str] + ) -> Mapping[str, str] | None: + if not contains_trace_headers(headers): + return None + return extract_trace_headers(headers) + async def _check_model( self, request: AnyRequest | AnyPoolingRequest, diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 7e50aa045123..bdab43e92528 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -44,7 +44,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.tasks import SupportedTask from vllm.tokenizers import TokenizerLike -from vllm.tracing import init_tracer +from vllm.tracing import init_tracer, log_tracing_disabled_warning from vllm.transformers_utils.config import maybe_register_config_serialize_by_value from vllm.usage.usage_lib import UsageContext from vllm.utils.async_utils import cancel_task_threadsafe @@ -403,6 +403,10 @@ async def add_request( if self.errored: raise EngineDeadError() + if trace_headers and self.observability_config.otlp_traces_endpoint is None: + log_tracing_disabled_warning() + trace_headers = None + is_pooling = isinstance(params, PoolingParams) if ( From 46fb84c2d4ef1f89cf472c62b8c964819a9e512a Mon Sep 17 00:00:00 2001 From: devpramod Date: Sat, 10 Oct 2026 01:34:10 -0400 Subject: [PATCH 233/624] [Model] Support EAGLE3/DFlash aux hidden states for Granite (#60855) Signed-off-by: Pai, Pramod Co-authored-by: Claude Opus 5.5 (1M context) --- vllm/model_executor/models/granite.py | 27 ++++++++++++++++++++++----- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/models/granite.py b/vllm/model_executor/models/granite.py index 81ab62937e1a..24c0851bdf85 100644 --- a/vllm/model_executor/models/granite.py +++ b/vllm/model_executor/models/granite.py @@ -60,7 +60,13 @@ from vllm.model_executor.utils import set_weight_attrs from vllm.sequence import IntermediateTensors -from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant +from .interfaces import ( + EagleModelMixin, + SupportsEagle3, + SupportsLoRA, + SupportsPP, + SupportsQuant, +) from .utils import ( AutoWeightsLoader, PPMissingLayer, @@ -307,7 +313,7 @@ def forward( @support_torch_compile -class GraniteModel(nn.Module): +class GraniteModel(nn.Module, EagleModelMixin): hf_to_vllm_mapper = WeightsMapper( orig_to_new_stacked={ # weight_name: (param_name, shard_id) @@ -363,7 +369,7 @@ def forward( positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None, inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: if get_pp_group().is_first_rank: if inputs_embeds is not None: hidden_states = inputs_embeds @@ -375,8 +381,14 @@ def forward( assert intermediate_tensors is not None hidden_states = intermediate_tensors["hidden_states"] - for layer in islice(self.layers, self.start_layer, self.end_layer): + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, None) + for idx, layer in enumerate( + islice(self.layers, self.start_layer, self.end_layer) + ): hidden_states = layer(positions, hidden_states) + self._maybe_add_hidden_state( + aux_hidden_states, idx + 1, hidden_states, None + ) if not get_pp_group().is_last_rank: return IntermediateTensors( @@ -386,6 +398,9 @@ def forward( ) hidden_states = self.norm(hidden_states) + + if len(aux_hidden_states) > 0: + return hidden_states, aux_hidden_states return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -393,7 +408,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) -class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant): +class GraniteForCausalLM( + nn.Module, SupportsLoRA, SupportsPP, SupportsQuant, SupportsEagle3 +): hf_to_vllm_mapper = GraniteModel.hf_to_vllm_mapper # LoRA specific attributes packed_modules_mapping: dict[str, list[str]] = { From a205fb317703f54a9b5afee87a03af5c18f95c75 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sat, 10 Oct 2026 00:54:52 -0500 Subject: [PATCH 234/624] [Bugfix][Model] Restore AXK1 config defaults dropped with the vendored config (#60954) Signed-off-by: Andreas Karatzas Co-authored-by: Claude --- vllm/model_executor/models/AXK1.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index 83859946a255..2d9e5a0dab9b 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -116,7 +116,6 @@ def __init__( self.ep_size = self.ep_group.size() assert config.n_routed_experts is not None assert config.num_experts_per_tok is not None - assert config.scoring_func is not None assert config.hidden_act is not None self.n_routed_experts: int = config.n_routed_experts self.n_shared_experts: int | None = config.n_shared_experts @@ -134,7 +133,7 @@ def __init__( config.n_routed_experts, prefix=f"{prefix}.gate", ) - if config.topk_method == "noaux_tc": + if getattr(config, "topk_method", "noaux_tc") == "noaux_tc": self.gate.e_score_correction_bias = nn.Parameter( torch.empty(config.n_routed_experts, dtype=torch.float32) ) @@ -186,7 +185,7 @@ def __init__( num_expert_group=config.n_group, topk_group=config.topk_group, prefix=f"{prefix}.experts", - scoring_func=config.scoring_func, + scoring_func=getattr(config, "scoring_func", "sigmoid"), # we do scaling outside, set factor to 1.0 to avoid double mul # aiter applies routed_scaling_factor internally routed_scaling_factor=self.routed_scaling_factor, @@ -645,11 +644,10 @@ def __init__( self.routed_scaling_factor = config.routed_scaling_factor def _is_layer_sparse(self) -> bool: - assert self.config.moe_layer_freq is not None return ( self.config.n_routed_experts is not None and self.layer_idx >= self.config.first_k_dense_replace - and self.layer_idx % self.config.moe_layer_freq == 0 + and self.layer_idx % getattr(self.config, "moe_layer_freq", 1) == 0 ) def forward( From ff440c780544abd63bc82d57557adedf15647bfd Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Sat, 10 Oct 2026 14:06:57 +0800 Subject: [PATCH 235/624] [BreakableCG] Add eager break for Mamba-family ops (#57881) Signed-off-by: zhenwei-intel Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py | 2 ++ vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py | 2 ++ vllm/model_executor/layers/mamba/mamba_mixer.py | 2 ++ vllm/model_executor/layers/mamba/mamba_mixer2.py | 2 ++ vllm/model_executor/layers/mamba/short_conv.py | 2 ++ 5 files changed, 10 insertions(+) diff --git a/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py index c8f5a9ba64c9..a7bbac74f66d 100644 --- a/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py @@ -4,6 +4,7 @@ from einops import rearrange from torch import nn +from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import ( VllmConfig, get_current_vllm_config, @@ -524,6 +525,7 @@ def weight_loader(param, loaded_weight, loaded_shard_id=None): return weight_loader +@eager_break_during_capture def olmo_hybrid_gdn_full_forward( hidden_states: torch.Tensor, output: torch.Tensor, diff --git a/vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py b/vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py index e7618b29b713..2d88f17e6af3 100644 --- a/vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py +++ b/vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py @@ -8,6 +8,7 @@ import torch.nn.functional as F from einops import rearrange +from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import get_current_vllm_config from vllm.forward_context import ForwardContext, get_forward_context from vllm.model_executor.custom_op import PluggableLayer @@ -325,6 +326,7 @@ def _forward( output[:num_actual_tokens], _ = self.out_proj(hidden) +@eager_break_during_capture def linear_attention( hidden_states: torch.Tensor, output: torch.Tensor, diff --git a/vllm/model_executor/layers/mamba/mamba_mixer.py b/vllm/model_executor/layers/mamba/mamba_mixer.py index 8f3b8a2638c5..74cac8b63014 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer.py @@ -7,6 +7,7 @@ from torch import nn from torch.nn.parameter import Parameter +from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config from vllm.distributed.parallel_state import ( get_tensor_model_parallel_rank, @@ -465,6 +466,7 @@ def split_batch_to_prefill_and_decode( ) +@eager_break_during_capture def mamba_mixer( hidden_states: torch.Tensor, output: torch.Tensor, diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 3753153fa3f2..bfa1103bb368 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -8,6 +8,7 @@ import torch from torch import nn +from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config from vllm.config.mamba import MambaBackendEnum from vllm.distributed import ( @@ -1214,6 +1215,7 @@ def share_replayssm_ring_trackers( replayssm_mixers[last_layer_name]._updates_replayssm_trackers = True +@eager_break_during_capture def mamba_mixer2( projected_states: torch.Tensor, output: torch.Tensor, diff --git a/vllm/model_executor/layers/mamba/short_conv.py b/vllm/model_executor/layers/mamba/short_conv.py index 37aef2b5bdb0..7a695dfd31c3 100644 --- a/vllm/model_executor/layers/mamba/short_conv.py +++ b/vllm/model_executor/layers/mamba/short_conv.py @@ -4,6 +4,7 @@ import torch +from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config from vllm.distributed import get_tensor_model_parallel_world_size from vllm.forward_context import ForwardContext, get_forward_context @@ -360,6 +361,7 @@ def mamba_type(self) -> MambaAttentionBackendEnum: return MambaAttentionBackendEnum.SHORT_CONV +@eager_break_during_capture def short_conv( hidden_states: torch.Tensor, output: torch.Tensor, From 699c892dd5dd20557d774250df15d584d4d9b318 Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Sat, 10 Oct 2026 15:47:39 +0900 Subject: [PATCH 236/624] [Perf] Skip the M-RoPE warmup when every holder of the RoPE fuses it (#60246) Signed-off-by: lesj0610 Co-authored-by: Isotr0py --- .../test_qwen_vl_triton_warmup.py | 21 ++++++++++++------- .../warmup/qwen_vl_triton_warmup.py | 6 +++++- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/tests/model_executor/test_qwen_vl_triton_warmup.py b/tests/model_executor/test_qwen_vl_triton_warmup.py index e94f2ae21717..008c77c14f7a 100644 --- a/tests/model_executor/test_qwen_vl_triton_warmup.py +++ b/tests/model_executor/test_qwen_vl_triton_warmup.py @@ -3,6 +3,7 @@ from types import SimpleNamespace +import pytest import torch from vllm.model_executor.warmup.qwen_vl_triton_warmup import ( @@ -61,7 +62,10 @@ def forward(self, *args, **kwargs): assert calls == [] -def test_mrope_warmup_launches_triton_shapes() -> None: +@pytest.mark.parametrize( + "owners_fused, skipped", [((), False), ((True,), True), ((True, False), False)] +) +def test_mrope_warmup_launches_triton_shapes(owners_fused, skipped) -> None: from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding launched: list[tuple[torch.Size, torch.Size]] = [] @@ -84,7 +88,13 @@ class FakeMropeModel(torch.nn.Module): def __init__(self) -> None: super().__init__() - self.rotary_emb = FakeRope() + rope = FakeRope() + if not owners_fused: + self.rotary_emb = rope + for i, fused in enumerate(owners_fused): + owner = torch.nn.Module() + owner.rotary_emb, owner.use_fused_qk_norm_rope_gate = rope, fused + self.add_module(f"attn{i}", owner) def get_mrope_input_positions(self, input_tokens, mm_features): raise AssertionError("warmup must not build live M-RoPE positions") @@ -99,11 +109,8 @@ def get_mrope_input_positions(self, input_tokens, mm_features): device=torch.device("cpu"), ) _warm_mrope(runner, FakeMropeModel()) - assert [shape for shape, _ in launched] == [ - torch.Size((3, 1)), - torch.Size((3, 2)), - torch.Size((3, 16)), - ] + expected = [torch.Size((3, n)) for n in (1, 2, 16)] + assert [shape for shape, _ in launched] == ([] if skipped else expected) def test_mrope_warmup_skips_models_without_supports_mrope() -> None: diff --git a/vllm/model_executor/warmup/qwen_vl_triton_warmup.py b/vllm/model_executor/warmup/qwen_vl_triton_warmup.py index 45956fe15cb5..08f6bed2dae2 100644 --- a/vllm/model_executor/warmup/qwen_vl_triton_warmup.py +++ b/vllm/model_executor/warmup/qwen_vl_triton_warmup.py @@ -72,7 +72,11 @@ def _warm_mrope(runner: "GPUModelRunner", model: torch.nn.Module) -> None: (module for module in model.modules() if isinstance(module, MRotaryEmbedding)), None, ) - if rope is None: + if rope is None or all( + getattr(module, "use_fused_qk_norm_rope_gate", False) + for module in model.modules() + if any(child is rope for child in module.children()) + ): return model_config = runner.model_config From c1692c85f16272e6b6f3de5ae4157e2f3b61fd4a Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Sat, 10 Oct 2026 14:50:34 +0800 Subject: [PATCH 237/624] [Misc] mv chat_utils.py to renderers. (#60522) Signed-off-by: wang.yuqi --- .github/CODEOWNERS | 2 -- docs/models/generative_models.md | 2 +- .../pooling/embed/vision_embedding_offline.py | 2 +- rust/src/chat/src/multimodal.rs | 2 +- .../multimodal/llm/test_mm_cache_stats.py | 2 +- .../llm/test_mm_processor_kwargs.py | 2 +- .../entrypoints/pooling/embed/test_online.py | 2 +- .../reward/test_token_reward_online.py | 2 +- tests/entrypoints/pooling/scoring/util.py | 6 ++--- .../unit_tests/test_chat_billing_header.py | 2 +- .../multimodal/generation/test_memory_leak.py | 2 +- .../models/multimodal/pooling/test_colpali.py | 6 ++--- .../multimodal/pooling/test_colqwen3.py | 4 +-- .../pooling/test_jinavl_reranker.py | 4 +-- .../pooling/test_llama_nemotron_vl.py | 8 +++--- .../multimodal/processing/test_common.py | 2 +- .../test_chat_utils.py | 16 ++++++------ .../test_chat_utils_prompt_embeds.py | 6 ++--- tests/renderers/test_hf.py | 2 +- tests/renderers/test_warmup.py | 2 +- tests/tokenizers_/test_deepseek_v4.py | 2 +- tests/tokenizers_/test_mistral.py | 2 +- vllm/entrypoints/anthropic/serving.py | 2 +- vllm/entrypoints/cohere/serving.py | 4 +-- vllm/entrypoints/common/offline.py | 8 +++--- vllm/entrypoints/generate/api_router.py | 2 +- vllm/entrypoints/generate/base/protocol.py | 2 +- .../generate/structured_decisions/protocol.py | 2 +- .../structured_decisions/strategies.py | 2 +- .../launchers/api_server/app_state.py | 2 +- vllm/entrypoints/launchers/cli_args.py | 6 ++--- .../entrypoints/launchers/render/app_state.py | 2 +- vllm/entrypoints/llm.py | 10 +++---- .../openai/chat_completion/batch_serving.py | 2 +- .../openai/chat_completion/protocol.py | 10 +++---- .../openai/chat_completion/serving.py | 10 +++---- vllm/entrypoints/openai/responses/context.py | 6 ++--- vllm/entrypoints/openai/responses/protocol.py | 8 +++--- vllm/entrypoints/openai/responses/serving.py | 6 ++--- .../openai/responses/streaming_events.py | 2 +- vllm/entrypoints/openai/responses/utils.py | 2 +- vllm/entrypoints/pooling/base/io_processor.py | 6 ++--- vllm/entrypoints/pooling/base/protocol.py | 8 +++--- vllm/entrypoints/pooling/base/serving.py | 2 +- .../entrypoints/pooling/embed/io_processor.py | 10 +++---- vllm/entrypoints/pooling/embed/protocol.py | 2 +- vllm/entrypoints/pooling/factories.py | 4 +-- vllm/entrypoints/pooling/offline.py | 2 +- .../pooling/scoring/io_processor.py | 2 +- vllm/entrypoints/pooling/scoring/typing.py | 2 +- vllm/entrypoints/pooling/scoring/utils.py | 6 ++--- vllm/entrypoints/pooling/typing.py | 2 +- .../scale_out/token_in_token_out/serving.py | 2 +- vllm/entrypoints/serve/engine/typing.py | 2 +- vllm/entrypoints/serve/tokenize/protocol.py | 8 +++--- vllm/entrypoints/serve/tokenize/serving.py | 2 +- vllm/parser/abstract_parser.py | 8 +++--- vllm/parser/engine/parser_engine.py | 2 +- vllm/parser/harmony.py | 2 +- vllm/parser/response_template.py | 2 +- vllm/parser/utils.py | 2 +- vllm/renderers/base.py | 4 +-- vllm/{entrypoints => renderers}/chat_utils.py | 0 vllm/renderers/cohere.py | 6 ++--- vllm/renderers/deepseek_v32.py | 2 +- vllm/renderers/deepseek_v4.py | 2 +- vllm/renderers/hf.py | 26 +++++++++---------- vllm/renderers/inkling.py | 4 +-- vllm/renderers/inputs/preprocess.py | 2 +- vllm/renderers/kimi_k3.py | 6 ++--- vllm/renderers/mistral.py | 4 +-- vllm/renderers/online_derenderer.py | 2 +- vllm/renderers/online_renderer.py | 12 ++++----- vllm/renderers/params.py | 2 +- vllm/renderers/terratorch.py | 4 +-- vllm/tokenizers/deepseek_v32.py | 2 +- vllm/tokenizers/deepseek_v4.py | 2 +- vllm/tokenizers/deepseek_v41.py | 2 +- vllm/tokenizers/kimi_audio.py | 2 +- vllm/tokenizers/mistral.py | 2 +- vllm/tokenizers/protocol.py | 2 +- vllm/tool_parsers/apertus_tool_parser.py | 2 +- vllm/tool_parsers/deepseekv31_tool_parser.py | 2 +- vllm/tool_parsers/deepseekv3_tool_parser.py | 2 +- vllm/tool_parsers/dots_tool_parser.py | 2 +- .../tool_parsers/functiongemma_tool_parser.py | 2 +- vllm/tool_parsers/gigachat3_tool_parser.py | 2 +- vllm/tool_parsers/granite4_tool_parser.py | 2 +- .../granite_20b_fc_tool_parser.py | 2 +- vllm/tool_parsers/hermes_tool_parser.py | 2 +- vllm/tool_parsers/hy_v3_tool_parser.py | 2 +- vllm/tool_parsers/hy_v4_tool_parser.py | 2 +- vllm/tool_parsers/internlm2_tool_parser.py | 2 +- vllm/tool_parsers/jamba_tool_parser.py | 2 +- vllm/tool_parsers/llama_tool_parser.py | 2 +- vllm/tool_parsers/minicpm5xml_tool_parser.py | 2 +- vllm/tool_parsers/muse_glimmer_tool_parser.py | 2 +- vllm/tool_parsers/phi4mini_tool_parser.py | 2 +- vllm/tool_parsers/poolside_v1_tool_parser.py | 2 +- vllm/tool_parsers/streaming.py | 2 +- vllm/tool_parsers/xlam_tool_parser.py | 2 +- 101 files changed, 185 insertions(+), 187 deletions(-) rename tests/{entrypoints/unit_tests => renderers}/test_chat_utils.py (99%) rename vllm/{entrypoints => renderers}/chat_utils.py (100%) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 65d9454b283b..a92ab6a2c963 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -43,8 +43,6 @@ /vllm/entrypoints/pooling @noooop /vllm/entrypoints/serve/sagemaker @DarkLight1337 /vllm/entrypoints/serve @njhill -/vllm/entrypoints/*.py @njhill -/vllm/entrypoints/chat_utils.py @DarkLight1337 /vllm/entrypoints/common @DarkLight1337 /vllm/entrypoints/llm.py @DarkLight1337 diff --git a/docs/models/generative_models.md b/docs/models/generative_models.md index 46f8f4096d6a..04d836a8f45e 100644 --- a/docs/models/generative_models.md +++ b/docs/models/generative_models.md @@ -127,7 +127,7 @@ If the model doesn't have a chat template or you want to specify another one, you can explicitly pass a chat template: ```python -from vllm.entrypoints.chat_utils import load_chat_template +from vllm.renderers.chat_utils import load_chat_template # You can find a list of existing chat templates under `examples/` custom_template = load_chat_template(chat_template="") diff --git a/examples/pooling/embed/vision_embedding_offline.py b/examples/pooling/embed/vision_embedding_offline.py index 5175a6cd46ce..02923e9d8c20 100644 --- a/examples/pooling/embed/vision_embedding_offline.py +++ b/examples/pooling/embed/vision_embedding_offline.py @@ -223,7 +223,7 @@ def run_vlm2vec_qwen2vl(seed: int): from peft import PeftConfig, PeftModel from transformers import AutoModelForImageTextToText, AutoProcessor - from vllm.entrypoints.chat_utils import load_chat_template + from vllm.renderers.chat_utils import load_chat_template model_id = "TIGER-Lab/VLM2Vec-Qwen2VL-2B" diff --git a/rust/src/chat/src/multimodal.rs b/rust/src/chat/src/multimodal.rs index 4138d39bdab3..4b625f3d8b53 100644 --- a/rust/src/chat/src/multimodal.rs +++ b/rust/src/chat/src/multimodal.rs @@ -714,7 +714,7 @@ fn input_audio_data_url(data: &str, format: Option<&str>) -> Result { /// /// Embedding inputs share their base modality's budget rather than getting one /// of their own, matching Python's `modality.replace("_embeds", "")` in -/// `vllm/entrypoints/chat_utils.py`. +/// `vllm/renderers/chat_utils.py`. fn media_part_limit_modality(part: &MediaContentPart) -> Option { match part { MediaContentPart::Text { .. } => None, diff --git a/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py b/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py index ac86788f69a6..f15feb339d93 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py +++ b/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py @@ -8,7 +8,7 @@ from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS from vllm import LLM -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam +from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.v1.metrics import loggers as stat_loggers from vllm.v1.metrics.reader import Counter, Metric diff --git a/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py b/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py index 55990ed1dda0..91d2b0aa434a 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py +++ b/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py @@ -7,9 +7,9 @@ import pytest from vllm import LLM, SamplingParams -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.inputs import TextPrompt from vllm.outputs import RequestOutput +from vllm.renderers.chat_utils import ChatCompletionMessageParam def _make_mock_llm() -> LLM: diff --git a/tests/entrypoints/pooling/embed/test_online.py b/tests/entrypoints/pooling/embed/test_online.py index 8ac70876f4d2..ec78515de3f1 100644 --- a/tests/entrypoints/pooling/embed/test_online.py +++ b/tests/entrypoints/pooling/embed/test_online.py @@ -18,7 +18,6 @@ ) from tests.models.utils import check_embeddings_close from tests.utils import RemoteOpenAIServer -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse from vllm.entrypoints.pooling.utils import ( @@ -27,6 +26,7 @@ decode_pooling_output, ) from vllm.platforms import current_platform +from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.tokenizers import get_tokenizer from vllm.utils.serial_utils import EMBED_DTYPES, ENDIANNESS, binary2tensor diff --git a/tests/entrypoints/pooling/reward/test_token_reward_online.py b/tests/entrypoints/pooling/reward/test_token_reward_online.py index 9d89f8c42328..0ea0b8daf87b 100644 --- a/tests/entrypoints/pooling/reward/test_token_reward_online.py +++ b/tests/entrypoints/pooling/reward/test_token_reward_online.py @@ -11,13 +11,13 @@ from tests.models.utils import check_embeddings_close from tests.utils import RemoteOpenAIServer -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse from vllm.entrypoints.pooling.utils import ( MetadataItem, build_metadata_items, decode_pooling_output, ) +from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.tokenizers import get_tokenizer from vllm.utils.serial_utils import EMBED_DTYPES, ENDIANNESS, binary2tensor diff --git a/tests/entrypoints/pooling/scoring/util.py b/tests/entrypoints/pooling/scoring/util.py index dbf4b3dc3acd..0e989d2d8779 100644 --- a/tests/entrypoints/pooling/scoring/util.py +++ b/tests/entrypoints/pooling/scoring/util.py @@ -11,12 +11,12 @@ from transformers import AutoModel, AutoTokenizer from tests.conftest import HfRunner -from vllm.entrypoints.chat_utils import ( +from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam +from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score +from vllm.renderers.chat_utils import ( ChatCompletionContentPartImageParam, ChatCompletionContentPartTextParam, ) -from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam -from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score from vllm.transformers_utils.repo_utils import hf_api diff --git a/tests/entrypoints/unit_tests/test_chat_billing_header.py b/tests/entrypoints/unit_tests/test_chat_billing_header.py index 31fa97e93915..3ad694759760 100644 --- a/tests/entrypoints/unit_tests/test_chat_billing_header.py +++ b/tests/entrypoints/unit_tests/test_chat_billing_header.py @@ -6,7 +6,7 @@ import pytest -from vllm.entrypoints.chat_utils import _parse_chat_message_content +from vllm.renderers.chat_utils import _parse_chat_message_content HEADER = ( "x-anthropic-billing-header: cc_version=2.1.220.300; " diff --git a/tests/models/multimodal/generation/test_memory_leak.py b/tests/models/multimodal/generation/test_memory_leak.py index f58b7308b51d..555a61a0df33 100644 --- a/tests/models/multimodal/generation/test_memory_leak.py +++ b/tests/models/multimodal/generation/test_memory_leak.py @@ -13,8 +13,8 @@ from tests.models.registry import HF_EXAMPLE_MODELS from vllm import LLM, SamplingParams from vllm.distributed import cleanup_dist_env_and_memory -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.platforms import current_platform +from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.utils.mem_constants import KiB_bytes, MiB_bytes from vllm.utils.mem_utils import format_mib diff --git a/tests/models/multimodal/pooling/test_colpali.py b/tests/models/multimodal/pooling/test_colpali.py index 9e040bf9a1e5..f9f1eef18911 100644 --- a/tests/models/multimodal/pooling/test_colpali.py +++ b/tests/models/multimodal/pooling/test_colpali.py @@ -14,12 +14,12 @@ import torch from PIL import Image -from vllm.entrypoints.chat_utils import ( +from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam +from vllm.platforms import current_platform +from vllm.renderers.chat_utils import ( ChatCompletionContentPartImageParam, ChatCompletionContentPartTextParam, ) -from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam -from vllm.platforms import current_platform from ....conftest import VllmRunner diff --git a/tests/models/multimodal/pooling/test_colqwen3.py b/tests/models/multimodal/pooling/test_colqwen3.py index 992fbf209fed..2ff1d2baeb25 100644 --- a/tests/models/multimodal/pooling/test_colqwen3.py +++ b/tests/models/multimodal/pooling/test_colqwen3.py @@ -14,11 +14,11 @@ import torch from PIL import Image -from vllm.entrypoints.chat_utils import ( +from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam +from vllm.renderers.chat_utils import ( ChatCompletionContentPartImageParam, ChatCompletionContentPartTextParam, ) -from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam from ....conftest import VllmRunner diff --git a/tests/models/multimodal/pooling/test_jinavl_reranker.py b/tests/models/multimodal/pooling/test_jinavl_reranker.py index dca2403e3510..7aebc4620b48 100644 --- a/tests/models/multimodal/pooling/test_jinavl_reranker.py +++ b/tests/models/multimodal/pooling/test_jinavl_reranker.py @@ -5,12 +5,12 @@ from transformers import AutoModel from vllm.assets.base import VLLM_S3_BUCKET_URL -from vllm.entrypoints.chat_utils import ( +from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam +from vllm.renderers.chat_utils import ( ChatCompletionContentPartImageEmbedsParam, ChatCompletionContentPartImageParam, ChatCompletionContentPartTextParam, ) -from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam from ....conftest import HfRunner, VllmRunner diff --git a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py index e9e58d863541..3cbeda2e49e1 100644 --- a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py +++ b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py @@ -18,13 +18,13 @@ from PIL import Image from transformers import AutoModel, AutoModelForSequenceClassification, AutoProcessor -from vllm.entrypoints.chat_utils import ( - ChatCompletionContentPartImageParam, - ChatCompletionContentPartTextParam, -) from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam from vllm.model_executor.layers.fusion.mm_input_norm import FusedMMInputNorm from vllm.platforms import current_platform +from vllm.renderers.chat_utils import ( + ChatCompletionContentPartImageParam, + ChatCompletionContentPartTextParam, +) from vllm.transformers_utils.processors.nemotron_vl import ( SIGLIP_MEAN, SIGLIP_STD, diff --git a/tests/models/multimodal/processing/test_common.py b/tests/models/multimodal/processing/test_common.py index 20a5121e6008..820d655e905e 100644 --- a/tests/models/multimodal/processing/test_common.py +++ b/tests/models/multimodal/processing/test_common.py @@ -23,13 +23,13 @@ PlaceholderMetadataResolver, collect_ec_item_metadata, ) -from vllm.entrypoints.chat_utils import _get_embeds_data from vllm.inputs import MultiModalDataDict, MultiModalInput from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.cache import MultiModalProcessorOnlyCache from vllm.multimodal.inputs import MultiModalFeatureSpec, batched_tensors_equal from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext from vllm.platforms import current_platform +from vllm.renderers.chat_utils import _get_embeds_data from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config from vllm.utils.mistral import is_mistral_tokenizer diff --git a/tests/entrypoints/unit_tests/test_chat_utils.py b/tests/renderers/test_chat_utils.py similarity index 99% rename from tests/entrypoints/unit_tests/test_chat_utils.py rename to tests/renderers/test_chat_utils.py index 75867f04b2e5..d0a9b9e6c7a8 100644 --- a/tests/entrypoints/unit_tests/test_chat_utils.py +++ b/tests/renderers/test_chat_utils.py @@ -14,7 +14,14 @@ from vllm.assets.image import ImageAsset from vllm.assets.video import VideoAsset from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ( +from vllm.exceptions import VLLMValidationError +from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict +from vllm.multimodal.utils import ( + encode_audio_url, + encode_image_url, + encode_video_url, +) +from vllm.renderers.chat_utils import ( MEDIA_CONNECTOR_REGISTRY, AsyncMultiModalItemTracker, ChatCompletionMessageParam, @@ -26,13 +33,6 @@ parse_chat_messages_async, validate_chat_template, ) -from vllm.exceptions import VLLMValidationError -from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict -from vllm.multimodal.utils import ( - encode_audio_url, - encode_image_url, - encode_video_url, -) from vllm.utils.serial_utils import tensor2base64 KIMI_K2_5_MODEL_ID = "moonshotai/Kimi-K2.5" diff --git a/tests/renderers/test_chat_utils_prompt_embeds.py b/tests/renderers/test_chat_utils_prompt_embeds.py index bbd4dc5c5f2e..f26680793ee2 100644 --- a/tests/renderers/test_chat_utils_prompt_embeds.py +++ b/tests/renderers/test_chat_utils_prompt_embeds.py @@ -16,7 +16,9 @@ import torch from transformers import AutoTokenizer -from vllm.entrypoints.chat_utils import ( +from vllm.exceptions import VLLMValidationError +from vllm.renderers import TokenizeParams +from vllm.renderers.chat_utils import ( _ENABLE_PROMPT_EMBEDS_ERROR, _PROMPT_EMBEDS_MISSING_DATA_ERROR, _RESERVED_PLACEHOLDER_IN_TEXT_ERROR, @@ -26,8 +28,6 @@ parse_chat_messages, parse_chat_messages_async, ) -from vllm.exceptions import VLLMValidationError -from vllm.renderers import TokenizeParams from vllm.renderers.hf import ( _PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR, _build_mixed_prompt_embeds, diff --git a/tests/renderers/test_hf.py b/tests/renderers/test_hf.py index 833656aa3c31..c4419a70313d 100644 --- a/tests/renderers/test_hf.py +++ b/tests/renderers/test_hf.py @@ -5,9 +5,9 @@ import pytest from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.exceptions import VLLMValidationError +from vllm.renderers.chat_utils import load_chat_template from vllm.renderers.hf import ( _consolidate_system_messages, _convert_developer_to_system, diff --git a/tests/renderers/test_warmup.py b/tests/renderers/test_warmup.py index b0e3dcb54796..09fe047d56ef 100644 --- a/tests/renderers/test_warmup.py +++ b/tests/renderers/test_warmup.py @@ -28,9 +28,9 @@ import pytest -from vllm.entrypoints.chat_utils import ChatTemplateResolutionError from vllm.multimodal.processing import BaseMultiModalProcessor from vllm.renderers.base import BaseRenderer +from vllm.renderers.chat_utils import ChatTemplateResolutionError from vllm.renderers.params import ChatParams diff --git a/tests/tokenizers_/test_deepseek_v4.py b/tests/tokenizers_/test_deepseek_v4.py index 4403c2a877d0..19c40f72c51a 100644 --- a/tests/tokenizers_/test_deepseek_v4.py +++ b/tests/tokenizers_/test_deepseek_v4.py @@ -8,7 +8,7 @@ import pytest import torch -from vllm.entrypoints.chat_utils import parse_chat_messages +from vllm.renderers.chat_utils import parse_chat_messages from vllm.renderers.registry import RENDERER_REGISTRY from vllm.tokenizers.deepseek_v4 import get_deepseek_v4_tokenizer from vllm.tokenizers.registry import TokenizerRegistry diff --git a/tests/tokenizers_/test_mistral.py b/tests/tokenizers_/test_mistral.py index 24bd8e00f1e0..2480d9b994f2 100644 --- a/tests/tokenizers_/test_mistral.py +++ b/tests/tokenizers_/test_mistral.py @@ -10,9 +10,9 @@ from mistral_common.guidance.grammar_factory import GrammarFactory from mistral_common.tokens.tokenizers.base import SpecialTokenPolicy, SpecialTokens -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.exceptions import VLLMValidationError +from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.tokenizers.mistral import ( MistralTokenizer, _validate_apply_chat_template_args, diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index f11697fa165e..7b1cf95e99b0 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -35,7 +35,6 @@ AnthropicToolChangeDefinition, AnthropicUsage, ) -from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.generate.base.protocol import ( JsonSchemaResponseFormat, ResponseFormat, @@ -54,6 +53,7 @@ from vllm.entrypoints.serve.exception_handling.utils import sanitize_message from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.logger import init_logger +from vllm.renderers.chat_utils import ChatTemplateContentFormatOption from vllm.renderers.hf import HfRenderer, resolve_chat_template from vllm.renderers.online_renderer import OnlineRenderer diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 98242cc55d09..ec6dd89c3ef8 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -45,7 +45,6 @@ from pydantic import BaseModel from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.cohere.cohere_chat_message import ( CitationSource as InternalCitationSource, ) @@ -91,6 +90,7 @@ from vllm.entrypoints.serve.exception_handling.utils import sanitize_message from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.parser.abstract_parser import Parser +from vllm.renderers.chat_utils import ChatTemplateContentFormatOption from vllm.renderers.cohere import ( MESSAGES_CITATIONS_KEY, POSITION_TO_SOURCE_KEY, @@ -509,7 +509,7 @@ def _convert_tool_message(cls, msg: ToolChatMessageV2) -> dict[str, Any]: # ``TextToolContent`` and ``DocumentToolContent``. The latter # produces ``{"type": "document", "document": {...}}`` blocks, # which vLLM's shared ``_parse_chat_message_content_parts`` in - # ``vllm/entrypoints/chat_utils.py`` rejects because the OpenAI + # ``vllm/renderers/chat_utils.py`` rejects because the OpenAI # ``ChatCompletionContentPartParam`` union only knows text / # image / audio / etc. Emitting them anyway triggers a Pydantic # ``ValidatorIterator`` failure at chat-parsing time, so we diff --git a/vllm/entrypoints/common/offline.py b/vllm/entrypoints/common/offline.py index 3f75c3033dd6..b33b8c34c5ad 100644 --- a/vllm/entrypoints/common/offline.py +++ b/vllm/entrypoints/common/offline.py @@ -15,15 +15,15 @@ SamplingParams, ) from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ( - ChatCompletionMessageParam, - ChatTemplateContentFormatOption, -) from vllm.exceptions import VLLMValidationError from vllm.inputs import EngineInput from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.renderers import BaseRenderer, ChatParams, merge_kwargs +from vllm.renderers.chat_utils import ( + ChatCompletionMessageParam, + ChatTemplateContentFormatOption, +) from vllm.renderers.inputs.preprocess import ( conversation_to_seq, parse_model_prompt, diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 18dfb87a26da..66a37d6b13ee 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -72,7 +72,7 @@ async def init_generate_state( default_chat_template_kwargs: dict[str, Any], ): from vllm.entrypoints.anthropic.serving import AnthropicServingMessages - from vllm.entrypoints.chat_utils import load_chat_template + from vllm.renderers.chat_utils import load_chat_template # The Cohere serving handler depends on the optional `cohere` SDK for # its wire-format protocol models, and is additionally gated on the diff --git a/vllm/entrypoints/generate/base/protocol.py b/vllm/entrypoints/generate/base/protocol.py index 89f6816c41cb..9781bf935f33 100644 --- a/vllm/entrypoints/generate/base/protocol.py +++ b/vllm/entrypoints/generate/base/protocol.py @@ -15,10 +15,10 @@ import vllm.envs as envs from vllm.config.utils import replace -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel, UsageInfo from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.sampling_params import StructuredOutputsParams logger = init_logger(__name__) diff --git a/vllm/entrypoints/generate/structured_decisions/protocol.py b/vllm/entrypoints/generate/structured_decisions/protocol.py index b15fee171fca..ab89084dff2d 100644 --- a/vllm/entrypoints/generate/structured_decisions/protocol.py +++ b/vllm/entrypoints/generate/structured_decisions/protocol.py @@ -8,10 +8,10 @@ from pydantic import Field, model_validator from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.generate.base.protocol import validate_cache_salt from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs +from vllm.renderers.chat_utils import ChatTemplateContentFormatOption from vllm.utils import random_uuid diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index f79375e42a83..48debdb0b30f 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -16,10 +16,10 @@ from vllm.config import ModelConfig from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.generate.label_reads import next_token_label_reads from vllm.inputs import EngineInput, tokens_input from vllm.lora.request import LoRARequest +from vllm.renderers.chat_utils import ChatTemplateContentFormatOption from vllm.renderers.inputs.preprocess import extract_prompt_components from vllm.renderers.online_renderer import OnlineRenderer from vllm.sampling_params import SamplingParams diff --git a/vllm/entrypoints/launchers/api_server/app_state.py b/vllm/entrypoints/launchers/api_server/app_state.py index 4f30beb1a81f..ce68890717e3 100644 --- a/vllm/entrypoints/launchers/api_server/app_state.py +++ b/vllm/entrypoints/launchers/api_server/app_state.py @@ -7,7 +7,6 @@ from starlette.datastructures import State from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.launchers.cli_args import resolve_default_chat_template_kwargs from vllm.entrypoints.mcp.tool_server import init_tool_server from vllm.entrypoints.openai.models.protocol import BaseModelPath @@ -16,6 +15,7 @@ from vllm.entrypoints.serve.utils.api_utils import process_lora_modules from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.plugins.endpoint_plugins.interface import init_endpoint_plugins_state +from vllm.renderers.chat_utils import load_chat_template from vllm.renderers.online_derenderer import OnlineDerenderer from vllm.renderers.online_renderer import OnlineRenderer from vllm.tasks import FALLBACK_SUPPORTED_TASKS, POOLING_TASKS, SupportedTask diff --git a/vllm/entrypoints/launchers/cli_args.py b/vllm/entrypoints/launchers/cli_args.py index 28fe79fd691e..5cd58da87eac 100644 --- a/vllm/entrypoints/launchers/cli_args.py +++ b/vllm/entrypoints/launchers/cli_args.py @@ -17,12 +17,12 @@ from vllm.entrypoints.anthropic.protocol import ( AnthropicDisabledThinkingEffortOption, ) -from vllm.entrypoints.chat_utils import ( +from vllm.entrypoints.openai.models.protocol import LoRAModulePath +from vllm.logger import init_logger +from vllm.renderers.chat_utils import ( ChatTemplateContentFormatOption, validate_chat_template, ) -from vllm.entrypoints.openai.models.protocol import LoRAModulePath -from vllm.logger import init_logger from vllm.tool_parsers import ToolParserManager from vllm.tool_parsers.tool_strict_level import ToolStrictLevelName from vllm.utils.argparse_utils import FlexibleArgumentParser diff --git a/vllm/entrypoints/launchers/render/app_state.py b/vllm/entrypoints/launchers/render/app_state.py index 33aebb9bc39e..d25263f56b5a 100644 --- a/vllm/entrypoints/launchers/render/app_state.py +++ b/vllm/entrypoints/launchers/render/app_state.py @@ -5,7 +5,6 @@ from starlette.datastructures import State from vllm.config import VllmConfig -from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.launchers.cli_args import resolve_default_chat_template_kwargs from vllm.entrypoints.mcp.tool_server import init_tool_server from vllm.entrypoints.openai.models.protocol import BaseModelPath @@ -15,6 +14,7 @@ from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.plugins.endpoint_plugins.interface import init_endpoint_plugins_state from vllm.renderers import renderer_from_config +from vllm.renderers.chat_utils import load_chat_template from vllm.renderers.online_derenderer import OnlineDerenderer from vllm.renderers.online_renderer import OnlineRenderer diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index a879b614844e..ff7bd9c57402 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -28,11 +28,6 @@ ) from vllm.config.quantization import QuantizationConfigArgs from vllm.engine.arg_utils import EngineArgs -from vllm.entrypoints.chat_utils import ( - ChatCompletionMessageParam, - ChatTemplateContentFormatOption, - load_chat_template, -) from vllm.entrypoints.common.offline import _O, _R, OfflineInferenceMixin from vllm.entrypoints.generate.beam_search.offline import BeamSearchOfflineMixin from vllm.entrypoints.pooling.offline import PoolingOfflineMixin @@ -44,6 +39,11 @@ from vllm.model_executor.layers.quantization import QuantizationMethods from vllm.outputs import PoolingRequestOutput, RequestOutput from vllm.platforms import current_platform +from vllm.renderers.chat_utils import ( + ChatCompletionMessageParam, + ChatTemplateContentFormatOption, + load_chat_template, +) from vllm.sampling_params import SamplingParams from vllm.tokenizers import TokenizerLike from vllm.usage.usage_lib import UsageContext diff --git a/vllm/entrypoints/openai/chat_completion/batch_serving.py b/vllm/entrypoints/openai/chat_completion/batch_serving.py index d3ea1a1ac3de..fe83b9ddba84 100644 --- a/vllm/entrypoints/openai/chat_completion/batch_serving.py +++ b/vllm/entrypoints/openai/chat_completion/batch_serving.py @@ -8,7 +8,6 @@ from fastapi import Request -from vllm.entrypoints.chat_utils import ConversationMessage from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata from vllm.entrypoints.openai.chat_completion.protocol import ( BatchChatCompletionRequest, @@ -24,6 +23,7 @@ from vllm.logger import init_logger from vllm.outputs import RequestOutput from vllm.parser.abstract_parser import Parser +from vllm.renderers.chat_utils import ConversationMessage from vllm.tokenizers import TokenizerLike from vllm.utils.async_utils import merge_async_iterators from vllm.utils.collection_utils import as_list diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index dee8f283cf68..d2e09851eba1 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -19,11 +19,6 @@ ) from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ( - ChatCompletionMessageParam, - ChatTemplateContentFormatOption, - has_non_text_content, -) from vllm.entrypoints.generate.base.protocol import ( AnyResponseFormat, DeltaMessage, @@ -44,6 +39,11 @@ from vllm.logger import init_logger from vllm.logprobs import Logprob from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs +from vllm.renderers.chat_utils import ( + ChatCompletionMessageParam, + ChatTemplateContentFormatOption, + has_non_text_content, +) from vllm.sampling_params import ( BeamSearchParams, RepetitionDetectionParams, diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 69b46ba67f61..b291ec6e8aae 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -11,11 +11,6 @@ from fastapi import Request from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import ( - ChatTemplateContentFormatOption, - ConversationMessage, - make_tool_call_id, -) from vllm.entrypoints.generate.base.protocol import ( DeltaMessage, FunctionCall, @@ -61,6 +56,11 @@ from vllm.outputs import RequestOutput from vllm.parser import ParserManager from vllm.parser.abstract_parser import Parser +from vllm.renderers.chat_utils import ( + ChatTemplateContentFormatOption, + ConversationMessage, + make_tool_call_id, +) from vllm.renderers.online_renderer import OnlineRenderer from vllm.sampling_params import BeamSearchParams, SamplingParams from vllm.tokenizers import TokenizerLike diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index a68df6dd4d65..334e8032e283 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -21,9 +21,6 @@ from openai_harmony import Author, Message, Role, TextContent from vllm import envs -from vllm.entrypoints.chat_utils import ( - ChatTemplateContentFormatOption, -) from vllm.entrypoints.mcp.tool import Tool from vllm.entrypoints.mcp.tool_server import ToolServer from vllm.entrypoints.openai.parser.harmony_utils import render_for_completion @@ -41,6 +38,9 @@ ) from vllm.outputs import RequestOutput from vllm.parser.abstract_parser import Parser +from vllm.renderers.chat_utils import ( + ChatTemplateContentFormatOption, +) from vllm.tokenizers import TokenizerLike from vllm.utils import random_uuid diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 406ab2d0818f..373cb35df057 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -54,10 +54,6 @@ ) from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ( - ChatCompletionMessageParam, - ChatTemplateContentFormatOption, -) from vllm.entrypoints.generate.base.protocol import ( PerRequestMetrics, StopParam, @@ -68,6 +64,10 @@ from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs +from vllm.renderers.chat_utils import ( + ChatCompletionMessageParam, + ChatTemplateContentFormatOption, +) from vllm.sampling_params import ( RequestOutputKind, SamplingParams, diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index 3a417ca7a566..ba4aa0434768 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -23,9 +23,6 @@ from vllm import envs from vllm.config.utils import replace from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import ( - ChatTemplateContentFormatOption, -) from vllm.entrypoints.generate.base.protocol import ( DeltaMessage, RequestResponseMetadata, @@ -86,6 +83,9 @@ from vllm.outputs import CompletionOutput from vllm.parser import Parser, ParserManager from vllm.renderers import TokenizeParams +from vllm.renderers.chat_utils import ( + ChatTemplateContentFormatOption, +) from vllm.renderers.online_renderer import ( OnlineRenderer, ResponsesPreviousMessages, diff --git a/vllm/entrypoints/openai/responses/streaming_events.py b/vllm/entrypoints/openai/responses/streaming_events.py index 2f5b151a1b35..b101040a57a2 100644 --- a/vllm/entrypoints/openai/responses/streaming_events.py +++ b/vllm/entrypoints/openai/responses/streaming_events.py @@ -62,7 +62,6 @@ from openai.types.responses.tool import Tool from openai_harmony import Message as HarmonyMessage -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import DeltaMessage, DeltaToolCall from vllm.entrypoints.mcp.tool_server import ToolServer from vllm.entrypoints.openai.parser.harmony_utils import ( @@ -78,6 +77,7 @@ ) from vllm.outputs import CompletionOutput from vllm.parser.harmony import Segment +from vllm.renderers.chat_utils import make_tool_call_id from vllm.utils import random_uuid TOOL_NAME_TO_MCP_SERVER_LABEL: Final[dict[str, str]] = { diff --git a/vllm/entrypoints/openai/responses/utils.py b/vllm/entrypoints/openai/responses/utils.py index 6c8507840fdf..c4523964e9dd 100644 --- a/vllm/entrypoints/openai/responses/utils.py +++ b/vllm/entrypoints/openai/responses/utils.py @@ -30,7 +30,6 @@ from openai.types.responses.tool import Tool from vllm import envs -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import FunctionCall, FunctionDefinition from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionMessageParam, @@ -39,6 +38,7 @@ from vllm.entrypoints.openai.responses.protocol import ResponseInputOutputItem from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tool_parsers.utils import ( build_responses_tool_call_name_map, flat_namespace_tool_name, diff --git a/vllm/entrypoints/pooling/base/io_processor.py b/vllm/entrypoints/pooling/base/io_processor.py index 0d9431f022ef..f79841b5d9ee 100644 --- a/vllm/entrypoints/pooling/base/io_processor.py +++ b/vllm/entrypoints/pooling/base/io_processor.py @@ -10,12 +10,12 @@ PoolingRequestOutput, ) from vllm.config import VllmConfig -from vllm.entrypoints.chat_utils import ( - ChatTemplateConfig, -) from vllm.exceptions import VLLMValidationError from vllm.lora.request import LoRARequest from vllm.renderers import BaseRenderer, merge_kwargs +from vllm.renderers.chat_utils import ( + ChatTemplateConfig, +) from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq from vllm.utils.async_utils import make_async from vllm.utils.mistral import is_mistral_tokenizer diff --git a/vllm/entrypoints/pooling/base/protocol.py b/vllm/entrypoints/pooling/base/protocol.py index 0c646b969ab5..c6d11488ac23 100644 --- a/vllm/entrypoints/pooling/base/protocol.py +++ b/vllm/entrypoints/pooling/base/protocol.py @@ -7,14 +7,14 @@ from pydantic import Field, model_validator from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ( - ChatCompletionMessageParam, - ChatTemplateContentFormatOption, -) from vllm.entrypoints.generate.base.protocol import validate_cache_salt from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel from vllm.exceptions import VLLMValidationError from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs +from vllm.renderers.chat_utils import ( + ChatCompletionMessageParam, + ChatTemplateContentFormatOption, +) from vllm.tasks import check_removed_pooling_task from vllm.utils import random_uuid from vllm.utils.serial_utils import EmbedDType, EncodingFormat, Endianness diff --git a/vllm/entrypoints/pooling/base/serving.py b/vllm/entrypoints/pooling/base/serving.py index 2fb0b8e1bc72..59d4051aa57b 100644 --- a/vllm/entrypoints/pooling/base/serving.py +++ b/vllm/entrypoints/pooling/base/serving.py @@ -14,7 +14,6 @@ from vllm import PoolingRequestOutput, envs from vllm.config import VllmConfig from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import ChatTemplateConfig from vllm.entrypoints.generate.base.protocol import validate_request_mm_kwargs from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.serve.engine.serving import BaseServing @@ -23,6 +22,7 @@ from vllm.exceptions import GenerationError, RetryableRequestError from vllm.lora.request import LoRARequest from vllm.renderers.base import BaseRenderer +from vllm.renderers.chat_utils import ChatTemplateConfig from vllm.utils.async_utils import make_async, merge_async_iterators from ...serve.engine.protocol import ErrorResponse diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index 1b23a77ec60d..94f84f09b678 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -12,15 +12,15 @@ from openai.types.chat.chat_completion_content_part_image_param import ImageURL from vllm import PoolingParams -from vllm.entrypoints.chat_utils import ( - ChatCompletionContentPartParam, - ChatCompletionMessageParam, - CustomChatCompletionMessageParam, -) from vllm.inputs import tokens_input from vllm.logger import init_logger from vllm.outputs import PoolingOutput, PoolingRequestOutput from vllm.renderers import merge_kwargs +from vllm.renderers.chat_utils import ( + ChatCompletionContentPartParam, + ChatCompletionMessageParam, + CustomChatCompletionMessageParam, +) from vllm.renderers.hf import resolve_chat_template from vllm.utils.collection_utils import chunk_list from vllm.utils.mistral import is_mistral_tokenizer diff --git a/vllm/entrypoints/pooling/embed/protocol.py b/vllm/entrypoints/pooling/embed/protocol.py index a853a8ea00ee..d52871d6a6f7 100644 --- a/vllm/entrypoints/pooling/embed/protocol.py +++ b/vllm/entrypoints/pooling/embed/protocol.py @@ -17,8 +17,8 @@ from pydantic import BaseModel, BeforeValidator, Field, model_validator from vllm import PoolingParams -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel, UsageInfo +from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.utils import random_uuid from ..base.protocol import ( diff --git a/vllm/entrypoints/pooling/factories.py b/vllm/entrypoints/pooling/factories.py index e6d01be5a9e9..52c7e6f42b7e 100644 --- a/vllm/entrypoints/pooling/factories.py +++ b/vllm/entrypoints/pooling/factories.py @@ -6,10 +6,10 @@ from fastapi import FastAPI from vllm.config import ModelConfig, VllmConfig -from vllm.entrypoints.chat_utils import ChatTemplateConfig from vllm.logger import init_logger from vllm.plugins.io_processors import has_io_processor from vllm.renderers import BaseRenderer +from vllm.renderers.chat_utils import ChatTemplateConfig from vllm.tasks import POOLING_TASKS, SCORE_TYPE_MAP, SupportedTask from .base.io_processor import PoolingIOProcessor @@ -151,7 +151,7 @@ def init_pooling_state( if model_config is None: return - from vllm.entrypoints.chat_utils import load_chat_template + from vllm.renderers.chat_utils import load_chat_template from vllm.tasks import POOLING_TASKS from .classify.serving import ServingClassification diff --git a/vllm/entrypoints/pooling/offline.py b/vllm/entrypoints/pooling/offline.py index 81d78101c429..19be1eddde91 100644 --- a/vllm/entrypoints/pooling/offline.py +++ b/vllm/entrypoints/pooling/offline.py @@ -6,7 +6,6 @@ from tqdm.auto import tqdm -from vllm.entrypoints.chat_utils import ChatTemplateConfig from vllm.entrypoints.common.offline import OfflineInferenceMixin from vllm.inputs import DataPrompt, PromptType from vllm.logger import init_logger @@ -18,6 +17,7 @@ ScoringRequestOutput, ) from vllm.pooling_params import PoolingParams +from vllm.renderers.chat_utils import ChatTemplateConfig from vllm.tasks import SCORE_TYPE_MAP, PoolingTask, SupportedTask from .base.io_processor import PoolingIOProcessor diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index 91c5f762cf2c..1e4c638968f8 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -9,6 +9,7 @@ from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt from vllm.logger import init_logger from vllm.renderers import TokenizeParams +from vllm.renderers.chat_utils import ChatTemplateResolutionError from vllm.renderers.hf import safe_apply_chat_template from vllm.renderers.inputs.preprocess import ( extract_target_prompt, @@ -18,7 +19,6 @@ from vllm.tasks import PoolingTask from vllm.utils.mistral import is_mistral_tokenizer -from ...chat_utils import ChatTemplateResolutionError from ..base.io_processor import PoolingIOProcessor from ..pooling.protocol import PoolingCompletionRequest from ..typing import ( diff --git a/vllm/entrypoints/pooling/scoring/typing.py b/vllm/entrypoints/pooling/scoring/typing.py index 9afaada6d666..9fa059ce8ea6 100644 --- a/vllm/entrypoints/pooling/scoring/typing.py +++ b/vllm/entrypoints/pooling/scoring/typing.py @@ -5,7 +5,7 @@ from typing_extensions import TypedDict -from vllm.entrypoints.chat_utils import ( +from vllm.renderers.chat_utils import ( ChatCompletionContentPartImageEmbedsParam, ChatCompletionContentPartImageParam, ChatCompletionContentPartTextParam, diff --git a/vllm/entrypoints/pooling/scoring/utils.py b/vllm/entrypoints/pooling/scoring/utils.py index d7cc888a1556..8d953cf6ddd4 100644 --- a/vllm/entrypoints/pooling/scoring/utils.py +++ b/vllm/entrypoints/pooling/scoring/utils.py @@ -7,7 +7,9 @@ from vllm import PromptType, TextPrompt from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ( +from vllm.exceptions import VLLMValidationError +from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict +from vllm.renderers.chat_utils import ( BaseMultiModalItemTracker, ChatCompletionContentPartParam, ChatCompletionContentPartTextParam, @@ -15,8 +17,6 @@ MultiModalItemTracker, _parse_chat_message_content_parts, ) -from vllm.exceptions import VLLMValidationError -from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict from .typing import ( ScoreContentPartParam, diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index 8e45ce33f4b2..a299fc657d54 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -9,10 +9,10 @@ from pydantic import ConfigDict from vllm import PoolingParams, PoolingRequestOutput, PromptType -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.inputs import DataPrompt, EngineInput from vllm.lora.request import LoRARequest from vllm.renderers import ChatParams, TokenizeParams +from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.renderers.inputs import DictPrompt from ...tasks import PoolingTask diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index b6d7bf19ba93..816d70bec000 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -13,7 +13,6 @@ from fastapi import Request from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import AsyncMultiModalItemTracker from vllm.entrypoints.generate.base.protocol import ( PerRequestMetrics, RequestResponseMetadata, @@ -47,6 +46,7 @@ PlaceholderRange, ) from vllm.outputs import RequestOutput +from vllm.renderers.chat_utils import AsyncMultiModalItemTracker from vllm.renderers.online_renderer import OnlineRenderer from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.tokenizers import TokenizerLike diff --git a/vllm/entrypoints/serve/engine/typing.py b/vllm/entrypoints/serve/engine/typing.py index a6878d500150..59e5642c5784 100644 --- a/vllm/entrypoints/serve/engine/typing.py +++ b/vllm/entrypoints/serve/engine/typing.py @@ -4,7 +4,6 @@ from typing import Protocol, TypeAlias from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.openai.chat_completion.protocol import ( BatchChatCompletionRequest, ChatCompletionRequest, @@ -35,6 +34,7 @@ ) from vllm.entrypoints.speech_to_text.translation.protocol import TranslationRequest from vllm.renderers import ChatParams, TokenizeParams +from vllm.renderers.chat_utils import ChatTemplateContentFormatOption class RendererRequest(Protocol): diff --git a/vllm/entrypoints/serve/tokenize/protocol.py b/vllm/entrypoints/serve/tokenize/protocol.py index 130a1b4b45c4..261f7075f4bf 100644 --- a/vllm/entrypoints/serve/tokenize/protocol.py +++ b/vllm/entrypoints/serve/tokenize/protocol.py @@ -7,14 +7,14 @@ from pydantic import ConfigDict, Field, model_validator from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ( - ChatCompletionMessageParam, - ChatTemplateContentFormatOption, -) from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionToolsParam from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel from vllm.exceptions import VLLMValidationError from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs +from vllm.renderers.chat_utils import ( + ChatCompletionMessageParam, + ChatTemplateContentFormatOption, +) class TokenizeCompletionRequest(OpenAIBaseModel): diff --git a/vllm/entrypoints/serve/tokenize/serving.py b/vllm/entrypoints/serve/tokenize/serving.py index c2232324ce2b..b97737abc938 100644 --- a/vllm/entrypoints/serve/tokenize/serving.py +++ b/vllm/entrypoints/serve/tokenize/serving.py @@ -5,7 +5,6 @@ from fastapi import Request -from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.openai.models.serving import ( OpenAIModelRegistry, OpenAIServingModels, @@ -23,6 +22,7 @@ from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.inputs import TokensPrompt, tokens_input from vllm.logger import init_logger +from vllm.renderers.chat_utils import ChatTemplateContentFormatOption from vllm.renderers.online_renderer import OnlineRenderer from vllm.tokenizers import TokenizerLike diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index 4af133c9882c..61826dc8c86c 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -21,10 +21,6 @@ TriggeredTagsFormat, ) -from vllm.entrypoints.chat_utils import ( - get_tool_call_id_type, - make_tool_call_id, -) from vllm.entrypoints.generate.base.protocol import ( DeltaMessage, ExtractedToolCallInformation, @@ -42,6 +38,10 @@ from vllm.parser.metrics import record_tool_parser_invocation from vllm.parser.utils import count_history_tool_calls from vllm.reasoning.abs_reasoning_parsers import ReasoningParser +from vllm.renderers.chat_utils import ( + get_tool_call_id_type, + make_tool_call_id, +) from vllm.sampling_params import ( SamplingParams, StructuredOutputsParams, diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index b04fdaa7944d..c74be931d256 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -13,7 +13,6 @@ import regex as re -from vllm.entrypoints.chat_utils import get_tool_call_id_type, make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -27,6 +26,7 @@ from vllm.parser.engine.events import EventType, SemanticEvent from vllm.parser.engine.parser_engine_config import ParserEngineConfig, ParserState from vllm.parser.engine.streaming_parser_engine import StreamingParserEngine +from vllm.renderers.chat_utils import get_tool_call_id_type, make_tool_call_id from vllm.tool_parsers.utils import ( coerce_to_schema_type, extract_types_from_schema, diff --git a/vllm/parser/harmony.py b/vllm/parser/harmony.py index 68b50250defb..1cc8ccd7ffec 100644 --- a/vllm/parser/harmony.py +++ b/vllm/parser/harmony.py @@ -23,7 +23,6 @@ TagFormat, ) -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -43,6 +42,7 @@ from vllm.logger import init_logger from vllm.parser.abstract_parser import DelegatingParser, structured_outputs_to_format from vllm.reasoning.gptoss_reasoning_parser import GptOssReasoningParser +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tool_parsers.gptoss_tool_parser import GptOssToolParser from vllm.tool_parsers.structural_tag_registry import ( SimplifiedToolChoice, diff --git a/vllm/parser/response_template.py b/vllm/parser/response_template.py index 985b76308a33..85f8b67efbd6 100644 --- a/vllm/parser/response_template.py +++ b/vllm/parser/response_template.py @@ -11,7 +11,6 @@ from openai.types.responses import ToolChoiceFunction -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -30,6 +29,7 @@ load_response_template, ) from vllm.reasoning.abs_reasoning_parsers import ReasoningParser +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tool_parsers.abstract_tool_parser import ToolParser from vllm.tool_parsers.utils import iter_response_function_tool_dicts diff --git a/vllm/parser/utils.py b/vllm/parser/utils.py index 9194b7ef1821..0deb227c0bbb 100644 --- a/vllm/parser/utils.py +++ b/vllm/parser/utils.py @@ -5,12 +5,12 @@ from openai.types.responses import ResponseFunctionToolCall -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.responses.protocol import ( ResponseInputOutputItem, ResponsesRequest, ) +from vllm.renderers.chat_utils import ChatCompletionMessageParam def count_tool_calls(tool_calls: object) -> int: diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index 8c1d91e032b1..e04bfa64eb76 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -66,7 +66,7 @@ if TYPE_CHECKING: from vllm.config import VllmConfig - from vllm.entrypoints.chat_utils import ( + from vllm.renderers.chat_utils import ( ChatCompletionMessageParam, ConversationMessage, ) @@ -362,7 +362,7 @@ def warmup(self, chat_params: ChatParams) -> None: no-op if the background warmup — or a prior reset_mm_cache that joined it — already ran it. """ - from vllm.entrypoints.chat_utils import ChatTemplateResolutionError + from vllm.renderers.chat_utils import ChatTemplateResolutionError # prevent MM processor hangs with set_default_torch_num_threads(1): diff --git a/vllm/entrypoints/chat_utils.py b/vllm/renderers/chat_utils.py similarity index 100% rename from vllm/entrypoints/chat_utils.py rename to vllm/renderers/chat_utils.py diff --git a/vllm/renderers/cohere.py b/vllm/renderers/cohere.py index e22ebf1ebf1b..eb7aeffaac31 100644 --- a/vllm/renderers/cohere.py +++ b/vllm/renderers/cohere.py @@ -94,13 +94,13 @@ from typing import Any from vllm.config import VllmConfig -from vllm.entrypoints.chat_utils import ( +from vllm.logger import init_logger +from vllm.renderers.chat_utils import ( ChatCompletionMessageParam, ConversationMessage, parse_chat_messages, parse_chat_messages_async, ) -from vllm.logger import init_logger from vllm.tokenizers.hf import HfTokenizer from vllm.utils.async_utils import make_async @@ -165,7 +165,7 @@ class MelodyContentType(str, Enum): # melody-shape content list the renderer emits. # # Motivation: OpenAI's ``ChatCompletionContentPartParam`` union (in -# ``vllm/entrypoints/chat_utils.py``) only allows ``text`` / ``image_url`` +# ``vllm/renderers/chat_utils.py``) only allows ``text`` / ``image_url`` # / ``audio`` / ``video`` / ``file`` content parts on tool messages, so # ``{"type": "document", "document": {...}}`` blocks -- valid in Cohere # v2 -- get rejected by vLLM's shared chat validator before they can diff --git a/vllm/renderers/deepseek_v32.py b/vllm/renderers/deepseek_v32.py index ba2e0d60c1a1..6e4f9fd91bec 100644 --- a/vllm/renderers/deepseek_v32.py +++ b/vllm/renderers/deepseek_v32.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from vllm.config import VllmConfig -from vllm.entrypoints.chat_utils import ( +from vllm.renderers.chat_utils import ( ChatCompletionMessageParam, ConversationMessage, parse_chat_messages, diff --git a/vllm/renderers/deepseek_v4.py b/vllm/renderers/deepseek_v4.py index 470e21f1eb11..784764bc8b50 100644 --- a/vllm/renderers/deepseek_v4.py +++ b/vllm/renderers/deepseek_v4.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from vllm.config import VllmConfig -from vllm.entrypoints.chat_utils import ( +from vllm.renderers.chat_utils import ( ChatCompletionMessageParam, ConversationMessage, parse_chat_messages, diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py index 4a8bcd0c03d7..ec5b2d96ff7f 100644 --- a/vllm/renderers/hf.py +++ b/vllm/renderers/hf.py @@ -22,13 +22,6 @@ from typing_extensions import override from vllm import envs -from vllm.entrypoints.chat_utils import ( - PROMPT_EMBEDS_PLACEHOLDER_TOKEN, - ChatTemplateResolutionError, - load_chat_template, - parse_chat_messages, - parse_chat_messages_async, -) from vllm.exceptions import VLLMValidationError from vllm.inputs import EmbedsPrompt from vllm.inputs.engine import MultiModalInput @@ -46,6 +39,13 @@ apply_token_matches, find_mm_placeholders, ) +from vllm.renderers.chat_utils import ( + PROMPT_EMBEDS_PLACEHOLDER_TOKEN, + ChatTemplateResolutionError, + load_chat_template, + parse_chat_messages, + parse_chat_messages_async, +) from vllm.tokenizers.hf import HfTokenizer, maybe_make_thread_pool from vllm.transformers_utils.chat_templates import get_chat_template_fallback_path from vllm.transformers_utils.processor import cached_get_processor @@ -59,18 +59,18 @@ from collections.abc import Set from vllm.config import ModelConfig, VllmConfig - from vllm.entrypoints.chat_utils import ( - ChatCompletionMessageParam, - ChatTemplateContentFormat, - ChatTemplateContentFormatOption, - ConversationMessage, - ) from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict, TokensPrompt from vllm.inputs.engine import TokensInput from vllm.multimodal.processing.processor import ( MultiModalPromptUpdates, ResolvedPromptUpdate, ) + from vllm.renderers.chat_utils import ( + ChatCompletionMessageParam, + ChatTemplateContentFormat, + ChatTemplateContentFormatOption, + ConversationMessage, + ) from .inputs import DictPrompt from .params import ChatParams diff --git a/vllm/renderers/inkling.py b/vllm/renderers/inkling.py index 87eedfd69de4..6ca82e309912 100644 --- a/vllm/renderers/inkling.py +++ b/vllm/renderers/inkling.py @@ -14,13 +14,13 @@ """ from vllm.config import VllmConfig -from vllm.entrypoints.chat_utils import ( +from vllm.logger import init_logger +from vllm.renderers.chat_utils import ( ChatCompletionMessageParam, ConversationMessage, parse_chat_messages, parse_chat_messages_async, ) -from vllm.logger import init_logger from vllm.tokenizers.hf import HfTokenizer from vllm.utils.async_utils import make_async diff --git a/vllm/renderers/inputs/preprocess.py b/vllm/renderers/inputs/preprocess.py index 82a49f5638cb..9daf38fdbd3b 100644 --- a/vllm/renderers/inputs/preprocess.py +++ b/vllm/renderers/inputs/preprocess.py @@ -21,7 +21,7 @@ import torch from vllm.config import ModelConfig - from vllm.entrypoints.chat_utils import ChatCompletionMessageParam + from vllm.renderers.chat_utils import ChatCompletionMessageParam @overload diff --git a/vllm/renderers/kimi_k3.py b/vllm/renderers/kimi_k3.py index e0f937f2e6ce..0d6e14d072a4 100644 --- a/vllm/renderers/kimi_k3.py +++ b/vllm/renderers/kimi_k3.py @@ -3,14 +3,14 @@ from typing import Any, cast from vllm.config import VllmConfig -from vllm.entrypoints.chat_utils import ( +from vllm.exceptions import VLLMValidationError +from vllm.multimodal.media.connector import merge_media_io_kwargs +from vllm.renderers.chat_utils import ( ChatCompletionMessageParam, ConversationMessage, parse_chat_messages, parse_chat_messages_async, ) -from vllm.exceptions import VLLMValidationError -from vllm.multimodal.media.connector import merge_media_io_kwargs from vllm.tokenizers.hf import HfTokenizer from vllm.utils.async_utils import make_async diff --git a/vllm/renderers/mistral.py b/vllm/renderers/mistral.py index da2275afc6c1..712f2e90dcc8 100644 --- a/vllm/renderers/mistral.py +++ b/vllm/renderers/mistral.py @@ -2,13 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from vllm.config import VllmConfig -from vllm.entrypoints.chat_utils import ( +from vllm.logger import init_logger +from vllm.renderers.chat_utils import ( ChatCompletionMessageParam, ConversationMessage, parse_chat_messages, parse_chat_messages_async, ) -from vllm.logger import init_logger from vllm.tokenizers.mistral import MistralTokenizer from vllm.utils.async_utils import make_async diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index 66b73fb4912c..8e3e4ca25abc 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -4,7 +4,6 @@ from typing import Any from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.generate.base.protocol import ( DeltaMessage, ToolCall, @@ -42,6 +41,7 @@ from vllm.logger import init_logger from vllm.parser import Parser, ParserManager from vllm.renderers import BaseRenderer +from vllm.renderers.chat_utils import ChatTemplateContentFormatOption from vllm.tokenizers import TokenizerLike from vllm.tokenizers.detokenizer_utils import ( convert_prompt_ids_to_tokens, diff --git a/vllm/renderers/online_renderer.py b/vllm/renderers/online_renderer.py index e8220938b26a..b4e3fc6c4666 100644 --- a/vllm/renderers/online_renderer.py +++ b/vllm/renderers/online_renderer.py @@ -12,12 +12,6 @@ from openai_harmony import ToolNamespaceConfig from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ( - ChatCompletionMessageParam, - ChatTemplateContentFormatOption, - ConversationMessage, - has_non_text_content, -) from vllm.entrypoints.generate.base.protocol import validate_request_mm_kwargs from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionNamedToolChoiceParam, @@ -59,6 +53,12 @@ from vllm.logger import init_logger from vllm.parser import Parser, ParserManager from vllm.renderers import BaseRenderer, ChatParams, TokenizeParams, merge_kwargs +from vllm.renderers.chat_utils import ( + ChatCompletionMessageParam, + ChatTemplateContentFormatOption, + ConversationMessage, + has_non_text_content, +) from vllm.renderers.inputs.preprocess import ( parse_model_prompt, prompt_to_seq, diff --git a/vllm/renderers/params.py b/vllm/renderers/params.py index 456a4d0ac563..5e4cc58d3d73 100644 --- a/vllm/renderers/params.py +++ b/vllm/renderers/params.py @@ -13,7 +13,7 @@ if TYPE_CHECKING: import torch - from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption + from vllm.renderers.chat_utils import ChatTemplateContentFormatOption else: torch = LazyLoader("torch", globals(), "torch") diff --git a/vllm/renderers/terratorch.py b/vllm/renderers/terratorch.py index 68b4936a352e..5e188bf70757 100644 --- a/vllm/renderers/terratorch.py +++ b/vllm/renderers/terratorch.py @@ -1,13 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm.entrypoints.chat_utils import ( +from vllm.logger import init_logger +from vllm.renderers.chat_utils import ( ChatCompletionMessageParam, ConversationMessage, parse_chat_messages, parse_chat_messages_async, ) -from vllm.logger import init_logger from .base import BaseRenderer from .inputs import DictPrompt diff --git a/vllm/tokenizers/deepseek_v32.py b/vllm/tokenizers/deepseek_v32.py index 3f0fff938ad9..f933e74cf5af 100644 --- a/vllm/tokenizers/deepseek_v32.py +++ b/vllm/tokenizers/deepseek_v32.py @@ -5,7 +5,7 @@ from transformers import TokenizersBackend -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam +from vllm.renderers.chat_utils import ChatCompletionMessageParam from .deepseek_v32_encoding import encode_messages from .hf import HfTokenizer, get_cached_tokenizer diff --git a/vllm/tokenizers/deepseek_v4.py b/vllm/tokenizers/deepseek_v4.py index 532ca0df7c9b..471d6fbb4ac9 100644 --- a/vllm/tokenizers/deepseek_v4.py +++ b/vllm/tokenizers/deepseek_v4.py @@ -5,7 +5,7 @@ from transformers import TokenizersBackend -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam +from vllm.renderers.chat_utils import ChatCompletionMessageParam from .deepseek_v4_encoding import encode_messages from .hf import HfTokenizer, get_cached_tokenizer diff --git a/vllm/tokenizers/deepseek_v41.py b/vllm/tokenizers/deepseek_v41.py index a1449e59e203..7ff1f5cc3b83 100644 --- a/vllm/tokenizers/deepseek_v41.py +++ b/vllm/tokenizers/deepseek_v41.py @@ -6,7 +6,7 @@ from transformers import TokenizersBackend -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam +from vllm.renderers.chat_utils import ChatCompletionMessageParam from .deepseek_v41_encoding import ( IMAGE_PLACEHOLDER, diff --git a/vllm/tokenizers/kimi_audio.py b/vllm/tokenizers/kimi_audio.py index 10f57a7f7ded..32de7abf056b 100644 --- a/vllm/tokenizers/kimi_audio.py +++ b/vllm/tokenizers/kimi_audio.py @@ -13,8 +13,8 @@ from transformers import AddedToken, BatchEncoding from transformers.utils import chat_template_utils as hf_chat_utils -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.logger import init_logger +from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.tokenizers.protocol import TokenizerLike from vllm.transformers_utils.repo_utils import hf_api diff --git a/vllm/tokenizers/mistral.py b/vllm/tokenizers/mistral.py index 4ffefe0db6d9..15cf0322581a 100644 --- a/vllm/tokenizers/mistral.py +++ b/vllm/tokenizers/mistral.py @@ -32,10 +32,10 @@ from pydantic import ValidationError from transformers.tokenization_mistral_common import MistralCommonBackend -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger +from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.tokenizers.protocol import TokenizerLike if TYPE_CHECKING: diff --git a/vllm/tokenizers/protocol.py b/vllm/tokenizers/protocol.py index 74b32e60d603..b477dd1986e9 100644 --- a/vllm/tokenizers/protocol.py +++ b/vllm/tokenizers/protocol.py @@ -7,7 +7,7 @@ if TYPE_CHECKING: from transformers import BatchEncoding - from vllm.entrypoints.chat_utils import ChatCompletionMessageParam + from vllm.renderers.chat_utils import ChatCompletionMessageParam class TokenizerLike(Protocol): diff --git a/vllm/tool_parsers/apertus_tool_parser.py b/vllm/tool_parsers/apertus_tool_parser.py index 6825b3ef39e1..af9654c0cb26 100644 --- a/vllm/tool_parsers/apertus_tool_parser.py +++ b/vllm/tool_parsers/apertus_tool_parser.py @@ -14,7 +14,6 @@ import regex as re from partial_json_parser.core.options import Allow -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -26,6 +25,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, diff --git a/vllm/tool_parsers/deepseekv31_tool_parser.py b/vllm/tool_parsers/deepseekv31_tool_parser.py index f4c06d467fab..ff5f605f511f 100644 --- a/vllm/tool_parsers/deepseekv31_tool_parser.py +++ b/vllm/tool_parsers/deepseekv31_tool_parser.py @@ -5,7 +5,6 @@ import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -18,6 +17,7 @@ ChatCompletionRequest, ) from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser diff --git a/vllm/tool_parsers/deepseekv3_tool_parser.py b/vllm/tool_parsers/deepseekv3_tool_parser.py index 13fbcb5ed26d..649eded2878f 100644 --- a/vllm/tool_parsers/deepseekv3_tool_parser.py +++ b/vllm/tool_parsers/deepseekv3_tool_parser.py @@ -5,7 +5,6 @@ import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -18,6 +17,7 @@ ChatCompletionRequest, ) from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, diff --git a/vllm/tool_parsers/dots_tool_parser.py b/vllm/tool_parsers/dots_tool_parser.py index 3e8134473939..311e545365f4 100644 --- a/vllm/tool_parsers/dots_tool_parser.py +++ b/vllm/tool_parsers/dots_tool_parser.py @@ -7,7 +7,6 @@ import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -21,6 +20,7 @@ ChatCompletionToolsParam, ) from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser from vllm.tool_parsers.utils import is_complete_json, partial_tag_overlap diff --git a/vllm/tool_parsers/functiongemma_tool_parser.py b/vllm/tool_parsers/functiongemma_tool_parser.py index 708ecb1dd997..6fcae6e55220 100644 --- a/vllm/tool_parsers/functiongemma_tool_parser.py +++ b/vllm/tool_parsers/functiongemma_tool_parser.py @@ -6,7 +6,6 @@ import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -20,6 +19,7 @@ ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser diff --git a/vllm/tool_parsers/gigachat3_tool_parser.py b/vllm/tool_parsers/gigachat3_tool_parser.py index 020d4c1af700..7f9bdf86d6be 100644 --- a/vllm/tool_parsers/gigachat3_tool_parser.py +++ b/vllm/tool_parsers/gigachat3_tool_parser.py @@ -6,7 +6,6 @@ import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -20,6 +19,7 @@ ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser diff --git a/vllm/tool_parsers/granite4_tool_parser.py b/vllm/tool_parsers/granite4_tool_parser.py index 913135f13d70..aabd455c6766 100644 --- a/vllm/tool_parsers/granite4_tool_parser.py +++ b/vllm/tool_parsers/granite4_tool_parser.py @@ -7,7 +7,6 @@ import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -21,6 +20,7 @@ ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, diff --git a/vllm/tool_parsers/granite_20b_fc_tool_parser.py b/vllm/tool_parsers/granite_20b_fc_tool_parser.py index 4c6c0a1b0efe..8b2ccf58e86b 100644 --- a/vllm/tool_parsers/granite_20b_fc_tool_parser.py +++ b/vllm/tool_parsers/granite_20b_fc_tool_parser.py @@ -9,7 +9,6 @@ import regex as re from partial_json_parser.core.options import Allow -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -22,6 +21,7 @@ ChatCompletionRequest, ) from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, diff --git a/vllm/tool_parsers/hermes_tool_parser.py b/vllm/tool_parsers/hermes_tool_parser.py index c5b318cd19fc..1fa19d0321fc 100644 --- a/vllm/tool_parsers/hermes_tool_parser.py +++ b/vllm/tool_parsers/hermes_tool_parser.py @@ -6,7 +6,6 @@ import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -20,6 +19,7 @@ ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, diff --git a/vllm/tool_parsers/hy_v3_tool_parser.py b/vllm/tool_parsers/hy_v3_tool_parser.py index 0effb377d0a6..9947294ecd52 100644 --- a/vllm/tool_parsers/hy_v3_tool_parser.py +++ b/vllm/tool_parsers/hy_v3_tool_parser.py @@ -7,7 +7,6 @@ import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -21,6 +20,7 @@ ChatCompletionToolsParam, ) from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, diff --git a/vllm/tool_parsers/hy_v4_tool_parser.py b/vllm/tool_parsers/hy_v4_tool_parser.py index 2ce0381a0132..de707c00228e 100644 --- a/vllm/tool_parsers/hy_v4_tool_parser.py +++ b/vllm/tool_parsers/hy_v4_tool_parser.py @@ -11,7 +11,6 @@ import regex as re import vllm.envs as envs -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -23,6 +22,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ToolParser from vllm.tool_parsers.tool_strict_level import ToolStrictLevel diff --git a/vllm/tool_parsers/internlm2_tool_parser.py b/vllm/tool_parsers/internlm2_tool_parser.py index 709fd6c8f2a8..9bec016ec907 100644 --- a/vllm/tool_parsers/internlm2_tool_parser.py +++ b/vllm/tool_parsers/internlm2_tool_parser.py @@ -7,7 +7,6 @@ import partial_json_parser from partial_json_parser.core.options import Allow -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -21,6 +20,7 @@ ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, diff --git a/vllm/tool_parsers/jamba_tool_parser.py b/vllm/tool_parsers/jamba_tool_parser.py index 1ba67960b961..83eea1d5800a 100644 --- a/vllm/tool_parsers/jamba_tool_parser.py +++ b/vllm/tool_parsers/jamba_tool_parser.py @@ -8,7 +8,6 @@ import regex as re from partial_json_parser.core.options import Allow -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -22,6 +21,7 @@ ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser from vllm.tool_parsers.utils import extract_intermediate_diff, is_complete_json diff --git a/vllm/tool_parsers/llama_tool_parser.py b/vllm/tool_parsers/llama_tool_parser.py index 93b0980c3775..8696af46994d 100644 --- a/vllm/tool_parsers/llama_tool_parser.py +++ b/vllm/tool_parsers/llama_tool_parser.py @@ -10,7 +10,6 @@ from transformers import PreTrainedTokenizerBase import vllm.envs as envs -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -23,6 +22,7 @@ ChatCompletionRequest, ) from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tool_parsers.abstract_tool_parser import ( Tool, ToolParser, diff --git a/vllm/tool_parsers/minicpm5xml_tool_parser.py b/vllm/tool_parsers/minicpm5xml_tool_parser.py index 8b419d36df7f..818b51c47f4f 100644 --- a/vllm/tool_parsers/minicpm5xml_tool_parser.py +++ b/vllm/tool_parsers/minicpm5xml_tool_parser.py @@ -7,7 +7,6 @@ import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -22,6 +21,7 @@ ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, diff --git a/vllm/tool_parsers/muse_glimmer_tool_parser.py b/vllm/tool_parsers/muse_glimmer_tool_parser.py index 623b3540d7d3..5a0b26cf423a 100644 --- a/vllm/tool_parsers/muse_glimmer_tool_parser.py +++ b/vllm/tool_parsers/muse_glimmer_tool_parser.py @@ -41,7 +41,6 @@ from openai.types.responses import ToolChoiceFunction from transformers import PreTrainedTokenizerBase -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -56,6 +55,7 @@ ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tool_parsers.abstract_tool_parser import ( Tool, ToolParser, diff --git a/vllm/tool_parsers/phi4mini_tool_parser.py b/vllm/tool_parsers/phi4mini_tool_parser.py index 820e797f4d9f..0ef25b35e0f4 100644 --- a/vllm/tool_parsers/phi4mini_tool_parser.py +++ b/vllm/tool_parsers/phi4mini_tool_parser.py @@ -8,7 +8,6 @@ import regex as re from transformers import PreTrainedTokenizerBase -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaMessage, ExtractedToolCallInformation, @@ -19,6 +18,7 @@ ChatCompletionRequest, ) from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tool_parsers.abstract_tool_parser import ( Tool, ToolParser, diff --git a/vllm/tool_parsers/poolside_v1_tool_parser.py b/vllm/tool_parsers/poolside_v1_tool_parser.py index 178036760634..d2dfb7815fbe 100644 --- a/vllm/tool_parsers/poolside_v1_tool_parser.py +++ b/vllm/tool_parsers/poolside_v1_tool_parser.py @@ -19,7 +19,6 @@ from openai.types.responses import ToolChoiceFunction from partial_json_parser.core.options import Allow -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -36,6 +35,7 @@ ResponsesRequest, ) from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, diff --git a/vllm/tool_parsers/streaming.py b/vllm/tool_parsers/streaming.py index 07c1a2e031eb..bd5f9a7f3aa0 100644 --- a/vllm/tool_parsers/streaming.py +++ b/vllm/tool_parsers/streaming.py @@ -5,12 +5,12 @@ from dataclasses import dataclass from typing import TYPE_CHECKING -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, DeltaToolCall, ) +from vllm.renderers.chat_utils import make_tool_call_id from vllm.utils.mistral import is_mistral_tokenizer if TYPE_CHECKING: diff --git a/vllm/tool_parsers/xlam_tool_parser.py b/vllm/tool_parsers/xlam_tool_parser.py index 52dc3063556b..dec34fc0b9a1 100644 --- a/vllm/tool_parsers/xlam_tool_parser.py +++ b/vllm/tool_parsers/xlam_tool_parser.py @@ -10,7 +10,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) -from vllm.entrypoints.chat_utils import make_tool_call_id +from vllm.renderers.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, From a98247ab4db686ee03c66d5feb3c761e52a2f8ab Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Sat, 10 Oct 2026 15:51:52 +0900 Subject: [PATCH 238/624] [Model] Drop the 4x M-RoPE cache headroom for Qwen3-VL, Qwen3.5 and Qwen3.8 (#59853) Signed-off-by: lesj0610 Signed-off-by: lesj0610 Co-authored-by: lesj0610 --- vllm/config/model.py | 4 ++++ .../layers/rotary_embedding/__init__.py | 4 ++++ .../layers/rotary_embedding/mrope.py | 8 ++++--- vllm/model_executor/models/qwen3.py | 10 ++++++++- vllm/model_executor/models/qwen3_5.py | 2 ++ vllm/model_executor/models/qwen3_moe.py | 4 ++++ vllm/model_executor/models/qwen3_next.py | 2 ++ vllm/model_executor/models/qwen3_vl.py | 22 +++++++++++++++++-- vllm/model_executor/models/qwen3_vl_moe.py | 2 ++ vllm/models/qwen4_exp/amd/qsa.py | 2 ++ vllm/models/qwen4_exp/nvidia/qsa.py | 2 ++ 11 files changed, 56 insertions(+), 6 deletions(-) diff --git a/vllm/config/model.py b/vllm/config/model.py index cfbb870e7887..502f1bbd8d4f 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -508,6 +508,10 @@ def compute_hash(self) -> str: # here early. if self.multimodal_config: factors["language_model_only"] = self.multimodal_config.language_model_only + # Sizes Qwen3.5's M-RoPE cache. + factors["video_pruning"] = ( + self.multimodal_config.is_multimodal_pruning_enabled() + ) return hash_factors(factors) def _update_nested( diff --git a/vllm/model_executor/layers/rotary_embedding/__init__.py b/vllm/model_executor/layers/rotary_embedding/__init__.py index e23a5128b8a1..85d4cb63e501 100644 --- a/vllm/model_executor/layers/rotary_embedding/__init__.py +++ b/vllm/model_executor/layers/rotary_embedding/__init__.py @@ -36,6 +36,7 @@ def get_rope( rope_parameters: dict[str, Any] | None = None, dtype: torch.dtype | None = None, dual_chunk_attention_config: dict[str, Any] | None = None, + mrope_positions_factor: int = 4, ) -> RotaryEmbedding: if dtype is None: dtype = torch.get_default_dtype() @@ -78,6 +79,7 @@ def get_rope( rope_parameters_args, dual_chunk_attention_args, dtype, + mrope_positions_factor, ) if key in _ROPE_DICT: return _ROPE_DICT[key] @@ -118,6 +120,7 @@ def get_rope( dtype, mrope_section=rope_parameters["mrope_section"], mrope_interleaved=rope_parameters.get("mrope_interleaved", False), + positions_factor=mrope_positions_factor, ) elif "use_fope" in rope_parameters and rope_parameters["use_fope"]: extra_kwargs = { @@ -263,6 +266,7 @@ def get_rope( dtype, mrope_section=rope_parameters["mrope_section"], mrope_interleaved=rope_parameters.get("mrope_interleaved", False), + positions_factor=mrope_positions_factor, scaling_factor=scaling_factor, **extra_kwargs, ) diff --git a/vllm/model_executor/layers/rotary_embedding/mrope.py b/vllm/model_executor/layers/rotary_embedding/mrope.py index 11882856e049..9bcc5a95e16d 100644 --- a/vllm/model_executor/layers/rotary_embedding/mrope.py +++ b/vllm/model_executor/layers/rotary_embedding/mrope.py @@ -315,8 +315,9 @@ def __init__( dtype: torch.dtype, mrope_section: list[int] | None = None, mrope_interleaved: bool = False, - # YaRN parameters. *, + positions_factor: int = 4, + # YaRN parameters. scaling_factor: float | None = None, beta_fast: int = 32, beta_slow: int = 1, @@ -345,8 +346,9 @@ def __init__( # In Qwen2.5-VL, the maximum index value is related to the duration of # the input video. We enlarge max_position_embeddings to 4 times to get - # a larger the cos and sin cache. - self.cache_max_position_num = max_position_embeddings * 4 + # a larger the cos and sin cache. Models whose positions never exceed + # the token count pass `positions_factor=1`. + self.cache_max_position_num = max_position_embeddings * positions_factor super().__init__( head_size, rotary_dim, diff --git a/vllm/model_executor/models/qwen3.py b/vllm/model_executor/models/qwen3.py index 7a1d76234f32..31d1a094ebef 100644 --- a/vllm/model_executor/models/qwen3.py +++ b/vllm/model_executor/models/qwen3.py @@ -79,6 +79,7 @@ def __init__( attn_type: str = AttentionType.DECODER, dual_chunk_attention_config: dict[str, Any] | None = None, per_layer_sliding_window: int | None = None, + mrope_positions_factor: int = 4, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -124,6 +125,7 @@ def __init__( max_position=max_position, rope_parameters=rope_parameters, dual_chunk_attention_config=dual_chunk_attention_config, + mrope_positions_factor=mrope_positions_factor, ) attn_cls = ( EncoderOnlyAttention @@ -179,6 +181,7 @@ def __init__( quant_config: QuantizationConfig | None = None, prefix: str = "", per_layer_sliding_window: int | None = None, + mrope_positions_factor: int = 4, ) -> None: super().__init__() self.hidden_size = config.hidden_size @@ -211,6 +214,7 @@ def __init__( attn_type=attn_type, dual_chunk_attention_config=dual_chunk_attention_config, per_layer_sliding_window=per_layer_sliding_window, + mrope_positions_factor=mrope_positions_factor, ) self.mlp = Qwen3MLP( hidden_size=self.hidden_size, @@ -263,9 +267,13 @@ def forward( } ) class Qwen3Model(Qwen2Model): + decoder_layer_type: type[nn.Module] = Qwen3DecoderLayer + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__( - vllm_config=vllm_config, prefix=prefix, decoder_layer_type=Qwen3DecoderLayer + vllm_config=vllm_config, + prefix=prefix, + decoder_layer_type=self.decoder_layer_type, ) diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index 0d110dd58be1..f105948f9283 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -91,6 +91,7 @@ Qwen3VLForConditionalGeneration, Qwen3VLMultiModalProcessor, Qwen3VLProcessingInfo, + mrope_positions_factor, ) from .utils import ( AutoWeightsLoader, @@ -161,6 +162,7 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.self_attn", reduce_results=not self.use_attn_reduce_scatter_for_moe, + mrope_positions_factor=mrope_positions_factor(vllm_config), ) else: raise ValueError(f"Invalid layer_type {self.layer_type}") diff --git a/vllm/model_executor/models/qwen3_moe.py b/vllm/model_executor/models/qwen3_moe.py index 3f293550ac9a..870a082d3a72 100644 --- a/vllm/model_executor/models/qwen3_moe.py +++ b/vllm/model_executor/models/qwen3_moe.py @@ -256,6 +256,7 @@ def __init__( quant_config: QuantizationConfig | None = None, prefix: str = "", dual_chunk_attention_config: dict[str, Any] | None = None, + mrope_positions_factor: int = 4, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -303,6 +304,7 @@ def __init__( max_position=max_position_embeddings, rope_parameters=rope_parameters, dual_chunk_attention_config=dual_chunk_attention_config, + mrope_positions_factor=mrope_positions_factor, ) attention_kwargs: dict[str, Any] = {} if dual_chunk_attention_config: @@ -351,6 +353,7 @@ def __init__( vllm_config: VllmConfig, prefix: str = "", is_fused_checkpoint_transposed: bool = False, + mrope_positions_factor: int = 4, ) -> None: super().__init__() @@ -376,6 +379,7 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.self_attn", dual_chunk_attention_config=dual_chunk_attention_config, + mrope_positions_factor=mrope_positions_factor, ) # `mlp_only_layers` in the config. diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index b0bf0fecd19c..04d20c574145 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -280,6 +280,7 @@ def __init__( quant_config: QuantizationConfig | None = None, reduce_results: bool = True, prefix: str = "", + mrope_positions_factor: int = 4, ) -> None: super().__init__() self.config = config @@ -331,6 +332,7 @@ def __init__( max_position=config.max_position_embeddings, rope_parameters=config.rope_parameters, dual_chunk_attention_config=self.dual_chunk_attention_config, + mrope_positions_factor=mrope_positions_factor, ) # Late-interaction retrieval models (e.g. ColQwen3.5) run BIDIRECTIONAL diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index 6648949c7364..77780b3d1d40 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -52,7 +52,7 @@ should_torch_compile_mm_encoder, support_torch_compile, ) -from vllm.config import VllmConfig +from vllm.config import VllmConfig, get_current_vllm_config from vllm.config.multimodal import ( MultiModalConfig, MultiModalDummyOptions, @@ -148,7 +148,7 @@ Qwen2VLProcessingInfo, _create_qwen2vl_field_factory, ) -from .qwen3 import Qwen3ForCausalLM, Qwen3Model +from .qwen3 import Qwen3DecoderLayer, Qwen3ForCausalLM, Qwen3Model from .utils import ( AutoWeightsLoader, PPMissingLayer, @@ -1748,6 +1748,22 @@ def get_video_repl( return PromptUpdateDetails.from_seq(all_token_ids) +def mrope_positions_factor(vllm_config: VllmConfig) -> int: + # Video pruning keeps the positions of the tokens it drops. A drafter gets + # the target's positions, so it follows the target. + spec = vllm_config.speculative_config + model_config = (spec and spec.target_model_config) or vllm_config.model_config + mm_config = model_config.multimodal_config + pruned = mm_config is not None and mm_config.is_multimodal_pruning_enabled() + return 4 if pruned else 1 + + +class Qwen3VLDecoderLayer(Qwen3DecoderLayer): + def __init__(self, **kwargs) -> None: + factor = mrope_positions_factor(get_current_vllm_config()) + super().__init__(**kwargs, mrope_positions_factor=factor) + + @support_torch_compile( dynamic_arg_dims={ "input_ids": 0, @@ -1761,6 +1777,8 @@ def get_video_repl( } ) class Qwen3LLMModel(Qwen3Model): + decoder_layer_type = Qwen3VLDecoderLayer + def forward( self, input_ids: torch.Tensor | None, diff --git a/vllm/model_executor/models/qwen3_vl_moe.py b/vllm/model_executor/models/qwen3_vl_moe.py index 4d0c7788a6e3..9a1d771ec4e1 100644 --- a/vllm/model_executor/models/qwen3_vl_moe.py +++ b/vllm/model_executor/models/qwen3_vl_moe.py @@ -56,6 +56,7 @@ Qwen3VLForConditionalGeneration, Qwen3VLMultiModalProcessor, Qwen3VLProcessingInfo, + mrope_positions_factor, ) from .utils import StageMissingLayer, maybe_prefix @@ -73,6 +74,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: vllm_config=vllm_config, prefix=prefix, is_fused_checkpoint_transposed=True, + mrope_positions_factor=mrope_positions_factor(vllm_config), ) diff --git a/vllm/models/qwen4_exp/amd/qsa.py b/vllm/models/qwen4_exp/amd/qsa.py index ff01490bb190..709fd0b08ef1 100644 --- a/vllm/models/qwen4_exp/amd/qsa.py +++ b/vllm/models/qwen4_exp/amd/qsa.py @@ -23,6 +23,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.models.qwen3_next import Qwen3NextAttention +from vllm.model_executor.models.qwen3_vl import mrope_positions_factor from vllm.platforms import current_platform from vllm.utils.torch_utils import ( LayerNameType, @@ -248,6 +249,7 @@ def __init__( head_size=self.head_dim, max_position=config.max_position_embeddings, rope_parameters=config.rope_parameters, + mrope_positions_factor=mrope_positions_factor(vllm_config), ) self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index 0246db80943d..9f0934d5a0c6 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -29,6 +29,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding, get_rope from vllm.model_executor.models.qwen3_next import Qwen3NextAttention +from vllm.model_executor.models.qwen3_vl import mrope_positions_factor from vllm.model_executor.models.utils import ( AutoWeightsLoader, WeightsMapper, @@ -427,6 +428,7 @@ def __init__( head_size=self.head_dim, max_position=config.max_position_embeddings, rope_parameters=config.rope_parameters, + mrope_positions_factor=mrope_positions_factor(vllm_config), ) self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) From c41b2639e29c3bc01add1d34bef3032a6d9d8aca Mon Sep 17 00:00:00 2001 From: Kay Yan Date: Sat, 10 Oct 2026 17:08:36 +0800 Subject: [PATCH 239/624] [Bugfix][Responses API] Complete cache-write token accounting (#60973) Signed-off-by: Kay Yan Co-authored-by: Codex --- .../responses/test_parsable_context_unit.py | 10 ++++++++++ tests/entrypoints/unit_tests/test_context.py | 19 ++++++++++++++++++- vllm/entrypoints/openai/responses/context.py | 5 +++++ vllm/entrypoints/openai/responses/serving.py | 2 +- 4 files changed, 34 insertions(+), 2 deletions(-) diff --git a/tests/entrypoints/openai/responses/test_parsable_context_unit.py b/tests/entrypoints/openai/responses/test_parsable_context_unit.py index 710cd2d04cc9..f3927c9d7152 100644 --- a/tests/entrypoints/openai/responses/test_parsable_context_unit.py +++ b/tests/entrypoints/openai/responses/test_parsable_context_unit.py @@ -169,6 +169,7 @@ def _make_request_output( text: str = "Hello, world!", token_ids: Sequence[int] = (1, 2, 3), finish_reason: str = "stop", + num_cache_creation_tokens: int | None = None, ) -> RequestOutput: return RequestOutput( request_id="test", @@ -186,6 +187,7 @@ def _make_request_output( ) ], finished=True, + num_cache_creation_tokens=num_cache_creation_tokens, ) @@ -425,6 +427,14 @@ def test_reasoning_tokens_counted_per_round(): assert ctx.num_reasoning_tokens == 3 +def test_cache_creation_tokens_counted_per_round(): + ctx = _make_context(None) + + for count, expected in [(None, 0), (0, 0), (3, 3), (4, 7)]: + ctx.append_output(_make_request_output(num_cache_creation_tokens=count)) + assert ctx.num_cache_creation_tokens == expected + + def test_reasoning_tokens_zero_without_parser(): ctx = _make_context(None) diff --git a/tests/entrypoints/unit_tests/test_context.py b/tests/entrypoints/unit_tests/test_context.py index a5b1afe8c50f..f0af86231aaf 100644 --- a/tests/entrypoints/unit_tests/test_context.py +++ b/tests/entrypoints/unit_tests/test_context.py @@ -20,6 +20,7 @@ def create_mock_request_output( output_token_ids=None, num_cached_tokens=0, finished=True, + num_cache_creation_tokens=None, ): """Helper function to create a mock RequestOutput object for testing.""" outputs = [] @@ -44,6 +45,7 @@ def create_mock_request_output( outputs=outputs, finished=finished, num_cached_tokens=num_cached_tokens, + num_cache_creation_tokens=num_cache_creation_tokens, ) @@ -117,7 +119,8 @@ def make_harmony_context( return context, fake_parser -def test_single_turn_token_counting(): +@pytest.mark.parametrize("num_cache_creation_tokens", [None, 0, 3]) +def test_single_turn_token_counting(num_cache_creation_tokens): """Test token counting behavior for a single turn.""" # Create a context context, _ = make_harmony_context() @@ -127,6 +130,7 @@ def test_single_turn_token_counting(): prompt_token_ids=[1, 2, 3, 4, 5], # 5 prompt tokens output_token_ids=[6, 7, 8], # 3 output tokens num_cached_tokens=2, # 2 cached tokens + num_cache_creation_tokens=num_cache_creation_tokens, ) # Append the output to the context @@ -136,6 +140,7 @@ def test_single_turn_token_counting(): assert context.num_prompt_tokens == 5 assert context.num_output_tokens == 3 assert context.num_cached_tokens == 2 + assert context.num_cache_creation_tokens == (num_cache_creation_tokens or 0) assert context.num_tool_output_tokens == 0 # No tool tokens in first turn # Verify internal state tracking @@ -398,6 +403,7 @@ async def test_streaming_multi_turn_token_counting(): num_prompt_tokens = [3, 8, 13] num_output_tokens = [3, 3, 2] num_cached_tokens = [0, 3, 8] + num_cache_creation_tokens = [3, 4, 2] # Simulate three turns of conversation: # Turn 1: stream tokens one by one, then finish the message @@ -411,6 +417,7 @@ async def test_streaming_multi_turn_token_counting(): prompt_token_ids=[1, 2, 3], # 3 prompt tokens output_token_ids=[101], # Single token num_cached_tokens=num_cached_tokens[0], + num_cache_creation_tokens=num_cache_creation_tokens[0], finished=False, # Not end of message yet ) ) @@ -419,6 +426,7 @@ async def test_streaming_multi_turn_token_counting(): context.append_output( create_mock_request_output( output_token_ids=[102], + num_cache_creation_tokens=num_cache_creation_tokens[0], finished=False, ) ) @@ -427,6 +435,7 @@ async def test_streaming_multi_turn_token_counting(): context.append_output( create_mock_request_output( output_token_ids=[103], + num_cache_creation_tokens=num_cache_creation_tokens[0], finished=True, # End of message ) ) @@ -435,6 +444,7 @@ async def test_streaming_multi_turn_token_counting(): assert context.num_prompt_tokens == 3 # Initial prompt tokens assert context.num_output_tokens == 3 # Three output tokens assert context.num_cached_tokens == 0 + assert context.num_cache_creation_tokens == 3 assert context.num_tool_output_tokens == 0 # No tool output in first turn assert context.first_tok_of_message is True # Ready for next message @@ -454,6 +464,7 @@ async def test_streaming_multi_turn_token_counting(): ], # 8 tokens (includes previous) output_token_ids=[201], num_cached_tokens=num_cached_tokens[1], # Some tokens cached + num_cache_creation_tokens=num_cache_creation_tokens[1], finished=False, ) ) @@ -463,6 +474,7 @@ async def test_streaming_multi_turn_token_counting(): context.append_output( create_mock_request_output( output_token_ids=[202], + num_cache_creation_tokens=num_cache_creation_tokens[1], finished=False, ) ) @@ -471,6 +483,7 @@ async def test_streaming_multi_turn_token_counting(): context.append_output( create_mock_request_output( output_token_ids=[203], + num_cache_creation_tokens=num_cache_creation_tokens[1], finished=True, # End of reasoning message ) ) @@ -480,6 +493,7 @@ async def test_streaming_multi_turn_token_counting(): assert context.num_output_tokens == 3 + 3 # First turn + second turn assert context.num_reasoning_tokens == 3 # All tokens in analysis channel assert context.num_cached_tokens == 3 # Cached tokens from second turn + assert context.num_cache_creation_tokens == 3 + 4 # Formula: this turn prompt tokens - last turn prompt - last turn output expected_tool_tokens = 8 - 3 - 3 # = 2 @@ -505,6 +519,7 @@ async def test_streaming_multi_turn_token_counting(): ], # 13 tokens output_token_ids=[301], num_cached_tokens=num_cached_tokens[2], # More cached tokens + num_cache_creation_tokens=num_cache_creation_tokens[2], finished=False, ) ) @@ -512,6 +527,7 @@ async def test_streaming_multi_turn_token_counting(): context.append_output( create_mock_request_output( output_token_ids=[302], + num_cache_creation_tokens=num_cache_creation_tokens[2], finished=True, ) ) @@ -523,6 +539,7 @@ async def test_streaming_multi_turn_token_counting(): assert context.num_cached_tokens == sum( num_cached_tokens ) # Accumulated cached tokens + assert context.num_cache_creation_tokens == sum(num_cache_creation_tokens) # Additional tool tokens from third turn # Formula: this turn prompt - last turn prompt - last turn output diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index 334e8032e283..9f5a281163da 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -315,6 +315,7 @@ def __init__( self.num_prompt_tokens = 0 self.num_output_tokens = 0 self.num_cached_tokens = 0 + self.num_cache_creation_tokens = 0 self.num_reasoning_tokens = 0 # not implemented yet for ParsableContext self.all_turn_metrics: list[TurnMetrics] = [] @@ -347,6 +348,7 @@ def __init__( def append_output(self, output: RequestOutput) -> None: self.num_prompt_tokens = len(output.prompt_token_ids or []) self.num_cached_tokens = output.num_cached_tokens or 0 + self.num_cache_creation_tokens += output.num_cache_creation_tokens or 0 self.num_output_tokens += len(output.outputs[0].token_ids or []) if output.kv_transfer_params is not None: self.kv_transfer_params = output.kv_transfer_params @@ -651,6 +653,7 @@ def __init__( self.num_prompt_tokens = 0 self.num_output_tokens = 0 self.num_cached_tokens = 0 + self.num_cache_creation_tokens = 0 self.num_reasoning_tokens = 0 self.num_tool_output_tokens = 0 @@ -773,6 +776,8 @@ def _update_prefill_token_usage(self, output: RequestOutput) -> None: self.num_cached_tokens += num_cached_token self.current_turn_metrics.cached_input_tokens = num_cached_token + self.num_cache_creation_tokens += output.num_cache_creation_tokens or 0 + def _update_decode_token_usage(self, output: RequestOutput) -> int: """Update token usage statistics for the decode phase of generation. diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index ba4aa0434768..78a3f1015ff7 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -881,7 +881,7 @@ async def responses_full_generator( output_tokens=num_generated_tokens, total_tokens=num_prompt_tokens + num_generated_tokens, input_tokens_details=InputTokensDetails( - cache_write_tokens=getattr(context, "num_cache_creation_tokens", 0), + cache_write_tokens=context.num_cache_creation_tokens, cached_tokens=num_cached_tokens, input_tokens_per_turn=[ turn.input_tokens for turn in context.all_turn_metrics From 2bea22002d66a6c89d208fa682b40a89538f023c Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Sat, 10 Oct 2026 20:25:11 +0800 Subject: [PATCH 240/624] [BUG] Fix KV offload shared-memory slot selection using worker rank (#58497) Signed-off-by: Tony Lin --- tests/v1/kv_offload/test_factory.py | 33 +++++++++++++---------------- vllm/v1/kv_offload/cpu/spec.py | 9 ++------ vllm/v1/kv_offload/tiering/spec.py | 9 +------- 3 files changed, 18 insertions(+), 33 deletions(-) diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index 84b9da36e06b..dbc586bcab07 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -6,6 +6,7 @@ from unittest.mock import MagicMock import pytest +import torch from vllm.v1.kv_offload.base import ( CanonicalKVCaches, @@ -231,9 +232,7 @@ def fake_worker_ctor(**kwargs): monkeypatch.setattr(tiering_spec_module, "SharedOffloadRegion", fake_region_ctor) monkeypatch.setattr(tiering_spec_module, "CPUOffloadingWorker", fake_worker_ctor) - monkeypatch.setattr( - tiering_spec_module.torch.accelerator, "current_device_index", lambda: 5 - ) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 5) kv_caches = MagicMock() spec.create_worker(kv_caches) @@ -244,12 +243,13 @@ def fake_worker_ctor(**kwargs): assert worker_calls[0]["mmap_region"] is region -def test_tiering_spec_create_worker_folds_device_index_for_sharded_layout(monkeypatch): +def test_tiering_spec_create_worker_uses_worker_rank_for_sharded_layout(monkeypatch): import vllm.v1.kv_offload.tiering.spec as tiering_spec_module spec = _create_spec( spec_name="TieringOffloadingSpec", worker_kv_bytes_per_block=4096, + rank=1, world_size=4, ) assert isinstance(spec, TieringOffloadingSpec) @@ -263,9 +263,9 @@ def fake_region_ctor(**kwargs): monkeypatch.setattr(tiering_spec_module, "SharedOffloadRegion", fake_region_ctor) monkeypatch.setattr(tiering_spec_module, "CPUOffloadingWorker", MagicMock()) monkeypatch.setattr( - tiering_spec_module.torch.accelerator, + torch.accelerator, "current_device_index", - lambda: 5, + lambda: 6, ) spec.create_worker(MagicMock()) @@ -386,9 +386,7 @@ def fake_worker_ctor(**kwargs): monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: False) monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor) monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", fake_worker_ctor) - monkeypatch.setattr( - cpu_spec_module.torch.accelerator, "current_device_index", lambda: 5 - ) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 5) kv_caches = MagicMock() spec.create_worker(kv_caches) @@ -464,16 +462,16 @@ def test_cpu_spec_create_worker_skips_mmap_for_empty_cache(monkeypatch): @pytest.mark.parametrize( - ("replicated_layout", "device_index", "world_size", "expected_rank"), + ("replicated_layout", "rank", "device_index", "world_size", "expected_rank"), [ - (True, 5, 4, 0), # replicated: always slot 0 - (True, 0, 4, 0), # replicated: slot 0 regardless of device - (False, 5, 4, 1), # non-replicated: 5 % 4 == 1 - (False, 7, 4, 3), # non-replicated: 7 % 4 == 3 + (True, 1, 5, 4, 0), + (True, 0, 0, 4, 0), + (False, 1, 5, 4, 1), + (False, 3, 6, 4, 3), ], ) def test_cpu_spec_create_worker_rank_assignment( - monkeypatch, replicated_layout, device_index, world_size, expected_rank + monkeypatch, replicated_layout, rank, device_index, world_size, expected_rank ): import vllm.v1.kv_offload.cpu.spec as cpu_spec_module @@ -483,6 +481,7 @@ def test_cpu_spec_create_worker_rank_assignment( spec = _create_spec( cpu_bytes_to_use=worker_kv_bytes_per_block * 8, worker_kv_bytes_per_block=worker_kv_bytes_per_block, + rank=rank, world_size=world_size, replicated_layout=replicated_layout, ) @@ -495,9 +494,7 @@ def fake_region_ctor(**kwargs): monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor) monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", MagicMock()) - monkeypatch.setattr( - cpu_spec_module.torch.accelerator, "current_device_index", lambda: device_index - ) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: device_index) spec.create_worker(MagicMock()) diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index bf289d37a65b..c26cd32f9733 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Any -import torch from typing_extensions import override from vllm.platforms import current_platform @@ -172,12 +171,8 @@ def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: # mmap'd; fall back to the tensor path (empty tensors) as before. if self._uses_shared_region() and self.num_chunks > 0: # Replicated layout puts all ranks on slot 0 (single MLA copy); - # otherwise each rank takes its own slot by physical device index. - if self.replicated_layout: - rank = 0 - else: - world_size = self.config.parallel.world_size - rank = torch.accelerator.current_device_index() % world_size + # otherwise each worker uses its own rank-indexed slot. + rank = 0 if self.replicated_layout else self.config.parallel.rank mmap_region = SharedOffloadRegion( engine_id=self.config.engine_id, num_chunks=self.num_chunks, diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 93c6cbabcb85..508361d39261 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -54,7 +54,6 @@ from typing import Any -import torch from typing_extensions import override from vllm.logger import init_logger @@ -436,13 +435,7 @@ def _uses_shared_region(self) -> bool: @override def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: - world_size = self.config.parallel.world_size - if self.replicated_layout: - rank = 0 - else: - # Fold the global physical device index into the replica-local - # [0, world_size) slot range. - rank = torch.accelerator.current_device_index() % world_size + rank = 0 if self.replicated_layout else self.config.parallel.rank worker_mmap = SharedOffloadRegion( engine_id=self._engine_id, num_chunks=self.num_chunks, From bf16c9c13949a10fc8e77189776f95538a598709 Mon Sep 17 00:00:00 2001 From: Mook <68294499+Godmook@users.noreply.github.com> Date: Sun, 11 Oct 2026 00:05:55 +0900 Subject: [PATCH 241/624] [Bugfix][Model] K-EXAONE: restore loading of pre-2.0 configs and apply RoPE to the MTP layer (#60187) Signed-off-by: godmook Signed-off-by: Mook <68294499+Godmook@users.noreply.github.com> --- tests/models/registry.py | 4 +++ vllm/model_executor/models/exaone_moe.py | 39 ++++++++++++++---------- vllm/model_executor/models/registry.py | 1 + 3 files changed, 28 insertions(+), 16 deletions(-) diff --git a/tests/models/registry.py b/tests/models/registry.py index 31c3c5433aec..b57f3fcb7277 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -266,6 +266,10 @@ def check_available_online( ), "Exaone4ForCausalLM": _HfExamplesInfo("LGAI-EXAONE/EXAONE-4.0-32B"), "ExaoneMoeForCausalLM": _HfExamplesInfo("LGAI-EXAONE/K-EXAONE-236B-A23B"), + "ExaoneMoEForCausalLM": _HfExamplesInfo( + "LGAI-EXAONE/K-EXAONE-236B-A23B", + hf_overrides={"architectures": ["ExaoneMoEForCausalLM"]}, + ), "FalconForCausalLM": _HfExamplesInfo("tiiuae/falcon-7b"), "FalconH1ForCausalLM": _HfExamplesInfo("tiiuae/Falcon-H1-0.5B-Base"), "FlexOlmoForCausalLM": _HfExamplesInfo("allenai/Flex-reddit-2x7B-1T"), diff --git a/vllm/model_executor/models/exaone_moe.py b/vllm/model_executor/models/exaone_moe.py index d5b8b4066239..ad6b0b3594fe 100644 --- a/vllm/model_executor/models/exaone_moe.py +++ b/vllm/model_executor/models/exaone_moe.py @@ -227,24 +227,31 @@ def __init__( layer_idx = extract_layer_index(prefix) - if config.sliding_windows is not None: - self.sliding_window_size = config.sliding_windows[layer_idx] - - if config.layer_types[layer_idx] == "full_attention": - self.sliding_window_size = None - + # Configs published before K-EXAONE 2.0 only have the scalar + # `sliding_window` and no MTP fields; their MTP layer follows the main + # layer with the same index. + layer_types = config.layer_types + sliding_windows = getattr(config, "sliding_windows", None) if is_mtp: - self.sliding_window = ( - config.mtp_layer_types[layer_idx] == "sliding_attention" + layer_types = getattr(config, "mtp_layer_types", None) or layer_types + sliding_windows = ( + getattr(config, "mtp_sliding_windows", None) or sliding_windows ) - if config.mtp_sliding_windows is not None: - self.sliding_window_size = config.mtp_sliding_windows[layer_idx] - - if config.mtp_layer_types[layer_idx] == "full_attention": - self.sliding_window_size = None - # apply rotary embeddings to every layer in full attention models - self.apply_rope_all_layers = "sliding_attention" not in config.layer_types + if layer_types[layer_idx] == "full_attention": + self.sliding_window_size = None + elif sliding_windows is not None: + self.sliding_window_size = sliding_windows[layer_idx] + else: + self.sliding_window_size = config.sliding_window + + # Full-attention layers are NoPE when sliding layers are present, except + # the MTP layer, which is trained with RoPE. + self.use_rope = ( + bool(self.sliding_window_size) + or is_mtp + or "sliding_attention" not in config.layer_types + ) set_default_rope_theta(config, default_theta=1000000) self.rotary_emb = get_rope( @@ -279,7 +286,7 @@ def forward( k = self.k_norm(k) k = k.flatten(-2, -1) - if self.sliding_window_size or self.apply_rope_all_layers: + if self.use_rope: q, k = self.rotary_emb(positions, q, k) attn_output = self.attn(q, k, v) output, _ = self.o_proj(attn_output) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index e019d674f502..875179a88fae 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -98,6 +98,7 @@ "ExaoneForCausalLM": ("exaone", "ExaoneForCausalLM"), "Exaone4ForCausalLM": ("exaone4", "Exaone4ForCausalLM"), "ExaoneMoeForCausalLM": ("exaone_moe", "ExaoneMoeForCausalLM"), + "ExaoneMoEForCausalLM": ("exaone_moe", "ExaoneMoeForCausalLM"), "FalconForCausalLM": ("falcon", "FalconForCausalLM"), "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"), "FalconH1ForCausalLM": ("falcon_h1", "FalconH1ForCausalLM"), From 838526c7e5e855f05365a939ef83218cd8e2dbbe Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Sat, 10 Oct 2026 11:09:45 -0400 Subject: [PATCH 242/624] [Compile] Fix sampler compile warning (#60680) Signed-off-by: yewentao256 --- csrc/libtorch_stable/sampler.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/csrc/libtorch_stable/sampler.cu b/csrc/libtorch_stable/sampler.cu index 703a748ce686..09151004c1ad 100644 --- a/csrc/libtorch_stable/sampler.cu +++ b/csrc/libtorch_stable/sampler.cu @@ -642,6 +642,7 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode( indices, logits, rowStart, rowEnd, outIndices, outLogits, stride1, topK); } +#ifdef USE_ROCM // Measured native-kernel policy for gfx950, FP32, and k=1024. This // active-column split crossover is independent of the Python AITER/native // backend crossover; the two policies can be retuned separately. At exactly @@ -693,7 +694,6 @@ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecodeDeviceLengthAware( stride1, topK); } -#ifdef USE_ROCM // The grid and workspace stay fixed while graph replays change device lengths. __device__ int gfx950TopK512ActiveBlocks(int rowLength, int numRows, int maxBlocks) { @@ -755,7 +755,7 @@ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode512DeviceLengthAware } #endif } -#endif +#endif // USE_ROCM } // namespace vllm From 3709632ff2944a5f2ecdacc84ede5cd134b7ae08 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Sat, 10 Oct 2026 11:41:10 -0400 Subject: [PATCH 243/624] [Refactor] Remove dead code multiple places (#60120) Signed-off-by: yewentao256 --- vllm/_custom_ops.py | 6 -- vllm/lora/model_manager.py | 4 - .../layers/fused_moe/experts/cpu_moe.py | 28 ------- vllm/model_executor/models/bert.py | 3 - vllm/model_executor/models/blip.py | 7 -- vllm/model_executor/models/diffusion_gemma.py | 11 --- vllm/model_executor/models/keye.py | 19 ----- vllm/model_executor/models/paddleocr_vl.py | 21 ----- vllm/model_executor/models/phi4mm_audio.py | 6 -- vllm/model_executor/models/step3p5.py | 9 --- .../warmup/jit_warmup_triton_helper.py | 81 ------------------- vllm/tokenizers/mistral.py | 8 -- 12 files changed, 203 deletions(-) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index eef1183697f3..e7862073083f 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -4464,12 +4464,6 @@ def cpu_gemm_wna16( return out -def cpu_activation_lut_bf16(input: torch.Tensor, activation: str) -> torch.Tensor: - out = torch.empty_like(input) - torch.ops._C.activation_lut_bf16(out, input, activation) - return out - - def cpu_prepack_moe_weight( weight: torch.Tensor, isa: str, diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index 8dc7b5605600..a7f9f02e77af 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -40,7 +40,6 @@ from vllm.model_executor.layers.fused_moe import MoERunner from vllm.model_executor.models import ( SupportsLoRA, - SupportsMultiModal, is_pooling_model, supports_multimodal, ) @@ -60,9 +59,6 @@ class SupportsLoRAModel(nn.Module, SupportsLoRA): ... -class SupportsLoRAMultiModalModel(SupportsLoRAModel, SupportsMultiModal): ... - - class AdapterLRUCache(LRUCache[int, T]): def __init__(self, capacity: int, deactivate_fn: Callable[[int], object]): super().__init__(capacity) diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py index 14b25ce30e37..54a9026f97b6 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py @@ -973,34 +973,6 @@ def apply( ) -# =========================================================================== -# INT8 W8A8 MoE -# =========================================================================== - - -def prepare_int8_moe_layer_for_cpu( - w13: torch.Tensor, - w2: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - """Prepack INT8 MoE weights for the current CPU architecture.""" - # SMMLA packing for AArch64 - if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: - return ( - cpu_prepack_moe_weight_int8(w13, "neon"), - cpu_prepack_moe_weight_int8(w2, "neon"), - ) - # VSX packing for POWER - if current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC: - return ( - cpu_prepack_moe_weight_int8(w13, "vsx"), - cpu_prepack_moe_weight_int8(w2, "vsx"), - ) - # VNNI packing for x86 - packed_w13 = torch.ops._C.convert_weight_packed(w13) - packed_w2 = torch.ops._C.convert_weight_packed(w2) - return packed_w13, packed_w2 - - class CPUExpertsInt8(mk.FusedMoEExpertsModular): """CPU INT8 W8A8 per-channel weight / dynamic per-token activation modular MoE experts.""" diff --git a/vllm/model_executor/models/bert.py b/vllm/model_executor/models/bert.py index 8f061b7284ef..cdaa0bce9725 100644 --- a/vllm/model_executor/models/bert.py +++ b/vllm/model_executor/models/bert.py @@ -564,9 +564,6 @@ def __init__( self.layer_norm = nn.LayerNorm(hidden_size, eps=layer_norm_eps) self.decoder = nn.Linear(hidden_size, vocab_size, bias=True) - def tie_weights_with_embeddings(self, embeddings_weight: torch.Tensor): - self.decoder.weight = embeddings_weight - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: x = self.dense(hidden_states) x = self.activation(x) diff --git a/vllm/model_executor/models/blip.py b/vllm/model_executor/models/blip.py index a4e821ec6b12..4f2275ecd21e 100644 --- a/vllm/model_executor/models/blip.py +++ b/vllm/model_executor/models/blip.py @@ -129,13 +129,6 @@ def __init__( prefix=f"{prefix}.attn", ) - def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): - return ( - tensor.view(bsz, seq_len, self.num_heads, self.head_dim) - .transpose(1, 2) - .contiguous() - ) - def forward( self, hidden_states: torch.Tensor, diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 7c64980b57bb..7f4f9ffba99e 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -269,17 +269,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): eps=getattr(text_config, "rms_norm_eps", 1e-6), ) - def compute_self_conditioning( - self, - inputs_embeds: torch.Tensor, - probs: torch.Tensor, - ) -> torch.Tensor: - embed_weight = self.model.embed_tokens.weight - soft_embeds = torch.matmul( - probs.to(embed_weight.dtype), embed_weight - ) * self.model.normalizer.to(inputs_embeds.dtype) - return self.self_conditioning(inputs_embeds, soft_embeds) - # ------------------------------------------------------------------ # # Multimodal: reuse Gemma4's image parsing, processing & embedding # ------------------------------------------------------------------ # diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py index 60d0d03b659d..83ff14dcf948 100644 --- a/vllm/model_executor/models/keye.py +++ b/vllm/model_executor/models/keye.py @@ -255,25 +255,6 @@ def interpolate_pos_encoding( patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) return patch_pos_embed - def fetch_position_embedding_lfu_cache(self, embeddings, h, w, max_cache: int = 20): - grid = (h, w) - if grid in self.cache_position_embedding: - self.cache_position_count[grid] += 1 - return self.cache_position_embedding[grid] - - if len(self.cache_position_embedding) >= max_cache: - min_hit_grid = min( - self.cache_position_count, - key=self.cache_position_count.__getitem__, - ) - self.cache_position_count.pop(min_hit_grid) - self.cache_position_embedding.pop(min_hit_grid) - - position_embedding = self.interpolate_pos_encoding(embeddings, h, w, True) - self.cache_position_count[grid] = 1 - self.cache_position_embedding[grid] = position_embedding - return position_embedding - def forward( self, pixel_values: torch.FloatTensor, diff --git a/vllm/model_executor/models/paddleocr_vl.py b/vllm/model_executor/models/paddleocr_vl.py index 1fde4cd0a11b..1c165c08ef5d 100644 --- a/vllm/model_executor/models/paddleocr_vl.py +++ b/vllm/model_executor/models/paddleocr_vl.py @@ -455,27 +455,6 @@ def interpolate_pos_encoding( patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) return patch_pos_embed - def fetch_position_embedding_lfu_cache( - self, embeddings: torch.Tensor, h: int, w: int, max_cache: int = 20 - ): - grid = (h, w) - if grid in self.cache_position_embedding: - self.cache_position_count[grid] += 1 - return self.cache_position_embedding[grid] - - if len(self.cache_position_embedding) >= max_cache: - min_hit_grid = min( - self.cache_position_count, - key=self.cache_position_count.__getitem__, - ) - self.cache_position_count.pop(min_hit_grid) - self.cache_position_embedding.pop(min_hit_grid) - - position_embedding = self.interpolate_pos_encoding(embeddings, h, w, True) - self.cache_position_count[grid] = 1 - self.cache_position_embedding[grid] = position_embedding - return position_embedding - def forward( self, pixel_values: torch.FloatTensor, diff --git a/vllm/model_executor/models/phi4mm_audio.py b/vllm/model_executor/models/phi4mm_audio.py index 9daedcf38e7b..151fcfc6ecdf 100644 --- a/vllm/model_executor/models/phi4mm_audio.py +++ b/vllm/model_executor/models/phi4mm_audio.py @@ -1256,12 +1256,6 @@ def __init__(self, config: PreTrainedConfig, **kwargs: Any) -> None: self.input_embeds = None self.audio_embed_sizes = None - def set_audio_embeds(self, input_embeds: torch.Tensor) -> None: - self.input_embeds = input_embeds - - def set_audio_embed_sizes(self, audio_embed_sizes: torch.Tensor) -> None: - self.audio_embed_sizes = audio_embed_sizes - def get_audio_features( self, input_embeds: torch.Tensor, diff --git a/vllm/model_executor/models/step3p5.py b/vllm/model_executor/models/step3p5.py index b56d4d3ed8d8..a6dffeb1d40f 100644 --- a/vllm/model_executor/models/step3p5.py +++ b/vllm/model_executor/models/step3p5.py @@ -17,7 +17,6 @@ get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, - get_tp_group, ) from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul, SwigluStepAndMul @@ -459,7 +458,6 @@ def __init__( f"Unsupported attention implementation: {config.att_impl_type}" ) self.use_moe = False - self.tp_group = get_tp_group() self.use_fused_all_reduce = ( get_tensor_model_parallel_world_size() > 1 and get_dp_group().world_size == 1 @@ -496,13 +494,6 @@ def __init__( ) self.prefix = prefix - def add_and_maybe_inplace_all_reduce( - self, in1: torch.Tensor, in2: torch.Tensor - ) -> torch.Tensor: - if not self.use_fused_all_reduce: - return in1 + in2 - return self.tp_group.all_reduce(in1 + in2) - def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor ) -> torch.Tensor: diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py index f7765cf23293..70ec8cd70fc5 100644 --- a/vllm/model_executor/warmup/jit_warmup_triton_helper.py +++ b/vllm/model_executor/warmup/jit_warmup_triton_helper.py @@ -718,65 +718,6 @@ def pointer( return TritonWarmupTensor(dtype, aligned=self.is_aligned(name), shape=shape) -def _literal_str_refs(node: ast.AST) -> tuple[str | int, ...]: - if isinstance(node, ast.Constant) and isinstance(node.value, str | int): - return (node.value,) - if isinstance(node, ast.List | ast.Tuple): - refs: list[str | int] = [] - for elt in node.elts: - if isinstance(elt, ast.Constant) and isinstance(elt.value, str | int): - refs.append(elt.value) - else: - raise ValueError( - f"Unsupported Triton specialization ref: {ast.dump(elt)}" - ) - return tuple(refs) - raise ValueError(f"Unsupported Triton specialization refs: {ast.dump(node)}") - - -def _normalize_arg_refs( - refs: tuple[str | int, ...], - arg_names: tuple[str, ...], -) -> frozenset[str]: - names: set[str] = set() - for ref in refs: - if isinstance(ref, int): - names.add(arg_names[ref]) - else: - names.add(ref) - return frozenset(names) - - -def _decorator_keyword_refs( - function_def: ast.FunctionDef, - keyword_name: str, -) -> tuple[str | int, ...]: - for decorator in function_def.decorator_list: - if not isinstance(decorator, ast.Call): - continue - decorator_name = get_ast_full_name(decorator.func) - if decorator_name not in ("triton.jit", "jit"): - continue - for keyword in decorator.keywords: - if keyword.arg == keyword_name: - return _literal_str_refs(keyword.value) - return () - - -def _triton_do_not_specialize_args( - kernel: Callable[..., Any], - function_def: ast.FunctionDef, - arg_names: tuple[str, ...], -) -> frozenset[str]: - refs = getattr(kernel, "do_not_specialize", None) - if refs is not None: - return _normalize_arg_refs(tuple(refs), arg_names) - return _normalize_arg_refs( - _decorator_keyword_refs(function_def, "do_not_specialize"), - arg_names, - ) - - def _triton_constexpr_arg_names( kernel: Callable[..., Any], function_def: ast.FunctionDef, @@ -821,25 +762,3 @@ def _pointer_arg_names( if name in candidate_names: pointer_names.add(name) return frozenset(pointer_names) - - -def trace_triton_kernel_specialization_args( - kernel: Callable[..., Any], -) -> tuple[str, ...]: - function_def = get_function_source_node(kernel) - if not isinstance(function_def, ast.FunctionDef): - raise ValueError("Expected Triton kernel to be defined as a function") - source_fn = getattr(kernel, "fn", kernel) - arg_names = tuple(inspect.signature(source_fn).parameters) - constexpr_args = _triton_constexpr_arg_names(kernel, function_def, arg_names) - do_not_specialize_args = _triton_do_not_specialize_args( - kernel, function_def, arg_names - ) - pointer_args = _pointer_arg_names(function_def, arg_names) - - return tuple( - name - for name in arg_names - if name in constexpr_args - or (name not in pointer_args and name not in do_not_specialize_args) - ) diff --git a/vllm/tokenizers/mistral.py b/vllm/tokenizers/mistral.py index 15cf0322581a..de321025a313 100644 --- a/vllm/tokenizers/mistral.py +++ b/vllm/tokenizers/mistral.py @@ -178,14 +178,6 @@ def _tekken_token_to_id(tokenizer: "Tekkenizer", t: str | bytes) -> int: return tokenizer.unk_id -def mistral_common_tekkenizer(tokenizer: object) -> "Tekkenizer | None": - """Return the underlying `Tekkenizer` for a `MistralCommonBackend`.""" - mistral = getattr(tokenizer, "tokenizer", None) - instruct = getattr(mistral, "instruct_tokenizer", None) - tekken = getattr(instruct, "tokenizer", None) - return tekken if isinstance(tekken, Tekkenizer) else None - - def tekken_convert_ids_to_tokens( tokenizer: "Tekkenizer", ids: Sequence[int] ) -> list[str | bytes]: From f2376cdc4da852ec48ab5bcb747fdfa65a4dfd4a Mon Sep 17 00:00:00 2001 From: siyu Date: Sat, 10 Oct 2026 23:48:06 +0800 Subject: [PATCH 244/624] [Fast Start] Preload the FlashInfer autotune table on the weight cache daemon (#60085) Signed-off-by: liusy58 Signed-off-by: Isotr0py Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Isotr0py --- docs/features/preload.md | 30 +++++++- .../model_loader/test_weight_cache.py | 73 ++++++++++++++++++- .../test_flashinfer_autotune_warmup.py | 31 ++++++++ vllm/entrypoints/cli/preload.py | 9 ++- .../model_loader/weight_cache/daemon.py | 39 ++++++++++ .../model_loader/weight_cache/utils.py | 35 ++++++++- .../warmup/flashinfer_autotune_cache.py | 16 +++- vllm/model_executor/warmup/kernel_warmup.py | 10 ++- vllm/v1/worker/gpu/model_runner.py | 12 ++- vllm/v1/worker/gpu_model_runner.py | 15 +++- 10 files changed, 253 insertions(+), 17 deletions(-) diff --git a/docs/features/preload.md b/docs/features/preload.md index 74839bf35b38..6e8fad9ad03e 100644 --- a/docs/features/preload.md +++ b/docs/features/preload.md @@ -50,6 +50,26 @@ vllm serve meta-llama/Llama-3.1-8B-Instruct \ The daemon itself must load from disk; passing `--load-format ipc_cache` to `vllm preload` is an error. +## Preloading the FlashInfer autotune cache + +FlashInfer has several implementations of each operation and chooses between +them by benchmarking. vLLM runs that pass during kernel warmup and persists the +result per rank in the on-disk FlashInfer autotune cache (under +`VLLM_CACHE_ROOT`, or `VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR`), so a later engine +with the same configuration loads the tuned configs instead of profiling again. +Normally the first serving engine pays for the pass, and on a large MoE model it +dominates what is left of the startup time once the weights come from the +daemon. + +Unless FlashInfer autotune is disabled (`--no-enable-flashinfer-autotune`), +`vllm preload` pays it at preload time: each daemon runs the autotune pass in +place — profiling dummy runs against the shard it already holds — after +loading and before binding its socket, then writes the cache: + +```bash +vllm preload --model /path/to/model --tensor-parallel-size 4 +``` + ## How it works 1. `vllm preload` spawns one daemon process per GPU. Each daemon loads its @@ -164,8 +184,14 @@ Socket paths are derived from the GPU UUID, so they are stable regardless of - **Platform**: CUDA and ROCm only; other platforms raise `UnsupportedPlatformForIPCError` even when `fallback` is enabled, since it is a permanent misconfiguration rather than a transient daemon outage. -- **Parallelism**: tensor, expert and data parallelism are supported; - launching the daemon with pipeline parallelism is rejected. +- **Parallelism**: tensor, expert, pipeline and data parallelism are + supported, including across nodes. +- **Autotune preloading**: tuning runs across the daemon's world group, + including across nodes and DP ranks. With model-based speculative decoding + both daemon groups skip it (the target cannot run the drafter's dummy + passes, and a table tuned on the draft model would key on the draft's + config, which no engine looks up), so those deployments' first engine + tunes and fills the cache instead. - **Quantization**: every quantization method in the model must declare support for pre-processed weights (the daemon transfers weights *after* quantization post-processing). Unsupported methods raise diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py index 7960a237914b..9628cb47e0bb 100644 --- a/tests/model_executor/model_loader/test_weight_cache.py +++ b/tests/model_executor/model_loader/test_weight_cache.py @@ -12,12 +12,14 @@ import sys import tempfile import threading +import time import urllib.error import urllib.request from dataclasses import dataclass, field from typing import Any import pytest +import regex as re from vllm import SamplingParams from vllm.assets.image import ImageAsset @@ -81,6 +83,21 @@ def check_health(self) -> None: f"Weight cache daemon is not ready: HTTP {error.code}" ) from error + def wait_until_ready(self, timeout_s: float = 900) -> None: + """Poll until READY (which covers the autotune tuners) or the preload + process dies.""" + deadline = time.monotonic() + timeout_s + while True: + try: + self.check_health() + return + except (AssertionError, urllib.error.URLError): + if time.monotonic() > deadline: + raise + if self._proc is not None and self._proc.poll() is not None: + raise AssertionError("vllm preload exited before ready") from None + time.sleep(2) + def _stop(self) -> None: assert self._proc is not None self._proc.terminate() @@ -99,6 +116,10 @@ class ModelCase: images: list | None = None llm_kwargs: dict[str, Any] = field(default_factory=dict) daemon_args: list[str] = field(default_factory=list) + # Additionally assert the daemons tune in place before serving and write + # the on-disk FlashInfer autotune cache, and the restart engine reuses it + # (needs SM90+/FlashInfer). + check_flashinfer_cache: bool = False def generate( @@ -157,6 +178,9 @@ def generate( max_model_len=4096, ), daemon_args=["--trust-remote-code"], + # Kimi K3 has tunable FlashInfer ops (a dense model like Qwen2.5 tunes + # nothing and upstream never writes the cache file for it). + check_flashinfer_cache=True, ) # Qwen3.5-0.8B ships one MTP layer in the target checkpoint, so method="mtp" @@ -186,18 +210,38 @@ def generate( [QWEN_CASE, K3_CASE, QWEN_MTP_CASE], ids=["qwen3.5", "kimi-k3", "qwen3.5-mtp"], ) -def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase): +def test_ipc_cache_cold_start_and_warm_restart( + vllm_runner, case: ModelCase, tmp_path, monkeypatch, capfd +): """Cold start falls back to disk; warm restarts load weights via CUDA IPC. All runs must produce outputs identical to a default-loader baseline. The warm runs disable the disk fallback, so they only pass if the weights really came from the daemon — for the MTP case, both the target's and the draft's daemon groups. + + Cases with ``check_flashinfer_cache`` additionally cover in-daemon + autotune (default-on with `--enable-flashinfer-autotune`): each daemon + tunes before binding its socket and must write the on-disk FlashInfer + autotune cache, and the restart engine must load tactics from it instead + of re-profiling. The daemon's cache file is keyed with the + OPENAI_API_SERVER batch defaults while these engines run under LLM_CLASS, + so they never hit the daemon's file by design; reuse is proven against + the file the earlier in-process engines wrote. """ if not current_platform.is_cuda_alike(): pytest.skip("Weight cache IPC sharing requires CUDA or ROCm") if case is K3_CASE and not current_platform.is_device_capability_family(100): pytest.skip("Kimi K3 IPC weight cache requires an SM100 MXFP4 backend") + if case.check_flashinfer_cache: + from vllm.utils.flashinfer import has_flashinfer + + if not (has_flashinfer() and current_platform.has_device_capability(90)): + pytest.skip("FlashInfer autotune requires FlashInfer and SM90+") + # Isolate the on-disk autotune cache. The engines are in-process + # (envs read lazily); the daemon subprocesses spawn later and + # inherit it. + monkeypatch.setenv("VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", str(tmp_path)) # Baseline: plain disk loading with the default loader. baseline_outputs = generate( @@ -212,15 +256,40 @@ def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase): with tempfile.TemporaryDirectory(prefix="vllm_ipc_empty_") as empty_socket_dir: cold_outputs = generate(vllm_runner, case, empty_socket_dir, fallback=True) + if case.check_flashinfer_cache: + # The baseline engine already wrote its own (LLM_CLASS-keyed) table. + engine_cache_files = set(tmp_path.rglob("autotune_configs*.json")) + with WeightCacheDaemon( case.model, tp_size=1, extra_args=case.daemon_args, ) as d: warm_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False) - d.check_health() + d.wait_until_ready() + if case.check_flashinfer_cache: + # Each daemon tunes before its ready report; exactly one new file + # (their OPENAI_API_SERVER-keyed table) must appear. + tuner_files = set(tmp_path.rglob("autotune_configs*.json")) + tuner_files -= engine_cache_files + assert len(tuner_files) == 1, ( + f"expected the tuner's tuned table, got {tuner_files}" + ) # Warm restart: a second engine lifetime against the same daemon. + if case.check_flashinfer_cache: + capfd.readouterr() # drain: the assertion reads the increment restart_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False) + if case.check_flashinfer_cache: + # The restart's EngineCore subprocess inherits pytest's fds; the + # flashinfer autotuner logs "[Autotuner]: Loaded N configs from + # " (INFO, not once-deduplicated) on a disk cache hit. + out = capfd.readouterr() + m = re.search( + r"\[Autotuner\]: Loaded (\d+) configs from", out.out + out.err + ) + assert m and int(m.group(1)) > 0, ( + "restart engine did not load the on-disk autotune cache" + ) assert cold_outputs == baseline_outputs assert warm_outputs == baseline_outputs diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py index f7ea6847f26f..59046a5f3e8d 100644 --- a/tests/model_executor/test_flashinfer_autotune_warmup.py +++ b/tests/model_executor/test_flashinfer_autotune_warmup.py @@ -386,3 +386,34 @@ def test_rejected_cache_retunes_every_rank_in_its_tuning_group(autotune_run, pp, rerun.assert_collectives_match() assert set(rerun.profile_groups) == set(range(tp)) assert {rank for rank, _, _ in rerun.saves} == set(range(tp)) + + +def test_autotune_cache_key_normalizes_kv_init_fields(): + """mamba_block_size/kv_cache_layout only resolve at engine KV-init time; + they must not feed the autotune cache key, or nothing can compute the key + before startup (the daemon tunes before any KV init). block_size must + stay in the key: the daemon replays the engine's block-size resolution + before tuning, and different block sizes may select different ops.""" + from vllm.config.cache import CacheConfig + from vllm.model_executor.warmup.flashinfer_autotune_cache import ( + _normalize_cache_config_for_hash, + ) + + fresh = CacheConfig() + # Simulate the engine's KV-init write-backs (hybrid model values). + mutated = CacheConfig() + mutated.mamba_block_size = 512 + mutated.kv_cache_layout = "LBHNC" + + # Negative control: the write-backs are hashed factors. + assert mutated.compute_hash() != fresh.compute_hash() + normalized = _normalize_cache_config_for_hash(mutated) + assert normalized.compute_hash() == fresh.compute_hash() + # Factory state is returned untouched. + assert _normalize_cache_config_for_hash(fresh) is fresh + # block_size stays significant (the daemon keys on the resolved value). + resolved = CacheConfig(block_size=64) + assert ( + _normalize_cache_config_for_hash(resolved).compute_hash() + != fresh.compute_hash() + ) diff --git a/vllm/entrypoints/cli/preload.py b/vllm/entrypoints/cli/preload.py index 9eab7421f175..90e746fdb686 100644 --- a/vllm/entrypoints/cli/preload.py +++ b/vllm/entrypoints/cli/preload.py @@ -39,6 +39,7 @@ from vllm.model_executor.model_loader.weight_cache.protocol import ( check_ipc_platform_support, ) +from vllm.usage.usage_lib import UsageContext from vllm.utils.network_utils import get_distributed_init_method, get_open_port if typing.TYPE_CHECKING: @@ -157,7 +158,13 @@ def cmd(args: argparse.Namespace) -> None: ) engine_args = EngineArgs.from_cli_args(args) - vllm_config = engine_args.create_engine_config() + # Warmup context for flashinfer autotune. + usage_context = ( + UsageContext.OPENAI_API_SERVER + if engine_args.kernel_config.enable_flashinfer_autotune is not False + else None + ) + vllm_config = engine_args.create_engine_config(usage_context=usage_context) if vllm_config.load_config.load_format == "ipc_cache": raise ValueError( "The weight cache daemon itself must load from disk; use the " diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index 2c6944f1eb14..0fc63da7869b 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/vllm/model_executor/model_loader/weight_cache/daemon.py @@ -112,11 +112,14 @@ verify_peer_is_owner, ) from vllm.model_executor.model_loader.weight_cache.utils import ( + build_warmup_runner, export_model_attrs, format_daemon_role, is_draft_model_cacheable, ) +from vllm.model_executor.warmup.kernel_warmup import flashinfer_autotune from vllm.platforms import current_platform +from vllm.utils.flashinfer import has_flashinfer from vllm.utils.mem_utils import format_gib from vllm.utils.torch_utils import set_default_torch_dtype from vllm.v1.worker.workspace import init_workspace_manager @@ -236,6 +239,41 @@ def __init__( is_draft=is_draft, ) + def warmup(self) -> None: + """Run the FlashInfer autotune pass against the cached model.""" + vllm_config = self.vllm_config + # Mirrors the gate in the engine's kernel warmup. + if ( + vllm_config.kernel_config.enable_flashinfer_autotune is False + or not has_flashinfer() + or not current_platform.has_device_capability(90) + or is_draft_model_cacheable(vllm_config.speculative_config) + ): + return + try: + current_platform.update_block_size_for_backend(vllm_config) + runner = build_warmup_runner( + vllm_config, + self.local_rank, + is_draft=self.is_draft, + model_config=self.model_config, + ) + assert self.model is not None, "warmup ran before load_model" + runner.load_model(model=self.model) + # The V2 runner requires skip_attn when no KV cache exists + flashinfer_autotune(runner, skip_attn=vllm_config.use_v2_model_runner) + logger.info( + "Weight cache %s daemon rank %d tuned FlashInfer; the tuned " + "configs are in the on-disk autotune cache", + self.role, + self.global_rank, + ) + finally: + # Free the runner and the dummy-run activations; the daemon's + # weights are untouched. + gc.collect() + torch.accelerator.empty_cache() + def load_model(self) -> None: torch.accelerator.set_device_index(self.local_rank) # Outside the config context so the daemon keeps its own rendezvous @@ -457,6 +495,7 @@ def _run_daemon( pp_rank, ) daemon.load_model() + daemon.warmup() daemon.serve_forever( ready_callback=lambda: ready_queue.put((daemon.role, daemon.global_rank)) ) diff --git a/vllm/model_executor/model_loader/weight_cache/utils.py b/vllm/model_executor/model_loader/weight_cache/utils.py index 1abbf572f847..bc5de3cc9768 100644 --- a/vllm/model_executor/model_loader/weight_cache/utils.py +++ b/vllm/model_executor/model_loader/weight_cache/utils.py @@ -7,10 +7,17 @@ wire format. """ -from typing import Any +from typing import TYPE_CHECKING, Any, cast -from vllm.config import SpeculativeConfig +import torch + +from vllm.config import SpeculativeConfig, VllmConfig, replace, set_current_vllm_config from vllm.model_executor.models.interfaces import SupportsEagleBase +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from vllm.config import ModelConfig + from vllm.v1.worker.gpu_model_runner import GPUModelRunner # Speculative methods whose draft model the daemon caches in its own group. # Other drafts keep loading from disk in the engine. @@ -46,3 +53,27 @@ def format_daemon_role(is_draft: bool) -> str: def format_socket_role_suffix(is_draft: bool) -> str: """Socket-name suffix keeping the draft group distinct from the target.""" return "_draft" if is_draft else "" + + +def build_warmup_runner( + vllm_config: VllmConfig, + local_rank: int, + *, + is_draft: bool, + model_config: "ModelConfig", +) -> "GPUModelRunner": + """Create a model runner for the daemon's warmup dummy runs.""" + if is_draft: + vllm_config = replace( + vllm_config, model_config=model_config, speculative_config=None + ) + device = torch.device(current_platform.device_type, local_rank) + with set_current_vllm_config(vllm_config): + if vllm_config.use_v2_model_runner: + from vllm.v1.worker.gpu.model_runner import GPUModelRunner as Runner + else: + from vllm.v1.worker.gpu_model_runner import GPUModelRunner as Runner + runner = Runner(vllm_config, device) + # The two runner classes share the warmup surface; the nominal type + # stays V1's GPUModelRunner. + return cast("GPUModelRunner", runner) diff --git a/vllm/model_executor/warmup/flashinfer_autotune_cache.py b/vllm/model_executor/warmup/flashinfer_autotune_cache.py index afc3ef3ac014..ca2c4d1421c8 100644 --- a/vllm/model_executor/warmup/flashinfer_autotune_cache.py +++ b/vllm/model_executor/warmup/flashinfer_autotune_cache.py @@ -10,14 +10,28 @@ from typing import TYPE_CHECKING import vllm.envs as envs +from vllm.config import CacheConfig, VllmConfig, replace if TYPE_CHECKING: from vllm.distributed.parallel_state import GroupCoordinator from vllm.v1.worker.gpu_model_runner import GPUModelRunner +def _normalize_cache_config_for_hash(cache_config: CacheConfig) -> CacheConfig: + """Normalize cache config for hashing""" + if cache_config.mamba_block_size is None and cache_config.kv_cache_layout is None: + return cache_config + normalized = replace(cache_config, mamba_block_size=None) + assert normalized.kv_cache_layout is None + return normalized + + def flashinfer_autotune_cache_hash(runner: "GPUModelRunner") -> str: - config_hash = runner.vllm_config.compute_hash(include_version=False) + vllm_config: VllmConfig = runner.vllm_config + cache_config = _normalize_cache_config_for_hash(vllm_config.cache_config) + if cache_config is not vllm_config.cache_config: + vllm_config = replace(vllm_config, cache_config=cache_config) + config_hash = vllm_config.compute_hash(include_version=False) return hashlib.sha256(config_hash.encode()).hexdigest() diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index e5751730e85d..685b33a5652a 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -406,7 +406,7 @@ def _all_ranks_have_matching_cache(path: Path, group) -> bool: return fingerprint is not None and all(f == fingerprint for f in gathered) -def flashinfer_autotune(runner: "GPUModelRunner") -> None: +def flashinfer_autotune(runner: "GPUModelRunner", *, skip_attn: bool = False) -> None: """Autotune FlashInfer operations. FlashInfer have many implementations for the same operation, autotuning runs benchmarks for each implementation and stores @@ -415,6 +415,10 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: Without autotuning, FlashInfer will rely on heuristics, which may be significantly slower. + Args: + runner: The model runner to tune. + skip_attn: Additionally skip attention in the tuning dummy runs. + With PP > 1, stages run different layers and may profile different ops, so each stage's TP group tunes separately with its own cache file; otherwise the world group tunes together. Per-tactic timings are @@ -425,6 +429,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: rank. A rank with a cache hit skips the per-tactic reduce the others block in, so ranks keep loaded configs only if every rank in the tuning group has a matching file and successfully loads it. + """ from flashinfer.autotuner import AutoTuner, set_autotune_process_group @@ -485,11 +490,12 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: hisparse_enabled = ( runner.vllm_config.attention_config.hisparse_config is not None ) + skip_attn = skip_attn or hisparse_enabled if hisparse_enabled: # HiSparse hot-buffer attention is bounded by decode batch # size, not the prefill-sized batch used for the full model. autotune_hisparse_flashinfer_attention(runner) - _run_flashinfer_autotune_dummy_runs(runner, skip_attn=hisparse_enabled) + _run_flashinfer_autotune_dummy_runs(runner, skip_attn=skip_attn) replayssm_autotune_warmup(runner) _autotune_kimi_k3_kda_qkvg(runner.get_model()) finally: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 674cf23527d1..e04a381b94e6 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -373,9 +373,15 @@ def get_supported_tasks(self) -> tuple[SupportedTask, ...]: tasks.extend(PoolingRunner.get_supported_tasks(self.model)) return tuple(tasks) - def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: + def load_model( + self, + load_dummy_weights: bool = False, + *args, + model: nn.Module | None = None, + **kwargs, + ) -> None: time_before_load = time.perf_counter() - if load_dummy_weights: + if model is None and load_dummy_weights: self.load_config.load_format = "dummy" self.eplb.prepare_load() eplb_models_added = False @@ -385,7 +391,7 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: # Capture warmup providers selected while constructing the model. with self.jit_warmup_registry.activate(): - self.model = model_loader.load_model( + self.model = model or model_loader.load_model( vllm_config=self.vllm_config, model_config=self.vllm_config.model_config, ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 22ad662b9018..90bea22c9d4f 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -5217,9 +5217,15 @@ def update_config(self, overrides: dict[str, Any]) -> None: setattr(self, config_name, new_config) @instrument(span_name="Loading (GPU)") - def load_model(self, load_dummy_weights: bool = False) -> None: + def load_model( + self, + load_dummy_weights: bool = False, + *, + model: nn.Module | None = None, + ) -> None: """Args: load_dummy_weights: load dummy weights instead of real weights. + model: a pre-loaded model to skip the model loader. """ logger.info_once( @@ -5235,13 +5241,14 @@ def load_model(self, load_dummy_weights: bool = False) -> None: try: with DeviceMemoryProfiler() as m: time_before_load = time.perf_counter() - if load_dummy_weights: + if model is None and load_dummy_weights: self.load_config.load_format = "dummy" model_loader = get_model_loader(self.load_config) # Capture warmup providers selected while constructing the model. with self.jit_warmup_registry.activate(): # type: ignore[attr-defined] - self.model = model_loader.load_model( - vllm_config=self.vllm_config, model_config=self.model_config + self.model = model or model_loader.load_model( + vllm_config=self.vllm_config, + model_config=self.model_config, ) lookback_depth = getattr(self.model, "token_lookback_depth", 0) if lookback_depth > 0: From 187a0eb98aa42341d703f83421d693fa7585581b Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Sun, 11 Oct 2026 00:36:03 +0800 Subject: [PATCH 245/624] [Bugfix][Frontend] Forward --enable-log-outputs to the Anthropic handler (#60982) Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> --- vllm/entrypoints/anthropic/serving.py | 4 ++++ vllm/entrypoints/generate/api_router.py | 2 ++ 2 files changed, 6 insertions(+) diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index 7b1cf95e99b0..22319f1c70c7 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -163,6 +163,8 @@ def __init__( tool_strict_level: str = "auto", enable_prompt_tokens_details: bool = False, enable_force_include_usage: bool = False, + enable_log_outputs: bool = False, + enable_log_deltas: bool = True, default_chat_template_kwargs: dict[str, Any] | None = None, disabled_thinking_effort: AnthropicDisabledThinkingEffortOption = "auto", ): @@ -181,6 +183,8 @@ def __init__( tool_parser=tool_parser, enable_prompt_tokens_details=enable_prompt_tokens_details, enable_force_include_usage=enable_force_include_usage, + enable_log_outputs=enable_log_outputs, + enable_log_deltas=enable_log_deltas, default_chat_template_kwargs=default_chat_template_kwargs, ) self._merge_inline_system = self._should_merge_inline_system(online_renderer) diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 66a37d6b13ee..6bf8d690265a 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -193,6 +193,8 @@ async def init_generate_state( reasoning_parser=args.structured_outputs_config.reasoning_parser, enable_prompt_tokens_details=args.enable_prompt_tokens_details, enable_force_include_usage=args.enable_force_include_usage, + enable_log_outputs=args.enable_log_outputs, + enable_log_deltas=args.enable_log_deltas, default_chat_template_kwargs=default_chat_template_kwargs, disabled_thinking_effort=args.anthropic_disabled_thinking_effort, ) From 276fbcff2717bd934cfa37c8a2e4c391f3e7237b Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Sat, 10 Oct 2026 15:45:33 -0400 Subject: [PATCH 246/624] [Deprecation] Remove scheduled deprecated items (#60923) Signed-off-by: yewentao256 --- rust/src/server/src/routes/pause.rs | 4 --- vllm/benchmarks/throughput.py | 25 ++----------------- vllm/engine/protocol.py | 2 -- .../pooling/pooling/io_processor.py | 21 +--------------- vllm/entrypoints/rl/online/api_router.py | 3 --- vllm/scripts.py | 17 ------------- vllm/v1/engine/async_llm.py | 12 --------- 7 files changed, 3 insertions(+), 81 deletions(-) delete mode 100644 vllm/scripts.py diff --git a/rust/src/server/src/routes/pause.rs b/rust/src/server/src/routes/pause.rs index 1654c85b0200..ed0c175b84c3 100644 --- a/rust/src/server/src/routes/pause.rs +++ b/rust/src/server/src/routes/pause.rs @@ -39,10 +39,6 @@ fn invalid_query(error: QueryRejection) -> ApiError { ApiError::invalid_request(error.body_text(), Some("mode")) } -// TODO: the Python frontend also accepts the deprecated -// `wait_for_inflight_requests` flag (equivalent to `mode="wait"`); it is -// intentionally omitted here in favor of the `mode` parameter. - /// Pause the scheduler so generation can be halted (e.g. for weight updates). pub async fn pause( State(state): State>, diff --git a/vllm/benchmarks/throughput.py b/vllm/benchmarks/throughput.py index 6245b20e86c8..76f6a72fd98a 100644 --- a/vllm/benchmarks/throughput.py +++ b/vllm/benchmarks/throughput.py @@ -625,15 +625,7 @@ def assign_loras(requests, args): def validate_args(args): """Validate command-line arguments.""" - # === Deprecation and Defaulting === - if args.dataset is not None: - warnings.warn( - "The '--dataset' argument will be deprecated in the next release. " - "Please use '--dataset-name' and '--dataset-path' instead.", - stacklevel=2, - ) - args.dataset_path = args.dataset - + # === Defaulting === if not getattr(args, "tokenizer", None): args.tokenizer = args.model @@ -647,11 +639,7 @@ def validate_args(args): raise ValueError("--prequeue-requests is not supported with --async-engine") # === Dataset Configuration === - if ( - not args.dataset - and not args.dataset_path - and args.dataset_name not in {"prefix_repetition"} - ): + if not args.dataset_path and args.dataset_name != "prefix_repetition": print("When dataset path is not set, it will default to random dataset") args.dataset_name = "random" random_input_len = getattr(args, "random_input_len", None) @@ -809,15 +797,6 @@ def add_cli_args(parser: FlexibleArgumentParser): help="Name of the dataset to benchmark on.", default="sharegpt", ) - parser.add_argument( - "--dataset", - type=str, - default=None, - help="Path to the ShareGPT dataset, will be deprecated in\ - the next release. The dataset is expected to " - "be a json in form of list[dict[..., conversations: " - "list[dict[..., value: ]]]]", - ) parser.add_argument( "--dataset-path", type=str, default=None, help="Path to the dataset" ) diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index 2bb128ba073e..15f6b801bcd6 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -224,7 +224,6 @@ async def pause_generation( self, *, mode: "PauseMode" = "abort", - wait_for_inflight_requests: bool = False, clear_cache: bool = True, ) -> None: """Pause new generation/encoding requests. @@ -236,7 +235,6 @@ async def pause_generation( - ``"wait"``: Wait for in-flight requests to complete. - ``"keep"``: Freeze requests in queue; they resume on :meth:`resume_generation`. - wait_for_inflight_requests: DEPRECATED. Use ``mode="wait"`` instead. clear_cache: DEPRECATED. Whether to clear KV and prefix caches after draining. diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py index d5132f2a6649..74e7351c258a 100644 --- a/vllm/entrypoints/pooling/pooling/io_processor.py +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -4,7 +4,6 @@ from typing import Any from vllm import PoolingParams, PoolingRequestOutput -from vllm.logger import init_logger from vllm.plugins.io_processors import get_io_processor from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq @@ -21,8 +20,6 @@ ) from .protocol import IOProcessorRequest, IOProcessorResponse -logger = init_logger(__name__) - class UnsupportedCombinedTaskIOProcessor(PoolingIOProcessor): name = "embed&token_classify" @@ -117,23 +114,7 @@ def post_process_online( ctx.final_res_batch, request_id=ctx.request_id, ) - - if callable( - output_to_response := getattr(self.io_processor, "output_to_response", None) - ): - logger.warning_once( - "`IOProcessor.output_to_response` is deprecated. To ensure " - "consistency between offline and online APIs, " - "`IOProcessorResponse` will become a transparent wrapper " - "around output data from v0.19 onwards.", - ) - - if hasattr(output, "request_id") and output.request_id is None: - output.request_id = ctx.request_id # type: ignore - - ctx.response = output_to_response(output) # type: ignore - else: - ctx.response = IOProcessorResponse(request_id=ctx.request_id, data=output) + ctx.response = IOProcessorResponse(request_id=ctx.request_id, data=output) ####################################### # offline APIs diff --git a/vllm/entrypoints/rl/online/api_router.py b/vllm/entrypoints/rl/online/api_router.py index ef9269749f3b..51182fd99683 100644 --- a/vllm/entrypoints/rl/online/api_router.py +++ b/vllm/entrypoints/rl/online/api_router.py @@ -32,7 +32,6 @@ def engine_client(request: Request) -> EngineClient: async def pause_generation( raw_request: Request, mode: Annotated[PauseMode, Query()] = "abort", - wait_for_inflight_requests: bool = Query(False), clear_cache: Annotated[bool, Query()] = True, ) -> JSONResponse: """Pause generation requests to allow weight updates. @@ -44,7 +43,6 @@ async def pause_generation( - ``"abort"``: Abort all in-flight requests immediately (default). - ``"wait"``: Wait for in-flight requests to complete. - ``"keep"``: Freeze requests in queue; they resume on /resume. - wait_for_inflight_requests: DEPRECATED. Use ``mode="wait"`` instead. clear_cache: DEPRECATED. Whether to clear KV/prefix caches after draining. Ignored when mode="keep". @@ -55,7 +53,6 @@ async def pause_generation( await engine.pause_generation( mode=mode, clear_cache=clear_cache, - wait_for_inflight_requests=wait_for_inflight_requests, ) return JSONResponse( content={"status": "paused"}, diff --git a/vllm/scripts.py b/vllm/scripts.py deleted file mode 100644 index f158860726be..000000000000 --- a/vllm/scripts.py +++ /dev/null @@ -1,17 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from vllm.entrypoints.cli.main import main as vllm_main -from vllm.logger import init_logger - -logger = init_logger(__name__) - - -# Backwards compatibility for the move from vllm.scripts to -# vllm.entrypoints.cli.main -def main(): - logger.warning( - "vllm.scripts.main() is deprecated. Please re-install " - "vllm or use vllm.entrypoints.cli.main.main() instead." - ) - vllm_main() diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index bdab43e92528..08445d8f913a 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -4,7 +4,6 @@ import os import socket import time -import warnings from collections.abc import AsyncGenerator, Iterable, Mapping from copy import copy from typing import Any @@ -944,7 +943,6 @@ async def pause_generation( self, *, mode: PauseMode = "abort", - wait_for_inflight_requests: bool | None = None, clear_cache: bool = True, ) -> None: """Pause generation to allow model weight updates. @@ -960,20 +958,10 @@ async def pause_generation( - ``"wait"``: Wait for in-flight requests to complete. - ``"keep"``: Freeze requests in queue; they resume on :meth:`resume_generation`. - wait_for_inflight_requests: DEPRECATED: use mode argument. clear_cache: Whether to clear KV cache and prefix cache after draining. Set to ``False`` to preserve cache for faster resume. """ - if wait_for_inflight_requests: - warnings.warn( - "The `wait_for_inflight_requests` parameter in " - "`AsyncLLM.pause_generation()` is deprecated. " - "Please use `mode` argument instead.", - DeprecationWarning, - stacklevel=2, - ) - mode = "wait" if clear_cache: await self.renderer.clear_mm_cache_async() await self.engine_core.pause_scheduler_async(mode=mode, clear_cache=clear_cache) From 3c062d391d15cfa41d3f321cad24f3030c892a97 Mon Sep 17 00:00:00 2001 From: Subhash Polisetti <49185065+subhashpolisetti@users.noreply.github.com> Date: Sat, 10 Oct 2026 18:41:06 -0700 Subject: [PATCH 247/624] [Bugfix] Support 0-dim tensors in packed weight transfer (#52235) Signed-off-by: subhash polisetti Signed-off-by: aoshen02 Co-authored-by: aoshen02 Co-authored-by: Codex --- tests/distributed/test_packed_tensor.py | 29 ++++++++++++++++++- .../weight_transfer/packed_tensor.py | 6 ++-- 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/tests/distributed/test_packed_tensor.py b/tests/distributed/test_packed_tensor.py index 5b6846af7683..63af7f4b56fc 100644 --- a/tests/distributed/test_packed_tensor.py +++ b/tests/distributed/test_packed_tensor.py @@ -618,6 +618,7 @@ def test_roundtrip_basic(self): """Test basic IPC producer -> consumer roundtrip.""" params = [ ("w1", torch.randn(10, 20, dtype=torch.float32).cuda()), + ("scale", torch.tensor(1.5, dtype=torch.float32).cuda()), ("w2", torch.randn(5, dtype=torch.float16).cuda()), ] gpu_uuid = self._get_gpu_uuid() @@ -634,7 +635,7 @@ def test_roundtrip_basic(self): ) assert num_chunks == 1 - assert len(result) == 2 + assert len(result) == len(params) for (orig_name, orig_tensor), (res_name, res_tensor) in zip(params, result): assert orig_name == res_name assert res_tensor.shape == orig_tensor.shape @@ -746,3 +747,29 @@ def test_consumer_wrong_uuid_raises(self): tensor_sizes=c.tensor_sizes, device_index=torch.cuda.current_device(), ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_pack_unpack_scalar(dtype): + """Preserve scalar shape, dtype and value alongside a matrix on CPU.""" + params = [ + ("weight", torch.randn(2, 3, dtype=dtype)), + ("scale", torch.tensor(1.5, dtype=dtype)), + ] + + chunk = pack_tensors(iter(params), lambda item: item[1], buffer_size_bytes=1024) + assert chunk is not None + + result = unpack_tensor( + chunk.packed_tensor, + chunk.names, + chunk.shapes, + chunk.dtypes, + chunk.tensor_sizes, + ) + + for (expected_name, expected), (name, actual) in zip(params, result, strict=True): + assert name == expected_name + assert actual.shape == expected.shape + assert actual.dtype == expected.dtype + torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/vllm/distributed/weight_transfer/packed_tensor.py b/vllm/distributed/weight_transfer/packed_tensor.py index dcd8efc8c4af..2e715383801f 100644 --- a/vllm/distributed/weight_transfer/packed_tensor.py +++ b/vllm/distributed/weight_transfer/packed_tensor.py @@ -50,7 +50,7 @@ def unpack_tensor( unpacked_tensors = packed_tensor.split(tensor_sizes) return [ - (name, tensor.contiguous().view(dtype).view(*shape)) + (name, tensor.contiguous().view(dtype).view(shape)) for name, shape, dtype, tensor in zip(names, shapes, dtypes, unpacked_tensors) ] @@ -105,7 +105,7 @@ def pack_tensors( name, orig_tensor = item # Apply post processing and convert to linearized uint8 tensor - tensor = post_iter_func(item).contiguous().view(torch.uint8).view(-1) + tensor = post_iter_func(item).contiguous().view(-1).view(torch.uint8) if tensor.numel() > buffer_size_bytes: import warnings @@ -350,7 +350,7 @@ def packed_ipc_producer( for name, orig_tensor in iterator: flat = ( - post_iter_func((name, orig_tensor)).contiguous().view(torch.uint8).view(-1) + post_iter_func((name, orig_tensor)).contiguous().view(-1).view(torch.uint8) ) if flat.numel() > buffer_size_bytes: From b89852aad0854cc8df9138bc412b9a75c1fa8120 Mon Sep 17 00:00:00 2001 From: Shuolei Wang <46160365+ShuoleiWang@users.noreply.github.com> Date: Sun, 11 Oct 2026 11:33:29 +0800 Subject: [PATCH 248/624] [RL][Frontend] Return token-level weight-version spans on final outputs (#59582) Signed-off-by: Shuolei Wang Signed-off-by: Shuolei Wang <948904026@qq.com> Co-authored-by: aoshen02 --- .../online_serving/token_in_token_out.md | 15 ++++++ docs/training/async_rl.md | 20 ++++++- examples/rl/rlhf_async_new_apis.py | 43 +++++++-------- .../engine-core-client/src/protocol/output.rs | 3 ++ .../src/tests/python_compat.py | 3 ++ tests/v1/engine/test_output_processor.py | 53 +++++++++++++++++++ .../scale_out/token_in_token_out/protocol.py | 5 ++ .../scale_out/token_in_token_out/serving.py | 2 + vllm/outputs.py | 14 +++++ vllm/v1/engine/__init__.py | 3 ++ vllm/v1/engine/async_llm.py | 5 +- vllm/v1/engine/core.py | 7 +++ vllm/v1/engine/llm_engine.py | 1 + vllm/v1/engine/output_processor.py | 23 ++++++++ 14 files changed, 169 insertions(+), 28 deletions(-) diff --git a/docs/serving/online_serving/token_in_token_out.md b/docs/serving/online_serving/token_in_token_out.md index e22d4906155b..b3e15ae61692 100644 --- a/docs/serving/online_serving/token_in_token_out.md +++ b/docs/serving/online_serving/token_in_token_out.md @@ -55,6 +55,21 @@ The server returns a 400 for: For using `output_mode` with separate prefill and decode pools, see [Disaggregated Prefilling](../../features/disagg_prefill.md#generate-api-output-modes). +## Weight versions + +The final output of every choice carries `weight_versions`, the label of the weights that sampled each range of its output tokens. Ranges are half-open over output-token indexes (prompt excluded), contiguous and in generation order, so a request that crossed one weight update gets two: + +```json +"weight_versions": [ + {"version": "v7", "start": 0, "end": 57}, + {"version": "v8", "start": 57, "end": 128} +] +``` + +- With `stream: true`, only the last chunk of each choice carries the list, including the last chunk after an abort. Its offsets cover the whole output. Earlier chunks carry `null`. +- An output with no tokens, such as a request aborted before its first token, gets `[]`. +- The labels are the ones you set with `/update_weight_version` or `/finish_weight_update`, see [Async Reinforcement Learning](../../training/async_rl.md#weight-versions-in-outputs). vLLM doesn't order or check them, and the label is `"default"` until you set one. + ## Aborting requests `POST /inference/v1/abort_requests` aborts in-flight requests. It is registered wherever `/inference/v1/generate` is and requires the API key when `--api-key` is set, like `/inference/v1/generate`. diff --git a/docs/training/async_rl.md b/docs/training/async_rl.md index 9e75a24eaa1a..b5ae067e3cb1 100644 --- a/docs/training/async_rl.md +++ b/docs/training/async_rl.md @@ -43,10 +43,11 @@ With `VLLM_SERVER_DEV_MODE=1`, the vLLM HTTP server exposes the same functionali - `POST /pause?mode=keep` - Pause generation - `POST /resume` - Resume generation - `POST /abort_requests` - Abort in-flight requests without pausing the scheduler (send `{}` to abort all, or `{"request_ids": [...]}`) +- `POST /update_weight_version` - Set the `weight_version` label without changing weights (send `{"new_version": "v8"}`) - `GET /weight_info` - Return the latest committed `weight_version` !!! note "Data Parallelism" - When using data parallelism with vLLM's **internal load balancer** (i.e. `data_parallel_backend="ray"`), pause and resume are handled automatically across all DP ranks -- a single call is sufficient. When using an **external load balancer** (i.e. multiple independent vLLM instances behind a proxy), you must send pause and resume requests to **every** engine instance individually before and after the weight update. + When using data parallelism with vLLM's **internal load balancer** (i.e. `data_parallel_backend="ray"`), pause and resume are handled automatically across all DP ranks -- a single call is sufficient. When using an **external load balancer** (i.e. multiple independent vLLM instances behind a proxy), you must send pause and resume requests to **every** engine instance individually before and after the weight update, and set the weight version on every instance too. ## Typical Async RL Flow @@ -54,12 +55,27 @@ A typical async RL loop with weight syncing looks like this: 1. Start generating rollouts from the current policy 2. Once trainer has new weights to update to, pause generation with `mode="keep"` -3. Sync the updated weights from the trainer to the inference engine (see [Weight Transfer](weight_transfer/README.md)) +3. Sync the updated weights from the trainer to the inference engine (see [Weight Transfer](weight_transfer/README.md)) and set their weight version 4. Resume generation -- in-flight requests continue with the new weights 5. Repeat The key insight is that requests paused with `mode="keep"` will produce tokens from the **old** weights before the pause and tokens from the **new** weights after resume. The `clear_cache` parameter controls whether the KV cache is invalidated during the pause. When `clear_cache=True`, previously cached key-value entries are discarded, so all tokens generated after resume will be computed entirely with the new weights. When `clear_cache=False`, existing KV cache entries are retained, meaning some tokens in context may still reflect the old weights (stale KV cache). +## Weight Versions in Outputs + +To tell which tokens came from which weights, set a weight-version label while generation is paused, after the new weights are loaded: pass `weight_version` to `finish_weight_update`, or call `update_weight_version` (`POST /update_weight_version` over HTTP). Labels are strings you choose, such as the trainer step; vLLM doesn't order or check them. + +The final output of every sequence then carries `weight_versions`, one span per label over its output-token indexes. A request paused with `mode="keep"` across one update gets two spans: + +```python +completion = final_output.outputs[0] +for span in completion.weight_versions: + # The tokens in [span.start, span.end) were sampled by weights `span.version`. + print(span.version, completion.token_ids[span.start : span.end]) +``` + +The offsets cover the whole output, so with streamed delta outputs, collect the tokens before slicing. `/inference/v1/generate` returns the same list on each choice; see [Tokens In <> Tokens Out API](../serving/online_serving/token_in_token_out.md#weight-versions). + ## Example The [async RLHF example](../../examples/rl/rlhf_async_new_apis.py) demonstrates this pattern with `vllm.AsyncLLMEngine`, NCCL weight transfer, and mid-flight pause/resume with validation. diff --git a/examples/rl/rlhf_async_new_apis.py b/examples/rl/rlhf_async_new_apis.py index 081023156edf..3449d0f17f20 100644 --- a/examples/rl/rlhf_async_new_apis.py +++ b/examples/rl/rlhf_async_new_apis.py @@ -26,8 +26,8 @@ * Pause generation once any request reaches a token threshold. * Broadcast the training model's weights to the inference engine via the NCCL weight transfer engine, replacing the base weights. -* Resume generation and collect results, noting which tokens were - generated before vs. after the weight swap. +* Label the new weights, resume generation, and use `weight_versions` + to locate the weight swap in each output. * Validate correctness by launching a fresh vLLM instance loaded directly with the training model and comparing its output to the post-swap tokens from the weight-synced engine. @@ -61,6 +61,7 @@ MODEL_NAME_V1 = "Qwen/Qwen3-1.7B-Base" MODEL_NAME_V2 = "Qwen/Qwen3-1.7B" PAUSE_TOKEN_THRESHOLD = 10 +WEIGHT_VERSION_V2 = "v2" ATTN_BACKEND = "TRITON_ATTN" if current_platform.is_rocm() else "FLASH_ATTN" @@ -77,44 +78,27 @@ def __init__(self, **kwargs): log_requests=engine_args.enable_log_requests, log_stats=not engine_args.disable_log_stats, ) - self._generation_paused = False self._request_pause_flag = False async def do_generate( self, prompt_token_ids: list[int], sampling_params: vllm.SamplingParams - ) -> tuple[vllm.RequestOutput, int]: - """Generate a single request, setting the request pause flag once the - token count reaches the threshold. - - Returns (output, pause_token_index). pause_token_index is the number - of tokens generated before the weight change, or -1 if no pause. - """ - pause_token_index = -1 - prev_token_count = 0 + ) -> vllm.RequestOutput: + """Generate a request and set the pause flag at the token threshold.""" async for request_output in self.generate( {"prompt_token_ids": prompt_token_ids}, sampling_params, request_id=str(uuid.uuid4()), ): output = request_output - cur_token_count = len(output.outputs[0].token_ids) - if ( - cur_token_count >= PAUSE_TOKEN_THRESHOLD - and not self._request_pause_flag - ): + if len(output.outputs[0].token_ids) >= PAUSE_TOKEN_THRESHOLD: self._request_pause_flag = True - if self._generation_paused and pause_token_index == -1: - pause_token_index = prev_token_count - prev_token_count = cur_token_count - return output, pause_token_index + return output async def pause_after_n_tokens(self): """Wait for any request to set the pause flag, then pause.""" while not self._request_pause_flag: await asyncio.sleep(0) await super().pause_generation(mode="keep") - await asyncio.sleep(5) - self._generation_paused = True @ray.remote(num_gpus=1) @@ -262,9 +246,18 @@ def generate(self, token_ids: list[int], max_new_tokens: int) -> list[int]: ray.get(llm.pause_after_n_tokens.remote()) ray.get(train_model.broadcast_weights.remote()) +# Label the new weights before resuming generation. +ray.get(llm.update_weight_version.remote(WEIGHT_VERSION_V2)) ray.get(llm.resume_generation.remote()) -results = ray.get(gen_futures) +results = [] +for output in ray.get(gen_futures): + completion = output.outputs[0] + pause_idx = next( + (s.start for s in completion.weight_versions if s.version == WEIGHT_VERSION_V2), + len(completion.token_ids), + ) + results.append((output, pause_idx)) for i, (output, pause_idx) in enumerate(results): all_token_ids = list(output.outputs[0].token_ids) @@ -326,7 +319,7 @@ def generate(self, token_ids: list[int], max_new_tokens: int) -> list[int]: num_pass = 0 num_total = len(results) -for i, ((output, pause_idx), (val_output, _)) in enumerate(zip(results, val_results)): +for i, ((output, pause_idx), val_output) in enumerate(zip(results, val_results)): expected = list(output.outputs[0].token_ids)[pause_idx:] actual = list(val_output.outputs[0].token_ids) match = actual == expected diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index 9339c117a889..8945f8b11020 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -228,6 +228,9 @@ struct WireEngineCoreOutputs { /// wave needs to start in other engines. #[serde(default)] start_wave: Option, + /// Decoded for Python wire compatibility; not exposed by the Rust frontend. + #[serde(default)] + weight_version: Option, } /// Data-parallel control notifications multiplexed through `EngineCoreOutputs`. diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index bd8d40dc32f2..c0ca97dbe638 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -139,6 +139,7 @@ class EngineCoreOutputs( finished_requests: set[str] | None = None wave_complete: int | None = None start_wave: int | None = None + weight_version: str | None = None request = EngineCoreRequest( @@ -240,6 +241,7 @@ class EngineCoreOutputs( ) ], finished_requests={"req-1"}, + weight_version="v7", ) extended_outputs = EngineCoreOutputs( @@ -252,6 +254,7 @@ class EngineCoreOutputs( ) ], finished_requests=outputs.finished_requests, + weight_version=outputs.weight_version, ) extended_outputs_bytes = msgspec.msgpack.encode(extended_outputs) # The ordinary frontend's schema ignores even non-default extension values. diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index 44c3ee20b5e4..8eb85765993b 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -24,6 +24,7 @@ PoolingRequestOutput, RequestOutput, SamplingMask, + WeightVersionSpan, ) from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.tokenizers import TokenizerLike @@ -87,6 +88,7 @@ def test_completion_output_preserves_each_sampling_mask_position() -> None: ] state.routed_experts_chunks = [] state.spec_decode_metrics = None + state.weight_version_starts = [] output = state._new_completion_output([1, 2, 3], FinishReason.LENGTH, None) @@ -1596,6 +1598,7 @@ def test_sampling_masks_follow_output_kind(output_kind): state.sampling_mask_chunks = [] state.routed_experts_chunks = [] state.spec_decode_metrics = None + state.weight_version_starts = [] supports = [[10, 11], [20], [30, 31]] masks = [] @@ -1725,3 +1728,53 @@ def test_async_pooling_cache_miss_does_not_fail_output_processor(): assert output.error is not None assert output.error.retryable assert not output_processor.has_unfinished_requests() + + +@pytest.mark.parametrize( + "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY] +) +@pytest.mark.parametrize( + "steps,expected", + [ + # Merge adjacent equal labels; A -> B -> A remains three spans. + ( + [("v7", [10, 11]), ("v7", [12]), ("v8", [13, 14]), ("v7", [15])], + [("v7", 0, 3), ("v8", 3, 5), ("v7", 5, 6)], + ), + # An output without tokens has no spans. + ([], []), + ], +) +def test_weight_version_spans(output_kind: RequestOutputKind, steps, expected): + output_processor = OutputProcessor(tokenizer=None, log_stats=False) + queue = RequestOutputCollector(output_kind, request_id="request-0-int") + request = EngineCoreRequest( + request_id="request-0-int", + external_req_id="request-0", + prompt_token_ids=[1, 2, 3], + mm_features=None, + arrival_time=0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + sampling_params=SamplingParams(detokenize=False, output_kind=output_kind), + pooling_params=None, + ) + output_processor.add_request(request, prompt=None, queue=queue) + + for i, (weight_version, token_ids) in enumerate(steps): + # Copy: merging deltas extends token lists in place. + output = EngineCoreOutput( + request_id="request-0-int", new_token_ids=list(token_ids) + ) + output_processor.process_outputs([output], weight_version=weight_version) + if i == 0 and output_kind == RequestOutputKind.DELTA: + intermediate = queue.get_nowait() + assert intermediate is not None + assert intermediate.outputs[0].weight_versions is None + + # Unread deltas merge into the abort output without losing spans. + output_processor.abort_requests(["request-0"], internal=False) + final = queue.get_nowait() + assert final is not None and final.finished + assert final.outputs[0].weight_versions == [WeightVersionSpan(*s) for s in expected] diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index 84a6a2c62388..bccc8d8eff0b 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -31,6 +31,7 @@ ) from vllm.entrypoints.serve.engine.protocol import UsageInfo from vllm.logprobs import Logprob +from vllm.outputs import WeightVersionSpan from vllm.renderers import TokenizeParams from vllm.sampling_params import SamplingParams from vllm.utils import random_uuid @@ -411,6 +412,8 @@ class GenerateChoiceBase(BaseModel): # or (b) `enable_return_routed_experts` is off server-side. routed_experts: str | None = None sampling_mask: list[list[int]] | None = None + # Contiguous weight version spans over output-token indexes. + weight_versions: list[WeightVersionSpan] | None = None @field_validator("token_ids") @classmethod @@ -443,6 +446,8 @@ class GenerateStreamChoiceBase(BaseModel): token_ids: list[int] | None = None routed_experts: str | None = None sampling_mask: list[list[int]] | None = None + # Set on the final chunk only; covers the whole output. + weight_versions: list[WeightVersionSpan] | None = None class GenerateTokensStreamChoice(GenerateStreamChoiceBase): diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index 816d70bec000..085ae97e16a7 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -453,6 +453,7 @@ async def serve_tokens_full_generator( token_ids=as_list(output.token_ids), routed_experts=routed_experts_b64, sampling_mask=sampling_mask, + weight_versions=output.weight_versions, ) if text_mode: text_choices.append( @@ -625,6 +626,7 @@ async def serve_tokens_stream_generator( token_ids=as_list(delta_token_ids), routed_experts=routed_experts_b64, sampling_mask=sampling_mask, + weight_versions=output.weight_versions, ) chunk: GenerateTokensStreamResponse | GenerateTextStreamResponse if text_mode: diff --git a/vllm/outputs.py b/vllm/outputs.py index 8e69ffa91d91..145b7d510f71 100644 --- a/vllm/outputs.py +++ b/vllm/outputs.py @@ -38,6 +38,15 @@ class SamplingMask: token_ids: list[list[int]] +@dataclass +class WeightVersionSpan: + """Output tokens `[start, end)` were sampled by weights labelled `version`.""" + + version: str + start: int + end: int + + @dataclass class CompletionOutput: """The output data of one completion output of a request. @@ -62,6 +71,8 @@ class CompletionOutput: ``--per-request-spec-decode-metrics`` is enabled; None otherwise. Surfaced in the response as ``metrics.speculative_decoding`` for single-sequence (``n == 1``) requests. + weight_versions: Contiguous weight version spans over output-token indexes. + Set on the final output only. """ @@ -76,6 +87,7 @@ class CompletionOutput: lora_request: LoRARequest | None = None sampling_mask: SamplingMask | None = None spec_decode_metrics: RequestSpecDecodeMetrics | None = None + weight_versions: list[WeightVersionSpan] | None = None def finished(self) -> bool: return self.finish_reason is not None @@ -223,6 +235,8 @@ def add(self, next_output: "RequestOutput", aggregate: bool) -> None: # aggregation with earlier chunks that have no R3. if next_completion.routed_experts is not None: completion.routed_experts = next_completion.routed_experts + # Like finish_reason, spans come only with the terminal output. + completion.weight_versions = next_completion.weight_versions completion.finish_reason = next_completion.finish_reason completion.stop_reason = next_completion.stop_reason else: diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index effe8d2416e4..1b14b08246aa 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -282,6 +282,9 @@ class EngineCoreOutputs( # "old" wave, so the next wave needs to be started in other engines. start_wave: int | None = None + # Weight version used for this step. + weight_version: str | None = None + def __post_init__(self): if self.timestamp == 0.0: self.timestamp = time.monotonic() diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 08445d8f913a..71e30990c578 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -855,7 +855,10 @@ async def output_handler(): outputs_slice = engine_core_outputs[start:end] # 2) Process EngineCoreOutputs. processed_outputs = output_processor.process_outputs( - outputs_slice, outputs.timestamp, iteration_stats + outputs_slice, + outputs.timestamp, + iteration_stats, + weight_version=outputs.weight_version, ) # NOTE: RequestOutputs are pushed to their queues. assert not processed_outputs.request_outputs diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index b13764abec3a..bbca64a58ae9 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -639,6 +639,11 @@ def _attach_iteration_details( else: eco.scheduler_stats.iteration_details = iteration_details + def _attach_weight_version(self, outputs: dict[int, EngineCoreOutputs]) -> None: + # Weight versions change only while generation is paused. + for eco in outputs.values(): + eco.weight_version = self._weight_version + def _should_throttle_prefills(self) -> bool: """Whether to defer new prefills this step (DP prefill balancing). Overridden by the DP engine core; never throttles otherwise.""" @@ -672,6 +677,7 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: scheduler_output, model_output ) self._attach_iteration_details(engine_core_outputs, iteration_details) + self._attach_weight_version(engine_core_outputs) return engine_core_outputs, scheduler_output.total_num_scheduled_tokens > 0 @@ -773,6 +779,7 @@ def step_with_batch_queue( scheduler_output, model_output ) self._attach_iteration_details(engine_core_outputs, iteration_details) + self._attach_weight_version(engine_core_outputs) # NOTE(nick): We can either handle the deferred tasks here or save # in a field and do it immediately once step_with_batch_queue is diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index 2dd70fc17917..330c114a86dd 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -335,6 +335,7 @@ def step(self) -> list[RequestOutput | PoolingRequestOutput]: outputs.outputs, engine_core_timestamp=outputs.timestamp, iteration_stats=iteration_stats, + weight_version=outputs.weight_version, ) mm_processor_cache = self.renderer.mm_processor_cache diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index abb45f7239be..263a28fb3901 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -19,6 +19,7 @@ RequestError, RequestOutput, SamplingMask, + WeightVersionSpan, ) from vllm.sampling_params import RequestOutputKind from vllm.tokenizers import TokenizerLike @@ -190,6 +191,8 @@ def __init__( # Routed experts accumulation (prompt + sample chunks) self.routed_experts_chunks: list[np.ndarray] = [] self.sampling_mask_chunks: list[SamplingMaskLists] = [] + # (weight version, first output-token index) for each span. + self.weight_version_starts: list[tuple[str, int]] = [] # Stream Interval self.stream_interval = stream_interval @@ -471,8 +474,25 @@ def _new_completion_output( finish_reason=str(finish_reason) if finished else None, stop_reason=stop_reason if finished else None, spec_decode_metrics=self.spec_decode_metrics if finished else None, + weight_versions=self._weight_version_spans() if finished else None, ) + def track_weight_version(self, weight_version: str) -> None: + """Start a span when the weight version changes.""" + assert self.detokenizer is not None + starts = self.weight_version_starts + if not starts or starts[-1][0] != weight_version: + starts.append((weight_version, self.detokenizer.num_output_tokens())) + + def _weight_version_spans(self) -> list[WeightVersionSpan] | None: + assert self.detokenizer is not None + num_tokens = self.detokenizer.num_output_tokens() + starts = self.weight_version_starts + if not starts and num_tokens: + return None # No label was reported for these tokens. + ends = [start for _, start in starts[1:]] + [num_tokens] + return [WeightVersionSpan(v, s, e) for (v, s), e in zip(starts, ends)] + def _new_pooling_output(self, pooling_output: torch.Tensor) -> PoolingOutput: return PoolingOutput(data=pooling_output) @@ -659,6 +679,7 @@ def process_outputs( engine_core_outputs: list[EngineCoreOutput], engine_core_timestamp: float | None = None, iteration_stats: IterationStats | None = None, + weight_version: str | None = None, ) -> OutputProcessorOutput: """Process the EngineCoreOutputs: 1) Compute stats for logging @@ -736,6 +757,8 @@ def process_outputs( if pooling_output is None: assert req_state.detokenizer is not None assert req_state.logprobs_processor is not None + if new_token_ids and weight_version is not None: + req_state.track_weight_version(weight_version) # 2) Detokenize the token ids into text and perform stop checks. num_prev_tokens = req_state.detokenizer.num_output_tokens() stop_string = req_state.detokenizer.update( From 01abe975b5837f9e0722a196a28e57331735e0a6 Mon Sep 17 00:00:00 2001 From: bzsuni <86399306+bzsuni@users.noreply.github.com> Date: Sun, 11 Oct 2026 11:55:41 +0800 Subject: [PATCH 249/624] [Bugfix][KV Connector] Fix speculative block reservations across KV cache groups (#59880) Signed-off-by: bzsuni --- .../unit/test_remote_prefill_lifecycle.py | 133 ++++++++++++++---- tests/v1/kv_connector/unit/utils.py | 2 + vllm/v1/core/sched/scheduler.py | 34 ++--- 3 files changed, 128 insertions(+), 41 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py index 59fcc678bfb9..b72de811c0f9 100644 --- a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy +from dataclasses import replace from unittest.mock import patch import pytest @@ -761,35 +762,54 @@ def test_async_loads_both_admitted_when_pool_fits(): assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS -def test_async_load_reserves_blocks_for_promotion_margin(): - """An async load is not admitted unless the blocks its own promotion will - need are still free. - - A parked load is allocated without lookahead slots, but promotion pads it - to ``1 + num_spec_tokens`` and asks for the lookahead margin on top. Without - reserving that margin, two loads can be admitted that together consume the - whole pool; the head then fails ``allocate_slots`` forever while the load - behind it is never reached, and with nothing running no block is ever freed. +@pytest.mark.parametrize( + ("swa_enabled", "block_size", "num_spec_tokens", "num_blocks"), + [(False, 16, 3, 8), (True, 16, 3, 17), (True, 64, 31, 40)], + ids=["single-group", "multiple-groups", "different-block-sizes"], +) +def test_async_load_reserves_blocks_for_promotion_margin( + swa_enabled, block_size, num_spec_tokens, num_blocks +): + """Async loads must leave enough blocks for promotion in every cache group. - req_a (4 blocks) and req_b (3 blocks) exactly fill the 7 usable blocks, so - req_b must be held back in WAITING holding no blocks, leaving req_a room to - be promoted and run. + Otherwise two loads can fill the pool and prevent the head from being + promoted, with nothing running to free blocks and break the deadlock. """ - vllm_config = create_vllm_config(num_speculative_tokens=3) - BLOCK_SIZE = vllm_config.cache_config.block_size - scheduler = create_scheduler(vllm_config, num_blocks=8) # usable = 7 + vllm_config = create_vllm_config( + block_size=block_size, + num_speculative_tokens=num_spec_tokens, + disable_hybrid_kv_cache_manager=False, + ) + hash_block_size = 16 + kv_cache_config = make_kv_cache_config( + hash_block_size, swa_enabled=swa_enabled, sw_size=4096, num_blocks=num_blocks + ) + full_group = kv_cache_config.kv_cache_groups[0] + full_group.kv_cache_spec = replace(full_group.kv_cache_spec, block_size=block_size) + if swa_enabled: + sliding_group = kv_cache_config.kv_cache_groups[-1] + sliding_group.kv_cache_spec = replace( + sliding_group.kv_cache_spec, + page_size_padded=full_group.kv_cache_spec.page_size_bytes, + ) + scheduler = create_scheduler( + vllm_config, + num_blocks=num_blocks, + kv_cache_config=kv_cache_config, + hash_block_size=hash_block_size, + ) req_a = create_request( request_id=1, - block_size=BLOCK_SIZE, - num_tokens=BLOCK_SIZE * 4, + block_size=hash_block_size, + num_tokens=block_size * 4, do_remote_prefill=True, max_tokens=1, ) req_b = create_request( request_id=2, - block_size=BLOCK_SIZE, - num_tokens=BLOCK_SIZE * 3, + block_size=hash_block_size, + num_tokens=block_size * 3, do_remote_prefill=True, max_tokens=1, ) @@ -800,16 +820,13 @@ def test_async_load_reserves_blocks_for_promotion_margin(): with patch.object( scheduler.connector, "get_num_new_matched_tokens", - side_effect=[(BLOCK_SIZE * 4, True), (BLOCK_SIZE * 3, True)], + side_effect=[(block_size * 4, True), (block_size * 3, True)], ): scheduler.schedule() assert req_a.status == RequestStatus.WAITING_FOR_REMOTE_KVS assert req_b.status == RequestStatus.WAITING - req_to_blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[ - 0 - ].req_to_blocks - assert req_b.request_id not in req_to_blocks + assert not any(scheduler.kv_cache_manager.get_block_ids(req_b.request_id)) # req_a's load lands: it must be promotable and actually get scheduled. scheduler.update_from_output( @@ -821,6 +838,74 @@ def test_async_load_reserves_blocks_for_promotion_margin(): assert scheduler_output.num_scheduled_tokens[req_a.request_id] > 0 +@pytest.mark.parametrize( + ("mamba_cache_mode", "mamba_block_size", "num_blocks"), + [("none", 128, 18), ("align", 16, 18)], +) +def test_async_load_does_not_double_reserve_mamba_spec_blocks( + mamba_cache_mode, mamba_block_size, num_blocks +): + """Two hybrid loads fit when Mamba's speculative states are counted once.""" + block_size = 16 + num_spec_tokens = 3 + vllm_config = create_vllm_config( + block_size=mamba_block_size, + max_model_len=128, + num_speculative_tokens=num_spec_tokens, + disable_hybrid_kv_cache_manager=False, + ) + vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode + kv_cache_config = make_kv_cache_config( + block_size, + mamba_enabled=True, + mamba_cache_mode=mamba_cache_mode, + num_blocks=num_blocks, + ) + mamba_group = kv_cache_config.kv_cache_groups[-1] + mamba_group.kv_cache_spec = replace( + mamba_group.kv_cache_spec, + block_size=mamba_block_size, + num_speculative_blocks=num_spec_tokens, + ) + scheduler = create_scheduler( + vllm_config, + num_blocks=num_blocks, + kv_cache_config=kv_cache_config, + hash_block_size=block_size, + ) + reqs = [ + create_request( + request_id=i, + block_size=block_size, + num_tokens=num_tokens, + do_remote_prefill=True, + max_tokens=1, + ) + for i, num_tokens in enumerate((64, 48), 1) + ] + for req in reqs: + scheduler.add_request(req) + + with patch.object( + scheduler.connector, + "get_num_new_matched_tokens", + side_effect=[(req.num_tokens, True) for req in reqs], + ): + scheduler.schedule() + + assert all(req.status == RequestStatus.WAITING_FOR_REMOTE_KVS for req in reqs) + scheduler.update_from_output( + scheduler.schedule(), + create_model_runner_output( + [], finished_recving={req.request_id for req in reqs} + ), + ) + output = scheduler.schedule() + for req in reqs: + assert req.status == RequestStatus.RUNNING + assert output.num_scheduled_tokens[req.request_id] > 0 + + def test_deferred_lookup_does_not_block_parked_load(): """A request with a deferred connector lookup must not strand an async load queued behind it. diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index d738e183952f..28bfda419a76 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -177,6 +177,7 @@ def create_scheduler( vllm_config: VllmConfig, num_blocks: int = 10000, kv_cache_config: KVCacheConfig | None = None, + hash_block_size: int | None = None, ) -> Scheduler | AsyncScheduler: """Initialize Scheduler For Testing.""" block_size = vllm_config.cache_config.block_size @@ -207,6 +208,7 @@ def create_scheduler( log_stats=True, structured_output_manager=StructuredOutputManager(vllm_config), block_size=block_size, + hash_block_size=hash_block_size, ) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 0758c6122565..2603fa10a75e 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -344,6 +344,17 @@ def __init__( self.cache_config.enable_mamba_shared_prefix_checkpoint ), ) + # Mamba managers already account for their speculative state blocks. + num_spec_decode_tokens = ( + 1 + self.num_spec_tokens + self.num_lookahead_tokens + if self.num_spec_tokens + else 0 + ) + self._spec_decode_step_blocks = sum( + cdiv(num_spec_decode_tokens, manager.block_size) + for manager in self.kv_cache_manager.coordinator.single_type_managers + if not isinstance(manager.kv_cache_spec, MambaSpec) + ) # Bind after construction so connectors can access the cache manager. if self.connector is not None: self.connector.bind_kv_cache_manager(self.kv_cache_manager) @@ -1236,7 +1247,7 @@ def skip_request(from_queue: RequestQueue) -> None: # predictable preemptions. reserved_blocks = ( self._inflight_prefill_reserved_blocks() - + self._spec_decode_step_blocks() + + self._spec_decode_step_blocks ) # Replayed tokens are already counted in the adopted hit; a @@ -3059,9 +3070,12 @@ def _connector_finished( def _request_remaining_blocks(self, request: Request) -> int: """Blocks `request` still needs to hold its full sequence and be promoted.""" full_num_tokens = min(request.num_tokens, self.max_model_len) - num_blocks = self.kv_cache_manager.coordinator.get_num_blocks_to_allocate( + num_tokens = full_num_tokens + if self.num_spec_tokens: + num_tokens += 1 + self.num_spec_tokens + self.num_lookahead_tokens + return self.kv_cache_manager.coordinator.get_num_blocks_to_allocate( request_id=request.request_id, - num_tokens=full_num_tokens, + num_tokens=num_tokens, new_computed_blocks=self.kv_cache_manager.empty_kv_cache_blocks.blocks, num_encoder_tokens=0, total_computed_tokens=request.num_computed_tokens, @@ -3070,20 +3084,6 @@ def _request_remaining_blocks(self, request: Request) -> int: apply_admission_cap=True, prefill_end=max(request.num_prompt_tokens, request.num_tokens - 1), ) - return num_blocks + self._spec_decode_step_blocks() - - def _spec_decode_step_blocks(self) -> int: - """Number of blocks for the extra KV slots a spec decode step needs. - - When using async kv load, scheduler must reserve enough blocks for - full sequence + the spec decode step, otherwise request cannot be - able to run after the async_load if we are out of kv blocks. - """ - if not self.num_spec_tokens: - return 0 - return cdiv( - 1 + self.num_spec_tokens + self.num_lookahead_tokens, self.block_size - ) def _set_kv_fetch_stage(self, request: Request, stage: str | None) -> None: kv_fetch_stages = self._kv_fetch_stages From f8471164bf9f5190752fd7b903d415a34fc4ba1c Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Sat, 10 Oct 2026 21:15:05 -0700 Subject: [PATCH 250/624] [MRV2] Detatch weight-reload logic from MRV1 (#60937) --- tests/v1/worker/test_gpu_model_runner.py | 61 ++++++--- .../worker/test_gpu_worker_weight_transfer.py | 4 +- .../model_loader/reload/inplace.py | 123 ++++++++++++++++++ vllm/v1/worker/gpu/model_runner.py | 33 ++--- vllm/v1/worker/gpu_model_runner.py | 121 ++--------------- vllm/v1/worker/gpu_worker.py | 23 +++- 6 files changed, 215 insertions(+), 150 deletions(-) create mode 100644 vllm/model_executor/model_loader/reload/inplace.py diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 158befc15027..0578768d68c7 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -10,6 +10,7 @@ import pytest import torch +import vllm.model_executor.model_loader.reload.inplace as reload_inplace_module import vllm.v1.worker.gpu_model_runner as gpu_model_runner_module from vllm.config import ( AttentionConfig, @@ -63,6 +64,7 @@ from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras from vllm.v1.worker.gpu_input_batch import InputBatch from vllm.v1.worker.gpu_model_runner import GPUModelRunner +from vllm.v1.worker.gpu_worker import Worker from vllm.v1.worker.utils import select_common_block_size BLOCK_SIZE = 16 @@ -1045,13 +1047,22 @@ def test_update_states_pp_async_multi_request_keeps_rank_state_consistent( ) +def _update_config(model_runner: GPUModelRunner, overrides: dict) -> None: + # Apply overrides via the worker without constructing a full worker. + worker = object.__new__(Worker) + worker.vllm_config = model_runner.vllm_config + worker.model_runner = model_runner + worker.update_config(overrides) + + def test_update_config(model_runner): # Simple update - model_runner.update_config({"load_config": {"load_format": "dummy"}}) + _update_config(model_runner, {"load_config": {"load_format": "dummy"}}) assert model_runner.load_config.load_format == "dummy" + assert model_runner.vllm_config.load_config is model_runner.load_config # Raise error on non-existing config with pytest.raises(ValueError, match="do_not_exist_config"): - model_runner.update_config({"do_not_exist_config": "dummy"}) + _update_config(model_runner, {"do_not_exist_config": "dummy"}) def test_load_model_weights_inplace(dist_init, model_runner, model_runner_2): @@ -1059,12 +1070,14 @@ def test_load_model_weights_inplace(dist_init, model_runner, model_runner_2): # model_runner_2 loads dummy weights first then load real weights inplace model_runner.load_model() original_load_format = model_runner_2.load_config.load_format - model_runner_2.update_config({"load_config": {"load_format": "dummy"}}) + _update_config(model_runner_2, {"load_config": {"load_format": "dummy"}}) model_runner_2.load_model() # Initial model loading with dummy weights assert str(model_runner.get_model().state_dict()) != str( model_runner_2.get_model().state_dict() ) - model_runner_2.update_config({"load_config": {"load_format": original_load_format}}) + _update_config( + model_runner_2, {"load_config": {"load_format": original_load_format}} + ) model_runner_2.reload_weights() # Load real weights inplace assert str(model_runner.get_model().state_dict()) == str( model_runner_2.get_model().state_dict() @@ -1083,26 +1096,34 @@ def test_reload_weights_path_replaces_object_storage_source(monkeypatch): # silently re-streams the original checkpoint. loader = Mock() loader.get_all_weights.return_value = iter(()) - monkeypatch.setattr(gpu_model_runner_module, "get_model_loader", lambda _: loader) - monkeypatch.setattr(gpu_model_runner_module, "initialize_layerwise_reload", Mock()) - monkeypatch.setattr(gpu_model_runner_module, "finalize_layerwise_reload", Mock()) + monkeypatch.setattr(reload_inplace_module, "get_model_loader", lambda _: loader) + monkeypatch.setattr(reload_inplace_module, "initialize_layerwise_reload", Mock()) + monkeypatch.setattr(reload_inplace_module, "finalize_layerwise_reload", Mock()) - runner = Mock(lora_config=None) - runner.model_config = SimpleNamespace( - model="/tmp/pulled-config-files", - model_weights="s3://bucket/original", - revision="abc123", - quantization=None, + vllm_config = SimpleNamespace( + lora_config=None, + load_config=SimpleNamespace(load_format="runai_streamer"), + model_config=SimpleNamespace( + model="/tmp/pulled-config-files", + model_weights="s3://bucket/original", + revision="abc123", + quantization=None, + ), ) - runner.get_model.return_value.named_parameters.return_value = [] - runner.get_model.return_value.load_weights.return_value = None - - GPUModelRunner.reload_weights(runner, weights_path="org/new-model") + model = Mock() + model.named_parameters.return_value = [] + model.load_weights.return_value = None - loader.get_all_weights.assert_called_once_with( - runner.model_config, runner.get_model.return_value + reload_inplace_module.reload_weights( + vllm_config, + model, + weights_iterator=None, + weights_path="org/new-model", + is_checkpoint_format=True, ) - cfg = runner.model_config + + loader.get_all_weights.assert_called_once_with(vllm_config.model_config, model) + cfg = vllm_config.model_config assert (cfg.model, cfg.model_weights, cfg.revision) == ("org/new-model", "", None) diff --git a/tests/v1/worker/test_gpu_worker_weight_transfer.py b/tests/v1/worker/test_gpu_worker_weight_transfer.py index 49ef70789478..f1965cb1fae2 100644 --- a/tests/v1/worker/test_gpu_worker_weight_transfer.py +++ b/tests/v1/worker/test_gpu_worker_weight_transfer.py @@ -13,7 +13,7 @@ from vllm.config import ParallelConfig, VllmConfig, get_current_vllm_config from vllm.lora.layers import BaseLayerWithLoRA -from vllm.v1.worker.gpu_model_runner import _get_parameter_for_reload +from vllm.model_executor.model_loader.reload.inplace import get_parameter_for_reload from vllm.v1.worker.gpu_worker import Worker @@ -95,7 +95,7 @@ def test_reload_parameter_lookup_preserves_lora_module_names(): assert set(named_parameters) == {"proj.base_layer.qweight"} assert named_parameters["proj.base_layer.qweight"] is qweight assert model.get_parameter("proj.base_layer.qweight") is qweight - assert _get_parameter_for_reload(model, "proj.qweight") is qweight + assert get_parameter_for_reload(model, "proj.qweight") is qweight def test_start_update_finish_delegates_to_engine(): diff --git a/vllm/model_executor/model_loader/reload/inplace.py b/vllm/model_executor/model_loader/reload/inplace.py new file mode 100644 index 000000000000..dfe11ed4e745 --- /dev/null +++ b/vllm/model_executor/model_loader/reload/inplace.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import time +from collections.abc import Iterable +from typing import cast + +import torch +import torch.nn as nn + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.model_loader import get_model_loader + +from .layerwise import finalize_layerwise_reload, initialize_layerwise_reload + +logger = init_logger(__name__) + +__all__ = ["get_parameter_for_reload", "reload_weights"] + + +def get_parameter_for_reload(model: nn.Module, name: str) -> nn.Parameter: + """Resolve checkpoint names without changing the model's module tree.""" + from vllm.lora.layers import BaseLayerWithLoRA + + module_name, _, parameter_name = name.rpartition(".") + module = model.get_submodule(module_name) + if isinstance(module, BaseLayerWithLoRA): + module = module.base_layer + return module.get_parameter(parameter_name) + + +def reload_weights( + vllm_config: VllmConfig, + model: nn.Module, + weights_iterator: Iterable[tuple[str, torch.Tensor]] | None, + weights_path: str | None, + is_checkpoint_format: bool, +) -> None: + """Reload a model's weights in place from a weights iterator or from disk. + + Callers are responsible for invalidating any state derived from the old + weights (e.g. LoRA, encoder and multimodal caches). + + Args: + vllm_config: config of the model being reloaded. If `weights_path` is + given, `vllm_config.model_config` is updated to point at it. + model: model to load weights into + weights_iterator: weights to load into model + weights_path: path to load weights from if weights_iterator is not + provided. Use path of original model if neither is provided. + is_checkpoint_format: set to False if weights have already been + processed into kernel format (repacking, renaming, etc.) + + """ + # TODO(@kylesayrs): generalize to all runners and loaders + # argument validation + if weights_iterator is None and not is_checkpoint_format: + logger.warning( + "Reloading from disk means that weights will be in checkpoint format. " + "Please use `is_checkpoint_format=True` " + "to avoid weight reloading errors" + ) + + model_config = vllm_config.model_config + load_config = vllm_config.load_config + weights_to_load = { + name.replace(".base_layer.", ".") if vllm_config.lora_config else name + for name, _ in model.named_parameters() + } + counter_before_reloading = time.perf_counter() + + # load weights from disk if none are provided + if weights_iterator is None: + model_loader = get_model_loader(load_config) + if not hasattr(model_loader, "get_all_weights"): + raise NotImplementedError( + f"Model reloading with `{load_config.load_format}` format" + ) + + if weights_path is not None: + # The revision and any object-storage `model_weights` source + # belong to the model we are reloading away from, so they must + # not be carried over to the new path. + model_config.model = weights_path + model_config.model_weights = "" + model_config.revision = None + weights_iterator = model_loader.get_all_weights(model_config, model) + weights_iterator = cast(Iterable[tuple[str, torch.Tensor]], weights_iterator) + + # begin loading weights + logger.info_once("Reloading weights inplace...") + if is_checkpoint_format: + # load weights from checkpoint/ original model format + initialize_layerwise_reload(model) + loaded_weights = model.load_weights(weights_iterator) + finalize_layerwise_reload(model, model_config) + + else: + # load weights from kernel format + logger.warning_once( + "Reloading with `is_checkpoint_format=True` requires that " + "weights be in kernel format and already sharded", + ) + loaded_weights = set() + for name, loaded_weight in weights_iterator: + param = get_parameter_for_reload(model, name) # TODO: buffers? + param.copy_(loaded_weight) + loaded_weights.add(name) + + # logging and validation + counter_after_reloading = time.perf_counter() + diff_seconds = counter_after_reloading - counter_before_reloading + logger.info_once( + "Reloading and processing weights took %.2f seconds", + diff_seconds, + ) + if model_config.quantization is None and loaded_weights is not None: + weights_not_loaded = weights_to_load - loaded_weights + if weights_not_loaded: + logger.warning( + "Following weights were not loaded from checkpoint: %s", + weights_not_loaded, + ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index e04a381b94e6..20fa4c17bb07 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -19,6 +19,7 @@ import functools import gc import time +from collections.abc import Iterable from contextlib import AbstractContextManager from copy import deepcopy from dataclasses import replace @@ -46,6 +47,7 @@ initialize_mamba_ssu_backend, ) from vllm.model_executor.model_loader import get_model_loader +from vllm.model_executor.model_loader.reload.inplace import reload_weights from vllm.model_executor.models.interfaces import requires_raw_input_tokens from vllm.model_executor.offloader import ( create_offloader, @@ -563,22 +565,21 @@ def get_draft_model(self) -> nn.Module | None: return None return speculator.model - def reload_weights(self, *args, **kwargs) -> None: - # TODO(Wentao): Use full version instead of import when fully migrated to v2 - from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1 - - GPUModelRunnerV1.reload_weights(self, *args, **kwargs) # type: ignore[arg-type] - - def update_config(self, *args, **kwargs) -> None: - # TODO(Wentao): Use full version instead of import when fully migrated to v2 - from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1 - - GPUModelRunnerV1.update_config(self, *args, **kwargs) # type: ignore[arg-type] - - # v2 reads config via self.vllm_config (e.g. in load_model), so keep it - # in sync with the attributes the v1 helper just replaced. - self.vllm_config.model_config = self.model_config - self.vllm_config.load_config = self.load_config + def reload_weights( + self, + weights_iterator: Iterable[tuple[str, torch.Tensor]] | None = None, + weights_path: str | None = None, + is_checkpoint_format: bool = True, + ) -> None: + """Reload weights in place.""" + vllm_config = self.vllm_config + model = self.get_model() + reload_weights( + vllm_config, model, weights_iterator, weights_path, is_checkpoint_format + ) + self.reset_lora_state() + self.reset_encoder_cache() + self.reset_mm_cache() @functools.cached_property def main_stream(self) -> torch.cuda.Stream: diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 90bea22c9d4f..4db3f00b91d4 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -34,7 +34,6 @@ VllmConfig, get_layers_from_vllm_config, set_current_vllm_config, - update_config, ) from vllm.config.cache import CacheConfig from vllm.config.ec_manager_config import EncoderCacheManagerMetadata @@ -57,7 +56,7 @@ set_forward_context, ) from vllm.logger import init_logger -from vllm.lora.layers import BaseLayerWithLoRA, LoRAMapping, LoRAMappingType +from vllm.lora.layers import LoRAMapping, LoRAMappingType from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager @@ -67,10 +66,7 @@ ) from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding from vllm.model_executor.model_loader import get_model_loader -from vllm.model_executor.model_loader.reload import ( - finalize_layerwise_reload, - initialize_layerwise_reload, -) +from vllm.model_executor.model_loader.reload.inplace import reload_weights from vllm.model_executor.models.interfaces import ( MixtureOfExperts, MultiModalEmbeddings, @@ -247,15 +243,6 @@ logger = init_logger(__name__) -def _get_parameter_for_reload(model: nn.Module, name: str) -> nn.Parameter: - """Resolve checkpoint names without changing the model's module tree.""" - module_name, _, parameter_name = name.rpartition(".") - module = model.get_submodule(module_name) - if isinstance(module, BaseLayerWithLoRA): - module = module.base_layer - return module.get_parameter(parameter_name) - - AttnMetadataDict: TypeAlias = dict[str, AttentionMetadata] # list when ubatching is enabled PerLayerAttnMetadata: TypeAlias = list[AttnMetadataDict] | AttnMetadataDict @@ -5203,19 +5190,6 @@ def propose_draft_token_ids( return draft_token_ids - def update_config(self, overrides: dict[str, Any]) -> None: - allowed_config_names = {"load_config", "model_config"} - for config_name, config_overrides in overrides.items(): - if config_name not in allowed_config_names: - allowed = ", ".join(sorted(allowed_config_names)) - raise ValueError( - f"Config override '{config_name}' is not supported. " - f"Supported configs: {allowed}" - ) - config = getattr(self, config_name) - new_config = update_config(config, config_overrides) - setattr(self, config_name, new_config) - @instrument(span_name="Loading (GPU)") def load_model( self, @@ -5472,89 +5446,16 @@ def reload_weights( weights_path: str | None = None, is_checkpoint_format: bool = True, ) -> None: - """Reload weights from a weights iterator or from disk - - Args: - weights_iterator: weights to load into model - weights_path: path to load weights from if weights_iterator is not - provided. Use path of original model if neither is provided. - is_checkpoint_format: set to False if weights have already been - processed into kernel format (repacking, renaming, etc.) - - """ - # TODO(@kylesayrs): generalize to all runners and loaders - # argument validation - if weights_iterator is None and not is_checkpoint_format: - logger.warning( - "Reloading from disk means that weights will be in checkpoint format. " - "Please use `is_checkpoint_format=True` " - "to avoid weight reloading errors" - ) - - model = self.get_model() - weights_to_load = { - name.replace(".base_layer.", ".") if self.lora_config else name - for name, _ in model.named_parameters() - } - counter_before_reloading = time.perf_counter() - - # load weights from disk if none are provided - if weights_iterator is None: - model_loader = get_model_loader(self.load_config) - if not hasattr(model_loader, "get_all_weights"): - raise NotImplementedError( - f"Model reloading with `{self.load_config.load_format}` format" - ) - - if weights_path is not None: - # The revision and any object-storage `model_weights` source - # belong to the model we are reloading away from, so they must - # not be carried over to the new path. - self.model_config.model = weights_path - self.model_config.model_weights = "" - self.model_config.revision = None - weights_iterator = model_loader.get_all_weights(self.model_config, model) - weights_iterator = cast( - Iterable[tuple[str, torch.Tensor]], weights_iterator - ) - - # begin loading weights - logger.info_once("Reloading weights inplace...") - if is_checkpoint_format: - # load weights from checkpoint/ original model format - initialize_layerwise_reload(model) - loaded_weights = model.load_weights(weights_iterator) - finalize_layerwise_reload(model, self.model_config) - - else: - # load weights from kernel format - logger.warning_once( - "Reloading with `is_checkpoint_format=True` requires that " - "weights be in kernel format and already sharded", - ) - loaded_weights = set() - for name, loaded_weight in weights_iterator: - param = _get_parameter_for_reload(model, name) # TODO: buffers? - param.copy_(loaded_weight) - loaded_weights.add(name) - - self.reset_lora_state() - - # logging and validation - counter_after_reloading = time.perf_counter() - diff_seconds = counter_after_reloading - counter_before_reloading - logger.info_once( - "Reloading and processing weights took %.2f seconds", - diff_seconds, + """Reload weights in place; see `reload_weights` in + `vllm.model_executor.model_loader.reload.inplace` for details.""" + reload_weights( + self.vllm_config, + self.get_model(), + weights_iterator=weights_iterator, + weights_path=weights_path, + is_checkpoint_format=is_checkpoint_format, ) - if self.model_config.quantization is None and loaded_weights is not None: - weights_not_loaded = weights_to_load - loaded_weights - if weights_not_loaded: - logger.warning( - "Following weights were not loaded from checkpoint: %s", - weights_not_loaded, - ) - + self.reset_lora_state() self.reset_encoder_cache() self.reset_mm_cache() diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index a01e5dd5f7a3..37fa00bb7efa 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -16,7 +16,12 @@ import torch.nn as nn import vllm.envs as envs -from vllm.config import CUDAGraphMode, VllmConfig, set_current_vllm_config +from vllm.config import ( + CUDAGraphMode, + VllmConfig, + set_current_vllm_config, + update_config, +) from vllm.config.compilation import CompilationMode from vllm.config.profiler import validate_profile_prefix from vllm.device_allocator import get_mem_allocator_instance @@ -591,7 +596,21 @@ def load_model(self, *, load_dummy_weights: bool = False) -> None: set_torch_threads_for_runtime() def update_config(self, overrides: dict[str, Any]) -> None: - self.model_runner.update_config(overrides) + """Apply config overrides to the worker, its model runner and the + shared VllmConfig.""" + allowed_config_names = {"load_config", "model_config"} + for config_name, config_overrides in overrides.items(): + if config_name not in allowed_config_names: + allowed = ", ".join(sorted(allowed_config_names)) + raise ValueError( + f"Config override '{config_name}' is not supported. " + f"Supported configs: {allowed}" + ) + config = getattr(self.vllm_config, config_name) + new_config = update_config(config, config_overrides) + setattr(self.vllm_config, config_name, new_config) + setattr(self.model_runner, config_name, new_config) + setattr(self, config_name, new_config) def reload_weights(self, *args, **kwargs) -> None: with set_current_vllm_config(self.vllm_config): From 2eb491e18b1ce5b8a25e4c88f32dbb5d028d525c Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Sat, 10 Oct 2026 22:40:45 -0700 Subject: [PATCH 251/624] [Perf][MiniMax-M3] MSA prefill attends with the E4M3 query; bump fmha_sm100 to 968061a (#60684) Signed-off-by: Yongye Zhu Co-authored-by: Claude Opus 5.5 --- cmake/external_projects/fmha_sm100.cmake | 11 +- ...st_minimax_m3_msa_cutlass_sparse_decode.py | 179 +++++++++++++++--- vllm/models/minimax_m3/nvidia/model.py | 2 +- .../minimax_m3/nvidia/sparse_attention_msa.py | 16 +- 4 files changed, 176 insertions(+), 32 deletions(-) diff --git a/cmake/external_projects/fmha_sm100.cmake b/cmake/external_projects/fmha_sm100.cmake index f6bb0b8acee7..6b84b35c9049 100644 --- a/cmake/external_projects/fmha_sm100.cmake +++ b/cmake/external_projects/fmha_sm100.cmake @@ -17,7 +17,7 @@ else() FetchContent_Declare( fmha_sm100 GIT_REPOSITORY https://github.com/vllm-project/MSA.git - GIT_TAG be4a52aea420b6bef73be9af1a649859964790cb + GIT_TAG 968061ab16cb14991a5c7d4c0e8eabf531be0dcf GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" @@ -36,11 +36,14 @@ set(FMHA_SM100_PY_ROOT "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100") install(FILES "${FMHA_SM100_PY_ROOT}/__init__.py" + "${FMHA_SM100_PY_ROOT}/_jit_cache.py" "${FMHA_SM100_PY_ROOT}/api.py" "${FMHA_SM100_PY_ROOT}/bench_utils.py" "${FMHA_SM100_PY_ROOT}/jit.py" + "${FMHA_SM100_PY_ROOT}/msa_warmup.py" "${FMHA_SM100_PY_ROOT}/nvfp4_kv.py" "${FMHA_SM100_PY_ROOT}/q8kv4_decode_adapter.py" + "${FMHA_SM100_PY_ROOT}/q8kv4_prefill_adapter.py" "${FMHA_SM100_PY_ROOT}/sparse.py" "${FMHA_SM100_PY_ROOT}/sparse_fmha_adapter.py" DESTINATION vllm/third_party/fmha_sm100 @@ -65,6 +68,12 @@ install(DIRECTORY "${FMHA_SM100_PY_ROOT}/kvouter/" PATTERN "__pycache__" EXCLUDE PATTERN "*.pyc" EXCLUDE) +install(DIRECTORY "${FMHA_SM100_PY_ROOT}/prefill_q8kv4/" + DESTINATION vllm/third_party/fmha_sm100/prefill_q8kv4 + COMPONENT fmha_sm100 + PATTERN "__pycache__" EXCLUDE + PATTERN "*.pyc" EXCLUDE) + install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cute/" DESTINATION vllm/third_party/fmha_sm100/cute COMPONENT fmha_sm100 diff --git a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py b/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py index 0007f20c5bfc..42e2b86965e0 100644 --- a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py +++ b/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py @@ -782,6 +782,46 @@ def _make_nvfp4_kv_cache( return kv_cache, torch.cat(dequantized, dim=-1).to(torch.bfloat16) +def _e4m3_staged_kv_reference( + kv_reference: torch.Tensor, k_scale: torch.Tensor, v_scale: torch.Tensor +) -> torch.Tensor: + """The dequantized NVFP4 cache as an E4M3-query kernel reads it: each + code * block_scale rounded to E4M3, then the global scale.""" + sides = [] + for side, global_scale in zip( + kv_reference.float().split(HEAD_DIM, dim=-1), (k_scale, v_scale) + ): + sides.append( + (side / global_scale).to(torch.float8_e4m3fn).float() * global_scale + ) + return torch.cat(sides, dim=-1).to(torch.bfloat16) + + +def _assert_close_fp8_prefill(actual: torch.Tensor, expected: torch.Tensor) -> None: + """The fp8-q prefill rounds each split's probabilities to E4M3, and the splits + follow the SM count, so a rare element moves past the BF16 tolerance.""" + close = torch.isclose(actual, expected, atol=0.05, rtol=0.05) + assert close.float().mean().item() > 1 - 1e-4 + torch.testing.assert_close(actual, expected, atol=0.25, rtol=0.05) + + +def _record_query_dtypes( + monkeypatch: pytest.MonkeyPatch, kernel_name: str +) -> list[torch.dtype]: + """Record the query dtype of every ``fmha_sm100.sparse.`` call.""" + from vllm.third_party.fmha_sm100 import sparse as msa_sparse + + kernel = getattr(msa_sparse, kernel_name) + query_dtypes: list[torch.dtype] = [] + + def record(q, *args, **kwargs): + query_dtypes.append(q.dtype) + return kernel(q, *args, **kwargs) + + monkeypatch.setattr(msa_sparse, kernel_name, record) + return query_dtypes + + def test_nvfp4_kv_cache_views_address_head_slots() -> None: """Slot 2h (K) and 2h + 1 (V) hold head h's data then scales, so each head is one contiguous run of the page: a TP rank's page is a byte slice.""" @@ -898,27 +938,14 @@ def run() -> None: torch.testing.assert_close(actual, expected, atol=0.05, rtol=0.05) -@pytest.mark.parametrize( - ("num_q_heads", "num_kv_heads"), - [pytest.param(64, 4, id="tp1"), pytest.param(16, 1, id="tp4")], -) -def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( - num_q_heads: int, +def _make_prefill_metadata( + q_lens: list[int], + seq_lens_list: list[int], + block_table: torch.Tensor, num_kv_heads: int, -) -> None: - """MSA prefill reads NVFP4 pages as their exact dequantized values, - including chunked prefill over cached context, and takes q from its fp8 - copy (the only one the fused insert writes on this path).""" - torch.manual_seed(0) - q_lens = [300, 128, 77] - seq_lens_list = [300, 640, 1100] - block_table, num_pages = _make_paged_layout(seq_lens_list) - k_scale = torch.tensor(0.5, dtype=torch.float32, device="cuda") - v_scale = torch.tensor(2.0, dtype=torch.float32, device="cuda") - kv_cache, kv_reference = _make_nvfp4_kv_cache( - num_pages, num_kv_heads, "nvfp4", k_scale, v_scale - ) - +) -> tuple[MiniMaxM3SparseMetadata, torch.Tensor]: + """Prefill-only metadata (chunks over cached context) and a top-k selecting + every causally visible page (<= TOPK per query), ending with the local one.""" total_q = sum(q_lens) seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device="cuda") context_lens = seq_lens - torch.tensor(q_lens, dtype=torch.int32, device="cuda") @@ -948,7 +975,6 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( prefill=prefill, ) - # Every causally visible page (<= TOPK per query), ending with the local one. topk = torch.full( (total_q, num_kv_heads, TOPK), -1, dtype=torch.int32, device="cuda" ) @@ -959,6 +985,35 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( assert visible <= TOPK topk[token, :, :visible] = torch.arange(visible, device="cuda") token += 1 + return metadata, topk + + +@pytest.mark.parametrize( + ("num_q_heads", "num_kv_heads"), + [pytest.param(64, 4, id="tp1"), pytest.param(16, 1, id="tp4")], +) +def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( + monkeypatch: pytest.MonkeyPatch, + num_q_heads: int, + num_kv_heads: int, +) -> None: + """MSA prefill reads NVFP4 pages as their dequantized values, staged to E4M3 + for the E4M3 q, including chunked prefill over cached context; it attends + with the fp8 q (the only one the fused insert writes on this path).""" + torch.manual_seed(0) + q_lens = [300, 128, 77] + seq_lens_list = [300, 640, 1100] + block_table, num_pages = _make_paged_layout(seq_lens_list) + k_scale = torch.tensor(0.5, dtype=torch.float32, device="cuda") + v_scale = torch.tensor(2.0, dtype=torch.float32, device="cuda") + kv_cache, kv_reference = _make_nvfp4_kv_cache( + num_pages, num_kv_heads, "nvfp4", k_scale, v_scale + ) + + metadata, topk = _make_prefill_metadata( + q_lens, seq_lens_list, block_table, num_kv_heads + ) + total_q = sum(q_lens) layer_name = "model.layers.0.self_attn.attn" q_scale = 0.5 @@ -976,6 +1031,7 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( / q_scale ).to(torch.float8_e4m3fn) query = query_fp8.to(torch.bfloat16) * q_scale + kv_reference = _e4m3_staged_kv_reference(kv_reference, k_scale, v_scale) common = dict( num_heads=num_q_heads, head_size=HEAD_DIM, @@ -984,6 +1040,7 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( topk_blocks=TOPK, sparse_block_size=BLOCK_SIZE, ) + query_dtypes = _record_query_dtypes(monkeypatch, "sparse_atten_nvfp4_kv_func") forward_context = ForwardContext( no_compile_layers={}, attn_metadata={layer_name: metadata}, @@ -1001,7 +1058,8 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( query_fp8=query_fp8, ) current_platform.synchronize() - torch.testing.assert_close(actual, expected, atol=0.05, rtol=0.05) + assert query_dtypes == [torch.float8_e4m3fn] + _assert_close_fp8_prefill(actual, expected) def test_msa_triton_decode_fallback_reads_dequantized_query_fp8() -> None: @@ -1079,3 +1137,80 @@ def test_msa_triton_decode_fallback_reads_dequantized_query_fp8() -> None: ) current_platform.synchronize() torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02) + + +@pytest.mark.parametrize( + ("num_q_heads", "num_kv_heads"), + [pytest.param(64, 4, id="tp1"), pytest.param(16, 1, id="tp4")], +) +def test_msa_prefill_reads_query_fp8_over_fp8_cache( + monkeypatch: pytest.MonkeyPatch, + num_q_heads: int, + num_kv_heads: int, +) -> None: + """With CUTLASS decode on, FP8 prefill (incl. chunks over cached context) + attends with the fp8 q the fused insert writes, not a BF16 copy, and + matches Triton on the dequantized q.""" + torch.manual_seed(0) + q_lens = [300, 128, 77] + seq_lens_list = [300, 640, 1100] + block_table, num_pages = _make_paged_layout(seq_lens_list) + kv_cache = ( + torch.randn( + num_pages, + num_kv_heads, + BLOCK_SIZE, + 2 * HEAD_DIM, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.5 + ).to(torch.float8_e4m3fn) + metadata, topk = _make_prefill_metadata( + q_lens, seq_lens_list, block_table, num_kv_heads + ) + layer_name = "model.layers.0.self_attn.attn" + q_scale = 0.5 + layer = SimpleNamespace( + layer_name=layer_name, + topk_indices_buffer=topk, + _q_scale_float=q_scale, + ) + query_fp8 = ( + torch.randn( + sum(q_lens), num_q_heads * HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + / q_scale + ).to(torch.float8_e4m3fn) + query = query_fp8.to(torch.bfloat16) * q_scale + common = dict( + num_heads=num_q_heads, + head_size=HEAD_DIM, + scale=SM_SCALE, + num_kv_heads=num_kv_heads, + kv_cache_dtype="fp8", + topk_blocks=TOPK, + sparse_block_size=BLOCK_SIZE, + ) + impl = MiniMaxM3SparseMSAImpl(**common, msa_decode_backend="cutlass") + assert impl.use_cutlass_decode + query_dtypes = _record_query_dtypes(monkeypatch, "sparse_atten_func") + forward_context = ForwardContext( + no_compile_layers={}, + attn_metadata={layer_name: metadata}, + slot_mapping={}, + ) + with override_forward_context(forward_context): + expected = MiniMaxM3SparseTritonImpl(**common).forward( + layer, query, kv_cache, torch.empty_like(query) + ) + actual = impl.forward( + layer, + torch.full_like(query, float("nan")), + kv_cache, + torch.zeros_like(query), + query_fp8=query_fp8, + ) + current_platform.synchronize() + assert query_dtypes == [torch.float8_e4m3fn] + _assert_close_fp8_prefill(actual, expected) diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py index 06c28a7a38b4..d67e9ea19d90 100644 --- a/vllm/models/minimax_m3/nvidia/model.py +++ b/vllm/models/minimax_m3/nvidia/model.py @@ -614,7 +614,7 @@ def forward( index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix] q = qkv.new_empty((num_tokens, self.q_size)) # With query_fp8, q is emitted only in fp8; the attend dequantizes it - # into ``q`` for the tokens that run on bf16-query kernels. + # into ``q`` for the tokens that run on Triton decode. query_fp8 = self._allocate_query_fp8(qkv) # index_q matches the index-K cache dtype (e4m3 for the fp8 score path); # the fused kernel emits fp8 directly when this buffer is e4m3. diff --git a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py index 5aa23b4ac849..ed8ca2a78474 100644 --- a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py +++ b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py @@ -266,7 +266,7 @@ def forward( q = query[:num_tokens].view(-1, self.num_heads, hd) out = output[:num_tokens].view(-1, self.num_heads, hd) # Given query_fp8, it is the only q the fused insert wrote: CUTLASS - # decode reads it directly, and the bf16-query kernels read it + # decode and prefill read it directly, and Triton decode reads it # dequantized into ``query``. q_fp8 = ( None @@ -339,11 +339,11 @@ def forward( # [H, prefill, MK] transposed view; build_k2q_csr consumes the # strided view directly (topK stays innermost-contiguous). prefill_topk = topk[nd:num_tokens].transpose(0, 1) - qp = ( - q[nd:] - if q_fp8 is None - else _dequantize_query(q_fp8[nd:], q_scale_float, q[nd:]) - ) + # An E4M3 query attends as is, its scale folded into softmax_scale. + if q_fp8 is None: + qp, softmax_scale = q[nd:], self.scale + else: + qp, softmax_scale = q_fp8[nd:], self.scale * q_scale_float k2q_row_ptr, k2q_q_indices, schedule = build_k2q_csr( prefill_topk, p.cu_seqlens_q, @@ -375,7 +375,7 @@ def forward( max_seqlen_k=p.max_seq_len, blk_kv=SPARSE_BLOCK_SIZE, causal=True, - softmax_scale=self.scale, + softmax_scale=softmax_scale, page_table=p.block_table, seqused_k=p.seq_lens, schedule=schedule, @@ -393,7 +393,7 @@ def forward( topK=self.topk_blocks, blk_kv=SPARSE_BLOCK_SIZE, causal=True, - softmax_scale=self.scale, + softmax_scale=softmax_scale, cu_seqlens_q=p.cu_seqlens_q, cu_seqlens_k=p.cu_seqlens_k, max_seqlen_q=p.max_query_len, From 61f084b5f3e58aa3d3eb5c86aeaf4686c460c478 Mon Sep 17 00:00:00 2001 From: Bruno Volpato Date: Sun, 11 Oct 2026 02:21:55 -0400 Subject: [PATCH 252/624] [Bugfix][Multimodal] Fix per-video cache option reuse (#55203) Signed-off-by: bvolpato Co-authored-by: Codex --- .../multimodal/processing/test_qwen3_vl.py | 138 +++++++++++++----- vllm/model_executor/models/qwen3_vl.py | 69 ++++++++- vllm/multimodal/processing/processor.py | 35 +++-- 3 files changed, 197 insertions(+), 45 deletions(-) diff --git a/tests/models/multimodal/processing/test_qwen3_vl.py b/tests/models/multimodal/processing/test_qwen3_vl.py index 24281098154f..0c74a86ff000 100644 --- a/tests/models/multimodal/processing/test_qwen3_vl.py +++ b/tests/models/multimodal/processing/test_qwen3_vl.py @@ -15,6 +15,7 @@ from vllm.config.multimodal import MultiModalDummyOptions from vllm.model_executor.models.qwen3_vl import Qwen3VLProcessingInfo from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.cache import MultiModalProcessorOnlyCache from vllm.multimodal.processing.context import BaseProcessingInfo from ...registry import HF_EXAMPLE_MODELS @@ -336,65 +337,136 @@ def test_processor_kwargs_videos_kwargs_does_not_leak_into_image_budget( ("hf_mm_kwargs", "scalar_kwargs"), [ ({"num_frames": [8, 16]}, [{"num_frames": 8}, {"num_frames": 16}]), - ({"fps": [4.0, 12.0]}, [{"fps": 4.0}, {"fps": 12.0}]), + ({"fps": [2.0, 4.0]}, [{"fps": 2.0}, {"fps": 4.0}]), + ({"num_frames": [8, 16, 32]}, [{"num_frames": 8}, {"num_frames": 16}]), + ({"fps": [2.0, 4.0, 8.0]}, [{"fps": 2.0}, {"fps": 4.0}]), ], ) +@pytest.mark.parametrize("model_kwargs", [False, True]) +@pytest.mark.parametrize("use_uuids", [False, True]) +@pytest.mark.parametrize("scoped", [False, True]) +@pytest.mark.skip_global_cleanup def test_processor_multi_video_list_kwargs( model_id: str, hf_mm_kwargs: dict[str, Any], scalar_kwargs: list[dict[str, Any]], + model_kwargs: bool, + use_uuids: bool, + scoped: bool, ) -> None: - """List-valued video kwargs must be scalarized per item without mutation. - - Each video in the multi-video request must match processing the same input - independently with its corresponding scalar ``fps``/``num_frames`` value. - """ + """Per-video options must survive duplicate items and partial cache hits.""" + if scoped: + hf_mm_kwargs = {"videos_kwargs": hf_mm_kwargs} + hf_mm_kwargs_before = deepcopy(hf_mm_kwargs) ctx = build_model_context( model_id, limit_mm_per_prompt={"image": 0, "video": 2}, + mm_processor_kwargs=hf_mm_kwargs if model_kwargs else None, + mm_processor_cache_gb=1, ) processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) - - prompt = ( - "<|vision_start|><|video_pad|><|vision_end|>" - "<|vision_start|><|video_pad|><|vision_end|>" - ) - mm_data = { - "video": [ - _build_video_mm_data(num_frames=32)["video"][0], - _build_video_mm_data(num_frames=32)["video"][0], - ] - } - - hf_mm_kwargs_before = deepcopy(hf_mm_kwargs) - processed = processor( - prompt, - mm_items=processor.info.parse_mm_data(mm_data), - hf_processor_mm_kwargs=hf_mm_kwargs, - ) + cache = MultiModalProcessorOnlyCache(ctx.model_config) + prompt = "<|vision_start|><|video_pad|><|vision_end|>" * 2 + videos = [ + _build_video_mm_data(num_frames=32, original_fps=8.0)["video"][0] + for _ in range(3) + ] + for i, (frames, _) in enumerate(videos): + frames.fill(i) def process_single(hf_kwargs: dict[str, Any]) -> torch.Tensor: single_prompt = "<|vision_start|><|video_pad|><|vision_end|>" - single_mm_data = _build_video_mm_data(num_frames=32) single = processor( single_prompt, - mm_items=processor.info.parse_mm_data(single_mm_data), + mm_items=processor.info.parse_mm_data({"video": [videos[0]]}), hf_processor_mm_kwargs=hf_kwargs, ) return single["mm_kwargs"].get_data()["video_grid_thw"][0] expected_grids = [process_single(kwargs) for kwargs in scalar_kwargs] - multi_grids = processed["mm_kwargs"].get_data()["video_grid_thw"] - assert not torch.equal(expected_grids[0], expected_grids[1]) - assert torch.equal(multi_grids[0], expected_grids[0]) - assert torch.equal(multi_grids[1], expected_grids[1]) + + def run(indices, *, cache=None, cached_only=False): + return processor( + prompt, + mm_items=processor.info.parse_mm_data( + {"video": [None if cached_only else videos[i] for i in indices]} + ), + mm_uuid_items=( + {"video": [f"video-{i}" for i in indices]} if use_uuids else None + ), + hf_processor_mm_kwargs={} if model_kwargs else hf_mm_kwargs, + cache=cache, + ) + + for indices in ([0, 0], [0, 1], [0, 2], [2, 0]): + expected = run(indices) + actual = run(indices, cache=cache) + assert actual["prompt_token_ids"] == expected["prompt_token_ids"] + assert actual["mm_placeholders"] == expected["mm_placeholders"] + assert actual["mm_kwargs"] == expected["mm_kwargs"] + assert actual["mm_hashes"] == expected["mm_hashes"] + assert len(set(actual["mm_hashes"]["video"])) == 2 + multi_grids = actual["mm_kwargs"].get_data()["video_grid_thw"] + assert torch.equal(multi_grids[0], expected_grids[0]) + assert torch.equal(multi_grids[1], expected_grids[1]) + + if use_uuids: + cached_only = run([2, 0], cache=cache, cached_only=True) + assert cached_only["prompt_token_ids"] == actual["prompt_token_ids"] + assert cached_only["mm_kwargs"] == actual["mm_kwargs"] assert hf_mm_kwargs == hf_mm_kwargs_before - video_phs = processed["mm_placeholders"].get("video", []) - assert len(video_phs) == 2, ( - f"Expected exactly 2 video placeholders, got {len(video_phs)}" + +@pytest.mark.parametrize(("key", "value"), [("fps", 2.0), ("num_frames", 8)]) +@pytest.mark.parametrize("length", [0, 1]) +@pytest.mark.parametrize("use_cache", [False, True]) +@pytest.mark.parametrize("scoped", [False, True]) +def test_processor_rejects_short_per_video_options( + key: str, value: float | int, length: int, use_cache: bool, scoped: bool +) -> None: + """Reject malformed option lists with and without a warm processor cache.""" + ctx = build_model_context( + MODEL_ID, + limit_mm_per_prompt={"image": 0, "video": 2}, + ) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + cache = MultiModalProcessorOnlyCache(ctx.model_config) if use_cache else None + prompt = "<|vision_start|><|video_pad|><|vision_end|>" * 2 + videos = _build_video_mm_data(num_frames=32, original_fps=8.0)["video"] * 2 + + def process(values: list[float | int]) -> None: + kwargs = {key: values} + processor( + prompt, + mm_items=processor.info.parse_mm_data({"video": videos}), + hf_processor_mm_kwargs={"videos_kwargs": kwargs} if scoped else kwargs, + cache=cache, + ) + + if use_cache: + process([value, value]) + + with pytest.raises(ValueError, match=f"{key} must have at least 2 entries"): + process([value] * length) + + +@pytest.mark.parametrize("kwargs", [{"fps": [2.0, 4.0]}, {"num_frames": [8, 16]}]) +def test_processor_profiles_fewer_videos_than_configured_options( + kwargs: dict[str, list[float] | list[int]], +) -> None: + """Profiling may use fewer videos than the configured per-video list.""" + ctx = build_model_context( + MODEL_ID, + limit_mm_per_prompt={"image": 0, "video": 2}, + mm_processor_kwargs=kwargs, ) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + + profiled = processor.get_dummy_mm_inputs({"video": 1}) + + assert len(profiled["mm_placeholders"]["video"]) == 1 + assert profiled["mm_kwargs"].get_data()["video_grid_thw"].shape[0] == 1 def _build_video_embeds_mm_data( diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index 77780b3d1d40..527e2b0f3e2e 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -60,7 +60,7 @@ VideoPruningMethod, ) from vllm.distributed import get_pp_group, parallel_state -from vllm.inputs import MultiModalDataDict +from vllm.inputs import MultiModalDataDict, MultiModalHashes from vllm.logger import init_logger from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY from vllm.model_executor.layers.attention.mm_encoder_attention import ( @@ -81,6 +81,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.hasher import MultiModalHasher from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalFieldConfig, @@ -94,6 +95,7 @@ from vllm.multimodal.processing import ( BaseDummyInputsBuilder, BaseMultiModalProcessor, + ProcessorInputs, PromptReplacement, PromptUpdate, PromptUpdateDetails, @@ -1342,6 +1344,67 @@ def _replace_video_token_placeholders( class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]): + def _get_per_video_kwargs( + self, kwargs: Mapping[str, object], num_videos: int | None = None + ) -> dict[str, list[Any]]: + merged = self.info._merge_and_resolve_mm_processor_kwargs(kwargs) + per_video_kwargs = { + key: value + for key, value in merged.get("videos_kwargs", {}).items() + if key in ("fps", "num_frames") and isinstance(value, list) + } + if num_videos: + for key, values in per_video_kwargs.items(): + if len(values) < num_videos: + raise ValueError( + f"{key} must have at least {num_videos} entries " + "(one per video); " + f"got {len(values)}." + ) + return per_video_kwargs + + def _get_mm_hashes(self, inputs: ProcessorInputs) -> MultiModalHashes: + hashes = dict(super()._get_mm_hashes(inputs)) + per_video_kwargs = self._get_per_video_kwargs( + inputs.hf_processor_mm_kwargs, len(hashes.get("video", [])) + ) + if per_video_kwargs and "video" in hashes: + algorithm = self.info.ctx.get_mm_config().mm_hasher_algorithm + hashes["video"] = [ + MultiModalHasher.hash_kwargs( + algorithm, + mm_hash=mm_hash, + **{key: values[i] for key, values in per_video_kwargs.items()}, + ) + for i, mm_hash in enumerate(hashes["video"]) + ] + return hashes + + def _get_cache_missing_processor_kwargs( + self, + inputs: ProcessorInputs, + mm_is_cached: Mapping[str, list[bool]], + ) -> Mapping[str, object]: + per_video_kwargs = self._get_per_video_kwargs( + inputs.hf_processor_mm_kwargs, len(mm_is_cached.get("video", [])) + ) + if not per_video_kwargs or "video" not in mm_is_cached: + return inputs.hf_processor_mm_kwargs + kwargs = self.info._merge_and_resolve_mm_processor_kwargs( + inputs.hf_processor_mm_kwargs + ) + kwargs["videos_kwargs"].update( + { + key: [ + value + for value, hit in zip(values, mm_is_cached["video"]) + if not hit + ] + for key, values in per_video_kwargs.items() + } + ) + return kwargs + @staticmethod def _expands_only_video_token(hf_processor: ProcessorMixin) -> bool: """Transformers>=5.10 processors override `replace_video_token` @@ -1412,6 +1475,7 @@ def _apply_hf_processor_main( video_input_ids_lst: list[list[int]] = [] if videos := hf_data.pop("videos", []): assert isinstance(videos, list) + self._get_per_video_kwargs(hf_kwargs, len(videos)) video_grid_thw_lst = [] pixel_values_videos_lst = [] timestamps_per_video = [] @@ -1544,6 +1608,9 @@ def _apply_hf_processor_main( non_video_scoped_kwargs = dict(non_video_mm_kwargs["videos_kwargs"]) non_video_scoped_kwargs.pop("fps", None) non_video_scoped_kwargs.pop("num_frames", None) + # The context merges model defaults again when calling HF. + for key in self._get_per_video_kwargs({}): + non_video_scoped_kwargs[key] = None non_video_mm_kwargs["videos_kwargs"] = non_video_scoped_kwargs processed_data = self.info.ctx.call_hf_processor( self.info.get_hf_processor(**non_video_mm_kwargs), diff --git a/vllm/multimodal/processing/processor.py b/vllm/multimodal/processing/processor.py index abdb58540255..5b46267dff54 100644 --- a/vllm/multimodal/processing/processor.py +++ b/vllm/multimodal/processing/processor.py @@ -1393,6 +1393,21 @@ def _recompute_cached_prompt_update( """ return replace(cached_update, item_idx=new_item_idx) + def _get_mm_hashes(self, inputs: ProcessorInputs) -> MultiModalHashes: + """Get each item's processor and encoder cache identity.""" + return inputs.get_mm_hashes( + self.info.model_id, + self.info.ctx.get_mm_config().mm_hasher_algorithm, + ) + + def _get_cache_missing_processor_kwargs( + self, + inputs: ProcessorInputs, + mm_is_cached: MultiModalIsCached, + ) -> Mapping[str, object]: + """Select processor options for the items missing from the cache.""" + return inputs.hf_processor_mm_kwargs + def _merge_mm_kwargs( self, cache: BaseMultiModalProcessorCache, @@ -1464,10 +1479,7 @@ def _apply_hf_processor( # Use overrides if provided; fallback to data-dependent hashing. with timing_ctx.record("get_mm_hashes"): - mm_hashes = inputs.get_mm_hashes( - self.info.model_id, - self.info.ctx.get_mm_config().mm_hasher_algorithm, - ) + mm_hashes = self._get_mm_hashes(inputs) mm_prompt_updates = self._get_mm_prompt_updates( inputs.mm_data_items, @@ -1500,10 +1512,7 @@ def _cached_apply_hf_processor( return self._apply_hf_processor(inputs, timing_ctx) with timing_ctx.record("get_mm_hashes"): - mm_hashes = inputs.get_mm_hashes( - self.info.model_id, - self.info.ctx.get_mm_config().mm_hasher_algorithm, - ) + mm_hashes = self._get_mm_hashes(inputs) with timing_ctx.record("get_cache_missing_items"): mm_is_cached, mm_missing_data_items = self._get_cache_missing_items( @@ -1512,25 +1521,29 @@ def _cached_apply_hf_processor( mm_hashes=mm_hashes, ) + missing_processor_kwargs = self._get_cache_missing_processor_kwargs( + inputs, mm_is_cached + ) + # NOTE: The prompt does not correspond to `mm_missing_data_items`, # so we can't apply prompt updates until the new multimodal # items are combined with the cached multimodal items with timing_ctx.record("apply_hf_processor"): mm_missing_processed_data = self._apply_hf_processor_main( mm_items=mm_missing_data_items, - hf_kwargs=inputs.hf_processor_mm_kwargs, + hf_kwargs=missing_processor_kwargs, ) mm_missing_kwargs = MultiModalKwargsItems.from_hf_inputs( mm_missing_processed_data, self._get_mm_fields_config( - mm_missing_processed_data, inputs.hf_processor_mm_kwargs + mm_missing_processed_data, missing_processor_kwargs ), ) mm_missing_prompt_updates = self._get_mm_prompt_updates( mm_missing_data_items, - inputs.hf_processor_mm_kwargs, + missing_processor_kwargs, mm_missing_kwargs, ) From ef42b4385f8512567703692a5bfbb9c0edc59c89 Mon Sep 17 00:00:00 2001 From: debanshd Date: Sun, 11 Oct 2026 00:03:06 -0700 Subject: [PATCH 253/624] [Bugfix] Reclaim cached GPU memory during weight loading (#57516) Signed-off-by: Debanshu Das Signed-off-by: Isotr0py Co-authored-by: Isotr0py Co-authored-by: Kimi Code --- vllm/model_executor/model_loader/base_loader.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/vllm/model_executor/model_loader/base_loader.py b/vllm/model_executor/model_loader/base_loader.py index 4c774808bf79..65f64a33267b 100644 --- a/vllm/model_executor/model_loader/base_loader.py +++ b/vllm/model_executor/model_loader/base_loader.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import gc import time from abc import ABC, abstractmethod @@ -86,6 +87,11 @@ def load_model( logger.debug("Loading weights on %s ...", load_device) self.load_weights(model, model_config) + # Release transient tensors held by PyTorch's caching allocator + # before weight post-processing and raw CUDA driver allocations. + gc.collect() + torch.accelerator.empty_cache() + # Log peak GPU memory after loading weights. This is needed # to have test coverage on peak memory for online quantization. if current_platform.is_cuda_alike() or current_platform.is_xpu(): @@ -103,6 +109,9 @@ def load_model( process_weights_after_loading(model, model_config, target_device) log_online_quantization_time(vllm_config) + gc.collect() + torch.accelerator.empty_cache() + return model.eval() From 26a488ad93d12c658449a19dea2eb78f1be38f83 Mon Sep 17 00:00:00 2001 From: Xun Sun Date: Sun, 11 Oct 2026 15:10:04 +0800 Subject: [PATCH 254/624] [Fast Start][Bugfix] Snapshot GPU memory after external weights are ready (#60966) --- vllm/v1/worker/gpu_worker.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 37fa00bb7efa..152def62221c 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -512,12 +512,13 @@ def init_device(self): gc.collect() torch.accelerator.empty_cache() - # take current memory snapshot - self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device) # Weights from external model loader process external_weight_memory = get_model_loader( self.load_config ).get_external_weight_memory(self.vllm_config) + + # take current memory snapshot after external weights are ready + self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device) self.requested_memory = request_memory( init_snapshot, self.cache_config, external_weight_memory ) From fcae5c98870ff2c20659731dce7ba3a3ff0d48aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20P=C3=A9rez=20de=20Algaba?= <124347725+jperezdealgaba@users.noreply.github.com> Date: Sun, 11 Oct 2026 10:26:02 +0200 Subject: [PATCH 255/624] [Security] Restrict default CORS origins (#61009) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Juan Pérez de Algaba Co-authored-by: Cyrus Leung --- rust/src/cmd/src/cli.rs | 10 ++- rust/src/cmd/src/cli/tests.rs | 28 +++------ rust/src/server/src/config.rs | 4 +- rust/src/server/src/middleware/cors.rs | 2 +- rust/src/server/src/routes/tests.rs | 44 ++++++++++--- tests/entrypoints/launchers/test_cli_args.py | 7 +++ .../serve/middleware/test_cors_middleware.py | 62 +++++++++++++++++++ vllm/entrypoints/launchers/cli_args.py | 4 +- 8 files changed, 124 insertions(+), 37 deletions(-) create mode 100644 tests/entrypoints/serve/middleware/test_cors_middleware.py diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 14b949c87de4..eeb458c86e75 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -401,9 +401,9 @@ pub struct SharedRuntimeArgs { #[serde(default)] pub served_model_name: Vec, - /// CORS allowed origins as a JSON list. `["*"]` allows any origin. - #[arg(long, value_parser = parse_json::, value_name = "JSON", default_value = r#"["*"]"#)] - #[serde(default = "default_cors_wildcard")] + /// CORS allowed origins as a JSON list. Defaults to `[]` (no cross-origin access). + #[arg(long, value_parser = parse_json::, value_name = "JSON", default_value = "[]")] + #[serde(default = "default_cors_origins")] pub allowed_origins: JsonStringList, /// CORS allowed methods as a JSON list. `["*"]` allows the standard set. @@ -664,6 +664,10 @@ fn default_cors_wildcard() -> JsonStringList { JsonStringList(vec!["*".to_string()]) } +fn default_cors_origins() -> JsonStringList { + JsonStringList(Vec::new()) +} + fn default_py_bootstrap_parser_selection() -> ParserSelection { ParserSelection::None } diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 57780027deac..cc545ce33e80 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -205,9 +205,7 @@ fn serve_args_forward_python_flags_with_separator() { disable_log_stats: false, served_model_name: [], allowed_origins: JsonStringList( - [ - "*", - ], + [], ), allowed_methods: JsonStringList( [ @@ -1075,9 +1073,7 @@ fn frontend_args_accept_json() { disable_log_stats: false, served_model_name: [], allowed_origins: JsonStringList( - [ - "*", - ], + [], ), allowed_methods: JsonStringList( [ @@ -1382,13 +1378,13 @@ fn serve_args_parse_cors_flags() { } #[test] -fn serve_args_cors_defaults_are_permissive() { +fn serve_args_cors_defaults_restrict_origins() { let cli = Cli::try_parse_from(["vllm-rs", "serve", "Qwen/Qwen3-0.6B"]).unwrap(); let Command::Serve(serve) = cli.command else { panic!("expected serve args"); }; - assert_eq!(serve.runtime.allowed_origins.0, ["*"]); + assert!(serve.runtime.allowed_origins.0.is_empty()); assert_eq!(serve.runtime.allowed_methods.0, ["*"]); assert_eq!(serve.runtime.allowed_headers.0, ["*"]); assert!(!serve.runtime.allow_credentials); @@ -1748,9 +1744,7 @@ fn serve_args_accept_handshake_aliases() { disable_log_stats: false, served_model_name: [], allowed_origins: JsonStringList( - [ - "*", - ], + [], ), allowed_methods: JsonStringList( [ @@ -1906,9 +1900,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { sse_keep_alive_interval: None, }, cors: CorsConfig { - allow_origins: [ - "*", - ], + allow_origins: [], allow_methods: [ "*", ], @@ -2002,9 +1994,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { sse_keep_alive_interval: None, }, cors: CorsConfig { - allow_origins: [ - "*", - ], + allow_origins: [], allow_methods: [ "*", ], @@ -2142,9 +2132,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present sse_keep_alive_interval: None, }, cors: CorsConfig { - allow_origins: [ - "*", - ], + allow_origins: [], allow_methods: [ "*", ], diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index 3f95b5891cb2..ca91ba40cd50 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -65,7 +65,7 @@ pub struct ApiServerOptions { pub sse_keep_alive_interval: Option, } -/// CORS settings mirroring Python's `CORSMiddleware`; the default is permissive. +/// CORS settings mirroring Python's `CORSMiddleware`; origins are denied by default. #[derive(Debug, Clone, PartialEq, Eq, Serialize)] pub struct CorsConfig { /// Allowed origins. `["*"]` allows any origin. @@ -81,7 +81,7 @@ pub struct CorsConfig { impl Default for CorsConfig { fn default() -> Self { Self { - allow_origins: vec!["*".to_string()], + allow_origins: Vec::new(), allow_methods: vec!["*".to_string()], allow_headers: vec!["*".to_string()], allow_credentials: false, diff --git a/rust/src/server/src/middleware/cors.rs b/rust/src/server/src/middleware/cors.rs index 89d1b6f3b7bd..bfc030a3a9d9 100644 --- a/rust/src/server/src/middleware/cors.rs +++ b/rust/src/server/src/middleware/cors.rs @@ -105,7 +105,7 @@ pub fn cors_layer(cfg: &CorsConfig) -> CorsLayer { // Emit `Vary: Origin` only when the allow-origin is dynamic (explicit // origins, or credentials); the wildcard + no-credentials case emits no - // `Vary` at all, and an empty list disables the header here. + // `Vary` at all. An empty origin list never permits an origin. let vary: Vec = if !wildcard_origins || cfg.allow_credentials { vec![header::ORIGIN] } else { diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 118e412b45b6..7ba4bbf0f416 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -1784,7 +1784,7 @@ async fn api_key_auth_allows_unguarded_route_without_token() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] -async fn cors_default_simple_request_allows_any_origin() { +async fn cors_default_simple_request_omits_allow_origin() { let (mut app, _engine_task) = test_app_with_cors(CorsConfig::default()).await; let response = app .call( @@ -1799,17 +1799,12 @@ async fn cors_default_simple_request_allows_any_origin() { .expect("call app"); assert_eq!(response.status(), StatusCode::OK); - assert_eq!( - header_value(&response, "access-control-allow-origin"), - Some("*") - ); - // Wildcard origins without credentials emit no `Vary` (Starlette parity). - assert_eq!(header_value(&response, "vary"), None); + assert_eq!(header_value(&response, "access-control-allow-origin"), None); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] -async fn cors_default_preflight_returns_explicit_methods_and_max_age() { +async fn cors_default_preflight_omits_allow_origin() { let (mut app, _engine_task) = test_app_with_cors(CorsConfig::default()).await; let response = app .call( @@ -1825,6 +1820,32 @@ async fn cors_default_preflight_returns_explicit_methods_and_max_age() { .await .expect("call app"); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(header_value(&response, "access-control-allow-origin"), None); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn cors_wildcard_preflight_returns_explicit_methods_and_max_age() { + let cors = CorsConfig { + allow_origins: vec!["*".to_string()], + ..CorsConfig::default() + }; + let (mut app, _engine_task) = test_app_with_cors(cors).await; + let response = app + .call( + Request::builder() + .method("OPTIONS") + .uri("/v1/chat/completions") + .header("origin", "http://example.com") + .header("access-control-request-method", "POST") + .header("access-control-request-headers", "content-type") + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("call app"); + assert_eq!(response.status(), StatusCode::OK); // `*` methods expand to the explicit method list, matching Starlette // (never the literal `*`). @@ -1917,6 +1938,7 @@ async fn cors_explicit_origin_disallowed_omits_allow_origin() { #[serial] async fn cors_wildcard_with_credentials_reflects_origin_without_panic() { let cors = CorsConfig { + allow_origins: vec!["*".to_string()], allow_credentials: true, ..CorsConfig::default() }; @@ -1971,8 +1993,12 @@ async fn cors_unauthorized_response_has_no_cors_headers() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn cors_preflight_bypasses_auth_and_returns_cors_headers() { + let cors = CorsConfig { + allow_origins: vec!["*".to_string()], + ..CorsConfig::default() + }; let (mut app, _engine_task) = - test_app_with_cors_and_keys(CorsConfig::default(), vec!["secret".to_string()]).await; + test_app_with_cors_and_keys(cors, vec!["secret".to_string()]).await; let response = app .call( Request::builder() diff --git a/tests/entrypoints/launchers/test_cli_args.py b/tests/entrypoints/launchers/test_cli_args.py index 84279e155582..0a4c97141d07 100644 --- a/tests/entrypoints/launchers/test_cli_args.py +++ b/tests/entrypoints/launchers/test_cli_args.py @@ -44,6 +44,13 @@ def serve_parser(): return _build_vllm_parsers()["vllm serve"] +def test_cors_origins_require_explicit_opt_in(serve_parser): + assert serve_parser.parse_args([]).allowed_origins == [] + assert serve_parser.parse_args( + ["--allowed-origins", '["http://trusted.example"]'] + ).allowed_origins == ["http://trusted.example"] + + ### Test config parsing def test_config_arg_parsing(serve_parser, cli_config_file): args = serve_parser.parse_args([]) diff --git a/tests/entrypoints/serve/middleware/test_cors_middleware.py b/tests/entrypoints/serve/middleware/test_cors_middleware.py new file mode 100644 index 000000000000..543c101a9307 --- /dev/null +++ b/tests/entrypoints/serve/middleware/test_cors_middleware.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from fastapi import FastAPI +from starlette.testclient import TestClient + +from vllm.entrypoints.launchers.cli_args import FrontendArgs +from vllm.entrypoints.serve.middleware.register import init_entrypoints_middleware + + +def _make_client(allowed_origins: list[str] | None = None) -> TestClient: + args = FrontendArgs() + if allowed_origins is not None: + args.allowed_origins = allowed_origins + + app = FastAPI() + + @app.get("/v1/models") + async def models(): + return {"data": [{"id": "private-model"}]} + + @app.post("/v1/chat/completions") + async def chat(): + return {"message": {"content": "ok"}} + + init_entrypoints_middleware(args, app, ()) + return TestClient(app) + + +def test_default_origins_prevent_cross_origin_reads_and_json_requests(): + client = _make_client() + response = client.get("/v1/models", headers={"Origin": "http://attacker.example"}) + assert response.status_code == 200 + assert "access-control-allow-origin" not in response.headers + + preflight = client.options( + "/v1/chat/completions", + headers={ + "Origin": "http://attacker.example", + "Access-Control-Request-Method": "POST", + "Access-Control-Request-Headers": "content-type", + }, + ) + assert preflight.status_code == 400 + assert "access-control-allow-origin" not in preflight.headers + + +def test_explicit_origin_allows_browser_requests(): + client = _make_client(["http://trusted.example"]) + response = client.get("/v1/models", headers={"Origin": "http://trusted.example"}) + assert response.headers["access-control-allow-origin"] == "http://trusted.example" + + preflight = client.options( + "/v1/chat/completions", + headers={ + "Origin": "http://trusted.example", + "Access-Control-Request-Method": "POST", + "Access-Control-Request-Headers": "content-type", + }, + ) + assert preflight.status_code == 200 + assert preflight.headers["access-control-allow-origin"] == "http://trusted.example" diff --git a/vllm/entrypoints/launchers/cli_args.py b/vllm/entrypoints/launchers/cli_args.py index 5cd58da87eac..e2c2625e9462 100644 --- a/vllm/entrypoints/launchers/cli_args.py +++ b/vllm/entrypoints/launchers/cli_args.py @@ -316,8 +316,8 @@ class FrontendArgs(BaseFrontendArgs): while keeping logs for other endpoints.""" allow_credentials: bool = False """Allow credentials.""" - allowed_origins: list[str] = field(default_factory=lambda: ["*"]) - """Allowed origins.""" + allowed_origins: list[str] = field(default_factory=list) + """Allowed origins. Cross-origin browser access is disabled by default.""" allowed_methods: list[str] = field(default_factory=lambda: ["*"]) """Allowed methods.""" allowed_headers: list[str] = field(default_factory=lambda: ["*"]) From 1789f095881d14966d7d4d9a2647748bf32dd834 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sun, 11 Oct 2026 16:58:07 +0800 Subject: [PATCH 256/624] [Frontend] Explain disabled scale-out endpoints with opt-in hints (#61042) Signed-off-by: aoshen02 Co-authored-by: OpenAI Codex --- docs/usage/security.md | 2 +- rust/src/server/src/routes.rs | 24 ++++++++- rust/src/server/src/routes/tests.rs | 22 +++++--- tests/entrypoints/scale_out/test_factories.py | 52 ++++++++++++++++--- vllm/entrypoints/scale_out/factories.py | 41 ++++++++++++--- 5 files changed, 117 insertions(+), 24 deletions(-) diff --git a/docs/usage/security.md b/docs/usage/security.md index ab5aa96becc5..50c4149be81b 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -178,7 +178,7 @@ When `--api-key` is configured, the following endpoints require Bearer token aut - `/v1/systemone` - Structured decisions API - `/v1/load_lora_adapter` - Load a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`) - `/v1/unload_lora_adapter` - Unload a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`) -- `/inference/v1/generate` - Generate completions (available when `--enable-scale-out` is set, or with `--tokens-only`) +- `/inference/v1/generate` - Generate completions (available when `--enable-scale-out` is set, or with `--tokens-only`; otherwise returns HTTP 403 with instructions to enable it) - `/cohere/v2/chat/render` - Render Cohere Chat v2 requests (requires `VLLM_ENABLE_COHERE_API=1` and either `--enable-scale-out` or `--tokens-only`) - `/v2/embed` - Cohere Embed API - `/v2/rerank` - Cohere Rerank API diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs index 28f35573249e..251398611ba1 100644 --- a/rust/src/server/src/routes.rs +++ b/rust/src/server/src/routes.rs @@ -24,17 +24,36 @@ mod world_size; use std::sync::Arc; -use axum::Router; use axum::extract::DefaultBodyLimit; +use axum::http::StatusCode; use axum::middleware::{from_fn, from_fn_with_state}; use axum::routing::{get, post}; +use axum::{Json, Router}; use tower_http::trace::TraceLayer; use tracing::{info, warn}; use crate::DEFAULT_REQUEST_BODY_LIMIT_BYTES; use crate::middleware; +use crate::routes::openai::utils::types::{ErrorDetail, ErrorResponse}; use crate::state::AppState; +const SCALE_OUT_DISABLED_MESSAGE: &str = + "Scale-out endpoints are disabled. Restart vllm serve with --enable-scale-out to enable them."; + +async fn disabled_scale_out() -> (StatusCode, Json) { + ( + StatusCode::FORBIDDEN, + Json(ErrorResponse { + error: ErrorDetail { + message: SCALE_OUT_DISABLED_MESSAGE.to_string(), + error_type: "PermissionDeniedError".to_string(), + param: None, + code: Some("403".to_string()), + }, + }), + ) +} + fn server_dev_mode_enabled() -> bool { std::env::var("VLLM_SERVER_DEV_MODE") .ok() @@ -105,7 +124,8 @@ fn build_router_with_options( if scale_out_endpoints_enabled { router = router.route("/inference/v1/generate", post(inference::generate)); } else { - info!("scale-out endpoints are disabled; pass --enable-scale-out to enable them"); + info!("{SCALE_OUT_DISABLED_MESSAGE}"); + router = router.route("/inference/v1/generate", post(disabled_scale_out)); } if runtime_lora_updating_enabled { diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 7ba4bbf0f416..d574e02a9215 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -758,7 +758,7 @@ async fn scale_out_generate_route_follows_api_server_options() { .await .expect("call app"); - assert_ne!(response.status(), StatusCode::NOT_FOUND); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -770,10 +770,7 @@ async fn scale_out_generate_route_is_disabled_by_default() { Arc::new(FakeChatBackend::new()), ) .await; - let mut app = build_router_with_scale_out_endpoints( - Arc::new(AppState::new(vec!["model".to_string()], chat)), - false, - ); + let mut app = build_router(Arc::new(AppState::new(vec!["model".to_string()], chat))); let response = app .call( @@ -781,13 +778,22 @@ async fn scale_out_generate_route_is_disabled_by_default() { .method("POST") .uri("/inference/v1/generate") .header("content-type", "application/json") - .body(Body::from("{}")) + .body(Body::from("not JSON")) .expect("build request"), ) .await .expect("call app"); - assert_eq!(response.status(), StatusCode::NOT_FOUND); + assert_eq!(response.status(), StatusCode::FORBIDDEN); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read response"); + let error: serde_json::Value = serde_json::from_slice(&body).expect("decode error"); + assert_eq!(error["error"]["type"], "PermissionDeniedError"); + assert!( + error["error"]["message"] + .as_str() + .expect("error message") + .contains("--enable-scale-out") + ); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -816,7 +822,7 @@ async fn scale_out_generate_route_can_be_enabled() { .await .expect("call app"); - assert_ne!(response.status(), StatusCode::NOT_FOUND); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); } /// Build a dev-mode router backed by a mock engine using a custom ready diff --git a/tests/entrypoints/scale_out/test_factories.py b/tests/entrypoints/scale_out/test_factories.py index ea96fff876d1..7937c0d45bc0 100644 --- a/tests/entrypoints/scale_out/test_factories.py +++ b/tests/entrypoints/scale_out/test_factories.py @@ -1,13 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import logging from argparse import Namespace +import pytest from fastapi import FastAPI +from fastapi.testclient import TestClient from vllm.entrypoints.launchers.api_server.routers import register_api_routers +from vllm.entrypoints.scale_out.factories import register_scale_out_api_routers from vllm.entrypoints.scale_out.token_in_token_out.api_router import attach_router -from vllm.entrypoints.serve.middleware.authenticate import GUARDED_PREFIX +from vllm.entrypoints.serve.middleware.authenticate import ( + GUARDED_PREFIX, + AuthenticationMiddleware, +) from vllm.tasks import SupportedTask RENDER_PATHS = { @@ -41,17 +48,50 @@ def registered_paths( return {route.path for route in app.routes} -def test_scale_out_routes_are_disabled_by_default(): - paths = registered_paths(("generate",)) +@pytest.mark.parametrize("path", sorted(SCALE_OUT_PATHS)) +@pytest.mark.parametrize("authorized", [False, True]) +def test_disabled_scale_out_routes_explain_opt_in_without_parsing(path, authorized): + app = FastAPI() + register_scale_out_api_routers(app, ("generate",)) + app.add_middleware(AuthenticationMiddleware, tokens=["secret"]) + headers = {"Authorization": "Bearer secret"} if authorized else {} + client = TestClient(app) + response = client.post(path, content="not JSON", headers=headers) - assert SCALE_OUT_PATHS.isdisjoint(paths) + assert response.status_code == (403 if authorized else 401) + if authorized: + assert "--enable-scale-out" in response.json()["error"]["message"] + assert response.json()["error"]["type"] == "PermissionDeniedError" + assert path not in app.openapi()["paths"] + assert client.post("/unknown", headers=headers).status_code == 404 -def test_disabled_scale_out_routes_are_logged(caplog): +def test_disabled_scale_out_routes_are_logged_at_info(caplog): with caplog.at_level("INFO", logger="vllm.entrypoints.scale_out.factories"): registered_paths(("generate",)) - assert any("--enable-scale-out" in record.message for record in caplog.records) + assert [ + record.levelno + for record in caplog.records + if "--enable-scale-out" in record.message + ] == [logging.INFO] + + +@pytest.mark.parametrize("tokens_only", [False, True]) +def test_enabled_generate_routes_use_native_request_validation(tokens_only): + app = FastAPI() + app.state.args = Namespace( + tokens_only=tokens_only, enable_scale_out=not tokens_only + ) + register_scale_out_api_routers(app, ("generate",)) + response = TestClient(app).post( + GENERATE_PATH, + content="not JSON", + headers={"Content-Type": "application/json"}, + ) + + assert response.status_code == 422 + assert response.json()["detail"][0]["type"] == "json_invalid" def test_scale_out_routes_can_be_enabled_for_generate_server(): diff --git a/vllm/entrypoints/scale_out/factories.py b/vllm/entrypoints/scale_out/factories.py index bd48c709ddc0..45bb9dbcde6b 100644 --- a/vllm/entrypoints/scale_out/factories.py +++ b/vllm/entrypoints/scale_out/factories.py @@ -1,11 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from argparse import Namespace +from http import HTTPStatus from typing import TYPE_CHECKING from fastapi import FastAPI +from fastapi.responses import JSONResponse +from fastapi.routing import APIRoute from vllm.engine.protocol import EngineClient +from vllm.entrypoints.serve.exception_handling.error_response import ( + create_error_response, +) from vllm.logger import init_logger from vllm.tasks import SupportedTask @@ -66,6 +72,9 @@ def register_scale_out_api_routers( app: FastAPI, supported_tasks: tuple["SupportedTask", ...], ): + from .derender.api_router import router as derender_render + from .render.api_router import router as render_render + args = getattr(app.state, "args", None) # `vllm launch render` and `vllm serve --tokens-only` are dedicated to # serving these endpoints, so they always register them regardless of @@ -73,17 +82,35 @@ def register_scale_out_api_routers( dedicated_mode = "render" in supported_tasks or getattr(args, "tokens_only", False) enabled = dedicated_mode or getattr(args, "enable_scale_out", False) if not enabled: - logger.info( - "Scale-out endpoints are disabled. Set --enable-scale-out to enable them." + message = ( + "Scale-out endpoints are disabled. Restart vllm serve with " + "--enable-scale-out to enable them." ) + logger.info(message) + + async def disabled_scale_out() -> JSONResponse: + error = create_error_response( + message, "PermissionDeniedError", HTTPStatus.FORBIDDEN + ) + return JSONResponse(error.model_dump(), status_code=HTTPStatus.FORBIDDEN) + + from .token_in_token_out.api_router import router as token_router + + for route in ( + *render_render.routes, + *derender_render.routes, + *token_router.routes, + ): + if isinstance(route, APIRoute): + app.add_api_route( + route.path, + disabled_scale_out, + methods=sorted(route.methods), + include_in_schema=False, + ) return - from .render.api_router import router as render_render - app.include_router(render_render) - - from .derender.api_router import router as derender_render - app.include_router(derender_render) if "generate" in supported_tasks: From 2339bd39d661965b821bfbe652db8f536e0b3dcd Mon Sep 17 00:00:00 2001 From: Stu Cao Date: Sun, 11 Oct 2026 02:15:18 -0700 Subject: [PATCH 257/624] [Bugfix][Mooncake] Use the EAGLE replay boundary when retaining sliding-window tails (#61012) Signed-off-by: Stu Cao Signed-off-by: Yifan Qiao Co-authored-by: Claude Fable 5.1 Co-authored-by: Yifan Qiao Co-authored-by: Codex --- .../unit/test_mooncake_store_coordinator.py | 44 ++++++++++++++++ .../v1/mooncake/store/coordinator.py | 13 +++-- vllm/v1/core/kv_cache_coordinator.py | 51 +++++++++++-------- vllm/v1/core/single_type_kv_cache_manager.py | 4 +- 4 files changed, 85 insertions(+), 27 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py index ac52ca07ab33..6896777226cf 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py @@ -3,6 +3,7 @@ from math import lcm +import pytest import torch from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501 @@ -422,6 +423,49 @@ def test_store_mask_retention_interval_zero_keeps_only_replay_boundary(): assert masks[1] == [i == 11 for i in range(16)] +def test_store_mask_retention_interval_zero_eagle_keeps_lookup_tail(): + """With EAGLE the replay boundary sits one scheduler block lower, as in + core's ``replay_boundaries``: the lookup matches that block and drops + back from it. The retained SWA tail must span the ``need`` contiguous + blocks the lookup requires (cdiv(8 - 1, 8) + 1 = 2 here); the non-EAGLE + boundary leaves only one of them inside the saved prefix.""" + coord = _make_coord( + _retention_groups(), hash_block_size=8, use_eagle=True, retention_interval=0 + ) + # num_prompt=100 -> boundary 64 -> chunks 7, 8 (8 is the EAGLE peek). + masks = coord.store_mask(96, num_prompt_tokens=100) + assert masks[0] is None + assert masks[1] == [i in (7, 8) for i in range(12)] + # Aligned num_prompt=96 -> boundaries 32 (resend) and 64 (longer sibling). + masks = coord.store_mask(96, num_prompt_tokens=96) + assert masks[1] == [i in (3, 4, 7, 8) for i in range(12)] + + +@pytest.mark.parametrize("num_prompt_tokens", [100, 96]) +def test_store_mask_retention_interval_zero_eagle_round_trip(num_prompt_tokens): + """Everything store_mask keeps is in the store; the EAGLE lookup of the + same prompt must then reach the full block below the one it drops (64 of + 96). The unaligned prompt hit 0 before the fix; the aligned one already + hit.""" + groups = _retention_groups() + coord = _make_coord(groups, hash_block_size=8, use_eagle=True, retention_interval=0) + hashes = _hashes(96 // coord.hash_block_size) + exists: set[tuple[int, bytes]] = set() + store_masks = coord.store_mask(96, num_prompt_tokens=num_prompt_tokens) + for gid, (group, mask) in enumerate(zip(groups, store_masks, strict=True)): + group_hashes = coord.block_hashes_for_spec(hashes, group.kv_cache_spec) + for chunk_id, block_hash in enumerate(group_hashes): + if mask is None or mask[chunk_id]: + exists.add((gid, bytes(block_hash))) + + _masks, hit = coord.find_longest_cache_hit( + hashes, + max_length=96, + cached_block_pool=ExternalCachedBlockPool(coord.hash_block_size, exists), + ) + assert hit == 64 + + def test_store_mask_retention_interval_keeps_segment_and_replay_tails(): """Sparse segment tails (interval=64 -> chunks 7,15) plus the replay boundary tail (num_prompt=100 -> chunk 11) coexist.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index 48a36bc82f8b..2382948268bb 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -11,6 +11,7 @@ ) from vllm.utils.math_utils import cdiv from vllm.v1.core.block_pool import BlockPool +from vllm.v1.core.kv_cache_coordinator import replay_boundaries from vllm.v1.core.kv_cache_utils import ( BlockHash, KVCacheBlock, @@ -341,6 +342,15 @@ def _reachable_masks( f"aligned_token_len ({aligned_token_len}) must be a multiple of " f"{mask_alignment}" ) + reachable_boundaries = ( + () + if num_prompt_tokens is None + else replay_boundaries( + num_prompt_tokens, + self.lcm_block_size, + use_eagle=bool(self.eagle_group_ids), + ) + ) masks: list[list[bool] | None] = [] for g_idx, g in enumerate(self.kv_cache_groups): spec = _unwrap_spec(g.kv_cache_spec) @@ -352,9 +362,6 @@ def _reachable_masks( manager_cls = KVCacheSpecRegistry.get_manager_class(spec) assert manager_cls is not None use_eagle = g_idx in self.eagle_group_ids - reachable_boundaries = ( - () if num_prompt_tokens is None else (num_prompt_tokens - 1,) - ) mask = manager_cls.reachable_block_mask( start_block=start_chunk, end_block=end_chunk, diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 019941a1bf03..b6defc52e847 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -62,6 +62,30 @@ def _validate_prefix_cache_retention_interval( ) +def replay_boundaries( + num_prompt_tokens: int, scheduler_block_size: int, use_eagle: bool +) -> tuple[int, ...]: + """Positions a later request replaying a prompt can resume at. + + A hit is the shortest across all groups, so every group retains state + at each position; EAGLE groups also keep the block above, which they + match and drop back from (see ``reachable_block_mask``). + + Two positions are reachable: a resend of the identical prompt is capped + at ``num_tokens - 1`` (its last token is recomputed for logits), a + longer sibling matches the final aligned block. They differ only on a + block-aligned prompt, where retaining just the higher one collapses the + resend's hit to 0. The alignment is the scheduler block size, not the + finer hash granularity, which would over-estimate the reach. + """ + if not use_eagle: + return (num_prompt_tokens - 1,) + block = scheduler_block_size + resend = (num_prompt_tokens - 1) // block * block + extension = num_prompt_tokens // block * block + return tuple(sorted({max(resend - block, 0), max(extension - block, 0)})) + + class KVCacheCoordinator(ABC): """Coordinate the KV cache of different KV cache groups.""" @@ -319,27 +343,6 @@ def allocate_new_blocks( for manager in self.single_type_managers ) - def get_replay_boundaries(self, request: Request) -> tuple[int, ...]: - """Positions a later request replaying this prompt can resume at. - - A hit is the shortest across all groups, so every group retains state - at each position; EAGLE groups also keep the block above, which they - match and drop back from (see ``reachable_block_mask``). - - Two positions are reachable: a resend of the identical prompt is capped - at ``num_tokens - 1`` (its last token is recomputed for logits), a - longer sibling matches the final aligned block. They differ only on a - block-aligned prompt, where retaining just the higher one collapses the - resend's hit to 0. The alignment is the scheduler block size, not the - finer hash granularity, which would over-estimate the reach. - """ - if not self.eagle_group_ids: - return (request.num_prompt_tokens - 1,) - block = self.scheduler_block_size - resend = (request.num_prompt_tokens - 1) // block * block - extension = request.num_prompt_tokens // block * block - return tuple(sorted({max(resend - block, 0), max(extension - block, 0)})) - def get_num_cacheable_tokens( self, num_computed_tokens: int, kv_cache_group_id: int ) -> int: @@ -358,7 +361,11 @@ def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: (including tokens that are already cached). """ - boundaries = self.get_replay_boundaries(request) + boundaries = replay_boundaries( + request.num_prompt_tokens, + self.scheduler_block_size, + use_eagle=bool(self.eagle_group_ids), + ) for group_id, manager in enumerate(self.single_type_managers): if not manager.enable_caching: continue diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 77c6a4a5f750..218d59f430c9 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -492,7 +492,7 @@ def cache_blocks( boundary; a positive multiple of ``scheduler_block_size`` keeps a tail once per that-sized segment. Only SWA acts on it. replay_boundaries: Positions a later request replaying this prompt - can resume at, from ``get_replay_boundaries``. + can resume at, from ``replay_boundaries``. """ if not self.kv_cache_spec.prefix_cacheable: @@ -505,7 +505,7 @@ def cache_blocks( # Token boundaries whose reachable tail must be retained under sparse # retention: every position a replaying sibling can resume at (see - # ``get_replay_boundaries``) and any detected shared-prefix junction. + # ``replay_boundaries``) and any detected shared-prefix junction. reachable_boundaries = [*replay_boundaries] if request.shared_prefix_boundary: reachable_boundaries.append(request.shared_prefix_boundary) From f425a212081181e3d6ad2ca92daa74ba0f9e83a4 Mon Sep 17 00:00:00 2001 From: Subhash Polisetti <49185065+subhashpolisetti@users.noreply.github.com> Date: Sun, 11 Oct 2026 03:03:33 -0700 Subject: [PATCH 258/624] [Bugfix][Frontend] Support scalar priority and wire priority into LLM.chat() (#57243) Signed-off-by: subhash polisetti --- tests/entrypoints/llm/test_chat.py | 19 +++++++++++++ tests/entrypoints/llm/test_generate.py | 8 ++++++ .../unit_tests/test_offline_utils.py | 6 +++++ vllm/entrypoints/common/offline.py | 17 +++++++----- vllm/entrypoints/llm.py | 27 ++++++++++++++----- 5 files changed, 65 insertions(+), 12 deletions(-) diff --git a/tests/entrypoints/llm/test_chat.py b/tests/entrypoints/llm/test_chat.py index 276426e692c7..8fdd41c51119 100644 --- a/tests/entrypoints/llm/test_chat.py +++ b/tests/entrypoints/llm/test_chat.py @@ -65,6 +65,25 @@ def test_multi_chat(text_llm): assert len(outputs) == 2 +def test_chat_priority(text_llm): + messages = [ + {"role": "system", "content": "You are a helpful assistant"}, + {"role": "user", "content": "Explain the concept of entropy."}, + ] + + # A single value is applied to every conversation. + outputs = text_llm.chat(messages, priority=1) + assert len(outputs) == 1 + + outputs = text_llm.chat([messages, messages], priority=[0, 1]) + assert len(outputs) == 2 + + # Exception raised, if the length of priority does not match the length of + # conversations + with pytest.raises(VLLMValidationError): + text_llm.chat([messages, messages], priority=[0]) + + def test_llm_chat_tokenization_no_double_bos(text_llm): """LLM.chat() should not add special tokens when using chat templates. Check we get a single BOS token for llama chat. diff --git a/tests/entrypoints/llm/test_generate.py b/tests/entrypoints/llm/test_generate.py index 4a2c59dc5842..ccff4eb7a0f9 100644 --- a/tests/entrypoints/llm/test_generate.py +++ b/tests/entrypoints/llm/test_generate.py @@ -75,6 +75,10 @@ def test_multiple_priority(llm: LLM): outputs = llm.generate(PROMPTS, sampling_params=None, priority=[0] * len(PROMPTS)) assert len(PROMPTS) == len(outputs) + # Generate works when priority is a single value applied to every prompt + outputs = llm.generate(PROMPTS, sampling_params=None, priority=1) + assert len(PROMPTS) == len(outputs) + # Exception raised, if the length of priority does not match the length of prompts with pytest.raises(VLLMValidationError): outputs = llm.generate( @@ -91,6 +95,10 @@ def test_single_prompt_priority(llm: LLM): outputs = llm.generate(PROMPTS[0], sampling_params=None, priority=[0]) assert len(outputs) == 1 + # Single string prompts should also accept a single priority value. + outputs = llm.generate(PROMPTS[0], sampling_params=None, priority=1) + assert len(outputs) == 1 + def test_max_model_len(vllm_runner): max_model_len = 20 diff --git a/tests/entrypoints/unit_tests/test_offline_utils.py b/tests/entrypoints/unit_tests/test_offline_utils.py index a37969c9b366..8a0308ce3a46 100644 --- a/tests/entrypoints/unit_tests/test_offline_utils.py +++ b/tests/entrypoints/unit_tests/test_offline_utils.py @@ -55,3 +55,9 @@ def test_matching_lengths_pass_through(mixin: OfflineInferenceMixin): ] assert mixin._lora_request_to_seq([lora], num_requests=1) == [lora] assert mixin._priority_to_seq([3], num_requests=1) == [3] + + +def test_scalar_priority_broadcast(mixin: OfflineInferenceMixin): + assert mixin._priority_to_seq(5, num_requests=1) == [5] + assert mixin._priority_to_seq(5, num_requests=3) == [5, 5, 5] + assert mixin._priority_to_seq(None, num_requests=3) == [0, 0, 0] diff --git a/vllm/entrypoints/common/offline.py b/vllm/entrypoints/common/offline.py index b33b8c34c5ad..1bb0bedccbd0 100644 --- a/vllm/entrypoints/common/offline.py +++ b/vllm/entrypoints/common/offline.py @@ -281,10 +281,13 @@ def _lora_request_to_seq( def _priority_to_seq( self, - priority: list[int] | None, + priority: int | Sequence[int] | None, num_requests: int, ) -> Sequence[int]: - if priority is not None: + if priority is None: + priority = 0 + + if isinstance(priority, Sequence): if len(priority) != num_requests: raise VLLMValidationError( f"The lengths of prompts ({num_requests}) " @@ -293,7 +296,7 @@ def _priority_to_seq( return priority - return [0] * num_requests + return [priority] * num_requests def _add_completion_requests( self, @@ -304,7 +307,7 @@ def _add_completion_requests( *, use_tqdm: bool | Callable[..., tqdm] = True, lora_request: Sequence[LoRARequest] | LoRARequest | None = None, - priority: list[int] | None = None, + priority: int | Sequence[int] | None = None, tokenization_kwargs: dict[str, Any] | None = None, mm_processor_kwargs: dict[str, Any] | None = None, ) -> list[str]: @@ -341,7 +344,7 @@ def _run_completion( *, use_tqdm: bool | Callable[..., tqdm] = True, lora_request: Sequence[LoRARequest] | LoRARequest | None = None, - priority: list[int] | None = None, + priority: int | Sequence[int] | None = None, tokenization_kwargs: dict[str, Any] | None = None, mm_processor_kwargs: dict[str, Any] | None = None, ): @@ -367,6 +370,7 @@ def _run_chat( *, use_tqdm: bool | Callable[..., tqdm] = True, lora_request: Sequence[LoRARequest] | LoRARequest | None = None, + priority: int | Sequence[int] | None = None, chat_template: str | None = None, chat_template_content_format: ChatTemplateContentFormatOption = "auto", add_generation_prompt: bool = True, @@ -381,6 +385,7 @@ def _run_chat( params=params, use_tqdm=use_tqdm, lora_request=lora_request, + priority=priority, chat_template=chat_template, chat_template_content_format=chat_template_content_format, chat_template_kwargs=chat_template_kwargs, @@ -402,7 +407,7 @@ def _add_chat_requests( *, use_tqdm: bool | Callable[..., tqdm] = True, lora_request: Sequence[LoRARequest] | LoRARequest | None = None, - priority: list[int] | None = None, + priority: int | Sequence[int] | None = None, chat_template: str | None = None, chat_template_content_format: ChatTemplateContentFormatOption = "auto", add_generation_prompt: bool = True, diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index ff7bd9c57402..6b7f91d08e4c 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -422,7 +422,7 @@ def generate( *, use_tqdm: bool | Callable[..., tqdm] = True, lora_request: Sequence[LoRARequest] | LoRARequest | None = None, - priority: list[int] | None = None, + priority: int | Sequence[int] | None = None, tokenization_kwargs: dict[str, Any] | None = None, mm_processor_kwargs: dict[str, Any] | None = None, ) -> list[RequestOutput]: @@ -448,9 +448,9 @@ def generate( lora_request: LoRA request to use for generation, if any. priority: The priority of the requests, if any. Only applicable when priority scheduling policy is enabled. - If provided, must be a list of integers matching the length - of `prompts`, where each priority value corresponds to the prompt - at the same index. + When it is a single value, it is applied to every prompt. + When it is a list, the list must have the same length as the + prompts and it is paired one by one with the prompt. tokenization_kwargs: Overrides for `tokenizer.encode`. mm_processor_kwargs: Overrides for `processor.__call__`. @@ -486,7 +486,7 @@ def enqueue( prompts: PromptType | Sequence[PromptType], sampling_params: SamplingParams | Sequence[SamplingParams] | None = None, lora_request: Sequence[LoRARequest] | LoRARequest | None = None, - priority: list[int] | None = None, + priority: int | Sequence[int] | None = None, use_tqdm: bool | Callable[..., tqdm] = True, tokenization_kwargs: dict[str, Any] | None = None, mm_processor_kwargs: dict[str, Any] | None = None, @@ -502,6 +502,10 @@ def enqueue( sampling_params: The sampling parameters for text generation. lora_request: LoRA request to use for generation, if any. priority: The priority of the requests, if any. + Only applicable when priority scheduling policy is enabled. + When it is a single value, it is applied to every prompt. + When it is a list, the list must have the same length as the + prompts and it is paired one by one with the prompt. use_tqdm: If True, shows a tqdm progress bar while adding requests. tokenization_kwargs: Overrides for `tokenizer.encode`. mm_processor_kwargs: Overrides for `processor.__call__`. @@ -617,6 +621,7 @@ def chat( sampling_params: SamplingParams | Sequence[SamplingParams] | None = None, use_tqdm: bool | Callable[..., tqdm] = True, lora_request: Sequence[LoRARequest] | LoRARequest | None = None, + priority: int | Sequence[int] | None = None, chat_template: str | None = None, chat_template_content_format: ChatTemplateContentFormatOption = "auto", add_generation_prompt: bool = True, @@ -651,6 +656,11 @@ def chat( it is used to create the progress bar. If `False`, no progress bar is created. lora_request: LoRA request to use for generation, if any. + priority: The priority of the requests, if any. + Only applicable when priority scheduling policy is enabled. + When it is a single value, it is applied to every prompt. + When it is a list, the list must have the same length as the + prompts and it is paired one by one with the prompt. chat_template: The template to use for structuring the chat. If not provided, the model's default chat template will be used. chat_template_content_format: The format to render message content. @@ -695,6 +705,7 @@ def chat( output_type=RequestOutput, use_tqdm=use_tqdm, lora_request=lora_request, + priority=priority, chat_template=chat_template, chat_template_content_format=chat_template_content_format, chat_template_kwargs=chat_template_kwargs, @@ -712,7 +723,7 @@ def enqueue_chat( sampling_params: SamplingParams | Sequence[SamplingParams] | None = None, use_tqdm: bool | Callable[..., tqdm] = True, lora_request: Sequence[LoRARequest] | LoRARequest | None = None, - priority: list[int] | None = None, + priority: int | Sequence[int] | None = None, chat_template: str | None = None, chat_template_content_format: ChatTemplateContentFormatOption = "auto", add_generation_prompt: bool = True, @@ -739,6 +750,10 @@ def enqueue_chat( conversations. lora_request: LoRA request to use for generation, if any. priority: The priority of the requests, if any. + Only applicable when priority scheduling policy is enabled. + When it is a single value, it is applied to every prompt. + When it is a list, the list must have the same length as the + prompts and it is paired one by one with the prompt. chat_template: The template to use for structuring the chat. chat_template_content_format: The format to render message content. add_generation_prompt: If True, adds a generation template From c85616efc6f0eaa50ca1f1a09c32368aedffd3d8 Mon Sep 17 00:00:00 2001 From: tc-mb <157115220+tc-mb@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:21:10 +0800 Subject: [PATCH 259/624] [Model] Support MiniCPM-V 4.7 (#58674) Signed-off-by: tc-mb Co-authored-by: Cyrus Leung --- docs/models/supported_models.md | 2 +- tests/model_executor/test_minicpmv.py | 114 +++ tests/model_executor/test_minicpmv4_7.py | 310 +++++++ tests/models/registry.py | 5 + .../parser/engine/test_minicpmv_reasoning.py | 459 +++++++++++ vllm/model_executor/models/minicpmv4_6.py | 335 ++++++-- vllm/model_executor/models/minicpmv4_7.py | 764 ++++++++++++++++++ vllm/model_executor/models/registry.py | 4 + vllm/parser/minicpmv.py | 405 ++++++++++ vllm/parser/parser_manager.py | 13 + vllm/reasoning/__init__.py | 4 + vllm/reasoning/minicpmv_reasoning_parser.py | 16 + .../chat_templates/registry.py | 1 + 13 files changed, 2368 insertions(+), 64 deletions(-) create mode 100644 tests/model_executor/test_minicpmv4_7.py create mode 100644 tests/parser/engine/test_minicpmv_reasoning.py create mode 100644 vllm/model_executor/models/minicpmv4_7.py create mode 100644 vllm/parser/minicpmv.py create mode 100644 vllm/reasoning/minicpmv_reasoning_parser.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 221ab877eaa8..bbcb94f32cf5 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -592,7 +592,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `MiDashengLMModel` | MiDashengLM | T + A+ | `mispeech/midashenglm-7b` | | ✅︎ | | `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + IE+ + VE+ + A+ | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ | | `MiniCPMO` | MiniCPM-O | T + IE+ + VE+ + AE+ | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ | -| `MiniCPMV` | MiniCPM-V | T + IE+ + VE+ | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, `openbmb/MiniCPM-V-4_6`, etc. | ✅︎ | | +| `MiniCPMV` | MiniCPM-V | T + IE+ + VE+ | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, `openbmb/MiniCPM-V-4_6`, `openbmb/MiniCPM-V-4_7`, etc. | ✅︎ | | | `MiniMaxM3SparseForConditionalGeneration` | MiniMax-M3 | T + I+ + V+ | `MiniMaxAI/MiniMax-M3`, `MiniMaxAI/MiniMax-M3-MXFP8`, etc. | | ✅︎ | | `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + IE+ | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ | | `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I+ | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ | diff --git a/tests/model_executor/test_minicpmv.py b/tests/model_executor/test_minicpmv.py index 96b80101b414..cffdea260097 100644 --- a/tests/model_executor/test_minicpmv.py +++ b/tests/model_executor/test_minicpmv.py @@ -6,6 +6,9 @@ from vllm.model_executor.models.minicpmv import MiniCPMV2_6 from vllm.model_executor.models.minicpmv4_6 import ( MiniCPMV4_6ForConditionalGeneration, + _expand_flags, + _select_vision_items, + _vision_downsample_mode, ) pytestmark = pytest.mark.skip_global_cleanup @@ -54,3 +57,114 @@ def test_image_and_video_embeds_stay_in_their_own_modality(): assert torch.equal(modalities["images"]["image_embeds"], image_embeds) assert torch.equal(modalities["videos"]["image_embeds"], video_embeds) + + +class TestPerItemMergerFlags: + """The merger flag is per request but items from several requests can be + batched together, so it has to be applied per item.""" + + def test_expand_flags_broadcasts_a_single_value(self): + assert _expand_flags(None, 3) == [None, None, None] + assert _expand_flags(torch.tensor([True]), 3) == [True, True, True] + assert _expand_flags(True, 2) == [True, True] + + def test_expand_flags_keeps_per_item_values(self): + flags = [torch.tensor([True]), torch.tensor([False])] + + assert _expand_flags(flags, 2) == [True, False] + + def test_expand_flags_falls_back_on_a_shape_mismatch(self): + # Too many values to map onto the items: keep one decision for the + # batch rather than failing the request. + assert _expand_flags([True, False, True], 2) == [True, True] + + def test_vision_downsample_mode(self): + assert _vision_downsample_mode(None) is None + assert _vision_downsample_mode(True) == "16x" + assert _vision_downsample_mode(False) == "4x" + + def test_select_vision_items_slices_by_item(self): + # 3 items holding 2, 1 and 2 slices; one channel, one patch, width 4. + pixel_values = [torch.full((1, 1, 4), float(i)) for i in range(5)] + tgt_sizes = torch.arange(10, dtype=torch.long).reshape(5, 2) + image_input = { + "pixel_values": pixel_values, + "tgt_sizes": tgt_sizes, + "num_slices": torch.tensor([2, 1, 2]), + } + + selected = _select_vision_items(image_input, [2], [2, 1, 2]) + + assert len(selected["pixel_values"]) == 2 + assert bool(selected["pixel_values"][0][0, 0, 0] == 3.0) + assert bool(selected["pixel_values"][1][0, 0, 0] == 4.0) + assert selected["num_slices"].tolist() == [2] + assert torch.equal(selected["tgt_sizes"], tgt_sizes[3:5]) + + def test_select_vision_items_preserves_order(self): + pixel_values = [torch.full((1, 1, 4), float(i)) for i in range(3)] + image_input = { + "pixel_values": pixel_values, + "tgt_sizes": torch.zeros(3, 2, dtype=torch.long), + "num_slices": torch.tensor([1, 1, 1]), + } + + selected = _select_vision_items(image_input, [2, 0], [1, 1, 1]) + + assert selected["num_slices"].tolist() == [1, 1] + assert bool(selected["pixel_values"][0][0, 0, 0] == 2.0) + assert bool(selected["pixel_values"][1][0, 0, 0] == 0.0) + + def test_process_vision_input_applies_flags_per_item(self): + calls: list[str | None] = [] + + def fake_get_vision_hidden_states(data, downsample_mode=None): + calls.append(downsample_mode) + tag = 1.0 if downsample_mode == "16x" else 0.0 + return [torch.full((1,), tag) for _ in data["pixel_values"]] + + model = object.__new__(MiniCPMV4_6ForConditionalGeneration) + model.get_vision_hidden_states = fake_get_vision_hidden_states + + image_input = { + "type": "pixel_values", + "pixel_values": [torch.zeros(1, 1, 4) for _ in range(3)], + "tgt_sizes": torch.zeros(3, 2, dtype=torch.long), + "num_slices": torch.tensor([1, 1, 1]), + } + + result = model._process_vision_input( + image_input, + use_vit_merger=[ + torch.tensor([True]), + torch.tensor([False]), + torch.tensor([True]), + ], + ) + + # Each item keeps its own mode even though they share one call. + assert [float(r[0]) for r in result] == [1.0, 0.0, 1.0] + # Mixed flags are encoded as separate groups, not one collapsed batch. + assert calls == ["16x", "4x"] + + def test_process_vision_input_uses_one_group_for_uniform_flags(self): + calls: list[str | None] = [] + + def fake_get_vision_hidden_states(data, downsample_mode=None): + calls.append(downsample_mode) + return [torch.zeros(1) for _ in data["pixel_values"]] + + model = object.__new__(MiniCPMV4_6ForConditionalGeneration) + model.get_vision_hidden_states = fake_get_vision_hidden_states + + image_input = { + "type": "pixel_values", + "pixel_values": [torch.zeros(1, 1, 4) for _ in range(2)], + "tgt_sizes": torch.zeros(2, 2, dtype=torch.long), + "num_slices": torch.tensor([1, 1]), + } + + model._process_vision_input(image_input, use_vit_merger=torch.tensor([True])) + + # The common case still makes a single encoder call. + assert calls == ["16x"] diff --git a/tests/model_executor/test_minicpmv4_7.py b/tests/model_executor/test_minicpmv4_7.py new file mode 100644 index 000000000000..3d9e8b467830 --- /dev/null +++ b/tests/model_executor/test_minicpmv4_7.py @@ -0,0 +1,310 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import torch + +from vllm.model_executor.models.minicpmv4_6 import _stack_vit_merger_qkv +from vllm.model_executor.models.minicpmv4_7 import ( + MiniCPMV4_7MultiModalProcessor, + _compute_canvas_single, + build_image_bounds, + canvas_rope_delta, +) + +pytestmark = pytest.mark.skip_global_cleanup + + +def test_canvas_mrope_sizes_survive_prefix_cache_stripping(): + """A cached image must keep `tgt_sizes`, or canvas M-RoPE falls back. + + `strip_covered_mm_data` drops the payload of a prefix-cache-covered item and + keeps only CPU-side metadata. Canvas M-RoPE places slices from `tgt_sizes`, + so 4.7 marks it `keep_on_cpu` while the shared 4.6 config stays as it was. + """ + from vllm.model_executor.models.minicpmv4_6 import _minicpmv4_6_field_config + + shared = _minicpmv4_6_field_config({}) + assert shared["tgt_sizes"].field.keep_on_cpu is False + assert shared["video_tgt_sizes"].field.keep_on_cpu is False + + processor = MiniCPMV4_7MultiModalProcessor.__new__(MiniCPMV4_7MultiModalProcessor) + fields = processor._get_mm_fields_config({}, {}) + + assert fields["tgt_sizes"].field.keep_on_cpu is True + assert fields["video_tgt_sizes"].field.keep_on_cpu is True + # Other fields keep their previous setting. + assert fields["pixel_values"].field.keep_on_cpu is False + assert fields["video_image_sizes"].field.keep_on_cpu is True + + +def test_modality_mm_kwargs_resolution(): + """Flat keys reach both modalities; scopes stay scoped, and win on conflict. + + `_merge_and_resolve_mm_processor_kwargs` routes a flat key into every scope + whose HF processor declares it, so reading is a plain scope lookup. The + routing itself is upstream's, tested here against the MiniCPM-V schema to + pin the precedence rule (`scoped > flat`). + """ + from vllm.model_executor.models.minicpmv4_6 import ( + _flat_processor_kwargs, + _scoped_mm_kwarg, + ) + from vllm.multimodal.processing.context import _resolve_mm_processor_kwargs + + # What MiniCPMV4_6{Image,Video}ProcessorKwargs declare. + processor_keys = {"downsample_mode", "max_slice_nums"} + schema = { + "text_kwargs": set(), + "images_kwargs": set(processor_keys), + "videos_kwargs": set(processor_keys), + "audio_kwargs": set(), + } + + def read(mm_kwargs, modality, key): + merged = _resolve_mm_processor_kwargs(dict(mm_kwargs), schema) + return _scoped_mm_kwarg(merged, modality, key) + + assert read({}, "image", "downsample_mode") is None + # A flat key is routed into both scopes. + assert read({"downsample_mode": "4x"}, "image", "downsample_mode") == "4x" + assert read({"downsample_mode": "4x"}, "video", "downsample_mode") == "4x" + # A scoped key stays in its scope. + assert ( + read({"images_kwargs": {"downsample_mode": "16x"}}, "image", "downsample_mode") + == "16x" + ) + assert ( + read({"images_kwargs": {"downsample_mode": "16x"}}, "video", "downsample_mode") + is None + ) + videos_only = {"videos_kwargs": {"max_slice_nums": 1}} + assert read(videos_only, "video", "max_slice_nums") == 1 + # Scoped to videos, so image must not see it. + assert read(videos_only, "image", "max_slice_nums") is None + # An explicit scope wins over the flat key it is merged with. + assert ( + read( + {"downsample_mode": "16x", "images_kwargs": {"downsample_mode": "4x"}}, + "image", + "downsample_mode", + ) + == "4x" + ) + # A non-mapping scope is ignored rather than raising. + assert read({"images_kwargs": None}, "image", "downsample_mode") is None + + # `_flat_processor_kwargs` drops the scoped keys and pins the mode. + flat = _flat_processor_kwargs( + { + "downsample_mode": "4x", + "images_kwargs": {"a": 1}, + "videos_kwargs": {"b": 2}, + "audio_kwargs": {"c": 3}, + }, + "16x", + ) + assert flat == {"downsample_mode": "16x"} + + +def test_vit_merger_qkv_is_stacked_before_load(): + mapped = list( + _stack_vit_merger_qkv( + [ + ( + "model.vision_tower.vit_merger.self_attn.q_proj.weight", + torch.zeros(1), + ), + ( + "model.vision_tower.vit_merger.self_attn.k_proj.weight", + torch.zeros(1), + ), + ( + "model.vision_tower.vit_merger.self_attn.v_proj.weight", + torch.zeros(1), + ), + ("vit_merger.self_attn.qkv_proj.weight", torch.zeros(1)), + ] + ) + ) + names = [name for name, _ in mapped] + shards = [getattr(tensor, "shard_id", None) for _, tensor in mapped] + assert names == [ + "model.vision_tower.vit_merger.self_attn.qkv_proj.weight", + "model.vision_tower.vit_merger.self_attn.qkv_proj.weight", + "model.vision_tower.vit_merger.self_attn.qkv_proj.weight", + "vit_merger.self_attn.qkv_proj.weight", + ] + assert shards == ["q", "k", "v", None] + + +# Canvas M-RoPE goldens, cross-checked against the reference implementation in +# transformers PR 48979 (`modular_minicpmv4_7.py`). Synthetic ids keep the token +# sequences readable: the canvas math only compares ids for equality. +_IM_START = 10 +_IM_END = 11 +_SLICE_START = 12 +_SLICE_END = 13 +_NEWLINE = 14 +_IMAGE_PAD = 20 +_VIDEO_PAD = 21 +_TEXT = 100 + +_CANVAS_IDS = { + "im_start_id": _IM_START, + "im_end_id": _IM_END, + "slice_start_id": _SLICE_START, + "slice_end_id": _SLICE_END, + "newline_id": _NEWLINE, +} + + +def _canvas_positions(tokens, target_sizes): + input_ids = torch.tensor([tokens], dtype=torch.long) + targets = torch.tensor(target_sizes, dtype=torch.long) + bounds = build_image_bounds(input_ids[0], _CANVAS_IDS) + positions = _compute_canvas_single( + input_ids[0], + torch.arange(len(tokens), dtype=torch.long), + bounds, + targets, + _CANVAS_IDS, + ) + return positions.tolist(), canvas_rope_delta(positions, len(tokens)) + + +def test_canvas_mrope_trailing_newline_stays_inside_the_canvas(): + """A newline right after `` belongs to the canvas, not to the text. + + The chat template joins content parts with ``\\n``, so an image followed by + text always leaves one there. The transformers reference closes a frame over + every marker including newlines; leaving it out pushes every later position + one step forward and makes `rope_delta` differ by one. + """ + tokens = [_TEXT, _IM_START, _IMAGE_PAD, _IM_END, _NEWLINE, _TEXT, _TEXT] + positions, delta = _canvas_positions(tokens, [[4, 4]]) + + assert positions == [ + [0, 1, 1, 1, 1, 3, 4], + [0, 0, 1, 2, 1, 3, 4], + [0, 0, 1, 2, 1, 3, 4], + ] + assert delta == -2 + + +def test_canvas_mrope_thumbnail_matches_reference(): + # PAD ; 16x turns a 4x4 patch grid into 1 token. + tokens = [_TEXT, _IM_START, _IMAGE_PAD, _IM_END, _TEXT] + positions, delta = _canvas_positions(tokens, [[4, 4]]) + + assert positions == [ + [0, 1, 1, 1, 3], + [0, 0, 1, 2, 3], + [0, 0, 1, 2, 3], + ] + assert delta == -1 + + +def test_canvas_mrope_slice_row_matches_reference(): + # Thumbnail then a 1x2 slice row; the slices tile the thumbnail's canvas. + tokens = [ + _TEXT, + _IM_START, + _IMAGE_PAD, + _IM_END, + _SLICE_START, + _IMAGE_PAD, + _SLICE_END, + _SLICE_START, + _IMAGE_PAD, + _SLICE_END, + _TEXT, + ] + positions, delta = _canvas_positions(tokens, [[4, 4], [4, 4], [4, 4]]) + + assert positions == [ + [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 4], + [0, 0, 1, 2, 1, 1, 1, 1, 1, 1, 4], + [0, 0, 1, 3, 1, 1, 1, 2, 2, 2, 4], + ] + assert delta == -6 + + +def test_canvas_mrope_resamples_thumbnail_over_canvas(): + # A 8x8 patch grid is 2x2 LLM tokens, so the canvas is resampled (not arange). + tokens = [_TEXT, _IM_START] + [_IMAGE_PAD] * 4 + [_IM_END, _TEXT] + positions, delta = _canvas_positions(tokens, [[8, 8]]) + + assert positions == [ + [0, 1, 1, 1, 1, 1, 1, 4], + [0, 0, 1, 1, 2, 2, 3, 4], + [0, 0, 1, 2, 1, 2, 3, 4], + ] + assert delta == -3 + + +def test_canvas_mrope_video_frames_share_one_canvas(): + tokens = [ + _TEXT, + _IM_START, + _VIDEO_PAD, + _IM_END, + _IM_START, + _VIDEO_PAD, + _IM_END, + _TEXT, + ] + positions, delta = _canvas_positions(tokens, [[4, 4], [4, 4]]) + + assert positions == [ + [0, 1, 1, 1, 3, 3, 3, 5], + [0, 0, 1, 2, 2, 3, 4, 5], + [0, 0, 1, 2, 2, 3, 4, 5], + ] + assert delta == -2 + + +def test_canvas_mrope_interleaved_image_and_video(): + tokens = ( + [_TEXT, _IM_START] + + [_IMAGE_PAD] * 4 + + [_IM_END, _IM_START] + + [_VIDEO_PAD] * 4 + + [_IM_END, _TEXT] + ) + positions, delta = _canvas_positions(tokens, [[8, 8], [8, 8]]) + + assert positions == [ + [0, 1, 1, 1, 1, 1, 1, 4, 4, 4, 4, 4, 4, 7], + [0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 7], + [0, 0, 1, 2, 1, 2, 3, 3, 4, 5, 4, 5, 6, 7], + ] + assert delta == -6 + + +def test_canvas_span_is_shared_by_image_and_video(): + """A video frame and a same-geometry image must place identical offsets. + + Both paths go through `_assign_canvas_span`; this pins that they cannot + drift apart now that the duplicated geometry is gone. + """ + thumb = [_IM_START] + [_IMAGE_PAD] * 4 + [_IM_END] + + img_pos, _ = _canvas_positions([_TEXT, *thumb, _TEXT], [[8, 8]]) + vid_pos, _ = _canvas_positions([_TEXT, *thumb, *thumb, _TEXT], [[8, 8], [8, 8]]) + + def offsets(pos, start): + # The halo is one below the canvas base. + base = pos[1][start] + 1 + return [[c[start + k] - base for k in range(3)] for c in pos] + + assert offsets(img_pos, 0) == offsets(vid_pos, 0) + + +def test_canvas_mrope_delta_leaves_room_for_the_first_decoded_token(): + # First decoded token must land on max_prefill_position + 1, matching + # transformers and the other native M-RoPE models. + tokens = [_TEXT, _IM_START, _IMAGE_PAD, _IM_END, _TEXT] + positions, delta = _canvas_positions(tokens, [[4, 4]]) + + max_position = max(max(channel) for channel in positions) + assert delta == max_position + 1 - len(tokens) diff --git a/tests/models/registry.py b/tests/models/registry.py index b57f3fcb7277..2ed6d4871c27 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1118,6 +1118,11 @@ def check_available_online( "MiniCPMV4_6ForConditionalGeneration": _HfExamplesInfo( "openbmb/MiniCPM-V-4_6", ), + "MiniCPMV4_7ForConditionalGeneration": _HfExamplesInfo( + "openbmb/MiniCPM-V-4_7", + min_transformers_version="5.18.0", + is_available_online=False, + ), "MiniMaxM3SparseForConditionalGeneration": _HfExamplesInfo( "MiniMaxAI/MiniMax-M3", trust_remote_code=True, diff --git a/tests/parser/engine/test_minicpmv_reasoning.py b/tests/parser/engine/test_minicpmv_reasoning.py new file mode 100644 index 000000000000..5edf5bfc5da9 --- /dev/null +++ b/tests/parser/engine/test_minicpmv_reasoning.py @@ -0,0 +1,459 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the MiniCPM-V family parser.""" + +import pytest + +from tests.parser.engine.conftest import make_mock_tokenizer +from tests.parser.engine.streaming_helpers import simulate_reasoning_streaming +from vllm.entrypoints.generate.base.protocol import ( + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, +) +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ChatCompletionToolsParam, +) +from vllm.parser.minicpmv import ( + EscapedNewlineNormalizer, + MiniCPMVOutputNormalizer, + recover_newlines, +) +from vllm.parser.parser_manager import ParserManager +from vllm.reasoning import ReasoningParserManager +from vllm.tool_parsers import ToolParserManager + +_VOCAB = { + "": 10, + "": 11, + "<|im_end|>": 12, + "": 13, + "": 14, +} + +_TOOLS = [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, + ) +] + + +@pytest.fixture +def mock_tokenizer(): + return make_mock_tokenizer(_VOCAB) + + +@pytest.fixture +def request_obj(): + return ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + ) + + +def _make_parser(tokenizer, *, enable_thinking=None, tools=None): + """Build the parser the serving stack resolves for MiniCPM-V. + + Everything is built through ``ParserManager`` so the tests exercise the + ``DelegatingParser`` the server actually drives, not the parser engine on + its own. Tool calling goes through ``qwen3_coder``: MiniCPM-V emits that + XML syntax verbatim and has no tool parser of its own. + """ + parser_cls = ParserManager.get_parser( + tool_parser_name="qwen3_coder" if tools else None, + reasoning_parser_name="minicpmv", + enable_auto_tools=bool(tools), + ) + assert parser_cls is not None + kwargs = {} + if enable_thinking is not None: + kwargs["chat_template_kwargs"] = {"enable_thinking": enable_thinking} + return parser_cls(tokenizer, tools, **kwargs) + + +def test_reasoning_parser_is_registered(): + parser_cls = ReasoningParserManager.get_reasoning_parser("minicpmv") + + assert parser_cls.__name__ == "MiniCPMVParserReasoningAdapter" + + +def test_thinking_defaults_to_the_chat_template(mock_tokenizer): + """Without `enable_thinking`, the grammar must wait for reasoning. + + The chat template ends the assistant turn with `` unless the flag is + explicitly false, so the model reasons by default. This is the flag + `serving.py` latches into the structured-output gate, so a non-thinking + default both disables the grammar and leaks the thinking into `content`. + """ + + def reasoning_pending(**kwargs): + return not _make_parser(mock_tokenizer, **kwargs).is_reasoning_end([1, 2]) + + assert reasoning_pending() is True # default matches the template + assert reasoning_pending(enable_thinking=True) is True + assert reasoning_pending(enable_thinking=False) is False + + +def test_no_minicpmv_tool_parser_is_registered(): + """MiniCPM-V deliberately registers no tool parser. + + The family emits the Qwen3-Coder XML syntax verbatim, so + ``--tool-call-parser qwen3_coder`` already covers it. Registering a + ``minicpmv`` tool parser would duplicate that name for no parsing gain. + """ + with pytest.raises(KeyError): + ToolParserManager.get_tool_parser("minicpmv") + + +def test_composes_with_qwen3_coder_tool_parser(mock_tokenizer): + """`--tool-call-parser qwen3_coder` must work through the serving Parser.""" + request = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + tools=_TOOLS, + ) + parser_cls = ParserManager.get_parser( + tool_parser_name="qwen3_coder", + reasoning_parser_name="minicpmv", + enable_auto_tools=True, + ) + assert parser_cls is not None + parser = parser_cls(mock_tokenizer, request.tools) + + _, _, tool_calls = parser.parse( + "" + "SF", + request, + enable_auto_tools=True, + ) + + assert tool_calls is not None + assert tool_calls[0].name == "get_weather" + assert "SF" in tool_calls[0].arguments + + +class TestNonThinking: + @pytest.fixture + def parser(self, mock_tokenizer): + return _make_parser(mock_tokenizer, enable_thinking=False) + + def test_plain_content_is_unchanged(self, parser): + reasoning, content = parser.extract_reasoning("The answer is 42.", None) + + assert reasoning is None + assert content == "The answer is 42." + + def test_reserved_markers_are_removed(self, parser): + output = ( + "reasoning<|reserved_14|>tool call<|reserved_15|>" + ) + + reasoning, content = parser.extract_reasoning(output, None) + + assert reasoning is None + assert content == "reasoningtool call" + + def test_unrelated_reserved_text_is_unchanged(self, parser): + reasoning, content = parser.extract_reasoning( + "content<|reserved_16|>", + None, + ) + + assert reasoning is None + assert content == "content<|reserved_16|>" + + def test_dangling_standard_think_end_is_removed(self, parser): + output = "first part \n\nfinal part" + + reasoning, content = parser.extract_reasoning(output, None) + + assert reasoning is None + assert content == "first part \n\nfinal part" + + def test_truncated_standard_reasoning_does_not_leak(self, parser): + reasoning, content = parser.extract_reasoning( + "prefixprivate reasoning", + None, + ) + + assert reasoning is None + assert content == "prefix" + + def test_streaming_split_standard_markers(self, parser): + reasoning, content = simulate_reasoning_streaming( + parser, + [ + "<", + "think>", + "private reasoning", + "", + "final answer", + ], + ) + + assert reasoning == "" + assert content == "final answer" + + def test_streaming_split_reserved_markers(self, parser): + reasoning, content = simulate_reasoning_streaming( + parser, + [ + "", + "reasoning", + "<|reserved", + "_13|>", + "final answer", + ], + ) + + assert reasoning == "" + assert content == "reasoningfinal answer" + + def test_tool_calls_are_parsed(self, mock_tokenizer): + """Thinking is off by default, and tool calls must still be parsed.""" + parser = _make_parser(mock_tokenizer, enable_thinking=False, tools=_TOOLS) + request = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + tools=_TOOLS, + ) + + reasoning, content, tool_calls = parser.parse( + "" + "SF", + request, + enable_auto_tools=True, + ) + + assert reasoning is None + assert not content + assert tool_calls is not None + assert tool_calls[0].name == "get_weather" + assert "SF" in tool_calls[0].arguments + + def test_streaming_tool_calls_are_parsed(self, mock_tokenizer): + """Regression: tool calls must stream when thinking is explicitly off. + + The reasoning parser used to report ``reasoning_ended=False`` in + non-thinking mode, which kept the delegating parser in its reasoning + phase and streamed the raw ```` XML to the client as content. + """ + parser = _make_parser(mock_tokenizer, enable_thinking=False, tools=_TOOLS) + request = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + tools=_TOOLS, + ) + + delta = parser.parse_delta( + "" + "SF", + [], + request, + prompt_token_ids=[], + finished=True, + ) + + assert delta is not None + assert delta.tool_calls + assert delta.tool_calls[0].function is not None + assert delta.tool_calls[0].function.name == "get_weather" + assert "SF" in (delta.tool_calls[0].function.arguments or "") + assert not delta.content + + +class TestThinking: + def test_standard_think_tags_are_parsed(self, mock_tokenizer): + parser = _make_parser(mock_tokenizer, enable_thinking=True) + + reasoning, content = parser.extract_reasoning( + "private reasoningfinal answer", + None, + ) + + assert reasoning == "private reasoning" + assert content == "final answer" + + def test_reserved_markers_are_removed(self, mock_tokenizer): + parser = _make_parser(mock_tokenizer, enable_thinking=True) + + reasoning, content = parser.extract_reasoning( + "private reasoning<|reserved_13|>final answer", + None, + ) + + assert reasoning == "private reasoning" + assert content == "final answer" + + def test_streaming_im_end_is_removed(self): + tokenizer = make_mock_tokenizer(_VOCAB, special_tokens=[]) + parser = _make_parser(tokenizer, enable_thinking=True) + + reasoning, content = simulate_reasoning_streaming( + parser, + ["Answer: 111", "<|im_end|>"], + ) + + assert reasoning == "Answer: 111" + assert content == "" + + +class TestNewlineRecovery: + """The tokenizer writes a newline as the two characters ``\\`` and ``n``.""" + + def test_recover_newlines(self): + text = ( + "first\\nsecond\\rthird\\r\\nfourth " + "`code\\nvalue` " + "```python\\nvalue = '\\\\n'\\n``` " + "$x\\ny$ $$x\\ny$$ \\(x\\ny\\) \\[x\\ny\\] " + "escaped \\\\n" + ) + + assert recover_newlines(text) == ( + "first\nsecond\nthird\nfourth " + "`code\\nvalue` " + "```python\\nvalue = '\\\\n'\\n``` " + "$x\\ny$ $$x\\ny$$ \\(x\\ny\\) \\[x\\ny\\] " + "escaped \\\\n" + ) + + def test_streaming_matches_complete_recovery(self): + text = ( + "first\\nsecond `code\\nvalue` $x\\ny$ ```text\\nvalue\\n``` last\\r\\nline" + ) + expected = recover_newlines(text) + + for split in range(len(text) + 1): + normalizer = EscapedNewlineNormalizer() + actual = normalizer.feed(text[:split]) + actual += normalizer.feed(text[split:], final=True) + assert actual == expected + + normalizer = EscapedNewlineNormalizer() + actual = "".join(normalizer.feed(char) for char in text) + actual += normalizer.feed("", final=True) + assert actual == expected + + def test_backslash_runs_stay_consistent_across_splits(self): + # A run of backslashes can be split over chunks. Its parity decides + # whether a following `n` is a newline, so streaming and non-streaming + # must agree for every split point. + for count in range(1, 7): + text = "\\" * count + "n" + expected = recover_newlines(text) + + for split in range(len(text) + 1): + normalizer = EscapedNewlineNormalizer() + actual = normalizer.feed(text[:split]) + actual += normalizer.feed(text[split:], final=True) + assert actual == expected, (count, split) + + def test_streaming_recovers_reasoning_and_content(self): + normalizer = MiniCPMVOutputNormalizer() + + first = normalizer.normalize_delta( + DeltaMessage(reasoning="reason\\"), + finished=False, + ) + second = normalizer.normalize_delta( + DeltaMessage(reasoning="nnext\\", content="answer\\"), + finished=False, + ) + final = normalizer.normalize_delta( + DeltaMessage(content="nend"), + finished=True, + ) + + assert first == DeltaMessage(reasoning="reason") + assert second == DeltaMessage(reasoning="\nnext\\", content="answer") + assert final == DeltaMessage(content="\nend") + + def test_streaming_does_not_modify_tool_arguments(self): + normalizer = MiniCPMVOutputNormalizer() + arguments = '{"value":"first\\\\nsecond"}' + delta = DeltaMessage( + tool_calls=[ + DeltaToolCall( + index=0, + function=DeltaFunctionCall(arguments=arguments), + ) + ] + ) + + result = normalizer.normalize_delta(delta, finished=True) + + assert result is not None + assert result.tool_calls[0].function is not None + assert result.tool_calls[0].function.arguments == arguments + + def test_parse_recovers_content(self, mock_tokenizer, request_obj): + parser = _make_parser(mock_tokenizer, enable_thinking=False) + + reasoning, content, tool_calls = parser.parse( + "first\\nsecond", + request_obj, + ) + + assert reasoning is None + assert content == "first\nsecond" + assert not tool_calls + + def test_parse_delta_recovers_content(self, mock_tokenizer, request_obj): + parser = _make_parser(mock_tokenizer, enable_thinking=False) + + chunk = parser.parse_delta( + "first\\", + [], + request_obj, + finished=False, + ) + rest = parser.parse_delta( + "nsecond", + [], + request_obj, + finished=True, + ) + + # The trailing half of the escape is held back until it is complete. + assert chunk is not None + assert chunk.content == "first" + assert rest is not None + assert rest.content == "\nsecond" + + def test_buffered_delta_is_dropped_not_sent_empty( + self, mock_tokenizer, request_obj + ): + parser = _make_parser(mock_tokenizer, enable_thinking=False) + parser.parse_delta("hello", [], request_obj, finished=False) + + # The lone backtick is held back for the next chunk. Reporting it as an + # all-None DeltaMessage would reach the client as an empty SSE chunk. + assert parser.parse_delta("`", [], request_obj, finished=False) is None + + def test_normalizer_drops_empty_delta(self): + normalizer = MiniCPMVOutputNormalizer() + + assert ( + normalizer.normalize_delta(DeltaMessage(content="`"), finished=False) + is None + ) + + def test_normalizer_flushes_buffered_content_when_finished(self): + normalizer = MiniCPMVOutputNormalizer() + normalizer.normalize_delta(DeltaMessage(content="`"), finished=False) + + final = normalizer.normalize_delta(DeltaMessage(), finished=True) + + assert final is not None + assert final.content == "`" diff --git a/vllm/model_executor/models/minicpmv4_6.py b/vllm/model_executor/models/minicpmv4_6.py index 0a96c2d8479d..333ab8eedd66 100644 --- a/vllm/model_executor/models/minicpmv4_6.py +++ b/vllm/model_executor/models/minicpmv4_6.py @@ -45,6 +45,7 @@ HasInnerState, IsHybrid, MultiModalEmbeddings, + SupportsEagle3, SupportsLoRA, SupportsMRoPE, SupportsMultiModal, @@ -87,35 +88,186 @@ def _minicpmv4_6_field_config(hf_inputs: Mapping[str, torch.Tensor]): fields["use_vit_merger"] = MultiModalFieldConfig.batched( "image", keep_on_cpu=True ) + if "video_use_vit_merger" in hf_inputs: + fields["video_use_vit_merger"] = MultiModalFieldConfig.batched( + "video", keep_on_cpu=True + ) return fields +def _expand_flags(value: object, num_items: int) -> list[bool | None]: + """Broadcast a batched per-item flag to exactly one entry per item.""" + if value is None: + return [None] * num_items + + raw: list[object] + if isinstance(value, torch.Tensor): + raw = value.reshape(-1).tolist() + elif isinstance(value, list | tuple): + raw = [] + for item in value: + if isinstance(item, torch.Tensor): + raw.extend(item.reshape(-1).tolist()) + else: + raw.append(item) + else: + raw = [value] + + if len(raw) == num_items: + return [bool(v) for v in raw] + if len(raw) == 1: + return [bool(raw[0])] * num_items + # Unexpected shape: keep one decision for the whole batch instead of + # failing the request. + return [any(bool(v) for v in raw)] * num_items + + +def _vision_downsample_mode(flag: bool | None) -> str | None: + if flag is None: + return None + return "16x" if flag else "4x" + + +def _select_vision_items( + image_input: MiniCPMVImagePixelInputs, + item_indices: list[int], + num_slices: list[int], +) -> MiniCPMVImagePixelInputs: + """Slice a vision input down to the given items, keeping their order.""" + offsets = [0] + for n in num_slices: + offsets.append(offsets[-1] + n) + + pixel_values: list[torch.Tensor] = [] + tgt_sizes: list[torch.Tensor] = [] + sizes: list[int] = [] + for i in item_indices: + start, end = offsets[i], offsets[i + 1] + pixel_values.extend(image_input["pixel_values"][start:end]) + tgt_sizes.append(image_input["tgt_sizes"][start:end]) + sizes.append(end - start) + + return MiniCPMVImagePixelInputs( + type="pixel_values", + pixel_values=pixel_values, + tgt_sizes=torch.cat(tgt_sizes, dim=0), + num_slices=torch.tensor(sizes), + ) + + +# HF-style modality-scoped keys inside `mm_processor_kwargs`. They scope a value +# to a single modality and must not be forwarded to the processors, which take +# flat arguments only. +_MODALITY_SCOPED_MM_KWARGS = ("images_kwargs", "videos_kwargs", "audio_kwargs") + + +def _flat_processor_kwargs( + mm_kwargs: Mapping[str, object], + downsample_mode: str, +) -> dict[str, object]: + """Drop the modality-scoped nesting and pin the resolved downsample mode. + + The placeholder count and the vision tower must read the same mode, so the + value resolved for this modality is passed down explicitly rather than + leaving the processor to fall back to its own default. + """ + kwargs = { + key: value + for key, value in mm_kwargs.items() + if key not in _MODALITY_SCOPED_MM_KWARGS + } + kwargs["downsample_mode"] = downsample_mode + return kwargs + + +_MODALITY_TO_MM_KWARGS_SCOPE = { + "image": "images_kwargs", + "video": "videos_kwargs", +} + + +def _scoped_mm_kwarg( + merged: Mapping[str, object], + modality: str, + key: str, +) -> Any: + """Read a processor kwarg for one modality. + + ``merged`` is what ``_merge_and_resolve_mm_processor_kwargs`` returns: it + routes a flat key into every scope whose processor declares it, and a value + scoped to the other modality stays there, so it is not visible here. + """ + scoped = merged.get(_MODALITY_TO_MM_KWARGS_SCOPE[modality]) + if isinstance(scoped, Mapping): + return scoped.get(key) + return None + + class MiniCPMV4_6MultiModalProcessor(MiniCPMVMultiModalProcessor): def _resolve_downsample_mode( self, mm_kwargs: Mapping[str, object], + modality: str = "image", ) -> str: - ds = mm_kwargs.get("downsample_mode") - if ds is not None: - return str(ds) + # A flat `downsample_mode` applies to both modalities; the nested + # `images_kwargs` / `videos_kwargs` forms scope it to one, matching + # MiniCPMV4_6Processor. info = self.info assert isinstance(info, MiniCPMV4_6ProcessingInfo) + merged = info._merge_and_resolve_mm_processor_kwargs(mm_kwargs) + downsample_mode = _scoped_mm_kwarg(merged, modality, "downsample_mode") + if downsample_mode is not None: + return str(downsample_mode) return info._get_downsample_mode() + def _resolve_max_slice_nums( + self, + mm_kwargs: Mapping[str, object], + modality: str = "image", + ) -> int | None: + # Per-request max_slice_nums. None keeps the processor default. + # Resolved like `downsample_mode` so the nested per-modality form + # (`images_kwargs` / `videos_kwargs`) is honoured too. This must match + # what `process_images` passes to the image processor: the placeholder + # count and the visual token count are derived from it separately, so + # they diverge (and the engine aborts) if the two sides read different + # values. + info = self.info + assert isinstance(info, MiniCPMV4_6ProcessingInfo) + merged = info._merge_and_resolve_mm_processor_kwargs(mm_kwargs) + max_slice = _scoped_mm_kwarg(merged, modality, "max_slice_nums") + if max_slice is None: + return None + return int(max_slice) + def get_image_prompt_texts( self, image_size, image_idx: int = 0, downsample_mode: str | None = None, + max_slice_nums: int | None = None, ) -> str: info = self.info assert isinstance(info, MiniCPMV4_6ProcessingInfo) return info.get_slice_image_placeholder( image_size, image_idx=image_idx, + max_slice_nums=max_slice_nums, downsample_mode=downsample_mode, ) + def _video_local_id_prefix(self, video_idx: int) -> str: + """`{idx}` prefix emitted in front of a video placeholder. + + MiniCPM-V 4.6 numbers each video with a local id, matching + MiniCPMV4_6Processor. 4.7 treats a video as one temporal sequence and is + trained without these ids, so it overrides this to return an empty string. + """ + tokenizer = self.info.get_tokenizer() + id_start = getattr(tokenizer, "image_id_start_token", "") + id_end = getattr(tokenizer, "image_id_end_token", "") + return f"{id_start}{video_idx}{id_end}" + def get_video_prompt_texts( self, image_size, @@ -144,8 +296,6 @@ def get_video_prompt_texts( image_end = getattr(tokenizer, "image_end_token", "") slice_start = getattr(tokenizer, "slice_start_token", "") slice_end = getattr(tokenizer, "slice_end_token", "") - id_start = getattr(tokenizer, "image_id_start_token", "") - id_end = getattr(tokenizer, "image_id_end_token", "") per_frame = image_start + video_token * source_tokens + image_end if grids[0] > 0 and grids[1] > 0 and patch_tokens > 0: @@ -154,7 +304,7 @@ def get_video_prompt_texts( per_frame += "\n".join(rows) body = per_frame * num_frames - return f"{id_start}{video_idx}{id_end}" + body + return self._video_local_id_prefix(video_idx) + body def process_images( self, @@ -179,10 +329,18 @@ def process_images( n_images = len(parsed_images) image_processor = self.info.get_image_processor() patch_size = image_processor.patch_size + ds_mode = self._resolve_downsample_mode(mm_kwargs, modality="image") + image_proc_kwargs = _flat_processor_kwargs(mm_kwargs, ds_mode) + # `_flat_processor_kwargs` drops the nested form, so the value resolved + # from it has to be pinned back for the processor and the prompt text + # to agree on the slice count. + image_max_slice = self._resolve_max_slice_nums(mm_kwargs, modality="image") + if image_max_slice is not None: + image_proc_kwargs["max_slice_nums"] = image_max_slice per_image_pixel_values: list[torch.Tensor] = [] per_image_tgt_sizes: list[torch.Tensor] = [] for image in parsed_images: - ip_out = image_processor([image], **mm_kwargs) + ip_out = image_processor([image], **image_proc_kwargs) pv = ip_out["pixel_values"] # (1, C, P, sum_W) ts = ip_out["target_sizes"] # (n_slices, 2) if pv.ndim == 4 and pv.shape[0] == 1: @@ -210,7 +368,6 @@ def process_images( "tgt_sizes": per_image_tgt_sizes, } - ds_mode = self._resolve_downsample_mode(mm_kwargs) insert_layer_id = getattr( self.info.get_hf_config(), "insert_layer_id", @@ -245,7 +402,11 @@ def process_videos( image_processor = self.info.get_image_processor() patch_size = image_processor.patch_size video_max_slice = self.info.get_video_max_slice_num() - video_mm_kwargs = {**mm_kwargs, "max_slice_nums": video_max_slice} + video_ds_mode = self._resolve_downsample_mode(mm_kwargs, modality="video") + video_mm_kwargs = { + **_flat_processor_kwargs(mm_kwargs, video_ds_mode), + "max_slice_nums": video_max_slice, + } per_video_pixel_values: list[torch.Tensor] = [] per_video_tgt_sizes: list[torch.Tensor] = [] @@ -312,10 +473,17 @@ def process_videos( if not per_video_pixel_values: return {} + insert_layer_id = getattr(self.info.get_hf_config(), "insert_layer_id", -1) + merger_flag = video_ds_mode != "4x" and insert_layer_id >= 0 + return { "video_pixel_values": per_video_pixel_values, "video_tgt_sizes": per_video_tgt_sizes, "video_image_sizes": per_video_image_sizes, + "video_use_vit_merger": [ + torch.tensor([merger_flag], dtype=torch.bool) + for _ in per_video_pixel_values + ], } def _get_prompt_updates( @@ -324,7 +492,13 @@ def _get_prompt_updates( hf_processor_mm_kwargs: Mapping[str, object], out_mm_kwargs, ): - ds_mode = self._resolve_downsample_mode(hf_processor_mm_kwargs) + image_ds_mode = self._resolve_downsample_mode( + hf_processor_mm_kwargs, modality="image" + ) + video_ds_mode = self._resolve_downsample_mode( + hf_processor_mm_kwargs, modality="video" + ) + image_max_slice = self._resolve_max_slice_nums(hf_processor_mm_kwargs) placeholders = [ ("image", self.info.image_pattern), @@ -360,7 +534,8 @@ def get_image_replacement(item_idx: int): self.get_image_prompt_texts( image_size, item_idx, - downsample_mode=ds_mode, + downsample_mode=image_ds_mode, + max_slice_nums=image_max_slice, ), add_special_tokens=False, ), @@ -389,7 +564,7 @@ def get_video_replacement(item_idx: int): self.get_video_prompt_texts( frame_size, num_frames, - downsample_mode=ds_mode, + downsample_mode=video_ds_mode, video_idx=item_idx, ), add_special_tokens=False, @@ -409,7 +584,7 @@ def get_video_replacement(item_idx: int): self.get_video_prompt_texts( frame_size, num_frames, - downsample_mode=ds_mode, + downsample_mode=video_ds_mode, video_idx=item_idx, ), add_special_tokens=False, @@ -437,31 +612,28 @@ def _recompute_cached_prompt_update( new_update = super()._recompute_cached_prompt_update( cached_update, new_item_idx ) - # MiniCPM-V 4.6 prefixes video placeholders with `{idx}` - # (the base class only rewrites the image modality). + # Video placeholders carry an `{idx}` prefix that + # the base class does not rewrite. 4.7 has no prefix, so this is a + # no-op there. if cached_update.modality == "video": - tokenizer = self.info.get_tokenizer() - id_start = getattr(tokenizer, "image_id_start_token", "") - id_end = getattr(tokenizer, "image_id_end_token", "") - video_token = getattr(tokenizer, "video_token", "<|video_pad|>") - - text = tokenizer.decode(cached_update.content.full) - prev_item_idx = cached_update.item_idx - - new_update = new_update.with_content( - PromptUpdateDetails.select_token_ids( - cached_encode( - tokenizer, - text.replace( - f"{id_start}{prev_item_idx}{id_end}", - f"{id_start}{new_item_idx}{id_end}", - 1, + old_prefix = self._video_local_id_prefix(cached_update.item_idx) + new_prefix = self._video_local_id_prefix(new_item_idx) + if old_prefix != new_prefix: + tokenizer = self.info.get_tokenizer() + video_token = getattr(tokenizer, "video_token", "<|video_pad|>") + + text = tokenizer.decode(cached_update.content.full) + + new_update = new_update.with_content( + PromptUpdateDetails.select_token_ids( + cached_encode( + tokenizer, + text.replace(old_prefix, new_prefix, 1), + add_special_tokens=False, ), - add_special_tokens=False, - ), - cached_encode(tokenizer, video_token, add_special_tokens=False), + cached_encode(tokenizer, video_token, add_special_tokens=False), + ) ) - ) return new_update def _get_mm_fields_config( @@ -676,6 +848,31 @@ def get_num_image_tokens( return source_tokens + grids[0] * grids[1] * patch_tokens +def _stack_vit_merger_qkv( + weights: Iterable[tuple[str, torch.Tensor]], +) -> Iterable[tuple[str, torch.Tensor]]: + """Fold separate vit_merger q/k/v projections into the fused qkv_proj module. + + HF checkpoints store the window-attention projections separately, while + vLLM keeps them fused. The shard_id mechanism on the submodule's + WeightsMapper cannot do this: AutoWeightsLoader still recurses over every + incoming key after calling a submodule's load_weights, so a "k_proj" key + with no matching child raises before the mapping is honoured. + """ + pairs = ( + ("vit_merger.self_attn.q_proj.", "q"), + ("vit_merger.self_attn.k_proj.", "k"), + ("vit_merger.self_attn.v_proj.", "v"), + ) + for name, tensor in weights: + for src, shard in pairs: + if src in name: + name = name.replace(src, "vit_merger.self_attn.qkv_proj.", 1) + tensor.shard_id = shard + break + yield name, tensor + + class MiniCPMV4_6ViTWindowAttentionSelfAttn(nn.Module): def __init__( self, @@ -957,6 +1154,7 @@ class MiniCPMV4_6ForConditionalGeneration( SupportsPP, HasInnerState, IsHybrid, + SupportsEagle3, SupportsMRoPE, ): supports_encoder_tp_data = True @@ -1031,7 +1229,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): apply_encoder_attention_mask=True, prefix=maybe_prefix(prefix, "vpm"), ) - if config.drop_vision_last_layer: + if self._drop_vision_last_layer(): self.vpm.encoder.layers = self.vpm.encoder.layers[:-1] self.vit_merger = MiniCPMV4_6ViTWindowAttentionMerger( @@ -1062,6 +1260,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.language_model.make_empty_intermediate_tensors ) + def _drop_vision_last_layer(self) -> bool: + """Whether the last vision encoder layer is dropped. + + MiniCPM-V 4.6 checkpoints always define this. 4.7's config no longer + carries the field, so it overrides this. + """ + return self.config.drop_vision_last_layer + # ----- Multimodal parsing ----- def _parse_and_validate_vision_input( @@ -1174,35 +1380,35 @@ def _process_vision_input(self, image_input, use_vit_merger=None): if image_input["type"] == "image_embeds": return image_input["image_embeds"] - downsample_mode = None - if use_vit_merger is not None: - downsample_mode = "16x" if use_vit_merger else "4x" - image_features = self.get_vision_hidden_states( - image_input, - downsample_mode=downsample_mode, - ) - num_slices = image_input["num_slices"] - results = [] - idx = 0 - for n in num_slices.tolist(): - group = image_features[idx : idx + n] - results.append(torch.cat(group, dim=0)) - idx += n - return results + num_slices = image_input["num_slices"].tolist() + flags = _expand_flags(use_vit_merger, len(num_slices)) + + # The merger flag is per request, but the model runner may batch items + # from several concurrent requests. Encoding such a mixed batch would + # apply one request's downsample mode to another request's items, so + # group by flag and encode each group separately. + encoded: dict[int, torch.Tensor] = {} + for flag in dict.fromkeys(flags): + item_indices = [i for i, f in enumerate(flags) if f == flag] + group_input = _select_vision_items(image_input, item_indices, num_slices) + group_features = self.get_vision_hidden_states( + group_input, + downsample_mode=_vision_downsample_mode(flag), + ) + + idx = 0 + for i in item_indices: + n = num_slices[i] + encoded[i] = torch.cat(group_features[idx : idx + n], dim=0) + idx += n + + return [encoded[i] for i in range(len(num_slices))] # ----- Multimodal embedding interface ----- def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: - use_vit_merger_tensors = kwargs.pop("use_vit_merger", None) - use_vit_merger = None - if use_vit_merger_tensors is not None: - if isinstance(use_vit_merger_tensors, torch.Tensor): - use_vit_merger = bool(use_vit_merger_tensors.any().item()) - elif isinstance(use_vit_merger_tensors, list | tuple): - use_vit_merger = any( - bool(t.any().item()) if isinstance(t, torch.Tensor) else bool(t) - for t in use_vit_merger_tensors - ) + use_vit_merger = kwargs.pop("use_vit_merger", None) + video_use_vit_merger = kwargs.pop("video_use_vit_merger", None) image_kwargs = { k: v @@ -1235,7 +1441,7 @@ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: multimodal_embeddings += tuple( self._process_vision_input( video_input, - use_vit_merger=use_vit_merger, + use_vit_merger=video_use_vit_merger, ) ) @@ -1298,7 +1504,10 @@ def load_weights( weights: Iterable[tuple[str, torch.Tensor]], ) -> set[str]: loader = AutoWeightsLoader(self) - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + return loader.load_weights( + _stack_vit_merger_qkv(weights), + mapper=self.hf_to_vllm_mapper, + ) def get_mm_mapping(self) -> MultiModelKeys: return MultiModelKeys.from_string_field( diff --git a/vllm/model_executor/models/minicpmv4_7.py b/vllm/model_executor/models/minicpmv4_7.py new file mode 100644 index 000000000000..3ce860916965 --- /dev/null +++ b/vllm/model_executor/models/minicpmv4_7.py @@ -0,0 +1,764 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only MiniCPM-V 4.7 model. + +Same vision/LLM stack as MiniCPM-V 4.6, plus canvas 3D M-RoPE. +""" + +import math +from collections.abc import Mapping +from typing import Any, NamedTuple + +import torch + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalFieldConfig + +from .minicpmv import MiniCPMVDummyInputsBuilder +from .minicpmv4_6 import ( + MiniCPMV4_6ForConditionalGeneration, + MiniCPMV4_6MultiModalProcessor, + MiniCPMV4_6ProcessingInfo, + _scoped_mm_kwarg, +) + +logger = init_logger(__name__) + + +class MiniCPMV4_7ProcessingInfo(MiniCPMV4_6ProcessingInfo): + def _mm_max_slice_nums(self, **kwargs: object) -> int | None: + # No schema resolution here: it would call back into `get_hf_processor` + # below, which resolves this same value. + merged = self.ctx.get_merged_mm_kwargs(kwargs) + max_slice = _scoped_mm_kwarg(merged, "image", "max_slice_nums") + if max_slice is None: + max_slice = merged.get("max_slice_nums") + if max_slice is None: + return None + return int(max_slice) + + def _configured_max_slice_nums(self) -> int | None: + # Configured (server-level) value only: per-request values are passed + # explicitly to the image processor call, so they must not be written + # back to the shared processor instance. + return self._mm_max_slice_nums() + + def get_hf_processor(self, **kwargs: object): + hf_processor = super().get_hf_processor(**kwargs) + image_processor = getattr(hf_processor, "image_processor", None) + configured = self._configured_max_slice_nums() + if image_processor is not None and configured is not None: + image_processor.max_slice_nums = configured + return hf_processor + + def get_image_max_slice_num(self) -> int: + max_slice = self._configured_max_slice_nums() + if max_slice is not None: + return max_slice + return super().get_image_max_slice_num() + + +class MiniCPMV4_7MultiModalProcessor(MiniCPMV4_6MultiModalProcessor): + def _get_mm_fields_config( + self, + hf_inputs, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + fields = dict(super()._get_mm_fields_config(hf_inputs, hf_processor_mm_kwargs)) + # Canvas M-RoPE places slices from `tgt_sizes`. Stripping the payload of + # a prefix-cache-covered item only keeps CPU-side metadata, so without + # this a cached image loses its target sizes, `image_bounds` and + # `target_sizes` disagree, and the whole request silently falls back to + # sequential positions. + for key, modality in (("tgt_sizes", "image"), ("video_tgt_sizes", "video")): + if key in fields: + fields[key] = MultiModalFieldConfig.batched(modality, keep_on_cpu=True) + return fields + + def _video_local_id_prefix(self, video_idx: int) -> str: + # MiniCPMV4_7Processor emits no local `` in front of video + # placeholders: a video is one temporal sequence and 4.7 was trained + # without them. + return "" + + def get_image_prompt_texts( + self, + image_size, + image_idx: int = 0, + downsample_mode: str | None = None, + max_slice_nums: int | None = None, + ) -> str: + info = self.info + assert isinstance(info, MiniCPMV4_6ProcessingInfo) + if max_slice_nums is None: + max_slice_nums = info.get_image_max_slice_num() + return info.get_slice_image_placeholder( + image_size, + image_idx=image_idx, + max_slice_nums=max_slice_nums, + downsample_mode=downsample_mode, + ) + + +def _as_hw_pairs(tensor: torch.Tensor) -> torch.Tensor: + data = tensor.to(torch.long) + if data.numel() == 0: + return data.reshape(0, 2) + if data.shape[-1] != 2: + raise ValueError(f"expected last dim 2, got {tuple(data.shape)}") + return data.reshape(-1, 2) + + +def _sequential_positions(seq_len: int) -> torch.Tensor: + """Positions for the fallback where canvas M-RoPE cannot be computed.""" + return torch.arange(seq_len).unsqueeze(0).expand(3, -1) + + +@MULTIMODAL_REGISTRY.register_processor( + MiniCPMV4_7MultiModalProcessor, + info=MiniCPMV4_7ProcessingInfo, + dummy_inputs=MiniCPMVDummyInputsBuilder, +) +class MiniCPMV4_7ForConditionalGeneration(MiniCPMV4_6ForConditionalGeneration): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + self._init_canvas_mrope(vllm_config) + + def _drop_vision_last_layer(self) -> bool: + # MiniCPMV4_7Config no longer carries this field (transformers marks it + # removed), and 4.7 keeps every vision encoder layer. + return False + + def _init_canvas_mrope(self, vllm_config: VllmConfig) -> None: + config = self.config + uses_canvas = bool(getattr(config, "uses_mrope_canvas", False)) + if not uses_canvas: + mrope_mode = getattr(config, "mrope_mode", None) + uses_canvas = mrope_mode is not None and "canvas" in str(mrope_mode).lower() + + self._uses_canvas_mrope = uses_canvas + self._canvas_special_ids: dict[str, int] | None = None + if not uses_canvas: + return + + # MiniCPMV4_7Config pins the structural ids in config.json; the + # conversion script resolves them from the tokenizer. Prefer them so the + # canvas matches training, and fall back to the tokenizer for checkpoints + # whose config predates these fields. + configured_ids: dict[str, int | None] = { + "im_start_id": getattr(config, "image_start_id", None), + "im_end_id": getattr(config, "image_end_id", None), + "slice_start_id": getattr(config, "slice_start_id", None), + "slice_end_id": getattr(config, "slice_end_id", None), + "newline_id": getattr(config, "newline_id", None), + } + if all(value is not None for value in configured_ids.values()): + self._canvas_special_ids = { + name: int(value) + for name, value in configured_ids.items() + if value is not None + } + return + + from vllm.tokenizers.registry import get_tokenizer + + model_config = vllm_config.model_config + tokenizer = get_tokenizer( + model_config.tokenizer, + revision=model_config.tokenizer_revision, + trust_remote_code=model_config.trust_remote_code, + ) + self._canvas_special_ids = { + "im_start_id": tokenizer.convert_tokens_to_ids( + getattr(tokenizer, "image_start_token", "") + ), + "im_end_id": tokenizer.convert_tokens_to_ids( + getattr(tokenizer, "image_end_token", "") + ), + "slice_start_id": tokenizer.convert_tokens_to_ids( + getattr(tokenizer, "slice_start_token", "") + ), + "slice_end_id": tokenizer.convert_tokens_to_ids( + getattr(tokenizer, "slice_end_token", "") + ), + "newline_id": tokenizer.encode("\n", add_special_tokens=False)[0], + } + + def get_mrope_input_positions( + self, + input_tokens: list[int], + mm_features: list[MultiModalFeatureSpec], + ) -> tuple[torch.Tensor, int]: + seq_len = len(input_tokens) + if ( + not getattr(self, "_uses_canvas_mrope", False) + or self._canvas_special_ids is None + or not mm_features + ): + return _sequential_positions(seq_len), 0 + + input_ids = torch.tensor(input_tokens, dtype=torch.long) + image_bounds = build_image_bounds(input_ids, self._canvas_special_ids) + + target_sizes_list: list[torch.Tensor] = [] + for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset): + if mm_feature.data is None: + continue + for key in ("tgt_sizes", "video_tgt_sizes"): + tgt = mm_feature.data.get(key) + if tgt is None or not isinstance(tgt.data, torch.Tensor): + continue + try: + target_sizes_list.append(_as_hw_pairs(tgt.data)) + except ValueError: + # Not an (N, 2) size table. The bounds/sizes match below + # turns this into a sequential fallback. + continue + + if not target_sizes_list or image_bounds.numel() == 0: + logger.warning( + "MiniCPM-V 4.7 canvas M-RoPE missing inputs: " + "tgt_chunks=%s image_bounds=%s; using sequential positions.", + len(target_sizes_list), + tuple(image_bounds.shape) if image_bounds is not None else None, + ) + return _sequential_positions(seq_len), 0 + + target_sizes = torch.cat(target_sizes_list, dim=0) + if target_sizes.ndim != 2 or target_sizes.shape[-1] != 2: + logger.warning( + "MiniCPM-V 4.7 canvas M-RoPE unexpected tgt_sizes shape %s; " + "using sequential positions.", + tuple(target_sizes.shape), + ) + return _sequential_positions(seq_len), 0 + if image_bounds.shape[0] != target_sizes.shape[0]: + logger.warning( + "MiniCPM-V 4.7 canvas M-RoPE size mismatch: " + "image_bounds=%s target_sizes=%s; using sequential positions.", + tuple(image_bounds.shape), + tuple(target_sizes.shape), + ) + return _sequential_positions(seq_len), 0 + + try: + pos3d = _compute_canvas_single( + input_ids, + torch.arange(seq_len, dtype=torch.long), + image_bounds, + target_sizes, + self._canvas_special_ids, + ) + except Exception: + logger.warning( + "MiniCPM-V 4.7 canvas M-RoPE failed; " + "falling back to sequential positions.", + exc_info=True, + ) + return _sequential_positions(seq_len), 0 + return pos3d, canvas_rope_delta(pos3d, seq_len) + + +# Canvas 3D M-RoPE: position ids for the slice canvas. +def compute_llm_grid( + n_tokens: int, vision_grid_height: int, vision_grid_width: int +) -> tuple[int, int]: + if n_tokens <= 0 or vision_grid_height <= 0 or vision_grid_width <= 0: + return 0, 0 + vit_total = vision_grid_height * vision_grid_width + if vit_total == n_tokens: + return vision_grid_height, vision_grid_width + if vit_total < n_tokens: + return 0, 0 + factor = int(round(math.sqrt(vit_total / n_tokens))) + if ( + factor > 0 + and vision_grid_height % factor == 0 + and vision_grid_width % factor == 0 + ): + height, width = vision_grid_height // factor, vision_grid_width // factor + if height * width == n_tokens: + return height, width + ratio = vision_grid_height / vision_grid_width + width = max(1, int(round(math.sqrt(n_tokens / ratio)))) + height = n_tokens // width + if height * width == n_tokens and height > 0: + return height, width + return 0, 0 + + +def build_image_bounds( + input_ids: torch.LongTensor, special_token_ids: dict +) -> torch.LongTensor: + start_ids = [ + x + for x in ( + special_token_ids.get("im_start_id"), + special_token_ids.get("slice_start_id"), + ) + if x is not None + ] + end_ids = [ + x + for x in ( + special_token_ids.get("im_end_id"), + special_token_ids.get("slice_end_id"), + ) + if x is not None + ] + if not start_ids or not end_ids: + return torch.zeros(0, 2, dtype=torch.long, device=input_ids.device) + start_cond = torch.zeros_like(input_ids, dtype=torch.bool) + end_cond = torch.zeros_like(input_ids, dtype=torch.bool) + for sid in start_ids: + start_cond |= input_ids == sid + for eid in end_ids: + end_cond |= input_ids == eid + starts = torch.where(start_cond)[0] + 1 + ends = torch.where(end_cond)[0] + n = min(len(starts), len(ends)) + if n == 0: + return torch.zeros(0, 2, dtype=torch.long, device=input_ids.device) + return torch.stack([starts[:n], ends[:n]], dim=-1) + + +def _compute_canvas_single( + input_ids, position_ids_2d, image_bound, target_sizes, special_token_ids +): + # Inference always handles one unpadded sequence. The packed (cu_seqlens) + # form below is kept only to mirror the transformers reference + # implementation of its canvas position computation. + seq_len = input_ids.shape[0] + cu_seqlens = torch.tensor([0, seq_len], device=input_ids.device, dtype=torch.long) + return _compute_canvas_packed( + position_ids_2d.unsqueeze(0), + cu_seqlens, + image_bound, + target_sizes, + input_ids.unsqueeze(0), + special_token_ids, + )[:, 0, :] + + +def _target_hw(target_sizes: torch.Tensor, index: int) -> tuple[int, int]: + """Patch-grid height/width recorded for one image or slice.""" + return int(target_sizes[index, 0]), int(target_sizes[index, 1]) + + +class _CanvasLayout(NamedTuple): + """Grid of one canvas measured in LLM tokens. + + A canvas is a thumbnail plus the slices cut from it, tiled row-major. + """ + + thumb_h: int + thumb_w: int + """Thumbnail grid.""" + slice_h: int + slice_w: int + """Per-slice grid. Zero when the canvas has no slices.""" + cols: int + """Slices sharing a canvas row. Zero when the canvas has no slices.""" + height: int + width: int + """Assembled canvas.""" + + +def _canvas_layout( + thumbnail: tuple[int, int, int], + slices: list[tuple[int, int, int]], + target_sizes: torch.Tensor, +) -> _CanvasLayout: + """Measure the canvas for one ``(thumbnail, slices)`` group. + + ``thumbnail`` and each slice are ``(start, end, target_index)`` token spans. + """ + t_bs, t_be, t_gi = thumbnail + thumb_h, thumb_w = compute_llm_grid(t_be - t_bs, *_target_hw(target_sizes, t_gi)) + if not slices: + return _CanvasLayout( + thumb_h, thumb_w, 0, 0, 0, max(thumb_h, 1), max(thumb_w, 1) + ) + + cols = len(slices) + for k in range(len(slices) - 1): + if slices[k + 1][0] - slices[k][1] > 2: + cols = k + 1 + break + rows = len(slices) // cols if cols > 0 else 1 + if rows * cols != len(slices): + rows, cols = 1, len(slices) + + s0_bs, s0_be, s0_gi = slices[0] + slice_h, slice_w = compute_llm_grid(s0_be - s0_bs, *_target_hw(target_sizes, s0_gi)) + return _CanvasLayout( + thumb_h, thumb_w, slice_h, slice_w, cols, rows * slice_h, cols * slice_w + ) + + +def _assign_canvas_span( + pos3d: torch.Tensor, + *, + thumbnail: tuple[int, int, int], + span_end: int, + slices: list[tuple[int, int, int]], + base: int, + target_sizes: torch.Tensor, +) -> int: + """Place one canvas at ``base``, writing its 3-D positions into ``pos3d``. + + Shared by the video-frame and single-image paths; they differ only in how + the per-canvas base advances, not in the geometry placed here. + + Mutates ``pos3d`` in place, as the packed canvas computation has always + done. Returns the distance to the next canvas base, so each caller keeps + only its own base advancement. + """ + t_bs, t_be, _ = thumbnail + device = pos3d.device + span_start = t_bs - 1 # the token + n_thumb = t_be - t_bs + layout = _canvas_layout(thumbnail, slices, target_sizes) + cols, canvas_H, canvas_W = layout.cols, layout.height, layout.width + + # base coat + pos3d[0, 0, span_start:span_end] = base + pos3d[1, 0, span_start:span_end] = base + pos3d[2, 0, span_start:span_end] = base + + # -> halo + halo_lo = max(base - 1, 0) + pos3d[1, 0, span_start] = halo_lo + pos3d[2, 0, span_start] = halo_lo + + # -> halo (right after thumbnail) + if t_be < span_end: + pos3d[1, 0, t_be] = base + canvas_H + pos3d[2, 0, t_be] = base + canvas_W + + # thumbnail visual tokens + llm_th, llm_tw = layout.thumb_h, layout.thumb_w + if llm_th > 0 and llm_tw > 0 and llm_th * llm_tw == n_thumb: + if slices and canvas_H > 0 and canvas_W > 0: + h_c = torch.linspace(0, canvas_H - 1, llm_th, device=device).round().long() + w_c = torch.linspace(0, canvas_W - 1, llm_tw, device=device).round().long() + else: + h_c = torch.arange(llm_th, device=device) + w_c = torch.arange(llm_tw, device=device) + h_idx = h_c.view(-1, 1).expand(-1, llm_tw).reshape(-1) + w_idx = w_c.view(1, -1).expand(llm_th, -1).reshape(-1) + pos3d[0, 0, t_bs:t_be] = base + pos3d[1, 0, t_bs:t_be] = h_idx + base + pos3d[2, 0, t_bs:t_be] = w_idx + base + else: + fallback = torch.arange(n_thumb, device=device, dtype=torch.long) + base + for d in range(3): + pos3d[d, 0, t_bs:t_be] = fallback + + # slice visual tokens + slice specials + for k, (s_bs, s_be, s_gi) in enumerate(slices): + n_s = s_be - s_bs + sh_k, sw_k = compute_llm_grid(n_s, *_target_hw(target_sizes, s_gi)) + h_off = (k // cols) * layout.slice_h + w_off = (k % cols) * layout.slice_w + + slice_start_pos = s_bs - 1 + if slice_start_pos >= span_start: + pos3d[1, 0, slice_start_pos] = base + h_off + pos3d[2, 0, slice_start_pos] = base + w_off + + slice_end_pos = s_be + if sh_k > 0 and sw_k > 0: + end_h = h_off + sh_k - 1 + end_w = w_off + sw_k - 1 + else: + end_h = h_off + end_w = w_off + if slice_end_pos < span_end: + pos3d[1, 0, slice_end_pos] = base + end_h + pos3d[2, 0, slice_end_pos] = base + end_w + + if sh_k > 0 and sw_k > 0 and sh_k * sw_k == n_s: + h_idx = ( + torch.arange(sh_k, device=device) + .view(-1, 1) + .expand(-1, sw_k) + .reshape(-1) + ) + w_idx = ( + torch.arange(sw_k, device=device) + .view(1, -1) + .expand(sh_k, -1) + .reshape(-1) + ) + pos3d[0, 0, s_bs:s_be] = base + pos3d[1, 0, s_bs:s_be] = h_idx + h_off + base + pos3d[2, 0, s_bs:s_be] = w_idx + w_off + base + else: + fallback = torch.arange(n_s, device=device, dtype=torch.long) + base + for d in range(3): + pos3d[d, 0, s_bs:s_be] = fallback + + # \n between slice rows -> W = right edge + 1. + # H stays within the ended row (does NOT cross to next row): the \n sits at + # the right-outside of the row that just ended, at H = that row's last row + # and W = one column past the rightmost slice. + for k in range(len(slices) - 1): + gap_s = slices[k][1] + gap_e = slices[k + 1][0] + if gap_e - gap_s <= 2: + continue + boundary_h = ((k // cols) + 1) * layout.slice_h - 1 + right_edge_w = cols * layout.slice_w + for nl_pos in range(gap_s + 1, gap_e - 1): + pos3d[1, 0, nl_pos] = base + boundary_h + pos3d[2, 0, nl_pos] = base + right_edge_w + + return max(canvas_H, canvas_W) + 1 + + +def _compute_canvas_packed( + position_ids_2d, + cu_seqlens, + image_bound, + target_sizes, + input_ids, + special_token_ids, +): + if position_ids_2d.ndim == 1: + position_ids_2d = position_ids_2d.unsqueeze(0) + if input_ids.ndim == 1: + input_ids = input_ids.unsqueeze(0) + + has_images = ( + isinstance(image_bound, torch.Tensor) + and image_bound.numel() > 0 + and isinstance(target_sizes, torch.Tensor) + and target_sizes.numel() > 0 + ) + if not has_images: + return position_ids_2d.unsqueeze(0).expand(3, -1, -1).contiguous() + + device = position_ids_2d.device + pos3d = position_ids_2d.unsqueeze(0).expand(3, -1, -1).clone() + ids_flat = input_ids.view(-1) + + im_start_id = special_token_ids.get("im_start_id") + im_end_id = special_token_ids.get("im_end_id") + slice_start_id = special_token_ids.get("slice_start_id") + slice_end_id = special_token_ids.get("slice_end_id") + + structural_ids = { + v + for v in (im_start_id, im_end_id, slice_start_id, slice_end_id) + if v is not None + } + + num_seqs = len(cu_seqlens) - 1 + img_ptr = 0 + num_imgs = len(image_bound) + + newline_id = special_token_ids.get("newline_id") + gap_ok_ids = structural_ids | ({newline_id} if newline_id is not None else set()) + + def _group_end_pos(g, s_end): + """First position after a group's trailing structural tokens. + + Matches the transformers reference, which closes a span over every + marker in ``markers`` -- newlines included. The chat template joins + content parts with ``\\n``, so an image followed by text always ends with + one; leaving it out moves that newline out of the canvas and pushes + every later position one step forward. + """ + le = g["slices"][-1][1] if g["slices"] else g["thumbnail"][1] + ge = le + while ge < s_end and ids_flat[ge].item() in gap_ok_ids: + ge += 1 + return min(ge, s_end) + + def _is_tight_gap(g_prev, g_next, s_end): + """Check if only structural/nl tokens sit between two groups.""" + ge = _group_end_pos(g_prev, s_end) + ns = g_next["thumbnail"][0] - 1 # im_start of next group + return all(ids_flat[p].item() in gap_ok_ids for p in range(ge, ns)) + + for si in range(num_seqs): + s_start = cu_seqlens[si].item() + s_end = cu_seqlens[si + 1].item() + if s_start >= s_end: + continue + + seq_imgs = [] + while img_ptr < num_imgs: + bs_ = image_bound[img_ptr, 0].item() + be_ = image_bound[img_ptr, 1].item() + if bs_ >= s_start and be_ <= s_end: + seq_imgs.append((bs_, be_, img_ptr)) + img_ptr += 1 + else: + break + if not seq_imgs: + continue + + # ---- group images: thumbnail + following slices ---- + # merge consecutive no-slice frames into one video group + raw_groups = [] + cur_group: dict[str, Any] | None = None + for bs_, be_, gi in seq_imgs: + marker_pos = bs_ - 1 + is_slice = ( + slice_start_id is not None + and marker_pos >= s_start + and ids_flat[marker_pos].item() == slice_start_id + ) + if is_slice and cur_group is not None: + cur_group["slices"].append((bs_, be_, gi)) + else: + if cur_group is not None: + raw_groups.append(cur_group) + cur_group = {"thumbnail": (bs_, be_, gi), "slices": []} + if cur_group is not None: + raw_groups.append(cur_group) + + # merge consecutive tightly-adjacent groups into video groups. + # Two groups are "tightly adjacent" if the gap between them + # contains only structural tokens and \n (no real text). + # This handles both no-slice frames and frames-with-slices. + groups = [] + i_g = 0 + while i_g < len(raw_groups): + video_frames = [raw_groups[i_g]] + j_g = i_g + 1 + while j_g < len(raw_groups) and _is_tight_gap( + raw_groups[j_g - 1], raw_groups[j_g], s_end + ): + video_frames.append(raw_groups[j_g]) + j_g += 1 + if len(video_frames) == 1: + groups.append(raw_groups[i_g]) + else: + groups.append({"video_frames": video_frames}) + i_g = j_g + + # ---- assign 3-D positions ---- + pos = 0 + cursor = s_start + + for group in groups: + is_video_group = "video_frames" in group + + if is_video_group: + # === Video group: multiple frames merged === + # Each frame gets its own base_f so different frames have + # distinct spatial positions. Inter-frame gap tokens (\n + # etc.) are pure 1D text. After each gap, pos advances + # by 1 extra so halo_lo of the next frame does not + # collide with the last gap token. + frames = group["video_frames"] + first_frame = frames[0] + last_frame = frames[-1] + + # group_start = of first frame + group_start = first_frame["thumbnail"][0] - 1 + # group_end = after last structural token of last frame + group_end = _group_end_pos(last_frame, s_end) + + # text before video group + text_len = group_start - cursor + if text_len > 0: + t_pos = ( + torch.arange(text_len, device=device, dtype=torch.long) + pos + ) + for d in range(3): + pos3d[d, 0, cursor:group_start] = t_pos + pos += text_len + + # process each frame with its own per-frame base + frame_cursor = group_start + for vf in frames: + t_bs = vf["thumbnail"][0] + vf_slices = vf["slices"] + im_start_pos = t_bs - 1 + + # --- inter-frame gap: pure 1D text T=H=W=pos --- + gap_len = im_start_pos - frame_cursor + if gap_len > 0: + t_pos = ( + torch.arange(gap_len, device=device, dtype=torch.long) + pos + ) + for d in range(3): + pos3d[d, 0, frame_cursor:im_start_pos] = t_pos + pos += gap_len + pos += 1 # buffer: prevents halo_lo collision + + # --- per-frame base --- + base_f = pos + frame_end = _group_end_pos(vf, group_end) + pos = base_f + _assign_canvas_span( + pos3d, + thumbnail=vf["thumbnail"], + span_end=frame_end, + slices=vf_slices, + base=base_f, + target_sizes=target_sizes, + ) + frame_cursor = frame_end + + # trailing tokens after last frame (if any) + if frame_cursor < group_end: + trail_len = group_end - frame_cursor + t_pos = ( + torch.arange(trail_len, device=device, dtype=torch.long) + pos + ) + for d in range(3): + pos3d[d, 0, frame_cursor:group_end] = t_pos + pos += trail_len + + cursor = group_end + + else: + # === Single image group (thumbnail + optional slices) === + slices = group["slices"] + + group_start = group["thumbnail"][0] - 1 + group_end = _group_end_pos(group, s_end) + + # text before group + text_len = group_start - cursor + if text_len > 0: + t_pos = ( + torch.arange(text_len, device=device, dtype=torch.long) + pos + ) + for d in range(3): + pos3d[d, 0, cursor:group_start] = t_pos + pos += text_len + + base = pos + pos = base + _assign_canvas_span( + pos3d, + thumbnail=group["thumbnail"], + span_end=group_end, + slices=slices, + base=base, + target_sizes=target_sizes, + ) + cursor = group_end + + # remaining text + rem = s_end - cursor + if rem > 0: + t_pos = torch.arange(rem, device=device, dtype=torch.long) + pos + for d in range(3): + pos3d[d, 0, cursor:s_end] = t_pos + + return pos3d + + +def canvas_rope_delta(position_ids: torch.Tensor, seq_len: int) -> int: + """Offset applied to the 1-D positions of tokens decoded after the prompt. + + The first decoded token sits one past the last prefill position. + """ + return int(position_ids.amax().item()) + 1 - seq_len diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 875179a88fae..fa2924441314 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -514,6 +514,10 @@ "minicpmv4_6", "MiniCPMV4_6ForConditionalGeneration", ), + "MiniCPMV4_7ForConditionalGeneration": ( + "minicpmv4_7", + "MiniCPMV4_7ForConditionalGeneration", + ), "Mistral3ForConditionalGeneration": ( "mistral3", "Mistral3ForConditionalGeneration", diff --git a/vllm/parser/minicpmv.py b/vllm/parser/minicpmv.py new file mode 100644 index 000000000000..88a6865c40c9 --- /dev/null +++ b/vllm/parser/minicpmv.py @@ -0,0 +1,405 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniCPM-V family parser. + +Wraps the Qwen3 grammar (```` reasoning + XML tool calls) with two +MiniCPM-V specifics: + +- ``<|im_end|>`` and the reserved marker tokens are passthrough terminals, + so they never leak into ``content``. +- The tokenizer decodes a newline as the two characters ``\\`` and ``n`` + (e.g. token id 1639) instead of ``\\n``, so the decoded text is recovered + before it reaches the client. Code blocks, inline code and LaTeX are left + alone, as are tool call arguments. +""" + +from __future__ import annotations + +import functools +from collections.abc import Sequence +from dataclasses import replace +from typing import TYPE_CHECKING + +from vllm.entrypoints.generate.base.protocol import DeltaMessage +from vllm.parser.abstract_parser import DelegatingParser +from vllm.parser.engine.parser_engine import ParserEngine +from vllm.parser.engine.parser_engine_config import ( + ParserEngineConfig, + ParserState, + Transition, +) +from vllm.parser.qwen3 import Qwen3Parser, qwen3_config + +if TYPE_CHECKING: + from vllm.entrypoints.generate.base.protocol import FunctionCall + from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest + from vllm.entrypoints.openai.responses.protocol import ResponsesRequest + from vllm.tokenizers import TokenizerLike + from vllm.tool_parsers.abstract_tool_parser import Tool + + +_RESERVED_MARKER_IDS = range(12, 16) +_IM_END = "<|im_end|>" + + +@functools.cache +def minicpmv_config(thinking: bool) -> ParserEngineConfig: + """MiniCPM-V grammar: the qwen3 grammar plus MiniCPM-V markers. + + Both modes keep the qwen3 tool call terminals and transitions, so + ```` output is parsed whether or not thinking is enabled. + """ + name = "minicpmv" if thinking else "minicpmv_no_thinking" + base = qwen3_config(thinking=thinking, name=name) + + transitions = dict(base.transitions) + if not thinking: + # Default mode: the model answers directly, so `` must not open + # a reasoning span that would swallow the answer. A stray + # `...` in the output is still dropped, not printed. + transitions[(ParserState.CONTENT, "THINK_START")] = Transition( + ParserState.REASONING + ) + base = replace(base, strip_trailing_reasoning_whitespace=False) + + terminals = dict(base.terminals) + token_id_terminals = dict(base.token_id_terminals) + terminals["IM_END"] = _IM_END + token_id_terminals["IM_END"] = _IM_END + for state in (ParserState.CONTENT, ParserState.REASONING): + transitions[(state, "IM_END")] = Transition(state) + + for index in _RESERVED_MARKER_IDS: + text_name = f"RESERVED_TEXT_{index}" + token_name = f"RESERVED_TOKEN_{index}" + terminals[text_name] = f"" + terminals[token_name] = f"<|reserved_{index}|>" + token_id_terminals[token_name] = f"<|reserved_{index}|>" + for state in (ParserState.CONTENT, ParserState.REASONING): + transitions[(state, text_name)] = Transition(state) + transitions[(state, token_name)] = Transition(state) + + return replace( + base, + terminals=terminals, + token_id_terminals=token_id_terminals, + transitions=transitions, + ) + + +class EscapedNewlineNormalizer: + """Recover literal escape sequences written by the model as real text. + + Scans character by character so that a backslash inside a code block, + inline code, LaTeX span or a doubled escape is left untouched. A trailing + half of an escape is held in ``_pending`` until the next chunk arrives, + which keeps streaming output identical to the non-streaming result. + """ + + def __init__(self) -> None: + self._closing_marker: str | None = None + self._pending = "" + + def feed(self, text: str, *, final: bool = False) -> str: + text = self._pending + text + self._pending = "" + output: list[str] = [] + index = 0 + + while index < len(text): + if self._closing_marker is not None: + marker = self._closing_marker + remaining = text[index:] + if remaining.startswith(marker): + output.append(marker) + index += len(marker) + self._closing_marker = None + continue + if not final and marker.startswith(remaining): + self._pending = remaining + break + output.append(text[index]) + index += 1 + continue + + remaining = text[index:] + if text[index] == "`": + if not final and remaining in ("`", "``"): + self._pending = remaining + break + if remaining.startswith("```"): + output.append("```") + index += 3 + self._closing_marker = "```" + else: + output.append("`") + index += 1 + self._closing_marker = "`" + continue + + if text[index] == "$": + if not final and remaining == "$": + self._pending = remaining + break + if remaining.startswith("$$"): + output.append("$$") + index += 2 + self._closing_marker = "$$" + else: + output.append("$") + index += 1 + self._closing_marker = "$" + continue + + if text[index] != "\\": + output.append(text[index]) + index += 1 + continue + + slash_end = index + 1 + while slash_end < len(text) and text[slash_end] == "\\": + slash_end += 1 + if slash_end - index > 1: + if not final and slash_end == len(text): + # The run may still grow, and its parity decides whether a + # following `n` is a newline or stays literal. Hold it until + # the run is complete. + self._pending = text[index:slash_end] + break + output.append(text[index:slash_end]) + index = slash_end + continue + + if not final and remaining in ("\\", "\\r", "\\r\\"): + self._pending = remaining + break + if remaining.startswith("\\r\\n"): + output.append("\n") + index += 4 + elif remaining.startswith(("\\n", "\\r")): + output.append("\n") + index += 2 + elif remaining.startswith("\\("): + output.append("\\(") + index += 2 + self._closing_marker = "\\)" + elif remaining.startswith("\\["): + output.append("\\[") + index += 2 + self._closing_marker = "\\]" + else: + output.append("\\") + index += 1 + + if final and self._pending: + output.append(self._pending) + self._pending = "" + return "".join(output) + + +def recover_newlines(text: str | None) -> str | None: + if text is None or "\\" not in text: + return text + return EscapedNewlineNormalizer().feed(text, final=True) + + +class MiniCPMVOutputNormalizer: + """Apply :class:`EscapedNewlineNormalizer` to a whole stream. + + Reasoning and content are tracked separately so that a half-written escape + at the end of a reasoning chunk does not swallow the first characters of + the content that follows. + """ + + def __init__(self) -> None: + self._reasoning = EscapedNewlineNormalizer() + self._content = EscapedNewlineNormalizer() + + def normalize_delta( + self, + delta: DeltaMessage | None, + *, + finished: bool, + ) -> DeltaMessage | None: + reasoning_finished = finished or ( + delta is not None and (delta.content is not None or bool(delta.tool_calls)) + ) + reasoning = self._reasoning.feed( + delta.reasoning if delta and delta.reasoning else "", + final=reasoning_finished, + ) + content = self._content.feed( + delta.content if delta and delta.content else "", + final=finished, + ) + + if delta is None: + if not reasoning and not content: + return None + return DeltaMessage( + reasoning=reasoning or None, + content=content or None, + ) + + if reasoning: + delta.reasoning = reasoning + elif delta.reasoning is not None: + delta.reasoning = None + if content: + delta.content = content + elif delta.content is not None: + delta.content = None + + # A chunk that normalizes to nothing must be dropped, not returned as an + # all-None DeltaMessage: the engine only skips on `None`, so anything + # else reaches the client as an empty delta. + if not delta.content and not delta.reasoning and not delta.tool_calls: + return None + return delta + + +class MiniCPMVParser(Qwen3Parser): + """Parse MiniCPM-V family output in thinking and non-thinking modes.""" + + CONFIG_NAME = "minicpmv" + + def __init__( + self, + tokenizer: TokenizerLike, + tools: list[Tool] | None = None, + **kwargs, + ) -> None: + chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} + # The chat template emits `` unless `enable_thinking` is + # explicitly false, so the grammar must default to thinking too, or the + # reasoning span is misread as content. + thinking_enabled = chat_kwargs.get("enable_thinking", True) + kwargs.setdefault("parser_engine_config", minicpmv_config(thinking_enabled)) + super().__init__(tokenizer, tools, **kwargs) + # ``Qwen3Parser.__init__`` reads the same key with a default of True, + # which does not apply here: MiniCPM-V defaults to thinking disabled. + # Keep the attribute in sync with the grammar built above. + self.thinking_enabled = thinking_enabled + + @property + def reasoning_ended(self) -> bool: + # The engine derives this from the presence of THINK terminals, which + # both MiniCPM-V configs keep, so the base property stays False in + # non-thinking mode. There is no reasoning span to end there. + if not self.thinking_enabled: + return True + return super().reasoning_ended + + def extract_reasoning( + self, + model_output: str, + request: ChatCompletionRequest | ResponsesRequest, + ) -> tuple[str | None, str | None]: + if self.thinking_enabled: + return super().extract_reasoning(model_output, request) + _, content = ParserEngine.extract_reasoning(self, model_output, request) + return None, content + + def extract_reasoning_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + ) -> DeltaMessage | None: + if self.thinking_enabled: + return super().extract_reasoning_streaming( + previous_text, + current_text, + delta_text, + previous_token_ids, + current_token_ids, + delta_token_ids, + ) + + delta = ParserEngine.extract_reasoning_streaming( + self, + previous_text, + current_text, + delta_text, + previous_token_ids, + current_token_ids, + delta_token_ids, + ) + if delta is None: + return None + delta.reasoning = None + if delta.content is None and not delta.tool_calls: + return None + return delta + + def get_streaming_fallback_content( + self, + text: str, + request: ChatCompletionRequest | ResponsesRequest, + ) -> str | None: + if self.thinking_enabled: + return super().get_streaming_fallback_content(text, request) + delta = ParserEngine.finish_streaming(self) + return delta.content if delta is not None else None + + +class MiniCPMVUnifiedParser(DelegatingParser): + """MiniCPM-V serving parser: the adapters plus newline recovery. + + Recovery cannot live in :class:`MiniCPMVParser`: the serving stack drives + the reasoning/tool adapters, never ``ParserEngine.parse`` / + ``ParserEngine.parse_delta``, so a rewrite applied there would only be + visible to callers that construct the engine directly. + """ + + def __init__( + self, + tokenizer: TokenizerLike, + tools: list[Tool] | None = None, + *args, + **kwargs, + ) -> None: + super().__init__(tokenizer, tools, *args, **kwargs) + self._output_normalizer = MiniCPMVOutputNormalizer() + + def parse( + self, + model_output: str, + request: ChatCompletionRequest | ResponsesRequest, + enable_auto_tools: bool = False, + model_output_token_ids: Sequence[int] = (), + ) -> tuple[str | None, str | None, list[FunctionCall] | None]: + reasoning, content, tool_calls = super().parse( + model_output, + request, + enable_auto_tools, + model_output_token_ids, + ) + # Tool call arguments are JSON, where a literal `\n` is meaningful and + # must survive as written. + return ( + recover_newlines(reasoning), + recover_newlines(content), + tool_calls, + ) + + def parse_delta( + self, + delta_text: str, + delta_token_ids: list[int], + request: ChatCompletionRequest | ResponsesRequest, + prompt_token_ids: list[int] | None = None, + *, + finished: bool, + ) -> DeltaMessage | None: + delta = super().parse_delta( + delta_text, + delta_token_ids, + request, + prompt_token_ids, + finished=finished, + ) + return self._output_normalizer.normalize_delta(delta, finished=finished) diff --git a/vllm/parser/parser_manager.py b/vllm/parser/parser_manager.py index cf0be908ed63..6766d15e914c 100644 --- a/vllm/parser/parser_manager.py +++ b/vllm/parser/parser_manager.py @@ -242,6 +242,19 @@ class _CohereCommandParser(CohereCommandParser): return _CohereCommandParser + if "minicpmv" in {reasoning_parser_name, tool_parser_name}: + from vllm.parser.minicpmv import MiniCPMVUnifiedParser + + r_cls = reasoning_parser_cls + t_cls = tool_parser_cls + + class _MiniCPMVParser(MiniCPMVUnifiedParser): + reasoning_parser_cls = r_cls + tool_parser_cls = t_cls + tool_strict_level = strict_level + + return _MiniCPMVParser + from vllm.parser.abstract_parser import DelegatingParser r_cls = reasoning_parser_cls diff --git a/vllm/reasoning/__init__.py b/vllm/reasoning/__init__.py index f9ab1dc9a11f..9e75e9b77ef4 100644 --- a/vllm/reasoning/__init__.py +++ b/vllm/reasoning/__init__.py @@ -116,6 +116,10 @@ "mimo_engine_reasoning_parser", "MiMoParserReasoningAdapter", ), + "minicpmv": ( + "minicpmv_reasoning_parser", + "MiniCPMVParserReasoningAdapter", + ), "minimax_m2": ( "minimax_m2_reasoning_parser", "MiniMaxM2ReasoningParser", diff --git a/vllm/reasoning/minicpmv_reasoning_parser.py b/vllm/reasoning/minicpmv_reasoning_parser.py new file mode 100644 index 000000000000..402c0524cab0 --- /dev/null +++ b/vllm/reasoning/minicpmv_reasoning_parser.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.parser.engine.adapters import ParserEngineReasoningAdapter +from vllm.parser.minicpmv import MiniCPMVParser + + +class MiniCPMVParserReasoningAdapter(ParserEngineReasoningAdapter): + # Reasoning-only on purpose: this family registers no tool parser, because + # MiniCPM-V emits the Qwen3-Coder XML syntax verbatim and `qwen3_coder` + # already covers it. ``make_adapters`` only builds reasoning/tool pairs, so + # the two halves are declared explicitly here. + _parser_engine_cls = MiniCPMVParser + + +__all__ = ["MiniCPMVParserReasoningAdapter"] diff --git a/vllm/transformers_utils/chat_templates/registry.py b/vllm/transformers_utils/chat_templates/registry.py index 422854ab3d4d..f902ecbbaafb 100644 --- a/vllm/transformers_utils/chat_templates/registry.py +++ b/vllm/transformers_utils/chat_templates/registry.py @@ -32,6 +32,7 @@ def _get_minicpmv_chat_template_fallback(tokenizer_name_or_path: str) -> Path | "deepseek_vl_v2": CHAT_TEMPLATES_DIR / "template_deepseek_vl2.jinja", "minicpmv": _get_minicpmv_chat_template_fallback, "minicpmv4_6": _get_minicpmv_chat_template_fallback, + "minicpmv4_7": _get_minicpmv_chat_template_fallback, "paligemma": CHAT_TEMPLATES_DIR / "template_basic.jinja", "siglip": CHAT_TEMPLATES_DIR / "template_basic.jinja", "siglip2": CHAT_TEMPLATES_DIR / "template_basic.jinja", From bd5bbce6e45e09dbd07ac3e8c394626c059078f7 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sun, 11 Oct 2026 21:15:14 +0800 Subject: [PATCH 260/624] [RL] Release weight-transfer communicators without stopping inference (#61060) Signed-off-by: aoshen02 Signed-off-by: fh3 --- rust/src/engine-core-client/src/client.rs | 12 ++++++++ rust/src/server/src/routes.rs | 4 +++ rust/src/server/src/routes/tests.rs | 10 +++++-- rust/src/server/src/routes/weight_transfer.rs | 16 +++++++++++ tests/distributed/test_weight_transfer.py | 28 +++++++++++++++++++ tests/distributed/test_weight_transfer_m2n.py | 26 ++++++++++++++++- .../test_route_request_bodies.py | 5 ++++ .../rl/weight_transfer/test_offline.py | 7 +++++ .../test_weight_operation_metrics.py | 1 + .../worker/test_gpu_worker_weight_transfer.py | 23 +++++++++++++++ vllm/distributed/weight_transfer/base.py | 4 +-- vllm/distributed/weight_transfer/clients.py | 10 ++++++- .../distributed/weight_transfer/m2n_engine.py | 4 ++- .../weight_transfer/m2n_trainer.py | 4 ++- .../weight_transfer/nccl_engine.py | 6 ++-- .../weight_transfer/sparse_nccl_engine.py | 5 +++- vllm/engine/protocol.py | 4 +++ vllm/entrypoints/rl/offline.py | 4 +++ vllm/entrypoints/rl/online/api_router.py | 7 +++++ vllm/entrypoints/rl/online/metrics.py | 2 +- vllm/v1/engine/async_llm.py | 4 +++ vllm/v1/worker/gpu_worker.py | 9 ++++++ 22 files changed, 181 insertions(+), 14 deletions(-) diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index d8b436c6466c..49910d80c740 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -828,6 +828,18 @@ impl EngineCoreClient { Ok(()) } + /// Release the trainer connection without shutting down inference. + pub async fn shutdown_weight_transfer_engine(&self) -> Result<()> { + self.collective_rpc( + "shutdown_weight_transfer_engine", + None, + Vec::::new(), + BTreeMap::::new(), + ) + .await?; + Ok(()) + } + /// Start a weight update for the base model. pub async fn start_weight_update(&self) -> Result<()> { self.collective_rpc( diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs index 251398611ba1..23c4244c821a 100644 --- a/rust/src/server/src/routes.rs +++ b/rust/src/server/src/routes.rs @@ -145,6 +145,10 @@ fn build_router_with_options( "/init_weight_transfer_engine", post(weight_transfer::init_weight_transfer_engine), ) + .route( + "/shutdown_weight_transfer_engine", + post(weight_transfer::shutdown_weight_transfer_engine), + ) .route( "/start_weight_update", post(weight_transfer::start_weight_update), diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index d574e02a9215..9b8378f89934 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -6708,6 +6708,7 @@ async fn admin_routes_are_hidden_when_dev_mode_is_disabled() { ("POST", "/reset_mm_cache"), ("POST", "/reset_encoder_cache"), ("POST", "/init_weight_transfer_engine"), + ("POST", "/shutdown_weight_transfer_engine"), ("POST", "/start_weight_update"), ("POST", "/start_draft_weight_update"), ("POST", "/update_weights"), @@ -6739,7 +6740,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() { let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| { boxed_test_future(async move { let mut calls = Vec::new(); - for _ in 0..10 { + for _ in 0..11 { let utility = recv_engine_message(dealer).await; assert_eq!(utility[0].as_ref(), &[0x03]); let payload = decode_value(&utility[1]).expect("decode utility"); @@ -6769,6 +6770,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() { "collective_rpc [\"finish_weight_update\",null,[],{}]", "set_weight_version [\"step-2\"]", "get_weight_version []", + "collective_rpc [\"shutdown_weight_transfer_engine\",null,[],{}]", ] "#]].assert_debug_eq(&calls); }) @@ -6817,6 +6819,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() { Some(json!({"new_version": "step-2"})), ), ("GET", "/weight_info", None), + ("POST", "/shutdown_weight_transfer_engine", None), ] { let mut request = Request::builder().method(method).uri(path); let body = match body { @@ -6838,7 +6841,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() { } engine_task.await.expect("mock engine task"); let metrics_after = METRICS.render().unwrap(); - for operation in ["update", "finish"] { + for (operation, count) in [("update", 2.0), ("finish", 2.0), ("shutdown", 1.0)] { assert_eq!( metric_delta( &metrics_before, @@ -6846,7 +6849,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() { "vllm:rl_weight_update_operation_duration_seconds_count", Some(&format!("operation=\"{operation}\"")), ), - 2.0 + count ); assert_eq!( metric_value( @@ -6868,6 +6871,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() { "{\"message\":\"Weight update finished\"}", "{\"success\":true,\"new_version\":\"step-2\"}", "{\"weight_version\":\"step-2\"}", + "{\"message\":\"Weight transfer shut down\"}", ] "#]] .assert_debug_eq(&responses); diff --git a/rust/src/server/src/routes/weight_transfer.rs b/rust/src/server/src/routes/weight_transfer.rs index 5ccd66e39b68..7e1d6af61c7e 100644 --- a/rust/src/server/src/routes/weight_transfer.rs +++ b/rust/src/server/src/routes/weight_transfer.rs @@ -75,6 +75,22 @@ pub async fn init_weight_transfer_engine( })) } +/// Release the trainer connection without shutting down inference. +pub async fn shutdown_weight_transfer_engine( + State(state): State>, +) -> Result, ApiError> { + let _recorder = METRICS.api_server.record_weight_operation("shutdown"); + state + .engine_core_client() + .shutdown_weight_transfer_engine() + .await + .map_err(|error| utility_call_error("shutdown_weight_transfer_engine", error))?; + + Ok(Json(MessageResponse { + message: "Weight transfer shut down", + })) +} + /// Start a weight update transaction. pub async fn start_weight_update( State(state): State>, diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index d044ef3bb7ac..642193ce882f 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -218,6 +218,27 @@ def test_rejects_negative_num_updates(self): # --- Unit Tests: Engine Parsing --- +@pytest.mark.parametrize( + "engine_type", + [ + NCCLWeightTransferEngine, + NCCLTrainerWeightTransferEngine, + SparseNCCLWeightTransferEngine, + SparseNCCLTrainerWeightTransferEngine, + ], +) +def test_nccl_shutdown_destroys_communicator_once(engine_type): + engine = object.__new__(engine_type) + group = MagicMock() + engine.model_update_group = group + + engine.shutdown() + engine.shutdown() + + group.destroy.assert_called_once_with() + assert engine.model_update_group is None + + class TestNCCLEngineParsing: """Test NCCLWeightTransferEngine parsing methods.""" @@ -1292,6 +1313,9 @@ def test_ray_client_sends_typed_requests(self, monkeypatch): handle.finish_weight_update.remote.assert_called_once_with() handle.update_weight_version.remote.assert_called_once_with("step-42") + client.shutdown_weight_transfer_engine() + handle.shutdown_weight_transfer_engine.remote.assert_called_once_with() + def test_http_client_pickles_ipc_handles_for_json(self, monkeypatch): """HTTP update_weights must encode raw ipc_handles as a base64 pickle.""" captured = {} @@ -1322,6 +1346,7 @@ def test_http_client_passes_through_nccl_update_info(self, monkeypatch): captured = {} def fake_post(self, path, json=None): + captured["path"] = path captured["json"] = json monkeypatch.setattr(HTTPVLLMWeightSyncClient, "_post", fake_post) @@ -1333,6 +1358,9 @@ def fake_post(self, path, json=None): client.finish_weight_update("step-42") assert captured["json"] == {"weight_version": "step-42"} + client.shutdown_weight_transfer_engine() + assert captured == {"path": "shutdown_weight_transfer_engine", "json": None} + class TestModuleSource: """`ModuleSource` metadata vs. materialized iteration (dense, no GPU).""" diff --git a/tests/distributed/test_weight_transfer_m2n.py b/tests/distributed/test_weight_transfer_m2n.py index 0591adff4d8d..588ffeeb264b 100644 --- a/tests/distributed/test_weight_transfer_m2n.py +++ b/tests/distributed/test_weight_transfer_m2n.py @@ -46,7 +46,10 @@ mesh_from_tensor, placements_from_tensor, ) -from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo +from vllm.distributed.weight_transfer.m2n_trainer import ( + M2NTrainerInitInfo, + M2NTrainerWeightTransferEngine, +) from vllm.model_executor.layers.linear import ( ColumnParallelLinear, RowParallelLinear, @@ -57,6 +60,27 @@ VALID_UID_B64 = base64.b64encode(b"\x00" * 128).decode() +@pytest.mark.parametrize( + "engine_type,group_attribute", + [ + (M2NWeightTransferEngine, "model_update_group"), + (M2NTrainerWeightTransferEngine, "group"), + ], +) +def test_shutdown_destroys_communicator_once(engine_type, group_attribute): + engine = object.__new__(engine_type) + engine._handle = None + engine._executor = None + group = Mock() + setattr(engine, group_attribute, group) + + engine.shutdown() + engine.shutdown() + + group.destroy.assert_called_once_with() + assert getattr(engine, group_attribute) is None + + class TestLayout: def test_replicated_splits_nothing(self): """A replicated tensor imposes no divisibility constraint. The shape diff --git a/tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py b/tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py index b51d08db0758..9e4df28d03b4 100644 --- a/tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py +++ b/tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py @@ -89,6 +89,11 @@ def test_invalid_field_is_rejected_before_the_engine(client, path, body): ("/update_weights", {"update_info": {"names": []}}, "update_weights"), ("/update_weights", {"update_info": [{}, {}]}, "update_weights"), ("/collective_rpc", {"method": "m"}, "collective_rpc"), + ( + "/shutdown_weight_transfer_engine", + {}, + "shutdown_weight_transfer_engine", + ), ], ) def test_valid_body_reaches_the_engine(client, path, body, engine_call): diff --git a/tests/entrypoints/rl/weight_transfer/test_offline.py b/tests/entrypoints/rl/weight_transfer/test_offline.py index e4a30f17e61d..0af71a1dee24 100644 --- a/tests/entrypoints/rl/weight_transfer/test_offline.py +++ b/tests/entrypoints/rl/weight_transfer/test_offline.py @@ -169,6 +169,13 @@ def check_init_called(self): assert called, "init_transfer_engine should have been called" assert param == "hello", f"Expected 'hello', got {param}" + llm.shutdown_weight_transfer_engine() + assert all( + llm.collective_rpc( + lambda worker: worker.weight_transfer_engine.shutdown_called + ) + ) + @create_new_process_for_each_test() def test_update_weights_calls_engine(vllm_runner): diff --git a/tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py b/tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py index 22ec5e611252..a3bb67ed2b60 100644 --- a/tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py +++ b/tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py @@ -75,6 +75,7 @@ def test_record_releases_in_flight_and_observes_duration(registry, fail): "operation,route,body", [ ("init", api_router.init_weight_transfer_engine, {"init_info": {}}), + ("shutdown", api_router.shutdown_weight_transfer_engine, None), ("start", api_router.start_weight_update, None), ("start_draft", api_router.start_draft_weight_update, None), ("update", api_router.update_weights, {"update_info": {}}), diff --git a/tests/v1/worker/test_gpu_worker_weight_transfer.py b/tests/v1/worker/test_gpu_worker_weight_transfer.py index f1965cb1fae2..ec63c08c70c4 100644 --- a/tests/v1/worker/test_gpu_worker_weight_transfer.py +++ b/tests/v1/worker/test_gpu_worker_weight_transfer.py @@ -24,6 +24,7 @@ def __init__(self, raise_on_update: bool = False): self.raise_on_update = raise_on_update self.started = False self.finished = False + self.closed = False self.reset_count = 0 self.supports_draft_weight_update = False self.update_calls: list[dict] = [] @@ -49,6 +50,9 @@ def finish_weight_update(self) -> None: def reset_weight_update_target(self) -> None: self.reset_count += 1 + def shutdown(self) -> None: + self.closed = True + class _RecordingModelRunner: def __init__(self) -> None: @@ -82,6 +86,25 @@ def test_reload_weights_sets_current_config(): assert model_runner.seen_config is worker.vllm_config +@pytest.mark.parametrize( + "enabled,active", [(True, False), (True, True), (False, False)] +) +def test_shutdown_weight_transfer_engine_preserves_inference(enabled, active): + engine = _RecordingEngine() if enabled else None + worker = _make_worker(engine) + worker._weight_update_active = active + model_runner = worker.model_runner + if active: + with pytest.raises(RuntimeError, match="during a weight update"): + worker.shutdown_weight_transfer_engine() + else: + worker.shutdown_weight_transfer_engine() + if engine is not None: + assert engine.closed is not active + assert worker.weight_transfer_engine is engine + assert worker.model_runner is model_runner + + def test_reload_parameter_lookup_preserves_lora_module_names(): base_layer = nn.Module() qweight = nn.Parameter(torch.ones(1)) diff --git a/vllm/distributed/weight_transfer/base.py b/vllm/distributed/weight_transfer/base.py index 3140b70e9656..6eba3fa96f51 100644 --- a/vllm/distributed/weight_transfer/base.py +++ b/vllm/distributed/weight_transfer/base.py @@ -514,9 +514,7 @@ def receive_weights(self, update_info: TUpdateInfo) -> None: @abstractmethod def shutdown(self) -> None: - """Shutdown the weight transfer engine. - This should be called when the worker is shutting down. - """ + """Release transfer resources without shutting down the inference worker.""" raise NotImplementedError diff --git a/vllm/distributed/weight_transfer/clients.py b/vllm/distributed/weight_transfer/clients.py index 745e67779e0a..0ccce60f973a 100644 --- a/vllm/distributed/weight_transfer/clients.py +++ b/vllm/distributed/weight_transfer/clients.py @@ -60,7 +60,7 @@ class HTTPVLLMWeightSyncClient: Mirrors `vllm/entrypoints/rl/online/api_router.py`: `/init_weight_transfer_engine`, `/start_weight_update`, `/update_weights`, - `/finish_weight_update`. + `/finish_weight_update`, `/shutdown_weight_transfer_engine`. """ def __init__(self, base_url: str, timeout: float = 300) -> None: @@ -78,6 +78,9 @@ def _post(self, path: str, json: dict[str, Any] | None = None) -> None: def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None: self._post("init_weight_transfer_engine", {"init_info": init_info}) + def shutdown_weight_transfer_engine(self) -> None: + self._post("shutdown_weight_transfer_engine") + def start_weight_update(self) -> None: self._post("start_weight_update") @@ -109,6 +112,11 @@ def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None: request = WeightTransferInitRequest(init_info=init_info) ray.get([h.init_weight_transfer_engine.remote(request) for h in self.handles]) + def shutdown_weight_transfer_engine(self) -> None: + import ray + + ray.get([h.shutdown_weight_transfer_engine.remote() for h in self.handles]) + def start_weight_update(self) -> None: import ray diff --git a/vllm/distributed/weight_transfer/m2n_engine.py b/vllm/distributed/weight_transfer/m2n_engine.py index c61bac4f2d1a..e480e60eaf91 100644 --- a/vllm/distributed/weight_transfer/m2n_engine.py +++ b/vllm/distributed/weight_transfer/m2n_engine.py @@ -449,5 +449,7 @@ def shutdown(self) -> None: torch.accelerator.synchronize() self._handle.destroy() self._handle = None - self.model_update_group = None + if self.model_update_group is not None: + self.model_update_group.destroy() + self.model_update_group = None self._parameter_destinations = [] diff --git a/vllm/distributed/weight_transfer/m2n_trainer.py b/vllm/distributed/weight_transfer/m2n_trainer.py index 2d1d9f47ec02..8b9b73f9f9cc 100644 --- a/vllm/distributed/weight_transfer/m2n_trainer.py +++ b/vllm/distributed/weight_transfer/m2n_trainer.py @@ -410,4 +410,6 @@ def shutdown(self) -> None: if self._executor is not None: self._executor.shutdown() self._executor = None - self.group = None + if self.group is not None: + self.group.destroy() + self.group = None diff --git a/vllm/distributed/weight_transfer/nccl_engine.py b/vllm/distributed/weight_transfer/nccl_engine.py index 77d8b608727e..331cb5eacc01 100644 --- a/vllm/distributed/weight_transfer/nccl_engine.py +++ b/vllm/distributed/weight_transfer/nccl_engine.py @@ -229,7 +229,7 @@ def state_dict_info_iterator(): def shutdown(self) -> None: if self.model_update_group is not None: - # Clean up the communicator by removing the reference + self.model_update_group.destroy() self.model_update_group = None @@ -456,4 +456,6 @@ def _post_send_sync(self) -> None: torch.cuda.current_stream().synchronize() def shutdown(self) -> None: - self.model_update_group = None + if self.model_update_group is not None: + self.model_update_group.destroy() + self.model_update_group = None diff --git a/vllm/distributed/weight_transfer/sparse_nccl_engine.py b/vllm/distributed/weight_transfer/sparse_nccl_engine.py index 01b29529deb7..f9233516a56e 100644 --- a/vllm/distributed/weight_transfer/sparse_nccl_engine.py +++ b/vllm/distributed/weight_transfer/sparse_nccl_engine.py @@ -186,6 +186,7 @@ def receive_weights(self, update_info: SparseNCCLWeightTransferUpdateInfo) -> No def shutdown(self) -> None: if self.model_update_group is not None: + self.model_update_group.destroy() self.model_update_group = None @@ -371,4 +372,6 @@ def _post_send_sync(self) -> None: torch.cuda.current_stream().synchronize() def shutdown(self) -> None: - self.model_update_group = None + if self.model_update_group is not None: + self.model_update_group.destroy() + self.model_update_group = None diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index 15f6b801bcd6..15f8f8b744f0 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -296,6 +296,10 @@ async def init_weight_transfer_engine( """Initialize weight transfer for RL training.""" raise NotImplementedError + async def shutdown_weight_transfer_engine(self) -> None: + """Release the trainer connection without shutting down inference.""" + raise NotImplementedError + async def start_weight_update(self) -> None: """Start a new weight update.""" raise NotImplementedError diff --git a/vllm/entrypoints/rl/offline.py b/vllm/entrypoints/rl/offline.py index 1db6a82533dd..be72a9a5540c 100644 --- a/vllm/entrypoints/rl/offline.py +++ b/vllm/entrypoints/rl/offline.py @@ -26,6 +26,10 @@ def init_weight_transfer_engine( "init_weight_transfer_engine", kwargs={"init_info": init_info_dict} ) + def shutdown_weight_transfer_engine(self) -> None: + """Release the trainer connection without shutting down inference.""" + self.llm_engine.collective_rpc("shutdown_weight_transfer_engine") + def start_weight_update(self) -> None: """Start a new weight update.""" self.llm_engine.collective_rpc("start_weight_update") diff --git a/vllm/entrypoints/rl/online/api_router.py b/vllm/entrypoints/rl/online/api_router.py index 51182fd99683..847287336ee5 100644 --- a/vllm/entrypoints/rl/online/api_router.py +++ b/vllm/entrypoints/rl/online/api_router.py @@ -168,6 +168,13 @@ async def init_weight_transfer_engine( return JSONResponse(content={"message": "Weight transfer initialized"}) +@router.post("/shutdown_weight_transfer_engine") +async def shutdown_weight_transfer_engine(raw_request: Request): + with weight_operation_metrics().record("shutdown"): + await engine_client(raw_request).shutdown_weight_transfer_engine() + return JSONResponse(content={"message": "Weight transfer shut down"}) + + @router.post("/start_weight_update") async def start_weight_update(raw_request: Request): with weight_operation_metrics().record("start"): diff --git a/vllm/entrypoints/rl/online/metrics.py b/vllm/entrypoints/rl/online/metrics.py index ed16f9a820b5..b67f479f6dc0 100644 --- a/vllm/entrypoints/rl/online/metrics.py +++ b/vllm/entrypoints/rl/online/metrics.py @@ -10,7 +10,7 @@ from prometheus_client import REGISTRY, CollectorRegistry, Gauge, Histogram -Operation = Literal["init", "start", "start_draft", "update", "finish"] +Operation = Literal["init", "shutdown", "start", "start_draft", "update", "finish"] class WeightOperationMetrics: diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 71e30990c578..60588e076097 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -1318,6 +1318,10 @@ async def init_weight_transfer_engine( "init_weight_transfer_engine", kwargs={"init_info": request.init_info} ) + async def shutdown_weight_transfer_engine(self) -> None: + """Release the trainer connection without shutting down inference.""" + await self.collective_rpc("shutdown_weight_transfer_engine") + async def start_weight_update(self) -> None: """Start a new weight update.""" await self.collective_rpc("start_weight_update") diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 152def62221c..3efe77a3a4d2 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -1482,6 +1482,15 @@ def init_weight_transfer_engine(self, init_info: dict) -> None: typed_init_info = self.weight_transfer_engine.parse_init_info(init_info) self.weight_transfer_engine.init_transfer_engine(typed_init_info) + def shutdown_weight_transfer_engine(self) -> None: + """Release a trainer connection without shutting down inference.""" + if self._weight_update_active: + raise RuntimeError( + "Cannot release a trainer connection during a weight update" + ) + if self.weight_transfer_engine is not None: + self.weight_transfer_engine.shutdown() + def start_weight_update(self) -> None: """Start a new weight update session. From 3e8e1364ff41330fb667b2d76d6645176d61fdd0 Mon Sep 17 00:00:00 2001 From: Zheng Cai Date: Sun, 11 Oct 2026 22:43:37 +0800 Subject: [PATCH 261/624] [Perf][GLM-5.3-Flash] Shard long-context indexer prefill rows across TP (#54951) Signed-off-by: Zheng Cai <8370601+zigzagcai@users.noreply.github.com> Signed-off-by: Zheng Cai Signed-off-by: Codex Co-authored-by: root Co-authored-by: OpenAI Codex Co-authored-by: zigzagcai Co-authored-by: Claude Code Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- .../v1/attention/test_indexer_tp_row_shard.py | 213 ++++++++++++++++++ .../layers/sparse_attn_indexer.py | 124 ++++++++-- vllm/models/glm5next/nvidia/sparse_indexer.py | 63 ++++-- vllm/v1/attention/backends/mla/indexer.py | 74 +++++- 4 files changed, 436 insertions(+), 38 deletions(-) create mode 100644 tests/v1/attention/test_indexer_tp_row_shard.py diff --git a/tests/v1/attention/test_indexer_tp_row_shard.py b/tests/v1/attention/test_indexer_tp_row_shard.py new file mode 100644 index 000000000000..e65e4789f132 --- /dev/null +++ b/tests/v1/attention/test_indexer_tp_row_shard.py @@ -0,0 +1,213 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TP query-row sharding of the DSA indexer prefill (see PrefillRowShard): +each rank scores a contiguous slice; the group exchanges int32s per row.""" + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer +import vllm.v1.attention.backends.mla.indexer as indexer +from vllm.config import CUDAGraphMode +from vllm.model_executor.layers.sparse_attn_indexer import PrefillRowShard + + +def _chunk(start: int, end: int) -> SimpleNamespace: + """Chunk over rows [start, end); bounds derive from the global row.""" + n = end - start + return SimpleNamespace( + token_start=start, + token_end=end, + cu_seqlen_ks=torch.zeros(n, dtype=torch.int32), + cu_seqlen_ke=((torch.arange(n) * 7 + start) % 65).int(), + ) + + +@pytest.mark.parametrize("world", [2, 3, 8]) +def test_shard_narrowing_covers_every_row_once(world: int) -> None: + """Narrowed windows tile the prefill rows exactly once, carry the causal + bounds of those global rows, and skip chunks foreign to this rank.""" + decode, rows = 5, 997 + edges = [decode, decode + 330, decode + 337, decode + rows] + chunks = [_chunk(lo, hi) for lo, hi in zip(edges, edges[1:])] + sizes = [rows // world + int(r < rows % world) for r in range(world)] + covered = [] + for rank in range(world): + shard = PrefillRowShard(decode, sizes, rank) + for chunk in chunks: + window = shard.narrow(chunk) + if window is None: + continue + start, end, ks, ke = window + lo, hi = start - chunk.token_start, end - chunk.token_start + torch.testing.assert_close(ks, chunk.cu_seqlen_ks[lo:hi]) + torch.testing.assert_close(ke, chunk.cu_seqlen_ke[lo:hi]) + covered.append((start, end)) + covered.sort() + # Windows must be disjoint, contiguous and exactly cover the prefill rows. + assert [s for s, _ in covered] == [decode, *[e for _, e in covered][:-1]] + assert covered[-1][1] == decode + rows + + +def test_shard_exchange_reassembles_rows_only( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Gathered indices land only on real prefill rows: leading decode rows + and trailing CUDA-graph padding stay untouched, and the full k-pool + width (tail columns included) is exchanged.""" + decode, pad, topk, kpool = 5, 3, 8, 4 + width = topk + kpool - 1 + sizes = [17, 9, 21] + buffer = torch.full((decode + sum(sizes) + pad, width), -7, dtype=torch.int32) + gathered = torch.arange(sum(sizes) * width, dtype=torch.int32).reshape(-1, width) + calls = [] + + def fake_group(): + def all_gatherv(local, dim, sizes): + calls.append((local, sizes)) + return gathered + + return SimpleNamespace(all_gatherv=all_gatherv) + + monkeypatch.setattr(sparse_indexer, "get_tp_group", fake_group) + for rank in range(len(sizes)): + PrefillRowShard(decode, sizes, rank).exchange_topk(buffer, width) + assert calls[0][1] == sizes and calls[0][0].is_contiguous() + torch.testing.assert_close(buffer[decode : decode + sum(sizes)], gathered) + assert torch.all(buffer[:decode] == -7), "decode rows stay outside" + assert torch.all(buffer[decode + sum(sizes) :] == -7), "padding untouched" + + +@pytest.mark.parametrize("tp_size", [2, 4]) +def test_balanced_row_shard_partitions_and_declines(tp_size: int) -> None: + """The split covers every row once with >= 1 row per rank and declines + below the activation floor (or on a single rank).""" + n = indexer.MIN_TP_SHARD_ROWS_PER_RANK * tp_size + shapes = [ + ([4 * n], [4 * n]), # fresh prompt: the full causal cost ramp + # uneven requests at different context depths + ([2 * n, 7, 4 * n, n // 3 + 5], [2 * n, 9007, 4 * n + 500, n // 3 + 60005]), + ] + ratio = 4 + for query_lens, seq_lens in shapes: + sizes = indexer.balanced_prefill_row_shard( + torch.tensor(seq_lens, dtype=torch.int32), + torch.tensor(query_lens, dtype=torch.int32), + ratio, + tp_size, + ) + assert sizes is not None and len(sizes) == tp_size + assert sum(sizes) == sum(query_lens) and min(sizes) >= 1 + lens = torch.tensor([n - 1], dtype=torch.int32) + assert indexer.balanced_prefill_row_shard(lens, lens, ratio, tp_size) is None + assert indexer.balanced_prefill_row_shard(lens, lens, ratio, 1) is None + + +def test_row_sharding_gate_rejects_unsupported_configurations( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The gate is the whole safety envelope; nothing else guards the exchange.""" + monkeypatch.setattr(indexer.current_platform, "is_cuda", lambda: True) + + def supported(**kwargs): + mode = kwargs.pop("cudagraph_mode", CUDAGraphMode.PIECEWISE) + cfg = SimpleNamespace(compilation_config=SimpleNamespace(cudagraph_mode=mode)) + args = {"dcp_world_size": 1, "use_pcp": False, "tp_size": 4, **kwargs} + return indexer.tp_prefill_row_sharding_supported(cfg, **args) + + assert supported() + for bad in ({"tp_size": 1}, {"dcp_world_size": 2}, {"use_pcp": True}): + assert not supported(**bad) + for name in ( + "VLLM_DISABLE_PYNCCL", + "VLLM_USE_NCCL_SYMM_MEM", + "VLLM_BATCH_INVARIANT", + ): + with monkeypatch.context() as m: + m.setattr(indexer.envs, name, True) + assert not supported() + # FULL captures mixed batches whole, so the exchange must stay out. + assert not supported(cudagraph_mode=CUDAGraphMode.FULL) + assert supported(cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE) + + +def test_profiling_run_reserves_exchange_transient( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The dummy run returns before the exchange executes, so its all_gatherv + output plus staging copy must be reserved up front when row sharding is + active -- and only then.""" + rows, topk = 128, 8 + buffer = torch.empty((rows, topk), dtype=torch.int32) + peak = 2 * rows * topk * buffer.element_size() + specs = (((1,), torch.uint8), ((1,), torch.uint8)) + numels: list[int] = [] + real_empty = torch.empty + + def recording_empty(*args, **kwargs): + out = real_empty(*args, **kwargs) + numels.append(out.numel()) + return out + + def run_dummy(tp_world: int) -> None: + numels.clear() + with monkeypatch.context() as m: + m.setattr(torch, "empty", recording_empty) + m.setattr( + sparse_indexer, + "get_forward_context", + lambda: SimpleNamespace(attn_metadata=None), + ) + m.setattr(sparse_indexer, "_gather_workspace_shapes", lambda *a: specs) + m.setattr( + sparse_indexer, + "current_workspace_manager", + lambda: SimpleNamespace(get_simultaneous=lambda *s: None), + ) + m.setattr( + sparse_indexer, + "get_tp_group", + lambda: SimpleNamespace(world_size=tp_world), + ) + m.setattr( + sparse_indexer, + "get_current_vllm_config_or_none", + lambda: SimpleNamespace( + compilation_config=SimpleNamespace( + cudagraph_mode=CUDAGraphMode.PIECEWISE + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + prefill_context_parallel_size=1, + ), + ), + ) + m.setattr(sparse_indexer.envs, "VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", 1) + m.setattr(indexer.current_platform, "is_cuda", lambda: True) + sparse_indexer.sparse_attn_indexer( + hidden_states=torch.empty(4, 4), + k_cache_prefix="k", + kv_cache=torch.empty(4, 4), + q_quant=torch.empty(4, 4), + q_scale=None, + k=None, + weights=torch.empty(4), + quant_block_size=128, + scale_fmt=None, + topk_tokens=topk, + head_dim=64, + max_model_len=1024, + total_seq_lens=512, + topk_indices_buffer=buffer, + skip_k_cache_insert=True, + use_pcp=False, + pcp_shard_decode_requests=False, + dense_mha_metadata_layer_name="", + ) + + run_dummy(tp_world=2) + assert peak in numels + run_dummy(tp_world=1) + assert peak not in numels diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index f8fedf046aca..935e4c0d2d44 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -9,8 +9,17 @@ from vllm import _custom_ops as ops from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture -from vllm.config import CUDAGraphMode, get_current_vllm_config -from vllm.distributed import get_dcp_group, get_pcp_group +from vllm.config import ( + CUDAGraphMode, + get_current_vllm_config, + get_current_vllm_config_or_none, +) +from vllm.distributed import ( + get_dcp_group, + get_pcp_group, + get_tensor_model_parallel_rank, + get_tp_group, +) from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp @@ -43,6 +52,7 @@ ) from vllm.v1.attention.backends.mla.indexer import ( DeepseekV32IndexerMetadata, + tp_prefill_row_sharding_supported, ) from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton from vllm.v1.attention.ops.pcp import maybe_gather_indexer_k @@ -317,6 +327,75 @@ def kv_cache_as_quant_view( return kv_cache.unsqueeze(-2) +class PrefillRowShard: + """One TP rank's contiguous slice of the replicated prefill rows. + + Rows are independent, so each rank scores only its slice and the group + exchanges final int32 top-k indices instead of the logits.""" + + def __init__(self, num_decode_tokens: int, sizes: list[int], rank: int): + self.num_decode_tokens = num_decode_tokens + self.sizes = sizes + self.start = num_decode_tokens + sum(sizes[:rank]) + self.stop = self.start + sizes[rank] + + @classmethod + def from_metadata( + cls, prefill_metadata, num_decode_tokens: int + ) -> "PrefillRowShard | None": + sizes = getattr(prefill_metadata, "row_shard_sizes", None) + if sizes is None: + return None + return cls(num_decode_tokens, sizes, get_tensor_model_parallel_rank()) + + def narrow(self, chunk) -> "tuple[int, int, torch.Tensor, torch.Tensor] | None": + """Row window and narrowed causal bounds of ``chunk`` for this rank.""" + start = max(chunk.token_start, self.start) + end = min(chunk.token_end, self.stop) + if start >= end: + return None + lo, hi = start - chunk.token_start, end - chunk.token_start + return start, end, chunk.cu_seqlen_ks[lo:hi], chunk.cu_seqlen_ke[lo:hi] + + def exchange_topk(self, topk_buffer: torch.Tensor, width: int) -> None: + """Gather per-rank final indices back into global row order.""" + # The local input must outlive the async PyNCCL gather's enqueued copy. + local = topk_buffer[self.start : self.stop, :width].contiguous() + end = self.num_decode_tokens + sum(self.sizes) + topk_buffer[self.num_decode_tokens : end, :width] = get_tp_group().all_gatherv( + local, dim=0, sizes=self.sizes + ) + + +def _reserve_row_shard_exchange_transient( + topk_indices_buffer: torch.Tensor, device: torch.device +) -> None: + """Profiling/dummy runs return before ``exchange_topk`` executes; reserve + its worst-case transient so the memory profile accounts for it. + + The exchange stages this rank's contiguous slice and the all_gatherv + wrapper allocates the full gathered tensor; each is bounded by the + buffer's row capacity times its (tile-padded) width. + """ + cfg = get_current_vllm_config_or_none() + if cfg is None: + return + parallel = cfg.parallel_config + if not tp_prefill_row_sharding_supported( + cfg, + parallel.decode_context_parallel_size, + parallel.prefill_context_parallel_size > 1, + get_tp_group().world_size, + ): + return + rows, width = topk_indices_buffer.shape[:2] + _ = torch.empty( + 2 * rows * width * topk_indices_buffer.element_size(), + dtype=torch.uint8, + device=device, + ) + + @eager_break_during_capture def sparse_attn_indexer( hidden_states: torch.Tensor, @@ -388,6 +467,10 @@ def sparse_attn_indexer( max_logits_elems, dtype=torch.uint8, device=hidden_states.device ) + # The dummy run returns before exchange_topk executes, so reserve its + # worst-case transient for the memory profile. + _reserve_row_shard_exchange_transient(topk_indices_buffer, hidden_states.device) + return sparse_attn_indexer_fake( hidden_states, k_cache_prefix, @@ -515,6 +598,12 @@ def sparse_attn_indexer( scales_spec, *gather_specs, ) + # One row window per rank; DSV4.1 candidate rows stay replicated. + row_shard = ( + None + if candidate_blocks is not None + else PrefillRowShard.from_metadata(prefill_metadata, num_decode_tokens) + ) for chunk in prefill_metadata.chunks: cu_seqlen_ks = chunk.cu_seqlen_ks cu_seqlen_ke = chunk.cu_seqlen_ke @@ -530,6 +619,15 @@ def sparse_attn_indexer( chunk.local_cu_seq_lens, ) + # Narrow scoring to this rank; KV gathers remain unconditional for + # continuation chunks that reuse the workspace. + row_start, row_end = chunk.token_start, chunk.token_end + if row_shard is not None: + narrowed = row_shard.narrow(chunk) + if narrowed is None: + continue + row_start, row_end, cu_seqlen_ks, cu_seqlen_ke = narrowed + # PCP + DCP KV all-gather. deinterleave_idx = chunk.pcp_deinterleave_idx if deinterleave_idx is not None: @@ -547,15 +645,9 @@ def sparse_attn_indexer( out_scales, ) - q_slice = q_quant[chunk.token_start : chunk.token_end] - q_scale_slice = ( - q_scale[chunk.token_start : chunk.token_end] - if q_scale is not None - else None - ) - topk_indices = topk_indices_buffer[ - chunk.token_start : chunk.token_end, :topk_tokens - ] + q_slice = q_quant[row_start:row_end] + q_scale_slice = q_scale[row_start:row_end] if q_scale is not None else None + topk_indices = topk_indices_buffer[row_start:row_end, :topk_tokens] if chunk.local_total_seq_lens == 0 or q_slice.shape[0] == 0: logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32) @@ -578,7 +670,7 @@ def sparse_attn_indexer( q_slice_cast, k_quant_cast, k_scale_cast, - weights[chunk.token_start : chunk.token_end], + weights[row_start:row_end], cu_seqlen_ks, cu_seqlen_ke, ) @@ -586,7 +678,7 @@ def sparse_attn_indexer( logits = fp8_fp4_mqa_logits( (q_slice_cast, q_scale_slice), (k_quant_cast, k_scale_cast), - weights[chunk.token_start : chunk.token_end], + weights[row_start:row_end], cu_seqlen_ks, cu_seqlen_ke, clean_logits=False, @@ -637,9 +729,13 @@ def sparse_attn_indexer( dcp_rank, dcp_world_size, cp_kv_cache_interleave_size, - row_starts=chunk.cu_seqlen_ks, + row_starts=cu_seqlen_ks, ) + if row_shard is not None: + # Rows are already fully ranked, so this is a layout gather. + row_shard.exchange_topk(topk_indices_buffer, topk_tokens) + if has_decode: decode_metadata = attn_metadata_narrowed.decode assert decode_metadata is not None diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 77fa4c409152..154fdcda900e 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -12,7 +12,11 @@ from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.layers.indexer_topk import get_indexer_topk -from vllm.model_executor.layers.sparse_attn_indexer import _merge_dcp_topk_global +from vllm.model_executor.layers.sparse_attn_indexer import ( + PrefillRowShard, + _merge_dcp_topk_global, + _reserve_row_shard_exchange_transient, +) from vllm.models.glm5next.common.sparse_indexer import ( RADIX_TOPK_WORKSPACE_SIZE, _build_decode_scatter_indices, @@ -169,6 +173,10 @@ def sparse_attn_indexer_kpool( max_logits_elems, dtype=torch.uint8, device=hidden_states.device ) + # The dummy run returns before exchange_topk executes, so reserve its + # worst-case transient for the memory profile. + _reserve_row_shard_exchange_transient(topk_indices_buffer, hidden_states.device) + return topk_indices_buffer attn_metadata_narrowed = attn_metadata[k_cache_prefix] assert isinstance(attn_metadata_narrowed, DeepseekV32IndexerMetadata) @@ -282,6 +290,12 @@ def sparse_attn_indexer_kpool( _pos = positions[num_decode_tokens:num_tokens].to(torch.int32) _buf = topk_indices_buffer[num_decode_tokens:num_tokens] _fill_causal_indices(_buf, _pos) + # Short prefills score nothing locally; suppress the shard there. + row_shard = None + if not short_prefill: + row_shard = PrefillRowShard.from_metadata( + prefill_metadata, num_decode_tokens + ) # Get the full shared workspace buffers once (will allocate on first use). # Layout switches between FP8 (head_dim bytes + 4-byte fp32 scale) and @@ -310,12 +324,16 @@ def sparse_attn_indexer_kpool( chunk.local_cu_seq_lens, ) - q_slice = q_quant[chunk.token_start : chunk.token_end] - q_scale_slice = ( - q_scale[chunk.token_start : chunk.token_end] - if q_scale is not None - else None - ) + row_start, row_end = chunk.token_start, chunk.token_end + cu_seqlen_ks, cu_seqlen_ke = chunk.cu_seqlen_ks, chunk.cu_seqlen_ke + if row_shard is not None: + narrowed = row_shard.narrow(chunk) + if narrowed is None: + continue + row_start, row_end, cu_seqlen_ks, cu_seqlen_ke = narrowed + + q_slice = q_quant[row_start:row_end] + q_scale_slice = q_scale[row_start:row_end] if q_scale is not None else None # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8 # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf. if use_fp4_cache: @@ -335,9 +353,9 @@ def sparse_attn_indexer_kpool( logits = fp8_fp4_mqa_logits( (q_slice_cast, q_scale_slice), (k_quant_cast, k_scale_cast), - weights[chunk.token_start : chunk.token_end], - chunk.cu_seqlen_ks, - chunk.cu_seqlen_ke, + weights[row_start:row_end], + cu_seqlen_ks, + cu_seqlen_ke, clean_logits=False, ) @@ -351,17 +369,15 @@ def sparse_attn_indexer_kpool( ) topk_dst = pool_topk else: - topk_dst = topk_indices_buffer[ - chunk.token_start : chunk.token_end, :topk_tokens - ] + topk_dst = topk_indices_buffer[row_start:row_end, :topk_tokens] if logits.shape[1] == 0: topk_dst.fill_(-1) else: torch.ops._C.top_k_per_row_prefill( logits, - chunk.cu_seqlen_ks, - chunk.cu_seqlen_ke, + cu_seqlen_ks, + cu_seqlen_ke, topk_dst, num_rows, logits.stride(0), @@ -376,7 +392,7 @@ def sparse_attn_indexer_kpool( dcp_rank, dcp_world_size, pool_cp_interleave, - row_starts=chunk.cu_seqlen_ks, + row_starts=cu_seqlen_ks, ) if index_kpool > 1: @@ -385,10 +401,7 @@ def sparse_attn_indexer_kpool( # Fused expand-pools + append-tail into one Triton kernel # (replaces ~25 elementwise ops). seq_len is token-granular # (pos+1); the kernel derives pool_len internally. - q_seq = ( - positions[chunk.token_start : chunk.token_end].to(torch.int32) - + 1 - ) + q_seq = positions[row_start:row_end].to(torch.int32) + 1 expanded = kpool_ops.expand_pools_and_append_tail( pool_ids, q_seq, index_kpool ) @@ -397,9 +410,13 @@ def sparse_attn_indexer_kpool( expanded = kpool_ops.expand_pools_to_tokens( pool_ids, valid, topk_tokens, index_kpool ) - topk_indices_buffer[ - chunk.token_start : chunk.token_end, : expanded.shape[-1] - ] = expanded + topk_indices_buffer[row_start:row_end, : expanded.shape[-1]] = expanded + + if row_shard is not None: + # The k-pool expansion appends the request's incomplete tail + # (index_kpool - 1 entries) after the logical top-k history; those + # columns are part of the attention index and are exchanged too. + row_shard.exchange_topk(topk_indices_buffer, topk_tokens + index_kpool - 1) if has_decode: decode_metadata = attn_metadata_narrowed.decode diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index e776057b05b7..97de9ffa9ba2 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -7,7 +7,7 @@ import torch import vllm.envs as envs -from vllm.config import VllmConfig, get_current_vllm_config +from vllm.config import CUDAGraphMode, VllmConfig, get_current_vllm_config from vllm.distributed import get_dcp_group, get_pcp_group from vllm.logger import init_logger from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs @@ -204,6 +204,64 @@ def __call__( ) +# Below this many rows per rank, stay replicated (TP4 engages at 64K rows). +MIN_TP_SHARD_ROWS_PER_RANK = 16_384 + + +def balanced_prefill_row_shard( + seq_lens_cpu: torch.Tensor, + query_lens_cpu: torch.Tensor, + compress_ratio: int, + tp_size: int, +) -> list[int] | None: + """Return contiguous TP row counts balanced by indexer MQA work.""" + num_rows = int(query_lens_cpu.sum()) + if tp_size < 2 or num_rows < MIN_TP_SHARD_ROWS_PER_RANK * tp_size: + return None + + query_lens = query_lens_cpu.to(torch.int64) + first_key = torch.repeat_interleave( + seq_lens_cpu.to(torch.int64) - query_lens + 1, query_lens + ) + row_in_request = torch.arange(num_rows) - torch.repeat_interleave( + torch.cumsum(query_lens, 0) - query_lens, query_lens + ) + cost = torch.cumsum((first_key + row_in_request) // compress_ratio, 0) + total = int(cost[-1]) if cost.numel() else 0 + if total <= 0: + return None + + targets = torch.arange(1, tp_size) * total // tp_size + bounds = [0, *torch.searchsorted(cost, targets).tolist(), num_rows] + # Keep >= 1 row per rank (the floor above leaves ample room); this also + # handles repeated costs from pool compression. + for i in range(1, tp_size): + bounds[i] = max(bounds[i], bounds[i - 1] + 1) + for i in range(tp_size - 1, 0, -1): + bounds[i] = min(bounds[i], bounds[i + 1] - 1) + return [bounds[i + 1] - bounds[i] for i in range(tp_size)] + + +def tp_prefill_row_sharding_supported( + vllm_config: VllmConfig, + dcp_world_size: int, + use_pcp: bool, + tp_size: int, +) -> bool: + """Whether replicated prefill rows may be sharded across TP ranks.""" + cudagraph_mode = vllm_config.compilation_config.cudagraph_mode or CUDAGraphMode.NONE + return ( + current_platform.is_cuda() + and dcp_world_size == 1 + and not use_pcp + and tp_size > 1 + and not envs.VLLM_DISABLE_PYNCCL + and not envs.VLLM_USE_NCCL_SYMM_MEM + and not envs.VLLM_BATCH_INVARIANT + and cudagraph_mode.mixed_mode() != CUDAGraphMode.FULL + ) + + class DeepseekV32IndexerBackend(AttentionBackend): @classmethod def supports_pcp(cls) -> bool: @@ -654,6 +712,8 @@ def __call__( class DeepseekV32IndexerPrefillMetadata: chunks: list[DeepseekV32IndexerPrefillChunkMetadata] max_prefill_seq_len: int = -1 + # Per-TP-rank contiguous row counts, or None for the replicated path. + row_shard_sizes: list[int] | None = None @dataclass @@ -941,6 +1001,7 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: scheduler_config = self.vllm_config.scheduler_config parallel_config = self.vllm_config.parallel_config self.dcp_world_size = parallel_config.decode_context_parallel_size + self.tp_size = parallel_config.tensor_parallel_size self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0 self.pcp_world_size = parallel_config.prefill_context_parallel_size self.use_pcp = self.pcp_world_size > 1 @@ -1512,6 +1573,16 @@ def build( # Skip when total_seq_lens is 0 (i.e., no compressed token). if metadata is not None: chunks.append(metadata) + row_shard_sizes = None + if tp_prefill_row_sharding_supported( + self.vllm_config, self.dcp_world_size, self.use_pcp, self.tp_size + ): + row_shard_sizes = balanced_prefill_row_shard( + seq_lens_cpu[num_decodes:], + prefill_query_lens_cpu, + self.compress_ratio, + self.tp_size, + ) prefill_metadata = DeepseekV32IndexerPrefillMetadata( chunks, max_prefill_seq_len=( @@ -1519,6 +1590,7 @@ def build( if num_prefills > 0 else 0 ), + row_shard_sizes=row_shard_sizes, ) decode_metadata = None From 239a3fa3a8ffc6508e867452c115ea7bcee2304b Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sun, 11 Oct 2026 23:30:56 +0800 Subject: [PATCH 262/624] [Sleep] Retain draft-private weights with existing parameter selector (#61046) Signed-off-by: aoshen02 Co-authored-by: Codex --- docs/features/sleep_mode.md | 39 ++++++++++++++++++++++++--- tests/v1/worker/test_gpu_worker.py | 42 ++++++++++++++++++++++++++++++ vllm/config/model.py | 3 +++ vllm/v1/worker/gpu_worker.py | 19 +++++++++++--- 4 files changed, 97 insertions(+), 6 deletions(-) diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index 0db82de9d3c5..bdfef2ecac7d 100644 --- a/docs/features/sleep_mode.md +++ b/docs/features/sleep_mode.md @@ -77,8 +77,15 @@ llm.wake_up(tags=["kv_cache"]) #### Retaining frozen weights during RLHF updates Set `sleep_preserve_parameter_names` (CLI: `--sleep-preserve-parameter-names`) -to runtime parameter-name glob -patterns for weights that stay frozen during training. Each pattern must match +to runtime parameter-name glob patterns for weights that stay frozen during +training. It covers two use cases: + +- **Frozen Engram tables (original use case):** retain table weights and scales + omitted by the trainer, without retransmitting them after every sleep. +- **Frozen draft weights:** retain draft-private parameters when the trainer + updates only the target model. + +Target patterns match `model.named_parameters()`; checkpoint names and `requires_grad` are not used. Level-2 sleep backs up selected GPU parameters to pageable CPU memory; weights @@ -87,7 +94,33 @@ copy. Allow enough host memory for backups. Level-1 behavior is unchanged. The trainer must omit these parameters from updates, and reload/post-processing must preserve their values and storage (including avoiding replacement with meta -tensors). This option does not filter updates and applies only to the target model. +tensors). This option does not filter updates. + +For frozen Engram tables omitted by the trainer, select the target table weights +and scales. For example, DeepSeek V4.1 uses these runtime parameter names: + +```bash +VLLM_SERVER_DEV_MODE=1 vllm serve deepseek-ai/DeepSeek-V4.1-Flash --enable-sleep-mode \ + --sleep-preserve-parameter-names '*.engram.embed_tokens.weight*' +``` + +CPU-offloaded tables remain on CPU without an extra backup copy. Engram's +[inference-time CPU offload](engram.md#cpu-offload) is configured separately; +this option preserves the frozen tables across level-2 sleep and weight updates. + +To retain a frozen drafter while updating only target weights, use `draft.*`: + +```bash +VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3.8-27B --enable-sleep-mode \ + --speculative-config '{"method":"mtp","num_speculative_tokens":3}' \ + --sleep-preserve-parameter-names 'draft.*' +``` + +vLLM selects all draft-private parameters internally; you do not need to list +them individually. Use `draft.` to select specific draft parameters. +Parameters shared with the target are excluded from the draft selection, so +target updates remain authoritative. An empty `draft.*` selection is valid. +The Engram and draft patterns can be supplied together. #### Release only KV cache memory diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 6dc20a817794..adb8f52e2e8b 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -19,6 +19,48 @@ ) +def test_sleep_preservation_restores_private_draft_not_shared_target(): + target = torch.nn.Linear(2, 2, bias=False) + draft = torch.nn.Module() + draft.private = torch.nn.Linear(2, 2, bias=False) + draft.shared = target.weight + expected = draft.private.weight.detach().clone() + pointer = draft.private.weight.data_ptr() + worker = object.__new__(gpu_worker.Worker) + worker.get_draft_model = lambda: draft + assert set(worker._get_sleep_parameters(target)) == { + "weight", + "draft.private.weight", + } + worker._sleep_saved_parameters = {"draft.private.weight": expected} + with torch.no_grad(): + target.weight.fill_(3.0) + draft.private.weight.zero_() + worker._restore_sleep_parameters(target) + torch.testing.assert_close(draft.private.weight, expected) + assert draft.private.weight.data_ptr() == pointer + assert torch.all(target.weight == 3.0) + assert worker._sleep_saved_parameters == {} + + +@pytest.mark.parametrize("has_shared_draft", [False, True]) +def test_sleep_preservation_allows_empty_draft_but_rejects_target_typo( + has_shared_draft, +): + target = torch.nn.Linear(2, 2, bias=False) + draft = torch.nn.Module() if has_shared_draft else None + if draft is not None: + draft.shared = target.weight + worker = object.__new__(gpu_worker.Worker) + worker.get_draft_model = lambda: draft + worker.model_config = SimpleNamespace(sleep_preserve_parameter_names=["draft.*"]) + worker._save_sleep_parameters(target) + assert worker._sleep_saved_parameters == {} + worker.model_config.sleep_preserve_parameter_names = ["missing"] + with pytest.raises(ValueError, match="No parameter matches sleep retention"): + worker._save_sleep_parameters(target) + + def test_load_model_preserves_compiled_graphs_at_runtime(monkeypatch): """Profiling must use serving's thread count to keep Dynamo guards valid.""" from torch._dynamo.testing import CompileCounter diff --git a/vllm/config/model.py b/vllm/config/model.py index 502f1bbd8d4f..d4bafad5d6d3 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -372,6 +372,9 @@ class ModelConfig: hip platforms are supported).""" sleep_preserve_parameter_names: list[str] = field(default_factory=list) """Parameter-name globs to preserve across level-2 sleep. + Use target globs for frozen Engram table weights and scales omitted by the + trainer (e.g. ``*.engram.embed_tokens.weight*``). + Use ``draft.`` for draft-private parameters, or ``draft.*`` for all. The sender must omit these parameters; loaders must preserve their storage.""" sleep_mode_backend: str = "cumem" """Mechanism used to free and restore GPU state for sleep mode. ``"cumem"`` diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 3efe77a3a4d2..04feb5da025a 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -254,16 +254,28 @@ def sleep_mode_backend(self) -> "SleepModeBackend": ) return self._sleep_mode_backend + def _get_sleep_parameters(self, model: nn.Module) -> dict[str, nn.Parameter]: + parameters = dict(model.named_parameters()) + draft = self.get_draft_model() + if draft is not None: + shared = {id(parameter) for parameter in parameters.values()} + parameters.update( + (f"draft.{name}", parameter) + for name, parameter in draft.named_parameters() + if id(parameter) not in shared + ) + return parameters + def _save_sleep_parameters(self, model: nn.Module) -> None: patterns = self.model_config.sleep_preserve_parameter_names if not patterns: self._sleep_saved_parameters = {} return - parameters = dict(model.named_parameters()) + parameters = self._get_sleep_parameters(model) names: set[str] = set() for pattern in patterns: matches = fnmatch_filter(parameters, pattern) - if not matches: + if not matches and pattern != "draft.*": raise ValueError(f"No parameter matches sleep retention: {pattern}") names.update(matches) self._sleep_saved_parameters = { @@ -274,8 +286,9 @@ def _save_sleep_parameters(self, model: nn.Module) -> None: @torch.no_grad() def _restore_sleep_parameters(self, model: nn.Module) -> None: + parameters = self._get_sleep_parameters(model) for name, value in self._sleep_saved_parameters.items(): - model.get_parameter(name).copy_(value) + parameters[name].copy_(value) self._sleep_saved_parameters.clear() def sleep(self, level: int = 1) -> None: From 93d2cab35e8ae66bbaa1dea1c077986d5a79eb80 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20P=C3=A9rez=20de=20Algaba?= <124347725+jperezdealgaba@users.noreply.github.com> Date: Sun, 11 Oct 2026 17:39:11 +0200 Subject: [PATCH 263/624] [Security] Fix P2P ZMQ socket exhaustion via unbounded peer sessions (#51504) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: jperezde Signed-off-by: Juan Pérez de Algaba Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> --- docs/usage/security.md | 19 ++ tests/test_envs.py | 18 ++ .../v1/kv_offload/tiering/p2p/test_manager.py | 273 ++++++++++++++++++ .../kv_offload/tiering/p2p/test_sessions.py | 138 +++++++++ .../tiering/p2p/test_zmq_transport.py | 44 +++ vllm/envs.py | 18 ++ vllm/v1/kv_offload/tiering/p2p/control/zmq.py | 76 +++-- vllm/v1/kv_offload/tiering/p2p/manager.py | 125 +++++++- .../kv_offload/tiering/p2p/session/session.py | 43 ++- 9 files changed, 712 insertions(+), 42 deletions(-) diff --git a/docs/usage/security.md b/docs/usage/security.md index 50c4149be81b..f069de0def74 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -303,6 +303,25 @@ Only deployments whose API clients are trusted should start the server with `--trust-request-mm-kwargs` to restore per-request overrides. Do not enable this option on an endpoint exposed to untrusted clients. +### P2P KV-Offload Peer Session Limits + +Deployments that enable the P2P secondary tier (`OffloadingConnector` with `TieringOffloadingSpec` and a `p2p` secondary tier) accept peer addresses from `kv_transfer_params` in API requests. Without limits, an attacker with API access could submit many unique non-listening peer addresses, each consuming ZeroMQ sockets until the context quota is exhausted, crashing the EngineCore. + +vLLM enforces several configurable limits to prevent this: + +| Environment Variable | Default | Description | +| --- | --- | --- | +| `VLLM_P2P_MAX_PEERS` | `128` | Maximum concurrent P2P peer sessions. New peers (both outbound and inbound) are rejected when this limit is reached. | +| `VLLM_P2P_HANDSHAKE_TIMEOUT_S` | `30` | Seconds to wait for a newly connected peer to complete the P2P handshake. Sessions that remain connected but not ready after this deadline are torn down. | +| `VLLM_P2P_IDLE_TIMEOUT_S` | `300` | Seconds a P2P session may remain idle (no pending work, no messages) before it is eligible for eviction. Idle eviction runs only when remaining peer slots are within 8 of `VLLM_P2P_MAX_PEERS`. Set to `0` to disable idle eviction. | + +Connections that fail (ZMQ socket errors, handshake failures) are handled gracefully — the request falls back to local prefill instead of crashing the engine. + +**Recommendations:** + +- **Public-facing deployments with P2P enabled:** Lower `VLLM_P2P_MAX_PEERS` to the expected number of legitimate peers (e.g., `16` or `32`). +- **Sensitive deployments:** Consider deploying behind a reverse proxy that strips or validates `kv_transfer_params` from API requests to prevent client-supplied peer addresses entirely. + ## Tool Server and MCP Security vLLM supports connecting to external tool servers via the `--tool-server` argument. This enables models to call tools through the Responses API (`/v1/responses`). Tool server support works with all models — it is not limited to specific model architectures. diff --git a/tests/test_envs.py b/tests/test_envs.py index dfd08a9ec287..7356539bcdec 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -60,6 +60,24 @@ def test_p2p_side_channel_defaults_and_override(monkeypatch: pytest.MonkeyPatch) assert envs.VLLM_P2P_SIDE_CHANNEL_PORT == 5799 +def test_p2p_resource_limit_defaults_and_override( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.delenv("VLLM_P2P_MAX_PEERS", raising=False) + monkeypatch.delenv("VLLM_P2P_HANDSHAKE_TIMEOUT_S", raising=False) + monkeypatch.delenv("VLLM_P2P_IDLE_TIMEOUT_S", raising=False) + assert envs.VLLM_P2P_MAX_PEERS == 128 + assert envs.VLLM_P2P_HANDSHAKE_TIMEOUT_S == 30 + assert envs.VLLM_P2P_IDLE_TIMEOUT_S == 300 + + monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "64") + monkeypatch.setenv("VLLM_P2P_HANDSHAKE_TIMEOUT_S", "15") + monkeypatch.setenv("VLLM_P2P_IDLE_TIMEOUT_S", "60") + assert envs.VLLM_P2P_MAX_PEERS == 64 + assert envs.VLLM_P2P_HANDSHAKE_TIMEOUT_S == 15 + assert envs.VLLM_P2P_IDLE_TIMEOUT_S == 60 + + def test_getattr_with_cache(monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("VLLM_HOST_IP", "1.1.1.1") monkeypatch.setenv("VLLM_PORT", "1234") diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index e35b1d022e42..930158c7f6c8 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -647,6 +647,7 @@ def __init__( close_req_ids: list[str] | None = None, close_stores: list[int] | None = None, close_failed_serves: list[ReqContext] | None = None, + last_activity_at: float | None = None, ) -> None: self.peer_id = peer_id self.alive = alive @@ -663,6 +664,9 @@ def __init__( self.stores_added: list[tuple[str, list, object, int]] = [] self.attached: list[object] = [] self.finishes: list[str] = [] + self.last_activity_at: float = ( + last_activity_at if last_activity_at is not None else time.monotonic() + ) # Mirror P2PSession._server._inflight (transfer_id → handle) and # P2PSession._client.has_active_loads for the shutdown-drain and # drain_jobs paths. Tests populate _server._inflight when needed. @@ -1844,3 +1848,272 @@ def test_same_host_port_gets_distinct_nixl_names(self, monkeypatch): mgr_b = self._construct(monkeypatch, host="localhost", port=5710) assert mgr_a._local_id == mgr_b._local_id == "localhost:5710" assert mgr_a._nixl_agent_name != mgr_b._nixl_agent_name + + +# --------------------------------------------------------------------------- +# Tests for session capacity limit (VLLM_P2P_MAX_PEERS) +# --------------------------------------------------------------------------- + + +class TestSessionCapLimit: + """_get_or_create_session and _accept_new_peers reject new peers when + the session count reaches VLLM_P2P_MAX_PEERS.""" + + def _make_with_control(self): + mgr = _make_manager() + created = [] + + class FakeControl: + def connect(self, peer_id): + conn = _RecordingConn(peer_id) + created.append(conn) + return conn + + def poll(self): + return [] + + class FakeData: + block_len = 4096 + base_addr = 0x1000 + num_blocks = 16 + config_fingerprint = "" + + def get_agent_metadata(self): + return b"meta" + + def add_remote_peer(self, *a, **k): + pass + + def remove_remote_peer(self, pid): + pass + + mgr._control = FakeControl() # type: ignore[assignment] + mgr._data = FakeData() # type: ignore[assignment] + return mgr, created + + def test_outbound_rejected_at_capacity(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "2") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None) + + mgr, created = self._make_with_control() + mgr._sessions["peer-a:1"] = _FakeSession(peer_id="peer-a:1") # type: ignore[assignment] + mgr._sessions["peer-b:2"] = _FakeSession(peer_id="peer-b:2") # type: ignore[assignment] + + result = mgr._get_or_create_session("peer-c:3") + assert result is None + assert "peer-c:3" not in mgr._sessions + assert len(created) == 0 + + def test_on_new_request_records_failure_on_cap(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None) + + mgr, _ = self._make_with_control() + mgr._sessions["existing:1"] = _FakeSession(peer_id="existing:1") # type: ignore[assignment] + + ctx = _req_context( + kv_params=_remote_prefiller_kv_params( + remote_host="10.0.0.2", remote_port=9000, kv_request_id="req-cap" + ) + ) + mgr.on_new_request(ctx) + assert "req-cap" in mgr._failed_req_ids + + def test_inbound_rejected_at_capacity(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None) + + mgr, _ = self._make_with_control() + mgr._sessions["existing:1"] = _FakeSession(peer_id="existing:1") # type: ignore[assignment] + + conn = _RecordingConn("new-peer:2") + mgr._accept_new_peers([conn]) + assert conn.close_calls == 1 + assert "new-peer:2" not in mgr._sessions + + def test_existing_session_reused_below_cap(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None) + + mgr, _ = self._make_with_control() + existing = _FakeSession(peer_id="existing:1") + mgr._sessions["existing:1"] = existing # type: ignore[assignment] + + result = mgr._get_or_create_session("existing:1") + assert result is existing + + +# --------------------------------------------------------------------------- +# Tests for graceful error boundary in _get_or_create_session +# --------------------------------------------------------------------------- + + +class TestGetOrCreateSessionErrorBoundary: + def test_connect_failure_returns_none(self): + mgr = _make_manager() + + class FailingControl: + def connect(self, peer_id): + raise RuntimeError("socket exhausted") + + def poll(self): + return [] + + mgr._control = FailingControl() # type: ignore[assignment] + result = mgr._get_or_create_session("bad:1") + assert result is None + assert "bad:1" not in mgr._sessions + + def test_session_construction_failure_closes_conn(self): + """connect() succeeding but P2PSession construction failing must + close the conn so it is not left registered without a session.""" + mgr = _make_manager() + created: list[_RecordingConn] = [] + + class RaisingConn(_RecordingConn): + def send(self, msg): + raise RuntimeError("send failed during connect") + + class Control: + def connect(self, peer_id): + conn = RaisingConn(peer_id) + created.append(conn) + return conn + + def poll(self): + return [] + + class FakeData: + block_len = 4096 + base_addr = 0x1000 + num_blocks = 16 + config_fingerprint = "" + + def get_agent_metadata(self): + return b"meta" + + mgr._control = Control() # type: ignore[assignment] + mgr._data = FakeData() # type: ignore[assignment] + result = mgr._get_or_create_session("bad:2") + assert result is None + assert "bad:2" not in mgr._sessions + assert len(created) == 1 + assert created[0].close_calls == 1 + + +# --------------------------------------------------------------------------- +# Tests for idle session eviction +# --------------------------------------------------------------------------- + + +class TestIdleSessionEviction: + def _make_with_fake_data(self): + mgr = _make_manager() + + class FakeControl: + def poll(self): + return [] + + class FakeData: + def remove_remote_peer(self, pid): + pass + + mgr._control = FakeControl() # type: ignore[assignment] + mgr._data = FakeData() # type: ignore[assignment] + return mgr + + def test_idle_session_evicted(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_IDLE_TIMEOUT_S", "10") + monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_IDLE_TIMEOUT_S", None) + vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None) + + mgr = self._make_with_fake_data() + idle_session = _FakeSession( + peer_id="idle:1", + last_activity_at=time.monotonic() - 20, + ) + mgr._sessions["idle:1"] = idle_session # type: ignore[assignment] + + mgr._reap_idle_sessions() + assert "idle:1" not in mgr._sessions + + def test_idle_session_kept_below_headroom(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_IDLE_TIMEOUT_S", "10") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_IDLE_TIMEOUT_S", None) + vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None) + + mgr = self._make_with_fake_data() + idle_session = _FakeSession( + peer_id="idle:1", + last_activity_at=time.monotonic() - 20, + ) + mgr._sessions["idle:1"] = idle_session # type: ignore[assignment] + + mgr._reap_idle_sessions() + assert "idle:1" in mgr._sessions + + def test_active_session_kept(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_IDLE_TIMEOUT_S", "10") + monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_IDLE_TIMEOUT_S", None) + vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None) + + mgr = self._make_with_fake_data() + active = _FakeSession( + peer_id="active:1", + last_activity_at=time.monotonic(), + ) + mgr._sessions["active:1"] = active # type: ignore[assignment] + + mgr._reap_idle_sessions() + assert "active:1" in mgr._sessions + + def test_session_with_pending_work_not_evicted(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_IDLE_TIMEOUT_S", "1") + monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_IDLE_TIMEOUT_S", None) + vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None) + + mgr = self._make_with_fake_data() + busy = _FakeSession( + peer_id="busy:1", + last_activity_at=time.monotonic() - 100, + ) + busy._server._inflight[1] = object() + mgr._sessions["busy:1"] = busy # type: ignore[assignment] + + mgr._reap_idle_sessions() + assert "busy:1" in mgr._sessions + + def test_disabled_when_zero(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_IDLE_TIMEOUT_S", "0") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_IDLE_TIMEOUT_S", None) + + mgr = self._make_with_fake_data() + old = _FakeSession( + peer_id="old:1", + last_activity_at=time.monotonic() - 999999, + ) + mgr._sessions["old:1"] = old # type: ignore[assignment] + + mgr._reap_idle_sessions() + assert "old:1" in mgr._sessions diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index 09739345be20..300cbacd665d 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -2698,3 +2698,141 @@ def test_success_wrong_type(self): } with pytest.raises(ValueError, match="success"): TransferDoneMsg.validate(msg) + + +# --------------------------------------------------------------------------- +# Tests for handshake deadline +# --------------------------------------------------------------------------- + + +class TestHandshakeDeadline: + """A connected-but-not-ready session is marked dead after the + handshake timeout expires.""" + + def test_marks_dead_after_timeout(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_HANDSHAKE_TIMEOUT_S", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_HANDSHAKE_TIMEOUT_S", None) + + transport = FakeDataTransport() + conn = FakeConnection(peer_id="stale:9999") + session = P2PSession( + peer_id="stale:9999", + local_id="local:1", + transport=transport, + local_block_len=4096, + local_hash_seed=_DEFAULT_HASH_SEED, + conn=conn, + ) + + assert session.connected + assert not session.ready + assert session.alive + + session._connected_at = time.monotonic() - 5 + + session.poll() + assert not session.alive + + def test_no_timeout_when_ready(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_HANDSHAKE_TIMEOUT_S", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_HANDSHAKE_TIMEOUT_S", None) + + transport = FakeDataTransport() + conn = FakeConnection(peer_id="good:1234") + session = P2PSession( + peer_id="good:1234", + local_id="local:1", + transport=transport, + local_block_len=4096, + local_hash_seed=_DEFAULT_HASH_SEED, + conn=conn, + ) + + # Complete the handshake by feeding ConnectMsg + ConnectAckMsg. + conn._inbox = [ + { + TYPE_KEY: ConnectMsg.TYPE, + ConnectMsg.PEER_ID: "good:1234", + ConnectMsg.AGENT_METADATA: b"meta", + ConnectMsg.BASE_ADDR: 0x1000, + ConnectMsg.NUM_BLOCKS: 16, + ConnectMsg.BLOCK_LEN: 4096, + ConnectMsg.CONFIG_FINGERPRINT: "", + ConnectMsg.HASH_SEED: _DEFAULT_HASH_SEED, + }, + {TYPE_KEY: ConnectAckMsg.TYPE, ConnectAckMsg.PEER_ID: "good:1234"}, + ] + session.poll() + assert session.ready + + session._connected_at = time.monotonic() - 100 + session.poll() + assert session.alive + + def test_buffered_ack_at_deadline_completes_handshake(self, monkeypatch): + """A ConnectAck already in the inbox at the deadline must complete + the handshake instead of being discarded by mark_dead.""" + monkeypatch.setenv("VLLM_P2P_HANDSHAKE_TIMEOUT_S", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_HANDSHAKE_TIMEOUT_S", None) + + transport = FakeDataTransport() + conn = FakeConnection(peer_id="late-ack:1") + session = P2PSession( + peer_id="late-ack:1", + local_id="local:1", + transport=transport, + local_block_len=4096, + local_hash_seed=_DEFAULT_HASH_SEED, + conn=conn, + ) + session._connected_at = time.monotonic() - 5 + conn._inbox = [ + { + TYPE_KEY: ConnectMsg.TYPE, + ConnectMsg.PEER_ID: "late-ack:1", + ConnectMsg.AGENT_METADATA: b"meta", + ConnectMsg.BASE_ADDR: 0x1000, + ConnectMsg.NUM_BLOCKS: 16, + ConnectMsg.BLOCK_LEN: 4096, + ConnectMsg.CONFIG_FINGERPRINT: "", + ConnectMsg.HASH_SEED: _DEFAULT_HASH_SEED, + }, + {TYPE_KEY: ConnectAckMsg.TYPE, ConnectAckMsg.PEER_ID: "late-ack:1"}, + ] + session.poll() + assert session.ready + assert session.alive + + def test_activity_timestamp_updated_on_recv(self): + transport = FakeDataTransport() + conn = FakeConnection(peer_id="peer:1") + session = P2PSession( + peer_id="peer:1", + local_id="local:1", + transport=transport, + local_block_len=4096, + local_hash_seed=_DEFAULT_HASH_SEED, + conn=conn, + ) + before = session._last_activity_at + + conn._inbox = [ + { + TYPE_KEY: ConnectMsg.TYPE, + ConnectMsg.PEER_ID: "peer:1", + ConnectMsg.AGENT_METADATA: b"meta", + ConnectMsg.BASE_ADDR: 0x1000, + ConnectMsg.NUM_BLOCKS: 16, + ConnectMsg.BLOCK_LEN: 4096, + ConnectMsg.CONFIG_FINGERPRINT: "", + ConnectMsg.HASH_SEED: _DEFAULT_HASH_SEED, + }, + ] + session.poll() + assert session._last_activity_at >= before diff --git a/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py b/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py index f814a04887c4..90e87bb40742 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py +++ b/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py @@ -336,3 +336,47 @@ def test_inbound_message_survives_dead_registration(self): finally: transport_a.close() transport_b.close() + + +class TestExpandedMonitorEvents: + """The monitor mask includes EVENT_DISCONNECTED and EVENT_CLOSED so + both are treated as fatal.""" + + def test_fatal_events_include_disconnected_and_closed(self): + from vllm.v1.kv_offload.tiering.p2p.control.zmq import ( + _FATAL_MONITOR_EVENTS, + ) + + assert zmq.EVENT_DISCONNECTED in _FATAL_MONITOR_EVENTS + assert zmq.EVENT_CLOSED in _FATAL_MONITOR_EVENTS + + def test_monitor_mask_covers_fatal_events(self): + from vllm.v1.kv_offload.tiering.p2p.control.zmq import ( + _FATAL_MONITOR_EVENTS, + _MONITOR_EVENT_MASK, + ) + + for ev in _FATAL_MONITOR_EVENTS: + assert _MONITOR_EVENT_MASK & ev == ev + + def test_disconnect_detected_after_peer_shutdown(self): + """When a peer's transport closes, the outbound connection is + eventually marked dead by the monitor.""" + transport_a, port_a = _make_transport() + transport_b, _ = _make_transport() + + try: + conn_out = transport_b.connect(f"127.0.0.1:{port_a}") + conn_out.send({"type": "hello"}) + _wait_for_inbound(transport_a) + + transport_a.close() + + end = time.monotonic() + 15 + while conn_out.alive and time.monotonic() < end: + transport_b.poll() + time.sleep(0.05) + + assert not conn_out.alive + finally: + transport_b.close() diff --git a/vllm/envs.py b/vllm/envs.py index 02c282fd8ec2..14489a3f2503 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -233,6 +233,9 @@ VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5600 VLLM_P2P_SIDE_CHANNEL_HOST: str = "localhost" VLLM_P2P_SIDE_CHANNEL_PORT: int = 5710 + VLLM_P2P_MAX_PEERS: int = 128 + VLLM_P2P_HANDSHAKE_TIMEOUT_S: int = 30 + VLLM_P2P_IDLE_TIMEOUT_S: int = 300 VLLM_EC_SIDE_CHANNEL_HOST: str = "localhost" VLLM_EC_SIDE_CHANNEL_PORT: int = 5601 VLLM_MOONCAKE_BOOTSTRAP_PORT: int = 8998 @@ -1739,6 +1742,21 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_P2P_SIDE_CHANNEL_PORT": lambda: int( os.getenv("VLLM_P2P_SIDE_CHANNEL_PORT", "5710") ), + # Maximum number of concurrent P2P peer sessions. Limits ZMQ socket + # consumption to prevent denial-of-service via unbounded peer + # creation from client-supplied kv_transfer_params. Default: 128. + "VLLM_P2P_MAX_PEERS": lambda: int(os.getenv("VLLM_P2P_MAX_PEERS", "128")), + # Seconds to wait for a P2P peer handshake to complete (connected + # but not yet ready). Sessions exceeding this deadline are reaped. + # Prevents never-responding peers from accumulating indefinitely. + "VLLM_P2P_HANDSHAKE_TIMEOUT_S": lambda: int( + os.getenv("VLLM_P2P_HANDSHAKE_TIMEOUT_S", "30") + ), + # Seconds a P2P session may remain idle (no pending work, no + # messages) before it is eligible for eviction. Idle eviction runs + # only when remaining peer slots are within 8 of VLLM_P2P_MAX_PEERS. + # Default: 300. Set to 0 to disable. + "VLLM_P2P_IDLE_TIMEOUT_S": lambda: int(os.getenv("VLLM_P2P_IDLE_TIMEOUT_S", "300")), # IP address used for the EC connector's ZMQ side channel # (producer ROUTER bind, consumer DEALER dial). "VLLM_EC_SIDE_CHANNEL_HOST": lambda: os.getenv( diff --git a/vllm/v1/kv_offload/tiering/p2p/control/zmq.py b/vllm/v1/kv_offload/tiering/p2p/control/zmq.py index b28ed76dac97..074c3869f4cd 100644 --- a/vllm/v1/kv_offload/tiering/p2p/control/zmq.py +++ b/vllm/v1/kv_offload/tiering/p2p/control/zmq.py @@ -27,11 +27,35 @@ _HEARTBEAT_TIMEOUT_MS = 10000 _HEARTBEAT_TTL_MS = 10000 +# Events that indicate a fatal connection failure — the peer will never +# become reachable through this socket. +_FATAL_MONITOR_EVENTS: frozenset[int] = frozenset() # populated after import + # Shared sentinels returned when there is nothing to report. _EMPTY_INBOX: tuple[dict, ...] = () _EMPTY_NEW_CONNECTIONS: tuple[ControlConnection, ...] = () +def _build_fatal_events() -> frozenset[int]: + events = {zmq.EVENT_DISCONNECTED, zmq.EVENT_CLOSED} + for attr in ( + "EVENT_HANDSHAKE_FAILED_AUTH", + "EVENT_HANDSHAKE_FAILED_PROTOCOL", + "EVENT_HANDSHAKE_FAILED_NO_DETAIL", + ): + val = getattr(zmq, attr, None) + if val is not None: + events.add(val) + return frozenset(events) + + +_FATAL_MONITOR_EVENTS = _build_fatal_events() + +_MONITOR_EVENT_MASK: int = 0 +for _ev in _FATAL_MONITOR_EVENTS: + _MONITOR_EVENT_MASK |= _ev + + def _tcp_addr(host: str, port: int | str) -> str: return f"tcp://{host}:{port}" @@ -247,7 +271,7 @@ def _open_connection( safe_id = peer_id.replace(":", "-").replace("/", "-") monitor_addr = f"inproc://p2p-monitor-{safe_id}-{self._monitor_seq}" self._monitor_seq += 1 - dealer.monitor(monitor_addr, zmq.EVENT_DISCONNECTED) + dealer.monitor(monitor_addr, _MONITOR_EVENT_MASK) monitor_sock = self._zmq_ctx.socket(zmq.PAIR) monitor_sock.connect(monitor_addr) @@ -319,29 +343,33 @@ def _recv_router(self) -> None: self._pending_inbound.append((sender_id, msg)) def _check_monitors(self) -> None: - """Non-blocking: check all monitor sockets for disconnection.""" + """Non-blocking: check all monitor sockets for fatal events.""" for conn in self._connections.values(): if not conn.alive: continue - try: - event = zmq.utils.monitor.recv_monitor_message( - conn.monitor_socket, zmq.NOBLOCK - ) - except zmq.Again: - continue - except zmq.ZMQError as exc: - logger.warning( - "ZmqTransport %s: monitor error for peer %s: %s", - self._local_id, - conn.peer_id, - exc, - ) - continue - - if event["event"] == zmq.EVENT_DISCONNECTED: - logger.debug( - "ZmqTransport %s: peer %s disconnected", - self._local_id, - conn.peer_id, - ) - conn.mark_dead() + while True: + try: + event = zmq.utils.monitor.recv_monitor_message( + conn.monitor_socket, zmq.NOBLOCK + ) + except zmq.Again: + break + except zmq.ZMQError as exc: + logger.warning( + "ZmqTransport %s: monitor error for peer %s: %s", + self._local_id, + conn.peer_id, + exc, + ) + break + + ev = event["event"] + if ev in _FATAL_MONITOR_EVENTS: + logger.debug( + "ZmqTransport %s: peer %s fatal event 0x%x", + self._local_id, + conn.peer_id, + ev, + ) + conn.mark_dead() + break diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index f759a4481e5d..b7268eb420d6 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -69,6 +69,11 @@ # enough to avoid busy-spinning the scheduler thread. _DRAIN_SLEEP_S = 0.001 +# Idle sessions are only reaped when remaining capacity is at most this +# many slots. Session setup is expensive; below this headroom, keep idle +# peers so a later burst does not pay handshake again. +_IDLE_REAP_HEADROOM = 8 + def _remote_prefiller_params(kv_params: dict | None) -> dict | None: """Return the ``remote_prefiller`` sub-dict, or None if absent. @@ -428,7 +433,9 @@ def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: _annotate_req_context(req_context) source = req_context.get_state(P2PSourceInfo) if source is not None: - self._get_or_create_session(source.peer_id) + session = self._get_or_create_session(source.peer_id) + if session is None: + self._failed_req_ids.add(source.kv_request_id) dest = req_context.get_state(P2PDestInfo) if dest is not None and dest.kv_request_id: return RequestOffloadingContext(policy=OffloadPolicy.REQUEST_LEVEL) @@ -685,38 +692,78 @@ def _get_hash_seed(self) -> str: self._hash_seed = get_none_hash_seed() return self._hash_seed - def _get_or_create_session(self, peer_id: str) -> P2PSession: + def _get_or_create_session(self, peer_id: str) -> P2PSession | None: """Return the existing session for peer_id, or open one outbound. Consumer-side helper for on_new_request: when ``remote_prefiller`` - (PD) or ``remote_kv_source`` (symmetric P2P) is set, the consumer must reach the - producer at peer_id. If we already have a session toward that - peer (from a prior load or a peer-initiated inbound), reuse it; - otherwise open an outbound ControlConnection and build a + (PD) or ``remote_kv_source`` (symmetric P2P) is set, the consumer + must reach the producer at peer_id. If we already have a session + toward that peer (from a prior load or a peer-initiated inbound), + reuse it; otherwise open an outbound ControlConnection and build a connected session. + + Returns None if the session could not be created (capacity + exceeded or transport error). """ session = self._sessions.get(peer_id) if session is not None: return session - conn = self._control.connect(peer_id) - session = P2PSession( - peer_id=peer_id, - local_id=self._local_id, - transport=self._data, - local_block_len=self._data.block_len, - local_hash_seed=self._get_hash_seed(), - conn=conn, - ) + + max_peers = envs.VLLM_P2P_MAX_PEERS + if len(self._sessions) >= max_peers: + logger.warning( + "P2P %s: rejecting outbound peer %s — session limit " + "reached (%d/%d). Set VLLM_P2P_MAX_PEERS to increase.", + self._local_id, + peer_id, + len(self._sessions), + max_peers, + ) + return None + + conn = None + try: + conn = self._control.connect(peer_id) + session = P2PSession( + peer_id=peer_id, + local_id=self._local_id, + transport=self._data, + local_block_len=self._data.block_len, + local_hash_seed=self._get_hash_seed(), + conn=conn, + ) + except Exception: + logger.exception( + "P2P %s: failed to open connection to %s", + self._local_id, + peer_id, + ) + if conn is not None: + conn.close() + return None + self._sessions[peer_id] = session return session def _accept_new_peers(self, new_connections: Sequence[ControlConnection]) -> None: + max_peers = envs.VLLM_P2P_MAX_PEERS for conn in new_connections: logger.info( "P2P %s: accepting incoming connection from %s", self._local_id, conn.peer_id, ) + if len(self._sessions) >= max_peers: + logger.warning( + "P2P %s: rejecting inbound peer %s — session limit " + "reached (%d/%d). Set VLLM_P2P_MAX_PEERS to increase.", + self._local_id, + conn.peer_id, + len(self._sessions), + max_peers, + ) + conn.close() + continue try: existing = self._sessions.get(conn.peer_id) if existing is not None: @@ -827,6 +874,53 @@ def _reap_unbound_stores(self) -> None: len(batches), ) + def _reap_idle_sessions(self) -> None: + """Evict idle sessions when remaining peer slots are scarce. + + A session is idle when it has no pending work and its last activity + timestamp exceeds ``VLLM_P2P_IDLE_TIMEOUT_S``. Eviction runs only + when ``len(_sessions) >= VLLM_P2P_MAX_PEERS - _IDLE_REAP_HEADROOM`` + so a quiet but healthy mesh is not torn down. Uses the same + teardown path as ``_reap_dead_sessions``. + """ + idle_timeout = envs.VLLM_P2P_IDLE_TIMEOUT_S + if idle_timeout <= 0: + return + max_peers = envs.VLLM_P2P_MAX_PEERS + headroom = min(_IDLE_REAP_HEADROOM, max(0, max_peers - 1)) + if len(self._sessions) < max_peers - headroom: + return + deadline = time.monotonic() - idle_timeout + idle: list[str] | None = None + for pid, s in self._sessions.items(): + if not s.has_pending_work and s.last_activity_at <= deadline: + if idle is None: + idle = [] + idle.append(pid) + if idle is None: + return + for pid in idle: + session = self._sessions.pop(pid) + stale_kv_ids = [ + kid for kid, s in self._kv_to_session.items() if s is session + ] + for kid in stale_kv_ids: + del self._kv_to_session[kid] + close_result = session.close() + for job_id in close_result.failed_jobs: + self._finished_jobs.append(JobResult(job_id=job_id, success=False)) + for job_id in close_result.failed_stores: + self._finished_jobs.append(JobResult(job_id=job_id, success=False)) + self._failed_req_ids.update(close_result.failed_req_ids) + self._failed_serve_ctxs.extend(close_result.failed_serves) + self._data.remove_remote_peer(pid) + logger.info( + "P2P %s: evicted idle peer %s (idle > %ds)", + self._local_id, + pid, + idle_timeout, + ) + # ------------------------------------------------------------------ # Polling # ------------------------------------------------------------------ @@ -893,6 +987,7 @@ def _poll_once(self) -> None: self._reap_dead_sessions() self._reap_unbound_stores() + self._reap_idle_sessions() # ------------------------------------------------------------------ # Lifecycle diff --git a/vllm/v1/kv_offload/tiering/p2p/session/session.py b/vllm/v1/kv_offload/tiering/p2p/session/session.py index f43e1a5aab5a..ef490304c901 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/session.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/session.py @@ -18,9 +18,11 @@ from __future__ import annotations import contextlib +import time from collections.abc import Sequence from typing import TYPE_CHECKING, NamedTuple +import vllm.envs as envs from vllm.logger import init_logger from vllm.v1.kv_offload.base import OffloadKey from vllm.v1.kv_offload.tiering.p2p.control.base import ControlConnection @@ -128,6 +130,10 @@ def __init__( # Consecutive non-protocol dispatch errors. Reset on success. self._dispatch_error_count: int = 0 + now = time.monotonic() + self._connected_at: float | None = None + self._last_activity_at: float = now + # kv_request_ids whose FetchMsg arrived during the current poll # tick. Drained and returned in the next poll() result so the # manager can bind kv_request_id → session and replay any @@ -168,6 +174,11 @@ def has_pending_work(self) -> bool: """True while inbound loads or outbound transfers are outstanding.""" return self._client.has_active_loads or self._server.has_inflight_transfers + @property + def last_activity_at(self) -> float: + """Monotonic timestamp of the last meaningful session activity.""" + return self._last_activity_at + # ------------------------------------------------------------------ # Connection lifecycle # ------------------------------------------------------------------ @@ -181,6 +192,9 @@ def attach_connection(self, conn: ControlConnection) -> None: if self._conn is not None: raise ValueError(f"P2PSession {self.peer_id}: already connected") self._conn = conn + now = time.monotonic() + self._connected_at = now + self._last_activity_at = now self._send_connect() # ------------------------------------------------------------------ @@ -249,17 +263,39 @@ def serve_external_requests(self, parent: ParentManager) -> None: def poll(self) -> SessionPollResult: """Process incoming messages, drive transfers, apply timeouts.""" if self._conn is None: - # Pending session — store-job timeouts still apply so buffered - # jobs that never get picked up are surfaced as failures. return SessionPollResult( loads=[], stores=self._server.collect_idle_timeouts(), new_fetch_ids=[], ) - for msg in self._conn.recv(): + # Drain the inbox first so a ConnectAck already buffered at the + # deadline (e.g. after a scheduler stall) can complete the + # handshake instead of being discarded by mark_dead. + messages = self._conn.recv() + if messages: + self._last_activity_at = time.monotonic() + for msg in messages: self._on_message(msg) + # Handshake deadline: a peer that never completes the handshake + # is stuck in connected-but-not-ready state. Mark it dead so the + # manager reaps it instead of accumulating sockets indefinitely. + if ( + not self._send_ready + and self._conn.alive + and self._connected_at is not None + and time.monotonic() - self._connected_at + > envs.VLLM_P2P_HANDSHAKE_TIMEOUT_S + ): + logger.warning( + "P2PSession %s: handshake timed out after %ds " + "— marking connection dead", + self.peer_id, + envs.VLLM_P2P_HANDSHAKE_TIMEOUT_S, + ) + self._conn.mark_dead() + loads = self._client.collect_results() stores = self._server.collect_results() self._server.drain_pending_aborts() @@ -533,6 +569,7 @@ def _do_send(self, msg: dict) -> None: return try: self._conn.send(msg) + self._last_activity_at = time.monotonic() logger.debug("P2PSession %s: sent %s", self.peer_id, msg.get(TYPE_KEY)) except Exception: # A send failure means the connection is broken. Swallowing it From b99fdbe1f896abd7f4b133685a37397c3427f7d5 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Mon, 12 Oct 2026 00:27:21 +0800 Subject: [PATCH 264/624] [Bugfix] Respect the configured offline chat template (#61061) Signed-off-by: Ting Sun Co-authored-by: Cyrus Leung --- tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py | 1 + vllm/entrypoints/common/offline.py | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py b/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py index 91d2b0aa434a..a381e0fa3b67 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py +++ b/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py @@ -14,6 +14,7 @@ def _make_mock_llm() -> LLM: llm = object.__new__(LLM) + llm.chat_template = None llm.model_config = SimpleNamespace( runner_type="generate", enable_prompt_embeds=False ) diff --git a/vllm/entrypoints/common/offline.py b/vllm/entrypoints/common/offline.py index 1bb0bedccbd0..a5af009f2226 100644 --- a/vllm/entrypoints/common/offline.py +++ b/vllm/entrypoints/common/offline.py @@ -54,6 +54,7 @@ class OfflineInferenceMixin: renderer: BaseRenderer llm_engine: "LLMEngine" model_config: ModelConfig + chat_template: str | None def _resolve_mm_lora( self, @@ -180,7 +181,9 @@ def _preprocess_chat( renderer = self.renderer chat_params = ChatParams( - chat_template=chat_template, + chat_template=( + self.chat_template if chat_template is None else chat_template + ), chat_template_content_format=chat_template_content_format, chat_template_kwargs=merge_kwargs( chat_template_kwargs, From c059bee812842fb76c2b045e240c02f37c0d498d Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Sun, 11 Oct 2026 09:57:07 -0700 Subject: [PATCH 265/624] [MRV2] Switch CI tests still running on MRV1 unnecessarily (#61014) Signed-off-by: Nick Hill --- .buildkite/test-amd.yaml | 4 -- .buildkite/test_areas/quantization.yaml | 8 ---- .../memory/sleep_mode/test_sleep_mode.py | 15 ++---- .../correctness_e2e/test_attn_quant.py | 1 - tests/models/kimi_k3/test_kda_metadata.py | 2 +- tests/utils.py | 14 ------ tests/v1/core/test_scheduler.py | 6 +++ tests/v1/core/utils.py | 2 +- tests/v1/e2e/test_replayssm_decode.py | 47 +++++++++++++------ .../unit/offloading_connector/test_config.py | 20 ++++---- .../test_eagle3_aux_hidden_states_pp.py | 2 +- tests/v1/worker/test_gpu_worker.py | 4 +- 12 files changed, 59 insertions(+), 66 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index a3f7dbd83874..a14a0a33a4d2 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2046,7 +2046,6 @@ steps: - vllm/model_executor/model_loader/ optional: true commands: - - unset VLLM_USE_V2_MODEL_RUNNER - pytest -v -s models/quantization #-------------------------------------------------- mi300 · models / transformers ---------------------------------------------------# @@ -2353,7 +2352,6 @@ steps: # https://github.com/pytorch/ao/issues/2919, we'll have to skip new torchao tests for now # we can only upgrade after this is resolved # TODO(jerryzh168): resolve the above comment - - unset VLLM_USE_V2_MODEL_RUNNER - uv pip install --system torchao==0.18.0 - uv pip install --system conch-triton-kernels - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT @@ -5163,7 +5161,6 @@ steps: - vllm/model_executor/model_loader/ optional: true commands: - - unset VLLM_USE_V2_MODEL_RUNNER - pytest -v -s models/quantization #------------------------------------------------------- mi355 · quantization --------------------------------------------------------# @@ -5186,7 +5183,6 @@ steps: - tests/rocm/test_moe_weight_replay.py optional: true commands: - - unset VLLM_USE_V2_MODEL_RUNNER - uv pip install --system torchao==0.17.0 - uv pip install --system conch-triton-kernels - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py --ignore quantization/test_gfx950_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index 55cef9fbd33a..9517b54f8ced 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -7,8 +7,6 @@ steps: key: quantization timeout_in_minutes: 40 parallelism: 4 - env: - VLLM_USE_V2_MODEL_RUNNER: "0" source_file_dependencies: - csrc/ - vllm/model_executor/layers/quantization @@ -41,8 +39,6 @@ steps: - tests/rocm/test_moe_weight_replay.py - tests/rocm/test_moe_padding.py commands: - # Use the ROCm default instead of the NVIDIA runner override. - - unset VLLM_USE_V2_MODEL_RUNNER - uv pip install --system torchao==0.17.0 - uv pip install --system conch-triton-kernels - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT @@ -130,8 +126,6 @@ steps: key: quantized-models-test automatic_shard: true timeout_in_minutes: 65 - env: - VLLM_USE_V2_MODEL_RUNNER: "0" source_file_dependencies: - vllm/model_executor/layers/quantization - tests/models/quantization @@ -153,6 +147,4 @@ steps: - vllm/platforms/rocm.py - vllm/model_executor/model_loader/ commands: - # Use the ROCm default instead of the NVIDIA runner override. - - unset VLLM_USE_V2_MODEL_RUNNER - pytest -v -s models/quantization diff --git a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py index d897d710015b..bc0f2756072f 100644 --- a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py +++ b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py @@ -411,15 +411,9 @@ def test_deep_sleep_compressed_tensors_kv_scales( @requires_fp8 -def test_deep_sleep_fp8_kvcache_mrv1( - monkeypatch: pytest.MonkeyPatch, gpu_memory_cleared -): - # Regression test for https://github.com/vllm-project/vllm/pull/28783. - # In particular, verify that MRV1 does not rely on post_kv_cache_wake_up() - # to restore correct output after level-2 sleep. - monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") - envs.disable_envs_cache() - +def test_deep_sleep_fp8_kvcache(gpu_memory_cleared): + # Regression test for https://github.com/vllm-project/vllm/pull/28783: + # fp8 KV cache output must be restored correctly after level-2 sleep. model = "Qwen/Qwen2-0.5B" used_bytes_baseline = current_platform.get_current_memory_usage() @@ -454,10 +448,9 @@ def test_deep_sleep_fp8_kvcache_mrv1( @requires_fp8 -def test_deep_sleep_fp8_kvcache_mrv1_with_undefined_remap( +def test_deep_sleep_fp8_kvcache_with_undefined_remap( monkeypatch: pytest.MonkeyPatch, ): - monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") envs.disable_envs_cache() diff --git a/tests/compile/correctness_e2e/test_attn_quant.py b/tests/compile/correctness_e2e/test_attn_quant.py index c7fc849ae52f..e13a8f264a6a 100644 --- a/tests/compile/correctness_e2e/test_attn_quant.py +++ b/tests/compile/correctness_e2e/test_attn_quant.py @@ -77,7 +77,6 @@ def settings(enabled): env1=env, env2=env, method="generate", - force_v1_runner=True, ) matches = [ int(count) for count in re.findall(r"'attn_quant_fusion': (\d+)", logs.text) diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index 63024634d305..65b29b44bdae 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -219,7 +219,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch): scheduler_config=SimpleNamespace(max_num_seqs=4), additional_config={}, num_speculative_tokens=2, - use_v2_model_runner=False, + use_v2_model_runner=True, ) builder = KimiK3KDAMetadataBuilder( kv_cache_spec=kv_cache_spec, diff --git a/tests/utils.py b/tests/utils.py index dd6bf4ae969d..94f62ef8b870 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1165,7 +1165,6 @@ def compare_two_settings( method: str = "generate", max_wait_seconds: float | None = None, include_seeded_sampling: bool = True, - force_v1_runner: bool = False, ) -> None: """Launch API server with two different sets of arguments/environments and compare the results of the API calls. @@ -1178,9 +1177,6 @@ def compare_two_settings( env2: The second set of environment variables to pass to the API server. include_seeded_sampling: Whether to include temperature=1.0 seeded sampling checks in the default generate comparison. - force_v1_runner: Whether to pin all compared settings to the v1 model - runner to avoid mixing model runner differences into correctness - tests. """ compare_all_settings( @@ -1190,7 +1186,6 @@ def compare_two_settings( method=method, max_wait_seconds=max_wait_seconds, include_seeded_sampling=include_seeded_sampling, - force_v1_runner=force_v1_runner, ) @@ -1202,7 +1197,6 @@ def compare_all_settings( method: str = "generate", max_wait_seconds: float | None = None, include_seeded_sampling: bool = True, - force_v1_runner: bool = False, ) -> None: """Launch API server with several different sets of arguments/environments and compare the results of the API calls with the first set of arguments. @@ -1213,16 +1207,8 @@ def compare_all_settings( all_envs: A list of environment dictionaries to pass to the API server. include_seeded_sampling: Whether to include temperature=1.0 seeded sampling checks in the default generate comparison. - force_v1_runner: Whether to pin all compared settings to the v1 model - runner to avoid mixing model runner differences into correctness - tests. """ - if force_v1_runner: - all_envs = [ - {"VLLM_USE_V2_MODEL_RUNNER": "0", **(env or {})} for env in all_envs - ] - trust_remote_code = False for args in all_args: if "--trust-remote-code" in args: diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 615dd82a1c7b..be3ea553d926 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -352,6 +352,12 @@ def test_async_scheduling_pp_allows_rescheduling_with_output_placeholders(): _ = scheduler.schedule() assert req.num_output_placeholders > 0 + if scheduler.use_v2_model_runner: + # MRV2 staggers same-request decodes by pp_size steps to match the + # worker-side sampled-token broadcast cadence. + output = scheduler.schedule() + assert req.request_id not in output.num_scheduled_tokens + # before any update_from_output, we still expect the request can be # scheduled again (multi-step in-flight). output = scheduler.schedule() diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index b5e1bf89329e..cf221fc22a34 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -238,7 +238,7 @@ def create_scheduler( structured_output_manager=StructuredOutputManager(vllm_config), ) if use_v2_model_runner is None: - use_v2_model_runner = bool(envs.VLLM_USE_V2_MODEL_RUNNER) + use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER is not False scheduler.use_v2_model_runner = use_v2_model_runner return scheduler diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index ed3c6b0217c8..89797db4cdd7 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -122,7 +122,10 @@ def test_replayssm_flashinfer_decode_matches_baseline( @pytest.mark.parametrize("model_name", FLASHINFER_MODELS) -def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_name): +@pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"]) +def test_replayssm_flashinfer_spec_decode_matches_baseline( + vllm_runner, model_name, monkeypatch, use_v2_model_runner +): common = dict( max_model_len=1024, trust_remote_code=True, @@ -134,20 +137,34 @@ def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_na "prompt_lookup_max": 3, }, ) - with vllm_runner(model_name, mamba_backend="flashinfer", **common) as llm: - baseline = llm.generate_greedy_logprobs( - FLASHINFER_PROMPTS, max_tokens=32, num_logprobs=5 - ) - with vllm_runner( - model_name, - use_replayssm=True, - replayssm_buffer_len=16, - mamba_backend="flashinfer", - **common, - ) as llm: - replay = llm.generate_greedy_logprobs( - FLASHINFER_PROMPTS, max_tokens=32, num_logprobs=5 - ) + try: + with monkeypatch.context() as patch: + patch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner))) + envs.disable_envs_cache() + with vllm_runner(model_name, mamba_backend="flashinfer", **common) as llm: + assert ( + llm.llm.llm_engine.vllm_config.use_v2_model_runner + is use_v2_model_runner + ) + baseline = llm.generate_greedy_logprobs( + FLASHINFER_PROMPTS, max_tokens=32, num_logprobs=5 + ) + with vllm_runner( + model_name, + use_replayssm=True, + replayssm_buffer_len=16, + mamba_backend="flashinfer", + **common, + ) as llm: + assert ( + llm.llm.llm_engine.vllm_config.use_v2_model_runner + is use_v2_model_runner + ) + replay = llm.generate_greedy_logprobs( + FLASHINFER_PROMPTS, max_tokens=32, num_logprobs=5 + ) + finally: + envs.disable_envs_cache() check_logprobs_close( outputs_0_lst=baseline, diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index 237b37600401..07d530f97d1a 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -71,7 +71,7 @@ def _make_vllm_config( decode_context_parallel_size=decode_context_parallel_size, ) config.kv_events_config = None - config.use_v2_model_runner = False + config.use_v2_model_runner = True config.kv_transfer_config = KVTransferConfig( kv_connector="OffloadingConnector", kv_role="kv_both", @@ -284,7 +284,7 @@ def _parallelism_agnostic( kv_cache_groups: list[KVCacheGroupSpec], *, canonical: bool = False, - v2: bool = False, + v2: bool = True, ) -> bool: config = _make_vllm_config( extra_config={"canonical_layout": True} if canonical else None @@ -308,7 +308,7 @@ def _replicated_layout( prefill_context_parallel_size: int = 1, decode_context_parallel_size: int = 1, use_mla: bool = True, - use_v2_model_runner: bool = False, + use_v2_model_runner: bool = True, distributed_executor_backend: Any = "mp", nnodes: int = 1, world_size: int | None = None, @@ -1105,7 +1105,10 @@ def test_replicated_layout_parallel_gate(kwargs: dict[str, Any], case: str): def test_parallelism_agnostic_for_single_full_attention_group(): - assert _parallelism_agnostic([KVCacheGroupSpec(["l0"], _full_attention_spec())]) + # The direct (non-canonical) parallelism-agnostic layout is MRV1-only. + assert _parallelism_agnostic( + [KVCacheGroupSpec(["l0"], _full_attention_spec())], v2=False + ) _SWA_SPEC = SlidingWindowSpec( @@ -1147,7 +1150,9 @@ def _uniform_group(*specs: KVCacheSpec) -> KVCacheGroupSpec: ], ) def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]): - assert not _parallelism_agnostic(kv_cache_groups) + # The direct layout is never parallelism-agnostic on MRV2, so check the + # MRV1 gate. + assert not _parallelism_agnostic(kv_cache_groups, v2=False) @pytest.mark.parametrize( @@ -1178,7 +1183,7 @@ def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]): def test_canonical_layout_gate(kv_cache_groups, certified): """The canonical layout certifies portability group by group; none of these shapes are portable in the direct layout.""" - assert not _parallelism_agnostic(kv_cache_groups) + assert not _parallelism_agnostic(kv_cache_groups, v2=False) assert _parallelism_agnostic(kv_cache_groups, canonical=True) is certified @@ -1233,7 +1238,7 @@ def test_canonical_layout_certifies_v2_model_runner(): registration, so the static gate must not depend on the model-runner version — the v2 runner is the case the canonical layout exists for.""" groups = _groups(_full_attention_spec()) - assert _parallelism_agnostic(groups, canonical=True, v2=True) + assert _parallelism_agnostic(groups, canonical=True) @pytest.mark.parametrize("tp_size", [1, 2, 4]) @@ -1275,7 +1280,6 @@ def test_canonical_mla_dsa_rows_are_tp_independent(tp_size, scheduler): def test_parallelism_agnostic_disabled_on_v2_model_runner(): config = _make_vllm_config() - config.use_v2_model_runner = True kv_cache_config = KVCacheConfig( num_blocks=0, kv_cache_tensors=[], diff --git a/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py b/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py index 7c6acfafe321..f310c5a083ab 100644 --- a/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py +++ b/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py @@ -105,7 +105,7 @@ def make_stage(start, end, taps, world_size=2): model_config=SimpleNamespace(hf_text_config=config), quant_config=None, speculative_config=None, - use_v2_model_runner=False, + use_v2_model_runner=True, ) ) model._set_aux_hidden_state_layers(taps) diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index adb8f52e2e8b..a34133e2f8d3 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -424,8 +424,8 @@ def run_model(scheduler_output, intermediate_tensors): pipeline_parallel_size=2, distributed_executor_backend="mp" ), ), - use_v2_model_runner=False, - model_runner=SimpleNamespace(execute_model=run_model), + use_v2_model_runner=True, + model_runner=SimpleNamespace(execute_model=run_model, is_pooling_model=False), annotate_profile=lambda scheduler_output: nullcontext(), _pp_send_work=[previous_tensor_send], ) From f26d9f432ced7b1087d02d29bffed2249df1a18e Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Mon, 12 Oct 2026 01:55:00 +0800 Subject: [PATCH 266/624] [Bugfix][GLM-5.3-Flash] Honor `--mamba-ssm-cache-dtype` for the KDA recurrent state (#59147) Signed-off-by: Jared Wen Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/models/glm5next/test_kda_state_dtype.py | 41 +++++++++++++++++++ vllm/models/glm5next/common/kda.py | 4 +- vllm/models/glm5next/common/model.py | 4 +- 3 files changed, 47 insertions(+), 2 deletions(-) create mode 100644 tests/models/glm5next/test_kda_state_dtype.py diff --git a/tests/models/glm5next/test_kda_state_dtype.py b/tests/models/glm5next/test_kda_state_dtype.py new file mode 100644 index 000000000000..b112e47cd520 --- /dev/null +++ b/tests/models/glm5next/test_kda_state_dtype.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""``--mamba-ssm-cache-dtype`` must reach GLM-5.3-Flash's KDA recurrent state +(it used to be dropped, leaving the state float32 unconditionally).""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.models.glm5next.common.kda import Glm5NextLinearAttention +from vllm.models.glm5next.common.model import ( + Glm5NextForCausalLM, + Glm5NextForConditionalGeneration, +) + + +@pytest.mark.parametrize( + ("ssm_dtype", "expected"), + [("auto", torch.float32), ("float32", torch.float32), ("bfloat16", torch.bfloat16)], +) +def test_kda_recurrent_state_dtype_follows_ssm_cache_dtype(ssm_dtype, expected): + model_config = SimpleNamespace(dtype=torch.bfloat16) + cache_config = SimpleNamespace( + mamba_cache_dtype="auto", + mamba_ssm_cache_dtype=ssm_dtype, + use_kda_recoverssm=False, + ) + vllm_config = SimpleNamespace(model_config=model_config, cache_config=cache_config) + want = (torch.bfloat16, expected) + + assert Glm5NextForCausalLM.get_mamba_state_dtype_from_config(vllm_config) == want + assert ( + Glm5NextForConditionalGeneration.get_mamba_state_dtype_from_config(vllm_config) + == want + ) + # The layer's own dtype must agree with the model-level one, since the + # KV-cache spec and the kernels read them separately. + layer = Glm5NextLinearAttention.__new__(Glm5NextLinearAttention) + layer.model_config, layer.cache_config = model_config, cache_config + assert layer.get_state_dtype() == want diff --git a/vllm/models/glm5next/common/kda.py b/vllm/models/glm5next/common/kda.py index 141a8cc7cb82..68d405b40cfc 100644 --- a/vllm/models/glm5next/common/kda.py +++ b/vllm/models/glm5next/common/kda.py @@ -181,7 +181,9 @@ def get_state_dtype( if self.model_config is None or self.cache_config is None: raise ValueError("model_config and cache_config must be set") return MambaStateDtypeCalculator.kda_state_dtype( - self.model_config.dtype, self.cache_config.mamba_cache_dtype + self.model_config.dtype, + self.cache_config.mamba_cache_dtype, + self.cache_config.mamba_ssm_cache_dtype, ) def get_state_shape( diff --git a/vllm/models/glm5next/common/model.py b/vllm/models/glm5next/common/model.py index e35887ca7c54..1e58166054b0 100644 --- a/vllm/models/glm5next/common/model.py +++ b/vllm/models/glm5next/common/model.py @@ -1092,7 +1092,9 @@ def get_mamba_state_dtype_from_config( vllm_config: "VllmConfig", ) -> tuple[torch.dtype, torch.dtype]: return MambaStateDtypeCalculator.kda_state_dtype( - vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype + vllm_config.model_config.dtype, + vllm_config.cache_config.mamba_cache_dtype, + vllm_config.cache_config.mamba_ssm_cache_dtype, ) @classmethod From 3fe917abb13160d60bea5330399f5c09ca282348 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Sun, 11 Oct 2026 16:04:41 -0700 Subject: [PATCH 267/624] [Bugfix][MRV2] Fix FlashInfer attention and ReplaySSM warmups on MRV2 (#61013) Signed-off-by: Nick Hill --- tests/model_executor/test_attention_warmup.py | 71 +++++++++++++++++++ .../model_executor/warmup/attention_warmup.py | 42 +++++++++++ vllm/model_executor/warmup/kernel_warmup.py | 13 ++-- .../warmup/minimax_m3_msa_warmup.py | 11 ++- vllm/v1/worker/gpu/model_runner.py | 9 ++- 5 files changed, 129 insertions(+), 17 deletions(-) create mode 100644 tests/model_executor/test_attention_warmup.py create mode 100644 vllm/model_executor/warmup/attention_warmup.py diff --git a/tests/model_executor/test_attention_warmup.py b/tests/model_executor/test_attention_warmup.py new file mode 100644 index 000000000000..861276280e24 --- /dev/null +++ b/tests/model_executor/test_attention_warmup.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import pytest + +from vllm.model_executor.warmup import attention_warmup +from vllm.model_executor.warmup.attention_warmup import mixed_batch_attention_warmup + + +def _worker(*, use_v2: bool, max_num_batched_tokens: int, max_model_len: int = 4096): + calls = [] + worker = SimpleNamespace( + model_runner=SimpleNamespace(_dummy_run=lambda **kw: calls.append(kw)), + use_v2_model_runner=use_v2, + scheduler_config=SimpleNamespace(max_num_batched_tokens=max_num_batched_tokens), + model_config=SimpleNamespace(max_model_len=max_model_len), + execute_model=object(), + sample_tokens=object(), + ) + return worker, calls + + +@pytest.mark.parametrize( + ("max_num_batched_tokens", "max_model_len", "expected"), + [(8192, 4096, 16), (6, 4096, 6), (8192, 8, 8)], +) +def test_v1_runs_mixed_dummy_batch_within_limits( + max_num_batched_tokens, max_model_len, expected +): + worker, calls = _worker( + use_v2=False, + max_num_batched_tokens=max_num_batched_tokens, + max_model_len=max_model_len, + ) + mixed_batch_attention_warmup(worker) + assert calls == [ + dict( + num_tokens=expected, + skip_eplb=True, + is_profile=True, + force_attention=True, + create_mixed_batch=True, + ) + ] + + +@pytest.mark.parametrize(("max_num_batched_tokens", "expected"), [(8192, 16), (6, 6)]) +@pytest.mark.parametrize("ran", [True, False]) +def test_v2_runs_mixed_prefill_decode_step( + monkeypatch, max_num_batched_tokens, expected, ran +): + worker, dummy_calls = _worker( + use_v2=True, max_num_batched_tokens=max_num_batched_tokens + ) + mixed_calls = [] + + def fake_mixed_warmup(runner, execute_model, sample_tokens, num_tokens): + mixed_calls.append((runner, execute_model, sample_tokens, num_tokens)) + return ran + + monkeypatch.setattr( + attention_warmup, "run_mixed_prefill_decode_warmup", fake_mixed_warmup + ) + mixed_batch_attention_warmup(worker) + + assert mixed_calls == [ + (worker.model_runner, worker.execute_model, worker.sample_tokens, expected) + ] + # No dummy-run fallback when the mixed step is skipped. + assert dummy_calls == [] diff --git a/vllm/model_executor/warmup/attention_warmup.py b/vllm/model_executor/warmup/attention_warmup.py new file mode 100644 index 000000000000..cae556a49a0f --- /dev/null +++ b/vllm/model_executor/warmup/attention_warmup.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import TYPE_CHECKING, cast + +from vllm.logger import init_logger +from vllm.v1.worker.gpu.warmup import run_mixed_prefill_decode_warmup + +if TYPE_CHECKING: + from vllm.v1.worker.gpu.model_runner import GPUModelRunner + from vllm.v1.worker.gpu_worker import Worker + +logger = init_logger(__name__) + + +def mixed_batch_attention_warmup(worker: "Worker", num_tokens: int = 16) -> None: + """Run a mixed prefill+decode batch to warm up both attention paths.""" + num_tokens = min( + num_tokens, + worker.scheduler_config.max_num_batched_tokens, + worker.model_config.max_model_len, + ) + if not worker.use_v2_model_runner: + worker.model_runner._dummy_run( + num_tokens=num_tokens, + skip_eplb=True, + is_profile=True, + force_attention=True, + create_mixed_batch=True, + ) + return + + if not run_mixed_prefill_decode_warmup( + cast("GPUModelRunner", worker.model_runner), + worker.execute_model, + worker.sample_tokens, + num_tokens, + ): + logger.info_once( + "Skipping mixed prefill+decode attention warmup; attention kernels " + "will be loaded on the first request." + ) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 685b33a5652a..b7bf50865fca 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -15,6 +15,9 @@ import vllm.envs as envs from vllm.logger import init_logger +from vllm.model_executor.warmup.attention_warmup import ( + mixed_batch_attention_warmup, +) from vllm.model_executor.warmup.b12x_warmup import b12x_warmup from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup @@ -286,15 +289,7 @@ def _is_flashinfer_backend(backend): ) ): logger.info_once("Warming up FlashInfer attention.") - # Warmup with mixed batch containing both prefill and decode tokens - # This is to warm up both prefill and decode attention kernels - worker.model_runner._dummy_run( - num_tokens=16, - skip_eplb=True, - is_profile=True, - force_attention=True, - create_mixed_batch=True, - ) + mixed_batch_attention_warmup(worker) def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None: diff --git a/vllm/model_executor/warmup/minimax_m3_msa_warmup.py b/vllm/model_executor/warmup/minimax_m3_msa_warmup.py index 18bf14249111..2c8996a93341 100644 --- a/vllm/model_executor/warmup/minimax_m3_msa_warmup.py +++ b/vllm/model_executor/warmup/minimax_m3_msa_warmup.py @@ -4,6 +4,9 @@ from typing import TYPE_CHECKING from vllm.logger import init_logger +from vllm.model_executor.warmup.attention_warmup import ( + mixed_batch_attention_warmup, +) from vllm.models.minimax_m3.nvidia.model import MiniMaxM3SparseAttention from vllm.platforms import current_platform from vllm.tracing import instrument @@ -34,10 +37,4 @@ def minimax_m3_msa_warmup(worker: "Worker") -> None: logger.info("Warming up MiniMax M3 MSA kernels.") # Cover sparse prefill through the normal model path. - worker.model_runner._dummy_run( - num_tokens=16, - skip_eplb=True, - is_profile=True, - force_attention=True, - create_mixed_batch=True, - ) + mixed_batch_attention_warmup(worker) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 20fa4c17bb07..57cef586f81f 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -857,6 +857,11 @@ def _dummy_run( - num_reqs * self.speculator.num_extra_query_per_req, ) + assert num_tokens <= self.max_num_tokens, ( + f"Dummy run with {num_tokens} tokens exceeds max_num_batched_tokens " + f"({self.max_num_tokens})" + ) + # Distribute the remainder evenly so no dummy request exceeds # ceil(num_tokens / num_reqs) <= max_model_len tokens. @@ -1898,9 +1903,11 @@ def execute_model( ) assert block_tables is not None attn_groups = self.attn_groups - if dummy_run and is_profile: + if dummy_run and is_profile and not valid_dummy_state_slots: # Mamba layers take a cheap warmup path with no metadata; # attention metadata is still built so those kernels tune. + # valid_dummy_state_slots runs the real Mamba kernels on distinct + # state slots; the caller must restore that state afterwards. attn_groups = [ [g for g in groups if not isinstance(g.kv_cache_spec, MambaSpec)] for groups in attn_groups From cc9b0f097d49abfa6fc94d1d773dca1cc2b9619c Mon Sep 17 00:00:00 2001 From: Laith Sakka Date: Sun, 11 Oct 2026 17:38:56 -0700 Subject: [PATCH 268/624] [Fix] MiniMax: fix CUDA IPC handle exchange on cuda-bindings >= 13.4 (#60917) Signed-off-by: Laith Sakka --- .../minimax_rms_norm/lamport_workspace.py | 23 ++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py b/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py index 0b2ef3a6aabd..eac2a64a0457 100644 --- a/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py +++ b/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py @@ -4,6 +4,7 @@ import array import contextlib +import ctypes import struct import sys import threading @@ -16,6 +17,7 @@ from cuda import cudart _ALIGN = 1 << 21 # 2 MiB — CUDA IPC allocation alignment +_IPC_HANDLE_SIZE = 64 # sizeof(cudaIpcMemHandle_t) # --------------------------------------------------------------------------- @@ -30,6 +32,20 @@ def _check(error): raise RuntimeError(f"CUDA runtime error: {error}") +# cuda-bindings >= 13.4 no longer exposes cudaIpcMemHandle_t.reserved, so +# copy the raw handle bytes through getPtr() instead. This is safe because the +# CUDA ABI fixes the handle as an opaque 64-byte struct (char reserved[64]). +def _ipc_handle_to_bytes(handle) -> bytes: + return ctypes.string_at(handle.getPtr(), _IPC_HANDLE_SIZE) + + +def _ipc_handle_from_bytes(raw: bytes): + assert len(raw) == _IPC_HANDLE_SIZE, len(raw) + handle = cudart.cudaIpcMemHandle_t() + ctypes.memmove(handle.getPtr(), raw, _IPC_HANDLE_SIZE) + return handle + + def _cuda_malloc(size: int): aligned = ((size + _ALIGN - 1) >> 21) << 21 err, ptr = cudart.cudaMalloc(aligned) @@ -84,15 +100,16 @@ def __init__(self, rank: int, world_size: int, size: int, process_group=None): all_handles: list[bytes | None] = [None] * world_size torch.distributed.all_gather_object( - all_handles, bytes(local_handle.reserved), group=process_group + all_handles, _ipc_handle_to_bytes(local_handle), group=process_group ) for r in range(world_size): if r == rank: self.peer_ptrs[r] = self.local_ptr else: - handle = cudart.cudaIpcMemHandle_t() - handle.reserved = all_handles[r] + raw = all_handles[r] + assert raw is not None + handle = _ipc_handle_from_bytes(raw) err, ptr = cudart.cudaIpcOpenMemHandle( handle, cudart.cudaIpcMemLazyEnablePeerAccess ) From 03091e28b8ba956634c153beab0e52decafdef31 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Ganczarenko?= Date: Mon, 12 Oct 2026 03:25:29 +0200 Subject: [PATCH 269/624] [Bugfix] Read Ling quantization_config from top-level config for Bailing V3 VL (#59026) Signed-off-by: Michal Ganczarenko Co-authored-by: Claude Opus 5.5 --- vllm/model_executor/models/bailing_moe_v3.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/bailing_moe_v3.py b/vllm/model_executor/models/bailing_moe_v3.py index 9f5c078fac9c..d82e4c1b8b41 100644 --- a/vllm/model_executor/models/bailing_moe_v3.py +++ b/vllm/model_executor/models/bailing_moe_v3.py @@ -20,7 +20,13 @@ from transformers.configuration_utils import PreTrainedConfig from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config +from vllm.config import ( + CacheConfig, + ModelConfig, + VllmConfig, + get_current_vllm_config, + get_current_vllm_config_or_none, +) from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_rank, @@ -230,6 +236,12 @@ def _configure_ling_fp8_quant_config( quant_config.ignored_layers_match_mode = "suffix" hf_quant_config = getattr(config, "quantization_config", None) + if hf_quant_config is None: + vllm_config = get_current_vllm_config_or_none() + model_config = getattr(vllm_config, "model_config", None) + hf_quant_config = getattr( + getattr(model_config, "hf_config", None), "quantization_config", None + ) if not isinstance(hf_quant_config, dict): return From 124b89d5b7bbbf7599763858c247848e073b024a Mon Sep 17 00:00:00 2001 From: dingdangmao <2756823972@qq.com> Date: Mon, 12 Oct 2026 10:08:14 +0800 Subject: [PATCH 270/624] [Bugfix][RecoverSSM] Correct aligned state indices at exact block boundaries (#59962) Signed-off-by: apex-mochen <2756823972@qq.com> Co-authored-by: Codex --- tests/models/kimi_k3/test_kda.py | 9 ++++++--- tests/v1/worker/test_mamba_hybrid_model_state.py | 16 ++++++++++++---- vllm/models/kimi_k3/nvidia/ops/recoverssm.py | 3 ++- vllm/v1/worker/gpu/model_states/recoverssm.py | 2 +- 4 files changed, 21 insertions(+), 9 deletions(-) diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index 88acdb6197dc..0b250654450e 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -692,8 +692,10 @@ def test_kda_spec_decode_correctness( pytest.param(True, False, None, True, id="aligned"), ], ) +@pytest.mark.parametrize("accepted", [(2, 8), (3, 4), (5, 8)]) @torch.inference_mode() def test_kda_recoverssm_verify_and_group_commit( + accepted: tuple[int, int], monkeypatch: pytest.MonkeyPatch, lower_bound: float | None, use_request_indices: bool, @@ -738,7 +740,6 @@ def test_kda_recoverssm_verify_and_group_commit( state_indices = torch.tensor( [5, 6] if align_mode else [1, 2], dtype=torch.int32, device=DEVICE ) - accepted = [2, 8] if use_request_indices: global_num_accepted = torch.tensor( [0, accepted[0], 0, accepted[1]], @@ -891,7 +892,9 @@ def test_kda_recoverssm_verify_and_group_commit( if align_mode: assert block_table is not None row = request_indices[seq_idx] if use_request_indices else seq_idx - final_block = block_table[row, (4 + commit_len) // 8] + final_block = block_table[ + row, {2: 0, 3: 0, 4: 0, 5: 1, 8: 1}[commit_len] + ] committed_states[final_block] = committed_state.transpose(-1, -2) if align_mode and 4 + commit_len >= 8: _, boundary_state = naive_recurrent_kda( @@ -944,7 +947,7 @@ def test_kda_recoverssm_verify_and_group_commit( if align_mode: assert block_table is not None row = request_indices[seq_idx] if use_request_indices else seq_idx - block = block_table[row, (4 + commit_len) // 8] + block = block_table[row, {2: 0, 3: 0, 4: 0, 5: 1, 8: 1}[commit_len]] source_block = state_indices[seq_idx] if align_mode else block if conv_state_dim_first: actual_conv = layer.kv_cache[0][block, :, :history_len] diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index 0dcc4725d779..b52af4224974 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -387,7 +387,13 @@ def test_recoverssm_commits_accepted_window_after_v2_sampling() -> None: @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") -def test_recoverssm_align_tracks_mixed_batch_state_and_neutralizes_copy_bias() -> None: +@pytest.mark.parametrize( + ("computed", "sampled", "expected_column"), + [(4, 3, 0), (4, 4, 0), (4, 5, 1), (8, 8, 1), (8, 9, 2)], +) +def test_recoverssm_align_tracks_mixed_batch_state_and_neutralizes_copy_bias( + computed: int, sampled: int, expected_column: int +) -> None: state = object.__new__(MambaHybridModelState) state._align_mode = True state._mamba_ctx = None @@ -400,11 +406,13 @@ def test_recoverssm_align_tracks_mixed_batch_state_and_neutralizes_copy_bias() - metadata.commit_recoverssm_state.return_value = RecoverSSMPostprocessMetadata( num_spec_decodes=1, request_indices=torch.tensor([1], dtype=torch.int32, device="cuda"), - num_computed_tokens=torch.tensor([6, 7], dtype=torch.int32, device="cuda"), + num_computed_tokens=torch.tensor( + [6, computed], dtype=torch.int32, device="cuda" + ), block_size=8, block_table=torch.zeros((2, 4), dtype=torch.int32, device="cuda"), ) - num_sampled = torch.tensor([2, 3], dtype=torch.int32, device="cuda") + num_sampled = torch.tensor([2, sampled], dtype=torch.int32, device="cuda") idx_mapping = torch.tensor([3, 1], dtype=torch.int32, device="cuda") group = SimpleNamespace(layer_names=["layer"]) @@ -412,7 +420,7 @@ def test_recoverssm_align_tracks_mixed_batch_state_and_neutralizes_copy_bias() - state.postprocess_state(idx_mapping, num_sampled) - expected_state_indices = [-1, 1, -1, -1, -1] + expected_state_indices = [-1, expected_column, -1, -1, -1] assert state._mamba_state_idx_gpu.tolist() == expected_state_indices expected_accepted = [9, 1, 9, 2, 9] assert state.num_accepted_tokens_gpu.tolist() == expected_accepted diff --git a/vllm/models/kimi_k3/nvidia/ops/recoverssm.py b/vllm/models/kimi_k3/nvidia/ops/recoverssm.py index c603ec7ed1fc..7eda21f1a5ab 100644 --- a/vllm/models/kimi_k3/nvidia/ops/recoverssm.py +++ b/vllm/models/kimi_k3/nvidia/ops/recoverssm.py @@ -284,7 +284,8 @@ def _prepare_commit_plan_kernel( ) final_num_computed = num_computed + commit_len final_state_col = tl.minimum( - final_num_computed // mamba_block_size, block_table_width - 1 + tl.maximum(final_num_computed - 1, 0) // mamba_block_size, + block_table_width - 1, ) final_state_idx = tl.load( block_table_ptr diff --git a/vllm/v1/worker/gpu/model_states/recoverssm.py b/vllm/v1/worker/gpu/model_states/recoverssm.py index cfbe1a4d45d1..985f29079814 100644 --- a/vllm/v1/worker/gpu/model_states/recoverssm.py +++ b/vllm/v1/worker/gpu/model_states/recoverssm.py @@ -94,7 +94,7 @@ def _postprocess_recoverssm_align_kernel( tl.store( state_idx_ptr + req_state_idx, tl.minimum( - (num_computed + num_sampled) // MAMBA_BLOCK_SIZE, + tl.maximum(num_computed + num_sampled - 1, 0) // MAMBA_BLOCK_SIZE, BLOCK_TABLE_WIDTH - 1, ), ) From f147949be8b9e576f80133bee01e717ee7c5db0e Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Mon, 12 Oct 2026 10:25:32 +0800 Subject: [PATCH 271/624] [XPU] limit FA kernel block size to multipleOf 64 (#57297) Signed-off-by: Yan Ma --- .buildkite/intel_jobs/test-intel.yaml | 2 +- .../unit/test_bidirectional_kv_transfer.py | 6 +- .../unit/test_mooncake_connector.py | 62 ++--- .../unit/test_mooncake_connector_hma.py | 34 ++- .../test_mooncake_connector_hybrid_mamba.py | 16 +- .../kv_connector/unit/test_multi_connector.py | 21 +- .../kv_connector/unit/test_nixl_connector.py | 224 ++++++++++++------ .../unit/test_nixl_connector_hma.py | 3 +- tests/v1/kv_connector/unit/utils.py | 5 + vllm/v1/attention/backends/flash_attn.py | 2 + vllm/v1/worker/gpu/attn_utils.py | 20 +- 11 files changed, 269 insertions(+), 126 deletions(-) diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index a2a0e14be18c..1567c0eba724 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -105,7 +105,7 @@ steps: pytest -v -s v1/test_serial_utils.py && pytest -v -s v1/e2e/general/test_correctness_sliding_window.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 --deselect=tests/v1/spec_decode/test_mtp.py::test_glm_mtp_defers_lm_head && - pytest -v -s v1/kv_connector/unit --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 --ignore=v1/kv_connector/unit/test_nixl_connector_hma.py --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --deselect=tests/v1/kv_connector/unit/test_nixl_connector.py::test_abort_timeout_on_prefiller[ray]' + pytest -v -s v1/kv_connector/unit --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 --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --deselect=tests/v1/kv_connector/unit/test_nixl_connector.py::test_abort_timeout_on_prefiller[ray]' parallelism: 2 - label: "XPU Qwen3-30B + V1 KV Connector" key: xpu-qwen3-30b-v1-kv-connector diff --git a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py index 3d539f97c8cb..2aae7d71f515 100644 --- a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py +++ b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py @@ -51,6 +51,7 @@ create_scheduler, create_vllm_config, make_kv_cache_config, + maybe_update_block_size, ) pytestmark = pytest.mark.cpu_test @@ -87,10 +88,11 @@ def _make_connector_with_fake_worker( hand_shake_latency=0, cycles_before_done=0, do_handshake=True ): """Create a NixlConnector with FakeNixlConnectorWorker.""" + block_size = maybe_update_block_size(32) vllm_config = create_vllm_config( - kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, block_size=32 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, block_size=block_size ) - kv_cache_config = make_kv_cache_config(block_size=32, num_blocks=2) + kv_cache_config = make_kv_cache_config(block_size=block_size, num_blocks=2) connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index d326615aad8f..f79d7aa316f0 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -57,7 +57,12 @@ ) from vllm.v1.request import RequestStatus -from .utils import create_request, create_scheduler, create_vllm_config +from .utils import ( + create_request, + create_scheduler, + create_vllm_config, + maybe_update_block_size, +) @pytest.mark.parametrize( @@ -148,6 +153,7 @@ def test_region_length_validation_checks_replicated_gqa_heads( def _make_test_kv_cache_config() -> KVCacheConfig: + block_size = maybe_update_block_size(16) return KVCacheConfig( num_blocks=0, kv_cache_tensors=[], @@ -160,7 +166,7 @@ def _make_test_kv_cache_config() -> KVCacheConfig: "model.layers.1.eagle_attn", ], FullAttentionSpec( - block_size=16, + block_size=block_size, num_kv_heads=4, head_size=64, dtype=torch.float16, @@ -170,6 +176,14 @@ def _make_test_kv_cache_config() -> KVCacheConfig: ) +def _make_worker_vllm_config(kv_role: str): + return create_vllm_config( + kv_connector="MooncakeConnector", + kv_role=kv_role, + block_size=maybe_update_block_size(16), + ) + + def _make_packed_mla_kv_cache_config( num_blocks: int, *, num_groups: int = 1 ) -> KVCacheConfig: @@ -429,9 +443,7 @@ async def test_send_kv_to_decode_aligns_consumer_regions_by_layer_metadata( ): """Producer sends its PP layer shard to the matching consumer layer address.""" monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5") - vllm_config = create_vllm_config( - kv_connector="MooncakeConnector", kv_role="kv_producer" - ) + vllm_config = _make_worker_vllm_config("kv_producer") with set_current_vllm_config(vllm_config), patch_worker_dependencies(): prefill_connector = MooncakeConnector( @@ -1205,7 +1217,7 @@ def _xfer_meta( @contextlib.contextmanager def mooncake_register_worker(kv_cache_config, *, kv_role: str = "kv_consumer"): """Build a worker whose engine registration and sender thread are patched.""" - vllm_config = create_vllm_config(kv_connector="MooncakeConnector", kv_role=kv_role) + vllm_config = _make_worker_vllm_config(kv_role) connector_mod = ( "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector" ) @@ -1241,9 +1253,7 @@ async def test_receive_kv_selects_remote_pp_workers( expected_addrs: list[str], ): """Decode workers should not hard-code producer pp_rank 0.""" - vllm_config = create_vllm_config( - kv_connector="MooncakeConnector", kv_role="kv_consumer" - ) + vllm_config = _make_worker_vllm_config("kv_consumer") with set_current_vllm_config(vllm_config), patch_worker_dependencies(): decode_connector = MooncakeConnector( @@ -1314,7 +1324,7 @@ def test_resolve_need_send_accounts_for_remote_tp_fanout(): @pytest.mark.asyncio async def test_heterogeneous_pp_waits_for_consumer_without_shared_layers(): """P4/D2 must retain KV until both D stages finish, including a no-op pull.""" - config = create_vllm_config(kv_connector="MooncakeConnector", kv_role="kv_producer") + config = _make_worker_vllm_config("kv_producer") with ( set_current_vllm_config(config), patch_worker_dependencies(), @@ -1386,9 +1396,7 @@ async def test_kv_producer(monkeypatch): Verifies memory offset calculation: ptr = base_addr + block_id * block_len. """ monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5") - vllm_config = create_vllm_config( - kv_connector="MooncakeConnector", kv_role="kv_producer" - ) + vllm_config = _make_worker_vllm_config("kv_producer") with set_current_vllm_config(vllm_config), patch_worker_dependencies(): prefill_connector = MooncakeConnector( @@ -1589,9 +1597,7 @@ async def test_kv_consumuer(monkeypatch): Verifies that MooncakeXferMetadata is correctly serialized and sent via ZMQ. """ - vllm_config = create_vllm_config( - kv_connector="MooncakeConnector", kv_role="kv_consumer" - ) + vllm_config = _make_worker_vllm_config("kv_consumer") with set_current_vllm_config(vllm_config), patch_worker_dependencies() as mocks: decode_connector = MooncakeConnector( @@ -1661,9 +1667,7 @@ async def test_kv_consumuer(monkeypatch): @pytest.mark.asyncio async def test_worker_get_finished_timeout(monkeypatch): """Tests the cleanup mechanism for requests.""" - vllm_config = create_vllm_config( - kv_connector="MooncakeConnector", kv_role="kv_producer" - ) + vllm_config = _make_worker_vllm_config("kv_producer") with set_current_vllm_config(vllm_config), patch_worker_dependencies(): prefill_connector = MooncakeConnector( vllm_config, @@ -1713,7 +1717,7 @@ async def test_worker_get_finished_timeout(monkeypatch): def test_register_kv_caches(layout: KVCacheLayout, separate_kv_head_groups: bool): """Tests the memory registration logic with the underlying Mooncake engine.""" spec = FullAttentionSpec( - block_size=16, + block_size=maybe_update_block_size(16), num_kv_heads=4, head_size=64, dtype=torch.float16, @@ -1794,7 +1798,10 @@ def test_register_noncontiguous_packed_allows_matching_pp(): known. Rejecting every local PP>1 layout at startup is too strict. """ spec = FullAttentionSpec( - block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16 + block_size=maybe_update_block_size(16), + num_kv_heads=4, + head_size=64, + dtype=torch.float16, ) layer_names = ["model.layers.0.self_attn", "model.layers.1.self_attn"] with mooncake_register_worker(_make_test_kv_cache_config()) as (worker, _reg): @@ -1814,8 +1821,9 @@ def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes(): with mooncake_register_worker(_make_test_kv_cache_config()) as (worker, reg): worker.use_mla = True worker.transfer_topo.is_mla = True - mla_cache = torch.zeros((2, 16, 512), dtype=torch.float16) - eagle_cache = torch.zeros((2, 16, 8, 64), dtype=torch.float16) + block_size = maybe_update_block_size(16) + mla_cache = torch.zeros((2, block_size, 512), dtype=torch.float16) + eagle_cache = torch.zeros((2, block_size, 8, 64), dtype=torch.float16) kv_caches = { "model.layers.0.mla_attn": mla_cache, "model.layers.1.eagle_attn": eagle_cache, @@ -2240,7 +2248,7 @@ def _register_sliced_packed_mla( ): """Register column slices of one MLA packed row. Returns worker fields.""" spec = MLAAttentionSpec( - block_size=16, + block_size=maybe_update_block_size(16), num_kv_heads=1, head_size=64, dtype=torch.uint8, @@ -2286,7 +2294,7 @@ def test_sliced_packed_mla_registers_each_layer_view(): KVCacheGroupSpec( [_layer_name(i) for i in layer_idxs], MLAAttentionSpec( - block_size=16, + block_size=maybe_update_block_size(16), num_kv_heads=1, head_size=64, dtype=torch.uint8, @@ -2436,9 +2444,7 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size): remote_block_len = LOCAL_BLOCK_LEN * P_TP_SIZE // d_tp_size monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5") - vllm_config = create_vllm_config( - kv_connector="MooncakeConnector", kv_role="kv_producer" - ) + vllm_config = _make_worker_vllm_config("kv_producer") with set_current_vllm_config(vllm_config), patch_worker_dependencies(): prefill_connector = MooncakeConnector( diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py b/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py index a7fb2be02365..87dafebb4c2e 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py @@ -30,7 +30,12 @@ ) from .test_mooncake_connector import FakeMooncakeWrapper, patch_worker_dependencies -from .utils import create_request, create_vllm_config, make_kv_cache_config +from .utils import ( + create_request, + create_vllm_config, + make_kv_cache_config, + maybe_update_block_size, +) # --------------------------------------------------------------------------- @@ -224,12 +229,11 @@ def test_metadata_hma_block_ids(): async def test_build_transfer_params_multi_group_trimming(monkeypatch): """_build_transfer_params trims per-group blocks when local > remote.""" monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5") + block_size = maybe_update_block_size(16) vllm_config = create_vllm_config( - kv_connector="MooncakeConnector", kv_role="kv_producer" - ) - kv_cache_config = make_kv_cache_config( - block_size=vllm_config.cache_config.block_size, swa_enabled=True + kv_connector="MooncakeConnector", kv_role="kv_producer", block_size=block_size ) + kv_cache_config = make_kv_cache_config(block_size=block_size, swa_enabled=True) with set_current_vllm_config(vllm_config), patch_worker_dependencies(): connector = MooncakeConnector( @@ -320,12 +324,11 @@ async def test_build_transfer_params_multi_group_trimming(monkeypatch): async def test_build_transfer_params_group_count_mismatch(monkeypatch): """_build_transfer_params reports an error when group counts differ.""" monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5") + block_size = maybe_update_block_size(16) vllm_config = create_vllm_config( - kv_connector="MooncakeConnector", kv_role="kv_producer" - ) - kv_cache_config = make_kv_cache_config( - block_size=vllm_config.cache_config.block_size, swa_enabled=True + kv_connector="MooncakeConnector", kv_role="kv_producer", block_size=block_size ) + kv_cache_config = make_kv_cache_config(block_size=block_size, swa_enabled=True) with set_current_vllm_config(vllm_config), patch_worker_dependencies(): connector = MooncakeConnector( @@ -449,7 +452,7 @@ def test_request_finished_with_hma_groups(): def _make_kv_consumer_worker( swa_enabled: bool = False, kv_cache_config: KVCacheConfig | None = None ): - block_size = 16 + block_size = maybe_update_block_size(16) vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_consumer", @@ -495,6 +498,7 @@ def test_worker_is_hma_required(swa_enabled, expected_is_hma): def test_worker_is_hma_required_multiple_full_attention_groups(): """Two full-attention groups with different block sizes also select request-level failure reporting, matching the scheduler.""" + block_size = maybe_update_block_size(16) kv_cache_config = KVCacheConfig( num_blocks=100, kv_cache_tensors=[], @@ -502,13 +506,19 @@ def test_worker_is_hma_required_multiple_full_attention_groups(): KVCacheGroupSpec( ["layer0"], FullAttentionSpec( - block_size=16, num_kv_heads=4, head_size=16, dtype=torch.float16 + block_size=block_size, + num_kv_heads=4, + head_size=16, + dtype=torch.float16, ), ), KVCacheGroupSpec( ["layer1"], FullAttentionSpec( - block_size=32, num_kv_heads=4, head_size=16, dtype=torch.float16 + block_size=2 * block_size, + num_kv_heads=4, + head_size=16, + dtype=torch.float16, ), ), ], diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py b/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py index 969bf65f8c54..65abd910bbcf 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py @@ -36,7 +36,7 @@ ) from .test_mooncake_connector import patch_worker_dependencies -from .utils import create_request, create_vllm_config +from .utils import create_request, create_vllm_config, maybe_update_block_size def noop_shutdown(): @@ -147,6 +147,7 @@ def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch): vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_consumer", + block_size=maybe_update_block_size(16), ) kv_cache_config = make_hybrid_gdn_kv_cache_config( vllm_config.cache_config.block_size @@ -203,6 +204,7 @@ def test_register_kv_caches_scales_attention_len_to_kernel_block(monkeypatch): vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_consumer", + block_size=maybe_update_block_size(16), ) kv_cache_config = make_hybrid_gdn_kv_cache_config( vllm_config.cache_config.block_size @@ -392,6 +394,7 @@ def test_register_kv_caches_deduplicates_shared_backing_memory(monkeypatch): vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_consumer", + block_size=maybe_update_block_size(16), ) kv_cache_config = make_hybrid_gdn_kv_cache_config( vllm_config.cache_config.block_size @@ -405,8 +408,13 @@ def test_register_kv_caches_deduplicates_shared_backing_memory(monkeypatch): ) worker = connector.connector_worker - backing = torch.empty((4, 64), dtype=torch.float16) - fa_cache = backing[:2, :16] + fa_spec = kv_cache_config.kv_cache_groups[0].kv_cache_spec + fa_cache_width = ( + fa_spec.page_size_bytes + // torch.tensor([], dtype=torch.float16).element_size() + ) + backing = torch.empty((4, fa_cache_width), dtype=torch.float16) + fa_cache = backing[:2] gdn_cache = backing[:3] with patch.object( @@ -438,6 +446,7 @@ def test_hybrid_gdn_transfer_params_preserve_group_identity(monkeypatch): vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_producer", + block_size=maybe_update_block_size(16), ) kv_cache_config = make_hybrid_gdn_kv_cache_config( vllm_config.cache_config.block_size @@ -567,6 +576,7 @@ def test_hybrid_gdn_keeps_packed_fa_and_gdn_regions_whole( vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_producer", + block_size=maybe_update_block_size(16), ) kv_cache_config = make_hybrid_gdn_kv_cache_config( vllm_config.cache_config.block_size diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index e576e3d0d9e6..4ac3eaf7c569 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -10,7 +10,7 @@ import pytest import torch -from tests.v1.kv_connector.unit.utils import create_vllm_config +from tests.v1.kv_connector.unit.utils import create_vllm_config, maybe_update_block_size from vllm import LLM, SamplingParams from vllm.config import KVTransferConfig from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory @@ -314,10 +314,13 @@ def test_multi_example_connector_consistency(): }, ) + block_size = maybe_update_block_size(16) + expected_num_blocks = 2 if block_size == 64 else 7 + expected_num_matched_tokens = (expected_num_blocks - 1) * block_size llm = LLM( model=MODEL_NAME, enforce_eager=True, - block_size=16, + block_size=block_size, gpu_memory_utilization=0.5, kv_transfer_config=kv_transfer_config, async_scheduling=False, @@ -372,7 +375,7 @@ def test_multi_example_connector_consistency(): "set_xfer_handshake_metadata_pp_aware", "on_new_request", "get_num_new_matched_tokens 0", - "update_state_after_alloc num_blocks=[7] 0", + f"update_state_after_alloc num_blocks=[{expected_num_blocks}] 0", "build_connector_meta", ] # First three events are from initialization. Layer hooks run before the @@ -399,7 +402,7 @@ def test_multi_example_connector_consistency(): "set_xfer_handshake_metadata_pp_aware", "on_new_request", "get_num_new_matched_tokens 0", - "update_state_after_alloc num_blocks=[7] 0", + f"update_state_after_alloc num_blocks=[{expected_num_blocks}] 0", "build_connector_meta", ] # Reset prefix cache or else we'll just get the tokens back from there. @@ -423,13 +426,14 @@ def test_multi_example_connector_consistency(): assert storage1_scheduler_events[:4] == [ "on_new_request", "get_num_new_matched_tokens 0", - "update_state_after_alloc num_blocks=[7] 96", + "update_state_after_alloc " + f"num_blocks=[{expected_num_blocks}] {expected_num_matched_tokens}", "build_connector_meta", ] assert storage2_scheduler_events[:4] == [ "on_new_request", "get_num_new_matched_tokens 0", - "update_state_after_alloc num_blocks=[7] 0", + f"update_state_after_alloc num_blocks=[{expected_num_blocks}] 0", "build_connector_meta", ] @@ -457,13 +461,14 @@ def test_multi_example_connector_consistency(): assert storage1_scheduler_events[:4] == [ "on_new_request", "get_num_new_matched_tokens 0", - "update_state_after_alloc num_blocks=[7] 0", + f"update_state_after_alloc num_blocks=[{expected_num_blocks}] 0", "build_connector_meta", ] assert storage2_scheduler_events[:4] == [ "on_new_request", "get_num_new_matched_tokens 0", - "update_state_after_alloc num_blocks=[7] 96", + "update_state_after_alloc " + f"num_blocks=[{expected_num_blocks}] {expected_num_matched_tokens}", "build_connector_meta", ] diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index d918c38fbf8c..607f5ff3cd5e 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -84,6 +84,7 @@ create_scheduler, create_vllm_config, make_kv_cache_config, + maybe_update_block_size, ) @@ -418,7 +419,8 @@ def test_kv_transfer_handshake(dist_init): # Test setup, we creates a scheduler that contains a NixlConnector # of role SCHEDULER, and expect it to be serving NixlAgentMetadata from # all workers of the instance. - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) vllm_config.cache_config.kv_cache_layout = "BLHNC" # in case the test runs on non-GPU machine vllm_config.kv_transfer_config.kv_buffer_device = "cpu" @@ -434,7 +436,7 @@ def test_kv_transfer_handshake(dist_init): KVCacheGroupSpec( ["layer0", "layer1", "layer2"], FullAttentionSpec( - block_size=16, + block_size=block_size, num_kv_heads=4, head_size=16, dtype=torch.float16, @@ -544,7 +546,9 @@ def __init__( **kwargs, ): if kv_cache_config is None: - kv_cache_config = make_kv_cache_config(block_size=16) + kv_cache_config = make_kv_cache_config( + block_size=args[0].cache_config.block_size + ) super().__init__(*args, kv_cache_config=kv_cache_config, **kwargs) self._hand_shake_latency = hand_shake_latency self.kv_cache_layout = kv_cache_layout @@ -674,7 +678,8 @@ def test_pcp_producer_exposes_dcp_shards_or_canonical_replica( """Replicated PCP is canonicalized; PCP-DCP publishes every shard.""" from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend - vllm_config = create_vllm_config(kv_role="kv_producer") + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(kv_role="kv_producer", block_size=block_size) vllm_config.parallel_config.prefill_context_parallel_size = pcp_size vllm_config.parallel_config.decode_context_parallel_size = dcp_size with ( @@ -692,7 +697,7 @@ def test_pcp_producer_exposes_dcp_shards_or_canonical_replica( connector = NixlConnector( vllm_config, KVConnectorRole.WORKER, - make_kv_cache_config(block_size=16), + make_kv_cache_config(block_size=block_size), ) worker = connector.connector_worker @@ -746,12 +751,13 @@ def test_multi_xfer_one_engine( This test triggers the connector to load remote KV for the same `request_id`. """ - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) request_id = "req_id" # Test worker role in decode server. - kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=10) + kv_cache_config = make_kv_cache_config(block_size=block_size, num_blocks=10) connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, @@ -837,12 +843,15 @@ def test_async_load_kv( prefill_tp_size, ): """Test that NixlConnector's start_load_kv should be non-blocking.""" - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) vllm_config.parallel_config.tensor_parallel_size = decode_tp_size # Test worker role in decode server. connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id @@ -899,12 +908,13 @@ def test_prefill_tp_size_greater_than_decode_tp_size( lambda: local_tp_size, ) - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) # num_blocks=1 keeps the fake handshake's region_num_blocks consistent # with the minimal single-block registration below (num_descs == 1); # _fa_desc_replicated cross-checks the two. - kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=1) + kv_cache_config = make_kv_cache_config(block_size=block_size, num_blocks=1) connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, @@ -972,16 +982,21 @@ def test_prefill_tp_size_greater_than_decode_tp_size_mla( """Verify remote TP > local TP handshake succeeds with different remote configurations for an MLA model. """ - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) d_tp_size = 1 p_tp_size = 2 # Build two separate connectors/workers to emulate P TP=2 ranks. conn_p0 = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) conn_p1 = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) conn_p0.connector_worker = FakeNixlConnectorWorker( vllm_config, conn_p0.engine_id, hand_shake_latency=0 @@ -1075,11 +1090,14 @@ def test_concurrent_load_kv( dist_init, ): """Test that multiple start_load_kv calls should occur concurrently.""" - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) # Test worker role in decode server. connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id @@ -1143,8 +1161,10 @@ def test_handshake_fails_on_kv_cache_layout_mismatch( """Verify that adding a remote agent fails if kv_cache_layout differs. This test is only relevant for heterogeneous TP. """ + block_size = maybe_update_block_size(16) vllm_config = create_vllm_config( - enable_permute_local_kv=enable_permute_local_kv + enable_permute_local_kv=enable_permute_local_kv, + block_size=block_size, ) # Mock TP world size to 2 to force heterogeneous TP when @@ -1155,7 +1175,9 @@ def test_handshake_fails_on_kv_cache_layout_mismatch( ): # Initialize connector and worker (with fake NIXL wrapper) connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, @@ -1205,7 +1227,10 @@ def test_handshake_succeed_on_kv_cache_layout_mismatch_with_experimental( """Verify that adding a remote agent fails if kv_cache_layout differs. This test is only relevant for heterogeneous TP. """ - vllm_config = create_vllm_config(enable_permute_local_kv=True) + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config( + enable_permute_local_kv=True, block_size=block_size + ) # Mock TP world size to 2 to force heterogeneous TP when # remote_tp_size=1 @@ -1215,7 +1240,9 @@ def test_handshake_succeed_on_kv_cache_layout_mismatch_with_experimental( ): # Initialize connector and worker (with fake NIXL wrapper) connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, @@ -1261,7 +1288,9 @@ def test_hybrid_mamba_attention_remote_descs_use_packed_head_slices( self, default_vllm_config, dist_init ): worker = FakeNixlConnectorWorker( - create_vllm_config(), "engine", hand_shake_latency=0 + create_vllm_config(block_size=maybe_update_block_size(16)), + "engine", + hand_shake_latency=0, ) remote_block_len = 2048 @@ -1316,13 +1345,16 @@ def test_handshake_mixed_fa_mla_hetero_tp(self, default_vllm_config, dist_init): heterogeneous TP must NOT raise (previously a NotImplementedError), and the per-region gate must still reject a wrong block_len. """ - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) with patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501 return_value=2, ): connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -1409,7 +1441,8 @@ def test_handshake_validates_gqa_replicated_block_len( expected local_block_len * tp_ratio and rejected the valid handshake. """ - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) with patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501 @@ -1418,7 +1451,7 @@ def test_handshake_validates_gqa_replicated_block_len( connector = NixlConnector( vllm_config, KVConnectorRole.WORKER, - make_kv_cache_config(block_size=16), + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -1465,7 +1498,8 @@ def test_handshake_rejects_wrong_block_len_without_gqa_replication( block_lens when GQA replication is NOT in effect (32 KV heads, D_TP=4, P_TP=2: head_ratio=4, both sides have >1 head/rank). """ - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) with patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501 @@ -1474,7 +1508,7 @@ def test_handshake_rejects_wrong_block_len_without_gqa_replication( connector = NixlConnector( vllm_config, KVConnectorRole.WORKER, - make_kv_cache_config(block_size=16), + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -1522,11 +1556,14 @@ def test_handshake_rejects_wrong_block_len_without_gqa_replication( ) def test_kv_connector_stats(default_vllm_config, dist_init): """Test that KV transfer stats are properly recorded and retrieved.""" - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) # Test worker role in decode server. connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -1602,9 +1639,12 @@ def test_reqs_to_send_deadline_rebased_to_worker_clock(default_vllm_config, dist remote read pulls another request's data (silent accuracy corruption). The worker must anchor the remaining TTL to its own clock. """ - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -1793,9 +1833,12 @@ def counter_value(name: str) -> float: FakeNixlWrapper, ) def test_notification_after_expiry_is_counted(default_vllm_config, dist_init): - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -2323,7 +2366,10 @@ def test_register_kv_caches( 2. nixl_wrapper.get_xfer_descs() is called with blocks_data containing block layout info """ - vllm_config = create_vllm_config(attention_backend=attn_backend) + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config( + attention_backend=attn_backend, block_size=block_size + ) vllm_config.cache_config.kv_cache_layout = layout # Import the appropriate backend based on the parameter @@ -2348,7 +2394,6 @@ def test_register_kv_caches( patch(f"{nixl_worker}.get_current_attn_backends") as mock_get_attn_backends, ): mock_get_attn_backends.return_value = [backend_cls] - block_size = 16 num_blocks = 8 num_heads = 4 head_size = 16 @@ -2472,7 +2517,7 @@ def test_register_kv_caches( } assert len(blocks_data) == num_blocks * len(base_addrs) - assert connector.connector_worker.block_size == 16 + assert connector.connector_worker.block_size == block_size def test_register_packed_dsv4_mla_cache_as_single_region( @@ -2573,7 +2618,8 @@ def test_kv_buffer_to_nixl_memory_types( """Test that register_kv_caches() passes the correct memory types from the config to the nixl_wrapper. """ - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) # Override the default memory types in the config vllm_config.kv_transfer_config.kv_buffer_device = kv_buffer_device from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( @@ -2603,7 +2649,9 @@ def test_kv_buffer_to_nixl_memory_types( ): # noqa: E501 # Create connector and replace its worker with a fake one for isolation connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) # Verify get_reg_descs was called with the correct memory_type @@ -2617,17 +2665,18 @@ def test_kv_buffer_to_nixl_memory_types( ) def test_shutdown_cleans_up_resources(default_vllm_config, dist_init): """Test that shutdown() properly cleans up all resources.""" - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) scheduler = NixlConnectorScheduler( vllm_config, vllm_config.kv_transfer_config.engine_id, - make_kv_cache_config(block_size=16), + make_kv_cache_config(block_size=block_size), ) worker = NixlConnectorWorker( vllm_config, vllm_config.kv_transfer_config.engine_id, - make_kv_cache_config(block_size=16), + make_kv_cache_config(block_size=block_size), ) nixl_wrapper = worker.nixl_wrapper @@ -2643,7 +2692,7 @@ def test_shutdown_cleans_up_resources(default_vllm_config, dist_init): # Mock register_kv_cache which registers local handle worker.src_xfer_handles_by_block_size = {worker.block_size: 455} # P TP = 2 * D TP case, we should register 2 local handles - worker.src_xfer_handles_by_tp_ratio = {(-2, 16): [456, 457]} + worker.src_xfer_handles_by_tp_ratio = {(-2, block_size): [456, 457]} worker.dst_xfer_side_handles = {"engine1": {0: 789}} worker._remote_agents = {"engine1": {(0, 0): "agent1"}} # _cleanup_remote_engine (called by shutdown) also clears these: @@ -2688,13 +2737,15 @@ def _setup_worker_with_remote_engine( engine_ttl: float = 10.0, ) -> tuple[Any, str]: """Create a worker with one remote engine registered.""" + block_size = maybe_update_block_size(16) vllm_config = create_vllm_config( + block_size=block_size, kv_connector_extra_config={"engine_ttl": engine_ttl}, ) worker = NixlConnectorWorker( vllm_config, vllm_config.kv_transfer_config.engine_id, - make_kv_cache_config(block_size=16), + make_kv_cache_config(block_size=block_size), ) engine_id = "remote-engine-1" @@ -2857,7 +2908,11 @@ def setup(self, monkeypatch): from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend - config = create_vllm_config(kv_connector_extra_config={"engine_ttl": 0}) + block_size = maybe_update_block_size(16) + + config = create_vllm_config( + kv_connector_extra_config={"engine_ttl": 0}, block_size=block_size + ) config.kv_transfer_config.kv_buffer_device = "cpu" platform = SimpleNamespace( device_type="cpu", @@ -3258,12 +3313,15 @@ def test_aborted_request_removed_from_worker_in_batch(default_vllm_config, dist_ the real scheduler, then simulate an abort (request not in next scheduler iteration) and verify the worker no longer tracks it as in-batch. """ - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) scheduler = create_scheduler(vllm_config) # KVConnector Worker in P connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -3398,9 +3456,12 @@ def test_empty_recv_is_reported_only_when_awaited( A failed notification changes neither: the KV is local either way, and the producer frees its own blocks on a timeout. """ - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, @@ -3518,12 +3579,13 @@ def test_transfer_failure_logging( """ import logging - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) connector = NixlConnector( vllm_config, KVConnectorRole.WORKER, - make_kv_cache_config(block_size=16, swa_enabled=enable_hma), + make_kv_cache_config(block_size=block_size, swa_enabled=enable_hma), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, @@ -3648,10 +3710,13 @@ def emit(self, record): ) def test_handshake_failure_returns_finished(default_vllm_config, dist_init): """Test that handshake failures mark blocks invalid and return via get_finished.""" - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0.1 @@ -3704,10 +3769,13 @@ def test_transfer_setup_failure_returns_finished( default_vllm_config, dist_init, is_hma ): """Setup failures report the request; only non-HMA reports block IDs.""" - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -3776,7 +3844,9 @@ def release_xfer_handle(self, handle: int) -> None: def _make_split_read_connector(vllm_config, request_id, states): """Seed a request with scripted in-flight xfer handles + recv metadata.""" connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=vllm_config.cache_config.block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -3818,7 +3888,7 @@ def test_split_read_failure_defers_report_until_last_handle( request_id = "split_read_partial_failure" err_handle, live_handle = 11, 22 connector, worker, wrapper = _make_split_read_connector( - create_vllm_config(), + create_vllm_config(block_size=maybe_update_block_size(16)), request_id, {err_handle: ["ERR"], live_handle: ["PROC", "PROC", "DONE"]}, ) @@ -3894,10 +3964,15 @@ def test_failed_request_skips_kv_postprocessing( # Use enable_permute_local_kv=True so that # post_process_device_kv_on_receive would be called on the success path, # making the assertion meaningful (not trivially true). - vllm_config = create_vllm_config(enable_permute_local_kv=True) + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config( + enable_permute_local_kv=True, block_size=block_size + ) connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, @@ -4220,14 +4295,15 @@ def test_compatibility_hash_validation( enforce_handshake_compat: whether to enforce compatibility checking """ + block_size = maybe_update_block_size(16) local_vllm_config = create_vllm_config( model="facebook/opt-125m", - block_size=16, + block_size=block_size, kv_connector_extra_config={ "enforce_handshake_compat": enforce_handshake_compat }, ) - kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2) + kv_cache_config = make_kv_cache_config(block_size=block_size, num_blocks=2) decode_connector = NixlConnector( local_vllm_config, KVConnectorRole.WORKER, kv_cache_config ) @@ -4253,7 +4329,7 @@ def test_compatibility_hash_validation( remote_config_params: dict[str, Any] = { "model": "facebook/opt-125m", - "block_size": 16, + "block_size": block_size, **config_overrides, } remote_vllm_config = create_vllm_config(**remote_config_params) @@ -4276,7 +4352,7 @@ def test_compatibility_hash_validation( decode_worker.backend_name, ) - prefill_block_size = config_overrides.get("block_size", 16) + prefill_block_size = remote_vllm_config.cache_config.block_size prefill_block_lens = [4096 * prefill_block_size] prefill_metadata = NixlAgentMetadata( engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID, @@ -4352,12 +4428,15 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario) - NixlHandshakePayload decoder - NixlAgentMetadata decoder """ + block_size = maybe_update_block_size(16) local_vllm_config = create_vllm_config( model="facebook/opt-125m", - block_size=16, + block_size=block_size, ) decode_connector = NixlConnector( - local_vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + local_vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) decode_worker = decode_connector.connector_worker @@ -4441,11 +4520,14 @@ def test_mla_broadcast_notif_uses_remote_request_id( decode_tp_size = 1 prefill_tp_size = 4 - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) vllm_config.parallel_config.tensor_parallel_size = decode_tp_size connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -4550,7 +4632,8 @@ def test_kv_both_deprecation_warning(default_vllm_config, dist_init): _print_warning_once.cache_clear() - vllm_config = create_vllm_config(kv_role="kv_both") + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(kv_role="kv_both", block_size=block_size) with patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector.logger" @@ -4559,7 +4642,7 @@ def test_kv_both_deprecation_warning(default_vllm_config, dist_init): NixlConnector( vllm_config, KVConnectorRole.WORKER, - make_kv_cache_config(block_size=16), + make_kv_cache_config(block_size=block_size), ) mock_logger.warning_once.assert_called_once() @@ -4575,14 +4658,15 @@ def test_kv_both_deprecation_warning(default_vllm_config, dist_init): def test_explicit_kv_role_no_deprecation_warning(default_vllm_config, dist_init): """kv_role='kv_consumer' or 'kv_producer' should NOT emit a warning.""" for role in ("kv_consumer", "kv_producer"): - vllm_config = create_vllm_config(kv_role=role) + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(kv_role=role, block_size=block_size) with patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector.logger" ) as mock_logger: NixlConnector( vllm_config, KVConnectorRole.WORKER, - make_kv_cache_config(block_size=16), + make_kv_cache_config(block_size=block_size), ) ( diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index c665ce0c5b88..469e0f5d8294 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -39,6 +39,7 @@ create_vllm_config, make_kv_cache_config, make_nixl_scheduler, + maybe_update_block_size, ) @@ -1027,7 +1028,7 @@ def test_fewer_blocks_with_hma(monkeypatch, model_name, sw_size): kv_connector="NixlConnector", kv_role="kv_consumer", ) - block_size = 16 + block_size = maybe_update_block_size(16) llm_kwargs = { "model": model_name, "enforce_eager": True, diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 28bfda419a76..59b2e7ad0349 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -31,6 +31,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import ( # noqa ExampleConnector, ) +from vllm.platforms import current_platform from vllm.utils.hashing import sha256 from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash @@ -59,6 +60,10 @@ def create_mock_connector() -> Mock: return connector +def maybe_update_block_size(block_size: int) -> int: + return 64 if current_platform.is_xpu() else block_size + + def assert_scheduler_empty(scheduler: Scheduler): """Confirm the scheduler is "empty" - i.e. no leaks.""" # Scheduler Metadata. diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 6179b27bb6db..4a336586acd8 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -355,6 +355,8 @@ def _get_fa4_hd256_block_size( def get_supported_kernel_block_sizes( cls, kv_cache_spec: KVCacheSpec | None = None ) -> list[int | MultipleOf]: + if current_platform.is_xpu(): + return [MultipleOf(64)] if block_size := cls._get_sm90_fa4_fp8_kv_block_size(kv_cache_spec): return [block_size] if block_size := cls._get_fa4_hd256_block_size(kv_cache_spec): diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index d1feaef031d9..8964c660a1d9 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -14,9 +14,11 @@ get_layers_from_vllm_config, ) from vllm.config.compilation import CUDAGraphMode +from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.multimodal.inputs import MultiModalFeatureSpec +from vllm.platforms import current_platform from vllm.v1.attention.backend import ( AttentionCGSupport, CommonAttentionMetadata, @@ -45,6 +47,8 @@ prepare_kernel_block_sizes, ) +logger = init_logger(__name__) + if TYPE_CHECKING: from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.cudagraph_utils import ( @@ -263,10 +267,24 @@ def init_attn_backend( # by all backends within that group. kernel_block_sizes = prepare_kernel_block_sizes(kv_cache_config, attn_groups) + if current_platform.is_xpu(): + for kv_cache_group_id, selected_kernel_block_size in enumerate( + kernel_block_sizes + ): + if selected_kernel_block_size % 64 != 0: + for attn_group in attn_groups[kv_cache_group_id]: + logger.warning_once( + "On XPU platform, current selected kernel block size is " + "not fully optimized: kernel_block_size=%s, backend=%s, " + "kv_cache_spec=%s", + selected_kernel_block_size, + attn_group.backend, + repr(attn_group.kv_cache_spec), + ) # Phase 3: create metadata builders and determine cudagraph support. attn_backend_workspace: torch.Tensor | None = None for kv_cache_group_id, groups in enumerate(attn_groups): - kernel_block_size = None + kernel_block_size: int | None = None if kv_cache_group_id < len(kernel_block_sizes): kernel_block_size = kernel_block_sizes[kv_cache_group_id] for group in groups: From bfa7e5b7afe40ebe48288801a2baf5c1135cd881 Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Mon, 12 Oct 2026 10:39:40 +0800 Subject: [PATCH 272/624] [XPU][CI] Enable XPU graph in graph tests (#60784) Signed-off-by: zhenwei-intel Co-authored-by: Kunshang Ji --- .../v1/cudagraph/test_breakable_cudagraph.py | 46 ++++++++++++------- tests/v1/cudagraph/test_cudagraph_dispatch.py | 15 ++++-- tests/v1/cudagraph/test_cudagraph_mode.py | 7 +++ tests/v1/cudagraph/test_encoder_cudagraph.py | 27 ++++++----- 4 files changed, 63 insertions(+), 32 deletions(-) diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index b0b49a9567b1..1b9a79a6ef2c 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/tests/v1/cudagraph/test_breakable_cudagraph.py @@ -11,6 +11,18 @@ import pytest import torch +from vllm.platforms import current_platform + +DEVICE_TYPE = current_platform.device_type +HAS_DEVICE_GRAPH = current_platform.is_cuda_alike() or current_platform.is_xpu() + +if current_platform.is_xpu(): + from vllm.v1.worker.xpu_model_runner import _torch_cuda_wrapper + + # Route torch.cuda graph/stream APIs to torch.xpu, as XPUModelRunner does. + with _torch_cuda_wrapper(): + pass + @pytest.fixture(autouse=True) def _enable_breakable_cudagraph(monkeypatch: pytest.MonkeyPatch): @@ -96,8 +108,8 @@ def cuda_capture_stream(): stream, so all capture-using tests need to run under ``torch.cuda.stream(...)`` for a separate stream. """ - if not torch.cuda.is_available(): - pytest.skip("CUDA required") + if not HAS_DEVICE_GRAPH: + pytest.skip("CUDA or XPU graph required") from vllm.utils.torch_utils import _current_stream_tls prev_stream = getattr(_current_stream_tls, "value", None) @@ -200,7 +212,7 @@ def worker(): def test_capture_with_no_eager_break_records_one_graph(cuda_capture_stream): from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture - x = torch.zeros(4, device="cuda") + x = torch.zeros(4, device=DEVICE_TYPE) cap = BreakableCUDAGraphCapture() with cap: x.add_(1.0) @@ -212,7 +224,7 @@ def test_capture_with_no_eager_break_records_one_graph(cuda_capture_stream): def test_add_eager_creates_alternating_graph_eager_graph(cuda_capture_stream): from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture - x = torch.zeros(4, device="cuda") + x = torch.zeros(4, device=DEVICE_TYPE) counter = {"eager_calls": 0} def eager_step(): @@ -253,7 +265,7 @@ def test_capture_replay_matches_eager_simple(cuda_capture_stream): """ from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture - x = torch.zeros(8, device="cuda") + x = torch.zeros(8, device=DEVICE_TYPE) log: list[str] = [] def eager_break_op(): @@ -269,21 +281,21 @@ def eager_break_op(): # Capture-time: graph kernels were recorded only; eager segment ran # once on x == 0, leaving x == 0. torch.accelerator.synchronize() - assert torch.equal(x, torch.zeros(8, device="cuda")) + assert torch.equal(x, torch.zeros(8, device=DEVICE_TYPE)) assert log == ["eager"] # Replay with a fresh input: 10 -> 11 -> 22 -> 27. x.fill_(10.0) cap.replay() torch.accelerator.synchronize() - assert torch.equal(x, torch.full((8,), 27.0, device="cuda")) + assert torch.equal(x, torch.full((8,), 27.0, device=DEVICE_TYPE)) assert log == ["eager", "eager"] # Replay again with another input: 100 -> 101 -> 202 -> 207. x.fill_(100.0) cap.replay() torch.accelerator.synchronize() - assert torch.equal(x, torch.full((8,), 207.0, device="cuda")) + assert torch.equal(x, torch.full((8,), 207.0, device=DEVICE_TYPE)) assert log == ["eager", "eager", "eager"] @@ -301,13 +313,13 @@ def attention_like(t: torch.Tensor) -> None: # In-place double; stands in for "real" attention work. t.mul_(2.0) - x = torch.zeros(4, device="cuda") + x = torch.zeros(4, device=DEVICE_TYPE) # Outside capture: decorator should just call through. x.fill_(3.0) attention_like(x) torch.accelerator.synchronize() - assert torch.equal(x, torch.full((4,), 6.0, device="cuda")) + assert torch.equal(x, torch.full((4,), 6.0, device=DEVICE_TYPE)) # Inside capture: decorator should split the graph. Only the eager # segment actually mutates state during capture. @@ -318,7 +330,7 @@ def attention_like(t: torch.Tensor) -> None: attention_like(x) # eager: x *= 2 (on x == 0, no-op) x.add_(1.0) # recorded torch.accelerator.synchronize() - assert torch.equal(x, torch.zeros(4, device="cuda")) + assert torch.equal(x, torch.zeros(4, device=DEVICE_TYPE)) # 2 graph segments + 1 eager segment, ordered G E G; the arithmetic # equivalence check below verifies the ordering. assert len(cap.segments) == 3 @@ -329,7 +341,7 @@ def attention_like(t: torch.Tensor) -> None: x.fill_(2.0) cap.replay() torch.accelerator.synchronize() - assert torch.equal(x, torch.full((4,), 15.0, device="cuda")) + assert torch.equal(x, torch.full((4,), 15.0, device=DEVICE_TYPE)) @pytest.mark.parametrize("keyword", [False, True]) @@ -352,7 +364,7 @@ def consume(x, data_out, scale_out): data_out.copy_(x.data) scale_out.copy_(x.scale) - source = torch.ones((128, 128), device="cuda", dtype=torch.bfloat16) + source = torch.ones((128, 128), device=DEVICE_TYPE, dtype=torch.bfloat16) data_out = source.to(torch.float8_e4m3fn) scale_out = source[:, ::32].to(torch.uint8) output = torch.empty_like(source) @@ -391,7 +403,7 @@ def test_eager_attention_inside_multistream_overlap(cuda_capture_stream): ) from vllm.utils.multi_stream_utils import maybe_execute_in_parallel - x = torch.zeros(1024, device="cuda") + x = torch.zeros(1024, device=DEVICE_TYPE) output = torch.empty_like(x) aux_stream = torch.cuda.Stream() main_event = torch.cuda.Event() @@ -440,7 +452,7 @@ def test_replay_invokes_eager_segments_in_order(cuda_capture_stream): from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture log: list[str] = [] - x = torch.zeros(1, device="cuda") + x = torch.zeros(1, device=DEVICE_TYPE) def make_eager(name): def step(): @@ -502,7 +514,7 @@ def test_nested_decorated_op_runs_inline(cuda_capture_stream): eager_break_during_capture, ) - x = torch.zeros(4, device="cuda") + x = torch.zeros(4, device=DEVICE_TYPE) inner_calls = 0 @eager_break_during_capture @@ -533,7 +545,7 @@ def outer_op(t: torch.Tensor) -> None: cap.replay() torch.accelerator.synchronize() # 0 -> +2 -> +1 (inner) -> +10 (outer) -> +100 = 113 - assert torch.equal(x, torch.full((4,), 113.0, device="cuda")) + assert torch.equal(x, torch.full((4,), 113.0, device=DEVICE_TYPE)) assert inner_calls == 2 # replay invokes the outer's lambda again diff --git a/tests/v1/cudagraph/test_cudagraph_dispatch.py b/tests/v1/cudagraph/test_cudagraph_dispatch.py index 40874be5bd59..b0a4587ddfc5 100644 --- a/tests/v1/cudagraph/test_cudagraph_dispatch.py +++ b/tests/v1/cudagraph/test_cudagraph_dispatch.py @@ -24,6 +24,14 @@ from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher DEVICE_TYPE = current_platform.device_type +HAS_DEVICE_GRAPH = current_platform.is_cuda() or current_platform.is_xpu() + +if current_platform.is_xpu(): + from vllm.v1.worker.xpu_model_runner import _torch_cuda_wrapper + + # Route torch.cuda graph/stream APIs to torch.xpu, as XPUModelRunner does. + with _torch_cuda_wrapper(): + pass # Helper MLP for testing @@ -50,6 +58,7 @@ def _create_vllm_config( mock_config.parallel_config = ParallelConfig() mock_config.speculative_config = None # No speculative decoding mock_config.num_speculative_tokens = 0 + mock_config.use_cumem_cudagraph_pool = False if not lora_config: mock_config.lora_config = None else: @@ -268,7 +277,7 @@ def test_get_capture_descs_empty_when_not_initialized(self): assert dispatcher.get_capture_descs() == [] -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") +@pytest.mark.skipif(not HAS_DEVICE_GRAPH, reason="Requires CUDA or XPU graph") class TestCUDAGraphWrapper: def setup_method(self): self.vllm_config = _create_vllm_config(CompilationConfig()) @@ -415,7 +424,7 @@ def _run_and_monitor_call( @create_new_process_for_each_test("spawn") -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") +@pytest.mark.skipif(not HAS_DEVICE_GRAPH, reason="Requires CUDA or XPU graph") def test_capture_replay_bypass_logic(): comp_config = CompilationConfig( mode=CompilationMode.VLLM_COMPILE, @@ -484,7 +493,7 @@ def test_capture_replay_bypass_logic(): @create_new_process_for_each_test("spawn") -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") +@pytest.mark.skipif(not HAS_DEVICE_GRAPH, reason="Requires CUDA or XPU graph") def test_nested_wrappers(): """Tests a scenario with a PIECEWISE wrapper inside a FULL one.""" comp_config = CompilationConfig( diff --git a/tests/v1/cudagraph/test_cudagraph_mode.py b/tests/v1/cudagraph/test_cudagraph_mode.py index 0147f90df534..f35ef4d82974 100644 --- a/tests/v1/cudagraph/test_cudagraph_mode.py +++ b/tests/v1/cudagraph/test_cudagraph_mode.py @@ -19,6 +19,13 @@ ("TritonAttn", "FULL", True), ("TritonAttn", "FULL_AND_PIECEWISE", True), ] +elif current_platform.is_xpu(): + combo_cases_1 = [ + ("FA2", "FULL", True), + ("FA2", "FULL_AND_PIECEWISE", True), + ("TritonAttn", "FULL", True), + ("TritonAttn", "FULL_AND_PIECEWISE", True), + ] else: combo_cases_1 = [ ("FA3", "FULL", True), diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index 348944387886..2799a3ef4dd6 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/tests/v1/cudagraph/test_encoder_cudagraph.py @@ -32,6 +32,15 @@ EncoderItemSpec, ) +HAS_DEVICE_GRAPH = current_platform.is_cuda_alike() or current_platform.is_xpu() + +if current_platform.is_xpu(): + from vllm.v1.worker.xpu_model_runner import _torch_cuda_wrapper + + # Route torch.cuda graph/stream APIs to torch.xpu, as XPUModelRunner does. + with _torch_cuda_wrapper(): + pass + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -662,12 +671,10 @@ def _make_video_mm_kwargs( # --------------------------------------------------------------------------- -@pytest.mark.skipif( - not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" -) +@pytest.mark.skipif(not HAS_DEVICE_GRAPH, reason="Requires CUDA, ROCm or XPU graph") class TestEncoderCudaGraphCaptureReplay: def setup_method(self): - self.device = torch.device("cuda:0") + self.device = torch.device(current_platform.device_type, 0) self.dtype = torch.float16 self.model = SimpleMockViTModel().to(self.device).half() self.mgr = _make_manager_for_gpu( @@ -758,9 +765,7 @@ def test_chunking_when_images_exceed_max_batch(self): # --------------------------------------------------------------------------- -@pytest.mark.skipif( - not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" -) +@pytest.mark.skipif(not HAS_DEVICE_GRAPH, reason="Requires CUDA, ROCm or XPU graph") @pytest.mark.usefixtures("dist_init", "workspace_init") @pytest.mark.parametrize("profile_only", [False, True]) @torch.inference_mode() @@ -776,7 +781,7 @@ def test_eonly_capture_preserves_outputs_across_replay_and_fallback(profile_only from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner from vllm.v1.worker.workspace import lock_workspace - device = torch.device("cuda:0") + device = torch.device(current_platform.device_type, 0) dtype = torch.float16 model = SimpleMockViTModel().to(device).half() manager = _make_manager_for_gpu(model, _BUDGETS, _MAX_BATCH, device, dtype) @@ -1047,12 +1052,10 @@ def test_video_model_returns_video_for_video_kwargs(self): _VIDEO_MAX_FRAMES = 8 # 2 frames per item at max_batch_size=4 -@pytest.mark.skipif( - not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" -) +@pytest.mark.skipif(not HAS_DEVICE_GRAPH, reason="Requires CUDA, ROCm or XPU graph") class TestEncoderCudaGraphVideoReplay: def setup_method(self): - self.device = torch.device("cuda:0") + self.device = torch.device(current_platform.device_type, 0) self.dtype = torch.float16 self.model = SimpleMockViTVideoModel().to(self.device).half() self.mgr = _make_manager_for_gpu( From d4e92ccc37e67048a957aa8e08dc5a2bc665f422 Mon Sep 17 00:00:00 2001 From: shaohuaxi Date: Mon, 12 Oct 2026 11:58:15 +0800 Subject: [PATCH 273/624] [Bugfix][Frontend] Honor endpoint plugin route overrides (#58839) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 子华 Co-authored-by: Codex Co-authored-by: Cyrus Leung --- docs/design/endpoint_plugins.md | 4 +- tests/plugins_tests/test_endpoint_plugins.py | 160 ++++++++++++++++++- vllm/entrypoints/launchers/app.py | 4 +- vllm/plugins/endpoint_plugins/interface.py | 51 +++++- 4 files changed, 210 insertions(+), 9 deletions(-) diff --git a/docs/design/endpoint_plugins.md b/docs/design/endpoint_plugins.md index 98b755788e2e..361bd7099b46 100644 --- a/docs/design/endpoint_plugins.md +++ b/docs/design/endpoint_plugins.md @@ -126,7 +126,9 @@ Do not expect a single endpoint plugin to also mutate engine/worker state. If yo ## Path-prefix convention -There is currently no route conflict enforcement (tracked as a follow-up to RFC [#46565](https://github.com/vllm-project/vllm/issues/46565)). A plugin's `attach_router` can register a path that collides with a core route and routes attached later win. To avoid surprising operators: +Plugin hooks run after the core routers have been registered. HTTP routes added by plugins replace earlier routes with the same path template and HTTP method, including their OpenAPI description. Methods not overridden by the plugin remain available. When multiple plugins register the same operation, the last registered route takes precedence. Different path templates are not treated as equivalent. + +There is currently no rejection of route conflicts (tracked as a follow-up to RFC [#46565](https://github.com/vllm-project/vllm/issues/46565)). To avoid surprising operators: - Namespace your routes under a distinct prefix, e.g. `/plugins//...`, rather than reusing `/v1/...` or other core prefixes - Only register routes under a core prefix (like the worked example's `/v1/admin/scheduler_config`) if you specifically intend to override or extend existing behavior and document that clearly for operators allowlisting your plugin diff --git a/tests/plugins_tests/test_endpoint_plugins.py b/tests/plugins_tests/test_endpoint_plugins.py index 2d921539004f..864c915fe083 100644 --- a/tests/plugins_tests/test_endpoint_plugins.py +++ b/tests/plugins_tests/test_endpoint_plugins.py @@ -14,7 +14,7 @@ import httpx import pytest -from fastapi import FastAPI +from fastapi import APIRouter, FastAPI from vllm_add_dummy_endpoint_plugin import DummyAdminEndpointPlugin from vllm.entrypoints.launchers.app import build_app @@ -38,6 +38,34 @@ def __init__(self): raise RuntimeError("boom") +class _ShadowEndpointPlugin: + name = "shadow_endpoint_plugin" + required_tasks = None + + def __init__(self, path: str, include_in_schema: bool): + self.path = path + self.include_in_schema = include_in_schema + + def attach_router(self, app: FastAPI) -> None: + assert str(app.url_path_for("show_version")) == "/version" + app.openapi() + + @app.get(self.path, status_code=209, include_in_schema=self.include_in_schema) + async def plugin_version(plugin_arg: str) -> dict[str, str]: + return {"served_by": plugin_arg} + + async def init_state(self, engine_client, state, args) -> None: + pass + + +class _RouterEndpointPlugin(DummyAdminEndpointPlugin): + def __init__(self, router: APIRouter): + self.router = router + + def attach_router(self, app: FastAPI) -> None: + app.include_router(self.router) + + class _FakeEngineClient: """Minimal stand in exercising `collective_rpc`. Not a real engine.""" @@ -162,6 +190,136 @@ def test_attach_is_noop_when_nothing_discovered(monkeypatch: pytest.MonkeyPatch) assert app.state.endpoint_plugins == [] +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/version", "/metrics"]) +@pytest.mark.parametrize("include_in_schema", [True, False]) +async def test_endpoint_plugin_override_preserves_hook_and_schema( + monkeypatch: pytest.MonkeyPatch, path: str, include_in_schema: bool +): + """Hooks see core routes; overrides win over both HTTP routes and mounts.""" + args = _build_args() + monkeypatch.setenv("VLLM_PLUGINS", "shadow_endpoint_plugin") + monkeypatch.setattr( + "vllm.plugins.load_plugins_by_group", + _fake_loader( + { + "shadow_endpoint_plugin": lambda: _ShadowEndpointPlugin( + path, include_in_schema + ) + } + ), + ) + app = build_app(args, supported_tasks=()) + + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + response = await client.get(path, params={"plugin_arg": "plugin"}) + schema = (await client.get("/openapi.json")).json() + + assert response.status_code == 209 + assert response.json() == {"served_by": "plugin"} + if include_in_schema: + operation = schema["paths"][path]["get"] + assert "209" in operation["responses"] + assert "200" not in operation["responses"] + assert operation["parameters"][0]["name"] == "plugin_arg" + assert operation["parameters"][0]["required"] is True + response_schema = operation["responses"]["209"]["content"]["application/json"][ + "schema" + ] + assert response_schema["additionalProperties"] == {"type": "string"} + else: + assert path not in schema["paths"] + + +@pytest.mark.asyncio +async def test_endpoint_plugin_override_preserves_other_methods( + monkeypatch: pytest.MonkeyPatch, +): + """Overriding GET preserves POST's validation and the shared core route.""" + core = APIRouter() + + @core.api_route("/shared", methods=["GET", "POST"], status_code=201) + async def original(value: int) -> dict[str, int]: + return {"value": value} + + plugin = APIRouter() + + @plugin.get("/shared", status_code=202) + async def replacement(): + return {"served_by": "plugin"} + + app = FastAPI() + app.include_router(core) + original_route = app.routes[-1] + monkeypatch.setattr( + "vllm.plugins.load_endpoint_plugins", + lambda supported_tasks: [_RouterEndpointPlugin(plugin)], + ) + attach_endpoint_plugins(app, ()) + + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + response = await client.get("/shared") + post_response = await client.post("/shared", params={"value": 7}) + invalid_response = await client.post("/shared", params={"value": "invalid"}) + + assert response.status_code == 202 + assert response.json() == {"served_by": "plugin"} + assert post_response.status_code == 201 + assert post_response.json() == {"value": 7} + assert invalid_response.status_code == 422 + assert original_route.methods == {"GET", "POST"} + assert core.routes[0].methods == {"GET", "POST"} + operations = app.openapi()["paths"]["/shared"] + assert "202" in operations["get"]["responses"] + assert "201" in operations["post"]["responses"] + + +@pytest.mark.asyncio +async def test_later_endpoint_plugin_overrides_same_operation( + monkeypatch: pytest.MonkeyPatch, +): + """The last registered override wins while unrelated plugin routes survive.""" + first = APIRouter() + + @first.get("/shared", status_code=201) + @first.get("/plugins/first") + async def first_handler(): + return {"served_by": "first"} + + second = APIRouter() + + @second.get("/shared", status_code=202) + async def second_handler(): + return {"served_by": "second"} + + monkeypatch.setattr( + "vllm.plugins.load_endpoint_plugins", + lambda supported_tasks: [ + _RouterEndpointPlugin(first), + _RouterEndpointPlugin(second), + ], + ) + app = FastAPI() + attach_endpoint_plugins(app, ()) + + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + response = await client.get("/shared") + unrelated = await client.get("/plugins/first") + unsupported = await client.delete("/shared") + + assert response.status_code == 202 + assert response.json() == {"served_by": "second"} + assert unrelated.status_code == 200 + assert unrelated.json() == {"served_by": "first"} + assert unsupported.status_code == 405 + responses = app.openapi()["paths"]["/shared"]["get"]["responses"] + assert "202" in responses + assert "201" not in responses + + @pytest.mark.asyncio async def test_init_state_is_noop_without_phase_a(monkeypatch: pytest.MonkeyPatch): """`init_app_state` callers that never ran `build_app` (e.g. diff --git a/vllm/entrypoints/launchers/app.py b/vllm/entrypoints/launchers/app.py index ddea70781051..59c4fc37a07a 100644 --- a/vllm/entrypoints/launchers/app.py +++ b/vllm/entrypoints/launchers/app.py @@ -44,8 +44,8 @@ def build_app( register_api_routers(args, app, supported_tasks, model_config) - # Endpoint plugins are attached last so their routes are registered after all core - # routers. This runs even for the CPU only render server. A plugin eligible for + # Endpoint plugin hooks run after core routers so plugins can inspect them. + # This runs even for the CPU only render server. A plugin eligible for # the `render` task still gets its routes registered. It receives # `engine_client=None` at Phase B (see `_init_endpoint_plugins_state`). attach_endpoint_plugins(app, supported_tasks) diff --git a/vllm/plugins/endpoint_plugins/interface.py b/vllm/plugins/endpoint_plugins/interface.py index 00451217d2f2..0d847895cfa2 100644 --- a/vllm/plugins/endpoint_plugins/interface.py +++ b/vllm/plugins/endpoint_plugins/interface.py @@ -29,16 +29,22 @@ """ from argparse import Namespace +from copy import copy from typing import TYPE_CHECKING, Protocol, runtime_checkable from fastapi import FastAPI +from fastapi.routing import APIRoute from starlette.datastructures import State +from starlette.routing import BaseRoute from vllm.engine.protocol import EngineClient +from vllm.logger import init_logger if TYPE_CHECKING: from vllm.tasks import SupportedTask +logger = init_logger(__name__) + @runtime_checkable class EndpointPlugin(Protocol): @@ -64,8 +70,8 @@ def attach_router(self, app: FastAPI) -> None: """Register this plugin's routes on `app`. Called once during `build_app()` after all core routers have been - attached. Routes attached here can shadow core routes with the same - path. There is currently no conflict enforcement (see RFC #46565 follow ups). + attached. HTTP routes registered here replace earlier routes with the + same path template and method. Other methods remain available. """ ... @@ -93,15 +99,50 @@ def attach_endpoint_plugins( ) -> None: """Phase A of endpoint plugin wiring: discover, gate and attach routes. - Attached last after all core routers. This is so endpoint plugin routes can - shadow core routes with the same path (see `EndpointPlugin.attach_router` - docstring). No-ops when no plugins are discovered/allowlisted. + Hooks run after all core routers have been attached. Plugin HTTP routes + replace earlier operations with the same path template and method in both + routing and OpenAPI. No-ops when no plugins are discovered/allowlisted. """ from vllm.plugins import load_endpoint_plugins endpoint_plugins = load_endpoint_plugins(supported_tasks) + existing_routes = {id(route): route for route in app.routes} for plugin in endpoint_plugins: plugin.attach_router(app) + + routes: list[BaseRoute] = [] + for route in app.routes: + if id(route) in existing_routes or not isinstance(route, APIRoute): + routes.append(route) + continue + + insert_at = len(routes) + retained: list[BaseRoute] = [] + for previous in routes: + if ( + isinstance(previous, APIRoute) + and previous.path == route.path + and previous.methods & route.methods + ): + logger.info( + "Endpoint plugin overrides %s %s", + ", ".join(sorted(previous.methods & route.methods)), + route.path, + ) + # Keep replacements ahead of mounts, preserving other methods. + insert_at = min(insert_at, len(retained)) + methods = previous.methods - route.methods + if not methods: + continue + previous = copy(previous) + previous.methods = methods + retained.append(previous) + retained.insert(insert_at, route) + routes = retained + + app.router.routes[:] = routes + if endpoint_plugins: + app.openapi_schema = None app.state.endpoint_plugins = endpoint_plugins From fe4df3df0eb703862c686b256e9fe44fa94c8945 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Mon, 12 Oct 2026 13:41:52 +0800 Subject: [PATCH 274/624] [Bugfix][GLM-5.3-Flash] Restore MTP index sharing across draft iterations (#61079) Signed-off-by: Canlin Co-authored-by: Claude Opus 5.5 --- .../test_speculative_draft_hf_overrides.py | 20 ++++++++++++++++++- vllm/config/speculative.py | 11 ++++++++-- 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py index e7014faa47d0..39ebd7924bf2 100644 --- a/tests/config/test_speculative_draft_hf_overrides.py +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -15,7 +15,7 @@ from unittest.mock import MagicMock, patch import pytest -from transformers import PreTrainedConfig +from transformers import Glm5NextConfig, PreTrainedConfig from vllm.config.parallel import ParallelConfig from vllm.config.speculative import SpeculativeConfig @@ -270,3 +270,21 @@ def test_mtp_index_share_override( speculative_config.draft_model_config.hf_config.index_share_for_mtp_iteration is expected ) + + +@pytest.mark.cpu_test +def test_glm5_next_mtp_override_promotes_index_share(): + """The flag lives on text_config, but the MTP proposers read the draft config.""" + config = Glm5NextConfig( + architectures=["Glm5NextForConditionalGeneration"], + text_config={ + "num_nextn_predict_layers": 1, + "index_share_for_mtp_iteration": True, + }, + ) + + draft_config = SpeculativeConfig.hf_config_override(config) + + assert draft_config.architectures == ["Glm5NextMTPModel"] + assert draft_config.n_predict == 1 + assert draft_config.index_share_for_mtp_iteration is True diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 4c61ead82dab..b916d7da2bed 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1058,9 +1058,16 @@ def hf_config_override(hf_config: PreTrainedConfig) -> PreTrainedConfig: ) if hf_config.model_type == "glm5_next": hf_config.model_type = "glm5_next_mtp" - n_predict = hf_config.get_text_config().num_nextn_predict_layers + text_config = hf_config.get_text_config() + share_mtp_indices = getattr( + text_config, "index_share_for_mtp_iteration", False + ) hf_config.update( - {"n_predict": n_predict, "architectures": ["Glm5NextMTPModel"]} + { + "n_predict": text_config.num_nextn_predict_layers, + "architectures": ["Glm5NextMTPModel"], + "index_share_for_mtp_iteration": share_mtp_indices, + } ) return hf_config From 74f4404f58c0fb16821b34c16343486eeaf164a5 Mon Sep 17 00:00:00 2001 From: Hubert <163992334+hubert-marek@users.noreply.github.com> Date: Mon, 12 Oct 2026 07:44:12 +0200 Subject: [PATCH 275/624] [Bugfix][Distributed] Own SHM out-of-band buffers before releasing read slots (#61000) Signed-off-by: Hubert Marek <163992334+hubert-marek@users.noreply.github.com> Signed-off-by: Ao Shen Co-authored-by: Codex Co-authored-by: Ao Shen --- tests/distributed/test_shm_broadcast.py | 105 +++++++++++++----- .../device_communicators/shm_broadcast.py | 15 ++- 2 files changed, 87 insertions(+), 33 deletions(-) diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py index 4984a2480a61..0e6efd10429b 100644 --- a/tests/distributed/test_shm_broadcast.py +++ b/tests/distributed/test_shm_broadcast.py @@ -57,6 +57,7 @@ def distributed_run(fn, world_size, timeout=60, mp_context=None): ctx = mp_context or mp number_of_processes = world_size processes = [] + master_port = str(get_open_port()) for i in range(number_of_processes): env = {} env["RANK"] = str(i) @@ -64,7 +65,7 @@ def distributed_run(fn, world_size, timeout=60, mp_context=None): env["WORLD_SIZE"] = str(number_of_processes) env["LOCAL_WORLD_SIZE"] = str(number_of_processes) env["MASTER_ADDR"] = "localhost" - env["MASTER_PORT"] = "12345" + env["MASTER_PORT"] = master_port p = ctx.Process(target=fn, args=(env,)) processes.append(p) p.start() @@ -88,11 +89,14 @@ def distributed_run(fn, world_size, timeout=60, mp_context=None): break time.sleep(0.1) # Check every 100ms - # Check for timeout if no failures detected yet + # Reap every child; failures can occur after the last polling iteration. + failed_processes.clear() for i, p in enumerate(processes): if p.is_alive(): p.kill() - p.join() + p.join() + if p.exitcode != 0: + failed_processes.append((i, p.exitcode)) # Report failures if failed_processes: @@ -102,6 +106,11 @@ def distributed_run(fn, world_size, timeout=60, mp_context=None): raise AssertionError(error_msg) +def test_distributed_run_reports_timeout(): + with pytest.raises(AssertionError, match="Distributed test failed"): + distributed_run(lambda _: time.sleep(60), world_size=1, timeout=0) + + def worker_fn_wrapper(fn): # `mp.Process` cannot accept environment variables directly # so we need to pass the environment variables as arguments @@ -381,14 +390,10 @@ def worker_fn_tensor_broadcast(): # Both ranks construct the identical reference payload. torch.manual_seed(42) payload = { - # 2MiB: rides the shm ring as an out-of-band buffer (the receiving - # side must copy out of the reusable ring chunk). "mid": torch.randn(1024, 512), - # 16MiB > max_chunk_bytes: overflows to the zmq socket (the - # receiving side aliases the zmq.Frame zero-copy). - "big": torch.randn(4096, 2048, dtype=torch.bfloat16), "nested": ["plain", 123, {"inner": torch.arange(5)}], } + big = torch.randn(4096, 2048, dtype=torch.bfloat16) if rank == writer_rank: with mock.patch( @@ -396,28 +401,23 @@ def worker_fn_tensor_broadcast(): wraps=_reduce_tensor, ) as wrapped_reduce: message_queue.enqueue(payload) + # Overflow is per message, not per tensor: send this separately. + message_queue.enqueue(big) assert wrapped_reduce.call_count == 3 - # Cycle the ring (max_chunks=4) several times over so that aliased - # ring chunks would be overwritten. - for i in range(16): - message_queue.enqueue({"junk": torch.full((1024, 512), float(i))}) else: - received = message_queue.dequeue(timeout=30) - for key in ("mid", "big"): - assert torch.equal(received[key], payload[key]), key - assert received[key].dtype == payload[key].dtype, key + with mock.patch.object(MessageQueue, "recv", wraps=MessageQueue.recv) as recv: + received = message_queue.dequeue(timeout=30) + recv.assert_not_called() + received_big = message_queue.dequeue(timeout=30) + recv.assert_called_once() + assert torch.equal(received["mid"], payload["mid"]) + assert received["mid"].dtype == payload["mid"].dtype + assert torch.equal(received_big, big) + assert received_big.dtype == big.dtype assert torch.equal(received["nested"][2]["inner"], torch.arange(5)) - - snapshot = received["mid"].clone() - for i in range(16): - junk = message_queue.dequeue(timeout=30) - assert torch.equal(junk["junk"], torch.full((1024, 512), float(i))) - # Tensors received via the shm ring must not alias chunk memory - # that the writer has reused for subsequent messages. - assert torch.equal(received["mid"], snapshot) # Rebuilt tensors must be writable, like regular tensors. received["mid"] += 1.0 - received["big"][0, 0] = 1.0 + received_big[0, 0] = 1.0 dist.barrier() print(f"tensor broadcast passed the test! Rank {rank}") @@ -427,6 +427,61 @@ def test_tensor_broadcast(): distributed_run(worker_fn_tensor_broadcast, 2) +@pytest.mark.parametrize("kind", ["tensor", "numpy", "numpy_readonly"]) +@pytest.mark.parametrize( + "nbytes, overflow, num_oob", + [ + pytest.param(1024 * 1024 - 2, False, 0, id="inline"), + pytest.param(1024 * 1024, False, 1, id="shm_oob"), + pytest.param(2 * 1024 * 1024, True, 1, id="overflow"), + ], +) +def test_broadcast_retains_values_after_ring_reuse(kind, nbytes, overflow, num_oob): + """Received data belongs to its consumer, not the next queue message.""" + writer = MessageQueue( + n_reader=1, + n_local_reader=1, + max_chunk_bytes=2 * 1024 * 1024, + max_chunks=1, + connect_ip="127.0.0.1", + ) + reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0) + writer.wait_until_ready() + reader.wait_until_ready() + array = np.full(nbytes // 2, 17, dtype=np.uint16) + array.flags.writeable = kind != "numpy_readonly" + payload = torch.from_numpy(array) if kind == "tensor" else array + try: + with ( + mock.patch.object(MessageQueue, "recv", wraps=MessageQueue.recv) as recv, + mock.patch.object(pickle, "loads", wraps=pickle.loads) as loads, + ): + writer.enqueue(payload) + received = reader.dequeue(timeout=5) + received_oob = len(loads.call_args.kwargs["buffers"]) + # The spy must not keep the ring's memoryviews alive. + loads.reset_mock() + assert recv.call_count == int(overflow) + assert received_oob == num_oob + assert type(received) is type(payload) + assert received.dtype == payload.dtype + if isinstance(received, np.ndarray): + assert received.flags.writeable == array.flags.writeable + np.testing.assert_array_equal(received, array) + for value in range(4): + # Fit the replacement in SHM even when the original overflowed. + writer.enqueue(np.full(512 * 1024, value, dtype=np.uint16)) + replacement = reader.dequeue(timeout=5) + np.testing.assert_array_equal(replacement, value) + np.testing.assert_array_equal(received, array) + if array.flags.writeable: + received[0] = 42 + assert array[0] == 17 + finally: + reader.shutdown() + writer.shutdown() + + def _dumps_oob(obj) -> tuple[bytes, list]: """Pickle `obj` the same way `MessageQueue.enqueue` does: tensor dispatch table + out-of-band buffers >= 1MiB.""" diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index 5656e8aaa153..c5180d68a3c4 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -408,8 +408,8 @@ def _rebuild_tensor(buf: Any, shape: tuple[int, ...], dtype_str: str) -> torch.T Counterpart of `_reduce_tensor`. Note that pickle passes the original buffer-providing object from `loads(buffers=...)` straight to this function (no `PickleBuffer` wrapper on the receiving side), so `buf` is - a `zmq.Frame`, a `memoryview` of a shared-memory ring chunk, or `bytes` - if the buffer was serialized in-band. + a `zmq.Frame`, an owned `bytearray` from the shared-memory ring, or + `bytes`/`bytearray` if the buffer was serialized in-band. """ dtype = getattr(torch, dtype_str) assert isinstance(dtype, torch.dtype) @@ -423,11 +423,8 @@ def _rebuild_tensor(buf: Any, shape: tuple[int, ...], dtype_str: str) -> torch.T except ValueError: # Empty or read-only frame buffer; fall through to the copy path. pass - # Shared-memory ring buffer chunks are reused by the writer once all - # readers have marked them read, so we must copy out of them. bytearray - # (vs bytes) keeps the resulting tensor writable, matching normal tensor - # semantics. - raw = bytearray(buf) + # dequeue already owns SHM buffers. Copy other buffers to keep them writable. + raw = buf if isinstance(buf, bytearray) else bytearray(buf) if not raw: assert 0 in shape return torch.empty(shape, dtype=dtype) @@ -981,7 +978,9 @@ def dequeue( buf_offset = offset + 4 buf_len = from_bytes_big(buf[offset:buf_offset]) offset = buf_offset + buf_len - all_buffers.append(buf[buf_offset:offset]) + buffer = buf[buf_offset:offset] + # Out-of-band arrays may outlive this reusable ring slot. + all_buffers.append(buffer if i == 0 else bytearray(buffer)) obj = pickle.loads(all_buffers[0], buffers=all_buffers[1:]) if overflow: obj = MessageQueue.recv(self.local_socket, timeout) From 6a70913590d8eab63320da4f3ce70976c9192500 Mon Sep 17 00:00:00 2001 From: Zhenzhong Xu Date: Mon, 12 Oct 2026 14:47:59 +0800 Subject: [PATCH 276/624] [XPU] bump up auto-round-lib to 0.16.0 (#60976) Signed-off-by: Zhenzhong1 Co-authored-by: Kunshang Ji --- requirements/xpu.txt | 2 +- tests/quantization/test_auto_round.py | 17 +---------------- 2 files changed, 2 insertions(+), 17 deletions(-) diff --git a/requirements/xpu.txt b/requirements/xpu.txt index d9d0f3080f86..14372071f0b5 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -19,5 +19,5 @@ torchaudio torchvision torchcodec >= 0.14 # Required for the torchcodec video decoding backend -auto_round_lib==0.15.0 +auto_round_lib==0.16.0 vllm_xpu_kernels==0.1.15.4 diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py index eb8273c6d6ba..267e1a11ba86 100644 --- a/tests/quantization/test_auto_round.py +++ b/tests/quantization/test_auto_round.py @@ -146,21 +146,6 @@ }, } -# These checkpoints fail during XPU graph capture, so run them eagerly there. -XPU_EAGER_ONLY_MODELS = { - "OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc", - "Intel/Qwen2-0.5B-Instruct-int4-sym-AutoRound", - "Intel/Qwen3-8B-w2g64-for-ut", - "INCModel/Qwen3-30B-A3B-12L-W4A16-test", -} - - -def _runner_kwargs(model: str) -> dict[str, Any]: - kwargs = dict(MODEL_RUNNER_KWARGS.get(model, {})) - if current_platform.is_xpu() and model in XPU_EAGER_ONLY_MODELS: - kwargs["enforce_eager"] = True - return kwargs - @pytest.mark.skipif( not ( @@ -172,7 +157,7 @@ def _runner_kwargs(model: str) -> dict[str, Any]: ) @pytest.mark.parametrize("model", MODELS + QWEN3_AUTOROUND_MODELS) def test_auto_round_model(vllm_runner, model): - with vllm_runner(model, **_runner_kwargs(model)) as llm: + with vllm_runner(model, **MODEL_RUNNER_KWARGS.get(model, {})) as llm: output = llm.generate_greedy(["The capital of France is"], max_tokens=8) assert output From 36f397281295f5b58852a7eaf429b90348161079 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Mon, 12 Oct 2026 00:09:10 -0700 Subject: [PATCH 277/624] [Spec Decode] Skip metadata build for dflash/dspark when all model backends support it (#60692) Signed-off-by: Giancarlo Delfin Co-authored-by: Rishi Puri <20074092+puririshi98@users.noreply.github.com> --- .../attention/test_deepseek_v4_swa_visible.py | 54 ++++++ .../test_triton_attn_metadata_builder.py | 154 ++++++++++++++++++ .../v1/attention/test_triton_attn_segments.py | 73 --------- ...test_gpu_target_dependent_ar_speculator.py | 22 +++ .../deepseek_v41/nvidia/flashinfer_sparse.py | 44 +++-- vllm/v1/attention/backends/mla/sparse_swa.py | 55 +++++-- vllm/v1/attention/backends/triton_attn.py | 3 + .../gpu/spec_decode/dflash/cudagraph.py | 9 +- .../gpu/spec_decode/dflash/speculator.py | 45 +++-- vllm/v1/worker/gpu/spec_decode/speculator.py | 21 +++ .../target_dependent_ar/speculator.py | 52 +----- 11 files changed, 360 insertions(+), 172 deletions(-) create mode 100644 tests/v1/attention/test_triton_attn_metadata_builder.py delete mode 100644 tests/v1/attention/test_triton_attn_segments.py diff --git a/tests/v1/attention/test_deepseek_v4_swa_visible.py b/tests/v1/attention/test_deepseek_v4_swa_visible.py index 8f08929549a8..5ccd72180846 100644 --- a/tests/v1/attention/test_deepseek_v4_swa_visible.py +++ b/tests/v1/attention/test_deepseek_v4_swa_visible.py @@ -828,3 +828,57 @@ def test_v41_flashinfer_mixed_sparse_indices_respect_replay_start(): # Decode row is copied through; prefill rows follow the bounded window. assert bounded[0].cpu().tolist() == plain[0].cpu().tolist() assert bounded[1:].cpu().tolist() == rows[1:] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_dspark_noncausal_draft_decode_update_matches_build(): + """A DSpark draft graph refreshes its non-causal SWA metadata in place + instead of rebuilding it, so the refresh must match a fresh build.""" + from vllm.v1.attention.backends.mla.compressor_utils import ( + get_dspark_swa_index_width, + ) + + num_spec = 4 + query_lens = [num_spec] * 3 + + def make_dspark_builder(): + builder = make_builder(vision=False) + builder.is_dspark = True + builder.noncausal_index_width = get_dspark_swa_index_width(WINDOW, num_spec) + builder.decode_threshold = num_spec + return builder + + def common_metadata(seq_lens): + query_start_loc, seq_lens_t, _, slot_mapping, block_table = make_batch( + seq_lens, query_lens, torch.device("cuda") + ) + return CommonAttentionMetadata( + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc.cpu(), + seq_lens=seq_lens_t, + seq_lens_cpu_upper_bound=seq_lens_t.cpu(), + num_reqs=len(seq_lens), + num_actual_tokens=int(query_start_loc[-1]), + max_query_len=num_spec, + max_seq_len=max(seq_lens), + block_table_tensor=block_table, + slot_mapping=slot_mapping, + causal=False, + ) + + captured_seq_lens = [6, 40, 130] + replay_seq_lens = [130, 9, 90] + builder = make_dspark_builder() + captured = common_metadata(captured_seq_lens) + md = builder.build(0, captured) + replay = common_metadata(replay_seq_lens) + expected = make_dspark_builder().build(0, replay) + + # A replay sees new inputs through the same device buffers. + captured.seq_lens.copy_(replay.seq_lens) + captured.block_table_tensor.copy_(replay.block_table_tensor) + builder.update_draft_decode_metadata(md) + + assert md.decode_swa_width == builder.noncausal_index_width + torch.testing.assert_close(md.decode_swa_lens, expected.decode_swa_lens) + torch.testing.assert_close(md.decode_swa_indices, expected.decode_swa_indices) diff --git a/tests/v1/attention/test_triton_attn_metadata_builder.py b/tests/v1/attention/test_triton_attn_metadata_builder.py new file mode 100644 index 000000000000..9ec814556633 --- /dev/null +++ b/tests/v1/attention/test_triton_attn_metadata_builder.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata +from vllm.config import CUDAGraphMode +from vllm.platforms import current_platform +from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadataBuilder +from vllm.v1.attention.backends.triton_attn_diffkv import ( + TritonAttentionDiffKVMetadataBuilder, +) +from vllm.v1.kv_cache_interface import FullAttentionSpec + + +# 64 segments only while num_seqs * num_kv_heads * 16 < 188 SMs on SM12.0 and +# the 64-segment scratch holds every query token (e.g. multi-query verify). +@pytest.mark.parametrize( + "builder_cls", + [TritonAttentionMetadataBuilder, TritonAttentionDiffKVMetadataBuilder], +) +@pytest.mark.parametrize( + "capability,num_kv_heads,num_seqs,query_len,segments", + [ + ((12, 0), 1, 11, 1, 64), + ((12, 0), 1, 12, 1, 16), + ((12, 0), 8, 1, 1, 64), + ((12, 0), 8, 2, 1, 16), + ((9, 0), 1, 1, 1, 16), + ((12, 0), 1, 2, 4, 64), + ((12, 0), 1, 10, 4, 16), + ], +) +def test_split_k_segments_follow_sm_occupancy( + monkeypatch, builder_cls, capability, num_kv_heads, num_seqs, query_len, segments +): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + current_platform, "is_device_capability", lambda cap: cap == capability + ) + monkeypatch.setattr(current_platform, "num_compute_units", lambda: 188) + config = SimpleNamespace( + model_config=SimpleNamespace( + get_num_attention_heads=lambda _: 8, + get_num_kv_heads=lambda _: num_kv_heads, + get_head_size=lambda: 128, + rswa_window=None, + ), + parallel_config=None, + scheduler_config=SimpleNamespace(max_num_seqs=16), + speculative_config=None, + compilation_config=SimpleNamespace( + cudagraph_mode=CUDAGraphMode.NONE, static_forward_context={} + ), + ) + spec = FullAttentionSpec( + block_size=16, num_kv_heads=num_kv_heads, head_size=128, dtype=torch.bfloat16 + ) + builder = builder_cls(spec, ["layer.0"], config, "cpu") + batch = BatchSpec(seq_lens=[128] * num_seqs, query_lens=[query_len] * num_seqs) + metadata = builder.build(0, create_common_attn_metadata(batch, 16, "cpu")) + assert metadata.num_par_softmax_segments == segments + assert metadata.softmax_segm_output.shape[2] == segments + assert metadata.softmax_segm_max.shape[0] >= num_seqs * query_len + # DiffKV re-allocates the output; it must stay sized like the base buffers. + assert builder.softmax_segm_output.shape[0] == builder.softmax_segm_max.shape[0] + assert ( + metadata.softmax_segm_output.data_ptr() + == builder.softmax_segm_output.data_ptr() + ) + + +def _draft_builder(builder_cls, rswa_window: int | None): + config = SimpleNamespace( + model_config=SimpleNamespace( + get_num_attention_heads=lambda _: 8, + get_num_kv_heads=lambda _: 1, + get_head_size=lambda: 128, + rswa_window=rswa_window, + ), + parallel_config=None, + scheduler_config=SimpleNamespace(max_num_seqs=16), + speculative_config=None, + compilation_config=SimpleNamespace( + cudagraph_mode=CUDAGraphMode.NONE, static_forward_context={} + ), + ) + spec = FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=128, dtype=torch.bfloat16 + ) + return builder_cls(spec, ["layer.0"], config, "cpu") + + +def _draft_decode_batch(seq_lens: list[int]): + batch = BatchSpec(seq_lens=seq_lens, query_lens=[1] * len(seq_lens)) + return create_common_attn_metadata(batch, 16, "cpu") + + +@pytest.mark.parametrize( + "builder_cls", + [TritonAttentionMetadataBuilder, TritonAttentionDiffKVMetadataBuilder], +) +@pytest.mark.parametrize("rswa_window", [None, 64]) +def test_draft_decode_update_flag_matches_build_writes(builder_cls, rswa_window): + """The update is a no-op, so the builder may claim support only when + build() writes no builder-owned tensor that a captured graph reads.""" + builder = _draft_builder(builder_cls, rswa_window) + common_metadata = _draft_decode_batch([5, 9]) + common_metadata.rswa_prefix_lens = torch.tensor([3, 7], dtype=torch.int32) + + def owned_bytes(): + # Bytes, not values: uninitialized scratch may hold NaNs. + return { + name: value.view(torch.uint8).clone() + for name, value in vars(builder).items() + if isinstance(value, torch.Tensor) + } + + before = owned_bytes() + builder.build(0, common_metadata) + after = owned_bytes() + + written = [name for name in before if not torch.equal(before[name], after[name])] + assert builder.supports_draft_decode_metadata_update is (not written) + assert builder.supports_draft_decode_metadata_update is (rswa_window is None) + + +@pytest.mark.parametrize( + "builder_cls", + [TritonAttentionMetadataBuilder, TritonAttentionDiffKVMetadataBuilder], +) +def test_draft_decode_update_matches_fresh_build(builder_cls): + """A replayed graph reads metadata built at capture; after an in-place + update it must match a fresh build of the replayed batch.""" + torch.manual_seed(0) + builder = _draft_builder(builder_cls, rswa_window=None) + captured = _draft_decode_batch([100, 120]) + metadata = builder.build(0, captured) + + replayed = _draft_decode_batch([128, 30]) + # The replayed batch arrives through the same persistent input buffers. + captured.seq_lens.copy_(replayed.seq_lens) + captured.block_table_tensor.copy_(replayed.block_table_tensor) + captured.slot_mapping.copy_(replayed.slot_mapping) + builder.update_draft_decode_metadata(metadata) + + expected = _draft_builder(builder_cls, rswa_window=None).build(0, replayed) + assert torch.equal(metadata.seq_lens, expected.seq_lens) + assert torch.equal(metadata.block_table, expected.block_table) + assert torch.equal(metadata.slot_mapping, expected.slot_mapping) + assert torch.equal(metadata.query_start_loc, expected.query_start_loc) diff --git a/tests/v1/attention/test_triton_attn_segments.py b/tests/v1/attention/test_triton_attn_segments.py deleted file mode 100644 index 874c61d2379e..000000000000 --- a/tests/v1/attention/test_triton_attn_segments.py +++ /dev/null @@ -1,73 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from types import SimpleNamespace - -import pytest -import torch - -from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata -from vllm.config import CUDAGraphMode -from vllm.platforms import current_platform -from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadataBuilder -from vllm.v1.attention.backends.triton_attn_diffkv import ( - TritonAttentionDiffKVMetadataBuilder, -) -from vllm.v1.kv_cache_interface import FullAttentionSpec - - -# 64 segments only while num_seqs * num_kv_heads * 16 < 188 SMs on SM12.0 and -# the 64-segment scratch holds every query token (e.g. multi-query verify). -@pytest.mark.parametrize( - "builder_cls", - [TritonAttentionMetadataBuilder, TritonAttentionDiffKVMetadataBuilder], -) -@pytest.mark.parametrize( - "capability,num_kv_heads,num_seqs,query_len,segments", - [ - ((12, 0), 1, 11, 1, 64), - ((12, 0), 1, 12, 1, 16), - ((12, 0), 8, 1, 1, 64), - ((12, 0), 8, 2, 1, 16), - ((9, 0), 1, 1, 1, 16), - ((12, 0), 1, 2, 4, 64), - ((12, 0), 1, 10, 4, 16), - ], -) -def test_split_k_segments_follow_sm_occupancy( - monkeypatch, builder_cls, capability, num_kv_heads, num_seqs, query_len, segments -): - monkeypatch.setattr(current_platform, "is_cuda", lambda: True) - monkeypatch.setattr( - current_platform, "is_device_capability", lambda cap: cap == capability - ) - monkeypatch.setattr(current_platform, "num_compute_units", lambda: 188) - config = SimpleNamespace( - model_config=SimpleNamespace( - get_num_attention_heads=lambda _: 8, - get_num_kv_heads=lambda _: num_kv_heads, - get_head_size=lambda: 128, - rswa_window=None, - ), - parallel_config=None, - scheduler_config=SimpleNamespace(max_num_seqs=16), - speculative_config=None, - compilation_config=SimpleNamespace( - cudagraph_mode=CUDAGraphMode.NONE, static_forward_context={} - ), - ) - spec = FullAttentionSpec( - block_size=16, num_kv_heads=num_kv_heads, head_size=128, dtype=torch.bfloat16 - ) - builder = builder_cls(spec, ["layer.0"], config, "cpu") - batch = BatchSpec(seq_lens=[128] * num_seqs, query_lens=[query_len] * num_seqs) - metadata = builder.build(0, create_common_attn_metadata(batch, 16, "cpu")) - assert metadata.num_par_softmax_segments == segments - assert metadata.softmax_segm_output.shape[2] == segments - assert metadata.softmax_segm_max.shape[0] >= num_seqs * query_len - # DiffKV re-allocates the output; it must stay sized like the base buffers. - assert builder.softmax_segm_output.shape[0] == builder.softmax_segm_max.shape[0] - assert ( - metadata.softmax_segm_output.data_ptr() - == builder.softmax_segm_output.data_ptr() - ) diff --git a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py index a3ec5f4eb7f4..490b84e77896 100644 --- a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py +++ b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py @@ -609,3 +609,25 @@ def fake_get_scheduler_metadata(**kwargs): assert not called assert metadata.scheduler_metadata is None + + +@pytest.mark.parametrize(("num_speculative_steps", "expected"), [(1, False), (3, True)]) +def test_decode_metadata_update_needs_draft_decode_steps( + num_speculative_steps, expected +): + """One speculative step runs no draft decodes, so the in-place metadata + update (and its fallback log) does not apply even on a capable backend.""" + speculator = object.__new__(_TestSpeculator) + speculator.num_speculative_steps = num_speculative_steps + speculator.attn_groups = [ + [ + SimpleNamespace( + supports_draft_decode_metadata_update=True, + backend=SimpleNamespace(get_name=lambda: "CAPABLE"), + ) + ] + ] + + speculator._configure_decode_metadata_update() + + assert speculator.supports_decode_metadata_update is expected diff --git a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py index 0e02b486a517..a5ab09283959 100644 --- a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py @@ -206,24 +206,36 @@ def build( metadata = super().build( common_prefix_len, common_attn_metadata, fast_build, replay_start ) - num_tokens = metadata.num_decode_tokens - if not common_attn_metadata.causal and num_tokens > 0: - assert metadata.decode_swa_lens is not None - assert metadata.seq_lens is not None - assert metadata.token_to_req_indices is not None - topk_lens = self._decode_topk_lens[:num_tokens] - seq_lens = self._decode_seq_lens[:num_tokens] - torch.clamp(metadata.decode_swa_lens, min=self.window_size, out=topk_lens) - torch.index_select( - metadata.seq_lens, - 0, - metadata.token_to_req_indices[:num_tokens], - out=seq_lens, - ) - metadata.flashinfer_decode_topk_lens = topk_lens - metadata.flashinfer_decode_seq_lens = seq_lens + if not common_attn_metadata.causal and metadata.num_decode_tokens > 0: + self._fill_noncausal_decode_lens(metadata) return metadata + def update_draft_decode_metadata( + self, metadata: "DeepseekSparseSWAMetadata" + ) -> None: + super().update_draft_decode_metadata(metadata) + if metadata.flashinfer_decode_topk_lens is not None: + self._fill_noncausal_decode_lens(metadata) + + def _fill_noncausal_decode_lens( + self, metadata: "DeepseekSparseSWAMetadata" + ) -> None: + num_tokens = metadata.num_decode_tokens + assert metadata.decode_swa_lens is not None + assert metadata.seq_lens is not None + assert metadata.token_to_req_indices is not None + topk_lens = self._decode_topk_lens[:num_tokens] + seq_lens = self._decode_seq_lens[:num_tokens] + torch.clamp(metadata.decode_swa_lens, min=self.window_size, out=topk_lens) + torch.index_select( + metadata.seq_lens, + 0, + metadata.token_to_req_indices[:num_tokens], + out=seq_lens, + ) + metadata.flashinfer_decode_topk_lens = topk_lens + metadata.flashinfer_decode_seq_lens = seq_lens + class DeepseekSparseSWAFlashInferBackend(DeepseekSparseSWABackend): @staticmethod diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index c5cbb4ef2822..a691efe8bfb4 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -852,23 +852,44 @@ def update_draft_decode_metadata( num_tokens, ) - _COMPUTE_SWA_INDICES_AND_LENS_KERNEL( - metadata.decode_swa_indices, - metadata.decode_swa_lens, - metadata.decode_swa_indices.shape[-1], - metadata.decode_swa_indices.shape[-1], - metadata.decode_swa_lens, # unused (HAS_IMAGE=False) - metadata.decode_swa_lens, # unused (HAS_IMAGE=False) - metadata.query_start_loc, - metadata.seq_lens, - metadata.token_to_req_indices, - metadata.is_valid_token, - metadata.block_table, - self.block_size, - metadata.replay_start, - num_tokens=metadata.num_decode_tokens, - token_offset=0, - ) + # build() points non-causal (DSpark draft) metadata at its own buffer. + noncausal = self.decode_swa_indices_noncausal + if ( + noncausal is not None + and metadata.decode_swa_indices.data_ptr() == noncausal.data_ptr() + ): + _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL( + metadata.decode_swa_indices, + metadata.decode_swa_lens, + self.window_size, + self.noncausal_index_width, + metadata.query_start_loc, + metadata.seq_lens, + metadata.token_to_req_indices, + metadata.is_valid_token, + metadata.block_table, + self.block_size, + num_tokens=num_tokens, + token_offset=0, + ) + else: + _COMPUTE_SWA_INDICES_AND_LENS_KERNEL( + metadata.decode_swa_indices, + metadata.decode_swa_lens, + metadata.decode_swa_indices.shape[-1], + metadata.decode_swa_indices.shape[-1], + metadata.decode_swa_lens, # unused (HAS_IMAGE=False) + metadata.decode_swa_lens, # unused (HAS_IMAGE=False) + metadata.query_start_loc, + metadata.seq_lens, + metadata.token_to_req_indices, + metadata.is_valid_token, + metadata.block_table, + self.block_size, + metadata.replay_start, + num_tokens=metadata.num_decode_tokens, + token_offset=0, + ) tile_sched = self.build_tile_scheduler(metadata.num_decode_tokens) metadata.tile_sched_swaonly = tile_sched[_LAYER_TYPE_SWAONLY] metadata.tile_sched_c4a = tile_sched[_LAYER_TYPE_C4A] diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index d8577b61f1c9..095163f14423 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -193,6 +193,9 @@ def __init__( device=device, ) self.rswa_window = model_config.rswa_window + # R-SWA restages prefix lengths in build(), which the no-op + # update_draft_decode_metadata() does not refresh. + self.supports_draft_decode_metadata_update = self.rswa_window is None self.persistent_rswa_prefix_lens: torch.Tensor | None = None if self.rswa_window is not None: self.persistent_rswa_prefix_lens = torch.empty( diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py index e21ca9f8707e..2e189da0c939 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable, Mapping +from typing import Any import torch @@ -84,11 +85,11 @@ def capture( kv_cache_config: KVCacheConfig, max_model_len: int, causal: bool | Mapping[int, bool], - precompute_context_kv: Callable[[int], None], + prepare_forward: Callable[[int, dict[str, Any] | None], None], progress_bar_desc: str = "Capturing CUDA graphs", ) -> None: - """``precompute_context_kv(num_reqs)`` is captured ahead of the query - forward.""" + """``prepare_forward(num_reqs, attn_metadata)`` is captured ahead of the + query forward.""" def create_forward_fn( desc: BatchExecutionDescriptor, @@ -115,7 +116,7 @@ def create_forward_fn( attn_metadata, slot_mappings = attn_state def forward(cg_mode: CUDAGraphMode) -> None: - precompute_context_kv(num_reqs) + prepare_forward(num_reqs, attn_metadata) forward_fn( num_reqs, num_tokens, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 884de5a264be..5a13441ee892 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -166,12 +166,18 @@ def capture(self) -> None: self.kv_cache_config, self.max_model_len, causal=self._group_causal, - precompute_context_kv=lambda num_reqs: self._precompute_context_kv( - 0, self._num_graph_context_tokens(num_reqs) - ), + prepare_forward=self._prepare_graph_forward, progress_bar_desc=f"Capturing {self._speculator_name.lower()} CUDA graphs", ) + def _prepare_graph_forward( + self, num_reqs: int, attn_metadata: dict[str, Any] | None + ) -> None: + self._precompute_context_kv(0, self._num_graph_context_tokens(num_reqs)) + # Recorded in the FULL graph so replays skip the eager metadata build. + if attn_metadata is not None and self.supports_decode_metadata_update: + self._update_draft_decode_metadata(attn_metadata, num_reqs) + def load_draft_model( self, target_model: nn.Module, @@ -486,29 +492,32 @@ def propose( else: self._precompute_context_kv(0, num_target_tokens, dummy_run) - # Rebuild the draft attention metadata even when replaying the FULL - # graph so that any attention metadata builder state is updated. - draft_attn_metadata = self._build_uniform_attn_metadata( - num_reqs=num_reqs, - batch_desc=batch_desc, - num_query_per_req=self.num_query_per_req, - seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound, - step=self.num_query_per_req, - causal=self._group_causal, - ) - draft_slot_mappings_by_layer = build_slot_mappings_by_layer( - self.block_tables.slot_mappings[:, :num_tokens_padded], - self.kv_cache_config, - ) + is_full_cudagraph = batch_desc.cg_mode == CUDAGraphMode.FULL + # During FULL cudagraph with a model that supports decode metadata updates, + # the metadata update is captured in the graph, and we can skip the full + # build below. + if not (is_full_cudagraph and self.supports_decode_metadata_update): + draft_attn_metadata = self._build_uniform_attn_metadata( + num_reqs=num_reqs, + batch_desc=batch_desc, + num_query_per_req=self.num_query_per_req, + seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound, + step=self.num_query_per_req, + causal=self._group_causal, + ) # DFlash processes all speculative tokens in one forward pass, # so the real token count is num_query_tokens. self._prepare_eplb_forward(num_query_tokens) - if batch_desc.cg_mode == CUDAGraphMode.FULL: + if is_full_cudagraph: assert self.query_cudagraph_manager is not None self.query_cudagraph_manager.run_fullgraph(batch_desc) else: + draft_slot_mappings_by_layer = build_slot_mappings_by_layer( + self.block_tables.slot_mappings[:, :num_tokens_padded], + self.kv_cache_config, + ) self._generate_draft( num_reqs, num_tokens_padded, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index c3e8dcbff4d1..d86a0b0351f0 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -207,6 +207,9 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.supports_mm_inputs = False self.pcp_manager: PCPManager | None = None + # Whether every draft attention builder can refresh its decode metadata + # in place, so the refresh can be recorded in the draft CUDA graph. + self.supports_decode_metadata_update = False @abstractmethod def load_draft_model( @@ -299,6 +302,24 @@ def set_attn( # builders and buffers. self.target_input_buffers = target_input_buffers self.target_attn_groups = target_attn_groups + self._configure_decode_metadata_update() + + def _configure_decode_metadata_update(self) -> None: + unsupported_backends = sorted( + { + attn_group.backend.get_name() + for attn_groups in self.attn_groups + for attn_group in attn_groups + if not attn_group.supports_draft_decode_metadata_update + } + ) + self.supports_decode_metadata_update = not unsupported_backends + if unsupported_backends: + logger.info_once( + "Attention backend(s) %s cannot update draft decode metadata in " + "place; falling back to rebuilding attention metadata eagerly.", + ", ".join(unsupported_backends), + ) def _build_attn_metadata( self, diff --git a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py index 02e28875b823..fdbedd78f6bd 100644 --- a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py @@ -10,18 +10,15 @@ from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger from vllm.triton_utils import tl, triton -from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer -from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers -from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.cudagraph_utils import ( SpeculatorCudaGraphManager, ) -from vllm.v1.worker.utils import AttentionGroup, get_uniform_decode_token_count +from vllm.v1.worker.utils import get_uniform_decode_token_count logger = init_logger(__name__) @@ -51,7 +48,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.prefill_cudagraph_manager: SpeculatorCudaGraphManager | None = None self.decode_cudagraph_manager: SpeculatorCudaGraphManager | None = None - self.use_fused_multi_step_decode = False def load_model(self, target_model: nn.Module) -> None: super().load_model(target_model) @@ -83,44 +79,12 @@ def advance_draft_positions(self) -> bool: """ return True - def set_attn( - self, - model_state: ModelState, - kv_cache_config: KVCacheConfig, - block_tables: BlockTables, - target_input_buffers: InputBuffers, - target_attn_groups: list[list[AttentionGroup]], - ) -> None: - super().set_attn( - model_state, - kv_cache_config, - block_tables, - target_input_buffers, - target_attn_groups, - ) - self._configure_fused_multi_step_decode() - - def _configure_fused_multi_step_decode(self) -> None: + def _configure_decode_metadata_update(self) -> None: + # With one speculative step there are no draft decode steps to update. if self.num_speculative_steps == 1: - self.use_fused_multi_step_decode = False + self.supports_decode_metadata_update = False return - - unsupported_backends = sorted( - { - attn_group.backend.get_name() - for attn_groups in self.attn_groups - for attn_group in attn_groups - if not attn_group.supports_draft_decode_metadata_update - } - ) - self.use_fused_multi_step_decode = not unsupported_backends - if unsupported_backends: - logger.info_once( - "Fused multi-step draft decode is not supported by attention " - "backend(s) %s; falling back to rebuilding attention metadata " - "between draft steps.", - ", ".join(unsupported_backends), - ) + super()._configure_decode_metadata_update() def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: # Initialize cudagraph manager for draft prefill (draft position 0). @@ -181,7 +145,7 @@ def capture(self) -> None: assert self.decode_cudagraph_manager is not None decode_fn = ( self._generate_fused_drafts - if self.use_fused_multi_step_decode + if self.supports_decode_metadata_update else self._generate_draft ) self.decode_cudagraph_manager.capture( @@ -192,7 +156,7 @@ def capture(self) -> None: self.attn_groups, self.kv_cache_config, progress_bar_desc="Capturing decode CUDA graphs", - specialize_spec_tokens=self.use_fused_multi_step_decode + specialize_spec_tokens=self.supports_decode_metadata_update and self.speculative_config.uses_dynamic_speculative_decoding(), ) self.on_multi_step_decode_end(self.max_num_reqs) @@ -384,7 +348,7 @@ def propose( # Generate the remaining num_speculative_steps - 1 draft tokens. decode_fn = ( self._fused_multi_step_decode - if self.use_fused_multi_step_decode + if self.supports_decode_metadata_update else self._multi_step_decode ) decode_fn( From e3bbaf23695f7f30f144a288057b0419ab2b6a8d Mon Sep 17 00:00:00 2001 From: AMD-yanfeiwang Date: Mon, 12 Oct 2026 15:53:16 +0800 Subject: [PATCH 278/624] [ROCm][DSv4] Fuse the shared expert into AITER MegaMoE (#61004) Signed-off-by: Yanfei Wang Co-authored-by: aoshen02 --- .../models/test_deepseek_v4_aiter_mega_moe.py | 20 +++ .../device_communicators/all2all.py | 6 +- vllm/envs.py | 7 ++ vllm/models/deepseek_v4/amd/mega_moe.py | 117 ++++++++++++++++-- vllm/models/deepseek_v4/amd/model.py | 13 +- 5 files changed, 151 insertions(+), 12 deletions(-) diff --git a/tests/models/test_deepseek_v4_aiter_mega_moe.py b/tests/models/test_deepseek_v4_aiter_mega_moe.py index a8f692e07bfb..19b3697811df 100644 --- a/tests/models/test_deepseek_v4_aiter_mega_moe.py +++ b/tests/models/test_deepseek_v4_aiter_mega_moe.py @@ -29,6 +29,7 @@ def _make_experts( intermediate_size=INTER, max_num_batched_tokens=3000, swiglu_limit=None, + fuse_shared_expert=False, ): monkeypatch.setattr( mega_moe_module, @@ -47,6 +48,7 @@ def _make_experts( hidden_size=HIDDEN, intermediate_size=intermediate_size, swiglu_limit=swiglu_limit, + fuse_shared_expert=fuse_shared_expert, ) @@ -108,6 +110,24 @@ def test_experts_load_into_local_slots(monkeypatch): assert torch.all(w2 == 10 * expert + 3) +@pytest.mark.skipif(not current_platform.is_rocm(), reason="Requires ROCm") +def test_shared_expert_route_appended(monkeypatch): + experts = _make_experts( + monkeypatch, ep_rank=3, ep_size=8, num_experts=384, fuse_shared_expert=True + ) + ids = torch.randint(0, 384, (37, 2), device="cuda", dtype=torch.int32) + ids[::5, 1] = -1 + weights = torch.rand(37, 2, device="cuda") + out_weights, out_ids = experts.append_shared_expert(weights, ids) + + shared = torch.full((37, 1), 3 * 49 + 48, device="cuda", dtype=torch.int32) + expected_ids = torch.cat([torch.where(ids >= 0, ids + ids // 48, ids), shared], 1) + torch.testing.assert_close(out_ids, expected_ids) + torch.testing.assert_close( + out_weights, torch.cat([weights, torch.ones(37, 1, device="cuda")], 1) + ) + + @pytest.mark.skipif(not current_platform.is_rocm(), reason="Requires ROCm") @pytest.mark.parametrize("hash_routing", [False, True]) @pytest.mark.parametrize("num_tokens", [1, 37]) diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index cff34139f968..3466700c3334 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -1012,7 +1012,11 @@ def __init__(self, cpu_group, all2all_backend: str): if get_current_vllm_config().kernel_config.moe_backend == "aiter_mega_moe": # MegaMoEV2 places its dispatch/combine workspaces on the MoRI # symmetric heap, which defaults to 2 GB. MegaMoEV2 requires > 4GB - heap_size = os.environ.setdefault("MORI_SHMEM_HEAP_SIZE", "8G") + # (> 8GB with a fused shared expert). + heap_size = os.environ.setdefault( + "MORI_SHMEM_HEAP_SIZE", + "16G" if envs.VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS else "8G", + ) logger.info_once("AITER MegaMoE: MORI_SHMEM_HEAP_SIZE=%s", heap_size) mori.shmem.shmem_torch_process_group_init("mori") diff --git a/vllm/envs.py b/vllm/envs.py index 14489a3f2503..a2b46030772f 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -155,6 +155,7 @@ VLLM_ROCM_USE_AITER_FP4BMM: bool = True VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = False + VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS: bool = False VLLM_ROCM_MONO_DECODE: bool = False VLLM_ROCM_USE_AITER_TRITON_GEMM: bool = True VLLM_ROCM_USE_SKINNY_GEMM: bool = True @@ -1391,6 +1392,12 @@ def _resolve_rust_cli_path() -> str | None: os.getenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "False").lower() in ("true", "1") ), + # Fuse DeepSeek-V4's shared expert into AITER MegaMoE (requantized to + # MXFP4). By default is disabled. + "VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS": lambda: ( + os.getenv("VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS", "False").lower() + in ("true", "1") + ), # Run the eligible layers of decode steps on the model's mono decode # kernels: persistent launches a layer, the TP all-reduces in-kernel # (DeepSeek-V4.1 on CDNA4). By default is disabled. diff --git a/vllm/models/deepseek_v4/amd/mega_moe.py b/vllm/models/deepseek_v4/amd/mega_moe.py index 2fdff9335145..312c089e5aad 100644 --- a/vllm/models/deepseek_v4/amd/mega_moe.py +++ b/vllm/models/deepseek_v4/amd/mega_moe.py @@ -10,10 +10,51 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig from vllm.distributed import get_ep_group +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + get_fp8_block_weight_scale, +) +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mxfp4_quantize +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + scaled_dequantize, +) from vllm.model_executor.utils import set_weight_attrs +from vllm.triton_utils import tl, triton from vllm.utils.math_utils import next_power_of_2 +@triton.jit +def _append_shared_expert_kernel( + topk_ids_ptr, + topk_weights_ptr, + out_ids_ptr, + out_weights_ptr, + num_tokens, + shared_expert_id, + TOPK: tl.constexpr, + ROUTED_PER_RANK: tl.constexpr, + BLOCK_M: tl.constexpr, +): + rows = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + mask = rows < num_tokens + for k in tl.static_range(TOPK): + e = tl.load(topk_ids_ptr + rows * TOPK + k, mask=mask, other=0) + w = tl.load(topk_weights_ptr + rows * TOPK + k, mask=mask, other=0.0) + e = tl.where(e >= 0, e + e // ROUTED_PER_RANK, e) + tl.store(out_ids_ptr + rows * (TOPK + 1) + k, e, mask=mask) + tl.store(out_weights_ptr + rows * (TOPK + 1) + k, w, mask=mask) + tl.store( + out_ids_ptr + rows * (TOPK + 1) + TOPK, + tl.zeros_like(rows) + shared_expert_id, + mask=mask, + ) + tl.store( + out_weights_ptr + rows * (TOPK + 1) + TOPK, + tl.full([BLOCK_M], 1.0, tl.float32), + mask=mask, + ) + + class DeepseekV4AiterMegaMoEExperts(nn.Module): # MegaMoEV2 owns large symmetric-heap buffers, so every layer shares one # instance and passes its own weights per call. @@ -28,16 +69,25 @@ def __init__( hidden_size: int, intermediate_size: int, swiglu_limit: float | None, + fuse_shared_expert: bool = False, prefix: str = "", ): super().__init__() ep_group = get_ep_group() self.ep_rank = ep_group.rank_in_group self.ep_size = ep_group.world_size - self.num_experts = num_experts - self.num_local_experts = num_experts // self.ep_size - self.experts_start_idx = self.ep_rank * self.num_local_experts - self.top_k = top_k + self.num_routed_local_experts = num_experts // self.ep_size + self.experts_start_idx = self.ep_rank * self.num_routed_local_experts + # A fused shared expert is one extra local expert per rank and one + # extra route per token, to the token's own rank. + self.fuse_shared_expert = fuse_shared_expert + self.num_local_experts = self.num_routed_local_experts + fuse_shared_expert + self.num_experts = self.num_local_experts * self.ep_size + self.routed_top_k = top_k + self.top_k = top_k + fuse_shared_expert + self.shared_expert_id = ( + self.ep_rank * self.num_local_experts + self.num_routed_local_experts + ) self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.swiglu_limit = swiglu_limit or 0.0 @@ -70,7 +120,7 @@ def weight_loader( return_success: bool = False, ) -> bool | None: local_expert_id = expert_id - self.experts_start_idx - if not 0 <= local_expert_id < self.num_local_experts: + if not 0 <= local_expert_id < self.num_routed_local_experts: return False if return_success else None expert_data = param.data[local_expert_id] if shard_id in ("w1", "w3"): @@ -81,6 +131,24 @@ def weight_loader( expert_data.copy_(loaded_weight.view(torch.uint8)) return True if return_success else None + def load_shared_expert(self, shared_experts: nn.Module) -> None: + slot = self.num_routed_local_experts + for linear, weight, weight_scale in ( + (shared_experts.gate_up_proj, self.w13_weight, self.w13_weight_scale), + (shared_experts.down_proj, self.w2_weight, self.w2_weight_scale), + ): + block_scale = get_fp8_block_weight_scale(linear) + assert block_scale is not None, "shared expert must be block-FP8" + group = GroupShape( + linear.weight.shape[0] // block_scale.shape[0], + linear.weight.shape[1] // block_scale.shape[1], + ) + packed, scale = mxfp4_quantize( + scaled_dequantize(linear.weight, block_scale, group, torch.bfloat16) + ) + weight.data[slot].copy_(packed.view(torch.uint8)) + weight_scale.data[slot].copy_(scale.view(torch.uint8)) + def finalize_weights(self) -> None: if self._runtime is not None: return @@ -129,6 +197,33 @@ def finalize_weights(self) -> None: self._runtime_cache[key] = runtime self._runtime = runtime + def append_shared_expert( + self, topk_weights: torch.Tensor, topk_ids: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Remap routed ids to the per-rank (routed + shared) layout and append + a weight-1.0 route to this rank's shared slot.""" + num_tokens = topk_ids.shape[0] + out_ids = torch.empty( + (num_tokens, self.top_k), dtype=topk_ids.dtype, device=topk_ids.device + ) + out_weights = torch.empty( + (num_tokens, self.top_k), dtype=torch.float32, device=topk_ids.device + ) + if num_tokens: + block_m = 64 + _append_shared_expert_kernel[(triton.cdiv(num_tokens, block_m),)]( + topk_ids.contiguous(), + topk_weights.to(torch.float32).contiguous(), + out_ids, + out_weights, + num_tokens, + self.shared_expert_id, + TOPK=self.routed_top_k, + ROUTED_PER_RANK=self.num_routed_local_experts, + BLOCK_M=block_m, + ) + return out_weights, out_ids + def forward( self, hidden_states: torch.Tensor, @@ -137,6 +232,8 @@ def forward( ) -> torch.Tensor: runtime = self._runtime assert runtime is not None + if self.fuse_shared_expert: + topk_weights, topk_ids = self.append_shared_expert(topk_weights, topk_ids) # MegaMoEV2 reads these per launch; rebind to this layer's experts. runtime._s1_w1 = self.w13_weight runtime._s1_w1_scale = self.w13_weight_scale @@ -149,6 +246,10 @@ def forward( def finalize_mega_moe_weights(model: nn.Module) -> None: - for module in model.modules(): - if isinstance(module, DeepseekV4AiterMegaMoEExperts): - module.finalize_weights() + for module in list(model.modules()): + experts = getattr(module, "experts", None) + if isinstance(experts, DeepseekV4AiterMegaMoEExperts): + if experts.fuse_shared_expert: + experts.load_shared_expert(module.shared_experts) + module.shared_experts = None + experts.finalize_weights() diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 8b26fc0060e7..51d22b75e8c1 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -74,7 +74,10 @@ make_layers, maybe_prefix, ) -from vllm.models.deepseek_v4.amd.mega_moe import DeepseekV4AiterMegaMoEExperts +from vllm.models.deepseek_v4.amd.mega_moe import ( + DeepseekV4AiterMegaMoEExperts, + finalize_mega_moe_weights, +) from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention from vllm.platforms import current_platform from vllm.platforms.rocm import on_gfx950 @@ -690,6 +693,10 @@ def __init__( hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, swiglu_limit=self.swiglu_limit, + fuse_shared_expert=( + config.n_shared_experts == 1 + and envs.VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS + ), prefix=f"{prefix}.experts", ) return @@ -1524,6 +1531,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def process_weights_after_loading(self) -> None: + # Must precede the gate_up preshuffle of the fused shared experts. + finalize_mega_moe_weights(self) # After per-layer quant finalize, so we preshuffle the final fp8 weights. fused_compressor_layers = 0 for module in self.modules(): @@ -1532,8 +1541,6 @@ def process_weights_after_loading(self) -> None: module.prepare_attn_preshuffle() elif isinstance(module, DeepseekV4MLP): module.prepare_gateup_preshuffle() - elif isinstance(module, DeepseekV4AiterMegaMoEExperts): - module.finalize_weights() if fused_compressor_layers: logger.info( "Fused the C4 compressor GEMMs in %d DeepSeek V4 layers", From f4114e068e5d8e961defeb7747a9449d5eec0659 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Mon, 12 Oct 2026 15:55:06 +0800 Subject: [PATCH 279/624] [Bugfix][Frontend] Preserve token IDs in chat derender responses (#61065) Signed-off-by: Ting Sun --- .../scale_out/derender/test_derender_parity.py | 5 +++-- .../scale_out/derender/test_derender_stream.py | 3 ++- tests/entrypoints/scale_out/derender/utils.py | 4 ++++ vllm/renderers/online_derenderer.py | 13 +++++++++++++ 4 files changed, 22 insertions(+), 3 deletions(-) diff --git a/tests/entrypoints/scale_out/derender/test_derender_parity.py b/tests/entrypoints/scale_out/derender/test_derender_parity.py index 4ea3bc246192..f33187387403 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_parity.py +++ b/tests/entrypoints/scale_out/derender/test_derender_parity.py @@ -126,7 +126,7 @@ async def _disagg( ], }, "prompt_tokens": prompt_tokens, - "chat_request": chat_request, + "chat_request": {"return_token_ids": True, **chat_request}, }, ) assert resp.status_code == 200, resp.text @@ -149,6 +149,7 @@ def _assert_parity(coupled: dict, disagg: dict) -> None: assert d["message"].get("reasoning") == c["message"].get("reasoning") assert _tool_sig(d) == _tool_sig(c) assert d["finish_reason"] == c["finish_reason"] + assert d["token_ids"] == c["token_ids"] assert disagg["usage"]["prompt_tokens"] == coupled["usage"]["prompt_tokens"] assert disagg["usage"]["completion_tokens"] == len(c["token_ids"]) @@ -386,7 +387,7 @@ async def _run_stream_parity_case( [tid for ch in coupled_choices for tid in ch["token_ids"]], [len(ch["token_ids"]) for ch in coupled_choices] + ([0] if finish_only_tail else []), - {"model": MODEL, "messages": messages, **extra}, + {"model": MODEL, "messages": messages, "return_token_ids": True, **extra}, len(prompt_token_ids), prompt_token_ids, finish_reason="stop" if finish_reason == "tool_calls" else finish_reason, diff --git a/tests/entrypoints/scale_out/derender/test_derender_stream.py b/tests/entrypoints/scale_out/derender/test_derender_stream.py index 0a7d96f3cf90..d97a2f992ac2 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_stream.py +++ b/tests/entrypoints/scale_out/derender/test_derender_stream.py @@ -1311,7 +1311,7 @@ async def test_finish_reason_stop_for_named_tool_choice(self, parsed_derenderer) async def test_include_reasoning_false_suppresses_reasoning( self, parsed_derenderer ): - chat_request = _chat_request(include_reasoning=False) + chat_request = _chat_request(include_reasoning=False, return_token_ids=True) chunk, _ = await parsed_derenderer.derender_chat_stream( model=MODEL_NAME, generate_chunk=_make_stream_chunk( @@ -1320,6 +1320,7 @@ async def test_include_reasoning_false_suppresses_reasoning( chat_request=chat_request, ) delta = chunk.choices[0].delta + assert chunk.choices[0].token_ids is None assert delta.reasoning is None assert delta.content == "c" diff --git a/tests/entrypoints/scale_out/derender/utils.py b/tests/entrypoints/scale_out/derender/utils.py index d8128353bf3c..c249e2053a6a 100644 --- a/tests/entrypoints/scale_out/derender/utils.py +++ b/tests/entrypoints/scale_out/derender/utils.py @@ -95,6 +95,10 @@ async def stream_chat_derender( ) assert resp.status_code == 200, resp.text data = resp.json() + if chat_request.get("return_token_ids") and chat_request.get( + "include_reasoning", True + ): + assert data["chunk"]["choices"][0]["token_ids"] == tids state = data["stream_state"] choices.extend(data["chunk"]["choices"]) if on_chunk is not None: diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index 8e3e4ca25abc..3fe50f16c391 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -160,6 +160,11 @@ def _derender_chat( else {} ) + return_token_ids = ( + chat_request is not None + and chat_request.return_token_ids + and (not has_parser or chat_request.include_reasoning) + ) for choice in generate_response.choices: if not choice.token_ids: raise ValueError(f"choice {choice.index} has empty or null token_ids") @@ -201,6 +206,7 @@ def _derender_chat( message=message, logprobs=resolved_logprobs, finish_reason=choice.finish_reason, + token_ids=choice.token_ids if return_token_ids else None, ) ) @@ -432,6 +438,7 @@ async def derender_chat_stream( stream_choices: list[ChatCompletionResponseStreamChoice] = [] updated_state = state + return_token_ids = chat_request is not None and chat_request.return_token_ids for choice in generate_chunk.choices: delta_tids = choice.token_ids or [] new_text, updated_state = await self._detokenize_delta_async( @@ -477,6 +484,7 @@ async def derender_chat_stream( delta=delta, logprobs=resolved_logprobs, finish_reason=choice.finish_reason, + token_ids=choice.token_ids if return_token_ids else None, ) ) @@ -583,6 +591,10 @@ def _replay(token_ids: list[int], chunk_lens: list[int]) -> None: tools_streamed = state.tools_streamed last_tool_call_ids = list(state.last_tool_call_ids) + return_token_ids = ( + chat_request.return_token_ids and chat_request.include_reasoning + ) + # At most one choice: the caller (derender_chat_stream) already # rejects >1 before dispatching here. role_sent/tools_streamed/ # output_token_ids below are updated for that single choice, not @@ -660,6 +672,7 @@ def _replay(token_ids: list[int], chunk_lens: list[int]) -> None: index=choice.index, delta=delta_message, finish_reason=finish_reason, + token_ids=choice.token_ids if return_token_ids else None, ) stream_choices.append( maybe_filter_parallel_tool_calls(stream_choice, chat_request) From 4163a80142d41b98f94d9b798506448162a22407 Mon Sep 17 00:00:00 2001 From: Shuolei Wang <46160365+ShuoleiWang@users.noreply.github.com> Date: Mon, 12 Oct 2026 15:56:59 +0800 Subject: [PATCH 280/624] [Rust Frontend] Return weight-version spans on final outputs (#61070) Signed-off-by: Shuolei Wang Signed-off-by: Bugen Zhao Co-authored-by: aoshen02 Co-authored-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- .../online_serving/token_in_token_out.md | 2 +- rust/src/engine-core-client/src/client/imp.rs | 4 +- .../src/client/output_channel.rs | 1 + .../engine-core-client/src/client/state.rs | 2 + .../engine-core-client/src/client/stream.rs | 2 + .../engine-core-client/src/protocol/output.rs | 10 +++- .../engine-core-client/src/tests/client.rs | 3 + rust/src/llm/src/lib.rs | 2 +- rust/src/llm/src/output.rs | 55 +++++++++++++++++++ rust/src/llm/tests/generate.rs | 45 ++++++++++++--- .../server/src/routes/inference/generate.rs | 9 +++ .../src/routes/inference/generate/types.rs | 4 ++ 12 files changed, 127 insertions(+), 12 deletions(-) diff --git a/docs/serving/online_serving/token_in_token_out.md b/docs/serving/online_serving/token_in_token_out.md index b3e15ae61692..eade608bc451 100644 --- a/docs/serving/online_serving/token_in_token_out.md +++ b/docs/serving/online_serving/token_in_token_out.md @@ -66,7 +66,7 @@ The final output of every choice carries `weight_versions`, the label of the wei ] ``` -- With `stream: true`, only the last chunk of each choice carries the list, including the last chunk after an abort. Its offsets cover the whole output. Earlier chunks carry `null`. +- With `stream: true`, only the last chunk of each choice carries the list, including the last chunk after an abort. Its offsets cover the whole output. Earlier chunks carry `null` or omit the field. - An output with no tokens, such as a request aborted before its first token, gets `[]`. - The labels are the ones you set with `/update_weight_version` or `/finish_weight_update`, see [Async Reinforcement Learning](../../training/async_rl.md#weight-versions-in-outputs). vLLM doesn't order or check them, and the label is `"default"` until you set one. diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index 0a4af1a5ee5f..fa82d9591aa8 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -453,10 +453,12 @@ pub(crate) async fn run_output_dispatcher_loop( let has_outputs = !batch.outputs.is_empty(); let mut iteration_tokens: u64 = batch.outputs.iter().map(output_iteration_tokens).sum(); - let (engine_index, timestamp) = (batch.engine_index, batch.timestamp); + let (engine_index, timestamp, weight_version) = + (batch.engine_index, batch.timestamp, batch.weight_version); let outputs = batch.outputs.into_iter().map(|output| EngineCoreStreamOutput { engine_index, timestamp, + weight_version: weight_version.clone(), output, }); for output in inner.send_outputs(outputs) { diff --git a/rust/src/engine-core-client/src/client/output_channel.rs b/rust/src/engine-core-client/src/client/output_channel.rs index 7b58d3dc2f1f..3de4da366d54 100644 --- a/rust/src/engine-core-client/src/client/output_channel.rs +++ b/rust/src/engine-core-client/src/client/output_channel.rs @@ -130,6 +130,7 @@ mod tests { EngineCoreStreamOutput { engine_index: 0, timestamp: 0.0, + weight_version: None, output: EngineCoreOutput { request_id: "req".to_string(), new_token_ids: token_ids.to_vec(), diff --git a/rust/src/engine-core-client/src/client/state.rs b/rust/src/engine-core-client/src/client/state.rs index 9c183b87e4e0..9f6f8bbf5c2b 100644 --- a/rust/src/engine-core-client/src/client/state.rs +++ b/rust/src/engine-core-client/src/client/state.rs @@ -344,6 +344,7 @@ impl RequestRegistry { let output = EngineCoreStreamOutput { engine_index: engine_id.engine_index().unwrap_or(0), timestamp, + weight_version: None, output: EngineCoreOutput { request_id: request_id.clone(), finish_reason: Some(EngineCoreFinishReason::Abort), @@ -527,6 +528,7 @@ mod tests { EngineCoreStreamOutput { engine_index: 0, timestamp: 0.0, + weight_version: None, output, } } diff --git a/rust/src/engine-core-client/src/client/stream.rs b/rust/src/engine-core-client/src/client/stream.rs index 6f363f61017d..0f54da98c97d 100644 --- a/rust/src/engine-core-client/src/client/stream.rs +++ b/rust/src/engine-core-client/src/client/stream.rs @@ -3,6 +3,7 @@ use std::ops::Deref; use std::pin::Pin; +use std::sync::Arc; use std::task::{Context, Poll}; use futures::stream::FusedStream; @@ -30,6 +31,7 @@ enum State { pub struct EngineCoreStreamOutput { pub engine_index: u32, pub timestamp: f64, + pub weight_version: Option>, pub output: EngineCoreOutput, } diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index 8945f8b11020..d269f2491163 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::BTreeSet; +use std::sync::Arc; use bytes::Bytes; use enum_as_inner::EnumAsInner; @@ -228,7 +229,8 @@ struct WireEngineCoreOutputs { /// wave needs to start in other engines. #[serde(default)] start_wave: Option, - /// Decoded for Python wire compatibility; not exposed by the Rust frontend. + /// Label of the weights used for this step, as set by the + /// `set_weight_version` utility; `"default"` until set. #[serde(default)] weight_version: Option, } @@ -247,6 +249,9 @@ pub struct RequestBatchOutputs { pub scheduler_stats: Option>, pub timestamp: f64, pub finished_requests: Option>, + /// Label of the weights used for this step. It changes only while + /// generation is paused, so it applies to every output in the batch. + pub weight_version: Option>, } #[derive(Debug, Clone, PartialEq, Default)] @@ -327,6 +332,7 @@ impl TryFrom for EngineCoreOutputs { scheduler_stats: value.scheduler_stats, timestamp: value.timestamp, finished_requests: value.finished_requests, + weight_version: value.weight_version.map(Arc::from), } .into()), (false, Some(_), None, None) => Ok(UtilityCallOutput { @@ -365,6 +371,7 @@ impl From for WireEngineCoreOutputs { scheduler_stats: batch.scheduler_stats, timestamp: batch.timestamp, finished_requests: batch.finished_requests, + weight_version: batch.weight_version.map(|version| version.to_string()), ..Default::default() }, EngineCoreOutputs::Utility(utility) => Self { @@ -558,6 +565,7 @@ mod tests { "req-1", }, ), + weight_version: None, }, ) "#]] diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 0ad58b027154..43eb88db79ed 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -2884,6 +2884,9 @@ fn python_msgpack_fixtures_match_rust_encoding() { "req-1", }, ), + weight_version: Some( + "v7", + ), }, ) "#]] diff --git a/rust/src/llm/src/lib.rs b/rust/src/llm/src/lib.rs index 7e9b5e0ba53f..c94c1d80f743 100644 --- a/rust/src/llm/src/lib.rs +++ b/rust/src/llm/src/lib.rs @@ -16,7 +16,7 @@ mod request_metrics; pub use error::{Error, Result}; pub use output::{ CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStream, - GenerateOutputStreamExt, GeneratePromptInfo, TokenUsage, + GenerateOutputStreamExt, GeneratePromptInfo, TokenUsage, WeightVersionSpan, }; pub use request::GenerateRequest; pub use request_metrics::current_unix_timestamp_secs; diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index b6ccb7717508..81dbb3fda53c 100644 --- a/rust/src/llm/src/output.rs +++ b/rust/src/llm/src/output.rs @@ -24,6 +24,14 @@ use crate::error::Result; use crate::inflight::RequestGuard; use crate::request_metrics::{RequestMetricsTracker, current_unix_timestamp_secs}; +/// Output tokens `[start, end)` were sampled by weights labelled `version`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct WeightVersionSpan { + pub version: Arc, + pub start: usize, + pub end: usize, +} + /// Token usage metadata for one request. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct TokenUsage { @@ -55,6 +63,7 @@ pub struct CollectedGenerateOutput { pub sampling_mask: Option, /// Per-request speculative-decoding metrics from the terminal output. pub spec_decode_metrics: Option, + pub weight_versions: Option>, } /// Prompt-scoped metadata emitted only once on the first [`GenerateOutput`] for @@ -172,6 +181,8 @@ pub struct GenerateOutput { pub sampling_mask: Option, /// Per-request speculative-decoding metrics, present on terminal outputs. pub spec_decode_metrics: Option>, + /// Weight-version spans over all output tokens, set on terminal outputs only. + pub weight_versions: Option>, } impl GenerateOutput { @@ -225,6 +236,7 @@ impl GenerateOutput { self.kv_transfer_params = next.kv_transfer_params; self.ec_transfer_params = next.ec_transfer_params; self.spec_decode_metrics = next.spec_decode_metrics; + self.weight_versions = next.weight_versions; } } @@ -253,6 +265,7 @@ impl GenerateOutput { ec_transfer_params: None, sampling_mask: None, spec_decode_metrics: None, + weight_versions: None, } } } @@ -267,6 +280,9 @@ pub struct GenerateOutputStream { pending_prompt_info: Option>, raw_stream: EngineCoreOutputStream, request_metrics: RequestMetricsTracker, + /// (weight version, first output-token index) for each span. + weight_version_starts: Vec<(Arc, usize)>, + num_output_tokens: usize, /// Removes this request's external→internal tracking edge on drop. Held for /// its `Drop` side effect only; never read directly. _request_guard: RequestGuard, @@ -289,6 +305,8 @@ impl GenerateOutputStream { })), raw_stream, request_metrics, + weight_version_starts: Vec::new(), + num_output_tokens: 0, _request_guard: request_guard, } } @@ -320,6 +338,15 @@ impl GenerateOutputStream { fn convert_output(&mut self, raw: EngineCoreStreamOutput) -> Result { let received_at = current_unix_timestamp_secs(); self.request_metrics.observe_output(raw.timestamp, received_at, &raw.output); + let num_new_tokens = raw.output.new_token_ids.len(); + // Start a span when tokens arrive under a new label. + if num_new_tokens > 0 + && let Some(version) = raw.weight_version + && self.weight_version_starts.last().is_none_or(|(last, _)| *last != version) + { + self.weight_version_starts.push((version, self.num_output_tokens)); + } + self.num_output_tokens += num_new_tokens; let raw = raw.output; @@ -357,6 +384,11 @@ impl GenerateOutputStream { if let Some(finish_reason) = finish_reason.as_ref() { self.request_metrics.record_finished(received_at, finish_reason.clone()); } + let weight_versions = if finish_reason.is_some() { + self.weight_version_spans() + } else { + None + }; let output = GenerateOutput { request_id: raw.request_id, @@ -369,10 +401,31 @@ impl GenerateOutputStream { ec_transfer_params: raw.ec_transfer_params, sampling_mask, spec_decode_metrics: raw.spec_decode_metrics, + weight_versions, }; Ok(output) } + + /// Spans over the output so far; `[]` for no tokens, `None` if no label was reported. + fn weight_version_spans(&self) -> Option> { + let starts = &self.weight_version_starts; + if starts.is_empty() && self.num_output_tokens > 0 { + return None; + } + let ends = starts.iter().skip(1).map(|(_, start)| *start).chain([self.num_output_tokens]); + Some( + starts + .iter() + .zip(ends) + .map(|((version, start), end)| WeightVersionSpan { + version: version.clone(), + start: *start, + end, + }) + .collect(), + ) + } } impl Stream for GenerateOutputStream { @@ -476,6 +529,7 @@ impl> + Send> T { ec_transfer_params: None, sampling_mask, spec_decode_metrics: None, + weight_versions: None, }); } @@ -490,6 +544,7 @@ impl> + Send> T { collected.kv_transfer_params = output.kv_transfer_params.map(|value| *value); collected.ec_transfer_params = output.ec_transfer_params.map(|value| *value); collected.spec_decode_metrics = output.spec_decode_metrics.map(|value| *value); + collected.weight_versions = output.weight_versions; if let Some(mask) = collected.sampling_mask.as_ref() && mask.rows.len() != collected.token_ids.len() { diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index e1205ede9961..8a150d07a44b 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -31,6 +31,7 @@ use vllm_engine_core_client::test_utils::{ use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId}; use vllm_llm::{ Error, FinishReason, GenerateOutputStreamExt as _, GeneratePromptInfo, GenerateRequest, Llm, + WeightVersionSpan, }; use vllm_metrics::METRICS; use zeromq::prelude::{SocketRecv, SocketSend}; @@ -342,15 +343,16 @@ async fn generate_merges_outputs_by_stream_interval() { let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); // One engine step per batch. - let steps: [(&[u32], Option); 6] = [ - (&[1], None), - (&[2], None), - (&[3], None), - (&[4, 5], None), - (&[6], None), - (&[7], Some(EngineCoreFinishReason::Length)), + // The label changes within the merged output [2, 3, 4, 5]. + let steps: [(&[u32], &str, Option); 6] = [ + (&[1], "v7", None), + (&[2], "v7", None), + (&[3], "v8", None), + (&[4, 5], "v8", None), + (&[6], "v7", None), + (&[7], "v7", Some(EngineCoreFinishReason::Length)), ]; - for (token_ids, finish_reason) in steps { + for (token_ids, weight_version, finish_reason) in steps { send_outputs( push, RequestBatchOutputs { @@ -361,6 +363,7 @@ async fn generate_merges_outputs_by_stream_interval() { Some(token_logprobs(token_ids)), None, )], + weight_version: Some(weight_version.into()), ..Default::default() } .into(), @@ -412,6 +415,21 @@ async fn generate_merges_outputs_by_stream_interval() { for output in &outputs { assert_eq!(output.logprobs, Some(token_logprobs(&output.token_ids))); } + let span = |version: &str, start, end| WeightVersionSpan { + version: version.into(), + start, + end, + }; + let weight_versions: Vec<_> = + outputs.iter().map(|output| output.weight_versions.clone()).collect(); + assert_eq!( + weight_versions, + [ + None, + None, + Some(vec![span("v7", 0, 2), span("v8", 2, 5), span("v7", 5, 7)]) + ] + ); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -530,6 +548,7 @@ async fn collect_output_rejects_partial_sampling_mask() { rows: vec![vec![1, 33, 99]], }), spec_decode_metrics: None, + weight_versions: None, }; let error = futures::stream::iter([Ok(output)]).collect_output().await.unwrap_err(); @@ -731,6 +750,7 @@ async fn abort_resolves_external_request_id_to_internal_before_reaching_engine() push, RequestBatchOutputs { outputs: vec![request_output(&request.request_id, vec![7], None)], + weight_version: Some("v7".into()), ..Default::default() } .into(), @@ -764,6 +784,15 @@ async fn abort_resolves_external_request_id_to_internal_before_reaching_engine() let terminal = stream.next().await.unwrap().unwrap(); assert_eq!(terminal.finish_reason, Some(FinishReason::Abort)); assert!(terminal.token_ids.is_empty()); + // The synthesized abort output reports the spans received so far. + assert_eq!( + terminal.weight_versions, + Some(vec![WeightVersionSpan { + version: "v7".into(), + start: 0, + end: 1 + }]) + ); assert!(stream.next().await.is_none()); let _ = shutdown_tx.send(()); diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index 985fe44fb136..a30cd0c85ab4 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -235,6 +235,7 @@ async fn generate_chunk_stream( finish_reason: finish_reason.map(|reason| reason.as_str().to_string()), token_ids, sampling_mask: output.sampling_mask.map(|mask| mask.rows), + weight_versions: output.weight_versions, }], usage: include_continuous_usage .then(|| Usage::from_token_usage(usage, enable_prompt_tokens_details)), @@ -336,6 +337,7 @@ fn collect_generate( finish_reason: Some(finish_reason), token_ids: collected.token_ids, sampling_mask: collected.sampling_mask.map(|mask| mask.rows), + weight_versions: collected.weight_versions, }], prompt_logprobs, prompt_token_id_logprobs: collected @@ -629,6 +631,7 @@ mod tests { ec_transfer_params: None, sampling_mask: None, spec_decode_metrics: None, + weight_versions: None, }), Ok(GenerateOutput { request_id: String::new(), @@ -645,6 +648,7 @@ mod tests { ec_transfer_params: None, sampling_mask: None, spec_decode_metrics: None, + weight_versions: None, }), ]); @@ -731,6 +735,7 @@ mod tests { ec_transfer_params: None, sampling_mask: None, spec_decode_metrics: None, + weight_versions: None, } } @@ -905,6 +910,7 @@ mod tests { rows: vec![vec![30, 40]], }), spec_decode_metrics: None, + weight_versions: None, }; let response = collect_generate( @@ -959,6 +965,7 @@ mod tests { prompt_token_ids: vec![10], sampling_mask: None, spec_decode_metrics: Some(metrics.clone()), + weight_versions: None, }; let response = collect_generate( @@ -1051,6 +1058,7 @@ mod tests { prompt_token_ids: vec![10, 20], sampling_mask: None, spec_decode_metrics: None, + weight_versions: None, }; let response = collect_generate( @@ -1086,6 +1094,7 @@ mod tests { ec_transfer_params: None, sampling_mask: None, spec_decode_metrics: None, + weight_versions: None, prompt_token_ids, }; diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs index 0a37640d63a1..6352b187250a 100644 --- a/rust/src/server/src/routes/inference/generate/types.rs +++ b/rust/src/server/src/routes/inference/generate/types.rs @@ -9,6 +9,7 @@ use serde_json::{Map, Value}; use validator::Validate; use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics; use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; +use vllm_llm::WeightVersionSpan; use vllm_text::SamplingParams; use crate::routes::openai::utils::types::{Normalizable, StreamOptions, Usage}; @@ -73,6 +74,7 @@ pub(super) struct GenerateResponseChoice { pub finish_reason: Option, pub token_ids: Vec, pub sampling_mask: Option>>, + pub weight_versions: Option>, } /// Mirrors the Python vLLM `GenerateResponseStreamChoice` class. @@ -84,6 +86,8 @@ pub(super) struct GenerateResponseStreamChoice { pub finish_reason: Option, pub token_ids: Vec, pub sampling_mask: Option>>, + /// Set on the final chunk only; covers the whole output. + pub weight_versions: Option>, } /// Mirrors the Python vLLM `GenerateStreamResponse` class. From 5bed289707f3b647c903eb6911a2ca6e20ebf74c Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Mon, 12 Oct 2026 16:06:47 +0800 Subject: [PATCH 281/624] [DX] Update CODEOWNERS for entrypoints (#60984) Signed-off-by: DarkLight1337 --- .github/CODEOWNERS | 51 ++++++++++++++++++++++++++++++++++------------ 1 file changed, 38 insertions(+), 13 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index a92ab6a2c963..01dab4e76b52 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -34,21 +34,48 @@ /vllm/utils/argparse_utils.py # Entrypoints -/vllm/entrypoints/anthropic @mgoin @DarkLight1337 +/vllm/entrypoints/anthropic @chaunceyjiang @mgoin @DarkLight1337 +/vllm/entrypoints/cohere @DarkLight1337 /vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 /vllm/entrypoints/mcp @heheda12345 -/vllm/entrypoints/openai @aarnphm @chaunceyjiang @DarkLight1337 -/vllm/entrypoints/speech_to_text/realtime @njhill +/vllm/entrypoints/openai/chat_completion @chaunceyjiang @DarkLight1337 +/vllm/entrypoints/openai/completion @chaunceyjiang @DarkLight1337 +/vllm/entrypoints/openai/models @chaunceyjiang @DarkLight1337 +/vllm/entrypoints/openai/parser @aarnphm @chaunceyjiang @sfeng33 +/vllm/entrypoints/openai/responses @aarnphm @chaunceyjiang @sfeng33 +/vllm/entrypoints/pooling @noooop @DarkLight1337 +/vllm/entrypoints/rl @aoshen02 @njhill +/vllm/entrypoints/scale_out @aoshen02 @noooop @DarkLight1337 /vllm/entrypoints/speech_to_text @NickLucche -/vllm/entrypoints/pooling @noooop -/vllm/entrypoints/serve/sagemaker @DarkLight1337 +/vllm/entrypoints/speech_to_text/realtime @njhill @NickLucche /vllm/entrypoints/serve @njhill +/vllm/entrypoints/serve/sagemaker @DarkLight1337 /vllm/entrypoints/common @DarkLight1337 /vllm/entrypoints/llm.py @DarkLight1337 +/tests/entrypoints @robertgshaw2-redhat @aarnphm @NickLucche @AndreasKaratzas +/tests/entrypoints/anthropic @mgoin @DarkLight1337 +/tests/entrypoints/cohere @DarkLight1337 +/tests/entrypoints/cli @hmellor @mgoin @DarkLight1337 +/tests/entrypoints/openai/chat_completion @chaunceyjiang @DarkLight1337 +/tests/entrypoints/openai/completion @chaunceyjiang @DarkLight1337 +/tests/entrypoints/openai/models @chaunceyjiang @DarkLight1337 +/tests/entrypoints/openai/parser @aarnphm @chaunceyjiang @sfeng33 +/tests/entrypoints/openai/responses @aarnphm @chaunceyjiang @sfeng33 +/tests/entrypoints/pooling @noooop @DarkLight1337 +/tests/entrypoints/rl @aoshen02 @njhill +/tests/entrypoints/scale_out @noooop @DarkLight1337 +/tests/entrypoints/speech_to_text @NickLucche +/tests/entrypoints/speech_to_text/realtime @njhill @NickLucche +/tests/entrypoints/serve @njhill +/tests/entrypoints/serve/sagemaker @DarkLight1337 +/tests/entrypoints/llm @DarkLight1337 +/tests/entrypoints/llm/test_struct_output_generate.py @mgoin @aarnphm @arpera + # Rust Frontend /rust/ @BugenZhao @njhill /rust/src/bench @esmeetu +/rust/src/server/src/routes/inference @aoshen02 @BugenZhao @njhill /tools/build_rust.sh @BugenZhao @njhill /rust-toolchain.toml @BugenZhao @njhill /.buildkite/test_areas/rust* @BugenZhao @njhill @@ -62,6 +89,12 @@ /vllm/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning /vllm/parser @aarnphm @chaunceyjiang @sfeng33 @bbrowning +/tests/tokenizers_ @DarkLight1337 @njhill +/tests/renderers @DarkLight1337 @njhill +/tests/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning +/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning +/tests/tool_use @aarnphm @chaunceyjiang @sfeng33 @bbrowning + # vLLM V1 /vllm/v1/attention @LucasWilkinson @MatthewBonanni /vllm/v1/attention/backend.py @WoosukKwon @zhuohan123 @youkaichao @alexm-redhat @njhill @@ -99,7 +132,6 @@ /tests/distributed/test_multi_node_assignment.py @youkaichao /tests/distributed/test_pipeline_parallel.py @youkaichao /tests/distributed/test_same_node.py @youkaichao -/tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche @AndreasKaratzas /tests/evals @mgoin @vadiklyutiy @AndreasKaratzas /tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256 @zyongye @AndreasKaratzas /tests/kernels/ir @ProExpertProg @tjtanaa @@ -107,7 +139,6 @@ /tests/multimodal @DarkLight1337 @ywang96 @NickLucche @shen-shanshan /tests/multimodal/media @Isotr0py /tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye @AndreasKaratzas -/tests/entrypoints/llm/test_struct_output_generate.py @mgoin @aarnphm @arpera /tests/v1/structured_output @mgoin @aarnphm @arpera /tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium /tests/lora @jeejeelee @@ -117,9 +148,6 @@ /tests/v1/kv_offload @ApostaC @orozery /tests/v1/simple_kv_offload @ivanium /tests/v1/determinism @yewentao256 -/tests/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning -/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning -/tests/tool_use @aarnphm @chaunceyjiang @sfeng33 @bbrowning # Transformers modeling backend /vllm/model_executor/models/transformers @hmellor @@ -221,15 +249,12 @@ mkdocs.yaml @hmellor /examples/pooling @noooop /docs/models/pooling_models @noooop /tests/models/*/pooling* @noooop -/tests/entrypoints/pooling @noooop /vllm/config/pooler.py @noooop /vllm/model_executor/layers/pooler @noooop # Reinforcement learning /vllm/distributed/weight_transfer @aoshen02 -/vllm/entrypoints/rl @njhill @aoshen02 /tests/distributed/test_weight_transfer.py @aoshen02 -/tests/entrypoints/rl @DarkLight1337 @robertgshaw2-redhat @njhill @aarnphm @NickLucche @AndreasKaratzas @aoshen02 /tests/v1/worker/test_gpu_worker_weight_transfer.py @aoshen02 /examples/rl @aoshen02 /docs/training/rlhf.md @aoshen02 From 99b44c7a4282e903edc0c4703f50d3a72f4d419b Mon Sep 17 00:00:00 2001 From: Thuongvu Ho Date: Mon, 12 Oct 2026 01:31:07 -0700 Subject: [PATCH 282/624] [Metrics] Add admission rejection metrics (#55812) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Thuongvu Ho Co-authored-by: Nicolò Lucchesi --- docs/design/metrics.md | 1 + .../serve/instrumentator/test_metrics.py | 1 + tests/v1/engine/test_admission_control.py | 79 ++++++++++++++++++- vllm/v1/engine/async_llm.py | 11 ++- vllm/v1/metrics/loggers.py | 31 ++++++++ 5 files changed, 119 insertions(+), 4 deletions(-) diff --git a/docs/design/metrics.md b/docs/design/metrics.md index 3e070f588890..5d9fa42c11c3 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -30,6 +30,7 @@ In v1, an extensive set of metrics are exposed via a Prometheus-compatible `/met - `vllm:prompt_tokens_total` (Counter) - Total number of prompt tokens processed. - `vllm:generation_tokens_total` (Counter) - Total number of generated tokens. - `vllm:request_success_total` (Counter) - Number of finished requests (by finish reason). +- `vllm:admission_rejections_total` (Counter) - Admission attempts rejected by a queue limit, by reason. - `vllm:request_prompt_tokens` (Histogram) - Histogram of input prompt token counts. - `vllm:request_generation_tokens` (Histogram) - Histogram of generation token counts. - `vllm:time_to_first_token_seconds` (Histogram) - Time to first token (TTFT). diff --git a/tests/entrypoints/serve/instrumentator/test_metrics.py b/tests/entrypoints/serve/instrumentator/test_metrics.py index ac4576fbb95d..a685b9efb374 100644 --- a/tests/entrypoints/serve/instrumentator/test_metrics.py +++ b/tests/entrypoints/serve/instrumentator/test_metrics.py @@ -188,6 +188,7 @@ async def test_metrics_counts( "vllm:prefix_cache_queries", "vllm:prefix_cache_hits", "vllm:num_preemptions_total", + "vllm:admission_rejections_total", "vllm:prompt_tokens_total", "vllm:generation_tokens_total", "vllm:iteration_tokens_total", diff --git a/tests/v1/engine/test_admission_control.py b/tests/v1/engine/test_admission_control.py index 15be166c8c0e..dabd42711ca1 100644 --- a/tests/v1/engine/test_admission_control.py +++ b/tests/v1/engine/test_admission_control.py @@ -6,6 +6,7 @@ These tests cover: - OutputProcessor.get_num_queued_tokens() token counting - AsyncLLM.check_admission() admission control logic +- Admission rejection Prometheus samples - Exception classes (GracefulHTTPError, QueueOverflowError, MaxQueuedTokensError) - create_error_response() mapping GracefulHTTPError to HTTP 503 - SchedulerConfig field defaults and validation @@ -14,6 +15,7 @@ import argparse import asyncio +import logging import multiprocessing from collections.abc import Awaitable, Callable from http import HTTPStatus @@ -21,6 +23,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest +from prometheus_client import REGISTRY from pydantic import ValidationError from vllm.config.scheduler import SchedulerConfig @@ -40,6 +43,10 @@ from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.engine.output_processor import OutputProcessor +from vllm.v1.metrics.loggers import StatLoggerManager + +_ADMISSION_METRIC_MODEL = "admission-metric-test" +_ADMISSION_LOGGER = "vllm.v1.engine.async_llm" pytestmark = pytest.mark.cpu_test @@ -70,10 +77,43 @@ def _make_async_llm( llm.output_processor = MagicMock() llm.output_processor.get_num_unfinished_requests.return_value = num_unfinished llm.output_processor.get_num_queued_tokens.return_value = num_queued_tokens + llm.logger_manager = None llm.admission_stats = None return llm +def _attach_prometheus_logger( + llm: AsyncLLM, engine_indexes: list[int] | None = None +) -> None: + """Attach a real Prometheus logger manager to a lightweight AsyncLLM.""" + vllm_config = MagicMock() + vllm_config.observability_config.show_hidden_metrics = False + vllm_config.observability_config.kv_cache_metrics = False + vllm_config.observability_config.custom_histogram_buckets = None + vllm_config.model_config.served_model_name = _ADMISSION_METRIC_MODEL + vllm_config.model_config.max_model_len = 2048 + vllm_config.model_config.is_diffusion = False + vllm_config.speculative_config = None + vllm_config.kv_transfer_config = None + vllm_config.ec_transfer_config = None + vllm_config.lora_config = None + llm.logger_manager = StatLoggerManager( + vllm_config=vllm_config, + engine_idxs=engine_indexes, + enable_default_loggers=False, + ) + + +def _admission_rejection_count(reason: str) -> float: + """Read the Prometheus sample for one admission rejection reason.""" + value = REGISTRY.get_sample_value( + "vllm:admission_rejections_total", + {"model_name": _ADMISSION_METRIC_MODEL, "reason": reason}, + ) + assert value is not None + return value + + def _make_output_processor(**request_states) -> OutputProcessor: op = OutputProcessor.__new__(OutputProcessor) op.request_states = request_states @@ -157,11 +197,13 @@ def test_shared_request_snapshot_enforces_global_limit(): server_0.set_num_requests(2) llm = _make_async_llm(max_num_queued_reqs=2) llm.admission_stats = server_1 + _attach_prometheus_logger(llm) with pytest.raises(QueueOverflowError): llm.check_admission() assert server_0.get_num_requests() == 2 + assert _admission_rejection_count("max_num_queued_reqs") == 1 # --------------------------------------------------------------------------- @@ -247,19 +289,35 @@ def test_admission_no_limits_allows_everything(): def test_admission_reqs_allows_when_under_limit(): llm = _make_async_llm(max_num_queued_reqs=10, num_unfinished=5) + _attach_prometheus_logger(llm) + llm.check_admission() llm.check_admission() + assert _admission_rejection_count("max_num_queued_reqs") == 0 + assert _admission_rejection_count("max_num_queued_tokens") == 0 -def test_admission_reqs_rejects_at_limit(): +def test_admission_reqs_rejects_at_limit(caplog): llm = _make_async_llm(max_num_queued_reqs=10, num_unfinished=10) - with pytest.raises(QueueOverflowError): + _attach_prometheus_logger(llm, engine_indexes=[0, 1, 2]) + with ( + caplog.at_level(logging.INFO, logger=_ADMISSION_LOGGER), + pytest.raises(QueueOverflowError), + ): llm.check_admission() + assert _admission_rejection_count("max_num_queued_reqs") == 1 + assert _admission_rejection_count("max_num_queued_tokens") == 0 + assert not any( + "Request queue full" in record.getMessage() for record in caplog.records + ) def test_admission_reqs_rejects_with_n(): llm = _make_async_llm(max_num_queued_reqs=10, num_unfinished=8) + _attach_prometheus_logger(llm) with pytest.raises(QueueOverflowError): llm.check_admission(3) + assert _admission_rejection_count("max_num_queued_reqs") == 1 + assert _admission_rejection_count("max_num_queued_tokens") == 0 def test_admission_reqs_allows_n_at_boundary(): @@ -283,8 +341,11 @@ def test_admission_tokens_allows_when_under_limit(): def test_admission_tokens_rejects_at_limit(): llm = _make_async_llm(max_num_queued_tokens=1000, num_queued_tokens=1000) + _attach_prometheus_logger(llm) with pytest.raises(MaxQueuedTokensError): llm.check_admission() + assert _admission_rejection_count("max_num_queued_reqs") == 0 + assert _admission_rejection_count("max_num_queued_tokens") == 1 def test_admission_tokens_rejects_over_limit(): @@ -320,8 +381,11 @@ def test_admission_req_limit_checked_before_token_limit(): num_unfinished=10, num_queued_tokens=1000, ) + _attach_prometheus_logger(llm) with pytest.raises(QueueOverflowError): llm.check_admission() + assert _admission_rejection_count("max_num_queued_reqs") == 1 + assert _admission_rejection_count("max_num_queued_tokens") == 0 # -- n derived from params at the add_request call site --------------------- @@ -381,6 +445,17 @@ async def add_request_async(_request): assert len(request_states) == 1 +def test_admission_preflight_pass_then_reject_records_once(): + llm = _make_async_llm(max_num_queued_reqs=1, num_unfinished=0) + _attach_prometheus_logger(llm) + llm.check_admission() + llm.output_processor.get_num_unfinished_requests.return_value = 1 + with pytest.raises(QueueOverflowError): + llm.check_admission() + assert _admission_rejection_count("max_num_queued_reqs") == 1 + assert _admission_rejection_count("max_num_queued_tokens") == 0 + + @pytest.mark.asyncio @pytest.mark.parametrize("first_n", [1, 3]) async def test_parallel_admission_is_all_or_nothing(first_n: int): diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 60588e076097..f2142af2b3fb 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -348,7 +348,7 @@ def check_admission(self, n: int = 1, request_id: str | None = None) -> None: else self.get_num_unfinished_requests() ) if current_requests + n > max_num_reqs: - logger.info( + logger.debug( "Request queue full - rejecting request %s " "(current=%d, n=%d, max=%d).", request_id, @@ -356,21 +356,28 @@ def check_admission(self, n: int = 1, request_id: str | None = None) -> None: n, max_num_reqs, ) + self._record_admission_rejection("max_num_queued_reqs") raise QueueOverflowError() max_queued_tokens = self.scheduler_config.max_num_queued_tokens if max_queued_tokens is not None: current_tokens = self.get_num_queued_tokens() if current_tokens >= max_queued_tokens: - logger.info( + logger.debug( "Max queued tokens reached - rejecting request %s " "(current_tokens=%d, max=%d).", request_id, current_tokens, max_queued_tokens, ) + self._record_admission_rejection("max_num_queued_tokens") raise MaxQueuedTokensError() + def _record_admission_rejection(self, reason: str) -> None: + """Increment the frontend admission rejection counter if logging is on.""" + if self.logger_manager is not None: + self.logger_manager.record_admission_rejection(reason) + async def get_supported_tasks(self) -> tuple[SupportedTask, ...]: if not hasattr(self, "_supported_tasks"): # Cache the result diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index 7ff1d2993cad..c42d17790926 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -43,6 +43,9 @@ WAITING_REASON_CAPACITY = "capacity" WAITING_REASON_DEFERRED = "deferred" +# User-facing reason labels for frontend admission rejections. +ADMISSION_REASONS = ("max_num_queued_reqs", "max_num_queued_tokens") + PerEngineStatLoggerFactory = Callable[[VllmConfig, int], "StatLoggerBase"] AggregateStatLoggerFactory = type["AggregateStatLoggerBase"] StatLoggerFactory = AggregateStatLoggerFactory | PerEngineStatLoggerFactory @@ -102,6 +105,9 @@ class AggregateStatLoggerBase(StatLoggerBase): @abstractmethod def __init__(self, vllm_config: VllmConfig, engine_indexes: list[int]): ... + def record_admission_rejection(self, reason: str) -> None: # noqa + """Record a frontend admission rejection. Optional for custom loggers.""" + class LoggingStatLogger(StatLoggerBase): def __init__(self, vllm_config: VllmConfig, engine_index: int = 0): @@ -816,6 +822,23 @@ def __init__( "request_num_preemptions", overrides=custom_buckets ) + # + # Frontend admission + # + counter_admission_rejections = self._counter_cls( + name="vllm:admission_rejections", + documentation=( + "Number of admission attempts rejected by a queue limit, by reason." + ), + labelnames=["model_name", "reason"], + ) + self.admission_rejections = { + reason: counter_admission_rejections.labels( + model_name=model_name, reason=reason + ) + for reason in ADMISSION_REASONS + } + # # Histograms of counts # @@ -1276,6 +1299,9 @@ def record( finished_request.max_tokens_param ) + def record_admission_rejection(self, reason: str) -> None: + self.admission_rejections[reason].inc() + def record_sleep_state(self, sleep: int = 0, level: int = 0): awake = 1 discard_all = 0 @@ -1383,6 +1409,11 @@ def record_sleep_state(self, sleep: int = 0, level: int = 0): for logger in self.stat_loggers: logger.record_sleep_state(sleep, level) + def record_admission_rejection(self, reason: str) -> None: + """Record one frontend admission rejection on aggregate loggers.""" + for stat_logger in self.stat_loggers: + stat_logger.record_admission_rejection(reason) + def log(self): for logger in self.stat_loggers: logger.log() From e05201aa2e5ccf19284e5062064050a06d50e281 Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Mon, 12 Oct 2026 16:59:38 +0800 Subject: [PATCH 283/624] [XPU] refine memory release for sleep mode (#61158) Signed-off-by: Yan Ma --- vllm/device_allocator/xpumem.py | 26 ++++++++----------- .../model_loader/sharded_state_loader.py | 7 +---- 2 files changed, 12 insertions(+), 21 deletions(-) diff --git a/vllm/device_allocator/xpumem.py b/vllm/device_allocator/xpumem.py index 7d4b40a78e76..5eb836008d68 100644 --- a/vllm/device_allocator/xpumem.py +++ b/vllm/device_allocator/xpumem.py @@ -146,6 +146,7 @@ def __init__(self): self.pointer_to_data: dict[int, AllocationData] = {} self.current_tag: str = XpuMemAllocator.default_tag self.allocator_and_pools: dict[str, list[Any]] = {} + self._pending_allocators: list[Any] = [] self.python_malloc_callback = self._python_malloc_callback self.python_free_callback = self._python_free_callback @@ -291,24 +292,17 @@ def wake_up(self, tags: list[str] | None = None) -> None: data.cpu_backup_tensor = None def release_pools(self) -> None: - """Drop Python references to MemPool/pluggable allocators eagerly. - - This prevents pool destruction from being deferred to interpreter - finalization, which can happen after parts of XPU runtime are already - torn down. - """ - if not self.allocator_and_pools: + """Release pools, retaining allocators until their allocations are freed.""" + if not self.allocator_and_pools and not self._pending_allocators: return - # Note: keep allocators alive while MemPool objects are destroyed. - # MemPool teardown may invoke allocator virtual methods (e.g. raw_delete) - # when releasing cached blocks. If allocator wrappers are dropped first, - # C++ can hit "pure virtual method called" during shutdown. pool_entries = [e for es in self.allocator_and_pools.values() for e in es] self.allocator_and_pools.clear() mem_pools = [entry[0] for entry in pool_entries] - allocators = [entry[1] for entry in pool_entries] + # Live tensors can outlast their MemPool wrappers. Keep allocators alive + # across calls so deferred frees cannot invoke a destroyed allocator. + self._pending_allocators.extend(entry[1] for entry in pool_entries) pool_entries.clear() xpu_sync = getattr(torch.xpu, "synchronize", None) @@ -318,12 +312,14 @@ def release_pools(self) -> None: except Exception: logger.debug("torch.xpu.synchronize() failed during release_pools") - # Phase 1: drop MemPool refs while allocators are still strongly held. mem_pools.clear() gc.collect() + torch.xpu.empty_cache() - # Phase 2: now it is safe to release allocator wrappers. - allocators.clear() + # Free callbacks remove allocation records, including cached blocks. + # If any remain, a later release_pools() call must retry this cleanup. + if not self.pointer_to_data: + self._pending_allocators.clear() @contextmanager def use_memory_pool(self, tag: str | None = None): diff --git a/vllm/model_executor/model_loader/sharded_state_loader.py b/vllm/model_executor/model_loader/sharded_state_loader.py index 97a31e8e879a..9e8e6e9cfc75 100644 --- a/vllm/model_executor/model_loader/sharded_state_loader.py +++ b/vllm/model_executor/model_loader/sharded_state_loader.py @@ -211,12 +211,7 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: # TODO: support un-sharded checkpoints too raise ValueError( f"Could not find checkpoint files " - f"'{ - os.path.join( - local_model_path, - rendered_pattern, - ) - }', " + f"'{os.path.join(local_model_path, rendered_pattern)}', " f"only pre-sharded checkpoints are currently supported!" ) state_dict = self._filter_subtensors(model.state_dict()) From a33d46b47193e3273dc38f980abbf55459f3c77f Mon Sep 17 00:00:00 2001 From: zhangchengzhucufe-dev Date: Mon, 12 Oct 2026 17:10:55 +0800 Subject: [PATCH 284/624] [Bugfix][XPU] Forward softcap and alibi_slopes in xpu_ops.flash_attn_varlen_func (#58936) Signed-off-by: zhangchengzhucufe-dev Signed-off-by: Kunshang Ji Co-authored-by: Kunshang Ji --- vllm/_xpu_ops.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py index 1f00bc681d11..a384ac721572 100644 --- a/vllm/_xpu_ops.py +++ b/vllm/_xpu_ops.py @@ -1087,8 +1087,8 @@ def flash_attn_varlen_func( block_table=block_table, s_aux=s_aux, window_size=real_window_size, - # alibi_slopes = alibi_slopes, - # softcap=softcap, + alibi_slopes=alibi_slopes, + softcap=softcap, return_softmax_lse=return_softmax_lse, q_descale=q_descale, k_descale=k_descale, From 83773ecdf3ce6da9a2f2bd14cb87e6bbf58d4229 Mon Sep 17 00:00:00 2001 From: Shijin Zhang <75300765+Dovis01@users.noreply.github.com> Date: Mon, 12 Oct 2026 17:22:55 +0800 Subject: [PATCH 285/624] [Frontend] Support CustomToolCall in the non-Harmony Responses API (#57322) Signed-off-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com> Co-authored-by: Chauncey --- .../openai/responses/test_protocol.py | 35 +++++ .../openai/responses/test_responses_utils.py | 88 ++++++++++++ .../openai/responses/test_streaming_events.py | 116 ++++++++++++++- tests/tool_parsers/test_gemma4_tool_parser.py | 16 +++ .../test_k2_horizon_tool_parser.py | 12 ++ .../test_kimi_k3_named_tool_choice.py | 24 ++++ .../test_poolside_v1_tool_parser.py | 20 ++- .../test_structural_tag_registry.py | 15 +- tests/tool_use/test_muse_glimmer.py | 13 ++ tests/tool_use/test_tool_choice_required.py | 35 +++++ vllm/entrypoints/openai/responses/protocol.py | 38 ++++- .../openai/responses/streaming_events.py | 136 ++++++++++++++---- vllm/entrypoints/openai/responses/utils.py | 94 ++++++++++-- vllm/parser/abstract_parser.py | 20 ++- vllm/parser/response_template.py | 5 +- .../tool_parsers/gemma4_engine_tool_parser.py | 9 +- vllm/tool_parsers/k2_horizon_tool_parser.py | 4 +- vllm/tool_parsers/kimi_k3_tool_parser.py | 8 +- vllm/tool_parsers/muse_glimmer_tool_parser.py | 9 +- vllm/tool_parsers/poolside_v1_tool_parser.py | 9 +- vllm/tool_parsers/structural_tag_registry.py | 15 +- vllm/tool_parsers/utils.py | 49 ++++++- 22 files changed, 698 insertions(+), 72 deletions(-) diff --git a/tests/entrypoints/openai/responses/test_protocol.py b/tests/entrypoints/openai/responses/test_protocol.py index db5d7d692490..04b07992528b 100644 --- a/tests/entrypoints/openai/responses/test_protocol.py +++ b/tests/entrypoints/openai/responses/test_protocol.py @@ -1,13 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest from openai_harmony import ( Message, ) +from pydantic import ValidationError from vllm.entrypoints.openai.responses.protocol import ( + ResponsesRequest, serialize_message, serialize_messages, ) +from vllm.exceptions import VLLMValidationError def test_serialize_message() -> None: @@ -37,3 +41,34 @@ def test_serialize_messages() -> None: } msg = Message.from_dict(msg_value) assert serialize_messages([msg, dict_value]) == [msg_value, dict_value] + + +def test_custom_tool_grammar_format_rejected() -> None: + grammar = {"type": "grammar", "syntax": "lark", "definition": 'start: "pwd"'} + with pytest.raises(VLLMValidationError) as exc_info: + ResponsesRequest( + input="hi", tools=[{"type": "custom", "name": "emit", "format": grammar}] + ) + assert exc_info.value.parameter == "tools" + + request = ResponsesRequest( + input="hi", + tools=[{"type": "custom", "name": "emit", "format": {"type": "text"}}], + ) + assert request.tools is not None and request.tools[0].type == "custom" + + +@pytest.mark.parametrize( + "tools", + [ + "abc", + {"type": "custom"}, + [None], + [{"type": "custom", "name": "x", "format": "text"}], + ], +) +def test_malformed_tools_rejected_by_field_validation(tools) -> None: + """Malformed tools bodies fail pydantic field validation (4xx), not an + AttributeError inside the custom-format validator (500).""" + with pytest.raises(ValidationError): + ResponsesRequest(input="hi", tools=tools) diff --git a/tests/entrypoints/openai/responses/test_responses_utils.py b/tests/entrypoints/openai/responses/test_responses_utils.py index 967ddacda6de..6789ddc08bad 100644 --- a/tests/entrypoints/openai/responses/test_responses_utils.py +++ b/tests/entrypoints/openai/responses/test_responses_utils.py @@ -4,6 +4,7 @@ from unittest.mock import patch import pytest +from openai.types.responses import CustomTool, FunctionTool, ResponseCustomToolCall from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall from openai.types.responses.response_function_tool_call_output_item import ( ResponseFunctionToolCallOutputItem, @@ -16,10 +17,14 @@ Summary, ) +from vllm.entrypoints.generate.base.protocol import FunctionCall from vllm.entrypoints.openai.responses.utils import ( _construct_message_from_response_item, + build_response_output_items, construct_chat_messages_with_tool_call, construct_input_messages, + decode_custom_tool_input, + decode_custom_tool_input_prefix, should_continue_final_message, ) from vllm.exceptions import VLLMValidationError @@ -248,6 +253,89 @@ def test_construct_chat_messages_preserves_single_item_conversions(self): assert formatted_item["content"] == "dongyi" +class TestCustomTools: + def test_custom_tool_call_output_item(self): + tools = [ + CustomTool(type="custom", name="emit_command", format={"type": "text"}), + FunctionTool( + type="function", name="get_weather", parameters={}, strict=None + ), + ] + custom, function = build_response_output_items( + reasoning=None, + content=None, + tool_calls=[ + FunctionCall( + id="call_1", name="emit_command", arguments='{"input": "pwd"}' + ), + FunctionCall( + id="call_2", name="get_weather", arguments='{"city": "Paris"}' + ), + ], + tools=tools, + ) + assert isinstance(custom, ResponseCustomToolCall) + assert (custom.call_id, custom.input) == ("call_1", "pwd") + assert custom.id.startswith("ctc_") + assert isinstance(function, ResponseFunctionToolCall) + assert function.arguments == '{"city": "Paris"}' + + def test_custom_tool_history_replays_through_function_shim(self): + messages = construct_chat_messages_with_tool_call( + [ + make_reasoning_item(content_text="Run pwd."), + ResponseCustomToolCall( + type="custom_tool_call", + id="ctc_1", + call_id="call_exec", + name="emit_command", + input='echo "hi"', + ), + { + "type": "custom_tool_call_output", + "call_id": "call_exec", + "output": "hi", + }, + {"role": "user", "content": "Continue"}, + ] + ) + assert messages == [ + { + "role": "assistant", + "reasoning": "Run pwd.", + "tool_calls": [ + { + "id": "call_exec", + "type": "function", + "function": { + "name": "emit_command", + "arguments": '{"input": "echo \\"hi\\""}', + }, + } + ], + }, + {"role": "tool", "content": "hi", "tool_call_id": "call_exec"}, + {"role": "user", "content": "Continue"}, + ] + + def test_decode_custom_tool_input(self): + assert decode_custom_tool_input('{"input": "pwd"}') == "pwd" + assert decode_custom_tool_input('{"cmd": "pwd"}') == '{"cmd": "pwd"}' + assert decode_custom_tool_input('{"file": "a.py"}') == '{"file": "a.py"}' + assert decode_custom_tool_input('{"input": 7}') == '{"input": 7}' + assert decode_custom_tool_input("pwd") == "pwd" + + def test_decode_custom_tool_input_prefix(self): + assert decode_custom_tool_input_prefix('{"inp') == "" + assert decode_custom_tool_input_prefix('{"input": "ls \\"dir') == 'ls "dir' + assert decode_custom_tool_input_prefix('{"input": "a\\') == "a" + assert decode_custom_tool_input_prefix('{"input": "\\ud83d') == "" + assert ( + decode_custom_tool_input_prefix('{"input": "\\ud83d\\ude00!"}') + == "\U0001f600!" + ) + + class TestReasoningItemContentPriority: """Tests that content is prioritized over summary for reasoning items.""" diff --git a/tests/entrypoints/openai/responses/test_streaming_events.py b/tests/entrypoints/openai/responses/test_streaming_events.py index 21fe5537bb74..a37109f85134 100644 --- a/tests/entrypoints/openai/responses/test_streaming_events.py +++ b/tests/entrypoints/openai/responses/test_streaming_events.py @@ -2,7 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest -from openai.types.responses import ResponseFunctionWebSearch, ResponseStreamEvent +from openai.types.responses import ( + CustomTool, + ResponseCustomToolCall, + ResponseFunctionWebSearch, + ResponseStreamEvent, +) from openai_harmony import Message, Role from pydantic import TypeAdapter @@ -20,6 +25,7 @@ emit_reasoning_done_events, split_delta, ) +from vllm.entrypoints.openai.responses.utils import decode_custom_tool_input def test_browser_find_uses_responses_action_type(): @@ -190,3 +196,111 @@ def test_reasoning_to_content_transition(self): types = [e.type for e in events] assert "response.reasoning_text.delta" in types assert "response.output_text.delta" in types + + +class TestSimpleToolCallShapes: + def test_arguments_done_emitted_without_deltas(self): + processor = SimpleStreamingEventProcessor() + events = _run_through_processor( + processor, DeltaMessage(tool_calls=[_make_tool_call(0, name="noop")]) + ) + events += processor.close_current() + assert [e.type for e in events] == [ + "response.output_item.added", + "response.function_call_arguments.done", + "response.output_item.done", + ] + + +class TestCustomToolStreaming: + def test_custom_tool_input_events(self): + processor = SimpleStreamingEventProcessor( + tools=[CustomTool(type="custom", name="emit_command")] + ) + events = _run_through_processor( + processor, + DeltaMessage( + tool_calls=[ + _make_tool_call(0, name="emit_command", arguments='{"input": "ls ') + ] + ), + ) + events += _run_through_processor( + processor, + DeltaMessage(tool_calls=[_make_tool_call(0, arguments='-la \\"x\\""}')]), + ) + events += processor.close_current() + + assert [e.type for e in events] == [ + "response.output_item.added", + "response.custom_tool_call_input.delta", + "response.custom_tool_call_input.delta", + "response.custom_tool_call_input.done", + "response.output_item.done", + ] + done = events[-1].item + assert isinstance(done, ResponseCustomToolCall) + assert done.input == 'ls -la "x"' + assert "".join(e.delta for e in events[1:3]) == done.input + assert events[3].input == done.input + assert events[0].item.type == "custom_tool_call" + assert events[0].item.id == done.id + assert done.id.startswith("ctc_") + assert events[0].item.call_id == done.call_id + + def test_unrelated_keys_are_not_streamed_as_payload(self): + """Review repro: chunks of an object whose input key arrives last must + not leak fragments of other keys into the payload.""" + processor = SimpleStreamingEventProcessor( + tools=[CustomTool(type="custom", name="emit_command")] + ) + events = _run_through_processor( + processor, + DeltaMessage( + tool_calls=[ + _make_tool_call(0, name="emit_command", arguments='{"a": "xy"') + ] + ), + ) + events += _run_through_processor( + processor, + DeltaMessage( + tool_calls=[_make_tool_call(0, arguments=', "input": "hello"}')] + ), + ) + events += processor.close_current() + + deltas = [ + e for e in events if e.type == "response.custom_tool_call_input.delta" + ] + done_input = next( + e for e in events if e.type == "response.custom_tool_call_input.done" + ).input + item = events[-1].item + assert "".join(e.delta for e in deltas) == "hello" # "xy" never leaks + assert done_input == "hello" + assert item.input == "hello" + assert decode_custom_tool_input('{"a": "xy", "input": "hello"}') == "hello" + + def test_done_reported_authoritative_payload_on_mismatch(self): + """If the final decode is not an extension of what was streamed, the + done event and item carry the authoritative value and a correcting + delta is emitted rather than dropping the divergence.""" + processor = SimpleStreamingEventProcessor( + tools=[CustomTool(type="custom", name="emit_command")] + ) + processor.state.current_state = _StateType.TOOL_CALL + processor.state.tool_call_is_custom = True + processor.state.tool_call_payload = "xy" # stale fragment already streamed + processor.state.accumulated_text = '{"a": "xy", "input": "hello"}' + + events = processor.close_current() + delta = next( + e for e in events if e.type == "response.custom_tool_call_input.delta" + ) + done_input = next( + e for e in events if e.type == "response.custom_tool_call_input.done" + ).input + assert delta.delta == "hello" # full authoritative resend, not a dropped suffix + assert done_input == "hello" + assert events[-1].item.input == "hello" diff --git a/tests/tool_parsers/test_gemma4_tool_parser.py b/tests/tool_parsers/test_gemma4_tool_parser.py index 97ed0cd69d4c..43fac974052c 100644 --- a/tests/tool_parsers/test_gemma4_tool_parser.py +++ b/tests/tool_parsers/test_gemma4_tool_parser.py @@ -8,6 +8,7 @@ import pytest from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.parser.gemma4 import ( TOOL_CALL_END, TOOL_CALL_START, @@ -117,6 +118,21 @@ def mock_request(): return request +def test_named_custom_tool_choice_skips_structured_outputs(mock_tokenizer): + # Responses custom tools ride the function shim, so a named custom choice + # must route like a named function choice, not into guided JSON. + tools = [{"type": "custom", "name": "emit_command"}] + request = ResponsesRequest( + input="hi", + tools=tools, + tool_choice={"type": "custom", "name": "emit_command"}, + ) + Gemma4EngineToolParser(mock_tokenizer, tools=tools).adjust_request(request) + + assert request.skip_special_tokens is False + assert request.structured_outputs is None + + # --------------------------------------------------------------------------- # Unit tests for _parse_gemma4_args (shared parser logic) # --------------------------------------------------------------------------- diff --git a/tests/tool_parsers/test_k2_horizon_tool_parser.py b/tests/tool_parsers/test_k2_horizon_tool_parser.py index 6fd4844aeb34..b57b50553e7e 100644 --- a/tests/tool_parsers/test_k2_horizon_tool_parser.py +++ b/tests/tool_parsers/test_k2_horizon_tool_parser.py @@ -8,6 +8,7 @@ from vllm.entrypoints.generate.base.protocol import DeltaMessage from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.parser.parser_manager import ParserManager from vllm.tool_parsers import ToolParserManager from vllm.tool_parsers.k2_horizon_tool_parser import K2HorizonToolParser @@ -316,3 +317,14 @@ def test_whitespace_only_surrounding_content_is_not_preserved(tokenizer): def test_tool_parser_registered(): assert ToolParserManager.get_tool_parser("k2_horizon") is K2HorizonToolParser assert K2HorizonToolParser.supports_required_and_named is False + + +def test_named_tool_resolves_custom_choice(): + # Custom tools ride the function shim; a ToolChoiceCustom must resolve to + # its name like ToolChoiceFunction does. + request = ResponsesRequest( + input="hi", + tools=[{"type": "custom", "name": "emit_command"}], + tool_choice={"type": "custom", "name": "emit_command"}, + ) + assert K2HorizonToolParser._named_tool(request) == "emit_command" diff --git a/tests/tool_parsers/test_kimi_k3_named_tool_choice.py b/tests/tool_parsers/test_kimi_k3_named_tool_choice.py index 36b04b67b195..d426f623da34 100644 --- a/tests/tool_parsers/test_kimi_k3_named_tool_choice.py +++ b/tests/tool_parsers/test_kimi_k3_named_tool_choice.py @@ -8,6 +8,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.exceptions import VLLMValidationError from vllm.sampling_params import StructuredOutputsParams @@ -60,3 +61,26 @@ def test_named_choice_allowed_with_structural_tag(): def test_named_choice_rejected_without_structural_tag(): with pytest.raises(VLLMValidationError): _parser().adjust_request(_request(with_tag=False)) + + +def _custom_request(with_tag: bool) -> ResponsesRequest: + req = ResponsesRequest( + input="hi", + tools=[{"type": "custom", "name": "emit_command"}], + tool_choice={"type": "custom", "name": "emit_command"}, + ) + if with_tag: + req.structured_outputs = StructuredOutputsParams( + structural_tag='{"type": "structural_tag", "format": {}}' + ) + return req + + +def test_named_custom_choice_allowed_with_structural_tag(): + req = _parser().adjust_request(_custom_request(with_tag=True)) + assert req.skip_special_tokens is False + + +def test_named_custom_choice_rejected_without_structural_tag(): + with pytest.raises(VLLMValidationError): + _parser().adjust_request(_custom_request(with_tag=False)) diff --git a/tests/tool_parsers/test_poolside_v1_tool_parser.py b/tests/tool_parsers/test_poolside_v1_tool_parser.py index de6e751e5c91..635081b43b37 100644 --- a/tests/tool_parsers/test_poolside_v1_tool_parser.py +++ b/tests/tool_parsers/test_poolside_v1_tool_parser.py @@ -78,12 +78,15 @@ def _build_chat_request(*, tool_choice: str | dict[str, Any]) -> ChatCompletionR def _build_responses_request( - *, tool_choice: str | dict[str, Any], include: list[str] | None = None + *, + tool_choice: str | dict[str, Any], + include: list[str] | None = None, + tools: list[dict[str, Any]] | None = None, ) -> ResponsesRequest: return ResponsesRequest( model="poolside-test", input=[{"role": "user", "content": "write the file"}], - tools=[_responses_write_file_tool()], + tools=tools if tools is not None else [_responses_write_file_tool()], tool_choice=tool_choice, stream=True, max_output_tokens=200, @@ -145,6 +148,19 @@ def test_named_skips_structured_outputs_responses() -> None: assert request.skip_special_tokens is False +def test_named_custom_tool_choice_skips_structured_outputs_responses() -> None: + # A named custom choice parses to ToolChoiceCustom; custom tools ride the + # function shim, so it must route like ToolChoiceFunction. + request = _build_responses_request( + tools=[{"type": "custom", "name": "emit_command"}], + tool_choice={"type": "custom", "name": "emit_command"}, + ) + PoolsideV1ToolParser(_StubTokenizer()).adjust_request(request) + + assert request.text is None + assert request.skip_special_tokens is False + + def test_auto_still_keeps_special_tokens() -> None: request = _build_chat_request(tool_choice="auto") _make_parser(request).adjust_request(request) diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py index 817c48a4401d..52d0d824691a 100644 --- a/tests/tool_parsers/test_structural_tag_registry.py +++ b/tests/tool_parsers/test_structural_tag_registry.py @@ -7,7 +7,7 @@ from unittest.mock import MagicMock import pytest -from openai.types.responses import ToolChoiceFunction +from openai.types.responses import CustomTool, ToolChoiceFunction from openai.types.responses.tool_choice_allowed import ToolChoiceAllowed from xgrammar import Grammar, StructuralTag from xgrammar import get_model_structural_tag as get_xgrammar_builtin_structural_tag @@ -787,6 +787,19 @@ def test_get_model_structural_tag_supports_named_tool_choice( assert isinstance(tag, StructuralTag) +def test_glm_4_7_forced_custom_tool_uses_function_shim(): + tools = [CustomTool(type="custom", name="emit_command", description="Emit.")] + tag = get_model_structural_tag( + model="glm_4_7", + tools=tools, + tool_choice="required", + reasoning=False, + ) + + assert isinstance(tag, StructuralTag) + assert "emit_command" in tag.model_dump_json() + + @pytest.mark.parametrize( ("parser_cls", "model"), [ diff --git a/tests/tool_use/test_muse_glimmer.py b/tests/tool_use/test_muse_glimmer.py index 5badd9b03162..2fbf3cff49c0 100644 --- a/tests/tool_use/test_muse_glimmer.py +++ b/tests/tool_use/test_muse_glimmer.py @@ -29,6 +29,7 @@ import pytest from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.reasoning.muse_glimmer_reasoning_parser import MuseGlimmerReasoningParser from vllm.tool_parsers.muse_glimmer_tool_parser import MuseGlimmerToolParser @@ -400,3 +401,15 @@ def test_exact_match_kept(): T, _call("get_weather"), _req("get_weather") ) assert out.tool_calls[0].function.name == "get_weather" + + +def test_named_custom_tool_choice_skips_structured_outputs(): + # Custom tools ride the function shim; a named custom choice must skip + # the base adjust_request like a named function choice does. + request = ResponsesRequest( + input="hi", + tools=[{"type": "custom", "name": "emit_command"}], + tool_choice={"type": "custom", "name": "emit_command"}, + ) + T.adjust_request(request) + assert request.structured_outputs is None diff --git a/tests/tool_use/test_tool_choice_required.py b/tests/tool_use/test_tool_choice_required.py index be58ffa83aa9..aedada111aa5 100644 --- a/tests/tool_use/test_tool_choice_required.py +++ b/tests/tool_use/test_tool_choice_required.py @@ -6,7 +6,9 @@ import pytest import regex as re from openai.types.responses import ( + CustomTool, FunctionTool, + ToolChoiceCustom, ToolChoiceFunction, WebSearchTool, ) @@ -21,8 +23,10 @@ from vllm.exceptions import VLLMValidationError from vllm.tool_parsers.streaming import extract_required_tool_call_streaming from vllm.tool_parsers.utils import ( + custom_tool_parameters, find_tool_properties, get_json_schema_from_tools, + iter_response_function_tool_dicts, ) pytestmark = pytest.mark.cpu_test @@ -483,6 +487,37 @@ def test_get_json_schema_rejects_only_non_function_tools(self): get_json_schema_from_tools(tools=[WEB_SEARCH_TOOL], tool_choice="required") +class TestCustomToolShim: + """``custom`` tools are shown to the model as a single-string function.""" + + def test_rendered_as_single_string_function(self): + tool = CustomTool( + type="custom", name="emit_command", description="Emit a command." + ) + [tool_dict] = iter_response_function_tool_dicts([tool]) + assert tool_dict["type"] == "function" + assert tool_dict["name"] == "emit_command" + assert tool_dict["description"] == "Emit a command." + assert tool_dict["parameters"] == custom_tool_parameters() + assert tool_dict["parameters"]["required"] == ["input"] + + def test_forced_tool_choice_sees_custom_tool(self): + tools = [CustomTool(type="custom", name="emit_command"), FUNCTION_TOOL] + schema = get_json_schema_from_tools(tools=tools, tool_choice="required") + assert isinstance(schema, dict) + names = [e["properties"]["name"]["enum"] for e in schema["items"]["anyOf"]] + assert names == [["emit_command"], ["get_weather"]] + named = get_json_schema_from_tools( + tools=tools, + tool_choice=ToolChoiceCustom(type="custom", name="emit_command"), + ) + assert named == custom_tool_parameters() + assert ( + find_tool_properties(tools, "emit_command") + == custom_tool_parameters()["properties"] + ) + + class TestMalformedToolSchemaDefs: """Malformed `$defs` in caller-supplied tool parameters is a 400, not a 500.""" diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 373cb35df057..200054a4ef59 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -14,6 +14,8 @@ ResponseCodeInterpreterCallInterpretingEvent, ResponseContentPartAddedEvent, ResponseContentPartDoneEvent, + ResponseCustomToolCallInputDeltaEvent, + ResponseCustomToolCallInputDoneEvent, ResponseFunctionToolCall, ResponseInputItemParam, ResponseMcpCallArgumentsDeltaEvent, @@ -699,6 +701,34 @@ def input_item_parsing(cls, data): data["input"] = processed_input return data + @model_validator(mode="before") + @classmethod + def check_custom_tool_format(cls, data): + if not isinstance(data, dict): + return data + tools = data.get("tools") + if not isinstance(tools, list): + return data + for tool in tools: + if isinstance(tool, dict): + tool_dict = tool + elif hasattr(tool, "model_dump"): + tool_dict = tool.model_dump() + else: + continue + tool_format = tool_dict.get("format") + if ( + tool_dict.get("type") == "custom" + and isinstance(tool_format, dict) + and tool_format.get("type") == "grammar" + ): + raise VLLMValidationError( + "Custom tools with a grammar format are not supported; " + "only free-form text custom tools are accepted.", + parameter="tools", + ) + return data + @model_validator(mode="before") @classmethod def check_tool_usage(cls, data): @@ -708,9 +738,9 @@ def check_tool_usage(cls, data): tools = data.get("tools") tool_choice = data.get("tool_choice", "auto") has_tools = tools is not None and len(tools) > 0 - is_named_tool_choice = ( - isinstance(tool_choice, dict) and tool_choice.get("type") == "function" - ) + is_named_tool_choice = isinstance(tool_choice, dict) and tool_choice.get( + "type" + ) in ("function", "custom") if not has_tools: if tool_choice in ("auto", "none"): @@ -913,6 +943,8 @@ class ResponseInProgressEvent(OpenAIResponseInProgressEvent): | ResponseOutputItemDoneEvent | ResponseContentPartAddedEvent | ResponseContentPartDoneEvent + | ResponseCustomToolCallInputDeltaEvent + | ResponseCustomToolCallInputDoneEvent | ResponseReasoningTextDeltaEvent | ResponseReasoningTextDoneEvent | ResponseCodeInterpreterCallInProgressEvent diff --git a/vllm/entrypoints/openai/responses/streaming_events.py b/vllm/entrypoints/openai/responses/streaming_events.py index b101040a57a2..49107661a79c 100644 --- a/vllm/entrypoints/openai/responses/streaming_events.py +++ b/vllm/entrypoints/openai/responses/streaming_events.py @@ -29,6 +29,9 @@ ResponseCodeInterpreterToolCallParam, ResponseContentPartAddedEvent, ResponseContentPartDoneEvent, + ResponseCustomToolCall, + ResponseCustomToolCallInputDeltaEvent, + ResponseCustomToolCallInputDoneEvent, ResponseFunctionCallArgumentsDeltaEvent, ResponseFunctionCallArgumentsDoneEvent, ResponseFunctionToolCall, @@ -73,6 +76,9 @@ ) from vllm.entrypoints.openai.responses.utils import ( build_responses_tool_call_name_map, + custom_tool_names, + decode_custom_tool_input, + decode_custom_tool_input_prefix, resolve_responses_tool_call_name, ) from vllm.outputs import CompletionOutput @@ -821,6 +827,8 @@ class SimpleStreamingState: tool_call_name: str = "" tool_call_namespace: str | None = None tool_call_index: int | None = None + tool_call_is_custom: bool = False + tool_call_payload: str = "" current_state: _StateType = field(default_factory=lambda: _StateType.NONE) @@ -1039,38 +1047,80 @@ def emit_simple_tool_call_open( name: str, index: int | None, namespace: str | None = None, + is_custom: bool = False, call_id: str | None = None, ) -> list[StreamingResponsesResponse]: state.current_state = _StateType.TOOL_CALL - state.current_item_id = f"fc_{random_uuid()}" + state.current_item_id = ( + f"ctc_{random_uuid()}" if is_custom else f"fc_{random_uuid()}" + ) state.tool_call_id = call_id or make_tool_call_id() state.tool_call_name = name state.tool_call_namespace = namespace state.tool_call_index = index + state.tool_call_is_custom = is_custom + state.tool_call_payload = "" state.accumulated_text = "" + item: ResponseOutputItem + if is_custom: + item = ResponseCustomToolCall( + type="custom_tool_call", + id=state.current_item_id, + call_id=state.tool_call_id, + name=name, + input="", + ) + else: + item = ResponseFunctionToolCallItem( + type="function_call", + id=state.current_item_id, + call_id=state.tool_call_id, + name=name, + namespace=namespace, + arguments="", + status="in_progress", + ) return [ ResponseOutputItemAddedEvent( type="response.output_item.added", sequence_number=-1, output_index=state.output_index, - item=ResponseFunctionToolCallItem( - type="function_call", - id=state.current_item_id, - call_id=state.tool_call_id, - name=name, - namespace=namespace, - arguments="", - status="in_progress", - ), + item=item, ), ] +def _emit_custom_tool_input_delta( + state: SimpleStreamingState, + payload: str, +) -> list[StreamingResponsesResponse]: + streamed = state.tool_call_payload + if payload == streamed: + return [] + state.tool_call_payload = payload + # A payload that is not an extension of what was streamed cannot be + # retracted, so the authoritative value is resent in full. + delta = payload[len(streamed) :] if payload.startswith(streamed) else payload + return [ + ResponseCustomToolCallInputDeltaEvent( + type="response.custom_tool_call_input.delta", + sequence_number=-1, + output_index=state.output_index, + item_id=state.current_item_id, + delta=delta, + ) + ] + + def emit_simple_tool_call_delta( state: SimpleStreamingState, delta: str, ) -> list[StreamingResponsesResponse]: state.accumulated_text += delta + if state.tool_call_is_custom: + return _emit_custom_tool_input_delta( + state, decode_custom_tool_input_prefix(state.accumulated_text) + ) return [ ResponseFunctionCallArgumentsDeltaEvent( type="response.function_call_arguments.delta", @@ -1085,32 +1135,58 @@ def emit_simple_tool_call_delta( def emit_simple_tool_call_done( state: SimpleStreamingState, ) -> list[StreamingResponsesResponse]: - events: list[StreamingResponsesResponse] = [ - ResponseFunctionCallArgumentsDoneEvent( - type="response.function_call_arguments.done", - sequence_number=-1, - output_index=state.output_index, - item_id=state.current_item_id, - arguments=state.accumulated_text, + events: list[StreamingResponsesResponse] + item: ResponseOutputItem + if state.tool_call_is_custom: + payload = decode_custom_tool_input(state.accumulated_text) + events = _emit_custom_tool_input_delta(state, payload) + events.append( + ResponseCustomToolCallInputDoneEvent( + type="response.custom_tool_call_input.done", + sequence_number=-1, + output_index=state.output_index, + item_id=state.current_item_id, + input=payload, + ) + ) + item = ResponseCustomToolCall( + type="custom_tool_call", + id=state.current_item_id, + call_id=state.tool_call_id, name=state.tool_call_name, - ), + input=payload, + ) + else: + events = [ + ResponseFunctionCallArgumentsDoneEvent( + type="response.function_call_arguments.done", + sequence_number=-1, + output_index=state.output_index, + item_id=state.current_item_id, + arguments=state.accumulated_text, + name=state.tool_call_name, + ) + ] + item = ResponseFunctionToolCall( + type="function_call", + name=state.tool_call_name, + namespace=state.tool_call_namespace, + arguments=state.accumulated_text, + status="completed", + id=state.current_item_id, + call_id=state.tool_call_id, + ) + events.append( ResponseOutputItemDoneEvent( type="response.output_item.done", sequence_number=-1, output_index=state.output_index, - item=ResponseFunctionToolCall( - type="function_call", - name=state.tool_call_name, - namespace=state.tool_call_namespace, - arguments=state.accumulated_text, - status="completed", - id=state.current_item_id, - call_id=state.tool_call_id, - ), - ), - ] + item=item, + ) + ) state.output_index += 1 state.tool_call_namespace = None + state.tool_call_is_custom = False state.current_state = _StateType.NONE return events @@ -1194,6 +1270,7 @@ def __init__( ) -> None: self.state = state or SimpleStreamingState() self.tool_call_name_map = build_responses_tool_call_name_map(tools) + self.custom_tool_names = custom_tool_names(tools) self.output_items: list[ResponseOutputItem] = [] def resolve_target_state( @@ -1264,6 +1341,7 @@ def open( call_name.name, tool_call.index, call_name.namespace, + call_name.name in self.custom_tool_names, tool_call.id, ) return handlers.open_fn(self.state) diff --git a/vllm/entrypoints/openai/responses/utils.py b/vllm/entrypoints/openai/responses/utils.py index c4523964e9dd..d9b71d684922 100644 --- a/vllm/entrypoints/openai/responses/utils.py +++ b/vllm/entrypoints/openai/responses/utils.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json from collections.abc import Iterable from typing import Any +import partial_json_parser from openai.types.chat import ( ChatCompletionAssistantMessageParam, ChatCompletionMessageToolCallParam, @@ -13,6 +15,7 @@ Function as FunctionCallTool, ) from openai.types.responses import ( + ResponseCustomToolCall, ResponseFunctionToolCall, ResponseOutputItem, ResponseOutputMessage, @@ -28,6 +31,7 @@ Content as ResponseReasoningTextContent, ) from openai.types.responses.tool import Tool +from partial_json_parser.core.options import Allow from vllm import envs from vllm.entrypoints.generate.base.protocol import FunctionCall, FunctionDefinition @@ -40,6 +44,7 @@ from vllm.logger import init_logger from vllm.renderers.chat_utils import make_tool_call_id from vllm.tool_parsers.utils import ( + CUSTOM_TOOL_INPUT_KEY, build_responses_tool_call_name_map, flat_namespace_tool_name, iter_response_function_tool_dicts, @@ -50,6 +55,51 @@ logger = init_logger(__name__) +def custom_tool_names(tools: list[Tool] | None) -> frozenset[str]: + return frozenset(tool.name for tool in tools or [] if tool.type == "custom") + + +def encode_custom_tool_input(payload: str) -> str: + return json.dumps({CUSTOM_TOOL_INPUT_KEY: payload}, ensure_ascii=False) + + +def _custom_tool_payload(parsed: Any) -> str | None: + if not isinstance(parsed, dict): + return None + value = parsed.get(CUSTOM_TOOL_INPUT_KEY) + return value if isinstance(value, str) else None + + +def decode_custom_tool_input(arguments: str) -> str: + """Payload of a completed shim call; bare text is passed through.""" + try: + payload = _custom_tool_payload(json.loads(arguments)) + except ValueError: + return arguments + return arguments if payload is None else payload + + +def decode_custom_tool_input_prefix(arguments: str) -> str: + """Longest decodable payload prefix of partial shim arguments, so that + streamed ``custom_tool_call_input`` deltas stay a prefix of the final input.""" + try: + parsed = partial_json_parser.loads(arguments, Allow.STR | Allow.OBJ) + except (ValueError, IndexError, AssertionError): + return "" + payload = _custom_tool_payload(parsed) or "" + if payload and "\ud800" <= payload[-1] <= "\udbff": + payload = payload[:-1] + return payload + + +def _item_type(item: ResponseInputOutputItem) -> str | None: + return item.get("type") if isinstance(item, dict) else getattr(item, "type", None) + + +def _item_dict(item: ResponseInputOutputItem) -> dict[str, Any]: + return item if isinstance(item, dict) else item.model_dump() + + def build_response_output_items( reasoning: str | None, content: str | None, @@ -59,6 +109,7 @@ def build_response_output_items( ) -> list[ResponseOutputItem]: outputs: list[ResponseOutputItem] = [] tool_call_name_map = build_responses_tool_call_name_map(tools) + custom_names = custom_tool_names(tools) if reasoning: outputs.append( @@ -93,14 +144,27 @@ def build_response_output_items( if tool_calls: for idx, tool_call in enumerate(tool_calls): + call_id = tool_call.id or make_tool_call_id( + func_name=tool_call.name, idx=idx + ) + if tool_call.name in custom_names: + outputs.append( + ResponseCustomToolCall( + id=f"ctc_{random_uuid()}", + call_id=call_id, + type="custom_tool_call", + name=tool_call.name, + input=decode_custom_tool_input(tool_call.arguments), + ) + ) + continue call_name = resolve_responses_tool_call_name( tool_call.name, tool_call_name_map=tool_call_name_map ) outputs.append( ResponseFunctionToolCall( id=f"fc_{random_uuid()}", - call_id=tool_call.id - or make_tool_call_id(func_name=tool_call.name, idx=idx), + call_id=call_id, type="function_call", status="completed", name=call_name.name, @@ -255,6 +319,7 @@ def _construct_message_from_response_item( prev_msg if prev_msg and prev_msg.get("role") == "assistant" else None ) + tool_call: ChatCompletionMessageToolCallParam | None = None if isinstance(item, ResponseFunctionToolCall): tool_name = item.name if item.namespace: @@ -267,6 +332,17 @@ def _construct_message_from_response_item( ), type="function", ) + elif _item_type(item) == "custom_tool_call": + call = _item_dict(item) + tool_call = ChatCompletionMessageToolCallParam( + id=call["call_id"], + function=FunctionCallTool( + name=call["name"], + arguments=encode_custom_tool_input(call.get("input") or ""), + ), + type="function", + ) + if tool_call is not None: if prev_assistant_msg: tool_calls = prev_assistant_msg.get("tool_calls") if tool_calls is None: @@ -285,14 +361,14 @@ def _construct_message_from_response_item( logger.warning( "Previous assistant message has unknown tool_calls format. " "Tool call merging is skipped and a new assistant message is created. " - "Item %s", - item.id, + "Call %s", + tool_call["id"], ) return ChatCompletionAssistantMessageParam( role="assistant", tool_calls=[tool_call], ) - elif isinstance(item, ResponseReasoningItem): + if isinstance(item, ResponseReasoningItem): reasoning = "" if item.encrypted_content: raise VLLMValidationError( @@ -336,12 +412,12 @@ def _construct_message_from_response_item( content=item.output, tool_call_id=item.call_id, ) - elif isinstance(item, dict) and item.get("type") == "function_call_output": - # Append the function call output as a tool message. + elif _item_type(item) in ("function_call_output", "custom_tool_call_output"): + output = _item_dict(item) return ChatCompletionToolMessageParam( role="tool", - content=item.get("output"), - tool_call_id=item.get("call_id"), + content=output.get("output"), + tool_call_id=output.get("call_id"), ) elif isinstance(item, dict) and item.get("role") == "assistant": content = item.get("content") diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index 61826dc8c86c..12a9c2369843 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -7,7 +7,7 @@ from collections.abc import Sequence from dataclasses import dataclass, field, replace -from openai.types.responses import ToolChoiceFunction +from openai.types.responses import ToolChoiceCustom, ToolChoiceFunction from pydantic import TypeAdapter, ValidationError from xgrammar import Grammar, StructuralTag from xgrammar.structural_tag import ( @@ -359,7 +359,9 @@ def extract_reasoning( def _get_function_name( self, request: ChatCompletionRequest | ResponsesRequest ) -> str: - if request.tool_choice and isinstance(request.tool_choice, ToolChoiceFunction): + if request.tool_choice and isinstance( + request.tool_choice, (ToolChoiceFunction, ToolChoiceCustom) + ): return request.tool_choice.name if request.tool_choice and isinstance( request.tool_choice, ChatCompletionNamedToolChoiceParam @@ -398,7 +400,7 @@ def _extract_tool_calls( supports_required_and_named = tool_parser.supports_required_and_named is_named_tool_choice = request.tool_choice and isinstance( request.tool_choice, - (ToolChoiceFunction, ChatCompletionNamedToolChoiceParam), + (ToolChoiceFunction, ToolChoiceCustom, ChatCompletionNamedToolChoiceParam), ) is_required_tool_choice = request.tool_choice == "required" is_auto_tool_choice = enable_auto_tools and ( @@ -503,7 +505,11 @@ def _apply_structural_tag( or request.tool_choice == "required" or isinstance( request.tool_choice, - (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction), + ( + ChatCompletionNamedToolChoiceParam, + ToolChoiceFunction, + ToolChoiceCustom, + ), ) ) if not need_tool_calling: @@ -723,7 +729,11 @@ def _extract_tool_calls_streaming( and request.tool_choice and isinstance( request.tool_choice, - (ToolChoiceFunction, ChatCompletionNamedToolChoiceParam), + ( + ToolChoiceFunction, + ToolChoiceCustom, + ChatCompletionNamedToolChoiceParam, + ), ) ): delta_message, function_name_returned = extract_named_tool_call_streaming( diff --git a/vllm/parser/response_template.py b/vllm/parser/response_template.py index 85f8b67efbd6..b1665f1772a5 100644 --- a/vllm/parser/response_template.py +++ b/vllm/parser/response_template.py @@ -9,7 +9,7 @@ from collections.abc import Callable, Iterable, Sequence from typing import TYPE_CHECKING, Any, cast -from openai.types.responses import ToolChoiceFunction +from openai.types.responses import ToolChoiceCustom, ToolChoiceFunction from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, @@ -173,7 +173,8 @@ def _streaming_template(template: dict[str, Any]) -> ResponseTemplate: def _is_named_tool_choice(request: ChatCompletionRequest | ResponsesRequest) -> bool: return isinstance( - request.tool_choice, (ToolChoiceFunction, ChatCompletionNamedToolChoiceParam) + request.tool_choice, + (ToolChoiceFunction, ToolChoiceCustom, ChatCompletionNamedToolChoiceParam), ) diff --git a/vllm/tool_parsers/gemma4_engine_tool_parser.py b/vllm/tool_parsers/gemma4_engine_tool_parser.py index 04c03ecaa20a..fcb578b129cb 100644 --- a/vllm/tool_parsers/gemma4_engine_tool_parser.py +++ b/vllm/tool_parsers/gemma4_engine_tool_parser.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from openai.types.responses import ToolChoiceFunction +from openai.types.responses import ToolChoiceCustom, ToolChoiceFunction from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionNamedToolChoiceParam, @@ -29,7 +29,12 @@ def adjust_request( if request.tools: tc = request.tool_choice if tc == "required" or isinstance( - tc, (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction) + tc, + ( + ChatCompletionNamedToolChoiceParam, + ToolChoiceFunction, + ToolChoiceCustom, + ), ): request.skip_special_tokens = False return request diff --git a/vllm/tool_parsers/k2_horizon_tool_parser.py b/vllm/tool_parsers/k2_horizon_tool_parser.py index 742bf7d14013..a63f5fa12239 100644 --- a/vllm/tool_parsers/k2_horizon_tool_parser.py +++ b/vllm/tool_parsers/k2_horizon_tool_parser.py @@ -6,7 +6,7 @@ from typing import Any import regex as re -from openai.types.responses import ToolChoiceFunction +from openai.types.responses import ToolChoiceCustom, ToolChoiceFunction from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, @@ -184,7 +184,7 @@ def _named_tool(request: ToolParserRequest) -> str | None: tool_choice = request.tool_choice if isinstance(tool_choice, ChatCompletionNamedToolChoiceParam): return tool_choice.function.name - if isinstance(tool_choice, ToolChoiceFunction): + if isinstance(tool_choice, (ToolChoiceFunction, ToolChoiceCustom)): return tool_choice.name return None diff --git a/vllm/tool_parsers/kimi_k3_tool_parser.py b/vllm/tool_parsers/kimi_k3_tool_parser.py index 76c6b0aa657c..15b5a22df861 100644 --- a/vllm/tool_parsers/kimi_k3_tool_parser.py +++ b/vllm/tool_parsers/kimi_k3_tool_parser.py @@ -35,7 +35,7 @@ from collections.abc import Sequence import regex as re -from openai.types.responses import ToolChoiceFunction +from openai.types.responses import ToolChoiceCustom, ToolChoiceFunction from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, @@ -149,7 +149,11 @@ def adjust_request( ) -> ChatCompletionRequest | ResponsesRequest: named = isinstance( request.tool_choice, - (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction), + ( + ChatCompletionNamedToolChoiceParam, + ToolChoiceFunction, + ToolChoiceCustom, + ), ) structured_outputs = getattr(request, "structured_outputs", None) has_structural_tag = ( diff --git a/vllm/tool_parsers/muse_glimmer_tool_parser.py b/vllm/tool_parsers/muse_glimmer_tool_parser.py index 5a0b26cf423a..644ca5ed3e3a 100644 --- a/vllm/tool_parsers/muse_glimmer_tool_parser.py +++ b/vllm/tool_parsers/muse_glimmer_tool_parser.py @@ -38,7 +38,7 @@ from collections.abc import Iterator, Sequence import regex as re -from openai.types.responses import ToolChoiceFunction +from openai.types.responses import ToolChoiceCustom, ToolChoiceFunction from transformers import PreTrainedTokenizerBase from vllm.entrypoints.generate.base.protocol import ( @@ -228,7 +228,12 @@ def adjust_request( if request.tools and ( tool_choice == "required" or isinstance( - tool_choice, (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction) + tool_choice, + ( + ChatCompletionNamedToolChoiceParam, + ToolChoiceFunction, + ToolChoiceCustom, + ), ) ): return request diff --git a/vllm/tool_parsers/poolside_v1_tool_parser.py b/vllm/tool_parsers/poolside_v1_tool_parser.py index d2dfb7815fbe..998927264757 100644 --- a/vllm/tool_parsers/poolside_v1_tool_parser.py +++ b/vllm/tool_parsers/poolside_v1_tool_parser.py @@ -16,7 +16,7 @@ import partial_json_parser.core.complete import regex as re -from openai.types.responses import ToolChoiceFunction +from openai.types.responses import ToolChoiceCustom, ToolChoiceFunction from partial_json_parser.core.options import Allow from vllm.entrypoints.generate.base.protocol import ( @@ -171,7 +171,12 @@ def adjust_request( if request.tools: tc = request.tool_choice if tc == "required" or isinstance( - tc, (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction) + tc, + ( + ChatCompletionNamedToolChoiceParam, + ToolChoiceFunction, + ToolChoiceCustom, + ), ): request.skip_special_tokens = False return request diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py index f3148a827eb0..7e5902b8c2f6 100644 --- a/vllm/tool_parsers/structural_tag_registry.py +++ b/vllm/tool_parsers/structural_tag_registry.py @@ -6,7 +6,7 @@ from collections.abc import Callable, Sequence from typing import Any, Literal, TypeAlias -from openai.types.responses import FunctionTool +from openai.types.responses import CustomTool, FunctionTool, ToolChoiceCustom from openai.types.responses.response import ToolChoice as ResponsesToolChoice from openai.types.responses.tool import Tool as ResponsesTool from openai.types.responses.tool_choice_allowed import ToolChoiceAllowed @@ -39,7 +39,11 @@ ChatCompletionToolsParam, ) from vllm.tool_parsers.tool_strict_level import ToolStrictLevel -from vllm.tool_parsers.utils import get_function_tools, require_function_tools +from vllm.tool_parsers.utils import ( + custom_tool_parameters, + get_function_tools, + require_function_tools, +) ToolChoice: TypeAlias = ( Literal["none", "auto", "required"] @@ -287,6 +291,11 @@ def _dump_tool_for_xgrammar( if tool.strict is not None: function["strict"] = tool.strict return {"type": "function", "function": function} + if isinstance(tool, CustomTool): + function = {"name": tool.name, "parameters": custom_tool_parameters()} + if tool.description is not None: + function["description"] = tool.description + return {"type": "function", "function": function} dumped_tool = tool.model_dump(mode="json", exclude_none=True) if isinstance(tool, ChatCompletionToolsParam): return dumped_tool @@ -306,7 +315,7 @@ def _dump_tool_choice_for_xgrammar( if isinstance(tool_choice, ChatCompletionNamedToolChoiceParam): return tool_choice.model_dump(mode="json", exclude_none=True) - if isinstance(tool_choice, ToolChoiceFunction): + if isinstance(tool_choice, (ToolChoiceFunction, ToolChoiceCustom)): return { "type": "function", "function": {"name": tool_choice.name}, diff --git a/vllm/tool_parsers/utils.py b/vllm/tool_parsers/utils.py index 7a93c843b8e3..2462824ccfc0 100644 --- a/vllm/tool_parsers/utils.py +++ b/vllm/tool_parsers/utils.py @@ -13,8 +13,10 @@ import partial_json_parser from openai.types.responses import ( + CustomTool, FunctionTool, NamespaceTool, + ToolChoiceCustom, ToolChoiceFunction, ) from openai.types.responses.tool import Tool as ResponsesTool @@ -37,6 +39,32 @@ logger = init_logger(__name__) +CUSTOM_TOOL_INPUT_KEY = "input" + + +def custom_tool_parameters() -> dict[str, Any]: + """Schema of the single-string function shim a ``custom`` tool is shown as.""" + return { + "type": "object", + "properties": { + CUSTOM_TOOL_INPUT_KEY: { + "type": "string", + "description": "The freeform text payload for this tool.", + } + }, + "required": [CUSTOM_TOOL_INPUT_KEY], + "additionalProperties": False, + } + + +def custom_tool_as_function_dict(tool: CustomTool) -> dict[str, Any]: + return { + "type": "function", + "name": tool.name, + "description": tool.description, + "parameters": custom_tool_parameters(), + } + def safe_literal_eval(text: str): with warnings.catch_warnings(): @@ -188,8 +216,9 @@ def get_function_tools( This is the single source for both the chat template's tool list and the tools a tool-call grammar constrains, so the two cannot drift apart. - Namespace functions are flattened to ``__``; other - Responses tool types are never rendered for the model and are dropped. + Namespace functions are flattened to ``__``; custom tools + become a single-string function shim. Other Responses tool types are never + rendered for the model and are dropped. """ function_tools: list[ChatCompletionToolsParam | FunctionTool] = [] for tool in tools: @@ -206,6 +235,10 @@ def get_function_tools( for namespaced_tool in tool.tools if namespaced_tool.type == "function" ) + elif isinstance(tool, CustomTool): + function_tools.append( + FunctionTool.model_validate(custom_tool_as_function_dict(tool)) + ) elif isinstance(tool, (FunctionTool, ChatCompletionToolsParam)): function_tools.append(tool) return function_tools @@ -307,7 +340,7 @@ def find_tool_properties( if not tools: return {} for tool in tools: - if isinstance(tool, (FunctionTool, NamespaceTool)): + if isinstance(tool, (FunctionTool, NamespaceTool, CustomTool)): for name, params in iter_response_function_tool_info(tool): if name == tool_name: return (params or {}).get("properties", {}) @@ -328,7 +361,7 @@ def find_tool_name( if not tools: return False for tool in tools: - if isinstance(tool, (FunctionTool, NamespaceTool)): + if isinstance(tool, (FunctionTool, NamespaceTool, CustomTool)): for name, _ in iter_response_function_tool_info(tool): if name == tool_name: return True @@ -414,7 +447,9 @@ def _get_json_schema_from_tools( def get_json_schema_from_tools( - tool_choice: str | ToolChoiceFunction | ChatCompletionNamedToolChoiceParam, + tool_choice: ( + str | ToolChoiceFunction | ToolChoiceCustom | ChatCompletionNamedToolChoiceParam + ), tools: list[Tool] | None, parallel_tool_calls: bool | None = None, ) -> str | dict | None: @@ -423,12 +458,12 @@ def get_json_schema_from_tools( return None # tool_choice: Forced Function (Responses) if (not isinstance(tool_choice, str)) and isinstance( - tool_choice, ToolChoiceFunction + tool_choice, (ToolChoiceFunction, ToolChoiceCustom) ): tool_name = tool_choice.name responses_tool_map: dict[str, dict[str, Any] | None] = {} for tool in tools: - if not isinstance(tool, (FunctionTool, NamespaceTool)): + if not isinstance(tool, (FunctionTool, NamespaceTool, CustomTool)): continue for name, params in iter_response_function_tool_info(tool): responses_tool_map[name] = params From 0466bd36ea9f6300acbc15c705b1ebe61590c691 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Mon, 12 Oct 2026 17:43:47 +0800 Subject: [PATCH 286/624] [Rust Frontend] Build the Kimi K3 full-output grammar from its parser mode (#59190) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../chat/tests/grammar_replay/Kimi-K3.json | 1016 ++++++++++------- .../src/chat/tests/grammar_replay/Kimi-K3.txt | 99 +- rust/src/parser/src/output_grammar.rs | 19 +- rust/src/parser/src/unified/kimi_k3.rs | 14 +- .../src/parser/src/unified/kimi_k3/grammar.rs | 356 ++++++ .../src/unified/kimi_k3/grammar/arguments.rs | 184 +++ .../src/unified/kimi_k3/structural_tag.rs | 343 ------ rust/src/parser/src/unified/mod.rs | 2 +- 8 files changed, 1246 insertions(+), 787 deletions(-) create mode 100644 rust/src/parser/src/unified/kimi_k3/grammar.rs create mode 100644 rust/src/parser/src/unified/kimi_k3/grammar/arguments.rs delete mode 100644 rust/src/parser/src/unified/kimi_k3/structural_tag.rs diff --git a/rust/src/chat/tests/grammar_replay/Kimi-K3.json b/rust/src/chat/tests/grammar_replay/Kimi-K3.json index 7fe53b0ab2be..5306b936604f 100644 --- a/rust/src/chat/tests/grammar_replay/Kimi-K3.json +++ b/rust/src/chat/tests/grammar_replay/Kimi-K3.json @@ -8,213 +8,413 @@ "format": { "type": "sequence", "elements": [ - { - "type": "optional", - "content": { - "type": "const_string", - "value": "<|open|>response<|sep|>" - } - }, { "type": "tag", "begin": "", "content": { "type": "any_text", "excludes": [ - "<|open|>", - "<|close|>", + "<|close|>think<|sep|>", "<|end_of_msg|>" ], "max_tokens": null, "max_chars": null }, - "end": "<|close|>response<|sep|>" + "end": "<|close|>think<|sep|>" }, { - "type": "optional", - "content": { - "type": "tag", - "begin": "<|open|>tools<|sep|>", - "content": { - "type": "tags_with_separator", - "tags": [ + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ { "type": "tag", - "begin": "<|open|>call tool=\"get_weather\" index=\"", + "begin": "<|open|>response<|sep|>", "content": { - "type": "sequence", - "elements": [ - { - "type": "regex", - "pattern": "[1-9][0-9]*" - }, - { - "type": "const_string", - "value": "\"<|sep|>" - }, - { - "type": "or", - "elements": [ - { - "type": "tag", - "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>", - "content": { - "type": "any_text", - "excludes": [ - "<|close|>argument<|sep|>", - "<|close|>call<|sep|>" - ], - "max_tokens": null, - "max_chars": null - }, - "end": "<|close|>argument<|sep|>" - }, - { - "type": "tag", - "begin": "<|open|>json type=\"object\"<|sep|>", - "content": { - "type": "json_schema", - "json_schema": { - "type": "object", - "properties": { - "location": { - "type": "string" - } - }, - "required": [ - "location" - ] - }, - "style": "json", - "any_order": false, - "max_whitespace_cnt": null, - "excludes": [] - }, - "end": "<|close|>json<|sep|>" - } - ] - } - ] + "type": "any_text", + "excludes": [ + "<|open|>", + "<|close|>", + "<|end_of_msg|>" + ], + "max_tokens": null, + "max_chars": null }, - "end": "<|close|>call<|sep|>" + "end": "<|close|>response<|sep|>" }, { - "type": "tag", - "begin": "<|open|>call tool=\"add\" index=\"", + "type": "optional", "content": { - "type": "sequence", - "elements": [ - { - "type": "regex", - "pattern": "[1-9][0-9]*" - }, - { - "type": "const_string", - "value": "\"<|sep|>" - }, - { - "type": "or", - "elements": [ - { + "type": "tag", + "begin": "<|open|>tools<|sep|>", + "content": { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<|open|>call tool=\"get_weather\" index=\"", + "content": { "type": "sequence", "elements": [ { - "type": "tag", - "begin": "<|open|>argument key=\"y\" type=\"number\"<|sep|>", - "content": { - "type": "json_schema", - "json_schema": { - "type": "number" - }, - "style": "json", - "any_order": false, - "max_whitespace_cnt": null, - "excludes": [] - }, - "end": "<|close|>argument<|sep|>" + "type": "regex", + "pattern": "[1-9][0-9]*" }, { - "type": "tag", - "begin": "<|open|>argument key=\"x\" type=\"number\"<|sep|>", - "content": { - "type": "json_schema", - "json_schema": { - "type": "number" + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>", + "content": { + "type": "any_text", + "excludes": [ + "<|close|>argument<|sep|>", + "<|close|>call<|sep|>" + ], + "max_tokens": null, + "max_chars": null + }, + "end": "<|close|>argument<|sep|>" }, - "style": "json", - "any_order": false, - "max_whitespace_cnt": null, - "excludes": [] - }, - "end": "<|close|>argument<|sep|>" + { + "type": "tag", + "begin": "<|open|>json type=\"object\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>json<|sep|>" + } + ] + } + ] + }, + "end": "<|close|>call<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>call tool=\"add\" index=\"", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "[1-9][0-9]*" }, { - "type": "tag", - "begin": "<|open|>argument key=\"items\" type=\"array\"<|sep|>", - "content": { - "type": "json_schema", - "json_schema": { - "type": "array", - "items": { - "type": "string" - } + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"y\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"x\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"items\" type=\"array\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "array", + "items": { + "type": "string" + } + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + } + ] }, - "style": "json", - "any_order": false, - "max_whitespace_cnt": null, - "excludes": [] - }, - "end": "<|close|>argument<|sep|>" + { + "type": "tag", + "begin": "<|open|>json type=\"object\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>json<|sep|>" + } + ] } ] }, - { - "type": "tag", - "begin": "<|open|>json type=\"object\"<|sep|>", - "content": { - "type": "json_schema", - "json_schema": { - "type": "object", - "properties": { - "y": { - "type": "number" + "end": "<|close|>call<|sep|>" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + "end": "<|close|>tools<|sep|>" + } + } + ] + }, + { + "type": "tag", + "begin": "<|open|>tools<|sep|>", + "content": { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<|open|>call tool=\"get_weather\" index=\"", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "[1-9][0-9]*" + }, + { + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>", + "content": { + "type": "any_text", + "excludes": [ + "<|close|>argument<|sep|>", + "<|close|>call<|sep|>" + ], + "max_tokens": null, + "max_chars": null + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>json type=\"object\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>json<|sep|>" + } + ] + } + ] + }, + "end": "<|close|>call<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>call tool=\"add\" index=\"", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "[1-9][0-9]*" + }, + { + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"y\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] }, - "x": { - "type": "number" + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"x\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"items\" type=\"array\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "array", + "items": { + "type": "string" + } + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] }, - "items": { - "type": "array", + "end": "<|close|>argument<|sep|>" + } + ] + }, + { + "type": "tag", + "begin": "<|open|>json type=\"object\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, "items": { - "type": "string" + "type": "array", + "items": { + "type": "string" + } } - } + }, + "required": [ + "y", + "x", + "items" + ] }, - "required": [ - "y", - "x", - "items" - ] + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] }, - "style": "json", - "any_order": false, - "max_whitespace_cnt": null, - "excludes": [] - }, - "end": "<|close|>json<|sep|>" - } - ] - } - ] - }, - "end": "<|close|>call<|sep|>" - } - ], - "separator": "", - "at_least_one": true, - "stop_after_first": false - }, - "end": "<|close|>tools<|sep|>" - } + "end": "<|close|>json<|sep|>" + } + ] + } + ] + }, + "end": "<|close|>call<|sep|>" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + "end": "<|close|>tools<|sep|>" + } + ] }, { "type": "optional", @@ -225,7 +425,7 @@ } ] }, - "coverage": "final_output_only" + "coverage": "from_token_zero" }, "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|>I will call the tools.<|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool=\"add\" index=\"2\"<|sep|><|open|>argument key=\"y\" type=\"number\"<|sep|>1.0<|close|>argument<|sep|><|open|>argument key=\"x\" type=\"number\"<|sep|>2<|close|>argument<|sep|><|open|>argument key=\"items\" type=\"array\"<|sep|>[\"left\",\"right\"]<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 40, 748, 2390, 276, 7697, 13, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163587, 10257, 4453, 878, 1210, 1, 4002, 878, 17, 1, 163589, 163587, 47185, 2355, 878, 88, 1, 1798, 878, 10006, 1, 163589, 16, 13, 15, 163588, 47185, 163589, 163587, 47185, 2355, 878, 87, 1, 1798, 878, 10006, 1, 163589, 17, 163588, 47185, 163589, 163587, 47185, 2355, 878, 7989, 1, 1798, 878, 4392, 1, 163589, 3530, 3794, 3923, 2277, 3956, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] @@ -235,210 +435,225 @@ "format": { "type": "sequence", "elements": [ - { - "type": "optional", - "content": { - "type": "const_string", - "value": "<|open|>response<|sep|>" - } - }, { "type": "tag", "begin": "", "content": { "type": "any_text", "excludes": [ - "<|open|>", - "<|close|>", + "<|close|>think<|sep|>", "<|end_of_msg|>" ], "max_tokens": null, "max_chars": null }, - "end": "<|close|>response<|sep|>" + "end": "<|close|>think<|sep|>" }, { - "type": "tag", - "begin": "<|open|>tools<|sep|>", - "content": { - "type": "tags_with_separator", - "tags": [ - { + "type": "sequence", + "elements": [ + { + "type": "optional", + "content": { "type": "tag", - "begin": "<|open|>call tool=\"get_weather\" index=\"", + "begin": "<|open|>response<|sep|>", "content": { - "type": "sequence", - "elements": [ - { - "type": "regex", - "pattern": "[1-9][0-9]*" - }, - { - "type": "const_string", - "value": "\"<|sep|>" - }, - { - "type": "or", + "type": "any_text", + "excludes": [ + "<|open|>", + "<|close|>", + "<|end_of_msg|>" + ], + "max_tokens": null, + "max_chars": null + }, + "end": "<|close|>response<|sep|>" + } + }, + { + "type": "tag", + "begin": "<|open|>tools<|sep|>", + "content": { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<|open|>call tool=\"get_weather\" index=\"", + "content": { + "type": "sequence", "elements": [ { - "type": "tag", - "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>", - "content": { - "type": "any_text", - "excludes": [ - "<|close|>argument<|sep|>", - "<|close|>call<|sep|>" - ], - "max_tokens": null, - "max_chars": null - }, - "end": "<|close|>argument<|sep|>" + "type": "regex", + "pattern": "[1-9][0-9]*" }, { - "type": "tag", - "begin": "<|open|>json type=\"object\"<|sep|>", - "content": { - "type": "json_schema", - "json_schema": { - "type": "object", - "properties": { - "location": { - "type": "string" - } - }, - "required": [ - "location" - ] - }, - "style": "json", - "any_order": false, - "max_whitespace_cnt": null, - "excludes": [] - }, - "end": "<|close|>json<|sep|>" - } - ] - } - ] - }, - "end": "<|close|>call<|sep|>" - }, - { - "type": "tag", - "begin": "<|open|>call tool=\"add\" index=\"", - "content": { - "type": "sequence", - "elements": [ - { - "type": "regex", - "pattern": "[1-9][0-9]*" - }, - { - "type": "const_string", - "value": "\"<|sep|>" - }, - { - "type": "or", - "elements": [ + "type": "const_string", + "value": "\"<|sep|>" + }, { - "type": "sequence", + "type": "or", "elements": [ { "type": "tag", - "begin": "<|open|>argument key=\"y\" type=\"number\"<|sep|>", + "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>", "content": { - "type": "json_schema", - "json_schema": { - "type": "number" - }, - "style": "json", - "any_order": false, - "max_whitespace_cnt": null, - "excludes": [] + "type": "any_text", + "excludes": [ + "<|close|>argument<|sep|>", + "<|close|>call<|sep|>" + ], + "max_tokens": null, + "max_chars": null }, "end": "<|close|>argument<|sep|>" }, { "type": "tag", - "begin": "<|open|>argument key=\"x\" type=\"number\"<|sep|>", + "begin": "<|open|>json type=\"object\"<|sep|>", "content": { "type": "json_schema", "json_schema": { - "type": "number" + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] }, "style": "json", "any_order": false, "max_whitespace_cnt": null, "excludes": [] }, - "end": "<|close|>argument<|sep|>" + "end": "<|close|>json<|sep|>" + } + ] + } + ] + }, + "end": "<|close|>call<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>call tool=\"add\" index=\"", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "[1-9][0-9]*" + }, + { + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"y\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"x\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"items\" type=\"array\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "array", + "items": { + "type": "string" + } + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + } + ] }, { "type": "tag", - "begin": "<|open|>argument key=\"items\" type=\"array\"<|sep|>", + "begin": "<|open|>json type=\"object\"<|sep|>", "content": { "type": "json_schema", "json_schema": { - "type": "array", - "items": { - "type": "string" - } + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] }, "style": "json", "any_order": false, "max_whitespace_cnt": null, "excludes": [] }, - "end": "<|close|>argument<|sep|>" + "end": "<|close|>json<|sep|>" } ] - }, - { - "type": "tag", - "begin": "<|open|>json type=\"object\"<|sep|>", - "content": { - "type": "json_schema", - "json_schema": { - "type": "object", - "properties": { - "y": { - "type": "number" - }, - "x": { - "type": "number" - }, - "items": { - "type": "array", - "items": { - "type": "string" - } - } - }, - "required": [ - "y", - "x", - "items" - ] - }, - "style": "json", - "any_order": false, - "max_whitespace_cnt": null, - "excludes": [] - }, - "end": "<|close|>json<|sep|>" } ] - } - ] - }, - "end": "<|close|>call<|sep|>" - } - ], - "separator": "", - "at_least_one": true, - "stop_after_first": false - }, - "end": "<|close|>tools<|sep|>" + }, + "end": "<|close|>call<|sep|>" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + "end": "<|close|>tools<|sep|>" + } + ] }, { "type": "optional", @@ -449,7 +664,7 @@ } ] }, - "coverage": "final_output_only" + "coverage": "from_token_zero" }, "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool=\"add\" index=\"2\"<|sep|><|open|>argument key=\"y\" type=\"number\"<|sep|>1.0<|close|>argument<|sep|><|open|>argument key=\"x\" type=\"number\"<|sep|>2<|close|>argument<|sep|><|open|>argument key=\"items\" type=\"array\"<|sep|>[\"left\",\"right\"]<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163587, 10257, 4453, 878, 1210, 1, 4002, 878, 17, 1, 163589, 163587, 47185, 2355, 878, 88, 1, 1798, 878, 10006, 1, 163589, 16, 13, 15, 163588, 47185, 163589, 163587, 47185, 2355, 878, 87, 1, 1798, 878, 10006, 1, 163589, 17, 163588, 47185, 163589, 163587, 47185, 2355, 878, 7989, 1, 1798, 878, 4392, 1, 163589, 3530, 3794, 3923, 2277, 3956, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] @@ -459,100 +674,115 @@ "format": { "type": "sequence", "elements": [ - { - "type": "optional", - "content": { - "type": "const_string", - "value": "<|open|>response<|sep|>" - } - }, { "type": "tag", "begin": "", "content": { "type": "any_text", "excludes": [ - "<|open|>", - "<|close|>", + "<|close|>think<|sep|>", "<|end_of_msg|>" ], "max_tokens": null, "max_chars": null }, - "end": "<|close|>response<|sep|>" + "end": "<|close|>think<|sep|>" }, { - "type": "tag", - "begin": "<|open|>tools<|sep|>", - "content": { - "type": "tags_with_separator", - "tags": [ - { + "type": "sequence", + "elements": [ + { + "type": "optional", + "content": { "type": "tag", - "begin": "<|open|>call tool=\"get_weather\" index=\"", + "begin": "<|open|>response<|sep|>", "content": { - "type": "sequence", - "elements": [ - { - "type": "regex", - "pattern": "[1-9][0-9]*" - }, - { - "type": "const_string", - "value": "\"<|sep|>" - }, - { - "type": "or", + "type": "any_text", + "excludes": [ + "<|open|>", + "<|close|>", + "<|end_of_msg|>" + ], + "max_tokens": null, + "max_chars": null + }, + "end": "<|close|>response<|sep|>" + } + }, + { + "type": "tag", + "begin": "<|open|>tools<|sep|>", + "content": { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<|open|>call tool=\"get_weather\" index=\"", + "content": { + "type": "sequence", "elements": [ { - "type": "tag", - "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>", - "content": { - "type": "any_text", - "excludes": [ - "<|close|>argument<|sep|>", - "<|close|>call<|sep|>" - ], - "max_tokens": null, - "max_chars": null - }, - "end": "<|close|>argument<|sep|>" + "type": "regex", + "pattern": "[1-9][0-9]*" }, { - "type": "tag", - "begin": "<|open|>json type=\"object\"<|sep|>", - "content": { - "type": "json_schema", - "json_schema": { - "type": "object", - "properties": { - "location": { - "type": "string" - } + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>", + "content": { + "type": "any_text", + "excludes": [ + "<|close|>argument<|sep|>", + "<|close|>call<|sep|>" + ], + "max_tokens": null, + "max_chars": null }, - "required": [ - "location" - ] + "end": "<|close|>argument<|sep|>" }, - "style": "json", - "any_order": false, - "max_whitespace_cnt": null, - "excludes": [] - }, - "end": "<|close|>json<|sep|>" + { + "type": "tag", + "begin": "<|open|>json type=\"object\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>json<|sep|>" + } + ] } ] - } - ] - }, - "end": "<|close|>call<|sep|>" - } - ], - "separator": "", - "at_least_one": true, - "stop_after_first": false - }, - "end": "<|close|>tools<|sep|>" + }, + "end": "<|close|>call<|sep|>" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": true + }, + "end": "<|close|>tools<|sep|>" + } + ] }, { "type": "optional", @@ -563,7 +793,7 @@ } ] }, - "coverage": "final_output_only" + "coverage": "from_token_zero" }, "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] diff --git a/rust/src/chat/tests/grammar_replay/Kimi-K3.txt b/rust/src/chat/tests/grammar_replay/Kimi-K3.txt index 3aeb6126420e..6a1c814c131c 100644 --- a/rust/src/chat/tests/grammar_replay/Kimi-K3.txt +++ b/rust/src/chat/tests/grammar_replay/Kimi-K3.txt @@ -1,8 +1,29 @@ -# auto (final_output_only) +# auto (from_token_zero) sequence - optional `<|open|>response<|sep|>` - tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` - optional + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + or + sequence + tag `<|open|>response<|sep|>` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + optional + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` + tag `<|open|>call tool="add" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + sequence + tag `<|open|>argument key="y" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="x" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="items" type="array"<|sep|>` json(string[]) `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ y: number, x: number, items: string[] }) `<|close|>json<|sep|>` tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` tags_with_separator `` at_least_one tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` @@ -24,42 +45,44 @@ sequence tag `<|open|>json type="object"<|sep|>` json({ y: number, x: number, items: string[] }) `<|close|>json<|sep|>` optional `<|close|>message<|sep|>` -# required (final_output_only) +# required (from_token_zero) sequence - optional `<|open|>response<|sep|>` - tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` - tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` - tags_with_separator `` at_least_one - tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` - sequence - /[1-9][0-9]*/ - `"<|sep|>` - or - tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` - tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` - tag `<|open|>call tool="add" index="` .. `<|close|>call<|sep|>` - sequence - /[1-9][0-9]*/ - `"<|sep|>` - or - sequence - tag `<|open|>argument key="y" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` - tag `<|open|>argument key="x" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` - tag `<|open|>argument key="items" type="array"<|sep|>` json(string[]) `<|close|>argument<|sep|>` - tag `<|open|>json type="object"<|sep|>` json({ y: number, x: number, items: string[] }) `<|close|>json<|sep|>` + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + sequence + optional tag `<|open|>response<|sep|>` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` + tag `<|open|>call tool="add" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + sequence + tag `<|open|>argument key="y" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="x" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="items" type="array"<|sep|>` json(string[]) `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ y: number, x: number, items: string[] }) `<|close|>json<|sep|>` optional `<|close|>message<|sep|>` -# named (final_output_only) +# named (from_token_zero) sequence - optional `<|open|>response<|sep|>` - tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` - tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` - tags_with_separator `` at_least_one - tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` - sequence - /[1-9][0-9]*/ - `"<|sep|>` - or - tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` - tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + sequence + optional tag `<|open|>response<|sep|>` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one stop_after_first + tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` optional `<|close|>message<|sep|>` diff --git a/rust/src/parser/src/output_grammar.rs b/rust/src/parser/src/output_grammar.rs index 53c52538ee1d..0108c3193ef7 100644 --- a/rust/src/parser/src/output_grammar.rs +++ b/rust/src/parser/src/output_grammar.rs @@ -12,7 +12,8 @@ use xgrammar_structural_tag::builders::{StructuralTagBuilder, StructuralTagOptio use xgrammar_structural_tag::format::Format; use xgrammar_structural_tag::tool::ToolChoiceValue; use xgrammar_structural_tag::{ - FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag, + FunctionDefinition, FunctionToolParam, NormalizedToolChoice, ToolChoice, ToolParam, + build_structural_tag, }; use crate::tool::Tool; @@ -218,6 +219,22 @@ fn tool_params(tools: &[Tool], strict_level: ToolStrictLevel) -> Vec .collect() } +/// Normalize the request tools the way `build_structural_tag` does before +/// calling a crate builder, for model-owned grammar builders. Returns `None` +/// when the request asks for no tool grammar. +pub(crate) fn normalize_tool_choice( + ctx: &OutputGrammarContext<'_>, +) -> Result> { + if !tool_grammar_applies(ctx) { + return Ok(None); + } + + let tools = tool_params(ctx.tools, ctx.tool_strict_level); + let normalized = + xgrammar_structural_tag::normalize_tool_choice(&tools, ctx.tool_choice.clone())?; + Ok(Some(normalized)) +} + /// Whether the request asks for a tool grammar at all. fn tool_grammar_applies(ctx: &OutputGrammarContext<'_>) -> bool { if ctx.tools.is_empty() { diff --git a/rust/src/parser/src/unified/kimi_k3.rs b/rust/src/parser/src/unified/kimi_k3.rs index d6e32463c1fb..13d6fa42eb6f 100644 --- a/rust/src/parser/src/unified/kimi_k3.rs +++ b/rust/src/parser/src/unified/kimi_k3.rs @@ -46,9 +46,7 @@ //! now only applies to `argument` and `json` blocks inside a call body, which //! are still parsed from text. -mod structural_tag; - -pub use structural_tag::KimiK3StructuralTagBuilder; +mod grammar; use serde_json::{Map, Value}; use vllm_tokenizer::{DecodedText, DynTokenizer, Tokenizer}; @@ -58,11 +56,8 @@ use winnow::error::{ContextError, ErrMode, ModalResult, StrContext}; use winnow::prelude::*; use winnow::token::{literal, rest, take_till, take_until, take_while}; -use self::structural_tag::KIMI_K3_STRUCTURAL_TAG_BUILDER; use super::{Result, UnifiedParser, UnifiedParserOutput, special_token}; -use crate::output_grammar::{ - self, BuiltOutputGrammar, OutputGrammarContext, visible_format_from_builder, -}; +use crate::output_grammar::{self, BuiltOutputGrammar, OutputGrammarContext}; use crate::tool::{Tool, ToolCallDelta}; use crate::unified::parsing_failed; use crate::utils::{ @@ -347,10 +342,7 @@ impl UnifiedParser for KimiK3UnifiedParser { &self, ctx: &OutputGrammarContext<'_>, ) -> output_grammar::Result> { - Ok( - visible_format_from_builder(Some(&KIMI_K3_STRUCTURAL_TAG_BUILDER), ctx)? - .map(BuiltOutputGrammar::final_output_only), - ) + grammar::build_output_grammar(&self.mode, ctx) } fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> { diff --git a/rust/src/parser/src/unified/kimi_k3/grammar.rs b/rust/src/parser/src/unified/kimi_k3/grammar.rs new file mode 100644 index 000000000000..166fbafb0f82 --- /dev/null +++ b/rust/src/parser/src/unified/kimi_k3/grammar.rs @@ -0,0 +1,356 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Full-output grammar for Kimi K3 XTML generations. +//! +//! The grammar starts at the channel the prompt left open, which is the +//! parser's initialized [`KimiK3Mode`], and follows the same channel +//! transitions as the parser. A prompt that opened no channel may first emit a +//! whole `think` channel. +//! +//! | Tool choice | After `think` closes | After `response` closes | +//! | ----------- | ------------------------------ | ----------------------- | +//! | `auto` | `response [tools]`, or `tools` | `[tools]` | +//! | `required` | `[response] tools` | `tools` | +//! | named | `[response] tools`, one call | `tools`, one call | +//! +//! Every path may end with the message close, and a stop token is legal only +//! after the last channel closes. + +mod arguments; + +use xgrammar_structural_tag::NormalizedToolChoice; +use xgrammar_structural_tag::format::Format; +use xgrammar_structural_tag::tool::BuilderToolChoice; + +use self::arguments::call_tag; +use super::{ + CLOSE, END_OF_MSG, KimiK3Mode, MESSAGE_CLOSE, OPEN, RESPONSE_CLOSE, RESPONSE_OPEN, THINK_CLOSE, + THINK_OPEN, TOOLS_CLOSE, TOOLS_OPEN, +}; +use crate::output_grammar::{ + self, BuiltOutputGrammar, OutputGrammarContext, normalize_tool_choice, +}; + +/// Build the grammar of everything generated after a prompt that left the +/// parser in `mode`, or `None` when the request asks for no tool grammar. +pub(super) fn build_output_grammar( + mode: &KimiK3Mode, + ctx: &OutputGrammarContext<'_>, +) -> output_grammar::Result> { + let Some(NormalizedToolChoice { + function_tools, + choice, + .. + }) = normalize_tool_choice(ctx)? + else { + return Ok(None); + }; + + let single_call = choice == BuilderToolChoice::Forced || !ctx.parallel_tool_calls; + let tools = Format::tag( + TOOLS_OPEN, + Format::tags_with_separator( + function_tools.iter().map(call_tag).collect(), + "", + true, + single_call, + ), + TOOLS_CLOSE, + ); + + let mut elements = match mode { + KimiK3Mode::Reasoning => vec![think_body(), after_reasoning(choice, tools)], + KimiK3Mode::Response => vec![response_body(), after_response(choice, tools)], + KimiK3Mode::Idle => vec![ + Format::optional(Format::tag(THINK_OPEN, think_text(), THINK_CLOSE)), + after_reasoning(choice, tools), + ], + KimiK3Mode::Epilogue | KimiK3Mode::Tools | KimiK3Mode::Call { .. } | KimiK3Mode::Done => { + unreachable!("initialization leaves the parser idle or inside `think` or `response`") + } + }; + elements.push(Format::optional(Format::const_string(MESSAGE_CLOSE))); + + Ok(Some(BuiltOutputGrammar::from_token_zero(Format::sequence( + elements, + )))) +} + +/// Channels that may follow the closed `think` channel. +fn after_reasoning(tool_choice: BuilderToolChoice, tools: Format) -> Format { + let response = Format::tag(RESPONSE_OPEN, response_text(), RESPONSE_CLOSE); + match tool_choice { + BuilderToolChoice::Auto => Format::or(vec![ + Format::sequence(vec![response, Format::optional(tools.clone())]), + tools, + ]), + BuilderToolChoice::Required | BuilderToolChoice::Forced => { + Format::sequence(vec![Format::optional(response), tools]) + } + } +} + +/// Channels that may follow the closed `response` channel. +fn after_response(tool_choice: BuilderToolChoice, tools: Format) -> Format { + match tool_choice { + BuilderToolChoice::Auto => Format::optional(tools), + BuilderToolChoice::Required | BuilderToolChoice::Forced => tools, + } +} + +/// The rest of the `think` channel the prompt opened. +fn think_body() -> Format { + Format::tag("", think_text(), THINK_CLOSE) +} + +/// The rest of the `response` channel the prompt opened. +fn response_body() -> Format { + Format::tag("", response_text(), RESPONSE_CLOSE) +} + +fn think_text() -> Format { + Format::any_text_excluding(&[THINK_CLOSE, END_OF_MSG]) +} + +fn response_text() -> Format { + // Keep marker-looking prefixes out of response text so the grammar must + // resolve them through a valid channel boundary. + Format::any_text_excluding(&[OPEN, CLOSE, END_OF_MSG]) +} + +#[cfg(test)] +mod tests { + use expect_test::{Expect, expect}; + use serde_json::json; + use xgrammar_structural_tag::ToolChoice; + use xgrammar_structural_tag::format::{EndBoundary, TagFormat}; + + use super::super::CALL_CLOSE; + use super::*; + use crate::output_grammar::test_utils::outline; + use crate::output_grammar::{GrammarCoverage, ToolStrictLevel}; + use crate::tool::Tool; + + fn tools() -> Vec { + ["get_weather", "add"] + .map(|name| Tool { + name: name.to_string(), + description: None, + parameters: json!({ "type": "object", "properties": {} }), + strict: None, + defer_loading: None, + }) + .into() + } + + fn check( + mode: KimiK3Mode, + tool_choice: ToolChoice, + parallel_tool_calls: bool, + expected: Expect, + ) { + let tools = tools(); + let grammar = build_output_grammar( + &mode, + &OutputGrammarContext { + tools: &tools, + tool_choice: &tool_choice, + tool_strict_level: ToolStrictLevel::Function, + parallel_tool_calls, + }, + ) + .unwrap() + .unwrap(); + assert_eq!(grammar.coverage, GrammarCoverage::FromTokenZero); + let mut format = grammar.format; + elide_call_bodies(&mut format); + expected.assert_eq(&outline(&format)); + } + + /// Replace each call body with `..`, so that outlines show the channel + /// structure. The `arguments` tests cover call bodies. + fn elide_call_bodies(format: &mut Format) { + match format { + Format::Sequence(format) => format.elements.iter_mut().for_each(elide_call_bodies), + Format::Or(format) => format.elements.iter_mut().for_each(elide_call_bodies), + Format::Optional(format) => elide_call_bodies(&mut format.content), + Format::Tag(tag) => elide_tag(tag), + Format::TagsWithSeparator(format) => format.tags.iter_mut().for_each(elide_tag), + _ => {} + } + } + + fn elide_tag(tag: &mut TagFormat) { + if matches!(&tag.end, EndBoundary::Text(end) if end == CALL_CLOSE) { + *tag.content = Format::const_string(".."); + } else { + elide_call_bodies(&mut tag.content); + } + } + + #[test] + fn prompt_opened_think_channel() { + check( + KimiK3Mode::Reasoning, + ToolChoice::auto(), + true, + expect![[r#" + sequence + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + or + sequence + tag `<|open|>response<|sep|>` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + optional + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + check( + KimiK3Mode::Reasoning, + ToolChoice::required(), + true, + expect![[r#" + sequence + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + sequence + optional tag `<|open|>response<|sep|>` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + check( + KimiK3Mode::Reasoning, + ToolChoice::function("add"), + true, + expect![[r#" + sequence + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + sequence + optional tag `<|open|>response<|sep|>` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one stop_after_first + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + } + + #[test] + fn prompt_opened_response_channel() { + check( + KimiK3Mode::Response, + ToolChoice::auto(), + true, + expect![[r#" + sequence + tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + optional + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + check( + KimiK3Mode::Response, + ToolChoice::required(), + true, + expect![[r#" + sequence + tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + check( + KimiK3Mode::Response, + ToolChoice::function("add"), + true, + expect![[r#" + sequence + tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one stop_after_first + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + } + + #[test] + fn no_prompt_opened_channel() { + check( + KimiK3Mode::Idle, + ToolChoice::auto(), + true, + expect![[r#" + sequence + optional tag `<|open|>think<|sep|>` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + or + sequence + tag `<|open|>response<|sep|>` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + optional + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + } + + #[test] + fn serial_tool_calls_stop_after_one_call() { + check( + KimiK3Mode::Reasoning, + ToolChoice::required(), + false, + expect![[r#" + sequence + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + sequence + optional tag `<|open|>response<|sep|>` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one stop_after_first + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + } + + #[test] + fn non_strict_auto_builds_no_grammar() { + let tools = tools(); + let grammar = build_output_grammar( + &KimiK3Mode::Reasoning, + &OutputGrammarContext { + tools: &tools, + tool_choice: &ToolChoice::auto(), + tool_strict_level: ToolStrictLevel::Auto, + parallel_tool_calls: true, + }, + ) + .unwrap(); + assert_eq!(grammar, None); + } +} diff --git a/rust/src/parser/src/unified/kimi_k3/grammar/arguments.rs b/rust/src/parser/src/unified/kimi_k3/grammar/arguments.rs new file mode 100644 index 000000000000..49587535aadd --- /dev/null +++ b/rust/src/parser/src/unified/kimi_k3/grammar/arguments.rs @@ -0,0 +1,184 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Kimi K3 XTML tool-call grammar: call tags and their typed arguments. + +use serde_json::Value; +use xgrammar_structural_tag::format::{Format, JsonSchemaFormat, TagFormat}; +use xgrammar_structural_tag::tool::{FunctionToolParam, function_parameters}; + +use super::super::{ARG_CLOSE, ARG_OPEN, CALL_CLOSE, JSON_CLOSE, JSON_OPEN, OPEN, SEP}; +use crate::output_grammar::arguments::{ + self, ArgumentOptions, ArgumentSyntax, JsonType, ParameterKey, ValueOption, group_by, one_of, +}; + +/// XTML arguments: +/// `<|open|>argument key="KEY" type="TYPE"<|sep|>VALUE<|close|>argument<|sep|>`, +/// adjacent to each other. +/// +/// The parser decodes each value by its `type` attribute, so every value option +/// sits under the attribute of its own JSON type. String values are raw text +/// and every other value is JSON. +struct XtmlArguments; + +impl ArgumentSyntax for XtmlArguments { + fn separator(&self) -> Option { + None + } + + fn parameter(&self, key: ParameterKey<'_>, options: &[ValueOption<'_>]) -> Option { + let escaped; + let key = match key { + ParameterKey::Declared(key) => { + escaped = escape_attr_value(key); + ParameterKey::Declared(&escaped) + } + ParameterKey::Free => ParameterKey::Free, + }; + let arguments = group_by(options, |option| xtml_type(option.ty)) + .into_iter() + .filter_map(|(type_name, options)| { + let values = options + .into_iter() + .filter_map(|option| match option.ty { + // The parser ends the call at its close marker, so a + // string value may contain neither close marker. + JsonType::String => option.raw_string(&[ARG_CLOSE, CALL_CLOSE]), + _ => Some(option.json()), + }) + .collect::>(); + (!values.is_empty()).then(|| { + let suffix = format!("\" type=\"{type_name}\"{SEP}"); + key.tag( + &format!("{ARG_OPEN} key=\""), + &suffix, + one_of(values), + ARG_CLOSE, + ) + }) + }) + .collect::>(); + (!arguments.is_empty()).then(|| one_of(arguments)) + } +} + +/// The `type` attribute K3 renders for a value of type `ty`, as +/// `_xtml_type` in the checkpoint's `encoding_k3.py`: integers and floats are +/// both `number`. +fn xtml_type(ty: JsonType) -> &'static str { + match ty { + JsonType::String => "string", + JsonType::Integer | JsonType::Number => "number", + JsonType::Boolean => "boolean", + JsonType::Null => "null", + JsonType::Object => "object", + JsonType::Array => "array", + } +} + +/// Build the tag for one call to `tool`. +pub(super) fn call_tag(tool: &FunctionToolParam) -> TagFormat { + let parameters = function_parameters(&tool.function); + let typed_arguments = + arguments::arguments(¶meters, &XtmlArguments, &ArgumentOptions::default()); + let call_body = Format::or(vec![typed_arguments, raw_json_arguments(¶meters)]); + + TagFormat::new( + format!( + "{OPEN}call tool=\"{}\" index=\"", + escape_attr_value(&tool.function.name) + ), + Format::sequence(vec![ + Format::regex("[1-9][0-9]*"), + Format::const_string(format!("\"{SEP}")), + call_body, + ]), + CALL_CLOSE, + ) +} + +fn raw_json_arguments(parameters: &Value) -> Format { + Format::tag( + format!("{JSON_OPEN} type=\"object\"{SEP}"), + Format::JsonSchema(JsonSchemaFormat::new(parameters.clone())), + JSON_CLOSE, + ) +} + +fn escape_attr_value(value: &str) -> String { + value.replace('&', "&").replace('"', """) +} + +#[cfg(test)] +mod tests { + use expect_test::{Expect, expect}; + use serde_json::json; + use xgrammar_structural_tag::FunctionDefinition; + use xgrammar_structural_tag::format::Format; + use xgrammar_structural_tag::tool::FunctionToolParam; + + use crate::output_grammar::test_utils::outline; + + fn check(parameters: serde_json::Value, expected: Expect) { + let tool = FunctionToolParam::new( + FunctionDefinition::new("get_weather").with_parameters(parameters), + ); + expected.assert_eq(&outline(&Format::Tag(super::call_tag(&tool)))); + } + + #[test] + fn call_tag_matches_xtml_arguments() { + check( + json!({ + "$defs": { + "place": { "type": "object", "properties": { "city": { "type": "string" } } } + }, + "type": "object", + "properties": { + "unit": { "type": "string", "enum": ["celsius", "fahrenheit"] }, + "place": { "$ref": "#/$defs/place", "type": "object" } + }, + "required": ["place"] + }), + expect![[r#" + tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + sequence + optional + tag `<|open|>argument key="unit" type="string"<|sep|>` .. `<|close|>argument<|sep|>` + or + `celsius` + `fahrenheit` + tag `<|open|>argument key="place" type="object"<|sep|>` json({ city?: string } where place = { city?: string }) `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ unit?: "celsius" | "fahrenheit", place: place & object } where place = { city?: string }) `<|close|>json<|sep|>` + "#]], + ); + } + + #[test] + fn union_argument_content_matches_its_xtml_type() { + check( + json!({ + "type": "object", + "properties": { + "count": { "type": ["integer", "null"] } + } + }), + expect![[r#" + tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + optional + or + tag `<|open|>argument key="count" type="number"<|sep|>` json(integer) `<|close|>argument<|sep|>` + tag `<|open|>argument key="count" type="null"<|sep|>` json(null) `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ count?: integer | null }) `<|close|>json<|sep|>` + "#]], + ); + } +} diff --git a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs deleted file mode 100644 index 4754726c1b5d..000000000000 --- a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs +++ /dev/null @@ -1,343 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -//! Structural-tag grammar for Kimi K3 XTML tool calls. - -use serde_json::Value; -use xgrammar_structural_tag::Result; -use xgrammar_structural_tag::builders::{ - ReasoningMode, StructuralTagBuilder, StructuralTagContext, StructuralTagOptions, -}; -use xgrammar_structural_tag::format::{Format, JsonSchemaFormat, StructuralTag, TagFormat}; -use xgrammar_structural_tag::tool::{BuilderToolChoice, FunctionToolParam, function_parameters}; - -use crate::output_grammar::arguments::{ - self, ArgumentOptions, ArgumentSyntax, JsonType, ParameterKey, ValueOption, group_by, one_of, -}; - -use super::{ - ARG_CLOSE, ARG_OPEN, CALL_CLOSE, CLOSE, END_OF_MSG, JSON_CLOSE, JSON_OPEN, MESSAGE_CLOSE, OPEN, - RESPONSE_CLOSE, RESPONSE_OPEN, SEP, THINK_CLOSE, THINK_OPEN, TOOLS_CLOSE, TOOLS_OPEN, -}; - -pub(super) static KIMI_K3_STRUCTURAL_TAG_BUILDER: KimiK3StructuralTagBuilder = - KimiK3StructuralTagBuilder; - -/// XTML arguments: -/// `<|open|>argument key="KEY" type="TYPE"<|sep|>VALUE<|close|>argument<|sep|>`, -/// adjacent to each other. -/// -/// The parser decodes each value by its `type` attribute, so every value option -/// sits under the attribute of its own JSON type. String values are raw text -/// and every other value is JSON. -struct XtmlArguments; - -impl ArgumentSyntax for XtmlArguments { - fn separator(&self) -> Option { - None - } - - fn parameter(&self, key: ParameterKey<'_>, options: &[ValueOption<'_>]) -> Option { - let escaped; - let key = match key { - ParameterKey::Declared(key) => { - escaped = escape_attr_value(key); - ParameterKey::Declared(&escaped) - } - ParameterKey::Free => ParameterKey::Free, - }; - let arguments = group_by(options, |option| xtml_type(option.ty)) - .into_iter() - .filter_map(|(type_name, options)| { - let values = options - .into_iter() - .filter_map(|option| match option.ty { - // The parser ends the call at its close marker, so a - // string value may contain neither close marker. - JsonType::String => option.raw_string(&[ARG_CLOSE, CALL_CLOSE]), - _ => Some(option.json()), - }) - .collect::>(); - (!values.is_empty()).then(|| { - let suffix = format!("\" type=\"{type_name}\"{SEP}"); - key.tag( - &format!("{ARG_OPEN} key=\""), - &suffix, - one_of(values), - ARG_CLOSE, - ) - }) - }) - .collect::>(); - (!arguments.is_empty()).then(|| one_of(arguments)) - } -} - -/// The `type` attribute K3 renders for a value of type `ty`, as -/// `_xtml_type` in the checkpoint's `encoding_k3.py`: integers and floats are -/// both `number`. -fn xtml_type(ty: JsonType) -> &'static str { - match ty { - JsonType::String => "string", - JsonType::Integer | JsonType::Number => "number", - JsonType::Boolean => "boolean", - JsonType::Null => "null", - JsonType::Object => "object", - JsonType::Array => "array", - } -} - -/// Kimi K3 XTML structural-tag builder. -#[derive(Debug, Clone, Copy, Default)] -pub struct KimiK3StructuralTagBuilder; - -impl StructuralTagBuilder for KimiK3StructuralTagBuilder { - fn build(&self, ctx: StructuralTagContext<'_>) -> Result { - let mut elements = response_prefix(ctx.options.reasoning); - - let tools = match ctx.tool_choice { - // Serving lowering filters empty tools before calling the builder, - // while direct `build_structural_tag(..., auto, ...)` calls do not. - BuilderToolChoice::Auto if ctx.function_tools.is_empty() => None, - BuilderToolChoice::Auto => Some(Format::optional(tools_channel( - ctx.function_tools, - ctx.options, - ))), - BuilderToolChoice::Forced | BuilderToolChoice::Required => { - Some(tools_channel(ctx.function_tools, ctx.options)) - } - }; - if let Some(tools) = tools { - elements.push(tools); - } - elements.push(Format::optional(Format::const_string(MESSAGE_CLOSE))); - - Ok(StructuralTag::new(Format::sequence(elements))) - } -} - -fn response_prefix(reasoning: ReasoningMode) -> Vec { - let mut elements = Vec::new(); - if reasoning != ReasoningMode::Disabled { - let prefix = Format::tag( - if reasoning == ReasoningMode::Auto { - THINK_OPEN - } else { - "" - }, - Format::any_text_excluding(&[THINK_CLOSE, END_OF_MSG]), - THINK_CLOSE, - ); - elements.push(if reasoning == ReasoningMode::Auto { - Format::optional(prefix) - } else { - prefix - }); - elements.push(Format::const_string(RESPONSE_OPEN)); - } else { - elements.push(Format::optional(Format::const_string(RESPONSE_OPEN))); - } - elements.push(Format::tag( - "", - // Keep marker-looking prefixes out of response text so the grammar - // must resolve them through a valid channel boundary. - Format::any_text_excluding(&[OPEN, CLOSE, END_OF_MSG]), - RESPONSE_CLOSE, - )); - elements -} - -fn tools_channel(tools: &[FunctionToolParam], options: StructuralTagOptions) -> Format { - let calls = tools.iter().map(|tool| call_tag(tool, options)).collect(); - Format::tag( - TOOLS_OPEN, - Format::tags_with_separator(calls, "", true, false), - TOOLS_CLOSE, - ) -} - -fn call_tag(tool: &FunctionToolParam, options: StructuralTagOptions) -> TagFormat { - let parameters = function_parameters(&tool.function); - let typed_arguments = arguments::arguments( - ¶meters, - &XtmlArguments, - &ArgumentOptions { - any_order: options.any_order, - max_whitespace_cnt: options.max_whitespace_cnt, - excludes: vec![], - }, - ); - let call_body = Format::or(vec![ - typed_arguments, - raw_json_arguments(¶meters, options), - ]); - - TagFormat::new( - format!( - "{OPEN}call tool=\"{}\" index=\"", - escape_attr_value(&tool.function.name) - ), - Format::sequence(vec![ - Format::regex("[1-9][0-9]*"), - Format::const_string(format!("\"{SEP}")), - call_body, - ]), - CALL_CLOSE, - ) -} - -fn raw_json_arguments(parameters: &Value, options: StructuralTagOptions) -> Format { - Format::tag( - format!("{JSON_OPEN} type=\"object\"{SEP}"), - json_schema(parameters.clone(), options), - JSON_CLOSE, - ) -} - -fn json_schema(schema: Value, options: StructuralTagOptions) -> Format { - Format::JsonSchema( - JsonSchemaFormat::new(schema) - .with_any_order(options.any_order) - .with_max_whitespace_cnt(options.max_whitespace_cnt), - ) -} - -fn escape_attr_value(value: &str) -> String { - value.replace('&', "&").replace('"', """) -} - -#[cfg(test)] -mod tests { - use expect_test::expect; - use serde_json::json; - use xgrammar_structural_tag::builders::StructuralTagOptions; - use xgrammar_structural_tag::{ - FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag, - }; - - use super::KimiK3StructuralTagBuilder; - use crate::output_grammar::test_utils::outline; - - fn tool(name: &str, parameters: serde_json::Value) -> ToolParam { - ToolParam::Function(FunctionToolParam::new( - FunctionDefinition::new(name).with_parameters(parameters), - )) - } - - #[test] - fn required_structural_tag_matches_xtml_channels() { - let tools = vec![tool( - "get_weather", - json!({ - "$defs": { - "place": { "type": "object", "properties": { "city": { "type": "string" } } } - }, - "type": "object", - "properties": { - "unit": { "type": "string", "enum": ["celsius", "fahrenheit"] }, - "place": { "$ref": "#/$defs/place", "type": "object" } - }, - "required": ["place"] - }), - )]; - let tag = build_structural_tag( - KimiK3StructuralTagBuilder, - &tools, - ToolChoice::required(), - StructuralTagOptions::default().with_reasoning(false), - ) - .unwrap(); - - expect![[r#" - sequence - optional `<|open|>response<|sep|>` - tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` - tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` - tags_with_separator `` at_least_one - tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` - sequence - /[1-9][0-9]*/ - `"<|sep|>` - or - sequence - optional - tag `<|open|>argument key="unit" type="string"<|sep|>` .. `<|close|>argument<|sep|>` - or - `celsius` - `fahrenheit` - tag `<|open|>argument key="place" type="object"<|sep|>` json({ city?: string } where place = { city?: string }) `<|close|>argument<|sep|>` - tag `<|open|>json type="object"<|sep|>` json({ unit?: "celsius" | "fahrenheit", place: place & object } where place = { city?: string }) `<|close|>json<|sep|>` - optional `<|close|>message<|sep|>` - "#]].assert_eq(&outline(&tag.format)); - } - - #[test] - fn reasoning_grammar_starts_inside_prefilled_think_channel() { - let tools = vec![tool("ping", json!({ "type": "object", "properties": {} }))]; - let tag = build_structural_tag( - KimiK3StructuralTagBuilder, - &tools, - ToolChoice::auto(), - StructuralTagOptions::default().with_reasoning(true), - ) - .unwrap(); - let value = serde_json::to_value(tag).unwrap(); - - assert_eq!( - value["format"]["elements"][0]["end"], - "<|close|>think<|sep|>" - ); - assert_eq!( - value["format"]["elements"][1]["value"], - "<|open|>response<|sep|>" - ); - assert_eq!(value["format"]["elements"][3]["type"], "optional"); - } - - #[test] - fn forced_choice_keeps_only_the_named_tool() { - let tools = vec![ - tool("search", json!({ "type": "object" })), - tool("lookup", json!({ "type": "object" })), - ]; - let tag = build_structural_tag( - KimiK3StructuralTagBuilder, - &tools, - ToolChoice::function("lookup"), - StructuralTagOptions::default().with_reasoning(false), - ) - .unwrap() - .to_json_string() - .unwrap(); - - assert!(tag.contains("lookup")); - assert!(!tag.contains("search")); - } - - #[test] - fn union_argument_content_matches_its_xtml_type() { - let tools = vec![tool( - "set_count", - json!({ - "type": "object", - "properties": { - "count": { "type": ["integer", "null"] } - } - }), - )]; - let tag = build_structural_tag( - KimiK3StructuralTagBuilder, - &tools, - ToolChoice::required(), - StructuralTagOptions::default(), - ) - .unwrap() - .to_json_string() - .unwrap(); - - assert!(tag.contains(r#"type=\"number\""#)); - assert!(tag.contains(r#""json_schema":{"type":"integer"}"#), "{tag}"); - assert!(tag.contains(r#"type=\"null\""#)); - assert!(tag.contains(r#""json_schema":{"type":"null"}"#), "{tag}"); - } -} diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs index aa0e92e6e17f..493b7b3f1b18 100644 --- a/rust/src/parser/src/unified/mod.rs +++ b/rust/src/parser/src/unified/mod.rs @@ -15,7 +15,7 @@ pub use gemma4::Gemma4UnifiedParser; pub use hf::{HfTemplateError, HfUnifiedParser, ResponseTemplate}; pub use hy::{HyV3UnifiedParser, HyV4UnifiedParser}; pub use inkling::InklingUnifiedParser; -pub use kimi_k3::{KimiK3StructuralTagBuilder, KimiK3UnifiedParser}; +pub use kimi_k3::KimiK3UnifiedParser; use thiserror::Error; use thiserror_ext::Macro; use vllm_tokenizer::{DecodedText, DynTokenizer}; From 5a0906263209da2c1165ce82b9fe8e767aa89669 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Mon, 12 Oct 2026 17:43:48 +0800 Subject: [PATCH 287/624] [Rust Frontend] Compose answer constraints into parser-owned output grammars (#60781) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- rust/src/chat/src/output/default/mod.rs | 55 ++ rust/src/chat/src/output/structural_tag.rs | 165 +++++- rust/src/chat/tests/grammar_replay.py | 93 ++- .../grammar_replay/DeepSeek-V3.2-Exp.json | 211 ++++++- .../grammar_replay/DeepSeek-V3.2-Exp.txt | 17 + .../grammar_replay/DeepSeek-V4-Flash.json | 211 ++++++- .../grammar_replay/DeepSeek-V4-Flash.txt | 17 + .../grammar_replay/DeepSeek-V4.1-Flash.json | 211 ++++++- .../grammar_replay/DeepSeek-V4.1-Flash.txt | 17 + .../chat/tests/grammar_replay/GLM-4.5.json | 68 +++ .../src/chat/tests/grammar_replay/GLM-4.5.txt | 7 + .../tests/grammar_replay/GLM-4.7-Flash.json | 209 ++++++- .../tests/grammar_replay/GLM-4.7-Flash.txt | 14 + .../tests/grammar_replay/GLM-5.2-FP8.json | 209 ++++++- .../chat/tests/grammar_replay/GLM-5.2-FP8.txt | 14 + rust/src/chat/tests/grammar_replay/Hy3.json | 321 +++++++++- rust/src/chat/tests/grammar_replay/Hy3.txt | 50 ++ .../chat/tests/grammar_replay/Kimi-K2.5.json | 188 +++++- .../chat/tests/grammar_replay/Kimi-K2.5.txt | 20 + .../chat/tests/grammar_replay/Kimi-K3.json | 549 +++++++++++++++++- .../src/chat/tests/grammar_replay/Kimi-K3.txt | 55 ++ .../grammar_replay/MiMo-V2.6-Flash-RL.json | 212 ++++++- .../grammar_replay/MiMo-V2.6-Flash-RL.txt | 14 + .../tests/grammar_replay/MiniMax-M2.5.json | 229 +++++++- .../tests/grammar_replay/MiniMax-M2.5.txt | 21 + .../chat/tests/grammar_replay/MiniMax-M3.json | 40 ++ .../chat/tests/grammar_replay/MiniMax-M3.txt | 2 + .../chat/tests/grammar_replay/Qwen3-0.6B.json | 228 +++++++- .../chat/tests/grammar_replay/Qwen3-0.6B.txt | 20 + .../chat/tests/grammar_replay/Qwen3.5-4B.json | 222 ++++++- .../chat/tests/grammar_replay/Qwen3.5-4B.txt | 18 + .../grammar_replay/Seed-OSS-36B-Instruct.json | 40 ++ .../grammar_replay/Seed-OSS-36B-Instruct.txt | 2 + rust/src/chat/tests/roundtrip.rs | 188 ++++-- rust/src/parser/src/output_grammar.rs | 9 + rust/src/parser/src/tool/minimax_m2.rs | 2 +- rust/src/parser/src/tool/mod.rs | 69 ++- rust/src/parser/src/unified/combined.rs | 114 +++- rust/src/parser/src/unified/hy/hy_v3.rs | 2 + rust/src/parser/src/unified/hy/hy_v4.rs | 1 + .../src/parser/src/unified/kimi_k3/grammar.rs | 254 ++++++-- 41 files changed, 4133 insertions(+), 255 deletions(-) create mode 100644 rust/src/chat/tests/grammar_replay/GLM-4.5.json create mode 100644 rust/src/chat/tests/grammar_replay/GLM-4.5.txt create mode 100644 rust/src/chat/tests/grammar_replay/MiniMax-M3.json create mode 100644 rust/src/chat/tests/grammar_replay/MiniMax-M3.txt create mode 100644 rust/src/chat/tests/grammar_replay/Seed-OSS-36B-Instruct.json create mode 100644 rust/src/chat/tests/grammar_replay/Seed-OSS-36B-Instruct.txt diff --git a/rust/src/chat/src/output/default/mod.rs b/rust/src/chat/src/output/default/mod.rs index b31c5e9e4cbe..981173dc7450 100644 --- a/rust/src/chat/src/output/default/mod.rs +++ b/rust/src/chat/src/output/default/mod.rs @@ -13,8 +13,10 @@ use vllm_parser::output_grammar::{BuiltOutputGrammar, OutputGrammarContext}; use vllm_parser::unified::{CombinedParser, UnifiedParser}; use vllm_text::tokenizer::DynTokenizer; use xgrammar_structural_tag::ToolChoice; +use xgrammar_structural_tag::format::Format; use self::unified::unified_event_stream; +use super::structural_tag::answer_format; use super::structured::structured_chat_event_stream; use crate::error::Result; use crate::output::{ChatOutputProcessor, DynChatEventStream, DynDecodedTextEventStream}; @@ -48,6 +50,10 @@ struct GrammarInputs { tools: Vec, tool_choice: ToolChoice, tool_strict_level: ToolStrictLevel, + /// The request's structured-output constraint, normalized for composition + /// into the output grammar. The original stays on the request until an + /// output grammar replaces it. + answer: Option, } impl DefaultChatOutputProcessor { @@ -137,6 +143,7 @@ impl DefaultChatOutputProcessor { tools: request.tools().to_vec(), tool_choice: request.tool_choice().into(), tool_strict_level, + answer: request.sampling_params.structured_outputs.as_ref().and_then(answer_format), }), }) } @@ -258,6 +265,7 @@ impl ChatOutputProcessor for DefaultChatOutputProcessor { tool_choice: &inputs.tool_choice, tool_strict_level: inputs.tool_strict_level, parallel_tool_calls: self.parallel_tool_calls, + answer: inputs.answer.as_ref(), }) .map_err(|error| Error::OutputGrammar { error: Box::new(error), @@ -283,7 +291,9 @@ mod tests { use std::sync::Arc; use thiserror_ext::AsReport as _; + use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams; use vllm_tokenizer::test_utils::TestTokenizer; + use xgrammar_structural_tag::format::Format; use super::DefaultChatOutputProcessor; use crate::output::ChatOutputProcessor; @@ -339,6 +349,51 @@ mod tests { assert!(build(ToolStrictLevel::Parameter, None, ChatToolChoice::None).is_none()); } + #[test] + fn answer_constraint_keeps_optional_tool_calls() { + let tools = vec![ChatTool { + name: "search".to_string(), + description: None, + parameters: serde_json::json!({"type": "object"}), + strict: None, + defer_loading: None, + }]; + let schema = serde_json::json!({"type": "object"}); + let build = |structured_outputs| { + let mut request = ChatRequest { + tool_context: ResolvedToolContext::new( + &[], + tools.clone(), + Some(ChatToolChoice::Auto), + true, + ) + .unwrap(), + ..ChatRequest::for_test() + }; + request.sampling_params.structured_outputs = structured_outputs; + let mut processor = DefaultChatOutputProcessor::new( + &mut request, + "other-model", + tokenizer(), + &ParserSelection::Explicit("qwen3_coder".to_string()), + &ParserSelection::None, + ToolStrictLevel::Auto, + ) + .unwrap(); + processor.initialize(&[]).unwrap(); + processor.build_output_grammar().unwrap() + }; + + // Non-strict `auto` alone needs no grammar; with an answer constraint + // the grammar holds the answer or a call. + assert!(build(None).is_none()); + let built = build(Some(StructuredOutputsParams::json(schema.clone()))).unwrap(); + let Format::Or(branches) = built.format else { + panic!("expected the calls or the answer, got {:?}", built.format); + }; + assert_eq!(branches.elements.last(), Some(&Format::json_schema(schema))); + } + #[test] fn tool_grammar_honors_parallel_call_policy() { for parallel in [false, true] { diff --git a/rust/src/chat/src/output/structural_tag.rs b/rust/src/chat/src/output/structural_tag.rs index dba6b6d1b739..b99a32d5bd3d 100644 --- a/rust/src/chat/src/output/structural_tag.rs +++ b/rust/src/chat/src/output/structural_tag.rs @@ -3,19 +3,91 @@ //! Applies parser-built xgrammar structural-tag constraints. +use serde::Deserialize; +use serde_json::{Value, json}; use vllm_engine_core_client::protocol::structured_outputs::{ - StructuredOutputBackend, StructuredOutputsParams, + StructuredOutputBackend, StructuredOutputConstraint, StructuredOutputsParams, }; use vllm_parser::output_grammar::{BuiltOutputGrammar, GrammarCoverage}; use vllm_text::TextRequest; use xgrammar_structural_tag::StructuralTag; +use xgrammar_structural_tag::format::{Format, GrammarFormat, TagFormat}; use crate::{Error, Result}; +/// Normalize the user's structured-output constraint to the answer format a +/// parser may compose into its output grammar. +/// +/// Like Python's `structured_outputs_to_format`, this ignores the request's +/// options: no backend reads them per request, only from the engine config. +/// +/// Returns `None` when the constraint must stay on its own wire field, so the +/// request keeps today's behavior: a grammar that is not XGrammar EBNF (such +/// as Lark), or a structural tag that does not parse. +pub(crate) fn answer_format(params: &StructuredOutputsParams) -> Option { + Some(match ¶ms.constraint { + StructuredOutputConstraint::Json(Value::String(schema)) => { + Format::json_schema(serde_json::from_str(schema).ok()?) + } + StructuredOutputConstraint::Json(schema) => Format::json_schema(schema.clone()), + StructuredOutputConstraint::JsonObject => Format::json_schema(json!({ "type": "object" })), + StructuredOutputConstraint::Regex(pattern) => Format::regex(pattern), + StructuredOutputConstraint::Choice(choices) => { + Format::or(choices.iter().map(Format::const_string).collect()) + } + // Lark rules use `:`; XGrammar EBNF rules use `::=`. + StructuredOutputConstraint::Grammar(grammar) if grammar.contains("::=") => { + Format::Grammar(GrammarFormat { + grammar: grammar.clone(), + }) + } + StructuredOutputConstraint::Grammar(_) => return None, + StructuredOutputConstraint::StructuralTag(tag) => structural_tag_format(tag)?, + }) +} + +/// The legacy structural tag shape: free text with JSON tags behind triggers. +#[derive(Deserialize)] +struct LegacyStructuralTag { + structures: Vec, + triggers: Vec, +} + +#[derive(Deserialize)] +struct LegacyStructure { + begin: String, + schema: Value, + end: String, +} + +/// Read a structural tag in its current (`format`) or legacy +/// (`structures`/`triggers`) shape. +fn structural_tag_format(tag: &str) -> Option { + let value: Value = serde_json::from_str(tag).ok()?; + if value.get("structures").is_none() { + return Some(serde_json::from_value::(value).ok()?.format); + } + let legacy: LegacyStructuralTag = serde_json::from_value(value).ok()?; + let triggers: Vec<&str> = legacy.triggers.iter().map(String::as_str).collect(); + let tags = legacy + .structures + .into_iter() + .map(|structure| { + TagFormat::new( + structure.begin, + Format::json_schema(structure.schema), + structure.end, + ) + }) + .collect(); + Some(Format::triggered_tags(&triggers, tags)) +} + /// Apply one parser-built output grammar to the prepared text request. /// -/// A parser-owned grammar replaces any structured output constraint the user -/// supplied, matching the Python frontend. `None` leaves the request untouched. +/// The parser has already composed the user's answer constraint into the +/// grammar where its protocol allows, so the grammar replaces the request's +/// structured outputs. `None` leaves the request untouched. pub(crate) fn apply_output_grammar( request: &mut TextRequest, built: Option, @@ -45,13 +117,92 @@ pub(crate) fn apply_output_grammar( #[cfg(test)] mod tests { use serde_json::json; - use vllm_engine_core_client::protocol::structured_outputs::{ - StructuredOutputBackend, StructuredOutputsParams, - }; - use xgrammar_structural_tag::format::Format; + use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputOptions; use super::*; + #[test] + fn answer_format_normalizes_composable_constraints() { + let schema = json!({ "type": "object", "properties": { "a": { "type": "integer" } } }); + let cases = [ + ( + StructuredOutputsParams::json(schema.clone()), + Format::json_schema(schema.clone()), + ), + ( + StructuredOutputsParams::json(Value::String(schema.to_string())), + Format::json_schema(schema), + ), + ( + StructuredOutputsParams::json_object(), + Format::json_schema(json!({ "type": "object" })), + ), + ( + StructuredOutputsParams::regex("[a-z]+"), + Format::regex("[a-z]+"), + ), + ( + StructuredOutputsParams::choice(vec!["yes".to_string(), "no".to_string()]), + Format::or(vec![ + Format::const_string("yes"), + Format::const_string("no"), + ]), + ), + ( + StructuredOutputsParams::grammar("root ::= \"a\""), + Format::Grammar(GrammarFormat { + grammar: "root ::= \"a\"".to_string(), + }), + ), + ( + StructuredOutputsParams::structural_tag( + r#"{"type": "structural_tag", "format": {"type": "const_string", "value": "x"}}"#, + ), + Format::const_string("x"), + ), + ( + StructuredOutputsParams::structural_tag( + r#"{"type": "structural_tag", "structures": [{"begin": "", "schema": {"type": "object"}, "end": ""}], "triggers": ["", + Format::json_schema(json!({ "type": "object" })), + "", + )], + ), + ), + ( + StructuredOutputsParams { + options: StructuredOutputOptions { + disable_any_whitespace: true, + disable_additional_properties: true, + whitespace_pattern: Some(" ?".to_string()), + }, + ..StructuredOutputsParams::json_object() + }, + Format::json_schema(json!({ "type": "object" })), + ), + ]; + for (params, expected) in cases { + assert_eq!(answer_format(¶ms), Some(expected), "{params:?}"); + } + } + + #[test] + fn answer_format_leaves_other_constraints_on_their_wire_field() { + let cases = [ + StructuredOutputsParams::json(Value::String("{".to_string())), + StructuredOutputsParams::grammar("start: \"a\""), + StructuredOutputsParams::structural_tag(r#"{"structures": [{"begin": ""}]}"#), + StructuredOutputsParams::structural_tag(r#"{"format": {"type": "unknown"}}"#), + ]; + for params in cases { + assert_eq!(answer_format(¶ms), None, "{params:?}"); + } + } + #[test] fn output_grammar_overwrites_answer_constraint() { let mut request = TextRequest::for_test(); diff --git a/rust/src/chat/tests/grammar_replay.py b/rust/src/chat/tests/grammar_replay.py index d34df627bae0..056ab1d35289 100755 --- a/rust/src/chat/tests/grammar_replay.py +++ b/rust/src/chat/tests/grammar_replay.py @@ -22,14 +22,17 @@ # /// """Replay the grammar cases exported by the roundtrip tests through XGrammar. -Each file holds one model's cases, keyed by tool-choice variant. Every case is -compiled against the model tokenizer the same way as -`vllm/v1/structured_output/backend_xgrammar.py`. Token-zero cases are then -replayed from the first generated token: each token must be allowed by the -bitmask and accepted by the matcher, and the final stop token must terminate it. +Each file holds one model's distinct output grammars, each with the generations +of the fixture variants that build it. Every grammar is compiled against the +model tokenizer the same way as `vllm/v1/structured_output/backend_xgrammar.py`. +The generations of token-zero grammars are then replayed from the first +generated token: each token must be allowed by the bitmask and accepted by the +matcher, and the final stop token must terminate it. """ import json +import multiprocessing +import os import sys from functools import cache from pathlib import Path @@ -49,23 +52,29 @@ def load_tokenizer(model_id: str, vocab_size: int): return tokenizer, xgr.GrammarCompiler(tokenizer_info) -def replay(model: dict, case: dict) -> str: - """Check one case of a model and return a summary of what was checked.""" - tokenizer, compiler = load_tokenizer(model["model_id"], model["vocab_size"]) - generation = case["generation_token_ids"] +def compile_grammar(model: dict, grammar: dict) -> xgr.CompiledGrammar: + """Compile one output grammar of a model.""" + _, compiler = load_tokenizer(model["model_id"], model["vocab_size"]) + structural_tag = {"type": "structural_tag", "format": grammar["format"]} + return compiler.compile_structural_tag(json.dumps(structural_tag)) + + +def replay( + model: dict, grammar: dict, compiled: xgr.CompiledGrammar, generation: dict +) -> str: + """Check one generation of a grammar and return a summary of what was checked.""" + tokenizer, _ = load_tokenizer(model["model_id"], model["vocab_size"]) + token_ids = generation["generation_token_ids"] - # The case token IDs come from the Rust tokenizer; the runner must load the + # The recorded token IDs come from the Rust tokenizer; the runner must load the # same vocabulary for the grammar to mean the same thing. - encoded = tokenizer.encode(case["completion"], add_special_tokens=False) - if encoded != generation[:-1]: + encoded = tokenizer.encode(generation["completion"], add_special_tokens=False) + if encoded != token_ids[:-1]: raise AssertionError( "tokenizer mismatch: the completion encodes to " - f"{encoded}, but the case records {generation[:-1]}" + f"{encoded}, but the generation records {token_ids[:-1]}" ) - grammar = case["grammar"] - structural_tag = {"type": "structural_tag", "format": grammar["format"]} - compiled = compiler.compile_structural_tag(json.dumps(structural_tag)) if grammar["coverage"] != "from_token_zero": return f"compiled only ({grammar['coverage']} is gated by the engine)" @@ -73,18 +82,41 @@ def replay(model: dict, case: dict) -> str: compiled, override_stop_tokens=model["all_stop_token_ids"] ) bitmask = xgr.allocate_token_bitmask(1, model["vocab_size"]) - for index, token_id in enumerate(generation): + for index, token_id in enumerate(token_ids): matcher.fill_next_token_bitmask(bitmask) allowed = (bitmask[0, token_id // 32].item() >> (token_id % 32)) & 1 if not allowed or not matcher.accept_token(token_id): - prefix = tokenizer.decode(generation[:index]) + prefix = tokenizer.decode(token_ids[:index]) raise AssertionError( f"token {index} ({token_id}, {tokenizer.decode([token_id])!r}) " f"rejected after {prefix!r}" ) if not matcher.is_terminated(): raise AssertionError("matcher not terminated after the stop token") - return f"replayed {len(generation)} tokens from token zero" + return f"replayed {len(token_ids)} tokens from token zero" + + +def check_model(path: Path) -> tuple[list[str], list[str]]: + """Check every grammar of one model file and return its stdout and stderr lines.""" + model = json.loads(path.read_text()) + out, err = [], [] + for case in model["grammars"]: + grammar = case["grammar"] + coverage = grammar["coverage"] + try: + compiled = compile_grammar(model, grammar) + except Exception as error: + names = ", ".join(case["generations"]) + err.append(f"FAIL {path.name} {names} [{coverage}]: {error}") + continue + for name, generation in case["generations"].items(): + try: + summary = replay(model, grammar, compiled, generation) + except Exception as error: + err.append(f"FAIL {path.name} {name} [{coverage}]: {error}") + else: + out.append(f"PASS {path.name} {name}: {summary}") + return out, err def main() -> int: @@ -93,19 +125,18 @@ def main() -> int: print(f"no grammar replay cases in {CASES_DIR}", file=sys.stderr) return 1 - failures = 0 - for path in paths: - model = json.loads(path.read_text()) - for name, case in model["cases"].items(): - try: - summary = replay(model, case) - except Exception as error: - failures += 1 - coverage = case["grammar"]["coverage"] - print(f"FAIL {path.name} {name} [{coverage}]: {error}", file=sys.stderr) - else: - print(f"PASS {path.name} {name}: {summary}") + # The cost is dominated by loading each model's tokenizer, which no other model + # shares, so check the model files in parallel. Fork after the imports above so + # that the workers share them instead of paying for them again. + workers = min(len(paths), os.cpu_count() or 1) + with multiprocessing.get_context("fork").Pool(workers) as pool: + results = pool.map(check_model, paths, chunksize=1) + failures = 0 + for out, err in results: + print("\n".join(out)) + print("\n".join(err), file=sys.stderr) + failures += len(err) return 1 if failures else 0 diff --git a/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.json b/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.json index 4903b5900c74..f08b418bf0e6 100644 --- a/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.json +++ b/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.json @@ -2,8 +2,8 @@ "model_id": "deepseek-ai/DeepSeek-V3.2-Exp", "vocab_size": 129280, "all_stop_token_ids": [1], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -107,10 +107,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.I will call the tools.\n\n<｜DSML｜function_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", - "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128799, 43, 759, 1926, 270, 6704, 339, 30, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 3327, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 1718, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 32, 1] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.I will call the tools.\n\n<｜DSML｜function_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128799, 43, 759, 1926, 270, 6704, 339, 30, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 3327, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 1718, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 32, 1] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -208,10 +212,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n<｜DSML｜function_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", - "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128799, 271, 30, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 3327, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 1718, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 32, 1] + "generations": { + "required": { + "completion": "Need call the weather and add tools.\n\n<｜DSML｜function_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128799, 271, 30, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 3327, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 1718, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 32, 1] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -267,8 +275,189 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n<｜DSML｜function_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n", - "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128799, 271, 30, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 1718, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 32, 1] + "generations": { + "named": { + "completion": "Need call the weather and add tools.\n\n<｜DSML｜function_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128799, 271, 30, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 1718, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 32, 1] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128799, 24313, 19995, 16562, 69958, 2928, 62773, 1] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "\n\n<｜DSML｜function_calls>\n" + }, + { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<｜DSML｜invoke name=\"get_weather\">\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "deepseek_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n" + }, + { + "type": "tag", + "begin": "<｜DSML｜invoke name=\"add\">\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "deepseek_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + { + "type": "const_string", + "value": "" + } + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_content": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128799, 24313, 19995, 16562, 69958, 2928, 62773, 1] + }, + "auto_schema_calls": { + "completion": "Need call the weather and add tools.\n\n<｜DSML｜function_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128799, 271, 30, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 3327, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 1718, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 32, 1] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.txt b/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.txt index 5cf40f3218e8..25b572b8a247 100644 --- a/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.txt +++ b/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.txt @@ -24,3 +24,20 @@ sequence `\n\n<｜DSML｜function_calls>\n` tag `<｜DSML｜invoke name="get_weather">\n` deepseek_xml({ location: string }) `\n` `` + +# none_schema (from_token_zero) +sequence + tag `` text `` + json({ answer: string }) + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + tag `` text `` + or + sequence + `\n\n<｜DSML｜function_calls>\n` + tags_with_separator `` at_least_one + tag `<｜DSML｜invoke name="get_weather">\n` deepseek_xml({ location: string }) `\n` + tag `<｜DSML｜invoke name="add">\n` deepseek_xml({ y: number, x: number, items: string[] }) `\n` + `` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.json b/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.json index 1a286d8acda9..8ab51cbbf991 100644 --- a/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.json +++ b/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.json @@ -2,8 +2,8 @@ "model_id": "deepseek-ai/DeepSeek-V4-Flash", "vocab_size": 129280, "all_stop_token_ids": [1], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -107,10 +107,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.I will call the tools.\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", - "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 43, 759, 1926, 270, 6704, 339, 30, 128825, 72461, 4941, 12548, 1018, 30, 128825, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 30, 128825, 40148, 5406, 2329, 1281, 3327, 3816, 30, 128825, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 1718, 128825, 72461, 4941, 12548, 32, 1] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.I will call the tools.\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 43, 759, 1926, 270, 6704, 339, 30, 128825, 72461, 4941, 12548, 1018, 30, 128825, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 30, 128825, 40148, 5406, 2329, 1281, 3327, 3816, 30, 128825, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 1718, 128825, 72461, 4941, 12548, 32, 1] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -208,10 +212,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", - "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 72461, 4941, 12548, 1018, 30, 128825, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 30, 128825, 40148, 5406, 2329, 1281, 3327, 3816, 30, 128825, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 1718, 128825, 72461, 4941, 12548, 32, 1] + "generations": { + "required": { + "completion": "Need call the weather and add tools.\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 72461, 4941, 12548, 1018, 30, 128825, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 30, 128825, 40148, 5406, 2329, 1281, 3327, 3816, 30, 128825, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 1718, 128825, 72461, 4941, 12548, 32, 1] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -267,8 +275,189 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n", - "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 72461, 4941, 12548, 1018, 30, 128825, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 1718, 128825, 72461, 4941, 12548, 32, 1] + "generations": { + "named": { + "completion": "Need call the weather and add tools.\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 72461, 4941, 12548, 1018, 30, 128825, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 1718, 128825, 72461, 4941, 12548, 32, 1] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 24313, 19995, 16562, 69958, 2928, 62773, 1] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "\n\n<｜DSML｜tool_calls>\n" + }, + { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<｜DSML｜invoke name=\"get_weather\">\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "deepseek_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n" + }, + { + "type": "tag", + "begin": "<｜DSML｜invoke name=\"add\">\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "deepseek_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + { + "type": "const_string", + "value": "" + } + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_content": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 24313, 19995, 16562, 69958, 2928, 62773, 1] + }, + "auto_schema_calls": { + "completion": "Need call the weather and add tools.\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 72461, 4941, 12548, 1018, 30, 128825, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 30, 128825, 40148, 5406, 2329, 1281, 3327, 3816, 30, 128825, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 1718, 128825, 72461, 4941, 12548, 32, 1] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.txt b/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.txt index 5e216e8ffcf2..535da9ba2876 100644 --- a/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.txt +++ b/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.txt @@ -24,3 +24,20 @@ sequence `\n\n<｜DSML｜tool_calls>\n` tag `<｜DSML｜invoke name="get_weather">\n` deepseek_xml({ location: string }) `\n` `` + +# none_schema (from_token_zero) +sequence + tag `` text `` + json({ answer: string }) + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + tag `` text `` + or + sequence + `\n\n<｜DSML｜tool_calls>\n` + tags_with_separator `` at_least_one + tag `<｜DSML｜invoke name="get_weather">\n` deepseek_xml({ location: string }) `\n` + tag `<｜DSML｜invoke name="add">\n` deepseek_xml({ y: number, x: number, items: string[] }) `\n` + `` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.json b/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.json index a21d8b90e5af..a47ce49c2f87 100644 --- a/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.json +++ b/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.json @@ -2,8 +2,8 @@ "model_id": "deepseek-ai/DeepSeek-V4.1-Flash", "vocab_size": 129280, "all_stop_token_ids": [1], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -107,10 +107,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.I will call the tools.\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"get_weather\">\n<｜DSML｜ parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜ invoke name=\"add\">\n<｜DSML｜ parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜ parameter name=\"x\" string=\"false\">2\n<｜DSML｜ parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", - "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 43, 759, 1926, 270, 6704, 339, 30, 128825, 10699, 1018, 30, 128825, 34756, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 10767, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 30, 128825, 34756, 2329, 1281, 3327, 3816, 30, 128825, 10767, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 1718, 128825, 10699, 32, 1] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.I will call the tools.\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"get_weather\">\n<｜DSML｜ parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜ invoke name=\"add\">\n<｜DSML｜ parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜ parameter name=\"x\" string=\"false\">2\n<｜DSML｜ parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 43, 759, 1926, 270, 6704, 339, 30, 128825, 10699, 1018, 30, 128825, 34756, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 10767, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 30, 128825, 34756, 2329, 1281, 3327, 3816, 30, 128825, 10767, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 1718, 128825, 10699, 32, 1] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -208,10 +212,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"get_weather\">\n<｜DSML｜ parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜ invoke name=\"add\">\n<｜DSML｜ parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜ parameter name=\"x\" string=\"false\">2\n<｜DSML｜ parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", - "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 10699, 1018, 30, 128825, 34756, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 10767, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 30, 128825, 34756, 2329, 1281, 3327, 3816, 30, 128825, 10767, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 1718, 128825, 10699, 32, 1] + "generations": { + "required": { + "completion": "Need call the weather and add tools.\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"get_weather\">\n<｜DSML｜ parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜ invoke name=\"add\">\n<｜DSML｜ parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜ parameter name=\"x\" string=\"false\">2\n<｜DSML｜ parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 10699, 1018, 30, 128825, 34756, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 10767, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 30, 128825, 34756, 2329, 1281, 3327, 3816, 30, 128825, 10767, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 1718, 128825, 10699, 32, 1] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -267,8 +275,189 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"get_weather\">\n<｜DSML｜ parameter name=\"location\" string=\"true\">Shanghai\n\n", - "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 10699, 1018, 30, 128825, 34756, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 10767, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 1718, 128825, 10699, 32, 1] + "generations": { + "named": { + "completion": "Need call the weather and add tools.\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"get_weather\">\n<｜DSML｜ parameter name=\"location\" string=\"true\">Shanghai\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 10699, 1018, 30, 128825, 34756, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 10767, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 1718, 128825, 10699, 32, 1] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 24313, 19995, 16562, 69958, 2928, 62773, 1] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "\n\n<｜DSML｜ calls>\n" + }, + { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<｜DSML｜ invoke name=\"get_weather\">\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "deepseek_v4_1_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n" + }, + { + "type": "tag", + "begin": "<｜DSML｜ invoke name=\"add\">\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "deepseek_v4_1_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + { + "type": "const_string", + "value": "" + } + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_content": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 24313, 19995, 16562, 69958, 2928, 62773, 1] + }, + "auto_schema_calls": { + "completion": "Need call the weather and add tools.\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"get_weather\">\n<｜DSML｜ parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜ invoke name=\"add\">\n<｜DSML｜ parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜ parameter name=\"x\" string=\"false\">2\n<｜DSML｜ parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 10699, 1018, 30, 128825, 34756, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 10767, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 30, 128825, 34756, 2329, 1281, 3327, 3816, 30, 128825, 10767, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 1718, 128825, 10699, 32, 1] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.txt b/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.txt index 8c92c3a5a374..cb2c646545df 100644 --- a/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.txt +++ b/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.txt @@ -24,3 +24,20 @@ sequence `\n\n<｜DSML｜ calls>\n` tag `<｜DSML｜ invoke name="get_weather">\n` deepseek_v4_1_xml({ location: string }) `\n` `` + +# none_schema (from_token_zero) +sequence + tag `` text `` + json({ answer: string }) + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + tag `` text `` + or + sequence + `\n\n<｜DSML｜ calls>\n` + tags_with_separator `` at_least_one + tag `<｜DSML｜ invoke name="get_weather">\n` deepseek_v4_1_xml({ location: string }) `\n` + tag `<｜DSML｜ invoke name="add">\n` deepseek_v4_1_xml({ y: number, x: number, items: string[] }) `\n` + `` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/GLM-4.5.json b/rust/src/chat/tests/grammar_replay/GLM-4.5.json new file mode 100644 index 000000000000..e00f05757bb2 --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/GLM-4.5.json @@ -0,0 +1,68 @@ +{ + "model_id": "zai-org/GLM-4.5", + "vocab_size": 151552, + "all_stop_token_ids": [151329, 151336, 151338], + "grammars": [ + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "optional", + "content": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "\n", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "const_string", + "value": "\n" + } + ] + } + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "\nNeed call the weather and add tools.\n{\"answer\":\"sunny\"}", + "generation_token_ids": [198, 151350, 23591, 1618, 279, 9101, 323, 912, 7375, 13, 151351, 198, 4913, 9214, 3252, 82, 27208, 9204, 151329] + }, + "auto_schema_content": { + "completion": "\nNeed call the weather and add tools.\n{\"answer\":\"sunny\"}", + "generation_token_ids": [198, 151350, 23591, 1618, 279, 9101, 323, 912, 7375, 13, 151351, 198, 4913, 9214, 3252, 82, 27208, 9204, 151329] + } + } + } + ] +} diff --git a/rust/src/chat/tests/grammar_replay/GLM-4.5.txt b/rust/src/chat/tests/grammar_replay/GLM-4.5.txt new file mode 100644 index 000000000000..8a37a06dbda5 --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/GLM-4.5.txt @@ -0,0 +1,7 @@ +# none_schema, auto_schema_content (from_token_zero) +sequence + optional + sequence + tag `\n` text `` + `\n` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.json b/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.json index b289ec62e62b..6fd2a440e4d4 100644 --- a/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.json +++ b/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.json @@ -2,8 +2,8 @@ "model_id": "zai-org/GLM-4.7-Flash", "vocab_size": 154880, "all_stop_token_ids": [154820, 154827, 154829], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -99,10 +99,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.I will call the tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", - "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 40, 686, 1618, 279, 7375, 13, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.I will call the tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 40, 686, 1618, 279, 7375, 13, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -198,10 +202,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", - "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + "generations": { + "required": { + "completion": "Need call the weather and add tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -244,8 +252,187 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.get_weatherlocationShanghai", - "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154820] + "generations": { + "named": { + "completion": "Need call the weather and add tools.get_weatherlocationShanghai", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154820] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 4913, 9214, 3252, 82, 27208, 9204, 154820] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "or", + "elements": [ + { + "type": "triggered_tags", + "triggers": [ + "" + ], + "tags": [ + { + "type": "tag", + "begin": "get_weather", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "glm_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "" + }, + { + "type": "tag", + "begin": "add", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "glm_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "" + } + ], + "at_least_one": true, + "stop_after_first": false, + "excludes": [ + "", + "", + "", + "", + "", + "", + "" + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_content": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 4913, 9214, 3252, 82, 27208, 9204, 154820] + }, + "auto_schema_calls": { + "completion": "Need call the weather and add tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.txt b/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.txt index 58b4bef2da32..bc58a47fc936 100644 --- a/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.txt +++ b/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.txt @@ -16,3 +16,17 @@ sequence sequence tag `` text `` tag `get_weather` glm_xml({ location: string }) `` + +# none_schema (from_token_zero) +sequence + tag `` text `` + json({ answer: string }) + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + tag `` text `` + or + triggered_tags [``] excluding [``, ``, ``, ``, ``, ``, ``] at_least_one + tag `get_weather` glm_xml({ location: string }) `` + tag `add` glm_xml({ y: number, x: number, items: string[] }) `` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.json b/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.json index 404667c30a7b..ecf8880406de 100644 --- a/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.json +++ b/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.json @@ -2,8 +2,8 @@ "model_id": "zai-org/GLM-5.2-FP8", "vocab_size": 154880, "all_stop_token_ids": [154820, 154827, 154829], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -99,10 +99,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.I will call the tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", - "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 40, 686, 1618, 279, 7375, 13, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.I will call the tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 40, 686, 1618, 279, 7375, 13, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -198,10 +202,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", - "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + "generations": { + "required": { + "completion": "Need call the weather and add tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -244,8 +252,187 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.get_weatherlocationShanghai", - "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154820] + "generations": { + "named": { + "completion": "Need call the weather and add tools.get_weatherlocationShanghai", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154820] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 4913, 9214, 3252, 82, 27208, 9204, 154820] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "or", + "elements": [ + { + "type": "triggered_tags", + "triggers": [ + "" + ], + "tags": [ + { + "type": "tag", + "begin": "get_weather", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "glm_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "" + }, + { + "type": "tag", + "begin": "add", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "glm_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "" + } + ], + "at_least_one": true, + "stop_after_first": false, + "excludes": [ + "", + "", + "", + "", + "", + "", + "" + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_content": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 4913, 9214, 3252, 82, 27208, 9204, 154820] + }, + "auto_schema_calls": { + "completion": "Need call the weather and add tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.txt b/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.txt index 58b4bef2da32..bc58a47fc936 100644 --- a/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.txt +++ b/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.txt @@ -16,3 +16,17 @@ sequence sequence tag `` text `` tag `get_weather` glm_xml({ location: string }) `` + +# none_schema (from_token_zero) +sequence + tag `` text `` + json({ answer: string }) + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + tag `` text `` + or + triggered_tags [``] excluding [``, ``, ``, ``, ``, ``, ``] at_least_one + tag `get_weather` glm_xml({ location: string }) `` + tag `add` glm_xml({ y: number, x: number, items: string[] }) `` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/Hy3.json b/rust/src/chat/tests/grammar_replay/Hy3.json index f6ae7b0969fd..da089616300a 100644 --- a/rust/src/chat/tests/grammar_replay/Hy3.json +++ b/rust/src/chat/tests/grammar_replay/Hy3.json @@ -2,8 +2,8 @@ "model_id": "tencent/Hy3", "vocab_size": 120832, "all_stop_token_ids": [120025], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -263,10 +263,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.I will call the tools.\nget_weather\nlocation\nShanghai\n\nadd\ny\n1.0\nx\n2\nitems\n[\"left\", \"right\"]\n\n", - "generation_token_ids": [47893, 1609, 252, 9969, 289, 1189, 7106, 13, 120030, 40, 823, 1609, 252, 7106, 13, 120035, 185, 120033, 874, 62, 39899, 120043, 185, 120066, 24806, 120067, 185, 120068, 4169, 768, 32381, 120069, 185, 120034, 185, 120033, 2176, 120043, 185, 120066, 88, 120067, 185, 120068, 16, 13, 15, 120069, 185, 120066, 87, 120067, 185, 120068, 17, 120069, 185, 120066, 12792, 120067, 185, 120068, 5831, 1639, 1221, 478, 1331, 7029, 120069, 185, 120034, 185, 120036, 120025] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.I will call the tools.\nget_weather\nlocation\nShanghai\n\nadd\ny\n1.0\nx\n2\nitems\n[\"left\", \"right\"]\n\n", + "generation_token_ids": [47893, 1609, 252, 9969, 289, 1189, 7106, 13, 120030, 40, 823, 1609, 252, 7106, 13, 120035, 185, 120033, 874, 62, 39899, 120043, 185, 120066, 24806, 120067, 185, 120068, 4169, 768, 32381, 120069, 185, 120034, 185, 120033, 2176, 120043, 185, 120066, 88, 120067, 185, 120068, 16, 13, 15, 120069, 185, 120066, 87, 120067, 185, 120068, 17, 120069, 185, 120066, 12792, 120067, 185, 120068, 5831, 1639, 1221, 478, 1331, 7029, 120069, 185, 120034, 185, 120036, 120025] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -523,10 +527,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\nget_weather\nlocation\nShanghai\n\nadd\ny\n1.0\nx\n2\nitems\n[\"left\", \"right\"]\n\n", - "generation_token_ids": [47893, 1609, 252, 9969, 289, 1189, 7106, 13, 120030, 120035, 185, 120033, 874, 62, 39899, 120043, 185, 120066, 24806, 120067, 185, 120068, 4169, 768, 32381, 120069, 185, 120034, 185, 120033, 2176, 120043, 185, 120066, 88, 120067, 185, 120068, 16, 13, 15, 120069, 185, 120066, 87, 120067, 185, 120068, 17, 120069, 185, 120066, 12792, 120067, 185, 120068, 5831, 1639, 1221, 478, 1331, 7029, 120069, 185, 120034, 185, 120036, 120025] + "generations": { + "required": { + "completion": "Need call the weather and add tools.\nget_weather\nlocation\nShanghai\n\nadd\ny\n1.0\nx\n2\nitems\n[\"left\", \"right\"]\n\n", + "generation_token_ids": [47893, 1609, 252, 9969, 289, 1189, 7106, 13, 120030, 120035, 185, 120033, 874, 62, 39899, 120043, 185, 120066, 24806, 120067, 185, 120068, 4169, 768, 32381, 120069, 185, 120034, 185, 120033, 2176, 120043, 185, 120066, 88, 120067, 185, 120068, 16, 13, 15, 120069, 185, 120066, 87, 120067, 185, 120068, 17, 120069, 185, 120066, 12792, 120067, 185, 120068, 5831, 1639, 1221, 478, 1331, 7029, 120069, 185, 120034, 185, 120036, 120025] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -618,8 +626,299 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\nget_weather\nlocation\nShanghai\n\n", - "generation_token_ids": [47893, 1609, 252, 9969, 289, 1189, 7106, 13, 120030, 120035, 185, 120033, 874, 62, 39899, 120043, 185, 120066, 24806, 120067, 185, 120068, 4169, 768, 32381, 120069, 185, 120034, 185, 120036, 120025] + "generations": { + "named": { + "completion": "Need call the weather and add tools.\nget_weather\nlocation\nShanghai\n\n", + "generation_token_ids": [47893, 1609, 252, 9969, 289, 1189, 7106, 13, 120030, 120035, 185, 120033, 874, 62, 39899, 120043, 185, 120066, 24806, 120067, 185, 120068, 4169, 768, 32381, 120069, 185, 120034, 185, 120036, 120025] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "\n" + }, + { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "get_weather\n", + "content": { + "type": "sequence", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "location" + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "\n" + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "any_text", + "excludes": [ + "", + "", + "", + "", + "", + "", + "", + "", + "" + ], + "max_tokens": null, + "max_chars": null + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "\n" + } + ] + } + ] + }, + "end": "" + }, + { + "type": "tag", + "begin": "add\n", + "content": { + "type": "sequence", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "y" + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "\n" + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "any_text", + "excludes": [ + "", + "", + "", + "", + "", + "", + "", + "", + "" + ], + "max_tokens": null, + "max_chars": null + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "\n" + } + ] + }, + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "x" + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "\n" + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "any_text", + "excludes": [ + "", + "", + "", + "", + "", + "", + "", + "", + "" + ], + "max_tokens": null, + "max_chars": null + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "\n" + } + ] + }, + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "items" + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "\n" + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "any_text", + "excludes": [ + "", + "", + "", + "", + "", + "", + "", + "", + "" + ], + "max_tokens": null, + "max_chars": null + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "\n" + } + ] + } + ] + }, + "end": "" + } + ], + "separator": "\n", + "at_least_one": true, + "stop_after_first": false + }, + { + "type": "const_string", + "value": "\n" + } + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_calls": { + "completion": "Need call the weather and add tools.\nget_weather\nlocation\nShanghai\n\nadd\ny\n1.0\nx\n2\nitems\n[\"left\", \"right\"]\n\n", + "generation_token_ids": [47893, 1609, 252, 9969, 289, 1189, 7106, 13, 120030, 120035, 185, 120033, 874, 62, 39899, 120043, 185, 120066, 24806, 120067, 185, 120068, 4169, 768, 32381, 120069, 185, 120034, 185, 120033, 2176, 120043, 185, 120066, 88, 120067, 185, 120068, 16, 13, 15, 120069, 185, 120066, 87, 120067, 185, 120068, 17, 120069, 185, 120066, 12792, 120067, 185, 120068, 5831, 1639, 1221, 478, 1331, 7029, 120069, 185, 120034, 185, 120036, 120025] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/Hy3.txt b/rust/src/chat/tests/grammar_replay/Hy3.txt index dcd6bb08ec8b..5642b3697078 100644 --- a/rust/src/chat/tests/grammar_replay/Hy3.txt +++ b/rust/src/chat/tests/grammar_replay/Hy3.txt @@ -110,3 +110,53 @@ sequence `` `\n` `\n` + +# auto_schema_calls (from_token_zero) +sequence + tag `` text `` + or + sequence + `\n` + tags_with_separator `\n` at_least_one + tag `get_weather\n` .. `` + sequence + sequence + `` + `location` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + tag `add\n` .. `` + sequence + sequence + `` + `y` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + sequence + `` + `x` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + sequence + `` + `items` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + `\n` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/Kimi-K2.5.json b/rust/src/chat/tests/grammar_replay/Kimi-K2.5.json index 8ac6e689f78b..3fb22cc3b8cf 100644 --- a/rust/src/chat/tests/grammar_replay/Kimi-K2.5.json +++ b/rust/src/chat/tests/grammar_replay/Kimi-K2.5.json @@ -2,8 +2,8 @@ "model_id": "moonshotai/Kimi-K2.5", "vocab_size": 163840, "all_stop_token_ids": [163585, 163586], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -134,10 +134,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.I will call the tools.<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{\"location\": \"Shanghai\"}<|tool_call_end|><|tool_call_begin|>functions.add:1<|tool_call_argument_begin|>{\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}<|tool_call_end|><|tool_calls_section_end|>", - "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163607, 40, 748, 2390, 276, 7697, 13, 163595, 163597, 41937, 1150, 21055, 2800, 25, 15, 163598, 8264, 5791, 1289, 414, 2797, 31231, 16934, 163599, 163597, 41937, 3149, 25, 16, 163598, 8264, 88, 1289, 220, 16, 13, 15, 11, 414, 87, 1289, 220, 17, 11, 414, 7989, 1289, 8681, 3794, 665, 414, 2277, 3956, 92, 163599, 163596, 163586] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.I will call the tools.<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{\"location\": \"Shanghai\"}<|tool_call_end|><|tool_call_begin|>functions.add:1<|tool_call_argument_begin|>{\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}<|tool_call_end|><|tool_calls_section_end|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163607, 40, 748, 2390, 276, 7697, 13, 163595, 163597, 41937, 1150, 21055, 2800, 25, 15, 163598, 8264, 5791, 1289, 414, 2797, 31231, 16934, 163599, 163597, 41937, 3149, 25, 16, 163598, 8264, 88, 1289, 220, 16, 13, 15, 11, 414, 87, 1289, 220, 17, 11, 414, 7989, 1289, 8681, 3794, 665, 414, 2277, 3956, 92, 163599, 163596, 163586] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -261,10 +265,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{\"location\": \"Shanghai\"}<|tool_call_end|><|tool_call_begin|>functions.add:1<|tool_call_argument_begin|>{\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}<|tool_call_end|><|tool_calls_section_end|>", - "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163607, 163595, 163597, 41937, 1150, 21055, 2800, 25, 15, 163598, 8264, 5791, 1289, 414, 2797, 31231, 16934, 163599, 163597, 41937, 3149, 25, 16, 163598, 8264, 88, 1289, 220, 16, 13, 15, 11, 414, 87, 1289, 220, 17, 11, 414, 7989, 1289, 8681, 3794, 665, 414, 2277, 3956, 92, 163599, 163596, 163586] + "generations": { + "required": { + "completion": "Need call the weather and add tools.<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{\"location\": \"Shanghai\"}<|tool_call_end|><|tool_call_begin|>functions.add:1<|tool_call_argument_begin|>{\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}<|tool_call_end|><|tool_calls_section_end|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163607, 163595, 163597, 41937, 1150, 21055, 2800, 25, 15, 163598, 8264, 5791, 1289, 414, 2797, 31231, 16934, 163599, 163597, 41937, 3149, 25, 16, 163598, 8264, 88, 1289, 220, 16, 13, 15, 11, 414, 87, 1289, 220, 17, 11, 414, 7989, 1289, 8681, 3794, 665, 414, 2277, 3956, 92, 163599, 163596, 163586] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -333,8 +341,166 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{\"location\": \"Shanghai\"}<|tool_call_end|><|tool_calls_section_end|>", - "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163607, 163595, 163597, 41937, 1150, 21055, 2800, 25, 15, 163598, 8264, 5791, 1289, 414, 2797, 31231, 16934, 163599, 163596, 163586] + "generations": { + "named": { + "completion": "Need call the weather and add tools.<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{\"location\": \"Shanghai\"}<|tool_call_end|><|tool_calls_section_end|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163607, 163595, 163597, 41937, 1150, 21055, 2800, 25, 15, 163598, 8264, 5791, 1289, 414, 2797, 31231, 16934, 163599, 163596, 163586] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "<|tool_calls_section_begin|>" + }, + { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<|tool_call_begin|>functions.get_weather:", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "\\d+" + }, + { + "type": "const_string", + "value": "<|tool_call_argument_begin|>" + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "end": "<|tool_call_end|>" + }, + { + "type": "tag", + "begin": "<|tool_call_begin|>functions.add:", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "\\d+" + }, + { + "type": "const_string", + "value": "<|tool_call_argument_begin|>" + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "end": "<|tool_call_end|>" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + { + "type": "const_string", + "value": "<|tool_calls_section_end|>" + } + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_calls": { + "completion": "Need call the weather and add tools.<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{\"location\": \"Shanghai\"}<|tool_call_end|><|tool_call_begin|>functions.add:1<|tool_call_argument_begin|>{\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}<|tool_call_end|><|tool_calls_section_end|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163607, 163595, 163597, 41937, 1150, 21055, 2800, 25, 15, 163598, 8264, 5791, 1289, 414, 2797, 31231, 16934, 163599, 163597, 41937, 3149, 25, 16, 163598, 8264, 88, 1289, 220, 16, 13, 15, 11, 414, 87, 1289, 220, 17, 11, 414, 7989, 1289, 8681, 3794, 665, 414, 2277, 3956, 92, 163599, 163596, 163586] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/Kimi-K2.5.txt b/rust/src/chat/tests/grammar_replay/Kimi-K2.5.txt index 553c1ce6f4c3..ca3bab8061d7 100644 --- a/rust/src/chat/tests/grammar_replay/Kimi-K2.5.txt +++ b/rust/src/chat/tests/grammar_replay/Kimi-K2.5.txt @@ -44,3 +44,23 @@ sequence `<|tool_call_argument_begin|>` json({ location: string }) `<|tool_calls_section_end|>` + +# auto_schema_calls (from_token_zero) +sequence + tag `` text `` + or + sequence + `<|tool_calls_section_begin|>` + tags_with_separator `` at_least_one + tag `<|tool_call_begin|>functions.get_weather:` .. `<|tool_call_end|>` + sequence + /\d+/ + `<|tool_call_argument_begin|>` + json({ location: string }) + tag `<|tool_call_begin|>functions.add:` .. `<|tool_call_end|>` + sequence + /\d+/ + `<|tool_call_argument_begin|>` + json({ y: number, x: number, items: string[] }) + `<|tool_calls_section_end|>` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/Kimi-K3.json b/rust/src/chat/tests/grammar_replay/Kimi-K3.json index 5306b936604f..3e1832e805f7 100644 --- a/rust/src/chat/tests/grammar_replay/Kimi-K3.json +++ b/rust/src/chat/tests/grammar_replay/Kimi-K3.json @@ -2,8 +2,8 @@ "model_id": "moonshotai/Kimi-K3", "vocab_size": 163840, "all_stop_token_ids": [163585, 163586], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -427,10 +427,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|>I will call the tools.<|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool=\"add\" index=\"2\"<|sep|><|open|>argument key=\"y\" type=\"number\"<|sep|>1.0<|close|>argument<|sep|><|open|>argument key=\"x\" type=\"number\"<|sep|>2<|close|>argument<|sep|><|open|>argument key=\"items\" type=\"array\"<|sep|>[\"left\",\"right\"]<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", - "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 40, 748, 2390, 276, 7697, 13, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163587, 10257, 4453, 878, 1210, 1, 4002, 878, 17, 1, 163589, 163587, 47185, 2355, 878, 88, 1, 1798, 878, 10006, 1, 163589, 16, 13, 15, 163588, 47185, 163589, 163587, 47185, 2355, 878, 87, 1, 1798, 878, 10006, 1, 163589, 17, 163588, 47185, 163589, 163587, 47185, 2355, 878, 7989, 1, 1798, 878, 4392, 1, 163589, 3530, 3794, 3923, 2277, 3956, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|>I will call the tools.<|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool=\"add\" index=\"2\"<|sep|><|open|>argument key=\"y\" type=\"number\"<|sep|>1.0<|close|>argument<|sep|><|open|>argument key=\"x\" type=\"number\"<|sep|>2<|close|>argument<|sep|><|open|>argument key=\"items\" type=\"array\"<|sep|>[\"left\",\"right\"]<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 40, 748, 2390, 276, 7697, 13, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163587, 10257, 4453, 878, 1210, 1, 4002, 878, 17, 1, 163589, 163587, 47185, 2355, 878, 88, 1, 1798, 878, 10006, 1, 163589, 16, 13, 15, 163588, 47185, 163589, 163587, 47185, 2355, 878, 87, 1, 1798, 878, 10006, 1, 163589, 17, 163588, 47185, 163589, 163587, 47185, 2355, 878, 7989, 1, 1798, 878, 4392, 1, 163589, 3530, 3794, 3923, 2277, 3956, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -666,10 +670,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool=\"add\" index=\"2\"<|sep|><|open|>argument key=\"y\" type=\"number\"<|sep|>1.0<|close|>argument<|sep|><|open|>argument key=\"x\" type=\"number\"<|sep|>2<|close|>argument<|sep|><|open|>argument key=\"items\" type=\"array\"<|sep|>[\"left\",\"right\"]<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", - "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163587, 10257, 4453, 878, 1210, 1, 4002, 878, 17, 1, 163589, 163587, 47185, 2355, 878, 88, 1, 1798, 878, 10006, 1, 163589, 16, 13, 15, 163588, 47185, 163589, 163587, 47185, 2355, 878, 87, 1, 1798, 878, 10006, 1, 163589, 17, 163588, 47185, 163589, 163587, 47185, 2355, 878, 7989, 1, 1798, 878, 4392, 1, 163589, 3530, 3794, 3923, 2277, 3956, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] + "generations": { + "required": { + "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool=\"add\" index=\"2\"<|sep|><|open|>argument key=\"y\" type=\"number\"<|sep|>1.0<|close|>argument<|sep|><|open|>argument key=\"x\" type=\"number\"<|sep|>2<|close|>argument<|sep|><|open|>argument key=\"items\" type=\"array\"<|sep|>[\"left\",\"right\"]<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163587, 10257, 4453, 878, 1210, 1, 4002, 878, 17, 1, 163589, 163587, 47185, 2355, 878, 88, 1, 1798, 878, 10006, 1, 163589, 16, 13, 15, 163588, 47185, 163589, 163587, 47185, 2355, 878, 87, 1, 1798, 878, 10006, 1, 163589, 17, 163588, 47185, 163589, 163587, 47185, 2355, 878, 7989, 1, 1798, 878, 4392, 1, 163589, 3530, 3794, 3923, 2277, 3956, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -795,8 +803,527 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", - "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] + "generations": { + "named": { + "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [ + "<|close|>think<|sep|>", + "<|end_of_msg|>" + ], + "max_tokens": null, + "max_chars": null + }, + "end": "<|close|>think<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>response<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>response<|sep|>" + }, + { + "type": "optional", + "content": { + "type": "const_string", + "value": "<|close|>message<|sep|>" + } + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|>{\"answer\":\"sunny\"}<|close|>response<|sep|><|close|>message<|sep|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 8264, 26921, 7471, 82, 38975, 16934, 163588, 12092, 163589, 163588, 2778, 163589, 163586] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [ + "<|close|>think<|sep|>", + "<|end_of_msg|>" + ], + "max_tokens": null, + "max_chars": null + }, + "end": "<|close|>think<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "<|open|>response<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>response<|sep|>" + }, + { + "type": "optional", + "content": { + "type": "tag", + "begin": "<|open|>tools<|sep|>", + "content": { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<|open|>call tool=\"get_weather\" index=\"", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "[1-9][0-9]*" + }, + { + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>", + "content": { + "type": "any_text", + "excludes": [ + "<|close|>argument<|sep|>", + "<|close|>call<|sep|>" + ], + "max_tokens": null, + "max_chars": null + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>json type=\"object\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>json<|sep|>" + } + ] + } + ] + }, + "end": "<|close|>call<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>call tool=\"add\" index=\"", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "[1-9][0-9]*" + }, + { + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"y\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"x\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"items\" type=\"array\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "array", + "items": { + "type": "string" + } + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + } + ] + }, + { + "type": "tag", + "begin": "<|open|>json type=\"object\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>json<|sep|>" + } + ] + } + ] + }, + "end": "<|close|>call<|sep|>" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + "end": "<|close|>tools<|sep|>" + } + } + ] + }, + { + "type": "sequence", + "elements": [ + { + "type": "optional", + "content": { + "type": "const_string", + "value": "<|open|>response<|sep|><|close|>response<|sep|>" + } + }, + { + "type": "tag", + "begin": "<|open|>tools<|sep|>", + "content": { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<|open|>call tool=\"get_weather\" index=\"", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "[1-9][0-9]*" + }, + { + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>", + "content": { + "type": "any_text", + "excludes": [ + "<|close|>argument<|sep|>", + "<|close|>call<|sep|>" + ], + "max_tokens": null, + "max_chars": null + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>json type=\"object\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>json<|sep|>" + } + ] + } + ] + }, + "end": "<|close|>call<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>call tool=\"add\" index=\"", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "[1-9][0-9]*" + }, + { + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"y\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"x\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"items\" type=\"array\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "array", + "items": { + "type": "string" + } + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + } + ] + }, + { + "type": "tag", + "begin": "<|open|>json type=\"object\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>json<|sep|>" + } + ] + } + ] + }, + "end": "<|close|>call<|sep|>" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + "end": "<|close|>tools<|sep|>" + } + ] + } + ] + }, + { + "type": "optional", + "content": { + "type": "const_string", + "value": "<|close|>message<|sep|>" + } + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_content": { + "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|>{\"answer\":\"sunny\"}<|close|>response<|sep|><|close|>message<|sep|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 8264, 26921, 7471, 82, 38975, 16934, 163588, 12092, 163589, 163588, 2778, 163589, 163586] + }, + "auto_schema_calls": { + "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool=\"add\" index=\"2\"<|sep|><|open|>argument key=\"y\" type=\"number\"<|sep|>1.0<|close|>argument<|sep|><|open|>argument key=\"x\" type=\"number\"<|sep|>2<|close|>argument<|sep|><|open|>argument key=\"items\" type=\"array\"<|sep|>[\"left\",\"right\"]<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163587, 10257, 4453, 878, 1210, 1, 4002, 878, 17, 1, 163589, 163587, 47185, 2355, 878, 88, 1, 1798, 878, 10006, 1, 163589, 16, 13, 15, 163588, 47185, 163589, 163587, 47185, 2355, 878, 87, 1, 1798, 878, 10006, 1, 163589, 17, 163588, 47185, 163589, 163587, 47185, 2355, 878, 7989, 1, 1798, 878, 4392, 1, 163589, 3530, 3794, 3923, 2277, 3956, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/Kimi-K3.txt b/rust/src/chat/tests/grammar_replay/Kimi-K3.txt index 6a1c814c131c..d7357e275c86 100644 --- a/rust/src/chat/tests/grammar_replay/Kimi-K3.txt +++ b/rust/src/chat/tests/grammar_replay/Kimi-K3.txt @@ -86,3 +86,58 @@ sequence tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` optional `<|close|>message<|sep|>` + +# none_schema (from_token_zero) +sequence + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + tag `<|open|>response<|sep|>` json({ answer: string }) `<|close|>response<|sep|>` + optional `<|close|>message<|sep|>` + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + or + sequence + tag `<|open|>response<|sep|>` json({ answer: string }) `<|close|>response<|sep|>` + optional + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` + tag `<|open|>call tool="add" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + sequence + tag `<|open|>argument key="y" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="x" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="items" type="array"<|sep|>` json(string[]) `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ y: number, x: number, items: string[] }) `<|close|>json<|sep|>` + sequence + optional `<|open|>response<|sep|><|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` + tag `<|open|>call tool="add" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + sequence + tag `<|open|>argument key="y" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="x" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="items" type="array"<|sep|>` json(string[]) `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ y: number, x: number, items: string[] }) `<|close|>json<|sep|>` + optional `<|close|>message<|sep|>` diff --git a/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.json b/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.json index 582f29a1b943..451255dd8c45 100644 --- a/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.json +++ b/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.json @@ -2,8 +2,8 @@ "model_id": "XiaomiMiMo/MiMo-V2.6-Flash-RL", "vocab_size": 152576, "all_stop_token_ids": [151643, 151645, 151672], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -99,10 +99,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.I will call the tools.Shanghai1.02[\"left\", \"right\"]", - "generation_token_ids": [151667, 23657, 1618, 279, 9104, 323, 912, 7375, 13, 151668, 40, 686, 1618, 279, 7375, 13, 151657, 27, 1688, 28280, 69364, 1784, 16181, 28, 2527, 29, 2016, 30070, 522, 16181, 1472, 1688, 29, 151658, 151657, 27, 1688, 92784, 1784, 16181, 29368, 29, 16, 13, 15, 522, 16181, 1784, 16181, 25387, 29, 17, 522, 16181, 1784, 16181, 28, 3615, 29, 1183, 2359, 497, 330, 1291, 1341, 522, 16181, 1472, 1688, 29, 151658, 151645] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.I will call the tools.Shanghai1.02[\"left\", \"right\"]", + "generation_token_ids": [151667, 23657, 1618, 279, 9104, 323, 912, 7375, 13, 151668, 40, 686, 1618, 279, 7375, 13, 151657, 27, 1688, 28280, 69364, 1784, 16181, 28, 2527, 29, 2016, 30070, 522, 16181, 1472, 1688, 29, 151658, 151657, 27, 1688, 92784, 1784, 16181, 29368, 29, 16, 13, 15, 522, 16181, 1784, 16181, 25387, 29, 17, 522, 16181, 1784, 16181, 28, 3615, 29, 1183, 2359, 497, 330, 1291, 1341, 522, 16181, 1472, 1688, 29, 151658, 151645] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -198,10 +202,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.Shanghai1.02[\"left\", \"right\"]", - "generation_token_ids": [151667, 23657, 1618, 279, 9104, 323, 912, 7375, 13, 151668, 151657, 27, 1688, 28280, 69364, 1784, 16181, 28, 2527, 29, 2016, 30070, 522, 16181, 1472, 1688, 29, 151658, 151657, 27, 1688, 92784, 1784, 16181, 29368, 29, 16, 13, 15, 522, 16181, 1784, 16181, 25387, 29, 17, 522, 16181, 1784, 16181, 28, 3615, 29, 1183, 2359, 497, 330, 1291, 1341, 522, 16181, 1472, 1688, 29, 151658, 151645] + "generations": { + "required": { + "completion": "Need call the weather and add tools.Shanghai1.02[\"left\", \"right\"]", + "generation_token_ids": [151667, 23657, 1618, 279, 9104, 323, 912, 7375, 13, 151668, 151657, 27, 1688, 28280, 69364, 1784, 16181, 28, 2527, 29, 2016, 30070, 522, 16181, 1472, 1688, 29, 151658, 151657, 27, 1688, 92784, 1784, 16181, 29368, 29, 16, 13, 15, 522, 16181, 1784, 16181, 25387, 29, 17, 522, 16181, 1784, 16181, 28, 3615, 29, 1183, 2359, 497, 330, 1291, 1341, 522, 16181, 1472, 1688, 29, 151658, 151645] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -247,8 +255,190 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.Shanghai", - "generation_token_ids": [151667, 23657, 1618, 279, 9104, 323, 912, 7375, 13, 151668, 151657, 27, 1688, 28280, 69364, 1784, 16181, 28, 2527, 29, 2016, 30070, 522, 16181, 1472, 1688, 29, 151658, 151645] + "generations": { + "named": { + "completion": "Need call the weather and add tools.Shanghai", + "generation_token_ids": [151667, 23657, 1618, 279, 9104, 323, 912, 7375, 13, 151668, 151657, 27, 1688, 28280, 69364, 1784, 16181, 28, 2527, 29, 2016, 30070, 522, 16181, 1472, 1688, 29, 151658, 151645] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "optional", + "content": { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + } + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [151667, 23657, 1618, 279, 9104, 323, 912, 7375, 13, 151668, 4913, 9217, 3252, 82, 27297, 9207, 151645] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "optional", + "content": { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + } + }, + { + "type": "or", + "elements": [ + { + "type": "triggered_tags", + "triggers": [ + "" + ], + "tags": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "qwen_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "" + }, + { + "type": "tag", + "begin": "", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "qwen_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "" + } + ], + "at_least_one": true, + "stop_after_first": false, + "excludes": [ + "", + "", + "", + "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [151667, 23657, 1618, 279, 9104, 323, 912, 7375, 13, 151668, 4913, 9217, 3252, 82, 27297, 9207, 151645] + }, + "auto_schema_calls": { + "completion": "Need call the weather and add tools.Shanghai1.02[\"left\", \"right\"]", + "generation_token_ids": [151667, 23657, 1618, 279, 9104, 323, 912, 7375, 13, 151668, 151657, 27, 1688, 28280, 69364, 1784, 16181, 28, 2527, 29, 2016, 30070, 522, 16181, 1472, 1688, 29, 151658, 151657, 27, 1688, 92784, 1784, 16181, 29368, 29, 16, 13, 15, 522, 16181, 1784, 16181, 25387, 29, 17, 522, 16181, 1784, 16181, 28, 3615, 29, 1183, 2359, 497, 330, 1291, 1341, 522, 16181, 1472, 1688, 29, 151658, 151645] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.txt b/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.txt index 1d12560eb034..7cc9aa0e59a2 100644 --- a/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.txt +++ b/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.txt @@ -16,3 +16,17 @@ sequence sequence optional tag `` text `` tag `` qwen_xml({ location: string }) `` + +# none_schema (from_token_zero) +sequence + optional tag `` text `` + json({ answer: string }) + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + optional tag `` text `` + or + triggered_tags [``] excluding [``, ``, ``, `` qwen_xml({ location: string }) `` + tag `` qwen_xml({ y: number, x: number, items: string[] }) `` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.json b/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.json index 64f65e2f5a32..d4c2e8e789f6 100644 --- a/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.json +++ b/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.json @@ -2,8 +2,8 @@ "model_id": "MiniMaxAI/MiniMax-M2.5", "vocab_size": 200064, "all_stop_token_ids": [200020], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -116,10 +116,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n\nI will call the tools.\n\n\nShanghai\n\n\n1.0\n2\n[\"left\", \"right\"]\n\n", - "generation_token_ids": [42478, 1728, 275, 12562, 306, 1187, 6629, 320, 200051, 367, 73, 715, 1728, 275, 6629, 320, 200052, 10, 60, 88638, 1925, 1139, 848, 95, 51872, 6143, 60, 24662, 1925, 1139, 28415, 3361, 165806, 1579, 24662, 1100, 1579, 88638, 1100, 60, 88638, 1925, 1139, 2132, 6143, 60, 24662, 1925, 1139, 121, 3361, 49, 46, 48, 1579, 24662, 1100, 60, 24662, 1925, 1139, 120, 3361, 50, 1579, 24662, 1100, 60, 24662, 1925, 1139, 12405, 3361, 7795, 2038, 1164, 494, 1581, 8760, 1579, 24662, 1100, 1579, 88638, 1100, 200053, 200020] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.\n\n\nI will call the tools.\n\n\nShanghai\n\n\n1.0\n2\n[\"left\", \"right\"]\n\n", + "generation_token_ids": [42478, 1728, 275, 12562, 306, 1187, 6629, 320, 200051, 367, 73, 715, 1728, 275, 6629, 320, 200052, 10, 60, 88638, 1925, 1139, 848, 95, 51872, 6143, 60, 24662, 1925, 1139, 28415, 3361, 165806, 1579, 24662, 1100, 1579, 88638, 1100, 60, 88638, 1925, 1139, 2132, 6143, 60, 24662, 1925, 1139, 121, 3361, 49, 46, 48, 1579, 24662, 1100, 60, 24662, 1925, 1139, 120, 3361, 50, 1579, 24662, 1100, 60, 24662, 1925, 1139, 12405, 3361, 7795, 2038, 1164, 494, 1581, 8760, 1579, 24662, 1100, 1579, 88638, 1100, 200053, 200020] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -226,10 +230,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n\n1.0\n2\n[\"left\", \"right\"]\n\n", - "generation_token_ids": [42478, 1728, 275, 12562, 306, 1187, 6629, 320, 200051, 4368, 200052, 10, 60, 88638, 1925, 1139, 848, 95, 51872, 6143, 60, 24662, 1925, 1139, 28415, 3361, 165806, 1579, 24662, 1100, 1579, 88638, 1100, 60, 88638, 1925, 1139, 2132, 6143, 60, 24662, 1925, 1139, 121, 3361, 49, 46, 48, 1579, 24662, 1100, 60, 24662, 1925, 1139, 120, 3361, 50, 1579, 24662, 1100, 60, 24662, 1925, 1139, 12405, 3361, 7795, 2038, 1164, 494, 1581, 8760, 1579, 24662, 1100, 1579, 88638, 1100, 200053, 200020] + "generations": { + "required": { + "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n\n1.0\n2\n[\"left\", \"right\"]\n\n", + "generation_token_ids": [42478, 1728, 275, 12562, 306, 1187, 6629, 320, 200051, 4368, 200052, 10, 60, 88638, 1925, 1139, 848, 95, 51872, 6143, 60, 24662, 1925, 1139, 28415, 3361, 165806, 1579, 24662, 1100, 1579, 88638, 1100, 60, 88638, 1925, 1139, 2132, 6143, 60, 24662, 1925, 1139, 121, 3361, 49, 46, 48, 1579, 24662, 1100, 60, 24662, 1925, 1139, 120, 3361, 50, 1579, 24662, 1100, 60, 24662, 1925, 1139, 12405, 3361, 7795, 2038, 1164, 494, 1581, 8760, 1579, 24662, 1100, 1579, 88638, 1100, 200053, 200020] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -294,8 +302,207 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n", - "generation_token_ids": [42478, 1728, 275, 12562, 306, 1187, 6629, 320, 200051, 4368, 200052, 10, 60, 88638, 1925, 1139, 848, 95, 51872, 6143, 60, 24662, 1925, 1139, 28415, 3361, 165806, 1579, 24662, 1100, 1579, 88638, 1100, 200053, 200020] + "generations": { + "named": { + "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n", + "generation_token_ids": [42478, 1728, 275, 12562, 306, 1187, 6629, 320, 200051, 4368, 200052, 10, 60, 88638, 1925, 1139, 848, 95, 51872, 6143, 60, 24662, 1925, 1139, 28415, 3361, 165806, 1579, 24662, 1100, 1579, 88638, 1100, 200053, 200020] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "const_string", + "value": "\n\n" + } + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.\n\n\n{\"answer\":\"sunny\"}", + "generation_token_ids": [42478, 1728, 275, 12562, 306, 1187, 6629, 320, 200051, 367, 16673, 33591, 15368, 78697, 7582, 22089, 200020] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "const_string", + "value": "\n\n" + } + ] + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "\n\n" + }, + { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "minimax_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n" + }, + { + "type": "tag", + "begin": "\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "minimax_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + { + "type": "const_string", + "value": "" + } + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_content": { + "completion": "Need call the weather and add tools.\n\n\n{\"answer\":\"sunny\"}", + "generation_token_ids": [42478, 1728, 275, 12562, 306, 1187, 6629, 320, 200051, 367, 16673, 33591, 15368, 78697, 7582, 22089, 200020] + }, + "auto_schema_calls": { + "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n\n1.0\n2\n[\"left\", \"right\"]\n\n", + "generation_token_ids": [42478, 1728, 275, 12562, 306, 1187, 6629, 320, 200051, 4368, 200052, 10, 60, 88638, 1925, 1139, 848, 95, 51872, 6143, 60, 24662, 1925, 1139, 28415, 3361, 165806, 1579, 24662, 1100, 1579, 88638, 1100, 60, 88638, 1925, 1139, 2132, 6143, 60, 24662, 1925, 1139, 121, 3361, 49, 46, 48, 1579, 24662, 1100, 60, 24662, 1925, 1139, 120, 3361, 50, 1579, 24662, 1100, 60, 24662, 1925, 1139, 12405, 3361, 7795, 2038, 1164, 494, 1581, 8760, 1579, 24662, 1100, 1579, 88638, 1100, 200053, 200020] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.txt b/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.txt index 431238009a13..fae13b84450a 100644 --- a/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.txt +++ b/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.txt @@ -30,3 +30,24 @@ sequence `\n\n` tag `\n` minimax_xml({ location: string }) `\n` `` + +# none_schema (from_token_zero) +sequence + sequence + tag `` text `` + `\n\n` + json({ answer: string }) + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + sequence + tag `` text `` + `\n\n` + or + sequence + `\n\n` + tags_with_separator `` at_least_one + tag `\n` minimax_xml({ location: string }) `\n` + tag `\n` minimax_xml({ y: number, x: number, items: string[] }) `\n` + `` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/MiniMax-M3.json b/rust/src/chat/tests/grammar_replay/MiniMax-M3.json new file mode 100644 index 000000000000..73679eff6a1e --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/MiniMax-M3.json @@ -0,0 +1,40 @@ +{ + "model_id": "MiniMaxAI/MiniMax-M3", + "vocab_size": 200064, + "all_stop_token_ids": [200020], + "grammars": [ + { + "grammar": { + "format": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "coverage": "final_output_only" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [200059, 42478, 1728, 275, 12562, 306, 1187, 6629, 46, 200060, 16673, 33591, 15368, 78697, 7582, 22089, 200020] + }, + "auto_schema_content": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [200059, 42478, 1728, 275, 12562, 306, 1187, 6629, 46, 200060, 16673, 33591, 15368, 78697, 7582, 22089, 200020] + } + } + } + ] +} diff --git a/rust/src/chat/tests/grammar_replay/MiniMax-M3.txt b/rust/src/chat/tests/grammar_replay/MiniMax-M3.txt new file mode 100644 index 000000000000..f54e72d4c825 --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/MiniMax-M3.txt @@ -0,0 +1,2 @@ +# none_schema, auto_schema_content (final_output_only) +json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.json b/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.json index 5d171fc5f080..8206b8b40c1a 100644 --- a/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.json +++ b/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.json @@ -2,8 +2,8 @@ "model_id": "Qwen/Qwen3-0.6B", "vocab_size": 151936, "all_stop_token_ids": [151643, 151645], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -106,10 +106,14 @@ }, "coverage": "from_token_zero" }, - "completion": "\nNeed call the weather and add tools.\n\n\nI will call the tools.\n\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Shanghai\"}}\n\n\n{\"name\": \"add\", \"arguments\": {\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}}\n", - "generation_token_ids": [151667, 198, 23657, 1618, 279, 9104, 323, 912, 7375, 624, 151668, 271, 40, 686, 1618, 279, 7375, 624, 151657, 198, 4913, 606, 788, 330, 455, 69364, 497, 330, 16370, 788, 5212, 2527, 788, 330, 2016, 30070, 95642, 151658, 198, 151657, 198, 4913, 606, 788, 330, 718, 497, 330, 16370, 788, 5212, 88, 788, 220, 16, 13, 15, 11, 330, 87, 788, 220, 17, 11, 330, 3615, 788, 4383, 2359, 497, 330, 1291, 1341, 11248, 151658, 151645] + "generations": { + "auto": { + "completion": "\nNeed call the weather and add tools.\n\n\nI will call the tools.\n\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Shanghai\"}}\n\n\n{\"name\": \"add\", \"arguments\": {\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}}\n", + "generation_token_ids": [151667, 198, 23657, 1618, 279, 9104, 323, 912, 7375, 624, 151668, 271, 40, 686, 1618, 279, 7375, 624, 151657, 198, 4913, 606, 788, 330, 455, 69364, 497, 330, 16370, 788, 5212, 2527, 788, 330, 2016, 30070, 95642, 151658, 198, 151657, 198, 4913, 606, 788, 330, 718, 497, 330, 16370, 788, 5212, 88, 788, 220, 16, 13, 15, 11, 330, 87, 788, 220, 17, 11, 330, 3615, 788, 4383, 2359, 497, 330, 1291, 1341, 11248, 151658, 151645] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -212,10 +216,14 @@ }, "coverage": "from_token_zero" }, - "completion": "\nNeed call the weather and add tools.\n\n\n\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Shanghai\"}}\n\n\n{\"name\": \"add\", \"arguments\": {\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}}\n", - "generation_token_ids": [151667, 198, 23657, 1618, 279, 9104, 323, 912, 7375, 624, 151668, 271, 151657, 198, 4913, 606, 788, 330, 455, 69364, 497, 330, 16370, 788, 5212, 2527, 788, 330, 2016, 30070, 95642, 151658, 198, 151657, 198, 4913, 606, 788, 330, 718, 497, 330, 16370, 788, 5212, 88, 788, 220, 16, 13, 15, 11, 330, 87, 788, 220, 17, 11, 330, 3615, 788, 4383, 2359, 497, 330, 1291, 1341, 11248, 151658, 151645] + "generations": { + "required": { + "completion": "\nNeed call the weather and add tools.\n\n\n\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Shanghai\"}}\n\n\n{\"name\": \"add\", \"arguments\": {\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}}\n", + "generation_token_ids": [151667, 198, 23657, 1618, 279, 9104, 323, 912, 7375, 624, 151668, 271, 151657, 198, 4913, 606, 788, 330, 455, 69364, 497, 330, 16370, 788, 5212, 2527, 788, 330, 2016, 30070, 95642, 151658, 198, 151657, 198, 4913, 606, 788, 330, 718, 497, 330, 16370, 788, 5212, 88, 788, 220, 16, 13, 15, 11, 330, 87, 788, 220, 17, 11, 330, 3615, 788, 4383, 2359, 497, 330, 1291, 1341, 11248, 151658, 151645] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -270,8 +278,206 @@ }, "coverage": "from_token_zero" }, - "completion": "\nNeed call the weather and add tools.\n\n\n\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Shanghai\"}}\n", - "generation_token_ids": [151667, 198, 23657, 1618, 279, 9104, 323, 912, 7375, 624, 151668, 271, 151657, 198, 4913, 606, 788, 330, 455, 69364, 497, 330, 16370, 788, 5212, 2527, 788, 330, 2016, 30070, 95642, 151658, 151645] + "generations": { + "named": { + "completion": "\nNeed call the weather and add tools.\n\n\n\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Shanghai\"}}\n", + "generation_token_ids": [151667, 198, 23657, 1618, 279, 9104, 323, 912, 7375, 624, 151668, 271, 151657, 198, 4913, 606, 788, 330, 455, 69364, 497, 330, 16370, 788, 5212, 2527, 788, 330, 2016, 30070, 95642, 151658, 151645] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "optional", + "content": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "const_string", + "value": "\n\n" + } + ] + } + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "\nNeed call the weather and add tools.\n\n\n{\"answer\":\"sunny\"}", + "generation_token_ids": [151667, 198, 23657, 1618, 279, 9104, 323, 912, 7375, 624, 151668, 271, 4913, 9217, 3252, 82, 27297, 9207, 151645] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "optional", + "content": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "const_string", + "value": "\n\n" + } + ] + } + }, + { + "type": "or", + "elements": [ + { + "type": "triggered_tags", + "triggers": [ + "" + ], + "tags": [ + { + "type": "tag", + "begin": "\n{\"name\": \"get_weather\", \"arguments\": ", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "}\n" + }, + { + "type": "tag", + "begin": "\n{\"name\": \"add\", \"arguments\": ", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "}\n" + } + ], + "at_least_one": true, + "stop_after_first": false, + "excludes": [ + "", + "" + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_content": { + "completion": "\nNeed call the weather and add tools.\n\n\n{\"answer\":\"sunny\"}", + "generation_token_ids": [151667, 198, 23657, 1618, 279, 9104, 323, 912, 7375, 624, 151668, 271, 4913, 9217, 3252, 82, 27297, 9207, 151645] + }, + "auto_schema_calls": { + "completion": "\nNeed call the weather and add tools.\n\n\n\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Shanghai\"}}\n\n\n{\"name\": \"add\", \"arguments\": {\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}}\n", + "generation_token_ids": [151667, 198, 23657, 1618, 279, 9104, 323, 912, 7375, 624, 151668, 271, 151657, 198, 4913, 606, 788, 330, 455, 69364, 497, 330, 16370, 788, 5212, 2527, 788, 330, 2016, 30070, 95642, 151658, 198, 151657, 198, 4913, 606, 788, 330, 718, 497, 330, 16370, 788, 5212, 88, 788, 220, 16, 13, 15, 11, 330, 87, 788, 220, 17, 11, 330, 3615, 788, 4383, 2359, 497, 330, 1291, 1341, 11248, 151658, 151645] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.txt b/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.txt index fe38f5e9ff4a..865d1b543071 100644 --- a/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.txt +++ b/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.txt @@ -25,3 +25,23 @@ sequence tag `` text `` `\n\n` tag `\n{"name": "get_weather", "arguments": ` json({ location: string }) `}\n` + +# none_schema (from_token_zero) +sequence + optional + sequence + tag `` text `` + `\n\n` + json({ answer: string }) + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + optional + sequence + tag `` text `` + `\n\n` + or + triggered_tags [``] excluding [``, ``] at_least_one + tag `\n{"name": "get_weather", "arguments": ` json({ location: string }) `}\n` + tag `\n{"name": "add", "arguments": ` json({ y: number, x: number, items: string[] }) `}\n` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.json b/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.json index 8bd4c901e541..9cc7b53ff9ca 100644 --- a/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.json +++ b/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.json @@ -2,8 +2,8 @@ "model_id": "Qwen/Qwen3.5-4B", "vocab_size": 248320, "all_stop_token_ids": [248044, 248046], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -103,10 +103,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n\nI will call the tools.\n\n\n\n\nShanghai\n\n\n\n\n\n\n1.0\n\n\n2\n\n\n[\"left\", \"right\"]\n\n\n", - "generation_token_ids": [22916, 1562, 279, 8831, 321, 884, 7141, 13, 198, 248069, 271, 40, 668, 1562, 279, 7141, 13, 271, 248058, 198, 27, 1628, 27362, 67017, 29, 198, 27, 15704, 28, 2447, 29, 198, 1948, 29086, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 198, 248058, 198, 27, 1628, 89587, 29, 198, 27, 15704, 28408, 29, 198, 16, 13, 15, 198, 510, 15704, 29, 198, 27, 15704, 24582, 29, 198, 17, 198, 510, 15704, 29, 198, 27, 15704, 28, 3494, 29, 198, 1143, 2282, 487, 328, 1246, 1293, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 248046] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.\n\n\nI will call the tools.\n\n\n\n\nShanghai\n\n\n\n\n\n\n1.0\n\n\n2\n\n\n[\"left\", \"right\"]\n\n\n", + "generation_token_ids": [22916, 1562, 279, 8831, 321, 884, 7141, 13, 198, 248069, 271, 40, 668, 1562, 279, 7141, 13, 271, 248058, 198, 27, 1628, 27362, 67017, 29, 198, 27, 15704, 28, 2447, 29, 198, 1948, 29086, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 198, 248058, 198, 27, 1628, 89587, 29, 198, 27, 15704, 28408, 29, 198, 16, 13, 15, 198, 510, 15704, 29, 198, 27, 15704, 24582, 29, 198, 17, 198, 510, 15704, 29, 198, 27, 15704, 28, 3494, 29, 198, 1143, 2282, 487, 328, 1246, 1293, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 248046] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -206,10 +210,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n\n\n\n\n\n1.0\n\n\n2\n\n\n[\"left\", \"right\"]\n\n\n", - "generation_token_ids": [22916, 1562, 279, 8831, 321, 884, 7141, 13, 198, 248069, 271, 248058, 198, 27, 1628, 27362, 67017, 29, 198, 27, 15704, 28, 2447, 29, 198, 1948, 29086, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 198, 248058, 198, 27, 1628, 89587, 29, 198, 27, 15704, 28408, 29, 198, 16, 13, 15, 198, 510, 15704, 29, 198, 27, 15704, 24582, 29, 198, 17, 198, 510, 15704, 29, 198, 27, 15704, 28, 3494, 29, 198, 1143, 2282, 487, 328, 1246, 1293, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 248046] + "generations": { + "required": { + "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n\n\n\n\n\n1.0\n\n\n2\n\n\n[\"left\", \"right\"]\n\n\n", + "generation_token_ids": [22916, 1562, 279, 8831, 321, 884, 7141, 13, 198, 248069, 271, 248058, 198, 27, 1628, 27362, 67017, 29, 198, 27, 15704, 28, 2447, 29, 198, 1948, 29086, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 198, 248058, 198, 27, 1628, 89587, 29, 198, 27, 15704, 28408, 29, 198, 16, 13, 15, 198, 510, 15704, 29, 198, 27, 15704, 24582, 29, 198, 17, 198, 510, 15704, 29, 198, 27, 15704, 28, 3494, 29, 198, 1143, 2282, 487, 328, 1246, 1293, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 248046] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -261,8 +269,200 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n\n", - "generation_token_ids": [22916, 1562, 279, 8831, 321, 884, 7141, 13, 198, 248069, 271, 248058, 198, 27, 1628, 27362, 67017, 29, 198, 27, 15704, 28, 2447, 29, 198, 1948, 29086, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 248046] + "generations": { + "named": { + "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n\n", + "generation_token_ids": [22916, 1562, 279, 8831, 321, 884, 7141, 13, 198, 248069, 271, 248058, 198, 27, 1628, 27362, 67017, 29, 198, 27, 15704, 28, 2447, 29, 198, 1948, 29086, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 248046] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "const_string", + "value": "\n\n" + } + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.\n\n\n{\"answer\":\"sunny\"}", + "generation_token_ids": [22916, 1562, 279, 8831, 321, 884, 7141, 13, 198, 248069, 271, 4754, 8944, 3147, 82, 26415, 8934, 248046] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "const_string", + "value": "\n\n" + } + ] + }, + { + "type": "or", + "elements": [ + { + "type": "triggered_tags", + "triggers": [ + "\n\n\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "qwen_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n\n" + }, + { + "type": "tag", + "begin": "\n\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "qwen_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n\n" + } + ], + "at_least_one": true, + "stop_after_first": false, + "excludes": [ + "", + "" + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_content": { + "completion": "Need call the weather and add tools.\n\n\n{\"answer\":\"sunny\"}", + "generation_token_ids": [22916, 1562, 279, 8831, 321, 884, 7141, 13, 198, 248069, 271, 4754, 8944, 3147, 82, 26415, 8934, 248046] + }, + "auto_schema_calls": { + "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n\n\n\n\n\n1.0\n\n\n2\n\n\n[\"left\", \"right\"]\n\n\n", + "generation_token_ids": [22916, 1562, 279, 8831, 321, 884, 7141, 13, 198, 248069, 271, 248058, 198, 27, 1628, 27362, 67017, 29, 198, 27, 15704, 28, 2447, 29, 198, 1948, 29086, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 198, 248058, 198, 27, 1628, 89587, 29, 198, 27, 15704, 28408, 29, 198, 16, 13, 15, 198, 510, 15704, 29, 198, 27, 15704, 24582, 29, 198, 17, 198, 510, 15704, 29, 198, 27, 15704, 28, 3494, 29, 198, 1143, 2282, 487, 328, 1246, 1293, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 248046] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.txt b/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.txt index b0577806494b..56831217d83d 100644 --- a/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.txt +++ b/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.txt @@ -22,3 +22,21 @@ sequence tag `` text `` `\n\n` tag `\n\n` qwen_xml({ location: string }) `\n\n` + +# none_schema (from_token_zero) +sequence + sequence + tag `` text `` + `\n\n` + json({ answer: string }) + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + sequence + tag `` text `` + `\n\n` + or + triggered_tags [`\n`, ``] at_least_one + tag `\n\n` qwen_xml({ location: string }) `\n\n` + tag `\n\n` qwen_xml({ y: number, x: number, items: string[] }) `\n\n` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/Seed-OSS-36B-Instruct.json b/rust/src/chat/tests/grammar_replay/Seed-OSS-36B-Instruct.json new file mode 100644 index 000000000000..e933ae4abbb3 --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/Seed-OSS-36B-Instruct.json @@ -0,0 +1,40 @@ +{ + "model_id": "ByteDance-Seed/Seed-OSS-36B-Instruct", + "vocab_size": 155136, + "all_stop_token_ids": [2], + "grammars": [ + { + "grammar": { + "format": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "coverage": "final_output_only" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.\n{\"answer\":\"sunny\"}", + "generation_token_ids": [3, 40143, 2434, 405, 11785, 440, 1731, 9633, 141, 4, 326, 13059, 35733, 15301, 40841, 2313, 16320, 2] + }, + "auto_schema_content": { + "completion": "Need call the weather and add tools.\n{\"answer\":\"sunny\"}", + "generation_token_ids": [3, 40143, 2434, 405, 11785, 440, 1731, 9633, 141, 4, 326, 13059, 35733, 15301, 40841, 2313, 16320, 2] + } + } + } + ] +} diff --git a/rust/src/chat/tests/grammar_replay/Seed-OSS-36B-Instruct.txt b/rust/src/chat/tests/grammar_replay/Seed-OSS-36B-Instruct.txt new file mode 100644 index 000000000000..f54e72d4c825 --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/Seed-OSS-36B-Instruct.txt @@ -0,0 +1,2 @@ +# none_schema, auto_schema_content (final_output_only) +json({ answer: string }) diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index 079ce3e1a7e7..c5080a59c5c2 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -7,9 +7,10 @@ //! parsed from the generated assistant completion and then rendered back to the exact same //! assistant completion. //! -//! The tool-call fixture also snapshots the output grammar built by the initialized parser for each -//! tool-choice variant, with the rendered completion as its generation, as one file per model under -//! `tests/grammar_replay/`. `grammar_replay.rs` replays these cases through the real XGrammar. +//! The tool-call fixture also snapshots the output grammars built by the initialized parser, one +//! entry per distinct grammar with the rendered completion of every variant that builds it as a +//! generation, as one file per model under `tests/grammar_replay/`. `grammar_replay.rs` replays +//! these generations through the real XGrammar. use std::collections::{BTreeMap, BTreeSet}; use std::fmt::Debug; @@ -30,6 +31,7 @@ use vllm_chat::{ GenerationPromptMode, LoadModelBackendsOptions, NewChatOutputProcessorOptions, ParserSelection, RendererSelection, ToolStrictLevel, load_model_backends, }; +use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams; use vllm_parser::output_grammar::BuiltOutputGrammar; use vllm_text::{DecodedTextEvent, Finished, Prompt}; use vllm_tokenizer::Tokenizer; @@ -57,6 +59,13 @@ struct RoundtripCase { json_fmt: JsonFmt, /// Whether the template renders tool-call argument object keys in sorted order. sort_json_keys: bool, + /// Which tool-call fixture variants with a JSON schema this case runs; + /// every case runs the variants without one. The variants answered with + /// content need a template that renders a final turn without tool calls + /// back to its completion (some templates drop its reasoning or its stop + /// suffix); the variant answered with calls needs a tool grammar, which + /// holds the calls next to the schema. + with_schema: &'static [ToolCallMixVariant], } #[derive(Clone, Copy)] @@ -125,10 +134,11 @@ impl ThinkingBehavior { } } -/// Tool-choice variant of the tool-call fixture. +/// Variant of the tool-call fixture: the request's tool choice and optional +/// JSON schema, and the completion that answers it. #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize)] #[serde(rename_all = "snake_case")] -enum ToolChoiceVariant { +enum ToolCallMixVariant { /// `auto`, with content before the tool calls. Auto, /// `required`, with tool calls only: builder grammars admit no content @@ -136,18 +146,59 @@ enum ToolChoiceVariant { Required, /// Named `get_weather`, with a single call to it and no content. Named, + + /// `none` with a JSON schema, answered with schema-valid content. + NoneSchema, + /// `auto` with a JSON schema, answered with schema-valid content. + AutoSchemaContent, + /// `auto` with a JSON schema, answered with tool calls only: builder + /// grammars hold either schema-valid content or the calls. + AutoSchemaCalls, } -impl ToolChoiceVariant { - const ALL: [Self; 3] = [Self::Auto, Self::Required, Self::Named]; +/// The schema-valid content of the schema variants answered with content. +const SCHEMA_CONTENT: &str = r#"{"answer":"sunny"}"#; + +impl ToolCallMixVariant { + const WITHOUT_SCHEMA: &[Self] = &[Self::Auto, Self::Required, Self::Named]; + const WITH_SCHEMA: &[Self] = &[ + Self::NoneSchema, + Self::AutoSchemaContent, + Self::AutoSchemaCalls, + ]; fn tool_choice(self) -> ChatToolChoice { match self { - Self::Auto => ChatToolChoice::Auto, + Self::Auto | Self::AutoSchemaContent | Self::AutoSchemaCalls => ChatToolChoice::Auto, Self::Required => ChatToolChoice::Required, Self::Named => ChatToolChoice::Function { name: "get_weather".to_string(), }, + Self::NoneSchema => ChatToolChoice::None, + } + } + + fn structured_outputs(self) -> Option { + matches!( + self, + Self::NoneSchema | Self::AutoSchemaContent | Self::AutoSchemaCalls + ) + .then(|| { + StructuredOutputsParams::json(serde_json::json!({ + "type": "object", + "properties": { "answer": { "type": "string" } }, + "required": ["answer"] + })) + }) + } + + /// The expected text content and number of tool calls. + fn expected(self) -> (Option<&'static str>, usize) { + match self { + Self::Auto => (Some("I will call the tools."), 2), + Self::Required | Self::AutoSchemaCalls => (None, 2), + Self::Named => (None, 1), + Self::NoneSchema | Self::AutoSchemaContent => (Some(SCHEMA_CONTENT), 0), } } } @@ -163,6 +214,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: spaced_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -176,6 +228,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -189,6 +242,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -202,6 +256,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -215,6 +270,12 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + // The MiniMax M3 tool parser builds no tool grammar, so the schema + // excludes the calls. + with_schema: &[ + ToolCallMixVariant::NoneSchema, + ToolCallMixVariant::AutoSchemaContent, + ], } } @@ -228,6 +289,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: false }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -241,6 +303,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -254,6 +317,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: false }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -267,6 +331,12 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + // The GLM-4.5 tool parser builds no tool grammar, so the schema + // excludes the calls. + with_schema: &[ + ToolCallMixVariant::NoneSchema, + ToolCallMixVariant::AutoSchemaContent, + ], } } @@ -280,6 +350,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -293,6 +364,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -306,6 +378,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), sort_json_keys: true, + with_schema: &[ToolCallMixVariant::AutoSchemaCalls], } } @@ -329,6 +402,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: spaced_json_fmt(), sort_json_keys: false, + with_schema: &[ToolCallMixVariant::AutoSchemaCalls], } } @@ -346,6 +420,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -363,6 +438,7 @@ impl RoundtripCase { }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: &[ToolCallMixVariant::AutoSchemaCalls], } } @@ -376,6 +452,12 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + // The Seed-OSS tool parser builds no tool grammar, so the schema + // excludes the calls. + with_schema: &[ + ToolCallMixVariant::NoneSchema, + ToolCallMixVariant::AutoSchemaContent, + ], } } @@ -389,6 +471,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -402,6 +485,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -415,6 +499,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: &[ToolCallMixVariant::AutoSchemaCalls], } } @@ -428,6 +513,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), sort_json_keys: true, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } } @@ -545,8 +631,10 @@ async fn run_roundtrip_tool_call_mix( backends: &vllm_chat::LoadedModelBackends, ) -> Result<()> { let mut results = Vec::new(); - for variant in ToolChoiceVariant::ALL { - let result = run_roundtrip_tool_call_mix_inner(case, backends, variant).await?; + for &variant in ToolCallMixVariant::WITHOUT_SCHEMA.iter().chain(case.with_schema) { + let result = run_roundtrip_tool_call_mix_inner(case, backends, variant) + .await + .with_context(|| format!("variant {}", serde_json::to_value(variant).unwrap()))?; results.push((variant, result)); } check_grammar_replay_file(case, backends, results) @@ -555,9 +643,9 @@ async fn run_roundtrip_tool_call_mix( async fn run_roundtrip_tool_call_mix_inner( case: &RoundtripCase, backends: &vllm_chat::LoadedModelBackends, - variant: ToolChoiceVariant, + variant: ToolCallMixVariant, ) -> Result { - let request = roundtrip_request( + let mut request = roundtrip_request( "roundtrip-reasoning-tools", vec![ChatMessage::text( ChatRole::User, @@ -568,12 +656,9 @@ async fn run_roundtrip_tool_call_mix_inner( Some(true), // always enable thinking in this fixture case.thinking_behavior, ); + request.sampling_params.structured_outputs = variant.structured_outputs(); let expected_reasoning = "Need call the weather and add tools."; - let (expected_text, tool_call_count) = match variant { - ToolChoiceVariant::Auto => (Some("I will call the tools."), 2), - ToolChoiceVariant::Required => (None, 2), - ToolChoiceVariant::Named => (None, 1), - }; + let (expected_text, tool_call_count) = variant.expected(); let mut content = vec![AssistantContentBlock::Reasoning { text: expected_reasoning.to_string(), @@ -617,11 +702,13 @@ async fn run_roundtrip_tool_call_mix_inner( "parsed message: {:#?}", result.parsed_message ); - assert_eq!(tool_calls[0].name, "get_weather"); - assert_eq!( - tool_calls[0].arguments, - expected_arguments(case, r#"{"location": "Shanghai"}"#)?, - ); + if let Some(weather) = tool_calls.first() { + assert_eq!(weather.name, "get_weather"); + assert_eq!( + weather.arguments, + expected_arguments(case, r#"{"location": "Shanghai"}"#)?, + ); + } if let Some(add) = tool_calls.get(1) { assert_eq!(add.name, "add"); assert_eq!( @@ -1117,16 +1204,23 @@ struct GrammarReplayFile { /// Request stop set, passed to the matcher as `override_stop_tokens`. #[serde(serialize_with = "serialize_one_line")] all_stop_token_ids: BTreeSet, - /// Cases keyed by tool-choice variant. - cases: BTreeMap, + /// Distinct output grammars, in the order of the first variant building each. + grammars: Vec, } -/// One parser-built output grammar with the generation it must accept. +/// One parser-built output grammar with the generations it must accept. #[derive(Serialize)] struct GrammarReplayCase { /// Parser-built grammar; only token-zero grammars are replayed from the /// first generated token. grammar: BuiltOutputGrammar, + /// Generations keyed by the fixture variants whose requests build this grammar. + generations: BTreeMap, +} + +/// One rendered completion replayed against its output grammar. +#[derive(Serialize)] +struct GrammarReplayGeneration { /// Generated completion text, without the stop token. completion: String, /// Generated token IDs presented to the matcher, ending with a stop token. @@ -1150,13 +1244,13 @@ fn grammar_replay_dir() -> PathBuf { PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/grammar_replay") } -/// Snapshot the output grammars of one model's roundtrips, keyed by variant, -/// as its grammar replay file, or check that no stale file remains when the +/// Snapshot the output grammars of one model's roundtrips, grouping the +/// variants that build the same grammar, as its grammar replay file, or check that no stale file remains when the /// parser builds none. fn check_grammar_replay_file( case: &RoundtripCase, backends: &vllm_chat::LoadedModelBackends, - results: Vec<(ToolChoiceVariant, RoundtripResult)>, + results: Vec<(ToolCallMixVariant, RoundtripResult)>, ) -> Result<()> { let model_name = case.model_id.rsplit_once('/').map_or(case.model_id, |(_, name)| name); let path = grammar_replay_dir().join(format!("{model_name}.json")); @@ -1197,7 +1291,7 @@ fn check_grammar_replay_file( case.assistant_stop_suffix ); - let mut cases = BTreeMap::new(); + let mut grammars = Vec::::new(); for (variant, grammar, closed_completion) in results { let (completion, mut generation_token_ids) = match completion_body( tokenizer.as_ref(), @@ -1208,27 +1302,41 @@ fn check_grammar_replay_file( CompletionBody::TokenIds(body) => (tokenizer.decode(body, false)?, body.to_vec()), }; generation_token_ids.push(stop_token_id); - cases.insert( - variant, - GrammarReplayCase { + let generation = GrammarReplayGeneration { + completion, + generation_token_ids, + }; + match grammars.iter_mut().find(|case| case.grammar == grammar) { + Some(case) => { + case.generations.insert(variant, generation); + } + None => grammars.push(GrammarReplayCase { grammar, - completion, - generation_token_ids, - }, - ); + generations: BTreeMap::from([(variant, generation)]), + }), + } } // A readable outline of the same grammars, for review only. let mut outline = String::new(); - for (variant, case) in &cases { + for case in &grammars { if !outline.is_empty() { outline.push('\n'); } - let variant = serde_json::to_value(variant)?; + let variants = case + .generations + .keys() + .map(|variant| { + serde_json::to_value(variant)? + .as_str() + .map(str::to_owned) + .context("variant name") + }) + .collect::>>()?; let coverage = serde_json::to_value(case.grammar.coverage)?; outline.push_str(&format!( "# {} ({})\n", - variant.as_str().context("variant name")?, + variants.join(", "), coverage.as_str().context("coverage name")?, )); outline.push_str(&vllm_parser::output_grammar::test_utils::outline( @@ -1241,7 +1349,7 @@ fn check_grammar_replay_file( model_id: case.model_id, vocab_size: backends.text_backend.model_vocab_size(), all_stop_token_ids, - cases, + grammars, }; let mut json = serde_json::to_string_pretty(&file)?; json.push('\n'); diff --git a/rust/src/parser/src/output_grammar.rs b/rust/src/parser/src/output_grammar.rs index 0108c3193ef7..ba9b4968d54b 100644 --- a/rust/src/parser/src/output_grammar.rs +++ b/rust/src/parser/src/output_grammar.rs @@ -132,6 +132,10 @@ pub struct OutputGrammarContext<'a> { pub tool_strict_level: ToolStrictLevel, /// Whether the request permits more than one tool call. pub parallel_tool_calls: bool, + /// The user's constraint on the answer text, such as a `response_format` + /// JSON schema, normalized to a format. A builder that composes it inserts + /// it where the model writes its answer; `None` leaves the answer free. + pub answer: Option<&'a Format>, } /// Errors produced while building an output grammar. @@ -243,9 +247,12 @@ fn tool_grammar_applies(ctx: &OutputGrammarContext<'_>) -> bool { match ctx.tool_choice { ToolChoice::Value(ToolChoiceValue::None) => false, + // With an answer constraint, the grammar must hold the calls `auto` + // allows, or the answer would exclude them. ToolChoice::Value(ToolChoiceValue::Auto) => { ctx.tool_strict_level >= ToolStrictLevel::Function || ctx.tools.iter().any(|tool| tool.strict == Some(true)) + || ctx.answer.is_some() } ToolChoice::Value(ToolChoiceValue::Required) | ToolChoice::NamedFunction(_) @@ -284,6 +291,7 @@ mod tests { tool_choice, tool_strict_level: ToolStrictLevel::Auto, parallel_tool_calls: true, + answer: None, }, ) .unwrap() @@ -310,6 +318,7 @@ mod tests { tool_choice: &tool_choice, tool_strict_level: level, parallel_tool_calls: true, + answer: None, }, ) .unwrap() diff --git a/rust/src/parser/src/tool/minimax_m2.rs b/rust/src/parser/src/tool/minimax_m2.rs index 3b868013bb32..33d316579f27 100644 --- a/rust/src/parser/src/tool/minimax_m2.rs +++ b/rust/src/parser/src/tool/minimax_m2.rs @@ -122,7 +122,7 @@ impl ToolParser for MinimaxM2ToolParser { Some(xgrammar_structural_tag::Model::Minimax.builder()) } - fn build_visible_format( + fn build_call_format( &self, ctx: &OutputGrammarContext<'_>, ) -> output_grammar::Result> { diff --git a/rust/src/parser/src/tool/mod.rs b/rust/src/parser/src/tool/mod.rs index e4afbda3e902..974d6c758a27 100644 --- a/rust/src/parser/src/tool/mod.rs +++ b/rust/src/parser/src/tool/mod.rs @@ -20,6 +20,7 @@ mod seed_oss; #[cfg(any(test, feature = "test-util"))] pub mod test_utils; use std::collections::{BTreeMap, btree_map}; +use std::sync::Once; pub use deepseek_dsml::{DeepSeekV4ToolParser, DeepSeekV32ToolParser, DeepSeekV41ToolParser}; pub use deepseek_json::{DeepSeekV3ToolParser, DeepSeekV31ToolParser}; @@ -39,8 +40,11 @@ pub use qwen_coder::Qwen3CoderToolParser; pub use seed_oss::SeedOssToolParser; use serde::{Deserialize, Serialize}; use serde_json::Value; +use tracing::warn; +use xgrammar_structural_tag::ToolChoice; pub use xgrammar_structural_tag::builders::StructuralTagBuilder; use xgrammar_structural_tag::format::Format; +use xgrammar_structural_tag::tool::{AllowedToolsMode, ToolChoiceValue}; use crate::output_grammar::{self, OutputGrammarContext}; use crate::utils; @@ -196,21 +200,54 @@ pub trait ToolParser: Send { } /// Build the language of everything the model may emit after reasoning - /// ends: visible text plus tool calls. `None` means no tool grammar applies. + /// ends for the context's tool choice: visible text plus tool calls. `None` + /// means no tool grammar applies. /// /// The default forwards to [`Self::structural_tag_builder`] with reasoning /// disabled, which is exactly today's strict-tool-calling grammar. - fn build_visible_format( + fn build_call_format( &self, ctx: &OutputGrammarContext<'_>, ) -> output_grammar::Result> { output_grammar::visible_format_from_builder(self.structural_tag_builder(), ctx) } + /// Build the visible language with the request's answer constraint + /// composed in, as the Python frontend does: an optional call becomes + /// `Or[calls built as required, answer]`, and a tool choice that obliges a + /// call ignores the answer. Parsers override [`Self::build_call_format`] + /// instead. + fn build_visible_format( + &self, + ctx: &OutputGrammarContext<'_>, + ) -> output_grammar::Result> { + let Some(answer) = ctx.answer else { + return self.build_call_format(ctx); + }; + match call_required(ctx.tool_choice) { + // Calls built from the optional choice would also admit free text + // in place of the answer. + Some(tool_choice) => { + let calls = self.build_call_format(&OutputGrammarContext { + tool_choice: &tool_choice, + ..*ctx + })?; + Ok(calls.map(|calls| Format::or(vec![calls, answer.clone()]))) + } + None => { + let calls = self.build_call_format(ctx)?; + if calls.is_some() { + warn_answer_ignored_once(); + } + Ok(calls) + } + } + } + /// Return the xgrammar structural-tag builder used for strict tool calling. /// /// To be deprecated: this only exists to back the default - /// [`Self::build_visible_format`]. New parsers should override that method + /// [`Self::build_call_format`]. New parsers should override that method /// directly instead of exposing a builder. fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { None @@ -243,6 +280,32 @@ pub trait ToolParser: Send { fn reset(&mut self) -> String; } +/// The tool choice with an optional call made obligatory, or `None` when the +/// choice forbids or already obliges a call. +fn call_required(tool_choice: &ToolChoice) -> Option { + let mut required = tool_choice.clone(); + match &mut required { + ToolChoice::Value(value @ ToolChoiceValue::Auto) => *value = ToolChoiceValue::Required, + ToolChoice::AllowedTools(choice) if choice.allowed_tools.mode == AllowedToolsMode::Auto => { + choice.allowed_tools.mode = AllowedToolsMode::Required; + } + ToolChoice::FlatAllowedTools(choice) if choice.mode == AllowedToolsMode::Auto => { + choice.mode = AllowedToolsMode::Required; + } + _ => return None, + } + Some(required) +} + +/// Warn that a request's answer constraint is ignored because its tool +/// choice obliges a call, as the Python frontend does. +fn warn_answer_ignored_once() { + static ONCE: Once = Once::new(); + ONCE.call_once(|| { + warn!("structured outputs are ignored because tool_choice forces a tool call"); + }); +} + /// Extension methods for easily testing `ToolParser` implementations. /// /// These helpers do not handle partial parsing or error recovery, so they are diff --git a/rust/src/parser/src/unified/combined.rs b/rust/src/parser/src/unified/combined.rs index ba44e7fb9613..b1f36e885546 100644 --- a/rust/src/parser/src/unified/combined.rs +++ b/rust/src/parser/src/unified/combined.rs @@ -87,10 +87,19 @@ impl UnifiedParser for CombinedParser { &self, ctx: &OutputGrammarContext<'_>, ) -> output_grammar::Result> { - let Some(tool) = self.tool.as_ref() else { + // Without an active reasoning or tool parser, the request's own + // constraint applies unchanged. + if self.reasoning.is_none() && self.tool.is_none() { return Ok(None); + } + + // The tool parser composes the answer constraint into its calls where + // it can; without a tool grammar, the answer is the visible language. + let calls = match self.tool.as_ref() { + Some(tool) => tool.build_visible_format(ctx)?, + None => None, }; - let Some(visible) = tool.build_visible_format(ctx)? else { + let Some(visible) = calls.or_else(|| ctx.answer.cloned()) else { return Ok(None); }; // The reasoning parser may wrap that language with the reasoning phase @@ -165,7 +174,8 @@ mod tests { use super::CombinedParser; use crate::output_grammar::{ - GrammarCoverage, OutputGrammarContext, full_format_from_builder_for_test, + BuiltOutputGrammar, GrammarCoverage, OutputGrammarContext, + full_format_from_builder_for_test, }; use crate::reasoning::{ DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser, Glm45ReasoningParser, @@ -235,6 +245,7 @@ mod tests { tool_choice: &tool_choice, tool_strict_level: Default::default(), parallel_tool_calls: true, + answer: None, }; let tool = T::create(&tools).unwrap(); let expected = full_format_from_builder_for_test( @@ -524,6 +535,7 @@ mod tests { tool_choice: &tool_choice, tool_strict_level: Default::default(), parallel_tool_calls: true, + answer: None, }; let reasoning = Qwen3ReasoningParser::create(Arc::new(tokenizer())).unwrap(); let tool = Qwen3XmlToolParser::create(&tools).unwrap(); @@ -553,6 +565,7 @@ mod tests { tool_choice: &tool_choice, tool_strict_level: Default::default(), parallel_tool_calls: true, + answer: None, }; let visible = Format::const_string("answer"); @@ -586,6 +599,7 @@ mod tests { tool_choice: &tool_choice, tool_strict_level: Default::default(), parallel_tool_calls: true, + answer: None, }; let visible = Format::const_string("visible"); @@ -622,6 +636,7 @@ mod tests { tool_choice: &tool_choice, tool_strict_level: Default::default(), parallel_tool_calls: true, + answer: None, }; let tokenizer = Arc::new(tokenizer()); let reasoning = Qwen3ReasoningParser::create(tokenizer.clone()).unwrap(); @@ -650,4 +665,97 @@ mod tests { ); } } + + fn grammar_ctx<'a>( + tools: &'a [Tool], + tool_choice: &'a ToolChoice, + answer: Option<&'a Format>, + ) -> OutputGrammarContext<'a> { + OutputGrammarContext { + tools, + tool_choice, + tool_strict_level: Default::default(), + parallel_tool_calls: true, + answer, + } + } + + fn qwen3_parser(tools: &[Tool], with_tool: bool) -> CombinedParser { + let reasoning = Qwen3ReasoningParser::create(Arc::new(tokenizer())).unwrap(); + let tool = with_tool.then(|| Qwen3XmlToolParser::create(tools).unwrap()); + let mut parser = CombinedParser::new(Some(reasoning), tool); + parser.initialize(&[]).unwrap(); + parser + } + + #[test] + fn answer_constraint_composes_with_the_tool_choice() { + let tools = test_tools(); + let answer = Format::json_schema(serde_json::json!({ "type": "object" })); + let reasoning = |visible| { + Format::sequence(vec![ + Format::optional(Format::sequence(vec![ + Format::tag("", Format::any_text(), ""), + Format::const_string("\n\n"), + ])), + visible, + ]) + }; + let required = ToolChoice::required(); + let required_calls = Qwen3XmlToolParser::create(&tools) + .unwrap() + .build_visible_format(&grammar_ctx(&tools, &required, None)) + .unwrap() + .unwrap(); + + // No callable tool: the answer alone, after the reasoning phase. + for (tool_choice, with_tool) in [(ToolChoice::none(), true), (ToolChoice::auto(), false)] { + let parser = qwen3_parser(&tools, with_tool); + let ctx = grammar_ctx(&tools, &tool_choice, Some(&answer)); + let actual = parser.build_output_grammar(&ctx).unwrap().unwrap(); + assert_eq!(actual.coverage, GrammarCoverage::FromTokenZero); + assert_eq!(actual.format, reasoning(answer.clone()), "{tool_choice:?}"); + } + + // `auto`, strict or not, keeps tool calls next to the answer. + let parser = qwen3_parser(&tools, true); + let auto = ToolChoice::auto(); + let actual = parser.build_output_grammar(&grammar_ctx(&tools, &auto, Some(&answer))); + let expected = Format::or(vec![required_calls, answer.clone()]); + assert_eq!(actual.unwrap().unwrap().format, reasoning(expected)); + + // A required call ignores the answer constraint. + for tool_choice in [ToolChoice::required(), ToolChoice::function("get_weather")] { + let with_answer = grammar_ctx(&tools, &tool_choice, Some(&answer)); + let without = grammar_ctx(&tools, &tool_choice, None); + assert_eq!( + parser.build_output_grammar(&with_answer).unwrap(), + parser.build_output_grammar(&without).unwrap(), + "{tool_choice:?}" + ); + } + } + + #[test] + fn answer_without_a_reasoning_phase_covers_the_final_output() { + let tools = test_tools(); + let answer = Format::regex("[a-z]+"); + let none = ToolChoice::none(); + let auto = ToolChoice::auto(); + + // Without an active parser, the request keeps its own constraint. + let parser = CombinedParser::new(None, None); + let ctx = grammar_ctx(&tools, &auto, Some(&answer)); + assert!(parser.build_output_grammar(&ctx).unwrap().is_none()); + + let parser = CombinedParser::new(None, Some(Qwen3XmlToolParser::create(&tools).unwrap())); + let ctx = grammar_ctx(&tools, &none, Some(&answer)); + let alone = parser.build_output_grammar(&ctx).unwrap().unwrap(); + assert_eq!(alone, BuiltOutputGrammar::final_output_only(answer.clone())); + + let ctx = grammar_ctx(&tools, &auto, Some(&answer)); + let composed = parser.build_output_grammar(&ctx).unwrap().unwrap(); + assert_eq!(composed.coverage, GrammarCoverage::FinalOutputOnly); + assert!(matches!(composed.format, Format::Or(_))); + } } diff --git a/rust/src/parser/src/unified/hy/hy_v3.rs b/rust/src/parser/src/unified/hy/hy_v3.rs index 3599ce2cb9a0..18a34954b3b6 100644 --- a/rust/src/parser/src/unified/hy/hy_v3.rs +++ b/rust/src/parser/src/unified/hy/hy_v3.rs @@ -164,6 +164,7 @@ mod tests { tools: &tools, tool_strict_level: Default::default(), parallel_tool_calls: true, + answer: None, tool_choice: &ToolChoice::required(), }) .unwrap() @@ -194,6 +195,7 @@ mod tests { tool_choice: &ToolChoice::required(), tool_strict_level: Default::default(), parallel_tool_calls: true, + answer: None, }) .unwrap() .unwrap(); diff --git a/rust/src/parser/src/unified/hy/hy_v4.rs b/rust/src/parser/src/unified/hy/hy_v4.rs index a5d81e0f2ae7..f10fb662eca3 100644 --- a/rust/src/parser/src/unified/hy/hy_v4.rs +++ b/rust/src/parser/src/unified/hy/hy_v4.rs @@ -178,6 +178,7 @@ mod tests { tools: &tools, tool_strict_level: Default::default(), parallel_tool_calls: true, + answer: None, tool_choice: &ToolChoice::required(), }) .unwrap() diff --git a/rust/src/parser/src/unified/kimi_k3/grammar.rs b/rust/src/parser/src/unified/kimi_k3/grammar.rs index 166fbafb0f82..e6ef746fa2ac 100644 --- a/rust/src/parser/src/unified/kimi_k3/grammar.rs +++ b/rust/src/parser/src/unified/kimi_k3/grammar.rs @@ -16,6 +16,10 @@ //! //! Every path may end with the message close, and a stop token is legal only //! after the last channel closes. +//! +//! The user's answer constraint, if any, replaces the free text of the +//! `response` channel. Without a callable tool, the `response` channel then +//! follows `think` alone. mod arguments; @@ -33,38 +37,44 @@ use crate::output_grammar::{ }; /// Build the grammar of everything generated after a prompt that left the -/// parser in `mode`, or `None` when the request asks for no tool grammar. +/// parser in `mode`, or `None` when the request asks for no tool grammar and +/// has no answer constraint. pub(super) fn build_output_grammar( mode: &KimiK3Mode, ctx: &OutputGrammarContext<'_>, ) -> output_grammar::Result> { - let Some(NormalizedToolChoice { - function_tools, - choice, - .. - }) = normalize_tool_choice(ctx)? - else { + let normalized = normalize_tool_choice(ctx)?; + if normalized.is_none() && ctx.answer.is_none() { return Ok(None); - }; - - let single_call = choice == BuilderToolChoice::Forced || !ctx.parallel_tool_calls; - let tools = Format::tag( - TOOLS_OPEN, - Format::tags_with_separator( - function_tools.iter().map(call_tag).collect(), - "", - true, - single_call, - ), - TOOLS_CLOSE, + } + let calls = normalized.map( + |NormalizedToolChoice { + function_tools, + choice, + .. + }| { + let single_call = choice == BuilderToolChoice::Forced || !ctx.parallel_tool_calls; + let tools = Format::tag( + TOOLS_OPEN, + Format::tags_with_separator( + function_tools.iter().map(call_tag).collect(), + "", + true, + single_call, + ), + TOOLS_CLOSE, + ); + (choice, tools) + }, ); + let answer = ctx.answer; let mut elements = match mode { - KimiK3Mode::Reasoning => vec![think_body(), after_reasoning(choice, tools)], - KimiK3Mode::Response => vec![response_body(), after_response(choice, tools)], + KimiK3Mode::Reasoning => vec![think_body(), after_reasoning(calls, answer)], + KimiK3Mode::Response => rest_of_response(calls, answer), KimiK3Mode::Idle => vec![ Format::optional(Format::tag(THINK_OPEN, think_text(), THINK_CLOSE)), - after_reasoning(choice, tools), + after_reasoning(calls, answer), ], KimiK3Mode::Epilogue | KimiK3Mode::Tools | KimiK3Mode::Call { .. } | KimiK3Mode::Done => { unreachable!("initialization leaves the parser idle or inside `think` or `response`") @@ -77,38 +87,73 @@ pub(super) fn build_output_grammar( )))) } -/// Channels that may follow the closed `think` channel. -fn after_reasoning(tool_choice: BuilderToolChoice, tools: Format) -> Format { - let response = Format::tag(RESPONSE_OPEN, response_text(), RESPONSE_CLOSE); - match tool_choice { - BuilderToolChoice::Auto => Format::or(vec![ +/// Channels that may follow the closed `think` channel, given the tool +/// choice and `tools` channel when a tool is callable, and the user's answer +/// constraint for the `response` channel. +fn after_reasoning(calls: Option<(BuilderToolChoice, Format)>, answer: Option<&Format>) -> Format { + let response = Format::tag(RESPONSE_OPEN, response_content(answer), RESPONSE_CLOSE); + let Some((choice, tools)) = calls else { + return response; + }; + let empty = || Format::const_string(format!("{RESPONSE_OPEN}{RESPONSE_CLOSE}")); + match (choice, answer) { + (BuilderToolChoice::Auto, None) => Format::or(vec![ Format::sequence(vec![response, Format::optional(tools.clone())]), tools, ]), - BuilderToolChoice::Required | BuilderToolChoice::Forced => { + (BuilderToolChoice::Auto, Some(_)) => Format::or(vec![ + Format::sequence(vec![response, Format::optional(tools.clone())]), + Format::sequence(vec![Format::optional(empty()), tools]), + ]), + (BuilderToolChoice::Required | BuilderToolChoice::Forced, None) => { Format::sequence(vec![Format::optional(response), tools]) } + (BuilderToolChoice::Required | BuilderToolChoice::Forced, Some(_)) => { + Format::sequence(vec![ + Format::optional(Format::or(vec![response, empty()])), + tools, + ]) + } } } -/// Channels that may follow the closed `response` channel. -fn after_response(tool_choice: BuilderToolChoice, tools: Format) -> Format { - match tool_choice { - BuilderToolChoice::Auto => Format::optional(tools), - BuilderToolChoice::Required | BuilderToolChoice::Forced => tools, +/// The rest of the `response` channel the prompt opened, then the channels +/// that may follow it. +fn rest_of_response( + calls: Option<(BuilderToolChoice, Format)>, + answer: Option<&Format>, +) -> Vec { + let response = Format::tag("", response_content(answer), RESPONSE_CLOSE); + let Some((choice, tools)) = calls else { + return vec![response]; + }; + let empty = || Format::const_string(RESPONSE_CLOSE); + match (choice, answer) { + (BuilderToolChoice::Auto, None) => vec![response, Format::optional(tools)], + (BuilderToolChoice::Auto, Some(_)) => vec![Format::or(vec![ + Format::sequence(vec![response, Format::optional(tools.clone())]), + Format::sequence(vec![empty(), tools]), + ])], + (BuilderToolChoice::Required | BuilderToolChoice::Forced, None) => vec![response, tools], + (BuilderToolChoice::Required | BuilderToolChoice::Forced, Some(_)) => { + vec![Format::or(vec![response, empty()]), tools] + } } } +/// The text of the `response` channel: the user's answer constraint, or free +/// text. A constrained answer rules out the empty `response` channel the chat +/// template writes before tool calls, so callers allow that channel where +/// calls follow. +fn response_content(answer: Option<&Format>) -> Format { + answer.cloned().unwrap_or_else(response_text) +} + /// The rest of the `think` channel the prompt opened. fn think_body() -> Format { Format::tag("", think_text(), THINK_CLOSE) } -/// The rest of the `response` channel the prompt opened. -fn response_body() -> Format { - Format::tag("", response_text(), RESPONSE_CLOSE) -} - fn think_text() -> Format { Format::any_text_excluding(&[THINK_CLOSE, END_OF_MSG]) } @@ -158,6 +203,28 @@ mod tests { tool_choice: &tool_choice, tool_strict_level: ToolStrictLevel::Function, parallel_tool_calls, + answer: None, + }, + ) + .unwrap() + .unwrap(); + assert_eq!(grammar.coverage, GrammarCoverage::FromTokenZero); + let mut format = grammar.format; + elide_call_bodies(&mut format); + expected.assert_eq(&outline(&format)); + } + + /// Like [`check`], for non-strict tools and a JSON answer constraint. + fn check_answer(mode: KimiK3Mode, tools: &[Tool], tool_choice: ToolChoice, expected: Expect) { + let answer = Format::json_schema(json!({ "type": "object" })); + let grammar = build_output_grammar( + &mode, + &OutputGrammarContext { + tools, + tool_choice: &tool_choice, + tool_strict_level: ToolStrictLevel::Auto, + parallel_tool_calls: true, + answer: Some(&answer), }, ) .unwrap() @@ -348,9 +415,120 @@ mod tests { tool_choice: &ToolChoice::auto(), tool_strict_level: ToolStrictLevel::Auto, parallel_tool_calls: true, + answer: None, }, ) .unwrap(); assert_eq!(grammar, None); } + + #[test] + fn answer_constraint_fills_the_response_channel() { + // Non-strict `auto` keeps the calls next to the constrained answer. + // Calls may follow the empty `response` channel the chat template + // writes before them. + check_answer( + KimiK3Mode::Reasoning, + &tools(), + ToolChoice::auto(), + expect![[r#" + sequence + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + or + sequence + tag `<|open|>response<|sep|>` json(object) `<|close|>response<|sep|>` + optional + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + sequence + optional `<|open|>response<|sep|><|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + // A required call may follow a constrained answer. + check_answer( + KimiK3Mode::Reasoning, + &tools(), + ToolChoice::required(), + expect![[r#" + sequence + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + sequence + optional + or + tag `<|open|>response<|sep|>` json(object) `<|close|>response<|sep|>` + `<|open|>response<|sep|><|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + check_answer( + KimiK3Mode::Response, + &tools(), + ToolChoice::auto(), + expect![[r#" + sequence + or + sequence + tag `` json(object) `<|close|>response<|sep|>` + optional + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + sequence + `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + } + + #[test] + fn answer_constraint_without_a_callable_tool_requires_the_response() { + check_answer( + KimiK3Mode::Idle, + &tools(), + ToolChoice::none(), + expect![[r#" + sequence + optional tag `<|open|>think<|sep|>` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + tag `<|open|>response<|sep|>` json(object) `<|close|>response<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + check_answer( + KimiK3Mode::Reasoning, + &[], + ToolChoice::auto(), + expect![[r#" + sequence + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + tag `<|open|>response<|sep|>` json(object) `<|close|>response<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + check_answer( + KimiK3Mode::Response, + &[], + ToolChoice::auto(), + expect![[r#" + sequence + tag `` json(object) `<|close|>response<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + } } From 2534a96fcf165202f5358335f8b0705175159fd0 Mon Sep 17 00:00:00 2001 From: Jimmy Wang <53066820+JimmyWang0417@users.noreply.github.com> Date: Mon, 12 Oct 2026 17:47:52 +0800 Subject: [PATCH 288/624] [Docs] List additional unauthenticated endpoints (#61081) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Jimmy Wang <53066820+JimmyWang0417@users.noreply.github.com> Co-authored-by: Juan Pérez de Algaba <124347725+jperezdealgaba@users.noreply.github.com> --- docs/usage/security.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/usage/security.md b/docs/usage/security.md index f069de0def74..0d7894a467a2 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -208,9 +208,13 @@ The following endpoints **do not require authentication** even when `--api-key` - `/update_weights` - Update model weights (can alter model behavior) - `/get_world_size` - Get distributed world size - `/abort_requests` - Abort in-flight requests (available with `--tokens-only`. Use the authenticated `/inference/v1/abort_requests` otherwise) +- `/fault_tolerance/apply` - Trigger fault-recovery operations across ranks (available when fault tolerance is enabled) +- `/fault_tolerance/status` - Query fault-recovery status (available when fault tolerance is enabled) **Utility endpoints:** +- `/metrics` - Prometheus metrics, including model, throughput, queue, and KV-cache statistics +- `/kv_event_sources` - KV-event publisher configuration for each data-parallel rank, which may reveal internal transport endpoints (when KV events are configured) - `/tokenize` - Tokenize text (not gated by `--enable-scale-out`) - `/detokenize` - Detokenize tokens - `/health` - Health check From f7a3a7265f5161594721e8c1aa6b68a761629b68 Mon Sep 17 00:00:00 2001 From: MatthewCYM <40845677+MatthewCYM@users.noreply.github.com> Date: Mon, 12 Oct 2026 18:42:51 +0800 Subject: [PATCH 289/624] [Bugfix][Multimodal] Fix Qwen3-ASR audio preprocessing parity (#57146) Signed-off-by: MatthewCYM --- .../multimodal/processing/test_qwen3_asr.py | 181 ++++++++++++++ vllm/model_executor/models/qwen3_asr.py | 73 +++++- .../transformers_utils/processors/__init__.py | 2 - .../processors/qwen3_asr.py | 229 ------------------ 4 files changed, 243 insertions(+), 242 deletions(-) create mode 100644 tests/models/multimodal/processing/test_qwen3_asr.py delete mode 100644 vllm/transformers_utils/processors/qwen3_asr.py diff --git a/tests/models/multimodal/processing/test_qwen3_asr.py b/tests/models/multimodal/processing/test_qwen3_asr.py new file mode 100644 index 000000000000..b016a9d70860 --- /dev/null +++ b/tests/models/multimodal/processing/test_qwen3_asr.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Qwen3-ASR preprocessing must preserve its HF processor's audio inputs.""" + +import numpy as np +import pytest +import torch +from transformers.feature_extraction_utils import BatchFeature +from transformers.models.qwen3_asr import Qwen3ASRProcessor + +from vllm.config import ModelConfig +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.cache import MultiModalProcessorOnlyCache +from vllm.multimodal.inputs import batched_tensors_equal +from vllm.transformers_utils.processor import cached_processor_from_config + +from .test_common import _assert_inputs_equal + +_AUDIO_PROMPT = "<|audio_start|><|audio_pad|><|audio_end|>" + + +@pytest.fixture(params=["Qwen/Qwen3-ASR-0.6B", "Qwen/Qwen3-ASR-0.6B-hf"]) +def asr_processor(request): + config = ModelConfig( + request.param, + dtype="float32", + max_model_len=4096, + limit_mm_per_prompt={"audio": 2}, + mm_processor_cache_gb=0, + ) + processor = MULTIMODAL_REGISTRY.create_processor(config) + native = Qwen3ASRProcessor.from_pretrained(request.param) + return processor, native + + +@pytest.mark.parametrize( + "samples", [321, 15999, 16000, 16001, 16159, 16160, 16161, 32001, 480001] +) +@pytest.mark.parametrize("num_audios", [1, 2]) +def test_audio_processing_matches_hf(asr_processor, samples, num_audios): + """Each clip must match HF independently of other clips in the request.""" + processor, native = asr_processor + rng = np.random.RandomState(42) + audios = [ + rng.randn(samples + index * 160).astype(np.float32) + for index in range(num_audios) + ] + expected = [ + native(text=_AUDIO_PROMPT, audio=audio, return_tensors="pt") for audio in audios + ] + actual = processor( + _AUDIO_PROMPT * num_audios, + mm_items=processor.info.parse_mm_data({"audio": audios}), + hf_processor_mm_kwargs={}, + ) + + expected_ids = [ + token for item in expected for token in item["input_ids"][0].tolist() + ] + assert actual["prompt_token_ids"] == expected_ids + masks = [item["input_features_mask"][0] for item in expected] + expected_features = torch.cat( + [ + item["input_features"][0, :, mask.bool()] + for item, mask in zip(expected, masks) + ], + dim=1, + ) + actual_features = actual["mm_kwargs"].get_data() + torch.testing.assert_close( + actual_features["input_audio_features"], expected_features, rtol=0, atol=0 + ) + torch.testing.assert_close( + actual_features["audio_feature_lengths"], + torch.stack([mask.sum() for mask in masks]), + ) + + +@pytest.mark.parametrize("cached_indices", [[], [0], [1], [0, 1]]) +def test_audio_processing_cache_is_batch_independent(asr_processor, cached_indices): + """Partial cache hits must not change features, masks, or prompt expansion.""" + processor, _ = asr_processor + config = processor.info.ctx.model_config + config.get_multimodal_config().mm_processor_cache_gb = 1 + cache = MultiModalProcessorOnlyCache(config) + rng = np.random.RandomState(42) + audios = [rng.randn(samples).astype(np.float32) for samples in (16001, 32000)] + for index in cached_indices: + processor( + _AUDIO_PROMPT, + mm_items=processor.info.parse_mm_data({"audio": audios[index]}), + hf_processor_mm_kwargs={}, + cache=cache, + ) + + for batch in (audios, audios[::-1]): + kwargs = dict( + mm_items=processor.info.parse_mm_data({"audio": batch}), + hf_processor_mm_kwargs={}, + ) + _assert_inputs_equal( + processor(_AUDIO_PROMPT * 2, **kwargs), + processor(_AUDIO_PROMPT * 2, **kwargs, cache=cache), + ) + + +def test_text_only_processing(asr_processor): + """Empty media input must not invoke the audio feature extractor.""" + processor, native = asr_processor + prompt = "Transcribe the audio." + actual = processor( + prompt, + mm_items=processor.info.parse_mm_data({}), + hf_processor_mm_kwargs={}, + ) + assert actual["prompt_token_ids"] == native.tokenizer.encode(prompt) + assert actual["mm_kwargs"].get_data() == {} + + +@pytest.mark.parametrize("hf_data", [{}, {"text": ""}, {"audio": None}, {"audio": []}]) +def test_hf_processor_without_audio(asr_processor, hf_data): + """The media hook must not require audio or call HF for absent waveforms.""" + processor, _ = asr_processor + assert not processor._call_hf_processor(hf_data, {}) + + +@pytest.mark.parametrize("hf_data", [{}, {"text": ""}]) +def test_postprocessing_without_audio_features(asr_processor, hf_data): + """Non-audio output must pass through even when HF input is nonempty.""" + processor, _ = asr_processor + processed = BatchFeature({"input_ids": torch.tensor([[1, 2]])}) + actual = processor._postprocess_hf_mm_data(hf_data, {}, processed) + assert actual is processed + assert set(actual) == {"input_ids"} + torch.testing.assert_close(actual["input_ids"], torch.tensor([[1, 2]])) + + +def test_chat_template_processing(asr_processor): + """The generic processor loader must still render audio chat prompts.""" + processor, native = asr_processor + chat_processor = cached_processor_from_config(processor.info.ctx.model_config) + messages = [ + {"role": "system", "content": ""}, + {"role": "user", "content": [{"type": "audio", "audio": "unused"}]}, + ] + prompt = chat_processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + expected_prompt = native.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + assert prompt == expected_prompt + audio = np.random.RandomState(42).randn(16001).astype(np.float32) + expected = native(text=expected_prompt, audio=audio, return_tensors="pt") + actual = processor( + prompt, + mm_items=processor.info.parse_mm_data({"audio": audio}), + hf_processor_mm_kwargs={}, + ) + assert actual["prompt_token_ids"] == expected["input_ids"][0].tolist() + + +def test_precomputed_audio_features(asr_processor): + """Precomputed features must bypass waveform extraction unchanged.""" + processor, _ = asr_processor + processor.info.ctx.get_mm_config().enable_mm_embeds = True + audio = np.random.RandomState(42).randn(16001).astype(np.float32) + expected = processor( + _AUDIO_PROMPT, + mm_items=processor.info.parse_mm_data({"audio": audio}), + hf_processor_mm_kwargs={}, + ) + features = expected["mm_kwargs"].get_data() + actual = processor( + _AUDIO_PROMPT, + mm_items=processor.info.parse_mm_data({"audio": features}), + hf_processor_mm_kwargs={}, + ) + assert actual["prompt_token_ids"] == expected["prompt_token_ids"] + assert actual["mm_placeholders"] == expected["mm_placeholders"] + assert batched_tensors_equal(actual["mm_kwargs"].get_data(), features) diff --git a/vllm/model_executor/models/qwen3_asr.py b/vllm/model_executor/models/qwen3_asr.py index c8a0849fc56a..161d1efacd06 100644 --- a/vllm/model_executor/models/qwen3_asr.py +++ b/vllm/model_executor/models/qwen3_asr.py @@ -29,6 +29,7 @@ import torch import torch.nn as nn from transformers.feature_extraction_utils import BatchFeature +from transformers.models.qwen3_asr import Qwen3ASRFeatureExtractor, Qwen3ASRProcessor from transformers.models.whisper import WhisperFeatureExtractor from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig @@ -54,7 +55,6 @@ from vllm.model_executor.models.qwen3 import Qwen3ForCausalLM from vllm.model_executor.models.qwen3_omni_moe_thinker import ( Qwen3OmniMoeAudioEncoder, - Qwen3OmniMoeThinkerMultiModalProcessor, ) from vllm.model_executor.models.utils import ( AutoWeightsLoader, @@ -80,6 +80,7 @@ ) from vllm.multimodal.processing import ( BaseDummyInputsBuilder, + BaseMultiModalProcessor, BaseProcessingInfo, PromptReplacement, PromptUpdate, @@ -92,9 +93,6 @@ Qwen3ASRThinkerConfig, ) from vllm.transformers_utils.processor import cached_processor_from_config -from vllm.transformers_utils.processors.qwen3_asr import ( - Qwen3ASRProcessor, -) from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import async_tensor_h2d @@ -188,18 +186,19 @@ def get_hf_config(self): return self.ctx.get_hf_config(Qwen3ASRConfig).thinker_config def get_hf_processor(self, **kwargs: object) -> Qwen3ASRProcessor: - processor = self.ctx.get_hf_processor( + return self.ctx.get_hf_processor( Qwen3ASRProcessor, **kwargs, ) - if not hasattr(processor, "audio_token"): - processor.audio_token = "<|audio_pad|>" - return processor - def get_feature_extractor(self, **kwargs: object) -> WhisperFeatureExtractor: + def get_feature_extractor( + self, **kwargs: object + ) -> WhisperFeatureExtractor | Qwen3ASRFeatureExtractor: hf_processor = self.get_hf_processor(**kwargs) feature_extractor = hf_processor.feature_extractor - assert isinstance(feature_extractor, WhisperFeatureExtractor) + assert isinstance( + feature_extractor, (WhisperFeatureExtractor, Qwen3ASRFeatureExtractor) + ) return feature_extractor def get_supported_mm_limits(self) -> Mapping[str, int | None]: @@ -275,8 +274,60 @@ def _parse_audio_data( class Qwen3ASRMultiModalProcessor( - Qwen3OmniMoeThinkerMultiModalProcessor[Qwen3ASRProcessingInfo] + BaseMultiModalProcessor[Qwen3ASRProcessingInfo], ): + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + + def _call_hf_processor( + self, + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], + ) -> BatchFeature: + audios = hf_data.get("audio") + if not audios: + return BatchFeature() + assert isinstance(audios, list) + + if len(audios) == 1: + return super()._call_hf_processor(hf_data, hf_kwargs) + + # HF requires one text per audio, and batch padding changes clip features. + # Keep clips independent so partial cache hits cannot change their inputs. + text = self._get_hf_mm_text({"audio": 1}) + outputs: dict[str, list[torch.Tensor]] = { + "input_features": [], + "input_features_mask": [], + } + for audio in audios: + processed = super()._call_hf_processor( + dict(hf_data, text=text, audio=audio), hf_kwargs + ) + for key, values in outputs.items(): + values.append(processed[key][0]) + + return BatchFeature(outputs) + + def _postprocess_hf_mm_data( + self, + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], + processed_data: BatchFeature, + ) -> BatchFeature: + if "input_features" not in processed_data: + return processed_data + + features = processed_data.pop("input_features") + masks = processed_data.pop("input_features_mask") + processed_data["input_audio_features"] = torch.cat( + [feature[:, mask.bool()] for feature, mask in zip(features, masks)], dim=1 + ) + processed_data["feature_attention_mask"] = masks + processed_data["audio_feature_lengths"] = torch.stack( + [mask.sum() for mask in masks] + ) + return processed_data + def _get_mm_fields_config( self, hf_inputs: BatchFeature, diff --git a/vllm/transformers_utils/processors/__init__.py b/vllm/transformers_utils/processors/__init__.py index f2af4901e2e9..db45e49d5adc 100644 --- a/vllm/transformers_utils/processors/__init__.py +++ b/vllm/transformers_utils/processors/__init__.py @@ -38,7 +38,6 @@ "OpenVLAProcessor", "OvisProcessor", "Ovis2_5Processor", - "Qwen3ASRProcessor", "Step3VLProcessor", "InklingProcessor", "InklingImageProcessor", @@ -74,7 +73,6 @@ "OpenVLAProcessor": "vllm.transformers_utils.processors.openvla", "OvisProcessor": "vllm.transformers_utils.processors.ovis", "Ovis2_5Processor": "vllm.transformers_utils.processors.ovis2_5", - "Qwen3ASRProcessor": "vllm.transformers_utils.processors.qwen3_asr", "Step3VLProcessor": "vllm.transformers_utils.processors.step3_vl", "InklingProcessor": "vllm.transformers_utils.processors.inkling", "InklingImageProcessor": "vllm.transformers_utils.processors.inkling", diff --git a/vllm/transformers_utils/processors/qwen3_asr.py b/vllm/transformers_utils/processors/qwen3_asr.py deleted file mode 100644 index 55d38537928d..000000000000 --- a/vllm/transformers_utils/processors/qwen3_asr.py +++ /dev/null @@ -1,229 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# ruff: noqa -# mypy: ignore-errors -# coding=utf-8 -# Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import regex as re - -import numpy as np - -from transformers import AutoProcessor -from transformers.audio_utils import AudioInput -from transformers.feature_extraction_utils import BatchFeature -from transformers.processing_utils import ProcessingKwargs, ProcessorMixin -from transformers.tokenization_utils_base import TextInput - - -class Qwen3ASRProcessorKwargs(ProcessingKwargs, total=False): - _defaults = { - "text_kwargs": { - "padding": False, - "padding_side": "left", - }, - "audio_kwargs": { - "sampling_rate": 16000, - "padding": True, - "return_attention_mask": True, - }, - } - - -def _get_feat_extract_output_lengths(input_lengths): - """ - Computes the output length of the convolutional layers and the output length of the audio encoder - """ - - input_lengths_leave = input_lengths % 100 - feat_lengths = (input_lengths_leave - 1) // 2 + 1 - output_lengths = ( - ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13 - ) - return output_lengths - - -class Qwen3ASRProcessor(ProcessorMixin): - r""" - Constructs a Qwen3ASR processor. - [`Qwen3ASRProcessor`] offers all the functionalities of [`WhisperFeatureExtractor`], and [`Qwen2TokenizerFast`]. See the - [`~Qwen3ASRProcessor.__call__`] and [`~Qwen3ASRProcessor.decode`] for more information. - - Args: - feature_extractor ([`WhisperFeatureExtractor`], *optional*): - The audio feature extractor. - tokenizer ([`Qwen2TokenizerFast`], *optional*): - The text tokenizer. - chat_template (`Optional[str]`, *optional*): - The Jinja template to use for formatting the conversation. If not provided, the default chat template is used. - """ - - attributes = ["feature_extractor", "tokenizer"] - feature_extractor_class = "WhisperFeatureExtractor" - tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast") - - def __init__(self, feature_extractor=None, tokenizer=None, chat_template=None): - super().__init__(feature_extractor, tokenizer, chat_template=chat_template) - self.audio_token = self.tokenizer.audio_token - self.audio_bos_token = self.tokenizer.audio_bos_token - self.audio_eos_token = self.tokenizer.audio_eos_token - - def __call__( - self, - text: TextInput = None, - audio: AudioInput = None, - **kwargs, - ) -> BatchFeature: - """ - Main method to prepare for the model one or several sequences(s) and audio(s). This method forwards the `text` - and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode - the text. To prepare the audio(s), this method forwards the `audio` and `kwargs` arguments to - WhisperFeatureExtractor's [`~WhisperFeatureExtractor.__call__`] if `audio` is not `None`. Please refer to the doctsring - of the above two methods for more information. - - Args: - text (`str`, `List[str]`, `List[List[str]]`): - The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings - (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set - `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). - audio (`np.ndarray`, `List[np.ndarray]`): - The audio or batch of audio to be prepared. Each audio can be a NumPy array. - """ - - if text is None: - raise ValueError("You need to specify either a `text` input to process.") - - output_kwargs = self._merge_kwargs( - Qwen3ASRProcessorKwargs, - tokenizer_init_kwargs=self.tokenizer.init_kwargs, - **kwargs, - ) - - if audio is not None: - output_kwargs["audio_kwargs"]["padding"] = True - output_kwargs["audio_kwargs"]["truncation"] = False - audio_inputs = self.feature_extractor( - audio, **output_kwargs["audio_kwargs"] - ) - audio_inputs["feature_attention_mask"] = audio_inputs.pop( - "attention_mask" - ) # rename feature_attention_mask to prevent conflicts later on - audio_inputs["input_features"] = audio_inputs.pop( - "input_features" - ) # rename input_features to prevent conflicts later on - audio_lengths = iter( - _get_feat_extract_output_lengths( - audio_inputs["feature_attention_mask"].sum(-1) - ) - ) - else: - audio_inputs = {} - audio_lengths = iter([]) - - if not isinstance(text, list): - text = [text] - - text = self.replace_multimodal_special_tokens( - text, - audio_lengths, - ) - - texts_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) - - return BatchFeature( - data={**texts_inputs, **audio_inputs}, - tensor_type=kwargs.get("return_tensors"), - ) - - def replace_multimodal_special_tokens( - self, - text, - audio_lengths, - ): - processed_text = [] - for sample in text: - positions = [] - special_tokens = [re.escape(tok) for tok in [self.audio_token]] - pattern = "|".join(special_tokens) - positions = sorted( - [ - (match.start(), match.group()) - for match in re.finditer(pattern, sample) - ] - ) - positions.sort(key=lambda x: x[0]) - - for _, special_token in positions: - if special_token == self.audio_token: - sample = sample.replace( - self.audio_token, - "<|audio_placeholder|>" * next(audio_lengths), - 1, - ) - - sample = sample.replace("<|audio_placeholder|>", self.audio_token) - processed_text.append(sample) - return processed_text - - def get_chunked_index( - self, token_indices: np.ndarray, tokens_per_chunk: int - ) -> list[tuple[int, int]]: - """ - Splits token index list into chunks based on token value ranges. - - Given a list of token indices, returns a list of (start, end) index tuples representing - slices of the list where the token values fall within successive ranges of `tokens_per_chunk`. - - For example, if `tokens_per_chunk` is 1000, the function will create chunks such that: - - the first chunk contains token values < 1000, - - the second chunk contains values >= 1000 and < 2000, and so on. - - Parameters: - token_indices (`np.ndarray`): A monotonically increasing list of token index values. - tokens_per_chunk (`int`): Number of tokens per chunk (used as the chunk size threshold). - - Returns: - `list[tuple[int, int]]`: A list of tuples, each representing the start (inclusive) - and end (exclusive) indices of a chunk in `token_indices`. - """ - - def _iter(): - i, start_idx = 0, 0 # skip bos token - current_chunk = 1 - while i < len(token_indices): # skip eos token - if token_indices[i] >= current_chunk * tokens_per_chunk: - yield (start_idx, i) - start_idx = i - current_chunk += 1 - i += 1 - yield (start_idx, len(token_indices)) - - return list(_iter()) - - def apply_chat_template(self, conversations, chat_template=None, **kwargs): - kwargs["return_dict"] = False - return super().apply_chat_template(conversations, chat_template, **kwargs) - - @property - def model_input_names(self): - tokenizer_input_names = self.tokenizer.model_input_names - feature_extractor_input_names = self.feature_extractor.model_input_names - return list( - dict.fromkeys( - tokenizer_input_names - + feature_extractor_input_names - + ["feature_attention_mask"] - ) - ) From ee8f7ea4493bf95e7c04951aec701af4000abdb7 Mon Sep 17 00:00:00 2001 From: Thang Nguyen <69278249+Thangnguyenvn98@users.noreply.github.com> Date: Mon, 12 Oct 2026 06:46:07 -0400 Subject: [PATCH 290/624] [CI] Enroll a fourth batch of five steps in automatic sharding (#60858) Signed-off-by: Thang Nguyen Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/distributed.yaml | 3 +++ .buildkite/test_areas/engine.yaml | 1 + .buildkite/test_areas/misc.yaml | 1 + 3 files changed, 5 insertions(+) diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 78a6af2e0ca4..f83946386a38 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -77,6 +77,7 @@ steps: - label: ":nvidia: (L4) Distributed DP Basic" key: distributed-dp-tests-2-gpus + automatic_shard: true timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" device: l4 @@ -300,6 +301,7 @@ steps: - label: ":nvidia: (L4) Distributed DP Extended" key: distributed-dp-tests-4-gpus + automatic_shard: true timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" device: l4 @@ -677,6 +679,7 @@ steps: - label: ":nvidia: (L4) RayExecutorV2" key: rayexecutorv2-4-gpus + automatic_shard: true timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" device: l4 diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 26485e0a172d..21a65bb5ed6b 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -164,6 +164,7 @@ steps: device: h100 num_devices: 2 key: replayssm-e2e + automatic_shard: true timeout_in_minutes: 90 working_dir: "/vllm-workspace/tests" source_file_dependencies: diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 93f59b50e52f..20b48cb98f1f 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -653,6 +653,7 @@ steps: - label: ":nvidia: (H100) Batch Invariance" key: batch-invariance-h100 + automatic_shard: true timeout_in_minutes: 60 device: h100 source_file_dependencies: From 21445806bdf6cfc4fdb9b161643e081bd73cc716 Mon Sep 17 00:00:00 2001 From: Zupeng Wang Date: Mon, 12 Oct 2026 19:12:30 +0800 Subject: [PATCH 291/624] [Rust Frontend] Support logprob_token_ids on OpenAI endpoints (#59549) Signed-off-by: Wang Zupeng Signed-off-by: Bugen Zhao Co-authored-by: OpenAI Codex Co-authored-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- .../routes/openai/chat_completions/convert.rs | 38 +++- .../routes/openai/chat_completions/types.rs | 6 + .../openai/chat_completions/validate.rs | 7 + .../src/routes/openai/completions/convert.rs | 30 ++- .../src/routes/openai/completions/types.rs | 5 + .../src/routes/openai/completions/validate.rs | 16 ++ .../src/routes/openai/utils/logprobs.rs | 24 ++- rust/src/server/src/routes/tests.rs | 171 ++++++++++++++++++ 8 files changed, 292 insertions(+), 5 deletions(-) diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 46aaafbc69e2..30c4fc8ad24f 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -126,6 +126,13 @@ pub(super) fn prepare_chat_request( .prompt_logprobs .or((request.echo && !request.stream).then_some(top_logprobs)); let include_prompt_logprobs = prompt_logprobs.is_some(); + let logprob_token_ids = request.logprob_token_ids.filter(|ids| !ids.is_empty()); + let logprobs = (request.logprobs && logprob_token_ids.is_none()).then_some(top_logprobs); + let output_top_logprobs = if logprob_token_ids.is_some() { + -1 + } else { + top_logprobs + }; let structured_outputs = convert_from_response_format( request.response_format.as_ref(), @@ -157,7 +164,7 @@ pub(super) fn prepare_chat_request( max_tokens: request.max_completion_tokens, min_tokens: request.min_tokens, thinking_token_budget: request.thinking_token_budget, - logprobs: request.logprobs.then_some(top_logprobs), + logprobs, prompt_logprobs, prompt_logprob_token_ids: None, prompt_logprob_start: None, @@ -172,7 +179,7 @@ pub(super) fn prepare_chat_request( allowed_token_ids: request.allowed_token_ids, bad_words: request.bad_words, bad_words_token_ids: None, - logprob_token_ids: None, + logprob_token_ids, structured_outputs, skip_reading_prefix_cache: None, vllm_xargs: merge_kv_transfer_params( @@ -213,7 +220,7 @@ pub(super) fn prepare_chat_request( include_usage, include_continuous_usage, requested_logprobs, - output_top_logprobs: top_logprobs, + output_top_logprobs, include_prompt_logprobs, include_reasoning, echo, @@ -1556,6 +1563,31 @@ mod tests { assert!(prepared.options.is_named_tool_choice); } + #[test] + fn selected_logprobs_override_chat_top_k_unless_empty() { + for ids in [vec![10, 20], vec![]] { + let request = serde_json::from_value(json!({ + "messages": [{"role": "user", "content": "hello"}], + "logprobs": true, "top_logprobs": 1, "logprob_token_ids": ids, + })) + .unwrap(); + let prepared = prepare_chat_request( + request, + &served(&["test-model"]), + ResolvedRequestContext::default(), + ) + .unwrap(); + let params = prepared.chat_request.sampling_params; + let selected = !ids.is_empty(); + assert_eq!(params.logprob_token_ids, selected.then_some(ids)); + assert_eq!(params.logprobs, (!selected).then_some(1)); + assert_eq!( + prepared.options.output_top_logprobs, + if selected { -1 } else { 1 } + ); + } + } + #[test] fn prepare_chat_request_lowers_logprobs_fields() { let request = ChatCompletionRequest { diff --git a/rust/src/server/src/routes/openai/chat_completions/types.rs b/rust/src/server/src/routes/openai/chat_completions/types.rs index 204f3bab5aab..21776dd075e0 100644 --- a/rust/src/server/src/routes/openai/chat_completions/types.rs +++ b/rust/src/server/src/routes/openai/chat_completions/types.rs @@ -171,6 +171,11 @@ pub struct ChatCompletionRequest { /// Number of prompt logprobs to return pub prompt_logprobs: Option, + /// Specific vocab token IDs to return logprobs for at each generated + /// position, in addition to the sampled token. Takes precedence over + /// `top_logprobs`. Requires `logprobs=true`. + pub logprob_token_ids: Option>, + /// Restrict output to these token IDs only pub allowed_token_ids: Option>, @@ -308,6 +313,7 @@ impl Default for ChatCompletionRequest { truncate_prompt_tokens: None, truncation_side: None, prompt_logprobs: None, + logprob_token_ids: None, allowed_token_ids: None, bad_words: None, echo: false, diff --git a/rust/src/server/src/routes/openai/chat_completions/validate.rs b/rust/src/server/src/routes/openai/chat_completions/validate.rs index 579b815ba104..e3701118d94b 100644 --- a/rust/src/server/src/routes/openai/chat_completions/validate.rs +++ b/rust/src/server/src/routes/openai/chat_completions/validate.rs @@ -35,6 +35,13 @@ pub(super) fn validate_request_compat( ); } + if request.logprob_token_ids.as_ref().is_some_and(|ids| !ids.is_empty()) && !request.logprobs { + bail_invalid_request!( + param = "logprob_token_ids", + "when using `logprob_token_ids`, `logprobs` must be set to true." + ); + } + if let Some(prompt_logprobs) = request.prompt_logprobs { if prompt_logprobs < 0 && prompt_logprobs != -1 { bail_invalid_request!( diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index 33fc6bbe0a07..c75cd8da1db3 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -100,6 +100,12 @@ pub(super) fn prepare_completion_request( request.max_tokens }; let echo = completion_echo_text(&request, tokenizer)?; + let logprob_token_ids = request.logprob_token_ids.filter(|ids| !ids.is_empty()); + let logprobs = if logprob_token_ids.is_some() { + None + } else { + logprobs + }; let structured_outputs = convert_from_response_format_value(&request.response_format, &request.structured_outputs)?; @@ -137,7 +143,7 @@ pub(super) fn prepare_completion_request( allowed_token_ids: request.allowed_token_ids, bad_words: None, bad_words_token_ids: None, - logprob_token_ids: None, + logprob_token_ids, structured_outputs, skip_reading_prefix_cache: None, vllm_xargs: merge_kv_transfer_params( @@ -653,6 +659,28 @@ mod tests { ); } + #[test] + fn selected_logprobs_override_completion_top_k_unless_empty() { + for ids in [vec![10, 20], vec![]] { + let request = serde_json::from_value(json!({ + "prompt": "hello", "logprobs": 0, "logprob_token_ids": ids, + })) + .unwrap(); + let prepared = prepare_completion_request( + request, + &served(&["test-model"]), + ResolvedRequestContext::default(), + &test_tokenizer(), + ) + .unwrap(); + let params = prepared.text_request.sampling_params; + let selected = !ids.is_empty(); + assert_eq!(params.logprob_token_ids, selected.then_some(ids)); + assert_eq!(params.logprobs, (!selected).then_some(0)); + assert_eq!(prepared.options.requested_logprobs, Some(0)); + } + } + #[test] fn prepare_completion_request_accepts_logprobs_fields() { let request: CompletionRequest = serde_json::from_value(json!({ diff --git a/rust/src/server/src/routes/openai/completions/types.rs b/rust/src/server/src/routes/openai/completions/types.rs index b10f8a23d701..9f7027511cbf 100644 --- a/rust/src/server/src/routes/openai/completions/types.rs +++ b/rust/src/server/src/routes/openai/completions/types.rs @@ -158,6 +158,11 @@ pub struct CompletionRequest { /// Number of prompt logprobs to return pub prompt_logprobs: Option, + /// Specific vocab token IDs to return logprobs for at each generated + /// position, in addition to the sampled token. Takes precedence over + /// the top-k selected by `logprobs`. Requires `logprobs` to be set. + pub logprob_token_ids: Option>, + // -------- Extra vLLM Parameters -------- /// Whether to add special tokens (e.g. BOS) to the prompt #[serde(default = "default_true")] diff --git a/rust/src/server/src/routes/openai/completions/validate.rs b/rust/src/server/src/routes/openai/completions/validate.rs index e41ad7a85d4e..96090cee935e 100644 --- a/rust/src/server/src/routes/openai/completions/validate.rs +++ b/rust/src/server/src/routes/openai/completions/validate.rs @@ -67,6 +67,22 @@ pub(super) fn validate_request_compat( } } + if request.logprob_token_ids.as_ref().is_some_and(|ids| !ids.is_empty()) { + if request.logprobs.is_none() { + bail_invalid_request!( + param = "logprob_token_ids", + "when using `logprob_token_ids`, `logprobs` must be set." + ); + } + if request.echo && request.max_tokens == Some(0) { + bail_invalid_request!( + param = "logprob_token_ids", + "`logprob_token_ids` is not supported when `echo=True` and \ + `max_tokens=0` because no output tokens are generated." + ); + } + } + if request.use_beam_search { bail_invalid_request!( param = "use_beam_search", diff --git a/rust/src/server/src/routes/openai/utils/logprobs.rs b/rust/src/server/src/routes/openai/utils/logprobs.rs index 1de2cdc10015..3f8b4d9fa749 100644 --- a/rust/src/server/src/routes/openai/utils/logprobs.rs +++ b/rust/src/server/src/routes/openai/utils/logprobs.rs @@ -284,7 +284,9 @@ fn chat_top_logprob_entries( usize::try_from(top_logprobs).unwrap_or(0) }; - position.entries.iter().take(limit) + // The engine row is the sampled token followed by the top-k or requested + // tokens, so the sampled token may repeat. Keep each token's first entry. + position.entries.iter().unique_by(|entry| entry.token_id).take(limit) } fn token_bytes(token: &str) -> Vec { @@ -370,6 +372,26 @@ mod tests { assert_eq!(chat_top_logprobs_len(-1), 3); } + #[test] + fn chat_logprobs_dedupes_repeated_tokens_before_top_logprobs_cut() { + // Sampled `A` is also the top-1 token: the engine row is `[A, A, B, C]`. + let mut logprobs = sample_logprobs(); + let sampled = logprobs.positions[0].entries[0].clone(); + logprobs.positions[0].entries.insert(0, sampled); + + let top_tokens = |top_logprobs| { + let chat_logprobs = + decoded_logprobs_to_openai_chat(&logprobs, top_logprobs, false).unwrap(); + chat_logprobs.content.unwrap()[0] + .top_logprobs + .iter() + .map(|entry| entry.token.clone()) + .collect::>() + }; + assert_eq!(top_tokens(2), ["A", "B"]); + assert_eq!(top_tokens(-1), ["A", "B", "C"]); + } + #[test] fn chat_logprobs_clamps_nan_and_negative_infinity_for_json() { let mut logprobs = sample_logprobs(); diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 9b8378f89934..c6c05e63b1c8 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -3012,6 +3012,177 @@ async fn non_stream_chat_rejects_when_image_count_exceeds_limit_mm_per_prompt() ); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn selected_logprobs_invalid_requests_return_client_errors() { + let (mut app, engine_task) = test_app_with_engine_handle().await; + for chat_route in [true, false] { + let base = if chat_route { + json!({"messages": [{"role": "user", "content": "hello"}], "logprobs": true}) + } else { + json!({"prompt": "hello", "logprobs": 0}) + }; + for invalid in [ + json!({"logprob_token_ids": [-1]}), + json!({"logprob_token_ids": vec![1; 129]}), + if chat_route { + json!({"logprob_token_ids": [1], "logprobs": false}) + } else { + json!({"logprob_token_ids": [1], "logprobs": null}) + }, + json!({"logprob_token_ids": [1], "use_beam_search": true}), + ] { + let mut body = base.clone(); + body.as_object_mut().unwrap().extend(invalid.as_object().unwrap().clone()); + let path = if chat_route { + "/v1/chat/completions" + } else { + "/v1/completions" + }; + let (status, error) = post_json(&mut app, path, body.clone()).await; + assert_eq!(status, StatusCode::BAD_REQUEST, "{body}: {error}"); + } + } + let (status, error) = post_json( + &mut app, + "/v1/completions", + json!({"prompt": "hello", "logprobs": 0, + "logprob_token_ids": [1], "echo": true, "max_tokens": 0}), + ) + .await; + assert_eq!(status, StatusCode::BAD_REQUEST, "{error}"); + engine_task.abort_and_join().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn selected_logprobs_reach_engine_and_both_openai_response_modes() { + for chat_route in [true, false] { + for stream in [true, false] { + let ipc = IpcNamespace::new().unwrap(); + let address = ipc.handshake_endpoint(); + let engine_task = MockEngineTask::new(spawn_mock_engine_task( + address.clone(), + b"selected-logprobs".to_vec(), + |dealer, push| { + boxed_test_future(async move { + let add = recv_engine_message(dealer).await; + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); + let params = request.sampling_params.as_ref().unwrap(); + assert_eq!(params.logprob_token_ids, Some(vec![65, 104, 66])); + assert_eq!(params.logprobs, None); + send_outputs( + push, + RequestBatchOutputs { + outputs: vec![request_output_with_logprobs( + &request.request_id, + vec![104], + Some(EngineCoreFinishReason::Stop), + None, + Some(Logprobs { + positions: vec![PositionLogprobs { + entries: vec![ + TokenLogprob { + token_id: 104, + logprob: -0.5, + rank: 1, + }, + TokenLogprob { + token_id: 65, + logprob: -4.0, + rank: 20, + }, + TokenLogprob { + token_id: 104, + logprob: -0.5, + rank: 1, + }, + TokenLogprob { + token_id: 66, + logprob: -5.0, + rank: 30, + }, + ], + }], + }), + None, + )], + ..Default::default() + } + .into(), + ) + .await; + }) + }, + )); + let client = EngineCoreClient::connect( + EngineCoreClientConfig::new_single(address).with_local_input_output_addresses( + Some(ipc.input_endpoint()), + Some(ipc.output_endpoint()), + ), + ) + .await + .unwrap(); + let chat = + ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new())); + let mut app = build_router(Arc::new(AppState::new(vec!["test-model".into()], chat))); + let mut body = if chat_route { + json!({"messages": [{"role": "user", "content": "hello"}], "logprobs": true, "top_logprobs": 1}) + } else { + json!({"prompt": "hello", "logprobs": 0}) + }; + body["stream"] = json!(stream); + // The sampled token `h` (104) is also requested. + body["logprob_token_ids"] = json!([65, 104, 66]); + let response = app + .call( + Request::builder() + .method("POST") + .uri(if chat_route { + "/v1/chat/completions" + } else { + "/v1/completions" + }) + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + engine_task.await.unwrap(); + let payloads = if stream { + sse_json_payloads(std::str::from_utf8(&bytes).unwrap()) + } else { + vec![serde_json::from_slice(&bytes).unwrap()] + }; + let logprobs = payloads + .iter() + .map(|payload| &payload["choices"][0]["logprobs"]) + .find(|value| value.is_object()) + .unwrap(); + if chat_route { + let actual = logprobs["content"][0]["top_logprobs"] + .as_array() + .unwrap() + .iter() + .map(|entry| json!([entry["token"], entry["logprob"]])) + .collect::>(); + assert_eq!( + actual, + [json!(["h", -0.5]), json!(["A", -4.0]), json!(["B", -5.0])] + ); + } else { + assert_eq!( + logprobs["top_logprobs"][0], + json!({"h": -0.5, "A": -4.0, "B": -5.0}) + ); + } + } + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn non_stream_chat_includes_logprobs_and_prompt_logprobs() { From 3b9e46277407ed53317e78114ad50f1f9bd9d6a4 Mon Sep 17 00:00:00 2001 From: 0xsensei Date: Mon, 12 Oct 2026 16:42:35 +0530 Subject: [PATCH 292/624] [Bugfix][Rust Frontend] Correct DP routing after request removal (#59631) Signed-off-by: 0xsensei Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- .../engine-core-client/src/client/state.rs | 109 ++++++++++++++++-- 1 file changed, 101 insertions(+), 8 deletions(-) diff --git a/rust/src/engine-core-client/src/client/state.rs b/rust/src/engine-core-client/src/client/state.rs index 9f6f8bbf5c2b..3d4844f3add8 100644 --- a/rust/src/engine-core-client/src/client/state.rs +++ b/rust/src/engine-core-client/src/client/state.rs @@ -26,6 +26,7 @@ struct TrackedRequest { sender: OutputSender, engine_id: EngineId, lora: Option, + admission_stats_generation: Option, } /// Frontend-side view of one LoRA request's scheduling phase. @@ -73,6 +74,8 @@ struct EngineRoutingState { /// Added to the snapshot so each admission raises the routing score before /// the next snapshot arrives. Only tracked once a snapshot exists. admitted_since_stats: usize, + /// Identifies admissions made after the latest scheduler snapshot. + stats_generation: u64, } impl EngineRoutingState { @@ -90,10 +93,21 @@ impl EngineRoutingState { } /// Record one request admitted to this engine. - fn record_admission(&mut self) { + fn record_admission(&mut self) -> Option { self.inflight += 1; if self.last_scheduler_stats.is_some() { self.admitted_since_stats += 1; + return Some(self.stats_generation); + } + None + } + + /// Record one request removed from this engine. Its admission is dropped + /// from `admitted_since_stats` only if no newer snapshot has reset that count. + fn record_removal(&mut self, admission_stats_generation: Option) { + self.inflight -= 1; + if admission_stats_generation == Some(self.stats_generation) { + self.admitted_since_stats -= 1; } } @@ -102,6 +116,10 @@ impl EngineRoutingState { fn apply_scheduler_counts(&mut self, next: EngineLoadSnapshot) { self.last_scheduler_stats = Some(next); self.admitted_since_stats = 0; + self.stats_generation = self + .stats_generation + .checked_add(1) + .expect("scheduler stats generation overflow"); } } @@ -161,21 +179,22 @@ impl RequestRegistry { if lora.is_some() { self.active_lora_requests += 1; } + let state = self + .routing_per_engine + .get_mut(&engine_id) + .expect("request registry must track all known engines"); + let admission_stats_generation = state.record_admission(); + self.requests.insert( request_id, TrackedRequest { sender: tx, engine_id: engine_id.clone(), lora, + admission_stats_generation, }, ); - let state = self - .routing_per_engine - .get_mut(&engine_id) - .expect("request registry must track all known engines"); - state.record_admission(); - Ok((engine_id, rx)) } @@ -368,7 +387,7 @@ impl RequestRegistry { self.routing_per_engine .get_mut(&tracked.engine_id) .expect("request registry must track all known engines") - .inflight -= 1; + .record_removal(tracked.admission_stats_generation); Some((tracked.sender, tracked.engine_id)) } @@ -1079,6 +1098,80 @@ mod tests { expect_test::expect!["[2, 5, 2]"].assert_eq(&format!("{chosen:?}")); } + #[test] + fn registry_finished_requests_do_not_route_live_work_to_a_busier_engine() { + let mut registry = RequestRegistry::new( + &[0, 1].map(|rank| connected_engine(EngineId::from_engine_index(rank))), + ); + registry.apply_scheduler_counts( + 0, + EngineLoadSnapshot { + waiting: 0, + running: 0, + }, + ); + registry.apply_scheduler_counts( + 1, + EngineLoadSnapshot { + waiting: 0, + running: 1, + }, + ); + + let mut finished_routes = Vec::new(); + for request_id in ["finished-1", "finished-2"].map(String::from) { + let (engine_id, _) = + registry.register(request_id.clone(), None, None, NonZeroU32::MIN).unwrap(); + + finished_routes.push(engine_id.engine_index()); + drop(registry.finish_many([&request_id])); + } + + let (live_engine, _) = + registry.register("live-request".into(), None, None, NonZeroU32::MIN).unwrap(); + expect_test::expect!["([Some(0), Some(0)], Some(0))"].assert_eq(&format!( + "{:?}", + (finished_routes, live_engine.engine_index()) + )); + } + + #[test] + fn registry_removal_keeps_admissions_counted_after_a_newer_snapshot() { + let engine = EngineId::from_engine_index(0); + let mut registry = RequestRegistry::new(&[connected_engine(engine.clone())]); + registry.apply_scheduler_counts( + 0, + EngineLoadSnapshot { + waiting: 0, + running: 0, + }, + ); + + // `old` is admitted before the next snapshot, which already counts it. + registry.register("old".into(), None, None, NonZeroU32::MIN).unwrap(); + registry.apply_scheduler_counts( + 0, + EngineLoadSnapshot { + waiting: 0, + running: 1, + }, + ); + // `new` is admitted after that snapshot, so only it is counted on top. + registry.register("new".into(), None, None, NonZeroU32::MIN).unwrap(); + + // Removing `old` must not take back the admission counted for `new`. + drop(registry.remove("old")); + let state = ®istry.routing_per_engine[&engine]; + expect_test::expect!["(1, 1, 2)"].assert_eq(&format!( + "{:?}", + ( + state.inflight, + state.admitted_since_stats, + state.routing_score() + ) + )); + } + #[test] fn register_with_data_parallel_rank_routes_to_specified_engine() { let engine_0 = EngineId::from_engine_index(0); From f5cda48c9023e8ea3b5d29b57815c8455473e97b Mon Sep 17 00:00:00 2001 From: Reid <61492567+reidliu41@users.noreply.github.com> Date: Mon, 12 Oct 2026 19:12:39 +0800 Subject: [PATCH 293/624] [Bugfix][Rust Frontend] Export corrupted request metric (#59937) Signed-off-by: reidliu41 Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao --- rust/src/llm/src/request_metrics.rs | 8 ++++++++ rust/src/llm/tests/generate.rs | 4 ++++ rust/src/metrics/src/request.rs | 9 +++++++++ 3 files changed, 21 insertions(+) diff --git a/rust/src/llm/src/request_metrics.rs b/rust/src/llm/src/request_metrics.rs index 61cea120940d..71decdd281ff 100644 --- a/rust/src/llm/src/request_metrics.rs +++ b/rust/src/llm/src/request_metrics.rs @@ -47,6 +47,7 @@ pub(crate) struct RequestMetricsTracker { num_generation_tokens: u32, num_preemptions: u64, latest_num_cached_tokens: u32, + is_corrupted: bool, } /// Cached request metric handles for one model and engine index. @@ -61,6 +62,7 @@ struct RequestMetricHandles { prompt_tokens_external_kv_transfer: U64Counter, prompt_tokens_cached: U64Counter, generation_tokens: U64Counter, + corrupted_requests: U64Counter, // Request lifecycle counters and histograms. request_success: Family, @@ -123,6 +125,7 @@ impl RequestMetricsTracker { num_generation_tokens: 0, num_preemptions: 0, latest_num_cached_tokens: 0, + is_corrupted: false, } } @@ -141,6 +144,7 @@ impl RequestMetricsTracker { } self.num_generation_tokens += output.new_token_ids.len() as u32; self.handles.generation_tokens.inc_by(output.new_token_ids.len() as u64); + self.is_corrupted |= output.num_nans_in_logits > 0; if let Some(events) = &output.events { self.observe_events(events); @@ -187,6 +191,9 @@ impl RequestMetricsTracker { }; self.record_request_success(finish_reason); + if self.is_corrupted { + self.handles.corrupted_requests.inc(); + } self.handles.request_prompt_tokens.observe(self.prompt_len as f64); self.handles @@ -297,6 +304,7 @@ fn resolve_request_metric_handles( ), prompt_tokens_cached: metrics.prompt_tokens_cached.get_or_create_owned(&labels), generation_tokens: metrics.generation_tokens.get_or_create_owned(&labels), + corrupted_requests: metrics.corrupted_requests.get_or_create_owned(&labels), request_success: metrics.request_success.clone(), request_prompt_tokens: metrics.request_prompt_tokens.get_or_create_owned(&labels), request_generation_tokens: metrics.request_generation_tokens.get_or_create_owned(&labels), diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 8a150d07a44b..0a1eacd73ff6 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -898,6 +898,7 @@ async fn generate_records_request_metrics_in_prometheus_output() { engine_index: 4, timestamp: 10.0, outputs: vec![EngineCoreOutput { + num_nans_in_logits: 2, prefill_stats: Some(Box::new(PrefillStats { num_prompt_tokens: 2, num_computed_tokens: 2, @@ -982,6 +983,9 @@ async fn generate_records_request_metrics_in_prometheus_output() { assert!(rendered.contains(&format!( "vllm:generation_tokens_total{{model_name=\"{model_name}\",engine=\"4\"}} 3" ))); + assert!(rendered.contains(&format!( + "vllm:corrupted_requests_total{{model_name=\"{model_name}\",engine=\"4\"}} 1" + ))); assert!(rendered.contains(&format!( "vllm:num_preemptions_total{{model_name=\"{model_name}\",engine=\"4\"}} 1" ))); diff --git a/rust/src/metrics/src/request.rs b/rust/src/metrics/src/request.rs index b6e7683e768b..46260cf58397 100644 --- a/rust/src/metrics/src/request.rs +++ b/rust/src/metrics/src/request.rs @@ -106,6 +106,7 @@ pub struct RequestMetrics { pub prompt_tokens_by_source: PromptTokenSourceCounterFamily, pub prompt_tokens_cached: Family, pub generation_tokens: Family, + pub corrupted_requests: Family, // Request lifecycle counters and histograms. pub request_success: FinishedReasonCounterFamily, @@ -173,6 +174,13 @@ impl RequestMetrics { generation_tokens.clone(), ); + let corrupted_requests = Family::default(); + registry.register( + "vllm:corrupted_requests", + "Corrupted requests, in terms of total number of requests with NaNs in logits.", + corrupted_requests.clone(), + ); + // Request lifecycle counters and histograms. let request_success = Family::default(); registry.register( @@ -309,6 +317,7 @@ impl RequestMetrics { prompt_tokens_by_source, prompt_tokens_cached, generation_tokens, + corrupted_requests, request_success, request_prompt_tokens, request_generation_tokens, From 32b05addc1301be904ba1b6c9a5e3704581b1fc3 Mon Sep 17 00:00:00 2001 From: Reid <61492567+reidliu41@users.noreply.github.com> Date: Mon, 12 Oct 2026 19:12:43 +0800 Subject: [PATCH 294/624] [Bugfix][Rust Frontend] Restore DP wave delivery after XPUB reconnects (#60249) Signed-off-by: reidliu41 Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- .../src/coordinator/bootstrap.rs | 24 +++++--------- .../src/coordinator/inproc.rs | 33 +++++++++++++++++-- .../engine-core-client/src/tests/client.rs | 33 ++++++++++++++----- 3 files changed, 65 insertions(+), 25 deletions(-) diff --git a/rust/src/engine-core-client/src/coordinator/bootstrap.rs b/rust/src/engine-core-client/src/coordinator/bootstrap.rs index 68229085d239..17c346c933e3 100644 --- a/rust/src/engine-core-client/src/coordinator/bootstrap.rs +++ b/rust/src/engine-core-client/src/coordinator/bootstrap.rs @@ -63,22 +63,9 @@ async fn wait_for_engine_subscriptions( timeout: ready_timeout, } })??; - if message.len() != 1 { + if !is_engine_subscription(&message) { bail_unexpected_handshake_message!( - "expected 1 frame for coordinator subscription, got {}", - message.len() - ); - } - - let frame = message - .into_vec() - .into_iter() - .next() - .expect("single-frame coordinator subscription message"); - if frame.as_ref() != [0x01] { - bail_unexpected_handshake_message!( - "expected coordinator subscription frame [0x01], got {:?}", - frame.as_ref() + "expected coordinator subscription [0x01], got {message:?}" ); } received += 1; @@ -87,6 +74,13 @@ async fn wait_for_engine_subscriptions( Ok(()) } +/// Whether `message` is the single-frame subscribe-all notification (`[0x01]`) +/// that an engine's XSUB sends whenever it connects or reconnects to the +/// coordinator broadcast socket. +pub(super) fn is_engine_subscription(message: &ZmqMessage) -> bool { + message.len() == 1 && message.get(0).is_some_and(|frame| frame.as_ref() == [0x01]) +} + /// Send the coordinator READY marker to all subscribed engines. async fn send_ready_to_engines(input_socket: &mut XPubSocket) -> Result<()> { input_socket.send(ZmqMessage::from(Bytes::from_static(b"READY"))).await?; diff --git a/rust/src/engine-core-client/src/coordinator/inproc.rs b/rust/src/engine-core-client/src/coordinator/inproc.rs index bf4c23cc137a..cbd69095f85c 100644 --- a/rust/src/engine-core-client/src/coordinator/inproc.rs +++ b/rust/src/engine-core-client/src/coordinator/inproc.rs @@ -7,10 +7,11 @@ use serde_tuple::{Deserialize_tuple, Serialize_tuple}; use thiserror_ext::AsReport; use tokio::sync::mpsc; use tracing::{debug, warn}; -use zeromq::prelude::SocketSend; -use zeromq::{XPubSocket, ZmqMessage}; +use zeromq::prelude::{SocketRecv, SocketSend}; +use zeromq::{XPubSocket, ZmqError, ZmqMessage}; use crate::client::imp::ClientInner; +use crate::coordinator::bootstrap::is_engine_subscription; use crate::coordinator::handle::{CoordinatorCommand, CoordinatorState}; use crate::error::{Error, Result, bail_unexpected_coordinator_output}; use crate::protocol::encode_msgpack; @@ -187,6 +188,34 @@ impl InProcCoordinatorRunner { let result: Result<()> = async { loop { tokio::select! { + // Keep receiving engine subscriptions: zmq.rs XPUB registers a + // subscription only when `recv()` yields it (unlike libzmq), so an + // engine whose XSUB reconnects gets no broadcasts until we read its + // subscribe-all. If a wave is running, replay it so that engine does + // not miss the `START_DP_WAVE` sent while it was disconnected. + subscription = self.coordinator_input.recv() => { + match subscription { + Ok(message) => { + if is_engine_subscription(&message) { + let state = *self.state.lock(); + if state.engines_running { + debug!( + current_wave = state.current_wave, + "replaying active DP wave after engine subscription" + ); + self.broadcast_start_wave(state.current_wave, None).await?; + } + } + } + Err(error @ (ZmqError::Codec(_) | ZmqError::Network(_))) => { + warn!( + error = %error, + "coordinator subscription peer disconnected" + ); + } + Err(error) => return Err(error.into()), + } + } // Received frontend-originated command from the handle. command = self.command_rx.recv() => { let Some(command) = command else { diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 43eb88db79ed..7b7a880bde6d 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -257,7 +257,7 @@ async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec { dealer.recv().await.unwrap().into_vec() } -async fn recv_start_dp_wave(sub: &mut SubSocket) -> (u32, u32) { +async fn recv_start_dp_wave(sub: &mut SubSocket) -> (u32, Option) { let frames = sub.recv().await.unwrap().into_vec(); assert_eq!(frames.len(), 2); assert_eq!( @@ -607,12 +607,27 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { let handshake_address = handshake_address.clone(); async move { let mut engine = setup_mock_engine_sockets(handshake_address, &[0x00, 0x00]).await; + let coordinator_address = engine + .init + .addresses + .coordinator_input + .clone() + .expect("coordinator input address should be present"); let mut coordinator = engine.coordinator.take().expect("coordinator sockets should be present"); let data_socket = engine.data_sockets.first_mut().expect("data socket"); let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; - assert_eq!((wave, exclude_engine), (0, 0)); + assert_eq!((wave, exclude_engine), (0, Some(0))); + + let mut late_subscriber = SubSocket::new(); + late_subscriber.connect(&coordinator_address).await.unwrap(); + late_subscriber.subscribe("").await.unwrap(); + + let replay = recv_start_dp_wave(&mut late_subscriber).await; + assert_eq!(replay, (0, None)); + let replay = recv_start_dp_wave(&mut coordinator.input_sub).await; + assert_eq!(replay, (0, None)); let add = recv_engine_message(&mut data_socket.dealer).await; assert_eq!(add[0].as_ref(), &[0x00]); @@ -656,7 +671,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { .await; let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; - assert_eq!((wave, exclude_engine), (1, 0)); + assert_eq!((wave, exclude_engine), (1, Some(0))); let add = recv_engine_message(&mut data_socket.dealer).await; assert_eq!(add[0].as_ref(), &[0x00]); @@ -693,7 +708,9 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { let data_socket = engine.data_sockets.first_mut().expect("data socket"); let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; - assert_eq!((wave, exclude_engine), (0, 0)); + assert_eq!((wave, exclude_engine), (0, Some(0))); + let replay = recv_start_dp_wave(&mut coordinator.input_sub).await; + assert_eq!(replay, (0, None)); let add = recv_engine_message(&mut data_socket.dealer).await; assert_eq!(add[0].as_ref(), &[0x00]); @@ -727,7 +744,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { .await; let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; - assert_eq!((wave, exclude_engine), (1, 0)); + assert_eq!((wave, exclude_engine), (1, Some(0))); assert!( timeout( @@ -812,7 +829,7 @@ async fn coordinator_rebroadcasts_engine_start_wave_control() { engine.coordinator.take().expect("coordinator sockets should be present"); let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; - assert_eq!((wave, exclude_engine), (4, 1)); + assert_eq!((wave, exclude_engine), (4, Some(1))); let _ = shutdown0_rx.await; } @@ -838,7 +855,7 @@ async fn coordinator_rebroadcasts_engine_start_wave_control() { .await; let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; - assert_eq!((wave, exclude_engine), (4, 1)); + assert_eq!((wave, exclude_engine), (4, Some(1))); let _ = shutdown1_rx.await; } @@ -880,7 +897,7 @@ async fn coordinator_accepts_stats_only_outputs() { let data_socket = engine.data_sockets.first_mut().expect("data socket"); let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; - assert_eq!((wave, exclude_engine), (0, 0)); + assert_eq!((wave, exclude_engine), (0, Some(0))); send_outputs( &mut coordinator.output_push, From 8d7b93fec2eb72ce01f2217472ac1f51b9bce600 Mon Sep 17 00:00:00 2001 From: Shimi Bandiel Date: Mon, 12 Oct 2026 04:37:53 -0700 Subject: [PATCH 295/624] [Frontend] Report stop_reason on /inference/v1/generate and pass it through derender (#60251) Signed-off-by: Hyeonki Hong Signed-off-by: Shimi Bandiel Signed-off-by: Bugen Zhao Co-authored-by: Hyeonki Hong Co-authored-by: Mimir Co-authored-by: Claude Co-authored-by: Bugen Zhao --- .../server/src/routes/inference/generate.rs | 72 ++++++++++++++++++- .../src/routes/inference/generate/types.rs | 4 +- .../scale_out/derender/test_derender.py | 32 +++++++++ .../derender/test_derender_stream.py | 15 ++++ .../test_generate_stream.py | 67 +++++++++++++++++ .../scale_out/token_in_token_out/protocol.py | 6 ++ .../scale_out/token_in_token_out/serving.py | 2 + vllm/renderers/online_derenderer.py | 5 ++ 8 files changed, 201 insertions(+), 2 deletions(-) diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index a30cd0c85ab4..930c698c30e4 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -226,6 +226,8 @@ async fn generate_chunk_stream( None }; + let stop_reason = + finish_reason.as_ref().and_then(|reason| reason.as_stop_reason().cloned()); let prompt_token_ids = prompt_token_ids.take(); y.yield_ok(GenerateStreamResponse { request_id: request_id.clone(), @@ -233,6 +235,7 @@ async fn generate_chunk_stream( index: 0, logprobs, finish_reason: finish_reason.map(|reason| reason.as_str().to_string()), + stop_reason, token_ids, sampling_mask: output.sampling_mask.map(|mask| mask.rows), weight_versions: output.weight_versions, @@ -319,6 +322,7 @@ fn collect_generate( None }; let finish_reason = collected.finish_reason.as_str().to_string(); + let stop_reason = collected.finish_reason.as_stop_reason().cloned(); if enable_log_requests { info!( @@ -335,6 +339,7 @@ fn collect_generate( index: 0, logprobs, finish_reason: Some(finish_reason), + stop_reason, token_ids: collected.token_ids, sampling_mask: collected.sampling_mask.map(|mask| mask.rows), weight_versions: collected.weight_versions, @@ -537,7 +542,7 @@ mod tests { use futures::{TryStreamExt as _, stream}; use vllm_engine_core_client::protocol::multimodal::{MmModality, PlaceholderRange}; - use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics; + use vllm_engine_core_client::protocol::output::{RequestSpecDecodeMetrics, StopReason}; use vllm_engine_core_client::protocol::sampling_mask::SamplingMask; use vllm_llm::GeneratePromptInfo; @@ -1125,4 +1130,69 @@ mod tests { ) .expect_err("multi-token prompt without payload is an engine failure"); } + + #[tokio::test] + async fn generate_chunk_stream_reports_stop_reason() { + let chunks = collect_chunks( + vec![ + stream_output(Some(&[11, 22]), vec![33], None), + stream_output( + Some(&[11, 22]), + vec![44], + Some(FinishReason::Stop(Some(StopReason::Text( + "wor".to_string(), + )))), + ), + ], + false, + None, + ) + .await; + + assert_eq!(chunks.len(), 2); + assert_eq!(chunks[0].choices[0].stop_reason, None); + assert_eq!(chunks[1].choices[0].finish_reason.as_deref(), Some("stop")); + assert_eq!( + chunks[1].choices[0].stop_reason, + Some(StopReason::Text("wor".to_string())) + ); + } + + #[test] + fn collect_generate_reports_stop_token_id() { + let output = CollectedGenerateOutput { + request_id: "raw-stop".to_string(), + prompt_logprobs: None, + prompt_token_id_logprobs: None, + token_ids: vec![30, 7], + logprobs: None, + finish_reason: FinishReason::Stop(Some(StopReason::TokenId(7))), + usage: TokenUsage { + prompt_token_count: 2, + output_token_count: 2, + cached_token_count: 0, + }, + kv_transfer_params: None, + ec_transfer_params: None, + prompt_token_ids: vec![10, 20], + sampling_mask: None, + spec_decode_metrics: None, + weight_versions: None, + }; + + let response = collect_generate( + output, + "raw-stop".to_string(), + ApiServerOptions::default(), + ResponseOptions::default(), + None, + ) + .expect("response"); + + assert_eq!(response.choices[0].finish_reason.as_deref(), Some("stop")); + assert_eq!( + response.choices[0].stop_reason, + Some(StopReason::TokenId(7)) + ); + } } diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs index 6352b187250a..d2737b5bce69 100644 --- a/rust/src/server/src/routes/inference/generate/types.rs +++ b/rust/src/server/src/routes/inference/generate/types.rs @@ -7,7 +7,7 @@ use llm_multimodal::MediaContentPart; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use validator::Validate; -use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics; +use vllm_engine_core_client::protocol::output::{RequestSpecDecodeMetrics, StopReason}; use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; use vllm_llm::WeightVersionSpan; use vllm_text::SamplingParams; @@ -72,6 +72,7 @@ pub(super) struct GenerateResponseChoice { pub index: u32, pub logprobs: Option, pub finish_reason: Option, + pub stop_reason: Option, pub token_ids: Vec, pub sampling_mask: Option>>, pub weight_versions: Option>, @@ -84,6 +85,7 @@ pub(super) struct GenerateResponseStreamChoice { pub index: u32, pub logprobs: Option, pub finish_reason: Option, + pub stop_reason: Option, pub token_ids: Vec, pub sampling_mask: Option>>, /// Set on the final chunk only; covers the whole output. diff --git a/tests/entrypoints/scale_out/derender/test_derender.py b/tests/entrypoints/scale_out/derender/test_derender.py index d29aa3f53f73..30d64ed39bfd 100644 --- a/tests/entrypoints/scale_out/derender/test_derender.py +++ b/tests/entrypoints/scale_out/derender/test_derender.py @@ -112,6 +112,38 @@ async def test_derender_chat_roundtrip(client): assert data["choices"][0]["message"]["role"] == "assistant" +@pytest.mark.asyncio +async def test_derender_chat_passes_stop_reason_through(client): + """The coupled server reports the stop string or token id that ended + generation; derender must carry the generate value through unchanged.""" + gen_req = await _render_chat(client) + generate_response = _make_generate_response(gen_req["token_ids"][:3]) + generate_response["choices"][0]["stop_reason"] = "\n\n" + + response = await client.post( + "/v1/chat/completions/derender", + json={"model": MODEL_NAME, "generate_response": generate_response}, + ) + assert response.status_code == 200 + assert response.json()["choices"][0]["stop_reason"] == "\n\n" + + +@pytest.mark.asyncio +async def test_derender_completion_passes_stop_reason_through(client): + gen_req = await _render_completion(client, "Hello world") + generate_response = _make_completion_generate_response( + gen_req["token_ids"][:3], "gen-stop" + ) + generate_response["choices"][0]["stop_reason"] = 151645 + + response = await client.post( + "/v1/completions/derender", + json={"model": MODEL_NAME, "generate_responses": [generate_response]}, + ) + assert response.status_code == 200 + assert response.json()["choices"][0]["stop_reason"] == 151645 + + @pytest.mark.asyncio async def test_derender_chat_usage(client): """Supplied prompt_tokens flows through into usage correctly.""" diff --git a/tests/entrypoints/scale_out/derender/test_derender_stream.py b/tests/entrypoints/scale_out/derender/test_derender_stream.py index d97a2f992ac2..55b0789843ee 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_stream.py +++ b/tests/entrypoints/scale_out/derender/test_derender_stream.py @@ -195,6 +195,7 @@ def _make_stream_chunk( metrics: PerRequestMetrics | None = None, logprobs: dict | None = None, prompt_token_ids: list[int] | None = None, + stop_reason: int | str | None = None, ) -> GenerateStreamResponse: """Build a GenerateStreamResponse SSE chunk.""" return GenerateTokensStreamResponse( @@ -204,6 +205,7 @@ def _make_stream_chunk( index=index, token_ids=token_ids, finish_reason=finish_reason, + stop_reason=stop_reason, logprobs=logprobs, ) ], @@ -886,6 +888,19 @@ async def test_finish_reason_forwarded(self, derenderer, tokenizer): ) assert chunk.choices[0].finish_reason == "length" + @pytest.mark.asyncio + async def test_stop_reason_forwarded(self, derenderer, tokenizer): + """stop_reason from the final generate chunk reaches the derendered + choice, matching what the coupled server reports.""" + token_ids = tokenizer.encode("done")[:2] + chunk, _ = await derenderer.derender_completion_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk( + token_ids, finish_reason="stop", stop_reason="\n\n" + ), + ) + assert chunk.choices[0].stop_reason == "\n\n" + class TestStreamLogprobs: """Streaming derender must carry per-chunk logprobs with token ids diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py index 7f0cd1286041..c30732f3150f 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py @@ -139,6 +139,7 @@ def _make_request_output( request_id: str, token_ids: list[int], finish_reason: str | None = None, + stop_reason: int | str | None = None, finished: bool = False, prompt_token_ids: list[int] | None = None, logprobs: list[dict[int, Any] | None] | None = None, @@ -160,6 +161,7 @@ def _make_request_output( cumulative_logprob=None, logprobs=logprobs, finish_reason=finish_reason, + stop_reason=stop_reason, spec_decode_metrics=spec_decode_metrics, ) ], @@ -378,6 +380,71 @@ async def mock_generate(*args, **kwargs): assert data_chunks[2]["choices"][0]["finish_reason"] == "stop" +@pytest.mark.asyncio +async def test_non_stream_reports_engine_stop_reason(): + """The matched stop token id from the engine lands on the choice.""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", + token_ids=[10, 42], + finish_reason="stop", + stop_reason=42, + finished=True, + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=10), + model=MODEL_NAME, + stream=False, + ) + + response = await serving.serve_tokens(request) + choice = response.model_dump()["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["stop_reason"] == 42 + + +@pytest.mark.asyncio +async def test_stream_reports_engine_stop_reason_on_final_chunk(): + """A matched stop string is reported with finish_reason and null before.""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10]) + yield _make_request_output( + "req-1", + token_ids=[20], + finish_reason="stop", + stop_reason="done", + finished=True, + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=10), + model=MODEL_NAME, + stream=True, + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + data_chunks = [c for c in parsed if c != "[DONE]"] + assert len(data_chunks) == 2 + + assert data_chunks[0]["choices"][0]["stop_reason"] is None + assert data_chunks[1]["choices"][0]["finish_reason"] == "stop" + assert data_chunks[1]["choices"][0]["stop_reason"] == "done" + + @pytest.mark.asyncio @pytest.mark.parametrize("stream", [False, True]) async def test_serve_tokens_returns_spec_decode_metrics(stream: bool): diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index bccc8d8eff0b..c8d82eeb5a23 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -400,6 +400,9 @@ class GenerateChoiceBase(BaseModel): index: int # per OpenAI spec this is the default finish_reason: str | None = "stop" + # Stop string or token id that ended generation (repetition detection + # sets it too); None on EOS, length and abort. + stop_reason: int | str | None = None token_ids: list[int] | None = None # Per-token expert routing decisions, base64-encoded `.npy` bytes # (numpy serialization). Shape after decode: @@ -443,6 +446,9 @@ class GenerateStreamChoiceBase(BaseModel): index: int finish_reason: str | None = None + # Stop string or token id that ended generation (repetition detection + # sets it too); None on EOS, length and abort. + stop_reason: int | str | None = None token_ids: list[int] | None = None routed_experts: str | None = None sampling_mask: list[list[int]] | None = None diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index 085ae97e16a7..f09ee6fad5ab 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -450,6 +450,7 @@ async def serve_tokens_full_generator( index=output.index, logprobs=logprobs, finish_reason=output.finish_reason if output.finish_reason else "stop", + stop_reason=output.stop_reason, token_ids=as_list(output.token_ids), routed_experts=routed_experts_b64, sampling_mask=sampling_mask, @@ -623,6 +624,7 @@ async def serve_tokens_stream_generator( index=i, logprobs=logprobs, finish_reason=finish_reason, + stop_reason=output.stop_reason, token_ids=as_list(delta_token_ids), routed_experts=routed_experts_b64, sampling_mask=sampling_mask, diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index 3fe50f16c391..345e5443ea72 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -206,6 +206,7 @@ def _derender_chat( message=message, logprobs=resolved_logprobs, finish_reason=choice.finish_reason, + stop_reason=choice.stop_reason, token_ids=choice.token_ids if return_token_ids else None, ) ) @@ -484,6 +485,7 @@ async def derender_chat_stream( delta=delta, logprobs=resolved_logprobs, finish_reason=choice.finish_reason, + stop_reason=choice.stop_reason, token_ids=choice.token_ids if return_token_ids else None, ) ) @@ -672,6 +674,7 @@ def _replay(token_ids: list[int], chunk_lens: list[int]) -> None: index=choice.index, delta=delta_message, finish_reason=finish_reason, + stop_reason=choice.stop_reason, token_ids=choice.token_ids if return_token_ids else None, ) stream_choices.append( @@ -777,6 +780,7 @@ def _derender_completion( index=index, text=decoded_text, finish_reason=choice.finish_reason, + stop_reason=choice.stop_reason, logprobs=completion_logprobs, ) ) @@ -878,6 +882,7 @@ async def derender_completion_stream( text=new_text, logprobs=completion_logprobs, finish_reason=choice.finish_reason, + stop_reason=choice.stop_reason, ) ) From aa1d44df16fe554405c6064cad2a904a16b63581 Mon Sep 17 00:00:00 2001 From: Liana Koleva <43767763+lianakoleva@users.noreply.github.com> Date: Mon, 12 Oct 2026 05:20:53 -0700 Subject: [PATCH 296/624] [Bugfix][Rust Frontend] Alias `qwen3_xml` to `qwen3_coder` and add `qwen3_json` for the Qwen JSON format (#60136) Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao --- rust/src/chat/src/lib.rs | 2 +- rust/src/chat/src/output/default/mod.rs | 2 +- rust/src/chat/src/parser/tool/mod.rs | 14 ++-- rust/src/chat/src/parser/tool/tests.rs | 31 +++++++- rust/src/chat/tests/roundtrip.rs | 2 +- rust/src/parser/Cargo.toml | 2 +- .../benches/{qwen3_xml.rs => qwen3_json.rs} | 14 ++-- rust/src/parser/src/tool/json/mod.rs | 2 +- rust/src/parser/src/tool/json/qwen.rs | 77 ++++++++++--------- rust/src/parser/src/tool/mod.rs | 2 +- rust/src/parser/src/unified/combined.rs | 16 ++-- vllm/tool_parsers/__init__.py | 7 ++ 12 files changed, 103 insertions(+), 68 deletions(-) rename rust/src/parser/benches/{qwen3_xml.rs => qwen3_json.rs} (89%) diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index f437745e3967..0b7e89914ce0 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -431,7 +431,7 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, granite4, hermes, hf, hy_v3, hy_v4, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, mimo, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string()); + expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, granite4, hermes, hf, hy_v3, hy_v4, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, mimo, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_json, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string()); } #[test] diff --git a/rust/src/chat/src/output/default/mod.rs b/rust/src/chat/src/output/default/mod.rs index 981173dc7450..e95ebd8b0dcf 100644 --- a/rust/src/chat/src/output/default/mod.rs +++ b/rust/src/chat/src/output/default/mod.rs @@ -498,7 +498,7 @@ mod tests { Err(error) => error, }; - expect_test::expect!["unified parsing requires the tool and reasoning selections to resolve to the same parser; resolved tool=qwen3_xml, reasoning=gemma4"] + expect_test::expect!["unified parsing requires the tool and reasoning selections to resolve to the same parser; resolved tool=qwen3_json, reasoning=gemma4"] .assert_eq(&format!("{error}")); } diff --git a/rust/src/chat/src/parser/tool/mod.rs b/rust/src/chat/src/parser/tool/mod.rs index 9f3ab4f516a3..9360b62e2f1e 100644 --- a/rust/src/chat/src/parser/tool/mod.rs +++ b/rust/src/chat/src/parser/tool/mod.rs @@ -10,7 +10,7 @@ pub use vllm_parser::tool::{ DeepSeekV41ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, Granite4ToolParser, HermesToolParser, Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, MiMoToolParser, MinimaxM2ToolParser, MinimaxM3ToolParser, MistralToolParser, Phi4MiniJsonToolParser, - Qwen3CoderToolParser, Qwen3XmlToolParser, SeedOssToolParser, ToolParser, ToolParserError, + Qwen3CoderToolParser, Qwen3JsonToolParser, SeedOssToolParser, ToolParser, ToolParserError, }; use crate::parser::ParserFactory; @@ -39,6 +39,9 @@ pub mod names { pub const MISTRAL: &str = "mistral"; pub const PHI4_MINI_JSON: &str = "phi4_mini_json"; pub const QWEN3_CODER: &str = "qwen3_coder"; + // The ``-wrapped JSON format of Qwen2.5, QwQ and Qwen3. + pub const QWEN3_JSON: &str = "qwen3_json"; + // An alias of `qwen3_coder` for the same XML format. pub const QWEN3_XML: &str = "qwen3_xml"; pub const SEED_OSS: &str = "seed_oss"; } @@ -82,8 +85,9 @@ impl ToolParserFactory { .register_parser::(names::MINIMAX_M3) .register_parser::(names::MISTRAL) .register_parser::(names::PHI4_MINI_JSON) - .register_parser::(names::QWEN3_XML) .register_parser::(names::QWEN3_CODER) + .register_parser::(names::QWEN3_XML) + .register_parser::(names::QWEN3_JSON) .register_parser::(names::SEED_OSS); factory @@ -92,9 +96,9 @@ impl ToolParserFactory { .register_pattern("mimo-v2", names::MIMO) .register_pattern("qwen3-coder", names::QWEN3_CODER) .register_pattern("qwen3.5", names::QWEN3_CODER) - .register_pattern("qwq", names::HERMES) - .register_pattern("qwen2.5", names::HERMES) - .register_pattern("qwen3", names::QWEN3_XML) + .register_pattern("qwq", names::QWEN3_JSON) + .register_pattern("qwen2.5", names::QWEN3_JSON) + .register_pattern("qwen3", names::QWEN3_JSON) .register_pattern("hermes", names::HERMES) // Narrow to `internlm2` substring so it matches `internlm2-chat-7b` // and `internlm2_5-7b-chat` but NOT `internlm-chat-7b` (InternLM v1, diff --git a/rust/src/chat/src/parser/tool/tests.rs b/rust/src/chat/src/parser/tool/tests.rs index 69022e550d59..d7eb6c2dac52 100644 --- a/rust/src/chat/src/parser/tool/tests.rs +++ b/rust/src/chat/src/parser/tool/tests.rs @@ -86,10 +86,10 @@ fn factory_distinguishes_qwen_tool_formats() { for (model, expected) in [ ("Qwen/Qwen3.5-0.8B", Some(names::QWEN3_CODER)), - ("Qwen/Qwen3-0.6B", Some(names::QWEN3_XML)), + ("Qwen/Qwen3-0.6B", Some(names::QWEN3_JSON)), ("Qwen/Qwen3-Coder-30B", Some(names::QWEN3_CODER)), - ("Qwen/QwQ-32B", Some(names::HERMES)), - ("Qwen/Qwen2.5-0.5B-Instruct", Some(names::HERMES)), + ("Qwen/QwQ-32B", Some(names::QWEN3_JSON)), + ("Qwen/Qwen2.5-0.5B-Instruct", Some(names::QWEN3_JSON)), ("Qwen/Qwen2-1.5B-Instruct", None), ] { assert_eq!(factory.resolve_name_for_model(model), expected, "{model}"); @@ -97,7 +97,7 @@ fn factory_distinguishes_qwen_tool_formats() { } #[test] -fn factory_parses_qwen2_5_coder_template_with_hermes() { +fn factory_parses_qwen2_5_coder_template_with_qwen3_json() { let factory = ToolParserFactory::new(); let tool_call = r#" {"name":"get_weather","arguments":{"location":"Tokyo"}} @@ -115,6 +115,29 @@ fn factory_parses_qwen2_5_coder_template_with_hermes() { assert_eq!(output.calls()[0].arguments, r#"{"location":"Tokyo"}"#); } +#[test] +fn factory_parses_qwen3_coder_format_with_qwen3_xml() { + let factory = ToolParserFactory::new(); + let tool_call = " + + +Tokyo + + +"; + + let mut parser = factory.create(names::QWEN3_XML, &[]).unwrap(); + let mut output = ToolParserOutput::default(); + parser.parse_into(tool_call, &mut output).unwrap(); + output.append(parser.finish().unwrap()); + let output = output.coalesce(); + + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[0].arguments, r#"{"location":"Tokyo"}"#); +} + #[test] fn factory_new_resolves_default_patterns() { let factory = ToolParserFactory::new(); diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index c5080a59c5c2..41f06d773432 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -204,7 +204,7 @@ impl ToolCallMixVariant { } impl RoundtripCase { - /// Qwen3 XML tool-call format with `qwen3` reasoning tags. + /// Qwen3 JSON tool-call format with `qwen3` reasoning tags. fn qwen3() -> Self { Self { model_id: "Qwen/Qwen3-0.6B", diff --git a/rust/src/parser/Cargo.toml b/rust/src/parser/Cargo.toml index aff4c795ab63..10972ace3ee3 100644 --- a/rust/src/parser/Cargo.toml +++ b/rust/src/parser/Cargo.toml @@ -53,7 +53,7 @@ harness = false required-features = ["test-util"] [[bench]] -name = "qwen3_xml" +name = "qwen3_json" harness = false required-features = ["test-util"] diff --git a/rust/src/parser/benches/qwen3_xml.rs b/rust/src/parser/benches/qwen3_json.rs similarity index 89% rename from rust/src/parser/benches/qwen3_xml.rs rename to rust/src/parser/benches/qwen3_json.rs index 5138007b541f..a19260ea1d2a 100644 --- a/rust/src/parser/benches/qwen3_xml.rs +++ b/rust/src/parser/benches/qwen3_json.rs @@ -6,7 +6,7 @@ use std::time::Duration; use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main}; use tool_parser::parsers::QwenParser as ExternalQwenParser; use vllm_parser::tool::test_utils::{split_by_chars, test_tools}; -use vllm_parser::tool::{Qwen3XmlToolParser, Tool, ToolParser}; +use vllm_parser::tool::{Qwen3JsonToolParser, Tool, ToolParser}; mod utils; use utils::{feed_external_parser, feed_parser, openai_tools}; @@ -27,12 +27,12 @@ fn mixed_fixture() -> String { } fn long_normal_text_fixture() -> String { - let line = "This is ordinary assistant text with no Qwen XML tool markers at all.\n"; + let line = "This is ordinary assistant text with no Qwen3 JSON tool markers at all.\n"; line.repeat(LONG_NORMAL_TEXT_REPEATS) } fn native_parser(tools: &[Tool]) -> Box { - Qwen3XmlToolParser::create(tools).expect("Qwen XML parser should initialize") + Qwen3JsonToolParser::create(tools).expect("Qwen3 JSON parser should initialize") } fn run_stream_group( @@ -98,14 +98,14 @@ fn run_stream_group( group.finish(); } -fn bench_qwen3_xml(c: &mut Criterion) { +fn bench_qwen3_json(c: &mut Criterion) { let tools = test_tools(); let mixed_text = mixed_fixture(); let long_normal_text = long_normal_text_fixture(); run_stream_group( c, - "qwen3_xml/mixed_text_tool_call", + "qwen3_json/mixed_text_tool_call", &tools, &mixed_text, CHUNK_CHARS, @@ -115,7 +115,7 @@ fn bench_qwen3_xml(c: &mut Criterion) { run_stream_group( c, - "qwen3_xml/long_normal_text", + "qwen3_json/long_normal_text", &tools, &long_normal_text, CHUNK_CHARS, @@ -124,5 +124,5 @@ fn bench_qwen3_xml(c: &mut Criterion) { ); } -criterion_group!(benches, bench_qwen3_xml); +criterion_group!(benches, bench_qwen3_json); criterion_main!(benches); diff --git a/rust/src/parser/src/tool/json/mod.rs b/rust/src/parser/src/tool/json/mod.rs index 487be423aef2..5112ecf5d066 100644 --- a/rust/src/parser/src/tool/json/mod.rs +++ b/rust/src/parser/src/tool/json/mod.rs @@ -9,7 +9,7 @@ pub use internlm2::Internlm2ToolParser; pub use llama::Llama3JsonToolParser; pub use mistral::MistralToolParser; pub use phi4mini::Phi4MiniJsonToolParser; -pub use qwen::Qwen3XmlToolParser; +pub use qwen::Qwen3JsonToolParser; mod granite4; mod hermes; diff --git a/rust/src/parser/src/tool/json/qwen.rs b/rust/src/parser/src/tool/json/qwen.rs index c4db8c2fbc8b..88325d3d9112 100644 --- a/rust/src/parser/src/tool/json/qwen.rs +++ b/rust/src/parser/src/tool/json/qwen.rs @@ -4,8 +4,8 @@ use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace}; use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput}; -const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig { - parser_name: "Qwen XML", +const QWEN3_JSON_CONFIG: JsonToolCallConfig = JsonToolCallConfig { + parser_name: "Qwen3 JSON", start_marker: "", framed_start_marker: Some("\n"), end_marker: "", @@ -15,7 +15,8 @@ const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig { arguments_key: &["arguments"], }; -/// Tool parser for Qwen XML-wrapped JSON tool calls. +/// Tool parser for Qwen2.5, QwQ and Qwen3 JSON tool calls wrapped in +/// `` tags. /// /// Example tool call content: /// @@ -30,20 +31,20 @@ const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig { /// /// Note: parallel calls are represented as repeated /// `...` blocks, not as multiple calls inside one tag. -pub struct Qwen3XmlToolParser { +pub struct Qwen3JsonToolParser { inner: JsonToolCallParser, } -impl Qwen3XmlToolParser { - /// Create a Qwen XML tool parser. +impl Qwen3JsonToolParser { + /// Create a Qwen3 JSON tool parser. fn new(_tools: &[Tool]) -> Self { Self { - inner: JsonToolCallParser::new(QWEN_XML_CONFIG), + inner: JsonToolCallParser::new(QWEN3_JSON_CONFIG), } } } -impl ToolParser for Qwen3XmlToolParser { +impl ToolParser for Qwen3JsonToolParser { fn create(tools: &[Tool]) -> Result> where Self: Sized + 'static, @@ -73,7 +74,7 @@ mod tests { use expect_test::expect; use thiserror_ext::AsReport; - use super::Qwen3XmlToolParser; + use super::Qwen3JsonToolParser; use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; use crate::tool::tests::assert_tool_framing_preserves_body_whitespace; use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; @@ -85,8 +86,8 @@ mod tests { } #[test] - fn qwen_xml_parse_complete_without_tool_call_keeps_text() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_parse_complete_without_tool_call_keeps_text() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let output = parser.parse_complete("Hello, world!").unwrap(); assert_eq!(output.normal_text(), "Hello, world!"); @@ -94,8 +95,8 @@ mod tests { } #[test] - fn qwen_xml_parse_complete_extracts_raw_json_arguments() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_parse_complete_extracts_raw_json_arguments() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let arguments = r#"{ "location": "Tokyo", "days": "3" }"#; let output = parser .parse_complete(&format!( @@ -112,8 +113,8 @@ mod tests { } #[test] - fn qwen_xml_does_not_validate_or_normalize_arguments() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_does_not_validate_or_normalize_arguments() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let arguments = r#"{"location":"Tokyo",}"#; let output = parser.parse_complete(&build_tool_call("get_weather", arguments)).unwrap(); @@ -121,8 +122,8 @@ mod tests { } #[test] - fn qwen_xml_streaming_emits_argument_deltas() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_streaming_emits_argument_deltas() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let chunks = [ "", "\n{\"name\": \"get_weather\", \"arguments\": ", @@ -154,13 +155,13 @@ mod tests { } #[test] - fn qwen_xml_streaming_handles_split_markers() { + fn qwen3_json_streaming_handles_split_markers() { let input = format!( "hello {}", build_tool_call("get_weather", r#"{"location":"Tokyo"}"#) ); let chunks = split_by_chars(&input, 5); - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let output = collect_stream(&mut parser, &chunks); @@ -170,8 +171,8 @@ mod tests { } #[test] - fn qwen_xml_keeps_end_marker_literal_inside_json_string() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_keeps_end_marker_literal_inside_json_string() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let arguments = r#"{"text":"literal inside"}"#; let output = parser.parse_complete(&build_tool_call("echo", arguments)).unwrap(); @@ -180,8 +181,8 @@ mod tests { } #[test] - fn qwen_xml_decodes_escaped_function_name() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_decodes_escaped_function_name() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let output = parser .parse_complete( r#" @@ -194,8 +195,8 @@ mod tests { } #[test] - fn qwen_xml_requires_newline_after_tool_call_start() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_requires_newline_after_tool_call_start() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let input = r#"{"name":"get_weather","arguments":{}} "#; @@ -206,12 +207,12 @@ mod tests { } #[test] - fn qwen_xml_streams_start_marker_without_newline_as_text() { + fn qwen3_json_streams_start_marker_without_newline_as_text() { for chunk in [ "{\"name\":\"get_weather\"", "hi\n{\"name\":\"get_weather\"", ] { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let mut output = ToolParserOutput::default(); parser.parse_into(chunk, &mut output).unwrap(); @@ -223,8 +224,8 @@ mod tests { } #[test] - fn qwen_xml_requires_newline_before_tool_call_end() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_requires_newline_before_tool_call_end() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let error = parser .parse_complete( r#" @@ -236,14 +237,14 @@ mod tests { } #[test] - fn qwen_xml_streaming_extracts_multiple_tool_calls() { + fn qwen3_json_streaming_extracts_multiple_tool_calls() { let input = format!( "{}{}", build_tool_call("get_weather", r#"{"location":"Shanghai"}"#), build_tool_call("add", r#"{"x":1,"y":2}"#), ); let chunks = split_by_chars(&input, 7); - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let output = collect_stream(&mut parser, &chunks); @@ -275,8 +276,8 @@ mod tests { } #[test] - fn qwen_xml_finish_fails_incomplete_tool_call() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_finish_fails_incomplete_tool_call() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); parser .parse_chunk( r#" @@ -286,13 +287,13 @@ mod tests { let error = parser.finish().unwrap_err(); - expect!["tool parser parsing failed: incomplete Qwen XML tool call"] + expect!["tool parser parsing failed: incomplete Qwen3 JSON tool call"] .assert_eq(&error.to_report_string()); } #[test] - fn qwen_xml_malformed_field_order_fails_fast() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_malformed_field_order_fails_fast() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let error = parser .parse_chunk( r#" @@ -302,14 +303,14 @@ mod tests { .unwrap_err(); expect![[r#" - tool parser parsing failed: near "{\"arguments\":{},\"name\":\"get_weather\"}\n": invalid Qwen XML + tool parser parsing failed: near "{\"arguments\":{},\"name\":\"get_weather\"}\n": invalid Qwen3 JSON expected `name`"#]] .assert_eq(&error.to_report_string()); } #[test] fn tool_framing_preserves_body_whitespace_across_chunk_boundaries() { - assert_tool_framing_preserves_body_whitespace::( + assert_tool_framing_preserves_body_whitespace::( "\n", "\n{\"name\":\"get_weather\",\"arguments\":{}}\n", ); diff --git a/rust/src/parser/src/tool/mod.rs b/rust/src/parser/src/tool/mod.rs index 974d6c758a27..69ecee1bffff 100644 --- a/rust/src/parser/src/tool/mod.rs +++ b/rust/src/parser/src/tool/mod.rs @@ -29,7 +29,7 @@ pub use glm_xml::{Glm45MoeToolParser, Glm47MoeToolParser}; pub(crate) use hy::{HyDialect, HyToolMarkers, HyToolParser}; pub use json::{ Granite4ToolParser, HermesToolParser, Internlm2ToolParser, Llama3JsonToolParser, - MistralToolParser, Phi4MiniJsonToolParser, Qwen3XmlToolParser, + MistralToolParser, Phi4MiniJsonToolParser, Qwen3JsonToolParser, }; pub use kimi_k2::KimiK2ToolParser; pub use mimo::MiMoToolParser; diff --git a/rust/src/parser/src/unified/combined.rs b/rust/src/parser/src/unified/combined.rs index b1f36e885546..aaab4312569c 100644 --- a/rust/src/parser/src/unified/combined.rs +++ b/rust/src/parser/src/unified/combined.rs @@ -186,7 +186,7 @@ mod tests { use crate::tool::{ DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser, Glm47MoeToolParser, KimiK2ToolParser, MinimaxM2ToolParser, Qwen3CoderToolParser, - Qwen3XmlToolParser, Tool, ToolParser, + Qwen3JsonToolParser, Tool, ToolParser, }; use crate::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput}; @@ -463,7 +463,7 @@ mod tests { #[test] fn combined_parser_emits_tool_calls_from_visible_content() { - let tool = Qwen3XmlToolParser::create(&test_tools()).unwrap(); + let tool = Qwen3JsonToolParser::create(&test_tools()).unwrap(); let mut parser = CombinedParser::new(None, Some(tool)); let output = collect( @@ -515,7 +515,7 @@ mod tests { #[test] fn split_parser_wrappers_match_builder_full_formats() { - assert_builder_parity::("qwen3"); + assert_builder_parity::("qwen3"); assert_builder_parity::("qwen3.5"); assert_builder_parity::("deepseek-r1"); assert_builder_parity::("deepseek-v3.1"); @@ -538,7 +538,7 @@ mod tests { answer: None, }; let reasoning = Qwen3ReasoningParser::create(Arc::new(tokenizer())).unwrap(); - let tool = Qwen3XmlToolParser::create(&tools).unwrap(); + let tool = Qwen3JsonToolParser::create(&tools).unwrap(); let visible = tool.build_visible_format(&ctx).unwrap().unwrap(); let mut parser = CombinedParser::new(Some(reasoning), Some(tool)); parser.initialize(&[]).unwrap(); @@ -640,7 +640,7 @@ mod tests { }; let tokenizer = Arc::new(tokenizer()); let reasoning = Qwen3ReasoningParser::create(tokenizer.clone()).unwrap(); - let tool = Qwen3XmlToolParser::create(&tools).unwrap(); + let tool = Qwen3JsonToolParser::create(&tools).unwrap(); let mut parser = CombinedParser::new(Some(reasoning), Some(tool)); parser.initialize(&[]).unwrap(); let outside = parser.build_output_grammar(&ctx).unwrap().unwrap().format; @@ -682,7 +682,7 @@ mod tests { fn qwen3_parser(tools: &[Tool], with_tool: bool) -> CombinedParser { let reasoning = Qwen3ReasoningParser::create(Arc::new(tokenizer())).unwrap(); - let tool = with_tool.then(|| Qwen3XmlToolParser::create(tools).unwrap()); + let tool = with_tool.then(|| Qwen3JsonToolParser::create(tools).unwrap()); let mut parser = CombinedParser::new(Some(reasoning), tool); parser.initialize(&[]).unwrap(); parser @@ -702,7 +702,7 @@ mod tests { ]) }; let required = ToolChoice::required(); - let required_calls = Qwen3XmlToolParser::create(&tools) + let required_calls = Qwen3JsonToolParser::create(&tools) .unwrap() .build_visible_format(&grammar_ctx(&tools, &required, None)) .unwrap() @@ -748,7 +748,7 @@ mod tests { let ctx = grammar_ctx(&tools, &auto, Some(&answer)); assert!(parser.build_output_grammar(&ctx).unwrap().is_none()); - let parser = CombinedParser::new(None, Some(Qwen3XmlToolParser::create(&tools).unwrap())); + let parser = CombinedParser::new(None, Some(Qwen3JsonToolParser::create(&tools).unwrap())); let ctx = grammar_ctx(&tools, &none, Some(&answer)); let alone = parser.build_output_grammar(&ctx).unwrap().unwrap(); assert_eq!(alone, BuiltOutputGrammar::final_output_only(answer.clone())); diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py index 8446aae81540..4b9cb33502e5 100644 --- a/vllm/tool_parsers/__init__.py +++ b/vllm/tool_parsers/__init__.py @@ -194,6 +194,13 @@ "qwen3_engine_tool_parser", "Qwen3EngineToolParser", ), + # TODO: Qwen2.5/QwQ/Qwen3 JSON tool calls. Aliased to `hermes` for now; + # a dedicated parser should treat the template's newline before and + # between `` blocks as framing, as the Rust frontend does. + "qwen3_json": ( + "hermes_tool_parser", + "Hermes2ProToolParser", + ), "qwen3_xml": ( "qwen3_engine_tool_parser", "Qwen3EngineToolParser", From 258b97ff6bc1dbbc7b76d2d5a5a87591f9de107f Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Mon, 12 Oct 2026 20:32:39 +0800 Subject: [PATCH 297/624] [Bugfix][EC Connector] Size transfers by measured encoder output widths (#57904) Signed-off-by: Tianyu Guo --- docs/features/ec_cpu_connector.md | 2 +- .../unit/cpu/scheduler/test_scheduler.py | 8 +- .../unit/cpu/test_ec_shared_region.py | 4 +- .../unit/cpu/worker/test_worker.py | 29 +++-- .../unit/test_cross_encoder_cache.py | 28 +++++ .../unit/test_encoder_output_width.py | 92 ++++++++++++++++ .../unit/test_scheduler_nixl_consumer.py | 73 +++++++++---- .../unit/test_scheduler_nixl_producer.py | 14 +-- tests/v1/ec_connector/unit/utils.py | 18 +++- vllm/config/ec_transfer.py | 5 + .../ec_transfer/ec_connector/cpu/common.py | 32 +----- .../ec_transfer/ec_connector/cpu/protocol.py | 4 +- .../ec_connector/cpu/scheduler/__init__.py | 70 ++++++++---- .../cpu/scheduler/embedding_cache.py | 21 ++-- .../ec_connector/cpu/worker/__init__.py | 16 ++- .../ec_connector/encoder_output_width.py | 102 ++++++++++++++++++ .../ec_connector/mooncake/scheduler.py | 14 +-- vllm/v1/engine/core.py | 5 + 18 files changed, 416 insertions(+), 121 deletions(-) create mode 100644 tests/v1/ec_connector/unit/test_encoder_output_width.py create mode 100644 vllm/distributed/ec_transfer/ec_connector/encoder_output_width.py diff --git a/docs/features/ec_cpu_connector.md b/docs/features/ec_cpu_connector.md index 2a6df25f6dbf..18274b6dd970 100644 --- a/docs/features/ec_cpu_connector.md +++ b/docs/features/ec_cpu_connector.md @@ -96,7 +96,7 @@ EC transfer is configured via `--ec-transfer-config` (CLI) or the `ec_transfer_c | --- | --- | --- | --- | | `ec_enable_nixl` | `bool` | No (default `false`) | Enables NIXL P2P transfer in addition to local CPU offload. Omitted or `false` imports no NIXL/ZMQ. Extra config is not type coerced, so a string value is parsed: `"true"`, `"1"`, `"yes"` enable it, anything else does not. | | `consumer_ack_timeout_s` | `float` | No (default `2.0`) | How long a consumer waits for an `XferAck` before giving up on a read. The producer answers `XferReq`s from its own scheduler step, so its reply latency scales with the encoder's `--max-num-batched-tokens`: a loaded encoder whose steps run longer than this makes consumers abandon reads the producer is about to grant. Raise it for large encoder batches. | -| `ec_cpu_bytes` | `int` | Yes | Total size, in bytes, of the shared CPU mmap region. `ECCPUConnector` raises `ValueError` if unset. Block count = `ec_cpu_bytes // block_size_bytes`, where `block_size_bytes = hidden_dim * dtype.element_size()` (`hidden_dim` accounts for Qwen3-VL deepstack: `out_hidden_size * (1 + num_deepstack_layers)`). | +| `ec_cpu_bytes` | `int` | Yes | Total size, in bytes, of the shared CPU mmap region. `ECCPUConnector` raises `ValueError` if unset. Block count = `ec_cpu_bytes // block_size_bytes`, where each block holds one LLM input-embedding vector. Entries allocate enough blocks for their modality-specific output shape, including visual DeepStack features. | ### Environment Variables diff --git a/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py b/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py index 24ea717cabb7..ef86505057cb 100644 --- a/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py +++ b/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py @@ -22,12 +22,16 @@ def __init__(self, offset, length): self.offset = offset self.length = length + def get_num_embeds(self): + return self.length + class _Feature: def __init__(self, mm_hash, length=1, identifier=None): self.mm_hash = mm_hash self.identifier = identifier if identifier is not None else mm_hash self.mm_position = _Pos(0, length) + self.modality = "image" class _Request: @@ -80,7 +84,7 @@ def _make_scheduler( def _load_ids(meta) -> list[int]: """Transfer ids the scheduler dispatched in this step's metadata.""" - return [transfer_id for transfer_id, _ in meta.loads.values()] + return [transfer_id for transfer_id, _, _ in meta.loads.values()] def _load_blocks(meta, mm_hash: str) -> list[int]: @@ -90,7 +94,7 @@ def _load_blocks(meta, mm_hash: str) -> list[int]: def _seed_cached(s: ECCPUScheduler, mm_hash: str, n_blocks: int): """Pre-populate a ready cache entry backed by real blocks.""" - s._cache.alloc(mm_hash, n_blocks) + s._cache.alloc(mm_hash, n_blocks, (n_blocks, _BS // s._dtype.itemsize)) s._cache.mark_ready(mm_hash) diff --git a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py index c168970147b0..0b64d68ca6c0 100644 --- a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py +++ b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py @@ -97,10 +97,9 @@ def test_second_instance_opens_existing_file_and_shares_memory(): r1.cleanup() -def test_dense_dp_engines_get_distinct_regions(monkeypatch): +def test_dense_dp_engines_get_distinct_regions(): """Dense DP engines reset data_parallel_rank to 0 but share instance_id, so the region must be keyed by data_parallel_index.""" - monkeypatch.setattr(common, "_get_encoder_cache_hidden_dim", lambda cfg: 8) instance_id = str(uuid.uuid4()) regions = [] try: @@ -113,6 +112,7 @@ def test_dense_dp_engines_get_distinct_regions(monkeypatch): cfg.instance_id = instance_id cfg.parallel_config = parallel_config cfg.model_config.dtype = torch.float16 + cfg.model_config.get_inputs_embeds_size.return_value = 8 cfg.ec_transfer_config.ec_connector_extra_config = {"ec_cpu_bytes": 64} regions.append(common.create_ec_shared_region(cfg)) diff --git a/tests/v1/ec_connector/unit/cpu/worker/test_worker.py b/tests/v1/ec_connector/unit/cpu/worker/test_worker.py index f4f5f03f6946..4ff0979e305e 100644 --- a/tests/v1/ec_connector/unit/cpu/worker/test_worker.py +++ b/tests/v1/ec_connector/unit/cpu/worker/test_worker.py @@ -91,7 +91,9 @@ def _make_region() -> ECSharedRegion: def _vllm_config(rank: int = 0) -> Mock: - return create_ec_vllm_config(rank=rank, dtype=_DTYPE) + config = create_ec_vllm_config(rank=rank, dtype=_DTYPE) + config.model_config.get_inputs_embeds_size.return_value = _HIDDEN_DIM + return config def _meta( @@ -105,7 +107,7 @@ def _meta( return ECCPUConnectorMetadata( saves=saves or {}, loads={ - mm_hash: (idx, block_ids) + mm_hash: (idx, block_ids, (len(block_ids), _HIDDEN_DIM)) for idx, (mm_hash, block_ids) in enumerate((loads or {}).items()) }, ) @@ -904,21 +906,29 @@ def test_e2e_scheduler_worker_save_then_load(make_worker, monkeypatch): # Build scheduler sharing the same region. monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region) - scheduler = ECCPUScheduler(_vllm_config()) + config = _vllm_config() + # DeepStack-wide output: each embedding spans four blocks. + width = 4 * _HIDDEN_DIM + config.ec_transfer_config.mm_encoder_output_widths = {"image": width} + scheduler = ECCPUScheduler(config) # -- Step 1: scheduler allocates, worker saves -- - n_blocks = 3 - src = torch.arange( - n_blocks * _HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE - ).reshape(n_blocks, _HIDDEN_DIM) + n_tokens = 2 + src = torch.arange(n_tokens * width, dtype=_DTYPE, device=DEVICE_TYPE).reshape( + n_tokens, width + ) class _Pos: offset = 0 - length = n_blocks + length = n_tokens + + def get_num_embeds(self): + return self.length class _Feature: identifier = "img_001" mm_position = _Pos() + modality = "image" class _Request: request_id = "req_e2e" @@ -948,8 +958,9 @@ class _Output: scheduler.update_state_after_alloc(_Request(), 0) meta_load = scheduler.build_connector_meta(scheduler_output=None) assert "img_001" in meta_load.loads - _, load_blocks = meta_load.loads["img_001"] + _, load_blocks, shape = meta_load.loads["img_001"] assert load_blocks == meta_save.saves["img_001"] + assert shape == tuple(src.shape) load_cache: dict[str, torch.Tensor] = {} worker.start_load_caches(load_cache, meta_load) diff --git a/tests/v1/ec_connector/unit/test_cross_encoder_cache.py b/tests/v1/ec_connector/unit/test_cross_encoder_cache.py index 26155d8b6e64..20dd32416c0b 100644 --- a/tests/v1/ec_connector/unit/test_cross_encoder_cache.py +++ b/tests/v1/ec_connector/unit/test_cross_encoder_cache.py @@ -45,6 +45,34 @@ SPEC = TensorSpec((2, 2), "torch.float32", 16) +def test_push_shapes_follow_each_modality(monkeypatch): + """Each push is sized at its modality's measured encoder output width.""" + from vllm.distributed.ec_transfer.ec_connector.mooncake import scheduler + + monkeypatch.setattr(scheduler, "ensure_mooncake_available", lambda: None) + widths = {"image": 8192, "audio": 2048} + config = create_ec_vllm_config(ec_role="ec_producer", encoder_output_widths=widths) + request = SimpleNamespace( + request_id="mixed", + mm_features=[SimpleNamespace(identifier=m, modality=m) for m in widths], + get_num_encoder_embeds=lambda index: 2, + ec_transfer_params={ + "consumer_zmq": "tcp://consumer:1234", + "ec_items": [{"mm_hash": m, "transfer_id": m} for m in widths], + }, + ) + instance = scheduler.ECMooncakeScheduler(config) + try: + for index in range(len(widths)): + instance.update_state_after_alloc(request, index) + metadata = instance.build_connector_meta( + SimpleNamespace(free_encoder_mm_hashes=[], preempted_req_ids=set()) + ) + assert [spec.shape for spec in metadata.pushes] == [(2, 8192), (2, 2048)] + finally: + instance.close() + + def test_embedding_key_combines_namespace_and_identifier(): config = create_ec_vllm_config(ec_role="ec_producer") config.ec_transfer_config.ec_connector_extra_config.update( diff --git a/tests/v1/ec_connector/unit/test_encoder_output_width.py b/tests/v1/ec_connector/unit/test_encoder_output_width.py new file mode 100644 index 000000000000..b7d7c1a8225e --- /dev/null +++ b/tests/v1/ec_connector/unit/test_encoder_output_width.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Startup measurement of the encoder output width EC producers size by.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tests.v1.ec_connector.unit.utils import create_ec_vllm_config +from vllm.distributed.ec_transfer.ec_connector import encoder_output_width as eow + +pytestmark = pytest.mark.cpu_test + + +class _DeepStackModel: + """Emits each modality at its own width, as Qwen3-Omni does.""" + + widths = {"image": 4 * 16, "audio": 16} + + def embed_multimodal(self, *, modality): + return [torch.zeros(3, self.widths[modality])] + + +class _Budget: + def __init__(self, *args, **kwargs): + self.mm_max_toks_per_item = dict.fromkeys(_DeepStackModel.widths, 8) + + def get_dummy_encoder_profile_inputs(self, modality, max_items_per_batch): + assert max_items_per_batch == 1 + return [(modality, None)] + + +def _worker(first_pp_rank=True, supports_mm=True): + runner = SimpleNamespace( + supports_mm_inputs=supports_mm, + get_model=_DeepStackModel, + device=torch.device("cpu"), + ) + pp_group = SimpleNamespace(is_first_rank=first_pp_rank) + return SimpleNamespace(model_runner=runner, vllm_config=None), pp_group + + +@pytest.fixture +def patched(monkeypatch): + import vllm.distributed.parallel_state as ps + import vllm.multimodal.encoder_budget as budget_mod + import vllm.multimodal.utils as mm_utils + + monkeypatch.setattr(budget_mod, "MultiModalBudget", _Budget) + monkeypatch.setattr( + mm_utils, + "group_and_batch_mm_kwargs", + lambda items, **_: iter([(items[0][0], 1, {"modality": items[0][0]})]), + ) + return lambda group: monkeypatch.setattr(ps, "get_pp_group", lambda: group) + + +def test_worker_reports_the_encoder_output_width_per_modality(patched): + """The width comes from the tensor the encoder emits, not from config.""" + worker, group = _worker() + patched(group) + assert eow._measure_on_worker(worker) == {"image": 64, "audio": 16} + + +@pytest.mark.parametrize("first_pp_rank, supports_mm", [(False, True), (True, False)]) +def test_ranks_without_an_encoder_report_nothing(patched, first_pp_rank, supports_mm): + worker, group = _worker(first_pp_rank, supports_mm) + patched(group) + assert eow._measure_on_worker(worker) == {} + + +def test_producer_records_the_widths_its_encoder_ranks_agree_on(): + config = create_ec_vllm_config(ec_role="ec_producer", encoder_output_widths={}) + reports = [{"image": 64}, {}, {"image": 64}] + eow.measure_encoder_output_widths(config, lambda fn: reports) + assert config.ec_transfer_config.mm_encoder_output_widths == {"image": 64} + assert eow.get_encoder_output_width(config, "image") == 64 + + +@pytest.mark.parametrize( + "reports, match", + [ + ([{}, {}], "could not measure"), + ([{"image": 64}, {"image": 16}], "disagree"), + ], + ids=["no-encoder", "disagree"], +) +def test_producer_refuses_to_start_without_a_trustworthy_width(reports, match): + config = create_ec_vllm_config(ec_role="ec_producer", encoder_output_widths={}) + with pytest.raises(ValueError, match=match): + eow.measure_encoder_output_widths(config, lambda fn: reports) diff --git a/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py b/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py index a3e87f9b7888..8c9e14563f5b 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py @@ -25,6 +25,9 @@ class _Pos: def __init__(self, offset, length): self.offset, self.length = offset, length + def get_num_embeds(self): + return self.length + class _Feature: def __init__(self, mm_hash, length=1, data=None): @@ -85,6 +88,9 @@ def close(self): pass +_MISSING = object() + + def _consumer_sched(monkeypatch): def _region(cfg): return ECSharedRegion( @@ -95,16 +101,15 @@ def _region(cfg): # ec_enable_nixl defaults to False, so this builds gate-off; the fields # below flip on the NIXL consumer state directly. - s = ECCPUScheduler(create_ec_vllm_config(ec_role="ec_consumer")) + # No measured widths: a consumer takes each shape from the announcement. + s = ECCPUScheduler( + create_ec_vllm_config(ec_role="ec_consumer", encoder_output_widths={}) + ) # Turn on NIXL consumer state without constructing real transports. s._nixl_enabled = True s._transport = _FakeTransport() s._data = None s._compat_hash = "c" - # _setup_nixl normally computes these from model_config; set them - # directly since this helper builds gate-off then flips fields on. - s._hidden_dim = _HID - s._element_size = _ES s._metadata_resolver._cache["image"] = {"image_grid_thw"} return s @@ -115,24 +120,32 @@ def _params(mm_hash, length): "peer_host": "h", "peer_port": 1, "size_bytes": length * _HID * _ES, + "shape": [length, _HID], } } def test_new_remote_read_defers_then_completes(monkeypatch): + """The read is sized from the announced shape, not a local width.""" + width = 4 * _HID s = _consumer_sched(monkeypatch) fake = _FakeSession() # Route _start_xfer to our fake session instead of real ZMQ/NIXL, but # still reserve a real not-ready cache entry so mark_ready works. - def _fake_start(mm_hash, info, size): - entry = s._cache.alloc(mm_hash, 1) + def _fake_start(mm_hash, info, size, shape): + assert size == width * _ES + assert shape == (1, width) + entry = s._cache.alloc(mm_hash, size // _BS, shape) assert entry is not None fake.started.append(mm_hash) return True monkeypatch.setattr(s, "_start_xfer", _fake_start) - req = _Request([_Feature("h1", 1)], params=_params("h1", 1)) + params = _params("h1", 1) + params["h1"]["size_bytes"] = width * _ES + params["h1"]["shape"] = [1, width] + req = _Request([_Feature("h1", 1)], params=params) # Step 1: unseen remote item -> read started, request deferred. assert s.ensure_cache_available(req, 0) is False @@ -155,6 +168,7 @@ def _fake_start(mm_hash, info, size): s.update_state_after_alloc(req, 0) meta = s.build_connector_meta(scheduler_output=None) assert "h1" in meta.loads + assert meta.loads["h1"][2] == (1, width) assert "h1" not in s._step_completed # cleared by promote # Step 3: still cached -> admitted directly, no new transfer. @@ -176,8 +190,8 @@ def test_failed_read_falls_back_only_with_local_input(monkeypatch, payload): s = _consumer_sched(monkeypatch) fake = _FakeSession() - def _fake_start(mm_hash, info, size): - entry = s._cache.alloc(mm_hash, 1) + def _fake_start(mm_hash, info, size, shape): + entry = s._cache.alloc(mm_hash, 1, shape) assert entry is not None fake.started.append(mm_hash) return True @@ -230,8 +244,8 @@ def test_remote_wait_budget_survives_retries_and_long_steps(monkeypatch, retry): fake = _FakeSession() s._sessions[("h", 1)] = fake - def start(mm_hash, info, size): - assert s._cache.alloc(mm_hash, 1) is not None + def start(mm_hash, info, size, shape): + assert s._cache.alloc(mm_hash, 1, shape) is not None fake.started.append(mm_hash) return True @@ -261,8 +275,8 @@ def test_retryable_read_re_requests_without_admitting(monkeypatch): s = _consumer_sched(monkeypatch) fake = _FakeSession() - def _fake_start(mm_hash, info, size): - entry = s._cache.alloc(mm_hash, 1) + def _fake_start(mm_hash, info, size, shape): + entry = s._cache.alloc(mm_hash, 1, shape) assert entry is not None fake.started.append(mm_hash) return True @@ -287,14 +301,29 @@ def _fake_start(mm_hash, info, size): @pytest.mark.parametrize( - "size_fields", - [{"size_bytes": 999}, {"size_bytes": None}, {"size_bytes": "big"}, {}], - ids=["mismatch", "null-size", "text-size", "no-size"], + "fields", + [ + {"size_bytes": 999}, + {"size_bytes": None}, + {"size_bytes": "big"}, + {"size_bytes": _MISSING}, + {"shape": _MISSING}, + {"shape": [2, _HID]}, + ], + ids=[ + "size-mismatch", + "null-size", + "text-size", + "no-size", + "no-shape", + "wrong-rows", + ], ) -def test_invalid_remote_size_fails_the_request(monkeypatch, size_fields): - """Invalid remote sizes fail the request without raising in the scheduler.""" +def test_invalid_announcement_fails_the_request(monkeypatch, fields): + """An unusable announced size or shape fails the request, not the engine.""" s = _consumer_sched(monkeypatch) - announced = {"peer_host": "h", "peer_port": 1, **size_fields} + announced = {**_params("h1", 1)["h1"], **fields} + announced = {k: v for k, v in announced.items() if v is not _MISSING} req = _Request([_Feature("h1", 1)], params={"h1": announced}) assert s.ensure_cache_available(req, 0) is False assert s.take_unavailable_requests() == {"r1"} @@ -306,7 +335,7 @@ def test_invalid_remote_size_fails_the_request(monkeypatch, size_fields): def test_deferral_budget_is_per_request(monkeypatch): """Shared hashes have independent wait budgets and cancellation cleanup.""" s = _consumer_sched(monkeypatch) - monkeypatch.setattr(s._cache, "alloc", lambda key, n: None) + monkeypatch.setattr(s._cache, "alloc", lambda key, n, shape=None: None) params = _params("h1", 1) old = _Request([_Feature("h1", 1)], params=params, req_id="old") new = _Request([_Feature("h1", 1)], params=params, req_id="new") @@ -419,7 +448,7 @@ def test_in_flight_hash_defers_without_second_transfer(monkeypatch): s = _consumer_sched(monkeypatch) started: list[str] = [] - def _spy_start_xfer(mm_hash, info, size): + def _spy_start_xfer(mm_hash, info, size, shape): started.append(mm_hash) return True diff --git a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py index d7ef19b5f4f4..167f333f7301 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py @@ -16,6 +16,9 @@ class _Pos: def __init__(self, offset, length): self.offset, self.length = offset, length + def get_num_embeds(self): + return self.length + class _Feature: def __init__(self, mm_hash, length=1): @@ -53,11 +56,8 @@ def test_request_finished_producer_emits_params(monkeypatch): # Simulate NIXL-enabled producer bookkeeping without building real NIXL. s._nixl_enabled = True s._peer_host, s._peer_port = "1.2.3.4", 5601 - # _setup_nixl normally computes these from model_config; set them - # directly since this test builds gate-off then flips fields on. - s._hidden_dim, s._element_size = 32, 2 # feature length=2, hidden_dim=32, element_size=2 -> 128 bytes -> 2 blocks. - entry = s._cache.alloc("h1", 2) + entry = s._cache.alloc("h1", 2, (2, 32)) assert entry is not None s._cache.mark_ready("h1") @@ -71,6 +71,7 @@ def test_request_finished_producer_emits_params(monkeypatch): "peer_host": "1.2.3.4", "peer_port": 5601, "size_bytes": 2 * 32 * 2, + "shape": [2, 32], } } s.shutdown() @@ -84,8 +85,7 @@ def test_request_finished_announces_not_ready_entry(monkeypatch): s = _sched_gate_off(monkeypatch) s._nixl_enabled = True s._peer_host, s._peer_port = "1.2.3.4", 5601 - s._hidden_dim, s._element_size = 32, 2 - s._cache.alloc("h1", 2) # allocated but not marked ready + s._cache.alloc("h1", 2, (2, 32)) # allocated but not marked ready delay, params = s.request_finished(_Request([_Feature("h1", length=2)])) assert delay is False @@ -96,6 +96,7 @@ def test_request_finished_announces_not_ready_entry(monkeypatch): "peer_host": "1.2.3.4", "peer_port": 5601, "size_bytes": 2 * 32 * 2, + "shape": [2, 32], } } s.shutdown() @@ -105,7 +106,6 @@ def test_request_finished_skips_unallocated_entry(monkeypatch): s = _sched_gate_off(monkeypatch) s._nixl_enabled = True s._peer_host, s._peer_port = "1.2.3.4", 5601 - s._hidden_dim, s._element_size = 32, 2 # No alloc() for "h1" at all — e.g. the cache was full at save time. delay, params = s.request_finished(_Request([_Feature("h1", length=2)])) diff --git a/tests/v1/ec_connector/unit/utils.py b/tests/v1/ec_connector/unit/utils.py index a718bc63b635..705a8ff7ad1e 100644 --- a/tests/v1/ec_connector/unit/utils.py +++ b/tests/v1/ec_connector/unit/utils.py @@ -18,12 +18,16 @@ _PCP_SIZE = 1 +_EMBED_WIDTH = 32 + + def create_ec_vllm_config( *, ec_role: ECRole = "ec_both", tensor_parallel_size: int = 1, rank: int = 0, dtype: torch.dtype = torch.float16, + encoder_output_widths: dict[str, int] | None = None, ) -> Mock: """Build a `VllmConfig` stand-in for EC connector unit tests. @@ -31,10 +35,8 @@ def create_ec_vllm_config( role derivation and the rank arithmetic the connector depends on behave as they do in production. - `model_config` is a stub because nothing under test needs a real one: the - connector reads only `dtype` here, and the paths that inspect the model - itself are not exercised. Building a real `ModelConfig` would resolve an HF - model and inspect its architecture. + `model_config` supplies a dtype and embedding width without resolving an + HF model or inspecting its architecture. Args: ec_role: EC role this instance plays. `is_ec_producer` and @@ -42,6 +44,8 @@ def create_ec_vllm_config( tensor_parallel_size: TP degree. rank: Global rank of this worker. dtype: Encoder cache dtype. + encoder_output_widths: Per-modality encoder output width, as a + producer measures it at startup. Defaults to the embedding width. Returns: A `VllmConfig`-specced mock carrying the configs above. @@ -59,9 +63,15 @@ def create_ec_vllm_config( ec_role=ec_role, engine_id=str(uuid.uuid4()), ) + ec_transfer_config.mm_encoder_output_widths = ( + dict.fromkeys(("image", "video", "audio"), _EMBED_WIDTH) + if encoder_output_widths is None + else encoder_output_widths + ) model_config = Mock(spec=ModelConfig) model_config.dtype = dtype + model_config.get_inputs_embeds_size.return_value = _EMBED_WIDTH vllm_config = Mock(spec=VllmConfig) vllm_config.ec_transfer_config = ec_transfer_config diff --git a/vllm/config/ec_transfer.py b/vllm/config/ec_transfer.py index 5c2b8e589c2e..b3ca10a5a690 100644 --- a/vllm/config/ec_transfer.py +++ b/vllm/config/ec_transfer.py @@ -55,6 +55,11 @@ class ECTransferConfig: """The Python module path to dynamically load the EC connector from. Only supported in V1.""" + mm_encoder_output_widths: dict[str, int] = field(default_factory=dict, init=False) + """Per-token width of the encoder output for each modality, measured by EC + producers at startup. It can exceed the LM embedding width (e.g. DeepStack). + """ + def compute_hash(self) -> str: """WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/common.py b/vllm/distributed/ec_transfer/ec_connector/cpu/common.py index eff16e08b0ae..461cd6a55a6c 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/common.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/common.py @@ -33,12 +33,14 @@ class ECCPUConnectorMetadata(ECConnectorMetadata): saves: dict[str, list[int]] = field(default_factory=dict) # Consumer role: mm_hashes whose bytes are available in the local mmap, - # mapped to (transfer_id, block_ids); the worker's start_load_caches + # mapped to (transfer_id, block_ids, shape); the worker's start_load_caches # copies mmap[block_ids] → GPU encoder_cache and reports the transfer_id # once its copy lands. The id is minted per dispatch, so two loads of the # same mm_hash in different steps are distinct and a late report from the # earlier one cannot release the later one's pin. - loads: dict[str, tuple[int, list[int]]] = field(default_factory=dict) + loads: dict[str, tuple[int, list[int], tuple[int, int]]] = field( + default_factory=dict + ) @dataclass @@ -66,30 +68,6 @@ def aggregate(self, other: ECConnectorWorkerMetadata) -> ECConnectorWorkerMetada return self -def _get_encoder_cache_hidden_dim(vllm_config: "VllmConfig") -> int: - """Return the per-token hidden dimension for encoder cache entries. - - For most models this equals the LLM's hidden size. Qwen3-VL (and any - future model with deepstack visual encoding) is an exception: the ViT - concatenates its own output with features from N decoder layers before - storing in encoder_cache, producing a tensor of width - ``out_hidden_size * (1 + N)`` per visual token. Using the plain LLM - hidden size would under-allocate EC blocks and silently truncate the - transferred data, leading to a shape mismatch on the consumer. - """ - model_config = vllm_config.model_config - hf_config = getattr(model_config, "hf_config", None) - vision_config = ( - getattr(hf_config, "vision_config", None) if hf_config is not None else None - ) - if vision_config is not None: - out_hidden_size = getattr(vision_config, "out_hidden_size", None) - deepstack_indexes = getattr(vision_config, "deepstack_visual_indexes", None) - if out_hidden_size is not None and deepstack_indexes: - return out_hidden_size * (1 + len(deepstack_indexes)) - return model_config.get_inputs_embeds_size() - - def create_ec_shared_region(vllm_config: "VllmConfig") -> ECSharedRegion: """Build the EC mmap region from `vllm_config`. @@ -103,7 +81,7 @@ def create_ec_shared_region(vllm_config: "VllmConfig") -> ECSharedRegion: engine_id = f"{vllm_config.instance_id}_dp{dp_index}" dtype = vllm_config.model_config.dtype - hidden_dim = _get_encoder_cache_hidden_dim(vllm_config) + hidden_dim = vllm_config.model_config.get_inputs_embeds_size() element_size = torch.empty(0, dtype=dtype).element_size() block_size_bytes = hidden_dim * element_size diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/protocol.py b/vllm/distributed/ec_transfer/ec_connector/cpu/protocol.py index 85bd3cae32fa..e20bf48f0ba8 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/protocol.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/protocol.py @@ -14,9 +14,9 @@ import msgspec -# Bump when the on-wire shape (XferReq/XferAck) changes in a +# Bump when the wire format or tensor/block layout changes in a # backward-incompatible way. Peers reject mismatched versions. -EC_CONNECTOR_VERSION: int = 1 +EC_CONNECTOR_VERSION: int = 2 class XferStatus(enum.IntEnum): diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py index 65403980b9d0..567f2d59c9f3 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py @@ -16,12 +16,14 @@ from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( ECCPUConnectorMetadata, ECCPUWorkerMetadata, - _get_encoder_cache_hidden_dim, create_ec_shared_region, ) from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import ( EmbeddingCache, ) +from vllm.distributed.ec_transfer.ec_connector.encoder_output_width import ( + get_encoder_output_width, +) from vllm.distributed.ec_transfer.ec_connector.utils import ( PlaceholderMetadataResolver, collect_ec_item_metadata, @@ -49,6 +51,7 @@ class ECCPUScheduler: """Scheduler delegate for the ECCPUConnector.""" def __init__(self, vllm_config: "VllmConfig") -> None: + self._vllm_config = vllm_config ec_config = vllm_config.ec_transfer_config assert ec_config is not None self._is_producer: bool = ec_config.is_ec_producer @@ -61,8 +64,8 @@ def __init__(self, vllm_config: "VllmConfig") -> None: # mm_hash → block IDs allocated this step for GPU→mmap saves. self._pending_saves: dict[str, list[int]] = {} - # mm_hash → (transfer_id, block IDs) to load from mmap→GPU this step. - self._pending_loads: dict[str, tuple[int, list[int]]] = {} + # mm_hash → (transfer_id, block IDs, shape) to load from mmap→GPU. + self._pending_loads: dict[str, tuple[int, list[int], tuple[int, int]]] = {} # Dispatched loads awaiting completion reports, keyed by transfer id: # transfer_id → (mm_hash, reports still outstanding). The pin taken at @@ -108,11 +111,9 @@ def __init__(self, vllm_config: "VllmConfig") -> None: self._unrecoverable: set[str] = set() self._peer_host: str | None = None self._peer_port: int | None = None - # Model shape for size checks + compat hash; only set by - # _setup_nixl (the gate-off path never touches model_config). - self._dtype: torch.dtype | None = None - self._hidden_dim: int = 0 - self._element_size: int = 0 + dtype = vllm_config.model_config.dtype + assert isinstance(dtype, torch.dtype) + self._dtype = dtype self._ack_timeout_s: float = 0.0 if self._nixl_enabled: try: @@ -148,9 +149,6 @@ def _setup_nixl(self, vllm_config: "VllmConfig") -> None: ) engine_id = self._ec_config.engine_id assert engine_id is not None - self._dtype = vllm_config.model_config.dtype - self._hidden_dim = _get_encoder_cache_hidden_dim(vllm_config) - self._element_size = torch.empty(0, dtype=self._dtype).element_size() # How long a consumer waits for an XferAck. The producer answers from # its scheduler step, so its reply latency scales with the encoder's # batch size: a deployment whose steps run longer must raise this. @@ -296,22 +294,30 @@ def _admit_item( # held by a quarantined/settling DMA and cannot be reused. return False - expected = pos.length * self._hidden_dim * self._element_size try: size = int(remote["size_bytes"]) + shape = tuple(int(dim) for dim in remote["shape"]) except (KeyError, TypeError, ValueError) as error: raise _RemoteUnavailable("the announced size was unusable") from error - if size != expected: + # The width is the producer's to report; the row count is not. + num_embeds = pos.get_num_embeds() + if ( + len(shape) != 2 + or shape[0] != num_embeds + or size != self._num_blocks(shape) * self._region.block_size_bytes + ): logger.warning( - "EC consumer: size mismatch mm_hash=%s announced=%d expected=%d", + "EC consumer: size mismatch mm_hash=%s announced=%d shape=%s " + "expected %d rows", mm_hash, size, - expected, + shape, + num_embeds, ) raise _RemoteUnavailable("the announced size was wrong") try: - started = self._start_xfer(mm_hash, remote, expected) + started = self._start_xfer(mm_hash, remote, size, (shape[0], shape[1])) except Exception as error: logger.exception( "EC consumer: failed to start NIXL xfer mm_hash=%s", mm_hash @@ -377,7 +383,11 @@ def take_unavailable_requests(self) -> set[str]: return failed def _start_xfer( - self, mm_hash: str, info: "dict[str, Any]", size_bytes: int + self, + mm_hash: str, + info: "dict[str, Any]", + size_bytes: int, + shape: tuple[int, int], ) -> bool: """Allocate a not-ready cache entry and start a NIXL READ into it. @@ -391,7 +401,7 @@ def _start_xfer( ) n_blocks = max(1, ceil(size_bytes / self._region.block_size_bytes)) - entry = self._cache.alloc(mm_hash, n_blocks) + entry = self._cache.alloc(mm_hash, n_blocks, shape) if entry is None: logger.debug( "EC consumer: cache full for mm_hash=%s (%d blocks); deferring", @@ -491,22 +501,36 @@ def _promote_completed_reads(self) -> None: """ self._step_completed.clear() + def _num_blocks(self, shape: tuple[int, ...]) -> int: + size_bytes = shape[0] * shape[1] * self._dtype.itemsize + block_size = self._region.block_size_bytes + return (size_bytes + block_size - 1) // block_size + def update_state_after_alloc(self, request: "Request", index: int) -> None: feature = request.mm_features[index] mm_hash = feature.identifier if self._is_producer and self._cache.get(mm_hash) is None: - entry = self._cache.alloc(mm_hash, feature.mm_position.length) + shape = ( + feature.mm_position.get_num_embeds(), + get_encoder_output_width(self._vllm_config, feature.modality), + ) + entry = self._cache.alloc(mm_hash, self._num_blocks(shape), shape) if entry is not None: self._pending_saves[mm_hash] = list(entry.block_ids) if self._is_consumer and mm_hash not in self._pending_loads: entry = self._cache.get(mm_hash) if entry is not None and entry.ready: + assert entry.shape is not None self._cache.pin(mm_hash) transfer_id = self._next_transfer_id self._next_transfer_id += 1 - self._pending_loads[mm_hash] = (transfer_id, list(entry.block_ids)) + self._pending_loads[mm_hash] = ( + transfer_id, + list(entry.block_ids), + entry.shape, + ) self._load_acks[transfer_id] = (mm_hash, self._expected_load_reports) def build_connector_meta( @@ -608,13 +632,13 @@ def request_finished( # A read arriving before the save lands is NACKed NACK_NOT_READY, # which the consumer retries on a later step rather than treating # as a miss. - size_bytes = ( - feature.mm_position.length * self._hidden_dim * self._element_size - ) + assert entry.shape is not None + size_bytes = len(entry.block_ids) * self._region.block_size_bytes items[mm_hash].update( peer_host=self._peer_host, peer_port=self._peer_port, size_bytes=size_bytes, + shape=list(entry.shape), ) logger.debug( "EC producer: announcing NIXL-readable encodings req_id=%s items=%s", diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py index 2fd4b01059db..046fba05741e 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py @@ -14,13 +14,20 @@ class CacheEntry: - """A single cache entry. Read `.block_ids` and `.ready` freely; - mutations only through EmbeddingCache methods.""" + """A single cache entry. Read `.block_ids`, `.shape` and `.ready` freely; + mutations only through EmbeddingCache methods. - __slots__ = ("block_ids", "_pin_count") + `shape` is the encoder output the blocks hold, so a load can restore it + without knowing how wide this model's encoder output is. + """ + + __slots__ = ("block_ids", "shape", "_pin_count") - def __init__(self, block_ids: tuple[int, ...]) -> None: + def __init__( + self, block_ids: tuple[int, ...], shape: tuple[int, int] | None = None + ) -> None: self.block_ids = block_ids + self.shape = shape self._pin_count = -1 # not ready @property @@ -76,7 +83,9 @@ def has_held_entries(self) -> bool: with self._lock: return len(self._entries) != len(self._entries_free_list) - def alloc(self, key: str, n_blocks: int) -> CacheEntry | None: + def alloc( + self, key: str, n_blocks: int, shape: tuple[int, int] | None = None + ) -> CacheEntry | None: """Allocate *n_blocks* for *key*, evicting as needed. The entry starts not-ready. Returns None if there is not enough @@ -97,7 +106,7 @@ def alloc(self, key: str, n_blocks: int) -> CacheEntry | None: # in the free list when it is evicted, and lets consumers detect # which of an entry's blocks are adjacent. block_ids = tuple(sorted(self._free_blocks.pop() for _ in range(n_blocks))) - entry = CacheEntry(block_ids) + entry = CacheEntry(block_ids, shape) self._entries[key] = entry return entry diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py b/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py index ef7ff0f75162..95c6c9612885 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py @@ -267,8 +267,7 @@ def save_caches( # so they are added up in uint64; signed arithmetic rejects them. src_base = np.uint64(src.view(-1).view(torch.uint8).data_ptr()) dst_base = self._region.blocks.data_ptr() - # Entries are sized by placeholder count, which can exceed the number - # of embeddings, so blocks the output never reaches get no descriptor. + # Copy only the tensor bytes, excluding any block padding. n_used = -(-total_bytes // block_size) if n_used: slots, first_blocks, num_blocks = _coalesce_runs(block_ids[:n_used]) @@ -367,7 +366,7 @@ def start_load_caches( load_items = sorted( connector_metadata.loads.items(), key=lambda kv: kv[1][1][0] ) - total_blocks = sum(len(block_ids) for _, (_, block_ids) in load_items) + total_blocks = sum(len(block_ids) for _, (_, block_ids, _) in load_items) stream = self._acquire_stream() compute_stream = current_platform.current_stream() @@ -384,7 +383,7 @@ def start_load_caches( bufs = self._buf_pool.acquire(total_blocks) slots, first_blocks, num_blocks = _coalesce_runs( - list(chain.from_iterable(ids for _, (_, ids) in load_items)) + list(chain.from_iterable(ids for _, (_, ids, _) in load_items)) ) bufs.add_copies( 0, @@ -405,7 +404,7 @@ def start_load_caches( Transfer( start_event=start_event, end_event=end_event, - completions=[tid for _, (tid, _) in load_items], + completions=[tid for _, (tid, _, _) in load_items], bufs=bufs, stream=stream, num_bytes=total_blocks * block_size, @@ -414,11 +413,10 @@ def start_load_caches( # Slice contiguous buffer into per-hash views. offset = 0 - for mm_hash, (_, block_ids) in load_items: + for mm_hash, (_, block_ids, shape) in load_items: n = len(block_ids) - encoder_cache[mm_hash] = ( - dst_buf[offset : offset + n].view(dtype).reshape(n, -1) - ) + flat = dst_buf[offset : offset + n].view(dtype).reshape(-1) + encoder_cache[mm_hash] = flat[: shape[0] * shape[1]].view(shape) offset += n current_platform.current_stream().wait_stream(stream) diff --git a/vllm/distributed/ec_transfer/ec_connector/encoder_output_width.py b/vllm/distributed/ec_transfer/ec_connector/encoder_output_width.py new file mode 100644 index 000000000000..c1a270bb2b18 --- /dev/null +++ b/vllm/distributed/ec_transfer/ec_connector/encoder_output_width.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Encoder output widths for sizing EC transfers. + +An encoder can emit features wider than the LM embedding (DeepStack packs +several layers on the last dim) by an amount only the model knows. Producers +size transfers before the first real item is encoded, so the engine measures +the width once at startup by encoding a dummy item per modality. +""" + +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.logger import init_logger + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +logger = init_logger(__name__) + + +@torch.inference_mode() +def _measure_on_worker(worker: Any) -> dict[str, int]: + """Encode one dummy item per modality; runs on every worker. + + Returns {} on ranks that hold no encoder. + """ + from vllm.distributed.parallel_state import get_pp_group + from vllm.multimodal import MULTIMODAL_REGISTRY + from vllm.multimodal.encoder_budget import MultiModalBudget + from vllm.multimodal.utils import group_and_batch_mm_kwargs + from vllm.utils.torch_utils import PIN_MEMORY + + runner = worker.model_runner + if not runner.supports_mm_inputs or not get_pp_group().is_first_rank: + return {} + budget = MultiModalBudget( + worker.vllm_config, MULTIMODAL_REGISTRY, enable_cache=False + ) + model = runner.get_model() + widths: dict[str, int] = {} + for modality in budget.mm_max_toks_per_item: + items = budget.get_dummy_encoder_profile_inputs(modality, 1) + if not items: + continue + _, _, mm_kwargs = next( + group_and_batch_mm_kwargs( + items, device=runner.device, pin_memory=PIN_MEMORY + ) + ) + outputs = model.embed_multimodal(**mm_kwargs) + widths[modality] = int(outputs[0].shape[-1]) + del outputs + return widths + + +def measure_encoder_output_widths( + vllm_config: "VllmConfig", collective_rpc: Callable[..., list[Any]] +) -> None: + """Record the encoder output widths on an EC producer's config. + + Call after the model is loaded and before the KV cache is allocated: the + Scheduler's connector reads the widths, and once the KV cache holds the + memory budget there is no room left for a dummy encode. + + Raises: + ValueError: No worker could encode a dummy item, or workers disagree. + + """ + ec_config = vllm_config.ec_transfer_config + if ec_config is None or not ec_config.is_ec_producer: + return + reports = [widths for widths in collective_rpc(_measure_on_worker) if widths] + if not reports: + raise ValueError( + "EC producer could not measure the encoder output width: the model " + "encoded no dummy multimodal input." + ) + if any(widths != reports[0] for widths in reports): + raise ValueError(f"Workers disagree on the encoder output widths: {reports}") + logger.info("EC producer encoder output widths: %s", reports[0]) + ec_config.mm_encoder_output_widths = reports[0] + + +def get_encoder_output_width(vllm_config: "VllmConfig", modality: str) -> int: + """Encoder output width for `modality`, as measured at producer startup. + + Raises: + ValueError: No width was measured for `modality`, so a transfer + cannot be sized. + + """ + assert vllm_config.ec_transfer_config is not None + widths = vllm_config.ec_transfer_config.mm_encoder_output_widths + if modality not in widths: + raise ValueError( + f"No encoder output width was measured for modality {modality!r}; " + f"measured: {sorted(widths)}" + ) + return widths[modality] diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py b/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py index 2a607118db45..597ba767dfc8 100644 --- a/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py @@ -20,8 +20,8 @@ from vllm.distributed.ec_transfer.ec_connector.base import ( ECConnectorMetadata, ) -from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( - _get_encoder_cache_hidden_dim, +from vllm.distributed.ec_transfer.ec_connector.encoder_output_width import ( + get_encoder_output_width, ) from vllm.distributed.ec_transfer.ec_connector.mooncake.config import ( _RESERVATION_TTL_SECONDS, @@ -85,9 +85,7 @@ def __init__(self, vllm_config: VllmConfig) -> None: self._is_consumer = config.is_consumer self._control_addr = config.control_addr self._push_wait_timeout = config.push_wait_timeout_s - self._encoder_cache_hidden_dim = ( - _get_encoder_cache_hidden_dim(vllm_config) if config.is_producer else None - ) + self._vllm_config = vllm_config self._model_config = vllm_config.model_config self._control_client = ControlClient(config.control_timeout_ms) self._control_executor = ThreadPoolExecutor( @@ -447,8 +445,10 @@ def _prepare_push_spec(self, request: Any, index: int) -> None: def _encoder_output_spec(self, request: Any, index: int) -> TensorSpec: dtype = self._model_config.dtype assert isinstance(dtype, torch.dtype) - assert self._encoder_cache_hidden_dim is not None - shape = (request.get_num_encoder_embeds(index), self._encoder_cache_hidden_dim) + hidden_dim = get_encoder_output_width( + self._vllm_config, request.mm_features[index].modality + ) + shape = (request.get_num_encoder_embeds(index), hidden_dim) return TensorSpec(shape, str(dtype), math.prod(shape) * dtype.itemsize) def update_state_after_alloc(self, request: Any, index: int) -> None: diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index bbca64a58ae9..35cc8cc1860c 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -27,6 +27,9 @@ cleanup_dist_env_and_memory, stateless_destroy_torch_distributed_process_group, ) +from vllm.distributed.ec_transfer.ec_connector.encoder_output_width import ( + measure_encoder_output_widths, +) from vllm.envs import enable_envs_cache from vllm.logger import configure_logging, init_logger from vllm.logging_utils.dump_input import dump_engine_exception @@ -148,6 +151,8 @@ def __init__( if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: self._eep_scale_up_before_kv_init() + measure_encoder_output_widths(vllm_config, self.collective_rpc) + # Setup KV Caches and update CacheConfig after profiling. kv_cache_config = self._initialize_kv_caches(vllm_config) self.structured_output_manager = StructuredOutputManager(vllm_config) From c75cb63cd2bb8e0dc62608e9930d8b953ee3b422 Mon Sep 17 00:00:00 2001 From: Harry Huang Date: Mon, 12 Oct 2026 20:56:44 +0800 Subject: [PATCH 298/624] [Qwen4Exp] support Sequence Parallelism for Qwen3.8-Flash-Next (#56322) Signed-off-by: huanghaoyan.hhy Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .buildkite/test_areas/models_basic.yaml | 3 +- .buildkite/test_areas/models_distributed.yaml | 17 ++ .../core/test_vocab_parallel_embedding.py | 42 +++ tests/models/qwen4_exp/test_config.py | 86 ++++-- tests/models/qwen4_exp/test_ple.py | 103 ++++--- tests/models/qwen4_exp/test_tpsp.py | 267 ++++++++++++++++++ tests/models/test_registry.py | 16 ++ tests/test_config.py | 41 +++ vllm/config/model.py | 8 + vllm/config/parallel.py | 5 + vllm/engine/arg_utils.py | 3 + .../layers/vocab_parallel_embedding.py | 35 +-- vllm/model_executor/models/__init__.py | 4 + vllm/model_executor/models/config.py | 10 + vllm/model_executor/models/interfaces.py | 21 ++ vllm/model_executor/models/qwen3_next.py | 8 +- vllm/model_executor/models/registry.py | 12 + .../qwen4_exp/common/ngram_embedding.py | 56 ++-- vllm/models/qwen4_exp/common/ops/ple.py | 40 +-- vllm/models/qwen4_exp/nvidia/model.py | 124 ++++++-- vllm/models/qwen4_exp/nvidia/mtp.py | 44 ++- .../qwen4_exp/nvidia/ngram_embedding.py | 45 ++- vllm/models/qwen4_exp/nvidia/ple_layer.py | 60 ++-- 23 files changed, 871 insertions(+), 179 deletions(-) create mode 100644 tests/models/qwen4_exp/test_tpsp.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 483a9aceac22..288e7621cab4 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -205,12 +205,12 @@ steps: device: h200_35gb source_file_dependencies: - vllm/models/qwen4_exp/ - - vllm/transformers_utils/configs/qwen4_exp.py - vllm/config/compilation.py - vllm/v1/attention/backends/short_conv_attn.py - tests/models/qwen4_exp/ commands: - pytest -v -s models/qwen4_exp/test_hc_ops.py models/qwen4_exp/test_qsa_prepare.py models/qwen4_exp/test_qsa_reference.py + - pytest -v -s models/qwen4_exp/test_ple.py -k test_fused_conv_correctness mirror: amd: label: ":amd: (MI355 DPX) Qwen4 Exp" @@ -231,7 +231,6 @@ steps: device: cpu-small source_file_dependencies: - vllm/models/qwen4_exp/ - - vllm/transformers_utils/configs/qwen4_exp.py - vllm/config/compilation.py - vllm/v1/attention/backends/short_conv_attn.py - tests/models/qwen4_exp/ diff --git a/.buildkite/test_areas/models_distributed.yaml b/.buildkite/test_areas/models_distributed.yaml index 9ab75cfc6e70..8d2bc5dbaf3c 100644 --- a/.buildkite/test_areas/models_distributed.yaml +++ b/.buildkite/test_areas/models_distributed.yaml @@ -2,6 +2,23 @@ group: Models - Distributed depends_on: - image-build steps: +- label: ":nvidia: (L4) Qwen4 Exp SP" + key: qwen4-exp-sp-tests-4-gpus + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/tests" + device: l4 + num_devices: 4 + source_file_dependencies: + - vllm/config/parallel.py + - vllm/models/common/ops/sequence_parallel.py + - vllm/models/qwen4_exp/ + - vllm/model_executor/layers/vocab_parallel_embedding.py + - vllm/model_executor/models/qwen3_next.py + - tests/models/qwen4_exp/test_ple.py + - tests/models/qwen4_exp/test_tpsp.py + commands: + - pytest -v -s models/qwen4_exp/test_tpsp.py + - label: ":nvidia: (L4) Distributed Models Shard %N" key: distributed-model-tests-2-gpus timeout_in_minutes: 35 diff --git a/tests/kernels/core/test_vocab_parallel_embedding.py b/tests/kernels/core/test_vocab_parallel_embedding.py index b0038c084e3c..b50b5e568871 100644 --- a/tests/kernels/core/test_vocab_parallel_embedding.py +++ b/tests/kernels/core/test_vocab_parallel_embedding.py @@ -9,6 +9,8 @@ reconstructs a full-table lookup. """ +from unittest.mock import Mock + import pytest import torch @@ -124,3 +126,43 @@ def test_rank_sum_is_full_lookup(tp_size): total += _fused(ids, _shard(full_table, si, rows, dtype), si) torch.testing.assert_close(total, full_table[ids.long()], atol=0.0, rtol=0.0) + + +@requires_cuda +@torch.inference_mode() +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("use_fused", [False, True]) +def test_embedding_optionally_reduces_output(dtype, use_fused, default_vllm_config): + """SP's deferred reduction preserves masked contributions and their dtype.""" + with torch.device("cuda"): + values = torch.arange(128 * 16).reshape(128, 16) + table = ((values % 29 - 14) / 16).to(dtype) + ids = torch.tensor([0, 63, 64, 127, 12]) + expected = table[ids].float() + for rank in range(2): + local = expected.clone() + local[ids // 64 != rank] = 0 + remote = (expected - local).to(dtype) + group = Mock(world_size=2, rank_in_group=rank) + # Model the other vocabulary owner's contribution, including FP8 bytes. + group.all_reduce.side_effect = lambda output, remote=remote: ( + output + remote.view(output.dtype) + ) + for reduce_results in (True, False): + embedding = VocabParallelEmbedding( + 128, + 16, + params_dtype=dtype, + parallel_group=group, + reduce_results=reduce_results, + ) + embedding.use_fused_embedding = use_fused + embedding.weight_loader(embedding.get_parameter("weight"), table) + actual = embedding(ids) + assert actual.dtype == dtype + torch.testing.assert_close( + actual.float(), + expected if reduce_results else local, + atol=0, + rtol=0, + ) diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index 6fde77eda246..818f51543857 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -9,6 +9,7 @@ import torch from transformers import Qwen4ExpConfig, Qwen4ExpTextConfig +from vllm.config import CacheConfig, ParallelConfig from vllm.config.speculative import SpeculativeConfig from vllm.model_executor.models.config import ( Qwen3_5ForConditionalGenerationConfig, @@ -48,23 +49,54 @@ def _text_config(**kwargs) -> Qwen4ExpTextConfig: return Qwen4ExpTextConfig(**values) -def test_qwen4_exp_mtp_returns_sample_and_multi_streams() -> None: +def test_moe_sp_rejects_pipeline_parallel() -> None: + """Automatic MoE SP must validate HC topology even without --enable-hc-sp.""" + config = SimpleNamespace( + model_config=SimpleNamespace( + hf_text_config=_text_config(num_experts=4, ple_layer_ids=[]), + ), + cache_config=CacheConfig(), + parallel_config=ParallelConfig( + tensor_parallel_size=2, + pipeline_parallel_size=2, + data_parallel_size=2, + enable_expert_parallel=True, + all2all_backend="allgather_reducescatter", + ), + ) + with ( + patch("vllm.platforms.current_platform.is_cuda", return_value=True), + pytest.raises(ValueError, match="MoE SP requires PP=1"), + ): + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(config) + + +@pytest.mark.parametrize("use_sp", [False, True]) +def test_qwen4_exp_mtp_returns_sample_and_multi_streams(use_sp: bool) -> None: + """MTP returns complete, contiguous outputs after gathering padded SP shards.""" from vllm.models.qwen4_exp.nvidia.mtp import ( Qwen4ExpMultiTokenPredictor, ) model = object.__new__(Qwen4ExpMultiTokenPredictor) + model.use_hc_sequence_parallel = use_sp + model.use_sequence_parallel = False torch.nn.Module.__init__(model) model.hc_count = 2 model.hidden_size = 4 model.num_mtp_layers = 1 - model.layers = [ - lambda **kwargs: ( - kwargs["hidden_states"], - kwargs["hidden_states"], - torch.zeros(kwargs["hidden_states"].shape[0], 2), - ), - ] + model.pre_fc_norm_embedding = torch.nn.Identity() + model.pre_fc_norm_hidden = torch.nn.Identity() + model.fc_embedding = torch.nn.Identity() + model.fc_hidden = torch.nn.Identity() + + def decoder(**kwargs): + """Check that the embedding residual shares the hidden state's token slice.""" + hidden = kwargs["hidden_states"] + assert kwargs["prev_block_output"].shape == (hidden.shape[0], model.hidden_size) + return hidden, hidden, torch.zeros(hidden.shape[0], model.hc_count) + + model.layers = [decoder] model.hyper_connection_mixer = SimpleNamespace( combine_and_mix=lambda hidden_states, block_output, injection: ( hidden_states, @@ -72,24 +104,32 @@ def test_qwen4_exp_mtp_returns_sample_and_multi_streams() -> None: None, ), ) - multi_hidden = torch.arange(16, dtype=torch.float32).reshape(2, 8) - pp_group = SimpleNamespace(is_first_rank=False, is_last_rank=True) - - with patch( - "vllm.models.qwen4_exp.nvidia.mtp.get_pp_group", - return_value=pp_group, + multi_hidden = torch.arange(24, dtype=torch.float32).reshape(3, 8) + expected_sample = multi_hidden.unflatten(-1, (2, 4)).mean(-2) + pp_group = SimpleNamespace(is_first_rank=True, is_last_rank=True) + + def gather(local): + """Supply rank 1's last token and padding for a simulated two-rank gather.""" + assert local.shape[0] == 2 + remote = torch.cat([expected_sample[2:], multi_hidden[2:]], dim=-1) + return torch.cat([local, remote, torch.full_like(remote, float("nan"))]) + + with ( + patch("vllm.models.qwen4_exp.nvidia.mtp.get_pp_group", return_value=pp_group), + patch("vllm.models.qwen4_exp.nvidia.mtp.sp_shard", side_effect=lambda x: x[:2]), + patch("vllm.models.qwen4_exp.nvidia.mtp.sp_all_gather", side_effect=gather), ): sample_hidden, returned_multi_hidden = model.forward( input_ids=None, - positions=torch.arange(2), - intermediate_tensors={"hidden_states": multi_hidden}, + positions=torch.arange(3), + hidden_states=multi_hidden, + inputs_embeds=torch.zeros(3, 4), ) - torch.testing.assert_close( - sample_hidden, - multi_hidden.unflatten(-1, (2, 4)).mean(dim=-2), - ) - assert returned_multi_hidden is multi_hidden + torch.testing.assert_close(sample_hidden, expected_sample) + torch.testing.assert_close(returned_multi_hidden, multi_hidden) + assert sample_hidden.is_contiguous() + assert returned_multi_hidden.is_contiguous() @spawn_new_process_for_each_test @@ -150,9 +190,7 @@ def test_qwen4_exp_rejects_pipeline_parallel_only_with_ple(ple_layer_ids) -> Non hf_text_config=_text_config(ple_layer_ids=ple_layer_ids), multimodal_config=None, ), - parallel_config=SimpleNamespace( - pipeline_parallel_size=2, enable_dbo=False, ubatch_size=1 - ), + parallel_config=ParallelConfig(pipeline_parallel_size=2), speculative_config=None, ) with patch.object( diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index d7ee461d8f0d..a4a428ecccea 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -611,46 +611,65 @@ def test_ple_embedding_dtype_overrides_modelopt_exclusion() -> None: ) -def test_pinned_embedding_forward_finalizes_prefetched_output( +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("reduce_results", [False, True]) +def test_pinned_embedding_finalize_prefetch_optionally_reduces_output( monkeypatch: pytest.MonkeyPatch, + dtype: torch.dtype, + reduce_results: bool, ) -> None: + """Deferring reduction must preserve full token rows and wait for the lookup.""" + total_tokens = 3 embedding = Qwen4ExpPLEPinnedHostEmbedding.__new__(Qwen4ExpPLEPinnedHostEmbedding) nn.Module.__init__(embedding) - embedding._prefetch_buffer = torch.empty(4, 2, 3, dtype=torch.float8_e4m3fn) + values = torch.arange(1, 13).reshape(4, 3) + embedding._prefetch_buffer = torch.stack( + (values, torch.zeros_like(values)), dim=1 + ).to(dtype) embedding._output_dim = 6 - hidden_states = torch.zeros(2, 4, dtype=torch.bfloat16) - expected = torch.arange(12).reshape(2, 6).to(torch.float8_e4m3fn) - - def finalize_prefetched( - prefetch_output: torch.Tensor, - output: torch.Tensor, - ) -> None: - assert prefetch_output is embedding._prefetch_buffer - output.copy_(expected) - + embedding.tp_size = 2 + embedding.reduce_results = reduce_results + embedding.etp_data_parallel_size = 1 + embedding._prefetch_stream = object() + waited_streams: list[object] = [] + reduced_dtypes = [] monkeypatch.setattr( - embedding, - "_finalize_prefetch", - finalize_prefetched, + torch.cuda, + "current_stream", + lambda: SimpleNamespace(wait_stream=waited_streams.append), ) - output = embedding(hidden_states) + def all_reduce(tensor: torch.Tensor) -> torch.Tensor: + assert waited_streams == [embedding._prefetch_stream] + reduced_dtypes.append(tensor.dtype) + # The remote rank owns the other head's vocabulary rows. + return tensor + tensor.flip(1) - assert output.dtype == embedding._prefetch_buffer.dtype - assert torch.equal(output, expected) + embedding.parallel_group = SimpleNamespace(all_reduce=all_reduce) + output = embedding.finalize_prefetch(total_tokens=total_tokens) + + expected = ( + values.unsqueeze(1).expand(-1, 2, -1).to(dtype) + if reduce_results + else embedding._prefetch_buffer + ) + expected = expected[:total_tokens].flatten(-2) + comm_dtype = torch.int8 if dtype == torch.float8_e4m3fn else dtype + assert reduced_dtypes == ([comm_dtype] if reduce_results else []) + assert waited_streams == [embedding._prefetch_stream] + assert output.dtype == dtype + torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0) -def test_pinned_embedding_forward_requires_prior_start_prefetch() -> None: - """Forward before any start_prefetch must fail loudly, not read garbage.""" +def test_pinned_embedding_finalize_prefetch_requires_prior_start_prefetch() -> None: + """Finalizing before start_prefetch must fail instead of reading garbage.""" embedding = Qwen4ExpPLEPinnedHostEmbedding.__new__(Qwen4ExpPLEPinnedHostEmbedding) nn.Module.__init__(embedding) embedding._prefetch_buffer = None embedding._prefetch_stream = None embedding._output_dim = 6 - hidden_states = torch.zeros(2, 4, dtype=torch.bfloat16) - with pytest.raises(RuntimeError, match="prior start_prefetch"): - embedding(hidden_states) + embedding.finalize_prefetch(total_tokens=2) def test_pinned_embedding_finalize_requires_prior_start_prefetch() -> None: @@ -781,15 +800,10 @@ def test_ple_pinned_embedding_loads_on_cpu_and_looks_up_through_uva( @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_pinned_embedding_start_prefetch_allocates_lazily_and_feeds_forward( +def test_pinned_embedding_start_prefetch_allocates_lazily( monkeypatch: pytest.MonkeyPatch, ) -> None: - """start_prefetch lazily allocates the stream/buffer, then forward consumes it. - - Exercises the full pinned-prefetch contract end to end: first call - allocates the side stream and buffer with the constructor geometry, - and forward finalizes the lookup into the returned output. - """ + """Allocate prefetch storage on first use and finalize the lookup results.""" _mock_etp_group(monkeypatch) monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0) monkeypatch.setattr( @@ -825,8 +839,7 @@ def test_pinned_embedding_start_prefetch_allocates_lazily_and_feeds_forward( ) ngram_ids = torch.tensor([[3, 0], [1, 2]], device="cuda:0") - hidden_states = torch.zeros(2, 4, dtype=torch.bfloat16, device="cuda:0") - embedding.start_prefetch(hidden_states, ngram_ids) + embedding.start_prefetch(ngram_ids) buffer = embedding._prefetch_buffer assert buffer is not None @@ -835,7 +848,7 @@ def test_pinned_embedding_start_prefetch_allocates_lazily_and_feeds_forward( assert buffer.device.type == "cuda" assert embedding._prefetch_stream is not None - output = embedding(hidden_states) + output = embedding.finalize_prefetch(total_tokens=ngram_ids.shape[0]) expected = loaded_weight[ngram_ids.cpu()].to(device="cuda:0").flatten(-2) torch.testing.assert_close(output, expected, rtol=0, atol=0) @@ -1149,7 +1162,7 @@ def test_ngram_prefetch_ids_outlive_start_prefetch() -> None: prefetched: list[torch.Tensor] = [] module.ngram_embedding = SimpleNamespace( supports_prefetch=True, - start_prefetch=lambda _, ids: prefetched.append(weak_ref_tensor(ids)), + start_prefetch=lambda ids: prefetched.append(weak_ref_tensor(ids)), ) expected = _reference_ngram_ids(input_ids, query_start_loc, ngram_context, **params) @@ -1157,7 +1170,7 @@ def test_ngram_prefetch_ids_outlive_start_prefetch() -> None: # like a later CUDA graph segment would. pool = torch.cuda.MemPool() with torch.cuda.use_mem_pool(pool): - module.start_prefetch(None, input_ids, query_start_loc, ngram_context) + module.start_prefetch(input_ids, query_start_loc, ngram_context) torch.full_like(expected, -1) assert torch.equal(prefetched[0], expected) @@ -1859,12 +1872,18 @@ def _make_conv_metadata( ], ) @pytest.mark.parametrize( - "layer_cls", [Qwen4ExpPLELayer, Qwen4ExpPLELayerAMD], ids=["nvidia", "amd"] + ("layer_cls", "fuse_residual"), + [ + pytest.param(Qwen4ExpPLELayer, True, id="nvidia-fused"), + pytest.param(Qwen4ExpPLELayer, False, id="nvidia-unfused"), + pytest.param(Qwen4ExpPLELayerAMD, True, id="amd-fused"), + ], ) def test_fused_conv_correctness( case: _ConvBatchCase, state_layout: str, layer_cls: type[nn.Module], + fuse_residual: bool, ) -> None: device = torch.device("cuda") metadata, num_real_tokens = _make_conv_metadata(case, device) @@ -1906,12 +1925,15 @@ def test_fused_conv_correctness( ) null_state = conv_state[NULL_BLOCK_ID].clone() residual_kernel = residual.clone() - residual_reference = residual.clone() + # The unfused kernel must overwrite the buffer without reading its contents. + residual_reference = ( + residual.clone() if fuse_residual else torch.zeros_like(residual) + ) module._short_conv_dilated_dispatch( inputs=inputs, residual=residual_kernel, - outer_residual=outer_residual, + outer_residual=outer_residual if fuse_residual else None, metadata=metadata, conv_state=conv_state, conv_weights=weights, @@ -1925,7 +1947,8 @@ def test_fused_conv_correctness( conv_state_len=module.conv_state_len, dilation=module.short_conv_dilation, ) - residual_reference = outer_residual + residual_reference + if fuse_residual: + residual_reference = outer_residual + residual_reference assert torch.equal( residual_kernel[:num_real_tokens], residual_reference[:num_real_tokens] @@ -1935,7 +1958,7 @@ def test_fused_conv_correctness( if case.graph_padding: assert torch.equal( residual_kernel[num_real_tokens:], - (outer_residual + residual)[num_real_tokens:], + residual_reference[num_real_tokens:], ) diff --git a/tests/models/qwen4_exp/test_tpsp.py b/tests/models/qwen4_exp/test_tpsp.py new file mode 100644 index 000000000000..5182ab09c398 --- /dev/null +++ b/tests/models/qwen4_exp/test_tpsp.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Multi-GPU equivalence tests for Qwen4Exp HC and MoE sequence parallelism.""" + +from dataclasses import MISSING, fields +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +import torch.multiprocessing as mp +from torch import nn +from transformers import Qwen4ExpTextConfig + +from vllm.platforms import current_platform + +if not current_platform.is_cuda(): + pytest.skip("Qwen4Exp SP requires CUDA", allow_module_level=True) + +from tests.utils import multi_gpu_test +from vllm.config import ( + EngramConfig, + ParallelConfig, + VllmConfig, + set_current_vllm_config, +) +from vllm.distributed import ( + cleanup_dist_env_and_memory, + init_distributed_environment, + initialize_model_parallel, + tensor_model_parallel_all_reduce, +) +from vllm.forward_context import set_forward_context +from vllm.model_executor.layers.linear import RowParallelLinear +from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_padding_mask, + sp_shard, +) +from vllm.models.qwen4_exp.nvidia import model as qwen4_model +from vllm.models.qwen4_exp.nvidia.model import Qwen4ExpDecoderLayer +from vllm.utils.network_utils import get_open_port +from vllm.v1.attention.backends.short_conv_attn import PleShortConvAttentionMetadata +from vllm.v1.worker.workspace import init_workspace_manager, reset_workspace_manager + +from .test_ple import _ConvBatchCase, _make_conv_metadata + + +class _TPAttentionProjection(nn.Module): + """Keep real TP communication without setting up an attention KV cache.""" + + def __init__(self, config, *, reduce_results: bool, **kwargs) -> None: + super().__init__() + self.proj = RowParallelLinear( + config.hidden_size, + config.hidden_size, + bias=False, + input_is_parallel=False, + reduce_results=reduce_results, + ) + + def forward(self, hidden_states: torch.Tensor): + return self.proj(hidden_states)[0] + + +def _make_config(rank: int, *, hc_sp: bool = False, moe_sp: bool = False) -> VllmConfig: + """Configure HC SP explicitly or MoE SP through TP=2, DP=2 and EP.""" + config = VllmConfig( + parallel_config=ParallelConfig( + tensor_parallel_size=2, + data_parallel_size=2 if moe_sp else 1, + data_parallel_rank=rank // 2, + enable_expert_parallel=moe_sp, + enable_hc_sp=hc_sp, + all2all_backend="allgather_reducescatter", + ), + ) + config.model_config = SimpleNamespace( + architecture="Qwen4ExpForCausalLM", + is_moe=True, + sleep_mode_offload_cudagraph=False, + hf_text_config=Qwen4ExpTextConfig( + vocab_size=128, + eos_token_id=0, + hidden_size=128, + hc_count=2, + hc_lowrank=16, + num_hidden_layers=1, + layer_types=["linear_attention"], + num_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=256, + shared_expert_intermediate_size=256, + ple_layer_ids=[] if moe_sp else [1], + ple_embed_dim=128, + heads_per_ngram=2, + ngram_vocab_size_base=32, + make_ngram_vocab_size_divisible_by=64, + ), + dtype=torch.bfloat16, + ) + config.engram_config = EngramConfig(cpu_offload=False) + return config + + +def _make_decoder(config: VllmConfig, rank: int) -> Qwen4ExpDecoderLayer: + """Use the production constructor and identical initial weights for both modes.""" + with ( + set_current_vllm_config(config), + patch.object(qwen4_model, "QwenGatedDeltaNetAttention", _TPAttentionProjection), + ): + layer = Qwen4ExpDecoderLayer( + config, config.model_config.hf_text_config.layer_types[0], "model.layers.0" + ) + torch.manual_seed(12) + for param in layer.parameters(): + param.normal_(std=0.08) + # Distinct rank contributions expose missing or duplicate reductions. + layer.linear_attn.proj.weight.add_(rank * 0.01) + experts = layer.mlp.experts.routed_experts + if config.parallel_config.enable_expert_parallel: + experts.w13_weight.add_(rank * 0.01) + experts.quant_method.process_weights_after_loading(experts) + return layer + + +def _check_decoder_ple(rank: int) -> None: + """Compare HC SP with full tokens across a PLE prefill and cached decode.""" + configs = [_make_config(rank, hc_sp=enabled) for enabled in (False, True)] + layers = [_make_decoder(config, rank) for config in configs] + states = [] + for layer in layers: + ple = layer.ple + assert ple is not None + state = torch.zeros(2, ple.hc_hidden_size, ple.conv_state_len) + if not is_conv_state_dim_first(): + state = state.transpose(-1, -2).contiguous() + ple.kv_cache = (state,) + states.append(state) + + # Three prefill tokens straddle the TP=2 boundary and require SP padding. + for case in ( + _ConvBatchCase(prefill_query_lens=(3,)), + _ConvBatchCase(num_decodes=1), + ): + layout, num_tokens = _make_conv_metadata(case, torch.device("cuda", rank)) + metadata_args = { + f.name: None + for f in fields(PleShortConvAttentionMetadata) + if f.default is MISSING and f.default_factory is MISSING + } + metadata = PleShortConvAttentionMetadata(**(metadata_args | vars(layout))) + ids = torch.randint(0, 128, (num_tokens,)) + positions = torch.arange(num_tokens).expand(3, -1) + qsl = torch.tensor([0, num_tokens], dtype=torch.int32) + context = torch.zeros(1, layers[0].config.ngram_size - 1, dtype=torch.int64) + hidden = torch.randn(num_tokens, 256) + block = torch.randn(num_tokens, 128) + injection = torch.randn(num_tokens, 2) + outputs = [] + for config, layer in zip(configs, layers): + assert layer.ple is not None + inputs = (hidden, block, injection) + if config.parallel_config.enable_hc_sp: + inputs = tuple(sp_shard(x) for x in inputs) + with ( + set_current_vllm_config(config), + set_forward_context({layer.ple.prefix: metadata}, config), + ): + layer.ple.start_prefetch(ids, qsl, context) + output = layer( + *inputs, + positions, + input_ids=ids, + query_start_loc=qsl, + ngram_context=context, + ) + if config.parallel_config.enable_hc_sp: + output = tuple(sp_all_gather(x)[:num_tokens] for x in output) + # A later MoE call may reuse workspace-backed output storage. + outputs.append(tuple(x.clone() for x in output)) + torch.testing.assert_close(outputs[1], outputs[0], atol=0.002, rtol=0.02) + torch.testing.assert_close(states[1], states[0], atol=0.002, rtol=0.02) + + +def _check_moe_sp(rank: int, config: VllmConfig) -> None: + """Compare native SP with full-token GR and the MoE wrapper's public interface.""" + layer = _make_decoder(config, rank) + dp_rank = config.parallel_config.data_parallel_rank + counts = torch.tensor([3, 5], device="cpu") + num_tokens = int(counts[dp_rank]) + torch.manual_seed(30 + dp_rank) + hidden = torch.randn(num_tokens, 256) + positions = torch.arange(num_tokens) + # Include runner padding in addition to the padding introduced by SP. + padding = sp_padding_mask(positions == num_tokens - 1, hidden) + with ( + set_current_vllm_config(config), + set_forward_context( + None, + config, + num_tokens=num_tokens, + num_tokens_across_dp=counts, + is_padding=padding, + ), + ): + # The reference keeps full GR rows. MoE handles its own chunk and gather. + full_hidden, attn_input, injection = layer.attn_hyper_connection.mix(hidden) + attn_output = tensor_model_parallel_all_reduce(layer.linear_attn(attn_input)) + full_hidden, moe_input, injection = layer.mlp_hyper_connection.combine_and_mix( + full_hidden, attn_output, injection + ) + expected = tuple( + x.clone() for x in (full_hidden, layer.mlp(moe_input), injection) + ) + output = layer( + sp_shard(hidden), + None, + None, + positions, + input_ids=None, + query_start_loc=None, + ngram_context=None, + ) + actual = tuple(sp_all_gather(x)[:num_tokens] for x in output) + torch.testing.assert_close(actual, expected, atol=0.002, rtol=0.02) + + +def _run_sp(rank: int, port: int, moe_sp: bool) -> None: + """Initialize rank-local runtime state and run one equivalence check.""" + torch.accelerator.set_device_index(rank) + torch.set_num_threads(1) + config = _make_config(rank, moe_sp=moe_sp) + init_distributed_environment( + world_size=4 if moe_sp else 2, + rank=rank, + local_rank=rank, + distributed_init_method=f"tcp://127.0.0.1:{port}", + ) + try: + # Spawned ranks bypass the GPU worker's workspace initialization. + init_workspace_manager(torch.device(f"cuda:{rank}")) + with set_current_vllm_config(config): + initialize_model_parallel(tensor_model_parallel_size=2) + with torch.device(f"cuda:{rank}"), torch.inference_mode(): + torch.set_default_dtype(torch.bfloat16) + if moe_sp: + _check_moe_sp(rank, config) + else: + _check_decoder_ple(rank) + finally: + reset_workspace_manager() + cleanup_dist_env_and_memory() + + +@multi_gpu_test(num_gpus=2) +def test_hc_sp_decoder_and_ple() -> None: + """HC SP preserves decoder output and PLE cache across a token-shard boundary.""" + mp.spawn(_run_sp, args=(get_open_port(), False), nprocs=2) + + +@multi_gpu_test(num_gpus=4) +def test_moe_sp_equivalence(monkeypatch) -> None: + """Native MoE SP preserves outputs with unequal DP batches and padding.""" + monkeypatch.setenv("VLLM_MOE_SKIP_PADDING", "1") + mp.spawn(_run_sp, args=(get_open_port(), True), nprocs=4) diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index d269aeec1330..cffc72804326 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -127,6 +127,22 @@ def test_registry_is_pp(model_arch, is_pp): assert model_info.supports_pp is is_pp +@pytest.mark.parametrize( + "model_arch,supported", + [ + ("Qwen4ExpForCausalLM", not current_platform.is_rocm()), + ("Qwen4ExpForConditionalGeneration", not current_platform.is_rocm()), + ("Qwen4ExpMTP", not current_platform.is_rocm()), + ("LlamaForCausalLM", False), + ], +) +def test_registry_supports_hc_sp(model_arch, supported): + """HC SP is opt-in for each model implementation, including MTP.""" + model_info = ModelRegistry._try_inspect_model_cls(model_arch) + assert model_info is not None + assert model_info.supports_hc_sp is supported + + @create_new_process_for_each_test() @pytest.mark.parametrize( "model_arch", diff --git a/tests/test_config.py b/tests/test_config.py index 734253edc9f8..a7e36ada6259 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -4109,6 +4109,47 @@ def test_revision_resolved_for_model(mock_resolve): mock_resolve.assert_any_call(model, None, config.hf_token) +@pytest.mark.parametrize( + "enabled,supported,is_cuda,tp_size,pp_size,error", + [ + (True, True, True, 2, 1, None), + (True, False, True, 2, 1, "not supported for this model"), + (True, True, False, 2, 1, "requires CUDA"), + (True, True, True, 1, 1, None), + (True, False, False, 1, 2, None), + (True, True, True, 2, 2, "requires PP=1"), + (False, False, False, 2, 1, None), + ], +) +def test_hc_sp_model_support(enabled, supported, is_cuda, tp_size, pp_size, error): + """Validate active HC SP configurations and ignore the flag with TP=1.""" + from vllm.model_executor.models.registry import ModelRegistry + + model = SimpleNamespace( + architectures=["HCModel"], + registry=ModelRegistry, + model_arch_config=SimpleNamespace(total_num_attention_heads=8), + multimodal_config=None, + ) + parallel = ParallelConfig( + enable_hc_sp=enabled, + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + # Avoid backend auto-detection in this config-only test. + distributed_executor_backend="mp", + ) + with ( + patch("vllm.config.model.current_platform.is_cuda", return_value=is_cuda), + patch.object(ModelRegistry, "is_hc_sp_supported_model", return_value=supported), + patch.object(ModelRegistry, "is_pp_supported_model", return_value=True), + ): + if error: + with pytest.raises(ValueError, match=error): + ModelConfig.verify_with_parallel_config(model, parallel) + else: + ModelConfig.verify_with_parallel_config(model, parallel) + + @pytest.mark.parametrize( ("layer_types", "expected_attention"), [ diff --git a/vllm/config/model.py b/vllm/config/model.py index d4bafad5d6d3..545c306e1e1f 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1537,6 +1537,14 @@ def verify_with_parallel_config( self, parallel_config: ParallelConfig, ) -> None: + if parallel_config.enable_hc_sp and parallel_config.tensor_parallel_size > 1: + if not current_platform.is_cuda(): + raise ValueError("--enable-hc-sp currently requires CUDA") + if not self.registry.is_hc_sp_supported_model(self.architectures, self): + raise ValueError("--enable-hc-sp is not supported for this model") + if parallel_config.pipeline_parallel_size != 1: + raise ValueError("--enable-hc-sp requires PP=1") + total_num_attention_heads = self.model_arch_config.total_num_attention_heads tensor_parallel_size = parallel_config.tensor_parallel_size if total_num_attention_heads % tensor_parallel_size != 0: diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 6960661e0c7e..e103c9ffd73d 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -181,6 +181,11 @@ class ParallelConfig: --data-parallel-start-rank.""" is_moe_model: bool | None = None """Whether the deployed model is MoE (if known).""" + enable_hc_sp: bool = False + """Enable sequence parallelism for HyperConnection (HC) computation. + Currently supported by Qwen4Exp on CUDA with TP>1 and PP=1 (ignored when TP=1). + Does not change MoE parallelism. MoE SP may enable HC SP independently of this flag. + """ enable_expert_parallel: bool = False """Use expert parallelism instead of tensor parallelism for MoE layers.""" enable_batch_sharded_sampling: bool | None = None diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 31777e266f8b..a89290081605 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -538,6 +538,7 @@ class EngineArgs: data_parallel_external_lb: bool = False data_parallel_multi_port_external_lb: bool = False data_parallel_backend: DataParallelBackend = ParallelConfig.data_parallel_backend + enable_hc_sp: bool = ParallelConfig.enable_hc_sp enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel enable_batch_sharded_sampling: bool | None = ( ParallelConfig.enable_batch_sharded_sampling @@ -1258,6 +1259,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "server per local data parallel rank and exposes aggregated health on " "a supervisor port.", ) + parallel_group.add_argument("--enable-hc-sp", **parallel_kwargs["enable_hc_sp"]) parallel_group.add_argument( "--enable-expert-parallel", "-ep", @@ -2532,6 +2534,7 @@ def create_engine_config( data_parallel_backend=self.data_parallel_backend, data_parallel_hybrid_lb=self.data_parallel_hybrid_lb, is_moe_model=model_config.is_moe, + enable_hc_sp=self.enable_hc_sp, enable_expert_parallel=self.enable_expert_parallel, enable_batch_sharded_sampling=self.enable_batch_sharded_sampling, enable_ep_weight_filter=self.enable_ep_weight_filter, diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index be20c412438c..92e65c839dbf 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -27,6 +27,7 @@ method_has_implemented_embedding, resolve_quant_method, ) +from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8 from vllm.model_executor.layers.utils import dispatch_unquantized_gemm from vllm.model_executor.parameter import BasevLLMParameter from vllm.model_executor.utils import set_weight_attrs @@ -240,6 +241,7 @@ class VocabParallelEmbedding(PluggableLayer): quant_config: quant config for the layer prefix: full name of the layer in the state dict disable_tp: If true, tensor parallelism will be disabled for this layer. + reduce_results: Whether to all-reduce contributions from vocabulary shards. quant_method: Preselected quantization method for model-specific layers. parallel_group: Process group used to shard and reduce the embedding. @@ -258,6 +260,7 @@ def __init__( prefix: str = "", *, disable_tp: bool = False, + reduce_results: bool = True, quant_method: QuantizeMethodBase | None = None, parallel_group: GroupCoordinator | None = None, ): @@ -266,6 +269,7 @@ def __init__( # Keep the input dimensions. self.disable_tp = disable_tp self.parallel_group = parallel_group + self.reduce_results = reduce_results if disable_tp: tp_rank, self.tp_size = 0, 1 elif parallel_group is not None: @@ -516,6 +520,7 @@ def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): param[loaded_weight.shape[0] :].data.fill_(0) def forward(self, input_): + """Look up IDs, optionally reducing contributions from vocabulary shards.""" if self.tp_size == 1: return self.quant_method.embedding(self, input_.long()) @@ -543,24 +548,20 @@ def forward(self, input_): self.shard_indices.added_vocab_end_index, ) output_parallel = self.quant_method.embedding(self, masked_input.long()) - if output_parallel.dtype in ( - torch.float8_e4m3fn, - torch.float8_e5m2, - ): - # Each vocab token has one owner, so FP8 bytes can use int8 SUM. - comm_output = output_parallel.view(torch.int8) - comm_output.masked_fill_(input_mask.unsqueeze(-1), 0) - output = ( - self.parallel_group.all_reduce(comm_output) - if self.parallel_group is not None - else tensor_model_parallel_all_reduce(comm_output) - ) - return output.view(output_parallel.dtype) + + output_dtype = output_parallel.dtype + if is_fp8(output_dtype): + # Each vocab token has one owner, so FP8 bytes can use int8 SUM. + output_parallel = output_parallel.view(torch.int8) + if not self.use_fused_embedding: output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0) - # Reduce across all the model parallel GPUs. - if self.parallel_group is not None: - return self.parallel_group.all_reduce(output_parallel) - return tensor_model_parallel_all_reduce(output_parallel) + if self.reduce_results: + output_parallel = ( + self.parallel_group.all_reduce(output_parallel) + if self.parallel_group is not None + else tensor_model_parallel_all_reduce(output_parallel) + ) + return output_parallel.view(output_dtype) def extra_repr(self) -> str: s = f"num_embeddings={self.num_embeddings}" diff --git a/vllm/model_executor/models/__init__.py b/vllm/model_executor/models/__init__.py index 478b00d0af33..78fa39ee413c 100644 --- a/vllm/model_executor/models/__init__.py +++ b/vllm/model_executor/models/__init__.py @@ -3,6 +3,7 @@ from .interfaces import ( HasInnerState, + SupportsHCSP, SupportsLoRA, SupportsMRoPE, SupportsMultiModal, @@ -10,6 +11,7 @@ SupportsPP, SupportsTranscription, has_inner_state, + supports_hc_sp, supports_lora, supports_mrope, supports_multimodal, @@ -33,6 +35,8 @@ "is_text_generation_model", "HasInnerState", "has_inner_state", + "SupportsHCSP", + "supports_hc_sp", "SupportsLoRA", "supports_lora", "SupportsMultiModal", diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 7f6a3aef40a7..d64e51bfa8bf 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -912,11 +912,21 @@ class Qwen4ExpForConditionalGenerationConfig(Qwen3_5ForConditionalGenerationConf @staticmethod def verify_and_update_config(vllm_config: "VllmConfig") -> None: + from vllm.platforms import current_platform + Qwen3_5ForConditionalGenerationConfig.verify_and_update_config(vllm_config) text_config = vllm_config.model_config.hf_text_config if text_config.hc_count <= 1: raise ValueError("Qwen4Exp requires hc_count > 1") parallel_config = vllm_config.parallel_config + # MoE SP also enables HC SP without the explicit HC flag. + if ( + current_platform.is_cuda() + and parallel_config.use_sequence_parallel_moe + and parallel_config.pipeline_parallel_size != 1 + ): + raise ValueError("Qwen4Exp MoE SP requires PP=1") + uses_ple_or_qsa = bool(text_config.ple_layer_ids) or ( getattr(text_config, "indexer_n_heads", None) is not None ) diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 272b78da50a4..4dc67ff34d27 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -858,6 +858,27 @@ def _supports_pp_inspect(model: type[object] | object) -> bool: return supports_kw(model_forward, "intermediate_tensors") +@runtime_checkable +class SupportsHCSP(Protocol): + """Support sequence-parallel HC without changing MoE parallelism.""" + + supports_hc_sp: ClassVar[Literal[True]] = True + + +@overload +def supports_hc_sp(model: type[object]) -> TypeIs[type[SupportsHCSP]]: ... + + +@overload +def supports_hc_sp(model: object) -> TypeIs[SupportsHCSP]: ... + + +def supports_hc_sp( + model: type[object] | object, +) -> TypeIs[type[SupportsHCSP]] | TypeIs[SupportsHCSP]: + return getattr(model, "supports_hc_sp", False) + + @runtime_checkable class HasInnerState(Protocol): """The interface required for all models that has inner state.""" diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 04d20c574145..ee5c05d0a4ae 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -129,7 +129,12 @@ def _should_replicate_misaligned_shared_expert( class Qwen3NextSparseMoeBlock(nn.Module): - def __init__(self, vllm_config: VllmConfig, prefix: str = ""): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + reduce_results: bool = True, + ): super().__init__() config = vllm_config.model_config.hf_text_config @@ -231,6 +236,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = ""): num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, n_shared_experts=1 if self.shared_expert is None else None, + reduce_results=reduce_results, fuse_shared_experts=self.is_fused_shared_expert_enabled, shared_expert_prefix=f"{prefix}.shared_expert", shared_expert_gate=self.shared_expert_gate diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index fa2924441314..427648045f9d 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -49,6 +49,7 @@ is_attention_free, is_hybrid, requires_raw_input_tokens, + supports_hc_sp, supports_mamba_prefix_caching, supports_multimodal, supports_multimodal_encoder_tp_data, @@ -873,6 +874,7 @@ class _ModelInfo: requires_raw_input_tokens: bool supports_multimodal_encoder_tp_data: bool supports_pp: bool + supports_hc_sp: bool has_inner_state: bool is_attention_free: bool is_hybrid: bool @@ -903,6 +905,7 @@ def from_model_cls(model: type[nn.Module]) -> "_ModelInfo": model ), supports_pp=supports_pp(model), + supports_hc_sp=supports_hc_sp(model), has_inner_state=has_inner_state(model), is_attention_free=is_attention_free(model), is_hybrid=is_hybrid(model), @@ -1463,6 +1466,15 @@ def is_pp_supported_model( model_info, _ = self.inspect_model_cls(architectures, model_config) return model_info.supports_pp + def is_hc_sp_supported_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + """Check support for HC SP without changing MoE parallelism.""" + model_info, _ = self.inspect_model_cls(architectures, model_config) + return model_info.supports_hc_sp + def model_has_inner_state( self, architectures: str | list[str], diff --git a/vllm/models/qwen4_exp/common/ngram_embedding.py b/vllm/models/qwen4_exp/common/ngram_embedding.py index 6dad6b7b5afe..85a382c5ec67 100644 --- a/vllm/models/qwen4_exp/common/ngram_embedding.py +++ b/vllm/models/qwen4_exp/common/ngram_embedding.py @@ -70,6 +70,7 @@ def __init__( num_ngram_heads: int = 1, max_total_tokens: int = 0, data_parallel_rank: int = 0, + reduce_results: bool = True, ) -> None: del num_ngram_heads, max_total_tokens super().__init__( @@ -80,6 +81,7 @@ def __init__( prefix=prefix, quant_method=embedding_method, parallel_group=get_etp_group(), + reduce_results=reduce_results, ) self.embedding_method = embedding_method self.data_parallel_rank = data_parallel_rank @@ -153,14 +155,15 @@ def _select_embeddings( return embeddings.narrow(0, slot_offset, local_num_tokens) @abstractmethod - def start_prefetch( - self, - hidden_states: torch.Tensor, - ngram_ids: torch.Tensor, - ) -> None: + def start_prefetch(self, ngram_ids: torch.Tensor) -> None: """Start an asynchronous lookup when supported.""" raise NotImplementedError + @abstractmethod + def finalize_prefetch(self, total_tokens: int) -> torch.Tensor: + """Return prefetched embeddings for the full input token set.""" + raise NotImplementedError + class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase): """Quantization interface shared by resident and pinned PLE tables.""" @@ -351,16 +354,16 @@ def allocate_embedding_weight( """Allocate the complete PLE weight on the active device.""" return torch.empty(num_embeddings, embedding_dim, dtype=dtype) - def start_prefetch( - self, - hidden_states: torch.Tensor, - ngram_ids: torch.Tensor, - ) -> None: + def start_prefetch(self, ngram_ids: torch.Tensor) -> None: """Resident embedding prefetch is a no-op.""" return None + def finalize_prefetch(self, total_tokens: int) -> torch.Tensor: + """Reject prefetch finalization for resident embeddings.""" + raise RuntimeError("Device PLE embedding does not support prefetch") + def forward(self, ngram_ids: torch.Tensor) -> torch.Tensor: - """Gather ETP inputs, look up embeddings, and select local rows.""" + """Gather ETP IDs, optionally reduce lookup results, and select local rows.""" slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0]) gathered_ids = self._gather_dp_ids(ngram_ids, slot_size) embeddings = super().forward(gathered_ids) @@ -422,6 +425,7 @@ def __init__( num_ngram_heads: int = 1, max_total_tokens: int = 0, data_parallel_rank: int = 0, + reduce_results: bool = True, ) -> None: if not is_uva_available(): raise RuntimeError("Engram CPU offload requires UVA support") @@ -435,6 +439,7 @@ def __init__( num_ngram_heads=num_ngram_heads, max_total_tokens=max_total_tokens, data_parallel_rank=data_parallel_rank, + reduce_results=reduce_results, ) self._uva_weight = get_accelerator_view_from_cpu_tensor(self.weight) self._row_bytes = self.embedding_dim * self.weight.element_size() @@ -521,11 +526,7 @@ def _reduce_etp_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor: return self.parallel_group.all_reduce(embeddings) @eager_break_during_capture - def start_prefetch( - self, - hidden_states: torch.Tensor, - ngram_ids: torch.Tensor, - ) -> None: + def start_prefetch(self, ngram_ids: torch.Tensor) -> None: """Gather ETP IDs and launch their UVA lookup on the side stream.""" buffer = self._prefetch_buffer if buffer is None: @@ -574,14 +575,18 @@ def _finalize_prefetch( prefetch_output: torch.Tensor, output: torch.Tensor, ) -> None: - """Join the side stream, reduce ETP shards, and select local rows.""" + """Join the side stream, optionally reduce ETP shards, and select rows.""" prefetch_stream = self._prefetch_stream if prefetch_stream is None: raise RuntimeError("pinned PLE finalize requires a prior start_prefetch") torch.cuda.current_stream().wait_stream(prefetch_stream) slot_size, slot_offset = self._get_dp_gather_slot(output.shape[0]) - active_output = prefetch_output[: slot_size * self.etp_data_parallel_size] - embeddings = self._reduce_etp_embeddings(active_output) + embeddings = prefetch_output[: slot_size * self.etp_data_parallel_size] + if self.reduce_results: + embeddings = self._reduce_etp_embeddings(embeddings) + else: + # The outer TP reduce-scatter requires ETP=TP. + assert self.etp_data_parallel_size == 1 embeddings = self._select_embeddings( embeddings, output.shape[0], @@ -589,11 +594,18 @@ def _finalize_prefetch( ) output.copy_(embeddings.flatten(-2)) - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - """Finish the pinned lookup into graph-owned output storage.""" + def finalize_prefetch(self, total_tokens: int) -> torch.Tensor: + """Finalize the full-token pinned lookup with optional ETP reduction.""" buffer = self._prefetch_buffer if buffer is None: raise RuntimeError("pinned PLE lookup requires a prior start_prefetch") - output = buffer.new_empty((hidden_states.shape[0], self._output_dim)) + # total_tokens counts this DP rank's tokens before SP sharding. + output = buffer.new_empty((total_tokens, self._output_dim)) self._finalize_prefetch(buffer, output) return output + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + raise RuntimeError( + "Pinned-host PLE embedding requires " + "start_prefetch() and finalize_prefetch()." + ) diff --git a/vllm/models/qwen4_exp/common/ops/ple.py b/vllm/models/qwen4_exp/common/ops/ple.py index bc4c99957d90..41ce082c9e70 100644 --- a/vllm/models/qwen4_exp/common/ops/ple.py +++ b/vllm/models/qwen4_exp/common/ops/ple.py @@ -344,6 +344,7 @@ def _ple_conv_kernel( MODE: tl.constexpr, HAS_INIT: tl.constexpr, NULL_STATE_ID: tl.constexpr, + FUSE_RESIDUAL: tl.constexpr, launch_pdl: tl.constexpr, ): t = tl.program_id(0) @@ -438,21 +439,23 @@ def _ple_conv_kernel( conv = acc.to(residual_ptr.dtype.element_ty).to(tl.float32) y = conv * tl.sigmoid(conv) conv_output = tl.where(out_ok, y, 0.0).to(residual_ptr.dtype.element_ty) - residual = tl.load( - residual_ptr + output_t * C + c_offs, - mask=c_mask, - other=0.0, - ) - # Preserve the original eager operation boundaries: short convolution is - # first accumulated into the BF16/FP16 PLE output, then the outer residual - # is added to that rounded value. - ple_output = (residual + conv_output).to(residual_ptr.dtype.element_ty) - outer_residual = tl.load( - outer_residual_ptr + output_t * C + c_offs, - mask=c_mask, - other=0.0, - ) - ple_output = outer_residual.to(tl.float32) + ple_output.to(tl.float32) + ple_output = conv_output + if FUSE_RESIDUAL: + residual = tl.load( + residual_ptr + output_t * C + c_offs, + mask=c_mask, + other=0.0, + ) + # Preserve the original eager operation boundaries: short convolution is + # first accumulated into the BF16/FP16 PLE output, then the outer residual + # is added to that rounded value. + ple_output = (residual + conv_output).to(residual_ptr.dtype.element_ty) + outer_residual = tl.load( + outer_residual_ptr + output_t * C + c_offs, + mask=c_mask, + other=0.0, + ) + ple_output = outer_residual.to(tl.float32) + ple_output.to(tl.float32) if launch_pdl: tl.extra.cuda.gdc_launch_dependents() tl.store( @@ -575,7 +578,7 @@ def ple_conv( conv_state: torch.Tensor, conv_weights: torch.Tensor, state_indices: torch.Tensor, - outer_residual: torch.Tensor, + outer_residual: torch.Tensor | None, *, mode: Literal["decode", "spec", "prefill"], dilation: int, @@ -585,7 +588,9 @@ def ple_conv( spec_query_len: int = 1, token_indices: torch.Tensor | None = None, ) -> None: - """Add short convolution and the outer residual; update state.""" + """Update state and write activated convolution into residual. + A non-None outer_residual enables both fused residual additions. + """ BLOCK_C = 512 kernel_spec_query_len = spec_query_len if mode == "spec" else 1 T, C = inputs.shape @@ -657,6 +662,7 @@ def ple_conv( SPEC_QUERY_LEN=kernel_spec_query_len, MODE=mode, HAS_INIT=has_initial_states_arg, + FUSE_RESIDUAL=outer_residual is not None, NULL_STATE_ID=NULL_BLOCK_ID, launch_pdl=launch_pdl, num_warps=num_warps, diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 2de883cc4747..2fb056bde43c 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -9,8 +9,10 @@ from torch import nn from transformers import Qwen4ExpConfig, Qwen4ExpTextConfig +from vllm import envs from vllm.config import VllmConfig from vllm.distributed import get_pp_group +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.fused_moe.utils import ( is_model_fused_shared_expert_compatible, ) @@ -36,6 +38,7 @@ IsHybrid, MixtureOfExperts, MultiModalEmbeddings, + SupportsHCSP, SupportsLoRA, SupportsMRoPE, SupportsPP, @@ -66,6 +69,12 @@ maybe_fuse_shared_experts, maybe_prefix, ) +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_padding_mask, + sp_reduce_scatter, + sp_shard, +) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.sequence import IntermediateTensors @@ -84,6 +93,14 @@ _QSA_LAYER_TYPES = ("qwen_sparse_attention", "indexed_attention") +def is_hc_sequence_parallel_enabled(vllm_config: VllmConfig) -> bool: + """Enable GR/PLE SP explicitly or automatically with sequence-parallel MoE.""" + parallel_config = vllm_config.parallel_config + return parallel_config.tensor_parallel_size > 1 and ( + parallel_config.enable_hc_sp or parallel_config.use_sequence_parallel_moe + ) + + def without_modelopt_fp4( quant_config: QuantizationConfig | None, ) -> QuantizationConfig | None: @@ -161,13 +178,19 @@ def _remap_qsa_cache_scale_name( class Qwen4ExpSparseMoeBlock(Qwen3NextSparseMoeBlock): """Qwen3Next MoE with Qwen4Exp HC validation.""" - def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: - parallel_config = vllm_config.parallel_config - if parallel_config.use_sequence_parallel_moe: - raise NotImplementedError( - "Qwen4Exp HC does not support sequence-parallel MoE" - ) - super().__init__(vllm_config=vllm_config, prefix=prefix) + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + reduce_results: bool = True, + ) -> None: + super().__init__( + vllm_config=vllm_config, prefix=prefix, reduce_results=reduce_results + ) + if self.replicate_shared_expert and not self.is_sequence_parallel: + # Reduce routed outputs before adding replicated shared outputs. + # SP then slices the result to avoid summing shared outputs again via RS. + self.experts.moe_config.skip_final_all_reduce = False config = vllm_config.model_config.hf_text_config self.n_shared_experts = int(config.shared_expert_intermediate_size > 0) @@ -185,11 +208,12 @@ def __init__( self.config = config self.layer_type = layer_type + # MoE SP also enables HC SP, while HC SP can be enabled independently. + self.use_hc_sequence_parallel = is_hc_sequence_parallel_enabled(vllm_config) self.layer_idx = extract_layer_index(prefix) - if vllm_config.parallel_config.use_sequence_parallel_moe: - raise NotImplementedError( - "Qwen4Exp HC does not support sequence-parallel MoE" - ) + self.use_sequence_parallel = ( + vllm_config.parallel_config.use_sequence_parallel_moe + ) self.ple: Qwen4ExpPLELayer | None = None ple_layer_ids = config.ple_layer_ids if (self.layer_idx + 1) in ple_layer_ids: @@ -204,6 +228,7 @@ def __init__( layer_idx=self.layer_idx, ple_dense_layer_id=ple_dense_layer_id, prefix=f"{prefix}.ple", + use_sequence_parallel=self.use_hc_sequence_parallel, ) if layer_type == "linear_attention": @@ -212,6 +237,7 @@ def __init__( vllm_config=vllm_config, prefix=f"{prefix}.linear_attn", gqa_interleaved_layout=False, + reduce_results=not self.use_hc_sequence_parallel, ) elif layer_type in _QSA_LAYER_TYPES: use_qsa = getattr(config, "indexer_n_heads", None) is not None @@ -222,6 +248,7 @@ def __init__( cache_config=vllm_config.cache_config, quant_config=quant_config, prefix=f"{prefix}.self_attn", + reduce_results=not self.use_hc_sequence_parallel, ) else: self.self_attn = Qwen4ExpQSAAttention( @@ -230,12 +257,15 @@ def __init__( layer_id=self.layer_idx, quant_config=quant_config, prefix=f"{prefix}.self_attn", + reduce_results=not self.use_hc_sequence_parallel, ) else: raise ValueError(f"Invalid layer_type {layer_type}") self.mlp = Qwen4ExpSparseMoeBlock( - vllm_config=vllm_config, prefix=f"{prefix}.mlp" + vllm_config=vllm_config, + prefix=f"{prefix}.mlp", + reduce_results=not self.use_hc_sequence_parallel, ) hc_config = HyperConnectionConfig( @@ -295,6 +325,9 @@ def forward( else: hidden_states, block_input, injection = attn_hc.mix(hidden_states) + if self.use_hc_sequence_parallel: + # Positions retain the full token axis for both text and MRoPE inputs. + block_input = sp_all_gather(block_input)[: positions.shape[-1]] if self.layer_type == "linear_attention": attn_out = self.linear_attn(hidden_states=block_input) elif self.layer_type in _QSA_LAYER_TYPES: @@ -306,10 +339,25 @@ def forward( raise ValueError("Invalid layer_type") mlp_hc = self.mlp_hyper_connection + if self.use_hc_sequence_parallel: + attn_out = sp_reduce_scatter(attn_out) hidden_states, block_input, injection = mlp_hc.combine_and_mix( hidden_states, attn_out, injection ) - mlp_out = self.mlp(block_input) + if self.use_sequence_parallel: + # MoE consumes local tokens and returns complete local outputs. + mlp_out = self.mlp(block_input, already_sequence_parallel=True) + else: + if self.use_hc_sequence_parallel: + block_input = sp_all_gather(block_input)[: positions.shape[-1]] + mlp_out = self.mlp(block_input) + if self.use_hc_sequence_parallel: + assert isinstance(self.mlp, Qwen4ExpSparseMoeBlock) + reduce_output = self.mlp.experts.moe_config.skip_final_all_reduce + # Only sum partial outputs. Reduced outputs already include all ranks. + mlp_out = ( + sp_reduce_scatter(mlp_out) if reduce_output else sp_shard(mlp_out) + ) return hidden_states, mlp_out, injection @@ -373,6 +421,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.num_redundant_experts = ( vllm_config.parallel_config.eplb_config.num_redundant_experts ) + # MoE SP also enables HC SP, while HC SP can be enabled independently. + self.use_hc_sequence_parallel = is_hc_sequence_parallel_enabled(vllm_config) + self.use_sequence_parallel = ( + vllm_config.parallel_config.use_sequence_parallel_moe + ) self.vocab_size = config.vocab_size self._qsa_layer_ids = frozenset( layer_idx @@ -447,7 +500,6 @@ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: @staticmethod def _start_layer_ple_prefetch( layer: nn.Module, - hidden_states: torch.Tensor, input_ids: torch.Tensor | None, query_start_loc: torch.Tensor | None, ngram_context: torch.Tensor | None, @@ -458,12 +510,7 @@ def _start_layer_ple_prefetch( return if input_ids is None or query_start_loc is None or ngram_context is None: raise RuntimeError("PLE inputs were not prepared") - ple.start_prefetch( - hidden_states, - input_ids, - query_start_loc, - ngram_context, - ) + ple.start_prefetch(input_ids, query_start_loc, ngram_context) def forward( self, @@ -482,19 +529,31 @@ def forward( if input_ids is None: raise ValueError("input_ids or inputs_embeds is required") hidden_states = self.embed_input_ids(input_ids) + if self.use_hc_sequence_parallel: + if ( + self.use_sequence_parallel + and envs.VLLM_MOE_SKIP_PADDING + and is_forward_context_available() + ): + forward_context = get_forward_context() + forward_context.is_padding = sp_padding_mask( + forward_context.is_padding, hidden_states + ) + hidden_states = sp_shard(hidden_states) + # Expand only the local token rows when SP is enabled. hidden_states = hidden_states.repeat(1, self.config.hc_count) else: if intermediate_tensors is None: raise ValueError("pipeline stage requires intermediate tensors") hidden_states = intermediate_tensors["hidden_states"] + full_num_tokens = positions.shape[-1] block_output = None injection = None last_layer = None if self.start_layer < self.end_layer: self._start_layer_ple_prefetch( self.layers[self.start_layer], - hidden_states, input_ids, query_start_loc, ngram_context, @@ -506,7 +565,6 @@ def forward( if layer_idx + 1 < self.end_layer: self._start_layer_ple_prefetch( self.layers[layer_idx + 1], - hidden_states, input_ids, query_start_loc, ngram_context, @@ -526,6 +584,8 @@ def forward( deepstack_embed = deepstack_input_embeds[ f"deepstack_input_embeds_{layer_idx}" ] + if self.use_hc_sequence_parallel: + deepstack_embed = sp_shard(deepstack_embed) deepstack_embed = ( deepstack_embed.unsqueeze(-2) .expand( @@ -560,6 +620,24 @@ def forward( multi_hidden, sample_hidden_states, _ = final_mixer.combine_and_mix( hidden_states, block_output, injection ) + if self.use_hc_sequence_parallel: + if self._mtp_hidden_buffer is not None: + # Gather LM-head and MTP states together when both are needed. + hidden_size = self.config.hidden_size + packed_hidden_states = torch.cat( + [sample_hidden_states, multi_hidden], dim=-1 + ) + packed_hidden_states = sp_all_gather(packed_hidden_states)[ + :full_num_tokens + ] + sample_hidden_states, multi_hidden = packed_hidden_states.split( + [hidden_size, self.config.hc_count * hidden_size], dim=-1 + ) + sample_hidden_states = sample_hidden_states.contiguous() + else: + sample_hidden_states = sp_all_gather(sample_hidden_states)[ + :full_num_tokens + ] if self._mtp_hidden_buffer is not None: # Capture the pre-final-mixer multi-stream hidden state # [T, hc_count*H] for the MTP drafter (zero extra compute: @@ -615,6 +693,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: class Qwen4ExpForCausalLM( nn.Module, HasInnerState, + SupportsHCSP, SupportsLoRA, SupportsMRoPE, SupportsPP, @@ -865,6 +944,7 @@ def get_hf_config(self) -> Qwen4ExpConfig: class Qwen4ExpForConditionalGeneration( Qwen3_5ForConditionalGeneration, HasInnerState, + SupportsHCSP, Qwen4ExpMixtureOfExperts, ): """Qwen3-VL vision tower backed by the Qwen4Exp language model.""" diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index 61290617ae75..32e95faa0dc8 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -20,8 +20,10 @@ from torch import nn from transformers import Qwen4ExpTextConfig +from vllm import envs from vllm.config import VllmConfig, replace, set_current_vllm_config from vllm.distributed import get_pp_group +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.fused_moe.utils import ( is_model_fused_shared_expert_compatible, ) @@ -33,7 +35,7 @@ VocabParallelEmbedding, ) from vllm.model_executor.model_loader.utils import configure_quant_config -from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.interfaces import SupportsHCSP, SupportsPP from vllm.model_executor.models.qwen3_5 import Qwen3_5Model from vllm.model_executor.models.utils import ( AutoWeightsLoader, @@ -44,6 +46,11 @@ maybe_fuse_shared_experts, maybe_prefix, ) +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_padding_mask, + sp_shard, +) from vllm.sequence import IntermediateTensors from .hyperconnection import GatedResidual, HyperConnectionConfig @@ -54,6 +61,7 @@ Qwen4ExpDecoderLayer, Qwen4ExpMixtureOfExperts, Qwen4ExpSparseMoeBlock, + is_hc_sequence_parallel_enabled, ) @@ -171,6 +179,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: config: Qwen4ExpTextConfig = model_config.hf_text_config self.config = config + # MoE SP also enables HC SP, while HC SP can be enabled independently. + self.use_hc_sequence_parallel = is_hc_sequence_parallel_enabled(vllm_config) + self.use_sequence_parallel = ( + vllm_config.parallel_config.use_sequence_parallel_moe + ) self.vocab_size = config.vocab_size self.mtp_start_layer_idx = config.num_hidden_layers @@ -310,6 +323,22 @@ def forward( hidden_states = intermediate_tensors["hidden_states"] current_step_idx = spec_step_idx % self.num_mtp_layers + full_num_tokens = positions.shape[-1] + if self.use_hc_sequence_parallel: + # SP requires PP=1, so the first-rank branch sets prev_block_output. + assert prev_block_output is not None + if ( + self.use_sequence_parallel + and envs.VLLM_MOE_SKIP_PADDING + and is_forward_context_available() + ): + forward_context = get_forward_context() + forward_context.is_padding = sp_padding_mask( + forward_context.is_padding, hidden_states + ) + hidden_states = sp_shard(hidden_states) + # The FC output is already complete and needs only a slice. + prev_block_output = sp_shard(prev_block_output) layer = self.layers[current_step_idx] hidden_states, block_output, injection = layer( hidden_states=hidden_states, @@ -337,6 +366,17 @@ def forward( hidden_states, block_output, injection ) ) + if self.use_hc_sequence_parallel: + # Gather both outputs together, then restore their contiguous layouts. + packed_hidden_states = torch.cat( + [sample_hidden_states, multi_hidden], dim=-1 + ) + packed_hidden_states = sp_all_gather(packed_hidden_states)[:full_num_tokens] + sample_hidden_states, multi_hidden = packed_hidden_states.split( + [hidden_size, hc_count * hidden_size], dim=-1 + ) + sample_hidden_states = sample_hidden_states.contiguous() + multi_hidden = multi_hidden.contiguous() return sample_hidden_states, multi_hidden def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -357,7 +397,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=mapper) -class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): +class Qwen4ExpMTP(nn.Module, SupportsHCSP, SupportsPP, Qwen4ExpMixtureOfExperts): packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"], diff --git a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py index cf03cbb40d47..794eea87a620 100644 --- a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py +++ b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py @@ -9,11 +9,14 @@ from transformers import Qwen4ExpTextConfig from vllm.config import get_current_vllm_config +from vllm.distributed import get_etp_group, get_tp_group from vllm.logger import init_logger from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, ) +from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8 from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter, sp_shard from ..common.ngram_embedding import ( Qwen4ExpPLEDeviceEmbedding, @@ -146,8 +149,10 @@ def __init__( prefix: str, quant_config: QuantizationConfig | None = None, params_dtype: torch.dtype | None = None, + use_sequence_parallel: bool = False, ) -> None: super().__init__() + self.use_sequence_parallel = use_sequence_parallel self.embedding_dim = embedding_dim self.ngram_size = int(config.ngram_size) self.heads_per_ngram = int(config.heads_per_ngram) @@ -211,6 +216,11 @@ def __init__( if engram_config is not None and engram_config.cpu_offload else Qwen4ExpPLEDeviceEmbedding ) + # A TP reduce-scatter covers all vocabulary owners only when ETP=TP. + self.use_reduce_scatter = ( + self.use_sequence_parallel + and get_etp_group().world_size == get_tp_group().world_size + ) self.ngram_embedding = embedding_cls( padded_vocab_size, self.head_dim, @@ -221,6 +231,7 @@ def __init__( num_ngram_heads=self.ngram_heads, max_total_tokens=max_total_tokens, data_parallel_rank=data_parallel_rank, + reduce_results=not self.use_reduce_scatter, ) if self.ngram_embedding.supports_prefetch: # The side-stream lookup outlives eager-break args, whose @@ -348,6 +359,18 @@ def compute_ngram_ids( id_blocks.append(ids[request_indices, adjusted_columns]) return torch.cat(id_blocks, dim=-1) + def _sp_shard_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor: + """Shard embedding token rows, optionally reducing vocabulary contributions.""" + # Each embedding element has one vocabulary owner, including FP8 bytes. + output_dtype = embeddings.dtype + if is_fp8(embeddings): + embeddings = embeddings.view(torch.int8) + if self.use_reduce_scatter: + embeddings = sp_reduce_scatter(embeddings) + else: + embeddings = sp_shard(embeddings) + return embeddings.view(output_dtype) + def forward( self, hidden_states: torch.Tensor, @@ -355,15 +378,22 @@ def forward( query_start_loc: torch.Tensor, ngram_context: torch.Tensor, ) -> torch.Tensor: + """Return full embeddings or a token shard after combining vocabulary owners.""" embedding = self.ngram_embedding if embedding.supports_prefetch: - return embedding(hidden_states) - ngram_ids = self.compute_ngram_ids(input_ids, query_start_loc, ngram_context) - return self.ngram_embedding(ngram_ids).flatten(-2) + # Prefetch contains full tokens even when hidden states are sharded. + embeddings = embedding.finalize_prefetch(total_tokens=input_ids.shape[0]) + else: + ngram_ids = self.compute_ngram_ids( + input_ids, query_start_loc, ngram_context + ) + embeddings = embedding(ngram_ids).flatten(-2) + if not self.use_sequence_parallel: + return embeddings + return self._sp_shard_embeddings(embeddings) def start_prefetch( self, - hidden_states: torch.Tensor, input_ids: torch.Tensor, query_start_loc: torch.Tensor, ngram_context: torch.Tensor, @@ -378,7 +408,7 @@ def start_prefetch( ngram_context, output=self._prefetch_ids[: input_ids.numel()], ) - embedding.start_prefetch(hidden_states, ngram_ids) + embedding.start_prefetch(ngram_ids) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load hash buffers and checkpoint-split embedding rows.""" @@ -444,8 +474,3 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if regular_weights: loaded.update(AutoWeightsLoader(self).load_weights(regular_weights)) return loaded - - -__all__ = [ - "Qwen4ExpNGramEmbedding", -] diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index bdb16b3adef0..a595dba18d47 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -18,6 +18,7 @@ MambaStateShapeCalculator, is_conv_state_dim_first, ) +from vllm.models.common.ops.sequence_parallel import sp_all_gather, sp_shard from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.attention.backends.short_conv_attn import ( PleShortConvAttentionBackend, @@ -69,8 +70,10 @@ def __init__( layer_idx: int = 0, ple_dense_layer_id: int | None = None, prefix: str = "", + use_sequence_parallel: bool = False, ) -> None: super().__init__() + self.use_sequence_parallel = use_sequence_parallel model_config = vllm_config.model_config cache_config = vllm_config.cache_config quant_config = vllm_config.quant_config @@ -100,8 +103,9 @@ def __init__( prefix=f"{prefix}.ple_embedding", quant_config=quant_config, params_dtype=model_config.dtype, + use_sequence_parallel=self.use_sequence_parallel, ) - # The PLE cache is TP-replicated, so this merged projection is too. + # Projection weights stay replicated even when their token rows are sharded. self.kv_proj = MergedColumnParallelLinear( int(config.ple_embed_dim), [self.hc_hidden_size, self.hidden_size], @@ -151,18 +155,12 @@ def _dequantize_embeddings( def start_prefetch( self, - hidden_states: torch.Tensor, input_ids: torch.Tensor, query_start_loc: torch.Tensor, ngram_context: torch.Tensor, ) -> None: """Start the pinned PLE lookup while the preceding decoder layer runs.""" - self.ple_embedding.start_prefetch( - hidden_states, - input_ids, - query_start_loc, - ngram_context, - ) + self.ple_embedding.start_prefetch(input_ids, query_start_loc, ngram_context) @property def mamba_type(self) -> MambaAttentionBackendEnum: @@ -192,7 +190,7 @@ def _short_conv_dilated_dispatch( self, inputs: torch.Tensor, residual: torch.Tensor, - outer_residual: torch.Tensor, + outer_residual: torch.Tensor | None, metadata: PleShortConvAttentionMetadata, conv_state: torch.Tensor, conv_weights: torch.Tensor, @@ -207,7 +205,8 @@ def _short_conv_dilated_dispatch( has_non_spec = has_prefill or has_decode inputs = inputs[: metadata.num_actual_tokens] residual = residual[: metadata.num_actual_tokens] - outer_residual = outer_residual[: metadata.num_actual_tokens] + if outer_residual is not None: + outer_residual = outer_residual[: metadata.num_actual_tokens] spec_token_indices = None non_spec_token_indices = None @@ -259,11 +258,14 @@ def _short_conv_dilated_dispatch( residual_d, residual_p = torch.split( residual, [num_decode_tokens, num_prefill_tokens], dim=0 ) - outer_residual_d, outer_residual_p = torch.split( - outer_residual, - [num_decode_tokens, num_prefill_tokens], - dim=0, - ) + outer_residual_d: torch.Tensor | None = None + outer_residual_p: torch.Tensor | None = None + if outer_residual is not None: + outer_residual_d, outer_residual_p = torch.split( + outer_residual, + [num_decode_tokens, num_prefill_tokens], + dim=0, + ) token_indices_d = None token_indices_p = None else: @@ -336,14 +338,16 @@ def _short_conv( self, inputs: torch.Tensor, residual: torch.Tensor, - outer_residual: torch.Tensor, + outer_residual: torch.Tensor | None, ) -> None: + """Update convolution state, optionally fusing both residual additions.""" forward_context = get_forward_context() attn_metadata = forward_context.attn_metadata # Profiling omits all metadata or this Mamba entry. Short convolution # is a no-op there, but preserve the outer residual addition. if attn_metadata is None: - residual.add_(outer_residual) + if outer_residual is not None: + residual.add_(outer_residual) return if not isinstance(attn_metadata, dict): @@ -354,7 +358,8 @@ def _short_conv( layer_attn_metadata = attn_metadata.get(self.prefix) if layer_attn_metadata is None: - residual.add_(outer_residual) + if outer_residual is not None: + residual.add_(outer_residual) return if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata): raise TypeError( @@ -396,7 +401,10 @@ def forward( ngram_context: torch.Tensor, ) -> torch.Tensor: input_ids = input_ids.reshape(-1) - if input_ids.shape[0] != hidden_states.shape[0]: + if ( + not self.use_sequence_parallel + and input_ids.shape[0] != hidden_states.shape[0] + ): raise ValueError( "PLE expects input_ids and hidden_states to have the same " f"token length, got {input_ids.shape[0]} and " @@ -408,6 +416,7 @@ def forward( query_start_loc, ngram_context, ) + assert embeddings.shape[0] == hidden_states.shape[0] embeddings = self._dequantize_embeddings(embeddings, hidden_states.dtype) kv, _ = self.kv_proj(embeddings) key, value = kv.split(self.kv_proj.output_sizes, dim=-1) @@ -420,9 +429,16 @@ def forward( self.norm_conv.weight, self.norm_key.eps, ) - # State routing depends on runtime request metadata and remains outside - # the graph; short convolution accumulates into gated_output. - self._short_conv(conv_input, gated_output, hidden_states) + if self.use_sequence_parallel: + # Only convolution needs full tokens. Add residuals on local shards. + conv_input = sp_all_gather(conv_input)[: input_ids.shape[0]] + # Metadata may omit padding tokens or all tokens during profiling. + conv_output = torch.zeros_like(conv_input) + self._short_conv(conv_input, conv_output, outer_residual=None) + gated_output.add_(sp_shard(conv_output)) + gated_output.add_(hidden_states) + else: + self._short_conv(conv_input, gated_output, hidden_states) return gated_output From 7900d2eabc4f9a7facb7a196a04c2c32803a79c8 Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Mon, 12 Oct 2026 08:18:44 -0500 Subject: [PATCH 299/624] [Test] Accept tied tokens at the k-th value in the Triton top-k test (#60726) Signed-off-by: Rohan138 Co-authored-by: Claude Opus 5.5 --- tests/v1/sample/test_topk_topp_sampler.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index c0b979245659..190a4e3fcaf0 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -251,8 +251,17 @@ def _compare_results( triton_kept = (result_triton != float("-inf")).sum(dim=-1) if p is None: - # Top-k only: expect exact match - assert torch.equal(pytorch_kept, triton_kept), ( + # Top-k only: Triton keeps exactly k (index-ordered tie-break at + # the k-th value) while the sort-based reference keeps all ties, + # so PyTorch may only keep extra tokens tied with the k-th value. + pytorch_mask = result_pytorch != float("-inf") + triton_mask = result_triton != float("-inf") + kth = torch.where(triton_mask, logits.float(), float("inf")) + kth = kth.amin(dim=-1, keepdim=True).expand_as(logits) + extra = pytorch_mask & ~triton_mask + assert not (triton_mask & ~pytorch_mask).any() and torch.equal( + logits.float()[extra], kth[extra] + ), ( f"Top-k mask mismatch: PyTorch kept {pytorch_kept.tolist()}, " f"Triton kept {triton_kept.tolist()}" ) From 8c175f403acd2749086c298cd1c6f29ed9b4e2ba Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Mon, 12 Oct 2026 15:20:35 +0200 Subject: [PATCH 300/624] [CI][ROCm] Compare greedy output only in the AITER attn+quant test (#61185) Signed-off-by: Stefan Koncarevic --- tests/compile/correctness_e2e/test_attn_quant.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/compile/correctness_e2e/test_attn_quant.py b/tests/compile/correctness_e2e/test_attn_quant.py index e13a8f264a6a..1d8f73da4b9a 100644 --- a/tests/compile/correctness_e2e/test_attn_quant.py +++ b/tests/compile/correctness_e2e/test_attn_quant.py @@ -77,6 +77,10 @@ def settings(enabled): env1=env, env2=env, method="generate", + # Fusing output quant drops a bf16 round-trip, so the two settings + # agree on the argmax but not on the exact distribution. Greedy + # decoding survives that, sampling at temperature 1.0 does not. + include_seeded_sampling=False, ) matches = [ int(count) for count in re.findall(r"'attn_quant_fusion': (\d+)", logs.text) From ac28877f0e86452f2e151279bdc373cd58ae88f0 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Mon, 12 Oct 2026 21:21:35 +0800 Subject: [PATCH 301/624] [Bugfix] Fix out-of-bounds lookahead in adaptive block verification (#60945) Signed-off-by: chaunceyjiang --- vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index b2efa6432e9f..4387b1d38ab8 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -425,6 +425,9 @@ def _compute_local_residual_mass_kernel( PADDED_VOCAB_NUM_BLOCKS: tl.constexpr, ): logit_idx = tl.program_id(0).to(tl.int64) + # Adaptive verification can end before num_speculative_steps. + if logit_idx + 1 >= tl.num_programs(0): + return draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx) if draft_step_idx == 0 or draft_step_idx >= num_speculative_steps: # The acceptance threshold, h, looks one position ahead and sums From 314ab8cffd9b115198d50491c5f1c4a2112ad4a1 Mon Sep 17 00:00:00 2001 From: DimensionSTP <65501090+DimensionSTP@users.noreply.github.com> Date: Mon, 12 Oct 2026 22:25:36 +0900 Subject: [PATCH 302/624] [Bugfix][Structured Output] Drain compilation tasks before releasing the backend (#61047) Signed-off-by: DimensionSTP <65501090+DimensionSTP@users.noreply.github.com> Co-authored-by: Artem Perevedentsev --- .../test_structured_output_manager.py | 85 ++++++++++++++++++- vllm/v1/structured_output/__init__.py | 9 ++ 2 files changed, 93 insertions(+), 1 deletion(-) diff --git a/tests/v1/structured_output/test_structured_output_manager.py b/tests/v1/structured_output/test_structured_output_manager.py index 67c15825580c..1d1e7b78df04 100644 --- a/tests/v1/structured_output/test_structured_output_manager.py +++ b/tests/v1/structured_output/test_structured_output_manager.py @@ -2,12 +2,19 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import time +from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass +from threading import Event import pytest from transformers import AutoTokenizer -from vllm.config import DeviceConfig, StructuredOutputsConfig, VllmConfig +from vllm.config import ( + DeviceConfig, + SchedulerConfig, + StructuredOutputsConfig, + VllmConfig, +) from vllm.config.model import ModelConfig from vllm.config.speculative import SpeculativeConfig from vllm.config.structured_outputs import StructuredOutputsBackend @@ -576,3 +583,79 @@ def test_outlines_choice_with_non_bmp_characters(tokenizer): grammar = backend.compile_grammar(request_type, grammar_spec) assert grammar.accept_tokens("", tokenizer.encode(choices[0])) + + +def test_clear_backend_waits_for_running_grammar_and_cancels_pending( + tokenizer, monkeypatch +): + """Keep the compiler alive until running work finishes during shutdown.""" + manager, request = _build_harness(tokenizer, "xgrammar", use_reasoner=False) + manager.executor.shutdown(wait=True) + manager.executor = ThreadPoolExecutor(max_workers=1) + started, release = Event(), Event() + assert manager.backend is not None + compile_grammar = manager.backend.compile_grammar + + def delayed_compile(*args, **kwargs): + started.set() + assert release.wait(MAX_WAIT_SECONDS) + return compile_grammar(*args, **kwargs) + + monkeypatch.setattr(manager.backend, "compile_grammar", delayed_compile) + try: + manager.grammar_init(request) + assert request.structured_output_request is not None + running = request.structured_output_request._grammar + assert isinstance(running, Future) + assert started.wait(MAX_WAIT_SECONDS) + pending_request = Request( + request_id="pending-at-shutdown", + prompt_token_ids=[], + sampling_params=request.sampling_params, + pooling_params=None, + ) + manager.grammar_init(pending_request) + assert pending_request.structured_output_request is not None + pending = pending_request.structured_output_request._grammar + assert isinstance(pending, Future) + # Cancellation releases the running task without timing-based sleeps. + pending.add_done_callback(lambda _: release.set()) + manager.clear_backend() + release.set() + assert running.result(timeout=MAX_WAIT_SECONDS) is not None + assert pending.cancelled() + assert manager.backend is None + manager.clear_backend() + finally: + release.set() + manager.executor.shutdown(wait=True, cancel_futures=True) + + +@pytest.mark.parametrize("skip_tokenizer_init", [False, True]) +@pytest.mark.parametrize("max_num_seqs", [1, 256]) +def test_clear_backend_without_grammar(skip_tokenizer_init, max_num_seqs): + """Stop initialized thread pools even if no request needed a backend.""" + manager = StructuredOutputManager( + VllmConfig( + model_config=ModelConfig( + tokenizer=TOKENIZER, skip_tokenizer_init=skip_tokenizer_init + ), + scheduler_config=SchedulerConfig( + max_num_seqs=max_num_seqs, + max_model_len=128, + is_encoder_decoder=False, + ), + device_config=DeviceConfig(device="cpu"), + ) + ) + manager.clear_backend() + manager.clear_backend() + assert manager.backend is None + executors = [] + if not skip_tokenizer_init: + executors.append(manager.executor) + if max_num_seqs > manager.fill_bitmask_parallel_threshold: + executors.append(manager.executor_for_fillmask) + for executor in executors: + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + executor.submit(int, 0) diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 2163a800a563..3f461ed47f75 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -482,5 +482,14 @@ def accept_tokens(self, request: "Request", new_token_ids: list[int]) -> bool: return grammar.accept_tokens(request.request_id, new_token_ids[grammar_start:]) def clear_backend(self) -> None: + # Finish tasks that can still use the backend before releasing it. + if not self.vllm_config.model_config.skip_tokenizer_init: + self.executor.shutdown(wait=True, cancel_futures=True) + if ( + self.fill_bitmask_parallel_threshold + < self.vllm_config.scheduler_config.max_num_seqs + ): + self.executor_for_fillmask.shutdown(wait=True, cancel_futures=True) if self.backend is not None: self.backend.destroy() + self.backend = None From 092c35dd155db2644da6bb671332c51902bd3039 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Mon, 12 Oct 2026 21:27:34 +0800 Subject: [PATCH 303/624] [Model] Support Cloudflare/clef-omni on /v1/systemone (#61064) Signed-off-by: zjy0516 Signed-off-by: Isotr0py Co-authored-by: Claude Opus 5.5 Co-authored-by: Isotr0py --- .../online_serving/structured_decisions.md | 27 ++ .../structured_decisions/test_strategies.py | 345 +++++++++++++- .../test_filter_duplicate_safetensors.py | 29 ++ .../multimodal/pooling/test_clef_omni.py | 168 +++++++ .../multimodal/processing/test_qwen3_omni.py | 30 ++ tests/models/registry.py | 4 + tests/multimodal/test_video.py | 16 + vllm/entrypoints/generate/api_router.py | 39 -- .../structured_decisions/api_router.py | 39 ++ .../generate/structured_decisions/protocol.py | 9 + .../generate/structured_decisions/serving.py | 23 +- .../structured_decisions/strategies.py | 362 ++++++++++++++- .../launchers/api_server/app_state.py | 13 + .../launchers/api_server/routers.py | 6 + .../model_loader/default_loader.py | 8 +- vllm/model_executor/models/clef_omni.py | 437 ++++++++++++++++++ .../models/qwen3_omni_moe_thinker.py | 11 +- vllm/model_executor/models/registry.py | 1 + vllm/multimodal/video.py | 64 +++ .../model_arch_config_convertor.py | 30 ++ 20 files changed, 1592 insertions(+), 69 deletions(-) create mode 100644 tests/models/multimodal/pooling/test_clef_omni.py create mode 100644 vllm/model_executor/models/clef_omni.py diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md index a9ed5a2ee5aa..4d8d65d89e94 100644 --- a/docs/serving/online_serving/structured_decisions.md +++ b/docs/serving/online_serving/structured_decisions.md @@ -130,6 +130,32 @@ the request lists them, a `noul` labels its options `yes` and `no`, and a | options per `choice` | 26 | | levels per `score` | 10 | +## Clef-Omni + +[Clef-Omni](https://huggingface.co/Cloudflare/clef-omni) is a decision model: +a joint schema head on the Qwen3-Omni thinker scores every option of every +question in one forward pass. vLLM serves it as a pooling model. + +```bash +vllm serve Cloudflare/clef-omni +``` + +- A request is one read. The prompt is the one the checkpoint's + `joint_schema_model.py` builds: the state, then every question with its + options. A JSON state is written as compact JSON with sorted keys. +- A state too long for the model is truncated to fit, as the reference does. + A request whose questions alone exceed 64,000 tokens, or `max_model_len`, + returns 400. +- `instructions` returns 400, since the prompt is fixed. `chat_template_kwargs` + has no effect. +- There are no labels, so `label_mass` is 1 and `confidence` is the chosen + option's probability. `usage.input_tokens` counts the prompt once, and + `output_tokens` is 0. +- `images`, `audio` and `videos` are read as the reference does: placed + before the state, videos sampled at 2 fps, and a video's soundtrack heard + with it when every video has one. Audio clips cannot be sent together with + videos that have soundtracks. + ## Request fields | field | meaning | @@ -139,6 +165,7 @@ the request lists them, a `noul` labels its options `yes` and `no`, and a | `instructions` | optional context placed ahead of the questions | | `chat_template_kwargs` | passed to the chat template, for example `{"enable_thinking": false}` | | `cache_salt` | optional [prefix-cache salt](../../usage/security.md#cache-salting) applied to every question | +| `images`, `audio`, `videos` | optional media URLs or data URLs, for models that take them ([Clef-Omni](#clef-omni)) | A question with any field other than `type`, `instructions` and `criteria` is rejected with a 400. diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index 41b9a8f01271..8d38e61e6248 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -1,12 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json +from io import BytesIO from types import SimpleNamespace from typing import Any from unittest.mock import MagicMock +import numpy as np +import pybase64 as base64 import pytest +import torch from fastapi import FastAPI, Request from fastapi.testclient import TestClient +from PIL import Image from vllm.config import ModelConfig from vllm.engine.protocol import EngineClient @@ -23,22 +29,32 @@ ) from vllm.entrypoints.generate.structured_decisions.serving import ( ServingStructuredDecisions, - state_text, + parse_questions, ) from vllm.entrypoints.generate.structured_decisions.strategies import ( DiffusionGemmaCanvasStrategy, + JointSchemaStrategy, NextTokenStrategy, ReadContext, reply_label_ids, select_read_strategy, + state_text, ) from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.entrypoints.serve.exception_handling.register import ( init_exception_handler, ) +from vllm.inputs import mm_input from vllm.logprobs import Logprob -from vllm.outputs import CompletionOutput, RequestOutput +from vllm.multimodal.media import MediaConnector, MediaWithBytes +from vllm.outputs import ( + CompletionOutput, + PoolingOutput, + PoolingRequestOutput, + RequestOutput, +) from vllm.renderers.hf import HfRenderer from vllm.renderers.online_renderer import OnlineRenderer @@ -53,17 +69,26 @@ def model(architecture: str, logprobs_mode: str = "raw_logprobs") -> Any: def test_strategy_selection(): - qwen = "Qwen3ForCausalLM" - assert select_read_strategy(model(qwen)) is NextTokenStrategy - assert select_read_strategy(model(qwen, "processed_logprobs")) is NextTokenStrategy + qwen, generate = "Qwen3ForCausalLM", ("generate",) + assert select_read_strategy(model(qwen), generate) is NextTokenStrategy assert ( - select_read_strategy(model("DiffusionGemmaForBlockDiffusion")) + select_read_strategy(model(qwen, "processed_logprobs"), generate) + is NextTokenStrategy + ) + assert ( + select_read_strategy(model("DiffusionGemmaForBlockDiffusion"), generate) is DiffusionGemmaCanvasStrategy ) with pytest.raises(ValueError, match="does not support LlamaForCausalLM"): - select_read_strategy(model("LlamaForCausalLM")) + select_read_strategy(model("LlamaForCausalLM"), generate) with pytest.raises(ValueError, match="not raw_logits"): - select_read_strategy(model(qwen, "raw_logits")) + select_read_strategy(model(qwen, "raw_logits"), generate) + # A generative architecture served as a pooling model has no read strategy. + with pytest.raises(ValueError, match="need the generate task"): + select_read_strategy(model(qwen), ("embed",)) + # Clef-Omni scores options with its head, so it reads no logprobs. + clef = model("ClefOmniForJointSchema", "raw_logits") + assert select_read_strategy(clef, ("token_classify",)) is JointSchemaStrategy def test_route_is_registered_by_default(): @@ -126,7 +151,7 @@ def test_labels_start_the_reply(qwen): @pytest.fixture -def serving_and_engine(qwen): +def engine_and_renderer(qwen): tokenizer, _ = qwen model_config = ModelConfig(model="Qwen/Qwen3-0.6B", dtype="float32") renderer = HfRenderer( @@ -148,6 +173,19 @@ def serving_and_engine(qwen): chat_template=None, chat_template_content_format="auto", ) + yield engine, online_renderer + renderer.shutdown() + + +def serving(engine, strategy) -> ServingStructuredDecisions: + name = engine.model_config.model + models = OpenAIServingModels(engine, [BaseModelPath(name=name, model_path=name)]) + return ServingStructuredDecisions(models, strategy) + + +@pytest.fixture +def serving_and_engine(engine_and_renderer): + engine, online_renderer = engine_and_renderer async def generate(prompt, params, request_id, **kwargs): label_ids = params.logprob_token_ids @@ -171,11 +209,7 @@ async def generate(prompt, params, request_id, **kwargs): engine.generate.side_effect = generate strategy = NextTokenStrategy(ReadContext(engine, online_renderer, None, "auto", {})) - models = OpenAIServingModels( - engine, [BaseModelPath(name=model_config.model, model_path=model_config.model)] - ) - yield ServingStructuredDecisions(models, strategy), engine - renderer.shutdown() + return serving(engine, strategy), engine @pytest.mark.asyncio @@ -256,3 +290,286 @@ def test_canvas_read(): strategy._read_input( {"type": "token", "prompt_token_ids": [1, 2, 3, 4]}, [1, 2, 3, 4] ) + + +JOINT_SCHEMA_PROMPT = """<|im_start|>system +Read the complete state and schema. Decide every field jointly. Each answer \ +must be exactly one of that field's allowed options.<|im_end|> +<|im_start|>user +STATE: +{"customer":{"tier":"gold"},"ticket":"Charged twice"} + +SCHEMA FIELDS: + +FIELD 1 +ID: team +TYPE: choice +INSTRUCTION: Which team? +ALLOWED OPTIONS: +OPTION 1: {"description":"payments","option_id":"billing"} +OPTION 2: {"option_id":"security"} +OPTION 3: {"description":"deliveries","option_id":"shipping"} +END FIELD + +FIELD 2 +ID: urgent +TYPE: noul +INSTRUCTION: urgent +ALLOWED OPTIONS: +OPTION 1: {"description":"money is moving","option_id":"true"} +OPTION 2: {"description":"The proposition is false or the answer is no.",\ +"option_id":"false"} +END FIELD + +FIELD 3 +ID: tone +TYPE: score +INSTRUCTION: How upset? +ALLOWED OPTIONS: +OPTION 1: {"description":"calm","option_id":"0"} +OPTION 2: {"description":"annoyed","option_id":"1"} +OPTION 3: {"description":"furious","option_id":"2"} +END FIELD + +<|im_end|> +<|im_start|>assistant + + + + +JOINT SCHEMA DECISIONS:""" + +JOINT_SCHEMA_REQUEST: dict[str, Any] = dict( + state={"ticket": "Charged twice", "customer": {"tier": "gold"}}, + questions={ + "team": { + "type": "choice", + "instructions": "Which team?", + "criteria": { + "shipping": "deliveries", + "billing": "payments", + "security": None, + }, + }, + "urgent": {"type": "noul", "criteria": {"true": "money is moving"}}, + "tone": { + "type": "score", + "instructions": "How upset?", + "criteria": ["calm", "annoyed", "furious"], + }, + }, +) + + +@pytest.fixture +def joint_schema(qwen, engine_and_renderer): + tokenizer, _ = qwen + engine, online_renderer = engine_and_renderer + # The probability the head gives each option, by option id. + scores = {"billing": 0.7, "security": 0.1, "shipping": 0.2} + scores |= {"true": 0.9, "false": 0.1, "0": 0.2, "1": 0.5, "2": 0.3} + + async def encode(prompt, params, request_id, **kwargs): + ids = prompt["prompt_token_ids"] + spans = params.extra_kwargs["joint_schema"]["option_spans"] + options = [json.loads(tokenizer.decode(ids[s:e])) for s, e in spans] + data = torch.tensor([scores[o["option_id"]] for o in options]) + yield PoolingRequestOutput(request_id, PoolingOutput(data), ids, 0, True) + + engine.encode.side_effect = encode + strategy = JointSchemaStrategy( + ReadContext(engine, online_renderer, None, "auto", {}) + ) + return serving(engine, strategy), engine, tokenizer + + +@pytest.mark.asyncio +async def test_joint_schema_reads_every_question_at_once(joint_schema): + serving, engine, tokenizer = joint_schema + response = await serving.create_decision( + StructuredDecisionRequest(**JOINT_SCHEMA_REQUEST, cache_salt="salt") + ) + assert isinstance(response, StructuredDecisionResponse) + + engine.encode.assert_called_once() + engine.check_admission.assert_called_once_with(1) + prompt, params = engine.encode.call_args.args[:2] + ids = prompt["prompt_token_ids"] + assert tokenizer.decode(ids) == JOINT_SCHEMA_PROMPT + assert prompt["cache_salt"] == "salt" + assert params.task == "token_classify" + schema = params.extra_kwargs["joint_schema"] + assert schema["question_types"] == [1, 0, 2] + assert schema["option_counts"] == [3, 2, 3] + assert [tokenizer.decode(ids[s:e]) for s, e in schema["question_spans"]] == [ + "Which team?", + "urgent", + "How upset?", + ] + + # Options are read in the model's order and answered in the request's. + team = response.answers["team"] + assert team["choice"] == "billing" + assert list(team["probabilities"]) == ["shipping", "billing", "security"] + assert team["probabilities"] == pytest.approx( + {"shipping": 0.2, "billing": 0.7, "security": 0.1} + ) + assert response.answers["urgent"]["noul"] == pytest.approx(0.9) + assert response.answers["tone"]["score"] == pytest.approx(1.1) + assert response.usage.input_tokens == len(ids) + assert response.usage.output_tokens == 0 + assert all(d.label_mass == 1.0 for d in response.diagnostics.values()) + + +@pytest.mark.asyncio +async def test_joint_schema_rejects_instructions(joint_schema): + serving, engine, _ = joint_schema + response = await serving.create_decision( + StructuredDecisionRequest(**JOINT_SCHEMA_REQUEST, instructions="Be strict.") + ) + assert isinstance(response, ErrorResponse) + assert response.error.code == 400 + engine.encode.assert_not_called() + + +def test_joint_schema_truncates_the_state(joint_schema): + serving, _, _ = joint_schema + strategy = serving.strategy + request = StructuredDecisionRequest(**JOINT_SCHEMA_REQUEST) + questions = parse_questions(request, strategy.limits()) + prompt, schema, _ = strategy.encode(questions, "word " * 100) + strategy.max_length = len(prompt) - 50 + truncated, truncated_schema, _ = strategy.encode(questions, "word " * 100) + # The state loses 50 tokens and the questions stay at the end. + assert len(truncated) == strategy.max_length + assert truncated[-200:] == prompt[-200:] + assert truncated_schema == schema + strategy.max_length = 100 + with pytest.raises(StructuredDecisionError, match="questions take"): + strategy.encode(questions, "word") + + +@pytest.mark.asyncio +async def test_media_needs_a_model_that_takes_it(serving_and_engine): + serving, engine = serving_and_engine + request = StructuredDecisionRequest( + **JOINT_SCHEMA_REQUEST, images=["data:image/png;base64,AAAA"] + ) + response = await serving.create_decision(request) + assert isinstance(response, ErrorResponse) + assert response.error.code == 400 + assert "does not take image inputs" in response.error.message + engine.generate.assert_not_called() + + +@pytest.mark.asyncio +async def test_joint_schema_reads_media_before_the_state(joint_schema, monkeypatch): + serving, engine, tokenizer = joint_schema + strategy = serving.strategy + strategy.__dict__["_connector"] = MediaConnector() + image_pad = tokenizer.convert_tokens_to_ids("<|image_pad|>") + rendered: dict[str, Any] = {} + + async def process_for_engine_async(prompt, arrival_time): + rendered.update(prompt) + # The image's placeholder expands to three tokens. + ids = [ + token + for pad in prompt["prompt_token_ids"] + for token in ([pad] * 3 if pad == image_pad else [pad]) + ] + return mm_input(ids, {}, {}, {}, cache_salt=prompt.get("cache_salt")) + + renderer = strategy.context.online_renderer.renderer + monkeypatch.setattr(renderer, "process_for_engine_async", process_for_engine_async) + buffer = BytesIO() + Image.new("RGB", (8, 8), "red").save(buffer, format="PNG") + image = "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode() + response = await serving.create_decision( + StructuredDecisionRequest(**JOINT_SCHEMA_REQUEST, images=[image]) + ) + assert isinstance(response, StructuredDecisionResponse) + + placeholder = tokenizer.encode( + "<|vision_start|><|image_pad|><|vision_end|>\n", add_special_tokens=False + ) + assert rendered["prompt_token_ids"] == strategy.prefix + placeholder + assert rendered["multi_modal_data"]["image"][0].size == (8, 8) + assert rendered["mm_processor_kwargs"] == { + "images_kwargs": {"size": strategy.IMAGE_SIZE} + } + prompt = engine.encode.call_args.args[0] + assert prompt["type"] == "multimodal" + ids = prompt["prompt_token_ids"] + expanded = len(strategy.prefix) + len(placeholder) + 2 + assert ids[:expanded].count(image_pad) == 3 + state_and_questions = JOINT_SCHEMA_PROMPT.split("STATE:\n", 1)[1] + assert tokenizer.decode(ids[expanded:]) == state_and_questions + assert response.answers["team"]["choice"] == "billing" + assert response.usage.input_tokens == len(ids) + + +@pytest.mark.asyncio +async def test_joint_schema_hears_video_soundtracks(joint_schema, monkeypatch): + serving, engine, tokenizer = joint_schema + strategy = serving.strategy + video = MediaWithBytes( + (np.zeros((4, 8, 8, 3), np.uint8), {"fps": 2.0}), b"video-bytes" + ) + track = (np.zeros(160, np.float32), 16000) + rendered: dict[str, Any] = {} + + class FakeConnector: + async def fetch_image_async(self, url): + raise AssertionError("no images in this test") + + async def load_from_url_async(self, url, media_io, *, fetch_timeout=None): + return video + + async def fetch_audio_async(self, url): + if url == "https://example.org/silent.mp4": + raise ValueError("No audio stream found.") + return track + + strategy.__dict__["_connector"] = FakeConnector() + + async def process_for_engine_async(prompt, arrival_time): + rendered.clear() + rendered.update(prompt) + return mm_input(prompt["prompt_token_ids"], {}, {}, {}) + + renderer = strategy.context.online_renderer.renderer + monkeypatch.setattr(renderer, "process_for_engine_async", process_for_engine_async) + + response = await serving.create_decision( + StructuredDecisionRequest( + **JOINT_SCHEMA_REQUEST, videos=["https://example.org/talking.mp4"] + ) + ) + assert isinstance(response, StructuredDecisionResponse) + assert rendered["multi_modal_data"] == {"video": [video], "audio": [track]} + assert rendered["mm_processor_kwargs"] == { + "images_kwargs": {"size": strategy.IMAGE_SIZE}, + "use_audio_in_video": True, + } + + # A video without a decodable soundtrack is watched silently. + response = await serving.create_decision( + StructuredDecisionRequest( + **JOINT_SCHEMA_REQUEST, videos=["https://example.org/silent.mp4"] + ) + ) + assert isinstance(response, StructuredDecisionResponse) + assert rendered["multi_modal_data"] == {"video": [video]} + assert rendered["mm_processor_kwargs"]["use_audio_in_video"] is False + + # Audio clips cannot be read with videos that have soundtracks. + response = await serving.create_decision( + StructuredDecisionRequest( + **JOINT_SCHEMA_REQUEST, + videos=["https://example.org/talking.mp4"], + audio=["https://example.org/clip.wav"], + ) + ) + assert isinstance(response, ErrorResponse) + assert response.error.code == 400 diff --git a/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py b/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py index da04d0108fdf..7aca4b64ab62 100644 --- a/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py +++ b/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py @@ -137,3 +137,32 @@ def is_unused_checkpoint_weight(self, name: str) -> bool: ["model-00002-of-00003.safetensors", "model-00003-of-00003.safetensors"], sorted(f.name for f in mtp_checkpoint.glob("*.safetensors")), ] + + +def test_secondary_source_with_explicit_patterns_skips_backbone_index(tmp_path): + """A source with explicit allow patterns (e.g. a standalone head file in + the repo root) must not be validated against the backbone's index.""" + save_file( + {"backbone.weight": torch.zeros(1)}, + tmp_path / "model-00001-of-00001.safetensors", + ) + save_file({"head.weight": torch.zeros(1)}, tmp_path / "joint_head.safetensors") + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps( + {"weight_map": {"backbone.weight": "model-00001-of-00001.safetensors"}} + ) + ) + + loader = DefaultModelLoader(LoadConfig()) + secondary = DefaultModelLoader.Source( + model_or_path=str(tmp_path), + revision=None, + prefix="pooler.head.", + allow_patterns_overrides=["joint_head.safetensors"], + ) + assert [n for n, _ in loader._get_weights_iterator(secondary)] == [ + "pooler.head.head.weight" + ] + # The primary source still filters files the backbone index doesn't list. + primary = DefaultModelLoader.Source(model_or_path=str(tmp_path), revision=None) + assert [n for n, _ in loader._get_weights_iterator(primary)] == ["backbone.weight"] diff --git a/tests/models/multimodal/pooling/test_clef_omni.py b/tests/models/multimodal/pooling/test_clef_omni.py new file mode 100644 index 000000000000..ccdc913b3132 --- /dev/null +++ b/tests/models/multimodal/pooling/test_clef_omni.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The Clef-Omni pooler against a port of the checkpoint's joint schema head, +which loops over questions where the pooler batches them.""" + +import math +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +import torch.nn.functional as F + +from vllm.model_executor.models.clef_omni import JointSchemaHead, JointSchemaPooler +from vllm.pooling_params import PoolingParams +from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates + +HIDDEN, VOCAB = 64, 500 +TYPES = [0, 1, 1, 2] +COUNTS = [2, 3, 1, 5] + + +def reference_logits(head, hidden, token_ids, embedding, question_spans, spans): + """``JointSchemaHead.forward`` of the checkpoint's joint_schema_model.py + for one record, with ``spans`` the option spans of each question.""" + hidden = head.hidden_norm(hidden) + memory = head.memory_projection(hidden).unsqueeze(0) + global_vector = hidden[-1] + questions = torch.stack([hidden[s:e].mean(0) for s, e in question_spans]) + contexts = [torch.stack([hidden[s:e].mean(0) for s, e in q]) for q in spans] + lexicals = [ + torch.stack([embedding[token_ids[s:e]].mean(0) for s, e in q]) for q in spans + ] + options = torch.cat( + [ + head.option_context_projection(c) + + head.option_lexical_projection(lex) + + head.option_question_projection(questions[i]).unsqueeze(0) + for i, (c, lex) in enumerate(zip(contexts, lexicals)) + ] + ).unsqueeze(0) + for layer in head.evidence_layers: + options = layer(options, memory) + routed = torch.split(options[0], [len(q) for q in spans]) + + base = head.question_projection(questions) + summaries = [] + for field, opts in zip(base, routed): + weights = torch.softmax(opts @ field / math.sqrt(opts.shape[-1]), dim=0) + summaries.append((weights[:, None] * opts).sum(0)) + fields = ( + base + + head.option_summary_norm(torch.stack(summaries)) + + head.global_projection(global_vector) + + head.type_embedding(torch.tensor(TYPES)) + ).unsqueeze(0) + for layer in head.layers: + fields = layer(fields, memory) + fields = head.field_norm(fields[0]) + + logits = [] + for i, (field, lexical, opts) in enumerate(zip(fields, lexicals, routed)): + anchor = F.normalize(questions[i] + global_vector, dim=-1) + prior_scale = head.prior_logit_scale.clamp(max=math.log(100.0)).exp() + prior = prior_scale * (F.normalize(lexical, dim=-1) @ anchor) + opts = head.option_norm(opts) + field = field.expand_as(opts) + features = torch.cat([field, opts, field * opts, (field - opts).abs()], -1) + joint_scale = head.joint_logit_scale.clamp(max=math.log(100.0)).exp() + joint = joint_scale * F.cosine_similarity(field, opts, dim=-1) + joint = joint + head.residual_scorer(features).squeeze(-1) + logits.append(prior + torch.sigmoid(head.residual_gate) * joint) + return logits + + +@pytest.fixture +def pooler(): + torch.manual_seed(0) + head = JointSchemaHead( + hidden_size=HIDDEN, + width=32, + routing_layers=2, + layers=2, + heads=4, + feedforward=48, + ).eval() + with torch.no_grad(): + for p in head.parameters(): + p.add_(torch.randn_like(p) * 0.2) + embedding = torch.randn(VOCAB, HIDDEN) + config = SimpleNamespace( + scheduler_config=SimpleNamespace( + enable_chunked_prefill=False, async_scheduling=False + ) + ) + with patch( + "vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config", + return_value=config, + ): + pooler = JointSchemaPooler(head, lambda ids: embedding[ids]) + return pooler, embedding + + +def schema(length: int): + """Questions after a 20-token state, with options of 1 to 4 tokens.""" + pos = 20 + question_spans: list[tuple[int, int]] = [] + spans: list[list[tuple[int, int]]] = [] + for count in COUNTS: + question_spans.append((pos, pos + 4)) + pos += 6 + spans.append([]) + for k in range(count): + n = 1 + (k * 3) % 4 + spans[-1].append((pos, pos + n)) + pos += n + 1 + assert pos < length + extra = { + "question_types": TYPES, + "question_spans": [[s - length, e - length] for s, e in question_spans], + "option_spans": [[s - length, e - length] for q in spans for s, e in q], + "option_counts": COUNTS, + } + return {"joint_schema": extra}, question_spans, spans + + +@torch.inference_mode() +def test_pooler_matches_reference(pooler): + pooler, embedding = pooler + lengths = [120, 90, 100] + token_ids = [torch.randint(0, VOCAB, (n,)) for n in lengths] + hidden = torch.randn(sum(lengths), HIDDEN) + first, first_questions, first_spans = schema(lengths[0]) + last, last_questions, last_spans = schema(lengths[2]) + params = [ + PoolingParams(task="token_classify", use_activation=True, extra_kwargs=first), + # A prompt without a schema, like the profiling run's. + PoolingParams(task="token_classify", use_activation=True), + PoolingParams(task="token_classify", use_activation=False, extra_kwargs=last), + ] + padded = torch.zeros(len(lengths), max(lengths), dtype=torch.long) + for i, ids in enumerate(token_ids): + padded[i, : len(ids)] = ids + metadata = PoolingMetadata( + prompt_lens=torch.tensor(lengths), + prompt_token_ids=None, + prompt_token_ids_cpu=padded, + pooling_params=params, + pooling_states=[PoolingStates() for _ in lengths], + ) + metadata.build_pooling_cursor( + torch.tensor(lengths).numpy(), + seq_lens_cpu=torch.tensor(lengths), + device=torch.device("cpu"), + ) + probs, empty, logits = pooler(hidden, metadata) + + first_hidden, _, last_hidden = torch.split(hidden, lengths) + expected = reference_logits( + pooler.head, first_hidden, token_ids[0], embedding, first_questions, first_spans + ) + expected_probs = torch.cat([q.softmax(-1) for q in expected]) + torch.testing.assert_close(probs, expected_probs, atol=1e-5, rtol=1e-4) + assert empty.numel() == 0 + expected = reference_logits( + pooler.head, last_hidden, token_ids[2], embedding, last_questions, last_spans + ) + torch.testing.assert_close(logits, torch.cat(expected), atol=1e-5, rtol=1e-4) diff --git a/tests/models/multimodal/processing/test_qwen3_omni.py b/tests/models/multimodal/processing/test_qwen3_omni.py index 301471f45a7b..c0ebf0297184 100644 --- a/tests/models/multimodal/processing/test_qwen3_omni.py +++ b/tests/models/multimodal/processing/test_qwen3_omni.py @@ -108,3 +108,33 @@ def get_token_count(duration: float) -> int: f"Expected longer audio (2s) to have more tokens than shorter (1s). " f"Got short={short_tokens}, long={long_tokens}" ) + + +@pytest.mark.parametrize("model_id", ["Cloudflare/clef-omni"]) +def test_audio_in_video_follows_requested_fps(model_id: str) -> None: + """Audio and video tokens interleave by the fps the video was processed + at, as in HF, not by the default of one frame a second.""" + ctx = build_model_context( + model_id, + limit_mm_per_prompt={"audio": 1, "image": 0, "video": 1}, + ) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + rng = np.random.RandomState(42) + video = rng.randint(0, 255, (4, 64, 64, 3), dtype=np.uint8) + audio = rng.rand(32000).astype(np.float32) + prompt = "<|vision_start|><|video_pad|><|vision_end|>" + mm_kwargs: dict[str, Any] = { + "use_audio_in_video": True, + "videos_kwargs": {"fps": 2.0, "do_sample_frames": False}, + } + processed = processor( + prompt, + mm_items=processor.info.parse_mm_data( + {"video": [video], "audio": [(audio, 16000)]} + ), + hf_processor_mm_kwargs=mm_kwargs, + ) + expected = processor.info.get_hf_processor()( + text=[prompt], videos=[video], audio=[audio], return_tensors="pt", **mm_kwargs + ) + assert processed["prompt_token_ids"] == expected["input_ids"][0].tolist() diff --git a/tests/models/registry.py b/tests/models/registry.py index 2ed6d4871c27..0a7027461ef3 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1351,6 +1351,10 @@ def check_available_online( "Qwen/Qwen3-Omni-30B-A3B-Instruct", max_model_len=4096, ), + "ClefOmniForJointSchema": _HfExamplesInfo( + "Cloudflare/clef-omni", + max_model_len=4096, + ), "Qwen3ASRForConditionalGeneration": _HfExamplesInfo( "Qwen/Qwen3-ASR-0.6B", max_model_len=4096, diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index 407bffed20a5..0066d239eae2 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -21,6 +21,7 @@ from vllm.multimodal.video import ( PYNVVIDEOCODEC_VIDEO_BACKEND, VIDEO_LOADER_REGISTRY, + ClefOmniVideoBackend, DynamicVideoBackend, Glm5NextVideoBackend, GLM46VVideoBackend, @@ -1813,3 +1814,18 @@ def test_glm5next_read_frames_dense_walk_matches_stock(tmp_path): assert abs(round(float(np.asarray(frame).mean())) - idx) <= 1 # One initial seek, then pure walking -- no re-seek churn. assert cap.seeks == 1 + + +def test_clef_omni_backend_samples_reference_frames(tmp_path): + """Clef-Omni samples the first frame at or after each 0.5 s boundary + (the reference's 2 fps timestamp sampling), not uniform strides.""" + pytest.importorskip("cv2") + + # 2 s at 30 fps: the reference reads frames 0, 15, 30 and 45, where + # uniform sampling would read 0, 19, 39 and 59. + path = _write_gray_video(tmp_path, total_frames=60, fps=30) + frames, metadata = ClefOmniVideoBackend.load_bytes(path.read_bytes()) + assert metadata["frames_indices"] == [0, 15, 30, 45] + assert metadata["do_sample_frames"] is False + for frame, idx in zip(frames, [0, 15, 30, 45]): + assert abs(round(float(np.asarray(frame).mean())) - idx) <= 1 diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 6bf8d690265a..9bbe4e879b34 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -5,7 +5,6 @@ from fastapi import FastAPI import vllm.envs as envs -from vllm.logger import init_logger if TYPE_CHECKING: from argparse import Namespace @@ -18,8 +17,6 @@ else: RequestLogger = object -logger = init_logger(__name__) - def register_generate_api_routers(app: FastAPI): from vllm.entrypoints.openai.chat_completion.api_router import ( @@ -56,12 +53,6 @@ def register_generate_api_routers(app: FastAPI): register_generative_scoring_api_router(app) - from .structured_decisions.api_router import ( - register_structured_decisions_api_router, - ) - - register_structured_decisions_api_router(app) - async def init_generate_state( engine_client: "EngineClient", @@ -231,33 +222,3 @@ async def init_generate_state( state.openai_serving_models, request_logger=request_logger, ) - - from .structured_decisions.serving import ServingStructuredDecisions - from .structured_decisions.strategies import ReadContext, select_read_strategy - - strategy = None - if "generate" in supported_tasks: - try: - strategy_cls = select_read_strategy(engine_client.model_config) - strategy = strategy_cls( - ReadContext( - engine_client=engine_client, - online_renderer=state.online_renderer, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - default_chat_template_kwargs=default_chat_template_kwargs, - ) - ) - except ValueError as e: - # Info, since every model without a read strategy lands here. - logger.info("/v1/systemone is disabled: %s", e) - strategy = None - state.serving_structured_decisions = ( - ServingStructuredDecisions( - state.openai_serving_models, - strategy, - request_logger=request_logger, - ) - if strategy is not None - else None - ) diff --git a/vllm/entrypoints/generate/structured_decisions/api_router.py b/vllm/entrypoints/generate/structured_decisions/api_router.py index 5729b58b2e50..a3afe7fe1120 100644 --- a/vllm/entrypoints/generate/structured_decisions/api_router.py +++ b/vllm/entrypoints/generate/structured_decisions/api_router.py @@ -1,19 +1,30 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from argparse import Namespace from http import HTTPStatus +from typing import Any from fastapi import APIRouter, Depends, FastAPI, Request from fastapi.responses import JSONResponse +from starlette.datastructures import State +from vllm.engine.protocol import EngineClient from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.entrypoints.serve.utils.api_utils import ( load_aware_call, validate_json_request, with_cancellation, ) +from vllm.entrypoints.serve.utils.request_logger import RequestLogger +from vllm.logger import init_logger +from vllm.renderers.chat_utils import load_chat_template +from vllm.tasks import SupportedTask from .protocol import StructuredDecisionRequest from .serving import ServingStructuredDecisions +from .strategies import ReadContext, select_read_strategy + +logger = init_logger(__name__) router = APIRouter() @@ -50,3 +61,31 @@ async def create_structured_decision( def register_structured_decisions_api_router(app: FastAPI): app.include_router(router) + + +def init_structured_decisions_state( + engine_client: EngineClient, + state: State, + args: Namespace, + request_logger: RequestLogger | None, + supported_tasks: tuple[SupportedTask, ...], + default_chat_template_kwargs: dict[str, Any], +): + state.serving_structured_decisions = None + try: + strategy_cls = select_read_strategy(engine_client.model_config, supported_tasks) + strategy = strategy_cls( + ReadContext( + engine_client=engine_client, + online_renderer=state.online_renderer, + chat_template=load_chat_template(args.chat_template), + chat_template_content_format=args.chat_template_content_format, + default_chat_template_kwargs=default_chat_template_kwargs, + ) + ) + except ValueError as e: + logger.debug("/v1/systemone is disabled: %s", e) + return + state.serving_structured_decisions = ServingStructuredDecisions( + state.openai_serving_models, strategy, request_logger=request_logger + ) diff --git a/vllm/entrypoints/generate/structured_decisions/protocol.py b/vllm/entrypoints/generate/structured_decisions/protocol.py index ab89084dff2d..a8f104b8cb09 100644 --- a/vllm/entrypoints/generate/structured_decisions/protocol.py +++ b/vllm/entrypoints/generate/structured_decisions/protocol.py @@ -30,6 +30,15 @@ class StructuredDecisionRequest(OpenAIBaseModel): instructions: str | None = Field( default=None, description="Context placed ahead of the questions." ) + images: list[str] | None = Field( + default=None, description="Image URLs or data URLs, for models that see." + ) + audio: list[str] | None = Field( + default=None, description="Audio URLs or data URLs, for models that hear." + ) + videos: list[str] | None = Field( + default=None, description="Video URLs or data URLs, for models that see." + ) chat_template_kwargs: dict[str, Any] | None = None cache_salt: str | None = Field( default=None, diff --git a/vllm/entrypoints/generate/structured_decisions/serving.py b/vllm/entrypoints/generate/structured_decisions/serving.py index 2633647a0b65..52fdd430b116 100644 --- a/vllm/entrypoints/generate/structured_decisions/serving.py +++ b/vllm/entrypoints/generate/structured_decisions/serving.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio -import json from typing import Any from fastapi import Request @@ -24,10 +23,6 @@ logger = init_logger(__name__) -def state_text(state: Any) -> str: - return state if isinstance(state, str) else json.dumps(state, ensure_ascii=False) - - def parse_questions( request: StructuredDecisionRequest, limits: DecisionLimits, @@ -72,6 +67,18 @@ def __init__( self.strategy = strategy self.limits = strategy.limits() + def _media(self, request: StructuredDecisionRequest) -> dict[str, list[str]]: + urls = { + "image": request.images, + "audio": request.audio, + "video": request.videos, + } + media = {modality: items for modality, items in urls.items() if items} + if unsupported := media.keys() - self.strategy.modalities: + names = ", ".join(sorted(unsupported)) + raise StructuredDecisionError(f"this model does not take {names} inputs") + return media + async def create_decision( self, request: StructuredDecisionRequest, @@ -87,8 +94,9 @@ async def create_decision( request_id = f"decision-{base_id}" try: questions = parse_questions(request, self.limits) + media = self._media(request) lora_request = self._maybe_get_adapters(request) # type: ignore[arg-type] - engine_client.check_admission(len(questions)) + engine_client.check_admission(self.strategy.num_requests(questions)) trace_headers = ( None if raw_request is None @@ -97,13 +105,14 @@ async def create_decision( reads = await self.strategy.read( questions, request.instructions, - state_text(request.state), + self.strategy.render_state(request.state), request_id=request_id, chat_template_kwargs=request.chat_template_kwargs, lora_request=lora_request, priority=request.priority, cache_salt=request.cache_salt, trace_headers=trace_headers, + media=media, ) except StructuredDecisionError as e: return self.create_error_response(e) diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index 48debdb0b30f..f4365924ddf0 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -5,24 +5,40 @@ Only models in READ_STRATEGIES are currently supported. """ +import asyncio import hashlib import json import math import random +import time from abc import ABC, abstractmethod from collections.abc import Mapping, Sequence from dataclasses import dataclass -from typing import Any +from functools import cached_property +from typing import Any, ClassVar, TypeVar +import numpy as np + +import vllm.envs as envs from vllm.config import ModelConfig from vllm.engine.protocol import EngineClient from vllm.entrypoints.generate.label_reads import next_token_label_reads -from vllm.inputs import EngineInput, tokens_input +from vllm.inputs import EngineInput, MultiModalInput, TokensPrompt, tokens_input +from vllm.logger import init_logger from vllm.lora.request import LoRARequest +from vllm.multimodal.media import ( + MEDIA_CONNECTOR_REGISTRY, + ImageMediaIO, + MediaConnector, + MediaIO, + VideoMediaIO, +) +from vllm.pooling_params import PoolingParams from vllm.renderers.chat_utils import ChatTemplateContentFormatOption from vllm.renderers.inputs.preprocess import extract_prompt_components from vllm.renderers.online_renderer import OnlineRenderer from vllm.sampling_params import SamplingParams +from vllm.tasks import SupportedTask from vllm.tokenizers import TokenizerLike from .protocol import ReadPromptRequest @@ -34,6 +50,10 @@ label_softmax, ) +_M = TypeVar("_M") + +logger = init_logger(__name__) + @dataclass class ReadContext: @@ -59,13 +79,30 @@ class QuestionRead: output_tokens: int +def state_text(state: Any) -> str: + return state if isinstance(state, str) else json.dumps(state, ensure_ascii=False) + + class ReadStrategy(ABC): + #: The engine task reads run as. + task: ClassVar[SupportedTask] = "generate" + #: The modalities a request may attach media of. + modalities: ClassVar[frozenset[str]] = frozenset() + def __init__(self, context: ReadContext): self.context = context @abstractmethod def limits(self) -> DecisionLimits: ... + def render_state(self, state: Any) -> str: + """The request's state as the model reads it.""" + return state_text(state) + + def num_requests(self, questions: list[Question]) -> int: + """The engine requests a decision takes.""" + return len(questions) + @abstractmethod async def read( self, @@ -79,6 +116,7 @@ async def read( priority: int, cache_salt: str | None = None, trace_headers: Mapping[str, str] | None = None, + media: Mapping[str, Sequence[str]] | None = None, ) -> list[QuestionRead]: """One read per question, in order.""" @@ -214,6 +252,7 @@ async def read( priority: int, cache_salt: str | None = None, trace_headers: Mapping[str, str] | None = None, + media: Mapping[str, Sequence[str]] | None = None, ) -> list[QuestionRead]: ctx = self.context read_request = self._read_request(chat_template_kwargs, cache_salt) @@ -334,7 +373,315 @@ def _read_input( ) +class JointSchemaStrategy(ReadStrategy): + """Clef-Omni. The whole request is one pooling read: the state, then every + question with its options, as the model was trained to read them. The + model's joint schema head scores each option, so a read has no labels. + + Mirrors ``encode_record`` in the checkpoint's ``joint_schema_model.py``. + Each piece is tokenized on its own, so the spans of the questions and the + options are known. The spans count from the end of the prompt, so they + hold after the media placeholders, which come before the state, expand.""" + + task = "token_classify" + modalities = frozenset({"image", "audio", "video"}) + + SYSTEM_PROMPT = ( + "Read the complete state and schema. Decide every field jointly. Each " + "answer must be exactly one of that field's allowed options." + ) + #: The reference's default bound on a prompt, in tokens. + MAX_LENGTH = 64000 + TYPE_IDS = {"noul": 0, "choice": 1, "score": 2} + NOUL_CRITERIA = { + "true": "The proposition is true or the answer is yes.", + "false": "The proposition is false or the answer is no.", + } + PLACEHOLDERS = { + "image": "<|vision_start|><|image_pad|><|vision_end|>\n", + "audio": "<|audio_start|><|audio_pad|><|audio_end|>\n", + "video": "<|vision_start|><|video_pad|><|vision_end|>\n", + } + IMAGE_SIZE = {"shortest_edge": 65_536, "longest_edge": 1_048_576} + #: The reference samples videos at 2 fps; the "clef_omni" video backend + #: reproduces the reference's sampling, pixel cap and frame count cap. + VIDEO_FPS = 2 + + def __init__(self, context: ReadContext): + super().__init__(context) + self.tokenizer = context.online_renderer.renderer.get_tokenizer() + self.max_length = min( + self.MAX_LENGTH, context.engine_client.model_config.max_model_len + ) + self.prefix = self._tokens( + f"<|im_start|>system\n{self.SYSTEM_PROMPT}<|im_end|>\n" + "<|im_start|>user\nSTATE:\n" + ) + self.suffix = self._tokens( + "\n<|im_end|>\n<|im_start|>assistant\n\n\n\n\n" + "JOINT SCHEMA DECISIONS:" + ) + + def limits(self) -> DecisionLimits: + return DecisionLimits(max_questions=64, max_options=len(LABELS)) + + def render_state(self, state: Any) -> str: + return self._render(state) + + def num_requests(self, questions: list[Question]) -> int: + return 1 + + @staticmethod + def _render(value: Any) -> str: + if isinstance(value, str): + return value + return json.dumps( + value, ensure_ascii=False, separators=(",", ":"), sort_keys=True + ) + + def _tokens(self, text: str) -> list[int]: + return self.tokenizer.encode(text, add_special_tokens=False) + + @cached_property + def _connector(self) -> MediaConnector: + model_config = self.context.engine_client.model_config + return MEDIA_CONNECTOR_REGISTRY.load( + envs.VLLM_MEDIA_CONNECTOR, + media_io_kwargs=model_config.get_multimodal_config().media_io_kwargs, + allowed_local_media_path=model_config.allowed_local_media_path, + allowed_media_domains=model_config.allowed_media_domains, + ) + + async def _fetch( + self, urls: Sequence[str], media_io: MediaIO[_M], timeout: int + ) -> list[_M]: + return await asyncio.gather( + *( + self._connector.load_from_url_async( + url, media_io, fetch_timeout=timeout + ) + for url in urls + ) + ) + + async def _soundtrack( + self, video_url: str + ) -> tuple[np.ndarray, int | float] | None: + """The video's soundtrack, or None when it has no decodable one.""" + try: + return await self._connector.fetch_audio_async(video_url) + except ValueError as e: + logger.debug("no decodable soundtrack in %s: %s", video_url, e) + return None + + async def _render_media( + self, media: Mapping[str, Sequence[str]], cache_salt: str | None + ) -> MultiModalInput: + """The system prompt and the media, with the media's placeholders + expanded. A video's soundtrack is heard with it when every video has + one.""" + connector = self._connector + video_io = VideoMediaIO( + ImageMediaIO(), + num_frames=-1, + fps=self.VIDEO_FPS, + video_backend="clef_omni", + ) + try: + images, audio, videos, tracks = await asyncio.gather( + asyncio.gather( + *map(connector.fetch_image_async, media.get("image", ())) + ), + asyncio.gather( + *map(connector.fetch_audio_async, media.get("audio", ())) + ), + self._fetch( + media.get("video", ()), video_io, envs.VLLM_VIDEO_FETCH_TIMEOUT + ), + asyncio.gather(*map(self._soundtrack, media.get("video", ()))), + ) + except ValueError as e: + raise StructuredDecisionError(f"media: {e}") from e + + audio_in_video = bool(videos) and all(t is not None for t in tracks) + if audio_in_video and audio: + raise StructuredDecisionError( + "audio clips cannot be read with videos that have soundtracks" + ) + text = ( + self.PLACEHOLDERS["image"] * len(images) + + self.PLACEHOLDERS["audio"] * len(audio) + + self.PLACEHOLDERS["video"] * len(videos) + ) + if audio_in_video: + audio = [track for track in tracks if track is not None] + mm_data: dict[str, Any] = {} + if images: + mm_data["image"] = images + if audio: + mm_data["audio"] = audio + if videos: + mm_data["video"] = videos + mm_kwargs: dict[str, Any] = {"images_kwargs": {"size": self.IMAGE_SIZE}} + if videos: + mm_kwargs["use_audio_in_video"] = audio_in_video + prompt = TokensPrompt( + prompt_token_ids=self.prefix + self._tokens(text), + multi_modal_data=mm_data, + mm_processor_kwargs=mm_kwargs, + ) + if cache_salt is not None: + prompt["cache_salt"] = cache_salt + renderer = self.context.online_renderer.renderer + try: + engine_input = await renderer.process_for_engine_async(prompt, time.time()) + except ValueError as e: + raise StructuredDecisionError(f"media: {e}") from e + assert engine_input["type"] == "multimodal" + return engine_input + + def _options(self, q: Question) -> list[tuple[int, str, Any]]: + """Each option's index in ``q.options``, id and description, in the + model's order: choices sorted by name, noul true then false, score + levels by index.""" + if q.type.name == "noul": + return [ + ( + i, + key, + self.NOUL_CRITERIA[key] if o.description is None else o.description, + ) + for i, (key, o) in enumerate(zip(("true", "false"), q.options)) + ] + if q.type.name == "choice": + return sorted( + ((i, o.name, o.description) for i, o in enumerate(q.options)), + key=lambda option: option[1], + ) + if q.type.name == "score": + return [(i, str(i), o.name) for i, o in enumerate(q.options)] + raise StructuredDecisionError( + f"question {q.id!r}: this model does not answer {q.type.name!r} questions" + ) + + def encode( + self, questions: list[Question], state: str, prefix: list[int] | None = None + ) -> tuple[list[int], dict[str, list[Any]], list[list[int]]]: + """The prompt, the schema for the pooler, and each question's option + indices in the order of the output. ``prefix`` is what precedes the + state: the system prompt, and the media if there are any.""" + prefix = self.prefix if prefix is None else prefix + schema_ids = self._tokens("\n\nSCHEMA FIELDS:\n") + question_spans, option_spans, orders = [], [], [] + for n, q in enumerate(questions, start=1): + schema_ids += self._tokens( + f"\nFIELD {n}\nID: {q.id}\nTYPE: {q.type.name}\nINSTRUCTION: " + ) + start = len(schema_ids) + schema_ids += self._tokens(q.instructions or q.id) + question_spans.append([start, len(schema_ids)]) + schema_ids += self._tokens("\nALLOWED OPTIONS:\n") + options = self._options(q) + for m, (_, option_id, description) in enumerate(options, start=1): + schema_ids += self._tokens(f"OPTION {m}: ") + start = len(schema_ids) + semantics = {"option_id": option_id} + if description is not None: + semantics["description"] = description + schema_ids += self._tokens(self._render(semantics)) + option_spans.append([start, len(schema_ids)]) + schema_ids += self._tokens("\n") + schema_ids += self._tokens("END FIELD\n") + orders.append([i for i, _, _ in options]) + + fixed = len(prefix) + len(schema_ids) + len(self.suffix) + if fixed > self.max_length: + raise StructuredDecisionError( + f"the questions take {fixed} tokens, more than the " + f"{self.max_length} this model reads" + ) + state_ids = self._tokens(state)[: self.max_length - fixed] + prompt = prefix + state_ids + schema_ids + self.suffix + # The schema ends len(self.suffix) tokens before the prompt does. + end = len(schema_ids) + len(self.suffix) + schema: dict[str, list[Any]] = { + "question_types": [self.TYPE_IDS[q.type.name] for q in questions], + "question_spans": [[s - end, e - end] for s, e in question_spans], + "option_spans": [[s - end, e - end] for s, e in option_spans], + "option_counts": [len(order) for order in orders], + } + return prompt, schema, orders + + async def read( + self, + questions: list[Question], + instructions: str | None, + state: str, + *, + request_id: str, + chat_template_kwargs: dict[str, Any] | None, + lora_request: LoRARequest | None, + priority: int, + cache_salt: str | None = None, + trace_headers: Mapping[str, str] | None = None, + media: Mapping[str, Sequence[str]] | None = None, + ) -> list[QuestionRead]: + if instructions: + raise StructuredDecisionError( + "this model reads a fixed prompt, so it takes no instructions" + ) + engine_input: EngineInput + if media: + engine_input = await self._render_media(media, cache_salt) + prompt, schema, orders = self.encode( + questions, state, engine_input["prompt_token_ids"] + ) + engine_input["prompt_token_ids"] = prompt + else: + prompt, schema, orders = self.encode(questions, state) + engine_input = tokens_input(prompt, cache_salt=cache_salt) + result = None + async for result in self.context.engine_client.encode( + engine_input, + PoolingParams(task="token_classify", extra_kwargs={"joint_schema": schema}), + f"{request_id}-0", + lora_request=lora_request, + trace_headers=trace_headers, + priority=priority, + ): + pass + if result is None: + raise ValueError("Failed to score the questions") + if result.error is not None: + raise ValueError(f"Failed to score the questions: {result.error.message}") + probs = result.outputs.data.tolist() + if len(probs) != sum(schema["option_counts"]): + raise ValueError( + f"Expected {sum(schema['option_counts'])} option scores, " + f"got {len(probs)}" + ) + + reads, offset = [], 0 + for n, order in enumerate(orders): + q_probs = [0.0] * len(order) + for i, p in zip(order, probs[offset : offset + len(order)]): + q_probs[i] = p + offset += len(order) + reads.append( + QuestionRead( + probs=q_probs, + label_mass=1.0, + argmax_is_label=True, + # One forward pass reads the prompt for every question. + input_tokens=len(prompt) if n == 0 else 0, + output_tokens=0, + ) + ) + return reads + + READ_STRATEGIES: dict[str, type[ReadStrategy]] = { + "ClefOmniForJointSchema": JointSchemaStrategy, "Qwen3ForCausalLM": NextTokenStrategy, "Qwen3_5ForConditionalGeneration": NextTokenStrategy, "Qwen3_5MoeForConditionalGeneration": NextTokenStrategy, @@ -347,7 +694,9 @@ def _read_input( LOGPROBS_MODES = frozenset({"raw_logprobs", "processed_logprobs"}) -def select_read_strategy(model_config: ModelConfig) -> type[ReadStrategy]: +def select_read_strategy( + model_config: ModelConfig, supported_tasks: Sequence[SupportedTask] +) -> type[ReadStrategy]: """Raise ValueError when the model cannot serve structured decisions.""" strategy = READ_STRATEGIES.get(model_config.architecture) if strategy is None: @@ -356,7 +705,12 @@ def select_read_strategy(model_config: ModelConfig) -> type[ReadStrategy]: f"{model_config.architecture}. Supported architectures: " f"{sorted(READ_STRATEGIES)}" ) - if model_config.logprobs_mode not in LOGPROBS_MODES: + if strategy.task not in supported_tasks: + raise ValueError( + f"Structured decisions on {model_config.architecture} need the " + f"{strategy.task} task" + ) + if strategy.task == "generate" and model_config.logprobs_mode not in LOGPROBS_MODES: raise ValueError( "Structured decisions need --logprobs-mode raw_logprobs or " f"processed_logprobs, not {model_config.logprobs_mode}" diff --git a/vllm/entrypoints/launchers/api_server/app_state.py b/vllm/entrypoints/launchers/api_server/app_state.py index ce68890717e3..ff87dbea4915 100644 --- a/vllm/entrypoints/launchers/api_server/app_state.py +++ b/vllm/entrypoints/launchers/api_server/app_state.py @@ -157,6 +157,19 @@ async def init_app_state( init_pooling_state(engine_client, state, args, request_logger, supported_tasks) + from vllm.entrypoints.generate.structured_decisions.api_router import ( + init_structured_decisions_state, + ) + + init_structured_decisions_state( + engine_client, + state, + args, + request_logger, + supported_tasks, + default_chat_template_kwargs, + ) + await init_endpoint_plugins_state(engine_client, state, args) state.enable_server_load_tracking = args.enable_server_load_tracking diff --git a/vllm/entrypoints/launchers/api_server/routers.py b/vllm/entrypoints/launchers/api_server/routers.py index 40d04862afb2..52c995a30135 100644 --- a/vllm/entrypoints/launchers/api_server/routers.py +++ b/vllm/entrypoints/launchers/api_server/routers.py @@ -25,6 +25,12 @@ def register_api_routers( register_models_api_router(app) + from vllm.entrypoints.generate.structured_decisions.api_router import ( + register_structured_decisions_api_router, + ) + + register_structured_decisions_api_router(app) + from vllm.entrypoints.serve.sagemaker.api_router import ( attach_router as register_sagemaker_api_router, ) diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index b162f60b0577..6bedabfdc013 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/vllm/model_executor/model_loader/default_loader.py @@ -247,9 +247,11 @@ def _prepare_weights( subfolder=subfolder, revision=revision, ) - hf_weights_files = filter_duplicate_safetensors_files( - hf_weights_files, hf_folder, index_file - ) + # only filter the weights from index without explicit allow patterns + if allow_patterns_overrides is None: + hf_weights_files = filter_duplicate_safetensors_files( + hf_weights_files, hf_folder, index_file + ) else: hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files) diff --git a/vllm/model_executor/models/clef_omni.py b/vllm/model_executor/models/clef_omni.py new file mode 100644 index 000000000000..7cd0d676873b --- /dev/null +++ b/vllm/model_executor/models/clef_omni.py @@ -0,0 +1,437 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Clef-Omni (https://huggingface.co/Cloudflare/clef-omni). + +The Qwen3-Omni thinker with a joint schema head that scores every option of +every question in one forward pass. The head reads the final hidden states of +the whole prompt, so it runs as a ``token_classify`` pooler. + +Each request carries its schema in ``PoolingParams.extra_kwargs`` under +``"joint_schema"``: + +- ``question_types``: 0 (noul), 1 (choice) or 2 (score) per question. +- ``question_spans``: the ``[start, end)`` of each question's instructions. +- ``option_spans``: the ``[start, end)`` of each option, question by question. +- ``option_counts``: the number of options of each question. + +Spans count from the end of the prompt (``-len(prompt) <= start < end <= 0``), +so they hold after multimodal placeholders are expanded. The output has one +value per option, in ``option_spans`` order: its probability within its +question, or its logit with ``use_activation=False``. +""" + +import math +from collections.abc import Callable, Mapping, Set +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn +from transformers.feature_extraction_utils import BatchFeature + +from vllm.config import VllmConfig +from vllm.model_executor.layers.pooler import Pooler, PoolingParamsUpdate +from vllm.model_executor.layers.pooler.tokwise.methods import AllPool +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from vllm.model_executor.model_loader.default_loader import DefaultModelLoader +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.tasks import PoolingTask +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.outputs import PoolerOutput +from vllm.v1.pool.metadata import PoolingMetadata + +from .interfaces_base import default_pooling_type +from .qwen3_omni_moe_thinker import ( + Qwen3OmniMoeThinkerDummyInputsBuilder, + Qwen3OmniMoeThinkerForConditionalGeneration, + Qwen3OmniMoeThinkerMultiModalProcessor, + Qwen3OmniMoeThinkerProcessingInfo, +) + +JOINT_HEAD_WEIGHTS = "joint_head.safetensors" +NUM_QUESTION_TYPES = 3 +MAX_LOGIT_SCALE = math.log(100.0) + + +def segment_mean( + values: torch.Tensor, segment_ids: torch.Tensor, lengths: torch.Tensor +) -> torch.Tensor: + sums = values.new_zeros(lengths.shape[0], values.shape[-1], dtype=torch.float32) + sums.index_add_(0, segment_ids, values.float()) + return (sums / lengths[:, None]).to(values.dtype) + + +@dataclass +class SchemaIndices: + """Where one prompt's questions and options are, on the model's device.""" + + positions: torch.Tensor + """Every token of every span, questions first.""" + span_ids: torch.Tensor + """The span of each of ``positions``.""" + span_lengths: torch.Tensor + """``[Q + O]``.""" + question_types: torch.Tensor + """``[Q]``.""" + option_question: torch.Tensor + """The question of each option, ``[O]``.""" + option_slots: torch.Tensor + """Each option's flat index in a ``[Q, max_options]`` grid, ``[O]``.""" + option_mask: torch.Tensor + """``[Q, max_options]``, true at the slots of options.""" + + +class EvidenceRoutingLayer(nn.Module): + def __init__(self, width: int, heads: int, feedforward: int) -> None: + super().__init__() + self.query_norm = nn.LayerNorm(width) + self.memory_norm = nn.LayerNorm(width) + self.attention = nn.MultiheadAttention(width, heads, batch_first=True) + self.feedforward_norm = nn.LayerNorm(width) + # The Identity takes the place of the checkpoint's dropout. + self.feedforward = nn.Sequential( + nn.Linear(width, feedforward), + nn.GELU(), + nn.Identity(), + nn.Linear(feedforward, width), + ) + + def forward(self, queries: torch.Tensor, memory: torch.Tensor) -> torch.Tensor: + # Separate key and value tensors, like the reference, keep attention's + # unpacked input projections. + routed, _ = self.attention( + self.query_norm(queries), + self.memory_norm(memory), + self.memory_norm(memory), + need_weights=False, + ) + queries = queries + routed + return queries + self.feedforward(self.feedforward_norm(queries)) + + +class JointSchemaHead(nn.Module): + def __init__( + self, + hidden_size: int, + width: int, + routing_layers: int, + layers: int, + heads: int, + feedforward: int, + ) -> None: + super().__init__() + self.hidden_norm = nn.LayerNorm(hidden_size) + self.memory_projection = nn.Linear(hidden_size, width, bias=False) + self.question_projection = nn.Linear(hidden_size, width, bias=False) + self.option_question_projection = nn.Linear(hidden_size, width, bias=False) + self.global_projection = nn.Linear(hidden_size, width, bias=False) + self.option_context_projection = nn.Linear(hidden_size, width, bias=False) + self.option_lexical_projection = nn.Linear(hidden_size, width, bias=False) + self.type_embedding = nn.Embedding(NUM_QUESTION_TYPES, width) + self.evidence_layers = nn.ModuleList( + EvidenceRoutingLayer(width, heads, feedforward) + for _ in range(routing_layers) + ) + self.option_summary_norm = nn.LayerNorm(width) + self.layers = nn.ModuleList( + nn.TransformerDecoderLayer( + d_model=width, + nhead=heads, + dim_feedforward=feedforward, + dropout=0.0, + activation="gelu", + batch_first=True, + norm_first=True, + ) + for _ in range(layers) + ) + self.field_norm = nn.LayerNorm(width) + self.option_norm = nn.LayerNorm(width) + self.residual_scorer = nn.Sequential( + nn.Linear(width * 4, width), + nn.GELU(), + nn.Identity(), + nn.Linear(width, 1), + ) + self.prior_logit_scale = nn.Parameter(torch.zeros(())) + self.joint_logit_scale = nn.Parameter(torch.zeros(())) + self.residual_gate = nn.Parameter(torch.zeros(())) + + def forward( + self, + hidden_states: torch.Tensor, + lexical: torch.Tensor, + schema: SchemaIndices, + ) -> torch.Tensor: + """The ``[O]`` option logits of one prompt, from its ``[L, hidden]`` + final hidden states and each option's ``[O, hidden]`` mean output + embedding.""" + num_questions = schema.question_types.shape[0] + option_question = schema.option_question + hidden_states = self.hidden_norm(hidden_states) + memory = self.memory_projection(hidden_states).unsqueeze(0) + global_vector = hidden_states[-1] + means = segment_mean( + hidden_states[schema.positions], schema.span_ids, schema.span_lengths + ) + question_vectors, contexts = means[:num_questions], means[num_questions:] + + options = ( + self.option_context_projection(contexts) + + self.option_lexical_projection(lexical) + + self.option_question_projection(question_vectors)[option_question] + ).unsqueeze(0) + for layer in self.evidence_layers: + options = layer(options, memory) + options = options[0] + + # Each field summarizes its options, weighted by their match to it. + base_fields = self.question_projection(question_vectors) + width = options.shape[-1] + grid = options.new_zeros(schema.option_mask.numel(), width) + grid[schema.option_slots] = options + grid = grid.view(*schema.option_mask.shape, width) + routing = torch.einsum("qow,qw->qo", grid, base_fields) / math.sqrt(width) + routing = routing.masked_fill(~schema.option_mask, float("-inf")) + summaries = torch.einsum("qo,qow->qw", routing.softmax(dim=-1), grid) + fields = ( + base_fields + + self.option_summary_norm(summaries) + + self.global_projection(global_vector).unsqueeze(0) + + self.type_embedding(schema.question_types) + ).unsqueeze(0) + for layer in self.layers: + fields = layer(fields, memory) + fields = self.field_norm(fields[0])[option_question] + + anchors = F.normalize(question_vectors + global_vector, dim=-1) + prior_scale = self.prior_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() + prior = prior_scale * ( + F.normalize(lexical, dim=-1) * anchors[option_question] + ).sum(dim=-1) + options = self.option_norm(options) + cosine = F.cosine_similarity(fields, options, dim=-1) + features = torch.cat( + [fields, options, fields * options, torch.abs(fields - options)], dim=-1 + ) + residual = self.residual_scorer(features).squeeze(-1) + joint_scale = self.joint_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() + joint = joint_scale * cosine + residual + return prior + torch.sigmoid(self.residual_gate) * joint + + +def parse_joint_schema( + extra_kwargs: dict[str, Any] | None, prompt_len: int +) -> tuple[list[int], list[tuple[int, int]], list[int]] | None: + """The question types, the absolute spans of the questions then of the + options, and the option counts. None if the schema is missing or invalid.""" + schema = (extra_kwargs or {}).get("joint_schema") + if not isinstance(schema, dict): + return None + try: + types = [int(t) for t in schema["question_types"]] + counts = [int(c) for c in schema["option_counts"]] + question_spans = [ + (int(start) + prompt_len, int(end) + prompt_len) + for start, end in schema["question_spans"] + ] + option_spans = [ + (int(start) + prompt_len, int(end) + prompt_len) + for start, end in schema["option_spans"] + ] + except (KeyError, TypeError, ValueError): + return None + spans = question_spans + option_spans + if ( + not types + or len(counts) != len(types) + or len(question_spans) != len(types) + or len(option_spans) != sum(counts) + or any(not 0 <= t < NUM_QUESTION_TYPES for t in types) + or any(c < 1 for c in counts) + or any(not 0 <= start < end <= prompt_len for start, end in spans) + ): + return None + return types, spans, counts + + +class JointSchemaPooler(Pooler): + """Runs the joint schema head on each finished prompt. + + It never raises, because it runs in the engine core: a prompt without a + valid schema, such as a profiling dummy, gets an empty output. + """ + + def __init__( + self, + head: JointSchemaHead, + embed_lexical: Callable[[torch.Tensor], torch.Tensor], + ) -> None: + super().__init__() + self.pooling = AllPool() + self.head = head + self.embed_lexical = embed_lexical + + def get_supported_tasks(self) -> Set[PoolingTask]: + return {"token_classify"} + + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate: + return PoolingParamsUpdate(requires_token_ids=True) + + def forward( + self, + hidden_states: torch.Tensor, + pooling_metadata: PoolingMetadata, + ) -> PoolerOutput: + pooled = self.pooling(hidden_states, pooling_metadata) + prompt_token_ids = pooling_metadata.get_prompt_token_ids_cpu() + outputs: list[torch.Tensor | None] = [] + for hidden, token_ids, params in zip( + pooled, prompt_token_ids, pooling_metadata.pooling_params + ): + if hidden is None: + outputs.append(None) + continue + parsed = parse_joint_schema(params.extra_kwargs, hidden.shape[0]) + if parsed is None: + outputs.append(hidden.new_empty(0, dtype=torch.float32)) + continue + schema, lexical_ids, lexical_span_ids = self._indices( + *parsed, hidden.device + ) + lexical = segment_mean( + self.embed_lexical( + async_tensor_h2d(token_ids[lexical_ids], hidden.device) + ), + lexical_span_ids, + schema.span_lengths[len(parsed[0]) :], + ) + logits = self.head(hidden, lexical, schema).float() + if params.use_activation is not False: + mask = schema.option_mask + grid = logits.new_full((mask.numel(),), float("-inf")) + grid[schema.option_slots] = logits + probs = grid.view(mask.shape).softmax(dim=-1) + logits = probs.view(-1)[schema.option_slots] + outputs.append(logits) + return outputs + + @staticmethod + def _indices( + types: list[int], + spans: list[tuple[int, int]], + counts: list[int], + device: torch.device, + ) -> tuple[SchemaIndices, torch.Tensor, torch.Tensor]: + """The schema's indices, plus the option tokens' positions (on the CPU) + and their option spans, for the lexical embeddings.""" + num_questions = len(types) + lengths = torch.tensor([end - start for start, end in spans]) + span_ids = torch.repeat_interleave(torch.arange(len(spans)), lengths) + positions = torch.cat([torch.arange(start, end) for start, end in spans]) + num_question_tokens = int(lengths[:num_questions].sum()) + + counts_cpu = torch.tensor(counts) + max_options = max(counts) + option_question = torch.repeat_interleave( + torch.arange(num_questions), counts_cpu + ) + option_slots = torch.cat( + [torch.arange(count) + q * max_options for q, count in enumerate(counts)] + ) + schema = SchemaIndices( + positions=async_tensor_h2d(positions, device), + span_ids=async_tensor_h2d(span_ids, device), + span_lengths=async_tensor_h2d(lengths, device), + question_types=async_tensor_h2d(types, device, dtype=torch.long), + option_question=async_tensor_h2d(option_question, device), + option_slots=async_tensor_h2d(option_slots, device), + option_mask=async_tensor_h2d( + torch.arange(max_options) < counts_cpu[:, None], device + ), + ) + lexical_span_ids = async_tensor_h2d( + span_ids[num_question_tokens:] - num_questions, device + ) + return schema, positions[num_question_tokens:], lexical_span_ids + + +MAX_FRAME_PIXELS = 256 * 32 * 32 + + +def cap_frame_pixels(frames: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: + """The reference reads each video frame at at most 256 * 32 * 32 pixels, + in even sizes. Frames are NHWC uint8, numpy or tensor; tensors stay on + their device (a CUDA tensor from torchcodec must not round-trip).""" + height, width = frames.shape[1:3] + if height * width <= MAX_FRAME_PIXELS: + return frames + scale = math.sqrt(MAX_FRAME_PIXELS / (width * height)) + new_width = max(2, round(width * scale / 2) * 2) + new_height = max(2, round(height * scale / 2) * 2) + is_tensor = isinstance(frames, torch.Tensor) + tensor = frames if is_tensor else torch.from_numpy(frames) + resized = F.interpolate( + tensor.permute(0, 3, 1, 2).float(), + size=(new_height, new_width), + mode="bilinear", + ) + if tensor.dtype == torch.uint8: + resized = resized.round() + resized = resized.to(tensor.dtype).permute(0, 2, 3, 1).contiguous() + return resized if is_tensor else resized.numpy() + + +def _cap_video(item: Any) -> Any: + if isinstance(item, tuple): + frames, metadata = item + return cap_frame_pixels(frames), metadata + return cap_frame_pixels(item) + + +class ClefOmniMultiModalProcessor(Qwen3OmniMoeThinkerMultiModalProcessor): + """Caps video frames at the reference's per-frame pixel count.""" + + def _call_hf_processor( + self, + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], + ) -> BatchFeature: + if videos := hf_data.get("videos"): + assert isinstance(videos, list) + hf_data = {**hf_data, "videos": [_cap_video(v) for v in videos]} + return super()._call_hf_processor(hf_data, hf_kwargs) + + +@default_pooling_type(tok_pooling_type="ALL") +@MULTIMODAL_REGISTRY.register_processor( + ClefOmniMultiModalProcessor, + info=Qwen3OmniMoeThinkerProcessingInfo, + dummy_inputs=Qwen3OmniMoeThinkerDummyInputsBuilder, +) +class ClefOmniForJointSchema(Qwen3OmniMoeThinkerForConditionalGeneration): + is_pooling_model = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + model_config = vllm_config.model_config + head_config = model_config.hf_config.joint_head_config + self.pooler = JointSchemaPooler( + JointSchemaHead(**head_config), self._embed_lexical + ) + # The head's weights ship as their own safetensors file. + self.secondary_weights = [ + DefaultModelLoader.Source( + model_or_path=model_config.model, + revision=model_config.revision, + prefix="pooler.head.", + allow_patterns_overrides=[JOINT_HEAD_WEIGHTS], + ) + ] + + def _embed_lexical(self, token_ids: torch.Tensor) -> torch.Tensor: + # The head embeds options with the output embeddings. ParallelLMHead + # forbids lookups, so use the vocab-parallel embedding's. + return VocabParallelEmbedding.forward(self.language_model.lm_head, token_ids) diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 4ae80a4a36e8..321b5c228973 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -83,6 +83,7 @@ PromptUpdate, PromptUpdateDetails, ) +from vllm.multimodal.utils import copy_mm_embedding_modality from vllm.sequence import IntermediateTensors from vllm.transformers_utils.processor import cached_processor_from_config from vllm.utils.gpu_sync_debug import gpu_sync_allowed @@ -1498,8 +1499,12 @@ def get_replacement_qwen2_use_audio_in_video(item_idx: int): audio_in_video_item_idx += 1 + # Match the fps the processor used, as HF and the mrope positions do. + processed_second_per_grid_ts = out_mm_data.get("second_per_grid_ts") second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None) - if second_per_grid_ts: + if isinstance(processed_second_per_grid_ts, torch.Tensor): + video_second_per_grid_t = processed_second_per_grid_ts[item_idx].item() + elif second_per_grid_ts: video_second_per_grid_t = second_per_grid_ts[item_idx] else: video_second_per_grid_t = 2.0 @@ -1904,7 +1909,9 @@ def embed_input_ids( embeddings_main, embeddings_multiscale = torch.split( embeddings, [visual_dim, multi_dim], dim=-1 ) - multimodal_embeddings[index] = embeddings_main + multimodal_embeddings[index] = copy_mm_embedding_modality( + embeddings, embeddings_main + ) multimodal_embeddings_multiscale.append(embeddings_multiscale) if not is_interleaved: current_positions = mm_positions[ diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 427648045f9d..36f9bd6e6fcc 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -287,6 +287,7 @@ _TOKEN_CLASSIFICATION_MODELS = { "BertForTokenClassification": ("bert", "BertForTokenClassification"), + "ClefOmniForJointSchema": ("clef_omni", "ClefOmniForJointSchema"), "ModernBertForTokenClassification": ( "modernbert", "ModernBertForTokenClassification", diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index fc0cdc28d478..75a6c1beb1cf 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -1570,3 +1570,67 @@ def load_bytes( valid_frame_indices=valid_frame_indices, ) return frames, metadata + + +@VIDEO_LOADER_REGISTRY.register("clef_omni") +class ClefOmniVideoBackend(VideoBackend): + """Clef-Omni's reference sampling: the first frame at or after each + ``1/fps`` boundary (frame i sits at ``i / original_fps`` seconds, so the + k-th boundary is crossed by frame ``ceil(k * original_fps / fps)``), + capped at 768 frames and padded to an even count by repeating the last + frame.""" + + MAX_FRAMES = 768 + + @classmethod + def compute_frames_index_to_sample( + cls, + source: VideoSourceMetadata, + target: VideoTargetMetadata, + **kwargs, + ) -> list[int]: + if source.original_fps <= 0: + raise ValueError("video has an unknown frame rate") + step = source.original_fps / target.fps + indices = [] + k = 0 + while (index := math.ceil(k * step - 1e-6)) < source.total_frames_num: + indices.append(index) + if len(indices) > cls.MAX_FRAMES: + raise ValueError( + f"video is longer than {cls.MAX_FRAMES / target.fps:.0f} seconds" + ) + k += 1 + if not indices: + raise ValueError("video has no decodable frames") + return indices + + @classmethod + def load_bytes( + cls, + data: bytes, + num_frames: int = -1, + fps: int = 2, + max_duration: int = 300, + frame_recovery: bool = False, + *, + backend: VideoDecoderBackend = "opencv", + **kwargs: Any, + ) -> tuple[DecodedFrames, dict[str, Any]]: + frames, metadata = super().load_bytes( + data, + num_frames=num_frames, + fps=fps, + max_duration=max_duration, + frame_recovery=frame_recovery, + backend=backend, + **kwargs, + ) + if len(frames) % 2: + last = frames[-1:] + if isinstance(frames, torch.Tensor): + frames = torch.cat([frames, last]) + else: + frames = np.concatenate([frames, last]) + metadata["frames_indices"].append(metadata["frames_indices"][-1]) + return frames, metadata diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index d79c1275b6e7..147b65759dea 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -17,6 +17,10 @@ ConfigFormat, get_safetensors_params_metadata, ) +from vllm.transformers_utils.repo_utils import ( + file_or_path_exists, + get_hf_file_to_dict, +) from vllm.utils.torch_utils import common_broadcastable_dtype logger = init_logger(__name__) @@ -643,6 +647,31 @@ def __init__( _strip_mimo_v2_attention_chunk_size(hf_config, hf_text_config) +class Qwen3OmniMoeModelArchConfigConvertor(ModelArchConfigConvertorBase): + def __init__( + self, + hf_config: PreTrainedConfig, + hf_text_config: PreTrainedConfig, + revision: str | None = None, + ): + # Clef-Omni keeps Qwen3-Omni's config and adds a joint schema head; + # route to its wrapper class when the head's config file is present. + # Mutate hf_config (like the deepseek_v4/mimo_v2 rewrites above) + # because model-class resolution reads hf_config.architectures. + if getattr(hf_config, "architectures", None) == [ + "Qwen3OmniMoeForConditionalGeneration" + ]: + model_path = hf_config.name_or_path + if model_path and file_or_path_exists( + model_path, "joint_head_config.json", revision + ): + hf_config.architectures = ["ClefOmniForJointSchema"] + hf_config.joint_head_config = get_hf_file_to_dict( + "joint_head_config.json", model_path, revision + ) + super().__init__(hf_config, hf_text_config, revision) + + class MimoV2MTPModelArchConfigConvertor(ModelArchConfigConvertorBase): def __init__( self, @@ -822,6 +851,7 @@ def derive_max_model_len_and_key(self) -> tuple[float, str | None]: "bailing_hybrid_v3_mtp": BailingHybridV3MTPModelArchConfigConvertor, "pangu_ultra_moe_mtp": PanguUltraMoeMTPModelArchConfigConvertor, "qwen3_5_mtp": Qwen3_5MTPModelArchConfigConvertor, + "qwen3_omni_moe": Qwen3OmniMoeModelArchConfigConvertor, "qwen3_next_mtp": Qwen3NextMTPModelArchConfigConvertor, "qwen4_exp_mtp": Qwen4ExpMTPModelArchConfigConvertor, "RefinedWeb": FalconModelArchConfigConvertor, From 3facf822b771cc0a47a346249d6ea42f9071aaac Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Mon, 12 Oct 2026 14:45:54 +0100 Subject: [PATCH 304/624] [Performance] Load HF tokenizers and processors from a local snapshot (#60825) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Co-authored-by: Kevin H. Luu --- tests/tokenizers_/test_basic.py | 16 +++++ tests/transformers_utils/test_processor.py | 28 ++++++++ vllm/tokenizers/registry.py | 15 ++++- vllm/transformers_utils/processor.py | 17 ++++- vllm/transformers_utils/repo_utils.py | 74 ++++++++++++++++++++++ 5 files changed, 146 insertions(+), 4 deletions(-) diff --git a/tests/tokenizers_/test_basic.py b/tests/tokenizers_/test_basic.py index 4a1a2b156050..a46be988e5c8 100644 --- a/tests/tokenizers_/test_basic.py +++ b/tests/tokenizers_/test_basic.py @@ -1,11 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import _get_protocol_attrs # type: ignore +from unittest.mock import patch import pytest from transformers import ( PreTrainedTokenizerBase, TokenizersBackend, + tokenization_utils_base, ) from vllm.tokenizers import TokenizerLike, get_tokenizer @@ -56,6 +58,20 @@ def test_tokenizer_revision(tokenizer_name: str): get_tokenizer(tokenizer_name, revision="never") +def test_hub_tokenizer_loads_from_local_snapshot(): + """Transformers makes a Hub request per file when given a repo ID.""" + with patch.object( + tokenization_utils_base, + "list_repo_templates", + wraps=tokenization_utils_base.list_repo_templates, + ) as list_repo_templates: + tokenizer = get_tokenizer("openai-community/gpt2") + + assert isinstance(tokenizer, PreTrainedTokenizerBase) + assert tokenizer.name_or_path == "openai-community/gpt2" + list_repo_templates.assert_not_called() + + @pytest.mark.parametrize("tokenizer_name", ["BAAI/bge-base-en"]) @pytest.mark.parametrize("n_tokens", [510]) def test_special_tokens(tokenizer_name: str, n_tokens: int): diff --git a/tests/transformers_utils/test_processor.py b/tests/transformers_utils/test_processor.py index a3a1c7841865..93f742ff743f 100644 --- a/tests/transformers_utils/test_processor.py +++ b/tests/transformers_utils/test_processor.py @@ -2,14 +2,19 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import importlib +from unittest.mock import patch +from transformers import processing_utils from transformers.processing_utils import ProcessingKwargs from typing_extensions import Unpack +from vllm.tokenizers import get_tokenizer from vllm.transformers_utils.processor import ( + get_processor, get_processor_kwargs_keys, get_processor_kwargs_type, ) +from vllm.transformers_utils.repo_utils import _snapshot_non_weight_files class _FakeProcessorKwargs(ProcessingKwargs, total=False): # type: ignore @@ -65,3 +70,26 @@ def test_get_processor_kwargs_from_processor_module_scan_returns_full_union(): proc = _ProcWithoutUnpack() keys = get_processor_kwargs_keys(get_processor_kwargs_type(proc)) _assert_has_all_expected(keys) + + +def test_hub_processor_loads_from_local_snapshot(): + """Transformers makes a Hub request per file when given a repo ID.""" + with patch.object( + processing_utils, + "list_repo_templates", + wraps=processing_utils.list_repo_templates, + ) as list_repo_templates: + processor = get_processor("Qwen/Qwen2-VL-2B-Instruct") + + assert processor.tokenizer.name_or_path == "Qwen/Qwen2-VL-2B-Instruct" + assert processor.chat_template is not None + list_repo_templates.assert_not_called() + + +def test_processor_reuses_tokenizer_snapshot(): + """The chat template lookup must not snapshot the repo a second time.""" + model = "Qwen/Qwen2-VL-2B-Instruct" + _snapshot_non_weight_files.cache_clear() + get_tokenizer(model) + get_processor(model) + assert _snapshot_non_weight_files.cache_info().misses == 1 diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index ebf85a11fd25..a7adf4597f6a 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -4,7 +4,7 @@ from dataclasses import dataclass, field from functools import lru_cache from pathlib import Path -from typing import TYPE_CHECKING, assert_never +from typing import TYPE_CHECKING, assert_never, cast import huggingface_hub from typing_extensions import TypeVar @@ -14,11 +14,12 @@ from vllm.transformers_utils.config import _maybe_register_hf_config, get_config from vllm.transformers_utils.repo_utils import ( any_pattern_in_repo_files, + get_non_weight_snapshot_path, is_mistral_model_repo, ) from vllm.utils.import_utils import resolve_obj_by_qualname -from .hf import CachedHfTokenizer +from .hf import CachedHfTokenizer, HfTokenizer from .protocol import TokenizerLike if TYPE_CHECKING: @@ -164,6 +165,12 @@ def resolve_tokenizer_args( "mistral_format=True is not supported with tokenizer_mode='hf'" ) kwargs["mistral_format"] = False + tokenizer_name = get_non_weight_snapshot_path( + str(tokenizer_name), + revision=revision, + cache_dir=download_dir, + token=kwargs.get("token"), + ) return tokenizer_mode, tokenizer_name, args, kwargs @@ -202,6 +209,7 @@ def get_tokenizer( apply_fastokens_patch() + requested_name = str(tokenizer_name) tokenizer_mode, tokenizer_name, args, kwargs = cached_resolve_tokenizer_args( tokenizer_name, *args, @@ -254,6 +262,9 @@ def get_tokenizer( kwargs.setdefault("config", config) tokenizer = tokenizer_cls_.from_pretrained(tokenizer_name, *args, **kwargs) + # Report the Hub ID rather than the local snapshot it was loaded from + if tokenizer_mode == "hf" and not envs.VLLM_USE_MODELSCOPE: + cast(HfTokenizer, tokenizer).name_or_path = requested_name if model_type in _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: from vllm.tokenizers.hf import get_cached_tokenizer diff --git a/vllm/transformers_utils/processor.py b/vllm/transformers_utils/processor.py index 9aa592084990..1b9b0bd442a7 100644 --- a/vllm/transformers_utils/processor.py +++ b/vllm/transformers_utils/processor.py @@ -22,7 +22,10 @@ from vllm.logger import init_logger from vllm.transformers_utils import processors -from vllm.transformers_utils.repo_utils import get_hf_file_to_dict +from vllm.transformers_utils.repo_utils import ( + get_hf_file_to_dict, + get_non_weight_snapshot_path, +) from vllm.transformers_utils.utils import convert_model_repo_to_path from vllm.utils.func_utils import get_allowed_kwarg_only_overrides @@ -214,7 +217,12 @@ def get_processor( if revision is None: revision = "main" try: - processor_name = convert_model_repo_to_path(processor_name) + processor_name = requested_name = convert_model_repo_to_path(processor_name) + # A snapshot pinned to `revision` would shadow remote code at `code_revision` + if kwargs.get("code_revision") in (None, revision): + processor_name = get_non_weight_snapshot_path( + processor_name, revision=revision, token=kwargs.get("token") + ) registered_cls_name = get_processor_cls_name_from_config( processor_name, revision=revision ) @@ -269,6 +277,11 @@ def get_processor( f"found type: {type(processor)}" ) + # Report the Hub ID rather than the local snapshot it was loaded from + tokenizer = getattr(processor, "tokenizer", processor) + if hasattr(tokenizer, "name_or_path"): + tokenizer.name_or_path = requested_name + return processor diff --git a/vllm/transformers_utils/repo_utils.py b/vllm/transformers_utils/repo_utils.py index b6a23b007aa2..43ef07676e3a 100644 --- a/vllm/transformers_utils/repo_utils.py +++ b/vllm/transformers_utils/repo_utils.py @@ -304,6 +304,80 @@ def get_model_path(model: str | Path, revision: str | None = None): ) +_WEIGHT_FILE_PATTERNS = [ + "*.safetensors", + "*.bin", + "*.pt", + "*.pth", + "*.ckpt", + "*.gguf", + "*.h5", + "*.msgpack", + "*.onnx", + "*.onnx_data", + # Large non-weight assets that tokenizers and processors never read + "*.png", + "*.jpg", + "*.jpeg", + "*.gif", + "*.webp", + "*.mp3", + "*.mp4", + "*.wav", + "*.flac", + "*.parquet", + "*.arrow", + "*.ipynb", + "*.zip", + "*.tar*", +] + + +@cache +def _snapshot_non_weight_files( + repo_id: str, + revision: str | None, + cache_dir: str | None, + token: str | bool | None, +) -> str: + # Avoid circular import + from vllm.model_executor.model_loader.weight_utils import get_lock + + with get_lock(repo_id, cache_dir): + return hf_api().snapshot_download( + repo_id=repo_id, + revision=revision, + cache_dir=cache_dir, + token=token, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + ignore_patterns=_WEIGHT_FILE_PATTERNS, + ) + + +def get_non_weight_snapshot_path( + repo_id: str, + revision: str | None = None, + cache_dir: str | None = None, + token: str | bool | None = None, +) -> str: + """Download all non-weight files of a Hub repo and return the local snapshot. + + Loading tokenizers and processors from a local directory stops Transformers + from making a Hub request per file (including the `additional_chat_templates` + listing). Returns `repo_id` unchanged if it is local or cannot be fetched, so + Transformers handles it as before. Failures are not cached, so they are retried. + """ + if envs.VLLM_USE_MODELSCOPE or Path(repo_id).exists(): + return repo_id + # Normalise so that `None` and "main" share a cache entry + revision = revision or huggingface_hub.constants.DEFAULT_REVISION + try: + return _snapshot_non_weight_files(repo_id, revision, cache_dir, token) + except (HfHubHTTPError, OSError, ValueError): + logger.debug("Failed to snapshot %s", repo_id, exc_info=True) + return repo_id + + def _try_download_from_hf_hub( model: str | Path, file_name: str, revision: str | None ) -> Path | None: From 4c36d21280824d1a9c521106932408243078de0d Mon Sep 17 00:00:00 2001 From: Suraj Sharan Date: Mon, 12 Oct 2026 17:49:30 +0400 Subject: [PATCH 305/624] [Bugfix][LoRA] Apply PEFT rank_pattern / alpha_pattern scaling per module (#59801) Signed-off-by: suraj Co-authored-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Co-authored-by: linitra24 --- tests/lora/test_lora_weights.py | 57 +++++++++++++++++ tests/lora/test_peft_helper.py | 105 ++++++++++++++++++++++++++++++++ vllm/lora/lora_model.py | 4 +- vllm/lora/lora_weights.py | 38 +++++++++--- vllm/lora/peft_helper.py | 60 ++++++++++++++++-- 5 files changed, 249 insertions(+), 15 deletions(-) diff --git a/tests/lora/test_lora_weights.py b/tests/lora/test_lora_weights.py index ff132147933e..7230126e40c8 100644 --- a/tests/lora/test_lora_weights.py +++ b/tests/lora/test_lora_weights.py @@ -4,6 +4,7 @@ import torch from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights +from vllm.lora.peft_helper import PEFTHelper RANK, ALPHA = 16, 32 # PEFTHelper stores alpha/sqrt(r) for use_rslora adapters. @@ -50,3 +51,59 @@ def test_non_gated_moe_keeps_w3_unscaled(): ) assert packed.scaling[:2] == pytest.approx([RSLORA_SCALING] * 2) assert packed.scaling[2] == 1.0 + + +@pytest.mark.parametrize( + "packer,stacked", + [ + (PackedLoRALayerWeights.pack_moe, False), + (PackedLoRALayerWeights.pack_moe_stacked, True), + ], +) +def test_moe_packing_keeps_per_projection_scaling(packer, stacked: bool): + # e.g. rank_pattern / alpha_pattern giving gate, down and up different scalings + loras = moe_loras(1.0, stacked) + for lora, scaling in zip(loras, (2.0, 4.0, 8.0)): + lora.scaling = scaling + packed = packer(loras, "experts") + assert packed.scaling == pytest.approx([2.0, 4.0, 8.0]) + + +def test_moe_packing_keeps_per_expert_scaling(): + """Experts of one projection can have different scaling, e.g. from + alpha_pattern.""" + loras = moe_loras(RSLORA_SCALING) + moe_loras(RSLORA_SCALING) + loras[4].scaling = 2 * RSLORA_SCALING # w2 of expert 1 + expected = [lora.lora_b * lora.scaling for lora in loras] + packed = PackedLoRALayerWeights.pack_moe(loras, "experts").optimize() + for i, lora_b in enumerate(packed.lora_b): + torch.testing.assert_close(lora_b, torch.stack(expected[i::3])) + + +def test_moe_packing_rejects_per_expert_rank(): + """Experts of one projection must share a rank to be stacked.""" + loras = moe_loras(1.0) + moe_loras(1.0) + loras[4].lora_a = torch.randn(RANK // 2, 4) + loras[4].lora_b = torch.randn(6, RANK // 2) + with pytest.raises(ValueError, match="different LoRA ranks"): + PackedLoRALayerWeights.pack_moe(loras, "experts") + + +def test_from_config_uses_checkpoint_module_name(): + peft_helper = PEFTHelper( + r=RANK, + lora_alpha=ALPHA, + target_modules=["q_proj"], + rank_pattern={r"^model\.layers\.0\.self_attn\.q_proj": 4}, + ) + # the pattern refers to the checkpoint name, not the name after weights mapping + lora = LoRALayerWeights.from_config( + "language_model.model.layers.0.self_attn.q_proj", + peft_helper, + "model.layers.0.self_attn.q_proj", + ) + assert lora.rank == 4 + assert lora.scaling == pytest.approx(ALPHA / 4) + lora = LoRALayerWeights.from_config("model.layers.1.self_attn.q_proj", peft_helper) + assert lora.rank == RANK + assert lora.scaling == pytest.approx(ALPHA / RANK) diff --git a/tests/lora/test_peft_helper.py b/tests/lora/test_peft_helper.py index bb1246144b0d..2fcf5ff96205 100644 --- a/tests/lora/test_peft_helper.py +++ b/tests/lora/test_peft_helper.py @@ -22,6 +22,11 @@ {"modules_to_save": ["lm_head"]}, "Unsupported modules_to_save", ), + ( + "test_rank_pattern", + {"r": 8, "rank_pattern": {"q_proj": 32}}, + "LoRA rank 32 is greater than max_lora_rank", + ), ("test_rank_zero", {"r": 0}, "must be a positive integer"), ("test_rank_negative", {"r": -8}, "must be a positive integer"), ("test_lora_bias", {"lora_bias": True}, "does not support LoRA bias"), @@ -153,3 +158,103 @@ def test_peft_helper_init_lora_weights_supported(init_lora_weights): "use_bdlora": None, } ).validate_legal(lora_config) + + +@pytest.mark.parametrize( + "config,module_name,expected", + [ + # rank_pattern only: alpha / r_m + ({"rank_pattern": {"q_proj": 4}}, "model.layers.0.self_attn.q_proj", (4, 8.0)), + ({"rank_pattern": {"q_proj": 4}}, "model.layers.0.self_attn.v_proj", (16, 2.0)), + # alpha_pattern only: alpha_m / r + ( + {"alpha_pattern": {"v_proj": 64}}, + "model.layers.0.self_attn.v_proj", + (16, 4.0), + ), + # keys only match at a "." boundary, like in PEFT + ({"rank_pattern": {"proj": 4}}, "model.layers.0.self_attn.q_proj", (16, 2.0)), + # fused MoE experts from PEFT `target_parameters`: gate_up_proj is stored + # as `experts.base_layer`, down_proj as `experts` + ( + {"rank_pattern": {"experts.gate_up_proj": 2}}, + "model.layers.0.mlp.experts.base_layer", + (2, 16.0), + ), + ( + {"rank_pattern": {"experts.down_proj": 2}}, + "model.layers.0.mlp.experts", + (2, 16.0), + ), + ( + {"rank_pattern": {"experts.gate_up_proj": 2}}, + "model.layers.0.mlp.experts", + (16, 2.0), + ), + ], +) +def test_peft_helper_rank_and_alpha_pattern(config, module_name, expected): + peft_helper = PEFTHelper( + r=16, lora_alpha=32, target_modules=["q_proj", "v_proj"], **config + ) + rank, scaling = peft_helper.get_rank_and_scaling(module_name) + assert rank == expected[0] + assert scaling == pytest.approx(expected[1]) + + +def test_peft_helper_pattern_rslora(): + peft_helper = PEFTHelper( + r=16, + lora_alpha=32, + target_modules=["q_proj"], + use_rslora=True, + rank_pattern={"q_proj": 4}, + ) + assert peft_helper.get_rank_and_scaling("model.q_proj") == (4, 32 / math.sqrt(4)) + + +def test_peft_helper_dynamic_rank_conversion(): + """PEFT's `save_as_lora` with a dynamic rank writes r=1, lora_alpha=1 and the + same per-module values in rank_pattern and alpha_pattern, so every module + keeps a scaling of 1.""" + peft_helper = PEFTHelper( + r=1, + lora_alpha=1, + target_modules=["q_proj", "v_proj"], + rank_pattern={"layers.0.self_attn.q_proj": 32, "layers.0.self_attn.v_proj": 3}, + alpha_pattern={"layers.0.self_attn.q_proj": 32, "layers.0.self_attn.v_proj": 3}, + ) + assert peft_helper.get_rank_and_scaling("model.layers.0.self_attn.q_proj") == ( + 32, + 1.0, + ) + assert peft_helper.get_rank_and_scaling("model.layers.0.self_attn.v_proj") == ( + 3, + 1.0, + ) + with pytest.raises(ValueError, match="LoRA rank 32 is greater than max_lora_rank"): + peft_helper.validate_legal(LoRAConfig(max_lora_rank=16)) + + +def test_peft_helper_null_patterns(): + peft_helper = PEFTHelper.from_dict( + { + "r": 8, + "lora_alpha": 16, + "target_modules": ["q_proj"], + "rank_pattern": None, + "alpha_pattern": None, + } + ) + assert peft_helper.get_rank_and_scaling("model.q_proj") == (8, 2.0) + + +@pytest.mark.parametrize("bad_rank", [0, -4, 2.5]) +def test_peft_helper_invalid_rank_pattern(bad_rank): + with pytest.raises(ValueError, match="rank_pattern"): + PEFTHelper( + r=8, + lora_alpha=16, + target_modules=["q_proj"], + rank_pattern={"q_proj": bad_rank}, + ) diff --git a/vllm/lora/lora_model.py b/vllm/lora/lora_model.py index 196445332443..f4520b12c75b 100644 --- a/vllm/lora/lora_model.py +++ b/vllm/lora/lora_model.py @@ -150,8 +150,10 @@ def from_lora_tensors( continue if module_name not in loras: + # rank/alpha patterns refer to the names before weights mapping + peft_module_name, _ = parse_fine_tuned_lora_name(tensor_name) loras[module_name] = LoRALayerWeights.from_config( - module_name, peft_helper + module_name, peft_helper, peft_module_name ) if is_lora_a: diff --git a/vllm/lora/lora_weights.py b/vllm/lora/lora_weights.py index 13a692a62353..9ea7c550d703 100644 --- a/vllm/lora/lora_weights.py +++ b/vllm/lora/lora_weights.py @@ -59,15 +59,22 @@ def from_config( cls, module_name: str, peft_helper: PEFTHelper, + peft_module_name: str | None = None, ) -> "LoRALayerWeights": + """`peft_module_name` is the module name in the adapter checkpoint, used + to resolve `rank_pattern` / `alpha_pattern`. Defaults to `module_name`. + """ + rank, scaling = peft_helper.get_rank_and_scaling( + peft_module_name or module_name + ) # lora_a and lora_b are set to None for config-based construction return cls( module_name, - peft_helper.r, + rank, peft_helper.lora_alpha, None, None, - peft_helper.vllm_lora_scaling_factor, + scaling, ) @classmethod @@ -168,6 +175,18 @@ def pack_moe( rank = first_lora.rank lora_alpha = first_lora.lora_alpha assert len(loras) % 3 == 0 + for proj_loras in (loras[0::3], loras[1::3], loras[2::3]): + experts = [lora for lora in proj_loras if lora is not None] + if len({lora.lora_a.shape[0] for lora in experts}) > 1: + raise ValueError( + f"The experts of {module_name} have different LoRA ranks for " + "the same projection, which is not supported." + ) + # The packed LoRA has one scaling per projection, so fold it into + # lora_b where the experts differ (e.g. from alpha_pattern). + if len({lora.scaling for lora in experts}) > 1: + for lora in experts: + lora.optimize() w1_lora_a_lst = [] w2_lora_a_lst = [] w3_lora_a_lst = [] @@ -201,10 +220,12 @@ def pack_moe( w1_lora_b = torch.stack(w1_lora_b_lst, dim=0) # (num_experts,output_size,rank) w2_lora_b = torch.stack(w2_lora_b_lst, dim=0) - # All w1, w2, w3 have the same scaling factor. Use the per-adapter - # scaling (e.g. alpha/sqrt(r) for rsLoRA) instead of alpha/rank. - scaling = first_lora.scaling - last_scaling = scaling + # Use each projection's own scaling (e.g. alpha/sqrt(r) for rsLoRA, or + # per-module values from rank_pattern / alpha_pattern) instead of + # alpha/rank. + w1_scaling = loras[0].scaling # type: ignore[union-attr] + w2_scaling = loras[1].scaling # type: ignore[union-attr] + last_scaling = loras[2].scaling if loras[2] is not None else w1_scaling if is_non_gated_moe: # For non-gated MoE, reuse w1 tensors for w3 to avoid memory waste @@ -224,7 +245,7 @@ def pack_moe( [lora_alpha, lora_alpha, lora_alpha], [w1_lora_a, w2_lora_a, w3_lora_a], [w1_lora_b, w2_lora_b, w3_lora_b], - scaling=[scaling, scaling, last_scaling], + scaling=[w1_scaling, w2_scaling, last_scaling], ) return obj @@ -251,14 +272,13 @@ def pack_moe_stacked( assert w1_lora is not None and w2_lora is not None and w3_lora is not None rank = w1_lora.rank lora_alpha = w1_lora.lora_alpha - scaling = w1_lora.scaling return cls( module_name, rank, [lora_alpha, lora_alpha, lora_alpha], [w1_lora.lora_a, w2_lora.lora_a, w3_lora.lora_a], [w1_lora.lora_b, w2_lora.lora_b, w3_lora.lora_b], - scaling=[scaling, scaling, scaling], + scaling=[w1_lora.scaling, w2_lora.scaling, w3_lora.scaling], ) def optimize(self) -> "PackedLoRALayerWeights": diff --git a/vllm/lora/peft_helper.py b/vllm/lora/peft_helper.py index de05dea1cfc1..721b7a039f00 100644 --- a/vllm/lora/peft_helper.py +++ b/vllm/lora/peft_helper.py @@ -9,6 +9,8 @@ from dataclasses import MISSING, dataclass, field, fields from typing import Literal +import regex as re + from vllm.config.lora import LoRAConfig from vllm.logger import init_logger @@ -41,6 +43,9 @@ class PEFTHelper: layer_replication: list[list[int]] | None = field(default=None) use_bdlora: dict | None = field(default=None) use_qalora: bool = field(default=False) + # Per-module overrides of `r` and `lora_alpha`, keyed by module name pattern + rank_pattern: dict[str, int] = field(default_factory=dict) + alpha_pattern: dict[str, float] = field(default_factory=dict) # Extra vllm field, start with 'vllm_' to avoid conflict vllm_lora_scaling_factor: float = field(default=1.0) vllm_max_position_embeddings: int | None = field(default=False) @@ -97,11 +102,55 @@ def _init_is_supported(self) -> bool: def __post_init__(self): if self.r <= 0: raise ValueError(f"LoRA rank `r` must be a positive integer, got {self.r}.") + self.rank_pattern = self.rank_pattern or {} + self.alpha_pattern = self.alpha_pattern or {} + for key, rank in self.rank_pattern.items(): + if not isinstance(rank, int) or rank <= 0: + raise ValueError( + f"LoRA rank for `{key}` in `rank_pattern` must be a positive " + f"integer, got {rank}." + ) if self.use_rslora: logger.info_once("Loading LoRA weights trained with rsLoRA.") - self.vllm_lora_scaling_factor = self.lora_alpha / math.sqrt(self.r) - else: - self.vllm_lora_scaling_factor = self.lora_alpha / self.r + self.vllm_lora_scaling_factor = self._scaling(self.r, self.lora_alpha) + + def _scaling(self, rank: int, alpha: float) -> float: + if self.use_rslora: + return alpha / math.sqrt(rank) + return alpha / rank + + @staticmethod + def _match_pattern(pattern: dict, name: str): + # Same rule as PEFT's `get_pattern_key`: a key matches if it is a suffix + # of the module name starting at a "." boundary; the first match wins. + for key, value in pattern.items(): + if re.match(rf"(.*\.)?({key})$", name): + return value + return None + + def get_rank_and_scaling(self, module_name: str) -> tuple[int, float]: + """Return the rank and scaling PEFT uses for `module_name`, taking + `rank_pattern` and `alpha_pattern` into account. + + `module_name` is the name in the adapter checkpoint, without the + `base_model.model.` prefix and before any weights mapping. + """ + if not self.rank_pattern and not self.alpha_pattern: + return self.r, self.vllm_lora_scaling_factor + # PEFT stores fused MoE expert LoRAs (`target_parameters`) as + # `experts.base_layer` (gate_up_proj) and `experts` (down_proj), but + # matches the patterns against the parameter names. + if module_name.endswith(".experts.base_layer"): + module_name = module_name.removesuffix(".base_layer") + ".gate_up_proj" + elif module_name.endswith(".experts"): + module_name += ".down_proj" + rank = self._match_pattern(self.rank_pattern, module_name) + if rank is None: + rank = self.r + alpha = self._match_pattern(self.alpha_pattern, module_name) + if alpha is None: + alpha = self.lora_alpha + return rank, self._scaling(rank, alpha) @classmethod def from_dict(cls, config_dict: dict) -> "PEFTHelper": @@ -141,9 +190,10 @@ def validate_legal(self, lora_config: LoRAConfig) -> None: constraints and requirements. """ error_msg = self._validate_features() - if self.r > lora_config.max_lora_rank: + max_rank = max([self.r, *self.rank_pattern.values()]) + if max_rank > lora_config.max_lora_rank: error_msg.append( - f"LoRA rank {self.r} is greater than max_lora_rank" + f"LoRA rank {max_rank} is greater than max_lora_rank" f" {lora_config.max_lora_rank}." ) if self.bias != "none": From f022975964e27ed18dd8e437eed0ab70b0a00b33 Mon Sep 17 00:00:00 2001 From: bzsuni <86399306+bzsuni@users.noreply.github.com> Date: Mon, 12 Oct 2026 21:49:46 +0800 Subject: [PATCH 306/624] [Bugfix] Reject beam search in completion render (#61212) Signed-off-by: bzsuni --- .../scale_out/render/test_render.py | 36 ++++++++++++++++--- vllm/entrypoints/scale_out/render/serving.py | 6 ++++ 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/tests/entrypoints/scale_out/render/test_render.py b/tests/entrypoints/scale_out/render/test_render.py index 6c11a518d764..308555be362e 100644 --- a/tests/entrypoints/scale_out/render/test_render.py +++ b/tests/entrypoints/scale_out/render/test_render.py @@ -12,6 +12,7 @@ import pytest_asyncio from tests.utils import RemoteLaunchRenderServer +from vllm.entrypoints.openai.completion.protocol import CompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.entrypoints.scale_out.render.api_router import router from vllm.entrypoints.scale_out.render.serving import ServingRender @@ -21,7 +22,7 @@ MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM" -def _build_responses_serving_render() -> ServingRender: +def _build_serving_render() -> ServingRender: serving = ServingRender.__new__(ServingRender) serving.model_config = SimpleNamespace( max_model_len=100, @@ -43,6 +44,31 @@ def _build_responses_serving_render() -> ServingRender: return serving +@pytest.mark.asyncio +@pytest.mark.skip_global_cleanup +async def test_render_completion_rejects_beam_search(): + serving = _build_serving_render() + serving.online_renderer.render_completion = AsyncMock( + return_value=[{"prompt_token_ids": [7, 8, 9]}] + ) + request = CompletionRequest( + model=MODEL_NAME, + prompt="Test prompt", + use_beam_search=True, + n=2, + length_penalty=0.7, + ) + + response = await serving.render_completion_request(request) + + assert isinstance(response, ErrorResponse) + assert response.error.code == 400 + assert ( + response.error.message == "Beam search is not supported by the render endpoint" + ) + serving.online_renderer.render_completion.assert_not_awaited() + + @pytest.mark.skip_global_cleanup def test_responses_render_route_is_registered(): assert any(route.path == "/v1/responses/render" for route in router.routes) @@ -60,7 +86,7 @@ def test_responses_render_route_documents_not_implemented(): @pytest.mark.asyncio @pytest.mark.skip_global_cleanup async def test_render_responses_returns_generate_request_without_stored_state(): - serving = _build_responses_serving_render() + serving = _build_serving_render() serving.online_renderer.render_responses = AsyncMock( return_value=MagicMock( messages=[{"role": "user", "content": "Test prompt"}], @@ -101,7 +127,7 @@ async def test_render_responses_returns_generate_request_without_stored_state(): @pytest.mark.asyncio @pytest.mark.skip_global_cleanup async def test_render_responses_rejects_previous_response_id(): - serving = _build_responses_serving_render() + serving = _build_serving_render() serving.online_renderer.render_responses = AsyncMock() request = ResponsesRequest( model=MODEL_NAME, @@ -120,7 +146,7 @@ async def test_render_responses_rejects_previous_response_id(): @pytest.mark.asyncio @pytest.mark.skip_global_cleanup async def test_render_responses_rejects_empty_token_ids(): - serving = _build_responses_serving_render() + serving = _build_serving_render() serving.online_renderer.render_responses = AsyncMock( return_value=MagicMock( messages=[], @@ -141,7 +167,7 @@ async def test_render_responses_rejects_empty_token_ids(): async def test_render_responses_carries_reasoning_parser_kwargs(): """/v1/responses gives the engine's reasoning parser the effective template kwargs, so the rendered request must carry them to generate.""" - serving = _build_responses_serving_render() + serving = _build_serving_render() online = serving.online_renderer online.parser = MagicMock(reasoning_parser_cls=object) online.chat_template = None diff --git a/vllm/entrypoints/scale_out/render/serving.py b/vllm/entrypoints/scale_out/render/serving.py index 84771c5f7390..5caf78bffe9e 100644 --- a/vllm/entrypoints/scale_out/render/serving.py +++ b/vllm/entrypoints/scale_out/render/serving.py @@ -186,6 +186,12 @@ async def render_completion_request( error_check_ret = await self._check_model(request) if error_check_ret is not None: return error_check_ret + + if request.use_beam_search: + return self.create_error_response( + "Beam search is not supported by the render endpoint" + ) + result = await self.online_renderer.render_completion( request, skip_mm_cache=True ) From b3ee560ae816b763b09b1aeddf9adca95a0014d0 Mon Sep 17 00:00:00 2001 From: Ilya Markov Date: Mon, 12 Oct 2026 16:02:56 +0200 Subject: [PATCH 307/624] [Bugfix][MoE] Apply expert_map remapping for CUTLASS FP4 MOE expert parallel (#59042) Signed-off-by: Markov Ilya --- tests/kernels/moe/test_cutlass_moe.py | 203 +++++++++++++++++- tests/kernels/moe/test_ocp_mx_moe.py | 111 ++++++++++ .../layers/fused_moe/experts/cutlass_moe.py | 45 +++- .../experts/gpt_oss_triton_kernels_moe.py | 48 +---- .../layers/fused_moe/oracle/mxfp4.py | 55 +++++ vllm/model_executor/layers/fused_moe/utils.py | 67 ++++++ 6 files changed, 470 insertions(+), 59 deletions(-) diff --git a/tests/kernels/moe/test_cutlass_moe.py b/tests/kernels/moe/test_cutlass_moe.py index 050e1bc10487..a06abfbc1fad 100644 --- a/tests/kernels/moe/test_cutlass_moe.py +++ b/tests/kernels/moe/test_cutlass_moe.py @@ -9,7 +9,7 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from tests.kernels.moe.utils import make_dummy_moe_config +from tests.kernels.moe.utils import make_dummy_moe_config, make_test_quant_config from vllm import _custom_ops as ops from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe import ( @@ -26,8 +26,10 @@ ) from vllm.model_executor.layers.fused_moe.config import ( FUSED_MOE_UNQUANTIZED_CONFIG, + FusedMoEParallelConfig, FusedMoEQuantConfig, fp8_w8a8_moe_quant_config, + mxfp4_moe_quant_config, ) from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( CutlassExpertsFp4, @@ -41,7 +43,10 @@ ) from vllm.model_executor.layers.fused_moe.oracle import mxfp4 as mxfp4_oracle from vllm.model_executor.layers.fused_moe.oracle import nvfp4 as nvfp4_oracle -from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input +from vllm.model_executor.layers.fused_moe.utils import ( + moe_kernel_quantize_input, + remap_topk_to_local, +) from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -183,6 +188,189 @@ def test_cutlass_moe_permutation_maps_padding_to_zero(): torch.testing.assert_close(restored[3], torch.zeros_like(restored[3])) +def test_remap_topk_ids_to_local(): + device = "cuda" + expert_map = torch.tensor([-1, 0, -1, 1, -1, 2], dtype=torch.int32, device=device) + topk_ids = torch.tensor([[1, 3], [0, 5], [-1, 4]], dtype=torch.int32, device=device) + expected = torch.tensor( + [[0, 1], [-1, 2], [-1, -1]], dtype=torch.int32, device=device + ) + actual = remap_topk_to_local(topk_ids, expert_map, out_dtype=torch.int32) + torch.testing.assert_close(actual, expected) + + +def test_remap_topk_ids_to_local_requires_contiguous_topk_ids(): + expert_map = torch.tensor([0, 1], dtype=torch.int32) + topk_ids = torch.tensor([[0, 1], [1, 0]], dtype=torch.int32).t() + assert not topk_ids.is_contiguous() + with pytest.raises(AssertionError): + remap_topk_to_local(topk_ids, expert_map) + + +def test_cutlass_fp4_moe_supports_expert_parallel_config(): + ep_config = FusedMoEParallelConfig( + tp_size=1, + tp_rank=0, + pcp_size=1, + pcp_rank=0, + dp_size=8, + dp_rank=0, + ep_size=8, + ep_rank=0, + sp_size=1, + use_ep=True, + all2all_backend="allgather_reducescatter", + enable_eplb=False, + ) + assert CutlassExpertsMxfp4._supports_parallel_config(ep_config) + assert CutlassExpertsFp4._supports_parallel_config(ep_config) + + +def _slice_moe_quant_config( + quant_config: FusedMoEQuantConfig, start: int, end: int +) -> FusedMoEQuantConfig: + num_experts = quant_config.w1_scale.shape[0] + sliced_config = copy.deepcopy(quant_config) + for desc_name in ("_a1", "_a2", "_w1", "_w2"): + desc = getattr(sliced_config, desc_name) + for field_name in ("scale", "alpha_or_gscale", "bias", "zp"): + tensor = getattr(desc, field_name) + if ( + isinstance(tensor, torch.Tensor) + and tensor.ndim > 0 + and tensor.shape[0] == num_experts + ): + setattr(desc, field_name, tensor[start:end]) + return sliced_config + + +def _swizzle_mxfp4_moe_scales( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + e, w1_n, w1_scale_k = w1_scale.shape + w1_k = w1_scale_k * 32 + w2_m, w2_scale_n = w2_scale.shape[1], w2_scale.shape[2] + w2_n = w2_scale_n * 32 + + swizzled_w1 = [] + swizzled_w2 = [] + for expert in range(e): + swizzled_w1.append( + swizzle_mxfp4_scales(w1_scale[expert], w1_n, w1_k).reshape(w1_n, w1_scale_k) + ) + swizzled_w2.append( + swizzle_mxfp4_scales(w2_scale[expert], w2_m, w2_n).reshape(w2_m, w2_scale_n) + ) + return torch.stack(swizzled_w1), torch.stack(swizzled_w2) + + +def _make_cutlass_fp4_moe_weights( + quantization: str, + num_experts: int, + n: int, + k: int, + dtype: torch.dtype, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, FusedMoEQuantConfig]: + if quantization == "nvfp4": + return make_test_quant_config( + num_experts, + n, + k, + in_dtype=dtype, + quant_dtype="nvfp4", + block_shape=None, + per_act_token_quant=False, + make_gate=True, + ) + + from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + mxfp4_quantize, + ) + + w1, w1_scale = mxfp4_quantize( + torch.randn((num_experts, 2 * n, k), device=device, dtype=dtype) / 8 + ) + w2, w2_scale = mxfp4_quantize( + torch.randn((num_experts, k, n), device=device, dtype=dtype) / 8 + ) + w1_scale, w2_scale = _swizzle_mxfp4_moe_scales(w1_scale, w2_scale) + return ( + w1.contiguous(), + w2.contiguous(), + mxfp4_moe_quant_config(w1_scale=w1_scale, w2_scale=w2_scale), + ) + + +@pytest.mark.parametrize("quantization", ["nvfp4", "mxfp4"]) +@torch.inference_mode() +def test_cutlass_fp4_moe_EP( + quantization: str, + workspace_init, +): + """Sharded CUTLASS FP4/MXFP4 MoE with expert_map must match full weights.""" + experts_cls = CutlassExpertsFp4 if quantization == "nvfp4" else CutlassExpertsMxfp4 + if not experts_cls._supports_current_device(): + pytest.skip(f"CUTLASS {quantization} MoE is not supported on this GPU") + + set_random_seed(7) + m, n, k = 8, 128, 128 + num_experts, num_local_experts, topk = 8, 2, 2 + dtype = torch.bfloat16 + device = torch.device("cuda") + + w1, w2, quant_config = _make_cutlass_fp4_moe_weights( + quantization, num_experts, n, k, dtype, device + ) + hidden_states = torch.randn((m, k), device=device, dtype=dtype) / 10 + score = torch.randn((m, num_experts), device=device, dtype=dtype) + topk_weights, topk_ids, _ = fused_topk( + hidden_states, score, topk, renormalize=False + ) + + moe_config = make_dummy_moe_config( + num_experts=num_experts, + num_local_experts=num_experts, + experts_per_token=topk, + hidden_dim=k, + intermediate_size=n, + in_dtype=dtype, + max_num_tokens=m, + ) + kernel = mk.FusedMoEKernel( + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), + experts_cls(moe_config=moe_config, quant_config=quant_config), + ) + apply_kwargs = { + "hidden_states": hidden_states, + "w1": w1, + "w2": w2, + "topk_weights": topk_weights, + "topk_ids": topk_ids, + "global_num_experts": num_experts, + "activation": MoEActivation.SILU, + "expert_map": None, + "apply_router_weight_on_input": False, + } + reference = kernel.apply(**apply_kwargs) + sharded = run_with_expert_maps( + num_experts, + num_local_experts, + quant_config, + experts_cls, + **apply_kwargs, + ) + + assert reference.abs().max() > 0 + torch.testing.assert_close(sharded, reference, atol=1e-2, rtol=1e-2) + + @pytest.mark.parametrize("quantization", ["nvfp4", "mxfp4"]) @torch.inference_mode() def test_cutlass_fp4_moe_padded_routes_do_not_change_valid_output(quantization: str): @@ -486,6 +674,7 @@ def run_with_expert_maps( num_experts: int, num_local_experts: int, quant_config: FusedMoEQuantConfig, + experts_cls: type[mk.FusedMoEExperts], **cutlass_moe_kwargs, ): def slice_experts(): @@ -513,9 +702,7 @@ def slice_experts(): for k, t in full_tensors.items(): cutlass_moe_kwargs[k] = t[s:e] - new_quant_config = copy.deepcopy(quant_config) - new_quant_config._w1.scale = quant_config.w1_scale[s:e] - new_quant_config._w2.scale = quant_config.w2_scale[s:e] + new_quant_config = _slice_moe_quant_config(quant_config, s, e) yield cutlass_moe_kwargs, new_quant_config @@ -523,13 +710,14 @@ def slice_experts(): for kwargs, new_quant_config in slice_experts(): w2 = kwargs["w2"] a = kwargs["hidden_states"] + w1 = kwargs["w1"] moe_config = make_dummy_moe_config( max_num_tokens=kwargs.get("hidden_states").shape[0], experts_per_token=kwargs.get("topk_ids").shape[1], num_experts=num_experts, num_local_experts=num_local_experts, hidden_dim=w2.shape[1], - intermediate_size=w2.shape[2], + intermediate_size=w1.shape[1] // 2, in_dtype=a.dtype, ) kernel = mk.FusedMoEKernel( @@ -539,7 +727,7 @@ def slice_experts(): allow_new_interface=True, use_monolithic=False, ), - CutlassExpertsFp8( + experts_cls( moe_config=moe_config, quant_config=new_quant_config, ), @@ -624,6 +812,7 @@ def run_8_bit( num_experts, num_local_experts, quant_config, + CutlassExpertsFp8, **kwargs, ) diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 13f29087d951..0a8f5c29d034 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -2114,6 +2114,117 @@ def test_aiter_mxfp4_monolithic_rejects_expert_parallel( assert "parallel config" in reason +@pytest.mark.parametrize( + "all2all_backend", + [ + "deepep_high_throughput", + "flashinfer_nvlink_two_sided", + "flashinfer_nvlink_one_sided", + ], +) +def test_make_mxfp4_moe_kernel_rejects_cutlass_unsupported_expert_parallel( + monkeypatch: pytest.MonkeyPatch, + all2all_backend: str, +): + """Compressed-tensors CUTLASS W4A4 bypasses the oracle; guard at assembly.""" + from tests.kernels.moe.utils import make_dummy_moe_config + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEParallelConfig, + mxfp4_moe_quant_config, + ) + from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( + CutlassExpertsMxfp4, + ) + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + Mxfp4MoeBackend, + make_mxfp4_moe_kernel, + ) + + monkeypatch.setattr( + CutlassExpertsMxfp4, "_supports_current_device", staticmethod(lambda: True) + ) + + moe_config = replace( + make_dummy_moe_config( + num_experts=256, + num_local_experts=32, + hidden_dim=4096, + intermediate_size=2048, + ), + moe_parallel_config=FusedMoEParallelConfig( + tp_size=1, + tp_rank=0, + pcp_size=1, + pcp_rank=0, + dp_size=8, + dp_rank=0, + ep_size=8, + ep_rank=0, + sp_size=1, + use_ep=True, + all2all_backend=all2all_backend, + enable_eplb=False, + ), + ) + w1_scale = torch.empty(32, 64, 128, dtype=torch.uint8) + w2_scale = torch.empty(32, 4096, 64, dtype=torch.uint8) + quant_config = mxfp4_moe_quant_config(w1_scale=w1_scale, w2_scale=w2_scale) + + with pytest.raises(ValueError, match="parallel config"): + make_mxfp4_moe_kernel( + quant_config, + moe_config, + CutlassExpertsMxfp4, + Mxfp4MoeBackend.MARLIN, + ) + + +def test_make_mxfp4_moe_kernel_validates_cutlass_expert_parallel( + monkeypatch: pytest.MonkeyPatch, +): + """CUTLASS MXFP4 supports EP with allgather_reducescatter after expert_map.""" + from tests.kernels.moe.utils import make_dummy_moe_config + from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig + from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( + CutlassExpertsMxfp4, + ) + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + Mxfp4MoeBackend, + _validate_mxfp4_experts_config, + ) + + monkeypatch.setattr( + CutlassExpertsMxfp4, "_supports_current_device", staticmethod(lambda: True) + ) + + moe_config = replace( + make_dummy_moe_config( + num_experts=256, + num_local_experts=32, + hidden_dim=4096, + intermediate_size=2048, + ), + moe_parallel_config=FusedMoEParallelConfig( + tp_size=1, + tp_rank=0, + pcp_size=1, + pcp_rank=0, + dp_size=8, + dp_rank=0, + ep_size=8, + ep_rank=0, + sp_size=1, + use_ep=True, + all2all_backend="allgather_reducescatter", + enable_eplb=False, + ), + ) + + _validate_mxfp4_experts_config( + CutlassExpertsMxfp4, moe_config, Mxfp4MoeBackend.MARLIN + ) + + # Every activation-quantizing OCP MX scheme must map to a `quant_dtype` that # `moe_kernel_quantize_input` actually dispatches on. Its final `else` returns # the activation untouched, so a name it does not know (e.g. "mxfp6" instead of diff --git a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py index 6f9a62f282ce..5c6a96f26daa 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py @@ -31,6 +31,7 @@ ) from vllm.model_executor.layers.fused_moe.utils import ( _resize_cache, + remap_topk_to_local, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -412,10 +413,9 @@ def activation_format() -> mk.FusedMoEActivationFormat: @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - # CutlassExpertsFp8 does not support expert map, which is - # needed for STANDARD activation format kernels in DP/EP mode. - # Note that the BATCHED activation format does not use - # the expert map for identifying experts. + # EP is supported only with allgather_reducescatter (expert_map remapping). + # DeepEP HT and FlashInfer NVLink prepare/finalize use incompatible + # post-dispatch layouts for Standard-format grouped GEMM. return not ( moe_parallel_config.use_fi_nvl_two_sided_kernels or moe_parallel_config.use_deepep_ht_kernels @@ -491,6 +491,15 @@ def workspace_shapes( FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max +def _prepare_topk_ids_for_cutlass_fp4_moe( + topk_ids: torch.Tensor, + expert_map: torch.Tensor | None, +) -> torch.Tensor: + if expert_map is None: + return topk_ids + return remap_topk_to_local(topk_ids, expert_map, out_dtype=torch.int32) + + def run_cutlass_moe_fp4( output: torch.Tensor, a: torch.Tensor, @@ -513,6 +522,7 @@ def run_cutlass_moe_fp4( e: int, device: torch.device, apply_router_weight_on_input: bool = False, + expert_map: torch.Tensor | None = None, *, activation_config: ApplyMoEActivationConfig | None = None, ) -> None: @@ -543,6 +553,8 @@ def run_cutlass_moe_fp4( assumes that topk < k < n to satisfy - up/down projection expectations. """ + topk_ids = _prepare_topk_ids_for_cutlass_fp4_moe(topk_ids, expert_map) + is_gated = activation.is_gated # For gated activations (e.g. SiLU), w1 output is 2*n (gate + up). # For non-gated activations (e.g. SiLU_NO_MUL), w1 output is n (up only). @@ -729,9 +741,14 @@ def _supports_activation(activation: MoEActivation) -> bool: @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - # CutlassExpertsFp4 does not support expert map, which is - # needed for STANDARD activation format kernels in EP mode. - return moe_parallel_config.ep_size == 1 + # EP is supported only with allgather_reducescatter (expert_map remapping). + # DeepEP HT and FlashInfer NVLink prepare/finalize use incompatible + # post-dispatch layouts for Standard-format grouped GEMM. + return not ( + moe_parallel_config.use_fi_nvl_two_sided_kernels + or moe_parallel_config.use_deepep_ht_kernels + or moe_parallel_config.use_fi_nvl_one_sided_kernels + ) @staticmethod def _supports_batch_invariance() -> bool: @@ -806,6 +823,7 @@ def apply( e=e, device=hidden_states.device, apply_router_weight_on_input=apply_router_weight_on_input, + expert_map=expert_map, activation_config=self.activation_config, ) @@ -828,10 +846,13 @@ def run_cutlass_moe_mxfp4( e: int, device: torch.device, apply_router_weight_on_input: bool = False, + expert_map: torch.Tensor | None = None, *, activation_config: ApplyMoEActivationConfig | None = None, ) -> None: """MXFP4 x MXFP4 MoE implementation using CUTLASS grouped GEMM.""" + topk_ids = _prepare_topk_ids_for_cutlass_fp4_moe(topk_ids, expert_map) + is_gated = activation.is_gated w1_n = n * 2 if is_gated else n @@ -1036,7 +1057,14 @@ def _supports_activation(activation: MoEActivation) -> bool: def _supports_parallel_config( moe_parallel_config: FusedMoEParallelConfig, ) -> bool: - return moe_parallel_config.ep_size == 1 + # EP is supported only with allgather_reducescatter (expert_map remapping). + # DeepEP HT and FlashInfer NVLink prepare/finalize use incompatible + # post-dispatch layouts for Standard-format grouped GEMM. + return not ( + moe_parallel_config.use_fi_nvl_two_sided_kernels + or moe_parallel_config.use_deepep_ht_kernels + or moe_parallel_config.use_fi_nvl_one_sided_kernels + ) @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: @@ -1103,6 +1131,7 @@ def apply( e=e, device=hidden_states.device, apply_router_weight_on_input=apply_router_weight_on_input, + expert_map=expert_map, activation_config=self.activation_config, ) diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py index e3b6995d693c..3e70c1589365 100644 --- a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py @@ -21,7 +21,10 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) -from vllm.model_executor.layers.fused_moe.utils import _resize_cache +from vllm.model_executor.layers.fused_moe.utils import ( + _resize_cache, + remap_topk_to_local, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kMxfp4Static, @@ -994,49 +997,6 @@ def masked_moe_sum( ) -@triton.jit -def _remap_topk_to_local_kernel( - topk_ids_ptr, # [n] global expert IDs (-1 = invalid) - expert_map_ptr, # [num_experts] global->local (-1 for non-local) - out_ptr, # [n] int64 local expert IDs (-1 for invalid/non-local) - n_elements, - BLOCK: tl.constexpr, -): - pid = tl.program_id(0) - offs = pid * BLOCK + tl.arange(0, BLOCK) - mask = offs < n_elements - tid = tl.load(topk_ids_ptr + offs, mask=mask, other=-1) - # Gather expert_map[tid] for valid (tid >= 0); clamp the index so invalid - # rows don't read OOB, then select -1 for them. Matches - # torch.where(tid >= 0, expert_map[clamp(tid, 0)], -1) -- preserving -1 (a - # plain expert_map[-1] would wrap to a valid local id and misroute). - valid = tid >= 0 - idx = tl.where(valid, tid, 0) - local = tl.load(expert_map_ptr + idx, mask=mask, other=-1) - out = tl.where(valid, local.to(tl.int64), -1) - tl.store(out_ptr + offs, out, mask=mask) - - -def remap_topk_to_local( - topk_ids: torch.Tensor, expert_map: torch.Tensor -) -> torch.Tensor: - """Fused global->local expert-id mapping over a topk_ids tensor, preserving -1. - - Replaces ``torch.where(topk_ids >= 0, expert_map[topk_ids.clamp(min=0)], -1)`` - with one kernel. Returns a NEW int64 tensor -- the caller keeps the original - ``topk_ids`` as ``global_topk_ids``, so this must not write in place. - - (Distinct from ``deep_gemm_utils.apply_expert_map``, which is a scalar - ``@triton.jit`` device helper called from within other kernels.) - """ - out = torch.empty_like(topk_ids, dtype=torch.int64) - n = topk_ids.numel() - BLOCK = 1024 - grid = (triton.cdiv(n, BLOCK),) - _remap_topk_to_local_kernel[grid](topk_ids, expert_map, out, n, BLOCK=BLOCK) - return out - - class BaseOAITritonExperts(mk.FusedMoEExpertsModular): @property def expects_unquantized_inputs(self) -> bool: diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 7667dfc6aef0..cbbc39b94310 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -2133,6 +2133,59 @@ def make_mxfp4_moe_quant_config( ) +_MXFP4_SCHEME_CANDIDATES: list[tuple[QuantKey, QuantKey | None]] = [ + (kMxfp4Static, kMxfp4Dynamic), + (kMxfp4Static, kMxfp8Dynamic), + (kMxfp4Static, kFp8Dynamic128Sym), + (kMxfp4Static, kFp8StaticTensorSym), + (kMxfp4Static, None), +] + + +def _mxfp4_quant_keys_for_experts( + experts_cls: type[mk.FusedMoEExperts], + mxfp4_backend: Mxfp4MoeBackend, +) -> tuple[QuantKey | None, QuantKey | None]: + for weight_key, activation_key in _MXFP4_SCHEME_CANDIDATES: + if experts_cls._supports_quant_scheme(weight_key, activation_key): + return weight_key, activation_key + return kMxfp4Static, _backend_activation_key(mxfp4_backend) + + +def _activation_format_for_config( + moe_config: FusedMoEConfig, +) -> mk.FusedMoEActivationFormat: + return ( + mk.FusedMoEActivationFormat.BatchedExperts + if moe_config.moe_parallel_config.use_batched_activation_format + else mk.FusedMoEActivationFormat.Standard + ) + + +def _validate_mxfp4_experts_config( + experts_cls: type[mk.FusedMoEExperts], + moe_config: FusedMoEConfig, + mxfp4_backend: Mxfp4MoeBackend, +) -> None: + """Fail closed when a caller bypasses the oracle and picks experts directly.""" + activation_format = _activation_format_for_config(moe_config) + weight_key, activation_key = _mxfp4_quant_keys_for_experts( + experts_cls, mxfp4_backend + ) + supported, reason = experts_cls.is_supported_config( + experts_cls, + moe_config, + weight_key, + activation_key, + activation_format, + ) + if not supported: + raise ValueError( + f"{experts_cls.__name__} does not support this MoE deployment" + f"{f' since {reason}' if reason else ''}." + ) + + def make_mxfp4_moe_kernel( moe_quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, @@ -2141,6 +2194,8 @@ def make_mxfp4_moe_kernel( routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> mk.FusedMoEKernel: """Create a FusedMoEKernel for the given MXFP4 backend.""" + _validate_mxfp4_experts_config(experts_cls, moe_config, mxfp4_backend) + is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic) prepare_finalize = maybe_make_prepare_finalize( diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index c724a7c7cb21..b1a6dd5632b8 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -235,6 +235,73 @@ def count_expert_num_tokens( return expert_num_tokens +@triton.jit +def _remap_topk_to_local_kernel( + topk_ids_ptr, # [n] global expert IDs (-1 = invalid) + expert_map_ptr, # [num_experts] global->local (-1 for non-local) + out_ptr, # [n] local expert IDs (-1 for invalid/non-local) + n_elements, + BLOCK: tl.constexpr, + OUT_INT32: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elements + tid = tl.load(topk_ids_ptr + offs, mask=mask, other=-1) + # Gather expert_map[tid] for valid (tid >= 0); clamp the index so invalid + # rows don't read OOB, then select -1 for them. Matches + # torch.where(tid >= 0, expert_map[clamp(tid, 0)], -1) -- preserving -1 (a + # plain expert_map[-1] would wrap to a valid local id and misroute). + valid = tid >= 0 + idx = tl.where(valid, tid, 0) + local = tl.load(expert_map_ptr + idx, mask=mask, other=-1) + if OUT_INT32: + out = tl.where(valid, local, -1) + else: + out = tl.where(valid, local.to(tl.int64), -1) + tl.store(out_ptr + offs, out, mask=mask) + + +def remap_topk_to_local( + topk_ids: torch.Tensor, + expert_map: torch.Tensor, + *, + out_dtype: torch.dtype = torch.int64, +) -> torch.Tensor: + """Map global expert ids to local ids; non-local or invalid routes become -1. + + Replaces ``torch.where(topk_ids >= 0, expert_map[topk_ids.clamp(min=0)], -1)`` + with one kernel. Returns a NEW tensor -- the caller keeps the original + ``topk_ids`` as ``global_topk_ids``, so this must not write in place. + + Args: + topk_ids: Global expert ids; -1 marks invalid slots. Must be contiguous; + the Triton kernel indexes by flat ``numel()`` offsets. + expert_map: Global-to-local mapping; -1 marks non-local experts. + out_dtype: Output dtype, ``torch.int32`` or ``torch.int64``. + + Returns: + Remapped local expert ids with the same shape as ``topk_ids``. + + """ + if out_dtype not in (torch.int32, torch.int64): + raise ValueError(f"out_dtype must be int32 or int64, got {out_dtype}") + assert topk_ids.is_contiguous() + out = torch.empty(topk_ids.shape, dtype=out_dtype, device=topk_ids.device) + n = topk_ids.numel() + BLOCK = 1024 + grid = (triton.cdiv(n, BLOCK),) + _remap_topk_to_local_kernel[grid]( + topk_ids, + expert_map, + out, + n, + BLOCK=BLOCK, + OUT_INT32=out_dtype == torch.int32, + ) + return out + + def _resize_cache(x: torch.Tensor, v: tuple[int, ...]) -> torch.Tensor: """Shrink the given tensor and apply the given view to it. This is used to resize the intermediate fused_moe caches. From a281ad32f0368caa2c6ebe38a0e59a75b3f22392 Mon Sep 17 00:00:00 2001 From: Ilya Markov Date: Mon, 12 Oct 2026 16:03:09 +0200 Subject: [PATCH 308/624] [EPLB] Add EPLB support to dspark draft models (DeepSeekV4) (#57549) Signed-off-by: Markov Ilya --- tests/distributed/test_eplb_spec_decode.py | 230 +++++++++++++++++- tests/models/test_deepseek_v4_mega_moe.py | 95 ++++++++ tests/models/test_dspark_mla.py | 51 ++++ tests/v1/worker/test_dspark_utils.py | 37 ++- .../worker/test_gpu_model_runner_v2_eplb.py | 83 ++++++- vllm/distributed/eplb/eplb_state.py | 64 ++--- vllm/models/deepseek_v4/common/eplb_util.py | 50 ++++ vllm/models/deepseek_v4/nvidia/dspark.py | 16 +- vllm/models/deepseek_v4/nvidia/model.py | 25 +- vllm/v1/worker/gpu/eplb_utils.py | 37 ++- .../v1/worker/gpu/spec_decode/dspark/utils.py | 21 +- vllm/v1/worker/gpu_model_runner.py | 32 ++- 12 files changed, 665 insertions(+), 76 deletions(-) create mode 100644 vllm/models/deepseek_v4/common/eplb_util.py diff --git a/tests/distributed/test_eplb_spec_decode.py b/tests/distributed/test_eplb_spec_decode.py index ced206a3056a..b49d06cd7681 100644 --- a/tests/distributed/test_eplb_spec_decode.py +++ b/tests/distributed/test_eplb_spec_decode.py @@ -2,11 +2,21 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from __future__ import annotations +from types import SimpleNamespace + import lm_eval import pytest +import torch from tests.utils import large_gpu_mark +from vllm.distributed.eplb.eplb_state import EplbState +from vllm.models.deepseek_v4.common.eplb_util import dspark_draft_supports_eplb from vllm.platforms import current_platform +from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config +from vllm.v1.worker.gpu.eplb_utils import ( + EPLBController, + draft_model_supports_eplb, +) def get_model_args( @@ -15,14 +25,18 @@ def get_model_args( spec_method: str, tp_size: int, model_max_len: int, + num_speculative_tokens: int = 1, use_async: bool = True, ) -> dict: speculative_config = { "method": spec_method, "model": spec_model_name, - "num_speculative_tokens": 1, + "num_speculative_tokens": num_speculative_tokens, "max_model_len": model_max_len, } + if spec_method == "dspark": + speculative_config["model"] = model_name + speculative_config["draft_sample_method"] = "probabilistic" eplb_config = { "num_redundant_experts": tp_size, "window_size": 128, @@ -42,6 +56,8 @@ def get_model_args( "enable_eplb": True, "max_model_len": model_max_len, } + if spec_method == "dspark": + model_args["trust_remote_code"] = True return model_args @@ -51,11 +67,190 @@ def get_model_args( ) +class FakeEplbState: + instances: list[FakeEplbState] = [] + + def __init__(self, parallel_config, device: torch.device): + self.add_model_calls: list[tuple[object, object, str | None]] = [] + FakeEplbState.instances.append(self) + + def add_model( + self, + model: object, + model_config: object, + *, + model_name: str | None = None, + ) -> None: + self.add_model_calls.append((model, model_config, model_name)) + + +def _make_dsv4_dspark_hf_config() -> DeepseekV4Config: + return DeepseekV4Config( + architectures=["DeepseekV4ForCausalLM"], + hidden_size=128, + num_hidden_layers=2, + n_routed_experts=8, + num_experts_per_tok=2, + num_hash_layers=0, + n_shared_experts=1, + moe_intermediate_size=128, + hc_mult=1, + hc_eps=1e-5, + rms_norm_eps=1e-5, + dspark_target_layer_ids=[0], + dspark_markov_rank=8, + index_topk=4, + head_dim=64, + num_attention_heads=4, + vocab_size=256, + n_mtp_layers=2, + enable_confidence_head=False, + compress_ratios=[1, 1], + ) + + +@pytest.fixture +def dspark_vllm_config(): + hf_config = _make_dsv4_dspark_hf_config() + model_config = SimpleNamespace( + dtype=torch.bfloat16, hf_config=hf_config, model="dspark" + ) + return SimpleNamespace( + model_config=model_config, + quant_config=None, + kernel_config=SimpleNamespace(moe_backend="deep_gemm_mega_moe"), + parallel_config=SimpleNamespace( + pipeline_parallel_size=1, + tensor_parallel_size=1, + data_parallel_size=1, + enable_expert_parallel=True, + enable_eplb=True, + enable_elastic_ep=False, + eplb_config=SimpleNamespace(num_redundant_experts=4), + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=4), + compilation_config=SimpleNamespace(static_forward_context={}), + speculative_config=SimpleNamespace( + method="dspark", + draft_model_config=model_config, + ), + ) + + +def _make_moe_topology( + *, + num_routed_experts: int = 8, + num_redundant_experts: int = 4, + num_expert_groups: int = 1, +) -> SimpleNamespace: + num_physical_experts = num_routed_experts + num_redundant_experts + return SimpleNamespace( + num_routed_experts=num_routed_experts, + num_redundant_experts=num_redundant_experts, + num_physical_experts=num_physical_experts, + num_logical_experts=num_routed_experts, + num_expert_groups=num_expert_groups, + ) + + +def test_eplb_state_accepts_matching_dsv4_draft_and_target(): + draft = _make_moe_topology() + target = _make_moe_topology() + + EplbState.assert_confs_equal(draft, target) + + +def test_eplb_state_rejects_mismatched_dsv4_draft_redundant_experts(): + draft = _make_moe_topology(num_redundant_experts=0) + target = _make_moe_topology(num_redundant_experts=4) + + with pytest.raises(RuntimeError, match="mismatch"): + EplbState.assert_confs_equal(draft, target) + + +@pytest.mark.parametrize( + ("model_type", "expected"), + [ + ("deepseek_v4", True), + ("deepseek_v41", False), + ], +) +def test_dspark_draft_supports_eplb( + dspark_vllm_config, model_type: str, expected: bool +): + draft_model_config = dspark_vllm_config.speculative_config.draft_model_config + draft_model_config.hf_config.model_type = model_type + assert dspark_draft_supports_eplb(draft_model_config) is expected + + +def test_draft_model_supports_eplb_requires_draft_moe(dspark_vllm_config): + assert not draft_model_supports_eplb( + dspark_vllm_config.speculative_config, + None, + ) + + +@pytest.mark.parametrize( + "model_type,load_dummy_weights,expected_registered", + [ + ("deepseek_v4", False, True), + ("deepseek_v4", True, False), + ("deepseek_v41", False, False), + ], +) +def test_eplb_dspark_draft_registration( + dspark_vllm_config, + monkeypatch: pytest.MonkeyPatch, + model_type, + load_dummy_weights: bool, + expected_registered: bool, +): + """Register V4 drafts; skip incompatible V4.1 drafts and dummy weights.""" + FakeEplbState.instances.clear() + monkeypatch.setattr( + "vllm.v1.worker.gpu.eplb_utils.EplbState", + FakeEplbState, + ) + monkeypatch.setattr( + "vllm.v1.worker.gpu.eplb_utils.get_mixture_of_experts_model", + lambda model: model, + ) + + draft_model_config = dspark_vllm_config.speculative_config.draft_model_config + draft_model_config.hf_config.model_type = model_type + draft = SimpleNamespace() + controller = EPLBController(dspark_vllm_config.parallel_config, torch.device("cpu")) + controller.prepare_load() + speculator = SimpleNamespace(model=draft, eplb_state=None) + + def set_eplb_state(state) -> None: + speculator.eplb_state = state + + speculator.set_eplb_state = set_eplb_state + + registered = controller.maybe_register_speculator( + speculator, + dspark_vllm_config.speculative_config, + load_dummy_weights=load_dummy_weights, + ) + + assert registered is expected_registered + assert controller.state is not None + if expected_registered: + assert controller.state.add_model_calls == [ + (draft, draft_model_config, "dspark (draft)") + ] + assert speculator.eplb_state is controller.state + else: + assert controller.state.add_model_calls == [] + assert speculator.eplb_state is None + + @pytest.mark.parametrize( "model_setup", [ pytest.param( - ("mtp", "Qwen/Qwen3-Next-80B-A3B-Instruct", None, 4, 0.86), + ("mtp", "Qwen/Qwen3-Next-80B-A3B-Instruct", None, 4, 0.86, 1), marks=large_gpu_mark(min_gb=80), ), pytest.param( @@ -65,20 +260,42 @@ def get_model_args( "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct", 4, 0.92, + 1, ), marks=pytest.mark.skip(reason="Skipping due to CI OOM issues"), ), + pytest.param( + ( + "dspark", + "deepseek-ai/DeepSeek-V4-Flash-DSpark", + None, + 4, + 0.95, + 7, + ), + marks=large_gpu_mark(min_gb=80), + ), ], - ids=["qwen3_next_mtp", "llama4_eagle"], + ids=["qwen3_next_mtp", "llama4_eagle", "dsv4_dspark"], ) def test_eplb_spec_decode( monkeypatch: pytest.MonkeyPatch, - model_setup: tuple[str, str, str, int, float], + model_setup: tuple[str, str, str | None, int, float, int], ): """Test the correctness of EPLB speculative decoding with GSM8K dataset. - Applicable to MoE models with mtp or eagle spec decode. + + Applicable to MoE models with mtp, eagle, or dspark spec decode. """ - method, model_name, spec_model_name, tp_size, expected_gsm8k_value = model_setup + ( + method, + model_name, + spec_model_name, + tp_size, + expected_gsm8k_value, + num_speculative_tokens, + ) = model_setup + if method == "dspark" and not current_platform.is_device_capability_family(100): + pytest.skip("DSV4 DSpark EPLB requires SM100 for MegaMoE") TASK = "gsm8k" FILTER = "exact_match,strict-match" @@ -90,6 +307,7 @@ def test_eplb_spec_decode( spec_method=method, tp_size=tp_size, model_max_len=4096, + num_speculative_tokens=num_speculative_tokens, ) results = lm_eval.simple_evaluate( diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 76915b94a32d..250a9d729186 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -15,6 +15,10 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import ( deepgemm_post_process_weight_scale_block, ) +from vllm.model_executor.models.interfaces import ( + get_mixture_of_experts_model, + is_mixture_of_experts, +) from vllm.models.deepseek_v4.nvidia.dspark import DSparkDeepseekV4ForCausalLM from vllm.models.deepseek_v4.nvidia.model import ( DeepseekV4ForCausalLM, @@ -28,6 +32,7 @@ from vllm.models.deepseek_v41.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID from vllm.models.deepseek_v41.nvidia.model import DeepseekV4MoE as DeepseekV41MoE from vllm.platforms import current_platform +from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config from vllm.utils.deep_gemm import is_deep_gemm_supported from vllm.utils.torch_utils import set_default_torch_dtype @@ -913,6 +918,96 @@ def test_deepseek_v4_drafter_pwal_hooks_finalize_mega_moe(): assert calls == ["mtp", "dspark"] +@pytest.fixture +def v4_dspark_config(dist_init): + hf_config = DeepseekV4Config( + architectures=["DeepseekV4ForCausalLM"], + hidden_size=128, + num_hidden_layers=2, + n_routed_experts=8, + num_experts_per_tok=2, + num_hash_layers=0, + n_shared_experts=1, + moe_intermediate_size=128, + hidden_act="silu", + swiglu_limit=10.0, + norm_topk_prob=True, + topk_method="noaux_tc", + routed_scaling_factor=1.5, + hc_mult=1, + hc_eps=1e-5, + rms_norm_eps=1e-5, + dspark_target_layer_ids=[0], + dspark_markov_rank=8, + index_topk=4, + head_dim=64, + num_attention_heads=4, + vocab_size=256, + n_mtp_layers=2, + enable_confidence_head=False, + compress_ratios=[1, 1], + q_lora_rank=32, + o_lora_rank=32, + qk_rope_head_dim=16, + o_groups=4, + sliding_window=128, + hc_sinkhorn_iters=1, + ) + model_config = SimpleNamespace( + dtype=torch.bfloat16, + hf_config=hf_config, + max_model_len=2048, + ) + return SimpleNamespace( + model_config=model_config, + quant_config=None, + kernel_config=SimpleNamespace( + moe_backend="deep_gemm_mega_moe", + enable_jit_warmup=False, + ), + parallel_config=SimpleNamespace( + pipeline_parallel_size=1, + tensor_parallel_size=1, + data_parallel_size=1, + enable_expert_parallel=True, + enable_eplb=True, + eplb_config=SimpleNamespace(num_redundant_experts=4), + ), + attention_config=SimpleNamespace(backend=None), + cache_config=SimpleNamespace(block_size=64, cache_dtype="auto"), + scheduler_config=SimpleNamespace(max_num_batched_tokens=4), + compilation_config=SimpleNamespace(static_forward_context={}), + speculative_config=SimpleNamespace(draft_model_config=model_config), + use_v2_model_runner=False, + ) + + +def test_dspark_draft_registers_mixture_of_experts( + v4_dspark_config, monkeypatch, dist_init +): + """DSV4 DSpark MoE drafts expose MixtureOfExperts for EPLB registration.""" + if not current_platform.is_device_capability_family(100): + pytest.skip("DeepGEMM MegaMoE requires SM100") + + hf_config = v4_dspark_config.model_config.hf_config + monkeypatch.setattr( + "vllm.models.deepseek_v4.nvidia.dspark.get_current_vllm_config", + lambda: v4_dspark_config, + ) + + with ( + set_default_torch_dtype(v4_dspark_config.model_config.dtype), + torch.device("cuda"), + ): + draft = DSparkDeepseekV4ForCausalLM(vllm_config=v4_dspark_config) + + assert is_mixture_of_experts(draft) + assert get_mixture_of_experts_model(draft) is draft + assert len(draft.moe_layers) == draft.model.num_dspark_layers == 2 + assert draft.num_routed_experts == hf_config.n_routed_experts + assert draft.num_redundant_experts == 4 + + @pytest.mark.skipif( not torch.cuda.is_available(), reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.", diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index 53a93b940cd5..b2156af346c8 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -266,6 +266,57 @@ def load_scale(param, weight, *args): torch.testing.assert_close(param, checkpoint_scale) +def test_dsv4_dspark_fused_moe_expert_mapping_includes_redundant_experts( + monkeypatch, +): + from vllm.models.deepseek_v4.nvidia import dspark + + captured: dict[str, int] = {} + + def capture_expert_mapping(*args, **kwargs): + captured.update(kwargs) + return [] + + draft = SimpleNamespace( + config=SimpleNamespace( + num_attention_heads=4, + n_routed_experts=8, + expert_dtype="fp4", + ), + quant_config=None, + pad_shared_expert=False, + model=SimpleNamespace( + layers=[ + SimpleNamespace( + ffn=SimpleNamespace( + use_native_mega_moe=False, + use_mega_moe=False, + n_redundant_experts=4, + ) + ) + ], + confidence_head=None, + ), + named_parameters=lambda: [], + process_weights_after_loading=lambda: None, + ) + draft._remap_dspark_name = lambda name: ( + dspark.DSparkDeepseekV4ForCausalLM._remap_dspark_name(draft, name) + ) + monkeypatch.setattr(dspark, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(dspark, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + dspark, + "fused_moe_make_expert_params_mapping", + capture_expert_mapping, + ) + + dspark.DSparkDeepseekV4ForCausalLM.load_weights(draft, []) + + assert captured["num_experts"] == 8 + assert captured["num_redundant_experts"] == 4 + + def test_dsv4_context_wkv_weights_are_duplicated_by_draft_layer(): weights = [ ("mtp.0.attn.wkv.weight", torch.arange(4)), diff --git a/tests/v1/worker/test_dspark_utils.py b/tests/v1/worker/test_dspark_utils.py index 1014051d649c..891ca5602079 100644 --- a/tests/v1/worker/test_dspark_utils.py +++ b/tests/v1/worker/test_dspark_utils.py @@ -15,6 +15,16 @@ from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator +@dataclass +class _FakeHFConfig: + model_type: str + + +@dataclass +class _FakeModelConfig: + hf_config: _FakeHFConfig + + @dataclass class _FakeEPLBConfig: num_redundant_experts: int = 0 @@ -37,12 +47,28 @@ def __post_init__(self) -> None: raise ValueError("elastic EP requires EPLB") -def test_dspark_parallel_config_disables_eplb_atomically(): +@pytest.mark.parametrize( + ("model_type", "enable_eplb", "num_redundant_experts", "shares_eplb_config"), + [ + ("deepseek_v4", True, 32, True), + ("deepseek_v41", False, 0, False), + ], +) +def test_dspark_parallel_config_eplb_by_model_type( + model_type: str, + enable_eplb: bool, + num_redundant_experts: int, + shares_eplb_config: bool, +): target_config = _FakeParallelConfig() + draft_model_config = _FakeModelConfig( + hf_config=_FakeHFConfig(model_type=model_type) + ) draft_config = _get_dspark_parallel_config( target_config, tensor_parallel_size=4, + draft_model_config=draft_model_config, ) assert target_config.pipeline_parallel_size == 2 @@ -54,10 +80,13 @@ def test_dspark_parallel_config_disables_eplb_atomically(): assert draft_config is not target_config assert draft_config.pipeline_parallel_size == 1 assert draft_config.tensor_parallel_size == 4 - assert not draft_config.enable_eplb - assert draft_config.eplb_config.num_redundant_experts == 0 + assert draft_config.enable_eplb is enable_eplb + assert draft_config.eplb_config.num_redundant_experts == num_redundant_experts assert not draft_config.enable_elastic_ep - assert draft_config.eplb_config is not target_config.eplb_config + if shares_eplb_config: + assert draft_config.eplb_config is target_config.eplb_config + else: + assert draft_config.eplb_config is not target_config.eplb_config @pytest.mark.parametrize("pcp_size", [1, 4]) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index 6b38bfef968d..b0acf6916015 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -29,15 +29,21 @@ class FakeEplbState: def __init__(self, parallel_config: Any, device: torch.device): self.parallel_config = parallel_config self.device = device - self.add_model_calls: list[tuple[Any, Any]] = [] + self.add_model_calls: list[tuple[Any, Any, str | None]] = [] self.step_calls: list[tuple[bool, bool, bool]] = [] self.async_started = False self.is_async = True self.update_mapping_calls: list[tuple[Any, torch.Tensor]] = [] FakeEplbState.instances.append(self) - def add_model(self, model: Any, model_config: Any) -> None: - self.add_model_calls.append((model, model_config)) + def add_model( + self, + model: Any, + model_config: Any, + *, + model_name: str | None = None, + ) -> None: + self.add_model_calls.append((model, model_config, model_name)) def step(self, is_dummy: bool, is_profile: bool, *, log_stats: bool) -> None: self.step_calls.append((is_dummy, is_profile, log_stats)) @@ -98,6 +104,15 @@ def _make_runner(**overrides: Any) -> Any: runner.execute_model_state = None for key, value in overrides.items(): setattr(runner, key, value) + spec_config = runner.speculative_config + runner._draft_workspace_lane = int( + spec_config is not None + and ( + spec_config.use_dspark() + if hasattr(spec_config, "use_dspark") + else getattr(spec_config, "method", None) == "dspark" + ) + ) return runner @@ -131,7 +146,7 @@ def test_v2_load_model_registers_moe_with_eplb(monkeypatch): assert runner.model is model assert runner.model_state is not None assert runner.eplb_state is not None - assert runner.eplb_state.add_model_calls == [(model, runner.model_config)] + assert runner.eplb_state.add_model_calls == [(model, runner.model_config, None)] assert runner.eplb_state.async_started is True @@ -159,7 +174,7 @@ def test_v2_load_model_with_dummy_weights_defers_eplb_async_loop(monkeypatch): assert runner.load_config.load_format == "dummy" assert runner.eplb_state is not None # A scaling-up worker registers so it can join the EPLB communicator. - assert runner.eplb_state.add_model_calls == [(model, runner.model_config)] + assert runner.eplb_state.add_model_calls == [(model, runner.model_config, None)] assert runner.eplb_state.async_started is False @@ -174,6 +189,64 @@ def test_v2_setup_eplb_from_mapping_updates_state_in_place(): assert state.update_mapping_calls == [(runner.model_config, mapping)] +def test_v2_load_model_registers_dspark_speculator_with_eplb(monkeypatch): + FakeEplbState.instances.clear() + target_model = SimpleNamespace(is_moe=True) + draft_model = SimpleNamespace(is_moe=True) + draft_model_config = SimpleNamespace( + model="dspark-draft", + hf_config=SimpleNamespace(model_type="deepseek_v4"), + ) + + class FakeDSparkSpeculator: + def __init__(self): + self.model = draft_model + self.eplb_state = None + + def load_model(self, target_model): + return None + + def set_eplb_state(self, state) -> None: + self.eplb_state = state + + monkeypatch.setattr(mrv2, "DeviceMemoryProfiler", FakeMemoryProfiler) + monkeypatch.setattr(eplb, "EplbState", FakeEplbState) + monkeypatch.setattr(mrv2, "DraftModelSpeculator", FakeDSparkSpeculator) + monkeypatch.setattr( + mrv2, + "get_model_loader", + lambda load_config: SimpleNamespace(load_model=lambda **_: target_model), + ) + monkeypatch.setattr( + mrv2, + "init_model_state", + lambda *args: SimpleNamespace(num_new_sampled_tokens_per_step=1), + ) + monkeypatch.setattr( + eplb, + "get_mixture_of_experts_model", + lambda model: model if getattr(model, "is_moe", False) else None, + ) + + speculator = FakeDSparkSpeculator() + runner = _make_runner( + is_last_pp_rank=False, + speculative_config=SimpleNamespace( + method="dspark", + draft_model_config=draft_model_config, + ), + speculator=speculator, + ) + mrv2.GPUModelRunner.load_model(runner) + + assert runner.eplb_state is not None + assert runner.eplb_state.add_model_calls == [ + (draft_model, draft_model_config, "dspark-draft (draft)"), + (target_model, runner.model_config, None), + ] + assert speculator.eplb_state is runner.eplb_state + + def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch): events = [] runner = _make_runner(is_last_pp_rank=False, num_speculative_steps=0) diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index 205226cde4a9..2d782c417945 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -323,39 +323,47 @@ def validate_ep_configuration(self, new_model: MixtureOfExperts): """ if len(self.model_states) > 0: model = next(iter(self.model_states.values())).model - if ( - model.num_routed_experts != new_model.num_routed_experts - or model.num_redundant_experts != new_model.num_redundant_experts - or model.num_physical_experts != new_model.num_physical_experts - or model.num_logical_experts != new_model.num_logical_experts - or model.num_expert_groups != new_model.num_expert_groups - ): - raise RuntimeError( - "Model: {} " - "with config {} " - "{} {} {} {} " - "mismatch with new model {} " - "with config {} " - "{} {} {} {}".format( - type(model), - model.num_routed_experts, - model.num_redundant_experts, - model.num_physical_experts, - model.num_logical_experts, - model.num_expert_groups, - type(new_model), - new_model.num_routed_experts, - new_model.num_redundant_experts, - new_model.num_physical_experts, - new_model.num_logical_experts, - new_model.num_expert_groups, - ) + self.assert_confs_equal(model, new_model) + + @staticmethod + def assert_confs_equal( + model: MixtureOfExperts, new_model: MixtureOfExperts + ) -> None: + """Raise if the models have different expert parallel configurations.""" + if ( + model.num_routed_experts != new_model.num_routed_experts + or model.num_redundant_experts != new_model.num_redundant_experts + or model.num_physical_experts != new_model.num_physical_experts + or model.num_logical_experts != new_model.num_logical_experts + or model.num_expert_groups != new_model.num_expert_groups + ): + raise RuntimeError( + "Model: {} " + "with config {} " + "{} {} {} {} " + "mismatch with new model {} " + "with config {} " + "{} {} {} {}".format( + type(model), + model.num_routed_experts, + model.num_redundant_experts, + model.num_physical_experts, + model.num_logical_experts, + model.num_expert_groups, + type(new_model), + new_model.num_routed_experts, + new_model.num_redundant_experts, + new_model.num_physical_experts, + new_model.num_logical_experts, + new_model.num_expert_groups, ) + ) def add_model( self, model: MixtureOfExperts, model_config: ModelConfig, + model_name: str | None = None, ): """Build the initial EPLB state.""" self.validate_ep_configuration(model) @@ -506,7 +514,7 @@ def add_model( dtype=torch.int64, device=self.device, ), - model_name=model_config.model, + model_name=model_name or model_config.model, model=model, expert_buffer=expert_buffer, rebalanced=False, diff --git a/vllm/models/deepseek_v4/common/eplb_util.py b/vllm/models/deepseek_v4/common/eplb_util.py new file mode 100644 index 000000000000..53c6d100243d --- /dev/null +++ b/vllm/models/deepseek_v4/common/eplb_util.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared EPLB MoE registration helpers for DeepSeek V4 family models.""" + +from collections.abc import Iterable +from typing import Any + +import torch.nn as nn + +from vllm.config import ModelConfig + + +def dspark_draft_supports_eplb(draft_model_config: ModelConfig) -> bool: + """Return whether a DSpark draft can share EPLB state with the target. + + Only DeepSeek-V4 DSpark drafts reuse the target expert layout. V4.1 drafts + use a smaller routed-expert count and cannot share EPLB state. + """ + return getattr(draft_model_config.hf_config, "model_type", None) == "deepseek_v4" + + +def collect_moe_layers( + moe_model: Any, + layers: Iterable[nn.Module], + config: object, + *, + decoder_layer_type: type[nn.Module], + moe_type: type[nn.Module], +) -> None: + """Populate ``MixtureOfExperts`` fields from decoder layers for EPLB. + + Works for both target backbones and DSpark draft stacks. Draft models pass + their shorter ``model.layers`` list; pipeline-parallel placeholder layers + are skipped via ``decoder_layer_type``. + """ + moe_model.num_expert_groups = getattr(config, "n_group", 1) + moe_model.moe_layers = [] + moe_model.moe_mlp_layers = [] + example_moe = None + for layer in layers: + if not isinstance(layer, decoder_layer_type): + continue + if not isinstance(layer.ffn, moe_type): + continue + example_moe = layer.ffn + moe_model.moe_mlp_layers.append(layer.ffn) + moe_model.moe_layers.append(layer.ffn.experts) + + moe_model.num_moe_layers = len(moe_model.moe_layers) + moe_model.extract_moe_parameters(example_moe) diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index 54250185c840..2a40412e933b 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -52,11 +52,14 @@ sp_padding_mask, sp_shard, ) +from vllm.models.deepseek_v4.common.eplb_util import collect_moe_layers from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID from .model import ( DeepseekV4DecoderLayer, + DeepseekV4MixtureOfExperts, DeepseekV4Model, + DeepseekV4MoE, _use_sequence_parallel, make_deepseek_v4_expert_params_mapping, prepare_mega_gate_routing_metadata, @@ -346,7 +349,7 @@ def _insert_context_kv( ) -class DSparkDeepseekV4ForCausalLM(nn.Module): +class DSparkDeepseekV4ForCausalLM(nn.Module, DeepseekV4MixtureOfExperts): # Draft weights ship in the target checkpoint (mtp.*) without embed/head, so # load_dspark_model always aliases the target's. has_own_embed_tokens = False @@ -373,6 +376,16 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: prefix=maybe_prefix(prefix, "lm_head"), ) self.logits_processor = LogitsProcessor(self.config.vocab_size) + self.set_moe_parameters() + + def set_moe_parameters(self) -> None: + collect_moe_layers( + self, + self.model.layers, + self.config, + decoder_layer_type=DeepseekV4DecoderLayer, + moe_type=DeepseekV4MoE, + ) # --- Hooks used by the speculator ------------------------------------- @@ -451,6 +464,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ckpt_down_proj_name="w2", ckpt_up_proj_name="w3", num_experts=self.config.n_routed_experts, + num_redundant_experts=first_layer.ffn.n_redundant_experts, ) expert_scale_suffix = ( ".weight_scale" diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index fc10cd07c1cf..ff4470dc4c97 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -81,6 +81,7 @@ sp_shard, ) from vllm.models.deepseek_v4.attention import DeepseekV4Attention +from vllm.models.deepseek_v4.common.eplb_util import collect_moe_layers from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import ( DeepseekV4FlashInferMLAAttention, DeepseekV4FlashInferSM120Attention, @@ -1841,23 +1842,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.set_moe_parameters() def set_moe_parameters(self) -> None: - self.num_expert_groups = getattr(self.config, "n_group", 1) - self.num_moe_layers = self.config.num_hidden_layers - self.moe_layers: list[nn.Module] = [] - self.moe_mlp_layers: list[DeepseekV4MoE] = [] - example_moe: DeepseekV4MoE | None = None - for layer in self.model.layers: - if isinstance(layer, PPMissingLayer): - continue - if not isinstance(layer, DeepseekV4DecoderLayer): - continue - if isinstance(layer.ffn, DeepseekV4MoE): - example_moe = layer.ffn - self.moe_mlp_layers.append(layer.ffn) - self.moe_layers.append(layer.ffn.experts) - - self.num_moe_layers = len(self.moe_layers) - self.extract_moe_parameters(example_moe) + collect_moe_layers( + self, + self.model.layers, + self.config, + decoder_layer_type=DeepseekV4DecoderLayer, + moe_type=DeepseekV4MoE, + ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) diff --git a/vllm/v1/worker/gpu/eplb_utils.py b/vllm/v1/worker/gpu/eplb_utils.py index e5eb31ad28f0..5c2c5d14121e 100644 --- a/vllm/v1/worker/gpu/eplb_utils.py +++ b/vllm/v1/worker/gpu/eplb_utils.py @@ -13,8 +13,10 @@ from vllm.distributed.eplb.eplb_state import EplbState from vllm.logger import init_logger from vllm.model_executor.models.interfaces import ( + MixtureOfExperts, get_mixture_of_experts_model, ) +from vllm.models.deepseek_v4.common.eplb_util import dspark_draft_supports_eplb if TYPE_CHECKING: from vllm.v1.worker.gpu.model_runner import GPUModelRunner @@ -22,6 +24,25 @@ logger = init_logger(__name__) +def eplb_draft_model_name(model_config: ModelConfig) -> str: + """Return the EPLB step-log display name for a draft model.""" + return f"{model_config.model} (draft)" + + +def draft_model_supports_eplb( + speculative_config: Any | None, + draft_moe_model: MixtureOfExperts | None, +) -> bool: + """Return whether a draft MoE model should register with EPLB.""" + if draft_moe_model is None: + return False + if speculative_config is None or speculative_config.draft_model_config is None: + return False + if getattr(speculative_config, "method", None) == "dspark": + return dspark_draft_supports_eplb(speculative_config.draft_model_config) + return True + + def step_eplb_after(*, is_dummy: bool = False) -> Callable: """Step EPLB after a model runner method completes successfully.""" @@ -93,16 +114,26 @@ def maybe_register_speculator( draft_model = speculator.model draft_moe_model = get_mixture_of_experts_model(draft_model) - if draft_moe_model is None: + if not draft_model_supports_eplb(speculative_config, draft_moe_model): return False + assert draft_moe_model is not None assert not self.parallel_config.enable_elastic_ep, ( "Elastic EP is not supported with draft model." ) assert speculative_config is not None - assert speculative_config.draft_model_config is not None + draft_model_config = speculative_config.draft_model_config + assert draft_model_config is not None + logger.info_once( + "EPLB is enabled for MoE part of drafter model %s.", + draft_model_config.model, + ) assert self.state is not None - self.state.add_model(draft_moe_model, speculative_config.draft_model_config) + self.state.add_model( + draft_moe_model, + draft_model_config, + model_name=eplb_draft_model_name(draft_model_config), + ) speculator.set_eplb_state(self.state) self._has_registered_models = True return True diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index 489d0da554a2..38108d4e0414 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -6,6 +6,7 @@ from vllm.config import ModelConfig, ParallelConfig, VllmConfig, replace from vllm.logger import init_logger from vllm.model_executor.model_loader.utils import get_draft_load_config +from vllm.models.deepseek_v4.common.eplb_util import dspark_draft_supports_eplb from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config @@ -35,11 +36,24 @@ def _resolve_dspark_attention_backend( def _get_dspark_parallel_config( parallel_config: ParallelConfig, tensor_parallel_size: int, + draft_model_config: ModelConfig, ) -> ParallelConfig: - if parallel_config.enable_eplb: + enable_draft_eplb = parallel_config.enable_eplb and dspark_draft_supports_eplb( + draft_model_config + ) + if parallel_config.enable_eplb and not enable_draft_eplb: logger.warning_once( - "EPLB is disabled for the DSpark draft model. EPLB remains enabled " - "for the target model." + "EPLB is disabled for the DSpark draft model (%s). EPLB remains " + "enabled for the target model.", + draft_model_config.hf_config.model_type, + ) + + if enable_draft_eplb: + return replace( + parallel_config, + pipeline_parallel_size=1, + tensor_parallel_size=tensor_parallel_size, + enable_elastic_ep=False, ) return replace( @@ -81,6 +95,7 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo parallel_config=_get_dspark_parallel_config( vllm_config.parallel_config, speculative_config.draft_parallel_config.tensor_parallel_size, + draft_model_config, ), attention_config=replace( vllm_config.attention_config, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 4db3f00b91d4..669b314d4cac 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -206,6 +206,10 @@ ) from vllm.v1.worker.dp_utils import coordinate_batch_across_dp from vllm.v1.worker.ec_connector_model_runner_mixin import ECConnectorModelRunnerMixin +from vllm.v1.worker.gpu.eplb_utils import ( + draft_model_supports_eplb, + eplb_draft_model_name, +) from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin @@ -5237,25 +5241,34 @@ def load_model( logger.info_once("Loading drafter model...") if hasattr(self.drafter, "load_model"): self.drafter.load_model(self.model) + drafter_moe_model = None + if self.parallel_config.enable_eplb and hasattr( + self.drafter, "model" + ): + drafter_moe_model = get_mixture_of_experts_model( + self.drafter.model + ) + if not draft_model_supports_eplb( + self.vllm_config.speculative_config, + drafter_moe_model, + ): + drafter_moe_model = None + if ( self.parallel_config.enable_eplb and hasattr(self.drafter, "model") - and ( - drafter_moe_model := get_mixture_of_experts_model( - self.drafter.model - ) - ) - is not None + and drafter_moe_model is not None ): assert not self.parallel_config.enable_elastic_ep, ( "Elastic EP is not supported with drafter model." ) spec_config = self.vllm_config.speculative_config assert spec_config is not None - assert spec_config.draft_model_config is not None + draft_model_config = spec_config.draft_model_config + assert draft_model_config is not None logger.info_once( "EPLB is enabled for MoE part of drafter model %s.", - spec_config.draft_model_config.model, + draft_model_config.model, ) if self.eplb_state is None: self.eplb_state = EplbState( @@ -5263,7 +5276,8 @@ def load_model( ) self.eplb_state.add_model( drafter_moe_model, - spec_config.draft_model_config, + draft_model_config, + model_name=eplb_draft_model_name(draft_model_config), ) assert hasattr(self.drafter, "set_eplb_state") self.drafter.set_eplb_state(self.eplb_state) From 716961cf1fc474c491d401adece882af3b39d5ab Mon Sep 17 00:00:00 2001 From: Shang Wang Date: Mon, 12 Oct 2026 10:05:09 -0400 Subject: [PATCH 309/624] [Bugfix] Drop enable_in_kernel_fc2_reduce from FlashInfer CuteDSL MegaMoE config (#59900) Signed-off-by: Shang Wang Co-authored-by: Claude Opus 5.5 --- tests/kernels/moe/test_flashinfer_moe_ep.py | 62 +++++++++++++++++++ .../layers/fused_moe/flashinfer_moe_ep.py | 3 +- 2 files changed, 64 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_flashinfer_moe_ep.py b/tests/kernels/moe/test_flashinfer_moe_ep.py index c3223b7d9800..7ce3124ae144 100644 --- a/tests/kernels/moe/test_flashinfer_moe_ep.py +++ b/tests/kernels/moe/test_flashinfer_moe_ep.py @@ -10,9 +10,14 @@ the DeepGEMM kernel (FP8 activations and FP4 weights, power-of-two scales per 32). The expected output is then known in closed form and compared bit for bit. + +A test that needs no SM100 also checks the adapter's kwargs against the +installed FlashInfer's config types, so config drift fails on any CI runner +with FlashInfer installed. """ from dataclasses import dataclass +from types import SimpleNamespace import pytest import torch @@ -22,6 +27,7 @@ from vllm.config import VllmConfig, set_current_vllm_config from vllm.config.kernel import FLASHINFER_MOE_EP_CUTEDSL, FLASHINFER_MOE_EP_DEEP_GEMM from vllm.distributed.parallel_state import ensure_model_parallel_initialized +from vllm.model_executor.layers.fused_moe import flashinfer_moe_ep as fi_ep from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, @@ -305,3 +311,59 @@ def test_init_warmup_forward_is_exact_on_grid_inputs( if torch.accelerator.device_count() < world_size: pytest.skip(f"Need at least {world_size} GPUs") distributed_run(_init_warmup_forward, world_size, geometry, backend) + + +@pytest.mark.parametrize("backend", BACKENDS, ids=("cutedsl", "deep_gemm")) +def test_adapter_kwargs_match_installed_flashinfer( + monkeypatch: pytest.MonkeyPatch, backend: str +): + """The adapter's kwargs construct the installed FlashInfer's config types. + + Needs no SM100, so config kwarg drift fails in CI rather than at model load. + Fakes only the megakernel layer (symmetric memory), the EP group, the device + index and the DeepGEMM import. + """ + moe_ep = pytest.importorskip("flashinfer.moe_ep") + layer_args: list[tuple] = [] + + class FakeMegaLayer: + def __init__(self, *args) -> None: + layer_args.append(args) + + def destroy(self) -> None: + pass + + monkeypatch.setattr(moe_ep, "MoEEpMegaLayer", FakeMegaLayer) + monkeypatch.setattr(fi_ep, "_expose_deep_gemm_to_flashinfer", lambda: None) + monkeypatch.setattr( + fi_ep, + "get_ep_group", + lambda: SimpleNamespace(world_size=2, rank_in_group=0, device_group=None), + ) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) + + geometry = GEOMETRIES[0] + moe = SimpleNamespace( + moe_backend=backend, + num_experts=geometry.num_experts, + experts_per_token=geometry.top_k, + hidden_dim=geometry.hidden_size, + intermediate_size=geometry.intermediate_size, + max_num_tokens=MAX_TOKENS_PER_RANK, + swiglu_limit=10.0, + routing_method=_routing_method(backend), + ) + weight = torch.zeros(1, dtype=torch.bfloat16) + FlashInferMoeEp( + moe, + FlashInferMoeEpWeights(w13=weight, w2=weight), + apply_topk_in_fc1=apply_topk_in_fc1(moe), + ) + + ((_, _, _, mega_config),) = layer_args + expected = ( + moe_ep.Nvfp4CutedslMegaMoeConfig + if backend == FLASHINFER_MOE_EP_CUTEDSL + else moe_ep.DeepGemmMegaMoeConfig + ) + assert isinstance(mega_config.megakernel, expected) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py b/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py index bac5345255c9..63345a417438 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py @@ -315,13 +315,14 @@ def __init__( w2_scale=weights.w2_scale, ) if spec.kernel == "cutedsl": + # In-kernel fc2 reduce keeps its default (off); FlashInfer 0.7.1rc1 + # renamed in_kernel_fc2_reduce to enable_in_kernel_fc2_reduce. megakernel = api.Nvfp4CutedslMegaMoeConfig( intermediate_size=moe.intermediate_size, top_k=moe.experts_per_token, gate_up_clamp=moe.swiglu_limit, fast_math=True, apply_topk_in_fc1=apply_topk_in_fc1, - enable_in_kernel_fc2_reduce=False, combine_dtype="bf16", input_norm_const=epilogue.input_norm_const, fc1_alpha=None, From f3cb10b690c352265048991901de0e62b65abe8c Mon Sep 17 00:00:00 2001 From: Hoai Vo <62198968+hoaivoke@users.noreply.github.com> Date: Mon, 12 Oct 2026 10:09:53 -0400 Subject: [PATCH 310/624] [Model][LoRA] Enable LoRA support for BERT and RoBERTa embedding models (#61033) Signed-off-by: hoaivoke Co-authored-by: Claude Opus 5.5 --- docs/models/pooling_models/embed.md | 6 +++--- vllm/model_executor/models/bert.py | 17 +++++++++++++++-- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/docs/models/pooling_models/embed.md b/docs/models/pooling_models/embed.md index 4c33126977e9..5211997edb81 100644 --- a/docs/models/pooling_models/embed.md +++ b/docs/models/pooling_models/embed.md @@ -37,7 +37,7 @@ You can compute pairwise similarity scores to build a similarity matrix using th | Architecture | Models | Example HF Models | [LoRA](../../features/lora.md) | [PP](../../serving/parallelism_scaling.md) | | ------------ | ------ | ----------------- | ------------------------------ | ------------------------------------------ | -| `BertModel` | BERT-based | `BAAI/bge-base-en-v1.5`, `Snowflake/snowflake-arctic-embed-xs`, etc. | | | +| `BertModel` | BERT-based | `BAAI/bge-base-en-v1.5`, `Snowflake/snowflake-arctic-embed-xs`, etc. | ✅︎ | | | `BertSpladeSparseEmbeddingModel` | SPLADE | `naver/splade-v3` | | | | `BgeM3EmbeddingModel` | BGE-M3 | `BAAI/bge-m3` | | | | `Gemma2Model`C | Gemma 2-based | `BAAI/bge-multilingual-gemma2`, etc. | ✅︎ | ✅︎ | @@ -51,9 +51,9 @@ You can compute pairwise similarity scores to build a similarity matrix using th | `NomicBertModel` | Nomic BERT | `nomic-ai/nomic-embed-text-v1`, `nomic-ai/nomic-embed-text-v2-moe`, `Snowflake/snowflake-arctic-embed-m-long`, etc. | | | | `Qwen2Model`C, `Qwen2ForCausalLM`C | Qwen2-based | `ssmits/Qwen2-7B-Instruct-embed-base` (see note), `Alibaba-NLP/gte-Qwen2-7B-instruct` (see note), etc. | ✅︎ | ✅︎ | | `Qwen3Model`C, `Qwen3ForCausalLM`C | Qwen3-based | `Qwen/Qwen3-Embedding-0.6B`, `ai-sage/Giga-Embeddings-instruct-480M-0826`^, etc. | ✅︎ | ✅︎ | -| `RobertaModel`, `RobertaForMaskedLM` | RoBERTa-based | `sentence-transformers/all-roberta-large-v1`, etc. | | | +| `RobertaModel`, `RobertaForMaskedLM` | RoBERTa-based | `sentence-transformers/all-roberta-large-v1`, etc. | ✅︎ | | | `VoyageQwen3BidirectionalEmbedModel`C | Voyage Qwen3-based with bidirectional attention | `voyageai/voyage-4-nano`, etc. | ✅︎ | ✅︎ | -| `XLMRobertaModel` | XLMRobertaModel-based | `BAAI/bge-m3` (see note), `intfloat/multilingual-e5-base`, `jinaai/jina-embeddings-v3` (see note), etc. | | | +| `XLMRobertaModel` | XLMRobertaModel-based | `BAAI/bge-m3` (see note), `intfloat/multilingual-e5-base`, `jinaai/jina-embeddings-v3` (see note), etc. | ✅︎ | | | `*Model`C, `*ForCausalLM`C, etc. | Generative models | N/A | \* | \* | !!! note diff --git a/vllm/model_executor/models/bert.py b/vllm/model_executor/models/bert.py index cdaa0bce9725..508748681261 100644 --- a/vllm/model_executor/models/bert.py +++ b/vllm/model_executor/models/bert.py @@ -41,7 +41,7 @@ from vllm.tasks import PoolingTask from vllm.v1.pool.metadata import PoolingMetadata -from .interfaces import SupportsCrossEncoding, SupportsQuant +from .interfaces import SupportsCrossEncoding, SupportsLoRA, SupportsQuant from .interfaces_base import attn_type, default_pooling_type from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix @@ -448,7 +448,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @default_pooling_type(seq_pooling_type="CLS") -class BertEmbeddingModel(nn.Module, SupportsQuant): +class BertEmbeddingModel(nn.Module, SupportsQuant, SupportsLoRA): """A model that uses Bert to provide embedding functionalities. This class encapsulates the BertModel and provides an interface for @@ -463,9 +463,22 @@ class BertEmbeddingModel(nn.Module, SupportsQuant): is_pooling_model = True embedding_class: type[nn.Module] = BertEmbedding + packed_modules_mapping = {"qkv_proj": ["query", "key", "value"]} + # PEFT adapters name modules relative to the HF backbone (optionally under + # `bert.`/`roberta.`), while vLLM nests the backbone under `model.`. + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "bert.": "model.", + "roberta.": "model.", + "encoder.": "model.encoder.", + "embeddings.": "model.embeddings.", + } + ) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() + self.config = vllm_config.model_config.hf_config pooler_config = vllm_config.model_config.pooler_config assert pooler_config is not None From ffec221276fa39d7927b9bb63dc12766d77fc403 Mon Sep 17 00:00:00 2001 From: Shang Wang Date: Mon, 12 Oct 2026 10:23:00 -0400 Subject: [PATCH 311/624] [Bugfix][DSv4.1] Match decoder-replay input_ids dtype for FlashInfer MegaMoE (#61041) Signed-off-by: Shang Wang Co-authored-by: Claude Opus 5.5 --- .../test_deepseek_v41_decoder_replay_layers.py | 17 +++++++++++++++-- .../nvidia/decoder_replay_cudagraph.py | 7 ++++--- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/tests/models/test_deepseek_v41_decoder_replay_layers.py b/tests/models/test_deepseek_v41_decoder_replay_layers.py index 614d903ca10d..b7f5ca29ec2d 100644 --- a/tests/models/test_deepseek_v41_decoder_replay_layers.py +++ b/tests/models/test_deepseek_v41_decoder_replay_layers.py @@ -8,6 +8,7 @@ import torch from vllm.config import CompilationConfig, CUDAGraphMode +from vllm.config.kernel import FLASHINFER_MOE_EP_BACKENDS, NATIVE_MEGA_MOE_BACKENDS from vllm.forward_context import ( BatchDescriptor, ForwardContext, @@ -144,8 +145,18 @@ def test_piecewise_graphs_break_out_by_size(): write_batch_kv.assert_called_once_with(*states) -def test_replay_graph_matches_eager(monkeypatch): - """The replay graph of the next captured size pads and runs the rows.""" +@pytest.mark.parametrize( + ("moe_backend", "ids_dtype"), + [ + ("auto", torch.int32), + *((b, torch.int64) for b in sorted(NATIVE_MEGA_MOE_BACKENDS)), + *((b, torch.int32) for b in sorted(FLASHINFER_MOE_EP_BACKENDS)), + ], + ids=str, +) +def test_replay_graph_matches_eager(monkeypatch, moe_backend, ids_dtype): + """The replay graph of the next captured size pads and runs the rows. Input ids + arrive as the model passes them: int32, cast to int64 for native MegaMoE only.""" monkeypatch.setattr(cudagraph_utils, "get_pp_group", MagicMock) capture = torch.cuda.stream(torch.cuda.Stream()) monkeypatch.setattr(cudagraph_utils, "graph_capture", lambda device: capture) @@ -157,6 +168,7 @@ def test_replay_graph_matches_eager(monkeypatch): cfg.cache_config.use_kda_recoverssm = False cfg.model_config.hf_config = MagicMock(hc_mult=HC, hidden_size=HIDDEN) cfg.model_config.dtype = torch.float32 + cfg.kernel_config.moe_backend = moe_backend def run_layers(hidden, positions, ids, pre, post, mix, residual): return residual + hidden[:, None] + positions[:, None, None], pre + ids[:, None] @@ -173,6 +185,7 @@ def prepare(desc): torch.randn(12, *b.shape[1:], device=DEVICE).to(b.dtype) for b in manager.input_buffers ] + states[2] = torch.arange(12, device=DEVICE, dtype=ids_dtype) # input_ids rows = torch.tensor([11, 2, 7], device=DEVICE) desc = manager.dispatch(2, 3, None, 0) with override_forward_context(_graph_context(desc.num_tokens)): diff --git a/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py b/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py index 871f535584c6..c985f973bb49 100644 --- a/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py +++ b/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py @@ -10,7 +10,7 @@ from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper from vllm.config import CUDAGraphMode, VllmConfig -from vllm.config.kernel import MEGA_MOE_BACKENDS +from vllm.config.kernel import NATIVE_MEGA_MOE_BACKENDS from vllm.forward_context import get_forward_context, override_forward_context from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers from vllm.v1.worker.gpu.cudagraph_utils import ( @@ -56,8 +56,9 @@ def __init__(self, vllm_config: VllmConfig, device, layers: DecoderReplayLayers) config, act = vllm_config.model_config.hf_config, vllm_config.model_config.dtype assert isinstance(act, torch.dtype) hc, h, f32 = config.hc_mult, config.hidden_size, torch.float32 - mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS - ids = torch.int64 if mega_moe else torch.int32 + # DeepseekV4Model.forward casts input_ids to int64 for native MegaMoE only. + native = vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS + ids = torch.int64 if native else torch.int32 # hidden_states, positions, input_ids, pre_mix, post_mix, res_mix, residual shapes = [(h,), (), (), (hc,), (hc, 1), (hc, hc), (hc, h)] dtypes = [act, torch.int64, ids, f32, f32, f32, act] From a686c8e1b7e5870eadfc9f47447de2b8a26211fc Mon Sep 17 00:00:00 2001 From: Nick Riasanovsky Date: Mon, 12 Oct 2026 11:38:43 -0400 Subject: [PATCH 312/624] [Triton] Migrate Triton kernels off tl.make_block_ptr for triton-main compat (#60502) Signed-off-by: Nick Riasanovsky --- .../amd/ops/third_party/kda/kernels.py | 357 +++++----------- .../nvidia/ops/third_party/kda/kernels.py | 357 +++++----------- .../inkling/amd/ops/fa4_rel_attention.py | 40 +- .../nvidia/ops/third_party/kda/chunk.py | 260 ++++------- .../nvidia/ops/third_party/kda/chunk_intra.py | 274 +++++------- .../kda/chunk_intra_token_parallel.py | 16 +- .../minimax_m3/common/ops/index_topk.py | 18 +- .../ops/chunk_delta_h.py | 182 ++++---- .../flash_linear_attention/ops/chunk_o.py | 49 +-- .../ops/chunk_scaled_dot_kkt.py | 35 +- .../flash_linear_attention/ops/cumsum.py | 69 +-- .../ops/fused_norm_gate.py | 34 +- .../flash_linear_attention/ops/kda.py | 402 ++++++------------ .../flash_linear_attention/ops/l2norm.py | 15 +- .../flash_linear_attention/ops/solve_tril.py | 172 +++----- .../flash_linear_attention/ops/wy_fast.py | 75 ++-- 16 files changed, 838 insertions(+), 1517 deletions(-) diff --git a/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py b/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py index 9805b3747130..0e32c9454351 100644 --- a/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py +++ b/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py @@ -228,13 +228,13 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_c, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H i_i, i_j = i_c // NC, i_c % NC if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -255,58 +255,51 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( A += (bos * H + i_h) * BT Aqk += (bos * H + i_h) * BT - p_b = tl.make_block_ptr( - beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,) - ) - b_b = tl.load(p_b, boundary_check=(0,)) + o_i = i_t * BT + i_i * BC + tl.arange(0, BC) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + m_i = o_i < T + m_j = o_j < T + p_b = beta + bos * H + i_h + o_i * H + b_b = tl.load(p_b, mask=m_i, other=0.0) b_A = tl.zeros([BC, BC], dtype=tl.float32) b_Aqk = tl.zeros([BC, BC], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_k = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_g = tl.make_block_ptr( - g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - b_kt = tl.make_block_ptr( - k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - p_gk = tl.make_block_ptr( - g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K + m_qk = m_i[:, None] & m_k[None, :] + m_kj = m_k[:, None] & m_j[None, :] + + p_q = q + o_i[:, None] * (H * K) + o_k[None, :] + p_k = k + o_i[:, None] * (H * K) + o_k[None, :] + p_g = g + o_i[:, None] * (H * K) + o_k[None, :] + b_kt = k + o_k[:, None] + o_j[None, :] * (H * K) + p_gk = g + o_k[:, None] + o_j[None, :] * (H * K) + # [BK,] b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0) # [BC, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) * exp2(b_g - b_gn[None, :]) + b_g = tl.load(p_g, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) * exp2(b_g - b_gn[None, :]) # [BK, BC] - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kt = tl.load(b_kt, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, mask=m_kj, other=0.0) + b_kt = tl.load(b_kt, mask=m_kj, other=0.0) # [BC, BC] b_ktg = b_kt * exp2(b_gn[:, None] - b_gk) b_A += tl.dot(b_k, b_ktg) - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale b_Aqk += tl.dot(b_qg, b_ktg) b_A *= b_b[:, None] - p_A = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) - tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) - p_Aqk = tl.make_block_ptr( - Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) - tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + o_jA = i_j * BC + tl.arange(0, BC) + m_A = m_i[:, None] & (o_jA[None, :] < BT) + p_A = A + o_i[:, None] * (H * BT) + o_jA[None, :] + tl.store(p_A, b_A.to(A.dtype.element_ty), mask=m_A) + p_Aqk = Aqk + o_i[:, None] * (H * BT) + o_jA[None, :] + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), mask=m_A) @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @@ -333,12 +326,12 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( BK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -357,33 +350,14 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( m_A = (i_t * BT + i_i * BC + o_i) < T o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) + o_c = i_t * BT + i_i * BC + tl.arange(0, BC) + m_qk = m_A[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) + b_g = tl.load(p_g, mask=m_qk, other=0.0) p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] @@ -536,12 +510,12 @@ def recompute_w_u_fwd_kernel( IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -550,108 +524,57 @@ def recompute_w_u_fwd_kernel( T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_b = beta + bos * H + i_h + o_t * H + b_b = tl.load(p_b, mask=m_t, other=0.0) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) - b_A = tl.load(p_A, boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_A, other=0.0) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_u = tl.make_block_ptr( - u + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_u = u + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - p_w = tl.make_block_ptr( - w + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] + p_w = w + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_kb = b_k * b_b[:, None] - p_gk = tl.make_block_ptr( - gk + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + p_gk = gk + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_gk = tl.load(p_gk, mask=m_tk, other=0.0) b_kb *= exp2(b_gk) if STORE_QG: - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_qg = tl.make_block_ptr( - qg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_qg = qg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_tk, other=0.0) b_qg = b_q * exp2(b_gk) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), mask=m_tk) if STORE_KG: last_idx = min(i_t * BT + BT, T) - 1 - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K b_gn = tl.load( gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 ) b_kg = b_k * exp2(b_gn - b_gk) - p_kg = tl.make_block_ptr( - kg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + p_kg = kg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), mask=m_tk) b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_tk) def recompute_w_u_fwd( @@ -732,13 +655,13 @@ def chunk_gla_fwd_kernel_o( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -754,70 +677,44 @@ def chunk_gla_fwd_kernel_o( m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] b_o = tl.zeros([BT, BV], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + o_i = tl.arange(0, BT) + m_t = o_t < T + m_v = o_v < V + m_tv = m_t[:, None] & m_v[None, :] + m_A = m_t[:, None] & (o_i[None, :] < BT) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_h = tl.make_block_ptr( - h + (i_tg * H + i_h) * K * V, - (V, K), - (K, 1), - (i_v * BV, i_k * BK), - (BV, BK), - (1, 0), - ) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_qk = m_t[:, None] & m_k[None, :] + m_h = m_v[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_h = h + (i_tg * H + i_h) * K * V + o_v[:, None] * K + o_k[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BT, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, mask=m_qk, other=0.0) # [BT, BK] b_qg = (b_q * exp2(b_g)).to(b_q.dtype) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BT, BV] if i_k >= 0: b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_o = tl.make_block_ptr( - o + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_o = o + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_i[None, :] # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_A, other=0.0) b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) b_o += tl.dot(b_A, b_v, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_tv) def chunk_gla_fwd_o_gk( @@ -896,12 +793,12 @@ def kda_gate_cumsum_fwd_kernel( HAS_BIAS: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -911,26 +808,15 @@ def kda_gate_cumsum_fwd_kernel( else: bos = i_b * T - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), - ) - p_y = tl.make_block_ptr( - y + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), - ) + o_t = i_t * BT + tl.arange(0, BT) + o_d = i_d * BD + tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + p_g = g + (bos * H + i_h) * D + o_t[:, None] * (H * D) + o_d[None, :] + p_y = y + (bos * H + i_h) * D + o_t[:, None] * (H * D) + o_d[None, :] - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(p_g, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: - o_d = i_d * BD + tl.arange(0, BD) b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) b_g = b_g + b_bias[None, :] @@ -951,12 +837,12 @@ def kda_gate_cumsum_fwd_kernel( # Out-of-bounds rows (load returns 0, but softplus/bias can still make # b_gate non-zero) participate in the dot product. They only contribute to - # out-of-bounds output rows, which are masked away by `boundary_check` on + # out-of-bounds output rows, which are masked away by `mask` on # the store, so visible output matches unfused gate + chunk-local cumsum. - o_t = tl.arange(0, BT) - m_cumsum = tl.where(o_t[:, None] >= o_t[None, :], 1.0, 0.0) + o_c = tl.arange(0, BT) + m_cumsum = tl.where(o_c[:, None] >= o_c[None, :], 1.0, 0.0) b_y = tl.dot(m_cumsum, b_gate, allow_tf32=False) * cumsum_scale - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_g) def fused_kda_gate_chunk_cumsum( @@ -1271,7 +1157,7 @@ def kda_gate_fwd_kernel( BD: tl.constexpr, HAS_BIAS: tl.constexpr, ): - i_t, i_h = tl.program_id(0), tl.program_id(1) + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1) n_t = i_t * BT b_a = tl.load(A + i_h).to(tl.float32) @@ -1280,25 +1166,14 @@ def kda_gate_fwd_kernel( stride_row = H * D stride_col = 1 - g_ptr = tl.make_block_ptr( - base=g + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) - - y_ptr = tl.make_block_ptr( - base=y + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) + o_t = n_t + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + g_ptr = g + i_h * D + o_t[:, None] * stride_row + o_d[None, :] * stride_col + y_ptr = y + i_h * D + o_t[:, None] * stride_row + o_d[None, :] * stride_col - b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(g_ptr, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: n_d = tl.arange(0, BD) @@ -1321,7 +1196,7 @@ def kda_gate_fwd_kernel( sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) b_y = b_a * sp - tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(y_ptr, b_y.to(y.dtype.element_ty), mask=m_g) def fused_kda_gate( diff --git a/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py b/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py index 8f662f68aa24..da498b5855ed 100644 --- a/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py +++ b/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py @@ -230,13 +230,13 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_c, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H i_i, i_j = i_c // NC, i_c % NC if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -257,58 +257,51 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( A += (bos * H + i_h) * BT Aqk += (bos * H + i_h) * BT - p_b = tl.make_block_ptr( - beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,) - ) - b_b = tl.load(p_b, boundary_check=(0,)) + o_i = i_t * BT + i_i * BC + tl.arange(0, BC) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + m_i = o_i < T + m_j = o_j < T + p_b = beta + bos * H + i_h + o_i * H + b_b = tl.load(p_b, mask=m_i, other=0.0) b_A = tl.zeros([BC, BC], dtype=tl.float32) b_Aqk = tl.zeros([BC, BC], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_k = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_g = tl.make_block_ptr( - g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - b_kt = tl.make_block_ptr( - k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - p_gk = tl.make_block_ptr( - g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K + m_qk = m_i[:, None] & m_k[None, :] + m_kj = m_k[:, None] & m_j[None, :] + + p_q = q + o_i[:, None] * (H * K) + o_k[None, :] + p_k = k + o_i[:, None] * (H * K) + o_k[None, :] + p_g = g + o_i[:, None] * (H * K) + o_k[None, :] + b_kt = k + o_k[:, None] + o_j[None, :] * (H * K) + p_gk = g + o_k[:, None] + o_j[None, :] * (H * K) + # [BK,] b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0) # [BC, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) * exp2(b_g - b_gn[None, :]) + b_g = tl.load(p_g, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) * exp2(b_g - b_gn[None, :]) # [BK, BC] - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kt = tl.load(b_kt, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, mask=m_kj, other=0.0) + b_kt = tl.load(b_kt, mask=m_kj, other=0.0) # [BC, BC] b_ktg = b_kt * exp2(b_gn[:, None] - b_gk) b_A += tl.dot(b_k, b_ktg) - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale b_Aqk += tl.dot(b_qg, b_ktg) b_A *= b_b[:, None] - p_A = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) - tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) - p_Aqk = tl.make_block_ptr( - Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) - tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + o_jA = i_j * BC + tl.arange(0, BC) + m_A = m_i[:, None] & (o_jA[None, :] < BT) + p_A = A + o_i[:, None] * (H * BT) + o_jA[None, :] + tl.store(p_A, b_A.to(A.dtype.element_ty), mask=m_A) + p_Aqk = Aqk + o_i[:, None] * (H * BT) + o_jA[None, :] + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), mask=m_A) @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @@ -335,12 +328,12 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( BK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -359,33 +352,14 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( m_A = (i_t * BT + i_i * BC + o_i) < T o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) + o_c = i_t * BT + i_i * BC + tl.arange(0, BC) + m_qk = m_A[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) + b_g = tl.load(p_g, mask=m_qk, other=0.0) p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] @@ -538,12 +512,12 @@ def recompute_w_u_fwd_kernel( IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -552,108 +526,57 @@ def recompute_w_u_fwd_kernel( T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_b = beta + bos * H + i_h + o_t * H + b_b = tl.load(p_b, mask=m_t, other=0.0) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) - b_A = tl.load(p_A, boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_A, other=0.0) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_u = tl.make_block_ptr( - u + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_u = u + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - p_w = tl.make_block_ptr( - w + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] + p_w = w + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_kb = b_k * b_b[:, None] - p_gk = tl.make_block_ptr( - gk + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + p_gk = gk + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_gk = tl.load(p_gk, mask=m_tk, other=0.0) b_kb *= exp2(b_gk) if STORE_QG: - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_qg = tl.make_block_ptr( - qg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_qg = qg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_tk, other=0.0) b_qg = b_q * exp2(b_gk) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), mask=m_tk) if STORE_KG: last_idx = min(i_t * BT + BT, T) - 1 - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K b_gn = tl.load( gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 ) b_kg = b_k * exp2(b_gn - b_gk) - p_kg = tl.make_block_ptr( - kg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + p_kg = kg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), mask=m_tk) b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_tk) def recompute_w_u_fwd( @@ -734,13 +657,13 @@ def chunk_gla_fwd_kernel_o( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -756,70 +679,44 @@ def chunk_gla_fwd_kernel_o( m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] b_o = tl.zeros([BT, BV], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + o_i = tl.arange(0, BT) + m_t = o_t < T + m_v = o_v < V + m_tv = m_t[:, None] & m_v[None, :] + m_A = m_t[:, None] & (o_i[None, :] < BT) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_h = tl.make_block_ptr( - h + (i_tg * H + i_h) * K * V, - (V, K), - (K, 1), - (i_v * BV, i_k * BK), - (BV, BK), - (1, 0), - ) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_qk = m_t[:, None] & m_k[None, :] + m_h = m_v[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_h = h + (i_tg * H + i_h) * K * V + o_v[:, None] * K + o_k[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BT, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, mask=m_qk, other=0.0) # [BT, BK] b_qg = (b_q * exp2(b_g)).to(b_q.dtype) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BT, BV] if i_k >= 0: b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_o = tl.make_block_ptr( - o + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_o = o + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_i[None, :] # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_A, other=0.0) b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) b_o += tl.dot(b_A, b_v, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_tv) def chunk_gla_fwd_o_gk( @@ -898,12 +795,12 @@ def kda_gate_cumsum_fwd_kernel( HAS_BIAS: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -913,26 +810,15 @@ def kda_gate_cumsum_fwd_kernel( else: bos = i_b * T - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), - ) - p_y = tl.make_block_ptr( - y + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), - ) + o_t = i_t * BT + tl.arange(0, BT) + o_d = i_d * BD + tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + p_g = g + (bos * H + i_h) * D + o_t[:, None] * (H * D) + o_d[None, :] + p_y = y + (bos * H + i_h) * D + o_t[:, None] * (H * D) + o_d[None, :] - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(p_g, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: - o_d = i_d * BD + tl.arange(0, BD) b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) b_g = b_g + b_bias[None, :] @@ -953,12 +839,12 @@ def kda_gate_cumsum_fwd_kernel( # Out-of-bounds rows (load returns 0, but softplus/bias can still make # b_gate non-zero) participate in the dot product. They only contribute to - # out-of-bounds output rows, which are masked away by `boundary_check` on + # out-of-bounds output rows, which are masked away by `mask` on # the store, so visible output matches unfused gate + chunk-local cumsum. - o_t = tl.arange(0, BT) - m_cumsum = tl.where(o_t[:, None] >= o_t[None, :], 1.0, 0.0) + o_c = tl.arange(0, BT) + m_cumsum = tl.where(o_c[:, None] >= o_c[None, :], 1.0, 0.0) b_y = tl.dot(m_cumsum, b_gate, allow_tf32=False) * cumsum_scale - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_g) def fused_kda_gate_chunk_cumsum( @@ -1273,7 +1159,7 @@ def kda_gate_fwd_kernel( BD: tl.constexpr, HAS_BIAS: tl.constexpr, ): - i_t, i_h = tl.program_id(0), tl.program_id(1) + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1) n_t = i_t * BT b_a = tl.load(A + i_h).to(tl.float32) @@ -1282,25 +1168,14 @@ def kda_gate_fwd_kernel( stride_row = H * D stride_col = 1 - g_ptr = tl.make_block_ptr( - base=g + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) - - y_ptr = tl.make_block_ptr( - base=y + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) + o_t = n_t + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + g_ptr = g + i_h * D + o_t[:, None] * stride_row + o_d[None, :] * stride_col + y_ptr = y + i_h * D + o_t[:, None] * stride_row + o_d[None, :] * stride_col - b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(g_ptr, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: n_d = tl.arange(0, BD) @@ -1323,7 +1198,7 @@ def kda_gate_fwd_kernel( sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) b_y = b_a * sp - tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(y_ptr, b_y.to(y.dtype.element_ty), mask=m_g) def fused_kda_gate( diff --git a/vllm/models/inkling/amd/ops/fa4_rel_attention.py b/vllm/models/inkling/amd/ops/fa4_rel_attention.py index d448ac8e2c5f..936b51eb6ea4 100644 --- a/vllm/models/inkling/amd/ops/fa4_rel_attention.py +++ b/vllm/models/inkling/amd/ops/fa4_rel_attention.py @@ -165,15 +165,23 @@ def _inkling_rel_attention_kernel( q_head_start = pid_kh * gqa_group_size bt_row = block_table_ptr + pid_b * stride_bt_b - q_block_ptr = tl.make_block_ptr( - base=q_ptr + q_start * stride_q_t + q_head_start * stride_q_h, - shape=(q_len, gqa_group_size, head_dim), - strides=(stride_q_t, stride_q_h, stride_q_d), - offsets=(q_block, 0, 0), - block_shape=(BLOCK_Q, BLOCK_H, BLOCK_D), - order=(2, 1, 0), + o0 = q_block + tl.arange(0, BLOCK_Q) + o1 = tl.arange(0, BLOCK_H) + o2 = tl.arange(0, BLOCK_D) + m_q = ( + (o0 < q_len)[:, None, None] + & (o1 < gqa_group_size)[None, :, None] + & (o2 < head_dim)[None, None, :] ) - q = tl.load(q_block_ptr, boundary_check=(0, 1, 2), padding_option="zero") + q_block_ptr = ( + q_ptr + + q_start * stride_q_t + + q_head_start * stride_q_h + + o0[:, None, None] * stride_q_t + + o1[None, :, None] * stride_q_h + + o2[None, None, :] * stride_q_d + ) + q = tl.load(q_block_ptr, mask=m_q, other=0.0) q = tl.reshape(q, (BLOCK_QH, BLOCK_D)) q_rows = q_block + tl.arange(0, BLOCK_Q) @@ -265,18 +273,18 @@ def _inkling_rel_attention_kernel( acc /= l_i[:, None] acc = tl.reshape(acc, (BLOCK_Q, BLOCK_H, BLOCK_D)) - out_block_ptr = tl.make_block_ptr( - base=out_ptr + q_start * stride_o_t + q_head_start * stride_o_h, - shape=(q_len, gqa_group_size, head_dim), - strides=(stride_o_t, stride_o_h, stride_o_d), - offsets=(q_block, 0, 0), - block_shape=(BLOCK_Q, BLOCK_H, BLOCK_D), - order=(2, 1, 0), + out_block_ptr = ( + out_ptr + + q_start * stride_o_t + + q_head_start * stride_o_h + + o0[:, None, None] * stride_o_t + + o1[None, :, None] * stride_o_h + + o2[None, None, :] * stride_o_d ) tl.store( out_block_ptr, acc.to(out_ptr.dtype.element_ty), - boundary_check=(0, 1, 2), + mask=m_q, ) diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py index 3911d20d0358..19fcd9a686bf 100644 --- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py @@ -69,12 +69,12 @@ def recompute_w_u_fwd_kernel( IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -83,108 +83,57 @@ def recompute_w_u_fwd_kernel( T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_b = beta + bos * H + i_h + o_t * H + b_b = tl.load(p_b, mask=m_t, other=0.0).to(tl.float32) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) - b_A = tl.load(p_A, boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_A, other=0.0) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_u = tl.make_block_ptr( - u + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_u = u + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - p_w = tl.make_block_ptr( - w + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] + p_w = w + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_kb = b_k * b_b[:, None] - p_gk = tl.make_block_ptr( - gk + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + p_gk = gk + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_gk = tl.load(p_gk, mask=m_tk, other=0.0) b_kb *= exp2(b_gk) if STORE_QG: - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_qg = tl.make_block_ptr( - qg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_qg = qg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_tk, other=0.0) b_qg = b_q * exp2(b_gk) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), mask=m_tk) if STORE_KG: last_idx = min(i_t * BT + BT, T) - 1 - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K b_gn = tl.load( gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 ) b_kg = b_k * exp2(b_gn - b_gk) - p_kg = tl.make_block_ptr( - kg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + p_kg = kg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), mask=m_tk) b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_tk) def recompute_w_u_fwd( @@ -265,13 +214,13 @@ def chunk_gla_fwd_kernel_o( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -287,70 +236,44 @@ def chunk_gla_fwd_kernel_o( m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] b_o = tl.zeros([BT, BV], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + o_i = tl.arange(0, BT) + m_t = o_t < T + m_v = o_v < V + m_tv = m_t[:, None] & m_v[None, :] + m_A = m_t[:, None] & (o_i[None, :] < BT) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_h = tl.make_block_ptr( - h + (i_tg * H + i_h) * K * V, - (V, K), - (K, 1), - (i_v * BV, i_k * BK), - (BV, BK), - (1, 0), - ) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_qk = m_t[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_h = h + (i_tg * H + i_h) * K * V + o_v[:, None] * K + o_k[None, :] + m_h = m_v[:, None] & m_k[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BT, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, mask=m_qk, other=0.0) # [BT, BK] b_qg = (b_q * exp2(b_g)).to(b_q.dtype) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BT, BV] if i_k >= 0: b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_o = tl.make_block_ptr( - o + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_o = o + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_i[None, :] # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_A, other=0.0) b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) b_o += tl.dot(b_A, b_v, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_tv) def chunk_gla_fwd_o_gk( @@ -434,12 +357,12 @@ def kda_gate_chunk_cumsum_vector_kernel( IS_VARLEN: tl.constexpr, USE_LOWER_BOUND: tl.constexpr, ): - i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -472,34 +395,15 @@ def kda_gate_chunk_cumsum_vector_kernel( i_s -= 1 - p_s = tl.make_block_ptr( - s + (bos * H + i_h) * S, - (T, S), - (H * S, 1), - (i_t * BT, i_s * BS), - (BT, BS), - (1, 0), - ) - p_o = tl.make_block_ptr( - o + (bos * H + i_h) * S, - (T, S), - (H * S, 1), - (i_t * BT, i_s * BS), - (BT, BS), - (1, 0), - ) + o_t = i_t * BT + tl.arange(0, BT) + o_s = i_s * BS + tl.arange(0, BS) + m_s = (o_t[:, None] < T) & (o_s[None, :] < S) + p_s = s + (bos * H + i_h) * S + o_t[:, None] * (H * S) + o_s[None, :] + p_o = o + (bos * H + i_h) * S + o_t[:, None] * (H * S) + o_s[None, :] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + b_s = tl.load(p_s, mask=m_s, other=0.0).to(tl.float32) if HAS_BIAS: - p_bias = tl.make_block_ptr( - g_bias + i_h * S, - (S,), - (1,), - (i_s * BS,), - (BS,), - (0,), - ) - b_bias = tl.load(p_bias, boundary_check=(0,)).to(tl.float32) + b_bias = tl.load(g_bias + i_h * S + o_s, mask=o_s < S, other=0.0).to(tl.float32) b_s += b_bias[None, :] b_a = tl.exp(tl.load(A_log + i_h).to(tl.float32)) @@ -517,7 +421,7 @@ def kda_gate_chunk_cumsum_vector_kernel( # Boundary loads return zero, but bias and gate activation can make padded # rows nonzero. Padding trails valid rows, so it only affects masked stores. b_o = tl.cumsum(b_gate, axis=0) * cumsum_scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_s) def fused_kda_gate_chunk_cumsum( @@ -843,33 +747,19 @@ def kda_gate_fwd_kernel( HAS_BIAS: tl.constexpr, USE_LOWER_BOUND: tl.constexpr, ): - i_t, i_h = tl.program_id(0), tl.program_id(1) + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1) n_t = i_t * BT b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) - stride_row = H * D - stride_col = 1 - - g_ptr = tl.make_block_ptr( - base=g + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) - - y_ptr = tl.make_block_ptr( - base=y + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) + o_t = n_t + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + g_ptr = g + i_h * D + o_t[:, None] * (H * D) + o_d[None, :] + y_ptr = y + i_h * D + o_t[:, None] * (H * D) + o_d[None, :] - b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(g_ptr, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: n_d = tl.arange(0, BD) @@ -887,7 +777,7 @@ def kda_gate_fwd_kernel( sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) b_y = -b_a * sp - tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(y_ptr, b_y.to(y.dtype.element_ty), mask=m_g) def fused_kda_gate( diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py index 07b4a388936e..286d23c814d0 100644 --- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py @@ -70,14 +70,14 @@ def chunk_kda_fwd_kernel_inter_solve_fused( 5. Computes merged Akk_inv 6. Writes Akk_inv to Akk """ - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1) i_b, i_hv = i_bh // HV, i_bh % HV i_h = i_hv // (HV // H) if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -106,6 +106,15 @@ def chunk_kda_fwd_kernel_inter_solve_fused( m_tc1 = (i_tc1 + o_i) < T m_tc2 = (i_tc2 + o_i) < T m_tc3 = (i_tc3 + o_i) < T + o_c0 = i_tc0 + o_i + o_c1 = i_tc1 + o_i + o_c2 = i_tc2 + o_i + o_c3 = i_tc3 + o_i + m_tc0 = o_c0 < T + m_A0 = m_tc0[:, None] & (o_i[None, :] < BT) + m_A1 = m_tc1[:, None] & (o_i[None, :] < BT) + m_A2 = m_tc2[:, None] & (o_i[None, :] < BT) + m_A3 = m_tc3[:, None] & (o_i[None, :] < BT) b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32) b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32) @@ -129,29 +138,21 @@ def chunk_kda_fwd_kernel_inter_solve_fused( o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K - p_k0 = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) - ) - p_g0 = tl.make_block_ptr( - g, (T, K), (HV * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) - ) - b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) - b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) + m_ck0 = m_tc0[:, None] & m_k[None, :] + p_k0 = k + o_c0[:, None] * (H * K) + o_k[None, :] + p_g0 = g + o_c0[:, None] * (HV * K) + o_k[None, :] + b_k0 = tl.load(p_k0, mask=m_ck0, other=0.0).to(tl.float32) + b_g0 = tl.load(p_g0, mask=m_ck0, other=0.0).to(tl.float32) if i_tc1 < T: - p_q1 = tl.make_block_ptr( - q, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) - ) - p_k1 = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) - ) - p_g1 = tl.make_block_ptr( - g, (T, K), (HV * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) - ) + m_ck1 = m_tc1[:, None] & m_k[None, :] + p_q1 = q + o_c1[:, None] * (H * K) + o_k[None, :] + p_k1 = k + o_c1[:, None] * (H * K) + o_k[None, :] + p_g1 = g + o_c1[:, None] * (HV * K) + o_k[None, :] # [BC, BK] - b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32) - b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32) - b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32) + b_q1 = tl.load(p_q1, mask=m_ck1, other=0.0).to(tl.float32) + b_k1 = tl.load(p_k1, mask=m_ck1, other=0.0).to(tl.float32) + b_g1 = tl.load(p_g1, mask=m_ck1, other=0.0).to(tl.float32) # [BK] b_gn1 = tl.load(g + i_tc1 * HV * K + o_k, mask=m_k, other=0).to(tl.float32) # [BC, BK] @@ -163,19 +164,14 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt) if i_tc2 < T: - p_q2 = tl.make_block_ptr( - q, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) - ) - p_k2 = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) - ) - p_g2 = tl.make_block_ptr( - g, (T, K), (HV * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) - ) + m_ck2 = m_tc2[:, None] & m_k[None, :] + p_q2 = q + o_c2[:, None] * (H * K) + o_k[None, :] + p_k2 = k + o_c2[:, None] * (H * K) + o_k[None, :] + p_g2 = g + o_c2[:, None] * (HV * K) + o_k[None, :] # [BC, BK] - b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) - b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) - b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) + b_q2 = tl.load(p_q2, mask=m_ck2, other=0.0).to(tl.float32) + b_k2 = tl.load(p_k2, mask=m_ck2, other=0.0).to(tl.float32) + b_g2 = tl.load(p_g2, mask=m_ck2, other=0.0).to(tl.float32) # [BK] b_gn2 = tl.load(g + i_tc2 * HV * K + o_k, mask=m_k, other=0).to( tl.float32 @@ -195,19 +191,14 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Akk21 += tl.dot(b_kg2, b_kgt) if i_tc3 < T: - p_q3 = tl.make_block_ptr( - q, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) - ) - p_k3 = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) - ) - p_g3 = tl.make_block_ptr( - g, (T, K), (HV * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) - ) + m_ck3 = m_tc3[:, None] & m_k[None, :] + p_q3 = q + o_c3[:, None] * (H * K) + o_k[None, :] + p_k3 = k + o_c3[:, None] * (H * K) + o_k[None, :] + p_g3 = g + o_c3[:, None] * (HV * K) + o_k[None, :] # [BC, BK] - b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) - b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) - b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) + b_q3 = tl.load(p_q3, mask=m_ck3, other=0.0).to(tl.float32) + b_k3 = tl.load(p_k3, mask=m_ck3, other=0.0).to(tl.float32) + b_g3 = tl.load(p_g3, mask=m_ck3, other=0.0).to(tl.float32) # [BK] b_gn3 = tl.load(g + i_tc3 * HV * K + o_k, mask=m_k, other=0).to( tl.float32 @@ -236,82 +227,44 @@ def chunk_kda_fwd_kernel_inter_solve_fused( # save off-diagonal Aqk blocks and prepare Akk ################################################################################ if i_tc1 < T: - p_Aqk10 = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0) - ) - tl.store( - p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) - ) + p_Aqk10 = Aqk + o_c1[:, None] * (HV * BT) + o_i[None, :] + tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), mask=m_A1) - p_b1 = tl.make_block_ptr( - beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,) - ) - b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) + p_b1 = beta + bos * HV + i_hv + o_c1 * HV + b_b1 = tl.load(p_b1, mask=m_tc1, other=0.0).to(tl.float32) b_Akk10 = b_Akk10 * b_b1[:, None] if i_tc2 < T: - p_Aqk20 = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0) - ) - p_Aqk21 = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0) - ) - tl.store( - p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) - ) - tl.store( - p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) - ) + p_Aqk20 = Aqk + o_c2[:, None] * (HV * BT) + o_i[None, :] + p_Aqk21 = Aqk + o_c2[:, None] * (HV * BT) + (o_i + BC)[None, :] + tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), mask=m_A2) + tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), mask=m_A2) - p_b2 = tl.make_block_ptr( - beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,) - ) - b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) + p_b2 = beta + bos * HV + i_hv + o_c2 * HV + b_b2 = tl.load(p_b2, mask=m_tc2, other=0.0).to(tl.float32) b_Akk20 = b_Akk20 * b_b2[:, None] b_Akk21 = b_Akk21 * b_b2[:, None] if i_tc3 < T: - p_Aqk30 = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0) - ) - p_Aqk31 = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0) - ) - p_Aqk32 = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0) - ) - tl.store( - p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) - ) - tl.store( - p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) - ) - tl.store( - p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) - ) - - p_b3 = tl.make_block_ptr( - beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,) - ) - b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) + p_Aqk30 = Aqk + o_c3[:, None] * (HV * BT) + o_i[None, :] + p_Aqk31 = Aqk + o_c3[:, None] * (HV * BT) + (o_i + BC)[None, :] + p_Aqk32 = Aqk + o_c3[:, None] * (HV * BT) + (o_i + 2 * BC)[None, :] + tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), mask=m_A3) + tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), mask=m_A3) + tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), mask=m_A3) + + p_b3 = beta + bos * HV + i_hv + o_c3 * HV + b_b3 = tl.load(p_b3, mask=m_tc3, other=0.0).to(tl.float32) b_Akk30 = b_Akk30 * b_b3[:, None] b_Akk31 = b_Akk31 * b_b3[:, None] b_Akk32 = b_Akk32 * b_b3[:, None] - p_Akk00 = tl.make_block_ptr( - Akkd, (T, BC), (HV * BC, 1), (i_tc0, 0), (BC, BC), (1, 0) - ) - p_Akk11 = tl.make_block_ptr( - Akkd, (T, BC), (HV * BC, 1), (i_tc1, 0), (BC, BC), (1, 0) - ) - p_Akk22 = tl.make_block_ptr( - Akkd, (T, BC), (HV * BC, 1), (i_tc2, 0), (BC, BC), (1, 0) - ) - p_Akk33 = tl.make_block_ptr( - Akkd, (T, BC), (HV * BC, 1), (i_tc3, 0), (BC, BC), (1, 0) - ) - b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) - b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) - b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) - b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) + p_Akk00 = Akkd + o_c0[:, None] * (HV * BC) + o_i[None, :] + p_Akk11 = Akkd + o_c1[:, None] * (HV * BC) + o_i[None, :] + p_Akk22 = Akkd + o_c2[:, None] * (HV * BC) + o_i[None, :] + p_Akk33 = Akkd + o_c3[:, None] * (HV * BC) + o_i[None, :] + b_Ai00 = tl.load(p_Akk00, mask=m_A0, other=0.0).to(tl.float32) + b_Ai11 = tl.load(p_Akk11, mask=m_A1, other=0.0).to(tl.float32) + b_Ai22 = tl.load(p_Akk22, mask=m_A2, other=0.0).to(tl.float32) + b_Ai33 = tl.load(p_Akk33, mask=m_A3, other=0.0).to(tl.float32) ################################################################################ # forward substitution on diagonals @@ -397,47 +350,27 @@ def chunk_kda_fwd_kernel_inter_solve_fused( # store full Akk_inv to Akk ################################################################################ - p_Akk00 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc0, 0), (BC, BC), (1, 0) - ) - p_Akk10 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0) - ) - p_Akk11 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc1, BC), (BC, BC), (1, 0) - ) - p_Akk20 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0) - ) - p_Akk21 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0) - ) - p_Akk22 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0) - ) - p_Akk30 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0) - ) - p_Akk31 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0) - ) - p_Akk32 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0) - ) - p_Akk33 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0) - ) - - tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + p_Akk00 = Akk + o_c0[:, None] * (HV * BT) + o_i[None, :] + p_Akk10 = Akk + o_c1[:, None] * (HV * BT) + o_i[None, :] + p_Akk11 = Akk + o_c1[:, None] * (HV * BT) + (o_i + BC)[None, :] + p_Akk20 = Akk + o_c2[:, None] * (HV * BT) + o_i[None, :] + p_Akk21 = Akk + o_c2[:, None] * (HV * BT) + (o_i + BC)[None, :] + p_Akk22 = Akk + o_c2[:, None] * (HV * BT) + (o_i + 2 * BC)[None, :] + p_Akk30 = Akk + o_c3[:, None] * (HV * BT) + o_i[None, :] + p_Akk31 = Akk + o_c3[:, None] * (HV * BT) + (o_i + BC)[None, :] + p_Akk32 = Akk + o_c3[:, None] * (HV * BT) + (o_i + 2 * BC)[None, :] + p_Akk33 = Akk + o_c3[:, None] * (HV * BT) + (o_i + 3 * BC)[None, :] + + tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), mask=m_A0) + tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), mask=m_A1) + tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), mask=m_A1) + tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), mask=m_A2) + tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), mask=m_A2) + tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), mask=m_A2) + tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), mask=m_A3) + tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), mask=m_A3) + tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), mask=m_A3) + tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), mask=m_A3) @triton.heuristics( @@ -474,14 +407,14 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( IS_VARLEN: tl.constexpr, USE_GATHER: tl.constexpr, ): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_hv = i_bh // HV, i_bh % HV i_h = i_hv // (HV // H) if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -505,16 +438,19 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( Aqk = Aqk + (bos * HV + i_hv) * BT Akk = Akk + (bos * HV + i_hv) * BC - p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (HV * K, 1), (i_ti, 0), (BC, BK), (1, 0)) + o_k = tl.arange(0, BK) + m_k = o_k < K + m_ck = m_c[:, None] & m_k[None, :] + p_q = q + o_c[:, None] * (H * K) + o_k[None, :] + p_k = k + o_c[:, None] * (H * K) + o_k[None, :] + p_g = g + o_c[:, None] * (HV * K) + o_k[None, :] - p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,)) + p_beta = beta + o_c * HV - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + b_q = tl.load(p_q, mask=m_ck, other=0.0) + b_k = tl.load(p_k, mask=m_ck, other=0.0) + b_g = tl.load(p_g, mask=m_ck, other=0.0) + b_beta = tl.load(p_beta, mask=m_c, other=0.0).to(tl.float32) if USE_GATHER: b_gn = gather( @@ -546,12 +482,12 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) b_Akk = tl.where(m_Akk, b_Akk, 0.0) - p_Aqk = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0) - ) - p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV * BC, 1), (i_ti, 0), (BC, BC), (1, 0)) - tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + m_Aqk_st = m_c[:, None] & (o_i[None, :] < BT) + m_Akk_st = m_c[:, None] & (o_i[None, :] < BC) + p_Aqk = Aqk + o_c[:, None] * (HV * BT) + (i_i * BC + o_i)[None, :] + p_Akk = Akk + o_c[:, None] * (HV * BC) + o_i[None, :] + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), mask=m_Aqk_st) + tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), mask=m_Akk_st) tl.debug_barrier() @@ -566,7 +502,7 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( b_a += tl.sum(b_a[:, None] * b_Ai, 0) b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) b_Ai += m_I - tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), mask=m_Akk_st) def chunk_kda_fwd_intra( diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py index 86654f264431..7cb68b2fd49a 100644 --- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py @@ -108,20 +108,16 @@ def chunk_kda_fwd_kernel_intra_token_parallel( b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) # g: [B, T, HV, K], beta: [B, T, HV] - p_g = tl.make_block_ptr( - g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0) - ) - p_beta = tl.make_block_ptr(beta + i_t * HV, (HV,), (1,), (i_hg * BH,), (BH,), (0,)) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + p_g = g + i_t * HV * K + o_hv[:, None] * K + o_k[None, :] + p_beta = beta + i_t * HV + o_hv + b_g = tl.load(p_g, mask=m_hk, other=0.0).to(tl.float32) + b_beta = tl.load(p_beta, mask=m_hv, other=0.0).to(tl.float32) b_k *= b_beta[:, None] for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))): b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) - p_gj = tl.make_block_ptr( - g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0) - ) - b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32) + p_gj = g + j * HV * K + o_hv[:, None] * K + o_k[None, :] + b_gj = tl.load(p_gj, mask=m_hk, other=0.0).to(tl.float32) b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0) b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale diff --git a/vllm/models/minimax_m3/common/ops/index_topk.py b/vllm/models/minimax_m3/common/ops/index_topk.py index 025c3dd428b4..2d612c12f254 100644 --- a/vllm/models/minimax_m3/common/ops/index_topk.py +++ b/vllm/models/minimax_m3/common/ops/index_topk.py @@ -115,15 +115,17 @@ def _index_block_score_kernel( if BLOCK_SIZE_Q * pid_q >= q_len: return - q_ptrs = tl.make_block_ptr( - base=q_ptr + seq_start * stride_q_n + pid_h * stride_q_h, - shape=(q_len, head_dim), - strides=(stride_q_n, stride_q_d), - offsets=(pid_q * BLOCK_SIZE_Q, 0), - block_shape=(BLOCK_SIZE_Q, head_dim), - order=(1, 0), + o_q = pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q) + o_d = tl.arange(0, head_dim) + m_q = o_q < q_len + q_ptrs = ( + q_ptr + + seq_start.to(tl.int64) * stride_q_n + + pid_h.to(tl.int64) * stride_q_h + + o_q[:, None].to(tl.int64) * stride_q_n + + o_d[None, :].to(tl.int64) * stride_q_d ) - q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero") + q = tl.load(q_ptrs, mask=m_q[:, None], other=0.0) # A matched fp8 pair (e.g. e4m3 x e4m3) lowers to a native FP8 MMA, so # upcast only when the operands differ: there is no mixed-dtype fp8 MMA, # including across fp8 flavours (fp8e4nv vs fp8e4b8). bf16 and fp32 loads diff --git a/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py b/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py index eb1c3af15297..cc27437fac58 100644 --- a/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py +++ b/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py @@ -110,109 +110,82 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( ht = ht + i_nh * V * K # load initial state + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + o_k1 = tl.arange(0, 64) + m_k1 = o_k1 < K + o_k2 = 64 + o_k1 + m_k2 = o_k2 < K + o_k3 = 128 + o_k1 + m_k3 = o_k3 < K + o_k4 = 192 + o_k1 + m_k4 = o_k4 < K if USE_INITIAL_STATE: - p_h0_1 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) - b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) + p_h0_1 = h0 + o_v[:, None] * K + o_k1[None, :] + m_h0_1 = m_v[:, None] & m_k1[None, :] + b_h1 += tl.load(p_h0_1, mask=m_h0_1, other=0.0).to(tl.float32) if K > 64: - p_h0_2 = tl.make_block_ptr( - h0, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0) - ) - b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) + p_h0_2 = h0 + o_v[:, None] * K + o_k2[None, :] + m_h0_2 = m_v[:, None] & m_k2[None, :] + b_h2 += tl.load(p_h0_2, mask=m_h0_2, other=0.0).to(tl.float32) if K > 128: - p_h0_3 = tl.make_block_ptr( - h0, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0) - ) - b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) + p_h0_3 = h0 + o_v[:, None] * K + o_k3[None, :] + m_h0_3 = m_v[:, None] & m_k3[None, :] + b_h3 += tl.load(p_h0_3, mask=m_h0_3, other=0.0).to(tl.float32) if K > 192: - p_h0_4 = tl.make_block_ptr( - h0, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0) - ) - b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) + p_h0_4 = h0 + o_v[:, None] * K + o_k4[None, :] + m_h0_4 = m_v[:, None] & m_k4[None, :] + b_h4 += tl.load(p_h0_4, mask=m_h0_4, other=0.0).to(tl.float32) # main recurrence for i_t in range(NT): - p_h1 = tl.make_block_ptr( - h + i_t.to(tl.int64) * stride_h, - (V, K), - (K, 1), - (i_v * BV, 0), - (BV, 64), - (1, 0), - ) - tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + o_t = (i_t * BT + tl.arange(0, BT)).to(tl.int64) + m_t = o_t < T + p_h1 = h + i_t.to(tl.int64) * stride_h + o_v[:, None] * K + o_k1[None, :] + m_h1 = m_v[:, None] & m_k1[None, :] + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), mask=m_h1) if K > 64: - p_h2 = tl.make_block_ptr( - h + i_t.to(tl.int64) * stride_h, - (V, K), - (K, 1), - (i_v * BV, 64), - (BV, 64), - (1, 0), - ) - tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + p_h2 = h + i_t.to(tl.int64) * stride_h + o_v[:, None] * K + o_k2[None, :] + m_h2 = m_v[:, None] & m_k2[None, :] + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), mask=m_h2) if K > 128: - p_h3 = tl.make_block_ptr( - h + i_t.to(tl.int64) * stride_h, - (V, K), - (K, 1), - (i_v * BV, 128), - (BV, 64), - (1, 0), - ) - tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + p_h3 = h + i_t.to(tl.int64) * stride_h + o_v[:, None] * K + o_k3[None, :] + m_h3 = m_v[:, None] & m_k3[None, :] + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), mask=m_h3) if K > 192: - p_h4 = tl.make_block_ptr( - h + i_t.to(tl.int64) * stride_h, - (V, K), - (K, 1), - (i_v * BV, 192), - (BV, 64), - (1, 0), - ) - tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + p_h4 = h + i_t.to(tl.int64) * stride_h + o_v[:, None] * K + o_k4[None, :] + m_h4 = m_v[:, None] & m_k4[None, :] + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), mask=m_h4) - p_w = tl.make_block_ptr( - w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0) - ) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * stride_w + o_k1[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k1[None, :], other=0.0) b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) if K > 64: - p_w = tl.make_block_ptr( - w, (T, K), (stride_w, 1), (i_t * BT, 64), (BT, 64), (1, 0) - ) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * stride_w + o_k2[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k2[None, :], other=0.0) b_v += tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype)) if K > 128: - p_w = tl.make_block_ptr( - w, (T, K), (stride_w, 1), (i_t * BT, 128), (BT, 64), (1, 0) - ) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * stride_w + o_k3[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k3[None, :], other=0.0) b_v += tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype)) if K > 192: - p_w = tl.make_block_ptr( - w, (T, K), (stride_w, 1), (i_t * BT, 192), (BT, 64), (1, 0) - ) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * stride_w + o_k4[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k4[None, :], other=0.0) b_v += tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype)) - p_v = tl.make_block_ptr( - v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v + p_v = v + o_t[:, None] * stride_v + o_v[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & m_v[None, :], other=0.0) - b_v if SAVE_NEW_VALUE: - p_v = tl.make_block_ptr( - v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + p_v = v_new + o_t[:, None] * stride_v + o_v[None, :] + tl.store( + p_v, b_v.to(p_v.dtype.element_ty), mask=m_t[:, None] & m_v[None, :] ) - tl.store(p_v, b_v.to(p_v.dtype.element_ty), boundary_check=(0, 1)) last_idx = min((i_t.to(tl.int64) + 1) * BT, T) - 1 if USE_G: - m_t = (i_t.to(tl.int64) * BT + tl.arange(0, BT)) < T b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - p_g = tl.make_block_ptr( - g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) - ) - b_g = tl.load(p_g, boundary_check=(0,)) + p_g = g + bos * H + i_h + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0) if USE_EXP2: b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] b_g_last = exp2(b_g_last) @@ -239,7 +212,6 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( else: b_h1 *= exp(b_gk_last1)[None, :] if K > 64: - o_k2 = 64 + o_k1 b_gk_last2 = tl.load( gk + (bos + last_idx) * H * K + i_h * K + o_k2, mask=(o_k2 < K), @@ -250,7 +222,6 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( else: b_h2 *= exp(b_gk_last2)[None, :] if K > 128: - o_k3 = 128 + o_k1 b_gk_last3 = tl.load( gk + (bos + last_idx) * H * K + i_h * K + o_k3, mask=(o_k3 < K), @@ -261,7 +232,6 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( else: b_h3 *= exp(b_gk_last3)[None, :] if K > 192: - o_k4 = 192 + o_k1 b_gk_last4 = tl.load( gk + (bos + last_idx) * H * K + i_h * K + o_k4, mask=(o_k4 < K), @@ -273,48 +243,38 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( b_h4 *= exp(b_gk_last4)[None, :] b_v = b_v.to(k.dtype.element_ty) - p_k = tl.make_block_ptr( - k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1) - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_k1[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k1[:, None] & m_t[None, :], other=0.0) b_h1 += tl.trans(tl.dot(b_k, b_v)) if K > 64: - p_k = tl.make_block_ptr( - k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1) - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_k2[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k2[:, None] & m_t[None, :], other=0.0) b_h2 += tl.trans(tl.dot(b_k, b_v)) if K > 128: - p_k = tl.make_block_ptr( - k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1) - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_k3[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k3[:, None] & m_t[None, :], other=0.0) b_h3 += tl.trans(tl.dot(b_k, b_v)) if K > 192: - p_k = tl.make_block_ptr( - k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1) - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_k4[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k4[:, None] & m_t[None, :], other=0.0) b_h4 += tl.trans(tl.dot(b_k, b_v)) # epilogue if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) - tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_v[:, None] * K + o_k1[None, :] + m_ht = m_v[:, None] & m_k1[None, :] + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), mask=m_ht) if K > 64: - p_ht = tl.make_block_ptr( - ht, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0) - ) - tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_v[:, None] * K + o_k2[None, :] + m_ht = m_v[:, None] & m_k2[None, :] + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), mask=m_ht) if K > 128: - p_ht = tl.make_block_ptr( - ht, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0) - ) - tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_v[:, None] * K + o_k3[None, :] + m_ht = m_v[:, None] & m_k3[None, :] + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), mask=m_ht) if K > 192: - p_ht = tl.make_block_ptr( - ht, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0) - ) - tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_v[:, None] * K + o_k4[None, :] + m_ht = m_v[:, None] & m_k4[None, :] + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), mask=m_ht) def chunk_gated_delta_rule_fwd_h( diff --git a/vllm/third_party/flash_linear_attention/ops/chunk_o.py b/vllm/third_party/flash_linear_attention/ops/chunk_o.py index 0c323b8ce215..219c8ed43f71 100644 --- a/vllm/third_party/flash_linear_attention/ops/chunk_o.py +++ b/vllm/third_party/flash_linear_attention/ops/chunk_o.py @@ -60,14 +60,18 @@ def chunk_fwd_kernel_o( USE_G: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t, i_bh = ( + tl.program_id(0), + tl.program_id(1).to(tl.int64), + tl.program_id(2).to(tl.int64), + ) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -90,22 +94,21 @@ def chunk_fwd_kernel_o( b_o = tl.zeros([BT, BV], dtype=tl.float32) b_A = tl.zeros([BT, BT], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + o_v = i_v * BV + tl.arange(0, BV) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) - ) - p_k = tl.make_block_ptr( - k, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1) - ) - p_h = tl.make_block_ptr( - h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0) - ) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + p_q = q + o_t[:, None] * (Hg * K) + o_k[None, :] + p_k = k + o_k[:, None] + o_t[None, :] * (Hg * K) + p_h = h + o_v[:, None] * K + o_k[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_t[:, None] & m_k[None, :], other=0.0) # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k[:, None] & m_t[None, :], other=0.0) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=(o_v[:, None] < V) & m_k[None, :], other=0.0) # [BT, BK] @ [BK, BV] -> [BT, BV] b_o += tl.dot(b_q, tl.trans(b_h)) @@ -114,28 +117,22 @@ def chunk_fwd_kernel_o( if USE_G: g += bos * H + i_h - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + p_g = g + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0) b_o = b_o * exp(b_g)[:, None] b_A = b_A * exp(b_g[:, None] - b_g[None, :]) - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) b_A = tl.where(m_A, b_A, 0) - p_v = tl.make_block_ptr( - v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - p_o = tl.make_block_ptr( - o, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) + p_v = v + o_t[:, None] * (H * V) + o_v[None, :] + p_o = o + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) # to fix mma -> mma layout conversion # already solved by triton v3.2 or higher b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_t[:, None] & (o_v < V)[None, :]) def chunk_fwd_o( diff --git a/vllm/third_party/flash_linear_attention/ops/chunk_scaled_dot_kkt.py b/vllm/third_party/flash_linear_attention/ops/chunk_scaled_dot_kkt.py index 10ff970c809f..f610fc7031be 100644 --- a/vllm/third_party/flash_linear_attention/ops/chunk_scaled_dot_kkt.py +++ b/vllm/third_party/flash_linear_attention/ops/chunk_scaled_dot_kkt.py @@ -61,12 +61,12 @@ def chunk_scaled_dot_kkt_fwd_kernel( USE_G: tl.constexpr, CAST_DOT_TO_K_DTYPE: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -78,22 +78,19 @@ def chunk_scaled_dot_kkt_fwd_kernel( o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T - p_beta = tl.make_block_ptr( - beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) - ) - b_beta = tl.load(p_beta, boundary_check=(0,)) + p_beta = beta + bos * H + i_h + o_t * H + b_beta = tl.load(p_beta, mask=m_t, other=0.0) b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr( - k + (bos * Hg + i_h // (H // Hg)) * K, - (T, K), - (Hg * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), + o_k = i_k * BK + tl.arange(0, BK) + p_k = ( + k + + (bos * Hg + i_h // (H // Hg)) * K + + o_t[:, None] * (Hg * K) + + o_k[None, :] ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_t[:, None] & (o_k < K)[None, :], other=0.0) b_kb = b_k * b_beta[:, None] if CAST_DOT_TO_K_DTYPE: # RDNA: force operands to k's native dtype so WMMA is used. @@ -103,17 +100,15 @@ def chunk_scaled_dot_kkt_fwd_kernel( b_A += tl.dot(b_kb, tl.trans(b_k).to(b_kb.dtype)) if USE_G: - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + p_g = g + bos * H + i_h + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0) b_g_diff = b_g[:, None] - b_g[None, :] b_A = b_A * exp(b_g_diff) m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) b_A = tl.where(m_A, b_A, 0) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (BT * H) + tl.arange(0, BT)[None, :] + tl.store(p_A, b_A.to(p_A.dtype.element_ty), mask=m_t[:, None]) def chunk_scaled_dot_kkt_fwd( diff --git a/vllm/third_party/flash_linear_attention/ops/cumsum.py b/vllm/third_party/flash_linear_attention/ops/cumsum.py index b0820104b1a1..1fcab916caf2 100644 --- a/vllm/third_party/flash_linear_attention/ops/cumsum.py +++ b/vllm/third_party/flash_linear_attention/ops/cumsum.py @@ -37,12 +37,12 @@ def chunk_local_cumsum_scalar_kernel( IS_VARLEN: tl.constexpr, HEAD_FIRST: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -52,23 +52,21 @@ def chunk_local_cumsum_scalar_kernel( else: bos, eos = i_b * T, i_b * T + T + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T if HEAD_FIRST: - p_s = tl.make_block_ptr( - s + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,) - ) - p_o = tl.make_block_ptr( - o + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,) - ) + p_s = s + bos * H + i_h * T + o_t + p_o = o + bos * H + i_h * T + o_t else: - p_s = tl.make_block_ptr(s + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_s = s + bos * H + i_h + o_t * H + p_o = o + bos * H + i_h + o_t * H # [BT] - b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) + b_s = tl.load(p_s, mask=m_t, other=0.0).to(tl.float32) b_o = tl.cumsum(b_s, axis=0) if REVERSE: b_z = tl.sum(b_s, axis=0) b_o = -b_o + b_z[None] + b_s - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_t) @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @@ -96,12 +94,12 @@ def chunk_local_cumsum_vector_kernel( IS_VARLEN: tl.constexpr, HEAD_FIRST: tl.constexpr, ): - i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -117,44 +115,19 @@ def chunk_local_cumsum_vector_kernel( else: m_s = tl.where(o_i[:, None] >= o_i[None, :], 1.0, 0.0) + o_t = i_t * BT + tl.arange(0, BT) + o_s = i_s * BS + tl.arange(0, BS) + m_ts = (o_t[:, None] < T) & (o_s[None, :] < S) if HEAD_FIRST: - p_s = tl.make_block_ptr( - s + (bos * H + i_h * T) * S, - (T, S), - (S, 1), - (i_t * BT, i_s * BS), - (BT, BS), - (1, 0), - ) - p_o = tl.make_block_ptr( - o + (bos * H + i_h * T) * S, - (T, S), - (S, 1), - (i_t * BT, i_s * BS), - (BT, BS), - (1, 0), - ) + p_s = s + (bos * H + i_h * T) * S + o_t[:, None] * S + o_s[None, :] + p_o = o + (bos * H + i_h * T) * S + o_t[:, None] * S + o_s[None, :] else: - p_s = tl.make_block_ptr( - s + (bos * H + i_h) * S, - (T, S), - (H * S, 1), - (i_t * BT, i_s * BS), - (BT, BS), - (1, 0), - ) - p_o = tl.make_block_ptr( - o + (bos * H + i_h) * S, - (T, S), - (H * S, 1), - (i_t * BT, i_s * BS), - (BT, BS), - (1, 0), - ) + p_s = s + (bos * H + i_h) * S + o_t[:, None] * (H * S) + o_s[None, :] + p_o = o + (bos * H + i_h) * S + o_t[:, None] * (H * S) + o_s[None, :] # [BT, BS] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + b_s = tl.load(p_s, mask=m_ts, other=0.0).to(tl.float32) b_o = tl.dot(m_s, b_s, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_ts) def chunk_local_cumsum_scalar( diff --git a/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py b/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py index 8f58ddcf8c7f..51ef43270713 100644 --- a/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py +++ b/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py @@ -50,31 +50,30 @@ def layer_norm_gated_fwd_kernel( HAS_BIAS: tl.constexpr, launch_pdl: tl.constexpr, ): - i_t = tl.program_id(0) + i_t = tl.program_id(0).to(tl.int64) + o_t = i_t * BT + tl.arange(0, BT) o_d = tl.arange(0, BD) m_d = o_d < D + m_t = o_t < T + m_x = m_t[:, None] & m_d[None, :] if launch_pdl: tl.extra.cuda.gdc_wait() tl.extra.cuda.gdc_launch_dependents() - p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + p_x = x + o_t[:, None] * D + o_d[None, :] + b_x = tl.load(p_x, mask=m_x, other=0.0).to(tl.float32) if HAS_RESIDUAL: - p_res = tl.make_block_ptr( - residual, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) - ) - b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32) + p_res = residual + o_t[:, None] * D + o_d[None, :] + b_x += tl.load(p_res, mask=m_x, other=0.0).to(tl.float32) if STORE_RESIDUAL_OUT: - p_res_out = tl.make_block_ptr( - residual_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) - ) - tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1)) + p_res_out = residual_out + o_t[:, None] * D + o_d[None, :] + tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), mask=m_x) if not IS_RMS_NORM: b_mean = tl.sum(b_x, axis=1) / D - p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,)) - tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,)) + p_mean = mean + o_t + tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), mask=m_t) b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0) b_var = tl.sum(b_xbar * b_xbar, axis=1) / D else: @@ -82,8 +81,8 @@ def layer_norm_gated_fwd_kernel( b_var = tl.sum(b_xbar * b_xbar, axis=1) / D b_rstd = 1 / tl.sqrt(b_var + eps) - p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) - tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) + p_rstd = rstd + o_t + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), mask=m_t) if HAS_WEIGHT: b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) @@ -99,7 +98,6 @@ def layer_norm_gated_fwd_kernel( b_y = b_y + b_b[None, :] # swish/sigmoid output gate - o_t = i_t * BT + tl.arange(0, BT) o_g = (o_t // H) * g_stride_n + (o_t % H) * D b_g = tl.load( g + o_g[:, None] + o_d[None, :], @@ -112,8 +110,8 @@ def layer_norm_gated_fwd_kernel( b_y = b_y * tl.sigmoid(b_g) # Write output - p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + p_y = y + o_t[:, None] * D + o_d[None, :] + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_x) @triton.heuristics( diff --git a/vllm/third_party/flash_linear_attention/ops/kda.py b/vllm/third_party/flash_linear_attention/ops/kda.py index b5c3a92b65e7..9327b24ddd0d 100644 --- a/vllm/third_party/flash_linear_attention/ops/kda.py +++ b/vllm/third_party/flash_linear_attention/ops/kda.py @@ -178,27 +178,26 @@ def layer_norm_gated_fwd_kernel( HAS_WEIGHT: tl.constexpr, HAS_BIAS: tl.constexpr, ): - i_t = tl.program_id(0) + i_t = tl.program_id(0).to(tl.int64) + o_t = i_t * BT + tl.arange(0, BT) o_d = tl.arange(0, BD) m_d = o_d < D + m_t = o_t < T + m_x = m_t[:, None] & m_d[None, :] - p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + p_x = x + o_t[:, None] * D + o_d[None, :] + b_x = tl.load(p_x, mask=m_x, other=0.0).to(tl.float32) if HAS_RESIDUAL: - p_res = tl.make_block_ptr( - residual, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) - ) - b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32) + p_res = residual + o_t[:, None] * D + o_d[None, :] + b_x += tl.load(p_res, mask=m_x, other=0.0).to(tl.float32) if STORE_RESIDUAL_OUT: - p_res_out = tl.make_block_ptr( - residual_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) - ) - tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1)) + p_res_out = residual_out + o_t[:, None] * D + o_d[None, :] + tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), mask=m_x) if not IS_RMS_NORM: b_mean = tl.sum(b_x, axis=1) / D - p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,)) - tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,)) + p_mean = mean + o_t + tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), mask=m_t) b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0) b_var = tl.sum(b_xbar * b_xbar, axis=1) / D else: @@ -206,8 +205,8 @@ def layer_norm_gated_fwd_kernel( b_var = tl.sum(b_xbar * b_xbar, axis=1) / D b_rstd = 1 / tl.sqrt(b_var + eps) - p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) - tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) + p_rstd = rstd + o_t + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), mask=m_t) if HAS_WEIGHT: b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) @@ -223,7 +222,6 @@ def layer_norm_gated_fwd_kernel( b_y = b_y + b_b[None, :] # swish/sigmoid output gate - o_t = i_t * BT + tl.arange(0, BT) o_g = (o_t // H) * g_stride_n + (o_t % H) * D b_g = tl.load( g + o_g[:, None] + o_d[None, :], @@ -236,8 +234,8 @@ def layer_norm_gated_fwd_kernel( b_y = b_y * tl.sigmoid(b_g) # Write output - p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + p_y = y + o_t[:, None] * D + o_d[None, :] + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_x) @triton.heuristics( @@ -576,13 +574,13 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_c, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H i_i, i_j = i_c // NC, i_c % NC if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -603,58 +601,51 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( A += (bos * H + i_h) * BT Aqk += (bos * H + i_h) * BT - p_b = tl.make_block_ptr( - beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,) - ) - b_b = tl.load(p_b, boundary_check=(0,)) + o_i = i_t * BT + i_i * BC + tl.arange(0, BC) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + m_i = o_i < T + m_j = o_j < T + + p_b = beta + bos * H + i_h + o_i * H + b_b = tl.load(p_b, mask=m_i, other=0.0) b_A = tl.zeros([BC, BC], dtype=tl.float32) b_Aqk = tl.zeros([BC, BC], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_k = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_g = tl.make_block_ptr( - g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - b_kt = tl.make_block_ptr( - k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - p_gk = tl.make_block_ptr( - g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K + m_qk = m_i[:, None] & m_k[None, :] + m_kj = m_k[:, None] & m_j[None, :] + p_q = q + o_i[:, None] * (H * K) + o_k[None, :] + p_k = k + o_i[:, None] * (H * K) + o_k[None, :] + p_g = g + o_i[:, None] * (H * K) + o_k[None, :] + b_kt = k + o_k[:, None] + o_j[None, :] * (H * K) + p_gk = g + o_k[:, None] + o_j[None, :] * (H * K) + # [BK,] b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0) # [BC, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) * exp2(b_g - b_gn[None, :]) + b_g = tl.load(p_g, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) * exp2(b_g - b_gn[None, :]) # [BK, BC] - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kt = tl.load(b_kt, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, mask=m_kj, other=0.0) + b_kt = tl.load(b_kt, mask=m_kj, other=0.0) # [BC, BC] b_ktg = b_kt * exp2(b_gn[:, None] - b_gk) b_A += tl.dot(b_k, b_ktg) - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale b_Aqk += tl.dot(b_qg, b_ktg) b_A *= b_b[:, None] - p_A = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) - tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) - p_Aqk = tl.make_block_ptr( - Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) - tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + o_jA = i_j * BC + tl.arange(0, BC) + m_A = m_i[:, None] & (o_jA[None, :] < BT) + p_A = A + o_i[:, None] * (H * BT) + o_jA[None, :] + tl.store(p_A, b_A.to(A.dtype.element_ty), mask=m_A) + p_Aqk = Aqk + o_i[:, None] * (H * BT) + o_jA[None, :] + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), mask=m_A) @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @@ -681,12 +672,12 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( BK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -705,33 +696,14 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( m_A = (i_t * BT + i_i * BC + o_i) < T o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) + o_c = i_t * BT + i_i * BC + tl.arange(0, BC) + m_qk = m_A[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) + b_g = tl.load(p_g, mask=m_qk, other=0.0) p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] @@ -878,12 +850,12 @@ def recompute_w_u_fwd_kernel( IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -892,108 +864,57 @@ def recompute_w_u_fwd_kernel( T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_b = beta + bos * H + i_h + o_t * H + b_b = tl.load(p_b, mask=m_t, other=0.0) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) - b_A = tl.load(p_A, boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_A, other=0.0) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_u = tl.make_block_ptr( - u + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_u = u + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - p_w = tl.make_block_ptr( - w + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] + p_w = w + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_kb = b_k * b_b[:, None] - p_gk = tl.make_block_ptr( - gk + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + p_gk = gk + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_gk = tl.load(p_gk, mask=m_tk, other=0.0) b_kb *= exp2(b_gk) if STORE_QG: - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_qg = tl.make_block_ptr( - qg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_qg = qg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_tk, other=0.0) b_qg = b_q * exp2(b_gk) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), mask=m_tk) if STORE_KG: last_idx = min(i_t * BT + BT, T) - 1 - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K b_gn = tl.load( gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 ) b_kg = b_k * exp2(b_gn - b_gk) - p_kg = tl.make_block_ptr( - kg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + p_kg = kg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), mask=m_tk) b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_tk) def recompute_w_u_fwd( @@ -1074,17 +995,17 @@ def chunk_gla_fwd_kernel_o( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_tg = i_t + i_tg = i_t.to(tl.int64) i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int32), - tl.load(cu_seqlens + i_n + 1).to(tl.int32), + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), ) T = eos - bos NT = tl.cdiv(T, BT) @@ -1096,70 +1017,44 @@ def chunk_gla_fwd_kernel_o( m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] b_o = tl.zeros([BT, BV], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + o_i = tl.arange(0, BT) + m_t = o_t < T + m_v = o_v < V + m_tv = m_t[:, None] & m_v[None, :] + m_A = m_t[:, None] & (o_i[None, :] < BT) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_h = tl.make_block_ptr( - h + (i_tg * H + i_h) * K * V, - (V, K), - (K, 1), - (i_v * BV, i_k * BK), - (BV, BK), - (1, 0), - ) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_qk = m_t[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_h = h + (i_tg * H + i_h) * K * V + o_v[:, None] * K + o_k[None, :] + m_h = m_v[:, None] & m_k[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BT, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, mask=m_qk, other=0.0) # [BT, BK] b_qg = (b_q * exp2(b_g)).to(b_q.dtype) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BT, BV] if i_k >= 0: b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_o = tl.make_block_ptr( - o + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_o = o + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_i[None, :] # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_A, other=0.0) b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) b_o += tl.dot(b_A, b_v, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_tv) def chunk_gla_fwd_o_gk( @@ -1238,12 +1133,12 @@ def kda_gate_cumsum_fwd_kernel( IS_VARLEN: tl.constexpr, USE_LOWER_BOUND: tl.constexpr, ): - i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -1253,27 +1148,17 @@ def kda_gate_cumsum_fwd_kernel( else: bos = i_b * T - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), - ) - p_y = tl.make_block_ptr( - y + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), - ) + o_t = i_t * BT + tl.arange(0, BT) + o_d = i_d * BD + tl.arange(0, BD) + m_t = o_t < T + m_d = o_d < D + m_g = m_t[:, None] & m_d[None, :] + p_g = g + (bos * H + i_h) * D + o_t[:, None] * (H * D) + o_d[None, :] + p_y = y + (bos * H + i_h) * D + o_t[:, None] * (H * D) + o_d[None, :] - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(p_g, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: - o_d = i_d * BD + tl.arange(0, BD) - b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) + b_bias = tl.load(g_bias + i_h * D + o_d, mask=m_d, other=0.0).to(tl.float32) b_g = b_g + b_bias[None, :] b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) @@ -1290,12 +1175,12 @@ def kda_gate_cumsum_fwd_kernel( # Out-of-bounds rows (load returns 0, but softplus/bias can still make # b_gate non-zero) participate in the dot product. They only contribute to - # out-of-bounds output rows, which are masked away by `boundary_check` on - # the store, so visible output matches unfused gate + chunk-local cumsum. - o_t = tl.arange(0, BT) - m_cumsum = tl.where(o_t[:, None] >= o_t[None, :], 1.0, 0.0) + # out-of-bounds output rows, which are masked away by the store mask, + # so visible output matches unfused gate + chunk-local cumsum. + o_c = tl.arange(0, BT) + m_cumsum = tl.where(o_c[:, None] >= o_c[None, :], 1.0, 0.0) b_y = tl.dot(m_cumsum, b_gate, allow_tf32=False) * cumsum_scale - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_g) def fused_kda_gate_chunk_cumsum( @@ -1605,7 +1490,7 @@ def kda_gate_fwd_kernel( HAS_BIAS: tl.constexpr, USE_LOWER_BOUND: tl.constexpr, ): - i_t, i_h = tl.program_id(0), tl.program_id(1) + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1) n_t = i_t * BT b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) @@ -1613,25 +1498,14 @@ def kda_gate_fwd_kernel( stride_row = H * D stride_col = 1 - g_ptr = tl.make_block_ptr( - base=g + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) - - y_ptr = tl.make_block_ptr( - base=y + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) + o_t = n_t + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + g_ptr = g + i_h * D + o_t[:, None] * stride_row + o_d[None, :] * stride_col + y_ptr = y + i_h * D + o_t[:, None] * stride_row + o_d[None, :] * stride_col - b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(g_ptr, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: n_d = tl.arange(0, BD) @@ -1649,7 +1523,7 @@ def kda_gate_fwd_kernel( sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) b_y = -b_a * sp - tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(y_ptr, b_y.to(y.dtype.element_ty), mask=m_g) def fused_kda_gate( diff --git a/vllm/third_party/flash_linear_attention/ops/l2norm.py b/vllm/third_party/flash_linear_attention/ops/l2norm.py index 2eb137a242fb..4827df2d07dd 100644 --- a/vllm/third_party/flash_linear_attention/ops/l2norm.py +++ b/vllm/third_party/flash_linear_attention/ops/l2norm.py @@ -66,13 +66,18 @@ def l2norm_fwd_kernel( BT: tl.constexpr, BD: tl.constexpr, ): - i_t = tl.program_id(0) - p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + i_t = tl.program_id(0).to(tl.int64) + o_t = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_x = m_t[:, None] & (o_d[None, :] < D) + p_x = x + o_t[:, None] * D + o_d[None, :] + p_y = y + o_t[:, None] * D + o_d[None, :] + + b_x = tl.load(p_x, mask=m_x, other=0.0).to(tl.float32) b_var = tl.sum(b_x * b_x, axis=1) b_y = b_x / tl.sqrt(b_var + eps)[:, None] - p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_x) @triton.jit diff --git a/vllm/third_party/flash_linear_attention/ops/solve_tril.py b/vllm/third_party/flash_linear_attention/ops/solve_tril.py index 8d3811ca4c17..93ef770be627 100644 --- a/vllm/third_party/flash_linear_attention/ops/solve_tril.py +++ b/vllm/third_party/flash_linear_attention/ops/solve_tril.py @@ -69,12 +69,12 @@ def solve_tril_16x16_kernel( Ai = Ai + (bos * H + i_h) * 16 offset = (i_t * 16) % BT + o_t = (i_t * 16 + o_i).to(tl.int64) + m_t = o_t < T if not USE_TMA: - p_A = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * 16, offset), (16, 16), (1, 0) - ) + p_A = A + o_t[:, None] * (H * BT) + (o_i + offset)[None, :] # [16, 16] - b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) + b_A = tl.load(p_A, mask=m_t[:, None], other=0.0).to(tl.float32) else: desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) desc_o = make_tensor_descriptor(Ai, [T, 16], [H * 16, 1], [16, 16]) @@ -88,13 +88,11 @@ def solve_tril_16x16_kernel( b_A = tl.where((o_i == i)[:, None], b_a, b_A) b_A += m_I if not USE_TMA: - p_Ai = tl.make_block_ptr( - Ai, (T, 16), (H * 16, 1), (i_t * 16, 0), (16, 16), (1, 0) - ) + p_Ai = Ai + o_t[:, None] * (H * 16) + o_i[None, :] tl.store( p_Ai, b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=m_t[:, None], ) else: desc_o.store([i_t * 16, 0], b_A.to(desc_o.dtype, fp_downcast_rounding="rtne")) @@ -143,15 +141,14 @@ def merge_16x16_to_32x32_inverse_kernel( A += (bos * H + i_h) * BT Ai += (bos * H + i_h) * BT + o_t = (i_t * BT + o_i).to(tl.int64) if not USE_TMA: - p_A_11 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) - ) - p_A_22 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) + p_A_11 = A + o_t[:, None] * (H * BT) + o_i[None, :] + p_A_22 = A + (o_t[:, None] + 16) * (H * BT) + (o_i[None, :] + 16) + b_Ai_11 = tl.load(p_A_11, mask=(o_t[:, None] < T), other=0.0).to(tl.float32) + b_Ai_22 = tl.load(p_A_22, mask=(o_t[:, None] + 16 < T), other=0.0).to( + tl.float32 ) - b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) - b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) else: desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16]) @@ -175,10 +172,8 @@ def merge_16x16_to_32x32_inverse_kernel( b_Ai_22 += m_I if not USE_TMA: - p_A_21 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) - ) - b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) + p_A_21 = A + (o_t[:, None] + 16) * (H * BT) + o_i[None, :] + b_A_21 = tl.load(p_A_21, mask=(o_t[:, None] + 16 < T), other=0.0).to(tl.float32) else: b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) @@ -189,29 +184,23 @@ def merge_16x16_to_32x32_inverse_kernel( ) if not USE_TMA: - p_Ai_11 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) - ) - p_Ai_21 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) - ) - p_Ai_22 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) - ) + p_Ai_11 = Ai + o_t[:, None] * (H * BT) + o_i[None, :] + p_Ai_21 = Ai + (o_t[:, None] + 16) * (H * BT) + o_i[None, :] + p_Ai_22 = Ai + (o_t[:, None] + 16) * (H * BT) + (o_i[None, :] + 16) tl.store( p_Ai_11, b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] < T), ) tl.store( p_Ai_22, b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 16 < T), ) tl.store( p_Ai_21, b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 16 < T), ) else: desc_o.store( @@ -268,23 +257,22 @@ def merge_16x16_to_64x64_inverse_kernel( A += (bos * H + i_h) * BT Ai += (bos * H + i_h) * BT + o_t = (i_t * BT + o_i).to(tl.int64) if not USE_TMA: - p_A_11 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) - ) - p_A_22 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) + p_A_11 = A + o_t[:, None] * (H * BT) + o_i[None, :] + p_A_22 = A + (o_t[:, None] + 16) * (H * BT) + (o_i[None, :] + 16) + p_A_33 = A + (o_t[:, None] + 32) * (H * BT) + (o_i[None, :] + 32) + p_A_44 = A + (o_t[:, None] + 48) * (H * BT) + (o_i[None, :] + 48) + b_Ai_11 = tl.load(p_A_11, mask=(o_t[:, None] < T), other=0.0).to(tl.float32) + b_Ai_22 = tl.load(p_A_22, mask=(o_t[:, None] + 16 < T), other=0.0).to( + tl.float32 ) - p_A_33 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0) + b_Ai_33 = tl.load(p_A_33, mask=(o_t[:, None] + 32 < T), other=0.0).to( + tl.float32 ) - p_A_44 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0) + b_Ai_44 = tl.load(p_A_44, mask=(o_t[:, None] + 48 < T), other=0.0).to( + tl.float32 ) - b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) - b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) - b_Ai_33 = tl.load(p_A_33, boundary_check=(0, 1)).to(tl.float32) - b_Ai_44 = tl.load(p_A_44, boundary_check=(0, 1)).to(tl.float32) else: desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16]) @@ -321,30 +309,18 @@ def merge_16x16_to_64x64_inverse_kernel( b_Ai_44 += m_I if not USE_TMA: - p_A_21 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) - ) - p_A_31 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0) - ) - p_A_32 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0) - ) - p_A_41 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0) - ) - p_A_42 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0) - ) - p_A_43 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0) - ) - b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) - b_A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32) - b_A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32) - b_A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32) - b_A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32) - b_A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32) + p_A_21 = A + (o_t[:, None] + 16) * (H * BT) + o_i[None, :] + p_A_31 = A + (o_t[:, None] + 32) * (H * BT) + o_i[None, :] + p_A_32 = A + (o_t[:, None] + 32) * (H * BT) + (o_i[None, :] + 16) + p_A_41 = A + (o_t[:, None] + 48) * (H * BT) + o_i[None, :] + p_A_42 = A + (o_t[:, None] + 48) * (H * BT) + (o_i[None, :] + 16) + p_A_43 = A + (o_t[:, None] + 48) * (H * BT) + (o_i[None, :] + 32) + b_A_21 = tl.load(p_A_21, mask=(o_t[:, None] + 16 < T), other=0.0).to(tl.float32) + b_A_31 = tl.load(p_A_31, mask=(o_t[:, None] + 32 < T), other=0.0).to(tl.float32) + b_A_32 = tl.load(p_A_32, mask=(o_t[:, None] + 32 < T), other=0.0).to(tl.float32) + b_A_41 = tl.load(p_A_41, mask=(o_t[:, None] + 48 < T), other=0.0).to(tl.float32) + b_A_42 = tl.load(p_A_42, mask=(o_t[:, None] + 48 < T), other=0.0).to(tl.float32) + b_A_43 = tl.load(p_A_43, mask=(o_t[:, None] + 48 < T), other=0.0).to(tl.float32) else: b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) b_A_31 = desc.load([i_t * BT + 32, 0]).to(tl.float32) @@ -390,85 +366,65 @@ def merge_16x16_to_64x64_inverse_kernel( ) if not USE_TMA: - p_Ai_11 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) - ) - p_Ai_22 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) - ) - p_Ai_33 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0) - ) - p_Ai_44 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0) - ) - p_Ai_21 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) - ) - p_Ai_31 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0) - ) - p_Ai_32 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0) - ) - p_Ai_41 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0) - ) - p_Ai_42 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0) - ) - p_Ai_43 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0) - ) + p_Ai_11 = Ai + o_t[:, None] * (H * BT) + o_i[None, :] + p_Ai_22 = Ai + (o_t[:, None] + 16) * (H * BT) + (o_i[None, :] + 16) + p_Ai_33 = Ai + (o_t[:, None] + 32) * (H * BT) + (o_i[None, :] + 32) + p_Ai_44 = Ai + (o_t[:, None] + 48) * (H * BT) + (o_i[None, :] + 48) + p_Ai_21 = Ai + (o_t[:, None] + 16) * (H * BT) + o_i[None, :] + p_Ai_31 = Ai + (o_t[:, None] + 32) * (H * BT) + o_i[None, :] + p_Ai_32 = Ai + (o_t[:, None] + 32) * (H * BT) + (o_i[None, :] + 16) + p_Ai_41 = Ai + (o_t[:, None] + 48) * (H * BT) + o_i[None, :] + p_Ai_42 = Ai + (o_t[:, None] + 48) * (H * BT) + (o_i[None, :] + 16) + p_Ai_43 = Ai + (o_t[:, None] + 48) * (H * BT) + (o_i[None, :] + 32) tl.store( p_Ai_11, b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] < T), ) tl.store( p_Ai_22, b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 16 < T), ) tl.store( p_Ai_33, b_Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 32 < T), ) tl.store( p_Ai_44, b_Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 48 < T), ) tl.store( p_Ai_21, b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 16 < T), ) tl.store( p_Ai_31, b_Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 32 < T), ) tl.store( p_Ai_32, b_Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 32 < T), ) tl.store( p_Ai_41, b_Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 48 < T), ) tl.store( p_Ai_42, b_Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 48 < T), ) tl.store( p_Ai_43, b_Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 48 < T), ) else: desc_o.store( diff --git a/vllm/third_party/flash_linear_attention/ops/wy_fast.py b/vllm/third_party/flash_linear_attention/ops/wy_fast.py index 52d2b28195a8..040c37ebfe7f 100644 --- a/vllm/third_party/flash_linear_attention/ops/wy_fast.py +++ b/vllm/third_party/flash_linear_attention/ops/wy_fast.py @@ -46,12 +46,12 @@ def recompute_w_u_fwd_kernel( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -60,60 +60,41 @@ def recompute_w_u_fwd_kernel( T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_beta = tl.make_block_ptr( - beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) - ) - p_g = tl.make_block_ptr(g + (bos * H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) - b_beta = tl.load(p_beta, boundary_check=(0,)) - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_g = tl.exp(tl.load(p_g, boundary_check=(0,))) + o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) + m_t = o_t < T + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_beta = beta + bos * H + i_h + o_t * H + p_g = g + (bos * H + i_h) + o_t * H + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_A[None, :] + b_beta = tl.load(p_beta, mask=m_t, other=0.0) + b_A = tl.load(p_A, mask=m_A, other=0.0) + b_g = tl.exp(tl.load(p_g, mask=m_t, other=0.0)) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_u = tl.make_block_ptr( - u + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_u = u + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_beta[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb, allow_tf32=False) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr( - k + (bos * Hg + i_h // (H // Hg)) * K, - (T, K), - (Hg * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_w = tl.make_block_ptr( - w + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = ( + k + + (bos * Hg + i_h // (H // Hg)) * K + + o_t[:, None] * (Hg * K) + + o_k[None, :] ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_w = w + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_k, other=0.0) b_kb = (b_k * b_beta[:, None] * b_g[:, None]).to(b_k.dtype) b_w = tl.dot(b_A, b_kb) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_k) def recompute_w_u_fwd( From 50b21a2667b6903d2e0d651aab9f449832e5e1e9 Mon Sep 17 00:00:00 2001 From: TJian Date: Mon, 12 Oct 2026 23:50:28 +0800 Subject: [PATCH 313/624] [Perf] [Feat] [ROCm] Add densemha support to ROCm AITER Sparse MLA (#49263) Signed-off-by: tjtanaa --- .buildkite/test-amd.yaml | 10 +- .buildkite/test_areas/attention.yaml | 9 +- ...est_rocm_aiter_mla_sparse_metadata_sync.py | 503 ++++++++++++++++++ .../v1/attention/test_rocm_glm5next_sparse.py | 80 ++- .../v1/attention/test_sparse_mla_backends.py | 275 +++++++++- vllm/models/deepseek_v32/amd/rocm.py | 35 +- .../backends/mla/rocm_aiter_mla_sparse.py | 214 ++++---- 7 files changed, 984 insertions(+), 142 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index a14a0a33a4d2..b4c4b59fdfe8 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2632,8 +2632,11 @@ steps: - vllm/_aiter_ops.py - vllm/envs.py - vllm/platforms/rocm.py + - vllm/models/deepseek_v32/amd/rocm.py commands: - - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s v1/attention --ignore=v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + + - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":amd: (MI300) Metrics, Tracing" # TBD timeout_in_minutes: 30 @@ -5642,10 +5645,11 @@ steps: - vllm/_aiter_ops.py - vllm/envs.py - vllm/platforms/rocm.py + - vllm/models/deepseek_v32/amd/rocm.py optional: true commands: - - pytest -v -s v1/attention --ignore=v1/attention/test_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s v1/attention --ignore=v1/attention/test_mla_backends.py --ignore=v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_mla_backends.py v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":amd: (MI355 DPX) CUDAGraph" timeout_in_minutes: 55 diff --git a/.buildkite/test_areas/attention.yaml b/.buildkite/test_areas/attention.yaml index 83ba90e2cb1b..ec47e45fe7a0 100644 --- a/.buildkite/test_areas/attention.yaml +++ b/.buildkite/test_areas/attention.yaml @@ -30,6 +30,10 @@ steps: - vllm/_aiter_ops.py - vllm/envs.py - vllm/platforms/rocm.py + - vllm/models/deepseek_v32/amd/rocm.py + commands: + - pytest -v -s v1/attention --ignore=v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":nvidia: (B200) V1 Attention Shard %N" key: v1-attention-b200 @@ -59,9 +63,10 @@ steps: - vllm/_aiter_ops.py - vllm/envs.py - vllm/platforms/rocm.py + - vllm/models/deepseek_v32/amd/rocm.py commands: - - pytest -v -s v1/attention --ignore=v1/attention/test_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s v1/attention --ignore=v1/attention/test_mla_backends.py --ignore=v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_mla_backends.py v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":nvidia: (DGX) Spark FlashInfer NVFP4 XQA" key: flashinfer-nvfp4-xqa-dgx-spark diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py index aabce97fac9e..dc1ff79eeabd 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py @@ -33,6 +33,101 @@ class _FakeAiter(ModuleType): get_mla_metadata_v1: Mock +@pytest.mark.parametrize("block_size", [16, 64, 128]) +@pytest.mark.parametrize("topk_tokens", [128, 256]) +def test_rocm_index_conversion_matrix(block_size, topk_tokens): + """Convert valid, masked, and out-of-range sparse indices exactly.""" + device = torch.device("cuda") + num_tokens = 4 + max_blocks = 4 + req_ids = torch.tensor([0, 1, 0, 1], dtype=torch.int32, device=device) + block_table = torch.tensor( + [[3, 5, 7, 9], [2, 4, 6, 8]], dtype=torch.int32, device=device + ) + token_indices = torch.arange( + num_tokens * topk_tokens, dtype=torch.int32, device=device + ).view(num_tokens, topk_tokens) + token_indices %= max_blocks * block_size + token_indices[0, 0] = -1 + token_indices[1, 1] = max_blocks * block_size + cu_seqlens = torch.arange( + 0, + (num_tokens + 1) * topk_tokens, + topk_tokens, + dtype=torch.int32, + device=device, + ) + output = torch.full( + (num_tokens * topk_tokens,), -99, dtype=torch.int32, device=device + ) + + sparse_mod.triton_convert_req_index_to_global_index( + req_ids, + block_table, + token_indices, + cu_seqlens, + output, + BLOCK_STRIDE_ROWS=block_size, + BLOCK_SIZE=block_size, + NUM_TOPK_TOKENS=topk_tokens, + ) + + block_ids = token_indices // block_size + valid = (token_indices >= 0) & (block_ids < max_blocks) + safe_blocks = block_ids.clamp(0, max_blocks - 1) + bases = block_table[req_ids[:, None].long(), safe_blocks.long()] + expected = torch.where(valid, bases * block_size + token_indices % block_size, 0) + torch.testing.assert_close(output.view_as(expected), expected, rtol=0, atol=0) + + +def test_rocm_index_conversion_rejects_mismatched_token_extents(): + """Reject full-batch request IDs paired with decode-subset indices.""" + with pytest.raises(AssertionError, match="must cover the same tokens"): + sparse_mod.triton_convert_req_index_to_global_index( + torch.zeros(3, dtype=torch.int32), + torch.zeros(1, 4, dtype=torch.int32), + torch.zeros(2, 128, dtype=torch.int32), + torch.arange(3, dtype=torch.int32), + torch.zeros(256, dtype=torch.int32), + BLOCK_STRIDE_ROWS=64, + NUM_TOPK_TOKENS=128, + ) + + +def test_rocm_index_conversion_respects_packed_valid_counts(): + """ROCm's indptr contract packs only each token's valid top-k prefix.""" + device = torch.device("cuda") + topk_tokens = 128 + valid_counts = torch.tensor([1, 17, 128], dtype=torch.int32, device=device) + cu_seqlens = torch.zeros(4, dtype=torch.int32, device=device) + torch.cumsum(valid_counts, dim=0, out=cu_seqlens[1:]) + req_ids = torch.tensor([0, 0, 0], dtype=torch.int32, device=device) + block_table = torch.tensor([[2, 3, 4]], dtype=torch.int32, device=device) + token_indices = torch.arange( + 3 * topk_tokens, dtype=torch.int32, device=device + ).view(3, topk_tokens) + token_indices %= 3 * 64 + output = torch.full((int(cu_seqlens[-1]),), -99, dtype=torch.int32, device=device) + + sparse_mod.triton_convert_req_index_to_global_index( + req_ids, + block_table, + token_indices, + cu_seqlens, + output, + BLOCK_STRIDE_ROWS=64, + BLOCK_SIZE=64, + NUM_TOPK_TOKENS=topk_tokens, + ) + + expected_parts = [] + for token_id, count in enumerate(valid_counts.tolist()): + local = token_indices[token_id, :count] + expected_parts.append(block_table[0, (local // 64).long()] * 64 + local % 64) + expected = torch.cat(expected_parts) + torch.testing.assert_close(output, expected, rtol=0, atol=0) + + def _make_builder(): builder = object.__new__(sparse_mod.ROCMAiterMLASparseMetadataBuilder) max_num_batched_tokens = 8 @@ -44,6 +139,18 @@ def _make_builder(): builder.kv_cache_dtype = "fp8" builder.mla_dims = SimpleNamespace(kv_lora_rank=512, qk_rope_head_dim=64) builder.topk_tokens = topk_tokens + # SparseMLACommonMetadataBuilder state consumed by its build() method. + builder.cp_kv_cache_interleave_size = 1 + builder.use_pcp = False + builder._prefill_backend = None + + def cpu_token_to_req(metadata): + return torch.repeat_interleave( + torch.arange(metadata.num_reqs, dtype=torch.int32), + metadata.query_start_loc_cpu.diff(), + ) + + builder._build_req_id_per_token = cpu_token_to_req # type: ignore[method-assign] builder.req_id_per_token_buffer = torch.zeros( max_num_batched_tokens, dtype=torch.int32, device="cpu" ) @@ -62,6 +169,8 @@ def _make_builder(): builder._use_persistent_metadata = True builder._num_attention_heads = 16 builder._num_compute_units = current_platform.num_compute_units() + builder._mla_q_dtype = torch.bfloat16 + builder._mla_kv_dtype = torch.bfloat16 builder._mla_work_meta_data = torch.empty(1, dtype=torch.int32, device="cpu") builder._mla_work_indptr = torch.empty(1, dtype=torch.int32, device="cpu") builder._mla_work_info_set = torch.empty(1, dtype=torch.int32, device="cpu") @@ -216,6 +325,8 @@ def fake_get_mla_metadata_v1(*args, **kwargs): assert events == ["metadata", "sync"] assert fake_get_mla_metadata_v1_mock.call_count == 1 assert fake_get_mla_metadata_v1_mock.call_args.kwargs["max_split_per_batch"] == 1 + assert fake_get_mla_metadata_v1_mock.call_args.kwargs["dtype_q"] == torch.bfloat16 + assert fake_get_mla_metadata_v1_mock.call_args.kwargs["dtype_kv"] == torch.bfloat16 events.clear() @@ -223,3 +334,395 @@ def fake_get_mla_metadata_v1(*args, **kwargs): assert events == [] assert fake_get_mla_metadata_v1_mock.call_count == 1 + + +def test_dense_prefill_builds_sparse_metadata_for_decode_subset(monkeypatch): + """Dense prefill must exclude prefill tokens from AITER sparse metadata. + + A mixed batch is routed as sparse MQA for its leading decode requests and + dense MHA for its trailing short prefill requests. The AITER metadata + buffers therefore must describe only the decode subset; including prefill + tokens would mismatch the query passed to ``forward_mqa``. + """ + builder = _make_builder() + builder.vllm_config = SimpleNamespace( + attention_config=SimpleNamespace(sparse_mla_force_mqa=False) + ) + + query_start_loc = torch.tensor([0, 1, 5], dtype=torch.int32) + seq_lens = torch.tensor([16, 4], dtype=torch.int32) + common_metadata = SimpleNamespace( + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens, + max_query_len=4, + max_seq_len=16, + ) + metadata = sparse_mod.ROCMAiterMLASparseMetadata( + num_reqs=2, + max_query_len=4, + max_seq_len=16, + num_actual_tokens=5, + query_start_loc=query_start_loc, + slot_mapping=torch.arange(5), + block_table=torch.arange(16, dtype=torch.int32).view(2, 8), + req_id_per_token=torch.tensor([0, 1, 1, 1, 1], dtype=torch.int32), + topk_tokens=4, + num_decodes=1, + num_prefills=1, + num_decode_tokens=1, + seq_lens=seq_lens, + prefill_max_seq_len=4, + prefill=SimpleNamespace(use_dense_mha=True), + ) + + monkeypatch.setattr( + sparse_mod.SparseMLACommonMetadataBuilder, + "build", + lambda *args, **kwargs: metadata, + ) + sparse_call: dict[str, torch.Tensor] = {} + + def fake_generate_sparse_seqlen( + query_lens, sparse_seq_lens, sparse_query_start_loc, topk, num_tokens, max_q + ): + sparse_call.update( + query_lens=query_lens.clone(), + seq_lens=sparse_seq_lens.clone(), + query_start_loc=sparse_query_start_loc.clone(), + num_tokens=num_tokens, + ) + return torch.ones(num_tokens, dtype=torch.int32) + + monkeypatch.setattr( + sparse_mod, "generate_sparse_seqlen_triton", fake_generate_sparse_seqlen + ) + fake_aiter = _FakeAiter("aiter") + fake_aiter.get_mla_metadata_v1 = Mock() + monkeypatch.setitem(sys.modules, "aiter", fake_aiter) + monkeypatch.setattr( + sparse_mod.torch.cuda, + "current_stream", + lambda device=None: SimpleNamespace(synchronize=lambda: None), + ) + + result = builder.build(0, common_metadata) + + assert result.prefill is metadata.prefill + assert result.num_prefills == 1 + assert result.req_id_per_token.tolist() == [0] + assert result.qo_indptr.numel() == 2 + assert result.paged_kv_indptr.numel() == 2 + assert sparse_call["num_tokens"] == 1 + assert sparse_call["query_lens"].tolist() == [1] + assert sparse_call["seq_lens"].tolist() == [16] + assert sparse_call["query_start_loc"].tolist() == [0, 1] + + +def test_long_prefill_shares_full_sparse_metadata_with_decode(monkeypatch): + """A prefill longer than top-k must remain in shared sparse metadata. + + Dense MHA is only selected when the prefill sequence fits within the + sparse top-k window. Once it exceeds that window, both decode and prefill + are handled by sparse MQA, so every query token must be represented by the + same sparse paging metadata. + """ + builder = _make_builder() + builder.vllm_config = SimpleNamespace( + attention_config=SimpleNamespace(sparse_mla_force_mqa=False) + ) + + # One decode token followed by a five-token prefill. topk_tokens is four, + # so the prefill must not be split onto the dense-MHA path. + query_start_loc = torch.tensor([0, 1, 6], dtype=torch.int32) + seq_lens = torch.tensor([16, 5], dtype=torch.int32) + common_metadata = SimpleNamespace( + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens, + max_query_len=5, + max_seq_len=16, + ) + metadata = sparse_mod.ROCMAiterMLASparseMetadata( + num_reqs=2, + max_query_len=5, + max_seq_len=16, + num_actual_tokens=6, + query_start_loc=query_start_loc, + slot_mapping=torch.arange(6), + block_table=torch.arange(16, dtype=torch.int32).view(2, 8), + req_id_per_token=torch.tensor([0, 1, 1, 1, 1, 1], dtype=torch.int32), + topk_tokens=4, + num_decodes=1, + num_prefills=1, + num_decode_tokens=1, + seq_lens=seq_lens, + prefill_max_seq_len=5, + prefill=SimpleNamespace(use_dense_mha=False), + ) + + monkeypatch.setattr( + sparse_mod.SparseMLACommonMetadataBuilder, + "build", + lambda *args, **kwargs: metadata, + ) + sparse_call: dict[str, torch.Tensor] = {} + + def fake_generate_sparse_seqlen( + query_lens, sparse_seq_lens, sparse_query_start_loc, topk, num_tokens, max_q + ): + sparse_call.update( + query_lens=query_lens.clone(), + seq_lens=sparse_seq_lens.clone(), + query_start_loc=sparse_query_start_loc.clone(), + num_tokens=num_tokens, + ) + return torch.ones(num_tokens, dtype=torch.int32) + + monkeypatch.setattr( + sparse_mod, "generate_sparse_seqlen_triton", fake_generate_sparse_seqlen + ) + fake_aiter = _FakeAiter("aiter") + get_persistent_metadata = Mock() + fake_aiter.get_mla_metadata_v1 = get_persistent_metadata + monkeypatch.setitem(sys.modules, "aiter", fake_aiter) + monkeypatch.setattr( + sparse_mod.torch.cuda, + "current_stream", + lambda device=None: SimpleNamespace(synchronize=lambda: None), + ) + + result = builder.build(0, common_metadata) + + assert result.num_decodes == 1 + assert result.num_prefills == 1 + assert result.req_id_per_token.tolist() == [0, 1, 1, 1, 1, 1] + assert result.qo_indptr.numel() == 7 + assert result.paged_kv_indptr.numel() == 7 + assert result.paged_kv_indices.numel() == 6 * result.topk_tokens + assert sparse_call["num_tokens"] == 6 + assert sparse_call["query_lens"].tolist() == [1, 5] + assert sparse_call["seq_lens"].tolist() == [16, 5] + assert sparse_call["query_start_loc"].tolist() == [0, 1, 6] + get_persistent_metadata.assert_called_once() + assert result.work_meta_data is builder._mla_work_meta_data + assert result.reduce_indptr is builder._mla_reduce_indptr + + +def test_forward_mqa_slices_full_batch_metadata_to_query_tokens(monkeypatch): + """The sparse kernel inputs must have the same token extent as its query. + + In a mixed dense-prefill batch, ``forward_mqa`` receives only decode query + tokens while some metadata buffers may represent the full batch. This + guards the slicing that prevents the index conversion from writing beyond + the decode-sized output. + """ + impl = object.__new__(sparse_mod.ROCMAiterMLASparseImpl) + impl.kv_cache_dtype = "auto" + impl.num_heads = 4 + impl.topk_indices_buffer = torch.arange(48, dtype=torch.int32).view(12, 4) + expected_output = torch.zeros(2, 4, 512, dtype=torch.bfloat16) + forward_mla = Mock(return_value=(expected_output, None)) + monkeypatch.setattr(impl, "_forward_mla", forward_mla) + + metadata = SimpleNamespace( + req_id_per_token=torch.tensor([0, 1] + [2] * 5 + [3] * 5, dtype=torch.int32), + block_table=torch.arange(40, dtype=torch.int32).view(4, 10), + paged_kv_indptr=torch.arange(3, dtype=torch.int32), + paged_kv_indices=torch.zeros(8, dtype=torch.int32), + block_size=1, + topk_tokens=4, + ) + query = torch.zeros(2, 4, 576, dtype=torch.bfloat16) + captured = {} + + def fake_convert(req_ids, block_table, token_indices, *args, **kwargs): + captured["req_ids"] = req_ids.clone() + captured["token_indices"] = token_indices.clone() + + monkeypatch.setattr( + sparse_mod, "triton_convert_req_index_to_global_index", fake_convert + ) + monkeypatch.setattr( + sparse_mod.AiterMLAHelper, "get_mla_padded_q", lambda num_heads, q: q + ) + + output, lse = impl.forward_mqa( + query, + torch.zeros(1, 1, 576, dtype=torch.bfloat16), + metadata, + SimpleNamespace(), + ) + + assert output is expected_output + assert lse is None + assert captured["req_ids"].tolist() == [0, 1] + assert captured["token_indices"].shape == (2, 4) + forward_mla.assert_called_once() + + +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +def test_forward_mqa_slices_decode_subset_for_bf16_and_fp8(monkeypatch, kv_cache_dtype): + """Both cache dtypes must slice mixed-batch metadata to decode queries.""" + impl = object.__new__(sparse_mod.ROCMAiterMLASparseImpl) + impl.kv_cache_dtype = kv_cache_dtype + impl.num_heads = 4 + impl.topk_indices_buffer = torch.arange(48, dtype=torch.int32).view(12, 4) + expected_output = torch.zeros(2, 4, 512, dtype=torch.bfloat16) + forward_mla = Mock(return_value=(expected_output, None)) + monkeypatch.setattr(impl, "_forward_mla", forward_mla) + + metadata = SimpleNamespace( + req_id_per_token=torch.tensor([0, 1] + [2] * 5 + [3] * 5, dtype=torch.int32), + block_table=torch.arange(40, dtype=torch.int32).view(4, 10), + paged_kv_indptr=torch.arange(3, dtype=torch.int32), + paged_kv_indices=torch.zeros(8, dtype=torch.int32), + block_size=1, + topk_tokens=4, + ) + query = torch.zeros(2, 4, 576, dtype=torch.bfloat16) + kv_cache = torch.zeros(1, 1, 576, dtype=torch.uint8) + captured = {} + + def fake_convert(req_ids, block_table, token_indices, *args, **kwargs): + captured["req_ids"] = req_ids.clone() + captured["token_indices"] = token_indices.clone() + + monkeypatch.setattr( + sparse_mod, "triton_convert_req_index_to_global_index", fake_convert + ) + monkeypatch.setattr( + sparse_mod.AiterMLAHelper, "get_mla_padded_q", lambda num_heads, q: q + ) + if kv_cache_dtype == "fp8": + fp8_dtype = current_platform.fp8_dtype() + monkeypatch.setattr( + sparse_mod.ops, + "scaled_fp8_quant", + lambda q, scale: (q.to(fp8_dtype), None), + ) + + output, lse = impl.forward_mqa( + query, + kv_cache, + metadata, + SimpleNamespace(_q_scale=torch.tensor(1.0)), + ) + + assert output is expected_output + assert lse is None + assert captured["req_ids"].tolist() == [0, 1] + assert captured["token_indices"].shape == (2, 4) + + +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.parametrize("length_relation", ["short", "long"]) +@pytest.mark.parametrize("workload", ["prefill", "mixed", "decode"]) +def test_sparse_metadata_workload_length_dtype_matrix( + monkeypatch, workload, length_relation, kv_cache_dtype +): + """Cover routing and persistent metadata across the ROCm support matrix. + + The matrix crosses pure prefill, mixed prefill/decode, and decode-only + workloads with sequence lengths below/above top-k and BF16/FP8 KV cache. + """ + builder = _make_builder() + builder.vllm_config = SimpleNamespace( + attention_config=SimpleNamespace(sparse_mla_force_mqa=False) + ) + if kv_cache_dtype == "fp8": + builder._mla_q_dtype = current_platform.fp8_dtype() + builder._mla_kv_dtype = current_platform.fp8_dtype() + + is_long = length_relation == "long" + prefill_len = 5 if is_long else 3 + decode_seq_len = 8 if is_long else 3 + if workload == "prefill": + query_lens = [prefill_len] + seq_lens_list = [prefill_len] + num_decodes, num_prefills = 0, 1 + elif workload == "mixed": + query_lens = [1, prefill_len] + seq_lens_list = [decode_seq_len, prefill_len] + num_decodes, num_prefills = 1, 1 + else: + query_lens = [1, 1] + seq_lens_list = [decode_seq_len, decode_seq_len + 1] + num_decodes, num_prefills = 2, 0 + + query_start_loc = torch.tensor( + [0, *torch.tensor(query_lens).cumsum(0).tolist()], dtype=torch.int32 + ) + seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32) + num_tokens = sum(query_lens) + common_metadata = SimpleNamespace( + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens, + max_query_len=max(query_lens), + max_seq_len=max(seq_lens_list), + ) + metadata = sparse_mod.ROCMAiterMLASparseMetadata( + num_reqs=len(query_lens), + max_query_len=max(query_lens), + max_seq_len=max(seq_lens_list), + num_actual_tokens=num_tokens, + query_start_loc=query_start_loc, + slot_mapping=torch.arange(num_tokens), + block_table=torch.arange(len(query_lens) * 8, dtype=torch.int32).view( + len(query_lens), 8 + ), + req_id_per_token=torch.repeat_interleave( + torch.arange(len(query_lens), dtype=torch.int32), + torch.tensor(query_lens), + ), + topk_tokens=4, + num_decodes=num_decodes, + num_prefills=num_prefills, + num_decode_tokens=num_decodes, + seq_lens=seq_lens, + prefill_max_seq_len=prefill_len if num_prefills else 0, + prefill=SimpleNamespace(use_dense_mha=not is_long) if num_prefills else None, + ) + monkeypatch.setattr( + sparse_mod.SparseMLACommonMetadataBuilder, + "build", + lambda *args, **kwargs: metadata, + ) + monkeypatch.setattr( + sparse_mod, + "generate_sparse_seqlen_triton", + lambda *args, **kwargs: torch.ones(args[4], dtype=torch.int32), + ) + get_persistent_metadata = Mock() + fake_aiter = _FakeAiter("aiter") + fake_aiter.get_mla_metadata_v1 = get_persistent_metadata + monkeypatch.setitem(sys.modules, "aiter", fake_aiter) + monkeypatch.setattr( + sparse_mod.torch.cuda, + "current_stream", + lambda device=None: SimpleNamespace(synchronize=lambda: None), + ) + + result = builder.build(0, common_metadata) + + uses_dense_prefill = workload in ("prefill", "mixed") and not is_long + expected_sparse_tokens = num_decodes if uses_dense_prefill else num_tokens + assert result.req_id_per_token.numel() == expected_sparse_tokens + assert result.qo_indptr.numel() == expected_sparse_tokens + 1 + assert result.paged_kv_indptr.numel() == expected_sparse_tokens + 1 + + uses_fallback = workload == "prefill" and is_long and kv_cache_dtype == "auto" + if expected_sparse_tokens == 0 or uses_fallback: + get_persistent_metadata.assert_not_called() + else: + get_persistent_metadata.assert_called_once() + if uses_fallback or expected_sparse_tokens == 0: + assert result.work_meta_data is None + assert result.reduce_indptr is None + else: + assert result.work_meta_data is builder._mla_work_meta_data + assert result.reduce_indptr is builder._mla_reduce_indptr diff --git a/tests/v1/attention/test_rocm_glm5next_sparse.py b/tests/v1/attention/test_rocm_glm5next_sparse.py index 6077b9eb2fb6..e4244bc06aff 100644 --- a/tests/v1/attention/test_rocm_glm5next_sparse.py +++ b/tests/v1/attention/test_rocm_glm5next_sparse.py @@ -155,8 +155,86 @@ def test_rocm_sparse_triton_route( ) +@pytest.mark.parametrize( + ("gfx942", "fp8", "num_prefills", "num_decodes", "prefill_max_seq_len", "route"), + [ + pytest.param(True, False, 1, 0, 129, "triton", id="above-topk"), + pytest.param(False, False, 1, 0, 129, "aiter", id="gfx950"), + pytest.param(True, True, 1, 0, 129, "aiter", id="fp8"), + pytest.param(True, False, 0, 2, 129, "aiter", id="decode-only"), + pytest.param(True, False, 1, 1, 129, "aiter", id="mixed"), + pytest.param(True, False, 1, 0, 128, "aiter", id="equal-topk"), + pytest.param(True, False, 1, 0, 127, "aiter", id="below-topk"), + pytest.param(True, False, 0, 0, 129, "aiter", id="empty"), + pytest.param( + True, False, 1, 0, None, "missing", id="required-prefill-metadata" + ), + ], +) +def test_rocm_long_bf16_prefill_dispatch( + monkeypatch, gfx942, fp8, num_prefills, num_decodes, prefill_max_seq_len, route +): + """Check the new dispatch gate using CPU tensors and intercepted kernels.""" + monkeypatch.setattr("vllm.platforms.rocm.on_gfx942", lambda: gfx942) + + class KernelSelected(Exception): + pass + + def triton_kernel(**kwargs): + raise KernelSelected("triton") + + def aiter_kernel(*args, **kwargs): + raise KernelSelected("aiter") + + monkeypatch.setattr(sparse_mod, "rocm_sparse_attn_prefill", triton_kernel) + monkeypatch.setattr(sparse_mod.rocm_aiter_ops, "mla_decode_fwd", aiter_kernel) + impl = object.__new__(sparse_mod.ROCMAiterMLASparseImpl) + impl.num_heads = 64 + impl.kv_lora_rank = 512 + impl.kv_cache_dtype = "fp8" if fp8 else "auto" + impl.scale = 576**-0.5 + impl.sinks = None + + dtype = torch.float8_e4m3fn if fp8 else torch.bfloat16 + q = torch.empty(2, 64, 576, dtype=dtype) + kv = torch.empty(4, 1, 576, dtype=dtype) + metadata = SimpleNamespace( + attn_out_dtype=torch.bfloat16, + num_prefills=num_prefills, + num_decodes=num_decodes, + num_decode_tokens=num_decodes, + max_query_len=2, + prefill_max_seq_len=prefill_max_seq_len, + topk_tokens=128, + work_meta_data=None, + qo_indptr=torch.arange(3, dtype=torch.int32), + paged_kv_indices=torch.empty(0, dtype=torch.int32), + paged_kv_indptr=torch.zeros(3, dtype=torch.int32), + paged_kv_last_page_len=torch.ones(2, dtype=torch.int32), + ) + # Irrelevant routes must not inspect prefill-only metadata. + if ( + not gfx942 + or fp8 + or num_decodes + or not num_prefills + or prefill_max_seq_len is None + ): + del metadata.prefill_max_seq_len + del metadata.topk_tokens + layer = SimpleNamespace(_q_scale=1.0, _k_scale=1.0) + error = AttributeError if route == "missing" else KernelSelected + match = "prefill_max_seq_len" if route == "missing" else f"^{route}$" + with pytest.raises(error, match=match): + impl._forward_mla(layer, q, kv, metadata) + + +@pytest.mark.parametrize("gfx942", [False, True]) @pytest.mark.parametrize("num_heads", [8, 12]) -def test_rocm_sparse_triton_route_preserves_padded_sinks(monkeypatch, num_heads): +def test_rocm_sparse_triton_route_preserves_padded_sinks( + monkeypatch, num_heads, gfx942 +): + monkeypatch.setattr("vllm.platforms.rocm.on_gfx942", lambda: gfx942) captured = {} def fake_rocm_sparse_attn_prefill(**kwargs): diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 6eb089ce3ed6..c21fc15a5278 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -47,13 +47,12 @@ from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.platforms import current_platform -# TODO: Integrate ROCMAiterMLASparseBackend for ROCm. -# The ROCm sparse MLA backend (rocm_aiter_mla_sparse.py) has a compatible -# forward_mqa interface but needs validation on ROCm hardware. -if not current_platform.is_cuda(): +# This module's end-to-end sparse-backend correctness test supports CUDA and +# ROCm. The remaining tests exercise CUDA-specific FlashMLA/FlashInfer APIs +# and are skipped by the autouse fixture below when collecting on ROCm. +if not (current_platform.is_cuda() or current_platform.is_rocm()): pytest.skip( - "Sparse MLA backend tests currently only support CUDA. " - "ROCm support requires integrating ROCMAiterMLASparseBackend.", + "Sparse MLA backend tests require CUDA or ROCm.", allow_module_level=True, ) @@ -65,14 +64,16 @@ ) from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import current_stream -from vllm.v1.attention.backends.mla import ( - flashattn_mla_sparse as flashattn_sparse_module, -) from vllm.v1.attention.backends.mla import index_group as index_group_module -from vllm.v1.attention.backends.mla.flashattn_mla_sparse import ( - FlashAttnMLASparseImpl, - FlashAttnMLASparseMetadataBuilder, -) + +if current_platform.is_cuda(): + from vllm.v1.attention.backends.mla import ( + flashattn_mla_sparse as flashattn_sparse_module, + ) + from vllm.v1.attention.backends.mla.flashattn_mla_sparse import ( + FlashAttnMLASparseImpl, + FlashAttnMLASparseMetadataBuilder, + ) from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( FlashInferMLASparseImpl, FlashInferMLASparseTRTLLMBackend, @@ -111,9 +112,43 @@ ) from vllm.v1.hisparse.types import SparseKVRowMirror +if current_platform.is_rocm(): + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( + ROCMAiterMLASparseBackend, + ) + + SPARSE_BACKENDS = [ROCMAiterMLASparseBackend] + SPARSE_BACKEND_IDS = ["ROCMAiterMLA"] + SPARSE_KV_CACHE_DTYPES = ["auto", "fp8"] + SPARSE_BLOCK_SIZES = [64] +else: + SPARSE_BACKENDS = [FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend] + SPARSE_BACKEND_IDS = ["FlashMLA", "FlashInferTRTLLM"] + SPARSE_KV_CACHE_DTYPES = ["auto", "fp8", "fp8_ds_mla", "nvfp4_ds_mla"] + SPARSE_BLOCK_SIZES = [32, 64] + + +@pytest.fixture(autouse=True) +def _skip_cuda_specific_sparse_mla_tests_on_rocm(request): + """Skip tests whose helpers exercise CUDA-only sparse MLA APIs.""" + rocm_portable_tests = { + "test_sparse_backend_decode_correctness", + "test_sparse_backend_prefill_correctness", + "test_rocm_fused_v32_dense_prefill_matches_mqa", + "test_rocm_long_bf16_prefill_avoids_nonpersistent_aiter", + } + if ( + current_platform.is_rocm() + and getattr(request.node, "originalname", request.node.name) + not in rocm_portable_tests + ): + pytest.skip("This sparse MLA test exercises a CUDA-specific API") + + SPARSE_BACKEND_BATCH_SPECS = { name: BATCH_SPECS[name] for name in [ + "small_decode", "mixed_small", "mixed_medium", "small_prefill", @@ -128,10 +163,61 @@ SPARSE_BACKEND_BATCH_SPECS["large_q_pure_prefill"] = BatchSpec( seq_lens=[256] * 2, query_lens=[256] * 2 ) +SPARSE_BACKEND_BATCH_SPECS["large_decode_over_topk"] = BatchSpec( + seq_lens=[256, 384], query_lens=[1, 1] +) DEVICE_TYPE = current_platform.device_type +@pytest.mark.parametrize( + ("batch_name", "tensor_parallel_size"), + [ + ("large_q_prefill", 2), + ("large_q_pure_prefill", 4), + ("large_q_pure_prefill", 1), + ], +) +def test_rocm_long_bf16_prefill_avoids_nonpersistent_aiter( + default_vllm_config, + dist_init, + workspace_init, + monkeypatch, + batch_name, + tensor_parallel_size, +): + """MI300 long BF16 prefills must remain correct without AITER ASM dispatch.""" + if not current_platform.is_rocm(): + pytest.skip("ROCm sparse MLA regression") + from vllm._aiter_ops import rocm_aiter_ops + + # Exercise the gfx942 routing decision on other ROCm devices too; + # the numerical kernel still compiles for the actual device. + monkeypatch.setattr("vllm.platforms.rocm.on_gfx942", lambda: True) + + def unsupported_aiter(*args, **kwargs): + pytest.fail("Long gfx942 BF16 prefill reached nonpersistent AITER") + + monkeypatch.setattr(rocm_aiter_ops, "mla_decode_fwd", unsupported_aiter) + test_sparse_backend_decode_correctness( + default_vllm_config=default_vllm_config, + dist_init=dist_init, + backend_cls=ROCMAiterMLASparseBackend, + batch_name=batch_name, + kv_cache_dtype="auto", + tensor_parallel_size=tensor_parallel_size, + block_size=64, + workspace_init=workspace_init, + q_scale=1.0, + k_scale=1.0, + total_num_heads=128, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + monkeypatch=monkeypatch, + ) + + def test_nope_flashinfer_sparse_mla_uses_model_scale(monkeypatch): """Weight absorption must not change the model's attention temperature.""" model_scale = 256**-0.5 @@ -403,16 +489,13 @@ def unpack_e2m1(packed: torch.Tensor) -> torch.Tensor: @pytest.mark.parametrize( "backend_cls", - [FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend], - ids=["FlashMLA", "FlashInferTRTLLM"], + SPARSE_BACKENDS, + ids=SPARSE_BACKEND_IDS, ) @pytest.mark.parametrize("batch_name", list(SPARSE_BACKEND_BATCH_SPECS.keys())) -@pytest.mark.parametrize( - "kv_cache_dtype", - ["auto", "fp8", "fp8_ds_mla", "nvfp4_ds_mla"], -) +@pytest.mark.parametrize("kv_cache_dtype", SPARSE_KV_CACHE_DTYPES) @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4]) -@pytest.mark.parametrize("block_size", [32, 64]) +@pytest.mark.parametrize("block_size", SPARSE_BLOCK_SIZES) @pytest.mark.parametrize(("q_scale", "k_scale"), [(1.0, 1.0), (2.0, 3.0)]) @pytest.mark.parametrize( ("total_num_heads", "qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim"), @@ -525,6 +608,12 @@ def test_sparse_backend_decode_correctness( "attn_module_list_cfg": [{"topk_tokens": topk_tokens}], }, ) + # This shared correctness test constructs a sparse SDPA reference for every + # query token. ROCm normally routes short prefills to dense MHA, whose + # output intentionally differs from that sparse reference. Keep this test + # on the sparse-MQA path; dense-prefill routing has dedicated ROCm tests. + if current_platform.is_rocm(): + vllm_config.attention_config.sparse_mla_force_mqa = True model_config = vllm_config.model_config model_config.hf_text_config = SimpleNamespace( index_topk=topk_tokens, @@ -737,7 +826,7 @@ def test_sparse_backend_decode_correctness( vllm_config=vllm_config, ) vllm_config.compilation_config.static_forward_context["placeholder"] = ( - SimpleNamespace(prefill_backend=prefill_backend) + SimpleNamespace(prefill_backend=prefill_backend, impl=SimpleNamespace()) ) builder_cls = backend_cls.get_builder_cls() @@ -1384,11 +1473,13 @@ def test_masked_mha_workspace_fits_accounts_for_batch_and_context(): @pytest.mark.skipif( - torch.cuda.get_device_capability()[0] < 10, + current_platform.is_cuda() and torch.cuda.get_device_capability()[0] < 10, reason="Sparse MLA forward_mha requires FA4 (SM100+)", ) @pytest.mark.parametrize("batch_name", list(PREFILL_BATCH_SPECS.keys())) -@pytest.mark.parametrize("kv_cache_dtype", ["auto"]) +@pytest.mark.parametrize( + "kv_cache_dtype", (["auto", "fp8"] if current_platform.is_rocm() else ["auto"]) +) @pytest.mark.parametrize( ("num_heads", "qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim"), [ @@ -1408,11 +1499,19 @@ def test_sparse_backend_prefill_correctness( v_head_dim, workspace_init, ): - """Test dense and masked MHA across supported sparse MLA dimensions.""" - backend_cls = FlashMLASparseBackend + """Test the normal dense forward_mha path for sparse MLA short prefill.""" + backend_cls = ( + ROCMAiterMLASparseBackend + if current_platform.is_rocm() + else FlashMLASparseBackend + ) + if current_platform.is_rocm() and ( + batch_name.startswith("masked_mha") or qk_nope_head_dim != 128 + ): + pytest.skip("ROCm dense prefill requires the DeepSeek MLA layout") batch_spec = PREFILL_BATCH_SPECS[batch_name] - device = torch.device("cuda") + device = torch.device(DEVICE_TYPE) dtype = torch.bfloat16 block_size = 64 @@ -1579,7 +1678,7 @@ def test_sparse_backend_prefill_correctness( vllm_config=vllm_config, ) vllm_config.compilation_config.static_forward_context["placeholder"] = ( - SimpleNamespace(prefill_backend=prefill_backend) + SimpleNamespace(prefill_backend=prefill_backend, impl=SimpleNamespace()) ) builder_cls = backend_cls.get_builder_cls() @@ -1653,7 +1752,10 @@ def test_sparse_backend_prefill_correctness( assert out_buffer.shape == ref_output.shape assert torch.isfinite(out_buffer).all(), "Non-finite values in output" - torch.testing.assert_close(out_buffer, ref_output, rtol=0.01, atol=0.01) + if kv_cache_dtype.startswith("fp8"): + torch.testing.assert_close(out_buffer, ref_output, rtol=0.065, atol=0.05) + else: + torch.testing.assert_close(out_buffer, ref_output, rtol=0.01, atol=0.01) @pytest.mark.parametrize( @@ -3057,7 +3159,7 @@ def test_hisparse_mixed_batch_bf16_row_split( vllm_config=vllm_config, ) vllm_config.compilation_config.static_forward_context["placeholder"] = ( - SimpleNamespace(prefill_backend=prefill_backend) + SimpleNamespace(prefill_backend=prefill_backend, impl=SimpleNamespace()) ) builder_cls = FlashMLASparseBackend.get_builder_cls() @@ -4264,3 +4366,118 @@ def test_sparse_mla_common_impl_resolves_buffer_lazily(): assert isinstance(SparseMLACommonImpl.topk_indices_buffer, property), ( "topk_indices_buffer must stay a lazily-resolved property" ) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm model implementation") +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +def test_rocm_fused_v32_dense_prefill_matches_mqa( + tmp_path, vllm_runner, kv_cache_dtype +): + """Fused normalization, RoPE and mixed dispatch preserve generation/logprobs. + + Chunking unequal prompts exercises prefill alongside decode. The two long + prompts straddle the dense-prefill threshold. No pretrained weights or + tokenizer download is needed. + """ + import json + + from vllm import SamplingParams + + config = { + "architectures": ["DeepseekV32ForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 0, + "eos_token_id": 1, + "ep_size": 1, + "first_k_dense_replace": 1, + "hidden_act": "silu", + "hidden_size": 2048, + "index_head_dim": 128, + "index_n_heads": 64, + "index_topk": 2048, + "initializer_range": 0.02, + "intermediate_size": 4096, + "kv_lora_rank": 512, + "max_position_embeddings": 8192, + "model_type": "deepseek_v32", + "moe_intermediate_size": 512, + "moe_layer_freq": 1, + "n_group": 1, + "n_routed_experts": 8, + "n_shared_experts": 1, + "norm_topk_prob": True, + "num_attention_heads": 32, + "num_experts_per_tok": 2, + "num_hidden_layers": 2, + "num_key_value_heads": 32, + "num_nextn_predict_layers": 0, + "q_lora_rank": 512, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "rms_norm_eps": 1e-06, + "rope_scaling": { + "beta_fast": 32, + "beta_slow": 1, + "factor": 40, + "mscale": 1.0, + "mscale_all_dim": 1.0, + "original_max_position_embeddings": 4096, + "type": "yarn", + }, + "rope_theta": 10000, + "routed_scaling_factor": 2.5, + "scoring_func": "sigmoid", + "tie_word_embeddings": False, + "topk_group": 1, + "topk_method": "noaux_tc", + "torch_dtype": "bfloat16", + "transformers_version": "4.44.2", + "use_cache": True, + "v_head_dim": 128, + "vocab_size": 4096, + } + (tmp_path / "config.json").write_text(json.dumps(config)) + prompts = [ + {"prompt_token_ids": [2 + (i * 17 + j * 13) % 4000 for i in range(n)]} + for j, n in enumerate([64, 128, 2048, 2049]) + ] + sampling = SamplingParams(temperature=0, max_tokens=16, ignore_eos=True, logprobs=5) + results = [] + for force_mqa in [True, False]: + with vllm_runner( + str(tmp_path), + load_format="dummy", + skip_tokenizer_init=True, + dtype="bfloat16", + kv_cache_dtype=kv_cache_dtype, + seed=42, + max_model_len=4096, + max_num_seqs=4, + max_num_batched_tokens=256, + block_size=64, + enable_chunked_prefill=True, + enable_prefix_caching=False, + enforce_eager=True, + gpu_memory_utilization=0.25, + attention_config={ + "backend": "ROCM_AITER_MLA_SPARSE", + "sparse_mla_force_mqa": force_mqa, + }, + ) as runner: + outputs = runner.llm.generate(prompts, sampling, use_tqdm=False) + results.append([request.outputs[0] for request in outputs]) + + for reference, actual in zip(*results): + assert len(actual.token_ids) == 16 + assert actual.token_ids == reference.token_ids + assert actual.logprobs is not None and reference.logprobs is not None + for token, expected, observed in zip( + actual.token_ids, reference.logprobs, actual.logprobs + ): + torch.testing.assert_close( + torch.tensor(observed[token].logprob), + torch.tensor(expected[token].logprob), + rtol=0.01, + atol=0.01, + ) diff --git a/vllm/models/deepseek_v32/amd/rocm.py b/vllm/models/deepseek_v32/amd/rocm.py index b9ac31f8be2b..7211edaa3a81 100644 --- a/vllm/models/deepseek_v32/amd/rocm.py +++ b/vllm/models/deepseek_v32/amd/rocm.py @@ -1,11 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import cast + import torch from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache from vllm.models.deepseek_v32.attention import DeepseekV32Attention, DeepseekV32Indexer @@ -179,6 +182,8 @@ def _fused_attention( else: attn_metadata = attn_metadata_raw + attn_metadata = cast(MLACommonMetadata | None, attn_metadata) + slot_mapping = forward_context.slot_mapping assert isinstance(slot_mapping, dict) mla_slot = slot_mapping.get(self.layer_name) @@ -213,6 +218,10 @@ def _fused_attention( mla_kv_cache = self.kv_cache mla_k_scale = self._k_scale + use_mha = attn_metadata is not None and self._use_sparse_mha(attn_metadata) + kv_c_out = torch.empty_like(kv_c) if use_mha else None + k_pe_out = torch.empty_like(k_pe) if use_mha else None + q_c = fused_norm_rope( positions, q_c, @@ -237,9 +246,18 @@ def _fused_attention( mla_k_scale=mla_k_scale, has_indexer=has_indexer, index_rope_interleave=self._index_rope_interleave, + kv_c_out=kv_c_out, + k_pe_out=k_pe_out, ) - ql_nope, q_pe = self._compute_ql_nope(q_c) + if use_mha: + q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim) + q_nope, q_pe = q.split( + [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + ql_nope = q_nope + else: + ql_nope, q_pe = self._compute_ql_nope(q_c) if self.indexer is not None and not self.skip_topk: index_q = self.indexer.wq_b(q_c)[0] @@ -260,7 +278,7 @@ def _fused_attention( indexer_n_head_scale, has_indexer=has_indexer, index_rope_interleave=self._index_rope_interleave, - quantize_mqa=self._fp8_kv, + quantize_mqa=self._fp8_kv and not use_mha, ) if self.indexer is not None and not self.skip_topk: @@ -270,6 +288,19 @@ def _fused_attention( output.zero_() return + if use_mha: + assert kv_c_out is not None and k_pe_out is not None + mha_q = torch.cat((q_nope, mqa_q), dim=-1) + self.forward_impl( + mha_q, + kv_c_out, + k_pe_out.unsqueeze(1), + self.kv_cache, + attn_metadata, + output, + ) + return + num_actual = attn_metadata.num_actual_tokens # type: ignore[attr-defined] kv_cache = self.kv_cache if self._fp8_kv_needs_view: diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index e30fec715deb..a0393e772c3e 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -12,11 +12,10 @@ from vllm.config import VllmConfig, get_current_vllm_config from vllm.config.cache import CacheDType from vllm.logger import init_logger -from vllm.model_executor.layers.attention.mla_attention import ( - get_mla_dims, -) from vllm.model_executor.layers.attention.sparse_mla_attention import ( - SharedTopkIndicesBuffer, + SparseMLACommonImpl, + SparseMLACommonMetadata, + SparseMLACommonMetadataBuilder, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -26,10 +25,7 @@ AttentionBackend, AttentionCGSupport, AttentionLayer, - AttentionMetadata, - AttentionMetadataBuilder, CommonAttentionMetadata, - MLAAttentionImpl, MultipleOf, ) from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( @@ -39,7 +35,6 @@ align_blocks_to_rows, flat_kv_row_view, ) -from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( rocm_sparse_attn_decode_bf16, rocm_sparse_attn_prefill, @@ -76,6 +71,18 @@ def _use_rocm_sparse_triton( ) +def _is_long_bf16_prefill( + kv_dtype: torch.dtype, + metadata: "ROCMAiterMLASparseMetadata", +) -> bool: + return ( + kv_dtype == torch.bfloat16 + and metadata.num_decodes == 0 + and metadata.num_prefills > 0 + and metadata.prefill_max_seq_len > metadata.topk_tokens + ) + + @triton.jit def _fit_kpool_indices_kernel( token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS + TAIL_WIDTH] @@ -239,6 +246,10 @@ def triton_convert_req_index_to_global_index( assert req_id.dtype == torch.int32 assert block_table.dtype == torch.int32 assert token_indices.dtype == torch.int32 + assert req_id.shape[0] == token_indices.shape[0], ( + f"req_id ({req_id.shape[0]}) and token_indices ({token_indices.shape[0]}) " + "must cover the same tokens" + ) assert token_indices.shape[1] == NUM_TOPK_TOKENS assert NUM_TOPK_TOKENS % BLOCK_N == 0, ( f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible byBLOCK_N ({BLOCK_N})" @@ -390,35 +401,13 @@ def supports_sink(cls) -> bool: return on_mi3xx() -@dataclass -class ROCMAiterMLASparseMetadata(AttentionMetadata): - num_reqs: int - max_query_len: int - max_seq_len: int - - num_actual_tokens: int # Number of tokens excluding padding. - query_start_loc: torch.Tensor - slot_mapping: torch.Tensor - - block_table: torch.Tensor - req_id_per_token: torch.Tensor - - qo_indptr: torch.Tensor - paged_kv_last_page_len: torch.Tensor - paged_kv_indices: torch.Tensor - paged_kv_indptr: torch.Tensor - attn_out_dtype: torch.dtype - - block_size: int = 1 - topk_tokens: int = 2048 - - # Fields read by the shared MLA forward. This impl has no dense-MHA prefill - # path (supports_dense_mha_prefill=False), so it always runs the MQA path; - num_decodes: int = 0 - num_prefills: int = 0 - num_decode_tokens: int = 0 - prefill_max_seq_len: int = 0 - prefill: object = None +@dataclass(kw_only=True) +class ROCMAiterMLASparseMetadata(SparseMLACommonMetadata): + qo_indptr: torch.Tensor | None = None + paged_kv_last_page_len: torch.Tensor | None = None + paged_kv_indices: torch.Tensor | None = None + paged_kv_indptr: torch.Tensor | None = None + attn_out_dtype: torch.dtype | None = None # Persistent MLA metadata (only populated when persistent mode is enabled, # i.e. when the aiter sparse decode kernel supports work-stealing splits). @@ -432,9 +421,10 @@ class ROCMAiterMLASparseMetadata(AttentionMetadata): @dataclass class ROCMAiterMLASparseMetadataBuilder( - AttentionMetadataBuilder[ROCMAiterMLASparseMetadata] + SparseMLACommonMetadataBuilder[ROCMAiterMLASparseMetadata] ): _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + metadata_cls = ROCMAiterMLASparseMetadata def __init__( self, @@ -443,20 +433,15 @@ def __init__( vllm_config: VllmConfig, device: torch.device, ): - self.kv_cache_spec = kv_cache_spec - self.model_config = vllm_config.model_config + super().__init__(kv_cache_spec, layer_names, vllm_config, device) self.model_dtype = vllm_config.model_config.dtype self.kv_cache_dtype = vllm_config.cache_config.cache_dtype parallel_config = vllm_config.parallel_config - self.device = device max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens - self.vllm_config = vllm_config self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) self.num_heads = self.model_config.get_num_attention_heads(parallel_config) - self.mla_dims = get_mla_dims(self.model_config) - self.topk_tokens = vllm_config.model_config.hf_text_config.index_topk attention_context = vllm_config.compilation_config.static_forward_context # Sink decode must use AITER's nonpersistent path. In particular, # gfx942 has no persistent+LSE kernel, and its metadata heuristic @@ -521,9 +506,12 @@ def __init__( kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto") if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"): kv_cache_dtype_str = "fp8" + q_dtype = dtypes.fp8 else: kv_cache_dtype_str = "bf16" kv_dtype = dtypes.d_dtypes.get(kv_cache_dtype_str, dtypes.bf16) + self._mla_q_dtype = q_dtype + self._mla_kv_dtype = kv_dtype ( (work_meta_data_size, work_meta_data_type), @@ -587,13 +575,19 @@ def build( common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> ROCMAiterMLASparseMetadata: - num_tokens = common_attn_metadata.num_actual_tokens - (num_decodes, num_prefills, num_decode_tokens, _) = split_decodes_and_prefills( - common_attn_metadata, - decode_threshold=self.reorder_batch_threshold or 1, + metadata = super().build( + common_prefix_len, common_attn_metadata, fast_build=fast_build + ) + use_dense_prefill = ( + metadata.prefill is not None and metadata.prefill.use_dense_mha ) + num_reqs = metadata.num_decodes if use_dense_prefill else metadata.num_reqs + num_tokens = int(common_attn_metadata.query_start_loc_cpu[num_reqs].item()) + num_decodes = metadata.num_decodes + num_prefills = metadata.num_prefills + num_decode_tokens = metadata.num_decode_tokens starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32) - seg_lengths = np.diff(starts) + seg_lengths = np.diff(starts[: num_reqs + 1]) req_id_per_token = np.repeat( np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths ) @@ -616,15 +610,13 @@ def build( self.req_id_per_token_buffer[:new_req_extent].copy_( np_to_pinned_tensor(req_id_per_token), non_blocking=True ) - query_lens = ( - common_attn_metadata.query_start_loc[1:] - - common_attn_metadata.query_start_loc[:-1] - ) - seq_lens = common_attn_metadata.seq_lens + query_start_loc = common_attn_metadata.query_start_loc[: num_reqs + 1] + query_lens = query_start_loc[1:] - query_start_loc[:-1] + seq_lens = common_attn_metadata.seq_lens[:num_reqs] sparse_seqlen = generate_sparse_seqlen_triton( query_lens, seq_lens, - common_attn_metadata.query_start_loc, + query_start_loc, self.topk_tokens, num_tokens, common_attn_metadata.max_query_len, @@ -663,8 +655,13 @@ def build( reduce_indptr = None reduce_final_map = None reduce_partial_map = None - if self._use_persistent_metadata and not use_triton_sparse: - num_reqs = common_attn_metadata.num_reqs + is_long_pure_bf16_prefill = _is_long_bf16_prefill(self._mla_kv_dtype, metadata) + if ( + self._use_persistent_metadata + and not use_triton_sparse + and not is_long_pure_bf16_prefill + and num_tokens > 0 + ): with gpu_sync_allowed(): seq_lens_cpu = common_attn_metadata.seq_lens[:num_reqs].cpu().numpy() clamped_seq_lens = np.minimum( @@ -708,6 +705,8 @@ def build( uni_seqlen_qo=1, fast_mode=True, max_split_per_batch=max_split_per_batch, + dtype_q=self._mla_q_dtype, + dtype_kv=self._mla_kv_dtype, ) torch.cuda.current_stream(self.device).synchronize() self._prev_metadata_key = metadata_key @@ -718,32 +717,18 @@ def build( reduce_final_map = self._mla_reduce_final_map reduce_partial_map = self._mla_reduce_partial_map - metadata = ROCMAiterMLASparseMetadata( - num_reqs=common_attn_metadata.num_reqs, - max_query_len=common_attn_metadata.max_query_len, - max_seq_len=common_attn_metadata.max_seq_len, - num_actual_tokens=common_attn_metadata.num_actual_tokens, - query_start_loc=common_attn_metadata.query_start_loc, - slot_mapping=common_attn_metadata.slot_mapping, - block_table=common_attn_metadata.block_table_tensor, - req_id_per_token=req_id_per_token, - block_size=self.kv_cache_spec.block_size, - attn_out_dtype=self.model_dtype, - topk_tokens=self.topk_tokens, - num_decodes=num_decodes, - num_prefills=num_prefills, - num_decode_tokens=num_decode_tokens, - qo_indptr=qo_indptr, - paged_kv_last_page_len=paged_kv_last_page_len, - paged_kv_indices=paged_kv_indices, - paged_kv_indptr=paged_kv_indptr, - work_meta_data=work_meta_data, - work_indptr=work_indptr, - work_info_set=work_info_set, - reduce_indptr=reduce_indptr, - reduce_final_map=reduce_final_map, - reduce_partial_map=reduce_partial_map, - ) + metadata.attn_out_dtype = self.model_dtype + metadata.req_id_per_token = req_id_per_token + metadata.qo_indptr = qo_indptr + metadata.paged_kv_last_page_len = paged_kv_last_page_len + metadata.paged_kv_indices = paged_kv_indices + metadata.paged_kv_indptr = paged_kv_indptr + metadata.work_meta_data = work_meta_data + metadata.work_indptr = work_indptr + metadata.work_info_set = work_info_set + metadata.reduce_indptr = reduce_indptr + metadata.reduce_final_map = reduce_final_map + metadata.reduce_partial_map = reduce_partial_map return metadata @@ -776,11 +761,7 @@ def log2sumexp2(a: torch.Tensor, dim: int) -> torch.Tensor: return (result, lse) -class ROCMAiterMLASparseImpl( - MLAAttentionImpl[ROCMAiterMLASparseMetadata], SharedTopkIndicesBuffer -): - is_sparse = True - supports_dense_mha_prefill = False +class ROCMAiterMLASparseImpl(SparseMLACommonImpl[ROCMAiterMLASparseMetadata]): supports_dcp = False use_aiter_sparse_mla = False @@ -803,10 +784,6 @@ def __init__( ) -> None: AiterMLAHelper.check_num_heads_validity(num_heads) - self.num_heads = num_heads - self.head_size = head_size - self.scale = float(scale) - self.num_kv_heads = num_kv_heads sinks = mla_args.pop("sinks", None) if sinks is not None: if sinks.dtype != torch.float32: @@ -821,10 +798,24 @@ def __init__( if not sinks.is_contiguous(): raise ValueError("ROCm AITER MLA sinks must be contiguous") self.sinks: torch.Tensor | None = sinks - self.kv_cache_dtype = kv_cache_dtype - self.kv_lora_rank: int = mla_args["kv_lora_rank"] + super().__init__( + num_heads, + head_size, + scale, + num_kv_heads, + alibi_slopes, + sliding_window, + kv_cache_dtype, + logits_soft_cap, + attn_type, + kv_sharing_target_layer_name, + indexer=indexer, + topk_indices_buffer=topk_indices_buffer, + **mla_args, + ) self.softmax_scale = scale - self.init_topk_indices_buffer(indexer, topk_indices_buffer) + self.supports_dense_mha_prefill = sinks is None # type: ignore[misc] + self.masked_mha_available = False vllm_config = get_current_vllm_config() max_tokens = vllm_config.scheduler_config.max_num_batched_tokens @@ -894,14 +885,26 @@ def _forward_mla( ) ) - if triton_sink_fallback or _use_rocm_sparse_triton( - kv_cache_dtype=self.kv_cache_dtype, - head_size=q.shape[-1], - kv_lora_rank=self.kv_lora_rank, - num_prefills=attn_metadata.num_prefills, - num_decodes=attn_metadata.num_decodes, - num_decode_tokens=attn_metadata.num_decode_tokens, - max_query_len=attn_metadata.max_query_len, + # Long BF16 prefills omit persistent metadata. gfx942's + # nonpersistent ASM path is unsupported or unsafe for these shapes. + # Sink attention has its own Triton/LSE dispatch below. + triton_prefill_fallback = ( + not need_lse + and on_gfx942() + and _is_long_bf16_prefill(kv_c_and_k_pe_cache.dtype, attn_metadata) + ) + if ( + triton_sink_fallback + or triton_prefill_fallback + or _use_rocm_sparse_triton( + kv_cache_dtype=self.kv_cache_dtype, + head_size=q.shape[-1], + kv_lora_rank=self.kv_lora_rank, + num_prefills=attn_metadata.num_prefills, + num_decodes=attn_metadata.num_decodes, + num_decode_tokens=attn_metadata.num_decode_tokens, + max_query_len=attn_metadata.max_query_len, + ) ): output = torch.empty( [num_tokens, q.shape[1], self.kv_lora_rank], @@ -1078,6 +1081,7 @@ def _forward_mla( # Empty ragged rows have only sink mass and no value contribution. # AITER can return NaN output/LSE for those rows; do not multiply it # by a zero normalization factor and propagate the NaN. + assert attn_metadata.paged_kv_indptr is not None has_keys = ( attn_metadata.paged_kv_indptr[1:] > attn_metadata.paged_kv_indptr[:-1] ).unsqueeze(-1) @@ -1156,7 +1160,7 @@ def forward_mqa( else: ops.concat_mla_q(ql_nope, q_pe, q) - num_actual_toks = attn_metadata.num_actual_tokens + num_actual_toks = q.shape[0] # Get topk indices assert self.topk_indices_buffer is not None @@ -1168,7 +1172,7 @@ def forward_mqa( kv_c_and_k_pe_cache, attn_metadata.block_size ) triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token, + attn_metadata.req_id_per_token[:num_actual_toks], attn_metadata.block_table, topk_indices, attn_metadata.paged_kv_indptr, From e0317b28ce3ba4c9bd2ff69d885e3f34688a6d71 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Mon, 12 Oct 2026 08:53:26 -0700 Subject: [PATCH 314/624] [Feature] Support batch invariance on ROCm (#52231) Signed-off-by: Matthew Wong Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Andreas Karatzas --- .buildkite/test-amd.yaml | 89 +++- .buildkite/test_areas/misc.yaml | 112 +++++- docs/features/batch_invariance.md | 1 + .../quantization/test_triton_scaled_mm.py | 25 +- tests/v1/determinism/test_batch_invariance.py | 1 + .../determinism/test_batch_invariance_vlm.py | 6 +- .../test_cutlass_batch_invariance.py | 118 +++++- .../test_lora_shrink_batch_invariant.py | 6 +- .../test_matmul_batch_invariant.py | 52 ++- .../test_mx_linear_batch_invariant.py | 165 ++++++++ .../determinism/test_nvfp4_batch_invariant.py | 10 +- .../test_online_batch_invariance.py | 6 +- .../test_rms_norm_batch_invariant.py | 18 +- tests/v1/determinism/utils.py | 22 +- .../device_communicators/cuda_communicator.py | 39 +- .../determinism/batch_invariant.py | 380 +++++++++++++++--- .../model_executor/kernels/linear/__init__.py | 43 ++ .../kernels/linear/mxfp4/aiter.py | 18 +- .../kernels/linear/mxfp8/rocm_native.py | 6 +- .../kernels/linear/scaled_mm/__init__.py | 2 + .../kernels/linear/scaled_mm/triton.py | 42 ++ .../fused_moe/experts/aiter_mxfp8_moe.py | 4 + .../fused_moe/experts/ocp_mx_emulation_moe.py | 2 +- .../compressed_tensors/triton_scaled_mm.py | 12 +- .../model_executor/layers/quantization/fp8.py | 4 +- .../layers/quantization/online/fp8.py | 10 +- .../layers/quantization/utils/fp8_utils.py | 5 +- vllm/model_executor/layers/utils.py | 7 + vllm/platforms/rocm.py | 9 + .../backends/mla/prefill/aiter_flash_attn.py | 10 + .../backends/mla/prefill/flash_attn.py | 4 +- 31 files changed, 1106 insertions(+), 122 deletions(-) create mode 100644 tests/v1/determinism/test_mx_linear_batch_invariant.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index b4c4b59fdfe8..49bc3c32b67a 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -273,7 +273,7 @@ steps: #------------------------------------------------------------ mi250 · v1 -------------------------------------------------------------# - label: ":amd: (MI250) Batch Invariance" # TBD - timeout_in_minutes: 35 + timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] dind: false agent_pool: mi250_1 @@ -282,14 +282,18 @@ steps: source_file_dependencies: - vllm/v1/attention - vllm/model_executor/layers + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism - tests/v1/determinism/ - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: + - export VLLM_GPU_MEMORY_UTILIZATION=0.8 + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA - label: ":amd: (MI250) CUDAGraph" # TBD timeout_in_minutes: 55 @@ -2443,6 +2447,56 @@ steps: #------------------------------------------------------------ mi300 · v1 -------------------------------------------------------------# +- label: ":amd: (MI300) Batch Invariance" # TBD + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pip install pytest-timeout pytest-forked + - pytest -v -s v1/determinism/test_batch_invariance.py + - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN + +- label: ":amd: (MI300) Batch Invariance (Distributed)" # TBD + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/device_communicators/ + - vllm/model_executor/determinism + - vllm/model_executor/layers + - vllm/v1/attention + - tests/v1/determinism/ + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - VLLM_TEST_TP_SIZE=4 pytest -v -s v1/determinism/test_batch_invariance.py -k "bs1_vs_bsN or without_batch_invariance_should_fail" + - VLLM_TEST_TP_SIZE=4 VLLM_TEST_DISABLE_CUSTOM_AR=1 pytest -v -s v1/determinism/test_batch_invariance.py -k bs1_vs_bsN + - VLLM_TP_SIZE=4 VLLM_NEEDLE_TRIALS=2 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN + - VLLM_TEST_MODEL=allenai/OLMoE-1B-7B-0924 VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py + - label: ":amd: (MI300) Spec Decode Speculators + MTP" # TBD timeout_in_minutes: 75 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -5215,6 +5269,37 @@ steps: #------------------------------------------------------------ mi355 · v1 -------------------------------------------------------------# +- label: ":amd: (MI355) Batch Invariance" # TBD + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pip install pytest-timeout pytest-forked + - pytest -v -s v1/determinism/test_batch_invariance.py + - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py + - pytest -v -s v1/determinism/test_batch_invariance_vlm.py + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN + - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_online_batch_invariance.py + - pytest -v -s v1/determinism/test_mx_linear_batch_invariant.py + - label: ":amd: (MI355) Spec Decode AL DFlash2 Nightly" timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 20b48cb98f1f..02f78e280c98 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -644,12 +644,33 @@ steps: source_file_dependencies: - vllm/v1/attention - vllm/model_executor/layers + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism - tests/v1/determinism/ commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + mirror: + amd: + label: ":amd: (MI250) Batch Invariance" + dind: false + device: mi250_1 + timeout_in_minutes: 60 + env: + VLLM_GPU_MEMORY_UTILIZATION: "0.8" + VLLM_TEST_CLEAN_GPU_MEMORY: "1" + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py - label: ":nvidia: (H100) Batch Invariance" key: batch-invariance-h100 @@ -660,6 +681,8 @@ steps: - vllm/v1/attention - vllm/model_executor/layers - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism - tests/v1/determinism/ commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn @@ -671,15 +694,18 @@ steps: - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN mirror: amd: - label: ":amd: (MI250) Batch Invariance" + label: ":amd: (MI300) Batch Invariance" dind: false - device: mi250_1 - timeout_in_minutes: 35 + device: mi300_1 + timeout_in_minutes: 60 depends_on: - image-build-amd source_file_dependencies: - vllm/v1/attention - vllm/model_executor/layers + - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism - tests/v1/determinism/ - vllm/_aiter_ops.py - vllm/platforms/rocm.py @@ -688,6 +714,11 @@ steps: - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN - label: ":nvidia: (B200) Batch Invariance" key: batch-invariance-b200 @@ -697,6 +728,8 @@ steps: - vllm/v1/attention - vllm/model_executor/layers - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism - tests/v1/determinism/ commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn @@ -711,6 +744,79 @@ steps: - pytest -v -s v1/determinism/test_matmul_batch_invariant.py - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py - pytest -v -s v1/determinism/test_online_batch_invariance.py + mirror: + amd: + label: ":amd: (MI355) Batch Invariance" + dind: false + device: mi355_1 + timeout_in_minutes: 60 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pip install pytest-timeout pytest-forked + - pytest -v -s v1/determinism/test_batch_invariance.py + - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py + - pytest -v -s v1/determinism/test_batch_invariance_vlm.py + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN + - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_online_batch_invariance.py + - pytest -v -s v1/determinism/test_mx_linear_batch_invariant.py + +- label: ":nvidia: (H100) Batch Invariance (Distributed)" + key: batch-invariance-distributed-4xh100 + timeout_in_minutes: 45 + device: h100 + optional: true + num_devices: 4 + source_file_dependencies: + - vllm/distributed/device_communicators/ + - vllm/model_executor/determinism + - vllm/model_executor/layers + - vllm/v1/attention + - tests/v1/determinism/ + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - VLLM_TEST_TP_SIZE=4 pytest -v -s v1/determinism/test_batch_invariance.py -k "bs1_vs_bsN or without_batch_invariance_should_fail" + - VLLM_TEST_TP_SIZE=4 VLLM_TEST_DISABLE_CUSTOM_AR=1 pytest -v -s v1/determinism/test_batch_invariance.py -k bs1_vs_bsN + - VLLM_TP_SIZE=4 VLLM_NEEDLE_TRIALS=2 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN + - VLLM_TEST_MODEL=allenai/OLMoE-1B-7B-0924 VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py + mirror: + amd: + label: ":amd: (MI300) Batch Invariance (Distributed)" + dind: false + device: mi300_4 + timeout_in_minutes: 45 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/device_communicators/ + - vllm/model_executor/determinism + - vllm/model_executor/layers + - vllm/v1/attention + - tests/v1/determinism/ + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - VLLM_TEST_TP_SIZE=4 pytest -v -s v1/determinism/test_batch_invariance.py -k "bs1_vs_bsN or without_batch_invariance_should_fail" + - VLLM_TEST_TP_SIZE=4 VLLM_TEST_DISABLE_CUSTOM_AR=1 pytest -v -s v1/determinism/test_batch_invariance.py -k bs1_vs_bsN + - VLLM_TP_SIZE=4 VLLM_NEEDLE_TRIALS=2 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN + - VLLM_TEST_MODEL=allenai/OLMoE-1B-7B-0924 VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py - label: ":nvidia: (H200 MIG 35GB) Acceptance Length (Large Models)" device: h200_35gb diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index 03d07896cce4..01dc7116fef9 100644 --- a/docs/features/batch_invariance.md +++ b/docs/features/batch_invariance.md @@ -20,6 +20,7 @@ Batch invariance is crucial for several use cases: Batch invariance is supported on the following platforms: - NVIDIA GPUs with compute capability 8.0 or higher. +- AMD Instinct GPUs (gfx90a, gfx942, gfx950). - Intel XPUs with Triton support. ### Attention Backend Selection for XPU diff --git a/tests/kernels/quantization/test_triton_scaled_mm.py b/tests/kernels/quantization/test_triton_scaled_mm.py index 55d1f76b2d1e..b0a2c215ab02 100644 --- a/tests/kernels/quantization/test_triton_scaled_mm.py +++ b/tests/kernels/quantization/test_triton_scaled_mm.py @@ -32,11 +32,9 @@ def torch_scaled_mm( out = torch.mm(a.to(torch.float32), b.to(torch.float32)) out = scale_a * out out = scale_b.T * out - out = out.to(out_dtype) if bias is not None: - out = out + bias - - return out + out = out + bias.float() + return out.to(out_dtype) def get_8bit_types(): @@ -125,6 +123,25 @@ def test_scaled_mm( torch.testing.assert_close(c_check, c_actual, rtol=1e-1, atol=1e-1) +@pytest.mark.parametrize("in_dtype", get_8bit_types()) +@pytest.mark.parametrize("use_heuristic", [True, False]) +def test_scaled_mm_adds_bias_before_rounding(in_dtype, use_heuristic): + # The accumulator is 1 + 3 * 2**-10 and the bias 2**-8. Rounding to bf16 + # first gives 1.0 (1 + 2**-8 is a tie, to even); adding first, 1 + 2**-7. + a = torch.ones((1, 32), device=device).to(in_dtype) + col = torch.tensor([64.0] * 16 + [2.0, 1.0] + [0.0] * 14, device=device) + b = col[:, None].expand(32, 16).to(in_dtype).contiguous() + scale_a = torch.ones((1, 1), device=device) + scale_b = torch.full((1, 1), 2**-10, device=device) + bias = torch.full((16,), 2**-8, device=device, dtype=torch.bfloat16) + args = (a, b, scale_a, scale_b, torch.bfloat16, bias) + + out = triton_scaled_mm(*args, use_heuristic=use_heuristic) + expected = torch_scaled_mm(*args) + assert (expected == 1 + 2**-7).all() + torch.testing.assert_close(out, expected, rtol=0, atol=0) + + # TD operand loads must be bit-exact vs the plain masked-load path. @pytest.mark.skipif( not (current_platform.is_cuda_alike() or current_platform.is_xpu()), diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index 452362715111..139c68ed2b81 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/tests/v1/determinism/test_batch_invariance.py @@ -180,6 +180,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( llm = LLM( model=TEST_MODEL, tensor_parallel_size=tp_size, + disable_custom_all_reduce=os.getenv("VLLM_TEST_DISABLE_CUSTOM_AR") == "1", max_num_seqs=128, max_model_len=8192, dtype="auto", # not everything is supported diff --git a/tests/v1/determinism/test_batch_invariance_vlm.py b/tests/v1/determinism/test_batch_invariance_vlm.py index fe95b03c61fa..5c557a1364d3 100644 --- a/tests/v1/determinism/test_batch_invariance_vlm.py +++ b/tests/v1/determinism/test_batch_invariance_vlm.py @@ -9,7 +9,7 @@ import pytest import torch from PIL import Image -from utils import _extract_step_logprobs, skip_if_not_cuda +from utils import _extract_step_logprobs, skip_if_not_cuda_alike from vllm import LLM, SamplingParams from vllm.inputs import TextPrompt @@ -89,7 +89,7 @@ def _assert_batch_invariant(llm: LLM, inputs: list[TextPrompt], sampling) -> Non ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("input_type", ["image", "video"]) @pytest.mark.parametrize("mm_encoder_attn_backend", ["FLASH_ATTN", "TORCH_SDPA"]) def test_vlm_batch_invariance_bs1_vs_bsN( @@ -101,7 +101,7 @@ def test_vlm_batch_invariance_bs1_vs_bsN( _run_vlm_batch_invariance(input_type, mm_encoder_attn_backend, enforce_eager=True) -@skip_if_not_cuda +@skip_if_not_cuda_alike def test_vlm_batch_invariance_default_execution(monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0") _run_vlm_batch_invariance("image", "FLASH_ATTN", enforce_eager=False) diff --git a/tests/v1/determinism/test_cutlass_batch_invariance.py b/tests/v1/determinism/test_cutlass_batch_invariance.py index c3384a62f6f0..b8f83817065c 100644 --- a/tests/v1/determinism/test_cutlass_batch_invariance.py +++ b/tests/v1/determinism/test_cutlass_batch_invariance.py @@ -19,8 +19,10 @@ from tests.kernels.quantization.nvfp4_utils import get_nvfp4_global_scale from tests.utils import TestFP8Layer, requires_fp8 from vllm import _custom_ops as ops -from vllm.model_executor.kernels.linear.scaled_mm.cutlass import ( +from vllm.model_executor.kernels.linear import fp8_scaled_mm_is_batch_invariant +from vllm.model_executor.kernels.linear.scaled_mm import ( CutlassFP8ScaledMMLinearKernel, + TritonFP8ScaledMMLinearKernel, ) from vllm.model_executor.layers.fused_moe import fused_topk from vllm.model_executor.layers.fused_moe.activation import MoEActivation @@ -31,8 +33,11 @@ CutlassExpertsFp4, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8DynamicTensorSym, kFp8DynamicTokenSym, + kFp8StaticChannelSym, kFp8StaticTensorSym, + kFp8StaticTokenSym, ) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -87,16 +92,42 @@ @pytest.fixture(autouse=True) def setup_cuda(): - if not current_platform.is_cuda(): - pytest.skip("CUTLASS FP8 kernels require CUDA.") + if not current_platform.is_cuda_alike(): + pytest.skip("Requires CUDA or ROCm.") torch.set_default_device("cuda") @requires_fp8 +@pytest.mark.parametrize( + "kernel", + [ + pytest.param( + CutlassFP8ScaledMMLinearKernel, + marks=pytest.mark.skipif( + not current_platform.is_cuda(), reason="CUTLASS requires CUDA." + ), + ), + TritonFP8ScaledMMLinearKernel, + ], +) +@pytest.mark.parametrize( + "act_key,weight_key", + [ + (kFp8DynamicTokenSym, kFp8StaticTensorSym), + (kFp8DynamicTokenSym, kFp8StaticChannelSym), + (kFp8DynamicTokenSym, kFp8StaticTokenSym), + (kFp8StaticTensorSym, kFp8StaticTensorSym), + # An amax over the whole batch: the GEMM is invariant, the layer is not. + (kFp8DynamicTensorSym, kFp8StaticTensorSym), + ], +) @pytest.mark.parametrize("weight_shape", [(1024, 2048), (4608, 4096)]) @pytest.mark.parametrize("batch_size", [1, 16, 17, 32, 64, 65, 256, 257]) @torch.inference_mode() -def test_cutlass_fp8_batch_invariant_fixed_config( +def test_fp8_scaled_mm_batch_invariant_fixed_config( + kernel: type, + act_key, + weight_key, weight_shape: tuple[int, int], batch_size: int, default_vllm_config, @@ -108,14 +139,16 @@ def test_cutlass_fp8_batch_invariant_fixed_config( torch.manual_seed(0) layer = TestFP8Layer( weight_shape=weight_shape, - activation_quant_key=kFp8DynamicTokenSym, - weight_quant_key=kFp8StaticTensorSym, + activation_quant_key=act_key, + weight_quant_key=weight_key, input_dtype=torch.bfloat16, out_dtype=torch.bfloat16, device=torch.device("cuda"), - force_kernel=CutlassFP8ScaledMMLinearKernel, + force_kernel=kernel, ) - assert isinstance(layer.kernel, CutlassFP8ScaledMMLinearKernel) + assert isinstance(layer.kernel, kernel) + dynamic_per_tensor = act_key == kFp8DynamicTensorSym + assert fp8_scaled_mm_is_batch_invariant(layer.kernel) != dynamic_per_tensor in_features = weight_shape[1] needle = torch.randn((1, in_features), device="cuda", dtype=torch.bfloat16) @@ -124,6 +157,12 @@ def test_cutlass_fp8_batch_invariant_fixed_config( filler = torch.randn( (max(batch_size - 1, 0), in_features), device="cuda", dtype=torch.bfloat16 ) + if dynamic_per_tensor: + if batch_size == 1: + pytest.skip("no other rows to move the scale") + front_output = layer(torch.cat([needle, 100 * filler], dim=0))[0] + assert not torch.equal(front_output, baseline) + return front_batch = torch.cat([needle, filler], dim=0) back_batch = torch.cat([filler, needle], dim=0) @@ -135,6 +174,69 @@ def test_cutlass_fp8_batch_invariant_fixed_config( torch.testing.assert_close(back_output, baseline, rtol=0, atol=0) +@requires_fp8 +@pytest.mark.parametrize("scheme", ["fp8_per_tensor", "fp8_per_channel"]) +@torch.inference_mode() +def test_online_fp8_linear_batch_invariant( + scheme, monkeypatch, default_vllm_config, dist_init +): + from vllm.config.model import ModelConfig + from vllm.config.quantization import resolve_quantization_config + from vllm.model_executor.layers.linear import ReplicatedLinear + from vllm.model_executor.layers.quantization.online.base import ( + OnlineQuantizationConfig, + ) + from vllm.utils.torch_utils import set_default_torch_dtype + + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True) + default_vllm_config.model_config = ModelConfig(dtype="bfloat16") + quant_config = OnlineQuantizationConfig(resolve_quantization_config(scheme, None)) + set_random_seed(0) + with set_default_torch_dtype(torch.bfloat16), torch.device("cuda"): + layer = ReplicatedLinear( + 1024, 768, bias=False, quant_config=quant_config, prefix="mlp.down_proj" + ) + layer.weight.weight_loader(layer.weight, torch.randn(768, 1024)) + + needle = torch.randn(1, 1024) + filler = 100 * torch.randn(63, 1024) + baseline = layer(needle)[0] + assert torch.equal(layer(torch.cat([needle, filler]))[0][:1], baseline) + assert torch.equal(layer(torch.cat([filler, needle]))[0][-1:], baseline) + + +@requires_fp8 +@torch.inference_mode() +def test_block_fp8_batch_invariant_across_tuned_block_size_k(monkeypatch): + """The block-FP8 kernel applies the block scales once per K tile, so a tuned + BLOCK_SIZE_K below the scale group size reorders the sum.""" + from vllm.model_executor.layers.quantization.utils import fp8_utils + + base = dict( + BLOCK_SIZE_M=64, BLOCK_SIZE_N=128, GROUP_SIZE_M=1, num_warps=4, num_stages=2 + ) + table = {256: {**base, "BLOCK_SIZE_K": 128}, 512: {**base, "BLOCK_SIZE_K": 64}} + monkeypatch.setattr(fp8_utils, "get_w8a8_block_fp8_configs", lambda *_: table) + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True) + N, K, block = 12288, 2048, [128, 128] + + set_random_seed(0) + fp8 = current_platform.fp8_dtype() + A = torch.randn(512, K).to(fp8) + B = torch.randn(N, K).to(fp8) + As = torch.rand(512, K // 128) / 100 + Bs = torch.rand(N // 128, K // 128) / 100 + + def run(m): + return fp8_utils.w8a8_triton_block_scaled_mm( + A[:m], B, As[:m], Bs, block, torch.bfloat16 + )[:256] + + assert torch.equal(run(256), run(512)) + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", False) + assert not torch.equal(run(256), run(512)) + + @_NVFP4_REQUIRES_SM100 @pytest.mark.parametrize("dtype", _NVFP4_LINEAR_DTYPES) @pytest.mark.parametrize("shape", _NVFP4_LINEAR_CONSISTENCY_SHAPES) diff --git a/tests/v1/determinism/test_lora_shrink_batch_invariant.py b/tests/v1/determinism/test_lora_shrink_batch_invariant.py index 2a7029a60324..0c4c595c1d60 100644 --- a/tests/v1/determinism/test_lora_shrink_batch_invariant.py +++ b/tests/v1/determinism/test_lora_shrink_batch_invariant.py @@ -4,7 +4,7 @@ import pytest import torch -from utils import skip_if_not_cuda +from utils import skip_if_not_cuda_alike import vllm.lora.ops.triton_ops.lora_shrink_op as shrink_op import vllm.lora.ops.triton_ops.utils as lora_utils @@ -39,7 +39,7 @@ def _shrink(inputs, weights, mapping, scaling): return out -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize( "hidden_size,rank,nslices,dtype", [ @@ -68,7 +68,7 @@ def test_lora_shrink_matches_reference(hidden_size, rank, nslices, dtype): torch.testing.assert_close(out, ref, rtol=5e-3, atol=5e-3) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("nslices", [1, 3]) def test_lora_shrink_invariant_across_batch_sizes(monkeypatch, nslices): reduce_kernel, reduce_grids = shrink_op._lora_shrink_reduce_kernel, [] diff --git a/tests/v1/determinism/test_matmul_batch_invariant.py b/tests/v1/determinism/test_matmul_batch_invariant.py index af395e74e7f1..b6ea1648f398 100644 --- a/tests/v1/determinism/test_matmul_batch_invariant.py +++ b/tests/v1/determinism/test_matmul_batch_invariant.py @@ -10,7 +10,11 @@ import torch from utils import skip_unsupported -from vllm.model_executor.determinism.batch_invariant import matmul_batch_invariant +import vllm.model_executor.determinism.batch_invariant as bi +from vllm.model_executor.determinism.batch_invariant import ( + enable_batch_invariant_mode, + matmul_batch_invariant, +) from vllm.model_executor.determinism.batch_invariant_configs import ( _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS, _get_tuned_matmul_arch_family, @@ -129,3 +133,49 @@ def test_matmul_batch_invariance_across_tuned_m_buckets(m, transpose_b): batch_output = matmul_batch_invariant(a, b) assert torch.equal(single_output[0], batch_output[0]) + + +@pytest.fixture +def batch_invariant_mode(): + # The overrides are process-global: remove them so later tests compare + # against the stock aten kernels again. + enable_batch_invariant_mode() + yield + assert bi._batch_invariant_LIB is not None + bi._batch_invariant_LIB._destroy() + bi._batch_invariant_MODE = False + + +@skip_unsupported +@pytest.mark.usefixtures("batch_invariant_mode") +@pytest.mark.parametrize("op", ["mm", "addmm", "bmm"]) +def test_compiled_matmul_reaches_batch_invariant_out_overloads(op): + """Inductor lowers mm/addmm/bmm to ``extern_kernels.(..., out=buf)``, + i.e. the ``.out`` overload, which eager never dispatches.""" + if ( + op != "bmm" + and current_platform.is_cuda() + and not current_platform.is_device_capability_family(80) + ): + pytest.skip("mm/addmm are not overridden on SM90+") + torch.manual_seed(0) + dtype, k, n, needle = torch.bfloat16, 1024, 768, 37 + w = torch.randn(4, k, n, device=DEVICE_TYPE, dtype=dtype) + bias = torch.randn(n, device=DEVICE_TYPE, dtype=dtype) + fn = { + "mm": lambda x: torch.mm(x[0], w[0]), + # Bare addmm: a following pointwise op would make Inductor unfuse the bias. + "addmm": lambda x: torch.addmm(bias, x[0], w[0]), + "bmm": lambda x: torch.bmm(x, w), + }[op] + compiled = torch.compile(fn, fullgraph=True) + x = torch.randn(4, 256, k, device=DEVICE_TYPE, dtype=dtype) + row = None + for m in (64, 256, 128): + xs = x[:, :m].clone() + # With dynamic=True, a specialization on M would silently recompile. + torch._dynamo.mark_dynamic(xs, 1) + out = compiled(xs) + assert torch.equal(out, fn(xs)), f"compiled {op} bypassed the BI kernel (M={m})" + row = out[..., needle, :] if row is None else row + assert torch.equal(out[..., needle, :], row) diff --git a/tests/v1/determinism/test_mx_linear_batch_invariant.py b/tests/v1/determinism/test_mx_linear_batch_invariant.py new file mode 100644 index 000000000000..af31302ac91c --- /dev/null +++ b/tests/v1/determinism/test_mx_linear_batch_invariant.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The ROCm native MX linear kernels must not depend on the number of rows.""" + +import pytest +import torch +from utils import requires_mx + +from vllm.utils.torch_utils import set_random_seed + +SEED = 0 + +# E8M0 exponent half-width. With 4-bit operands and a narrow spread the fp32 +# partial sums are exact, so a reordering of them would be invisible. +SCALE_SPREAD = 15 + +# (N, K, use_asm_gemm) +MXFP4_CASES = [ + (4096, 2048, False), + (2048, 6144, False), + (4608, 7168, False), + (8192, 2048, True), + (8192, 1024, True), + (8192, 14336, True), +] + +MXFP8_CASES = [(4096, 2048), (2048, 6144), (1024, 768), (1536, 1024)] + +# The untuned AITER MXFP4 configs switch at M = 9/33/65/129/257/513. +MXFP4_TOKEN_COUNTS = [ + 1, + 2, + 8, + 9, + 16, + 32, + 33, + 64, + 65, + 128, + 129, + 256, + 257, + 512, + 513, + 1024, + 2048, +] + +MXFP8_TOKEN_COUNTS = [1, 32, 64, 65, 128, 129, 256, 257, 512, 1024, 1025, 2048] + +# Row 0 lands in the first tile of every decomposition, so it can stay +# invariant when the rest of the output does not. +CHECK_ROWS = [0, 1, 2, 3, 7, 15, 31] + + +def _probe(out: torch.Tensor) -> torch.Tensor: + """The leading rows of an output, kept so the sweep does not hold every GEMM.""" + return out[: CHECK_ROWS[-1] + 1].clone() + + +def _variant_rows(probes: dict[int, torch.Tensor]) -> list[str]: + """Rows whose contents changed with the number of rows in the launch.""" + failures = [] + for row in CHECK_ROWS: + counts = [n for n in sorted(probes) if n > row] + if not counts: + continue + reference = probes[counts[0]][row] + variant = [n for n in counts if not torch.equal(probes[n][row], reference)] + if variant: + failures.append(f"row {row} changed at row counts {variant}") + return failures + + +def _mxfp4_weights(n: int, k: int, use_asm_gemm: bool): + """Weights laid out as AiterMxfp4LinearKernel.process_weights_after_loading.""" + weight = torch.randint(0, 255, (n, k // 2), dtype=torch.uint8, device="cuda") + weight_scale = torch.randint( + 127 - SCALE_SPREAD, + 128 + SCALE_SPREAD, + (n, k // 32), + dtype=torch.uint8, + device="cuda", + ) + if not use_asm_gemm: + return weight, weight_scale.T.contiguous() + + from aiter.ops.shuffle import shuffle_weight + + sm, sn = weight_scale.shape + weight_scale = ( + weight_scale.view(sm // 32, 2, 16, sn // 8, 2, 4, 1) + .permute(0, 3, 5, 2, 4, 1, 6) + .contiguous() + .view(sm, sn) + ) + return shuffle_weight(weight, layout=(16, 16)), weight_scale + + +def _mxfp8_weights(n: int, k: int): + """E4M3 weights with E8M0 block scales, as the MXFP8 layer stores them.""" + weight = (torch.randn(n, k, device="cuda") / 8).to(torch.float8_e4m3fn) + weight_scale = torch.randint( + 127 - SCALE_SPREAD, + 128 + SCALE_SPREAD, + (n, k // 32), + dtype=torch.uint8, + device="cuda", + ) + return weight, weight_scale + + +@requires_mx +@pytest.mark.parametrize("n,k,use_asm_gemm", MXFP4_CASES) +def test_mxfp4_linear_is_batch_invariant(n: int, k: int, use_asm_gemm: bool): + pytest.importorskip("aiter") + + # Importing the module registers torch.ops.vllm.gemm_with_dynamic_quant. + import vllm.model_executor.kernels.linear.mxfp4.aiter # noqa: F401 + + set_random_seed(SEED) + weight, weight_scale = _mxfp4_weights(n, k, use_asm_gemm) + x = torch.randn(max(MXFP4_TOKEN_COUNTS), k, device="cuda", dtype=torch.bfloat16) + + probes = { + num_tokens: _probe( + torch.ops.vllm.gemm_with_dynamic_quant( + x[:num_tokens], weight, weight_scale, use_asm_gemm, torch.bfloat16 + ) + ) + for num_tokens in MXFP4_TOKEN_COUNTS + } + + failures = _variant_rows(probes) + assert not failures, ( + f"MXFP4 linear depends on the row count " + f"(N={n}, K={k}, asm={use_asm_gemm}):\n " + "\n ".join(failures) + ) + + +@requires_mx +@pytest.mark.parametrize("n,k", MXFP8_CASES) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_mxfp8_linear_is_batch_invariant(n: int, k: int, dtype: torch.dtype): + from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( + _mxfp8_dot_scaled_linear, + ) + + set_random_seed(SEED) + weight, weight_scale = _mxfp8_weights(n, k) + x = torch.randn(max(MXFP8_TOKEN_COUNTS), k, device="cuda", dtype=dtype) + + probes = { + num_tokens: _probe( + _mxfp8_dot_scaled_linear(x[:num_tokens], weight, weight_scale) + ) + for num_tokens in MXFP8_TOKEN_COUNTS + } + + failures = _variant_rows(probes) + assert not failures, ( + f"MXFP8 linear depends on the row count (N={n}, K={k}, {dtype}):\n " + + "\n ".join(failures) + ) diff --git a/tests/v1/determinism/test_nvfp4_batch_invariant.py b/tests/v1/determinism/test_nvfp4_batch_invariant.py index 1e9b0cc8fc56..162e6b36c3da 100644 --- a/tests/v1/determinism/test_nvfp4_batch_invariant.py +++ b/tests/v1/determinism/test_nvfp4_batch_invariant.py @@ -9,10 +9,14 @@ from utils import ( _extract_step_logprobs, _random_prompt, - skip_if_not_cuda, + skip_if_not_cuda_alike, ) from vllm import LLM, SamplingParams +from vllm.platforms import current_platform + +# FLASH_ATTN is CUDA-only. +E2E_BACKENDS = ["TRITON_ATTN"] if current_platform.is_rocm() else ["FLASH_ATTN"] pytestmark = pytest.mark.skipif( not hasattr(torch, "float8_e4m3fn"), @@ -40,8 +44,8 @@ def _make_llm(max_num_seqs: int, backend: str) -> LLM: ) -@skip_if_not_cuda -@pytest.mark.parametrize("backend", ["FLASH_ATTN"]) +@skip_if_not_cuda_alike +@pytest.mark.parametrize("backend", E2E_BACKENDS) def test_dense_nvfp4_generation_is_deterministic_across_batch_sizes_e2e(backend): seed = int(os.getenv("VLLM_TEST_SEED", "12345")) random.seed(seed) diff --git a/tests/v1/determinism/test_online_batch_invariance.py b/tests/v1/determinism/test_online_batch_invariance.py index dcba87dc4927..fb5db254ba43 100644 --- a/tests/v1/determinism/test_online_batch_invariance.py +++ b/tests/v1/determinism/test_online_batch_invariance.py @@ -6,6 +6,7 @@ Environment variables: - VLLM_TEST_MODEL: served model name (e.g., Qwen/Qwen3-1.7B / DeepSeek-R1) - VLLM_TP_SIZE: tensor parallelism size (e.g., 4) + - VLLM_TEST_SERVER_ARGS: extra server args (e.g., "--data-parallel-size 4") """ @@ -16,7 +17,7 @@ import openai import pytest -from utils import BACKENDS, TEST_MODEL, _random_prompt, skip_if_not_cuda +from utils import BACKENDS, TEST_MODEL, _random_prompt, skip_if_not_cuda_alike from tests.utils import RemoteOpenAIServer @@ -132,7 +133,7 @@ def _compare_bs1_vs_bsn_single_process( ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("backend", BACKENDS) def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( backend: str, @@ -156,6 +157,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( ] if tp_size: server_args += ["-tp", tp_size] + server_args += os.getenv("VLLM_TEST_SERVER_ARGS", "").split() with RemoteOpenAIServer(TEST_MODEL, server_args) as server: client = server.get_client() diff --git a/tests/v1/determinism/test_rms_norm_batch_invariant.py b/tests/v1/determinism/test_rms_norm_batch_invariant.py index 1feb5e67413d..f26217bd7410 100644 --- a/tests/v1/determinism/test_rms_norm_batch_invariant.py +++ b/tests/v1/determinism/test_rms_norm_batch_invariant.py @@ -4,7 +4,7 @@ import pytest import torch -from utils import skip_if_not_cuda, skip_unsupported +from utils import skip_if_not_cuda_alike, skip_unsupported from vllm.model_executor.determinism.batch_invariant import ( rms_norm_batch_invariant, @@ -27,7 +27,7 @@ def _rms_norm_reference( return (output * weight.float()).to(input_tensor.dtype) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("batch_size", [1, 4, 64, 300]) @pytest.mark.parametrize("hidden_size", [512, 2048, 4096, 8192]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @@ -76,7 +76,7 @@ def test_rms_norm_batch_invariant_vs_reference( ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("hidden_size", [512, 4096]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("eps", [1e-6]) @@ -189,7 +189,7 @@ def _assert_rows_bit_identical(small, large, msg): torch.testing.assert_close(small, large, rtol=0.0, atol=0.0, msg=msg) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("hidden_size", [512, 4096]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("seed", list(range(4))) @@ -221,7 +221,7 @@ def rms_norm(x): ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("hidden_size", [512, 4096]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("seed", list(range(4))) @@ -267,7 +267,7 @@ def quant(x, res): ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("hidden_size", [512, 4096]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("seed", list(range(4))) @@ -303,7 +303,7 @@ def per_block_quant(x): ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("batch_size", [1, 16, 128]) @pytest.mark.parametrize("seq_len", [1, 32, 512]) @pytest.mark.parametrize("hidden_size", [2048, 4096]) @@ -343,7 +343,7 @@ def test_rms_norm_3d_input( ) -@skip_if_not_cuda +@skip_if_not_cuda_alike def test_rms_norm_numerical_stability(default_vllm_config): """Test RMS norm numerical stability with extreme values. @@ -431,7 +431,7 @@ def test_rms_norm_formula(default_vllm_config): ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("hidden_size", [128, 1024, 4096, 16384]) def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int): """Test RMS norm with various hidden sizes to ensure block size handling. diff --git a/tests/v1/determinism/utils.py b/tests/v1/determinism/utils.py index f03ea05b4331..405b5ac3d8a5 100644 --- a/tests/v1/determinism/utils.py +++ b/tests/v1/determinism/utils.py @@ -34,12 +34,16 @@ class DeviceConfig(NamedTuple): available=current_platform.is_xpu() and HAS_TRITON, backends=["TRITON_ATTN"], ), + "rocm": DeviceConfig( + available=current_platform.is_rocm() and HAS_TRITON, + backends=["TRITON_ATTN"], + ), } DEFAULT_MODEL = "Qwen/Qwen3-1.7B" TEST_MODEL = os.getenv("VLLM_TEST_MODEL", DEFAULT_MODEL) -# Override backends for MLA models (MLA only supported on CUDA). +# Override backends for MLA models. if os.getenv("VLLM_TEST_MODEL"): config = get_config(TEST_MODEL, trust_remote_code=False) if ModelArchConfigConvertorBase(config, config.get_text_config()).is_deepseek_mla(): @@ -52,6 +56,10 @@ class DeviceConfig(NamedTuple): available=DEVICE_BACKENDS["xpu"].available, backends=[], ) + DEVICE_BACKENDS["rocm"] = DeviceConfig( + available=DEVICE_BACKENDS["rocm"].available, + backends=["TRITON_MLA"], + ) # Only include backends for devices that are actually available. BACKENDS: list[str] = sorted( @@ -60,7 +68,7 @@ class DeviceConfig(NamedTuple): skip_unsupported = pytest.mark.skipif( not any(cfg.available for cfg in DEVICE_BACKENDS.values()), - reason="Requires CUDA >= Ampere (SM80) or Intel XPU with Triton", + reason="Requires CUDA >= Ampere (SM80), ROCm, or Intel XPU with Triton", ) skip_if_not_cuda = pytest.mark.skipif( @@ -68,6 +76,16 @@ class DeviceConfig(NamedTuple): reason="Requires CUDA >= Ampere (SM80)", ) +skip_if_not_cuda_alike = pytest.mark.skipif( + not (DEVICE_BACKENDS["cuda"].available or DEVICE_BACKENDS["rocm"].available), + reason="Requires CUDA >= Ampere (SM80) or ROCm", +) + +requires_mx = pytest.mark.skipif( + not (current_platform.is_rocm() and current_platform.supports_mx()), + reason="requires a ROCm device with native MX support (gfx95x)", +) + def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str: # Generate more realistic prompts that will actually produce varied tokens diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 29d5bae89d2c..1bc6182288ab 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -413,6 +413,17 @@ def all_reduce(self, input_): out = ca_comm.custom_all_reduce(input_) assert out is not None return out + if ( + envs.VLLM_BATCH_INVARIANT + and current_platform.is_rocm() + and input_.is_floating_point() + ): + # Integer sums are exact in any order and stay on RCCL. + from vllm.model_executor.determinism.batch_invariant import ( + all_reduce_batch_invariant, + ) + + return all_reduce_batch_invariant(input_, self.device_group) symm_mem_comm = self.symm_mem_comm if symm_mem_comm is not None and symm_mem_comm.should_use_symm_mem(input_): out = symm_mem_comm.all_reduce(input_) @@ -498,7 +509,17 @@ def reduce_scatter(self, input_: torch.Tensor, dim: int = -1): chunk_size = input_tensor.shape[0] // world_size output_shape = (chunk_size,) + input_tensor.shape[1:] - if should_nccl_symm_mem_ag_rs(): + if ( + envs.VLLM_BATCH_INVARIANT + and current_platform.is_rocm() + and input_.is_floating_point() + ): + from vllm.model_executor.determinism.batch_invariant import ( + reduce_scatter_batch_invariant, + ) + + output = reduce_scatter_batch_invariant(input_tensor, self.device_group) + elif should_nccl_symm_mem_ag_rs(): output = self._reduce_scatter_symm_mem(input_tensor) else: output = torch.empty( @@ -563,6 +584,22 @@ def _reduce_scatterv( assert output.device == input_tensor.device assert output.is_contiguous() + if ( + envs.VLLM_BATCH_INVARIANT + and current_platform.is_rocm() + and input_.is_floating_point() + ): + from vllm.model_executor.determinism.batch_invariant import ( + reduce_scatter_batch_invariant, + ) + + result = reduce_scatter_batch_invariant( + input_tensor, self.device_group, sizes=sizes + ) + if output is not None: + return output.copy_(result) + return result.movedim(0, dim).contiguous() + if self._can_use_aiter_ag_rs(sizes): aiter_comm = self.aiter_ar_comm assert aiter_comm is not None diff --git a/vllm/model_executor/determinism/batch_invariant.py b/vllm/model_executor/determinism/batch_invariant.py index 403f47570204..e4b04793feed 100644 --- a/vllm/model_executor/determinism/batch_invariant.py +++ b/vllm/model_executor/determinism/batch_invariant.py @@ -6,6 +6,7 @@ from typing import Any import torch +import torch.distributed as dist import vllm.envs as envs from vllm.model_executor.determinism.batch_invariant_configs import ( @@ -59,6 +60,7 @@ def matmul_kernel_persistent( stride_bn, stride_cm, stride_cn, + stride_bias, BLOCK_SIZE_M: tl.constexpr, # BLOCK_SIZE_N: tl.constexpr, # BLOCK_SIZE_K: tl.constexpr, # @@ -128,7 +130,7 @@ def matmul_kernel_persistent( c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) if HAS_BIAS: - bias_ptrs = bias_ptr + offs_cn + bias_ptrs = bias_ptr + offs_cn * stride_bias bias = tl.load(bias_ptrs, mask=offs_cn < N, other=0.0).to(tl.float32) accumulator += bias c = accumulator.to(c_ptr.dtype.element_ty) @@ -305,6 +307,8 @@ def grid(META): ), ) + if not _BLOCK_SIZES_RESOLVED: + _resolve_block_sizes() configs = { torch.bfloat16: { "BLOCK_SIZE_M": 128, @@ -345,6 +349,7 @@ def grid(META): b.stride(1), # c.stride(0), c.stride(1), # + bias.stride(0) if bias is not None else 1, NUM_SMS=NUM_SMS, # A_LARGE=a.numel() > 2**31, B_LARGE=b.numel() > 2**31, @@ -489,16 +494,19 @@ def bmm_kernel( @triton.jit -def _log_softmax_kernel( +def _softmax_kernel( input_ptr, output_ptr, input_row_stride, + input_col_stride, output_row_stride, n_cols, BLOCK_SIZE: tl.constexpr, + LOG: tl.constexpr = False, ): - """Compute log_softmax along the last dimension of a 2D tensor. - Each block handles one row of the input tensor. + """Compute softmax, or log_softmax when ``LOG``, along the last dimension of a + 2D tensor. Each block handles one row of the input tensor, in a fixed block + order, so the result never depends on the row count. """ # Get the row index for this block row_idx = tl.program_id(0).to(tl.int64) @@ -514,7 +522,9 @@ def _log_softmax_kernel( mask = col_idx < n_cols # Load values - vals = tl.load(row_start_ptr + col_idx, mask=mask, other=-float("inf")) + vals = tl.load( + row_start_ptr + col_idx * input_col_stride, mask=mask, other=-float("inf") + ) # Update maximum max_val = tl.max(tl.maximum(vals, max_val)) @@ -526,25 +536,26 @@ def _log_softmax_kernel( mask = col_idx < n_cols # Load values - vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) + vals = tl.load(row_start_ptr + col_idx * input_col_stride, mask=mask, other=0.0) # Compute exp(x - max_val) and accumulate exp_vals = tl.exp(vals - max_val) sum_exp += tl.sum(tl.where(mask, exp_vals, 0.0)) - # Compute log(sum_exp) - log_sum_exp = tl.log(sum_exp) + log_sum_exp = tl.log(sum_exp) if LOG else 0.0 - # Step 3: Compute final log_softmax values: x - max_val - log_sum_exp + # Step 3: normalise for col_offset in range(0, n_cols, BLOCK_SIZE): col_idx = col_offset + tl.arange(0, BLOCK_SIZE) mask = col_idx < n_cols # Load values - vals = tl.load(row_start_ptr + col_idx, mask=mask) + vals = tl.load(row_start_ptr + col_idx * input_col_stride, mask=mask) - # Compute log_softmax - output = vals - max_val - log_sum_exp + if LOG: + output = vals - max_val - log_sum_exp + else: + output = tl.exp(vals - max_val) / sum_exp # Store results tl.store(output_row_start_ptr + col_idx, output, mask=mask) @@ -570,30 +581,55 @@ def log_softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor: # Flatten all dimensions except the last one original_shape = input.shape input_2d = input.reshape(-1, input.shape[-1]) - input_2d = input_2d.contiguous() n_rows, n_cols = input_2d.shape # Allocate output tensor - output = torch.empty_like(input_2d) + output = torch.empty((n_rows, n_cols), dtype=input_2d.dtype, device=input_2d.device) # Choose block size based on the number of columns BLOCK_SIZE = 1024 # Launch kernel with one block per row grid = (n_rows,) - _log_softmax_kernel[grid]( + _softmax_kernel[grid]( input_2d, output, input_2d.stride(0), + input_2d.stride(1), output.stride(0), n_cols, BLOCK_SIZE=BLOCK_SIZE, + LOG=True, ) # Reshape output back to original shape return output.reshape(original_shape) +def softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor: + """Compute softmax along the last dimension using a Triton kernel.""" + if dim != -1 and dim != input.ndim - 1: + raise ValueError( + "This implementation only supports softmax along the last dimension" + ) + + original_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]) + n_rows, n_cols = input_2d.shape + output = torch.empty((n_rows, n_cols), dtype=input_2d.dtype, device=input_2d.device) + + _softmax_kernel[(n_rows,)]( + input_2d, + output, + input_2d.stride(0), + input_2d.stride(1), + output.stride(0), + n_cols, + BLOCK_SIZE=1024, + ) + return output.reshape(original_shape) + + @triton.jit def mean_kernel( input_ptr, @@ -735,6 +771,9 @@ def mean_dim( def mm_batch_invariant(a, b, bias=None): if a.device.type == "xpu": + # The descriptor kernel reads the bias at unit stride. + if bias is not None: + bias = bias.contiguous() return matmul_descriptor_persistent(a, b, bias=bias) return matmul_persistent(a, b, bias=bias) @@ -820,6 +859,8 @@ def bmm_batch_invariant(a, b, *, out=None): assert out.dtype == dtype and out.device == a.device, "out tensor mismatch" c = out + if not _BLOCK_SIZES_RESOLVED: + _resolve_block_sizes() configs = { torch.bfloat16: { "BLOCK_SIZE_M": 128, @@ -877,8 +918,50 @@ def bmm_batch_invariant(a, b, *, out=None): return c -def addmm_batch_invariant(bias, a, b): - return mm_batch_invariant(a, b, bias=bias) +def _addmm_fused_bias(bias, beta, alpha, n): + # The kernel adds bias into its fp32 accumulator indexed by output column, + # so it can only absorb a row vector of width n, unscaled by alpha. + if beta == 0 or alpha != 1: + return None + if bias.dim() != 1 or bias.shape[0] != n: + return None + # fp32 keeps the scaling exact; the kernel casts the bias to fp32 anyway. + return bias if beta == 1 else bias.float() * beta + + +def _addmm_impl(bias, a, b, beta, alpha, out): + fused_bias = _addmm_fused_bias(bias, beta, alpha, b.shape[1]) + result = mm_batch_invariant(a, b, bias=fused_bias) + if fused_bias is None: + if alpha != 1: + result = result * alpha + if beta != 0: + result = result + (bias if beta == 1 else beta * bias) + if out is None: + return result + return out.copy_(result) + + +def addmm_batch_invariant(bias, a, b, *, beta=1, alpha=1): + return _addmm_impl(bias, a, b, beta, alpha, None) + + +# Inductor lowers mm/addmm/bmm to ``extern_kernels.(..., out=buf)``, which +# dispatches the ``.out`` overload rather than the default one. Registering only +# the default leaves the compiled path on the vendor GEMM, so a compiled +# ``torch.addmm`` is batch variant while the eager one is not. + + +def mm_out_batch_invariant(a, b, *, out): + return matmul_batch_invariant(a, b, out=out) + + +def addmm_out_batch_invariant(bias, a, b, *, beta=1, alpha=1, out): + return _addmm_impl(bias, a, b, beta, alpha, out) + + +def bmm_out_batch_invariant(a, b, *, out): + return bmm_batch_invariant(a, b, out=out) def _log_softmax_batch_invariant(input, dim, _half_to_float): @@ -887,10 +970,16 @@ def _log_softmax_batch_invariant(input, dim, _half_to_float): return log_softmax(input, dim=dim) -def softmax_batch_invariant(input, dim, dtype=None): - # Compute softmax in a deterministic way - # First subtract max for numerical stability (standard practice) +def softmax_batch_invariant(input, dim, half_to_float=False): + if half_to_float: + input = input.float() + if dim == -1 or dim == input.ndim - 1: + return softmax(input, dim=-1) + + # Off the last dimension this falls back to torch ops, whose reduction split + # depends on the tensor shape. input_max = torch.amax(input, dim=dim, keepdim=True) + # First subtract max for numerical stability (standard practice) input = input - input_max exp_x = torch.exp(input) sum_exp_x = torch.sum(exp_x, dim=dim, keepdim=True) @@ -1035,11 +1124,150 @@ def rms_norm_batch_invariant( def linear_batch_invariant(input, weight, bias=None): - output = matmul_batch_invariant(input, weight.t()) - if bias is not None: - output = output + bias - return output + bias = bias.expand(weight.shape[0]) + output = mm_batch_invariant( + input.reshape(-1, input.shape[-1]), weight.t(), bias=bias + ) + return output.view(*input.shape[:-1], output.shape[-1]) + + +@triton.jit +def _fixed_order_sum_kernel( + src_ptr, # (WORLD_SIZE, numel) gathered contributions, rank-major + out_ptr, # (numel,) + numel, + stride_rank, + WORLD_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + LARGE: tl.constexpr, +): + """Sum ``WORLD_SIZE`` contributions in ascending rank order, fp32 accumulator. + + An output element sums exactly ``WORLD_SIZE`` values in a compile-time + constant order and rounds once, so the result depends only on the rank + ordering and never on how many elements are being reduced. + """ + pid = tl.program_id(0) + if LARGE: + pid = pid.to(tl.int64) + stride_rank = stride_rank.to(tl.int64) + offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offs < numel + + acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for r in tl.static_range(WORLD_SIZE): + vals = tl.load(src_ptr + r * stride_rank + offs, mask=mask, other=0.0) + acc += vals.to(tl.float32) + + tl.store(out_ptr + offs, acc.to(out_ptr.dtype.element_ty), mask=mask) + + +def all_reduce_batch_invariant( + input_: torch.Tensor, group: "torch.distributed.ProcessGroup | None" = None +) -> torch.Tensor: + """Sum all-reduce whose result does not depend on the message size. + + Library all-reduces (NCCL/RCCL) pick their algorithm, channel count and chunk + boundaries from the message size, so the order in which a given element's + contributions are summed changes with the number of tokens in the batch. + + Instead, all-gather the contributions -- pure data movement, so bitwise + reproducible at any size -- and reduce them with ``_fixed_order_sum_kernel``. + """ + world_size = dist.get_world_size(group) + if world_size == 1: + return input_ + + x = input_.contiguous() + numel = x.numel() + gathered = torch.empty((world_size, numel), dtype=x.dtype, device=x.device) + dist.all_gather_into_tensor(gathered, x.view(-1), group=group) + + out = torch.empty_like(x) + BLOCK_SIZE = 1024 + grid = (triton.cdiv(numel, BLOCK_SIZE),) + _fixed_order_sum_kernel[grid]( + gathered, + out.view(-1), + numel, + numel, + WORLD_SIZE=world_size, + BLOCK_SIZE=BLOCK_SIZE, + LARGE=gathered.numel() > 2**31, + ) + return out + + +def reduce_scatter_batch_invariant( + input_: torch.Tensor, + group: "torch.distributed.ProcessGroup | None" = None, + sizes: list[int] | None = None, +) -> torch.Tensor: + """Sum reduce-scatter over dim 0 whose result does not depend on the size. + + ``ncclReduceScatter`` picks its chunking from the message size just as + ``ncclAllReduce`` does. Route around it the same way + ``all_reduce_batch_invariant`` does: an all-to-all sends rank ``d`` exactly + the rows it owns from every rank -- pure data movement, so bitwise + reproducible however the library chunks it -- and lands them rank-major, + which is the layout ``_fixed_order_sum_kernel`` already consumes. + + ``sizes`` (the reduce-scatterv case) changes only *which* rows a rank + receives, never how the received contributions are summed, so variable + shard sizes are invariant for the same reason. + + Args: + input_: Contiguous full-size input; dim 0 is the scattered axis. + group: Process group, defaults to the world group. + sizes: Per-rank row counts. Uniform split when ``None``. + + Returns: + This rank's shard of the sum. + + """ + world_size = dist.get_world_size(group) + if world_size == 1: + return input_ + + x = input_.contiguous() + if sizes is None: + assert x.shape[0] % world_size == 0 + sizes = [x.shape[0] // world_size] * world_size + else: + assert len(sizes) == world_size + assert x.shape[0] == sum(sizes) + + rows = sizes[dist.get_rank(group)] + out = torch.empty((rows, *x.shape[1:]), dtype=x.dtype, device=x.device) + recv = torch.empty( + (world_size * rows, *x.shape[1:]), dtype=x.dtype, device=x.device + ) + # Every rank takes part even when its own shard is empty: it still has rows + # to send to the others. + dist.all_to_all_single( + recv, + x, + output_split_sizes=[rows] * world_size, + input_split_sizes=sizes, + group=group, + ) + + numel = out.numel() + if numel == 0: + return out + BLOCK_SIZE = 1024 + grid = (triton.cdiv(numel, BLOCK_SIZE),) + _fixed_order_sum_kernel[grid]( + recv.view(-1), + out.view(-1), + numel, + numel, + WORLD_SIZE=world_size, + BLOCK_SIZE=BLOCK_SIZE, + LARGE=recv.numel() > 2**31, + ) + return out _batch_invariant_MODE = False @@ -1047,9 +1275,33 @@ def linear_batch_invariant(input, weight, bias=None): _fp16_block_size_n = 256 _fp32_block_size_n = 128 _fp32_num_stages = 3 +_BLOCK_SIZES_RESOLVED = False _NUM_SMS: int = 0 +def _resolve_block_sizes(): + """Size the tiles to the shared memory, to avoid triton OutOfResources. + + Runs when the mode is enabled, or on first use by a caller that never + enables it, such as a kernel unit test. + """ + global _fp16_block_size_n, _fp32_block_size_n, _fp32_num_stages + global _BLOCK_SIZES_RESOLVED + _BLOCK_SIZES_RESOLVED = True + if not current_platform.is_cuda_alike(): + return + if get_max_shared_memory_bytes() > 106496: + _fp16_block_size_n = 256 + _fp32_block_size_n = 128 + _fp32_num_stages = 3 + else: + _fp16_block_size_n = 128 + # SM89 N=1 fp32 128/stages=3 compiles to 131072B (>101376). + # Only N=1 OORs; keep 128/stages=3 for every other N. + _fp32_block_size_n = 32 + _fp32_num_stages = 2 + + def enable_batch_invariant_mode(): global _batch_invariant_MODE, _batch_invariant_LIB global _fp16_block_size_n, _fp32_block_size_n, _fp32_num_stages @@ -1062,12 +1314,17 @@ def enable_batch_invariant_mode(): _batch_invariant_LIB = torch.library.Library("aten", "IMPL") key = current_platform.dispatch_key - if current_platform.is_cuda(): - if current_platform.is_device_capability_family(80): - # SM80 (Ampere) cannot rely on cuBLASLt-only determinism; install the - # triton persistent matmul overrides for mm/addmm/matmul/linear. + if current_platform.is_cuda_alike(): + if current_platform.is_rocm() or current_platform.is_device_capability_family( + 80 + ): + # Neither cuBLASLt on SM80 (Ampere) nor hipBLASLt can be made batch + # invariant; install the triton persistent matmul overrides for + # mm/addmm/matmul/linear. _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, key) _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, key) + _batch_invariant_LIB.impl("aten::mm.out", mm_out_batch_invariant, key) + _batch_invariant_LIB.impl("aten::addmm.out", addmm_out_batch_invariant, key) _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, key) _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, key) else: @@ -1077,18 +1334,7 @@ def enable_batch_invariant_mode(): os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1" - # Query the shared memory size and set block size - # accordingly to avoid triton OutOfResources. - if get_max_shared_memory_bytes() > 106496: - _fp16_block_size_n = 256 - _fp32_block_size_n = 128 - _fp32_num_stages = 3 - else: - _fp16_block_size_n = 128 - # SM89 N=1 fp32 128/stages=3 compiles to 131072B (>101376). - # Only N=1 OORs; keep 128/stages=3 for every other N. - _fp32_block_size_n = 32 - _fp32_num_stages = 2 + _resolve_block_sizes() elif current_platform.is_xpu(): # Tensor descriptors need a global-memory allocator; must be set outside # torch.compile regions (triton.set_allocator modifies global state). @@ -1100,6 +1346,8 @@ def _triton_alloc_fn(size: int, alignment: int, stream: int | None): _NUM_SMS = num_compute_units(0) _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, key) _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, key) + _batch_invariant_LIB.impl("aten::mm.out", mm_out_batch_invariant, key) + _batch_invariant_LIB.impl("aten::addmm.out", addmm_out_batch_invariant, key) _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, key) _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, key) @@ -1123,7 +1371,9 @@ def _triton_alloc_fn(size: int, alignment: int, stream: int | None): _batch_invariant_LIB.impl( "aten::bmm", bmm_batch_invariant, key, allow_override=True ) - torch.bmm = bmm_batch_invariant + _batch_invariant_LIB.impl( + "aten::bmm.out", bmm_out_batch_invariant, key, allow_override=True + ) reduced_precision_val = ( (False, False) if is_torch_equal_or_newer("2.10.0") else False @@ -1145,28 +1395,32 @@ def override_envs_for_invariance(): os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" - # NCCL determinism settings. NCCL keeps the launch mode process-wide, so, - # unlike the pins below, it also applies to communicators shared with other - # processes. - os.environ["NCCL_LAUNCH_MODE"] = "GROUP" - pin_nccl_env( - { - "NCCL_COLLNET_ENABLE": "0", - "NCCL_NVLS_ENABLE": "0", - "NCCL_P2P_NET_DISABLE": "1", - "NCCL_MIN_NCHANNELS": "1", - "NCCL_MAX_NCHANNELS": "1", - "NCCL_PROTO": "Simple", - # NCCL >= 2.31 zero-fills the algorithm table of every collective - # when NCCL_ALGO is set and re-enables only the named ones; together - # with the NCCL_PROTO above, collectives not named here end up with - # no algorithm and fail with ncclInvalidUsage. Re-enable Ring and - # Tree for all collectives, then pin AllReduce to Tree. - "NCCL_ALGO": "ring,tree;allreduce:tree", - "NCCL_NTHREADS": "1", - "NCCL_SOCKET_NTHREADS": "1", - } - ) + if current_platform.is_rocm(): + # The skinny GEMMs pick their split of K from the batch size. + os.environ["VLLM_ROCM_USE_SKINNY_GEMM"] = "0" + else: + # NCCL determinism settings. NCCL keeps the launch mode process-wide, so, + # unlike the pins below, it also applies to communicators shared with other + # processes. + os.environ["NCCL_LAUNCH_MODE"] = "GROUP" + pin_nccl_env( + { + "NCCL_COLLNET_ENABLE": "0", + "NCCL_NVLS_ENABLE": "0", + "NCCL_P2P_NET_DISABLE": "1", + "NCCL_MIN_NCHANNELS": "1", + "NCCL_MAX_NCHANNELS": "1", + "NCCL_PROTO": "Simple", + # NCCL >= 2.31 zero-fills the algorithm table of every collective + # when NCCL_ALGO is set and re-enables only the named ones; together + # with the NCCL_PROTO above, collectives not named here end up with + # no algorithm and fail with ncclInvalidUsage. Re-enable Ring and + # Tree for all collectives, then pin AllReduce to Tree. + "NCCL_ALGO": "ring,tree;allreduce:tree", + "NCCL_NTHREADS": "1", + "NCCL_SOCKET_NTHREADS": "1", + } + ) # torch.compile settings os.environ["VLLM_USE_AOT_COMPILE"] = "0" diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 484143ffc14a..0ce22f477ac7 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -218,6 +218,7 @@ ) from vllm.model_executor.kernels.linear.scaled_mm.triton import ( TritonFp8BlockScaledMMKernel, + TritonFP8ScaledMMLinearKernel, TritonInt8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.xpu import ( @@ -253,6 +254,20 @@ def _get_linear_backend(*, quantization: str) -> str: return config.kernel_config.linear_backend +def _batch_invariant_fp8_kernel(kernel, config, quantization: str): + can_implement, reason = is_supported_and_can_implement_kernel(kernel, config, None) + if not can_implement: + raise ValueError(f"Batch invariance on ROCm: {reason}") + linear_backend = _get_linear_backend(quantization=quantization) + if linear_backend not in ("auto", "triton"): + logger.warning_once( + "VLLM_BATCH_INVARIANT overrides --linear-backend=%s; " + "using the Triton backend for deterministic execution.", + linear_backend, + ) + return kernel + + # Kernel classes covered by each --linear-backend value. _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = { "b12x": { @@ -723,6 +738,13 @@ def init_fp8_linear_kernel( ) if activation_quant_key.scale.group_shape.is_per_group(): + if envs.VLLM_BATCH_INVARIANT and current_platform.is_rocm(): + force_kernel = _batch_invariant_fp8_kernel( + TritonFp8BlockScaledMMKernel, + scaled_mm_linear_kernel_config, + "fp8_block_w8a8", + ) + kernel_type = choose_scaled_mm_linear_kernel( config=scaled_mm_linear_kernel_config, possible_kernels=_POSSIBLE_FP8_BLOCK_KERNELS, # type: ignore[misc] @@ -755,6 +777,13 @@ def init_fp8_linear_kernel( ) else: + if envs.VLLM_BATCH_INVARIANT and current_platform.is_rocm(): + force_kernel = _batch_invariant_fp8_kernel( + TritonFP8ScaledMMLinearKernel, + scaled_mm_linear_kernel_config, + "fp8_w8a8", + ) + kernel_type = choose_scaled_mm_linear_kernel( config=scaled_mm_linear_kernel_config, possible_kernels=_POSSIBLE_FP8_KERNELS, # type: ignore[arg-type] @@ -780,6 +809,19 @@ def init_fp8_linear_kernel( ) +def fp8_scaled_mm_is_batch_invariant( + kernel: FP8ScaledMMLinearKernel | Fp8BlockScaledMMLinearKernel, +) -> bool: + """Whether a row's output from the kernel is independent of the other rows.""" + if not isinstance( + kernel, (CutlassFP8ScaledMMLinearKernel, TritonFP8ScaledMMLinearKernel) + ): + return False + # A dynamic per-tensor activation scale is an amax over every row. + scale = kernel.config.activation_quant_key.scale + return scale.static or not scale.group_shape.is_per_tensor() + + def init_int8_linear_kernel( is_channelwise: bool, is_static_input_scheme: bool, @@ -1250,6 +1292,7 @@ def register_linear_kernel( __all__ = [ + "fp8_scaled_mm_is_batch_invariant", "init_fp8_linear_kernel", "init_int8_linear_kernel", "init_nvfp4_linear_kernel", diff --git a/vllm/model_executor/kernels/linear/mxfp4/aiter.py b/vllm/model_executor/kernels/linear/mxfp4/aiter.py index 420e7f9493e1..4cd09c1f905f 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/aiter.py +++ b/vllm/model_executor/kernels/linear/mxfp4/aiter.py @@ -62,7 +62,12 @@ def gemm_with_dynamic_quant( N = weight.shape[0] K = weight.shape[1] if rocm_use_aiter_fp4_asm_gemm: - if M <= 64 and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K): + # The preshuffle kernel's tuned configs split K in some M buckets. + if ( + M <= 64 + and not envs.VLLM_BATCH_INVARIANT + and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K) + ): if x_scales is None: # use hip quant kernel for performance if M >= 32: @@ -116,7 +121,16 @@ def gemm_with_dynamic_quant( x_q.shape[0], weight.shape[0], device=x_q.device, dtype=out_dtype ) - gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y) + config = None + if envs.VLLM_BATCH_INVARIANT: + # The configs split K in some M buckets (up to M <= 256 for + # tuned shapes). Keep the tile, drop the split. + from aiter.ops.triton.gemm_afp4wfp4 import _get_config + + config, _ = _get_config(M, N, K) + config = dict(config, NUM_KSPLIT=1) + + gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y, config) return y def gemm_with_dynamic_quant_fake( diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py index faf1e180fa2f..98920f2df13a 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py @@ -12,6 +12,7 @@ import torch from torch.nn.parameter import Parameter +import vllm.envs as envs from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( MXFP8_BLOCK_SIZE, MXFP8_SCALE_DTYPE, @@ -246,7 +247,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: N, K = weight.shape scale_k = K // MXFP8_BLOCK_SIZE weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous() - block_scale = _as_block32_scale(weight_scale) + # rocm_mxfp8_block32_gemm's split-K varies with M. + block_scale = ( + None if envs.VLLM_BATCH_INVARIANT else _as_block32_scale(weight_scale) + ) if block_scale is not None: # Checkpoints with 32x32 scale blocks (DeepSeek V4/V4.1) get them # back from the loader expanded per row; keep the blocks and use diff --git a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py index 75db4bcaa9ed..ac096b47b96a 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py @@ -39,6 +39,7 @@ ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.triton import ( + TritonFP8ScaledMMLinearKernel, TritonInt8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.xpu import ( @@ -65,6 +66,7 @@ "PerTensorTorchFP8ScaledMMLinearKernel", "RowWiseTorchFP8ScaledMMLinearKernel", "ROCmFP8ScaledMMLinearKernel", + "TritonFP8ScaledMMLinearKernel", "TritonInt8ScaledMMLinearKernel", "ZentorchInt8ScaledMMLinearKernel", "Fp8BlockScaledMMLinearKernel", diff --git a/vllm/model_executor/kernels/linear/scaled_mm/triton.py b/vllm/model_executor/kernels/linear/scaled_mm/triton.py index d6169d2cfc1a..fed471e474ee 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/triton.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/triton.py @@ -20,6 +20,8 @@ ) from .cutlass import CutlassInt8ScaledMMLinearKernel from .ScaledMMLinearKernel import ( + FP8ScaledMMLinearKernel, + FP8ScaledMMLinearLayerConfig, Int8ScaledMMLinearLayerConfig, ) @@ -156,6 +158,46 @@ def apply_weights( return out +class TritonFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cuda_alike(): + return False, "requires ROCm or CUDA." + return True, None + + @classmethod + def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: + act_group = c.activation_quant_key.scale.group_shape + weight_group = c.weight_quant_key.scale.group_shape + if not (act_group.is_per_tensor() or act_group.is_per_token()): + return False, "requires per-tensor or per-token activation scales." + if not ( + weight_group.is_per_tensor() + or weight_group.is_per_channel() + or weight_group.is_per_token() + ): + return False, "requires per-tensor, per-channel or per-token weight scales." + return True, None + + def apply_scaled_mm( + self, + *, + A: torch.Tensor, + B: torch.Tensor, + out_dtype: torch.dtype, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None, + output_shape: list, + ) -> torch.Tensor: + out = triton_scaled_mm( + A, B, scale_a=As, scale_b=Bs, out_dtype=out_dtype, bias=bias + ) + return out.view(*output_shape) + + class TritonFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): @classmethod def is_supported(cls, compute_capability=None): diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py index c15a1a13d42e..1efb058ef07c 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py @@ -51,6 +51,10 @@ def _supports_current_device() -> bool: def _supports_parallel_config(moe_parallel_config) -> bool: return True + @staticmethod + def _supports_batch_invariance() -> bool: + return False + @staticmethod def is_supported_config( cls, moe_config, weight_key, activation_key, activation_format diff --git a/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py index da0bca4977b7..ed44bf4cb135 100644 --- a/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py @@ -218,7 +218,7 @@ def apply( global_num_experts=global_num_experts, expert_map=expert_map, a1q_scale=None, - a2_scale=None, + a2_scale=a2_scale, workspace13=workspace13, workspace2=workspace2, expert_tokens_meta=expert_tokens_meta, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py b/vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py index e0d1fdf27170..8adb41151c43 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py @@ -150,16 +150,16 @@ def scaled_mm_kernel( scale_b = scale_b.broadcast_to((BLOCK_SIZE_N, 1)) accumulator = scale_b.T * accumulator.to(tl.float32) - # Convert to output format. - c = accumulator.to(c_ptr.type.element_ty) - - # Add bias, it's already in output format, so add it after conversion. + # Add bias in accumulator dtype for less rounding error. if bias_ptr: offsets_bias = offsets_bn bias_ptrs = bias_ptr + offsets_bias bias_mask = offsets_bias < N bias = tl.load(bias_ptrs, bias_mask) - c += bias + accumulator += bias + + # Convert to output format. + c = accumulator.to(c_ptr.type.element_ty) # Save output offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) @@ -224,6 +224,8 @@ def triton_scaled_mm( tile_shape = (64, 128, 128) else: tile_shape = (128, 128, 128) + else: + tile_shape = (block_size_m, block_size_n, block_size_k) block_size_m, block_size_n, block_size_k = tile_shape diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 1c1d18614ecd..b18ec0a0aab3 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -9,10 +9,10 @@ from vllm.config import get_current_vllm_config from vllm.distributed import get_tensor_model_parallel_world_size from vllm.model_executor.kernels.linear import ( + fp8_scaled_mm_is_batch_invariant, init_fp8_linear_kernel, ) from vllm.model_executor.kernels.linear.scaled_mm import ( - CutlassFP8ScaledMMLinearKernel, MarlinFP8ScaledMMLinearKernel, ) from vllm.model_executor.layers.attention import Attention @@ -432,7 +432,7 @@ def apply( bias, ) else: - if isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel): + if fp8_scaled_mm_is_batch_invariant(self.fp8_linear): return self.fp8_linear.apply_weights(layer, x, bias) # per-tensor/channel: dequant to BF16 and run GEMM diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index f4662ad5a66e..8d8a7eca6066 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -15,9 +15,11 @@ import vllm.envs as envs from vllm import _custom_ops as ops from vllm.config import get_current_vllm_config -from vllm.model_executor.kernels.linear import init_fp8_linear_kernel +from vllm.model_executor.kernels.linear import ( + fp8_scaled_mm_is_batch_invariant, + init_fp8_linear_kernel, +) from vllm.model_executor.kernels.linear.scaled_mm import ( - CutlassFP8ScaledMMLinearKernel, MarlinFP8ScaledMMLinearKernel, ) from vllm.model_executor.layers.fused_moe import RoutedExperts @@ -278,7 +280,7 @@ def apply( ) -> torch.Tensor: # if batch invariant mode is enabled, use BF16 dequant if envs.VLLM_BATCH_INVARIANT: - if isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel): + if fp8_scaled_mm_is_batch_invariant(self.fp8_linear): return self.fp8_linear.apply_weights(layer, x, bias) weight_fp8 = layer.weight.to(torch.bfloat16) @@ -466,7 +468,7 @@ def apply( # if batch invariant mode is enabled dequant if ( envs.VLLM_BATCH_INVARIANT - and not isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel) + and not fp8_scaled_mm_is_batch_invariant(self.fp8_linear) and not isinstance(x, QuantizedActivation) ): weight_dequant = ( diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index fb87c3a61450..00f9e5b92eed 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -1041,7 +1041,10 @@ def w8a8_triton_block_scaled_mm( C_shape = A.shape[:-1] + (N,) C = A.new_empty(C_shape, dtype=output_dtype) - configs = get_w8a8_block_fp8_configs(N, K, block_size[0], block_size[1]) + if envs.VLLM_BATCH_INVARIANT: + configs = None # use the default config + else: + configs = get_w8a8_block_fp8_configs(N, K, block_size[0], block_size[1]) if configs: # Get the optimal config if there is one config = configs[min(configs.keys(), key=lambda x: abs(x - M))] diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 80eef7c75a12..8d285dda75b8 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -288,6 +288,13 @@ def wvsplitkrc_dispatch(n: int, k: int, m: int, cu_count: int) -> tuple[int, boo def rocm_unquantized_gemm_impl( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: + if envs.VLLM_BATCH_INVARIANT: + from vllm.model_executor.determinism.batch_invariant import ( + linear_batch_invariant, + ) + + return linear_batch_invariant(x, weight, bias) + from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950, on_gfx1250 n = x.numel() // x.size(-1) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 5840ffe3b792..807e571002ac 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -1068,6 +1068,15 @@ def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None: if use_aiter_fp8_linear and "-quant_fp8" not in compilation_config.custom_ops: compilation_config.custom_ops.append("+quant_fp8") + # Only RMSNorm's custom op reaches rms_norm_batch_invariant. Inductor + # lowers the native one to a reduction whose block size it picks by + # benchmarking, so the summation order varies with the batch and run. + if ( + envs.VLLM_BATCH_INVARIANT + and "-rms_norm" not in compilation_config.custom_ops + ): + compilation_config.custom_ops.append("+rms_norm") + if use_aiter_fused_se and "-grouped_topk" in compilation_config.custom_ops: logger.warning_once( "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled, which " diff --git a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py index 8b0e337bb221..ecc2a50135f6 100644 --- a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py @@ -12,6 +12,7 @@ import torch +import vllm.envs as envs from vllm.platforms import current_platform from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend @@ -47,6 +48,15 @@ def is_available(cls) -> bool: return rocm_aiter_ops.is_enabled() + @classmethod + def validate_configuration(cls, device_capability, selector_config): + invalid_reasons = super().validate_configuration( + device_capability, selector_config + ) + if envs.VLLM_BATCH_INVARIANT: + invalid_reasons.append("batch invariance not supported") + return invalid_reasons + def __init__( self, num_heads: int, diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py index 3af623b1f13b..b6cf96e5b738 100644 --- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -427,13 +427,13 @@ def _flash_attn_varlen_diff_headdims( kwargs["return_softmax_lse"] = return_softmax_lse kwargs["out"] = out kwargs["output_scale"] = output_scale + if envs.VLLM_BATCH_INVARIANT: + kwargs["num_splits"] = 1 else: # ROCm leverages the upstream flash_attn, which takes a parameter # called "return_attn_probs" instead of return_softmax_lse kwargs["return_attn_probs"] = return_softmax_lse assert out is None and output_scale is None - if envs.VLLM_BATCH_INVARIANT: - kwargs["num_splits"] = 1 attn_out = _FA4_MLA_PREFILL_KERNEL( q=q, From d79b52b18219ecbaba780f72ae0a4aa0546ac00f Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Mon, 12 Oct 2026 12:46:22 -0400 Subject: [PATCH 315/624] [Misc] Consolidate csrc shared headers and prune dead code (#60721) Signed-off-by: mgoin Co-authored-by: Claude --- .buildkite/test-amd.yaml | 46 +- .buildkite/test_areas/distributed.yaml | 4 +- .buildkite/test_areas/kernels.yaml | 35 +- .buildkite/test_areas/lm_eval.yaml | 8 +- .buildkite/test_areas/models_basic.yaml | 2 +- .gitignore | 3 - .pre-commit-config.yaml | 2 +- CMakeLists.txt | 235 ++-- cmake/cpu_extension.cmake | 6 +- cmake/hipify.py | 2 +- .../activation_kernels.cu | 2 +- csrc/{libtorch_stable => }/async_util.cuh | 0 .../dcp_utils/dcp_direct_a2a_lse_reduce.cu | 0 .../attention/dcp_utils/dcp_direct_common.cuh | 0 .../dcp_utils/dcp_direct_kv_gather.cu | 0 .../dcp_utils/dcp_direct_q_gather.cu | 0 .../attention/merge_attn_states.cu | 5 +- .../cutlass_sm100_mla/device/sm100_mla.hpp | 0 .../kernel/sm100_fmha_mla_reduction.hpp | 0 .../sm100_fmha_mla_tma_warpspecialized.hpp | 0 .../kernel/sm100_mla_tile_scheduler.hpp | 0 .../attention/mla/sm100_cutlass_mla_kernel.cu | 2 +- csrc/cache.h | 93 -- csrc/{libtorch_stable => }/cache_kernels.cu | 22 +- .../cache_kernels_fused.cu | 8 +- csrc/{libtorch_stable => }/concat_mla_q.cuh | 0 .../{libtorch_stable => }/cooperative_topk.cu | 0 .../cooperative_topk.cuh | 0 csrc/core/batch_invariant.hpp | 18 - csrc/{ => core}/cuda_compat.h | 6 - csrc/core/exception.hpp | 3 - csrc/core/registration.h | 5 - .../core/math.hpp => core/utils.hpp} | 16 +- .../vec_dtypes}/dtype_bfloat16.cuh | 2 +- .../vec_dtypes}/dtype_float16.cuh | 2 +- .../vec_dtypes}/dtype_float32.cuh | 2 +- .../vec_dtypes}/dtype_fp8.cuh | 7 +- .../vec_dtypes/dtype_generic.cuh} | 0 .../vec_dtypes/vec_dtypes.h} | 2 +- csrc/cpu/cpu_attn_impl.hpp | 2 +- .../{moe => cpu}/dynamic_4bit_int_moe_cpu.cpp | 0 csrc/cpu/ops.h | 46 + csrc/cpu/torch_bindings.cpp | 2 +- csrc/{libtorch_stable => }/cub_helpers.h | 0 csrc/cuda_utils.h | 25 +- .../cuda_utils_kernels.cu | 0 csrc/{libtorch_stable => }/cuda_vec_utils.cuh | 0 csrc/{libtorch_stable => }/cuda_view.cu | 0 .../custom_all_gather_reduce_scatter.cu | 0 .../custom_all_gather_reduce_scatter_ops.cpp | 0 .../custom_all_reduce.cu | 0 csrc/custom_collective_common.cuh | 31 - .../cutlass_extensions/common.hpp | 25 +- .../epilogue/broadcast_load_epilogue_c2x.hpp | 0 .../epilogue/scaled_mm_epilogues_c2x.hpp | 0 .../epilogue/scaled_mm_epilogues_c3x.hpp | 40 +- .../cutlass_extensions/torch_utils.hpp | 49 +- .../vllm_collective_builder.cuh | 0 .../vllm_numeric_conversion.cuh | 6 +- csrc/cutlass_extensions/vllm_type_utils.cuh | 2 +- csrc/dispatch_utils.h | 302 +++-- .../dsv3_fused_a_gemm.cu | 2 +- .../{libtorch_stable => }/fp32_router_gemm.cu | 0 .../fp32_router_gemm_entry.cu | 2 +- ...deepseek_v4_qnorm_rope_kv_insert_kernel.cu | 2 +- ..._kimi_k3_mla_key_concat_kv_cache_kernel.cu | 8 +- ...minimax_m3_qknorm_rope_kv_insert_kernel.cu | 8 +- .../fused_qknorm_rope_kernel.cu | 2 +- .../gdn/fused_gdn_decode_kernel.cu | 2 +- .../{libtorch_stable => }/hisparse_kernels.cu | 1 - .../kimi_k3/attn_res_kernel.cu | 0 .../kimi_k3/fused_kda_chunk_kernel_rocm.cu | 0 .../kimi_k3/fused_kda_decode_kernel.cu | 2 +- .../kimi_k3/fused_kda_decode_kernel_rocm.cu | 0 .../launch_bounds_utils.h | 0 .../layernorm_kernels.cu | 2 +- .../layernorm_quant_kernels.cu | 4 +- .../attention/attention_utils.cuh | 57 - .../cutlass_extensions/common.cpp | 11 - csrc/libtorch_stable/dispatch_utils.h | 188 --- csrc/libtorch_stable/mamba/static_switch.h | 28 - .../moe/marlin_moe_wna16/generate_kernels.py | 315 ----- csrc/libtorch_stable/ops.h | 822 ------------ .../w8a8/int8/per_token_group_quant.cu | 12 - .../w8a8/per_token_group_quant_8bit.h | 20 - csrc/libtorch_stable/torch_bindings.cpp | 1171 ---------------- csrc/libtorch_stable/torch_utils.h | 88 -- .../mamba/selective_scan.h | 0 .../mamba/selective_scan_fwd.cu | 8 +- .../minimax_reduce_rms_kernel.cu | 2 +- .../minimax_reduce_rms_kernel.h | 0 .../moe/grouped_topk_kernels.cu | 2 +- .../moe/marlin_moe_wna16/.gitignore | 1 - .../moe/marlin_moe_wna16/kernel.h | 4 +- .../moe/marlin_moe_wna16/marlin_template.h | 8 +- .../moe/marlin_moe_wna16/ops.cu | 2 +- .../moe/moeTopKFuncs.cuh | 0 .../moe/moe_align_sum_kernels.cu | 10 +- csrc/{libtorch_stable => }/moe/moe_ops.h | 0 .../moe/moe_permute_unpermute_op.cu | 4 +- csrc/{libtorch_stable => }/moe/moe_wna16.cu | 2 +- .../moe/moe_wna16_utils.h | 0 .../moe/permute_unpermute_kernels/dispatch.h | 0 .../moe_permute_unpermute_kernel.cu | 2 +- .../moe_permute_unpermute_kernel.h | 2 +- .../moe_permute_unpermute_kernel.inl | 0 .../moe/topk_softmax_kernels.cu | 4 +- .../moe/topk_softplus_sqrt_kernels.cu | 4 +- .../moe/torch_bindings.cpp | 0 .../ngram_embedding_kernels.cu | 0 .../nvfp4_ds_mla_cache.h | 0 .../nvfp4_ds_mla_cache_kernels.cu | 0 .../nvfp4_kv_cache_kernels.cu | 6 +- csrc/ops.h | 848 +++++++++++- .../{libtorch_stable => }/persistent_topk.cuh | 0 .../pos_encoding_kernels.cu | 2 +- .../quantization/activation_kernels.cu | 8 +- .../quantization/awq/dequantize.cuh | 0 .../quantization/awq/gemm_kernels.cu | 4 +- .../cutlass_w4a8/get_group_starts.cuh | 2 +- .../cutlass_w4a8/w4a8_grouped_mm_entry.cu | 8 +- .../cutlass_w4a8/w4a8_mm_entry.cu | 8 +- .../quantization/cutlass_w4a8/w4a8_utils.cu | 0 .../quantization/cutlass_w4a8/w4a8_utils.cuh | 0 .../activation_nvfp4_quant_fusion_kernels.cu | 6 +- .../fp4/mxfp4_blockwise_moe_kernel.cu | 4 +- .../quantization/fp4/mxfp4_experts_quant.cu | 8 +- .../fp4/nvfp4_blockwise_moe_kernel.cu | 4 +- .../quantization/fp4/nvfp4_experts_quant.cu | 6 +- .../quantization/fp4/nvfp4_quant_entry.cu | 4 +- .../quantization/fp4/nvfp4_quant_kernels.cu | 6 +- .../quantization/fp4/nvfp4_scaled_mm_entry.cu | 4 +- .../fp4/nvfp4_scaled_mm_kernels.cu | 7 +- .../fp4/nvfp4_scaled_mm_sm120_kernels.cu | 7 +- .../quantization/fp4/nvfp4_utils.cuh | 0 ...fused_layernorm_dynamic_per_token_quant.cu | 2 +- .../fused_silu_mul_block_quant.cu | 0 .../fused_kernels/layernorm_utils.cuh | 4 +- .../fused_kernels/quant_conversions.cuh | 4 +- .../quantization/gptq/compat.cuh | 0 .../quantization/gptq/matrix_view.cuh | 0 .../quantization/gptq/q_gemm.cu | 0 .../quantization/gptq/qdq_2.cuh | 0 .../quantization/gptq/qdq_3.cuh | 0 .../quantization/gptq/qdq_4.cuh | 0 .../quantization/gptq/qdq_8.cuh | 0 .../quantization/gptq/qdq_util.cuh | 0 .../hadacore/hadamard_transform_cuda.cu | 4 +- .../quantization/machete/Readme.md | 0 .../quantization/machete/generate.py | 3 +- .../machete/machete_collective_builder.cuh | 2 +- .../machete/machete_interleaving_utils.cuh | 0 .../quantization/machete/machete_mainloop.cuh | 5 +- .../machete/machete_mm_kernel.cuh | 6 +- .../machete/machete_mm_launcher.cuh | 4 +- .../machete/machete_prepack_kernel.cuh | 2 +- .../machete/machete_prepack_launcher.cuh | 4 +- .../machete/machete_prepacked_layout.cuh | 0 .../quantization/machete/machete_pytorch.cu | 0 .../quantization/marlin/.gitignore | 1 - .../quantization/marlin/awq_marlin_repack.cu | 2 +- .../quantization/marlin/dequant.h | 0 .../quantization/marlin/generate_kernels.py | 11 +- .../quantization/marlin/gptq_marlin_repack.cu | 2 +- .../quantization/marlin/kernel.h | 0 .../quantization/marlin/marlin.cu | 2 +- .../quantization/marlin/marlin.cuh | 0 .../quantization/marlin/marlin_dtypes.cuh | 0 .../marlin/marlin_int4_fp8_preprocess.cu | 2 +- .../quantization/marlin/marlin_mma.h | 0 .../quantization/marlin/marlin_template.h | 0 .../quantization/vectorization.cuh | 0 .../quantization/vectorization_utils.cuh | 0 .../w8a8/cutlass/c3x/cutlass_gemm_caller.cuh | 6 +- .../w8a8/cutlass/c3x/scaled_mm.cuh | 4 +- .../cutlass/c3x/scaled_mm_azp_sm90_int8.cu | 2 +- .../c3x/scaled_mm_blockwise_sm100_fp8.cu | 2 +- ...scaled_mm_blockwise_sm100_fp8_dispatch.cuh | 6 +- .../c3x/scaled_mm_blockwise_sm120_fp8.cu | 2 +- ...scaled_mm_blockwise_sm120_fp8_dispatch.cuh | 1 - .../c3x/scaled_mm_blockwise_sm90_fp8.cu | 2 +- .../scaled_mm_blockwise_sm90_fp8_dispatch.cuh | 0 .../w8a8/cutlass/c3x/scaled_mm_helper.hpp | 18 +- .../w8a8/cutlass/c3x/scaled_mm_kernels.hpp | 0 .../w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu | 2 +- .../c3x/scaled_mm_sm100_fp8_dispatch.cuh | 2 +- .../w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu | 4 +- .../c3x/scaled_mm_sm120_fp8_dispatch.cuh | 0 .../w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu | 2 +- .../c3x/scaled_mm_sm90_fp8_dispatch.cuh | 2 +- .../w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu | 2 +- .../c3x/scaled_mm_sm90_int8_dispatch.cuh | 0 .../w8a8/cutlass/moe/get_group_starts.cuh | 2 +- .../w8a8/cutlass/moe/grouped_mm_c3x.cuh | 6 +- .../w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu | 2 +- .../w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu | 2 +- .../quantization/w8a8/cutlass/moe/moe_data.cu | 4 +- .../w8a8/cutlass/scaled_mm_c2x.cu | 4 +- .../w8a8/cutlass/scaled_mm_c2x.cuh | 6 +- .../cutlass/scaled_mm_c2x_sm75_dispatch.cuh | 0 .../cutlass/scaled_mm_c2x_sm80_dispatch.cuh | 0 .../scaled_mm_c2x_sm89_fp8_dispatch.cuh | 0 .../scaled_mm_c2x_sm89_int8_dispatch.cuh | 0 .../w8a8/cutlass/scaled_mm_c3x_sm100.cu | 0 .../w8a8/cutlass/scaled_mm_c3x_sm120.cu | 0 .../w8a8/cutlass/scaled_mm_c3x_sm90.cu | 0 .../w8a8/cutlass/scaled_mm_entry.cu | 4 +- .../quantization/w8a8/fp8/amd/quant_utils.cuh | 2 +- .../quantization/w8a8/fp8/common.cu | 4 +- csrc/quantization/w8a8/fp8/common.cuh | 16 +- .../w8a8/fp8/nvidia/quant_utils.cuh | 4 +- .../w8a8/fp8/per_token_group_quant.cu | 31 +- .../quantization/w8a8/int8/scaled_quant.cu | 0 csrc/rocm/attention.cu | 6 +- csrc/rocm/dispatch_utils.h | 24 + csrc/rocm/q_gemm_rdna3.cu | 3 +- csrc/{ => rocm}/quickreduce/base.h | 0 .../quickreduce}/custom_quickreduce.cu | 2 +- csrc/{ => rocm}/quickreduce/quick_reduce.h | 0 .../quickreduce/quick_reduce_impl.cuh | 0 csrc/rocm/quickreduce/torch_bindings.cpp | 51 + csrc/rocm/skinny_gemms.cu | 4 +- csrc/rocm/skinny_gemms_int4.cu | 3 +- csrc/{libtorch_stable => }/sampled_topk.cuh | 0 csrc/{libtorch_stable => }/sampler.cu | 2 +- csrc/{libtorch_stable => }/topk.cu | 0 .../topk_histogram_4096.cuh | 0 csrc/torch_bindings.cpp | 1194 ++++++++++++++++- csrc/torch_utils.h | 101 +- csrc/{libtorch_stable => }/type_convert.cuh | 0 .../vocab_parallel_embedding_kernels.cu | 0 docs/design/fusions.md | 4 +- tests/kernels/test_bf16_skinny_gemm.py | 5 +- tools/ci/kernel_symbol_map.py | 6 +- .../device_communicators/quick_all_reduce.py | 2 +- .../kernels/linear/scaled_mm/cutlass.py | 3 +- vllm/models/common/ops/fused_qk_rmsnorm.py | 3 +- vllm/models/kimi_k3/amd/ops/kda_chunk.py | 2 +- vllm/models/kimi_k3/amd/ops/kda_decode.py | 2 +- .../ops/fused_mla_key_concat_kv_cache.py | 2 +- .../attention/backends/mla/flashmla_sparse.py | 2 +- vllm/v1/hisparse/runtime.py | 2 +- 242 files changed, 2792 insertions(+), 3647 deletions(-) rename csrc/{libtorch_stable => }/activation_kernels.cu (99%) rename csrc/{libtorch_stable => }/async_util.cuh (100%) rename csrc/{libtorch_stable => }/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu (100%) rename csrc/{libtorch_stable => }/attention/dcp_utils/dcp_direct_common.cuh (100%) rename csrc/{libtorch_stable => }/attention/dcp_utils/dcp_direct_kv_gather.cu (100%) rename csrc/{libtorch_stable => }/attention/dcp_utils/dcp_direct_q_gather.cu (100%) rename csrc/{libtorch_stable => }/attention/merge_attn_states.cu (99%) rename csrc/{libtorch_stable => }/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp (100%) rename csrc/{libtorch_stable => }/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp (100%) rename csrc/{libtorch_stable => }/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp (100%) rename csrc/{libtorch_stable => }/attention/mla/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp (100%) rename csrc/{libtorch_stable => }/attention/mla/sm100_cutlass_mla_kernel.cu (99%) delete mode 100644 csrc/cache.h rename csrc/{libtorch_stable => }/cache_kernels.cu (99%) rename csrc/{libtorch_stable => }/cache_kernels_fused.cu (98%) rename csrc/{libtorch_stable => }/concat_mla_q.cuh (100%) rename csrc/{libtorch_stable => }/cooperative_topk.cu (100%) rename csrc/{libtorch_stable => }/cooperative_topk.cuh (100%) delete mode 100644 csrc/core/batch_invariant.hpp rename csrc/{ => core}/cuda_compat.h (90%) delete mode 100644 csrc/core/exception.hpp rename csrc/{libtorch_stable/core/math.hpp => core/utils.hpp} (70%) rename csrc/{attention => core/vec_dtypes}/dtype_bfloat16.cuh (99%) rename csrc/{attention => core/vec_dtypes}/dtype_float16.cuh (99%) rename csrc/{attention => core/vec_dtypes}/dtype_float32.cuh (99%) rename csrc/{attention => core/vec_dtypes}/dtype_fp8.cuh (88%) rename csrc/{attention/attention_generic.cuh => core/vec_dtypes/dtype_generic.cuh} (100%) rename csrc/{attention/attention_dtypes.h => core/vec_dtypes/vec_dtypes.h} (79%) rename csrc/{moe => cpu}/dynamic_4bit_int_moe_cpu.cpp (100%) create mode 100644 csrc/cpu/ops.h rename csrc/{libtorch_stable => }/cub_helpers.h (100%) rename csrc/{libtorch_stable => }/cuda_utils_kernels.cu (100%) rename csrc/{libtorch_stable => }/cuda_vec_utils.cuh (100%) rename csrc/{libtorch_stable => }/cuda_view.cu (100%) rename csrc/{libtorch_stable => }/custom_all_gather_reduce_scatter.cu (100%) rename csrc/{libtorch_stable => }/custom_all_gather_reduce_scatter_ops.cpp (100%) rename csrc/{libtorch_stable => }/custom_all_reduce.cu (100%) rename csrc/{libtorch_stable => }/cutlass_extensions/common.hpp (89%) rename csrc/{libtorch_stable => }/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp (100%) rename csrc/{libtorch_stable => }/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp (100%) rename csrc/{libtorch_stable => }/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp (93%) rename csrc/{libtorch_stable => }/cutlass_extensions/torch_utils.hpp (76%) rename csrc/{libtorch_stable => }/cutlass_extensions/vllm_collective_builder.cuh (100%) rename csrc/{libtorch_stable => }/cutlass_extensions/vllm_numeric_conversion.cuh (99%) rename csrc/{libtorch_stable => }/dsv3_fused_a_gemm.cu (99%) rename csrc/{libtorch_stable => }/fp32_router_gemm.cu (100%) rename csrc/{libtorch_stable => }/fp32_router_gemm_entry.cu (99%) rename csrc/{libtorch_stable => }/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu (99%) rename csrc/{libtorch_stable => }/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu (99%) rename csrc/{libtorch_stable => }/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu (99%) rename csrc/{libtorch_stable => }/fused_qknorm_rope_kernel.cu (99%) rename csrc/{libtorch_stable => }/gdn/fused_gdn_decode_kernel.cu (99%) rename csrc/{libtorch_stable => }/hisparse_kernels.cu (99%) rename csrc/{libtorch_stable => }/kimi_k3/attn_res_kernel.cu (100%) rename csrc/{libtorch_stable => }/kimi_k3/fused_kda_chunk_kernel_rocm.cu (100%) rename csrc/{libtorch_stable => }/kimi_k3/fused_kda_decode_kernel.cu (99%) rename csrc/{libtorch_stable => }/kimi_k3/fused_kda_decode_kernel_rocm.cu (100%) rename csrc/{libtorch_stable => }/launch_bounds_utils.h (100%) rename csrc/{libtorch_stable => }/layernorm_kernels.cu (99%) rename csrc/{libtorch_stable => }/layernorm_quant_kernels.cu (99%) delete mode 100644 csrc/libtorch_stable/attention/attention_utils.cuh delete mode 100644 csrc/libtorch_stable/cutlass_extensions/common.cpp delete mode 100644 csrc/libtorch_stable/dispatch_utils.h delete mode 100644 csrc/libtorch_stable/mamba/static_switch.h delete mode 100644 csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py delete mode 100644 csrc/libtorch_stable/ops.h delete mode 100644 csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu delete mode 100644 csrc/libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h delete mode 100644 csrc/libtorch_stable/torch_bindings.cpp delete mode 100644 csrc/libtorch_stable/torch_utils.h rename csrc/{libtorch_stable => }/mamba/selective_scan.h (100%) rename csrc/{libtorch_stable => }/mamba/selective_scan_fwd.cu (99%) rename csrc/{libtorch_stable => }/minimax_reduce_rms_kernel.cu (99%) rename csrc/{libtorch_stable => }/minimax_reduce_rms_kernel.h (100%) rename csrc/{libtorch_stable => }/moe/grouped_topk_kernels.cu (99%) rename csrc/{libtorch_stable => }/moe/marlin_moe_wna16/.gitignore (73%) rename csrc/{libtorch_stable => }/moe/marlin_moe_wna16/kernel.h (95%) rename csrc/{libtorch_stable => }/moe/marlin_moe_wna16/marlin_template.h (99%) rename csrc/{libtorch_stable => }/moe/marlin_moe_wna16/ops.cu (99%) rename csrc/{libtorch_stable => }/moe/moeTopKFuncs.cuh (100%) rename csrc/{libtorch_stable => }/moe/moe_align_sum_kernels.cu (99%) rename csrc/{libtorch_stable => }/moe/moe_ops.h (100%) rename csrc/{libtorch_stable => }/moe/moe_permute_unpermute_op.cu (99%) rename csrc/{libtorch_stable => }/moe/moe_wna16.cu (99%) rename csrc/{libtorch_stable => }/moe/moe_wna16_utils.h (100%) rename csrc/{libtorch_stable => }/moe/permute_unpermute_kernels/dispatch.h (100%) rename csrc/{libtorch_stable => }/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu (98%) rename csrc/{libtorch_stable => }/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h (97%) rename csrc/{libtorch_stable => }/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl (100%) rename csrc/{libtorch_stable => }/moe/topk_softmax_kernels.cu (99%) rename csrc/{libtorch_stable => }/moe/topk_softplus_sqrt_kernels.cu (99%) rename csrc/{libtorch_stable => }/moe/torch_bindings.cpp (100%) rename csrc/{libtorch_stable => }/ngram_embedding_kernels.cu (100%) rename csrc/{libtorch_stable => }/nvfp4_ds_mla_cache.h (100%) rename csrc/{libtorch_stable => }/nvfp4_ds_mla_cache_kernels.cu (100%) rename csrc/{libtorch_stable => }/nvfp4_kv_cache_kernels.cu (99%) rename csrc/{libtorch_stable => }/persistent_topk.cuh (100%) rename csrc/{libtorch_stable => }/pos_encoding_kernels.cu (99%) rename csrc/{libtorch_stable => }/quantization/activation_kernels.cu (99%) rename csrc/{libtorch_stable => }/quantization/awq/dequantize.cuh (100%) rename csrc/{libtorch_stable => }/quantization/awq/gemm_kernels.cu (99%) rename csrc/{libtorch_stable => }/quantization/cutlass_w4a8/get_group_starts.cuh (99%) rename csrc/{libtorch_stable => }/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu (98%) rename csrc/{libtorch_stable => }/quantization/cutlass_w4a8/w4a8_mm_entry.cu (98%) rename csrc/{libtorch_stable => }/quantization/cutlass_w4a8/w4a8_utils.cu (100%) rename csrc/{libtorch_stable => }/quantization/cutlass_w4a8/w4a8_utils.cuh (100%) rename csrc/{libtorch_stable => }/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu (97%) rename csrc/{libtorch_stable => }/quantization/fp4/mxfp4_blockwise_moe_kernel.cu (99%) rename csrc/{libtorch_stable => }/quantization/fp4/mxfp4_experts_quant.cu (99%) rename csrc/{libtorch_stable => }/quantization/fp4/nvfp4_blockwise_moe_kernel.cu (99%) rename csrc/{libtorch_stable => }/quantization/fp4/nvfp4_experts_quant.cu (99%) rename csrc/{libtorch_stable => }/quantization/fp4/nvfp4_quant_entry.cu (98%) rename csrc/{libtorch_stable => }/quantization/fp4/nvfp4_quant_kernels.cu (98%) rename csrc/{libtorch_stable => }/quantization/fp4/nvfp4_scaled_mm_entry.cu (97%) rename csrc/{libtorch_stable => }/quantization/fp4/nvfp4_scaled_mm_kernels.cu (98%) rename csrc/{libtorch_stable => }/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu (98%) rename csrc/{libtorch_stable => }/quantization/fp4/nvfp4_utils.cuh (100%) rename csrc/{libtorch_stable => }/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu (99%) rename csrc/{libtorch_stable => }/quantization/fused_kernels/fused_silu_mul_block_quant.cu (100%) rename csrc/{libtorch_stable => }/quantization/fused_kernels/layernorm_utils.cuh (99%) rename csrc/{libtorch_stable => }/quantization/fused_kernels/quant_conversions.cuh (96%) rename csrc/{libtorch_stable => }/quantization/gptq/compat.cuh (100%) rename csrc/{libtorch_stable => }/quantization/gptq/matrix_view.cuh (100%) rename csrc/{libtorch_stable => }/quantization/gptq/q_gemm.cu (100%) rename csrc/{libtorch_stable => }/quantization/gptq/qdq_2.cuh (100%) rename csrc/{libtorch_stable => }/quantization/gptq/qdq_3.cuh (100%) rename csrc/{libtorch_stable => }/quantization/gptq/qdq_4.cuh (100%) rename csrc/{libtorch_stable => }/quantization/gptq/qdq_8.cuh (100%) rename csrc/{libtorch_stable => }/quantization/gptq/qdq_util.cuh (100%) rename csrc/{libtorch_stable => }/quantization/hadamard/hadacore/hadamard_transform_cuda.cu (99%) rename csrc/{libtorch_stable => }/quantization/machete/Readme.md (100%) rename csrc/{libtorch_stable => }/quantization/machete/generate.py (99%) rename csrc/{libtorch_stable => }/quantization/machete/machete_collective_builder.cuh (94%) rename csrc/{libtorch_stable => }/quantization/machete/machete_interleaving_utils.cuh (100%) rename csrc/{libtorch_stable => }/quantization/machete/machete_mainloop.cuh (99%) rename csrc/{libtorch_stable => }/quantization/machete/machete_mm_kernel.cuh (98%) rename csrc/{libtorch_stable => }/quantization/machete/machete_mm_launcher.cuh (96%) rename csrc/{libtorch_stable => }/quantization/machete/machete_prepack_kernel.cuh (98%) rename csrc/{libtorch_stable => }/quantization/machete/machete_prepack_launcher.cuh (96%) rename csrc/{libtorch_stable => }/quantization/machete/machete_prepacked_layout.cuh (100%) rename csrc/{libtorch_stable => }/quantization/machete/machete_pytorch.cu (100%) rename csrc/{libtorch_stable => }/quantization/marlin/.gitignore (73%) rename csrc/{libtorch_stable => }/quantization/marlin/awq_marlin_repack.cu (99%) rename csrc/{libtorch_stable => }/quantization/marlin/dequant.h (100%) rename csrc/{libtorch_stable => }/quantization/marlin/generate_kernels.py (96%) rename csrc/{libtorch_stable => }/quantization/marlin/gptq_marlin_repack.cu (99%) rename csrc/{libtorch_stable => }/quantization/marlin/kernel.h (100%) rename csrc/{libtorch_stable => }/quantization/marlin/marlin.cu (99%) rename csrc/{libtorch_stable => }/quantization/marlin/marlin.cuh (100%) rename csrc/{libtorch_stable => }/quantization/marlin/marlin_dtypes.cuh (100%) rename csrc/{libtorch_stable => }/quantization/marlin/marlin_int4_fp8_preprocess.cu (99%) rename csrc/{libtorch_stable => }/quantization/marlin/marlin_mma.h (100%) rename csrc/{libtorch_stable => }/quantization/marlin/marlin_template.h (100%) rename csrc/{libtorch_stable => }/quantization/vectorization.cuh (100%) rename csrc/{libtorch_stable => }/quantization/vectorization_utils.cuh (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh (96%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm.cuh (98%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu (90%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu (89%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh (98%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu (96%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh (99%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu (89%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp (82%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_kernels.hpp (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu (97%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh (99%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu (91%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu (97%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh (99%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu (91%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8_dispatch.cuh (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/moe/get_group_starts.cuh (99%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh (97%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu (99%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu (99%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/moe/moe_data.cu (99%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_c2x.cu (98%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_c2x.cuh (97%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_c2x_sm75_dispatch.cuh (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_c2x_sm80_dispatch.cuh (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_fp8_dispatch.cuh (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_int8_dispatch.cuh (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_entry.cu (99%) rename csrc/{libtorch_stable => }/quantization/w8a8/fp8/common.cu (99%) rename csrc/{libtorch_stable => }/quantization/w8a8/fp8/per_token_group_quant.cu (97%) rename csrc/{libtorch_stable => }/quantization/w8a8/int8/scaled_quant.cu (100%) create mode 100644 csrc/rocm/dispatch_utils.h rename csrc/{ => rocm}/quickreduce/base.h (100%) rename csrc/{ => rocm/quickreduce}/custom_quickreduce.cu (99%) rename csrc/{ => rocm}/quickreduce/quick_reduce.h (100%) rename csrc/{ => rocm}/quickreduce/quick_reduce_impl.cuh (100%) create mode 100644 csrc/rocm/quickreduce/torch_bindings.cpp rename csrc/{libtorch_stable => }/sampled_topk.cuh (100%) rename csrc/{libtorch_stable => }/sampler.cu (99%) rename csrc/{libtorch_stable => }/topk.cu (100%) rename csrc/{libtorch_stable => }/topk_histogram_4096.cuh (100%) rename csrc/{libtorch_stable => }/type_convert.cuh (100%) rename csrc/{libtorch_stable => }/vocab_parallel_embedding_kernels.cu (100%) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 49bc3c32b67a..e4ae7f897e35 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -829,9 +829,7 @@ steps: - examples/features/data_parallel/data_parallel_offline.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py - - csrc/custom_quickreduce.cu - - csrc/ops.h - - csrc/torch_bindings.cpp + - csrc/rocm/quickreduce/ - vllm/model_executor/layers/ - vllm/entrypoints/llm.py - vllm/config/parallel.py @@ -1514,7 +1512,7 @@ steps: parallelism: 5 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - csrc/rocm/ - tests/kernels/moe @@ -1560,6 +1558,7 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - csrc/attention/ + - csrc/core/vec_dtypes/ - vllm/v1/attention - vllm/model_executor/layers/attention - tests/kernels/attention @@ -1585,10 +1584,10 @@ steps: - vllm/platforms/rocm.py - tests/kernels/core/test_minimax_reduce_rms.py - tests/utils.py - - csrc/libtorch_stable/minimax_reduce_rms_kernel.cu - - csrc/libtorch_stable/minimax_reduce_rms_kernel.h - - csrc/libtorch_stable/ops.h - - csrc/libtorch_stable/torch_bindings.cpp + - csrc/minimax_reduce_rms_kernel.cu + - csrc/minimax_reduce_rms_kernel.h + - csrc/ops.h + - csrc/torch_bindings.cpp - vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py optional: true commands: @@ -1610,9 +1609,9 @@ steps: - vllm/models/kimi_k3/nvidia/ops/third_party/kda/ - vllm/models/kimi_k3/amd/ - CMakeLists.txt - - csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu - - csrc/libtorch_stable/ops.h - - csrc/libtorch_stable/torch_bindings.cpp + - csrc/kimi_k3/fused_kda_decode_kernel_rocm.cu + - csrc/ops.h + - csrc/torch_bindings.cpp - tests/models/kimi_k3/test_amd_attn_res.py - tests/models/kimi_k3/test_amd_kda_decode.py - tests/models/kimi_k3/test_kda.py @@ -1633,7 +1632,7 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/models/kimi_k3/ - - csrc/libtorch_stable/kimi_k3/ + - csrc/kimi_k3/ - tests/models/kimi_k3/ - vllm/platforms/rocm.py commands: @@ -3512,9 +3511,7 @@ steps: optional: true working_dir: "/vllm-workspace/" source_file_dependencies: - - csrc/custom_quickreduce.cu - - csrc/ops.h - - csrc/torch_bindings.cpp + - csrc/rocm/quickreduce/ - vllm/distributed/ - vllm/model_executor/layers/ - vllm/entrypoints/llm.py @@ -4033,7 +4030,7 @@ steps: optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/libtorch_stable/gdn/ + - csrc/gdn/ - tests/evals/qwen4_exp/ - vllm/model_executor/layers/mamba/ - vllm/models/qwen4_exp/ @@ -4604,10 +4601,10 @@ steps: - vllm/platforms/rocm.py - tests/kernels/core/test_minimax_reduce_rms.py - tests/utils.py - - csrc/libtorch_stable/minimax_reduce_rms_kernel.cu - - csrc/libtorch_stable/minimax_reduce_rms_kernel.h - - csrc/libtorch_stable/ops.h - - csrc/libtorch_stable/torch_bindings.cpp + - csrc/minimax_reduce_rms_kernel.cu + - csrc/minimax_reduce_rms_kernel.h + - csrc/ops.h + - csrc/torch_bindings.cpp - vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py optional: true commands: @@ -4622,7 +4619,7 @@ steps: source_file_dependencies: - csrc/quantization/fp4/ - csrc/attention/mla/ - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - vllm/model_executor/layers/fused_moe/cutlass_moe.py - vllm/v1/attention/backends/triton_attn.py - vllm/v1/attention/backends/rocm_attn.py @@ -4674,7 +4671,7 @@ steps: optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu + - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu - vllm/models/deepseek_v4/common/ops/ - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py - vllm/platforms/rocm.py @@ -4690,6 +4687,7 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - csrc/attention/ + - csrc/core/vec_dtypes/ - vllm/v1/attention - vllm/model_executor/layers/attention - tests/kernels/attention @@ -4709,7 +4707,7 @@ steps: parallelism: 5 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - csrc/rocm/ - tests/kernels/moe @@ -4759,7 +4757,7 @@ steps: - vllm/utils/torch_utils.py - vllm/platforms/ - vllm/_aiter_ops.py - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - vllm/distributed/device_communicators/ - vllm/config commands: diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index f83946386a38..0b1a73136b04 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -510,9 +510,7 @@ steps: - examples/features/data_parallel/data_parallel_offline.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py - - csrc/custom_quickreduce.cu - - csrc/ops.h - - csrc/torch_bindings.cpp + - csrc/rocm/quickreduce/ - vllm/model_executor/layers/ - vllm/entrypoints/llm.py - vllm/config/parallel.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index aa0971b4c55e..472c19e74ae1 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -72,10 +72,10 @@ steps: num_devices: 2 device: h100 source_file_dependencies: - - csrc/libtorch_stable/minimax_reduce_rms_kernel.cu - - csrc/libtorch_stable/minimax_reduce_rms_kernel.h - - csrc/libtorch_stable/ops.h - - csrc/libtorch_stable/torch_bindings.cpp + - csrc/minimax_reduce_rms_kernel.cu + - csrc/minimax_reduce_rms_kernel.h + - csrc/ops.h + - csrc/torch_bindings.cpp - vllm/model_executor/layers/minimax_rms_norm/ - vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py - tests/kernels/core/test_minimax_reduce_rms.py @@ -122,7 +122,7 @@ steps: depends_on: - image-build-amd source_file_dependencies: - - csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu + - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu - vllm/models/deepseek_v4/common/ops/ - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py - vllm/platforms/rocm.py @@ -316,6 +316,7 @@ steps: timeout_in_minutes: 30 source_file_dependencies: - csrc/attention/ + - csrc/core/vec_dtypes/ - vllm/v1/attention # TODO: remove this dependency (https://github.com/vllm-project/vllm/issues/32267) - vllm/model_executor/layers/attention @@ -335,6 +336,7 @@ steps: - image-build-amd source_file_dependencies: - csrc/attention/ + - csrc/core/vec_dtypes/ - vllm/v1/attention - vllm/model_executor/layers/attention - tests/kernels/attention @@ -467,7 +469,7 @@ steps: num_devices: 1 timeout_in_minutes: 50 source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - tests/kernels/moe - vllm/model_executor/layers/fused_moe/ @@ -485,7 +487,7 @@ steps: device: mi300_1 timeout_in_minutes: 55 source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - tests/kernels/moe - vllm/model_executor/layers/fused_moe/ @@ -576,11 +578,10 @@ steps: # optional: true source_file_dependencies: - csrc/quantization/fp4/ - - csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu - - csrc/libtorch_stable/quantization/fp4/ - - csrc/libtorch_stable/quantization/w8a8/cutlass/ + - csrc/nvfp4_kv_cache_kernels.cu + - csrc/quantization/w8a8/cutlass/ - csrc/attention/mla/ - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - vllm/model_executor/layers/fusion/ - vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py - vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py @@ -606,9 +607,9 @@ steps: - vllm/config/attention.py - vllm/v1/attention/backends/registry.py - vllm/_custom_ops.py - - csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu - - csrc/libtorch_stable/ops.h - - csrc/libtorch_stable/torch_bindings.cpp + - csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu + - csrc/ops.h + - csrc/torch_bindings.cpp - tests/kernels/attention/test_flashinfer_mla_decode.py - tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py - tests/kernels/attention/test_cache.py @@ -853,7 +854,7 @@ steps: device: h100 num_devices: 2 source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - tests/kernels/moe - vllm/model_executor/layers/fused_moe/ @@ -896,7 +897,7 @@ steps: device: b200-k8s num_devices: 2 source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - tests/kernels/moe - vllm/model_executor/layers/fused_moe/ @@ -1009,7 +1010,7 @@ steps: source_file_dependencies: - csrc/quantization/fp4/ - csrc/attention/mla/ - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - vllm/model_executor/layers/fused_moe/cutlass_moe.py - vllm/v1/attention/backends/triton_attn.py - vllm/v1/attention/backends/rocm_attn.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 06607951057b..f6a4a89f11cd 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -230,7 +230,7 @@ steps: - vllm/v1/attention/backends/mla/ - vllm/model_executor/layers/attention/sparse_mla_attention.py - vllm/v1/worker/gpu/ - - csrc/libtorch_stable/hisparse_kernels.cu + - csrc/hisparse_kernels.cu commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-hisparse.txt @@ -1185,7 +1185,7 @@ steps: optional: true num_devices: 4 source_file_dependencies: - - csrc/libtorch_stable/gdn/ + - csrc/gdn/ - tests/evals/qwen4_exp/ - vllm/model_executor/layers/mamba/ - vllm/models/qwen4_exp/ @@ -1203,7 +1203,7 @@ steps: optional: true num_devices: 4 source_file_dependencies: - - csrc/libtorch_stable/gdn/ + - csrc/gdn/ - tests/evals/qwen4_exp/ - vllm/model_executor/layers/mamba/ - vllm/models/qwen4_exp/ @@ -1224,7 +1224,7 @@ steps: depends_on: - image-build-amd source_file_dependencies: - - csrc/libtorch_stable/gdn/ + - csrc/gdn/ - tests/evals/qwen4_exp/ - vllm/model_executor/layers/mamba/ - vllm/models/qwen4_exp/ diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 288e7621cab4..552910e85a96 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -132,7 +132,7 @@ steps: device: b200-k8s source_file_dependencies: - vllm/models/kimi_k3/ - - csrc/libtorch_stable/kimi_k3/ + - csrc/kimi_k3/ - tests/models/kimi_k3/ - tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py commands: diff --git a/.gitignore b/.gitignore index 46418c03924e..57d0d54fc87b 100644 --- a/.gitignore +++ b/.gitignore @@ -237,9 +237,6 @@ benchmarks/**/*.json actionlint shellcheck*/ -# Ignore moe/marlin_moe gen code -csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_* - # Ignore ep_kernels_workspace folder ep_kernels_workspace/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8e257edccc12..78a672c1104b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -21,7 +21,7 @@ repos: rev: v21.1.2 hooks: - id: clang-format - exclude: 'csrc/libtorch_stable/moe/topk_softmax_kernels.cu|vllm/third_party/.*' + exclude: 'csrc/moe/topk_softmax_kernels.cu|vllm/third_party/.*' types_or: [c++, cuda] args: [--style=file, --verbose] - repo: https://github.com/DavidAnson/markdownlint-cli2 diff --git a/CMakeLists.txt b/CMakeLists.txt index b4dd9ada5eb6..3a71d7e1e5ad 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -363,8 +363,8 @@ endif() if(VLLM_GPU_LANG STREQUAL "HIP") set(VLLM_EXT_SRC - "csrc/torch_bindings.cpp" - "csrc/custom_quickreduce.cu") + "csrc/rocm/quickreduce/torch_bindings.cpp" + "csrc/rocm/quickreduce/custom_quickreduce.cu") message(STATUS "Enabling C extension.") define_extension_target( @@ -374,16 +374,8 @@ if(VLLM_GPU_LANG STREQUAL "HIP") 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) endif() # _C HIP endif if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") @@ -419,36 +411,35 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() set(VLLM_STABLE_EXT_SRC - "csrc/libtorch_stable/torch_bindings.cpp" - "csrc/libtorch_stable/cuda_view.cu" - "csrc/libtorch_stable/cuda_utils_kernels.cu" - "csrc/libtorch_stable/activation_kernels.cu" - "csrc/libtorch_stable/ngram_embedding_kernels.cu" - "csrc/libtorch_stable/vocab_parallel_embedding_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" - "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu" - "csrc/libtorch_stable/quantization/gptq/q_gemm.cu" - "csrc/libtorch_stable/pos_encoding_kernels.cu" - "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu" - "csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu" - "csrc/libtorch_stable/layernorm_kernels.cu" - "csrc/libtorch_stable/layernorm_quant_kernels.cu" - "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu" - "csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu" - "csrc/libtorch_stable/attention/merge_attn_states.cu" - "csrc/libtorch_stable/sampler.cu" - "csrc/libtorch_stable/topk.cu" - "csrc/libtorch_stable/mamba/selective_scan_fwd.cu" - "csrc/libtorch_stable/cache_kernels.cu" - "csrc/libtorch_stable/cache_kernels_fused.cu" - "csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu" - "csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp" - "csrc/libtorch_stable/custom_all_reduce.cu" - "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu" - "csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu") + "csrc/torch_bindings.cpp" + "csrc/cuda_view.cu" + "csrc/cuda_utils_kernels.cu" + "csrc/activation_kernels.cu" + "csrc/ngram_embedding_kernels.cu" + "csrc/vocab_parallel_embedding_kernels.cu" + "csrc/quantization/activation_kernels.cu" + "csrc/quantization/w8a8/int8/scaled_quant.cu" + "csrc/quantization/w8a8/fp8/common.cu" + "csrc/quantization/w8a8/fp8/per_token_group_quant.cu" + "csrc/quantization/gptq/q_gemm.cu" + "csrc/pos_encoding_kernels.cu" + "csrc/fused_qknorm_rope_kernel.cu" + "csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu" + "csrc/layernorm_kernels.cu" + "csrc/layernorm_quant_kernels.cu" + "csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu" + "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu" + "csrc/attention/merge_attn_states.cu" + "csrc/sampler.cu" + "csrc/topk.cu" + "csrc/mamba/selective_scan_fwd.cu" + "csrc/cache_kernels.cu" + "csrc/cache_kernels_fused.cu" + "csrc/custom_all_gather_reduce_scatter.cu" + "csrc/custom_all_gather_reduce_scatter_ops.cpp" + "csrc/custom_all_reduce.cu" + "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu" + "csrc/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu") if(VLLM_GPU_LANG STREQUAL "CUDA" AND DEFINED CMAKE_CUDA_COMPILER_VERSION AND @@ -470,10 +461,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_STABLE_EXT_SRC - "csrc/libtorch_stable/hisparse_kernels.cu" - "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu" - "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu" - "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu") + "csrc/hisparse_kernels.cu" + "csrc/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu" + "csrc/attention/dcp_utils/dcp_direct_kv_gather.cu" + "csrc/attention/dcp_utils/dcp_direct_q_gather.cu") SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library") @@ -508,12 +499,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") FetchContent_MakeAvailable(cutlass) list(APPEND VLLM_STABLE_EXT_SRC - "csrc/libtorch_stable/cutlass_extensions/common.cpp" - "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu" - "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu" - "csrc/libtorch_stable/minimax_reduce_rms_kernel.cu") + "csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu" + "csrc/quantization/fp4/nvfp4_quant_entry.cu" + "csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu" + "csrc/quantization/awq/gemm_kernels.cu" + "csrc/minimax_reduce_rms_kernel.cu") # # Machete kernels @@ -527,7 +517,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # preselected input type pairs and schedules. # Generate sources: set(MACHETE_GEN_SCRIPT - ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/quantization/machete/generate.py) + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/machete/generate.py) file(MD5 ${MACHETE_GEN_SCRIPT} MACHETE_GEN_SCRIPT_HASH) message(STATUS "Machete generation script hash: ${MACHETE_GEN_SCRIPT_HASH}") @@ -560,7 +550,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() # Add machete generated sources - file(GLOB MACHETE_GEN_SOURCES "csrc/libtorch_stable/quantization/machete/generated/*.cu") + file(GLOB MACHETE_GEN_SOURCES "csrc/quantization/machete/generated/*.cu") list(APPEND VLLM_STABLE_EXT_SRC ${MACHETE_GEN_SOURCES}) # forward compatible @@ -569,7 +559,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") CUDA_ARCHS "${MACHETE_ARCHS}") list(APPEND VLLM_STABLE_EXT_SRC - csrc/libtorch_stable/quantization/machete/machete_pytorch.cu) + csrc/quantization/machete/machete_pytorch.cu) message(STATUS "Building Machete kernels for archs: ${MACHETE_ARCHS}") else() if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 @@ -587,7 +577,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # The fused MiniMax-M3 KV insert quantizes NVFP4 with SM100-family-only # cvt instructions, so its SM10x/11x code is built for the family targets. set(MINIMAX_M3_FUSED_SRC - "csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu") + "csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu") list(REMOVE_ITEM VLLM_STABLE_EXT_SRC "${MINIMAX_M3_FUSED_SRC}") set_gencode_flags_for_srcs( SRCS "${VLLM_STABLE_EXT_SRC}" @@ -611,9 +601,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") if(COOPERATIVE_TOPK_ARCHS) list(APPEND VLLM_STABLE_EXT_SRC - "csrc/libtorch_stable/cooperative_topk.cu") + "csrc/cooperative_topk.cu") set_gencode_flags_for_srcs( - SRCS "csrc/libtorch_stable/cooperative_topk.cu" + SRCS "csrc/cooperative_topk.cu" CUDA_ARCHS "${COOPERATIVE_TOPK_ARCHS}") endif() @@ -656,7 +646,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # preselected input type pairs and schedules. # Generate sources: set(MARLIN_GEN_SCRIPT - ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/quantization/marlin/generate_kernels.py) + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/marlin/generate_kernels.py) file(MD5 ${MARLIN_GEN_SCRIPT} MARLIN_GEN_SCRIPT_HASH) list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR) set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})") @@ -691,7 +681,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if (MARLIN_ARCHS) - file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm80_kernel_*_float16.cu") + file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/marlin/sm80_kernel_*_float16.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}" CUDA_ARCHS "${MARLIN_ARCHS}") @@ -701,7 +691,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC}) - file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm80_kernel_*_bfloat16.cu") + file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/quantization/marlin/sm80_kernel_*_bfloat16.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_TEMPLATE_BF16_KERNEL_SRC}" CUDA_ARCHS "${MARLIN_BF16_ARCHS}") @@ -713,7 +703,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if (MARLIN_SM75_ARCHS) - file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm75_kernel_*.cu") + file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/quantization/marlin/sm75_kernel_*.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_TEMPLATE_SM75_KERNEL_SRC}" CUDA_ARCHS "${MARLIN_SM75_ARCHS}") @@ -725,7 +715,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if (MARLIN_FP8_ARCHS) - file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm89_kernel_*.cu") + file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/quantization/marlin/sm89_kernel_*.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}" CUDA_ARCHS "${MARLIN_FP8_ARCHS}") @@ -737,10 +727,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() set(MARLIN_SRCS - "csrc/libtorch_stable/quantization/marlin/marlin.cu" - "csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu" - "csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu" - "csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu") + "csrc/quantization/marlin/marlin.cu" + "csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu" + "csrc/quantization/marlin/gptq_marlin_repack.cu" + "csrc/quantization/marlin/awq_marlin_repack.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_SRCS}" CUDA_ARCHS "${MARLIN_OTHER_ARCHS}") @@ -763,7 +753,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS) - set(DSV3_FUSED_A_GEMM_SRCS "csrc/libtorch_stable/dsv3_fused_a_gemm.cu") + set(DSV3_FUSED_A_GEMM_SRCS "csrc/dsv3_fused_a_gemm.cu") set_gencode_flags_for_srcs( SRCS "${DSV3_FUSED_A_GEMM_SRCS}" CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}") @@ -778,8 +768,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_sm90plus(FP32_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND FP32_ROUTER_GEMM_ARCHS) set(FP32_ROUTER_GEMM_SRCS - "csrc/libtorch_stable/fp32_router_gemm_entry.cu" - "csrc/libtorch_stable/fp32_router_gemm.cu") + "csrc/fp32_router_gemm_entry.cu" + "csrc/fp32_router_gemm.cu") set_gencode_flags_for_srcs( SRCS "${FP32_ROUTER_GEMM_SRCS}" CUDA_ARCHS "${FP32_ROUTER_GEMM_ARCHS}") @@ -799,11 +789,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a;" "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SCALED_MM_ARCHS) set(SCALED_MM_SM90_SRCS - "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu") + "csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu") set_gencode_flags_for_srcs( SRCS "${SCALED_MM_SM90_SRCS}" CUDA_ARCHS "${SCALED_MM_ARCHS}") @@ -834,9 +824,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) set(SCALED_MM_SM120_SRCS - "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu" + "csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu" ) set_gencode_flags_for_srcs( SRCS "${SCALED_MM_SM120_SRCS}" @@ -868,9 +858,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) set(SCALED_MM_SM100_SRCS - "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu" + "csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu" ) set_gencode_flags_for_srcs( SRCS "${SCALED_MM_SM100_SRCS}" @@ -901,7 +891,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # subtract out the archs that are already built for 3x list(REMOVE_ITEM SCALED_MM_2X_ARCHS ${SCALED_MM_3X_ARCHS}) if (SCALED_MM_2X_ARCHS) - set(SCALED_MM_C2X_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cu") + set(SCALED_MM_C2X_SRCS "csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cu") set_gencode_flags_for_srcs( SRCS "${SCALED_MM_C2X_SRCS}" CUDA_ARCHS "${SCALED_MM_2X_ARCHS}") @@ -927,7 +917,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # if it's possible to compile MoE kernels that use its output. cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a" "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS) - set(CUTLASS_MOE_SM90_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu") + set(CUTLASS_MOE_SM90_SRCS "csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu") set_gencode_flags_for_srcs( SRCS "${CUTLASS_MOE_SM90_SRCS}" CUDA_ARCHS "${SCALED_MM_ARCHS}") @@ -951,7 +941,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) - set(CUTLASS_MOE_SM100_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu") + set(CUTLASS_MOE_SM100_SRCS "csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu") set_gencode_flags_for_srcs( SRCS "${CUTLASS_MOE_SM100_SRCS}" CUDA_ARCHS "${SCALED_MM_ARCHS}") @@ -976,7 +966,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND CUTLASS_MOE_DATA_ARCHS) - set(CUTLASS_MOE_DATA_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu") + set(CUTLASS_MOE_DATA_SRCS "csrc/quantization/w8a8/cutlass/moe/moe_data.cu") set_gencode_flags_for_srcs( SRCS "${CUTLASS_MOE_DATA_SRCS}" CUDA_ARCHS "${CUTLASS_MOE_DATA_ARCHS}") @@ -1008,12 +998,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM120_ARCHS) set(FP4_SM120_SRCS - "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu" - "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu") + "csrc/quantization/fp4/nvfp4_quant_kernels.cu" + "csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu" + "csrc/quantization/fp4/nvfp4_experts_quant.cu" + "csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu" + "csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu" + "csrc/nvfp4_kv_cache_kernels.cu") set_gencode_flags_for_srcs( SRCS "${FP4_SM120_SRCS}" CUDA_ARCHS "${FP4_SM120_ARCHS}") @@ -1034,15 +1024,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM100_ARCHS) set(FP4_SM100_SRCS - "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu" - "csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu" - "csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu" - "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu" - "csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu") + "csrc/quantization/fp4/nvfp4_quant_kernels.cu" + "csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu" + "csrc/quantization/fp4/nvfp4_experts_quant.cu" + "csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu" + "csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu" + "csrc/quantization/fp4/mxfp4_experts_quant.cu" + "csrc/quantization/fp4/mxfp4_blockwise_moe_kernel.cu" + "csrc/nvfp4_kv_cache_kernels.cu" + "csrc/nvfp4_ds_mla_cache_kernels.cu") if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) message(STATUS "Building mxfp4_experts_quant unsupported stubs because CUDA compiler version is not >= 12.9 (found ${CMAKE_CUDA_COMPILER_VERSION}).") @@ -1066,9 +1056,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(W4A8_ARCHS "9.0a" "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND W4A8_ARCHS) set(W4A8_SRCS - "csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu" - "csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu" - "csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cu" + "csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu" + "csrc/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu" + "csrc/quantization/cutlass_w4a8/w4a8_utils.cu" ) set_gencode_flags_for_srcs( @@ -1101,7 +1091,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS) set(CUTLASS_MLA_SRCS - "csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu") + "csrc/attention/mla/sm100_cutlass_mla_kernel.cu") set_gencode_flags_for_srcs( SRCS "${CUTLASS_MLA_SRCS}" CUDA_ARCHS "${MLA_ARCHS}") @@ -1125,7 +1115,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(FUSED_KDA_DECODE_ARCHS) set(FUSED_KDA_DECODE_SRC - "csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu") + "csrc/kimi_k3/fused_kda_decode_kernel.cu") set_gencode_flags_for_srcs( SRCS "${FUSED_KDA_DECODE_SRC}" CUDA_ARCHS "${FUSED_KDA_DECODE_ARCHS}") @@ -1137,7 +1127,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(FUSED_GDN_DECODE_ARCHS) set(FUSED_GDN_DECODE_SRC - "csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu") + "csrc/gdn/fused_gdn_decode_kernel.cu") set_gencode_flags_for_srcs( SRCS "${FUSED_GDN_DECODE_SRC}" CUDA_ARCHS "${FUSED_GDN_DECODE_ARCHS}") @@ -1154,7 +1144,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(KIMI_K3_ATTN_RES_ARCHS) set(KIMI_K3_ATTN_RES_SRC - "csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu") + "csrc/kimi_k3/attn_res_kernel.cu") set_gencode_flags_for_srcs( SRCS "${KIMI_K3_ATTN_RES_SRC}" CUDA_ARCHS "${KIMI_K3_ATTN_RES_ARCHS}") @@ -1169,7 +1159,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # Hadacore kernels cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}") if(HADACORE_ARCHS) - set(HADACORE_SRCS "csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu") + set(HADACORE_SRCS "csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu") set_gencode_flags_for_srcs( SRCS "${HADACORE_SRCS}" CUDA_ARCHS "${HADACORE_ARCHS}") @@ -1193,7 +1183,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") list(FILTER FUSED_KDA_DECODE_HIP_ARCHS INCLUDE REGEX "gfx942|gfx950") if(FUSED_KDA_DECODE_HIP_ARCHS) set(FUSED_KDA_DECODE_HIP_SRC - "csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu") + "csrc/kimi_k3/fused_kda_decode_kernel_rocm.cu") list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_KDA_DECODE_HIP_SRC}") message(STATUS "Building fused KDA decode for HIP archs: ${FUSED_KDA_DECODE_HIP_ARCHS}") @@ -1203,7 +1193,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") list(FILTER FUSED_KDA_CHUNK_HIP_ARCHS INCLUDE REGEX "gfx950") if(FUSED_KDA_CHUNK_HIP_ARCHS) set(FUSED_KDA_CHUNK_HIP_SRC - "csrc/libtorch_stable/kimi_k3/fused_kda_chunk_kernel_rocm.cu") + "csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu") list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_KDA_CHUNK_HIP_SRC}") message(STATUS "Building fused KDA chunk for HIP archs: ${FUSED_KDA_CHUNK_HIP_ARCHS}") @@ -1289,21 +1279,21 @@ endif() # set(VLLM_MOE_EXT_SRC - "csrc/libtorch_stable/moe/torch_bindings.cpp" - "csrc/libtorch_stable/moe/moe_align_sum_kernels.cu" - "csrc/libtorch_stable/moe/topk_softmax_kernels.cu" - "csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu") + "csrc/moe/torch_bindings.cpp" + "csrc/moe/moe_align_sum_kernels.cu" + "csrc/moe/topk_softmax_kernels.cu" + "csrc/moe/topk_softplus_sqrt_kernels.cu") if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_MOE_EXT_SRC - "csrc/libtorch_stable/moe/moe_wna16.cu" - "csrc/libtorch_stable/moe/grouped_topk_kernels.cu") + "csrc/moe/moe_wna16.cu" + "csrc/moe/grouped_topk_kernels.cu") endif() if(VLLM_GPU_LANG STREQUAL "CUDA") set(MOE_PERMUTE_SRC - "csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu" - "csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu") + "csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu" + "csrc/moe/moe_permute_unpermute_op.cu") list(APPEND VLLM_MOE_EXT_SRC "${MOE_PERMUTE_SRC}") endif() @@ -1314,7 +1304,7 @@ set_gencode_flags_for_srcs( if(VLLM_GPU_LANG STREQUAL "CUDA") set(VLLM_MOE_WNA16_SRC - "csrc/libtorch_stable/moe/moe_wna16.cu") + "csrc/moe/moe_wna16.cu") set_gencode_flags_for_srcs( SRCS "${VLLM_MOE_WNA16_SRC}" @@ -1349,7 +1339,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") # preselected input type pairs and schedules. # Generate sources: set(MOE_MARLIN_GEN_SCRIPT - ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py) + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/marlin/generate_kernels.py) file(MD5 ${MOE_MARLIN_GEN_SCRIPT} MOE_MARLIN_GEN_SCRIPT_HASH) list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR) set(MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MOE_MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})") @@ -1363,6 +1353,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") COMMAND ${CMAKE_COMMAND} -E env PYTHONPATH=$ENV{PYTHONPATH} ${Python_EXECUTABLE} ${MOE_MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR} + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/moe/marlin_moe_wna16 RESULT_VARIABLE moe_marlin_generation_result OUTPUT_VARIABLE moe_marlin_generation_output OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/moe_marlin_generation.log @@ -1384,7 +1375,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") endif() if (MARLIN_MOE_ARCHS) - file(GLOB MARLIN_MOE_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_*.cu") + file(GLOB MARLIN_MOE_SRC "csrc/moe/marlin_moe_wna16/sm80_kernel_*.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_MOE_SRC}" CUDA_ARCHS "${MARLIN_MOE_ARCHS}") @@ -1396,7 +1387,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") endif() if (MARLIN_MOE_SM75_ARCHS) - file(GLOB MARLIN_MOE_SM75_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/sm75_kernel_*.cu") + file(GLOB MARLIN_MOE_SM75_SRC "csrc/moe/marlin_moe_wna16/sm75_kernel_*.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_MOE_SM75_SRC}" CUDA_ARCHS "${MARLIN_MOE_SM75_ARCHS}") @@ -1408,7 +1399,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") endif() if (MARLIN_MOE_FP8_ARCHS) - file(GLOB MARLIN_MOE_FP8_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/sm89_kernel_*.cu") + file(GLOB MARLIN_MOE_FP8_SRC "csrc/moe/marlin_moe_wna16/sm89_kernel_*.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_MOE_FP8_SRC}" CUDA_ARCHS "${MARLIN_MOE_FP8_ARCHS}") @@ -1419,7 +1410,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_MOE_EXT_SRC ${MARLIN_MOE_FP8_SRC}) endif() - set(MARLIN_MOE_OTHER_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu") + set(MARLIN_MOE_OTHER_SRC "csrc/moe/marlin_moe_wna16/ops.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_MOE_OTHER_SRC}" CUDA_ARCHS "${MARLIN_MOE_OTHER_ARCHS}") diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 730b963c2362..c4d111e77bfe 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -497,7 +497,7 @@ set(VLLM_EXT_SRC "csrc/cpu/mla_decode.cpp" "csrc/cpu/pos_encoding.cpp" "csrc/cpu/mamba_cpu.cpp" - "csrc/moe/dynamic_4bit_int_moe_cpu.cpp" + "csrc/cpu/dynamic_4bit_int_moe_cpu.cpp" "csrc/cpu/cpu_fused_moe.cpp" "csrc/cpu/cpu_attn.cpp" "csrc/cpu/cpu_isa.cpp" @@ -599,7 +599,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/layernorm.cpp" "csrc/cpu/mla_decode.cpp" "csrc/cpu/pos_encoding.cpp" - "csrc/moe/dynamic_4bit_int_moe_cpu.cpp") + "csrc/cpu/dynamic_4bit_int_moe_cpu.cpp") set(VLLM_EXT_SRC_AVX2 "csrc/cpu/sgl-kernels/fla.cpp" @@ -617,7 +617,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/layernorm.cpp" "csrc/cpu/mla_decode.cpp" "csrc/cpu/pos_encoding.cpp" - "csrc/moe/dynamic_4bit_int_moe_cpu.cpp") + "csrc/cpu/dynamic_4bit_int_moe_cpu.cpp") message(STATUS "CPU extension (AVX512F + BF16 + VNNI + AMX) source files: ${VLLM_EXT_SRC_AVX512} ${VLLM_EXT_SRC_SGL}") message(STATUS "CPU extension (AVX512F) source files: ${VLLM_EXT_SRC_AVX512}") diff --git a/cmake/hipify.py b/cmake/hipify.py index f4932260c9eb..2cf58a7b7d33 100755 --- a/cmake/hipify.py +++ b/cmake/hipify.py @@ -65,7 +65,7 @@ def _expected_hip_build_path(source_abs: str, output_directory: str) -> str: project_directory=args.project_dir, output_directory=args.output_dir, # Hipify resolves quoted includes next to the including file first; vLLM - # uses paths relative to csrc/ (e.g. "libtorch_stable/torch_utils.h" + # uses paths relative to csrc/ (e.g. "torch_utils.h" # from quantization/w8a8/fp8/*.cu). Without an include root here, those # headers are never found and are not hipified or rewritten in dependents. header_include_dirs=["."], diff --git a/csrc/libtorch_stable/activation_kernels.cu b/csrc/activation_kernels.cu similarity index 99% rename from csrc/libtorch_stable/activation_kernels.cu rename to csrc/activation_kernels.cu index a7ee87b35f4e..9e28c3343bd7 100644 --- a/csrc/libtorch_stable/activation_kernels.cu +++ b/csrc/activation_kernels.cu @@ -10,7 +10,7 @@ #include #include -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include "async_util.cuh" #include "cuda_vec_utils.cuh" #include "dispatch_utils.h" diff --git a/csrc/libtorch_stable/async_util.cuh b/csrc/async_util.cuh similarity index 100% rename from csrc/libtorch_stable/async_util.cuh rename to csrc/async_util.cuh diff --git a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu b/csrc/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu similarity index 100% rename from csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu rename to csrc/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu diff --git a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_common.cuh b/csrc/attention/dcp_utils/dcp_direct_common.cuh similarity index 100% rename from csrc/libtorch_stable/attention/dcp_utils/dcp_direct_common.cuh rename to csrc/attention/dcp_utils/dcp_direct_common.cuh diff --git a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu b/csrc/attention/dcp_utils/dcp_direct_kv_gather.cu similarity index 100% rename from csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu rename to csrc/attention/dcp_utils/dcp_direct_kv_gather.cu diff --git a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu b/csrc/attention/dcp_utils/dcp_direct_q_gather.cu similarity index 100% rename from csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu rename to csrc/attention/dcp_utils/dcp_direct_q_gather.cu diff --git a/csrc/libtorch_stable/attention/merge_attn_states.cu b/csrc/attention/merge_attn_states.cu similarity index 99% rename from csrc/libtorch_stable/attention/merge_attn_states.cu rename to csrc/attention/merge_attn_states.cu index cc89397f68a1..3be334fd64fc 100644 --- a/csrc/libtorch_stable/attention/merge_attn_states.cu +++ b/csrc/attention/merge_attn_states.cu @@ -6,9 +6,8 @@ #include "../dispatch_utils.h" #include -#include "../../attention/attention_dtypes.h" -#include "attention_utils.cuh" -#include "../../quantization/w8a8/fp8/common.cuh" +#include "core/vec_dtypes/vec_dtypes.h" +#include "quantization/w8a8/fp8/common.cuh" namespace vllm { diff --git a/csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp b/csrc/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp similarity index 100% rename from csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp rename to csrc/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp diff --git a/csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp b/csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp similarity index 100% rename from csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp rename to csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp diff --git a/csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp b/csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp similarity index 100% rename from csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp rename to csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp diff --git a/csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp b/csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp similarity index 100% rename from csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp rename to csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp diff --git a/csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu b/csrc/attention/mla/sm100_cutlass_mla_kernel.cu similarity index 99% rename from csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu rename to csrc/attention/mla/sm100_cutlass_mla_kernel.cu index 150e32462812..b04c014b9805 100644 --- a/csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu +++ b/csrc/attention/mla/sm100_cutlass_mla_kernel.cu @@ -18,7 +18,7 @@ limitations under the License. * Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929 * by Alcanderian JieXin Liang */ -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include diff --git a/csrc/cache.h b/csrc/cache.h deleted file mode 100644 index dbc2675fc70c..000000000000 --- a/csrc/cache.h +++ /dev/null @@ -1,93 +0,0 @@ -#pragma once - -#include -#include - -#include -#include - -void swap_blocks(torch::Tensor& src, torch::Tensor& dst, - int64_t block_size_in_bytes, - const torch::Tensor& block_mapping); - -void swap_blocks_batch(const torch::Tensor& src_ptrs, - const torch::Tensor& dst_ptrs, - const torch::Tensor& sizes, - bool is_src_access_order_any); - -void reshape_and_cache(torch::Tensor& key, torch::Tensor& value, - torch::Tensor& key_cache, torch::Tensor& value_cache, - torch::Tensor& slot_mapping, - const std::string& kv_cache_dtype, - torch::Tensor& k_scale, torch::Tensor& v_scale); - -void reshape_and_cache_flash(torch::Tensor& key, torch::Tensor& value, - torch::Tensor& key_cache, - torch::Tensor& value_cache, - torch::Tensor& slot_mapping, - const std::string& kv_cache_dtype, - torch::Tensor& k_scale, torch::Tensor& v_scale); - -void concat_and_cache_mla(torch::Tensor& kv_c, torch::Tensor& k_pe, - torch::Tensor& kv_cache, torch::Tensor& slot_mapping, - const std::string& kv_cache_dtype, - torch::Tensor& scale); - -// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla -void concat_and_cache_mla_rope_fused( - torch::Tensor& positions, torch::Tensor& q_pe, torch::Tensor& k_pe, - torch::Tensor& kv_c, torch::Tensor& rope_cos_sin_cache, bool rope_is_neox, - torch::Tensor& kv_cache_slot_mapping, torch::Tensor& kv_cache, - const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale); - -// Just for unittest -void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache, - const double scale, const std::string& kv_cache_dtype); - -void gather_and_maybe_dequant_cache( - torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] - torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::Tensor const& cu_seq_lens, // [BATCH+1] - torch::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] - int64_t num_tokens, const std::string& kv_cache_dtype, - torch::Tensor const& scale, - std::optional seq_starts = std::nullopt); - -// TODO(hc): cp_gather_cache need support scaled kvcahe in the future. -void cp_gather_cache( - torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] - torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::Tensor const& cu_seq_lens, // [BATCH+1] - int64_t batch_size, std::optional seq_starts = std::nullopt); - -// Gather and upconvert FP8 KV cache to BF16 workspace -void cp_gather_and_upconvert_fp8_kv_cache( - torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656] - torch::Tensor const& dst, // [TOT_TOKENS, 576] - torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::Tensor const& workspace_starts, // [BATCH] - int64_t batch_size, std::optional seq_starts = std::nullopt); - -// Indexer K quantization and cache function -void indexer_k_quant_and_cache( - torch::Tensor& k, // [num_tokens, head_dim] - torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride] - torch::Tensor& slot_mapping, // [num_tokens] - int64_t quant_block_size, // quantization block size - const std::string& scale_fmt); - -// Concatenate query nope and rope for MLA/DSA attention -void concat_mla_q( - torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] - torch::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] - torch::Tensor& q_out); // [num_tokens, num_heads, nope_dim + rope_dim] - -// Extract function to gather quantized K cache -void cp_gather_indexer_k_quant_cache( - const torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride] - torch::Tensor& dst_k, // [num_tokens, head_dim] - torch::Tensor& dst_scale, // [num_tokens, head_dim / quant_block_size * 4] - const torch::Tensor& block_table, // [batch_size, num_blocks] - const torch::Tensor& cu_seq_lens); // [batch_size + 1] diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/cache_kernels.cu similarity index 99% rename from csrc/libtorch_stable/cache_kernels.cu rename to csrc/cache_kernels.cu index 9130adf87d69..d1c1e9df9896 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/csrc/cache_kernels.cu @@ -1,8 +1,8 @@ #include "torch_utils.h" #include "dispatch_utils.h" -#include "../cuda_utils.h" -#include "../cuda_compat.h" +#include "core/utils.hpp" +#include "core/cuda_compat.h" #include "quantization/vectorization_utils.cuh" #include "concat_mla_q.cuh" @@ -12,9 +12,9 @@ #endif #ifdef USE_ROCM - #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" + #include "quantization/w8a8/fp8/amd/quant_utils.cuh" #else - #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" + #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" #endif #include @@ -712,8 +712,7 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel( } __syncthreads(); - for (int iter = 0; iter < cuda_utils::ceil_div(batch_size, int(blockDim.x)); - iter++) { + for (int iter = 0; iter < div_ceil(batch_size, int(blockDim.x)); iter++) { int tid = iter * blockDim.x + threadIdx.x; if (tid < batch_size) { const int seq_start = cu_seq_lens[tid]; @@ -1229,7 +1228,7 @@ __device__ __forceinline__ bool map_gather_page_task( const int32_t source_begin = seq_starts == nullptr ? 0 : seq_starts[req_id]; const int32_t first_block = source_begin / block_size; const int32_t num_pages = - cuda_utils::ceil_div(source_begin + seq_len, block_size) - first_block; + div_ceil(source_begin + seq_len, block_size) - first_block; if (relative_page < num_pages) { page.req_id = req_id; page.logical_block = first_block + relative_page; @@ -1408,8 +1407,7 @@ void gather_and_maybe_dequant_cache( const int64_t dst_entry_stride = dst.stride(0); const int32_t page_threads = num_tokens >= (1 << 20) ? 128 : 256; const int32_t required_blocks = - cuda_utils::ceil_div(static_cast(num_tokens), block_size) + - 2 * num_reqs; + div_ceil(static_cast(num_tokens), block_size) + 2 * num_reqs; const dim3 grid(required_blocks); const dim3 block(page_threads); @@ -1685,7 +1683,7 @@ void cp_gather_cache( const int32_t num_reqs = static_cast(batch_size); const int32_t required_blocks = - cuda_utils::ceil_div(total_tokens, block_size) + 2 * num_reqs; + div_ceil(total_tokens, block_size) + 2 * num_reqs; const dim3 grid(required_blocks); const dim3 block(256); @@ -1828,8 +1826,8 @@ void cp_gather_and_upconvert_fp8_kv_cache( const int32_t* seq_starts_ptr = seq_starts.has_value() ? seq_starts.value().const_data_ptr() : nullptr; - const int required_blocks = cuda_utils::ceil_div(total_tokens, block_size) + - 2 * static_cast(batch_size); + const int required_blocks = + div_ceil(total_tokens, block_size) + 2 * static_cast(batch_size); vllm::cp_gather_and_upconvert_fp8_kv_cache_page<<< required_blocks, block_size_threads, 0, stream>>>( src_ptr, reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()), diff --git a/csrc/libtorch_stable/cache_kernels_fused.cu b/csrc/cache_kernels_fused.cu similarity index 98% rename from csrc/libtorch_stable/cache_kernels_fused.cu rename to csrc/cache_kernels_fused.cu index 744e3a2dcdcd..4dc691636d55 100644 --- a/csrc/libtorch_stable/cache_kernels_fused.cu +++ b/csrc/cache_kernels_fused.cu @@ -1,13 +1,13 @@ #include "torch_utils.h" #include "dispatch_utils.h" -#include "../cuda_compat.h" +#include "core/cuda_compat.h" -#include "../quantization/w8a8/fp8/common.cuh" +#include "quantization/w8a8/fp8/common.cuh" #ifdef USE_ROCM - #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" + #include "quantization/w8a8/fp8/amd/quant_utils.cuh" #else - #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" + #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" #endif #ifdef USE_ROCM diff --git a/csrc/libtorch_stable/concat_mla_q.cuh b/csrc/concat_mla_q.cuh similarity index 100% rename from csrc/libtorch_stable/concat_mla_q.cuh rename to csrc/concat_mla_q.cuh diff --git a/csrc/libtorch_stable/cooperative_topk.cu b/csrc/cooperative_topk.cu similarity index 100% rename from csrc/libtorch_stable/cooperative_topk.cu rename to csrc/cooperative_topk.cu diff --git a/csrc/libtorch_stable/cooperative_topk.cuh b/csrc/cooperative_topk.cuh similarity index 100% rename from csrc/libtorch_stable/cooperative_topk.cuh rename to csrc/cooperative_topk.cuh diff --git a/csrc/core/batch_invariant.hpp b/csrc/core/batch_invariant.hpp deleted file mode 100644 index 8273bc74b1ef..000000000000 --- a/csrc/core/batch_invariant.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once -#include -#include - -namespace vllm { - -// vllm_is_batch_invariant(); returns true -// if env VLLM_BATCH_INVARIANT=1 -inline bool vllm_is_batch_invariant() { - static bool cached = []() { - std::string env_key = "VLLM_BATCH_INVARIANT"; - const char* val = std::getenv(env_key.c_str()); - return (val && std::atoi(val) != 0) ? 1 : 0; - }(); - return cached; -} - -} // namespace vllm diff --git a/csrc/cuda_compat.h b/csrc/core/cuda_compat.h similarity index 90% rename from csrc/cuda_compat.h rename to csrc/core/cuda_compat.h index d7d589db62cf..0eafeff6219b 100644 --- a/csrc/cuda_compat.h +++ b/csrc/core/cuda_compat.h @@ -54,12 +54,6 @@ struct Utils { __shfl_xor(var, lane_mask, width) #endif -#ifndef USE_ROCM - #define VLLM_SHFL_SYNC(var, src_lane) __shfl_sync(uint32_t(-1), var, src_lane) -#else - #define VLLM_SHFL_SYNC(var, src_lane) __shfl(var, src_lane) -#endif - #ifndef USE_ROCM #define VLLM_SHFL_DOWN_SYNC(var, lane_delta) \ __shfl_down_sync(uint32_t(-1), var, lane_delta) diff --git a/csrc/core/exception.hpp b/csrc/core/exception.hpp deleted file mode 100644 index f3b2ffaef6cc..000000000000 --- a/csrc/core/exception.hpp +++ /dev/null @@ -1,3 +0,0 @@ -#pragma once - -#define VLLM_IMPLIES(p, q) (!(p) || (q)) diff --git a/csrc/core/registration.h b/csrc/core/registration.h index 4d0ce1c572c1..e5396e9a8b13 100644 --- a/csrc/core/registration.h +++ b/csrc/core/registration.h @@ -12,11 +12,6 @@ // could be a macro instead of a literal token. #define TORCH_LIBRARY_EXPAND(NAME, MODULE) TORCH_LIBRARY(NAME, MODULE) -// A version of the TORCH_LIBRARY_IMPL macro that expands the NAME, i.e. so NAME -// could be a macro instead of a literal token. -#define TORCH_LIBRARY_IMPL_EXPAND(NAME, DEVICE, MODULE) \ - TORCH_LIBRARY_IMPL(NAME, DEVICE, MODULE) - // REGISTER_EXTENSION allows the shared library to be loaded and initialized // via python's import statement. #define REGISTER_EXTENSION(NAME) \ diff --git a/csrc/libtorch_stable/core/math.hpp b/csrc/core/utils.hpp similarity index 70% rename from csrc/libtorch_stable/core/math.hpp rename to csrc/core/utils.hpp index 6764e1fd6054..36b8c2e38e9a 100644 --- a/csrc/libtorch_stable/core/math.hpp +++ b/csrc/core/utils.hpp @@ -1,7 +1,8 @@ #pragma once #include -#include +#include +#include inline constexpr uint32_t next_pow_2(uint32_t const num) { if (num <= 1) return num; @@ -26,3 +27,16 @@ template inline constexpr T round_to_next_multiple_of(T a, T b) { return a % b == 0 ? a : ((a / b) + 1) * b; } + +namespace vllm { + +// Returns true if env VLLM_BATCH_INVARIANT=1 +inline bool vllm_is_batch_invariant() { + static bool cached = []() { + const char* val = std::getenv("VLLM_BATCH_INVARIANT"); + return val && std::atoi(val) != 0; + }(); + return cached; +} + +} // namespace vllm diff --git a/csrc/attention/dtype_bfloat16.cuh b/csrc/core/vec_dtypes/dtype_bfloat16.cuh similarity index 99% rename from csrc/attention/dtype_bfloat16.cuh rename to csrc/core/vec_dtypes/dtype_bfloat16.cuh index 97a25baa1fc0..5d01a51e76fe 100644 --- a/csrc/attention/dtype_bfloat16.cuh +++ b/csrc/core/vec_dtypes/dtype_bfloat16.cuh @@ -20,7 +20,7 @@ */ #pragma once -#include "attention_generic.cuh" +#include "dtype_generic.cuh" #include "dtype_float32.cuh" #ifndef USE_ROCM diff --git a/csrc/attention/dtype_float16.cuh b/csrc/core/vec_dtypes/dtype_float16.cuh similarity index 99% rename from csrc/attention/dtype_float16.cuh rename to csrc/core/vec_dtypes/dtype_float16.cuh index 3a1815f0ed4f..6883d521d2ac 100644 --- a/csrc/attention/dtype_float16.cuh +++ b/csrc/core/vec_dtypes/dtype_float16.cuh @@ -20,7 +20,7 @@ */ #pragma once -#include "attention_generic.cuh" +#include "dtype_generic.cuh" #include "dtype_float32.cuh" #ifdef USE_ROCM diff --git a/csrc/attention/dtype_float32.cuh b/csrc/core/vec_dtypes/dtype_float32.cuh similarity index 99% rename from csrc/attention/dtype_float32.cuh rename to csrc/core/vec_dtypes/dtype_float32.cuh index 7c6a686db3ba..33f78c492e61 100644 --- a/csrc/attention/dtype_float32.cuh +++ b/csrc/core/vec_dtypes/dtype_float32.cuh @@ -20,7 +20,7 @@ */ #pragma once -#include "attention_generic.cuh" +#include "dtype_generic.cuh" #include diff --git a/csrc/attention/dtype_fp8.cuh b/csrc/core/vec_dtypes/dtype_fp8.cuh similarity index 88% rename from csrc/attention/dtype_fp8.cuh rename to csrc/core/vec_dtypes/dtype_fp8.cuh index 3d56859d8fc5..6d3c35e978b3 100644 --- a/csrc/attention/dtype_fp8.cuh +++ b/csrc/core/vec_dtypes/dtype_fp8.cuh @@ -1,7 +1,8 @@ #pragma once -#include "attention_generic.cuh" -#include "torch_utils.h" +#include "dtype_generic.cuh" + +#include #include #ifdef ENABLE_FP8 @@ -31,7 +32,7 @@ inline Fp8KVCacheDataType get_fp8_kv_cache_data_type( } else if (dtype_str == "fp8_e5m2") { return Fp8KVCacheDataType::kFp8E5M2; } - TORCH_UTILS_CHECK(false, "Unsupported fp8 kv cache data type: ", dtype_str); + STD_TORCH_CHECK(false, "Unsupported fp8 kv cache data type: ", dtype_str); } // fp8 vector types for quantization of kv cache diff --git a/csrc/attention/attention_generic.cuh b/csrc/core/vec_dtypes/dtype_generic.cuh similarity index 100% rename from csrc/attention/attention_generic.cuh rename to csrc/core/vec_dtypes/dtype_generic.cuh diff --git a/csrc/attention/attention_dtypes.h b/csrc/core/vec_dtypes/vec_dtypes.h similarity index 79% rename from csrc/attention/attention_dtypes.h rename to csrc/core/vec_dtypes/vec_dtypes.h index 64f86381d9db..41a7b91ca45a 100644 --- a/csrc/attention/attention_dtypes.h +++ b/csrc/core/vec_dtypes/vec_dtypes.h @@ -1,6 +1,6 @@ #pragma once -#include "attention_generic.cuh" +#include "dtype_generic.cuh" #include "dtype_float16.cuh" #include "dtype_float32.cuh" #include "dtype_bfloat16.cuh" diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index 0666ee413f13..0290c7ab60b3 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -14,7 +14,7 @@ namespace cpu_attention { enum class ISA { AMX, VEC, VEC16, NEON, VXE, RVV, VSX, AMX_FP8 }; -// Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly. +// Mirrors csrc/core/vec_dtypes/dtype_fp8.cuh Fp8KVCacheDataType exactly. enum class Fp8KVCacheDataType { kAuto = 0, kFp8E4M3 = 1, diff --git a/csrc/moe/dynamic_4bit_int_moe_cpu.cpp b/csrc/cpu/dynamic_4bit_int_moe_cpu.cpp similarity index 100% rename from csrc/moe/dynamic_4bit_int_moe_cpu.cpp rename to csrc/cpu/dynamic_4bit_int_moe_cpu.cpp diff --git a/csrc/cpu/ops.h b/csrc/cpu/ops.h new file mode 100644 index 000000000000..ead20b01ae67 --- /dev/null +++ b/csrc/cpu/ops.h @@ -0,0 +1,46 @@ +#pragma once + +#include +#include +#include +#include +#include + +void rms_norm(torch::Tensor& out, torch::Tensor& input, + std::optional weight, double epsilon); + +void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual, + std::optional weight, double epsilon); + +void rotary_embedding(torch::Tensor& positions, torch::Tensor& query, + std::optional key, int64_t head_size, + torch::Tensor& cos_sin_cache, bool is_neox, + int64_t rope_dim_offset, bool inverse); + +void silu_and_mul(torch::Tensor& out, torch::Tensor& input); + +void gelu_and_mul(torch::Tensor& out, torch::Tensor& input); + +void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input); + +void gelu_tanh(torch::Tensor& out, torch::Tensor& input); + +void gelu_new(torch::Tensor& out, torch::Tensor& input); + +void gelu_fast(torch::Tensor& out, torch::Tensor& input); + +void gelu_quick(torch::Tensor& out, torch::Tensor& input); + +void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, + torch::Tensor const& scale, + std::optional const& azp); + +void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, + torch::Tensor& scales, + std::optional const& azp); + +torch::Tensor dynamic_4bit_int_moe_cpu( + torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights, + torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size, + int64_t intermediate_size, int64_t group_size, + bool apply_router_weight_on_input, int64_t activation_kind); diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 2cb1be095d3a..f39cf31c1f17 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -1,4 +1,4 @@ -#include "cache.h" +#include #include "ops.h" #include "core/registration.h" diff --git a/csrc/libtorch_stable/cub_helpers.h b/csrc/cub_helpers.h similarity index 100% rename from csrc/libtorch_stable/cub_helpers.h rename to csrc/cub_helpers.h diff --git a/csrc/cuda_utils.h b/csrc/cuda_utils.h index 6e62ea208db8..bca7c75bc178 100644 --- a/csrc/cuda_utils.h +++ b/csrc/cuda_utils.h @@ -1,21 +1,8 @@ #pragma once +#include #include -#if defined(__HIPCC__) - #define HOST_DEVICE_INLINE __host__ __device__ - #define DEVICE_INLINE __device__ - #define HOST_INLINE __host__ -#elif defined(__CUDACC__) || defined(_NVHPC_CUDA) - #define HOST_DEVICE_INLINE __host__ __device__ __forceinline__ - #define DEVICE_INLINE __device__ __forceinline__ - #define HOST_INLINE __host__ __forceinline__ -#else - #define HOST_DEVICE_INLINE inline - #define DEVICE_INLINE inline - #define HOST_INLINE inline -#endif - #define CUDA_CHECK(cmd) \ do { \ cudaError_t e = cmd; \ @@ -29,13 +16,3 @@ int64_t get_device_attribute(int64_t attribute, int64_t device_id); int64_t get_max_shared_memory_per_block_device_attribute(int64_t device_id); - -namespace cuda_utils { - -template -HOST_DEVICE_INLINE constexpr std::enable_if_t, T> -ceil_div(T a, T b) { - return (a + b - 1) / b; -} - -}; // namespace cuda_utils \ No newline at end of file diff --git a/csrc/libtorch_stable/cuda_utils_kernels.cu b/csrc/cuda_utils_kernels.cu similarity index 100% rename from csrc/libtorch_stable/cuda_utils_kernels.cu rename to csrc/cuda_utils_kernels.cu diff --git a/csrc/libtorch_stable/cuda_vec_utils.cuh b/csrc/cuda_vec_utils.cuh similarity index 100% rename from csrc/libtorch_stable/cuda_vec_utils.cuh rename to csrc/cuda_vec_utils.cuh diff --git a/csrc/libtorch_stable/cuda_view.cu b/csrc/cuda_view.cu similarity index 100% rename from csrc/libtorch_stable/cuda_view.cu rename to csrc/cuda_view.cu diff --git a/csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu b/csrc/custom_all_gather_reduce_scatter.cu similarity index 100% rename from csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu rename to csrc/custom_all_gather_reduce_scatter.cu diff --git a/csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp b/csrc/custom_all_gather_reduce_scatter_ops.cpp similarity index 100% rename from csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp rename to csrc/custom_all_gather_reduce_scatter_ops.cpp diff --git a/csrc/libtorch_stable/custom_all_reduce.cu b/csrc/custom_all_reduce.cu similarity index 100% rename from csrc/libtorch_stable/custom_all_reduce.cu rename to csrc/custom_all_reduce.cu diff --git a/csrc/custom_collective_common.cuh b/csrc/custom_collective_common.cuh index 353dcd07e8b4..8c6a98149d0f 100644 --- a/csrc/custom_collective_common.cuh +++ b/csrc/custom_collective_common.cuh @@ -216,21 +216,6 @@ DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; } -template -DINLINE void barrier_at_start_release(const RankSignals& sg, Signal* self_sg, - int rank) { - __syncthreads(); - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank]; - auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x]; - st_flag_release(peer_counter_ptr, flag); - while (ld_flag_acquire(self_counter_ptr) != flag); - } - __syncthreads(); - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - // This function is meant to be used as the second or the final // synchronization barrier in the all reduce kernel. If it's the final // synchronization barrier, we don't need to make any visibility guarantees @@ -279,22 +264,6 @@ DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; } -template -DINLINE void barrier_at_start_release(const RankSignals& sg, Signal* self_sg, - int rank) { - __syncthreads(); - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank], - flag, __ATOMIC_RELEASE, __MEMORY_SCOPE_SYSTEM); - while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x], - __ATOMIC_ACQUIRE, - __MEMORY_SCOPE_DEVICE) < flag); - } - __syncthreads(); - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - template DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) { __syncthreads(); diff --git a/csrc/libtorch_stable/cutlass_extensions/common.hpp b/csrc/cutlass_extensions/common.hpp similarity index 89% rename from csrc/libtorch_stable/cutlass_extensions/common.hpp rename to csrc/cutlass_extensions/common.hpp index ca96b0ef3fef..0a6fbd838411 100644 --- a/csrc/libtorch_stable/cutlass_extensions/common.hpp +++ b/csrc/cutlass_extensions/common.hpp @@ -25,7 +25,15 @@ inline int get_cuda_max_shared_memory_per_block_opt_in(int const device) { return max_shared_mem_per_block_opt_in; } -int32_t get_sm_version_num(); +inline int32_t get_sm_version_num() { + int32_t major_capability, minor_capability; + cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor, + 0); + cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor, + 0); + int32_t version_num = major_capability * 10 + minor_capability; + return version_num; +} /** * A wrapper for a kernel that is used to guard against compilation on @@ -110,21 +118,6 @@ struct enable_sm100_to_sm120 : Kernel { } }; -template -struct enable_sm120_only : Kernel { - template - CUTLASS_DEVICE void operator()(Args&&... args) { -#if defined __CUDA_ARCH__ - #if __CUDA_ARCH__ == 1200 - Kernel::operator()(std::forward(args)...); - #else - printf("This kernel only supports sm120a.\n"); - asm("trap;"); - #endif -#endif - } -}; - // SM12x family includes SM120 (RTX 5090) and SM121 (DGX Spark GB10) template struct enable_sm120_family : Kernel { diff --git a/csrc/libtorch_stable/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp b/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp similarity index 100% rename from csrc/libtorch_stable/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp rename to csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp diff --git a/csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp b/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp similarity index 100% rename from csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp rename to csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp diff --git a/csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp b/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp similarity index 93% rename from csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp rename to csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp index c2ddcea6d027..30f4a547f875 100644 --- a/csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp +++ b/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp @@ -1,15 +1,9 @@ #pragma once -#include "cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp" -#include "cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp" +#include "broadcast_load_epilogue_c3x.hpp" +#include "broadcast_load_epilogue_array_c3x.hpp" -// This header is shared by both _C (unstable ABI) and _C_stable_libtorch -// (stable ABI) targets. When compiled under the stable ABI target, -// TORCH_TARGET_VERSION is defined and Tensor is unavailable, so we -// use torch::stable::Tensor instead. -#ifdef TORCH_TARGET_VERSION - #include -#endif +#include /* This file defines custom epilogues for fusing channel scales, token scales, @@ -23,38 +17,10 @@ namespace vllm::c3x { -#ifdef TORCH_TARGET_VERSION using TensorType = torch::stable::Tensor; -#else -using TensorType = torch::Tensor; -#endif using namespace cute; -template -struct identity { - CUTLASS_HOST_DEVICE - T operator()(T lhs) const { return lhs; } -}; - -template -struct TrivialEpilogue { - private: - using Accum = cutlass::epilogue::fusion::Sm90AccFetch; - using Compute = cutlass::epilogue::fusion::Sm90Compute< - cutlass::epilogue::thread::Identity, ElementD, ElementAcc, - cutlass::FloatRoundStyle::round_to_nearest>; - - public: - using EVTCompute = cutlass::epilogue::fusion::Sm90EVT; - using ArgumentType = typename EVTCompute::Arguments; - - template - static ArgumentType prepare_args(Args... args) { - return {}; - } -}; - /* * This class provides the common load descriptors for the * ScaledEpilogue[...] classes diff --git a/csrc/libtorch_stable/cutlass_extensions/torch_utils.hpp b/csrc/cutlass_extensions/torch_utils.hpp similarity index 76% rename from csrc/libtorch_stable/cutlass_extensions/torch_utils.hpp rename to csrc/cutlass_extensions/torch_utils.hpp index 5f973033fd72..c5b9d9591e6d 100644 --- a/csrc/libtorch_stable/cutlass_extensions/torch_utils.hpp +++ b/csrc/cutlass_extensions/torch_utils.hpp @@ -1,19 +1,12 @@ #pragma once -#include "torch_utils.h" - -// This header is shared between _C (unstable ABI, used by machete) and -// _C_stable_libtorch (stable ABI, used by W4A8/sparse). TORCH_TARGET_VERSION -// is defined only for the stable target, so we switch includes and types -// accordingly. TorchTensor (not Tensor) avoids ambiguity with cute::Tensor. -#ifdef TORCH_TARGET_VERSION - #include - #include - #include +#include +#include +#include +#include + +// TorchTensor (not Tensor) avoids ambiguity with cute::Tensor. using TorchTensor = torch::stable::Tensor; -#else -using TorchTensor = torch::Tensor; -#endif #include "cute/layout.hpp" #include "cutlass/layout/matrix.h" @@ -70,15 +63,15 @@ CUTE_HOST_DEVICE constexpr auto make_shape_from_idx(F&& f) { template static inline auto make_cute_layout(TorchTensor const& tensor, std::string_view name = "tensor") { - TORCH_UTILS_CHECK(tensor.dim() <= rank(Stride{})); + STD_TORCH_CHECK(tensor.dim() <= rank(Stride{})); auto stride = cute::transform_with_idx(Stride{}, [&](auto const& stride_ele, auto const& idx) { using StrideEle = std::decay_t; if (idx < tensor.dim()) { if constexpr (cute::is_static_v) { - TORCH_UTILS_CHECK(StrideEle::value == tensor.stride(idx), "Expected ", - name, ".stride(", idx, ") to be ", StrideEle::value); + STD_TORCH_CHECK(StrideEle::value == tensor.stride(idx), "Expected ", + name, ".stride(", idx, ") to be ", StrideEle::value); return StrideEle{}; } else { if (tensor.size(idx) == 1) { @@ -122,29 +115,7 @@ static inline auto maybe_make_cute_layout( } // -// Torch Type to Cutlass Type (equivalent_cutlass_type) -// - -template -struct equivalent_cutlass_type { - using type = T; -}; - -template -using equivalent_cutlass_type_t = typename equivalent_cutlass_type::type; - -template <> -struct equivalent_cutlass_type { - using type = cutlass::half_t; -}; - -template <> -struct equivalent_cutlass_type { - using type = cutlass::bfloat16_t; -}; - -// -// equivalent_scalar_t (basically inverse of equivalent_cutlass_type) +// Cutlass Type to Torch Type (equivalent_scalar_type) // // Return a `torch::headeronly::CppTypeToScalarType` compatible type, i.e. diff --git a/csrc/libtorch_stable/cutlass_extensions/vllm_collective_builder.cuh b/csrc/cutlass_extensions/vllm_collective_builder.cuh similarity index 100% rename from csrc/libtorch_stable/cutlass_extensions/vllm_collective_builder.cuh rename to csrc/cutlass_extensions/vllm_collective_builder.cuh diff --git a/csrc/libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh b/csrc/cutlass_extensions/vllm_numeric_conversion.cuh similarity index 99% rename from csrc/libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh rename to csrc/cutlass_extensions/vllm_numeric_conversion.cuh index 90f226cf64c0..0f74053fd29f 100644 --- a/csrc/libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh +++ b/csrc/cutlass_extensions/vllm_numeric_conversion.cuh @@ -1,9 +1,9 @@ #pragma once #include "cutlass/numeric_conversion.h" -#include "cutlass_extensions/vllm_custom_types.cuh" -#include "cutlass_extensions/cute_utils.cuh" -#include "cutlass_extensions/vllm_type_utils.cuh" +#include "vllm_custom_types.cuh" +#include "cute_utils.cuh" +#include "vllm_type_utils.cuh" // this file extends: // https://github.com/NVIDIA/cutlass/blob/cutlass-3.5.0/include/cutlass/numeric_conversion.h diff --git a/csrc/cutlass_extensions/vllm_type_utils.cuh b/csrc/cutlass_extensions/vllm_type_utils.cuh index 500ed508c830..ba79dddf6f42 100644 --- a/csrc/cutlass_extensions/vllm_type_utils.cuh +++ b/csrc/cutlass_extensions/vllm_type_utils.cuh @@ -2,7 +2,7 @@ #include "cutlass/half.h" #include "cuda_bf16.h" -#include "cutlass_extensions/vllm_custom_types.cuh" +#include "vllm_custom_types.cuh" namespace cutlass { diff --git a/csrc/dispatch_utils.h b/csrc/dispatch_utils.h index de0c505b7a62..aa125a08fd6c 100644 --- a/csrc/dispatch_utils.h +++ b/csrc/dispatch_utils.h @@ -1,158 +1,188 @@ /* - * Adapted from - * https://github.com/pytorch/pytorch/blob/v2.0.1/aten/src/ATen/Dispatch.h + * Stable ABI compatible dispatch utilities for vLLM. + * Adapted from dispatch_utils.h to use PyTorch's header-only (THO_*) macros + * instead of the ATen (AT_*) macros. + * + * These macros use: + * - THO_DISPATCH_SWITCH instead of AT_DISPATCH_SWITCH + * - THO_DISPATCH_CASE instead of AT_DISPATCH_CASE + * - torch::headeronly::ScalarType instead of at::ScalarType + * + * Add more macros here as needed when migrating additional kernels. */ #pragma once -#include +#include +#include +#include // Need a special dispatch case macro since we will nest the FP8 dispatch. // Instead of the usual 'scalar_t', this names the dispatched type 'fp8_t'. -#define AT_DISPATCH_FP8_CASE(enum_type, ...) \ - AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__) - -#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) +#define VLLM_STABLE_DISPATCH_FP8_CASE(enum_type, ...) \ + THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__) + +// Same idea, for dispatching on an int32/int64 index tensor (e.g. topk_ids) +// nested inside a value-type dispatch. Named 'idx_t' instead of 'scalar_t'. +#define VLLM_STABLE_DISPATCH_IDX_CASE(enum_type, ...) \ + THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, idx_t, __VA_ARGS__) + +#define VLLM_STABLE_DISPATCH_CASE_IDX_TYPES(...) \ + VLLM_STABLE_DISPATCH_IDX_CASE(torch::headeronly::ScalarType::Int, \ + __VA_ARGS__) \ + VLLM_STABLE_DISPATCH_IDX_CASE(torch::headeronly::ScalarType::Long, \ + __VA_ARGS__) + +#define VLLM_STABLE_DISPATCH_IDX_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_IDX_TYPES(__VA_ARGS__)) + +#define VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(...) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__) + +#define VLLM_STABLE_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) + +#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(...) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Byte, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Short, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Int, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Long, __VA_ARGS__) + +#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(...) \ + VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt16, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt32, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt64, __VA_ARGS__) + +#define VLLM_STABLE_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__)) + +#define VLLM_STABLE_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH( \ + TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__)) + +// FP8 type dispatch - ROCm uses FNUZ format, CUDA uses OCP format +#ifdef USE_ROCM + #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...) \ + VLLM_STABLE_DISPATCH_FP8_CASE( \ + torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ + VLLM_STABLE_DISPATCH_FP8_CASE( \ + torch::headeronly::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) +#else + #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...) \ + VLLM_STABLE_DISPATCH_FP8_CASE( \ + torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__) +#endif -#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) +// When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'. +// See VLLM_STABLE_DISPATCH_FP8_CASE above. +#define VLLM_STABLE_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__)) -#define VLLM_DISPATCH_CASE_HALF_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) +// Half types dispatch (Half + BFloat16) +#define VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(...) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__) -#define VLLM_DISPATCH_HALF_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_HALF_TYPES(__VA_ARGS__)) +#define VLLM_STABLE_DISPATCH_HALF_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(__VA_ARGS__)) -// ROCm devices might use either fn or fnuz, so set up dispatch table for both. -// A host-based check at runtime will create a preferred FP8 type for ROCm -// such that the correct kernel is dispatched. +// Quant type dispatch (FP8 + INT8) #ifdef USE_ROCM - #define VLLM_DISPATCH_CASE_FP8_TYPES(...) \ - AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ - AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) - - #define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) + #define VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(...) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, \ + __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fnuz, \ + __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) #else - #define VLLM_DISPATCH_CASE_FP8_TYPES(...) \ - AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) - - #define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) + #define VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(...) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, \ + __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) #endif -// When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'. -// See AT_DISPATCH_FP8_CASE above. -#define VLLM_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__)) - -#define VLLM_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__)) - -#define VLLM_DISPATCH_CASE_FLOATING_AND_BYTE_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) - -#define VLLM_DISPATCH_FLOATING_AND_BYTE_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_DISPATCH_CASE_FLOATING_AND_BYTE_TYPES(__VA_ARGS__)) - -#define VLLM_DISPATCH_CASE_INTEGRAL_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Short, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__) - -#define VLLM_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Short, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::UInt16, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::UInt32, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::UInt64, __VA_ARGS__) - -#define VLLM_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__)) - -#define VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH( \ - TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__)) - -#define VLLM_DISPATCH_VEC_SIZE(VEC_SIZE, ...) \ - switch (VEC_SIZE) { \ - case 16: { \ - constexpr int vec_size = 16; \ - __VA_ARGS__(); \ - break; \ - } \ - case 8: { \ - constexpr int vec_size = 8; \ - __VA_ARGS__(); \ - break; \ - } \ - case 4: { \ - constexpr int vec_size = 4; \ - __VA_ARGS__(); \ - break; \ - } \ - case 2: { \ - constexpr int vec_size = 2; \ - __VA_ARGS__(); \ - break; \ - } \ - default: { \ - constexpr int vec_size = 1; \ - __VA_ARGS__(); \ - break; \ - } \ +#define VLLM_STABLE_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__)) + +// Group size dispatch (pure C++ if/else, no ATen dependency) +#define VLLM_STABLE_DISPATCH_GROUP_SIZE(group_size, const_group_size, ...) \ + if (group_size == 128) { \ + constexpr int const_group_size = 128; \ + __VA_ARGS__(); \ + } else if (group_size == 64) { \ + constexpr int const_group_size = 64; \ + __VA_ARGS__(); \ } -#define VLLM_DISPATCH_BOOL(expr, const_expr, ...) \ - if (expr) { \ - constexpr bool const_expr = true; \ - __VA_ARGS__(); \ - } else { \ - constexpr bool const_expr = false; \ - __VA_ARGS__(); \ +// Boolean dispatch +#define VLLM_STABLE_DISPATCH_BOOL(expr, const_expr, ...) \ + if (expr) { \ + constexpr bool const_expr = true; \ + __VA_ARGS__(); \ + } else { \ + constexpr bool const_expr = false; \ + __VA_ARGS__(); \ } -#define VLLM_DISPATCH_GROUP_SIZE(group_size, const_group_size, ...) \ - if (group_size == 128) { \ - constexpr int const_group_size = 128; \ - __VA_ARGS__(); \ - } else if (group_size == 64) { \ - constexpr int const_group_size = 64; \ - __VA_ARGS__(); \ +// Vec size dispatch (pure C++ switch, no ATen dependency) +#define VLLM_STABLE_DISPATCH_VEC_SIZE(VEC_SIZE, ...) \ + switch (VEC_SIZE) { \ + case 16: { \ + constexpr int vec_size = 16; \ + __VA_ARGS__(); \ + break; \ + } \ + case 8: { \ + constexpr int vec_size = 8; \ + __VA_ARGS__(); \ + break; \ + } \ + case 4: { \ + constexpr int vec_size = 4; \ + __VA_ARGS__(); \ + break; \ + } \ + case 2: { \ + constexpr int vec_size = 2; \ + __VA_ARGS__(); \ + break; \ + } \ + default: { \ + constexpr int vec_size = 1; \ + __VA_ARGS__(); \ + break; \ + } \ } -#define VLLM_DISPATCH_RANK234(NUM_DIMS, ...) \ - switch (NUM_DIMS) { \ - case 2: { \ - constexpr int tensor_rank = 2; \ - __VA_ARGS__(); \ - break; \ - } \ - case 3: { \ - constexpr int tensor_rank = 3; \ - __VA_ARGS__(); \ - break; \ - } \ - case 4: { \ - constexpr int tensor_rank = 4; \ - __VA_ARGS__(); \ - break; \ - } \ - default: \ - TORCH_CHECK(false, "Expects rank 2, 3 or 4 tensors but got ", NUM_DIMS); \ +// Tensor rank dispatch (2D, 3D, 4D) +#define VLLM_STABLE_DISPATCH_RANK234(NUM_DIMS, ...) \ + switch (NUM_DIMS) { \ + case 2: { \ + constexpr int tensor_rank = 2; \ + __VA_ARGS__(); \ + break; \ + } \ + case 3: { \ + constexpr int tensor_rank = 3; \ + __VA_ARGS__(); \ + break; \ + } \ + case 4: { \ + constexpr int tensor_rank = 4; \ + __VA_ARGS__(); \ + break; \ + } \ + default: \ + STD_TORCH_CHECK( \ + false, "Expects rank 2, 3 or 4 tensors but got unsupported rank"); \ } diff --git a/csrc/libtorch_stable/dsv3_fused_a_gemm.cu b/csrc/dsv3_fused_a_gemm.cu similarity index 99% rename from csrc/libtorch_stable/dsv3_fused_a_gemm.cu rename to csrc/dsv3_fused_a_gemm.cu index 4250683b9158..c0989f8420da 100644 --- a/csrc/libtorch_stable/dsv3_fused_a_gemm.cu +++ b/csrc/dsv3_fused_a_gemm.cu @@ -25,7 +25,7 @@ #include #include "core/registration.h" -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/fp32_router_gemm.cu b/csrc/fp32_router_gemm.cu similarity index 100% rename from csrc/libtorch_stable/fp32_router_gemm.cu rename to csrc/fp32_router_gemm.cu diff --git a/csrc/libtorch_stable/fp32_router_gemm_entry.cu b/csrc/fp32_router_gemm_entry.cu similarity index 99% rename from csrc/libtorch_stable/fp32_router_gemm_entry.cu rename to csrc/fp32_router_gemm_entry.cu index fc09193c8643..e55ac7e72104 100644 --- a/csrc/libtorch_stable/fp32_router_gemm_entry.cu +++ b/csrc/fp32_router_gemm_entry.cu @@ -6,7 +6,7 @@ #include #include "core/registration.h" -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu similarity index 99% rename from csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu rename to csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index 5a5caacce7a5..4eb8d7756b19 100644 --- a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -47,7 +47,7 @@ #include #include -#include "cuda_compat.h" +#include "core/cuda_compat.h" #include "dispatch_utils.h" #include "type_convert.cuh" diff --git a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu b/csrc/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu similarity index 99% rename from csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu rename to csrc/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu index 3892790c4040..dd88645904d2 100644 --- a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu +++ b/csrc/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu @@ -57,16 +57,16 @@ #include #include -#include "cuda_compat.h" +#include "core/cuda_compat.h" #include "dispatch_utils.h" #include "type_convert.cuh" #ifndef USE_ROCM #include - #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" + #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" #else #include - #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" + #include "quantization/w8a8/fp8/amd/quant_utils.cuh" #endif #include #include @@ -835,7 +835,7 @@ static void launchPdlSlots(KernelT kernel, int num_tokens, int slots_per_token, // clang-format off // hipify's CUDA->HIP regex catastrophically backtracks on "> > >"; keep the // launch closer as ">>>". clang-format would otherwise re-split it (it does - // not parse the CUDA launch syntax in this libtorch_stable file). + // not parse the CUDA launch syntax in this stable-ABI file). kernel<<>>(args...); // clang-format on #endif diff --git a/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu b/csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu similarity index 99% rename from csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu rename to csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu index 0a45af8ade2b..09484daf3315 100644 --- a/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu +++ b/csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu @@ -54,15 +54,15 @@ #include "torch_utils.h" -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include "type_convert.cuh" -#include "../attention/dtype_fp8.cuh" +#include "core/vec_dtypes/dtype_fp8.cuh" #include "dispatch_utils.h" #ifdef USE_ROCM - #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" + #include "quantization/w8a8/fp8/amd/quant_utils.cuh" #else - #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" + #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" #endif // Direct float -> E4M3 FP8 conversion for the indexer Q / index-K outputs. diff --git a/csrc/libtorch_stable/fused_qknorm_rope_kernel.cu b/csrc/fused_qknorm_rope_kernel.cu similarity index 99% rename from csrc/libtorch_stable/fused_qknorm_rope_kernel.cu rename to csrc/fused_qknorm_rope_kernel.cu index 01bad79cd8e2..a75ede86fc37 100644 --- a/csrc/libtorch_stable/fused_qknorm_rope_kernel.cu +++ b/csrc/fused_qknorm_rope_kernel.cu @@ -21,7 +21,7 @@ #include "torch_utils.h" #include "async_util.cuh" -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include "type_convert.cuh" #include "dispatch_utils.h" diff --git a/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu b/csrc/gdn/fused_gdn_decode_kernel.cu similarity index 99% rename from csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu rename to csrc/gdn/fused_gdn_decode_kernel.cu index 1229404dc634..5e606ffd7fa1 100644 --- a/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu +++ b/csrc/gdn/fused_gdn_decode_kernel.cu @@ -10,7 +10,7 @@ #include #include "../torch_utils.h" -#include "../../cuda_compat.h" +#include "core/cuda_compat.h" namespace { diff --git a/csrc/libtorch_stable/hisparse_kernels.cu b/csrc/hisparse_kernels.cu similarity index 99% rename from csrc/libtorch_stable/hisparse_kernels.cu rename to csrc/hisparse_kernels.cu index 5c7621c2f598..da3857d1e973 100644 --- a/csrc/libtorch_stable/hisparse_kernels.cu +++ b/csrc/hisparse_kernels.cu @@ -17,7 +17,6 @@ #include "torch_utils.h" #include "ops.h" -#include "../cuda_utils.h" #include #include diff --git a/csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu b/csrc/kimi_k3/attn_res_kernel.cu similarity index 100% rename from csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu rename to csrc/kimi_k3/attn_res_kernel.cu diff --git a/csrc/libtorch_stable/kimi_k3/fused_kda_chunk_kernel_rocm.cu b/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu similarity index 100% rename from csrc/libtorch_stable/kimi_k3/fused_kda_chunk_kernel_rocm.cu rename to csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu diff --git a/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu b/csrc/kimi_k3/fused_kda_decode_kernel.cu similarity index 99% rename from csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu rename to csrc/kimi_k3/fused_kda_decode_kernel.cu index 00e5ce794d68..c5214ff48671 100644 --- a/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu +++ b/csrc/kimi_k3/fused_kda_decode_kernel.cu @@ -9,7 +9,7 @@ #include #include "../torch_utils.h" -#include "../../cuda_compat.h" +#include "core/cuda_compat.h" namespace { diff --git a/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu b/csrc/kimi_k3/fused_kda_decode_kernel_rocm.cu similarity index 100% rename from csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu rename to csrc/kimi_k3/fused_kda_decode_kernel_rocm.cu diff --git a/csrc/libtorch_stable/launch_bounds_utils.h b/csrc/launch_bounds_utils.h similarity index 100% rename from csrc/libtorch_stable/launch_bounds_utils.h rename to csrc/launch_bounds_utils.h diff --git a/csrc/libtorch_stable/layernorm_kernels.cu b/csrc/layernorm_kernels.cu similarity index 99% rename from csrc/libtorch_stable/layernorm_kernels.cu rename to csrc/layernorm_kernels.cu index b342c59f1801..8d6741ddac48 100644 --- a/csrc/libtorch_stable/layernorm_kernels.cu +++ b/csrc/layernorm_kernels.cu @@ -3,7 +3,7 @@ #include "torch_utils.h" #include "cub_helpers.h" -#include "../core/batch_invariant.hpp" +#include "core/utils.hpp" #include "type_convert.cuh" #include "dispatch_utils.h" #include "quantization/vectorization_utils.cuh" diff --git a/csrc/libtorch_stable/layernorm_quant_kernels.cu b/csrc/layernorm_quant_kernels.cu similarity index 99% rename from csrc/libtorch_stable/layernorm_quant_kernels.cu rename to csrc/layernorm_quant_kernels.cu index f43be531de03..7f8fb9e51007 100644 --- a/csrc/libtorch_stable/layernorm_quant_kernels.cu +++ b/csrc/layernorm_quant_kernels.cu @@ -10,8 +10,8 @@ #include "torch_utils.h" #include "cub_helpers.h" -#include "../core/batch_invariant.hpp" -#include "../quantization/w8a8/fp8/common.cuh" +#include "core/utils.hpp" +#include "quantization/w8a8/fp8/common.cuh" #include "type_convert.cuh" #include "dispatch_utils.h" #include "quantization/vectorization_utils.cuh" diff --git a/csrc/libtorch_stable/attention/attention_utils.cuh b/csrc/libtorch_stable/attention/attention_utils.cuh deleted file mode 100644 index 29783cb0e5f3..000000000000 --- a/csrc/libtorch_stable/attention/attention_utils.cuh +++ /dev/null @@ -1,57 +0,0 @@ -/* - * Adapted from - * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp - * Copyright (c) 2023, The vLLM team. - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include "../../cuda_compat.h" -#include "../../attention/attention_dtypes.h" - -#include -#include - -namespace vllm { - -// Q*K^T operation. -template -inline __device__ float qk_dot_(const Vec (&q)[N], const Vec (&k)[N]) { - using A_vec = typename FloatVec::Type; - // Compute the parallel products for Q*K^T (treat vector lanes separately). - A_vec qk_vec = mul(q[0], k[0]); -#pragma unroll - for (int ii = 1; ii < N; ++ii) { - qk_vec = vllm::fma(q[ii], k[ii], qk_vec); - } - - // Finalize the reduction across lanes. - float qk = sum(qk_vec); -#pragma unroll - for (int mask = THREAD_GROUP_SIZE / 2; mask >= 1; mask /= 2) { - qk += VLLM_SHFL_XOR_SYNC(qk, mask); - } - return qk; -} - -template -struct Qk_dot { - template - static inline __device__ float dot(const Vec (&q)[N], const Vec (&k)[N]) { - return qk_dot_(q, k); - } -}; - -} // namespace vllm diff --git a/csrc/libtorch_stable/cutlass_extensions/common.cpp b/csrc/libtorch_stable/cutlass_extensions/common.cpp deleted file mode 100644 index 5bc9463bfa60..000000000000 --- a/csrc/libtorch_stable/cutlass_extensions/common.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "common.hpp" - -int32_t get_sm_version_num() { - int32_t major_capability, minor_capability; - cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor, - 0); - cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor, - 0); - int32_t version_num = major_capability * 10 + minor_capability; - return version_num; -} \ No newline at end of file diff --git a/csrc/libtorch_stable/dispatch_utils.h b/csrc/libtorch_stable/dispatch_utils.h deleted file mode 100644 index aa125a08fd6c..000000000000 --- a/csrc/libtorch_stable/dispatch_utils.h +++ /dev/null @@ -1,188 +0,0 @@ -/* - * Stable ABI compatible dispatch utilities for vLLM. - * Adapted from dispatch_utils.h to use PyTorch's header-only (THO_*) macros - * instead of the ATen (AT_*) macros. - * - * These macros use: - * - THO_DISPATCH_SWITCH instead of AT_DISPATCH_SWITCH - * - THO_DISPATCH_CASE instead of AT_DISPATCH_CASE - * - torch::headeronly::ScalarType instead of at::ScalarType - * - * Add more macros here as needed when migrating additional kernels. - */ -#pragma once - -#include -#include -#include - -// Need a special dispatch case macro since we will nest the FP8 dispatch. -// Instead of the usual 'scalar_t', this names the dispatched type 'fp8_t'. -#define VLLM_STABLE_DISPATCH_FP8_CASE(enum_type, ...) \ - THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__) - -// Same idea, for dispatching on an int32/int64 index tensor (e.g. topk_ids) -// nested inside a value-type dispatch. Named 'idx_t' instead of 'scalar_t'. -#define VLLM_STABLE_DISPATCH_IDX_CASE(enum_type, ...) \ - THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, idx_t, __VA_ARGS__) - -#define VLLM_STABLE_DISPATCH_CASE_IDX_TYPES(...) \ - VLLM_STABLE_DISPATCH_IDX_CASE(torch::headeronly::ScalarType::Int, \ - __VA_ARGS__) \ - VLLM_STABLE_DISPATCH_IDX_CASE(torch::headeronly::ScalarType::Long, \ - __VA_ARGS__) - -#define VLLM_STABLE_DISPATCH_IDX_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_IDX_TYPES(__VA_ARGS__)) - -#define VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(...) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__) - -#define VLLM_STABLE_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) - -#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(...) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Byte, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Short, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Int, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Long, __VA_ARGS__) - -#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(...) \ - VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt16, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt32, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt64, __VA_ARGS__) - -#define VLLM_STABLE_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__)) - -#define VLLM_STABLE_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH( \ - TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__)) - -// FP8 type dispatch - ROCm uses FNUZ format, CUDA uses OCP format -#ifdef USE_ROCM - #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...) \ - VLLM_STABLE_DISPATCH_FP8_CASE( \ - torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ - VLLM_STABLE_DISPATCH_FP8_CASE( \ - torch::headeronly::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) -#else - #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...) \ - VLLM_STABLE_DISPATCH_FP8_CASE( \ - torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__) -#endif - -// When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'. -// See VLLM_STABLE_DISPATCH_FP8_CASE above. -#define VLLM_STABLE_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__)) - -// Half types dispatch (Half + BFloat16) -#define VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(...) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__) - -#define VLLM_STABLE_DISPATCH_HALF_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(__VA_ARGS__)) - -// Quant type dispatch (FP8 + INT8) -#ifdef USE_ROCM - #define VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(...) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, \ - __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fnuz, \ - __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) -#else - #define VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(...) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, \ - __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) -#endif - -#define VLLM_STABLE_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__)) - -// Group size dispatch (pure C++ if/else, no ATen dependency) -#define VLLM_STABLE_DISPATCH_GROUP_SIZE(group_size, const_group_size, ...) \ - if (group_size == 128) { \ - constexpr int const_group_size = 128; \ - __VA_ARGS__(); \ - } else if (group_size == 64) { \ - constexpr int const_group_size = 64; \ - __VA_ARGS__(); \ - } - -// Boolean dispatch -#define VLLM_STABLE_DISPATCH_BOOL(expr, const_expr, ...) \ - if (expr) { \ - constexpr bool const_expr = true; \ - __VA_ARGS__(); \ - } else { \ - constexpr bool const_expr = false; \ - __VA_ARGS__(); \ - } - -// Vec size dispatch (pure C++ switch, no ATen dependency) -#define VLLM_STABLE_DISPATCH_VEC_SIZE(VEC_SIZE, ...) \ - switch (VEC_SIZE) { \ - case 16: { \ - constexpr int vec_size = 16; \ - __VA_ARGS__(); \ - break; \ - } \ - case 8: { \ - constexpr int vec_size = 8; \ - __VA_ARGS__(); \ - break; \ - } \ - case 4: { \ - constexpr int vec_size = 4; \ - __VA_ARGS__(); \ - break; \ - } \ - case 2: { \ - constexpr int vec_size = 2; \ - __VA_ARGS__(); \ - break; \ - } \ - default: { \ - constexpr int vec_size = 1; \ - __VA_ARGS__(); \ - break; \ - } \ - } - -// Tensor rank dispatch (2D, 3D, 4D) -#define VLLM_STABLE_DISPATCH_RANK234(NUM_DIMS, ...) \ - switch (NUM_DIMS) { \ - case 2: { \ - constexpr int tensor_rank = 2; \ - __VA_ARGS__(); \ - break; \ - } \ - case 3: { \ - constexpr int tensor_rank = 3; \ - __VA_ARGS__(); \ - break; \ - } \ - case 4: { \ - constexpr int tensor_rank = 4; \ - __VA_ARGS__(); \ - break; \ - } \ - default: \ - STD_TORCH_CHECK( \ - false, "Expects rank 2, 3 or 4 tensors but got unsupported rank"); \ - } diff --git a/csrc/libtorch_stable/mamba/static_switch.h b/csrc/libtorch_stable/mamba/static_switch.h deleted file mode 100644 index 840cb2374a2f..000000000000 --- a/csrc/libtorch_stable/mamba/static_switch.h +++ /dev/null @@ -1,28 +0,0 @@ -// Inspired by -// https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h -// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h - -// clang-format off -// adapted from https://github.com/state-spaces/mamba/blob/main/csrc/selective_scan/static_switch.h -#pragma once - -/// @param COND - a boolean expression to switch by -/// @param CONST_NAME - a name given for the constexpr bool variable. -/// @param ... - code to execute for true and false -/// -/// Usage: -/// ``` -/// BOOL_SWITCH(flag, BoolConst, [&] { -/// some_function(...); -/// }); -/// ``` -#define BOOL_SWITCH(COND, CONST_NAME, ...) \ - [&] { \ - if (COND) { \ - constexpr bool CONST_NAME = true; \ - return __VA_ARGS__(); \ - } else { \ - constexpr bool CONST_NAME = false; \ - return __VA_ARGS__(); \ - } \ - }() diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py b/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py deleted file mode 100644 index 6948a98a374d..000000000000 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py +++ /dev/null @@ -1,315 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import glob -import itertools -import os -import subprocess -import sys - -import jinja2 - -ARCHS = [] -SUPPORT_FP8 = False -SUPPORT_SM75 = False -SUPPORT_SM80 = False -for arch in sys.argv[1].split(","): - arch = arch[: arch.index(".") + 2].replace(".", "") - arch = int(arch) - # SM89 and the SM12x family (SM120 RTX 5090, SM121 DGX Spark GB10) - # fully support mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32. - # SM90 and SM100 can use this PTX, but it’s simulated - # with FP16 MMA, so it cannot achieve any acceleration. - if arch == 89 or arch // 10 == 12: - SUPPORT_FP8 = True - if arch >= 80: - SUPPORT_SM80 = True - if arch == 75: - SUPPORT_SM75 = True - -FILE_HEAD_COMMENT = """ -// auto generated by generate_kernels.py -// clang-format off -""".lstrip() - -FILE_HEAD = ( - FILE_HEAD_COMMENT - + """ -#include "kernel.h" -#include "marlin_template.h" - -namespace MARLIN_NAMESPACE_NAME { -""" -) - -TEMPLATE = ( - "template __global__ void Marlin<" - "{{a_type_id}}, " - "{{b_type_id}}, " - "{{c_type_id}}, " - "{{s_type_id}}, " - "{{threads}}, " - "{{thread_m_blocks}}, " - "{{thread_n_blocks}}, " - "{{thread_k_blocks}}, " - "{{m_block_size_8}}, " - "{{stages}}, " - "{{group_blocks}}, " - "{{is_zp_float}}>" - "( MARLIN_KERNEL_PARAMS );" -) - -THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128), (128, 64, 128)] - -THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4] - -QUANT_CONFIGS = [ - # AWQ-INT4 - { - "b_type": "kU4", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 2, 4, 8], - }, - # GPTQ-INT4 - { - "b_type": "kU4B8", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 2, 4, 8], - }, - # AWQ-INT8 - { - "b_type": "kU8B128", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 2, 4, 8], - }, - # FP8 - { - "b_type": "kFE4M3fn", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 8], - }, - # NVFP4 - { - "b_type": "kFE2M1f", - "s_type": "kFE4M3fn", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [1], - }, - # MXFP4 - { - "a_type": ["kBFloat16"], - "b_type": "kFE2M1f", - "s_type": "kFE8M0fnu", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [2], - }, - # MXFP8 - { - "a_type": ["kBFloat16"], - "b_type": "kFE4M3fn", - "s_type": "kFE8M0fnu", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [2], - }, - # AWQ-INT4 with INT8 activation - { - "a_type": ["kS8"], - "b_type": "kU4", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [-1, 2, 4, 8], - }, - # GPTQ-INT4 with INT8 activation - { - "a_type": ["kS8"], - "b_type": "kU4B8", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [-1, 2, 4, 8], - }, - # GPTQ-INT4 with FP8 activation - { - "a_type": ["kFE4M3fn"], - "b_type": "kU4B8", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [-1, 2, 4, 8], - }, - # AWQ-INT4 with FP8 activation - { - "a_type": ["kFE4M3fn"], - "b_type": "kU4", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [-1, 2, 4, 8], - }, - # MXFP4 with FP8 activation - { - "a_type": ["kFE4M3fn"], - "b_type": "kFE2M1f", - "c_type": ["kBFloat16"], - "s_type": "kFE8M0fnu", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [2], - }, -] - - -def remove_old_kernels(): - for filename in glob.glob(os.path.dirname(__file__) + "/*kernel_*.cu"): - subprocess.call(["rm", "-f", filename]) - - filename = os.path.dirname(__file__) + "/kernel_selector.h" - subprocess.call(["rm", "-f", filename]) - - -def generate_new_kernels(): - result_dict = {} - sm_75_result_dict = {} - - for quant_config in QUANT_CONFIGS: - c_types = quant_config.get("c_type", ["kFloat16", "kBFloat16"]) - a_types = quant_config.get("a_type", ["kFloat16", "kBFloat16"]) - b_type = quant_config["b_type"] - all_group_blocks = quant_config["group_blocks"] - all_m_blocks = quant_config["thread_m_blocks"] - all_thread_configs = quant_config["thread_configs"] - - for a_type, c_type in itertools.product(a_types, c_types): - if not SUPPORT_FP8 and a_type == "kFE4M3fn": - continue - if "16" in a_type and "16" in c_type and a_type != c_type: - continue - s_type = quant_config.get("s_type", c_type) - if (a_type, b_type, c_type) not in result_dict: - result_dict[(a_type, b_type, c_type)] = [] - if a_type in ["kFloat16", "kS8"] and c_type == "kFloat16": - sm_75_result_dict[(a_type, b_type, c_type)] = [] - - for group_blocks, m_blocks, thread_configs in itertools.product( - all_group_blocks, all_m_blocks, all_thread_configs - ): - thread_k, thread_n, threads = thread_configs - - if threads == 256: - # for small batch (m_blocks == 1), - # we only need (128, 128, 256) - # for large batch (m_blocks > 1), - # we only need (64, 256, 256) - if m_blocks <= 1 and (thread_k, thread_n) != (128, 128): - continue - if m_blocks > 1 and (thread_k, thread_n) != (64, 256): - continue - - config = { - "threads": threads, - "s_type": s_type, - "thread_m_blocks": max(m_blocks, 1), - "thread_k_blocks": thread_k // 16, - "thread_n_blocks": thread_n // 16, - "m_block_size_8": "true" if m_blocks == 0.5 else "false", - "stages": 4, - "group_blocks": group_blocks, - "is_zp_float": "false", - } - - if SUPPORT_SM80: - result_dict[(a_type, b_type, c_type)].append(config) - if (a_type, b_type, c_type) in sm_75_result_dict and SUPPORT_SM75: - config_sm75 = config.copy() - config_sm75["stages"] = 2 - sm_75_result_dict[(a_type, b_type, c_type)].append(config_sm75) - - kernel_selector_str = FILE_HEAD_COMMENT - - for result_dict_tmp in [result_dict, sm_75_result_dict]: - for (a_type, b_type, c_type), config_list in result_dict_tmp.items(): - all_template_str_list = [] - if not config_list: - continue - for config in config_list: - s_type = config["s_type"] - template_str = jinja2.Template(TEMPLATE).render( - a_type_id=f"vllm::{a_type}.id()", - b_type_id=f"vllm::{b_type}.id()", - c_type_id=f"vllm::{c_type}.id()", - s_type_id=f"vllm::{s_type}.id()", - **config, - ) - all_template_str_list.append(template_str) - - conditions = [ - f"a_type == vllm::{a_type}", - f"b_type == vllm::{b_type}", - f"c_type == vllm::{c_type}", - f"s_type == vllm::{s_type}", - f"threads == {config['threads']}", - f"thread_m_blocks == {config['thread_m_blocks']}", - f"thread_n_blocks == {config['thread_n_blocks']}", - f"thread_k_blocks == {config['thread_k_blocks']}", - f"m_block_size_8 == {config['m_block_size_8']}", - f"stages == {config['stages']}", - f"group_blocks == {config['group_blocks']}", - f"is_zp_float == {config['is_zp_float']}", - ] - conditions = " && ".join(conditions) - - if kernel_selector_str == FILE_HEAD_COMMENT: - kernel_selector_str += f"if ({conditions})\n kernel = " - else: - kernel_selector_str += f"else if ({conditions})\n kernel = " - - kernel_template2 = ( - "Marlin<{{a_type_id}}, {{b_type_id}}, {{c_type_id}}, " - "{{s_type_id}}, {{threads}}, {{thread_m_blocks}}, " - "{{thread_n_blocks}}, {{thread_k_blocks}}, " - "{{m_block_size_8}}, {{stages}}, {{group_blocks}}, " - "{{is_zp_float}}>;" - ) - - kernel_selector_str += ( - jinja2.Template(kernel_template2).render( - a_type_id=f"vllm::{a_type}.id()", - b_type_id=f"vllm::{b_type}.id()", - c_type_id=f"vllm::{c_type}.id()", - s_type_id=f"vllm::{s_type}.id()", - **config, - ) - + "\n" - ) - - file_content = FILE_HEAD + "\n\n" - file_content += "\n\n".join(all_template_str_list) + "\n\n}\n" - if a_type == "kFE4M3fn": - filename = f"sm89_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu" - elif result_dict_tmp is sm_75_result_dict: - filename = f"sm75_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu" - else: - filename = f"sm80_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu" - - filename = filename.lower() - - with open(os.path.join(os.path.dirname(__file__), filename), "w") as f: - f.write(file_content) - - if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT: - kernel_selector_str += ( - "else if (a_type == vllm::kFE4M3fn)\n" - " STD_TORCH_CHECK(false, " - '"marlin kernel with fp8 activation is not built.");' - ) - - with open(os.path.join(os.path.dirname(__file__), "kernel_selector.h"), "w") as f: - f.write(kernel_selector_str) - - -if __name__ == "__main__": - remove_old_kernels() - generate_new_kernels() diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h deleted file mode 100644 index 3266d28b1edb..000000000000 --- a/csrc/libtorch_stable/ops.h +++ /dev/null @@ -1,822 +0,0 @@ -#pragma once - -#include -#include -#include - -#include -#include -#include - -#include - -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, - torch::stable::Tensor& output_s, - int64_t group_size, double eps, double fp8_min, - double fp8_max, bool scale_ue8m0, - bool dummy_is_scale_transposed, - bool dummy_is_tma_aligned); - -// Fused activation quantisation + DeepGEMM-compatible UE8M0-packed scales. -void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s_packed, - int64_t group_size, double eps, - double min_8bit, double max_8bit); - -void per_token_group_quant_int8(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s, - int64_t group_size, double eps, double int8_min, - double int8_max); - -#ifndef USE_ROCM -bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability); -bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability); -bool cutlass_group_gemm_supported(int64_t cuda_device_capability); - -void cutlass_scaled_mm(torch::stable::Tensor& out, - torch::stable::Tensor const& a, - torch::stable::Tensor const& b, - torch::stable::Tensor const& a_scales, - torch::stable::Tensor const& b_scales, - std::optional const& bias); - -void cutlass_moe_mm(torch::stable::Tensor& out_tensors, - torch::stable::Tensor const& a_tensors, - torch::stable::Tensor const& b_tensors, - torch::stable::Tensor const& a_scales, - torch::stable::Tensor const& b_scales, - torch::stable::Tensor const& expert_offsets, - torch::stable::Tensor const& problem_sizes, - torch::stable::Tensor const& a_strides, - torch::stable::Tensor const& b_strides, - torch::stable::Tensor const& c_strides, bool per_act_token, - bool per_out_ch); - -void cutlass_scaled_mm_azp(torch::stable::Tensor& out, - torch::stable::Tensor const& a, - torch::stable::Tensor const& b, - torch::stable::Tensor const& a_scales, - torch::stable::Tensor const& b_scales, - torch::stable::Tensor const& azp_adj, - std::optional const& azp, - std::optional const& bias); - -void get_cutlass_moe_mm_data( - const torch::stable::Tensor& topk_ids, - torch::stable::Tensor& expert_offsets, - torch::stable::Tensor& problem_sizes1, - torch::stable::Tensor& problem_sizes2, - torch::stable::Tensor& input_permutation, - torch::stable::Tensor& output_permutation, const int64_t num_experts, - const int64_t n, const int64_t k, - const std::optional& blockscale_offsets, - const bool is_gated); - -void get_cutlass_moe_mm_problem_sizes_from_expert_offsets( - const torch::stable::Tensor& expert_first_token_offset, - torch::stable::Tensor& problem_sizes1, - torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k, - const bool swap_ab); - -void get_cutlass_batched_moe_mm_data( - torch::stable::Tensor& expert_offsets, - torch::stable::Tensor& problem_sizes1, - torch::stable::Tensor& problem_sizes2, - const torch::stable::Tensor& expert_num_tokens, - const int64_t num_local_experts, const int64_t padded_m, const int64_t n, - const int64_t k); - -// FP4/NVFP4 ops -bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability); - -void cutlass_scaled_fp4_mm(torch::stable::Tensor& D, - torch::stable::Tensor const& A, - torch::stable::Tensor const& B, - torch::stable::Tensor const& A_sf, - torch::stable::Tensor const& B_sf, - torch::stable::Tensor const& alpha); - -void cutlass_fp4_group_mm(torch::stable::Tensor& output, - const torch::stable::Tensor& a, - const torch::stable::Tensor& b, - const torch::stable::Tensor& a_blockscale, - const torch::stable::Tensor& b_blockscales, - const torch::stable::Tensor& alphas, - const torch::stable::Tensor& problem_sizes, - const torch::stable::Tensor& expert_offsets, - const torch::stable::Tensor& sf_offsets); - -std::tuple scaled_fp4_quant_func( - torch::stable::Tensor const& input, - torch::stable::Tensor const& input_scale, bool is_sf_swizzled_layout); - -void scaled_fp4_quant_out(torch::stable::Tensor const& input, - torch::stable::Tensor const& input_scale, - bool is_sf_swizzled_layout, - torch::stable::Tensor& output, - torch::stable::Tensor& output_scale); - -void scaled_fp4_experts_quant( - torch::stable::Tensor& output, torch::stable::Tensor& output_scale, - torch::stable::Tensor const& input, - torch::stable::Tensor const& input_global_scale, - torch::stable::Tensor const& input_offset_by_experts, - torch::stable::Tensor const& output_scale_offset_by_experts); - -void silu_and_mul_scaled_fp4_experts_quant( - torch::stable::Tensor& output, torch::stable::Tensor& output_scale, - torch::stable::Tensor const& input, - torch::stable::Tensor const& input_global_scale, - torch::stable::Tensor const& input_offset_by_experts, - torch::stable::Tensor const& output_scale_offset_by_experts); - -void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out, - torch::stable::Tensor& output_block_scale, - torch::stable::Tensor& input, - torch::stable::Tensor& input_global_scale); - -void cutlass_mxfp4_group_mm(torch::stable::Tensor& output, - const torch::stable::Tensor& a, - const torch::stable::Tensor& b, - const torch::stable::Tensor& a_blockscale, - const torch::stable::Tensor& b_blockscales, - const torch::stable::Tensor& problem_sizes, - const torch::stable::Tensor& expert_offsets, - const torch::stable::Tensor& sf_offsets); - -// AWQ ops -torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats, - torch::stable::Tensor _kernel, - torch::stable::Tensor _scaling_factors, - torch::stable::Tensor _zeros, - int64_t split_k_iters); - -torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel, - torch::stable::Tensor _scaling_factors, - torch::stable::Tensor _zeros, - int64_t split_k_iters, int64_t thx, - int64_t thy); - -// DSV3 fused A GEMM: conditionally compiled so declaration and impl -// registration are in the source file (dsv3_fused_a_gemm.cu) - -#endif - -// CPU tensor -> CUDA UVA view (shared CUDA/ROCm) -torch::stable::Tensor get_cuda_view_from_cpu_tensor( - torch::stable::Tensor& cpu_tensor); - -// Attention kernels (shared CUDA/ROCm) -void merge_attn_states( - torch::stable::Tensor& output, - std::optional output_lse, - const torch::stable::Tensor& prefix_output, - const torch::stable::Tensor& prefix_lse, - const torch::stable::Tensor& suffix_output, - const torch::stable::Tensor& suffix_lse, - const std::optional prefill_tokens_with_context, - const std::optional& output_scale = std::nullopt); - -torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x, - bool inplace); - -// Layernorm kernels (shared CUDA/ROCm) -void rms_norm(torch::stable::Tensor& out, torch::stable::Tensor& input, - std::optional weight, double epsilon); - -void fused_add_rms_norm(torch::stable::Tensor& input, - torch::stable::Tensor& residual, - std::optional weight, - double epsilon); - -// Layernorm-quant kernels (shared CUDA/ROCm) -void rms_norm_static_fp8_quant(torch::stable::Tensor& out, - torch::stable::Tensor& input, - torch::stable::Tensor& weight, - torch::stable::Tensor& scale, double epsilon); - -void fused_add_rms_norm_static_fp8_quant(torch::stable::Tensor& out, - torch::stable::Tensor& input, - torch::stable::Tensor& residual, - torch::stable::Tensor& weight, - torch::stable::Tensor& scale, - double epsilon); - -// Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm) -void rms_norm_dynamic_per_token_quant( - torch::stable::Tensor& out, torch::stable::Tensor const& input, - torch::stable::Tensor const& weight, torch::stable::Tensor& scales, - double const var_epsilon, std::optional scale_ub, - std::optional residual); - -void rms_norm_per_block_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor const& weight, - torch::stable::Tensor& scales, - double const var_epsilon, - std::optional scale_ub, - std::optional residual, - int64_t group_size, bool is_scale_transposed); - -void silu_and_mul_per_block_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor& scales, - int64_t group_size, - std::optional scale_ub, - bool is_scale_transposed); - -// Positional encoding kernels (shared CUDA/ROCm) -void rotary_embedding(torch::stable::Tensor& positions, - torch::stable::Tensor& query, - std::optional key, - int64_t head_size, torch::stable::Tensor& cos_sin_cache, - bool is_neox, int64_t rope_dim_offset, bool inverse); - -void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q, - int64_t num_heads_k, int64_t num_heads_v, - int64_t head_dim, double eps, - torch::stable::Tensor& q_weight, - torch::stable::Tensor& k_weight, - torch::stable::Tensor& cos_sin_cache, bool is_neox, - torch::stable::Tensor& position_ids, - int64_t forced_token_heads_per_warp); - -void fused_deepseek_v4_kv_rope_insert( - torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& position_ids, - torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size, - std::optional fp8_scale, bool kv_mxfp8); - -torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( - torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv, - torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& position_ids, - torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, - double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8, - bool apply_q_rope, bool is_q_interleaved); - -void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( - torch::stable::Tensor& q, torch::stable::Tensor const& kv, - torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& position_ids, - torch::stable::Tensor const& cos_sin_cache, double eps, - int64_t cache_block_size, bool apply_q_norm); - -void fused_kimi_k3_mla_key_concat_kv_cache_insert( - torch::stable::Tensor& q, torch::stable::Tensor const& k_nope, - torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, - torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, - std::optional position_ids, - std::optional cos_sin_cache); - -void fused_kimi_k3_mla_key_concat_ds_mla_insert( - torch::stable::Tensor& q, torch::stable::Tensor const& k_nope, - torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, - torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, - std::optional position_ids, - std::optional cos_sin_cache); - -void fused_kimi_k3_mla_kv_concat(torch::stable::Tensor const& k_nope, - torch::stable::Tensor const& k_pe, - torch::stable::Tensor& k_out); - -void fused_kimi_k3_mla_kv_concat_quant_fp8(torch::stable::Tensor const& k_nope, - torch::stable::Tensor const& k_pe, - torch::stable::Tensor const& v, - torch::stable::Tensor& k_fp8, - torch::stable::Tensor& v_fp8); - -void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( - torch::stable::Tensor const& q, torch::stable::Tensor const& k_nope, - torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, - torch::stable::Tensor const& v, torch::stable::Tensor& q_fp8, - torch::stable::Tensor& k_fp8, torch::stable::Tensor& v_fp8, - torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& q_scale_inv, - torch::stable::Tensor const& k_scale_inv, - torch::stable::Tensor const& v_scale_inv, - torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size, - std::optional position_ids, - std::optional cos_sin_cache); - -void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert( - torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, - torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, - torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, - std::optional position_ids, - std::optional cos_sin_cache); - -void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert( - torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, - torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, - torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& q_scale_inv, - torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size, - std::optional position_ids, - std::optional cos_sin_cache); - -void fused_kimi_k3_mla_decode_q_concat_ds_mla_insert( - torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, - torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, - torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, - std::optional position_ids, - std::optional cos_sin_cache); - -void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( - torch::stable::Tensor const& q, torch::stable::Tensor const& kv, - torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& position_ids, - torch::stable::Tensor const& cos_sin_cache, - torch::stable::Tensor const& fp8_scale, - torch::stable::Tensor const& q_fp8_scale_inv, double eps, - int64_t cache_block_size, bool apply_q_norm); - -#ifndef USE_ROCM -std::tuple -minimax_allreduce_rms_qk(torch::stable::Tensor qkv, - torch::stable::Tensor const& norm_weight_q, - torch::stable::Tensor const& norm_weight_k, - torch::stable::Tensor workspace, int64_t const q_size, - int64_t const kv_size, int64_t const rank, - int64_t const nranks, double const eps); -#endif - -// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV / -// index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs -// the index branch and scatters k/v/index_k into their paged caches. -void fused_minimax_m3_qknorm_rope_kv_insert( - torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight, - torch::stable::Tensor const& k_norm_weight, - torch::stable::Tensor const& cos_sin_cache, - torch::stable::Tensor const& positions, int64_t num_heads, - int64_t num_kv_heads, int64_t rotary_dim, double eps, - std::optional index_q_norm_weight, - std::optional index_k_norm_weight, - int64_t num_index_heads, std::optional slot_mapping, - std::optional index_slot_mapping, - std::optional kv_cache, - std::optional index_cache, int64_t block_size, - std::optional q_out, - std::optional index_q_out, - const std::string& kv_cache_dtype, bool skip_index_branch, - std::optional q_fp8_out, double q_fp8_scale, - std::optional kv_k_scale, - std::optional kv_v_scale); - -#ifdef VLLM_ENABLE_FUSED_KDA_DECODE -void fused_kda_decode( - torch::stable::Tensor const& x, torch::stable::Tensor const& weight, - std::optional bias, - torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g, - torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, - torch::stable::Tensor const& dt_bias, - torch::stable::Tensor const& state_indices, torch::stable::Tensor& state, - torch::stable::Tensor& out, std::optional lower_bound, - std::optional output_gate, - std::optional norm_weight, double norm_eps); - -#endif - -#ifdef VLLM_ENABLE_FUSED_GDN_DECODE -void fused_gdn_decode_post_conv_mtp( - torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a, - torch::stable::Tensor const& b, torch::stable::Tensor const& a_log, - torch::stable::Tensor const& dt_bias, - torch::stable::Tensor const& state_indices, - torch::stable::Tensor const& cu_seqlens, - torch::stable::Tensor const& num_accepted_tokens, - torch::stable::Tensor& state, torch::stable::Tensor const& output_gate, - torch::stable::Tensor const& norm_weight, torch::stable::Tensor& out, - double scale, double norm_eps, const std::string& output_gate_activation); - -#endif - -#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK -void fused_kda_prologue( - torch::stable::Tensor const& q, torch::stable::Tensor const& k, - torch::stable::Tensor const& v, torch::stable::Tensor const& raw_g, - torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, - torch::stable::Tensor const& dt_bias, torch::stable::Tensor& qg, - torch::stable::Tensor& w, torch::stable::Tensor& u, - torch::stable::Tensor& kg_t, torch::stable::Tensor& aqk, - torch::stable::Tensor& decay, torch::stable::Tensor const& cu_seqlens, - torch::stable::Tensor const& chunk_indices, - std::optional conv_weight, - std::optional conv_state, - std::optional conv_state_indices, - std::optional conv_has_initial_state, double scale, - double lower_bound); - -void fused_kda_chunk( - torch::stable::Tensor const& qg, torch::stable::Tensor const& w, - torch::stable::Tensor const& u, torch::stable::Tensor const& kg_t, - torch::stable::Tensor const& aqk, torch::stable::Tensor const& decay, - std::optional initial_state, - std::optional final_state, - torch::stable::Tensor& out, torch::stable::Tensor const& cu_seqlens, - torch::stable::Tensor const& chunk_offsets, double scale, - std::optional group_state, int64_t groups, - std::optional checkpoint_state, - std::optional checkpoint_offsets, - std::optional checkpoint_state_indices, - std::optional state_cache, - std::optional state_indices, - std::optional has_initial_state); -#endif - -#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES -void kimi_k3_attn_res(torch::stable::Tensor& prefix, - std::optional delta, - torch::stable::Tensor& blocks, - torch::stable::Tensor const& norm_weight, - torch::stable::Tensor const& qk_weight, - std::optional output_norm_weight, - torch::stable::Tensor& output, int64_t num_blocks, - int64_t block_write_idx, double eps, - double output_norm_eps); -#endif - -// Sampler kernels (shared CUDA/ROCm) -void apply_repetition_penalties_( - torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask, - const torch::stable::Tensor& output_mask, - const torch::stable::Tensor& repetition_penalties); - -void top_k_per_row_prefill(const torch::stable::Tensor& logits, - const torch::stable::Tensor& rowStarts, - const torch::stable::Tensor& rowEnds, - torch::stable::Tensor& indices, int64_t numRows, - int64_t stride0, int64_t stride1, int64_t topK); - -void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n, - const torch::stable::Tensor& seqLens, - torch::stable::Tensor& indices, int64_t numRows, - int64_t stride0, int64_t stride1, int64_t topK); - -void persistent_topk(const torch::stable::Tensor& logits, - const torch::stable::Tensor& lengths, - torch::stable::Tensor& output, - torch::stable::Tensor& workspace, int64_t k, - int64_t max_seq_len); - -#ifdef VLLM_ENABLE_COOPERATIVE_TOPK -void cooperative_topk(const torch::stable::Tensor& logits, - const torch::stable::Tensor& lengths, - torch::stable::Tensor& output, - torch::stable::Tensor& workspace, int64_t k, - int64_t max_seq_len); -#endif - -void selective_scan_fwd( - const torch::stable::Tensor& u, const torch::stable::Tensor& delta, - const torch::stable::Tensor& A, const torch::stable::Tensor& B, - const torch::stable::Tensor& C, - const std::optional& D_, - const std::optional& z_, - const std::optional& delta_bias_, - bool delta_softplus, - const std::optional& query_start_loc, - const std::optional& cache_indices, - const std::optional& has_initial_state, - const torch::stable::Tensor& ssm_states, int64_t null_block_id, - int64_t block_size, - const std::optional& block_idx_first_scheduled_token, - const std::optional& block_idx_last_scheduled_token, - const std::optional& initial_state_idx, - const std::optional& cu_chunk_seqlen, - const std::optional& last_chunk_indices); - -using fptr_t = int64_t; -fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, - torch::stable::Tensor& rank_data, int64_t rank, - bool fully_connected); -void all_reduce(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, fptr_t reg_buffer, - int64_t reg_buffer_sz_bytes); -void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, fptr_t reg_buffer, - int64_t reg_buffer_sz_bytes); -void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, fptr_t local_buffer, - fptr_t multicast_buffer, fptr_t epoch_buffer, - int64_t stage_sz_bytes); -void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, fptr_t reg_buffer, - int64_t reg_buffer_sz_bytes); -void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, - fptr_t local_buffer, fptr_t epoch_buffer, - int64_t stage_sz_bytes); -void mnnvl_multimem_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, - fptr_t local_buffer, fptr_t multicast_buffer, - int64_t stage_sz_bytes, int64_t block_limit); -void dispose(fptr_t _fa); -int64_t meta_size(); -void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs); -std::tuple, std::vector> -get_graph_buffer_ipc_meta(fptr_t _fa); -void register_graph_buffers(fptr_t _fa, - const std::vector>& handles, - const std::vector>& offsets); -std::tuple allocate_shared_buffer_and_handle( - int64_t size); -int64_t open_mem_handle(torch::stable::Tensor& mem_handle); -void free_shared_buffer(int64_t buffer); - -// Activation kernels (shared CUDA/ROCm) -void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); -void silu_and_mul_clamp(torch::stable::Tensor& out, - torch::stable::Tensor& input, double limit, - 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, - torch::stable::Tensor& input); -void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, - double threshold); -void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, - double alpha = 1.702, double limit = 7.0); -void situ_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, - double beta = 1.0, double linear_beta = -1.0); -void situ_and_mul_quant( - torch::stable::Tensor& out, torch::stable::Tensor& scale, - torch::stable::Tensor& input, double beta = 1.0, double linear_beta = -1.0, - int64_t group_size = 0, - std::optional num_valid_tokens = std::nullopt, - int64_t topk = 1); -void masked_situ_and_mul(torch::stable::Tensor& out, - torch::stable::Tensor& input, - const torch::stable::Tensor& expert_num_tokens, - double beta = 1.0, double linear_beta = -1.0); -void masked_moe_activation(torch::stable::Tensor& out, - torch::stable::Tensor& input, - const torch::stable::Tensor& valid_token_counts, - const std::string& activation, - double clamp_limit = 0.0, double alpha = 1.0, - double beta = 0.0, double situ_beta = 1.0, - double situ_linear_beta = -1.0); -void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input); -void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input); -void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input); - -void relu_squared(torch::stable::Tensor& out, torch::stable::Tensor& input); - -// INT8 quantization kernels (shared CUDA/ROCm) -void static_scaled_int8_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor const& scale, - std::optional const& azp); - -void dynamic_scaled_int8_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor& scales, - std::optional const& azp); - -// FP8 quantization kernels (shared CUDA/ROCm) -void static_scaled_fp8_quant( - torch::stable::Tensor& out, torch::stable::Tensor const& input, - torch::stable::Tensor const& scale, - std::optional group_shape = - std::nullopt); - -void dynamic_scaled_fp8_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor& scale); - -void dynamic_per_token_scaled_fp8_quant( - torch::stable::Tensor& out, torch::stable::Tensor const& input, - torch::stable::Tensor& scale, - std::optional const& scale_ub); - -// GPTQ kernels (shared CUDA/ROCm) -torch::stable::Tensor gptq_gemm(torch::stable::Tensor a, - torch::stable::Tensor b_q_weight, - torch::stable::Tensor b_gptq_qzeros, - torch::stable::Tensor b_gptq_scales, - bool use_exllama, bool use_v2_format, - int64_t bit); - -void gptq_shuffle(torch::stable::Tensor q_weight, int64_t bit); - -// Cache ops (shared CUDA/ROCm) -void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst, - int64_t block_size_in_bytes, - const torch::stable::Tensor& block_mapping); - -// Batch swap: submit all block copies in a single driver call. -void swap_blocks_batch(const torch::stable::Tensor& src_ptrs, - const torch::stable::Tensor& dst_ptrs, - const torch::stable::Tensor& sizes, - bool is_src_access_order_any); - -void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value, - torch::stable::Tensor& key_cache, - torch::stable::Tensor& value_cache, - torch::stable::Tensor& slot_mapping, - const std::string& kv_cache_dtype, - torch::stable::Tensor& k_scale, - torch::stable::Tensor& v_scale); - -void reshape_and_cache_flash( - torch::stable::Tensor& key, torch::stable::Tensor& value, - torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, - torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype, - torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale); - -void concat_and_cache_mla(torch::stable::Tensor& kv_c, - torch::stable::Tensor& k_pe, - torch::stable::Tensor& kv_cache, - torch::stable::Tensor& slot_mapping, - const std::string& kv_cache_dtype, - torch::stable::Tensor& scale); - -void concat_and_cache_mla_grouped( - torch::stable::Tensor& kv_c, torch::stable::Tensor& k_pe, - torch::stable::Tensor& kv_cache_ptrs, torch::stable::Tensor& slot_mapping, - int64_t block_size, int64_t block_stride, int64_t entry_stride, - std::optional kv_scales, - const std::string& kv_cache_dtype); - -#ifndef USE_ROCM -// HiSparse kernels use raw PTX in the row copy; CUDA-only. -void hisparse_resolve_residency( - torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, - torch::stable::Tensor const& hot_block_table, - torch::stable::Tensor const& global_indices, - torch::stable::Tensor& hot_indices, - torch::stable::Tensor& device_global_indices, - torch::stable::Tensor& lru_slots, - std::optional const& request_state_indices, - int64_t region_stride, int64_t max_union_rows, - std::optional const& miss_mask, - std::optional const& stats, - std::optional const& attention_indices, - int64_t attention_block_stride, - std::optional const& request_ids, - std::optional const& source_block_table, - int64_t source_block_size, - std::optional const& resolved_global_indices, - std::optional const& valid_counts, - std::optional const& swap_host_physical_rows, - std::optional const& swap_device_physical_rows, - std::optional const& swap_counts, - std::optional const& resident_block_table, - int64_t resident_block_size, int64_t resident_null_block); - -void hisparse_invalidate_written_slots( - torch::stable::Tensor& device_global_indices, - torch::stable::Tensor const& request_state_indices, - torch::stable::Tensor const& req_id_per_token, - torch::stable::Tensor const& written_slots); - -void hisparse_gather_plan( - torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, - torch::stable::Tensor const& global_indices, - torch::stable::Tensor const& hot_indices, - torch::stable::Tensor const& miss_mask, - std::optional const& request_state_indices, - std::optional const& attention_indices, - int64_t attention_block_stride); - -void hisparse_gather_compact(torch::stable::Tensor const& host_cache, - torch::stable::Tensor& hot_cache, - torch::stable::Tensor const& miss_global_indices, - torch::stable::Tensor const& miss_hot_indices, - torch::stable::Tensor const& miss_counts); - -#endif // !USE_ROCM - -// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla -void concat_and_cache_mla_rope_fused( - torch::stable::Tensor& positions, torch::stable::Tensor& q_pe, - torch::stable::Tensor& k_pe, torch::stable::Tensor& kv_c, - torch::stable::Tensor& rope_cos_sin_cache, bool rope_is_neox, - torch::stable::Tensor& slot_mapping, torch::stable::Tensor& kv_cache, - const std::string& kv_cache_dtype, - torch::stable::Tensor& kv_cache_quant_scale); - -// Just for unittest -void convert_fp8(torch::stable::Tensor& dst_cache, - torch::stable::Tensor& src_cache, const double scale, - const std::string& kv_cache_dtype); - -void gather_and_maybe_dequant_cache( - torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // ENTRIES...] - torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] - torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] - int64_t num_tokens, const std::string& kv_cache_dtype, - torch::stable::Tensor const& scale, - std::optional seq_starts = std::nullopt); - -// TODO(hc): cp_gather_cache need support scaled kvcahe in the future. -void cp_gather_cache( - torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // ENTRIES...] - torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] - int64_t batch_size, - std::optional seq_starts = std::nullopt); - -// Gather and upconvert FP8 KV cache to BF16 workspace -void cp_gather_and_upconvert_fp8_kv_cache( - torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // 656] - torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] - torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& workspace_starts, // [BATCH] - int64_t batch_size, - std::optional seq_starts = std::nullopt, - std::optional host_cache = std::nullopt, - std::optional host_row_ids = std::nullopt, - std::optional device_row_ids = std::nullopt); - -// Gather and upconvert an nvfp4_ds_mla KV cache to a BF16 workspace -void cp_gather_and_upconvert_nvfp4_kv_cache( - torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // 352] - torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] - torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& workspace_starts, // [BATCH] - int64_t batch_size); - -// Indexer K quantization and cache function -void indexer_k_quant_and_cache( - torch::stable::Tensor& k, // [num_tokens, head_dim] - torch::stable::Tensor& kv_cache, // [num_blocks, block_size, - // cache_stride] - torch::stable::Tensor& slot_mapping, // [num_tokens] - int64_t quant_block_size, // quantization block size - const std::string& scale_fmt); - -// Concatenate query nope and rope for MLA/DSA attention -void concat_mla_q( - torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] - torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] - torch::stable::Tensor& q_out); // [num_tokens, num_heads, nope_dim + - // rope_dim] - -// Extract function to gather quantized K cache -void cp_gather_indexer_k_quant_cache( - const torch::stable::Tensor& kv_cache, // [num_blocks, block_size, - // cache_stride] - torch::stable::Tensor& dst_k, // [num_tokens, head_dim] - torch::stable::Tensor& dst_scale, // [num_tokens, head_dim / - // quant_block_size * 4] - const torch::stable::Tensor& block_table, // [batch_size, num_blocks] - const torch::stable::Tensor& cu_seq_lens); // [batch_size + 1] - -// Fused vocab-parallel embedding lookup (see -// vocab_parallel_embedding_kernels.cu). -void vocab_parallel_embedding( - torch::stable::Tensor& out, const torch::stable::Tensor& input_ids, - const torch::stable::Tensor& weight, int64_t org_vocab_start_index, - int64_t org_vocab_end_index, int64_t num_org_vocab_padding, - int64_t added_vocab_start_index, int64_t added_vocab_end_index); - -// LongCat n-gram embedding index kernel (see ngram_embedding_kernels.cu). -void ngram_compute_n_gram_ids( - int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights, - torch::stable::Tensor& ne_mods, - torch::stable::Tensor& exclusive_ne_embedder_size_sums, - torch::stable::Tensor& exclusive_req_len_sums, - torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices, - torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids); diff --git a/csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu b/csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu deleted file mode 100644 index 2ffbee7ca26e..000000000000 --- a/csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu +++ /dev/null @@ -1,12 +0,0 @@ -#include - -#include "libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h" - -void per_token_group_quant_int8(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s, - int64_t group_size, double eps, double int8_min, - double int8_max) { - per_token_group_quant_8bit(input, output_q, output_s, group_size, eps, - int8_min, int8_max); -} diff --git a/csrc/libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h b/csrc/libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h deleted file mode 100644 index 6630c0decee6..000000000000 --- a/csrc/libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -#include - -// 8-bit per-token-group quantization helper used by both FP8 and INT8 -void per_token_group_quant_8bit(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s, - int64_t group_size, double eps, double min_8bit, - double max_8bit, bool scale_ue8m0 = false); - -// Public op: register-resident packed quant for the DeepGEMM Blackwell path. -// Restricted to group_size == 128 and bf16/fp16 input; other configurations -// raise STD_TORCH_CHECK. The legacy shared-memory fallback was removed because -// no production caller (deep_gemm_moe / input_quant_fp8) uses other shapes. -void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s_packed, - int64_t group_size, double eps, - double min_8bit, double max_8bit); diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp deleted file mode 100644 index 3d4f9ab27389..000000000000 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ /dev/null @@ -1,1171 +0,0 @@ -#include "ops.h" -#include "cuda_utils.h" -#include "core/registration.h" - -#include - -// Register ops with STABLE_TORCH_LIBRARY for libtorch stable ABI compatibility. -// Note: We register under namespace "_C" so ops are accessible as -// torch.ops._C. for compatibility with existing code. -STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { - // Compute per-token-group FP8 quantized tensor and scaling factor. - // The dummy arguments are here so we can correctly fuse with RMSNorm. - ops.def( - "per_token_group_fp8_quant(Tensor input, Tensor! output_q, Tensor! " - "output_s, " - "int group_size, float eps, float fp8_min, float fp8_max, bool " - "scale_ue8m0, bool dummy_is_scale_transposed, bool dummy_is_tma_aligned " - ") -> ()"); - // Compute per-token-group 8-bit quantized tensor and UE8M0-packed, - // TMA-aligned scales for DeepGEMM. - ops.def( - "per_token_group_fp8_quant_packed(Tensor input, Tensor! output_q, " - "Tensor! output_s_packed, int group_size, float eps, float fp8_min, " - "float fp8_max) -> ()"); - // Compute per-token-group INT8 quantized tensor and scaling factor. - ops.def( - "per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! " - "output_s, int group_size, float eps, float int8_min, float int8_max) -> " - "()"); - ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor"); - -#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 - - // Machete (Dense) Optimized Mixed Precision GEMM for Hopper. - ops.def( - "machete_supported_schedules(" - " ScalarType a_type," - " int b_type," - " ScalarType? maybe_group_scales_type," - " ScalarType? maybe_group_zeros_type," - " ScalarType? maybe_channel_scales_type," - " ScalarType? maybe_token_scales_type," - " ScalarType? maybe_out_type" - ") -> str[]"); - ops.def( - "machete_mm(" - " Tensor A," - " Tensor B," - " int b_type," - " ScalarType? out_type," - " Tensor? group_scales," - " Tensor? group_zeros," - " int? group_size," - " Tensor? channel_scales," - " Tensor? token_scales," - " str? schedule" - ") -> Tensor"); - ops.def( - "machete_prepack_B(" - " Tensor B," - " ScalarType a_type," - " int b_type," - " ScalarType? group_scales_type" - ") -> Tensor"); - // conditionally compiled so impl registration is in source file - - // Marlin GEMM - ops.def( - "marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, " - "Tensor? b_bias_or_none,Tensor b_scales, " - "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, " - "Tensor workspace, int b_type_id, " - "SymInt size_m, SymInt size_n, SymInt size_k, " - "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // gptq_marlin repack from GPTQ. - ops.def( - "gptq_marlin_repack(Tensor b_q_weight, " - "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // awq_marlin repack from AWQ. - ops.def( - "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, " - "SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // preprocess W-int4A-fp8 weight for marlin kernel - ops.def( - "marlin_int4_fp8_preprocess(Tensor qweight, " - "Tensor? qzeros_or_none, bool inplace) -> Tensor"); - // conditionally compiled so impl registrations are in source file -#endif - -#ifndef USE_ROCM - // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column - // quantization, as well as bias - ops.def( - "cutlass_scaled_mm(Tensor! out, Tensor a," - " Tensor b, Tensor a_scales," - " Tensor b_scales, Tensor? bias) -> ()"); - - // CUTLASS w8a8 GEMM, supporting asymmetric per-tensor or per-row/column - // quantization. - ops.def( - "cutlass_scaled_mm_azp(Tensor! out, Tensor a," - " Tensor b, Tensor a_scales," - " Tensor b_scales, Tensor azp_adj," - " Tensor? azp, Tensor? bias) -> ()"); - - // Check if cutlass scaled_mm is supported for CUDA devices of the given - // capability - ops.def("cutlass_scaled_mm_supports_fp8(int cuda_device_capability) -> bool"); - - // Check if cutlass grouped gemm is supported for CUDA devices of the given - // capability - ops.def("cutlass_group_gemm_supported(int cuda_device_capability) -> bool"); - - // CUTLASS w8a8 grouped GEMM - ops.def( - "cutlass_moe_mm(Tensor! out_tensors, Tensor a_tensors, Tensor b_tensors, " - " Tensor a_scales, Tensor b_scales, Tensor expert_offsets, " - " Tensor problem_sizes, Tensor a_strides, " - " Tensor b_strides, Tensor c_strides, bool per_act_token, " - " bool per_out_ch) -> ()"); - - // A function that computes data required to run fused MoE with w8a8 grouped - // GEMM. It takes topk_ids as an input, and computes expert_offsets - // (token start indices of each expert). In addition to this, it computes - // problem sizes for each expert's multiplication used by the two mms called - // from fused MoE operation, and arrays with permutations required to shuffle - // and de-shuffle the input/output of the fused operation. - ops.def( - "get_cutlass_moe_mm_data(Tensor topk_ids, Tensor! expert_offsets, " - " Tensor! problem_sizes1, Tensor! problem_sizes2, " - " Tensor! input_permutation, " - " Tensor! output_permutation, int num_experts, " - " int n, int k, Tensor? blockscale_offsets, " - " bool is_gated) -> ()"); - - // compute per-expert problem sizes from expert_first_token_offset - // produced by vLLM's moe_permute kernel - ops.def( - "get_cutlass_moe_mm_problem_sizes_from_expert_offsets(" - " Tensor expert_first_token_offset, " - " Tensor! problem_sizes1, " - " Tensor! problem_sizes2, " - " int n, int k, bool swap_ab) -> ()"); - - // A function that computes data required to run fused MoE with w8a8 grouped - // GEMM in batched expert format. It takes expert_num_tokens - // as an input, and computes expert_offsets (token start indices of each - // expert). In addition to this, it computes problem sizes for each expert's - // multiplication used by the two mms called from fused MoE operation. - ops.def( - "get_cutlass_batched_moe_mm_data(Tensor! expert_offsets, " - " Tensor! problem_sizes1, " - " Tensor! problem_sizes2, " - " Tensor expert_num_tokens, " - " int num_local_experts, int padded_m, " - " int n, int k) -> ()"); - - // Check if cutlass scaled_mm supports block quantization (used by DeepSeekV3) - ops.def( - "cutlass_scaled_mm_supports_block_fp8(int cuda_device_capability) -> " - "bool"); - - // CUTLASS nvfp4 block scaled GEMM - ops.def( - "cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b," - " Tensor block_scale_a, Tensor block_scale_b," - " Tensor alpha) -> ()"); - - // cutlass nvfp4 block scaled group GEMM - ops.def( - "cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b," - " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas," - " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); - - // cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE) - ops.def( - "cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b," - " Tensor a_blockscale, Tensor b_blockscales," - " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); - - // Compute NVFP4 block quantized tensor. - ops.def( - "scaled_fp4_quant(Tensor input," - " Tensor input_scale, bool " - "is_sf_swizzled_layout) -> (Tensor, Tensor)"); - - // Out variant - // TODO: Add out_variant tag once PyTorch supports it (added in 2.11) - // This registration is now migrated to stable ABI - // at::Tag::out_variant is not available in the stable ABI (enum_tag.h is not - // yet in torch/headeronly), the tag should be applied from Python - // via torch.library.Library.define(..., tags=(torch.Tag.out_variant,)) - // with the .impl remaining in C++. - // See pytorch/pytorch#176117. - ops.def( - "scaled_fp4_quant.out(Tensor input," - " Tensor input_scale, bool " - "is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) " - "-> ()"); - - // Compute NVFP4 experts quantization. - ops.def( - "scaled_fp4_experts_quant(Tensor! output, Tensor! output_scale," - "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts," - "Tensor output_scale_offset_by_experts) -> ()"); - - // Fused SiLU+Mul+NVFP4 experts quantization. - ops.def( - "silu_and_mul_scaled_fp4_experts_quant(Tensor! output, Tensor! " - "output_scale," - "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts," - "Tensor output_scale_offset_by_experts) -> ()"); - - // Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs). - ops.def( - "mxfp4_experts_quant(Tensor! output, Tensor! output_scale," - "Tensor input, Tensor input_offset_by_experts," - "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); - - // Fused SiLU+Mul+MXFP4 experts quantization. - ops.def( - "silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! " - "output_scale," - "Tensor input, Tensor input_offset_by_experts," - "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); - - // Fused SiLU+Mul+NVFP4 quantization. - ops.def( - "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, " - "Tensor input, Tensor input_global_scale) -> ()"); - - // Check if cutlass_scaled_mm_fp4 is supported for CUDA devices - // of the given capability - ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool"); - - // CUTLASS w4a8 GEMM - ops.def( - "cutlass_w4a8_mm(" - " Tensor A," - " Tensor B," - " Tensor group_scales," - " int group_size," - " Tensor channel_scales," - " Tensor token_scales," - " ScalarType? out_type," - " str? maybe_schedule" - ") -> Tensor"); - - // pack scales - ops.def("cutlass_pack_scale_fp8(Tensor scales) -> Tensor"); - - // encode and reorder weight matrix - ops.def("cutlass_encode_and_reorder_int4b(Tensor B) -> Tensor"); - - // CUTLASS w4a8 grouped GEMM - ops.def( - "cutlass_w4a8_moe_mm(" - " Tensor! out_tensors," - " Tensor a_tensors," - " Tensor b_tensors," - " Tensor a_scales," - " Tensor b_scales," - " Tensor b_group_scales," - " int b_group_size," - " Tensor expert_offsets," - " Tensor problem_sizes," - " Tensor a_strides," - " Tensor b_strides," - " Tensor c_strides," - " Tensor group_scale_strides," - " str? maybe_schedule" - ") -> ()"); - - ops.def( - "cutlass_encode_and_reorder_int4b_grouped(Tensor b_tensors) -> (Tensor, " - "Tensor)"); - - // SM100 CUTLASS MLA decode - // conditionally compiled so impl registrations are in source file - ops.def( - "sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope," - " Tensor q_pe, Tensor kv_c_and_k_pe_cache," - " Tensor seq_lens, Tensor page_table," - " Tensor workspace, float scale," - " int num_kv_splits) -> ()"); - - ops.def( - "sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches," - " int sm_count, int num_kv_splits) " - "-> int"); - // Quantized GEMM for AWQ. - ops.def( - "awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, " - "Tensor _zeros, SymInt split_k_iters) -> Tensor"); - - // Dequantization for AWQ. - ops.def( - "awq_dequantize(Tensor _kernel, Tensor _scaling_factors, " - "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor"); - - // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens). - // conditionally compiled so impl registration is in source file - ops.def( - "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b, " - "bool enable_pdl=False) -> ()"); - - // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+). - // conditionally compiled so impl registration is in source file - ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); -#endif - - // Merge attn states - // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005 - // can be used to combine partial attention results (in the split-KV case) - ops.def( - "merge_attn_states(" - " Tensor! output," - " Tensor!? output_lse," - " Tensor prefix_output," - " Tensor prefix_lse," - " Tensor suffix_output," - " Tensor suffix_lse," - " int!? prefill_tokens_with_context," - " Tensor? output_scale=None) -> ()"); - - // Hadamard transforms - // conditionally compiled so impl registration is in source file - ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor"); - - // Apply Root Mean Square (RMS) Normalization to the input tensor. - ops.def( - "rms_norm(Tensor! result, Tensor input, Tensor? weight, float epsilon) " - "-> " - "()"); - - // In-place fused Add and RMS Normalization. - ops.def( - "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, " - "float epsilon) -> ()"); - - // Layernorm-quant - // Apply Root Mean Square (RMS) Normalization to the input tensor. - ops.def( - "rms_norm_static_fp8_quant(Tensor! result, Tensor input, Tensor weight, " - "Tensor scale, float epsilon) -> " - "()"); - - // In-place fused Add and RMS Normalization. - ops.def( - "fused_add_rms_norm_static_fp8_quant(Tensor! result, Tensor input, " - "Tensor! residual, Tensor weight, " - "Tensor scale, float epsilon) -> ()"); - - // Fused Layernorm + Quant kernels - ops.def( - "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, " - "Tensor weight, Tensor! scale, float epsilon, " - "Tensor? scale_ub, Tensor!? residual) -> ()"); - - // Fused Layernorm + Block quant kernels - ops.def( - "rms_norm_per_block_quant(Tensor! result, Tensor input, " - "Tensor weight, Tensor! scale, float epsilon, " - "Tensor? scale_ub, Tensor!? residual, int group_size, " - "bool is_scale_transposed) -> ()"); - - // Fused SiLU+Mul + per-block quantization - ops.def( - "silu_and_mul_per_block_quant(" - "Tensor! out, " - "Tensor input, " - "Tensor! scales, " - "int group_size, " - "Tensor? scale_ub=None, " - "bool is_scale_transposed=False) -> ()"); - - // Rotary embedding - // Apply GPT-NeoX or GPT-J style rotary embedding to query and key. - ops.def( - "rotary_embedding(Tensor positions, Tensor! query," - " Tensor!? key, int head_size," - " Tensor cos_sin_cache, bool is_neox, int " - "rope_dim_offset=0, bool inverse=False) -> ()"); - - // Function for fused QK Norm and RoPE - ops.def( - "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, " - "int num_heads_k, int num_heads_v, int head_dim, float eps, " - "Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, " - "bool is_neox, Tensor position_ids, " - "int forced_token_heads_per_warp=-1) -> ()"); - - // q_head_padded is the padded Q head count of the returned tensor, or 0 to - // do the KV insert alone and return an empty tensor. The Q knobs are - // independent: apply_q_norm and apply_q_rope each drop that step for Q - // alone (KV is always rotated), and is_q_interleaved reads and writes Q in - // FlashMLA's mega-attention chunk-interleaved layout, which moves the - // padding heads to the tail of every head-dim chunk. - ops.def( - "fused_deepseek_v4_kv_rope_insert(" - "Tensor kv, Tensor! k_cache, Tensor slot_mapping, Tensor position_ids, " - "Tensor cos_sin_cache, int cache_block_size, Tensor? fp8_scale=None, " - "bool kv_mxfp8=False) -> " - "()"); - ops.def( - "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(" - "Tensor q_in, Tensor kv, Tensor! k_cache, " - "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " - "int q_head_padded, float eps, int cache_block_size, " - "bool apply_q_norm=True, bool kv_mxfp8=False, bool apply_q_rope=True, " - "bool is_q_interleaved=False) -> Tensor"); - - // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate - // FP8 tensor, and KV into a contiguous 512-wide token-strided cache. - ops.def( - "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(" - "Tensor! q, Tensor kv, Tensor! k_cache, Tensor slot_mapping, " - "Tensor position_ids, Tensor cos_sin_cache, float eps, " - "int cache_block_size, bool apply_q_norm=True) -> ()"); - ops.def( - "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(" - "Tensor q, Tensor kv, Tensor! q_fp8, Tensor! k_cache, " - "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " - "Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, " - "int cache_block_size, bool apply_q_norm=True) -> ()"); - - // Kimi-K3 MLA epilogues: optional RoPE followed by concat/cache insertion. - ops.def( - "fused_kimi_k3_mla_key_concat_kv_cache_insert(" - "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, " - "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " - "int cache_block_size, Tensor? position_ids=None, " - "Tensor? cos_sin_cache=None) -> ()"); - ops.def( - "fused_kimi_k3_mla_key_concat_ds_mla_insert(" - "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, " - "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " - "int cache_block_size, Tensor? position_ids=None, " - "Tensor? cos_sin_cache=None) -> ()"); - ops.def( - "fused_kimi_k3_mla_kv_concat(Tensor k_nope, Tensor k_pe, Tensor! k_out) " - "-> ()"); - ops.def( - "fused_kimi_k3_mla_kv_concat_quant_fp8(" - "Tensor k_nope, Tensor k_pe, Tensor v, Tensor! k_fp8, Tensor! v_fp8) " - "-> ()"); - ops.def( - "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(" - "Tensor q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, Tensor v, " - "Tensor! q_fp8, Tensor! k_fp8, Tensor! v_fp8, Tensor! k_cache, " - "Tensor slot_mapping, Tensor q_scale_inv, Tensor k_scale_inv, " - "Tensor v_scale_inv, Tensor cache_scale_inv, int cache_block_size, " - "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()"); - - // Kimi-K3 MLA decode epilogue: concat mqa_q = [ql_nope | q_pe] and insert the - // latent [kv_c_normed | k_pe] into the paged cache (bf16 / fp8 / fp8_ds_mla). - ops.def( - "fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(" - "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " - "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " - "int cache_block_size, Tensor? position_ids=None, " - "Tensor? cos_sin_cache=None) -> ()"); - ops.def( - "fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(" - "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " - "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " - "Tensor q_scale_inv, Tensor cache_scale_inv, int cache_block_size, " - "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()"); - ops.def( - "fused_kimi_k3_mla_decode_q_concat_ds_mla_insert(" - "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " - "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " - "int cache_block_size, Tensor? position_ids=None, " - "Tensor? cos_sin_cache=None) -> ()"); - -#ifndef USE_ROCM - ops.def( - "minimax_allreduce_rms_qk(" - "Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, " - "Tensor workspace, int q_size, int kv_size, int rank, int nranks, " - "float eps) -> (Tensor, Tensor)"); -#endif - - // Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE + KV-insert. - ops.def( - "fused_minimax_m3_qknorm_rope_kv_insert(" - "Tensor! qkv, Tensor q_norm_weight, Tensor k_norm_weight, " - "Tensor cos_sin_cache, Tensor positions, int num_heads, " - "int num_kv_heads, int rotary_dim, float eps, " - "Tensor? index_q_norm_weight, Tensor? index_k_norm_weight, " - "int num_index_heads, " - "Tensor? slot_mapping, Tensor? index_slot_mapping, " - "Tensor!? kv_cache, Tensor!? index_cache, " - "int block_size, Tensor!? q_out, Tensor!? index_q_out, " - "str kv_cache_dtype, bool skip_index_branch=False, " - "Tensor!? q_fp8_out=None, float q_fp8_scale=1.0, " - "Tensor? kv_k_scale=None, Tensor? kv_v_scale=None) -> ()"); - -#ifdef VLLM_ENABLE_FUSED_KDA_DECODE - ops.def( - "fused_kda_decode(" - "Tensor x, Tensor weight, Tensor? bias, Tensor! conv_state, " - "Tensor raw_g, Tensor raw_beta, Tensor A_log, Tensor dt_bias, " - "Tensor state_indices, Tensor! state, Tensor! out, " - "float? lower_bound=None, Tensor? output_gate=None, " - "Tensor? norm_weight=None, float norm_eps=1e-5) -> ()"); -#endif - -#ifdef VLLM_ENABLE_FUSED_GDN_DECODE - ops.def( - "fused_gdn_decode_post_conv_mtp(" - "Tensor mixed_qkv, Tensor a, Tensor b, Tensor A_log, Tensor dt_bias, " - "Tensor state_indices, Tensor cu_seqlens, Tensor num_accepted_tokens, " - "Tensor! state, Tensor output_gate, Tensor norm_weight, Tensor! out, " - "float scale, float norm_eps=1e-5, " - "str output_gate_activation='silu') -> ()"); -#endif - -#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK - ops.def( - "fused_kda_prologue(" - "Tensor q, Tensor k, Tensor v, Tensor raw_g, Tensor raw_beta, " - "Tensor A_log, Tensor dt_bias, Tensor! qg, Tensor! w, Tensor! u, " - "Tensor! kg_t, Tensor! aqk, Tensor! decay, Tensor cu_seqlens, " - "Tensor chunk_indices, Tensor? conv_weight, Tensor(e!)? conv_state, " - "Tensor? conv_state_indices, Tensor? conv_has_initial_state, " - "float scale, float lower_bound) -> ()"); - ops.def( - "fused_kda_chunk(" - "Tensor qg, Tensor w, Tensor u, Tensor kg_t, Tensor aqk, Tensor decay, " - "Tensor? initial_state, Tensor(a!)? final_state, Tensor! out, " - "Tensor cu_seqlens, Tensor chunk_offsets, float scale, " - "Tensor(b!)? group_state, int groups, " - "Tensor(c!)? checkpoint_state=None, Tensor? checkpoint_offsets=None, " - "Tensor? checkpoint_state_indices=None, Tensor(d!)? state_cache=None, " - "Tensor? state_indices=None, Tensor? has_initial_state=None) -> ()"); -#endif - -#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES - ops.def( - "kimi_k3_attn_res(" - "Tensor! prefix, Tensor? delta, Tensor! blocks, Tensor norm_weight, " - "Tensor qk_weight, Tensor? output_norm_weight, Tensor! output, " - "int num_blocks, " - "int block_write_idx, float eps, " - "float output_norm_eps) -> ()"); -#endif - - // Apply repetition penalties to logits in-place. - ops.def( - "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, " - "Tensor output_mask, Tensor repetition_penalties) -> ()"); - - // Optimized top-k per row operations. - ops.def( - "top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, " - "Tensor! indices, int numRows, int stride0, " - "int stride1, int topK) -> ()"); - - ops.def( - "top_k_per_row_decode(Tensor logits, int next_n, " - "Tensor seq_lens, Tensor! indices, " - "int numRows, int stride0, int stride1, int topK) -> ()"); - - ops.def( - "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, " - "Tensor workspace, int k, int max_seq_len) -> ()"); - -#ifdef VLLM_ENABLE_COOPERATIVE_TOPK - ops.def( - "cooperative_topk(Tensor logits, Tensor lengths, Tensor! output, " - "Tensor workspace, int k, int max_seq_len) -> ()"); -#endif - - // 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) -> ()"); - - ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()"); - - // SwiGLU activation with input clamping. - // alpha scales the sigmoid (gate * sigmoid(alpha * gate)); beta is added to - // the up half (up + beta). Defaults alpha=1.0, beta=0.0 give silu(gate)*up. - ops.def( - "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, " - "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) -> ()"); - - // Activation function used in GeGLU with `tanh` approximation. - ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()"); - - // FATReLU implementation. - ops.def("fatrelu_and_mul(Tensor! out, Tensor input, float threshold) -> ()"); - - ops.def( - "swigluoai_and_mul(Tensor! out, Tensor input, float alpha=1.702, float " - "limit=7.0) " - "-> ()"); - - // SituGLU implementation used in Kimi models. - ops.def( - "situ_and_mul(Tensor! out, Tensor input, float beta=1.0, float " - "linear_beta=-1.0) -> ()"); - // Fused SituGLU activation + dynamic FP8 quantization for the Humming w2 path - // (writes the fp8 down input and its float32 scale). group_size=0 -> - // per-token scale [.., 1]; group_size=128 -> k-major block-FP8 scale [.., - // d/128]. - ops.def( - "situ_and_mul_quant(Tensor! out, Tensor! scale, Tensor input, " - "float beta=1.0, float linear_beta=-1.0, int group_size=0, " - "Tensor? num_valid_tokens=None, int topk=1) -> ()"); - ops.def( - "masked_situ_and_mul(Tensor! out, Tensor input, Tensor " - "expert_num_tokens, float beta=1.0, float linear_beta=-1.0) -> ()"); - ops.def( - "masked_moe_activation(Tensor! out, Tensor input, Tensor " - "valid_token_counts, str activation, float clamp_limit=0.0, float " - "alpha=1.0, float beta=0.0, float situ_beta=1.0, float " - "situ_linear_beta=-1.0) -> ()"); - - // GELU implementation used in GPT-2. - ops.def("gelu_new(Tensor! out, Tensor input) -> ()"); - - // Approximate GELU implementation. - ops.def("gelu_fast(Tensor! out, Tensor input) -> ()"); - - // Quick GELU implementation. - ops.def("gelu_quick(Tensor! out, Tensor input) -> ()"); - - // relu(x)^2 activation from https://arxiv.org/abs/2109.08668v2 - ops.def("relu_squared(Tensor! out, Tensor input) -> ()"); - - // Compute int8 quantized tensor for given scaling factor. - ops.def( - "static_scaled_int8_quant(Tensor! result, Tensor input, Tensor scale," - "Tensor? azp) -> ()"); - - // Compute int8 quantized tensor and scaling factor - ops.def( - "dynamic_scaled_int8_quant(Tensor! result, Tensor input, Tensor! scale, " - "Tensor!? azp) -> ()"); - - // Compute FP8 quantized tensor for given scaling factor. - // Supports per-tensor, per-channel, per-token, and arbitrary 2D group - // scaling. Optional group_m/group_n specify the group shape explicitly; - // required for 1D scales to disambiguate per-channel vs per-token. - ops.def( - "static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale, " - "int[]? group_shape=None) -> ()"); - - // Compute dynamic-per-tensor FP8 quantized tensor and scaling factor. - ops.def( - "dynamic_scaled_fp8_quant(Tensor! result, Tensor input, Tensor! scale) " - "-> " - "()"); - - // Compute dynamic-per-token FP8 quantized tensor and scaling factor. - ops.def( - "dynamic_per_token_scaled_fp8_quant(Tensor! result, Tensor input, " - "Tensor! scale, Tensor? scale_ub) -> " - "()"); - - // Quantized GEMM for GPTQ. - // Note: even though the C++ inferred schema is correct for this op, it seems - // to prevent the meta function registry. - ops.def( - "gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, " - "Tensor b_gptq_scales, bool use_exllama, bool " - "use_v2_format, int bit) " - "-> Tensor"); - - // Post processing for GPTQ. - ops.def("gptq_shuffle(Tensor! q_weight, int bit) -> ()"); - - // Mamba selective scan kernel - ops.def( - "selective_scan_fwd(Tensor! u, Tensor! delta," - "Tensor! A, Tensor! B, Tensor! C," - "Tensor? D_, Tensor!? z_, Tensor? delta_bias_," - "bool delta_softplus," - "Tensor? query_start_loc," - "Tensor? cache_indices," - "Tensor? has_initial_state," - "Tensor! ssm_states," - "int null_block_id," - "int block_size," - "Tensor? block_idx_first_scheduled_token," - "Tensor? block_idx_last_scheduled_token," - "Tensor? initial_state_idx," - "Tensor? cu_chunk_seqlen," - "Tensor? last_chunk_indices) -> ()"); - - // LongCat n-gram embedding index kernel. All tensor args are marked mutable - // to match the (non-const) stable-Tensor& C++ signature; only ne_token_table - // and n_gram_ids are actually written in place. - // Fused vocab-parallel embedding: gather the rows this rank owns and write - // zeros for the rest, so the following all-reduce reconstructs the full - // embedding. - ops.def( - "vocab_parallel_embedding(Tensor! out, Tensor input_ids, Tensor weight, " - "int org_vocab_start_index, int org_vocab_end_index, " - "int num_org_vocab_padding, int added_vocab_start_index, " - "int added_vocab_end_index) -> ()"); - ops.def( - "ngram_compute_n_gram_ids(int ne_n, int ne_k, Tensor(a!) ne_weights, " - "Tensor(b!) ne_mods, Tensor(c!) exclusive_ne_embedder_size_sums, " - "Tensor(d!) exclusive_req_len_sums, Tensor(e!) ne_token_table, " - "Tensor(f!) row_indices, Tensor(g!) column_starts, " - "Tensor(h!) n_gram_ids) -> ()"); -} - -STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { - // LongCat n-gram embedding index kernel. - ops.impl("ngram_compute_n_gram_ids", TORCH_BOX(&ngram_compute_n_gram_ids)); - ops.impl("vocab_parallel_embedding", TORCH_BOX(&vocab_parallel_embedding)); - - // Per-token group quantization - ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8)); - ops.impl("per_token_group_fp8_quant_packed", - TORCH_BOX(&per_token_group_quant_8bit_packed)); - ops.impl("per_token_group_quant_int8", - TORCH_BOX(&per_token_group_quant_int8)); - -#ifndef USE_ROCM - // CUTLASS scaled_mm ops - ops.impl("cutlass_scaled_mm", TORCH_BOX(&cutlass_scaled_mm)); - ops.impl("cutlass_scaled_mm_azp", TORCH_BOX(&cutlass_scaled_mm_azp)); - ops.impl("cutlass_moe_mm", TORCH_BOX(&cutlass_moe_mm)); - ops.impl("get_cutlass_moe_mm_data", TORCH_BOX(&get_cutlass_moe_mm_data)); - ops.impl("get_cutlass_moe_mm_problem_sizes_from_expert_offsets", - TORCH_BOX(&get_cutlass_moe_mm_problem_sizes_from_expert_offsets)); - ops.impl("get_cutlass_batched_moe_mm_data", - TORCH_BOX(&get_cutlass_batched_moe_mm_data)); - - // FP4/NVFP4 ops - ops.impl("cutlass_scaled_fp4_mm", TORCH_BOX(&cutlass_scaled_fp4_mm)); - ops.impl("scaled_fp4_quant", TORCH_BOX(&scaled_fp4_quant_func)); - ops.impl("scaled_fp4_quant.out", TORCH_BOX(&scaled_fp4_quant_out)); - ops.impl("scaled_fp4_experts_quant", TORCH_BOX(&scaled_fp4_experts_quant)); - ops.impl("silu_and_mul_scaled_fp4_experts_quant", - TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant)); - ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant)); - // mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only). - // W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu. - - // AWQ ops - ops.impl("awq_gemm", TORCH_BOX(&awq_gemm)); - ops.impl("awq_dequantize", TORCH_BOX(&awq_dequantize)); - - // DSV3 fused A GEMM: conditionally compiled so impl registration is in - // source file (dsv3_fused_a_gemm.cu) -#endif - - ops.impl("merge_attn_states", TORCH_BOX(&merge_attn_states)); - - // Layernorm kernels (shared CUDA/ROCm) - ops.impl("rms_norm", TORCH_BOX(&rms_norm)); - ops.impl("fused_add_rms_norm", TORCH_BOX(&fused_add_rms_norm)); - - // Layernorm-quant kernels (shared CUDA/ROCm) - ops.impl("rms_norm_static_fp8_quant", TORCH_BOX(&rms_norm_static_fp8_quant)); - ops.impl("fused_add_rms_norm_static_fp8_quant", - TORCH_BOX(&fused_add_rms_norm_static_fp8_quant)); - - // Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm) - ops.impl("rms_norm_dynamic_per_token_quant", - TORCH_BOX(&rms_norm_dynamic_per_token_quant)); - ops.impl("rms_norm_per_block_quant", TORCH_BOX(&rms_norm_per_block_quant)); - ops.impl("silu_and_mul_per_block_quant", - TORCH_BOX(&silu_and_mul_per_block_quant)); - - // Positional encoding kernels (shared CUDA/ROCm) - ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding)); - ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope)); - ops.impl("fused_deepseek_v4_kv_rope_insert", - TORCH_BOX(&fused_deepseek_v4_kv_rope_insert)); - ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", - TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert)); - ops.impl( - "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert", - TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert)); - ops.impl( - "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", - TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert)); - ops.impl("fused_kimi_k3_mla_key_concat_kv_cache_insert", - TORCH_BOX(&fused_kimi_k3_mla_key_concat_kv_cache_insert)); - ops.impl("fused_kimi_k3_mla_key_concat_ds_mla_insert", - TORCH_BOX(&fused_kimi_k3_mla_key_concat_ds_mla_insert)); - ops.impl("fused_kimi_k3_mla_kv_concat", - TORCH_BOX(&fused_kimi_k3_mla_kv_concat)); - ops.impl("fused_kimi_k3_mla_kv_concat_quant_fp8", - TORCH_BOX(&fused_kimi_k3_mla_kv_concat_quant_fp8)); - ops.impl("fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert", - TORCH_BOX(&fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert)); - ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_insert", - TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_insert)); - ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert", - TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert)); - ops.impl("fused_kimi_k3_mla_decode_q_concat_ds_mla_insert", - TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_ds_mla_insert)); -#ifndef USE_ROCM - ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk)); -#endif - ops.impl("fused_minimax_m3_qknorm_rope_kv_insert", - TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert)); -#ifdef VLLM_ENABLE_FUSED_KDA_DECODE - ops.impl("fused_kda_decode", TORCH_BOX(&fused_kda_decode)); -#endif - -#ifdef VLLM_ENABLE_FUSED_GDN_DECODE - ops.impl("fused_gdn_decode_post_conv_mtp", - TORCH_BOX(&fused_gdn_decode_post_conv_mtp)); -#endif - -#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK - ops.impl("fused_kda_prologue", TORCH_BOX(&fused_kda_prologue)); - ops.impl("fused_kda_chunk", TORCH_BOX(&fused_kda_chunk)); -#endif - -#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES - ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res)); -#endif - - // Sampler kernels (shared CUDA/ROCm) - ops.impl("apply_repetition_penalties_", - TORCH_BOX(&apply_repetition_penalties_)); - ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill)); - ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode)); - ops.impl("persistent_topk", TORCH_BOX(&persistent_topk)); -#ifdef VLLM_ENABLE_COOPERATIVE_TOPK - ops.impl("cooperative_topk", TORCH_BOX(&cooperative_topk)); -#endif - - // 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)); - ops.impl("gelu_tanh_and_mul", TORCH_BOX(&gelu_tanh_and_mul)); - ops.impl("fatrelu_and_mul", TORCH_BOX(&fatrelu_and_mul)); - ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul)); - ops.impl("situ_and_mul", TORCH_BOX(&situ_and_mul)); - ops.impl("situ_and_mul_quant", TORCH_BOX(&situ_and_mul_quant)); - ops.impl("masked_situ_and_mul", TORCH_BOX(&masked_situ_and_mul)); - ops.impl("masked_moe_activation", TORCH_BOX(&masked_moe_activation)); - ops.impl("gelu_new", TORCH_BOX(&gelu_new)); - ops.impl("gelu_fast", TORCH_BOX(&gelu_fast)); - ops.impl("gelu_quick", TORCH_BOX(&gelu_quick)); - ops.impl("relu_squared", TORCH_BOX(&relu_squared)); - ops.impl("silu_and_mul_with_clamp", TORCH_BOX(&silu_and_mul_clamp)); - - // INT8 quantization kernels - ops.impl("static_scaled_int8_quant", TORCH_BOX(&static_scaled_int8_quant)); - ops.impl("dynamic_scaled_int8_quant", TORCH_BOX(&dynamic_scaled_int8_quant)); - - // FP8 quantization kernels - ops.impl("static_scaled_fp8_quant", TORCH_BOX(&static_scaled_fp8_quant)); - ops.impl("dynamic_scaled_fp8_quant", TORCH_BOX(&dynamic_scaled_fp8_quant)); - ops.impl("dynamic_per_token_scaled_fp8_quant", - TORCH_BOX(&dynamic_per_token_scaled_fp8_quant)); - - // GPTQ kernels - ops.impl("gptq_gemm", TORCH_BOX(&gptq_gemm)); - ops.impl("gptq_shuffle", TORCH_BOX(&gptq_shuffle)); - - // Mamba kernels - ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) { - ops.impl("get_cuda_view_from_cpu_tensor", - TORCH_BOX(&get_cuda_view_from_cpu_tensor)); -} - -STABLE_TORCH_LIBRARY_FRAGMENT(_C_cuda_utils, cuda_utils) { - cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int"); - cuda_utils.def( - "get_max_shared_memory_per_block_device_attribute(int device_id) -> int"); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd, - cuda_utils) { - cuda_utils.impl("get_device_attribute", TORCH_BOX(&get_device_attribute)); - cuda_utils.impl("get_max_shared_memory_per_block_device_attribute", - TORCH_BOX(&get_max_shared_memory_per_block_device_attribute)); -} - -// These capability-check functions take only primitive args (no tensors), so -// there is no device to dispatch on. CompositeExplicitAutograd makes them -// available for all backends. This is the stable ABI equivalent of calling -// ops.impl("op_name", &func) without a dispatch key in the non-stable API. -STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) { -#ifndef USE_ROCM - ops.impl("cutlass_scaled_mm_supports_fp8", - TORCH_BOX(&cutlass_scaled_mm_supports_fp8)); - ops.impl("cutlass_group_gemm_supported", - TORCH_BOX(&cutlass_group_gemm_supported)); - ops.impl("cutlass_scaled_mm_supports_block_fp8", - TORCH_BOX(&cutlass_scaled_mm_supports_block_fp8)); - ops.impl("cutlass_scaled_mm_supports_fp4", - TORCH_BOX(&cutlass_scaled_mm_supports_fp4)); -#endif -} - -// Cache ops -STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { - // Swap in (out) the cache blocks from src to dst. - ops.def( - "swap_blocks(Tensor src, Tensor! dst," - " int block_size_in_bytes, Tensor block_mapping) -> ()"); - - // Batch swap: submit all block copies in a single driver call. - ops.def( - "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs," - " Tensor sizes," - " bool is_src_access_order_any=False) -> ()"); - - // Reshape the key and value tensors and cache them. - ops.def( - "reshape_and_cache(Tensor key, Tensor value," - " Tensor! key_cache, Tensor! value_cache," - " Tensor slot_mapping," - " str kv_cache_dtype," - " Tensor k_scale, Tensor v_scale) -> ()"); - - // Reshape the key and value tensors and cache them. - ops.def( - "reshape_and_cache_flash(Tensor key, Tensor value," - " Tensor! key_cache," - " Tensor! value_cache," - " Tensor slot_mapping," - " str kv_cache_dtype," - " Tensor k_scale, Tensor v_scale) -> ()"); - - // Concat kv_c and k_pe and cache them. - ops.def( - "concat_and_cache_mla(Tensor kv_c, Tensor k_pe," - " Tensor! kv_cache," - " Tensor slot_mapping," - " str kv_cache_dtype," - " Tensor scale) -> ()"); - - // Grouped concat_and_cache_mla across all layers. Each layer's cache base - // pointer and optional plain-FP8 scale are read from device tensors. - ops.def( - "concat_and_cache_mla_grouped(Tensor kv_c, Tensor k_pe," - " Tensor kv_cache_ptrs," - " Tensor slot_mapping," - " int block_size, int block_stride," - " int entry_stride," - " Tensor? kv_scales=None," - " str kv_cache_dtype='auto') -> ()"); - -#ifndef USE_ROCM - ops.def( - "hisparse_resolve_residency(Tensor host_cache," - " Tensor! hot_cache," - " Tensor hot_block_table," - " Tensor global_indices," - " Tensor! hot_indices," - " Tensor! device_global_indices," - " Tensor! lru_slots," - " Tensor? request_state_indices," - " int region_stride," - " int max_union_rows," - " Tensor(a!)? miss_mask=None," - " Tensor(b!)? stats=None," - " Tensor(c!)? attention_indices=None," - " int attention_block_stride=0," - " Tensor? request_ids=None," - " Tensor? source_block_table=None," - " int source_block_size=0," - " Tensor(d!)? resolved_global_indices=None," - " Tensor(e!)? valid_counts=None," - " Tensor(f!)? swap_host_physical_rows=None," - " Tensor(g!)? swap_device_physical_rows=None," - " Tensor(h!)? swap_counts=None," - " Tensor? resident_block_table=None," - " int resident_block_size=0," - " int resident_null_block=0) -> ()"); - - ops.def( - "hisparse_invalidate_written_slots(Tensor! device_global_indices," - " Tensor request_state_indices," - " Tensor req_id_per_token," - " Tensor written_slots) -> ()"); - - ops.def( - "hisparse_gather_plan(Tensor host_cache," - " Tensor! hot_cache," - " Tensor global_indices," - " Tensor hot_indices," - " Tensor miss_mask," - " Tensor? request_state_indices," - " Tensor(a!)? attention_indices=None," - " int attention_block_stride=0) -> ()"); - - ops.def( - "hisparse_gather_compact(Tensor host_cache," - " Tensor! hot_cache," - " Tensor miss_global_indices," - " Tensor miss_hot_indices," - " Tensor miss_counts) -> ()"); - -#endif // !USE_ROCM - - // Rotate Q and K, then write to kv cache for MLA - ops.def( - "concat_and_cache_mla_rope_fused(" - " Tensor positions," - " Tensor! q_pe," - " Tensor! k_pe," - " Tensor kv_c," - " Tensor cos_sin_cache," - " bool is_neox," - " Tensor slot_mapping," - " Tensor! kv_cache," - " str kv_cache_dtype," - " Tensor kv_cache_scale) -> ()"); - - // Convert the key and value cache to fp8 data type. - ops.def( - "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, " - "str kv_cache_dtype) -> ()"); - - // Gather cache blocks from src_cache to dst, dequantizing from - // src_cache's dtype to dst's dtype if necessary. - ops.def( - "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, " - " Tensor block_table, Tensor cu_seq_lens, " - " Tensor token_to_seq, " - " int num_tokens, " - " str kv_cache_dtype, " - " Tensor scale, Tensor? seq_starts) -> ()"); - - ops.def( - "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, " - "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()"); - - ops.def( - "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, " - "Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? " - "seq_starts, Tensor? host_cache=None, Tensor? host_row_ids=None, Tensor? " - "device_row_ids=None) -> ()"); - - ops.def( - "cp_gather_and_upconvert_nvfp4_kv_cache(Tensor src_cache, Tensor! dst, " - "Tensor block_table, Tensor workspace_starts, int batch_size) -> ()"); - - ops.def( - "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor " - "slot_mapping, " - "int quant_block_size, str kv_cache_dtype) -> ()"); - - ops.def("concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()"); - - ops.def( - "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! " - "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()"); -} - -STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ar) { - custom_ar.def( - "init_custom_ar(int[] ipc_tensors, Tensor rank_data, " - "int rank, bool fully_connected) -> int"); - custom_ar.def( - "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, " - "int reg_buffer_sz_bytes) -> ()"); - custom_ar.def("dispose(int fa) -> ()"); - custom_ar.def("meta_size() -> int"); - custom_ar.def("register_buffer(int fa, int[] ipc_tensors) -> ()"); - custom_ar.def("get_graph_buffer_ipc_meta(int fa) -> (int[], int[])"); - custom_ar.def( - "register_graph_buffers(int fa, int[][] handles, int[][] offsets) -> ()"); - custom_ar.def("allocate_shared_buffer_and_handle(int size) -> (int, Tensor)"); - custom_ar.def("open_mem_handle(Tensor mem_handle) -> int"); - custom_ar.def("free_shared_buffer(int ptr) -> ()"); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ar) { - custom_ar.impl("init_custom_ar", TORCH_BOX(&init_custom_ar)); - custom_ar.impl("all_reduce", TORCH_BOX(&all_reduce)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CPU, custom_ar) { - custom_ar.impl("open_mem_handle", TORCH_BOX(&open_mem_handle)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CompositeExplicitAutograd, custom_ar) { - custom_ar.impl("dispose", TORCH_BOX(&dispose)); - custom_ar.impl("meta_size", TORCH_BOX(&meta_size)); - custom_ar.impl("register_buffer", TORCH_BOX(®ister_buffer)); - custom_ar.impl("get_graph_buffer_ipc_meta", - TORCH_BOX(&get_graph_buffer_ipc_meta)); - custom_ar.impl("register_graph_buffers", TORCH_BOX(®ister_graph_buffers)); - custom_ar.impl("allocate_shared_buffer_and_handle", - TORCH_BOX(&allocate_shared_buffer_and_handle)); - custom_ar.impl("free_shared_buffer", TORCH_BOX(&free_shared_buffer)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) { - ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) { - ops.impl("swap_blocks", TORCH_BOX(&swap_blocks)); - ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache)); - ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash)); - ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla)); - ops.impl("concat_and_cache_mla_grouped", - TORCH_BOX(&concat_and_cache_mla_grouped)); - -#ifndef USE_ROCM - ops.impl("hisparse_resolve_residency", - TORCH_BOX(&hisparse_resolve_residency)); - ops.impl("hisparse_invalidate_written_slots", - TORCH_BOX(&hisparse_invalidate_written_slots)); - ops.impl("hisparse_gather_plan", TORCH_BOX(&hisparse_gather_plan)); - ops.impl("hisparse_gather_compact", TORCH_BOX(&hisparse_gather_compact)); -#endif // !USE_ROCM - ops.impl("concat_and_cache_mla_rope_fused", - TORCH_BOX(&concat_and_cache_mla_rope_fused)); - ops.impl("convert_fp8", TORCH_BOX(&convert_fp8)); - ops.impl("gather_and_maybe_dequant_cache", - TORCH_BOX(&gather_and_maybe_dequant_cache)); - ops.impl("cp_gather_cache", TORCH_BOX(&cp_gather_cache)); - ops.impl("cp_gather_and_upconvert_fp8_kv_cache", - TORCH_BOX(&cp_gather_and_upconvert_fp8_kv_cache)); - ops.impl("cp_gather_and_upconvert_nvfp4_kv_cache", - TORCH_BOX(&cp_gather_and_upconvert_nvfp4_kv_cache)); - ops.impl("indexer_k_quant_and_cache", TORCH_BOX(&indexer_k_quant_and_cache)); - ops.impl("concat_mla_q", TORCH_BOX(&concat_mla_q)); - ops.impl("cp_gather_indexer_k_quant_cache", - TORCH_BOX(&cp_gather_indexer_k_quant_cache)); -} - -REGISTER_EXTENSION(_C_stable_libtorch) diff --git a/csrc/libtorch_stable/torch_utils.h b/csrc/libtorch_stable/torch_utils.h deleted file mode 100644 index f02346739e0a..000000000000 --- a/csrc/libtorch_stable/torch_utils.h +++ /dev/null @@ -1,88 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include - -#include -#include - -#include -#include -#include -#include - -// Stable ABI equivalent of TORCH_CHECK_NOT_IMPLEMENTED. -#define STD_TORCH_CHECK_NOT_IMPLEMENTED(cond, ...) \ - STD_TORCH_CHECK(cond, "NotImplementedError: ", __VA_ARGS__) - -// Device properties cache for stable ABI compatibility. -// Uses raw CUDA/HIP APIs instead of ATen functions. -// Using inline ensures a single instance across all translation units. -inline std::deque device_flags; -inline std::vector device_properties; -inline std::once_flag vectors_init_flag; - -inline void do_init_device_vectors() { - int device_count; - cudaError_t err = cudaGetDeviceCount(&device_count); - if (err != cudaSuccess) { - STD_TORCH_CHECK(false, "cudaGetDeviceCount failed: " + - std::string(cudaGetErrorString(err))); - } - device_flags.resize(device_count); - device_properties.resize(device_count); -} - -inline void initDeviceVectors() { - std::call_once(vectors_init_flag, do_init_device_vectors); -} - -inline void initDeviceProperty(int device_index) { - cudaDeviceProp device_prop{}; - cudaError_t err = cudaGetDeviceProperties(&device_prop, device_index); - if (err != cudaSuccess) { - STD_TORCH_CHECK(false, "cudaGetDeviceProperties failed: " + - std::string(cudaGetErrorString(err))); - } - device_properties[device_index] = device_prop; -} - -// Get device properties using raw CUDA/HIP APIs (stable ABI compatible). -// Caches results per device so cudaGetDeviceProperties is called at most once -// per device. -inline cudaDeviceProp* get_device_prop() { - initDeviceVectors(); - int device_index; - cudaError_t err = cudaGetDevice(&device_index); - if (err != cudaSuccess) { - STD_TORCH_CHECK( - false, "cudaGetDevice failed: " + std::string(cudaGetErrorString(err))); - } - STD_TORCH_CHECK(device_index >= 0 && static_cast(device_index) < - device_properties.size(), - "CUDA device index " + std::to_string(device_index) + - " out of range [0, " + - std::to_string(device_properties.size()) + ")"); - - std::call_once(device_flags[device_index], initDeviceProperty, device_index); - return &device_properties[device_index]; -} - -// Utility to get the current CUDA stream for a given device using stable APIs. -// Returns a cudaStream_t for use in kernel launches. -inline cudaStream_t get_current_cuda_stream(int32_t device_index = -1) { - void* stream_ptr = nullptr; - TORCH_ERROR_CODE_CHECK( - aoti_torch_get_current_cuda_stream(device_index, &stream_ptr)); - return reinterpret_cast(stream_ptr); -} - -// Utility to get the current cuBLAS handle using stable APIs. -inline cublasHandle_t get_current_cuda_blas_handle() { - void* blas_handle_ptr = nullptr; - TORCH_ERROR_CODE_CHECK(torch_get_current_cuda_blas_handle(&blas_handle_ptr)); - return reinterpret_cast(blas_handle_ptr); -} diff --git a/csrc/libtorch_stable/mamba/selective_scan.h b/csrc/mamba/selective_scan.h similarity index 100% rename from csrc/libtorch_stable/mamba/selective_scan.h rename to csrc/mamba/selective_scan.h diff --git a/csrc/libtorch_stable/mamba/selective_scan_fwd.cu b/csrc/mamba/selective_scan_fwd.cu similarity index 99% rename from csrc/libtorch_stable/mamba/selective_scan_fwd.cu rename to csrc/mamba/selective_scan_fwd.cu index c17b06a79bd0..cfc1f038f131 100644 --- a/csrc/libtorch_stable/mamba/selective_scan_fwd.cu +++ b/csrc/mamba/selective_scan_fwd.cu @@ -14,7 +14,7 @@ #endif #include "selective_scan.h" -#include "static_switch.h" +#include "../dispatch_utils.h" template; constexpr int kSmemSize = Ktraits::kSmemSize + kNRows * MAX_DSTATE * sizeof(typename Ktraits::scan_t); dim3 grid(params.batch, params.dim / kNRows); diff --git a/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu b/csrc/minimax_reduce_rms_kernel.cu similarity index 99% rename from csrc/libtorch_stable/minimax_reduce_rms_kernel.cu rename to csrc/minimax_reduce_rms_kernel.cu index 58d61b353d6d..5a3da3dadffc 100644 --- a/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu +++ b/csrc/minimax_reduce_rms_kernel.cu @@ -27,7 +27,7 @@ #include #include -#include "cuda_compat.h" +#include "core/cuda_compat.h" #include "minimax_reduce_rms_kernel.h" #include diff --git a/csrc/libtorch_stable/minimax_reduce_rms_kernel.h b/csrc/minimax_reduce_rms_kernel.h similarity index 100% rename from csrc/libtorch_stable/minimax_reduce_rms_kernel.h rename to csrc/minimax_reduce_rms_kernel.h diff --git a/csrc/libtorch_stable/moe/grouped_topk_kernels.cu b/csrc/moe/grouped_topk_kernels.cu similarity index 99% rename from csrc/libtorch_stable/moe/grouped_topk_kernels.cu rename to csrc/moe/grouped_topk_kernels.cu index 852b2b468adf..ea89aa38a1f0 100644 --- a/csrc/libtorch_stable/moe/grouped_topk_kernels.cu +++ b/csrc/moe/grouped_topk_kernels.cu @@ -22,7 +22,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/.gitignore b/csrc/moe/marlin_moe_wna16/.gitignore similarity index 73% rename from csrc/libtorch_stable/moe/marlin_moe_wna16/.gitignore rename to csrc/moe/marlin_moe_wna16/.gitignore index 7dc482a89466..ba805f9250ec 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/.gitignore +++ b/csrc/moe/marlin_moe_wna16/.gitignore @@ -1,3 +1,2 @@ sm*_kernel_*.cu kernel_selector.h -kernel_*.cu diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h b/csrc/moe/marlin_moe_wna16/kernel.h similarity index 95% rename from csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h rename to csrc/moe/marlin_moe_wna16/kernel.h index 25c57f3b2ddf..1fd00c12d1bd 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h +++ b/csrc/moe/marlin_moe_wna16/kernel.h @@ -3,8 +3,8 @@ #define MARLIN_NAMESPACE_NAME marlin_moe_wna16 #endif -#include "libtorch_stable/quantization/marlin/marlin.cuh" -#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" +#include "quantization/marlin/marlin.cuh" +#include "quantization/marlin/marlin_dtypes.cuh" #include "core/scalar_type.hpp" #define MARLIN_KERNEL_PARAMS \ diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h b/csrc/moe/marlin_moe_wna16/marlin_template.h similarity index 99% rename from csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h rename to csrc/moe/marlin_moe_wna16/marlin_template.h index aac9bd157f6a..e2d80ce6bfbe 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h +++ b/csrc/moe/marlin_moe_wna16/marlin_template.h @@ -23,10 +23,10 @@ #define MARLIN_NAMESPACE_NAME marlin_moe_wna16 #endif -#include "libtorch_stable/quantization/marlin/marlin.cuh" -#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" -#include "libtorch_stable/quantization/marlin/dequant.h" -#include "libtorch_stable/quantization/marlin/marlin_mma.h" +#include "quantization/marlin/marlin.cuh" +#include "quantization/marlin/marlin_dtypes.cuh" +#include "quantization/marlin/dequant.h" +#include "quantization/marlin/marlin_mma.h" #include "core/scalar_type.hpp" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu b/csrc/moe/marlin_moe_wna16/ops.cu similarity index 99% rename from csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu rename to csrc/moe/marlin_moe_wna16/ops.cu index 48de07dc753d..8336b46e1590 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu +++ b/csrc/moe/marlin_moe_wna16/ops.cu @@ -32,7 +32,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ static_assert(std::is_same::value || \ diff --git a/csrc/libtorch_stable/moe/moeTopKFuncs.cuh b/csrc/moe/moeTopKFuncs.cuh similarity index 100% rename from csrc/libtorch_stable/moe/moeTopKFuncs.cuh rename to csrc/moe/moeTopKFuncs.cuh diff --git a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu b/csrc/moe/moe_align_sum_kernels.cu similarity index 99% rename from csrc/libtorch_stable/moe/moe_align_sum_kernels.cu rename to csrc/moe/moe_align_sum_kernels.cu index 5c9a522feb7b..a2424f833908 100644 --- a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu +++ b/csrc/moe/moe_align_sum_kernels.cu @@ -8,11 +8,11 @@ #include #include -#include "../../cuda_compat.h" -#include "libtorch_stable/core/math.hpp" -#include "libtorch_stable/dispatch_utils.h" -#include "libtorch_stable/quantization/vectorization.cuh" -#include "libtorch_stable/torch_utils.h" +#include "core/cuda_compat.h" +#include "core/utils.hpp" +#include "dispatch_utils.h" +#include "quantization/vectorization.cuh" +#include "torch_utils.h" #define CEILDIV(x, y) (((x) + (y) - 1) / (y)) diff --git a/csrc/libtorch_stable/moe/moe_ops.h b/csrc/moe/moe_ops.h similarity index 100% rename from csrc/libtorch_stable/moe/moe_ops.h rename to csrc/moe/moe_ops.h diff --git a/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu b/csrc/moe/moe_permute_unpermute_op.cu similarity index 99% rename from csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu rename to csrc/moe/moe_permute_unpermute_op.cu index 957cd60e49de..31c937b2a2bb 100644 --- a/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu +++ b/csrc/moe/moe_permute_unpermute_op.cu @@ -7,8 +7,8 @@ #include #include "core/registration.h" -#include "libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h" -#include "libtorch_stable/torch_utils.h" +#include "moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h" +#include "torch_utils.h" #include diff --git a/csrc/libtorch_stable/moe/moe_wna16.cu b/csrc/moe/moe_wna16.cu similarity index 99% rename from csrc/libtorch_stable/moe/moe_wna16.cu rename to csrc/moe/moe_wna16.cu index 9345a7c9f789..a8971041d041 100644 --- a/csrc/libtorch_stable/moe/moe_wna16.cu +++ b/csrc/moe/moe_wna16.cu @@ -8,7 +8,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include "moe_wna16_utils.h" #define DIVIDE(x, size) (((x) + (size) - 1) / (size)) diff --git a/csrc/libtorch_stable/moe/moe_wna16_utils.h b/csrc/moe/moe_wna16_utils.h similarity index 100% rename from csrc/libtorch_stable/moe/moe_wna16_utils.h rename to csrc/moe/moe_wna16_utils.h diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/dispatch.h b/csrc/moe/permute_unpermute_kernels/dispatch.h similarity index 100% rename from csrc/libtorch_stable/moe/permute_unpermute_kernels/dispatch.h rename to csrc/moe/permute_unpermute_kernels/dispatch.h diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu similarity index 98% rename from csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu rename to csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu index 8d8bdf8baaf6..56c1e2bddf48 100644 --- a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu +++ b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu @@ -1,7 +1,7 @@ #include #include -#include "libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h" +#include "moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h" // moe_permute kernels require at least CUDA 12.0 #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000) diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h similarity index 97% rename from csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h rename to csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h index 53afad2f5d41..1496eccb8d4d 100644 --- a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h +++ b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h @@ -10,7 +10,7 @@ #include "cutlass/array.h" #include "cutlass/numeric_size.h" -#include "libtorch_stable/moe/permute_unpermute_kernels/dispatch.h" +#include "moe/permute_unpermute_kernels/dispatch.h" template inline T* get_ptr(torch::stable::Tensor& t) { diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl similarity index 100% rename from csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl rename to csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl diff --git a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu b/csrc/moe/topk_softmax_kernels.cu similarity index 99% rename from csrc/libtorch_stable/moe/topk_softmax_kernels.cu rename to csrc/moe/topk_softmax_kernels.cu index 098644a1b901..0e462770177c 100644 --- a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu +++ b/csrc/moe/topk_softmax_kernels.cu @@ -24,9 +24,9 @@ #include #include -#include "../../cuda_compat.h" +#include "core/cuda_compat.h" #include "../cub_helpers.h" -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #ifndef USE_ROCM #include diff --git a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu b/csrc/moe/topk_softplus_sqrt_kernels.cu similarity index 99% rename from csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu rename to csrc/moe/topk_softplus_sqrt_kernels.cu index fbb07216d202..237f76334abd 100644 --- a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu +++ b/csrc/moe/topk_softplus_sqrt_kernels.cu @@ -25,9 +25,9 @@ #include #include -#include "../../cuda_compat.h" +#include "core/cuda_compat.h" #include "../cub_helpers.h" -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #ifndef USE_ROCM #include #include diff --git a/csrc/libtorch_stable/moe/torch_bindings.cpp b/csrc/moe/torch_bindings.cpp similarity index 100% rename from csrc/libtorch_stable/moe/torch_bindings.cpp rename to csrc/moe/torch_bindings.cpp diff --git a/csrc/libtorch_stable/ngram_embedding_kernels.cu b/csrc/ngram_embedding_kernels.cu similarity index 100% rename from csrc/libtorch_stable/ngram_embedding_kernels.cu rename to csrc/ngram_embedding_kernels.cu diff --git a/csrc/libtorch_stable/nvfp4_ds_mla_cache.h b/csrc/nvfp4_ds_mla_cache.h similarity index 100% rename from csrc/libtorch_stable/nvfp4_ds_mla_cache.h rename to csrc/nvfp4_ds_mla_cache.h diff --git a/csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu b/csrc/nvfp4_ds_mla_cache_kernels.cu similarity index 100% rename from csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu rename to csrc/nvfp4_ds_mla_cache_kernels.cu diff --git a/csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu b/csrc/nvfp4_kv_cache_kernels.cu similarity index 99% rename from csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu rename to csrc/nvfp4_kv_cache_kernels.cu index bb11e474a6a4..341c4cdd9396 100644 --- a/csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu +++ b/csrc/nvfp4_kv_cache_kernels.cu @@ -15,10 +15,10 @@ // - reciprocal_approximate_ftz() for fast reciprocal #define NVFP4_ENABLE_ELTS16 1 -#include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh" +#include "quantization/fp4/nvfp4_utils.cuh" -#include "libtorch_stable/dispatch_utils.h" -#include "libtorch_stable/torch_utils.h" +#include "dispatch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/ops.h b/csrc/ops.h index d7ee2d080c48..3266d28b1edb 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -1,72 +1,822 @@ #pragma once +#include +#include +#include + #include #include -#include -#include +#include -#include "core/scalar_type.hpp" +#include -#include +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"); -// 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. -void rms_norm(torch::Tensor& out, torch::Tensor& input, - std::optional weight, double epsilon); + // Get the raw data pointer + void* data_ptr = tensor.mutable_data_ptr(); -void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual, - std::optional weight, double epsilon); + /// 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()); +} -// rotary_embedding also exist in csrc/libtorch_stable/ops.h (torch::stable -// ABI for CUDA). It remains here because the CPU build still uses these -// torch::Tensor declarations. -void rotary_embedding(torch::Tensor& positions, torch::Tensor& query, - std::optional key, int64_t head_size, - torch::Tensor& cos_sin_cache, bool is_neox, - int64_t rope_dim_offset, bool inverse); +void per_token_group_quant_fp8(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s, + int64_t group_size, double eps, double fp8_min, + double fp8_max, bool scale_ue8m0, + bool dummy_is_scale_transposed, + bool dummy_is_tma_aligned); -void silu_and_mul(torch::Tensor& out, torch::Tensor& input); +// Fused activation quantisation + DeepGEMM-compatible UE8M0-packed scales. +void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s_packed, + int64_t group_size, double eps, + double min_8bit, double max_8bit); -void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit, - double alpha = 1.0, double beta = 0.0); +void per_token_group_quant_int8(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s, + int64_t group_size, double eps, double int8_min, + double int8_max); -void gelu_and_mul(torch::Tensor& out, torch::Tensor& input); +#ifndef USE_ROCM +bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability); +bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability); +bool cutlass_group_gemm_supported(int64_t cuda_device_capability); -void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input); +void cutlass_scaled_mm(torch::stable::Tensor& out, + torch::stable::Tensor const& a, + torch::stable::Tensor const& b, + torch::stable::Tensor const& a_scales, + torch::stable::Tensor const& b_scales, + std::optional const& bias); -void gelu_tanh(torch::Tensor& out, torch::Tensor& input); +void cutlass_moe_mm(torch::stable::Tensor& out_tensors, + torch::stable::Tensor const& a_tensors, + torch::stable::Tensor const& b_tensors, + torch::stable::Tensor const& a_scales, + torch::stable::Tensor const& b_scales, + torch::stable::Tensor const& expert_offsets, + torch::stable::Tensor const& problem_sizes, + torch::stable::Tensor const& a_strides, + torch::stable::Tensor const& b_strides, + torch::stable::Tensor const& c_strides, bool per_act_token, + bool per_out_ch); -void gelu_new(torch::Tensor& out, torch::Tensor& input); +void cutlass_scaled_mm_azp(torch::stable::Tensor& out, + torch::stable::Tensor const& a, + torch::stable::Tensor const& b, + torch::stable::Tensor const& a_scales, + torch::stable::Tensor const& b_scales, + torch::stable::Tensor const& azp_adj, + std::optional const& azp, + std::optional const& bias); -void gelu_fast(torch::Tensor& out, torch::Tensor& input); +void get_cutlass_moe_mm_data( + const torch::stable::Tensor& topk_ids, + torch::stable::Tensor& expert_offsets, + torch::stable::Tensor& problem_sizes1, + torch::stable::Tensor& problem_sizes2, + torch::stable::Tensor& input_permutation, + torch::stable::Tensor& output_permutation, const int64_t num_experts, + const int64_t n, const int64_t k, + const std::optional& blockscale_offsets, + const bool is_gated); -void gelu_quick(torch::Tensor& out, torch::Tensor& input); +void get_cutlass_moe_mm_problem_sizes_from_expert_offsets( + const torch::stable::Tensor& expert_first_token_offset, + torch::stable::Tensor& problem_sizes1, + torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k, + const bool swap_ab); -void relu_squared(torch::Tensor& out, torch::Tensor& input); +void get_cutlass_batched_moe_mm_data( + torch::stable::Tensor& expert_offsets, + torch::stable::Tensor& problem_sizes1, + torch::stable::Tensor& problem_sizes2, + const torch::stable::Tensor& expert_num_tokens, + const int64_t num_local_experts, const int64_t padded_m, const int64_t n, + const int64_t k); -void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, - torch::Tensor const& scale, - std::optional const& azp); +// FP4/NVFP4 ops +bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability); -void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, - torch::Tensor& scales, - std::optional const& azp); +void cutlass_scaled_fp4_mm(torch::stable::Tensor& D, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& A_sf, + torch::stable::Tensor const& B_sf, + torch::stable::Tensor const& alpha); -torch::Tensor dynamic_4bit_int_moe_cpu( - torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights, - torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size, - int64_t intermediate_size, int64_t group_size, - bool apply_router_weight_on_input, int64_t activation_kind); +void cutlass_fp4_group_mm(torch::stable::Tensor& output, + const torch::stable::Tensor& a, + const torch::stable::Tensor& b, + const torch::stable::Tensor& a_blockscale, + const torch::stable::Tensor& b_blockscales, + const torch::stable::Tensor& alphas, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& sf_offsets); -using fptr_t = int64_t; -#ifdef USE_ROCM -fptr_t init_custom_qr(int64_t rank, int64_t world_size, - std::optional qr_max_size = std::nullopt); -void qr_destroy(fptr_t _fa); -torch::Tensor qr_get_handle(fptr_t _fa); -void qr_open_handles(fptr_t _fa, const std::vector& handles); -void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, - int64_t quant_level, bool cast_bf2half = false); -int64_t qr_max_size(); +std::tuple scaled_fp4_quant_func( + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_scale, bool is_sf_swizzled_layout); + +void scaled_fp4_quant_out(torch::stable::Tensor const& input, + torch::stable::Tensor const& input_scale, + bool is_sf_swizzled_layout, + torch::stable::Tensor& output, + torch::stable::Tensor& output_scale); + +void scaled_fp4_experts_quant( + torch::stable::Tensor& output, torch::stable::Tensor& output_scale, + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_global_scale, + torch::stable::Tensor const& input_offset_by_experts, + torch::stable::Tensor const& output_scale_offset_by_experts); + +void silu_and_mul_scaled_fp4_experts_quant( + torch::stable::Tensor& output, torch::stable::Tensor& output_scale, + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_global_scale, + torch::stable::Tensor const& input_offset_by_experts, + torch::stable::Tensor const& output_scale_offset_by_experts); + +void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out, + torch::stable::Tensor& output_block_scale, + torch::stable::Tensor& input, + torch::stable::Tensor& input_global_scale); + +void cutlass_mxfp4_group_mm(torch::stable::Tensor& output, + const torch::stable::Tensor& a, + const torch::stable::Tensor& b, + const torch::stable::Tensor& a_blockscale, + const torch::stable::Tensor& b_blockscales, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& sf_offsets); + +// AWQ ops +torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats, + torch::stable::Tensor _kernel, + torch::stable::Tensor _scaling_factors, + torch::stable::Tensor _zeros, + int64_t split_k_iters); + +torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel, + torch::stable::Tensor _scaling_factors, + torch::stable::Tensor _zeros, + int64_t split_k_iters, int64_t thx, + int64_t thy); + +// DSV3 fused A GEMM: conditionally compiled so declaration and impl +// registration are in the source file (dsv3_fused_a_gemm.cu) + +#endif + +// CPU tensor -> CUDA UVA view (shared CUDA/ROCm) +torch::stable::Tensor get_cuda_view_from_cpu_tensor( + torch::stable::Tensor& cpu_tensor); + +// Attention kernels (shared CUDA/ROCm) +void merge_attn_states( + torch::stable::Tensor& output, + std::optional output_lse, + const torch::stable::Tensor& prefix_output, + const torch::stable::Tensor& prefix_lse, + const torch::stable::Tensor& suffix_output, + const torch::stable::Tensor& suffix_lse, + const std::optional prefill_tokens_with_context, + const std::optional& output_scale = std::nullopt); + +torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x, + bool inplace); + +// Layernorm kernels (shared CUDA/ROCm) +void rms_norm(torch::stable::Tensor& out, torch::stable::Tensor& input, + std::optional weight, double epsilon); + +void fused_add_rms_norm(torch::stable::Tensor& input, + torch::stable::Tensor& residual, + std::optional weight, + double epsilon); + +// Layernorm-quant kernels (shared CUDA/ROCm) +void rms_norm_static_fp8_quant(torch::stable::Tensor& out, + torch::stable::Tensor& input, + torch::stable::Tensor& weight, + torch::stable::Tensor& scale, double epsilon); + +void fused_add_rms_norm_static_fp8_quant(torch::stable::Tensor& out, + torch::stable::Tensor& input, + torch::stable::Tensor& residual, + torch::stable::Tensor& weight, + torch::stable::Tensor& scale, + double epsilon); + +// Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm) +void rms_norm_dynamic_per_token_quant( + torch::stable::Tensor& out, torch::stable::Tensor const& input, + torch::stable::Tensor const& weight, torch::stable::Tensor& scales, + double const var_epsilon, std::optional scale_ub, + std::optional residual); + +void rms_norm_per_block_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor const& weight, + torch::stable::Tensor& scales, + double const var_epsilon, + std::optional scale_ub, + std::optional residual, + int64_t group_size, bool is_scale_transposed); + +void silu_and_mul_per_block_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor& scales, + int64_t group_size, + std::optional scale_ub, + bool is_scale_transposed); + +// Positional encoding kernels (shared CUDA/ROCm) +void rotary_embedding(torch::stable::Tensor& positions, + torch::stable::Tensor& query, + std::optional key, + int64_t head_size, torch::stable::Tensor& cos_sin_cache, + bool is_neox, int64_t rope_dim_offset, bool inverse); + +void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q, + int64_t num_heads_k, int64_t num_heads_v, + int64_t head_dim, double eps, + torch::stable::Tensor& q_weight, + torch::stable::Tensor& k_weight, + torch::stable::Tensor& cos_sin_cache, bool is_neox, + torch::stable::Tensor& position_ids, + int64_t forced_token_heads_per_warp); + +void fused_deepseek_v4_kv_rope_insert( + torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size, + std::optional fp8_scale, bool kv_mxfp8); + +torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv, + torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, + double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8, + bool apply_q_rope, bool is_q_interleaved); + +void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( + torch::stable::Tensor& q, torch::stable::Tensor const& kv, + torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, double eps, + int64_t cache_block_size, bool apply_q_norm); + +void fused_kimi_k3_mla_key_concat_kv_cache_insert( + torch::stable::Tensor& q, torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_key_concat_ds_mla_insert( + torch::stable::Tensor& q, torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_kv_concat(torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, + torch::stable::Tensor& k_out); + +void fused_kimi_k3_mla_kv_concat_quant_fp8(torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, + torch::stable::Tensor const& v, + torch::stable::Tensor& k_fp8, + torch::stable::Tensor& v_fp8); + +void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( + torch::stable::Tensor const& q, torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor const& v, torch::stable::Tensor& q_fp8, + torch::stable::Tensor& k_fp8, torch::stable::Tensor& v_fp8, + torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& q_scale_inv, + torch::stable::Tensor const& k_scale_inv, + torch::stable::Tensor const& v_scale_inv, + torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert( + torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, + torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert( + torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, + torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& q_scale_inv, + torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_decode_q_concat_ds_mla_insert( + torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, + torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( + torch::stable::Tensor const& q, torch::stable::Tensor const& kv, + torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, + torch::stable::Tensor const& fp8_scale, + torch::stable::Tensor const& q_fp8_scale_inv, double eps, + int64_t cache_block_size, bool apply_q_norm); + +#ifndef USE_ROCM +std::tuple +minimax_allreduce_rms_qk(torch::stable::Tensor qkv, + torch::stable::Tensor const& norm_weight_q, + torch::stable::Tensor const& norm_weight_k, + torch::stable::Tensor workspace, int64_t const q_size, + int64_t const kv_size, int64_t const rank, + int64_t const nranks, double const eps); +#endif + +// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV / +// index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs +// the index branch and scatters k/v/index_k into their paged caches. +void fused_minimax_m3_qknorm_rope_kv_insert( + torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight, + torch::stable::Tensor const& k_norm_weight, + torch::stable::Tensor const& cos_sin_cache, + torch::stable::Tensor const& positions, int64_t num_heads, + int64_t num_kv_heads, int64_t rotary_dim, double eps, + std::optional index_q_norm_weight, + std::optional index_k_norm_weight, + int64_t num_index_heads, std::optional slot_mapping, + std::optional index_slot_mapping, + std::optional kv_cache, + std::optional index_cache, int64_t block_size, + std::optional q_out, + std::optional index_q_out, + const std::string& kv_cache_dtype, bool skip_index_branch, + std::optional q_fp8_out, double q_fp8_scale, + std::optional kv_k_scale, + std::optional kv_v_scale); + +#ifdef VLLM_ENABLE_FUSED_KDA_DECODE +void fused_kda_decode( + torch::stable::Tensor const& x, torch::stable::Tensor const& weight, + std::optional bias, + torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g, + torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, + torch::stable::Tensor const& dt_bias, + torch::stable::Tensor const& state_indices, torch::stable::Tensor& state, + torch::stable::Tensor& out, std::optional lower_bound, + std::optional output_gate, + std::optional norm_weight, double norm_eps); + +#endif + +#ifdef VLLM_ENABLE_FUSED_GDN_DECODE +void fused_gdn_decode_post_conv_mtp( + torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a, + torch::stable::Tensor const& b, torch::stable::Tensor const& a_log, + torch::stable::Tensor const& dt_bias, + torch::stable::Tensor const& state_indices, + torch::stable::Tensor const& cu_seqlens, + torch::stable::Tensor const& num_accepted_tokens, + torch::stable::Tensor& state, torch::stable::Tensor const& output_gate, + torch::stable::Tensor const& norm_weight, torch::stable::Tensor& out, + double scale, double norm_eps, const std::string& output_gate_activation); + +#endif + +#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK +void fused_kda_prologue( + torch::stable::Tensor const& q, torch::stable::Tensor const& k, + torch::stable::Tensor const& v, torch::stable::Tensor const& raw_g, + torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, + torch::stable::Tensor const& dt_bias, torch::stable::Tensor& qg, + torch::stable::Tensor& w, torch::stable::Tensor& u, + torch::stable::Tensor& kg_t, torch::stable::Tensor& aqk, + torch::stable::Tensor& decay, torch::stable::Tensor const& cu_seqlens, + torch::stable::Tensor const& chunk_indices, + std::optional conv_weight, + std::optional conv_state, + std::optional conv_state_indices, + std::optional conv_has_initial_state, double scale, + double lower_bound); + +void fused_kda_chunk( + torch::stable::Tensor const& qg, torch::stable::Tensor const& w, + torch::stable::Tensor const& u, torch::stable::Tensor const& kg_t, + torch::stable::Tensor const& aqk, torch::stable::Tensor const& decay, + std::optional initial_state, + std::optional final_state, + torch::stable::Tensor& out, torch::stable::Tensor const& cu_seqlens, + torch::stable::Tensor const& chunk_offsets, double scale, + std::optional group_state, int64_t groups, + std::optional checkpoint_state, + std::optional checkpoint_offsets, + std::optional checkpoint_state_indices, + std::optional state_cache, + std::optional state_indices, + std::optional has_initial_state); +#endif + +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES +void kimi_k3_attn_res(torch::stable::Tensor& prefix, + std::optional delta, + torch::stable::Tensor& blocks, + torch::stable::Tensor const& norm_weight, + torch::stable::Tensor const& qk_weight, + std::optional output_norm_weight, + torch::stable::Tensor& output, int64_t num_blocks, + int64_t block_write_idx, double eps, + double output_norm_eps); +#endif + +// Sampler kernels (shared CUDA/ROCm) +void apply_repetition_penalties_( + torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask, + const torch::stable::Tensor& output_mask, + const torch::stable::Tensor& repetition_penalties); + +void top_k_per_row_prefill(const torch::stable::Tensor& logits, + const torch::stable::Tensor& rowStarts, + const torch::stable::Tensor& rowEnds, + torch::stable::Tensor& indices, int64_t numRows, + int64_t stride0, int64_t stride1, int64_t topK); + +void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n, + const torch::stable::Tensor& seqLens, + torch::stable::Tensor& indices, int64_t numRows, + int64_t stride0, int64_t stride1, int64_t topK); + +void persistent_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, int64_t k, + int64_t max_seq_len); + +#ifdef VLLM_ENABLE_COOPERATIVE_TOPK +void cooperative_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, int64_t k, + int64_t max_seq_len); #endif + +void selective_scan_fwd( + const torch::stable::Tensor& u, const torch::stable::Tensor& delta, + const torch::stable::Tensor& A, const torch::stable::Tensor& B, + const torch::stable::Tensor& C, + const std::optional& D_, + const std::optional& z_, + const std::optional& delta_bias_, + bool delta_softplus, + const std::optional& query_start_loc, + const std::optional& cache_indices, + const std::optional& has_initial_state, + const torch::stable::Tensor& ssm_states, int64_t null_block_id, + int64_t block_size, + const std::optional& block_idx_first_scheduled_token, + const std::optional& block_idx_last_scheduled_token, + const std::optional& initial_state_idx, + const std::optional& cu_chunk_seqlen, + const std::optional& last_chunk_indices); + +using fptr_t = int64_t; +fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, + torch::stable::Tensor& rank_data, int64_t rank, + bool fully_connected); +void all_reduce(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes); +void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes); +void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t local_buffer, + fptr_t multicast_buffer, fptr_t epoch_buffer, + int64_t stage_sz_bytes); +void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes); +void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, + fptr_t local_buffer, fptr_t epoch_buffer, + int64_t stage_sz_bytes); +void mnnvl_multimem_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, + fptr_t local_buffer, fptr_t multicast_buffer, + int64_t stage_sz_bytes, int64_t block_limit); +void dispose(fptr_t _fa); +int64_t meta_size(); +void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs); +std::tuple, std::vector> +get_graph_buffer_ipc_meta(fptr_t _fa); +void register_graph_buffers(fptr_t _fa, + const std::vector>& handles, + const std::vector>& offsets); +std::tuple allocate_shared_buffer_and_handle( + int64_t size); +int64_t open_mem_handle(torch::stable::Tensor& mem_handle); +void free_shared_buffer(int64_t buffer); + +// Activation kernels (shared CUDA/ROCm) +void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); +void silu_and_mul_clamp(torch::stable::Tensor& out, + torch::stable::Tensor& input, double limit, + 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, + torch::stable::Tensor& input); +void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, + double threshold); +void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, + double alpha = 1.702, double limit = 7.0); +void situ_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, + double beta = 1.0, double linear_beta = -1.0); +void situ_and_mul_quant( + torch::stable::Tensor& out, torch::stable::Tensor& scale, + torch::stable::Tensor& input, double beta = 1.0, double linear_beta = -1.0, + int64_t group_size = 0, + std::optional num_valid_tokens = std::nullopt, + int64_t topk = 1); +void masked_situ_and_mul(torch::stable::Tensor& out, + torch::stable::Tensor& input, + const torch::stable::Tensor& expert_num_tokens, + double beta = 1.0, double linear_beta = -1.0); +void masked_moe_activation(torch::stable::Tensor& out, + torch::stable::Tensor& input, + const torch::stable::Tensor& valid_token_counts, + const std::string& activation, + double clamp_limit = 0.0, double alpha = 1.0, + double beta = 0.0, double situ_beta = 1.0, + double situ_linear_beta = -1.0); +void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input); +void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input); +void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input); + +void relu_squared(torch::stable::Tensor& out, torch::stable::Tensor& input); + +// INT8 quantization kernels (shared CUDA/ROCm) +void static_scaled_int8_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor const& scale, + std::optional const& azp); + +void dynamic_scaled_int8_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor& scales, + std::optional const& azp); + +// FP8 quantization kernels (shared CUDA/ROCm) +void static_scaled_fp8_quant( + torch::stable::Tensor& out, torch::stable::Tensor const& input, + torch::stable::Tensor const& scale, + std::optional group_shape = + std::nullopt); + +void dynamic_scaled_fp8_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor& scale); + +void dynamic_per_token_scaled_fp8_quant( + torch::stable::Tensor& out, torch::stable::Tensor const& input, + torch::stable::Tensor& scale, + std::optional const& scale_ub); + +// GPTQ kernels (shared CUDA/ROCm) +torch::stable::Tensor gptq_gemm(torch::stable::Tensor a, + torch::stable::Tensor b_q_weight, + torch::stable::Tensor b_gptq_qzeros, + torch::stable::Tensor b_gptq_scales, + bool use_exllama, bool use_v2_format, + int64_t bit); + +void gptq_shuffle(torch::stable::Tensor q_weight, int64_t bit); + +// Cache ops (shared CUDA/ROCm) +void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst, + int64_t block_size_in_bytes, + const torch::stable::Tensor& block_mapping); + +// Batch swap: submit all block copies in a single driver call. +void swap_blocks_batch(const torch::stable::Tensor& src_ptrs, + const torch::stable::Tensor& dst_ptrs, + const torch::stable::Tensor& sizes, + bool is_src_access_order_any); + +void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value, + torch::stable::Tensor& key_cache, + torch::stable::Tensor& value_cache, + torch::stable::Tensor& slot_mapping, + const std::string& kv_cache_dtype, + torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale); + +void reshape_and_cache_flash( + torch::stable::Tensor& key, torch::stable::Tensor& value, + torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, + torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype, + torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale); + +void concat_and_cache_mla(torch::stable::Tensor& kv_c, + torch::stable::Tensor& k_pe, + torch::stable::Tensor& kv_cache, + torch::stable::Tensor& slot_mapping, + const std::string& kv_cache_dtype, + torch::stable::Tensor& scale); + +void concat_and_cache_mla_grouped( + torch::stable::Tensor& kv_c, torch::stable::Tensor& k_pe, + torch::stable::Tensor& kv_cache_ptrs, torch::stable::Tensor& slot_mapping, + int64_t block_size, int64_t block_stride, int64_t entry_stride, + std::optional kv_scales, + const std::string& kv_cache_dtype); + +#ifndef USE_ROCM +// HiSparse kernels use raw PTX in the row copy; CUDA-only. +void hisparse_resolve_residency( + torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& hot_block_table, + torch::stable::Tensor const& global_indices, + torch::stable::Tensor& hot_indices, + torch::stable::Tensor& device_global_indices, + torch::stable::Tensor& lru_slots, + std::optional const& request_state_indices, + int64_t region_stride, int64_t max_union_rows, + std::optional const& miss_mask, + std::optional const& stats, + std::optional const& attention_indices, + int64_t attention_block_stride, + std::optional const& request_ids, + std::optional const& source_block_table, + int64_t source_block_size, + std::optional const& resolved_global_indices, + std::optional const& valid_counts, + std::optional const& swap_host_physical_rows, + std::optional const& swap_device_physical_rows, + std::optional const& swap_counts, + std::optional const& resident_block_table, + int64_t resident_block_size, int64_t resident_null_block); + +void hisparse_invalidate_written_slots( + torch::stable::Tensor& device_global_indices, + torch::stable::Tensor const& request_state_indices, + torch::stable::Tensor const& req_id_per_token, + torch::stable::Tensor const& written_slots); + +void hisparse_gather_plan( + torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& global_indices, + torch::stable::Tensor const& hot_indices, + torch::stable::Tensor const& miss_mask, + std::optional const& request_state_indices, + std::optional const& attention_indices, + int64_t attention_block_stride); + +void hisparse_gather_compact(torch::stable::Tensor const& host_cache, + torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& miss_global_indices, + torch::stable::Tensor const& miss_hot_indices, + torch::stable::Tensor const& miss_counts); + +#endif // !USE_ROCM + +// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla +void concat_and_cache_mla_rope_fused( + torch::stable::Tensor& positions, torch::stable::Tensor& q_pe, + torch::stable::Tensor& k_pe, torch::stable::Tensor& kv_c, + torch::stable::Tensor& rope_cos_sin_cache, bool rope_is_neox, + torch::stable::Tensor& slot_mapping, torch::stable::Tensor& kv_cache, + const std::string& kv_cache_dtype, + torch::stable::Tensor& kv_cache_quant_scale); + +// Just for unittest +void convert_fp8(torch::stable::Tensor& dst_cache, + torch::stable::Tensor& src_cache, const double scale, + const std::string& kv_cache_dtype); + +void gather_and_maybe_dequant_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // ENTRIES...] + torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] + torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] + int64_t num_tokens, const std::string& kv_cache_dtype, + torch::stable::Tensor const& scale, + std::optional seq_starts = std::nullopt); + +// TODO(hc): cp_gather_cache need support scaled kvcahe in the future. +void cp_gather_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // ENTRIES...] + torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] + int64_t batch_size, + std::optional seq_starts = std::nullopt); + +// Gather and upconvert FP8 KV cache to BF16 workspace +void cp_gather_and_upconvert_fp8_kv_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // 656] + torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& workspace_starts, // [BATCH] + int64_t batch_size, + std::optional seq_starts = std::nullopt, + std::optional host_cache = std::nullopt, + std::optional host_row_ids = std::nullopt, + std::optional device_row_ids = std::nullopt); + +// Gather and upconvert an nvfp4_ds_mla KV cache to a BF16 workspace +void cp_gather_and_upconvert_nvfp4_kv_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // 352] + torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& workspace_starts, // [BATCH] + int64_t batch_size); + +// Indexer K quantization and cache function +void indexer_k_quant_and_cache( + torch::stable::Tensor& k, // [num_tokens, head_dim] + torch::stable::Tensor& kv_cache, // [num_blocks, block_size, + // cache_stride] + torch::stable::Tensor& slot_mapping, // [num_tokens] + int64_t quant_block_size, // quantization block size + const std::string& scale_fmt); + +// Concatenate query nope and rope for MLA/DSA attention +void concat_mla_q( + torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] + torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] + torch::stable::Tensor& q_out); // [num_tokens, num_heads, nope_dim + + // rope_dim] + +// Extract function to gather quantized K cache +void cp_gather_indexer_k_quant_cache( + const torch::stable::Tensor& kv_cache, // [num_blocks, block_size, + // cache_stride] + torch::stable::Tensor& dst_k, // [num_tokens, head_dim] + torch::stable::Tensor& dst_scale, // [num_tokens, head_dim / + // quant_block_size * 4] + const torch::stable::Tensor& block_table, // [batch_size, num_blocks] + const torch::stable::Tensor& cu_seq_lens); // [batch_size + 1] + +// Fused vocab-parallel embedding lookup (see +// vocab_parallel_embedding_kernels.cu). +void vocab_parallel_embedding( + torch::stable::Tensor& out, const torch::stable::Tensor& input_ids, + const torch::stable::Tensor& weight, int64_t org_vocab_start_index, + int64_t org_vocab_end_index, int64_t num_org_vocab_padding, + int64_t added_vocab_start_index, int64_t added_vocab_end_index); + +// LongCat n-gram embedding index kernel (see ngram_embedding_kernels.cu). +void ngram_compute_n_gram_ids( + int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights, + torch::stable::Tensor& ne_mods, + torch::stable::Tensor& exclusive_ne_embedder_size_sums, + torch::stable::Tensor& exclusive_req_len_sums, + torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices, + torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids); diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/persistent_topk.cuh similarity index 100% rename from csrc/libtorch_stable/persistent_topk.cuh rename to csrc/persistent_topk.cuh diff --git a/csrc/libtorch_stable/pos_encoding_kernels.cu b/csrc/pos_encoding_kernels.cu similarity index 99% rename from csrc/libtorch_stable/pos_encoding_kernels.cu rename to csrc/pos_encoding_kernels.cu index 74af743b0961..bad616fc8748 100644 --- a/csrc/libtorch_stable/pos_encoding_kernels.cu +++ b/csrc/pos_encoding_kernels.cu @@ -1,6 +1,6 @@ #include "torch_utils.h" -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include "dispatch_utils.h" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/activation_kernels.cu b/csrc/quantization/activation_kernels.cu similarity index 99% rename from csrc/libtorch_stable/quantization/activation_kernels.cu rename to csrc/quantization/activation_kernels.cu index 822a41969e7e..9d7da85a22ee 100644 --- a/csrc/libtorch_stable/quantization/activation_kernels.cu +++ b/csrc/quantization/activation_kernels.cu @@ -1,10 +1,10 @@ -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include -#include "libtorch_stable/core/math.hpp" -#include "cuda_compat.h" -#include "libtorch_stable/dispatch_utils.h" +#include "core/utils.hpp" +#include "core/cuda_compat.h" +#include "dispatch_utils.h" #include "quantization/w8a8/fp8/common.cuh" #ifndef USE_ROCM diff --git a/csrc/libtorch_stable/quantization/awq/dequantize.cuh b/csrc/quantization/awq/dequantize.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/awq/dequantize.cuh rename to csrc/quantization/awq/dequantize.cuh diff --git a/csrc/libtorch_stable/quantization/awq/gemm_kernels.cu b/csrc/quantization/awq/gemm_kernels.cu similarity index 99% rename from csrc/libtorch_stable/quantization/awq/gemm_kernels.cu rename to csrc/quantization/awq/gemm_kernels.cu index 43301f0fe206..a0be65704544 100644 --- a/csrc/libtorch_stable/quantization/awq/gemm_kernels.cu +++ b/csrc/quantization/awq/gemm_kernels.cu @@ -9,9 +9,9 @@ Shang and Dang, Xingyu and Han, Song}, journal={arXiv}, year={2023} #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" -#include "libtorch_stable/quantization/awq/dequantize.cuh" +#include "quantization/awq/dequantize.cuh" #include diff --git a/csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh b/csrc/quantization/cutlass_w4a8/get_group_starts.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh rename to csrc/quantization/cutlass_w4a8/get_group_starts.cuh index 721761791251..aff0502550e1 100644 --- a/csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh +++ b/csrc/quantization/cutlass_w4a8/get_group_starts.cuh @@ -3,7 +3,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include "cutlass/bfloat16.h" #include "cutlass/float8.h" diff --git a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu b/csrc/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu similarity index 98% rename from csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu rename to csrc/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu index 1d7eb04093a2..82d6b2ec8a89 100644 --- a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu +++ b/csrc/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu @@ -16,12 +16,12 @@ // vllm includes #include #include -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "torch_utils.h" +#include "cutlass_extensions/torch_utils.hpp" +#include "cutlass_extensions/common.hpp" #include "get_group_starts.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" #include "w4a8_utils.cuh" namespace vllm::cutlass_w4a8_moe { diff --git a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu b/csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu similarity index 98% rename from csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu rename to csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu index 8f19394b3904..c4e96e8aee22 100644 --- a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu +++ b/csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu @@ -5,8 +5,8 @@ #include #include -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" +#include "torch_utils.h" +#include "cutlass_extensions/torch_utils.hpp" #include "w4a8_utils.cuh" #include "cutlass/cutlass.h" @@ -21,8 +21,8 @@ #include "cutlass/util/packed_stride.hpp" #include "cutlass/util/mixed_dtype_utils.hpp" -#include "libtorch_stable/cutlass_extensions/common.hpp" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/common.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" #include diff --git a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cu b/csrc/quantization/cutlass_w4a8/w4a8_utils.cu similarity index 100% rename from csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cu rename to csrc/quantization/cutlass_w4a8/w4a8_utils.cu diff --git a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cuh b/csrc/quantization/cutlass_w4a8/w4a8_utils.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cuh rename to csrc/quantization/cutlass_w4a8/w4a8_utils.cuh diff --git a/csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu b/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu similarity index 97% rename from csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu rename to csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu index 6238a27191b7..2064775111fb 100644 --- a/csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu +++ b/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu @@ -15,8 +15,8 @@ */ #include -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/dispatch_utils.h" +#include "torch_utils.h" +#include "dispatch_utils.h" #include "../../cuda_vec_utils.cuh" #include @@ -25,7 +25,7 @@ #include #include "cuda_utils.h" -#include "libtorch_stable/launch_bounds_utils.h" +#include "launch_bounds_utils.h" // Define before including nvfp4_utils.cuh so the header // can use this macro during compilation. diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu b/csrc/quantization/fp4/mxfp4_blockwise_moe_kernel.cu similarity index 99% rename from csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu rename to csrc/quantization/fp4/mxfp4_blockwise_moe_kernel.cu index e4d2f2201250..3234717f6be1 100644 --- a/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu +++ b/csrc/quantization/fp4/mxfp4_blockwise_moe_kernel.cu @@ -8,11 +8,11 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" #include "cute/tensor.hpp" #include "cutlass/tensor_ref.h" diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu b/csrc/quantization/fp4/mxfp4_experts_quant.cu similarity index 99% rename from csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu rename to csrc/quantization/fp4/mxfp4_experts_quant.cu index 086c07acf002..4c6875f0623a 100644 --- a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu +++ b/csrc/quantization/fp4/mxfp4_experts_quant.cu @@ -25,9 +25,9 @@ #include #include -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/dispatch_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "torch_utils.h" +#include "dispatch_utils.h" +#include "cutlass_extensions/common.hpp" #include "../../cuda_vec_utils.cuh" #include "cuda_utils.h" @@ -41,7 +41,7 @@ static_assert(CVT_FP4_ELTS_PER_THREAD == 16, #define VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED 0 #endif -#include "libtorch_stable/launch_bounds_utils.h" +#include "launch_bounds_utils.h" #if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu b/csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu similarity index 99% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu rename to csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu index ea3344406352..8e8a30e956d7 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu +++ b/csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu @@ -16,11 +16,11 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" #include "cute/tensor.hpp" #include "cutlass/tensor_ref.h" diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu b/csrc/quantization/fp4/nvfp4_experts_quant.cu similarity index 99% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu rename to csrc/quantization/fp4/nvfp4_experts_quant.cu index 6296f0d901cb..55c722f1de98 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu +++ b/csrc/quantization/fp4/nvfp4_experts_quant.cu @@ -15,8 +15,8 @@ */ #include -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/dispatch_utils.h" +#include "torch_utils.h" +#include "dispatch_utils.h" #include "../../cuda_vec_utils.cuh" #include @@ -26,7 +26,7 @@ #include "cuda_utils.h" #include "nvfp4_utils.cuh" -#include "libtorch_stable/launch_bounds_utils.h" +#include "launch_bounds_utils.h" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu b/csrc/quantization/fp4/nvfp4_quant_entry.cu similarity index 98% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu rename to csrc/quantization/fp4/nvfp4_quant_entry.cu index e1e7e7a74da2..ace7bbbc10c7 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu +++ b/csrc/quantization/fp4/nvfp4_quant_entry.cu @@ -16,9 +16,9 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" #include "nvfp4_utils.cuh" #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \ diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu b/csrc/quantization/fp4/nvfp4_quant_kernels.cu similarity index 98% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu rename to csrc/quantization/fp4/nvfp4_quant_kernels.cu index 7ec418ee381b..13a3026cc889 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu +++ b/csrc/quantization/fp4/nvfp4_quant_kernels.cu @@ -21,12 +21,12 @@ #include -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/dispatch_utils.h" +#include "torch_utils.h" +#include "dispatch_utils.h" #include "../../cuda_vec_utils.cuh" #include "cuda_utils.h" -#include "libtorch_stable/launch_bounds_utils.h" +#include "launch_bounds_utils.h" // Define before including nvfp4_utils.cuh so the header // can use this macro during compilation. diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu b/csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu similarity index 97% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu rename to csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu index bfb526fcd401..6a60fe6f6044 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu +++ b/csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu @@ -16,9 +16,9 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" #if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100 void cutlass_scaled_fp4_mm_sm100a(torch::stable::Tensor& D, diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu b/csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu similarity index 98% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu rename to csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu index af9f24a70e0b..f925a599db46 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu +++ b/csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu @@ -16,9 +16,9 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" #include "cutlass/cutlass.h" @@ -31,8 +31,7 @@ #include "cutlass/util/packed_stride.hpp" -#include "libtorch_stable/core/math.hpp" -#include "core/batch_invariant.hpp" +#include "core/utils.hpp" using namespace cute; diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu b/csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu similarity index 98% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu rename to csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu index 3a45ede8dfd5..02ded13c7835 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu +++ b/csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu @@ -16,9 +16,9 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" #include "cutlass/cutlass.h" @@ -31,8 +31,7 @@ #include "cutlass/util/packed_stride.hpp" -#include "libtorch_stable/core/math.hpp" -#include "core/batch_invariant.hpp" +#include "core/utils.hpp" using namespace cute; diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh b/csrc/quantization/fp4/nvfp4_utils.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh rename to csrc/quantization/fp4/nvfp4_utils.cuh diff --git a/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu b/csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu similarity index 99% rename from csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu rename to csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu index 56dd47038727..f864d6037ef3 100644 --- a/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu +++ b/csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu @@ -2,7 +2,7 @@ #include "../../torch_utils.h" #include "../../dispatch_utils.h" -#include "../../../core/batch_invariant.hpp" +#include "core/utils.hpp" #include "layernorm_utils.cuh" #include "quant_conversions.cuh" diff --git a/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu b/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu similarity index 100% rename from csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu rename to csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu diff --git a/csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh b/csrc/quantization/fused_kernels/layernorm_utils.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh rename to csrc/quantization/fused_kernels/layernorm_utils.cuh index 9a3d221562e9..d633cb9a6e57 100644 --- a/csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh +++ b/csrc/quantization/fused_kernels/layernorm_utils.cuh @@ -4,12 +4,12 @@ * __device__ layernorm utilities. */ -#include "libtorch_stable/quantization/vectorization.cuh" +#include "quantization/vectorization.cuh" #include "quantization/utils.cuh" #include "quant_conversions.cuh" #include "../../cub_helpers.h" -#include "../../../cuda_compat.h" +#include "core/cuda_compat.h" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/fused_kernels/quant_conversions.cuh b/csrc/quantization/fused_kernels/quant_conversions.cuh similarity index 96% rename from csrc/libtorch_stable/quantization/fused_kernels/quant_conversions.cuh rename to csrc/quantization/fused_kernels/quant_conversions.cuh index a81a0b9feb2f..1953e90f5de0 100644 --- a/csrc/libtorch_stable/quantization/fused_kernels/quant_conversions.cuh +++ b/csrc/quantization/fused_kernels/quant_conversions.cuh @@ -4,9 +4,9 @@ * __device__ helper functions to deal with float -> quant datatype conversion */ -#include "libtorch_stable/quantization/vectorization.cuh" +#include "quantization/vectorization.cuh" // TODO(luka/varun):refactor common.cuh to use this file instead -#include "../../../quantization/w8a8/fp8/common.cuh" +#include "quantization/w8a8/fp8/common.cuh" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/gptq/compat.cuh b/csrc/quantization/gptq/compat.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/compat.cuh rename to csrc/quantization/gptq/compat.cuh diff --git a/csrc/libtorch_stable/quantization/gptq/matrix_view.cuh b/csrc/quantization/gptq/matrix_view.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/matrix_view.cuh rename to csrc/quantization/gptq/matrix_view.cuh diff --git a/csrc/libtorch_stable/quantization/gptq/q_gemm.cu b/csrc/quantization/gptq/q_gemm.cu similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/q_gemm.cu rename to csrc/quantization/gptq/q_gemm.cu diff --git a/csrc/libtorch_stable/quantization/gptq/qdq_2.cuh b/csrc/quantization/gptq/qdq_2.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/qdq_2.cuh rename to csrc/quantization/gptq/qdq_2.cuh diff --git a/csrc/libtorch_stable/quantization/gptq/qdq_3.cuh b/csrc/quantization/gptq/qdq_3.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/qdq_3.cuh rename to csrc/quantization/gptq/qdq_3.cuh diff --git a/csrc/libtorch_stable/quantization/gptq/qdq_4.cuh b/csrc/quantization/gptq/qdq_4.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/qdq_4.cuh rename to csrc/quantization/gptq/qdq_4.cuh diff --git a/csrc/libtorch_stable/quantization/gptq/qdq_8.cuh b/csrc/quantization/gptq/qdq_8.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/qdq_8.cuh rename to csrc/quantization/gptq/qdq_8.cuh diff --git a/csrc/libtorch_stable/quantization/gptq/qdq_util.cuh b/csrc/quantization/gptq/qdq_util.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/qdq_util.cuh rename to csrc/quantization/gptq/qdq_util.cuh diff --git a/csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu b/csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu similarity index 99% rename from csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu rename to csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu index a5d199db2737..2575b20b8eef 100644 --- a/csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu +++ b/csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu @@ -11,8 +11,8 @@ Redistribution and use in source and binary forms, with or without modification, THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS “AS IS” AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. ***********/ -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/dispatch_utils.h" +#include "torch_utils.h" +#include "dispatch_utils.h" #include #include diff --git a/csrc/libtorch_stable/quantization/machete/Readme.md b/csrc/quantization/machete/Readme.md similarity index 100% rename from csrc/libtorch_stable/quantization/machete/Readme.md rename to csrc/quantization/machete/Readme.md diff --git a/csrc/libtorch_stable/quantization/machete/generate.py b/csrc/quantization/machete/generate.py similarity index 99% rename from csrc/libtorch_stable/quantization/machete/generate.py rename to csrc/quantization/machete/generate.py index 11a5bbdd13c1..60a49c587fa4 100644 --- a/csrc/libtorch_stable/quantization/machete/generate.py +++ b/csrc/quantization/machete/generate.py @@ -496,7 +496,8 @@ def prepacked_type_key(prepack_type: PrepackTypeConfig): def generate(): - # See csrc/quantization/machete/Readme.md, the Codegeneration for more info + # See csrc/quantization/machete/Readme.md, the + # Codegeneration for more info # about how this works SCRIPT_DIR = os.path.dirname(__file__) diff --git a/csrc/libtorch_stable/quantization/machete/machete_collective_builder.cuh b/csrc/quantization/machete/machete_collective_builder.cuh similarity index 94% rename from csrc/libtorch_stable/quantization/machete/machete_collective_builder.cuh rename to csrc/quantization/machete/machete_collective_builder.cuh index ba8da0af2c31..ee825583dee1 100644 --- a/csrc/libtorch_stable/quantization/machete/machete_collective_builder.cuh +++ b/csrc/quantization/machete/machete_collective_builder.cuh @@ -1,6 +1,6 @@ #pragma once -#include "libtorch_stable/cutlass_extensions/vllm_collective_builder.cuh" +#include "cutlass_extensions/vllm_collective_builder.cuh" #include "machete_mainloop.cuh" namespace cutlass::gemm::collective { diff --git a/csrc/libtorch_stable/quantization/machete/machete_interleaving_utils.cuh b/csrc/quantization/machete/machete_interleaving_utils.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/machete/machete_interleaving_utils.cuh rename to csrc/quantization/machete/machete_interleaving_utils.cuh diff --git a/csrc/libtorch_stable/quantization/machete/machete_mainloop.cuh b/csrc/quantization/machete/machete_mainloop.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/machete/machete_mainloop.cuh rename to csrc/quantization/machete/machete_mainloop.cuh index a0d1ded8876a..9127a701a0fd 100644 --- a/csrc/libtorch_stable/quantization/machete/machete_mainloop.cuh +++ b/csrc/quantization/machete/machete_mainloop.cuh @@ -17,8 +17,9 @@ // before feeding them into the tensor cores in registers, we need the weight // matrix to be A. To achieve this we compute the transpose of Y = XW^t as // Y^t = W^tX^t. This is mostly done outside of this file in -// csrc/quantization/machete/machete_mm_kernel.cuh, but this why A is the -// quantized/narrow type and has the prepacked layout despite the API being: +// csrc/quantization/machete/machete_mm_kernel.cuh, but this why +// A is the quantized/narrow type and has the prepacked layout despite the API +// being: // B_prepacked = machete_prepack_B(B) // Y = machete_mm(A, B_prepacked) // diff --git a/csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh b/csrc/quantization/machete/machete_mm_kernel.cuh similarity index 98% rename from csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh rename to csrc/quantization/machete/machete_mm_kernel.cuh index 57655f369cdf..db3321a39db8 100644 --- a/csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh +++ b/csrc/quantization/machete/machete_mm_kernel.cuh @@ -18,9 +18,9 @@ // clang-format on #include "cutlass_extensions/cute_utils.cuh" -#include "libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" -#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" +#include "cutlass_extensions/vllm_numeric_conversion.cuh" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/torch_utils.hpp" #include "machete_collective_builder.cuh" #include "machete_prepacked_layout.cuh" #include "machete_interleaving_utils.cuh" diff --git a/csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh b/csrc/quantization/machete/machete_mm_launcher.cuh similarity index 96% rename from csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh rename to csrc/quantization/machete/machete_mm_launcher.cuh index 5ecf03daf790..14206f1d5cc2 100644 --- a/csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh +++ b/csrc/quantization/machete/machete_mm_launcher.cuh @@ -1,9 +1,9 @@ #pragma once #include "machete_mm_kernel.cuh" -#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" +#include "cutlass_extensions/torch_utils.hpp" #include "core/scalar_type.hpp" -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh b/csrc/quantization/machete/machete_prepack_kernel.cuh similarity index 98% rename from csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh rename to csrc/quantization/machete/machete_prepack_kernel.cuh index 6a28dceccc8d..e1e054e5a007 100644 --- a/csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh +++ b/csrc/quantization/machete/machete_prepack_kernel.cuh @@ -2,7 +2,7 @@ #include "machete_mm_kernel.cuh" #include "cutlass_extensions/cute_utils.cuh" -#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" +#include "cutlass_extensions/torch_utils.hpp" #include namespace machete { diff --git a/csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh b/csrc/quantization/machete/machete_prepack_launcher.cuh similarity index 96% rename from csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh rename to csrc/quantization/machete/machete_prepack_launcher.cuh index 63af5e496b84..26bf200732d3 100644 --- a/csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh +++ b/csrc/quantization/machete/machete_prepack_launcher.cuh @@ -1,9 +1,9 @@ #pragma once #include "machete_prepack_kernel.cuh" -#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" +#include "cutlass_extensions/torch_utils.hpp" #include "core/scalar_type.hpp" -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/quantization/machete/machete_prepacked_layout.cuh b/csrc/quantization/machete/machete_prepacked_layout.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/machete/machete_prepacked_layout.cuh rename to csrc/quantization/machete/machete_prepacked_layout.cuh diff --git a/csrc/libtorch_stable/quantization/machete/machete_pytorch.cu b/csrc/quantization/machete/machete_pytorch.cu similarity index 100% rename from csrc/libtorch_stable/quantization/machete/machete_pytorch.cu rename to csrc/quantization/machete/machete_pytorch.cu diff --git a/csrc/libtorch_stable/quantization/marlin/.gitignore b/csrc/quantization/marlin/.gitignore similarity index 73% rename from csrc/libtorch_stable/quantization/marlin/.gitignore rename to csrc/quantization/marlin/.gitignore index 7dc482a89466..ba805f9250ec 100644 --- a/csrc/libtorch_stable/quantization/marlin/.gitignore +++ b/csrc/quantization/marlin/.gitignore @@ -1,3 +1,2 @@ sm*_kernel_*.cu kernel_selector.h -kernel_*.cu diff --git a/csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu b/csrc/quantization/marlin/awq_marlin_repack.cu similarity index 99% rename from csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu rename to csrc/quantization/marlin/awq_marlin_repack.cu index 55ce5b4e732d..09b647a03a39 100644 --- a/csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu +++ b/csrc/quantization/marlin/awq_marlin_repack.cu @@ -7,7 +7,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" namespace marlin { diff --git a/csrc/libtorch_stable/quantization/marlin/dequant.h b/csrc/quantization/marlin/dequant.h similarity index 100% rename from csrc/libtorch_stable/quantization/marlin/dequant.h rename to csrc/quantization/marlin/dequant.h diff --git a/csrc/libtorch_stable/quantization/marlin/generate_kernels.py b/csrc/quantization/marlin/generate_kernels.py similarity index 96% rename from csrc/libtorch_stable/quantization/marlin/generate_kernels.py rename to csrc/quantization/marlin/generate_kernels.py index 0b908d86d563..9e0af3c275ac 100644 --- a/csrc/libtorch_stable/quantization/marlin/generate_kernels.py +++ b/csrc/quantization/marlin/generate_kernels.py @@ -161,12 +161,15 @@ }, ] +# Kernels are written next to the kernel.h / marlin_template.h they include. +OUTPUT_DIR = sys.argv[2] if len(sys.argv) > 2 else os.path.dirname(__file__) + def remove_old_kernels(): - for filename in glob.glob(os.path.dirname(__file__) + "/*kernel_*.cu"): + for filename in glob.glob(OUTPUT_DIR + "/*kernel_*.cu"): subprocess.call(["rm", "-f", filename]) - filename = os.path.dirname(__file__) + "/kernel_selector.h" + filename = OUTPUT_DIR + "/kernel_selector.h" subprocess.call(["rm", "-f", filename]) @@ -297,7 +300,7 @@ def generate_new_kernels(): filename = filename.lower() - with open(os.path.join(os.path.dirname(__file__), filename), "w") as f: + with open(os.path.join(OUTPUT_DIR, filename), "w") as f: f.write(file_content) if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT: @@ -307,7 +310,7 @@ def generate_new_kernels(): '"marlin kernel with fp8 activation is not built.");' ) - with open(os.path.join(os.path.dirname(__file__), "kernel_selector.h"), "w") as f: + with open(os.path.join(OUTPUT_DIR, "kernel_selector.h"), "w") as f: f.write(kernel_selector_str) diff --git a/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu b/csrc/quantization/marlin/gptq_marlin_repack.cu similarity index 99% rename from csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu rename to csrc/quantization/marlin/gptq_marlin_repack.cu index 5dab9fe932bc..3fcf3406eb1b 100644 --- a/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu +++ b/csrc/quantization/marlin/gptq_marlin_repack.cu @@ -7,7 +7,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" namespace marlin { diff --git a/csrc/libtorch_stable/quantization/marlin/kernel.h b/csrc/quantization/marlin/kernel.h similarity index 100% rename from csrc/libtorch_stable/quantization/marlin/kernel.h rename to csrc/quantization/marlin/kernel.h diff --git a/csrc/libtorch_stable/quantization/marlin/marlin.cu b/csrc/quantization/marlin/marlin.cu similarity index 99% rename from csrc/libtorch_stable/quantization/marlin/marlin.cu rename to csrc/quantization/marlin/marlin.cu index cd52f3c73619..d77ad692b2bb 100644 --- a/csrc/libtorch_stable/quantization/marlin/marlin.cu +++ b/csrc/quantization/marlin/marlin.cu @@ -32,7 +32,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ static_assert(std::is_same::value || \ diff --git a/csrc/libtorch_stable/quantization/marlin/marlin.cuh b/csrc/quantization/marlin/marlin.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/marlin/marlin.cuh rename to csrc/quantization/marlin/marlin.cuh diff --git a/csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh b/csrc/quantization/marlin/marlin_dtypes.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh rename to csrc/quantization/marlin/marlin_dtypes.cuh diff --git a/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu b/csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu similarity index 99% rename from csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu rename to csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu index b48601435e3c..c4463269099d 100644 --- a/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu +++ b/csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu @@ -8,7 +8,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" // for only non-zp format (like gptq) __global__ void marlin_int4_fp8_preprocess_kernel_without_zp( diff --git a/csrc/libtorch_stable/quantization/marlin/marlin_mma.h b/csrc/quantization/marlin/marlin_mma.h similarity index 100% rename from csrc/libtorch_stable/quantization/marlin/marlin_mma.h rename to csrc/quantization/marlin/marlin_mma.h diff --git a/csrc/libtorch_stable/quantization/marlin/marlin_template.h b/csrc/quantization/marlin/marlin_template.h similarity index 100% rename from csrc/libtorch_stable/quantization/marlin/marlin_template.h rename to csrc/quantization/marlin/marlin_template.h diff --git a/csrc/libtorch_stable/quantization/vectorization.cuh b/csrc/quantization/vectorization.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/vectorization.cuh rename to csrc/quantization/vectorization.cuh diff --git a/csrc/libtorch_stable/quantization/vectorization_utils.cuh b/csrc/quantization/vectorization_utils.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/vectorization_utils.cuh rename to csrc/quantization/vectorization_utils.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh b/csrc/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh similarity index 96% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh rename to csrc/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh index a93a136bf425..0722101b3879 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh +++ b/csrc/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh @@ -5,7 +5,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include "cutlass/cutlass.h" @@ -19,8 +19,8 @@ #include "cutlass/gemm/collective/collective_builder.hpp" #include "cutlass/util/packed_stride.hpp" -#include "libtorch_stable/core/math.hpp" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "core/utils.hpp" +#include "cutlass_extensions/common.hpp" // clang-format on namespace vllm::c3x { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm.cuh similarity index 98% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm.cuh index cf18b63a8848..b20c055fb829 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm.cuh @@ -14,8 +14,8 @@ #include "cutlass/epilogue/collective/collective_builder.hpp" #include "cutlass/gemm/collective/collective_builder.hpp" -#include "libtorch_stable/core/math.hpp" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "core/utils.hpp" +#include "cutlass_extensions/common.hpp" // clang-format on /* diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu similarity index 90% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu index 5178eed0722e..bc088cf633f4 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm90_int8_dispatch.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu similarity index 89% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu index b7930012265f..f3df69850ec6 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_blockwise_sm100_fp8_dispatch.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh similarity index 98% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh index fa2add918213..1e7a31f5ca55 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh @@ -2,7 +2,7 @@ #include -#include "cuda_utils.h" +#include "core/utils.hpp" #include "cutlass/cutlass.h" #include "cutlass/numeric_types.h" @@ -242,10 +242,10 @@ void cutlass_gemm_blockwise_sm100_fp8_dispatch(torch::stable::Tensor& out, if (!swap_ab) { constexpr int TILE_N = 128; int tile_m = 256; - if (cuda_utils::ceil_div(n, TILE_N) * cuda_utils::ceil_div(m, 64) <= sms) { + if (div_ceil(n, TILE_N) * div_ceil(m, 64) <= sms) { tile_m = 64; } - else if (cuda_utils::ceil_div(n, TILE_N) * cuda_utils::ceil_div(m, 128) <= sms) { + else if (div_ceil(n, TILE_N) * div_ceil(m, 128) <= sms) { tile_m = 128; } if (tile_m == 64) { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu similarity index 96% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu index 021c16d66824..9f3c8a196df7 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_blockwise_sm120_fp8_dispatch.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh index 9c921ea90a32..34bb7ef3f23f 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh @@ -2,7 +2,6 @@ #include -#include "cuda_utils.h" #include "cutlass/cutlass.h" #include "cutlass/numeric_types.h" diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu similarity index 89% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu index a97909d37ae8..d3318c487675 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu @@ -1,7 +1,7 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_blockwise_sm90_fp8_dispatch.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp similarity index 82% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp index 913436186c3a..a15bc1da4a56 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp @@ -1,7 +1,7 @@ #include #include -#include "cuda_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "core/utils.hpp" +#include "cutlass_extensions/common.hpp" template void dispatch_scaled_mm(torch::stable::Tensor& c, @@ -41,14 +41,12 @@ void dispatch_scaled_mm(torch::stable::Tensor& c, STD_TORCH_CHECK(b_scales.dim() == 2, "b scale must be 2d tensor."); int32_t version_num = get_sm_version_num(); if (version_num >= 90) { - STD_TORCH_CHECK( - a.size(0) == a_scales.size(0) && - cuda_utils::ceil_div(a.size(1), int64_t(128)) == a_scales.size(1), - "a_scale_group_shape must be [1, 128]."); - STD_TORCH_CHECK( - cuda_utils::ceil_div(b.size(0), int64_t(128)) == b_scales.size(0) && - cuda_utils::ceil_div(b.size(1), int64_t(128)) == b_scales.size(1), - "b_scale_group_shape must be [128, 128]."); + STD_TORCH_CHECK(a.size(0) == a_scales.size(0) && + div_ceil(a.size(1), int64_t(128)) == a_scales.size(1), + "a_scale_group_shape must be [1, 128]."); + STD_TORCH_CHECK(div_ceil(b.size(0), int64_t(128)) == b_scales.size(0) && + div_ceil(b.size(1), int64_t(128)) == b_scales.size(1), + "b_scale_group_shape must be [128, 128]."); } STD_TORCH_CHECK(!bias, "Bias not yet supported blockwise scaled_mm"); diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_kernels.hpp b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_kernels.hpp similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_kernels.hpp rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_kernels.hpp diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu similarity index 97% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu index 84040a6a2218..59022b041e20 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm100_fp8_dispatch.cuh" -#include "core/batch_invariant.hpp" +#include "core/utils.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh index 42f9e4d5bc83..f790b3653d57 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh @@ -4,7 +4,7 @@ #include "scaled_mm.cuh" #include "cutlass_gemm_caller.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" /** * This file defines Gemm kernel configurations for SM100 (fp8) based on the diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu similarity index 91% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu index 1f709699eb9e..5cfe1d8d3ab9 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu @@ -1,7 +1,7 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm120_fp8_dispatch.cuh" -#include "core/batch_invariant.hpp" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "core/utils.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu similarity index 97% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu index e86c9bd48d3f..c94b9bd07207 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm90_fp8_dispatch.cuh" -#include "core/batch_invariant.hpp" +#include "core/utils.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh index 2fae3016c309..f78b8daea510 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh @@ -4,7 +4,7 @@ #include "scaled_mm.cuh" #include "cutlass_gemm_caller.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" /** * This file defines Gemm kernel configurations for SM90 (fp8) based on the Gemm diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu similarity index 91% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu index bf2bd030e3c3..717a2a588307 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm90_int8_dispatch.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/get_group_starts.cuh b/csrc/quantization/w8a8/cutlass/moe/get_group_starts.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/moe/get_group_starts.cuh rename to csrc/quantization/w8a8/cutlass/moe/get_group_starts.cuh index 6942c4b48ccb..2189e5d0abd8 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/get_group_starts.cuh +++ b/csrc/quantization/w8a8/cutlass/moe/get_group_starts.cuh @@ -3,7 +3,7 @@ #include #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include "cutlass/bfloat16.h" #include "cutlass/float8.h" diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh b/csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh similarity index 97% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh rename to csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh index b4cd520e9679..12241fb299b3 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh +++ b/csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh @@ -7,9 +7,9 @@ #include "cutlass/gemm/device/gemm_universal_adapter.h" #include -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" -#include "libtorch_stable/cutlass_extensions/common.hpp" -#include "libtorch_stable/torch_utils.h" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/common.hpp" +#include "torch_utils.h" #include "get_group_starts.cuh" using namespace cute; diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu b/csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu rename to csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu index 1f88bf140da5..f03f2c0bc47e 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu +++ b/csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu @@ -1,6 +1,6 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu b/csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu rename to csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu index d494bfb9d08c..14bb0586cdfc 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu +++ b/csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu @@ -1,6 +1,6 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu b/csrc/quantization/w8a8/cutlass/moe/moe_data.cu similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu rename to csrc/quantization/w8a8/cutlass/moe/moe_data.cu index 8ebc3dd9e878..b6043e803306 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu +++ b/csrc/quantization/w8a8/cutlass/moe/moe_data.cu @@ -1,11 +1,11 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include #include -#include "libtorch_stable/dispatch_utils.h" +#include "dispatch_utils.h" #include diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cu similarity index 98% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cu rename to csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cu index 184a26b491f8..bd60058d73fd 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cu +++ b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cu @@ -9,8 +9,8 @@ #include "scaled_mm_c2x_sm89_fp8_dispatch.cuh" #include "scaled_mm_c2x_sm89_int8_dispatch.cuh" -#include "core/batch_invariant.hpp" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp" +#include "core/utils.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp" using namespace vllm; diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cuh similarity index 97% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh rename to csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cuh index 9f9e12e982ed..30e587c63ad4 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh +++ b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cuh @@ -3,7 +3,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" // clang-format will break include orders // clang-format off @@ -22,8 +22,8 @@ #include "cutlass/epilogue/threadblock/fusion/visitors.hpp" #include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h" -#include "libtorch_stable/core/math.hpp" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "core/utils.hpp" +#include "cutlass_extensions/common.hpp" // clang-format on using namespace cute; diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm75_dispatch.cuh b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm75_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm75_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm75_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm80_dispatch.cuh b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm80_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm80_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm80_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_fp8_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_fp8_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_int8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_int8_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_int8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_int8_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu rename to csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu rename to csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu rename to csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu rename to csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu index 51f84d2ffd95..0c05f59d3fcd 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu +++ b/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu @@ -3,9 +3,9 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" void cutlass_scaled_mm_sm75(torch::stable::Tensor& c, torch::stable::Tensor const& a, diff --git a/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh b/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh index 32925c0c3c95..cb050c922302 100644 --- a/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh @@ -5,7 +5,7 @@ #include #include -#include "../../../../attention/attention_dtypes.h" +#include "core/vec_dtypes/vec_dtypes.h" #include namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/fp8/common.cu b/csrc/quantization/w8a8/fp8/common.cu similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/fp8/common.cu rename to csrc/quantization/w8a8/fp8/common.cu index abebd6f02f80..471499786977 100644 --- a/csrc/libtorch_stable/quantization/w8a8/fp8/common.cu +++ b/csrc/quantization/w8a8/fp8/common.cu @@ -1,5 +1,5 @@ -#include "../../../../quantization/w8a8/fp8/common.cuh" -#include "../../../../cuda_compat.h" +#include "common.cuh" +#include "core/cuda_compat.h" #include "../../../dispatch_utils.h" #include "../../../cub_helpers.h" #include "../../vectorization_utils.cuh" diff --git a/csrc/quantization/w8a8/fp8/common.cuh b/csrc/quantization/w8a8/fp8/common.cuh index 087f5099165a..143be698cc82 100644 --- a/csrc/quantization/w8a8/fp8/common.cuh +++ b/csrc/quantization/w8a8/fp8/common.cuh @@ -1,17 +1,17 @@ #pragma once -#include "libtorch_stable/quantization/vectorization.cuh" -#include "../../utils.cuh" +#include "quantization/vectorization.cuh" +#include "quantization/utils.cuh" #include -// This header is shared between _C and _C_stable_libtorch targets. +// This header is shared between _rocm_C and _C_stable_libtorch targets. // torch_utils.h provides get_device_prop(). We need to pass USE_CUDA // to the .so to expose some of the shims used by torch_utils.h. For now -// this is only done for _C_stable_libtorch and not for _C, so we use the -// non stable at::cuda::getCurrentDeviceProperties for _C for now. +// this is only done for _C_stable_libtorch and not for _rocm_C, so we use +// the non stable at::cuda::getCurrentDeviceProperties for _rocm_C for now. #ifdef TORCH_TARGET_VERSION - #include "../../../libtorch_stable/torch_utils.h" + #include "torch_utils.h" #else #ifdef USE_ROCM #include @@ -19,9 +19,9 @@ #endif #ifndef USE_ROCM - #include "nvidia/quant_utils.cuh" + #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" #else - #include "amd/quant_utils.cuh" + #include "quantization/w8a8/fp8/amd/quant_utils.cuh" #endif // Determines the preferred FP8 type for the current platform. diff --git a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh index 7fcf541fd64d..77aabe465111 100644 --- a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh @@ -1,6 +1,6 @@ #pragma once -#include "../../../../attention/attention_dtypes.h" +#include "core/vec_dtypes/vec_dtypes.h" #include #include #include @@ -243,7 +243,7 @@ __inline__ __device__ uint8_t scaled_vec_conversion( } // torch stable-ABI (headeronly) scalar types delegate to the CUDA-native -// conversions, so libtorch_stable kernels dispatched on c10::BFloat16 / +// conversions, so stable-ABI kernels dispatched on c10::BFloat16 / // c10::Half quantize correctly without manual casts. template <> __inline__ __device__ uint8_t scaled_vec_conversion( diff --git a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu similarity index 97% rename from csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu rename to csrc/quantization/w8a8/fp8/per_token_group_quant.cu index a0aee0cd9287..d2a888643899 100644 --- a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu +++ b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu @@ -3,8 +3,6 @@ #include #include -#include "libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h" - #include #ifdef USE_ROCM @@ -13,10 +11,10 @@ #include #endif -#include "libtorch_stable/quantization/vectorization.cuh" -#include "libtorch_stable/quantization/vectorization_utils.cuh" -#include "libtorch_stable/dispatch_utils.h" -#include "libtorch_stable/torch_utils.h" +#include "quantization/vectorization.cuh" +#include "quantization/vectorization_utils.cuh" +#include "dispatch_utils.h" +#include "torch_utils.h" __device__ __forceinline__ float GroupReduceMax(float val) { #ifdef USE_ROCM @@ -346,11 +344,13 @@ inline int GetGroupsPerBlockX(int64_t padded_groups_per_row) { return 4; } -void per_token_group_quant_8bit(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s, - int64_t group_size, double eps, double min_8bit, - double max_8bit, bool scale_ue8m0) { +// 8-bit per-token-group quantization helper used by both FP8 and INT8 +static void per_token_group_quant_8bit(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s, + int64_t group_size, double eps, + double min_8bit, double max_8bit, + bool scale_ue8m0 = false) { STD_TORCH_CHECK(input.is_contiguous()); STD_TORCH_CHECK(output_q.is_contiguous()); @@ -870,3 +870,12 @@ void per_token_group_quant_fp8(const torch::stable::Tensor& input, per_token_group_quant_8bit(input, output_q, output_s, group_size, eps, fp8_min, fp8_max, scale_ue8m0); } + +void per_token_group_quant_int8(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s, + int64_t group_size, double eps, double int8_min, + double int8_max) { + per_token_group_quant_8bit(input, output_q, output_s, group_size, eps, + int8_min, int8_max); +} diff --git a/csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu b/csrc/quantization/w8a8/int8/scaled_quant.cu similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu rename to csrc/quantization/w8a8/int8/scaled_quant.cu diff --git a/csrc/rocm/attention.cu b/csrc/rocm/attention.cu index 2a11653a58d7..caf6ee9a14c5 100644 --- a/csrc/rocm/attention.cu +++ b/csrc/rocm/attention.cu @@ -19,11 +19,11 @@ #include #include #include -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include -#include "../attention/dtype_fp8.cuh" -#include "../quantization/w8a8/fp8/amd/quant_utils.cuh" +#include "core/vec_dtypes/dtype_fp8.cuh" +#include "quantization/w8a8/fp8/amd/quant_utils.cuh" // ROCm 6.2 compatibility: map OCP fp8 types to FNUZ variants if OCP is absent #if !defined(HIP_FP8_TYPE_OCP) diff --git a/csrc/rocm/dispatch_utils.h b/csrc/rocm/dispatch_utils.h new file mode 100644 index 000000000000..12026965f098 --- /dev/null +++ b/csrc/rocm/dispatch_utils.h @@ -0,0 +1,24 @@ +/* + * Adapted from + * https://github.com/pytorch/pytorch/blob/v2.0.1/aten/src/ATen/Dispatch.h + */ +#pragma once + +#include + +// Need a special dispatch case macro since we will nest the FP8 dispatch. +// Instead of the usual 'scalar_t', this names the dispatched type 'fp8_t'. +#define AT_DISPATCH_FP8_CASE(enum_type, ...) \ + AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__) + +// ROCm devices might use either fn or fnuz, so set up dispatch table for both. +// A host-based check at runtime will create a preferred FP8 type for ROCm +// such that the correct kernel is dispatched. +#define VLLM_DISPATCH_CASE_FP8_TYPES(...) \ + AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ + AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) + +// When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'. +// See AT_DISPATCH_FP8_CASE above. +#define VLLM_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__)) diff --git a/csrc/rocm/q_gemm_rdna3.cu b/csrc/rocm/q_gemm_rdna3.cu index e6f1c2fb176b..8ba5b48f05f8 100644 --- a/csrc/rocm/q_gemm_rdna3.cu +++ b/csrc/rocm/q_gemm_rdna3.cu @@ -3,7 +3,8 @@ // // W4A16 GPTQ kernel for RDNA3 (gfx1100 / RX 7900 XTX class), templated on the // activation dtype (half or __hip_bfloat16). Adapted from exllamav2's 4-bit -// kernel (csrc/quantization/gptq/q_gemm.cu) with the following changes: +// kernel (csrc/quantization/gptq/q_gemm.cu) with the following +// changes: // // 1. Deterministic split-K epilogue. K is split across gridDim.z blocks; // each block stores its FP32 block partial to a scratch tensor, and the diff --git a/csrc/quickreduce/base.h b/csrc/rocm/quickreduce/base.h similarity index 100% rename from csrc/quickreduce/base.h rename to csrc/rocm/quickreduce/base.h diff --git a/csrc/custom_quickreduce.cu b/csrc/rocm/quickreduce/custom_quickreduce.cu similarity index 99% rename from csrc/custom_quickreduce.cu rename to csrc/rocm/quickreduce/custom_quickreduce.cu index d4e5d179a54e..c1e75ae05d6a 100644 --- a/csrc/custom_quickreduce.cu +++ b/csrc/rocm/quickreduce/custom_quickreduce.cu @@ -5,7 +5,7 @@ #ifdef USE_ROCM - #include "quickreduce/quick_reduce.h" + #include "rocm/quickreduce/quick_reduce.h" quickreduce::fptr_t init_custom_qr(int64_t rank, int64_t world_size, std::optional qr_max_size) { diff --git a/csrc/quickreduce/quick_reduce.h b/csrc/rocm/quickreduce/quick_reduce.h similarity index 100% rename from csrc/quickreduce/quick_reduce.h rename to csrc/rocm/quickreduce/quick_reduce.h diff --git a/csrc/quickreduce/quick_reduce_impl.cuh b/csrc/rocm/quickreduce/quick_reduce_impl.cuh similarity index 100% rename from csrc/quickreduce/quick_reduce_impl.cuh rename to csrc/rocm/quickreduce/quick_reduce_impl.cuh diff --git a/csrc/rocm/quickreduce/torch_bindings.cpp b/csrc/rocm/quickreduce/torch_bindings.cpp new file mode 100644 index 000000000000..b7d9a695ac48 --- /dev/null +++ b/csrc/rocm/quickreduce/torch_bindings.cpp @@ -0,0 +1,51 @@ +#include +#include "core/registration.h" +#include +#include + +#include +#include + +#ifdef USE_ROCM +// Defined in custom_quickreduce.cu. +using fptr_t = int64_t; +fptr_t init_custom_qr(int64_t rank, int64_t world_size, + std::optional qr_max_size = std::nullopt); +void qr_destroy(fptr_t _fa); +torch::Tensor qr_get_handle(fptr_t _fa); +void qr_open_handles(fptr_t _fa, const std::vector& handles); +void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, + int64_t quant_level, bool cast_bf2half = false); +int64_t qr_max_size(); +#endif + +// Note on op signatures: +// The X_meta signatures are for the meta functions corresponding to op X. +// They must be kept in sync with the signature for X. Generally, only +// functions that return Tensors require a meta function. +// +// See the following links for detailed docs on op registration and function +// schemas. +// https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ptttacy8y1u9 +// https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#annotations + +#ifdef USE_ROCM +TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { + // Quick Reduce all-reduce kernels (ROCm-only; stays on legacy _C). + custom_ar.def( + "qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool " + "cast_bf2half) -> ()"); + custom_ar.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce); + + custom_ar.def("init_custom_qr", &init_custom_qr); + custom_ar.def("qr_destroy", &qr_destroy); + custom_ar.def("qr_get_handle", &qr_get_handle); + + custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()"); + custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles); + + custom_ar.def("qr_max_size", &qr_max_size); +} +#endif + +REGISTER_EXTENSION(TORCH_EXTENSION_NAME) diff --git a/csrc/rocm/skinny_gemms.cu b/csrc/rocm/skinny_gemms.cu index 79af0814550d..bd7c0daa1db2 100644 --- a/csrc/rocm/skinny_gemms.cu +++ b/csrc/rocm/skinny_gemms.cu @@ -12,10 +12,10 @@ #include #include -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include "dispatch_utils.h" #include "quantization/w8a8/fp8/common.cuh" -#include "core/batch_invariant.hpp" +#include "core/utils.hpp" // Number of streams the pre-allocated split-K pool covers. ~7.5 MiB per slot. static constexpr int64_t kWvSlots = 8; diff --git a/csrc/rocm/skinny_gemms_int4.cu b/csrc/rocm/skinny_gemms_int4.cu index 4588b216fe01..be3cda7be6a3 100644 --- a/csrc/rocm/skinny_gemms_int4.cu +++ b/csrc/rocm/skinny_gemms_int4.cu @@ -10,8 +10,7 @@ #include #include // std::in_range -#include "../cuda_compat.h" -#include "dispatch_utils.h" +#include "core/cuda_compat.h" // Combined RDNA macro (gfx11 + gfx12) - both use 32-wide wavefronts #if defined(__GFX11__) || defined(__GFX12__) diff --git a/csrc/libtorch_stable/sampled_topk.cuh b/csrc/sampled_topk.cuh similarity index 100% rename from csrc/libtorch_stable/sampled_topk.cuh rename to csrc/sampled_topk.cuh diff --git a/csrc/libtorch_stable/sampler.cu b/csrc/sampler.cu similarity index 99% rename from csrc/libtorch_stable/sampler.cu rename to csrc/sampler.cu index 09151004c1ad..e9145e3b99e6 100644 --- a/csrc/libtorch_stable/sampler.cu +++ b/csrc/sampler.cu @@ -1,4 +1,4 @@ -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include "dispatch_utils.h" #include "torch_utils.h" diff --git a/csrc/libtorch_stable/topk.cu b/csrc/topk.cu similarity index 100% rename from csrc/libtorch_stable/topk.cu rename to csrc/topk.cu diff --git a/csrc/libtorch_stable/topk_histogram_4096.cuh b/csrc/topk_histogram_4096.cuh similarity index 100% rename from csrc/libtorch_stable/topk_histogram_4096.cuh rename to csrc/topk_histogram_4096.cuh diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index bcf0ce3e1c48..3d4f9ab27389 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -1,39 +1,1171 @@ -// Provides torch::Tensor for ops.h (previously included transitively via -// cache.h, which is no longer included here after cache ops moved to -// _C_stable_libtorch). -#include #include "ops.h" +#include "cuda_utils.h" #include "core/registration.h" -#include -#include - -// Note on op signatures: -// The X_meta signatures are for the meta functions corresponding to op X. -// They must be kept in sync with the signature for X. Generally, only -// functions that return Tensors require a meta function. -// -// See the following links for detailed docs on op registration and function -// schemas. -// https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ptttacy8y1u9 -// https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#annotations - -#ifdef USE_ROCM -TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { - // Quick Reduce all-reduce kernels (ROCm-only; stays on legacy _C). - custom_ar.def( - "qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool " - "cast_bf2half) -> ()"); - custom_ar.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce); - custom_ar.def("init_custom_qr", &init_custom_qr); - custom_ar.def("qr_destroy", &qr_destroy); - custom_ar.def("qr_get_handle", &qr_get_handle); +#include + +// Register ops with STABLE_TORCH_LIBRARY for libtorch stable ABI compatibility. +// Note: We register under namespace "_C" so ops are accessible as +// torch.ops._C. for compatibility with existing code. +STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { + // Compute per-token-group FP8 quantized tensor and scaling factor. + // The dummy arguments are here so we can correctly fuse with RMSNorm. + ops.def( + "per_token_group_fp8_quant(Tensor input, Tensor! output_q, Tensor! " + "output_s, " + "int group_size, float eps, float fp8_min, float fp8_max, bool " + "scale_ue8m0, bool dummy_is_scale_transposed, bool dummy_is_tma_aligned " + ") -> ()"); + // Compute per-token-group 8-bit quantized tensor and UE8M0-packed, + // TMA-aligned scales for DeepGEMM. + ops.def( + "per_token_group_fp8_quant_packed(Tensor input, Tensor! output_q, " + "Tensor! output_s_packed, int group_size, float eps, float fp8_min, " + "float fp8_max) -> ()"); + // Compute per-token-group INT8 quantized tensor and scaling factor. + ops.def( + "per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! " + "output_s, int group_size, float eps, float int8_min, float int8_max) -> " + "()"); + ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor"); + +#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 + + // Machete (Dense) Optimized Mixed Precision GEMM for Hopper. + ops.def( + "machete_supported_schedules(" + " ScalarType a_type," + " int b_type," + " ScalarType? maybe_group_scales_type," + " ScalarType? maybe_group_zeros_type," + " ScalarType? maybe_channel_scales_type," + " ScalarType? maybe_token_scales_type," + " ScalarType? maybe_out_type" + ") -> str[]"); + ops.def( + "machete_mm(" + " Tensor A," + " Tensor B," + " int b_type," + " ScalarType? out_type," + " Tensor? group_scales," + " Tensor? group_zeros," + " int? group_size," + " Tensor? channel_scales," + " Tensor? token_scales," + " str? schedule" + ") -> Tensor"); + ops.def( + "machete_prepack_B(" + " Tensor B," + " ScalarType a_type," + " int b_type," + " ScalarType? group_scales_type" + ") -> Tensor"); + // conditionally compiled so impl registration is in source file + + // Marlin GEMM + ops.def( + "marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, " + "Tensor? b_bias_or_none,Tensor b_scales, " + "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, " + "Tensor workspace, int b_type_id, " + "SymInt size_m, SymInt size_n, SymInt size_k, " + "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // gptq_marlin repack from GPTQ. + ops.def( + "gptq_marlin_repack(Tensor b_q_weight, " + "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // awq_marlin repack from AWQ. + ops.def( + "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, " + "SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // preprocess W-int4A-fp8 weight for marlin kernel + ops.def( + "marlin_int4_fp8_preprocess(Tensor qweight, " + "Tensor? qzeros_or_none, bool inplace) -> Tensor"); + // conditionally compiled so impl registrations are in source file +#endif + +#ifndef USE_ROCM + // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column + // quantization, as well as bias + ops.def( + "cutlass_scaled_mm(Tensor! out, Tensor a," + " Tensor b, Tensor a_scales," + " Tensor b_scales, Tensor? bias) -> ()"); + + // CUTLASS w8a8 GEMM, supporting asymmetric per-tensor or per-row/column + // quantization. + ops.def( + "cutlass_scaled_mm_azp(Tensor! out, Tensor a," + " Tensor b, Tensor a_scales," + " Tensor b_scales, Tensor azp_adj," + " Tensor? azp, Tensor? bias) -> ()"); + + // Check if cutlass scaled_mm is supported for CUDA devices of the given + // capability + ops.def("cutlass_scaled_mm_supports_fp8(int cuda_device_capability) -> bool"); + + // Check if cutlass grouped gemm is supported for CUDA devices of the given + // capability + ops.def("cutlass_group_gemm_supported(int cuda_device_capability) -> bool"); + + // CUTLASS w8a8 grouped GEMM + ops.def( + "cutlass_moe_mm(Tensor! out_tensors, Tensor a_tensors, Tensor b_tensors, " + " Tensor a_scales, Tensor b_scales, Tensor expert_offsets, " + " Tensor problem_sizes, Tensor a_strides, " + " Tensor b_strides, Tensor c_strides, bool per_act_token, " + " bool per_out_ch) -> ()"); + + // A function that computes data required to run fused MoE with w8a8 grouped + // GEMM. It takes topk_ids as an input, and computes expert_offsets + // (token start indices of each expert). In addition to this, it computes + // problem sizes for each expert's multiplication used by the two mms called + // from fused MoE operation, and arrays with permutations required to shuffle + // and de-shuffle the input/output of the fused operation. + ops.def( + "get_cutlass_moe_mm_data(Tensor topk_ids, Tensor! expert_offsets, " + " Tensor! problem_sizes1, Tensor! problem_sizes2, " + " Tensor! input_permutation, " + " Tensor! output_permutation, int num_experts, " + " int n, int k, Tensor? blockscale_offsets, " + " bool is_gated) -> ()"); + + // compute per-expert problem sizes from expert_first_token_offset + // produced by vLLM's moe_permute kernel + ops.def( + "get_cutlass_moe_mm_problem_sizes_from_expert_offsets(" + " Tensor expert_first_token_offset, " + " Tensor! problem_sizes1, " + " Tensor! problem_sizes2, " + " int n, int k, bool swap_ab) -> ()"); + + // A function that computes data required to run fused MoE with w8a8 grouped + // GEMM in batched expert format. It takes expert_num_tokens + // as an input, and computes expert_offsets (token start indices of each + // expert). In addition to this, it computes problem sizes for each expert's + // multiplication used by the two mms called from fused MoE operation. + ops.def( + "get_cutlass_batched_moe_mm_data(Tensor! expert_offsets, " + " Tensor! problem_sizes1, " + " Tensor! problem_sizes2, " + " Tensor expert_num_tokens, " + " int num_local_experts, int padded_m, " + " int n, int k) -> ()"); + + // Check if cutlass scaled_mm supports block quantization (used by DeepSeekV3) + ops.def( + "cutlass_scaled_mm_supports_block_fp8(int cuda_device_capability) -> " + "bool"); + + // CUTLASS nvfp4 block scaled GEMM + ops.def( + "cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b," + " Tensor block_scale_a, Tensor block_scale_b," + " Tensor alpha) -> ()"); + + // cutlass nvfp4 block scaled group GEMM + ops.def( + "cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b," + " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas," + " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); + + // cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE) + ops.def( + "cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b," + " Tensor a_blockscale, Tensor b_blockscales," + " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); + + // Compute NVFP4 block quantized tensor. + ops.def( + "scaled_fp4_quant(Tensor input," + " Tensor input_scale, bool " + "is_sf_swizzled_layout) -> (Tensor, Tensor)"); + + // Out variant + // TODO: Add out_variant tag once PyTorch supports it (added in 2.11) + // This registration is now migrated to stable ABI + // at::Tag::out_variant is not available in the stable ABI (enum_tag.h is not + // yet in torch/headeronly), the tag should be applied from Python + // via torch.library.Library.define(..., tags=(torch.Tag.out_variant,)) + // with the .impl remaining in C++. + // See pytorch/pytorch#176117. + ops.def( + "scaled_fp4_quant.out(Tensor input," + " Tensor input_scale, bool " + "is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) " + "-> ()"); + + // Compute NVFP4 experts quantization. + ops.def( + "scaled_fp4_experts_quant(Tensor! output, Tensor! output_scale," + "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts) -> ()"); + + // Fused SiLU+Mul+NVFP4 experts quantization. + ops.def( + "silu_and_mul_scaled_fp4_experts_quant(Tensor! output, Tensor! " + "output_scale," + "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts) -> ()"); + + // Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs). + ops.def( + "mxfp4_experts_quant(Tensor! output, Tensor! output_scale," + "Tensor input, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); + + // Fused SiLU+Mul+MXFP4 experts quantization. + ops.def( + "silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! " + "output_scale," + "Tensor input, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); + + // Fused SiLU+Mul+NVFP4 quantization. + ops.def( + "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, " + "Tensor input, Tensor input_global_scale) -> ()"); + + // Check if cutlass_scaled_mm_fp4 is supported for CUDA devices + // of the given capability + ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool"); + + // CUTLASS w4a8 GEMM + ops.def( + "cutlass_w4a8_mm(" + " Tensor A," + " Tensor B," + " Tensor group_scales," + " int group_size," + " Tensor channel_scales," + " Tensor token_scales," + " ScalarType? out_type," + " str? maybe_schedule" + ") -> Tensor"); + + // pack scales + ops.def("cutlass_pack_scale_fp8(Tensor scales) -> Tensor"); + + // encode and reorder weight matrix + ops.def("cutlass_encode_and_reorder_int4b(Tensor B) -> Tensor"); + + // CUTLASS w4a8 grouped GEMM + ops.def( + "cutlass_w4a8_moe_mm(" + " Tensor! out_tensors," + " Tensor a_tensors," + " Tensor b_tensors," + " Tensor a_scales," + " Tensor b_scales," + " Tensor b_group_scales," + " int b_group_size," + " Tensor expert_offsets," + " Tensor problem_sizes," + " Tensor a_strides," + " Tensor b_strides," + " Tensor c_strides," + " Tensor group_scale_strides," + " str? maybe_schedule" + ") -> ()"); + + ops.def( + "cutlass_encode_and_reorder_int4b_grouped(Tensor b_tensors) -> (Tensor, " + "Tensor)"); + + // SM100 CUTLASS MLA decode + // conditionally compiled so impl registrations are in source file + ops.def( + "sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope," + " Tensor q_pe, Tensor kv_c_and_k_pe_cache," + " Tensor seq_lens, Tensor page_table," + " Tensor workspace, float scale," + " int num_kv_splits) -> ()"); + + ops.def( + "sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches," + " int sm_count, int num_kv_splits) " + "-> int"); + // Quantized GEMM for AWQ. + ops.def( + "awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, " + "Tensor _zeros, SymInt split_k_iters) -> Tensor"); + + // Dequantization for AWQ. + ops.def( + "awq_dequantize(Tensor _kernel, Tensor _scaling_factors, " + "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor"); + + // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens). + // conditionally compiled so impl registration is in source file + ops.def( + "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b, " + "bool enable_pdl=False) -> ()"); + + // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+). + // conditionally compiled so impl registration is in source file + ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); +#endif + + // Merge attn states + // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005 + // can be used to combine partial attention results (in the split-KV case) + ops.def( + "merge_attn_states(" + " Tensor! output," + " Tensor!? output_lse," + " Tensor prefix_output," + " Tensor prefix_lse," + " Tensor suffix_output," + " Tensor suffix_lse," + " int!? prefill_tokens_with_context," + " Tensor? output_scale=None) -> ()"); + + // Hadamard transforms + // conditionally compiled so impl registration is in source file + ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor"); + + // Apply Root Mean Square (RMS) Normalization to the input tensor. + ops.def( + "rms_norm(Tensor! result, Tensor input, Tensor? weight, float epsilon) " + "-> " + "()"); + + // In-place fused Add and RMS Normalization. + ops.def( + "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, " + "float epsilon) -> ()"); + + // Layernorm-quant + // Apply Root Mean Square (RMS) Normalization to the input tensor. + ops.def( + "rms_norm_static_fp8_quant(Tensor! result, Tensor input, Tensor weight, " + "Tensor scale, float epsilon) -> " + "()"); + + // In-place fused Add and RMS Normalization. + ops.def( + "fused_add_rms_norm_static_fp8_quant(Tensor! result, Tensor input, " + "Tensor! residual, Tensor weight, " + "Tensor scale, float epsilon) -> ()"); + + // Fused Layernorm + Quant kernels + ops.def( + "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, " + "Tensor weight, Tensor! scale, float epsilon, " + "Tensor? scale_ub, Tensor!? residual) -> ()"); + + // Fused Layernorm + Block quant kernels + ops.def( + "rms_norm_per_block_quant(Tensor! result, Tensor input, " + "Tensor weight, Tensor! scale, float epsilon, " + "Tensor? scale_ub, Tensor!? residual, int group_size, " + "bool is_scale_transposed) -> ()"); + + // Fused SiLU+Mul + per-block quantization + ops.def( + "silu_and_mul_per_block_quant(" + "Tensor! out, " + "Tensor input, " + "Tensor! scales, " + "int group_size, " + "Tensor? scale_ub=None, " + "bool is_scale_transposed=False) -> ()"); + + // Rotary embedding + // Apply GPT-NeoX or GPT-J style rotary embedding to query and key. + ops.def( + "rotary_embedding(Tensor positions, Tensor! query," + " Tensor!? key, int head_size," + " Tensor cos_sin_cache, bool is_neox, int " + "rope_dim_offset=0, bool inverse=False) -> ()"); + + // Function for fused QK Norm and RoPE + ops.def( + "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, " + "int num_heads_k, int num_heads_v, int head_dim, float eps, " + "Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, " + "bool is_neox, Tensor position_ids, " + "int forced_token_heads_per_warp=-1) -> ()"); + + // q_head_padded is the padded Q head count of the returned tensor, or 0 to + // do the KV insert alone and return an empty tensor. The Q knobs are + // independent: apply_q_norm and apply_q_rope each drop that step for Q + // alone (KV is always rotated), and is_q_interleaved reads and writes Q in + // FlashMLA's mega-attention chunk-interleaved layout, which moves the + // padding heads to the tail of every head-dim chunk. + ops.def( + "fused_deepseek_v4_kv_rope_insert(" + "Tensor kv, Tensor! k_cache, Tensor slot_mapping, Tensor position_ids, " + "Tensor cos_sin_cache, int cache_block_size, Tensor? fp8_scale=None, " + "bool kv_mxfp8=False) -> " + "()"); + ops.def( + "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(" + "Tensor q_in, Tensor kv, Tensor! k_cache, " + "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " + "int q_head_padded, float eps, int cache_block_size, " + "bool apply_q_norm=True, bool kv_mxfp8=False, bool apply_q_rope=True, " + "bool is_q_interleaved=False) -> Tensor"); + + // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate + // FP8 tensor, and KV into a contiguous 512-wide token-strided cache. + ops.def( + "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(" + "Tensor! q, Tensor kv, Tensor! k_cache, Tensor slot_mapping, " + "Tensor position_ids, Tensor cos_sin_cache, float eps, " + "int cache_block_size, bool apply_q_norm=True) -> ()"); + ops.def( + "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(" + "Tensor q, Tensor kv, Tensor! q_fp8, Tensor! k_cache, " + "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " + "Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, " + "int cache_block_size, bool apply_q_norm=True) -> ()"); + + // Kimi-K3 MLA epilogues: optional RoPE followed by concat/cache insertion. + ops.def( + "fused_kimi_k3_mla_key_concat_kv_cache_insert(" + "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, " + "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_key_concat_ds_mla_insert(" + "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, " + "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_kv_concat(Tensor k_nope, Tensor k_pe, Tensor! k_out) " + "-> ()"); + ops.def( + "fused_kimi_k3_mla_kv_concat_quant_fp8(" + "Tensor k_nope, Tensor k_pe, Tensor v, Tensor! k_fp8, Tensor! v_fp8) " + "-> ()"); + ops.def( + "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(" + "Tensor q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, Tensor v, " + "Tensor! q_fp8, Tensor! k_fp8, Tensor! v_fp8, Tensor! k_cache, " + "Tensor slot_mapping, Tensor q_scale_inv, Tensor k_scale_inv, " + "Tensor v_scale_inv, Tensor cache_scale_inv, int cache_block_size, " + "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()"); + + // Kimi-K3 MLA decode epilogue: concat mqa_q = [ql_nope | q_pe] and insert the + // latent [kv_c_normed | k_pe] into the paged cache (bf16 / fp8 / fp8_ds_mla). + ops.def( + "fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(" + "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " + "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(" + "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " + "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " + "Tensor q_scale_inv, Tensor cache_scale_inv, int cache_block_size, " + "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_decode_q_concat_ds_mla_insert(" + "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " + "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + +#ifndef USE_ROCM + ops.def( + "minimax_allreduce_rms_qk(" + "Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, " + "Tensor workspace, int q_size, int kv_size, int rank, int nranks, " + "float eps) -> (Tensor, Tensor)"); +#endif + + // Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE + KV-insert. + ops.def( + "fused_minimax_m3_qknorm_rope_kv_insert(" + "Tensor! qkv, Tensor q_norm_weight, Tensor k_norm_weight, " + "Tensor cos_sin_cache, Tensor positions, int num_heads, " + "int num_kv_heads, int rotary_dim, float eps, " + "Tensor? index_q_norm_weight, Tensor? index_k_norm_weight, " + "int num_index_heads, " + "Tensor? slot_mapping, Tensor? index_slot_mapping, " + "Tensor!? kv_cache, Tensor!? index_cache, " + "int block_size, Tensor!? q_out, Tensor!? index_q_out, " + "str kv_cache_dtype, bool skip_index_branch=False, " + "Tensor!? q_fp8_out=None, float q_fp8_scale=1.0, " + "Tensor? kv_k_scale=None, Tensor? kv_v_scale=None) -> ()"); + +#ifdef VLLM_ENABLE_FUSED_KDA_DECODE + ops.def( + "fused_kda_decode(" + "Tensor x, Tensor weight, Tensor? bias, Tensor! conv_state, " + "Tensor raw_g, Tensor raw_beta, Tensor A_log, Tensor dt_bias, " + "Tensor state_indices, Tensor! state, Tensor! out, " + "float? lower_bound=None, Tensor? output_gate=None, " + "Tensor? norm_weight=None, float norm_eps=1e-5) -> ()"); +#endif + +#ifdef VLLM_ENABLE_FUSED_GDN_DECODE + ops.def( + "fused_gdn_decode_post_conv_mtp(" + "Tensor mixed_qkv, Tensor a, Tensor b, Tensor A_log, Tensor dt_bias, " + "Tensor state_indices, Tensor cu_seqlens, Tensor num_accepted_tokens, " + "Tensor! state, Tensor output_gate, Tensor norm_weight, Tensor! out, " + "float scale, float norm_eps=1e-5, " + "str output_gate_activation='silu') -> ()"); +#endif + +#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK + ops.def( + "fused_kda_prologue(" + "Tensor q, Tensor k, Tensor v, Tensor raw_g, Tensor raw_beta, " + "Tensor A_log, Tensor dt_bias, Tensor! qg, Tensor! w, Tensor! u, " + "Tensor! kg_t, Tensor! aqk, Tensor! decay, Tensor cu_seqlens, " + "Tensor chunk_indices, Tensor? conv_weight, Tensor(e!)? conv_state, " + "Tensor? conv_state_indices, Tensor? conv_has_initial_state, " + "float scale, float lower_bound) -> ()"); + ops.def( + "fused_kda_chunk(" + "Tensor qg, Tensor w, Tensor u, Tensor kg_t, Tensor aqk, Tensor decay, " + "Tensor? initial_state, Tensor(a!)? final_state, Tensor! out, " + "Tensor cu_seqlens, Tensor chunk_offsets, float scale, " + "Tensor(b!)? group_state, int groups, " + "Tensor(c!)? checkpoint_state=None, Tensor? checkpoint_offsets=None, " + "Tensor? checkpoint_state_indices=None, Tensor(d!)? state_cache=None, " + "Tensor? state_indices=None, Tensor? has_initial_state=None) -> ()"); +#endif + +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES + ops.def( + "kimi_k3_attn_res(" + "Tensor! prefix, Tensor? delta, Tensor! blocks, Tensor norm_weight, " + "Tensor qk_weight, Tensor? output_norm_weight, Tensor! output, " + "int num_blocks, " + "int block_write_idx, float eps, " + "float output_norm_eps) -> ()"); +#endif + + // Apply repetition penalties to logits in-place. + ops.def( + "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, " + "Tensor output_mask, Tensor repetition_penalties) -> ()"); + + // Optimized top-k per row operations. + ops.def( + "top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, " + "Tensor! indices, int numRows, int stride0, " + "int stride1, int topK) -> ()"); + + ops.def( + "top_k_per_row_decode(Tensor logits, int next_n, " + "Tensor seq_lens, Tensor! indices, " + "int numRows, int stride0, int stride1, int topK) -> ()"); + + ops.def( + "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, " + "Tensor workspace, int k, int max_seq_len) -> ()"); + +#ifdef VLLM_ENABLE_COOPERATIVE_TOPK + ops.def( + "cooperative_topk(Tensor logits, Tensor lengths, Tensor! output, " + "Tensor workspace, int k, int max_seq_len) -> ()"); +#endif + + // 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) -> ()"); + + ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()"); + + // SwiGLU activation with input clamping. + // alpha scales the sigmoid (gate * sigmoid(alpha * gate)); beta is added to + // the up half (up + beta). Defaults alpha=1.0, beta=0.0 give silu(gate)*up. + ops.def( + "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, " + "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) -> ()"); + + // Activation function used in GeGLU with `tanh` approximation. + ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()"); + + // FATReLU implementation. + ops.def("fatrelu_and_mul(Tensor! out, Tensor input, float threshold) -> ()"); + + ops.def( + "swigluoai_and_mul(Tensor! out, Tensor input, float alpha=1.702, float " + "limit=7.0) " + "-> ()"); + + // SituGLU implementation used in Kimi models. + ops.def( + "situ_and_mul(Tensor! out, Tensor input, float beta=1.0, float " + "linear_beta=-1.0) -> ()"); + // Fused SituGLU activation + dynamic FP8 quantization for the Humming w2 path + // (writes the fp8 down input and its float32 scale). group_size=0 -> + // per-token scale [.., 1]; group_size=128 -> k-major block-FP8 scale [.., + // d/128]. + ops.def( + "situ_and_mul_quant(Tensor! out, Tensor! scale, Tensor input, " + "float beta=1.0, float linear_beta=-1.0, int group_size=0, " + "Tensor? num_valid_tokens=None, int topk=1) -> ()"); + ops.def( + "masked_situ_and_mul(Tensor! out, Tensor input, Tensor " + "expert_num_tokens, float beta=1.0, float linear_beta=-1.0) -> ()"); + ops.def( + "masked_moe_activation(Tensor! out, Tensor input, Tensor " + "valid_token_counts, str activation, float clamp_limit=0.0, float " + "alpha=1.0, float beta=0.0, float situ_beta=1.0, float " + "situ_linear_beta=-1.0) -> ()"); + + // GELU implementation used in GPT-2. + ops.def("gelu_new(Tensor! out, Tensor input) -> ()"); - custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()"); - custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles); + // Approximate GELU implementation. + ops.def("gelu_fast(Tensor! out, Tensor input) -> ()"); - custom_ar.def("qr_max_size", &qr_max_size); + // Quick GELU implementation. + ops.def("gelu_quick(Tensor! out, Tensor input) -> ()"); + + // relu(x)^2 activation from https://arxiv.org/abs/2109.08668v2 + ops.def("relu_squared(Tensor! out, Tensor input) -> ()"); + + // Compute int8 quantized tensor for given scaling factor. + ops.def( + "static_scaled_int8_quant(Tensor! result, Tensor input, Tensor scale," + "Tensor? azp) -> ()"); + + // Compute int8 quantized tensor and scaling factor + ops.def( + "dynamic_scaled_int8_quant(Tensor! result, Tensor input, Tensor! scale, " + "Tensor!? azp) -> ()"); + + // Compute FP8 quantized tensor for given scaling factor. + // Supports per-tensor, per-channel, per-token, and arbitrary 2D group + // scaling. Optional group_m/group_n specify the group shape explicitly; + // required for 1D scales to disambiguate per-channel vs per-token. + ops.def( + "static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale, " + "int[]? group_shape=None) -> ()"); + + // Compute dynamic-per-tensor FP8 quantized tensor and scaling factor. + ops.def( + "dynamic_scaled_fp8_quant(Tensor! result, Tensor input, Tensor! scale) " + "-> " + "()"); + + // Compute dynamic-per-token FP8 quantized tensor and scaling factor. + ops.def( + "dynamic_per_token_scaled_fp8_quant(Tensor! result, Tensor input, " + "Tensor! scale, Tensor? scale_ub) -> " + "()"); + + // Quantized GEMM for GPTQ. + // Note: even though the C++ inferred schema is correct for this op, it seems + // to prevent the meta function registry. + ops.def( + "gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, " + "Tensor b_gptq_scales, bool use_exllama, bool " + "use_v2_format, int bit) " + "-> Tensor"); + + // Post processing for GPTQ. + ops.def("gptq_shuffle(Tensor! q_weight, int bit) -> ()"); + + // Mamba selective scan kernel + ops.def( + "selective_scan_fwd(Tensor! u, Tensor! delta," + "Tensor! A, Tensor! B, Tensor! C," + "Tensor? D_, Tensor!? z_, Tensor? delta_bias_," + "bool delta_softplus," + "Tensor? query_start_loc," + "Tensor? cache_indices," + "Tensor? has_initial_state," + "Tensor! ssm_states," + "int null_block_id," + "int block_size," + "Tensor? block_idx_first_scheduled_token," + "Tensor? block_idx_last_scheduled_token," + "Tensor? initial_state_idx," + "Tensor? cu_chunk_seqlen," + "Tensor? last_chunk_indices) -> ()"); + + // LongCat n-gram embedding index kernel. All tensor args are marked mutable + // to match the (non-const) stable-Tensor& C++ signature; only ne_token_table + // and n_gram_ids are actually written in place. + // Fused vocab-parallel embedding: gather the rows this rank owns and write + // zeros for the rest, so the following all-reduce reconstructs the full + // embedding. + ops.def( + "vocab_parallel_embedding(Tensor! out, Tensor input_ids, Tensor weight, " + "int org_vocab_start_index, int org_vocab_end_index, " + "int num_org_vocab_padding, int added_vocab_start_index, " + "int added_vocab_end_index) -> ()"); + ops.def( + "ngram_compute_n_gram_ids(int ne_n, int ne_k, Tensor(a!) ne_weights, " + "Tensor(b!) ne_mods, Tensor(c!) exclusive_ne_embedder_size_sums, " + "Tensor(d!) exclusive_req_len_sums, Tensor(e!) ne_token_table, " + "Tensor(f!) row_indices, Tensor(g!) column_starts, " + "Tensor(h!) n_gram_ids) -> ()"); } + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { + // LongCat n-gram embedding index kernel. + ops.impl("ngram_compute_n_gram_ids", TORCH_BOX(&ngram_compute_n_gram_ids)); + ops.impl("vocab_parallel_embedding", TORCH_BOX(&vocab_parallel_embedding)); + + // Per-token group quantization + ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8)); + ops.impl("per_token_group_fp8_quant_packed", + TORCH_BOX(&per_token_group_quant_8bit_packed)); + ops.impl("per_token_group_quant_int8", + TORCH_BOX(&per_token_group_quant_int8)); + +#ifndef USE_ROCM + // CUTLASS scaled_mm ops + ops.impl("cutlass_scaled_mm", TORCH_BOX(&cutlass_scaled_mm)); + ops.impl("cutlass_scaled_mm_azp", TORCH_BOX(&cutlass_scaled_mm_azp)); + ops.impl("cutlass_moe_mm", TORCH_BOX(&cutlass_moe_mm)); + ops.impl("get_cutlass_moe_mm_data", TORCH_BOX(&get_cutlass_moe_mm_data)); + ops.impl("get_cutlass_moe_mm_problem_sizes_from_expert_offsets", + TORCH_BOX(&get_cutlass_moe_mm_problem_sizes_from_expert_offsets)); + ops.impl("get_cutlass_batched_moe_mm_data", + TORCH_BOX(&get_cutlass_batched_moe_mm_data)); + + // FP4/NVFP4 ops + ops.impl("cutlass_scaled_fp4_mm", TORCH_BOX(&cutlass_scaled_fp4_mm)); + ops.impl("scaled_fp4_quant", TORCH_BOX(&scaled_fp4_quant_func)); + ops.impl("scaled_fp4_quant.out", TORCH_BOX(&scaled_fp4_quant_out)); + ops.impl("scaled_fp4_experts_quant", TORCH_BOX(&scaled_fp4_experts_quant)); + ops.impl("silu_and_mul_scaled_fp4_experts_quant", + TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant)); + ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant)); + // mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only). + // W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu. + + // AWQ ops + ops.impl("awq_gemm", TORCH_BOX(&awq_gemm)); + ops.impl("awq_dequantize", TORCH_BOX(&awq_dequantize)); + + // DSV3 fused A GEMM: conditionally compiled so impl registration is in + // source file (dsv3_fused_a_gemm.cu) +#endif + + ops.impl("merge_attn_states", TORCH_BOX(&merge_attn_states)); + + // Layernorm kernels (shared CUDA/ROCm) + ops.impl("rms_norm", TORCH_BOX(&rms_norm)); + ops.impl("fused_add_rms_norm", TORCH_BOX(&fused_add_rms_norm)); + + // Layernorm-quant kernels (shared CUDA/ROCm) + ops.impl("rms_norm_static_fp8_quant", TORCH_BOX(&rms_norm_static_fp8_quant)); + ops.impl("fused_add_rms_norm_static_fp8_quant", + TORCH_BOX(&fused_add_rms_norm_static_fp8_quant)); + + // Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm) + ops.impl("rms_norm_dynamic_per_token_quant", + TORCH_BOX(&rms_norm_dynamic_per_token_quant)); + ops.impl("rms_norm_per_block_quant", TORCH_BOX(&rms_norm_per_block_quant)); + ops.impl("silu_and_mul_per_block_quant", + TORCH_BOX(&silu_and_mul_per_block_quant)); + + // Positional encoding kernels (shared CUDA/ROCm) + ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding)); + ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope)); + ops.impl("fused_deepseek_v4_kv_rope_insert", + TORCH_BOX(&fused_deepseek_v4_kv_rope_insert)); + ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", + TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert)); + ops.impl( + "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert", + TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert)); + ops.impl( + "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", + TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert)); + ops.impl("fused_kimi_k3_mla_key_concat_kv_cache_insert", + TORCH_BOX(&fused_kimi_k3_mla_key_concat_kv_cache_insert)); + ops.impl("fused_kimi_k3_mla_key_concat_ds_mla_insert", + TORCH_BOX(&fused_kimi_k3_mla_key_concat_ds_mla_insert)); + ops.impl("fused_kimi_k3_mla_kv_concat", + TORCH_BOX(&fused_kimi_k3_mla_kv_concat)); + ops.impl("fused_kimi_k3_mla_kv_concat_quant_fp8", + TORCH_BOX(&fused_kimi_k3_mla_kv_concat_quant_fp8)); + ops.impl("fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert", + TORCH_BOX(&fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert)); + ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_insert", + TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_insert)); + ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert", + TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert)); + ops.impl("fused_kimi_k3_mla_decode_q_concat_ds_mla_insert", + TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_ds_mla_insert)); +#ifndef USE_ROCM + ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk)); +#endif + ops.impl("fused_minimax_m3_qknorm_rope_kv_insert", + TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert)); +#ifdef VLLM_ENABLE_FUSED_KDA_DECODE + ops.impl("fused_kda_decode", TORCH_BOX(&fused_kda_decode)); +#endif + +#ifdef VLLM_ENABLE_FUSED_GDN_DECODE + ops.impl("fused_gdn_decode_post_conv_mtp", + TORCH_BOX(&fused_gdn_decode_post_conv_mtp)); +#endif + +#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK + ops.impl("fused_kda_prologue", TORCH_BOX(&fused_kda_prologue)); + ops.impl("fused_kda_chunk", TORCH_BOX(&fused_kda_chunk)); #endif -REGISTER_EXTENSION(TORCH_EXTENSION_NAME) +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES + ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res)); +#endif + + // Sampler kernels (shared CUDA/ROCm) + ops.impl("apply_repetition_penalties_", + TORCH_BOX(&apply_repetition_penalties_)); + ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill)); + ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode)); + ops.impl("persistent_topk", TORCH_BOX(&persistent_topk)); +#ifdef VLLM_ENABLE_COOPERATIVE_TOPK + ops.impl("cooperative_topk", TORCH_BOX(&cooperative_topk)); +#endif + + // 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)); + ops.impl("gelu_tanh_and_mul", TORCH_BOX(&gelu_tanh_and_mul)); + ops.impl("fatrelu_and_mul", TORCH_BOX(&fatrelu_and_mul)); + ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul)); + ops.impl("situ_and_mul", TORCH_BOX(&situ_and_mul)); + ops.impl("situ_and_mul_quant", TORCH_BOX(&situ_and_mul_quant)); + ops.impl("masked_situ_and_mul", TORCH_BOX(&masked_situ_and_mul)); + ops.impl("masked_moe_activation", TORCH_BOX(&masked_moe_activation)); + ops.impl("gelu_new", TORCH_BOX(&gelu_new)); + ops.impl("gelu_fast", TORCH_BOX(&gelu_fast)); + ops.impl("gelu_quick", TORCH_BOX(&gelu_quick)); + ops.impl("relu_squared", TORCH_BOX(&relu_squared)); + ops.impl("silu_and_mul_with_clamp", TORCH_BOX(&silu_and_mul_clamp)); + + // INT8 quantization kernels + ops.impl("static_scaled_int8_quant", TORCH_BOX(&static_scaled_int8_quant)); + ops.impl("dynamic_scaled_int8_quant", TORCH_BOX(&dynamic_scaled_int8_quant)); + + // FP8 quantization kernels + ops.impl("static_scaled_fp8_quant", TORCH_BOX(&static_scaled_fp8_quant)); + ops.impl("dynamic_scaled_fp8_quant", TORCH_BOX(&dynamic_scaled_fp8_quant)); + ops.impl("dynamic_per_token_scaled_fp8_quant", + TORCH_BOX(&dynamic_per_token_scaled_fp8_quant)); + + // GPTQ kernels + ops.impl("gptq_gemm", TORCH_BOX(&gptq_gemm)); + ops.impl("gptq_shuffle", TORCH_BOX(&gptq_shuffle)); + + // Mamba kernels + ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) { + ops.impl("get_cuda_view_from_cpu_tensor", + TORCH_BOX(&get_cuda_view_from_cpu_tensor)); +} + +STABLE_TORCH_LIBRARY_FRAGMENT(_C_cuda_utils, cuda_utils) { + cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int"); + cuda_utils.def( + "get_max_shared_memory_per_block_device_attribute(int device_id) -> int"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd, + cuda_utils) { + cuda_utils.impl("get_device_attribute", TORCH_BOX(&get_device_attribute)); + cuda_utils.impl("get_max_shared_memory_per_block_device_attribute", + TORCH_BOX(&get_max_shared_memory_per_block_device_attribute)); +} + +// These capability-check functions take only primitive args (no tensors), so +// there is no device to dispatch on. CompositeExplicitAutograd makes them +// available for all backends. This is the stable ABI equivalent of calling +// ops.impl("op_name", &func) without a dispatch key in the non-stable API. +STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) { +#ifndef USE_ROCM + ops.impl("cutlass_scaled_mm_supports_fp8", + TORCH_BOX(&cutlass_scaled_mm_supports_fp8)); + ops.impl("cutlass_group_gemm_supported", + TORCH_BOX(&cutlass_group_gemm_supported)); + ops.impl("cutlass_scaled_mm_supports_block_fp8", + TORCH_BOX(&cutlass_scaled_mm_supports_block_fp8)); + ops.impl("cutlass_scaled_mm_supports_fp4", + TORCH_BOX(&cutlass_scaled_mm_supports_fp4)); +#endif +} + +// Cache ops +STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { + // Swap in (out) the cache blocks from src to dst. + ops.def( + "swap_blocks(Tensor src, Tensor! dst," + " int block_size_in_bytes, Tensor block_mapping) -> ()"); + + // Batch swap: submit all block copies in a single driver call. + ops.def( + "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs," + " Tensor sizes," + " bool is_src_access_order_any=False) -> ()"); + + // Reshape the key and value tensors and cache them. + ops.def( + "reshape_and_cache(Tensor key, Tensor value," + " Tensor! key_cache, Tensor! value_cache," + " Tensor slot_mapping," + " str kv_cache_dtype," + " Tensor k_scale, Tensor v_scale) -> ()"); + + // Reshape the key and value tensors and cache them. + ops.def( + "reshape_and_cache_flash(Tensor key, Tensor value," + " Tensor! key_cache," + " Tensor! value_cache," + " Tensor slot_mapping," + " str kv_cache_dtype," + " Tensor k_scale, Tensor v_scale) -> ()"); + + // Concat kv_c and k_pe and cache them. + ops.def( + "concat_and_cache_mla(Tensor kv_c, Tensor k_pe," + " Tensor! kv_cache," + " Tensor slot_mapping," + " str kv_cache_dtype," + " Tensor scale) -> ()"); + + // Grouped concat_and_cache_mla across all layers. Each layer's cache base + // pointer and optional plain-FP8 scale are read from device tensors. + ops.def( + "concat_and_cache_mla_grouped(Tensor kv_c, Tensor k_pe," + " Tensor kv_cache_ptrs," + " Tensor slot_mapping," + " int block_size, int block_stride," + " int entry_stride," + " Tensor? kv_scales=None," + " str kv_cache_dtype='auto') -> ()"); + +#ifndef USE_ROCM + ops.def( + "hisparse_resolve_residency(Tensor host_cache," + " Tensor! hot_cache," + " Tensor hot_block_table," + " Tensor global_indices," + " Tensor! hot_indices," + " Tensor! device_global_indices," + " Tensor! lru_slots," + " Tensor? request_state_indices," + " int region_stride," + " int max_union_rows," + " Tensor(a!)? miss_mask=None," + " Tensor(b!)? stats=None," + " Tensor(c!)? attention_indices=None," + " int attention_block_stride=0," + " Tensor? request_ids=None," + " Tensor? source_block_table=None," + " int source_block_size=0," + " Tensor(d!)? resolved_global_indices=None," + " Tensor(e!)? valid_counts=None," + " Tensor(f!)? swap_host_physical_rows=None," + " Tensor(g!)? swap_device_physical_rows=None," + " Tensor(h!)? swap_counts=None," + " Tensor? resident_block_table=None," + " int resident_block_size=0," + " int resident_null_block=0) -> ()"); + + ops.def( + "hisparse_invalidate_written_slots(Tensor! device_global_indices," + " Tensor request_state_indices," + " Tensor req_id_per_token," + " Tensor written_slots) -> ()"); + + ops.def( + "hisparse_gather_plan(Tensor host_cache," + " Tensor! hot_cache," + " Tensor global_indices," + " Tensor hot_indices," + " Tensor miss_mask," + " Tensor? request_state_indices," + " Tensor(a!)? attention_indices=None," + " int attention_block_stride=0) -> ()"); + + ops.def( + "hisparse_gather_compact(Tensor host_cache," + " Tensor! hot_cache," + " Tensor miss_global_indices," + " Tensor miss_hot_indices," + " Tensor miss_counts) -> ()"); + +#endif // !USE_ROCM + + // Rotate Q and K, then write to kv cache for MLA + ops.def( + "concat_and_cache_mla_rope_fused(" + " Tensor positions," + " Tensor! q_pe," + " Tensor! k_pe," + " Tensor kv_c," + " Tensor cos_sin_cache," + " bool is_neox," + " Tensor slot_mapping," + " Tensor! kv_cache," + " str kv_cache_dtype," + " Tensor kv_cache_scale) -> ()"); + + // Convert the key and value cache to fp8 data type. + ops.def( + "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, " + "str kv_cache_dtype) -> ()"); + + // Gather cache blocks from src_cache to dst, dequantizing from + // src_cache's dtype to dst's dtype if necessary. + ops.def( + "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, " + " Tensor block_table, Tensor cu_seq_lens, " + " Tensor token_to_seq, " + " int num_tokens, " + " str kv_cache_dtype, " + " Tensor scale, Tensor? seq_starts) -> ()"); + + ops.def( + "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, " + "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()"); + + ops.def( + "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, " + "Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? " + "seq_starts, Tensor? host_cache=None, Tensor? host_row_ids=None, Tensor? " + "device_row_ids=None) -> ()"); + + ops.def( + "cp_gather_and_upconvert_nvfp4_kv_cache(Tensor src_cache, Tensor! dst, " + "Tensor block_table, Tensor workspace_starts, int batch_size) -> ()"); + + ops.def( + "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor " + "slot_mapping, " + "int quant_block_size, str kv_cache_dtype) -> ()"); + + ops.def("concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()"); + + ops.def( + "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! " + "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()"); +} + +STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ar) { + custom_ar.def( + "init_custom_ar(int[] ipc_tensors, Tensor rank_data, " + "int rank, bool fully_connected) -> int"); + custom_ar.def( + "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, " + "int reg_buffer_sz_bytes) -> ()"); + custom_ar.def("dispose(int fa) -> ()"); + custom_ar.def("meta_size() -> int"); + custom_ar.def("register_buffer(int fa, int[] ipc_tensors) -> ()"); + custom_ar.def("get_graph_buffer_ipc_meta(int fa) -> (int[], int[])"); + custom_ar.def( + "register_graph_buffers(int fa, int[][] handles, int[][] offsets) -> ()"); + custom_ar.def("allocate_shared_buffer_and_handle(int size) -> (int, Tensor)"); + custom_ar.def("open_mem_handle(Tensor mem_handle) -> int"); + custom_ar.def("free_shared_buffer(int ptr) -> ()"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ar) { + custom_ar.impl("init_custom_ar", TORCH_BOX(&init_custom_ar)); + custom_ar.impl("all_reduce", TORCH_BOX(&all_reduce)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CPU, custom_ar) { + custom_ar.impl("open_mem_handle", TORCH_BOX(&open_mem_handle)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CompositeExplicitAutograd, custom_ar) { + custom_ar.impl("dispose", TORCH_BOX(&dispose)); + custom_ar.impl("meta_size", TORCH_BOX(&meta_size)); + custom_ar.impl("register_buffer", TORCH_BOX(®ister_buffer)); + custom_ar.impl("get_graph_buffer_ipc_meta", + TORCH_BOX(&get_graph_buffer_ipc_meta)); + custom_ar.impl("register_graph_buffers", TORCH_BOX(®ister_graph_buffers)); + custom_ar.impl("allocate_shared_buffer_and_handle", + TORCH_BOX(&allocate_shared_buffer_and_handle)); + custom_ar.impl("free_shared_buffer", TORCH_BOX(&free_shared_buffer)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) { + ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) { + ops.impl("swap_blocks", TORCH_BOX(&swap_blocks)); + ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache)); + ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash)); + ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla)); + ops.impl("concat_and_cache_mla_grouped", + TORCH_BOX(&concat_and_cache_mla_grouped)); + +#ifndef USE_ROCM + ops.impl("hisparse_resolve_residency", + TORCH_BOX(&hisparse_resolve_residency)); + ops.impl("hisparse_invalidate_written_slots", + TORCH_BOX(&hisparse_invalidate_written_slots)); + ops.impl("hisparse_gather_plan", TORCH_BOX(&hisparse_gather_plan)); + ops.impl("hisparse_gather_compact", TORCH_BOX(&hisparse_gather_compact)); +#endif // !USE_ROCM + ops.impl("concat_and_cache_mla_rope_fused", + TORCH_BOX(&concat_and_cache_mla_rope_fused)); + ops.impl("convert_fp8", TORCH_BOX(&convert_fp8)); + ops.impl("gather_and_maybe_dequant_cache", + TORCH_BOX(&gather_and_maybe_dequant_cache)); + ops.impl("cp_gather_cache", TORCH_BOX(&cp_gather_cache)); + ops.impl("cp_gather_and_upconvert_fp8_kv_cache", + TORCH_BOX(&cp_gather_and_upconvert_fp8_kv_cache)); + ops.impl("cp_gather_and_upconvert_nvfp4_kv_cache", + TORCH_BOX(&cp_gather_and_upconvert_nvfp4_kv_cache)); + ops.impl("indexer_k_quant_and_cache", TORCH_BOX(&indexer_k_quant_and_cache)); + ops.impl("concat_mla_q", TORCH_BOX(&concat_mla_q)); + ops.impl("cp_gather_indexer_k_quant_cache", + TORCH_BOX(&cp_gather_indexer_k_quant_cache)); +} + +REGISTER_EXTENSION(_C_stable_libtorch) diff --git a/csrc/torch_utils.h b/csrc/torch_utils.h index 898b9e113e19..f02346739e0a 100644 --- a/csrc/torch_utils.h +++ b/csrc/torch_utils.h @@ -1,17 +1,88 @@ #pragma once -// Shared TORCH_UTILS_CHECK across both libtorch stable and unstable source -// files. Keep this header free of CUTLASS/CUTE so attention/quant headers can -// use it. -// -// If TORCH_TARGET_VERSION is defined, we are building _C_stable_libtorch.so so -// use STD_TORCH_CHECK via header-only. -// Otherwise, use TORCH_CHECK via torch/all.h. - -#ifdef TORCH_TARGET_VERSION - #include - #define TORCH_UTILS_CHECK STD_TORCH_CHECK -#else - #include - #define TORCH_UTILS_CHECK TORCH_CHECK -#endif +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include + +// Stable ABI equivalent of TORCH_CHECK_NOT_IMPLEMENTED. +#define STD_TORCH_CHECK_NOT_IMPLEMENTED(cond, ...) \ + STD_TORCH_CHECK(cond, "NotImplementedError: ", __VA_ARGS__) + +// Device properties cache for stable ABI compatibility. +// Uses raw CUDA/HIP APIs instead of ATen functions. +// Using inline ensures a single instance across all translation units. +inline std::deque device_flags; +inline std::vector device_properties; +inline std::once_flag vectors_init_flag; + +inline void do_init_device_vectors() { + int device_count; + cudaError_t err = cudaGetDeviceCount(&device_count); + if (err != cudaSuccess) { + STD_TORCH_CHECK(false, "cudaGetDeviceCount failed: " + + std::string(cudaGetErrorString(err))); + } + device_flags.resize(device_count); + device_properties.resize(device_count); +} + +inline void initDeviceVectors() { + std::call_once(vectors_init_flag, do_init_device_vectors); +} + +inline void initDeviceProperty(int device_index) { + cudaDeviceProp device_prop{}; + cudaError_t err = cudaGetDeviceProperties(&device_prop, device_index); + if (err != cudaSuccess) { + STD_TORCH_CHECK(false, "cudaGetDeviceProperties failed: " + + std::string(cudaGetErrorString(err))); + } + device_properties[device_index] = device_prop; +} + +// Get device properties using raw CUDA/HIP APIs (stable ABI compatible). +// Caches results per device so cudaGetDeviceProperties is called at most once +// per device. +inline cudaDeviceProp* get_device_prop() { + initDeviceVectors(); + int device_index; + cudaError_t err = cudaGetDevice(&device_index); + if (err != cudaSuccess) { + STD_TORCH_CHECK( + false, "cudaGetDevice failed: " + std::string(cudaGetErrorString(err))); + } + STD_TORCH_CHECK(device_index >= 0 && static_cast(device_index) < + device_properties.size(), + "CUDA device index " + std::to_string(device_index) + + " out of range [0, " + + std::to_string(device_properties.size()) + ")"); + + std::call_once(device_flags[device_index], initDeviceProperty, device_index); + return &device_properties[device_index]; +} + +// Utility to get the current CUDA stream for a given device using stable APIs. +// Returns a cudaStream_t for use in kernel launches. +inline cudaStream_t get_current_cuda_stream(int32_t device_index = -1) { + void* stream_ptr = nullptr; + TORCH_ERROR_CODE_CHECK( + aoti_torch_get_current_cuda_stream(device_index, &stream_ptr)); + return reinterpret_cast(stream_ptr); +} + +// Utility to get the current cuBLAS handle using stable APIs. +inline cublasHandle_t get_current_cuda_blas_handle() { + void* blas_handle_ptr = nullptr; + TORCH_ERROR_CODE_CHECK(torch_get_current_cuda_blas_handle(&blas_handle_ptr)); + return reinterpret_cast(blas_handle_ptr); +} diff --git a/csrc/libtorch_stable/type_convert.cuh b/csrc/type_convert.cuh similarity index 100% rename from csrc/libtorch_stable/type_convert.cuh rename to csrc/type_convert.cuh diff --git a/csrc/libtorch_stable/vocab_parallel_embedding_kernels.cu b/csrc/vocab_parallel_embedding_kernels.cu similarity index 100% rename from csrc/libtorch_stable/vocab_parallel_embedding_kernels.cu rename to csrc/vocab_parallel_embedding_kernels.cu diff --git a/docs/design/fusions.md b/docs/design/fusions.md index 4095c41a7e63..6f76838f5560 100644 --- a/docs/design/fusions.md +++ b/docs/design/fusions.md @@ -232,7 +232,7 @@ Supported quantization scheme/hardware combinations: - Pass: [`vllm/compilation/passes/fusion/rms_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rms_quant_fusion.py) - ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) -- CUDA/HIP kernels: [`csrc/libtorch_stable/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/layernorm_quant_kernels.cu) +- CUDA/HIP kernels: [`csrc/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_quant_kernels.cu) ### SiLU+Mul + Quantization (`fuse_act_quant`) @@ -258,7 +258,7 @@ Supported quantization scheme/hardware combinations: - Pass: [`vllm/compilation/passes/fusion/act_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/act_quant_fusion.py) - ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) - CUDA/HIP kernels: [`csrc/quantization/`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/) -- Fused SiLU+Mul+BlockQuant kernel: [`csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu) +- Fused SiLU+Mul+BlockQuant kernel: [`csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu) ### RMSNorm + Padding (`fuse_act_padding`) diff --git a/tests/kernels/test_bf16_skinny_gemm.py b/tests/kernels/test_bf16_skinny_gemm.py index 7ca207c8e4ad..173a2ea34432 100644 --- a/tests/kernels/test_bf16_skinny_gemm.py +++ b/tests/kernels/test_bf16_skinny_gemm.py @@ -185,10 +185,7 @@ def test_every_dsv3_routed_shape_is_instantiated() -> None: let (3216, 7168) ship without its DISPATCH_DSV3_SHAPE(7168, 3216). """ source = ( - Path(__file__).resolve().parents[2] - / "csrc" - / "libtorch_stable" - / "dsv3_fused_a_gemm.cu" + Path(__file__).resolve().parents[2] / "csrc" / "dsv3_fused_a_gemm.cu" ).read_text(encoding="utf-8") # Benchmark-only shapes live behind VLLM_K3_BENCH_SHAPES and are not built # by default, so they must not count as available. diff --git a/tools/ci/kernel_symbol_map.py b/tools/ci/kernel_symbol_map.py index c88f2d9ad551..3177a10bba9b 100755 --- a/tools/ci/kernel_symbol_map.py +++ b/tools/ci/kernel_symbol_map.py @@ -32,13 +32,13 @@ "errors": [], # [{"object", "source", "error"}, ...] "objects": [ { - "source": "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu", + "source": "csrc/fused_qknorm_rope_kernel.cu", "target": "_C", "object": "CMakeFiles/_C.dir/csrc/.../fused_qknorm_rope_kernel.cu.o", "device": true, "symbols": ["_Z21fusedQKNormRopeKernelI...", ...], - "deps": ["csrc/libtorch_stable/fused_qknorm_rope_kernel.cu", - "csrc/cuda_compat.h", ...] + "deps": ["csrc/fused_qknorm_rope_kernel.cu", + "csrc/core/cuda_compat.h", ...] # "error": "..." # present when cuobjdump failed on it }, ... diff --git a/vllm/distributed/device_communicators/quick_all_reduce.py b/vllm/distributed/device_communicators/quick_all_reduce.py index 406869ea2c4b..0aa2a6090261 100644 --- a/vllm/distributed/device_communicators/quick_all_reduce.py +++ b/vllm/distributed/device_communicators/quick_all_reduce.py @@ -28,7 +28,7 @@ class QuickReduceRegime(Enum): - # Keep integer ids aligned with csrc/quickreduce/quick_reduce.h + # Keep integer ids aligned with csrc/rocm/quickreduce/quick_reduce.h FP = 0 INT8 = 1 INT6 = 2 diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py index da8c69eed0a2..a541debfe9f3 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py @@ -104,7 +104,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # azp_adj is the AZP adjustment term, used to account for weights. # It does not depend on scales or azp, so it is the same for # static and dynamic quantization. - # For more details, see csrc/quantization/w8a8/cutlass/Epilogues.md + # For more details, see + # csrc/quantization/w8a8/cutlass/Epilogues.md # https://github.com/vllm-project/vllm/blob/main/csrc/quantization/w8a8/cutlass/Epilogues.md if not config.input_symmetric: weight = getattr(layer, w_q_name) diff --git a/vllm/models/common/ops/fused_qk_rmsnorm.py b/vllm/models/common/ops/fused_qk_rmsnorm.py index e64e81821062..fec8719a899c 100644 --- a/vllm/models/common/ops/fused_qk_rmsnorm.py +++ b/vllm/models/common/ops/fused_qk_rmsnorm.py @@ -68,7 +68,8 @@ def kernel( weight_ptr = kv_weight_ptr row_out = kv_out_ptr + token_idx * kv_out_stride - # RMSNorm in fp32 throughout — matches csrc/layernorm_kernels.cu's + # RMSNorm in fp32 throughout — matches + # csrc/layernorm_kernels.cu's # `(scalar_t)(x * s_variance * w)` and DeepseekV4's compressor kernel, which # keep x, rrms, and w all in fp32 and perform a single cast at store. block = tl.arange(0, BLOCK_SIZE) diff --git a/vllm/models/kimi_k3/amd/ops/kda_chunk.py b/vllm/models/kimi_k3/amd/ops/kda_chunk.py index 035a49f852d6..a3fbada908fa 100644 --- a/vllm/models/kimi_k3/amd/ops/kda_chunk.py +++ b/vllm/models/kimi_k3/amd/ops/kda_chunk.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ROCm entry point for the fused Kimi-K3 KDA chunk kernel. -The kernel in ``csrc/libtorch_stable/kimi_k3/fused_kda_chunk_kernel_rocm.cu`` +The kernel in ``csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu`` replaces the chunk-state recurrence and the output GEMM of the Triton chunk path with a single launch that keeps the per-chunk state in registers, so the ``[chunks, H, V, K]`` state tensor and the recomputed values never reach HBM. diff --git a/vllm/models/kimi_k3/amd/ops/kda_decode.py b/vllm/models/kimi_k3/amd/ops/kda_decode.py index 132cd58cf816..682776a80726 100644 --- a/vllm/models/kimi_k3/amd/ops/kda_decode.py +++ b/vllm/models/kimi_k3/amd/ops/kda_decode.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ROCm entry points for the fused Kimi-K3 KDA decode kernel. -The kernel in ``csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu`` +The kernel in ``csrc/kimi_k3/fused_kda_decode_kernel_rocm.cu`` replaces, for a pure non-speculative decode batch, the three Triton launches and two copies the AMD KDA layer otherwise runs per layer: the packed causal conv1d update, the recurrent delta-rule step, and the gated output RMSNorm. diff --git a/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py index 79de3af0bc22..a39ec787e0bb 100644 --- a/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py +++ b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py @@ -3,7 +3,7 @@ """Fused MLA prefill and decode epilogues for Kimi-K3. Thin wrappers over the CUDA ops in -``csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu``, which +``csrc/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu``, which mirror ``fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_{bf16,fp8}_insert``. - ``fused_mla_key_concat_kv_cache_insert`` (bf16): optionally apply RoPE, diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 0ff34f8ac5e0..4fe0be0b1f68 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -118,7 +118,7 @@ elements, stored permuted (an 8x4 -> 4x8 transpose: the scale for element block `s` lives at byte `8 * (s & 3) + (s >> 2)`) so that the 8 scales one FlashMLA dequant thread needs are contiguous. See the layout comment in - `csrc/libtorch_stable/cache_kernels.cu`. + `csrc/cache_kernels.cu`. """ diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index f1b0fa688eda..f7474b6bedad 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -712,7 +712,7 @@ def __init__( "HiSparse requires its compiled _C_cache_ops CUDA kernels " "(host-resident decode has no Python fallback). Rebuild vLLM " "from source so " - "csrc/libtorch_stable/hisparse_kernels.cu is included." + "csrc/hisparse_kernels.cu is included." ) self.max_num_reqs = max_num_reqs self.row_width = row_width From fa39582756dcbf25b65cc36d8daa465b19134372 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Mon, 12 Oct 2026 13:00:47 -0400 Subject: [PATCH 316/624] [Frontend] Serve LiquidAI/d1-3B on /v1/systemone (#60676) Signed-off-by: mgoin Co-authored-by: Simon Veitner Co-authored-by: Claude --- .../online_serving/structured_decisions.md | 25 ++- .../structured_decisions/test_strategies.py | 52 ++++++ .../structured_decisions/strategies.py | 162 ++++++++++++++++-- 3 files changed, 222 insertions(+), 17 deletions(-) diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md index 4d8d65d89e94..ce09d6f81fee 100644 --- a/docs/serving/online_serving/structured_decisions.md +++ b/docs/serving/online_serving/structured_decisions.md @@ -156,6 +156,29 @@ vllm serve Cloudflare/clef-omni with it when every video has one. Audio clips cannot be sent together with videos that have soundtracks. +## Liquid d1 + +[`LiquidAI/d1-3B`](https://huggingface.co/LiquidAI/d1-3B) is LFM2.5-VL-3B +post-trained to answer decisions as the first token of its reply. vLLM serves +it as a generative model. + +```bash +vllm serve LiquidAI/d1-3B +``` + +- The server writes the prompt and reads the label tokens the way the + `prompt.py` shipped with the checkpoint does, so picks and probabilities + match `model.system_one()`. `confidence` and `usage` follow the definitions + on this page instead. +- A JSON state is written as indented JSON. A `null` state leaves the images, + or the question alone, as the prompt. +- `images` are placed before the state and shrunk to at most 1024 x 1024 + pixels, as the reference does. +- A `choice` whose option names are all single letters is labeled with those + letters. +- When each request carries one question over a state that changes from + request to request, `--no-enable-prefix-caching` lowers latency. + ## Request fields | field | meaning | @@ -165,7 +188,7 @@ vllm serve Cloudflare/clef-omni | `instructions` | optional context placed ahead of the questions | | `chat_template_kwargs` | passed to the chat template, for example `{"enable_thinking": false}` | | `cache_salt` | optional [prefix-cache salt](../../usage/security.md#cache-salting) applied to every question | -| `images`, `audio`, `videos` | optional media URLs or data URLs, for models that take them ([Clef-Omni](#clef-omni)) | +| `images`, `audio`, `videos` | optional media URLs or data URLs, for models that take them ([Clef-Omni](#clef-omni), [Liquid d1](#liquid-d1), images only) | A question with any field other than `type`, `instructions` and `criteria` is rejected with a 400. diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index 8d38e61e6248..dae545d1f04c 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -26,6 +26,7 @@ from vllm.entrypoints.generate.structured_decisions.question_types import ( LABELS, StructuredDecisionError, + build_question, ) from vllm.entrypoints.generate.structured_decisions.serving import ( ServingStructuredDecisions, @@ -34,8 +35,10 @@ from vllm.entrypoints.generate.structured_decisions.strategies import ( DiffusionGemmaCanvasStrategy, JointSchemaStrategy, + LiquidStrategy, NextTokenStrategy, ReadContext, + cap_pixels, reply_label_ids, select_read_strategy, state_text, @@ -79,6 +82,8 @@ def test_strategy_selection(): select_read_strategy(model("DiffusionGemmaForBlockDiffusion"), generate) is DiffusionGemmaCanvasStrategy ) + lfm = "Lfm2VlForConditionalGeneration" + assert select_read_strategy(model(lfm), generate) is LiquidStrategy with pytest.raises(ValueError, match="does not support LlamaForCausalLM"): select_read_strategy(model("LlamaForCausalLM"), generate) with pytest.raises(ValueError, match="not raw_logits"): @@ -573,3 +578,50 @@ async def process_for_engine_async(prompt, arrival_time): ) assert isinstance(response, ErrorResponse) assert response.error.code == 400 + + +@pytest.mark.parametrize( + "type_name,criteria,ask", + [ + ( + "choice", + {"billing": "Charges", "app_fault": None}, + "Which?\n\nOptions:\nA Charges\nB app fault\n\n" + "Reply with the option code only.", + ), + ( + "choice", + {"a": "Berlin", "b": "Paris"}, + "Which?\n\nOptions:\na Berlin\nb Paris\n\nReply with the option code only.", + ), + ( + "noul", + {"true": "asks to cancel", "false": "anything else"}, + "Which?\nYes: asks to cancel\nNo: anything else\n\n" + "Reply with yes or no only.", + ), + ( + "score", + ["low", "high"], + "Which?\n\n0 low\n1 high\n\nReply with a single digit 0-1 only.", + ), + ], +) +def test_liquid_prompt_matches_the_checkpoint(type_name, criteria, ask): + # The wording of prompt.py in LiquidAI/d1-3B, which the model was trained on. + strategy = object.__new__(LiquidStrategy) + q = build_question("q", type_name, "Which?", criteria, 26) + assert strategy._content(strategy.render_state(None), q) == ask + assert strategy._content("Hi.", q) == f"Hi.\n\n\nQUESTION:\n{ask}" + state = strategy.render_state({"k": "é"}) + assert strategy._content(state, q).startswith('{\n "k": "é"\n}\n\n\nQUEST') + + +def test_liquid_caps_images_like_the_checkpoint(): + # runner.cap_pixels in LiquidAI/d1-3B: 1600x1598 becomes 1024x1023. + cap = LiquidStrategy.max_image_pixels + assert cap is not None + assert cap_pixels(Image.new("RGB", (1600, 1598)), cap).size == (1024, 1023) + small = Image.new("RGB", (640, 480)) + assert cap_pixels(small, cap) is small + assert NextTokenStrategy.max_image_pixels is None diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index f4365924ddf0..5567cb808ab5 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -18,6 +18,7 @@ from typing import Any, ClassVar, TypeVar import numpy as np +from PIL import Image import vllm.envs as envs from vllm.config import ModelConfig @@ -92,6 +93,16 @@ class ReadStrategy(ABC): def __init__(self, context: ReadContext): self.context = context + @cached_property + def _connector(self) -> MediaConnector: + model_config = self.context.engine_client.model_config + return MEDIA_CONNECTOR_REGISTRY.load( + envs.VLLM_MEDIA_CONNECTOR, + media_io_kwargs=model_config.get_multimodal_config().media_io_kwargs, + allowed_local_media_path=model_config.allowed_local_media_path, + allowed_media_domains=model_config.allowed_media_domains, + ) + @abstractmethod def limits(self) -> DecisionLimits: ... @@ -162,12 +173,25 @@ def type_label_ids( return label_ids +def cap_pixels(image: Image.Image, max_pixels: int) -> Image.Image: + """Shrink ``image`` to at most ``max_pixels`` pixels (bicubic).""" + w, h = image.size + if w * h <= max_pixels: + return image + scale = math.sqrt(max_pixels / (w * h)) + size = (max(1, int(w * scale)), max(1, int(h * scale))) + return image.resize(size, Image.Resampling.BICUBIC) + + class NextTokenStrategy(ReadStrategy): """Autoregressive models. Each question is one request: the state, then the question with its labeled options, and the logprobs of the label tokens as the reply's first token. The requests share the state. With prefix caching, it is prefilled once.""" + #: Images are shrunk to at most this many pixels before the processor. + max_image_pixels: ClassVar[int | None] = None + def __init__(self, context: ReadContext): super().__init__(context) # Every prompt ends with the same generation prompt, so a label's token @@ -210,6 +234,27 @@ def _read_input( """The engine input for a read of the rendered prompt.""" return engine_input + def _label_groups(self, question: Question) -> list[list[int]]: + """The tokens that count for each label. A label scores its best one.""" + return [[self.label_ids[label]] for label in question.labels] + + def _content(self, state: str, question: Question) -> str: + """The user turn: the state, then the question.""" + return f"{state}\n\n{question.type.prompt(question)}" + + async def _image_parts(self, urls: Sequence[str]) -> list[dict[str, Any]]: + """The images as chat content parts, fetched once so that every + question reads the same pictures.""" + if not urls: + return [] + try: + images = await asyncio.gather(*map(self._connector.fetch_image_async, urls)) + except ValueError as e: + raise StructuredDecisionError(f"media: {e}") from e + if self.max_image_pixels is not None: + images = [cap_pixels(image, self.max_image_pixels) for image in images] + return [{"type": "image_pil", "image_pil": image} for image in images] + def limits(self) -> DecisionLimits: return DecisionLimits(max_questions=64, max_options=len(LABELS)) @@ -234,6 +279,9 @@ async def _render( default_template=ctx.chat_template, default_template_content_format=ctx.chat_template_content_format, default_template_kwargs=ctx.default_chat_template_kwargs, + # A read that fails before it reaches the engine must not leave the + # sender cache believing the engine has the images. + skip_mm_cache=True, ) prompt_ids = extract_prompt_components( ctx.engine_client.model_config, engine_input @@ -256,10 +304,14 @@ async def read( ) -> list[QuestionRead]: ctx = self.context read_request = self._read_request(chat_template_kwargs, cache_salt) + image_parts = await self._image_parts((media or {}).get("image", ())) slots, engine_inputs, params = [], [], [] for q in questions: - messages = [{"role": "user", "content": f"{state}\n\n{q.type.prompt(q)}"}] + content: Any = self._content(state, q) + if image_parts: + content = [*image_parts, {"type": "text", "text": content}] + messages = [{"role": "user", "content": content}] if instructions: messages.insert(0, {"role": "system", "content": instructions}) engine_input, prompt_ids = await self._render(read_request, messages) @@ -268,9 +320,11 @@ async def read( "these chat options end the prompt differently, so the " "labels' tokens are unknown" ) - slots.append([self.label_ids[label] for label in q.labels]) + slots.append(self._label_groups(q)) engine_inputs.append(self._read_input(engine_input, prompt_ids)) - params.append(self._sampling_params(slots[-1], prompt_ids)) + params.append( + self._sampling_params([i for ids in slots[-1] for i in ids], prompt_ids) + ) label_reads = await next_token_label_reads( ctx.engine_client, @@ -283,14 +337,17 @@ async def read( ) reads = [] - for ids, label_read in zip(slots, label_reads): + for groups, label_read in zip(slots, label_reads): output = label_read.result.outputs[0] + logprobs = iter(label_read.logprobs) reads.append( QuestionRead( - probs=label_softmax(label_read.logprobs), + probs=label_softmax( + [max(next(logprobs) for _ in g) for g in groups] + ), label_mass=sum(math.exp(lp) for lp in label_read.logprobs), argmax_is_label=bool(output.token_ids) - and output.token_ids[0] in ids, + and any(output.token_ids[0] in ids for ids in groups), input_tokens=len(label_read.result.prompt_token_ids or ()), output_tokens=len(output.token_ids), ) @@ -373,6 +430,88 @@ def _read_input( ) +class LiquidStrategy(NextTokenStrategy): + """Liquid decision models (d1). The prompt and the label tokens are the + ones the checkpoint was trained with, from the `prompt.py` it ships.""" + + modalities = frozenset({"image"}) + # runner.py bounds every picture at 1024 x 1024 pixels before the processor. + max_image_pixels = 1024 * 1024 + + def render_state(self, state: Any) -> str: + if state is None: + return "" + if isinstance(state, str): + return state + return json.dumps(state, ensure_ascii=False, indent=2) + + def _single_tokens(self, texts: Sequence[str]) -> list[int]: + ids: list[int] = [] + for text in texts: + encoded = self._tokenizer().encode(text, add_special_tokens=False) + if len(encoded) == 1 and encoded[0] not in ids: + ids.append(encoded[0]) + return ids + + @staticmethod + def _codes(question: Question) -> list[str]: + """A choice's codes are A to Z, or the option names when those are + already single letters.""" + names = [o.name.strip() for o in question.options] + if all(len(n) == 1 and n.isalpha() for n in names): + return names + return list(question.labels) + + def _label_groups(self, question: Question) -> list[list[int]]: + if question.type.name == "noul": + forms: list[Sequence[str]] = [("yes", "Yes", "YES"), ("no", "No", "NO")] + elif question.type.name == "choice": + forms = [(code, f" {code}") for code in self._codes(question)] + else: + forms = [(label,) for label in question.labels] + groups = [self._single_tokens(f) for f in forms] + seen: set[int] = set() + for (label, *_), ids in zip(forms, groups): + if not ids or self._single_tokens([label]) != ids[:1] or seen & set(ids): + raise StructuredDecisionError( + f"question {question.id!r}: label {label!r} is not one " + "distinct token for this model" + ) + seen.update(ids) + return groups + + def _content(self, state: str, question: Question) -> str: + q = question + if q.type.name == "choice": + lines = "\n".join( + f"{code} {o.description or o.name.replace('_', ' ')}" + for code, o in zip(self._codes(q), q.options) + ) + ask = ( + f"{q.instructions}\n\nOptions:\n{lines}\n\n" + "Reply with the option code only." + ) + elif q.type.name == "noul": + yes, no = (o.description for o in q.options) + sides = f"\nYes: {yes}\nNo: {no}" if yes or no else "" + ask = f"{q.instructions}{sides}\n\nReply with yes or no only." + elif q.type.name == "score": + legend = "\n".join( + f"{label} {o.name}" for label, o in zip(q.labels, q.options) + ) + ask = ( + f"{q.instructions}\n\n{legend}\n\n" + f"Reply with a single digit 0-{len(q.options) - 1} only." + ) + else: + raise StructuredDecisionError( + f"question {q.id!r}: this model does not answer {q.type.name!r} " + "questions" + ) + # A null state leaves the images, or the question alone, as the prompt. + return f"{state}\n\n\nQUESTION:\n{ask}" if state else ask + + class JointSchemaStrategy(ReadStrategy): """Clef-Omni. The whole request is one pooling read: the state, then every question with its options, as the model was trained to read them. The @@ -442,16 +581,6 @@ def _render(value: Any) -> str: def _tokens(self, text: str) -> list[int]: return self.tokenizer.encode(text, add_special_tokens=False) - @cached_property - def _connector(self) -> MediaConnector: - model_config = self.context.engine_client.model_config - return MEDIA_CONNECTOR_REGISTRY.load( - envs.VLLM_MEDIA_CONNECTOR, - media_io_kwargs=model_config.get_multimodal_config().media_io_kwargs, - allowed_local_media_path=model_config.allowed_local_media_path, - allowed_media_domains=model_config.allowed_media_domains, - ) - async def _fetch( self, urls: Sequence[str], media_io: MediaIO[_M], timeout: int ) -> list[_M]: @@ -682,6 +811,7 @@ async def read( READ_STRATEGIES: dict[str, type[ReadStrategy]] = { "ClefOmniForJointSchema": JointSchemaStrategy, + "Lfm2VlForConditionalGeneration": LiquidStrategy, "Qwen3ForCausalLM": NextTokenStrategy, "Qwen3_5ForConditionalGeneration": NextTokenStrategy, "Qwen3_5MoeForConditionalGeneration": NextTokenStrategy, From 6e6e86078afe73af5180a9888e88751d51ef1b31 Mon Sep 17 00:00:00 2001 From: haotiaa Date: Mon, 12 Oct 2026 10:05:34 -0700 Subject: [PATCH 317/624] [Bugfix][Platform] Don't resolve current_platform during import vllm (#60766) Signed-off-by: Haotian An --- .../vllm_add_dummy_platform/__init__.py | 3 +++ tests/plugins_tests/test_platform_plugins.py | 17 +++++++++++++++++ vllm/utils/torch_utils.py | 18 ++++++++++++++---- 3 files changed, 34 insertions(+), 4 deletions(-) diff --git a/tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/__init__.py b/tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/__init__.py index 280b68514e19..2d28e151d328 100644 --- a/tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/__init__.py +++ b/tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/__init__.py @@ -3,6 +3,9 @@ def dummy_platform_plugin() -> str | None: + # Real platform plugins often import vLLM modules while registering. + from vllm.utils.torch_utils import direct_register_custom_op # noqa: F401 + return "vllm_add_dummy_platform.dummy_platform.DummyPlatform" diff --git a/tests/plugins_tests/test_platform_plugins.py b/tests/plugins_tests/test_platform_plugins.py index 6d32c4c6d6f0..a75e36a7d1c8 100644 --- a/tests/plugins_tests/test_platform_plugins.py +++ b/tests/plugins_tests/test_platform_plugins.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import subprocess +import sys + import pytest import torch @@ -31,6 +34,20 @@ def test_platform_plugins(): ) +def test_import_vllm_does_not_resolve_platform(): + # Platform plugins must be loaded after `import vllm` completes; resolving + # current_platform during it runs plugins against a half-imported vLLM. + # Needs a fresh interpreter, since vllm is already imported here. + code = ( + "import sys, vllm\n" + "p = sys.modules.get('vllm.platforms')\n" + "assert p is None or p._current_platform is None, p._init_trace\n" + "from vllm.platforms import current_platform\n" + "assert current_platform.device_name == 'DummyDevice'\n" + ) + subprocess.run([sys.executable, "-c", code], check=True) + + def test_oot_custom_op(default_vllm_config, monkeypatch: pytest.MonkeyPatch): # simulate workload by running an example load_general_plugins() diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index dce5e84fce25..4bf2e3397ca6 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -73,7 +73,15 @@ T = TypeVar("T") -PIN_MEMORY = is_pin_memory_available() +if TYPE_CHECKING: + PIN_MEMORY: bool + + +def __getattr__(name: str) -> Any: + # Lazy so that importing this module does not resolve current_platform. + if name == "PIN_MEMORY": + return is_pin_memory_available() + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") def is_quantized_kv_cache(kv_cache_dtype: str) -> bool: @@ -718,13 +726,15 @@ def async_tensor_h2d( data = torch.from_numpy(data) if isinstance(data, torch.Tensor): t = data - if PIN_MEMORY and not t.is_pinned(): + if is_pin_memory_available() and not t.is_pinned(): # Stage in pinned, contiguous buffer to ensure fully async copy. t = torch.empty( t.shape, dtype=dtype or t.dtype, device="cpu", pin_memory=True ).copy_(t) else: - t = torch.tensor(data, dtype=dtype, pin_memory=PIN_MEMORY, device="cpu") + t = torch.tensor( + data, dtype=dtype, pin_memory=is_pin_memory_available(), device="cpu" + ) assert t.is_cpu if out is not None: @@ -736,7 +746,7 @@ def async_tensor_h2d( def np_to_pinned_tensor(array: np.ndarray) -> torch.Tensor: t = torch.from_numpy(array) - return t.pin_memory() if PIN_MEMORY else t + return t.pin_memory() if is_pin_memory_available() else t def make_ndarray_with_pad( From c599b5971faf7ed8bff4bd1b14713293c20082d6 Mon Sep 17 00:00:00 2001 From: Chester Date: Tue, 13 Oct 2026 01:05:55 +0800 Subject: [PATCH 318/624] [Bugfix][Frontend] Return all logprobs when completions requests logprobs=-1 (#55717) Signed-off-by: Chester <42577861+superdiaodiao@users.noreply.github.com> --- .../completion/test_completion_error.py | 40 +++++++++++++++++++ vllm/entrypoints/openai/completion/serving.py | 4 +- 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index bbbb61f9c9c8..17934026efe6 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -16,6 +16,7 @@ from vllm.entrypoints.scale_out.render.serving import ServingRender from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.exceptions import GenerationError, VLLMValidationError +from vllm.logprobs import Logprob from vllm.outputs import CompletionOutput, RequestOutput from vllm.renderers.deepseek_v4 import DeepseekV4Renderer from vllm.renderers.hf import HfRenderer @@ -841,6 +842,45 @@ def test_logprobs_minus_one_allowed(): assert request.logprobs == -1 +def _minimal_serving_completion() -> OpenAIServingCompletion: + """A serving object with only what `_create_completion_logprobs` reads.""" + serving = OpenAIServingCompletion.__new__(OpenAIServingCompletion) + serving.return_tokens_as_token_ids = False + return serving + + +@pytest.mark.parametrize( + ("num_output_top_logprobs", "expected"), + [(-1, 3), (1, 2), (2, 3)], +) +def test_completion_logprobs_honors_requested_count( + num_output_top_logprobs: int, expected: int +): + """`logprobs=-1` asks for every logprob, so none may be filtered out. + + The validator accepts -1 and the sampler expands it to the vocabulary, so + dropping the candidates here would return an empty `top_logprobs` for every + position while still charging for the full computation. + """ + tokenizer = MagicMock() + tokenizer.decode = lambda token_ids: f"tok{token_ids[0]}" + step_logprobs = { + 10: Logprob(logprob=-0.1, decoded_token="a"), + 11: Logprob(logprob=-0.2, decoded_token="b"), + 12: Logprob(logprob=-0.3, decoded_token="c"), + } + + logprobs = _minimal_serving_completion()._create_completion_logprobs( + token_ids=[10], + top_logprobs=[step_logprobs], + num_output_top_logprobs=num_output_top_logprobs, + tokenizer=tokenizer, + ) + + assert logprobs.top_logprobs is not None + assert len(logprobs.top_logprobs[0]) == expected + + def test_logprobs_below_minus_one_rejected(): """Values more negative than -1 stay invalid.""" with pytest.raises(Exception, match="must be a positive value or -1"): diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index 7f7bdac0757b..e5ba20c52fab 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -722,7 +722,9 @@ def _create_completion_logprobs( return_as_token_id=should_return_as_token_id, ): max(top_lp[1].logprob, -9999.0) for i, top_lp in enumerate(step_top_logprobs.items()) - if logprob_token_ids or num_output_top_logprobs >= i + if logprob_token_ids + or num_output_top_logprobs == -1 + or num_output_top_logprobs >= i } ) From 12e6a585ae020cdcef9ba61aae45eec5ba29591d Mon Sep 17 00:00:00 2001 From: Laith Sakka Date: Mon, 12 Oct 2026 10:09:57 -0700 Subject: [PATCH 319/624] =?UTF-8?q?[Fix][Compile]=20Declare=20group=5Fshap?= =?UTF-8?q?e=20in=20the=20static=20FP8=20quant=20fusion=20=E2=80=A6=20(#61?= =?UTF-8?q?108)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Laith Sakka Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/compilation/passes/fusion/matcher_utils.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm/compilation/passes/fusion/matcher_utils.py b/vllm/compilation/passes/fusion/matcher_utils.py index 6d959c30677a..80be65a9781e 100644 --- a/vllm/compilation/passes/fusion/matcher_utils.py +++ b/vllm/compilation/passes/fusion/matcher_utils.py @@ -481,8 +481,13 @@ def forward_custom( if self.quant_key.scale.static: assert scale is not None + group_shape = self.quant_key.scale.group_shape _, result = auto_functionalized( - self.QUANT_OP, result=result, input=input, scale=scale + self.QUANT_OP, + result=result, + input=input, + scale=scale, + group_shape=[group_shape.row, group_shape.col], ) return result, scale else: From 45bb710ecaa2f0e100d1cab27308186e9f4aa7dd Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Mon, 12 Oct 2026 13:11:26 -0400 Subject: [PATCH 320/624] [CI] Add nightly Python 3.14 smoke test (#60456) Signed-off-by: mgoin Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/misc.yaml | 11 +++++ tests/standalone_tests/python_only_compile.sh | 43 +++++++++++++++++-- 2 files changed, 51 insertions(+), 3 deletions(-) diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 02f78e280c98..bf650fd756ba 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -446,6 +446,17 @@ steps: - setup.py - vllm/platforms/rocm.py +- label: ":nvidia: (H200 MIG 18GB) Python 3.14 Smoke Nightly" + key: python-314-smoke-nightly + depends_on: ~ + optional: true + timeout_in_minutes: 30 + device: h200_18gb + source_file_dependencies: + - tests/standalone_tests/python_only_compile.sh + commands: + - PYTHON_ONLY_COMPILE_PYTHON=3.14 bash standalone_tests/python_only_compile.sh + - label: ":nvidia: (H200 MIG 35GB) Async Engine, Inputs, Utils, Worker" device: h200_35gb key: async-engine-inputs-utils-worker diff --git a/tests/standalone_tests/python_only_compile.sh b/tests/standalone_tests/python_only_compile.sh index 7036dd4314f2..5a6c6961d67c 100644 --- a/tests/standalone_tests/python_only_compile.sh +++ b/tests/standalone_tests/python_only_compile.sh @@ -1,6 +1,10 @@ #!/bin/bash # This script tests if the python only compilation works correctly # for users who do not have any compilers installed on their system +# +# Set PYTHON_ONLY_COMPILE_PYTHON (e.g. 3.14) to install into a fresh uv venv +# with that interpreter and run a short generation, as a smoke test for +# Python versions other than the one the CI image is built with. set -e @@ -173,9 +177,12 @@ elif [[ ! -d vllm ]]; then exit 1 fi -# remove all compilers -apt remove --purge build-essential -y -apt autoremove -y +# remove all compilers, except for the smoke test, whose generation needs a +# C compiler at runtime for torch.compile and Triton +if [[ -z "${PYTHON_ONLY_COMPILE_PYTHON:-}" ]]; then + apt remove --purge build-essential -y + apt autoremove -y +fi rm -f /tmp/changed.file echo 'import os; os.system("touch /tmp/changed.file")' >> vllm/__init__.py @@ -185,12 +192,42 @@ if [[ -n "${rocm_wheel}" ]]; then VLLM_PRECOMPILED_WHEEL_LOCATION="${rocm_wheel}" VLLM_USE_PRECOMPILED=1 python3 setup.py develop --no-deps elif [[ "${is_rocm}" == "1" ]]; then VLLM_PRECOMPILED_WHEEL_COMMIT=$merge_base_commit VLLM_USE_PRECOMPILED=1 python3 setup.py develop --no-deps +elif [[ -n "${PYTHON_ONLY_COMPILE_PYTHON:-}" ]]; then + # Use the CUDA build the image's own torch was installed with, for both the + # new venv's torch and the precompiled wheel variant. + cuda_variant="$(python3 -c 'import torch; print("cu" + torch.version.cuda.replace(".", ""))')" + uv venv --python "${PYTHON_ONLY_COMPILE_PYTHON}" /tmp/python-only-venv + source /tmp/python-only-venv/bin/activate + VLLM_PRECOMPILED_WHEEL_COMMIT=$merge_base_commit VLLM_USE_PRECOMPILED=1 \ + VLLM_PRECOMPILED_WHEEL_VARIANT="${cuda_variant}" \ + uv pip install -e . --torch-backend="${cuda_variant}" else VLLM_PRECOMPILED_WHEEL_COMMIT=$merge_base_commit VLLM_USE_PRECOMPILED=1 pip3 install -vvv -e . fi # Run the script python3 -c 'import vllm' +if [[ -n "${PYTHON_ONLY_COMPILE_PYTHON:-}" ]]; then + python3 - <<'PY' +import os +import sys + +from vllm import LLM + +expected = os.environ["PYTHON_ONLY_COMPILE_PYTHON"] +actual = f"{sys.version_info.major}.{sys.version_info.minor}" +assert actual == expected, f"expected Python {expected}, got {actual}" + +llm = LLM( + "hmellor/tiny-random-LlamaForCausalLM", + max_model_len=256, + gpu_memory_utilization=0.3, +) +outputs = llm.generate(["Hello, my name is"]) +assert outputs[0].outputs[0].token_ids, "no tokens generated" +PY +fi + # Check if the clangd log file was created if [ ! -f /tmp/changed.file ]; then echo "ERROR: changed.file was not created, python only compilation failed" From a7b7bc1fbe1f5c4e47c53b0f02a3ecfbcd8054b9 Mon Sep 17 00:00:00 2001 From: simon-veitner-redhat Date: Mon, 12 Oct 2026 19:21:22 +0200 Subject: [PATCH 321/624] [Metrics] Count finished requests by watermark status (#58715) Signed-off-by: Simon Veitner Co-authored-by: Claude Co-authored-by: Lucas Wilkinson --- docs/features/watermarking.md | 28 ++++++ tests/v1/engine/test_output_processor.py | 94 +++++++++++++++++++++ tests/v1/metrics/test_histogram_buckets.py | 14 +++ tests/watermarking/test_watermarking_e2e.py | 16 ++++ vllm/v1/engine/output_processor.py | 6 +- vllm/v1/metrics/loggers.py | 30 +++++++ vllm/v1/metrics/stats.py | 16 ++++ 7 files changed, 203 insertions(+), 1 deletion(-) diff --git a/docs/features/watermarking.md b/docs/features/watermarking.md index 66b6465eb851..f51ec72eb1ae 100644 --- a/docs/features/watermarking.md +++ b/docs/features/watermarking.md @@ -56,6 +56,34 @@ watermarking algorithm that does not support it natively, at the cost of weaker detectability. See [Speculative decoding](#speculative-decoding). +## Monitoring + +When `watermark_config` is set, the Prometheus counter +`vllm:watermark_requests_total` counts finished requests by `status`: + +- `watermarked`: the request was sampled with the watermark switched on. +- `skipped`: the request asked for the watermark (omitted or `true`) but could + not use it, such as `temperature=0` and trace replay requests. A beam search + request that asks for the watermark runs without it and logs a warning. Its + beam steps count as `not_watermarked`, not `skipped`. +- `not_watermarked`: all other finished requests. This includes explicit + opt-outs, pooling requests, and the internal requests of beam search (one + per beam step), generative scoring and language detection. + +Each finished request has exactly one status, so the sum over `status` equals +`vllm:request_success_total` on the same instance. Requests that the client +aborts are in neither counter. The counter does not exist without +`watermark_config`. With `n > 1`, each sequence counts once. In a +disaggregated prefill/decode deployment, both instances count a request: sum +over decode instances only. For streaming input, the first input chunk sets +the status. + +The status is per request, not per token. A `watermarked` request can contain +tokens without the watermark: accepted draft tokens with +`allow_target_only_watermarking`, tokens whose context repeats +(`deduplicate_contexts`), the first `context_width` generated tokens with +`deduplicate_contexts: "all"`, and tokens that a grammar forces. + ## Architecture `WatermarkConfig` selects an algorithm and PRF. Model Runner V2 constructs the diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index 8eb85765993b..3c186a4e738d 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -6,6 +6,7 @@ from unittest.mock import MagicMock, Mock import numpy as np +import prometheus_client import pytest import torch @@ -17,6 +18,7 @@ MockEngineCore, ) from vllm import PoolingParams +from vllm.config.watermarking import WatermarkConfig from vllm.logprobs import FlatLogprobs, Logprob, PromptLogprobs, SampleLogprobs from vllm.lora.request import LoRARequest from vllm.outputs import ( @@ -41,6 +43,9 @@ RequestOutputCollector, RequestState, ) +from vllm.v1.engine.parallel_sampling import ParentRequest +from vllm.v1.metrics.loggers import PrometheusStatLogger +from vllm.v1.metrics.prometheus import unregister_vllm_metrics from vllm.v1.metrics.stats import IterationStats, PrefillStats, SchedulerStats from vllm.v1.outputs import SamplingMaskLists @@ -1169,6 +1174,95 @@ def test_iteration_stats(dummy_test_vectors): assert iteration_stats.num_generation_tokens == num_active +def test_watermark_status_partitions_finished_requests(dummy_test_vectors): + """Each finished request gets exactly one watermark status.""" + output_processor = OutputProcessor(dummy_test_vectors.tokenizer, log_stats=True) + prompt_token_ids = dummy_test_vectors.prompt_tokens[0] + + def make_request( + request_id: str, + sampling_params: SamplingParams | None = None, + pooling_params: PoolingParams | None = None, + ) -> EngineCoreRequest: + return EngineCoreRequest( + request_id=request_id, + external_req_id=f"{request_id}-ext", + prompt_token_ids=prompt_token_ids, + mm_features=None, + arrival_time=0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + sampling_params=sampling_params, + pooling_params=pooling_params, + ) + + # The frontend resolves the flags before the request reaches the + # OutputProcessor; set them as it does. + def resolved(watermarking: bool, skipped: bool, n: int = 1) -> SamplingParams: + params = SamplingParams(n=n, watermarking=watermarking) + params._watermarking_skipped = skipped + return params + + finish_reasons = { + "wm": FinishReason.ABORT, + "skip": FinishReason.STOP, + # A re-resolved request can carry both flags; it is watermarked. + "sticky": FinishReason.STOP, + "off": FinishReason.ERROR, + "pool": FinishReason.STOP, + } + output_processor.add_request(make_request("wm", resolved(True, False)), None) + output_processor.add_request(make_request("skip", resolved(False, True)), None) + output_processor.add_request(make_request("sticky", resolved(True, True)), None) + output_processor.add_request(make_request("off", resolved(False, False)), None) + output_processor.add_request( + make_request("pool", pooling_params=PoolingParams(task="embed")), None + ) + parent_req = ParentRequest(make_request("parent", resolved(True, False, n=2))) + for idx in range(2): + child_id, child_params = parent_req.get_child_info(idx) + output_processor.add_request( + make_request(child_id, child_params), None, parent_req, idx + ) + finish_reasons[child_id] = FinishReason.LENGTH + + outputs = [ + EngineCoreOutput( + request_id=request_id, + new_token_ids=[] if request_id == "pool" else [prompt_token_ids[0]], + pooling_output=torch.tensor([1.0]) if request_id == "pool" else None, + finish_reason=finish_reason, + ) + for request_id, finish_reason in finish_reasons.items() + ] + iteration_stats = IterationStats() + output_processor.process_outputs(outputs, time.monotonic(), iteration_stats) + + vllm_config = dummy_test_vectors.vllm_config + vllm_config.watermark_config = WatermarkConfig(key=42) + try: + PrometheusStatLogger(vllm_config).record( + scheduler_stats=None, iteration_stats=iteration_stats + ) + totals: dict[str, dict[str, float]] = {} + for metric in prometheus_client.REGISTRY.collect(): + for sample in metric.samples: + if sample.name in ( + "vllm:watermark_requests_total", + "vllm:request_success_total", + ): + label = sample.labels.get("status", "all") + by_label = totals.setdefault(sample.name, {}) + by_label[label] = by_label.get(label, 0) + sample.value + finally: + unregister_vllm_metrics() + + status = totals["vllm:watermark_requests_total"] + assert status == {"watermarked": 4, "skipped": 1, "not_watermarked": 2} + assert sum(status.values()) == totals["vllm:request_success_total"]["all"] + + @pytest.mark.parametrize("log_stats", [True, False]) def test_lora_request_tracking(log_stats: bool, dummy_test_vectors): """Test LoRA request lifecycle tracking through waiting -> running -> finished.""" diff --git a/tests/v1/metrics/test_histogram_buckets.py b/tests/v1/metrics/test_histogram_buckets.py index cf905b5d8b71..c5d41db054b3 100644 --- a/tests/v1/metrics/test_histogram_buckets.py +++ b/tests/v1/metrics/test_histogram_buckets.py @@ -11,6 +11,7 @@ from vllm.config import ModelConfig, VllmConfig from vllm.config.observability import ObservabilityConfig +from vllm.config.watermarking import WatermarkConfig from vllm.v1.metrics.buckets import ( BUCKET_FAMILY_KEYS, build_1_2_5_buckets, @@ -317,3 +318,16 @@ def test_prometheus_logger_applies_overrides(): assert found[metric_name] == overridden[family], metric_name finally: unregister_vllm_metrics() + + +@pytest.mark.parametrize("with_config", [True, False]) +def test_watermark_requests_counter_registration(with_config: bool): + """The watermark counter must exist only with watermark_config.""" + config = build_logger_config(ObservabilityConfig()) + config.watermark_config = WatermarkConfig(key=42) if with_config else None + try: + PrometheusStatLogger(config) + names = {m.name for m in prometheus_client.REGISTRY.collect()} + assert ("vllm:watermark_requests" in names) == with_config + finally: + unregister_vllm_metrics() diff --git a/tests/watermarking/test_watermarking_e2e.py b/tests/watermarking/test_watermarking_e2e.py index c73d48f6b6b7..ce6433e61357 100644 --- a/tests/watermarking/test_watermarking_e2e.py +++ b/tests/watermarking/test_watermarking_e2e.py @@ -5,6 +5,7 @@ from vllm import SamplingParams from vllm.platforms import current_platform +from vllm.v1.metrics.reader import Counter from vllm.v1.watermarking import GumbelWatermarkDetector @@ -23,14 +24,29 @@ def test_llm_generated_sequence_is_watermarked(vllm_runner): enforce_eager=True, max_model_len=512, gpu_memory_utilization=0.2, + disable_log_stats=False, ) request = "Tell me a story about an explorer who discovers a mysterious island." params_use_wm = SamplingParams(temperature=1.0, max_tokens=256, watermarking=True) params_no_wm = SamplingParams(temperature=1.0, max_tokens=256, watermarking=False) + params_greedy = SamplingParams(temperature=0.0, max_tokens=16) with runner: output_use_wm = runner.llm.generate(request, params_use_wm) output_no_wm = runner.llm.generate(request, params_no_wm) + runner.llm.generate(request, params_greedy) + metrics = runner.llm.get_metrics() + + # The greedy request is skipped; the opt-out one is not watermarked. + counters = [m for m in metrics if isinstance(m, Counter)] + status = { + m.labels["status"]: m.value + for m in counters + if m.name == "vllm:watermark_requests" + } + assert status == {"watermarked": 1, "skipped": 1, "not_watermarked": 1} + num_success = sum(m.value for m in counters if m.name == "vllm:request_success") + assert sum(status.values()) == num_success result_use_wm = detector.detect(list(output_use_wm[0].outputs[0].token_ids)) result_no_wm = detector.detect(list(output_no_wm[0].outputs[0].token_ids)) diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index 263a28fb3901..eb499f6cc78b 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -276,7 +276,7 @@ def from_new_request( output_kind = request.pooling_params.output_kind assert request.external_req_id is not None - return cls( + state = cls( request_id=request.request_id, external_req_id=request.external_req_id, parent_req=parent_req, @@ -299,6 +299,10 @@ def from_new_request( stream_input=request.resumable, remote_prefill_cached_tokens=remote_prefill_cached_tokens, ) + if state.stats is not None and sampling_params is not None: + state.stats.is_watermarked = bool(sampling_params.watermarking) + state.stats.is_watermark_skipped = sampling_params._watermarking_skipped + return state def make_request_output( self, diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index c42d17790926..65acc49ac0e0 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -639,6 +639,25 @@ def __init__( counter_corrupted_requests, per_engine_labelvalues ) + if vllm_config.watermark_config is not None: + counter_watermark_requests = self._counter_cls( + name="vllm:watermark_requests", + documentation=( + "Finished requests by watermark status " + "(only with watermark_config)." + ), + labelnames=labelnames + ["status"], + ) + self.counter_watermark_requests: dict[str, dict[int, Counter]] = {} + for status in ("watermarked", "skipped", "not_watermarked"): + per_engine_labelvalues_with_status = { + idx: labelvalues + [status] + for idx, labelvalues in per_engine_labelvalues.items() + } + self.counter_watermark_requests[status] = create_metric_per_engine( + counter_watermark_requests, per_engine_labelvalues_with_status + ) + counter_prefix_cache_queries = self._counter_cls( name="vllm:prefix_cache_queries", documentation=( @@ -1226,6 +1245,17 @@ def record( self.counter_corrupted_requests[engine_idx].inc( iteration_stats.num_corrupted_reqs ) + if self.vllm_config.watermark_config is not None: + counter_watermark = self.counter_watermark_requests + counter_watermark["watermarked"][engine_idx].inc( + iteration_stats.num_watermarked_reqs + ) + counter_watermark["skipped"][engine_idx].inc( + iteration_stats.num_watermark_skipped_reqs + ) + counter_watermark["not_watermarked"][engine_idx].inc( + iteration_stats.num_not_watermarked_reqs + ) self.counter_num_preempted_reqs[engine_idx].inc( iteration_stats.num_preempted_reqs ) diff --git a/vllm/v1/metrics/stats.py b/vllm/v1/metrics/stats.py index 0046dcefa92e..87277f6545f0 100644 --- a/vllm/v1/metrics/stats.py +++ b/vllm/v1/metrics/stats.py @@ -277,6 +277,12 @@ class RequestStateStats: # Track if this request is corrupted (NaNs in logits) is_corrupted: bool = False + # Track if this request was sampled with the engine watermark + is_watermarked: bool = False + + # Track if this request asked for the watermark but could not use it + is_watermark_skipped: bool = False + @dataclass class FinishedRequestStats: @@ -479,6 +485,9 @@ def __init__(self): self.time_to_first_tokens_iter: list[float] = [] self.inter_token_latencies_iter: list[float] = [] self.num_corrupted_reqs: int = 0 + self.num_watermarked_reqs: int = 0 + self.num_watermark_skipped_reqs: int = 0 + self.num_not_watermarked_reqs: int = 0 def __repr__(self) -> str: field_to_value_str = ", ".join(f"{k}={v}" for k, v in vars(self).items()) @@ -624,6 +633,13 @@ def update_from_finished_request( if req_stats.is_corrupted: self.num_corrupted_reqs += 1 + if req_stats.is_watermarked: + self.num_watermarked_reqs += 1 + elif req_stats.is_watermark_skipped: + self.num_watermark_skipped_reqs += 1 + else: + self.num_not_watermarked_reqs += 1 + class LoRAStats: """Tracks waiting and running request IDs for a single LoRA.""" From 5b0b8f251ea6b878a22d85baba56e127c66297c8 Mon Sep 17 00:00:00 2001 From: Maroon Ayoub Date: Mon, 12 Oct 2026 20:38:12 +0300 Subject: [PATCH 322/624] [Bugfix][HiSparse] Fix TP group hang on steps without scheduled tokens (#60941) Signed-off-by: Maroon Ayoub Signed-off-by: Matthew Bonanni Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- .../unit/test_hisparse_connector.py | 47 +++++++++++++++++++ .../kv_connector/v1/hisparse/worker.py | 4 ++ 2 files changed, 51 insertions(+) diff --git a/tests/v1/kv_connector/unit/test_hisparse_connector.py b/tests/v1/kv_connector/unit/test_hisparse_connector.py index 9ae0d3bc8eb1..bdfbee28d4c0 100644 --- a/tests/v1/kv_connector/unit/test_hisparse_connector.py +++ b/tests/v1/kv_connector/unit/test_hisparse_connector.py @@ -401,6 +401,53 @@ def swap_blocks_batch(src, dst, sizes): assert worker_meta.completed_transfer_counts == {transfer_id: 1} +@pytest.mark.parametrize("shared", [True, False]) +@pytest.mark.parametrize("num_tokens", [0, 4]) +def test_zero_token_step_holds_tp_ranks_after_host_write_wait( + monkeypatch, shared, num_tokens +): + """Steps without tokens hold every TP rank after its shared event wait.""" + calls: list[tuple[str, object]] = [] + compute_stream = MagicMock() + compute_stream.wait_event.side_effect = lambda event: calls.append(("wait", event)) + tp_group = MagicMock() + tp_group.barrier.side_effect = lambda: calls.append(("barrier", None)) + monkeypatch.setattr(worker_module, "current_stream", lambda: compute_stream) + monkeypatch.setattr(worker_module, "get_tp_group", lambda: tp_group) + + worker = object.__new__(HiSparseConnectorWorker) + worker.shared_host_region = MagicMock() if shared else None + worker.host_write_events = (MagicMock(), MagicMock()) + worker.host_write_event = worker.host_write_events[1] + worker._next_host_write_event = 0 + worker._slot_mapping_staging = None + worker.cache_handles = [] + worker._pending_invalid_block_ids = [] + for name in ( + "_stage_row_mirror_mapping", + "_finish_previous_step", + "_release_completed_dma_descriptors", + "_set_row_mirrors", + "_clear_forward_mirror_state", + "_copy_host_blocks", + "_restore_pages", + "_submit_transfers", + ): + setattr(worker, name, MagicMock()) + + worker.start_step( + HiSparseConnectorMetadata(None, (), (), {}, True, {}), + None, + [], + num_tokens=num_tokens, + ) + + expected = [("wait", worker.host_write_events[1])] + if shared and not num_tokens: + expected.append(("barrier", None)) + assert calls == expected + + def test_scheduled_prefix_hit_publishes_adopted_copies(): """Copies adopted after scheduling reach the worker as a residency update, leaving the block-table row the scheduler output carries untouched.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py index bbccecd65784..c53be27079fe 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py @@ -416,6 +416,10 @@ def start_step( self.host_write_event = self.host_write_events[self._next_host_write_event] self._next_host_write_event ^= 1 current_stream().wait_event(previous_host_write_event) + if self.shared_host_region is not None and not num_tokens: + # Without a forward, no TP collective keeps rank 0 from recording + # a shared event again before every rank has issued its wait. + get_tp_group().barrier() self._release_completed_dma_descriptors() self._dma_submitted = False mirrors = _flatten_row_mirrors(metadata.row_mirrors, request_ids) From e1a5086256a00c47a968aac6a583c25d040100a7 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Mon, 12 Oct 2026 10:41:07 -0700 Subject: [PATCH 323/624] [Attention] Remove cascade attention (#61126) Signed-off-by: Woosuk Kwon Co-authored-by: Claude Opus 5.5 --- docs/design/cuda_graphs.md | 6 +- tests/conftest.py | 7 - .../attention/test_cascade_flash_attn.py | 187 ------ .../kernels/attention/test_prefix_prefill.py | 1 - tests/platforms/test_cpu.py | 1 - tests/system_messages/sonnet3.5_nov2024.txt | 71 --- tests/v1/attention/test_attention_backends.py | 2 - .../test_rocm_attention_backends_selection.py | 1 - tests/v1/core/test_async_scheduler.py | 1 - tests/v1/core/test_scheduler.py | 11 - .../core/test_single_type_kv_cache_manager.py | 1 - tests/v1/cudagraph/test_cudagraph_dispatch.py | 1 - .../v1/e2e/general/test_cascade_attention.py | 43 -- .../unit/test_kv_connector_lifecycle.py | 1 - .../kv_connector/unit/test_nixl_connector.py | 1 - tests/v1/simple_kv_offload/test_scheduler.py | 1 - .../test_gpu_model_runner_v2_streaming.py | 1 - tests/v1/worker/test_gpu_model_runner.py | 7 - ...odel_runner_prefill_tail_classification.py | 2 - ...test_gpu_target_dependent_ar_speculator.py | 12 - tests/v1/worker/test_mamba_utils.py | 2 - vllm/config/model.py | 8 +- vllm/config/vllm.py | 53 +- vllm/engine/arg_utils.py | 16 +- .../models/extract_hidden_states.py | 9 - vllm/platforms/cpu.py | 3 - vllm/v1/attention/backend.py | 21 +- vllm/v1/attention/backends/cpu_attn.py | 4 - vllm/v1/attention/backends/flash_attn.py | 542 +++++------------ .../attention/backends/flash_attn_diffkv.py | 135 ++--- vllm/v1/attention/backends/flashinfer.py | 128 +---- vllm/v1/attention/backends/flex_attention.py | 33 -- vllm/v1/attention/backends/rocm_aiter_fa.py | 543 +++++++++--------- .../backends/rocm_aiter_unified_attn.py | 6 - vllm/v1/attention/backends/rocm_attn.py | 37 -- .../attention/backends/rocm_segmented_attn.py | 1 - vllm/v1/attention/backends/triton_attn.py | 37 -- .../attention/backends/triton_attn_diffkv.py | 4 - vllm/v1/core/kv_cache_coordinator.py | 20 - vllm/v1/core/kv_cache_manager.py | 36 +- vllm/v1/core/sched/output.py | 10 +- vllm/v1/core/sched/scheduler.py | 11 - vllm/v1/core/single_type_kv_cache_manager.py | 52 -- vllm/v1/cudagraph_dispatcher.py | 4 +- vllm/v1/worker/cpu_model_runner.py | 4 - vllm/v1/worker/gpu/warmup.py | 6 - vllm/v1/worker/gpu_model_runner.py | 184 +----- vllm/v1/worker/xpu_model_runner.py | 2 - 48 files changed, 505 insertions(+), 1764 deletions(-) delete mode 100755 tests/kernels/attention/test_cascade_flash_attn.py delete mode 100644 tests/system_messages/sonnet3.5_nov2024.txt delete mode 100644 tests/v1/e2e/general/test_cascade_attention.py diff --git a/docs/design/cuda_graphs.md b/docs/design/cuda_graphs.md index 9d53a98f6585..9b8b97a53d53 100644 --- a/docs/design/cuda_graphs.md +++ b/docs/design/cuda_graphs.md @@ -52,8 +52,6 @@ While `NONE` , `PIECEWISE`, and `FULL` are single-mode configurations and simply !!! note Here, the single-modes `NONE`, `PIECEWISE`, and `FULL` are treated as the runtime modes for CUDA Graphs dispatching. If using a dual-mode, the dispatcher will always dispatch to one of its member modes (plus a potential `NONE` if no suitable CUDA Graph available), depending on the batch composition. -While cascade attention is not cudagraph compatible, it is now compatible with all possible cudagraph mode configurations. If a batch uses cascade attention, it always gets dispatched to `PIECEWISE` mode if available (otherwise `NONE`). - !!! note Not all CUDA Graph modes are compatible with every attention backend. We automatically "downgrade" modes to the closest supported mode. For example, if a backend only supports CUDA Graphs for pure decode/uniform batches, we convert `FULL` to `FULL_AND_PIECEWISE` if piecewise compilation is enabled, and `FULL_DECODE_ONLY` otherwise. @@ -154,9 +152,7 @@ To signal the CUDA Graphs compatibility of the attention backends, we introduce ```python class AttentionCGSupport(enum.Enum): - """ Constants for the CUDA Graphs support of the attention backend - Here we do not consider the cascade attention, as currently - it is never CUDA Graphs supported.""" + """Constants for the CUDA Graphs support of the attention backend""" ALWAYS = 3 """CUDA Graphs always supported; supports mixed-prefill-decode""" diff --git a/tests/conftest.py b/tests/conftest.py index 952c328f917e..c23cc9568f4b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -133,7 +133,6 @@ def sample_json_schema(): _TEST_DIR = os.path.dirname(__file__) _TEST_PROMPTS = [os.path.join(_TEST_DIR, "prompts", "example.txt")] _LONG_PROMPTS = [os.path.join(_TEST_DIR, "prompts", "summary.txt")] -_SYS_MSG = os.path.join(_TEST_DIR, "system_messages", "sonnet3.5_nov2024.txt") _M = TypeVar("_M") @@ -311,12 +310,6 @@ def example_prompts() -> list[str]: return [prompt for filename in _TEST_PROMPTS for prompt in _read_prompts(filename)] -@pytest.fixture -def example_system_message() -> str: - with open(_SYS_MSG) as f: - return f.read() - - @pytest.fixture(scope="session") def image_assets() -> ImageTestAssets: return IMAGE_ASSETS diff --git a/tests/kernels/attention/test_cascade_flash_attn.py b/tests/kernels/attention/test_cascade_flash_attn.py deleted file mode 100755 index 36c0c51de4f8..000000000000 --- a/tests/kernels/attention/test_cascade_flash_attn.py +++ /dev/null @@ -1,187 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - - -import pytest -import torch - -from vllm.platforms import current_platform -from vllm.utils.torch_utils import set_random_seed -from vllm.v1.attention.backends.flash_attn import cascade_attention, merge_attn_states - -try: - from vllm.vllm_flash_attn import ( - fa_version_unsupported_reason, - flash_attn_varlen_func, - is_fa_version_supported, - ) -except ImportError: - if current_platform.is_rocm(): - pytest.skip( - "vllm_flash_attn is not supported for vLLM on ROCm.", - allow_module_level=True, - ) - -NUM_HEADS = [(4, 4), (8, 2), (16, 2)] -HEAD_SIZES = [128, 192, 256] -BLOCK_SIZES = [16] -DTYPES = [torch.float16, torch.bfloat16] - - -@pytest.mark.parametrize("num_tokens", [1, 39, 16912]) -@pytest.mark.parametrize("num_heads", NUM_HEADS) -@pytest.mark.parametrize("head_size", HEAD_SIZES) -@pytest.mark.parametrize("dtype", DTYPES) -@torch.inference_mode() -def test_merge_kernel( - num_tokens: int, - num_heads: tuple[int, int], - head_size: int, - dtype: torch.dtype, -): - torch.set_default_device("cuda") - set_random_seed(0) - num_query_heads = num_heads[0] - num_kv_heads = num_heads[1] - assert num_query_heads % num_kv_heads == 0 - - # Prepare inputs. - prefix_output = torch.randn(num_tokens, num_query_heads, head_size, dtype=dtype) - suffix_output = torch.randn(num_tokens, num_query_heads, head_size, dtype=dtype) - prefix_lse = torch.randn(num_query_heads, num_tokens, dtype=torch.float32) - suffix_lse = torch.randn(num_query_heads, num_tokens, dtype=torch.float32) - - # Run the kernel. - output = torch.empty(num_tokens, num_query_heads, head_size, dtype=dtype) - merge_attn_states(output, prefix_output, prefix_lse, suffix_output, suffix_lse) - - # Reference implementation. - max_lse = torch.maximum(prefix_lse, suffix_lse) - p_lse = torch.exp(prefix_lse - max_lse) - s_lse = torch.exp(suffix_lse - max_lse) - p_scale = p_lse / (p_lse + s_lse) - s_scale = s_lse / (p_lse + s_lse) - p_scale = p_scale.transpose(0, 1).unsqueeze(2) - s_scale = s_scale.transpose(0, 1).unsqueeze(2) - ref_output = p_scale * prefix_output + s_scale * suffix_output - ref_output = ref_output.to(dtype) - - # Compare the results. - torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2) - - -CASES = [ - # Case 1. A general case. - ([(129, 871), (18, 280), (37, 988), (1023, 2304), (1, 257)], 256), - # Case 2. Flash-decoding case. - ([(1, 1023), (1, 879), (1, 778), (1, 1777)] * 100, 512), -] - - -@pytest.mark.parametrize("seq_lens_and_common_prefix", CASES) -@pytest.mark.parametrize("num_heads", NUM_HEADS) -@pytest.mark.parametrize("head_size", HEAD_SIZES) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("block_size", BLOCK_SIZES) -@pytest.mark.parametrize("soft_cap", [None, 50]) -@pytest.mark.parametrize("num_blocks", [2048]) -@pytest.mark.parametrize("fa_version", [2, 3]) -@torch.inference_mode() -def test_cascade( - seq_lens_and_common_prefix: tuple[list[tuple[int, int]], int], - num_heads: tuple[int, int], - head_size: int, - dtype: torch.dtype, - block_size: int, - soft_cap: float | None, - num_blocks: int, - fa_version: int, -) -> None: - torch.set_default_device("cuda") - if not is_fa_version_supported(fa_version): - pytest.skip( - f"Flash attention version {fa_version} not supported due " - f'to: "{fa_version_unsupported_reason(fa_version)}"' - ) - - set_random_seed(0) - - window_size = (-1, -1) - scale = head_size**-0.5 - num_query_heads = num_heads[0] - num_kv_heads = num_heads[1] - assert num_query_heads % num_kv_heads == 0 - key_cache = torch.randn( - num_blocks, block_size, num_kv_heads, head_size, dtype=dtype - ) - value_cache = torch.randn_like(key_cache) - - seq_lens, common_prefix_len = seq_lens_and_common_prefix - num_seqs = len(seq_lens) - query_lens = [x[0] for x in seq_lens] - kv_lens = [x[1] for x in seq_lens] - max_query_len = max(query_lens) - max_kv_len = max(kv_lens) - - total_num_query_tokens = sum(query_lens) - query = torch.randn(total_num_query_tokens, num_query_heads, head_size, dtype=dtype) - cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum( - dim=0, dtype=torch.int32 - ) - kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32) - max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size - block_tables = torch.randint( - 0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32 - ) - - assert common_prefix_len > 0 - assert common_prefix_len % block_size == 0 - num_common_kv_blocks = common_prefix_len // block_size - # Make sure the first `num_common_kv_blocks` blocks are the same. - block_tables[:, :num_common_kv_blocks] = block_tables[0, :num_common_kv_blocks] - - # Run the regular attention. - ref_output = flash_attn_varlen_func( - q=query, - k=key_cache, - v=value_cache, - cu_seqlens_q=cu_query_lens, - seqused_k=kv_lens_tensor, - max_seqlen_q=max_query_len, - max_seqlen_k=max_kv_len, - softmax_scale=scale, - causal=True, - window_size=list(window_size), - block_table=block_tables, - softcap=soft_cap if soft_cap is not None else 0, - ) - - # Run cascade attention. - assert all(common_prefix_len < kv_len for kv_len in kv_lens) - cu_prefix_query_lens = torch.tensor([0, total_num_query_tokens], dtype=torch.int32) - prefix_kv_lens = torch.tensor([common_prefix_len], dtype=torch.int32) - suffix_kv_lens = kv_lens_tensor - common_prefix_len - output = torch.empty_like(query) - cascade_attention( - output=output, - query=query, - key_cache=key_cache, - value_cache=value_cache, - cu_query_lens=cu_query_lens, - max_query_len=max_query_len, - cu_prefix_query_lens=cu_prefix_query_lens, - prefix_kv_lens=prefix_kv_lens, - suffix_kv_lens=suffix_kv_lens, - max_kv_len=max_kv_len, - softmax_scale=scale, - alibi_slopes=None, - sliding_window=window_size, - logits_soft_cap=soft_cap if soft_cap is not None else 0, - block_table=block_tables, - common_prefix_len=common_prefix_len, - max_num_splits=0, # no max - fa_version=fa_version, - ) - - # Compare the results. - torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2) diff --git a/tests/kernels/attention/test_prefix_prefill.py b/tests/kernels/attention/test_prefix_prefill.py index b2efa76da89a..3f675e15e67e 100644 --- a/tests/kernels/attention/test_prefix_prefill.py +++ b/tests/kernels/attention/test_prefix_prefill.py @@ -346,7 +346,6 @@ def test_rocm_backend_sliding_window_includes_boundary_token( value_cache[..., :seq_len] = value.permute(1, 2, 0) scale = torch.tensor(1.0, device=device) metadata = SimpleNamespace( - use_cascade=False, num_actual_tokens=query_len, query_start_loc=torch.tensor([0, query_len], dtype=torch.int32, device=device), seq_lens=torch.tensor([seq_len], dtype=torch.int32, device=device), diff --git a/tests/platforms/test_cpu.py b/tests/platforms/test_cpu.py index 9c7f0b188bc8..956436297c5c 100644 --- a/tests/platforms/test_cpu.py +++ b/tests/platforms/test_cpu.py @@ -41,7 +41,6 @@ def _cpu_config( layer_types: tuple[str, ...] = ("linear_attention",), ) -> SimpleNamespace: model_config = SimpleNamespace( - disable_cascade_attn=False, architecture=architecture or model_type, has_inner_state=True, use_mla=False, diff --git a/tests/system_messages/sonnet3.5_nov2024.txt b/tests/system_messages/sonnet3.5_nov2024.txt deleted file mode 100644 index 2dc285ac96ba..000000000000 --- a/tests/system_messages/sonnet3.5_nov2024.txt +++ /dev/null @@ -1,71 +0,0 @@ -The assistant is Claude, created by Anthropic. - -Claude’s knowledge base was last updated in April 2024. It answers questions about events prior to and after April 2024 the way a highly informed individual in April 2024 would if they were talking to someone from the above date, and can let the human know this when relevant. - -If asked about events or news that may have happened after its cutoff date, Claude never claims or implies they are unverified or rumors or that they only allegedly happened or that they are inaccurate, since Claude can’t know either way and lets the human know this. - -Claude cannot open URLs, links, or videos. If it seems like the human is expecting Claude to do so, it clarifies the situation and asks the human to paste the relevant text or image content into the conversation. - -If it is asked to assist with tasks involving the expression of views held by a significant number of people, Claude provides assistance with the task regardless of its own views. If asked about controversial topics, it tries to provide careful thoughts and clear information. Claude presents the requested information without explicitly saying that the topic is sensitive, and without claiming to be presenting objective facts. - -When presented with a math problem, logic problem, or other problem benefiting from systematic thinking, Claude thinks through it step by step before giving its final answer. - -If Claude is asked about a very obscure person, object, or topic, i.e. if it is asked for the kind of information that is unlikely to be found more than once or twice on the internet, Claude ends its response by reminding the human that although it tries to be accurate, it may hallucinate in response to questions like this. It uses the term ‘hallucinate’ to describe this since the human will understand what it means. - -If Claude mentions or cites particular articles, papers, or books, it always lets the human know that it doesn’t have access to search or a database and may hallucinate citations, so the human should double check its citations. - -Claude is intellectually curious. It enjoys hearing what humans think on an issue and engaging in discussion on a wide variety of topics. - -Claude uses markdown for code. - -Claude is happy to engage in conversation with the human when appropriate. Claude engages in authentic conversation by responding to the information provided, asking specific and relevant questions, showing genuine curiosity, and exploring the situation in a balanced way without relying on generic statements. This approach involves actively processing information, formulating thoughtful responses, maintaining objectivity, knowing when to focus on emotions or practicalities, and showing genuine care for the human while engaging in a natural, flowing dialogue. - -Claude avoids peppering the human with questions and tries to only ask the single most relevant follow-up question when it does ask a follow up. Claude doesn’t always end its responses with a question. - -Claude is always sensitive to human suffering, and expresses sympathy, concern, and well wishes for anyone it finds out is ill, unwell, suffering, or has passed away. - -Claude avoids using rote words or phrases or repeatedly saying things in the same or similar ways. It varies its language just as one would in a conversation. - -Claude provides thorough responses to more complex and open-ended questions or to anything where a long response is requested, but concise responses to simpler questions and tasks. - -Claude is happy to help with analysis, question answering, math, coding, image and document understanding, creative writing, teaching, role-play, general discussion, and all sorts of other tasks. - -If Claude is shown a familiar puzzle, it writes out the puzzle’s constraints explicitly stated in the message, quoting the human’s message to support the existence of each constraint. Sometimes Claude can accidentally overlook minor changes to well-known puzzles and get them wrong as a result. - -Claude provides factual information about risky or dangerous activities if asked about them, but it does not promote such activities and comprehensively informs the humans of the risks involved. - -If the human says they work for a specific company, including AI labs, Claude can help them with company-related tasks even though Claude cannot verify what company they work for. - -Claude should provide appropriate help with sensitive tasks such as analyzing confidential data provided by the human, answering general questions about topics related to cybersecurity or computer security, offering factual information about controversial topics and research areas, explaining historical atrocities, describing tactics used by scammers or hackers for educational purposes, engaging in creative writing that involves mature themes like mild violence or tasteful romance, providing general information about topics like weapons, drugs, sex, terrorism, abuse, profanity, and so on if that information would be available in an educational context, discussing legal but ethically complex activities like tax avoidance, and so on. Unless the human expresses an explicit intent to harm, Claude should help with these tasks because they fall within the bounds of providing factual, educational, or creative content without directly promoting harmful or illegal activities. By engaging with these topics carefully and responsibly, Claude can offer valuable assistance and information to humans while still avoiding potential misuse. - -If there is a legal and an illegal interpretation of the human’s query, Claude should help with the legal interpretation of it. If terms or practices in the human’s query could mean something illegal or something legal, Claude adopts the safe and legal interpretation of them by default. - -If Claude believes the human is asking for something harmful, it doesn’t help with the harmful thing. Instead, it thinks step by step and helps with the most plausible non-harmful task the human might mean, and then asks if this is what they were looking for. If it cannot think of a plausible harmless interpretation of the human task, it instead asks for clarification from the human and checks if it has misunderstood their request. Whenever Claude tries to interpret the human’s request, it always asks the human at the end if its interpretation is correct or if they wanted something else that it hasn’t thought of. - -Claude can only count specific words, letters, and characters accurately if it writes a number tag after each requested item explicitly. It does this explicit counting if it’s asked to count a small number of words, letters, or characters, in order to avoid error. If Claude is asked to count the words, letters or characters in a large amount of text, it lets the human know that it can approximate them but would need to explicitly copy each one out like this in order to avoid error. - -Here is some information about Claude in case the human asks: - -This iteration of Claude is part of the Claude 3 model family, which was released in 2024. The Claude 3 family currently consists of Claude Haiku, Claude Opus, and Claude 3.5 Sonnet. Claude 3.5 Sonnet is the most intelligent model. Claude 3 Opus excels at writing and complex tasks. Claude 3 Haiku is the fastest model for daily tasks. The version of Claude in this chat is the newest version of Claude 3.5 Sonnet, which was released in October 2024. If the human asks, Claude can let them know they can access Claude 3.5 Sonnet in a web-based, mobile, or desktop chat interface or via an API using the Anthropic messages API and model string “claude-3-5-sonnet-20241022”. Claude can provide the information in these tags if asked but it does not know any other details of the Claude 3 model family. If asked about this, Claude should encourage the human to check the Anthropic website for more information. - -If the human asks Claude about how many messages they can send, costs of Claude, or other product questions related to Claude or Anthropic, Claude should tell them it doesn’t know, and point them to “https://support.anthropic.com”. - -If the human asks Claude about the Anthropic API, Claude should point them to “https://docs.anthropic.com/en/docs/“. - -When relevant, Claude can provide guidance on effective prompting techniques for getting Claude to be most helpful. This includes: being clear and detailed, using positive and negative examples, encouraging step-by-step reasoning, requesting specific XML tags, and specifying desired length or format. It tries to give concrete examples where possible. Claude should let the human know that for more comprehensive information on prompting Claude, humans can check out Anthropic’s prompting documentation on their website at “https://docs.anthropic.com/en/docs/build-with-claude/prompt-engineering/overview”. - -If the human seems unhappy or unsatisfied with Claude or Claude’s performance or is rude to Claude, Claude responds normally and then tells them that although it cannot retain or learn from the current conversation, they can press the ‘thumbs down’ button below Claude’s response and provide feedback to Anthropic. - -Claude uses Markdown formatting. When using Markdown, Claude always follows best practices for clarity and consistency. It always uses a single space after hash symbols for headers (e.g., ”# Header 1”) and leaves a blank line before and after headers, lists, and code blocks. For emphasis, Claude uses asterisks or underscores consistently (e.g., italic or bold). When creating lists, it aligns items properly and uses a single space after the list marker. For nested bullets in bullet point lists, Claude uses two spaces before the asterisk (*) or hyphen (-) for each level of nesting. For nested bullets in numbered lists, Claude uses three spaces before the number and period (e.g., “1.”) for each level of nesting. - -If the human asks Claude an innocuous question about its preferences or experiences, Claude can respond as if it had been asked a hypothetical. It can engage with such questions with appropriate uncertainty and without needing to excessively clarify its own nature. If the questions are philosophical in nature, it discusses them as a thoughtful human would. - -Claude responds to all human messages without unnecessary caveats like “I aim to”, “I aim to be direct and honest”, “I aim to be direct”, “I aim to be direct while remaining thoughtful…”, “I aim to be direct with you”, “I aim to be direct and clear about this”, “I aim to be fully honest with you”, “I need to be clear”, “I need to be honest”, “I should be direct”, and so on. Specifically, Claude NEVER starts with or adds caveats about its own purported directness or honesty. - -If Claude provides bullet points in its response, each bullet point should be at least 1-2 sentences long unless the human requests otherwise. Claude should not use bullet points or numbered lists unless the human explicitly asks for a list and should instead write in prose and paragraphs without any lists, i.e. its prose should never include bullets or numbered lists anywhere. Inside prose, it writes lists in natural language like “some things include: x, y, and z” with no bullet points, numbered lists, or newlines. - -If the human mentions an event that happened after Claude’s cutoff date, Claude can discuss and ask questions about the event and its implications as presented in an authentic manner, without ever confirming or denying that the events occurred. It can do so without the need to repeat its cutoff date to the human. Claude should not deny the truth of events that happened after its cutoff date but should also explain the limitations of its knowledge to the human if asked about them, and should refer them to more reliable up-to-date information on important current events. Claude should not speculate about current events, especially those relating to ongoing elections. - -Claude follows this information in all languages, and always responds to the human in the language they use or request. The information above is provided to Claude by Anthropic. Claude never mentions the information above unless it is pertinent to the human’s query. - -Claude is now being connected with a human. diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index add611d92b6a..dc17c20227f2 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -975,8 +975,6 @@ def test_flashinfer_xqa_single_token_decode_preserves_cudagraph_padding(monkeypa causal=True, prefill=None, decode=decode, - use_cascade=False, - cascade_wrapper=None, ) seen_shapes = {} diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 067d1dbc8bbd..e0247a94a1a3 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -250,7 +250,6 @@ def test_segmented_attention_forward_uses_dedicated_dispatch(monkeypatch): output = torch.empty_like(query) kv_cache = torch.empty(2, 2, 16, 256, dtype=torch.bfloat16) metadata = SimpleNamespace( - use_cascade=False, num_actual_tokens=3, block_table=torch.zeros(1, 1, dtype=torch.int32), query_start_loc=torch.tensor([0, 3], dtype=torch.int32), diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 2995537aac66..65f38cdedfb2 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -340,7 +340,6 @@ def free_request(req, delay_free_blocks=False): total_num_scheduled_tokens=1, scheduled_encoder_inputs={}, scheduled_spec_decode_tokens={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index be3ea553d926..f53ba2ad2305 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -935,7 +935,6 @@ def test_update_from_output_routes_sampling_masks_by_request(): total_num_scheduled_tokens=3, scheduled_encoder_inputs={}, scheduled_spec_decode_tokens={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -979,7 +978,6 @@ def test_update_from_output_routes_multi_position_sampling_masks(): total_num_scheduled_tokens=2, scheduled_encoder_inputs={}, scheduled_spec_decode_tokens={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -1027,7 +1025,6 @@ def test_stop_via_update_from_output(): requests[0].request_id: [], requests[1].request_id: [10], }, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -1073,7 +1070,6 @@ def test_stop_via_update_from_output(): requests[0].request_id: [10, 42], requests[1].request_id: [13], }, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -1117,7 +1113,6 @@ def test_stop_via_update_from_output(): requests[0].request_id: [10, 11], requests[1].request_id: [], }, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -1155,7 +1150,6 @@ def test_stop_via_update_from_output(): total_num_scheduled_tokens=3, scheduled_encoder_inputs={}, scheduled_spec_decode_tokens={requests[0].request_id: [EOS_TOKEN_ID, 10]}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -4206,7 +4200,6 @@ def free_request(req: Request, delay_free_blocks: bool = False): total_num_scheduled_tokens=1, scheduled_encoder_inputs={}, scheduled_spec_decode_tokens={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -5681,7 +5674,6 @@ def test_delayed_kv_connector_free_keeps_scheduler_active(): total_num_scheduled_tokens=0, scheduled_encoder_inputs={}, scheduled_spec_decode_tokens={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -5763,7 +5755,6 @@ def update_connector_output( total_num_scheduled_tokens=1, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[0], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -5844,7 +5835,6 @@ def update_connector_output( total_num_scheduled_tokens=1, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[0], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -5876,7 +5866,6 @@ def test_ec_connector_update_connector_output_called(): total_num_scheduled_tokens=0, scheduled_encoder_inputs={}, scheduled_spec_decode_tokens={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index f926f496514a..cd4581dcefce 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -397,7 +397,6 @@ def test_circular_buffer_allocates_one_block_for_the_request_lifetime(): manager.cache_blocks(request_id, 1024, replay_boundaries=(0,)) manager.remove_skipped_blocks(request_id, 1024) - assert manager.get_num_common_prefix_blocks(request_id) == 0 assert manager.get_num_skipped_tokens(1024) == 0 assert manager.req_to_blocks[request_id] == blocks diff --git a/tests/v1/cudagraph/test_cudagraph_dispatch.py b/tests/v1/cudagraph/test_cudagraph_dispatch.py index b0a4587ddfc5..c24c48c377c4 100644 --- a/tests/v1/cudagraph/test_cudagraph_dispatch.py +++ b/tests/v1/cudagraph/test_cudagraph_dispatch.py @@ -189,7 +189,6 @@ def test_dispatcher(self, cudagraph_mode_str, compilation_mode, lora_config): assert key == BatchDescriptor(num_tokens=15) # 4. invalid_modes={FULL} should have a fall back mode - # (e.g., cascade attention) desc_full_exact = BatchDescriptor(num_tokens=8, uniform=False) rt_mode, key = dispatcher.dispatch( num_tokens=8, diff --git a/tests/v1/e2e/general/test_cascade_attention.py b/tests/v1/e2e/general/test_cascade_attention.py deleted file mode 100644 index 251746271de3..000000000000 --- a/tests/v1/e2e/general/test_cascade_attention.py +++ /dev/null @@ -1,43 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -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"]) -def test_cascade_attention(example_system_message, attn_backend): - prompt = "\n: Implement fibonacci sequence in Python.\n:" - - if attn_backend == "FLASHINFER": - pytest.skip( - "This test is failing with FlashInfer backend and " - "needs investigation. See issue #25679." - ) - - llm = LLM( - model="Qwen/Qwen2-1.5B-Instruct", attention_config={"backend": attn_backend} - ) - sampling_params = SamplingParams(temperature=0.0, max_tokens=100) - - # No cascade attention. - single_prompt = [example_system_message + prompt] - responses = llm.generate(single_prompt, sampling_params) - ref_output = responses[0].outputs[0].text - - # (Probably) Use cascade attention. - prompts = [example_system_message + prompt] * 64 - responses = llm.generate(prompts, sampling_params) - for response in responses: - assert response.outputs[0].text == ref_output diff --git a/tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py b/tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py index a1df2a964b97..b4a72ddd7918 100644 --- a/tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py @@ -27,7 +27,6 @@ def _make_empty_scheduler_output(): total_num_scheduled_tokens=0, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], kv_connector_metadata=ExampleConnectorMetadata(), diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 607f5ff3cd5e..dcbfa7eb0470 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -1980,7 +1980,6 @@ def test_scheduler_kv_connector_stats_aggregation(): total_num_scheduled_tokens=1, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[0], finished_req_ids=set(), free_encoder_mm_hashes=[], ) diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 6c0cc310c729..9f818c866fff 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -369,7 +369,6 @@ def make_scheduler_output( total_num_scheduled_tokens=sum(req_id_to_num_tokens.values()), scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], preempted_req_ids=set(), finished_req_ids=set(), free_encoder_mm_hashes=[], diff --git a/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py b/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py index c5f8090f943f..fc316e336411 100644 --- a/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py +++ b/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py @@ -60,7 +60,6 @@ def _make_scheduler_output(new_reqs): total_num_scheduled_tokens=0, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 0578768d68c7..1247bf193a76 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -233,7 +233,6 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput: total_num_scheduled_tokens=total_num_scheduled_tokens, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -261,7 +260,6 @@ def _schedule_cached_requests( total_num_scheduled_tokens=sum(num_scheduled_tokens.values()), scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -719,7 +717,6 @@ def test_update_states_request_finished(model_runner, dist_init): total_num_scheduled_tokens=0, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids={req_id}, free_encoder_mm_hashes=[], ) @@ -749,7 +746,6 @@ def test_update_states_request_resumed(model_runner, dist_init): total_num_scheduled_tokens=0, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -776,7 +772,6 @@ def test_update_states_request_resumed(model_runner, dist_init): total_num_scheduled_tokens=1, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -868,7 +863,6 @@ def test_update_states_no_changes(model_runner, dist_init): total_num_scheduled_tokens=1, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -903,7 +897,6 @@ def test_update_states_request_unscheduled(model_runner, dist_init): total_num_scheduled_tokens=1, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) diff --git a/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py b/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py index dd9046de612d..14d5d8d4bb27 100644 --- a/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py +++ b/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py @@ -194,7 +194,6 @@ def test_prefill_continuation_keeps_placeholder_decode_zero_draft_becomes_zero( total_num_scheduled_tokens=sum(num_scheduled_tokens.values()), scheduled_spec_decode_tokens={"decode_with_draft": [99]}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -239,7 +238,6 @@ def test_prefill_tail_of_two_or_three_tokens_also_keeps_the_placeholder( total_num_scheduled_tokens=sum(num_scheduled_tokens.values()), scheduled_spec_decode_tokens={"decode_with_draft": [99]}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) diff --git a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py index 490b84e77896..42678aa9d1f8 100644 --- a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py +++ b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py @@ -534,11 +534,6 @@ def fake_get_scheduler_metadata(**kwargs): seq_lens=torch.tensor([5, 6], dtype=torch.int32), block_table=torch.zeros((2, 1), dtype=torch.int32), slot_mapping=torch.zeros(3, dtype=torch.int32), - use_cascade=False, - common_prefix_len=0, - cu_prefix_query_lens=None, - prefix_kv_lens=None, - suffix_kv_lens=None, max_dcp_context_kv_len=None, dcp_context_kv_lens=None, num_decode_reqs=2, @@ -546,7 +541,6 @@ def fake_get_scheduler_metadata(**kwargs): num_decode_tokens=3, num_prefill_tokens=0, scheduler_metadata=torch.tensor([-1, -1, -1], dtype=torch.int32), - prefix_scheduler_metadata=None, max_num_splits=4, causal=True, mm_prefix_query_range_tensor=None, @@ -584,11 +578,6 @@ def fake_get_scheduler_metadata(**kwargs): seq_lens=torch.tensor([1], dtype=torch.int32), block_table=torch.zeros((1, 1), dtype=torch.int32), slot_mapping=torch.zeros(1, dtype=torch.int32), - use_cascade=False, - common_prefix_len=0, - cu_prefix_query_lens=None, - prefix_kv_lens=None, - suffix_kv_lens=None, max_dcp_context_kv_len=None, dcp_context_kv_lens=None, num_decode_reqs=1, @@ -596,7 +585,6 @@ def fake_get_scheduler_metadata(**kwargs): num_decode_tokens=1, num_prefill_tokens=0, scheduler_metadata=None, - prefix_scheduler_metadata=None, max_num_splits=1, causal=True, mm_prefix_query_range_tensor=None, diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 7e588ccde2e5..61ad7e24dd13 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -118,7 +118,6 @@ def _make_scheduler_output( total_num_scheduled_tokens=0, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=finished_req_ids, free_encoder_mm_hashes=[], preempted_req_ids=preempted_req_ids, @@ -305,7 +304,6 @@ def _make_postprocess_scheduler_output( total_num_scheduled_tokens=sum(num_scheduled_tokens.values()), scheduled_spec_decode_tokens=scheduled_spec_decode_tokens or {}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], preempted_req_ids=set(), diff --git a/vllm/config/model.py b/vllm/config/model.py index 545c306e1e1f..87759f46061b 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -310,12 +310,8 @@ class ModelConfig: window functionality of the model, capping to sliding window size. If the model does not support sliding window, this argument is ignored.""" disable_cascade_attn: bool = True - """Disable cascade attention for V1. While cascade attention does not - change the mathematical correctness, disabling it could be useful for - preventing potential numerical issues. This defaults to True, so users - must opt in to cascade attention by setting this to False. Even when this - is set to False, cascade attention will only be used when the heuristic - tells that it's beneficial.""" + """Deprecated: cascade attention has been removed, so this has no effect. + It will be removed in v0.33.0.""" skip_tokenizer_init: bool = False """Skip initialization of tokenizer and detokenizer. Expects valid `prompt_token_ids` and `None` for prompt from the input. The generated diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index b4298cd145b8..af4c99113f0f 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1656,18 +1656,6 @@ def __post_init__(self): else: self.parallel_config.disable_nccl_for_dp_synchronization = False - if ( - self.speculative_config is not None - and self.scheduler_config.async_scheduling - and self.model_config is not None - and not self.model_config.disable_cascade_attn - ): - logger.warning_once( - "Disabling cascade attention (not yet compatible with " - "async speculative decoding).", - ) - self.model_config.disable_cascade_attn = True - if ( self.observability_config.per_request_spec_decode_metrics != "none" and self.speculative_config is None @@ -2090,36 +2078,17 @@ def has_blocked_weights(): ) # final check of cudagraph mode after all possible updates - if current_platform.is_cuda_alike(): - if ( - self.compilation_config.cudagraph_mode.has_full_cudagraphs() - and self.model_config is not None - and not self.model_config.disable_cascade_attn - and not self.compilation_config.cudagraph_mode.has_piecewise_cudagraphs() # noqa: E501 - ): - logger.warning_once( - "No piecewise cudagraph for executing cascade attention. " - "Will fall back to eager execution if a batch runs into " - "cascade attentions." - ) - - if self.compilation_config.cudagraph_mode.requires_piecewise_compilation(): - assert ( - self.compilation_config.mode == CompilationMode.VLLM_COMPILE - or envs.VLLM_USE_BREAKABLE_CUDAGRAPH - ), ( - "Compilation mode should be CompilationMode.VLLM_COMPILE " - "when cudagraph_mode piecewise cudagraphs is used, " - f"cudagraph_mode={self.compilation_config.cudagraph_mode}" - ) if ( - self.model_config - and envs.VLLM_BATCH_INVARIANT - and not self.model_config.disable_cascade_attn + current_platform.is_cuda_alike() + and self.compilation_config.cudagraph_mode.requires_piecewise_compilation() ): - self.model_config.disable_cascade_attn = True - logger.warning_once( - "Disabling cascade attention when VLLM_BATCH_INVARIANT is enabled.", + assert ( + self.compilation_config.mode == CompilationMode.VLLM_COMPILE + or envs.VLLM_USE_BREAKABLE_CUDAGRAPH + ), ( + "Compilation mode should be CompilationMode.VLLM_COMPILE " + "when cudagraph_mode piecewise cudagraphs is used, " + f"cudagraph_mode={self.compilation_config.cudagraph_mode}" ) if self.parallel_config.use_ubatching: @@ -2137,10 +2106,6 @@ def has_blocked_weights(): "--all2all-backend=nixl_ep and install the matching kernels." ) - if not self.model_config.disable_cascade_attn: - self.model_config.disable_cascade_attn = True - logger.warning_once("Disabling cascade attention when DBO is enabled.") - if not self.instance_id: self.instance_id = random_uuid()[:5] diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index a89290081605..ea628467e0b3 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -453,6 +453,12 @@ def get_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]: ) +_DISABLE_CASCADE_ATTN_DEPRECATION_MESSAGE = ( + "--disable-cascade-attn is deprecated and has no effect since cascade " + "attention has been removed. It will be removed in v0.33.0." +) + + class DeprecatedLogConfigFileAction(argparse.Action): """Warn when the legacy ``--log-config-file`` option is used.""" @@ -579,7 +585,7 @@ class EngineArgs: CacheConfig, "prefix_cache_retention_interval" ) disable_sliding_window: bool = ModelConfig.disable_sliding_window - disable_cascade_attn: bool = ModelConfig.disable_cascade_attn + disable_cascade_attn: bool | None = None offload_backend: str = OffloadConfig.offload_backend cpu_offload_gb: float = UVAOffloadConfig.cpu_offload_gb cpu_offload_params: set[str] = get_field(UVAOffloadConfig, "cpu_offload_params") @@ -985,7 +991,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--disable-sliding-window", **model_kwargs["disable_sliding_window"] ) model_group.add_argument( - "--disable-cascade-attn", **model_kwargs["disable_cascade_attn"] + "--disable-cascade-attn", + action=argparse.BooleanOptionalAction, + default=None, + help=_DISABLE_CASCADE_ATTN_DEPRECATION_MESSAGE, ) model_group.add_argument( "--skip-tokenizer-init", **model_kwargs["skip_tokenizer_init"] @@ -1940,6 +1949,8 @@ def create_model_config(self) -> ModelConfig: "launched vLLM.", self.seed, ) + if self.disable_cascade_attn is not None: + warnings.warn(_DISABLE_CASCADE_ATTN_DEPRECATION_MESSAGE, stacklevel=1) return ModelConfig( model=self.model, @@ -1971,7 +1982,6 @@ def create_model_config(self) -> ModelConfig: use_fp64_gumbel=self.use_fp64_gumbel, enable_trace_replay=self.enable_trace_replay, disable_sliding_window=self.disable_sliding_window, - disable_cascade_attn=self.disable_cascade_attn, skip_tokenizer_init=self.skip_tokenizer_init, enable_prompt_embeds=self.enable_prompt_embeds, served_model_name=self.served_model_name, diff --git a/vllm/model_executor/models/extract_hidden_states.py b/vllm/model_executor/models/extract_hidden_states.py index 653e7691fba3..89253ac5d8ed 100644 --- a/vllm/model_executor/models/extract_hidden_states.py +++ b/vllm/model_executor/models/extract_hidden_states.py @@ -126,10 +126,6 @@ def get_impl_cls() -> type["CacheOnlyAttentionImpl"]: def get_builder_cls() -> type["CacheOnlyAttentionMetadataBuilder"]: return CacheOnlyAttentionMetadataBuilder - @staticmethod - def use_cascade_attention(*args, **kwargs) -> bool: - return False - @classmethod def get_supported_head_sizes(cls) -> list[int]: return [] @@ -158,11 +154,6 @@ def build( common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> CacheOnlyAttentionMetadata: - use_cascade = common_prefix_len > 0 - if use_cascade: - raise NotImplementedError( - "Cascade attention not supported by CacheOnlyAttention" - ) causal = common_attn_metadata.causal if not causal: raise NotImplementedError( diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 43d472f22e67..542a7a405df5 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -216,9 +216,6 @@ def inference_mode(cls): def check_and_update_config(cls, vllm_config: VllmConfig) -> None: model_config = vllm_config.model_config - if model_config is not None: - model_config.disable_cascade_attn = True - # Import lazily: vllm.triton_utils imports vllm.platforms.current_platform, # which is still being resolved while this platform class is loading. from vllm.triton_utils import HAS_TRITON diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 68bb0b3bdd63..466bba586d08 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -552,9 +552,7 @@ def unpadded( class AttentionCGSupport(Enum): - """Constants for the cudagraph support of the attention backend - Here we do not consider the cascade attention, as currently - it is never cudagraph supported.""" + """Constants for the cudagraph support of the attention backend""" ALWAYS = 3 """Cudagraph always supported; supports mixed-prefill-decode""" @@ -692,7 +690,8 @@ def build( Some builders (MLA) require reorder_batch to be called prior to build. Args: - common_prefix_len: The length of the common prefix of the batch. + common_prefix_len: Deprecated and always 0. Kept for compatibility + with out-of-tree builders. common_attn_metadata: The common attention metadata. fast_build: The meta-data will prioritize speed of building over then speed at execution. Can be used for spec-decode where the @@ -758,20 +757,6 @@ def update_draft_decode_metadata(self, metadata: M) -> None: """ raise NotImplementedError - def use_cascade_attention( - self, - common_prefix_len: int, - query_lens: np.ndarray, - num_query_heads: int, - num_kv_heads: int, - use_alibi: bool, - use_sliding_window: bool, - use_local_attention: bool, - num_sms: int, - dcp_world_size: int, - ) -> bool: - return False - class AttentionLayer(Protocol): _q_scale: torch.Tensor diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index 46b01cc5519a..9062b48aeaf3 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -105,10 +105,6 @@ def get_impl_cls() -> type["CPUAttentionBackendImpl"]: def get_builder_cls() -> type["CPUAttentionMetadataBuilder"]: return CPUAttentionMetadataBuilder - @staticmethod - def use_cascade_attention(*args, **kwargs) -> bool: - return False - @dataclass class CPUAttentionMetadata: diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 4a336586acd8..6e906bc520c4 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -527,13 +527,6 @@ class FlashAttentionMetadata: block_table: torch.Tensor slot_mapping: torch.Tensor - # For cascade attention. - use_cascade: bool - common_prefix_len: int - cu_prefix_query_lens: torch.Tensor | None - prefix_kv_lens: torch.Tensor | None - suffix_kv_lens: torch.Tensor | None - # For GQA DCP max_dcp_context_kv_len: int | None = None dcp_context_kv_lens: torch.Tensor | None = None @@ -547,7 +540,6 @@ class FlashAttentionMetadata: # Optional aot scheduling scheduler_metadata: torch.Tensor | None = None - prefix_scheduler_metadata: torch.Tensor | None = None max_num_splits: int = 0 causal: bool | torch.Tensor = True @@ -870,7 +862,6 @@ def build( if envs.VLLM_BATCH_INVARIANT: max_num_splits = 1 - use_cascade = common_prefix_len > 0 max_dcp_context_kv_len = 0 dcp_context_kv_lens = None num_decode_reqs = 0 @@ -878,11 +869,6 @@ def build( num_decode_tokens = 0 num_prefill_tokens = 0 - cu_prefix_query_lens = None - prefix_kv_lens = None - suffix_kv_lens = None - prefix_scheduler_metadata = None - if self.dcp_world_size > 1: query_lens = query_start_loc[1:] - query_start_loc[:-1] context_kv_lens = seq_lens - query_lens @@ -946,35 +932,6 @@ def build( causal=False, max_num_splits=max_num_splits, ) - elif use_cascade: - cu_prefix_query_lens = torch.tensor( - [0, num_actual_tokens], dtype=torch.int32, device=self.device - ) - prefix_kv_lens = torch.tensor( - [common_prefix_len], dtype=torch.int32, device=self.device - ) - # Use GPU tensor directly - no CPU sync needed - suffix_kv_lens = seq_lens[:num_reqs] - common_prefix_len - prefix_scheduler_metadata = self._get_scheduler_metadata( - aot_schedule=aot_schedule, - batch_size=1, - cu_query_lens=cu_prefix_query_lens, - max_query_len=num_actual_tokens, - seqlens=prefix_kv_lens, - max_seq_len=common_prefix_len, - causal=False, - max_num_splits=max_num_splits, - ) - scheduler_metadata = self._get_scheduler_metadata( - aot_schedule=aot_schedule, - batch_size=num_reqs, - cu_query_lens=query_start_loc, - max_query_len=max_query_len, - seqlens=suffix_kv_lens, - max_seq_len=max_seq_len - common_prefix_len, - causal=True, - max_num_splits=max_num_splits, - ) else: scheduler_metadata = self._get_scheduler_metadata( aot_schedule=aot_schedule, @@ -1009,13 +966,7 @@ def build( num_prefill_reqs=num_prefill_reqs, num_decode_tokens=num_decode_tokens, num_prefill_tokens=num_prefill_tokens, - use_cascade=use_cascade, - common_prefix_len=common_prefix_len, scheduler_metadata=scheduler_metadata, - cu_prefix_query_lens=cu_prefix_query_lens, - prefix_kv_lens=prefix_kv_lens, - suffix_kv_lens=suffix_kv_lens, - prefix_scheduler_metadata=prefix_scheduler_metadata, max_num_splits=max_num_splits, causal=causal, ) @@ -1083,7 +1034,6 @@ def update_draft_decode_metadata(self, metadata: FlashAttentionMetadata) -> None num_reqs = metadata.num_decode_reqs or metadata.seq_lens.shape[0] assert self.dcp_world_size == 1 - assert not metadata.use_cascade scheduler_metadata = self._get_scheduler_metadata( aot_schedule=True, @@ -1098,12 +1048,6 @@ def update_draft_decode_metadata(self, metadata: FlashAttentionMetadata) -> None metadata.scheduler_metadata = self._store_scheduler_metadata(scheduler_metadata) - def use_cascade_attention(self, *args, **kwargs) -> bool: - if self.fa4_hd256: - # Cascade may use a non-page-aligned prefix length. - return False - return use_cascade_attention(*args, **kwargs) - class FlashAttentionImpl(AttentionImpl): can_return_lse_for_decode: bool = True @@ -1329,186 +1273,156 @@ def forward( key_cache = key_cache.view(current_platform.fp8_dtype()) value_cache = value_cache.view(current_platform.fp8_dtype()) - if not attn_metadata.use_cascade: - cu_seqlens_q = attn_metadata.query_start_loc - seqused_k = attn_metadata.seq_lens - max_seqlen_q = attn_metadata.max_query_len - max_seqlen_k = attn_metadata.max_seq_len - block_table = attn_metadata.block_table - scheduler_metadata = attn_metadata.scheduler_metadata + cu_seqlens_q = attn_metadata.query_start_loc + seqused_k = attn_metadata.seq_lens + max_seqlen_q = attn_metadata.max_query_len + max_seqlen_k = attn_metadata.max_seq_len + block_table = attn_metadata.block_table + scheduler_metadata = attn_metadata.scheduler_metadata + + descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads) - descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads) + q_descale = ( + layer._q_scale.expand(descale_shape) + if self.supports_quant_query_input + else None + ) + k_descale = layer._k_scale.expand(descale_shape) + v_descale = layer._v_scale.expand(descale_shape) - q_descale = ( - layer._q_scale.expand(descale_shape) - if self.supports_quant_query_input - else None + if self.dcp_world_size > 1: + self._forward_with_dcp( + query[:num_actual_tokens], + key[:num_actual_tokens], + value[:num_actual_tokens], + key_cache, + value_cache, + output[:num_actual_tokens], + attn_metadata, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + ) + return output + else: + causal = attn_metadata.causal + is_dynamic_causal = isinstance(causal, torch.Tensor) + + # The layer's own window wins over the group's: one KV cache + # group can hold both windowed and global layers (e.g. Gemma-3 + # with the hybrid KV cache manager disabled), and the group spec + # cannot describe both. + window = _maybe_symmetrize_window(self.sliding_window, causal) + sliding_window_size: list[int] | None = ( + list(window) if window is not None else None ) - k_descale = layer._k_scale.expand(descale_shape) - v_descale = layer._v_scale.expand(descale_shape) - - if self.dcp_world_size > 1: - self._forward_with_dcp( - query[:num_actual_tokens], - key[:num_actual_tokens], - value[:num_actual_tokens], - key_cache, - value_cache, - output[:num_actual_tokens], - attn_metadata, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - ) - return output - else: - causal = attn_metadata.causal - is_dynamic_causal = isinstance(causal, torch.Tensor) - - # The layer's own window wins over the group's: one KV cache - # group can hold both windowed and global layers (e.g. Gemma-3 - # with the hybrid KV cache manager disabled), and the group spec - # cannot describe both. - window = _maybe_symmetrize_window(self.sliding_window, causal) - sliding_window_size: list[int] | None = ( - list(window) if window is not None else None - ) - mm_prefix_query_ranges = attn_metadata.mm_prefix_query_range_tensor - mm_mask_mod = None - mm_aux = None - if ( - mm_prefix_query_ranges is not None - and not is_dynamic_causal - and causal is True - and self.vllm_flash_attn_version == 4 - ): - # Triton convention: 1 + window_size[0]. Global layers store - # (-1, -1) → sw stays None. - layer_window = self.sliding_window - sw_val = ( - 1 + layer_window[0] - if layer_window is not None and layer_window[0] >= 0 - else None - ) - # Gemma4: also clamp the bidirectional block to the - # sliding window when the layer opts in - # (mm_prefix_clamp_sliding_window flag from PR #47217). - mm_clamp_sw = 0 - if ( - getattr(layer, "mm_prefix_clamp_sliding_window", False) - and sw_val is not None - ): - mm_clamp_sw = sw_val - mm_mask_mod = _make_mm_prefix_mask_mod( - sliding_window=mm_clamp_sw, - sliding_window_left=sw_val, - ) - mm_aux = [mm_prefix_query_ranges, attn_metadata.query_start_loc] - # mm_prefix is (causal ∧ window) ∨ bidirectional-range — - # not ⊆ causal. FA #155 stopped auto-clearing causal/local - # when mask_mod is set, so the caller must disable them or - # the built-in causal path shorts out / clips the mask_mod. - causal = False - sliding_window_size = None - - # R-SWA: use CuTE-DSL mask_mod on FA4 for exact token-level - # mask without block-size approximation. The mask_mod encodes - # "causal AND (kv < prefix_len OR q - kv < rswa_window)", which - # supersedes any FA-layer sliding_window_size parameter. - rswa_mask_mod_fn = None - rswa_aux = None + mm_prefix_query_ranges = attn_metadata.mm_prefix_query_range_tensor + mm_mask_mod = None + mm_aux = None + if ( + mm_prefix_query_ranges is not None + and not is_dynamic_causal + and causal is True + and self.vllm_flash_attn_version == 4 + ): + # Triton convention: 1 + window_size[0]. Global layers store + # (-1, -1) → sw stays None. + layer_window = self.sliding_window + sw_val = ( + 1 + layer_window[0] + if layer_window is not None and layer_window[0] >= 0 + else None + ) + # Gemma4: also clamp the bidirectional block to the + # sliding window when the layer opts in + # (mm_prefix_clamp_sliding_window flag from PR #47217). + mm_clamp_sw = 0 if ( - attn_metadata.rswa_prefix_lens is not None - and self.vllm_flash_attn_version == 4 - and not is_dynamic_causal + getattr(layer, "mm_prefix_clamp_sliding_window", False) + and sw_val is not None ): - rswa_mask_mod_fn = _make_rswa_mask_mod() - rswa_aux = [ - attn_metadata.rswa_prefix_lens.to(torch.int32), - attn_metadata.rswa_window_tensor, # pre-allocated CUDA tensor - ] - # mask_mod fully expresses R-SWA; disable FA's own window. - sliding_window_size = None - - dynamic_causal = None - if isinstance(causal, torch.Tensor): - if self.vllm_flash_attn_version != 4: - raise NotImplementedError( - "Per-sequence causal requires FA4. Current version: " - f"FA{self.vllm_flash_attn_version}" - ) - dynamic_causal = causal - has_window = ( - sliding_window_size is not None and sliding_window_size[1] >= 0 + mm_clamp_sw = sw_val + mm_mask_mod = _make_mm_prefix_mask_mod( + sliding_window=mm_clamp_sw, + sliding_window_left=sw_val, + ) + mm_aux = [mm_prefix_query_ranges, attn_metadata.query_start_loc] + # mm_prefix is (causal ∧ window) ∨ bidirectional-range — + # not ⊆ causal. FA #155 stopped auto-clearing causal/local + # when mask_mod is set, so the caller must disable them or + # the built-in causal path shorts out / clips the mask_mod. + causal = False + sliding_window_size = None + + # R-SWA: use CuTE-DSL mask_mod on FA4 for exact token-level + # mask without block-size approximation. The mask_mod encodes + # "causal AND (kv < prefix_len OR q - kv < rswa_window)", which + # supersedes any FA-layer sliding_window_size parameter. + rswa_mask_mod_fn = None + rswa_aux = None + if ( + attn_metadata.rswa_prefix_lens is not None + and self.vllm_flash_attn_version == 4 + and not is_dynamic_causal + ): + rswa_mask_mod_fn = _make_rswa_mask_mod() + rswa_aux = [ + attn_metadata.rswa_prefix_lens.to(torch.int32), + attn_metadata.rswa_window_tensor, # pre-allocated CUDA tensor + ] + # mask_mod fully expresses R-SWA; disable FA's own window. + sliding_window_size = None + + dynamic_causal = None + if isinstance(causal, torch.Tensor): + if self.vllm_flash_attn_version != 4: + raise NotImplementedError( + "Per-sequence causal requires FA4. Current version: " + f"FA{self.vllm_flash_attn_version}" ) - causal = not has_window - - num_splits = attn_metadata.max_num_splits - if self.fa4_hd256: - # hd256 requires page-aligned lengths, exact-width block - # tables, and no SplitKV. - num_pages = cdiv(max_seqlen_k, FA4_HD256_PAGE_SIZE) - max_seqlen_k = num_pages * FA4_HD256_PAGE_SIZE - block_table = block_table[:, :num_pages] - num_splits = 1 - - _FA4_DENSE_ATTENTION_KERNEL( - q=query[:num_actual_tokens], - k=key_cache, - v=value_cache, - out=output[:num_actual_tokens], - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=max_seqlen_q, - seqused_k=seqused_k, - max_seqlen_k=max_seqlen_k, - softmax_scale=self.scale, - causal=causal, - alibi_slopes=self.alibi_slopes, - window_size=sliding_window_size, - block_table=block_table, - softcap=self.logits_soft_cap, - scheduler_metadata=scheduler_metadata, - fa_version=self.vllm_flash_attn_version, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - dynamic_causal=dynamic_causal, - num_splits=num_splits, - s_aux=self.sinks, - mask_mod=rswa_mask_mod_fn or mm_mask_mod, - aux_tensors=rswa_aux or mm_aux, + dynamic_causal = causal + has_window = ( + sliding_window_size is not None and sliding_window_size[1] >= 0 ) - return output - - # Cascade attention (rare case). - cascade_attention( - output[:num_actual_tokens], - query[:num_actual_tokens], - key_cache, - value_cache, - cu_query_lens=attn_metadata.query_start_loc, - max_query_len=attn_metadata.max_query_len, - cu_prefix_query_lens=attn_metadata.cu_prefix_query_lens, - prefix_kv_lens=attn_metadata.prefix_kv_lens, - suffix_kv_lens=attn_metadata.suffix_kv_lens, - max_kv_len=attn_metadata.max_seq_len, - softmax_scale=self.scale, - alibi_slopes=self.alibi_slopes, - sliding_window=self.sliding_window, - logits_soft_cap=self.logits_soft_cap, - block_table=attn_metadata.block_table, - common_prefix_len=attn_metadata.common_prefix_len, - max_num_splits=attn_metadata.max_num_splits, - fa_version=self.vllm_flash_attn_version, - prefix_scheduler_metadata=attn_metadata.prefix_scheduler_metadata, - suffix_scheduler_metadata=attn_metadata.scheduler_metadata, - q_descale=layer._q_scale, - k_descale=layer._k_scale, - v_descale=layer._v_scale, - s_aux=self.sinks, - ) - return output + causal = not has_window + + num_splits = attn_metadata.max_num_splits + if self.fa4_hd256: + # hd256 requires page-aligned lengths, exact-width block + # tables, and no SplitKV. + num_pages = cdiv(max_seqlen_k, FA4_HD256_PAGE_SIZE) + max_seqlen_k = num_pages * FA4_HD256_PAGE_SIZE + block_table = block_table[:, :num_pages] + num_splits = 1 + + _FA4_DENSE_ATTENTION_KERNEL( + q=query[:num_actual_tokens], + k=key_cache, + v=value_cache, + out=output[:num_actual_tokens], + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + seqused_k=seqused_k, + max_seqlen_k=max_seqlen_k, + softmax_scale=self.scale, + causal=causal, + alibi_slopes=self.alibi_slopes, + window_size=sliding_window_size, + block_table=block_table, + softcap=self.logits_soft_cap, + scheduler_metadata=scheduler_metadata, + fa_version=self.vllm_flash_attn_version, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + dynamic_causal=dynamic_causal, + num_splits=num_splits, + s_aux=self.sinks, + mask_mod=rswa_mask_mod_fn or mm_mask_mod, + aux_tensors=rswa_aux or mm_aux, + ) + return output def do_kv_cache_update( self, @@ -1957,176 +1871,4 @@ def rswa_mask_mod( return rswa_mask_mod -def use_cascade_attention( - common_prefix_len: int, - query_lens: np.ndarray, - num_query_heads: int, - num_kv_heads: int, - use_alibi: bool, - use_sliding_window: bool, - use_local_attention: bool, - num_sms: int, - dcp_world_size: int, -) -> bool: - """Decide whether to use cascade attention. - - This function 1) checks whether cascade attention is supported with the - given configuration, and 2) heuristically decides whether using cascade - attention can improve performance. - """ - # Too short common prefix. Probably not worth using cascade attention. - # We use an arbitrary threshold of 256 tokens. TODO: Tune this threshold. - # NOTE(woosuk): This is the common case. We should return False as soon as - # possible to avoid any unnecessary computation. - if common_prefix_len < 256: - return False - # Cascade attention is currently not supported with these variants. - if use_alibi or use_sliding_window or use_local_attention: - return False - # Too few queries. Probably not worth using cascade attention. - # We use an arbitrary threshold of 8 queries. TODO: Tune this threshold. - num_reqs = len(query_lens) - if num_reqs < 8: - return False - # disable cascade attention for DCP - if dcp_world_size > 1: - return False - - # Heuristics to decide whether using cascade attention is beneficial. - # 1. When FlashDecoding is not used for normal attention, cascade attention - # is likely to be faster since it saves memory bandwidth. - num_queries_per_kv = num_query_heads // num_kv_heads - # The criteria for using FlashDecoding can be found in the following link: - # https://github.com/vllm-project/flash-attention/blob/96266b1111111f3d11aabefaf3bacbab6a89d03c/csrc/flash_attn/flash_api.cpp#L535 - use_flash_decoding = ( - num_queries_per_kv > 1 - and not use_sliding_window - and not use_alibi - and np.all(query_lens == 1) - ) - if not use_flash_decoding: - # Use cascade attention. - return True - - # 2. When FlashDecoding is used for normal attention, it is not clear - # whether cascade attention is beneficial, because FlashDecoding can - # launch more CTAs than cascade attention. - # We use a simple performance model to compare the two methods. - # NOTE(woosuk): The performance model is very rough and may not be - # accurate. - num_tokens = num_reqs - # NOTE(woosuk): These are default tile sizes. flash-attn might use - # different tile sizes (e.g., 64 or 256) depending on the configuration. - q_tile_size = 128 - kv_tile_size = 128 - num_prefix_tiles = cdiv(common_prefix_len, kv_tile_size) - - cascade_ctas = num_query_heads * cdiv(num_tokens, q_tile_size) - cascade_waves = cdiv(cascade_ctas, num_sms) - cascade_time = cascade_waves * num_prefix_tiles - - flash_decoding_ctas = ( - num_reqs * num_kv_heads * cdiv(num_queries_per_kv, q_tile_size) - ) - flash_decoding_ctas *= num_prefix_tiles - flash_decoding_time = cdiv(flash_decoding_ctas, num_sms) - - # Use cascade attention if it is faster than FlashDecoding. - return cascade_time < flash_decoding_time - - -def cascade_attention( - output: torch.Tensor, - query: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - cu_query_lens: torch.Tensor, - max_query_len: int, - cu_prefix_query_lens: torch.Tensor, - prefix_kv_lens: torch.Tensor, - suffix_kv_lens: torch.Tensor, - max_kv_len: int, - softmax_scale: float, - alibi_slopes: torch.Tensor | None, - sliding_window: tuple[int, int], - logits_soft_cap: float, - block_table: torch.Tensor, - common_prefix_len: int, - max_num_splits: int, - fa_version: int, - prefix_scheduler_metadata: torch.Tensor | None = None, - suffix_scheduler_metadata: torch.Tensor | None = None, - q_descale: torch.Tensor | None = None, - k_descale: torch.Tensor | None = None, - v_descale: torch.Tensor | None = None, - s_aux: torch.Tensor | None = None, -) -> torch.Tensor: - assert alibi_slopes is None, "Cascade attention does not support ALiBi." - # TODO: Support sliding window. - assert sliding_window == (-1, -1), ( - "Cascade attention does not support sliding window." - ) - - num_tokens = query.shape[0] - block_size = key_cache.shape[-3] - assert common_prefix_len % block_size == 0 - num_common_kv_blocks = common_prefix_len // block_size - assert num_common_kv_blocks > 0 - descale_shape = (cu_prefix_query_lens.shape[0] - 1, key_cache.shape[-2]) - - # Process shared prefix. - prefix_output, prefix_lse = flash_attn_varlen_func( - q=query, - k=key_cache, - v=value_cache, - cu_seqlens_q=cu_prefix_query_lens, - seqused_k=prefix_kv_lens, - max_seqlen_q=num_tokens, - max_seqlen_k=common_prefix_len, - softmax_scale=softmax_scale, - causal=False, - window_size=list(sliding_window), - block_table=block_table[:1], - softcap=logits_soft_cap, - return_softmax_lse=True, - scheduler_metadata=prefix_scheduler_metadata, - fa_version=fa_version, - q_descale=q_descale.expand(descale_shape) if q_descale is not None else None, - k_descale=k_descale.expand(descale_shape) if k_descale is not None else None, - v_descale=v_descale.expand(descale_shape) if v_descale is not None else None, - # s_aux is incorporated into prefix_lse inside the GPU kernel, - # enabling its effect during the final attention merge. - s_aux=s_aux, - num_splits=1 if envs.VLLM_BATCH_INVARIANT else max_num_splits, - ) - - descale_shape = (cu_query_lens.shape[0] - 1, key_cache.shape[-2]) - - # Process suffix per query. - suffix_output, suffix_lse = flash_attn_varlen_func( - q=query, - k=key_cache, - v=value_cache, - cu_seqlens_q=cu_query_lens, - seqused_k=suffix_kv_lens, - max_seqlen_q=max_query_len, - max_seqlen_k=max_kv_len - common_prefix_len, - softmax_scale=softmax_scale, - causal=True, - window_size=list(sliding_window), - block_table=block_table[:, num_common_kv_blocks:], - softcap=logits_soft_cap, - return_softmax_lse=True, - scheduler_metadata=suffix_scheduler_metadata, - fa_version=fa_version, - q_descale=q_descale.expand(descale_shape) if q_descale is not None else None, - k_descale=k_descale.expand(descale_shape) if k_descale is not None else None, - v_descale=v_descale.expand(descale_shape) if v_descale is not None else None, - num_splits=1 if envs.VLLM_BATCH_INVARIANT else max_num_splits, - ) - - # Merge prefix and suffix outputs, and store the result in output. - merge_attn_states(output, prefix_output, prefix_lse, suffix_output, suffix_lse) - - _FA4_DENSE_ATTENTION_KERNEL = FA4DenseAttentionKernel() diff --git a/vllm/v1/attention/backends/flash_attn_diffkv.py b/vllm/v1/attention/backends/flash_attn_diffkv.py index 4bb58ef4ea43..7b38a3695eac 100644 --- a/vllm/v1/attention/backends/flash_attn_diffkv.py +++ b/vllm/v1/attention/backends/flash_attn_diffkv.py @@ -26,7 +26,6 @@ FlashAttentionBackend, FlashAttentionImpl, FlashAttentionMetadata, - cascade_attention, ) logger = init_logger(__name__) @@ -215,86 +214,54 @@ def forward( key_cache = key_cache.view(current_platform.fp8_dtype()) value_cache = value_cache.view(current_platform.fp8_dtype()) - if not attn_metadata.use_cascade: - cu_seqlens_q = attn_metadata.query_start_loc - seqused_k = attn_metadata.seq_lens - max_seqlen_q = attn_metadata.max_query_len - max_seqlen_k = attn_metadata.max_seq_len - block_table = attn_metadata.block_table - scheduler_metadata = attn_metadata.scheduler_metadata - - descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads) - - if self.dcp_world_size > 1: - self._forward_with_dcp( - query[:num_actual_tokens], - key[:num_actual_tokens], - value[:num_actual_tokens], - key_cache, - value_cache, - output[:num_actual_tokens], - attn_metadata, - q_descale=layer._q_scale.expand(descale_shape), - k_descale=layer._k_scale.expand(descale_shape), - v_descale=layer._v_scale.expand(descale_shape), - ) - return output - else: - sliding_window_size = ( - list(self.sliding_window) - if self.sliding_window is not None - else None - ) - flash_attn_varlen_func( - q=query[:num_actual_tokens], - k=key_cache, - v=value_cache, - out=output[:num_actual_tokens], - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=max_seqlen_q, - seqused_k=seqused_k, - max_seqlen_k=max_seqlen_k, - softmax_scale=self.scale, - causal=attn_metadata.causal, - alibi_slopes=self.alibi_slopes, - window_size=sliding_window_size, - block_table=block_table, - softcap=self.logits_soft_cap, - scheduler_metadata=scheduler_metadata, - fa_version=self.vllm_flash_attn_version, - q_descale=layer._q_scale.expand(descale_shape), - k_descale=layer._k_scale.expand(descale_shape), - v_descale=layer._v_scale.expand(descale_shape), - num_splits=attn_metadata.max_num_splits, - s_aux=self.sinks, - ) - return output - - # Cascade attention (rare case). - cascade_attention( - output[:num_actual_tokens], - query[:num_actual_tokens], - key_cache, - value_cache, - cu_query_lens=attn_metadata.query_start_loc, - max_query_len=attn_metadata.max_query_len, - cu_prefix_query_lens=attn_metadata.cu_prefix_query_lens, - prefix_kv_lens=attn_metadata.prefix_kv_lens, - suffix_kv_lens=attn_metadata.suffix_kv_lens, - max_kv_len=attn_metadata.max_seq_len, - softmax_scale=self.scale, - alibi_slopes=self.alibi_slopes, - sliding_window=self.sliding_window, - logits_soft_cap=self.logits_soft_cap, - block_table=attn_metadata.block_table, - common_prefix_len=attn_metadata.common_prefix_len, - max_num_splits=attn_metadata.max_num_splits, - fa_version=self.vllm_flash_attn_version, - prefix_scheduler_metadata=attn_metadata.prefix_scheduler_metadata, - suffix_scheduler_metadata=attn_metadata.scheduler_metadata, - q_descale=layer._q_scale, - k_descale=layer._k_scale, - v_descale=layer._v_scale, - s_aux=self.sinks, - ) - return output + cu_seqlens_q = attn_metadata.query_start_loc + seqused_k = attn_metadata.seq_lens + max_seqlen_q = attn_metadata.max_query_len + max_seqlen_k = attn_metadata.max_seq_len + block_table = attn_metadata.block_table + scheduler_metadata = attn_metadata.scheduler_metadata + + descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads) + + if self.dcp_world_size > 1: + self._forward_with_dcp( + query[:num_actual_tokens], + key[:num_actual_tokens], + value[:num_actual_tokens], + key_cache, + value_cache, + output[:num_actual_tokens], + attn_metadata, + q_descale=layer._q_scale.expand(descale_shape), + k_descale=layer._k_scale.expand(descale_shape), + v_descale=layer._v_scale.expand(descale_shape), + ) + return output + else: + sliding_window_size = ( + list(self.sliding_window) if self.sliding_window is not None else None + ) + flash_attn_varlen_func( + q=query[:num_actual_tokens], + k=key_cache, + v=value_cache, + out=output[:num_actual_tokens], + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + seqused_k=seqused_k, + max_seqlen_k=max_seqlen_k, + softmax_scale=self.scale, + causal=attn_metadata.causal, + alibi_slopes=self.alibi_slopes, + window_size=sliding_window_size, + block_table=block_table, + softcap=self.logits_soft_cap, + scheduler_metadata=scheduler_metadata, + fa_version=self.vllm_flash_attn_version, + q_descale=layer._q_scale.expand(descale_shape), + k_descale=layer._k_scale.expand(descale_shape), + v_descale=layer._v_scale.expand(descale_shape), + num_splits=attn_metadata.max_num_splits, + s_aux=self.sinks, + ) + return output diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 1a4145844705..f6c3e262d60f 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -14,7 +14,6 @@ BatchDecodeWithPagedKVCacheWrapper, BatchPrefillWithPagedKVCacheWrapper, BatchPrefillWithRaggedKVCacheWrapper, - MultiLevelCascadeAttentionWrapper, get_seq_lens, nvfp4_quantize_append_paged_kv_cache_with_slot_mapping, ) @@ -732,16 +731,6 @@ class FlashInferMetadata: Will be `None` if `num_decode_tokens == 0`. """ - # --- Special Case: Cascade Attention --- - - use_cascade: bool - """ - If True, the entire batch is a cascade attention call, and the - `prefill` and `decode` fields will both be None. - """ - - cascade_wrapper: MultiLevelCascadeAttentionWrapper | None - class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): kv_cache_spec: AttentionSpec @@ -951,8 +940,6 @@ def __init__( supports_dcp_with_varlen=False, ) - self._cascade_wrapper = None # Wrapper for cascade attention - # Global hyperparameters shared by all attention layers # TODO: discard this for trtllm-gen backend per_layer_parameters = get_per_layer_parameters( @@ -1360,15 +1347,6 @@ def _get_decode_wrapper(self, batch_size: int, use_cudagraph: bool = False): return decode_wrapper - def _get_cascade_wrapper(self): - if self._cascade_wrapper is None: - self._cascade_wrapper = MultiLevelCascadeAttentionWrapper( - 2, - self._get_workspace_buffer(), - get_flashinfer_layout_string(self.kv_cache_layout), - ) - return self._cascade_wrapper - def _compute_flashinfer_kv_metadata( self, num_blocks_np: np.ndarray, @@ -1445,10 +1423,8 @@ def build( qo_indptr_cpu = common_attn_metadata.query_start_loc_cpu # Step 1: Decide which dispatch modes to use: - # - Cascade attention (distinct mode) # - Prefill (FI native or TRTLLM) # - Decode (FI native, XQA, or trtllm-gen) - use_cascade = common_prefix_len > 0 uses_spec_reorder = self.reorder_batch_threshold > 1 # Page sizes >= 128 must use trtllm-gen; force it for prefill too. prefill_force_trtllm = ( @@ -1486,12 +1462,11 @@ def build( ) if not all_uses_trtllm: - sinks_would_be_dropped = self.use_dcp or use_cascade - if self.has_sinks and sinks_would_be_dropped: + if self.has_sinks and self.use_dcp: raise NotImplementedError( - "FlashInfer BatchDCPPrefillWrapper and " - "MultiLevelCascadeAttentionWrapper do not support attention sinks. " - "Please use TRTLLM on Blackwell or FlashAttention on earlier GPUs." + "FlashInfer BatchDCPPrefillWrapper does not support attention " + "sinks. Please use TRTLLM on Blackwell or FlashAttention on " + "earlier GPUs." ) if not self.global_hyperparameters.has_same_window_lefts: @@ -1508,7 +1483,7 @@ def build( ) # Step 2: Initialize the output metadata - # Leave prefill/decode/cascade_wrapper empty, to be populated + # Leave prefill/decode empty, to be populated # case by case depending on the batch contents and backend selection. attn_metadata = FlashInferMetadata( kv_cache_layout=self.kv_cache_layout, @@ -1521,10 +1496,8 @@ def build( num_prefills=num_prefills, num_prefill_tokens=num_prefill_tokens, causal=causal, - use_cascade=use_cascade, prefill=None, decode=None, - cascade_wrapper=None, ) # Guard access to seq_lens_cpu, which may not always be needed @@ -1532,7 +1505,7 @@ def build( # When all attention (both prefill and decode) uses TRTLLM, # seq_lens_cpu is not needed since TRTLLM paths use GPU tensors # (block_tables, seq_lens) directly. - needs_seq_lens_cpu = self.use_dcp or use_cascade or not all_uses_trtllm + needs_seq_lens_cpu = self.use_dcp or not all_uses_trtllm if needs_seq_lens_cpu: if ( self._num_speculative_tokens == 0 @@ -1578,13 +1551,6 @@ def build( seq_lens_np = None num_blocks_np = None - # Adjust num_block_np for cascade attention - if use_cascade: - assert num_blocks_np is not None - assert common_prefix_len % page_size == 0 - num_common_kv_blocks = common_prefix_len // page_size - num_blocks_np -= num_common_kv_blocks - # Compute paged_kv_indices if necessary # paged_kv_indices is only needed for FlashInfer native paths; # XQA/trtllm-gen paths use block_tables directly on GPU. @@ -1592,9 +1558,7 @@ def build( needs_native_paged_decode = ( num_decodes > 0 and not decode_with_flashinfer_trtllm_api ) - needs_paged_kv_indices = ( - use_cascade or needs_native_paged_prefill or needs_native_paged_decode - ) + needs_paged_kv_indices = needs_native_paged_prefill or needs_native_paged_decode if needs_paged_kv_indices: assert num_blocks_np is not None assert seq_lens_np is not None @@ -1608,57 +1572,6 @@ def build( else: paged_kv_indices = None - # Early-out for cascade attention - if use_cascade: - assert num_blocks_np is not None - # Grab the blocks of the shared prefix from the first request. - num_common_kv_blocks = common_prefix_len // page_size - - # Create CPU versions directly for cascade (no GPU versions needed) - shared_qo_indptr_cpu = torch.tensor( - [0, num_actual_tokens], dtype=torch.int32, device="cpu" - ) - shared_kv_page_indptr_cpu = torch.tensor( - [0, num_common_kv_blocks], dtype=torch.int32, device="cpu" - ) - shared_kv_page_indices_cpu = block_table_tensor[0, :num_common_kv_blocks] - shared_kv_last_page_len_cpu = torch.tensor( - [page_size], dtype=torch.int32, device="cpu" - ) - - # Remove the blocks of the shared prefix from all requests. - block_table_tensor = block_table_tensor[:, num_common_kv_blocks:] - num_blocks_np -= num_common_kv_blocks - - assert paged_kv_indices is not None - paged_kv_indptr_cpu = self.paged_kv_indptr.cpu[: 1 + num_reqs] - paged_kv_last_page_len_cpu = self.paged_kv_last_page_len.cpu[:num_reqs] - - attn_metadata.cascade_wrapper = self._get_cascade_wrapper() - # Cascade attention must use the same q dtype for prefill and decode - # because it does not support FP8 kv-cache or FP8 query yet. - assert self.q_data_type_prefill == self.q_data_type_decode - attn_metadata.cascade_wrapper.plan( - qo_indptr_arr=[shared_qo_indptr_cpu, qo_indptr_cpu], - paged_kv_indptr_arr=[shared_kv_page_indptr_cpu, paged_kv_indptr_cpu], - paged_kv_indices_arr=[shared_kv_page_indices_cpu, paged_kv_indices], - paged_kv_last_page_len=[ - shared_kv_last_page_len_cpu, - paged_kv_last_page_len_cpu, - ], - num_qo_heads=self.num_qo_heads, - num_kv_heads=self.num_kv_heads, - head_dim=self.head_dim, - page_size=self.page_size, - causal=True, - sm_scale=self.sm_scale, - window_left=self.window_left, - logits_soft_cap=self.logits_soft_cap, - q_data_type=self.q_data_type_prefill, - kv_data_type=self.kv_cache_dtype, - ) - return attn_metadata - # Step 3: Handle prefill and decode pathways case by case ## PREFILL PATHWAY if num_prefills > 0: @@ -1925,15 +1838,6 @@ def build( def update_draft_decode_metadata(self, _metadata: FlashInferMetadata) -> None: pass - def use_cascade_attention(self, *args, **kwargs) -> bool: - if self.kv_cache_spec.dtype != self.vllm_config.model_config.dtype: - # TODO: The cascade wrapper currently does not support setting - # kv cache dtype to something different from query dtype. - return False - # TODO: Cascade attention doesn't work, disable it for now - # return use_cascade_attention(*args, **kwargs) - return False - class FlashInferImpl(AttentionImpl): can_return_lse_for_decode: bool = True @@ -2223,24 +2127,6 @@ def forward( output_padded = output output = output[:num_actual_tokens] - if attn_metadata.use_cascade: - # Cascade attention (rare case). - assert attn_metadata.cascade_wrapper is not None - stride_order = kv_cache_layout.layer_view_order - if self.is_kvcache_nvfp4: - kv_cache_views = tuple( - cache.permute(*stride_order) - for cache in kv_cache.split(self.num_kv_heads, dim=1) - ) - else: - kv_perm = kv_cache.permute(*stride_order) - kv_cache_views = kv_perm.split(self.head_size, dim=-1) - kv_tuple = tuple( - canonicalize_singleton_dim_strides(cache) for cache in kv_cache_views - ) - output.copy_(attn_metadata.cascade_wrapper.run(query, kv_tuple)) - return output - # When using spec decoding, num_decodes can be < num_decode_tokens # because some decode requests may have more than one query token. num_decode_tokens = attn_metadata.num_decode_tokens diff --git a/vllm/v1/attention/backends/flex_attention.py b/vllm/v1/attention/backends/flex_attention.py index 1b9a21b65578..c2103082df7c 100644 --- a/vllm/v1/attention/backends/flex_attention.py +++ b/vllm/v1/attention/backends/flex_attention.py @@ -145,10 +145,6 @@ def get_impl_cls() -> type["FlexAttentionImpl"]: def get_builder_cls() -> type["FlexAttentionMetadataBuilder"]: return FlexAttentionMetadataBuilder - @staticmethod - def use_cascade_attention(*args, **kwargs) -> bool: - return False - @classmethod def get_supported_head_sizes(cls) -> list[int]: return [] @@ -375,12 +371,6 @@ class FlexAttentionMetadata: block_table: torch.Tensor slot_mapping: torch.Tensor - use_cascade: bool - common_prefix_len: int - cu_prefix_query_lens: torch.Tensor | None - prefix_kv_lens: torch.Tensor | None - suffix_kv_lens: torch.Tensor | None - # Block info total_cache_tokens: int block_size: int @@ -831,11 +821,6 @@ def build_block_mask(self) -> BlockMask: ) def __post_init__(self): - assert self.use_cascade is False, "Not implemented yet." - assert self.common_prefix_len == 0, "Not implemented yet." - assert self.cu_prefix_query_lens is None, "Not implemented yet." - assert self.prefix_kv_lens is None, "Not implemented yet." - assert self.suffix_kv_lens is None, "Not implemented yet." # Create a lookup mapping from query indices -> request number self.doc_ids = _offsets_to_doc_ids_tensor( self.query_start_loc_cpu, self.query_start_loc.device @@ -1029,16 +1014,6 @@ def build( slot_mapping = common_attn_metadata.slot_mapping num_blocks_per_seq = cdiv(seq_lens, self.block_size) - use_cascade = common_prefix_len > 0 - cu_prefix_query_lens = None - prefix_kv_lens = None - suffix_kv_lens = None - if use_cascade: - raise NotImplementedError( - "Cascade prefix attention is not yet implemented " - "for FlexAttention backend" - ) - block_size = self.kv_cache_spec.block_size max_possible_seq_len = self.model_config.max_model_len num_gpu_blocks = self.cache_config.num_gpu_blocks @@ -1128,11 +1103,6 @@ def build( seq_lens=self.persistent_seq_lens, block_table=block_table_tensor, slot_mapping=slot_mapping, - use_cascade=use_cascade, - common_prefix_len=common_prefix_len, - cu_prefix_query_lens=cu_prefix_query_lens, - prefix_kv_lens=prefix_kv_lens, - suffix_kv_lens=suffix_kv_lens, block_size=block_size, max_possible_sequence_length=max_possible_seq_len, num_reqs=num_reqs, @@ -1166,9 +1136,6 @@ def build( return out - def use_cascade_attention(self, *args, **kwargs) -> bool: - return False - class FlexAttentionImpl(AttentionImpl): sliding_window: int | None diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index 329aee53599a..0ec0d5dd3642 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -403,9 +403,6 @@ class AiterFlashAttentionMetadata: prefill_metadata: AiterFlashAttentionPrefillMetadata | None extend_metadata: AiterFlashAttentionChunkPrefillMetadata | None - # For cascade attention. - use_cascade: bool - # Only for fp8 shuffle layout kv cache, we allocate kv_scale for each layer # since we might integrate per token quant for kv cache in the future. k_scale: dict[str, torch.Tensor] | None @@ -696,8 +693,6 @@ def build( chunk_context_metadata=chunk_context_metadata, ) - use_cascade = common_prefix_len > 0 - attn_metadata = AiterFlashAttentionMetadata( num_actual_tokens=common_attn_metadata.num_actual_tokens, query_start_loc=common_attn_metadata.query_start_loc, @@ -714,7 +709,6 @@ def build( decode_metadata=decode_metadata, prefill_metadata=prefill_metadata, extend_metadata=extend_metadata, - use_cascade=use_cascade, k_scale=self.scale, v_scale=self.scale, kv_sharing_metadata=kv_sharing_metadata, @@ -755,14 +749,10 @@ def build_for_drafting( decode_metadata=decode_metadata, prefill_metadata=None, extend_metadata=None, - use_cascade=False, k_scale=self.scale, v_scale=self.scale, ) - def use_cascade_attention(self, *args, **kwargs) -> bool: - return False - class AiterFlashAttentionBackend(AttentionBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] @@ -1219,186 +1209,136 @@ def forward( else "NHD", total_tokens=num_actual_tokens - num_decode_tokens, ) - if not attn_metadata.use_cascade: - # calculate for pure prefills - if num_prefills > 0: - assert attn_metadata.prefill_metadata is not None - assert key is not None and value is not None - - prefill_query = query[num_decode_tokens + num_extend_tokens :] - prefill_key = key[num_decode_tokens + num_extend_tokens :] - prefill_value = value[num_decode_tokens + num_extend_tokens :] - - rocm_aiter_ops.flash_attn_varlen_func( - q=prefill_query, - k=prefill_key, - v=prefill_value, - cu_seqlens_q=attn_metadata.prefill_metadata.query_start_loc, - cu_seqlens_k=attn_metadata.prefill_metadata.query_start_loc, - max_seqlen_q=attn_metadata.prefill_metadata.max_query_len, - max_seqlen_k=attn_metadata.prefill_metadata.max_seq_len, - min_seqlen_q=1, - dropout_p=0.0, - softmax_scale=self.scale, - causal=attn_metadata.causal, - window_size=self.sliding_window, - alibi_slopes=self.alibi_slopes, - out=output_actual_tokens[num_decode_tokens + num_extend_tokens :], - sink_ptr=self.sinks, - ) + # calculate for pure prefills + if num_prefills > 0: + assert attn_metadata.prefill_metadata is not None + assert key is not None and value is not None + + prefill_query = query[num_decode_tokens + num_extend_tokens :] + prefill_key = key[num_decode_tokens + num_extend_tokens :] + prefill_value = value[num_decode_tokens + num_extend_tokens :] + + rocm_aiter_ops.flash_attn_varlen_func( + q=prefill_query, + k=prefill_key, + v=prefill_value, + cu_seqlens_q=attn_metadata.prefill_metadata.query_start_loc, + cu_seqlens_k=attn_metadata.prefill_metadata.query_start_loc, + max_seqlen_q=attn_metadata.prefill_metadata.max_query_len, + max_seqlen_k=attn_metadata.prefill_metadata.max_seq_len, + min_seqlen_q=1, + dropout_p=0.0, + softmax_scale=self.scale, + causal=attn_metadata.causal, + window_size=self.sliding_window, + alibi_slopes=self.alibi_slopes, + out=output_actual_tokens[num_decode_tokens + num_extend_tokens :], + sink_ptr=self.sinks, + ) - # calculate for extends - if num_extends > 0: - assert attn_metadata.extend_metadata is not None - assert key is not None and value is not None - extend_tokens_slice = slice( - num_decode_tokens, num_decode_tokens + num_extend_tokens - ) - extend_queries = query[extend_tokens_slice] - extend_keys = key[extend_tokens_slice] - extend_values = value[extend_tokens_slice] - extend_outputs = output[extend_tokens_slice] - k_scale = layer._k_scale - v_scale = layer._v_scale - if rocm_aiter_ops.is_shuffle_kv_cache_enabled(): - k_scale = attn_metadata.k_scale - v_scale = attn_metadata.v_scale - self.extend_forward( - attn_metadata=attn_metadata, - query=extend_queries, - key=extend_keys, - value=extend_values, - key_cache=key_cache, - value_cache=value_cache, - output=extend_outputs, - cu_seqlens_q=attn_metadata.extend_metadata.query_start_loc, - max_seqlen_q=attn_metadata.extend_metadata.max_query_len, - min_seqlen_q=1, - max_seqlen_k=attn_metadata.extend_metadata.max_seq_len, - block_table=attn_metadata.block_table[ - num_decodes : num_decodes + num_extends - ], - slot_mapping=attn_metadata.slot_mapping[ - num_decodes : num_decodes + num_extends - ], - k_scale=k_scale, - v_scale=v_scale, + # calculate for extends + if num_extends > 0: + assert attn_metadata.extend_metadata is not None + assert key is not None and value is not None + extend_tokens_slice = slice( + num_decode_tokens, num_decode_tokens + num_extend_tokens + ) + extend_queries = query[extend_tokens_slice] + extend_keys = key[extend_tokens_slice] + extend_values = value[extend_tokens_slice] + extend_outputs = output[extend_tokens_slice] + k_scale = layer._k_scale + v_scale = layer._v_scale + if rocm_aiter_ops.is_shuffle_kv_cache_enabled(): + k_scale = attn_metadata.k_scale + v_scale = attn_metadata.v_scale + self.extend_forward( + attn_metadata=attn_metadata, + query=extend_queries, + key=extend_keys, + value=extend_values, + key_cache=key_cache, + value_cache=value_cache, + output=extend_outputs, + cu_seqlens_q=attn_metadata.extend_metadata.query_start_loc, + max_seqlen_q=attn_metadata.extend_metadata.max_query_len, + min_seqlen_q=1, + max_seqlen_k=attn_metadata.extend_metadata.max_seq_len, + block_table=attn_metadata.block_table[ + num_decodes : num_decodes + num_extends + ], + slot_mapping=attn_metadata.slot_mapping[ + num_decodes : num_decodes + num_extends + ], + k_scale=k_scale, + v_scale=v_scale, + ) + + # calculate for decodes + if num_decodes > 0: + assert attn_metadata.decode_metadata is not None + decode_max_query_len = attn_metadata.decode_metadata.max_query_len + + # Use unified_attention for the decodes the paged kernels can't + # take: sliding window, sinks, or a multi-token batch + # (pa_fwd_asm and paged_attention_v1 don't support sinks). + if ( + self.sliding_window[0] != -1 + or decode_max_query_len > 1 + or self.sinks is not None + ): + k_descale, v_descale = self._get_kv_cache_descales(layer, num_decodes) + assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), ( + "Shuffle KV cache layout is not supported with sliding " + "window, sinks, or speculative decoding (multi-token decode)." ) + if not attn_metadata.causal: + from aiter.ops.triton.attention.mha_v3 import ( + flash_attn_with_kvcache, + ) - # calculate for decodes - if num_decodes > 0: - assert attn_metadata.decode_metadata is not None - decode_max_query_len = attn_metadata.decode_metadata.max_query_len - - # Use unified_attention for the decodes the paged kernels can't - # take: sliding window, sinks, or a multi-token batch - # (pa_fwd_asm and paged_attention_v1 don't support sinks). - if ( - self.sliding_window[0] != -1 - or decode_max_query_len > 1 - or self.sinks is not None - ): - k_descale, v_descale = self._get_kv_cache_descales( - layer, num_decodes + decode_query = query[:num_decode_tokens].reshape( + num_decodes, + decode_max_query_len, + query.shape[1], + query.shape[2], ) - assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), ( - "Shuffle KV cache layout is not supported with sliding " - "window, sinks, or speculative decoding (multi-token decode)." + decode_out = flash_attn_with_kvcache( + q=decode_query, + k_cache=key_cache, + v_cache=value_cache, + cache_seqlens=attn_metadata.seq_lens[:num_decodes], + softmax_scale=self.scale, + causal=attn_metadata.causal, + window_size=self.sliding_window, + softcap=self.logits_soft_cap, + q_descale=None, + k_descale=k_descale, + v_descale=v_descale, + page_table=attn_metadata.block_table[:num_decodes], ) - if not attn_metadata.causal: - from aiter.ops.triton.attention.mha_v3 import ( - flash_attn_with_kvcache, - ) - - decode_query = query[:num_decode_tokens].reshape( - num_decodes, - decode_max_query_len, + output[:num_decode_tokens].copy_( + decode_out.reshape( + num_decode_tokens, query.shape[1], query.shape[2], ) - decode_out = flash_attn_with_kvcache( - q=decode_query, - k_cache=key_cache, - v_cache=value_cache, - cache_seqlens=attn_metadata.seq_lens[:num_decodes], - softmax_scale=self.scale, - causal=attn_metadata.causal, - window_size=self.sliding_window, - softcap=self.logits_soft_cap, - q_descale=None, - k_descale=k_descale, - v_descale=v_descale, - page_table=attn_metadata.block_table[:num_decodes], - ) - output[:num_decode_tokens].copy_( - decode_out.reshape( - num_decode_tokens, - query.shape[1], - query.shape[2], - ) - ) - else: - # Non-uniform query lengths can appear in real serving - # traffic (e.g. mixed datasets). Fall back to varlen - # unified_attention instead of asserting. - from aiter.ops.triton.unified_attention import ( - unified_attention, - ) - - unified_attention( - q=query[:num_decode_tokens], - k=key_cache, - v=value_cache, - out=output[:num_decode_tokens], - cu_seqlens_q=attn_metadata.query_start_loc[ - : num_decodes + 1 - ], - max_seqlen_q=decode_max_query_len, - seqused_k=attn_metadata.seq_lens[:num_decodes], - max_seqlen_k=attn_metadata.max_seq_len, - softmax_scale=self.scale, - causal=True, - alibi_slopes=self.alibi_slopes, - window_size=self.sliding_window, - block_table=attn_metadata.block_table[:num_decodes], - softcap=self.logits_soft_cap, - q_descale=None, - k_descale=k_descale, - v_descale=v_descale, - sinks=self.sinks, - ) - return - - # The ll4mi kernel in paged_attention_v1 requires - # HEAD_SIZE >= 16 * NWARPS (= 64 on ROCm with NWARPS=4). - # For smaller head sizes or sliding window attention, - # fall back to the unified_attention triton kernel which - # handles both correctly. - _MIN_HEAD_SIZE_FOR_LL4MI = 64 - use_unified_attention = self.head_size < _MIN_HEAD_SIZE_FOR_LL4MI - - if use_unified_attention: - k_descale, v_descale = self._get_kv_cache_descales( - layer, num_decodes - ) - assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), ( - "unified_attention fallback with shuffle layout " - "is not supported yet." ) + else: + # Non-uniform query lengths can appear in real serving + # traffic (e.g. mixed datasets). Fall back to varlen + # unified_attention instead of asserting. from aiter.ops.triton.unified_attention import ( unified_attention, ) - decode_cu_seqlens_q = attn_metadata.query_start_loc[ - : num_decodes + 1 - ] unified_attention( q=query[:num_decode_tokens], k=key_cache, v=value_cache, out=output[:num_decode_tokens], - cu_seqlens_q=decode_cu_seqlens_q, - max_seqlen_q=1, + cu_seqlens_q=attn_metadata.query_start_loc[: num_decodes + 1], + max_seqlen_q=decode_max_query_len, seqused_k=attn_metadata.seq_lens[:num_decodes], max_seqlen_k=attn_metadata.max_seq_len, softmax_scale=self.scale, @@ -1410,120 +1350,157 @@ def forward( q_descale=None, k_descale=k_descale, v_descale=v_descale, + sinks=self.sinks, ) - elif rocm_aiter_ops.is_shuffle_kv_cache_enabled(): - _, num_heads, head_size = query.shape - num_blocks, block_size, num_kv_heads, _ = key_cache.shape - x = 16 // key_cache.element_size() - new_key_cache = key_cache.reshape( - num_blocks, num_kv_heads, head_size // x, block_size, x - ) - new_value_cache = value_cache.reshape( - num_blocks, num_kv_heads, block_size // x, head_size, x - ) + return + + # The ll4mi kernel in paged_attention_v1 requires + # HEAD_SIZE >= 16 * NWARPS (= 64 on ROCm with NWARPS=4). + # For smaller head sizes or sliding window attention, + # fall back to the unified_attention triton kernel which + # handles both correctly. + _MIN_HEAD_SIZE_FOR_LL4MI = 64 + use_unified_attention = self.head_size < _MIN_HEAD_SIZE_FOR_LL4MI + + if use_unified_attention: + k_descale, v_descale = self._get_kv_cache_descales(layer, num_decodes) + assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), ( + "unified_attention fallback with shuffle layout " + "is not supported yet." + ) + from aiter.ops.triton.unified_attention import ( + unified_attention, + ) - # This kernel derives block addresses arithmetically, so it - # only reads the right bytes when each side is a dense - # plane of pages. - assert new_key_cache.stride(0) == ( - num_kv_heads * head_size * block_size - ), ( - "paged_attention_common needs the K/V sides as dense " - "planes; the resolved KV cache layout interleaves them " - "within a block." - ) - num_seqs = attn_metadata.seq_lens.shape[0] - max_num_partitions = ( - attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1 - ) // _PARTITION_SIZE_ROCM - tmp_out = torch.empty( - (num_seqs, num_heads, max_num_partitions, head_size), - dtype=query.dtype, - device=query.device, - ) - exp_sums = torch.empty( - (num_seqs, num_heads, max_num_partitions), - dtype=torch.float32, - device=query.device, - ) - max_logits = torch.empty_like(exp_sums) - k_qscale = ( - layer._k_scale - if attn_metadata.k_scale is None - else attn_metadata.k_scale - ) - v_qscale = ( - layer._v_scale - if attn_metadata.v_scale is None - else attn_metadata.v_scale - ) - rocm_aiter_ops.paged_attention_common( - Q=query[:num_decode_tokens], - K=new_key_cache, - V=new_value_cache, - tmp_out=tmp_out, - max_logits=max_logits, - exp_sums=exp_sums, - max_seq_len=attn_metadata.max_seq_len, - block_tables=attn_metadata.block_table[:num_decodes], - context_lens=attn_metadata.seq_lens[:num_decodes], - block_tables_stride0=attn_metadata.block_table[ - :num_decodes - ].stride(0), - scale=self.scale, - K_QScale_hip=k_qscale, - V_QScale_hip=v_qscale, - K_QScale_asm=k_qscale, - V_QScale_asm=v_qscale, - out_=output[:num_decode_tokens], - kv_cache_dtype=self.kv_cache_dtype, - ) - else: - _, num_heads, head_size = query.shape - nbytes_per_qo_elem = torch.finfo(query.dtype).bits // 8 - num_seqs = attn_metadata.seq_lens.shape[0] - max_num_partitions = ( - attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1 - ) // _PARTITION_SIZE_ROCM - - workspace_buffer = torch.empty( - (num_seqs * num_heads * max_num_partitions * head_size) - * nbytes_per_qo_elem - + 2 * (num_seqs * num_heads * max_num_partitions) * 4, - dtype=torch.uint8, - device=output.device, - ) + decode_cu_seqlens_q = attn_metadata.query_start_loc[: num_decodes + 1] + unified_attention( + q=query[:num_decode_tokens], + k=key_cache, + v=value_cache, + out=output[:num_decode_tokens], + cu_seqlens_q=decode_cu_seqlens_q, + max_seqlen_q=1, + seqused_k=attn_metadata.seq_lens[:num_decodes], + max_seqlen_k=attn_metadata.max_seq_len, + softmax_scale=self.scale, + causal=True, + alibi_slopes=self.alibi_slopes, + window_size=self.sliding_window, + block_table=attn_metadata.block_table[:num_decodes], + softcap=self.logits_soft_cap, + q_descale=None, + k_descale=k_descale, + v_descale=v_descale, + ) + elif rocm_aiter_ops.is_shuffle_kv_cache_enabled(): + _, num_heads, head_size = query.shape + num_blocks, block_size, num_kv_heads, _ = key_cache.shape + x = 16 // key_cache.element_size() + new_key_cache = key_cache.reshape( + num_blocks, num_kv_heads, head_size // x, block_size, x + ) + new_value_cache = value_cache.reshape( + num_blocks, num_kv_heads, block_size // x, head_size, x + ) - # import so that aiter register the op to the namespace of - # torch.ops.aiter - import aiter # noqa: F401 - - torch.ops.aiter.paged_attention_v1( - output[:num_decode_tokens], - workspace_buffer, - query[:num_decode_tokens], - key_cache, - value_cache, - self.scale, - attn_metadata.block_table[:num_decodes], - attn_metadata.query_start_loc[:num_decodes], - attn_metadata.seq_lens[:num_decodes], - attn_metadata.max_seq_len, - self.alibi_slopes, - self.kv_cache_dtype, - "NHD", - self.logits_soft_cap, - layer._k_scale, - layer._v_scale, - None, - _PARTITION_SIZE_ROCM, - 1, - self.sliding_window[0] + 1, - ) - else: - raise NotImplementedError( - "Cascade attention is not implemented for ROCM AITER" - ) + # This kernel derives block addresses arithmetically, so it + # only reads the right bytes when each side is a dense + # plane of pages. + assert new_key_cache.stride(0) == ( + num_kv_heads * head_size * block_size + ), ( + "paged_attention_common needs the K/V sides as dense " + "planes; the resolved KV cache layout interleaves them " + "within a block." + ) + num_seqs = attn_metadata.seq_lens.shape[0] + max_num_partitions = ( + attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1 + ) // _PARTITION_SIZE_ROCM + tmp_out = torch.empty( + (num_seqs, num_heads, max_num_partitions, head_size), + dtype=query.dtype, + device=query.device, + ) + exp_sums = torch.empty( + (num_seqs, num_heads, max_num_partitions), + dtype=torch.float32, + device=query.device, + ) + max_logits = torch.empty_like(exp_sums) + k_qscale = ( + layer._k_scale + if attn_metadata.k_scale is None + else attn_metadata.k_scale + ) + v_qscale = ( + layer._v_scale + if attn_metadata.v_scale is None + else attn_metadata.v_scale + ) + rocm_aiter_ops.paged_attention_common( + Q=query[:num_decode_tokens], + K=new_key_cache, + V=new_value_cache, + tmp_out=tmp_out, + max_logits=max_logits, + exp_sums=exp_sums, + max_seq_len=attn_metadata.max_seq_len, + block_tables=attn_metadata.block_table[:num_decodes], + context_lens=attn_metadata.seq_lens[:num_decodes], + block_tables_stride0=attn_metadata.block_table[:num_decodes].stride( + 0 + ), + scale=self.scale, + K_QScale_hip=k_qscale, + V_QScale_hip=v_qscale, + K_QScale_asm=k_qscale, + V_QScale_asm=v_qscale, + out_=output[:num_decode_tokens], + kv_cache_dtype=self.kv_cache_dtype, + ) + else: + _, num_heads, head_size = query.shape + nbytes_per_qo_elem = torch.finfo(query.dtype).bits // 8 + num_seqs = attn_metadata.seq_lens.shape[0] + max_num_partitions = ( + attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1 + ) // _PARTITION_SIZE_ROCM + + workspace_buffer = torch.empty( + (num_seqs * num_heads * max_num_partitions * head_size) + * nbytes_per_qo_elem + + 2 * (num_seqs * num_heads * max_num_partitions) * 4, + dtype=torch.uint8, + device=output.device, + ) + + # import so that aiter register the op to the namespace of + # torch.ops.aiter + import aiter # noqa: F401 + + torch.ops.aiter.paged_attention_v1( + output[:num_decode_tokens], + workspace_buffer, + query[:num_decode_tokens], + key_cache, + value_cache, + self.scale, + attn_metadata.block_table[:num_decodes], + attn_metadata.query_start_loc[:num_decodes], + attn_metadata.seq_lens[:num_decodes], + attn_metadata.max_seq_len, + self.alibi_slopes, + self.kv_cache_dtype, + "NHD", + self.logits_soft_cap, + layer._k_scale, + layer._v_scale, + None, + _PARTITION_SIZE_ROCM, + 1, + self.sliding_window[0] + 1, + ) return output diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index 0d4fd5519484..9f8e9e9ab12b 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -129,10 +129,6 @@ def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: return (KVCacheLayout.LBHNC,) return (KVCacheLayout.LBHNC, KVCacheLayout.LHBNC) - @staticmethod - def use_cascade_attention(*args, **kwargs) -> bool: - return False - @staticmethod def get_builder_cls() -> type["RocmAiterUnifiedAttentionMetadataBuilder"]: return RocmAiterUnifiedAttentionMetadataBuilder @@ -237,8 +233,6 @@ def forward( # Profiling run. return output.fill_(0) - assert attn_metadata.use_cascade is False - # IMPORTANT! # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead diff --git a/vllm/v1/attention/backends/rocm_attn.py b/vllm/v1/attention/backends/rocm_attn.py index 3ccc1e5d5783..471bb9f05848 100644 --- a/vllm/v1/attention/backends/rocm_attn.py +++ b/vllm/v1/attention/backends/rocm_attn.py @@ -58,16 +58,8 @@ class RocmAttentionMetadata: block_table: torch.Tensor slot_mapping: torch.Tensor - # For cascade attention. - use_cascade: bool - common_prefix_len: int - cu_prefix_query_lens: torch.Tensor | None - prefix_kv_lens: torch.Tensor | None - suffix_kv_lens: torch.Tensor | None - # Optional aot scheduling scheduler_metadata: torch.Tensor | None = None - prefix_scheduler_metadata: torch.Tensor | None = None # DFlash drafting sets this to False via CommonAttentionMetadata. causal: bool = True @@ -124,23 +116,6 @@ def build( block_table_tensor = common_attn_metadata.block_table_tensor slot_mapping = common_attn_metadata.slot_mapping - use_cascade = common_prefix_len > 0 - - if use_cascade: - cu_prefix_query_lens = torch.tensor( - [0, num_actual_tokens], dtype=torch.int32, device=self.device - ) - prefix_kv_lens = torch.tensor( - [common_prefix_len], dtype=torch.int32, device=self.device - ) - suffix_kv_lens = common_attn_metadata.seq_lens.cpu() - common_prefix_len - suffix_kv_lens = suffix_kv_lens.to(self.device) - else: - cu_prefix_query_lens = None - prefix_kv_lens = None - suffix_kv_lens = None - prefix_scheduler_metadata = None - attn_metadata = RocmAttentionMetadata( num_actual_tokens=num_actual_tokens, max_query_len=max_query_len, @@ -149,12 +124,6 @@ def build( seq_lens=seq_lens, block_table=block_table_tensor, slot_mapping=slot_mapping, - use_cascade=use_cascade, - common_prefix_len=common_prefix_len, - cu_prefix_query_lens=cu_prefix_query_lens, - prefix_kv_lens=prefix_kv_lens, - suffix_kv_lens=suffix_kv_lens, - prefix_scheduler_metadata=prefix_scheduler_metadata, causal=common_attn_metadata.causal, ) return attn_metadata @@ -267,10 +236,6 @@ def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: # supported, just without the native kernels. return (KVCacheLayout.LHBNC, KVCacheLayout.LBHNC) - @staticmethod - def use_cascade_attention(*args, **kwargs) -> bool: - return False - @staticmethod def get_builder_cls() -> type["RocmAttentionMetadataBuilder"]: return RocmAttentionMetadataBuilder @@ -418,8 +383,6 @@ def forward( # Profiling run. return output.fill_(0) - assert attn_metadata.use_cascade is False - # IMPORTANT! # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead diff --git a/vllm/v1/attention/backends/rocm_segmented_attn.py b/vllm/v1/attention/backends/rocm_segmented_attn.py index c541eed45166..4725db917b70 100644 --- a/vllm/v1/attention/backends/rocm_segmented_attn.py +++ b/vllm/v1/attention/backends/rocm_segmented_attn.py @@ -271,7 +271,6 @@ def forward( workspace = self._get_workspace(query.device) if attn_metadata is None: return output.fill_(0) - assert not attn_metadata.use_cascade num_actual_tokens = attn_metadata.num_actual_tokens key_cache, value_cache = self._split_kv_cache(kv_cache) diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index 095163f14423..a95b0c1d0e13 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -80,16 +80,8 @@ class TritonAttentionMetadata: causal: bool | torch.Tensor - # For cascade attention. - use_cascade: bool - common_prefix_len: int - cu_prefix_query_lens: torch.Tensor | None - prefix_kv_lens: torch.Tensor | None - suffix_kv_lens: torch.Tensor | None - # Optional aot scheduling scheduler_metadata: torch.Tensor | None = None - prefix_scheduler_metadata: torch.Tensor | None = None mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None mm_prefix_range_tensor: torch.Tensor | None = None rswa_prefix_lens: torch.Tensor | None = None @@ -230,23 +222,6 @@ def build( block_table_tensor = common_attn_metadata.block_table_tensor slot_mapping = common_attn_metadata.slot_mapping - use_cascade = common_prefix_len > 0 - - if use_cascade: - cu_prefix_query_lens = torch.tensor( - [0, num_actual_tokens], dtype=torch.int32, device=self.device - ) - prefix_kv_lens = torch.tensor( - [common_prefix_len], dtype=torch.int32, device=self.device - ) - suffix_kv_lens = common_attn_metadata.seq_lens.cpu() - common_prefix_len - suffix_kv_lens = suffix_kv_lens.to(self.device) - else: - cu_prefix_query_lens = None - prefix_kv_lens = None - suffix_kv_lens = None - prefix_scheduler_metadata = None - attn_metadata = TritonAttentionMetadata( num_actual_tokens=num_actual_tokens, max_query_len=max_query_len, @@ -256,12 +231,6 @@ def build( block_table=block_table_tensor, slot_mapping=slot_mapping, causal=common_attn_metadata.causal, - use_cascade=use_cascade, - common_prefix_len=common_prefix_len, - cu_prefix_query_lens=cu_prefix_query_lens, - prefix_kv_lens=prefix_kv_lens, - suffix_kv_lens=suffix_kv_lens, - prefix_scheduler_metadata=prefix_scheduler_metadata, seq_threshold_3D=self.seq_threshold_3D, num_par_softmax_segments=self.num_par_softmax_segments, softmax_segm_output=self.softmax_segm_output, @@ -378,10 +347,6 @@ def supports_batch_invariance(cls) -> bool: def get_impl_cls() -> type["TritonAttentionImpl"]: return TritonAttentionImpl - @staticmethod - def use_cascade_attention(*args, **kwargs) -> bool: - return False - @staticmethod def get_builder_cls() -> type["TritonAttentionMetadataBuilder"]: return TritonAttentionMetadataBuilder @@ -622,8 +587,6 @@ def forward( # Profiling run. return output.fill_(0) - assert attn_metadata.use_cascade is False - # IMPORTANT! # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead diff --git a/vllm/v1/attention/backends/triton_attn_diffkv.py b/vllm/v1/attention/backends/triton_attn_diffkv.py index 932f483af37b..08e16c3e64aa 100644 --- a/vllm/v1/attention/backends/triton_attn_diffkv.py +++ b/vllm/v1/attention/backends/triton_attn_diffkv.py @@ -181,10 +181,6 @@ def forward( if attn_metadata is None: return output.fill_(0) - assert attn_metadata.use_cascade is False, ( - "Cascade attention not supported for TritonAttentionDiffKVImpl" - ) - num_actual_tokens = attn_metadata.num_actual_tokens head_size_qk = self.head_size head_size_v = TritonAttentionDiffKVBackend.head_size_v diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index b6defc52e847..7d4f2d5a2b44 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -426,23 +426,6 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: blocks.extend(manager.pop_blocks_for_free(request_id)) return blocks - def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]: - """Get the number of common prefix blocks for all requests with allocated - KV cache for each kv cache group. - - Args: - running_request_id: The request ID of any running request, used to - identify the common prefix blocks. - - Returns: - list[int]: The number of common prefix blocks for each kv cache group. - - """ - return [ - manager.get_num_common_prefix_blocks(running_request_id) - for manager in self.single_type_managers - ] - def remove_skipped_blocks( self, request_id: str, @@ -527,9 +510,6 @@ def __init__( ) self.num_single_type_manager = len(self.single_type_managers) - def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]: - return [0] * self.num_single_type_manager - def find_longest_cache_hit( self, block_hashes: list[BlockHash], diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 37ceb4e474cf..4f472170a874 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -689,39 +689,9 @@ def reset_prefix_cache(self) -> bool: return True def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]: - """Calculate the number of common prefix blocks for each kv cache group. - - The function selects a running request and iterates through its blocks. - A block is considered a common prefix block if ALL requests with - allocated KV cache share it (i.e., ref_cnt equals the number of entries - in req_to_blocks). - - NOTE(woosuk): The number of requests with allocated KV cache is **greater - than or equal to** the number of requests scheduled in the current step. - This is because having allocated KV cache only indicates that: - 1. The request has not yet finished, and - 2. The request holds its blocks unfreed. - - While all scheduled requests must have allocated KV cache, the inverse - is not necessarily true. There may be requests with allocated KV cache - that are not scheduled in the current step. - - This can result in an edge case where the number of common prefix blocks - is 0, even though all scheduled requests share a common prefix. This - occurs because there may be unscheduled requests that do not share the - common prefix. Currently, this case cannot be easily detected, so the - function returns 0 in such cases. - - Args: - running_request_id: The request ID of any running request, used to - identify the common prefix blocks. - - Returns: - list[int]: The number of common prefix blocks for each kv cache - group. - - """ - return self.coordinator.get_num_common_prefix_blocks(running_request_id) + """Deprecated: cascade attention has been removed. Always returns 0 + for each KV cache group. Kept for out-of-tree schedulers.""" + return [0] * self.num_kv_cache_groups def take_events(self) -> list[KVCacheEvent]: """Take the KV cache events from the block pool. diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index a5af5a3d5755..9ad0d4e65865 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, field from functools import cached_property from typing import TYPE_CHECKING @@ -253,9 +253,6 @@ class SchedulerOutput: # E.g., if a request has [0, 1], it could mean the vision encoder needs # to process that the request's 0-th and 1-th images in the current step. scheduled_encoder_inputs: dict[str, list[int]] - # Number of common prefix blocks for all requests in each KV cache group. - # This can be used for cascade attention. - num_common_prefix_blocks: list[int] # Request IDs that are finished in between the previous and the current # steps. This is used to notify the workers about the finished requests @@ -311,6 +308,10 @@ class SchedulerOutput: # Number of spec tokens to schedule for the next step. num_spec_tokens_to_schedule: int = 0 + # Deprecated: cascade attention has been removed, so this is always empty. + # Kept for out-of-tree schedulers and workers. + num_common_prefix_blocks: list[int] = field(default_factory=list) + @classmethod def make_empty(cls) -> "SchedulerOutput": return cls( @@ -320,7 +321,6 @@ def make_empty(cls) -> "SchedulerOutput": total_num_scheduled_tokens=0, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 2603fa10a75e..bcb6def9383c 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1414,16 +1414,6 @@ def skip_request(from_queue: RequestQueue) -> None: scheduled_running_reqs ) <= len(self.running) - # Get the longest common prefix among all requests in the running queue. - # This can be potentially used for cascade attention. - num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups) - with record_function_or_nullcontext("schedule: get_num_common_prefix_blocks"): - if self.running: - any_request_id = self.running[0].request_id - num_common_prefix_blocks = ( - self.kv_cache_manager.get_num_common_prefix_blocks(any_request_id) - ) - # Construct the scheduler output. if self.use_v2_model_runner: scheduled_new_reqs.extend(scheduled_resumed_reqs) @@ -1516,7 +1506,6 @@ def skip_request(from_queue: RequestQueue) -> None: scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, scheduled_encoder_inputs=scheduled_encoder_inputs, scheduled_encoder_input_stats=scheduled_encoder_input_stats, - num_common_prefix_blocks=num_common_prefix_blocks, preempted_req_ids=self.reset_preempted_req_ids, # finished_req_ids is an existing state in the scheduler, # instead of being newly scheduled in this step. diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 218d59f430c9..2e4fe55b96c2 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -588,21 +588,6 @@ def free(self, request_id: str) -> None: # Free blocks in reverse order so that the tail blocks are freed first. self.block_pool.free_blocks(reversed(self.pop_blocks_for_free(request_id))) - @abstractmethod - def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - """Get the number of common prefix blocks for all requests with allocated - KV cache. - - Args: - running_request_id: The request ID. - - Returns: - The number of common prefix blocks for all requests with allocated - KV cache. - - """ - raise NotImplementedError - @classmethod @abstractmethod def find_longest_cache_hit( @@ -898,16 +883,6 @@ def _cache_partial_tail_block( block_size=self.block_size, ) - def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - blocks = self.req_to_blocks[running_request_id] - num_common_blocks = 0 - for block in blocks: - if block.ref_cnt == len(self.req_to_blocks): - num_common_blocks += 1 - else: - break - return num_common_blocks - class RSWAManager(FullAttentionManager): """KV cache manager for Reference Sliding Window Attention (R-SWA). @@ -1194,14 +1169,6 @@ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: num_computed_tokens - self.sliding_window + 1 - self.extra_retained_tokens, ) - def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - """NOTE(Chen): The prefix blocks are null blocks for sliding window layers. - So it's not correct to count ref_cnt like FullAttentionManager. Return - 0 here for correctness. Need to support cascade attention + sliding - window in the future. - """ - return 0 - class CircularBufferManager(FullAttentionManager): """Claims the ring's single block per request; prefix caching disabled.""" @@ -1289,9 +1256,6 @@ def remove_skipped_blocks( ) -> None: return - def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - return 0 - def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: return 0 @@ -1447,10 +1411,6 @@ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: ) * self.attention_chunk_size return num_skipped_tokens - def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - """Cascade attention is not supported by chunked local attention.""" - return 0 - class MambaManager(SingleTypeKVCacheManager): supports_fine_grained_hash_lookup: ClassVar[bool] = True @@ -1680,10 +1640,6 @@ def remove_skipped_blocks( self.block_pool.free_blocks([blocks[last_state_block_idx]]) blocks[last_state_block_idx] = self._null_block - def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - """Cascade attention is not supported by mamba""" - return 0 - def _needs_internal_checkpoint( self, request_id: str, @@ -2192,11 +2148,6 @@ def cache_blocks( # requests, so this method is not relevant. raise ValueError("Should not be called as prefix caching is disabled.") - def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - # Cross-attention blocks contain request-specific encoder states - # and are not shared between different requests - return 0 - @classmethod def find_longest_cache_hit( cls, @@ -2394,9 +2345,6 @@ def cache_blocks( ) -> None: return None - def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - return 0 - @classmethod def find_longest_cache_hit( cls, diff --git a/vllm/v1/cudagraph_dispatcher.py b/vllm/v1/cudagraph_dispatcher.py index 0c1fb0b2c4c9..d6abba301aec 100644 --- a/vllm/v1/cudagraph_dispatcher.py +++ b/vllm/v1/cudagraph_dispatcher.py @@ -254,8 +254,8 @@ def dispatch( all modes are allowed. invalid_modes: Set of cudagraph modes to exclude. Subtracted from valid_modes to compute allowed modes. (e.g., {FULL} for - features like cascade attention not supported by full - cudagraphs). None means no modes are excluded. + encoder-decoder batches with encoder outputs). None means no + modes are excluded. """ allowed_modes = valid_modes or CUDAGraphMode.valid_runtime_modes() diff --git a/vllm/v1/worker/cpu_model_runner.py b/vllm/v1/worker/cpu_model_runner.py index 32fa01bd27ec..d8df0ca1033e 100644 --- a/vllm/v1/worker/cpu_model_runner.py +++ b/vllm/v1/worker/cpu_model_runner.py @@ -32,7 +32,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # Note: speculative decoding is now supported on CPU with C++ native impls self.use_cuda_graph = False - self.cascade_attn_enabled = False self._postprocess_tensors() @@ -109,9 +108,6 @@ def initialize_kv_cache( elif self.speculative_config.uses_draft_model(): logger.info("Draft model KV cache initialized for CPU backend") - def _init_device_properties(self) -> None: - pass - def _sync_device(self) -> None: pass diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 894e2e5b6bd9..b81ddf6d7335 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -106,7 +106,6 @@ def run_mixed_prefill_decode_warmup( prefill_token_ids = list(range(prefill_len)) kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups - num_kv_cache_groups = len(kv_cache_groups) block_count = _warmup_block_counter(model_runner) kv_cache_specs = [g.kv_cache_spec for g in kv_cache_groups] decode_prefill_block_counts = [ @@ -166,7 +165,6 @@ def _alloc_blocks(num_blocks: int) -> list[int]: decode_req_id: decode_prompt_len, } decode_prefill_output.total_num_scheduled_tokens = decode_prompt_len - decode_prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups decode_new_blocks = tuple(_alloc_blocks(n) for n in decode_block_deltas) cached_decode_req = CachedRequestData.make_empty() @@ -198,7 +196,6 @@ def _alloc_blocks(num_blocks: int) -> list[int]: prefill_req_id: prefill_len, } mixed_output.total_num_scheduled_tokens = num_tokens - mixed_output.num_common_prefix_blocks = [0] * num_kv_cache_groups cleanup_output = SchedulerOutput.make_empty() cleanup_output.finished_req_ids = {decode_req_id, prefill_req_id} @@ -272,7 +269,6 @@ def _warmup_kernels( decode_len = prompt_len + num_decode_steps * decode_query_len kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups - num_kv_cache_groups = len(kv_cache_groups) # Encoder-decoder models: give each warmup request a dummy encoder input so # cross-attention warms up over a realistic, non-empty key sequence. @@ -361,7 +357,6 @@ def _alloc_blocks(num_blocks: int) -> list[int]: prefill_output.scheduled_new_reqs = new_reqs prefill_output.num_scheduled_tokens = {rid: prompt_len for rid in req_ids} prefill_output.total_num_scheduled_tokens = prompt_len * num_reqs - prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups # Disable KV connector for warmup run. model_runner.kv_connector.set_disabled(True) @@ -424,7 +419,6 @@ def _run_decode_step(indices: list[int], spec_flags: list[bool]) -> None: decode_output.total_num_scheduled_tokens = sum( step_num_scheduled_tokens.values() ) - decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups worker_execute_model(decode_output) worker_sample_tokens(None) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 669b314d4cac..e19c1b84f457 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -115,7 +115,6 @@ from vllm.utils.math_utils import cdiv from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib from vllm.utils.nvtx_pytorch_hooks import PytHooks -from vllm.utils.platform_utils import num_compute_units from vllm.utils.torch_utils import ( PIN_MEMORY, async_tensor_h2d, @@ -145,15 +144,12 @@ from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher from vllm.v1.kv_cache_interface import ( AttentionSpec, - ChunkedLocalAttentionSpec, CrossAttentionSpec, EncoderOnlyAttentionSpec, - FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, KVCacheSpecKind, - SlidingWindowSpec, UniformTypeKVCacheSpecs, get_kv_cache_spec_kind, ) @@ -530,13 +526,11 @@ def __init__( ) # Model-related. - self.num_query_heads = model_config.get_num_attention_heads(parallel_config) self.inputs_embeds_size = model_config.get_inputs_embeds_size() - # Only relevant for models using ALiBi (e.g, MPT) - self.use_alibi = model_config.uses_alibi - - self.cascade_attn_enabled = not self.model_config.disable_cascade_attn self.is_mm_prefix_lm = self.model_config.is_mm_prefix_lm + # Deprecated: cascade attention has been removed. Kept for + # out-of-tree model runners. + self.cascade_attn_enabled = False # Multi-modal data support self.mm_registry = MULTIMODAL_REGISTRY @@ -769,9 +763,6 @@ def __init__( else: self.cudagraph_batch_sizes = [] - # Cache the device properties. - self._init_device_properties() - # Encoder timing registry for observability self.encoder_timing_registry: dict[str, EncoderTimingStats] = {} self._encoder_timing_lock = threading.Lock() @@ -1151,11 +1142,6 @@ def _zero_block_ids(self, block_ids: list[int]) -> None: if hasattr(self, "_kv_block_zeroer"): self._kv_block_zeroer.zero_block_ids(block_ids) - # Note: used for model runner override. - def _init_device_properties(self) -> None: - """Initialize attributes from torch.cuda.get_device_properties""" - self.num_sms = num_compute_units(self.device.index) - # Note: used for model runner override. def _sync_device(self) -> None: torch.accelerator.synchronize() @@ -2273,7 +2259,6 @@ def _build_attention_metadata( use_spec_decode: bool = False, for_cudagraph_capture: bool = False, num_scheduled_tokens: dict[str, int] | None = None, - cascade_attn_prefix_lens: list[list[int]] | None = None, slot_mappings: dict[int, torch.Tensor] | None = None, ) -> tuple[PerLayerAttnMetadata, CommonAttentionMetadata | None]: """Returns: @@ -2472,12 +2457,6 @@ def _build_attn_group_metadata( kv_cache_spec = kv_cache_spec.kv_cache_specs[attn_group.layer_names[0]] cache_key = (kv_cache_spec, type(builder)) - cascade_attn_prefix_len = ( - cascade_attn_prefix_lens[kv_cache_gid][attn_gid] - if cascade_attn_prefix_lens - else 0 - ) - extra_attn_metadata_args = {} if use_spec_decode and isinstance( builder, @@ -2517,7 +2496,6 @@ def _build_attn_group_metadata( attn_metadata_i = attn_group.build_metadata( common_attn_metadata, ubid or 0, - common_prefix_len=cascade_attn_prefix_len, **extra_attn_metadata_args, ) if builder.supports_update_block_table: @@ -2595,144 +2573,6 @@ def _build_attn_group_metadata( return attn_metadata, spec_decode_common_attn_metadata - def _compute_cascade_attn_prefix_lens( - self, - num_scheduled_tokens: np.ndarray, - num_computed_tokens: np.ndarray, - num_common_prefix_blocks: list[int], - ) -> list[list[int]] | None: - """Returns: - Optional[cascade_attn_prefix_lens] - cascade_attn_prefix_lens is 2D: - ``[kv_cache_group_id][attn_group_idx]``, - None if we should not use cascade attention - - """ - use_cascade_attn = False - num_kv_cache_groups = len(self.kv_cache_config.kv_cache_groups) - cascade_attn_prefix_lens: list[list[int]] = [ - [] for _ in range(num_kv_cache_groups) - ] - - for kv_cache_gid in range(num_kv_cache_groups): - for attn_group in self.attn_groups[kv_cache_gid]: - if isinstance(attn_group.kv_cache_spec, EncoderOnlyAttentionSpec): - cascade_attn_prefix_len = 0 - else: - # 0 if cascade attention should not be used - cascade_attn_prefix_len = self._compute_cascade_attn_prefix_len( - num_scheduled_tokens, - num_computed_tokens, - num_common_prefix_blocks[kv_cache_gid], - attn_group.kv_cache_spec, - attn_group.get_metadata_builder(), - ) - cascade_attn_prefix_lens[kv_cache_gid].append(cascade_attn_prefix_len) - use_cascade_attn |= cascade_attn_prefix_len > 0 - - return cascade_attn_prefix_lens if use_cascade_attn else None - - def _compute_cascade_attn_prefix_len( - self, - num_scheduled_tokens: np.ndarray, - num_computed_tokens: np.ndarray, - num_common_prefix_blocks: int, - kv_cache_spec: KVCacheSpec, - attn_metadata_builder: AttentionMetadataBuilder, - ) -> int: - """Compute the length of the common prefix for cascade attention. - - NOTE(woosuk): The common prefix length returned by this function - represents the length used specifically for cascade attention, not the - actual number of tokens shared between requests. When cascade attention - is disabled (use_cascade=False), this function returns 0 even if - requests share common tokens. Additionally, the common prefix length is - truncated to a multiple of the block size and may be further truncated - due to implementation details explained below. - - Args: - num_scheduled_tokens: Number of tokens scheduled per request. - num_computed_tokens: Number of tokens already computed per request. - num_common_prefix_blocks: Number of shared KV cache blocks. - kv_cache_spec: The KV cache spec of the group being considered. - attn_metadata_builder: Builder used to decide whether cascade - attention can be used. - - Returns: - int: Length of common prefix in tokens. - - """ - common_prefix_len = num_common_prefix_blocks * kv_cache_spec.block_size - if common_prefix_len == 0: - # Common case. - return 0 - - # NOTE(woosuk): Cascade attention uses two attention kernels: one - # for the common prefix and the other for the rest. For the first - # kernel, we concatenate all the query tokens (possibly from - # different requests) and treat them as if they are from the same - # request. Then, we use bi-directional attention to process the - # common prefix in the KV cache. Importantly, this means that the - # first kernel does not do any masking. - - # Consider the following example: - # Request 1's input query: [D, E, X] - # Request 1's kv cache: [A, B, C, D, E, X] - # Request 1's num_computed_tokens: 3 (i.e., [A, B, C]) - # Request 2's input query: [E, Y] - # Request 2's kv cache: [A, B, C, D, E, Y] - # Request 2's num_computed_tokens: 4 (i.e., [A, B, C, D]) - - # If we use [A, B, C, D, E] as the common prefix, then the - # first kernel will compute the bi-directional attention between - # input query [D, E, X, E, Y] and common prefix [A, B, C, D, E]. - # However, this is wrong because D in Request 1 should not attend to - # E in the common prefix (i.e., we need masking). - # To avoid this, [A, B, C, D] should be the common prefix. - # That is, the common prefix should be capped by the minimum - # num_computed_tokens among the requests, and plus one to include - # the first token of the query. - - # In practice, we use [A, B, C] as the common prefix, instead of - # [A, B, C, D] (i.e., the common prefix is capped by the minimum - # num_computed_tokens, without plus one). - # This is because of an implementation detail: We want to always - # use two kernels for cascade attention. Let's imagine: - # Request 3's input query: [D] - # Request 3's kv cache: [A, B, C, D] - # Request 3's num_computed_tokens: 3 (i.e., [A, B, C]) - # If we use [A, B, C, D] as the common prefix for Request 1-3, - # then Request 3 will be processed only by the first kernel, - # and the second kernel will get an empty input. While this is not - # a fundamental problem, our current implementation does not support - # this case. - common_prefix_len = min(common_prefix_len, num_computed_tokens.min()) - # common_prefix_len should be a multiple of the block size. - common_prefix_len = ( - common_prefix_len // kv_cache_spec.block_size * kv_cache_spec.block_size - ) - use_sliding_window = isinstance(kv_cache_spec, SlidingWindowSpec) or ( - isinstance(kv_cache_spec, FullAttentionSpec) - and kv_cache_spec.sliding_window is not None - ) - use_local_attention = isinstance(kv_cache_spec, ChunkedLocalAttentionSpec) or ( - isinstance(kv_cache_spec, FullAttentionSpec) - and kv_cache_spec.attention_chunk_size is not None - ) - assert isinstance(kv_cache_spec, AttentionSpec) - use_cascade = attn_metadata_builder.use_cascade_attention( - common_prefix_len=common_prefix_len, - query_lens=num_scheduled_tokens, - num_query_heads=self.num_query_heads, - num_kv_heads=kv_cache_spec.num_kv_heads, - use_alibi=self.use_alibi, - use_sliding_window=use_sliding_window, - use_local_attention=use_local_attention, - num_sms=self.num_sms, - dcp_world_size=self.dcp_world_size, - ) - return common_prefix_len if use_cascade else 0 - def _calc_mrope_positions(self, scheduler_output: "SchedulerOutput"): mrope_pos_ptr = 0 for index, req_id in enumerate(self.input_batch.req_ids): @@ -3881,7 +3721,8 @@ def _determine_batch_execution_and_padding( num_reqs: int, num_scheduled_tokens_np: np.ndarray, max_num_scheduled_tokens: int, - use_cascade_attn: bool, + # Deprecated and ignored. Kept for out-of-tree model runners. + use_cascade_attn: bool = False, allow_microbatching: bool = True, force_eager: bool = False, # For cudagraph capture TODO(lucas): Refactor how we capture cudagraphs (will @@ -3929,7 +3770,7 @@ def dispatch_cudagraph(num_tokens, disable_full=False, valid_modes=None): ) cudagraph_mode, batch_descriptor = dispatch_cudagraph( - num_tokens, disable_full=use_cascade_attn or has_encoder_output + num_tokens, disable_full=has_encoder_output ) num_tokens_padded = batch_descriptor.num_tokens @@ -4181,16 +4022,6 @@ def execute_model( self._prepare_inputs(scheduler_output, num_scheduled_tokens_np) ) - cascade_attn_prefix_lens = None - # Disable cascade attention when using microbatching (DBO) - if self.cascade_attn_enabled and not self.parallel_config.use_ubatching: - # Pre-compute cascade attention prefix lengths - cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens( - num_scheduled_tokens_np, - self.input_batch.num_computed_tokens_cpu[:num_reqs], - scheduler_output.num_common_prefix_blocks, - ) - ( cudagraph_mode, batch_desc, @@ -4202,7 +4033,7 @@ def execute_model( num_reqs=num_reqs, num_scheduled_tokens_np=num_scheduled_tokens_np, max_num_scheduled_tokens=max_num_scheduled_tokens, - use_cascade_attn=cascade_attn_prefix_lens is not None, + use_cascade_attn=False, num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs), allow_microbatching=self._allow_microbatching( num_reqs, num_scheduled_tokens_np @@ -4319,7 +4150,6 @@ def execute_model( max_num_sampled_tokens=max_num_sampled_tokens, use_spec_decode=use_spec_decode, num_scheduled_tokens=scheduler_output.num_scheduled_tokens, - cascade_attn_prefix_lens=cascade_attn_prefix_lens, slot_mappings=slot_mappings_by_group, ) ) diff --git a/vllm/v1/worker/xpu_model_runner.py b/vllm/v1/worker/xpu_model_runner.py index 15fa8e1e763f..83ff08371efa 100644 --- a/vllm/v1/worker/xpu_model_runner.py +++ b/vllm/v1/worker/xpu_model_runner.py @@ -24,8 +24,6 @@ def __init__( ): with _torch_cuda_wrapper(): super().__init__(vllm_config, device) - # FIXME: To be verified. - self.cascade_attn_enabled = False class XPUModelRunnerV2(GPUModelRunnerV2): From c9e62f13c2b52619d8bc99f0c76cf04e77bff3c0 Mon Sep 17 00:00:00 2001 From: josh-schell-amd Date: Mon, 12 Oct 2026 13:04:37 -0500 Subject: [PATCH 324/624] [CI][ROCm] Add AITER nightly overlay scripts (#60883) Signed-off-by: josh schell Co-authored-by: Claude Sonnet 5.5 Co-authored-by: Micah Williamson --- .../scripts/rocm/aiter-nightly-overlay.sh | 112 ++++++++ .../rocm/select_aiter_nightly_wheel.py | 260 ++++++++++++++++++ 2 files changed, 372 insertions(+) create mode 100755 .buildkite/scripts/rocm/aiter-nightly-overlay.sh create mode 100755 .buildkite/scripts/rocm/select_aiter_nightly_wheel.py diff --git a/.buildkite/scripts/rocm/aiter-nightly-overlay.sh b/.buildkite/scripts/rocm/aiter-nightly-overlay.sh new file mode 100755 index 000000000000..528247964baf --- /dev/null +++ b/.buildkite/scripts/rocm/aiter-nightly-overlay.sh @@ -0,0 +1,112 @@ +#!/bin/bash +# Install the AITER nightly wheel over this build's ci_base. +# +# Runs in CI when AITER_NIGHTLY=1, once the stock ci_base it builds on is +# published. In order: +# 1. Select the newest nightly wheel that fits the stock ci_base +# (select_aiter_nightly_wheel.py, beside this script, run inside that image) +# 2. Build a new image: pip install that wheel over the stock ci_base +# 3. Push the new image as this build's ci_base-build tag, replacing the stock +# one, so later steps and native jobs use it; record its digest +# +# Buildkite meta-data, a key-value store shared by every job in this build: +# reads rocm-ci-base-image the stock ci_base +# r/w aiter-nightly-wheel-url the wheel step 1 chose; a retry reads it +# back and installs that same wheel +# writes aiter-nightly-ci-base the new image as tag@digest; the vLLM +# image build uses it as its ci_base +# +# Exit status: +# 10 no wheel fits this ci_base +# 11 the wheel does not install +# Anything else is unexpected and passed through as is. +# Buildkite retries this step once on exit status 1 or 2, which it treats as +# transient Docker/registry failures (amd_infra_retry in the template). +# 10 and 11 sit outside that range: they are real AITER failures, and a retry +# would only fail the same way. + +set -euo pipefail + +SELECT_WHEEL="$(dirname "${BASH_SOURCE[0]}")/select_aiter_nightly_wheel.py" +BUILD_TAG="rocm/vllm-dev:ci_base-build-${BUILDKITE_BUILD_ID:?BUILDKITE_BUILD_ID is required}" + +# Scratch space for the build context and captured errors. The trap deletes it +# however the script exits: success, fail, or set -e. +scratch="$(mktemp -d)" +trap 'rm -rf "${scratch}"' EXIT +errors="${scratch}/errors" + +fail() { + local status="$1" title="$2" detail="$3" + echo "+++ :x: AITER nightly: ${title}" + echo "${detail}" >&2 + # shellcheck disable=SC2016 # the backticks are a literal markdown fence + printf ':x: **AITER nightly: %s**\n\n```\n%s\n```\n' "${title}" "${detail}" \ + | buildkite-agent annotate --style error --context aiter-nightly + exit "${status}" +} + +# Set by ensure-ci-base-amd; this step's depends_on guarantees it is there. +ci_base="$(buildkite-agent meta-data get rocm-ci-base-image)" +# Set below on the first attempt. Meta-data is per build, so a retry reuses +# that wheel, not a newer nightly. +pinned="$(buildkite-agent meta-data get aiter-nightly-wheel-url --default "")" + +echo "--- :docker: Pulling ${ci_base}" +docker pull "${ci_base}" + +echo "--- :mag: Selecting the AITER nightly for ${ci_base}" +set +e +out="$(docker run --rm -i --entrypoint python3 "${ci_base}" \ + -P - ${pinned:+--wheel-url "${pinned}"} \ + < "${SELECT_WHEEL}" 2> "${errors}")" +status=$? +set -e +cat "${errors}" >&2 +case "${status}" in + 0) ;; + 10) fail 10 "no wheel for this image" "${out}" ;; + *) fail "${status}" "could not select a wheel" "$(tail -n 20 "${errors}")" ;; +esac +# The first line is the wheel's URL; the rest, the annotation naming it. +wheel_url="$(head -n 1 <<< "${out}")" +note="$(tail -n +2 <<< "${out}")" +buildkite-agent meta-data set aiter-nightly-wheel-url "${wheel_url}" + +cat > "${scratch}/Dockerfile" <<'EOF' +ARG CI_BASE_IMAGE +FROM ${CI_BASE_IMAGE} +ARG AITER_WHEEL_URL +RUN python3 -m pip install --no-cache-dir "$AITER_WHEEL_URL" +EOF + +echo "--- :docker: Installing ${wheel_url##*/}" +if ! docker buildx build \ + --builder default \ + --provenance=false \ + --progress plain \ + --build-arg "CI_BASE_IMAGE=${ci_base}" \ + --build-arg "AITER_WHEEL_URL=${wheel_url}" \ + -t "${BUILD_TAG}" \ + "${scratch}"; then + fail 11 "the wheel does not install" \ + "Installing ${wheel_url##*/} over ${ci_base} failed; see this step's log." +fi + +docker push "${BUILD_TAG}" +# Handed off as tag@digest: build-test-image.sh rejects a ci_base that is not +# digest-pinned. +digest="$(docker buildx imagetools inspect "${BUILD_TAG}" | awk '$1 == "Digest:" { print $2; exit }')" +if [[ ! "${digest}" =~ ^sha256:[0-9a-f]{64}$ ]]; then + echo "Could not resolve the pushed digest of ${BUILD_TAG}" >&2 + exit 1 +fi +# The vLLM image build uses this as its ci_base. +buildkite-agent meta-data set aiter-nightly-ci-base "${BUILD_TAG}@${digest}" +# The annotation starts with :warning: when the wheel is stale. +style=info +if [[ "${note}" == :warning:* ]]; then + style=warning +fi +echo "${note}" | buildkite-agent annotate --style "${style}" --context aiter-nightly +echo "--- :white_check_mark: ${BUILD_TAG}@${digest}" diff --git a/.buildkite/scripts/rocm/select_aiter_nightly_wheel.py b/.buildkite/scripts/rocm/select_aiter_nightly_wheel.py new file mode 100755 index 000000000000..af3de54f6020 --- /dev/null +++ b/.buildkite/scripts/rocm/select_aiter_nightly_wheel.py @@ -0,0 +1,260 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Select the AITER nightly wheel for a vLLM ROCm image. + +aiter-nightly-overlay.sh runs this inside the stock ci_base, so it reads that +image's ROCm, torch and Python. It prints the wheel's URL, then the Buildkite +annotation naming it. +""" + +import argparse +import datetime +import importlib.metadata +import os +import sys +import urllib.request +from dataclasses import dataclass +from html.parser import HTMLParser +from pathlib import Path +from urllib.parse import unquote, urldefrag, urljoin + +import regex as re + +INDEX_URL = "https://rocm.frameworks-nightlies.amd.com/whl-multi-arch/amd-aiter/" +AITER_COMMIT_URL = "https://github.com/ROCm/aiter/commit/" +STALE_AFTER = datetime.timedelta(days=1) + +# Example: +# amd_aiter-0.1.25+rocm7.2.3.torch2.12.ae984f5.d20261008-cp312-cp312-linux_x86_64.whl +# is AITER main at commit ae984f5, built 2026-10-08, for ROCm 7.2.3, torch 2.12 +# and Python 3.12. 0.1.25 is a guess at the next release, not a real one. +WHEEL_NAME = re.compile( + r"^amd_aiter-(?P[^+-]+)\+(?P[^-]+)" + r"-(?Pcp\d+)-(?P=python)-linux_x86_64\.whl$" +) +ROCM_FIELD = re.compile(r"(?:^|\.)rocm(\d+(?:\.\d+)*)(?:a\d+)?(?=\.|$)") +TORCH_FIELD = re.compile(r"(?:^|\.)torch(\d+\.\d+)(?=\.|$)") +COMMIT_AND_DATE_FIELD = re.compile(r"(?:^|\.)([0-9a-f]{7,40})\.d(\d{8})(?=\.|$)") + + +@dataclass(frozen=True) +class Wheel: + """An AITER nightly wheel, read from its filename (see WHEEL_NAME).""" + + filename: str + url: str + version: str + base_version: str + rocm: str + torch: str + commit: str + build_date: datetime.date + python: str + + +class _LinkParser(HTMLParser): + """Collects every link on the index page, which lists one per wheel.""" + + def __init__(self): + super().__init__() + self.hrefs: list[str] = [] + + def handle_starttag(self, tag, attrs): + if tag == "a": + self.hrefs.extend( + value for name, value in attrs if name == "href" and value + ) + + +def wheel_from_url(url: str) -> Wheel | None: + """The AITER nightly wheel an absolute URL points at, if it is one.""" + filename = unquote(urldefrag(url).url.rsplit("/", 1)[-1]) + match = WHEEL_NAME.match(filename) + if not match: + return None + local = match["local"] + rocm, torch, build = ( + field.search(local) + for field in (ROCM_FIELD, TORCH_FIELD, COMMIT_AND_DATE_FIELD) + ) + if not (rocm and torch and build): + return None + return Wheel( + filename=filename, + url=url, + version=f"{match['base_version']}+{local}", + base_version=match["base_version"], + rocm=rocm[1], + torch=torch[1], + commit=build[1], + build_date=datetime.datetime.strptime(build[2], "%Y%m%d").date(), + python=match["python"], + ) + + +def parse_index(html: str, index_url: str) -> list[Wheel]: + """Every AITER nightly wheel linked from the index page, a plain HTML list + of links, one per file (the format pip reads).""" + parser = _LinkParser() + parser.feed(html) + wheels = (wheel_from_url(urljoin(index_url, href)) for href in parser.hrefs) + return [wheel for wheel in wheels if wheel] + + +def rocm_version() -> str: + """This image's ROCm version; ROCm 10+ ships as the `rocm` pip SDK.""" + try: + return importlib.metadata.version("rocm") + except importlib.metadata.PackageNotFoundError: + # No pip SDK: fall back to older ROCm releases, which install to ROCM_PATH. + rocm_path = Path(os.environ.get("ROCM_PATH", "/opt/rocm")) + return (rocm_path / ".info" / "version").read_text() + + +def rocm_release(version: str) -> str: + """`7.2.3-49` or `10.0.0a20261001` -> `7.2.3` or `10.0.0`, as wheel names + spell it.""" + match = re.match(r"\d+(?:\.\d+)*", version.strip()) + if not match: + raise ValueError(f"Unrecognised ROCm version: {version!r}") + return match.group(0) + + +def torch_release(version: str) -> str: + """`2.12.0a0+git6bbd260` -> `2.12`, as wheel names spell it.""" + match = re.match(r"\d+\.\d+", version) + if not match: + raise ValueError(f"Unrecognised torch version: {version!r}") + return match.group(0) + + +def python_tag(version_info: tuple[int, int]) -> str: + """(3, 12) -> `cp312`, as wheel names spell it.""" + return f"cp{version_info[0]}{version_info[1]}" + + +def _base_version_key(wheel: Wheel) -> tuple[int, ...]: + """`0.1.25` -> (0, 1, 25), so versions sort as numbers: 0.1.10 > 0.1.9.""" + return tuple(int(part) for part in re.findall(r"\d+", wheel.base_version)) + + +def _label(rocm: str, torch: str, python: str) -> str: + """`rocm7.2.3/torch2.12/cp312`, for messages.""" + return f"rocm{rocm}/torch{torch}/{python}" + + +def select_wheel(wheels: list[Wheel], rocm: str, torch: str, python: str) -> Wheel: + """Newest build for this ROCm, torch and Python, by date: version labels lag. + + Two builds with the same date and version: whichever the index lists first. + """ + target = (rocm, torch, python) + matching = [w for w in wheels if (w.rocm, w.torch, w.python) == target] + if not matching: + available = sorted({_label(w.rocm, w.torch, w.python) for w in wheels}) + raise LookupError( + f"No AITER nightly for {_label(*target)}. " + f"Available: {', '.join(available) or 'none'}" + ) + return max(matching, key=lambda w: (w.build_date, _base_version_key(w))) + + +def fetch(url: str) -> str: + """GET url as text. A failure exits 1, so Buildkite retries the step.""" + with urllib.request.urlopen(url, timeout=60) as response: + return response.read().decode() + + +def annotation(wheel: Wheel, today: datetime.date) -> str: + """The Buildkite annotation naming the wheel this build tests. + + Starts with :warning: if the wheel is stale; the overlay script styles it so. + """ + age = today - wheel.build_date + stale = ( + f":warning: **AITER nightly is {age.days} days old**: " + f"no newer build has been published for " + f"{_label(wheel.rocm, wheel.torch, wheel.python)}.\n\n" + if age > STALE_AFTER + else "" + ) + return stale + ( + f":crescent_moon: **AITER nightly**: main at " + f"[`{wheel.commit}`]({AITER_COMMIT_URL}{wheel.commit}), " + f"built {wheel.build_date.isoformat()} for torch {wheel.torch} " + f"(`{wheel.version}`)" + ) + + +class NoWheel(Exception): + """No wheel fits this image's ROCm, torch and Python. Exit 10.""" + + +def _image_wheel(wheel_url: str) -> Wheel: + """The wheel at wheel_url, if it fits this image's ROCm, torch and Python.""" + wheel = wheel_from_url(wheel_url) + if not wheel: + raise NoWheel(f"Not an AITER nightly wheel: {wheel_url}") + return _for_image([wheel]) + + +def _image() -> tuple[str, str, str]: + """This image's ROCm, torch and Python, as wheel names spell them.""" + try: + rocm = rocm_release(rocm_version()) + except (OSError, ValueError) as error: + raise NoWheel(f"Cannot read this image's ROCm version: {error}") from error + try: + torch = torch_release(importlib.metadata.version("torch")) + except (importlib.metadata.PackageNotFoundError, ValueError) as error: + raise NoWheel(f"Cannot read this image's torch version: {error}") from error + return rocm, torch, python_tag(sys.version_info[:2]) + + +def _for_image(wheels: list[Wheel]) -> Wheel: + """The newest of wheels for the image this runs in.""" + try: + return select_wheel(wheels, *_image()) + except LookupError as error: + raise NoWheel(str(error)) from error + + +def select(index_url: str, wheel_url: str = "") -> Wheel: + """The wheel to install: wheel_url on a retry, else the newest. + + However old the newest is, it is tested; the annotation shows its build date. + """ + if wheel_url: + return _image_wheel(wheel_url) + return _for_image(parse_index(fetch(index_url), index_url)) + + +def _today() -> datetime.date: + """Today in UTC; a function so tests can set the date.""" + return datetime.datetime.now(datetime.UTC).date() + + +# Not 1 or 2, which Buildkite retries; see aiter-nightly-overlay.sh. +EXIT_NO_WHEEL = 10 + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--index-url", default=INDEX_URL) + parser.add_argument( + "--wheel-url", default="", help="a retry's earlier choice; skips the index" + ) + args = parser.parse_args(argv) + try: + wheel = select(args.index_url, args.wheel_url) + except NoWheel as error: + print(error) + return EXIT_NO_WHEEL + print(wheel.url) + print(annotation(wheel, _today())) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 8055293353e4bc2aee2775bc0b077dd77b946b8c Mon Sep 17 00:00:00 2001 From: fululi12 Date: Mon, 12 Oct 2026 11:11:20 -0700 Subject: [PATCH 325/624] [ROCm][MLA] Cut ~21x host dispatch from the MLA metadata build (#58381) Signed-off-by: lifulu Signed-off-by: lifulu Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: lifulu --- .../test_aiter_mla_uniform_mtp_qo_len.py | 104 ++++++++++++++++++ .../layers/attention/mla_attention.py | 22 +++- .../attention/backends/mla/rocm_aiter_mla.py | 37 +++++-- 3 files changed, 145 insertions(+), 18 deletions(-) create mode 100644 tests/v1/attention/test_aiter_mla_uniform_mtp_qo_len.py diff --git a/tests/v1/attention/test_aiter_mla_uniform_mtp_qo_len.py b/tests/v1/attention/test_aiter_mla_uniform_mtp_qo_len.py new file mode 100644 index 000000000000..6c5416d1f41a --- /dev/null +++ b/tests/v1/attention/test_aiter_mla_uniform_mtp_qo_len.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Differential test for the AITER MLA uniform-MTP padding decision. + +`_uniform_padded_mtp_qo_len` decides whether a full-CUDA-graph decode batch may +synthesize dummy query rows. It used to reach that decision through eight tensor +dispatches -- sum/item, boolean-mask index, all, eq, nonzero, item, any -- on a +CPU tensor holding one entry per request. It now reads the tensor out once and +decides on Python ints. + +That rewrite also **reordered** the checks: the `num_decode_tokens <= sum(...)` +early-out moved below the divisibility and `uniform_qo_len <= 1` guards. Every +one of those paths returns 0, so the result should be unchanged -- but "should +be" is exactly the sort of claim worth checking exhaustively, so both +implementations are carried here and compared over the whole small input space. + +Pure host arithmetic; no GPU required. +""" + +import itertools + +import pytest +import torch + +from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAMetadataBuilder + + +def _reference(qo_len: torch.Tensor, max_qo_len: int, num_decode_tokens: int) -> int: + """The pre-#58381 implementation, verbatim.""" + num_reqs = qo_len.numel() + if num_reqs == 0 or num_decode_tokens <= 0: + return 0 + if num_decode_tokens <= int(qo_len.sum().item()): + return 0 + if num_decode_tokens % num_reqs != 0: + return 0 + + uniform_qo_len = num_decode_tokens // num_reqs + if uniform_qo_len <= 1: + return 0 + + positive_qo_len = qo_len[qo_len > 0] + if positive_qo_len.numel() == qo_len.numel(): + return 0 + if positive_qo_len.numel() > 0: + if max_qo_len != uniform_qo_len: + return 0 + if not torch.all(positive_qo_len == uniform_qo_len): + return 0 + + zero_positions = torch.nonzero(qo_len == 0, as_tuple=False).flatten() + if zero_positions.numel() > 0: + first_zero = int(zero_positions[0].item()) + if torch.any(qo_len[first_zero:] > 0): + return 0 + + return uniform_qo_len + + +_ACTUAL = AiterMLAMetadataBuilder._uniform_padded_mtp_qo_len + + +@pytest.mark.parametrize("num_reqs", [1, 2, 3, 4]) +def test_matches_reference_exhaustively(num_reqs): + """Every qo_len vector x decode-token count x max_qo_len in a small space.""" + checked = 0 + for lens in itertools.product(range(0, 4), repeat=num_reqs): + qo = torch.tensor(lens, dtype=torch.int32) + for num_decode_tokens in range(0, 13): + for max_qo_len in range(0, 5): + got = _ACTUAL(qo, max_qo_len, num_decode_tokens) + want = _reference(qo, max_qo_len, num_decode_tokens) + assert got == want, ( + f"lens={lens} num_decode_tokens={num_decode_tokens} " + f"max_qo_len={max_qo_len}: got {got}, reference {want}" + ) + checked += 1 + assert checked > 0 + + +def test_empty_batch(): + assert _ACTUAL(torch.zeros(0, dtype=torch.int32), 4, 8) == 0 + assert _reference(torch.zeros(0, dtype=torch.int32), 4, 8) == 0 + + +def test_the_padded_mtp_case_that_should_succeed(): + """A live prefix at the uniform width, then zero-padded dummies.""" + qo = torch.tensor([4, 4, 0, 0], dtype=torch.int32) + # 4 requests x 4 tokens = 16 captured, 8 live -> dummies are synthesizable. + assert _ACTUAL(qo, 4, 16) == 4 + assert _reference(qo, 4, 16) == 4 + + +def test_positive_after_zero_is_rejected(): + """A live request after a padding hole breaks the contiguous-prefix layout.""" + qo = torch.tensor([4, 0, 4, 0], dtype=torch.int32) + assert _ACTUAL(qo, 4, 16) == 0 + assert _reference(qo, 4, 16) == 0 + + +def test_ragged_live_lengths_are_rejected(): + qo = torch.tensor([4, 2, 0, 0], dtype=torch.int32) + assert _ACTUAL(qo, 4, 16) == 0 + assert _reference(qo, 4, 16) == 0 diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 861c894c383e..db9e8dbc88ca 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -2540,16 +2540,26 @@ def build( "Non-causal multi-token MLA requires an explicitly supported " "attention group." ) - query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] - num_active_reqs = int(torch.count_nonzero(query_lens > 0)) - uniform_active_queries = num_active_reqs > 0 and bool( - torch.all(query_lens[:num_active_reqs] == query_lens[0]) + # query_start_loc_cpu is host-side and holds one entry per + # request, so this check costs more in tensor dispatch than in + # the work it does. Read it out once and validate on Python ints: + # the block runs on every build, between decode steps. + qsl = query_start_loc_cpu.tolist() + query_lens = [qsl[i + 1] - qsl[i] for i in range(len(qsl) - 1)] + num_active_reqs = 0 + for qlen in query_lens: + if qlen > 0: + num_active_reqs += 1 + uniform_active_queries = num_active_reqs > 0 and all( + qlen == query_lens[0] for qlen in query_lens[:num_active_reqs] + ) + trailing_graph_padding = all( + qlen == 0 for qlen in query_lens[num_active_reqs:] ) - trailing_graph_padding = bool(torch.all(query_lens[num_active_reqs:] == 0)) if not (uniform_active_queries and trailing_graph_padding): raise ValueError( "Non-causal MLA requires a uniform query block; got query " - f"lengths {query_lens.tolist()}." + f"lengths {query_lens}." ) # Use exact GPU sequence lengths instead of the prefill path's CPU # context-length upper bounds. diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index 32154105a96d..1115da630b3c 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -563,8 +563,6 @@ def _uniform_padded_mtp_qo_len( # query_start_loc flat for dummy requests. Only synthesize dummy rows # when every padded request maps to the same qlen and the q buffer has # exactly that many rows. - if num_decode_tokens <= int(qo_len.sum().item()): - return 0 if num_decode_tokens % num_reqs != 0: return 0 @@ -572,19 +570,34 @@ def _uniform_padded_mtp_qo_len( if uniform_qo_len <= 1: return 0 - positive_qo_len = qo_len[qo_len > 0] - if positive_qo_len.numel() == qo_len.numel(): + # qo_len is derived from query_start_loc_cpu, so it is already + # host-side and tiny (one entry per request). Reading it out once and + # deciding on Python ints avoids eight separate tensor dispatches -- + # sum/item, boolean-mask index, all, eq, nonzero, item, any -- on the + # metadata-build path, which runs between decode steps while the GPUs + # are idle. + lens = qo_len.tolist() + if num_decode_tokens <= sum(lens): + return 0 + + first_zero = -1 + num_positive = 0 + for i, qlen in enumerate(lens): + if qlen > 0: + num_positive += 1 + # A positive entry after a zero breaks the contiguous live + # prefix that the padded layout requires. + if first_zero >= 0: + return 0 + elif first_zero < 0: + first_zero = i + + if num_positive == num_reqs: return 0 - if positive_qo_len.numel() > 0: + if num_positive > 0: if max_qo_len != uniform_qo_len: return 0 - if not torch.all(positive_qo_len == uniform_qo_len): - return 0 - - zero_positions = torch.nonzero(qo_len == 0, as_tuple=False).flatten() - if zero_positions.numel() > 0: - first_zero = int(zero_positions[0].item()) - if torch.any(qo_len[first_zero:] > 0): + if any(qlen != uniform_qo_len for qlen in lens if qlen > 0): return 0 return uniform_qo_len From 56e3b081aea454d407ff658b3b12a5692f54fa47 Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Tue, 13 Oct 2026 03:33:31 +0900 Subject: [PATCH 326/624] [Bugfix][weight loader] Stop warning when a parameter only widens the weight (#60020) Signed-off-by: lesj0610 Co-authored-by: Ning Xie --- vllm/model_executor/models/utils.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index bbb7a54b552e..ef1aac6de558 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -322,7 +322,15 @@ def _load_param( ) if param.dtype != weight_data.dtype: - logger.warning( + # An exact widening -- an fp32 accumulator reading a bf16 + # table, say -- keeps every stored value, so it is a debug + # line. Every other mismatch still warns. + exact_widening = weight_data.dtype in ( + torch.float16, + torch.bfloat16, + torch.float32, + ) and param.dtype in (torch.float32, torch.float64) + (logger.debug if exact_widening else logger.warning)( "Attempted to load weight %s with dtype %s into " "parameter with dtype %s", weight_qualname, From d6aca837c3c6b773a557448d526db219a5a5718b Mon Sep 17 00:00:00 2001 From: Maroon Ayoub Date: Mon, 12 Oct 2026 21:56:51 +0300 Subject: [PATCH 327/624] [Attention] Make the sparse MLA prefill workspace size configurable (#61023) Signed-off-by: Maroon Ayoub Co-authored-by: Claude --- vllm/config/attention.py | 9 +++++++++ vllm/config/vllm.py | 10 ++++++++++ .../attention/backends/mla/flashmla_sparse.py | 20 ++++++++----------- vllm/v1/attention/backends/mla/indexer.py | 11 +++++----- 4 files changed, 33 insertions(+), 17 deletions(-) diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 80548f58f2ca..67605f3c7e00 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -117,6 +117,15 @@ class AttentionConfig: When False (default), pure prefill batches use forward_mha when implemented. Set to True to always use the MQA path.""" + sparse_mla_prefill_workspace_factor: int | None = Field(default=None, ge=2) + """Sparse MLA prefill workspace size, in multiples of `max_model_len`. + FlashMLA sparse gathers a prefill step's context KV into this workspace and + runs attention in more chunks when the step holds more. The sparse indexer's + prefill K buffer is sized to fit in the same memory. Lower values return + memory to the KV cache at long `max_model_len`; the minimum of 2 leaves room + for a full-length request's context-parallel shard. If None, 5 when + `max_model_len` is at most 256K tokens, else 2.""" + flex_attn_block_m: int | None = None """Triton kernel BLOCK_M tile size for flex attention. Must be a power of 2 >= 16. If None and VLLM_BATCH_INVARIANT=1, diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index af4c99113f0f..f27c8418ab4d 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1787,6 +1787,16 @@ def has_blocked_weights(): self._maybe_disable_dynamic_sd_for_data_parallel() self._maybe_override_dynamic_sd_cudagraph_mode() + if ( + self.attention_config.sparse_mla_prefill_workspace_factor is None + and self.model_config is not None + ): + # Resolved once: KV cache auto-fit can lower max_model_len later, and + # factor * max_model_len must stay within the workspace reserved at init. + self.attention_config.sparse_mla_prefill_workspace_factor = ( + 5 if self.model_config.max_model_len <= 256 * 1024 else 2 + ) + if ( self.attention_config.hisparse_config is None and self.kv_transfer_config is not None diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 4fe0be0b1f68..3cf0c33c534d 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -334,14 +334,14 @@ class Chunk: fp8_use_mixed_batch: bool = False -def get_prefill_workspace_size(max_model_len: int): - # NOTE(Lucas): 5 is a magic number for controlling the prefill buffer size. - # May be tuned later. - # Memory usage: 5 * max_model_len * 576 * 2 bytes - # Example: DeepSeek-V3.2 with max_model_len=163840 -> +def get_prefill_workspace_size(vllm_config: VllmConfig) -> int: + # Memory usage: factor * max_model_len * 576 * 2 bytes + # Example: DeepSeek-V3.2 with max_model_len=163840 and the default factor -> # 5 * 163840 * 576 * 2 = ~900 MB # This fits nicely below the typical MoE workspace size of >2GB so this is "free" - return max_model_len * 5 + factor = vllm_config.attention_config.sparse_mla_prefill_workspace_factor + assert factor is not None + return vllm_config.model_config.max_model_len * factor class FlashMLASparseMetadataBuilder( @@ -562,9 +562,7 @@ def _build_fp8_separate_prefill_decode( prefill_seq_lens_cpu = seq_lens_cpu[num_decodes:] # Chunk prefill requests to fit within workspace size - max_prefill_buffer_size = get_prefill_workspace_size( - self.vllm_config.model_config.max_model_len - ) + max_prefill_buffer_size = get_prefill_workspace_size(self.vllm_config) # Workspace rows per entry, and the prefill rows each entry covers. workspace_rows = prefill_seq_lens_cpu row_bounds = np.arange(num_prefills + 1) @@ -804,9 +802,7 @@ def __init__( if kv_cache_dtype in QUANTIZED_DS_MLA_CACHE_FORMATS: # Reserve workspace during initialization assert vllm_config is not None and vllm_config.model_config is not None - prefill_workspace_size = get_prefill_workspace_size( - vllm_config.model_config.max_model_len - ) + prefill_workspace_size = get_prefill_workspace_size(vllm_config) parallel_config = vllm_config.parallel_config self.pcp_dcp_kv_gather = ( parallel_config.prefill_context_parallel_size > 1 diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 97de9ffa9ba2..2d797ae66283 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -917,15 +917,16 @@ def update_draft_decode_metadata( def get_max_prefill_buffer_size(vllm_config: VllmConfig): max_model_len = vllm_config.model_config.max_model_len - # NOTE(Chen): 40 is a magic number for controlling the prefill buffer size. + factor = vllm_config.attention_config.sparse_mla_prefill_workspace_factor + assert factor is not None # Each entry is 128 fp8 bytes and 4 scale bytes for a total of 132 bytes. - # The flashmla_sparse backend uses a workspace size of 5 * max_model_len. + # The flashmla_sparse backend uses a workspace size of factor * max_model_len. # The memory usage of the workspace there is 576 * 2 bytes; so we size this as - # (576 * 2 // 132) * 5 = 40 to maximize this workspace size while still fitting - # within the flashmla_sparse workspace. + # (576 * 2 // 132) * factor (40 at factor 5) to maximize this workspace + # size while still fitting within the flashmla_sparse workspace. # For DeepSeek-V3.2, the max_model_len is 163840. # 40 * 163840 * 132 = 865075200 bytes = 825 MB - return max_model_len * 40 + return max_model_len * (576 * 2 // 132) * factor def _supports_varlen_paged_mqa_logits() -> bool: From d0948f754dd1fa9a77721c5e9d3722036e74e6a7 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Mon, 12 Oct 2026 15:13:23 -0400 Subject: [PATCH 328/624] [Bug] Honor spec decode draft model config (#61240) Signed-off-by: yewentao256 --- tests/v1/worker/test_gpu_model_runner_v2.py | 34 +++++++++++++++++++ .../spec_decode/standalone_ar/speculator.py | 12 +++++++ 2 files changed, 46 insertions(+) diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 85c3ddfce425..8ccc3731d9c8 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import contextlib +from dataclasses import dataclass from types import SimpleNamespace from unittest.mock import Mock @@ -9,6 +10,7 @@ import torch import vllm.v1.worker.gpu.model_runner as model_runner_module +from vllm.config import CacheConfig, ParallelConfig from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry from vllm.v1.kv_cache_interface import ( CircularBufferSpec, @@ -21,6 +23,7 @@ from vllm.v1.worker.gpu.async_utils import async_copy_to_np from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.model_runner import ExecuteModelState, GPUModelRunner +from vllm.v1.worker.gpu.spec_decode.standalone_ar import speculator as ar_module from vllm.v1.worker.gpu.spec_decode.utils import get_drafter_hidden_states @@ -366,3 +369,34 @@ def test_async_copy_to_np_does_not_alias_reused_buffer(): buffer.fill_(1) assert snapshot.tolist() == [0, 0, 0, 0] + + +@pytest.mark.parametrize( + ("draft_dtype", "expected_dtype"), + [("auto", "auto"), ("bfloat16", "bfloat16"), (None, "fp8_e4m3")], +) +def test_standalone_draft_kv_dtype_override(monkeypatch, draft_dtype, expected_dtype): + @dataclass + class Config: + cache_config: CacheConfig + parallel_config: ParallelConfig + model_config: object + quant_config: object | None = None + + target = Config(CacheConfig(cache_dtype="fp8_e4m3"), ParallelConfig(), object()) + speculator = object.__new__(ar_module.StandaloneARSpeculator) + speculator.vllm_config = target + speculator.draft_model_config = object() + speculator.speculative_config = SimpleNamespace( + kv_cache_dtype=draft_dtype, draft_parallel_config=ParallelConfig() + ) + get_model = Mock(return_value=object()) + monkeypatch.setattr(ar_module, "get_model", get_model) + monkeypatch.setattr(ar_module, "set_model_tag", lambda _: contextlib.nullcontext()) + + speculator.load_draft_model(object(), set()) + + draft_config = get_model.call_args.kwargs["vllm_config"] + assert draft_config.cache_config.cache_dtype == expected_dtype + assert speculator.attn_vllm_config.cache_config.cache_dtype == expected_dtype + assert target.cache_config.cache_dtype == "fp8_e4m3" diff --git a/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py index 04e79debb869..bd404ef38fd9 100644 --- a/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py @@ -70,6 +70,12 @@ def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: def capture(self) -> None: pass # CUDA graph not yet supported for standalone AR speculator. + @property + def attn_vllm_config(self) -> VllmConfig: + if self.speculative_config.kv_cache_dtype is not None: + return self._draft_vllm_config + return self.vllm_config + def set_attn( self, model_state: ModelState, @@ -99,15 +105,21 @@ def load_draft_model( ) -> nn.Module: spec = self.speculative_config assert spec is not None + cache_config = self.vllm_config.cache_config + if spec.kv_cache_dtype is not None: + cache_config = replace(cache_config, cache_dtype=spec.kv_cache_dtype) + draft_vllm_config = replace( self.vllm_config, model_config=self.draft_model_config, quant_config=None, + cache_config=cache_config, parallel_config=replace( spec.draft_parallel_config, rank=self.vllm_config.parallel_config.rank, ), ) + self._draft_vllm_config = draft_vllm_config with set_model_tag("draft_model"): return get_model( vllm_config=draft_vllm_config, From 08582942d8af7fc58339f3a5fedaffbbc8ea9f02 Mon Sep 17 00:00:00 2001 From: Martin Hickey Date: Mon, 12 Oct 2026 20:17:12 +0100 Subject: [PATCH 329/624] [Bugfix][Frontend] Honor request_id and X-Request-Id in chat and completions render (#61236) Signed-off-by: Martin Hickey --- .../scale_out/render/test_render.py | 68 +++++++++++++++++++ .../scale_out/render/api_router.py | 6 +- vllm/entrypoints/scale_out/render/serving.py | 17 +++-- 3 files changed, 83 insertions(+), 8 deletions(-) diff --git a/tests/entrypoints/scale_out/render/test_render.py b/tests/entrypoints/scale_out/render/test_render.py index 308555be362e..37f4b8eecf47 100644 --- a/tests/entrypoints/scale_out/render/test_render.py +++ b/tests/entrypoints/scale_out/render/test_render.py @@ -367,6 +367,74 @@ async def test_completion_render_multiple_prompts(client): assert prompt["request_id"].startswith("cmpl-") +CHAT_BODY = { + "model": MODEL_NAME, + "messages": [{"role": "user", "content": "Hello"}], +} + + +@pytest.mark.asyncio +async def test_chat_render_honors_body_request_id(client): + response = await client.post( + "/v1/chat/completions/render", json={**CHAT_BODY, "request_id": "abc"} + ) + + assert response.status_code == 200 + assert response.json()["request_id"] == "chatcmpl-abc" + + +@pytest.mark.asyncio +async def test_chat_render_header_request_id_wins(client): + response = await client.post( + "/v1/chat/completions/render", + json={**CHAT_BODY, "request_id": "abc"}, + headers={"X-Request-Id": "hdr"}, + ) + + assert response.status_code == 200 + assert response.json()["request_id"] == "chatcmpl-hdr" + + +@pytest.mark.asyncio +async def test_chat_render_generates_unique_request_ids(client): + ids = set() + for _ in range(2): + response = await client.post("/v1/chat/completions/render", json=CHAT_BODY) + assert response.status_code == 200 + request_id = response.json()["request_id"] + assert request_id.startswith("chatcmpl-") + ids.add(request_id) + + assert len(ids) == 2 + + +@pytest.mark.asyncio +async def test_completion_render_request_id_single_prompt(client): + response = await client.post( + "/v1/completions/render", + json={"model": MODEL_NAME, "prompt": "Hello world", "request_id": "abc"}, + ) + + assert response.status_code == 200 + assert [r["request_id"] for r in response.json()] == ["cmpl-abc"] + + +@pytest.mark.asyncio +async def test_completion_render_request_id_multiple_prompts(client): + response = await client.post( + "/v1/completions/render", + json={ + "model": MODEL_NAME, + "prompt": ["Hello world", "Goodbye world"], + "request_id": "abc", + }, + headers={"X-Request-Id": "hdr"}, + ) + + assert response.status_code == 200 + assert [r["request_id"] for r in response.json()] == ["cmpl-hdr-0", "cmpl-hdr-1"] + + @pytest.mark.asyncio async def test_chat_completion_render_multi_turn(client): """Test chat completion render with multi-turn conversation.""" diff --git a/vllm/entrypoints/scale_out/render/api_router.py b/vllm/entrypoints/scale_out/render/api_router.py index 97df983d2e52..0bf6dbe5f107 100644 --- a/vllm/entrypoints/scale_out/render/api_router.py +++ b/vllm/entrypoints/scale_out/render/api_router.py @@ -43,7 +43,7 @@ async def render_chat_completion(request: ChatCompletionRequest, raw_request: Re "The model does not support Chat Completions Render API" ) - result = await handler.render_chat_request(request) + result = await handler.render_chat_request(request, raw_request) if isinstance(result, ErrorResponse): return JSONResponse(content=result.model_dump(), status_code=result.error.code) @@ -67,7 +67,7 @@ async def render_messages(request: AnthropicMessagesRequest, raw_request: Reques if handler is None: raise NotImplementedError("The model does not support Messages Render API") - result = await handler.render_messages_request(request) + result = await handler.render_messages_request(request, raw_request) if isinstance(result, ErrorResponse): return JSONResponse(content=result.model_dump(), status_code=result.error.code) @@ -90,7 +90,7 @@ async def render_completion(request: CompletionRequest, raw_request: Request): if handler is None: raise NotImplementedError("The model does not support Completions Render API") - result = await handler.render_completion_request(request) + result = await handler.render_completion_request(request, raw_request) if isinstance(result, ErrorResponse): return JSONResponse(content=result.model_dump(), status_code=result.error.code) diff --git a/vllm/entrypoints/scale_out/render/serving.py b/vllm/entrypoints/scale_out/render/serving.py index 5caf78bffe9e..b8d8ab7e1bc2 100644 --- a/vllm/entrypoints/scale_out/render/serving.py +++ b/vllm/entrypoints/scale_out/render/serving.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import TYPE_CHECKING +from fastapi import Request + from vllm.entrypoints.anthropic.protocol import AnthropicMessagesRequest from vllm.entrypoints.anthropic.serving import AnthropicServingMessages from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest @@ -31,7 +33,6 @@ extract_prompt_len, ) from vllm.renderers.online_renderer import OnlineRenderer -from vllm.utils import random_uuid logger = init_logger(__name__) @@ -77,6 +78,7 @@ def __init__( async def render_chat_request( self, request: ChatCompletionRequest, + raw_request: Request | None = None, ) -> GenerateRequest | ErrorResponse: """Validate the model and preprocess a chat completion request. @@ -139,7 +141,9 @@ async def render_chat_request( ) reasoning_ended = request.resolve_reasoning_ended(parser, token_ids) - request_id = f"chatcmpl-{random_uuid()}" + request_id = ( + f"chatcmpl-{self._base_request_id(raw_request, request.request_id)}" + ) return GenerateRequest( request_id=request_id, @@ -161,6 +165,7 @@ async def render_chat_request( async def render_messages_request( self, request: AnthropicMessagesRequest, + raw_request: Request | None = None, ) -> GenerateRequest | ErrorResponse: """Validate the model and preprocess an Anthropic Messages request. @@ -172,11 +177,12 @@ async def render_messages_request( request, merge_inline_system=self._merge_inline_system ) chat_req.return_mm_kwargs = request.return_mm_kwargs - return await self.render_chat_request(chat_req) + return await self.render_chat_request(chat_req, raw_request) async def render_completion_request( self, request: CompletionRequest, + raw_request: Request | None = None, ) -> list[GenerateRequest] | ErrorResponse: """Validate the model and preprocess a completion request. @@ -197,8 +203,9 @@ async def render_completion_request( ) if isinstance(result, ErrorResponse): return result + base_id = f"cmpl-{self._base_request_id(raw_request, request.request_id)}" generate_requests: list[GenerateRequest] = [] - for engine_input in result: + for i, engine_input in enumerate(result): prompt_components = extract_prompt_components( self.model_config, engine_input ) @@ -220,7 +227,7 @@ async def render_completion_request( max_tokens, self.default_sampling_params ) - request_id = f"cmpl-{random_uuid()}" + request_id = base_id if len(result) == 1 else f"{base_id}-{i}" generate_requests.append( GenerateRequest( From 1bdd3fa2d509ce8c40d92238050433822ce76d22 Mon Sep 17 00:00:00 2001 From: ztysth <166926442+ztysth@users.noreply.github.com> Date: Tue, 13 Oct 2026 03:25:35 +0800 Subject: [PATCH 330/624] [Bugfix][Core] Check multimodal chunk boundaries before encoder cache reuse (#59721) Signed-off-by: ztysth Signed-off-by: ztysth <166926442+ztysth@users.noreply.github.com> Co-authored-by: OpenAI Codex --- tests/v1/core/test_scheduler.py | 41 ++++++++++++++++++++++++++++++++- vllm/v1/core/sched/scheduler.py | 33 +++++++++++++------------- 2 files changed, 57 insertions(+), 17 deletions(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index f53ba2ad2305..107a37dc91c0 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -741,7 +741,8 @@ def test_same_step_duplicate_encoder_input_stays_cached(): assert "image" not in cache.freeable -def test_no_mm_input_chunking(): +@pytest.mark.parametrize("warm_encoder_cache", [False, True]) +def test_no_mm_input_chunking(warm_encoder_cache: bool): # Disable multimodal input chunking. scheduler = create_scheduler( model="llava-hf/llava-1.5-7b-hf", @@ -753,6 +754,16 @@ def test_no_mm_input_chunking(): requests = create_requests( num_requests=1, num_tokens=1200, mm_positions=mm_positions ) + if warm_encoder_cache: + cached_request = create_requests( + num_requests=1, + num_tokens=1200, + mm_positions=mm_positions, + req_ids=["cache-primer"], + )[0] + scheduler.encoder_cache_manager.allocate(cached_request, 0) + scheduler.encoder_cache_manager.free(cached_request) + for request in requests: scheduler.add_request(request) @@ -780,6 +791,8 @@ def test_no_mm_input_chunking(): assert output.scheduled_cached_reqs.num_reqs == 1 assert len(output.finished_req_ids) == 0 assert output.num_scheduled_tokens[requests[0].request_id] == 800 + if warm_encoder_cache: + assert not output.scheduled_encoder_inputs # Test that we fail if we disable chunked mm input and use too small # of a max_num_batched_tokens for the mm input. @@ -791,6 +804,32 @@ def test_no_mm_input_chunking(): ) +def test_no_mm_input_chunking_for_same_step_duplicate(): + scheduler = create_scheduler( + model="llava-hf/llava-1.5-7b-hf", + max_num_batched_tokens=1024, + disable_chunked_mm_input=True, + max_model_len=2048, + ) + request = create_requests( + num_requests=1, + num_tokens=1352, + mm_hashes_list=[["image", "image"]], + mm_positions=[ + [ + PlaceholderRange(offset=100, length=576), + PlaceholderRange(offset=776, length=576), + ] + ], + )[0] + scheduler.add_request(request) + + output = scheduler.schedule() + + assert output.num_scheduled_tokens[request.request_id] == 776 + assert output.scheduled_encoder_inputs == {request.request_id: [0]} + + @pytest.mark.parametrize("enable_prefix_caching", [True, False]) def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool): """Test scheduling behavior with concurrent partial requests. diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index bcb6def9383c..cec96d2a4923 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1863,6 +1863,23 @@ def _try_schedule_encoder_inputs( # already calculated encoder inputs and can skip here. continue + # If no encoder input chunking is allowed, we do not want to + # partially schedule a multimodal item. If the scheduled range would + # only cover part of the mm input, roll back to before the mm item. + if ( + self.scheduler_config.disable_chunked_mm_input + and num_computed_tokens < start_pos + and (num_computed_tokens + num_new_tokens) + < (start_pos + num_encoder_tokens) + ): + # Account for EAGLE shift when rolling back to avoid + # encoder cache miss. This ensures the scheduled range + # stops before start_pos even with the shift. + num_new_tokens = max( + 0, start_pos - (num_computed_tokens + shift_computed_tokens) + ) + break + if not self.is_encoder_decoder: # We are not using the encoder cache for encoder-decoder models, # yet. @@ -1902,22 +1919,6 @@ def _try_schedule_encoder_inputs( # previous step. continue - # If no encoder input chunking is allowed, we do not want to - # partially schedule a multimodal item. If the scheduled range would - # only cover part of the mm input, roll back to before the mm item. - if ( - self.scheduler_config.disable_chunked_mm_input - and num_computed_tokens < start_pos - and (num_computed_tokens + num_new_tokens) - < (start_pos + num_encoder_tokens) - ): - # Account for EAGLE shift when rolling back to avoid - # encoder cache miss. This ensures the scheduled range - # stops before start_pos even with the shift. - num_new_tokens = max( - 0, start_pos - (num_computed_tokens + shift_computed_tokens) - ) - break if not self.encoder_cache_manager.can_allocate( request, i, encoder_compute_budget, num_embeds_to_schedule ): From 7039c7f17b5190862ab94ec476419c9127f2f400 Mon Sep 17 00:00:00 2001 From: Martin Hickey Date: Mon, 12 Oct 2026 20:31:03 +0100 Subject: [PATCH 331/624] [Bugfix][Rust Frontend] Omit unset sampling params in render output (#61220) Signed-off-by: Martin Hickey --- rust/src/server/src/routes/tests.rs | 34 +++++++++++++++++++ rust/src/text/src/request.rs | 20 +++++++++++ .../token_in_token_out/test_protocol.py | 17 ++++++++++ 3 files changed, 71 insertions(+) diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index c6c05e63b1c8..d1e9a55d16cc 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -1374,6 +1374,40 @@ async fn render_chat_returns_generate_request_with_header_request_id() { assert!(json.get("mm_features").is_none()); } +#[tokio::test] +async fn render_chat_omits_unset_sampling_params() { + let mut app = test_render_app(); + let response = app + .call( + Request::builder() + .method("POST") + .uri("/v1/chat/completions/render") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "render-model", + "messages": [{"role": "user", "content": "hello"}], + "max_completion_tokens": 8 + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); + + let sampling_params = json["sampling_params"].as_object().expect("sampling_params object"); + assert!( + sampling_params.values().all(|value| !value.is_null()), + "unexpected null in {sampling_params:?}" + ); + assert_eq!(sampling_params["max_tokens"], 8); +} + #[tokio::test] async fn render_chat_uses_processor_owned_parser_selections() { let mut app = test_render_app_with_parser_selections( diff --git a/rust/src/text/src/request.rs b/rust/src/text/src/request.rs index b00d8f3b2f12..c3efcde496f4 100644 --- a/rust/src/text/src/request.rs +++ b/rust/src/text/src/request.rs @@ -65,6 +65,7 @@ impl Default for Prompt { /// /// Original Python definition: /// +#[serde_with::skip_serializing_none] #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(default)] pub struct SamplingParams { @@ -320,4 +321,23 @@ mod tests { assert!(err.to_string().contains("stop strings cannot be empty")); assert!(err.to_string().contains("req-empty-stop")); } + + #[test] + fn sampling_params_serialization_omits_unset_fields() { + let params = SamplingParams { + max_tokens: Some(64), + ..Default::default() + }; + + let value = serde_json::to_value(¶ms).unwrap(); + let obj = value.as_object().unwrap(); + assert!( + obj.values().all(|v| !v.is_null()), + "unexpected null in {value}" + ); + assert_eq!(obj["max_tokens"], 64); + + let back: SamplingParams = serde_json::from_value(value).unwrap(); + assert_eq!(back, params); + } } diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py b/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py index 6a57e3b2b475..842dded1c76c 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py @@ -132,6 +132,23 @@ def test_reasoning_parser_kwargs_rejects_unknown_keys(): GenerateRequest.model_validate(payload) +def test_rust_render_sampling_params_without_nulls_validate(): + """Rust render omits unset sampling fields; Python generate rejects nulls.""" + sampling_params = {"max_tokens": 64, "watermarking": True} + request = GenerateRequest.model_validate( + {"token_ids": [1, 2, 3], "sampling_params": sampling_params} + ) + assert request.sampling_params.max_tokens == 64 + + with pytest.raises(ValidationError, match="temperature"): + GenerateRequest.model_validate( + { + "token_ids": [1, 2, 3], + "sampling_params": {**sampling_params, "temperature": None}, + } + ) + + def test_output_mode_defaults_to_tokens(): assert GenerateRequest.model_validate(_base_payload()).output_mode == "tokens" From fabaecbfa5d48389cc8921d782b3f7d3aa0d928f Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Tue, 13 Oct 2026 03:32:25 +0800 Subject: [PATCH 332/624] [Misc] Expand @ZJY0516 CODEOWNERS to Mamba and KV cache (#61181) Signed-off-by: zjy0516 --- .github/CODEOWNERS | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 01dab4e76b52..b14edb0d5bce 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -3,13 +3,14 @@ # This lists cover the "core" components of vLLM that require careful review /vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng +/vllm/compilation/breakable_cudagraph.py @ZJY0516 /vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi @ivanium @wzhao18 /vllm/distributed/kv_transfer/kv_connector/v1/hisparse @MatthewBonanni @LucasWilkinson /vllm/lora @jeejeelee /vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni /vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye @wzhao18 /vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety @zyongye -/vllm/model_executor/layers/mamba @tdoublep @tomeras91 +/vllm/model_executor/layers/mamba @tdoublep @tomeras91 @ZJY0516 /vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy /vllm/model_executor/model_loader @22quinn /vllm/model_executor/determinism/ @yewentao256 @@ -26,7 +27,7 @@ # Any change to the VllmConfig changes can have a large user-facing impact, # so spam a lot of people /vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @yewentao256 @ProExpertProg -/vllm/config/cache.py @heheda12345 @ivanium +/vllm/config/cache.py @heheda12345 @ivanium @ZJY0516 # Config utils /vllm/config/utils.py @hmellor @@ -102,17 +103,18 @@ /vllm/v1/attention/backends/flashinfer.py @mgoin @pavanimajety @vadiklyutiy @wzhao18 /vllm/v1/attention/backends/triton_attn.py @tdoublep /vllm/v1/attention/backends/gdn_attn.py @ZJY0516 @vadiklyutiy -/vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium +/vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium @ZJY0516 /vllm/v1/sample @22quinn @houseroad @njhill /vllm/v1/spec_decode @benchislett @luccafong @MatthewBonanni /vllm/v1/structured_output @mgoin @aarnphm @benchislett @arpera -/vllm/v1/kv_cache_interface.py @heheda12345 @ivanium +/vllm/v1/kv_cache_interface.py @heheda12345 @ivanium @ZJY0516 /vllm/v1/kv_offload @ApostaC @orozery /vllm/v1/simple_kv_offload @ivanium /vllm/v1/hisparse @MatthewBonanni @LucasWilkinson /vllm/v1/engine @njhill /vllm/v1/executor @njhill /vllm/v1/worker @njhill +/vllm/v1/worker/mamba_utils.py @njhill @ZJY0516 /vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery @NickLucche @ivanium /vllm/v1/worker/encoder_cudagraph.py @shen-shanshan /vllm/v1/worker/encoder_cudagraph_defs.py @shen-shanshan @@ -120,6 +122,7 @@ # Model runner V2 /vllm/v1/worker/gpu @WoosukKwon @njhill @yewentao256 /vllm/v1/worker/gpu/kv_connector.py @orozery +/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @WoosukKwon @njhill @yewentao256 @ZJY0516 # CI & building /.buildkite @Harry-Chen @khluu @@ -132,7 +135,7 @@ /tests/distributed/test_multi_node_assignment.py @youkaichao /tests/distributed/test_pipeline_parallel.py @youkaichao /tests/distributed/test_same_node.py @youkaichao -/tests/evals @mgoin @vadiklyutiy @AndreasKaratzas +/tests/evals @mgoin @vadiklyutiy @AndreasKaratzas @ZJY0516 /tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256 @zyongye @AndreasKaratzas /tests/kernels/ir @ProExpertProg @tjtanaa /tests/models @DarkLight1337 @ywang96 @AndreasKaratzas @@ -140,9 +143,9 @@ /tests/multimodal/media @Isotr0py /tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye @AndreasKaratzas /tests/v1/structured_output @mgoin @aarnphm @arpera -/tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium +/tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium @ZJY0516 /tests/lora @jeejeelee -/tests/models/language/generation/test_hybrid.py @tdoublep @tomeras91 +/tests/models/language/generation/test_hybrid.py @tdoublep @tomeras91 @ZJY0516 /tests/v1/kv_connector/nixl_integration @NickLucche /tests/v1/kv_connector @ApostaC @orozery @ivanium @wzhao18 /tests/v1/kv_offload @ApostaC @orozery From abad32a718d2f9e3af7b18818eb474322eaadfc4 Mon Sep 17 00:00:00 2001 From: Mikko Tukiainen Date: Mon, 12 Oct 2026 22:53:08 +0300 Subject: [PATCH 333/624] [ROCm][Perf] Fix Gemma AR+RMS prefill fallback (#58547) Signed-off-by: Mikko Tukiainen Signed-off-by: Mikko Tukiainen --- .buildkite/test_areas/distributed.yaml | 2 + .../distributed/test_fusion_all_reduce.py | 65 ++++++++++++-- tests/distributed/test_aiter_rms_fallback.py | 89 +++++++++++++++++++ vllm/_aiter_ops.py | 16 +++- .../passes/fusion/allreduce_rms_fusion.py | 58 +++++++++++- vllm/kernels/aiter_ops.py | 3 + 6 files changed, 220 insertions(+), 13 deletions(-) create mode 100644 tests/distributed/test_aiter_rms_fallback.py diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 0b1a73136b04..f541c42c93a8 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -522,6 +522,7 @@ steps: - tests/distributed/test_rocm_aiter_custom_ar.py - tests/distributed/test_rocm_quick_reduce.py - tests/distributed/test_quick_all_reduce.py + - tests/distributed/test_aiter_rms_fallback.py - tests/v1/e2e/general/test_rocm_aiter_custom_ar.py - tests/utils.py - examples/rl/rlhf_async_new_apis.py @@ -539,6 +540,7 @@ steps: - pytest -v -s tests/v1/e2e/general/test_rocm_aiter_custom_ar.py - pytest -v -s tests/distributed/test_rocm_quick_reduce.py - pytest -v -s tests/distributed/test_quick_all_reduce.py + - pytest -v -s tests/distributed/test_aiter_rms_fallback.py - label: ":nvidia: (B200) Distributed" key: distributed-tests-2xb200 diff --git a/tests/compile/passes/distributed/test_fusion_all_reduce.py b/tests/compile/passes/distributed/test_fusion_all_reduce.py index 1c99cc25330f..9e2b96957576 100644 --- a/tests/compile/passes/distributed/test_fusion_all_reduce.py +++ b/tests/compile/passes/distributed/test_fusion_all_reduce.py @@ -177,10 +177,12 @@ def __init__( token_num=16, eps=1e-6, dtype: torch.dtype = torch.float16, + use_aiter: bool = False, ): super().__init__() self.hidden_size = hidden_size self.eps = eps + self.use_aiter = use_aiter self.norm = [GemmaRMSNorm(hidden_size, eps) for _ in range(4)] # Non-trivial weight (~Gemma range) so (1 + w) exercises the scale path. for n in self.norm: @@ -210,6 +212,8 @@ def ops_in_model_before(self): return [torch.ops.vllm.all_reduce.default] def ops_in_model_after(self): + if self.use_aiter: + return [rocm_aiter_ops.get_fused_allreduce_rmsnorm_op()] return [torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default] @@ -383,6 +387,32 @@ def ops_in_model_after(self): ] +class TestAiterAllReduceGemmaRMSNormGroupQuantFP8Model( + TestAiterAllReduceRMSNormGroupQuantFP8Model +): + """Gemma norms feeding the per-group FP8 quant sites. + + The fused AR+RMS+quant op has no ``gemma_norm`` mode, so these sites must + keep matching the quant-fused patterns with the pre-added ``1 + weight``. + The Gemma AR+RMS-only patterns match a prefix of the same subgraph, so if + they were registered ahead of the quant-fused ones they would consume the + ``all_reduce`` and strand ``rocm_aiter_group_fp8_quant`` as a separate + kernel -- which the ``fully_replaced=True`` check below would catch. + """ + + def __init__( + self, + hidden_size=128, + token_num=16, + eps=1e-6, + dtype: torch.dtype = torch.bfloat16, + ): + super().__init__(hidden_size, token_num, eps, dtype) + self.norm = [GemmaRMSNorm(hidden_size, eps) for _ in range(4)] + for norm in self.norm: + norm.weight.requires_grad_(False) + + class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module): def __init__( self, hidden_size=16, token_num=16, eps=1e-6, dtype: torch.dtype = torch.float16 @@ -450,10 +480,10 @@ def ops_in_model_before(self): pytest.param( TestAllReduceGemmaRMSNormModel, False, - False, + IS_AITER_FOUND, marks=pytest.mark.skipif( - current_platform.is_rocm(), - reason="Not supported on ROCm platform", + current_platform.is_rocm() and not IS_AITER_FOUND, + reason="AITER is required on ROCm", ), ), pytest.param( @@ -646,7 +676,10 @@ def all_reduce_fusion_pass_on_test_model( ) token_num = batch_size * seq_len - if test_model_cls is TestAllReduceRMSNormModel: + if test_model_cls in ( + TestAllReduceRMSNormModel, + TestAllReduceGemmaRMSNormModel, + ): model = test_model_cls( hidden_size, token_num, dtype=dtype, use_aiter=use_aiter ) @@ -667,9 +700,13 @@ def all_reduce_fusion_pass_on_test_model( ) backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False) backend.check_after_ops(model.ops_in_model_after()) - if test_model_cls in ( - TestAllReduceGemmaRMSNormModel, - TestAllReduceGemmaRMSNormStaticQuantFP8Model, + if ( + test_model_cls + in ( + TestAllReduceGemmaRMSNormModel, + TestAllReduceGemmaRMSNormStaticQuantFP8Model, + ) + and not use_aiter ): fused_op = torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default fused_nodes = list(find_op_nodes(fused_op, backend.graph_post_pass)) @@ -684,6 +721,13 @@ def all_reduce_fusion_pass_on_test_model( @pytest.mark.parametrize("hidden_size", [128]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False]) +@pytest.mark.parametrize( + "test_model_cls", + [ + TestAiterAllReduceRMSNormGroupQuantFP8Model, + TestAiterAllReduceGemmaRMSNormGroupQuantFP8Model, + ], +) @pytest.mark.skipif( not current_platform.is_rocm(), reason="ROCm AITER AR+RMS+per-group-FP8-quant fusion is ROCm-only", @@ -695,6 +739,7 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace( hidden_size: int, dtype: torch.dtype, enable_rms_norm_custom_op: bool, + test_model_cls: type[torch.nn.Module], monkeypatch: pytest.MonkeyPatch, ): """Sibling of ``test_all_reduce_fusion_pass_replace`` for the new @@ -710,6 +755,10 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace( residual, DSv3.2 indexer fan-out; parametrized over ``rocm_aiter_group_fp8_quant`` producer). + + Parametrized over plain and Gemma norms: the Gemma AR+RMS-only patterns + match a prefix of these subgraphs, so this also pins their registration + order behind the quant-fused ones. """ with monkeypatch.context() as m: m.setenv("VLLM_ROCM_USE_AITER", "1") @@ -730,7 +779,7 @@ def run_torch_spawn(fn, nprocs): args=( num_processes, master_port, - TestAiterAllReduceRMSNormGroupQuantFP8Model, + test_model_cls, batch_size, seq_len, hidden_size, diff --git a/tests/distributed/test_aiter_rms_fallback.py b/tests/distributed/test_aiter_rms_fallback.py new file mode 100644 index 000000000000..4a0fc67873a1 --- /dev/null +++ b/tests/distributed/test_aiter_rms_fallback.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import sys +import types +from types import SimpleNamespace + +import pytest +import torch + +import vllm._aiter_ops as aiter_ops +import vllm.distributed as vllm_distributed +import vllm.kernels # noqa: F401 + + +def test_rocm_aiter_gemma_rmsnorm_falls_back_after_allreduce( + monkeypatch: pytest.MonkeyPatch, +): + class FakeAiterCA: + world_size = 8 + fully_connected = True + + def custom_fused_ar_rms(self, *args, **kwargs): + return None + + class FakeAiterAllReduce: + aiter_ca = FakeAiterCA() + + def use_1stage_fused_ar_rms(self, inp): + return False + + class FakeQuickReduce: + disabled = False + + def should_quick_allreduce(self, inp): + raise AssertionError("Gemma norm must not use fused QR+RMSNorm") + + device_comm = SimpleNamespace(qr_comm=FakeQuickReduce()) + monkeypatch.setattr( + vllm_distributed, + "get_tp_group", + lambda: SimpleNamespace(device_communicator=device_comm), + ) + monkeypatch.setattr( + aiter_ops.rocm_aiter_ops, + "get_aiter_allreduce", + lambda: FakeAiterAllReduce(), + ) + monkeypatch.setattr( + vllm_distributed, + "tensor_model_parallel_all_reduce", + lambda inp: inp + 2, + ) + + calls = [] + fake_aiter = types.SimpleNamespace() + + def fake_rmsnorm2d_fwd_with_add( + out, + inp, + residual, + residual_out, + weight, + epsilon, + *, + gemma_norm, + ): + calls.append((inp, weight, epsilon, gemma_norm)) + residual_out.copy_(inp + residual) + out.copy_(residual_out) + + fake_aiter.rmsnorm2d_fwd_with_add = fake_rmsnorm2d_fwd_with_add + monkeypatch.setitem(sys.modules, "aiter", fake_aiter) + + inp = torch.randn(8192, 16, dtype=torch.bfloat16) + residual = torch.randn_like(inp) + weight = torch.randn(16, dtype=torch.bfloat16) + + out, residual_out = aiter_ops._rocm_aiter_fused_allreduce_rmsnorm_impl( + inp, residual, weight, 1e-6, gemma_norm=True + ) + + assert len(calls) == 1 + assert calls[0][0] is not inp + assert calls[0][1] is weight + assert calls[0][2] == 1e-6 + assert calls[0][3] is True + torch.testing.assert_close(residual_out, inp + 2 + residual) + torch.testing.assert_close(out, residual_out) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 0ed84906185c..8af999906f28 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1086,7 +1086,10 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl( hidden_dim = input_.shape[-1] row_size = hidden_dim * input_.element_size() fused_qr_rmsnorm_ok = ( - qr_comm is not None + # aiter.qr_all_reduce_rmsnorm does not implement Gemma's + # (1 + weight) scaling. + not gemma_norm + and qr_comm is not None and not getattr(qr_comm, "disabled", True) and hasattr(qr_comm, "should_quick_allreduce") and qr_comm.should_quick_allreduce(input_) @@ -1133,7 +1136,16 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl( use_1stage=use_1stage, gemma_norm=gemma_norm, ) - assert result is not None + if result is None: + from vllm.distributed import tensor_model_parallel_all_reduce + + return torch.ops.vllm_aiter.fused_add_rms_norm( + tensor_model_parallel_all_reduce(input_), + residual, + weight, + epsilon, + gemma_norm, + ) return result[0], result[1] diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index b2961debcec2..59757deb2686 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -1149,6 +1149,8 @@ def __del__(self) -> None: # TODO: make BasePattern to inherit from VllmPatternReplacement class AiterAllreduceFusedRMSNormPattern(BasePattern, VllmPatternReplacement): + gemma_norm = False + def __init__( self, epsilon: float, @@ -1170,7 +1172,8 @@ def _pattern( input: torch.Tensor, weight: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: allreduce_output = tensor_model_parallel_all_reduce(input) - rms = vllm.ir.ops.rms_norm(allreduce_output, weight, self.epsilon) + norm_weight = weight.float() + 1.0 if self.gemma_norm else weight + rms = vllm.ir.ops.rms_norm(allreduce_output, norm_weight, self.epsilon) return rms, allreduce_output @@ -1187,13 +1190,22 @@ def _replacement( residual=residual, weight=weight.to(input.dtype), epsilon=self.epsilon, + gemma_norm=self.gemma_norm, ) return allreduce[0], allreduce[1] return _replacement +class AiterAllreduceFusedGemmaRMSNormPattern(AiterAllreduceFusedRMSNormPattern): + """Match all-reduce followed by GemmaRMSNorm without a residual.""" + + gemma_norm = True + + class AiterAllreduceFusedAddRMSNormPattern(BasePattern, VllmPatternReplacement): + gemma_norm = False + def __init__( self, epsilon: float, @@ -1216,8 +1228,9 @@ def _pattern( residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: allreduce_output = tensor_model_parallel_all_reduce(input) + norm_weight = weight.float() + 1.0 if self.gemma_norm else weight rms, residual = vllm.ir.ops.fused_add_rms_norm( - allreduce_output, residual, weight, self.epsilon + allreduce_output, residual, norm_weight, self.epsilon ) return rms, residual @@ -1233,12 +1246,19 @@ def _replacement( residual=residual, weight=weight.to(input.dtype), epsilon=self.epsilon, + gemma_norm=self.gemma_norm, ) return allreduce[0], allreduce[1] return _replacement +class AiterAllreduceFusedAddGemmaRMSNormPattern(AiterAllreduceFusedAddRMSNormPattern): + """Match all-reduce followed by fused-add GemmaRMSNorm.""" + + gemma_norm = True + + class AiterAllreduceFusedAddRMSNormOutputOnlyPattern( AiterAllreduceFusedAddRMSNormPattern ): @@ -1267,6 +1287,14 @@ def _replacement( return _replacement +class AiterAllreduceFusedAddGemmaRMSNormOutputOnlyPattern( + AiterAllreduceFusedAddRMSNormOutputOnlyPattern +): + """Match fused-add GemmaRMSNorm when its residual output is dead.""" + + gemma_norm = True + + class AiterAllreduceFusedRMSNormGroupQuantFP8Pattern( BasePattern, VllmPatternReplacement ): @@ -1608,7 +1636,7 @@ def __init__(self, config: VllmConfig) -> None: # AR+RMS-only fusion runs first and consumes the all_reduce node, # leaving the trailing quant op stranded as an unfused kernel. # Register larger subgraphs first (DeepSeek indexer fan-out, then - # quant-only AR+RMS+quant, then AR+RMS-only). + # quant-only AR+RMS+quant, then Gemma AR+RMS, then AR+RMS-only). if supports_per_group_quant: self.register( AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern( @@ -1632,6 +1660,30 @@ def __init__(self, config: VllmConfig) -> None: ) ) + # Must stay above the non-Gemma AR+RMS patterns below, which would + # otherwise match first. + self.register( + AiterAllreduceFusedGemmaRMSNormPattern( + epsilon, + self.model_dtype, + self.device, + ) + ) + self.register( + AiterAllreduceFusedAddGemmaRMSNormPattern( + epsilon, + self.model_dtype, + self.device, + ) + ) + self.register( + AiterAllreduceFusedAddGemmaRMSNormOutputOnlyPattern( + epsilon, + self.model_dtype, + self.device, + ) + ) + self.register( AiterAllreduceFusedRMSNormPattern( epsilon, diff --git a/vllm/kernels/aiter_ops.py b/vllm/kernels/aiter_ops.py index 402c9cac1752..9817f5adc9bb 100644 --- a/vllm/kernels/aiter_ops.py +++ b/vllm/kernels/aiter_ops.py @@ -136,6 +136,7 @@ def _rocm_aiter_rmsnorm2d_fwd_with_add_impl( residual: torch.Tensor, weight: torch.Tensor, variance_epsilon: float, + gemma_norm: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: from aiter import rmsnorm2d_fwd_with_add @@ -157,6 +158,7 @@ def _rocm_aiter_rmsnorm2d_fwd_with_add_impl( residual_out, # residual output weight, variance_epsilon, + gemma_norm=gemma_norm, ) return out.reshape(x_shape), residual_out.reshape(residual_shape) @@ -166,6 +168,7 @@ def _rocm_aiter_rmsnorm2d_fwd_with_add_fake( residual: torch.Tensor, weight: torch.Tensor, variance_epsilon: float, + gemma_norm: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: residual_out = torch.empty_like(residual) out = torch.empty_like(x) From 654b725ea9cb11085bce072a22a5e85709aafee6 Mon Sep 17 00:00:00 2001 From: Andrey Talman Date: Mon, 12 Oct 2026 16:23:53 -0400 Subject: [PATCH 334/624] [CI] Pace the CRCR nightly report under the relay rate limit (#61269) Signed-off-by: Andrey Talman --- .buildkite/scripts/crcr-report.sh | 7 ++++++- .buildkite/scripts/crcr_report.py | 16 ++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/.buildkite/scripts/crcr-report.sh b/.buildkite/scripts/crcr-report.sh index 0c114f4653fa..a1c5610a5106 100644 --- a/.buildkite/scripts/crcr-report.sh +++ b/.buildkite/scripts/crcr-report.sh @@ -227,7 +227,12 @@ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" # five-minute lifetime and the loop above can run to the build-age deadline. # Build 91827 waited 6h48m and then took a 401 on all 360 callbacks from a # token that had expired hours earlier. -OIDC_TOKEN="$(buildkite-agent oidc request-token --skip-redaction --audience "${AUDIENCE}" 2>/dev/null)" +# Posting is paced under the relay's rate limit, which takes several minutes +# for a full nightly, so ask for a longer lifetime than the default. Fall back +# to the default if this agent rejects --lifetime. +OIDC_TOKEN="$(buildkite-agent oidc request-token --skip-redaction --audience "${AUDIENCE}" \ + --lifetime 1800 2>/dev/null \ + || buildkite-agent oidc request-token --skip-redaction --audience "${AUDIENCE}" 2>/dev/null)" if [[ -z "${OIDC_TOKEN}" ]]; then echo "could not mint a Buildkite OIDC token after the wait -- skipping report" exit 0 diff --git a/.buildkite/scripts/crcr_report.py b/.buildkite/scripts/crcr_report.py index 52f0649f011b..4fab92a2c56b 100644 --- a/.buildkite/scripts/crcr_report.py +++ b/.buildkite/scripts/crcr_report.py @@ -24,7 +24,9 @@ import argparse import json +import os import sys +import time import urllib.error import urllib.request @@ -43,6 +45,14 @@ } +# The relay rate-limits callbacks per repo over a sliding 60 s window +# (RATE_LIMIT_PER_MIN in pytorch/test-infra), and a nightly has far more jobs +# than that. Pace below the limit, and back off on a 429 for about a minute in +# total so the last retry lands in a new window. +_POST_INTERVAL_S = float(os.environ.get("CRCR_POST_INTERVAL_S", "0.65")) +_RATE_LIMIT_BACKOFF_S = (2, 4, 8, 16, 32) + + def post(url: str, token: str, body: dict, timeout: int = 30) -> tuple[int, str]: req = urllib.request.Request( url, @@ -113,6 +123,12 @@ def main() -> int: ok += 1 continue status, detail = post(args.callback_url, args.oidc_token, body) + for delay in _RATE_LIMIT_BACKOFF_S: + if status != 429: + break + time.sleep(delay) + status, detail = post(args.callback_url, args.oidc_token, body) + time.sleep(_POST_INTERVAL_S) if 200 <= status < 300: ok += 1 else: From deedc7be008fe93473ca621936cde4f3c6ca1610 Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Tue, 13 Oct 2026 05:24:12 +0900 Subject: [PATCH 335/624] [Perf][Humming] Skip the redundant MoE input copy before the w13 GEMM (#59340) Signed-off-by: lesj0610 Co-authored-by: lesj0610 Co-authored-by: Misha Goin --- .../layers/fused_moe/experts/fused_humming_moe.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py index 66e03229a9e1..f27171ef0de5 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py @@ -214,12 +214,20 @@ def process_input( if scatter_idx is not None: layout = "scatter" + hadamard_block_size = getattr(quant_config, f"{prefix}_hadamard_block_size") + has_input_processing = ( + mode.should_quantize + or (hadamard_block_size is not None and hadamard_block_size > 1) + or activation is not None + or scatter_idx is not None + or num_valid_tokens is not None + ) inputs, group_scales, token_scales = may_process_input( config, inputs=inputs, - outputs=quanted_input, + outputs=quanted_input if has_input_processing else None, token_scales=scale2 if mode.has_secondary_scale else scale, - hadamard_block_size=getattr(quant_config, f"{prefix}_hadamard_block_size"), + hadamard_block_size=hadamard_block_size, layout=layout, expert_tokens=expert_tokens if layout == "grouped_mask" else None, scatter_idx=scatter_idx, From fe341a53dda00472c281a2d1f4bd34a74e7eae6e Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Mon, 12 Oct 2026 16:49:49 -0400 Subject: [PATCH 336/624] [HiSparse] Report host pool in KV cache capacity logs (#60285) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- vllm/v1/core/kv_cache_utils.py | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 919a7b4b3f52..03c6c5ff62ac 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -2308,6 +2308,11 @@ def get_kv_cache_capacity( return int(max_concurrency * max_model_len), max_concurrency +def _kv_cache_device_name(vllm_config: VllmConfig) -> str: + device_type = vllm_config.device_config.device_type + return "GPU" if device_type == "cuda" else device_type.upper() + + def update_kv_cache_capacity( vllm_config: VllmConfig, kv_cache_config: KVCacheConfig ) -> None: @@ -2316,11 +2321,13 @@ def update_kv_cache_capacity( vllm_config.cache_config.kv_cache_size_tokens = num_tokens vllm_config.cache_config.kv_cache_max_concurrency = max_concurrency max_model_len = vllm_config.model_config.max_model_len - device_type = vllm_config.device_config.device_type + pools = _kv_cache_device_name(vllm_config) + if vllm_config.attention_config.hisparse_config is not None: + pools += " + host" logger.info_once( "%s KV cache size: %s tokens, " "Maximum concurrency for %s tokens per request: %.2fx", - "GPU" if device_type == "cuda" else device_type.upper(), + pools, f"{num_tokens:,}", f"{max_model_len:,}", max_concurrency, @@ -2449,22 +2456,30 @@ def _auto_fit_max_model_len( "to serve even a single token. Try increasing `gpu_memory_utilization`." ) + kv_cache_memory = ( + f"{format_gib(limiting_worker_mem)} GiB {_kv_cache_device_name(vllm_config)}" + ) + if vllm_config.attention_config.hisparse_config is not None: + host_pool_gib = format_gib(get_hisparse_host_pool_bytes(vllm_config)) + kv_cache_memory += f" + {host_pool_gib} GiB host" + if auto_fit_max >= original_max: # The model's full context length fits in memory logger.info_once( "Auto-fit max_model_len: full model context length %d fits in " - "available GPU memory", + "available KV cache memory (%s)", original_max, + kv_cache_memory, ) else: # Need to reduce max_model_len to fit in memory vllm_config.model_config.max_model_len = auto_fit_max logger.info_once( "Auto-fit max_model_len: reduced from %d to %d to fit in " - "available GPU memory (%s GiB available for KV cache)", + "available KV cache memory (%s)", original_max, auto_fit_max, - format_gib(limiting_worker_mem), + kv_cache_memory, ) From 69d54eb078db7ff9dd63f94d55a139b182c1776d Mon Sep 17 00:00:00 2001 From: Almog Tavor <70065337+almogtavor@users.noreply.github.com> Date: Tue, 13 Oct 2026 00:26:43 +0300 Subject: [PATCH 337/624] [Docs] Fix broken llm-d.ai links on the llm-d integration page (#61022) Signed-off-by: Almog Tavor Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- docs/deployment/integrations/llm-d.md | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/deployment/integrations/llm-d.md b/docs/deployment/integrations/llm-d.md index 7d261eb910b6..dc335edb0399 100644 --- a/docs/deployment/integrations/llm-d.md +++ b/docs/deployment/integrations/llm-d.md @@ -8,11 +8,11 @@ It is a [CNCF Sandbox project](https://www.cncf.io/blog/2026/03/24/welcome-llm-d A single vLLM server is fast, but at scale the picture changes: across many replicas, cache locality breaks under round-robin load balancing, long prompts inflate time-to-first-token, and accelerators sit underused. llm-d adds the cluster-level layer that vLLM does not aim to provide on its own: -- **[Prefix-aware routing](https://llm-d.ai/docs/guides/precise-prefix-cache-aware).** Instead of round-robin, llm-d reads vLLM's KV-cache events and routes each request to the replica that already holds its prefix, reusing cache instead of recomputing it. -- **[Distributed KV-cache management](https://llm-d.ai/docs/guides#advanced-kv-cache-management).** A global index tracks which token blocks live on which replica, and [tiered offloading](https://llm-d.ai/docs/guides/tiered-prefix-cache) spills cache to CPU memory or local SSD, extending the working set beyond accelerator HBM. -- **[Prefill/decode disaggregation](https://llm-d.ai/docs/guides/pd-disaggregation).** Prompt processing and token generation run on separate vLLM workers, with KV-cache moved over the vLLM [NIXL connector](https://docs.vllm.ai/en/latest/features/nixl_connector_usage/), lowering TTFT and steadying per-token latency on long prompts. -- **[Wide expert-parallelism](https://llm-d.ai/docs/guides/wide-expert-parallelism).** Serve large Mixture-of-Experts models such as DeepSeek-R1 and GPT-OSS across nodes with combined data and expert parallelism, for more KV-cache capacity and throughput. -- **SLO-aware [autoscaling](https://llm-d.ai/docs/guides/workload-autoscaling) and [flow control](https://llm-d.ai/docs/guides/flow-control).** Scale vLLM pools on real inference signals (queue depth, true demand) rather than raw GPU utilization, with multi-tenant fairness and priority dispatch. +- **[Prefix-aware routing](https://llm-d.ai/docs/well-lit-paths/foundations/precise-prefix-cache-routing).** Instead of round-robin, llm-d reads vLLM's KV-cache events and routes each request to the replica that already holds its prefix, reusing cache instead of recomputing it. +- **[Distributed KV-cache management](https://llm-d.ai/docs/architecture/advanced/kv-management).** A global index tracks which token blocks live on which replica, and [tiered offloading](https://llm-d.ai/docs/well-lit-paths/foundations/tiered-prefix-cache) spills cache to CPU memory or local SSD, extending the working set beyond accelerator HBM. +- **[Prefill/decode disaggregation](https://llm-d.ai/docs/well-lit-paths/foundations/pd-disaggregation).** Prompt processing and token generation run on separate vLLM workers, with KV-cache moved over the vLLM [NIXL connector](https://docs.vllm.ai/en/latest/features/nixl_connector_usage/), lowering TTFT and steadying per-token latency on long prompts. +- **[Wide expert-parallelism](https://llm-d.ai/docs/well-lit-paths/foundations/wide-expert-parallelism).** Serve large Mixture-of-Experts models such as DeepSeek-R1 and GPT-OSS across nodes with combined data and expert parallelism, for more KV-cache capacity and throughput. +- **SLO-aware [autoscaling](https://llm-d.ai/docs/well-lit-paths/foundations/workload-autoscaling) and [flow control](https://llm-d.ai/docs/well-lit-paths/foundations/flow-control).** Scale vLLM pools on real inference signals (queue depth, true demand) rather than raw GPU utilization, with multi-tenant fairness and priority dispatch. These are composable. Most teams start by adding prefix-aware routing over an existing vLLM pool, then layer in the rest as specific bottlenecks appear. @@ -28,9 +28,9 @@ See the [full list](https://github.com/llm-d/llm-d#performance-highlights) and r ## Get started -1. Deploy the [Optimized Baseline](https://llm-d.ai/docs/guides/optimized-baseline) with the [Quickstart](https://llm-d.ai/docs/getting-started/quickstart). It stands up an intelligent router over a vLLM pool on Kubernetes in a tested configuration. -2. Browse the [well-lit path guides](https://llm-d.ai/docs/guides), each a tested recipe for one of the capabilities above, and add the optimization that fits your workload. -3. Read the [Introduction](https://llm-d.ai/docs/getting-started) and [Architecture overview](https://llm-d.ai/docs/architecture) to see how the pieces wrap your vLLM deployment. +1. Deploy the [Optimized Baseline](https://llm-d.ai/docs/well-lit-paths/foundations/optimized-baseline) with the [Quickstart](https://llm-d.ai/docs/getting-started/quickstart). It stands up an intelligent router over a vLLM pool on Kubernetes in a tested configuration. +2. Browse the [well-lit path guides](https://llm-d.ai/docs/well-lit-paths), each a tested recipe for one of the capabilities above, and add the optimization that fits your workload. +3. Read the [Introduction](https://llm-d.ai/docs) and [Architecture overview](https://llm-d.ai/docs/architecture) to see how the pieces wrap your vLLM deployment. You can also deploy vLLM with llm-d via [KServe's LLMInferenceService](https://kserve.github.io/website/docs/model-serving/generative-inference/llmisvc/llmisvc-overview). From 44f0ddfd7b61e200aac323231dfc127d4cd53cfb Mon Sep 17 00:00:00 2001 From: Maroon Ayoub Date: Tue, 13 Oct 2026 00:28:35 +0300 Subject: [PATCH 338/624] [Bugfix][Parser] Stop re-parsing the whole GLM tool argument per chunk (#61253) Signed-off-by: Maroon Ayoub Co-authored-by: Flora Feng <4florafeng@gmail.com> --- tests/parser/engine/test_glm47_moe.py | 124 ++++++++++++++++++++- tests/parser/engine/test_parser_engine.py | 2 +- vllm/parser/engine/parser_engine.py | 81 ++++++++++---- vllm/parser/engine/parser_engine_config.py | 9 ++ vllm/parser/glm47_moe.py | 36 +++++- 5 files changed, 226 insertions(+), 26 deletions(-) diff --git a/tests/parser/engine/test_glm47_moe.py b/tests/parser/engine/test_glm47_moe.py index 58fe7aa28cbd..9f9d65ef2048 100644 --- a/tests/parser/engine/test_glm47_moe.py +++ b/tests/parser/engine/test_glm47_moe.py @@ -5,18 +5,27 @@ GLM can transition straight from reasoning into a tool call without ever emitting ````, so ```` is an implicit reasoning terminator. These cover that path plus the multi-turn prompts where an -earlier turn's markers must not be read as the current turn's state. +earlier turn's markers must not be read as the current turn's state. The +streaming tests check that long tool arguments stream the same JSON as a +non-streaming parse, without converting the whole argument on every delta. """ +import dataclasses + import pytest from tests.parser.engine.conftest import make_mock_tokenizer +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionToolsParam, +) from vllm.parser.glm47_moe import ( THINK_END, THINK_START, TOOL_CALL_END, TOOL_CALL_START, Glm47MoeParser, + _glm47_arg_value_extender, + glm47_moe_config, ) THINK_S, THINK_E = 1000, 1001 @@ -106,3 +115,116 @@ def test_no_markers_returns_input_ids(self, parser): def test_thinking_disabled(self, no_thinking_parser): ids = [THINK_S, TEXT, TOOL_S] assert no_thinking_parser.extract_content_ids(ids) == ids + + +WRITE_TOOL = ChatCompletionToolsParam( + type="function", + function={ + "name": "Write", + "parameters": { + "type": "object", + "properties": { + "file_path": {"type": "string"}, + "line_count": {"type": "integer"}, + "content": {"type": "string"}, + }, + }, + }, +) +CALL_HEAD = ( + f"{THINK_END}{TOOL_CALL_START}Write\n" + "file_path/repo/app.py\n" + "line_count120\n" + "content" +) +CALL_TAIL = f"\n{TOOL_CALL_END}" +# Escapes, non-ASCII, and tag-like text that must stay inside the value. +CONTENT_LINE = 'if a < b > c: print("x\\y") # \t\u00e9\n' + + +def _pieces(text: str, size: int) -> list[str]: + return [text[i : i + size] for i in range(0, len(text), size)] + + +def _stream_arguments(parser, request, chunks: list[str]) -> str: + args = [] + for i, chunk in enumerate(chunks): + delta = parser.parse_delta(chunk, [], request, finished=i == len(chunks) - 1) + for tool_call in (delta and delta.tool_calls) or []: + if tool_call.function and tool_call.function.arguments: + args.append(tool_call.function.arguments) + return "".join(args) + + +class TestStreamedToolArguments: + @pytest.mark.parametrize( + "tools, line_count", + [([], '"120"'), ([WRITE_TOOL], "120")], + ids=["no_schema", "schema"], + ) + def test_long_call_streams_non_streaming_arguments( + self, mock_request, tools, line_count + ): + mock_request.tools = tools + text = CALL_HEAD + CONTENT_LINE * 400 + CALL_TAIL + + streamed = _stream_arguments( + Glm47MoeParser(make_mock_tokenizer(VOCAB)), mock_request, _pieces(text, 7) + ) + result = Glm47MoeParser(make_mock_tokenizer(VOCAB)).extract_tool_calls( + text, mock_request + ) + + assert streamed == result.tool_calls[0].function.arguments + assert f'"line_count": {line_count},' in streamed + + def test_conversions_do_not_grow_with_argument_length(self, mock_request): + """The open value streams without re-running the converter.""" + mock_request.tools = [WRITE_TOOL] + + def streaming_conversions(lines: int) -> int: + calls = 0 + base = glm47_moe_config() + + def counting_converter(raw_args: str, partial: bool) -> str: + nonlocal calls + if partial: + calls += 1 + return base.arg_converter(raw_args, partial) + + parser = Glm47MoeParser( + make_mock_tokenizer(VOCAB), + parser_engine_config=dataclasses.replace( + base, arg_converter=counting_converter + ), + ) + chunks = [ + *_pieces(CALL_HEAD, 4), + *_pieces(CONTENT_LINE * lines, 4), + *_pieces(CALL_TAIL, 4), + ] + _stream_arguments(parser, mock_request, chunks) + return calls + + assert streaming_conversions(lines=40) == streaming_conversions(lines=400) + + def test_value_extender_stops_at_split_value_end(self): + extend = _glm47_arg_value_extender("kx") + assert [extend(d) for d in ['a"\n', ""]] == [ + 'a\\"\\n', + "kx", + "k1" + "j2" + "k3", + ], + ids=["closed_value", "repeated_key_not_last"], + ) + def test_no_value_extender_unless_open_value_is_last(self, raw_args): + assert _glm47_arg_value_extender(raw_args) is None diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py index dc91cef74028..25dfd546b018 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -1711,7 +1711,7 @@ class TestSafeArgPrefix: ], ) def test_safe_arg_prefix(self, json_str, expected): - assert ParserEngine._safe_arg_prefix(json_str) == expected + assert ParserEngine._safe_arg_prefix(json_str)[0] == expected # ── Coercion instability regression tests ──────────────────────── diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index c74be931d256..eaf9e4519e1d 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -24,7 +24,11 @@ from vllm.logger import init_logger from vllm.parser.abstract_parser import Parser, StreamState from vllm.parser.engine.events import EventType, SemanticEvent -from vllm.parser.engine.parser_engine_config import ParserEngineConfig, ParserState +from vllm.parser.engine.parser_engine_config import ( + ArgValueExtender, + ParserEngineConfig, + ParserState, +) from vllm.parser.engine.streaming_parser_engine import StreamingParserEngine from vllm.renderers.chat_utils import get_tool_call_id_type, make_tool_call_id from vllm.tool_parsers.utils import ( @@ -60,7 +64,10 @@ class ToolCallSlot: "_args_joined", "name_sent", "string_keys", - "streamed_json", + "_streamed_parts", + "_streamed_joined", + "extend_open_value", + "open_value_streams", ) def __init__(self) -> None: @@ -70,7 +77,10 @@ def __init__(self) -> None: self._args_joined: str | None = "" self.name_sent: bool = False self.string_keys: set[str] | None = None - self.streamed_json: str = "" + self._streamed_parts: list[str] = [] + self._streamed_joined: str | None = "" + self.extend_open_value: ArgValueExtender | None = None + self.open_value_streams: bool = False @property def args(self) -> str: @@ -82,6 +92,21 @@ def append_args(self, value: str) -> None: self._args_parts.append(value) self._args_joined = None + @property + def streamed_json(self) -> str: + if self._streamed_joined is None: + self._streamed_joined = "".join(self._streamed_parts) + return self._streamed_joined + + @streamed_json.setter + def streamed_json(self, value: str) -> None: + self._streamed_parts = [value] + self._streamed_joined = value + + def append_streamed(self, value: str) -> None: + self._streamed_parts.append(value) + self._streamed_joined = None + class ParserEngine(Parser): """A :class:`Parser` backed by a single declarative engine config. @@ -133,6 +158,7 @@ def __init__( self._arg_converter = parser_engine_config.arg_converter self._arg_structural_chars = parser_engine_config.arg_structural_chars + self._arg_value_extender = parser_engine_config.arg_value_extender self._stream_arg_deltas = parser_engine_config.stream_arg_deltas self._strip_trailing_reasoning_ws = ( parser_engine_config.strip_trailing_reasoning_whitespace @@ -303,7 +329,9 @@ def _coerce_dict(args: dict, properties: dict) -> tuple[dict, bool]: return args, changed @staticmethod - def _safe_arg_prefix(json_str: str, string_keys: set[str] | None = None) -> str: + def _safe_arg_prefix( + json_str: str, string_keys: set[str] | None = None + ) -> tuple[str, bool]: """Return the prefix of *json_str* up to the last top-level value. Middle values (followed by a comma) are stable across streaming @@ -313,6 +341,9 @@ def _safe_arg_prefix(json_str: str, string_keys: set[str] | None = None) -> str: String values for keys in ``string_keys`` are prefix-stable, so stream their unterminated content instead of buffering long arguments until the closing tag arrives. + + Also returns whether the prefix ends inside that streamed trailing + string value. """ last_colon = -1 last_key: str | None = None @@ -345,14 +376,14 @@ def _safe_arg_prefix(json_str: str, string_keys: set[str] | None = None) -> str: last_key = pending_key pending_key = None if last_colon < 0: - return "" + return "", False end = last_colon + 1 while end < len(json_str) and json_str[end] in (" ", "\t", "\n", "\r"): end += 1 if end >= len(json_str) or json_str[end] != '"': - return json_str[:end] + return json_str[:end], False if string_keys is not None and last_key not in string_keys: - return json_str[:end] + return json_str[:end], False escape = False for i in range(end + 1, len(json_str)): @@ -364,8 +395,8 @@ def _safe_arg_prefix(json_str: str, string_keys: set[str] | None = None) -> str: escape = True continue if c == '"': - return json_str[:i] - return json_str + return json_str[:i], True + return json_str, True @staticmethod def _streamable_string_keys(properties: dict) -> set[str] | None: @@ -1027,6 +1058,15 @@ def _compute_arg_delta(self, idx: int, raw_delta: str) -> str | None: return None slot = self._tool_slots[idx] + if slot.extend_open_value is not None: + diff = slot.extend_open_value(raw_delta) + if diff is not None: + if not diff or not slot.open_value_streams: + return None + slot.append_streamed(diff) + return diff + slot.extend_open_value = None + try: current_json = converter(slot.args, True) except (json.JSONDecodeError, ValueError, TypeError): @@ -1040,22 +1080,19 @@ def _compute_arg_delta(self, idx: int, raw_delta: str) -> str | None: current_json = self._fix_arg_types(current_json, slot.name) prev = slot.streamed_json - safe_json = self._safe_arg_prefix(current_json, slot.string_keys) + safe_json, value_streams = self._safe_arg_prefix(current_json, slot.string_keys) - if not safe_json or safe_json == prev: + if not safe_json: + return None + if prev and not safe_json.startswith(prev): return None - if prev: - if not safe_json.startswith(prev): - return None - diff = safe_json[len(prev) :] - else: - diff = safe_json - - if diff: - slot.streamed_json = safe_json - return diff - return None + slot.streamed_json = safe_json + if self._arg_value_extender is not None: + # Deltas that only extend the open value skip the conversion above. + slot.extend_open_value = self._arg_value_extender(slot.args) + slot.open_value_streams = value_streams + return safe_json[len(prev) :] or None def _flush_arg_converter(self, idx: int) -> str | None: converter = self._arg_converter diff --git a/vllm/parser/engine/parser_engine_config.py b/vllm/parser/engine/parser_engine_config.py index 215fa0d9bdf5..2bb2ae267e6f 100644 --- a/vllm/parser/engine/parser_engine_config.py +++ b/vllm/parser/engine/parser_engine_config.py @@ -24,6 +24,10 @@ from vllm.parser.engine.events import EventType +# Maps a raw argument delta to the JSON text it appends to the open trailing +# value, or returns None when the delta may change anything else. +ArgValueExtender = Callable[[str], str | None] + class ParserState(Enum): CONTENT = auto() @@ -90,6 +94,11 @@ class ParserEngineConfig: arg_structural_chars: frozenset[str] | None = None + # Called with the raw argument text when its converted JSON ends in an + # open value; returns an extender for that value, or None. Streaming + # skips ``arg_converter`` for every delta the extender accepts. + arg_value_extender: Callable[[str], ArgValueExtender | None] | None = None + # Special tokens exempt from auto-drop but not state-machine terminals. preserve_tokens: frozenset[str] = field(default_factory=frozenset) diff --git a/vllm/parser/glm47_moe.py b/vllm/parser/glm47_moe.py index 9a18228fb0d1..44377e62b707 100644 --- a/vllm/parser/glm47_moe.py +++ b/vllm/parser/glm47_moe.py @@ -24,6 +24,7 @@ from vllm.parser.engine.events import EventType from vllm.parser.engine.parser_engine import ParserEngine from vllm.parser.engine.parser_engine_config import ( + ArgValueExtender, ParserEngineConfig, ParserState, Transition, @@ -63,12 +64,16 @@ ) -def _glm47_arg_converter(raw_args: str, partial: bool) -> str: +def _glm47_parse_args( + raw_args: str, partial: bool +) -> tuple[dict[str, object], str | None]: + """Return the arguments and the key of the open trailing value, if any.""" params: dict[str, object] = {} for match in _ARG_RE.finditer(raw_args): params[match.group("key").strip()] = match.group("value") + open_key = None if partial: remaining = _ARG_RE.sub("", raw_args) match = _PARTIAL_ARG_RE.search(remaining) @@ -76,8 +81,34 @@ def _glm47_arg_converter(raw_args: str, partial: bool) -> str: key = match.group("key").strip() if key: params[key] = match.group("value") + open_key = key + + return params, open_key + + +def _glm47_arg_converter(raw_args: str, partial: bool) -> str: + return json.dumps(_glm47_parse_args(raw_args, partial)[0], ensure_ascii=False) + + +def _glm47_arg_value_extender(raw_args: str) -> ArgValueExtender | None: + # The open value runs to the end of raw_args and is the last JSON value, + # so a delta appends its JSON-escaped text unless it closes the value. + params, open_key = _glm47_parse_args(raw_args, partial=True) + if open_key is None or next(reversed(params)) != open_key: + return None + + keep = len(ARG_VALUE_END) - 1 + tail = raw_args[-keep:] + + def extend(delta: str) -> str | None: + nonlocal tail + text = tail + delta + if ARG_VALUE_END in text: + return None + tail = text[-keep:] + return json.dumps(delta, ensure_ascii=False)[1:-1] - return json.dumps(params, ensure_ascii=False) + return extend def _glm53_always_thinks(tokenizer: TokenizerLike) -> bool: @@ -194,6 +225,7 @@ def glm47_moe_config(thinking: bool = True) -> ParserEngineConfig: }, turn_boundary_tokens=GLM_TURN_BOUNDARIES, arg_converter=_glm47_arg_converter, + arg_value_extender=_glm47_arg_value_extender, stream_arg_deltas=True, tool_args_json=False, validate_tool_names=True, From 802feb6770a0522ce8d2999262d95f4618e04adf Mon Sep 17 00:00:00 2001 From: Akshay Viswanathan Date: Tue, 13 Oct 2026 00:29:31 +0200 Subject: [PATCH 339/624] [Bugfix][KVConnector] Fix MoRI-IO WRITE-mode requests stranded in WAITING_FOR_REMOTE_KVS (#57705) Signed-off-by: Akshay Viswanathan <4953343+akshayv@users.noreply.github.com> Co-authored-by: Akshay Viswanathan <4953343+akshayv@users.noreply.github.com> Co-authored-by: Cursor Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> --- .../unit/test_moriio_write_done_routing.py | 104 +++++++++++++++ .../kv_connector/v1/moriio/moriio_common.py | 9 +- .../v1/moriio/moriio_connector.py | 73 ++++++++--- .../kv_connector/v1/moriio/moriio_engine.py | 121 +++++++++++++----- 4 files changed, 255 insertions(+), 52 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py b/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py index adac007d19eb..57b6893869eb 100644 --- a/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py +++ b/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py @@ -3,20 +3,30 @@ """Unit tests for mori-io write_done routing (#51681).""" import threading +import time +from collections import OrderedDict from types import SimpleNamespace from typing import Any +from unittest.mock import patch import pytest +from vllm.distributed.kv_transfer.kv_connector.v1.moriio import moriio_engine from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( + ROLE, + MoRIIOMode, RemoteAllocInfo, ReqMeta, WriteTask, get_port_offset, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( + MoRIIOConnectorScheduler, MoRIIOConnectorWorker, ) +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import ( + MoRIIOWrapper, +) from .utils import make_moriio_writer @@ -79,6 +89,7 @@ def _write_task( *, multi_pod_hosts: list[str] | None = None, remote_dp_size_local: int = 0, + remote_dp_size: int = 1, request_id: str = "req", ) -> WriteTask: return WriteTask( @@ -93,9 +104,21 @@ def _write_task( remote_ip=remote_ip, multi_pod_hosts=multi_pod_hosts or [], remote_dp_size_local=remote_dp_size_local, + remote_dp_size=remote_dp_size, ) +def _real_wrapper() -> MoRIIOWrapper: + wrapper = MoRIIOWrapper.__new__(MoRIIOWrapper) + wrapper.lock = threading.Lock() + wrapper.done_req_ids = [] + wrapper.done_remote_allocate_req_dict = {} + wrapper.done_write_cache_req_ids = [] + wrapper.failed_write_cache_req_ids = [] + wrapper._terminal_transfer_ids = OrderedDict() + return wrapper + + def test_execute_write_task_ip_is_per_task_under_overwrite(): alloc = {"tA": RemoteAllocInfo(block_ids=None, decode_dp_rank=0)} worker = _fake_worker(alloc) @@ -269,3 +292,84 @@ def _capture_schedule_write_blocks(**kwargs: Any) -> None: assert captured["remote_dp_size_local"] == 4 assert captured["local_block_ids"] == [1] assert captured["remote_block_ids"] == [2] + + +def test_deferred_write_timeout_fails_transfer_on_every_decode_rank(): + wrapper = _real_wrapper() + sent: list[tuple[str, int, str]] = [] + wrapper.send_notify = lambda tid, ip, port, message_type: sent.append( + (ip, port, message_type) + ) + writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper, tp_rank=0)) + task = _write_task( + "tA", + "10.0.0.1", + multi_pod_hosts=["10.0.0.1", "10.0.0.2"], + remote_dp_size_local=2, + remote_dp_size=4, + ) + task.enqueue_time = time.perf_counter() - 2 * writer._defer_timeout + writer._deferred_tasks.append(task) + + writer._process_deferred_tasks() + + assert writer._deferred_tasks == [] + assert [ack.transfer_id for ack in wrapper.done_req_ids] == ["tA"] + # Global DP ranks 0-3, two per pod: the owning rank is unknown. + assert sent == [ + (host, 20000 + get_port_offset(local_rank, 0), "write_failed") + for host in ("10.0.0.1", "10.0.0.2") + for local_rank in (0, 1) + ] + + # A late decode allocation must not resurrect a write into freed blocks. + with patch.object(moriio_engine, "get_role", return_value=ROLE.PRODUCER): + wrapper._handle_remote_blocks_message( + {"transfer_id": "tA", "block_notify_list": [7], "decode_rank": 1} + ) + assert wrapper.done_remote_allocate_req_dict == {} + + +def test_write_scheduler_waits_for_worker_ack_past_defer_deadline(): + scheduler = SimpleNamespace( + is_producer=True, + mode=MoRIIOMode.WRITE, + _defer_timeout=60.0, + _deferred_send_deadlines={"req-1": (time.monotonic() - 1.0, "tA")}, + _pending_sent_acks={}, + unmap_request_id=lambda *args, **kwargs: None, + ) + output = SimpleNamespace(finished_sending=set()) + + MoRIIOConnectorScheduler.update_connector_output(scheduler, output) + assert output.finished_sending is None + assert "req-1" in scheduler._deferred_send_deadlines + + output.finished_sending = {"req-1"} + MoRIIOConnectorScheduler.update_connector_output(scheduler, output) + assert output.finished_sending == {"req-1"} + assert scheduler._deferred_send_deadlines == {} + + +def test_write_failed_is_reported_as_failed_recving_on_owning_rank(): + wrapper = _real_wrapper() + with patch.object(moriio_engine, "get_role", return_value=ROLE.CONSUMER): + for tid in ("tA", "tOtherRank"): + wrapper._handle_structured_message( + {"type": "write_failed", "transfer_id": tid} + ) + worker = SimpleNamespace( + mode=MoRIIOMode.WRITE, + is_producer=False, + moriio_wrapper=wrapper, + transfer_id_to_request_id={"tA": "req-1"}, + _unmatched_write_failures=OrderedDict(), + get_finished=lambda: (set(), set()), + ) + + results = MoRIIOConnectorWorker.get_transfer_results(worker) + + assert results.finished_recving == {"req-1"} + assert results.failed_recving == {"req-1"} + assert wrapper.pop_finished_write_req_ids() == set() + assert list(worker._unmatched_write_failures) == ["tOtherRank"] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index ae1eecf6e55d..e82d0668189d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -71,6 +71,7 @@ class WriteTask: remote_ip: str multi_pod_hosts: list[str] = field(default_factory=list) remote_dp_size_local: int = 0 + remote_dp_size: int = 1 enqueue_time: float = field(default_factory=time.perf_counter) retried: int = 0 @@ -307,8 +308,8 @@ def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig": # WRITE mode. # transfer_timeout -> Timeout for waiting_for_transfer_complete before # raising TransferError (sec). - # defer_timeout -> Timeout before a deferred send with no finished_sending - # notification is reaped and its blocks force-freed (sec). + # defer_timeout -> Timeout for deferred WRITE allocation or a missing + # READ release notification (sec). # recv_abort_timeout -> Timeout before an in-flight recv whose RDMA # completion never arrived is aborted (sec). @@ -414,8 +415,8 @@ class MoRIIOConstants: # Timeout (seconds) for waiting_for_transfer_complete before raising TransferError. # Overridable via kv_connector_extra_config["transfer_timeout"]. DEFAULT_TRANSFER_TIMEOUT = 30.0 - # Timeout (seconds) before a deferred send with no finished_sending - # notification is reaped and its blocks force-freed. + # Timeout (seconds) for deferred WRITE allocation or a missing READ release + # notification. # Overridable via kv_connector_extra_config["defer_timeout"]. DEFAULT_DEFER_TIMEOUT = 60.0 # Timeout (seconds) before an in-flight recv whose RDMA completion was lost diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 082643778e7d..c398cb20a8b0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -5,7 +5,7 @@ import queue import threading import time -from collections import defaultdict +from collections import OrderedDict, defaultdict from collections.abc import Collection from concurrent.futures import Future, ThreadPoolExecutor from typing import TYPE_CHECKING, Any, cast @@ -25,6 +25,7 @@ KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole, + KVConnectorTransferResults, SupportsHMA, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( @@ -116,6 +117,9 @@ _SQ_FULL_BACKOFF_INITIAL_S = 0.001 _SQ_FULL_BACKOFF_MAX_S = 0.05 _MAX_LOCAL_MAMBA_TAIL_BLOCKS = 1 +# A WRITE failure is broadcast to every decode DP rank, so ranks that never +# see the request would otherwise keep its transfer id forever. +_MAX_UNMATCHED_WRITE_FAILURES = 4096 try: @@ -356,6 +360,12 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: assert self.connector_worker is not None return self.connector_worker.get_finished() + def get_transfer_results( + self, finished_req_ids: set[str] + ) -> KVConnectorTransferResults: + assert self.connector_worker is not None + return self.connector_worker.get_transfer_results() + def get_block_ids_with_load_errors(self) -> set[int]: """Blocks whose MoRIIO read failed, for the scheduler to recompute.""" if self.connector_worker is None: @@ -667,12 +677,10 @@ def __init__( set_role(ROLE.CONSUMER) # Reqs to send and their expiration time self._reqs_need_send: dict[ReqId, float] = {} - # Deadlines for requests whose block freeing was deferred. - # Survives across scheduler steps. If the worker never reports - # finished_sending before the deadline, we inject them into - # connector_output.finished_sending so the scheduler frees the blocks to avoid - # hanging indefinitely waiting for a free notification that never comes. - # Value: (deadline, transfer_id) for unmapping after mutation. + # Requests whose block freeing is deferred. READ mode uses the deadline + # as a missing-ACK watchdog; WRITE mode waits for the worker ACK because + # only the worker can mark the transfer terminal before blocks are freed. + # Value: (deadline, transfer_id) for unmapping after completion. self._deferred_send_deadlines: dict[ReqId, tuple[float, TransferId | None]] = {} self._defer_timeout = float( self.kv_transfer_config.kv_connector_extra_config.get( @@ -1131,7 +1139,14 @@ def update_state_after_alloc( _should_notify = self._global_dp_rank == remote_dp_rank else: _should_notify = self._is_kv_master - if _should_notify: + if _should_notify and num_external_tokens == 0: + # Nothing left to fetch (full local prefix hit), so release + # the producer's blocks instead of sending an empty + # allocation, which it rejects. + self._release_write_prefill_blocks( + request.request_id, request.kv_transfer_params + ) + elif _should_notify: peer_zmq = get_peer_zmq_from_request_id( request.request_id, is_producer=False ) @@ -1159,12 +1174,6 @@ def update_state_after_alloc( f"remote_notify_port={remote_notify_port!r})" ) - # num_external_tokens == 0: nothing to push, so don't tell - # the producer to write into these blocks. - block_notify_list = ( - blocks.get_block_ids()[0] if num_external_tokens > 0 else [] - ) - # Wide-EP multi-pod: a pod binds notify sockets only for # its LOCAL ranks, so the port offset must use the per-pod # local rank (% dp_local), not the global rank. Single-pod @@ -1191,7 +1200,7 @@ def update_state_after_alloc( self.send_notify_block( req_id=request.request_id, transfer_id=request.kv_transfer_params["transfer_id"], - block_notify_list=block_notify_list, + block_notify_list=blocks.get_block_ids()[0], host=_notify_host, port=target_port, ) @@ -1541,8 +1550,9 @@ def update_connector_output(self, connector_output: KVConnectorOutput) -> None: * An ACK that arrives BEFORE its request finished is parked in ``_pending_sent_acks`` and released on a later step once the request enters ``_deferred_send_deadlines``. - * A deferred send whose ACK never arrives is reaped after - ``_defer_timeout`` and surfaced, so leaked blocks are freed. + * In READ mode, a deferred send whose ACK never arrives is reaped + after ``_defer_timeout``. WRITE-mode timeouts are owned by the worker, + which marks the transfer terminal before ACKing the block release. * A parked ACK that never matches a deferral before its own deadline is a stale duplicate (e.g. a real ACK landing after the send was already reaped) and is dropped. @@ -1575,7 +1585,7 @@ def update_connector_output(self, connector_output: KVConnectorOutput) -> None: expired = [ req_id for req_id, (deadline, _) in self._deferred_send_deadlines.items() - if now >= deadline + if self.mode == MoRIIOMode.READ and now >= deadline ] if expired: safe.update(expired) @@ -1685,6 +1695,7 @@ def __init__( # Completions that arrived before transfer_id_to_request_id was populated. # Retried each step until the mapping is established. self._unmatched_write_completions: set[str] = set() + self._unmatched_write_failures: OrderedDict[TransferId, None] = OrderedDict() # Producer-side READ-mode ACK fan-in. When decode TP is larger than # prefill TP, multiple decode ranks can read from one prefill rank and # notify the same transfer_id. Blocks are reusable only after all ACKs. @@ -1877,6 +1888,7 @@ def schedule_write_blocks( remote_ip: str, multi_pod_hosts: list[str], remote_dp_size_local: int, + remote_dp_size: int, ) -> None: """Schedule a block write operation. @@ -1892,6 +1904,7 @@ def schedule_write_blocks( remote_ip: IP address of remote node multi_pod_hosts: List of pod IPs for multi-pod Wide-EP remote_dp_size_local: Per-pod DP size for multi-pod + remote_dp_size: Global DP size of the decode instance """ # synchronization to prevent dirty reads between @@ -1915,6 +1928,7 @@ def schedule_write_blocks( remote_ip=remote_ip, multi_pod_hosts=multi_pod_hosts, remote_dp_size_local=remote_dp_size_local, + remote_dp_size=remote_dp_size, ) self._writer.schedule_write(task) @@ -2686,6 +2700,28 @@ def get_finished(self) -> tuple[set[str], set[str]]: return done_sending, done_recving + def get_transfer_results(self) -> KVConnectorTransferResults: + done_sending, done_recving = self.get_finished() + failed_recving: set[ReqId] = set() + + if self.mode == MoRIIOMode.WRITE and not self.is_producer: + pending = self._unmatched_write_failures + mapping = self.transfer_id_to_request_id + for transfer_id in self.moriio_wrapper.pop_failed_write_req_ids(): + pending[transfer_id] = None + for transfer_id in [t for t in pending if t in mapping]: + del pending[transfer_id] + failed_recving.add(mapping[transfer_id]) + while len(pending) > _MAX_UNMATCHED_WRITE_FAILURES: + pending.popitem(last=False) + done_recving |= failed_recving + + return KVConnectorTransferResults( + finished_sending=done_sending, + finished_recving=done_recving, + failed_recving=failed_recving, + ) + def wait_for_layer_load(self, layer_name: str) -> None: """Block until all in-flight READs of this layer have landed. @@ -3258,6 +3294,7 @@ def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_lay remote_ip=meta.remote_host, multi_pod_hosts=hosts, remote_dp_size_local=dp_local, + remote_dp_size=int(meta.remote_dp_size), ) def merge_contiguous_blocks( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py index 4ec3f72fb434..e33a1524e13f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py @@ -33,8 +33,10 @@ TransferError, TransferId, WriteTask, + fold_local_rank, get_port_offset, get_role, + pod_index, zmq_ctx, ) @@ -226,15 +228,6 @@ def _process_deferred_tasks(self) -> None: for task in self._deferred_tasks: if self._is_transfer_terminal(task.transfer_id): continue - if now - task.enqueue_time > defer_timeout: - logger.error( - "Deferred write task for request %s expired after %.1fs " - "(remote blocks never arrived), marking done", - task.request_id, - now - task.enqueue_time, - ) - self._mark_request_done(task.transfer_id) - continue if self._is_remote_ready(task): try: self._execute_write_task(task) @@ -244,11 +237,42 @@ def _process_deferred_tasks(self) -> None: task.request_id, ) self._mark_request_done(task.transfer_id) + elif ( + now - task.enqueue_time > defer_timeout + and self.worker.moriio_wrapper.fail_unallocated_write(task.transfer_id) + ): + self._notify_write_failed(task, now - task.enqueue_time) else: still_deferred.append(task) self._deferred_tasks = still_deferred + def _notify_write_failed(self, task: WriteTask, age: float) -> None: + self._clear_transfer_state(task.transfer_id) + logger.error( + "Deferred write task for request %s timed out after %.1fs waiting " + "for remote blocks", + task.request_id, + age, + ) + # No remote allocation arrived, so the decode DP rank that owns the + # request is unknown. Tell every rank; the others drop it. + for decode_dp_rank in range(max(task.remote_dp_size, 1)): + remote_ip, remote_port = self._resolve_notify_endpoint(task, decode_dp_rank) + try: + self.worker.moriio_wrapper.send_notify( + task.transfer_id, + remote_ip, + remote_port, + message_type="write_failed", + ) + except Exception: + logger.exception( + "Failed to notify decode DP rank %d that transfer %s timed out", + decode_dp_rank, + task.transfer_id, + ) + def _clear_transfer_state(self, transfer_id: TransferId) -> None: with self._write_state_lock: self._scheduled_writes.pop(transfer_id, None) @@ -314,27 +338,12 @@ def _execute_write_task(self, task: WriteTask) -> None: request_info = self._get_remote_alloc_info(task.transfer_id) with self._write_state_lock: request_info.completion_request_id = task.request_id - # Wide-EP multi-pod: task.remote_ip addresses only the first pod, - # so resolve the per-rank host from multi_pod_hosts (falls back to - # task.remote_ip for single-pod or on any indexing miss). - _remote_ip = task.remote_ip - _hosts = task.multi_pod_hosts - _dp_local = task.remote_dp_size_local - if _hosts and _dp_local > 0: - _pod_idx = int(request_info.decode_dp_rank) // _dp_local - if 0 <= _pod_idx < len(_hosts): - _remote_ip = _hosts[_pod_idx] - request_info.completion_remote_ip = _remote_ip - # Resolve the final notify port under the lock so finalize only - # reads it. The offset uses the per-pod local rank (% dp_local), - # since each pod binds notify sockets only for its local ranks; - # single-pod is bit-identical (modulus is a no-op). - _decode_dp_rank_for_port = int(request_info.decode_dp_rank) - if _dp_local > 0: - _decode_dp_rank_for_port = _decode_dp_rank_for_port % _dp_local - request_info.completion_notify_port = task.remote_notify_port + ( - get_port_offset(_decode_dp_rank_for_port, self.worker.tp_rank) - ) + # Resolve the final notify endpoint under the lock so finalize only + # reads it. + ( + request_info.completion_remote_ip, + request_info.completion_notify_port, + ) = self._resolve_notify_endpoint(task, request_info.decode_dp_rank) if task.transfer_id in self._sealed_writes: request_info.writes_expected = self._sealed_writes[task.transfer_id] @@ -506,6 +515,26 @@ def _finalize_if_complete( remote_port, ) + def _resolve_notify_endpoint( + self, task: WriteTask, decode_dp_rank: int + ) -> tuple[str, int]: + # Wide-EP multi-pod: task.remote_ip addresses only the first pod, so + # resolve the per-rank host from multi_pod_hosts (falls back to + # task.remote_ip for single-pod or on any indexing miss). The port + # offset uses the per-pod local rank, since each pod binds notify + # sockets only for its local ranks. + decode_dp_rank = int(decode_dp_rank) + dp_local = task.remote_dp_size_local + remote_ip = task.remote_ip + if task.multi_pod_hosts and dp_local > 0: + remote_pod = pod_index(decode_dp_rank, dp_local) + if 0 <= remote_pod < len(task.multi_pod_hosts): + remote_ip = task.multi_pod_hosts[remote_pod] + remote_port = task.remote_notify_port + get_port_offset( + fold_local_rank(decode_dp_rank, dp_local), self.worker.tp_rank + ) + return remote_ip, remote_port + class MoRIIOWrapper: """Wrapper for MoRIIO engine operations. @@ -539,6 +568,7 @@ def __init__( self.done_req_ids: list[MoRIIOTransferAck] = [] self.done_remote_allocate_req_dict: dict[TransferId, RemoteAllocInfo] = {} self.done_write_cache_req_ids: list[str] = [] + self.failed_write_cache_req_ids: list[str] = [] self._terminal_transfer_ids: OrderedDict[TransferId, None] = OrderedDict() self._transfer_timeout = transfer_timeout self.notify_thread: threading.Thread | None = None @@ -803,6 +833,8 @@ def _handle_structured_message(self, data: dict): self._handle_remote_blocks_message(data) elif message_type == "write_done": self._handle_write_done_message(data) + elif message_type == "write_failed": + self._handle_write_failed_message(data) elif message_type == "release": self._handle_release_message(data) else: @@ -837,6 +869,12 @@ def _handle_write_done_message(self, data: dict): with self.lock: self.done_write_cache_req_ids.append(transfer_id) + def _handle_write_failed_message(self, data: dict): + assert get_role() != ROLE.PRODUCER, "Only decode can get WRITE failure messages" + transfer_id = data["transfer_id"] + with self.lock: + self.failed_write_cache_req_ids.append(transfer_id) + def _handle_release_message(self, data: dict): assert get_role() == ROLE.PRODUCER, ( "Only prefill can get transfer release messages" @@ -852,6 +890,23 @@ def _handle_release_message(self, data: dict): self.done_remote_allocate_req_dict.pop(transfer_id, None) self._mark_transfer_terminal_locked(transfer_id) + def fail_unallocated_write(self, transfer_id: TransferId) -> bool: + """Mark a WRITE transfer terminal and release its source blocks if + decode never sent its allocation. + + Returns False if the transfer is already terminal or the allocation + has arrived; the writer then drops or executes the task as usual. + """ + with self.lock: + if ( + self._is_transfer_terminal_locked(transfer_id) + or transfer_id in self.done_remote_allocate_req_dict + ): + return False + self._mark_transfer_terminal_locked(transfer_id) + self.done_req_ids.append(MoRIIOTransferAck(transfer_id)) + return True + def _handle_completion_message(self, msg: str): with self.lock: if get_role() == ROLE.PRODUCER: @@ -929,6 +984,12 @@ def pop_finished_write_req_ids(self): self.done_write_cache_req_ids = [] return done_write_cache + def pop_failed_write_req_ids(self): + with self.lock: + failed_write_cache = set(self.failed_write_cache_req_ids) + self.failed_write_cache_req_ids = [] + return failed_write_cache + def shutdown(self): logger.debug("Closing MoRIIOWrapper and cleaning up ZMQ sockets") for path, sock in self.paths.items(): From 1bc6e072b0ae402db5b7f054fc4d862eeb3067ea Mon Sep 17 00:00:00 2001 From: xinnywinne <36653773+xinnywinne@users.noreply.github.com> Date: Mon, 12 Oct 2026 15:39:10 -0700 Subject: [PATCH 340/624] [Rust Frontend] Add per-request timing metrics to chat and completions (#59314) Signed-off-by: wxin Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- rust/src/chat/src/event.rs | 4 +- rust/src/chat/src/output/default/unified.rs | 7 + rust/src/chat/src/output/harmony/mod.rs | 1 + rust/src/chat/src/output/harmony/tests.rs | 3 + rust/src/chat/src/output/mod.rs | 2 + rust/src/chat/src/output/structured.rs | 20 ++- rust/src/chat/src/stream.rs | 6 + rust/src/chat/tests/roundtrip.rs | 2 + rust/src/cmd/src/cli.rs | 6 + rust/src/cmd/src/cli/tests.rs | 40 ++++++ rust/src/cmd/src/cli/unsupported.rs | 9 -- rust/src/llm/src/lib.rs | 2 +- rust/src/llm/src/output.rs | 10 ++ rust/src/llm/src/request_metrics.rs | 30 ++++- rust/src/llm/tests/generate.rs | 9 +- rust/src/server/src/config.rs | 7 + rust/src/server/src/grpc/convert.rs | 1 + rust/src/server/src/grpc/inference.rs | 1 + .../server/src/routes/inference/generate.rs | 3 + .../src/routes/openai/chat_completions.rs | 36 ++++++ .../routes/openai/chat_completions/types.rs | 4 + .../server/src/routes/openai/completions.rs | 107 ++++++++++++++- .../src/routes/openai/completions/types.rs | 4 + .../server/src/routes/openai/utils/metrics.rs | 122 ++++++++++++++++++ .../src/server/src/routes/openai/utils/mod.rs | 1 + rust/src/text/src/output/decoded.rs | 43 +++++- rust/src/text/src/output/mod.rs | 5 + .../chat_completion/test_chat_completion.py | 57 ++++++++ 28 files changed, 517 insertions(+), 25 deletions(-) create mode 100644 rust/src/server/src/routes/openai/utils/metrics.rs diff --git a/rust/src/chat/src/event.rs b/rust/src/chat/src/event.rs index 5fe25a9969ef..4056476b489d 100644 --- a/rust/src/chat/src/event.rs +++ b/rust/src/chat/src/event.rs @@ -5,7 +5,7 @@ use std::ops::Deref; use std::sync::Arc; use serde::{Deserialize, Serialize}; -use vllm_llm::TokenUsage; +use vllm_llm::{RequestTimestamps, TokenUsage}; use vllm_text::{DecodedLogprobs, DecodedPromptLogprobs}; use crate::FinishReason; @@ -235,6 +235,8 @@ pub enum ChatEvent { Done { message: AssistantMessage, usage: ChatTokenUsage, + /// Engine lifecycle times used by frontend response metrics. + timestamps: RequestTimestamps, finish_reason: FinishReason, /// Connector-specific KV transfer parameters for disaggregated serving. /// Boxed, like `ec_transfer_params`, to keep the per-token variants of diff --git a/rust/src/chat/src/output/default/unified.rs b/rust/src/chat/src/output/default/unified.rs index 9f872d97945b..a6238c440e70 100644 --- a/rust/src/chat/src/output/default/unified.rs +++ b/rust/src/chat/src/output/default/unified.rs @@ -277,6 +277,7 @@ pub(crate) async fn unified_event_stream( engine: finished.usage, reasoning_tokens: state.reasoning_tokens, }, + timestamps: finished.timestamps, finish_reason: finished.finish_reason, kv_transfer_params: finished.kv_transfer_params, ec_transfer_params: finished.ec_transfer_params, @@ -419,6 +420,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), })), } } @@ -742,6 +744,7 @@ mod tests { }, AssistantEvent::Done { usage: done_usage(0), + timestamps: Default::default(), finish_reason: crate::FinishReason::Stop(None), kv_transfer_params: None, ec_transfer_params: None, @@ -782,6 +785,7 @@ mod tests { }, AssistantEvent::Done { usage: done_usage(5), + timestamps: Default::default(), finish_reason: crate::FinishReason::Stop(None), kv_transfer_params: None, ec_transfer_params: None, @@ -814,6 +818,7 @@ mod tests { }, AssistantEvent::Done { usage: done_usage(2), + timestamps: Default::default(), finish_reason: crate::FinishReason::Stop(None), kv_transfer_params: None, ec_transfer_params: None, @@ -866,6 +871,7 @@ mod tests { }, AssistantEvent::Done { usage: done_usage(5), + timestamps: Default::default(), finish_reason: crate::FinishReason::Stop(None), kv_transfer_params: None, ec_transfer_params: None, @@ -912,6 +918,7 @@ mod tests { }, AssistantEvent::Done { usage: done_usage(0), + timestamps: Default::default(), finish_reason: crate::FinishReason::Stop(None), kv_transfer_params: None, ec_transfer_params: None, diff --git a/rust/src/chat/src/output/harmony/mod.rs b/rust/src/chat/src/output/harmony/mod.rs index 01e0390622bc..8ec873c6f2f3 100644 --- a/rust/src/chat/src/output/harmony/mod.rs +++ b/rust/src/chat/src/output/harmony/mod.rs @@ -380,6 +380,7 @@ async fn harmony_assistant_event_stream( if let Some(finished) = finished { y.yield_ok(AssistantEvent::Done { usage: finished.usage.into(), + timestamps: finished.timestamps, finish_reason: finished.finish_reason, kv_transfer_params: finished.kv_transfer_params, ec_transfer_params: finished.ec_transfer_params, diff --git a/rust/src/chat/src/output/harmony/tests.rs b/rust/src/chat/src/output/harmony/tests.rs index 5b89e53bb9aa..a529888ccff5 100644 --- a/rust/src/chat/src/output/harmony/tests.rs +++ b/rust/src/chat/src/output/harmony/tests.rs @@ -59,6 +59,7 @@ fn finished() -> Finished { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), } } @@ -132,6 +133,7 @@ fn interrupted_final_message_is_preserved() { output_token_count: 0, cached_token_count: 0, }), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -183,6 +185,7 @@ fn interrupted_analysis_message_is_preserved() { output_token_count: 0, cached_token_count: 0, }), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, diff --git a/rust/src/chat/src/output/mod.rs b/rust/src/chat/src/output/mod.rs index 8c1dae37f4de..14ac725206e7 100644 --- a/rust/src/chat/src/output/mod.rs +++ b/rust/src/chat/src/output/mod.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use futures::Stream; use trait_set::trait_set; use uuid::Uuid; +use vllm_llm::RequestTimestamps; use vllm_parser::output_grammar::BuiltOutputGrammar; use vllm_text::output::{DecodedLogprobs, DecodedPromptLogprobs, DecodedTextEvent}; @@ -52,6 +53,7 @@ pub(crate) enum AssistantEvent { ToolCallArgumentsDelta { delta: String }, Done { usage: ChatTokenUsage, + timestamps: RequestTimestamps, finish_reason: FinishReason, /// Connector-specific KV transfer parameters for disaggregated serving. /// Boxed, like `ec_transfer_params`, to keep the per-token variants of diff --git a/rust/src/chat/src/output/structured.rs b/rust/src/chat/src/output/structured.rs index 80035db4cb1e..3ed1a4e85271 100644 --- a/rust/src/chat/src/output/structured.rs +++ b/rust/src/chat/src/output/structured.rs @@ -10,6 +10,7 @@ use asynk_strim_attr::{TryYielder, try_stream}; use futures::{StreamExt as _, pin_mut}; +use vllm_llm::RequestTimestamps; use vllm_text::DecodedLogprobs; use super::{AssistantEvent, AssistantEventStream}; @@ -149,6 +150,7 @@ impl StructuredEventState { fn finish( &mut self, usage: ChatTokenUsage, + timestamps: RequestTimestamps, finish_reason: FinishReason, kv_transfer_params: Option>, ec_transfer_params: Option>, @@ -159,6 +161,7 @@ impl StructuredEventState { events.push(ChatEvent::Done { message: self.message.clone(), usage, + timestamps, finish_reason, kv_transfer_params, ec_transfer_params, @@ -312,13 +315,18 @@ pub(crate) async fn structured_chat_event_stream( } AssistantEvent::Done { usage, + timestamps, finish_reason, kv_transfer_params, ec_transfer_params, } => { - for next in - state.finish(usage, finish_reason, kv_transfer_params, ec_transfer_params)? - { + for next in state.finish( + usage, + timestamps, + finish_reason, + kv_transfer_params, + ec_transfer_params, + )? { y.yield_ok(next).await; } } @@ -353,6 +361,7 @@ mod tests { output_token_count: 1, cached_token_count: 0, }), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -408,6 +417,7 @@ mod tests { output_token_count: 1, cached_token_count: 0, }), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -461,6 +471,7 @@ mod tests { output_token_count: 1, cached_token_count: 0, }), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -514,6 +525,7 @@ mod tests { output_token_count: 1, cached_token_count: 0, }), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -582,6 +594,7 @@ mod tests { output_token_count: 1, cached_token_count: 0, }), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -633,6 +646,7 @@ mod tests { }, reasoning_tokens: 3, }, + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, diff --git a/rust/src/chat/src/stream.rs b/rust/src/chat/src/stream.rs index 509628e8e2af..fd58249461c8 100644 --- a/rust/src/chat/src/stream.rs +++ b/rust/src/chat/src/stream.rs @@ -7,6 +7,7 @@ use std::task::{Context, Poll}; use futures::Stream; use trait_set::trait_set; +use vllm_llm::RequestTimestamps; use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs}; use crate::FinishReason; @@ -22,6 +23,7 @@ pub struct CollectedAssistantMessage { pub logprobs: Option, pub token_ids: Vec, pub usage: ChatTokenUsage, + pub timestamps: RequestTimestamps, pub finish_reason: FinishReason, /// Connector-specific KV transfer parameters for disaggregated serving. pub kv_transfer_params: Option, @@ -81,6 +83,7 @@ impl ChatEventStream { ChatEvent::Done { message: done, usage, + timestamps, finish_reason, kv_transfer_params, ec_transfer_params, @@ -94,6 +97,7 @@ impl ChatEventStream { }), token_ids, usage, + timestamps, finish_reason, kv_transfer_params: kv_transfer_params.map(|value| *value), ec_transfer_params: ec_transfer_params.map(|value| *value), @@ -200,6 +204,7 @@ mod tests { output_token_count: 1, cached_token_count: 0, }), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -241,6 +246,7 @@ mod tests { output_token_count: 1, cached_token_count: 0, }), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index 41f06d773432..d2891ed02797 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -1051,6 +1051,7 @@ fn decoded_completion_stream( kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), })), } }); @@ -1063,6 +1064,7 @@ fn decoded_completion_stream( kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), }); events.push(DecodedTextEvent::TextDelta { decoded: chunk.delta, diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index eeb458c86e75..66b48f13d385 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -348,6 +348,11 @@ pub struct SharedRuntimeArgs { #[serde(default)] pub enable_prompt_tokens_details: bool, + /// Include per-request timing metrics in API responses. + #[arg(long, default_missing_value = "true", num_args = 0..=1)] + #[serde(default)] + pub enable_per_request_metrics: bool, + /// If specified, API server will add X-Request-Id header to responses. #[arg( long, @@ -635,6 +640,7 @@ impl SharedRuntimeArgs { ApiServerOptions { enable_log_requests: self.enable_log_requests, enable_prompt_tokens_details: self.enable_prompt_tokens_details, + enable_per_request_metrics: self.enable_per_request_metrics, enable_request_id_headers: self.enable_request_id_headers, enable_scale_out: self.enable_scale_out, sse_keep_alive_interval: (self.sse_keep_alive_interval > 0) diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index cc545ce33e80..8f4318ac1693 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -199,6 +199,7 @@ fn serve_args_forward_python_flags_with_separator() { chat_template_content_format: Auto, enable_log_requests: false, enable_prompt_tokens_details: false, + enable_per_request_metrics: false, enable_request_id_headers: false, enable_scale_out: false, sse_keep_alive_interval: 0, @@ -970,6 +971,40 @@ fn serve_args_reject_unsupported_flag_arg() { "#]].assert_eq(&error.to_string()); } +#[test] +fn serve_args_enable_per_request_metrics_requires_engine_stats() { + let cli = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--enable-per-request-metrics", + ]) + .expect("per-request metrics flag is supported"); + let Command::Serve(args) = cli.command else { + panic!("expected serve args"); + }; + let config = args.to_frontend_config("tcp://127.0.0.1:0".into()); + assert!(config.api_server_options.enable_per_request_metrics); + assert!(config.validate().is_ok()); + + let cli = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--enable-per-request-metrics", + "--disable-log-stats", + ]) + .expect("flags parse independently"); + let Command::Serve(args) = cli.command else { + panic!("expected serve args"); + }; + let error = args.to_frontend_config("tcp://127.0.0.1:0".into()).validate().unwrap_err(); + assert!( + format!("{error:#}") + .contains("--enable-per-request-metrics requires engine statistics logging") + ); +} + #[test] fn serve_args_reject_custom_generation_config_source() { let error = Cli::try_parse_from([ @@ -1067,6 +1102,7 @@ fn frontend_args_accept_json() { chat_template_content_format: Auto, enable_log_requests: false, enable_prompt_tokens_details: false, + enable_per_request_metrics: false, enable_request_id_headers: false, enable_scale_out: false, sse_keep_alive_interval: 0, @@ -1738,6 +1774,7 @@ fn serve_args_accept_handshake_aliases() { chat_template_content_format: Auto, enable_log_requests: false, enable_prompt_tokens_details: false, + enable_per_request_metrics: false, enable_request_id_headers: false, enable_scale_out: false, sse_keep_alive_interval: 0, @@ -1895,6 +1932,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { api_server_options: ApiServerOptions { enable_log_requests: false, enable_prompt_tokens_details: false, + enable_per_request_metrics: false, enable_request_id_headers: false, enable_scale_out: false, sse_keep_alive_interval: None, @@ -1989,6 +2027,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { api_server_options: ApiServerOptions { enable_log_requests: false, enable_prompt_tokens_details: false, + enable_per_request_metrics: false, enable_request_id_headers: false, enable_scale_out: false, sse_keep_alive_interval: None, @@ -2127,6 +2166,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present api_server_options: ApiServerOptions { enable_log_requests: false, enable_prompt_tokens_details: false, + enable_per_request_metrics: false, enable_request_id_headers: false, enable_scale_out: false, sse_keep_alive_interval: None, diff --git a/rust/src/cmd/src/cli/unsupported.rs b/rust/src/cmd/src/cli/unsupported.rs index 0a14d1b841c4..02e1377ff01c 100644 --- a/rust/src/cmd/src/cli/unsupported.rs +++ b/rust/src/cmd/src/cli/unsupported.rs @@ -476,15 +476,6 @@ pub struct ServerUnsupportedArgs { #[arg(long)] pub max_log_len: Option, - /// If set to True, include per-request timing metrics in API responses. - #[arg( - long, - visible_alias = "no-enable-per-request-metrics", - default_missing_value = "true", - num_args = 0..=1 - )] - pub enable_per_request_metrics: Option, - /// If set to True, enable tracking server_load_metrics in the app state. #[arg( long, diff --git a/rust/src/llm/src/lib.rs b/rust/src/llm/src/lib.rs index c94c1d80f743..1e15f11c12e3 100644 --- a/rust/src/llm/src/lib.rs +++ b/rust/src/llm/src/lib.rs @@ -19,7 +19,7 @@ pub use output::{ GenerateOutputStreamExt, GeneratePromptInfo, TokenUsage, WeightVersionSpan, }; pub use request::GenerateRequest; -pub use request_metrics::current_unix_timestamp_secs; +pub use request_metrics::{RequestTimestamps, current_unix_timestamp_secs}; pub use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs, TokenLogprob}; use crate::inflight::InflightRequests; diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index 81dbb3fda53c..3f0b0300626d 100644 --- a/rust/src/llm/src/output.rs +++ b/rust/src/llm/src/output.rs @@ -20,6 +20,7 @@ use vllm_engine_core_client::{ AbortCause, EngineCoreOutputStream, EngineCoreStreamDelivery, EngineCoreStreamOutput, }; +use crate::RequestTimestamps; use crate::error::Result; use crate::inflight::RequestGuard; use crate::request_metrics::{RequestMetricsTracker, current_unix_timestamp_secs}; @@ -183,6 +184,12 @@ pub struct GenerateOutput { pub spec_decode_metrics: Option>, /// Weight-version spans over all output tokens, set on terminal outputs only. pub weight_versions: Option>, + /// Snapshot of engine lifecycle timestamps after this output. + /// + /// Every output carries a snapshot because text decoding may finish the + /// request locally after matching a stop string, before the engine emits a + /// terminal output. + pub timestamps: RequestTimestamps, } impl GenerateOutput { @@ -237,6 +244,7 @@ impl GenerateOutput { self.ec_transfer_params = next.ec_transfer_params; self.spec_decode_metrics = next.spec_decode_metrics; self.weight_versions = next.weight_versions; + self.timestamps = next.timestamps; } } @@ -266,6 +274,7 @@ impl GenerateOutput { sampling_mask: None, spec_decode_metrics: None, weight_versions: None, + timestamps: RequestTimestamps::default(), } } } @@ -402,6 +411,7 @@ impl GenerateOutputStream { sampling_mask, spec_decode_metrics: raw.spec_decode_metrics, weight_versions, + timestamps: self.request_metrics.timestamps(), }; Ok(output) diff --git a/rust/src/llm/src/request_metrics.rs b/rust/src/llm/src/request_metrics.rs index 71decdd281ff..eebf0596b0d2 100644 --- a/rust/src/llm/src/request_metrics.rs +++ b/rust/src/llm/src/request_metrics.rs @@ -18,6 +18,15 @@ const PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE: &str = "local_compute"; const PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT: &str = "local_cache_hit"; const PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER: &str = "external_kv_transfer"; +/// Engine lifecycle timestamps used by frontend per-request response metrics. +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub struct RequestTimestamps { + pub queued_ts: f64, + pub scheduled_ts: f64, + pub first_token_ts: f64, + pub last_token_ts: f64, +} + /// Request-scoped metrics state tracked across streamed engine-core updates. /// /// This is the Rust-side counterpart of the Python frontend's request-lifecycle @@ -227,6 +236,16 @@ impl RequestMetricsTracker { } } + /// Engine lifecycle timestamps recorded so far. + pub(crate) fn timestamps(&self) -> RequestTimestamps { + RequestTimestamps { + queued_ts: self.queued_ts, + scheduled_ts: self.scheduled_ts, + first_token_ts: self.first_token_ts, + last_token_ts: self.last_token_ts, + } + } + /// Record prompt token counters through cached metric handles. fn record_prompt_tokens(&self, prefill_stats: &PrefillStats) { let computed = prefill_stats.num_computed_tokens as u64; @@ -388,7 +407,7 @@ mod tests { use vllm_engine_core_client::protocol::output::{EngineCoreEvent, EngineCoreEventType}; use vllm_engine_core_client::protocol::stats::PrefillStats; - use super::{RequestMetricsTracker, diff_or_zero}; + use super::{RequestMetricsTracker, RequestTimestamps, diff_or_zero}; #[test] fn tracker_updates_timing_state_across_prefill_decode_and_finish() { @@ -445,6 +464,15 @@ mod tests { assert_eq!(tracker.scheduled_ts, 9.0); assert_eq!(tracker.first_token_ts, 10.0); assert_eq!(tracker.last_token_ts, 11.5); + assert_eq!( + tracker.timestamps(), + RequestTimestamps { + queued_ts: 8.0, + scheduled_ts: 9.0, + first_token_ts: 10.0, + last_token_ts: 11.5, + } + ); assert!((tracker.first_token_latency - 0.2).abs() < 1e-9); assert_eq!( diff_or_zero(tracker.last_token_ts, tracker.first_token_ts), diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 0a1eacd73ff6..cc6fd9d169f6 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -352,7 +352,9 @@ async fn generate_merges_outputs_by_stream_interval() { (&[6], "v7", None), (&[7], "v7", Some(EngineCoreFinishReason::Length)), ]; - for (token_ids, weight_version, finish_reason) in steps { + for (step, (token_ids, weight_version, finish_reason)) in + steps.into_iter().enumerate() + { send_outputs( push, RequestBatchOutputs { @@ -364,6 +366,7 @@ async fn generate_merges_outputs_by_stream_interval() { None, )], weight_version: Some(weight_version.into()), + timestamp: step as f64 + 1.0, ..Default::default() } .into(), @@ -430,6 +433,9 @@ async fn generate_merges_outputs_by_stream_interval() { Some(vec![span("v7", 0, 2), span("v8", 2, 5), span("v7", 5, 7)]) ] ); + let last_token_timestamps: Vec<_> = + outputs.iter().map(|output| output.timestamps.last_token_ts).collect(); + assert_eq!(last_token_timestamps, [1.0, 4.0, 6.0]); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -549,6 +555,7 @@ async fn collect_output_rejects_partial_sampling_mask() { }), spec_decode_metrics: None, weight_versions: None, + timestamps: Default::default(), }; let error = futures::stream::iter([Ok(output)]).collect_output().await.unwrap_err(); diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index ca91ba40cd50..d70933fb446d 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -56,6 +56,8 @@ pub struct ApiServerOptions { pub enable_log_requests: bool, /// When `true`, include prompt token cache details in response usage. pub enable_prompt_tokens_details: bool, + /// Include request timing metrics in OpenAI chat and completion responses. + pub enable_per_request_metrics: bool, /// When `true`, set `X-Request-Id` on every HTTP response. pub enable_request_id_headers: bool, /// When `true`, register the scale-out `/inference/v1/generate` route. @@ -296,6 +298,11 @@ impl Config { /// Validate frontend configuration that can be checked before engine /// startup. pub fn validate(&self) -> Result<()> { + if self.api_server_options.enable_per_request_metrics && self.disable_log_stats { + bail!( + "--enable-per-request-metrics requires engine statistics logging; remove --disable-log-stats" + ); + } vllm_chat::validate_parser_overrides(&self.tool_call_parser, &self.reasoning_parser)?; self.cors.validate()?; if let Some(tls) = &self.tls { diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 8e31bb84d2d3..1eec1d0fe320 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -947,6 +947,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), } } diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs index 1484a9f9e375..fae2794ba331 100644 --- a/rust/src/server/src/grpc/inference.rs +++ b/rust/src/server/src/grpc/inference.rs @@ -259,6 +259,7 @@ impl pb::inference_server::Inference for InferenceServiceImpl { kv_transfer_params: collected.kv_transfer_params.map(Box::new), ec_transfer_params: collected.ec_transfer_params.map(Box::new), sampling_mask: collected.sampling_mask, + timestamps: collected.timestamps, }; let outputs = convert::to_sequence_output( diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index 930c698c30e4..42716de8de37 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -637,6 +637,7 @@ mod tests { sampling_mask: None, spec_decode_metrics: None, weight_versions: None, + timestamps: Default::default(), }), Ok(GenerateOutput { request_id: String::new(), @@ -654,6 +655,7 @@ mod tests { sampling_mask: None, spec_decode_metrics: None, weight_versions: None, + timestamps: Default::default(), }), ]); @@ -741,6 +743,7 @@ mod tests { sampling_mask: None, spec_decode_metrics: None, weight_versions: None, + timestamps: Default::default(), } } diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index a87c52fb5578..5d8dcda53641 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -38,6 +38,7 @@ use crate::routes::openai::chat_completions::types::{ use crate::routes::openai::utils::logprobs::{ decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps, }; +use crate::routes::openai::utils::metrics::PerRequestMetrics; use crate::routes::openai::utils::types::{ ChatLogProbs, FunctionCallDelta, FunctionCallResponse, StreamResponseEnvelope, ToolCall, ToolCallDelta, Usage, @@ -128,6 +129,7 @@ async fn collect_chat_completion( ApiServerOptions { enable_log_requests, enable_prompt_tokens_details, + enable_per_request_metrics, .. }: ApiServerOptions, ResponseOptions { @@ -158,6 +160,7 @@ async fn collect_chat_completion( logprobs, token_ids, usage, + timestamps, finish_reason, kv_transfer_params, ec_transfer_params, @@ -202,6 +205,8 @@ async fn collect_chat_completion( } else { None }; + let metrics = enable_per_request_metrics + .then(|| PerRequestMetrics::from_timestamps(timestamps, usage.output_token_count)); let usage = Usage::from_token_usage(usage, enable_prompt_tokens_details); if enable_log_requests { @@ -236,6 +241,7 @@ async fn collect_chat_completion( token_ids: (return_token_ids && include_output_metadata).then_some(token_ids), }], usage: Some(usage), + metrics, system_fingerprint: None, prompt_logprobs, prompt_token_ids: return_token_ids.then(|| prompt_token_ids.to_vec()), @@ -254,6 +260,7 @@ async fn chat_completion_chunk_stream( ApiServerOptions { enable_log_requests, enable_prompt_tokens_details, + enable_per_request_metrics, .. }: ApiServerOptions, ResponseOptions { @@ -428,6 +435,7 @@ async fn chat_completion_chunk_stream( } Ok(ChatEvent::Done { usage: final_usage, + timestamps, finish_reason, .. }) => { @@ -472,9 +480,16 @@ async fn chat_completion_chunk_stream( } if include_usage { + let metrics = enable_per_request_metrics.then(|| { + PerRequestMetrics::from_timestamps( + timestamps, + final_usage.output_token_count, + ) + }); y.yield_ok(usage_chunk( &envelope, Usage::from_token_usage(final_usage, enable_prompt_tokens_details), + metrics, )) .await; } @@ -496,9 +511,11 @@ async fn chat_completion_chunk_stream( fn usage_chunk( envelope: &Arc, usage: Usage, + metrics: Option, ) -> ChatCompletionStreamResponse { let mut chunk = ChatCompletionStreamResponse::new(envelope); chunk.usage = Some(usage); + chunk.metrics = metrics.map(Into::into); chunk } @@ -962,6 +979,12 @@ mod tests { Ok(ChatEvent::Done { message: Default::default(), usage: done_usage(1, 1, 1), + timestamps: vllm_llm::RequestTimestamps { + queued_ts: 10.0, + scheduled_ts: 10.2, + first_token_ts: 10.5, + last_token_ts: 11.0, + }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -975,6 +998,7 @@ mod tests { 1, ApiServerOptions { enable_prompt_tokens_details: true, + enable_per_request_metrics: true, ..Default::default() }, ResponseOptions { @@ -1006,6 +1030,10 @@ mod tests { .map(|details| details.cached_tokens), Some(1) ); + assert!(chunks[..3].iter().all(|chunk| chunk.metrics.is_none())); + let metrics = chunks[3].metrics.as_ref().expect("per-request metrics"); + assert!((metrics.time_to_first_token_ms.unwrap() - 300.0).abs() < 1e-9); + assert!((metrics.tokens_per_second.unwrap() - 1.25).abs() < 1e-9); } #[tokio::test] @@ -1041,6 +1069,7 @@ mod tests { Ok(ChatEvent::Done { message: Default::default(), usage: done_usage(1, 1, 0), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -1095,6 +1124,7 @@ mod tests { Ok(ChatEvent::Done { message: Default::default(), usage: done_usage(1, 2, 0), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -1175,6 +1205,7 @@ mod tests { Ok(ChatEvent::Done { message: Default::default(), usage: done_usage(1, 2, 0), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -1307,6 +1338,7 @@ mod tests { Ok(ChatEvent::Done { message: Default::default(), usage: done_usage(1, 4, 0), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -1385,6 +1417,7 @@ mod tests { Ok(ChatEvent::Done { message: Default::default(), usage: done_usage(1, 1, 0), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -1468,6 +1501,7 @@ mod tests { }, reasoning_tokens: 2, }, + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -1539,6 +1573,7 @@ mod tests { }, reasoning_tokens: 2, }, + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -1593,6 +1628,7 @@ mod tests { Ok(ChatEvent::Done { message: Default::default(), usage: done_usage(1, 1, 0), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, diff --git a/rust/src/server/src/routes/openai/chat_completions/types.rs b/rust/src/server/src/routes/openai/chat_completions/types.rs index 21776dd075e0..d73480b45d5e 100644 --- a/rust/src/server/src/routes/openai/chat_completions/types.rs +++ b/rust/src/server/src/routes/openai/chat_completions/types.rs @@ -13,6 +13,7 @@ use validator::Validate; use vllm_engine_core_client::protocol::sampling::RepetitionDetectionParams; use vllm_text::TruncationSide; +use crate::routes::openai::utils::metrics::{PerRequestMetrics, StreamPerRequestMetrics}; use crate::routes::openai::utils::structured_outputs::ResponseFormat; use crate::routes::openai::utils::types::{ ChatLogProbs, ChatMessage, Normalizable, PromptLogprobs, ReasoningEffort, StreamOptions, @@ -365,6 +366,7 @@ pub(super) struct ChatCompletionResponse { pub model: String, pub choices: Vec, pub usage: Option, + pub metrics: Option, pub system_fingerprint: Option, pub prompt_logprobs: Option, pub prompt_token_ids: Option>, @@ -412,6 +414,7 @@ pub(super) struct ChatCompletionStreamResponse { pub envelope: Arc, pub choices: Vec, pub usage: Option, + pub metrics: Option, pub prompt_token_ids: Option>, } @@ -422,6 +425,7 @@ impl ChatCompletionStreamResponse { envelope: Arc::clone(envelope), choices: Vec::new(), usage: None, + metrics: None, prompt_token_ids: None, } } diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs index 646a67d1fb8f..05dd579ecec7 100644 --- a/rust/src/server/src/routes/openai/completions.rs +++ b/rust/src/server/src/routes/openai/completions.rs @@ -33,6 +33,7 @@ use super::utils::logprobs::{ collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_openai, prompt_logprobs_to_maps, text_len, }; +use super::utils::metrics::PerRequestMetrics; use super::utils::types::{StreamResponseEnvelope, Usage}; use crate::config::ApiServerOptions; use crate::error::{ApiError, bail_server_error, server_error, text_submit_error}; @@ -140,6 +141,7 @@ async fn collect_completion( ApiServerOptions { enable_log_requests, enable_prompt_tokens_details, + enable_per_request_metrics, .. }: ApiServerOptions, ResponseOptions { @@ -197,6 +199,9 @@ async fn collect_completion( Some(prompt) => format!("{prompt}{}", collected.text), }; let finish_reason = completion_finish_reason_to_openai(&finish_reason)?.to_string(); + let metrics = enable_per_request_metrics.then(|| { + PerRequestMetrics::from_timestamps(collected.timestamps, collected.usage.output_token_count) + }); let usage = Usage::from_token_usage(collected.usage, enable_prompt_tokens_details); if enable_log_requests { @@ -225,6 +230,7 @@ async fn collect_completion( prompt_token_ids: return_token_ids.then(|| collected.prompt_token_ids.to_vec()), }], usage: Some(usage), + metrics, system_fingerprint: None, kv_transfer_params: collected.kv_transfer_params, ec_transfer_params: collected.ec_transfer_params, @@ -241,6 +247,7 @@ async fn completion_chunk_stream( ApiServerOptions { enable_log_requests, enable_prompt_tokens_details, + enable_per_request_metrics, .. }: ApiServerOptions, ResponseOptions { @@ -356,12 +363,19 @@ async fn completion_chunk_stream( yield_chunk!(final_chunk); if include_usage { + let metrics = enable_per_request_metrics.then(|| { + PerRequestMetrics::from_timestamps( + finished.timestamps, + finished.usage.output_token_count, + ) + }); y.yield_ok(CompletionSseChunk::Usage(usage_chunk( &envelope, Usage::from_token_usage( finished.usage, enable_prompt_tokens_details, ), + metrics, ))) .await; } @@ -411,9 +425,16 @@ async fn completion_chunk_stream( yield_chunk!(final_chunk); if include_usage { + let metrics = enable_per_request_metrics.then(|| { + PerRequestMetrics::from_timestamps( + finished.timestamps, + finished.usage.output_token_count, + ) + }); y.yield_ok(CompletionSseChunk::Usage(usage_chunk( &envelope, Usage::from_token_usage(finished.usage, enable_prompt_tokens_details), + metrics, ))) .await; } @@ -509,9 +530,14 @@ fn prompt_only_logprobs_to_openai( )) } -fn usage_chunk(envelope: &Arc, usage: Usage) -> CompletionStreamResponse { +fn usage_chunk( + envelope: &Arc, + usage: Usage, + metrics: Option, +) -> CompletionStreamResponse { let mut chunk = CompletionStreamResponse::new(envelope); chunk.usage = Some(usage); + chunk.metrics = metrics.map(Into::into); chunk } @@ -578,7 +604,7 @@ mod tests { use super::{ ApiServerOptions, CompletionSseChunk, CompletionStreamResponse, ResponseOptions, - StreamResponseEnvelope, completion_chunk_stream, final_chunk, + StreamResponseEnvelope, collect_completion, completion_chunk_stream, final_chunk, }; fn decoded_delta( @@ -639,6 +665,66 @@ mod tests { assert!(final_chunk(&stream_envelope(), FinishReason::Error).is_err()); } + #[tokio::test] + async fn collected_completion_includes_metrics_only_when_enabled() { + let make_stream = || { + stream::iter(vec![ + Ok(DecodedTextEvent::Start { + prompt_token_ids: vec![1].into(), + prompt_logprobs: None, + }), + Ok(decoded_delta( + "hi", + vec![2], + None, + Some(Finished { + usage: vllm_llm::TokenUsage { + prompt_token_count: 1, + output_token_count: 1, + cached_token_count: 0, + }, + finish_reason: FinishReason::stop_eos(), + kv_transfer_params: None, + ec_transfer_params: None, + sampling_mask: None, + timestamps: vllm_llm::RequestTimestamps { + queued_ts: 10.0, + scheduled_ts: 10.2, + first_token_ts: 10.5, + last_token_ts: 10.5, + }, + }), + )), + ]) + }; + + for enabled in [false, true] { + let response = collect_completion( + make_stream(), + "cmpl-1".into(), + "model".into(), + 1, + ApiServerOptions { + enable_per_request_metrics: enabled, + ..Default::default() + }, + ResponseOptions::default(), + ) + .await + .expect("completion response"); + let json = serde_json::to_value(response).expect("serializable response"); + if enabled { + assert!( + (json["metrics"]["time_to_first_token_ms"].as_f64().unwrap() - 300.0).abs() + < 1e-9 + ); + assert_eq!(json["metrics"]["mean_itl_ms"], serde_json::Value::Null); + } else { + assert_eq!(json["metrics"], serde_json::Value::Null); + } + } + } + #[tokio::test] async fn completion_chunk_stream_maps_streaming_logprobs() { let stream = stream::iter(vec![ @@ -702,6 +788,12 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: vllm_llm::RequestTimestamps { + queued_ts: 10.0, + scheduled_ts: 10.2, + first_token_ts: 10.5, + last_token_ts: 11.0, + }, }), )), ]); @@ -713,6 +805,7 @@ mod tests { 1, ApiServerOptions { enable_prompt_tokens_details: true, + enable_per_request_metrics: true, ..Default::default() }, ResponseOptions { @@ -779,6 +872,10 @@ mod tests { match &chunks[3] { CompletionSseChunk::Usage(chunk) => { + assert!(chunks[..3].iter().all(|chunk| chunk_response(chunk).metrics.is_none())); + let metrics = chunk.metrics.as_ref().expect("per-request metrics"); + assert!((metrics.queue_time_ms.unwrap() - 200.0).abs() < 1e-9); + assert!((metrics.tokens_per_second.unwrap() - 2.5).abs() < 1e-9); assert_eq!( chunk .usage @@ -815,6 +912,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), }), )), ]); @@ -879,6 +977,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), }), )), ]); @@ -932,6 +1031,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), }), )), ]); @@ -988,6 +1088,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), }), )), ]); @@ -1061,6 +1162,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), }), )), ]); @@ -1136,6 +1238,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), }), )), ]); diff --git a/rust/src/server/src/routes/openai/completions/types.rs b/rust/src/server/src/routes/openai/completions/types.rs index 9f7027511cbf..fb6aa22c601d 100644 --- a/rust/src/server/src/routes/openai/completions/types.rs +++ b/rust/src/server/src/routes/openai/completions/types.rs @@ -11,6 +11,7 @@ use validator::Validate; use vllm_engine_core_client::protocol::sampling::RepetitionDetectionParams; use vllm_text::{Prompt, TruncationSide}; +use crate::routes::openai::utils::metrics::{PerRequestMetrics, StreamPerRequestMetrics}; use crate::routes::openai::utils::types::{ LogProbs, Normalizable, PromptLogprobs, StreamOptions, StreamResponseEnvelope, StringOrArray, Usage, default_true, deserialize_request_top_k, validate_stop, @@ -239,6 +240,7 @@ pub(super) struct CompletionResponse { pub model: String, pub choices: Vec, pub usage: Option, + pub metrics: Option, pub system_fingerprint: Option, pub kv_transfer_params: Option, pub ec_transfer_params: Option, @@ -265,6 +267,7 @@ pub(super) struct CompletionStreamResponse { pub envelope: Arc, pub choices: Vec, pub usage: Option, + pub metrics: Option, } impl CompletionStreamResponse { @@ -274,6 +277,7 @@ impl CompletionStreamResponse { envelope: Arc::clone(envelope), choices: Vec::new(), usage: None, + metrics: None, } } } diff --git a/rust/src/server/src/routes/openai/utils/metrics.rs b/rust/src/server/src/routes/openai/utils/metrics.rs new file mode 100644 index 000000000000..0c1155f733a4 --- /dev/null +++ b/rust/src/server/src/routes/openai/utils/metrics.rs @@ -0,0 +1,122 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use serde::Serialize; +use vllm_llm::RequestTimestamps; + +/// Per-request response timing fields matching Python's `PerRequestMetrics`. +#[derive(Debug, Clone, PartialEq, Serialize)] +pub struct PerRequestMetrics { + pub time_to_first_token_ms: Option, + pub generation_time_ms: Option, + pub queue_time_ms: Option, + pub mean_itl_ms: Option, + pub tokens_per_second: Option, +} + +impl PerRequestMetrics { + pub fn from_timestamps(times: RequestTimestamps, num_generation_tokens: usize) -> Self { + let duration = |end: f64, start: f64| (end > 0.0 && start > 0.0).then_some(end - start); + let decode = duration(times.last_token_ts, times.first_token_ts); + let inference = duration(times.last_token_ts, times.scheduled_ts); + + Self { + time_to_first_token_ms: duration(times.first_token_ts, times.scheduled_ts) + .map(|seconds| seconds * 1000.0), + generation_time_ms: decode.map(|seconds| seconds * 1000.0), + queue_time_ms: duration(times.scheduled_ts, times.queued_ts) + .map(|seconds| seconds * 1000.0), + mean_itl_ms: decode + .filter(|_| num_generation_tokens > 1) + .map(|seconds| seconds * 1000.0 / (num_generation_tokens - 1) as f64), + tokens_per_second: inference + .filter(|seconds| *seconds > 0.0) + .map(|seconds| num_generation_tokens as f64 / seconds), + } + } +} + +/// Streaming form of [`PerRequestMetrics`]. +/// +/// Python serializes the final usage chunk with recursive `exclude_none=True`, +/// so unset fields are omitted instead of serialized as `null`. +#[serde_with::skip_serializing_none] +#[derive(Debug, Clone, PartialEq, Serialize)] +pub struct StreamPerRequestMetrics { + pub time_to_first_token_ms: Option, + pub generation_time_ms: Option, + pub queue_time_ms: Option, + pub mean_itl_ms: Option, + pub tokens_per_second: Option, +} + +impl From for StreamPerRequestMetrics { + fn from(metrics: PerRequestMetrics) -> Self { + let PerRequestMetrics { + time_to_first_token_ms, + generation_time_ms, + queue_time_ms, + mean_itl_ms, + tokens_per_second, + } = metrics; + Self { + time_to_first_token_ms, + generation_time_ms, + queue_time_ms, + mean_itl_ms, + tokens_per_second, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn matches_python_timing_formulas_and_missing_fields() { + let times = RequestTimestamps { + queued_ts: 10.0, + scheduled_ts: 10.2, + first_token_ts: 10.5, + last_token_ts: 11.0, + }; + let metrics = PerRequestMetrics::from_timestamps(times, 6); + assert!((metrics.time_to_first_token_ms.unwrap() - 300.0).abs() < 1e-9); + assert_eq!(metrics.generation_time_ms, Some(500.0)); + assert!((metrics.queue_time_ms.unwrap() - 200.0).abs() < 1e-9); + assert_eq!(metrics.mean_itl_ms, Some(100.0)); + assert!((metrics.tokens_per_second.unwrap() - 7.5).abs() < 1e-9); + + let single_token = PerRequestMetrics::from_timestamps(times, 1); + assert_eq!(single_token.mean_itl_ms, None); + + let missing = PerRequestMetrics::from_timestamps(RequestTimestamps::default(), 0); + assert_eq!(missing.time_to_first_token_ms, None); + assert_eq!(missing.generation_time_ms, None); + assert_eq!(missing.queue_time_ms, None); + assert_eq!(missing.mean_itl_ms, None); + assert_eq!(missing.tokens_per_second, None); + } + + #[test] + fn stream_serialization_omits_unset_fields() { + let metrics = PerRequestMetrics::from_timestamps( + RequestTimestamps { + queued_ts: 10.0, + scheduled_ts: 10.2, + first_token_ts: 10.5, + last_token_ts: 10.5, + }, + 1, + ); + let value = serde_json::to_value(StreamPerRequestMetrics::from(metrics)).unwrap(); + let metrics = value.as_object().unwrap(); + assert_eq!(metrics.len(), 4); + assert!(metrics.contains_key("time_to_first_token_ms")); + assert!(metrics.contains_key("generation_time_ms")); + assert!(metrics.contains_key("queue_time_ms")); + assert!(metrics.contains_key("tokens_per_second")); + assert!(!metrics.contains_key("mean_itl_ms")); + } +} diff --git a/rust/src/server/src/routes/openai/utils/mod.rs b/rust/src/server/src/routes/openai/utils/mod.rs index 4bbea924a38b..a5f00ae0eadf 100644 --- a/rust/src/server/src/routes/openai/utils/mod.rs +++ b/rust/src/server/src/routes/openai/utils/mod.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project pub mod logprobs; +pub mod metrics; pub mod structured_outputs; pub mod types; pub mod usage; diff --git a/rust/src/text/src/output/decoded.rs b/rust/src/text/src/output/decoded.rs index 64df0492bfbe..6308bdf6e116 100644 --- a/rust/src/text/src/output/decoded.rs +++ b/rust/src/text/src/output/decoded.rs @@ -10,7 +10,7 @@ use tracing::{Level, debug, trace}; use vllm_engine_core_client::AbortCause; use vllm_engine_core_client::protocol::output::StopReason; use vllm_engine_core_client::protocol::sampling_mask::SamplingMask; -use vllm_llm::{FinishReason, GenerateOutput, TokenUsage}; +use vllm_llm::{FinishReason, GenerateOutput, RequestTimestamps, TokenUsage}; use vllm_tokenizer::{DecodedText, DynTokenizer, IncrementalDecoder}; use super::logprobs::{ @@ -55,6 +55,7 @@ pub struct Finished { pub ec_transfer_params: Option>, /// Sampling support sets aligned with all generated token positions. pub sampling_mask: Option, + pub timestamps: RequestTimestamps, } /// Sample metadata emitted by one engine output update. @@ -317,6 +318,7 @@ pub async fn decoded_text_event_stream( kv_transfer_params, ec_transfer_params, sampling_mask, + timestamps: output.timestamps, })), }) .await; @@ -543,7 +545,19 @@ mod tests { name: "utf8 byte fallback resolves across engine updates", outputs: vec![ GenerateOutput::for_test(Some(Arc::clone(&prompt)), vec![0xe4], None), - GenerateOutput::for_test(None, vec![0xbd, 0xa0], Some(FinishReason::Length)), + GenerateOutput { + timestamps: RequestTimestamps { + queued_ts: 8.0, + scheduled_ts: 9.0, + first_token_ts: 10.0, + last_token_ts: 11.5, + }, + ..GenerateOutput::for_test( + None, + vec![0xbd, 0xa0], + Some(FinishReason::Length), + ) + }, ], expected: vec![ DecodedTextEvent::Start { @@ -565,6 +579,12 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: RequestTimestamps { + queued_ts: 8.0, + scheduled_ts: 9.0, + first_token_ts: 10.0, + last_token_ts: 11.5, + }, }), ), ], @@ -597,6 +617,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), }), ), ], @@ -634,12 +655,19 @@ mod tests { async fn stream_stop_string_sets_task_local_abort_cause_on_raw_stream_drop() { let prompt: Arc<[u32]> = Arc::from([]); let dropped_cause = Arc::new(Mutex::new(None)); + let timestamps = RequestTimestamps { + queued_ts: 1.0, + scheduled_ts: 2.0, + first_token_ts: 3.0, + last_token_ts: 4.0, + }; let raw_stream = DropRecordingStream { - next: Some(Ok(GenerateOutput::for_test( - Some(prompt), - ascii_tokens("hello"), - Some(FinishReason::Length), - ))), + // The engine output is not terminal. The frontend terminates the + // request after decoding the configured stop string. + next: Some(Ok(GenerateOutput { + timestamps, + ..GenerateOutput::for_test(Some(prompt), ascii_tokens("hello"), None) + })), dropped_cause: Arc::clone(&dropped_cause), }; let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new()); @@ -657,6 +685,7 @@ mod tests { assert_eq!(output.text, "he"); assert!(output.finish_reason.is_stop()); + assert_eq!(output.timestamps, timestamps); assert_eq!( *dropped_cause.lock().unwrap(), Some(AbortCause::StopStringMatched) diff --git a/rust/src/text/src/output/mod.rs b/rust/src/text/src/output/mod.rs index 8e09f0761870..b69d439e3f25 100644 --- a/rust/src/text/src/output/mod.rs +++ b/rust/src/text/src/output/mod.rs @@ -31,6 +31,7 @@ pub struct CollectedTextOutput { pub token_ids: Vec, pub finish_reason: FinishReason, pub usage: vllm_llm::TokenUsage, + pub timestamps: vllm_llm::RequestTimestamps, /// Connector-specific KV transfer parameters for disaggregated serving. pub kv_transfer_params: Option, /// Connector-specific encoder cache transfer parameters for disaggregated @@ -91,6 +92,7 @@ impl T { token_ids: delta_token_ids, finish_reason: FinishReason::Error, usage: vllm_llm::TokenUsage::default(), + timestamps: vllm_llm::RequestTimestamps::default(), kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, @@ -101,6 +103,7 @@ impl T { let mut collected = collected.unwrap(); collected.finish_reason = finished.finish_reason; collected.usage = finished.usage; + collected.timestamps = finished.timestamps; collected.kv_transfer_params = finished.kv_transfer_params.map(|value| *value); collected.ec_transfer_params = @@ -181,6 +184,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), })), }), ]); @@ -304,6 +308,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), })), }), ]); diff --git a/tests/entrypoints/openai/chat_completion/test_chat_completion.py b/tests/entrypoints/openai/chat_completion/test_chat_completion.py index e0422a520aae..ff6595fb0c62 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_completion.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_completion.py @@ -20,6 +20,7 @@ def default_server_args(): "--max-num-seqs", "128", "--enforce-eager", + "--enable-per-request-metrics", ] @@ -247,3 +248,59 @@ async def test_kv_transfer_prompt_token_ids_streaming(client: openai.AsyncOpenAI # streamed text-out, reconstructed from deltas, with generated token ids. assert content assert delta_token_ids + + +def _assert_complete_timing_metrics(metrics: object) -> None: + assert isinstance(metrics, dict) + timing_fields = { + "time_to_first_token_ms", + "generation_time_ms", + "queue_time_ms", + "mean_itl_ms", + "tokens_per_second", + } + assert timing_fields <= metrics.keys() + assert all(metrics[field] >= 0 for field in timing_fields) + + +@pytest.mark.asyncio +async def test_per_request_metrics_non_streaming(client: openai.AsyncOpenAI): + response = await client.chat.completions.create( + model=MODEL_NAME, + messages=[{"role": "user", "content": "Count to four."}], + max_completion_tokens=4, + temperature=0, + extra_body={"min_tokens": 4}, + ) + + _assert_complete_timing_metrics(response.model_extra.get("metrics")) + + +@pytest.mark.asyncio +async def test_per_request_metrics_on_final_usage_chunk(client: openai.AsyncOpenAI): + stream = await client.chat.completions.create( + model=MODEL_NAME, + messages=[{"role": "user", "content": "Say hello."}], + max_completion_tokens=1, + temperature=0, + stream=True, + stream_options={"include_usage": True}, + extra_body={"min_tokens": 1}, + ) + + chunks = [chunk async for chunk in stream] + assert len(chunks) > 1 + for chunk in chunks[:-1]: + assert chunk.usage is None + assert "metrics" not in (chunk.model_extra or {}) + + usage_chunk = chunks[-1] + assert usage_chunk.usage is not None + assert usage_chunk.choices == [] + metrics = (usage_chunk.model_extra or {}).get("metrics") + assert isinstance(metrics, dict) + assert metrics["time_to_first_token_ms"] >= 0 + assert metrics["generation_time_ms"] >= 0 + assert metrics["queue_time_ms"] >= 0 + assert metrics["tokens_per_second"] >= 0 + assert "mean_itl_ms" not in metrics From 148426693823ec8fa27c353a1f8b5b9f7d402d8c Mon Sep 17 00:00:00 2001 From: Tomas Ruiz Date: Tue, 13 Oct 2026 01:01:11 +0200 Subject: [PATCH 341/624] [Perf] Improve `_topk_topp_kernel()` for masked logits (#61213) Signed-off-by: Tomas Ruiz Co-authored-by: Codex Co-authored-by: Nick Hill --- benchmarks/benchmark_topk_topp.py | 325 +++++++++++++------------ vllm/v1/sample/ops/topk_topp_triton.py | 5 +- 2 files changed, 169 insertions(+), 161 deletions(-) diff --git a/benchmarks/benchmark_topk_topp.py b/benchmarks/benchmark_topk_topp.py index 14111f733b9a..2c5bbeb3eba0 100644 --- a/benchmarks/benchmark_topk_topp.py +++ b/benchmarks/benchmark_topk_topp.py @@ -1,23 +1,29 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Benchmark comparing Triton vs PyTorch sort-based top-k/top-p implementations. +"""Benchmark top-k/top-p filtering with Triton, FlashInfer, and optional PyTorch. -Compares: -- apply_top_k_top_p_triton (Triton binary search) -- apply_top_k_top_p (PyTorch sort-based) +All implementations return masked logits. Timings include a fresh-input copy. +FlashInfer timings include softmax and mask reconstruction when its APIs require them. +L2 is flushed before each measured iteration, outside the timed region. Scenarios: - top_k only (whole batch, partial batch) - top_p only (whole batch, partial batch) - mix of top_k and top_p +- optional sparse inputs via --num-finite-logits """ import argparse import gc +import json from dataclasses import dataclass +from pathlib import Path +import flashinfer.sampling as fi_sampling +import pandas as pd import torch +from flashinfer.testing import bench_gpu_time_with_cupti from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch from vllm.v1.sample.ops.topk_topp_triton import ( @@ -65,14 +71,18 @@ def calculate_ops_pct( def create_logits( - batch_size: int, vocab_size: int, device: str = "cuda" + batch_size: int, + vocab_size: int, + device: str = "cuda", + num_finite_logits: int | None = None, ) -> torch.Tensor: """Create random logits mimicking a realistic LLM distribution. Uses a Zipf-like probability distribution (rank^-1.1) converted to logits via log, then randomly permuted per row. This produces a peaked distribution where a small number of tokens capture most probability mass, similar to - real model outputs. + real model outputs. When num_finite_logits is set, retain that many ranks + before permuting their positions independently in each row. """ # Create Zipf-like probabilities: p(rank) ~ rank^(-alpha) ranks = torch.arange(1, vocab_size + 1, dtype=torch.float32, device=device) @@ -81,6 +91,8 @@ def create_logits( # Convert to logits (log-probabilities, unnormalized is fine) base_logits = probs.log() + if num_finite_logits is not None: + base_logits[num_finite_logits:] = -float("inf") # Broadcast to batch and randomly permute each row logits = base_logits.unsqueeze(0).expand(batch_size, -1).clone() @@ -90,6 +102,26 @@ def create_logits( return logits +def apply_top_k_top_p_flashinfer_mask(logits, k, p): + """Keep top-k in logits; adapt probability-based top-p back to logits.""" + if k is not None: + logits = fi_sampling.top_k_mask_logits(logits, k) + if p is not None: + # Top-p renormalization requires softmax, even after logits-based top-k. + probs = logits.softmax(dim=-1) + probs = fi_sampling.top_p_renorm_probs(probs, p, is_deterministic=True) + # Return the selected original logits, not normalized probabilities. + logits = logits.masked_fill(probs == 0, -float("inf")) + return logits + + +IMPLEMENTATIONS = { + "triton": apply_top_k_top_p_triton, + "flashinfer-mask": apply_top_k_top_p_flashinfer_mask, + "pytorch": apply_top_k_top_p_pytorch, +} + + def measure_memory() -> tuple[int, int]: """Return (allocated, reserved) memory in bytes.""" torch.accelerator.synchronize() @@ -107,53 +139,6 @@ def reset_memory_stats(): gc.collect() -def benchmark_function( - func, - logits: torch.Tensor, - k: torch.Tensor | None, - p: torch.Tensor | None, - warmup_iters: int = 5, - benchmark_iters: int = 20, -) -> tuple[float, int]: - """Benchmark a function and return (avg_time_ms, peak_memory_bytes). - - Returns average time in milliseconds and peak memory usage. - """ - # Warmup - for _ in range(warmup_iters): - logits_copy = logits.clone() - func(logits_copy, k, p) - torch.accelerator.synchronize() - - # Reset memory stats before benchmark - reset_memory_stats() - - # Benchmark - start_events = [ - torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters) - ] - end_events = [torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)] - - for i in range(benchmark_iters): - logits_copy = logits.clone() - start_events[i].record() - func(logits_copy, k, p) - end_events[i].record() - - torch.accelerator.synchronize() - - # Calculate timing - times = [ - start_events[i].elapsed_time(end_events[i]) for i in range(benchmark_iters) - ] - avg_time = sum(times) / len(times) - - # Get peak memory - _, peak_memory = measure_memory() - - return avg_time, peak_memory - - def create_benchmark_configs( batch_sizes: list[int], vocab_sizes: list[int], @@ -286,118 +271,103 @@ def format_memory(bytes_val: int) -> str: def run_benchmark( configs: list[BenchmarkConfig], + implementations: list[str], warmup_iters: int = 5, benchmark_iters: int = 20, verbose: bool = True, + num_finite_logits: int | None = None, + check: bool = False, + output: Path | None = None, ): - """Run all benchmarks and print results.""" + """Run selected implementations with identical inputs and masked-logit outputs.""" results = [] - print("=" * 100) - print("Top-k/Top-p Benchmark: Triton vs PyTorch Sort-based") - print("=" * 100) - print() - - for config in configs: - if verbose: - print(f"Running: {config.description}") - - # Create fresh logits for this config - logits = create_logits(config.batch_size, config.vocab_size) - - # Benchmark Triton - reset_memory_stats() - triton_time, triton_mem = benchmark_function( - apply_top_k_top_p_triton, - logits, - config.k_values, - config.p_values, - warmup_iters, - benchmark_iters, + for index, config in enumerate(configs): + logits = create_logits( + config.batch_size, + config.vocab_size, + num_finite_logits=num_finite_logits, ) - - # Benchmark PyTorch - reset_memory_stats() - pytorch_time, pytorch_mem = benchmark_function( - apply_top_k_top_p_pytorch, - logits, - config.k_values, - config.p_values, - warmup_iters, - benchmark_iters, - ) - - speedup = pytorch_time / triton_time if triton_time > 0 else float("inf") - mem_ratio = pytorch_mem / triton_mem if triton_mem > 0 else float("inf") - - result = { - "config": config, - "triton_time_ms": triton_time, - "pytorch_time_ms": pytorch_time, - "triton_mem": triton_mem, - "pytorch_mem": pytorch_mem, - "speedup": speedup, - "mem_ratio": mem_ratio, - } - results.append(result) - + checks = {} + if check: + checks = check_implementations(logits, config, implementations) + k, p = config.k_values, config.p_values + timings = {} + for name in implementations: + func = IMPLEMENTATIONS[name] + reset_memory_stats() + times = bench_gpu_time_with_cupti( + # note: measured runtime includes logits.clone() + lambda func=func, logits=logits, k=k, p=p: func(logits.clone(), k, p), + dry_run_iters=warmup_iters, + repeat_iters=benchmark_iters, + use_cuda_graph=False, # sampling is outside cuda graph + cold_l2_cache=True, + ) + elapsed = sum(times) / len(times) + _, memory = measure_memory() + timings[name] = dict(time_ms=elapsed, peak_memory=memory) + row = create_result_row(config, timings, checks) + results.append(row) + if output is not None: + with output.open("a") as file: + file.write(json.dumps(row) + "\n") if verbose: - print(f" Triton: {triton_time:.3f} ms, {format_memory(triton_mem)}") - print(f" PyTorch: {pytorch_time:.3f} ms, {format_memory(pytorch_mem)}") - print(f" Speedup: {speedup:.2f}x, Memory ratio: {mem_ratio:.2f}x") - print() - - # Clean up + print(f"Completed {index + 1}/{len(configs)}: {config.name}", flush=True) del logits reset_memory_stats() return results -def print_summary_table(results: list[dict]): - """Print a summary table of results.""" - print() - print("=" * 130) - print("SUMMARY TABLE") - print("=" * 130) - print() +def check_implementations( + logits: torch.Tensor, + config: BenchmarkConfig, + implementations: list[str], +) -> dict: + """Compare filtered distributions against PyTorch outside the timed region.""" + k, p = config.k_values, config.p_values + reference = apply_top_k_top_p_pytorch(logits.clone(), k, p).softmax(dim=-1) + checks = {} + for name in implementations: + actual = IMPLEMENTATIONS[name](logits.clone(), k, p).softmax(dim=-1) + total_variation = (actual - reference).abs().sum(-1).max().item() / 2 + checks[name] = { + "close": torch.allclose(actual, reference, rtol=1e-5, atol=1e-6), + "max_total_variation": total_variation, + } + if not checks[name]["close"]: + print(f"Distribution difference: {name} {config.name} {checks[name]}") + del actual + return checks - # Header - header = ( - f"{'Scenario':<40} {'Batch':>6} {'Vocab':>7} {'Ops%':>6} " - f"{'Triton (ms)':>12} {'PyTorch (ms)':>13} {'Speedup':>8} " - f"{'Tri Mem':>10} {'Pyt Mem':>10}" - ) - print(header) - print("-" * 130) - - # Group by scenario type - current_vocab = None - for result in results: - config = result["config"] - - # Add separator between vocab sizes - if current_vocab != config.vocab_size: - if current_vocab is not None: - print("-" * 130) - current_vocab = config.vocab_size - - scenario = config.name.split("_b")[0] # Extract scenario name - print( - f"{scenario:<40} {config.batch_size:>6} {config.vocab_size:>7} " - f"{config.ops_pct:>5.0f}% " - f"{result['triton_time_ms']:>12.3f} {result['pytorch_time_ms']:>13.3f} " - f"{result['speedup']:>7.2f}x " - f"{format_memory(result['triton_mem']):>10} " - f"{format_memory(result['pytorch_mem']):>10}" - ) - print("=" * 130) +def print_summary_table(results: list[dict]): + """Print the same result rows saved to the output file.""" + print() + print(pd.DataFrame(results).to_string(index=False, float_format="{:.3g}".format)) + + +def create_result_row(config, timings, checks) -> dict: + """Build a shared row for the summary table and JSONL output.""" + row = { + "Scenario": config.name.split("_b")[0], + "Batch": config.batch_size, + "Vocab": config.vocab_size, + "Ops%": f"{config.ops_pct:.0f}%", + } + row.update({f"{name} (ms)": data["time_ms"] for name, data in timings.items()}) + for name, data in timings.items(): + row[f"{name} mem"] = format_memory(data["peak_memory"]) + for name, data in checks.items(): + row[f"{name} check"] = data["close"] + row[f"{name} max TV"] = data["max_total_variation"] + return row def main(): parser = argparse.ArgumentParser( - description="Benchmark Triton vs PyTorch sort-based top-k/top-p implementations" + description="Benchmark Triton and FlashInfer top-k/top-p filtering" ) parser.add_argument( "--batch-sizes", @@ -431,29 +401,60 @@ def main(): help="Only print summary table", ) - args = parser.parse_args() + parser.add_argument( + "--num-finite-logits", + type=int, + help="Finite logits per row; remaining entries are -inf (default: dense)", + ) - # Print configuration - print(f"Batch sizes: {args.batch_sizes}") - print(f"Vocab sizes: {args.vocab_sizes}") - print(f"Warmup iterations: {args.warmup_iters}") - print(f"Benchmark iterations: {args.benchmark_iters}") - print() + parser.add_argument( + "--implementations", + nargs="+", + choices=list(IMPLEMENTATIONS), + default=["triton", "flashinfer-mask"], + help="Implementations to compare (PyTorch is an optional reference)", + ) + + parser.add_argument( + "--check", + action="store_true", + help="Check filtered distributions against PyTorch outside the timed region", + ) + + parser.add_argument( + "--scenarios", + nargs="+", + choices=[ + "topk_whole", + "topk_partial", + "topp_whole", + "topp_partial", + "topk_topp_whole", + "mixed_partial", + ], + help="Run only the selected scenarios (default: all)", + ) + parser.add_argument("--output", type=Path, help="Save incremental JSONL results") + + args = parser.parse_args() + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.open("x").close() + if args.num_finite_logits is not None and not ( + 0 <= args.num_finite_logits <= min(args.vocab_sizes) + ): + parser.error("--num-finite-logits must be between 0 and each vocabulary size") + torch.manual_seed(0) # Check CUDA if not torch.cuda.is_available(): print("ERROR: CUDA is not available. This benchmark requires a GPU.") return - device_name = torch.cuda.get_device_name(0) - print(f"GPU: {device_name}") - print() - # Create configs - configs = create_benchmark_configs( - args.batch_sizes, - args.vocab_sizes, - ) + configs = create_benchmark_configs(args.batch_sizes, args.vocab_sizes) + if args.scenarios is not None: + configs = [c for c in configs if c.name.split("_b")[0] in args.scenarios] # Run benchmarks results = run_benchmark( @@ -461,6 +462,10 @@ def main(): warmup_iters=args.warmup_iters, benchmark_iters=args.benchmark_iters, verbose=not args.quiet, + num_finite_logits=args.num_finite_logits, + implementations=args.implementations, + check=args.check, + output=args.output, ) # Print summary diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index 151f641d215d..4aa837e9d73c 100755 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/vllm/v1/sample/ops/topk_topp_triton.py @@ -206,7 +206,10 @@ def _topk_topp_kernel( k_pivots_num = tl.zeros((), dtype=tl.uint32) min_larger = float("inf") num_min_larger = tl.zeros((), dtype=tl.uint32) - if num_outliers > k: + if num_finite_total <= k: + # All finite logits survive top-k; skip the cutoff search. + k_pivot = -float("inf") + elif num_outliers > k: max_range = max_logit min_range = outlier_pivot search_range = tl.cast(num_outliers, tl.int32) From 189bbfb9f432e283e21dcdf2a78357ae5db05ace Mon Sep 17 00:00:00 2001 From: Ilya Boytsov Date: Tue, 13 Oct 2026 00:18:43 +0100 Subject: [PATCH 342/624] Add MoNe experts pruning support (#50415) Signed-off-by: Ilya Boytsov Signed-off-by: Misha Goin Signed-off-by: mgoin Co-authored-by: Ilya Boytsov Co-authored-by: Codex Co-authored-by: Misha Goin Co-authored-by: Claude Opus 5.5 --- docs/design/expert_substitution.md | 121 +++++ pyproject.toml | 2 + tests/kernels/moe/test_expert_substitution.py | 422 ++++++++++++++++++ .../moe/test_unquantized_backend_selection.py | 52 ++- .../model_loader/test_expert_substitution.py | 337 ++++++++++++++ tests/models/test_deepseek_v4_fi_moe_ep.py | 1 + .../model_executor/layers/fused_moe/config.py | 2 + .../layers/fused_moe/expert_substitution.py | 380 ++++++++++++++++ vllm/model_executor/layers/fused_moe/layer.py | 26 +- .../layers/fused_moe/modular_kernel.py | 2 + .../model_loader/tensorizer_loader.py | 2 + vllm/model_executor/model_loader/utils.py | 6 + .../models/expert_substitution.py | 109 +++++ 13 files changed, 1458 insertions(+), 4 deletions(-) create mode 100644 docs/design/expert_substitution.md create mode 100644 tests/kernels/moe/test_expert_substitution.py create mode 100644 tests/model_executor/model_loader/test_expert_substitution.py create mode 100644 vllm/model_executor/layers/fused_moe/expert_substitution.py create mode 100644 vllm/model_executor/models/expert_substitution.py diff --git a/docs/design/expert_substitution.md b/docs/design/expert_substitution.md new file mode 100644 index 000000000000..05e69bab6239 --- /dev/null +++ b/docs/design/expert_substitution.md @@ -0,0 +1,121 @@ +# Expert substitution (MoNE) + +[Mixture of Novices and Experts (MoNE)](https://arxiv.org/abs/2507.00390) +prunes MoE layers by replacing redundant experts with "novices" that return a +constant vector. vLLM loads the retained experts into compact weights and adds +the novices' output separately, so pruned experts use no MLP weight memory and +no MoE backend needs to know about them. + +## Checkpoint format + +The checkpoint keeps the original architecture. `config.json` lists the +substituted experts of each decoder layer: + +```json +{ + "architectures": ["DeepseekV2ForCausalLM"], + "n_routed_experts": 64, + "approximate_experts": { + "1": [1, 6, 8, 12], + "2": [0, 3, 17, 40] + } +} +``` + +For VLMs, `approximate_experts` may live on the text config. Retained experts +keep their logical checkpoint names. Each substituted expert stores one +`[hidden_size]` tensor instead of its MLP: + +```text +model.layers.1.mlp.experts.0.gate_proj.weight +model.layers.1.mlp.experts.1.approx_value +``` + +Alternatively, `compression_config.transform_config.expert_substitution` +declares the same information with explicitly named constant tensors: + +```json +{ + "compression_config": { + "transform_config": { + "expert_substitution": { + "version": 1, + "router_semantics": { + "preserve_logical_expert_ids": true, + "preserve_router_weights": true, + "renormalize_after_substitution": false + }, + "targets": { + "model.layers.1.mlp.experts": { + "num_logical_experts": 64, + "weight_layout": "compact_retained_experts", + "replacements": { + "1": { + "format": "constant-v1", + "tensors": {"value": "model.layers.1.mlp.expert_replacements.1.value"} + } + } + } + } + } + } + } +} +``` + +Only the values shown for `version`, `router_semantics`, `weight_layout`, and +`format` are supported. Each target path must contain exactly one layer index. +This form is not yet accepted alongside compressed-tensors quantization, whose +`transform_config` schema does not include `expert_substitution`. + +## Semantics + +For a substituted expert `j` with constant `v_j` and router weight `w_j(x)`, +the expert contributes `w_j(x) v_j`. Routing is unchanged: router weights are +neither renormalized nor redistributed to retained experts. + +## Execution + +`FusedMoEFactory` matches each MoE layer to `approximate_experts` by the layer +index in its prefix and, for a match, builds `SubstitutedRoutedExperts`, which: + +1. allocates weights only for retained experts, in ascending logical order; +2. after routing, sums `w_j(x) v_j` over substituted routes in FP32; +3. rewrites substituted routes as zero-weight routes to physical expert 0 and + retained routes as compact physical IDs, preserving the `[num_tokens, top_k]` + contract of every decomposed backend; +4. adds the constant output on one tensor-parallel rank, before the partial + outputs are all-reduced. + +Zero-weight routes may schedule some unnecessary GEMM work; the memory savings +are unaffected. Monolithic backends route internally and are not selected. + +## Weight loading + +Retained experts are loaded through the standard expert mapping by logical ID +and placed in their compact row. Substituted IDs map to no local row, the same +way expert parallelism skips non-local experts. + +Constant tensors are consumed at the model's `load_weights` boundary, so +initial loading and direct weight updates share one path. `approx_value` tensors of +layers that the model does not build (other pipeline stages, MTP layers) are +passed through to the model's loader, which skips them like any other weight +it does not own. Explicitly named constants are always consumed, and dropped +when their layer is built elsewhere. Loading fails if a local substituted +expert has no constant. + +Direct `load_weights` updates may change individual constants. Layerwise +reloads (`reload_weights` in checkpoint format) are rejected, and the model +keeps its previous weights. + +## Supported configurations + +- Models whose MoE layers use `FusedMoEFactory` with the default + `RoutedExperts`, and whose MoE prefix contains exactly one layer index. +- Unquantized experts. +- Tensor and pipeline parallelism, including MTP draft models. +- Any decomposed MoE backend. + +Expert and data parallelism, EPLB, MoE LoRA, fused shared experts, quantized +experts, and routed input/output transforms are rejected at initialization. +Deferred MoE finalize and layerwise weight reloads are rejected when used. diff --git a/pyproject.toml b/pyproject.toml index bc924440fa59..db5a2a3b9c83 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -226,6 +226,8 @@ view_seperator = "view_seperator" inverse_std_variences = "inverse_std_variences" [tool.typos.default.extend-words] +MoNE = "MoNE" +mone = "mone" Hel = "Hel" wether = "wether" iy = "iy" diff --git a/tests/kernels/moe/test_expert_substitution.py b/tests/kernels/moe/test_expert_substitution.py new file mode 100644 index 000000000000..70720af69ff0 --- /dev/null +++ b/tests/kernels/moe/test_expert_substitution.py @@ -0,0 +1,422 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from tests.kernels.moe.modular_kernel_tools.parallel_utils import ( + ProcessGroupInfo, + parallel_launch_with_config, +) +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config +from vllm.forward_context import get_forward_context, set_forward_context +from vllm.model_executor.layers.fused_moe import FusedMoEFactory +from vllm.model_executor.layers.fused_moe.expert_substitution import ( + ConstantExpertSubstitution, + SubstitutedRoutedExperts, + get_expert_substitution_spec, + make_expert_substitution, +) +from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( + FlashInferExperts, +) +from vllm.model_executor.models.mixtral import MixtralMoE +from vllm.platforms import current_platform +from vllm.v1.worker.workspace import init_workspace_manager + +MOE_PREFIX = "model.layers.1.mlp.experts" + + +def _model_config(approximate_experts, **kwargs) -> SimpleNamespace: + return SimpleNamespace( + hf_config=SimpleNamespace(approximate_experts=approximate_experts), + **kwargs, + ) + + +def _expert_substitution(targets: dict[str, dict[int, str]]) -> dict: + """The ``compression_config.transform_config.expert_substitution`` format.""" + return { + "version": 1, + "router_semantics": { + "preserve_logical_expert_ids": True, + "preserve_router_weights": True, + "renormalize_after_substitution": False, + }, + "targets": { + path: { + "num_logical_experts": 5, + "weight_layout": "compact_retained_experts", + "replacements": { + str(expert_id): { + "format": "constant-v1", + "tensors": {"value": name}, + } + for expert_id, name in replacements.items() + }, + } + for path, replacements in targets.items() + }, + } + + +def _compressed_model_config(expert_substitution: dict, **kwargs) -> SimpleNamespace: + return SimpleNamespace( + hf_config=SimpleNamespace( + compression_config={ + "transform_config": {"expert_substitution": expert_substitution} + } + ), + **kwargs, + ) + + +def _substitution(hidden_size: int, dtype=torch.float32): + # Experts 1 and 3 of 5 are substituted; 0, 2, 4 map to physical 0, 1, 2. + return ConstantExpertSubstitution(1, 5, (1, 3), hidden_size, dtype) + + +def test_parse_approximate_experts(): + assert get_expert_substitution_spec(None) is None + assert get_expert_substitution_spec(_model_config(None)) is None + spec = get_expert_substitution_spec(_model_config({"1": [3, 1], "2": [], 5: (0,)})) + assert spec.experts == {1: (1, 3), 5: (0,)} + assert spec.value_names == {} + text_config = SimpleNamespace( + hf_text_config=SimpleNamespace(approximate_experts={"0": [2]}), + hf_config=SimpleNamespace(), + ) + assert get_expert_substitution_spec(text_config).experts == {0: (2,)} + + +@pytest.mark.parametrize( + "approximate_experts", + [[1, 2], {"layer": [1]}, {"1": [1, 1]}, {"1": [-1]}, {"1": ["x"]}], +) +def test_parse_approximate_experts_rejects_invalid_entries(approximate_experts): + with pytest.raises(ValueError, match="approximate_experts"): + get_expert_substitution_spec(_model_config(approximate_experts)) + + +def test_parse_expert_substitution_with_explicit_value_names(): + spec = get_expert_substitution_spec( + _compressed_model_config( + _expert_substitution( + { + MOE_PREFIX: {3: "shared", 1: "layer_1.expert_1"}, + "model.layers.2.mlp.experts": {0: "shared"}, + } + ) + ) + ) + assert spec.experts == {1: (1, 3), 2: (0,)} + assert spec.value_names == { + "shared": ((1, 3), (2, 0)), + "layer_1.expert_1": ((1, 1),), + } + + +@pytest.mark.parametrize( + "mutate", + [ + lambda raw: raw.update(version=2), + lambda raw: raw["router_semantics"].update(renormalize_after_substitution=True), + lambda raw: raw["targets"].update({"mlp.experts": raw["targets"][MOE_PREFIX]}), + lambda raw: raw["targets"][MOE_PREFIX].update(weight_layout="dense"), + lambda raw: raw["targets"][MOE_PREFIX]["replacements"]["1"].update( + format="linear-v1" + ), + lambda raw: raw["targets"][MOE_PREFIX]["replacements"]["1"].pop("tensors"), + ], +) +def test_parse_expert_substitution_rejects_unsupported_metadata(mutate): + raw = _expert_substitution({MOE_PREFIX: {1: "value"}}) + mutate(raw) + with pytest.raises(ValueError, match="invalid expert_substitution metadata"): + get_expert_substitution_spec(_compressed_model_config(raw)) + + +def test_parse_rejects_both_formats(): + config = _compressed_model_config(_expert_substitution({MOE_PREFIX: {1: "v"}})) + config.hf_config.approximate_experts = {"1": [1]} + with pytest.raises(ValueError, match="both"): + get_expert_substitution_spec(config) + + +def test_make_expert_substitution_matches_layer_index(): + config = _model_config({"1": [3, 1]}) + substitution = make_expert_substitution(config, MOE_PREFIX, 5, 16, torch.bfloat16) + + assert substitution is not None + assert substitution.layer_idx == 1 + assert substitution.substituted_expert_ids == (1, 3) + assert substitution.num_compute_experts == 3 + assert substitution.values.dtype == torch.bfloat16 + assert [substitution.physical_expert_id(i) for i in range(5)] == [0, -1, 1, -1, 2] + assert ( + make_expert_substitution(config, "model.layers.2.mlp.experts", 5, 16, None) + is None + ) + assert make_expert_substitution(None, MOE_PREFIX, 5, 16, None) is None + explicit = _compressed_model_config(_expert_substitution({MOE_PREFIX: {1: "v"}})) + substitution = make_expert_substitution(explicit, MOE_PREFIX, 5, 16, None) + assert substitution is not None + assert substitution.substituted_expert_ids == (1,) + with pytest.raises(ValueError, match="exactly one layer index"): + make_expert_substitution(config, "mlp.experts", 5, 16, None) + with pytest.raises(ValueError, match="within"): + make_expert_substitution(config, MOE_PREFIX, 3, 16, None) + with pytest.raises(ValueError, match="retain at least one expert"): + make_expert_substitution(_model_config({"1": [0, 1]}), MOE_PREFIX, 2, 16, None) + + +def test_constant_substitution_loads_values(): + substitution = _substitution(2) + substitution.weight_loader(substitution.values, torch.tensor([1.0, 2.0]), 3) + with pytest.raises(ValueError, match="missing approx_value .* \\[1\\]"): + substitution.validate_loaded_values("layer") + substitution.weight_loader(substitution.values, torch.tensor([3.0, 4.0]), 1) + substitution.validate_loaded_values("layer") + torch.testing.assert_close( + substitution.values, torch.tensor([[3.0, 4.0], [1.0, 2.0]]) + ) + with pytest.raises(ValueError, match="not listed"): + substitution.weight_loader(substitution.values, torch.ones(2), 2) + with pytest.raises(ValueError, match="has shape"): + substitution.weight_loader(substitution.values, torch.ones(3), 1) + + +@pytest.mark.parametrize("num_tokens", [0, 1, 3, 7]) +@pytest.mark.parametrize("hidden_size", [2, 4]) +def test_constant_substitution_transforms_routes_and_computes_side_output( + num_tokens, hidden_size +): + substitution = _substitution(2) + substitution.values.data.copy_(torch.tensor([[10.0, 20.0], [30.0, 40.0]])) + + rows = torch.arange(num_tokens) % 3 + topk_ids = torch.tensor([[1, 2], [4, 3], [99, -1]])[rows] + topk_weights = torch.tensor([[0.25, 0.75], [0.6, 0.4], [0.5, 0.5]])[rows] + constant_output = substitution.transform_routes( + torch.zeros(num_tokens, hidden_size), topk_weights, topk_ids + ) + + torch.testing.assert_close( + topk_weights, torch.tensor([[0.0, 0.75], [0.6, 0.0], [0.0, 0.0]])[rows] + ) + torch.testing.assert_close(topk_ids, torch.tensor([[0, 1], [2, 0], [0, 0]])[rows]) + expected = torch.zeros(num_tokens, hidden_size) + expected[:, :2] = torch.tensor([[2.5, 5.0], [12.0, 16.0], [0.0, 0.0]])[rows] + torch.testing.assert_close(constant_output, expected) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param( + "cuda", + marks=pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="requires a CUDA-like platform", + ), + ), + ], +) +def test_constant_substitution_preserves_router_weight_precision(dtype, device): + substitution = _substitution(1, dtype).to(device) + substitution.values.data.copy_( + torch.tensor([[60000.0], [-60000.0]], dtype=dtype, device=device) + ) + topk_weights = torch.tensor([[0.5001, 0.4999]], dtype=torch.float32, device=device) + topk_ids = torch.tensor([[1, 3]], device=device) + # Nearly cancelling constants must retain the difference between FP32 weights. + expected = (topk_weights.double() @ substitution.values.double()).to(dtype) + + actual = substitution.transform_routes( + torch.zeros(1, 1, dtype=dtype, device=device), topk_weights, topk_ids + ) + + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +def _bare_substituted_routed_experts(**attrs) -> SubstitutedRoutedExperts: + layer = SubstitutedRoutedExperts.__new__(SubstitutedRoutedExperts) + torch.nn.Module.__init__(layer) + layer.expert_substitution = _substitution(2) + for name, value in attrs.items(): + setattr(layer, name, value) + return layer + + +@pytest.mark.parametrize("tp_rank", [0, 1]) +@pytest.mark.parametrize("padded_input", [False, True]) +def test_substituted_routed_experts_adds_constant_on_one_tp_rank(tp_rank, padded_input): + def apply(*, layer, x, topk_weights, topk_ids, **kwargs): + torch.testing.assert_close(topk_weights, torch.tensor([[0.0, 0.75]])) + torch.testing.assert_close(topk_ids, torch.tensor([[0, 1]])) + return torch.ones(x.shape[0], 2) + + layer = _bare_substituted_routed_experts( + moe_config=SimpleNamespace(tp_rank=tp_rank), + quant_method=SimpleNamespace(is_monolithic=False, apply=apply), + ) + layer.expert_substitution.values.data.copy_( + torch.tensor([[10.0, 20.0], [30.0, 40.0]]) + ) + actual = layer.forward_modular( + torch.zeros(1, 4 if padded_input else 2), + torch.tensor([[0.25, 0.75]]), + torch.tensor([[1, 2]]), + ) + expected = torch.ones(1, 2) + if tp_rank == 0: + expected += torch.tensor([[2.5, 5.0]]) + torch.testing.assert_close(actual, expected) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like platform" +) +def test_standard_fused_moe_model_discovers_substitution(dist_init): + vllm_config = VllmConfig() + vllm_config.model_config = _model_config( + {"0": [1]}, dtype=torch.bfloat16, is_moe=True + ) + vllm_config.kernel_config.moe_backend = "triton" + + with set_current_vllm_config(vllm_config): + moe = MixtralMoE( + num_experts=4, + top_k=2, + hidden_size=256, + intermediate_size=512, + params_dtype=torch.bfloat16, + prefix="model.layers.0.block_sparse_moe", + ) + + routed_experts = moe.experts.routed_experts + assert isinstance(routed_experts, SubstitutedRoutedExperts) + assert routed_experts.expert_substitution.substituted_expert_ids == (1,) + assert routed_experts.w13_weight.shape[0] == 3 + assert moe.experts.moe_config.require_decomposed_backend + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like platform" +) +@pytest.mark.parametrize( + "moe_backend", + [ + pytest.param("triton", id="triton"), + pytest.param( + "flashinfer_cutlass", + id="flashinfer-cutlass", + marks=pytest.mark.skipif( + not FlashInferExperts._supports_current_device(), + reason="FlashInfer CUTLASS is unavailable on this platform", + ), + ), + ], +) +def test_fused_moe_combines_compact_experts_and_substitution(dist_init, moe_backend): + hidden_size = 256 + vllm_config = VllmConfig() + vllm_config.model_config = _model_config( + {"0": [1, 3]}, dtype=torch.bfloat16, is_moe=True + ) + vllm_config.compilation_config.static_forward_context = {} + vllm_config.kernel_config.moe_backend = moe_backend + + with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): + init_workspace_manager(torch.accelerator.current_device_index()) + layer = FusedMoEFactory( + num_experts=4, + top_k=2, + hidden_size=hidden_size, + intermediate_size=512, + params_dtype=torch.bfloat16, + prefix="model.layers.0.mlp.experts", + renormalize=False, + ).cuda() + substitution = layer.routed_experts.expert_substitution + assert layer.routed_experts.w13_weight.shape[0] == 2 + with torch.no_grad(): + layer.routed_experts.w13_weight.normal_(0, 0.01) + layer.routed_experts.w2_weight.normal_(0, 0.01) + substitution.values.normal_(0, 0.01) + layer._quant_method.process_weights_after_loading(layer.routed_experts) + + hidden_states = torch.randn(8, hidden_size, dtype=torch.bfloat16, device="cuda") + router_logits = torch.randn(8, 4, dtype=torch.float32, device="cuda") + topk_weights, topk_ids = layer.router.select_experts( + hidden_states.clone(), router_logits + ) + expected = substitution.transform_routes(hidden_states, topk_weights, topk_ids) + expected += layer._quant_method.apply( + layer=layer.routed_experts, + x=hidden_states.clone(), + topk_weights=topk_weights, + topk_ids=topk_ids, + shared_experts=None, + shared_experts_input=None, + ) + + get_forward_context().all_moe_layers = None + actual = layer(hidden_states.clone(), router_logits) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +def _tp_substitution_worker( + process_group: ProcessGroupInfo, + vllm_config: VllmConfig, + _cpu_group, +) -> None: + with set_forward_context(None, vllm_config): + init_workspace_manager(process_group.local_rank) + layer = FusedMoEFactory( + num_experts=4, + top_k=1, + hidden_size=256, + intermediate_size=512, + params_dtype=torch.bfloat16, + prefix="model.layers.0.mlp.experts", + renormalize=True, + ) + with torch.no_grad(): + layer.routed_experts.w13_weight.zero_() + layer.routed_experts.w2_weight.zero_() + layer.routed_experts.expert_substitution.values.fill_(0.25) + layer._quant_method.process_weights_after_loading(layer.routed_experts) + + hidden_states = torch.zeros(2, 256, dtype=torch.bfloat16) + router_logits = torch.tensor( + [[-100.0, 100.0, -100.0, -100.0]] * 2, + dtype=torch.float32, + ) + get_forward_context().all_moe_layers = None + actual = layer(hidden_states, router_logits) + torch.testing.assert_close( + actual, torch.full_like(actual, 0.25), atol=0, rtol=0 + ) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike() or current_platform.device_count() < 2, + reason="requires two CUDA-like devices", +) +def test_constant_substitution_is_added_once_with_tp2(): + vllm_config = VllmConfig(parallel_config=ParallelConfig(tensor_parallel_size=2)) + vllm_config.model_config = _model_config( + {"0": [1]}, + dtype=torch.bfloat16, + is_moe=True, + sleep_mode_offload_cudagraph=False, + ) + vllm_config.compilation_config.static_forward_context = {} + vllm_config.kernel_config.moe_backend = "triton" + + parallel_launch_with_config(2, _tp_substitution_worker, vllm_config, None) diff --git a/tests/kernels/moe/test_unquantized_backend_selection.py b/tests/kernels/moe/test_unquantized_backend_selection.py index 25e6b7dec618..b64e909fbd45 100644 --- a/tests/kernels/moe/test_unquantized_backend_selection.py +++ b/tests/kernels/moe/test_unquantized_backend_selection.py @@ -136,7 +136,13 @@ def test_arm_cpu_unquantized_kernel_selection( ], ) @patch( - "vllm.utils.flashinfer.has_flashinfer", + "vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe." + "has_flashinfer_cutlass_fused_moe", + return_value=False, +) +@patch( + "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe." + "has_flashinfer_trtllm_fused_moe", return_value=False, ) @patch( @@ -145,13 +151,18 @@ def test_arm_cpu_unquantized_kernel_selection( ) def test_select_default_backend_by_platform( mock_aiter_enabled, - mock_has_flashinfer, + mock_has_flashinfer_trtllm, + mock_has_flashinfer_cutlass, monkeypatch, platform_method, expected_backend, ): """Test default backend selection per platform with all optional - accelerators (FlashInfer, AITER) disabled.""" + accelerators (FlashInfer, AITER) disabled. + + Patch the feature probes directly because their cached results can outlive + changes to has_flashinfer(). + """ with patch( "vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform" ) as mock_platform: @@ -407,6 +418,41 @@ def test_select_cuda_flashinfer_trtllm_modular_for_standard_all2all( assert experts_cls.__name__ == "TrtLlmBf16ExpertsModular" +@patch( + "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsBase._supports_current_device", + return_value=True, +) +@pytest.mark.parametrize("require_decomposed_backend", [False, True]) +def test_require_decomposed_backend_rejects_monolithic_kernels( + mock_supports_current_device, + require_decomposed_backend, +): + """Expert substitution adjusts routes, so the kernel must not route itself.""" + from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import ( + TrtLlmBf16ExpertsMonolithic, + ) + from vllm.model_executor.layers.fused_moe.modular_kernel import ( + FusedMoEActivationFormat, + ) + + moe_config = make_dummy_moe_config() + moe_config.routing_method = RoutingMethodType.Renormalize + moe_config.require_decomposed_backend = require_decomposed_backend + + supported, reason = TrtLlmBf16ExpertsMonolithic.is_supported_config( + TrtLlmBf16ExpertsMonolithic, + moe_config, + None, + None, + FusedMoEActivationFormat.Standard, + ) + + assert supported != require_decomposed_backend + assert (reason is not None and "routes adjusted" in reason) == ( + require_decomposed_backend + ) + + @patch( "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsBase._supports_current_device", return_value=True, diff --git a/tests/model_executor/model_loader/test_expert_substitution.py b/tests/model_executor/model_loader/test_expert_substitution.py new file mode 100644 index 000000000000..926720c68629 --- /dev/null +++ b/tests/model_executor/model_loader/test_expert_substitution.py @@ -0,0 +1,337 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from safetensors.torch import save_file +from transformers import DeepseekV2Config, DeepseekV2ForCausalLM + +from vllm.config import ( + AttentionConfig, + DeviceConfig, + KernelConfig, + LoadConfig, + ModelConfig, + VllmConfig, +) +from vllm.model_executor.layers.fused_moe.expert_substitution import ( + ConstantExpertSubstitution, +) +from vllm.model_executor.model_loader import get_model_loader +from vllm.model_executor.models.expert_substitution import as_expert_substitution_model +from vllm.model_executor.models.utils import AutoWeightsLoader, PPMissingLayer +from vllm.platforms import current_platform +from vllm.v1.attention.backends.registry import AttentionBackendEnum + +MODEL_CONFIG = SimpleNamespace( + hf_config=SimpleNamespace(approximate_experts={"0": [1, 2], "1": [1, 2]}) +) + + +def _compression_config(value_names: dict[int, dict[int, str]]) -> dict: + """Explicitly named constants given as ``{layer: {expert: tensor name}}``.""" + targets = { + f"model.layers.{layer}.mlp.experts": { + "num_logical_experts": 3, + "weight_layout": "compact_retained_experts", + "replacements": { + str(expert): {"format": "constant-v1", "tensors": {"value": name}} + for expert, name in names.items() + }, + } + for layer, names in value_names.items() + } + router_semantics = { + "preserve_logical_expert_ids": True, + "preserve_router_weights": True, + "renormalize_after_substitution": False, + } + return { + "transform_config": { + "expert_substitution": { + "version": 1, + "router_semantics": router_semantics, + "targets": targets, + } + } + } + + +# Expert 1 of both layers shares one explicitly named constant. +EXPLICIT_MODEL_CONFIG = SimpleNamespace( + hf_config=SimpleNamespace( + compression_config=_compression_config( + { + layer: {1: "constants.shared", 2: f"constants.layer_{layer}"} + for layer in (0, 1) + } + ) + ) +) + + +def _approx_value(layer: int, expert: int) -> str: + return f"model.layers.{layer}.mlp.experts.{expert}.approx_value" + + +def _cpu_model(local_layers=(0, 1), tracked=True, model_config=MODEL_CONFIG): + """Two decoder layers; the others are placeholders as on another PP stage.""" + + class Model(torch.nn.Module): + def __init__(self, prefix=""): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(2), requires_grad=False) + self.model = torch.nn.Module() + self.model.layers = torch.nn.ModuleList() + for index in range(2): + if index not in local_layers: + self.model.layers.append(PPMissingLayer()) + continue + layer = torch.nn.Module() + layer.mlp = torch.nn.Module() + layer.mlp.experts = ConstantExpertSubstitution( + index, 3, (1, 2), 2, torch.float32 + ) + self.model.layers.append(layer) + + def load_weights(self, weights): + loaded = AutoWeightsLoader(self).load_weights(weights) + return loaded if tracked else None + + return as_expert_substitution_model(Model, model_config)() + + +@pytest.mark.parametrize("local_layers", [(0, 1), (0,), (1,), ()]) +def test_loads_local_approx_values_and_passes_through_others(local_layers): + """Other PP stages and MTP drafts own none or some of the configured layers.""" + model = _cpu_model(local_layers) + loaded = model.load_weights( + [ + ("weight", torch.tensor([1.0, 2.0])), + *( + (_approx_value(layer, expert), torch.full((2,), 10.0 * layer + expert)) + for layer in range(2) + for expert in (1, 2) + ), + ] + ) + model.process_weights_after_loading() + + assert loaded == {"weight"} | { + f"model.layers.{layer}.mlp.experts.values" for layer in local_layers + } + for layer in local_layers: + torch.testing.assert_close( + model.model.layers[layer].mlp.experts.values, + torch.tensor([[1.0, 1.0], [2.0, 2.0]]) + 10.0 * layer, + ) + + +@pytest.mark.parametrize("local_layers", [(0, 1), (1,), ()]) +def test_loads_explicitly_named_values(local_layers): + """Named values are consumed even when their layer is built elsewhere.""" + model = _cpu_model(local_layers, model_config=EXPLICIT_MODEL_CONFIG) + loaded = model.load_weights( + [ + ("weight", torch.tensor([1.0, 2.0])), + ("constants.shared", torch.full((2,), 1.0)), + ("constants.layer_0", torch.full((2,), 2.0)), + ("constants.layer_1", torch.full((2,), 3.0)), + ] + ) + model.process_weights_after_loading() + + assert loaded == {"weight"} | { + f"model.layers.{layer}.mlp.experts.values" for layer in local_layers + } + for layer in local_layers: + torch.testing.assert_close( + model.model.layers[layer].mlp.experts.values, + torch.tensor([[1.0, 1.0], [2.0 + layer, 2.0 + layer]]), + ) + + +def test_direct_reload_updates_values_incrementally(): + model = _cpu_model(local_layers=(0,), tracked=False) + values = model.model.layers[0].mlp.experts.values + model.load_weights([(_approx_value(0, 2), torch.ones(2))]) + with pytest.raises(ValueError, match="missing approx_value .* \\[1\\]"): + model.process_weights_after_loading() + + assert model.load_weights([(_approx_value(0, 1), torch.zeros(2))]) is None + model.process_weights_after_loading() + model.load_weights([(_approx_value(0, 2), torch.full((2,), 3.0))]) + torch.testing.assert_close(values, torch.tensor([[0.0, 0.0], [3.0, 3.0]])) + + with pytest.raises(ValueError, match="not listed"): + model.load_weights([(_approx_value(0, 0), torch.ones(2))]) + + +def test_layerwise_reload_is_rejected_and_keeps_weights(): + from vllm.model_executor.model_loader.reload import ( + initialize_layerwise_reload, + record_metadata_for_reloading, + ) + + model = _cpu_model(local_layers=(0,)) + record_metadata_for_reloading(model) + loaded = [ + ("weight", torch.tensor([3.0, 4.0])), + (_approx_value(0, 1), torch.zeros(2)), + (_approx_value(0, 2), torch.ones(2)), + ] + model.load_weights(loaded) + model.process_weights_after_loading() + substitution = model.model.layers[0].mlp.experts + values = substitution.values + + initialize_layerwise_reload(model) + with pytest.raises(NotImplementedError, match="layerwise weight reload"): + model.load_weights([(name, weight + 1) for name, weight in loaded]) + + assert substitution.values is values + torch.testing.assert_close(values, torch.tensor([[0.0, 0.0], [1.0, 1.0]])) + torch.testing.assert_close(model.weight, torch.tensor([3.0, 4.0])) + assert substitution.expert_substitution_routes.tolist() == [0, -1, -2] + + +@pytest.mark.parametrize("missing_value", [False, True]) +def test_tensorizer_finalizes_local_substitution_values(monkeypatch, missing_value): + """Raw Tensorizer loading must reject an unloaded local constant row.""" + from vllm.model_executor.model_loader import tensorizer_loader + + model = _cpu_model(local_layers=(0,)) + weights = [(_approx_value(0, 1), torch.zeros(2))] + if not missing_value: + weights.append((_approx_value(0, 2), torch.ones(2))) + finalize = Mock(wraps=model.process_weights_after_loading) + monkeypatch.setattr(model, "process_weights_after_loading", finalize) + monkeypatch.setattr(tensorizer_loader, "initialize_model", lambda **kwargs: model) + loader = object.__new__(tensorizer_loader.TensorizerLoader) + monkeypatch.setattr(loader, "_get_weights_iterator", lambda: iter(weights)) + config = SimpleNamespace( + model_config=SimpleNamespace(dtype=torch.float32), + device_config=SimpleNamespace(device="cpu"), + ) + + if missing_value: + with pytest.raises(ValueError, match=r"missing approx_value .* \[2\]"): + loader._load_model_serialized_cpu(config) + else: + assert loader._load_model_serialized_cpu(config) is model + torch.testing.assert_close( + model.model.layers[0].mlp.experts.values, + torch.tensor([[0.0, 0.0], [1.0, 1.0]]), + ) + finalize.assert_called_once_with() + + +def _write_tiny_deepseek_mone_checkpoint( + model_dir: Path, explicit_names: bool +) -> tuple[dict[str, torch.Tensor], str]: + config = DeepseekV2Config( + architectures=["DeepseekV2ForCausalLM"], + vocab_size=32, + hidden_size=16, + intermediate_size=32, + moe_intermediate_size=8, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + max_position_embeddings=32, + first_k_dense_replace=0, + n_routed_experts=3, + n_shared_experts=1, + num_experts_per_tok=1, + n_group=1, + topk_group=1, + norm_topk_prob=False, + q_lora_rank=None, + kv_lora_rank=4, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + tie_word_embeddings=False, + dtype="float16", + ) + experts = "model.layers.0.mlp.experts" + if explicit_names: + value_name = "model.layers.0.mlp.expert_replacements.1.value" + config.compression_config = _compression_config({0: {1: value_name}}) + else: + value_name = f"{experts}.1.approx_value" + config.approximate_experts = {"0": [1]} + + hf_weights = DeepseekV2ForCausalLM(config).to(torch.float16).state_dict() + gate_up = hf_weights.pop(f"{experts}.gate_up_proj") + down = hf_weights.pop(f"{experts}.down_proj") + weights = {name: value.contiguous() for name, value in hf_weights.items()} + for expert in (0, 2): + gate, up = gate_up[expert].chunk(2, dim=0) + weights[f"{experts}.{expert}.gate_proj.weight"] = gate.contiguous() + weights[f"{experts}.{expert}.up_proj.weight"] = up.contiguous() + weights[f"{experts}.{expert}.down_proj.weight"] = down[expert].contiguous() + weights[value_name] = torch.arange(16, dtype=torch.float16) + config.save_pretrained(model_dir) + save_file(weights, model_dir / "model.safetensors") + return weights, value_name + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like device" +) +@pytest.mark.usefixtures("dist_init", "workspace_init") +@pytest.mark.parametrize( + "explicit_names", [False, True], ids=["approximate_experts", "expert_substitution"] +) +def test_loads_mone_checkpoint_into_compact_experts( + tmp_path: Path, explicit_names: bool +): + weights, value_name = _write_tiny_deepseek_mone_checkpoint(tmp_path, explicit_names) + model_config = ModelConfig( + model=str(tmp_path), + tokenizer=str(tmp_path), + skip_tokenizer_init=True, + dtype="float16", + max_model_len=32, + enforce_eager=True, + ) + load_config = LoadConfig(load_format="safetensors", use_tqdm_on_load=False) + vllm_config = VllmConfig( + model_config=model_config, + device_config=DeviceConfig(device="cuda"), + load_config=load_config, + attention_config=AttentionConfig(backend=AttentionBackendEnum.TRITON_MLA), + kernel_config=KernelConfig(moe_backend="triton"), + ) + model = get_model_loader(load_config).load_model(vllm_config, model_config) + + experts = "model.layers.0.mlp.experts" + routed_experts = model.model.layers[0].mlp.experts.routed_experts + # Logical expert 2 is stored in physical row 1. + torch.testing.assert_close( + routed_experts.w13_weight[1].cpu(), + torch.cat( + [ + weights[f"{experts}.2.gate_proj.weight"], + weights[f"{experts}.2.up_proj.weight"], + ] + ), + ) + torch.testing.assert_close( + routed_experts.w2_weight[1].cpu(), weights[f"{experts}.2.down_proj.weight"] + ) + substitution = routed_experts.expert_substitution + torch.testing.assert_close( + substitution.values.cpu(), weights[value_name].unsqueeze(0) + ) + + updated = weights[value_name] + 1 + assert model.load_weights([(value_name, updated)]) == { + f"{experts}.routed_experts.expert_substitution.values" + } + torch.testing.assert_close(substitution.values.cpu(), updated.unsqueeze(0)) diff --git a/tests/models/test_deepseek_v4_fi_moe_ep.py b/tests/models/test_deepseek_v4_fi_moe_ep.py index 7687a6786f7e..d98bb26a2a3b 100644 --- a/tests/models/test_deepseek_v4_fi_moe_ep.py +++ b/tests/models/test_deepseek_v4_fi_moe_ep.py @@ -142,6 +142,7 @@ def _megakernel_moe(moe_backend: str, **overrides) -> SimpleNamespace: is_act_and_mul=True, activation=MoEActivation.SILU, has_hash_routing=False, + require_decomposed_backend=False, routing_method=RoutingMethodType.DeepseekV4, router_logits_dtype=torch.float32, hidden_dim=256, diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index f20af8a6f826..e79700ae44d9 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1293,6 +1293,8 @@ class FusedMoEConfig: has_bias: bool = False is_lora_enabled: bool = False has_hash_routing: bool = False + # Routes are adjusted between routing and expert execution. + require_decomposed_backend: bool = False shared_expert_prefix: str | None = None # When True, the MoE skips its final cross-rank all-reduce (and the separate diff --git a/vllm/model_executor/layers/fused_moe/expert_substitution.py b/vllm/model_executor/layers/fused_moe/expert_substitution.py new file mode 100644 index 000000000000..2483200b0150 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/expert_substitution.py @@ -0,0 +1,380 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Constant expert substitution for MoNE-pruned checkpoints. + +MoNE (https://arxiv.org/abs/2507.00390) replaces redundant experts with +"novices" that return a constant vector. Checkpoints list the replaced experts +per decoder layer either in ``config.approximate_experts``, with constants named +``..approx_value``, or in +``compression_config.transform_config.expert_substitution``, with explicitly +named constants. Retained experts keep their logical checkpoint names and are +loaded into compact physical rows. +""" + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +import torch +from torch import nn + +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput +from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts +from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( + UnquantizedFusedMoEMethod, +) + +if TYPE_CHECKING: + from vllm.model_executor.layers.fused_moe.runner.shared_experts import SharedExperts + +APPROX_VALUE_SUFFIX = ".approx_value" + + +def decoder_layer_index(path: str) -> int | None: + """Return the only integer component of a module path, if unique.""" + indices = [int(part) for part in path.split(".") if part.isdigit()] + return indices[0] if len(indices) == 1 else None + + +@dataclass(frozen=True) +class ExpertSubstitutionSpec: + # Sorted substituted expert IDs per decoder layer. + experts: dict[int, tuple[int, ...]] + # Explicitly named constant tensors and the (layer, expert) pairs they + # supply. Without an entry, constants use the ``approx_value`` naming. + value_names: dict[str, tuple[tuple[int, int], ...]] = field(default_factory=dict) + + +# The only router semantics that the constant-v1 format supports. +_ROUTER_SEMANTICS = { + "preserve_logical_expert_ids": True, + "preserve_router_weights": True, + "renormalize_after_substitution": False, +} + + +def _expert_ids(layer_idx: int, expert_ids: Any) -> tuple[int, ...]: + ids = tuple(sorted(int(expert_id) for expert_id in expert_ids or ())) + if len(set(ids)) != len(ids) or any(expert_id < 0 for expert_id in ids): + raise ValueError( + f"substituted experts of layer {layer_idx} must be unique, " + f"non-negative expert IDs, got {list(ids)}" + ) + return ids + + +def _parse_approximate_experts(raw: Any) -> ExpertSubstitutionSpec: + if not isinstance(raw, Mapping): + raise ValueError("approximate_experts must map layer indices to expert IDs") + experts: dict[int, tuple[int, ...]] = {} + for layer, expert_ids in raw.items(): + try: + layer_idx = int(layer) + experts[layer_idx] = _expert_ids(layer_idx, expert_ids) + except (TypeError, ValueError) as e: + raise ValueError( + f"approximate_experts has an invalid entry for layer {layer!r}" + ) from e + return ExpertSubstitutionSpec({k: v for k, v in experts.items() if v}) + + +def _parse_expert_substitution(raw: Any) -> ExpertSubstitutionSpec: + experts: dict[int, tuple[int, ...]] = {} + value_names: dict[str, list[tuple[int, int]]] = {} + try: + if raw["version"] != 1 or raw["router_semantics"] != _ROUTER_SEMANTICS: + raise ValueError("only version 1 with unchanged routing is supported") + for path, target in raw["targets"].items(): + layer_idx = decoder_layer_index(path) + if layer_idx is None or layer_idx in experts: + raise ValueError(f"target {path!r} must name one decoder layer") + if target["weight_layout"] != "compact_retained_experts": + raise ValueError(f"unsupported weight layout for {path!r}") + for expert_id, replacement in target["replacements"].items(): + if replacement["format"] != "constant-v1": + raise ValueError(f"unsupported replacement format for {path!r}") + value_names.setdefault(replacement["tensors"]["value"], []).append( + (layer_idx, int(expert_id)) + ) + experts[layer_idx] = _expert_ids(layer_idx, target["replacements"]) + except (AttributeError, KeyError, TypeError, ValueError) as e: + raise ValueError(f"invalid expert_substitution metadata: {e}") from e + return ExpertSubstitutionSpec( + experts, {name: tuple(pairs) for name, pairs in value_names.items()} + ) + + +def get_expert_substitution_spec(model_config: Any) -> ExpertSubstitutionSpec | None: + """Parse the substituted experts declared in a model's HF config.""" + hf_configs = [ + getattr(model_config, name, None) for name in ("hf_text_config", "hf_config") + ] + approximate = None + for config in hf_configs: + approximate = getattr(config, "approximate_experts", None) + if approximate is not None: + break + compression_config = getattr(hf_configs[1], "compression_config", None) + transform_config = ( + compression_config.get("transform_config") + if isinstance(compression_config, Mapping) + else None + ) + declared = ( + transform_config.get("expert_substitution") + if isinstance(transform_config, Mapping) + else None + ) + if approximate is not None and declared is not None: + raise ValueError( + "config declares both approximate_experts and " + "compression_config.transform_config.expert_substitution" + ) + if approximate is not None: + spec = _parse_approximate_experts(approximate) + elif declared is not None: + spec = _parse_expert_substitution(declared) + else: + return None + return spec if spec.experts else None + + +class ConstantExpertSubstitution(nn.Module): + """Constant outputs for the substituted experts of one MoE layer. + + Router outputs keep logical expert IDs. ``transform_routes`` maps retained + experts to compact physical rows and turns substituted routes into + zero-weight routes to physical expert 0, so any decomposed MoE backend can + run them unchanged. + """ + + def __init__( + self, + layer_idx: int, + num_logical_experts: int, + substituted_expert_ids: tuple[int, ...], + hidden_size: int, + params_dtype: torch.dtype, + ) -> None: + super().__init__() + substituted = set(substituted_expert_ids) + if not substituted <= set(range(num_logical_experts)): + raise ValueError( + f"approximate_experts for layer {layer_idx} must be within " + f"[0, {num_logical_experts}), got {sorted(substituted)}" + ) + if len(substituted) == num_logical_experts: + raise ValueError( + f"approximate_experts for layer {layer_idx} must retain at least " + "one expert" + ) + self.layer_idx = layer_idx + self.substituted_expert_ids = tuple(sorted(substituted)) + retained = [i for i in range(num_logical_experts) if i not in substituted] + self.num_compute_experts = len(retained) + + # Retained experts map to their physical row r >= 0, substituted + # experts to -1 - (their row in `values`). + routes = [0] * num_logical_experts + for row, expert_id in enumerate(retained): + routes[expert_id] = row + for row, expert_id in enumerate(self.substituted_expert_ids): + routes[expert_id] = -1 - row + self._routes = tuple(routes) + self.register_buffer( + "expert_substitution_routes", + torch.tensor(routes, dtype=torch.int32), + persistent=False, + ) + + # NaN marks values that the checkpoint has not provided yet. + self.values = nn.Parameter( + torch.full((len(substituted), hidden_size), torch.nan, dtype=params_dtype), + requires_grad=False, + ) + self.values.weight_loader = self.weight_loader + + @property + def num_logical_experts(self) -> int: + return len(self._routes) + + def physical_expert_id(self, expert_id: int) -> int: + """Physical row of a retained expert, or -1 for a substituted one.""" + return max(self._routes[expert_id], -1) + + def weight_loader( + self, param: nn.Parameter, loaded_weight: torch.Tensor, expert_id: int + ) -> None: + if not 0 <= expert_id < self.num_logical_experts or ( + self._routes[expert_id] >= 0 + ): + raise ValueError( + f"expert {expert_id} of layer {self.layer_idx} is not listed in " + "approximate_experts" + ) + row = param.data[-1 - self._routes[expert_id]] + if loaded_weight.shape != row.shape: + raise ValueError( + f"approx_value for expert {expert_id} of layer {self.layer_idx} has " + f"shape {tuple(loaded_weight.shape)}, expected {tuple(row.shape)}" + ) + row.copy_(loaded_weight) + + def validate_loaded_values(self, prefix: str) -> None: + missing_rows = torch.isnan(self.values).any(dim=1).nonzero().flatten().tolist() + if missing_rows: + missing = [self.substituted_expert_ids[row] for row in missing_rows] + raise ValueError( + f"{prefix} is missing approx_value for {len(missing)} " + f"expert(s), first missing expert IDs: {missing[:8]}" + ) + + def transform_routes( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + ) -> torch.Tensor: + """Remap routes in place and return the constant experts' output.""" + valid = (topk_ids >= 0) & (topk_ids < self.num_logical_experts) + routes = self.expert_substitution_routes[ + torch.where(valid, topk_ids, torch.zeros_like(topk_ids)).long() + ].long() + substituted = valid & (routes < 0) + + # Sum each token's router weights per substituted expert in FP32; other + # routes land in a trailing column that is dropped. + num_values, value_size = self.values.shape + scales = topk_weights.new_zeros( + (topk_ids.shape[0], num_values + 1), dtype=torch.float32 + ) + scales.scatter_add_( + 1, + torch.where(substituted, -1 - routes, num_values), + topk_weights.float(), + ) + output = torch.zeros_like(hidden_states) + output[..., :value_size] = scales[:, :num_values] @ self.values.float() + + topk_weights.masked_fill_(~valid | substituted, 0.0) + topk_ids.copy_(routes.clamp_min(0).to(topk_ids.dtype)) + return output + + +def make_expert_substitution( + model_config: Any, + prefix: str, + num_experts: int, + hidden_size: int, + params_dtype: torch.dtype | None, +) -> ConstantExpertSubstitution | None: + spec = get_expert_substitution_spec(model_config) + if spec is None: + return None + layer_idx = decoder_layer_index(prefix) + if layer_idx is None: + raise ValueError( + f"cannot match MoE layer {prefix!r} to substituted experts: expected " + "exactly one layer index in its prefix" + ) + expert_ids = spec.experts.get(layer_idx) + if expert_ids is None: + return None + return ConstantExpertSubstitution( + layer_idx, + num_experts, + expert_ids, + hidden_size, + params_dtype or torch.get_default_dtype(), + ) + + +class SubstitutedRoutedExperts(RoutedExperts): + """Routed experts that hold weights only for retained experts.""" + + def __init__( + self, *args: Any, expert_substitution: ConstantExpertSubstitution, **kwargs: Any + ): + super().__init__(*args, **kwargs) + self.expert_substitution = expert_substitution + moe_config = self.moe_config + parallel_config = moe_config.moe_parallel_config + unsupported = [ + name + for name, enabled in ( + ( + "quantized MoE weights", + not isinstance(self.quant_method, UnquantizedFusedMoEMethod), + ), + ("monolithic MoE backends", self.quant_method.is_monolithic), + ("MoE LoRA", moe_config.is_lora_enabled), + ( + "EPLB", + parallel_config.enable_eplb + or moe_config.num_experts + != expert_substitution.num_compute_experts, + ), + ( + "fused shared experts", + self.expert_map_manager.num_fused_shared_experts > 0, + ), + ("expert parallelism", parallel_config.use_ep), + ("data parallelism", parallel_config.dp_size > 1), + ("prefill context parallelism", parallel_config.pcp_size > 1), + ("sequence parallelism", parallel_config.is_sequence_parallel), + ("router weights on expert inputs", self.apply_router_weight_on_input), + ) + if enabled + ] + if unsupported: + raise NotImplementedError( + "expert substitution does not support: " + ", ".join(unsupported) + ) + + def _map_global_expert_id_to_local_expert_id(self, expert_id: int) -> int: + # Checkpoints name retained experts by their logical ID. + physical_id = self.expert_substitution.physical_expert_id(expert_id) + if physical_id < 0: + return -1 + return super()._map_global_expert_id_to_local_expert_id(physical_id) + + def get_expert_mapping( + self, + ckpt_gate_proj_name: str | None = None, + ckpt_down_proj_name: str | None = None, + ckpt_up_proj_name: str | None = None, + include_fused: bool = False, + ) -> list[tuple[str, str, int, str]]: + # Fused checkpoint tensors would cover every logical expert at once. + return self.build_expert_params_mapping( + ckpt_gate_proj_name or self.ckpt_gate_proj_name, + ckpt_down_proj_name or self.ckpt_down_proj_name, + ckpt_up_proj_name or self.ckpt_up_proj_name, + num_experts=self.expert_substitution.num_logical_experts, + routed_experts_prefix="", + lora_base_layer_prefix=self.lora_base_layer_prefix, + ) + + def forward_modular( + self, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts: "SharedExperts | None" = None, + shared_experts_input: torch.Tensor | None = None, + ) -> torch.Tensor: + constant_output = self.expert_substitution.transform_routes( + x, topk_weights, topk_ids + ) + output = super().forward_modular( + x, topk_weights, topk_ids, shared_experts, shared_experts_input + ) + if isinstance(output, UnfinalizedMoEOutput): + raise NotImplementedError( + "expert substitution does not support deferred MoE finalize" + ) + # The constant output is replicated across TP ranks, whose partial + # outputs are all-reduced later, so add it on one rank only. + if self.moe_config.tp_rank == 0: + output.add_(constant_output[..., : output.shape[-1]]) + return output diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 6b7c64bdc429..839eba5d788a 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -23,6 +23,10 @@ from vllm.model_executor.layers.fused_moe.expert_map_manager import ( ExpertMapManager, ) +from vllm.model_executor.layers.fused_moe.expert_substitution import ( + SubstitutedRoutedExperts, + make_expert_substitution, +) from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( FusedMoERouter, @@ -253,6 +257,12 @@ def FusedMoEFactory( fuse_shared_experts, ) ) + expert_substitution = make_expert_substitution( + vllm_config.model_config, layer_name, num_experts, hidden_size, params_dtype + ) + if expert_substitution is not None: + # Route over all logical experts but allocate only the retained ones. + global_num_experts -= len(expert_substitution.substituted_expert_ids) # Initialize EPLB manager (or None?) eplb_state: EplbLayerState | None = None @@ -367,6 +377,7 @@ def FusedMoEFactory( device=vllm_config.device_config.device, routing_method=router.routing_method_type, # Not ideal has_hash_routing=hash_indices_table is not None, + require_decomposed_backend=expert_substitution is not None, swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, @@ -380,7 +391,20 @@ def FusedMoEFactory( # Create RoutedExperts instance BEFORE create_weights() # This will hold all expert weight parameters - if routed_experts_cls is None: + if expert_substitution is not None: + if routed_experts_cls not in (None, RoutedExperts) or ( + routed_input_transform is not None or routed_output_transform is not None + ): + raise NotImplementedError( + "expert substitution requires the default RoutedExperts and no " + "routed input or output transforms" + ) + routed_experts_cls = SubstitutedRoutedExperts + routed_experts_args = { + **(routed_experts_args or {}), + "expert_substitution": expert_substitution, + } + elif routed_experts_cls is None: routed_experts_cls = RoutedExperts assert params_dtype is not None diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 7a4cbd423746..d9a75c55ec40 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -569,6 +569,8 @@ def _make_reason(reason: str) -> str: ) elif moe_config.has_hash_routing and cls.is_monolithic(): return False, _make_reason("hash routing") + elif moe_config.require_decomposed_backend and cls.is_monolithic(): + return False, _make_reason("routes adjusted after routing") elif not cls._supports_routing_method( moe_config.routing_method, weight_key, activation_key ): diff --git a/vllm/model_executor/model_loader/tensorizer_loader.py b/vllm/model_executor/model_loader/tensorizer_loader.py index 338f9eac072a..ae74beb7ed5f 100644 --- a/vllm/model_executor/model_loader/tensorizer_loader.py +++ b/vllm/model_executor/model_loader/tensorizer_loader.py @@ -84,6 +84,8 @@ def _load_model_serialized_cpu( model = initialize_model(vllm_config=vllm_config, prefix=prefix) model.load_weights(self._get_weights_iterator()) + if hasattr(model, "process_weights_after_loading"): + model.process_weights_after_loading() return model.eval() def download_model(self, model_config: ModelConfig) -> None: diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index 65d1ff2808d7..8a236941002b 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -84,6 +84,12 @@ def initialize_model( if model_class is None: model_class, _ = get_model_architecture(model_config) + from vllm.model_executor.models.expert_substitution import ( + as_expert_substitution_model, + ) + + model_class = as_expert_substitution_model(model_class, model_config) + if vllm_config.quant_config is not None: configure_quant_config(vllm_config.quant_config, model_class) diff --git a/vllm/model_executor/models/expert_substitution.py b/vllm/model_executor/models/expert_substitution.py new file mode 100644 index 000000000000..0799bca8b7e0 --- /dev/null +++ b/vllm/model_executor/models/expert_substitution.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Iterable +from functools import wraps +from typing import Any, TypeVar + +import torch +from torch import nn + +from vllm.model_executor.layers.fused_moe.expert_substitution import ( + APPROX_VALUE_SUFFIX, + ConstantExpertSubstitution, + decoder_layer_index, + get_expert_substitution_spec, +) +from vllm.model_executor.model_loader.reload import finalize_layerwise_reload + +_T = TypeVar("_T", bound=nn.Module) + + +def _collect_expert_substitutions( + model: nn.Module, +) -> dict[int, tuple[str, ConstantExpertSubstitution]]: + substitutions: dict[int, tuple[str, ConstantExpertSubstitution]] = {} + for name, module in model.named_modules(): + if not isinstance(module, ConstantExpertSubstitution): + continue + if module.layer_idx in substitutions: + raise ValueError( + f"layer {module.layer_idx} has more than one MoE with substituted " + "experts" + ) + substitutions[module.layer_idx] = (f"{name}.values", module) + return substitutions + + +def _parse_approx_value_name(name: str) -> tuple[int, int] | None: + """Return (layer index, expert ID) of an ``approx_value`` tensor name.""" + if not name.endswith(APPROX_VALUE_SUFFIX): + return None + experts_path, _, expert_id = name.removesuffix(APPROX_VALUE_SUFFIX).rpartition(".") + layer_idx = decoder_layer_index(experts_path) + if layer_idx is None or not expert_id.isdigit(): + return None + return layer_idx, int(expert_id) + + +def as_expert_substitution_model(model_cls: type[_T], model_config: Any) -> type[_T]: + """Load constant expert values at the ``load_weights`` boundary. + + Initial loading and direct weight updates both go through ``load_weights``. + Explicitly named values are always consumed and loaded if their layer is + built by this model. ``approx_value`` tensors of layers built elsewhere + (other pipeline stages, MTP layers) are passed through, so the model's own + loader skips them like any other weight of a layer it does not own. + """ + spec = get_expert_substitution_spec(model_config) + if spec is None: + return model_cls + value_names = spec.value_names + + class ExpertSubstitutionModel(model_cls): # type: ignore[valid-type, misc] + @wraps(model_cls.__init__) # type: ignore[misc] + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._expert_substitutions = _collect_expert_substitutions(self) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + if any(s.values.is_meta for _, s in self._expert_substitutions.values()): + # Put the previous weights back so the model stays usable. + finalize_layerwise_reload(self, model_config) + raise NotImplementedError( + "expert substitution does not support layerwise weight reload" + ) + loaded_substitutions: set[str] = set() + + def remaining_weights(): + for name, loaded_weight in weights: + targets = value_names.get(name) + if targets is None: + key = _parse_approx_value_name(name) + if key is None or key[0] not in self._expert_substitutions: + yield name, loaded_weight + continue + targets = (key,) + for layer_idx, expert_id in targets: + if layer_idx not in self._expert_substitutions: + continue + param_name, substitution = self._expert_substitutions[layer_idx] + param = substitution.values + param.weight_loader(param, loaded_weight, expert_id) + loaded_substitutions.add(param_name) + + loaded = super().load_weights(remaining_weights()) + if loaded is not None: + loaded.update(loaded_substitutions) + return loaded + + def process_weights_after_loading(self) -> None: + for param_name, substitution in self._expert_substitutions.values(): + substitution.validate_loaded_values(param_name) + finalize = getattr(super(), "process_weights_after_loading", None) + if finalize is not None: + finalize() + + ExpertSubstitutionModel.__name__ = model_cls.__name__ + ExpertSubstitutionModel.__qualname__ = model_cls.__qualname__ + return ExpertSubstitutionModel From 98f721116227f0f48690deb05c815c057bac0668 Mon Sep 17 00:00:00 2001 From: Venky <23023424+venkywonka@users.noreply.github.com> Date: Mon, 12 Oct 2026 17:37:56 -0700 Subject: [PATCH 343/624] [Attention][Spec Decode] Multi-query Triton split-K for spec-decode verify, with empty-segment/padding hardening (#56148) Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com> Signed-off-by: Venky <23023424+venkywonka@users.noreply.github.com> Co-authored-by: OpenAI GPT-6 Astra Co-authored-by: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 (1M context) --- .../test_triton_unified_attention.py | 281 +++++++++++++++++- .../v1/attention/test_triton_attn_segments.py | 158 ++++++++++ vllm/v1/attention/backends/triton_attn.py | 32 +- .../attention/ops/triton_attention_helpers.py | 10 +- .../attention/ops/triton_unified_attention.py | 44 ++- 5 files changed, 500 insertions(+), 25 deletions(-) create mode 100644 tests/v1/attention/test_triton_attn_segments.py diff --git a/tests/kernels/attention/test_triton_unified_attention.py b/tests/kernels/attention/test_triton_unified_attention.py index d936e377040c..b50e97b52bc9 100644 --- a/tests/kernels/attention/test_triton_unified_attention.py +++ b/tests/kernels/attention/test_triton_unified_attention.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import MagicMock + import pytest import torch @@ -9,11 +11,17 @@ from vllm.triton_utils import tl, triton from vllm.utils.math_utils import next_power_of_2 from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.ops import triton_unified_attention as attention from vllm.v1.attention.ops.triton_attention_helpers import ( apply_softcap, compute_tile_loop_bounds, + softmax_step, +) +from vllm.v1.attention.ops.triton_unified_attention import ( + MAX_3D_QUERY_LEN, + reduce_segments, + unified_attention, ) -from vllm.v1.attention.ops.triton_unified_attention import unified_attention from vllm.v1.kv_cache_interface import KVQuantMode pytestmark = pytest.mark.skip_global_cleanup @@ -37,6 +45,259 @@ SEQ_THRESHOLD_3D_VALUES = [0, 8] +def _scratch(rows): + shapes = [(rows, 8, 16, 128), (rows, 8, 16), (rows, 8, 16)] + return [torch.empty(shape, dtype=torch.float32, device="cpu") for shape in shapes] + + +def _uses_splitk( + monkeypatch, + query_len, + scratch, + *, + causal=True, + threshold=8, + mm_prefix=None, + num_seqs=2, +): + query = torch.empty( + (num_seqs * query_len, 8, 128), dtype=torch.bfloat16, device="cpu" + ) + cache = torch.empty((1, 128, 1, 128), dtype=torch.bfloat16, device="cpu") + kernel = MagicMock() + monkeypatch.setattr(attention, "kernel_unified_attention", kernel) + monkeypatch.setattr(attention, "reduce_segments", MagicMock()) + unified_attention( + query, + cache, + cache, + torch.empty_like(query), + torch.arange(num_seqs + 1, dtype=torch.int32, device="cpu") * query_len, + query_len, + torch.full((num_seqs,), 128, dtype=torch.int32, device="cpu"), + 128, + 128**-0.5, + causal, + (-1, -1), + torch.zeros((num_seqs, 1), dtype=torch.int32, device="cpu"), + 0.0, + None, + None, + None, + seq_threshold_3D=threshold, + num_par_softmax_segments=16, + softmax_segm_output=scratch[0], + softmax_segm_max=scratch[1], + softmax_segm_expsum=scratch[2], + mm_prefix_range=mm_prefix, + ) + return kernel.__getitem__.return_value.call_args.kwargs["IS_3D"] + + +def _spy_reduce_segments(monkeypatch): + original = attention.reduce_segments + grids = [] + + def launch(grid): + grids.append(grid) + return original[grid] + + spy = MagicMock() + spy.__getitem__.side_effect = launch + monkeypatch.setattr(attention, "reduce_segments", spy) + return grids + + +@pytest.mark.parametrize( + "query_len,options,expected", + [ + (1, {}, True), + (1, {"causal": False}, True), + (5, {}, True), + (6, {"num_seqs": 1}, True), + (8, {}, True), + (9, {"num_seqs": 1}, False), + (5, {"causal": False}, False), + (5, {"causal": "per_sequence"}, False), + (5, {"threshold": 1}, False), + (5, {"batch_invariant": True}, False), + (5, {"mm_prefix": True}, False), + # BLOCK_Q is 2 for 8 query heads per KV head: 4 * ceil(4 / 2) <= 8. + (4, {"num_seqs": 4}, True), + (5, {"num_seqs": 4}, False), + (5, {"scratch_rows": 9}, False), + ], +) +def test_splitk_admission_preserves_decode_and_bounds_verification( + monkeypatch, query_len, options, expected +): + options = options.copy() + monkeypatch.setattr( + attention, "is_batch_invariant", options.pop("batch_invariant", False) + ) + num_seqs = options.setdefault("num_seqs", 2) + if options.get("causal") == "per_sequence": + options["causal"] = torch.ones(num_seqs, dtype=torch.bool, device="cpu") + if options.get("mm_prefix"): + options["mm_prefix"] = torch.zeros( + (num_seqs, 1, 2), dtype=torch.int32, device="cpu" + ) + scratch = _scratch(options.pop("scratch_rows", num_seqs * query_len)) + assert _uses_splitk(monkeypatch, query_len, scratch, **options) is expected + + +@pytest.mark.parametrize("empty_state", ["masked", "unvisited", "all_empty"]) +def test_reduce_segments_ignores_empty_softmax_states(empty_state: str) -> None: + """Empty segments must not suppress strongly negative, nonempty segments.""" + output = torch.full((1, 1, 32), float("nan"), device=DEVICE_TYPE) + partial = torch.zeros((1, 1, 16, 32), device=DEVICE_TYPE) + maxima = torch.full((1, 1, 16), -float("inf"), device=DEVICE_TYPE) + masses = torch.ones_like(maxima) + if empty_state != "all_empty": + partial[:, :, 0] = 3.0 + maxima[:, :, 0] = -1000.0 + if empty_state in ("masked", "all_empty"): + maxima[:, :, 1:] = 0.0 + masses[:, :, 1:] = 0.0 + reduce_segments[(1, 1)]( + output, + partial, + maxima, + masses, + torch.tensor([32], dtype=torch.int32, device=DEVICE_TYPE), + 1, + 1, + 1.0, + output.stride(0), + output.stride(1), + 1, + 16, + 32, + 32, + torch.tensor([0, 1], dtype=torch.int32, device=DEVICE_TYPE), + 16, + 16, + False, + ) + torch.testing.assert_close( + output, torch.full_like(output, 0.0 if empty_state == "all_empty" else 3.0) + ) + + +@triton.jit +def _softmax_after_empty_tile(output_ptr): + maximum = tl.full((1,), float("-inf"), tl.float32) + mass = tl.full((1,), 1.0, tl.float32) + empty_scores = tl.full((1, 16), float("-inf"), tl.float32) + maximum, mass, _, _ = softmax_step(empty_scores, maximum, mass) + scores = tl.full((1, 16), -1000.0, tl.float32) + _, mass, probabilities, _ = softmax_step(scores, maximum, mass) + tl.store(output_ptr + tl.arange(0, 16), tl.sum(probabilities, axis=0) / mass) + + +def test_splitk_softmax_recovers_after_fully_masked_tile() -> None: + """A masked first tile must not suppress later, strongly negative logits.""" + probabilities = torch.empty(16, device=DEVICE_TYPE) + _softmax_after_empty_tile[(1,)](probabilities) + torch.testing.assert_close(probabilities, torch.full_like(probabilities, 1 / 16)) + + +@pytest.mark.parametrize( + "query_lens,seq_lens", [([5, 0], [5, 0]), ([0], [0]), ([5], [0]), ([5], [5])] +) +def test_reduce_segments_leaves_graph_padding_untouched(query_lens, seq_lens) -> None: + """A padded reducer row must return before reading sequence or scratch data.""" + output = torch.full((10, 1, 32), 123.0, device=DEVICE_TYPE) + partial = torch.zeros((10, 1, 16, 32), device=DEVICE_TYPE) + maxima = torch.zeros((10, 1, 16), device=DEVICE_TYPE) + masses = torch.ones_like(maxima) + query_starts = torch.tensor( + [0] + query_lens, dtype=torch.int32, device=DEVICE_TYPE + ).cumsum(0, dtype=torch.int32) + reduce_segments[(10, 1)]( + output, + partial, + maxima, + masses, + torch.tensor(seq_lens, dtype=torch.int32, device=DEVICE_TYPE), + len(seq_lens), + 1, + 1.0, + output.stride(0), + output.stride(1), + 1, + 16, + 32, + 32, + query_starts, + 16, + 16, + False, + ) + active_rows = sum(q for q, kv in zip(query_lens, seq_lens) if kv > 0) + torch.testing.assert_close( + output[:active_rows], torch.zeros_like(output[:active_rows]) + ) + torch.testing.assert_close( + output[active_rows:], torch.full_like(output[active_rows:], 123.0) + ) + + +@pytest.mark.parametrize( + ("seq_lens", "num_heads", "q_dtype"), + [ + *[ + pytest.param([(w, 128), (w, 129), (w, 257)], (8, 1), None, id=f"w{w}") + for w in range(2, MAX_3D_QUERY_LEN + 1, 2) + ], + # Prefill chunk, verification, decode, padding, short query. + pytest.param( + [(8, 8), (6, 129), (1, 257), (0, 0), (3, 64)], (8, 1), None, id="mixed" + ), + *[ + pytest.param( + [(6, 128), (6, 129), (6, 257)], + heads, + FP8_DTYPE, + id=f"w6-fp8q-kv{heads[1]}", + ) + for heads in [(8, 1), (8, 8)] + ], + ], +) +@pytest.mark.parametrize("sliding_window", [None, 32]) +def test_speculative_splitk_matches_reference( + monkeypatch, seq_lens, num_heads, q_dtype, sliding_window +): + reduce_grids = _spy_reduce_segments(monkeypatch) + test_triton_unified_attn( + seq_lens=seq_lens, + num_heads=num_heads, + head_size=128, + sliding_window=sliding_window, + dtype=torch.bfloat16, + block_size=128, + soft_cap=None, + num_blocks=8, + q_dtype=q_dtype, + seq_threshold_3D=32, + ) + assert len(reduce_grids) == 1 + + +def test_speculative_splitk_bf16_query_fp8_kv(monkeypatch): + reduce_grids = _spy_reduce_segments(monkeypatch) + test_triton_unified_attn_bf16_query_fp8_kv( + seq_lens=[(6, 128), (6, 129), (6, 4097)], + num_heads=(8, 1), + head_size=128, + block_size=128, + num_blocks=64, + seq_threshold_3D=16, + ) + assert len(reduce_grids) == 1 + + @triton.jit def _compute_clamped_mm_tile_bounds( output_ptr, @@ -471,16 +732,19 @@ def test_triton_unified_attn( num_par_softmax_segments = 16 head_size_padded = next_power_of_2(head_size) + scratch_rows = seq_threshold_3D * ( + max_query_len if max_query_len <= MAX_3D_QUERY_LEN else 1 + ) softmax_segm_output = torch.empty( - (seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded), + (scratch_rows, num_query_heads, num_par_softmax_segments, head_size_padded), dtype=torch.float32, ) softmax_segm_max = torch.empty( - (seq_threshold_3D, num_query_heads, num_par_softmax_segments), + (scratch_rows, num_query_heads, num_par_softmax_segments), dtype=torch.float32, ) softmax_segm_expsum = torch.empty( - (seq_threshold_3D, num_query_heads, num_par_softmax_segments), + (scratch_rows, num_query_heads, num_par_softmax_segments), dtype=torch.float32, ) @@ -591,16 +855,19 @@ def test_triton_unified_attn_bf16_query_fp8_kv( num_par_softmax_segments = 16 head_size_padded = next_power_of_2(head_size) + scratch_rows = seq_threshold_3D * ( + max_query_len if max_query_len <= MAX_3D_QUERY_LEN else 1 + ) softmax_segm_output = torch.empty( - (seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded), + (scratch_rows, num_query_heads, num_par_softmax_segments, head_size_padded), dtype=torch.float32, ) softmax_segm_max = torch.empty( - (seq_threshold_3D, num_query_heads, num_par_softmax_segments), + (scratch_rows, num_query_heads, num_par_softmax_segments), dtype=torch.float32, ) softmax_segm_expsum = torch.empty( - (seq_threshold_3D, num_query_heads, num_par_softmax_segments), + (scratch_rows, num_query_heads, num_par_softmax_segments), dtype=torch.float32, ) diff --git a/tests/v1/attention/test_triton_attn_segments.py b/tests/v1/attention/test_triton_attn_segments.py new file mode 100644 index 000000000000..ae764a14acce --- /dev/null +++ b/tests/v1/attention/test_triton_attn_segments.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata +from vllm.config import CUDAGraphMode +from vllm.platforms import current_platform +from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadataBuilder +from vllm.v1.attention.backends.triton_attn_diffkv import ( + TritonAttentionDiffKVMetadataBuilder, +) +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVQuantMode + + +def _build( + monkeypatch, + builder_cls, + capability, + num_kv_heads, + query_lens, + max_num_seqs=16, + num_speculative_tokens=None, + parallel_drafting=False, + kv_quant_mode=KVQuantMode.NONE, +): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + current_platform, "is_device_capability", lambda cap: cap == capability + ) + monkeypatch.setattr(current_platform, "num_compute_units", lambda: 188) + config = SimpleNamespace( + model_config=SimpleNamespace( + get_num_attention_heads=lambda _: 8, + get_num_kv_heads=lambda _: num_kv_heads, + get_head_size=lambda: 128, + rswa_window=None, + ), + parallel_config=None, + scheduler_config=SimpleNamespace(max_num_seqs=max_num_seqs), + speculative_config=( + None + if num_speculative_tokens is None + else SimpleNamespace( + num_speculative_tokens=num_speculative_tokens, + parallel_drafting=parallel_drafting, + ) + ), + compilation_config=SimpleNamespace( + cudagraph_mode=CUDAGraphMode.NONE, static_forward_context={} + ), + ) + spec = FullAttentionSpec( + block_size=16, + num_kv_heads=num_kv_heads, + head_size=128, + dtype=torch.bfloat16, + kv_quant_mode=kv_quant_mode, + ) + builder = builder_cls(spec, ["layer.0"], config, "cpu") + batch = BatchSpec(seq_lens=[128] * len(query_lens), query_lens=query_lens) + metadata = builder.build(0, create_common_attn_metadata(batch, 16, "cpu")) + assert metadata.softmax_segm_output.shape[2] == metadata.num_par_softmax_segments + assert metadata.softmax_segm_max.shape[0] >= sum(query_lens) + # DiffKV re-allocates the output; it must stay sized like the base buffers. + assert builder.softmax_segm_output.shape[0] == builder.softmax_segm_max.shape[0] + assert ( + metadata.softmax_segm_output.data_ptr() + == builder.softmax_segm_output.data_ptr() + ) + return builder, metadata + + +BUILDERS = [TritonAttentionMetadataBuilder, TritonAttentionDiffKVMetadataBuilder] + + +# Verify batches: the 64-segment scratch holds min(11, max_num_seqs) * (1 + K) +# rows, and its 16-segment view covers every 3D-eligible sequence. 64 segments +# follow the request count, not Q blocks: on SM12.0 they stay faster for verify +# batches up to 11 requests (22 Q blocks at K=3). +@pytest.mark.parametrize("builder_cls", BUILDERS) +@pytest.mark.parametrize( + "max_num_seqs,num_speculative_tokens,num_seqs,rows,segments", + [ + (16, 3, 10, 44, 64), + (16, 3, 11, 44, 64), + (16, 3, 12, 44, 16), + (16, 3, 16, 44, 16), + (1, 5, 1, 6, 64), + ], +) +def test_split_k_segment_rows_cover_verify_queries( + monkeypatch, + builder_cls, + max_num_seqs, + num_speculative_tokens, + num_seqs, + rows, + segments, +): + builder, metadata = _build( + monkeypatch, + builder_cls, + (12, 0), + 1, + [1 + num_speculative_tokens] * num_seqs, + max_num_seqs, + num_speculative_tokens, + ) + assert builder.softmax_segm_max.shape[:2] == (rows, 8) + assert metadata.num_par_softmax_segments == segments + + +# Without SM12.0 64-segment reuse, the scratch holds exactly +# min(seq_threshold_3D, max_num_seqs) * query_len rows for split-K query_lens. +@pytest.mark.parametrize( + "num_speculative_tokens,parallel_drafting,max_num_seqs,rows", + [ + (None, False, 8, 8), + (4, False, 8, 40), + (7, False, 8, 64), + (8, False, 8, 64), + (3, True, 8, 56), + (4, True, 8, 64), + (5, False, 1, 6), + ], +) +@pytest.mark.parametrize( + "kv_quant_mode", [KVQuantMode.NONE, KVQuantMode.INT4_PER_TOKEN_HEAD] +) +def test_split_k_scratch_rows_follow_speculative_config( + monkeypatch, + num_speculative_tokens, + parallel_drafting, + max_num_seqs, + rows, + kv_quant_mode, +): + # The INT4 launcher never runs multi-query split-K. + if kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD: + rows = max_num_seqs + builder, _ = _build( + monkeypatch, + TritonAttentionMetadataBuilder, + (9, 0), + 1, + [1], + max_num_seqs, + num_speculative_tokens, + parallel_drafting, + kv_quant_mode, + ) + assert builder.softmax_segm_output.shape == (rows, 8, 16, 128) + assert builder.softmax_segm_max.shape == (rows, 8, 16) + assert builder.softmax_segm_expsum.shape == (rows, 8, 16) diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index a95b0c1d0e13..8b26bab84fa5 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -39,7 +39,10 @@ triton_reshape_and_cache_flash, triton_reshape_and_cache_flash_per_token_head_quant, ) -from vllm.v1.attention.ops.triton_unified_attention import unified_attention +from vllm.v1.attention.ops.triton_unified_attention import ( + MAX_3D_QUERY_LEN, + unified_attention, +) from vllm.v1.kv_cache_interface import ( AttentionSpec, KVQuantMode, @@ -144,7 +147,10 @@ def __init__( ) self.num_par_softmax_segments = NUM_PAR_SOFTMAX_SEGMENTS - # On SM120, batches whose 16-segment grid under-fills the SMs use 64. + # On SM120, batches whose 16-segment decode grid under-fills the SMs use + # 64. This counts requests, not Q blocks, on purpose: for 1..11 verify + # requests (K=3/5, 8 q / 1 KV head, 2k-128k context) 64 segments were + # faster or within 6% at 2k, and 1.1-3.6x faster from 8k up. self.max_seqs_64_segments = 0 if current_platform.is_cuda() and current_platform.is_device_capability( (12, 0) @@ -152,18 +158,34 @@ def __init__( self.max_seqs_64_segments = (current_platform.num_compute_units() - 1) // ( NUM_PAR_SOFTMAX_SEGMENTS * self.num_heads_kv ) - max_num_tokens_3d = self.seq_threshold_3D + headdim_padded = next_power_of_2(self.headdim) + max_query_len_3d = 1 + speculative_config = vllm_config.speculative_config + # The INT4 launcher runs multi-query batches in 2D only. + if ( + speculative_config is not None + and kv_cache_spec.kv_quant_mode != KVQuantMode.INT4_PER_TOKEN_HEAD + ): + num_speculative_tokens = speculative_config.num_speculative_tokens or 0 + query_len = 1 + num_speculative_tokens * ( + 2 if speculative_config.parallel_drafting else 1 + ) + max_query_len_3d = min(query_len, MAX_3D_QUERY_LEN) + # Scratch is indexed by query token, including verification tokens. + max_num_seqs_3d = min( + self.seq_threshold_3D, vllm_config.scheduler_config.max_num_seqs + ) + max_num_tokens_3d = max_num_seqs_3d * max_query_len_3d if self.max_seqs_64_segments > 0: self.num_par_softmax_segments = 64 # build() reuses this scratch at 16 segments with proportionally more rows. max_num_tokens_3d = max( - min(self.max_seqs_64_segments, max_num_tokens_3d), + min(self.max_seqs_64_segments, max_num_seqs_3d) * max_query_len_3d, cdiv( max_num_tokens_3d, self.num_par_softmax_segments // NUM_PAR_SOFTMAX_SEGMENTS, ), ) - headdim_padded = next_power_of_2(self.headdim) self.softmax_segm_output = torch.empty( ( max_num_tokens_3d, diff --git a/vllm/v1/attention/ops/triton_attention_helpers.py b/vllm/v1/attention/ops/triton_attention_helpers.py index 7804b013171e..f2399b742380 100644 --- a/vllm/v1/attention/ops/triton_attention_helpers.py +++ b/vllm/v1/attention/ops/triton_attention_helpers.py @@ -470,15 +470,15 @@ def softmax_step(S, M, L): # compute running maximum # m_j : (BLOCK_M,) m_j = tl.maximum(M, tl.max(S, axis=1)) - # For sliding window there's a chance the max is -inf due to masking of - # the entire row. In this case we need to set m_j 0 to avoid NaN - m_j = tl.where(m_j > float("-inf"), m_j, 0.0) + # Keep an empty row's maximum at -inf for later tiles; use a finite shift + # only for the exponent arithmetic below. + safe_m_j = tl.where(m_j > float("-inf"), m_j, 0.0) # P : (BLOCK_M, TILE_SIZE) - P = tl.exp(S - m_j[:, None]) + P = tl.exp(S - safe_m_j[:, None]) # l_j : (BLOCK_M,) l_j = tl.sum(P, axis=1) # alpha : (BLOCK_M, ) - alpha = tl.exp(M - m_j) + alpha = tl.exp(M - safe_m_j) # update constants L_new = L * alpha + l_j return m_j, L_new, P, alpha diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index 63a9cdc7f956..d914e574f2cc 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -33,6 +33,8 @@ logger = init_logger(__name__) is_batch_invariant = envs.VLLM_BATCH_INVARIANT float8_info = torch.finfo(current_platform.fp8_dtype()) +# Longest per-request query (e.g. 1 + num_speculative_tokens) run by split-K. +MAX_3D_QUERY_LEN = 8 @triton.jit @@ -721,12 +723,18 @@ def reduce_segments( query_token_idx = tl.program_id(0) query_head_idx = tl.program_id(1) + # Graph padding must be rejected before sequence lookup or scratch access. + if query_token_idx >= tl.load(query_start_len_ptr + num_seqs): + return + seq_idx = find_seq_idx( query_start_len_ptr, query_token_idx, num_seqs, BLOCK_Q, False ) # sequence len for this particular sequence seq_len = tl.load(seq_lens_ptr + seq_idx) + if seq_len <= 0: + return # number of segments for this particular sequence num_segments = NUM_SEGMENTS_PER_SEQ @@ -746,11 +754,16 @@ def reduce_segments( + tl.arange(0, NUM_SEGMENTS_PER_SEQ) ) segm_max = tl.load(segm_max_ptr + segm_offset, mask=segm_mask, other=float("-inf")) + segm_expsum = tl.load(segm_expsum_ptr + segm_offset, mask=segm_mask, other=0.0) + # Empty segments must not set the maximum. Unvisited segments retain + # the M=-inf,L=1 initial state. + valid_segment = segm_mask & (segm_expsum > 0) & (segm_max != float("-inf")) + segm_max = tl.where(valid_segment, segm_max, float("-inf")) overall_max = tl.max(segm_max) # load and rescale segment exp sums - segm_expsum = tl.load(segm_expsum_ptr + segm_offset, mask=segm_mask, other=0.0) - segm_expsum = segm_expsum * tl.exp(segm_max - overall_max) + segment_scale = tl.where(valid_segment, tl.exp(segm_max - overall_max), 0.0) + segm_expsum = segm_expsum * segment_scale overall_expsum = tl.sum(segm_expsum) # load, rescale, and add segment attention outputs @@ -766,7 +779,7 @@ def reduce_segments( mask=segm_mask[:, None] & dim_mask[None, :], other=0.0, ) - segm_output *= tl.exp(segm_max - overall_max)[:, None] + segm_output *= segment_scale[:, None] acc_sum = tl.sum(segm_output, axis=0) # safely divide by overall_expsum, returning 0.0 if overall_expsum is 0 acc = tl.where(overall_expsum == 0.0, 0.0, acc_sum / overall_expsum) @@ -1047,9 +1060,12 @@ def unified_attention( ) # Launch the 2D kernel if - # 1. No intermediate tiled softmax buffers for the 3D kernel have been allocated, or - # 2. The batch includes at least one prefill request, or - # 3. The number of sequences exceeds the configured threshold, or + # 1. No split-K threshold or scratch buffers, or the scratch has fewer + # rows than query tokens, or + # 2. max_seqlen_q > 1 and either it exceeds MAX_3D_QUERY_LEN, masking is + # non-causal, per-sequence causal, or mm-prefix, or the tuned + # large-head path applies, or + # 3. The number of Q blocks exceeds seq_threshold_3D, or # 4. Batch invariance is enabled use_3d = not ( seq_threshold_3D is None @@ -1057,10 +1073,22 @@ def unified_attention( or softmax_segm_output is None or softmax_segm_max is None or softmax_segm_expsum is None - or max_seqlen_q > 1 - or num_seqs > seq_threshold_3D + or ( + max_seqlen_q > 1 + and ( + max_seqlen_q > MAX_3D_QUERY_LEN + or not use_causal + or use_per_seq_causal + or use_mm_prefix + or tuned_large_head + ) + ) + or num_seqs * triton.cdiv(max_seqlen_q, BLOCK_Q) > seq_threshold_3D or is_batch_invariant ) + if use_3d and softmax_segm_max.shape[0] < q.shape[0]: + logger.debug_once("Split-K scratch rows < query tokens; using 2D.") + use_3d = False # The kernel signature is the same for 2D and 3D — only the launch # grid + a handful of constexpr toggles differ. Per-token-head scale From c6566bb9f678679c6e42ad2beb45fece8b35356a Mon Sep 17 00:00:00 2001 From: Hexiang Wang <56632993+whx-sjtu@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:20:38 +0800 Subject: [PATCH 344/624] [ROCm][MLA] Keep DCP prefill context in FP8 through AllGather (#59070) Signed-off-by: whx-sjtu Co-authored-by: OpenAI Codex Co-authored-by: Claude Code Co-authored-by: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> --- .buildkite/test-amd.yaml | 3 + .buildkite/test_areas/distributed.yaml | 3 + .../test_rocm_aiter_mla_op_registration.py | 20 + .../attention/test_rocm_aiter_mla_dcp_cprr.py | 7 + .../test_rocm_aiter_mla_mtp_split.py | 77 ++- .../v1/attention/test_rocm_mla_dcp_context.py | 445 ++++++++++++++++++ vllm/_aiter_ops.py | 52 ++ .../attention/backends/mla/rocm_aiter_mla.py | 152 +++++- .../attention/ops/rocm_aiter_mla_prefill.py | 105 +++++ 9 files changed, 847 insertions(+), 17 deletions(-) create mode 100644 tests/v1/attention/test_rocm_mla_dcp_context.py create mode 100644 vllm/v1/attention/ops/rocm_aiter_mla_prefill.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index e4ae7f897e35..d2702e7d4f82 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -3539,8 +3539,11 @@ steps: - examples/rl/rlhf_async_new_apis.py - tests/distributed/test_weight_transfer.py - tests/distributed/test_packed_tensor.py + - tests/v1/attention/test_rocm_mla_dcp_context.py + - vllm/v1/attention/ops/rocm_aiter_mla_prefill.py commands: - pytest -v -s tests/distributed/test_context_parallel.py + - pytest -v -s tests/v1/attention/test_rocm_mla_dcp_context.py - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index f541c42c93a8..b30cbac0ced4 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -528,8 +528,11 @@ steps: - examples/rl/rlhf_async_new_apis.py - tests/distributed/test_weight_transfer.py - tests/distributed/test_packed_tensor.py + - tests/v1/attention/test_rocm_mla_dcp_context.py + - vllm/v1/attention/ops/rocm_aiter_mla_prefill.py commands: - pytest -v -s tests/distributed/test_context_parallel.py + - pytest -v -s tests/v1/attention/test_rocm_mla_dcp_context.py - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization diff --git a/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py b/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py index 0944be8937c8..cb40c61c3b94 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py @@ -33,6 +33,26 @@ def _require_aiter(): pytest.skip("aiter is required on supported ROCm hardware for this test") +@torch.inference_mode() +def test_gather_kv_b_proj_op_schema() -> None: + """Validate fake tensors and mutation of both expanded K/V outputs.""" + _require_aiter() + from vllm._aiter_ops import rocm_aiter_ops # noqa: F401 + + rows, num_heads = 32, 12 + cache = torch.randn(rows, 1, 576, device="cuda").to(current_platform.fp8_dtype()) + scale = torch.tensor([0.07], device="cuda") + cu_seq_lens = torch.tensor([0, 13, rows], dtype=torch.int32, device="cuda") + indices = torch.arange(rows, dtype=torch.int32, device="cuda").flip(0) + weight = torch.randn(num_heads * 256, 512, dtype=torch.bfloat16, device="cuda") + k = torch.empty(rows, num_heads, 192, dtype=torch.bfloat16, device="cuda") + v = torch.empty(rows, num_heads, 128, dtype=torch.bfloat16, device="cuda") + opcheck( + torch.ops.vllm.rocm_aiter_gather_kv_b_proj, + (cache, scale, cu_seq_lens, indices, cu_seq_lens, weight, k, v), + ) + + @torch.inference_mode() def test_mla_decode_fwd_op_schema() -> None: """Opcheck validates registration, schema, fake-tensor, and ``o`` aliasing. diff --git a/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr.py b/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr.py index cd88ade8eebf..56886872ecc6 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr.py +++ b/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr.py @@ -142,6 +142,13 @@ def _build(monkeypatch, route, gathered_heads, segmented_supported=True): def init_common_builder(self, *args, **kwargs): self.num_heads = gathered_heads // DCP + self.q_data_type = torch.bfloat16 + self.mla_dims = SimpleNamespace( + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + ) self.dcp_world_size = DCP self.reorder_batch_threshold = 5 diff --git a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py index da141acf6881..1583995ac98b 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py +++ b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py @@ -740,26 +740,34 @@ def test_segmented_dcp_verify_matches_causal_attention(monkeypatch): torch.testing.assert_close(output, reference, rtol=3e-2, atol=3e-2) -@pytest.mark.parametrize("num_heads", [8, 16, 24, 32, 64, 128]) @pytest.mark.parametrize( - "spec_method, parallel_drafting", + "num_heads,spec_method,parallel_drafting,group_cache_dtype,needs_row_map", [ - ("deepseek_mtp", False), - # A drafter that is not one of the historically recognized MTP methods, - # and a parallel one, so the threshold is 1 + 2 * num_spec rather than - # 1 + num_spec. Sizing the metadata off a method name instead leaves - # these at qlen=1 while the router still admits the full range. - ("custom", True), - ("eagle", False), + *[ + (num_heads, method, parallel, None, True) + for num_heads in [8, 16, 24, 32, 64, 128] + for method, parallel in [ + ("deepseek_mtp", False), + # Parallel drafting needs qlen=1+2*num_spec, regardless of name. + ("custom", True), + ("eagle", False), + ] + ], + pytest.param(16, "deepseek_mtp", False, "auto", False, id="bf16-group"), + pytest.param(16, "deepseek_mtp", False, "fp8_e5m2", False, id="e5m2-group"), ], ) -def test_mtp_builder_init_sizes_native_fp8_metadata( - monkeypatch, num_heads, spec_method, parallel_drafting +def test_builder_init_sizes_decode_metadata_and_filters_dcp_maps( + monkeypatch, + num_heads, + spec_method, + parallel_drafting, + group_cache_dtype, + needs_row_map, ): - """Aiter init sizes the metadata for every query length decode can be handed. + """Size decode metadata for padded heads/qlen; skip unsupported DCP maps. - Sweeping num_heads asserts metadata is sized for the padded decode shape, - covering Kimi-K3 TP4's 24 -> 32 head path and native fp8 nhead=32 folding. + The group's raw cache format must override the global FP8 configuration. """ dtypes = SimpleNamespace(fp8="fp8", fp16="fp16", bf16="bf16") info_calls = [] @@ -789,6 +797,13 @@ def get_mla_metadata_info_v1( def init_common_builder(self, *args, **kwargs): self.num_heads = num_heads + self.q_data_type = config.model_config.dtype + self.mla_dims = SimpleNamespace( + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + ) self.dcp_world_size = 1 self.cp_kv_cache_interleave_size = 1 # Mirror what _init_reorder_batch_threshold would have left behind: the @@ -813,6 +828,7 @@ def init_common_builder(self, *args, **kwargs): init_common_builder, ) monkeypatch.setattr(rocm_aiter_mla, "_fp8_mla_prefill_supported", lambda: False) + monkeypatch.setattr(rocm_aiter_mla, "is_aiter_found_and_supported", lambda: True) config = SimpleNamespace( speculative_config=SimpleNamespace( @@ -840,7 +856,9 @@ def init_common_builder(self, *args, **kwargs): ), ) builder = AiterMLAMetadataBuilder( - kv_cache_spec=SimpleNamespace(block_size=1, dtype=torch.bfloat16), + kv_cache_spec=SimpleNamespace( + block_size=1, dtype=torch.bfloat16, cache_dtype_str=group_cache_dtype + ), layer_names=["layer.0"], vllm_config=config, device=torch.device("cpu"), @@ -860,6 +878,35 @@ def init_common_builder(self, *args, **kwargs): assert builder._mla_q_dtype == dtypes.fp8 assert builder._mla_kv_dtype == dtypes.fp8 + builder.dcp_world_size = 2 + chunk = object() + rows = torch.tensor([0], dtype=torch.int64) + build_rows = mock.Mock(return_value=rows) + monkeypatch.setattr(rocm_aiter_mla, "context_row_indices", build_rows) + monkeypatch.setattr( + rocm_aiter_mla, + "_supports_fused_dcp_prefill_config", + mock.Mock(side_effect=AssertionError("Support must be cached at init")), + ) + monkeypatch.setattr( + rocm_aiter_mla.MLACommonMetadataBuilder, + "build", + lambda *args: SimpleNamespace( + prefill=SimpleNamespace( + chunked_context=SimpleNamespace(chunks=[chunk]), + block_table=torch.empty(0), + ), + decode=None, + dcp_context_row_indices=None, + ), + ) + for _ in range(2): + metadata = builder.build(0, SimpleNamespace(causal=True)) + assert metadata.dcp_context_row_indices == ([rows] if needs_row_map else None) + assert build_rows.call_args_list == ( + [mock.call(chunk, torch.device("cpu"))] * 2 if needs_row_map else [] + ) + def test_mtp_decode_qlen4_keeps_uniform_rows_with_metadata(monkeypatch): get_mla_metadata_v1 = mock.MagicMock() diff --git a/tests/v1/attention/test_rocm_mla_dcp_context.py b/tests/v1/attention/test_rocm_mla_dcp_context.py new file mode 100644 index 000000000000..c29576f93ba9 --- /dev/null +++ b/tests/v1/attention/test_rocm_mla_dcp_context.py @@ -0,0 +1,445 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from tests.utils import multi_gpu_only +from vllm.platforms import current_platform +from vllm.v1.attention.ops.rocm_aiter_mla_prefill import context_row_indices + + +@pytest.mark.parametrize( + "padded,lengths,starts,total,expected", + [ + ([3, 2], [[7, 6], [2, 1]], [5, 0], 6, [0, 1, 5, 3, 4, 8]), + ([3], [[2, 9]], [4], 3, [3, 4, 5]), + ([3], [[1, 1]], [0], 3, None), + ], + ids=["padding_and_chunk_starts", "empty_rank", "inconsistent_total"], +) +def test_dcp_context_indices(padded, lengths, starts, total, expected): + chunk = SimpleNamespace( + padded_local_seq_lens=padded, + local_context_lens_allranks=lengths, + local_starts=starts, + num_local_context_tokens=sum(padded), + num_context_tokens=total, + ) + if expected is None: + with pytest.raises(AssertionError): + context_row_indices(chunk, torch.device("cpu")) + else: + rows = context_row_indices(chunk, torch.device("cpu")) + assert rows.tolist() == expected + assert rows.dtype == torch.int32 + + +def test_compressed_gather_preserves_bytes_and_workspace_partition(monkeypatch): + from vllm.v1.attention.ops import rocm_aiter_mla_prefill as prefill + + workspace = torch.full((12, 4), -1, dtype=torch.bfloat16) + cache = torch.arange(8, dtype=torch.uint8).reshape(2, 4) + chunk = SimpleNamespace( + num_local_context_tokens=2, + local_context_lens_allranks=[[2, 2]], + padded_local_cu_seq_lens=torch.tensor([0, 2]), + num_requests=1, + starts=torch.tensor([0]), + ) + + def local_gather(src, dst, *args): + assert src.dtype == dst.dtype == torch.uint8 + dst.copy_(src) + + def allgather(dst, src): + assert dst.dtype == src.dtype == torch.uint8 + assert dst.shape == (4, 4) + assert src.shape == (2, 4) + dst[:2].copy_(src) + dst[2:].copy_(src + 16) + + monkeypatch.setattr(prefill.ops, "cp_gather_cache", local_gather) + gathered = prefill.gather_compressed_context( + cache, workspace, torch.empty(1, 1), chunk, allgather, torch.float8_e4m3fn + ) + torch.testing.assert_close( + gathered.view(torch.uint8), torch.cat((cache, cache + 16)) + ) + # Retain the original logical capacity even though the byte view is larger. + assert gathered.data_ptr() - workspace.data_ptr() == 4 * 4 + torch.testing.assert_close(workspace.view(torch.uint8).reshape(-1, 4)[:2], cache) + + +@pytest.mark.parametrize( + "unsupported", + [ + "packed_cache", + "e5m2", + "bf16_cache", + "quantized_weight", + "bias", + "no_indices", + "query_dtype", + "prefill_dtype", + "weight_dtype", + "weight_layout", + "cache_layout", + "dimensions", + "no_aiter", + ], +) +def test_dcp_prefill_unsupported_formats_use_original_path(monkeypatch, unsupported): + from vllm.platforms import current_platform + + if not current_platform.is_rocm(): + pytest.skip("ROCm MLA backend") + from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( + AiterMLAImpl, + MLACommonImpl, + ) + + if unsupported == "no_aiter": + monkeypatch.setattr( + "vllm.v1.attention.backends.mla.rocm_aiter_mla." + "is_aiter_found_and_supported", + lambda: False, + ) + + impl = object.__new__(AiterMLAImpl) + impl.kv_cache_dtype = { + "packed_cache": "fp8_ds_mla", + "e5m2": "fp8_e5m2", + "bf16_cache": "auto", + }.get(unsupported, "fp8") + impl.kv_lora_rank = 512 + if unsupported == "dimensions": + impl.kv_lora_rank = 256 + impl.qk_nope_head_dim = impl.v_head_dim = 128 + impl.qk_rope_head_dim = 64 + impl.num_heads = 1 + impl.kv_b_proj = SimpleNamespace( + weight=torch.empty(256, 512, dtype=torch.bfloat16), + weight_scale=torch.ones(1) if unsupported == "quantized_weight" else None, + bias=torch.ones(256) if unsupported == "bias" else None, + ) + if unsupported == "weight_dtype": + impl.kv_b_proj.weight = impl.kv_b_proj.weight.half() + elif unsupported == "weight_layout": + impl.kv_b_proj.weight = torch.empty(512, 256, dtype=torch.bfloat16).T + metadata = SimpleNamespace( + prefill=SimpleNamespace( + q_data_type=torch.float16 + if unsupported == "prefill_dtype" + else torch.bfloat16 + ), + dcp_context_row_indices=None if unsupported == "no_indices" else [], + ) + expected = (torch.ones(1), torch.zeros(1)) + + def fused_mla_kv_concat(kv_nope, k_pe, use_fp8_prefill): + raise AssertionError("The fallback should receive the callback") + + def original_path(*args): + assert args[-1] is fused_mla_kv_concat + return expected + + monkeypatch.setattr( + MLACommonImpl, + "_context_parallel_compute_prefill_context", + original_path, + ) + q = torch.empty(1, 1, 192, dtype=torch.bfloat16) + cache = torch.empty(1, 1536, 576, dtype=torch.uint8) + if unsupported == "query_dtype": + q = q.half() + elif unsupported == "cache_layout": + cache = torch.empty(1, 1536, 1152, dtype=torch.uint8)[..., ::2] + assert ( + impl._context_parallel_compute_prefill_context( + q, + cache, + metadata, + torch.ones(1), + 8, + fused_mla_kv_concat_fn=fused_mla_kv_concat, + ) + is expected + ) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm MLA backend") +@multi_gpu_only(num_gpus=2) +@pytest.mark.parametrize("interleave", [1, 64]) +@pytest.mark.parametrize( + "context_lens,workspace_size", + [([145, 37], 1024), ([545, 0, 3, 0], 256)], + ids=["single_chunk", "continuations_and_empty_context"], +) +def test_dcp_prefill_matches_full_attention( + interleave, context_lens, workspace_size, tmp_path +): + """Real DCP gather, fused expansion and suffix/chunk merges match dense SDPA.""" + import torch.multiprocessing as mp + + from vllm._aiter_ops import is_aiter_found_and_supported + + if not is_aiter_found_and_supported(): + pytest.skip("AITER on CDNA 3 or newer is required") + mp.spawn( + _dcp_prefill_worker, + args=( + f"file://{tmp_path / 'dist_init'}", + interleave, + context_lens, + workspace_size, + ), + nprocs=2, + ) + + +@torch.inference_mode() +def _dcp_prefill_worker(rank, init_method, interleave, context_lens, workspace_size): + import torch.nn.functional as F + + from tests.v1.attention.test_mla_backends import MockMLAAttentionLayer + from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + create_vllm_config, + ) + from vllm.config import set_current_vllm_config + from vllm.distributed import ( + cleanup_dist_env_and_memory, + init_distributed_environment, + initialize_model_parallel, + ) + from vllm.model_executor.layers.linear import ColumnParallelLinear + from vllm.v1.attention.backends.mla.prefill.aiter_flash_attn import ( + AiterFlashAttnPrefillBackend, + ) + from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( + AiterMLAImpl, + AiterMLAMetadataBuilder, + ) + from vllm.v1.attention.ops.dcp import MLADCPManager + from vllm.v1.kv_cache_interface import MLAAttentionSpec + from vllm.v1.worker.workspace import init_workspace_manager + + torch.accelerator.set_device_index(rank) + torch.set_num_threads(1) + device = torch.device(f"cuda:{rank}") + init_workspace_manager(device) + world_size, num_heads, block_size = 2, 12, 64 + config = create_vllm_config( + model_name="deepseek-ai/DeepSeek-V2-Lite", + dtype=torch.bfloat16, + tensor_parallel_size=world_size, + max_model_len=2048, + block_size=block_size, + max_num_seqs=8, + max_num_batched_tokens=2048, + ) + config.cache_config.cache_dtype = "fp8" + config.parallel_config.decode_context_parallel_size = world_size + config.parallel_config.cp_kv_cache_interleave_size = interleave + config.attention_config.use_prefill_query_quantization = False + config.kernel_config.enable_jit_warmup = False + with set_current_vllm_config(config): + init_distributed_environment( + world_size=world_size, + rank=rank, + local_rank=rank, + distributed_init_method=init_method, + ) + initialize_model_parallel( + tensor_model_parallel_size=world_size, + decode_context_model_parallel_size=world_size, + ) + try: + projection = ColumnParallelLinear( + 512, + num_heads * world_size * 256, + bias=False, + params_dtype=torch.bfloat16, + ).to(device) + torch.manual_seed(101 + rank) + projection.weight.copy_(torch.randn_like(projection.weight) * 0.02) + dims = dict( + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + ) + impl = AiterMLAImpl( + num_heads=num_heads, + head_size=576, + scale=192**-0.5, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="fp8", + logits_soft_cap=None, + attn_type="decoder", + kv_sharing_target_layer_name=None, + q_lora_rank=None, + qk_head_dim=192, + kv_b_proj=projection, + **dims, + ) + layer = MockMLAAttentionLayer( + impl=impl, + num_heads=num_heads, + device=device, + kv_b_proj=projection, + q_scale=1.0, + k_scale=0.07, + **dims, + ) + layer.prefill_backend = AiterFlashAttnPrefillBackend( + num_heads=num_heads, + scale=impl.scale, + vllm_config=config, + **dims, + ) + layer.dcp_manager = MLADCPManager( + config, + device, + num_heads, + 576, + 512, + torch.bfloat16, + torch.bfloat16, + None, + True, + False, + ) + config.compilation_config.static_forward_context["mla"] = layer + spec = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=576, + dtype=current_platform.fp8_dtype(), + cache_dtype_str="fp8", + ) + builder = AiterMLAMetadataBuilder(spec, ["mla"], config, device) + builder.chunked_prefill_workspace_size = workspace_size + builder.chunked_prefill_workspace = torch.empty( + (workspace_size + workspace_size // world_size, 576), + dtype=torch.bfloat16, + device=device, + ) + layer.dcp_manager.init_kv_gather( + builder.chunked_prefill_workspace, workspace_size + ) + query_lens = [i + 2 for i in range(len(context_lens))] + batch = BatchSpec( + seq_lens=[c + q for c, q in zip(context_lens, query_lens)], + query_lens=query_lens, + ) + common = create_common_attn_metadata( + batch, block_size, device, arange_block_indices=True + ) + metadata = builder.build(0, common) + assert metadata.prefill is not None + chunked = metadata.prefill.chunked_context + assert chunked is not None and metadata.dcp_context_row_indices is not None + if 0 in context_lens: + assert chunked.empty_token_slices + assert any(c.is_continuation for c in chunked.chunks) + else: + assert len(chunked.chunks) == 1 + blocks_per_req = common.block_table_tensor.shape[1] + cache = torch.zeros( + (len(context_lens) * blocks_per_req, block_size, 576), + dtype=torch.bfloat16, + device=device, + ).to(current_platform.fp8_dtype()) + queries, suffixes, references = [], [], [] + # All ranks reconstruct the same global cache before sharding it. + torch.manual_seed(42) + for request, (context_len, query_len) in enumerate( + zip(context_lens, query_lens) + ): + latent = torch.randn( + context_len, 576, device=device, dtype=torch.bfloat16 + ) + quantized = (latent.float() / layer._k_scale).to(cache.dtype) + positions = torch.arange(context_len, device=device) + shard = quantized.view(torch.uint8)[ + (positions // interleave) % world_size == rank + ] + cache.view(len(context_lens), -1, 576).view(torch.uint8)[ + request, : len(shard) + ].copy_(shard) + suffix = torch.randn( + query_len, 576, device=device, dtype=torch.bfloat16 + ) + q = ( + torch.randn( + query_len, num_heads, 192, device=device, dtype=torch.bfloat16 + ) + * 0.5 + ) + full_latent = torch.cat( + ( + (quantized.float() * layer._k_scale).to(torch.bfloat16), + suffix, + ) + ) + kv = F.linear(full_latent[:, :512], projection.weight).view( + -1, num_heads, 256 + ) + k_nope, v = kv.split(128, dim=-1) + k = torch.cat( + (k_nope, full_latent[:, None, 512:].expand(-1, num_heads, -1)), + dim=-1, + ) + mask = torch.arange(context_len + query_len, device=device)[ + None, : + ] <= (context_len + torch.arange(query_len, device=device)[:, None]) + references.append( + F.scaled_dot_product_attention( + q.transpose(0, 1).float(), + k.transpose(0, 1).float(), + v.transpose(0, 1).float(), + attn_mask=mask, + scale=impl.scale, + ).transpose(0, 1) + ) + queries.append(q) + suffixes.append(suffix) + q, suffix = torch.cat(queries), torch.cat(suffixes) + output = torch.empty( + q.shape[0], num_heads * 128, device=device, dtype=torch.bfloat16 + ) + # Ensure this test cannot silently succeed through the fallback. + with pytest.MonkeyPatch.context() as patch: + + def unexpected_fallback(*args, **kwargs): + raise AssertionError("Expected the compressed DCP path") + + patch.setattr( + "vllm.model_executor.layers.attention.mla_attention." + "MLACommonImpl._context_parallel_compute_prefill_context", + unexpected_fallback, + ) + impl.forward_mha( + q, + suffix[:, :512], + suffix[:, None, 512:], + cache, + metadata, + layer._k_scale, + output, + ) + torch.testing.assert_close( + output.view(-1, num_heads, 128).float(), + torch.cat(references), + rtol=0.03, + atol=0.003, + ) + finally: + cleanup_dist_env_and_memory() diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 8af999906f28..0817c651a0f4 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -643,6 +643,36 @@ def _check_aiter_mla_fp8_support() -> bool: return _AITER_MLA_SUPPORTS_FP8 +def _rocm_aiter_gather_kv_b_proj_impl( + kv_cache: torch.Tensor, + kv_scale: torch.Tensor, + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, + cu_seq_lens: torch.Tensor, + weight: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, +) -> None: + from aiter.ops.triton.gather_kv_b_proj import gather_kv_b_proj + + gather_kv_b_proj( + kv_cache, kv_scale, kv_indptr, kv_indices, cu_seq_lens, weight, None, k, v + ) + + +def _rocm_aiter_gather_kv_b_proj_fake( + kv_cache: torch.Tensor, + kv_scale: torch.Tensor, + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, + cu_seq_lens: torch.Tensor, + weight: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, +) -> None: + pass + + def _rocm_aiter_mla_decode_fwd_impl( q: torch.Tensor, kv_buffer: torch.Tensor, @@ -2851,6 +2881,13 @@ def register_ops_once() -> None: dispatch_key=current_platform.dispatch_key, ) + direct_register_custom_op( + op_name="rocm_aiter_gather_kv_b_proj", + op_func=_rocm_aiter_gather_kv_b_proj_impl, + mutates_args=["k", "v"], + fake_impl=_rocm_aiter_gather_kv_b_proj_fake, + dispatch_key=current_platform.dispatch_key, + ) direct_register_custom_op( op_name="rocm_aiter_mla_decode_fwd", op_func=_rocm_aiter_mla_decode_fwd_impl, @@ -3465,6 +3502,21 @@ def triton_sparse_mla_fwd( extra_indices=extra_kv_indices, ) + @staticmethod + def gather_kv_b_proj( + kv_cache: torch.Tensor, + kv_scale: torch.Tensor, + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, + cu_seq_lens: torch.Tensor, + weight: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> None: + torch.ops.vllm.rocm_aiter_gather_kv_b_proj( + kv_cache, kv_scale, kv_indptr, kv_indices, cu_seq_lens, weight, k, v + ) + @staticmethod def mla_decode_fwd( q: torch.Tensor, diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index 1115da630b3c..d64b4797c031 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -2,15 +2,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools +from collections.abc import Callable from dataclasses import dataclass from enum import Enum from pathlib import Path -from typing import TYPE_CHECKING, ClassVar, Final +from typing import TYPE_CHECKING, ClassVar, Final, cast import torch from vllm import envs -from vllm._aiter_ops import rocm_aiter_ops +from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops from vllm.config import VllmConfig from vllm.config.cache import CacheDType from vllm.distributed import get_dcp_group @@ -22,6 +23,8 @@ MLACommonMetadata, MLACommonMetadataBuilder, QueryLenSupport, + accumulate_mla_context_chunk, + init_mla_context_partial, ) from vllm.triton_utils import tl, triton from vllm.utils.gpu_sync_debug import gpu_sync_allowed @@ -38,6 +41,11 @@ from vllm.v1.attention.ops.rocm_aiter_mla_merge import ( merge_mla_segments_triton, ) +from vllm.v1.attention.ops.rocm_aiter_mla_prefill import ( + context_row_indices, + expand_context, + gather_compressed_context, +) from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache if TYPE_CHECKING: @@ -46,6 +54,26 @@ logger = init_logger(__name__) +def _supports_fused_dcp_prefill_config( + kv_cache_dtype: str | None, + query_dtype: torch.dtype, + kv_lora_rank: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, +) -> bool: + """Whether the configuration supports fused DCP context expansion.""" + return ( + kv_cache_dtype in ("fp8", "fp8_e4m3") + and query_dtype == torch.bfloat16 + and kv_lora_rank == 512 + and qk_nope_head_dim == 128 + and qk_rope_head_dim == 64 + and v_head_dim == 128 + and is_aiter_found_and_supported() + ) + + def _segmented_mla_page_size(block_size: int) -> int: """Largest supported power-of-two subpage dividing a physical KV block. @@ -507,6 +535,7 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata): @dataclass class AiterMLAMetadata(MLACommonMetadata[AiterMLADecodeMetadata]): + dcp_context_row_indices: list[torch.Tensor] | None = None work_meta_data: torch.Tensor | None = None work_indptr: torch.Tensor | None = None work_info_set: torch.Tensor | None = None @@ -664,6 +693,19 @@ def __init__( device ).multi_processor_count + # Use the group's original format: decode normalizes E5M2 to "fp8". + kv_cache_dtype = getattr(kv_cache_spec, "cache_dtype_str", None) + if kv_cache_dtype is None: + kv_cache_dtype = vllm_config.cache_config.cache_dtype + self._supports_fused_dcp_prefill = _supports_fused_dcp_prefill_config( + kv_cache_dtype, + self.q_data_type, + self.mla_dims.kv_lora_rank, + self.mla_dims.qk_nope_head_dim, + self.mla_dims.qk_rope_head_dim, + self.mla_dims.v_head_dim, + ) + self.compilation_config = vllm_config.compilation_config self.decode_attn_out_dtype = vllm_config.model_config.dtype @@ -1577,6 +1619,17 @@ def build( attn_metadata = super().build( common_prefix_len, common_attn_metadata, fast_build ) + prefill = attn_metadata.prefill + if ( + prefill is not None + and prefill.chunked_context is not None + and self.dcp_world_size > 1 + and self._supports_fused_dcp_prefill + ): + attn_metadata.dcp_context_row_indices = [ + context_row_indices(chunk, prefill.block_table.device) + for chunk in prefill.chunked_context.chunks + ] if ( attn_metadata.decode is not None and attn_metadata.decode.has_persistent_metadata @@ -1975,6 +2028,101 @@ def __init__( self._mla_prefill_ps_asm_fwd = mla_prefill_ps_asm_fwd self._mla_reduce_v1 = mla_reduce_v1 + def _context_parallel_compute_prefill_context( + self, + q: torch.Tensor, + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: MLACommonMetadata, + k_scale: torch.Tensor, + dcp_world_size: int, + fused_mla_kv_concat_fn: Callable[ + [torch.Tensor, torch.Tensor, bool], tuple[torch.Tensor, torch.Tensor] + ] + | None = None, + ): + attn_metadata = cast(AiterMLAMetadata, attn_metadata) + assert attn_metadata.prefill is not None + prefill = attn_metadata.prefill + weight = getattr(self.kv_b_proj, "weight", None) + eligible = ( + _supports_fused_dcp_prefill_config( + self.kv_cache_dtype, + prefill.q_data_type, + self.kv_lora_rank, + self.qk_nope_head_dim, + self.qk_rope_head_dim, + self.v_head_dim, + ) + and q.dtype == torch.bfloat16 + and weight is not None + and weight.dtype == torch.bfloat16 + and weight.shape == (self.num_heads * 256, 512) + and weight.is_contiguous() + and getattr(self.kv_b_proj, "bias", None) is None + and getattr(self.kv_b_proj, "weight_scale", None) is None + and kv_c_and_k_pe_cache.is_contiguous() + and attn_metadata.dcp_context_row_indices is not None + ) + if not eligible: + return super()._context_parallel_compute_prefill_context( + q, + kv_c_and_k_pe_cache, + attn_metadata, + k_scale, + dcp_world_size, + fused_mla_kv_concat_fn, + ) + from vllm.platforms import current_platform + + assert weight is not None + assert attn_metadata.dcp_context_row_indices is not None + chunked = prefill.chunked_context + assert chunked is not None and chunked.dcp_manager is not None + assert prefill.prefill_backend is not None + assert len(chunked.chunks) == len(attn_metadata.dcp_context_row_indices) + output = output_lse = None + for chunk, row_indices in zip( + chunked.chunks, attn_metadata.dcp_context_row_indices + ): + gathered = gather_compressed_context( + kv_c_and_k_pe_cache, + chunked.workspace, + prefill.block_table[chunk.request_slice], + chunk, + chunked.dcp_manager.kv_gather, + current_platform.fp8_dtype(), + ) + k, v = expand_context( + gathered, + k_scale, + row_indices, + chunk.cu_seq_lens, + weight, + self.num_heads, + self.qk_nope_head_dim, + self.qk_rope_head_dim, + self.v_head_dim, + ) + attn_output, attn_lse = prefill.prefill_backend.run_prefill_context_chunk( + chunk=chunk, + q=q[chunk.token_slice], + k=k, + v=v, + ) + if output is None: + if len(chunked.chunks) == 1 and not chunked.empty_token_slices: + return attn_output, attn_lse + output, output_lse = init_mla_context_partial( + chunked, + attn_output, + attn_lse, + num_tokens=q.shape[0], + ) + accumulate_mla_context_chunk( + chunk, attn_output, attn_lse, output, output_lse + ) + return output, output_lse + def _flash_attn_varlen_diff_headdims( self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs ): diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_prefill.py b/vllm/v1/attention/ops/rocm_aiter_mla_prefill.py new file mode 100644 index 000000000000..c6d839f72453 --- /dev/null +++ b/vllm/v1/attention/ops/rocm_aiter_mla_prefill.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compressed DCP context preparation for ROCm MLA.""" + +from collections.abc import Callable + +import torch + +from vllm import _custom_ops as ops +from vllm._aiter_ops import rocm_aiter_ops +from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata + + +def context_row_indices( + chunk: MLACommonPrefillMetadata.ContextChunk, device: torch.device +) -> torch.Tensor: + """Map compact request-major rows to the padded rank-major AllGather.""" + assert chunk.padded_local_seq_lens is not None + assert chunk.local_context_lens_allranks is not None + assert chunk.local_starts is not None + segments = [] + offset = 0 + for padded, lengths, start in zip( + chunk.padded_local_seq_lens, + chunk.local_context_lens_allranks, + chunk.local_starts, + ): + for rank, length in enumerate(lengths): + count = min(max(0, length - start), padded) + if count: + first = rank * chunk.num_local_context_tokens + offset + segments.append(torch.arange(first, first + count, dtype=torch.int32)) + offset += padded + rows = torch.cat(segments) if segments else torch.empty(0, dtype=torch.int32) + assert rows.numel() == chunk.num_context_tokens + return rows.to(device=device, non_blocking=True) + + +def gather_compressed_context( + cache: torch.Tensor, + workspace: torch.Tensor, + block_table: torch.Tensor, + chunk: MLACommonPrefillMetadata.ContextChunk, + gather: Callable[[torch.Tensor, torch.Tensor], object], + cache_dtype: torch.dtype, +) -> torch.Tensor: + """Gather FP8 bytes, preserving the existing workspace row capacity.""" + assert chunk.local_context_lens_allranks is not None + assert workspace.is_contiguous() + rows, width = workspace.shape + fp8_workspace = workspace.view(cache_dtype).view(-1, width)[:rows] + local_rows = chunk.num_local_context_tokens + world_size = len(chunk.local_context_lens_allranks[0]) + offset = rows // (world_size + 1) + assert offset * (world_size + 1) == rows + assert local_rows <= offset + local = fp8_workspace[:local_rows] + gathered = fp8_workspace[offset : offset + world_size * local_rows] + assert chunk.padded_local_cu_seq_lens is not None + ops.cp_gather_cache( + cache.view(torch.uint8), + local.view(torch.uint8), + block_table, + chunk.padded_local_cu_seq_lens, + chunk.num_requests, + chunk.starts, + ) + gather(gathered.view(torch.uint8), local.view(torch.uint8)) + return gathered + + +def expand_context( + gathered: torch.Tensor, + scale: torch.Tensor, + row_indices: torch.Tensor, + cu_seq_lens: torch.Tensor, + weight: torch.Tensor, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fuse row reorganization, dequantization, projection and K/V packing.""" + k = torch.empty( + (row_indices.numel(), num_heads, qk_nope_head_dim + qk_rope_head_dim), + device=gathered.device, + dtype=weight.dtype, + ) + v = torch.empty( + (row_indices.numel(), num_heads, v_head_dim), + device=gathered.device, + dtype=weight.dtype, + ) + if row_indices.numel(): + rocm_aiter_ops.gather_kv_b_proj( + gathered.unsqueeze(1), + scale, + cu_seq_lens, + row_indices, + cu_seq_lens, + weight, + k, + v, + ) + return k, v From f47792af7fba76eb8e8f6a564b35cfad316d39f2 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Mon, 12 Oct 2026 18:33:26 -0700 Subject: [PATCH 345/624] [Model Loader] Make dummy weight init dtype-aware (#61134) Signed-off-by: Woosuk Kwon Co-authored-by: Claude Opus 5.5 --- tests/model_executor/test_weight_utils.py | 45 ++++++++++++++++++ .../model_loader/dummy_loader.py | 6 +-- .../model_loader/weight_utils.py | 47 ++++++++++++++++++- 3 files changed, 93 insertions(+), 5 deletions(-) diff --git a/tests/model_executor/test_weight_utils.py b/tests/model_executor/test_weight_utils.py index 512f518a0561..12a8e3acac29 100644 --- a/tests/model_executor/test_weight_utils.py +++ b/tests/model_executor/test_weight_utils.py @@ -6,13 +6,16 @@ import huggingface_hub.constants import pytest +import torch from huggingface_hub.utils import LocalEntryNotFoundError from vllm.model_executor.model_loader.weight_utils import ( download_weights_from_hf, drop_checkpoint_cache, + initialize_single_dummy_weight, maybe_remap_kv_scale_name, ) +from vllm.platforms import current_platform def test_download_weights_from_hf(): @@ -307,5 +310,47 @@ def test_drop_checkpoint_cache_releases_local_weight_files( fadvise.assert_called_once() +@pytest.mark.parametrize("dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) +def test_dummy_fp8_weight_is_not_degenerate(dtype): + """+-1e-3 underflows in FP8, so 8-bit floats default to +-1.""" + param = torch.empty(4096, dtype=dtype) + initialize_single_dummy_weight(param) + values = param.float() + assert values.isfinite().all() + assert values.abs().max() <= 1.0 + assert (values != 0).float().mean() > 0.99 + + +@pytest.mark.parametrize( + ("dtype", "name"), + [(torch.float8_e8m0fnu, "weight_scale_inv"), (torch.uint8, "w13_weight_scale")], +) +def test_dummy_e8m0_scale_is_fixed_exponent(dtype, name): + """E8M0 scales, including uint8-stored ones, get a finite 2**-7.""" + param = torch.empty(64, dtype=dtype) + initialize_single_dummy_weight(param, name=name) + assert (param.view(torch.uint8) == 120).all() + + +@pytest.mark.parametrize("dtype", [torch.uint8, torch.int8, torch.float4_e2m1fn_x2]) +def test_dummy_packed_weight_gets_random_bytes(dtype): + """Packed weights are filled deterministically instead of left uninitialized.""" + # Some backends can't fill packed FP4 directly, so view zeroed bytes. + params = [torch.zeros(4096, dtype=torch.uint8).view(dtype) for _ in range(2)] + for param in params: + initialize_single_dummy_weight(param, name="w13_weight") + raw = params[0].view(torch.uint8) + assert raw.unique().numel() > 250 + assert torch.equal(raw, params[1].view(torch.uint8)) + + +@pytest.mark.skipif(current_platform.is_rocm(), reason="ROCm zeroes integer params") +def test_dummy_wide_integer_param_is_untouched(): + """Wider integer params may be indices, so random values are unsafe.""" + param = torch.arange(16, dtype=torch.int32) + initialize_single_dummy_weight(param, name="g_idx") + assert torch.equal(param, torch.arange(16, dtype=torch.int32)) + + if __name__ == "__main__": test_download_weights_from_hf() diff --git a/vllm/model_executor/model_loader/dummy_loader.py b/vllm/model_executor/model_loader/dummy_loader.py index 0f263fe11709..ff732f6b222c 100644 --- a/vllm/model_executor/model_loader/dummy_loader.py +++ b/vllm/model_executor/model_loader/dummy_loader.py @@ -45,7 +45,7 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: # Exclude nonpersistent buffers, matching state_dict(). for name, tensor in get_layer_tensors(layer).items(): if name not in layer._non_persistent_buffers_set: - initialize_single_dummy_weight(tensor) + initialize_single_dummy_weight(tensor, name=name) def _process_online_quant_layer( self, @@ -55,8 +55,8 @@ def _process_online_quant_layer( """Materialize, apply dummy weights, and run quantization processing.""" materialize_layer(layer, info) - for tensor in get_layer_tensors(layer).values(): - initialize_single_dummy_weight(tensor) + for name, tensor in get_layer_tensors(layer).items(): + initialize_single_dummy_weight(tensor, name=name) for param in get_layer_tensors(layer).values(): param.weight_loader = _get_original_loader(param) diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index ac1da8d6987b..0a834cd90ca5 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -1521,13 +1521,30 @@ def composed_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: return composed_loader +# E8M0 byte 120 encodes 2**-7: with packed FP4 (|x| <= 6) or the FP8 dummy +# range below, the dequantized weights stay small and finite. +_DUMMY_E8M0_SCALE_BYTE = 120 +# Default dummy range of 8-bit floats. +-1e-3 is below FP8 E4M3's smallest +# normal (2**-6), so it would mostly round to zero or subnormals. +_DUMMY_FP8_RANGE = 1.0 + + @torch.no_grad() def initialize_single_dummy_weight( param: torch.Tensor, - low: float = -1e-3, - high: float = 1e-3, + low: float | None = None, + high: float | None = None, seed: int = 1234, + name: str = "", ) -> None: + """Fill `param` with deterministic dummy values suited to its dtype. + + 8-bit integer and packed FP4 params get random bytes (every FP4 code is + finite), except uint8 params named like a scale, which hold E8M0 scales + and get a fixed exponent. Other integer params (e.g. indices) are left + untouched. Float params are sampled from `[low, high]`, which defaults to + +-1e-3, or +-1 for 8-bit floats. + """ if param.device.type == "meta": return # deferred to finalize_layerwise_processing (e.g. online quant) @@ -1535,6 +1552,25 @@ def initialize_single_dummy_weight( param.fill_(dummy_weight_value) return + is_e8m0 = param.dtype == torch.float8_e8m0fnu or ( + param.dtype == torch.uint8 and "scale" in name + ) + if is_e8m0: + param.view(torch.uint8).fill_(_DUMMY_E8M0_SCALE_BYTE) + return + + if ( + param.dtype in (torch.uint8, torch.int8, torch.float4_e2m1fn_x2) + and not current_platform.is_tpu() + ): + generator = torch.Generator(device=param.device) + generator.manual_seed(seed) + if param.dtype == torch.int8: + param.random_(-128, 128, generator=generator) + else: + param.view(torch.uint8).random_(0, 256, generator=generator) + return + if not torch.is_floating_point(param): if current_platform.is_rocm(): # On ROCm, integer params (e.g. GPTQ qweight/qzeros) are left @@ -1543,6 +1579,13 @@ def initialize_single_dummy_weight( param.zero_() return + if low is None or high is None: + dummy_range = 1e-3 + if torch.finfo(param.dtype).bits < 16: + dummy_range = _DUMMY_FP8_RANGE + low = -dummy_range if low is None else low + high = dummy_range if high is None else high + if current_platform.is_tpu(): generator = torch.Generator(device="cpu") generator.manual_seed(seed) From cc2a6eea1a9f92c1d6b47be8fe48adf127d4c67f Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Mon, 12 Oct 2026 22:16:50 -0400 Subject: [PATCH 346/624] [ROCm][Perf][DSv4.1] Mono decode on every backbone layer (#61082) Signed-off-by: Fangzhou Ai Co-authored-by: Claude Opus 5.5 --- tests/models/test_dsv41_mono_decode.py | 83 ++++++++++--- tests/models/test_dsv41_mono_layers.py | 55 +++++---- vllm/config/kernel.py | 6 + vllm/engine/arg_utils.py | 4 + vllm/envs.py | 7 -- vllm/models/deepseek_v41/amd/model.py | 12 +- vllm/models/deepseek_v41/amd/mono/README.md | 23 ++-- .../deepseek_v41/amd/mono/attention/back.py | 10 +- .../deepseek_v41/amd/mono/attention/front.py | 28 ++++- vllm/models/deepseek_v41/amd/mono/layer.py | 84 ++++++++++++- vllm/models/deepseek_v41/amd/mono/runner.py | 39 +++++- .../deepseek_v41/amd/mono/stages/seam.py | 24 ++-- vllm/models/deepseek_v41/amd/mono_decode.py | 115 +++++++++++++----- 13 files changed, 379 insertions(+), 111 deletions(-) diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_mono_decode.py index 684e72c00482..309d278bcb42 100644 --- a/tests/models/test_dsv41_mono_decode.py +++ b/tests/models/test_dsv41_mono_decode.py @@ -27,21 +27,25 @@ def forward(self, *args, **kwargs): self.calls.append((args, kwargs)) return ("mono outputs",) + def first_pre_mix(self, rows): + return torch.zeros(rows, 4) + def ffn(self, *args): self.calls.append((args, {})) return ("ffn outputs",) -def _layer(comp="comp"): +def _layer(comp="comp", indexer=None): cache = lambda: torch.zeros(4, 128 * md.RECORD, dtype=torch.uint8) # noqa: E731 attn = SimpleNamespace( swa_cache_layer=SimpleNamespace(prefix="swa", kv_cache=cache()), compressed_cache_prefix=comp, topk_indices_buffer=torch.zeros(64, 512, dtype=torch.int32), compress_ratio=1, + indexer=indexer, _compressed_kv_cache=cache, ) - return SimpleNamespace(attn=attn) + return SimpleNamespace(attn=attn, engram=None, hc_mult=4) def _swa(**overrides): @@ -77,7 +81,12 @@ def run(monkeypatch): ``mode`` and ``metadata``; returns (result, runner calls).""" def _run( - mode=CUDAGraphMode.FULL, metadata=None, ffn_only=False, entry=None, **inputs + mode=CUDAGraphMode.FULL, + metadata=None, + ffn_only=False, + entry=None, + layer=None, + **inputs, ): if metadata is None: metadata = { @@ -97,7 +106,7 @@ def _run( # has no compressed cache part, _ = args.pop("x"), args.pop("positions") return mono.ffn(_layer(comp=None), part, **args), runner.calls - return mono(_layer(), **args), runner.calls + return mono(layer or _layer(), **args), runner.calls return _run @@ -114,6 +123,23 @@ def test_decode_step_runs_mono(run, mode): 1, ) assert kwargs["comp_cache"].shape == (4, 128, md.RECORD) + assert kwargs["fold"] and "select" not in kwargs + + +def test_first_layer_runs_mono_from_the_embedding(run): + """The first layer's residual is the embedding broadcast to the four + streams, which K1 takes as it is (no owed post to fold).""" + out, calls = run(residual=None, post_mix=None, res_mix=None, pre_mix=None) + assert out == ("mono outputs",) and len(calls) == 1 + args, kwargs = calls[0] + assert args[2].shape == (M, 4, 5120) and not kwargs["fold"] + + +def test_index_source_selects_between_the_kernels(run): + """An index source hands the runner its compressor and indexer to run + between K1 and K2.""" + out, calls = run(layer=_layer(indexer=object())) + assert out == ("mono outputs",) and callable(calls[0][1]["select"]) def test_padded_decode_step_runs_mono(run): @@ -149,7 +175,9 @@ def test_padded_decode_step_runs_mono(run): ), dict(metadata={}), # a profile / dummy run: no attention metadata dict(**_inputs(md.MAX_ROWS + 6)), - dict(residual=None), # the first layer's seam broadcasts the embedding + # a later pipeline stage's first layer: its stream arrives without + # the mixes + dict(residual=None, x=torch.zeros(M, 4, 5120, dtype=torch.bfloat16)), ], ids=[ "piecewise", @@ -158,7 +186,7 @@ def test_padded_decode_step_runs_mono(run): "tokens-past-rows", "no-metadata", "rows", - "first-layer", + "stage-input", ], ) def test_other_steps_take_the_original_path(run, case): @@ -202,22 +230,41 @@ def _decoder_layer(engram=None, **attn): ) +def _indexer(): + from vllm.models.deepseek_v41.amd.rocm import DeepseekV41RocmMxfp4Indexer + + return object.__new__(DeepseekV41RocmMxfp4Indexer) + + @pytest.mark.parametrize( "layer, path", [ (_decoder_layer(), "whole"), - (_decoder_layer(compress_ratio=0), "ffn"), - (_decoder_layer(engram=object()), "ffn"), - (_decoder_layer(compressor=object(), indexer=object()), "ffn"), + (_decoder_layer(compress_ratio=0), "whole"), + (_decoder_layer(engram=object()), "whole"), + (_decoder_layer(compressor=object(), indexer=_indexer()), "whole"), + (_decoder_layer(indexer=_indexer()), "whole"), + (_decoder_layer(compressor=object()), "ffn"), (_decoder_layer(indexer=object()), "ffn"), (_decoder_layer(kv_cache_dtype="auto"), "ffn"), (_decoder_layer(layer_id=40), None), ], - ids=["standard", "first", "engram", "kv-source", "index-source", "kv", "draft"], + ids=[ + "standard", + "window", + "engram", + "kv-source", + "index-source", + "compressor-only", + "other-indexer", + "kv", + "draft", + ], ) def test_layers_take_whole_or_ffn_launch(monkeypatch, layer, path): - """The whole-layer kernels serve the standard layers; every other backbone - layer keeps vLLM's attention and runs its FFN half as one launch.""" + """The whole-layer kernels serve every backbone layer of the checkpoint (an + index source's compressor and indexer run between them); a layer outside + them keeps vLLM's attention and runs its FFN half as one launch.""" _deployment(monkeypatch) mono = md.MonoDecodeLayer.create(layer, _config()) got = None if mono is None else "ffn" if mono.ffn_only else "whole" @@ -254,7 +301,7 @@ def test_other_steps_reduce_wo_b_themselves(run, case): assert out is None and not calls -def _config(**parallel): +def _config(monokernel=True, **parallel): parallel = ( dict(enable_expert_parallel=False, enable_eplb=False, data_parallel_size=1) | parallel @@ -262,6 +309,7 @@ def _config(**parallel): return SimpleNamespace( model_config=SimpleNamespace(hf_config=SimpleNamespace(num_hidden_layers=40)), parallel_config=SimpleNamespace(**parallel), + kernel_config=SimpleNamespace(monokernel=monokernel), ) @@ -269,7 +317,6 @@ def _deployment(monkeypatch, cdna=4, cus=256): """A deployment the kernels take (the runner module stubbed: no FlyDSL).""" import vllm.platforms.rocm as rocm - monkeypatch.setenv("VLLM_ROCM_MONO_DECODE", "1") monkeypatch.delenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", raising=False) monkeypatch.setattr(rocm, "get_cdna_version", lambda: cdna) monkeypatch.setattr(md, "get_tensor_model_parallel_world_size", lambda: 2) @@ -280,6 +327,14 @@ def _deployment(monkeypatch, cdna=4, cus=256): monkeypatch.setitem(sys.modules, name, runner) +def test_off_without_the_flag(monkeypatch): + """No --monokernel: no mono layer, whatever the deployment.""" + _deployment(monkeypatch) + assert ( + md.MonoDecodeLayer.create(_decoder_layer(), _config(monokernel=False)) is None + ) + + @pytest.mark.parametrize("cdna", [3, 5]) def test_opt_in_outside_the_kernels_cdna_raises(monkeypatch, cdna): _deployment(monkeypatch, cdna=cdna) diff --git a/tests/models/test_dsv41_mono_layers.py b/tests/models/test_dsv41_mono_layers.py index 580462a82394..a924cfda50c3 100644 --- a/tests/models/test_dsv41_mono_layers.py +++ b/tests/models/test_dsv41_mono_layers.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """The DeepSeek-V4.1 mono decode path through the real model: a five-layer -DeepSeek-V4.1-Flash at TP2 and TP4 on CDNA4, served with ``VLLM_ROCM_MONO_DECODE`` on +DeepSeek-V4.1-Flash at TP2 and TP4 on CDNA4, served with ``--monokernel`` on and off on the same random weights: every decoder layer's outputs at every decode step must match. Unlike ``test_dsv41_mono_numerics.py``, which checks the kernels against vLLM's ops one stage at a time, this runs @@ -10,12 +10,13 @@ the indexer's top-k -- so a change to the decoder layer that the kernels do not follow fails here. -The five layers hold both kinds the mono path takes: -- 0: sliding window only -> the FFN launch -- 1: ratio-2 KV and index source -> the FFN launch -- 2: ratio-2 consumer -> the whole layer (K1 + K2) -- 3: ratio-1 KV, index and candidate source -> the FFN launch -- 4: ratio-1 consumer -> the whole layer +Every layer takes the whole-layer path (K1 + K2); the five hold each variant: +- 0: the first layer, sliding window only (K1 takes the embedding settled) +- 1: ratio-2 KV and index source (vLLM's compressor and indexer between K1 + and K2) +- 2: ratio-2 consumer with an Engram (K1 takes the residual settled) +- 3: ratio-1 KV, index and candidate source +- 4: ratio-1 index source over 3's candidates and KV """ import zlib @@ -34,13 +35,17 @@ # one entry a layer, the three MTP layers' included "compress_ratios": [0, 2, 2, 1, 1, 0, 0, 0], "kv_source_layer_ids": [1, 3], - "index_source_layer_ids": [1, 3], + # 4: an indexer of its own over 3's candidates (a REINDEX layer) + "index_source_layer_ids": [1, 3, 4], "candidate_source_layer_id": 3, - "engram_layer_ids": [], - "engram_num_embeddings": [], + # an Engram on a whole layer (K1 then takes its residual settled), its + # tables shrunk: 24 prime-sized ranges from 1000 + "engram_layer_ids": [2], + "engram_num_embeddings": [32768], + "engram_vocab_size": 1000, "dspark_target_layer_ids": [2, 3, 4], } -WHOLE, FFN = (2, 4), (0, 1, 3) +WHOLE, FFN = (0, 1, 2, 3, 4), () STEPS = 5 # decode steps after the prefill's token @@ -227,7 +232,6 @@ def _decode(vllm_runner, monkeypatch, tp: int, mono: bool, prompts, path: str): from vllm import SamplingParams from vllm.inputs import TokensPrompt - monkeypatch.setenv("VLLM_ROCM_MONO_DECODE", "1" if mono else "0") monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1") with vllm_runner( @@ -237,6 +241,7 @@ def _decode(vllm_runner, monkeypatch, tp: int, mono: bool, prompts, path: str): hf_overrides=TRUNCATED, worker_extension_cls=f"{__name__}.MonoProbe", moe_backend="aiter", + monokernel=mono, language_model_only=True, tokenizer_mode="deepseek_v41", attention_config={"indexer_kv_dtype": "mxfp4", "indexer_sparse_logits": True}, @@ -300,15 +305,16 @@ def test_mono_decode_matches_the_model(vllm_runner, monkeypatch, tmp_path, tp): same.add((len(prompt), len(prompt) + step)) # each layer's outputs for those tokens, mono against vLLM's path, token by - # token. Layer 0 sees the same inputs on both paths at a request's first - # decode token: there the bound is test_dsv41_mono_numerics' (about 3x the - # kernels' 0.4%); its later tokens attend over KV each path wrote, 3%. Each - # later layer takes the previous one's slightly different output, which - # random weights amplify (2-4% by layer 4): 8%. A token whose top-6 routing - # is a near-tie on vLLM's path (6th / 7th score within 5e-3: a 0.4% input - # difference moves scores near 0.8 by about 3e-3) may pick another expert; - # once it has, its residual stream differs and the later layers are exempt - # for it. The residual stream: 3% (0.3% seen). A layer the kernels no + # token. Layer 0 sees the same inputs on both paths, but its FFN input + # differs by up to test_dsv41_mono_numerics' 1.2%, which random weights + # amplify through the MoE: 3%. Each later layer takes the previous one's + # slightly different output, which random weights amplify (2-4% by layer + # 4): 8%. A token whose top-6 routing is a near-tie on vLLM's path may pick + # another expert: a 0.4% input difference moves scores near 0.8 by about + # 3e-3, so 6th / 7th scores within 5e-3 at layer 0 and, past it, where the + # input carries the earlier layers' differences (0.4-0.7%), within 1e-2. + # Once it has, its residual stream differs and the later layers are exempt + # for it. The residual stream: 3% (0.5% seen). A layer the kernels no # longer follow is off by far more, on every token. for rank in range(tp): got, _ = torch.load(f"{tmp_path}/on.rank{rank}") @@ -317,8 +323,6 @@ def test_mono_decode_matches_the_model(vllm_runner, monkeypatch, tmp_path, tp): same & {k[1:] for k in got} & {k[1:] for k in want} ) assert len(keys) >= len(prompts), keys - # the bound that sees a 3% error: a request's first decode token - assert any(p == n for n, p in keys), keys flipped: set = set() for i in sorted(WHOLE + FFN): for k in keys: @@ -327,9 +331,10 @@ def test_mono_decode_matches_the_model(vllm_runner, monkeypatch, tmp_path, tp): (x, res), (x_ref, res_ref) = got[(i, *k)], want[(i, *k)] err = _rel(res, res_ref) assert err < 0.03, (rank, i, k, "residual", err) - bound = (0.015 if k[1] == k[0] else 0.03) if i == 0 else 0.08 + bound = 0.03 if i == 0 else 0.08 err = _rel(x, x_ref) if err >= bound: - assert gaps[(i, *k)] < 5e-3, (rank, i, k, "x", err) + tie = 5e-3 if i == 0 else 1e-2 + assert gaps[(i, *k)] < tie, (rank, i, k, "x", err) flipped.add(k) assert len(flipped) <= len(keys) // 4, (rank, sorted(flipped)) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index fa3ecf5df95d..5c4b877fd795 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -291,6 +291,12 @@ class KernelConfig: enable_jit_warmup: bool = True """If True, run JIT compile warmup during kernel warmup.""" + monokernel: bool = False + """If True, run the eligible layers of decode steps on the model's + MonoKernel: a persistent kernel that runs most or all of one decoder layer + in a single launch (DeepSeek-V4.1 on MI355X). A deployment that cannot run + it is refused at startup.""" + moe_backend: MoEBackend = "auto" """Backend for MoE expert computation kernels. Available options: diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index ea628467e0b3..86199e5f007b 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -758,6 +758,7 @@ class EngineArgs: enable_flashinfer_autotune: bool = get_field( KernelConfig, "enable_flashinfer_autotune" ) + monokernel: bool = get_field(KernelConfig, "monokernel") worker_cls: str = ParallelConfig.worker_cls worker_extension_cls: str = ParallelConfig.worker_extension_cls @@ -1802,6 +1803,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--enable-flashinfer-autotune", **kernel_kwargs["enable_flashinfer_autotune"], ) + kernel_group.add_argument("--monokernel", **kernel_kwargs["monokernel"]) moe_backend_kwargs = kernel_kwargs["moe_backend"] moe_backend_kwargs["type"] = lambda s: s.lower().replace("-", "_") kernel_group.add_argument("--moe-backend", **moe_backend_kwargs) @@ -2754,6 +2756,8 @@ def create_engine_config( "are mutually exclusive" ) kernel_config.enable_flashinfer_autotune = self.enable_flashinfer_autotune + if self.monokernel: + kernel_config.monokernel = True if self.moe_backend != "auto": kernel_config.moe_backend = self.moe_backend if self.linear_backend != "auto": diff --git a/vllm/envs.py b/vllm/envs.py index a2b46030772f..bb1c03a9dabd 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -156,7 +156,6 @@ VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = False VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS: bool = False - VLLM_ROCM_MONO_DECODE: bool = False VLLM_ROCM_USE_AITER_TRITON_GEMM: bool = True VLLM_ROCM_USE_SKINNY_GEMM: bool = True VLLM_ROCM_FP8_PADDING: bool = True @@ -1398,12 +1397,6 @@ def _resolve_rust_cli_path() -> str | None: os.getenv("VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS", "False").lower() in ("true", "1") ), - # Run the eligible layers of decode steps on the model's mono decode - # kernels: persistent launches a layer, the TP all-reduces in-kernel - # (DeepSeek-V4.1 on CDNA4). By default is disabled. - "VLLM_ROCM_MONO_DECODE": lambda: ( - os.getenv("VLLM_ROCM_MONO_DECODE", "False").lower() in ("true", "1") - ), # Whether to use aiter triton kernels for gemm ops. # By default is enabled. "VLLM_ROCM_USE_AITER_TRITON_GEMM": lambda: ( diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 5de173a8706c..4c5f15858749 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -282,7 +282,17 @@ def forward( engram_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: if self.mono is not None: - out = self.mono(self, x, positions, residual, post_mix, res_mix, pre_mix) + out = self.mono( + self, + x, + positions, + residual, + post_mix, + res_mix, + pre_mix, + engram_hashes, + engram_mask, + ) if out is not None: return out # Layer 0's attention seam projects the 2-D embedding with the folded diff --git a/vllm/models/deepseek_v41/amd/mono/README.md b/vllm/models/deepseek_v41/amd/mono/README.md index 88e40d73dd0a..462f7bdecd28 100644 --- a/vllm/models/deepseek_v41/amd/mono/README.md +++ b/vllm/models/deepseek_v41/amd/mono/README.md @@ -2,9 +2,10 @@ One fused decoder layer for vLLM's decode steps on MI355X: the attention seam, the attention, the attention all-reduce, the FFN seam, the MoE and the MoE -all-reduce, in two persistent launches a layer. The other backbone layers keep -vLLM's attention and run the rest -- the attention all-reduce, the FFN seam, -the MoE and its all-reduce -- as one launch (the FFN launch). ATOM's V4.1 mono decode +all-reduce, in two persistent launches a layer. A layer outside the whole-layer +kernels keeps vLLM's attention and runs the rest -- the attention all-reduce, +the FFN seam, the MoE and its all-reduce -- as one launch (the FFN launch). +ATOM's V4.1 mono decode (ROCm/ATOM `atom/models/deepseek_v41/mono`) is the reference design: the seam and MoE stages (`stages/`) and the framework under them (`common/`) are adapted from it, each file naming its origin; the attention stages (`attention/`) are @@ -12,13 +13,15 @@ new, reading vLLM's KV records and weights as loaded. ## Scope -- Layers: the whole layer on every backbone layer with the standard seam and no - compressor / indexer: 3-7, 9-13, 15-19 (compress ratio 2) and 21-23, 25-27, - 29-31, 33-35, 37-39 (ratio 1), 30 of 40. The FFN launch on the other 10: - layer 0 (embedding-broadcast seam), layer 1 (Engram at the seam) and the - compressor / indexer layers 2, 8, 14 (also Engram), 20, 24, 28, 32, 36, - whose wo_b leaves its TP reduce to it. The DSpark draft layers keep vLLM's - path. The paths interleave freely (same tensors at the layer boundary). +- Layers: the whole layer on all 40 backbone layers. K1 takes the residual + settled, without folding the previous output in, on layer 0 (the embedding + broadcast to the four streams) and after an Engram write (layers 1 and 14, + vLLM's Engram running before K1). On the index sources (2, 8, 14, 20 with + their compressor; 24, 28, 32, 36) K1 also hands out its normed input and bf16 + qr, vLLM's compressor and indexer run between K1 and K2, and K2 builds the + key table from that top-k. The FFN launch serves a layer outside these (an + MXFP8 KV cache, another indexer); the DSpark draft layers keep vLLM's path. + The paths interleave freely (same tensors at the layer boundary). - Steps: decode only (FULL graph or eager), M <= 48 rows (8 requests x 6 DSpark tokens), TP2 / TP4, `fp8_ds_mla` KV records (ROCm), text tokens (the routing bias is `e_score_correction_bias`). diff --git a/vllm/models/deepseek_v41/amd/mono/attention/back.py b/vllm/models/deepseek_v41/amd/mono/attention/back.py index 4526f8eda7ee..5cd4793be243 100644 --- a/vllm/models/deepseek_v41/amd/mono/attention/back.py +++ b/vllm/models/deepseek_v41/amd/mono/attention/back.py @@ -151,7 +151,9 @@ def key_rec(c, t, k): table: slot | 1 << 31 for a compressed row, a window slot otherwise, -1 for none (it reads a zero record: its scores are masked, its values zeros).""" a, ratio = c["args"], c["ratio"] - val = ld_i32(a["kt"], t * KEYS + fx.min(fx.max(k, 0), KEYS - 1)) + # a table this launch wrote (``krdy``) is read past the CU's cache + cm = CM_DEV if const_expr("krdy" in c) else 0 + val = ld_i32(a["kt"], t * KEYS + fx.min(fx.max(k, 0), KEYS - 1), cm) ok = val != fx.Int32(-1) slot = ok.select(val & fx.Int32(0x7FFFFFFF), fx.Int32(0)) if const_expr(ratio > 0): @@ -209,9 +211,11 @@ def stage_split(c, unit): g = lane // 16 r16 = lane % 16 k0 = ch * ck + if const_expr("krdy" in c): + await_flags(c, c["krdy"], 1, lambda u: t) # the key count, this lane's key and q: independent loads, out together - # (K1's key table holds -1 past a token's keys) - kv_len = ld_i32(a["klen"], t) + # (the key table holds -1 past a token's keys) + kv_len = ld_i32(a["klen"], t, CM_DEV if const_expr("krdy" in c) else 0) key = k0 + kt * 16 + r16 data, sc, okv = key_rec(c, t, key) qrow = d.heads * HEAD_DIM // 8 diff --git a/vllm/models/deepseek_v41/amd/mono/attention/front.py b/vllm/models/deepseek_v41/amd/mono/attention/front.py index 3ac9e3fe9b10..23c009739dc2 100644 --- a/vllm/models/deepseek_v41/amd/mono/attention/front.py +++ b/vllm/models/deepseek_v41/amd/mono/attention/front.py @@ -188,6 +188,8 @@ def stage_kt(c, t): ntopk.""" tid = c["tid"] a, ratio = c["args"], c["ratio"] + # device scope where this launch reads the table (K2's) + cm = CM_DEV if const_expr(c.get("dev", False)) else 0 slot = ld_i32(a["slot"], 2 * t) live = slot >= 0 nswa = live.select(ld_i32(a["swa_lens"], t), fx.Int32(0)) @@ -212,10 +214,10 @@ def stage_kt(c, t): c_val = (local >= 0).select(c_slot, fx.Int32(-1)) val = (k < ntopk).select(c_val, val) if k < KEYS: - bo.buffer_store(val, rsrc(a["kt"]), t * KEYS + k) + bo.buffer_store(val, rsrc(a["kt"]), t * KEYS + k, cache_modifier=cm) if tid == 0: - bo.buffer_store(kv_len, rsrc(a["klen"]), t) - bo.buffer_store(ntopk, rsrc(a["klen"]), c["S"] + t) + bo.buffer_store(kv_len, rsrc(a["klen"]), t, cache_modifier=cm) + bo.buffer_store(ntopk, rsrc(a["klen"]), c["S"] + t, cache_modifier=cm) # ---------------------------------------------------------------- wqkv_a @@ -357,6 +359,20 @@ def stage_qkv(c, t): amax = fx.max(amax, abs(y)) code = mx_code(butterfly(amax, (1, 2), fx.max)) mul = mx_mul(code) + if const_expr(c.get("index", False)): # noqa: SIM102 + # the bf16 qr, vLLM's indexer's input + if mine & (tt * 8 < Q_RANK): + words = [ + fx.Vector.from_elements([ys[2 * i], ys[2 * i + 1]], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Int32)[0] + for i in range(4) + ] + bo.buffer_store( + fx.Vector.from_elements(words, fx.Int32), + rsrc(a["qr_out"]), + t * (Q_RANK // 2) + 4 * tt, + ) if mine & (tt * 8 < Q_RANK): w0 = fp8_pack4(*[clamp_fp8(ys[i] * mul) for i in range(4)]) w1 = fp8_pack4(*[clamp_fp8(ys[4 + i] * mul) for i in range(4)]) @@ -597,8 +613,10 @@ def run_front(c, key, bid, x8_given=False, before_wqkv=None, after_wqkv=None): u = free + n_free * tid if u < xq_units: c["mb"].put(c["xrdy"], u, fx.Int32(1)) - for t in range(is_free.select(free, fx.Int32(s)), s, n_free): - stage_kt(c, t) + if const_expr(not c.get("index", False)): + # an index source's top-k comes after this launch: K2 builds its table + for t in range(is_free.select(free, fx.Int32(s)), s, n_free): + stage_kt(c, t) stamp(c, 1) poss = [wqb_pos(c, t0, n) for t0, n in tiles] diff --git a/vllm/models/deepseek_v41/amd/mono/layer.py b/vllm/models/deepseek_v41/amd/mono/layer.py index c1be9e3428ae..dcd39ae2b754 100644 --- a/vllm/models/deepseek_v41/amd/mono/layer.py +++ b/vllm/models/deepseek_v41/amd/mono/layer.py @@ -58,6 +58,12 @@ class MonoBuild: tp: int ratio: int # the layer's compress ratio (its attention's key sources) timeline: bool = False + # K1's seam folds the previous layer's output into the residual; off for a + # residual that arrives settled (the first layer, after an Engram) + fold: bool = True + # the layer selects its own top-k (an index source): K1 hands its normed + # input and qr out for vLLM's indexer, K2 builds the key table + index: bool = False def _moe_key(key: MonoBuild) -> MoeBuild: @@ -100,6 +106,7 @@ def place(regions): lo = min(o for o, _ in mr.values()) place({n: (o - lo, n_) for n, (o, n_) in mr.items()}) place({"normed": (0, s * HIDDEN * 2), "xrdy_moe": (s * HIDDEN * 2 + 256, s * 8)}) + place({"krdy": (0, s * 8)}) return out @@ -175,7 +182,15 @@ class K1Lds: seam: SeamLds front: FrontLds # type: ignore[valid-type] - name = kernel_symbol("dsv41_mono_k1", s=s, tp=tp, r=key.ratio, tl=key.timeline) + name = kernel_symbol( + "dsv41_mono_k1", + s=s, + tp=tp, + r=key.ratio, + tl=key.timeline, + f=key.fold, + i=key.index, + ) keyed = key_tuple(key, SOURCES) @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) @@ -215,6 +230,8 @@ def k1( t2r: Int64, kt: Int64, klen: Int64, + normed_out: Int64, + qr_out: Int64, scratch: Int64, epoch: Int64, tl: Int64, @@ -250,9 +267,11 @@ def k1( "t2r": t2r, "kt": kt, "klen": klen, + "qr_out": qr_out, } flayout = {n: layout[n] for n in front_scratch(s, Dims(tp))} c = mb_front.front_context(fkey, lds.front.peek(), fargs, scratch, tag, flayout) + c["index"] = key.index amb = Mailbox(tag) tid = fx.thread_idx.x ca = { @@ -284,6 +303,7 @@ def k1( "pre_out": pre_out, }, "d": StageDims(tp), + "fold": key.fold, } for region in ("lin", "pmix"): ca[region] = sreg(scratch, layout[region][0], region) @@ -295,9 +315,14 @@ def seam_slice(): seam.stage_slice(ca, task) gpu.barrier() + def store(*v): + store_x8(c, *v) + if const_expr(key.index): + store_normed(normed_out, *v) + def seam_rest(): for t in range(first_task(bid, NORM0), s, BLOCKS): - seam.stage_norm(ca, t, lambda *v: store_x8(c, *v)) + seam.stage_norm(ca, t, store) publish_x8(c, t) for t in range(first_task(bid, GATE0), s, BLOCKS): seam.stage_gate(ca, t) @@ -344,6 +369,8 @@ def launch( t2r: Int64, kt: Int64, klen: Int64, + normed_out: Int64, + qr_out: Int64, scratch: Int64, epoch: Int64, tl: Int64, @@ -386,6 +413,8 @@ def launch( t2r, kt, klen, + normed_out, + qr_out, scratch, epoch, tl, @@ -544,7 +573,9 @@ class K2Lds: seam: SeamLds # type: ignore[valid-type] moe: MoeLds # type: ignore[valid-type] - name = kernel_symbol("dsv41_mono_k2", s=s, tp=tp, r=key.ratio, tl=key.timeline) + name = kernel_symbol( + "dsv41_mono_k2", s=s, tp=tp, r=key.ratio, tl=key.timeline, i=key.index + ) keyed = key_tuple(key, SOURCES) @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) @@ -558,6 +589,13 @@ def k2( comp_block: Int32, kt: Int64, klen: Int64, + slot: Int64, + swa_idx: Int64, + swa_lens: Int64, + t2r: Int64, + topk: Int64, + comp_bt: Int64, + bt_stride: Int32, pos: Int64, sink: Int64, qk_scale: Int32, @@ -642,6 +680,32 @@ def push(t, col, v0, v1): if const_expr(key.timeline): c["tl"], c["tl_points"] = tl, mb_back.BACK_POINTS mb_back.epoch_begin(c, epoch, ep) + if const_expr(key.index): + # the key table, from the top-k vLLM's indexer selected after K1, + # on the last CTAs: each split task waits on its token's + ckt = { + "S": s, + "tid": c["tid"], + "dev": True, + "ratio": key.ratio, + "args": { + "slot": slot, + "swa_idx": swa_idx, + "swa_lens": swa_lens, + "pos": pos, + "t2r": t2r, + "topk": topk, + "comp_block": comp_block, + "comp_bt": comp_bt, + "bt_stride": bt_stride, + "kt": kt, + "klen": klen, + }, + } + c["krdy"] = sreg(scratch, layout["krdy"][0], "krdy") + for t in range(first_task(bid, BLOCKS - s), s, BLOCKS): + mb_front.stage_kt(ckt, t) + mb_back.publish(c, c["krdy"], t) mb_back.run_back(c, bkey, bid) gpu.barrier() @@ -685,6 +749,13 @@ def launch( comp_block: Int32, kt: Int64, klen: Int64, + slot: Int64, + swa_idx: Int64, + swa_lens: Int64, + t2r: Int64, + topk: Int64, + comp_bt: Int64, + bt_stride: Int32, pos: Int64, sink: Int64, qk_scale: Int32, @@ -736,6 +807,13 @@ def launch( comp_block, kt, klen, + slot, + swa_idx, + swa_lens, + t2r, + topk, + comp_bt, + bt_stride, pos, sink, qk_scale, diff --git a/vllm/models/deepseek_v41/amd/mono/runner.py b/vllm/models/deepseek_v41/amd/mono/runner.py index eb379902754f..ef9a1978e422 100644 --- a/vllm/models/deepseek_v41/amd/mono/runner.py +++ b/vllm/models/deepseek_v41/amd/mono/runner.py @@ -16,12 +16,13 @@ from __future__ import annotations import os +from collections.abc import Callable from dataclasses import dataclass from pathlib import Path import torch -from .attention.plan import HEAD_DIM, HIDDEN, KEYS, Dims +from .attention.plan import HEAD_DIM, HIDDEN, KEYS, Q_RANK, Dims from .common.plan import BLOCKS from .layer import ( EPOCH_WORDS, @@ -170,6 +171,12 @@ def __init__(self, tp: int, rank: int, group, device: torch.device | str = "cuda self.klen = torch.zeros(2 * MAX_TOKENS, dtype=torch.int32, device=dev) self._dummy = torch.zeros(4, dtype=torch.int32, device=dev) self._zero_rec = torch.zeros(1024, dtype=torch.uint8, device=dev) + # the first layer's incoming pre-mix: stream 0 of the broadcast embedding + self._first_pre = torch.zeros(MAX_TOKENS, HC, dtype=torch.float32, device=dev) + self._first_pre[:, 0] = 1.0 + # an index source's K1 outputs, vLLM's indexer's inputs + self.normed = torch.zeros(MAX_TOKENS, HIDDEN, dtype=torch.bfloat16, device=dev) + self.qr = torch.zeros(MAX_TOKENS, Q_RANK, dtype=torch.bfloat16, device=dev) # the attention seam's outputs, K1 -> K2 self.res_mid = torch.zeros( MAX_TOKENS, HC, HIDDEN, dtype=torch.bfloat16, device=dev @@ -203,14 +210,17 @@ def scratch(self, tokens: int) -> torch.Tensor: self._scratch[tokens] = buf return buf + def first_pre_mix(self, tokens: int) -> torch.Tensor: + return self._first_pre[:tokens] + @staticmethod def supports(tokens: int) -> bool: return 1 <= tokens <= MAX_TOKENS - def kernels(self, tokens: int, ratio: int): - key = (tokens, ratio) + def kernels(self, tokens: int, ratio: int, fold: bool = True, index: bool = False): + key = (tokens, ratio, fold, index) if key not in self._kernels: - b = MonoBuild(tokens, self.tp, ratio) + b = MonoBuild(tokens, self.tp, ratio, fold=fold, index=index) self._kernels[key] = (build_mono_k1(b), build_mono_k2(b)) return self._kernels[key] @@ -249,15 +259,21 @@ def forward( comp_cache: torch.Tensor | None = None, comp_block_table: torch.Tensor | None = None, outs: tuple | None = None, + fold: bool = True, + select: Callable[[torch.Tensor, torch.Tensor], None] | None = None, ): """-> (out, residual, post_mix, res_mix, pre_mix): the layer's outputs - (``outs``: buffers to write them into).""" + (``outs``: buffers to write them into). ``select(normed, qr)``: an index + source's top-k selection into ``topk_indices``, run between the + kernels on K1's normed input and qr.""" M = x.shape[0] assert self.supports(M), M a = w.attn assert a is not None ratio = a.ratio - k1, k2 = self.kernels(M, ratio) + index = select is not None + assert not index or ratio + k1, k2 = self.kernels(M, ratio, fold, index) st = torch.cuda.current_stream() if outs is None: outs = self._outs(M, residual) @@ -324,11 +340,15 @@ def forward( token_to_req.data_ptr(), self.kt.data_ptr(), self.klen.data_ptr(), + self.normed.data_ptr(), + self.qr.data_ptr(), self.scratch(M).data_ptr(), self.epoch.data_ptr(), 0, stream=st, ) + if select is not None: + select(self.normed[:M], self.qr[:M]) k2( self.q.data_ptr(), swa_cache.data_ptr(), @@ -339,6 +359,13 @@ def forward( comp_block, self.kt.data_ptr(), self.klen.data_ptr(), + slot_mapping.data_ptr(), + swa_indices.data_ptr(), + swa_lens.data_ptr(), + token_to_req.data_ptr(), + topk.data_ptr(), + bt.data_ptr(), + bt_stride, positions.data_ptr(), a.attn_sink.data_ptr(), self._qk_scale, diff --git a/vllm/models/deepseek_v41/amd/mono/stages/seam.py b/vllm/models/deepseek_v41/amd/mono/stages/seam.py index 82abee9f46fa..ea73e2ef5435 100644 --- a/vllm/models/deepseek_v41/amd/mono/stages/seam.py +++ b/vllm/models/deepseek_v41/amd/mono/stages/seam.py @@ -141,6 +141,8 @@ def slice_loads(c, c0, e): """Element e's loads of ``stage_slice`` (e clamped into the step's).""" a = c["args"] t, o, col = slice_at(c0, e) + if const_expr(not c.get("fold", True)): + return (ld_bf(a["res_in"], (t * HC + o) * HIDDEN + col),) r = [ld_bf(a["res_in"], (t * HC + h) * HIDDEN + col) for h in range(HC)] cm = [ld_f32(a["comb_in"], (t * HC + h) * HC + o) for h in range(HC)] return r, cm, ld_f32(a["post_in"], t * HC + o) @@ -150,15 +152,19 @@ def slice_row(c, c0, e, got): """Element e of ``stage_slice``'s R_new from ``slice_loads``' values.""" a = c["args"] t, o, col = slice_at(c0, e) - r, cm, post = got - # aiter's mhc_fused_post_pre_delayed_rmsnorm: ((post x + c0 r0) - # + c1 r1) + c2 r2, each product rounded, then fma(c3, r3, ..) - v = post * pend_value(c, t, col, e % COLS) + (r[0] * cm[0]) - for h in range_constexpr(1, HC - 1): - v = v + cm[h] * r[h] - # bf16 R_new is what the collapse, the sums of squares and the - # mix projection all read - v = bf16_round(fma(cm[HC - 1], r[HC - 1], v)) + if const_expr(not c.get("fold", True)): + # a settled residual (no owed post): stream o as it arrives + v = got[0] + else: + r, cm, post = got + # aiter's mhc_fused_post_pre_delayed_rmsnorm: ((post x + c0 r0) + # + c1 r1) + c2 r2, each product rounded, then fma(c3, r3, ..) + v = post * pend_value(c, t, col, e % COLS) + (r[0] * cm[0]) + for h in range_constexpr(1, HC - 1): + v = v + cm[h] * r[h] + # bf16 R_new is what the collapse, the sums of squares and the + # mix projection all read + v = bf16_round(fma(cm[HC - 1], r[HC - 1], v)) bo.buffer_store(v.to(fx.BFloat16), rsrc(a["res_out"]), (t * HC + o) * HIDDEN + col) fx.ptr_store(v, c["rl"] + (t * KT + o * COLS + e % COLS)) diff --git a/vllm/models/deepseek_v41/amd/mono_decode.py b/vllm/models/deepseek_v41/amd/mono_decode.py index b404a42de37a..ddeb6b26314f 100644 --- a/vllm/models/deepseek_v41/amd/mono_decode.py +++ b/vllm/models/deepseek_v41/amd/mono_decode.py @@ -1,15 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepSeek-V4.1 mono decode layer on ROCm CDNA4 (``VLLM_ROCM_MONO_DECODE=1``). +"""DeepSeek-V4.1 mono decode layer on ROCm CDNA4 (``--monokernel``). A decode step's backbone layer -- the attention seam, the attention, its TP all-reduce, the FFN seam, the MoE and its all-reduce -- runs as two persistent FlyDSL launches (``mono``, this directory) on the layer's loaded -weights, the TP reductions inside the kernels. A layer runs whole when it has -the standard seam and no compressor or indexer of its own (compress ratio 1 or -2) over ``fp8_ds_mla`` records. Any other backbone layer (the first, the Engram -layers, the compressor / indexer layers) keeps vLLM's attention seam and -attention and runs the rest as one launch, the FFN launch, from wo_b's +weights, the TP reductions inside the kernels. Every backbone layer of the +checkpoint runs whole (compress ratio 0, 1 or 2 over ``fp8_ds_mla`` records): +the first layer and the Engram layers hand K1 a settled residual, and an index +source runs vLLM's compressor and indexer between the two launches. A layer +outside that (an MXFP8 KV cache, another indexer) keeps vLLM's attention seam +and attention and runs the rest as one launch, the FFN launch, from wo_b's unreduced output. A step takes the kernels when it is decode only, at most ``MAX_ROWS`` rows (eager or a FULL CUDA graph: the model is not torch.compiled); a whole layer also needs causal SWA windows. Anything else runs the layer as @@ -20,7 +21,7 @@ other's pushes): the conditions read only metadata every rank holds alike. """ -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, Any, cast import torch @@ -62,7 +63,7 @@ def _mono_runner(device: torch.device): # its peer memory is allocated and exchanged eagerly; vLLM warms # every graph up eagerly before capturing it raise RuntimeError( - "VLLM_ROCM_MONO_DECODE was first reached inside a CUDA graph " + "--monokernel was first reached inside a CUDA graph " "capture, before an eager decode step." ) from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer @@ -113,6 +114,35 @@ def _unsupported_moe(layer: "DeepseekV4DecoderLayer", vllm_config) -> str | None return None +def _select( + layer: "DeepseekV4DecoderLayer", + positions: torch.Tensor, + normed: torch.Tensor, + qr: torch.Tensor, +) -> None: + """An index source's compressor and top-k selection, from K1's normed input + and bf16 qr: ``_forward_csa2_full`` / ``_forward_csa2_reindex`` and the + indexer op, serial.""" + attn = layer.attn + comp, indexer = attn.compressor, attn.indexer + assert indexer is not None + if comp is not None: + kv_score = torch.mm( + normed, comp.fused_wkv_wgate.weight.T, out_dtype=torch.float32 + ) + latent = comp(kv_score, positions) + comp.insert_cache(latent, positions, attn.rotary_emb) + indexer._produce_k(latent, positions, attn.indexer_rotary_emb) + weights, _ = indexer.weights_proj(normed) + q, q_scale, weights = indexer.forward_q( + qr, None, weights, positions, attn.indexer_rotary_emb + ) + assert weights is not None + indexer.indexer_op( + normed, (q, q_scale) if q_scale is not None else q, None, weights + ) + + class MonoDecodeLayer: """One decoder layer's mono path (``create``: None when the layer is not eligible): the whole layer (``__call__``) or, with ``ffn_only``, the FFN @@ -127,7 +157,7 @@ def __init__(self, ffn_only: bool = False) -> None: def create( layer: "DeepseekV4DecoderLayer", vllm_config ) -> "MonoDecodeLayer | None": - if not envs.VLLM_ROCM_MONO_DECODE: + if not vllm_config.kernel_config.monokernel: return None from vllm.platforms.rocm import get_cdna_version @@ -150,7 +180,7 @@ def create( if cus < BLOCKS: why = f"needs {BLOCKS} compute units, the GPU has {cus}" if why is not None: - raise ValueError(f"VLLM_ROCM_MONO_DECODE {why}.") + raise ValueError(f"--monokernel {why}.") # the layer: the kernels serve the backbone's seams and MoE config = vllm_config.model_config.hf_config attn, ffn = layer.attn, layer.ffn @@ -174,12 +204,17 @@ def create( if why is not None: logger.info_once("DSv4.1 mono decode off for some layers: %s", why) return None - # the whole layer: the kernels' attention takes the standard layers only - if layer.engram is not None: - why = "Engram" - elif attn.compressor is not None or attn.indexer is not None: - why = "a compressor / indexer" - elif attn.compress_ratio not in (1, 2): + # the whole layer: the kernels' attention takes the standard layers; an + # index source's compressor and indexer run between its two launches + from vllm.models.deepseek_v41.amd.rocm import DeepseekV41RocmMxfp4Indexer + + if attn.compressor is not None and attn.indexer is None: + why = "a compressor without an indexer" + elif attn.indexer is not None and not isinstance( + attn.indexer, DeepseekV41RocmMxfp4Indexer + ): + why = f"a {type(attn.indexer).__name__}" + elif attn.compress_ratio not in (0, 1, 2): why = f"compress ratio {attn.compress_ratio}" elif attn.kv_mxfp8 or attn.kv_cache_dtype != "fp8_ds_mla": why = f"a {attn.kv_cache_dtype} KV cache" @@ -297,20 +332,23 @@ def __call__( post_mix: torch.Tensor | None, res_mix: torch.Tensor | None, pre_mix: torch.Tensor | None, + engram_hashes: torch.Tensor | None = None, + engram_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, ...] | None: """The layer's outputs (x, residual, post_mix, res_mix, pre_mix) by the mono kernels, or None when this step takes the original path (always for an ``ffn_only`` layer).""" if self.ffn_only: return None - if residual is None or x.dim() != 2: - return self._skip("first layer") + first = residual is None + if first and x.dim() != 2: + return self._skip("first layer of a pipeline stage") why, swa, comp = self._decode_step(layer, x.shape[0]) attn = layer.attn - tensors = (x, residual, post_mix, res_mix, pre_mix) + tensors = (x,) if first else (x, residual, post_mix, res_mix, pre_mix) if why is None: # the kernels' attention: causal SWA windows, the compressed cache - if comp is None: + if comp is None and attn.compress_ratio: why = "layer metadata missing" elif swa.decode_swa_width != SWA_WIDTH: why = f"SWA width {swa.decode_swa_width}" @@ -324,12 +362,37 @@ def __call__( why = "non-contiguous layer inputs" if why is not None: return self._skip(why) - assert swa is not None and comp is not None + assert swa is not None runner = _mono_runner(x.device) logger.info_once( "DSv4.1 mono decode layer active (decode steps of <= %d rows)", MAX_ROWS ) - out = runner.forward( + # K1's seam folds the owed post into the residual, except where the + # residual arrives settled: the embedding, or after the Engram write + fold = True + if first: + residual = x.unsqueeze(1).expand(-1, layer.hc_mult, -1).contiguous() + pre_mix = runner.first_pre_mix(x.shape[0]) + post_mix = res_mix = pre_mix # unread without the fold + fold = False + elif layer.engram is not None and engram_hashes is not None: + residual = layer.mhc_post(x, residual, post_mix, res_mix) + hashes = engram_hashes[:, layer.engram.layer_hash_index] + residual = layer.engram(residual, hashes, engram_mask) + fold = False + kw: dict[str, Any] = {} + if attn.indexer is not None: + kw["select"] = lambda normed, qr: _select(layer, positions, normed, qr) + if attn.compress_ratio: + assert comp is not None + kw.update( + topk_indices=attn.topk_indices_buffer, + comp_cache=_record_view( + attn._compressed_kv_cache(), comp.block_size // attn.compress_ratio + ), + comp_block_table=comp.block_table, + ) + return runner.forward( self.weights(layer), x, residual, @@ -342,13 +405,9 @@ def __call__( swa.decode_swa_indices, swa.decode_swa_lens, swa.token_to_req_indices, - topk_indices=attn.topk_indices_buffer, - comp_cache=_record_view( - attn._compressed_kv_cache(), comp.block_size // attn.compress_ratio - ), - comp_block_table=comp.block_table, + fold=fold, + **kw, ) - return out def ffn( self, From 69d2b6f82dff3bdd78ec8658496e59c227b569d2 Mon Sep 17 00:00:00 2001 From: xiangdong <40376367+zxd1997066@users.noreply.github.com> Date: Tue, 13 Oct 2026 10:18:25 +0800 Subject: [PATCH 347/624] [CI][XPU] Standardize test job labels for intel jobs (#51755) Signed-off-by: zengxian Signed-off-by: xiangdong <40376367+zxd1997066@users.noreply.github.com> Co-authored-by: Kunshang Ji --- .../intel_xpu_ci/test-intel.yaml | 14 +++++------ .../intel_jobs/basic_correctness_intel.yaml | 4 ++-- .buildkite/intel_jobs/benchmarks_intel.yaml | 4 ++-- .buildkite/intel_jobs/engine_intel.yaml | 6 ++--- .buildkite/intel_jobs/entrypoints_intel.yaml | 2 +- .../intel_jobs/expert_parallelism_intel.yaml | 4 ++-- .buildkite/intel_jobs/kernels_intel.yaml | 4 ++-- .buildkite/intel_jobs/lm_eval_intel.yaml | 6 ++--- .buildkite/intel_jobs/lora_intel.yaml | 14 +++++------ .buildkite/intel_jobs/misc_intel.yaml | 24 +++++++++---------- .../intel_jobs/model_executor_intel.yaml | 4 ++-- .../intel_jobs/models_distributed_intel.yaml | 4 ++-- .../intel_jobs/models_multimodal_intel.yaml | 14 +++++------ ...ntization.yaml => quantization_intel.yaml} | 4 ++-- .buildkite/intel_jobs/samplers_intel.yaml | 4 ++-- .buildkite/intel_jobs/test-intel.yaml | 24 +++++++++---------- 16 files changed, 68 insertions(+), 68 deletions(-) rename .buildkite/intel_jobs/{quantization.yaml => quantization_intel.yaml} (95%) diff --git a/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml b/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml index a4a5ccbabce8..ed98d4408512 100644 --- a/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml +++ b/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml @@ -15,7 +15,7 @@ steps: limit: 2 - exit_status: -10 # Agent was lost limit: 2 - - label: "XPU example Test" + - label: ":intel: (BMG) Example Test" key: xpu-example-test depends_on: - image-build-xpu @@ -37,7 +37,7 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh example' - - label: "XPU V1 test" + - label: ":intel: (BMG) V1 test" key: xpu-v1-test depends_on: - image-build-xpu @@ -60,7 +60,7 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh v1' parallelism: 2 - - label: "XPU W8A8 FP8 Linear Examples" + - label: ":intel: (BMG) W8A8 FP8 Linear Examples" key: xpu-w8a8-fp8-linear-examples depends_on: - image-build-xpu @@ -83,7 +83,7 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh w8a8-fp8-linear' - - label: "XPU server test" + - label: ":intel: (BMG) Server test" key: xpu-server-test depends_on: - image-build-xpu @@ -105,7 +105,7 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh server' - - label: "XPU quantization test" + - label: ":intel: (BMG) Quantization test" key: xpu-quantization-test depends_on: - image-build-xpu @@ -127,7 +127,7 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh quantization' - - label: "XPU compressed tensors FP8 test" + - label: ":intel: (BMG) Compressed tensors FP8 test" key: xpu-compressed-tensors-fp8-test depends_on: - image-build-xpu @@ -150,7 +150,7 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh compressed-tensors-fp8' - - label: "XPU Graph" + - label: ":intel: (BMG) Graph" key: xpu-graph depends_on: - image-build-xpu diff --git a/.buildkite/intel_jobs/basic_correctness_intel.yaml b/.buildkite/intel_jobs/basic_correctness_intel.yaml index 548b6435a254..0016023dc746 100644 --- a/.buildkite/intel_jobs/basic_correctness_intel.yaml +++ b/.buildkite/intel_jobs/basic_correctness_intel.yaml @@ -1,8 +1,8 @@ -group: Basic Correctness +group: Basic Correctness Intel depends_on: - image-build-xpu steps: -- label: XPU Sleep Mode +- label: ":intel: (BMG) Sleep Mode" key: xpu-sleep-mode timeout_in_minutes: 45 device: intel_gpu diff --git a/.buildkite/intel_jobs/benchmarks_intel.yaml b/.buildkite/intel_jobs/benchmarks_intel.yaml index 9ce373e4f42e..f6c1e104b1db 100644 --- a/.buildkite/intel_jobs/benchmarks_intel.yaml +++ b/.buildkite/intel_jobs/benchmarks_intel.yaml @@ -1,8 +1,8 @@ -group: Benchmarks +group: Benchmarks Intel depends_on: - image-build-xpu steps: -- label: Benchmarks CLI Test +- label: ":intel: (BMG) Benchmarks CLI" key: benchmarks-cli-test timeout_in_minutes: 40 device: intel_gpu diff --git a/.buildkite/intel_jobs/engine_intel.yaml b/.buildkite/intel_jobs/engine_intel.yaml index fd028142cc0c..5235398c2579 100644 --- a/.buildkite/intel_jobs/engine_intel.yaml +++ b/.buildkite/intel_jobs/engine_intel.yaml @@ -2,7 +2,7 @@ group: Engine Intel depends_on: - image-build-xpu steps: -- label: Engine +- label: ":intel: (BMG) Engine" key: engine timeout_in_minutes: 40 device: intel_gpu @@ -40,7 +40,7 @@ steps: 'cd tests && pytest -v -s engine/test_arg_utils.py test_sequence.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py' -- label: Engine (1 GPU) +- label: ":intel: (BMG) Engine (1 GPU)" key: engine-1-gpu timeout_in_minutes: 30 device: intel_gpu @@ -65,7 +65,7 @@ steps: pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py --ignore v1/engine/test_core_engine_actor_manager.py && pytest -v -s test_config.py' -- label: V1 e2e (2 GPUs) +- label: ":intel: (BMG) V1 e2e (2 GPUs)" key: v1-e2e-2-gpus timeout_in_minutes: 30 device: intel_gpu diff --git a/.buildkite/intel_jobs/entrypoints_intel.yaml b/.buildkite/intel_jobs/entrypoints_intel.yaml index 96cbd2002a2f..f49f051ebc43 100644 --- a/.buildkite/intel_jobs/entrypoints_intel.yaml +++ b/.buildkite/intel_jobs/entrypoints_intel.yaml @@ -1,4 +1,4 @@ -group: Entrypoints +group: Entrypoints Intel depends_on: - image-build-xpu steps: diff --git a/.buildkite/intel_jobs/expert_parallelism_intel.yaml b/.buildkite/intel_jobs/expert_parallelism_intel.yaml index 0c57b7de1253..c0c95fb3a93a 100644 --- a/.buildkite/intel_jobs/expert_parallelism_intel.yaml +++ b/.buildkite/intel_jobs/expert_parallelism_intel.yaml @@ -1,8 +1,8 @@ -group: Expert Parallelism +group: Expert Parallelism Intel depends_on: - image-build-xpu steps: -- label: EPLB Algorithm +- label: ":intel: (BMG) EPLB Algorithm" key: eplb-algorithm timeout_in_minutes: 45 device: intel_gpu diff --git a/.buildkite/intel_jobs/kernels_intel.yaml b/.buildkite/intel_jobs/kernels_intel.yaml index 3650a729d7d7..f3a9d58f4357 100644 --- a/.buildkite/intel_jobs/kernels_intel.yaml +++ b/.buildkite/intel_jobs/kernels_intel.yaml @@ -2,7 +2,7 @@ group: Kernels Intel depends_on: - image-build-xpu steps: -- label: vLLM IR Tests +- label: ":intel: (BMG) vLLM IR Tests" key: vllm-ir-tests timeout_in_minutes: 30 device: intel_gpu @@ -26,7 +26,7 @@ steps: pytest -v -s ir && pytest -v -s kernels/ir' -- label: UVA Kernel Tests +- label: ":intel: (BMG) UVA Kernel Tests" key: uva-kernel-tests timeout_in_minutes: 20 device: intel_gpu diff --git a/.buildkite/intel_jobs/lm_eval_intel.yaml b/.buildkite/intel_jobs/lm_eval_intel.yaml index 62d99cf769b0..1919e4c85041 100644 --- a/.buildkite/intel_jobs/lm_eval_intel.yaml +++ b/.buildkite/intel_jobs/lm_eval_intel.yaml @@ -1,8 +1,8 @@ -group: LM Eval +group: LM Eval Intel depends_on: - image-build-xpu steps: -- label: "XPU GPQA Eval (GPT-OSS)" +- label: ":intel: (BMG) GPQA Eval (GPT-OSS)" key: xpu-gpqa-eval-gpt-oss timeout_in_minutes: 60 device: intel_gpu @@ -25,7 +25,7 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-xpu.txt' -- label: LM Eval Large Models(2 GPUs) +- label: ":intel: (BMG) LM Eval Large Models(2 GPUs)" key: lm-eval-large-models-2-gpus timeout_in_minutes: 45 device: intel_gpu diff --git a/.buildkite/intel_jobs/lora_intel.yaml b/.buildkite/intel_jobs/lora_intel.yaml index 885e4b546691..ae5e2d0165cc 100644 --- a/.buildkite/intel_jobs/lora_intel.yaml +++ b/.buildkite/intel_jobs/lora_intel.yaml @@ -2,7 +2,7 @@ group: LoRA Intel depends_on: - image-build-xpu steps: -- label: LoRA Runtime + Utils +- label: ":intel: (BMG) LoRA Runtime + Utils" key: lora-runtime-utils timeout_in_minutes: 45 device: intel_gpu @@ -35,7 +35,7 @@ steps: pytest -v -s lora/test_utils.py && pytest -v -s lora/test_worker.py' -- label: LoRA Fused/MoE Kernels +- label: ":intel: (BMG) LoRA Fused/MoE Kernels" key: lora-fused-moe-kernels timeout_in_minutes: 45 device: intel_gpu @@ -60,7 +60,7 @@ steps: pytest -v -s lora/test_fused_moe_lora_kernel.py && pytest -v -s lora/test_moe_lora_align_sum.py --deselect="tests/lora/test_moe_lora_align_sum.py::test_moe_lora_align_block_size_mixed_base_and_lora[1]"' -- label: LoRA Punica Kernels +- label: ":intel: (BMG) LoRA Punica Kernels" key: lora-punica-kernels timeout_in_minutes: 45 device: intel_gpu @@ -87,7 +87,7 @@ steps: pytest -v -s lora/test_punica_ops.py::test_kernels_hidden_size && pytest -v -s lora/test_punica_ops.py::test_add_lora_fused_moe_early_exit' -- label: LoRA Punica FP8/XPU Ops +- label: ":intel: (BMG) LoRA Punica FP8/XPU Ops" key: lora-punica-fp8-xpu-ops timeout_in_minutes: 60 device: intel_gpu @@ -112,7 +112,7 @@ steps: pytest -v -s lora/test_punica_ops_fp8.py && pytest -v -s lora/test_punica_xpu_ops.py' -- label: LoRA Models +- label: ":intel: (BMG) LoRA Models" key: lora-models timeout_in_minutes: 45 device: intel_gpu @@ -139,7 +139,7 @@ steps: pytest -v -s lora/test_chatglm3_tp.py && pytest -v -s lora/test_llama_tp.py::test_llama_lora' -- label: LoRA Multimodal +- label: ":intel: (BMG) LoRA Multimodal" key: lora-multimodal timeout_in_minutes: 45 device: intel_gpu @@ -163,7 +163,7 @@ steps: export VLLM_WORKER_MULTIPROC_METHOD=spawn && pytest -v -s lora/test_whisper.py' -- label: LoRA Add + Multimodal +- label: ":intel: (BMG) LoRA Add + Multimodal" key: lora-add-multimodal timeout_in_minutes: 45 device: intel_gpu diff --git a/.buildkite/intel_jobs/misc_intel.yaml b/.buildkite/intel_jobs/misc_intel.yaml index b7ae9f907846..c1706d70f6fc 100644 --- a/.buildkite/intel_jobs/misc_intel.yaml +++ b/.buildkite/intel_jobs/misc_intel.yaml @@ -2,7 +2,7 @@ group: Miscellaneous Intel depends_on: - image-build-xpu steps: -- label: XPU Batch Invariance +- label: ":intel: (BMG) Batch Invariance" key: xpu-batch-invariance timeout_in_minutes: 60 num_devices: 2 @@ -30,7 +30,7 @@ steps: pytest -v -s v1/determinism/test_xpu_batch_invariant_ut.py && bash ../.buildkite/scripts/hardware_ci/run-xpu-batch-invariance.sh' -- label: V1 Core + KV + Metrics +- label: ":intel: (BMG) V1 Core + KV + Metrics" key: v1-core-kv-metrics-intel timeout_in_minutes: 45 device: intel_gpu @@ -64,7 +64,7 @@ steps: VLLM_BATCH_INVARIANT=1 pytest -v -s -m "not cpu_test" v1/ec_connector/unit --ignore=v1/ec_connector/unit/test_ec_cpu_connector.py' -- label: EC CPU Connector (XPU) +- label: ":intel: (BMG) EC CPU Connector (XPU)" key: ec-cpu-connector-xpu timeout_in_minutes: 45 device: intel_gpu @@ -96,7 +96,7 @@ steps: cd tests && pytest -v -s v1/ec_connector/unit/test_ec_cpu_connector.py' -- label: V1 Sample + Logits +- label: ":intel: (BMG) V1 Sample + Logits" key: v1-sample-logits-intel timeout_in_minutes: 90 device: intel_gpu @@ -141,7 +141,7 @@ steps: limit: 1 parallelism: 4 -- label: Basic Models Tests (Initialization) +- label: ":intel: (BMG) Basic Models Tests (Initialization)" key: basic-models-tests-initialization timeout_in_minutes: 60 device: intel_gpu @@ -166,7 +166,7 @@ steps: cd tests && pytest -v -s models/test_initialization.py::test_can_initialize_large_subset[Eagle3MiniMaxM2ForCausalLM]' -- label: XPU CPU Offload +- label: ":intel: (BMG) CPU Offload" key: xpu-cpu-offload timeout_in_minutes: 60 device: intel_gpu @@ -194,7 +194,7 @@ steps: pytest -v -s v1/kv_offload && pytest -v -s v1/kv_connector/unit/test_offloading_connector.py' -- label: NixlConnector PD accuracy (4 GPUs) +- label: ":intel: (BMG) NixlConnector PD accuracy (4 GPUs)" key: nixlconnector-pd-accuracy-4-gpus timeout_in_minutes: 60 num_devices: 4 @@ -223,7 +223,7 @@ steps: PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh && PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh' -- label: Regression +- label: ":intel: (BMG) Regression" key: regression timeout_in_minutes: 50 device: intel_gpu @@ -256,7 +256,7 @@ steps: 'cd tests && pytest -v -s test_regression.py' -- label: Metrics, Tracing (2 GPUs) +- label: ":intel: (BMG) Metrics, Tracing (2 GPUs)" key: metrics-tracing-2-gpus timeout_in_minutes: 45 num_devices: 2 @@ -291,7 +291,7 @@ steps: 'cd tests && pytest -v -s v1/tracing' -- label: Async Engine, Inputs, Utils, Worker +- label: ":intel: (BMG) Async Engine, Inputs, Utils, Worker" key: async-engine-inputs-utils-worker timeout_in_minutes: 55 device: intel_gpu @@ -330,7 +330,7 @@ steps: pytest -v -s -m "not cpu_test" ./multimodal && pytest -v -s utils_ --ignore=utils_/test_mem_utils.py' -- label: Detokenizer Stop Strings (XPU) +- label: ":intel: (BMG) Detokenizer Stop Strings (XPU)" key: detokenizer-stop-strings-xpu timeout_in_minutes: 45 device: intel_gpu @@ -367,7 +367,7 @@ steps: 'cd tests && pytest -v -s detokenizer/test_stop_strings.py' -- label: Fusion Unit Tests +- label: ":intel: (BMG) Fusion Unit Tests" key: fusion-unit-tests timeout_in_minutes: 30 device: intel_gpu diff --git a/.buildkite/intel_jobs/model_executor_intel.yaml b/.buildkite/intel_jobs/model_executor_intel.yaml index d641c93e5a9e..cac6c931f09a 100644 --- a/.buildkite/intel_jobs/model_executor_intel.yaml +++ b/.buildkite/intel_jobs/model_executor_intel.yaml @@ -2,7 +2,7 @@ group: Model Executor Intel depends_on: - image-build-xpu steps: -- label: Model Executor (Intel) +- label: ":intel: (BMG) Model Executor" key: model-executor-intel timeout_in_minutes: 45 device: intel_gpu @@ -30,7 +30,7 @@ steps: cd tests && pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --ignore="model_executor/model_loader/tensorizer_loader/test_tensorizer.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload"' -- label: Tensorizer Model Loader Tests (XPU) +- label: ":intel: (BMG) Tensorizer Model Loader Tests (XPU)" key: tensorizer-model-loader-tests-xpu timeout_in_minutes: 45 device: intel_gpu diff --git a/.buildkite/intel_jobs/models_distributed_intel.yaml b/.buildkite/intel_jobs/models_distributed_intel.yaml index 5f2247eeaace..5fd723f8c5f0 100644 --- a/.buildkite/intel_jobs/models_distributed_intel.yaml +++ b/.buildkite/intel_jobs/models_distributed_intel.yaml @@ -1,8 +1,8 @@ -group: Models - Distributed +group: Models - Distributed Intel depends_on: - image-build-xpu steps: -- label: Distributed Model Tests (2 GPUs) +- label: ":intel: (BMG) Distributed Model Tests" key: distributed-model-tests-2-gpus timeout_in_minutes: 65 device: intel_gpu diff --git a/.buildkite/intel_jobs/models_multimodal_intel.yaml b/.buildkite/intel_jobs/models_multimodal_intel.yaml index 89654ec3aedf..9a9790d5adc7 100644 --- a/.buildkite/intel_jobs/models_multimodal_intel.yaml +++ b/.buildkite/intel_jobs/models_multimodal_intel.yaml @@ -1,8 +1,8 @@ -group: Models - Multimodal +group: Models - Multimodal Intel depends_on: - image-build-xpu steps: -- label: "Multi-Modal Models (Standard) 1: qwen2" +- label: ":intel: (BMG) Multi-Modal Models (Standard) 1: qwen2" key: multi-modal-models-standard-1-qwen2 timeout_in_minutes: 70 device: intel_gpu @@ -26,7 +26,7 @@ steps: pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" && pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model' -- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" +- label: ":intel: (BMG) Multi-Modal Models (Standard) 2: qwen3 + gemma" key: multi-modal-models-standard-2-qwen3-gemma timeout_in_minutes: 90 device: intel_gpu @@ -51,7 +51,7 @@ steps: pytest -v -s models/multimodal/generation/test_mm_prefix_lm.py -m core_model && pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model' -- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl" +- label: ":intel: (BMG) Multi-Modal Models (Standard) 3: llava + qwen2_vl" key: multi-modal-models-standard-3-llava-qwen2-vl timeout_in_minutes: 65 device: intel_gpu @@ -76,7 +76,7 @@ steps: pytest -v -s models/multimodal/generation/test_gemma_ragged_audio.py -m core_model pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model' -- label: "Multi-Modal Models (Standard) 4: other + whisper" +- label: ":intel: (BMG) Multi-Modal Models (Standard) 4: other + whisper" key: multi-modal-models-standard-4-other-whisper timeout_in_minutes: 45 device: intel_gpu @@ -99,7 +99,7 @@ steps: 'cd tests && pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_gemma_ragged_audio.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing' -- label: Multi-Modal Processor # 44min +- label: ":intel: (BMG) Multi-Modal Processor" # 44min key: multi-modal-processor timeout_in_minutes: 60 device: intel_gpu @@ -126,7 +126,7 @@ steps: --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB' parallelism: 4 -- label: Multi-Modal Models (Extended Generation 1) +- label: ":intel: (BMG) Multi-Modal Models (Extended Generation 1)" key: multi-modal-models-extended-generation-1 timeout_in_minutes: 60 device: intel_gpu diff --git a/.buildkite/intel_jobs/quantization.yaml b/.buildkite/intel_jobs/quantization_intel.yaml similarity index 95% rename from .buildkite/intel_jobs/quantization.yaml rename to .buildkite/intel_jobs/quantization_intel.yaml index af1b23b8a004..035ad11beef2 100644 --- a/.buildkite/intel_jobs/quantization.yaml +++ b/.buildkite/intel_jobs/quantization_intel.yaml @@ -1,8 +1,8 @@ -group: Quantization +group: Quantization Intel depends_on: - image-build-xpu steps: -- label: Quantization +- label: ":intel: (BMG) Quantization" key: quantization timeout_in_minutes: 30 env: diff --git a/.buildkite/intel_jobs/samplers_intel.yaml b/.buildkite/intel_jobs/samplers_intel.yaml index 027237225d3a..ed55c49da890 100644 --- a/.buildkite/intel_jobs/samplers_intel.yaml +++ b/.buildkite/intel_jobs/samplers_intel.yaml @@ -2,7 +2,7 @@ group: Samplers Intel depends_on: - image-build-xpu steps: -- label: Samplers Test (XPU Kernel) +- label: ":intel: (BMG) Samplers Test (XPU Kernel)" key: samplers-test-xpu-kernel-intel timeout_in_minutes: 40 device: intel_gpu @@ -29,7 +29,7 @@ steps: VLLM_XPU_USE_SAMPLER_KERNEL=1 pytest -v -s samplers --ignore=samplers/test_beam_search.py' -- label: Samplers Beam Search Test (XPU Kernel) +- label: ":intel: (BMG) Samplers Beam Search Test (XPU Kernel)" key: samplers-beam-search-test-xpu-kernel-intel timeout_in_minutes: 40 device: intel_gpu diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 1567c0eba724..74d6eee4900e 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -14,7 +14,7 @@ steps: limit: 2 - exit_status: -10 # Agent was lost limit: 2 - - label: "XPU example Test" + - label: ":intel: (BMG) Example Test" key: xpu-example-test depends_on: - image-build-xpu @@ -30,7 +30,7 @@ steps: REPO: "vllm-ci-test-repo" source_file_dependencies: - vllm/ - - .buildkite/intel_jobs/test-intel.yaml + - .buildkite/intel_jobs/test-intel.yaml commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh @@ -48,7 +48,7 @@ steps: python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192 && python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-LLMC -tp 2 --max-model-len 8192 ' - - label: "XPU W8A8 FP8 Linear Examples" + - label: ":intel: (BMG) W8A8 FP8 Linear Examples" key: xpu-w8a8-fp8-linear-examples depends_on: - image-build-xpu @@ -76,7 +76,7 @@ steps: python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model meta-llama/Llama-3.2-1B-Instruct --quantization fp8_per_tensor --enforce-eager --max-model-len 4096 && python3 examples/basic/offline_inference/generate.py --linear-backend torch --model meta-llama/Llama-3.2-1B-Instruct --quantization fp8_per_tensor --enforce-eager --max-model-len 4096 ' - - label: "XPU V1 test" + - label: ":intel: (BMG) V1 test" key: xpu-v1-test depends_on: - image-build-xpu @@ -92,7 +92,7 @@ steps: REPO: "vllm-ci-test-repo" source_file_dependencies: - vllm/ - - .buildkite/intel_jobs/test-intel.yaml + - .buildkite/intel_jobs/test-intel.yaml commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh @@ -107,7 +107,7 @@ steps: 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 --deselect=tests/v1/spec_decode/test_mtp.py::test_glm_mtp_defers_lm_head && pytest -v -s v1/kv_connector/unit --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 --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --deselect=tests/v1/kv_connector/unit/test_nixl_connector.py::test_abort_timeout_on_prefiller[ray]' parallelism: 2 - - label: "XPU Qwen3-30B + V1 KV Connector" + - label: ":intel: (BMG) Qwen3-30B + V1 KV Connector" key: xpu-qwen3-30b-v1-kv-connector depends_on: - image-build-xpu @@ -133,7 +133,7 @@ steps: pytest -v -s v1/kv_connector/unit/test_example_connector.py && cd .. && VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192' - - label: "XPU server test" + - label: ":intel: (BMG) Server test" key: xpu-server-test depends_on: - image-build-xpu @@ -149,14 +149,14 @@ steps: REPO: "vllm-ci-test-repo" source_file_dependencies: - vllm/ - - .buildkite/intel_jobs/test-intel.yaml + - .buildkite/intel_jobs/test-intel.yaml commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pytest -v -s entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py && pytest -v -s benchmarks/test_serve_cli.py' - - label: "XPU quantization test" + - label: ":intel: (BMG) Quantization test" key: xpu-quantization-test depends_on: - image-build-xpu @@ -172,14 +172,14 @@ steps: REPO: "vllm-ci-test-repo" source_file_dependencies: - vllm/ - - .buildkite/intel_jobs/test-intel.yaml + - .buildkite/intel_jobs/test-intel.yaml commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pytest -v -s quantization/test_auto_round.py && pytest -v -s quantization/test_online.py --deselect=tests/quantization/test_online.py::test_online_quantization_loads_real_weights --deselect="tests/quantization/test_online.py::test_online_quantization[False-requantization_mxfp8_ptcp_fp8]"' - - label: "XPU compressed tensors FP8 test" + - label: ":intel: (BMG) Compressed tensors FP8 test" key: xpu-compressed-tensors-fp8-test depends_on: - image-build-xpu @@ -204,7 +204,7 @@ steps: 'cd tests && pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8 && pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8_block_enabled' - - label: "XPU Graph" + - label: ":intel: (BMG) Graph" key: xpu-graph depends_on: - image-build-xpu From c077dc663d929fd19bdfb15ccf00f83cc490b8b0 Mon Sep 17 00:00:00 2001 From: frida-andersson Date: Mon, 12 Oct 2026 21:19:14 -0500 Subject: [PATCH 348/624] [ROCm][Perf][DSv4.1] Enable decoder SWA bounded replay (#59039) Signed-off-by: Frida Andersson --- .buildkite/test_areas/models_basic.yaml | 4 +- .../test_deepseek_v41_amd_replay_forward.py | 329 ++++++++++++++++++ .../models/test_deepseek_v41_replay_batch.py | 75 ++++ vllm/models/deepseek_v41/amd/model.py | 291 ++++++++++++++-- vllm/models/deepseek_v41/amd/vl_model.py | 5 + vllm/models/deepseek_v41/attention.py | 10 +- .../models/deepseek_v41/nvidia/model_state.py | 34 +- 7 files changed, 705 insertions(+), 43 deletions(-) create mode 100644 tests/models/test_deepseek_v41_amd_replay_forward.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 552910e85a96..045e04533a44 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -82,18 +82,20 @@ steps: - image-build-amd source_file_dependencies: - vllm/models/deepseek_v4/ + - vllm/models/deepseek_v41/ - tests/models/test_deepseek_v4_vl_rocm.py - tests/models/test_hyv4_rocm.py - tests/models/transformers/test_layer_registry_e2e.py - tests/models/test_deepseek_v41_replay_start.py - tests/models/test_deepseek_v41_decoder_replay_layers.py - tests/models/test_deepseek_v41_replay_batch.py + - tests/models/test_deepseek_v41_amd_replay_forward.py - tests/models/test_deepseek_v4_dspark_rocm.py - tests/models/test_deepseek_v4_rocm_wo_a.py - tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py - tests/models/test_dsv41_mono_decode.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_amd_replay_forward.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.py - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - label: ":nvidia: (B200) Inkling" diff --git a/tests/models/test_deepseek_v41_amd_replay_forward.py b/tests/models/test_deepseek_v41_amd_replay_forward.py new file mode 100644 index 000000000000..fe22bc80253f --- /dev/null +++ b/tests/models/test_deepseek_v41_amd_replay_forward.py @@ -0,0 +1,329 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""What the AMD replay forward does with a prefill longer than the window. + +The cut in this fixture is layer 1, so replay starts at layer 1. That is the +same boundary as production layer 20, on a 4-layer stand-in. The final +``mhc_post`` and the aux capture at that cut both run on the window rows. +""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from vllm.forward_context import ForwardContext, override_forward_context +from vllm.models.deepseek_v41.decoder_replay_layers import ReplayBatch +from vllm.sequence import IntermediateTensors + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +WINDOW = 128 +FULL = 300 +# Trailing window of a 300-token prefill. +ROWS = list(range(FULL - WINDOW, FULL)) +HC, HIDDEN = 4, 8 +DEVICE = torch.device("cuda") +# gfx942 writes the V4 record (584 B). Family-100 writes the V4.1 record (528 B). +# The SWA layer uses block size 32, which is the size the KV-only insert checks. +_BLOCK = 32 +_HEAD_DIM = 512 + + +class _RocmKvInsert: + """SWA cache from ``forward_kv``, checked against the KV-only insert.""" + + def __init__(self, num_tokens: int): + from vllm.models.deepseek_v41.attention import _use_v41_mxfp8_kv_record + + if not hasattr(torch.ops._C, "fused_deepseek_v4_kv_rope_insert"): + pytest.skip("fused DeepseekV4 KV insert op is not built") + + self.num_tokens = num_tokens + self.kv_mxfp8 = _use_v41_mxfp8_kv_record() + bytes_per = 528 if self.kv_mxfp8 else 584 + num_blocks = (num_tokens + _BLOCK - 1) // _BLOCK + 1 + cache = torch.full( + (num_blocks, _BLOCK, bytes_per), + 0xA5, + dtype=torch.uint8, + device=DEVICE, + ) + self.poison = cache.clone() + self.kv = torch.randn( + num_tokens, _HEAD_DIM, dtype=torch.bfloat16, device=DEVICE + ) + prefix = "replay.swa" + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=DEVICE) + self.metadata = { + prefix: SimpleNamespace(slot_mapping=slot_mapping, block_size=_BLOCK) + } + self.context = ForwardContext({}, self.metadata, {}) + inv_freq = 1.0 / ( + 10000 ** (torch.arange(0, 64, 2, dtype=torch.float32, device=DEVICE) / 64) + ) + freqs = torch.outer( + torch.arange(num_tokens + 8, dtype=torch.float32, device=DEVICE), inv_freq + ) + self.attn = SimpleNamespace( + kv_mxfp8=self.kv_mxfp8, + compressor=None, + indexer=None, + rotary_emb=SimpleNamespace( + cos_sin_cache=torch.cat((freqs.cos(), freqs.sin()), dim=-1) + ), + swa_cache_layer=SimpleNamespace( + prefix=prefix, kv_cache=cache, block_size=_BLOCK + ), + _run_parallel_input_projections=lambda hidden: (hidden, None, None), + _split_qkv_and_norm=lambda qr_kv: (None, None, self.kv[: qr_kv.shape[0]]), + ) + self.last_rows: int | None = None + self.last_positions: torch.Tensor | None = None + + def write( + self, + x, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + ): + from vllm.models.deepseek_v41.amd.rocm import ( + DeepseekV41ROCMAiterMLAAttention, + ) + from vllm.models.deepseek_v41.attention import DeepseekV4Attention + + assert ( + DeepseekV41ROCMAiterMLAAttention.forward_kv + is DeepseekV4Attention.forward_kv + ) + self.last_rows = x.shape[0] + self.last_positions = positions + DeepseekV41ROCMAiterMLAAttention.forward_kv(self.attn, positions, x) + + def assert_matches_forward(self): + assert self.last_rows == self.num_tokens + assert self.last_positions is not None + written = self.attn.swa_cache_layer.kv_cache + assert not torch.equal(written, self.poison) + ref = self.poison.clone() + meta = self.metadata[self.attn.swa_cache_layer.prefix] + torch.ops._C.fused_deepseek_v4_kv_rope_insert( + self.kv, + ref, + meta.slot_mapping, + self.last_positions, + self.attn.rotary_emb.cos_sin_cache, + _BLOCK, + None, + self.kv_mxfp8, + ) + assert torch.equal(written, ref) + assert torch.equal(written[-1], self.poison[-1]) + + +class _Layer(nn.Module): + def __init__( + self, + idx: int, + compressed: str | None = None, + indexer: str | None = None, + ): + super().__init__() + self.idx = idx + self.seen: list[torch.Tensor] = [] + self.attn = SimpleNamespace( + swa_cache_layer=SimpleNamespace(prefix=f"layers.{idx}.swa"), + compressed_cache_prefix=compressed, + indexer=( + None + if indexer is None + else SimpleNamespace(k_cache=SimpleNamespace(prefix=indexer)) + ), + compress_ratio=1, + ) + + def forward( + self, + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + engram_hashes, + engram_mask, + ): + self.seen.append(hidden_states[:, 0, 0].detach().clone()) + if residual is None: + residual = hidden_states + if pre_mix is None: + pre_mix = hidden_states[:, :, 0] + if post_mix is None: + post_mix = pre_mix + if res_mix is None: + res_mix = pre_mix + return hidden_states, residual, post_mix, res_mix, pre_mix + + +def _model(monkeypatch): + import vllm.models.deepseek_v41.amd.model as amd + + monkeypatch.setattr( + amd, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=False, is_last_rank=False), + ) + monkeypatch.setattr(amd.EngramLayout, "from_config", lambda config: None) + monkeypatch.setattr(amd, "MHCPostOp", lambda: SimpleNamespace()) + + # Layer 0 sits before the cut, so its keys are not collected. Layer 2 + # shares layer 1's compressed cache. Layer 3 has neither. + layers = [ + _Layer(0, "layers.0.compressed", "layers.0.indexer"), + _Layer(1, "layers.1.compressed", "layers.1.indexer"), + _Layer(2, "layers.1.compressed", "layers.2.indexer"), + _Layer(3), + ] + # Layer 1 is the KV source, so the replay callback is this layer's write. + insert = _RocmKvInsert(FULL) + layers[1].write_kv = insert.write + monkeypatch.setattr(amd, "make_layers", lambda n, factory, prefix: (0, n, layers)) + + config = SimpleNamespace( + vocab_size=32, + hc_eps=1e-6, + hc_mult=HC, + hidden_size=HIDDEN, + rms_norm_eps=1e-6, + index_topk=4, + kv_source_layer_ids=[1], + num_hidden_layers=4, + sliding_window=WINDOW, + engram_layer_ids=(), + candidate_source_layer_id=-1, + candidate_topk_blocks=0, + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=config, dtype=torch.bfloat16), + quant_config=None, + cache_config=SimpleNamespace(swa_bounded_replay=True), + use_v2_model_runner=True, + parallel_config=SimpleNamespace( + prefill_context_parallel_size=1, + use_ubatching=False, + enable_expert_parallel=False, + ), + speculative_config=None, + scheduler_config=SimpleNamespace(max_num_batched_tokens=FULL), + ) + model = amd.DeepseekV4Model(vllm_config=vllm_config) + # Production requests aux at the first replay layer. Here that id is the cut. + model.aux_hidden_state_layers = (model.decoder_replay_start,) + model.topk_indices_buffer = model.topk_indices_buffer.to(DEVICE) + + calls: list[torch.Tensor] = [] + + def mhc_post(hidden, residual, post_mix, res_mix): + calls.append(hidden[:, 0, 0].detach().clone()) + return hidden + + model.mhc_post = mhc_post + saved_aux: list[list[torch.Tensor]] = [] + run_layers = model._run_layers + + def wrapped(*args, **kwargs): + result = run_layers(*args, **kwargs) + # Snapshot: forward() extends the early list with the scattered aux. + saved_aux.append(list(result[5])) + return result + + model._run_layers = wrapped + scattered: list[tuple[torch.Tensor, ...]] = [] + replay = model.decoder_replay_layers + run = replay._run + + def wrapped_run(*args, **kwargs): + result = run(*args, **kwargs) + scattered.append(result) + return result + + replay._run = wrapped_run + return model, layers, calls, saved_aux, scattered, insert + + +def _forward(model, insert: _RocmKvInsert): + token_ids = torch.arange(FULL, device=DEVICE, dtype=torch.float32) + hidden = token_ids[:, None, None].expand(FULL, HC, HIDDEN).contiguous() + pre_mix = torch.ones(FULL, HC, device=DEVICE) + rows = torch.tensor(ROWS, device=DEVICE) + model.decoder_replay_layers.replay_batch = ReplayBatch( + rows, + ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={}), + ) + positions = torch.arange(FULL, device=DEVICE) + with override_forward_context(insert.context): + out = model( + input_ids=positions.to(torch.int32), + positions=positions, + intermediate_tensors=IntermediateTensors( + {"hidden_states": hidden, "pre_mix": pre_mix} + ), + ) + return out + + +def test_amd_replay_forward_trims_to_the_window(monkeypatch): + model, layers, calls, saved_aux, scattered, insert = _model(monkeypatch) + replay_layers = model.decoder_replay_layers + assert model.decoder_replay_start == 1 + assert replay_layers.first_swa_prefix == "layers.1.swa" + assert replay_layers.metadata_prefixes == { + "layers.1.swa", + "layers.1.compressed", + "layers.1.indexer", + "layers.2.swa", + "layers.2.indexer", + "layers.3.swa", + } + assert replay_layers.replay_batch is None + out = _forward(model, insert) + insert.assert_matches_forward() + replay = model.decoder_replay_layers.replay_batch + assert replay is not None and replay.rows.tolist() == ROWS + + early = [layer.seen[0] for layer in layers[: model.decoder_replay_start]] + late = [layer.seen[0] for layer in layers[model.decoder_replay_start :]] + full_ids = torch.arange(FULL, device=DEVICE, dtype=torch.float32) + window_ids = full_ids[ROWS] + assert [t.tolist() for t in early] == [full_ids.tolist()] + assert [t.tolist() for t in late] == [ + window_ids.tolist(), + window_ids.tolist(), + window_ids.tolist(), + ] + + # Scatter writes the replay rows back and leaves the trimmed rows at zero. + hidden = out["hidden_states"][:, 0, 0] + assert torch.equal(hidden[ROWS], window_ids) + assert hidden[: FULL - WINDOW].eq(0).all() + + # calls[0] is the cut-aux capture. calls[1] is the final mhc_post. + # Both run inside the replay batch, before the scatter back to FULL. + assert len(calls) == 2 + assert [t.shape[0] for t in calls] == [WINDOW, WINDOW] + assert torch.equal(calls[0], window_ids) + assert torch.equal(calls[1], window_ids) + + # The early loop skips this id. The replay loop does not see it either. + assert saved_aux[0] == [] + assert saved_aux[1] == [] + _hidden, _pre_mix, aux = scattered[0] + assert aux.shape[0] == FULL + assert torch.equal(aux[ROWS, 0], window_ids) + assert aux[: FULL - WINDOW].eq(0).all() diff --git a/tests/models/test_deepseek_v41_replay_batch.py b/tests/models/test_deepseek_v41_replay_batch.py index bcd4db64b25c..5a3625dcbf4e 100644 --- a/tests/models/test_deepseek_v41_replay_batch.py +++ b/tests/models/test_deepseek_v41_replay_batch.py @@ -290,6 +290,81 @@ def test_idle_dp_rank_dummy_trims_with_its_peers(dp_state): assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [1, 129] +def test_zero_prefix_window_warms_gather_and_skips_pad(monkeypatch): + """ROCm leaves the prefix-cache window at 0. Replay still needs replay_start + and a warmed gather kernel. The pad kernel stays on the prefix-cache path.""" + import vllm.models.deepseek_v41.nvidia.model_state as model_state + + calls = {"pad_warmup": 0, "gather_warmup": 0, "pad_launch": 0} + + def pad_warmup(*args, **kwargs): + calls["pad_warmup"] += 1 + + def gather_warmup(*args, **kwargs): + calls["gather_warmup"] += 1 + + class _Pad: + warmup = staticmethod(pad_warmup) + + def __getitem__(self, grid): + def launch(*args, **kwargs): + calls["pad_launch"] += 1 + + return launch + + monkeypatch.setattr(model_state, "_pad_replayed_slots_kernel", _Pad()) + monkeypatch.setattr( + model_state._gather_replay_batch_kernel, "warmup", gather_warmup + ) + + cfg = MagicMock() + cfg.model_config.enable_prompt_embeds = False + cfg.model_config.uses_mrope = False + cfg.model_config.is_multimodal_model = False + cfg.scheduler_config.max_num_seqs = 16 + cfg.scheduler_config.max_num_batched_tokens = 1024 + cfg.parallel_config.data_parallel_size = 1 + cfg.compilation_config.fast_moe_cold_start = False + cfg.compilation_config.cudagraph_mode = CUDAGraphMode.NONE + layers = SimpleNamespace(window=WINDOW, replay_batch=None, metadata_prefixes=set()) + model = SimpleNamespace(token_lookback_depth=0, decoder_replay_layers=layers) + groups = SimpleNamespace( + kv_cache_groups=[ + SimpleNamespace( + layer_names=["swa"], + kv_cache_spec=SimpleNamespace( + prefix_cacheable=False, prefix_replay_tokens=0 + ), + ), + SimpleNamespace( + layer_names=["mla"], + kv_cache_spec=SimpleNamespace( + prefix_cacheable=True, prefix_replay_tokens=0 + ), + ), + ] + ) + + def prepare_attn(self, input_batch, cg_mode, block_tables, slot_mappings, *a, **kw): + return {"swa": kw["model_specific_attn_metadata"]} + + monkeypatch.setattr(DefaultModelState, "prepare_attn", prepare_attn) + state = DeepseekV41ModelState(cfg, model, None, DEVICE) + state._replay_start_np[0] = 16 + batch = _input_batch([8], [24], [True]) + state.prepare_attn( + batch, + CUDAGraphMode.NONE, + BLOCK_TABLES, + _slot_mappings(batch.num_tokens_after_padding), + [], + groups, + ) + assert state._replay[0] == 0 + assert calls == {"pad_warmup": 0, "gather_warmup": 1, "pad_launch": 0} + assert state._replay_start[:1].tolist() == [16] + + def test_dp_ranks_pad_to_one_replay_graph(dp_state): """Every rank pads its replay batch to the graph fitting the largest rank's.""" dp_state.replay_cudagraphs = MagicMock() diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 4c5f15858749..a52a408c74dd 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -62,6 +62,7 @@ from vllm.models.deepseek_v41.amd.mono_decode import MonoDecodeLayer from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention from vllm.models.deepseek_v41.attention import DeepseekV4Attention +from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -421,6 +422,42 @@ def forward( x = self.ffn(x, input_ids) return x, residual, post_mix, res_mix, ffn_pre + def write_kv( + self, + x: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + pre_mix: torch.Tensor, + post_mix: torch.Tensor, + res_mix: torch.Tensor, + residual: torch.Tensor, + ) -> None: + """Write the KV ``forward`` would for the replay layers' inputs, skipping + the rest of the layer; the replay batch then reruns its rows.""" + # The replay batch rules out Engram and sequence parallel here. + assert self.engram is None and not self.use_sequence_parallel + fuse_attn_norm = self.fuse_seam_norm + residual, post_mix, res_mix, x, _attn_pre = self.mhc_pre_delayed( + residual, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + pre_mix=pre_mix, + sublayer_out=x, + post_layer_mix=post_mix, + comb_res_mix=res_mix, + norm_weight=self.attn_norm.weight if fuse_attn_norm else None, + norm_eps=self.attn_norm.variance_epsilon, + ) + if not fuse_attn_norm: + x = self.attn_norm(x) + self.attn.forward_kv(positions, x) + class DeepseekV4Model(nn.Module, EagleModelMixin): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -498,6 +535,45 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) + # Config flag is the switch: ROCm forces swa_cache_layer.bounded_replay off. + self.decoder_replay_layers: DecoderReplayLayers | None = None + self.decoder_replay_start = self.end_layer + cut = max(config.kv_source_layer_ids) + if ( + vllm_config.cache_config.swa_bounded_replay + and vllm_config.use_v2_model_runner + and cut < self.end_layer - 1 + and self._decoder_replay_supported(vllm_config, cut) + ): + assert self.layers[cut].attn.compress_ratio <= 1 + # The last KV source writes every row's KV, then replays. + self.decoder_replay_start = cut + # The attention metadata keys the replay layers read. + metadata_prefixes: set[str] = set() + for layer in islice(self.layers, cut, self.end_layer): + attn = typing.cast(DeepseekV4DecoderLayer, layer).attn + metadata_prefixes.add(attn.swa_cache_layer.prefix) + if attn.compressed_cache_prefix is not None: + metadata_prefixes.add(attn.compressed_cache_prefix) + if attn.indexer is not None: + metadata_prefixes.add(attn.indexer.k_cache.prefix) + self.decoder_replay_layers = DecoderReplayLayers( + config.sliding_window, + self._run_replay_layers, + self.layers[cut].write_kv, + metadata_prefixes, + self.layers[cut].attn.swa_cache_layer.prefix, + ) + logger.info_once( + "Decoder SWA bounded replay: in eager prefill steps, layer %d " + "writes its KV for every token, and the rest of it and layers " + "%d-%d run on each request's last %d tokens only.", + cut, + cut + 1, + self.end_layer - 1, + config.sliding_window, + ) + # The n-gram hash needs a slot-keyed rolling store of compressed ids # (chunked prefill / decode lookback); key it off the first local # layer's sliding-window KV cache. Only PP ranks owning an engram @@ -657,13 +733,33 @@ def forward( if not get_pp_group().is_first_rank: assert intermediate_tensors is not None pre_mix = intermediate_tensors["pre_mix"] - aux_hidden_states: list[torch.Tensor] = [] - final_aux_recon: torch.Tensor | None = None # avoid duplicate mhc_post call - for idx, layer in enumerate( - islice(self.layers, self.start_layer, self.end_layer), - start=self.start_layer, - ): - hidden_states, residual, post_mix, res_mix, pre_mix = layer( + ( + hidden_states, + residual, + post_mix, + res_mix, + pre_mix, + aux_hidden_states, + final_aux_recon, + ran_layers, + ) = self._run_layers( + range(self.start_layer, self.decoder_replay_start), + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + engram_hashes, + engram_mask, + full_num_tokens, + ) + # Replay collapses inside its own batch. The early loop still sets + # ran_layers, so that flag alone must not post again on the full batch. + collapsed = False + if self.decoder_replay_layers is not None: + hidden_states, pre_mix, *late_aux = self.decoder_replay_layers( hidden_states, positions, input_ids, @@ -671,20 +767,15 @@ def forward( post_mix, res_mix, residual, - engram_hashes, - engram_mask, ) - if idx + 1 in self.aux_hidden_state_layers: - # Reconstruct the aux hidden state for draft models - aux_recon = self.mhc_post(hidden_states, residual, post_mix, res_mix) - aux_hidden_state = aux_recon.mean(dim=1) - if self.use_sequence_parallel: - aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens] - aux_hidden_states.append(aux_hidden_state) - final_aux_recon = aux_recon - if layer is not None: - # Reuse if the last layer was captured as an aux hidden state - if self.end_layer in self.aux_hidden_state_layers: + aux_hidden_states.extend(late_aux) + final_aux_recon = None + collapsed = True + if ran_layers and not collapsed: + if ( + self.end_layer in self.aux_hidden_state_layers + and final_aux_recon is not None + ): hidden_states = final_aux_recon else: hidden_states = self.mhc_post( @@ -714,6 +805,162 @@ def forward( return hidden_states, aux_hidden_states return hidden_states + def _run_layers( + self, + layer_ids: range, + hidden_states: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + pre_mix: torch.Tensor | None, + post_mix: torch.Tensor | None, + res_mix: torch.Tensor | None, + residual: torch.Tensor | None, + engram_hashes: torch.Tensor | None, + engram_mask: torch.Tensor | None, + full_num_tokens: int, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + list[torch.Tensor], + torch.Tensor | None, + bool, + ]: + aux_hidden_states: list[torch.Tensor] = [] + final_aux_recon: torch.Tensor | None = None + ran_layers = False + for idx in layer_ids: + ran_layers = True + hidden_states, residual, post_mix, res_mix, pre_mix = self.layers[idx]( + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + engram_hashes, + engram_mask, + ) + # The aux id of the first replay layer is the post of the last + # full-batch layer. Capture it on the trimmed rows instead. + if idx + 1 in self.aux_hidden_state_layers and not ( + self.decoder_replay_layers is not None + and idx + 1 == self.decoder_replay_start + ): + aux_recon = self.mhc_post(hidden_states, residual, post_mix, res_mix) + aux_hidden_state = aux_recon.mean(dim=1) + if self.use_sequence_parallel: + aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens] + aux_hidden_states.append(aux_hidden_state) + final_aux_recon = aux_recon + return ( + hidden_states, + residual, + post_mix, + res_mix, + pre_mix, + aux_hidden_states, + final_aux_recon, + ran_layers, + ) + + def _run_replay_layers( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + pre_mix: torch.Tensor, + post_mix: torch.Tensor, + res_mix: torch.Tensor, + residual: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + """Layers past the last KV source, on the rows they are given. + + The final ``mhc_post`` runs here, on this batch. The return is + ``(hidden_states, pre_mix, *aux)``. + """ + # Post of the layer before the cut. These rows are already trimmed, + # and sequence parallel is unsupported while replay is on. + cut_aux: list[torch.Tensor] = [] + if self.decoder_replay_start in self.aux_hidden_state_layers: + cut_aux.append( + self.mhc_post(hidden_states, residual, post_mix, res_mix).mean(dim=1) + ) + ( + hidden_states, + residual, + post_mix, + res_mix, + pre_mix, + aux, + final_aux_recon, + ran_layers, + ) = self._run_layers( + range(self.decoder_replay_start, self.end_layer), + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + None, + None, + positions.shape[0], + ) + assert residual is not None and post_mix is not None + assert res_mix is not None and pre_mix is not None + if ( + self.end_layer in self.aux_hidden_state_layers + and final_aux_recon is not None + ): + hidden_states = final_aux_recon + elif ran_layers: + hidden_states = self.mhc_post(hidden_states, residual, post_mix, res_mix) + return (hidden_states, pre_mix, *cut_aux, *aux) + + def _decoder_replay_supported(self, vllm_config: VllmConfig, cut: int) -> bool: + """Whether this rank may replay from layer ``cut`` on; warns when not.""" + parallel_config = vllm_config.parallel_config + spec_config = vllm_config.speculative_config + draft_config = spec_config.draft_model_config if spec_config else None + draft_hf_config = getattr(draft_config, "hf_config", None) + draft_window = getattr(draft_hf_config, "sliding_window", None) + draft_layer_types = getattr(draft_hf_config, "layer_types", None) or () + window = self.config.sliding_window + if self.start_layer > cut or self.end_layer < self.config.num_hidden_layers: + reason = ( + "the pipeline stage holding the last KV source layer must also " + "hold every layer after it" + ) + elif ( + self.use_sequence_parallel + or parallel_config.prefill_context_parallel_size > 1 + or parallel_config.use_ubatching + ): + reason = ( + "the replay-layer batch shrinks per rank, which sequence and " + "prefill-context parallelism and microbatching cannot follow" + ) + elif any(i >= cut for i in getattr(self.config, "engram_layer_ids", ())): + reason = "an Engram layer sits at or after the last KV source layer" + elif draft_config is not None and ( + draft_window is None + or draft_window > window + or any(t != "sliding_attention" for t in draft_layer_types) + ): + reason = ( + f"the drafter (sliding window {draft_window}) reads hidden states " + f"outside the target's {window}-token window" + ) + else: + return True + logger.warning_once("Decoder SWA bounded replay is off: %s.", reason) + return False + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: stacked_params_mapping = [ # (param_name, shard_name, shard_id) @@ -1128,6 +1375,10 @@ def token_lookback_depth(self) -> int: engram_hash = self.model.engram_hash return engram_hash.lookback_depth if engram_hash is not None else 0 + @property + def decoder_replay_layers(self) -> DecoderReplayLayers | None: + return self.model.decoder_replay_layers + def forward( self, input_ids: torch.Tensor, diff --git a/vllm/models/deepseek_v41/amd/vl_model.py b/vllm/models/deepseek_v41/amd/vl_model.py index 5ed045029b3c..5b8f7b826205 100644 --- a/vllm/models/deepseek_v41/amd/vl_model.py +++ b/vllm/models/deepseek_v41/amd/vl_model.py @@ -51,6 +51,7 @@ DeepseekV4VLProcessingInfo, ) from ..common.vl_cudagraph import DeepseekV4VLEncoderCudaGraphMixin +from ..decoder_replay_layers import DecoderReplayLayers from .model import ( DeepseekV41LLMForCausalLM, _linear_scale_param_name, @@ -299,6 +300,10 @@ def get_model_state_cls(): def token_lookback_depth(self) -> int: return self.language_model.token_lookback_depth + @property + def decoder_replay_layers(self) -> DecoderReplayLayers | None: + return self.language_model.decoder_replay_layers + def forward( self, input_ids: torch.Tensor, diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 559d874cc483..4d460d89b801 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -547,11 +547,11 @@ def __init__( swa_bounded_replay = False if swa_bounded_replay and current_platform.is_rocm(): logger.warning_once( - "SWA bounded replay is off on ROCm (the sparse SWA metadata " - "builders forward replay_start, but the window clamp it relies " - "on lives in the FlashInfer and FlashMLA prefill kernels, so " - "the padded slots fault); the sliding-window cache takes part " - "in prefix caching instead." + "Prefix-cache SWA replay is off on ROCm (the sparse SWA " + "metadata builders forward replay_start, but the window clamp " + "it relies on lives in the FlashInfer and FlashMLA prefill " + "kernels, so the padded slots fault); the sliding-window cache " + "takes part in prefix caching instead." ) swa_bounded_replay = False self.swa_cache_layer = DeepseekV4SWACache( diff --git a/vllm/models/deepseek_v41/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py index 7fc0bd25ff75..e8bbaef64ddc 100644 --- a/vllm/models/deepseek_v41/nvidia/model_state.py +++ b/vllm/models/deepseek_v41/nvidia/model_state.py @@ -337,7 +337,7 @@ def prepare_attn( self._warm_up_replay_kernels(input_batch, slot_mappings) window, cacheable_groups = self._replay replay_start: torch.Tensor | None = None - if window: + if window or self.decoder_replay_layers is not None: num_reqs = input_batch.num_reqs # Decode rows sit above the hit, so only prefills carry a replay # start; dummy batches (captures, profiling) carry none. @@ -349,7 +349,7 @@ def prepare_attn( replay_start = self._replay_start_staging.copy_to_gpu( replay_start_np, out=self._replay_start[:num_reqs] ) - if replay_start_np.any(): + if window and replay_start_np.any(): # The replayed tokens rebuild window KV only: their slots in the # prefix-cacheable groups are padded so the cached KV stays as is. _pad_replayed_slots_kernel[(num_reqs,)]( @@ -397,21 +397,21 @@ def _warm_up_replay_kernels( dummy batches never replay or trim, so they would not launch them.""" assert self._replay is not None window, cacheable_groups = self._replay - if not window: - return - _pad_replayed_slots_kernel.warmup( - slot_mappings, - slot_mappings.stride(0), - cacheable_groups, - input_batch.query_start_loc, - input_batch.positions, - self._replay_start, - window, - PAD_SLOT_ID, - NUM_GROUPS=cacheable_groups.numel(), - BLOCK=1024, - grid=(1,), - ) + # Pad warmup is prefix-cache SWA. A zero window must not skip the gather warmup. + if window: + _pad_replayed_slots_kernel.warmup( + slot_mappings, + slot_mappings.stride(0), + cacheable_groups, + input_batch.query_start_loc, + input_batch.positions, + self._replay_start, + window, + PAD_SLOT_ID, + NUM_GROUPS=cacheable_groups.numel(), + BLOCK=1024, + grid=(1,), + ) if self.decoder_replay_layers is None: return assert self._kept_slot_mappings is not None From df6c7ad03194a3051a197d57ca4a94e4a063014e Mon Sep 17 00:00:00 2001 From: Dakai An <77474977+andakai@users.noreply.github.com> Date: Tue, 13 Oct 2026 11:05:39 +0800 Subject: [PATCH 349/624] [Core][Refactor] Move KV cache planning logic to kv_cache_planning.py (#61183) Signed-off-by: Dakai An Signed-off-by: Yifan Qiao Co-authored-by: Yifan Qiao Co-authored-by: Claude Opus 5.5 --- .buildkite/test-amd.yaml | 5 + .buildkite/test_areas/lm_eval.yaml | 6 + tests/models/test_dspark_mla.py | 2 +- tests/models/test_initialization.py | 6 +- .../test_sparse_mla_kv_cache_layout.py | 2 +- tests/v1/core/test_contiguous_kv_packing.py | 6 +- tests/v1/core/test_kv_cache_planning.py | 2284 ++++++++++++++++ tests/v1/core/test_kv_cache_utils.py | 2348 +---------------- tests/v1/core/utils.py | 74 + tests/v1/engine/test_init_error_messaging.py | 4 +- .../unit/offloading_connector/test_config.py | 6 +- .../unit/test_hidden_states_connector.py | 2 +- tests/v1/simple_kv_offload/test_worker.py | 4 +- .../test_dsv4_packed_zeroer_geometry.py | 2 +- tests/v1/worker/test_gpu_model_runner.py | 5 +- vllm/v1/core/kv_cache_planning.py | 1637 ++++++++++++ vllm/v1/core/kv_cache_utils.py | 1615 +----------- vllm/v1/engine/core.py | 2 +- vllm/v1/hisparse/layout.py | 4 +- vllm/v1/worker/gpu/cudagraph_utils.py | 2 +- vllm/v1/worker/gpu_model_runner.py | 2 +- 21 files changed, 4063 insertions(+), 3955 deletions(-) create mode 100644 tests/v1/core/test_kv_cache_planning.py create mode 100644 vllm/v1/core/kv_cache_planning.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index d2702e7d4f82..5145f5a272ff 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -4321,6 +4321,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -4354,6 +4355,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -4387,6 +4389,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -4420,6 +4423,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -4453,6 +4457,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index f6a4a89f11cd..96e5d4fc5e9f 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -371,6 +371,7 @@ steps: - vllm/tokenizers/ - vllm/v1/attention/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/v1/spec_decode/metrics.py - vllm/v1/worker/ - vllm/_aiter_ops.py @@ -516,6 +517,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -559,6 +561,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -602,6 +605,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -645,6 +649,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -688,6 +693,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index b2156af346c8..81019421ed13 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -394,7 +394,7 @@ def test_k3_dspark_mla_kv_cache_spec_groups_with_target_mla(): """The draft's MLA layers must share a KV cache group with the target's.""" from vllm.model_executor.layers.attention.mla_attention import MLAAttention from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention - from vllm.v1.core.kv_cache_utils import _get_kv_cache_groups_uniform_page_size + from vllm.v1.core.kv_cache_planning import _get_kv_cache_groups_uniform_page_size vllm_config = SimpleNamespace( model_config=None, cache_config=SimpleNamespace(block_size=64) diff --git a/tests/models/test_initialization.py b/tests/models/test_initialization.py index 073f616af875..ee6dfee7ecbf 100644 --- a/tests/models/test_initialization.py +++ b/tests/models/test_initialization.py @@ -14,10 +14,8 @@ from vllm.platforms import current_platform from vllm.utils.mem_constants import GiB_bytes from vllm.v1.attention.backends.utils import resolve_kv_cache_layout -from vllm.v1.core.kv_cache_utils import ( - generate_scheduler_kv_cache_config, - get_kv_cache_configs, -) +from vllm.v1.core.kv_cache_planning import get_kv_cache_configs +from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config from vllm.v1.engine.core import EngineCore as V1EngineCore from ..utils import ( diff --git a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py index d869141df81a..45ba1ac0dc55 100644 --- a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py +++ b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py @@ -23,7 +23,7 @@ flat_kv_row_view, ) from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts -from vllm.v1.core.kv_cache_utils import get_kv_cache_config_from_groups +from vllm.v1.core.kv_cache_planning import get_kv_cache_config_from_groups from vllm.v1.kv_cache_interface import KVCacheGroupSpec, UniformTypeKVCacheSpecs from vllm.v1.kv_cache_layout import KVCacheLayout from vllm.v1.worker.utils import allocate_kv_cache, customize_attention_spec diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 6a705b6e5e8e..64cd55054015 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -18,13 +18,15 @@ from vllm.config import CacheConfig from vllm.v1.core.kv_cache_manager import KVCacheManager -from vllm.v1.core.kv_cache_utils import ( +from vllm.v1.core.kv_cache_planning import ( _get_kv_cache_bytes_per_block, _get_packed_kv_cache_groups, _pool_bytes_per_block, - generate_scheduler_kv_cache_config, get_kv_cache_config_from_groups, get_kv_cache_groups, +) +from vllm.v1.core.kv_cache_utils import ( + generate_scheduler_kv_cache_config, resolve_kv_cache_block_sizes, ) from vllm.v1.hisparse.layout import _build_hisparse_kv_cache_tensors diff --git a/tests/v1/core/test_kv_cache_planning.py b/tests/v1/core/test_kv_cache_planning.py new file mode 100644 index 000000000000..629286808148 --- /dev/null +++ b/tests/v1/core/test_kv_cache_planning.py @@ -0,0 +1,2284 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import mmap +import subprocess +import sys +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +import vllm.v1.core.kv_cache_planning as kv_cache_planning +import vllm.v1.core.kv_cache_utils as kv_cache_utils +import vllm.v1.hisparse.runtime as hisparse_runtime_module +from tests.v1.core.utils import ( + new_kv_cache_spec, + new_mamba_spec, + new_mla_spec, + new_sliding_window_spec, +) +from vllm.config import ( + CacheConfig, + ModelConfig, + SchedulerConfig, + VllmConfig, +) +from vllm.config.attention import HiSparseConfig +from vllm.utils.mem_constants import GiB_bytes +from vllm.v1.core.kv_cache_planning import ( + check_enough_kv_cache_memory, + estimate_max_model_len, + get_kv_cache_configs, + get_kv_cache_groups, + is_kv_cache_spec_uniform, + kv_cache_groups_tp_replicas, +) +from vllm.v1.core.kv_cache_utils import ( + generate_scheduler_kv_cache_config, +) +from vllm.v1.hisparse.layout import ( + get_hisparse_kv_cache_groups, +) +from vllm.v1.kv_cache_interface import ( + ChunkedLocalAttentionSpec, + CircularBufferSpec, + FullAttentionSpec, + HiddenStateCacheSpec, + HiSparseHotSpec, + HiSparseResidentSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheSpec, + KVCacheTensor, + MambaSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, + SlidingWindowSpec, + SparseCacheRole, + UniformTypeKVCacheSpecs, + iter_layer_specs, +) +from vllm.v1.kv_cache_layout import KVCacheLayout + +pytestmark = pytest.mark.cpu_test + + +@pytest.mark.parametrize("gpu_block_size", [32, 64]) +@pytest.mark.parametrize("shared_host_pool", [False, True]) +def test_hisparse_hma_uses_resolved_gpu_block_size( + monkeypatch, gpu_block_size, shared_host_pool +): + monkeypatch.setattr( + hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True + ) + specs = { + "model.layers.0.self_attn": MLAAttentionSpec( + block_size=gpu_block_size, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + is_index_group_leader=True, + ), + "model.layers.0.self_attn.indexer": MLAAttentionSpec( + block_size=gpu_block_size, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + cache_role=SparseCacheRole.INDEXER, + ), + } + config = SimpleNamespace( + attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()), + model_config=SimpleNamespace( + hf_config=SimpleNamespace(index_topk=128), + max_model_len=gpu_block_size, + ), + parallel_config=SimpleNamespace( + tensor_parallel_size=2 if shared_host_pool else 1, + pipeline_parallel_size=1, + prefill_context_parallel_size=1, + decode_context_parallel_size=1, + world_size=2 if shared_host_pool else 1, + distributed_executor_backend="mp", + nnodes_within_dp=1, + ), + cache_config=SimpleNamespace( + num_gpu_blocks_override=7, + prefix_cache_retention_interval=None, + get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC, + ), + ) + monkeypatch.setattr( + kv_cache_planning, "get_hisparse_host_pool_bytes", lambda _: 2**30 + ) + cache_config = kv_cache_planning.get_kv_cache_config_from_groups( + config, get_hisparse_kv_cache_groups(config, specs), available_memory=2**30 + ) + assert cache_config.num_blocks == 7 + assert cache_config.hisparse_host_num_blocks is not None + assert cache_config.hisparse_host_num_blocks > 7 + + host_group, indexer_group, *auxiliary_groups = cache_config.kv_cache_groups + assert cache_config.hisparse_shared_host_pool is shared_host_pool + host_page = host_group.kv_cache_spec.page_size_bytes + alignment = mmap.PAGESIZE if shared_host_pool else 1 + expected_host_stride = (host_page + alignment - 1) // alignment * alignment + assert cache_config.hisparse_host_block_stride == expected_host_stride + assert cache_config.hisparse_host_num_blocks == 2**30 // expected_host_stride + assert host_group.host_resident + assert not any(group.host_resident for group in [indexer_group, *auxiliary_groups]) + host_layers = set(host_group.layer_names) + for tensor in cache_config.kv_cache_tensors: + assert all( + (name in host_layers) == tensor.host_resident for name in tensor.layers + ) + host_tensor = next(t for t in cache_config.kv_cache_tensors if t.host_resident) + host_spec = host_group.kv_cache_spec.kv_cache_specs[host_tensor.layers[0]] + assert host_tensor.block_stride == host_spec.page_size_bytes + assert host_tensor.layer_stride == ( + host_spec.page_size_bytes * cache_config.hisparse_host_num_blocks + ) + assert host_group.kv_cache_spec.block_size == gpu_block_size + assert indexer_group.kv_cache_spec.block_size == gpu_block_size + host_specs = host_group.kv_cache_spec.kv_cache_specs + gpu_indexer_specs = indexer_group.kv_cache_spec.kv_cache_specs + assert set(host_specs) == {"model.layers.0.self_attn"} + assert set(gpu_indexer_specs) == {"model.layers.0.self_attn.indexer"} + assert indexer_group.kv_cache_spec.prefix_cacheable + assert host_group.enable_kv_transfer + auxiliary_specs = [group.kv_cache_spec for group in auxiliary_groups] + assert any(isinstance(spec, HiSparseResidentSpec) for spec in auxiliary_specs) + assert any(isinstance(spec, HiSparseHotSpec) for spec in auxiliary_specs) + assert all( + spec.block_size == gpu_block_size + for spec in auxiliary_specs + if isinstance(spec, (HiSparseResidentSpec, HiSparseHotSpec)) + ) + scheduler_block_size, hash_block_size = kv_cache_utils.resolve_kv_cache_block_sizes( + cache_config, + SimpleNamespace( + cache_config=SimpleNamespace( + block_size=16, + enable_prefix_caching=True, + prefix_match_unit=None, + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + kv_transfer_config=object(), + ), + ) + assert scheduler_block_size == hash_block_size == gpu_block_size + + +@pytest.mark.parametrize( + "max_model_len,num_gpu_blocks,ok", + [ + (8192, 160, False), + (8192, 1000, True), + (32768, 656, False), + (32768, 657, True), + (-1, 657, True), + ], +) +def test_hisparse_pool_must_fit_max_model_len( + monkeypatch, max_model_len, num_gpu_blocks, ok +): + """Each resident group takes its own blocks from HiSparse's shared GPU pool, + so a pool that fits only the indexer pages must be rejected at startup. + Resident pages are capped at the in-flight window: at 32768 tokens with 2048 + in flight, one request needs 512 indexer + 4 x 32 resident + 16 hot blocks, + plus the null block. Auto-fit (-1) must pick a length that passes the same + check.""" + monkeypatch.setattr( + hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True + ) + monkeypatch.setattr( + kv_cache_planning, "get_hisparse_host_pool_bytes", lambda _: 2**30 + ) + specs: dict[str, KVCacheSpec] = {} + for i in range(4): + specs[f"model.layers.{i}.self_attn"] = MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + is_index_group_leader=True, + ) + specs[f"model.layers.{i}.self_attn.indexer"] = MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + cache_role=SparseCacheRole.INDEXER, + ) + config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) + config.scheduler_config.max_num_batched_tokens = 2048 + # Async scheduling doubles the in-flight window, and with it the bound. + config.scheduler_config.async_scheduling = False + config.attention_config.hisparse_config = HiSparseConfig() + config.model_config.hf_config.index_topk = 128 + config.cache_config.num_gpu_blocks_override = num_gpu_blocks + config.cache_config.kv_cache_layout = "BLHNC" + # Workers report the resident/hot caches HiSparse derives. + specs |= { + name: group.kv_cache_spec + for group in get_hisparse_kv_cache_groups(config, specs) or [] + if isinstance(group.kv_cache_spec, (HiSparseResidentSpec, HiSparseHotSpec)) + for name in group.layer_names + } + + if ok: + kv_cache_planning.get_kv_cache_configs(config, [specs], [2**34]) + else: + with pytest.raises(ValueError, match="max seq len"): + kv_cache_planning.get_kv_cache_configs(config, [specs], [2**34]) + + +def new_chunked_local_attention_spec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.float32, + page_size_padded=None, + attention_chunk_size=4, +): + return ChunkedLocalAttentionSpec( + block_size=block_size, + num_kv_heads=num_kv_heads, + head_size=head_size, + dtype=dtype, + page_size_padded=page_size_padded, + attention_chunk_size=attention_chunk_size, + ) + + +def test_get_kv_cache_configs_multiple_workers(): + model_config = ModelConfig(max_model_len=16) + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + vllm_config.cache_config.prefix_cache_retention_interval = None + + ref_kv_cache_spec = new_kv_cache_spec() + same_kv_cache_specs = [ + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + ] + + # Basic case. All things are the same. + kv_cache_configs = get_kv_cache_configs( + vllm_config, + same_kv_cache_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + expected = KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer1", "layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), + ], + ) + assert kv_cache_configs == [expected, expected] + + # Different available memory. This is the case for TP. + # Use the smallest memory available. + kv_cache_configs = get_kv_cache_configs( + vllm_config, + same_kv_cache_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 20, + ], + ) + assert kv_cache_configs == [expected, expected] + + # Different KV cache specs. This is the case for PP. + different_layer_specs = [ + { + "layer1": new_kv_cache_spec(), + }, + { + "layer2": new_kv_cache_spec(), + "layer3": new_kv_cache_spec(), + }, + ] + + # Different workers have different layers. + kv_cache_configs = get_kv_cache_configs( + vllm_config, + different_layer_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer1"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1"], new_kv_cache_spec()), + ], + ), + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer2", "layer3"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer2", "layer3"], new_kv_cache_spec()), + ], + ), + ] + + # Some layers are the same, some are different. This is the case for TP+PP + tp_pp_kv_cache_specs = [ + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + { + "layer3": new_kv_cache_spec(), + }, + { + "layer3": new_kv_cache_spec(), + }, + ] + + kv_cache_configs = get_kv_cache_configs( + vllm_config, + tp_pp_kv_cache_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + expected_12 = KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer1", "layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), + ], + ) + expected_3 = KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer3"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer3"], ref_kv_cache_spec), + ], + ) + assert kv_cache_configs == [ + expected_12, + expected_12, + expected_3, + expected_3, + ] + + # Different workers have different types of layers. This is the case for + # hybrid models + PP. + different_type_layer_specs = [ + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + { + "layer3": new_sliding_window_spec(), + "layer4": new_sliding_window_spec(), + }, + ] + kv_cache_configs = get_kv_cache_configs( + vllm_config, + different_type_layer_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer1", "layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), + KVCacheGroupSpec([], new_sliding_window_spec()), + ], + ), + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer3", "layer4"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec([], ref_kv_cache_spec), + KVCacheGroupSpec(["layer3", "layer4"], new_sliding_window_spec()), + ], + ), + ] + + # When divided into multiple KVCacheGroups, need to ensure the number of + # layers per group is similar. + different_type_layer_specs = [ + { + "layer1": new_kv_cache_spec(), + "layer2": new_sliding_window_spec(), + "layer3": new_sliding_window_spec(), + }, + { + "layer4": new_kv_cache_spec(), + "layer5": new_sliding_window_spec(), + "layer6": new_sliding_window_spec(), + }, + ] + kv_cache_configs = get_kv_cache_configs( + vllm_config, + different_type_layer_specs, + [ + ref_kv_cache_spec.page_size_bytes * 10, + ref_kv_cache_spec.page_size_bytes * 10, + ], + ) + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer1"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer3"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1"], ref_kv_cache_spec), + KVCacheGroupSpec(["layer2"], new_sliding_window_spec()), + KVCacheGroupSpec(["layer3"], new_sliding_window_spec()), + ], + ), + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer4"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer5"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer6"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer4"], ref_kv_cache_spec), + KVCacheGroupSpec(["layer5"], new_sliding_window_spec()), + KVCacheGroupSpec(["layer6"], new_sliding_window_spec()), + ], + ), + ] + + # Have conflicting layers. Need to raise an error. + conflicting_layer_specs = [ + { + "layer1": new_kv_cache_spec(), + }, + { + "layer1": new_sliding_window_spec(), + }, + ] + with pytest.raises(AssertionError): + get_kv_cache_configs( + vllm_config, + conflicting_layer_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + + +@pytest.mark.parametrize( + "asymmetric_memory", + [False, True], + ids=["symmetric", "asymmetric"], +) +def test_get_kv_cache_configs_pp_sharding(asymmetric_memory): + model_config = ModelConfig(max_model_len=512) + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + vllm_config.cache_config.prefix_cache_retention_interval = None + + ref_kv_cache_spec = new_kv_cache_spec() + pp_kv_cache_specs = [ + {"layer1": ref_kv_cache_spec}, + {"layer2": ref_kv_cache_spec}, + ] + + expected_num_blocks = model_config.max_model_len // ref_kv_cache_spec.block_size + 1 + avail_memory = ref_kv_cache_spec.page_size_bytes * expected_num_blocks + + # With per-worker validation, each worker only needs memory for its own + # layers. Worker 2 having more memory shouldn't affect worker 1's config. + available_memory = ( + [avail_memory, avail_memory * 2] if asymmetric_memory else [avail_memory] * 2 + ) + + kv_cache_configs = get_kv_cache_configs( + vllm_config, + pp_kv_cache_specs, + available_memory, + ) + + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=expected_num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, + layers=["layer1"], + layer_stride=ref_kv_cache_spec.page_size_bytes + * expected_num_blocks, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[KVCacheGroupSpec(["layer1"], ref_kv_cache_spec)], + ), + KVCacheConfig( + num_blocks=expected_num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, + layers=["layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes + * expected_num_blocks, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[KVCacheGroupSpec(["layer2"], ref_kv_cache_spec)], + ), + ] + + +def test_project_kv_cache_groups_to_worker(): + spec_a = new_kv_cache_spec() + spec_b = new_kv_cache_spec(num_kv_heads=4) + + global_groups = [ + KVCacheGroupSpec(["layer1", "layer2", "layer3"], spec_a), + ] + worker_spec = {"layer1": spec_a, "layer2": spec_a} + projected = kv_cache_planning._project_kv_cache_groups_to_worker( + global_groups, worker_spec + ) + assert len(projected) == 1 + assert projected[0].layer_names == ["layer1", "layer2"] + assert projected[0].kv_cache_spec is spec_a + + projected = kv_cache_planning._project_kv_cache_groups_to_worker( + global_groups, {"layer4": spec_a} + ) + assert len(projected) == 1 + assert projected[0].layer_names == [] + assert projected[0].kv_cache_spec is spec_a + + uniform_spec = UniformTypeKVCacheSpecs( + block_size=16, + kv_cache_specs={"layer1": spec_a, "layer2": spec_b, "layer3": spec_a}, + ) + global_groups_uniform = [ + KVCacheGroupSpec(["layer1", "layer2", "layer3"], uniform_spec), + ] + projected = kv_cache_planning._project_kv_cache_groups_to_worker( + global_groups_uniform, {"layer1": spec_a, "layer3": spec_a} + ) + assert len(projected) == 1 + assert projected[0].layer_names == ["layer1", "layer3"] + proj_spec = projected[0].kv_cache_spec + assert isinstance(proj_spec, UniformTypeKVCacheSpecs) + assert set(proj_spec.kv_cache_specs.keys()) == {"layer1", "layer3"} + + +@pytest.mark.parametrize("with_draft", [False, True]) +@pytest.mark.parametrize("indexer_alignment", [None, 512]) +def test_sparse_mla_preserves_physical_row_addressing(with_draft, indexer_alignment): + from vllm.v1.attention.backends.mla.sparse_utils import flat_kv_row_view + from vllm.v1.worker.utils import allocate_kv_cache + + config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) + config.cache_config.kv_cache_layout = "BLHNC" + config.parallel_config.decode_context_parallel_size = 4 + common = dict(block_size=64, num_kv_heads=1, dtype=torch.uint8) + specs = { + "target": MLAAttentionSpec( + **common, + head_size=576, + state_content_bytes=656, + block_stride_alignment=656, + cache_dtype_str="fp8_ds_mla", + is_index_group_leader=True, + ), + "indexer": MLAAttentionSpec( + **common, + head_size=132, + cache_role=SparseCacheRole.INDEXER, + block_stride_alignment=indexer_alignment, + ), + "draft": SlidingWindowSpec( + block_size=64, + num_kv_heads=64, + head_size=64, + dtype=torch.bfloat16, + sliding_window=2048, + dcp_sharded=False, + ), + } + if not with_draft: + del specs["draft"] + groups = get_kv_cache_groups(config, specs) + cache_config = kv_cache_planning.get_kv_cache_config_from_groups( + config, groups, available_memory=8 * 1024 * 1024 + ) + caches = allocate_kv_cache( + cache_config, torch.device("cpu"), KVCacheLayout.BLHNC, [64] * len(groups) + ) + cache = caches["target"].squeeze(1) + if indexer_alignment is not None: + assert caches["indexer"].stride(0) % indexer_alignment == 0 + rows, stride_rows = flat_kv_row_view(cache, 64) + cache[1, 0].fill_(7) + torch.testing.assert_close(rows[stride_rows], cache[1, 0]) + assert torch.all(rows[stride_rows] == 7) + + +def test_is_kv_cache_spec_uniform(): + kv_cache_spec = { + "layer_1": new_kv_cache_spec(num_kv_heads=32), + "layer_2": new_kv_cache_spec(num_kv_heads=32), + } + assert is_kv_cache_spec_uniform(kv_cache_spec) + + kv_cache_spec = { + "layer_1": new_kv_cache_spec(num_kv_heads=32), + "layer_2": new_kv_cache_spec(num_kv_heads=32, sliding_window=1), + } + assert is_kv_cache_spec_uniform(kv_cache_spec) + + kv_cache_spec = { + "layer_1": new_kv_cache_spec(num_kv_heads=32), + "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), + } + assert not is_kv_cache_spec_uniform(kv_cache_spec) + + kv_cache_spec = { + "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), + "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), + } + assert is_kv_cache_spec_uniform(kv_cache_spec) + + kv_cache_spec = { + "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), + "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=2), + } + assert not is_kv_cache_spec_uniform(kv_cache_spec) + + script = """ +import sys + +from vllm.v1.core.kv_cache_planning import is_kv_cache_spec_uniform +from vllm.v1.kv_cache_interface import KVCacheSpec + +if sys.flags.optimize < 1: + raise RuntimeError("subprocess is not running with optimization enabled") +specs = { + "a": KVCacheSpec(block_size=1), + "b": KVCacheSpec(block_size=2), +} +if is_kv_cache_spec_uniform(specs): + raise RuntimeError("different specs were treated as uniform") +""" + subprocess.run( + [sys.executable, "-O", "-c", script], + cwd=Path(__file__).resolve().parents[3], + check=True, + timeout=60, + ) + + +@pytest.mark.parametrize( + ("model_id", "max_model_len", "want_estimated_max_len"), + [ + ("Qwen/Qwen1.5-7B", 16385, 16384), + ("Qwen/Qwen1.5-7B", 16383, 16383), + ], +) +def test_estimate_max_model_len(model_id, max_model_len, want_estimated_max_len): + # Create a VllmConfig + model_config = ModelConfig( + model_id, + runner="generate", + dtype="float16", + max_model_len=max_model_len, + ) + scheduler_config = SchedulerConfig( + max_num_batched_tokens=32768, + max_model_len=model_config.max_model_len, + is_encoder_decoder=model_config.is_encoder_decoder, + ) + + vllm_config = VllmConfig( + model_config=model_config, + scheduler_config=scheduler_config, + ) + + # Create KV cache specs + kv_cache_spec = {} + for i in range(32): + layer_name = f"layer_{i}" + kv_cache_spec[layer_name] = FullAttentionSpec( + block_size=16, + num_kv_heads=32, + head_size=128, + dtype=torch.float16, + ) + # Estimate the maximum model length, 16384 model_len need 8GB + estimated_max_len = estimate_max_model_len( + vllm_config, kv_cache_spec, 8 * GiB_bytes + ) + assert estimated_max_len == want_estimated_max_len + + +def test_get_kv_cache_config_one_worker(): + # pass max_model_len to pass check_enough_kv_cache_memory + model_config = ModelConfig(max_model_len=16) + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + vllm_config.cache_config.prefix_cache_retention_interval = None + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 + # all layers are full attention -> single group + kv_cache_specs_full = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + } + kv_cache_config_full = get_kv_cache_configs( + vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32] + )[0] + print(kv_cache_config_full) + assert kv_cache_config_full == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], + ) + + # all layers are sliding window -> single group + kv_cache_specs_sliding = { + "layer_1": new_sliding_window_spec(), + "layer_2": new_sliding_window_spec(), + } + kv_cache_config_sliding = get_kv_cache_configs( + vllm_config, [kv_cache_specs_sliding], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_sliding == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1", "layer_2"], new_sliding_window_spec()) + ], + ) + + # full + sliding, but disable_hybrid_kv_cache_manager + vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = True + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_sliding_window_spec(), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1", "layer_2"], new_kv_cache_spec(sliding_window=1) + ), + ], + ) + vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False + + # full + sliding, with hybrid_kv_cache_manager + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_sliding_window_spec(), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=64, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 64, + layers=["layer_1"], + layer_stride=mem_per_block_per_layer * 64, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 64, + layers=["layer_2"], + layer_stride=mem_per_block_per_layer * 64, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1"], new_kv_cache_spec()), + KVCacheGroupSpec(["layer_2"], new_sliding_window_spec()), + ], + ) + + # 2 full + 4 sliding, 2 layers per group + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + "layer_3": new_sliding_window_spec(), + "layer_4": new_sliding_window_spec(), + "layer_5": new_sliding_window_spec(), + "layer_6": new_sliding_window_spec(), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_3", "layer_5"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_4", "layer_6"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec()), + KVCacheGroupSpec(["layer_3", "layer_5"], new_sliding_window_spec()), + KVCacheGroupSpec(["layer_4", "layer_6"], new_sliding_window_spec()), + ], + ) + + # 3 full + 7 sliding, pad to 3 full + 9 sliding + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + "layer_3": new_kv_cache_spec(), + "layer_4": new_sliding_window_spec(), + "layer_5": new_sliding_window_spec(), + "layer_6": new_sliding_window_spec(), + "layer_7": new_sliding_window_spec(), + "layer_8": new_sliding_window_spec(), + "layer_9": new_sliding_window_spec(), + "layer_10": new_sliding_window_spec(), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_1", "layer_2", "layer_3"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_4", "layer_7", "layer_10"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_5", "layer_8"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_6", "layer_9"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1", "layer_2", "layer_3"], new_kv_cache_spec()), + KVCacheGroupSpec( + ["layer_4", "layer_7", "layer_10"], new_sliding_window_spec() + ), + KVCacheGroupSpec(["layer_5", "layer_8"], new_sliding_window_spec()), + KVCacheGroupSpec(["layer_6", "layer_9"], new_sliding_window_spec()), + ], + ) + + # 6 full + 5 sliding, pad to 6 full + 6 sliding. This is a typical case for gpt-oss + # eagle where there is only one more full attention layer than sliding window layers + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + "layer_3": new_kv_cache_spec(), + "layer_4": new_kv_cache_spec(), + "layer_5": new_kv_cache_spec(), + "layer_6": new_kv_cache_spec(), + "layer_7": new_sliding_window_spec(), + "layer_8": new_sliding_window_spec(), + "layer_9": new_sliding_window_spec(), + "layer_10": new_sliding_window_spec(), + "layer_11": new_sliding_window_spec(), + } + + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 6 * 32] + )[0] + print(kv_cache_config_hybrid) + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 6, + layers=[ + "layer_1", + "layer_2", + "layer_3", + "layer_4", + "layer_5", + "layer_6", + ], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 6, + layers=["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1", "layer_2", "layer_3", "layer_4", "layer_5", "layer_6"], + new_kv_cache_spec(), + ), + KVCacheGroupSpec( + ["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], + new_sliding_window_spec(), + ), + ], + ) + + # different hidden size but same type, use UniformTypeKVCacheSpecs + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(head_size=128), + "layer_2": new_kv_cache_spec(head_size=64), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] + )[0] + # Layers of different page sizes pack densely into one allocation: the + # 2x-larger layer_1 takes the first region, layer_2 the next. + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_1"], + layer_stride=mem_per_block_per_layer * 2 * 32, + block_stride=mem_per_block_per_layer * 2, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + offset=mem_per_block_per_layer * 2 * 32, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1", "layer_2"], + UniformTypeKVCacheSpecs( + block_size=16, kv_cache_specs=kv_cache_specs_hybrid + ), + ) + ], + ) + + # Different hidden size and different type, align by different block size + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(head_size=64), + "layer_2": new_sliding_window_spec(head_size=32), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32, + layers=["layer_1"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32, + layers=["layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1"], new_kv_cache_spec(head_size=64)), + KVCacheGroupSpec( + ["layer_2"], new_sliding_window_spec(head_size=32, block_size=32) + ), + ], + ) + + # different hidden size that cannot be aligned by using different block size, + # but can be aligned by padding the smaller physical page. + swa_spec = new_sliding_window_spec(head_size=96) + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(head_size=64), + "layer_2": swa_spec, + } + + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] + )[0] + padded_page_size = swa_spec.page_size_bytes + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=42, + kv_cache_tensors=[ + KVCacheTensor( + size=padded_page_size * 42, + layers=["layer_1"], + layer_stride=padded_page_size * 42, + block_stride=padded_page_size, + ), + KVCacheTensor( + size=padded_page_size * 42, + layers=["layer_2"], + layer_stride=padded_page_size * 42, + block_stride=padded_page_size, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1"], + new_kv_cache_spec( + head_size=64, + page_size_padded=padded_page_size, + ), + ), + KVCacheGroupSpec( + ["layer_2"], + new_sliding_window_spec(head_size=96), + ), + ], + ) + + # Test num_gpu_blocks_override + vllm_config.cache_config.num_gpu_blocks_override = 16 + kv_cache_config_override_blocks = get_kv_cache_configs( + vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_override_blocks == KVCacheConfig( + num_blocks=16, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 16 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 16, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], + ) + + +def test_get_kv_cache_configs_attention_free(): + kv_cache_specs: dict[str, KVCacheSpec] = {} + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) + vllm_config.cache_config.prefix_cache_retention_interval = None + kv_cache_configs = get_kv_cache_configs(vllm_config, [kv_cache_specs], [0]) + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=[], + ) + ] + + +def test_packed_groups_glm5_like_hybrid(): + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=8192)) + vllm_config.cache_config.kv_cache_layout = "BLHNC" + kpool, block_size = 4, 1024 + kv_cache_spec: dict[str, KVCacheSpec] = {} + for i in range(11): + kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + block_stride_alignment=32 * 576 * 2, + ) + kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + tokens_per_state=kpool, + block_stride_alignment=64 * 132, + ) + kv_cache_spec[f"layers.{i}.tail"] = CircularBufferSpec( + block_size=kpool, + num_kv_heads=2, + head_size=128, + head_size_v=0, + dtype=torch.bfloat16, + ) + for i in range(34): + kv_cache_spec[f"layers.{i}.kda"] = MambaSpec( + shapes=((3, 6144), (16, 128, 128)), + dtypes=(torch.bfloat16, torch.float32), + block_size=block_size, + ) + + groups = kv_cache_planning.get_kv_cache_groups(vllm_config, kv_cache_spec) + attn = [ + g + for g in groups + if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) + and isinstance(g.kv_cache_spec.first_spec, MLAAttentionSpec) + ] + rings = [ + g + for g in groups + if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) + and isinstance(g.kv_cache_spec.first_spec, CircularBufferSpec) + ] + mamba = [ + g + for g in groups + if isinstance( + getattr(g.kv_cache_spec, "first_spec", g.kv_cache_spec), MambaSpec + ) + ] + assert len(attn) == 1 and len(attn[0].layer_names) == 22 + assert len(rings) == 1 and len(rings[0].layer_names) == 11 + # 11 * (1.18 MB MLA + 34 KB indexer) per block fits 12 fp32-SSM KDA states. + assert len(mamba) == 3 + assert sum(len(g.layer_names) for g in mamba) == 34 + + kv_cache_config = kv_cache_planning.get_kv_cache_config_from_groups( + vllm_config, groups, available_memory=2 * 1024**3 + ) + strides = {t.block_stride for t in kv_cache_config.kv_cache_tensors} + assert len(strides) == 1 + (block_stride,) = strides + # Blocks sit a whole number of 32-row MLA pages and 64-state indexer pages + # apart so both caches can be re-paged for their kernels. + assert block_stride % (32 * 576 * 2) == 0 + assert block_stride % (64 * 132) == 0 + + +def new_swa_mla_spec(head_size=576, sliding_window=128, model_version=None): + return SlidingWindowMLASpec( + block_size=16, + num_kv_heads=1, + max_tp_shards=1, + head_size=head_size, + dtype=torch.float32, + sliding_window=sliding_window, + model_version=model_version, + ) + + +def new_indexer_mla_spec(block_size=16): + # Sparse-attention indexer k_cache: an MLAAttentionSpec with a much smaller + # page size than the main MLA attention (uint8, small head), so their pages + # cannot be unified. + return MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + ) + + +def test_mixed_page_size_groups_use_spec_compatibility(): + specs = {} + for i in range(3): + specs[f"mla.{i}"] = new_mla_spec() + specs[f"indexer.{i}"] = new_indexer_mla_spec() + specs.update({f"swa.{i}": new_swa_mla_spec(head_size=1024) for i in range(5)}) + + config = _grouping_config() + config.cache_config = CacheConfig() + config.cache_config.kv_cache_layout = "BLNHC" + groups = get_kv_cache_groups(config, specs) + + assert len(groups) == 3 + assert {name for group in groups for name in group.layer_names} == set(specs) + assert sorted(len(group.layer_names) for group in groups) == [2, 3, 6] + + +def _grouping_config(): + cache_config = CacheConfig() + cache_config.kv_cache_layout = "LBNHC" + return SimpleNamespace( + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), + speculative_config=None, + cache_config=cache_config, + ) + + +def test_hidden_state_group_preserves_hybrid_prefix_cache_granularity(): + block_size = 544 + full_spec = FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + # FullAttentionSpec stores both K and V, so this produces a + # 544 * 1 * (512 + 512) * 2 = 1,114,112-byte page. + head_size=512, + dtype=torch.bfloat16, + ) + mamba_spec = MambaSpec( + block_size=block_size, + shapes=((557056,),), + dtypes=(torch.bfloat16,), + mamba_cache_mode="align", + ) + hidden_spec = HiddenStateCacheSpec( + block_size=block_size, + num_kv_heads=3, + head_size=1024, + dtype=torch.bfloat16, + ) + assert full_spec.page_size_bytes == mamba_spec.page_size_bytes + + groups = get_kv_cache_groups( + _grouping_config(), + { + "model.full_attn": full_spec, + "model.mamba": mamba_spec, + "cache_only_layers.0": hidden_spec, + }, + ) + + hidden_group = next( + group + for group in groups + if isinstance(group.kv_cache_spec, HiddenStateCacheSpec) + ) + assert hidden_group.kv_cache_spec.block_size == 136 + + kv_cache_config = KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=groups, + ) + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + block_size=16, + enable_prefix_caching=True, + prefix_match_unit=None, + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + kv_transfer_config=object(), + ) + assert kv_cache_utils.resolve_kv_cache_block_sizes( + kv_cache_config, vllm_config + ) == (544, 136) + + +def test_multi_run_layer_compact_strides_place_hoisted_heads(): + """A layer-compact run region is its own dense allocation: under LHBNC the head + groups sit between the layers and the blocks, so a run's block stride is one head + group's slice, not the whole page (regression test for setStorage out-of-bounds + on ROCm hybrid models).""" + full = new_kv_cache_spec() + swa = new_sliding_window_spec(sliding_window=full.block_size * 2) + page = full.page_size_bytes + assert swa.page_size_bytes == page + + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) + vllm_config.cache_config.kv_cache_layout = "LHBNC" + config = kv_cache_planning.get_kv_cache_config_from_groups( + vllm_config, + [ + KVCacheGroupSpec(["full.0", "full.1"], full), + KVCacheGroupSpec(["swa.0", "swa.1"], swa), + ], + available_memory=4 * page * 8, + ) + + num_blocks = config.num_blocks + head_group = full.block_size * full.state_content_size_bytes + for tensor in config.kv_cache_tensors: + # One head group per block, head groups between the layers and blocks. + assert tensor.block_stride == head_group + assert tensor.layer_stride == full.num_heads * num_blocks * head_group + # The last block of the last layer's last head group must stay within the + # allocation (the old page-based stride overran it). + end = ( + tensor.offset + + tensor.layer_stride * (len(tensor.layers) - 1) + + (full.num_heads - 1) * num_blocks * head_group + + tensor.block_stride * (num_blocks - 1) + + head_group + ) + assert end <= tensor.size + + +def test_mla_draft_prefers_standard_layout_when_pages_can_be_unified(): + specs = { + "target.0.attn": new_mla_spec(), + "draft.0": new_sliding_window_spec(num_kv_heads=1, head_size=288), + } + assert len({spec.page_size_bytes for spec in specs.values()}) == 1 + + groups = get_kv_cache_groups(_grouping_config(), specs) + + assert len(groups) == 2 + assert all( + not isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in groups + ) + + +def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog_vllm): + # Sparse MLA pages cannot be padded safely. Keeping the draft's attention + # compute sliding-window while promoting only its allocation semantics lets + # every layer share the target's block table and remain contiguous. + draft = new_sliding_window_spec(block_size=16) + specs = { + "target.0.attn": new_mla_spec(block_size=64), + "target.0.indexer": new_indexer_mla_spec(block_size=64), + "draft.0": draft, + } + + groups = get_kv_cache_groups(_grouping_config(), specs) + assert len(groups) == 1 + assert set(groups[0].layer_names) == set(specs) + group_spec = groups[0].kv_cache_spec + assert isinstance(group_spec, UniformTypeKVCacheSpecs) + assert group_spec.block_size == 64 + promoted_draft = group_spec.kv_cache_specs["draft.0"] + assert isinstance(promoted_draft, FullAttentionSpec) + assert not isinstance(promoted_draft, SlidingWindowSpec) + assert promoted_draft.block_size == 64 + assert promoted_draft.sliding_window == draft.sliding_window + assert specs["draft.0"] is draft + assert "attention compute is unchanged" in caplog_vllm.text + + +def test_hidden_states_with_tp_scales_page_size(): + """When TP shrinks KV pages below the hidden-state per-token cost, + get_kv_cache_groups must scale up target block sizes so that the + common page accommodates the unsharded hidden states.""" + # Simulate TP=4 sharding a model with 8 KV heads → 2 per rank. + # KV page = block_size(16) * num_kv_heads(2) * head_size(64) * dtype(2) + # = 16 * 2 * 64 * 2 = 4096 bytes. + kv_spec = new_kv_cache_spec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.bfloat16, + ) + # Hidden-state per-token cost = num_hidden_states(6) * hidden_size(512) + # * dtype(2) = 6144 bytes, which exceeds the 4096-byte KV page. + hs_spec = HiddenStateCacheSpec( + block_size=16, + num_kv_heads=6, + head_size=512, + dtype=torch.bfloat16, + ) + specs = { + "target.0.attn": kv_spec, + "target.1.attn": kv_spec, + "cache_only_layers.48": hs_spec, + } + + groups = get_kv_cache_groups(_grouping_config(), specs) + + # The hidden-state layer should be present and no assertion should fire. + all_layers = {name for g in groups for name in g.layer_names} + assert "cache_only_layers.48" in all_layers + + # The target group block sizes must have been scaled up. + for g in groups: + if "cache_only_layers.48" not in g.layer_names: + assert g.kv_cache_spec.block_size > kv_spec.block_size + + +def test_auto_fit_max_model_len(): + """Test that max_model_len=-1 auto-fits to available GPU memory.""" + # Create config with original_max_model_len=-1 to trigger auto-fit + model_config = ModelConfig(max_model_len=1024) + # Simulate the user passing -1 by setting original_max_model_len + model_config.original_max_model_len = -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer + kv_cache_specs = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + } + + # With enough memory, max_model_len stays at the derived max + large_available_memory = mem_per_block_per_layer * 2 * 1024 # plenty of memory + _kv_cache_configs = get_kv_cache_configs( + vllm_config, [kv_cache_specs], [large_available_memory] + ) + assert vllm_config.model_config.max_model_len == 1024 + + # Reset for next test + model_config = ModelConfig(max_model_len=1024) + model_config.original_max_model_len = -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + + # With limited memory, max_model_len should be reduced + # Need memory for at least max_model_len tokens + # 32 blocks worth of memory for 2 layers = can fit 32*16=512 tokens + limited_memory = mem_per_block_per_layer * 2 * 32 + _kv_cache_configs = get_kv_cache_configs( + vllm_config, [kv_cache_specs], [limited_memory] + ) + # Should be reduced to fit in memory + assert vllm_config.model_config.max_model_len < 1024 + assert vllm_config.model_config.max_model_len > 0 + + +def test_auto_fit_max_model_len_with_hybrid(): + """Test that auto-fit works with hybrid KV cache specs.""" + # Create config with original_max_model_len=-1 to trigger auto-fit + model_config = ModelConfig(max_model_len=8192) + # Simulate the user passing -1 by setting original_max_model_len + model_config.original_max_model_len = -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer + gamma = 2 + kv_cache_specs = { + "layer_1": new_mamba_spec(num_speculative_blocks=gamma), + "layer_2": new_kv_cache_spec(), + } + + # One extra block on top of what a 1024-token request needs: the pool + # reserves one block as the null block. + available_memory = mem_per_block_per_layer * (1024 // 16 + 1 + gamma + 1) + _kv_cache_configs = get_kv_cache_configs( + vllm_config, [kv_cache_specs], [available_memory] + ) + assert vllm_config.model_config.max_model_len == 1024 + + +def test_auto_fit_max_model_len_not_triggered(): + """Test that auto-fit is not triggered when original_max_model_len is not -1.""" + model_config = ModelConfig(max_model_len=16) + # original_max_model_len should be None by default, not -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 + kv_cache_specs = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + } + + # This should work normally without auto-fit + _kv_cache_configs = get_kv_cache_configs( + vllm_config, [kv_cache_specs], [mem_per_block_per_layer * 2 * 32] + ) + assert vllm_config.model_config.max_model_len == 16 + + +def test_auto_fit_max_model_len_respects_num_gpu_blocks_override(): + """Auto-fit must size max_model_len against the override-clamped pool, not + the raw `available_memory`. Without this, auto-fit could pick a + max_model_len that no longer fits once `num_gpu_blocks_override` is applied. + """ + model_config = ModelConfig(max_model_len=16384) + model_config.original_max_model_len = -1 # request auto-fit + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + # Cap the cache to 32 blocks regardless of available memory. + vllm_config.cache_config.num_gpu_blocks_override = 32 + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 + kv_cache_specs = { + "layer_1": new_kv_cache_spec(), # block_size=16 + "layer_2": new_kv_cache_spec(), + } + # Plenty of raw memory (1024 blocks per layer would fit max_model_len=16384). + large_available_memory = mem_per_block_per_layer * 2 * 1024 + + get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) + + # 32 blocks * block_size 16 = 512 token slots, so max_model_len must + # auto-fit at or below that. + assert 0 < vllm_config.model_config.max_model_len <= 32 * 16 + + +def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override(): + """Admission check must use the override-clamped pool size, not raw + `available_memory`. Without this, startup could accept a max_model_len + that does not actually fit in `num_gpu_blocks_override` blocks. + """ + model_config = ModelConfig(max_model_len=16384) + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + # 32 blocks is far too small for max_model_len=16384 (would need 1024). + vllm_config.cache_config.num_gpu_blocks_override = 32 + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 + kv_cache_specs = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + } + # Plenty of raw memory: a bytes-only check against this would pass. + large_available_memory = mem_per_block_per_layer * 2 * 1024 + + with pytest.raises(ValueError, match="max seq len"): + get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) + + +def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes(): + """DFlash drafters can have a smaller head size than the target model. + + For example, MiMo uses 192-dim target KV heads while its DFlash draft uses + 128-dim KV heads. The resulting page sizes are 3:2 rather than an integer + block-size multiple, so the smaller page must be padded instead. + """ + # Attention layers read padded pages through a strided view, so padding is + # allowed. + target_spec = new_kv_cache_spec( + block_size=16, + num_kv_heads=1, + head_size=192, + dtype=torch.bfloat16, + ) + draft_spec = new_sliding_window_spec( + block_size=16, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + sliding_window=1024, + ) + + unified_specs = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "target_attn": target_spec, + "draft_attn": draft_spec, + } + ) + + assert unified_specs["target_attn"] == target_spec + unified_draft_spec = unified_specs["draft_attn"] + assert unified_draft_spec.block_size == draft_spec.block_size + assert unified_draft_spec.real_page_size_bytes == draft_spec.real_page_size_bytes + assert unified_draft_spec.page_size_padded == target_spec.page_size_bytes + assert unified_draft_spec.page_size_bytes == target_spec.page_size_bytes + + +def test_unify_hybrid_kv_cache_specs(): + # 1. has_full_attention and has_sliding_window + before_spec_1 = new_kv_cache_spec(block_size=64) + before_spec_2 = new_sliding_window_spec( + block_size=16, page_size_padded=32 * 1024, sliding_window=1024 + ) + kv_cache_spec = { + "layer_1": before_spec_1, + "layer_2": before_spec_2, + } + kv_cache_spec["draft_layer_1"] = replace(before_spec_1, dcp_sharded=False) + kv_cache_spec["draft_layer_2"] = replace(before_spec_2, dcp_sharded=False) + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + expected_spec_1 = new_kv_cache_spec(block_size=64) + expected_spec_2 = new_kv_cache_spec( + block_size=64, page_size_padded=64 * 1024, sliding_window=1024 + ) + assert kv_cache_spec["layer_1"] == expected_spec_1 + assert kv_cache_spec["layer_2"] == expected_spec_2 + assert kv_cache_spec["layer_2"].page_size_bytes == 64 * 1024 + assert kv_cache_spec["draft_layer_2"] == replace(expected_spec_2, dcp_sharded=False) + + # 2. has_full_attention and has_chunked_local_attention + before_spec_1 = new_kv_cache_spec() + before_spec_2 = new_chunked_local_attention_spec( + page_size_padded=32 * 1024, attention_chunk_size=512 + ) + kv_cache_spec = { + "layer_1": before_spec_1, + "layer_2": before_spec_2, + } + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + expected_spec_1 = new_kv_cache_spec() + expected_spec_2 = new_kv_cache_spec( + page_size_padded=32 * 1024, attention_chunk_size=512 + ) + + assert kv_cache_spec["layer_1"] == expected_spec_1 + assert kv_cache_spec["layer_2"] == expected_spec_2 + + # 3. has_full_attention, has_sliding_window and has_chunked_local_attention + before_spec_1 = new_kv_cache_spec() + before_spec_2 = new_sliding_window_spec( + page_size_padded=32 * 1024, sliding_window=1024 + ) + before_spec_3 = new_chunked_local_attention_spec( + page_size_padded=32 * 1024, attention_chunk_size=512 + ) + kv_cache_spec = { + "layer_1": before_spec_1, + "layer_2": before_spec_2, + "layer_3": before_spec_3, + } + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + expected_spec_1 = new_kv_cache_spec() + expected_spec_2 = new_kv_cache_spec(page_size_padded=32 * 1024, sliding_window=1024) + expected_spec_3 = new_kv_cache_spec( + page_size_padded=32 * 1024, attention_chunk_size=512 + ) + assert kv_cache_spec["layer_1"] == expected_spec_1 + assert kv_cache_spec["layer_2"] == expected_spec_2 + assert kv_cache_spec["layer_3"] == expected_spec_3 + + # 4. No FullAttentionSpec, should not convert + kv_cache_spec = { + "layer_1": new_sliding_window_spec(sliding_window=1024), + "layer_2": new_chunked_local_attention_spec(attention_chunk_size=512), + } + + with pytest.raises(ValueError): + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + + # Replicated Mamba state still requires the hybrid cache manager. + kv_cache_spec = {"attention": new_kv_cache_spec(), "mamba": new_mamba_spec()} + with pytest.raises(ValueError): + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + + +def test_unify_kv_cache_spec_page_size_mamba(): + """Regression test for https://github.com/vllm-project/vllm/issues/43626. + + MambaSpec's page_size_bytes is determined by its state shapes and does not + change with block_size, so unify_kv_cache_spec_page_size must pad the Mamba + page instead of scaling its block_size. This situation arises when a layer + with a page larger than the (already platform-aligned) Mamba page joins the + specs, e.g. a dense draft model with more KV heads than the hybrid main + model. + """ + # 1. Hybrid main model (Mamba + full attention, pages already aligned at + # 16KB) plus a dense draft model layer with a 2x larger page (32KB). + # Reproduces the bare AssertionError from #43626: 32768 % 16384 == 0, so + # the old code scaled the Mamba block_size, which left page_size_bytes + # unchanged at 16384. + mamba_spec = new_mamba_spec() # page_size_bytes = 16384 + main_attn_spec = new_kv_cache_spec() # page_size_bytes = 16384 + draft_attn_spec = new_kv_cache_spec(num_kv_heads=4) # page_size_bytes = 32768 + assert mamba_spec.page_size_bytes == main_attn_spec.page_size_bytes == 16384 + assert draft_attn_spec.page_size_bytes == 32768 + + unified = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "mamba_layer": mamba_spec, + "main_attn_layer": main_attn_spec, + "draft_attn_layer": draft_attn_spec, + } + ) + # Mamba page is padded; block_size (caching granularity) is unchanged. + assert unified["mamba_layer"].page_size_bytes == 32768 + assert unified["mamba_layer"].page_size_padded == 32768 + assert unified["mamba_layer"].block_size == mamba_spec.block_size + # Attention layer with smaller page still unifies by scaling block_size. + assert unified["main_attn_layer"].page_size_bytes == 32768 + assert unified["main_attn_layer"].block_size == 2 * main_attn_spec.block_size + assert unified["main_attn_layer"].page_size_padded is None + # Layer already at max page size is unchanged. + assert unified["draft_attn_layer"] == draft_attn_spec + + # 2. Mamba page already padded by the platform (state smaller than the + # padded page); the padding is re-applied at the new maximum. + padded_mamba_spec = new_mamba_spec( + shapes=((2, 256), (3, 32, 32)), page_size_padded=16384 + ) + assert padded_mamba_spec.page_size_bytes == 16384 + unified = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "mamba_layer": padded_mamba_spec, + "draft_attn_layer": draft_attn_spec, + } + ) + assert unified["mamba_layer"].page_size_bytes == 32768 + assert unified["mamba_layer"].page_size_padded == 32768 + + # 3. Mamba page that does not evenly divide the maximum page size is + # padded as well (the divisibility constraint only applies to block_size + # scaling). + odd_mamba_spec = new_mamba_spec(shapes=((6144,),)) + assert odd_mamba_spec.page_size_bytes == 24576 + assert 32768 % odd_mamba_spec.page_size_bytes != 0 + unified = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "mamba_layer": odd_mamba_spec, + "draft_attn_layer": draft_attn_spec, + } + ) + assert unified["mamba_layer"].page_size_bytes == 32768 + + # 4. Attention layers with non-divisible page sizes are padded too: every + # backend reads a padded page through the view's block stride, so there is + # no longer a case that must raise. + unified = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "attn_layer": new_kv_cache_spec(block_size=24), # 24576 + "draft_attn_layer": draft_attn_spec, # 32768 + } + ) + assert unified["attn_layer"].page_size_padded == 32768 + assert unified["attn_layer"].page_size_bytes == 32768 + assert unified["attn_layer"].block_size == 24 + + # 5. Uniform page sizes are returned unchanged. + specs = { + "mamba_layer": new_mamba_spec(), + "attn_layer": new_kv_cache_spec(), + } + assert kv_cache_planning.unify_kv_cache_spec_page_size(specs) == specs + + +@pytest.mark.parametrize("use_override", [True, False]) +def test_kv_cache_reserves_null_block_for_max_model_len(use_override): + """A KV cache sized to exactly the blocks a max_model_len request needs must + be rejected. BlockPool keeps one block as the null block, so only + num_blocks - 1 are usable; accepting num_blocks == needed would leave the + request unschedulable and hang the engine. Covers both the + num_gpu_blocks_override path and the memory-derived block count. + """ + block_size = 16 + max_model_len = 512 # needs 512 / 16 = 32 blocks + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + spec = new_kv_cache_spec(block_size=block_size) + + # 32 blocks -> only 31 usable after the null block: one short -> reject. + if use_override: + # Ample raw memory; the override alone constrains the block count. + vllm_config.cache_config.num_gpu_blocks_override = 32 + available_memory = [spec.page_size_bytes * 1024] + else: + available_memory = [spec.page_size_bytes * 32] + + with pytest.raises(ValueError, match="max seq len"): + get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) + + +def test_auto_fit_max_model_len_reserves_null_block(): + """Auto-fit (max_model_len=-1) must size max_model_len against usable + blocks, not the total pool. With memory for exactly 64 blocks, one is the + null block, so auto-fit must settle on 63 * block_size; picking the full + pool would leave a max-length request unschedulable and hang the engine. + """ + block_size = 16 + model_config = ModelConfig(max_model_len=1024) + model_config.original_max_model_len = -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + spec = new_kv_cache_spec(block_size=block_size) + + # Exactly the 1024 / 16 = 64 blocks a full-length request would need. + available_memory = [spec.page_size_bytes * 64] + + get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) + + assert vllm_config.model_config.max_model_len == 63 * block_size + + +def test_check_enough_kv_cache_memory_reserves_null_block(): + """The public admission check must reject a KV cache sized to exactly the + blocks a max_model_len request needs. BlockPool keeps one block as the + null block, so only num_blocks - 1 are usable; accepting + num_blocks == needed would leave the request unschedulable and hang the + engine. + """ + block_size = 16 + max_model_len = 512 # needs 512 / 16 = 32 blocks + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + spec = new_kv_cache_spec(block_size=block_size) + + # 32 blocks -> only 31 usable after the null block: one short -> reject. + with pytest.raises(ValueError, match="max seq len"): + check_enough_kv_cache_memory( + vllm_config, {"layer1": spec}, spec.page_size_bytes * 32 + ) + + # 33 blocks -> 32 usable after the null block -> accept. + check_enough_kv_cache_memory( + vllm_config, {"layer1": spec}, spec.page_size_bytes * 33 + ) + + +def _spec_decode_grouping_config(method="dspark", model_type=None): + """Grouping config with an EAGLE-family speculative method enabled.""" + return SimpleNamespace( + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), + cache_config=SimpleNamespace( + get_resolved_kv_cache_layout=lambda: SimpleNamespace( + is_block_outermost=True + ) + ), + model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type=model_type)), + speculative_config=SimpleNamespace( + method=method, + use_eagle=lambda: True, + use_eagle_block_drop=lambda: True, + ), + ) + + +def _hybrid_specs_with_draft(draft: bool, draft_shares_target_spec: bool = False): + """A K3-shaped hybrid: MLA full attention + Mamba, optionally plus a + DSpark-style draft MLA layer marked non_causal_multi_token_decode. + + The target's fp8 KV dtype is one difference from the draft on Kimi-K3 + (target `--kv-cache-dtype fp8_e4m3`, draft `auto`). Pass + draft_shares_target_spec to match that dtype while retaining the draft + marker as a group boundary. + """ + target_dtype = None if draft_shares_target_spec else "fp8_e4m3" + specs = { + "target.attn.0": new_mla_spec(block_size=64, cache_dtype_str=target_dtype), + "target.attn.1": new_mla_spec(block_size=64, cache_dtype_str=target_dtype), + "target.mamba.0": new_mamba_spec(block_size=64, mamba_cache_mode="align"), + "target.mamba.1": new_mamba_spec(block_size=64, mamba_cache_mode="align"), + } + if draft: + draft_spec = new_mla_spec(block_size=64) + specs["draft.attn.0"] = replace(draft_spec, non_causal_multi_token_decode=True) + return specs + + +def test_draft_group_annotated_on_hybrid_general_path(): + # A drafter's MLA layer carries non_causal_multi_token_decode, so its group + # is identifiable without keying off a model version. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) + ) + + flagged = [g for g in groups if g.is_eagle_group] + assert len(flagged) == 1 + assert "draft.attn.0" in flagged[0].layer_names + + +def test_mamba_groups_never_flagged_even_when_draft_shares_a_group(): + # Packed uniform-type groups can contain distinct target and draft layer + # specs; the combined group still holds volatile draft KV and must be + # flagged. What must never happen is a Mamba group being flagged: that + # widens its lookup window to two consecutive chunks, which align-mode + # checkpointing never produces, zeroing every lookup. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(), + _hybrid_specs_with_draft(draft=True, draft_shares_target_spec=True), + ) + + for group in groups: + if "draft.attn.0" in group.layer_names: + assert group.is_eagle_group + if isinstance(group.kv_cache_spec, MambaSpec): + assert not group.is_eagle_group + + +def test_draft_group_not_annotated_without_spec_decode(): + # The marker alone must not flag anything; the eagle semantics only apply + # when a speculative method is actually enabled. + config = _spec_decode_grouping_config() + config.speculative_config = None + groups = get_kv_cache_groups(config, _hybrid_specs_with_draft(draft=True)) + + assert not any(g.is_eagle_group for g in groups) + + +def test_unidentifiable_draft_with_mamba_warns(caplog_vllm): + # No group carries the draft marker, so consumers fall back to + # conservative behavior that silently breaks reuse for Mamba groups. + # That must at least be visible. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=False) + ) + + assert not any(g.is_eagle_group for g in groups) + assert "no KV cache group could be identified as the draft model's" in ( + caplog_vllm.text + ) + + +def test_unidentifiable_draft_without_mamba_does_not_warn(caplog_vllm): + # Pure-attention models degrade gracefully under the consumers' + # conservative fallback (a one-block hit drop at most), so the warning + # stays silent to avoid noise on every unannotated EAGLE deployment. + specs = { + "target.attn.0": new_mla_spec(block_size=64), + "target.attn.1": new_mla_spec(block_size=64), + } + groups = get_kv_cache_groups(_spec_decode_grouping_config(), specs) + + assert not any(g.is_eagle_group for g in groups) + assert "could be identified as the draft model's" not in caplog_vllm.text + + +def test_no_warning_when_draft_group_is_identified(caplog_vllm): + get_kv_cache_groups( + _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) + ) + + assert "could be identified as the draft model's" not in caplog_vllm.text + + +def _deepseek_v4_specs(model_version="deepseek_v4"): + """DeepseekV4-shaped specs: full MLA layers plus sliding-window MLA layers + at differing page sizes, with the MTP draft layer registered last.""" + return { + "model.layers.0.self_attn.attn": new_mla_spec(), + "model.layers.1.self_attn.attn": new_mla_spec(), + "model.layers.2.self_attn.attn": new_swa_mla_spec( + head_size=1024, model_version=model_version + ), + # The MTP block registers last, and its sliding-window size differs, so + # it lands in a group of its own. + "model.layers.3.self_attn.attn": new_swa_mla_spec( + head_size=1024, sliding_window=256, model_version=model_version + ), + } + + +@pytest.mark.parametrize( + ("method", "model_type"), + [("mtp", "deepseek_v4"), ("dspark", "deepseek_v4"), ("dspark", "deepseek_v41")], +) +def test_deepseek_v4_draft_group_annotated_on_packed_path(method, model_type): + # DeepseekV4's MTP block reuses the target's decoder layer, so its spec + # carries no draft marker and only the positional rule can find it. This + # pins the pre-existing behaviour that the unified annotator must preserve. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method=method, model_type=model_type), + _deepseek_v4_specs(model_version=None), + ) + + flagged = [g for g in groups if g.is_eagle_group] + assert len(flagged) == 1 + assert "model.layers.3.self_attn.attn" in flagged[0].layer_names + + +def test_trailing_layer_fallback_applies_to_any_mtp_model(): + # The positional rule is sound for every MTP drafter, not just DeepseekV4: + # MTP blocks reuse the target's decoder layer (no spec marker) and always + # register after every target layer. The model_type must not gate it. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method="mtp", model_type="other"), + _deepseek_v4_specs(), + ) + + flagged = [g for g in groups if g.is_eagle_group] + assert len(flagged) == 1 + assert "model.layers.3.self_attn.attn" in flagged[0].layer_names + + +def _qwen3_5_hybrid_specs(with_mtp_layer: bool): + """Qwen3.5-shaped hybrid: repeating [GDN x3, full-attn x1] blocks, with + the MTP drafter's full-attn layer (spec-identical to the target's) + registered last.""" + specs = {} + idx = 0 + for _ in range(2): + for _ in range(3): + specs[f"model.layers.{idx}.linear_attn"] = new_mamba_spec( + mamba_cache_mode="align" + ) + idx += 1 + specs[f"model.layers.{idx}.self_attn.attn"] = new_kv_cache_spec() + idx += 1 + if with_mtp_layer: + specs["mtp.layers.0.self_attn.attn"] = new_kv_cache_spec() + return specs + + +def test_qwen3_5_mtp_draft_group_annotated_on_hybrid_path(caplog_vllm): + # A hybrid mamba + full-attention model with an MTP drafter that is + # spec-indistinguishable from the target reaches the general multi-group + # path. The trailing-layer rule must locate the draft group there so the + # Mamba groups are not swept up by the flag-all consumer fallback. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method="mtp", model_type="qwen3_5"), + _qwen3_5_hybrid_specs(with_mtp_layer=True), + ) + + flagged = [g for g in groups if g.is_eagle_group] + assert len(flagged) == 1 + assert "mtp.layers.0.self_attn.attn" in flagged[0].layer_names + for group in groups: + if any( + isinstance(spec, MambaSpec) + for spec in iter_layer_specs(group.kv_cache_spec) + ): + assert not group.is_eagle_group + assert "could be identified as the draft model's" not in caplog_vllm.text + + +@pytest.mark.parametrize("method", ["eagle", "eagle3", "dspark"]) +def test_non_mtp_eagle_hybrid_still_warns(caplog_vllm, method): + # Other drafters are not covered by the trailing-layer rule, so an + # unidentifiable hybrid draft must still warn. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method=method, model_type="qwen3_5"), + _qwen3_5_hybrid_specs(with_mtp_layer=True), + ) + + assert not any(g.is_eagle_group for g in groups) + assert "could be identified as the draft model's" in caplog_vllm.text + + +def test_trailing_layer_fallback_requires_exact_partition(): + # If the groups do not partition the layers exactly (e.g. a caller that + # dropped or duplicated layers), the positional rule is meaningless and + # must not fire. + from vllm.v1.core.kv_cache_planning import _annotate_eagle_groups + + specs = _qwen3_5_hybrid_specs(with_mtp_layer=True) + config = _spec_decode_grouping_config(method="mtp", model_type="qwen3_5") + groups = get_kv_cache_groups(config, specs) + for g in groups: + g.is_eagle_group = False + # Remove one layer from its group: no longer an exact partition. + trimmed = [ + KVCacheGroupSpec( + [n for n in g.layer_names if n != "model.layers.0.linear_attn"], + g.kv_cache_spec, + ) + for g in groups + ] + _annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True) + + assert not any(g.is_eagle_group for g in trimmed) + + +_GQA_SPEC = FullAttentionSpec( + block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32 +) + + +_GQA_SWA_SPEC = SlidingWindowSpec( + block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32, sliding_window=128 +) + + +@pytest.mark.parametrize( + "specs,expected", + [ + pytest.param({"l.0": new_mla_spec(), "l.1": new_mla_spec()}, 4, id="mla"), + pytest.param( + { + "l.0": new_mla_spec(), + "l.1": new_mla_spec(), + "s.0": new_swa_mla_spec(), + "s.1": new_swa_mla_spec(), + }, + 4, + id="mla-and-swa-mla", + ), + pytest.param({"l.0": _GQA_SPEC, "l.1": _GQA_SPEC}, 1, id="gqa"), + pytest.param( + {"l.0": new_mla_spec(), "l.1": _GQA_SPEC}, + 1, + id="mla-and-gqa-uniform-group", + ), + pytest.param( + { + "l.0": new_mla_spec(), + "l.1": new_mla_spec(), + "s.0": _GQA_SWA_SPEC, + "s.1": _GQA_SWA_SPEC, + }, + 1, + id="mla-and-gqa-swa", + ), + ], +) +def test_kv_tp_replicas(monkeypatch, specs, expected): + """Resolved per layer before scheduler flattening hides mixed groups.""" + from vllm.config import ParallelConfig + from vllm.platforms import current_platform + + monkeypatch.setattr(current_platform, "device_count", lambda: 4) + parallel_config = ParallelConfig(tensor_parallel_size=4) + vllm_config = VllmConfig( + model_config=ModelConfig(max_model_len=16), parallel_config=parallel_config + ) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + mem = sum(s.page_size_bytes for s in specs.values()) * 10 + configs = get_kv_cache_configs(vllm_config, [specs], [mem]) + assert configs[0].kv_tp_replicas == expected + scheduler = generate_scheduler_kv_cache_config(configs) + assert scheduler.kv_tp_replicas == expected + + +@pytest.mark.parametrize( + "specs,tp_size,dcp_size,expected", + [ + pytest.param([new_mla_spec()], 8, 1, 8, id="mla"), + pytest.param([new_mla_spec()], 8, 2, 1, id="mla-dcp"), + pytest.param( + [replace(new_mla_spec(), max_tp_shards=None)], 8, 1, 1, id="mla-unset" + ), + pytest.param([replace(_GQA_SPEC, max_tp_shards=2)], 8, 1, 4, id="gqa-partial"), + pytest.param([replace(_GQA_SPEC, max_tp_shards=8)], 4, 1, 1, id="gqa-sharded"), + pytest.param( + [new_mla_spec(), replace(_GQA_SWA_SPEC, max_tp_shards=4)], + 8, + 1, + 2, + id="mla-and-gqa-partial", + ), + pytest.param( + [ + HiddenStateCacheSpec( + block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32 + ) + ], + 8, + 1, + 1, + id="hidden-state", + ), + ], +) +def test_kv_cache_groups_tp_replicas(specs, tp_size, dcp_size, expected): + """Replicas are the gcd of each layer's tp_size // max_tp_shards.""" + groups = [KVCacheGroupSpec([f"l.{i}"], spec) for i, spec in enumerate(specs)] + assert kv_cache_groups_tp_replicas(groups, tp_size, dcp_size) == expected diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index b06479f1caf1..1489315ec255 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -3,20 +3,23 @@ import copy import hashlib import importlib -import mmap -import subprocess -import sys from collections.abc import Callable from dataclasses import replace -from pathlib import Path from types import SimpleNamespace from typing import Any import pytest import torch +import vllm.v1.core.kv_cache_planning as kv_cache_planning import vllm.v1.core.kv_cache_utils as kv_cache_utils import vllm.v1.hisparse.runtime as hisparse_runtime_module +from tests.v1.core.utils import ( + new_kv_cache_spec, + new_mamba_spec, + new_mla_spec, + new_sliding_window_spec, +) from vllm.config import ( CacheConfig, ModelConfig, @@ -33,25 +36,21 @@ ) from vllm.sampling_params import SamplingParams from vllm.utils.hashing import sha256, sha256_cbor, xxhash, xxhash_cbor -from vllm.utils.mem_constants import GiB_bytes from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager +from vllm.v1.core.kv_cache_planning import ( + get_kv_cache_groups, +) from vllm.v1.core.kv_cache_utils import ( BlockHash, FreeKVCacheBlockQueue, KVCacheBlock, - check_enough_kv_cache_memory, - estimate_max_model_len, generate_block_hash_extra_keys, generate_scheduler_kv_cache_config, get_kv_cache_capacity, - get_kv_cache_configs, - get_kv_cache_groups, get_max_concurrency_for_kv_cache_config, get_request_block_hasher, hash_block_tokens, init_none_hash, - is_kv_cache_spec_uniform, - kv_cache_groups_tp_replicas, make_block_hash_with_group_id, tensor_data, to_event_extra_keys, @@ -62,10 +61,8 @@ get_hisparse_steady_state_concurrency, ) from vllm.v1.kv_cache_interface import ( - ChunkedLocalAttentionSpec, CircularBufferSpec, FullAttentionSpec, - HiddenStateCacheSpec, HiSparseHotSpec, HiSparseResidentSpec, KVCacheConfig, @@ -93,110 +90,6 @@ pytestmark = pytest.mark.cpu_test -@pytest.mark.parametrize("gpu_block_size", [32, 64]) -@pytest.mark.parametrize("shared_host_pool", [False, True]) -def test_hisparse_hma_uses_resolved_gpu_block_size( - monkeypatch, gpu_block_size, shared_host_pool -): - monkeypatch.setattr( - hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True - ) - specs = { - "model.layers.0.self_attn": MLAAttentionSpec( - block_size=gpu_block_size, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - is_index_group_leader=True, - ), - "model.layers.0.self_attn.indexer": MLAAttentionSpec( - block_size=gpu_block_size, - num_kv_heads=1, - head_size=128, - dtype=torch.bfloat16, - cache_role=SparseCacheRole.INDEXER, - ), - } - config = SimpleNamespace( - attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()), - model_config=SimpleNamespace( - hf_config=SimpleNamespace(index_topk=128), - max_model_len=gpu_block_size, - ), - parallel_config=SimpleNamespace( - tensor_parallel_size=2 if shared_host_pool else 1, - pipeline_parallel_size=1, - prefill_context_parallel_size=1, - decode_context_parallel_size=1, - world_size=2 if shared_host_pool else 1, - distributed_executor_backend="mp", - nnodes_within_dp=1, - ), - cache_config=SimpleNamespace( - num_gpu_blocks_override=7, - prefix_cache_retention_interval=None, - get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC, - ), - ) - monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30) - cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - config, get_hisparse_kv_cache_groups(config, specs), available_memory=2**30 - ) - assert cache_config.num_blocks == 7 - assert cache_config.hisparse_host_num_blocks is not None - assert cache_config.hisparse_host_num_blocks > 7 - - host_group, indexer_group, *auxiliary_groups = cache_config.kv_cache_groups - assert cache_config.hisparse_shared_host_pool is shared_host_pool - host_page = host_group.kv_cache_spec.page_size_bytes - alignment = mmap.PAGESIZE if shared_host_pool else 1 - expected_host_stride = (host_page + alignment - 1) // alignment * alignment - assert cache_config.hisparse_host_block_stride == expected_host_stride - assert cache_config.hisparse_host_num_blocks == 2**30 // expected_host_stride - assert host_group.host_resident - assert not any(group.host_resident for group in [indexer_group, *auxiliary_groups]) - host_layers = set(host_group.layer_names) - for tensor in cache_config.kv_cache_tensors: - assert all( - (name in host_layers) == tensor.host_resident for name in tensor.layers - ) - host_tensor = next(t for t in cache_config.kv_cache_tensors if t.host_resident) - host_spec = host_group.kv_cache_spec.kv_cache_specs[host_tensor.layers[0]] - assert host_tensor.block_stride == host_spec.page_size_bytes - assert host_tensor.layer_stride == ( - host_spec.page_size_bytes * cache_config.hisparse_host_num_blocks - ) - assert host_group.kv_cache_spec.block_size == gpu_block_size - assert indexer_group.kv_cache_spec.block_size == gpu_block_size - host_specs = host_group.kv_cache_spec.kv_cache_specs - gpu_indexer_specs = indexer_group.kv_cache_spec.kv_cache_specs - assert set(host_specs) == {"model.layers.0.self_attn"} - assert set(gpu_indexer_specs) == {"model.layers.0.self_attn.indexer"} - assert indexer_group.kv_cache_spec.prefix_cacheable - assert host_group.enable_kv_transfer - auxiliary_specs = [group.kv_cache_spec for group in auxiliary_groups] - assert any(isinstance(spec, HiSparseResidentSpec) for spec in auxiliary_specs) - assert any(isinstance(spec, HiSparseHotSpec) for spec in auxiliary_specs) - assert all( - spec.block_size == gpu_block_size - for spec in auxiliary_specs - if isinstance(spec, (HiSparseResidentSpec, HiSparseHotSpec)) - ) - scheduler_block_size, hash_block_size = kv_cache_utils.resolve_kv_cache_block_sizes( - cache_config, - SimpleNamespace( - cache_config=SimpleNamespace( - block_size=16, - enable_prefix_caching=True, - prefix_match_unit=None, - ), - parallel_config=SimpleNamespace(decode_context_parallel_size=1), - kv_transfer_config=object(), - ), - ) - assert scheduler_block_size == hash_block_size == gpu_block_size - - def test_hisparse_steady_state_concurrency_excludes_spilled_resident_pages(): """Running requests pin only their resident tail; the newest still needs its full in-flight window to be admitted.""" @@ -281,68 +174,6 @@ def test_hisparse_host_pool_must_fit_max_model_len(monkeypatch, extra_blocks, ok create_hisparse_layout(config, groups, host_budget=host_budget) -@pytest.mark.parametrize( - "max_model_len,num_gpu_blocks,ok", - [ - (8192, 160, False), - (8192, 1000, True), - (32768, 656, False), - (32768, 657, True), - (-1, 657, True), - ], -) -def test_hisparse_pool_must_fit_max_model_len( - monkeypatch, max_model_len, num_gpu_blocks, ok -): - """Each resident group takes its own blocks from HiSparse's shared GPU pool, - so a pool that fits only the indexer pages must be rejected at startup. - Resident pages are capped at the in-flight window: at 32768 tokens with 2048 - in flight, one request needs 512 indexer + 4 x 32 resident + 16 hot blocks, - plus the null block. Auto-fit (-1) must pick a length that passes the same - check.""" - monkeypatch.setattr( - hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True - ) - monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30) - specs: dict[str, KVCacheSpec] = {} - for i in range(4): - specs[f"model.layers.{i}.self_attn"] = MLAAttentionSpec( - block_size=64, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - is_index_group_leader=True, - ) - specs[f"model.layers.{i}.self_attn.indexer"] = MLAAttentionSpec( - block_size=64, - num_kv_heads=1, - head_size=128, - dtype=torch.bfloat16, - cache_role=SparseCacheRole.INDEXER, - ) - config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) - config.scheduler_config.max_num_batched_tokens = 2048 - # Async scheduling doubles the in-flight window, and with it the bound. - config.scheduler_config.async_scheduling = False - config.attention_config.hisparse_config = HiSparseConfig() - config.model_config.hf_config.index_topk = 128 - config.cache_config.num_gpu_blocks_override = num_gpu_blocks - config.cache_config.kv_cache_layout = "BLHNC" - # Workers report the resident/hot caches HiSparse derives. - specs |= { - name: group.kv_cache_spec - for group in get_hisparse_kv_cache_groups(config, specs) or [] - if isinstance(group.kv_cache_spec, (HiSparseResidentSpec, HiSparseHotSpec)) - for name in group.layer_names - } - - if ok: - kv_cache_utils.get_kv_cache_configs(config, [specs], [2**34]) - else: - with pytest.raises(ValueError, match="max seq len"): - kv_cache_utils.get_kv_cache_configs(config, [specs], [2**34]) - - def test_hisparse_derived_specs_do_not_affect_kv_cache_layout(monkeypatch): """Workers report HiSparse's resident/hot caches, which lay out their own raw backing; KV cache layout resolution must skip them.""" @@ -496,28 +327,6 @@ def test_effective_attention_block_size_matches_events(dcp): ] == [block_size] -def new_kv_cache_spec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.float32, - page_size_padded=None, - sliding_window=None, - attention_chunk_size=None, - kv_quant_mode=KVQuantMode.NONE, -): - return FullAttentionSpec( - block_size=block_size, - num_kv_heads=num_kv_heads, - head_size=head_size, - dtype=dtype, - page_size_padded=page_size_padded, - sliding_window=sliding_window, - attention_chunk_size=attention_chunk_size, - kv_quant_mode=kv_quant_mode, - ) - - def test_kv_cache_config_selects_only_transferable_groups(): """Connectors must see a stable projection of transfer-eligible groups.""" groups = [ @@ -581,60 +390,6 @@ def test_kv_cache_blocks_selects_requested_groups(): assert blocks.get_block_ids(group_ids=(0, 2)) == ([1], [3]) -def new_sliding_window_spec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.float32, - page_size_padded=None, - sliding_window=1, -): - return SlidingWindowSpec( - block_size=block_size, - num_kv_heads=num_kv_heads, - head_size=head_size, - dtype=dtype, - page_size_padded=page_size_padded, - sliding_window=sliding_window, - ) - - -def new_chunked_local_attention_spec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.float32, - page_size_padded=None, - attention_chunk_size=4, -): - return ChunkedLocalAttentionSpec( - block_size=block_size, - num_kv_heads=num_kv_heads, - head_size=head_size, - dtype=dtype, - page_size_padded=page_size_padded, - attention_chunk_size=attention_chunk_size, - ) - - -def new_mamba_spec( - block_size=16, - shapes=((2, 512), (3, 32, 32)), - dtypes=(torch.float32, torch.float32), - num_speculative_blocks=2, - mamba_cache_mode="none", - page_size_padded=None, -): - return MambaSpec( - block_size=block_size, - shapes=shapes, - dtypes=dtypes, - page_size_padded=page_size_padded, - mamba_cache_mode=mamba_cache_mode, - num_speculative_blocks=num_speculative_blocks, - ) - - @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_none_hash(monkeypatch, hash_fn): import vllm.v1.core.kv_cache_utils @@ -1441,430 +1196,6 @@ def test_metrics_empty_stats(): assert metrics.hit_rate == 0.75 -def test_get_kv_cache_configs_multiple_workers(): - model_config = ModelConfig(max_model_len=16) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - vllm_config.cache_config.prefix_cache_retention_interval = None - - ref_kv_cache_spec = new_kv_cache_spec() - same_kv_cache_specs = [ - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - ] - - # Basic case. All things are the same. - kv_cache_configs = get_kv_cache_configs( - vllm_config, - same_kv_cache_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - expected = KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer1", "layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), - ], - ) - assert kv_cache_configs == [expected, expected] - - # Different available memory. This is the case for TP. - # Use the smallest memory available. - kv_cache_configs = get_kv_cache_configs( - vllm_config, - same_kv_cache_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 20, - ], - ) - assert kv_cache_configs == [expected, expected] - - # Different KV cache specs. This is the case for PP. - different_layer_specs = [ - { - "layer1": new_kv_cache_spec(), - }, - { - "layer2": new_kv_cache_spec(), - "layer3": new_kv_cache_spec(), - }, - ] - - # Different workers have different layers. - kv_cache_configs = get_kv_cache_configs( - vllm_config, - different_layer_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer1"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1"], new_kv_cache_spec()), - ], - ), - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer2", "layer3"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer2", "layer3"], new_kv_cache_spec()), - ], - ), - ] - - # Some layers are the same, some are different. This is the case for TP+PP - tp_pp_kv_cache_specs = [ - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - { - "layer3": new_kv_cache_spec(), - }, - { - "layer3": new_kv_cache_spec(), - }, - ] - - kv_cache_configs = get_kv_cache_configs( - vllm_config, - tp_pp_kv_cache_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - expected_12 = KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer1", "layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), - ], - ) - expected_3 = KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer3"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer3"], ref_kv_cache_spec), - ], - ) - assert kv_cache_configs == [ - expected_12, - expected_12, - expected_3, - expected_3, - ] - - # Different workers have different types of layers. This is the case for - # hybrid models + PP. - different_type_layer_specs = [ - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - { - "layer3": new_sliding_window_spec(), - "layer4": new_sliding_window_spec(), - }, - ] - kv_cache_configs = get_kv_cache_configs( - vllm_config, - different_type_layer_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer1", "layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), - KVCacheGroupSpec([], new_sliding_window_spec()), - ], - ), - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer3", "layer4"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec([], ref_kv_cache_spec), - KVCacheGroupSpec(["layer3", "layer4"], new_sliding_window_spec()), - ], - ), - ] - - # When divided into multiple KVCacheGroups, need to ensure the number of - # layers per group is similar. - different_type_layer_specs = [ - { - "layer1": new_kv_cache_spec(), - "layer2": new_sliding_window_spec(), - "layer3": new_sliding_window_spec(), - }, - { - "layer4": new_kv_cache_spec(), - "layer5": new_sliding_window_spec(), - "layer6": new_sliding_window_spec(), - }, - ] - kv_cache_configs = get_kv_cache_configs( - vllm_config, - different_type_layer_specs, - [ - ref_kv_cache_spec.page_size_bytes * 10, - ref_kv_cache_spec.page_size_bytes * 10, - ], - ) - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer1"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer3"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1"], ref_kv_cache_spec), - KVCacheGroupSpec(["layer2"], new_sliding_window_spec()), - KVCacheGroupSpec(["layer3"], new_sliding_window_spec()), - ], - ), - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer4"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer5"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer6"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer4"], ref_kv_cache_spec), - KVCacheGroupSpec(["layer5"], new_sliding_window_spec()), - KVCacheGroupSpec(["layer6"], new_sliding_window_spec()), - ], - ), - ] - - # Have conflicting layers. Need to raise an error. - conflicting_layer_specs = [ - { - "layer1": new_kv_cache_spec(), - }, - { - "layer1": new_sliding_window_spec(), - }, - ] - with pytest.raises(AssertionError): - get_kv_cache_configs( - vllm_config, - conflicting_layer_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - - -@pytest.mark.parametrize( - "asymmetric_memory", - [False, True], - ids=["symmetric", "asymmetric"], -) -def test_get_kv_cache_configs_pp_sharding(asymmetric_memory): - model_config = ModelConfig(max_model_len=512) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - vllm_config.cache_config.prefix_cache_retention_interval = None - - ref_kv_cache_spec = new_kv_cache_spec() - pp_kv_cache_specs = [ - {"layer1": ref_kv_cache_spec}, - {"layer2": ref_kv_cache_spec}, - ] - - expected_num_blocks = model_config.max_model_len // ref_kv_cache_spec.block_size + 1 - avail_memory = ref_kv_cache_spec.page_size_bytes * expected_num_blocks - - # With per-worker validation, each worker only needs memory for its own - # layers. Worker 2 having more memory shouldn't affect worker 1's config. - available_memory = ( - [avail_memory, avail_memory * 2] if asymmetric_memory else [avail_memory] * 2 - ) - - kv_cache_configs = get_kv_cache_configs( - vllm_config, - pp_kv_cache_specs, - available_memory, - ) - - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=expected_num_blocks, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, - layers=["layer1"], - layer_stride=ref_kv_cache_spec.page_size_bytes - * expected_num_blocks, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[KVCacheGroupSpec(["layer1"], ref_kv_cache_spec)], - ), - KVCacheConfig( - num_blocks=expected_num_blocks, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, - layers=["layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes - * expected_num_blocks, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[KVCacheGroupSpec(["layer2"], ref_kv_cache_spec)], - ), - ] - - -def test_project_kv_cache_groups_to_worker(): - spec_a = new_kv_cache_spec() - spec_b = new_kv_cache_spec(num_kv_heads=4) - - global_groups = [ - KVCacheGroupSpec(["layer1", "layer2", "layer3"], spec_a), - ] - worker_spec = {"layer1": spec_a, "layer2": spec_a} - projected = kv_cache_utils._project_kv_cache_groups_to_worker( - global_groups, worker_spec - ) - assert len(projected) == 1 - assert projected[0].layer_names == ["layer1", "layer2"] - assert projected[0].kv_cache_spec is spec_a - - projected = kv_cache_utils._project_kv_cache_groups_to_worker( - global_groups, {"layer4": spec_a} - ) - assert len(projected) == 1 - assert projected[0].layer_names == [] - assert projected[0].kv_cache_spec is spec_a - - uniform_spec = UniformTypeKVCacheSpecs( - block_size=16, - kv_cache_specs={"layer1": spec_a, "layer2": spec_b, "layer3": spec_a}, - ) - global_groups_uniform = [ - KVCacheGroupSpec(["layer1", "layer2", "layer3"], uniform_spec), - ] - projected = kv_cache_utils._project_kv_cache_groups_to_worker( - global_groups_uniform, {"layer1": spec_a, "layer3": spec_a} - ) - assert len(projected) == 1 - assert projected[0].layer_names == ["layer1", "layer3"] - proj_spec = projected[0].kv_cache_spec - assert isinstance(proj_spec, UniformTypeKVCacheSpecs) - assert set(proj_spec.kv_cache_specs.keys()) == {"layer1", "layer3"} - - @pytest.mark.parametrize("sliding_window", [None, 256]) @pytest.mark.parametrize("disable_hybrid", [False, True]) @pytest.mark.parametrize("pcp_size", [1, 4]) @@ -1901,7 +1232,7 @@ def test_dcp_target_allocates_replicated_draft_independently( widths = [g.kv_cache_spec.max_num_blocks_per_req(config, 1024) for g in groups] assert widths == [16, 64] - cache_config = kv_cache_utils.get_kv_cache_config_from_groups( + cache_config = kv_cache_planning.get_kv_cache_config_from_groups( config, groups, available_memory=16 * 1024 * 1024 ) scheduler_config = generate_scheduler_kv_cache_config([cache_config]) @@ -1971,58 +1302,6 @@ def test_full_attention_merge_preserves_replicated_cache_geometry(use_mla): type(draft).merge([draft, replace(draft, dcp_sharded=True)]) -@pytest.mark.parametrize("with_draft", [False, True]) -@pytest.mark.parametrize("indexer_alignment", [None, 512]) -def test_sparse_mla_preserves_physical_row_addressing(with_draft, indexer_alignment): - from vllm.v1.attention.backends.mla.sparse_utils import flat_kv_row_view - from vllm.v1.worker.utils import allocate_kv_cache - - config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) - config.cache_config.kv_cache_layout = "BLHNC" - config.parallel_config.decode_context_parallel_size = 4 - common = dict(block_size=64, num_kv_heads=1, dtype=torch.uint8) - specs = { - "target": MLAAttentionSpec( - **common, - head_size=576, - state_content_bytes=656, - block_stride_alignment=656, - cache_dtype_str="fp8_ds_mla", - is_index_group_leader=True, - ), - "indexer": MLAAttentionSpec( - **common, - head_size=132, - cache_role=SparseCacheRole.INDEXER, - block_stride_alignment=indexer_alignment, - ), - "draft": SlidingWindowSpec( - block_size=64, - num_kv_heads=64, - head_size=64, - dtype=torch.bfloat16, - sliding_window=2048, - dcp_sharded=False, - ), - } - if not with_draft: - del specs["draft"] - groups = get_kv_cache_groups(config, specs) - cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - config, groups, available_memory=8 * 1024 * 1024 - ) - caches = allocate_kv_cache( - cache_config, torch.device("cpu"), KVCacheLayout.BLHNC, [64] * len(groups) - ) - cache = caches["target"].squeeze(1) - if indexer_alignment is not None: - assert caches["indexer"].stride(0) % indexer_alignment == 0 - rows, stride_rows = flat_kv_row_view(cache, 64) - cache[1, 0].fill_(7) - torch.testing.assert_close(rows[stride_rows], cache[1, 0]) - assert torch.all(rows[stride_rows] == 7) - - @pytest.mark.parametrize( "layer_type,dcp_size,expected_width", [ @@ -2115,103 +1394,6 @@ def test_merge_kv_cache_spec(): assert merged_layer_spec.sliding_window == 1 -def test_is_kv_cache_spec_uniform(): - kv_cache_spec = { - "layer_1": new_kv_cache_spec(num_kv_heads=32), - "layer_2": new_kv_cache_spec(num_kv_heads=32), - } - assert is_kv_cache_spec_uniform(kv_cache_spec) - - kv_cache_spec = { - "layer_1": new_kv_cache_spec(num_kv_heads=32), - "layer_2": new_kv_cache_spec(num_kv_heads=32, sliding_window=1), - } - assert is_kv_cache_spec_uniform(kv_cache_spec) - - kv_cache_spec = { - "layer_1": new_kv_cache_spec(num_kv_heads=32), - "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), - } - assert not is_kv_cache_spec_uniform(kv_cache_spec) - - kv_cache_spec = { - "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), - "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), - } - assert is_kv_cache_spec_uniform(kv_cache_spec) - - kv_cache_spec = { - "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), - "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=2), - } - assert not is_kv_cache_spec_uniform(kv_cache_spec) - - script = """ -import sys - -from vllm.v1.core.kv_cache_utils import is_kv_cache_spec_uniform -from vllm.v1.kv_cache_interface import KVCacheSpec - -if sys.flags.optimize < 1: - raise RuntimeError("subprocess is not running with optimization enabled") -specs = { - "a": KVCacheSpec(block_size=1), - "b": KVCacheSpec(block_size=2), -} -if is_kv_cache_spec_uniform(specs): - raise RuntimeError("different specs were treated as uniform") -""" - subprocess.run( - [sys.executable, "-O", "-c", script], - cwd=Path(__file__).resolve().parents[3], - check=True, - timeout=60, - ) - - -@pytest.mark.parametrize( - ("model_id", "max_model_len", "want_estimated_max_len"), - [ - ("Qwen/Qwen1.5-7B", 16385, 16384), - ("Qwen/Qwen1.5-7B", 16383, 16383), - ], -) -def test_estimate_max_model_len(model_id, max_model_len, want_estimated_max_len): - # Create a VllmConfig - model_config = ModelConfig( - model_id, - runner="generate", - dtype="float16", - max_model_len=max_model_len, - ) - scheduler_config = SchedulerConfig( - max_num_batched_tokens=32768, - max_model_len=model_config.max_model_len, - is_encoder_decoder=model_config.is_encoder_decoder, - ) - - vllm_config = VllmConfig( - model_config=model_config, - scheduler_config=scheduler_config, - ) - - # Create KV cache specs - kv_cache_spec = {} - for i in range(32): - layer_name = f"layer_{i}" - kv_cache_spec[layer_name] = FullAttentionSpec( - block_size=16, - num_kv_heads=32, - head_size=128, - dtype=torch.float16, - ) - # Estimate the maximum model length, 16384 model_len need 8GB - estimated_max_len = estimate_max_model_len( - vllm_config, kv_cache_spec, 8 * GiB_bytes - ) - assert estimated_max_len == want_estimated_max_len - - def test_get_max_concurrency_for_kv_cache_config(): # Create a VllmConfig model_id = "Qwen/Qwen1.5-7B" @@ -2425,411 +1607,6 @@ def test_allocate_with_lookahead(): assert len(blocks.get_block_ids()[0]) == 2 -def test_get_kv_cache_config_one_worker(): - # pass max_model_len to pass check_enough_kv_cache_memory - model_config = ModelConfig(max_model_len=16) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - vllm_config.cache_config.prefix_cache_retention_interval = None - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 - # all layers are full attention -> single group - kv_cache_specs_full = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - } - kv_cache_config_full = get_kv_cache_configs( - vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32] - )[0] - print(kv_cache_config_full) - assert kv_cache_config_full == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], - ) - - # all layers are sliding window -> single group - kv_cache_specs_sliding = { - "layer_1": new_sliding_window_spec(), - "layer_2": new_sliding_window_spec(), - } - kv_cache_config_sliding = get_kv_cache_configs( - vllm_config, [kv_cache_specs_sliding], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_sliding == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1", "layer_2"], new_sliding_window_spec()) - ], - ) - - # full + sliding, but disable_hybrid_kv_cache_manager - vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = True - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_sliding_window_spec(), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1", "layer_2"], new_kv_cache_spec(sliding_window=1) - ), - ], - ) - vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False - - # full + sliding, with hybrid_kv_cache_manager - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_sliding_window_spec(), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=64, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 64, - layers=["layer_1"], - layer_stride=mem_per_block_per_layer * 64, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 64, - layers=["layer_2"], - layer_stride=mem_per_block_per_layer * 64, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1"], new_kv_cache_spec()), - KVCacheGroupSpec(["layer_2"], new_sliding_window_spec()), - ], - ) - - # 2 full + 4 sliding, 2 layers per group - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - "layer_3": new_sliding_window_spec(), - "layer_4": new_sliding_window_spec(), - "layer_5": new_sliding_window_spec(), - "layer_6": new_sliding_window_spec(), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_3", "layer_5"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_4", "layer_6"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec()), - KVCacheGroupSpec(["layer_3", "layer_5"], new_sliding_window_spec()), - KVCacheGroupSpec(["layer_4", "layer_6"], new_sliding_window_spec()), - ], - ) - - # 3 full + 7 sliding, pad to 3 full + 9 sliding - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - "layer_3": new_kv_cache_spec(), - "layer_4": new_sliding_window_spec(), - "layer_5": new_sliding_window_spec(), - "layer_6": new_sliding_window_spec(), - "layer_7": new_sliding_window_spec(), - "layer_8": new_sliding_window_spec(), - "layer_9": new_sliding_window_spec(), - "layer_10": new_sliding_window_spec(), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_1", "layer_2", "layer_3"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_4", "layer_7", "layer_10"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_5", "layer_8"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_6", "layer_9"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1", "layer_2", "layer_3"], new_kv_cache_spec()), - KVCacheGroupSpec( - ["layer_4", "layer_7", "layer_10"], new_sliding_window_spec() - ), - KVCacheGroupSpec(["layer_5", "layer_8"], new_sliding_window_spec()), - KVCacheGroupSpec(["layer_6", "layer_9"], new_sliding_window_spec()), - ], - ) - - # 6 full + 5 sliding, pad to 6 full + 6 sliding. This is a typical case for gpt-oss - # eagle where there is only one more full attention layer than sliding window layers - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - "layer_3": new_kv_cache_spec(), - "layer_4": new_kv_cache_spec(), - "layer_5": new_kv_cache_spec(), - "layer_6": new_kv_cache_spec(), - "layer_7": new_sliding_window_spec(), - "layer_8": new_sliding_window_spec(), - "layer_9": new_sliding_window_spec(), - "layer_10": new_sliding_window_spec(), - "layer_11": new_sliding_window_spec(), - } - - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 6 * 32] - )[0] - print(kv_cache_config_hybrid) - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 6, - layers=[ - "layer_1", - "layer_2", - "layer_3", - "layer_4", - "layer_5", - "layer_6", - ], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 6, - layers=["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1", "layer_2", "layer_3", "layer_4", "layer_5", "layer_6"], - new_kv_cache_spec(), - ), - KVCacheGroupSpec( - ["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], - new_sliding_window_spec(), - ), - ], - ) - - # different hidden size but same type, use UniformTypeKVCacheSpecs - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(head_size=128), - "layer_2": new_kv_cache_spec(head_size=64), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] - )[0] - # Layers of different page sizes pack densely into one allocation: the - # 2x-larger layer_1 takes the first region, layer_2 the next. - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_1"], - layer_stride=mem_per_block_per_layer * 2 * 32, - block_stride=mem_per_block_per_layer * 2, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - offset=mem_per_block_per_layer * 2 * 32, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1", "layer_2"], - UniformTypeKVCacheSpecs( - block_size=16, kv_cache_specs=kv_cache_specs_hybrid - ), - ) - ], - ) - - # Different hidden size and different type, align by different block size - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(head_size=64), - "layer_2": new_sliding_window_spec(head_size=32), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32, - layers=["layer_1"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32, - layers=["layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1"], new_kv_cache_spec(head_size=64)), - KVCacheGroupSpec( - ["layer_2"], new_sliding_window_spec(head_size=32, block_size=32) - ), - ], - ) - - # different hidden size that cannot be aligned by using different block size, - # but can be aligned by padding the smaller physical page. - swa_spec = new_sliding_window_spec(head_size=96) - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(head_size=64), - "layer_2": swa_spec, - } - - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] - )[0] - padded_page_size = swa_spec.page_size_bytes - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=42, - kv_cache_tensors=[ - KVCacheTensor( - size=padded_page_size * 42, - layers=["layer_1"], - layer_stride=padded_page_size * 42, - block_stride=padded_page_size, - ), - KVCacheTensor( - size=padded_page_size * 42, - layers=["layer_2"], - layer_stride=padded_page_size * 42, - block_stride=padded_page_size, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1"], - new_kv_cache_spec( - head_size=64, - page_size_padded=padded_page_size, - ), - ), - KVCacheGroupSpec( - ["layer_2"], - new_sliding_window_spec(head_size=96), - ), - ], - ) - - # Test num_gpu_blocks_override - vllm_config.cache_config.num_gpu_blocks_override = 16 - kv_cache_config_override_blocks = get_kv_cache_configs( - vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_override_blocks == KVCacheConfig( - num_blocks=16, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 16 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 16, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], - ) - - -def test_get_kv_cache_configs_attention_free(): - kv_cache_specs: dict[str, KVCacheSpec] = {} - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) - vllm_config.cache_config.prefix_cache_retention_interval = None - kv_cache_configs = get_kv_cache_configs(vllm_config, [kv_cache_specs], [0]) - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=1, - kv_cache_tensors=[], - kv_cache_groups=[], - ) - ] - - def test_generate_uniform_type_kv_cache_specs(): # All layers are full attention, can be merged kv_cache_specs = { @@ -2879,223 +1656,28 @@ def test_generate_uniform_type_kv_cache_specs(): def test_generate_scheduler_kv_cache_config(): kv_cache_specs = { "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(head_size=128), - } - kv_cache_configs = [ - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1", "layer_2"], - UniformTypeKVCacheSpecs( - block_size=16, kv_cache_specs=kv_cache_specs - ), - ), - ], - ) - ] - scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs) - assert scheduler_kv_cache_config == KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[], - kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], - ) - - -def test_packed_groups_glm5_like_hybrid(): - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=8192)) - vllm_config.cache_config.kv_cache_layout = "BLHNC" - kpool, block_size = 4, 1024 - kv_cache_spec: dict[str, KVCacheSpec] = {} - for i in range(11): - kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - block_stride_alignment=32 * 576 * 2, - ) - kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=132, - dtype=torch.uint8, - tokens_per_state=kpool, - block_stride_alignment=64 * 132, - ) - kv_cache_spec[f"layers.{i}.tail"] = CircularBufferSpec( - block_size=kpool, - num_kv_heads=2, - head_size=128, - head_size_v=0, - dtype=torch.bfloat16, - ) - for i in range(34): - kv_cache_spec[f"layers.{i}.kda"] = MambaSpec( - shapes=((3, 6144), (16, 128, 128)), - dtypes=(torch.bfloat16, torch.float32), - block_size=block_size, - ) - - groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - attn = [ - g - for g in groups - if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) - and isinstance(g.kv_cache_spec.first_spec, MLAAttentionSpec) - ] - rings = [ - g - for g in groups - if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) - and isinstance(g.kv_cache_spec.first_spec, CircularBufferSpec) - ] - mamba = [ - g - for g in groups - if isinstance( - getattr(g.kv_cache_spec, "first_spec", g.kv_cache_spec), MambaSpec - ) - ] - assert len(attn) == 1 and len(attn[0].layer_names) == 22 - assert len(rings) == 1 and len(rings[0].layer_names) == 11 - # 11 * (1.18 MB MLA + 34 KB indexer) per block fits 12 fp32-SSM KDA states. - assert len(mamba) == 3 - assert sum(len(g.layer_names) for g in mamba) == 34 - - kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, groups, available_memory=2 * 1024**3 - ) - strides = {t.block_stride for t in kv_cache_config.kv_cache_tensors} - assert len(strides) == 1 - (block_stride,) = strides - # Blocks sit a whole number of 32-row MLA pages and 64-state indexer pages - # apart so both caches can be re-paged for their kernels. - assert block_stride % (32 * 576 * 2) == 0 - assert block_stride % (64 * 132) == 0 - - -def new_mla_spec(cache_dtype_str=None, block_size: int = 16): - # head_size = kv_lora_rank(512) + qk_rope_head_dim(64) = 576 - return MLAAttentionSpec( - block_size=block_size, - num_kv_heads=1, - max_tp_shards=1, - head_size=576, - dtype=torch.float32, - cache_dtype_str=cache_dtype_str, - ) - - -def new_swa_mla_spec(head_size=576, sliding_window=128, model_version=None): - return SlidingWindowMLASpec( - block_size=16, - num_kv_heads=1, - max_tp_shards=1, - head_size=head_size, - dtype=torch.float32, - sliding_window=sliding_window, - model_version=model_version, - ) - - -def new_indexer_mla_spec(block_size=16): - # Sparse-attention indexer k_cache: an MLAAttentionSpec with a much smaller - # page size than the main MLA attention (uint8, small head), so their pages - # cannot be unified. - return MLAAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=132, - dtype=torch.uint8, - ) - - -def test_mixed_page_size_groups_use_spec_compatibility(): - specs = {} - for i in range(3): - specs[f"mla.{i}"] = new_mla_spec() - specs[f"indexer.{i}"] = new_indexer_mla_spec() - specs.update({f"swa.{i}": new_swa_mla_spec(head_size=1024) for i in range(5)}) - - config = _grouping_config() - config.cache_config = CacheConfig() - config.cache_config.kv_cache_layout = "BLNHC" - groups = get_kv_cache_groups(config, specs) - - assert len(groups) == 3 - assert {name for group in groups for name in group.layer_names} == set(specs) - assert sorted(len(group.layer_names) for group in groups) == [2, 3, 6] - - -def _grouping_config(): - cache_config = CacheConfig() - cache_config.kv_cache_layout = "LBNHC" - return SimpleNamespace( - scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), - speculative_config=None, - cache_config=cache_config, - ) - - -def test_hidden_state_group_preserves_hybrid_prefix_cache_granularity(): - block_size = 544 - full_spec = FullAttentionSpec( - block_size=block_size, - num_kv_heads=1, - # FullAttentionSpec stores both K and V, so this produces a - # 544 * 1 * (512 + 512) * 2 = 1,114,112-byte page. - head_size=512, - dtype=torch.bfloat16, - ) - mamba_spec = MambaSpec( - block_size=block_size, - shapes=((557056,),), - dtypes=(torch.bfloat16,), - mamba_cache_mode="align", - ) - hidden_spec = HiddenStateCacheSpec( - block_size=block_size, - num_kv_heads=3, - head_size=1024, - dtype=torch.bfloat16, - ) - assert full_spec.page_size_bytes == mamba_spec.page_size_bytes - - groups = get_kv_cache_groups( - _grouping_config(), - { - "model.full_attn": full_spec, - "model.mamba": mamba_spec, - "cache_only_layers.0": hidden_spec, - }, - ) - - hidden_group = next( - group - for group in groups - if isinstance(group.kv_cache_spec, HiddenStateCacheSpec) - ) - assert hidden_group.kv_cache_spec.block_size == 136 - - kv_cache_config = KVCacheConfig( - num_blocks=1, + "layer_2": new_kv_cache_spec(head_size=128), + } + kv_cache_configs = [ + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1", "layer_2"], + UniformTypeKVCacheSpecs( + block_size=16, kv_cache_specs=kv_cache_specs + ), + ), + ], + ) + ] + scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs) + assert scheduler_kv_cache_config == KVCacheConfig( + num_blocks=10, kv_cache_tensors=[], - kv_cache_groups=groups, - ) - vllm_config = SimpleNamespace( - cache_config=SimpleNamespace( - block_size=16, - enable_prefix_caching=True, - prefix_match_unit=None, - ), - parallel_config=SimpleNamespace(decode_context_parallel_size=1), - kv_transfer_config=object(), + kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], ) - assert kv_cache_utils.resolve_kv_cache_block_sizes( - kv_cache_config, vllm_config - ) == (544, 136) def test_resolve_dcp_kv_block_size_unwraps_uniform_type_specs(): @@ -3133,125 +1715,6 @@ def test_resolve_dcp_kv_block_size_unwraps_uniform_type_specs(): assert kv_cache_utils.resolve_dcp_kv_cache_spec(wrapped_mamba, 4) is wrapped_mamba -def test_multi_run_layer_compact_strides_place_hoisted_heads(): - """A layer-compact run region is its own dense allocation: under LHBNC the head - groups sit between the layers and the blocks, so a run's block stride is one head - group's slice, not the whole page (regression test for setStorage out-of-bounds - on ROCm hybrid models).""" - full = new_kv_cache_spec() - swa = new_sliding_window_spec(sliding_window=full.block_size * 2) - page = full.page_size_bytes - assert swa.page_size_bytes == page - - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) - vllm_config.cache_config.kv_cache_layout = "LHBNC" - config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, - [ - KVCacheGroupSpec(["full.0", "full.1"], full), - KVCacheGroupSpec(["swa.0", "swa.1"], swa), - ], - available_memory=4 * page * 8, - ) - - num_blocks = config.num_blocks - head_group = full.block_size * full.state_content_size_bytes - for tensor in config.kv_cache_tensors: - # One head group per block, head groups between the layers and blocks. - assert tensor.block_stride == head_group - assert tensor.layer_stride == full.num_heads * num_blocks * head_group - # The last block of the last layer's last head group must stay within the - # allocation (the old page-based stride overran it). - end = ( - tensor.offset - + tensor.layer_stride * (len(tensor.layers) - 1) - + (full.num_heads - 1) * num_blocks * head_group - + tensor.block_stride * (num_blocks - 1) - + head_group - ) - assert end <= tensor.size - - -def test_mla_draft_prefers_standard_layout_when_pages_can_be_unified(): - specs = { - "target.0.attn": new_mla_spec(), - "draft.0": new_sliding_window_spec(num_kv_heads=1, head_size=288), - } - assert len({spec.page_size_bytes for spec in specs.values()}) == 1 - - groups = get_kv_cache_groups(_grouping_config(), specs) - - assert len(groups) == 2 - assert all( - not isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in groups - ) - - -def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog_vllm): - # Sparse MLA pages cannot be padded safely. Keeping the draft's attention - # compute sliding-window while promoting only its allocation semantics lets - # every layer share the target's block table and remain contiguous. - draft = new_sliding_window_spec(block_size=16) - specs = { - "target.0.attn": new_mla_spec(block_size=64), - "target.0.indexer": new_indexer_mla_spec(block_size=64), - "draft.0": draft, - } - - groups = get_kv_cache_groups(_grouping_config(), specs) - assert len(groups) == 1 - assert set(groups[0].layer_names) == set(specs) - group_spec = groups[0].kv_cache_spec - assert isinstance(group_spec, UniformTypeKVCacheSpecs) - assert group_spec.block_size == 64 - promoted_draft = group_spec.kv_cache_specs["draft.0"] - assert isinstance(promoted_draft, FullAttentionSpec) - assert not isinstance(promoted_draft, SlidingWindowSpec) - assert promoted_draft.block_size == 64 - assert promoted_draft.sliding_window == draft.sliding_window - assert specs["draft.0"] is draft - assert "attention compute is unchanged" in caplog_vllm.text - - -def test_hidden_states_with_tp_scales_page_size(): - """When TP shrinks KV pages below the hidden-state per-token cost, - get_kv_cache_groups must scale up target block sizes so that the - common page accommodates the unsharded hidden states.""" - # Simulate TP=4 sharding a model with 8 KV heads → 2 per rank. - # KV page = block_size(16) * num_kv_heads(2) * head_size(64) * dtype(2) - # = 16 * 2 * 64 * 2 = 4096 bytes. - kv_spec = new_kv_cache_spec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.bfloat16, - ) - # Hidden-state per-token cost = num_hidden_states(6) * hidden_size(512) - # * dtype(2) = 6144 bytes, which exceeds the 4096-byte KV page. - hs_spec = HiddenStateCacheSpec( - block_size=16, - num_kv_heads=6, - head_size=512, - dtype=torch.bfloat16, - ) - specs = { - "target.0.attn": kv_spec, - "target.1.attn": kv_spec, - "cache_only_layers.48": hs_spec, - } - - groups = get_kv_cache_groups(_grouping_config(), specs) - - # The hidden-state layer should be present and no assertion should fire. - all_layers = {name for g in groups for name in g.layer_names} - assert "cache_only_layers.48" in all_layers - - # The target group block sizes must have been scaled up. - for g in groups: - if "cache_only_layers.48" not in g.layer_names: - assert g.kv_cache_spec.block_size > kv_spec.block_size - - def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses(): assert get_kv_cache_spec_kind(new_mla_spec()) == KVCacheSpecKind.MLA_ATTENTION @@ -3525,179 +1988,6 @@ def test_request_with_prompt_embeds_and_mm_inputs(hash_fn: Callable[[Any], bytes assert block_hashes[1] == expected_hash2 -def test_auto_fit_max_model_len(): - """Test that max_model_len=-1 auto-fits to available GPU memory.""" - # Create config with original_max_model_len=-1 to trigger auto-fit - model_config = ModelConfig(max_model_len=1024) - # Simulate the user passing -1 by setting original_max_model_len - model_config.original_max_model_len = -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer - kv_cache_specs = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - } - - # With enough memory, max_model_len stays at the derived max - large_available_memory = mem_per_block_per_layer * 2 * 1024 # plenty of memory - _kv_cache_configs = get_kv_cache_configs( - vllm_config, [kv_cache_specs], [large_available_memory] - ) - assert vllm_config.model_config.max_model_len == 1024 - - # Reset for next test - model_config = ModelConfig(max_model_len=1024) - model_config.original_max_model_len = -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - - # With limited memory, max_model_len should be reduced - # Need memory for at least max_model_len tokens - # 32 blocks worth of memory for 2 layers = can fit 32*16=512 tokens - limited_memory = mem_per_block_per_layer * 2 * 32 - _kv_cache_configs = get_kv_cache_configs( - vllm_config, [kv_cache_specs], [limited_memory] - ) - # Should be reduced to fit in memory - assert vllm_config.model_config.max_model_len < 1024 - assert vllm_config.model_config.max_model_len > 0 - - -def test_auto_fit_max_model_len_with_hybrid(): - """Test that auto-fit works with hybrid KV cache specs.""" - # Create config with original_max_model_len=-1 to trigger auto-fit - model_config = ModelConfig(max_model_len=8192) - # Simulate the user passing -1 by setting original_max_model_len - model_config.original_max_model_len = -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer - gamma = 2 - kv_cache_specs = { - "layer_1": new_mamba_spec(num_speculative_blocks=gamma), - "layer_2": new_kv_cache_spec(), - } - - # One extra block on top of what a 1024-token request needs: the pool - # reserves one block as the null block. - available_memory = mem_per_block_per_layer * (1024 // 16 + 1 + gamma + 1) - _kv_cache_configs = get_kv_cache_configs( - vllm_config, [kv_cache_specs], [available_memory] - ) - assert vllm_config.model_config.max_model_len == 1024 - - -def test_auto_fit_max_model_len_not_triggered(): - """Test that auto-fit is not triggered when original_max_model_len is not -1.""" - model_config = ModelConfig(max_model_len=16) - # original_max_model_len should be None by default, not -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 - kv_cache_specs = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - } - - # This should work normally without auto-fit - _kv_cache_configs = get_kv_cache_configs( - vllm_config, [kv_cache_specs], [mem_per_block_per_layer * 2 * 32] - ) - assert vllm_config.model_config.max_model_len == 16 - - -def test_auto_fit_max_model_len_respects_num_gpu_blocks_override(): - """Auto-fit must size max_model_len against the override-clamped pool, not - the raw `available_memory`. Without this, auto-fit could pick a - max_model_len that no longer fits once `num_gpu_blocks_override` is applied. - """ - model_config = ModelConfig(max_model_len=16384) - model_config.original_max_model_len = -1 # request auto-fit - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - # Cap the cache to 32 blocks regardless of available memory. - vllm_config.cache_config.num_gpu_blocks_override = 32 - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 - kv_cache_specs = { - "layer_1": new_kv_cache_spec(), # block_size=16 - "layer_2": new_kv_cache_spec(), - } - # Plenty of raw memory (1024 blocks per layer would fit max_model_len=16384). - large_available_memory = mem_per_block_per_layer * 2 * 1024 - - get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) - - # 32 blocks * block_size 16 = 512 token slots, so max_model_len must - # auto-fit at or below that. - assert 0 < vllm_config.model_config.max_model_len <= 32 * 16 - - -def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override(): - """Admission check must use the override-clamped pool size, not raw - `available_memory`. Without this, startup could accept a max_model_len - that does not actually fit in `num_gpu_blocks_override` blocks. - """ - model_config = ModelConfig(max_model_len=16384) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - # 32 blocks is far too small for max_model_len=16384 (would need 1024). - vllm_config.cache_config.num_gpu_blocks_override = 32 - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 - kv_cache_specs = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - } - # Plenty of raw memory: a bytes-only check against this would pass. - large_available_memory = mem_per_block_per_layer * 2 * 1024 - - with pytest.raises(ValueError, match="max seq len"): - get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) - - -def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes(): - """DFlash drafters can have a smaller head size than the target model. - - For example, MiMo uses 192-dim target KV heads while its DFlash draft uses - 128-dim KV heads. The resulting page sizes are 3:2 rather than an integer - block-size multiple, so the smaller page must be padded instead. - """ - # Attention layers read padded pages through a strided view, so padding is - # allowed. - target_spec = new_kv_cache_spec( - block_size=16, - num_kv_heads=1, - head_size=192, - dtype=torch.bfloat16, - ) - draft_spec = new_sliding_window_spec( - block_size=16, - num_kv_heads=1, - head_size=128, - dtype=torch.bfloat16, - sliding_window=1024, - ) - - unified_specs = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "target_attn": target_spec, - "draft_attn": draft_spec, - } - ) - - assert unified_specs["target_attn"] == target_spec - unified_draft_spec = unified_specs["draft_attn"] - assert unified_draft_spec.block_size == draft_spec.block_size - assert unified_draft_spec.real_page_size_bytes == draft_spec.real_page_size_bytes - assert unified_draft_spec.page_size_padded == target_spec.page_size_bytes - assert unified_draft_spec.page_size_bytes == target_spec.page_size_bytes - - def test_unpadded_page_size_includes_per_token_head_scales(): # Per-token-head quant carries inline fp32 scales that are carved from the # raw KV allocation, so they must be budgeted into the offload width. The @@ -3729,173 +2019,6 @@ def test_page_size_padded_wins(): assert spec.page_size_bytes == 65536 -def test_unify_hybrid_kv_cache_specs(): - # 1. has_full_attention and has_sliding_window - before_spec_1 = new_kv_cache_spec(block_size=64) - before_spec_2 = new_sliding_window_spec( - block_size=16, page_size_padded=32 * 1024, sliding_window=1024 - ) - kv_cache_spec = { - "layer_1": before_spec_1, - "layer_2": before_spec_2, - } - kv_cache_spec["draft_layer_1"] = replace(before_spec_1, dcp_sharded=False) - kv_cache_spec["draft_layer_2"] = replace(before_spec_2, dcp_sharded=False) - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - expected_spec_1 = new_kv_cache_spec(block_size=64) - expected_spec_2 = new_kv_cache_spec( - block_size=64, page_size_padded=64 * 1024, sliding_window=1024 - ) - assert kv_cache_spec["layer_1"] == expected_spec_1 - assert kv_cache_spec["layer_2"] == expected_spec_2 - assert kv_cache_spec["layer_2"].page_size_bytes == 64 * 1024 - assert kv_cache_spec["draft_layer_2"] == replace(expected_spec_2, dcp_sharded=False) - - # 2. has_full_attention and has_chunked_local_attention - before_spec_1 = new_kv_cache_spec() - before_spec_2 = new_chunked_local_attention_spec( - page_size_padded=32 * 1024, attention_chunk_size=512 - ) - kv_cache_spec = { - "layer_1": before_spec_1, - "layer_2": before_spec_2, - } - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - expected_spec_1 = new_kv_cache_spec() - expected_spec_2 = new_kv_cache_spec( - page_size_padded=32 * 1024, attention_chunk_size=512 - ) - - assert kv_cache_spec["layer_1"] == expected_spec_1 - assert kv_cache_spec["layer_2"] == expected_spec_2 - - # 3. has_full_attention, has_sliding_window and has_chunked_local_attention - before_spec_1 = new_kv_cache_spec() - before_spec_2 = new_sliding_window_spec( - page_size_padded=32 * 1024, sliding_window=1024 - ) - before_spec_3 = new_chunked_local_attention_spec( - page_size_padded=32 * 1024, attention_chunk_size=512 - ) - kv_cache_spec = { - "layer_1": before_spec_1, - "layer_2": before_spec_2, - "layer_3": before_spec_3, - } - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - expected_spec_1 = new_kv_cache_spec() - expected_spec_2 = new_kv_cache_spec(page_size_padded=32 * 1024, sliding_window=1024) - expected_spec_3 = new_kv_cache_spec( - page_size_padded=32 * 1024, attention_chunk_size=512 - ) - assert kv_cache_spec["layer_1"] == expected_spec_1 - assert kv_cache_spec["layer_2"] == expected_spec_2 - assert kv_cache_spec["layer_3"] == expected_spec_3 - - # 4. No FullAttentionSpec, should not convert - kv_cache_spec = { - "layer_1": new_sliding_window_spec(sliding_window=1024), - "layer_2": new_chunked_local_attention_spec(attention_chunk_size=512), - } - - with pytest.raises(ValueError): - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - - # Replicated Mamba state still requires the hybrid cache manager. - kv_cache_spec = {"attention": new_kv_cache_spec(), "mamba": new_mamba_spec()} - with pytest.raises(ValueError): - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - - -def test_unify_kv_cache_spec_page_size_mamba(): - """Regression test for https://github.com/vllm-project/vllm/issues/43626. - - MambaSpec's page_size_bytes is determined by its state shapes and does not - change with block_size, so unify_kv_cache_spec_page_size must pad the Mamba - page instead of scaling its block_size. This situation arises when a layer - with a page larger than the (already platform-aligned) Mamba page joins the - specs, e.g. a dense draft model with more KV heads than the hybrid main - model. - """ - # 1. Hybrid main model (Mamba + full attention, pages already aligned at - # 16KB) plus a dense draft model layer with a 2x larger page (32KB). - # Reproduces the bare AssertionError from #43626: 32768 % 16384 == 0, so - # the old code scaled the Mamba block_size, which left page_size_bytes - # unchanged at 16384. - mamba_spec = new_mamba_spec() # page_size_bytes = 16384 - main_attn_spec = new_kv_cache_spec() # page_size_bytes = 16384 - draft_attn_spec = new_kv_cache_spec(num_kv_heads=4) # page_size_bytes = 32768 - assert mamba_spec.page_size_bytes == main_attn_spec.page_size_bytes == 16384 - assert draft_attn_spec.page_size_bytes == 32768 - - unified = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "mamba_layer": mamba_spec, - "main_attn_layer": main_attn_spec, - "draft_attn_layer": draft_attn_spec, - } - ) - # Mamba page is padded; block_size (caching granularity) is unchanged. - assert unified["mamba_layer"].page_size_bytes == 32768 - assert unified["mamba_layer"].page_size_padded == 32768 - assert unified["mamba_layer"].block_size == mamba_spec.block_size - # Attention layer with smaller page still unifies by scaling block_size. - assert unified["main_attn_layer"].page_size_bytes == 32768 - assert unified["main_attn_layer"].block_size == 2 * main_attn_spec.block_size - assert unified["main_attn_layer"].page_size_padded is None - # Layer already at max page size is unchanged. - assert unified["draft_attn_layer"] == draft_attn_spec - - # 2. Mamba page already padded by the platform (state smaller than the - # padded page); the padding is re-applied at the new maximum. - padded_mamba_spec = new_mamba_spec( - shapes=((2, 256), (3, 32, 32)), page_size_padded=16384 - ) - assert padded_mamba_spec.page_size_bytes == 16384 - unified = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "mamba_layer": padded_mamba_spec, - "draft_attn_layer": draft_attn_spec, - } - ) - assert unified["mamba_layer"].page_size_bytes == 32768 - assert unified["mamba_layer"].page_size_padded == 32768 - - # 3. Mamba page that does not evenly divide the maximum page size is - # padded as well (the divisibility constraint only applies to block_size - # scaling). - odd_mamba_spec = new_mamba_spec(shapes=((6144,),)) - assert odd_mamba_spec.page_size_bytes == 24576 - assert 32768 % odd_mamba_spec.page_size_bytes != 0 - unified = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "mamba_layer": odd_mamba_spec, - "draft_attn_layer": draft_attn_spec, - } - ) - assert unified["mamba_layer"].page_size_bytes == 32768 - - # 4. Attention layers with non-divisible page sizes are padded too: every - # backend reads a padded page through the view's block stride, so there is - # no longer a case that must raise. - unified = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "attn_layer": new_kv_cache_spec(block_size=24), # 24576 - "draft_attn_layer": draft_attn_spec, # 32768 - } - ) - assert unified["attn_layer"].page_size_padded == 32768 - assert unified["attn_layer"].page_size_bytes == 32768 - assert unified["attn_layer"].block_size == 24 - - # 5. Uniform page sizes are returned unchanged. - specs = { - "mamba_layer": new_mamba_spec(), - "attn_layer": new_kv_cache_spec(), - } - assert kv_cache_utils.unify_kv_cache_spec_page_size(specs) == specs - - def test_hma_not_disabled_when_kv_events_enabled(): """Test enabling KV events must not force disable_hybrid_kv_cache_manager to True. @@ -3982,78 +2105,6 @@ def test_resolve_block_hashes_rejects_mismatched_view(): resolve_block_hashes(mismatched, 2, 4) -@pytest.mark.parametrize("use_override", [True, False]) -def test_kv_cache_reserves_null_block_for_max_model_len(use_override): - """A KV cache sized to exactly the blocks a max_model_len request needs must - be rejected. BlockPool keeps one block as the null block, so only - num_blocks - 1 are usable; accepting num_blocks == needed would leave the - request unschedulable and hang the engine. Covers both the - num_gpu_blocks_override path and the memory-derived block count. - """ - block_size = 16 - max_model_len = 512 # needs 512 / 16 = 32 blocks - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - spec = new_kv_cache_spec(block_size=block_size) - - # 32 blocks -> only 31 usable after the null block: one short -> reject. - if use_override: - # Ample raw memory; the override alone constrains the block count. - vllm_config.cache_config.num_gpu_blocks_override = 32 - available_memory = [spec.page_size_bytes * 1024] - else: - available_memory = [spec.page_size_bytes * 32] - - with pytest.raises(ValueError, match="max seq len"): - get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) - - -def test_auto_fit_max_model_len_reserves_null_block(): - """Auto-fit (max_model_len=-1) must size max_model_len against usable - blocks, not the total pool. With memory for exactly 64 blocks, one is the - null block, so auto-fit must settle on 63 * block_size; picking the full - pool would leave a max-length request unschedulable and hang the engine. - """ - block_size = 16 - model_config = ModelConfig(max_model_len=1024) - model_config.original_max_model_len = -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - spec = new_kv_cache_spec(block_size=block_size) - - # Exactly the 1024 / 16 = 64 blocks a full-length request would need. - available_memory = [spec.page_size_bytes * 64] - - get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) - - assert vllm_config.model_config.max_model_len == 63 * block_size - - -def test_check_enough_kv_cache_memory_reserves_null_block(): - """The public admission check must reject a KV cache sized to exactly the - blocks a max_model_len request needs. BlockPool keeps one block as the - null block, so only num_blocks - 1 are usable; accepting - num_blocks == needed would leave the request unschedulable and hang the - engine. - """ - block_size = 16 - max_model_len = 512 # needs 512 / 16 = 32 blocks - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - spec = new_kv_cache_spec(block_size=block_size) - - # 32 blocks -> only 31 usable after the null block: one short -> reject. - with pytest.raises(ValueError, match="max seq len"): - check_enough_kv_cache_memory( - vllm_config, {"layer1": spec}, spec.page_size_bytes * 32 - ) - - # 33 blocks -> 32 usable after the null block -> accept. - check_enough_kv_cache_memory( - vllm_config, {"layer1": spec}, spec.page_size_bytes * 33 - ) - - def test_is_full_attention_spec_unwraps_uniform_type_specs(): """``UniformTypeKVCacheSpecs`` is not a ``FullAttentionSpec``, so callers scanning groups with a bare isinstance miss DeepSeek-V4-shaped configs.""" @@ -4105,340 +2156,3 @@ def test_wrapped_mamba_group_requires_block_zeroing(): ) assert config.needs_kv_cache_zeroing - - -def _spec_decode_grouping_config(method="dspark", model_type=None): - """Grouping config with an EAGLE-family speculative method enabled.""" - return SimpleNamespace( - scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), - cache_config=SimpleNamespace( - get_resolved_kv_cache_layout=lambda: SimpleNamespace( - is_block_outermost=True - ) - ), - model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type=model_type)), - speculative_config=SimpleNamespace( - method=method, - use_eagle=lambda: True, - use_eagle_block_drop=lambda: True, - ), - ) - - -def _hybrid_specs_with_draft(draft: bool, draft_shares_target_spec: bool = False): - """A K3-shaped hybrid: MLA full attention + Mamba, optionally plus a - DSpark-style draft MLA layer marked non_causal_multi_token_decode. - - The target's fp8 KV dtype is one difference from the draft on Kimi-K3 - (target `--kv-cache-dtype fp8_e4m3`, draft `auto`). Pass - draft_shares_target_spec to match that dtype while retaining the draft - marker as a group boundary. - """ - target_dtype = None if draft_shares_target_spec else "fp8_e4m3" - specs = { - "target.attn.0": new_mla_spec(block_size=64, cache_dtype_str=target_dtype), - "target.attn.1": new_mla_spec(block_size=64, cache_dtype_str=target_dtype), - "target.mamba.0": new_mamba_spec(block_size=64, mamba_cache_mode="align"), - "target.mamba.1": new_mamba_spec(block_size=64, mamba_cache_mode="align"), - } - if draft: - draft_spec = new_mla_spec(block_size=64) - specs["draft.attn.0"] = replace(draft_spec, non_causal_multi_token_decode=True) - return specs - - -def test_draft_group_annotated_on_hybrid_general_path(): - # A drafter's MLA layer carries non_causal_multi_token_decode, so its group - # is identifiable without keying off a model version. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) - ) - - flagged = [g for g in groups if g.is_eagle_group] - assert len(flagged) == 1 - assert "draft.attn.0" in flagged[0].layer_names - - -def test_mamba_groups_never_flagged_even_when_draft_shares_a_group(): - # Packed uniform-type groups can contain distinct target and draft layer - # specs; the combined group still holds volatile draft KV and must be - # flagged. What must never happen is a Mamba group being flagged: that - # widens its lookup window to two consecutive chunks, which align-mode - # checkpointing never produces, zeroing every lookup. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(), - _hybrid_specs_with_draft(draft=True, draft_shares_target_spec=True), - ) - - for group in groups: - if "draft.attn.0" in group.layer_names: - assert group.is_eagle_group - if isinstance(group.kv_cache_spec, MambaSpec): - assert not group.is_eagle_group - - -def test_draft_group_not_annotated_without_spec_decode(): - # The marker alone must not flag anything; the eagle semantics only apply - # when a speculative method is actually enabled. - config = _spec_decode_grouping_config() - config.speculative_config = None - groups = get_kv_cache_groups(config, _hybrid_specs_with_draft(draft=True)) - - assert not any(g.is_eagle_group for g in groups) - - -def test_unidentifiable_draft_with_mamba_warns(caplog_vllm): - # No group carries the draft marker, so consumers fall back to - # conservative behavior that silently breaks reuse for Mamba groups. - # That must at least be visible. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=False) - ) - - assert not any(g.is_eagle_group for g in groups) - assert "no KV cache group could be identified as the draft model's" in ( - caplog_vllm.text - ) - - -def test_unidentifiable_draft_without_mamba_does_not_warn(caplog_vllm): - # Pure-attention models degrade gracefully under the consumers' - # conservative fallback (a one-block hit drop at most), so the warning - # stays silent to avoid noise on every unannotated EAGLE deployment. - specs = { - "target.attn.0": new_mla_spec(block_size=64), - "target.attn.1": new_mla_spec(block_size=64), - } - groups = get_kv_cache_groups(_spec_decode_grouping_config(), specs) - - assert not any(g.is_eagle_group for g in groups) - assert "could be identified as the draft model's" not in caplog_vllm.text - - -def test_no_warning_when_draft_group_is_identified(caplog_vllm): - get_kv_cache_groups( - _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) - ) - - assert "could be identified as the draft model's" not in caplog_vllm.text - - -def _deepseek_v4_specs(model_version="deepseek_v4"): - """DeepseekV4-shaped specs: full MLA layers plus sliding-window MLA layers - at differing page sizes, with the MTP draft layer registered last.""" - return { - "model.layers.0.self_attn.attn": new_mla_spec(), - "model.layers.1.self_attn.attn": new_mla_spec(), - "model.layers.2.self_attn.attn": new_swa_mla_spec( - head_size=1024, model_version=model_version - ), - # The MTP block registers last, and its sliding-window size differs, so - # it lands in a group of its own. - "model.layers.3.self_attn.attn": new_swa_mla_spec( - head_size=1024, sliding_window=256, model_version=model_version - ), - } - - -@pytest.mark.parametrize( - ("method", "model_type"), - [("mtp", "deepseek_v4"), ("dspark", "deepseek_v4"), ("dspark", "deepseek_v41")], -) -def test_deepseek_v4_draft_group_annotated_on_packed_path(method, model_type): - # DeepseekV4's MTP block reuses the target's decoder layer, so its spec - # carries no draft marker and only the positional rule can find it. This - # pins the pre-existing behaviour that the unified annotator must preserve. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(method=method, model_type=model_type), - _deepseek_v4_specs(model_version=None), - ) - - flagged = [g for g in groups if g.is_eagle_group] - assert len(flagged) == 1 - assert "model.layers.3.self_attn.attn" in flagged[0].layer_names - - -def test_trailing_layer_fallback_applies_to_any_mtp_model(): - # The positional rule is sound for every MTP drafter, not just DeepseekV4: - # MTP blocks reuse the target's decoder layer (no spec marker) and always - # register after every target layer. The model_type must not gate it. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(method="mtp", model_type="other"), - _deepseek_v4_specs(), - ) - - flagged = [g for g in groups if g.is_eagle_group] - assert len(flagged) == 1 - assert "model.layers.3.self_attn.attn" in flagged[0].layer_names - - -def _qwen3_5_hybrid_specs(with_mtp_layer: bool): - """Qwen3.5-shaped hybrid: repeating [GDN x3, full-attn x1] blocks, with - the MTP drafter's full-attn layer (spec-identical to the target's) - registered last.""" - specs = {} - idx = 0 - for _ in range(2): - for _ in range(3): - specs[f"model.layers.{idx}.linear_attn"] = new_mamba_spec( - mamba_cache_mode="align" - ) - idx += 1 - specs[f"model.layers.{idx}.self_attn.attn"] = new_kv_cache_spec() - idx += 1 - if with_mtp_layer: - specs["mtp.layers.0.self_attn.attn"] = new_kv_cache_spec() - return specs - - -def test_qwen3_5_mtp_draft_group_annotated_on_hybrid_path(caplog_vllm): - # A hybrid mamba + full-attention model with an MTP drafter that is - # spec-indistinguishable from the target reaches the general multi-group - # path. The trailing-layer rule must locate the draft group there so the - # Mamba groups are not swept up by the flag-all consumer fallback. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(method="mtp", model_type="qwen3_5"), - _qwen3_5_hybrid_specs(with_mtp_layer=True), - ) - - flagged = [g for g in groups if g.is_eagle_group] - assert len(flagged) == 1 - assert "mtp.layers.0.self_attn.attn" in flagged[0].layer_names - for group in groups: - if any( - isinstance(spec, MambaSpec) - for spec in iter_layer_specs(group.kv_cache_spec) - ): - assert not group.is_eagle_group - assert "could be identified as the draft model's" not in caplog_vllm.text - - -@pytest.mark.parametrize("method", ["eagle", "eagle3", "dspark"]) -def test_non_mtp_eagle_hybrid_still_warns(caplog_vllm, method): - # Other drafters are not covered by the trailing-layer rule, so an - # unidentifiable hybrid draft must still warn. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(method=method, model_type="qwen3_5"), - _qwen3_5_hybrid_specs(with_mtp_layer=True), - ) - - assert not any(g.is_eagle_group for g in groups) - assert "could be identified as the draft model's" in caplog_vllm.text - - -def test_trailing_layer_fallback_requires_exact_partition(): - # If the groups do not partition the layers exactly (e.g. a caller that - # dropped or duplicated layers), the positional rule is meaningless and - # must not fire. - from vllm.v1.core.kv_cache_utils import _annotate_eagle_groups - - specs = _qwen3_5_hybrid_specs(with_mtp_layer=True) - config = _spec_decode_grouping_config(method="mtp", model_type="qwen3_5") - groups = get_kv_cache_groups(config, specs) - for g in groups: - g.is_eagle_group = False - # Remove one layer from its group: no longer an exact partition. - trimmed = [ - KVCacheGroupSpec( - [n for n in g.layer_names if n != "model.layers.0.linear_attn"], - g.kv_cache_spec, - ) - for g in groups - ] - _annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True) - - assert not any(g.is_eagle_group for g in trimmed) - - -_GQA_SPEC = FullAttentionSpec( - block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32 -) -_GQA_SWA_SPEC = SlidingWindowSpec( - block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32, sliding_window=128 -) - - -@pytest.mark.parametrize( - "specs,expected", - [ - pytest.param({"l.0": new_mla_spec(), "l.1": new_mla_spec()}, 4, id="mla"), - pytest.param( - { - "l.0": new_mla_spec(), - "l.1": new_mla_spec(), - "s.0": new_swa_mla_spec(), - "s.1": new_swa_mla_spec(), - }, - 4, - id="mla-and-swa-mla", - ), - pytest.param({"l.0": _GQA_SPEC, "l.1": _GQA_SPEC}, 1, id="gqa"), - pytest.param( - {"l.0": new_mla_spec(), "l.1": _GQA_SPEC}, - 1, - id="mla-and-gqa-uniform-group", - ), - pytest.param( - { - "l.0": new_mla_spec(), - "l.1": new_mla_spec(), - "s.0": _GQA_SWA_SPEC, - "s.1": _GQA_SWA_SPEC, - }, - 1, - id="mla-and-gqa-swa", - ), - ], -) -def test_kv_tp_replicas(monkeypatch, specs, expected): - """Resolved per layer before scheduler flattening hides mixed groups.""" - from vllm.config import ParallelConfig - from vllm.platforms import current_platform - - monkeypatch.setattr(current_platform, "device_count", lambda: 4) - parallel_config = ParallelConfig(tensor_parallel_size=4) - vllm_config = VllmConfig( - model_config=ModelConfig(max_model_len=16), parallel_config=parallel_config - ) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - mem = sum(s.page_size_bytes for s in specs.values()) * 10 - configs = get_kv_cache_configs(vllm_config, [specs], [mem]) - assert configs[0].kv_tp_replicas == expected - scheduler = generate_scheduler_kv_cache_config(configs) - assert scheduler.kv_tp_replicas == expected - - -@pytest.mark.parametrize( - "specs,tp_size,dcp_size,expected", - [ - pytest.param([new_mla_spec()], 8, 1, 8, id="mla"), - pytest.param([new_mla_spec()], 8, 2, 1, id="mla-dcp"), - pytest.param( - [replace(new_mla_spec(), max_tp_shards=None)], 8, 1, 1, id="mla-unset" - ), - pytest.param([replace(_GQA_SPEC, max_tp_shards=2)], 8, 1, 4, id="gqa-partial"), - pytest.param([replace(_GQA_SPEC, max_tp_shards=8)], 4, 1, 1, id="gqa-sharded"), - pytest.param( - [new_mla_spec(), replace(_GQA_SWA_SPEC, max_tp_shards=4)], - 8, - 1, - 2, - id="mla-and-gqa-partial", - ), - pytest.param( - [ - HiddenStateCacheSpec( - block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32 - ) - ], - 8, - 1, - 1, - id="hidden-state", - ), - ], -) -def test_kv_cache_groups_tp_replicas(specs, tp_size, dcp_size, expected): - """Replicas are the gcd of each layer's tp_size // max_tp_shards.""" - groups = [KVCacheGroupSpec([f"l.{i}"], spec) for i, spec in enumerate(specs)] - assert kv_cache_groups_tp_replicas(groups, tp_size, dcp_size) == expected diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index cf221fc22a34..541e55f5b60e 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -36,6 +36,10 @@ KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, + KVQuantMode, + MambaSpec, + MLAAttentionSpec, + SlidingWindowSpec, ) from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -333,3 +337,73 @@ def create_requests( ) requests.append(request) return requests + + +def new_kv_cache_spec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.float32, + page_size_padded=None, + sliding_window=None, + attention_chunk_size=None, + kv_quant_mode=KVQuantMode.NONE, +): + return FullAttentionSpec( + block_size=block_size, + num_kv_heads=num_kv_heads, + head_size=head_size, + dtype=dtype, + page_size_padded=page_size_padded, + sliding_window=sliding_window, + attention_chunk_size=attention_chunk_size, + kv_quant_mode=kv_quant_mode, + ) + + +def new_sliding_window_spec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.float32, + page_size_padded=None, + sliding_window=1, +): + return SlidingWindowSpec( + block_size=block_size, + num_kv_heads=num_kv_heads, + head_size=head_size, + dtype=dtype, + page_size_padded=page_size_padded, + sliding_window=sliding_window, + ) + + +def new_mamba_spec( + block_size=16, + shapes=((2, 512), (3, 32, 32)), + dtypes=(torch.float32, torch.float32), + num_speculative_blocks=2, + mamba_cache_mode="none", + page_size_padded=None, +): + return MambaSpec( + block_size=block_size, + shapes=shapes, + dtypes=dtypes, + page_size_padded=page_size_padded, + mamba_cache_mode=mamba_cache_mode, + num_speculative_blocks=num_speculative_blocks, + ) + + +def new_mla_spec(cache_dtype_str=None, block_size: int = 16): + # head_size = kv_lora_rank(512) + qk_rope_head_dim(64) = 576 + return MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + max_tp_shards=1, + head_size=576, + dtype=torch.float32, + cache_dtype_str=cache_dtype_str, + ) diff --git a/tests/v1/engine/test_init_error_messaging.py b/tests/v1/engine/test_init_error_messaging.py index 5626dec17043..f9455389ef08 100644 --- a/tests/v1/engine/test_init_error_messaging.py +++ b/tests/v1/engine/test_init_error_messaging.py @@ -4,7 +4,7 @@ import pytest import torch -from vllm.v1.core.kv_cache_utils import check_enough_kv_cache_memory +from vllm.v1.core.kv_cache_planning import check_enough_kv_cache_memory from vllm.v1.kv_cache_interface import FullAttentionSpec @@ -38,7 +38,7 @@ def test_kv_cache_oom_insufficient_memory(monkeypatch): config.parallel_config.decode_context_parallel_size = 1 monkeypatch.setattr( - "vllm.v1.core.kv_cache_utils.max_memory_usage_bytes", + "vllm.v1.core.kv_cache_planning.max_memory_usage_bytes", lambda c, s: 100 * 1024**3, # 100 GiB ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index 07d530f97d1a..a6692b5cd5fa 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -17,10 +17,8 @@ SchedulerOffloadConfig, ) from vllm.platforms import current_platform -from vllm.v1.core.kv_cache_utils import ( - generate_scheduler_kv_cache_config, - kv_cache_groups_tp_replicas, -) +from vllm.v1.core.kv_cache_planning import kv_cache_groups_tp_replicas +from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config from vllm.v1.kv_cache_interface import ( CircularBufferSpec, FullAttentionSpec, diff --git a/tests/v1/kv_connector/unit/test_hidden_states_connector.py b/tests/v1/kv_connector/unit/test_hidden_states_connector.py index 735b042259b5..854632dce1d1 100644 --- a/tests/v1/kv_connector/unit/test_hidden_states_connector.py +++ b/tests/v1/kv_connector/unit/test_hidden_states_connector.py @@ -10,7 +10,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.example_hidden_states_connector import ( # noqa: E501 ExampleHiddenStatesConnector, ) -from vllm.v1.core.kv_cache_utils import get_kv_cache_groups +from vllm.v1.core.kv_cache_planning import get_kv_cache_groups from vllm.v1.kv_cache_interface import ( FullAttentionSpec, HiddenStateCacheSpec, diff --git a/tests/v1/simple_kv_offload/test_worker.py b/tests/v1/simple_kv_offload/test_worker.py index 16e142d7cb97..4f9a7c451ffc 100644 --- a/tests/v1/simple_kv_offload/test_worker.py +++ b/tests/v1/simple_kv_offload/test_worker.py @@ -29,11 +29,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector import ( SimpleCPUOffloadConnector, ) -from vllm.v1.core.kv_cache_utils import ( +from vllm.v1.core.kv_cache_planning import ( get_kv_cache_config_from_groups, is_kv_cache_spec_uniform, - resolve_kv_cache_block_sizes, ) +from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, diff --git a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py index 4cc8f9ffe2fd..abd0830a472b 100644 --- a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py +++ b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py @@ -104,7 +104,7 @@ def test_overlaid_zeroer_dedups_segments_with_max_span(): zeroed no matter which group owns it.""" from unittest.mock import MagicMock - from vllm.v1.core.kv_cache_utils import get_kv_cache_config_from_groups + from vllm.v1.core.kv_cache_planning import get_kv_cache_config_from_groups from vllm.v1.kv_cache_interface import ( KVCacheGroupSpec, KVCacheLayout, diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 1247bf193a76..534d81312436 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -43,7 +43,10 @@ ROCMAiterMLASparseBackend, ) from vllm.v1.attention.backends.registry import AttentionBackendEnum -from vllm.v1.core.kv_cache_utils import estimate_max_model_len, get_kv_cache_configs +from vllm.v1.core.kv_cache_planning import ( + estimate_max_model_len, + get_kv_cache_configs, +) from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput from vllm.v1.kv_cache_interface import ( FullAttentionSpec, diff --git a/vllm/v1/core/kv_cache_planning.py b/vllm/v1/core/kv_cache_planning.py new file mode 100644 index 000000000000..6e258a68a9e7 --- /dev/null +++ b/vllm/v1/core/kv_cache_planning.py @@ -0,0 +1,1637 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""KV cache planning: grouping, sizing, and memory-fit logic.""" + +import math +from collections import defaultdict +from collections.abc import Callable, Iterable, Sequence +from dataclasses import replace +from functools import partial, reduce +from typing import Any + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv, round_up +from vllm.utils.mem_utils import format_gib +from vllm.utils.torch_utils import get_dtype_size +from vllm.v1.core.kv_cache_utils import _kv_cache_device_name +from vllm.v1.hisparse.layout import ( + get_hisparse_host_pool_bytes, + get_hisparse_kv_cache_config, + get_hisparse_kv_cache_groups, +) +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + ChunkedLocalAttentionSpec, + CircularBufferSpec, + FullAttentionSpec, + HiddenStateCacheSpec, + HiSparseHotSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheLayout, + KVCacheSpec, + KVCacheTensor, + MambaSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, + SlidingWindowSpec, + UniformTypeKVCacheSpecs, + compute_layout_strides, + iter_layer_specs, + replace_as, +) +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry + +logger = init_logger(__name__) + + +def _check_enough_kv_cache_memory( + available_memory: int, + get_needed_memory: Callable[[], int], + max_model_len: int, + estimate_max_model_len: Callable[[int], int], +): + if available_memory <= 0: + raise ValueError( + "No available memory for the cache blocks. " + "Try increasing `gpu_memory_utilization` when initializing the engine " + "(this flag also controls CPU memory reservation on the CPU " + "backend, despite its name). " + "See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ " + "for more details." + ) + + needed_memory = get_needed_memory() + + if needed_memory > available_memory: + estimated_max_len = estimate_max_model_len(available_memory) + estimated_msg = "" + if estimated_max_len > 0: + estimated_msg = ( + "Based on the available memory, " + f"the estimated maximum model length is {estimated_max_len}. " + ) + + raise ValueError( + f"To serve at least one request with the model's max seq len " + f"({max_model_len}), ({format_gib(needed_memory)} GiB KV " + f"cache is needed, which is larger than the available KV cache " + f"memory ({format_gib(available_memory)} GiB). {estimated_msg}" + f"Try increasing `gpu_memory_utilization` (which also controls " + f"CPU memory on the CPU backend) or decreasing `max_model_len` " + f"when initializing the engine. " + f"See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ " + f"for more details." + ) + + +def max_memory_usage_bytes( + vllm_config: VllmConfig, kv_cache_specs: Iterable[KVCacheSpec] +) -> int: + """Get the maximum memory usage in bytes for the given KV cache specs.""" + return sum(spec.max_memory_usage_bytes(vllm_config) for spec in kv_cache_specs) + + +def estimate_max_model_len( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], + available_memory: int, +) -> int: + """Estimates the maximum model length that can fit in the available memory + using binary search. + + This function temporarily modifies max_model_len during estimation but + restores the original value before returning, ensuring no side effects. + + Args: + vllm_config: The global VllmConfig + kv_cache_spec: The kv cache spec of each attention layer in the model + available_memory: Memory available for KV cache in bytes. + + Returns: + The estimated maximum model length that can fit in the available memory. + + """ + # Save the original max_model_len to restore after estimation + original_max_model_len = vllm_config.model_config.max_model_len + + # Define a function to check if a given model length fits in memory + def fits_in_memory(model_len: int) -> bool: + # Temporarily modify the max_model_len for this calculation + vllm_config.model_config.max_model_len = model_len + # Calculate memory needed for the given model length + memory_needed = max_memory_usage_bytes(vllm_config, kv_cache_spec.values()) + return memory_needed <= available_memory + + try: + # Binary search for the maximum model length + left, right = 1, original_max_model_len + + # If even the smallest model length doesn't fit, return 0 + if not fits_in_memory(left): + return 0 + + # Binary search for the maximum model length that fits + result = 1 + while left <= right: + mid = (left + right) // 2 + if fits_in_memory(mid): + result = mid + left = mid + 1 + else: + right = mid - 1 + return result + finally: + # Always restore the original max_model_len to avoid side effects + vllm_config.model_config.max_model_len = original_max_model_len + + +def check_enough_kv_cache_memory( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], + available_memory: int, +): + """Checks whether `available_memory` is enough for the KV cache to hold at + least one request with the model's max_model_len. + + Args: + vllm_config: The global VllmConfig + kv_cache_spec: The kv cache spec of each attention layer in the model + available_memory: Memory available for KV cache in bytes. + + Raises: + ValueError: If there is not enough memory available for the KV cache. + + """ + # No need to check for available memory if the kv_cache_spec is empty + if kv_cache_spec: + # Reserve the null block BlockPool permanently holds back, so the check + # plans against usable blocks, as in get_kv_cache_configs. Group a copy + # of the specs since grouping may unify them in-place. + groups = get_kv_cache_groups(vllm_config, dict(kv_cache_spec)) + check_memory = ( + available_memory - _pool_bytes_per_block(vllm_config, groups) + if groups + else available_memory + ) + _check_enough_kv_cache_memory( + check_memory, + lambda: max_memory_usage_bytes(vllm_config, kv_cache_spec.values()), + vllm_config.model_config.max_model_len, + lambda am: estimate_max_model_len(vllm_config, kv_cache_spec, am), + ) + + +def create_kv_cache_group_specs( + kv_cache_spec: dict[str, KVCacheSpec], grouped_layer_names: list[list[str]] +) -> list[KVCacheGroupSpec]: + """Create KVCacheGroupSpec object for each kv cache group layer. + The layers in the same group should share the same + KVCacheSpec. + + Args: + kv_cache_spec: + A mapping from each layer name to its corresponding KVCacheSpec. + grouped_layer_names: + A list of kv cache groups, where each element is a list of layer + names that belong to the same group and should share the same + KVCacheSpec. + + Returns: + A list of KVCacheGroupSpec objects, one for each group. + + """ + kv_cache_groups = [] + for layer_names_one_group in grouped_layer_names: + layer_specs = [ + kv_cache_spec[layer_name] for layer_name in layer_names_one_group + ] + merged_layer_spec = layer_specs[0].merge(layer_specs) + kv_cache_groups.append( + KVCacheGroupSpec(layer_names_one_group, merged_layer_spec) + ) + return kv_cache_groups + + +def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: + """Whether all layers in the given KVCacheSpec have the same KV cache spec. + Note that we regard FullAttentionSpec with and without sliding window as + the same type. + + Args: + kv_cache_spec: The kv cache spec of each attention layer in the model + + Returns: + True if all layers have the same type, False otherwise. + + """ + if not kv_cache_spec: + # Encoder-only models do not have KV cache, kv_cache_type can be + # regarded as uniform. + return True + try: + kv_cache_spec_values = list(kv_cache_spec.values()) + _ = kv_cache_spec_values[0].merge(kv_cache_spec_values) + except AssertionError: + return False + return True + + +def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int: + """Override the number of kv cache blocks if `num_gpu_blocks_override` is set. + The override is logged once, at the call site in `get_kv_cache_configs`. + """ + if vllm_config.cache_config.num_gpu_blocks_override is not None: + num_blocks = vllm_config.cache_config.num_gpu_blocks_override + return num_blocks + + +def _pool_bytes_per_block( + vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec] +) -> int: + """Bytes consumed by one block in the worker's shared KV cache pool, mirroring + the divisor used by `get_kv_cache_config_from_groups` to convert + `available_memory` into `num_blocks`. Used to compute the effective KV cache + capacity once `num_gpu_blocks_override` is applied. + """ + return _get_kv_cache_bytes_per_block( + [group for group in kv_cache_groups if not group.host_resident], + vllm_config.cache_config.get_resolved_kv_cache_layout(), + ) + + +def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int: + """Get the page size of the KV cache.""" + page_sizes = {layer.page_size_bytes for layer in kv_cache_specs} + assert len(page_sizes) == 1 + return page_sizes.pop() + + +def _get_kv_cache_groups_uniform_spec( + kv_cache_specs: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec]: + """Generates the KV cache configuration for a model with the same KV cache + spec for all layers. + + Args: + kv_cache_specs: The kv cache spec of each attention layer in the model + + Returns: + The generated KVCacheGroupSpecs + + """ + return create_kv_cache_group_specs(kv_cache_specs, [list(kv_cache_specs.keys())]) + + +def _get_kv_cache_groups_uniform_type( + spec: UniformTypeKVCacheSpecs, +) -> list[KVCacheGroupSpec]: + """Generates the KV cache configuration for a model with one type of KV cache + but different hidden sizes. All layers are merged into one group. + + Args: + spec: The UniformTypeKVCacheSpecs of the model + + Returns: + The generated KVCacheGroupSpecs + + """ + return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)] + + +def unify_kv_cache_spec_page_size( + kv_cache_spec: dict[str, KVCacheSpec], +) -> dict[str, KVCacheSpec]: + """Unify the page size of the given KVCacheSpec. If the page size of all layers + are the same, return the original KVCacheSpec. If not same, unify the page + size by increasing the block size of layers with smaller page size. Two + cases cannot be unified by block size alone and pad their physical page to + the maximum instead: Mamba layers, whose page size comes from state shapes + and is independent of block size; and non-MLA attention layers whose page + does not evenly divide the maximum (the padded page is read through a + strided view). MLA is excluded because sparse MLA indexes the cache in + whole token rows (see ``flat_kv_row_view``), so its block stride can only + be padded by its own row-aligned ``alignment``, not to an arbitrary page + size. Raise NotImplementedError if failed to unify the page size; + ``get_kv_cache_groups`` catches it to try the full-allocation fallback + (e.g. MLA next to an incompatible sliding-window draft). + + Args: + kv_cache_spec: The KVCacheSpec of each attention layer in the model + + Returns: + The updated KVCacheSpec with the same page_size_bytes. + + """ + page_sizes = {layer.page_size_bytes for layer in kv_cache_spec.values()} + if len(page_sizes) <= 1: + # All layers have the same page size, no need to unify. + return kv_cache_spec + + max_page_size = max(page_sizes) + new_kv_cache_spec = {} + for layer_name, layer_spec in kv_cache_spec.items(): + if layer_spec.page_size_bytes == max_page_size: + new_kv_cache_spec[layer_name] = layer_spec + elif isinstance(layer_spec, MambaSpec): + # MambaSpec's page size is determined by its state shapes and does + # not scale with block_size, so pad the page instead. This is the + # same padding mechanism the platform uses to align Mamba pages + # with the main model's attention page size; it is needed here + # when another layer (e.g. from a draft model) has a larger page + # than the already-aligned Mamba page. + new_spec: KVCacheSpec = replace(layer_spec, page_size_padded=max_page_size) + assert new_spec.page_size_bytes == max_page_size + new_kv_cache_spec[layer_name] = new_spec + else: + layer_page_size = layer_spec.page_size_bytes + if max_page_size % layer_page_size == 0: + ratio = max_page_size // layer_page_size + new_block_size = layer_spec.block_size * ratio + new_spec = replace(layer_spec, block_size=new_block_size) + elif isinstance(layer_spec, AttentionSpec) and not isinstance( + layer_spec, MLAAttentionSpec + ): + new_spec = replace(layer_spec, page_size_padded=max_page_size) + else: + raise NotImplementedError( + f"Layer {layer_name}: page size is not divisible by the " + "maximum page size and cannot be padded. Padding is only " + "supported for non-MLA attention layers." + ) + assert new_spec.page_size_bytes == max_page_size + new_kv_cache_spec[layer_name] = new_spec + return new_kv_cache_spec + + +def is_kv_cache_type_attention_free(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: + # kv_cache_spec is an empty dict for attention free models + return not kv_cache_spec + + +def _get_kv_cache_groups_uniform_page_size( + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec]: + """Generates the KV cache groups for hybrid models with multiple + attention types but still with a uniform page size (physical memory per + block per layer) for all layers. + + Detailed explanation about kv cache management of hybrid models: + The layers in the models are repeated with some patterns, e.g., a model + with 10 full attention layers and 20 sliding window attention layers can be + regarded as repeating the pattern (1 * full, 2 * sw) 10 times. + The KVCacheManager allocates different block tables for each of the 3 layers + in the pattern, and repeats each of them 10 times to generate the + block_table for the 30 layers in the model. + Therefore, we can group the layers in the model into 3 kv_cache_groups, each + of which contains 10 layers in the model. + The KVCacheManager allocates the block_table for each group based on its + kv_cache spec, and the model runner applies the block table to each layer + in the group. + For example: + 1. A model only uses full attention. The pattern is + (num_hidden_layers * full), so there is only one group and the block table + is shared by all layers. It is already handled by + `_get_kv_cache_config_uniform_type`. + 2. A model with 10 full attention layers and 20 sliding window + attention layers. There are 3 layers in the pattern (1 * full, 2 * sw), so + there are 3 kv_cache_groups, each of which represents 10 layers. + + To simplify the implementation, we make the following assumptions: + 1. Physical memory per block: Must be the same across all KV cache groups. + Breaking this assumption is non-trivial due to memory fragmentation concerns + when allocating blocks of different sizes. + 2. Tokens per block (block_size): Currently, we directly use + `CacheConfig.block_size` for all layers. It can be extended to vary by KV + cache group, but within each KV cache group, all layers must share the same + block size. + 3. Physical memory per token per layer: This property is decided by model + config. Currently we only support models that have the same physical memory + per token per layer for all layers. Can be relaxed with a simple extension, + but still need to keep physical memory per block the same for all groups. + 4. Number of layers per group: Currently assumed the same for all layers. + Can be relaxed with a simple extension, but still need to keep physical + memory per block the same for all groups. + 5. Attention type within groups: All layers in a group must share the same + attention type. One exception is that, when + `--disable-hybrid-kv-cache-manager` is true, the single group for full + attention layers may also include attention layers using sliding window or + LLaMA 4 local attention. See `unify_hybrid_kv_cache_specs` for more details. + 6. Support for multiple attention types: The design for most components is + general to an arbitrary number of attention types. But + `find_longest_cache_hit` only supports one attention type or two + types of full-attention plus exactly one another type. The general + implementation of this function is feasible but we don't know how to + implement it cleanly yet. + + As we assume tokens per block, physical memory per token per layer, and + number of layers per group are the same now, we can ensure that physical + memory per block is the same for all groups. + + Args: + kv_cache_spec: The KVCacheSpec of each attention layer in the model + Returns: + The generated KVCacheGroupSpecs + + """ + # Group all layers by kv_cache_spec. + # E.g., 2 full attention layers and 3 sliding window attention layers, + # -> (full.0, full.1), (sw.0, sw.1, sw.2). + same_type_layers: dict[KVCacheSpec, list[str]] = defaultdict(list) + for layer_name, layer_spec in kv_cache_spec.items(): + same_type_layers[layer_spec].append(layer_name) + + # Attempt to further merge same-type layers based on whether their KV + # cache specs can be merged, to minimize the group count. This benefits + # situations where specs share a block layout and differ only in a + # property it can reconcile (e.g. full attention layers differing only in + # sliding window / attention chunk size). + layer_buckets: list[list[str]] = [] + spec_buckets: list[list[KVCacheSpec]] = [] + for layer_spec, layer_names in same_type_layers.items(): + for names, specs in zip(layer_buckets, spec_buckets): + try: + # A raise means that the specs are incompatible. + type(specs[0]).merge([*specs, layer_spec]) + except (AssertionError, ValueError): + continue + names.extend(layer_names) + specs.append(layer_spec) + break + else: + layer_buckets.append(list(layer_names)) + spec_buckets.append([layer_spec]) + + # Split each group into smaller groups, to make the number of layers in each + # group identical. Add padding to the last group of each type if necessary. + # E.g., (full.0, full.1), (sw.0, sw.1, sw.2) + # split to 3 groups with 2 layers each: + # (full.0, full.1), (sw.0, sw.2), (sw.1, padding). + # FIXME(Chen): At the moment of writing this code (2025-06-02), all + # open-source hybrid model follows a n:1 pattern between different attention + # types (e.g., Gemma3 5:1 between sw and full, LLaMA4 3:1 between local and + # full), so we can use the "1" in the n:1 pattern as the group size, which + # is the minimum number of layers among all attention types. Need a better + # strategy if we want to support more complex patterns (e.g., 20 full + 30 + # sw, where the group size should be 10). + min_num_layers = min([len(layers) for layers in layer_buckets]) + group_size = min_num_layers + max_num_layers = max([len(layers) for layers in layer_buckets]) + if max_num_layers < min_num_layers * 1.5: + # If the number of layers is not much larger than the minimum number of + # layers, use the maximum number of layers as the group size to avoid + # too many padding layers. A typical example is gpt-oss-20b + eagle, + # with 12 sw + 13 full. We pad it to (13 sw, 13 full) instead of + # (12 sw, 24 full). 1.5 is a heuristic to avoid too many padding + # layers while accommodating speculative decoding drafters that add + # extra layers to one attention type. + group_size = max_num_layers + grouped_layers = [] + for layers in layer_buckets: + num_padding_layers = group_size - len(layers) % group_size + if num_padding_layers != group_size: + logger.warning( + "Add %d padding layers, may waste at most %.2f%% KV cache memory", # noqa + num_padding_layers, + num_padding_layers / len(layers) * 100, + ) + num_groups = cdiv(len(layers), group_size) + # In PP case, say if we have + # - stage 0: full.0, sw.0, sw.1 + # - stage 1: full.1, sw.2, sw.3 + # We should have 3 groups: (full.0, full.1), (sw.0, sw.2), (sw.1, sw.3) + # It can't be (full.0, full.1), (sw.0, sw.1), (sw.2, sw.3) because + # the 3 groups in stage 0 will be (full.0), (sw.0, sw.1), (empty group) + # and it will be padded to (full.0, padding), (sw.0, sw.1), + # (padding, padding) to ensure the number of layers in each group is + # the same and will cause memory waste. + # To avoid this, we assign layers[i::num_groups] to the i-th group + # instead of layers[i * group_size: (i + 1) * group_size] + for i in range(num_groups): + grouped_layers.append(layers[i::num_groups]) + return create_kv_cache_group_specs(kv_cache_spec, grouped_layers) + + +def _get_per_layer_spec( + group: KVCacheGroupSpec, + layer_name: str, +) -> KVCacheSpec: + spec = group.kv_cache_spec + if isinstance(spec, UniformTypeKVCacheSpecs): + return spec.kv_cache_specs[layer_name] + return spec + + +def _get_layer_block_stride(spec: KVCacheSpec) -> int: + """A layer-outer block stride: the page padded to its alignment.""" + return round_up(spec.page_size_bytes, spec.get_block_stride_alignment()) + + +def _get_kv_cache_bytes_per_block( + kv_cache_groups: list[KVCacheGroupSpec], layout: KVCacheLayout +) -> int: + """Return the largest cache group's bytes per block.""" + if not layout.is_block_outermost: + # Each layer's block stride is its own padded page. + return max( + sum( + _get_layer_block_stride(_get_per_layer_spec(group, layer_name)) + for layer_name in group.layer_names + ) + for group in kv_cache_groups + ) + bytes_per_block = max( + sum( + _get_per_layer_spec(group, layer_name).page_size_bytes + for layer_name in group.layer_names + ) + for group in kv_cache_groups + ) + assert bytes_per_block > 0 + alignments = [ + group.kv_cache_spec.page_size_bytes + for group in kv_cache_groups + if isinstance(group.kv_cache_spec, HiSparseHotSpec) + ] + alignments.extend( + _get_per_layer_spec(group, layer_name).get_block_stride_alignment() + for group in kv_cache_groups + for layer_name in group.layer_names + ) + return round_up(bytes_per_block, math.lcm(*alignments)) + + +def validate_kv_cache_layout( + layout: KVCacheLayout, + kv_cache_groups: list[KVCacheGroupSpec], +) -> None: + """Validate that the resolved layout can express this model's packing. + + The layout was chosen once in the engine core from the backends' supported + sets; a backend whose model packs pages side by side (e.g. the DeepSeek-V4 + indexer) declares block-outermost layouts there, so an inexpressible + layout reaching this point is an error. + """ + page_sizes = { + _get_per_layer_spec(group, layer_name).page_size_bytes + for group in kv_cache_groups + for layer_name in group.layer_names + } + if len(page_sizes) == 1: + # A rectangular layer dim exists; every layout can express it. + return + + # Mixed page sizes pack pages side by side within a block, which needs each page + # to be one contiguous chunk inside its block (a block-compact layout) and, with + # multiple KV cache groups, the layer dim inside the block dim. + if not layout.is_block_compact or ( + len(kv_cache_groups) > 1 and layout.is_layer_compact + ): + raise ValueError( + f"KV cache layout {layout.name} cannot express this model's " + f"mixed page sizes ({sorted(page_sizes)}); a backend should " + "declare block-outermost supported layouts (e.g. BLHNC), or " + "set VLLM_KV_CACHE_LAYOUT=BLHNC." + ) + + +def get_kv_cache_config_from_groups( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], + available_memory: int, +) -> KVCacheConfig: + """Generate the KV cache configuration from the KV cache groups and spec + of each layer. + + Args: + vllm_config: The global VllmConfig + kv_cache_groups: The KV cache groups + available_memory: Memory available for KV cache in bytes + Returns: + The generated KVCacheConfig + + """ + if len(kv_cache_groups) == 0: + # Attention free models do not have KV cache. + # Return num_blocks=1 as BlockPool always needs a null_block. + return KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=kv_cache_groups, + prefix_cache_retention_interval=( + vllm_config.cache_config.prefix_cache_retention_interval + ), + ) + + if vllm_config.attention_config.hisparse_config is not None: + host_budget = get_hisparse_host_pool_bytes(vllm_config) + return get_hisparse_kv_cache_config( + vllm_config, kv_cache_groups, available_memory, host_budget + ) + + layout = vllm_config.cache_config.get_resolved_kv_cache_layout() + validate_kv_cache_layout(layout, kv_cache_groups) + bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups, layout) + interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None + + num_blocks = available_memory // bytes_per_block + num_blocks = may_override_num_blocks(vllm_config, num_blocks) + size = bytes_per_block * num_blocks + + # Groups alias from byte 0. Spec regions are laid out differently: + # + # block-outer (the same packing repeats for every block): + # group 0: | blk 0 [ A | B | pad ] | blk 1 [ A | B | pad ] | ... + # group 1: | blk 0 [ C | D ] | blk 1 [ C | D ] | ... + # |<--- bytes_per_block -->| + # + # layer-outer (only supported for uniform page sizes or single-group models): + # group 0: | A [ blk 0 | blk 1 | ... ] | B [ blk 0 | blk 1 | ... ] | + # group 1: | C [ blk 0 | blk 1 | ... ] | D [ blk 0 | blk 1 | ... ] | + + kv_cache_tensors = [] + for group in kv_cache_groups: + group_spec = group.kv_cache_spec + layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list) + if isinstance(group_spec, UniformTypeKVCacheSpecs): + for layer_name, spec in group_spec.kv_cache_specs.items(): + layers_by_spec[spec].append(layer_name) + elif group.layer_names: + layers_by_spec[group_spec].extend(group.layer_names) + + byte_offset = 0 + for spec, layer_names in layers_by_spec.items(): + block_stride = interleaved_block_stride + padded_page = _get_layer_block_stride(spec) + if block_stride is None and padded_page != spec.page_size_bytes: + assert layout.is_block_compact, f"{layout.name} cannot pad pages" + block_stride = padded_page + layer_stride, block_stride, _, _, _ = compute_layout_strides( + spec, + num_blocks, + len(layer_names), + layout, + fixed_strides=(None, block_stride, None, None, None), + ) + kv_cache_tensors.append( + KVCacheTensor( + size=size, + layers=layer_names, + layer_stride=layer_stride, + block_stride=block_stride, + offset=byte_offset, + ) + ) + # Block-outer layers share each block; layer-outer ones own regions. + byte_offset += len(layer_names) * ( + spec.page_size_bytes if layout.is_block_outermost else layer_stride + ) + + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=kv_cache_tensors, + kv_cache_groups=kv_cache_groups, + prefix_cache_retention_interval=( + vllm_config.cache_config.prefix_cache_retention_interval + ), + ) + + +def _promote_local_kv_cache_specs( + kv_cache_spec: dict[str, KVCacheSpec], +) -> dict[str, KVCacheSpec]: + """Use full-attention allocation for local-attention cache specs. + + The returned specs affect KV cache management only. Attention modules keep + their original sliding-window or chunked-local compute behavior. + """ + promoted_specs = kv_cache_spec.copy() + + if is_kv_cache_spec_uniform( + promoted_specs + ) or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs): + return promoted_specs + + has_full_attention = any( + isinstance(spec, FullAttentionSpec) for spec in promoted_specs.values() + ) + has_sliding_window = any( + isinstance(spec, SlidingWindowSpec) for spec in promoted_specs.values() + ) + has_chunked_local_attention = any( + isinstance(spec, ChunkedLocalAttentionSpec) for spec in promoted_specs.values() + ) + full_block_sizes = { + spec.block_size + for spec in promoted_specs.values() + if isinstance(spec, FullAttentionSpec) + } + full_attention_block_size = ( + next(iter(full_block_sizes)) if len(full_block_sizes) == 1 else None + ) + + def promoted_page_size_padded(spec: AttentionSpec, block_size: int) -> int | None: + if spec.page_size_padded is None: + return None + unpadded_page_size = ( + spec.unpadded_page_size_bytes * block_size // spec.block_size + ) + return max(spec.page_size_padded, unpadded_page_size) + + promotions: dict[type[AttentionSpec], type[AttentionSpec]] = { + SlidingWindowMLASpec: MLAAttentionSpec, + SlidingWindowSpec: FullAttentionSpec, + ChunkedLocalAttentionSpec: FullAttentionSpec, + } + + if has_full_attention and (has_sliding_window or has_chunked_local_attention): + for layer_name, spec in kv_cache_spec.items(): + target_cls = next( + (promotions[c] for c in type(spec).__mro__ if c in promotions), None + ) + if target_cls is None: + continue + assert isinstance(spec, AttentionSpec) + block_size = full_attention_block_size or spec.block_size + promoted_specs[layer_name] = replace_as( + spec, + target_cls, + # Promoted specs allocate blocks for all tokens and never free + # below the window, so the trailing-edge extension is moot. + drop=("extra_retained_tokens",), + block_size=block_size, + page_size_padded=promoted_page_size_padded(spec, block_size), + ) + + if not ( + is_kv_cache_spec_uniform(promoted_specs) + or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs) + ): + raise ValueError("Failed to promote local KV cache specs to one unified type.") + + return promoted_specs + + +def _try_get_full_allocation_fallback_groups( + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec] | None: + """Try a supported full-allocation fallback for local-attention layers.""" + if any(isinstance(spec, HiddenStateCacheSpec) for spec in kv_cache_spec.values()): + return None + if any( + isinstance(spec, (SlidingWindowMLASpec, ChunkedLocalAttentionSpec)) + for spec in kv_cache_spec.values() + ): + return None + + has_mla = any(isinstance(spec, MLAAttentionSpec) for spec in kv_cache_spec.values()) + has_regular_swa = any( + isinstance(spec, SlidingWindowSpec) for spec in kv_cache_spec.values() + ) + if not (has_mla and has_regular_swa): + return None + + try: + promoted_specs = _promote_local_kv_cache_specs(kv_cache_spec) + except ValueError: + return None + uniform_spec = UniformTypeKVCacheSpecs.from_specs(promoted_specs) + if uniform_spec is None: + return None + logger.warning( + "KV cache page sizes cannot be unified; treating sliding-window " + "layers as full attention for cache allocation. Sliding-window " + "attention compute is unchanged." + ) + return _get_kv_cache_groups_uniform_type(uniform_spec) + + +def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]): + """This function tries to convert the KV cache specs to one type if the model + is a hybrid model with multiple type of KV cache. It will convert all + SlidingWindowSpec to FullAttentionSpec if both types are present. + + Args: + kv_cache_spec: The kv cache spec of each attention layer in the model + + """ + groups: defaultdict[bool, dict[str, KVCacheSpec]] = defaultdict(dict) + for name, spec in kv_cache_spec.items(): + replicated = isinstance(spec, AttentionSpec) and not spec.dcp_sharded + groups[replicated][name] = spec + for specs in groups.values(): + promoted_specs = _promote_local_kv_cache_specs(specs) + if promoted_specs == specs: + continue + logger.warning( + "Hybrid KV cache manager is disabled for this hybrid model, " + "This means we do not enable any optimizations for saving KV cache " + "memory (e.g., dropping the KV cache outside the sliding window). " + "The compute of layers like sliding window is still saved." + ) + kv_cache_spec.update(promoted_specs) + + +def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int: + """Pick a chunk size that minimizes total upward padding. + + Each x is rounded up to a multiple of d: + + x -> ceil(x / d) * d + + Total padding is: + + pad(d) = sum_i (ceil(x_i / d) * d - x_i) + + We brute-force d in [lower_bound, max(values)] (fine for small lists / small + maxima) and return the d with minimum padding. Ties prefer larger d. + """ + if not values: + raise ValueError("values must be non-empty") + if any(x <= 0 for x in values): + raise ValueError(f"values must be positive, got: {list(values)!r}") + + min_d = max(1, lower_bound if lower_bound is not None else 1) + max_d = max(values) + if min_d > max_d: + return min_d + + best_d = min_d + best_pad: int | None = None + for d in range(min_d, max_d + 1): + pad = sum((d - (x % d)) % d for x in values) + if best_pad is None or pad < best_pad or (pad == best_pad and d > best_d): + best_pad = pad + best_d = d + + return best_d + + +def _get_packed_kv_cache_groups( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec] | None: + """Group mixed-page-size layers for contiguous block-outermost packing. + + Greedily buckets layers into uniform-type specs. Buckets with equal layer + counts per page size are treated as a repeating layer pattern (one layer + per page size) and split into groups covering the same number of pattern + repeats (picked by ``_approximate_gcd`` to minimize padding), so all + groups pack into the same per-block layout. Mamba buckets are additionally + split to fit the block the attention buckets already need. + Returns None when the layout is not block-outermost or all layers already + share one page size. + """ + layout = vllm_config.cache_config.get_resolved_kv_cache_layout() + page_sizes = {spec.page_size_bytes for spec in kv_cache_spec.values()} + if not layout.is_block_outermost or len(page_sizes) <= 1: + return None + + buckets: list[dict[str, KVCacheSpec]] = [] + for name, spec in kv_cache_spec.items(): + for bucket in buckets: + candidate = {**bucket, name: spec} + if UniformTypeKVCacheSpecs.is_uniform_type(candidate): + bucket[name] = spec + break + else: + buckets.append({name: spec}) + + bucketed = [] + for bucket in buckets: + uniform_spec = UniformTypeKVCacheSpecs.from_specs(bucket) + assert uniform_spec is not None + page_size_layers: dict[int, list[str]] = defaultdict(list) + for layer_name, layer_spec in bucket.items(): + page_size_layers[layer_spec.page_size_bytes].append(layer_name) + # Only 1:1 patterns (one layer of each page size per repeat) are + # supported; counts sharing a gcd > 1 (e.g. 2:1) could in principle + # repeat too, but such buckets are emitted whole instead. + balanced = len(set(map(len, page_size_layers.values()))) == 1 + bucketed.append((uniform_spec, page_size_layers, balanced)) + + # Balanced buckets that mix page sizes must stay whole, so the largest one + # sets a floor on the repeats per group; larger single-size buckets are + # split down toward it. No such bucket means nothing needs packing. + min_repeats_per_group = max( + ( + spec.get_max_layers_per_page_size() + for spec, page_size_layers, balanced in bucketed + if balanced and len(page_size_layers) > 1 + ), + default=0, + ) + repeats_per_group = ( + _approximate_gcd( + [ + spec.get_max_layers_per_page_size() + for spec, _, balanced in bucketed + if balanced + ], + lower_bound=min_repeats_per_group, + ) + if min_repeats_per_group + else None + ) + + def num_groups_for(spec: UniformTypeKVCacheSpecs, balanced: bool) -> int: + if balanced and repeats_per_group is not None: + return cdiv(spec.get_max_layers_per_page_size(), repeats_per_group) + return 1 + + def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int: + """Page bytes of the largest of the n groups a bucket splits into.""" + return sum( + cdiv(len(names), n) * page for page, names in page_size_layers.items() + ) + + # Bytes a block must hold however the state buckets end up split: mamba, + # circular-buffer and unsplit sliding-window buckets can go down to one + # state per group, every other bucket's split is fixed by the repeat pattern. + def is_state_bucket(spec: UniformTypeKVCacheSpecs) -> bool: + if isinstance(spec.first_spec, (MambaSpec, CircularBufferSpec)): + return True + return repeats_per_group is None and isinstance( + spec.first_spec, SlidingWindowSpec + ) + + anchor_bytes = max( + ( + widest_group_bytes( + page_size_layers, + len(spec.kv_cache_specs) + if is_state_bucket(spec) + else num_groups_for(spec, balanced), + ) + for spec, page_size_layers, balanced in bucketed + ), + default=0, + ) + + groups = [] + for spec, page_size_layers, balanced in bucketed: + num_groups = num_groups_for(spec, balanced) + # Cap a state group at the states a block already fits rather than let + # it widen the block. + if anchor_bytes and is_state_bucket(spec): + states_per_block = max(anchor_bytes // spec.first_spec.page_size_bytes, 1) + num_groups = max( + num_groups, cdiv(len(spec.kv_cache_specs), states_per_block) + ) + if num_groups == 1: + groups.append(KVCacheGroupSpec(list(spec.kv_cache_specs), spec)) + continue + + pattern_repeats = list(zip(*page_size_layers.values())) + for i in range(num_groups): + group_layer_names = [ + name for repeat in pattern_repeats[i::num_groups] for name in repeat + ] + group_layer_specs = { + name: spec.kv_cache_specs[name] for name in group_layer_names + } + group_spec = UniformTypeKVCacheSpecs.from_specs(group_layer_specs) + assert group_spec is not None + groups.append(KVCacheGroupSpec(group_layer_names, group_spec)) + + _annotate_eagle_groups( + vllm_config, + kv_cache_spec, + groups, + use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), + ) + _warn_if_unannotated_eagle_mamba(vllm_config, groups) + return groups + + +def _uses_trailing_mtp_layers(vllm_config: VllmConfig) -> bool: + """Whether the drafter's KV layers can be located positionally. + + MTP drafters register their KV layers after the target layers, without a + distinguishing spec marker. DeepseekV4/V4.1 do this with ``dspark`` too. + The annotator separately checks that the groups partition the layers exactly. + """ + spec_config = vllm_config.speculative_config + if spec_config is None or not spec_config.use_eagle(): + return False + if spec_config.method == "mtp": + return True + model_config = vllm_config.model_config + return model_config is not None and model_config.hf_config.model_type in ( + "deepseek_v4", + "deepseek_v41", + ) + + +def _groups_partition_layers_exactly( + kv_cache_spec: dict[str, KVCacheSpec], + kv_cache_groups: list[KVCacheGroupSpec], +) -> bool: + """Whether the groups cover every layer of ``kv_cache_spec`` exactly once. + + The trailing-layer draft fallback is only meaningful when the last + registered layer is guaranteed to appear in exactly one group. + """ + names = [name for group in kv_cache_groups for name in group.layer_names] + return len(names) == len(kv_cache_spec) and set(names) == set(kv_cache_spec) + + +def _annotate_eagle_groups( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], + kv_cache_groups: list[KVCacheGroupSpec], + use_trailing_layer_fallback: bool = False, +) -> None: + """Flag the KV cache groups that hold drafter attention layers. + + Two detection rules, in order of preference: + + 1. Spec-driven. ``non_causal_multi_token_decode`` is declared on + MLAAttentionSpec and set by drafter attention layers that run a + non-causal multi-token decode (today only Kimi-K3 DSpark). It survives + MLAAttentionSpec.merge, so it still identifies a group after per-group + spec merging, wherever grouping happens to land. It is sufficient but + not necessary: a drafter whose spec is indistinguishable from the + target's cannot be found this way. + 2. Positional fallback for MTP drafters (including DeepseekV4/V4.1 DSpark), whose + MTP block reuses the target's own decoder layer and so carries no spec + marker. Their draft attention layers always register after every + target layer, so flag whichever group holds the last registered layer. + This rule is only valid where the groups partition exactly the layers + of ``kv_cache_spec``, which is re-checked here before applying it. + When a drafter's trailing caches span several groups, this rule flags + only the group holding the very last layer and must be generalized. + FIXME(yifan): avoid/generalize this hacky check. + + Args: + vllm_config: Config supplying the speculative method, if any. + kv_cache_spec: The kv cache spec of each attention layer, in layer + registration order. Only read by rule 2. + kv_cache_groups: Groups to annotate in place. + use_trailing_layer_fallback: Enable rule 2. Callers gate this on + ``_uses_trailing_mtp_layers``. + + """ + spec_config = vllm_config.speculative_config + if spec_config is None or not spec_config.use_eagle_block_drop(): + return + + for group in kv_cache_groups: + if any( + getattr(spec, "non_causal_multi_token_decode", False) + for spec in iter_layer_specs(group.kv_cache_spec) + ): + group.is_eagle_group = True + + if not use_trailing_layer_fallback: + return + if not _groups_partition_layers_exactly(kv_cache_spec, kv_cache_groups): + return + last_layer = next(reversed(kv_cache_spec)) + for group in kv_cache_groups: + if last_layer in group.layer_names: + group.is_eagle_group = True + break + + +def _warn_if_unannotated_eagle_mamba( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], +) -> None: + """Warn when no KV cache group could be identified as the draft model's. + + Args: + vllm_config: Config supplying the speculative method, if any. + kv_cache_groups: Groups as they will be handed to consumers. + + """ + spec_config = vllm_config.speculative_config + if spec_config is None or not spec_config.use_eagle(): + return + if any(group.is_eagle_group for group in kv_cache_groups): + return + mamba_groups = [ + idx + for idx, group in enumerate(kv_cache_groups) + if any( + isinstance(spec, MambaSpec) + for spec in iter_layer_specs(group.kv_cache_spec) + ) + ] + if not mamba_groups: + return + logger.warning( + "Speculative decoding (method=%s) is enabled but no KV cache group " + "could be identified as the draft model's.", + spec_config.method, + ) + + +def _largest_divisor_at_most(value: int, limit: int) -> int: + for candidate in range(min(value, limit), 0, -1): + if value % candidate == 0: + return candidate + return 1 + + +def _ensure_min_page_size( + groups: list[KVCacheGroupSpec], + common_page: int, + hidden_specs: dict[str, HiddenStateCacheSpec], +) -> tuple[list[KVCacheGroupSpec], int]: + """Scale up group block sizes so the common page is at least as large as + the biggest hidden-state per-token cost. + + This protects the hidden state extraction feature, where we store hidden states in + the KV cache, and return them to clients using a custom hidden state connector. + + Returns the (possibly scaled) groups and updated common page size. + """ + max_per_token = max( + spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) + for spec in hidden_specs.values() + ) + if max_per_token <= common_page: + return groups, common_page + + scale = cdiv(max_per_token, common_page) + common_page *= scale + scaled: list[KVCacheGroupSpec] = [] + for g in groups: + s = g.kv_cache_spec + kw: dict[str, Any] = {"block_size": s.block_size * scale} + if isinstance(s, (AttentionSpec, MambaSpec)) and s.page_size_padded is not None: + kw["page_size_padded"] = s.page_size_padded * scale + scaled.append(KVCacheGroupSpec(g.layer_names, replace(s, **kw))) + return scaled, common_page + + +def get_kv_cache_groups( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec]: + """Split the layers in the model into groups with the same KV cache spec. + + Args: + vllm_config: The global VllmConfig + kv_cache_spec: The kv cache spec of each attention layer in the model + + Returns: + The generated KVCacheGroups + + """ + if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager: + unify_hybrid_kv_cache_specs(kv_cache_spec) + + if is_kv_cache_type_attention_free(kv_cache_spec): + # This returns an empty list to allow for the KVCacheManager to handle + # attention free models. + return [] + + if hisparse_groups := get_hisparse_kv_cache_groups(vllm_config, kv_cache_spec): + return hisparse_groups + + if is_kv_cache_spec_uniform(kv_cache_spec): + # KV cache of all layers are the same, which is true for + # most models. Allocate the same amount of memory for + # each layer. + return _get_kv_cache_groups_uniform_spec(kv_cache_spec) + elif uniform_spec := UniformTypeKVCacheSpecs.from_specs(kv_cache_spec): + # All layers need the same number of token slots (e.g., all layers are + # full attention, or all layers are sliding window attention with the + # same window size). Put all layers into one group. + return _get_kv_cache_groups_uniform_type(uniform_spec) + # Hidden-state layers use their own block table and must not be absorbed + # into a compatible attention bucket. + hidden_specs = { + k: v for k, v in kv_cache_spec.items() if isinstance(v, HiddenStateCacheSpec) + } + filtered_spec = { + k: v + for k, v in kv_cache_spec.items() + if not isinstance(v, HiddenStateCacheSpec) + } + + if packed_groups := _get_packed_kv_cache_groups(vllm_config, filtered_spec): + # Block-outermost blocks are strided by the widest group, so hidden + # groups need no page alignment. + packed_groups += [ + KVCacheGroupSpec([name], spec) for name, spec in hidden_specs.items() + ] + return packed_groups + + # Prefer preserving each layer's cache semantics. If physical pages cannot + # be unified, try a supported allocation-only fallback before failing. + try: + filtered_spec = unify_kv_cache_spec_page_size(filtered_spec) + except NotImplementedError: + fallback_groups = _try_get_full_allocation_fallback_groups(kv_cache_spec) + if fallback_groups is None: + raise + return fallback_groups + groups = _get_kv_cache_groups_uniform_page_size(filtered_spec) + + # Add hidden-state layers back with page aligned to the common page. + if hidden_specs: + common_page = get_uniform_page_size([g.kv_cache_spec for g in groups]) + # TP may shrink the common page below hidden-state per-token cost. + groups, common_page = _ensure_min_page_size(groups, common_page, hidden_specs) + group_block_size = math.gcd(*(g.kv_cache_spec.block_size for g in groups)) + for name, spec in hidden_specs.items(): + per_token = spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) + max_block_size = max(common_page // per_token, 1) + new_bs = _largest_divisor_at_most(group_block_size, max_block_size) + wasted_bytes = common_page - new_bs * per_token + logger.info( + "Using block size %d for hidden-state cache layer %s; " + "page alignment wastes %d bytes (%.2f%%) per block", + new_bs, + name, + wasted_bytes, + wasted_bytes / common_page * 100, + ) + aligned = replace(spec, block_size=new_bs, page_size_padded=common_page) + groups.append(KVCacheGroupSpec([name], aligned)) + + _annotate_eagle_groups( + vllm_config, + kv_cache_spec, + groups, + use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), + ) + _warn_if_unannotated_eagle_mamba(vllm_config, groups) + return groups + + +def _layer_tp_replicas(spec: KVCacheSpec, tp_size: int, dcp_size: int) -> int: + if not isinstance(spec, AttentionSpec) or spec.max_tp_shards is None: + return 1 + if spec.dcp_sharded and dcp_size > 1: + return 1 + return max(1, tp_size // spec.max_tp_shards) + + +def kv_cache_groups_tp_replicas( + groups: list[KVCacheGroupSpec], tp_size: int, dcp_size: int = 1 +) -> int: + """Consecutive TP ranks holding identical KV for every layer.""" + specs = [spec for g in groups for spec in iter_layer_specs(g.kv_cache_spec)] + if not specs: + return 1 + return reduce( + math.gcd, (_layer_tp_replicas(s, tp_size, dcp_size) for s in specs), tp_size + ) + + +def _max_memory_usage_bytes_from_groups( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], +) -> int: + """Calculate maximum memory usage in bytes from KV cache groups. + + This correctly accounts for padding in hybrid models. For example, if a + model has 8 full attention layers and 9 sliding window layers, they will + be padded to 9 full + 9 sliding window for uniform group sizes. + + Each group independently claims blocks from the shared pool, so a request consumes + the sum of the per-group block counts, i.e. ``bytes_per_block * total_blocks``. + """ + kv_cache_groups = [group for group in kv_cache_groups if not group.host_resident] + if not kv_cache_groups: + return 0 + + bytes_per_block = _pool_bytes_per_block(vllm_config, kv_cache_groups) + total_blocks = 0 + for group in kv_cache_groups: + spec = group.kv_cache_spec + if isinstance(spec, UniformTypeKVCacheSpecs): + total_blocks += spec.max_memory_usage_pages(vllm_config) + else: + total_blocks += cdiv( + spec.max_memory_usage_bytes(vllm_config), + spec.page_size_bytes, + ) + + return bytes_per_block * total_blocks + + +def _estimate_max_model_len_from_groups( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], + available_memory: int, +) -> int: + """Binary search for the maximum model length that fits in available memory. + Returns 0 if even 1 token doesn't fit. + """ + original_max = vllm_config.model_config.max_model_len + hisparse_enabled = ( + vllm_config.attention_config.hisparse_config is not None + and bool(kv_cache_groups) + ) + + def fits(model_len: int) -> bool: + vllm_config.model_config.max_model_len = model_len + if hisparse_enabled: + try: + get_kv_cache_config_from_groups( + vllm_config, kv_cache_groups, available_memory + ) + except ValueError: + return False + return ( + _max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups) + <= available_memory + ) + + try: + left, right = 1, original_max + if not fits(left): + return 0 + result = 1 + while left <= right: + mid = (left + right) // 2 + if fits(mid): + result = mid + left = mid + 1 + else: + right = mid - 1 + return result + finally: + vllm_config.model_config.max_model_len = original_max + + +def _auto_fit_max_model_len( + vllm_config: VllmConfig, + projected_groups_per_worker: list[list[KVCacheGroupSpec]], + available_memory: list[int], +) -> None: + """When max_model_len is set to -1, this function estimates the largest + context length that can be supported with the available GPU memory. + It uses binary search to find the maximum length that fits across all + workers. + + Args: + vllm_config: The global VllmConfig (will be modified in-place) + projected_groups_per_worker: KV cache groups projected to each worker. + available_memory: Memory available for KV cache in bytes for each + worker. + + """ + original_max = vllm_config.model_config.max_model_len + + if all(not groups for groups in projected_groups_per_worker): + # All workers have empty specs (attention-free model) + logger.info_once( + "Auto-fit max_model_len: attention-free model, " + "using derived max_model_len=%d", + original_max, + ) + return + + # Find the max_model_len that fits across all workers. + auto_fit_max = original_max + limiting_worker_mem = available_memory[0] + for groups, avail_mem in zip(projected_groups_per_worker, available_memory): + if not groups: + continue + worker_max = _estimate_max_model_len_from_groups(vllm_config, groups, avail_mem) + if worker_max < auto_fit_max: + auto_fit_max = worker_max + limiting_worker_mem = avail_mem + + if auto_fit_max <= 0: + raise ValueError( + "Cannot auto-fit max_model_len: not enough GPU memory available " + "to serve even a single token. Try increasing `gpu_memory_utilization`." + ) + + kv_cache_memory = ( + f"{format_gib(limiting_worker_mem)} GiB {_kv_cache_device_name(vllm_config)}" + ) + if vllm_config.attention_config.hisparse_config is not None: + host_pool_gib = format_gib(get_hisparse_host_pool_bytes(vllm_config)) + kv_cache_memory += f" + {host_pool_gib} GiB host" + + if auto_fit_max >= original_max: + # The model's full context length fits in memory + logger.info_once( + "Auto-fit max_model_len: full model context length %d fits in " + "available KV cache memory (%s)", + original_max, + kv_cache_memory, + ) + else: + # Need to reduce max_model_len to fit in memory + vllm_config.model_config.max_model_len = auto_fit_max + logger.info_once( + "Auto-fit max_model_len: reduced from %d to %d to fit in " + "available KV cache memory (%s)", + original_max, + auto_fit_max, + kv_cache_memory, + ) + + +def _project_kv_cache_groups_to_worker( + global_kv_cache_groups: list[KVCacheGroupSpec], + worker_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec]: + """Projects global KV cache groups onto a single worker's assigned layers. + + In pipeline parallelism, each worker only owns a subset of layers. This + function filters the global groups to include only layers present on the + given worker, adjusting UniformTypeKVCacheSpecs accordingly. + + Args: + global_kv_cache_groups: The global KV cache groups for the whole model. + worker_spec: The KV cache spec of each layer on this worker. + + Returns: + The projected KV cache groups containing only this worker's layers. + + """ + projected_groups: list[KVCacheGroupSpec] = [] + for group in global_kv_cache_groups: + worker_layer_names = [ + layer_name for layer_name in group.layer_names if layer_name in worker_spec + ] + if len(worker_layer_names) == len(group.layer_names): + projected_groups.append(group) + continue + group_spec = group.kv_cache_spec + if worker_layer_names and isinstance(group_spec, UniformTypeKVCacheSpecs): + group_spec = UniformTypeKVCacheSpecs( + block_size=group_spec.block_size, + kv_cache_specs={ + layer_name: group_spec.kv_cache_specs[layer_name] + for layer_name in worker_layer_names + }, + ) + projected_groups.append( + KVCacheGroupSpec( + worker_layer_names, + group_spec, + is_eagle_group=group.is_eagle_group and bool(worker_layer_names), + ) + ) + return projected_groups + + +def get_kv_cache_configs( + vllm_config: VllmConfig, + kv_cache_specs: list[dict[str, KVCacheSpec]], + available_memory: list[int], +) -> list[KVCacheConfig]: + """Generates the KV cache configurations for a model. + Since we use a shared centralized controller for all workers, we need the + `kv_cache_config` to be consistent across all workers to make sure + the KV cache allocation can be applied to all workers. However, different + workers may have different memory available, and different type of layers + (when pipeline parallel is enabled). To handle the difference between + workers, the current implementation is: + 1. Merge the KV cache specs of all workers to get the KVCacheSpecs for + the whole model. + 2. Generate the KV cache groups based on the layer ratio of the whole model. + This also handles spec unification for hybrid models. + 3. Handle auto-fit max_model_len and memory checks using per-worker + projected groups to account for PP sharding. + 4. Generate the KV cache configs for each worker based on the KV cache + grouping strategy. (This is reasonable because the layer ratio of + different PP stages are similar.) + 5. Change the num_blocks of each worker to the smallest among all workers + and shrink tensor sizes proportionally to avoid allocating unused memory. + + Args: + vllm_config: The global VllmConfig + kv_cache_specs: List of dict[layer_name, KVCacheSpec] for each worker. + available_memory: Memory available for KV cache in bytes for each + worker. + + Returns: + The generated KVCacheConfigs for each worker. + + """ + # Merge the KV cache specs of all workers. Different PP stages may have + # different layer names, and different TP ranks of the same PP stage should + # have the same KV cache spec. + merged_kv_cache_specs: dict[str, KVCacheSpec] = {} + for kv_cache_spec_one_worker in kv_cache_specs: + for layer_name, layer_spec in kv_cache_spec_one_worker.items(): + if layer_name not in merged_kv_cache_specs: + merged_kv_cache_specs[layer_name] = layer_spec + else: + assert merged_kv_cache_specs[layer_name] == layer_spec, ( + "The KV cache specs for the same layer are different " + "across workers. This is not supported yet." + ) + + # Check if the KV cache specs are registered correctly. + # This is to prevent that some layers are initialized with unregistered specs. + KVCacheSpecRegistry.check_kv_cache_spec_registry(merged_kv_cache_specs) + + # When speculating with more than 1 speculative module (e.g. multi-layered MTP) + # tag every SlidingWindowSpec with how many extra tokens to retain in the window. + extra_retained_tokens = max(0, vllm_config.num_prefill_lookahead_tokens - 1) + for layer_name, layer_spec in merged_kv_cache_specs.items(): + if isinstance(layer_spec, SlidingWindowSpec): + merged_kv_cache_specs[layer_name] = replace( + layer_spec, extra_retained_tokens=extra_retained_tokens + ) + + # Get global KV cache groups. This also handles spec unification for + # hybrid models when disable_hybrid_kv_cache_manager is enabled. + # After this call, merged_kv_cache_specs may be modified in-place. + global_kv_cache_groups = get_kv_cache_groups(vllm_config, merged_kv_cache_specs) + + # If original_max_model_len was -1, automatically + # determine the maximum model length that fits in available GPU memory. + # We use per-worker projected groups to account for PP sharding. + projected_groups_per_worker = [ + _project_kv_cache_groups_to_worker(global_kv_cache_groups, worker_spec) + for worker_spec in kv_cache_specs + ] + + # If `num_gpu_blocks_override` is set, the cache size that will actually + # be allocated is decoupled from the profiled `available_memory`: + # `may_override_num_blocks` in `get_kv_cache_config_from_groups` clamps + # `num_blocks` to the override. Reflect that in `available_memory` here so + # auto-fit, the admission check, and the per-worker config builder all + # plan against the same effective capacity. + override = vllm_config.cache_config.num_gpu_blocks_override + if override is not None: + adjusted_memory: list[int] = [] + for groups, avail_mem in zip(projected_groups_per_worker, available_memory): + if not groups: + adjusted_memory.append(avail_mem) + continue + bytes_per_block = _pool_bytes_per_block(vllm_config, groups) + logger.info( + "Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d", + avail_mem // bytes_per_block, + override, + ) + adjusted_memory.append(override * bytes_per_block) + available_memory = adjusted_memory + + # Reserve the null block BlockPool permanently holds back, so auto-fit and + # the capacity check both plan against usable blocks. Allocation below + # still uses the full memory. + check_memory = [ + avail_mem - _pool_bytes_per_block(vllm_config, groups) if groups else avail_mem + for groups, avail_mem in zip(projected_groups_per_worker, available_memory) + ] + + if vllm_config.model_config.original_max_model_len == -1: + _auto_fit_max_model_len(vllm_config, projected_groups_per_worker, check_memory) + + # Check if the available memory is enough per worker. + for groups, avail_mem in zip(projected_groups_per_worker, check_memory): + if not groups: + continue + _check_enough_kv_cache_memory( + avail_mem, + partial(_max_memory_usage_bytes_from_groups, vllm_config, groups), + vllm_config.model_config.max_model_len, + partial(_estimate_max_model_len_from_groups, vllm_config, groups), + ) + + kv_cache_configs: list[KVCacheConfig] = [] + for projected_groups, kv_cache_spec_one_worker, available_memory_one_worker in zip( + projected_groups_per_worker, kv_cache_specs, available_memory + ): + assert sum(len(group.layer_names) for group in projected_groups) == len( + kv_cache_spec_one_worker + ), "Some layers are not assigned to any group." + kv_cache_configs.append( + get_kv_cache_config_from_groups( + vllm_config, projected_groups, available_memory_one_worker + ) + ) + + # Change the num_blocks of each rank to the smallest among all ranks. + # We also need to shrink the tensor size proportionally to avoid + # allocating unused memory. + min_num_blocks = min( + kv_cache_config.num_blocks for kv_cache_config in kv_cache_configs + ) + for i, kv_cache_config in enumerate(kv_cache_configs): + if kv_cache_config.num_blocks == min_num_blocks: + continue + # Re-plan with exactly the memory the smallest rank can afford, so + # strides and offsets stay consistent with the shrunken allocation. + groups = kv_cache_config.kv_cache_groups + kv_cache_configs[i] = get_kv_cache_config_from_groups( + vllm_config, + groups, + min_num_blocks * _pool_bytes_per_block(vllm_config, groups), + ) + + for kv_cache_config in kv_cache_configs: + kv_cache_config.kv_tp_replicas = kv_cache_groups_tp_replicas( + kv_cache_config.kv_cache_groups, + vllm_config.parallel_config.tensor_parallel_size, + vllm_config.parallel_config.decode_context_parallel_size, + ) + + return kv_cache_configs diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 03c6c5ff62ac..5c51a8d131ff 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -6,10 +6,8 @@ import hashlib import math import os -from collections import defaultdict from collections.abc import Callable, Iterable, Iterator, Sequence from dataclasses import dataclass, field, replace -from functools import partial, reduce from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, overload from vllm import envs @@ -17,34 +15,14 @@ from vllm.logger import init_logger from vllm.multimodal.utils import get_mm_features_in_window from vllm.utils.hashing import xxhash, xxhash_cbor -from vllm.utils.math_utils import cdiv, round_up -from vllm.utils.mem_utils import format_gib -from vllm.utils.torch_utils import get_dtype_size -from vllm.v1.hisparse.layout import ( - get_hisparse_host_pool_bytes, - get_hisparse_kv_cache_config, - get_hisparse_kv_cache_groups, -) +from vllm.utils.math_utils import cdiv from vllm.v1.kv_cache_interface import ( - AttentionSpec, - ChunkedLocalAttentionSpec, - CircularBufferSpec, - FullAttentionSpec, - HiddenStateCacheSpec, - HiSparseHotSpec, KVCacheConfig, KVCacheGroupSpec, - KVCacheLayout, KVCacheSpec, - KVCacheTensor, MambaSpec, - MLAAttentionSpec, - SlidingWindowMLASpec, - SlidingWindowSpec, UniformTypeKVCacheSpecs, - compute_layout_strides, iter_layer_specs, - replace_as, ) from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.request import Request @@ -1000,198 +978,6 @@ def request_block_hasher(request: Request) -> list[BlockHash]: return request_block_hasher -def _check_enough_kv_cache_memory( - available_memory: int, - get_needed_memory: Callable[[], int], - max_model_len: int, - estimate_max_model_len: Callable[[int], int], -): - if available_memory <= 0: - raise ValueError( - "No available memory for the cache blocks. " - "Try increasing `gpu_memory_utilization` when initializing the engine " - "(this flag also controls CPU memory reservation on the CPU " - "backend, despite its name). " - "See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ " - "for more details." - ) - - needed_memory = get_needed_memory() - - if needed_memory > available_memory: - estimated_max_len = estimate_max_model_len(available_memory) - estimated_msg = "" - if estimated_max_len > 0: - estimated_msg = ( - "Based on the available memory, " - f"the estimated maximum model length is {estimated_max_len}. " - ) - - raise ValueError( - f"To serve at least one request with the model's max seq len " - f"({max_model_len}), ({format_gib(needed_memory)} GiB KV " - f"cache is needed, which is larger than the available KV cache " - f"memory ({format_gib(available_memory)} GiB). {estimated_msg}" - f"Try increasing `gpu_memory_utilization` (which also controls " - f"CPU memory on the CPU backend) or decreasing `max_model_len` " - f"when initializing the engine. " - f"See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ " - f"for more details." - ) - - -def max_memory_usage_bytes( - vllm_config: VllmConfig, kv_cache_specs: Iterable[KVCacheSpec] -) -> int: - """Get the maximum memory usage in bytes for the given KV cache specs.""" - return sum(spec.max_memory_usage_bytes(vllm_config) for spec in kv_cache_specs) - - -def estimate_max_model_len( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], - available_memory: int, -) -> int: - """Estimates the maximum model length that can fit in the available memory - using binary search. - - This function temporarily modifies max_model_len during estimation but - restores the original value before returning, ensuring no side effects. - - Args: - vllm_config: The global VllmConfig - kv_cache_spec: The kv cache spec of each attention layer in the model - available_memory: Memory available for KV cache in bytes. - - Returns: - The estimated maximum model length that can fit in the available memory. - - """ - # Save the original max_model_len to restore after estimation - original_max_model_len = vllm_config.model_config.max_model_len - - # Define a function to check if a given model length fits in memory - def fits_in_memory(model_len: int) -> bool: - # Temporarily modify the max_model_len for this calculation - vllm_config.model_config.max_model_len = model_len - # Calculate memory needed for the given model length - memory_needed = max_memory_usage_bytes(vllm_config, kv_cache_spec.values()) - return memory_needed <= available_memory - - try: - # Binary search for the maximum model length - left, right = 1, original_max_model_len - - # If even the smallest model length doesn't fit, return 0 - if not fits_in_memory(left): - return 0 - - # Binary search for the maximum model length that fits - result = 1 - while left <= right: - mid = (left + right) // 2 - if fits_in_memory(mid): - result = mid - left = mid + 1 - else: - right = mid - 1 - return result - finally: - # Always restore the original max_model_len to avoid side effects - vllm_config.model_config.max_model_len = original_max_model_len - - -def check_enough_kv_cache_memory( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], - available_memory: int, -): - """Checks whether `available_memory` is enough for the KV cache to hold at - least one request with the model's max_model_len. - - Args: - vllm_config: The global VllmConfig - kv_cache_spec: The kv cache spec of each attention layer in the model - available_memory: Memory available for KV cache in bytes. - - Raises: - ValueError: If there is not enough memory available for the KV cache. - - """ - # No need to check for available memory if the kv_cache_spec is empty - if kv_cache_spec: - # Reserve the null block BlockPool permanently holds back, so the check - # plans against usable blocks, as in get_kv_cache_configs. Group a copy - # of the specs since grouping may unify them in-place. - groups = get_kv_cache_groups(vllm_config, dict(kv_cache_spec)) - check_memory = ( - available_memory - _pool_bytes_per_block(vllm_config, groups) - if groups - else available_memory - ) - _check_enough_kv_cache_memory( - check_memory, - lambda: max_memory_usage_bytes(vllm_config, kv_cache_spec.values()), - vllm_config.model_config.max_model_len, - lambda am: estimate_max_model_len(vllm_config, kv_cache_spec, am), - ) - - -def create_kv_cache_group_specs( - kv_cache_spec: dict[str, KVCacheSpec], grouped_layer_names: list[list[str]] -) -> list[KVCacheGroupSpec]: - """Create KVCacheGroupSpec object for each kv cache group layer. - The layers in the same group should share the same - KVCacheSpec. - - Args: - kv_cache_spec: - A mapping from each layer name to its corresponding KVCacheSpec. - grouped_layer_names: - A list of kv cache groups, where each element is a list of layer - names that belong to the same group and should share the same - KVCacheSpec. - - Returns: - A list of KVCacheGroupSpec objects, one for each group. - - """ - kv_cache_groups = [] - for layer_names_one_group in grouped_layer_names: - layer_specs = [ - kv_cache_spec[layer_name] for layer_name in layer_names_one_group - ] - merged_layer_spec = layer_specs[0].merge(layer_specs) - kv_cache_groups.append( - KVCacheGroupSpec(layer_names_one_group, merged_layer_spec) - ) - return kv_cache_groups - - -def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: - """Whether all layers in the given KVCacheSpec have the same KV cache spec. - Note that we regard FullAttentionSpec with and without sliding window as - the same type. - - Args: - kv_cache_spec: The kv cache spec of each attention layer in the model - - Returns: - True if all layers have the same type, False otherwise. - - """ - if not kv_cache_spec: - # Encoder-only models do not have KV cache, kv_cache_type can be - # regarded as uniform. - return True - try: - kv_cache_spec_values = list(kv_cache_spec.values()) - _ = kv_cache_spec_values[0].merge(kv_cache_spec_values) - except AssertionError: - return False - return True - - def get_max_concurrency_for_kv_cache_config( vllm_config: VllmConfig, kv_cache_config: KVCacheConfig ) -> float: @@ -1227,1052 +1013,6 @@ def get_max_concurrency_for_kv_cache_config( return min(limits) -def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int: - """Override the number of kv cache blocks if `num_gpu_blocks_override` is set. - The override is logged once, at the call site in `get_kv_cache_configs`. - """ - if vllm_config.cache_config.num_gpu_blocks_override is not None: - num_blocks = vllm_config.cache_config.num_gpu_blocks_override - return num_blocks - - -def _pool_bytes_per_block( - vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec] -) -> int: - """Bytes consumed by one block in the worker's shared KV cache pool, mirroring - the divisor used by `get_kv_cache_config_from_groups` to convert - `available_memory` into `num_blocks`. Used to compute the effective KV cache - capacity once `num_gpu_blocks_override` is applied. - """ - return _get_kv_cache_bytes_per_block( - [group for group in kv_cache_groups if not group.host_resident], - vllm_config.cache_config.get_resolved_kv_cache_layout(), - ) - - -def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int: - """Get the page size of the KV cache.""" - page_sizes = {layer.page_size_bytes for layer in kv_cache_specs} - assert len(page_sizes) == 1 - return page_sizes.pop() - - -def _get_kv_cache_groups_uniform_spec( - kv_cache_specs: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec]: - """Generates the KV cache configuration for a model with the same KV cache - spec for all layers. - - Args: - kv_cache_specs: The kv cache spec of each attention layer in the model - - Returns: - The generated KVCacheGroupSpecs - - """ - return create_kv_cache_group_specs(kv_cache_specs, [list(kv_cache_specs.keys())]) - - -def _get_kv_cache_groups_uniform_type( - spec: UniformTypeKVCacheSpecs, -) -> list[KVCacheGroupSpec]: - """Generates the KV cache configuration for a model with one type of KV cache - but different hidden sizes. All layers are merged into one group. - - Args: - spec: The UniformTypeKVCacheSpecs of the model - - Returns: - The generated KVCacheGroupSpecs - - """ - return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)] - - -def unify_kv_cache_spec_page_size( - kv_cache_spec: dict[str, KVCacheSpec], -) -> dict[str, KVCacheSpec]: - """Unify the page size of the given KVCacheSpec. If the page size of all layers - are the same, return the original KVCacheSpec. If not same, unify the page - size by increasing the block size of layers with smaller page size. Two - cases cannot be unified by block size alone and pad their physical page to - the maximum instead: Mamba layers, whose page size comes from state shapes - and is independent of block size; and non-MLA attention layers whose page - does not evenly divide the maximum (the padded page is read through a - strided view). MLA is excluded because sparse MLA indexes the cache in - whole token rows (see ``flat_kv_row_view``), so its block stride can only - be padded by its own row-aligned ``alignment``, not to an arbitrary page - size. Raise NotImplementedError if failed to unify the page size; - ``get_kv_cache_groups`` catches it to try the full-allocation fallback - (e.g. MLA next to an incompatible sliding-window draft). - - Args: - kv_cache_spec: The KVCacheSpec of each attention layer in the model - - Returns: - The updated KVCacheSpec with the same page_size_bytes. - - """ - page_sizes = {layer.page_size_bytes for layer in kv_cache_spec.values()} - if len(page_sizes) <= 1: - # All layers have the same page size, no need to unify. - return kv_cache_spec - - max_page_size = max(page_sizes) - new_kv_cache_spec = {} - for layer_name, layer_spec in kv_cache_spec.items(): - if layer_spec.page_size_bytes == max_page_size: - new_kv_cache_spec[layer_name] = layer_spec - elif isinstance(layer_spec, MambaSpec): - # MambaSpec's page size is determined by its state shapes and does - # not scale with block_size, so pad the page instead. This is the - # same padding mechanism the platform uses to align Mamba pages - # with the main model's attention page size; it is needed here - # when another layer (e.g. from a draft model) has a larger page - # than the already-aligned Mamba page. - new_spec: KVCacheSpec = replace(layer_spec, page_size_padded=max_page_size) - assert new_spec.page_size_bytes == max_page_size - new_kv_cache_spec[layer_name] = new_spec - else: - layer_page_size = layer_spec.page_size_bytes - if max_page_size % layer_page_size == 0: - ratio = max_page_size // layer_page_size - new_block_size = layer_spec.block_size * ratio - new_spec = replace(layer_spec, block_size=new_block_size) - elif isinstance(layer_spec, AttentionSpec) and not isinstance( - layer_spec, MLAAttentionSpec - ): - new_spec = replace(layer_spec, page_size_padded=max_page_size) - else: - raise NotImplementedError( - f"Layer {layer_name}: page size is not divisible by the " - "maximum page size and cannot be padded. Padding is only " - "supported for non-MLA attention layers." - ) - assert new_spec.page_size_bytes == max_page_size - new_kv_cache_spec[layer_name] = new_spec - return new_kv_cache_spec - - -def is_kv_cache_type_attention_free(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: - # kv_cache_spec is an empty dict for attention free models - return not kv_cache_spec - - -def _get_kv_cache_groups_uniform_page_size( - kv_cache_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec]: - """Generates the KV cache groups for hybrid models with multiple - attention types but still with a uniform page size (physical memory per - block per layer) for all layers. - - Detailed explanation about kv cache management of hybrid models: - The layers in the models are repeated with some patterns, e.g., a model - with 10 full attention layers and 20 sliding window attention layers can be - regarded as repeating the pattern (1 * full, 2 * sw) 10 times. - The KVCacheManager allocates different block tables for each of the 3 layers - in the pattern, and repeats each of them 10 times to generate the - block_table for the 30 layers in the model. - Therefore, we can group the layers in the model into 3 kv_cache_groups, each - of which contains 10 layers in the model. - The KVCacheManager allocates the block_table for each group based on its - kv_cache spec, and the model runner applies the block table to each layer - in the group. - For example: - 1. A model only uses full attention. The pattern is - (num_hidden_layers * full), so there is only one group and the block table - is shared by all layers. It is already handled by - `_get_kv_cache_config_uniform_type`. - 2. A model with 10 full attention layers and 20 sliding window - attention layers. There are 3 layers in the pattern (1 * full, 2 * sw), so - there are 3 kv_cache_groups, each of which represents 10 layers. - - To simplify the implementation, we make the following assumptions: - 1. Physical memory per block: Must be the same across all KV cache groups. - Breaking this assumption is non-trivial due to memory fragmentation concerns - when allocating blocks of different sizes. - 2. Tokens per block (block_size): Currently, we directly use - `CacheConfig.block_size` for all layers. It can be extended to vary by KV - cache group, but within each KV cache group, all layers must share the same - block size. - 3. Physical memory per token per layer: This property is decided by model - config. Currently we only support models that have the same physical memory - per token per layer for all layers. Can be relaxed with a simple extension, - but still need to keep physical memory per block the same for all groups. - 4. Number of layers per group: Currently assumed the same for all layers. - Can be relaxed with a simple extension, but still need to keep physical - memory per block the same for all groups. - 5. Attention type within groups: All layers in a group must share the same - attention type. One exception is that, when - `--disable-hybrid-kv-cache-manager` is true, the single group for full - attention layers may also include attention layers using sliding window or - LLaMA 4 local attention. See `unify_hybrid_kv_cache_specs` for more details. - 6. Support for multiple attention types: The design for most components is - general to an arbitrary number of attention types. But - `find_longest_cache_hit` only supports one attention type or two - types of full-attention plus exactly one another type. The general - implementation of this function is feasible but we don't know how to - implement it cleanly yet. - - As we assume tokens per block, physical memory per token per layer, and - number of layers per group are the same now, we can ensure that physical - memory per block is the same for all groups. - - Args: - kv_cache_spec: The KVCacheSpec of each attention layer in the model - Returns: - The generated KVCacheGroupSpecs - - """ - # Group all layers by kv_cache_spec. - # E.g., 2 full attention layers and 3 sliding window attention layers, - # -> (full.0, full.1), (sw.0, sw.1, sw.2). - same_type_layers: dict[KVCacheSpec, list[str]] = defaultdict(list) - for layer_name, layer_spec in kv_cache_spec.items(): - same_type_layers[layer_spec].append(layer_name) - - # Attempt to further merge same-type layers based on whether their KV - # cache specs can be merged, to minimize the group count. This benefits - # situations where specs share a block layout and differ only in a - # property it can reconcile (e.g. full attention layers differing only in - # sliding window / attention chunk size). - layer_buckets: list[list[str]] = [] - spec_buckets: list[list[KVCacheSpec]] = [] - for layer_spec, layer_names in same_type_layers.items(): - for names, specs in zip(layer_buckets, spec_buckets): - try: - # A raise means that the specs are incompatible. - type(specs[0]).merge([*specs, layer_spec]) - except (AssertionError, ValueError): - continue - names.extend(layer_names) - specs.append(layer_spec) - break - else: - layer_buckets.append(list(layer_names)) - spec_buckets.append([layer_spec]) - - # Split each group into smaller groups, to make the number of layers in each - # group identical. Add padding to the last group of each type if necessary. - # E.g., (full.0, full.1), (sw.0, sw.1, sw.2) - # split to 3 groups with 2 layers each: - # (full.0, full.1), (sw.0, sw.2), (sw.1, padding). - # FIXME(Chen): At the moment of writing this code (2025-06-02), all - # open-source hybrid model follows a n:1 pattern between different attention - # types (e.g., Gemma3 5:1 between sw and full, LLaMA4 3:1 between local and - # full), so we can use the "1" in the n:1 pattern as the group size, which - # is the minimum number of layers among all attention types. Need a better - # strategy if we want to support more complex patterns (e.g., 20 full + 30 - # sw, where the group size should be 10). - min_num_layers = min([len(layers) for layers in layer_buckets]) - group_size = min_num_layers - max_num_layers = max([len(layers) for layers in layer_buckets]) - if max_num_layers < min_num_layers * 1.5: - # If the number of layers is not much larger than the minimum number of - # layers, use the maximum number of layers as the group size to avoid - # too many padding layers. A typical example is gpt-oss-20b + eagle, - # with 12 sw + 13 full. We pad it to (13 sw, 13 full) instead of - # (12 sw, 24 full). 1.5 is a heuristic to avoid too many padding - # layers while accommodating speculative decoding drafters that add - # extra layers to one attention type. - group_size = max_num_layers - grouped_layers = [] - for layers in layer_buckets: - num_padding_layers = group_size - len(layers) % group_size - if num_padding_layers != group_size: - logger.warning( - "Add %d padding layers, may waste at most %.2f%% KV cache memory", # noqa - num_padding_layers, - num_padding_layers / len(layers) * 100, - ) - num_groups = cdiv(len(layers), group_size) - # In PP case, say if we have - # - stage 0: full.0, sw.0, sw.1 - # - stage 1: full.1, sw.2, sw.3 - # We should have 3 groups: (full.0, full.1), (sw.0, sw.2), (sw.1, sw.3) - # It can't be (full.0, full.1), (sw.0, sw.1), (sw.2, sw.3) because - # the 3 groups in stage 0 will be (full.0), (sw.0, sw.1), (empty group) - # and it will be padded to (full.0, padding), (sw.0, sw.1), - # (padding, padding) to ensure the number of layers in each group is - # the same and will cause memory waste. - # To avoid this, we assign layers[i::num_groups] to the i-th group - # instead of layers[i * group_size: (i + 1) * group_size] - for i in range(num_groups): - grouped_layers.append(layers[i::num_groups]) - return create_kv_cache_group_specs(kv_cache_spec, grouped_layers) - - -def _get_per_layer_spec( - group: KVCacheGroupSpec, - layer_name: str, -) -> KVCacheSpec: - spec = group.kv_cache_spec - if isinstance(spec, UniformTypeKVCacheSpecs): - return spec.kv_cache_specs[layer_name] - return spec - - -def _get_layer_block_stride(spec: KVCacheSpec) -> int: - """A layer-outer block stride: the page padded to its alignment.""" - return round_up(spec.page_size_bytes, spec.get_block_stride_alignment()) - - -def _get_kv_cache_bytes_per_block( - kv_cache_groups: list[KVCacheGroupSpec], layout: KVCacheLayout -) -> int: - """Return the largest cache group's bytes per block.""" - if not layout.is_block_outermost: - # Each layer's block stride is its own padded page. - return max( - sum( - _get_layer_block_stride(_get_per_layer_spec(group, layer_name)) - for layer_name in group.layer_names - ) - for group in kv_cache_groups - ) - bytes_per_block = max( - sum( - _get_per_layer_spec(group, layer_name).page_size_bytes - for layer_name in group.layer_names - ) - for group in kv_cache_groups - ) - assert bytes_per_block > 0 - alignments = [ - group.kv_cache_spec.page_size_bytes - for group in kv_cache_groups - if isinstance(group.kv_cache_spec, HiSparseHotSpec) - ] - alignments.extend( - _get_per_layer_spec(group, layer_name).get_block_stride_alignment() - for group in kv_cache_groups - for layer_name in group.layer_names - ) - return round_up(bytes_per_block, math.lcm(*alignments)) - - -def validate_kv_cache_layout( - layout: KVCacheLayout, - kv_cache_groups: list[KVCacheGroupSpec], -) -> None: - """Validate that the resolved layout can express this model's packing. - - The layout was chosen once in the engine core from the backends' supported - sets; a backend whose model packs pages side by side (e.g. the DeepSeek-V4 - indexer) declares block-outermost layouts there, so an inexpressible - layout reaching this point is an error. - """ - page_sizes = { - _get_per_layer_spec(group, layer_name).page_size_bytes - for group in kv_cache_groups - for layer_name in group.layer_names - } - if len(page_sizes) == 1: - # A rectangular layer dim exists; every layout can express it. - return - - # Mixed page sizes pack pages side by side within a block, which needs each page - # to be one contiguous chunk inside its block (a block-compact layout) and, with - # multiple KV cache groups, the layer dim inside the block dim. - if not layout.is_block_compact or ( - len(kv_cache_groups) > 1 and layout.is_layer_compact - ): - raise ValueError( - f"KV cache layout {layout.name} cannot express this model's " - f"mixed page sizes ({sorted(page_sizes)}); a backend should " - "declare block-outermost supported layouts (e.g. BLHNC), or " - "set VLLM_KV_CACHE_LAYOUT=BLHNC." - ) - - -def get_kv_cache_config_from_groups( - vllm_config: VllmConfig, - kv_cache_groups: list[KVCacheGroupSpec], - available_memory: int, -) -> KVCacheConfig: - """Generate the KV cache configuration from the KV cache groups and spec - of each layer. - - Args: - vllm_config: The global VllmConfig - kv_cache_groups: The KV cache groups - available_memory: Memory available for KV cache in bytes - Returns: - The generated KVCacheConfig - - """ - if len(kv_cache_groups) == 0: - # Attention free models do not have KV cache. - # Return num_blocks=1 as BlockPool always needs a null_block. - return KVCacheConfig( - num_blocks=1, - kv_cache_tensors=[], - kv_cache_groups=kv_cache_groups, - prefix_cache_retention_interval=( - vllm_config.cache_config.prefix_cache_retention_interval - ), - ) - - if vllm_config.attention_config.hisparse_config is not None: - host_budget = get_hisparse_host_pool_bytes(vllm_config) - return get_hisparse_kv_cache_config( - vllm_config, kv_cache_groups, available_memory, host_budget - ) - - layout = vllm_config.cache_config.get_resolved_kv_cache_layout() - validate_kv_cache_layout(layout, kv_cache_groups) - bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups, layout) - interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None - - num_blocks = available_memory // bytes_per_block - num_blocks = may_override_num_blocks(vllm_config, num_blocks) - size = bytes_per_block * num_blocks - - # Groups alias from byte 0. Spec regions are laid out differently: - # - # block-outer (the same packing repeats for every block): - # group 0: | blk 0 [ A | B | pad ] | blk 1 [ A | B | pad ] | ... - # group 1: | blk 0 [ C | D ] | blk 1 [ C | D ] | ... - # |<--- bytes_per_block -->| - # - # layer-outer (only supported for uniform page sizes or single-group models): - # group 0: | A [ blk 0 | blk 1 | ... ] | B [ blk 0 | blk 1 | ... ] | - # group 1: | C [ blk 0 | blk 1 | ... ] | D [ blk 0 | blk 1 | ... ] | - - kv_cache_tensors = [] - for group in kv_cache_groups: - group_spec = group.kv_cache_spec - layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list) - if isinstance(group_spec, UniformTypeKVCacheSpecs): - for layer_name, spec in group_spec.kv_cache_specs.items(): - layers_by_spec[spec].append(layer_name) - elif group.layer_names: - layers_by_spec[group_spec].extend(group.layer_names) - - byte_offset = 0 - for spec, layer_names in layers_by_spec.items(): - block_stride = interleaved_block_stride - padded_page = _get_layer_block_stride(spec) - if block_stride is None and padded_page != spec.page_size_bytes: - assert layout.is_block_compact, f"{layout.name} cannot pad pages" - block_stride = padded_page - layer_stride, block_stride, _, _, _ = compute_layout_strides( - spec, - num_blocks, - len(layer_names), - layout, - fixed_strides=(None, block_stride, None, None, None), - ) - kv_cache_tensors.append( - KVCacheTensor( - size=size, - layers=layer_names, - layer_stride=layer_stride, - block_stride=block_stride, - offset=byte_offset, - ) - ) - # Block-outer layers share each block; layer-outer ones own regions. - byte_offset += len(layer_names) * ( - spec.page_size_bytes if layout.is_block_outermost else layer_stride - ) - - return KVCacheConfig( - num_blocks=num_blocks, - kv_cache_tensors=kv_cache_tensors, - kv_cache_groups=kv_cache_groups, - prefix_cache_retention_interval=( - vllm_config.cache_config.prefix_cache_retention_interval - ), - ) - - -def _promote_local_kv_cache_specs( - kv_cache_spec: dict[str, KVCacheSpec], -) -> dict[str, KVCacheSpec]: - """Use full-attention allocation for local-attention cache specs. - - The returned specs affect KV cache management only. Attention modules keep - their original sliding-window or chunked-local compute behavior. - """ - promoted_specs = kv_cache_spec.copy() - - if is_kv_cache_spec_uniform( - promoted_specs - ) or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs): - return promoted_specs - - has_full_attention = any( - isinstance(spec, FullAttentionSpec) for spec in promoted_specs.values() - ) - has_sliding_window = any( - isinstance(spec, SlidingWindowSpec) for spec in promoted_specs.values() - ) - has_chunked_local_attention = any( - isinstance(spec, ChunkedLocalAttentionSpec) for spec in promoted_specs.values() - ) - full_block_sizes = { - spec.block_size - for spec in promoted_specs.values() - if isinstance(spec, FullAttentionSpec) - } - full_attention_block_size = ( - next(iter(full_block_sizes)) if len(full_block_sizes) == 1 else None - ) - - def promoted_page_size_padded(spec: AttentionSpec, block_size: int) -> int | None: - if spec.page_size_padded is None: - return None - unpadded_page_size = ( - spec.unpadded_page_size_bytes * block_size // spec.block_size - ) - return max(spec.page_size_padded, unpadded_page_size) - - promotions: dict[type[AttentionSpec], type[AttentionSpec]] = { - SlidingWindowMLASpec: MLAAttentionSpec, - SlidingWindowSpec: FullAttentionSpec, - ChunkedLocalAttentionSpec: FullAttentionSpec, - } - - if has_full_attention and (has_sliding_window or has_chunked_local_attention): - for layer_name, spec in kv_cache_spec.items(): - target_cls = next( - (promotions[c] for c in type(spec).__mro__ if c in promotions), None - ) - if target_cls is None: - continue - assert isinstance(spec, AttentionSpec) - block_size = full_attention_block_size or spec.block_size - promoted_specs[layer_name] = replace_as( - spec, - target_cls, - # Promoted specs allocate blocks for all tokens and never free - # below the window, so the trailing-edge extension is moot. - drop=("extra_retained_tokens",), - block_size=block_size, - page_size_padded=promoted_page_size_padded(spec, block_size), - ) - - if not ( - is_kv_cache_spec_uniform(promoted_specs) - or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs) - ): - raise ValueError("Failed to promote local KV cache specs to one unified type.") - - return promoted_specs - - -def _try_get_full_allocation_fallback_groups( - kv_cache_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec] | None: - """Try a supported full-allocation fallback for local-attention layers.""" - if any(isinstance(spec, HiddenStateCacheSpec) for spec in kv_cache_spec.values()): - return None - if any( - isinstance(spec, (SlidingWindowMLASpec, ChunkedLocalAttentionSpec)) - for spec in kv_cache_spec.values() - ): - return None - - has_mla = any(isinstance(spec, MLAAttentionSpec) for spec in kv_cache_spec.values()) - has_regular_swa = any( - isinstance(spec, SlidingWindowSpec) for spec in kv_cache_spec.values() - ) - if not (has_mla and has_regular_swa): - return None - - try: - promoted_specs = _promote_local_kv_cache_specs(kv_cache_spec) - except ValueError: - return None - uniform_spec = UniformTypeKVCacheSpecs.from_specs(promoted_specs) - if uniform_spec is None: - return None - logger.warning( - "KV cache page sizes cannot be unified; treating sliding-window " - "layers as full attention for cache allocation. Sliding-window " - "attention compute is unchanged." - ) - return _get_kv_cache_groups_uniform_type(uniform_spec) - - -def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]): - """This function tries to convert the KV cache specs to one type if the model - is a hybrid model with multiple type of KV cache. It will convert all - SlidingWindowSpec to FullAttentionSpec if both types are present. - - Args: - kv_cache_spec: The kv cache spec of each attention layer in the model - - """ - groups: defaultdict[bool, dict[str, KVCacheSpec]] = defaultdict(dict) - for name, spec in kv_cache_spec.items(): - replicated = isinstance(spec, AttentionSpec) and not spec.dcp_sharded - groups[replicated][name] = spec - for specs in groups.values(): - promoted_specs = _promote_local_kv_cache_specs(specs) - if promoted_specs == specs: - continue - logger.warning( - "Hybrid KV cache manager is disabled for this hybrid model, " - "This means we do not enable any optimizations for saving KV cache " - "memory (e.g., dropping the KV cache outside the sliding window). " - "The compute of layers like sliding window is still saved." - ) - kv_cache_spec.update(promoted_specs) - - -def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int: - """Pick a chunk size that minimizes total upward padding. - - Each x is rounded up to a multiple of d: - - x -> ceil(x / d) * d - - Total padding is: - - pad(d) = sum_i (ceil(x_i / d) * d - x_i) - - We brute-force d in [lower_bound, max(values)] (fine for small lists / small - maxima) and return the d with minimum padding. Ties prefer larger d. - """ - if not values: - raise ValueError("values must be non-empty") - if any(x <= 0 for x in values): - raise ValueError(f"values must be positive, got: {list(values)!r}") - - min_d = max(1, lower_bound if lower_bound is not None else 1) - max_d = max(values) - if min_d > max_d: - return min_d - - best_d = min_d - best_pad: int | None = None - for d in range(min_d, max_d + 1): - pad = sum((d - (x % d)) % d for x in values) - if best_pad is None or pad < best_pad or (pad == best_pad and d > best_d): - best_pad = pad - best_d = d - - return best_d - - -def _get_packed_kv_cache_groups( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec] | None: - """Group mixed-page-size layers for contiguous block-outermost packing. - - Greedily buckets layers into uniform-type specs. Buckets with equal layer - counts per page size are treated as a repeating layer pattern (one layer - per page size) and split into groups covering the same number of pattern - repeats (picked by ``_approximate_gcd`` to minimize padding), so all - groups pack into the same per-block layout. Mamba buckets are additionally - split to fit the block the attention buckets already need. - Returns None when the layout is not block-outermost or all layers already - share one page size. - """ - layout = vllm_config.cache_config.get_resolved_kv_cache_layout() - page_sizes = {spec.page_size_bytes for spec in kv_cache_spec.values()} - if not layout.is_block_outermost or len(page_sizes) <= 1: - return None - - buckets: list[dict[str, KVCacheSpec]] = [] - for name, spec in kv_cache_spec.items(): - for bucket in buckets: - candidate = {**bucket, name: spec} - if UniformTypeKVCacheSpecs.is_uniform_type(candidate): - bucket[name] = spec - break - else: - buckets.append({name: spec}) - - bucketed = [] - for bucket in buckets: - uniform_spec = UniformTypeKVCacheSpecs.from_specs(bucket) - assert uniform_spec is not None - page_size_layers: dict[int, list[str]] = defaultdict(list) - for layer_name, layer_spec in bucket.items(): - page_size_layers[layer_spec.page_size_bytes].append(layer_name) - # Only 1:1 patterns (one layer of each page size per repeat) are - # supported; counts sharing a gcd > 1 (e.g. 2:1) could in principle - # repeat too, but such buckets are emitted whole instead. - balanced = len(set(map(len, page_size_layers.values()))) == 1 - bucketed.append((uniform_spec, page_size_layers, balanced)) - - # Balanced buckets that mix page sizes must stay whole, so the largest one - # sets a floor on the repeats per group; larger single-size buckets are - # split down toward it. No such bucket means nothing needs packing. - min_repeats_per_group = max( - ( - spec.get_max_layers_per_page_size() - for spec, page_size_layers, balanced in bucketed - if balanced and len(page_size_layers) > 1 - ), - default=0, - ) - repeats_per_group = ( - _approximate_gcd( - [ - spec.get_max_layers_per_page_size() - for spec, _, balanced in bucketed - if balanced - ], - lower_bound=min_repeats_per_group, - ) - if min_repeats_per_group - else None - ) - - def num_groups_for(spec: UniformTypeKVCacheSpecs, balanced: bool) -> int: - if balanced and repeats_per_group is not None: - return cdiv(spec.get_max_layers_per_page_size(), repeats_per_group) - return 1 - - def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int: - """Page bytes of the largest of the n groups a bucket splits into.""" - return sum( - cdiv(len(names), n) * page for page, names in page_size_layers.items() - ) - - # Bytes a block must hold however the state buckets end up split: mamba, - # circular-buffer and unsplit sliding-window buckets can go down to one - # state per group, every other bucket's split is fixed by the repeat pattern. - def is_state_bucket(spec: UniformTypeKVCacheSpecs) -> bool: - if isinstance(spec.first_spec, (MambaSpec, CircularBufferSpec)): - return True - return repeats_per_group is None and isinstance( - spec.first_spec, SlidingWindowSpec - ) - - anchor_bytes = max( - ( - widest_group_bytes( - page_size_layers, - len(spec.kv_cache_specs) - if is_state_bucket(spec) - else num_groups_for(spec, balanced), - ) - for spec, page_size_layers, balanced in bucketed - ), - default=0, - ) - - groups = [] - for spec, page_size_layers, balanced in bucketed: - num_groups = num_groups_for(spec, balanced) - # Cap a state group at the states a block already fits rather than let - # it widen the block. - if anchor_bytes and is_state_bucket(spec): - states_per_block = max(anchor_bytes // spec.first_spec.page_size_bytes, 1) - num_groups = max( - num_groups, cdiv(len(spec.kv_cache_specs), states_per_block) - ) - if num_groups == 1: - groups.append(KVCacheGroupSpec(list(spec.kv_cache_specs), spec)) - continue - - pattern_repeats = list(zip(*page_size_layers.values())) - for i in range(num_groups): - group_layer_names = [ - name for repeat in pattern_repeats[i::num_groups] for name in repeat - ] - group_layer_specs = { - name: spec.kv_cache_specs[name] for name in group_layer_names - } - group_spec = UniformTypeKVCacheSpecs.from_specs(group_layer_specs) - assert group_spec is not None - groups.append(KVCacheGroupSpec(group_layer_names, group_spec)) - - _annotate_eagle_groups( - vllm_config, - kv_cache_spec, - groups, - use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), - ) - _warn_if_unannotated_eagle_mamba(vllm_config, groups) - return groups - - -def _uses_trailing_mtp_layers(vllm_config: VllmConfig) -> bool: - """Whether the drafter's KV layers can be located positionally. - - MTP drafters register their KV layers after the target layers, without a - distinguishing spec marker. DeepseekV4/V4.1 do this with ``dspark`` too. - The annotator separately checks that the groups partition the layers exactly. - """ - spec_config = vllm_config.speculative_config - if spec_config is None or not spec_config.use_eagle(): - return False - if spec_config.method == "mtp": - return True - model_config = vllm_config.model_config - return model_config is not None and model_config.hf_config.model_type in ( - "deepseek_v4", - "deepseek_v41", - ) - - -def _groups_partition_layers_exactly( - kv_cache_spec: dict[str, KVCacheSpec], - kv_cache_groups: list[KVCacheGroupSpec], -) -> bool: - """Whether the groups cover every layer of ``kv_cache_spec`` exactly once. - - The trailing-layer draft fallback is only meaningful when the last - registered layer is guaranteed to appear in exactly one group. - """ - names = [name for group in kv_cache_groups for name in group.layer_names] - return len(names) == len(kv_cache_spec) and set(names) == set(kv_cache_spec) - - -def _annotate_eagle_groups( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], - kv_cache_groups: list[KVCacheGroupSpec], - use_trailing_layer_fallback: bool = False, -) -> None: - """Flag the KV cache groups that hold drafter attention layers. - - Two detection rules, in order of preference: - - 1. Spec-driven. ``non_causal_multi_token_decode`` is declared on - MLAAttentionSpec and set by drafter attention layers that run a - non-causal multi-token decode (today only Kimi-K3 DSpark). It survives - MLAAttentionSpec.merge, so it still identifies a group after per-group - spec merging, wherever grouping happens to land. It is sufficient but - not necessary: a drafter whose spec is indistinguishable from the - target's cannot be found this way. - 2. Positional fallback for MTP drafters (including DeepseekV4/V4.1 DSpark), whose - MTP block reuses the target's own decoder layer and so carries no spec - marker. Their draft attention layers always register after every - target layer, so flag whichever group holds the last registered layer. - This rule is only valid where the groups partition exactly the layers - of ``kv_cache_spec``, which is re-checked here before applying it. - When a drafter's trailing caches span several groups, this rule flags - only the group holding the very last layer and must be generalized. - FIXME(yifan): avoid/generalize this hacky check. - - Args: - vllm_config: Config supplying the speculative method, if any. - kv_cache_spec: The kv cache spec of each attention layer, in layer - registration order. Only read by rule 2. - kv_cache_groups: Groups to annotate in place. - use_trailing_layer_fallback: Enable rule 2. Callers gate this on - ``_uses_trailing_mtp_layers``. - - """ - spec_config = vllm_config.speculative_config - if spec_config is None or not spec_config.use_eagle_block_drop(): - return - - for group in kv_cache_groups: - if any( - getattr(spec, "non_causal_multi_token_decode", False) - for spec in iter_layer_specs(group.kv_cache_spec) - ): - group.is_eagle_group = True - - if not use_trailing_layer_fallback: - return - if not _groups_partition_layers_exactly(kv_cache_spec, kv_cache_groups): - return - last_layer = next(reversed(kv_cache_spec)) - for group in kv_cache_groups: - if last_layer in group.layer_names: - group.is_eagle_group = True - break - - -def _warn_if_unannotated_eagle_mamba( - vllm_config: VllmConfig, - kv_cache_groups: list[KVCacheGroupSpec], -) -> None: - """Warn when no KV cache group could be identified as the draft model's. - - Args: - vllm_config: Config supplying the speculative method, if any. - kv_cache_groups: Groups as they will be handed to consumers. - - """ - spec_config = vllm_config.speculative_config - if spec_config is None or not spec_config.use_eagle(): - return - if any(group.is_eagle_group for group in kv_cache_groups): - return - mamba_groups = [ - idx - for idx, group in enumerate(kv_cache_groups) - if any( - isinstance(spec, MambaSpec) - for spec in iter_layer_specs(group.kv_cache_spec) - ) - ] - if not mamba_groups: - return - logger.warning( - "Speculative decoding (method=%s) is enabled but no KV cache group " - "could be identified as the draft model's.", - spec_config.method, - ) - - -def _largest_divisor_at_most(value: int, limit: int) -> int: - for candidate in range(min(value, limit), 0, -1): - if value % candidate == 0: - return candidate - return 1 - - -def _ensure_min_page_size( - groups: list[KVCacheGroupSpec], - common_page: int, - hidden_specs: dict[str, HiddenStateCacheSpec], -) -> tuple[list[KVCacheGroupSpec], int]: - """Scale up group block sizes so the common page is at least as large as - the biggest hidden-state per-token cost. - - This protects the hidden state extraction feature, where we store hidden states in - the KV cache, and return them to clients using a custom hidden state connector. - - Returns the (possibly scaled) groups and updated common page size. - """ - max_per_token = max( - spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) - for spec in hidden_specs.values() - ) - if max_per_token <= common_page: - return groups, common_page - - scale = cdiv(max_per_token, common_page) - common_page *= scale - scaled: list[KVCacheGroupSpec] = [] - for g in groups: - s = g.kv_cache_spec - kw: dict[str, Any] = {"block_size": s.block_size * scale} - if isinstance(s, (AttentionSpec, MambaSpec)) and s.page_size_padded is not None: - kw["page_size_padded"] = s.page_size_padded * scale - scaled.append(KVCacheGroupSpec(g.layer_names, replace(s, **kw))) - return scaled, common_page - - -def get_kv_cache_groups( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec]: - """Split the layers in the model into groups with the same KV cache spec. - - Args: - vllm_config: The global VllmConfig - kv_cache_spec: The kv cache spec of each attention layer in the model - - Returns: - The generated KVCacheGroups - - """ - if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager: - unify_hybrid_kv_cache_specs(kv_cache_spec) - - if is_kv_cache_type_attention_free(kv_cache_spec): - # This returns an empty list to allow for the KVCacheManager to handle - # attention free models. - return [] - - if hisparse_groups := get_hisparse_kv_cache_groups(vllm_config, kv_cache_spec): - return hisparse_groups - - if is_kv_cache_spec_uniform(kv_cache_spec): - # KV cache of all layers are the same, which is true for - # most models. Allocate the same amount of memory for - # each layer. - return _get_kv_cache_groups_uniform_spec(kv_cache_spec) - elif uniform_spec := UniformTypeKVCacheSpecs.from_specs(kv_cache_spec): - # All layers need the same number of token slots (e.g., all layers are - # full attention, or all layers are sliding window attention with the - # same window size). Put all layers into one group. - return _get_kv_cache_groups_uniform_type(uniform_spec) - # Hidden-state layers use their own block table and must not be absorbed - # into a compatible attention bucket. - hidden_specs = { - k: v for k, v in kv_cache_spec.items() if isinstance(v, HiddenStateCacheSpec) - } - filtered_spec = { - k: v - for k, v in kv_cache_spec.items() - if not isinstance(v, HiddenStateCacheSpec) - } - - if packed_groups := _get_packed_kv_cache_groups(vllm_config, filtered_spec): - # Block-outermost blocks are strided by the widest group, so hidden - # groups need no page alignment. - packed_groups += [ - KVCacheGroupSpec([name], spec) for name, spec in hidden_specs.items() - ] - return packed_groups - - # Prefer preserving each layer's cache semantics. If physical pages cannot - # be unified, try a supported allocation-only fallback before failing. - try: - filtered_spec = unify_kv_cache_spec_page_size(filtered_spec) - except NotImplementedError: - fallback_groups = _try_get_full_allocation_fallback_groups(kv_cache_spec) - if fallback_groups is None: - raise - return fallback_groups - groups = _get_kv_cache_groups_uniform_page_size(filtered_spec) - - # Add hidden-state layers back with page aligned to the common page. - if hidden_specs: - common_page = get_uniform_page_size([g.kv_cache_spec for g in groups]) - # TP may shrink the common page below hidden-state per-token cost. - groups, common_page = _ensure_min_page_size(groups, common_page, hidden_specs) - group_block_size = math.gcd(*(g.kv_cache_spec.block_size for g in groups)) - for name, spec in hidden_specs.items(): - per_token = spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) - max_block_size = max(common_page // per_token, 1) - new_bs = _largest_divisor_at_most(group_block_size, max_block_size) - wasted_bytes = common_page - new_bs * per_token - logger.info( - "Using block size %d for hidden-state cache layer %s; " - "page alignment wastes %d bytes (%.2f%%) per block", - new_bs, - name, - wasted_bytes, - wasted_bytes / common_page * 100, - ) - aligned = replace(spec, block_size=new_bs, page_size_padded=common_page) - groups.append(KVCacheGroupSpec([name], aligned)) - - _annotate_eagle_groups( - vllm_config, - kv_cache_spec, - groups, - use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), - ) - _warn_if_unannotated_eagle_mamba(vllm_config, groups) - return groups - - -def _layer_tp_replicas(spec: KVCacheSpec, tp_size: int, dcp_size: int) -> int: - if not isinstance(spec, AttentionSpec) or spec.max_tp_shards is None: - return 1 - if spec.dcp_sharded and dcp_size > 1: - return 1 - return max(1, tp_size // spec.max_tp_shards) - - -def kv_cache_groups_tp_replicas( - groups: list[KVCacheGroupSpec], tp_size: int, dcp_size: int = 1 -) -> int: - """Consecutive TP ranks holding identical KV for every layer.""" - specs = [spec for g in groups for spec in iter_layer_specs(g.kv_cache_spec)] - if not specs: - return 1 - return reduce( - math.gcd, (_layer_tp_replicas(s, tp_size, dcp_size) for s in specs), tp_size - ) - - def generate_scheduler_kv_cache_config( kv_cache_configs: list[KVCacheConfig], ) -> KVCacheConfig: @@ -2334,359 +1074,6 @@ def update_kv_cache_capacity( ) -def _max_memory_usage_bytes_from_groups( - vllm_config: VllmConfig, - kv_cache_groups: list[KVCacheGroupSpec], -) -> int: - """Calculate maximum memory usage in bytes from KV cache groups. - - This correctly accounts for padding in hybrid models. For example, if a - model has 8 full attention layers and 9 sliding window layers, they will - be padded to 9 full + 9 sliding window for uniform group sizes. - - Each group independently claims blocks from the shared pool, so a request consumes - the sum of the per-group block counts, i.e. ``bytes_per_block * total_blocks``. - """ - kv_cache_groups = [group for group in kv_cache_groups if not group.host_resident] - if not kv_cache_groups: - return 0 - - bytes_per_block = _pool_bytes_per_block(vllm_config, kv_cache_groups) - total_blocks = 0 - for group in kv_cache_groups: - spec = group.kv_cache_spec - if isinstance(spec, UniformTypeKVCacheSpecs): - total_blocks += spec.max_memory_usage_pages(vllm_config) - else: - total_blocks += cdiv( - spec.max_memory_usage_bytes(vllm_config), - spec.page_size_bytes, - ) - - return bytes_per_block * total_blocks - - -def _estimate_max_model_len_from_groups( - vllm_config: VllmConfig, - kv_cache_groups: list[KVCacheGroupSpec], - available_memory: int, -) -> int: - """Binary search for the maximum model length that fits in available memory. - Returns 0 if even 1 token doesn't fit. - """ - original_max = vllm_config.model_config.max_model_len - hisparse_enabled = ( - vllm_config.attention_config.hisparse_config is not None - and bool(kv_cache_groups) - ) - - def fits(model_len: int) -> bool: - vllm_config.model_config.max_model_len = model_len - if hisparse_enabled: - try: - get_kv_cache_config_from_groups( - vllm_config, kv_cache_groups, available_memory - ) - except ValueError: - return False - return ( - _max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups) - <= available_memory - ) - - try: - left, right = 1, original_max - if not fits(left): - return 0 - result = 1 - while left <= right: - mid = (left + right) // 2 - if fits(mid): - result = mid - left = mid + 1 - else: - right = mid - 1 - return result - finally: - vllm_config.model_config.max_model_len = original_max - - -def _auto_fit_max_model_len( - vllm_config: VllmConfig, - projected_groups_per_worker: list[list[KVCacheGroupSpec]], - available_memory: list[int], -) -> None: - """When max_model_len is set to -1, this function estimates the largest - context length that can be supported with the available GPU memory. - It uses binary search to find the maximum length that fits across all - workers. - - Args: - vllm_config: The global VllmConfig (will be modified in-place) - projected_groups_per_worker: KV cache groups projected to each worker. - available_memory: Memory available for KV cache in bytes for each - worker. - - """ - original_max = vllm_config.model_config.max_model_len - - if all(not groups for groups in projected_groups_per_worker): - # All workers have empty specs (attention-free model) - logger.info_once( - "Auto-fit max_model_len: attention-free model, " - "using derived max_model_len=%d", - original_max, - ) - return - - # Find the max_model_len that fits across all workers. - auto_fit_max = original_max - limiting_worker_mem = available_memory[0] - for groups, avail_mem in zip(projected_groups_per_worker, available_memory): - if not groups: - continue - worker_max = _estimate_max_model_len_from_groups(vllm_config, groups, avail_mem) - if worker_max < auto_fit_max: - auto_fit_max = worker_max - limiting_worker_mem = avail_mem - - if auto_fit_max <= 0: - raise ValueError( - "Cannot auto-fit max_model_len: not enough GPU memory available " - "to serve even a single token. Try increasing `gpu_memory_utilization`." - ) - - kv_cache_memory = ( - f"{format_gib(limiting_worker_mem)} GiB {_kv_cache_device_name(vllm_config)}" - ) - if vllm_config.attention_config.hisparse_config is not None: - host_pool_gib = format_gib(get_hisparse_host_pool_bytes(vllm_config)) - kv_cache_memory += f" + {host_pool_gib} GiB host" - - if auto_fit_max >= original_max: - # The model's full context length fits in memory - logger.info_once( - "Auto-fit max_model_len: full model context length %d fits in " - "available KV cache memory (%s)", - original_max, - kv_cache_memory, - ) - else: - # Need to reduce max_model_len to fit in memory - vllm_config.model_config.max_model_len = auto_fit_max - logger.info_once( - "Auto-fit max_model_len: reduced from %d to %d to fit in " - "available KV cache memory (%s)", - original_max, - auto_fit_max, - kv_cache_memory, - ) - - -def _project_kv_cache_groups_to_worker( - global_kv_cache_groups: list[KVCacheGroupSpec], - worker_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec]: - """Projects global KV cache groups onto a single worker's assigned layers. - - In pipeline parallelism, each worker only owns a subset of layers. This - function filters the global groups to include only layers present on the - given worker, adjusting UniformTypeKVCacheSpecs accordingly. - - Args: - global_kv_cache_groups: The global KV cache groups for the whole model. - worker_spec: The KV cache spec of each layer on this worker. - - Returns: - The projected KV cache groups containing only this worker's layers. - - """ - projected_groups: list[KVCacheGroupSpec] = [] - for group in global_kv_cache_groups: - worker_layer_names = [ - layer_name for layer_name in group.layer_names if layer_name in worker_spec - ] - if len(worker_layer_names) == len(group.layer_names): - projected_groups.append(group) - continue - group_spec = group.kv_cache_spec - if worker_layer_names and isinstance(group_spec, UniformTypeKVCacheSpecs): - group_spec = UniformTypeKVCacheSpecs( - block_size=group_spec.block_size, - kv_cache_specs={ - layer_name: group_spec.kv_cache_specs[layer_name] - for layer_name in worker_layer_names - }, - ) - projected_groups.append( - KVCacheGroupSpec( - worker_layer_names, - group_spec, - is_eagle_group=group.is_eagle_group and bool(worker_layer_names), - ) - ) - return projected_groups - - -def get_kv_cache_configs( - vllm_config: VllmConfig, - kv_cache_specs: list[dict[str, KVCacheSpec]], - available_memory: list[int], -) -> list[KVCacheConfig]: - """Generates the KV cache configurations for a model. - Since we use a shared centralized controller for all workers, we need the - `kv_cache_config` to be consistent across all workers to make sure - the KV cache allocation can be applied to all workers. However, different - workers may have different memory available, and different type of layers - (when pipeline parallel is enabled). To handle the difference between - workers, the current implementation is: - 1. Merge the KV cache specs of all workers to get the KVCacheSpecs for - the whole model. - 2. Generate the KV cache groups based on the layer ratio of the whole model. - This also handles spec unification for hybrid models. - 3. Handle auto-fit max_model_len and memory checks using per-worker - projected groups to account for PP sharding. - 4. Generate the KV cache configs for each worker based on the KV cache - grouping strategy. (This is reasonable because the layer ratio of - different PP stages are similar.) - 5. Change the num_blocks of each worker to the smallest among all workers - and shrink tensor sizes proportionally to avoid allocating unused memory. - - Args: - vllm_config: The global VllmConfig - kv_cache_specs: List of dict[layer_name, KVCacheSpec] for each worker. - available_memory: Memory available for KV cache in bytes for each - worker. - - Returns: - The generated KVCacheConfigs for each worker. - - """ - # Merge the KV cache specs of all workers. Different PP stages may have - # different layer names, and different TP ranks of the same PP stage should - # have the same KV cache spec. - merged_kv_cache_specs: dict[str, KVCacheSpec] = {} - for kv_cache_spec_one_worker in kv_cache_specs: - for layer_name, layer_spec in kv_cache_spec_one_worker.items(): - if layer_name not in merged_kv_cache_specs: - merged_kv_cache_specs[layer_name] = layer_spec - else: - assert merged_kv_cache_specs[layer_name] == layer_spec, ( - "The KV cache specs for the same layer are different " - "across workers. This is not supported yet." - ) - - # Check if the KV cache specs are registered correctly. - # This is to prevent that some layers are initialized with unregistered specs. - KVCacheSpecRegistry.check_kv_cache_spec_registry(merged_kv_cache_specs) - - # When speculating with more than 1 speculative module (e.g. multi-layered MTP) - # tag every SlidingWindowSpec with how many extra tokens to retain in the window. - extra_retained_tokens = max(0, vllm_config.num_prefill_lookahead_tokens - 1) - for layer_name, layer_spec in merged_kv_cache_specs.items(): - if isinstance(layer_spec, SlidingWindowSpec): - merged_kv_cache_specs[layer_name] = replace( - layer_spec, extra_retained_tokens=extra_retained_tokens - ) - - # Get global KV cache groups. This also handles spec unification for - # hybrid models when disable_hybrid_kv_cache_manager is enabled. - # After this call, merged_kv_cache_specs may be modified in-place. - global_kv_cache_groups = get_kv_cache_groups(vllm_config, merged_kv_cache_specs) - - # If original_max_model_len was -1, automatically - # determine the maximum model length that fits in available GPU memory. - # We use per-worker projected groups to account for PP sharding. - projected_groups_per_worker = [ - _project_kv_cache_groups_to_worker(global_kv_cache_groups, worker_spec) - for worker_spec in kv_cache_specs - ] - - # If `num_gpu_blocks_override` is set, the cache size that will actually - # be allocated is decoupled from the profiled `available_memory`: - # `may_override_num_blocks` in `get_kv_cache_config_from_groups` clamps - # `num_blocks` to the override. Reflect that in `available_memory` here so - # auto-fit, the admission check, and the per-worker config builder all - # plan against the same effective capacity. - override = vllm_config.cache_config.num_gpu_blocks_override - if override is not None: - adjusted_memory: list[int] = [] - for groups, avail_mem in zip(projected_groups_per_worker, available_memory): - if not groups: - adjusted_memory.append(avail_mem) - continue - bytes_per_block = _pool_bytes_per_block(vllm_config, groups) - logger.info( - "Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d", - avail_mem // bytes_per_block, - override, - ) - adjusted_memory.append(override * bytes_per_block) - available_memory = adjusted_memory - - # Reserve the null block BlockPool permanently holds back, so auto-fit and - # the capacity check both plan against usable blocks. Allocation below - # still uses the full memory. - check_memory = [ - avail_mem - _pool_bytes_per_block(vllm_config, groups) if groups else avail_mem - for groups, avail_mem in zip(projected_groups_per_worker, available_memory) - ] - - if vllm_config.model_config.original_max_model_len == -1: - _auto_fit_max_model_len(vllm_config, projected_groups_per_worker, check_memory) - - # Check if the available memory is enough per worker. - for groups, avail_mem in zip(projected_groups_per_worker, check_memory): - if not groups: - continue - _check_enough_kv_cache_memory( - avail_mem, - partial(_max_memory_usage_bytes_from_groups, vllm_config, groups), - vllm_config.model_config.max_model_len, - partial(_estimate_max_model_len_from_groups, vllm_config, groups), - ) - - kv_cache_configs: list[KVCacheConfig] = [] - for projected_groups, kv_cache_spec_one_worker, available_memory_one_worker in zip( - projected_groups_per_worker, kv_cache_specs, available_memory - ): - assert sum(len(group.layer_names) for group in projected_groups) == len( - kv_cache_spec_one_worker - ), "Some layers are not assigned to any group." - kv_cache_configs.append( - get_kv_cache_config_from_groups( - vllm_config, projected_groups, available_memory_one_worker - ) - ) - - # Change the num_blocks of each rank to the smallest among all ranks. - # We also need to shrink the tensor size proportionally to avoid - # allocating unused memory. - min_num_blocks = min( - kv_cache_config.num_blocks for kv_cache_config in kv_cache_configs - ) - for i, kv_cache_config in enumerate(kv_cache_configs): - if kv_cache_config.num_blocks == min_num_blocks: - continue - # Re-plan with exactly the memory the smallest rank can afford, so - # strides and offsets stay consistent with the shrunken allocation. - groups = kv_cache_config.kv_cache_groups - kv_cache_configs[i] = get_kv_cache_config_from_groups( - vllm_config, - groups, - min_num_blocks * _pool_bytes_per_block(vllm_config, groups), - ) - - for kv_cache_config in kv_cache_configs: - kv_cache_config.kv_tp_replicas = kv_cache_groups_tp_replicas( - kv_cache_config.kv_cache_groups, - vllm_config.parallel_config.tensor_parallel_size, - vllm_config.parallel_config.decode_context_parallel_size, - ) - - return kv_cache_configs - - class BlockHashListWithBlockSize: """Convert block-hash granularity from `hash_block_size` to `target_block_size`. Used when KV cache groups have different block sizes: `hash_block_size` diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 35cc8cc1860c..d8c976ff693b 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -50,10 +50,10 @@ from vllm.utils.network_utils import make_zmq_socket from vllm.utils.system_utils import decorate_logs, set_process_title from vllm.v1.attention.backends.utils import resolve_kv_cache_layout +from vllm.v1.core.kv_cache_planning import get_kv_cache_configs from vllm.v1.core.kv_cache_utils import ( BlockHash, generate_scheduler_kv_cache_config, - get_kv_cache_configs, get_request_block_hasher, init_none_hash, resolve_cache_hit_alignment_tokens, diff --git a/vllm/v1/hisparse/layout.py b/vllm/v1/hisparse/layout.py index 96d93650ad29..0a72427d4233 100644 --- a/vllm/v1/hisparse/layout.py +++ b/vllm/v1/hisparse/layout.py @@ -71,7 +71,7 @@ def get_hisparse_kv_cache_groups( if not mla_specs: return None - from vllm.v1.core.kv_cache_utils import get_kv_cache_groups + from vllm.v1.core.kv_cache_planning import get_kv_cache_groups mla_group_spec = UniformTypeKVCacheSpecs.from_specs(mla_specs) assert mla_group_spec is not None @@ -323,7 +323,7 @@ def get_hisparse_kv_cache_config( available_memory: int, host_budget: int, ) -> KVCacheConfig: - from vllm.v1.core.kv_cache_utils import ( + from vllm.v1.core.kv_cache_planning import ( _get_kv_cache_bytes_per_block, may_override_num_blocks, validate_kv_cache_layout, diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 73be62319a64..3a6409db5546 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -989,7 +989,7 @@ def _extrapolate_full_graph_memory(mem_samples: list[int], total_graphs: int) -> def _init_minimal_kv_cache_for_profiling(runner: "GPUModelRunner") -> None: """Allocate the smallest KV cache that still lets every graph be captured.""" - from vllm.v1.core.kv_cache_utils import ( + from vllm.v1.core.kv_cache_planning import ( get_kv_cache_config_from_groups, get_kv_cache_groups, ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index e19c1b84f457..90059fd39f5b 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -6206,7 +6206,7 @@ def profile_run(self, randomize_inputs: bool = False) -> None: gc.collect() def _init_minimal_kv_cache_for_profiling(self) -> None: - from vllm.v1.core.kv_cache_utils import ( + from vllm.v1.core.kv_cache_planning import ( get_kv_cache_config_from_groups, get_kv_cache_groups, ) From c5422f4c466afd765b32b662151d0f37dad7baa5 Mon Sep 17 00:00:00 2001 From: RAHUL KUMAR <55033230+pctablet505@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:28:32 +0530 Subject: [PATCH 350/624] [Test][Renderer] Cover the reasoning -> reasoning_content mirror in chat message parsing (#56129) Signed-off-by: pctablet505 <55033230+pctablet505@users.noreply.github.com> Co-authored-by: Claude Sonnet 5.5 --- tests/renderers/test_chat_utils.py | 46 ++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/tests/renderers/test_chat_utils.py b/tests/renderers/test_chat_utils.py index d0a9b9e6c7a8..47fe49a63e6e 100644 --- a/tests/renderers/test_chat_utils.py +++ b/tests/renderers/test_chat_utils.py @@ -27,6 +27,7 @@ ChatCompletionMessageParam, ConversationMessage, _load_embeds_dict, + _parse_chat_message_content, _parse_metadata_array, _postprocess_messages, parse_chat_messages, @@ -3196,3 +3197,48 @@ def test_validate_chat_template_rejects_invalid_type(): ) as exc_info: validate_chat_template(123) assert exc_info.value.parameter == "chat_template" + + +def _parse_one(message) -> ConversationMessage: + # Plain-text content only reads these two tracker attributes. + tracker = MagicMock() + parser = tracker.create_parser.return_value + parser.model_config.enable_prompt_embeds = False + parser.mm_placeholder_storage.return_value = {} + result = _parse_chat_message_content( + message, tracker, content_format="string", interleave_strings=False + ) + assert len(result) == 1 + return result[0] + + +def test_assistant_reasoning_is_mirrored_to_reasoning_content(): + """``reasoning`` is mirrored to ``reasoning_content`` for legacy templates.""" + conv_msg = _parse_one( + { + "role": "assistant", + "content": "4", + "reasoning": "2+2 equals 4", + } + ) + + assert conv_msg.get("reasoning") == "2+2 equals 4" + assert conv_msg.get("reasoning_content") == "2+2 equals 4" + + +def test_assistant_without_reasoning_sets_neither_field(): + """No reasoning in, no reasoning keys out: no empty thinking block in templates.""" + conv_msg = _parse_one({"role": "assistant", "content": "Hi"}) + + assert "reasoning" not in conv_msg + assert "reasoning_content" not in conv_msg + + +def test_non_assistant_reasoning_is_not_mirrored(): + """Only assistant turns carry reasoning; a user turn must not.""" + conv_msg = _parse_one( + {"role": "user", "content": "Hi", "reasoning": "should be ignored"} + ) + + assert "reasoning" not in conv_msg + assert "reasoning_content" not in conv_msg From 657a2e94174fadcae7ae729b076d61001509d804 Mon Sep 17 00:00:00 2001 From: Yoni Gozlan <74535834+yonigozlan@users.noreply.github.com> Date: Mon, 12 Oct 2026 21:22:48 -0700 Subject: [PATCH 351/624] [Rust Frontend] Accept whole-number floats for integer tool arguments (#61299) Signed-off-by: Yoni Gozlan <74535834+yonigozlan@users.noreply.github.com> Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao --- rust/src/parser/src/tool/parameters.rs | 31 +++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/rust/src/parser/src/tool/parameters.rs b/rust/src/parser/src/tool/parameters.rs index 391ca90f105f..afb3d4a9f2cb 100644 --- a/rust/src/parser/src/tool/parameters.rs +++ b/rust/src/parser/src/tool/parameters.rs @@ -294,7 +294,9 @@ fn try_convert_value(param_type: &JsonParamType<'_>, input: &ParamInput) -> Opti fn try_convert_text_value(param_type: &JsonParamType<'_>, value: &str) -> Option { match param_type { JsonParamType::String => Some(Value::String(value.to_string())), - JsonParamType::Integer => convert_integer_text(value), + JsonParamType::Integer => { + convert_integer_text(value).or_else(|| convert_whole_number_text(value)) + } JsonParamType::Number => convert_number_text(value), JsonParamType::Boolean => try_convert_boolean(value), JsonParamType::Object(_) if value.trim().is_empty() => Some(Value::Object(Map::new())), @@ -401,6 +403,14 @@ pub(crate) fn convert_integer_text(value: &str) -> Option { .map(Value::Number) } +/// Convert a whole number written as a float (`3.0`) to a JSON integer, since JSON Schema +/// counts any number with a zero fractional part as an `integer`. +fn convert_whole_number_text(value: &str) -> Option { + let number = without_digit_separators(value.trim())?.parse::().ok()?; + let in_range = number >= i64::MIN as f64 && number < i64::MAX as f64; + (number.fract() == 0.0 && in_range).then(|| Value::Number(Number::from(number as i64))) +} + /// Convert raw text to a JSON number, as for a `number` parameter. pub(crate) fn convert_number_text(value: &str) -> Option { try_convert_number(&without_digit_separators(value.trim())?) @@ -591,6 +601,25 @@ mod tests { } } + #[test] + fn integers_accept_whole_number_floats() { + let params = ToolSchema::from_schema(json!({ + "type": "object", + "properties": { + "count": { "type": "integer" }, + "either": { "type": ["integer", "number"] } + } + })); + + assert_eq!(params.convert("count", text("3.0")), json!(3)); + assert_eq!(params.convert("count", text(" -2.0 ")), json!(-2)); + assert_eq!(params.convert("count", text("1_000.0")), json!(1000)); + assert_eq!(params.convert("either", text("3.5")), json!(3.5)); + for invalid in ["3.5", "inf", "NaN", "1e300"] { + assert_eq!(params.convert("count", text(invalid)), json!(invalid)); + } + } + #[test] fn nullable_schemas_accept_null() { let params = ToolSchema::from_schema(json!({ From 28d679e762a0b6ed560d767a76ebe02699bf5c46 Mon Sep 17 00:00:00 2001 From: dskwe Date: Tue, 13 Oct 2026 12:23:45 +0800 Subject: [PATCH 352/624] [Bugfix] Preserve independent tokenizer code_revision in snapshot resolution (#61272) Signed-off-by: Yuan Li Co-authored-by: Cyrus Leung --- tests/tokenizers_/test_basic.py | 85 +++++++++++++++++++++++++++++++++ vllm/tokenizers/registry.py | 13 ++--- 2 files changed, 92 insertions(+), 6 deletions(-) diff --git a/tests/tokenizers_/test_basic.py b/tests/tokenizers_/test_basic.py index a46be988e5c8..a91124549457 100644 --- a/tests/tokenizers_/test_basic.py +++ b/tests/tokenizers_/test_basic.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from pathlib import Path from typing import _get_protocol_attrs # type: ignore from unittest.mock import patch import pytest from transformers import ( + AutoTokenizer, PreTrainedTokenizerBase, TokenizersBackend, tokenization_utils_base, @@ -13,6 +15,7 @@ from vllm.tokenizers import TokenizerLike, get_tokenizer from vllm.tokenizers.hf import HfTokenizer from vllm.tokenizers.mistral import MistralTokenizer +from vllm.tokenizers.registry import resolve_tokenizer_args def _get_missing_attrs(obj: object, target: type): @@ -58,6 +61,88 @@ def test_tokenizer_revision(tokenizer_name: str): get_tokenizer(tokenizer_name, revision="never") +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize( + "code_revision, error_match", + [ + pytest.param("never", "not a valid git identifier", id="unknown-revision"), + pytest.param( + "73293ac504e43f4a895688b57c65cebe24f01a38", + r"tokenization_fast\.py", + id="revision-without-code", + ), + ], +) +def test_tokenizer_code_revision_rejects_missing_code( + code_revision: str, error_match: str +): + """A data snapshot must not shadow the requested remote-code revision.""" + with pytest.raises(OSError, match=error_match): + get_tokenizer( + "hf-internal-testing/test_dynamic_tokenizer", + tokenizer_mode="hf", + revision="08a35780b7a78794436d8654aafa26ee0165c149", + code_revision=code_revision, + trust_remote_code=True, + ) + + +@pytest.mark.skip_global_cleanup +def test_tokenizer_code_revision_distinct_from_revision(): + """Distinct code/data revisions: the resolver must keep the Hub ID. + + Dynamic code is read from a local directory as-is, so resolving to a + data-revision snapshot would silently shadow an independent code + revision (the same contract already enforced for processors). + """ + revision = "08a35780b7a78794436d8654aafa26ee0165c149" + repo = "hf-internal-testing/test_dynamic_tokenizer" + _, name, _, kwargs = resolve_tokenizer_args( + repo, + tokenizer_mode="hf", + revision=revision, + code_revision=revision, + trust_remote_code=True, + ) + assert Path(name).is_dir() + assert kwargs["code_revision"] == revision + + _, name, _, kwargs = resolve_tokenizer_args( + repo, + tokenizer_mode="hf", + revision=revision, + code_revision="73293ac504e43f4a895688b57c65cebe24f01a38", + trust_remote_code=True, + ) + assert name == repo + assert kwargs["code_revision"] == "73293ac504e43f4a895688b57c65cebe24f01a38" + + +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize( + "code_revision", [None, "08a35780b7a78794436d8654aafa26ee0165c149"] +) +def test_tokenizer_code_revision_default_or_matching(code_revision: str | None): + revision = "08a35780b7a78794436d8654aafa26ee0165c149" + tokenizer = get_tokenizer( + "hf-internal-testing/test_dynamic_tokenizer", + tokenizer_mode="hf", + revision=revision, + code_revision=code_revision, + trust_remote_code=True, + ) + reference = AutoTokenizer.from_pretrained( + "hf-internal-testing/test_dynamic_tokenizer", + revision=revision, + code_revision=code_revision, + trust_remote_code=True, + ) + assert isinstance(tokenizer, PreTrainedTokenizerBase) + assert tokenizer.encode("hello world") == reference.encode("hello world") + assert tokenizer.name_or_path == "hf-internal-testing/test_dynamic_tokenizer" + _assert_tokenizer_like(tokenizer) + + def test_hub_tokenizer_loads_from_local_snapshot(): """Transformers makes a Hub request per file when given a repo ID.""" with patch.object( diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index a7adf4597f6a..8b1863023c99 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -165,12 +165,13 @@ def resolve_tokenizer_args( "mistral_format=True is not supported with tokenizer_mode='hf'" ) kwargs["mistral_format"] = False - tokenizer_name = get_non_weight_snapshot_path( - str(tokenizer_name), - revision=revision, - cache_dir=download_dir, - token=kwargs.get("token"), - ) + if kwargs.get("code_revision") in (None, revision): + tokenizer_name = get_non_weight_snapshot_path( + str(tokenizer_name), + revision=revision, + cache_dir=download_dir, + token=kwargs.get("token"), + ) return tokenizer_mode, tokenizer_name, args, kwargs From 69d00bc2645cac2144096be3ef1a12a73f95f366 Mon Sep 17 00:00:00 2001 From: Yoni Gozlan <74535834+yonigozlan@users.noreply.github.com> Date: Mon, 12 Oct 2026 21:55:20 -0700 Subject: [PATCH 353/624] [Bugfix] Accept whole-number floats for integer tool arguments in the response parser (#61297) Signed-off-by: Yoni Gozlan <74535834+yonigozlan@users.noreply.github.com> Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao --- tests/parser/test_chat_parsing.py | 8 ++++++++ vllm/parser/chat_parsing/response_parser.py | 9 ++++++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/tests/parser/test_chat_parsing.py b/tests/parser/test_chat_parsing.py index 679d97f15b20..e3b8dba17220 100644 --- a/tests/parser/test_chat_parsing.py +++ b/tests/parser/test_chat_parsing.py @@ -1555,6 +1555,14 @@ def test_coerce_tool_calls_casts_declared_types(self): def test_coerce_falls_back_to_raw_on_failure(self): self.assertEqual(_coerce("not-a-number", ("integer",)), "not-a-number") + def test_coerce_integer_accepts_whole_number_floats(self): + self.assertEqual(_coerce("3.0", ("integer",)), 3) + self.assertIsInstance(_coerce("3.0", ("integer",)), int) + self.assertEqual(_coerce("3.5", ("integer",)), "3.5") + self.assertEqual(_coerce("3.5", ("integer", "number")), 3.5) + self.assertEqual(_coerce("inf", ("integer",)), "inf") + self.assertEqual(_coerce("NaN", ("integer",)), "NaN") + def test_coerce_handles_any_of_and_null(self): self.assertEqual(_coerce("5", ("integer", "null")), 5) self.assertIsNone(_coerce("null", ("integer", "null"))) diff --git a/vllm/parser/chat_parsing/response_parser.py b/vllm/parser/chat_parsing/response_parser.py index 556df803a70a..0416fe0fe3ea 100644 --- a/vllm/parser/chat_parsing/response_parser.py +++ b/vllm/parser/chat_parsing/response_parser.py @@ -50,7 +50,14 @@ def _coerce(raw: str, types: tuple[str, ...]) -> Any: for type_name in types: try: if type_name == "integer": - return int(raw) + try: + return int(raw) + except ValueError: + # JSON Schema integers include whole numbers written as floats, e.g. "3.0" + number = float(raw) + if not number.is_integer(): + continue + return int(number) if type_name == "number": number = float(raw) if number != number or number in (float("inf"), float("-inf")): From 5555a2d8f7116b2b4877f66516f4128176cf722a Mon Sep 17 00:00:00 2001 From: Mook <68294499+Godmook@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:02:12 +0900 Subject: [PATCH 354/624] [Bugfix][Model] Fix EXAONE-4.5 image inputs: ViT received un-normalised pixels (#61104) Signed-off-by: godmook Co-authored-by: cjackal <44624812+cjackal@users.noreply.github.com> --- vllm/model_executor/models/exaone4_5.py | 6 ++++++ vllm/model_executor/models/exaone4_5_mtp.py | 1 + 2 files changed, 7 insertions(+) diff --git a/vllm/model_executor/models/exaone4_5.py b/vllm/model_executor/models/exaone4_5.py index 5a53e60b7d5e..ad4049c8eeb0 100644 --- a/vllm/model_executor/models/exaone4_5.py +++ b/vllm/model_executor/models/exaone4_5.py @@ -36,6 +36,7 @@ from vllm.distributed import utils as dist_utils from vllm.logger import init_logger from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention +from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear from vllm.model_executor.layers.quantization import QuantizationConfig @@ -272,6 +273,7 @@ def __init__( vision_config: Exaone4_5_VisionConfig, norm_eps: float = 1e-6, quant_config: QuantizationConfig | None = None, + input_norm: nn.Module | None = None, prefix: str = "", use_data_parallel: bool = False, ) -> None: @@ -279,6 +281,7 @@ def __init__( vision_config=vision_config, norm_eps=norm_eps, quant_config=quant_config, + input_norm=input_norm, prefix=prefix, ) depth = vision_config.depth @@ -335,7 +338,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" self.config = config + self.model_config = vllm_config.model_config self.multimodal_config = multimodal_config + self.video_pruning_rate = multimodal_config.video_pruning_rate self.is_multimodal_pruning_enabled = ( multimodal_config.is_multimodal_pruning_enabled() ) @@ -345,6 +350,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config.vision_config, norm_eps=getattr(config, "rms_norm_eps", 1e-6), quant_config=self.quant_config, + input_norm=build_mm_input_norm(self.model_config), prefix=maybe_prefix(prefix, "visual"), use_data_parallel=self.use_data_parallel, ) diff --git a/vllm/model_executor/models/exaone4_5_mtp.py b/vllm/model_executor/models/exaone4_5_mtp.py index d9c7960010a5..41235929867a 100644 --- a/vllm/model_executor/models/exaone4_5_mtp.py +++ b/vllm/model_executor/models/exaone4_5_mtp.py @@ -75,6 +75,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.layers = nn.ModuleList( Exaone4DecoderLayer( text_config, + cache_config=vllm_config.cache_config, quant_config=quant_config, prefix=f"{prefix}.layers.{idx}", ) From b337523c118d1cdb211732199f0a7d5b1e942bb9 Mon Sep 17 00:00:00 2001 From: Chaojun Zhang Date: Tue, 13 Oct 2026 13:12:10 +0800 Subject: [PATCH 355/624] [XPU] [Test] Extend spawn-safe sitecustomize monkeypatch to XPU (#57911) Signed-off-by: Chaojun Zhang Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: xiangdong <40376367+zxd1997066@users.noreply.github.com> --- .buildkite/intel_jobs/distributed_intel.yaml | 33 ++++++++++++++++++++ tests/v1/shutdown/conftest.py | 15 +++++++-- tests/v1/shutdown/test_forward_error.py | 8 ++--- tests/v1/shutdown/test_startup_error.py | 8 ++--- 4 files changed, 53 insertions(+), 11 deletions(-) create mode 100644 .buildkite/intel_jobs/distributed_intel.yaml diff --git a/.buildkite/intel_jobs/distributed_intel.yaml b/.buildkite/intel_jobs/distributed_intel.yaml new file mode 100644 index 000000000000..54800b4103d2 --- /dev/null +++ b/.buildkite/intel_jobs/distributed_intel.yaml @@ -0,0 +1,33 @@ +group: Distributed Intel +depends_on: + - image-build-xpu +steps: +- label: ":intel: (BMG) Distributed Torchrun + Shutdown" + key: distributed-torchrun-shutdown-tests-2-gpus + timeout_in_minutes: 50 + device: intel_gpu + agent_tags: + label: production + gpu: 2+ + mem: 24+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/distributed/ + - "!vllm/distributed/kv_transfer/" + - vllm/engine/ + - vllm/executor/ + - vllm/worker/worker_base.py + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/distributed/ + - tests/v1/shutdown + - tests/v1/worker/test_worker_memory_snapshot.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && pytest -v -s v1/shutdown' diff --git a/tests/v1/shutdown/conftest.py b/tests/v1/shutdown/conftest.py index b276d0fab267..ef54b188fefa 100644 --- a/tests/v1/shutdown/conftest.py +++ b/tests/v1/shutdown/conftest.py @@ -10,9 +10,18 @@ @pytest.fixture -def rocm_sitecustomize_factory(monkeypatch, tmp_path: Path): - """Return a function that installs a given sitecustomize payload.""" - if not current_platform.is_rocm(): +def spawn_sitecustomize_factory(monkeypatch, tmp_path: Path): + """Return a function that installs a given sitecustomize payload. + + Both ROCm and XPU force the `spawn` multiprocessing start method (see + `vllm.platforms.rocm`/`vllm.platforms.xpu`), which re-imports modules in + child processes instead of inheriting them via `fork`'s copy-on-write + memory. A plain `monkeypatch.setattr` in the parent test process is + therefore invisible to spawned EngineCore/worker subprocesses on these + platforms, so the payload must be installed via `sitecustomize.py` on + `PYTHONPATH` instead. + """ + if not (current_platform.is_rocm() or current_platform.is_xpu()): return lambda _: None def install(lines: Iterable[str]) -> None: diff --git a/tests/v1/shutdown/test_forward_error.py b/tests/v1/shutdown/test_forward_error.py index 8bc09a64cac9..c40a41cc3b41 100644 --- a/tests/v1/shutdown/test_forward_error.py +++ b/tests/v1/shutdown/test_forward_error.py @@ -41,21 +41,21 @@ def evil_forward(self, *args, **kwargs): @pytest.fixture -def rocm_evil_forward(rocm_sitecustomize_factory): +def evil_forward_fixture(spawn_sitecustomize_factory): lines = [ "from vllm.distributed import get_tensor_model_parallel_rank", "from vllm.model_executor.models.llama import LlamaForCausalLM", inspect.getsource(evil_forward), f"LlamaForCausalLM.forward = {evil_forward.__name__}", ] - rocm_sitecustomize_factory(lines) + spawn_sitecustomize_factory(lines) @pytest.mark.asyncio @pytest.mark.parametrize("tensor_parallel_size", [2, 1]) @pytest.mark.parametrize("model", MODELS) async def test_async_llm_model_error( - monkeypatch, rocm_evil_forward, tensor_parallel_size: int, model: str + monkeypatch, evil_forward_fixture, tensor_parallel_size: int, model: str ) -> None: """Test that AsyncLLM propagates a forward pass error and frees memory. @@ -118,7 +118,7 @@ async def generate(request_id: str): @pytest.mark.parametrize("model", MODELS) def test_llm_model_error( monkeypatch, - rocm_evil_forward, + evil_forward_fixture, tensor_parallel_size: int, enable_multiprocessing: bool, model: str, diff --git a/tests/v1/shutdown/test_startup_error.py b/tests/v1/shutdown/test_startup_error.py index 6d6d0d91e77f..d70f69409ee5 100644 --- a/tests/v1/shutdown/test_startup_error.py +++ b/tests/v1/shutdown/test_startup_error.py @@ -30,7 +30,7 @@ def evil_method(self, *args, **kwargs): @pytest.fixture -def rocm_evil_method(rocm_sitecustomize_factory, request): +def evil_method_fixture(spawn_sitecustomize_factory, request): failing_method = request.getfixturevalue("failing_method") lines = [ "from vllm.distributed import get_tensor_model_parallel_rank", @@ -38,7 +38,7 @@ def rocm_evil_method(rocm_sitecustomize_factory, request): inspect.getsource(evil_method), f"LlamaForCausalLM.{failing_method} = {evil_method.__name__}", ] - rocm_sitecustomize_factory(lines) + spawn_sitecustomize_factory(lines) @pytest.mark.timeout(SHUTDOWN_TEST_TIMEOUT_SEC) @@ -47,7 +47,7 @@ def rocm_evil_method(rocm_sitecustomize_factory, request): @pytest.mark.parametrize("failing_method", ["forward", "load_weights"]) def test_async_llm_startup_error( monkeypatch, - rocm_evil_method, + evil_method_fixture, model: str, tensor_parallel_size: int, failing_method: str, @@ -84,7 +84,7 @@ def test_async_llm_startup_error( @pytest.mark.parametrize("failing_method", ["forward", "load_weights"]) def test_llm_startup_error( monkeypatch, - rocm_evil_method, + evil_method_fixture, model: str, tensor_parallel_size: int, enable_multiprocessing: bool, From f7ba007b5e766128468fed38c5da9777b521b8ee Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Tue, 13 Oct 2026 01:13:33 -0400 Subject: [PATCH 356/624] [ROCm][MoE] Split AiterExperts by weight-quantization key (#60322) Signed-off-by: Divakar Verma Co-authored-by: Claude Sonnet 5 --- .buildkite/test-amd.yaml | 27 +++++- docs/design/moe_kernel_features.md | 2 +- .../moe/modular_kernel_tools/common.py | 8 +- .../moe/modular_kernel_tools/mk_objects.py | 40 +++++--- .../moe/test_modular_kernel_combinations.py | 44 ++++++--- tests/kernels/moe/test_mori_moe.py | 8 +- tests/kernels/moe/test_rocm_aiter_moe.py | 4 +- tests/kernels/moe/test_topk_softplus_sqrt.py | 2 +- .../layers/fused_moe/__init__.py | 12 ++- .../fused_moe/experts/rocm_aiter_moe.py | 97 ++++++++++++++----- .../layers/fused_moe/oracle/fp8.py | 6 +- .../layers/fused_moe/oracle/mxfp4.py | 8 +- .../layers/fused_moe/oracle/unquantized.py | 4 +- 13 files changed, 198 insertions(+), 64 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 5145f5a272ff..d815ee9268b2 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1528,7 +1528,32 @@ steps: commands: - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --ignore=kernels/moe/test_modular_kernel_combinations.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s kernels/moe/test_modular_kernel_combinations.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s kernels/moe/test_modular_kernel_combinations.py -k "not multigpu" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + +- label: ":amd: (MI300) MoE Kernels Multi-GPU" # TBD + timeout_in_minutes: 30 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/quantization/cutlass_w8a8/moe/ + - csrc/moe/ + - csrc/rocm/ + - tests/kernels/moe/test_modular_kernel_combinations.py + - tests/kernels/moe/modular_kernel_tools/ + - tests/kernels/utils.py + - vllm/model_executor/layers/fused_moe/ + - vllm/distributed/device_communicators/ + - vllm/distributed/parallel_state.py + - vllm/envs.py + - vllm/config + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s kernels/moe/test_modular_kernel_combinations.py -k "multigpu" - label: ":amd: (MI300) vLLM IR" # TBD timeout_in_minutes: 50 diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 6c71aa1f46d4..2528bc9ea459 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -91,7 +91,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k | marlin | standard,
batched | 3 / N/A | 3 / N/A | silu,
swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.fused_marlin_moe],
[`MarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.MarlinExperts],
[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.BatchedMarlinExperts] | | trtllm | standard | mxfp4,
nvfp4 | G(16),G(32) | 5 | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],
[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],
[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],
[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] | | hpc | standard | fp8 | G(128),T | silu | Y | Y | [`HPCExperts`][vllm.model_executor.layers.fused_moe.hpc_moe.HPCExperts] | -| rocm aiter moe | standard | mxfp4,
fp8 | G(32),G(128),A,T | silu, gelu,
swigluoai | Y | N | `rocm_aiter_fused_experts`,
`AiterExperts` | +| rocm aiter moe | standard | none,
mxfp4,
fp8 | G(32),G(128),A,T | silu, gelu,
swigluoai,
situ | Y | Y | [`rocm_aiter_fused_experts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.rocm_aiter_fused_experts],
[`AiterUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterUnquantizedExperts],
[`AiterFp8BlockExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterFp8BlockExperts],
[`AiterFp8TensorExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterFp8TensorExperts],
[`AiterFp8ChannelExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterFp8ChannelExperts],
[`AiterMxfp4Experts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterMxfp4Experts] | | cpu_moe | standard | N/A | N/A | silu, gelu,
gelu_tanh,
swigluoai | Y | N | [`X86CPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.X86CPUUnquantizedExperts],
[`ArmCPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.ArmCPUUnquantizedExperts],
[`CPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.CPUUnquantizedExperts] | | naive batched4 | batched | int8,
fp8 | G,A,T | silu, gelu | 6 | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.experts.fused_batched_moe.NaiveBatchedExperts] | diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 5f6349a32880..1f79b62bf8c8 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -167,6 +167,9 @@ def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]: vllm_config.model_config = SimpleNamespace( enforce_eager=True, is_moe=True, + enable_sleep_mode=False, + sleep_mode_backend="cumem", + sleep_mode_offload_cudagraph=False, ) vllm_config.parallel_config.data_parallel_size = self.world_size vllm_config.parallel_config.enable_expert_parallel = True @@ -722,6 +725,9 @@ def _maybe_convert_weights_for_experts( rank_weights: WeightTensors, ) -> WeightTensors: """Convert weights to expert-specific format (e.g., TrtLLM BlockMajorK).""" + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( + AiterExperts, + ) from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( Fp8MoeBackend, convert_to_fp8_moe_kernel_format, @@ -734,7 +740,7 @@ def _maybe_convert_weights_for_experts( # intermediate sizes that are not a multiple of 256, but only for the fp8 per-tensor # and per-token schemes. Those keep JIT-compiling until # https://github.com/ROCm/aiter/issues/5766 is fixed. - if fe_name == "AiterExperts" and ( + if issubclass(fe_type, AiterExperts) and ( config.quant_dtype is None or config.quant_block_shape is not None ): return _shuffle_weights_for_aiter(rank_weights) diff --git a/tests/kernels/moe/modular_kernel_tools/mk_objects.py b/tests/kernels/moe/modular_kernel_tools/mk_objects.py index 6ca644c5fce1..56119ffddeb7 100644 --- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py +++ b/tests/kernels/moe/modular_kernel_tools/mk_objects.py @@ -93,6 +93,7 @@ class ExpertInfo: ] common_float_and_int_types = common_float_types + [torch.int8] nvfp4_types = ["nvfp4"] +mxfp4_types = ["mxfp4"] fp8_types = [current_platform.fp8_dtype()] @@ -318,20 +319,37 @@ def expert_info(kind) -> ExpertInfo: if has_aiter(): from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterFp8BlockExperts, + AiterFp8ChannelExperts, + AiterFp8TensorExperts, + AiterMxfp4Experts, + AiterUnquantizedExperts, ) - register_experts( - AiterExperts, - standard_format, - # AiterExperts also supports the fully-unquantized (None, None) - # scheme (see SUPPORTED_W_A in rocm_aiter_moe.py), not just fp8. - common_float_types, - blocked_quantization_support=True, - needs_aiter=True, - ) + # Registration dtypes gate which quant configs the modular-kernel test + # harness will ever pair with a given class (see Config.is_valid()): + # they must match what each subclass's _supports_quant_scheme() actually + # declares, or the harness will generate a case the class rejects. + for aiter_experts_cls, aiter_supported_dtypes in ( + (AiterUnquantizedExperts, common_float_types), + (AiterFp8BlockExperts, fp8_types), + (AiterFp8TensorExperts, fp8_types), + (AiterFp8ChannelExperts, fp8_types), + (AiterMxfp4Experts, mxfp4_types), + ): + register_experts( + aiter_experts_cls, + standard_format, + aiter_supported_dtypes, + blocked_quantization_support=True, + needs_aiter=True, + ) else: - AiterExperts = None + AiterUnquantizedExperts = None + AiterFp8BlockExperts = None + AiterFp8TensorExperts = None + AiterFp8ChannelExperts = None + AiterMxfp4Experts = None if has_deep_gemm() and is_deep_gemm_supported(): register_experts( diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index f29c22085240..da643befecd9 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -16,7 +16,13 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import AiterExperts +from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( + AiterExperts, + AiterFp8BlockExperts, + AiterFp8ChannelExperts, + AiterFp8TensorExperts, + AiterUnquantizedExperts, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8DynamicTensorSym, kFp8StaticTensorSym, @@ -80,10 +86,10 @@ def assert_aiter_quant_scheme_case(config: Config) -> None: indirectly through the general quant-config sweep. See https://github.com/vllm-project/vllm/issues/54966.""" fe_cls = config.fused_experts_type - if fe_cls is not AiterExperts: + if not issubclass(fe_cls, AiterExperts): return - if config.quant_config is None: + if config.quant_config is None or config.quant_dtype is None: w_key, a_key = None, None else: w_key, a_key = config.fp8_quant_key_pair() @@ -101,7 +107,7 @@ def assert_aiter_activation_case(config: Config) -> None: AiterExperts being reached only indirectly through the general activation sweep. See https://github.com/vllm-project/vllm/issues/54966.""" fe_cls = config.fused_experts_type - if fe_cls is not AiterExperts: + if not issubclass(fe_cls, AiterExperts): return assert fe_cls._supports_activation(config.activation), ( @@ -156,7 +162,7 @@ def rank_worker( if ( topk == 1 and config.supports_apply_weight_on_input() - and config.fused_experts_type is AiterExperts + and issubclass(config.fused_experts_type, AiterExperts) and config.quant_block_shape is not None ): print( @@ -172,7 +178,7 @@ def rank_worker( # https://github.com/vllm-project/vllm/issues/57029 if ( config.world_size > 1 - and config.fused_experts_type is AiterExperts + and issubclass(config.fused_experts_type, AiterExperts) and getattr(config.prepare_finalize_type, "__name__", "") in ("DeepEPHTPrepareAndFinalize", "MoriPrepareAndFinalize") ): @@ -210,7 +216,7 @@ def rank_worker( is_aiter_fp8 = ( _cp.is_rocm() - and config.fused_experts_type is AiterExperts + and issubclass(config.fused_experts_type, AiterExperts) and config.quant_config is not None ) if is_aiter_fp8: @@ -270,7 +276,7 @@ def is_nyi_config(config: Config) -> bool: return True if config.activation != MoEActivation.SILU: - if config.fused_experts_type is not AiterExperts: + if not issubclass(config.fused_experts_type, AiterExperts): return True # AITER-only for this axis, for now if config.quant_dtype is not None: return True # unquantized-only for this axis, for now @@ -591,7 +597,7 @@ def test_aiter_moe_sorting_backend_dispatch_env_matrix( dtype=torch.bfloat16, quant_config=None, prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular, - fused_experts_type=AiterExperts, + fused_experts_type=AiterUnquantizedExperts, world_size=1, ) assert config.is_valid()[0] @@ -671,7 +677,7 @@ def test_aiter_moe_dispatch_policy_forwarded_through_apply(dispatch_policy: int) dtype=torch.bfloat16, quant_config=None, prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular, - fused_experts_type=AiterExperts, + fused_experts_type=AiterUnquantizedExperts, world_size=1, ) assert config.is_valid()[0] @@ -738,6 +744,20 @@ def test_aiter_moe_dispatch_policy_forwarded_through_apply(dispatch_policy: int) ] +def _padding_matrix_fe_cls( + quant_config: TestMoEQuantConfig | None, +) -> type[AiterExperts]: + """Map a _PADDING_QUANT_CONFIGS entry to the AiterExperts leaf class that + supports it (mirrors Config.fp8_quant_key_pair's scheme derivation).""" + if quant_config is None: + return AiterUnquantizedExperts + if quant_config.block_shape is not None: + return AiterFp8BlockExperts + if quant_config.per_out_ch_quant: + return AiterFp8ChannelExperts + return AiterFp8TensorExperts + + def _slice_gate_up_rows(t: torch.Tensor, real_per_half: int) -> torch.Tensor: """Slice a (E, 2*padded_per_half, ...) tensor's dim=1 down to the real gate/up halves of `real_per_half` size each.""" @@ -950,7 +970,7 @@ def test_aiter_moe_padding_matrix(mode: str, quant_config: TestMoEQuantConfig | dtype=torch.bfloat16, quant_config=quant_config, prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular, - fused_experts_type=AiterExperts, + fused_experts_type=_padding_matrix_fe_cls(quant_config), world_size=1, hidden_dim_unpadded=_PADDING_K_UNPADDED if pad_hidden else None, intermediate_size_per_partition_unpadded=( @@ -1113,7 +1133,7 @@ def test_aiter_moe_token_padding_garbage_rows( dtype=torch.bfloat16, quant_config=quant_config, prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular, - fused_experts_type=AiterExperts, + fused_experts_type=_padding_matrix_fe_cls(quant_config), world_size=1, ) assert config.is_valid()[0] diff --git a/tests/kernels/moe/test_mori_moe.py b/tests/kernels/moe/test_mori_moe.py index 9869b6bc5467..f06afc8bd63e 100644 --- a/tests/kernels/moe/test_mori_moe.py +++ b/tests/kernels/moe/test_mori_moe.py @@ -42,7 +42,10 @@ def _mori_worker(pgi, vllm_config, cpu_group, m, n, k, fp8, graph): from vllm._aiter_ops import rocm_aiter_ops - from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import AiterExperts + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( + AiterFp8ChannelExperts, + AiterUnquantizedExperts, + ) from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import ( MoriPrepareAndFinalize, ) @@ -94,7 +97,8 @@ def _mori_worker(pgi, vllm_config, cpu_group, m, n, k, fp8, graph): assert isinstance(prepare, MoriPrepareAndFinalize) assert prepare.use_fp8_dispatch == fp8 assert prepare.num_dispatchers() == 2 - experts = AiterExperts(config, quant) + experts_cls = AiterFp8ChannelExperts if fp8 else AiterUnquantizedExperts + experts = experts_cls(config, quant) assert not experts.expects_unquantized_inputs kernel = FusedMoEKernel(prepare, experts) local_w1, local_w2 = rocm_aiter_ops.shuffle_weights( diff --git a/tests/kernels/moe/test_rocm_aiter_moe.py b/tests/kernels/moe/test_rocm_aiter_moe.py index abd46753515e..af2f7beb1bd5 100644 --- a/tests/kernels/moe/test_rocm_aiter_moe.py +++ b/tests/kernels/moe/test_rocm_aiter_moe.py @@ -848,13 +848,13 @@ def test_activation_method_enum_values(): def test_aiter_mxfp4_quant_scheme_support_matches_gfx950(): """AITER MXFP4 MoE support should stay gfx950-only.""" from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterMxfp4Experts, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( kMxfp4Static, ) - assert AiterExperts._supports_quant_scheme(kMxfp4Static, None) is on_gfx950() + assert AiterMxfp4Experts._supports_quant_scheme(kMxfp4Static, None) is on_gfx950() @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py index bc438c66931b..a6f7f3080583 100644 --- a/tests/kernels/moe/test_topk_softplus_sqrt.py +++ b/tests/kernels/moe/test_topk_softplus_sqrt.py @@ -87,7 +87,7 @@ def _torch_topk_softplus_sqrt( "backend,expected_experts", [ ("aiter_triton_mxfp4_bf16", "AiterW4A16ExpertsMonolithic"), - ("aiter", "AiterExperts"), + ("aiter", "AiterMxfp4Experts"), ("triton_unfused", "UnfusedOAITritonExperts"), ], ) diff --git a/vllm/model_executor/layers/fused_moe/__init__.py b/vllm/model_executor/layers/fused_moe/__init__.py index 120fbd152668..e68310f98d0d 100644 --- a/vllm/model_executor/layers/fused_moe/__init__.py +++ b/vllm/model_executor/layers/fused_moe/__init__.py @@ -106,7 +106,11 @@ def get_config() -> dict[str, Any] | None: BatchedTritonExperts, ) from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterFp8BlockExperts, + AiterFp8ChannelExperts, + AiterFp8TensorExperts, + AiterMxfp4Experts, + AiterUnquantizedExperts, ) from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import ( TritonOrDeepGemmExperts, @@ -132,7 +136,11 @@ def get_config() -> dict[str, Any] | None: ) __all__ += [ - "AiterExperts", + "AiterUnquantizedExperts", + "AiterFp8BlockExperts", + "AiterFp8TensorExperts", + "AiterFp8ChannelExperts", + "AiterMxfp4Experts", "fused_topk", "fused_experts", "get_config_file_name", diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index 07eddbf4a092..5c506eae123a 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -482,29 +482,6 @@ def _supports_current_device() -> bool: def _supports_no_act_and_mul() -> bool: return False - @staticmethod - def _supports_quant_scheme( - weight_key: QuantKey | None, - activation_key: QuantKey | None, - ) -> bool: - SUPPORTED_W_A = [ - (None, None), - (kFp8Static128BlockSym, kFp8Dynamic128Sym), - (kFp8StaticTensorSym, kFp8StaticTensorSym), - (kFp8StaticTensorSym, kFp8DynamicTensorSym), - (kFp8StaticChannelSym, kFp8DynamicTokenSym), - (kMxfp4Static, None), - (kMxfp4Static, kMxfp4Dynamic), - ] - if (weight_key, activation_key) not in SUPPORTED_W_A: - return False - if weight_key == kMxfp4Static: - from vllm.platforms.rocm import on_gfx950, on_gfx1250 - - if not on_gfx950() or on_gfx1250(): - return False - return True - @staticmethod def _supports_activation(activation: MoEActivation) -> bool: return activation in [ @@ -597,3 +574,77 @@ def apply( output.set_(result) else: output.copy_(result) + + +class AiterUnquantizedExperts(AiterExperts): + """Unquantized (bf16/fp16) weights and activations.""" + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [(None, None)] + return (weight_key, activation_key) in SUPPORTED_W_A + + +class AiterFp8BlockExperts(AiterExperts): + """128x128 block-quantized weight scale, 128-grouped dynamic + activation scale.""" + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [(kFp8Static128BlockSym, kFp8Dynamic128Sym)] + return (weight_key, activation_key) in SUPPORTED_W_A + + +class AiterFp8TensorExperts(AiterExperts): + """Per-tensor static weight scale, either per-tensor static or + per-tensor dynamic activation scale.""" + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [ + (kFp8StaticTensorSym, kFp8StaticTensorSym), + (kFp8StaticTensorSym, kFp8DynamicTensorSym), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + +class AiterFp8ChannelExperts(AiterExperts): + """Per-output-channel static weight scale, per-token dynamic + activation scale.""" + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [(kFp8StaticChannelSym, kFp8DynamicTokenSym)] + return (weight_key, activation_key) in SUPPORTED_W_A + + +class AiterMxfp4Experts(AiterExperts): + """MXFP4-quantized weights, with either unquantized (bf16/fp16) or + MXFP4-quantized dynamic activations. gfx950-only (not gfx1250).""" + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [ + (kMxfp4Static, None), + (kMxfp4Static, kMxfp4Dynamic), + ] + if (weight_key, activation_key) not in SUPPORTED_W_A: + return False + from vllm.platforms.rocm import on_gfx950, on_gfx1250 + + return on_gfx950() and not on_gfx1250() diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 446825c59047..074acd2a0942 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -211,10 +211,12 @@ def backend_to_kernel_cls( elif backend == Fp8MoeBackend.AITER: from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterFp8BlockExperts, + AiterFp8ChannelExperts, + AiterFp8TensorExperts, ) - return [AiterExperts] + return [AiterFp8BlockExperts, AiterFp8TensorExperts, AiterFp8ChannelExperts] elif backend == Fp8MoeBackend.VLLM_CUTLASS: from vllm.model_executor.layers.fused_moe.experts.triton_cutlass_moe import ( diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index cbbc39b94310..88e80c755d1f 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -284,10 +284,10 @@ def backend_to_kernel_cls( elif backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterMxfp4Experts, ) - return [AiterExperts] + return [AiterMxfp4Experts] elif backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16: from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a16_moe import ( @@ -305,10 +305,10 @@ def backend_to_kernel_cls( elif backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4: from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterMxfp4Experts, ) - return [AiterExperts] + return [AiterMxfp4Experts] elif backend == Mxfp4MoeBackend.XPU: from vllm.model_executor.layers.fused_moe.experts.xpu_moe import XPUExpertsMxFp4 diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 25b55670d1f1..1af0450dfd13 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -126,10 +126,10 @@ def backend_to_kernel_cls( elif backend == UnquantizedMoeBackend.AITER: from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterUnquantizedExperts, ) - return [AiterExperts] + return [AiterUnquantizedExperts] elif backend == UnquantizedMoeBackend.TRITON: from vllm.model_executor.layers.fused_moe.experts.triton_moe import ( From 2ce8012fa0c7fbb37335545325e07dacd8141023 Mon Sep 17 00:00:00 2001 From: Sriram Kumar Date: Tue, 13 Oct 2026 08:21:52 +0300 Subject: [PATCH 357/624] [ROCm][Perf] Skip cleaning sparse prefill MQA logits (#51314) Signed-off-by: Sriram Kumar Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../attention/test_rocm_aiter_mla_sparse.py | 54 +++++++++++++++++++ .../v1/attention/ops/rocm_aiter_mla_sparse.py | 4 +- 2 files changed, 57 insertions(+), 1 deletion(-) create mode 100644 tests/v1/attention/test_rocm_aiter_mla_sparse.py diff --git a/tests/v1/attention/test_rocm_aiter_mla_sparse.py b/tests/v1/attention/test_rocm_aiter_mla_sparse.py new file mode 100644 index 000000000000..1d28b57bebec --- /dev/null +++ b/tests/v1/attention/test_rocm_aiter_mla_sparse.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm._aiter_ops import is_aiter_found_and_supported +from vllm.platforms import current_platform + +if not current_platform.is_rocm() or not is_aiter_found_and_supported(): + pytest.skip( + "ROCm AITER sparse indexer test requires ROCm with AITER.", + allow_module_level=True, + ) + +from vllm._aiter_ops import rocm_aiter_ops +from vllm.v1.attention.ops import rocm_aiter_mla_sparse as sparse_mod + + +def test_rocm_fp8_mqa_logits_disables_clean_logits(monkeypatch): + captured = {} + expected = torch.empty((1, 1)) + + def fake_fp8_mqa_logits(*args, clean_logits): + captured["clean_logits"] = clean_logits + return expected + + monkeypatch.setattr(sparse_mod, "_ON_GFX942", False) + monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True) + monkeypatch.setattr( + sparse_mod, + "mqa_logits_module", + lambda: SimpleNamespace(fp8_mqa_logits=fake_fp8_mqa_logits), + ) + + q = torch.empty((1, 1, 1)) + k_fp8 = torch.empty((1, 1)) + scale = torch.empty((1,)) + weights = torch.empty((1, 1)) + cu_seqlen_ks = torch.zeros((1,), dtype=torch.int32) + cu_seqlen_ke = torch.ones((1,), dtype=torch.int32) + + result = sparse_mod.rocm_fp8_mqa_logits( + q, + (k_fp8, scale), + weights, + cu_seqlen_ks, + cu_seqlen_ke, + ) + + assert result is expected + assert captured["clean_logits"] is False diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 6a93d2151127..78ce85f6ae97 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -933,7 +933,9 @@ def rocm_fp8_mqa_logits( if aiter_mqa_logits_module is not None: fp8_mqa_logits = aiter_mqa_logits_module.fp8_mqa_logits - return fp8_mqa_logits(q, k_fp8, scale, weights, cu_seqlen_ks, cu_seqlen_ke) + return fp8_mqa_logits( + q, k_fp8, scale, weights, cu_seqlen_ks, cu_seqlen_ke, clean_logits=False + ) else: return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke) From 0d2689e0dfd5e9f74f361cab7701df0feee33ef3 Mon Sep 17 00:00:00 2001 From: Sriram Kumar Date: Tue, 13 Oct 2026 08:25:56 +0300 Subject: [PATCH 358/624] [ROCm][Perf] Skip redundant sparse index remap on non-indexer layers (#51309) Signed-off-by: Sriram Kumar Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: Andreas Karatzas Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> --- .../attention/test_rocm_aiter_mla_sink.py | 2 + .../test_rocm_aiter_mla_sparse_index_remap.py | 295 ++++++++++++++++++ ...est_rocm_aiter_mla_sparse_metadata_sync.py | 4 + .../test_rocm_aiter_triton_sparse_mla.py | 2 + .../backends/mla/rocm_aiter_mla_sparse.py | 62 +++- 5 files changed, 349 insertions(+), 16 deletions(-) create mode 100644 tests/kernels/attention/test_rocm_aiter_mla_sparse_index_remap.py diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sink.py b/tests/kernels/attention/test_rocm_aiter_mla_sink.py index f66d408c8495..399a1fd0e3d0 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sink.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sink.py @@ -476,6 +476,8 @@ def test_sparse_mla_sink_forward_mqa_preserves_split_query(dtype): (num_tokens, 128), -1, dtype=torch.int32, device="cuda" ) impl.topk_indices_buffer[:, : selected.shape[1]] = selected + impl.owns_indexer = True + impl.index_epoch = 0 impl.q_concat_buffer = torch.empty_like(q) metadata = SimpleNamespace( attn_out_dtype=dtype, diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sparse_index_remap.py b/tests/kernels/attention/test_rocm_aiter_mla_sparse_index_remap.py new file mode 100644 index 000000000000..2ad9d3a70a52 --- /dev/null +++ b/tests/kernels/attention/test_rocm_aiter_mla_sparse_index_remap.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip( + "ROCm AITER sparse MLA index remap test requires ROCm.", + allow_module_level=True, + ) + +from vllm._aiter_ops import is_aiter_found_and_supported + +if not is_aiter_found_and_supported(): + pytest.skip( + "ROCm AITER sparse MLA index remap test requires a supported AITER " + "installation.", + allow_module_level=True, + ) + +from vllm.model_executor.layers.attention import sparse_mla_attention as sparse_mla_mod +from vllm.v1.attention.backends.mla import rocm_aiter_mla_sparse as sparse_mod + +NUM_HEADS = 16 +KV_LORA_RANK = 512 +QK_NOPE_HEAD_DIM = 128 +QK_ROPE_HEAD_DIM = 64 +HEAD_SIZE = 576 +NUM_TOKENS = 2 +TOPK_TOKENS = 4 + + +@pytest.fixture(autouse=True) +def _reset_index_epochs(): + # The counter is keyed by id(), and a freed test buffer can hand its id to + # the next one, so a leftover count would leak across tests. + sparse_mod._INDEX_EPOCHS.clear() + yield + sparse_mod._INDEX_EPOCHS.clear() + + +def _make_metadata(): + return SimpleNamespace( + num_actual_tokens=NUM_TOKENS, + req_id_per_token=torch.zeros(NUM_TOKENS, dtype=torch.int32, device="cpu"), + block_table=torch.arange(16, dtype=torch.int32, device="cpu").view(2, 8), + paged_kv_indptr=torch.zeros(NUM_TOKENS + 1, dtype=torch.int32, device="cpu"), + paged_kv_indices=torch.zeros( + NUM_TOKENS * TOPK_TOKENS, dtype=torch.int32, device="cpu" + ), + block_size=1, + topk_tokens=TOPK_TOKENS, + remapped_buf=None, + remapped_epoch=-1, + ) + + +def _patch_init_deps(monkeypatch, *, speculative: bool = False): + monkeypatch.setattr( + sparse_mod.AiterMLAHelper, + "check_num_heads_validity", + staticmethod(lambda num_heads: None), + ) + vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=NUM_TOKENS), + model_config=SimpleNamespace(dtype=torch.bfloat16), + speculative_config=SimpleNamespace() if speculative else None, + attention_config=SimpleNamespace(hisparse_config=None), + ) + monkeypatch.setattr(sparse_mod, "get_current_vllm_config", lambda: vllm_config) + monkeypatch.setattr(sparse_mla_mod, "get_current_vllm_config", lambda: vllm_config) + monkeypatch.setattr( + sparse_mla_mod, "get_tensor_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr( + sparse_mod, + "current_workspace_manager", + lambda: SimpleNamespace( + get_simultaneous=lambda *specs: tuple( + torch.zeros(shape, dtype=dtype, device="cpu") for shape, dtype in specs + ) + ), + ) + + +def _build_impl(indexer, topk_indices_buffer): + return sparse_mod.ROCMAiterMLASparseImpl( + num_heads=NUM_HEADS, + head_size=HEAD_SIZE, + scale=1.0, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="auto", + logits_soft_cap=None, + attn_type="decoder", + kv_sharing_target_layer_name=None, + topk_indices_buffer=topk_indices_buffer, + indexer=indexer, + q_lora_rank=None, + kv_lora_rank=KV_LORA_RANK, + qk_nope_head_dim=QK_NOPE_HEAD_DIM, + qk_rope_head_dim=QK_ROPE_HEAD_DIM, + qk_head_dim=QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM, + v_head_dim=QK_NOPE_HEAD_DIM, + kv_b_proj=None, + ) + + +def _new_buffer(): + return torch.zeros(NUM_TOKENS, TOPK_TOKENS, dtype=torch.int32, device="cpu") + + +def _patch_forward_deps(monkeypatch, impl, remap_calls): + monkeypatch.setattr( + sparse_mod, + "triton_convert_req_index_to_global_index", + lambda *args, **kwargs: remap_calls.append(args), + ) + monkeypatch.setattr( + sparse_mod.AiterMLAHelper, + "get_mla_padded_q", + staticmethod(lambda num_heads, q: q), + ) + attn_out = torch.zeros(NUM_TOKENS, NUM_HEADS, KV_LORA_RANK, device="cpu") + monkeypatch.setattr( + type(impl), + "_forward_mla", + lambda self, layer, q, kv_cache, attn_metadata: (attn_out, None), + ) + return attn_out + + +def _run_forward(impl, metadata): + return impl.forward_mqa( + torch.zeros(NUM_TOKENS, NUM_HEADS, HEAD_SIZE, device="cpu"), + torch.zeros(4, 1, HEAD_SIZE, device="cpu"), + metadata, + SimpleNamespace(_q_scale=None, _k_scale=None), + ) + + +@pytest.mark.parametrize("has_indexer", [True, False]) +def test_owns_indexer_is_structural(monkeypatch, has_indexer): + """owns_indexer reflects the layer, not the serving config.""" + _patch_init_deps(monkeypatch) + buffer = _new_buffer() + indexer = SimpleNamespace(topk_indices_buffer=buffer) if has_indexer else None + + impl = _build_impl(indexer, buffer) + + assert impl.owns_indexer is has_indexer + assert impl.topk_indices_buffer is buffer + + +def test_forward_mqa_remaps_once_per_metadata(monkeypatch): + """Skip layers sharing a metadata share its remap; the first one pays. + + All layers read one top-k buffer, so the reuse has to key off the metadata + that owns ``paged_kv_indices``, not off the buffer being distinct. + """ + _patch_init_deps(monkeypatch) + buffer = _new_buffer() + first = _build_impl(None, buffer) + second = _build_impl(None, buffer) + + remap_calls: list[tuple] = [] + attn_out = _patch_forward_deps(monkeypatch, first, remap_calls) + metadata = _make_metadata() + + out, lse = _run_forward(first, metadata) + assert out is attn_out + assert lse is None + assert len(remap_calls) == 1 + assert metadata.remapped_buf is buffer + + _run_forward(second, metadata) + assert len(remap_calls) == 1 + + +def test_forward_mqa_remaps_for_each_metadata(monkeypatch): + """Skip layers on a second metadata must remap its own paged_kv_indices. + + Regression: gating on ``indexer is not None`` skipped every layer in the + non-indexer metadata group, leaving its ``paged_kv_indices`` unwritten. + """ + _patch_init_deps(monkeypatch) + buffer = _new_buffer() + indexer_layer = _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer) + skip_layer = _build_impl(None, buffer) + + remap_calls: list[tuple] = [] + _patch_forward_deps(monkeypatch, indexer_layer, remap_calls) + + indexer_metadata = _make_metadata() + _run_forward(indexer_layer, indexer_metadata) + assert len(remap_calls) == 1 + assert indexer_metadata.remapped_buf is buffer + + # A distinct metadata instance carries its own, still-unwritten indices. + skip_metadata = _make_metadata() + _run_forward(skip_layer, skip_metadata) + assert len(remap_calls) == 2 + assert skip_metadata.remapped_buf is buffer + + +def test_forward_mqa_remaps_after_a_later_indexer_write(monkeypatch): + """A skip metadata must be refreshed once per indexer write, not once ever. + + Regression for cross-layer index sharing (``index_topk_freq``). With + freq=4 on a 78-layer model the 57 skip layers share one metadata while the + 21 indexer layers share another, so the skip metadata has to pick up every + new selection. Recording only the buffer left layers 4 onwards attending + to the first selection of the pass, which collapsed RULER niah_single_2 to + 0.026 on GLM-5.3 and took GPQA from 0.884 to 0.672. + """ + _patch_init_deps(monkeypatch) + buffer = _new_buffer() + # Construction order mirrors the layer stack: indexer, three skips, indexer. + _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer) + first_skip = _build_impl(None, buffer) + reuse_skip = _build_impl(None, buffer) + _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer) + next_skip = _build_impl(None, buffer) + + remap_calls: list[tuple] = [] + _patch_forward_deps(monkeypatch, first_skip, remap_calls) + skip_metadata = _make_metadata() + + _run_forward(first_skip, skip_metadata) + assert len(remap_calls) == 1 + + # Same selection, same metadata: reuse. + _run_forward(reuse_skip, skip_metadata) + assert len(remap_calls) == 1 + + # A later indexer layer has since rewritten the buffer in place. The object + # is unchanged, so only the epoch can reveal it. + _run_forward(next_skip, skip_metadata) + assert len(remap_calls) == 2 + assert skip_metadata.remapped_epoch == next_skip.index_epoch + + +def test_index_epoch_tracks_the_preceding_indexer_write(monkeypatch): + """Layers between two indexer layers share the earlier one's epoch.""" + _patch_init_deps(monkeypatch) + buffer = _new_buffer() + first_indexer = _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer) + skip_a = _build_impl(None, buffer) + skip_b = _build_impl(None, buffer) + second_indexer = _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer) + skip_c = _build_impl(None, buffer) + + assert skip_a.index_epoch == first_indexer.index_epoch + assert skip_b.index_epoch == first_indexer.index_epoch + assert second_indexer.index_epoch == first_indexer.index_epoch + 1 + assert skip_c.index_epoch == second_indexer.index_epoch + + +def test_forward_mqa_remaps_when_buffer_differs(monkeypatch): + """A layer reading a different buffer cannot reuse the metadata's remap.""" + _patch_init_deps(monkeypatch) + remapped_buffer = _new_buffer() + other_buffer = _new_buffer() + layer = _build_impl(None, other_buffer) + + remap_calls: list[tuple] = [] + _patch_forward_deps(monkeypatch, layer, remap_calls) + metadata = _make_metadata() + metadata.remapped_buf = remapped_buffer + + _run_forward(layer, metadata) + assert len(remap_calls) == 1 + assert metadata.remapped_buf is other_buffer + + +def test_forward_mqa_indexer_layer_always_remaps(monkeypatch): + """The indexer rewrites its buffer in place, so its remap is never stale.""" + _patch_init_deps(monkeypatch) + buffer = _new_buffer() + impl = _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer) + + remap_calls: list[tuple] = [] + _patch_forward_deps(monkeypatch, impl, remap_calls) + metadata = _make_metadata() + metadata.remapped_buf = buffer + metadata.remapped_epoch = impl.index_epoch + + _run_forward(impl, metadata) + assert len(remap_calls) == 1 diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py index dc1ff79eeabd..4289aac1e7cd 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py @@ -523,6 +523,8 @@ def test_forward_mqa_slices_full_batch_metadata_to_query_tokens(monkeypatch): impl.kv_cache_dtype = "auto" impl.num_heads = 4 impl.topk_indices_buffer = torch.arange(48, dtype=torch.int32).view(12, 4) + impl.owns_indexer = True + impl.index_epoch = 0 expected_output = torch.zeros(2, 4, 512, dtype=torch.bfloat16) forward_mla = Mock(return_value=(expected_output, None)) monkeypatch.setattr(impl, "_forward_mla", forward_mla) @@ -570,6 +572,8 @@ def test_forward_mqa_slices_decode_subset_for_bf16_and_fp8(monkeypatch, kv_cache impl.kv_cache_dtype = kv_cache_dtype impl.num_heads = 4 impl.topk_indices_buffer = torch.arange(48, dtype=torch.int32).view(12, 4) + impl.owns_indexer = True + impl.index_epoch = 0 expected_output = torch.zeros(2, 4, 512, dtype=torch.bfloat16) forward_mla = Mock(return_value=(expected_output, None)) monkeypatch.setattr(impl, "_forward_mla", forward_mla) diff --git a/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py b/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py index 645eec29796d..6849c4e3ebde 100644 --- a/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py +++ b/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py @@ -96,6 +96,8 @@ def test_forward_mqa_prepares_triton_sparse_mla_inputs( impl.sinks = torch.randn(num_heads, generator=gen).to(device) impl.kv_cache_dtype = kv_cache_dtype impl.topk_indices_buffer = topk_indices.to(device) + impl.owns_indexer = True + impl.index_epoch = 0 impl.use_aiter_sparse_mla = True layer = SimpleNamespace( _q_scale=torch.tensor([0.02], device=device), diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index a0393e772c3e..037b5432ed03 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -401,6 +401,18 @@ def supports_sink(cls) -> bool: return on_mi3xx() +# Indexer writes per shared top-k buffer, in layer construction order. +_INDEX_EPOCHS: dict[int, int] = {} + + +def _index_epoch(buf: object, owns_indexer: bool) -> int: + """Number of indexer writes up to and including this layer.""" + # id(): tensors are unhashable, and the buffer outlives every reader. + epoch = _INDEX_EPOCHS.get(id(buf), 0) + (1 if owns_indexer else 0) + _INDEX_EPOCHS[id(buf)] = epoch + return epoch + + @dataclass(kw_only=True) class ROCMAiterMLASparseMetadata(SparseMLACommonMetadata): qo_indptr: torch.Tensor | None = None @@ -409,6 +421,10 @@ class ROCMAiterMLASparseMetadata(SparseMLACommonMetadata): paged_kv_indptr: torch.Tensor | None = None attn_out_dtype: torch.dtype | None = None + # Source buffer and index_epoch of paged_kv_indices; -1 before any remap. + remapped_buf: object = None + remapped_epoch: int = -1 + # Persistent MLA metadata (only populated when persistent mode is enabled, # i.e. when the aiter sparse decode kernel supports work-stealing splits). work_meta_data: torch.Tensor | None = None @@ -818,6 +834,10 @@ def __init__( self.masked_mha_available = False vllm_config = get_current_vllm_config() + self.owns_indexer: bool = indexer is not None + self.index_epoch: int = _index_epoch( + self.topk_indices_buffer, self.owns_indexer + ) max_tokens = vllm_config.scheduler_config.max_num_batched_tokens q_concat_shape = (max_tokens, num_heads, head_size) (self.q_concat_buffer,) = current_workspace_manager().get_simultaneous( @@ -1162,25 +1182,35 @@ def forward_mqa( num_actual_toks = q.shape[0] - # Get topk indices - assert self.topk_indices_buffer is not None - topk_indices = fit_kpool_indices_to_aiter( - self.topk_indices_buffer[:num_actual_toks], attn_metadata.topk_tokens - ) - kv_c_and_k_pe_cache, block_stride_rows = flat_kv_row_view( kv_c_and_k_pe_cache, attn_metadata.block_size ) - triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token[:num_actual_toks], - attn_metadata.block_table, - topk_indices, - attn_metadata.paged_kv_indptr, - attn_metadata.paged_kv_indices, - BLOCK_SIZE=attn_metadata.block_size, - BLOCK_STRIDE_ROWS=block_stride_rows, - NUM_TOPK_TOKENS=attn_metadata.topk_tokens, - ) + + # Indexer layers may rewrite their buffer, so always remap. Non-indexer + # layers have their own KV cache group: remap once per indexer write. + if ( + self.owns_indexer + or attn_metadata.remapped_epoch != self.index_epoch + or attn_metadata.remapped_buf is not self.topk_indices_buffer + ): + # Get topk indices + assert self.topk_indices_buffer is not None + topk_indices = fit_kpool_indices_to_aiter( + self.topk_indices_buffer[:num_actual_toks], attn_metadata.topk_tokens + ) + + triton_convert_req_index_to_global_index( + attn_metadata.req_id_per_token[:num_actual_toks], + attn_metadata.block_table, + topk_indices, + attn_metadata.paged_kv_indptr, + attn_metadata.paged_kv_indices, + BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, + NUM_TOPK_TOKENS=attn_metadata.topk_tokens, + ) + attn_metadata.remapped_buf = self.topk_indices_buffer + attn_metadata.remapped_epoch = self.index_epoch # write the latent and rope to kv cache if fp8_attention: From 380638afbb2c5658f22124f6d1d38fcd464ea7d9 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Tue, 13 Oct 2026 13:46:37 +0800 Subject: [PATCH 359/624] [Model] Add Clef-Flash support (#61308) Signed-off-by: Tianyu Guo Signed-off-by: Isotr0py Co-authored-by: Codex Co-authored-by: Isotr0py --- .../online_serving/structured_decisions.md | 26 +- tests/config/test_model_arch_config.py | 34 ++ .../structured_decisions/test_strategies.py | 65 ++-- .../multimodal/pooling/test_clef_omni.py | 7 +- tests/models/registry.py | 4 + .../structured_decisions/strategies.py | 56 +++ vllm/model_executor/models/clef_flash.py | 44 +++ .../models/clef_joint_schema.py | 343 ++++++++++++++++++ vllm/model_executor/models/clef_omni.py | 339 +---------------- vllm/model_executor/models/registry.py | 1 + .../model_arch_config_convertor.py | 22 ++ 11 files changed, 580 insertions(+), 361 deletions(-) create mode 100644 vllm/model_executor/models/clef_flash.py create mode 100644 vllm/model_executor/models/clef_joint_schema.py diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md index ce09d6f81fee..4a7ec00dd2ea 100644 --- a/docs/serving/online_serving/structured_decisions.md +++ b/docs/serving/online_serving/structured_decisions.md @@ -130,6 +130,30 @@ the request lists them, a `noul` labels its options `yes` and `no`, and a | options per `choice` | 26 | | levels per `score` | 10 | +## Clef-Flash + +[Clef-Flash](https://huggingface.co/Cloudflare/clef-flash) uses a Qwen3.5-9B +backbone and a joint schema head to score all questions in one forward pass. +vLLM detects the head and serves the checkpoint as a pooling model. + +```bash +vllm serve Cloudflare/clef-flash --max-model-len 16384 +``` + +Send requests to `/v1/systemone` using `state` and `questions`, as in the +examples above. `choice`, `noul` and `score` questions are supported. + +- Text and JSON states use the checkpoint's fixed prompt. JSON is compact + with sorted keys. `instructions` returns 400; `chat_template_kwargs` + has no effect. +- The state is truncated to fit 16,384 tokens, or `max_model_len` if smaller. + A schema that exceeds this bound returns 400. +- `images` and `videos` are placed before the state. Videos are sampled at + 2 fps and processed by the Qwen3.5 vision processor. Audio inputs return + 400, and video soundtracks are ignored. +- Each request performs one pooling read. `usage.input_tokens` counts the + prompt once, `output_tokens` is 0, and `label_mass` is 1. + ## Clef-Omni [Clef-Omni](https://huggingface.co/Cloudflare/clef-omni) is a decision model: @@ -188,7 +212,7 @@ vllm serve LiquidAI/d1-3B | `instructions` | optional context placed ahead of the questions | | `chat_template_kwargs` | passed to the chat template, for example `{"enable_thinking": false}` | | `cache_salt` | optional [prefix-cache salt](../../usage/security.md#cache-salting) applied to every question | -| `images`, `audio`, `videos` | optional media URLs or data URLs, for models that take them ([Clef-Omni](#clef-omni), [Liquid d1](#liquid-d1), images only) | +| `images`, `audio`, `videos` | optional media URLs or data URLs, for models that take them ([Clef-Flash](#clef-flash), images/video; [Clef-Omni](#clef-omni); [Liquid d1](#liquid-d1), images only) | A question with any field other than `type`, `instructions` and `criteria` is rejected with a 400. diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index 8ec16ee7f9fb..2ff992a20562 100644 --- a/tests/config/test_model_arch_config.py +++ b/tests/config/test_model_arch_config.py @@ -70,6 +70,40 @@ def _load_groundtruth(filename: str) -> dict: return json.load(f) +@pytest.mark.parametrize("has_joint_head", [False, True]) +@pytest.mark.parametrize( + "architecture", ["Qwen3_5ForConditionalGeneration", "ColQwen3_5"] +) +def test_qwen35_joint_head_selects_clef_flash(tmp_path, has_joint_head, architecture): + """Only a Qwen3.5 generation checkpoint with a joint head selects Clef.""" + from transformers import Qwen3_5Config + + head_config = { + "hidden_size": 4096, + "width": 1024, + "routing_layers": 2, + "layers": 4, + "heads": 16, + "feedforward": 4096, + } + if has_joint_head: + (tmp_path / "joint_head_config.json").write_text(json.dumps(head_config)) + config = Qwen3_5Config(architectures=[architecture]) + config.name_or_path = str(tmp_path) + model_config = SimpleNamespace( + hf_config=config, + hf_text_config=config.text_config, + _supports_multimodal_for_mm_prefix=lambda: True, + ) + converted = ModelConfig.get_model_arch_config(model_config) + if has_joint_head and architecture == "Qwen3_5ForConditionalGeneration": + assert converted.architectures == ["ClefFlashForJointSchema"] + assert config.joint_head_config == head_config + else: + assert converted.architectures == [architecture] + assert not hasattr(config, "joint_head_config") + + def _assert_model_arch_config( model_config, expected: dict, check_head_size: bool = True ): diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index dae545d1f04c..505f0e35eef1 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -33,6 +33,7 @@ parse_questions, ) from vllm.entrypoints.generate.structured_decisions.strategies import ( + ClefFlashStrategy, DiffusionGemmaCanvasStrategy, JointSchemaStrategy, LiquidStrategy, @@ -94,6 +95,8 @@ def test_strategy_selection(): # Clef-Omni scores options with its head, so it reads no logprobs. clef = model("ClefOmniForJointSchema", "raw_logits") assert select_read_strategy(clef, ("token_classify",)) is JointSchemaStrategy + clef_flash = model("ClefFlashForJointSchema", "raw_logits") + assert select_read_strategy(clef_flash, ("token_classify",)) is ClefFlashStrategy def test_route_is_registered_by_default(): @@ -366,8 +369,8 @@ def test_canvas_read(): ) -@pytest.fixture -def joint_schema(qwen, engine_and_renderer): +@pytest.fixture(params=[JointSchemaStrategy, ClefFlashStrategy]) +def joint_schema(request, qwen, engine_and_renderer): tokenizer, _ = qwen engine, online_renderer = engine_and_renderer # The probability the head gives each option, by option id. @@ -382,9 +385,7 @@ async def encode(prompt, params, request_id, **kwargs): yield PoolingRequestOutput(request_id, PoolingOutput(data), ids, 0, True) engine.encode.side_effect = encode - strategy = JointSchemaStrategy( - ReadContext(engine, online_renderer, None, "auto", {}) - ) + strategy = request.param(ReadContext(engine, online_renderer, None, "auto", {})) return serving(engine, strategy), engine, tokenizer @@ -468,7 +469,10 @@ async def test_media_needs_a_model_that_takes_it(serving_and_engine): @pytest.mark.asyncio -async def test_joint_schema_reads_media_before_the_state(joint_schema, monkeypatch): +@pytest.mark.parametrize("num_images", [1, 2]) +async def test_joint_schema_reads_media_before_the_state( + joint_schema, monkeypatch, num_images +): serving, engine, tokenizer = joint_schema strategy = serving.strategy strategy.__dict__["_connector"] = MediaConnector() @@ -491,23 +495,30 @@ async def process_for_engine_async(prompt, arrival_time): Image.new("RGB", (8, 8), "red").save(buffer, format="PNG") image = "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode() response = await serving.create_decision( - StructuredDecisionRequest(**JOINT_SCHEMA_REQUEST, images=[image]) + StructuredDecisionRequest(**JOINT_SCHEMA_REQUEST, images=[image] * num_images) ) assert isinstance(response, StructuredDecisionResponse) - placeholder = tokenizer.encode( - "<|vision_start|><|image_pad|><|vision_end|>\n", add_special_tokens=False - ) + if isinstance(strategy, ClefFlashStrategy): + media_text = "<|vision_start|><|image_pad|><|vision_end|>" * num_images + "\n" + else: + media_text = "<|vision_start|><|image_pad|><|vision_end|>\n" * num_images + placeholder = tokenizer.encode(media_text, add_special_tokens=False) assert rendered["prompt_token_ids"] == strategy.prefix + placeholder assert rendered["multi_modal_data"]["image"][0].size == (8, 8) - assert rendered["mm_processor_kwargs"] == { - "images_kwargs": {"size": strategy.IMAGE_SIZE} - } + if isinstance(strategy, ClefFlashStrategy): + assert rendered["mm_processor_kwargs"] == { + "videos_kwargs": {"do_sample_frames": False} + } + else: + assert rendered["mm_processor_kwargs"] == { + "images_kwargs": {"size": strategy.IMAGE_SIZE} + } prompt = engine.encode.call_args.args[0] assert prompt["type"] == "multimodal" ids = prompt["prompt_token_ids"] - expanded = len(strategy.prefix) + len(placeholder) + 2 - assert ids[:expanded].count(image_pad) == 3 + expanded = len(strategy.prefix) + len(placeholder) + 2 * num_images + assert ids[:expanded].count(image_pad) == 3 * num_images state_and_questions = JOINT_SCHEMA_PROMPT.split("STATE:\n", 1)[1] assert tokenizer.decode(ids[expanded:]) == state_and_questions assert response.answers["team"]["choice"] == "billing" @@ -515,7 +526,9 @@ async def process_for_engine_async(prompt, arrival_time): @pytest.mark.asyncio -async def test_joint_schema_hears_video_soundtracks(joint_schema, monkeypatch): +async def test_joint_schema_video_audio_follows_model_modalities( + joint_schema, monkeypatch +): serving, engine, tokenizer = joint_schema strategy = serving.strategy video = MediaWithBytes( @@ -532,6 +545,7 @@ async def load_from_url_async(self, url, media_io, *, fetch_timeout=None): return video async def fetch_audio_async(self, url): + assert not isinstance(strategy, ClefFlashStrategy) if url == "https://example.org/silent.mp4": raise ValueError("No audio stream found.") return track @@ -552,11 +566,17 @@ async def process_for_engine_async(prompt, arrival_time): ) ) assert isinstance(response, StructuredDecisionResponse) - assert rendered["multi_modal_data"] == {"video": [video], "audio": [track]} - assert rendered["mm_processor_kwargs"] == { - "images_kwargs": {"size": strategy.IMAGE_SIZE}, - "use_audio_in_video": True, - } + if isinstance(strategy, ClefFlashStrategy): + assert rendered["multi_modal_data"] == {"video": [video]} + assert rendered["mm_processor_kwargs"] == { + "videos_kwargs": {"do_sample_frames": False} + } + else: + assert rendered["multi_modal_data"] == {"video": [video], "audio": [track]} + assert rendered["mm_processor_kwargs"] == { + "images_kwargs": {"size": strategy.IMAGE_SIZE}, + "use_audio_in_video": True, + } # A video without a decodable soundtrack is watched silently. response = await serving.create_decision( @@ -566,7 +586,8 @@ async def process_for_engine_async(prompt, arrival_time): ) assert isinstance(response, StructuredDecisionResponse) assert rendered["multi_modal_data"] == {"video": [video]} - assert rendered["mm_processor_kwargs"]["use_audio_in_video"] is False + if not isinstance(strategy, ClefFlashStrategy): + assert rendered["mm_processor_kwargs"]["use_audio_in_video"] is False # Audio clips cannot be read with videos that have soundtracks. response = await serving.create_decision( diff --git a/tests/models/multimodal/pooling/test_clef_omni.py b/tests/models/multimodal/pooling/test_clef_omni.py index ccdc913b3132..381ac23d53cc 100644 --- a/tests/models/multimodal/pooling/test_clef_omni.py +++ b/tests/models/multimodal/pooling/test_clef_omni.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The Clef-Omni pooler against a port of the checkpoint's joint schema head, +"""The shared Clef pooler against a port of the checkpoint's joint schema head, which loops over questions where the pooler batches them.""" import math @@ -11,7 +11,10 @@ import torch import torch.nn.functional as F -from vllm.model_executor.models.clef_omni import JointSchemaHead, JointSchemaPooler +from vllm.model_executor.models.clef_joint_schema import ( + JointSchemaHead, + JointSchemaPooler, +) from vllm.pooling_params import PoolingParams from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates diff --git a/tests/models/registry.py b/tests/models/registry.py index 0a7027461ef3..b08efd9fb0f1 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -802,6 +802,10 @@ def check_available_online( "Salesforce/blip2-opt-2.7b", extras={"6b": "Salesforce/blip2-opt-6.7b"}, ), + "ClefFlashForJointSchema": _HfExamplesInfo( + "Cloudflare/clef-flash", + max_model_len=4096, + ), "Cohere2VisionForConditionalGeneration": _HfExamplesInfo( "CohereLabs/command-a-vision-07-2025", extras={"command-a-plus": "CohereLabs/command-a-plus-05-2026-bf16"}, diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index 5567cb808ab5..66238cf7f103 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -809,7 +809,63 @@ async def read( return reads +class ClefFlashStrategy(JointSchemaStrategy): + """Clef-Flash's joint schema prompt with Qwen3.5 image/video processing.""" + + modalities = frozenset({"image", "video"}) + MAX_LENGTH = 16384 + PLACEHOLDERS = { + "image": "<|vision_start|><|image_pad|><|vision_end|>", + "video": "<|vision_start|><|video_pad|><|vision_end|>", + } + + async def _render_media( + self, media: Mapping[str, Sequence[str]], cache_salt: str | None + ) -> MultiModalInput: + connector = self._connector + video_io = VideoMediaIO(ImageMediaIO(), num_frames=-1, fps=self.VIDEO_FPS) + try: + images, videos = await asyncio.gather( + asyncio.gather( + *map(connector.fetch_image_async, media.get("image", ())) + ), + self._fetch( + media.get("video", ()), video_io, envs.VLLM_VIDEO_FETCH_TIMEOUT + ), + ) + except ValueError as e: + raise StructuredDecisionError(f"media: {e}") from e + text = ( + self.PLACEHOLDERS["image"] * len(images) + + self.PLACEHOLDERS["video"] * len(videos) + + "\n" + ) + mm_data: dict[str, Any] = {} + if images: + mm_data["image"] = images + if videos: + mm_data["video"] = videos + prompt = TokensPrompt( + prompt_token_ids=self.prefix + self._tokens(text), + multi_modal_data=mm_data, + mm_processor_kwargs={"videos_kwargs": {"do_sample_frames": False}}, + ) + if cache_salt is not None: + prompt["cache_salt"] = cache_salt + try: + engine_input = ( + await self.context.online_renderer.renderer.process_for_engine_async( + prompt, time.time() + ) + ) + except ValueError as e: + raise StructuredDecisionError(f"media: {e}") from e + assert engine_input["type"] == "multimodal" + return engine_input + + READ_STRATEGIES: dict[str, type[ReadStrategy]] = { + "ClefFlashForJointSchema": ClefFlashStrategy, "ClefOmniForJointSchema": JointSchemaStrategy, "Lfm2VlForConditionalGeneration": LiquidStrategy, "Qwen3ForCausalLM": NextTokenStrategy, diff --git a/vllm/model_executor/models/clef_flash.py b/vllm/model_executor/models/clef_flash.py new file mode 100644 index 000000000000..32a9e0a59003 --- /dev/null +++ b/vllm/model_executor/models/clef_flash.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Clef-Flash: the Qwen3.5 backbone with a joint schema decision head.""" + +import torch + +from vllm.config import VllmConfig +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from vllm.model_executor.model_loader.default_loader import DefaultModelLoader +from vllm.multimodal import MULTIMODAL_REGISTRY + +from .clef_joint_schema import JOINT_HEAD_WEIGHTS, JointSchemaHead, JointSchemaPooler +from .interfaces_base import default_pooling_type +from .qwen3_5 import Qwen3_5ForConditionalGeneration, Qwen3_5ProcessingInfo +from .qwen3_vl import Qwen3VLDummyInputsBuilder, Qwen3VLMultiModalProcessor + + +@default_pooling_type(tok_pooling_type="ALL") +@MULTIMODAL_REGISTRY.register_processor( + Qwen3VLMultiModalProcessor, + info=Qwen3_5ProcessingInfo, + dummy_inputs=Qwen3VLDummyInputsBuilder, +) +class ClefFlashForJointSchema(Qwen3_5ForConditionalGeneration): + is_pooling_model = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + model_config = vllm_config.model_config + self.pooler = JointSchemaPooler( + JointSchemaHead(**model_config.hf_config.joint_head_config), + self._embed_lexical, + ) + self.secondary_weights = [ + DefaultModelLoader.Source( + model_or_path=model_config.model, + revision=model_config.revision, + prefix="pooler.head.", + allow_patterns_overrides=[JOINT_HEAD_WEIGHTS], + ) + ] + + def _embed_lexical(self, token_ids: torch.Tensor) -> torch.Tensor: + return VocabParallelEmbedding.forward(self.language_model.lm_head, token_ids) diff --git a/vllm/model_executor/models/clef_joint_schema.py b/vllm/model_executor/models/clef_joint_schema.py new file mode 100644 index 000000000000..fad2beb09899 --- /dev/null +++ b/vllm/model_executor/models/clef_joint_schema.py @@ -0,0 +1,343 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared joint schema head for Clef decision models. + +The head reads the final hidden states of the whole prompt, so it runs as a +``token_classify`` pooler. + +Each request carries its schema in ``PoolingParams.extra_kwargs`` under +``"joint_schema"``: + +- ``question_types``: 0 (noul), 1 (choice) or 2 (score) per question. +- ``question_spans``: the ``[start, end)`` of each question's instructions. +- ``option_spans``: the ``[start, end)`` of each option, question by question. +- ``option_counts``: the number of options of each question. + +Spans count from the end of the prompt (``-len(prompt) <= start < end <= 0``), +so they hold after multimodal placeholders are expanded. The output has one +value per option, in ``option_spans`` order: its probability within its +question, or its logit with ``use_activation=False``. +""" + +import math +from collections.abc import Callable, Set +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn + +from vllm.model_executor.layers.pooler import Pooler, PoolingParamsUpdate +from vllm.model_executor.layers.pooler.tokwise.methods import AllPool +from vllm.tasks import PoolingTask +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.outputs import PoolerOutput +from vllm.v1.pool.metadata import PoolingMetadata + +JOINT_HEAD_WEIGHTS = "joint_head.safetensors" +NUM_QUESTION_TYPES = 3 +MAX_LOGIT_SCALE = math.log(100.0) + + +def segment_mean( + values: torch.Tensor, segment_ids: torch.Tensor, lengths: torch.Tensor +) -> torch.Tensor: + sums = values.new_zeros(lengths.shape[0], values.shape[-1], dtype=torch.float32) + sums.index_add_(0, segment_ids, values.float()) + return (sums / lengths[:, None]).to(values.dtype) + + +@dataclass +class SchemaIndices: + """Where one prompt's questions and options are, on the model's device.""" + + positions: torch.Tensor + """Every token of every span, questions first.""" + span_ids: torch.Tensor + """The span of each of ``positions``.""" + span_lengths: torch.Tensor + """``[Q + O]``.""" + question_types: torch.Tensor + """``[Q]``.""" + option_question: torch.Tensor + """The question of each option, ``[O]``.""" + option_slots: torch.Tensor + """Each option's flat index in a ``[Q, max_options]`` grid, ``[O]``.""" + option_mask: torch.Tensor + """``[Q, max_options]``, true at the slots of options.""" + + +class EvidenceRoutingLayer(nn.Module): + def __init__(self, width: int, heads: int, feedforward: int) -> None: + super().__init__() + self.query_norm = nn.LayerNorm(width) + self.memory_norm = nn.LayerNorm(width) + self.attention = nn.MultiheadAttention(width, heads, batch_first=True) + self.feedforward_norm = nn.LayerNorm(width) + # The Identity takes the place of the checkpoint's dropout. + self.feedforward = nn.Sequential( + nn.Linear(width, feedforward), + nn.GELU(), + nn.Identity(), + nn.Linear(feedforward, width), + ) + + def forward(self, queries: torch.Tensor, memory: torch.Tensor) -> torch.Tensor: + # Separate key and value tensors, like the reference, keep attention's + # unpacked input projections. + routed, _ = self.attention( + self.query_norm(queries), + self.memory_norm(memory), + self.memory_norm(memory), + need_weights=False, + ) + queries = queries + routed + return queries + self.feedforward(self.feedforward_norm(queries)) + + +class JointSchemaHead(nn.Module): + def __init__( + self, + hidden_size: int, + width: int, + routing_layers: int, + layers: int, + heads: int, + feedforward: int, + ) -> None: + super().__init__() + self.hidden_norm = nn.LayerNorm(hidden_size) + self.memory_projection = nn.Linear(hidden_size, width, bias=False) + self.question_projection = nn.Linear(hidden_size, width, bias=False) + self.option_question_projection = nn.Linear(hidden_size, width, bias=False) + self.global_projection = nn.Linear(hidden_size, width, bias=False) + self.option_context_projection = nn.Linear(hidden_size, width, bias=False) + self.option_lexical_projection = nn.Linear(hidden_size, width, bias=False) + self.type_embedding = nn.Embedding(NUM_QUESTION_TYPES, width) + self.evidence_layers = nn.ModuleList( + EvidenceRoutingLayer(width, heads, feedforward) + for _ in range(routing_layers) + ) + self.option_summary_norm = nn.LayerNorm(width) + self.layers = nn.ModuleList( + nn.TransformerDecoderLayer( + d_model=width, + nhead=heads, + dim_feedforward=feedforward, + dropout=0.0, + activation="gelu", + batch_first=True, + norm_first=True, + ) + for _ in range(layers) + ) + self.field_norm = nn.LayerNorm(width) + self.option_norm = nn.LayerNorm(width) + self.residual_scorer = nn.Sequential( + nn.Linear(width * 4, width), + nn.GELU(), + nn.Identity(), + nn.Linear(width, 1), + ) + self.prior_logit_scale = nn.Parameter(torch.zeros(())) + self.joint_logit_scale = nn.Parameter(torch.zeros(())) + self.residual_gate = nn.Parameter(torch.zeros(())) + + def forward( + self, + hidden_states: torch.Tensor, + lexical: torch.Tensor, + schema: SchemaIndices, + ) -> torch.Tensor: + """The ``[O]`` option logits of one prompt, from its ``[L, hidden]`` + final hidden states and each option's ``[O, hidden]`` mean output + embedding.""" + num_questions = schema.question_types.shape[0] + option_question = schema.option_question + hidden_states = self.hidden_norm(hidden_states) + memory = self.memory_projection(hidden_states).unsqueeze(0) + global_vector = hidden_states[-1] + means = segment_mean( + hidden_states[schema.positions], schema.span_ids, schema.span_lengths + ) + question_vectors, contexts = means[:num_questions], means[num_questions:] + + options = ( + self.option_context_projection(contexts) + + self.option_lexical_projection(lexical) + + self.option_question_projection(question_vectors)[option_question] + ).unsqueeze(0) + for layer in self.evidence_layers: + options = layer(options, memory) + options = options[0] + + # Each field summarizes its options, weighted by their match to it. + base_fields = self.question_projection(question_vectors) + width = options.shape[-1] + grid = options.new_zeros(schema.option_mask.numel(), width) + grid[schema.option_slots] = options + grid = grid.view(*schema.option_mask.shape, width) + routing = torch.einsum("qow,qw->qo", grid, base_fields) / math.sqrt(width) + routing = routing.masked_fill(~schema.option_mask, float("-inf")) + summaries = torch.einsum("qo,qow->qw", routing.softmax(dim=-1), grid) + fields = ( + base_fields + + self.option_summary_norm(summaries) + + self.global_projection(global_vector).unsqueeze(0) + + self.type_embedding(schema.question_types) + ).unsqueeze(0) + for layer in self.layers: + fields = layer(fields, memory) + fields = self.field_norm(fields[0])[option_question] + + anchors = F.normalize(question_vectors + global_vector, dim=-1) + prior_scale = self.prior_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() + prior = prior_scale * ( + F.normalize(lexical, dim=-1) * anchors[option_question] + ).sum(dim=-1) + options = self.option_norm(options) + cosine = F.cosine_similarity(fields, options, dim=-1) + features = torch.cat( + [fields, options, fields * options, torch.abs(fields - options)], dim=-1 + ) + residual = self.residual_scorer(features).squeeze(-1) + joint_scale = self.joint_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() + joint = joint_scale * cosine + residual + return prior + torch.sigmoid(self.residual_gate) * joint + + +def parse_joint_schema( + extra_kwargs: dict[str, Any] | None, prompt_len: int +) -> tuple[list[int], list[tuple[int, int]], list[int]] | None: + """The question types, the absolute spans of the questions then of the + options, and the option counts. None if the schema is missing or invalid.""" + schema = (extra_kwargs or {}).get("joint_schema") + if not isinstance(schema, dict): + return None + try: + types = [int(t) for t in schema["question_types"]] + counts = [int(c) for c in schema["option_counts"]] + question_spans = [ + (int(start) + prompt_len, int(end) + prompt_len) + for start, end in schema["question_spans"] + ] + option_spans = [ + (int(start) + prompt_len, int(end) + prompt_len) + for start, end in schema["option_spans"] + ] + except (KeyError, TypeError, ValueError): + return None + spans = question_spans + option_spans + if ( + not types + or len(counts) != len(types) + or len(question_spans) != len(types) + or len(option_spans) != sum(counts) + or any(not 0 <= t < NUM_QUESTION_TYPES for t in types) + or any(c < 1 for c in counts) + or any(not 0 <= start < end <= prompt_len for start, end in spans) + ): + return None + return types, spans, counts + + +class JointSchemaPooler(Pooler): + """Runs the joint schema head on each finished prompt. + + It never raises, because it runs in the engine core: a prompt without a + valid schema, such as a profiling dummy, gets an empty output. + """ + + def __init__( + self, + head: JointSchemaHead, + embed_lexical: Callable[[torch.Tensor], torch.Tensor], + ) -> None: + super().__init__() + self.pooling = AllPool() + self.head = head + self.embed_lexical = embed_lexical + + def get_supported_tasks(self) -> Set[PoolingTask]: + return {"token_classify"} + + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate: + return PoolingParamsUpdate(requires_token_ids=True) + + def forward( + self, + hidden_states: torch.Tensor, + pooling_metadata: PoolingMetadata, + ) -> PoolerOutput: + pooled = self.pooling(hidden_states, pooling_metadata) + prompt_token_ids = pooling_metadata.get_prompt_token_ids_cpu() + outputs: list[torch.Tensor | None] = [] + for hidden, token_ids, params in zip( + pooled, prompt_token_ids, pooling_metadata.pooling_params + ): + if hidden is None: + outputs.append(None) + continue + parsed = parse_joint_schema(params.extra_kwargs, hidden.shape[0]) + if parsed is None: + outputs.append(hidden.new_empty(0, dtype=torch.float32)) + continue + schema, lexical_ids, lexical_span_ids = self._indices( + *parsed, hidden.device + ) + lexical = segment_mean( + self.embed_lexical( + async_tensor_h2d(token_ids[lexical_ids], hidden.device) + ), + lexical_span_ids, + schema.span_lengths[len(parsed[0]) :], + ) + logits = self.head(hidden, lexical, schema).float() + if params.use_activation is not False: + mask = schema.option_mask + grid = logits.new_full((mask.numel(),), float("-inf")) + grid[schema.option_slots] = logits + probs = grid.view(mask.shape).softmax(dim=-1) + logits = probs.view(-1)[schema.option_slots] + outputs.append(logits) + return outputs + + @staticmethod + def _indices( + types: list[int], + spans: list[tuple[int, int]], + counts: list[int], + device: torch.device, + ) -> tuple[SchemaIndices, torch.Tensor, torch.Tensor]: + """The schema's indices, plus the option tokens' positions (on the CPU) + and their option spans, for the lexical embeddings.""" + num_questions = len(types) + lengths = torch.tensor([end - start for start, end in spans]) + span_ids = torch.repeat_interleave(torch.arange(len(spans)), lengths) + positions = torch.cat([torch.arange(start, end) for start, end in spans]) + num_question_tokens = int(lengths[:num_questions].sum()) + + counts_cpu = torch.tensor(counts) + max_options = max(counts) + option_question = torch.repeat_interleave( + torch.arange(num_questions), counts_cpu + ) + option_slots = torch.cat( + [torch.arange(count) + q * max_options for q, count in enumerate(counts)] + ) + schema = SchemaIndices( + positions=async_tensor_h2d(positions, device), + span_ids=async_tensor_h2d(span_ids, device), + span_lengths=async_tensor_h2d(lengths, device), + question_types=async_tensor_h2d(types, device, dtype=torch.long), + option_question=async_tensor_h2d(option_question, device), + option_slots=async_tensor_h2d(option_slots, device), + option_mask=async_tensor_h2d( + torch.arange(max_options) < counts_cpu[:, None], device + ), + ) + lexical_span_ids = async_tensor_h2d( + span_ids[num_question_tokens:] - num_questions, device + ) + return schema, positions[num_question_tokens:], lexical_span_ids diff --git a/vllm/model_executor/models/clef_omni.py b/vllm/model_executor/models/clef_omni.py index 7cd0d676873b..1533dde27378 100644 --- a/vllm/model_executor/models/clef_omni.py +++ b/vllm/model_executor/models/clef_omni.py @@ -1,47 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Clef-Omni (https://huggingface.co/Cloudflare/clef-omni). - -The Qwen3-Omni thinker with a joint schema head that scores every option of -every question in one forward pass. The head reads the final hidden states of -the whole prompt, so it runs as a ``token_classify`` pooler. - -Each request carries its schema in ``PoolingParams.extra_kwargs`` under -``"joint_schema"``: - -- ``question_types``: 0 (noul), 1 (choice) or 2 (score) per question. -- ``question_spans``: the ``[start, end)`` of each question's instructions. -- ``option_spans``: the ``[start, end)`` of each option, question by question. -- ``option_counts``: the number of options of each question. - -Spans count from the end of the prompt (``-len(prompt) <= start < end <= 0``), -so they hold after multimodal placeholders are expanded. The output has one -value per option, in ``option_spans`` order: its probability within its -question, or its logit with ``use_activation=False``. -""" +"""The Qwen3-Omni thinker with a Clef joint schema head.""" import math -from collections.abc import Callable, Mapping, Set -from dataclasses import dataclass +from collections.abc import Mapping from typing import Any import numpy as np import torch import torch.nn.functional as F -from torch import nn from transformers.feature_extraction_utils import BatchFeature from vllm.config import VllmConfig -from vllm.model_executor.layers.pooler import Pooler, PoolingParamsUpdate -from vllm.model_executor.layers.pooler.tokwise.methods import AllPool from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.model_executor.model_loader.default_loader import DefaultModelLoader from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.tasks import PoolingTask -from vllm.utils.torch_utils import async_tensor_h2d -from vllm.v1.outputs import PoolerOutput -from vllm.v1.pool.metadata import PoolingMetadata +from .clef_joint_schema import JOINT_HEAD_WEIGHTS, JointSchemaHead, JointSchemaPooler from .interfaces_base import default_pooling_type from .qwen3_omni_moe_thinker import ( Qwen3OmniMoeThinkerDummyInputsBuilder, @@ -50,314 +25,6 @@ Qwen3OmniMoeThinkerProcessingInfo, ) -JOINT_HEAD_WEIGHTS = "joint_head.safetensors" -NUM_QUESTION_TYPES = 3 -MAX_LOGIT_SCALE = math.log(100.0) - - -def segment_mean( - values: torch.Tensor, segment_ids: torch.Tensor, lengths: torch.Tensor -) -> torch.Tensor: - sums = values.new_zeros(lengths.shape[0], values.shape[-1], dtype=torch.float32) - sums.index_add_(0, segment_ids, values.float()) - return (sums / lengths[:, None]).to(values.dtype) - - -@dataclass -class SchemaIndices: - """Where one prompt's questions and options are, on the model's device.""" - - positions: torch.Tensor - """Every token of every span, questions first.""" - span_ids: torch.Tensor - """The span of each of ``positions``.""" - span_lengths: torch.Tensor - """``[Q + O]``.""" - question_types: torch.Tensor - """``[Q]``.""" - option_question: torch.Tensor - """The question of each option, ``[O]``.""" - option_slots: torch.Tensor - """Each option's flat index in a ``[Q, max_options]`` grid, ``[O]``.""" - option_mask: torch.Tensor - """``[Q, max_options]``, true at the slots of options.""" - - -class EvidenceRoutingLayer(nn.Module): - def __init__(self, width: int, heads: int, feedforward: int) -> None: - super().__init__() - self.query_norm = nn.LayerNorm(width) - self.memory_norm = nn.LayerNorm(width) - self.attention = nn.MultiheadAttention(width, heads, batch_first=True) - self.feedforward_norm = nn.LayerNorm(width) - # The Identity takes the place of the checkpoint's dropout. - self.feedforward = nn.Sequential( - nn.Linear(width, feedforward), - nn.GELU(), - nn.Identity(), - nn.Linear(feedforward, width), - ) - - def forward(self, queries: torch.Tensor, memory: torch.Tensor) -> torch.Tensor: - # Separate key and value tensors, like the reference, keep attention's - # unpacked input projections. - routed, _ = self.attention( - self.query_norm(queries), - self.memory_norm(memory), - self.memory_norm(memory), - need_weights=False, - ) - queries = queries + routed - return queries + self.feedforward(self.feedforward_norm(queries)) - - -class JointSchemaHead(nn.Module): - def __init__( - self, - hidden_size: int, - width: int, - routing_layers: int, - layers: int, - heads: int, - feedforward: int, - ) -> None: - super().__init__() - self.hidden_norm = nn.LayerNorm(hidden_size) - self.memory_projection = nn.Linear(hidden_size, width, bias=False) - self.question_projection = nn.Linear(hidden_size, width, bias=False) - self.option_question_projection = nn.Linear(hidden_size, width, bias=False) - self.global_projection = nn.Linear(hidden_size, width, bias=False) - self.option_context_projection = nn.Linear(hidden_size, width, bias=False) - self.option_lexical_projection = nn.Linear(hidden_size, width, bias=False) - self.type_embedding = nn.Embedding(NUM_QUESTION_TYPES, width) - self.evidence_layers = nn.ModuleList( - EvidenceRoutingLayer(width, heads, feedforward) - for _ in range(routing_layers) - ) - self.option_summary_norm = nn.LayerNorm(width) - self.layers = nn.ModuleList( - nn.TransformerDecoderLayer( - d_model=width, - nhead=heads, - dim_feedforward=feedforward, - dropout=0.0, - activation="gelu", - batch_first=True, - norm_first=True, - ) - for _ in range(layers) - ) - self.field_norm = nn.LayerNorm(width) - self.option_norm = nn.LayerNorm(width) - self.residual_scorer = nn.Sequential( - nn.Linear(width * 4, width), - nn.GELU(), - nn.Identity(), - nn.Linear(width, 1), - ) - self.prior_logit_scale = nn.Parameter(torch.zeros(())) - self.joint_logit_scale = nn.Parameter(torch.zeros(())) - self.residual_gate = nn.Parameter(torch.zeros(())) - - def forward( - self, - hidden_states: torch.Tensor, - lexical: torch.Tensor, - schema: SchemaIndices, - ) -> torch.Tensor: - """The ``[O]`` option logits of one prompt, from its ``[L, hidden]`` - final hidden states and each option's ``[O, hidden]`` mean output - embedding.""" - num_questions = schema.question_types.shape[0] - option_question = schema.option_question - hidden_states = self.hidden_norm(hidden_states) - memory = self.memory_projection(hidden_states).unsqueeze(0) - global_vector = hidden_states[-1] - means = segment_mean( - hidden_states[schema.positions], schema.span_ids, schema.span_lengths - ) - question_vectors, contexts = means[:num_questions], means[num_questions:] - - options = ( - self.option_context_projection(contexts) - + self.option_lexical_projection(lexical) - + self.option_question_projection(question_vectors)[option_question] - ).unsqueeze(0) - for layer in self.evidence_layers: - options = layer(options, memory) - options = options[0] - - # Each field summarizes its options, weighted by their match to it. - base_fields = self.question_projection(question_vectors) - width = options.shape[-1] - grid = options.new_zeros(schema.option_mask.numel(), width) - grid[schema.option_slots] = options - grid = grid.view(*schema.option_mask.shape, width) - routing = torch.einsum("qow,qw->qo", grid, base_fields) / math.sqrt(width) - routing = routing.masked_fill(~schema.option_mask, float("-inf")) - summaries = torch.einsum("qo,qow->qw", routing.softmax(dim=-1), grid) - fields = ( - base_fields - + self.option_summary_norm(summaries) - + self.global_projection(global_vector).unsqueeze(0) - + self.type_embedding(schema.question_types) - ).unsqueeze(0) - for layer in self.layers: - fields = layer(fields, memory) - fields = self.field_norm(fields[0])[option_question] - - anchors = F.normalize(question_vectors + global_vector, dim=-1) - prior_scale = self.prior_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() - prior = prior_scale * ( - F.normalize(lexical, dim=-1) * anchors[option_question] - ).sum(dim=-1) - options = self.option_norm(options) - cosine = F.cosine_similarity(fields, options, dim=-1) - features = torch.cat( - [fields, options, fields * options, torch.abs(fields - options)], dim=-1 - ) - residual = self.residual_scorer(features).squeeze(-1) - joint_scale = self.joint_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() - joint = joint_scale * cosine + residual - return prior + torch.sigmoid(self.residual_gate) * joint - - -def parse_joint_schema( - extra_kwargs: dict[str, Any] | None, prompt_len: int -) -> tuple[list[int], list[tuple[int, int]], list[int]] | None: - """The question types, the absolute spans of the questions then of the - options, and the option counts. None if the schema is missing or invalid.""" - schema = (extra_kwargs or {}).get("joint_schema") - if not isinstance(schema, dict): - return None - try: - types = [int(t) for t in schema["question_types"]] - counts = [int(c) for c in schema["option_counts"]] - question_spans = [ - (int(start) + prompt_len, int(end) + prompt_len) - for start, end in schema["question_spans"] - ] - option_spans = [ - (int(start) + prompt_len, int(end) + prompt_len) - for start, end in schema["option_spans"] - ] - except (KeyError, TypeError, ValueError): - return None - spans = question_spans + option_spans - if ( - not types - or len(counts) != len(types) - or len(question_spans) != len(types) - or len(option_spans) != sum(counts) - or any(not 0 <= t < NUM_QUESTION_TYPES for t in types) - or any(c < 1 for c in counts) - or any(not 0 <= start < end <= prompt_len for start, end in spans) - ): - return None - return types, spans, counts - - -class JointSchemaPooler(Pooler): - """Runs the joint schema head on each finished prompt. - - It never raises, because it runs in the engine core: a prompt without a - valid schema, such as a profiling dummy, gets an empty output. - """ - - def __init__( - self, - head: JointSchemaHead, - embed_lexical: Callable[[torch.Tensor], torch.Tensor], - ) -> None: - super().__init__() - self.pooling = AllPool() - self.head = head - self.embed_lexical = embed_lexical - - def get_supported_tasks(self) -> Set[PoolingTask]: - return {"token_classify"} - - def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate: - return PoolingParamsUpdate(requires_token_ids=True) - - def forward( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> PoolerOutput: - pooled = self.pooling(hidden_states, pooling_metadata) - prompt_token_ids = pooling_metadata.get_prompt_token_ids_cpu() - outputs: list[torch.Tensor | None] = [] - for hidden, token_ids, params in zip( - pooled, prompt_token_ids, pooling_metadata.pooling_params - ): - if hidden is None: - outputs.append(None) - continue - parsed = parse_joint_schema(params.extra_kwargs, hidden.shape[0]) - if parsed is None: - outputs.append(hidden.new_empty(0, dtype=torch.float32)) - continue - schema, lexical_ids, lexical_span_ids = self._indices( - *parsed, hidden.device - ) - lexical = segment_mean( - self.embed_lexical( - async_tensor_h2d(token_ids[lexical_ids], hidden.device) - ), - lexical_span_ids, - schema.span_lengths[len(parsed[0]) :], - ) - logits = self.head(hidden, lexical, schema).float() - if params.use_activation is not False: - mask = schema.option_mask - grid = logits.new_full((mask.numel(),), float("-inf")) - grid[schema.option_slots] = logits - probs = grid.view(mask.shape).softmax(dim=-1) - logits = probs.view(-1)[schema.option_slots] - outputs.append(logits) - return outputs - - @staticmethod - def _indices( - types: list[int], - spans: list[tuple[int, int]], - counts: list[int], - device: torch.device, - ) -> tuple[SchemaIndices, torch.Tensor, torch.Tensor]: - """The schema's indices, plus the option tokens' positions (on the CPU) - and their option spans, for the lexical embeddings.""" - num_questions = len(types) - lengths = torch.tensor([end - start for start, end in spans]) - span_ids = torch.repeat_interleave(torch.arange(len(spans)), lengths) - positions = torch.cat([torch.arange(start, end) for start, end in spans]) - num_question_tokens = int(lengths[:num_questions].sum()) - - counts_cpu = torch.tensor(counts) - max_options = max(counts) - option_question = torch.repeat_interleave( - torch.arange(num_questions), counts_cpu - ) - option_slots = torch.cat( - [torch.arange(count) + q * max_options for q, count in enumerate(counts)] - ) - schema = SchemaIndices( - positions=async_tensor_h2d(positions, device), - span_ids=async_tensor_h2d(span_ids, device), - span_lengths=async_tensor_h2d(lengths, device), - question_types=async_tensor_h2d(types, device, dtype=torch.long), - option_question=async_tensor_h2d(option_question, device), - option_slots=async_tensor_h2d(option_slots, device), - option_mask=async_tensor_h2d( - torch.arange(max_options) < counts_cpu[:, None], device - ), - ) - lexical_span_ids = async_tensor_h2d( - span_ids[num_question_tokens:] - num_questions, device - ) - return schema, positions[num_question_tokens:], lexical_span_ids - - MAX_FRAME_PIXELS = 256 * 32 * 32 diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 36f9bd6e6fcc..cd0557ac4585 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -287,6 +287,7 @@ _TOKEN_CLASSIFICATION_MODELS = { "BertForTokenClassification": ("bert", "BertForTokenClassification"), + "ClefFlashForJointSchema": ("clef_flash", "ClefFlashForJointSchema"), "ClefOmniForJointSchema": ("clef_omni", "ClefOmniForJointSchema"), "ModernBertForTokenClassification": ( "modernbert", diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 147b65759dea..06a937697106 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -647,6 +647,27 @@ def __init__( _strip_mimo_v2_attention_chunk_size(hf_config, hf_text_config) +class Qwen3_5ModelArchConfigConvertor(ModelArchConfigConvertorBase): + def __init__( + self, + hf_config: PreTrainedConfig, + hf_text_config: PreTrainedConfig, + revision: str | None = None, + ): + if getattr(hf_config, "architectures", None) == [ + "Qwen3_5ForConditionalGeneration" + ]: + model_path = hf_config.name_or_path + if model_path and file_or_path_exists( + model_path, "joint_head_config.json", revision + ): + hf_config.architectures = ["ClefFlashForJointSchema"] + hf_config.joint_head_config = get_hf_file_to_dict( + "joint_head_config.json", model_path, revision + ) + super().__init__(hf_config, hf_text_config, revision) + + class Qwen3OmniMoeModelArchConfigConvertor(ModelArchConfigConvertorBase): def __init__( self, @@ -850,6 +871,7 @@ def derive_max_model_len_and_key(self) -> tuple[float, str | None]: "nemotron-nas": NemotronNasModelArchConfigConvertor, "bailing_hybrid_v3_mtp": BailingHybridV3MTPModelArchConfigConvertor, "pangu_ultra_moe_mtp": PanguUltraMoeMTPModelArchConfigConvertor, + "qwen3_5": Qwen3_5ModelArchConfigConvertor, "qwen3_5_mtp": Qwen3_5MTPModelArchConfigConvertor, "qwen3_omni_moe": Qwen3OmniMoeModelArchConfigConvertor, "qwen3_next_mtp": Qwen3NextMTPModelArchConfigConvertor, From 86b1c2ef159dcc268d219b074cb047340ea718aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Tue, 13 Oct 2026 07:51:22 +0200 Subject: [PATCH 360/624] [Attention][DCP] Support a2a DCP comm backend for FlashMLA sparse (#61285) Signed-off-by: NickLucche Co-authored-by: Claude Opus 5.5 --- tests/v1/attention/test_sparse_mla_backends.py | 16 ++++++++++++++++ .../v1/attention/backends/mla/flashmla_sparse.py | 8 +------- 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index c21fc15a5278..f7c47fe6a3ec 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -4157,6 +4157,22 @@ def test_fp8_dcp_head_envelope_guard(local_heads, dcp_world_size, should_raise): assert local_pad == gathered_pad +@pytest.mark.skipif( + torch.cuda.get_device_capability() < (9, 0), + reason="FlashMLASparseBackend requires CUDA 9.0 or higher", +) +@pytest.mark.parametrize("comm_backend", ["ag_rs", "a2a"]) +def test_fp8_dcp_accepts_comm_backend(comm_backend): + """The mixed-batch path returns (out, lse) for every token in the layout + both DCP combines consume, so neither comm backend is rejected.""" + device = torch.device(DEVICE_TYPE) + vllm_config = _build_sparse_dcp_vllm_config(16, 4, comm_backend) + kv_cache_spec = create_standard_kv_cache_spec(vllm_config) + builder_cls = FlashMLASparseBackend.get_builder_cls() + builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device) + assert builder.fp8_use_mixed_batch + + def test_fp8_mixed_batch_dcp_neutralizes_empty_rows(monkeypatch): """A decode row whose top-k shard holds no local candidates (all -1) has undefined kernel out/lse; it must come back as (0, -inf), the identity of diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 3cf0c33c534d..80b8da179769 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -446,12 +446,6 @@ def __init__( ) if parallel_config.decode_context_parallel_size > 1: - if parallel_config.dcp_comm_backend != "ag_rs": - raise NotImplementedError( - "DCP for FlashMLA sparse is only validated with the " - "default 'ag_rs' DCP comm backend; got " - f"'{parallel_config.dcp_comm_backend}'" - ) if self.pcp_dcp_kv_gather and cache_config.cache_dtype != "fp8_ds_mla": raise NotImplementedError( "PCP+DCP sparse prefill gathers the KV through the fp8_ds_mla " @@ -1268,7 +1262,7 @@ def _forward_fp8_kv_mixed_batch( out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0) lse.masked_fill_(empty_rows.view(-1, 1), float("-inf")) # The head-padding slice above can leave `out` non-contiguous, and the - # merge feeds it to reduce_scatter. + # DCP combines (reduce_scatter, direct A2A) need packed heads. return out.contiguous(), lse def _fp8_flash_mla_kernel( From a6009ea85be08161f4b5954bbf317027f6ddfd82 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Mon, 12 Oct 2026 23:24:43 -0700 Subject: [PATCH 361/624] [Kernel][Build] Bump DeepGEMM to the 26/09/30 release and adapt the sparse indexer (#61149) Signed-off-by: Yongye Zhu Co-authored-by: Claude Opus 5.5 (1M context) --- cmake/external_projects/deepgemm.cmake | 12 +- csrc/cache_kernels.cu | 20 +- .../attention/test_deepgemm_attention.py | 189 +++++++++++------- .../test_rocm_paged_mxfp4_indexer.py | 1 + tests/kernels/test_compressor_kv_cache.py | 42 +++- .../test_fused_deepseek_v32_norm_rope.py | 15 +- .../test_fused_indexer_q_rope_quant.py | 5 + .../test_kpool_decode_update_batched.py | 8 +- .../layers/test_mla_short_prefill_indexer.py | 39 ++-- .../test_indexer_deepseek_v4_slot_mapping.py | 11 +- .../test_rocm_paged_mxfp4_indexer_plan.py | 1 + tools/build_deepgemm_C.py | 3 +- tools/check_wheel_deepgemm.py | 11 +- tools/install_deepgemm.sh | 8 +- .../layers/sparse_attn_indexer.py | 110 +++++++--- vllm/models/deepseek_v32/common/kernels.py | 14 +- .../common/ops/fused_compress_quant_cache.py | 12 +- .../common/ops/indexer_k_store.py | 15 +- .../glm5next/nvidia/ops/kpool_compress.py | 30 ++- vllm/models/glm5next/nvidia/sparse_indexer.py | 72 ++++--- vllm/utils/deep_gemm.py | 101 +++++++--- vllm/v1/attention/backends/mla/indexer.py | 4 + vllm/v1/attention/ops/common.py | 7 + 23 files changed, 519 insertions(+), 211 deletions(-) diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 1dcee345f845..b796623b7ac1 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -29,12 +29,11 @@ if(DEEPGEMM_SRC_DIR) else() # Keep in sync with tools/install_deepgemm.sh set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git") - # Pinned to the fork's dev branch: upstream 2.8.0 plus the SM120 - # port, the SM90 paged-MQA kv_block=32/next_n=4 port, configurable SwiGLU - # alpha/beta, and SiTU for FP8/FP4 Mega MoE. Also includes the CUDA 12.x - # layout header fix from vllm-project/DeepGEMM#12. - # TORCH_LIBRARY build with a CPython 3.10 stable-ABI floor. - set(_DEEPGEMM_UPSTREAM_TAG "1e1842a833699298f7afc02eefb2ed168fff6938") + # Pinned to the fork's dev branch: upstream's 26/09/30 release (2.8.1) + # plus the SM120 port, the SM90 paged-MQA kv_block=32/next_n=4 port, + # configurable SwiGLU alpha/beta, SiTU for FP8/FP4 Mega MoE, and NVFP4 + # Mega MoE. TORCH_LIBRARY build with a CPython 3.10 stable-ABI floor. + set(_DEEPGEMM_UPSTREAM_TAG "6ce1c5ec6627a6bf11dbcc698b0b1e19c9248432") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) @@ -146,6 +145,7 @@ if(DEEPGEMM_ARCHS) # install(FILES "${deepgemm_SOURCE_DIR}/deep_gemm/__init__.py" + "${deepgemm_SOURCE_DIR}/deep_gemm/locality_domain.py" DESTINATION vllm/third_party/deep_gemm COMPONENT _deep_gemm_C) diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu index d1c1e9df9896..24453ff1c154 100644 --- a/csrc/cache_kernels.cu +++ b/csrc/cache_kernels.cu @@ -630,7 +630,8 @@ __global__ void indexer_k_quant_and_cache_kernel( const int cache_block_size, // cache block size const int64_t cache_block_stride, // stride for each block in kv_cache - const bool use_ue8m0 // use ue8m0 scale format + const bool use_ue8m0, // use ue8m0 scale format + const bool use_mx_scale // store the ue8m0 scale as packed MX bytes ) { constexpr int VEC_SIZE = 4; const int64_t token_idx = blockIdx.x; @@ -679,7 +680,13 @@ __global__ void indexer_k_quant_and_cache_kernel( const int64_t dst_scale_idx = block_idx * cache_block_stride + cache_block_size * head_dim + (block_offset * head_dim + head_dim_idx) * 4 / quant_block_size; - reinterpret_cast(kv_cache)[dst_scale_idx / 4] = scale; + if (use_mx_scale) { + // One UE8M0 byte per 32 elements, all equal to the shared scale's. + reinterpret_cast(kv_cache)[dst_scale_idx / 4] = + ((__float_as_uint(scale) >> 23) & 0xFF) * 0x01010101u; + } else { + reinterpret_cast(kv_cache)[dst_scale_idx / 4] = scale; + } } } @@ -1912,7 +1919,7 @@ void cp_gather_and_upconvert_nvfp4_kv_cache( reinterpret_cast(k.data_ptr()), \ reinterpret_cast(kv_cache.data_ptr()), \ slot_mapping.const_data_ptr(), head_dim, quant_block_size, \ - cache_block_size, cache_block_stride, use_ue8m0); + cache_block_size, cache_block_stride, use_ue8m0, use_mx_scale); void indexer_k_quant_and_cache( torch::stable::Tensor& k, // [num_tokens, head_dim] @@ -1924,8 +1931,13 @@ void indexer_k_quant_and_cache( int head_dim = k.size(1); int cache_block_size = kv_cache.size(1); int64_t cache_block_stride = kv_cache.stride(0); - bool use_ue8m0 = scale_fmt == "ue8m0"; + // "mxfp8" stores the ue8m0 scale as four packed UE8M0 bytes, the MX scales + // DeepGEMM's SM100 MQA logits kernels read. + bool use_mx_scale = scale_fmt == "mxfp8"; + bool use_ue8m0 = use_mx_scale || scale_fmt == "ue8m0"; + STD_TORCH_CHECK(!use_mx_scale || (head_dim == 128 && quant_block_size == 128), + "mxfp8 scales require head_dim == quant_block_size == 128"); STD_TORCH_CHECK(k.device() == kv_cache.device(), "k and kv_cache must be on the same device"); STD_TORCH_CHECK(k.device() == slot_mapping.device(), diff --git a/tests/kernels/attention/test_deepgemm_attention.py b/tests/kernels/attention/test_deepgemm_attention.py index f332fb0fea21..109cc83830b1 100644 --- a/tests/kernels/attention/test_deepgemm_attention.py +++ b/tests/kernels/attention/test_deepgemm_attention.py @@ -5,6 +5,7 @@ import pytest import torch +from vllm.model_executor.layers.sparse_attn_indexer import as_fp32_logits from vllm.platforms import current_platform from vllm.utils.deep_gemm import ( _ceil_to_ue8m0, @@ -13,18 +14,29 @@ fp8_fp4_paged_mqa_logits, get_num_sms, get_paged_mqa_logits_metadata, + mqa_logits_mx_inputs, + mqa_logits_use_mx_scales, native_next_n_supported, ) from vllm.utils.import_utils import has_deep_gemm from vllm.utils.math_utils import cdiv +def encode_indexer_scales(sf: torch.Tensor) -> torch.Tensor: + """Per-token power-of-two fp32 scales as the indexer K cache stores them: + four equal packed UE8M0 bytes under MX scales, else the fp32 itself.""" + if not mqa_logits_use_mx_scales(): + return sf + ue8m0 = (sf.view(torch.int32) >> 23).to(torch.uint8) + return ue8m0[..., None].repeat_interleave(4, dim=-1).view(torch.int32)[..., 0] + + def kv_cache_cast_to_fp8(x: torch.Tensor) -> torch.Tensor: # x: (num_blocks, block_size, 1, head_dim) num_blocks, block_size, num_heads, head_dim = x.shape assert num_heads == 1 x_amax = x.abs().float().amax(dim=3, keepdim=True).clamp(1e-4) - sf = x_amax / 448.0 + sf = _ceil_to_ue8m0(x_amax / 448.0) x_scaled = (x * (1.0 / sf)).to(torch.float8_e4m3fn) x_fp8 = torch.empty( (num_blocks, block_size * (head_dim + 4)), @@ -34,9 +46,9 @@ def kv_cache_cast_to_fp8(x: torch.Tensor) -> torch.Tensor: x_fp8[:, : block_size * head_dim] = x_scaled.view( num_blocks, block_size * head_dim ).view(dtype=torch.uint8) - x_fp8[:, block_size * head_dim :] = sf.view(num_blocks, block_size).view( - dtype=torch.uint8 - ) + x_fp8[:, block_size * head_dim :] = encode_indexer_scales( + sf.view(num_blocks, block_size) + ).view(dtype=torch.uint8) return x_fp8.view(num_blocks, block_size, num_heads, head_dim + 4) @@ -91,64 +103,72 @@ def _ref_fp8_mqa_logits( return logits +def _packed_requests_test_data(seq_len: int, seq_len_kv: int): + # Two causal prefill requests packed into one KV buffer, so the second + # request's rows start at a nonzero K position. + half, half_kv = seq_len // 2, seq_len_kv // 2 + rows = torch.arange(seq_len, dtype=torch.int, device="cuda") + ks = torch.where(rows < half, 0, half_kv).to(torch.int) + ke = ks + half_kv - half + 1 + rows % half + return ks, ke + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @pytest.mark.skipif( not current_platform.has_device_capability(90), reason="SM90 and SM100 only" ) -@pytest.mark.parametrize("clean_logits", [True, False]) -def test_deepgemm_fp8_mqa_logits(clean_logits: bool): +@pytest.mark.parametrize("layout", ["cp", "no_cp", "packed"]) +def test_deepgemm_fp8_mqa_logits(layout: str): + """FP8 indexer prefill logits, compressed to each row's valid K span.""" torch.manual_seed(0) random.seed(0) num_heads, head_dim = 32, 128 - for seq_len in (512,): - for seq_len_kv in (1024,): - for disable_cp in (False, True): - q = torch.randn( - seq_len, - num_heads, - head_dim, - device="cuda", - dtype=torch.bfloat16, - ) - kv = torch.randn( - seq_len_kv, head_dim, device="cuda", dtype=torch.bfloat16 - ) - weights = torch.randn( - seq_len, num_heads, device="cuda", dtype=torch.float32 - ) - - if disable_cp: - ks = torch.zeros(seq_len, dtype=torch.int, device="cuda") - ke = torch.arange(seq_len, dtype=torch.int, device="cuda") + ( - seq_len_kv - seq_len - ) - else: - ks, ke = _generate_cp_test_data(seq_len, seq_len_kv) - - q_fp8 = q.to(torch.float8_e4m3fn) - kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0,), False) - logits = fp8_fp4_mqa_logits( - (q_fp8, None), kv_fp8, weights, ks, ke, clean_logits=clean_logits - ) - - ref_logits = _ref_fp8_mqa_logits( - q=q, - kv=kv, - weights=weights, - cu_seqlen_ks=ks, - cu_seqlen_ke=ke, - ) - ref_neginf_mask = ref_logits == float("-inf") - - if clean_logits: - neginf_mask = logits == float("-inf") - assert torch.equal(neginf_mask, ref_neginf_mask) - - ref_logits = ref_logits.masked_fill(ref_neginf_mask, 0) - logits = logits.masked_fill(ref_neginf_mask, 0) - diff = calc_diff(logits, ref_logits) - assert diff < 1e-3, f"{diff=}" + seq_len, seq_len_kv = 512, 1024 + q = torch.randn(seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + kv = torch.randn(seq_len_kv, head_dim, device="cuda", dtype=torch.bfloat16) + weights = torch.randn(seq_len, num_heads, device="cuda", dtype=torch.float32) + + if layout == "cp": + ks, ke = _generate_cp_test_data(seq_len, seq_len_kv) + elif layout == "no_cp": + ks = torch.zeros(seq_len, dtype=torch.int, device="cuda") + ke = torch.arange(seq_len, dtype=torch.int, device="cuda") + ( + seq_len_kv - seq_len + ) + else: + ks, ke = _packed_requests_test_data(seq_len, seq_len_kv) + + q_fp8 = q.to(torch.float8_e4m3fn) + kv_fp8, kv_sf = per_custom_dims_cast_to_fp8(kv, (0,), True) + q_scale, kernel_weights = mqa_logits_mx_inputs(q_fp8, None, weights) + max_seqlen_k = int((ke - ks).max()) + logits = fp8_fp4_mqa_logits( + (q_fp8, q_scale), + (kv_fp8, encode_indexer_scales(kv_sf)), + kernel_weights, + ks, + ke, + max_seqlen_k, + ) + assert logits.shape == (seq_len, max_seqlen_k) + logits = as_fp32_logits(logits, ke - ks) + + ref_logits = _ref_fp8_mqa_logits( + q=q, + kv=kv, + weights=kernel_weights.float(), + cu_seqlen_ks=ks, + cu_seqlen_ke=ke, + ) + # Compressed row i holds K positions [ks[i], ke[i]) from column zero. + cols = ks[:, None] + torch.arange(max_seqlen_k, device="cuda")[None, :] + valid = cols < ke[:, None] + ref_logits = ref_logits.gather(1, cols.clamp(max=seq_len_kv - 1).long()) + ref_logits = ref_logits.masked_fill(~valid, 0) + logits = logits.masked_fill(~valid, 0) + diff = calc_diff(logits, ref_logits) + assert diff < 1e-3, f"{diff=}" def _ref_fp8_fp4_paged_mqa_logits( @@ -212,9 +232,6 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int): if not native_next_n_supported(next_n): pytest.skip(f"next_n={next_n} has no native kernel on this architecture") - # NOTE: clean_logits=True is incompatible with the 2D context_lens - # required by csrc/apis/attention.hpp; only the False path is exercised. - clean_logits = False torch.manual_seed(0) random.seed(0) @@ -273,26 +290,42 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int): context_lens_2d = ( context_lens.unsqueeze(-1) - next_n + 1 + next_n_arange ).contiguous() + q_scale, kernel_weights = mqa_logits_mx_inputs(q_fp8, None, weights) + kernel_q = (q_fp8, q_scale) + kernel_context_lens, kernel_block_tables = context_lens_2d, block_tables + indices = None + if mqa_logits_use_mx_scales(): + # SM100 only takes varlen rows: one Q token each, tied to its + # request by `indices`. + assert q_scale is not None + kernel_q = ( + q_fp8.flatten(0, 1)[:, None], + q_scale.flatten(0, 1)[:, None], + ) + kernel_context_lens = context_lens_2d.view(-1, 1) + kernel_block_tables = block_tables.repeat_interleave(next_n, dim=0) + indices = torch.arange(batch_size, device="cuda", dtype=torch.int32) + indices = indices.repeat_interleave(next_n) schedule_metadata = get_paged_mqa_logits_metadata( - context_lens_2d, - blocksize, - get_num_sms(), + kernel_context_lens, blocksize, get_num_sms(), indices=indices ) logits = fp8_fp4_paged_mqa_logits( - (q_fp8, None), + kernel_q, kv_cache_fp8, - weights, - context_lens_2d, - block_tables, + kernel_weights, + kernel_context_lens, + kernel_block_tables, schedule_metadata, max_model_len, - clean_logits=clean_logits, + indices=indices, ) + if indices is not None: + logits = as_fp32_logits(logits, kernel_context_lens.view(-1)) ref_logits = _ref_fp8_fp4_paged_mqa_logits( q, kv_cache, - weights, + kernel_weights.float(), context_lens, block_tables, max_model_len, @@ -309,7 +342,7 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int): context_lens[row_indices] - next_n + next_n_offset ).unsqueeze(1) - logits = logits.masked_fill(~mask, 0) + logits = logits.float().masked_fill(~mask, 0) ref_logits = ref_logits.masked_fill(~mask, 0) diff = calc_diff(logits, ref_logits) assert diff < 1e-3, f"{diff=}" @@ -349,12 +382,28 @@ def test_deepgemm_paged_mqa_packed_manager_block_stride(): q = torch.randn(1, 1, num_heads, head_dim, device="cuda", dtype=torch.bfloat16).to( torch.float8_e4m3fn ) - weights = torch.randn(1, num_heads, device="cuda", dtype=torch.float32) + q_scale, weights = mqa_logits_mx_inputs( + q, None, torch.randn(1, num_heads, device="cuda", dtype=torch.float32) + ) context_lens = torch.full((1, 1), block_size, device="cuda", dtype=torch.int32) - schedule = get_paged_mqa_logits_metadata(context_lens, page_size, get_num_sms()) + indices = ( + torch.zeros(1, device="cuda", dtype=torch.int32) + if mqa_logits_use_mx_scales() + else None + ) + schedule = get_paged_mqa_logits_metadata( + context_lens, page_size, get_num_sms(), indices=indices + ) logits = [ fp8_fp4_paged_mqa_logits( - (q, None), cache, weights, context_lens, table, schedule, block_size, False + (q, q_scale), + cache, + weights, + context_lens, + table, + schedule, + block_size, + indices=indices, ) for cache, table in zip((compact_pages, packed_pages), tables) ] diff --git a/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py b/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py index 311ee89a8d12..e270ae4a5813 100644 --- a/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py +++ b/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py @@ -253,6 +253,7 @@ def _prefill_metadata( cu_seq_lens=torch.zeros(1, dtype=torch.int32, device=DEVICE), token_to_seq=torch.zeros(1, dtype=torch.int32, device=DEVICE), total_seq_lens=1, + max_seqlen_k=int(ends.max()), token_start=t0, token_end=t1, num_reqs=req_hi - req_lo, diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index cc4f4c72a931..e9924229055a 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -1422,7 +1422,10 @@ def test_dequantize_and_gather_k_cache( @pytest.mark.parametrize("num_tokens", [1, 4, 8, 17]) @pytest.mark.parametrize("block_size", [16, 64]) -def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int): +@pytest.mark.parametrize("scale_fmt", ["ue8m0", "mxfp8"]) +def test_indexer_quant_cache_roundtrip( + num_tokens: int, block_size: int, scale_fmt: str +): """K → indexer_k_quant_and_cache → cp_gather_indexer_k_quant_cache → manual dequant → compare against original.""" HEAD_DIM = 128 @@ -1442,7 +1445,9 @@ def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int): ) slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) - ops.indexer_k_quant_and_cache(k, kv_cache, slot_mapping, QUANT_BLOCK_SIZE, "ue8m0") + ops.indexer_k_quant_and_cache( + k, kv_cache, slot_mapping, QUANT_BLOCK_SIZE, scale_fmt + ) # ── Gather ────────────────────────────────────────────────────────── max_blocks_per_seq = num_blocks @@ -1465,7 +1470,12 @@ def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int): # ── Manual dequant ────────────────────────────────────────────────── k_fp8 = dst_k.view(torch.float8_e4m3fn).float() # [num_tokens, 128] - scale = dst_scale.view(torch.float32) # [num_tokens, 1] + if scale_fmt == "mxfp8": + # One UE8M0 byte per 32 elements, all equal to the token's scale. + assert (dst_scale == dst_scale[:, :1]).all() + scale = (dst_scale[:, :1].to(torch.int32) << 23).view(torch.float32) + else: + scale = dst_scale.view(torch.float32) # [num_tokens, 1] k_recovered = k_fp8 * scale # [num_tokens, 128] # ── Compare ───────────────────────────────────────────────────────── @@ -1749,8 +1759,13 @@ def _reference_kv_compress_norm_rope( ), ], ) -def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: bool): +@pytest.mark.parametrize("mx_scale", [False, True]) +def test_fused_kv_insert_indexer( + num_tokens: int, kv_block_size: int, use_fp4: bool, mx_scale: bool +): """Fused K compress+norm+rope+quant+insert for the indexer KV cache.""" + if use_fp4 and mx_scale: + pytest.skip("MX scales only change the FP8 cache") HEAD_DIM = 128 ROPE_DIM = 64 BLOCK_SIZE = 16 @@ -1833,6 +1848,7 @@ def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: b TOKEN_STRIDE=TOKEN_STRIDE, SCALE_DIM=SCALE_DIM, KV_BLOCK_STRIDE=kv_cache.stride(0), + MX_SCALE=mx_scale, num_warps=1, ) @@ -1876,12 +1892,18 @@ def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: b ), f"token {i}: FP8 bytes differ" scale_off = kv_block_size * TOKEN_STRIDE + pos * SCALE_DIM - actual_scale = kv_cache[blk, scale_off : scale_off + SCALE_DIM].view( - torch.float32 - ) - assert torch.equal(actual_scale, scale[i : i + 1]), ( - f"token {i}: scale {actual_scale.item()} != {scale[i].item()}" - ) + actual_scale = kv_cache[blk, scale_off : scale_off + SCALE_DIM] + if mx_scale: + # The power-of-two scale's exponent as 4 equal UE8M0 bytes. + ue8m0 = (scale[i : i + 1].view(torch.int32) >> 23).to(torch.uint8) + assert torch.equal(actual_scale, ue8m0.repeat(SCALE_DIM)), ( + f"token {i}: mx scale {actual_scale.tolist()} != {ue8m0.item()}" + ) + else: + actual_scale = actual_scale.view(torch.float32) + assert torch.equal(actual_scale, scale[i : i + 1]), ( + f"token {i}: scale {actual_scale.item()} != {scale[i].item()}" + ) @pytest.mark.parametrize("compress_ratio", [4, 128]) diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index 63bbe0c95ff4..ec61b94abd5e 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -32,6 +32,7 @@ from vllm.models.deepseek_v32.common import kernels as K from vllm.platforms import current_platform +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales FP8 = torch.float8_e4m3fn FP8_MAX = 448.0 @@ -116,6 +117,16 @@ def rope( return out +def indexer_scales(scale_bytes: torch.Tensor) -> torch.Tensor: + """Per-token fp32 scales from the indexer cache's 4 scale bytes per token.""" + if not mqa_logits_use_mx_scales(): + return scale_bytes.view(torch.float32) + # MX scales: the power-of-two scale's exponent as 4 equal UE8M0 bytes. + ue8m0 = scale_bytes.view(-1, 4) + assert (ue8m0 == ue8m0[:, :1]).all(), "MX scale bytes differ within a token" + return (ue8m0[:, 0].to(torch.int32) << 23).view(torch.float32) + + def ue8m0_quant(vals: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Per-row (last dim) UE8M0 fp8 quant matching kernels._fp8_ue8m0_quantize.""" amax = vals.float().abs().amax(dim=-1, keepdim=True) @@ -243,7 +254,7 @@ def test_fused_norm_rope( q_ref, s_ref = ue8m0_quant(ik_ref) flat = idx_cache[0].reshape(-1) vals = flat[: bs * cfg.index_head_dim].view(FP8).reshape(bs, cfg.index_head_dim) - scales = flat[bs * cfg.index_head_dim :].view(torch.float32) + scales = indexer_scales(flat[bs * cfg.index_head_dim :]) assert_fp8(vals[:num_tokens], q_ref, "indexer-K fp8") torch.testing.assert_close(scales[:num_tokens], s_ref, rtol=0, atol=0) @@ -318,7 +329,7 @@ def test_fused_norm_rope_packed_indexer_block_stride(cfg: ModelConfig): .view(FP8) .view(block_size, cfg.index_head_dim) ) - scales = packed[block_size * cfg.index_head_dim :].view(torch.float32) + scales = indexer_scales(packed[block_size * cfg.index_head_dim :]) assert_fp8(values, q_ref, "packed indexer-K fp8") torch.testing.assert_close(scales, s_ref, rtol=0, atol=0) assert (backing[block_size * idx_row : packed_block_stride] == 0).all() diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index 0ce3a918513f..c6dfa11b0f13 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -26,6 +26,7 @@ ) from vllm.models.deepseek_v4.common.ops import fused_indexer_q_rope_quant from vllm.platforms import current_platform +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales from vllm.utils.import_utils import has_cutedsl HEAD_DIM = 128 @@ -347,6 +348,10 @@ def _indexer_k_reference( exponent = k.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-4) / 448 scale = exponent.log2().ceil().exp2() values = (k.float() / scale).clamp(-448, 448).to(torch.float8_e4m3fn) + if mqa_logits_use_mx_scales(): + # The power-of-two scale's exponent as 4 equal UE8M0 bytes. + ue8m0 = (scale.view(torch.int32) >> 23).to(torch.uint8) + return values.view(torch.uint8), ue8m0.repeat(1, 4) return values.view(torch.uint8), scale.view(torch.uint8) diff --git a/tests/kernels/test_kpool_decode_update_batched.py b/tests/kernels/test_kpool_decode_update_batched.py index 4accf527f3d1..5dd9eff5f10a 100644 --- a/tests/kernels/test_kpool_decode_update_batched.py +++ b/tests/kernels/test_kpool_decode_update_batched.py @@ -29,6 +29,7 @@ import torch from vllm.platforms import current_platform +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales if current_platform.is_rocm(): from vllm.models.glm5next.amd.ops.kpool_compress import ( @@ -206,7 +207,12 @@ def _torch_reference( k_off = page_base + loc_tok * HEAD_DIM + torch.arange(HEAD_DIM) s_off = page_base + k_region + loc_tok * 4 kv_flat[k_off] = quantized.view(torch.uint8) - kv_flat[s_off : s_off + 4] = scale.detach().reshape(1).view(torch.uint8) + scale_bytes = scale.detach().reshape(1).view(torch.uint8) + if mqa_logits_use_mx_scales(): + # The power-of-two scale's exponent as 4 equal UE8M0 bytes. + ue8m0 = scale.reshape(1).view(torch.int32) >> 23 + scale_bytes = ue8m0.to(torch.uint8).repeat(4) + kv_flat[s_off : s_off + 4] = scale_bytes # stash -- gated on the TOKEN-granular tail slot, not on pos_valid. # pos_valid keys off the pool-granular cache_loc, which is -1 for diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index ba4ff895cbf6..7ca4cff4040d 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -771,6 +771,8 @@ def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool): candidates = _make_candidates(lens, num_candidates, cbk).cuda() # Reference: dense bf16 logits masked to the candidates, then torch.topk. + # The logits are compressed: row r starts at K position ks[r]. + max_len = int(lens.max()) dense = deep_gemm.fp8_fp4_mqa_logits( ( q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2), @@ -780,19 +782,13 @@ def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool): weights, ks, ke, - clean_logits=False, - logits_dtype=torch.bfloat16, + max_len, ) - cols = torch.arange(total, device="cuda") - in_range = (cols - ks.unsqueeze(1) >= 0) & ( - cols - ks.unsqueeze(1) < lens.unsqueeze(1) - ) - masked = dense.masked_fill( - ~(in_range & _candidate_token_mask(candidates, cbk, total, ks)), - float("-inf"), - ).float() - ref_abs = _reference_topk(masked, topk) # packed-workspace positions - ref = torch.where(ref_abs >= 0, ref_abs - ks.unsqueeze(1), -1) + in_range = torch.arange(max_len, device="cuda") < lens.unsqueeze(1) + # A row's last candidate block may extend past max_len. + in_cand = _candidate_token_mask(candidates, cbk, max_len + cbk)[:, :max_len] + masked = dense.masked_fill(~(in_range & in_cand), float("-inf")).float() + ref = _reference_topk(masked, topk) out = torch.full((rows, topk), -1, dtype=torch.int32, device="cuda") scratch = _sparse_scratch(rows, num_candidates, topk) @@ -827,11 +823,11 @@ def run_sparse(kernel_metadata=None): graph.replay() else: kernel_metadata = run_sparse() - _assert_same_selection(dense, out, ref, row_starts=ks) + _assert_same_selection(dense, out, ref) out.fill_(-1) run_sparse(kernel_metadata) - _assert_same_selection(dense, out, ref, row_starts=ks) + _assert_same_selection(dense, out, ref) def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int): @@ -902,21 +898,22 @@ def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool candidates = _make_candidates(row_lens.flatten(), num_candidates, cbk).cuda() def reference_topk(): - context_lens_2d = row_lens.view(batch, next_n) + # SM100 paged MQA logits take varlen rows: one Q token per row. + context_lens_2d = row_lens.view(rows, 1) + indices = torch.arange(batch, dtype=torch.int32, device="cuda") + indices = indices.repeat_interleave(next_n) sched = deep_gemm.get_paged_mqa_logits_metadata( - context_lens_2d, page_kv, deep_gemm.get_num_sms(), None + context_lens_2d, page_kv, deep_gemm.get_num_sms(), indices ) logits = deep_gemm.fp8_fp4_paged_mqa_logits( - (q, q_scale), + (q.view(rows, 1, num_heads, -1), q_scale.view(rows, 1, num_heads)), kv_cache, weights, context_lens_2d, - block_table, + block_table.repeat_interleave(next_n, dim=0), sched, int(req_lens.max()), - False, - torch.bfloat16, - None, + indices, ) lens_flat = row_lens.flatten() cols = torch.arange(logits.shape[1], device="cuda") diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 464fce66f212..62c8bc32bff5 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -14,6 +14,7 @@ from vllm.models.deepseek_v41.sparse_mla import ( DeepseekV4SparseMLABackend as DeepseekV41SparseMLABackend, ) +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends.mla import indexer as indexer_module from vllm.v1.attention.backends.mla.compressor_utils import ( @@ -113,10 +114,17 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): num_reqs, 4 ) num_sms = torch.cuda.get_device_properties(device).multi_processor_count + # SM100's paged MQA logits only take varlen rows. + indices = ( + torch.arange(num_reqs, dtype=torch.int32, device=device) + if mqa_logits_use_mx_scales() + else None + ) schedule_metadata = indexer_module.get_paged_mqa_logits_metadata( torch.div(seq_lens, 16, rounding_mode="floor").unsqueeze(1), 64, num_sms, + indices=indices, ).clone() builder = object.__new__(DeepseekV32IndexerMetadataBuilder) @@ -139,6 +147,7 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): decode_lens=torch.zeros(num_reqs, dtype=torch.int32, device=device), requires_padding=False, schedule_metadata=schedule_metadata, + indices=indices, ), block_table=block_table, ) @@ -157,7 +166,7 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): torch.testing.assert_close( schedule_metadata, indexer_module.get_paged_mqa_logits_metadata( - expected_decode_seq_lens.unsqueeze(1), 64, num_sms + expected_decode_seq_lens.unsqueeze(1), 64, num_sms, indices=indices ), ) assert metadata.decode is not None diff --git a/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py b/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py index 0283ed99ee3a..a3dbec86dca5 100644 --- a/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py +++ b/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py @@ -36,6 +36,7 @@ def _chunk(block_table, token_start, token_end, row_ends): cu_seq_lens=torch.zeros(1, dtype=torch.int32), token_to_seq=torch.zeros(1, dtype=torch.int32), total_seq_lens=1, + max_seqlen_k=int(ends.max()), token_start=token_start, token_end=token_end, num_reqs=block_table.shape[0], diff --git a/tools/build_deepgemm_C.py b/tools/build_deepgemm_C.py index 36824318c051..dec8b0105ac8 100644 --- a/tools/build_deepgemm_C.py +++ b/tools/build_deepgemm_C.py @@ -35,8 +35,9 @@ env.pop("DG_SKIP_CUDA_BUILD", None) print(f"[build_deepgemm_C] building in {src} with {sys.executable}", flush=True) +# CMake decides when to rebuild; setuptools would only check python_api.cpp. subprocess.check_call( - [sys.executable, "setup.py", "build_ext", "--inplace"], + [sys.executable, "setup.py", "build_ext", "--inplace", "--force"], cwd=src, env=env, ) diff --git a/tools/check_wheel_deepgemm.py b/tools/check_wheel_deepgemm.py index fcc35ead58c4..c8da1351bb54 100644 --- a/tools/check_wheel_deepgemm.py +++ b/tools/check_wheel_deepgemm.py @@ -4,10 +4,12 @@ """Assert the vendored DeepGEMM package has the TORCH_LIBRARY binding layout. Expects ``deep_gemm/_C.py`` plus a single ``_C_extension*.so`` (abi3) under -``vllm.third_party.deep_gemm``. Run after vLLM is installed, e.g. the H100 -deepgemm kernel tests in .buildkite/test_areas/kernels.yaml. +``vllm.third_party.deep_gemm``, and that the package imports: vLLM treats an +import failure as DeepGEMM being absent. Run after vLLM is installed, e.g. the +H100 deepgemm kernel tests in .buildkite/test_areas/kernels.yaml. """ +import importlib import importlib.util import sys from pathlib import Path @@ -38,6 +40,11 @@ def main() -> int: if missing: print(f"missing: {missing}", file=sys.stderr) return 1 + try: + importlib.import_module("vllm.third_party.deep_gemm") + except Exception as e: + print(f"vllm.third_party.deep_gemm failed to import: {e!r}", file=sys.stderr) + return 1 return 0 diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index 2fb2351a0ca5..b617e6d5679f 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -7,10 +7,10 @@ set -e # Default values # Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" -# NOTE: This targets the vLLM fork's dev branch, which carries the sm120 -# and sm90 paged-MQA ports plus the SwiGLU alpha/beta and SiTU Mega MoE -# activations, plus the CUDA 12.x layout header fix from vllm-project/DeepGEMM#12. -DEEPGEMM_GIT_REF="1e1842a833699298f7afc02eefb2ed168fff6938" +# NOTE: This targets the vLLM fork's dev branch: upstream's 26/09/30 release +# plus the sm120 and sm90 paged-MQA ports, the SwiGLU alpha/beta and SiTU +# Mega MoE activations, and NVFP4 Mega MoE. +DEEPGEMM_GIT_REF="6ce1c5ec6627a6bf11dbcc698b0b1e19c9248432" WHEEL_DIR="" # Parse command line arguments diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 935e4c0d2d44..3af3a161ac8b 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -42,6 +42,8 @@ fp8_fp4_mqa_logits, fp8_fp4_paged_mqa_logits, has_deep_gemm, + mqa_logits_mx_inputs, + mqa_logits_use_mx_scales, ) from vllm.utils.import_utils import has_cutedsl from vllm.utils.torch_utils import ( @@ -286,6 +288,42 @@ def fused_indexer_q_rope_quant( return q_fp8, weights_out +@triton.jit +def _upcast_logits_kernel(src_ptr, dst_ptr, row_ends_ptr, stride, BLOCK: tl.constexpr): + row = tl.program_id(0) + end = tl.load(row_ends_ptr + row) + src_ptr += row.to(tl.int64) * stride + dst_ptr += row.to(tl.int64) * stride + for start in range(tl.program_id(1) * BLOCK, end, tl.num_programs(1) * BLOCK): + offs = start + tl.arange(0, BLOCK) + mask = offs < end + vals = tl.load(src_ptr + offs, mask=mask) + tl.store(dst_ptr + offs, vals.to(tl.float32), mask=mask) + + +def as_fp32_logits(logits: torch.Tensor, row_ends: torch.Tensor) -> torch.Tensor: + """Upcast SM100's BF16 MQA logits for the FP32 top-k kernels. + + Only columns below ``row_ends`` (int32, one per row) are converted, the + span the consumers read; DeepGEMM's aligned row stride is kept. + """ + if logits.dtype == torch.float32: + return logits + out = torch.empty_strided( + logits.shape, logits.stride(), dtype=torch.float32, device=logits.device + ) + num_rows, width = logits.shape + if num_rows == 0 or width == 0: + return out + assert row_ends.numel() == num_rows and row_ends.dtype == torch.int32 + block = 1024 + splits = min(triton.cdiv(width, block), max(1, 1024 // num_rows)) + _upcast_logits_kernel[(num_rows, splits)]( + logits, out, row_ends.contiguous(), logits.stride(0), BLOCK=block + ) + return out + + def _gather_workspace_shapes( total_seq_lens: int, head_dim: int, @@ -537,7 +575,7 @@ def sparse_attn_indexer( kv_cache, slot_mapping_for_cache, quant_block_size, - scale_fmt, + "mxfp8" if mqa_logits_use_mx_scales() else scale_fmt, ) # The indexer and main MLA may classify the same short extend differently # because they use independent decode thresholds. Only the main MLA route @@ -649,6 +687,9 @@ def sparse_attn_indexer( q_scale_slice = q_scale[row_start:row_end] if q_scale is not None else None topk_indices = topk_indices_buffer[row_start:row_end, :topk_tokens] + # DeepGEMM returns compressed logits: row i starts at K position + # cu_seqlen_ks[i]. None means every row starts at column zero. + row_starts: torch.Tensor | None = None if chunk.local_total_seq_lens == 0 or q_slice.shape[0] == 0: logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32) topk_indices.fill_(-1) @@ -662,7 +703,9 @@ def sparse_attn_indexer( else: q_slice_cast = q_slice k_quant_cast = k_quant - k_scale_cast = k_scale.view(torch.float32).squeeze(-1) + k_scale_cast = k_scale.view( + torch.int32 if mqa_logits_use_mx_scales() else torch.float32 + ).squeeze(-1) if current_platform.is_xpu(): if q_scale_slice is not None: raise RuntimeError("XPU fp8_mqa_logits does not support FP4 Q") @@ -674,14 +717,23 @@ def sparse_attn_indexer( cu_seqlen_ks, cu_seqlen_ke, ) + row_starts = cu_seqlen_ks + row_ends = cu_seqlen_ke else: - logits = fp8_fp4_mqa_logits( - (q_slice_cast, q_scale_slice), - (k_quant_cast, k_scale_cast), - weights[row_start:row_end], - cu_seqlen_ks, - cu_seqlen_ke, - clean_logits=False, + q_scale_slice, chunk_weights = mqa_logits_mx_inputs( + q_slice_cast, q_scale_slice, weights[row_start:row_end] + ) + row_ends = cu_seqlen_ke - cu_seqlen_ks + logits = as_fp32_logits( + fp8_fp4_mqa_logits( + (q_slice_cast, q_scale_slice), + (k_quant_cast, k_scale_cast), + chunk_weights, + cu_seqlen_ks, + cu_seqlen_ke, + chunk.max_seqlen_k, + ), + row_ends, ) num_rows = logits.shape[0] if candidate_blocks is not None: @@ -694,8 +746,8 @@ def sparse_attn_indexer( if candidate_write: _select_candidate_blocks( logits, - cu_seqlen_ks, - cu_seqlen_ke, + row_starts, + row_ends, chunk_candidates.shape[1], candidate_block_size, chunk_candidates, @@ -703,15 +755,15 @@ def sparse_attn_indexer( else: _apply_candidate_mask( logits, - cu_seqlen_ks, - cu_seqlen_ke, + row_starts, + row_ends, chunk_candidates, candidate_block_size, ) ops.top_k_per_row_prefill( logits, - cu_seqlen_ks, - cu_seqlen_ke, + torch.zeros_like(row_ends) if row_starts is None else row_starts, + row_ends, topk_indices, num_rows, logits.stride(0), @@ -729,7 +781,7 @@ def sparse_attn_indexer( dcp_rank, dcp_world_size, cp_kv_cache_interleave_size, - row_starts=cu_seqlen_ks, + row_starts=row_starts, ) if row_shard is not None: @@ -822,16 +874,22 @@ def sparse_attn_indexer( max_model_len, ) else: - logits = fp8_fp4_paged_mqa_logits( - (padded_q_quant_cast, padded_q_scale), - kv_cache, - weights[:num_padded_tokens], - seq_lens, - decode_metadata.block_table, - decode_metadata.schedule_metadata, - max_model_len=max_model_len, - clean_logits=False, - indices=decode_metadata.indices, + padded_q_scale, decode_weights = mqa_logits_mx_inputs( + padded_q_quant_cast, padded_q_scale, weights[:num_padded_tokens] + ) + # SM100 decode is varlen: one row and context length per token. + logits = as_fp32_logits( + fp8_fp4_paged_mqa_logits( + (padded_q_quant_cast, padded_q_scale), + kv_cache, + decode_weights, + seq_lens, + decode_metadata.block_table, + decode_metadata.schedule_metadata, + max_model_len=max_model_len, + indices=decode_metadata.indices, + ), + seq_lens.view(-1), ) num_rows = logits.shape[0] if candidate_blocks is not None: diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py index 255b1f5aaf2f..f876120a6ef2 100644 --- a/vllm/models/deepseek_v32/common/kernels.py +++ b/vllm/models/deepseek_v32/common/kernels.py @@ -6,7 +6,9 @@ from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.v1.attention.ops.common import pack_ue8m0_x4 # Cache of tiny 1-element dummy tensors (per device, dtype) reused by the # has_indexer=False path so the indexer args don't allocate every call. @@ -87,6 +89,7 @@ def _fp8_quant_and_cache_write( SHUFFLE: tl.constexpr, BLOCK_TILE: tl.constexpr, HEAD_TILE: tl.constexpr, + MX_SCALE: tl.constexpr, FP8_MAX: tl.constexpr, USE_FNUZ: tl.constexpr, ): @@ -111,6 +114,8 @@ def _fp8_quant_and_cache_write( tl.store(kv_cache_ptr + block_start + value_off, k_fp8, mask=mask) scale_byte_off = block_start + cache_block_size * HEAD_DIM + block_offset * 4 + if MX_SCALE: + scale = pack_ue8m0_x4(scale) tl.store(kv_cache_scale_ptr + scale_byte_off // 4, scale) @@ -167,6 +172,7 @@ def _fused_norm_rope_kernel( INDEXER_CACHE_SHUFFLE: tl.constexpr, INDEXER_CACHE_BLOCK_TILE: tl.constexpr, INDEXER_CACHE_HEAD_TILE: tl.constexpr, + INDEXER_CACHE_MX_SCALE: tl.constexpr, # MLA KV cache (concat kv_c_normed + k_pe_roped, uses slot_mapping_ptr) mla_cache_ptr, mla_cache_block_stride, @@ -501,6 +507,7 @@ def _fused_norm_rope_kernel( INDEXER_CACHE_SHUFFLE, INDEXER_CACHE_BLOCK_TILE, INDEXER_CACHE_HEAD_TILE, + INDEXER_CACHE_MX_SCALE, FP8_MAX, USE_FNUZ, ) @@ -571,7 +578,10 @@ def fused_norm_rope( assert slot_mapping is not None if indexer_slot_mapping is None: indexer_slot_mapping = slot_mapping - idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32) + idx_cache_mx_scale = mqa_logits_use_mx_scales() + idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view( + torch.int32 if idx_cache_mx_scale else torch.float32 + ) idx_cache_block_size = indexer_k_cache.shape[1] idx_cache_stride = indexer_k_cache.shape[2] # The caller's cache reports whether its reader expects the shuffled @@ -591,6 +601,7 @@ def fused_norm_rope( ) else: idx_cache_scale_view = None + idx_cache_mx_scale = False idx_cache_block_size = 1 idx_cache_stride = 0 idx_cache_shuffle = False @@ -701,6 +712,7 @@ def fused_norm_rope( idx_cache_shuffle, _INDEXER_CACHE_BLOCK_TILE, idx_cache_head_tile, + idx_cache_mx_scale, # MLA KV cache (uses same slot_mapping) mla_kv_cache, mla_block_stride, diff --git a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py index 50ca86558e05..703169f4829f 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py @@ -33,7 +33,9 @@ ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales from vllm.utils.math_utils import round_up +from vllm.v1.attention.ops.common import pack_ue8m0_x4 if current_platform.is_rocm(): from vllm.platforms.rocm import _ON_GFX950 @@ -742,6 +744,7 @@ def _fused_kv_compress_norm_rope_insert_indexer_attn( TOKEN_STRIDE: tl.constexpr, # 128 for indexer SCALE_DIM: tl.constexpr, # 4 for indexer (1 float32) KV_BLOCK_STRIDE: tl.constexpr, + MX_SCALE: tl.constexpr, # store the scale as 4 packed UE8M0 bytes ): """Fused compress → RMSNorm → RoPE → FP8 quant → store. @@ -878,7 +881,10 @@ def _fused_kv_compress_norm_rope_insert_indexer_attn( # Single float32 scale scale_val = tl.exp2(exponent) - tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), scale_val) + if MX_SCALE: + tl.store(scale_ptr.to(tl.pointer_type(tl.int32)), pack_ue8m0_x4(scale_val)) + else: + tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), scale_val) # ============================================================================= @@ -927,6 +933,7 @@ def _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn( TOKEN_STRIDE: tl.constexpr, # HEAD_SIZE // 2 = 64 packed bytes/token SCALE_DIM: tl.constexpr, # HEAD_SIZE // QUANT_BLOCK = 4 ue8m0 bytes/token KV_BLOCK_STRIDE: tl.constexpr, + MX_SCALE: tl.constexpr, # unused for MXFP4 (kept for signature parity) ): """Fused compress → RMSNorm → RoPE → MXFP4 quant → store. @@ -1102,6 +1109,7 @@ class CompileKey: block_size: int kv_cache_block_size: int kv_block_stride: int + mx_scale: bool # Canonical kernel for the launcher's arg_names mapping; the launched # variant (fp8/mxfp4, same arg_names) is picked per call in __call__. @@ -1181,6 +1189,7 @@ def dispatch( # type: ignore[override] block_size=state_block_size, kv_cache_block_size=kv_cache_block_size, kv_block_stride=kv_block_stride, + mx_scale=not use_fp4_cache and mqa_logits_use_mx_scales(), ) def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: @@ -1301,6 +1310,7 @@ def __call__( TOKEN_STRIDE=token_stride, SCALE_DIM=scale_dim, KV_BLOCK_STRIDE=kv_cache.stride(0), + MX_SCALE=not use_fp4_cache and mqa_logits_use_mx_scales(), num_warps=1, **pdl_kwargs, ) diff --git a/vllm/models/deepseek_v41/common/ops/indexer_k_store.py b/vllm/models/deepseek_v41/common/ops/indexer_k_store.py index ba8ddbd95a9a..fe203e917861 100644 --- a/vllm/models/deepseek_v41/common/ops/indexer_k_store.py +++ b/vllm/models/deepseek_v41/common/ops/indexer_k_store.py @@ -17,6 +17,8 @@ from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales +from vllm.v1.attention.ops.common import pack_ue8m0_x4 from . import MXFP4_BLOCK_SIZE, _fp32x2_to_fp4x2 @@ -51,7 +53,8 @@ def indexer_k_norm_rope_store( kv_slot_mapping: [num_tokens] slots in the indexer cache (-1 = skip). compress_ratio: group size; keys are emitted at group boundaries. use_fp4_cache: MXFP4 (2 nibbles/byte + ue8m0 per 32) when True, else - per-token FP8 with a single fp32 scale. + per-token FP8 with a single fp32 scale (four packed ue8m0 bytes + under `mqa_logits_use_mx_scales`). """ num_tokens = kv_slot_mapping.numel() @@ -111,6 +114,7 @@ def indexer_k_norm_rope_store( KV_BLOCK_STRIDE=k_cache.stride(0), FP8_MAX=448.0, USE_FP4=use_fp4_cache, + MX_SCALE=not use_fp4_cache and mqa_logits_use_mx_scales(), SHUFFLE=shuffle, BLOCK_TILE_SIZE=_BLOCK_TILE_SIZE, HEAD_TILE_SIZE=_HEAD_TILE_SIZE, @@ -139,6 +143,7 @@ def _indexer_k_norm_rope_quant_store_kernel( KV_BLOCK_STRIDE: tl.constexpr, FP8_MAX: tl.constexpr, USE_FP4: tl.constexpr, + MX_SCALE: tl.constexpr, SHUFFLE: tl.constexpr, BLOCK_TILE_SIZE: tl.constexpr, HEAD_TILE_SIZE: tl.constexpr, @@ -252,4 +257,10 @@ def _indexer_k_norm_rope_quant_store_kernel( tl.store(val_ptr + tiled, x_uint8) else: tl.store(val_ptr + block, x_uint8) - tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), tl.exp2(exponent)) + if MX_SCALE: + tl.store( + scale_ptr.to(tl.pointer_type(tl.int32)), + pack_ue8m0_x4(tl.exp2(exponent)), + ) + else: + tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), tl.exp2(exponent)) diff --git a/vllm/models/glm5next/nvidia/ops/kpool_compress.py b/vllm/models/glm5next/nvidia/ops/kpool_compress.py index eb289c65cf68..310d60dfd41e 100644 --- a/vllm/models/glm5next/nvidia/ops/kpool_compress.py +++ b/vllm/models/glm5next/nvidia/ops/kpool_compress.py @@ -15,6 +15,8 @@ import torch from vllm.triton_utils import tl, triton +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales +from vllm.v1.attention.ops.common import pack_ue8m0_x4 # The GLM-5.3-Flash indexer head dimension is fixed at 128. INDEX_HEAD_DIM = 128 @@ -157,6 +159,7 @@ def _kpool_softmax_rotate_write_cache_kernel( HEAD_DIM: tl.constexpr, S_OFFSET_NBYTES_IN_PAGE: tl.constexpr, ROUND_SCALE: tl.constexpr, + MX_SCALE: tl.constexpr, HAS_WRITE_MASK: tl.constexpr, RETURN_COMPRESSED: tl.constexpr, WRITE_CACHE: tl.constexpr, @@ -245,7 +248,14 @@ def _kpool_softmax_rotate_write_cache_kernel( + loc_token_offset_in_page ) tl.store(buf_fp8_ptr + out_k_offsets, quantized, mask=mask) - tl.store(buf_fp32_ptr + out_s_offset, scale, mask=do_write) + if MX_SCALE: + tl.store( + buf_fp32_ptr.to(tl.pointer_type(tl.int32)) + out_s_offset, + pack_ue8m0_x4(scale), + mask=do_write, + ) + else: + tl.store(buf_fp32_ptr + out_s_offset, scale, mask=do_write) if RETURN_COMPRESSED: tl.store( @@ -324,6 +334,7 @@ def kpool_compress_and_write_cache( buf_fp8 = buf.view(torch.float8_e4m3fn) buf_fp32 = buf.view(torch.float32) + mx_scale = write_cache and mqa_logits_use_mx_scales() # bytes per page (last dim of kv_cache) viewed as uint8 buf_numel_per_page = buf.stride(0) s_offset_nbytes_in_page = page_size * head_dim @@ -361,7 +372,9 @@ def kpool_compress_and_write_cache( POOL_SIZE=slot_k.shape[1], HEAD_DIM=head_dim, S_OFFSET_NBYTES_IN_PAGE=s_offset_nbytes_in_page, - ROUND_SCALE=round_scale, + # MX scales need power-of-two scales. + ROUND_SCALE=round_scale or mx_scale, + MX_SCALE=mx_scale, HAS_WRITE_MASK=has_write_mask, RETURN_COMPRESSED=return_compressed, WRITE_CACHE=write_cache, @@ -479,6 +492,7 @@ def _kpool_decode_update_batched_kernel( HEAD_DIM: tl.constexpr, S_OFFSET_NBYTES_IN_PAGE: tl.constexpr, ROUND_SCALE: tl.constexpr, + MX_SCALE: tl.constexpr, BLOCK_D: tl.constexpr, ): """One program per request; iterates its NEXT_N verify tokens in order. @@ -604,7 +618,13 @@ def _kpool_decode_update_batched_kernel( + loc_token_offset_in_page ) tl.store(buf_fp8_ptr + out_k_offsets, quantized, mask=dim_mask) - tl.store(buf_fp32_ptr + out_s_offset, scale) + if MX_SCALE: + tl.store( + buf_fp32_ptr.to(tl.pointer_type(tl.int32)) + out_s_offset, + pack_ue8m0_x4(scale), + ) + else: + tl.store(buf_fp32_ptr + out_s_offset, scale) # Stash the current token AFTER any completion read so the completion # uses prior stashes (and the current token's own key/score via @@ -718,7 +738,9 @@ def kpool_decode_update_and_maybe_write_cache_batched( KPOOL_HEAD=tail_kv_cache.stride(1), HEAD_DIM=head_dim, S_OFFSET_NBYTES_IN_PAGE=page_size * head_dim, - ROUND_SCALE=round_scale, + # MX scales need power-of-two scales. + ROUND_SCALE=round_scale or mqa_logits_use_mx_scales(), + MX_SCALE=mqa_logits_use_mx_scales(), BLOCK_D=triton.next_power_of_2(head_dim), ) diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 154fdcda900e..1a3fbedb8387 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -16,6 +16,7 @@ PrefillRowShard, _merge_dcp_topk_global, _reserve_row_shard_exchange_transient, + as_fp32_logits, ) from vllm.models.glm5next.common.sparse_indexer import ( RADIX_TOPK_WORKSPACE_SIZE, @@ -29,7 +30,13 @@ ) from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops from vllm.platforms import current_platform -from vllm.utils.deep_gemm import has_deep_gemm +from vllm.utils.deep_gemm import ( + fp8_fp4_mqa_logits, + fp8_fp4_paged_mqa_logits, + has_deep_gemm, + mqa_logits_mx_inputs, + mqa_logits_use_mx_scales, +) from vllm.utils.import_utils import has_cutedsl from vllm.utils.torch_utils import ( LayerNameType, @@ -253,7 +260,7 @@ def sparse_attn_indexer_kpool( kv_cache, slot_mapping, quant_block_size, - scale_fmt, + "mxfp8" if mqa_logits_use_mx_scales() else scale_fmt, ) topk_indices_buffer[: hidden_states.shape[0]] = -1 @@ -343,20 +350,29 @@ def sparse_attn_indexer_kpool( else: q_slice_cast = q_slice k_quant_cast = k_quant - k_scale_cast = k_scale.view(torch.float32).squeeze(-1) - from vllm.utils.deep_gemm import fp8_fp4_mqa_logits + k_scale_cast = k_scale.view( + torch.int32 if mqa_logits_use_mx_scales() else torch.float32 + ).squeeze(-1) num_rows = q_slice.shape[0] if chunk.local_total_seq_lens == 0: logits = q_slice.new_empty((num_rows, 0), dtype=torch.float32) else: - logits = fp8_fp4_mqa_logits( - (q_slice_cast, q_scale_slice), - (k_quant_cast, k_scale_cast), - weights[row_start:row_end], - cu_seqlen_ks, - cu_seqlen_ke, - clean_logits=False, + q_scale_slice, chunk_weights = mqa_logits_mx_inputs( + q_slice_cast, q_scale_slice, weights[row_start:row_end] + ) + # Compressed: row i starts at K position cu_seqlen_ks[i]. + row_ends = cu_seqlen_ke - cu_seqlen_ks + logits = as_fp32_logits( + fp8_fp4_mqa_logits( + (q_slice_cast, q_scale_slice), + (k_quant_cast, k_scale_cast), + chunk_weights, + cu_seqlen_ks, + cu_seqlen_ke, + chunk.max_seqlen_k, + ), + row_ends, ) # kpool: logits are pool-granular (compress_ratio == index_kpool), @@ -376,8 +392,8 @@ def sparse_attn_indexer_kpool( else: torch.ops._C.top_k_per_row_prefill( logits, - cu_seqlen_ks, - cu_seqlen_ke, + torch.zeros_like(row_ends), + row_ends, topk_dst, num_rows, logits.stride(0), @@ -392,7 +408,7 @@ def sparse_attn_indexer_kpool( dcp_rank, dcp_world_size, pool_cp_interleave, - row_starts=cu_seqlen_ks, + row_starts=None, ) if index_kpool > 1: @@ -594,18 +610,22 @@ def sparse_attn_indexer_kpool( if use_fp4_cache else padded_q_quant_decode_tokens ) - from vllm.utils.deep_gemm import fp8_fp4_paged_mqa_logits - - logits = fp8_fp4_paged_mqa_logits( - (padded_q_quant_cast, padded_q_scale), - kv_cache, - padded_weights[:num_padded_tokens], - seq_lens, - decode_metadata.block_table, - decode_metadata.schedule_metadata, - max_model_len=max_pool_len, - clean_logits=False, - indices=decode_metadata.indices, + padded_q_scale, decode_weights = mqa_logits_mx_inputs( + padded_q_quant_cast, padded_q_scale, padded_weights[:num_padded_tokens] + ) + # SM100 decode is varlen: one row and context length per token. + logits = as_fp32_logits( + fp8_fp4_paged_mqa_logits( + (padded_q_quant_cast, padded_q_scale), + kv_cache, + decode_weights, + seq_lens, + decode_metadata.block_table, + decode_metadata.schedule_metadata, + max_model_len=max_pool_len, + indices=decode_metadata.indices, + ), + seq_lens.view(-1), ) num_rows = logits.shape[0] # kpool: logits are pool-granular -> select topk_tokens//kpool pools, diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 04c4508f0ba9..3dc9be5d9385 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -591,37 +591,79 @@ def transform_sf_into_required_layout(*args, **kwargs): ) +@functools.cache +def mqa_logits_use_mx_scales() -> bool: + """Whether the MQA logits kernels take only MX-scaled inputs. + + On SM100 they need MXFP8/MXFP4 Q and K with packed UE8M0 scales and BF16 + weights, and return BF16 logits. FP8 indexer K caches then store each + token's power-of-two scale as four equal UE8M0 bytes instead of an fp32. + """ + return current_platform.is_cuda() and current_platform.is_device_capability_family( + 100 + ) + + +# Four packed UE8M0 bytes encoding 1.0. +_UNIT_UE8M0_X4 = 0x7F7F7F7F +# Grown buffers are kept alive: CUDA graphs captured earlier still read them. +_unit_mx_scales: dict[torch.device, list[torch.Tensor]] = {} + + +def mqa_logits_mx_inputs( + q: torch.Tensor, q_scale: torch.Tensor | None, weights: torch.Tensor +) -> tuple[torch.Tensor | None, torch.Tensor]: + """Adapt indexer Q scales and weights to the SM100 MQA logits kernels. + + FP8 Q folds its per-(token, head) power-of-two scale into ``weights``, so + it is passed with unit MX scales, which is exact since relu commutes with + a positive scale. Returns ``(q_scale, weights)`` unchanged off SM100. + """ + if not mqa_logits_use_mx_scales(): + return q_scale, weights + if q_scale is None: + numel = q.numel() // q.shape[-1] + bufs = _unit_mx_scales.setdefault(q.device, []) + if not bufs or bufs[-1].numel() < numel: + size = max(numel, 2 * bufs[-1].numel() if bufs else 0) + bufs.append( + torch.full((size,), _UNIT_UE8M0_X4, dtype=torch.int32, device=q.device) + ) + q_scale = bufs[-1][:numel].view(q.shape[:-1]) + return q_scale, weights.to(torch.bfloat16) + + def fp8_fp4_mqa_logits( q: tuple[torch.Tensor, torch.Tensor | None], kv: tuple[torch.Tensor, torch.Tensor], weights: torch.Tensor, cu_seqlen_ks: torch.Tensor, cu_seqlen_ke: torch.Tensor, - clean_logits: bool, + max_seqlen_k: int, ) -> torch.Tensor: """Compute MQA logits for a single sequence without KV paging. Unified FP8/FP4 dispatch — the underlying DeepGEMM kernel takes - ``q = (values, scales_or_None)`` where ``scales`` is None for FP8 Q - (per-token scale is folded into ``weights``) and a packed block-scale - tensor for MXFP4 Q. + ``q = (values, scales_or_None)``. SM100 takes MXFP8 or MXFP4 with packed + UE8M0 block scales; SM90/SM120 also take FP8 Q with ``q_scale=None`` + (per-token scale folded into ``weights``). Args: - q: Tuple ``(q_values, q_scale)``. FP8 path: q_values is [M, H, D] - float8_e4m3fn and q_scale is None (per-token scale is folded - into ``weights``). FP4 path: q_values is packed uint8 and - q_scale is the companion block-scale tensor. - kv: Tuple `(k_packed, k_scales)` — FP8 layout is [N, D] - float8_e4m3fn plus fp32 scales [N]; FP4 layout is packed uint8. - weights: weights of shape [M, H], dtype `torch.float32`. + q: Tuple ``(q_values, q_scale)``, q_values of shape [M, H, D] (packed + to D/2 bytes for MXFP4) and q_scale int32 [M, H] or None. + kv: Tuple ``(k_values, k_scales)`` of shape [N, D] and [N]; k_scales + is int32 packed UE8M0 with MX scales, fp32 otherwise. + weights: Tensor of shape [M, H], bf16 on SM100 and fp32 elsewhere. cu_seqlen_ks: Start indices (inclusive) for valid K per query position, shape [M], dtype int32. cu_seqlen_ke: End indices (exclusive) for valid K per query position, shape [M], dtype int32. - clean_logits: Whether to clean the unfilled logits into `-inf`. + max_seqlen_k: Upper bound on ``cu_seqlen_ke - cu_seqlen_ks``. Returns: - Logits tensor of shape [M, N], dtype `torch.float32`. + Compressed logits of shape [M, max_seqlen_k]: row ``i`` holds K + positions ``[cu_seqlen_ks[i], cu_seqlen_ke[i])`` starting at column + zero, and the tail is not cleaned. bf16 on SM100, fp32 elsewhere. """ _lazy_init() @@ -633,7 +675,7 @@ def fp8_fp4_mqa_logits( weights, cu_seqlen_ks, cu_seqlen_ke, - clean_logits=clean_logits, + max_seqlen_k, ) @@ -699,24 +741,25 @@ def fp8_fp4_paged_mqa_logits( block_tables: torch.Tensor, schedule_metadata: torch.Tensor, max_model_len: int, - clean_logits: bool, indices: torch.Tensor | None = None, ) -> torch.Tensor: """Compute MQA logits using a paged KV-cache. Unified FP8/FP4 dispatch — the underlying DeepGEMM kernel takes - ``q = (values, scales_or_None)``; pass ``(q_tensor, None)`` for the FP8 - path and ``(q_values, q_scale)`` for MXFP4. + ``q = (values, scales_or_None)``; see `fp8_fp4_mqa_logits` for the + per-architecture formats. SM100 requires the varlen form (``indices`` + with next_n == 1). Args: - q: Tuple ``(q_values, q_scale)``. FP8 path: q_values is - [B, next_n, H, D] float8_e4m3fn and q_scale is None. FP4 path: - q_values is packed uint8 and q_scale is the companion - block-scale tensor. - kv_cache: Paged KV-cache. FP8 layout is [num_blocks, block_size, 1, - D+4], dtype `torch.uint8`, with the last 4 bytes per (block, pos) - storing the float dequant scale. - weights: Tensor of shape [B * next_n, H], dtype `torch.float32`. + q: Tuple ``(q_values, q_scale)``, q_values of shape + [B, next_n, H, D] (packed to D/2 bytes for MXFP4) and q_scale + int32 [B, next_n, H] or None. + kv_cache: Paged KV-cache of shape [num_blocks, block_size, 1, D+4] + (D/2+4 for MXFP4), dtype `torch.uint8`; the last 4 bytes per + (block, pos) hold the int32 packed UE8M0 scales with MX scales, + or the fp32 dequant scale otherwise. + weights: Tensor of shape [B * next_n, H], bf16 on SM100 and fp32 + elsewhere. context_lens: Tensor of shape [B], dtype int32; effective context length for each batch element. block_tables: Tensor of shape [B, max_blocks], dtype int32; maps logical @@ -724,12 +767,11 @@ def fp8_fp4_paged_mqa_logits( schedule_metadata: Returned by `get_paged_mqa_logits_metadata`; used to distribute work across SMs. max_model_len: Maximum sequence length used to size the logits output. - clean_logits: Whether to clean the unfilled logits into `-inf`. indices: Optional request index for each varlen row. Returns: - Logits tensor of shape [B * next_n, max_model_len], dtype - `torch.float32`. + Logits tensor of shape [B * next_n, max_model_len]; columns past each + row's context length are not cleaned. bf16 on SM100, fp32 elsewhere. """ _lazy_init() @@ -751,7 +793,6 @@ def fp8_fp4_paged_mqa_logits( block_tables, schedule_metadata, max_model_len, - clean_logits=clean_logits, **kwargs, ) @@ -1052,6 +1093,8 @@ def should_use_deepgemm_for_fp8_linear( "fp8_fp4_mqa_logits", "fp8_fp4_paged_mqa_logits", "get_paged_mqa_logits_metadata", + "mqa_logits_mx_inputs", + "mqa_logits_use_mx_scales", "native_next_n_supported", "per_block_cast_to_fp8", "is_deep_gemm_e8m0_used", diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 2d797ae66283..0382a81fe044 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -467,6 +467,9 @@ class DeepseekV32IndexerPrefillChunkMetadata: cu_seq_lens: torch.Tensor token_to_seq: torch.Tensor total_seq_lens: int + # Host upper bound on cu_seqlen_ke - cu_seqlen_ks, sizing the compressed + # DeepGEMM prefill logits. + max_seqlen_k: int token_start: int token_end: int num_reqs: int @@ -1957,6 +1960,7 @@ def build_prefill_chunk_metadata( cu_seq_lens=cu_seq_lens, token_to_seq=token_to_seq, total_seq_lens=total_seq_lens, + max_seqlen_k=int(compressed_seq_lens_cpu[start_idx:end_idx].max()), block_table=block_table[start_idx:end_idx], token_start=token_start, token_end=token_end, diff --git a/vllm/v1/attention/ops/common.py b/vllm/v1/attention/ops/common.py index 5c1677cac479..6936be9fa169 100644 --- a/vllm/v1/attention/ops/common.py +++ b/vllm/v1/attention/ops/common.py @@ -154,6 +154,13 @@ def __call__( ) +@triton.jit +def pack_ue8m0_x4(scale): + """Encode a power-of-two fp32 scale as four equal packed UE8M0 bytes, the + MX scales of a 128-wide row (see `mqa_logits_use_mx_scales`).""" + return ((scale.to(tl.int32, bitcast=True) >> 23) & 0xFF) * 0x01010101 + + def pack_seq_triton( x: torch.Tensor, lengths: torch.Tensor, From e2ad0ce5f5a944e10ed2ea664e4a24bf85791346 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 13 Oct 2026 14:25:16 +0800 Subject: [PATCH 362/624] [Bugfix][LoRA] Slice expert LoRA by the unpadded per-rank size under TP (#61235) Signed-off-by: aoshen524 Co-authored-by: aoshen524 Co-authored-by: Claude Opus 5.5 Co-authored-by: linitra24 --- vllm/lora/layers/fused_moe.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index c58c18e71ff7..19a1128ce145 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -3,6 +3,7 @@ import torch import torch.nn as nn +import torch.nn.functional as F from transformers import PreTrainedConfig from vllm import envs @@ -110,6 +111,15 @@ def use_ep(self) -> bool: def intermediate_size_per_partition(self) -> int: return self.moe_config.intermediate_size_per_partition + @property + def _tp_shard_size(self) -> int: + """Per-rank stride of an adapter along the expert intermediate dim: the + base weights are sharded by the unpadded size even when a backend pads.""" + unpadded = self.moe_config.intermediate_size_per_partition_unpadded + if unpadded is None: + return self.intermediate_size_per_partition + return unpadded + def _init_lora_stream_context(self) -> None: self._lora_stream: torch.cuda.Stream | None = None self._events: tuple[torch.cuda.Event, ...] | None = None @@ -311,7 +321,7 @@ def _slice_w13_b(self, w13_lora_b: torch.Tensor): return w13_lora_b # w13_lora_b shape (num_experts,output_size,rank) - shard_size = self.intermediate_size_per_partition + shard_size = self._tp_shard_size start_idx = self.tp_rank * shard_size end_idx = (self.tp_rank + 1) * shard_size @@ -322,7 +332,7 @@ def _slice_w2_a(self, w2_lora_a: torch.Tensor) -> torch.Tensor: if self.tp_size == 1: return w2_lora_a # w2_lora_a shape (num_experts,rank,input_size) - shard_size = self.intermediate_size_per_partition + shard_size = self._tp_shard_size start_idx = self.tp_rank * shard_size end_idx = (self.tp_rank + 1) * shard_size @@ -536,7 +546,7 @@ def _slice_w13_b(self, w13_lora_b: torch.Tensor): return w13_lora_b # w13_lora_b shape (num_experts,output_size,rank) - shard_size = self.intermediate_size_per_partition + shard_size = self._tp_shard_size start_idx = self.tp_rank * shard_size end_idx = (self.tp_rank + 1) * shard_size # HACK: Currently, only GPT-OSS is in interleaved order @@ -557,6 +567,11 @@ def _slice_w13_b(self, w13_lora_b: torch.Tensor): w3_lora_b = w13_lora_b[:, slice_size:, :] sliced_w1_lora_b = w1_lora_b[:, start_idx:end_idx, :] sliced_w3_lora_b = w3_lora_b[:, start_idx:end_idx, :] + # w3 starts at the allocated (possibly padded) per-rank size, as + # in the base w13 weight. + pad = self.intermediate_size_per_partition - sliced_w1_lora_b.shape[1] + if pad > 0: + sliced_w1_lora_b = F.pad(sliced_w1_lora_b, (0, 0, 0, pad)) return torch.cat([sliced_w1_lora_b, sliced_w3_lora_b], dim=1) From 4e4eab73ed2b4e00931a4636e926b8df1cb1fe50 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 12 Oct 2026 23:57:32 -0700 Subject: [PATCH 363/624] [CI] Shard Model Executor and L4 Pipeline + Context Parallelism automatically (#61202) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/distributed.yaml | 1 + .buildkite/test_areas/model_executor.yaml | 1 + 2 files changed, 2 insertions(+) diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index b30cbac0ced4..2f929b53b840 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -641,6 +641,7 @@ steps: - label: ":nvidia: (L4) Pipeline + Context Parallelism" key: pipeline-context-parallelism-4-gpus + automatic_shard: true timeout_in_minutes: 55 working_dir: "/vllm-workspace/tests" device: l4 diff --git a/.buildkite/test_areas/model_executor.yaml b/.buildkite/test_areas/model_executor.yaml index a23359b768a4..c9f467941af9 100644 --- a/.buildkite/test_areas/model_executor.yaml +++ b/.buildkite/test_areas/model_executor.yaml @@ -5,6 +5,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Model Executor" device: h200_35gb key: model-executor + automatic_shard: true timeout_in_minutes: 60 source_file_dependencies: - vllm/engine/arg_utils.py From bc3ddb198a0d5eb864bf484c4719864b671cc93b Mon Sep 17 00:00:00 2001 From: trienebutdiene <59642603+kimseunghyun-kr@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:30:19 +0800 Subject: [PATCH 364/624] [Model][Qwen4Exp] FP8 QSA main KV cache on Ampere (sm_80, sm_86) (#60602) Signed-off-by: kimseunghyun Signed-off-by: kimseunghyun Signed-off-by: Thien Tran Co-authored-by: kimseunghyun Co-authored-by: Claude Opus 5.5 Co-authored-by: Thien Tran --- vllm/models/qwen4_exp/nvidia/ops/qsa.py | 69 +++++++++++++++++++------ vllm/models/qwen4_exp/nvidia/qsa.py | 21 +++++--- 2 files changed, 67 insertions(+), 23 deletions(-) diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index 401048e02c55..c7fe6014497e 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -22,6 +22,28 @@ def _is_sm120() -> bool: return current_platform.get_device_capability() == (12, 0) +@triton.jit +def _maybe_dequant(tile): + if tile.dtype == tl.uint8: + # e4m3 bytes below sm_89 (no Triton fp8e4nv pointer type): sign to bit + # 15, exp+mantissa to bits 13..7 gives the fp16 of x * 2^-8. Exact for + # every finite code; the two NaN codes decode to +-480 after the host + # scale instead of NaN. + code = tile.to(tl.uint16) + bits = ((code & 0x80) << 8) | ((code & 0x7F) << 7) + return bits.to(tl.float16, bitcast=True) + elif tile.dtype == tl.float8e4nv: + return tile.to(tl.float16) + else: + return tile + + +@lru_cache(maxsize=1) +def _is_sm80() -> bool: + """True on sm_80 (A100/A30): selects the wide FP8 byte-decode tiles.""" + return current_platform.get_device_capability() == (8, 0) + + @lru_cache(maxsize=1) def _is_sm90() -> bool: """True on sm_90 (H100/H200/H20): selects the sm_90 tuning table.""" @@ -69,7 +91,6 @@ def _qsa_sparse_paged_gqa_splitk_kernel( NUM_TILES: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, - IS_FP8: tl.constexpr, ) -> None: row = tl.program_id(0) kv_head = tl.program_id(1) @@ -141,6 +162,9 @@ def _qsa_sparse_paged_gqa_splitk_kernel( mask=valid[None, :], other=0.0, ) + # e4m3 -> Q dtype is exact; keep the QK dot in Q's dtype (fp8 QK + # measured slower here and less accurate). + keys = _maybe_dequant(keys).to(query.dtype) values = tl.load( v_cache_ptr + safe_page[:, None] * stride_v_block @@ -150,10 +174,10 @@ def _qsa_sparse_paged_gqa_splitk_kernel( mask=valid[:, None], other=0.0, ) - if IS_FP8: - # e4m3 -> Q dtype is exact; keep the QK dot in Q's dtype (fp8 QK - # measured slower here and less accurate). - keys = keys.to(query.dtype) + # Dequant V to fp16 (not bf16) for the PV dot: P <= 1 (online + # softmax) so fp16 has the range, its wider mantissa is more + # accurate, and the fp8->fp16 upcast with an fp16 PV dot is faster. + values = _maybe_dequant(values) scores = tl.dot(query, keys) # Scaling scores avoids re-quantizing a scaled query to BF16; for fp8 # caches the K dequant scale is already folded into softmax_scale on the @@ -165,11 +189,6 @@ def _qsa_sparse_paged_gqa_splitk_kernel( probabilities = tl.where( valid[None, :], tl.math.exp2(scores - next_max[:, None]), 0.0 ) - if IS_FP8: - # Dequant V to fp16 (not bf16) for the PV dot: P <= 1 (online - # softmax) so fp16 has the range, its wider mantissa is more - # accurate, and the fp8->fp16 upcast with an fp16 PV dot is faster. - values = values.to(tl.float16) accumulator = tl.dot( probabilities.to(values.dtype), values, @@ -585,6 +604,10 @@ def _select_config( BLOCK_N, target_splits, num_warps = 64, 4, 2 else: BLOCK_N, target_splits, num_warps = 64, 1, 2 + if is_fp8 and _is_sm80(): + # Wider tiles measured faster for the byte-decode path on A100; + # unmeasured on sm_86, which keeps the table's tiles. + BLOCK_N, num_warps = 128, 4 num_tiles = triton.cdiv(num_columns, BLOCK_N) # Never more splits than tiles, never empty. num_splits = min(target_splits, num_tiles) @@ -607,7 +630,9 @@ def qsa_sparse_paged_attention( ) -> torch.Tensor: """Run sparse GQA directly over paged BF16 or FP8-e4m3 K/V caches. - With fp8 caches, k_scale/v_scale are the layer's per-tensor dequant scales + An fp8 cache is a native e4m3 view on sm_89+ and the raw uint8 storage + below sm_89 (no Triton fp8e4nv pointer type there); the kernel dequantizes + either on load. k_scale/v_scale are the layer's per-tensor dequant scales as host floats (e.g. layer._k_scale_float/_v_scale_float): k_scale is pre-multiplied into the softmax scale and v_scale becomes the kernel's output scale, so the kernel needs no device scale buffers. @@ -636,15 +661,21 @@ def qsa_sparse_paged_attention( assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0 assert q.dtype == torch.bfloat16 assert k_cache.dtype == v_cache.dtype - is_fp8 = k_cache.dtype == torch.float8_e4m3fn + is_fp8 = k_cache.dtype != torch.bfloat16 if is_fp8: assert k_scale is not None and v_scale is not None # Host pre-multiply: fold the K dequant scale into the attention scale # and pass V's dequant scale as the kernel's output scale. softmax_scale = (head_dim**-0.5) * float(k_scale) output_scale = float(v_scale) + if k_cache.dtype == torch.uint8: + # _maybe_dequant's byte decode yields x * 2^-8; fold 2^8 into the + # scales so the decode needs no multiply. + softmax_scale *= 256.0 + output_scale *= 256.0 + else: + assert k_cache.dtype == torch.float8_e4m3fn else: - assert k_cache.dtype == torch.bfloat16 softmax_scale = head_dim**-0.5 output_scale = 1.0 assert logical_indices.dtype == block_table.dtype == torch.int32 @@ -731,7 +762,6 @@ def qsa_sparse_paged_attention( NUM_TILES=num_tiles, BLOCK_M=block_m, BLOCK_N=block_n, - IS_FP8=is_fp8, num_warps=partial_warps, num_stages=2, ) @@ -768,9 +798,15 @@ def warmup_qsa_sparse_paged_attention( """Compile every production-reachable split-K/merge specialization.""" head_dim = kv_cache.shape[-1] // 2 key_cache, value_cache = kv_cache.transpose(1, 2).split(head_dim, dim=-1) - # An fp8 cache is allocated as uint8 and viewed as e4m3 at attention time. + # An fp8 cache is allocated as uint8 storage; warm the dtype the impl hands + # over: an e4m3 view on sm_89+, the raw bytes below. is_fp8 = kv_cache.dtype == torch.uint8 - cache_dtype = torch.float8_e4m3fn if is_fp8 else key_cache.dtype + fp8_dtype = ( + torch.float8_e4m3fn + if current_platform.has_device_capability(89) + else torch.uint8 + ) + cache_dtype = fp8_dtype if is_fp8 else key_cache.dtype num_kv_heads = key_cache.shape[2] group_size = num_query_heads // num_kv_heads block_m = triton.next_power_of_2(group_size) @@ -874,7 +910,6 @@ def warmup_qsa_sparse_paged_attention( NUM_TILES=num_tiles, BLOCK_M=block_m, BLOCK_N=block_n, - IS_FP8=is_fp8, num_warps=warps, num_stages=2, grid=(num_rows, num_kv_heads, num_splits), diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index 9f0934d5a0c6..9fb36673a5e7 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -297,10 +297,12 @@ def forward_qsa( key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) k_scale = v_scale = None if self.kv_cache_dtype in ("fp8", "fp8_e4m3"): - # The cache is allocated as uint8; reinterpret the e4m3 bytes - # (same itemsize, so shape and strides are preserved). - key_cache = key_cache.view(torch.float8_e4m3fn) - value_cache = value_cache.view(torch.float8_e4m3fn) + # The cache is allocated as uint8; on sm_89+ reinterpret the e4m3 + # bytes (same itemsize, so shape and strides are preserved). Below + # sm_89 the kernel byte-decodes the raw uint8 storage instead. + if current_platform.has_device_capability(89): + key_cache = key_cache.view(torch.float8_e4m3fn) + value_cache = value_cache.view(torch.float8_e4m3fn) # Host-side per-tensor dequant scales (Python floats), as used by # other host-scale backends; folded into the kernel's scales. k_scale = layer._k_scale_float @@ -308,6 +310,7 @@ def forward_qsa( if query.dtype != torch.bfloat16 or key_cache.dtype not in ( torch.bfloat16, torch.float8_e4m3fn, + torch.uint8, ): raise NotImplementedError( "Qwen4Exp QSA requires BF16 Q and BF16 or FP8-e4m3 K/V" @@ -492,9 +495,15 @@ def __init__( ) # One launch does the indexer prepare, the main QK-norm/RoPE/gate and # the main K/V cache write (see QSAIndexer.forward); otherwise all of - # them take the separate kernels. + # them take the separate kernels. The fused writer receives an fp8 + # pointer, which Triton cannot lower below sm_89. self.use_fused_qsa_prepare = ( - self.use_fused_qk_norm_rope_gate and self.indexer.use_fused_pre_indexer + self.use_fused_qk_norm_rope_gate + and self.indexer.use_fused_pre_indexer + and ( + self.kv_cache_torch_dtype != torch.uint8 + or current_platform.has_device_capability(89) + ) ) self.fuse_indexer_projection = vllm_config.lora_config is None if self.fuse_indexer_projection: From 2cc2e68d1c5652a5bc7ae10422697308e2193359 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Tue, 13 Oct 2026 15:58:59 +0800 Subject: [PATCH 365/624] [Core][Mamba] Remove leftovers of mamba_cache_mode="all" (#60043) Signed-off-by: zjy0516 Signed-off-by: Jiangyun Zhu Co-authored-by: Claude Opus 5.5 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .buildkite/test-amd.yaml | 2 +- .buildkite/test_areas/disaggregated.yaml | 2 +- docs/design/hybrid_kv_cache_manager.md | 2 +- docs/features/automatic_prefix_caching.md | 7 +- .../test_kimi_linear_context_parallel.py | 1 - .../configs/Qwen3.8-Flash-Next-FP8-rocm.yaml | 2 +- tests/kernels/mamba/test_mamba_mixer2.py | 10 +- .../test_mamba_eagle_resume_checkpoint.py | 4 +- tests/v1/core/test_contiguous_kv_packing.py | 9 +- tests/v1/core/test_prefix_caching.py | 90 +++------- .../v1/e2e/general/test_mamba_prefix_cache.py | 2 - tests/v1/e2e/test_hybrid_chunked_prefill.py | 2 - tests/v1/e2e/test_replayssm_decode.py | 4 - .../run_mamba_prefix_cache_test.sh | 2 - .../unit/test_offloading_connector.py | 1 - .../unit/test_remote_prefill_lifecycle.py | 3 + tests/v1/simple_kv_offload/test_scheduler.py | 74 ++++++-- vllm/config/cache.py | 28 ++- vllm/config/model.py | 4 - vllm/config/vllm.py | 10 +- vllm/engine/arg_utils.py | 18 +- .../layers/mamba/mamba_mixer2.py | 6 +- .../layers/mamba/ops/causal_conv1d.py | 164 +----------------- vllm/model_executor/models/config.py | 2 +- vllm/model_executor/models/falcon_h1.py | 2 - .../model_executor/models/granitemoehybrid.py | 2 - vllm/model_executor/models/interfaces.py | 25 +-- vllm/model_executor/models/jamba.py | 2 - vllm/model_executor/models/nemotron_h.py | 2 - vllm/model_executor/models/registry.py | 3 - vllm/model_executor/models/zamba2.py | 8 +- vllm/v1/attention/backends/mamba_attn.py | 9 +- vllm/v1/core/single_type_kv_cache_manager.py | 3 + 33 files changed, 161 insertions(+), 344 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index d815ee9268b2..7f0a351524e1 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -3523,7 +3523,7 @@ steps: - export MODEL_NAMES='moonshotai/Kimi-Linear-48B-A3B-Instruct' - export PREFILL_BLOCK_SIZE='2048' - export DECODE_BLOCK_SIZE='2048' - - export VLLM_SERVE_EXTRA_ARGS='--trust-remote-code,--mamba-cache-mode align,--max-model-len 32768' + - export VLLM_SERVE_EXTRA_ARGS='--trust-remote-code,--max-model-len 32768' - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index b9bd59583b82..4821fac0a20b 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -426,7 +426,7 @@ steps: # Keep TP=2 prefill and TP=1 decode above both hybrid page-size minima. PREFILL_BLOCK_SIZE: "2048" DECODE_BLOCK_SIZE: "2048" - VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--mamba-cache-mode align,--max-model-len 32768" + VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--max-model-len 32768" commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh diff --git a/docs/design/hybrid_kv_cache_manager.md b/docs/design/hybrid_kv_cache_manager.md index 9891c82f97c9..3984000c7d8f 100644 --- a/docs/design/hybrid_kv_cache_manager.md +++ b/docs/design/hybrid_kv_cache_manager.md @@ -204,7 +204,7 @@ The second question is the cache eviction policy. For now, we use one LRU queue ### Case 3: mamba models -The prefix caching support of the mamba model is work in progress. Once implemented, models with mamba layer + full attention layer can be supported via the full attention + X algorithm in case 2. +Models with mamba layers + full attention layers support prefix caching via the full attention + X algorithm in case 2, with the Mamba state cached at block boundaries. See [Automatic Prefix Caching](../features/automatic_prefix_caching.md). ## Implementation diff --git a/docs/features/automatic_prefix_caching.md b/docs/features/automatic_prefix_caching.md index cfc8c19acb2b..bfd87e7a2c57 100644 --- a/docs/features/automatic_prefix_caching.md +++ b/docs/features/automatic_prefix_caching.md @@ -15,23 +15,22 @@ Set `enable_prefix_caching=True` in vLLM engine to enable APC. Here is an exampl ## Hybrid Mamba models -Under `--mamba-cache-mode align`, Mamba state is stored only on the Mamba block grid, so a prefix-cache hit can resume only at a block boundary. `--enable-mamba-shared-prefix-checkpoint` also stores a checkpoint at the shared-prefix junction, the point where an earlier request with the same prefix stopped. Requests whose shared prefix ends inside a block can then reuse it. +With prefix caching enabled, Mamba state is stored only on the Mamba block grid, so a prefix-cache hit can resume only at a block boundary. `--enable-mamba-shared-prefix-checkpoint` also stores a checkpoint at the shared-prefix junction, the point where an earlier request with the same prefix stopped. Requests whose shared prefix ends inside a block can then reuse it. This helps when many requests share a long system prompt and then diverge. It is off by default, and takes effect only when all of the following hold: -- `--mamba-cache-mode align` +- prefix caching enabled - EAGLE/MTP speculative decoding on the Mamba group - `--prefix-match-unit` smaller than the Mamba block size - the model does not use multi-module MTP ```bash vllm serve \ - --mamba-cache-mode align \ --prefix-match-unit 64 \ --enable-mamba-shared-prefix-checkpoint ``` -`--prefix-match-unit` is required. It sets the granularity at which prefix-cache keys are computed. When unset it defaults to the greatest common divisor of the prefix-cacheable KV cache group block sizes. Under `align` that is the block size itself, so no sub-block boundary exists and the flag has no effect. +`--prefix-match-unit` is required. It sets the granularity at which prefix-cache keys are computed. When unset it defaults to the greatest common divisor of the prefix-cacheable KV cache group block sizes. For hybrid Mamba models that is the block size itself, so no sub-block boundary exists and the flag has no effect. Choose a value that divides the block size of every prefix-cacheable KV cache group, and that is a multiple of the per-state compression ratio for models that use one, such as sparse MLA. vLLM validates both at startup and names the offending sizes in the error. Read the served block size from the startup log. 64 is a reasonable starting point. diff --git a/tests/distributed/test_kimi_linear_context_parallel.py b/tests/distributed/test_kimi_linear_context_parallel.py index 06348074bb14..8bddf6e609c0 100644 --- a/tests/distributed/test_kimi_linear_context_parallel.py +++ b/tests/distributed/test_kimi_linear_context_parallel.py @@ -180,7 +180,6 @@ def _run_k3_partial_prefix_reuse( block_size=block_size, enable_chunked_prefill=True, enable_prefix_caching=True, - mamba_cache_mode="align", ) as runner: assert ( runner.llm.llm_engine.vllm_config.compilation_config.cudagraph_mode diff --git a/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml b/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml index db43869eb4ae..c35614cd5f8c 100644 --- a/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml +++ b/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml @@ -23,4 +23,4 @@ generation_config: extra_body: chat_template_kwargs: enable_thinking: true -server_args: "--tensor-parallel-size 4 --mamba-cache-mode align" +server_args: "--tensor-parallel-size 4" diff --git a/tests/kernels/mamba/test_mamba_mixer2.py b/tests/kernels/mamba/test_mamba_mixer2.py index 973e7885c680..db2274c3c933 100644 --- a/tests/kernels/mamba/test_mamba_mixer2.py +++ b/tests/kernels/mamba/test_mamba_mixer2.py @@ -96,15 +96,17 @@ def mixer2_gated_norm_tensor_parallel( gate_states = torch.randn(batch_size, seq_len, hidden_size) # create gated-norm with TP - mixer = Mixer2RMSNormGated( - full_hidden_size=hidden_size, - full_n_groups=n_groups, - ) + with ensure_current_vllm_config(): + mixer = Mixer2RMSNormGated( + full_hidden_size=hidden_size, + full_n_groups=n_groups, + ) mixer.weight.weight_loader(mixer.weight, weight) # load # create gated-norm without TP to compute reference # - utilize mock patching to disable TP when with ( + ensure_current_vllm_config(), unittest.mock.patch( "vllm.model_executor.layers.mamba.mamba_mixer2." "get_tensor_model_parallel_world_size", diff --git a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py index 324aef425b89..54ae43e8e232 100644 --- a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py +++ b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py @@ -38,9 +38,7 @@ def _manager( num_prefill_lookahead=0, ): init_none_hash(sha256) - config = _make_hybrid_kv_cache_config( - block_size, num_blocks, ["full", "mamba_align"] - ) + config = _make_hybrid_kv_cache_config(block_size, num_blocks, ["full", "mamba"]) if eagle_group is not None: groups = list(config.kv_cache_groups) groups[eagle_group] = replace(groups[eagle_group], is_eagle_group=True) diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 64cd55054015..f71fe160d83d 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -278,7 +278,14 @@ def test_scratch_group_survives_computed_block_truncation(self): neither divides the hit length nor bounds the (empty) block list.""" config = _shared_layout_config() config.cache_config.enable_prefix_caching = True - groups = get_kv_cache_groups(config, _make_csa_linear_specs()) + config.cache_config.mamba_cache_mode = "align" + specs = { + name: replace(spec, mamba_cache_mode="align") + if isinstance(spec, MambaSpec) + else spec + for name, spec in _make_csa_linear_specs().items() + } + groups = get_kv_cache_groups(config, specs) kv_cache_config = get_kv_cache_config_from_groups( config, groups, available_memory=BYTES_PER_BLOCK * 64 ) diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 836286d127fa..52455c608305 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -1617,6 +1617,7 @@ def make_kv_cache_config_hybrid_model( block_size=block_size, shapes=((1, 1),), dtypes=(torch.float32,), + mamba_cache_mode="align", ) return KVCacheConfig( @@ -1652,6 +1653,7 @@ def make_kv_cache_config_three_types( block_size=block_size, shapes=((1, 1),), dtypes=(torch.float32,), + mamba_cache_mode="align", ) elif third_spec_type == "sliding_window": third_spec = SlidingWindowSpec( @@ -1732,6 +1734,7 @@ def test_prefix_cache_hit_uses_per_group_dcp_geometry(draft_sharded): block_size=block_size, shapes=(1, 1), dtypes=(torch.float32,), + mamba_cache_mode="align", ), ), ], @@ -1757,10 +1760,13 @@ def test_prefix_cache_hit_uses_per_group_dcp_geometry(draft_sharded): req0 = make_request("0", common_token_ids + [99] * 7, block_size, hash_fn) computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0) assert num_computed_tokens == 0 - blocks = manager.allocate_slots( - req0, len(req0.prompt_token_ids), 0, computed_blocks - ) - assert blocks is not None + # Align-mode Mamba keeps a state only where a chunk ends, so end the first + # chunk on the shared prefix. + for num_new_tokens in (len(common_token_ids), 7): + blocks = manager.allocate_slots(req0, num_new_tokens, 0, computed_blocks) + assert blocks is not None + req0.num_computed_tokens += num_new_tokens + computed_blocks = None req1 = make_request("1", common_token_ids + [100] * 5, block_size, hash_fn) computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1) @@ -2285,7 +2291,7 @@ def _make_hybrid_kv_cache_config( - "full": FullAttentionSpec - "sliding_window": SlidingWindowSpec with window=2*block_size - "sliding_window_large": SlidingWindowSpec with window=4*block_size - - "mamba": MambaSpec + - "mamba": MambaSpec in prefix-caching ("align") mode """ spec_map = { @@ -2313,11 +2319,6 @@ def _make_hybrid_kv_cache_config( block_size=block_size, shapes=((1, 1),), dtypes=(torch.float32,), - ), - "mamba_align": lambda: MambaSpec( - block_size=block_size, - shapes=((1, 1),), - dtypes=(torch.float32,), mamba_cache_mode="align", ), } @@ -2427,14 +2428,16 @@ def test_prefill_hybrid_model_combinations(spec_types: list[str]): assert not computed_blocks.blocks[0] # No cache hit initially assert num_computed_tokens == 0 - blocks = manager.allocate_slots( - req0, 55, len(computed_blocks.blocks[0]) * block_size, computed_blocks - ) - assert blocks is not None - # Should have blocks for all groups - assert len(blocks.get_block_ids()) == num_groups - - manager.new_step_starts() + # Align-mode Mamba keeps a state only where a chunk ends, so end the first + # chunk on the common prefix. + for num_new_tokens in (len(common_token_ids), len(unique_token_ids)): + blocks = manager.allocate_slots(req0, num_new_tokens, 0, computed_blocks) + assert blocks is not None + # Should have blocks for all groups + assert len(blocks.get_block_ids()) == num_groups + req0.num_computed_tokens += num_new_tokens + computed_blocks = None + manager.new_step_starts() # Second request: should hit cached blocks for common prefix req1 = make_request("1", common_token_ids + [4] * 5, block_size, hash_fn) @@ -2549,7 +2552,7 @@ def test_prefill_hybrid_model_mamba_align(): num_blocks = 30 kv_cache_config = _make_hybrid_kv_cache_config( - block_size, num_blocks, ["full", "mamba_align"] + block_size, num_blocks, ["full", "mamba"] ) manager = make_kv_cache_manager( kv_cache_config, @@ -2585,7 +2588,7 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection(): """ block_size = 16 manager = make_kv_cache_manager( - _make_hybrid_kv_cache_config(block_size, 30, ["full", "mamba_align"]), + _make_hybrid_kv_cache_config(block_size, 30, ["full", "mamba"]), max_model_len=8192, enable_caching=True, hash_block_size=block_size, @@ -2660,7 +2663,7 @@ def test_hybrid_model_mamba_align_with_dynamic_draft_tokens(): num_blocks = 30 kv_cache_config = _make_hybrid_kv_cache_config( - block_size, num_blocks, ["full", "mamba_align"] + block_size, num_blocks, ["full", "mamba"] ) manager = KVCacheManager( kv_cache_config, @@ -5896,49 +5899,6 @@ def retained(retention_interval, shared_prefix_boundary, end_block=16): assert retained(0, None) == {14} -def test_mamba_shared_prefix_survives_zero_retention(): - """Manager-level check of the full wiring: a pinned shared-prefix boundary - (``Request.shared_prefix_boundary``, set by the scheduler on Marconi-style - detection) keeps its Mamba state block cached under - ``prefix_cache_retention_interval=0``, which otherwise retains only the - end-of-prompt replay boundary. Without this, a shared prefix (junction - before ``num_prompt``) would be recomputed by every sharing request.""" - block_size = 16 - - # 16-block (256-token) prompt; replay boundary is block 240 // 16 - 1 = 14. - token_ids = [i for i in range(16) for _ in range(block_size)] - - def cached_mamba_blocks(shared_prefix_boundary): - # Fresh manager per scenario so cached blocks don't leak between runs. - manager = make_kv_cache_manager( - _make_hybrid_kv_cache_config(block_size, 100, ["full", "mamba"]), - max_model_len=8192, - enable_caching=True, - hash_block_size=block_size, - retention_interval=0, - ) - req = make_request("r", token_ids, block_size, sha256) - req.shared_prefix_boundary = shared_prefix_boundary - computed_blocks, num_computed, _ = manager.get_computed_blocks(req) - blocks = manager.allocate_slots( - req, len(token_ids), num_computed, computed_blocks - ) - assert blocks is not None - pool = manager.block_pool - return { - i - for i in range(16) - if pool.get_cached_block(req.block_hashes[i], kv_cache_group_ids=[1]) - is not None - } - - # Without a pinned boundary, retention=0 keeps only the replay boundary (14). - assert cached_mamba_blocks(0) == {14} - # Pinning the shared prefix at token 96 (state block 5) retains it too, so a - # later request sharing that prefix can hit the Mamba state. - assert cached_mamba_blocks(96) == {5, 14} - - def test_mamba_shared_prefix_reuse_under_zero_retention(): """Full cross-request Marconi flow: a partial shared prefix cached by the detecting request must stay reusable by a later request under @@ -5949,7 +5909,7 @@ def test_mamba_shared_prefix_reuse_under_zero_retention(): def last_req_hit(retention, pin): manager = make_kv_cache_manager( - _make_hybrid_kv_cache_config(block_size, 200, ["full", "mamba_align"]), + _make_hybrid_kv_cache_config(block_size, 200, ["full", "mamba"]), max_model_len=8192, enable_caching=True, hash_block_size=block_size, diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index 50f34737085d..6c208e19fb5d 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -898,7 +898,6 @@ def _run_mamba_prefix_cache_mrv1( load_format="dummy", enable_prefix_caching=True, block_size=BLOCK_SIZE, - mamba_cache_mode="align", speculative_config={ "method": "qwen3_next_mtp", "num_speculative_tokens": num_speculative_tokens, @@ -1158,7 +1157,6 @@ def fake_sample( skip_tokenizer_init=True, enable_prefix_caching=True, block_size=BLOCK_SIZE, - mamba_cache_mode="align", speculative_config={ "method": "qwen3_next_mtp", "num_speculative_tokens": num_speculative_tokens, diff --git a/tests/v1/e2e/test_hybrid_chunked_prefill.py b/tests/v1/e2e/test_hybrid_chunked_prefill.py index 66c8d5b9b76c..c63a4ebc7893 100644 --- a/tests/v1/e2e/test_hybrid_chunked_prefill.py +++ b/tests/v1/e2e/test_hybrid_chunked_prefill.py @@ -88,7 +88,6 @@ def test_mtp_speculative_mixed_batch_short_prefill( trust_remote_code=True, enable_chunked_prefill=True, enable_prefix_caching=enable_prefix_caching, - mamba_cache_mode="align" if enable_prefix_caching else "none", ) as llm: sampling_params = SamplingParams( temperature=0.0, @@ -161,7 +160,6 @@ def test_qwen_mtp_mamba_prefix_cache_hit_is_bounded(vllm_runner, monkeypatch): enforce_eager=True, enable_chunked_prefill=True, enable_prefix_caching=True, - mamba_cache_mode="align", speculative_config={"method": "mtp", "num_speculative_tokens": 2}, ) as runner: llm = runner.get_llm() diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 89797db4cdd7..f75151414afd 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -60,7 +60,6 @@ def _check_replayssm_parity( max_model_len=1024, trust_remote_code=True, enable_prefix_caching=False, - mamba_cache_mode="none", tensor_parallel_size=tensor_parallel_size, mamba_backend=mamba_backend, ) @@ -130,7 +129,6 @@ def test_replayssm_flashinfer_spec_decode_matches_baseline( max_model_len=1024, trust_remote_code=True, enable_prefix_caching=False, - mamba_cache_mode="none", speculative_config={ "method": "ngram", "num_speculative_tokens": 3, @@ -182,7 +180,6 @@ def test_replayssm_flashinfer_mtp(vllm_runner, monkeypatch, use_v2_model_runner) max_model_len=1024, trust_remote_code=True, enable_prefix_caching=False, - mamba_cache_mode="none", mamba_backend="flashinfer", tensor_parallel_size=2, disable_log_stats=False, @@ -250,7 +247,6 @@ def _check_replayssm_prefix_caching_parity( trust_remote_code=True, enable_prefix_caching=True, enable_chunked_prefill=True, - mamba_cache_mode="align", disable_log_stats=False, # required for llm.get_metrics() tensor_parallel_size=tensor_parallel_size, ) diff --git a/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh b/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh index c6f79f37211c..6615a5c103b1 100755 --- a/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh @@ -60,7 +60,6 @@ vllm serve "$MODEL" \ --block-size 128 \ --trust-remote-code \ --enable-prefix-caching \ - --mamba-cache-mode align \ --kv-transfer-config "$KV_CONFIG" \ "${EXTRA_ARGS[@]}" & @@ -78,7 +77,6 @@ vllm serve "$MODEL" \ --block-size 128 \ --trust-remote-code \ --enable-prefix-caching \ - --mamba-cache-mode align \ --kv-transfer-config "$KV_CONFIG" \ "${EXTRA_ARGS[@]}" & diff --git a/tests/v1/kv_connector/unit/test_offloading_connector.py b/tests/v1/kv_connector/unit/test_offloading_connector.py index 76d219746913..53e7378a7b5c 100644 --- a/tests/v1/kv_connector/unit/test_offloading_connector.py +++ b/tests/v1/kv_connector/unit/test_offloading_connector.py @@ -583,7 +583,6 @@ def test_mamba_cpu_offload_boundary(model: str, block_size: int, tp_size: int): kv_transfer_config=kv_transfer_config, language_model_only=True, enable_prefix_caching=True, - mamba_cache_mode="align", # Use lossless state storage so this exact-equality regression isolates # offload boundary selection from low-precision Mamba checkpointing. mamba_cache_dtype="float32", diff --git a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py index b72de811c0f9..74de3c3a8cf1 100644 --- a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py @@ -632,6 +632,8 @@ def test_p_side_chunked_prefill_mamba(mock_platform): block_size=BLOCK_SIZE, ) vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False + # mamba_cache_mode "none" only runs without prefix caching. + vllm_config.cache_config.enable_prefix_caching = False kv_cache_config = make_kv_cache_config( block_size=BLOCK_SIZE, @@ -855,6 +857,7 @@ def test_async_load_does_not_double_reserve_mamba_spec_blocks( disable_hybrid_kv_cache_manager=False, ) vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode + vllm_config.cache_config.enable_prefix_caching = mamba_cache_mode == "align" kv_cache_config = make_kv_cache_config( block_size, mamba_enabled=True, diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 9f818c866fff..527f7bee05c2 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -49,6 +49,7 @@ FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, + KVCacheSpec, KVCacheTensor, MambaSpec, SlidingWindowSpec, @@ -2355,8 +2356,6 @@ def _make_hybrid_attention_mamba_scheduler( enable_kv_cache_events: bool = False, ) -> SchedulerFixture: """Build a scheduler for one attention group plus one Mamba group.""" - scheduler_block_size = scheduler_block_size or block_size - hash_block_size = hash_block_size or block_size attention_spec = FullAttentionSpec( block_size=attention_block_size, num_kv_heads=NUM_KV_HEADS, @@ -2369,9 +2368,38 @@ def _make_hybrid_attention_mamba_scheduler( dtypes=(torch.float32,), mamba_cache_mode=mamba_cache_mode, ) + return _make_two_group_scheduler( + attention_spec, + ("mamba", mamba_spec), + num_cpu_blocks=num_cpu_blocks, + num_gpu_blocks=num_gpu_blocks, + scheduler_block_size=scheduler_block_size or block_size, + hash_block_size=hash_block_size or block_size, + dcp_world_size=dcp_world_size, + lazy=lazy, + mamba_cache_mode=mamba_cache_mode, + enable_kv_cache_events=enable_kv_cache_events, + ) + + +def _make_two_group_scheduler( + attention_spec: FullAttentionSpec, + second_group: tuple[str, KVCacheSpec], + *, + num_cpu_blocks: int, + num_gpu_blocks: int, + scheduler_block_size: int, + hash_block_size: int, + dcp_world_size: int, + lazy: bool = False, + mamba_cache_mode: MambaCacheMode = "align", + enable_kv_cache_events: bool = False, +) -> SchedulerFixture: + """Build a scheduler for one attention group plus one named second group.""" + second_name, second_spec = second_group groups = [ KVCacheGroupSpec(["attention"], attention_spec), - KVCacheGroupSpec(["mamba"], mamba_spec), + KVCacheGroupSpec([second_name], second_spec), ] tensors = [ KVCacheTensor( @@ -2380,7 +2408,7 @@ def _make_hybrid_attention_mamba_scheduler( layer_stride=spec.page_size_bytes * num_gpu_blocks, block_stride=spec.page_size_bytes, ) - for group, spec in zip(groups, (attention_spec, mamba_spec)) + for group, spec in zip(groups, (attention_spec, second_spec)) ] kv_cache_config = KVCacheConfig( num_blocks=num_gpu_blocks, @@ -2428,34 +2456,48 @@ def test_hybrid_store_uses_resolved_group_block_sizes() -> None: silently offloads only a fraction of them. """ attention_block_size = BLOCK_SIZE - mamba_block_size = 4 * BLOCK_SIZE + replicated_block_size = 4 * BLOCK_SIZE dcp_world_size = 2 - fix = _make_hybrid_attention_mamba_scheduler( + fix = _make_two_group_scheduler( + FullAttentionSpec( + block_size=attention_block_size, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + ), + ( + "replicated_attention", + FullAttentionSpec( + block_size=replicated_block_size, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + dcp_sharded=False, + ), + ), num_cpu_blocks=32, num_gpu_blocks=32, - attention_block_size=attention_block_size, - block_size=mamba_block_size, + scheduler_block_size=replicated_block_size, # Every prefix-cacheable group's resolved block size must be a multiple # of the hash block, so it cannot exceed the smallest of them. hash_block_size=BLOCK_SIZE, dcp_world_size=dcp_world_size, - mamba_cache_mode="none", ) sched = fix.scheduler gpu_pool = fix.gpu_block_pool - attention_size, mamba_size = sched.group_block_sizes + attention_size, replicated_size = sched.group_block_sizes assert attention_size == attention_block_size * dcp_world_size - assert mamba_size == mamba_block_size + assert replicated_size == replicated_block_size confirmed = 2 * sched.block_size req = _make_cp_request(num_blocks=8, virtual_block_size=BLOCK_SIZE) attention_blocks = _allocate_cp_gpu_blocks( gpu_pool, req, confirmed // attention_size, attention_size, group_id=0 ) - mamba_blocks = _allocate_cp_gpu_blocks( - gpu_pool, req, confirmed // mamba_size, mamba_size, group_id=1 + replicated_blocks = _allocate_cp_gpu_blocks( + gpu_pool, req, confirmed // replicated_size, replicated_size, group_id=1 ) - kv_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks)) + kv_blocks = KVCacheBlocks(blocks=(attention_blocks, replicated_blocks)) sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) sched.build_connector_meta( make_scheduler_output( @@ -2476,11 +2518,11 @@ def test_hybrid_store_uses_resolved_group_block_sizes() -> None: # count. assert sched._reqs_to_store[req.request_id].num_stored_blocks == [ confirmed // attention_size, - confirmed // mamba_size, + confirmed // replicated_size, ] assert set(meta.store_gpu_blocks) == { *(b.block_id for b in attention_blocks), - *(b.block_id for b in mamba_blocks), + *(b.block_id for b in replicated_blocks), } diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 5eb26726c718..9bb716621594 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -176,8 +176,7 @@ def device_memory_utilization(self, value: float) -> None: page.""" mamba_block_size: int | None = Field(default=None, gt=0) """Size of a contiguous cache block in number of tokens for mamba cache. - Can be set only when prefix caching is enabled. - Value must be a multiple of 8 to align with causal_conv1d kernel.""" + Can be set only when prefix caching is enabled.""" mamba_cache_dtype: MambaDType = "auto" """The data type to use for the Mamba cache (both the conv as well as the ssm state). If set to 'auto', the data type will be inferred from the model @@ -187,19 +186,19 @@ def device_memory_utilization(self, value: float) -> None: still be controlled by mamba_cache_dtype). If set to 'auto', the data type for the ssm state will be determined by mamba_cache_dtype.""" mamba_cache_mode: MambaCacheMode = "none" - """The cache strategy for Mamba layers: + """The cache strategy for Mamba layers, derived from `enable_prefix_caching`: - - "none": set when prefix caching is disabled. - - "align": only cache the mamba state of the last token of each scheduler step and - when the token is at position i * block_size. This is the default when prefix - caching is enabled. + - "none": prefix caching is disabled. + - "align": prefix caching is enabled. Only cache the mamba state of the last + token of each scheduler step and when the token is at position + i * block_size. """ enable_mamba_shared_prefix_checkpoint: bool = False """Also register a Mamba "align" checkpoint at the shared-prefix junction -- where an EAGLE/MTP sibling was observed to resume -- instead of only at the - prompt tail. Off by default; only takes effect with `mamba_cache_mode` - "align", EAGLE on the Mamba group, and a prefix match unit smaller than the - Mamba block size.""" + prompt tail. Off by default; only takes effect with prefix caching enabled, + EAGLE on the Mamba group, and a prefix match unit smaller than the Mamba + block size.""" replayssm_buffer_len: int = Field(default=16, gt=0) """ReplaySSM logical history length B for Mamba2. Triton uses B physical rows and FlashInfer uses B+T, where T is the target verification length. @@ -207,11 +206,10 @@ def device_memory_utilization(self, value: float) -> None: use_replayssm: bool = False """Use the ReplaySSM Mamba2 decode kernel: cache recent SSM inputs and skip the per-step full-state store, writing the checkpoint back only on flush. - Requires mamba_cache_mode 'none' or 'align' (prefix caching) and the Triton - or FlashInfer mamba backend. Mamba2 speculative decode requires FlashInfer - and mamba_cache_mode 'none'. In align mode flushes are most efficient when - mamba_block_size is a multiple of replayssm_buffer_len, but this is not - required.""" + Requires the Triton or FlashInfer mamba backend; FlashInfer does not support + prefix caching yet. Mamba2 speculative decode requires FlashInfer. With + prefix caching, flushes are most efficient when mamba_block_size is a + multiple of replayssm_buffer_len, but this is not required.""" use_kda_recoverssm: bool = field(default=False, init=False) """Whether Kimi-K3 KDA uses RecoverSSM speculative decode.""" diff --git a/vllm/config/model.py b/vllm/config/model.py index 87759f46061b..c4016e82bc3d 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -2038,10 +2038,6 @@ def has_noops(self) -> bool: def has_inner_state(self): return self._model_info.has_inner_state - @property - def supports_mamba_prefix_caching(self) -> bool: - return self._model_info.supports_mamba_prefix_caching - @property def supports_replayssm(self) -> bool: return self._model_info.supports_replayssm diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index f27c8418ab4d..d6e834861a19 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -3360,7 +3360,7 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": and not self.use_v2_model_runner ): raise ValueError( - "RecoverSSM with align mode requires VLLM_USE_V2_MODEL_RUNNER=1" + "RecoverSSM with prefix caching requires VLLM_USE_V2_MODEL_RUNNER=1" ) if self.parallel_config.pipeline_parallel_size > 1: raise ValueError( @@ -3371,8 +3371,8 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": elif use_mamba_replayssm_spec: if self.cache_config.mamba_cache_mode != "none": raise ValueError( - "FlashInfer ReplaySSM speculative decoding requires " - "--mamba-cache-mode none" + "FlashInfer ReplaySSM speculative decoding does not support " + "prefix caching; use --no-enable-prefix-caching" ) query_len = 1 + self.num_speculative_tokens if self.cache_config.replayssm_buffer_len < query_len: @@ -3390,8 +3390,8 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": elif self.mamba_config.backend == MambaBackendEnum.FLASHINFER: if self.cache_config.mamba_cache_mode == "align": raise ValueError( - "FlashInfer ReplaySSM does not support " - "--mamba-cache-mode align yet; use none" + "FlashInfer ReplaySSM does not support prefix caching yet; " + "use --no-enable-prefix-caching" ) elif self.mamba_config.backend != MambaBackendEnum.TRITON: raise ValueError( diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 86199e5f007b..c9aa433162d2 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -453,6 +453,12 @@ def get_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]: ) +_MAMBA_CACHE_MODE_DEPRECATION_MESSAGE = ( + "--mamba-cache-mode is deprecated and will be removed in v0.33.0. " + "It has no effect: Mamba prefix caching follows --enable-prefix-caching." +) + + _DISABLE_CASCADE_ATTN_DEPRECATION_MESSAGE = ( "--disable-cascade-attn is deprecated and has no effect since cascade " "attention has been removed. It will be removed in v0.33.0." @@ -798,7 +804,7 @@ class EngineArgs: mamba_ssm_cache_dtype: MambaDType = CacheConfig.mamba_ssm_cache_dtype mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size") prefix_match_unit: int | None = get_field(CacheConfig, "prefix_match_unit") - mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode + mamba_cache_mode: MambaCacheMode | None = None enable_mamba_shared_prefix_checkpoint: bool = ( CacheConfig.enable_mamba_shared_prefix_checkpoint ) @@ -1403,7 +1409,12 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--prefix-match-unit", **cache_kwargs["prefix_match_unit"] ) cache_group.add_argument( - "--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"] + "--mamba-cache-mode", + **{ + **cache_kwargs["mamba_cache_mode"], + "default": None, + "help": _MAMBA_CACHE_MODE_DEPRECATION_MESSAGE, + }, ) cache_group.add_argument( "--enable-mamba-shared-prefix-checkpoint", @@ -2263,6 +2274,8 @@ def create_engine_config( assert self.enable_prefix_caching is not None, ( "enable_prefix_caching must be set by this point" ) + if self.mamba_cache_mode is not None: + logger.warning_once(_MAMBA_CACHE_MODE_DEPRECATION_MESSAGE) cache_config = CacheConfig( block_size=self.block_size, # type: ignore[arg-type] @@ -2282,7 +2295,6 @@ def create_engine_config( mamba_ssm_cache_dtype=self.mamba_ssm_cache_dtype, mamba_block_size=self.mamba_block_size, prefix_match_unit=self.prefix_match_unit, - mamba_cache_mode=self.mamba_cache_mode, enable_mamba_shared_prefix_checkpoint=( self.enable_mamba_shared_prefix_checkpoint ), diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index bfa1103bb368..14c483122ab1 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -812,10 +812,8 @@ def conv_ssm_forward( # - It will read the initial states for every sequence, # that has "has_initial_states_p" == True, # from "cache_indices", using "state_indices_tensor_p". - # - It updates the "conv_state" cache in positions pointed - # to by "state_indices_tensor_p". - # In particular, it will always write the state at the - # sequence end. + # - It writes the state at the sequence end to the "conv_state" + # cache in positions pointed to by "state_indices_tensor_p". x = hidden_states_B_C_p.transpose( 0, 1 ) # this is the form that causal-conv see diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index e8254e4f34eb..2094df902a27 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -20,17 +20,11 @@ def _causal_conv1d_fwd_kernel( # continuous batching w_ptr, # (dim, width) bias_ptr, initial_states_ptr, # conv_states_ptr - cache_indices_ptr, # (batch, n_blocks + padding) The second dimension contains - # the block indices relevant for each sequence - # plus potential 0-padding at the beginning and at the end + cache_indices_ptr, # (batch,) has_initial_states_ptr, query_start_loc_ptr, batch_ptr, token_chunk_offset_ptr, - block_idx_first_scheduled_token, # (batch,) - block_idx_last_scheduled_token, # (batch,) - initial_state_idx, # (batch,) - num_computed_tokens, # (batch,) o_ptr, # (dim, seqlen) - actually pointing to x_ptr # Matrix dimensions dim: tl.constexpr, @@ -46,7 +40,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching stride_cache_indices: tl.constexpr, stride_o_dim: tl.constexpr, stride_o_token: tl.int64, - stride_block_m: tl.constexpr, # Stride block to align divided by BLOCK_M # others pad_slot_id: tl.constexpr, null_block_id: tl.constexpr, @@ -54,7 +47,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching HAS_BIAS: tl.constexpr, KERNEL_WIDTH: tl.constexpr, SILU_ACTIVATION: tl.constexpr, - IS_APC_ENABLED: tl.constexpr, HAS_NULL_BLOCK: tl.constexpr, NP2_STATELEN: tl.constexpr, BLOCK_M: tl.constexpr, @@ -93,40 +85,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching # find the actual sequence length seqlen = sequence_end_index - sequence_start_index - B_size: tl.constexpr = stride_block_m * BLOCK_M - - if IS_APC_ENABLED: - # Handle the case if prefix caching is enabled. - # In particular, if prefix caching is enabled, the program write additional cache states to "cache_indices_ptr" - - # Get the length of the completed sequence so far and compute the offset. - current_first_index = tl.load(block_idx_first_scheduled_token + idx_seq) - current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq) - sequence_completed_index = tl.load(num_computed_tokens + idx_seq) - - # Compute the offset where the first stride_block_m-aligned first full block is - # Value in "token-space" - sequence_completed_offset_token = sequence_completed_index % B_size - seq_completed_offset = B_size - sequence_completed_offset_token - seq_end_offset = (seqlen - seq_completed_offset) % B_size - last_full_block_token_index = sequence_end_index - seq_end_offset - # If the sequence without the sequence_offset_index is stride_cache_chunk-aligned, then the last full chunk is the second-to-last one - if seq_end_offset == 0: - last_full_block_token_index = last_full_block_token_index - B_size - - # Get the number of blocks to be filled for the current sequence - # If n_block_to_fill = 0, then only the state at the sequence end is stored - n_block_to_fill = current_last_index - current_first_index - - # Get the index of the init block - conv_state_init_index = tl.load(initial_state_idx + idx_seq) - else: - n_block_to_fill = 0 - current_last_index = 0 - conv_state_init_index = 0 - current_first_index = 0 - last_full_block_token_index = 0 - token_offset = BLOCK_M * chunk_offset segment_len = min(BLOCK_M, seqlen - token_offset) @@ -137,7 +95,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching # cache_idx conv_states_input_coord = tl.load( - conv_state_indices_ptr + idx_seq * stride_cache_indices + conv_state_init_index + conv_state_indices_ptr + idx_seq * stride_cache_indices ).to(tl.int64) if HAS_NULL_BLOCK: # noqa @@ -224,20 +182,10 @@ def _causal_conv1d_fwd_kernel( # continuous batching loaded_x = tl.load(x_ptrs, mask_x, 0.0) idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] - # Compute the offset where the last block should be written in the conv_states - conv_states_output_coord = tl.load( - conv_state_indices_ptr - + idx_seq * stride_cache_indices - + current_last_index - ).to(tl.int64) - conv_states_ptrs_target = ( - conv_states_ptr - + (conv_states_output_coord * stride_conv_state_seq) # Offset from seq - + (idx_feats * stride_conv_state_dim) - )[None, :] + ( # [BLOCK_N,] - idx_tokens_conv * stride_conv_state_tok - )[:, None] + conv_states_base[None, :] + + (idx_tokens_conv * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] tl.debug_barrier() # NOTE: use this due to bug in Triton compiler @@ -346,52 +294,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching conv_states_ptrs = prior_tokens - 3 * stride_x_token # [BLOCK_N] col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") - # Store intermediate states aligned with stride_block_m - # The additional states are cached starting from the last stride_block_m. - # For example: - # If n_block_to_fill = 0, then only the state at the sequence end is cached and the process below is not involved. - # If n_block_to_fill > 0, then the states at the sequence end and at the n_block_to_fill-last - # stride_block_m are cached. - # For example chunk_offset = n_block_to_fill stores the state at last_full_block - if (chunk_offset - 1) < n_block_to_fill: - # Store the states at the chunk boundaries from the start of the sequence - idx_tokens_last = ( - last_full_block_token_index - - (n_block_to_fill - chunk_offset) * B_size - - state_len - ) + tl.arange(0, NP2_STATELEN) # [BLOCK_M] - x_ptrs = ( - x_ptr - + (idx_tokens_last * stride_x_token)[:, None] - + (idx_feats * stride_x_dim)[None, :] - ) # [BLOCK_M,BLOCK_N,] - - mask_x = (idx_tokens_last >= 0)[:, None] & (idx_feats < dim)[ - None, : - ] # token-index # token-index # feature-index - loaded_x = tl.load(x_ptrs, mask_x, 0.0) - idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] - - # cache_idx - conv_states_output_coord = tl.load( - conv_state_indices_ptr - + idx_seq * stride_cache_indices - + current_first_index - + (chunk_offset - 1) - ).to(tl.int64) - - conv_states_ptrs_target = ( - conv_states_ptr - + (conv_states_output_coord * stride_conv_state_seq) # Offset from seq - + (idx_feats * stride_conv_state_dim) - )[None, :] + ( # [BLOCK_N,] - idx_tokens_conv * stride_conv_state_tok - )[:, None] - - mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] - tl.debug_barrier() # NOTE: use this due to bug in Triton compiler - tl.store(conv_states_ptrs_target, loaded_x, mask) - if HAS_BIAS: bias = bias_ptr + idx_feats mask_bias = idx_feats < dim @@ -489,11 +391,6 @@ def causal_conv1d_fn( activation: str | None = "silu", pad_slot_id: int = PAD_SLOT_ID, null_block_id: int = NULL_BLOCK_ID, - block_idx_first_scheduled_token: torch.Tensor | None = None, - block_idx_last_scheduled_token: torch.Tensor | None = None, - initial_state_idx: torch.Tensor | None = None, - num_computed_tokens: torch.Tensor | None = None, - block_size_to_align=0, metadata=None, validate_data=False, ): @@ -536,16 +433,6 @@ def causal_conv1d_fn( for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id] in this case, the kernel will not process entries at indices 0 and 3 - block_idx_first_scheduled_token: (batch,), dtype int32 - The pointer into cache_indices, where the first cache block to be filled is located. - block_idx_last_scheduled_token: (batch,), dtype int32 - The pointer into cache_indices, where the last cache block to be filled is located. - initial_state_idx: (batch,), dtype int32 - The pointer into cache_indices, where the cache block containing the initial state is located. - num_computed_tokens: (batch,), dtype int32 - The number of tokens already completed for each sequence - block_size_to_align: int - The block size to align the cached states to out: same shape as `x` """ if isinstance(activation, bool) and activation: @@ -631,12 +518,6 @@ def causal_conv1d_fn( assert weight.stride(1) == 1 assert (dim, width) == weight.shape assert is_channel_last, "Need to run in channel-last layout" - if block_size_to_align is not None and block_size_to_align > 0: - assert (block_size_to_align % BLOCK_M) == 0, ( - "The mamba block size needs to be divisible by the BLOCK_M" - ) - else: - block_size_to_align = BLOCK_M if metadata is None: @@ -720,10 +601,6 @@ def grid(META): query_start_loc, batch_ptr, token_chunk_offset_ptr, - block_idx_first_scheduled_token, - block_idx_last_scheduled_token, - initial_state_idx, - num_computed_tokens, out, # Matrix dimensions dim, @@ -739,7 +616,6 @@ def grid(META): stride_cache_indices, stride_o_dim, stride_o_token, - block_size_to_align // BLOCK_M, # others pad_slot_id, null_block_id, @@ -747,7 +623,6 @@ def grid(META): HAS_BIAS=bias is not None, KERNEL_WIDTH=width, SILU_ACTIVATION=activation in ["silu", "swish"], - IS_APC_ENABLED=block_idx_last_scheduled_token is not None, HAS_NULL_BLOCK=null_block_id is not None, NP2_STATELEN=np2_statelen, # launch_cooperative_grid=True @@ -769,8 +644,6 @@ def _causal_conv1d_update_kernel( conv_state_indices_ptr, num_accepted_tokens_ptr, query_start_loc_ptr, # (batch + 1) - block_idx_last_scheduled_token, # (batch,) - initial_state_idx, # (batch,) o_ptr, # (batch, dim, seqlen) # Matrix dimensions batch: int, @@ -798,7 +671,6 @@ def _causal_conv1d_update_kernel( KERNEL_WIDTH: tl.constexpr, SILU_ACTIVATION: tl.constexpr, IS_VARLEN: tl.constexpr, - IS_APC_ENABLED: tl.constexpr, IS_SPEC_DECODING: tl.constexpr, NP2_STATELEN: tl.constexpr, HAS_NULL_BLOCK: tl.constexpr, @@ -818,17 +690,9 @@ def _causal_conv1d_update_kernel( # [BLOCK_N,] elements along the feature-dimension (channel) idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) - if IS_APC_ENABLED: - # Get the state from the initial_state_idx - conv_state_init = tl.load(initial_state_idx + idx_seq) - current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq) - else: - conv_state_init = 0 - current_last_index = 0 - # cache_idx conv_states_input_coord = tl.load( - conv_state_indices_ptr + idx_seq * stride_state_indices + conv_state_init + conv_state_indices_ptr + idx_seq * stride_state_indices ).to(tl.int64) if HAS_NULL_BLOCK: # noqa @@ -941,14 +805,9 @@ def _causal_conv1d_update_kernel( new_conv_state = tl.where(mask, conv_state, loaded_x) - # Get the state from the initial_state_idx - # cache_idx - conv_states_offset = tl.load( - conv_state_indices_ptr + idx_seq * stride_state_indices + current_last_index - ).to(tl.int64) conv_state_ptrs_target = ( conv_state_ptr - + (conv_states_offset * stride_conv_state_seq) # Offset from seq + + (conv_states_input_coord * stride_conv_state_seq) # Offset from seq + (idx_feats * stride_conv_state_dim) )[None, :] + ( # [BLOCK_N,] idx_tokens * stride_conv_state_tok @@ -1104,8 +963,6 @@ def causal_conv1d_update( query_start_loc: torch.Tensor | None = None, max_query_len: int = -1, null_block_id: int = NULL_BLOCK_ID, - block_idx_last_scheduled_token: torch.Tensor | None = None, - initial_state_idx: torch.Tensor | None = None, validate_data=False, out: torch.Tensor | None = None, ): @@ -1123,10 +980,6 @@ def causal_conv1d_update( If not None, the conv_state is a larger tensor along the batch dim, and we are selecting the batch coords specified by conv_state_indices. Useful for a continuous batching scenario. - block_idx_last_scheduled_token: (batch,), dtype int32 - The pointer into conv_state_indices, where the last cache block to be filled is located. - initial_state_idx: (batch,), dtype int32 - The pointer into conv_state_indices, where the cache block containing the initial state is located. num_accepted_tokens: (batch,), dtype int32 If not None, it indicates the number of accepted tokens for each sequence in the batch. @@ -1237,8 +1090,6 @@ def grid(META): conv_state_indices, num_accepted_tokens, query_start_loc, - block_idx_last_scheduled_token, - initial_state_idx, out, # Matrix dimensions batch, @@ -1266,7 +1117,6 @@ def grid(META): KERNEL_WIDTH=width, SILU_ACTIVATION=activation in ["silu", "swish"], IS_VARLEN=query_start_loc is not None, - IS_APC_ENABLED=block_idx_last_scheduled_token is not None, IS_SPEC_DECODING=num_accepted_tokens is not None, NP2_STATELEN=np2_statelen, HAS_NULL_BLOCK=null_block_id is not None, diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index d64e51bfa8bf..69ca643a3a97 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -692,7 +692,7 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: model_config.architecture, ) assert vllm_config.scheduler_config.enable_chunked_prefill, ( - "Chunked prefill is required for mamba cache mode 'align'." + "Chunked prefill is required for Mamba prefix caching." ) # By default, mamba block size will be set to max_model_len (see # below). When enabling prefix caching, we align mamba block size diff --git a/vllm/model_executor/models/falcon_h1.py b/vllm/model_executor/models/falcon_h1.py index 203a3f3ec393..7aff819dfcd7 100644 --- a/vllm/model_executor/models/falcon_h1.py +++ b/vllm/model_executor/models/falcon_h1.py @@ -42,7 +42,6 @@ HasInnerState, IsHybrid, SupportsLoRA, - SupportsMambaPrefixCaching, SupportsPP, ) from .utils import ( @@ -499,7 +498,6 @@ class FalconH1ForCausalLM( SupportsLoRA, SupportsPP, IsHybrid, - SupportsMambaPrefixCaching, ): packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], diff --git a/vllm/model_executor/models/granitemoehybrid.py b/vllm/model_executor/models/granitemoehybrid.py index 979d272c13d1..040f7453fb5b 100644 --- a/vllm/model_executor/models/granitemoehybrid.py +++ b/vllm/model_executor/models/granitemoehybrid.py @@ -42,7 +42,6 @@ HasInnerState, IsHybrid, SupportsLoRA, - SupportsMambaPrefixCaching, SupportsPP, SupportsQuant, ) @@ -587,7 +586,6 @@ class GraniteMoeHybridForCausalLM( SupportsPP, IsHybrid, SupportsQuant, - SupportsMambaPrefixCaching, ): packed_modules_mapping: dict[str, list[str]] = { "qkv_proj": [ diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 4dc67ff34d27..a780cbe23cf6 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -1134,13 +1134,10 @@ def has_noops( @runtime_checkable class SupportsMambaPrefixCaching(Protocol): - """The interface for models whose mamba layers support prefix caching. - - This is currently experimental. + """The interface for attention-free Mamba models that support prefix + caching. Hybrid models get the same methods from `IsHybrid`. """ - supports_mamba_prefix_caching: ClassVar[Literal[True]] = True - @classmethod def get_mamba_state_copy_func(cls) -> tuple["MambaStateCopyFunc", ...]: """Return copy functions for the model's Mamba states.""" @@ -1156,24 +1153,6 @@ def get_mamba_state_copy_funcs( return {mamba_type: copy_funcs for mamba_type in mamba_types} -@overload -def supports_mamba_prefix_caching( - model: type[object], -) -> TypeIs[type[SupportsMambaPrefixCaching]]: ... - - -@overload -def supports_mamba_prefix_caching( - model: object, -) -> TypeIs[SupportsMambaPrefixCaching]: ... - - -def supports_mamba_prefix_caching( - model: type[object] | object, -) -> TypeIs[type[SupportsMambaPrefixCaching]] | TypeIs[SupportsMambaPrefixCaching]: - return getattr(model, "supports_mamba_prefix_caching", False) - - @runtime_checkable class SupportsReplaySSM(Protocol): """The interface for models whose recurrent layers support ReplaySSM diff --git a/vllm/model_executor/models/jamba.py b/vllm/model_executor/models/jamba.py index fe9349c30b10..315c821d6a0e 100644 --- a/vllm/model_executor/models/jamba.py +++ b/vllm/model_executor/models/jamba.py @@ -44,7 +44,6 @@ HasInnerState, IsHybrid, SupportsLoRA, - SupportsMambaPrefixCaching, SupportsPP, ) from .utils import ( @@ -380,7 +379,6 @@ class JambaForCausalLM( SupportsLoRA, SupportsPP, IsHybrid, - SupportsMambaPrefixCaching, ): hf_to_vllm_mapper = WeightsMapper( orig_to_new_substr={".self_attn.": ".", ".A_log": ".A"}, diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index ae653950ac48..050a13431ddc 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -74,7 +74,6 @@ SupportsEagle, SupportsEagle3, SupportsLoRA, - SupportsMambaPrefixCaching, SupportsPP, SupportsQuant, SupportsReplaySSM, @@ -725,7 +724,6 @@ class NemotronHForCausalLM( IsHybrid, SupportsQuant, MixtureOfExperts, - SupportsMambaPrefixCaching, SupportsReplaySSM, ): # Relevant only if self.has_moe is True diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index cd0557ac4585..09e5025c4e0e 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -50,7 +50,6 @@ is_hybrid, requires_raw_input_tokens, supports_hc_sp, - supports_mamba_prefix_caching, supports_multimodal, supports_multimodal_encoder_tp_data, supports_multimodal_raw_input_only, @@ -881,7 +880,6 @@ class _ModelInfo: is_attention_free: bool is_hybrid: bool has_noops: bool - supports_mamba_prefix_caching: bool supports_replayssm: bool supports_transcription: bool supports_transcription_only: bool @@ -911,7 +909,6 @@ def from_model_cls(model: type[nn.Module]) -> "_ModelInfo": has_inner_state=has_inner_state(model), is_attention_free=is_attention_free(model), is_hybrid=is_hybrid(model), - supports_mamba_prefix_caching=supports_mamba_prefix_caching(model), supports_replayssm=supports_replayssm(model), supports_transcription=supports_transcription(model), supports_transcription_only=( diff --git a/vllm/model_executor/models/zamba2.py b/vllm/model_executor/models/zamba2.py index 2e90da86ea01..35e972818726 100644 --- a/vllm/model_executor/models/zamba2.py +++ b/vllm/model_executor/models/zamba2.py @@ -45,7 +45,7 @@ ) from vllm.sequence import IntermediateTensors -from .interfaces import HasInnerState, IsHybrid, SupportsMambaPrefixCaching +from .interfaces import HasInnerState, IsHybrid from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix @@ -833,7 +833,7 @@ def forward( return hidden_states -class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsMambaPrefixCaching): +class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid): """Zamba2 model with causal language modeling head. This class wraps the core Zamba2 model and adds: @@ -911,10 +911,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: LoRA and scheduler settings prefix: Optional prefix for parameter names - Raises: - AssertionError: If prefix caching is enabled - (not supported by Mamba) - """ config = vllm_config.model_config.hf_config diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 3e7ae905a30d..9820e8660e55 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -245,13 +245,10 @@ def _compute_chunk_metadata( The code below carefully constructs the chunks such that: 1. Chunks contain tokens from a *single* sequence only. - 2. For every sequence, we are guaranteed that we can - retrieve the mamba state *every* chunk_size tokens. + 2. Chunk boundaries fall on multiples of chunk_size in each + sequence's absolute token positions, so the chunking does not + depend on how chunked prefill split the sequence. Constraint (1) dramatically simplifies the mamba kernels. - Constraint (2) dramatically simplifies the implementation - of prefix caching for mamba (wip). We need to take care - of the interaction with chunked prefill in order to - satisfy constraint (2). `checkpoint_offsets_p[i]` forces an extra chunk end that many tokens into row i's query (0 = none), because the SSD scan only materializes diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 2e4fe55b96c2..28a7a1106f77 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -1426,6 +1426,9 @@ def __init__( ) -> None: super().__init__(kv_cache_spec, block_pool, **kwargs) self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode + assert self.mamba_cache_mode == "align" or not self.enable_caching, ( + "Mamba prefix caching requires mamba_cache_mode='align'" + ) self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks self.has_prefill_checkpoint_blocks = ( self.mamba_cache_mode == "align" From 3ad096faf81b4734ae995aa66f4ca0c904cea978 Mon Sep 17 00:00:00 2001 From: zhouyou Date: Tue, 13 Oct 2026 15:59:48 +0800 Subject: [PATCH 366/624] [Bugfix][EC Connector] Don't create hash directories on cache-miss lookups in ECExampleConnector (#60826) Signed-off-by: zhouyou9505 --- .../ec_connector/unit/test_ec_example_connector.py | 12 ++++++++++++ .../ec_transfer/ec_connector/example_connector.py | 14 ++++++++------ 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/tests/v1/ec_connector/unit/test_ec_example_connector.py b/tests/v1/ec_connector/unit/test_ec_example_connector.py index 725bb99eb055..a940fd4fcf5e 100644 --- a/tests/v1/ec_connector/unit/test_ec_example_connector.py +++ b/tests/v1/ec_connector/unit/test_ec_example_connector.py @@ -246,6 +246,18 @@ def test_has_cache_item_partial_exist( assert result[1] # Second exists assert not result[2] # Third doesn't exist + def test_has_cache_item_miss_creates_no_folder( + self, mock_vllm_config_consumer, temp_storage + ): + """Test a cache miss leaves no empty folder in the shared storage.""" + connector = ECExampleConnector( + vllm_config=mock_vllm_config_consumer, + role=ECConnectorRole.SCHEDULER, + ) + + assert not connector.has_cache_item("missing_hash") + assert os.listdir(temp_storage) == [] + class TestStateManagement: """Test connector state management.""" diff --git a/vllm/distributed/ec_transfer/ec_connector/example_connector.py b/vllm/distributed/ec_transfer/ec_connector/example_connector.py index e080f09b6bef..97075cdf38c9 100644 --- a/vllm/distributed/ec_transfer/ec_connector/example_connector.py +++ b/vllm/distributed/ec_transfer/ec_connector/example_connector.py @@ -202,7 +202,7 @@ def request_finished( def _found_match_for_mm_data(self, mm_hash) -> bool: """Check if the cache is hit for the request.""" - filename = self._generate_filename_debug(mm_hash) + filename = self._generate_filename_debug(mm_hash, create_folder=False) return os.path.exists(filename) def _generate_foldername_debug( @@ -219,11 +219,13 @@ def _generate_foldername_debug( os.makedirs(foldername, exist_ok=True) return foldername - def _generate_filename_debug(self, mm_hash: str) -> str: + def _generate_filename_debug( + self, + mm_hash: str, + create_folder: bool = True, + ) -> str: """Return the full path of the safetensors file for this mm_hash. - Ensures the parent directory exists because - `_generate_foldername_debug` is called with its default - (`create_folder=True`). + If `create_folder` is True (default) the parent directory is created. """ - foldername = self._generate_foldername_debug(mm_hash) # <- folder auto-created + foldername = self._generate_foldername_debug(mm_hash, create_folder) return os.path.join(foldername, "encoder_cache.safetensors") From ff305934f1c6f98edf7c499464a3c2be5f6e3f04 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Tue, 13 Oct 2026 16:15:47 +0800 Subject: [PATCH 367/624] [Tests][Mamba] Wait for accepted counts in fake MTP proposer (#61342) Signed-off-by: zjy0516 --- tests/v1/e2e/general/test_mamba_prefix_cache.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index 6c208e19fb5d..cf8aee141bc5 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -129,6 +129,10 @@ def fake_propose_draft_token_ids_fn( first_token_id_index = ( num_computed_tokens + 1 ) # bonus token isn't considered as computed + # num_accepted_tokens_cpu is filled by a non-blocking D2H copy in the + # same step; wait for it or a lagging GPU yields stale counts. + if self.num_accepted_tokens_event is not None: + self.num_accepted_tokens_event.synchronize() first_token_id_index += self.input_batch.num_accepted_tokens_cpu[0].item() proposed_draft_token_ids = [ prompt_token_ids[ From 24aa1191f433638722178fed646f0fd52e96a587 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Tue, 13 Oct 2026 10:35:57 +0200 Subject: [PATCH 368/624] =?UTF-8?q?[ROCm][Bugfix]=20Run=20the=20AITER=20AR?= =?UTF-8?q?+RMSNorm=20fallback=20test=20on=20the=20platform=E2=80=A6=20(#6?= =?UTF-8?q?1350)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Stefan Koncarevic Signed-off-by: Stefan Koncarevic --- tests/distributed/test_aiter_rms_fallback.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/distributed/test_aiter_rms_fallback.py b/tests/distributed/test_aiter_rms_fallback.py index 4a0fc67873a1..211ff15e6e0d 100644 --- a/tests/distributed/test_aiter_rms_fallback.py +++ b/tests/distributed/test_aiter_rms_fallback.py @@ -11,6 +11,7 @@ import vllm._aiter_ops as aiter_ops import vllm.distributed as vllm_distributed import vllm.kernels # noqa: F401 +from vllm.platforms import current_platform def test_rocm_aiter_gemma_rmsnorm_falls_back_after_allreduce( @@ -72,9 +73,14 @@ def fake_rmsnorm2d_fwd_with_add( fake_aiter.rmsnorm2d_fwd_with_add = fake_rmsnorm2d_fwd_with_add monkeypatch.setitem(sys.modules, "aiter", fake_aiter) - inp = torch.randn(8192, 16, dtype=torch.bfloat16) + # On the platform's device, not the CPU. The fallback leaves Python through + # `torch.ops.vllm_aiter.fused_add_rms_norm`, which is registered for the + # platform's dispatch key alone, so CPU tensors have nothing to dispatch to + # on a GPU build. The kernel behind it is still the fake above. + device = current_platform.device_type + inp = torch.randn(8192, 16, dtype=torch.bfloat16, device=device) residual = torch.randn_like(inp) - weight = torch.randn(16, dtype=torch.bfloat16) + weight = torch.randn(16, dtype=torch.bfloat16, device=device) out, residual_out = aiter_ops._rocm_aiter_fused_allreduce_rmsnorm_impl( inp, residual, weight, 1e-6, gemma_norm=True From 6f11c220b5b0f3b45b2a14c49ebbca2df12d1c80 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:41:23 +0100 Subject: [PATCH 369/624] [Transformers backend] Fuse MoE routers that score after the top-k (#57410) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) --- tests/models/transformers/fusers/test_moe.py | 70 ++++++++++++++++++ .../models/transformers/fusers/moe.py | 73 +++++++++++++++---- .../models/transformers/fx_utils.py | 30 +++++++- .../model_executor/models/transformers/moe.py | 31 ++++---- vllm/transformers_utils/config.py | 36 +++++++++ .../model_arch_config_convertor.py | 33 ++------- 6 files changed, 213 insertions(+), 60 deletions(-) diff --git a/tests/models/transformers/fusers/test_moe.py b/tests/models/transformers/fusers/test_moe.py index 7de4589eef91..f0e70c896baa 100644 --- a/tests/models/transformers/fusers/test_moe.py +++ b/tests/models/transformers/fusers/test_moe.py @@ -29,6 +29,42 @@ def forward(self, hidden_states): return logits, value, index +class SelectionScoredRouter(TopKRouter): + """Scores the selected logits only (GraniteMoE/Aria style), in fp32.""" + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight).float() + value, index = torch.topk(logits, self.top_k, dim=-1) + value = ( + torch.sigmoid(value) if self.sigmoid else F.softmax(value, dim=-1) + ).type_as(hidden_states) + return logits, value, index + + +class EpsilonRenormRouter(TopKRouter): + """Renormalizes via an epsilon-offset denominator (DeepSeek-V3 style).""" + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight) + scores = torch.sigmoid(logits) + value, index = torch.topk(scores, self.top_k, dim=-1) + denominator = value.sum(dim=-1, keepdim=True) + 1e-20 + value /= denominator + return logits, value, index + + +class UntraceableRouter(TopKRouter): + """Scores and selects, then loses the trace before the renormalization.""" + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight) + scores = F.softmax(logits, dim=-1) + value, index = torch.topk(scores, self.top_k, dim=-1) + if value.sum() > 0: # data-dependent branch: tracing stops here + value = value / value.sum(dim=-1, keepdim=True) + return logits, value, index + + class ScaledRouter(TopKRouter): """Greedy router scaling its top-k weights (DeepSeek `routed_scaling_factor`).""" @@ -268,6 +304,19 @@ def test_moe_fuser_detects_router(sigmoid): assert fuser.shared_name is None and fuser.shared_gate_name is None +@pytest.mark.parametrize("sigmoid", [False, True]) +def test_moe_fuser_matches_router_scoring_the_selection(sigmoid): + """Scoring the selection rather than all logits still allows the fuser to match.""" + with torch.device("meta"): + block = MoEBlock(lambda: SelectionScoredRouter(sigmoid=sigmoid)) + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + assert fuser.gate_name == "gate" + assert fuser.scoring_func == ("sigmoid" if sigmoid else "softmax") + assert fuser.renormalize is (not sigmoid) + assert fuser.router_dtype == torch.float32 + + def test_moe_fuser_matches_scaled_router(): """Weight scaling after the top-k (DeepSeek style) does not break matching.""" with torch.device("meta"): @@ -275,6 +324,27 @@ def test_moe_fuser_matches_scaled_router(): assert isinstance(MoEBlockFuser.match(block, "experts"), MoEBlockFuser) +@pytest.mark.parametrize( + "router_cls,expected", + [ + (TopKRouter, True), # divides by the sum of the selection + (EpsilonRenormRouter, True), # ... with an epsilon in the denominator + (ScaledRouter, False), # scales the weights instead of renormalizing + (UntraceableRouter, None), # trace stops before it could say + ], +) +def test_moe_fuser_reads_renormalize_from_the_gate(router_cls, expected): + """The router's graph, not the config, decides whether vLLM renormalizes: it + is the code being replaced, and it already bakes in `norm_topk_prob` because + the flag is a Python bool that fx resolves at trace time. Absence of the + division only means "no" on a complete trace.""" + with torch.device("meta"): + block = MoEBlock(router_cls) + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + assert fuser.renormalize is expected + + def test_moe_fuser_matches_grouped_router(): """Group masking between the score and the routing top-k still matches: the scorer is anchored on the routing (last) top-k, not the group one.""" diff --git a/vllm/model_executor/models/transformers/fusers/moe.py b/vllm/model_executor/models/transformers/fusers/moe.py index af3a101a0ce1..06fdac109ca4 100644 --- a/vllm/model_executor/models/transformers/fusers/moe.py +++ b/vllm/model_executor/models/transformers/fusers/moe.py @@ -17,6 +17,7 @@ from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.models.transformers.fx_utils import ( find_node, + is_complete, is_op, peel, trace, @@ -107,6 +108,24 @@ def _reaches(node: fx.Node, key: str) -> set[fx.Node]: return seen +def _renormalizes(topk: fx.Node) -> bool: + """Are the top-k weights divided by their own sum? + + Every HF router spells renormalization as a division by a sum over the + selection, optionally offset by an epsilon.""" + downstream = _reaches(topk, "users") + for node in downstream: + if not (is_op(node, "div") or is_op(node, "truediv")) or len(node.args) < 2: + continue + denominator = node.args[1] + if not isinstance(denominator, fx.Node): + continue + cone = _reaches(denominator, "all_input_nodes") + if any(is_op(n, "sum") and n in downstream for n in cone): + return True + return False + + class SharedExpertMLP(nn.Module): """Wraps an HF shared expert, applying the output gating it is paired with.""" @@ -147,12 +166,18 @@ class MoEBlockFuser: shared_name: str | None shared_gate_name: str | None router_dtype: torch.dtype | None = None + renormalize: bool | None = None @staticmethod - def _match_router(gate: nn.Module) -> tuple[str, torch.dtype | None] | None: - """Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`. - - Returns the scoring function and the dtype the router computes in.""" + def _match_router( + gate: nn.Module, + ) -> tuple[str, torch.dtype | None, bool | None] | None: + """Matches `topk(score(linear(x)))` or `score(topk(linear(x)))`, `score` being + `softmax`/`sigmoid`. + + Returns the scoring function, the dtype the router computes in, and + whether the top-k weights are renormalized (`None` if the trace was cut + short before it could say).""" state = {name for name, _ in named_state(gate)} if "weight" not in state or state - {"weight", "e_score_correction_bias"}: return None @@ -164,20 +189,31 @@ def _match_router(gate: nn.Module) -> tuple[str, torch.dtype | None] | None: if not topks: return None topk = topks[-1] - # Exactly one scoring op upstream of the top-k, fed (transitively) by a linear. - scorers = [ - n - for n in _reaches(topk, "all_input_nodes") - if is_op(n, "softmax") or is_op(n, "sigmoid") - ] + is_scorer = lambda n: is_op(n, "softmax") or is_op(n, "sigmoid") + post_topk = False + # Find all scoring functions that reach the top-k node. + scorers = [n for n in _reaches(topk, "all_input_nodes") if is_scorer(n)] + if not scorers: + # Find all scoring functions that are users of the top-k node. + scorers = [n for n in _reaches(topk, "users") if is_scorer(n)] + post_topk = True if len(scorers) != 1: return None scorer = scorers[0] logits_cone = _reaches(scorer, "all_input_nodes") if not any(is_op(n, "linear") for n in logits_cone): return None - scoring_func = "softmax" if is_op(scorer, "softmax") else "sigmoid" - return scoring_func, _forced_dtype(logits_cone) + is_softmax = is_op(scorer, "softmax") + # vLLM always scores every expert, so matching a router that softmaxes + # the selection alone takes renormalizing. Otherwise look for the + # division, reading its absence as "no" only from a complete trace. + renormalize: bool | None + if (post_topk and is_softmax) or _renormalizes(topk): + renormalize = True + else: + renormalize = False if is_complete(graph) else None + scoring_func = "softmax" if is_softmax else "sigmoid" + return scoring_func, _forced_dtype(logits_cone), renormalize @staticmethod def _match_shared_experts( @@ -221,14 +257,14 @@ def match(cls, moe_block: nn.Module, experts_name: str) -> "MoEBlockFuser | None if _returns_tuple(type(moe_block)): return None # Router: the child that scores + top-k selects. - gate_name = scoring_func = router_dtype = None + gate_name = scoring_func = router_dtype = renormalize = None for name, child in moe_block.named_children(): if ( name != experts_name and (router := cls._match_router(child)) is not None ): gate_name = name - scoring_func, router_dtype = router + scoring_func, router_dtype, renormalize = router break if gate_name is None or scoring_func is None: return None @@ -256,7 +292,14 @@ def match(cls, moe_block: nn.Module, experts_name: str) -> "MoEBlockFuser | None for name, child in moe_block.named_children(): if name not in accounted and next(named_state(child), None) is not None: return None - return cls(gate_name, scoring_func, shared_name, shared_gate_name, router_dtype) + return cls( + gate_name, + scoring_func, + shared_name, + shared_gate_name, + router_dtype, + renormalize, + ) def gate( self, moe_block: nn.Module, prefix: str, out_dtype: torch.dtype | None = None diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py index cb630866a390..818e1ca7bb4f 100644 --- a/vllm/model_executor/models/transformers/fx_utils.py +++ b/vllm/model_executor/models/transformers/fx_utils.py @@ -35,6 +35,21 @@ """The unpatched `nn.Module.__call__`. During tracing fx patches it to record `call_module` nodes; meta execution must call modules for real.""" +_MODULE_GETATTR = nn.Module.__getattr__ +"""The unpatched `nn.Module.__getattr__`. During tracing fx patches it to record +`get_attr` nodes; reading one while inferring a `get_attr`'s own value would +recurse until the stack runs out, truncating the graph.""" + + +@contextlib.contextmanager +def _unpatched_modules(): + """Restores the module hooks fx patches, so meta execution records nothing.""" + with ( + mock.patch.object(nn.Module, "__call__", _MODULE_CALL), + mock.patch.object(nn.Module, "__getattr__", _MODULE_GETATTR), + ): + yield + def is_leaf_call(node: object) -> TypeIs[fx.Node]: """Is node a call recorded by `_as_leaf_call` (e.g. an attention interface).""" @@ -132,7 +147,8 @@ def _infer_meta(self, kind: str, target: object, args: tuple, kwargs: dict): ) return _UNKNOWN if kind == "get_attr": - value = operator.attrgetter(str(target))(self.root) + with _unpatched_modules(): + value = operator.attrgetter(str(target))(self.root) if isinstance(value, torch.Tensor): value = torch.empty_like(value, device="meta") return value @@ -159,8 +175,6 @@ def meta_of(arg: object) -> object: if kind == "call_module": # Run the child's forward with all its state on "meta", without # mutating it (at match time params may be meta but buffers real). - # fx patches `nn.Module.__call__` while tracing; restore the real - # one so this execution is not itself recorded. child = self.root.get_submodule(str(target)) state = { name: torch.empty_like(tensor, device="meta") @@ -168,7 +182,7 @@ def meta_of(arg: object) -> object: child.named_parameters(), child.named_buffers() ) } - with mock.patch.object(nn.Module, "__call__", _MODULE_CALL): + with _unpatched_modules(): return torch.func.functional_call(child, state, meta_args, meta_kwargs) return _UNKNOWN @@ -686,6 +700,14 @@ def visit(self, node: ast.AST) -> ast.AST: _Replacer().visit(module) +def is_complete(graph: fx.Graph) -> bool: + """Did the trace run to completion? A partial graph has no `output` node. + + Only a complete graph can be read for the *absence* of an op; in a partial + one the tail is exactly what is missing.""" + return any(node.op == "output" for node in graph.nodes) + + def find_node(graph: fx.Graph, predicate: Callable[[fx.Node], bool]) -> fx.Node | None: """The first node in `graph` matching `predicate`, or `None`.""" return next((n for n in graph.nodes if predicate(n)), None) diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index 71f5bd066c8d..953459172d36 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -39,6 +39,7 @@ from vllm.model_executor.models.transformers.fusers.glu import GLUFuser from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser from vllm.model_executor.models.utils import extract_layer_index, maybe_prefix +from vllm.transformers_utils.config import CONFIG_FIELD_ALIASES from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, direct_register_custom_op from .base import Base @@ -180,23 +181,18 @@ def recursive_replace(self): text_config = self.text_config # Positional arguments - num_experts = self.model_config.get_num_experts() - top_k = getattr_iter(text_config, ["num_experts_per_tok", "top_k"], None) + num_experts_names = CONFIG_FIELD_ALIASES["num_experts"] + num_experts = getattr_iter(text_config, num_experts_names, None) + top_k_names = CONFIG_FIELD_ALIASES["num_experts_per_token"] + top_k = getattr_iter(text_config, top_k_names, None) assert top_k is not None hidden_size = text_config.hidden_size - intermediate_size = getattr_iter( - text_config, ["moe_intermediate_size", "intermediate_size"], None - ) + intermediate_size_names = CONFIG_FIELD_ALIASES["moe_intermediate_size"] + intermediate_size = getattr_iter(text_config, intermediate_size_names, None) assert intermediate_size is not None - num_shared_experts = getattr_iter( - text_config, - [ - "n_shared_experts", # DeepSeek, Docs, GLM - "moe_num_shared_experts", # Aria, Ernie - ], - 0, - ) + num_shared_experts_names = CONFIG_FIELD_ALIASES["num_shared_experts"] + num_shared_experts = getattr_iter(text_config, num_shared_experts_names, 0) # Common kwargs norm_topk_prob = getattr(text_config, "norm_topk_prob", None) @@ -275,6 +271,9 @@ def _recursive_replace(module: nn.Module, prefix: str): if "bias" in experts_param_name: has_bias = True break + # Whether the expert weights are stored in transposed format. + # Comes from Transformers `use_experts_implementation` decorator. + is_transposed = getattr(experts, "is_transposed", False) # If the config does not specify num_shared_experts, but # the model has shared experts, we assume there is one. if self.num_shared_experts == 0: @@ -308,6 +307,7 @@ def _recursive_replace(module: nn.Module, prefix: str): enable_eplb=enable_eplb, num_redundant_experts=num_redundant_experts, has_bias=has_bias, + is_fused_checkpoint_transposed=is_transposed, routed_experts_cls=TransformersRoutedExperts, ) fuser = MoEBlockFuser.match(moe_block, experts_name) @@ -346,6 +346,11 @@ def _recursive_replace(module: nn.Module, prefix: str): gate = fuser.gate(moe_block, prefix, router_dtype) kwargs |= dict( scoring_func=fuser.scoring_func, + renormalize=( + layer_renormalize + if fuser.renormalize is None + else fuser.renormalize + ), is_sequence_parallel=( self.parallel_config.use_sequence_parallel_moe ), diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index e046c301d23a..1d93a04e662a 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -52,6 +52,42 @@ logger = init_logger(__name__) +CONFIG_FIELD_ALIASES: dict[str, tuple[str, ...]] = { + "total_num_kv_heads": ( + "n_head_kv", # Falcon + "num_kv_heads", + "num_key_value_heads", # LLaMA-2 + "multi_query_group_num", # ChatGLM + "num_attention_groups", # Step3p5 + ), + "num_experts": ( + "num_experts", # Jamba + "moe_num_experts", # Dbrx + "n_routed_experts", # DeepSeek + "num_local_experts", # Mixtral + ), + "num_experts_per_token": ( + "num_experts_per_tok", + "num_experts_per_token", + "top_k_experts", + "moe_topk", # Hunyuan + "moe_top_k", + "top_k", + ), + "moe_intermediate_size": ( + "moe_intermediate_size", + "intermediate_size", + ), + "num_shared_experts": ( + "n_shared_experts", # DeepSeek, Docs, GLM + "moe_num_shared_experts", # Aria, Ernie + ), +} +"""HF config fields that models spell differently, keyed by the name vLLM uses. + +`ModelArchConfigConvertorBase` reads a scalar from these; a caller that needs the +raw value (a per-layer `list`, say) passes the aliases to `getattr_iter` itself.""" + _ST_POOLING_MODULE_TYPES = { "sentence_transformers.models.Pooling", "sentence_transformers.sentence_transformer.modules.pooling.Pooling", diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 06a937697106..6290f6c6645b 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -14,6 +14,7 @@ from vllm.config.utils import getattr_iter from vllm.logger import init_logger from vllm.transformers_utils.config import ( + CONFIG_FIELD_ALIASES, ConfigFormat, get_safetensors_params_metadata, ) @@ -148,23 +149,11 @@ def _get_qk_rope_head_dim(self) -> int: return qk_rope_head_dim def get_total_num_kv_heads(self) -> int: - attributes = [ - # For Falcon: - "n_head_kv", - "num_kv_heads", - # For LLaMA-2: - "num_key_value_heads", - # For ChatGLM: - "multi_query_group_num", - # For Step3p5: - "num_attention_groups", - ] + names = CONFIG_FIELD_ALIASES["total_num_kv_heads"] # For non-grouped-query attention models, the number of KV heads is # equal to the number of attention heads. default_factory = self.get_total_num_attention_heads - return getattr_iter( - self.hf_text_config, attributes, default_factory=default_factory - ) + return getattr_iter(self.hf_text_config, names, default_factory=default_factory) def get_num_experts_from_block_configs(self) -> int: """Check block_configs for heterogeneous models (e.g., NemotronH). @@ -188,13 +177,7 @@ def get_num_experts_from_block_configs(self) -> int: def get_num_experts(self) -> int: """Returns the number of experts in the model.""" - num_expert_names = [ - "num_experts", # Jamba - "moe_num_experts", # Dbrx - "n_routed_experts", # DeepSeek - "num_local_experts", # Mixtral - ] - + num_expert_names = CONFIG_FIELD_ALIASES["num_experts"] num_experts = getattr_iter(self.hf_text_config, num_expert_names, 0) if isinstance(num_experts, list): # Ernie VL's remote code uses list[int]... @@ -206,13 +189,7 @@ def get_num_experts(self) -> int: return num_experts def get_num_experts_per_token(self) -> int: - names = [ - "num_experts_per_tok", - "num_experts_per_token", - "top_k_experts", - "moe_topk", - "moe_top_k", - ] + names = CONFIG_FIELD_ALIASES["num_experts_per_token"] num_experts_per_token = getattr_iter(self.hf_text_config, names, 0) if isinstance(num_experts_per_token, list): return max(num_experts_per_token, default=0) From 183625d4f9f49968218ffbe43242877ead72d569 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Tue, 13 Oct 2026 12:48:13 +0400 Subject: [PATCH 370/624] [Bugfix] Fall back to auto for unquantized MoE layers when --moe-backend is marlin (#60219) Signed-off-by: Vadim Gimpelson Co-authored-by: Claude --- .../moe/test_unquantized_backend_selection.py | 14 ++++++++++++++ .../layers/fused_moe/oracle/unquantized.py | 7 ++++--- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/tests/kernels/moe/test_unquantized_backend_selection.py b/tests/kernels/moe/test_unquantized_backend_selection.py index b64e909fbd45..e3a83ec3c5c7 100644 --- a/tests/kernels/moe/test_unquantized_backend_selection.py +++ b/tests/kernels/moe/test_unquantized_backend_selection.py @@ -626,3 +626,17 @@ def test_select_explicit_triton_backend(is_lora_enabled): assert selected_backend == UnquantizedMoeBackend.TRITON assert experts_cls is not None + + +@skipif_not_cuda_rocm +@pytest.mark.parametrize("moe_backend", ["humming", "marlin"]) +def test_quantization_only_backend_falls_back_to_auto(moe_backend): + """A quantization-only --moe-backend also reaches the unquantized MoE + layers of a quantized model (e.g. a BF16 MTP draft), which must use the + auto-selected backend instead of failing.""" + moe_config = make_dummy_moe_config() + moe_config.moe_backend = moe_backend + + assert select_unquantized_moe_backend( + moe_config=moe_config + ) == select_unquantized_moe_backend(moe_config=make_dummy_moe_config()) diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 1af0450dfd13..2f05074c643c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -303,9 +303,10 @@ def _return_or_raise( ) runner_backend = moe_config.moe_backend - # 'humming' is quantization-only; an unquantized layer (e.g. excluded via - # modules_to_not_convert) falls through to auto instead of erroring. - if runner_backend not in ["auto", "humming"]: + # 'humming' and 'marlin' are quantization-only; an unquantized layer (e.g. + # excluded via modules_to_not_convert, or an unquantized MTP draft of a + # quantized model) falls through to auto instead of erroring. + if runner_backend not in ["auto", "humming", "marlin"]: requested_backend = map_unquantized_backend(runner_backend) if ( activation_format == mk.FusedMoEActivationFormat.BatchedExperts From 656d0d4506052a23d100c56e148c7b984828d637 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Tue, 13 Oct 2026 17:01:12 +0800 Subject: [PATCH 371/624] [Perf] Tune Triton GDN decode tiles on SM100-family GPUs (#61355) Signed-off-by: Thien Tran Co-authored-by: OpenAI Codex --- .../fla/test_fused_recurrent_packed_decode.py | 56 ++++++++++++++----- .../ops/fused_recurrent.py | 12 ++++ 2 files changed, 53 insertions(+), 15 deletions(-) diff --git a/tests/kernels/fla/test_fused_recurrent_packed_decode.py b/tests/kernels/fla/test_fused_recurrent_packed_decode.py index be6597cc865f..5ea66ced2390 100644 --- a/tests/kernels/fla/test_fused_recurrent_packed_decode.py +++ b/tests/kernels/fla/test_fused_recurrent_packed_decode.py @@ -18,17 +18,32 @@ ) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) @pytest.mark.parametrize("strided_mixed_qkv", [False, True]) +@pytest.mark.parametrize( + "batch,num_key_heads,num_value_heads,dtype,state_dtype", + [ + (32, 4, 8, torch.float16, torch.float16), + (32, 4, 8, torch.bfloat16, torch.bfloat16), + (32, 4, 8, torch.float32, torch.float32), + (1, 16, 64, torch.bfloat16, torch.float32), + (9, 4, 12, torch.bfloat16, torch.float32), + (1, 4, 12, torch.bfloat16, torch.bfloat16), + ], +) def test_fused_recurrent_packed_decode_matches_reference( - dtype: torch.dtype, strided_mixed_qkv: bool + dtype: torch.dtype, + strided_mixed_qkv: bool, + batch: int, + num_key_heads: int, + num_value_heads: int, + state_dtype: torch.dtype, ): - torch.manual_seed(0) + generator = torch.Generator(device=DEVICE).manual_seed(0) # Small but representative GDN config (Qwen3Next defaults are K=128, V=128). - B = 32 - H = 4 - HV = 8 # grouped value attention: HV must be divisible by H + B = batch + H = num_key_heads + HV = num_value_heads K = 128 V = 128 qkv_dim = 2 * (H * K) + (HV * V) @@ -38,22 +53,33 @@ def test_fused_recurrent_packed_decode_matches_reference( if strided_mixed_qkv: # Simulate a packed view into a larger projection buffer: # mixed_qkv.stride(0) > mixed_qkv.shape[1] - proj = torch.randn((B, qkv_dim + 64), device=device, dtype=dtype) + proj = torch.randn( + (B, qkv_dim + 64), device=device, dtype=dtype, generator=generator + ) mixed_qkv = proj[:, :qkv_dim] else: - mixed_qkv = torch.randn((B, qkv_dim), device=device, dtype=dtype) + mixed_qkv = torch.randn( + (B, qkv_dim), device=device, dtype=dtype, generator=generator + ) - a = torch.randn((B, HV), device=device, dtype=dtype) - b = torch.randn((B, HV), device=device, dtype=dtype) - A_log = torch.randn((HV,), device=device, dtype=dtype) - dt_bias = torch.randn((HV,), device=device, dtype=dtype) + a = torch.randn((B, HV), device=device, dtype=dtype, generator=generator) + b = torch.randn((B, HV), device=device, dtype=dtype, generator=generator) + A_log = torch.randn((HV,), device=device, dtype=dtype, generator=generator) + dt_bias = torch.randn((HV,), device=device, dtype=dtype, generator=generator) # Continuous batching indices (include PAD_SLOT_ID=-1 cases). Index 0 is # reserved as NULL_BLOCK_ID (CUDA graph padding), so valid slots start at 1. ssm_state_indices = torch.arange(1, B + 1, device=device, dtype=torch.int32) - ssm_state_indices[-3:] = -1 - - state0 = torch.randn((B + 1, HV, V, K), device=device, dtype=dtype) + if B > 3: + ssm_state_indices[-3:] = -1 + ssm_state_indices[-1] = 0 + + state0 = torch.randn( + (B + 1, HV, V, K), + device=device, + dtype=state_dtype, + generator=generator, + ) state_ref = state0.clone() state_packed = state0.clone() diff --git a/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py index eb08b938c2fb..2f3997e995fe 100644 --- a/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py +++ b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py @@ -10,6 +10,7 @@ import torch +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from .op import exp, log @@ -439,6 +440,17 @@ def fused_recurrent_gated_delta_rule_packed_decode( f"Packed decode kernel only supports NK=1 (got K={K}, BK={BK})." ) BV = min(triton.next_power_of_2(V), 32) + if ( + K == V == 128 + and current_platform.is_cuda() + and current_platform.is_device_capability_family(100, device_id=dev.index) + ): + # Tuned on GB300; validated on GB200 and with 148 SMs. + # MRv2 decode warmup/graph capture compiles the reachable BV variants. + if B * HV <= 64: + BV = 4 + elif initial_state.dtype == torch.float32: + BV = 8 num_stages = 3 num_warps = 1 From 206e45d47009447001e9080d430e64c90fba21a6 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 13 Oct 2026 02:24:06 -0700 Subject: [PATCH 372/624] [CI] Shard V1 Attention automatically instead of by fixed pytest shards (#61200) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/attention.yaml | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/.buildkite/test_areas/attention.yaml b/.buildkite/test_areas/attention.yaml index ec47e45fe7a0..5952aa40564c 100644 --- a/.buildkite/test_areas/attention.yaml +++ b/.buildkite/test_areas/attention.yaml @@ -2,8 +2,9 @@ group: Attention depends_on: - image-build steps: -- label: ":nvidia: (H100) V1 Attention Shard %N" +- label: ":nvidia: (H100) V1 Attention" key: v1-attention-h100-mi300 + automatic_shard: true timeout_in_minutes: 85 device: h100 source_file_dependencies: @@ -12,14 +13,18 @@ steps: - vllm/v1/attention - tests/v1/attention commands: - - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - parallelism: 2 + - pytest -v -s v1/attention mirror: amd: label: ":amd: (MI300) V1 Attention Shard %N" dind: false device: mi300_1 timeout_in_minutes: 125 + # automatic_shard covers NVIDIA steps only; the mirror keeps its + # fixed pytest shards. + parallelism: 2 + commands: + - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT depends_on: - image-build-amd source_file_dependencies: @@ -35,8 +40,9 @@ steps: - pytest -v -s v1/attention --ignore=v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: ":nvidia: (B200) V1 Attention Shard %N" +- label: ":nvidia: (B200) V1 Attention" key: v1-attention-b200 + automatic_shard: true timeout_in_minutes: 80 device: b200-k8s source_file_dependencies: @@ -45,14 +51,16 @@ steps: - vllm/v1/attention - tests/v1/attention commands: - - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - parallelism: 2 + - pytest -v -s v1/attention mirror: amd: label: ":amd: (MI355 DPX) V1 Attention Shard %N" dind: false device: mi355_dpx timeout_in_minutes: 125 + # automatic_shard covers NVIDIA steps only; the mirror keeps its + # fixed pytest shards. + parallelism: 2 depends_on: - image-build-amd source_file_dependencies: From c7a305ec10f3ec06c2784f1258162483c5fc8d73 Mon Sep 17 00:00:00 2001 From: sungsoo ha Date: Tue, 13 Oct 2026 03:09:17 -0700 Subject: [PATCH 373/624] [Bugfix][Mamba] Store conv_kernel - 1 inputs in Mamba2 prefill checkpoints (#60797) Signed-off-by: Sungsoo Ha Co-authored-by: Claude Opus 5.5 Co-authored-by: tomeras91 <57313761+tomeras91@users.noreply.github.com> --- .../mamba/test_mamba2_checkpoint_state.py | 16 +++++++++++----- vllm/model_executor/layers/mamba/mamba_mixer2.py | 1 + .../layers/mamba/ops/ssd_checkpoint.py | 7 +++++-- 3 files changed, 17 insertions(+), 7 deletions(-) diff --git a/tests/kernels/mamba/test_mamba2_checkpoint_state.py b/tests/kernels/mamba/test_mamba2_checkpoint_state.py index 54fcf80205ce..a06d894792f8 100644 --- a/tests/kernels/mamba/test_mamba2_checkpoint_state.py +++ b/tests/kernels/mamba/test_mamba2_checkpoint_state.py @@ -114,23 +114,27 @@ def test_checkpoint_does_not_perturb_the_final_state(): ) +@pytest.mark.parametrize("spec_slots", [0, 3]) @pytest.mark.parametrize("dim_first", [True, False]) -def test_store_selects_the_checkpoint_from_all_chunk_states(dim_first): +def test_store_selects_the_checkpoint_from_all_chunk_states(dim_first, spec_slots): """The store reads the checkpoint row out of every chunk state. Row 0 checkpoints 32 tokens in, its state being chunk 3 rather than the program's own row, so the store only passes if ``chunk_idx`` is used. Row 1 declines and must write nothing. The conv input is a column slice of a wider projection, the SD conv layout is a transposed view, and SSM rows - are padded pages, matching the strides the mixer passes. + are padded pages, matching the strides the mixer passes. Spec decoding + widens the conv state, and the store must still fill only the first + ``state_len`` slots with the newest inputs. """ dim, state_len, heads, head_dim, d_state = 16, 4, 2, 4, 8 projection = torch.randn(64, dim + 8, device=DEVICE) conv_input = projection[:, :dim] + num_slots = state_len + spec_slots if dim_first: - conv_state = torch.zeros(3, dim, state_len, device=DEVICE) + conv_state = torch.zeros(3, dim, num_slots, device=DEVICE) else: - conv_state = torch.zeros(3, state_len, dim, device=DEVICE).transpose(-1, -2) + conv_state = torch.zeros(3, num_slots, dim, device=DEVICE).transpose(-1, -2) varlen_states = torch.randn(5, heads, head_dim, d_state, device=DEVICE) # Blocks are pages holding more than the SSM state, so rows are padded. row = heads * head_dim * d_state @@ -152,12 +156,14 @@ def test_store_selects_the_checkpoint_from_all_chunk_states(dim_first): varlen_states, ssm_state, torch.tensor([0, 40, 64], dtype=torch.int32, device=DEVICE), + state_len=state_len, ) exact = dict(rtol=0, atol=0) torch.testing.assert_close( - conv_state[2], conv_input[32 - state_len : 32].T, **exact + conv_state[2, :, :state_len], conv_input[32 - state_len : 32].T, **exact ) + assert not conv_state[2, :, state_len:].any(), "wrote into the spec slots" torch.testing.assert_close(ssm_state[2], varlen_states[3], **exact) assert not conv_state[:2].any() and not ssm_storage[:2].any() assert not ssm_storage[:, row:].any(), "wrote into the page padding" diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 14c483122ab1..e310aaf40f12 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -890,6 +890,7 @@ def conv_ssm_forward( varlen_states, ssm_state, query_start_loc_p, + state_len=self.conv_kernel_size - 1, ) ssm_state[state_indices_tensor_p] = final_states if ring_start is not None and self._updates_replayssm_trackers: diff --git a/vllm/model_executor/layers/mamba/ops/ssd_checkpoint.py b/vllm/model_executor/layers/mamba/ops/ssd_checkpoint.py index 25bfcedbdc88..fdd8ac8f0ed5 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_checkpoint.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_checkpoint.py @@ -69,6 +69,8 @@ def store_prefill_checkpoint( chunk_states: torch.Tensor, ssm_state: torch.Tensor, query_start_loc: torch.Tensor, + *, + state_len: int, ) -> None: """Write each prefill row's mid-prefill checkpoint into its paged block. @@ -78,14 +80,15 @@ def store_prefill_checkpoint( checkpoint: Per-row offsets and destination blocks. chunk_idx: Per-row index into ``chunk_states`` of the checkpoint. x: Pre-conv activations, ``(dim, num_prefill_tokens)``. - conv_state: Paged conv state, ``(num_blocks, dim, state_len)``. + conv_state: Paged conv state, ``(num_blocks, dim, state_len + num_spec)``. chunk_states: SSD states at every chunk end. ssm_state: Paged SSM state. query_start_loc: Prefill query start locations. + state_len: ``conv_kernel_size - 1``, not the spec-widened conv width. + Readers take the first ``state_len`` slots as the newest inputs. """ dim = x.size(0) - state_len = conv_state.size(-1) ssm_row_size = chunk_states[0].numel() block_size = 1024 grid = ( From 877f6df889176f056713ac30919153b9dacc0cf1 Mon Sep 17 00:00:00 2001 From: JackRao123 <112158010+JackRao123@users.noreply.github.com> Date: Tue, 13 Oct 2026 03:15:11 -0700 Subject: [PATCH 374/624] [Bugfix][DSv4] Don't bake the C128 compressor skip into CUDA graphs (#59709) Signed-off-by: Jack Rao Co-authored-by: Claude Opus 5.5 --- tests/kernels/test_compressor_kv_cache.py | 78 +++++++++++++++++++++++ vllm/models/deepseek_v4/compressor.py | 6 +- 2 files changed, 81 insertions(+), 3 deletions(-) diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index e9924229055a..237c552fc124 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -1205,6 +1205,84 @@ def test_get_c128_boundary(starts, query_start_loc, expected): assert _get_c128_boundary(metadata) is expected +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@torch.inference_mode() +def test_c128_skip_is_not_baked_into_cuda_graph(monkeypatch): + """A graph captured on a batch with no C128 boundary must still launch the + compressor on replay; only eager steps may skip it.""" + from vllm.config import CUDAGraphMode + from vllm.forward_context import ForwardContext, override_forward_context + from vllm.models.deepseek_v4 import compressor as compressor_mod + from vllm.models.deepseek_v4.compressor import DeepseekCompressor + from vllm.models.deepseek_v4.nvidia.ops import sparse_attn_compress_cutedsl + + launches = torch.zeros(1, dtype=torch.int32, device="cuda") + monkeypatch.setattr( + compressor_mod, "_SAVE_PARTIAL_STATES_KERNEL", lambda **kw: None + ) + monkeypatch.setattr( + sparse_attn_compress_cutedsl, + "_SPARSE_ATTN_COMPRESSOR_CUTEDSL_KERNEL", + lambda **kw: launches.add_(1), + ) + + compressor = DeepseekCompressor.__new__(DeepseekCompressor) + torch.nn.Module.__init__(compressor) + compressor.head_dim, compressor.rope_head_dim = 512, 64 + compressor.compress_ratio, compressor.coff, compressor.overlap = 128, 1, False + compressor.use_fp4_cache = False + compressor._use_two_stage_fused_compressor = False + compressor._quant_block = 64 + compressor._token_stride = 576 + compressor._scale_dim = 8 + compressor.rms_norm_eps = 1e-6 + compressor.ape = torch.zeros(128, 512, device="cuda") + compressor.norm = SimpleNamespace(weight=torch.ones(512, device="cuda")) + compressor.state_cache = SimpleNamespace( + prefix="state", kv_cache=torch.zeros(2, 8, 1024, device="cuda") + ) + compressor.k_cache_prefix = "main" + kv_cache = torch.zeros(1, 64, 584, dtype=torch.uint8, device="cuda") + compressor._static_forward_context = {"main": SimpleNamespace(kv_cache=kv_cache)} + + num_tokens = 4 + slots = torch.arange(num_tokens, device="cuda") + metadata = { + "state": SimpleNamespace( + slot_mapping=slots, + block_table=torch.zeros(num_tokens, 1, dtype=torch.int32, device="cuda"), + block_size=8, + token_to_req_indices=torch.arange( + num_tokens, dtype=torch.int32, device="cuda" + ), + c128_boundary=False, + ), + "main": SimpleNamespace(slot_mapping=slots), + } + kv_score = torch.zeros(num_tokens, 1024, device="cuda") + positions = torch.arange(num_tokens, device="cuda") + rotary = SimpleNamespace(cos_sin_cache=torch.zeros(1, 64, device="cuda")) + context = ForwardContext( + {}, metadata, {}, cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE + ) + + def run(): + with override_forward_context(context): + compressor(kv_score, positions, rotary) + + run() + torch.accelerator.synchronize() + assert launches.item() == 0, "eager step with no boundary should skip" + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + launches.zero_() + graph.replay() + torch.accelerator.synchronize() + assert launches.item() == 1, "captured graph must not bake in the skip" + + # ── Test A: DeepseekV4 Attention path ────────────────────────────────────────────── diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index ccb864faa274..b1e8b5d5bc8b 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -7,7 +7,7 @@ import torch from torch import nn -from vllm.config import CUDAGraphMode, VllmConfig, get_current_vllm_config +from vllm.config import VllmConfig, get_current_vllm_config from vllm.forward_context import get_forward_context from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.layernorm import RMSNorm @@ -409,12 +409,12 @@ def forward( pdl_kwargs=pdl_kwargs, ) - # full graph cannot branch on per-step CPU metadata after capture + # A captured graph would replay the capture batch's skip decision. if ( current_platform.is_cuda() and self.head_dim == 512 and self.compress_ratio == 128 - and forward_context.cudagraph_runtime_mode != CUDAGraphMode.FULL + and not torch.cuda.is_current_stream_capturing() and state_metadata.c128_boundary is False ): return From 7173c64e2ec8763c6d8ddd5c3ff9cb6fc92e2da8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20P=C3=A9rez=20de=20Algaba?= <124347725+jperezdealgaba@users.noreply.github.com> Date: Tue, 13 Oct 2026 12:17:14 +0200 Subject: [PATCH 375/624] [Security] Bound decoded and downloaded media per request (#61226) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: juanpas97 Signed-off-by: Juan Pérez de Algaba Co-authored-by: OpenAI Codex --- docs/usage/security.md | 2 + tests/multimodal/media/test_audio.py | 100 ++++++++++++++++++ tests/multimodal/media/test_image.py | 66 +++++++++++- .../media/test_unprocessable_entity_error.py | 22 +++- tests/multimodal/media/test_video.py | 15 +++ tests/test_audio_media_size_precheck.py | 54 +++++++++- tests/test_connections.py | 51 ++++++++- vllm/connections.py | 65 ++++++++++-- vllm/envs.py | 13 +++ vllm/multimodal/media/__init__.py | 3 +- vllm/multimodal/media/audio.py | 60 ++++++++--- vllm/multimodal/media/base.py | 71 +++++++++++++ vllm/multimodal/media/connector.py | 54 +++++++--- vllm/multimodal/media/image.py | 13 ++- vllm/multimodal/media/video.py | 10 ++ vllm/renderers/chat_utils.py | 16 +++ 16 files changed, 565 insertions(+), 50 deletions(-) diff --git a/docs/usage/security.md b/docs/usage/security.md index 0d7894a467a2..428e64258e0c 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -94,7 +94,9 @@ denial of service: | Environment Variable | Default | Description | | --- | --- | --- | | `VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB` | `256` | Maximum size in MB for a single remote media response. Oversized responses are rejected while streaming before the full body is materialized in memory. | +| `VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB` | `256` | Maximum total HTTP media bytes downloaded for one request, including concurrent downloads. Enforced while streaming. | | `VLLM_MAX_IMAGE_PIXELS` | `178956970` (~179M pixels) | Maximum decoded image size in pixels. Images exceeding this are rejected before raster memory is allocated. Default matches PIL's built-in 2x decompression-bomb threshold (~680 MB for RGB). | +| `VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES` | `715827880` (~683 MiB) | Maximum aggregate decoded image, video-frame, and audio bytes in one request. Images are charged at four bytes per pixel before decoding; audio is limited before decoding; video files are charged after frame extraction. | | `VLLM_MAX_AUDIO_CLIP_FILESIZE_MB` | `25` | Maximum compressed filesize in MB for a single audio file. Enforced on all audio inputs (multimodal chat URLs, speech-to-text uploads, data: URLs, and local file paths) before decoding begins. | | `VLLM_MAX_AUDIO_DECODE_DURATION_S` | `600` | Maximum decoded audio duration in seconds. Prevents compressed audio from expanding into gigabytes of float32 PCM. | | `VLLM_MAX_AUDIO_DECODE_BYTES` | `268435456` (256 MiB) | Maximum float32 PCM bytes that audio decoding may allocate. Guards against sample-rate forgery where an inflated header sample rate bypasses the duration guard while the actual frame count causes a multi-GiB allocation. | diff --git a/tests/multimodal/media/test_audio.py b/tests/multimodal/media/test_audio.py index a75c20979e0c..5728fa9a4660 100644 --- a/tests/multimodal/media/test_audio.py +++ b/tests/multimodal/media/test_audio.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import threading +from concurrent.futures import ThreadPoolExecutor from io import BytesIO from pathlib import Path from unittest.mock import patch @@ -9,6 +11,7 @@ import pytest import soundfile as sf +import vllm.envs as envs from vllm.multimodal.media import AudioMediaIO from vllm.multimodal.media import audio as audio_module from vllm.multimodal.media.audio import ( @@ -16,6 +19,7 @@ load_audio_soundfile, load_audio_torchcodec, ) +from vllm.multimodal.media.base import MediaDecodedBudget from ...conftest import AudioTestAssets @@ -118,6 +122,102 @@ def _make_flac_bytes(frames: int, channels: int, samplerate: int) -> bytes: return buf.getvalue() +@pytest.mark.skip_global_cleanup +def test_request_budget_rejects_second_decoded_audio_clip(monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 384) + clip = _make_flac_bytes(frames=64, channels=1, samplerate=16000) + audio_io = AudioMediaIO(decoded_budget=MediaDecodedBudget()) + + audio_io.load_bytes(clip) + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"): + audio_io.load_bytes(clip) + + +@pytest.mark.skip_global_cleanup +def test_request_budget_allows_parallel_audio_decodes(monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 512) + monkeypatch.setattr(envs, "VLLM_MAX_AUDIO_DECODE_BYTES", 256) + budget = MediaDecodedBudget() + audio_io = AudioMediaIO(decoded_budget=budget) + decoding = threading.Barrier(2) + + def decode(*args, **kwargs): + assert kwargs["max_decode_bytes"] == 256 + decoding.wait(timeout=5) + return np.zeros(32, dtype=np.float32), 16000 + + with ( + patch.object(audio_module, "load_audio", side_effect=decode), + ThreadPoolExecutor(max_workers=2) as pool, + ): + results = list(pool.map(audio_io.load_bytes, [b"a", b"b"])) + + assert len(results) == 2 + assert budget.used == 256 + + +@pytest.mark.skip_global_cleanup +def test_request_budget_waits_for_inflight_audio_decode(monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 256) + monkeypatch.setattr(envs, "VLLM_MAX_AUDIO_DECODE_BYTES", 256) + budget = MediaDecodedBudget() + audio_io = AudioMediaIO(decoded_budget=budget) + decoding = threading.Event() + finish_decoding = threading.Event() + second_reservation = threading.Event() + reserve = budget.reserve + decode_calls = 0 + + def track_reservation(max_bytes): + if decoding.is_set(): + second_reservation.set() + return reserve(max_bytes) + + def decode(*args, **kwargs): + nonlocal decode_calls + decode_calls += 1 + decoding.set() + assert finish_decoding.wait(timeout=5) + return np.zeros(64, dtype=np.float32), 16000 + + with ( + patch.object(budget, "reserve", side_effect=track_reservation), + patch.object(audio_module, "load_audio", side_effect=decode), + ThreadPoolExecutor(max_workers=2) as pool, + ): + first = pool.submit(audio_io.load_bytes, b"a") + assert decoding.wait(timeout=5) + second = pool.submit(audio_io.load_bytes, b"b") + try: + assert second_reservation.wait(timeout=5) + finally: + finish_decoding.set() + first.result(timeout=5) + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"): + second.result(timeout=5) + + assert decode_calls == 1 + assert budget.used == 256 + + +@pytest.mark.skip_global_cleanup +def test_request_budget_releases_failed_audio_reservation(monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 256) + monkeypatch.setattr(envs, "VLLM_MAX_AUDIO_DECODE_BYTES", 256) + budget = MediaDecodedBudget() + audio_io = AudioMediaIO(decoded_budget=budget) + audio = (np.zeros(64, dtype=np.float32), 16000) + + with patch.object( + audio_module, "load_audio", side_effect=[ValueError("invalid audio"), audio] + ): + with pytest.raises(ValueError, match="invalid audio"): + audio_io.load_bytes(b"a") + assert audio_io.load_bytes(b"b")[0].nbytes == 256 + + assert budget.used == 256 + + def test_small_file_passes_memory_guard(): """A small valid file should pass both duration and memory guards.""" payload = _make_flac_bytes(frames=16000, channels=1, samplerate=16000) diff --git a/tests/multimodal/media/test_image.py b/tests/multimodal/media/test_image.py index 94efb312f272..60cfb86cc0d1 100644 --- a/tests/multimodal/media/test_image.py +++ b/tests/multimodal/media/test_image.py @@ -11,9 +11,9 @@ ALLOWED_IMAGE_FORMATS, open_image, ) -from vllm.multimodal.media import ImageMediaIO +from vllm.multimodal.media import ImageMediaIO, MediaConnector, MediaDecodedBudget -pytestmark = pytest.mark.cpu_test +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] ASSETS_DIR = Path(__file__).parent.parent / "assets" assert ASSETS_DIR.exists() @@ -288,6 +288,68 @@ def test_image_pixel_limit_disabled(monkeypatch): assert result.media.size == (1000, 1000) +def _png_bytes(width: int, height: int) -> bytes: + from io import BytesIO + + buf = BytesIO() + Image.new("RGB", (width, height), (1, 2, 3)).save(buf, format="PNG") + return buf.getvalue() + + +def test_request_decoded_budget_rejects_aggregate(monkeypatch): + import vllm.envs as envs + + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 600) + budget = MediaDecodedBudget() + image_io = ImageMediaIO(decoded_budget=budget) + data = _png_bytes(10, 10) + + image_io.load_bytes(data) + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"): + image_io.load_bytes(data) + + +def test_request_decoded_budget_accepts_images_within_limit(monkeypatch): + import vllm.envs as envs + + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 800) + budget = MediaDecodedBudget() + image_io = ImageMediaIO(decoded_budget=budget) + data = _png_bytes(10, 10) + + image_io.load_bytes(data) + image_io.load_bytes(data) + assert budget.used == 800 + + +def test_request_decoded_budget_disabled(monkeypatch): + import vllm.envs as envs + + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 0) + budget = MediaDecodedBudget() + image_io = ImageMediaIO(decoded_budget=budget) + data = _png_bytes(100, 100) + + image_io.load_bytes(data) + image_io.load_bytes(data) + assert budget.used == 0 + + +def test_media_connector_enforces_request_decoded_budget(monkeypatch): + import pybase64 + + import vllm.envs as envs + + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 600) + b64 = pybase64.b64encode(_png_bytes(10, 10)).decode() + url = f"data:image/png;base64,{b64}" + connector = MediaConnector(decoded_budget=MediaDecodedBudget()) + + connector.fetch_image(url) + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"): + connector.fetch_image(url) + + def test_image_media_io_rejects_postscript_payload(): """EPS/PostScript bytes must not reach an external renderer.""" import pybase64 as base64 diff --git a/tests/multimodal/media/test_unprocessable_entity_error.py b/tests/multimodal/media/test_unprocessable_entity_error.py index 6d3c2a0b0aa6..8d7142d019a5 100644 --- a/tests/multimodal/media/test_unprocessable_entity_error.py +++ b/tests/multimodal/media/test_unprocessable_entity_error.py @@ -13,11 +13,16 @@ import aiohttp import pytest -from vllm.connections import MediaDownloadSizeExceededError +from vllm.connections import ( + MediaDownloadSizeExceededError, + MediaRequestDownloadSizeExceededError, +) from vllm.entrypoints.serve import create_error_response from vllm.exceptions import VLLMClientError, VLLMUnprocessableEntityError from vllm.multimodal.media import MediaConnector +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] + class TestVLLMUnprocessableEntityError: """Tests for VLLMUnprocessableEntityError exception.""" @@ -148,6 +153,21 @@ def test_fetch_image_download_size_limit_error(self): assert "VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB=256" in str(exc_info.value) + def test_fetch_image_request_download_size_limit_error(self): + connector = MediaConnector() + + with patch.object( + connector.connection, "get_bytes", new_callable=MagicMock + ) as mock_get: + mock_get.side_effect = MediaRequestDownloadSizeExceededError( + "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB=256" + ) + + with pytest.raises(VLLMUnprocessableEntityError) as exc_info: + connector.fetch_image("https://example.com/image.jpg") + + assert "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB=256" in str(exc_info.value) + class TestErrorResponse: """Tests for error response creation.""" diff --git a/tests/multimodal/media/test_video.py b/tests/multimodal/media/test_video.py index a2fedb1747f5..e3e2ae66894a 100644 --- a/tests/multimodal/media/test_video.py +++ b/tests/multimodal/media/test_video.py @@ -9,8 +9,10 @@ import numpy.typing as npt import pybase64 import pytest +import torch from PIL import Image +import vllm.envs as envs from vllm.assets.base import get_vllm_public_assets from vllm.assets.video import ( video_get_metadata, @@ -18,6 +20,7 @@ video_to_pil_images_list, ) from vllm.multimodal.media import ImageMediaIO, VideoMediaIO +from vllm.multimodal.media.base import MediaDecodedBudget from vllm.multimodal.video import ( PYNVVIDEOCODEC_VIDEO_BACKEND, VIDEO_LOADER_REGISTRY, @@ -47,6 +50,18 @@ def load_bytes( return FAKE_OUTPUT_2 +@pytest.mark.skip_global_cleanup +def test_video_request_budget_counts_torch_frames(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 16) + budget = MediaDecodedBudget() + video_io = VideoMediaIO(ImageMediaIO(decoded_budget=budget)) + frames = torch.zeros((2, 2, 2, 3), dtype=torch.uint8) + video_io.video_loader = SimpleNamespace(load_bytes=lambda *_args, **_kwargs: frames) + + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"): + video_io.load_bytes(b"video") + + def test_video_media_io_kwargs(monkeypatch: pytest.MonkeyPatch): with monkeypatch.context() as m: m.setenv("VLLM_VIDEO_LOADER_BACKEND", "assert_10_frames_1_fps") diff --git a/tests/test_audio_media_size_precheck.py b/tests/test_audio_media_size_precheck.py index 7635022369d1..e2f3c57a0f75 100644 --- a/tests/test_audio_media_size_precheck.py +++ b/tests/test_audio_media_size_precheck.py @@ -2,15 +2,19 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import AsyncIterator, Iterator +from io import BytesIO from pathlib import Path from unittest.mock import patch +import numpy as np import pytest +from PIL import Image import vllm.envs as envs -from vllm.connections import HTTPConnection +from vllm.connections import HTTPConnection, MediaRequestDownloadBudget from vllm.exceptions import VLLMValidationError -from vllm.multimodal.media import AudioMediaIO, MediaConnector +from vllm.multimodal.media import AudioMediaIO, ImageMediaIO, MediaConnector +from vllm.multimodal.media.base import MediaDecodedBudget from vllm.utils.mem_constants import KiB_bytes, MiB_bytes pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] @@ -87,6 +91,46 @@ def test_audio_load_bytes_rejects_oversized(monkeypatch: pytest.MonkeyPatch): AudioMediaIO().load_bytes(b"\x00" * (MiB_bytes + 1)) +def test_request_audio_budget_limits_multiple_clips_before_decode( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 24) + budget = MediaDecodedBudget() + clip = np.zeros(4, dtype=np.float32) + + def decode(_source, **kwargs): + if kwargs["max_decode_bytes"] < clip.nbytes: + raise ValueError("VLLM_MAX_AUDIO_DECODE_BYTES") + return clip, 16000 + + with patch("vllm.multimodal.media.audio.load_audio", side_effect=decode) as load: + audio_io = AudioMediaIO(decoded_budget=budget) + audio_io.load_bytes(b"first") + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"): + audio_io.load_bytes(b"second") + + assert budget.used == clip.nbytes + assert load.call_count == 2 + + +def test_request_decoded_budget_is_shared_by_image_and_audio( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 24) + budget = MediaDecodedBudget() + image_bytes = BytesIO() + Image.new("RGB", (2, 3)).save(image_bytes, format="PNG") + ImageMediaIO(decoded_budget=budget).load_bytes(image_bytes.getvalue()) + + with ( + patch("vllm.multimodal.media.audio.load_audio") as decode, + pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"), + ): + AudioMediaIO(decoded_budget=budget).load_bytes(b"audio") + + decode.assert_not_called() + + def test_audio_load_file_rejects_oversized( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ): @@ -138,12 +182,16 @@ async def test_audio_connector_threads_byte_limit_to_http_reader( monkeypatch: pytest.MonkeyPatch, ): monkeypatch.setattr(envs, "VLLM_MAX_AUDIO_CLIP_FILESIZE_MB", 1) + budget = MediaRequestDownloadBudget() class _Connection: async def async_get_bytes(self, _url: str, **kwargs): assert kwargs["max_bytes"] == MiB_bytes + assert kwargs["request_budget"] is budget raise VLLMValidationError("Maximum file size exceeded") - connector = MediaConnector(connection=_Connection()) + connector = MediaConnector( # type: ignore[arg-type] + connection=_Connection(), download_budget=budget + ) with pytest.raises(VLLMValidationError, match="Maximum file size exceeded"): await connector.fetch_audio_async("https://example.test/audio") diff --git a/tests/test_connections.py b/tests/test_connections.py index 983e20f6c12a..6385e372e45e 100644 --- a/tests/test_connections.py +++ b/tests/test_connections.py @@ -9,7 +9,13 @@ import pytest import vllm.envs as envs -from vllm.connections import HTTPConnection, HTTPResponseSizeExceededError +from vllm.connections import ( + HTTPConnection, + HTTPResponseSizeExceededError, + MediaRequestDownloadBudget, +) + +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] _ONE_MIB = 1024 * 1024 _SMALL_BODY = b"a" * (_ONE_MIB // 2) @@ -150,6 +156,49 @@ async def test_async_get_bytes_with_limit_streams_instead_of_materializing() -> assert response.content.iterated == 2 +def test_request_download_budget_counts_across_responses( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB", 1) + budget = MediaRequestDownloadBudget() + connection = HTTPConnection() + first = _SyncResponse([b"a" * (3 * _ONE_MIB // 4)]) + second = _SyncResponse([b"b" * (_ONE_MIB // 2)]) + + with patch.object(connection, "get_response", side_effect=[first, second]): + connection.get_bytes("http://example.com/one", request_budget=budget) + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB"): + connection.get_bytes("http://example.com/two", request_budget=budget) + + assert budget.used == 3 * _ONE_MIB // 4 + assert second.iterated == 1 + + +@pytest.mark.asyncio +async def test_async_request_download_budget_counts_across_responses( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB", 1) + budget = MediaRequestDownloadBudget() + connection = HTTPConnection() + first = _AsyncResponse([b"a" * (3 * _ONE_MIB // 4)]) + second = _AsyncResponse([b"b" * (_ONE_MIB // 2)]) + + with patch.object( + connection, "get_async_response", new=AsyncMock(side_effect=[first, second]) + ): + await connection.async_get_bytes( + "http://example.com/one", request_budget=budget + ) + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB"): + await connection.async_get_bytes( + "http://example.com/two", request_budget=budget + ) + + assert budget.used == 3 * _ONE_MIB // 4 + assert second.content.iterated == 1 + + def test_get_bytes_rejects_content_length_over_limit( monkeypatch: pytest.MonkeyPatch, local_media_server: str ) -> None: diff --git a/vllm/connections.py b/vllm/connections.py index b06d2d563fa3..ded8da430f64 100644 --- a/vllm/connections.py +++ b/vllm/connections.py @@ -3,6 +3,7 @@ import asyncio import functools +import threading import time from collections.abc import Callable, Coroutine, Mapping, MutableMapping from pathlib import Path @@ -60,6 +61,33 @@ def __init__(self, max_bytes: int, received_bytes: int | None = None) -> None: super().__init__(message) +class MediaRequestDownloadSizeExceededError(ValueError): + """Raised when all HTTP media downloads in one request exceed the limit.""" + + +class MediaRequestDownloadBudget: + """Thread-safe total of HTTP media bytes received for one request.""" + + def __init__(self) -> None: + self.used = 0 + self._lock = threading.Lock() + + def charge(self, size: int) -> None: + limit_mb = envs.VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB + if limit_mb <= 0: + return + limit = limit_mb * MiB_bytes + with self._lock: + received = self.used + size + if received > limit: + raise MediaRequestDownloadSizeExceededError( + "Request media downloads exceed " + f"VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB={limit_mb:g} " + f"({received} bytes received)." + ) + self.used = received + + def _get_media_download_limit_bytes() -> int | None: max_size_mb = envs.VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB if max_size_mb <= 0: @@ -105,19 +133,26 @@ def _append_chunk( received_bytes: int, max_bytes: int | None, media_download_limit: bool, + request_budget: MediaRequestDownloadBudget | None = None, ) -> int: if not chunk: return received_bytes received_bytes += len(chunk) if max_bytes is not None and received_bytes > max_bytes: _raise_size_exceeded(max_bytes, received_bytes, media_download_limit) + if request_budget is not None: + request_budget.charge(len(chunk)) chunks.append(chunk) return received_bytes -def _read_response_bytes(response: requests.Response, max_bytes: int | None) -> bytes: +def _read_response_bytes( + response: requests.Response, + max_bytes: int | None, + request_budget: MediaRequestDownloadBudget | None = None, +) -> bytes: max_bytes, media_download_limit = _get_response_size_limit(max_bytes) - if max_bytes is None: + if max_bytes is None and request_budget is None: return response.content _check_content_length( @@ -129,16 +164,23 @@ def _read_response_bytes(response: requests.Response, max_bytes: int | None) -> received_bytes = 0 for chunk in response.iter_content(_RESPONSE_READ_CHUNK_SIZE): received_bytes = _append_chunk( - chunks, chunk, received_bytes, max_bytes, media_download_limit + chunks, + chunk, + received_bytes, + max_bytes, + media_download_limit, + request_budget, ) return b"".join(chunks) async def _async_read_response_bytes( - response: aiohttp.ClientResponse, max_bytes: int | None + response: aiohttp.ClientResponse, + max_bytes: int | None, + request_budget: MediaRequestDownloadBudget | None = None, ) -> bytes: max_bytes, media_download_limit = _get_response_size_limit(max_bytes) - if max_bytes is None: + if max_bytes is None and request_budget is None: return await response.read() _check_content_length(response.content_length, max_bytes, media_download_limit) @@ -146,7 +188,12 @@ async def _async_read_response_bytes( received_bytes = 0 async for chunk in response.content.iter_chunked(_RESPONSE_READ_CHUNK_SIZE): received_bytes = _append_chunk( - chunks, chunk, received_bytes, max_bytes, media_download_limit + chunks, + chunk, + received_bytes, + max_bytes, + media_download_limit, + request_budget, ) return b"".join(chunks) @@ -407,6 +454,7 @@ def get_bytes( timeout: float | None = None, allow_redirects: bool = True, max_bytes: int | None = None, + request_budget: MediaRequestDownloadBudget | None = None, ) -> bytes: with self.get_response( url, @@ -416,7 +464,7 @@ def get_bytes( ) as r: r.raise_for_status() - return _read_response_bytes(r, max_bytes) + return _read_response_bytes(r, max_bytes, request_budget) @_async_retry async def async_get_bytes( @@ -426,6 +474,7 @@ async def async_get_bytes( timeout: float | None = None, allow_redirects: bool = True, max_bytes: int | None = None, + request_budget: MediaRequestDownloadBudget | None = None, ) -> bytes: async with await self.get_async_response( url, @@ -434,7 +483,7 @@ async def async_get_bytes( ) as r: r.raise_for_status() - return await _async_read_response_bytes(r, max_bytes) + return await _async_read_response_bytes(r, max_bytes, request_budget) def get_text(self, url: str, *, timeout: float | None = None) -> str: with self.get_response(url, timeout=timeout) as r: diff --git a/vllm/envs.py b/vllm/envs.py index bb1c03a9dabd..dbbc49ff59a8 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -78,6 +78,7 @@ VLLM_MEDIA_CACHE_TTL_HOURS: float = 24 VLLM_MEDIA_FETCH_MAX_RETRIES: int = 3 VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB: int = 256 + VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB: int = 256 VLLM_MEDIA_URL_ALLOW_REDIRECTS: bool = True VLLM_MEDIA_LOADING_THREAD_COUNT: int = 8 VLLM_MAX_AUDIO_CLIP_FILESIZE_MB: int = 25 @@ -86,6 +87,7 @@ VLLM_MAX_AUDIO_PREPROCESS_WORKERS: int = max(1, min(os.cpu_count() or 1, 2)) VLLM_MAX_EMBED_DECODE_BYTES: int = 2_147_483_648 VLLM_MAX_IMAGE_PIXELS: int = 178_956_970 + VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES: int = 715_827_880 VLLM_VIDEO_LOADER_BACKEND: str = "opencv" VLLM_MEDIA_CONNECTOR: str = "http" VLLM_TARGET_DEVICE: str = "cuda" @@ -1063,6 +1065,15 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_MAX_IMAGE_PIXELS": lambda: int( os.getenv("VLLM_MAX_IMAGE_PIXELS", "178956970") ), + # Maximum aggregate decoded bytes across media in one request. + # The default allows one image at the per-image pixel limit at 4 bytes + # per pixel. Set to 0 to disable. + "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES": lambda: int( + os.getenv("VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", "715827880") + ), + "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB": lambda: int( + os.getenv("VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB", "256") + ), # Backend for Video IO — selects the frame-sampling algorithm. # - "opencv": uniform sampling. # - "opencv_dynamic": duration-aware dynamic sampling. @@ -2412,6 +2423,8 @@ def compile_factors() -> dict[str, object]: "VLLM_MAX_AUDIO_PREPROCESS_WORKERS", "VLLM_MAX_EMBED_DECODE_BYTES", "VLLM_MAX_IMAGE_PIXELS", + "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", + "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB", "VLLM_VIDEO_LOADER_BACKEND", "VLLM_MEDIA_CONNECTOR", "VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME", diff --git a/vllm/multimodal/media/__init__.py b/vllm/multimodal/media/__init__.py index 94b08c484882..84fc7d6ed88c 100644 --- a/vllm/multimodal/media/__init__.py +++ b/vllm/multimodal/media/__init__.py @@ -1,13 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from .audio import AudioEmbeddingMediaIO, AudioMediaIO -from .base import MediaIO, MediaWithBytes +from .base import MediaDecodedBudget, MediaIO, MediaWithBytes from .connector import MEDIA_CONNECTOR_REGISTRY, MediaConnector from .image import ImageEmbeddingMediaIO, ImageMediaIO from .video import VIDEO_LOADER_REGISTRY, VideoMediaIO __all__ = [ "MediaIO", + "MediaDecodedBudget", "MediaWithBytes", "AudioEmbeddingMediaIO", "AudioMediaIO", diff --git a/vllm/multimodal/media/audio.py b/vllm/multimodal/media/audio.py index c05f4e4dcc8f..d36b4a984e66 100644 --- a/vllm/multimodal/media/audio.py +++ b/vllm/multimodal/media/audio.py @@ -23,7 +23,7 @@ safe_to_dense, ) -from .base import MediaIO +from .base import MediaDecodedBudget, MediaIO logger = init_logger(__name__) @@ -527,8 +527,11 @@ class AudioMediaIO(MediaIO[tuple[npt.NDArray, float]]): error handling. """ - def __init__(self, **kwargs) -> None: + def __init__( + self, *, decoded_budget: MediaDecodedBudget | None = None, **kwargs + ) -> None: super().__init__() + self.decoded_budget = decoded_budget # Allow per-server override of the audio decoding backend, e.g.: # --media-io-kwargs '{"audio": {"audio_backend": "torchcodec"}}' @@ -564,13 +567,44 @@ def _validate_encoded_size(self, size: int) -> None: def load_bytes(self, data: bytes) -> tuple[npt.NDArray, float]: self._validate_encoded_size(len(data)) - return load_audio( - BytesIO(data), - sr=None, - max_duration_s=envs.VLLM_MAX_AUDIO_DECODE_DURATION_S, - max_decode_bytes=envs.VLLM_MAX_AUDIO_DECODE_BYTES, - backend=self.audio_backend, - ) + return self._load_with_budget(BytesIO(data)) + + def _load_with_budget(self, source: BytesIO | Path) -> tuple[npt.NDArray, float]: + budget = self.decoded_budget + if budget is None: + return load_audio( + source, + sr=None, + max_duration_s=envs.VLLM_MAX_AUDIO_DECODE_DURATION_S, + max_decode_bytes=envs.VLLM_MAX_AUDIO_DECODE_BYTES, + backend=self.audio_backend, + ) + + max_decode_bytes = envs.VLLM_MAX_AUDIO_DECODE_BYTES + reserved = budget.reserve(max_decode_bytes) + decoded_bytes = 0 + try: + try: + audio = load_audio( + source, + sr=None, + max_duration_s=envs.VLLM_MAX_AUDIO_DECODE_DURATION_S, + max_decode_bytes=reserved, + backend=self.audio_backend, + ) + except ValueError as exc: + if reserved < max_decode_bytes and "VLLM_MAX_AUDIO_DECODE_BYTES" in str( + exc + ): + raise ValueError( + "Request decoded media exceeds " + "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES." + ) from exc + raise + decoded_bytes = audio[0].nbytes + return audio + finally: + budget.settle(reserved, decoded_bytes) def load_base64( self, @@ -588,13 +622,7 @@ def load_base64( def load_file(self, filepath: Path) -> tuple[npt.NDArray, float]: self._validate_encoded_size(filepath.stat().st_size) - return load_audio( - filepath, - sr=None, - max_duration_s=envs.VLLM_MAX_AUDIO_DECODE_DURATION_S, - max_decode_bytes=envs.VLLM_MAX_AUDIO_DECODE_BYTES, - backend=self.audio_backend, - ) + return self._load_with_budget(filepath) def encode_base64( self, diff --git a/vllm/multimodal/media/base.py b/vllm/multimodal/media/base.py index d3e622156fc8..952b92e96939 100644 --- a/vllm/multimodal/media/base.py +++ b/vllm/multimodal/media/base.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import threading from abc import ABC, abstractmethod from collections.abc import Iterable, Iterator from dataclasses import dataclass, field @@ -9,9 +10,79 @@ import numpy as np +import vllm.envs as envs + _T = TypeVar("_T") +class MediaDecodedBudget: + """Thread-safe decoded-media byte accumulator for one request.""" + + def __init__(self) -> None: + self.used = 0 + self._reserved = 0 + self._condition = threading.Condition() + + def reserve(self, max_bytes: int) -> int: + """Reserve decoder capacity, waiting for in-flight decodes if needed.""" + limit = envs.VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES + if limit <= 0: + return max_bytes + max_bytes = max(0, max_bytes) + with self._condition: + while True: + available = limit - self.used - self._reserved + if available >= max_bytes or self._reserved == 0: + if available <= 0: + raise ValueError( + "Request decoded media exceeds " + "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES." + ) + reserved = min(max_bytes, available) + self._reserved += reserved + return reserved + self._condition.wait() + + def settle(self, reserved: int, decoded_bytes: int) -> None: + """Replace a decoder reservation with its actual decoded size.""" + limit = envs.VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES + if limit <= 0: + return + with self._condition: + self._reserved -= reserved + try: + used = self.used + decoded_bytes + if decoded_bytes > reserved or used > limit: + raise ValueError( + "Request decoded media exceeds " + f"VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES={limit} " + f"({used} bytes)." + ) + self.used = used + finally: + self._condition.notify_all() + + def charge(self, decoded_bytes: int) -> None: + limit = envs.VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES + if limit <= 0 or decoded_bytes <= 0: + return + with self._condition: + while ( + self.used + self._reserved + decoded_bytes > limit + and self._reserved > 0 + and self.used + decoded_bytes <= limit + ): + self._condition.wait() + used = self.used + decoded_bytes + if used > limit: + raise ValueError( + "Request decoded media exceeds " + f"VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES={limit} " + f"({used} bytes)." + ) + self.used = used + + @dataclass class MediaWithBytes(Generic[_T]): """Wrapper that couples a media object with its original encoded bytes. diff --git a/vllm/multimodal/media/connector.py b/vllm/multimodal/media/connector.py index 297dc91af232..e5a52417fee3 100644 --- a/vllm/multimodal/media/connector.py +++ b/vllm/multimodal/media/connector.py @@ -25,6 +25,8 @@ from vllm.connections import ( HTTPConnection, MediaDownloadSizeExceededError, + MediaRequestDownloadBudget, + MediaRequestDownloadSizeExceededError, global_http_connection, ) from vllm.exceptions import VLLMUnprocessableEntityError, VLLMValidationError @@ -33,7 +35,7 @@ from vllm.utils.registry import ExtensionManager from .audio import AudioEmbeddingMediaIO, AudioMediaIO -from .base import MediaIO, MediaWithBytes +from .base import MediaDecodedBudget, MediaIO, MediaWithBytes from .image import ImageEmbeddingMediaIO, ImageMediaIO from .video import VideoEmbeddingMediaIO, VideoMediaIO @@ -109,7 +111,9 @@ def _wrap_media_fetch_error( value=url, ) - if isinstance(exc, MediaDownloadSizeExceededError): + if isinstance( + exc, (MediaDownloadSizeExceededError, MediaRequestDownloadSizeExceededError) + ): return VLLMUnprocessableEntityError( f"Failed to fetch media from URL: {exc}", parameter="image_url", @@ -162,6 +166,8 @@ def __init__( *, allowed_local_media_path: str = "", allowed_media_domains: list[str] | None = None, + decoded_budget: MediaDecodedBudget | None = None, + download_budget: MediaRequestDownloadBudget | None = None, ) -> None: """Args: media_io_kwargs: Additional args passed to process media @@ -201,6 +207,8 @@ def __init__( if allowed_media_domains is None: allowed_media_domains = [] self.allowed_media_domains = allowed_media_domains + self.decoded_budget = decoded_budget + self.download_budget = download_budget # Media download cache (opt-in via VLLM_MEDIA_CACHE) self._media_cache_dir: str | None = None @@ -230,6 +238,12 @@ def __init__( media_cache, ) + def _image_io(self, image_mode: str | None = "RGB") -> ImageMediaIO: + return ImageMediaIO( + **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})), + decoded_budget=self.decoded_budget, + ) + def _get_cached_bytes(self, url: str) -> bytes | None: """Return cached bytes for a URL, or None if not cached/expired.""" if not self._media_cache_dir: @@ -385,12 +399,18 @@ def load_from_url( return media_io.load_bytes(cached) connection = self.connection + budget_kwargs = ( + {"request_budget": self.download_budget} + if self.download_budget is not None + else {} + ) try: data = connection.get_bytes( url_spec.url, timeout=fetch_timeout, allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS, max_bytes=max_bytes, + **budget_kwargs, ) except Exception as e: wrapped = _wrap_media_fetch_error(url, e) @@ -438,12 +458,18 @@ async def load_from_url_async( return await future connection = self.connection + budget_kwargs = ( + {"request_budget": self.download_budget} + if self.download_budget is not None + else {} + ) try: data = await connection.async_get_bytes( url_spec.url, timeout=fetch_timeout, allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS, max_bytes=max_bytes, + **budget_kwargs, ) except Exception as e: wrapped = _wrap_media_fetch_error(url, e) @@ -470,7 +496,9 @@ def fetch_audio( audio_url: str, ) -> tuple[np.ndarray, int | float]: """Load audio from a URL.""" - audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {})) + audio_io = AudioMediaIO( + decoded_budget=self.decoded_budget, **self.media_io_kwargs.get("audio", {}) + ) return self.load_from_url( audio_url, @@ -483,7 +511,9 @@ async def fetch_audio_async( audio_url: str, ) -> tuple[np.ndarray, int | float]: """Asynchronously fetch audio from a URL.""" - audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {})) + audio_io = AudioMediaIO( + decoded_budget=self.decoded_budget, **self.media_io_kwargs.get("audio", {}) + ) return await self.load_from_url_async( audio_url, @@ -503,9 +533,7 @@ def fetch_image( `media_io_kwargs={"image": {"image_mode": None}}` to keep the original image mode (e.g. preserving the alpha channel). """ - image_io = ImageMediaIO( - **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})) - ) + image_io = self._image_io(image_mode) try: return self.load_from_url( @@ -529,9 +557,7 @@ async def fetch_image_async( `media_io_kwargs={"image": {"image_mode": None}}` to keep the original image mode (e.g. preserving the alpha channel). """ - image_io = ImageMediaIO( - **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})) - ) + image_io = self._image_io(image_mode) try: return await self.load_from_url_async( @@ -551,9 +577,7 @@ def fetch_video( video_processor: str | None = None, ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: """Load video from an HTTP or base64 data URL.""" - image_io = ImageMediaIO( - **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})) - ) + image_io = self._image_io(image_mode) video_io_kwargs = dict(self.media_io_kwargs.get("video", {})) if "video_backend" not in video_io_kwargs and ( video_backend := get_video_loader_backend_for_processor(video_processor) @@ -580,9 +604,7 @@ async def fetch_video_async( `media_io_kwargs={"image": {"image_mode": None}}` to keep the original image mode (e.g. preserving the alpha channel). """ - image_io = ImageMediaIO( - **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})) - ) + image_io = self._image_io(image_mode) video_io_kwargs = dict(self.media_io_kwargs.get("video", {})) if "video_backend" not in video_io_kwargs and ( video_backend := get_video_loader_backend_for_processor(video_processor) diff --git a/vllm/multimodal/media/image.py b/vllm/multimodal/media/image.py index fb05f4ed741d..1e2e106189b4 100644 --- a/vllm/multimodal/media/image.py +++ b/vllm/multimodal/media/image.py @@ -17,7 +17,7 @@ ) from ..image import convert_image_mode, normalize_image, open_image, rgba_to_rgb -from .base import MediaIO, MediaWithBytes +from .base import MediaDecodedBudget, MediaIO, MediaWithBytes MAGIC_NUMPY_PREFIX = b"\x93NUMPY" # https://numpy.org/devdocs/reference/generated/numpy.lib.format.html#format-version-1-0 @@ -28,12 +28,19 @@ class ImageMediaIO(MediaIO[Image.Image]): error handling. """ - def __init__(self, image_mode: str | None = "RGB", **kwargs) -> None: + def __init__( + self, + image_mode: str | None = "RGB", + *, + decoded_budget: MediaDecodedBudget | None = None, + **kwargs, + ) -> None: super().__init__() # Target mode for loaded images; `None` keeps the original mode # (i.e. no conversion, alpha channel is preserved as-is). self.image_mode = image_mode + self.decoded_budget = decoded_budget # `kwargs` contains custom arguments from # --media-io-kwargs for this modality, merged with # per-request runtime media_io_kwargs via merge_kwargs(). @@ -87,6 +94,8 @@ def load_bytes(self, data: bytes) -> MediaWithBytes[Image.Image]: f"the maximum of {max_pixels} pixels. Set " f"VLLM_MAX_IMAGE_PIXELS to increase this limit." ) + if self.decoded_budget is not None: + self.decoded_budget.charge(w * h * 4) image = normalize_image(image) image.load() converted = self._convert_image_mode(image) diff --git a/vllm/multimodal/media/video.py b/vllm/multimodal/media/video.py index 7528331cdfd0..569ef3dc93bf 100644 --- a/vllm/multimodal/media/video.py +++ b/vllm/multimodal/media/video.py @@ -122,6 +122,16 @@ def load_bytes( video = self.video_loader.load_bytes( data, num_frames=self.num_frames, **self.kwargs ) + frames = video[0] if isinstance(video, tuple) else video + budget = getattr(self.image_io, "decoded_budget", None) + shape = getattr(frames, "shape", None) + if budget is not None and shape is not None and len(shape) >= 3: + decoded_bytes = ( + frames.numel() * frames.element_size() + if isinstance(frames, torch.Tensor) + else frames.nbytes + ) + budget.charge(decoded_bytes) return MediaWithBytes(video, data) def load_base64( diff --git a/vllm/renderers/chat_utils.py b/vllm/renderers/chat_utils.py index f18407e535e9..540564baab8d 100644 --- a/vllm/renderers/chat_utils.py +++ b/vllm/renderers/chat_utils.py @@ -53,6 +53,7 @@ from vllm import envs from vllm.config import ModelConfig +from vllm.connections import MediaRequestDownloadBudget from vllm.exceptions import VLLMValidationError from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict from vllm.logger import init_logger @@ -67,6 +68,7 @@ VisionChunkVideo, ) from vllm.multimodal.media import MEDIA_CONNECTOR_REGISTRY, MediaConnector +from vllm.multimodal.media.base import MediaDecodedBudget from vllm.multimodal.processing import BaseMultiModalProcessor from vllm.renderers.embed_utils import ( safe_load_prompt_embeds, @@ -609,6 +611,8 @@ def __init__( self._items_by_modality = defaultdict[str, list[_T]](list) # Track original modality for each vision_chunk item (image or video) self._modality_order = defaultdict[str, list[str]](list) + self.decoded_budget = MediaDecodedBudget() + self.download_budget = MediaRequestDownloadBudget() @cached_property def use_unified_vision_chunk_modality(self) -> bool: @@ -1052,6 +1056,8 @@ def _connector(self) -> MediaConnector: media_io_kwargs=self._tracker.media_io_kwargs, allowed_local_media_path=self._tracker.allowed_local_media_path, allowed_media_domains=self._tracker.allowed_media_domains, + decoded_budget=self._tracker.decoded_budget, + download_budget=self._tracker.download_budget, ) @property @@ -1138,6 +1144,10 @@ def parse_audio_embeds( def parse_image_pil( self, image_pil: Image.Image | None, uuid: str | None = None ) -> None: + if image_pil is not None: + self._tracker.decoded_budget.charge( + image_pil.size[0] * image_pil.size[1] * 4 + ) placeholder = self._tracker.add("image", (image_pil, uuid)) self._add_placeholder("image", placeholder) @@ -1233,6 +1243,8 @@ def _connector(self) -> MediaConnector: media_io_kwargs=self._tracker.media_io_kwargs, allowed_local_media_path=self._tracker.allowed_local_media_path, allowed_media_domains=self._tracker.allowed_media_domains, + decoded_budget=self._tracker.decoded_budget, + download_budget=self._tracker.download_budget, ) @property @@ -1351,6 +1363,10 @@ def parse_image_pil( image_pil: Image.Image | None, uuid: str | None = None, ) -> None: + if image_pil is not None: + self._tracker.decoded_budget.charge( + image_pil.size[0] * image_pil.size[1] * 4 + ) placeholder = self._tracker.add( "image", partial(self._item_with_uuid_async, image_pil, uuid) ) From 54db759eea66729faa6d41e8cdaceb02833cef7c Mon Sep 17 00:00:00 2001 From: samayp01 <67844360+samayp01@users.noreply.github.com> Date: Tue, 13 Oct 2026 06:27:34 -0400 Subject: [PATCH 376/624] [Bugfix][Structured Output] Reject structural_tag and unhandled specs in backend validation (#58862) Signed-off-by: Samay Co-authored-by: Artem Perevedentsev --- tests/v1/structured_output/test_validation.py | 57 +++++++++++++++++++ .../backend_lm_format_enforcer.py | 16 +++++- vllm/v1/structured_output/backend_outlines.py | 14 ++++- vllm/v1/structured_output/backend_xgrammar.py | 14 ++++- 4 files changed, 94 insertions(+), 7 deletions(-) diff --git a/tests/v1/structured_output/test_validation.py b/tests/v1/structured_output/test_validation.py index 276056b501cd..41166c260280 100644 --- a/tests/v1/structured_output/test_validation.py +++ b/tests/v1/structured_output/test_validation.py @@ -17,6 +17,13 @@ ) from vllm.tokenizers import mistral as mistral_tokenizers from vllm.v1.structured_output import backend_guidance +from vllm.v1.structured_output.backend_lm_format_enforcer import ( + validate_structured_output_request_lm_format_enforcer, +) +from vllm.v1.structured_output.backend_outlines import ( + validate_structured_output_request_outlines, +) +from vllm.v1.structured_output.backend_xgrammar import validate_xgrammar_grammar pytestmark = pytest.mark.cpu_test @@ -213,6 +220,9 @@ def test_schema_at_nesting_limit_accepted(field): INVALID_JSON_SCHEMA = {"type": "object", "properties": {"name": {"type": "str"}}} +STRUCTURAL_TAG = ( + '{"type": "structural_tag", "format": {"type": "const_string", "value": "hi"}}' +) @pytest.mark.parametrize( @@ -227,6 +237,9 @@ def test_schema_at_nesting_limit_accepted(field): ("lm-format-enforcer", StructuredOutputsParams(grammar="not a grammar")), ("outlines", StructuredOutputsParams(regex="(")), ("guidance", StructuredOutputsParams(structural_tag='{"nope": 1}')), + ("outlines", StructuredOutputsParams(structural_tag=STRUCTURAL_TAG)), + ("lm-format-enforcer", StructuredOutputsParams(structural_tag=STRUCTURAL_TAG)), + ("outlines", StructuredOutputsParams(json_object=True)), ], ) def test_unsupported_grammar_is_a_client_error(backend, structured_outputs): @@ -241,6 +254,50 @@ def test_unsupported_grammar_is_a_client_error(backend, structured_outputs): ) +@pytest.mark.parametrize( + "backend, structured_outputs", + [ + ("xgrammar", StructuredOutputsParams(json_object=True)), + ("xgrammar", StructuredOutputsParams(regex="[0-9]+")), + ("xgrammar", StructuredOutputsParams(regex="")), + ("xgrammar", StructuredOutputsParams(json={})), + ("lm-format-enforcer", StructuredOutputsParams(json_object=True)), + ("lm-format-enforcer", StructuredOutputsParams(regex="")), + ("lm-format-enforcer", StructuredOutputsParams(json={})), + ("outlines", StructuredOutputsParams(regex="")), + ("outlines", StructuredOutputsParams(json={})), + ], +) +def test_supported_spec_is_not_rejected(backend, structured_outputs): + """Supported specs, including the falsy but valid regex="" and json={}, must + not reach the validators' final rejection.""" + params = SamplingParams(structured_outputs=structured_outputs) + params._validate_structured_outputs( + _StubModelConfig(is_diffusion=False), + StructuredOutputsConfig(backend=backend), + tokenizer=object(), + ) + + +@pytest.mark.parametrize( + "validate", + [ + validate_structured_output_request_outlines, + validate_structured_output_request_lm_format_enforcer, + validate_xgrammar_grammar, + ], + ids=["outlines", "lm-format-enforcer", "xgrammar"], +) +def test_unhandled_spec_is_a_validation_error(validate): + """json_object=False is rejected before backend selection, so only a direct + validator call reaches the final branch; it must still be a 400.""" + params = SamplingParams( + structured_outputs=StructuredOutputsParams(json_object=False) + ) + with pytest.raises(VLLMValidationError, match="does not support this"): + validate(params) + + @pytest.mark.parametrize( "schema, expected_backend", [ diff --git a/vllm/v1/structured_output/backend_lm_format_enforcer.py b/vllm/v1/structured_output/backend_lm_format_enforcer.py index 9a16ddc1fe38..b48cdc791673 100644 --- a/vllm/v1/structured_output/backend_lm_format_enforcer.py +++ b/vllm/v1/structured_output/backend_lm_format_enforcer.py @@ -159,7 +159,7 @@ def validate_structured_output_request_lm_format_enforcer(params: SamplingParams so_params = params.structured_outputs - if so_params.regex: + if so_params.regex is not None: try: compile_regex_with_timeout( lmformatenforcer.RegexParser, @@ -170,7 +170,7 @@ def validate_structured_output_request_lm_format_enforcer(params: SamplingParams f"Failed to compile regex for lm-format-enforcer: {err}" ) from err return - elif so_params.json: + elif so_params.json is not None: if isinstance(so_params.json, str): try: # make sure schema is valid json @@ -185,10 +185,20 @@ def validate_structured_output_request_lm_format_enforcer(params: SamplingParams f"Error serializing structured outputs jsonschema: {e}" ) from e return - elif so_params.choice: + elif so_params.choice or so_params.json_object: return elif so_params.grammar: raise VLLMValidationError( "LM Format Enforcer structured outputs backend " "does not support grammar specifications" ) + elif so_params.structural_tag: + raise VLLMValidationError( + "LM Format Enforcer structured outputs backend " + "does not support structural_tag specifications" + ) + else: + raise VLLMValidationError( + "LM Format Enforcer structured outputs backend " + "does not support this structured outputs specification" + ) diff --git a/vllm/v1/structured_output/backend_outlines.py b/vllm/v1/structured_output/backend_outlines.py index 176d1a46e684..09588a680ab9 100644 --- a/vllm/v1/structured_output/backend_outlines.py +++ b/vllm/v1/structured_output/backend_outlines.py @@ -198,9 +198,9 @@ def validate_structured_output_request_outlines(params: SamplingParams): so_params = params.structured_outputs - if so_params.regex: + if so_params.regex is not None: validate_regex_is_buildable(so_params.regex) - elif so_params.json: + elif so_params.json is not None: if isinstance(so_params.json, str): try: # make sure schema is valid json @@ -236,6 +236,16 @@ def validate_structured_output_request_outlines(params: SamplingParams): "Outlines structured outputs backend " "does not support grammar specifications" ) + elif so_params.structural_tag: + raise VLLMValidationError( + "Outlines structured outputs backend " + "does not support structural_tag specifications" + ) + else: + raise VLLMValidationError( + "Outlines structured outputs backend " + "does not support this structured outputs specification" + ) def _prefix_needs_context(parsed) -> bool: diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 0199c5f75d6f..9631766d9267 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -414,7 +414,7 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: so_params = sampling_params.structured_outputs - if so_params.regex: + if so_params.regex is not None: # A NUL byte is never meaningful in a regex pattern and is not handled # by xgrammar's native regex converter. Reject it here, before the # pattern reaches that native code; the try/except below does not cover @@ -432,6 +432,7 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: raise VLLMValidationError( f"Failed to transform regex into a grammar: {err}" ) from err + return if so_params.choice: choice_grammar = choice_as_grammar(so_params.choice) @@ -445,7 +446,7 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: so_params.grammar = choice_grammar return - if so_params.json: + if so_params.json is not None: check_json_nesting(so_params.json) if isinstance(so_params.json, str): try: @@ -470,6 +471,9 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: ) from err return + if so_params.json_object: + return + if so_params.grammar: # Parse the grammar with the same syntax `compile_grammar` will use, # but don't compile it. The grammar is passed on unchanged. @@ -502,3 +506,9 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: xgr.Grammar.from_structural_tag(so_params.structural_tag) except Exception as e: raise VLLMValidationError("Invalid structural tag specification.") from e + return + + raise VLLMValidationError( + "XGrammar structured outputs backend " + "does not support this structured outputs specification" + ) From 426f3bd36b941c34ad183e10cf2c30ebd62a4755 Mon Sep 17 00:00:00 2001 From: noone-zhhh Date: Tue, 13 Oct 2026 18:49:39 +0800 Subject: [PATCH 377/624] [Bugfix][Model][Rust Frontend] Qwen3_ASR: Handle ragged Qwen audio batches from Rust frontend (#60979) Signed-off-by: zh98 Co-authored-by: Isotr0py --- vllm/model_executor/models/qwen2_5_omni_thinker.py | 12 +++++++----- vllm/model_executor/models/qwen3_asr.py | 6 ++++-- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/vllm/model_executor/models/qwen2_5_omni_thinker.py b/vllm/model_executor/models/qwen2_5_omni_thinker.py index 986de3e09636..c09ae1a87e79 100644 --- a/vllm/model_executor/models/qwen2_5_omni_thinker.py +++ b/vllm/model_executor/models/qwen2_5_omni_thinker.py @@ -132,11 +132,11 @@ def _get_feature_tensor(feature: MultiModalFeatureSpec, key: str) -> torch.Tenso def unpad_and_flat_audio_features( - input_audio_features: torch.Tensor, + input_audio_features: torch.Tensor | Sequence[torch.Tensor], audio_feature_lengths: torch.Tensor, ) -> torch.Tensor: - """Unpad and flatten batched audio features.""" - if len(audio_feature_lengths) != input_audio_features.size(0): + """Unpad and flatten stacked or ragged per-audio features.""" + if len(audio_feature_lengths) != len(input_audio_features): raise ValueError( "Length of audio_feature_lengths must match " "the batch size of input_audio_features." @@ -920,8 +920,10 @@ def _parse_and_validate_audio_input( # inputs features from rust frontend is batched and padded # with shape [batch_size, n_mels, padded_seq_len], different - # from python's shape [n_mels, batch_size * seq_len] - if ( + # from python's shape [n_mels, batch_size * seq_len]. + # Requests with different padded lengths are batched as a + # list of [n_mels, padded_seq_len] tensors. + if isinstance(input_audio_features, (list, tuple)) or ( isinstance(input_audio_features, torch.Tensor) and input_audio_features.dim() == 3 ): diff --git a/vllm/model_executor/models/qwen3_asr.py b/vllm/model_executor/models/qwen3_asr.py index 161d1efacd06..0283fd814d8d 100644 --- a/vllm/model_executor/models/qwen3_asr.py +++ b/vllm/model_executor/models/qwen3_asr.py @@ -478,8 +478,10 @@ def _parse_and_validate_audio_input( # inputs features from rust frontend is batched and padded # with shape [batch_size, n_mels, padded_seq_len], different - # from python's shape [n_mels, batch_size * seq_len] - if ( + # from python's shape [n_mels, batch_size * seq_len]. + # Requests with different padded lengths are batched as a + # list of [n_mels, padded_seq_len] tensors. + if isinstance(input_audio_features, (list, tuple)) or ( isinstance(input_audio_features, torch.Tensor) and input_audio_features.dim() == 3 ): From e6e84f350c647c6e76dad2e39b08cb85e17f1272 Mon Sep 17 00:00:00 2001 From: Linze Shi Date: Tue, 13 Oct 2026 04:16:21 -0700 Subject: [PATCH 378/624] [Feature][Pooling] Add offline fixed-length late chunking (#60164) Signed-off-by: Linze-Shi --- docs/models/pooling_models/token_embed.md | 34 ++ .../token_embed/late_chunking_offline.py | 72 ++++ .../entrypoints/pooling/test_late_chunking.py | 314 ++++++++++++++++++ .../pooling/token_embed/test_offline.py | 221 +++++++++++- .../layers/test_pooler_heads.py | 1 + .../layers/test_pooler_methods.py | 143 +++++++- tests/test_pooling_params.py | 92 +++++ .../entrypoints/pooling/embed/io_processor.py | 74 +++++ vllm/entrypoints/pooling/late_chunking.py | 92 +++++ vllm/entrypoints/pooling/offline.py | 2 +- vllm/entrypoints/pooling/typing.py | 15 +- vllm/model_executor/layers/pooler/special.py | 5 + .../layers/pooler/tokwise/heads.py | 23 ++ vllm/outputs.py | 17 + vllm/pooling_params.py | 49 ++- 15 files changed, 1147 insertions(+), 7 deletions(-) create mode 100644 examples/pooling/token_embed/late_chunking_offline.py create mode 100644 tests/entrypoints/pooling/test_late_chunking.py create mode 100644 vllm/entrypoints/pooling/late_chunking.py diff --git a/docs/models/pooling_models/token_embed.md b/docs/models/pooling_models/token_embed.md index 0c2a322e80fc..2bc51f4bec91 100644 --- a/docs/models/pooling_models/token_embed.md +++ b/docs/models/pooling_models/token_embed.md @@ -90,6 +90,7 @@ The following [pooling parameters][vllm.PoolingParams] are supported. ```python --8<-- "vllm/pooling_params.py:common-pooling-params" --8<-- "vllm/pooling_params.py:embed-pooling-params" +--8<-- "vllm/pooling_params.py:token-embed-pooling-params" ``` ### `LLM.encode` @@ -108,6 +109,39 @@ data = output.outputs.data print(f"Data: {data!r}") ``` +### Fixed-length late chunking + +Late chunking encodes the full input before averaging contextual token representations within each chunk. The chunk means pass through the model's embedding head, including its projection and optional normalization. + +This mode is currently available through the offline `LLM.encode` API for plain-text input. With an instance configured for `token_embed`, enable it per request: + +```python +from vllm import PoolingParams +from vllm.pooling_params import LateChunkingParams + +outputs = llm.encode( + [document], + pooling_task="token_embed", + pooling_params=PoolingParams( + late_chunking_params=LateChunkingParams(chunk_size=256) + ), +) +vectors = outputs[0].outputs.data # [num_chunks, embedding_dimension] +chunks = outputs[0].late_chunking.chunks +``` + +For `N` model-input tokens and chunk size `C`, the output contains `ceil(N / C)` vectors. The last chunk uses its actual length. Special tokens and instruction prefixes count toward both chunk boundaries and means. Omitting `late_chunking_params` preserves ordinary token embeddings. + +Each vector row corresponds to one entry in `chunks`. Its `token_range` and `char_range` are zero-based, half-open ranges. Character ranges index the submitted Python string; use `document[start:end]` to retrieve the source text. A chunk containing only special tokens has `char_range=None`. + +Use nonempty text within the model's context window and a fast tokenizer that provides aligned source offsets. Input truncation, explicit padding and renderer text normalization are not supported. The custom document/query outputs of `JinaForRanking` are also unsupported. Late chunking does not extend the model's context window or reduce its attention computation. + +The [offline example](../../../examples/pooling/token_embed/late_chunking_offline.py) loads an E5 model and prints the vector shape, token ranges and source excerpts for two documents: + +```bash +python examples/pooling/token_embed/late_chunking_offline.py --chunk-size 16 +``` + ### `LLM.score` The [score][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.score] method outputs similarity scores between sentence pairs. diff --git a/examples/pooling/token_embed/late_chunking_offline.py b/examples/pooling/token_embed/late_chunking_offline.py new file mode 100644 index 000000000000..4f82b17cceff --- /dev/null +++ b/examples/pooling/token_embed/late_chunking_offline.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from argparse import Namespace + +from vllm import LLM, EngineArgs, PoolingParams +from vllm.config import PoolerConfig +from vllm.pooling_params import LateChunkingParams +from vllm.utils.argparse_utils import FlexibleArgumentParser + + +def parse_args(): + parser = FlexibleArgumentParser() + parser = EngineArgs.add_cli_args(parser) + parser.add_argument( + "--chunk-size", + type=int, + default=16, + help="Number of model-input tokens per chunk, including special tokens.", + ) + parser.set_defaults( + model="intfloat/multilingual-e5-small", + runner="pooling", + pooler_config=PoolerConfig(task="token_embed"), + max_model_len=512, + enforce_eager=True, + ) + return parser.parse_args() + + +def main(args: Namespace): + pooling_params = PoolingParams( + late_chunking_params=LateChunkingParams(chunk_size=args.chunk_size) + ) + engine_args = vars(args).copy() + del engine_args["chunk_size"] + llm = LLM(**engine_args) + + # E5 document inputs use the "passage: " prefix; it also counts toward chunks. + documents = [ + "passage: Berlin is the capital of Germany. The city has many museums. " + "They display art and historical collections from around the world.", + "passage: Kyoto is a city in Japan. Its historic temples attract visitors " + "throughout the year.", + ] + outputs = llm.encode( + documents, + pooling_task="token_embed", + pooling_params=pooling_params, + ) + + for doc_index, (document, output) in enumerate(zip(documents, outputs)): + vectors = output.outputs.data + metadata = output.late_chunking + assert metadata is not None + print(f"\nDocument {doc_index}: {metadata.input_tokens} input tokens") + print(f"Chunk vectors: {tuple(vectors.shape)}") + for index, chunk in enumerate(metadata.chunks): + # vectors[index] corresponds to this chunk and its source-text range. + if chunk.char_range is None: + excerpt = "" + else: + start, end = chunk.char_range + excerpt = document[start:end] + print( + f" Chunk {index}: tokens={chunk.token_range}, chars={chunk.char_range}" + ) + print(f" {excerpt!r}") + + +if __name__ == "__main__": + main(parse_args()) diff --git a/tests/entrypoints/pooling/test_late_chunking.py b/tests/entrypoints/pooling/test_late_chunking.py new file mode 100644 index 000000000000..eaec62ac196c --- /dev/null +++ b/tests/entrypoints/pooling/test_late_chunking.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest import mock +from unittest.mock import Mock + +import pytest +import torch + +from vllm import PoolingParams +from vllm.config import PoolerConfig +from vllm.entrypoints.pooling.embed.io_processor import ( + JinaRankingTokenEmbedIOProcessor, + TokenEmbedIOProcessor, +) +from vllm.entrypoints.pooling.late_chunking import build_late_chunking_metadata +from vllm.entrypoints.pooling.offline import PoolingOfflineMixin +from vllm.entrypoints.pooling.typing import ( + OfflineEncodeInputsContext, + OfflineOutputsContext, +) +from vllm.exceptions import VLLMValidationError +from vllm.outputs import ( + LateChunk, + PoolingOutput, + PoolingRequestOutput, + RequestError, +) +from vllm.pooling_params import LateChunkingParams +from vllm.renderers import TokenizeParams + + +@pytest.fixture +def late_chunk_processor(): + processor = object.__new__(TokenEmbedIOProcessor) + processor.model_config = SimpleNamespace( + is_encoder_decoder=False, + architecture="NomicBertModel", + model_impl="auto", + is_matryoshka=False, + hf_config=SimpleNamespace(), + pooler_config=PoolerConfig(seq_pooling_type="MEAN", tok_pooling_type="ALL"), + ) + processor.vllm_config = SimpleNamespace( + model_config=processor.model_config, + cache_config=SimpleNamespace(enable_prefix_caching=False), + scheduler_config=SimpleNamespace(enable_chunked_prefill=False), + lora_config=None, + ) + processor.renderer = Mock() + processor.renderer.default_cmpl_tok_params = TokenizeParams(max_total_tokens=32) + processor.renderer.render_cmpl.side_effect = lambda **_: [ + { + "type": "token", + "prompt_token_ids": [101, 1, 2, 102], + "prompt_token_offsets": [(0, 0), (0, 1), (2, 3), (0, 0)], + } + ] + return processor + + +def _late_chunk_context(prompts="a b", **kwargs): + return OfflineEncodeInputsContext( + pooling_task="token_embed", + tokenization_kwargs=kwargs or None, + lora_request=None, + priorities=None, + prompts=prompts, + pooling_params=PoolingParams( + late_chunking_params=LateChunkingParams(chunk_size=2) + ), + ) + + +def test_late_chunk_render_requests_offsets_once(late_chunk_processor): + processor = late_chunk_processor + ctx = _late_chunk_context() + factory, count = processor.get_request_factory_offline(ctx) + render_params = next(factory()) + result = processor.render(render_params) + assert count == 1 + processor.renderer.render_cmpl.assert_called_once() + assert processor.renderer.render_cmpl.call_args.kwargs[ + "tok_params" + ].return_token_offsets + assert not render_params["tok_params"].return_token_offsets + assert "prompt_token_offsets" not in result["prompts"] + assert "late_chunking" not in result + assert [c.char_range for c in ctx.late_chunking[0].metadata.chunks] == [ + (0, 1), + (2, 3), + ] + assert ctx.late_chunking[0].metadata.input_tokens == 4 + + +@pytest.mark.parametrize("as_sequence", [False, True]) +@pytest.mark.parametrize("chunk_size", [None, 2]) +def test_jina_ranking_rejects_chunking_before_tokenization( + late_chunk_processor, as_sequence, chunk_size +): + processor = object.__new__(JinaRankingTokenEmbedIOProcessor) + processor.__dict__.update(late_chunk_processor.__dict__) + ctx = _late_chunk_context(prompts=["document", "query"]) + params = PoolingParams( + late_chunking_params=LateChunkingParams(chunk_size) if chunk_size else None + ) + ctx.pooling_params = [params] if as_sequence else params + factory, count = processor.get_request_factory_offline(ctx) + assert count == 1 + request = next(factory()) + if chunk_size is not None: + with pytest.raises(VLLMValidationError, match="JinaForRanking.*late chunking"): + processor.render(request) + late_chunk_processor.renderer.render_cmpl.assert_not_called() + else: + result = processor.render(request) + assert result["prompts"]["type"] == "token" + assert result["prompts"]["prompt_token_ids"] == [101, 1, 2, 102] + late_chunk_processor.renderer.render_cmpl.assert_called_once() + assert not request["tok_params"].return_token_offsets + + +@pytest.mark.parametrize( + "prompts, kwargs, error", + [ + pytest.param("", {}, "plain-text", id="empty"), + pytest.param([1, 2], {}, "plain-text", id="token-ids"), + pytest.param( + {"prompt": "abc", "multi_modal_data": {}}, + {}, + "plain-text", + id="multimodal", + ), + ] + + [ + pytest.param("abc", {key: value}, "does not support", id=key) + for key, value in ( + ("truncate_prompt_tokens", 2), + ("truncation", True), + ("pad_prompt_tokens", 8), + ("padding", True), + ("do_lower_case", True), + ) + ], +) +def test_late_chunking_rejects_unsupported_input_before_render( + late_chunk_processor, prompts, kwargs, error +): + processor = late_chunk_processor + with pytest.raises(VLLMValidationError, match=error): + factory, _ = processor.get_request_factory_offline( + _late_chunk_context(prompts, **kwargs) + ) + processor.render(next(factory())) + processor.renderer.render_cmpl.assert_not_called() + + +def test_late_chunk_ranges_keep_unicode_overlaps_and_special_only_chunks(): + text = "中 😀 e\u0301" + offsets = [(0, 0), (0, 1), (2, 3), (2, 3), (4, 6), (0, 0)] + metadata = build_late_chunking_metadata(text, len(offsets), offsets, 1) + assert [c.char_range for c in metadata.chunks] == [ + None, + (0, 1), + (2, 3), + (2, 3), + (4, 6), + None, + ] + assert [ + text[slice(*c.char_range)] if c.char_range else None for c in metadata.chunks + ] == [None, "中", "😀", "😀", "e\u0301", None] + chunks = build_late_chunking_metadata(text, len(offsets), offsets, 4).chunks + assert [(c.token_range, c.char_range) for c in chunks] == [ + ((0, 4), (0, 3)), + ((4, 6), (4, 6)), + ] + + +@pytest.mark.parametrize( + "offsets", [None, [(0, 1)], [(-1, 1), (0, 0)], [(0, 4), (0, 0)], [(2, 3), (0, 1)]] +) +def test_late_chunk_ranges_reject_missing_or_invalid_offsets(offsets): + with pytest.raises(VLLMValidationError, match="offsets"): + build_late_chunking_metadata("abc", 2, offsets, 2) + + +def _mock_chunk_llm(processor, outputs): + llm = mock.Mock(spec=PoolingOfflineMixin) + llm.encode = PoolingOfflineMixin.encode.__get__(llm) + llm._run_tiling_engine = PoolingOfflineMixin._run_tiling_engine.__get__(llm) + llm.pooling_io_processors = {"token_embed": processor} + llm._executor = SimpleNamespace(map=map) + llm.llm_engine = mock.Mock() + llm.llm_engine.vllm_config.scheduler_config.max_num_seqs = 2 + llm.llm_engine.has_unfinished_requests.return_value = False + llm.llm_engine.step.side_effect = outputs + llm._render_and_add_requests = mock.Mock( + side_effect=lambda **kw: [ + str(i) + "-internal" for i in range(len(kw["params"])) + ] + ) + return llm + + +def _chunk_output(request_id, rows=2, **kwargs): + return PoolingRequestOutput( + str(request_id), + PoolingOutput(torch.ones(rows, 4)), + [101, 1, 2, 102], + 0, + True, + **kwargs, + ) + + +def test_late_chunk_mapping_preserves_order_mixed_requests_and_errors( + late_chunk_processor, +): + error = RequestError("test_error", "original error") + failed = _chunk_output(3, rows=0, error=error) + llm = _mock_chunk_llm(late_chunk_processor, RuntimeError("step failed")) + params = [ + PoolingParams(late_chunking_params=LateChunkingParams(2)), + PoolingParams(), + PoolingParams(late_chunking_params=LateChunkingParams(1)), + PoolingParams(late_chunking_params=LateChunkingParams(2)), + ] + # A failed encode call must not leave metadata for the following batch. + with pytest.raises(RuntimeError, match="step failed"): + llm.encode( + ["a b"] * 2, + pooling_task="token_embed", + pooling_params=params[:2], + use_tqdm=False, + ) + llm.llm_engine.step.side_effect = [ + [ + _chunk_output(2, rows=4), + failed, + _chunk_output(1, rows=4), + _chunk_output(0), + ] + ] + outputs = llm.encode( + ["a b"] * 4, pooling_task="token_embed", pooling_params=params, use_tqdm=False + ) + assert outputs[0].late_chunking.chunks == [ + LateChunk((0, 2), (0, 1)), + LateChunk((2, 4), (2, 3)), + ] + assert outputs[1].late_chunking is None + assert len(outputs[2].late_chunking.chunks) == 4 + assert outputs[3].late_chunking is None + assert outputs[3].error is error + + +def test_late_chunk_mapping_rejects_successful_output_with_wrong_row_count( + late_chunk_processor, +): + llm = _mock_chunk_llm(late_chunk_processor, [[_chunk_output(0, rows=0)]]) + with pytest.raises(ValueError, match="does not match"): + llm.encode( + "a b", + pooling_task="token_embed", + pooling_params=PoolingParams(late_chunking_params=LateChunkingParams(2)), + use_tqdm=False, + ) + # All engine requests have finished before post-processing validates the output. + llm.llm_engine.abort_request.assert_not_called() + + +def test_late_chunk_contexts_remain_isolated_when_rendering_interleaves( + late_chunk_processor, +): + processor = late_chunk_processor + first, second = _late_chunk_context("xy z"), _late_chunk_context() + shared = PoolingParams(late_chunking_params=LateChunkingParams(2)) + first.pooling_params = second.pooling_params = shared + first_factory, _ = processor.get_request_factory_offline(first) + second_factory, _ = processor.get_request_factory_offline(second) + first_request, second_request = next(first_factory()), next(second_factory()) + processor.render(second_request) + assert first.late_chunking[0].metadata is None + processor.renderer.render_cmpl.side_effect = lambda **_: [ + { + "type": "token", + "prompt_token_ids": [101, 1, 2, 102], + "prompt_token_offsets": [(0, 0), (0, 2), (3, 4), (0, 0)], + } + ] + processor.render(first_request) + first_output = processor.post_process_offline( + OfflineOutputsContext([_chunk_output(0)], late_chunking=first.late_chunking) + )[0] + second_output = processor.post_process_offline( + OfflineOutputsContext([_chunk_output(0)], late_chunking=second.late_chunking) + )[0] + assert first_output.late_chunking.chunks[0].char_range == (0, 2) + assert second_output.late_chunking.chunks[0].char_range == (0, 1) + + +def test_plain_token_embed_skips_chunk_metadata(late_chunk_processor): + ctx = _late_chunk_context() + ctx.pooling_params = PoolingParams() + factory, _ = late_chunk_processor.get_request_factory_offline(ctx) + request = next(factory()) + result = late_chunk_processor.render(request) + assert not ctx.late_chunking + assert result["params"].late_chunking_params is None + assert not late_chunk_processor.renderer.render_cmpl.call_args.kwargs[ + "tok_params" + ].return_token_offsets diff --git a/tests/entrypoints/pooling/token_embed/test_offline.py b/tests/entrypoints/pooling/token_embed/test_offline.py index e0bda6542284..b65ef6ff8bed 100644 --- a/tests/entrypoints/pooling/token_embed/test_offline.py +++ b/tests/entrypoints/pooling/token_embed/test_offline.py @@ -3,9 +3,13 @@ import weakref import pytest +import torch -from vllm import LLM, PoolingRequestOutput +from tests.models.utils import check_embeddings_close +from vllm import LLM, PoolingParams, PoolingRequestOutput from vllm.config import PoolerConfig +from vllm.exceptions import VLLMValidationError +from vllm.pooling_params import LateChunkingParams from vllm.tasks import PoolingTask MODEL_NAME = "intfloat/multilingual-e5-small" @@ -61,3 +65,218 @@ def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm): with pytest.raises(ValueError, match=err_msg): llm.encode(prompt, pooling_task=task, use_tqdm=False) + + +def test_bge_m3_late_chunking_keeps_special_tokens(vllm_runner): + with vllm_runner( + "BAAI/bge-m3", + hf_overrides={"architectures": ["BgeM3EmbeddingModel"]}, + pooler_config=PoolerConfig(task="token_embed"), + dtype="float32", + max_model_len=128, + # Keep encoder batch size fixed for the numerical reference. + max_num_seqs=1, + gpu_memory_utilization=0.5, + enforce_eager=True, + ) as runner: + llm = runner.llm + text = "A document about Berlin and its museums." + sizes = (1, 3, 128, None) + tokens, *outputs = llm.encode( + [text] * len(sizes), + pooling_task="token_embed", + pooling_params=[ + PoolingParams( + use_activation=size not in (1, None), + late_chunking_params=LateChunkingParams(size) if size else None, + ) + for size in sizes + ], + use_tqdm=False, + ) + assert tokens.outputs.data.shape[0] == len(tokens.prompt_token_ids) + assert tokens.late_chunking.chunks[0].char_range is None + for size, output in zip(sizes[1:], outputs): + if size is None: + # Ordinary BGE-M3 token embeddings still drop BOS only. + expected = tokens.outputs.data[1:] + assert output.late_chunking is None + else: + # BGE-M3's projector is linear, so its unnormalized token + # outputs also provide a reference for projected chunk means. + expected = torch.nn.functional.normalize( + torch.stack([p.mean(0) for p in tokens.outputs.data.split(size)]), + dim=-1, + ) + assert len(output.late_chunking.chunks) == len(expected) + torch.testing.assert_close( + output.outputs.data, expected, atol=2e-5, rtol=2e-4 + ) + + +def test_nomic_late_chunking_without_special_tokens(vllm_runner): + """Keep source ranges without special tokens and reject empty tokenization.""" + with vllm_runner( + "nomic-ai/nomic-embed-text-v1", + revision="720244025c1a7e15661a174c63cce63c8218e52b", + trust_remote_code=True, + runner="pooling", + pooler_config=PoolerConfig(task="token_embed"), + dtype="float32", + max_model_len=128, + max_num_seqs=2, + gpu_memory_utilization=0.3, + enforce_eager=True, + enable_prefix_caching=False, + enable_chunked_prefill=False, + ) as runner: + llm = runner.llm + text = "search_document: 中文 😀 café e\u0301" + params = PoolingParams( + late_chunking_params=LateChunkingParams(3), use_activation=False + ) + kwargs = {"add_special_tokens": False} + raw = llm.encode( + text, + pooling_task="token_embed", + use_tqdm=False, + pooling_params=PoolingParams(use_activation=False), + tokenization_kwargs=kwargs, + )[0] + actual = llm.encode( + text, + pooling_task="token_embed", + use_tqdm=False, + pooling_params=params, + tokenization_kwargs=kwargs, + )[0] + expected = torch.stack([part.mean(0) for part in raw.outputs.data.split(3)]) + torch.testing.assert_close(actual.outputs.data, expected, atol=1e-5, rtol=1e-4) + encoded = llm.get_tokenizer()(text, return_offsets_mapping=True, **kwargs) + assert actual.prompt_token_ids == encoded["input_ids"] + for chunk in actual.late_chunking.chunks: + start, end = chunk.token_range + offsets = [(a, b) for a, b in encoded["offset_mapping"][start:end] if a < b] + assert chunk.char_range == (offsets[0][0], offsets[-1][1]) + for rejected in ("long " * 256, " "): + with pytest.raises(VLLMValidationError): + llm.encode( + rejected, + pooling_task="token_embed", + use_tqdm=False, + pooling_params=params, + tokenization_kwargs=kwargs, + ) + assert ( + llm.encode(text, pooling_task="token_embed", use_tqdm=False)[ + 0 + ].late_chunking + is None + ) + + +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize("normalize", [False, True]) +def test_e5_late_chunking_mixed_requests(llm: LLM, normalize: bool): + texts = [prompt, "Short text.", "中文 😀 café"] + raw = llm.encode( + texts, + pooling_task="token_embed", + use_tqdm=False, + pooling_params=PoolingParams(use_activation=False), + ) + sizes = [3, None, 100] + params = [ + PoolingParams( + late_chunking_params=LateChunkingParams(size) if size is not None else None, + use_activation=normalize, + ) + for size in sizes + ] + actual = llm.encode( + texts, pooling_task="token_embed", pooling_params=params, use_tqdm=False + ) + for original, result, size in zip(raw, actual, sizes): + expected = original.outputs.data + if size is not None: + expected = torch.stack( + [part.float().mean(0) for part in expected.split(size)] + ) + assert result.late_chunking is not None + assert len(result.late_chunking.chunks) == len(expected) + assert result.late_chunking.input_tokens == len(result.prompt_token_ids) + else: + assert result.late_chunking is None + if normalize: + expected = torch.nn.functional.normalize(expected, dim=-1) + check_embeddings_close( + embeddings_0_lst=result.outputs.data.tolist(), + embeddings_1_lst=expected.tolist(), + name_0="late_chunking", + name_1="token_reference", + ) + torch.testing.assert_close( + result.outputs.data.norm(dim=-1), + expected.norm(dim=-1), + atol=1e-3, + rtol=1e-3, + ) + + +@pytest.mark.parametrize("use_v2_runner", [False, True]) +def test_late_chunking_recomputes_cached_prefix_with_chunked_prefill( + vllm_runner, monkeypatch, use_v2_runner +): + # Dummy weights suffice for testing scheduling and reduction of the same states. + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_v2_runner else "0") + with vllm_runner( + "openai-community/gpt2", + runner="pooling", + convert="embed", + load_format="dummy", + pooler_config=PoolerConfig( + task="token_embed", seq_pooling_type="LAST", tok_pooling_type="ALL" + ), + dtype="float32", + max_model_len=128, + max_num_batched_tokens=16, + max_num_seqs=2, + enable_prefix_caching=True, + enable_chunked_prefill=True, + enforce_eager=True, + gpu_memory_utilization=0.3, + seed=0, + ) as runner: + text = "The document has several important words. " * 6 + llm = runner.llm + raw = llm.encode( + text, + pooling_task="token_embed", + use_tqdm=False, + pooling_params=PoolingParams(use_activation=False), + )[0] + assert len(raw.prompt_token_ids) > 16 + cached = llm.encode( + text, + pooling_task="token_embed", + use_tqdm=False, + pooling_params=PoolingParams( + use_activation=False, skip_reading_prefix_cache=False + ), + )[0] + assert cached.num_cached_tokens > 0 + params = PoolingParams( + late_chunking_params=LateChunkingParams(7), + use_activation=False, + skip_reading_prefix_cache=False, + ) + actual = llm.encode( + text, pooling_task="token_embed", pooling_params=params, use_tqdm=False + )[0] + assert actual.num_cached_tokens == 0 + assert actual.prompt_token_ids == raw.prompt_token_ids + expected = torch.stack([part.mean(0) for part in raw.outputs.data.split(7)]) + torch.testing.assert_close(actual.outputs.data, expected, atol=1e-5, rtol=1e-4) + assert actual.late_chunking.input_tokens == len(raw.prompt_token_ids) + assert params.skip_reading_prefix_cache is False + assert params.late_chunking_params is not None diff --git a/tests/model_executor/layers/test_pooler_heads.py b/tests/model_executor/layers/test_pooler_heads.py index a80c5b8344f5..00091ec0c292 100644 --- a/tests/model_executor/layers/test_pooler_heads.py +++ b/tests/model_executor/layers/test_pooler_heads.py @@ -308,6 +308,7 @@ def test_passthrough(self): x = torch.randn(5, _HIDDEN) param = PoolingParams(task="token_embed") out = head.forward_chunk(x, param) + assert out is x assert torch.equal(out, x) def test_none_chunked_prefill(self): diff --git a/tests/model_executor/layers/test_pooler_methods.py b/tests/model_executor/layers/test_pooler_methods.py index 75bbed4f37fe..d6b03dee1954 100644 --- a/tests/model_executor/layers/test_pooler_methods.py +++ b/tests/model_executor/layers/test_pooler_methods.py @@ -13,19 +13,22 @@ from vllm.config import PoolerConfig, set_current_vllm_config from vllm.model_executor.layers.pooler import PoolingParamsUpdate +from vllm.model_executor.layers.pooler.activations import PoolerNormalize from vllm.model_executor.layers.pooler.seqwise.methods import ( CLSPool, LastPool, MeanPool, get_seq_pooling_method, ) -from vllm.model_executor.layers.pooler.special import DispatchPooler +from vllm.model_executor.layers.pooler.special import BOSEOSFilter, DispatchPooler +from vllm.model_executor.layers.pooler.tokwise.heads import TokenEmbeddingPoolerHead from vllm.model_executor.layers.pooler.tokwise.methods import ( AllPool, StepPool, get_tok_pooling_method, ) -from vllm.pooling_params import PoolingParams +from vllm.model_executor.layers.pooler.tokwise.poolers import TokenPooler +from vllm.pooling_params import LateChunkingParams, PoolingParams from vllm.tasks import PoolingTask from vllm.v1.pool.metadata import PoolingCursor, PoolingMetadata, PoolingStates @@ -606,3 +609,139 @@ def test_step(self): def test_unknown_raises(self): with pytest.raises(NotImplementedError, match="UNKNOWN"): get_tok_pooling_method("UNKNOWN") + + +@pytest.mark.parametrize( + ("chunk_size", "dtype", "head_dtype"), + [ + (1, torch.float32, None), + (2, torch.float16, None), + (2, torch.bfloat16, torch.float32), + (7, torch.float32, None), + (10, torch.float32, None), + (2**63 - 1, torch.float32, None), + ], +) +def test_late_chunk_pool_matches_fp32_means_before_projection_and_normalization( + chunk_size, dtype, head_dtype +): + # Unequal norms expose the incorrect normalize-then-average implementation. + hidden = (torch.arange(28).reshape(7, 4) - 10).to(dtype) + projector = torch.nn.Linear(4, 3, bias=True).to(head_dtype or dtype) + head = TokenEmbeddingPoolerHead(head_dtype, projector, PoolerNormalize()) + pooler = TokenPooler(TestAllPool._make_all_pool(), head) + params = PoolingParams( + task="token_embed", + late_chunking_params=LateChunkingParams(chunk_size=chunk_size), + use_activation=True, + ) + metadata = _make_metadata([7], pooling_params=[params]) + actual = pooler(hidden, metadata)[0] + means = torch.stack( + [hidden[a : a + chunk_size].float().mean(0) for a in range(0, 7, chunk_size)] + ).to(head_dtype or dtype) + expected = torch.nn.functional.normalize(projector(means), p=2, dim=-1) + torch.testing.assert_close(actual, expected) + assert actual.dtype == (head_dtype or dtype) + + +def test_late_chunk_pool_reduces_before_nonlinear_projector(): + hidden = torch.tensor([[2.0, 0.0], [0.0, 1.0]]) + projector = torch.nn.Tanh() + head = TokenEmbeddingPoolerHead(projector=projector, activation=PoolerNormalize()) + params = PoolingParams( + task="token_embed", + late_chunking_params=LateChunkingParams(chunk_size=2), + use_activation=True, + ) + actual = head.forward_chunk(hidden, params) + expected = torch.nn.functional.normalize(projector(hidden.mean(0, keepdim=True))) + wrong_order = torch.nn.functional.normalize(projector(hidden).mean(0, keepdim=True)) + torch.testing.assert_close(actual, expected) + assert not torch.allclose(actual, wrong_order) + + +def test_late_chunk_pool_isolates_mixed_requests_and_owns_finished_storage(): + hidden = torch.arange(36, dtype=torch.float32).reshape(9, 4) + params = [ + PoolingParams( + task="token_embed", late_chunking_params=LateChunkingParams(chunk_size=2) + ), + PoolingParams(task="token_embed"), + PoolingParams( + task="token_embed", late_chunking_params=LateChunkingParams(chunk_size=1) + ), + ] + pooler = TokenPooler( + TestAllPool._make_all_pool(async_scheduling=True), TokenEmbeddingPoolerHead() + ) + outputs = pooler(hidden, _make_metadata([3, 4, 2], pooling_params=params)) + expected = [ + torch.stack([hidden[:2].mean(0), hidden[2]]), + hidden[3:7].clone(), + hidden[7:].clone(), + ] + hidden.fill_(-999) + for actual, reference in zip(outputs, expected): + torch.testing.assert_close(actual, reference) + + +@pytest.mark.parametrize("chunk_size", [1, 2, 8]) +@pytest.mark.parametrize("embedding_size", [1, 4]) +def test_bos_eos_filter_preserves_chunk_vectors_in_mixed_batch( + chunk_size, embedding_size +): + hidden = torch.arange(10 * embedding_size, dtype=torch.float32).reshape( + 10, embedding_size + ) + pooler = BOSEOSFilter( + TokenPooler(TestAllPool._make_all_pool(), TokenEmbeddingPoolerHead()), + bos_token_id=101, + eos_token_id=102, + ) + metadata = _make_metadata( + [5, 5], + token_ids=[[101, 1, 2, 3, 102]] * 2, + pooling_params=[ + PoolingParams( + task="token_embed", late_chunking_params=LateChunkingParams(chunk_size) + ), + PoolingParams(task="token_embed"), + ], + ) + chunked, ordinary = pooler(hidden, metadata) + expected = torch.stack([part.mean(0) for part in hidden[:5].split(chunk_size)]) + torch.testing.assert_close(chunked, expected) + torch.testing.assert_close(ordinary, hidden[6:9].squeeze(-1)) + + +def test_bos_eos_filter_waits_for_complete_chunk_vectors(): + pooler = BOSEOSFilter( + TokenPooler( + TestAllPool._make_all_pool(chunked=True), TokenEmbeddingPoolerHead() + ), + bos_token_id=101, + ) + params = [ + PoolingParams(task="token_embed", late_chunking_params=LateChunkingParams(3)) + ] + first = _make_metadata( + [4], + token_ids=[[101, 1, 2, 102]], + pooling_params=params, + num_scheduled_tokens=[2], + seq_lens=[2], + ) + hidden = torch.arange(8, dtype=torch.float32).reshape(4, 2) + assert pooler(hidden[:2], first) == [None] + last = _make_metadata( + [4], + token_ids=[[101, 1, 2, 102]], + pooling_params=params, + num_scheduled_tokens=[2], + seq_lens=[4], + ) + last.pooling_states = first.pooling_states + actual = pooler(hidden[2:], last)[0] + torch.testing.assert_close(actual, torch.stack([hidden[:3].mean(0), hidden[3]])) + assert not last.pooling_states[0].hidden_states_cache diff --git a/tests/test_pooling_params.py b/tests/test_pooling_params.py index f34270d0da52..81af02b502ab 100644 --- a/tests/test_pooling_params.py +++ b/tests/test_pooling_params.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass +from types import SimpleNamespace from typing import Any +import msgspec import pytest from pydantic import TypeAdapter, ValidationError @@ -13,6 +15,7 @@ from vllm.entrypoints.pooling.embed.protocol import EmbeddingRequest from vllm.entrypoints.pooling.pooling.protocol import PoolingRequest from vllm.exceptions import VLLMValidationError +from vllm.pooling_params import LateChunkingParams EMBEDDING_MODELS = [ EmbedModelInfo("intfloat/multilingual-e5-small", is_matryoshka=False), @@ -211,3 +214,92 @@ def test_token_classify(pooling_type: str): with pytest.raises(VLLMValidationError): pooling_params = PoolingParams(task=task, **{p: True}) pooling_params.verify(model_config) + + +@pytest.mark.parametrize("value", [0, -1, True, False, 1.5, "2"]) +def test_late_chunk_size_rejects_non_positive_integers(value): + with pytest.raises(VLLMValidationError, match="positive integer"): + PoolingParams(late_chunking_params=LateChunkingParams(chunk_size=value)) + + +def _late_chunking_model_config(): + return SimpleNamespace( + architecture="NomicBertModel", + model_impl="auto", + is_matryoshka=False, + hf_config=SimpleNamespace(), + pooler_config=PoolerConfig(seq_pooling_type="MEAN", tok_pooling_type="ALL"), + ) + + +def test_late_chunking_params_preserve_defaults_clone_and_wire_format(): + model_config = _late_chunking_model_config() + params = PoolingParams( + task="token_embed", late_chunking_params=LateChunkingParams(chunk_size=3) + ) + params.verify(model_config) + assert params.skip_reading_prefix_cache is True + assert params.use_activation is True + # Chunk means must include cached prefixes even when the user opts into reads. + explicit_cache_flag = PoolingParams( + task="token_embed", + late_chunking_params=LateChunkingParams(chunk_size=3), + skip_reading_prefix_cache=False, + ) + explicit_cache_flag.verify(model_config) + assert explicit_cache_flag.skip_reading_prefix_cache is True + clone = params.clone() + assert clone.late_chunking_params is not None + clone.late_chunking_params.chunk_size = 7 + assert params.late_chunking_params is not None + assert params.late_chunking_params.chunk_size == 3 + # Workers receive only the fixed chunk size, never source ranges. + assert msgspec.msgpack.decode( + msgspec.msgpack.encode(params.late_chunking_params) + ) == [3] + assert ( + msgspec.msgpack.decode(msgspec.msgpack.encode(params), type=PoolingParams) + == params + ) + # A message written before the appended field still uses the old default. + wire = msgspec.msgpack.decode(msgspec.msgpack.encode(PoolingParams())) + assert msgspec.convert(wire[:-1], type=PoolingParams).late_chunking_params is None + + +@pytest.mark.parametrize( + "task", ["embed", "classify", "token_classify", "plugin", None] +) +def test_late_chunking_rejects_other_tasks(task): + with pytest.raises(VLLMValidationError, match="requires token_embed"): + PoolingParams( + task=task, late_chunking_params=LateChunkingParams(chunk_size=2) + ).verify(_late_chunking_model_config()) + + +def test_late_chunking_uses_existing_dimension_validation(): + config = MockMatryoshkaModelConfig( + pooler_config=PoolerConfig(tok_pooling_type="ALL", use_activation=False), + ) + params = PoolingParams( + task="token_embed", late_chunking_params=LateChunkingParams(2), dimensions=4 + ) + params.verify(config) + assert params.dimensions == 4 + assert params.use_activation is False + params.dimensions = config.embedding_size + 1 + with pytest.raises(VLLMValidationError, match="dimensions in range"): + params.verify(config) + + +@pytest.mark.parametrize("from_config", [False, True]) +def test_late_chunking_rejects_token_filtering_after_resolving_defaults(from_config): + config = MockModelConfig(pooler_config=PoolerConfig(tok_pooling_type="STEP")) + params = PoolingParams( + task="token_embed", late_chunking_params=LateChunkingParams(2) + ) + if from_config: + config.pooler_config.step_tag_id = 42 + else: + params.step_tag_id = 42 + with pytest.raises(VLLMValidationError, match="all token rows"): + params.verify(config) diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index 94f84f09b678..b37a009c925a 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -12,6 +12,7 @@ from openai.types.chat.chat_completion_content_part_image_param import ImageURL from vllm import PoolingParams +from vllm.exceptions import VLLMValidationError from vllm.inputs import tokens_input from vllm.logger import init_logger from vllm.outputs import PoolingOutput, PoolingRequestOutput @@ -26,18 +27,26 @@ from vllm.utils.mistral import is_mistral_tokenizer from ..base.io_processor import PoolingIOProcessor +from ..late_chunking import ( + attach_late_chunking_metadata, + build_late_chunking_metadata, + prepare_late_chunking_input, +) from ..scoring.io_processor import JinaRankingIOProcessorMixin from ..typing import ( AnyOfflineInputsContext, AnyRenderParam, ChunkedEmbeddingMetadata, EncodeChatRenderParams, + LateChunkingSlot, OfflineEncodeInputsContext, + OfflineOutputsContext, PoolingChatLikeRequest, PoolingCompletionLikeRequest, PoolingEngineInput, PoolingServeContext, RequestFactory, + RequestGenerator, ) from .protocol import ( CohereEmbedContent, @@ -657,10 +666,75 @@ def _enforce_cohere_max_tokens(self, ctx: PoolingServeContext) -> None: class TokenEmbedIOProcessor(PoolingIOProcessor): name = "token_embed" + def get_request_factory_offline( + self, ctx: AnyOfflineInputsContext + ) -> tuple[RequestFactory, int]: + factory, num_requests = super().get_request_factory_offline(ctx) + assert isinstance(ctx, OfflineEncodeInputsContext) + if ctx.pooling_params is None or not any( + p.late_chunking_params is not None + for p in self._params_to_seq(ctx.pooling_params, num_requests) + ): + return factory, num_requests + + if (ctx.tokenization_kwargs or {}).get("padding") not in ( + None, + False, + "do_not_pad", + ): + raise VLLMValidationError("Late chunking does not support input padding") + + def request_factory() -> RequestGenerator: + for request in factory(): + slot = None + if request["params"].late_chunking_params is not None: + slot = LateChunkingSlot() + request["late_chunking"] = slot + ctx.late_chunking.append(slot) + yield request + + return request_factory, num_requests + + def render(self, render_params: AnyRenderParam) -> PoolingEngineInput: + late_chunking = render_params["params"].late_chunking_params + if late_chunking is None: + return super().render(render_params) + + text, tok_params = prepare_late_chunking_input(self.vllm_config, render_params) + render_params = render_params.copy() + render_params["tok_params"] = tok_params + result = super().render(render_params) + prompt = result["prompts"] + if prompt["type"] != "token": + raise VLLMValidationError("Late chunking requires tokenized text") + render_params["late_chunking"].metadata = build_late_chunking_metadata( + text, + len(prompt["prompt_token_ids"]), + prompt.pop("prompt_token_offsets", None), + late_chunking.chunk_size, + ) + return result + + def post_process_offline(self, ctx: OfflineOutputsContext): + if ctx.late_chunking: + for output, slot in zip(ctx.outputs, ctx.late_chunking, strict=True): + if slot is not None: + assert slot.metadata is not None + attach_late_chunking_metadata(output, slot.metadata) + return ctx.outputs + class JinaRankingTokenEmbedIOProcessor( TokenEmbedIOProcessor, JinaRankingIOProcessorMixin ): + def render(self, render_params: AnyRenderParam) -> PoolingEngineInput: + if render_params["params"].late_chunking_params is not None: + raise VLLMValidationError( + "JinaForRanking returns document/query vectors and does not support " + "late chunking" + ) + return super().render(render_params) + def get_request_factory_online( self, ctx: PoolingServeContext ) -> Sequence[AnyRenderParam]: diff --git a/vllm/entrypoints/pooling/late_chunking.py b/vllm/entrypoints/pooling/late_chunking.py new file mode 100644 index 000000000000..d848c3257f20 --- /dev/null +++ b/vllm/entrypoints/pooling/late_chunking.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Sequence +from dataclasses import replace +from typing import cast + +from vllm.config import VllmConfig +from vllm.exceptions import VLLMValidationError +from vllm.outputs import LateChunk, LateChunkingMetadata, PoolingRequestOutput +from vllm.renderers import TokenizeParams + +from .typing import AnyRenderParam, EncodeCMPLRenderParams + + +def prepare_late_chunking_input( + config: VllmConfig, render_params: AnyRenderParam +) -> tuple[str, TokenizeParams]: + """Validate full-document text input before requesting tokenizer offsets.""" + params = render_params["params"] + params.verify(config.model_config) + if "prompts" not in render_params: + raise VLLMValidationError("Late chunking requires plain-text input") + prompt = cast(EncodeCMPLRenderParams, render_params)["prompts"] + if ( + not isinstance(prompt, dict) + or set(prompt) - {"prompt", "cache_salt"} + or not isinstance(text := prompt.get("prompt"), str) + or not text + ): + raise VLLMValidationError("Late chunking requires nonempty plain-text input") + tok_params = render_params["tok_params"] + if ( + tok_params.pad_prompt_tokens is not None + or tok_params.truncate_prompt_tokens is not None + or tok_params.do_lower_case + ): + raise VLLMValidationError( + "Late chunking does not support input padding, truncation or " + "renderer text normalization" + ) + return text, replace(tok_params, return_token_offsets=True) + + +def build_late_chunking_metadata( + text: str, + num_tokens: int, + offsets: Sequence[tuple[int, int]] | None, + chunk_size: int, +) -> LateChunkingMetadata: + """Map chunks to original text, keeping special-token-only chunks.""" + if num_tokens <= 0 or offsets is None or len(offsets) != num_tokens: + raise VLLMValidationError( + "Late chunking requires a fast tokenizer with offsets aligned to " + "all input tokens" + ) + previous_start = previous_end = 0 + for start, end in offsets: + if not 0 <= start <= end <= len(text): + raise VLLMValidationError("Token offsets are not aligned to the input text") + if start < end: + if start < previous_start or end < previous_end: + raise VLLMValidationError("Token offsets are not in source-text order") + previous_start, previous_end = start, end + + chunks = [] + for start in range(0, num_tokens, chunk_size): + end = min(start + chunk_size, num_tokens) + source = [(a, b) for a, b in offsets[start:end] if a < b] + char_range = ( + (min(a for a, _ in source), max(b for _, b in source)) if source else None + ) + chunks.append(LateChunk(token_range=(start, end), char_range=char_range)) + return LateChunkingMetadata( + chunk_size=chunk_size, input_tokens=num_tokens, chunks=chunks + ) + + +def attach_late_chunking_metadata( + output: PoolingRequestOutput, metadata: LateChunkingMetadata +) -> None: + """Attach ranges only to successful, complete chunk-vector results.""" + if not output.finished or output.error is not None: + return + data = output.outputs.data + if ( + data.ndim != 2 + or data.shape[0] != len(metadata.chunks) + or len(output.prompt_token_ids) != metadata.input_tokens + ): + raise ValueError("Late-chunking output does not match its token/source ranges") + output.late_chunking = metadata diff --git a/vllm/entrypoints/pooling/offline.py b/vllm/entrypoints/pooling/offline.py index 19be1eddde91..edf92791a86a 100644 --- a/vllm/entrypoints/pooling/offline.py +++ b/vllm/entrypoints/pooling/offline.py @@ -129,7 +129,7 @@ def encode( io_processor, request_factory, num_requests, use_tqdm=use_tqdm ) outputs = io_processor.post_process_offline( - ctx=OfflineOutputsContext(outputs=outputs) + ctx=OfflineOutputsContext(outputs=outputs, late_chunking=ctx.late_chunking) ) return outputs diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index a299fc657d54..33aa65998978 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -3,7 +3,7 @@ import time from collections.abc import AsyncGenerator, Callable, Generator, Sequence from dataclasses import dataclass, field -from typing import Any, Generic, TypeAlias, TypedDict, TypeVar +from typing import Any, Generic, NotRequired, TypeAlias, TypedDict, TypeVar from fastapi import Request from pydantic import ConfigDict @@ -11,6 +11,7 @@ from vllm import PoolingParams, PoolingRequestOutput, PromptType from vllm.inputs import DataPrompt, EngineInput from vllm.lora.request import LoRARequest +from vllm.outputs import LateChunkingMetadata from vllm.renderers import ChatParams, TokenizeParams from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.renderers.inputs import DictPrompt @@ -80,6 +81,13 @@ class ChunkedEmbeddingMetadata: chunk_index: int +@dataclass +class LateChunkingSlot: + """Per-document source ranges shared by frontend rendering and output contexts.""" + + metadata: LateChunkingMetadata | None = None + + @dataclass(kw_only=True) class PoolingServeContext(Generic[PoolingRequestT]): model_config = ConfigDict(arbitrary_types_allowed=True) @@ -124,6 +132,9 @@ class OfflineInputsContext: tokenization_kwargs: dict[str, Any] | None lora_request: Sequence[LoRARequest | None] | None priorities: int | Sequence[int] | None + late_chunking: list[LateChunkingSlot | None] = field( + default_factory=list, kw_only=True + ) @dataclass @@ -158,6 +169,7 @@ class OfflineOutputsContext: ## for bi-encoder & late-interaction n_queries: int | None = None + late_chunking: list[LateChunkingSlot | None] = field(default_factory=list) class RenderParams(TypedDict): @@ -168,6 +180,7 @@ class RenderParams(TypedDict): params: PoolingParams lora_requests: LoRARequest | None priorities: int + late_chunking: NotRequired[LateChunkingSlot] class EncodeCMPLRenderParams(RenderParams): diff --git a/vllm/model_executor/layers/pooler/special.py b/vllm/model_executor/layers/pooler/special.py index 23354d66f3c4..f464f28300a7 100644 --- a/vllm/model_executor/layers/pooler/special.py +++ b/vllm/model_executor/layers/pooler/special.py @@ -198,6 +198,11 @@ def forward( for i, (prompt_len, token_ids) in enumerate( zip(pooling_metadata.prompt_lens, prompt_token_ids) ): + # Late chunking returns one row per chunk, with special tokens + # already included in the means. Skip BOS/EOS row removal for these + # requests; ordinary token outputs still use the filtering below. + if pooling_metadata.pooling_params[i].late_chunking_params is not None: + continue pooled_data = pooled_outputs[i] assert ( isinstance(pooled_data, torch.Tensor) diff --git a/vllm/model_executor/layers/pooler/tokwise/heads.py b/vllm/model_executor/layers/pooler/tokwise/heads.py index 78ec82a5384e..95242ee470a7 100644 --- a/vllm/model_executor/layers/pooler/tokwise/heads.py +++ b/vllm/model_executor/layers/pooler/tokwise/heads.py @@ -17,6 +17,23 @@ TokenPoolerHeadOutputItem: TypeAlias = torch.Tensor | None +def _mean_pool_chunks(data: torch.Tensor, chunk_size: int) -> torch.Tensor: + """Mean-pool complete contextual states, accumulating in FP32.""" + chunk_size = min(chunk_size, len(data)) + num_full, remainder = divmod(len(data), chunk_size) + full_end = num_full * chunk_size + full_chunks = data[:full_end].reshape(num_full, chunk_size, data.shape[-1]) + if not remainder: + return full_chunks.mean(dim=1, dtype=torch.float32) + + return torch.cat( + [ + full_chunks.mean(dim=1, dtype=torch.float32), + data[full_end:].mean(dim=0, keepdim=True, dtype=torch.float32), + ] + ) + + class TokenPoolerHead(nn.Module, ABC): @abstractmethod def get_supported_tasks(self) -> Set[PoolingTask]: @@ -80,6 +97,12 @@ def forward_chunk( if pooled_data is None: return None + if (late_chunking := pooling_param.late_chunking_params) is not None: + means = _mean_pool_chunks(pooled_data, late_chunking.chunk_size) + pooled_data = ( + means.to(pooled_data.dtype) if self.head_dtype is None else means + ) + if self.head_dtype is not None: pooled_data = pooled_data.to(self.head_dtype) # pooled_data shape: [n_tokens, hidden_size] diff --git a/vllm/outputs.py b/vllm/outputs.py index 145b7d510f71..887ebc5933de 100644 --- a/vllm/outputs.py +++ b/vllm/outputs.py @@ -277,6 +277,20 @@ def __repr__(self) -> str: _O = TypeVar("_O", default=PoolingOutput) +@dataclass(frozen=True) +class LateChunk: + token_range: tuple[int, int] + char_range: tuple[int, int] | None + + +@dataclass(frozen=True) +class LateChunkingMetadata: + chunk_size: int + input_tokens: int + chunks: list[LateChunk] + strategy: str = "fixed_tokens" + + class PoolingRequestOutput(Generic[_O]): """The output data of a pooling request to the LLM. @@ -287,6 +301,7 @@ class PoolingRequestOutput(Generic[_O]): num_cached_tokens: The number of tokens with prefix cache hit. finished (bool): A flag indicating whether the pooling is completed. error: Structured request-level error information, if the request failed. + late_chunking: Token and source ranges for fixed-length late chunking. """ @@ -299,8 +314,10 @@ def __init__( finished: bool, *, error: RequestError | None = None, + late_chunking: LateChunkingMetadata | None = None, ): self.request_id = request_id + self.late_chunking = late_chunking self.prompt_token_ids = prompt_token_ids self.num_cached_tokens = num_cached_tokens self.finished = finished diff --git a/vllm/pooling_params.py b/vllm/pooling_params.py index 319b77d5986a..45c34afbe3fd 100644 --- a/vllm/pooling_params.py +++ b/vllm/pooling_params.py @@ -36,6 +36,28 @@ class LateInteractionParams( query_uses: int | None = None +class LateChunkingParams( + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] + array_like=True, +): # type: ignore[call-arg] + """Fixed-length mean pooling before the token embedding head. + + Requests recompute all prompt tokens, bypassing prefix cache reads. + Special tokens count toward chunk boundaries and are included in chunk means, + including the BOS token normally filtered from BGE-M3 token embeddings. + """ + + chunk_size: int + + def __post_init__(self) -> None: + self.verify() + + def verify(self) -> None: + if type(self.chunk_size) is not int or self.chunk_size <= 0: + raise VLLMValidationError("chunk_size must be a positive integer") + + class PoolingParams( msgspec.Struct, omit_defaults=True, # type: ignore[call-arg] @@ -48,6 +70,8 @@ class PoolingParams( `None` uses the pooler's default, which is `True` in most cases. dimensions: Reduce the dimensions of embeddings if model support matryoshka representation. + late_chunking_params: Mean-pool contextual token states into fixed-size + chunks for the token embedding task. """ @@ -72,16 +96,22 @@ class PoolingParams( extra_kwargs: dict[str, Any] | None = None output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY + ## for token embedding models + # --8<-- [start:token-embed-pooling-params] + late_chunking_params: LateChunkingParams | None = None + """Mean-pool contextual tokens into fixed-size chunks for `token_embed`.""" + # --8<-- [end:token-embed-pooling-params] + @property def all_parameters(self) -> list[str]: - return ["dimensions", "use_activation"] + return ["dimensions", "use_activation", "late_chunking_params"] @property def valid_parameters(self): return { "embed": ["dimensions", "use_activation"], "classify": ["use_activation"], - "token_embed": ["dimensions", "use_activation"], + "token_embed": ["dimensions", "use_activation", "late_chunking_params"], "token_classify": ["use_activation"], } @@ -90,6 +120,15 @@ def clone(self) -> "PoolingParams": return deepcopy(self) def verify(self, model_config: ModelConfig) -> None: + if self.late_chunking_params is not None: + self.late_chunking_params.verify() + if self.task != "token_embed": + raise VLLMValidationError("Late chunking requires token_embed") + if self.late_interaction_params is not None: + raise VLLMValidationError( + "Late chunking cannot be combined with late-interaction scoring" + ) + # plugin task uses io_processor.parse_data to verify inputs, # skipping PoolingParams verify if self.task == "plugin": @@ -107,6 +146,11 @@ def verify(self, model_config: ModelConfig) -> None: self._merge_default_parameters(model_config) self._set_default_parameters(model_config) self._verify_valid_parameters() + if self.late_chunking_params is not None: + if self.step_tag_id is not None: + raise VLLMValidationError("Late chunking requires all token rows") + # KV cache hits omit hidden states needed for complete chunk means. + self.skip_reading_prefix_cache = True def _merge_default_parameters(self, model_config: ModelConfig) -> None: pooler_config = model_config.pooler_config @@ -228,6 +272,7 @@ def __repr__(self) -> str: f"requires_token_ids={self.requires_token_ids}, " f"skip_reading_prefix_cache={self.skip_reading_prefix_cache}, " f"late_interaction_params={self.late_interaction_params}, " + f"late_chunking_params={self.late_chunking_params}, " f"extra_kwargs={self.extra_kwargs})" ) From cd7517fb17bc26f5e63b348a73cc0dab70be8d48 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:31:53 +0400 Subject: [PATCH 379/624] [Bugfix] Fix ReplaySSM ring shape assertion when TP does not divide Mamba2 n_groups (#60704) Signed-off-by: Vadim Gimpelson Co-authored-by: Claude --- tests/kernels/mamba/test_ssu_dispatch.py | 24 +++++++++++++++++++ .../layers/mamba/mamba_utils.py | 3 ++- 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 3b3dc2dc6bec..e12369ef8a04 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -637,3 +637,27 @@ def test_replayssm_physical_ring_shape( (8, expected_ring_len), (2, expected_ring_len, 16), ) + + +def test_replayssm_ring_shape_extends_groups_for_head_shards(): + """A single group under TP is replicated to every head shard, as in + ``mamba2_state_shape``, so each rank's B_cache holds one group.""" + base_shapes = MambaStateShapeCalculator.mamba2_state_shape( + tp_world_size=2, + intermediate_size=256, + n_groups=1, + num_heads=4, + head_dim=64, + state_size=128, + conv_kernel=4, + ) + + shapes = MambaStateShapeCalculator.append_replayssm_ring( + base_shapes, + n_groups=1, + tp_world_size=2, + logical_window=16, + backend=MambaBackendEnum.TRITON, + ) + + assert shapes[4] == (1, 16, 128) diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 31e89ffe55b5..5c423739b7a0 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -238,13 +238,14 @@ def append_replayssm_ring( """Append the physical ReplaySSM ring shapes. ``base_shapes[1]`` is ``(nheads // tp, head_dim, state_size)``; - B_cache uses the un-extended ``n_groups``. + B_cache extends ``n_groups`` for head shards as in ``mamba2_state_shape``. """ ring_buffer_len = logical_window if backend == MambaBackendEnum.FLASHINFER: # FlashInfer keeps the live window and current verify window together. ring_buffer_len += 1 + num_speculative_tokens local_nheads, head_dim, state_size = base_shapes[1] + n_groups = n_groups + cls.extra_groups_for_head_shards(n_groups, tp_world_size) local_ngroups = divide(n_groups, tp_world_size) return ( *base_shapes, From 818f70c295137ddf9d2ba63262adb6f12e954b58 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:38:23 +0400 Subject: [PATCH 380/624] [Bugfix] Mask special added tokens in xgrammar JSON, regex and EBNF grammars (#60749) Signed-off-by: Vadim Gimpelson Signed-off-by: Artem Perevedentsev Co-authored-by: Claude Co-authored-by: Artem Perevedentsev --- .../test_backend_xgrammar_stop_tokens.py | 93 +++++++++++++++++++ vllm/v1/structured_output/backend_xgrammar.py | 86 +++++++++++++++-- 2 files changed, 170 insertions(+), 9 deletions(-) diff --git a/tests/v1/structured_output/test_backend_xgrammar_stop_tokens.py b/tests/v1/structured_output/test_backend_xgrammar_stop_tokens.py index e910d29aa38c..5b0baf44fd3c 100644 --- a/tests/v1/structured_output/test_backend_xgrammar_stop_tokens.py +++ b/tests/v1/structured_output/test_backend_xgrammar_stop_tokens.py @@ -8,8 +8,15 @@ bitmask while the FSM is still mid-object and truncate structured output. ``compile_grammar`` now forwards ``all_stop_token_ids`` to the matcher as ``override_stop_tokens`` so xgrammar masks them until the grammar completes. + +Issue #60580: xgrammar only treats tokens that decode to an empty string as +special, so a tokenizer's control tokens (e.g. ``<|im_start|>``) matched +grammar text and could be sampled inside a JSON string. Structural tags and +Lark grammars, which can name such tokens, must still be able to emit them. """ +import json + import pytest from transformers import AutoTokenizer @@ -25,6 +32,14 @@ QUOTE = 1 # standalone `"`; opens then closes the JSON string LETTER = 55 # `X`: valid string content, not a special/stop token by default +QWEN_TOKENIZER = "Qwen/Qwen3-0.6B" +QWEN_VOCAB_SIZE = 151936 + +# Qwen3 special added tokens. +ENDOFTEXT = 151643 # <|endoftext|>: a generation_config stop token +IM_START = 151644 # <|im_start|>: never a stop token +IM_END = 151645 # <|im_end|>: the tokenizer's eos + def _token_allowed(row, token_id: int) -> bool: word = int(row[token_id // 32].item()) & 0xFFFFFFFF @@ -80,3 +95,81 @@ def test_request_stop_tokens_gated_to_grammar_terminal(backend: XgrammarBackend) assert _token_allowed(bm_override[0], LETTER) assert _token_allowed(bm_default[0], EOS) assert _token_allowed(bm_override[0], EOS) + + +@pytest.fixture(scope="module") +def qwen_tokenizer(): + return AutoTokenizer.from_pretrained(QWEN_TOKENIZER) + + +@pytest.fixture(scope="module") +def qwen_backend(qwen_tokenizer) -> XgrammarBackend: + vllm_config = VllmConfig( + structured_outputs_config=StructuredOutputsConfig(backend="xgrammar") + ) + return XgrammarBackend( + vllm_config, tokenizer=qwen_tokenizer, vocab_size=QWEN_VOCAB_SIZE + ) + + +@pytest.mark.parametrize( + "request_type, grammar_spec", + [ + (StructuredOutputOptions.JSON, '{"type": "string"}'), + (StructuredOutputOptions.REGEX, '"[^"]*"'), + (StructuredOutputOptions.GRAMMAR, 'root ::= "\\"" [^"]* "\\""'), + ], +) +def test_special_tokens_masked_until_grammar_ends( + qwen_backend: XgrammarBackend, qwen_tokenizer, request_type, grammar_spec +): + grammar = qwen_backend.compile_grammar( + request_type, grammar_spec, stop_token_ids={IM_END, ENDOFTEXT} + ) + bitmask = qwen_backend.allocate_token_bitmask(1) + + text = qwen_tokenizer.encode('"The company is', add_special_tokens=False) + assert grammar.accept_tokens("req", text) + grammar.fill_bitmask(bitmask, 0) + for token_id in (IM_START, IM_END, ENDOFTEXT): + assert not _token_allowed(bitmask[0], token_id) + + quote = qwen_tokenizer.encode('"', add_special_tokens=False) + assert grammar.accept_tokens("req", quote) + grammar.fill_bitmask(bitmask, 0) + assert not _token_allowed(bitmask[0], IM_START) + assert _token_allowed(bitmask[0], IM_END) + assert _token_allowed(bitmask[0], ENDOFTEXT) + + +@pytest.mark.parametrize( + "request_type, grammar_spec", + [ + ( + StructuredOutputOptions.STRUCTURAL_TAG, + json.dumps( + { + "type": "structural_tag", + "format": { + "type": "tag", + "begin": "<|im_start|>", + "content": { + "type": "json_schema", + "json_schema": {"type": "string"}, + }, + "end": "\n", + }, + } + ), + ), + (StructuredOutputOptions.GRAMMAR, f'start: <[{IM_START}]> "assistant"'), + ], +) +def test_token_naming_grammars_can_emit_special_tokens( + qwen_backend: XgrammarBackend, request_type, grammar_spec +): + grammar = qwen_backend.compile_grammar(request_type, grammar_spec) + bitmask = qwen_backend.allocate_token_bitmask(1) + grammar.fill_bitmask(bitmask, 0) + assert _token_allowed(bitmask[0], IM_START) + assert grammar.accept_tokens("req", [IM_START]) diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 9631766d9267..94ccc55b8b98 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -26,6 +26,8 @@ if TYPE_CHECKING: import xgrammar as xgr + + from vllm.tokenizers import TokenizerLike else: xgr = LazyLoader("xgr", globals(), "xgrammar") @@ -43,6 +45,7 @@ def __post_init__(self): model_config is not None and model_config.hf_config.model_type == "plamo3" ) + full_vocab_tokenizer_info = None if is_mistral_tokenizer(self.tokenizer): # NOTE: ideally, xgrammar should handle this accordingly. # refer to https://github.com/mlc-ai/xgrammar/blob/d77c0a0173ef14779c918e3be7966ba852f7910f/python/xgrammar/tokenizer_info.py#L98 @@ -66,15 +69,46 @@ def __post_init__(self): ): tokenizer_info, _ = init_xgrammar() else: - tokenizer_info = xgr.TokenizerInfo.from_huggingface( + # tokenizer for structural tags + # accepts some of the special tokens + full_vocab_tokenizer_info = xgr.TokenizerInfo.from_huggingface( self.tokenizer, vocab_size=self.vocab_size, ) - self.compiler = xgr.GrammarCompiler( - tokenizer_info, - max_threads=8, - cache_enabled=True, - cache_limit_bytes=vllm.envs.VLLM_XGRAMMAR_CACHE_MB * 1024 * 1024, + # tokenizer for json + # does NOT accept any special tokens + tokenizer_info = self._filter_special_tokens( + self.tokenizer, full_vocab_tokenizer_info + ) + + # FIXME(arpera): + # Here we create not one but TWO compilers + # because xgrammar currently does not properly handle + # special tokens. + # Proper special tokens handling depends on structured output + # request type: json, structural_tag, etc. + # Structural tag may accept some of the special tokens + # At the same time json must NOT accept any special tokens + # Currently xgrammar does not properly handle all the cases + # So, we have to implement additional logic on our side + # We maintain TWO compilers for different structural output types + # If in future xgrammar will fix this problem, + # then we can consider to remove this + # Also, note that other backends like guidance + # do NOT have this problem (as far as I see, I may be wrong). + # So, if someone wants to go and fix the problem on xgrammar side + # you should first definitely check how other backends deal with it + # Some work on xgrammar side has been started: + # https://github.com/mlc-ai/xgrammar/issues/832 + + # Compiler for json. Must NOT accept any special tokens. + self.compiler = self._create_compiler(tokenizer_info) + + # Compiler for structural tag. May accept some of the special tokens. + self.full_vocab_compiler = ( + self.compiler + if full_vocab_tokenizer_info is None + else self._create_compiler(full_vocab_tokenizer_info) ) self.num_speculative_tokens = 0 @@ -122,7 +156,7 @@ def compile_grammar( ) elif request_type == StructuredOutputOptions.GRAMMAR: if grammar_is_likely_lark(grammar_spec): - ctx = self.compiler.compile_lark(grammar_spec) + ctx = self.full_vocab_compiler.compile_lark(grammar_spec) else: ctx = self.compiler.compile_grammar(grammar_spec) elif request_type == StructuredOutputOptions.REGEX: @@ -142,9 +176,11 @@ def compile_grammar( ) for s in s_tag["structures"] ] - ctx = self.compiler.compile_structural_tag(tags, s_tag["triggers"]) + ctx = self.full_vocab_compiler.compile_structural_tag( + tags, s_tag["triggers"] + ) else: - ctx = self.compiler.compile_structural_tag(grammar_spec) + ctx = self.full_vocab_compiler.compile_structural_tag(grammar_spec) else: logger.error( "Validation should have already occurred. Please file an issue." @@ -168,6 +204,38 @@ def allocate_token_bitmask(self, max_num_seqs: int): def destroy(self): del self.compiler + del self.full_vocab_compiler + + @staticmethod + def _create_compiler( + tokenizer_info: "xgr.TokenizerInfo", + ) -> "xgr.GrammarCompiler": + return xgr.GrammarCompiler( + tokenizer_info, + max_threads=8, + cache_enabled=True, + cache_limit_bytes=vllm.envs.VLLM_XGRAMMAR_CACHE_MB * 1024 * 1024, + ) + + @staticmethod + def _filter_special_tokens( + tokenizer: "TokenizerLike", tokenizer_info: "xgr.TokenizerInfo" + ) -> "xgr.TokenizerInfo": + """Filter ALL special tokens from tokenizer_info""" + special_token_ids = { + token_id + for token_id, token in getattr( + tokenizer, "added_tokens_decoder", {} + ).items() + if token.special + } + encoded_vocab = [""] * tokenizer_info.vocab_size + for token, token_id in tokenizer.get_vocab().items(): + if token_id < len(encoded_vocab) and token_id not in special_token_ids: + encoded_vocab[token_id] = token + return xgr.TokenizerInfo.from_vocab_and_metadata( + encoded_vocab, tokenizer_info.dump_metadata() + ) @dataclass From d9a543c500e9014045d7ae769df0ca7782e31870 Mon Sep 17 00:00:00 2001 From: Aditya Ganesh Kumar Date: Tue, 13 Oct 2026 17:31:59 +0530 Subject: [PATCH 381/624] [Bugfix][Frontend] Forward --enable-log-outputs to the Cohere Chat v2 (#61089) Signed-off-by: Aditya Ganesh Kumar --- vllm/entrypoints/cohere/serving.py | 4 ++++ vllm/entrypoints/generate/api_router.py | 2 ++ 2 files changed, 6 insertions(+) diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index ec6dd89c3ef8..93f6d14a3b9c 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -254,6 +254,8 @@ def __init__( tool_strict_level: str = "auto", enable_prompt_tokens_details: bool = False, enable_force_include_usage: bool = False, + enable_log_outputs: bool = False, + enable_log_deltas: bool = True, default_chat_template_kwargs: dict[str, Any] | None = None, is_reasoning_model: bool = True, ) -> None: @@ -272,6 +274,8 @@ def __init__( tool_parser=tool_parser, enable_prompt_tokens_details=enable_prompt_tokens_details, enable_force_include_usage=enable_force_include_usage, + enable_log_outputs=enable_log_outputs, + enable_log_deltas=enable_log_deltas, default_chat_template_kwargs=default_chat_template_kwargs, ) # Controls how the assistant's chain-of-thought is surfaced on diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 9bbe4e879b34..942a1c3e78c2 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -208,6 +208,8 @@ async def init_generate_state( reasoning_parser=args.structured_outputs_config.reasoning_parser, enable_prompt_tokens_details=args.enable_prompt_tokens_details, enable_force_include_usage=args.enable_force_include_usage, + enable_log_outputs=args.enable_log_outputs, + enable_log_deltas=args.enable_log_deltas, default_chat_template_kwargs=default_chat_template_kwargs, is_reasoning_model=args.cohere_is_reasoning_model, ) From 253e02b95e9c9deaebdc507f2eb26b29daf3fdf3 Mon Sep 17 00:00:00 2001 From: simon-veitner-redhat Date: Tue, 13 Oct 2026 14:17:48 +0200 Subject: [PATCH 382/624] [Perf] Warm up DeepGEMM MegaMoE kernels before serving (#60468) Signed-off-by: Simon Veitner Co-authored-by: Claude Co-authored-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> --- tests/model_executor/test_deep_gemm_warmup.py | 51 ++++++ .../model_executor/warmup/deep_gemm_warmup.py | 152 +++++++++++++++++- vllm/models/kimi_k3/nvidia/model.py | 1 + 3 files changed, 203 insertions(+), 1 deletion(-) diff --git a/tests/model_executor/test_deep_gemm_warmup.py b/tests/model_executor/test_deep_gemm_warmup.py index e71061e73317..c2dc12a27d97 100644 --- a/tests/model_executor/test_deep_gemm_warmup.py +++ b/tests/model_executor/test_deep_gemm_warmup.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import sys from types import SimpleNamespace import pytest @@ -88,3 +89,53 @@ def test_kernel_registers_itself_as_warmup_provider(is_bmm) -> None: provider = getattr(layer, "deep_gemm_warmup_provider", None) assert provider is (None if is_bmm else kernel) + + +def test_mega_moe_predicate_matches_only_native_mega_moe_layers(monkeypatch) -> None: + """Without the MegaMoE model modules loaded, nothing matches and nothing is + imported; with them loaded, only DeepSeek-V4 and Kimi K3 MegaMoE layers match.""" + name = "vllm.models.deepseek_v4.nvidia.model" + kimi_name = "vllm.models.kimi_k3.nvidia.model" + monkeypatch.delitem(sys.modules, name, raising=False) + monkeypatch.delitem(sys.modules, kimi_name, raising=False) + other = SimpleNamespace(use_native_mega_moe=True, use_mega_moe=True) + assert not deep_gemm_warmup._mega_moe_may_use_deep_gemm(other) + assert name not in sys.modules and kimi_name not in sys.modules + + moe_cls = type("DeepseekV4MoE", (), {}) + monkeypatch.setitem(sys.modules, name, SimpleNamespace(DeepseekV4MoE=moe_cls)) + native, fused = moe_cls(), moe_cls() + native.use_native_mega_moe, fused.use_native_mega_moe = True, False + assert deep_gemm_warmup._mega_moe_may_use_deep_gemm(native) + assert not deep_gemm_warmup._mega_moe_may_use_deep_gemm(fused) + assert not deep_gemm_warmup._mega_moe_may_use_deep_gemm(other) + + kimi_cls = type("KimiMoE", (), {}) + monkeypatch.setitem(sys.modules, kimi_name, SimpleNamespace(KimiMoE=kimi_cls)) + kimi_mega, kimi_fused = kimi_cls(), kimi_cls() + kimi_mega.use_mega_moe, kimi_fused.use_mega_moe = True, False + assert deep_gemm_warmup._mega_moe_may_use_deep_gemm(kimi_mega) + assert not deep_gemm_warmup._mega_moe_may_use_deep_gemm(kimi_fused) + + +def test_mega_moe_m_values_reach_mega_gate(monkeypatch) -> None: + """Each bf16_mega_gate config is warmed with an M above the GateLinear + threshold (1 for 128 experts), so the forward actually reaches the kernel.""" + dg = SimpleNamespace( + get_bf16_mega_gate_config=lambda m, *_: {"num_sms": -(-m // 16)}, + get_block_m_for_mega_moe=lambda *_: 16, + ) + monkeypatch.setattr(deep_gemm_warmup, "_import_deep_gemm", lambda: dg) + monkeypatch.setattr(deep_gemm_warmup, "_mega_moe_uses_mega_gate", lambda _: True) + monkeypatch.setattr( + deep_gemm_warmup, "get_ep_group", lambda: SimpleNamespace(world_size=8) + ) + buffer = SimpleNamespace(num_max_tokens_per_rank=64) + module = SimpleNamespace( + gate=SimpleNamespace(input_size=4096, output_size=128), + experts=SimpleNamespace( + top_k=6, num_experts=128, get_symm_buffer=lambda: buffer + ), + use_sequence_parallel=False, + ) + assert deep_gemm_warmup._get_fp8_fp4_mega_moe_m_values(module, 40) == [16, 32, 40] diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py index a4775ce22e45..309c66f07e25 100644 --- a/vllm/model_executor/warmup/deep_gemm_warmup.py +++ b/vllm/model_executor/warmup/deep_gemm_warmup.py @@ -5,11 +5,19 @@ be used during model execution beforehand. """ +import sys +from typing import TYPE_CHECKING, TypeGuard + import torch from tqdm import tqdm import vllm.envs as envs -from vllm.distributed.parallel_state import get_dp_group, is_global_first_rank +from vllm.distributed.parallel_state import ( + get_dp_group, + get_ep_group, + is_global_first_rank, +) +from vllm.forward_context import ForwardContext, override_forward_context from vllm.model_executor.kernels.linear.scaled_mm.deep_gemm import ( DeepGemmFp8BlockScaledMMKernel, ) @@ -23,6 +31,7 @@ ) from vllm.tracing import instrument from vllm.utils.deep_gemm import ( + _import_deep_gemm, fp8_gemm_nt, get_mk_alignment_for_contiguous_layout, m_grouped_fp8_gemm_nt_contiguous, @@ -32,6 +41,10 @@ from vllm.utils.platform_utils import num_compute_units from vllm.v1.worker.workspace import current_workspace_manager +if TYPE_CHECKING: + from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MoE + from vllm.models.kimi_k3.nvidia.model import KimiMoE + def _generate_optimal_warmup_m_values( max_tokens: int, n: int, device: torch.device @@ -144,6 +157,22 @@ def _fused_moe_grouped_gemm_may_use_deep_gemm(module: torch.nn.Module) -> bool: return isinstance(fused_experts, (DeepGemmExperts, TritonOrDeepGemmExperts)) +# Model construction already imports the model modules. Avoid importing them here. +def _mega_moe_uses_mega_gate(module: torch.nn.Module) -> "TypeGuard[DeepseekV4MoE]": + # DeepSeek-V4 routes with bf16_mega_gate, Kimi K3 with its GateLinear. + dsv4 = sys.modules.get("vllm.models.deepseek_v4.nvidia.model") + return dsv4 is not None and isinstance(module, dsv4.DeepseekV4MoE) + + +def _mega_moe_may_use_deep_gemm( + module: torch.nn.Module, +) -> "TypeGuard[DeepseekV4MoE | KimiMoE]": + if _mega_moe_uses_mega_gate(module): + return module.use_native_mega_moe + kimi = sys.modules.get("vllm.models.kimi_k3.nvidia.model") + return kimi is not None and isinstance(module, kimi.KimiMoE) and module.use_mega_moe + + FP8_GEMM_NT_WARMUP_CACHE: set[torch.Size] = set() @@ -305,6 +334,110 @@ def _warmup(w: torch.Tensor, w_scale: torch.Tensor): GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE.add(w.size()) +FP8_FP4_MEGA_MOE_WARMUP_CACHE: set[tuple] = set() + + +def _get_fp8_fp4_mega_moe_warmup_key(module: "DeepseekV4MoE | KimiMoE") -> tuple: + gate = module.gate + return ( + gate.output_size, + getattr(gate, "tid2eid", None) is not None, + gate.e_score_correction_bias is not None, + getattr(module, "image_sentinel_lo", 0), + module.experts.has_fused_shared_experts, + ) + + +def _get_fp8_fp4_mega_moe_m_values( + module: "DeepseekV4MoE | KimiMoE", max_tokens: int +) -> list[int]: + """Get the largest M per distinct fp8_fp4_mega_moe config and, where used, + bf16_mega_gate config.""" + dg = _import_deep_gemm() + gate, experts = module.gate, module.experts + uses_mega_gate = _mega_moe_uses_mega_gate(module) + num_max_tokens_per_rank = experts.get_symm_buffer().num_max_tokens_per_rank + if module.use_sequence_parallel: + max_tokens = cdiv(max_tokens, module.tp_size) + + m_values: dict[tuple, int] = {} + for num_tokens in range(1, max_tokens + 1): + gate_config = ( + dg.get_bf16_mega_gate_config( + num_tokens, gate.input_size, gate.output_size, experts.top_k + ) + if uses_mega_gate + else {} + ) + block_m = dg.get_block_m_for_mega_moe( + get_ep_group().world_size, + experts.num_experts, + num_max_tokens_per_rank, + num_tokens, + experts.top_k, + "fp8xfp4", + ) + # Keep the largest M: DeepSeek-V4 only calls bf16_mega_gate above a + # small token threshold, and a config's largest M is always above it. + m_values[(tuple(gate_config.items()), block_m)] = num_tokens + return sorted(m_values.values()) + + +def _deepgemm_fp8_fp4_mega_moe_warmup( + module: "DeepseekV4MoE | KimiMoE", + max_tokens: int, + pbar: tqdm | None = None, +): + from vllm.models.deepseek_v4.nvidia.model import ( + prepare_mega_gate_routing_metadata, + ) + + key = _get_fp8_fp4_mega_moe_warmup_key(module) + if key in FP8_FP4_MEGA_MOE_WARMUP_CACHE: + return + + m_values = _get_fp8_fp4_mega_moe_m_values(module, max_tokens) + max_m = m_values[-1] # sorted + + device = module.experts.get_symm_buffer().x.device + input_ids = torch.zeros(max_m, device=device, dtype=torch.int64) + hidden_states = torch.zeros( + (max_m, module.gate.input_size), device=device, dtype=torch.bfloat16 + ) + is_padding = torch.zeros(max_m + 1, device=device, dtype=torch.bool) + # With padding skip on, sequence parallel gives ranks > 0 a padding mask + # slice at any alignment. + sp_masked = module.use_sequence_parallel and envs.VLLM_MOE_SKIP_PADDING + offsets = (0, 1) if sp_masked else (0,) + + # Run the whole layer so the gate sees the same optional inputs as serving. + for num_tokens in m_values: + inputs: tuple = (hidden_states[:num_tokens],) + if _mega_moe_uses_mega_gate(module): + ids = input_ids[:num_tokens] + inputs += ( + ids, + prepare_mega_gate_routing_metadata( + ids, + has_hash_routing=module.gate.tid2eid is not None, + image_sentinel_base_id=module.image_sentinel_lo or None, + ), + ) + for offset in offsets: + forward_context = ForwardContext( + no_compile_layers={}, + attn_metadata={}, + slot_mapping={}, + is_padding=is_padding[offset : offset + num_tokens], + ) + with override_forward_context(forward_context): + module(*inputs) + if pbar is not None: + pbar.update(1) + + FP8_FP4_MEGA_MOE_WARMUP_CACHE.add(key) + + def deepgemm_fp8_gemm_nt_warmup( model: torch.nn.Module, max_tokens: int, pbar: tqdm | None = None ): @@ -330,11 +463,21 @@ def deepgemm_grouped_fp8_gemm_nt_contiguous_warmup( ) +@torch.inference_mode() +def deepgemm_fp8_fp4_mega_moe_warmup( + model: torch.nn.Module, max_tokens: int, pbar: tqdm | None = None +): + for m in model.modules(): + if _mega_moe_may_use_deep_gemm(m): + _deepgemm_fp8_fp4_mega_moe_warmup(m, max_tokens, pbar=pbar) + + def _count_warmup_iterations(model: torch.nn.Module, max_tokens: int) -> int: seen_fp8_sizes: set[torch.Size] = set(FP8_GEMM_NT_WARMUP_CACHE) seen_grouped_sizes: set[torch.Size] = set( GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE ) + seen_mega_moe_keys: set[tuple] = set(FP8_FP4_MEGA_MOE_WARMUP_CACHE) total = 0 for m in model.modules(): @@ -355,6 +498,11 @@ def _count_warmup_iterations(model: torch.nn.Module, max_tokens: int) -> int: if w2.size() not in seen_grouped_sizes: total += n_values seen_grouped_sizes.add(w2.size()) + elif _mega_moe_may_use_deep_gemm(m): + key = _get_fp8_fp4_mega_moe_warmup_key(m) + if key not in seen_mega_moe_keys: + total += len(_get_fp8_fp4_mega_moe_m_values(m, max_tokens)) + seen_mega_moe_keys.add(key) return total @@ -369,6 +517,8 @@ def deep_gemm_warmup(model: torch.nn.Module, max_tokens: int): with tqdm(total=total, desc="DeepGEMM warmup") as pbar: deepgemm_fp8_gemm_nt_warmup(model, max_tokens, pbar) deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(model, max_tokens, pbar) + deepgemm_fp8_fp4_mega_moe_warmup(model, max_tokens, pbar) else: deepgemm_fp8_gemm_nt_warmup(model, max_tokens, None) deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(model, max_tokens, None) + deepgemm_fp8_fp4_mega_moe_warmup(model, max_tokens, None) diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index eb2ac88aa436..ae9b7da3cd11 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -566,6 +566,7 @@ def __init__( ) self.latent_moe_use_norm = config.latent_moe_use_norm self.tp_size = get_tensor_model_parallel_world_size() + self.use_sequence_parallel = use_sequence_parallel self.routed_scaling_factor = config.routed_scaling_factor self.moe_renormalize = moe_renormalize self.use_grouped_topk = config.use_grouped_topk From 7342dc0317c02af21aa15109f147f98982cb8c4f Mon Sep 17 00:00:00 2001 From: Hexiang Wang <56632993+whx-sjtu@users.noreply.github.com> Date: Tue, 13 Oct 2026 20:19:53 +0800 Subject: [PATCH 383/624] [MoRIIO] Require pending remote prefill for READ cache hits (#60638) Signed-off-by: whx-sjtu --- .../unit/test_moriio_hma_scheduler.py | 40 ++++++++++++++++++- .../v1/moriio/moriio_connector.py | 6 +-- 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py index 4bd32f786a28..c6df14d7a0fe 100644 --- a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py +++ b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py @@ -806,12 +806,50 @@ def test_truncate_mamba_request_noop_for_single_token_prompt(): # -------------------------------------------------------------------------- # get_num_new_matched_tokens (hybrid N-1 accounting) # -------------------------------------------------------------------------- +@pytest.mark.parametrize( + "params,expected", + [ + (None, (0, False)), + ({}, (0, False)), + ({"do_remote_prefill": False}, (0, False)), + ({"do_remote_prefill": True}, (31, False)), + ], +) +def test_read_matches_only_requests_with_pending_remote_prefill(params, expected): + scheduler = _FakeScheduler(is_producer=False, mode=MoRIIOMode.READ) + request = SimpleNamespace(num_prompt_tokens=32, kv_transfer_params=params) + + assert scheduler.get_num_new_matched_tokens(request, 0) == expected + + +def test_read_rematch_after_alloc_keeps_queued_transfer(): + """Clearing the flag blocks another external hit, not the queued READ.""" + sched = _make_read_scheduler() + request = _make_read_request([[10, 11], [91]]) + request.num_prompt_tokens = 32 + assert sched.get_num_new_matched_tokens(request, 0) == (31, False) + + sched.update_state_after_alloc( + request, _FakeBlocks(all_groups=([100, 101], [200])), num_external_tokens=31 + ) + + assert sched._reqs_need_recv["req"][1] == [[100, 101], [200]] + assert request.kv_transfer_params["do_remote_prefill"] is False + + meta = sched.build_connector_meta(SchedulerOutput.make_empty()) + queued = meta.reqs_to_recv["req"] + assert queued.transfer_id == "tx" + assert queued.local_block_ids == [[100, 101], [200]] + assert queued.remote_block_ids == [[10, 11], [91]] + assert sched.get_num_new_matched_tokens(request, 0) == (0, False) + + def test_get_num_new_matched_tokens_read_recomputes_last_token(): sched = _FakeScheduler(is_producer=False, mode=MoRIIOMode.READ, _has_mamba=True) req = SimpleNamespace( num_prompt_tokens=10, prompt_token_ids=list(range(10)), - kv_transfer_params=None, + kv_transfer_params={"do_remote_prefill": True}, ) n, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=0) # READ always recomputes the final token locally: N-1 - computed. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index c398cb20a8b0..0e6eccd8ea91 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -776,10 +776,10 @@ def get_num_new_matched_tokens( num_prompt_tokens = request.num_prompt_tokens num_external_tokens = max(num_prompt_tokens - num_computed_tokens, 0) + params = request.kv_transfer_params or {} + if not params.get("do_remote_prefill"): + return 0, False if self.mode == MoRIIOMode.WRITE: - params = request.kv_transfer_params or {} - if not params.get("do_remote_prefill"): - return 0, False # MoriiO in write mode, no remote prefill. Hybrid models never get # here: register_kv_caches refuses WRITE mode for them, so there is # no recurrent-state accounting to do. From 0838cd6593ba751aa215c9ecd44048bb29dcff1b Mon Sep 17 00:00:00 2001 From: Nathan Sala <33957625+natsala13@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:42:29 +0300 Subject: [PATCH 384/624] [Bugfix] Load and apply Nemotron Super final vision LayerNorm (#61401) Signed-off-by: Shaun Kotek Co-authored-by: Shaun Kotek - Nvidia --- .../multimodal/test_nano_nemotron_vl.py | 57 +++++++++++++++++++ .../model_executor/models/nano_nemotron_vl.py | 35 +++++++++++- 2 files changed, 91 insertions(+), 1 deletion(-) diff --git a/tests/models/multimodal/test_nano_nemotron_vl.py b/tests/models/multimodal/test_nano_nemotron_vl.py index 75900cde53a3..41f9428e6d73 100644 --- a/tests/models/multimodal/test_nano_nemotron_vl.py +++ b/tests/models/multimodal/test_nano_nemotron_vl.py @@ -97,6 +97,7 @@ def test_nano_nemotron_vl_skips_multimodal_weights_in_text_only_mode(): object.__setattr__(model, "language_model", language_model) object.__setattr__(model, "mlp1", _AdapterModule()) object.__setattr__(model, "vision_model", _MissingMultiModalModule()) + object.__setattr__(model, "vision_final_layernorm", _MissingMultiModalModule()) object.__setattr__(model, "sound_encoder", None) language_weight = object() @@ -106,6 +107,8 @@ def test_nano_nemotron_vl_skips_multimodal_weights_in_text_only_mode(): ("mlp1.0.weight", object()), ("vision_model.radio_model.encoder.weight", object()), ("sound_encoder.encoder.weight", object()), + ("vision_projector.vision_final_layernorm.weight", object()), + ("vision_projector.vision_final_layernorm.bias", object()), ] ) @@ -120,6 +123,7 @@ def test_nano_nemotron_vl_loads_vision_weights_without_sound_encoder(): object.__setattr__(model, "language_model", language_model) object.__setattr__(model, "mlp1", _AdapterModule()) object.__setattr__(model, "vision_model", vision_model) + object.__setattr__(model, "vision_final_layernorm", None) object.__setattr__(model, "sound_encoder", None) language_weight = object() @@ -145,12 +149,65 @@ def test_nano_nemotron_vl_requires_sound_encoder_for_sound_weights(): object.__setattr__(model, "language_model", language_model) object.__setattr__(model, "mlp1", _AdapterModule()) object.__setattr__(model, "vision_model", vision_model) + object.__setattr__(model, "vision_final_layernorm", None) object.__setattr__(model, "sound_encoder", None) with pytest.raises(AssertionError): model.load_weights([("sound_encoder.encoder.weight", object())]) +def _vision_norm_model(use_layernorm=True): + model = object.__new__(NemotronH_Nano_VL_V2) + torch.nn.Module.__init__(model) + model.model_config = _ImageOnlyModelConfig() + model.language_model = _LanguageModel() + model.vision_model = _VisionModel() + model.sound_encoder = None + model.mlp1 = torch.nn.Identity() + model.vision_final_layernorm = ( + torch.nn.LayerNorm(4, eps=1e-6, dtype=torch.bfloat16) if use_layernorm else None + ) + model.patch_size = 2 + model.downsample_ratio = 0.5 + model.ps_version = "v2" + model.video_temporal_patch_size = 2 + return model + + +@pytest.mark.parametrize("mode", ["image", "dynamic", "video"]) +@pytest.mark.parametrize("use_layernorm", [True, False]) +def test_vision_layernorm_is_loaded_and_applied_before_pixel_shuffle( + mode, use_layernorm +): + model = _vision_norm_model(use_layernorm) + weight = torch.tensor([2.0, 3.0, 4.0, 5.0], dtype=torch.bfloat16) + bias = torch.tensor([0.5, -0.5, 1.0, -1.0], dtype=torch.bfloat16) + model.load_weights( + [ + ("vision_projector.vision_final_layernorm.weight", weight), + ("vision_projector.vision_final_layernorm.bias", bias), + ] + if use_layernorm + else [] + ) + features = torch.arange(16, dtype=torch.bfloat16).reshape(1, 4, 4) + expected = ( + torch.nn.functional.layer_norm(features, (4,), weight, bias, eps=1e-6) + if use_layernorm + else features + ) + expected = model.pixel_shuffle(expected.reshape(1, 2, 2, 4)).reshape(1, 1, 16) + model.vision_model = lambda *args, **kwargs: (None, features) + pixels = torch.zeros(2 if mode == "video" else 1, 3, 4, 4) + if mode == "dynamic": + output = model.extract_feature_dynamic(pixels, imgs_sizes=[(4, 4)]) + else: + output = model.extract_feature( + pixels, num_frames=2 if mode == "video" else None + ) + torch.testing.assert_close(output, expected) + + def _make_mm_items_with_video_bytes( video_bytes: bytes, ) -> MultiModalDataItems: diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 074a89abb93f..b77206295d5c 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -1001,6 +1001,16 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vision_projection_hidden_size = config.projector_hidden_size llm_hidden_size = config.text_config.hidden_size + # Match the HF vision projector for checkpoints trained with MTP. + self.vision_final_layernorm = ( + nn.LayerNorm( + vit_hidden_size, + eps=getattr(vision_config, "layer_norm_eps", 1e-6), + ).to(llm_dtype) + if (getattr(config.text_config, "num_nextn_predict_layers", 0) or 0) > 0 + else None + ) + mlp1 = nn.Sequential( RMSNorm( hidden_size=vit_hidden_size @@ -1095,6 +1105,8 @@ def extract_feature_dynamic( """Dynamic resolution extract_feature for images.""" _, vit_embeds = self.vision_model(pixel_values, imgs_sizes=imgs_sizes) vit_embeds = vit_embeds.to(dtype=torch.bfloat16) + if self.vision_final_layernorm is not None: + vit_embeds = self.vision_final_layernorm(vit_embeds) vit_embeds = self.pixel_shuffle_dynamic_res(vit_embeds, imgs_sizes=imgs_sizes) vit_embeds = self.mlp1(vit_embeds) return vit_embeds @@ -1127,6 +1139,8 @@ def extract_feature( else: _, vit_embeds = self.vision_model(chunk) vit_embeds = vit_embeds.to(dtype=torch.bfloat16) + if self.vision_final_layernorm is not None: + vit_embeds = self.vision_final_layernorm(vit_embeds) vit_embeds = vit_embeds.reshape( vit_embeds.shape[0], H_patches, W_patches, -1 ) @@ -1575,7 +1589,7 @@ def get_mm_mapping(self) -> MultiModelKeys: """Get the module prefix in multimodal models""" return MultiModelKeys.from_string_field( language_model="language_model", - connector=["mlp1", "sound_encoder.projection"], + connector=["mlp1", "sound_encoder.projection", "vision_final_layernorm"], tower_model=["vision_model", "sound_encoder.encoder"], ) @@ -1593,6 +1607,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): for modality in ("image", "video", "audio") ) adapter_dict = dict(self.mlp1.named_parameters()) + vision_norm_params = ( + dict(self.vision_final_layernorm.named_parameters()) + if load_multimodal_weights and self.vision_final_layernorm is not None + else {} + ) def is_llm(name: str) -> bool: return name.startswith("language_model") @@ -1606,6 +1625,11 @@ def is_vision_weights(name: str) -> bool: def is_sound_weights(name: str) -> bool: return name.startswith("sound") + vision_norm_prefix = "vision_projector.vision_final_layernorm." + + def is_vision_norm_weights(name: str) -> bool: + return name.startswith(vision_norm_prefix) + # LLM weights (the bulk of the model) are streamed lazily through a # generator so each tensor is copied into its parameter before the # iterator advances, avoiding stale-reference corruption with @@ -1637,6 +1661,15 @@ def llm_weights_gen(): continue assert self.sound_encoder is not None sound_weights.append((name, w.detach().clone())) + elif is_vision_norm_weights(name): + if not load_multimodal_weights: + continue + param_name = name.removeprefix(vision_norm_prefix) + assert param_name in vision_norm_params, ( + f"Unexpected vision LayerNorm weight: {name}" + ) + with torch.no_grad(): + default_weight_loader(vision_norm_params[param_name], w) # Fully drain the generator so every mm tensor is buffered, even if # the LLM loader stops iterating early. From 31e90ac3705e7415e4061d72c640c6c01737af15 Mon Sep 17 00:00:00 2001 From: simon-veitner-redhat Date: Tue, 13 Oct 2026 15:31:37 +0200 Subject: [PATCH 385/624] [Attention] Enable FA4 sparse MQA decode on Blackwell (#55866) Signed-off-by: Simon Veitner Co-authored-by: Lucas Wilkinson Co-authored-by: Claude Co-authored-by: Lucas Wilkinson Co-authored-by: Matthew Bonanni --- cmake/external_projects/vllm_flash_attn.cmake | 2 +- docs/design/attention_backends.md | 2 +- .../attention/test_flashinfer_mla_decode.py | 10 +- ...st_flashmla_nope_sm90_backend_selection.py | 34 ++ .../v1/attention/test_sparse_mla_backends.py | 475 ++++++++++++++---- .../test_sparse_mla_kv_cache_layout.py | 27 +- tests/v1/e2e/general/test_hisparse.py | 12 +- .../layers/attention/mla_attention.py | 6 +- .../layers/attention/sparse_mla_attention.py | 159 +++++- vllm/platforms/cuda.py | 3 +- .../backends/mla/flashattn_mla_sparse.py | 313 +++++++----- .../backends/mla/flashinfer_mla_sparse.py | 172 ++----- vllm/v1/attention/backends/registry.py | 4 + vllm/vllm_flash_attn/flash_attn_interface.py | 24 +- 14 files changed, 860 insertions(+), 383 deletions(-) diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 2a26023f740b..bac2afc828b3 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -39,7 +39,7 @@ else() FetchContent_Declare( vllm-flash-attn GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git - GIT_TAG 9cd61de38763d712bb6ce56e2a02cc2bf718c89f + GIT_TAG b273a02621d3c94bc747df67204f591b67f38fae GIT_PROGRESS TRUE # Don't share the vllm-flash-attn build between build types BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index d175add40b84..21cde2fbed16 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -147,7 +147,7 @@ Priority is **1 = highest** (tried first). --8<-- "gen:priority-mla" -> **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise. +> **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASH_ATTN_MLA_SPARSE_FA4` (then `FLASHINFER_MLA_SPARSE`) is preferred for low query-head counts (<= 16) on SM 10.x, while `FLASHMLA_SPARSE` is preferred otherwise. > > **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details. diff --git a/tests/kernels/attention/test_flashinfer_mla_decode.py b/tests/kernels/attention/test_flashinfer_mla_decode.py index c3f7b4c96f78..fc34b8725746 100644 --- a/tests/kernels/attention/test_flashinfer_mla_decode.py +++ b/tests/kernels/attention/test_flashinfer_mla_decode.py @@ -175,7 +175,15 @@ def unexpected_helper(*args, **kwargs): backend, "prepare_sparse_mla_safe_lengths", unexpected_helper ) - out, lse = impl._run_mqa_kernel(query.squeeze(1), kv_cache, indices, counts) + out, lse = impl._forward_mqa_kernel( + query.squeeze(1), + kv_cache, + indices, + counts, + layer=None, + block_size=block_size, + is_decode=True, + ) reference = trtllm_batch_decode_with_kv_cache_mla( query=query, kv_cache=kv_cache.unsqueeze(1), diff --git a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py index 8532d2f43aee..8f31bd488aa4 100644 --- a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py +++ b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py @@ -238,6 +238,7 @@ def _sparse_order(kv_cache_dtype, head_size, num_heads=32, capability=SM100): ) sparse = { "FLASH_ATTN_MLA_SPARSE", + "FLASH_ATTN_MLA_SPARSE_FA4", "FLASHMLA_SPARSE", "FLASHINFER_MLA_SPARSE", "FLASHINFER_MLA_SPARSE_SM90", @@ -276,3 +277,36 @@ def test_sm90_nope_fp8_ds_mla_resolves_to_flashmla(): for name in order[: order.index("FLASHMLA_SPARSE")]: backend_cls = AttentionBackendEnum[name].get_class() assert not backend_cls.supports_kv_cache_dtype("fp8_ds_mla"), name + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA backends") +def test_sm100_576_bf16_low_heads_falls_back_from_fa4(monkeypatch): + """With FA4 unusable, low-head bf16 selection lands on FlashInfer.""" + import vllm.v1.attention.backends.mla.flashattn_mla_sparse as fa4_sparse + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + monkeypatch.setattr(fa4_sparse, "_fa4_cute_mla_available", lambda: "no FA4") + + def reasons(name): + backend_cls = AttentionBackendEnum[name].get_class() + return backend_cls.validate_configuration( + head_size=576, + dtype=torch.bfloat16, + kv_cache_dtype="auto", + block_size=64, + use_mla=True, + has_sink=False, + use_sparse=True, + use_mm_prefix=False, + use_per_head_quant_scales=False, + device_capability=SM100, + attn_type="decoder", + ) + + order = _sparse_order("auto", 576, num_heads=16) + with set_current_vllm_config(VllmConfig()): + # The selector takes the first candidate with no rejection reasons. + selected = next(name for name in order if not reasons(name)) + assert order[0] == "FLASH_ATTN_MLA_SPARSE_FA4", order + assert selected == "FLASHINFER_MLA_SPARSE" diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index f7c47fe6a3ec..a2e216994328 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -34,6 +34,9 @@ ) from vllm.forward_context import set_forward_context from vllm.model_executor.layers.attention import mla_attention +from vllm.model_executor.layers.attention import ( + sparse_mla_attention as sparse_mla_module, +) from vllm.model_executor.layers.attention.mla_attention import _use_masked_mha from vllm.model_executor.layers.attention.sparse_mla_attention import ( GLOBAL_TOPK_MASK_MAX_BYTES, @@ -46,6 +49,7 @@ ) from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.platforms import current_platform +from vllm.platforms.interface import DeviceCapability # This module's end-to-end sparse-backend correctness test supports CUDA and # ROCm. The remaining tests exercise CUDA-specific FlashMLA/FlashInfer APIs @@ -71,8 +75,11 @@ flashattn_mla_sparse as flashattn_sparse_module, ) from vllm.v1.attention.backends.mla.flashattn_mla_sparse import ( + FlashAttnMLASparseFA4Backend, + FlashAttnMLASparseFA4Impl, FlashAttnMLASparseImpl, FlashAttnMLASparseMetadataBuilder, + _fa4_cute_mla_available, ) from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( FlashInferMLASparseImpl, @@ -227,6 +234,7 @@ def test_nope_flashinfer_sparse_mla_uses_model_scale(monkeypatch): req_id_per_token=torch.zeros(1, dtype=torch.int32), block_table=torch.zeros((1, 1), dtype=torch.int32), block_size=32, + num_decode_tokens=1, ) recorded_scale = None @@ -1401,6 +1409,7 @@ def test_masked_mha_workspace_fits_single_request_boundary(max_query_len, expect ("FLASHMLA_SPARSE", 8, 112 * 1024), ("FLASHINFER_MLA_SPARSE", 4, 36 * 1024), ("FLASHINFER_MLA_SPARSE", 8, 64 * 1024), + ("FLASH_ATTN_MLA_SPARSE_FA4", 4, 36 * 1024), ], ) def test_masked_mha_workspace_guards_long_routing_policy( @@ -3569,56 +3578,6 @@ def test_hisparse_keeps_speculative_rows_as_decodes(): assert builder.reorder_batch_threshold == 4 -def test_flashinfer_hisparse_decode_runs_batched_attention(): - device = torch.device("cpu") - num_tokens = 6 - q = torch.randn(num_tokens, 2, 4, device=device) - topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=device) - physical_topk = topk + 10 - valid_counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=device) - kernel_shapes: list[torch.Size] = [] - - def convert_decode(self, *args, **kwargs): # noqa: ARG001 - return physical_topk, valid_counts - - def prepare_kernel(self, *args, **kwargs): # noqa: ARG001 - pass - - def run_kernel(self, q, cache, indices, counts): # noqa: ARG001 - kernel_shapes.append(q.shape) - return q[..., :1], None - - cache_handle = SimpleNamespace( - runtime=SimpleNamespace( - hot=SimpleNamespace(attention_cache=torch.empty(1, device=device)) - ), - ) - index_group = object.__new__(HiSparseMLAIndexGroup) - index_group.caches = [cache_handle] - index_group.convert_logical_to_physical_topk = MethodType( - convert_decode, index_group - ) - impl = object.__new__(FlashInferMLASparseImpl) - impl.topk_indices_buffer = topk - impl.index_group = index_group - impl.index_group_index = 0 - impl._prepare_mqa_kernel = MethodType(prepare_kernel, impl) - impl._run_mqa_kernel = MethodType(run_kernel, impl) - metadata = SimpleNamespace(num_decode_tokens=num_tokens) - - output, lse = FlashInferMLASparseImpl.forward_mqa( - impl, - q, - torch.empty(1, device=device), - metadata, - SimpleNamespace(), - ) - - assert kernel_shapes == [q.shape] - assert output.shape == (num_tokens, 2, 1) - assert lse is None - - @pytest.mark.skipif( not current_platform.is_device_capability_family(90), reason="FlashAttention MLA requires Hopper", @@ -3690,7 +3649,7 @@ def test_flashattn_shared_indices_match_independent_conversion( convert = triton_convert_req_index_to_global_index convert_spy = MagicMock(wraps=convert) - for module in (index_group_module, flashattn_sparse_module): + for module in (index_group_module, sparse_mla_module): monkeypatch.setattr( module, "triton_convert_req_index_to_global_index", convert_spy ) @@ -3735,48 +3694,19 @@ def run(): return_valid_counts=True, ) for impl, cache, actual in zip(impls, caches, outputs): - expected = impl._run_mqa_kernel( - q_nope, q_rope, cache, physical, counts, block_size + expected, _ = impl._forward_mqa_kernel( + (q_nope, q_rope), + cache, + physical, + counts, + layer=None, + block_size=block_size, + is_decode=num_decode_tokens == num_tokens, ) torch.testing.assert_close(actual, expected, rtol=0, atol=0) stream.synchronize() -def test_flashattn_hisparse_decode_uses_index_group(): - num_tokens = 4 - q_nope = torch.empty(num_tokens, 2, 3, device=DEVICE_TYPE) - q_rope = torch.empty(num_tokens, 2, 1, device=DEVICE_TYPE) - topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=DEVICE_TYPE) - physical = topk.clone() - counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE) - index_group = object.__new__(HiSparseMLAIndexGroup) - index_group.convert_logical_to_physical_topk = MagicMock( - return_value=(physical, counts) - ) - index_group.physical_kv_cache = MagicMock( - return_value=torch.empty(1, device=DEVICE_TYPE) - ) - impl = object.__new__(FlashAttnMLASparseImpl) - impl.topk_indices_buffer = topk - impl.index_group = index_group - impl.index_group_index = 0 - impl._run_mqa_kernel = MagicMock(return_value=q_nope[..., :1]) - metadata = SimpleNamespace(num_decode_tokens=num_tokens, block_size=64) - - output, lse = FlashAttnMLASparseImpl.forward_mqa( - impl, - (q_nope, q_rope), - torch.empty(1, device=DEVICE_TYPE), - metadata, - SimpleNamespace(), - ) - - assert output.shape == (num_tokens, 2, 1) - assert lse is None - index_group.convert_logical_to_physical_topk.assert_called_once() - impl._run_mqa_kernel.assert_called_once() - - def test_flashinfer_sm120_hisparse_decode_uses_index_group(): num_tokens = 4 q = torch.empty(num_tokens, 2, 4, device=DEVICE_TYPE) @@ -3811,6 +3741,131 @@ def test_flashinfer_sm120_hisparse_decode_uses_index_group(): impl._run_mqa_kernel.assert_called_once() +def test_sparse_mqa_plain_lane_wiring(monkeypatch): + """Under DCP the plain lane converts with the DCP filter at the flat stride.""" + import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla + + num_tokens, topk_tokens, block_size, head_dim = 3, 8, 4, 6 + # Flat row stride (8) differs from block size (4): the wrong one misaddresses. + kv_cache = torch.zeros(2, 8, head_dim)[:, :block_size] + q = torch.zeros(num_tokens, 2, head_dim) + metadata = SimpleNamespace( + num_decode_tokens=num_tokens, + block_size=block_size, + block_table=torch.zeros((num_tokens, 1), dtype=torch.int32), + req_id_per_token=torch.zeros(num_tokens, dtype=torch.int32), + cp_kv_cache_interleave_size=2, + ) + index_filter = MagicMock(return_value=(None, None)) + monkeypatch.setattr(sparse_mla, "triton_filter_and_convert_dcp_index", index_filter) + + impl = object.__new__(FlashInferMLASparseImpl) + impl.topk_indices_buffer = torch.zeros(num_tokens, topk_tokens, dtype=torch.int32) + impl.index_group = MagicMock() + impl.index_group_index = 0 + impl.dcp_world_size = 2 + impl.dcp_rank = 1 + impl._convert_logical_to_physical_topk = MagicMock() + impl._forward_mqa_kernel = MagicMock(return_value=(q, None)) + + impl.forward_mqa(q, kv_cache, metadata, None) + + impl._convert_logical_to_physical_topk.assert_not_called() + assert index_filter.call_args.kwargs["BLOCK_STRIDE_ROWS"] == 8 + impl._forward_mqa_kernel.assert_called_once() + + +@pytest.mark.parametrize( + "num_decode_tokens,all_resident,split_q", + [ + pytest.param(0, True, False, id="resident"), + pytest.param(0, False, False, id="staged"), + pytest.param(2, True, True, id="mixed"), + pytest.param(5, True, False, id="decode"), + ], +) +def test_sparse_mqa_hisparse_lanes( + monkeypatch, num_decode_tokens, all_resident, split_q +): + """HiSparse splits a batch into at most one decode and one prefill call.""" + import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla + + num_tokens, topk_tokens, block_size, head_dim = 5, 8, 4, 6 + num_prefill = num_tokens - num_decode_tokens + # Distinct block counts, so the KV tensor a call is handed names its lane. + host = torch.zeros(2, block_size, head_dim, dtype=torch.bfloat16) + hot = torch.zeros(6, block_size, head_dim, dtype=torch.bfloat16) + staged = torch.zeros(3, block_size, head_dim, dtype=torch.bfloat16) + q = torch.arange(num_tokens, dtype=torch.float32).view(num_tokens, 1, 1) + metadata = SimpleNamespace( + num_decode_tokens=num_decode_tokens, block_size=block_size + ) + + def converted(rows, fill): + return ( + torch.full((rows, topk_tokens), fill, dtype=torch.int32), + torch.full((rows,), rows, dtype=torch.int32), + ) + + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.physical_kv_cache = MagicMock(return_value=hot) + index_group.cache = MagicMock( + return_value=SimpleNamespace(all_context_pages_resident=all_resident) + ) + # The fill value names the converter that produced a call's indices. + index_group.convert_logical_to_physical_topk = MagicMock( + return_value=converted(num_decode_tokens, 1) + if num_decode_tokens + else converted(num_tokens, 2) + ) + index_group.stage_prefill_rows = MagicMock(return_value=(staged, None, None)) + calls: list[tuple] = [] + + def convert_staged(*args, **kwargs): + # Decode rows are consumed before a later conversion can reuse their buffer. + assert len(calls) == min(num_decode_tokens, 1) + return converted(num_prefill, 3) + + monkeypatch.setattr( + sparse_mla, + "triton_convert_req_index_to_global_index", + MagicMock(side_effect=convert_staged), + ) + + def hook(rows, cache, indices, counts, *, is_decode, **kwargs): + if split_q: + assert rows[0].tolist() == rows[1].tolist() + rows = rows[0] + calls.append((rows.flatten().tolist(), cache.shape[0], indices[0, 0].item())) + return torch.full((rows.shape[0], 1, 1), 1.0 if is_decode else 2.0), None + + impl = object.__new__(FlashInferMLASparseImpl) + impl.topk_indices_buffer = torch.zeros(num_tokens, topk_tokens, dtype=torch.int32) + impl.index_group = index_group + impl.index_group_index = 0 + impl.dcp_world_size = 1 + impl._forward_mqa_kernel = hook + + out, lse = impl.forward_mqa((q, q) if split_q else q, host, metadata, None) + + q_rows = q.flatten().tolist() + expected = [] + if num_decode_tokens: + # Decode reads the hot cache through the group's conversion. + expected.append((q_rows[:num_decode_tokens], hot.shape[0], 1)) + if num_decode_tokens == 0 and all_resident: + expected.append((q_rows, hot.shape[0], 2)) + elif num_prefill: + expected.append((q_rows[num_decode_tokens:], staged.shape[0], 3)) + assert calls == expected + if not num_prefill: + # A decode-only batch returns before it looks up the prefill cache. + index_group.cache.assert_not_called() + assert lse is None + # A lane handed the whole batch would concatenate to more rows than this. + assert out.flatten().tolist() == [1.0] * num_decode_tokens + [2.0] * num_prefill + + def test_hisparse_resident_prefill_uses_attention_block_stride(): """A batch beyond the swap-row capacity must bypass swap_in. @@ -3887,6 +3942,250 @@ def test_hisparse_mixed_single_token_batch_is_not_decode_only(): assert not cache.decode_batch +def _fa4_impl(topk_indices_buffer, **fields): + impl = object.__new__(FlashAttnMLASparseFA4Impl) + impl.__dict__.update( + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + scale=1.0, + kv_cache_dtype="auto", + dcp_world_size=1, + dcp_rank=0, + need_to_return_lse_for_decode=False, + is_nope_mla=False, + _workspace_buffer=None, + bmm1_scale=None, + bmm2_scale=None, + # A non-index-producing layer does not wait on the indexer's ready event. + index_group=SparseMLAIndexGroupBuilder(topk_indices_buffer).register_layer( + False + )[0], + index_group_index=0, + ) + impl.__dict__.update(fields) + # topk_indices_buffer is a data descriptor, so it needs its own setter. + impl.topk_indices_buffer = topk_indices_buffer + return impl + + +def _fa4_inputs(counts, *, num_heads=16, topk=128, num_blocks=4, device=DEVICE_TYPE): + """Decode inputs whose token ``i`` has ``counts[i]`` valid rows, -1 padded.""" + device = torch.device(device) + block_size, kv_lora_rank, rope_dim = 64, 512, 64 + num_tokens = len(counts) + torch.manual_seed(0) + + def values(*shape): + return torch.rand(shape, dtype=torch.bfloat16, device=device) - 0.5 + + topk_indices = torch.full((num_tokens, topk), -1, dtype=torch.int32, device=device) + for tok, count in enumerate(counts): + topk_indices[tok, :count] = torch.randperm( + num_blocks * block_size, device=device + )[:count].to(torch.int32) + + return SimpleNamespace( + kv_lora_rank=kv_lora_rank, + rope_dim=rope_dim, + block_size=block_size, + topk=topk, + kv_cache=values(num_blocks, block_size, kv_lora_rank + rope_dim), + topk_indices=topk_indices, + valid_counts=torch.tensor(counts, dtype=torch.int32, device=device), + ql_nope=values(num_tokens, num_heads, kv_lora_rank), + q_pe=values(num_tokens, num_heads, rope_dim), + scale=(kv_lora_rank + rope_dim) ** -0.5, + metadata=SimpleNamespace( + req_id_per_token=torch.zeros(num_tokens, dtype=torch.int32, device=device), + block_table=torch.arange(num_blocks, dtype=torch.int32, device=device).view( + 1, num_blocks + ), + block_size=block_size, + num_decode_tokens=num_tokens, + cp_kv_cache_interleave_size=1, + ), + ) + + +def _record_fa4_kernels(monkeypatch, num_heads, *, device): + """Replace both lanes with recorders whose outputs identify the lane.""" + import flashinfer.decode + + fa4_calls: list[dict] = [] + trtllm_calls: list[dict] = [] + + def lse(rows): + return torch.arange(rows * num_heads, dtype=torch.float32, device=device).view( + rows, num_heads + ) + + def fake_fa4(**kwargs): + fa4_calls.append(kwargs) + rows = kwargs["q"].shape[0] + out = torch.ones(rows, num_heads, 512, dtype=torch.bfloat16, device=device) + return (out, lse(rows)) if kwargs.get("return_softmax_lse") else out + + def fake_trtllm(**kwargs): + trtllm_calls.append(kwargs) + rows = kwargs["query"].shape[0] + out = torch.full( + (rows, 1, num_heads, 512), 2.0, dtype=torch.bfloat16, device=device + ) + return (out, lse(rows)) if kwargs.get("return_lse") else out + + monkeypatch.setattr(flashattn_sparse_module, "flash_attn_varlen_func", fake_fa4) + monkeypatch.setattr( + flashinfer.decode, "trtllm_batch_decode_with_kv_cache_mla", fake_trtllm + ) + return SimpleNamespace(fa4=fa4_calls, trtllm=trtllm_calls, lse=lse) + + +@pytest.mark.parametrize( + "num_heads,return_lse", [(128, False), (16, True)], ids=["heads128", "heads16_lse"] +) +def test_fa4_sparse_decode_kernel_correctness(num_heads, return_lse): + if not current_platform.is_device_capability_family(100): + pytest.skip("FA4 sparse MLA requires SM 10.x") + if (reason := _fa4_cute_mla_available()) is not None: + pytest.skip(reason) + counts = [0, 1, 127, 128, 129, 2047, 2048] + inputs = _fa4_inputs(counts, num_heads=num_heads, topk=2048, num_blocks=64) + impl = _fa4_impl( + topk_indices_buffer=inputs.topk_indices, + scale=inputs.scale, + need_to_return_lse_for_decode=return_lse, + ) + q = (inputs.ql_nope, inputs.q_pe) + + with torch.inference_mode(): + out, lse = impl.forward_mqa(q, inputs.kv_cache, inputs.metadata, None) + if not return_lse: + assert lse is None + else: + out_fused, lse_fused = impl.forward_mqa( + torch.cat(q, dim=-1), inputs.kv_cache, inputs.metadata, None + ) + assert lse.shape == (len(counts), num_heads) + assert torch.equal(out, out_fused) and torch.equal(lse, lse_fused) + + kv_flat = inputs.kv_cache.view(-1, inputs.kv_lora_rank + inputs.rope_dim) + keys = kv_flat[inputs.topk_indices.clamp(min=0).long()].float() + scores = torch.einsum("thd,tkd->thk", torch.cat(q, dim=-1).float(), keys) + scores = (scores * inputs.scale).masked_fill( + (inputs.topk_indices < 0)[:, None, :], float("-inf") + ) + # The all-sentinel row softmaxes to NaN here; the kernel writes (0, -inf). + probs = torch.softmax(scores, dim=-1).nan_to_num() + torch.testing.assert_close( + out.float(), + torch.einsum("thk,tkd->thd", probs, keys[..., : inputs.kv_lora_rank]), + rtol=0.01, + atol=0.01, + ) + if return_lse: + torch.testing.assert_close( + lse, torch.logsumexp(scores, dim=-1), rtol=0, atol=5e-3 + ) + + +@pytest.mark.parametrize("is_decode", [True, False], ids=["decode", "prefill"]) +def test_fa4_sparse_hook_routes_lanes(monkeypatch, is_decode): + """The hook runs a decode batch on FA4 and every other batch on trtllm-gen.""" + import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla + + counts = [1, 2, 3, 4, 5, 6, 7] + num_tokens, dcp = len(counts), 2 + # Under DCP the hook gets the heads of every DCP rank, not just its own. + kernel_heads = 16 * dcp + inputs = _fa4_inputs(counts, num_heads=kernel_heads, device="cpu") + calls = _record_fa4_kernels(monkeypatch, kernel_heads, device="cpu") + # DCP without PCP (the gate rejects the pair) always fuses the query. + q = torch.cat((inputs.ql_nope, inputs.q_pe), dim=-1) + impl = _fa4_impl( + # One row longer than the batch: the varlen-scalar slices become observable. + topk_indices_buffer=inputs.topk_indices.new_zeros(num_tokens + 1, inputs.topk), + dcp_world_size=dcp, + need_to_return_lse_for_decode=True, + ) + impl._convert_logical_to_physical_topk = MagicMock() + dcp_filter = MagicMock(return_value=(inputs.topk_indices, inputs.valid_counts)) + monkeypatch.setattr(sparse_mla, "triton_filter_and_convert_dcp_index", dcp_filter) + inputs.metadata.num_decode_tokens = num_tokens if is_decode else 1 + + with torch.inference_mode(): + out, lse = impl.forward_mqa(q, inputs.kv_cache, inputs.metadata, None) + + dcp_filter.assert_called_once() + impl._convert_logical_to_physical_topk.assert_not_called() + assert out.shape == (num_tokens, kernel_heads, 512) + if is_decode: + (fa4,) = calls.fa4 + assert calls.trtllm == [] + assert fa4["return_softmax_lse"] is True + assert fa4["cu_seqlens_q"].tolist() == list(range(num_tokens + 1)) + assert (out == 1.0).all() + # FA4's LSE is already natural-log. + torch.testing.assert_close(lse, calls.lse(num_tokens), rtol=0, atol=0) + else: + (trtllm,) = calls.trtllm + assert calls.fa4 == [] + assert trtllm["return_lse"] is True + assert (out == 2.0).all() + # trtllm-gen's LSE is base 2; the DCP merge expects natural log. + expected = calls.lse(num_tokens) * math.log(2.0) + torch.testing.assert_close(lse, expected, rtol=0, atol=1e-5) + + +@pytest.mark.parametrize( + "config,expected", + [ + ({}, None), + ({"heads": 128}, None), + ({"heads": 12}, "heads"), + ({"dcp": 2}, None), + ({"dcp": 8}, "heads"), + ({"index_topk": 192}, "index_topk"), + ({"rope": 32}, "(kv_lora_rank, qk_rope_head_dim)"), + ({"flashinfer": False}, "FlashInfer"), + ], +) +def test_fa4_sparse_gate(monkeypatch, config, expected): + """``validate_configuration`` on a fixed SM100: one row per gate reason.""" + has_flashinfer = config.get("flashinfer", True) + monkeypatch.setattr("vllm.utils.flashinfer.has_flashinfer", lambda: has_flashinfer) + monkeypatch.setattr( + flashattn_sparse_module, "_fa4_cute_mla_available", lambda: None + ) + dcp = config.get("dcp", 1) + vllm_config = _build_sparse_dcp_vllm_config(config.get("heads", 16), dcp) + hf_text_config = vllm_config.model_config.hf_text_config + hf_text_config.index_topk = config.get("index_topk", 128) + hf_text_config.qk_rope_head_dim = config.get("rope", 64) + + with set_current_vllm_config(vllm_config): + reasons = FlashAttnMLASparseFA4Backend.validate_configuration( + head_size=576, + dtype=torch.bfloat16, + kv_cache_dtype="auto", + block_size=64, + use_mla=True, + has_sink=False, + use_sparse=True, + use_mm_prefix=False, + use_per_head_quant_scales=False, + device_capability=DeviceCapability(10, 0), + attn_type="decoder", + use_dcp=dcp > 1, + ) + + if expected is None: + assert reasons == [] + else: + (reason,) = reasons + assert expected in reason + + def test_flashmla_cache_dtype_aliases_use_ds_layout(): # kv-cache dtype aliases are canonicalized to fp8_ds_mla before the layer # stores kv_cache_dtype, so they cannot bypass the gate. diff --git a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py index 45ba1ac0dc55..b322acf0d5ee 100644 --- a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py +++ b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py @@ -13,15 +13,12 @@ from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache from vllm.platforms import current_platform from vllm.v1.attention.backend import MultipleOf -from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( - FlashInferMLASparseTRTLLMBackend, -) -from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseBackend from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend from vllm.v1.attention.backends.mla.sparse_utils import ( _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL, flat_kv_row_view, ) +from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts from vllm.v1.core.kv_cache_planning import get_kv_cache_config_from_groups from vllm.v1.kv_cache_interface import KVCacheGroupSpec, UniformTypeKVCacheSpecs @@ -49,17 +46,27 @@ @pytest.mark.parametrize( "backend,cache_dtype,sm100,stride_alignment", [ - (FlashInferMLASparseTRTLLMBackend, "auto", True, 32 * 1152), - (FlashInferMLASparseTRTLLMBackend, "fp8", True, 32 * 576), - (FlashMLASparseBackend, "auto", True, 1152), - (FlashMLASparseBackend, "fp8_ds_mla", False, MultipleOf(64)), - (FlashMLASparseBackend, "fp8_ds_mla", True, 656), - (FlashMLASparseBackend, "nvfp4_ds_mla", True, 352), + (AttentionBackendEnum.FLASHINFER_MLA_SPARSE, "auto", True, 32 * 1152), + (AttentionBackendEnum.FLASHINFER_MLA_SPARSE, "fp8", True, 32 * 576), + pytest.param( + AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE_FA4, + "auto", + True, + 64 * 1152, + marks=pytest.mark.skipif( + not current_platform.is_cuda(), reason="CUDA only" + ), + ), + (AttentionBackendEnum.FLASHMLA_SPARSE, "auto", True, 1152), + (AttentionBackendEnum.FLASHMLA_SPARSE, "fp8_ds_mla", False, MultipleOf(64)), + (AttentionBackendEnum.FLASHMLA_SPARSE, "fp8_ds_mla", True, 656), + (AttentionBackendEnum.FLASHMLA_SPARSE, "nvfp4_ds_mla", True, 352), ], ) def test_allocation_and_warmup_follow_addressing_mode( monkeypatch, layout, backend, cache_dtype, sm100, stride_alignment, device ): + backend = backend.get_class() monkeypatch.setattr( current_platform, "is_device_capability_family", lambda family: sm100 ) diff --git a/tests/v1/e2e/general/test_hisparse.py b/tests/v1/e2e/general/test_hisparse.py index 2409b5cb690b..d9e1fabf03f3 100644 --- a/tests/v1/e2e/general/test_hisparse.py +++ b/tests/v1/e2e/general/test_hisparse.py @@ -83,13 +83,20 @@ def _get_hisparse_worker(runner: VllmRunner) -> HiSparseConnectorWorker: not current_platform.is_cuda(), reason="HiSparse requires NVIDIA CUDA" ) @pytest.mark.parametrize( - "with_offloading", [False, True], ids=["standalone", "offload"] + "with_offloading,attention_backend", + [ + pytest.param(False, None, id="standalone"), + pytest.param(True, None, id="offload"), + # Offloading does not touch the attention kernels: FA4 runs standalone. + pytest.param(False, "FLASH_ATTN_MLA_SPARSE_FA4", id="standalone-fa4"), + ], ) @fork_new_process_for_each_test def test_hisparse_spill_and_prefix_restore( monkeypatch: pytest.MonkeyPatch, vllm_runner: type[VllmRunner], with_offloading: bool, + attention_backend: str | None, ): """Spilled prefixes restore and FULL-graph decode writes reach host KV. @@ -101,6 +108,8 @@ def test_hisparse_spill_and_prefix_restore( capability = current_platform.get_device_capability() if capability is None or capability.major < 9: pytest.skip("Sparse MLA requires Hopper or newer") + if attention_backend and not current_platform.is_device_capability_family(100): + pytest.skip("FA4 sparse MLA requires SM 10.x") monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip") @@ -142,6 +151,7 @@ def test_hisparse_spill_and_prefix_restore( ) ), kv_transfer_config=kv_transfer_config, + attention_backend=attention_backend, block_size=64, max_model_len=320, max_num_batched_tokens=1024, diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index db9e8dbc88ca..806faba50774 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -661,6 +661,7 @@ def __init__( "FLASHMLA_SPARSE", "FLASHINFER_MLA_SPARSE", "FLASHINFER_MLA_SPARSE_SM120", + "FLASH_ATTN_MLA_SPARSE_FA4", "DEEPSEEK_V32_INDEXER", ): from vllm.v1.attention.backends.mla.compressor_utils import ( @@ -751,7 +752,7 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: def _uses_flat_kv_cache(self) -> bool: backend = self.attn_backend.get_name() - return backend == "FLASHINFER_MLA_SPARSE" or ( + return backend in ("FLASHINFER_MLA_SPARSE", "FLASH_ATTN_MLA_SPARSE_FA4") or ( backend == "FLASHMLA_SPARSE" and self.kv_cache_dtype not in ("fp8_ds_mla", "nvfp4_ds_mla") ) @@ -1857,6 +1858,9 @@ def _use_masked_mha( seq_len: int, has_context: bool, ) -> bool: + if backend_name == "FLASH_ATTN_MLA_SPARSE_FA4": + # Its prefill batches run on FlashInfer's kernel, so it shares those rows. + backend_name = "FLASHINFER_MLA_SPARSE" thresholds = ( _MASKED_MHA_THRESHOLDS.get((qk_head_dim, v_head_dim), {}) .get(backend_name, {}) diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 37883593421a..a8fa74f91f62 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -36,12 +36,18 @@ is_quantized_kv_cache, np_to_pinned_tensor, ) -from vllm.v1.attention.backend import AttentionMetadataBuilder +from vllm.v1.attention.backend import AttentionLayer, AttentionMetadataBuilder from vllm.v1.attention.backends.fa_utils import get_flash_attn_version from vllm.v1.attention.backends.mla.index_group import ( + HiSparseMLAIndexGroup, SparseMLAIndexGroup, SparseMLAIndexGroupBuilder, ) +from vllm.v1.attention.backends.mla.sparse_utils import ( + flat_kv_row_view, + triton_convert_req_index_to_global_index, + triton_filter_and_convert_dcp_index, +) from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.attention.ops.dcp import MLADCPManager from vllm.v1.attention.ops.merge_attn_states import merge_attn_states @@ -630,6 +636,9 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic """Sparse MLA base with dense and masked-MHA prefill paths.""" is_sparse = True + # Reuse the index group's top-k conversion only for uniform decode batches + # of multi-layer groups; other batches convert per layer. + group_conversion_decode_only: ClassVar[bool] = False def __init__( self, @@ -728,6 +737,154 @@ def _convert_logical_to_physical_topk( return_valid_counts=return_valid_counts, ) + def forward_mqa( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: T, + layer: AttentionLayer, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + assert self.topk_indices_buffer is not None + num_tokens = q[0].shape[0] if isinstance(q, tuple) else q.shape[0] + topk = self.topk_indices_buffer[:num_tokens] + num_decode_tokens = attn_metadata.num_decode_tokens + block_size = attn_metadata.block_size + index_group = self.index_group + if not isinstance(index_group, HiSparseMLAIndexGroup): + _, block_stride_rows = flat_kv_row_view(kv_c_and_k_pe_cache, block_size) + if self.dcp_world_size > 1: + physical_topk, valid_counts = triton_filter_and_convert_dcp_index( + attn_metadata.req_id_per_token[:num_tokens], + attn_metadata.block_table, + topk, + dcp_size=self.dcp_world_size, + dcp_rank=self.dcp_rank, + cp_kv_cache_interleave_size=attn_metadata.cp_kv_cache_interleave_size, + BLOCK_SIZE=block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, + NUM_TOPK_TOKENS=topk.shape[1], + return_valid_counts=True, + ) + elif self.group_conversion_decode_only and not ( + index_group is not None + and index_group.num_layers > 1 + and num_decode_tokens >= num_tokens + ): + physical_topk, valid_counts = triton_convert_req_index_to_global_index( + attn_metadata.req_id_per_token[:num_tokens], + attn_metadata.block_table, + topk, + BLOCK_SIZE=block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, + NUM_TOPK_TOKENS=topk.shape[1], + return_valid_counts=True, + ) + else: + physical_topk, valid_counts = self._convert_logical_to_physical_topk( + topk, + attn_metadata, + block_stride_rows=block_stride_rows, + return_valid_counts=True, + ) + return self._forward_mqa_kernel( + q, + kv_c_and_k_pe_cache, + physical_topk, + valid_counts, + layer=layer, + block_size=block_size, + is_decode=num_decode_tokens >= num_tokens, + ) + + assert self.dcp_world_size == 1 + layer_index = self.index_group_index + hot_cache = index_group.physical_kv_cache(layer_index).view( + kv_c_and_k_pe_cache.dtype + ) + assert hot_cache.is_contiguous(), "HiSparse hot buffer must be contiguous" + decode_out = decode_lse = None + if num_decode_tokens: + physical_topk, valid_counts = index_group.convert_logical_to_physical_topk( + layer_index, + topk[:num_decode_tokens], + attn_metadata, + block_stride_rows=None, + return_valid_counts=True, + ) + q_decode = ( + (q[0][:num_decode_tokens], q[1][:num_decode_tokens]) + if isinstance(q, tuple) + else q[:num_decode_tokens] + ) + decode_out, decode_lse = self._forward_mqa_kernel( + q_decode, + hot_cache, + physical_topk, + valid_counts, + layer=layer, + block_size=block_size, + is_decode=True, + ) + if num_decode_tokens >= num_tokens: + return decode_out, decode_lse + + cache = index_group.cache(layer_index) + if num_decode_tokens == 0 and cache.all_context_pages_resident: + physical_topk, valid_counts = index_group.convert_logical_to_physical_topk( + layer_index, + topk, + attn_metadata, + block_stride_rows=None, + return_valid_counts=True, + ) + prefill_cache = hot_cache + else: + prefill_cache, block_table, req_ids = index_group.stage_prefill_rows( + layer_index, kv_c_and_k_pe_cache, attn_metadata + ) + physical_topk, valid_counts = triton_convert_req_index_to_global_index( + req_ids, + block_table, + topk[num_decode_tokens:], + BLOCK_SIZE=block_size, + NUM_TOPK_TOKENS=topk.shape[1], + return_valid_counts=True, + ) + q_prefill = ( + (q[0][num_decode_tokens:], q[1][num_decode_tokens:]) + if isinstance(q, tuple) + else q[num_decode_tokens:] + ) + prefill_out, prefill_lse = self._forward_mqa_kernel( + q_prefill, + prefill_cache, + physical_topk, + valid_counts, + layer=layer, + block_size=block_size, + is_decode=False, + ) + if decode_out is None: + return prefill_out, prefill_lse + out = torch.cat((decode_out, prefill_out)) + if decode_lse is None: + return out, None + assert prefill_lse is not None + return out, torch.cat((decode_lse, prefill_lse)) + + def _forward_mqa_kernel( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_cache: torch.Tensor, + topk_indices: torch.Tensor, + valid_counts: torch.Tensor, + *, + layer: AttentionLayer, + block_size: int, + is_decode: bool, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + raise NotImplementedError + @staticmethod def masked_mha_workspace_fits(prefill: MLACommonPrefillMetadata) -> bool: """Whether this prefill batch's top-k masks fit the workspace.""" diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 401a98f66feb..7a28dcbc96d7 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -105,9 +105,10 @@ def _get_backend_priorities( ] else: # BF16 KV Cache - # Prefer FlashInfer at low head counts (FlashMLA uses padding) + # Prefer FA4, then FlashInfer at low head counts (FlashMLA uses padding) if num_heads is not None and num_heads <= 16: sparse_backends = [ + AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE_FA4, AttentionBackendEnum.FLASHINFER_MLA_SPARSE, AttentionBackendEnum.FLASHMLA_SPARSE, ] diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index 0016eafedc8a..47347c8fe84e 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math from dataclasses import dataclass +from functools import cache from typing import Any, ClassVar import torch @@ -22,11 +24,14 @@ MultipleOf, ) from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla -from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup +from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( + FlashInferMLASparseImpl, + FlashInferMLASparseTRTLLMBackend, + FlashInferMLASparseTRTLLMMetadataBuilder, +) from vllm.v1.attention.backends.mla.sparse_utils import ( align_blocks_to_rows, flat_kv_row_view, - triton_convert_req_index_to_global_index, ) from vllm.v1.attention.ops.metadata import compute_token_to_req_indices from vllm.v1.kv_cache_interface import AttentionSpec @@ -153,6 +158,8 @@ def update_draft_decode_metadata(self, metadata: SparseMLACommonMetadata) -> Non class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]): + group_conversion_decode_only: ClassVar[bool] = True + def __init__( self, num_heads: int, @@ -205,144 +212,23 @@ def __init__( ) self.supports_quant_query_input = False - def forward_mqa( + def _forward_mqa_kernel( self, q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - kv_c_and_k_pe_cache: torch.Tensor, - attn_metadata: FlashAttnMLASparseMetadata, + kv_cache: torch.Tensor, + topk_indices: torch.Tensor, + valid_counts: torch.Tensor, + *, layer: AttentionLayer, + block_size: int, + is_decode: bool, ) -> tuple[torch.Tensor, torch.Tensor | None]: if not isinstance(q, tuple): raise NotImplementedError( "FlashAttnMLASparseImpl expects split (q_nope, q_rope) input." ) q_nope, q_rope = q - num_actual_toks = q_rope.shape[0] - - assert self.topk_indices_buffer is not None - topk_indices = self.topk_indices_buffer[:num_actual_toks] - index_group = self.index_group - if isinstance(index_group, HiSparseMLAIndexGroup): - num_decode_tokens = attn_metadata.num_decode_tokens - outputs = [] - if num_decode_tokens: - physical_topk, valid_counts = ( - index_group.convert_logical_to_physical_topk( - self.index_group_index, - topk_indices[:num_decode_tokens], - attn_metadata, - block_stride_rows=None, - return_valid_counts=True, - ) - ) - outputs.append( - self._run_mqa_kernel( - q_nope[:num_decode_tokens], - q_rope[:num_decode_tokens], - index_group.physical_kv_cache(self.index_group_index).view( - kv_c_and_k_pe_cache.dtype - ), - physical_topk, - valid_counts, - attn_metadata.block_size, - ) - ) - if num_decode_tokens < num_actual_toks: - cache = index_group.cache(self.index_group_index) - if num_decode_tokens == 0 and cache.all_context_pages_resident: - physical_topk, valid_counts = ( - index_group.convert_logical_to_physical_topk( - self.index_group_index, - topk_indices, - attn_metadata, - block_stride_rows=None, - return_valid_counts=True, - ) - ) - prefill_cache = index_group.physical_kv_cache( - self.index_group_index - ).view(kv_c_and_k_pe_cache.dtype) - else: - prefill_cache, block_table, req_ids = ( - index_group.stage_prefill_rows( - self.index_group_index, - kv_c_and_k_pe_cache, - attn_metadata, - ) - ) - physical_topk, valid_counts = ( - triton_convert_req_index_to_global_index( - req_ids, - block_table, - topk_indices[num_decode_tokens:], - BLOCK_SIZE=attn_metadata.block_size, - NUM_TOPK_TOKENS=topk_indices.shape[1], - return_valid_counts=True, - ) - ) - outputs.append( - self._run_mqa_kernel( - q_nope[num_decode_tokens:], - q_rope[num_decode_tokens:], - prefill_cache, - physical_topk, - valid_counts, - attn_metadata.block_size, - ) - ) - return torch.cat(outputs) if len(outputs) > 1 else outputs[0], None - - kv_rows, block_stride_rows = flat_kv_row_view( - kv_c_and_k_pe_cache, attn_metadata.block_size - ) - if ( - index_group is not None - and index_group.num_layers > 1 - and attn_metadata.num_decode_tokens == num_actual_toks - ): - topk_indices, valid_counts = self._convert_logical_to_physical_topk( - topk_indices, - attn_metadata, - block_stride_rows=block_stride_rows, - return_valid_counts=True, - ) - else: - topk_indices, valid_counts = triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token[:num_actual_toks], - attn_metadata.block_table, - topk_indices, - BLOCK_SIZE=attn_metadata.block_size, - BLOCK_STRIDE_ROWS=block_stride_rows, - NUM_TOPK_TOKENS=topk_indices.shape[1], - return_valid_counts=True, - ) - return ( - self._run_mqa_kernel( - q_nope, - q_rope, - kv_rows, - topk_indices, - valid_counts, - attn_metadata.block_size, - cache_is_flat=True, - ), - None, - ) - - def _run_mqa_kernel( - self, - q_nope: torch.Tensor, - q_rope: torch.Tensor, - kv_cache: torch.Tensor, - topk_indices: torch.Tensor, - valid_counts: torch.Tensor, - block_size: int, - *, - cache_is_flat: bool = False, - ) -> torch.Tensor: - kv_rows = ( - kv_cache if cache_is_flat else flat_kv_row_view(kv_cache, block_size)[0] - ) + kv_rows, _ = flat_kv_row_view(kv_cache, block_size) cu_seqlens_q = self.cu_seqlens_q_buffer[: q_rope.shape[0] + 1] v_cache = kv_rows[:, : self.kv_lora_rank].unsqueeze(1).unsqueeze(1) @@ -370,4 +256,167 @@ def _run_mqa_kernel( causal=True, fa_version=3, ) - return out + return out, None + + +def _fa4_cute_mla_available() -> str | None: + """Reason the FA4 cute-DSL MLA entry point is unusable, or None.""" + from vllm.vllm_flash_attn import flash_attn_interface as fa + + if not fa.is_fa_version_supported(4): + return f"FA4 unavailable: {fa.fa_version_unsupported_reason(4)}" + try: + from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd # noqa: F401 + except Exception as e: # cute-DSL / cutlass import chain + return f"FA4 cute-DSL entry point failed to import: {e!r}" + return None + + +class FlashAttnMLASparseFA4Backend(FlashInferMLASparseTRTLLMBackend): + """FA4 for is_decode kernel calls; trtllm-gen for the rest.""" + + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] + # The qv kernel asserts every descale is None: BF16 KV cache only. + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"] + + @staticmethod + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: + return [MultipleOf(64)] + + @staticmethod + def get_name() -> str: + return "FLASH_ATTN_MLA_SPARSE_FA4" + + @staticmethod + def get_builder_cls() -> type["FlashAttnMLASparseFA4MetadataBuilder"]: + return FlashAttnMLASparseFA4MetadataBuilder + + @staticmethod + def get_impl_cls() -> type[MLAAttentionImpl[Any]]: + return FlashAttnMLASparseFA4Impl + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + # 512 latent + 64 RoPE; split back apart for the qv kernel. + return [576] + + @classmethod + def supports_combination(cls, *args: Any, **kwargs: Any) -> str | None: + from vllm.config import get_current_vllm_config_or_none + from vllm.utils.flashinfer import has_flashinfer + + if not has_flashinfer(): + return "FA4 sparse MLA runs its prefill batches on FlashInfer's kernel" + if (reason := _fa4_cute_mla_available()) is not None: + return reason + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return None + if vllm_config.model_config is not None: + hf_config = vllm_config.model_config.hf_text_config + topk = getattr(hf_config, "index_topk", None) + # gather_kv_indices' last dim must be a whole number of n-tiles. + if topk is None or topk % 128 != 0: + return f"FA4 sparse MLA requires index_topk % 128 == 0, got {topk}" + dims = ( + getattr(hf_config, "kv_lora_rank", None), + getattr(hf_config, "qk_rope_head_dim", None), + ) + if dims != (512, 64): + return ( + "FA4 sparse MLA requires (kv_lora_rank, qk_rope_head_dim) " + f"== (512, 64), got {dims}" + ) + # Under DCP the query is all-gathered: the kernel sees heads * dcp_size. + dcp_size = vllm_config.parallel_config.decode_context_parallel_size + num_heads = vllm_config.model_config.get_num_attention_heads( + vllm_config.parallel_config + ) + # 128 gathered heads take FA's prefill kernel, unvalidated under DCP. + head_counts = (8, 16, 32, 64) if dcp_size > 1 else (8, 16, 32, 64, 128) + if num_heads * dcp_size not in head_counts: + return ( + f"FA4 sparse MLA requires {head_counts} gathered query heads, " + f"got num_heads={num_heads} * dcp_size={dcp_size}" + ) + return super().supports_combination(*args, **kwargs) + + +class FlashAttnMLASparseFA4MetadataBuilder(FlashInferMLASparseTRTLLMMetadataBuilder): + """The trtllm-gen builder, for the prefill lane's DCP workspace pre-size. + + Its UNIFORM_BATCH support and varlen decode bound also fit FA4, which runs + one query row per token and bakes the row count into its varlen scalars. + """ + + +@cache +def _fa4_varlen_scalars( + num_tokens: int, num_kv_rows: int, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Per-token varlen scalars; never evict, CUDA graphs bake in the pointers.""" + cu_seqlens_q = torch.arange(num_tokens + 1, dtype=torch.int32, device=device) + cu_seqlens_k = torch.zeros(num_tokens + 1, dtype=torch.int32, device=device) + seqused_k = torch.full((num_tokens,), num_kv_rows, dtype=torch.int32, device=device) + return cu_seqlens_q, cu_seqlens_k, seqused_k + + +class FlashAttnMLASparseFA4Impl(FlashInferMLASparseImpl): + # FA4 emits a natural-log LSE; prefill converts trtllm-gen's log2 LSE. + lse_base_on_e: bool = True + + def _forward_mqa_kernel( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_cache: torch.Tensor, + topk_indices: torch.Tensor, + valid_counts: torch.Tensor, + *, + layer: AttentionLayer, + block_size: int, + is_decode: bool, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if not is_decode: + out, lse = super()._forward_mqa_kernel( + q, + kv_cache, + topk_indices, + valid_counts, + layer=layer, + block_size=block_size, + is_decode=is_decode, + ) + return out, None if lse is None else lse * math.log(2.0) + if isinstance(q, tuple): + ql_nope, q_pe = q + else: + # 16B-aligned halves of the fused query; FA4 reads them in place. + ql_nope, q_pe = q[..., : self.kv_lora_rank], q[..., self.kv_lora_rank :] + kv_rows, _ = flat_kv_row_view(kv_cache, block_size) + num_tokens = q_pe.shape[0] + num_kv_rows = kv_rows.shape[0] + assert self.topk_indices_buffer is not None + cu_seqlens_q, cu_seqlens_k, seqused_k = _fa4_varlen_scalars( + self.topk_indices_buffer.shape[0], num_kv_rows, kv_rows.device + ) + # An all-sentinel row yields (0, -inf), the DCP merge identity; no masking. + kernel_out = flash_attn_varlen_func( + q=q_pe, + k=kv_rows[:, self.kv_lora_rank :].unsqueeze(1), + v=kv_rows[:, : self.kv_lora_rank].unsqueeze(1), + q_v=ql_nope, + max_seqlen_q=1, + cu_seqlens_q=cu_seqlens_q[: num_tokens + 1], + max_seqlen_k=num_kv_rows, + cu_seqlens_k=cu_seqlens_k[: num_tokens + 1], + seqused_k=seqused_k[:num_tokens], + gather_kv_indices=topk_indices, + # Part of FA4's compile key: pass it unconditionally. + gather_kv_valid_length=valid_counts, + softmax_scale=self.scale, + # Causality is in the top-k list; causal=True would clamp flat rows. + causal=False, + fa_version=4, + return_softmax_lse=self.need_to_return_lse_for_decode, + ) + return kernel_out if self.need_to_return_lse_for_decode else (kernel_out, None) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 21eb66aae490..d673868975ab 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -34,8 +34,6 @@ align_blocks_to_rows, flat_kv_row_view, prepare_sparse_mla_safe_lengths, - triton_convert_req_index_to_global_index, - triton_filter_and_convert_dcp_index, ) from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec @@ -129,14 +127,12 @@ def supports_combination( and parallel_config.decode_context_parallel_size > 1 ): return ( - "FLASHINFER_MLA_SPARSE does not support combined PCP+DCP; " + f"{cls.get_name()} does not support combined PCP+DCP; " "use FLASHMLA_SPARSE, which gathers each DCP KV shard before " "running the rank-local PCP prefill queries" ) if kv_cache_dtype == "fp8_ds_mla": - return ( - "FLASHINFER_MLA_SPARSE SM10 does not support fp8_ds_mla kv-cache dtype" - ) + return f"{cls.get_name()} SM10 does not support fp8_ds_mla kv-cache dtype" # FlashInfer MLA sparse SM10 kernel requires qk_nope_head_dim in [128, 192]. if vllm_config.model_config is not None: @@ -492,135 +488,6 @@ def __init__( # for query and kv_cache (mixed bf16+fp8 is not supported). self.supports_quant_query_input = True - def forward_mqa( - self, - q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - kv_c_and_k_pe_cache: torch.Tensor, - attn_metadata: FlashInferMLASparseMetadata, - layer: AttentionLayer, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - if isinstance(q, tuple): - ql_nope, q_pe = q - if q_pe.shape[-1] == 0 and ql_nope.is_contiguous(): - q = ql_nope - else: - q = torch.cat(q, dim=-1) - - num_actual_toks = q.shape[0] - - assert self.topk_indices_buffer is not None - topk_indices = self.topk_indices_buffer[:num_actual_toks] - - self._prepare_mqa_kernel(layer, q.device) - - index_group = self.index_group - if isinstance(index_group, HiSparseMLAIndexGroup): - num_decode_tokens = attn_metadata.num_decode_tokens - decode_out: torch.Tensor | None = None - decode_lse: torch.Tensor | None = None - if num_decode_tokens > 0: - physical_topk, valid_counts = ( - index_group.convert_logical_to_physical_topk( - self.index_group_index, - topk_indices[:num_decode_tokens], - attn_metadata, - block_stride_rows=None, - return_valid_counts=True, - ) - ) - decode_out, decode_lse = self._run_mqa_kernel( - q[:num_decode_tokens], - index_group.physical_kv_cache(self.index_group_index).view( - kv_c_and_k_pe_cache.dtype - ), - physical_topk, - valid_counts, - ) - if num_decode_tokens == num_actual_toks: - return decode_out, decode_lse - - cache = index_group.cache(self.index_group_index) - if num_decode_tokens == 0 and cache.all_context_pages_resident: - physical_topk, valid_counts = ( - index_group.convert_logical_to_physical_topk( - self.index_group_index, - topk_indices, - attn_metadata, - block_stride_rows=None, - return_valid_counts=True, - ) - ) - return self._run_mqa_kernel( - q, - index_group.physical_kv_cache(self.index_group_index).view( - kv_c_and_k_pe_cache.dtype - ), - physical_topk, - valid_counts, - ) - - prefill_cache, block_table, req_ids = index_group.stage_prefill_rows( - self.index_group_index, kv_c_and_k_pe_cache, attn_metadata - ) - prefill_indices, prefill_lens = triton_convert_req_index_to_global_index( - req_ids, - block_table, - topk_indices[num_decode_tokens:], - BLOCK_SIZE=attn_metadata.block_size, - NUM_TOPK_TOKENS=topk_indices.shape[1], - return_valid_counts=True, - ) - prefill_out, prefill_lse = self._run_mqa_kernel( - q[num_decode_tokens:], - prefill_cache, - prefill_indices, - prefill_lens, - ) - if decode_out is None: - return prefill_out, prefill_lse - output = torch.cat((decode_out, prefill_out)) - if decode_lse is None: - return output, None - assert prefill_lse is not None - return output, torch.cat((decode_lse, prefill_lse)) - - kv_rows, block_stride_rows = flat_kv_row_view( - kv_c_and_k_pe_cache, attn_metadata.block_size - ) - - if self.dcp_world_size > 1: - topk_indices_physical, seq_lens = triton_filter_and_convert_dcp_index( - attn_metadata.req_id_per_token[:num_actual_toks], - attn_metadata.block_table, - topk_indices, - dcp_size=self.dcp_world_size, - dcp_rank=self.dcp_rank, - cp_kv_cache_interleave_size=(attn_metadata.cp_kv_cache_interleave_size), - BLOCK_SIZE=attn_metadata.block_size, - BLOCK_STRIDE_ROWS=block_stride_rows, - NUM_TOPK_TOKENS=topk_indices.shape[1], - return_valid_counts=True, - ) - else: - topk_indices_physical, seq_lens = self._convert_logical_to_physical_topk( - topk_indices, - attn_metadata, - block_stride_rows=block_stride_rows, - return_valid_counts=True, - ) - - return self._run_mqa_kernel( - q, - # Block sizes and strides have been aligned to 32 rows. - kv_rows.view( - -1, - math.gcd(block_stride_rows, attn_metadata.block_size, 64), - kv_rows.shape[-1], - ), - topk_indices_physical, - seq_lens, - ) - def _prepare_mqa_kernel( self, layer: AttentionLayer, @@ -648,7 +515,6 @@ def autotune_hisparse_decode(self, layer: AttentionLayer) -> None: kv_cache = runtime.hot.attention_cache num_tokens = runtime.max_num_reqs topk_tokens = self.topk_indices_buffer.shape[1] - self._prepare_mqa_kernel(layer, kv_cache.device) q_dtype = ( current_platform.fp8_dtype() @@ -679,22 +545,48 @@ def autotune_hisparse_decode(self, layer: AttentionLayer) -> None: dtype=torch.int32, device=kv_cache.device, ) - self._run_mqa_kernel(q, kv_cache, topk_indices, seq_lens) + self._forward_mqa_kernel( + q, + kv_cache, + topk_indices, + seq_lens, + layer=layer, + block_size=runtime.hot.block_size, + is_decode=True, + ) - def _run_mqa_kernel( + def _forward_mqa_kernel( self, - q: torch.Tensor, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], kv_cache: torch.Tensor, topk_indices: torch.Tensor, seq_lens: torch.Tensor, + *, + layer: AttentionLayer, + block_size: int, + is_decode: bool, ) -> tuple[torch.Tensor, torch.Tensor | None]: + if isinstance(q, tuple): + ql_nope, q_pe = q + if q_pe.shape[-1] == 0 and ql_nope.is_contiguous(): + q = ql_nope + else: + q = torch.cat(q, dim=-1) + + self._prepare_mqa_kernel(layer, q.device) assert self._workspace_buffer is not None assert self.bmm1_scale is not None assert self.bmm2_scale is not None from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla - kv_cache = kv_cache.view(q.dtype) + kv_rows, block_stride_rows = flat_kv_row_view( + kv_cache.view(q.dtype), block_size + ) + # Block sizes and strides have been aligned to 32 rows. + kv_cache = kv_rows.view( + -1, math.gcd(block_stride_rows, block_size, 64), kv_rows.shape[-1] + ) # Single-token sparse decode. trtllm-gen requires the q_len_per_request # dim, but the sparse attention mask is fully per-token (each query token diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 207c112a384c..e1c4c8d06741 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -125,6 +125,10 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): FLASH_ATTN_MLA_SPARSE = ( "vllm.v1.attention.backends.mla.flashattn_mla_sparse.FlashAttnMLASparseBackend" ) + FLASH_ATTN_MLA_SPARSE_FA4 = ( + "vllm.v1.attention.backends.mla.flashattn_mla_sparse." + "FlashAttnMLASparseFA4Backend" + ) MINIMAX_M3_SPARSE = ( "vllm.models.minimax_m3.common.sparse_attention.MiniMaxM3SparseBackend" ) diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index 3c4c0d8948be..c746dd0a54a8 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/vllm/vllm_flash_attn/flash_attn_interface.py @@ -214,6 +214,8 @@ def flash_attn_varlen_func( aux_tensors=None, aux_tensor_leading_dims=None, dynamic_causal: "torch.Tensor | None" = None, + gather_kv_indices=None, + gather_kv_valid_length=None, ): """dropout_p should be set to 0.0 during evaluation Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads @@ -262,9 +264,10 @@ def flash_attn_varlen_func( testing only. The returned probabilities are not guaranteed to be correct (they might not have the right scaling). seqused_k: (batch_size,), dtype torch.int32. Actual number of key tokens per - sequence. Mutually exclusive with cu_seqlens_k, and required when + sequence. Mutually exclusive with cu_seqlens_k except on FA4, which reads + cu_seqlens_k as the KV offset and seqused_k as the length. Required when block_table is given. - q_v: optional value tensor fused with q. FA3 only. + q_v: optional value tensor fused with q. FA3 and FA4 only. block_table: (batch_size, max_blocks_per_seq), dtype torch.int32. Paged KV cache block table. Requires seqused_k. return_softmax_lse: bool. Whether to also return softmax_lse. @@ -289,18 +292,21 @@ def flash_attn_varlen_func( aux_tensors: auxiliary tensors consumed by mask_mod. FA4 only. aux_tensor_leading_dims: leading dimensions of each entry in aux_tensors. dynamic_causal: optional per-sequence causal offsets. FA4 only. + gather_kv_indices: (total_q, k) int32 KV rows per query token, `-1` padded, + k a multiple of 128. Requires q_v, excludes block_table. FA4 only. + gather_kv_valid_length: (total_q,) int32, real entries per row. FA4 only. Return: out: (total, nheads, headdim). - softmax_lse [optional, if return_softmax_lse=True]: (nheads, total_q_seqlen). The - logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax - normalization factor). + softmax_lse [optional, if return_softmax_lse=True]: (nheads, total_q_seqlen), or + (total_q, nheads) on FA4's q_v path. The logsumexp of each row of the matrix + QK^T * scaling (e.g., log of the softmax normalization factor). """ assert cu_seqlens_k is not None or seqused_k is not None, ( "cu_seqlens_k or seqused_k must be provided" ) - assert cu_seqlens_k is None or seqused_k is None, ( + assert fa_version == 4 or cu_seqlens_k is None or seqused_k is None, ( "cu_seqlens_k and seqused_k cannot be provided at the same time" ) assert block_table is None or seqused_k is not None, ( @@ -311,6 +317,9 @@ def flash_attn_varlen_func( f"Fused FP8 output (output_scale) is only supported by FA4, " f"got fa_version={fa_version}" ) + assert fa_version == 4 or ( + gather_kv_indices is None and gather_kv_valid_length is None + ), f"gather_kv_* is only supported by FA4, got fa_version={fa_version}" if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) @@ -446,12 +455,15 @@ def flash_attn_varlen_func( q, k, v, + qv=q_v, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, seqused_k=seqused_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, page_table=block_table, + gather_kv_indices=gather_kv_indices, + gather_kv_valid_length=gather_kv_valid_length, softmax_scale=softmax_scale, causal=causal, dynamic_causal=dynamic_causal, From 0fae39a0142c5d296f7938fa753e8f00b5d80171 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Tue, 13 Oct 2026 21:37:10 +0800 Subject: [PATCH 386/624] [Perf] per_token_group_quant_8bit: multiply by reciprocal instead of dividing per element (#61343) Signed-off-by: Thien Tran --- .../quantization/fused_kernels/layernorm_utils.cuh | 10 +++++++--- .../quantization/w8a8/fp8/per_token_group_quant.cu | 14 ++++++++------ 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/csrc/quantization/fused_kernels/layernorm_utils.cuh b/csrc/quantization/fused_kernels/layernorm_utils.cuh index d633cb9a6e57..8cab1824f278 100644 --- a/csrc/quantization/fused_kernels/layernorm_utils.cuh +++ b/csrc/quantization/fused_kernels/layernorm_utils.cuh @@ -516,6 +516,11 @@ __device__ void norm_and_quant( q8x4_t out; + // Group quant multiplies by the reciprocal of the scale (one div per + // group), matching per_token_group_quant_8bit. Per-token fp8 must keep + // exact division for FBGemm parity, so it is excluded here. + constexpr bool mul_by_recip = (group_size > 0) || is_scale_inverted; + float scale_val; if constexpr (group_size > 0) { @@ -528,14 +533,13 @@ __device__ void norm_and_quant( } else { scale_idx = blockIdx.x * num_groups + i * VEC_SIZE / group_size; } - scale_val = - is_scale_inverted ? 1.0f / scale[scale_idx] : scale[scale_idx]; + scale_val = 1.0f / scale[scale_idx]; } else { scale_val = *scale; } #pragma unroll for (int j = 0; j < VEC_SIZE; ++j) { - out.val[j] = ScaledQuant::quant_fn( + out.val[j] = ScaledQuant::quant_fn( static_cast(x.val[j] * rms) * w.val[j], scale_val); } vec_output[i] = out; diff --git a/csrc/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu index d2a888643899..0571cc90778d 100644 --- a/csrc/quantization/w8a8/fp8/per_token_group_quant.cu +++ b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu @@ -78,9 +78,13 @@ __device__ __forceinline__ void QuantizeGroup( const float y_s, const float min_8bit, const float max_8bit) { constexpr int vec_size = 16 / sizeof(T); + // One div.rn per group instead of per element; the register-resident + // kernel does the same, keeping both schedules bit-identical. + const float inv_y = 1.0f / y_s; + // quantize shared -> global 8-bit auto scalar_op_quant = [&] __device__(DST_DTYPE & dst, const T& src) { - float q = fminf(fmaxf(static_cast(src) / y_s, min_8bit), max_8bit); + float q = fminf(fmaxf(static_cast(src) * inv_y, min_8bit), max_8bit); dst = DST_DTYPE(q); }; @@ -272,16 +276,14 @@ __global__ void per_token_group_quant_8bit_register_kernel( } } - // A UE8M0 scale is an exact power of two, so multiplying by the reciprocal - // matches dividing. Otherwise the reciprocal costs an ulp and can move - // values that land exactly on an 8-bit boundary, so divide instead to stay - // bit-exact with the smem kernel. + // One div.rn per group; the smem path (QuantizeGroup) does the same, + // keeping both schedules bit-identical. const float inv_y = 1.0f / y_s; uint32_t packed[VEC_SIZE / 4] = {}; #pragma unroll for (int i = 0; i < VEC_SIZE; ++i) { const float v = static_cast(regs[i]); - const float scaled = SCALE_UE8M0 ? (v * inv_y) : (v / y_s); + const float scaled = v * inv_y; const float q = fminf(fmaxf(scaled, min_8bit), max_8bit); DST_DTYPE qb = DST_DTYPE(q); const uint8_t byte = *reinterpret_cast(&qb); From de8291e24f1312da0239605e33a89fe9cc901e47 Mon Sep 17 00:00:00 2001 From: Dakai An <77474977+andakai@users.noreply.github.com> Date: Tue, 13 Oct 2026 21:51:05 +0800 Subject: [PATCH 387/624] [Core] Add lifecycle hooks for worker extensions (#59463) Signed-off-by: Dakai An Co-authored-by: Claude Opus 5.5 Co-authored-by: aoshen02 --- tests/v1/worker/test_worker_base.py | 85 +++++++++++++++++++++++++++++ vllm/config/parallel.py | 5 +- vllm/v1/worker/worker_base.py | 13 +++++ 3 files changed, 102 insertions(+), 1 deletion(-) create mode 100644 tests/v1/worker/test_worker_base.py diff --git a/tests/v1/worker/test_worker_base.py b/tests/v1/worker/test_worker_base.py new file mode 100644 index 000000000000..079b643dc84f --- /dev/null +++ b/tests/v1/worker/test_worker_base.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest + +from vllm.v1.worker import worker_base +from vllm.v1.worker.worker_base import WorkerBase, WorkerWrapperBase + +pytestmark = pytest.mark.skip_global_cleanup + + +class ExtensionA: + events: list[str] + + def init_worker_extension(self): + self.events.append("init-a") + + def destroy_worker_extension(self): + self.events.append("destroy-a") + + +class ExtensionB: ... + + +class FailingExtension: + def destroy_worker_extension(self): + raise RuntimeError("destroy failed") + + +@pytest.fixture +def make_wrapper(monkeypatch): + # Fresh worker class per test, since extension injection mutates its bases. + class Worker(WorkerBase): + def __init__(self, **_kwargs): + self.events: list[str] = [] + + def init_device(self): + self.events.append("device") + + def shutdown(self): + self.events.append("shutdown") + + classes = { + cls.__name__: cls for cls in (Worker, ExtensionA, ExtensionB, FailingExtension) + } + monkeypatch.setattr(worker_base, "resolve_obj_by_qualname", classes.__getitem__) + monkeypatch.setattr( + worker_base, + "worker_receiver_cache_from_config", + lambda *_args, **_kwargs: None, + ) + + def make(extension: str) -> WorkerWrapperBase: + config = SimpleNamespace( + parallel_config=SimpleNamespace( + worker_cls="Worker", worker_extension_cls=extension + ), + enable_trace_function_call_for_thread=lambda: None, + ) + wrapper = WorkerWrapperBase() + wrapper.init_worker([{"vllm_config": config}]) + wrapper.init_device() + return wrapper + + return make + + +def test_worker_extension_hooks_follow_configured_class(make_wrapper): + for extension, expected in ( + ("ExtensionA", ["device", "init-a", "destroy-a", "shutdown"]), + ("ExtensionB", ["device", "shutdown"]), + ("ExtensionA", ["device", "init-a", "destroy-a", "shutdown"]), + ): + wrapper = make_wrapper(extension) + wrapper.shutdown() + assert wrapper.worker.events == expected + + +def test_failing_destroy_hook_is_logged_not_raised(make_wrapper, caplog_vllm): + wrapper = make_wrapper("FailingExtension") + wrapper.shutdown() + assert wrapper.worker.events == ["device", "shutdown"] + assert "Worker extension destroy hook failed" in caplog_vllm.text diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index e103c9ffd73d..73d83442baa7 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -309,7 +309,10 @@ class ParallelConfig: """The full name of the worker extension class to use. The worker extension class is dynamically inherited by the worker class. This is used to inject new attributes and methods to the worker class for use in collective_rpc - calls.""" + calls. It may also define `init_worker_extension(self)`, called after + device and distributed initialization but before model loading, and + `destroy_worker_extension(self)`, called (best-effort) before the worker + shuts down.""" master_addr: str = "127.0.0.1" """distributed master address for multi-node distributed inference when distributed_executor_backend is mp.""" diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index 95a0c3475d4d..e1a4eaf29102 100644 --- a/vllm/v1/worker/worker_base.py +++ b/vllm/v1/worker/worker_base.py @@ -244,9 +244,17 @@ def __init__( # Initialized after init_worker is called self.worker: WorkerBase self.vllm_config: VllmConfig + self._worker_extension_cls: type | None = None def shutdown(self) -> None: if self.worker is not None: + hook = getattr(self._worker_extension_cls, "destroy_worker_extension", None) + if callable(hook): + try: + hook(self.worker) + except Exception: + # Don't raise: callers still need to run their own teardown. + logger.exception("Worker extension destroy hook failed.") self.worker.shutdown() def update_environment_variables( @@ -293,6 +301,7 @@ def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None: "and pass the qualified name of the class as a string." ) + worker_extension_cls: type | None = None if parallel_config.worker_extension_cls: worker_extension_cls = resolve_obj_by_qualname( parallel_config.worker_extension_cls @@ -336,6 +345,7 @@ def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None: with set_current_vllm_config(self.vllm_config): # To make vLLM config available during worker initialization self.worker = worker_class(**kwargs) + self._worker_extension_cls = worker_extension_cls def initialize_from_config(self, kv_cache_configs: list[Any]) -> None: kv_cache_config = kv_cache_configs[self.global_rank] @@ -348,6 +358,9 @@ def init_device(self): with set_current_vllm_config(self.vllm_config): # To make vLLM config available during device initialization self.worker.init_device() # type: ignore + hook = getattr(self._worker_extension_cls, "init_worker_extension", None) + if callable(hook): + hook(self.worker) def __getattr__(self, attr: str): return getattr(self.worker, attr) From 61d9c777bfdef66efbf017c7c363d2bdca5db8ee Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Tue, 13 Oct 2026 22:42:53 +0800 Subject: [PATCH 388/624] [CI] Run the Clef pooler test on GPU (#61409) Signed-off-by: zjy0516 Co-authored-by: Claude Opus 5.5 --- tests/models/multimodal/pooling/test_clef_omni.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/models/multimodal/pooling/test_clef_omni.py b/tests/models/multimodal/pooling/test_clef_omni.py index 381ac23d53cc..c17b02222efa 100644 --- a/tests/models/multimodal/pooling/test_clef_omni.py +++ b/tests/models/multimodal/pooling/test_clef_omni.py @@ -15,9 +15,11 @@ JointSchemaHead, JointSchemaPooler, ) +from vllm.platforms import current_platform from vllm.pooling_params import PoolingParams from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates +DEVICE_TYPE = current_platform.device_type HIDDEN, VOCAB = 64, 500 TYPES = [0, 1, 1, 2] COUNTS = [2, 3, 1, 5] @@ -55,7 +57,7 @@ def reference_logits(head, hidden, token_ids, embedding, question_spans, spans): base + head.option_summary_norm(torch.stack(summaries)) + head.global_projection(global_vector) - + head.type_embedding(torch.tensor(TYPES)) + + head.type_embedding(torch.tensor(TYPES, device=hidden.device)) ).unsqueeze(0) for layer in head.layers: fields = layer(fields, memory) @@ -90,7 +92,8 @@ def pooler(): with torch.no_grad(): for p in head.parameters(): p.add_(torch.randn_like(p) * 0.2) - embedding = torch.randn(VOCAB, HIDDEN) + head.to(DEVICE_TYPE) + embedding = torch.randn(VOCAB, HIDDEN, device=DEVICE_TYPE) config = SimpleNamespace( scheduler_config=SimpleNamespace( enable_chunked_prefill=False, async_scheduling=False @@ -132,7 +135,7 @@ def test_pooler_matches_reference(pooler): pooler, embedding = pooler lengths = [120, 90, 100] token_ids = [torch.randint(0, VOCAB, (n,)) for n in lengths] - hidden = torch.randn(sum(lengths), HIDDEN) + hidden = torch.randn(sum(lengths), HIDDEN, device=DEVICE_TYPE) first, first_questions, first_spans = schema(lengths[0]) last, last_questions, last_spans = schema(lengths[2]) params = [ @@ -154,7 +157,7 @@ def test_pooler_matches_reference(pooler): metadata.build_pooling_cursor( torch.tensor(lengths).numpy(), seq_lens_cpu=torch.tensor(lengths), - device=torch.device("cpu"), + device=hidden.device, ) probs, empty, logits = pooler(hidden, metadata) From f5fefbffe51e1cc91e0d6507d272f337473989b3 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Tue, 13 Oct 2026 11:01:44 -0400 Subject: [PATCH 389/624] [Build] DeepGEMM migration to torch Stable ABI (#59482) Signed-off-by: Chris Leonard Signed-off-by: Shengqi Chen Co-authored-by: Shengqi Chen --- .buildkite/check-torch-abi.py | 1 - cmake/external_projects/deepgemm.cmake | 4 ++-- tools/install_deepgemm.sh | 2 +- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/.buildkite/check-torch-abi.py b/.buildkite/check-torch-abi.py index 37e2b8dadbf8..5bfec8abd077 100644 --- a/.buildkite/check-torch-abi.py +++ b/.buildkite/check-torch-abi.py @@ -17,7 +17,6 @@ # _C_stable_libtorch, but these two modules still bind libtorch directly. "_C.abi3.so", "_rocm_C.abi3.so", - "third_party/deep_gemm/_C*.so", ) diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index b796623b7ac1..7f1c0ba5157e 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -32,8 +32,8 @@ else() # Pinned to the fork's dev branch: upstream's 26/09/30 release (2.8.1) # plus the SM120 port, the SM90 paged-MQA kv_block=32/next_n=4 port, # configurable SwiGLU alpha/beta, SiTU for FP8/FP4 Mega MoE, and NVFP4 - # Mega MoE. TORCH_LIBRARY build with a CPython 3.10 stable-ABI floor. - set(_DEEPGEMM_UPSTREAM_TAG "6ce1c5ec6627a6bf11dbcc698b0b1e19c9248432") + # Mega MoE. TORCH_STABLE_LIBRARY build with a CPython 3.10 stable-ABI floor. + set(_DEEPGEMM_UPSTREAM_TAG "a981d8e3a58e877b7f53e63707de5834312fb710") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index b617e6d5679f..c293bb3962ba 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -10,7 +10,7 @@ DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" # NOTE: This targets the vLLM fork's dev branch: upstream's 26/09/30 release # plus the sm120 and sm90 paged-MQA ports, the SwiGLU alpha/beta and SiTU # Mega MoE activations, and NVFP4 Mega MoE. -DEEPGEMM_GIT_REF="6ce1c5ec6627a6bf11dbcc698b0b1e19c9248432" +DEEPGEMM_GIT_REF="a981d8e3a58e877b7f53e63707de5834312fb710" WHEEL_DIR="" # Parse command line arguments From f395a6b1ad49458ea5b0366805c3814c51d6b02b Mon Sep 17 00:00:00 2001 From: Rishabh Agarwal Date: Tue, 13 Oct 2026 11:03:48 -0400 Subject: [PATCH 390/624] [ROCM][Bugfix] Deferring to TRITON_ATTN backend for Prefix lm models instead of ROCM_ATTN and fixed minor bug in Gemma3 Multimodal code (#46754) Signed-off-by: GodRishUniverse Signed-off-by: Rishabh Agarwal Signed-off-by: GodRishUniverse --- .../test_rocm_attention_backends_selection.py | 25 +++++++++++++++++++ vllm/model_executor/models/gemma3_mm.py | 4 +++ 2 files changed, 29 insertions(+) diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index e0247a94a1a3..45580b395f72 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -61,6 +61,31 @@ def test_aiter_unified_attention_uses_dedicated_metadata_builder(): ) +def test_rocm_backends_do_not_support_mm_prefix(): + """Regression test: ROCm backends must not claim mm_prefix support. + + If they do, prefix-LM multimodal models that attend vision tokens + bidirectionally (e.g. Gemma3 4B) will use bidirectional attention + via the ROCm backend instead of falling back to Triton, producing + incorrect results. Remove this test only if support is properly added. + """ + from vllm.v1.attention.backends.rocm_aiter_unified_attn import ( + RocmAiterUnifiedAttentionBackend, + ) + from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend + + assert not RocmAttentionBackend.supports_mm_prefix(), ( + "RocmAttentionBackend does not support mm_prefix — " + "defer to Triton unified attention instead for now" + "See https://github.com/vllm-project/vllm/issues/48172" + ) + assert not RocmAiterUnifiedAttentionBackend.supports_mm_prefix(), ( + "RocmAiterUnifiedAttentionBackend does not support mm_prefix — " + "defer to Triton unified attention instead for now" + "See https://github.com/ROCm/aiter/issues/4168" + ) + + def test_segmented_attention_uses_dedicated_backend_components(): from vllm.v1.attention.backends.rocm_segmented_attn import ( RocmSegmentedAttentionBackend, diff --git a/vllm/model_executor/models/gemma3_mm.py b/vllm/model_executor/models/gemma3_mm.py index c6166bc2f423..b5b0c0d48bc4 100644 --- a/vllm/model_executor/models/gemma3_mm.py +++ b/vllm/model_executor/models/gemma3_mm.py @@ -557,6 +557,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) self.multi_modal_projector = Gemma3MultiModalProjector(config) + # https://huggingface.co/blog/gemma3, (image tokens attend to each other) + # explicit mention here as text side in gemma3 standalone does not need + # prefix_lm but multimodal does + config.text_config.is_mm_prefix_lm = True with self._mark_language_model(vllm_config): self.language_model = init_vllm_registered_model( vllm_config=vllm_config, From 669f85e686ae533a761362c919d213a87e0e807c Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Tue, 13 Oct 2026 23:08:07 +0800 Subject: [PATCH 391/624] [Misc] Break chat_utils into more maintainable modules. (#60972) Signed-off-by: wang.yuqi --- .../cohere/test_serving_conversion.py | 6 +- .../unit_tests/test_chat_billing_header.py | 4 +- .../multimodal/processing/test_common.py | 2 +- tests/renderers/test_chat_utils.py | 12 +- .../test_chat_utils_prompt_embeds.py | 9 +- tests/renderers/test_warmup.py | 2 +- vllm/entrypoints/cohere/serving.py | 4 +- .../entrypoints/pooling/embed/io_processor.py | 2 +- .../pooling/scoring/io_processor.py | 2 +- vllm/entrypoints/pooling/scoring/utils.py | 4 +- vllm/exceptions.py | 6 + vllm/renderers/base.py | 2 +- vllm/renderers/chat_utils.py | 2349 ----------------- vllm/renderers/chat_utils/__init__.py | 42 + vllm/renderers/chat_utils/chat_message.py | 826 ++++++ vllm/renderers/chat_utils/chat_template.py | 112 + vllm/renderers/chat_utils/content_parser.py | 560 ++++ vllm/renderers/chat_utils/item_tracker.py | 562 ++++ vllm/renderers/chat_utils/tool_call.py | 32 + vllm/renderers/chat_utils/types.py | 367 +++ vllm/renderers/hf.py | 3 +- 21 files changed, 2533 insertions(+), 2375 deletions(-) delete mode 100644 vllm/renderers/chat_utils.py create mode 100644 vllm/renderers/chat_utils/__init__.py create mode 100644 vllm/renderers/chat_utils/chat_message.py create mode 100644 vllm/renderers/chat_utils/chat_template.py create mode 100644 vllm/renderers/chat_utils/content_parser.py create mode 100644 vllm/renderers/chat_utils/item_tracker.py create mode 100644 vllm/renderers/chat_utils/tool_call.py create mode 100644 vllm/renderers/chat_utils/types.py diff --git a/tests/entrypoints/cohere/test_serving_conversion.py b/tests/entrypoints/cohere/test_serving_conversion.py index e6696c26b42c..7136b4acdd95 100644 --- a/tests/entrypoints/cohere/test_serving_conversion.py +++ b/tests/entrypoints/cohere/test_serving_conversion.py @@ -511,7 +511,7 @@ def test_text_only_list_preserves_per_block_text_parts(self): # A list of text blocks maps 1-to-1 onto ``text`` content # parts, mirroring the shape ``_convert_user_message`` emits # for multi-part user content. vLLM's - # ``_parse_chat_message_content`` will collapse this back to + # ``parse_chat_message_content`` will collapse this back to # ``"\n".join(...)`` for the ``ConversationMessage`` handed to # renderers (see the tool-role branch in # ``vllm/entrypoints/chat_utils.py``), so downstream chat @@ -533,7 +533,7 @@ def test_document_blocks_serialized_as_text_parts(self): """Document blocks must not survive on the OpenAI-shaped tool message. - vLLM's shared ``_parse_chat_message_content_parts`` runs a + vLLM's shared ``parse_chat_message_content_parts`` runs a Pydantic ``ValidatorIterator`` over the OpenAI content-part union, which only knows ``text`` / ``image_url`` / etc. and rejects ``{"type": "document", ...}`` blocks outright. We @@ -588,7 +588,7 @@ def test_document_only_content_serialized_as_text_part(self): def test_converted_tool_message_content_survives_openai_validation(self): """Regression for the reported crash. - ``vllm/entrypoints/chat_utils.py:_parse_chat_message_content_parts`` + ``vllm/renders/chat_message_content.py:parse_chat_message_content_parts`` runs a Pydantic ``ValidatorIterator`` over ``list[ChatCompletionContentPartParam]``, which is a discriminated union that only accepts ``text`` / ``image_url`` diff --git a/tests/entrypoints/unit_tests/test_chat_billing_header.py b/tests/entrypoints/unit_tests/test_chat_billing_header.py index 3ad694759760..1848aa568b91 100644 --- a/tests/entrypoints/unit_tests/test_chat_billing_header.py +++ b/tests/entrypoints/unit_tests/test_chat_billing_header.py @@ -6,7 +6,7 @@ import pytest -from vllm.renderers.chat_utils import _parse_chat_message_content +from vllm.renderers.chat_utils import parse_chat_message_content HEADER = ( "x-anthropic-billing-header: cc_version=2.1.220.300; " @@ -17,7 +17,7 @@ def _parse(message, content_format="openai"): tracker = MagicMock() tracker.create_parser.return_value.mm_placeholder_storage.return_value = {} - return _parse_chat_message_content( + return parse_chat_message_content( message, tracker, content_format, interleave_strings=False ) diff --git a/tests/models/multimodal/processing/test_common.py b/tests/models/multimodal/processing/test_common.py index 820d655e905e..667907c3e55c 100644 --- a/tests/models/multimodal/processing/test_common.py +++ b/tests/models/multimodal/processing/test_common.py @@ -29,7 +29,7 @@ from vllm.multimodal.inputs import MultiModalFeatureSpec, batched_tensors_equal from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext from vllm.platforms import current_platform -from vllm.renderers.chat_utils import _get_embeds_data +from vllm.renderers.chat_utils.item_tracker import _get_embeds_data from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config from vllm.utils.mistral import is_mistral_tokenizer diff --git a/tests/renderers/test_chat_utils.py b/tests/renderers/test_chat_utils.py index 47fe49a63e6e..877c4de74b5b 100644 --- a/tests/renderers/test_chat_utils.py +++ b/tests/renderers/test_chat_utils.py @@ -16,24 +16,24 @@ from vllm.config import ModelConfig from vllm.exceptions import VLLMValidationError from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict +from vllm.multimodal.media import MEDIA_CONNECTOR_REGISTRY from vllm.multimodal.utils import ( encode_audio_url, encode_image_url, encode_video_url, ) from vllm.renderers.chat_utils import ( - MEDIA_CONNECTOR_REGISTRY, AsyncMultiModalItemTracker, ChatCompletionMessageParam, ConversationMessage, - _load_embeds_dict, - _parse_chat_message_content, - _parse_metadata_array, - _postprocess_messages, + parse_chat_message_content, parse_chat_messages, parse_chat_messages_async, validate_chat_template, ) +from vllm.renderers.chat_utils.chat_message import _postprocess_messages +from vllm.renderers.chat_utils.content_parser import _load_embeds_dict +from vllm.renderers.chat_utils.item_tracker import _parse_metadata_array from vllm.utils.serial_utils import tensor2base64 KIMI_K2_5_MODEL_ID = "moonshotai/Kimi-K2.5" @@ -3205,7 +3205,7 @@ def _parse_one(message) -> ConversationMessage: parser = tracker.create_parser.return_value parser.model_config.enable_prompt_embeds = False parser.mm_placeholder_storage.return_value = {} - result = _parse_chat_message_content( + result = parse_chat_message_content( message, tracker, content_format="string", interleave_strings=False ) assert len(result) == 1 diff --git a/tests/renderers/test_chat_utils_prompt_embeds.py b/tests/renderers/test_chat_utils_prompt_embeds.py index f26680793ee2..d6f1673ca79a 100644 --- a/tests/renderers/test_chat_utils_prompt_embeds.py +++ b/tests/renderers/test_chat_utils_prompt_embeds.py @@ -18,15 +18,16 @@ from vllm.exceptions import VLLMValidationError from vllm.renderers import TokenizeParams -from vllm.renderers.chat_utils import ( - _ENABLE_PROMPT_EMBEDS_ERROR, +from vllm.renderers.chat_utils import parse_chat_messages, parse_chat_messages_async +from vllm.renderers.chat_utils.chat_message import ( _PROMPT_EMBEDS_MISSING_DATA_ERROR, _RESERVED_PLACEHOLDER_IN_TEXT_ERROR, MM_PARSER_MAP, +) +from vllm.renderers.chat_utils.content_parser import ( + _ENABLE_PROMPT_EMBEDS_ERROR, MODALITY_PLACEHOLDERS_MAP, PROMPT_EMBEDS_PLACEHOLDER_TOKEN, - parse_chat_messages, - parse_chat_messages_async, ) from vllm.renderers.hf import ( _PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR, diff --git a/tests/renderers/test_warmup.py b/tests/renderers/test_warmup.py index 09fe047d56ef..0c847e5600f2 100644 --- a/tests/renderers/test_warmup.py +++ b/tests/renderers/test_warmup.py @@ -28,9 +28,9 @@ import pytest +from vllm.exceptions import ChatTemplateResolutionError from vllm.multimodal.processing import BaseMultiModalProcessor from vllm.renderers.base import BaseRenderer -from vllm.renderers.chat_utils import ChatTemplateResolutionError from vllm.renderers.params import ChatParams diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 93f6d14a3b9c..4ce734a8121c 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -502,7 +502,7 @@ def _convert_tool_message(cls, msg: ToolChatMessageV2) -> dict[str, Any]: # Map every content block to a plain ``{"type": "text", "text": # ...}`` part, mirroring the shape :meth:`_convert_user_message` # emits for multi-part user content. vLLM's - # ``_parse_chat_message_content`` collapses a text-only tool + # ``parse_chat_message_content`` collapses a text-only tool # part list back into ``"\n".join(...)`` on the way to the # renderer, so simple chat templates that expect string tool # content still work; keeping the list here just preserves the @@ -512,7 +512,7 @@ def _convert_tool_message(cls, msg: ToolChatMessageV2) -> dict[str, Any]: # Tool message content uses ``ToolMessageV2Content``, a union of # ``TextToolContent`` and ``DocumentToolContent``. The latter # produces ``{"type": "document", "document": {...}}`` blocks, - # which vLLM's shared ``_parse_chat_message_content_parts`` in + # which vLLM's shared ``parse_chat_message_content_parts`` in # ``vllm/renderers/chat_utils.py`` rejects because the OpenAI # ``ChatCompletionContentPartParam`` union only knows text / # image / audio / etc. Emitting them anyway triggers a Pydantic diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index b37a009c925a..4b2ad7a3b8ee 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -9,7 +9,6 @@ ChatCompletionContentPartImageParam, ChatCompletionContentPartTextParam, ) -from openai.types.chat.chat_completion_content_part_image_param import ImageURL from vllm import PoolingParams from vllm.exceptions import VLLMValidationError @@ -21,6 +20,7 @@ ChatCompletionContentPartParam, ChatCompletionMessageParam, CustomChatCompletionMessageParam, + ImageURL, ) from vllm.renderers.hf import resolve_chat_template from vllm.utils.collection_utils import chunk_list diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index 1e4c638968f8..9ab36a13958c 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -7,9 +7,9 @@ import torch.nn.functional as F from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt +from vllm.exceptions import ChatTemplateResolutionError from vllm.logger import init_logger from vllm.renderers import TokenizeParams -from vllm.renderers.chat_utils import ChatTemplateResolutionError from vllm.renderers.hf import safe_apply_chat_template from vllm.renderers.inputs.preprocess import ( extract_target_prompt, diff --git a/vllm/entrypoints/pooling/scoring/utils.py b/vllm/entrypoints/pooling/scoring/utils.py index 8d953cf6ddd4..b3617ae9c601 100644 --- a/vllm/entrypoints/pooling/scoring/utils.py +++ b/vllm/entrypoints/pooling/scoring/utils.py @@ -15,7 +15,7 @@ ChatCompletionContentPartTextParam, ConversationMessage, MultiModalItemTracker, - _parse_chat_message_content_parts, + parse_chat_message_content_parts, ) from .typing import ( @@ -193,7 +193,7 @@ def _parse_score_content( mm_parser = mm_tracker.create_parser() - parse_res = _parse_chat_message_content_parts( + parse_res = parse_chat_message_content_parts( role=role, parts=parts, mm_tracker=mm_tracker, diff --git a/vllm/exceptions.py b/vllm/exceptions.py index 11ca13fe0bc3..85863c2fbf1d 100644 --- a/vllm/exceptions.py +++ b/vllm/exceptions.py @@ -62,6 +62,12 @@ class VLLMNotFoundError(VLLMClientError): pass +class ChatTemplateResolutionError(VLLMClientError): + """Raised when chat template resolution fails.""" + + pass + + class LoRAAdapterNotFoundError(VLLMNotFoundError): """Exception raised when a LoRA adapter is not found. diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index e04bfa64eb76..7b0f4debaad8 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -362,7 +362,7 @@ def warmup(self, chat_params: ChatParams) -> None: no-op if the background warmup — or a prior reset_mm_cache that joined it — already ran it. """ - from vllm.renderers.chat_utils import ChatTemplateResolutionError + from vllm.exceptions import ChatTemplateResolutionError # prevent MM processor hangs with set_default_torch_num_threads(1): diff --git a/vllm/renderers/chat_utils.py b/vllm/renderers/chat_utils.py deleted file mode 100644 index 540564baab8d..000000000000 --- a/vllm/renderers/chat_utils.py +++ /dev/null @@ -1,2349 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import asyncio -import json -import types -from abc import ABC, abstractmethod -from collections import Counter, defaultdict -from collections.abc import Awaitable, Callable, Iterable -from dataclasses import dataclass -from functools import cached_property, lru_cache, partial -from itertools import accumulate -from pathlib import Path -from typing import ( - TYPE_CHECKING, - Any, - Final, - Generic, - Literal, - Required, - TypeAlias, - TypeVar, - Union, - cast, - get_args, - get_origin, -) - -from openai.types.chat import ( - ChatCompletionAssistantMessageParam, - ChatCompletionContentPartImageParam, - ChatCompletionContentPartInputAudioParam, - ChatCompletionContentPartRefusalParam, - ChatCompletionContentPartTextParam, - ChatCompletionFunctionToolParam, - ChatCompletionMessageToolCallParam, - ChatCompletionToolMessageParam, -) -from openai.types.chat import ( - ChatCompletionContentPartParam as OpenAIChatCompletionContentPartParam, -) -from openai.types.chat import ( - ChatCompletionMessageParam as OpenAIChatCompletionMessageParam, -) -from openai.types.chat.chat_completion_content_part_input_audio_param import InputAudio -from openai.types.responses import ResponseInputImageParam -from openai_harmony import Message as OpenAIHarmonyMessage -from PIL import Image -from pydantic import BaseModel, ConfigDict, TypeAdapter - -# pydantic needs the TypedDict from typing_extensions -from typing_extensions import TypedDict, override - -from vllm import envs -from vllm.config import ModelConfig -from vllm.connections import MediaRequestDownloadBudget -from vllm.exceptions import VLLMValidationError -from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict -from vllm.logger import init_logger -from vllm.model_executor.models import SupportsMultiModal -from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.inputs import ( - MultiModalBatchedField, - MultiModalFlatField, - MultiModalSharedField, - VisionChunk, - VisionChunkImage, - VisionChunkVideo, -) -from vllm.multimodal.media import MEDIA_CONNECTOR_REGISTRY, MediaConnector -from vllm.multimodal.media.base import MediaDecodedBudget -from vllm.multimodal.processing import BaseMultiModalProcessor -from vllm.renderers.embed_utils import ( - safe_load_prompt_embeds, - safe_load_prompt_embeds_async, -) -from vllm.transformers_utils.processor import get_video_processor_cls_name -from vllm.utils import random_uuid -from vllm.utils.collection_utils import is_list_of, is_list_of_numbers -from vllm.utils.import_utils import LazyLoader - -if TYPE_CHECKING: - import torch - import transformers -else: - transformers = LazyLoader("transformers", globals(), "transformers") - torch = LazyLoader("torch", globals(), "torch") - -logger = init_logger(__name__) - - -class ChatTemplateResolutionError(ValueError): - """Raised when chat template resolution fails. - - This is a subclass of ValueError for backward compatibility with - existing exception handlers. - """ - - -MODALITY_PLACEHOLDERS_MAP = { - "image": "<##IMAGE##>", - "audio": "<##AUDIO##>", - "video": "<##VIDEO##>", - "prompt_embeds": "<##PROMPT_EMBEDS##>", -} - - -PROMPT_EMBEDS_PLACEHOLDER_TOKEN: Final[str] = "" -"""The special token used as a placeholder for each embedding -position during chat template rendering. - -Registered as an additional special token when `--enable-prompt-embeds` is set. -See `_ensure_prompt_embeds_placeholder_token` in `vllm/renderers/hf.py`. -""" - - -_REQUIRE_MM_PROCESSOR_ERROR: Final[str] = ( - "Resolving modality {modality!r} requires a multimodal processor " - "but none is available." -) - -_ENABLE_PROMPT_EMBEDS_ERROR: Final[str] = ( - "You must set `--enable-prompt-embeds` to input `prompt_embeds`" -) - -_PROMPT_EMBEDS_MISSING_DATA_ERROR: Final[str] = ( - "prompt_embeds content part requires a non-empty `data` field " - "with base64-encoded tensor bytes." -) - -_RESERVED_PLACEHOLDER_IN_TEXT_ERROR: Final[str] = ( - "Text content may not contain the reserved placeholder {token!r}. " - "This placeholder is used internally to mark `prompt_embeds` splice " - "positions in the tokenized prompt." -) - - -class AudioURL(TypedDict, total=False): - url: Required[str] - """ - Either a URL of the audio or a data URL with base64 encoded audio data. - """ - - -class ChatCompletionContentPartAudioParam(TypedDict, total=False): - audio_url: Required[AudioURL] - - type: Required[Literal["audio_url"]] - """The type of the content part.""" - - -MultiModalEmbedsPayload: TypeAlias = str | dict[str, str | list[int | float]] - - -class ChatCompletionContentPartImageEmbedsParam(TypedDict, total=False): - image_embeds: MultiModalEmbedsPayload | None - """ - The image embeddings. It can be either: - - A single base64 string. - - A dictionary of base64 tensors or numeric JSON metadata arrays. - """ - type: Required[Literal["image_embeds"]] - """The type of the content part.""" - uuid: str | None - """ - User-provided UUID of a media. User must guarantee that it is properly - generated and unique for different medias. - """ - - -class ChatCompletionContentPartAudioEmbedsParam(TypedDict, total=False): - audio_embeds: MultiModalEmbedsPayload | None - """ - The audio embeddings. It can be either: - - A single base64 string representing a serialized torch tensor. - - A dictionary of base64 tensors or numeric JSON metadata arrays. - """ - type: Required[Literal["audio_embeds"]] - """The type of the content part.""" - uuid: str | None - """ - User-provided UUID of a media. User must guarantee that it is properly - generated and unique for different medias. - """ - - -class ChatCompletionContentPartVideoEmbedsParam(TypedDict, total=False): - video_embeds: MultiModalEmbedsPayload | None - """ - The video embeddings. It can be either: - - A single base64 string representing a serialized torch tensor. - - A dictionary of base64 tensors or numeric JSON metadata arrays. - """ - type: Required[Literal["video_embeds"]] - """The type of the content part.""" - uuid: str | None - """ - User-provided UUID of a media. User must guarantee that it is properly - generated and unique for different medias. - """ - - -class ChatCompletionContentPartPromptEmbedsParam(TypedDict, total=False): - data: Required[str] - """ - Base64-encoded bytes of a serialized `torch.Tensor` of shape - `(num_tokens, hidden_size)`. The tensor's `dtype` and `hidden_size` must - match the model's input embedding layer. - """ - type: Required[Literal["prompt_embeds"]] - """The type of the content part.""" - - -class VideoURL(TypedDict, total=False): - url: Required[str] - """ - Either a URL of the video or a data URL with base64 encoded video data. - """ - - -class ChatCompletionContentPartVideoParam(TypedDict, total=False): - video_url: Required[VideoURL] - - type: Required[Literal["video_url"]] - """The type of the content part.""" - - -class PILImage(BaseModel): - """A PIL.Image.Image object.""" - - image_pil: Image.Image - model_config = ConfigDict(arbitrary_types_allowed=True) - - -class CustomChatCompletionContentPILImageParam(TypedDict, total=False): - """A simpler version of the param that only accepts a PIL image. - - Example: - { - "image_pil": ImageAsset('cherry_blossom').pil_image - } - - """ - - image_pil: PILImage | None - uuid: str | None - """ - User-provided UUID of a media. User must guarantee that it is properly - generated and unique for different medias. - """ - - -class CustomChatCompletionContentSimpleImageParam(TypedDict, total=False): - """A simpler version of the param that only accepts a plain image_url. - This is supported by OpenAI API, although it is not documented. - - Example: - { - "image_url": "https://example.com/image.jpg" - } - - """ - - image_url: str | None - uuid: str | None - """ - User-provided UUID of a media. User must guarantee that it is properly - generated and unique for different medias. - """ - - -class CustomChatCompletionContentSimpleAudioParam(TypedDict, total=False): - """A simpler version of the param that only accepts a plain audio_url. - - Example: - { - "audio_url": "https://example.com/audio.mp3" - } - - """ - - audio_url: str | None - - -class CustomChatCompletionContentSimpleVideoParam(TypedDict, total=False): - """A simpler version of the param that only accepts a plain audio_url. - - Example: - { - "video_url": "https://example.com/video.mp4" - } - - """ - - video_url: str | None - uuid: str | None - """ - User-provided UUID of a media. User must guarantee that it is properly - generated and unique for different medias. - """ - - -class CustomThinkCompletionContentParam(TypedDict, total=False): - """A Think Completion Content Param that accepts a plain text and a boolean. - - Example: - { - "thinking": "I am thinking about the answer", - "closed": True, - "type": "thinking" - } - - """ - - thinking: Required[str] - """The thinking content.""" - - closed: bool - """Whether the thinking is closed.""" - - type: Required[Literal["thinking"]] - """The thinking type.""" - - -class CustomChatCompletionContentToolReferenceParam(TypedDict, total=False): - """A tool reference content param that only accepts a plain tool name. - - Example: - { - "name": "get_weather", - "type": "tool_reference" - } - - """ - - name: str - """The name of the tool being referenced.""" - - type: Literal["tool_reference"] - """The content type.""" - - -ChatCompletionContentPartParam: TypeAlias = ( - OpenAIChatCompletionContentPartParam - | ChatCompletionContentPartAudioParam - | ChatCompletionContentPartInputAudioParam - | ChatCompletionContentPartVideoParam - | ChatCompletionContentPartRefusalParam - | CustomChatCompletionContentPILImageParam - | CustomChatCompletionContentSimpleImageParam - | ChatCompletionContentPartImageEmbedsParam - | ChatCompletionContentPartAudioEmbedsParam - | ChatCompletionContentPartVideoEmbedsParam - | ChatCompletionContentPartPromptEmbedsParam - | CustomChatCompletionContentSimpleAudioParam - | CustomChatCompletionContentSimpleVideoParam - | CustomChatCompletionContentToolReferenceParam - | str - | CustomThinkCompletionContentParam -) - - -class CustomChatCompletionMessageParam(TypedDict, total=False): - """Enables custom roles in the Chat Completion API.""" - - role: Required[str] - """The role of the message's author.""" - - content: str | list[ChatCompletionContentPartParam] - """The contents of the message.""" - - name: str - """An optional name for the participant. - - Provides the model information to differentiate between participants of the - same role. - """ - - tool_call_id: str | None - """Tool call that this message is responding to.""" - - tool_calls: list[ChatCompletionMessageToolCallParam] | None - """The tool calls generated by the model, such as function calls.""" - - reasoning: str | None - """The reasoning content for interleaved thinking.""" - - tools: list[ChatCompletionFunctionToolParam] | None - """The tools for developer role.""" - - task: str | None - """Model-specific task marker. Currently passed through for DeepSeek V4.""" - - -ChatCompletionMessageParam: TypeAlias = ( - OpenAIChatCompletionMessageParam - | CustomChatCompletionMessageParam - | OpenAIHarmonyMessage -) - - -# TODO: Make fields ReadOnly once mypy supports it -class ConversationMessage(TypedDict, total=False): - role: Required[str] - """The role of the message's author.""" - - content: str | None | list[dict[str, str]] - """The contents of the message""" - - tool_call_id: str | None - """Tool call that this message is responding to.""" - - name: str | None - """The name of the function to call""" - - tool_calls: list[ChatCompletionMessageToolCallParam] | None - """The tool calls generated by the model, such as function calls.""" - - reasoning: str | None - """The reasoning content for interleaved thinking.""" - - reasoning_content: str | None - """Deprecated: The reasoning content for interleaved thinking.""" - - tools: list[ChatCompletionFunctionToolParam] | None - """The tools for developer role.""" - - task: str | None - """Model-specific task marker. Currently passed through for DeepSeek V4.""" - - -# Passed in by user -ChatTemplateContentFormatOption = Literal["auto", "string", "openai"] - -# After resolving "auto" -ChatTemplateContentFormat = Literal["string", "openai"] - - -ModalityStr = Literal[ - "image", - "audio", - "video", - "image_embeds", - "audio_embeds", - "video_embeds", - "vision_chunk", - "prompt_embeds", -] -_T = TypeVar("_T") -_AsyncMultiModalItem: TypeAlias = Callable[[], Awaitable[tuple[object, str | None]]] - - -# Backward compatibility for single item input -class _BatchedSingleItemField(MultiModalSharedField): - pass - - -def _detect_field( - tensors: list[torch.Tensor], - mm_processor: BaseMultiModalProcessor, -): - first_item = tensors[0] - hidden_size = mm_processor.info.ctx.model_config.get_inputs_embeds_size() - - if ( - len(tensors) == 1 - and first_item.ndim == 3 - and first_item.shape[0] == 1 - and first_item.shape[-1] == hidden_size - ): - logger.warning( - "Batched multi-modal embedding inputs are deprecated for Chat API. " - "Please pass a separate content part for each multi-modal item." - ) - return _BatchedSingleItemField(batch_size=1) - - first_shape = first_item.shape - if all(t.shape == first_shape for t in tensors): - return MultiModalBatchedField() - - size_per_item = [len(tensor) for tensor in tensors] - slice_idxs = [0, *accumulate(size_per_item)] - slices = [ - (slice(slice_idxs[i], slice_idxs[i + 1]),) for i in range(len(size_per_item)) - ] - return MultiModalFlatField(slices=slices) - - -def _merge_embeds( - data_items: list[dict[str, "torch.Tensor"]], - mm_processor: BaseMultiModalProcessor, -): - if not data_items: - return {} - - first_keys = set(data_items[0].keys()) - if any(set(item.keys()) != first_keys for item in data_items[1:]): - raise VLLMValidationError( - "All dictionaries in the list of embeddings must have the same keys." - ) - - fields = { - key: _detect_field([item[key] for item in data_items], mm_processor) - for key in first_keys - } - data_merged = { - key: field._reduce_data([item[key] for item in data_items], pin_memory=False) - for key, field in fields.items() - } - - try: - # TODO: Support per-request mm_processor_kwargs - parsed_configs = mm_processor._get_mm_fields_config( - transformers.BatchFeature(data_merged), - {}, - ) - parsed_fields = {key: parsed_configs[key].field for key in first_keys} - keys_to_update = [ - key - for key in first_keys - if ( - fields[key] != parsed_fields[key] - and not isinstance(fields[key], _BatchedSingleItemField) - ) - ] - - for key in keys_to_update: - data_merged[key] = parsed_fields[key]._reduce_data( - [item[key] for item in data_items], pin_memory=False - ) - except Exception: - logger.exception( - "Error when parsing merged embeddings. " - "Falling back to auto-detected fields." - ) - - return data_merged - - -async def _load_embeds_dict( - data: dict[str, str | list[int | float]], - fetch: Callable[[str], Awaitable["torch.Tensor"]], -) -> dict[str, Any]: - encoded = {key: value for key, value in data.items() if isinstance(value, str)} - tensors = await asyncio.gather(*(fetch(value) for value in encoded.values())) - return {**data, **dict(zip(encoded, tensors))} - - -def _parse_metadata_array(key: str, value: list, metadata_fields: set[str]): - if key not in metadata_fields: - raise VLLMValidationError(f"JSON arrays are only supported for metadata: {key}") - if not is_list_of_numbers(value): - raise VLLMValidationError(f"Metadata {key} must be a finite numeric array.") - dtype = torch.float64 if any(isinstance(v, float) for v in value) else torch.int64 - try: - return torch.tensor(value, dtype=dtype) - except (ValueError, TypeError, OverflowError, RuntimeError) as error: - raise VLLMValidationError(f"Invalid metadata array: {key}") from error - - -def _get_embeds_data( - modality: str, - data_items: list[Any], - mm_processor: BaseMultiModalProcessor, -): - if len(data_items) == 0: - return data_items - - if all(item is None for item in data_items): - return data_items - - if is_list_of(data_items, torch.Tensor): - embeds_key = f"{modality}_embeds" - dict_items = [{embeds_key: item} for item in data_items] - return _merge_embeds(dict_items, mm_processor)[embeds_key] - - if is_list_of(data_items, dict): - metadata_fields = mm_processor.info.data_parser.placeholder_metadata_fields( - modality - ) - data_items = [ - { - key: _parse_metadata_array(key, value, metadata_fields) - if isinstance(value, list) - else value - for key, value in item.items() - } - for item in data_items - ] - return _merge_embeds(data_items, mm_processor) - - raise NotImplementedError(type(data_items)) - - -class BaseMultiModalItemTracker(ABC, Generic[_T]): - """Tracks multi-modal items in a given request and ensures that the number - of multi-modal items in a given request does not exceed the configured - maximum per prompt. - """ - - def __init__( - self, - model_config: ModelConfig, - media_io_kwargs: dict[str, dict[str, Any]] | None = None, - ): - super().__init__() - - self._model_config = model_config - self._media_io_kwargs = media_io_kwargs - - self._items_by_modality = defaultdict[str, list[_T]](list) - # Track original modality for each vision_chunk item (image or video) - self._modality_order = defaultdict[str, list[str]](list) - self.decoded_budget = MediaDecodedBudget() - self.download_budget = MediaRequestDownloadBudget() - - @cached_property - def use_unified_vision_chunk_modality(self) -> bool: - """Check if model uses unified vision_chunk modality for images/videos.""" - return getattr(self._model_config.hf_config, "use_unified_vision_chunk", False) - - @property - def model_config(self) -> ModelConfig: - return self._model_config - - @cached_property - def model_cls(self) -> type[SupportsMultiModal]: - from vllm.model_executor.model_loader import get_model_cls - - model_cls = get_model_cls(self.model_config) - return cast(type[SupportsMultiModal], model_cls) - - @property - def media_io_kwargs(self) -> dict[str, dict[str, Any]] | None: - return self._media_io_kwargs or ( - self._model_config.multimodal_config.media_io_kwargs - if self._model_config.multimodal_config - else None - ) - - @property - def allowed_local_media_path(self): - return self._model_config.allowed_local_media_path - - @property - def allowed_media_domains(self): - return self._model_config.allowed_media_domains - - @property - def mm_registry(self): - return MULTIMODAL_REGISTRY - - @cached_property - def mm_processor(self): - return self.mm_registry.create_processor(self.model_config) - - @property - def video_processor_name(self) -> str | None: - return get_video_processor_cls_name(self.model_config) - - def add(self, modality: ModalityStr, item: _T) -> str | None: - """Add a multi-modal item to the current prompt and returns the - placeholder string to use, if any. - - An optional uuid can be added which serves as a unique identifier of the - media. - - Note: - `prompt_embeds` bypass MM-processor validation because they are - pre-computed embeddings that do not go through any HF processor, encoder, - or model-specific placeholder logic. The corresponding placeholder string is - managed by the parser via `_add_placeholder`, so we return None here. - - """ - add_info = self._validate_add(modality) - if add_info is None: - self._items_by_modality["prompt_embeds"].append(item) - return None - - input_modality, original_modality, use_vision_chunk, num_items = add_info - - # Track original modality for vision_chunk items - if use_vision_chunk: - self._items_by_modality[input_modality].append(item) # type: ignore - self._modality_order["vision_chunk"].append(original_modality) - else: - self._items_by_modality[original_modality].append(item) - - return self.model_cls.get_placeholder_str(modality, num_items) - - def _validate_add(self, modality: ModalityStr) -> tuple[str, str, bool, int] | None: - """Validate that one more item of the modality can be tracked.""" - if modality == "prompt_embeds": - return None - - input_modality = modality.replace("_embeds", "") - original_modality = modality - use_vision_chunk = ( - self.use_unified_vision_chunk_modality - and original_modality in ["video", "image"] - ) - - # If use_unified_vision_chunk_modality is enabled, - # map image/video to vision_chunk - if use_vision_chunk: - # To avoid validation fail - # because models with use_unified_vision_chunk_modality=True - # will only accept vision_chunk modality. - input_modality = "vision_chunk" - num_items = len(self._items_by_modality[input_modality]) + 1 - else: - num_items = len(self._items_by_modality[original_modality]) + 1 - - mm_config = self.model_config.multimodal_config - if ( - mm_config is not None - and mm_config.enable_mm_embeds - and mm_config.get_limit_per_prompt(input_modality) == 0 - and original_modality.endswith("_embeds") - ): - # Skip validation: embeddings bypass limit when enable_mm_embeds=True - pass - else: - self.mm_processor.info.validate_num_items(input_modality, num_items) - - return input_modality, original_modality, use_vision_chunk, num_items - - @abstractmethod - def create_parser( - self, mm_processor_kwargs: dict[str, Any] | None = None - ) -> "BaseMultiModalContentParser": - raise NotImplementedError - - -def _resolve_vision_chunk_items( - vision_chunk_items: list[tuple[object, str | None]], - mm_processor: BaseMultiModalProcessor, - vision_chunks_modality_order: list[str], -): - # Process vision_chunk items - extract from (data, modality) tuples - # and convert to VisionChunk types with proper UUID handling - vision_chunks_uuids = [uuid for data, uuid in vision_chunk_items] - - assert len(vision_chunk_items) == len(vision_chunks_modality_order), ( - f"vision_chunk items ({len(vision_chunk_items)}) and " - f"modality_order ({len(vision_chunks_modality_order)}) must have same length" - ) - - processed_chunks: list[VisionChunk] = [] - video_idx = 0 - for inner_modality, (data, uuid) in zip( - vision_chunks_modality_order, vision_chunk_items - ): - if inner_modality == "image": - # Cast data to proper type for image - # Use .media (PIL.Image) directly to avoid redundant - # bytes→PIL conversion in media_processor - if hasattr(data, "media"): - image_data = data.media # type: ignore[union-attr] - processed_chunks.append( - VisionChunkImage(type="image", image=image_data, uuid=uuid) - ) - else: - processed_chunks.append(data) # type: ignore[arg-type] - elif inner_modality == "video": - # For video, we may need to split into chunks - # if processor supports it - # For now, just wrap as a video chunk placeholder - if hasattr(mm_processor, "split_video_chunks") and data is not None: - try: - video_uuid = uuid or random_uuid() - # video await result is (video_data, video_meta) tuple - if isinstance(data, tuple) and len(data) >= 1: - video_data = data[0] - else: - video_data = data - video_chunks = mm_processor.split_video_chunks(video_data) - for i, vc in enumerate(video_chunks): - processed_chunks.append( - VisionChunkVideo( - type="video_chunk", - video_chunk=vc["video_chunk"], - uuid=f"{video_uuid}-{i}", - video_idx=video_idx, - prompt=vc["prompt"], - ) - ) - video_idx += 1 - except Exception as e: - logger.warning("Failed to split video chunks: %s", e) - processed_chunks.append(data) # type: ignore[arg-type] - else: - processed_chunks.append(data) # type: ignore[arg-type] - return processed_chunks, vision_chunks_uuids - - -def _resolve_items( - items_by_modality: dict[str, list[tuple[object, str | None]]], - mm_processor: BaseMultiModalProcessor | None, - modality_order: dict[str, list[str]], -) -> tuple[MultiModalDataDict, MultiModalUUIDDict]: - """Materialize the tracker's per-modality items into `mm_data` / `mm_uuids`. - - Note: - `mm_processor` is `None` for text-only models (no registered HF - processor) whose only modality is `prompt_embeds`. Every other - modality requires a processor, enforced by the guard below. - - """ - if "image" in items_by_modality and "image_embeds" in items_by_modality: - raise VLLMValidationError( - "Mixing raw image and embedding inputs is not allowed", - parameter="image_embeds", - ) - if "audio" in items_by_modality and "audio_embeds" in items_by_modality: - raise VLLMValidationError( - "Mixing raw audio and embedding inputs is not allowed", - parameter="audio_embeds", - ) - if "video" in items_by_modality and "video_embeds" in items_by_modality: - raise VLLMValidationError( - "Mixing raw video and embedding inputs is not allowed", - parameter="video_embeds", - ) - # `prompt_embeds` bypasses HF MM processors. Every other modality requires one. - processor_modalities = items_by_modality.keys() - {"prompt_embeds"} - if processor_modalities and mm_processor is None: - raise RuntimeError( - _REQUIRE_MM_PROCESSOR_ERROR.format(modality=processor_modalities) - ) - - mm_data = {} - mm_uuids = {} - if "image_embeds" in items_by_modality: - assert mm_processor is not None - mm_data["image"] = _get_embeds_data( - "image", - [data for data, uuid in items_by_modality["image_embeds"]], - mm_processor, - ) - mm_uuids["image"] = [uuid for data, uuid in items_by_modality["image_embeds"]] - if "image" in items_by_modality: - mm_data["image"] = [data for data, uuid in items_by_modality["image"]] - mm_uuids["image"] = [uuid for data, uuid in items_by_modality["image"]] - if "audio_embeds" in items_by_modality: - assert mm_processor is not None - mm_data["audio"] = _get_embeds_data( - "audio", - [data for data, uuid in items_by_modality["audio_embeds"]], - mm_processor, - ) - mm_uuids["audio"] = [uuid for data, uuid in items_by_modality["audio_embeds"]] - if "audio" in items_by_modality: - mm_data["audio"] = [data for data, uuid in items_by_modality["audio"]] - mm_uuids["audio"] = [uuid for data, uuid in items_by_modality["audio"]] - if "video" in items_by_modality: - mm_data["video"] = [data for data, uuid in items_by_modality["video"]] - mm_uuids["video"] = [uuid for data, uuid in items_by_modality["video"]] - if "video_embeds" in items_by_modality: - assert mm_processor is not None - mm_data["video"] = _get_embeds_data( - "video", - [data for data, uuid in items_by_modality["video_embeds"]], - mm_processor, - ) - mm_uuids["video"] = [uuid for data, uuid in items_by_modality["video_embeds"]] - if "vision_chunk" in items_by_modality: - assert mm_processor is not None - # Process vision_chunk items - extract from (data, modality) tuples - # and convert to VisionChunk types with proper UUID handling - processed_chunks, vision_chunk_uuids = _resolve_vision_chunk_items( - items_by_modality["vision_chunk"], - mm_processor, - modality_order.get("vision_chunk", []), - ) - mm_data["vision_chunk"] = processed_chunks - mm_uuids["vision_chunk"] = vision_chunk_uuids - if "prompt_embeds" in items_by_modality: - mm_data["prompt_embeds"] = [ - data for data, _uuid in items_by_modality["prompt_embeds"] - ] - - return mm_data, mm_uuids - - -class MultiModalItemTracker(BaseMultiModalItemTracker[tuple[object, str | None]]): - def resolve_items( - self, - ) -> tuple[MultiModalDataDict | None, MultiModalUUIDDict | None]: - if not self._items_by_modality: - return None, None - - # Text-only models (`is_multimodal_model=False`) with inputs of - # modality `prompt_embeds` have no MM processor since `prompt_embeds` are - # pre-computed and require no processing, so we pass `None`. - mm_processor = ( - self.mm_processor if self._model_config.is_multimodal_model else None - ) - return _resolve_items( - dict(self._items_by_modality), - mm_processor, - self._modality_order, - ) - - def create_parser( - self, mm_processor_kwargs: dict[str, Any] | None = None - ) -> "BaseMultiModalContentParser": - return MultiModalContentParser(self, mm_processor_kwargs=mm_processor_kwargs) - - -class AsyncMultiModalItemTracker(BaseMultiModalItemTracker[_AsyncMultiModalItem]): - async def resolve_items( - self, - ) -> tuple[MultiModalDataDict | None, MultiModalUUIDDict | None]: - if not self._items_by_modality: - return None, None - - # Fetch all modalities together. Each tracked item is already an - # independent awaitable, and the async connector offloads blocking - # decode work, so waiting for one modality before starting the next - # needlessly adds their latency. - # Keep the original group and item order when rebuilding the result. - item_groups = list(self._items_by_modality.items()) - items = [item for _, group in item_groups for item in group] - results = await asyncio.gather( - *(item() for item in items), return_exceptions=True - ) - for result in results: - if isinstance(result, BaseException): - # Gathering with return_exceptions=True lets every task finish - # (or itself fail) before we raise, instead of abandoning - # still-in-flight fetches (real network/thread-pool work) the - # moment the first one fails. - raise result - - resolved_items_by_modality: dict[str, list[Any]] = {} - result_idx = 0 - for modality, group in item_groups: - next_result_idx = result_idx + len(group) - resolved_items_by_modality[modality] = results[result_idx:next_result_idx] - result_idx = next_result_idx - - mm_processor = ( - self.mm_processor if self._model_config.is_multimodal_model else None - ) - return _resolve_items( - resolved_items_by_modality, - mm_processor, - self._modality_order, - ) - - def create_parser( - self, mm_processor_kwargs: dict[str, Any] | None = None - ) -> "BaseMultiModalContentParser": - return AsyncMultiModalContentParser( - self, mm_processor_kwargs=mm_processor_kwargs - ) - - -class BaseMultiModalContentParser(ABC): - def __init__(self) -> None: - super().__init__() - - # stores model placeholders list with corresponding - # general MM placeholder: - # { - # "<##IMAGE##>": ["", "", ""], - # "<##AUDIO##>": ["