diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 029ccbbffa81..4672fcb2c2d8 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -11,9 +11,12 @@ steps: - "!vllm/distributed/kv_transfer/" - tests/models/test_initialization.py - tests/models/registry.py + - tests/models/test_beam_construction.py commands: # Run a subset of model initialization tests - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset + # Exercise Beam construction and generation without a checkpoint download. + - pytest -v -s models/test_beam_construction.py mirror: amd: label: ":amd: (MI355 DPX) Basic Models (Initialization)" @@ -257,6 +260,8 @@ steps: - tests/models/test_qwen3_5_mtp_config.py - tests/models/test_alice_ai.py - tests/models/test_alice_ai_config.py + - tests/models/test_beam_weight_mapping.py + - tests/models/test_beam_numerics.py - tests/models/transformers/fusers/ - tests/models/transformers/test_dynamic_rope.py - tests/model_executor/layers/test_activation.py @@ -267,7 +272,7 @@ steps: - tests/models/transformers/test_layer_registry.py 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/ models/transformers/test_dynamic_rope.py 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 models/test_alice_ai.py models/test_alice_ai_config.py + - 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/ models/transformers/test_dynamic_rope.py 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 models/test_alice_ai.py models/test_alice_ai_config.py models/test_beam_weight_mapping.py models/test_beam_numerics.py - label: ':amd: (MI355) DeepSeek-V4.1 Mono Decode Numerics' key: dsv41-mono-decode-numerics-amd diff --git a/csrc/fp32_router_gemm.cu b/csrc/fp32_router_gemm.cu index 9719adb513f5..832874f914a6 100644 --- a/csrc/fp32_router_gemm.cu +++ b/csrc/fp32_router_gemm.cu @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project // -// Router GEMM: activation(T) x weight(fp32) -> fp32, M<=32, for the +// Router GEMM: activation(T) x weight(fp32) -> fp32, M<=48, for the // supported (E, H) pairs listed at the bottom of this file. // Supports bf16 or fp32 activation; weight is always fp32. @@ -214,8 +214,26 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a, // production-fidelity harness (CUDA-graph replay, per-layer cold weights). // Only enabled on the Blackwell family where it was validated; Hopper and // other shapes / fp32 activation keep the legacy geometry. - if constexpr (std::is_same_v && kNumExperts == 128 && - kHiddenDim == 6144) { + if constexpr (std::is_same_v && kNumExperts == 256 && + kHiddenDim == 5120) { + if (!isBlackwellFamily()) { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + return; + } + if constexpr (kNumTokens > 16 && kNumTokens % 8 == 0) { + constexpr int block_size = kNumTokens == 24 ? 160 : 128; + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else { + constexpr int block_size = kNumTokens == 1 ? 64 : 128; + constexpr int token_groups = kNumTokens % 2 == 0 ? 2 : 1; + launchFp32RouterGemm(output, mat_a, mat_b, + stream); + } + } else if constexpr (std::is_same_v && + kNumExperts == 128 && kHiddenDim == 6144) { // MiniMax-M3. Legacy 128/1 only fills 128 blocks and pays the same // accumulator register cliffs; B300 sweep: // even M in [6, 10] : BS=384, EPB=1, 2 token groups (1.26-1.43x) @@ -268,8 +286,8 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a, } // --------------------------------------------------------------------------- -// Explicit instantiations: M=1..32, for both input types, for the supported -// (E, H) pairs: (256, 3072) [MiniMax-M2/M2.5] and (128, 6144) [MiniMax-M3]. +// Explicit instantiations for both input types. +// Supported (E, H) pairs: (256, 3072), (128, 6144), and (256, 5120). // --------------------------------------------------------------------------- #define INSTANTIATE(T, M, E, H) \ @@ -315,5 +333,31 @@ INSTANTIATE_ALL(__nv_bfloat16, 256, 3072) INSTANTIATE_ALL(float, 128, 6144) INSTANTIATE_ALL(__nv_bfloat16, 128, 6144) +INSTANTIATE_ALL(float, 256, 5120) +INSTANTIATE_ALL(__nv_bfloat16, 256, 5120) + +#define INSTANTIATE_LARGE(T, E, H) \ + INSTANTIATE(T, 33, E, H) \ + INSTANTIATE(T, 34, E, H) \ + INSTANTIATE(T, 35, E, H) \ + INSTANTIATE(T, 36, E, H) \ + INSTANTIATE(T, 37, E, H) \ + INSTANTIATE(T, 38, E, H) \ + INSTANTIATE(T, 39, E, H) \ + INSTANTIATE(T, 40, E, H) \ + INSTANTIATE(T, 41, E, H) \ + INSTANTIATE(T, 42, E, H) \ + INSTANTIATE(T, 43, E, H) \ + INSTANTIATE(T, 44, E, H) \ + INSTANTIATE(T, 45, E, H) \ + INSTANTIATE(T, 46, E, H) \ + INSTANTIATE(T, 47, E, H) \ + INSTANTIATE(T, 48, E, H) + +INSTANTIATE_LARGE(float, 256, 5120) +INSTANTIATE_LARGE(__nv_bfloat16, 256, 5120) + +#undef INSTANTIATE_LARGE + #undef INSTANTIATE_ALL #undef INSTANTIATE diff --git a/csrc/fp32_router_gemm_entry.cu b/csrc/fp32_router_gemm_entry.cu index e55ac7e72104..a6efc32831a1 100644 --- a/csrc/fp32_router_gemm_entry.cu +++ b/csrc/fp32_router_gemm_entry.cu @@ -25,10 +25,11 @@ inline int getSMVersion() { static constexpr int FP32_MAX_TOKENS = 32; // Supported (hidden_dim, num_experts) pairs (must match the instantiations in -// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3. +// fp32_router_gemm.cu). static inline bool fp32_router_gemm_supported(int hidden_dim, int num_experts) { return (hidden_dim == 3072 && num_experts == 256) || - (hidden_dim == 6144 && num_experts == 128); + (hidden_dim == 6144 && num_experts == 128) || + (hidden_dim == 5120 && num_experts == 256); } // Forward declarations — 4 template params must match fp32_router_gemm.cu @@ -60,8 +61,7 @@ struct Fp32LoopUnroller { invokeFp32RouterGemm( output, mat_a, mat_b, stream); } else { - throw std::invalid_argument( - "fp32_router_gemm: num_tokens must be in [1, 32]"); + throw std::invalid_argument("fp32_router_gemm: unsupported num_tokens"); } } }; @@ -77,6 +77,9 @@ void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens, } else if (num_experts == 128 && hidden_dim == 6144) { Fp32LoopUnroller::unroll( num_tokens, output, mat_a, mat_b, stream); + } else if (num_experts == 256 && hidden_dim == 5120) { + Fp32LoopUnroller::unroll(num_tokens, output, + mat_a, mat_b, stream); } else { throw std::invalid_argument( "fp32_router_gemm: unsupported (hidden_dim, num_experts) pair"); @@ -111,9 +114,11 @@ void fp32_router_gemm( STD_TORCH_CHECK( fp32_router_gemm_supported(hidden_dim, num_experts), "fp32_router_gemm: supported (hidden_dim, num_experts) pairs are " - "(3072, 256) and (6144, 128)"); - STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS, - "fp32_router_gemm: num_tokens must be in [0, 32]"); + "(3072, 256), (6144, 128), and (5120, 256)"); + const int max_tokens = + hidden_dim == 5120 && num_experts == 256 ? 48 : FP32_MAX_TOKENS; + STD_TORCH_CHECK(num_tokens <= max_tokens, + "fp32_router_gemm: num_tokens exceeds supported range"); STD_TORCH_CHECK( mat_a.scalar_type() == torch::headeronly::ScalarType::Float || mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16, diff --git a/docs/features/tool_calling.md b/docs/features/tool_calling.md index dd56db0d579c..9980429324ce 100644 --- a/docs/features/tool_calling.md +++ b/docs/features/tool_calling.md @@ -201,6 +201,30 @@ checkpoint's output format. !!! note With `tool_choice="auto"`, structural-tag constraints require both `VLLM_ENFORCE_STRICT_TOOL_CALLING=true` (the default) and at least one tool with `strict: true`, or a server-side floor set via `--tool-strict-level`. When these conditions are met and the selected parser supports structural tags, vLLM constrains the tool-call envelope and pins the argument schema of each tool that sets `strict: true` (or of every tool under `--tool-strict-level parameter`). Otherwise, vLLM extracts tool calls from raw text, so arguments may occasionally be malformed or violate the function's parameter schema. +### Beam Models (`beam`) + +Beam checkpoints describe their output format with `response_template` metadata, so the `beam` parser reads it like the `hf` parser does: + +```bash +vllm serve \ + --enable-auto-tool-choice \ + --tool-call-parser beam \ + --reasoning-parser beam +``` + +Unlike `hf`, the `beam` parser also constrains its output. For strict tools, required or named tool choice, and structured outputs, it derives a +structural-tag grammar from the template. The grammar follows the template's tool call format, pins the arguments of strict tools when their +encoding can be derived, allows a single call with `parallel_tool_calls=false`, and frames a structured output as the answer. Raw string arguments +keep the `format` or `pattern` of their schema, as XGrammar keeps them in JSON strings, unless a match could end the value early. With +`tool_choice="none"`, the grammar allows only reasoning and an answer, so a model that still sees the tools in its prompt cannot call them. When no +grammar can be derived, required and named tool choice are rejected, and strict tools under `tool_choice="auto"` are left unconstrained with a +warning. + +It also constrains the output of requests without strict tools, forced tool choice, or structured outputs to the template's format: tool calls with +free arguments or an answer, or just the rest of an answer the prompt already opened. This always-on grammar needs the `xgrammar` structured +outputs backend, the default, and is disabled with `VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT=0`. From the message start, every `beam` grammar opens with +reasoning that holds visible text, unless `enable_thinking` is false. + ### Hermes Models (`hermes`) All Nous Research Hermes-series models newer than Hermes 2 Pro should be supported. diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 0b69ac8f4753..2166c3a8eb08 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -379,6 +379,7 @@ th { | `BailingMoeV2ForCausalLM` | Ling | `inclusionAI/Ling-mini-2.0`, etc. | ✅︎ | ✅︎ | | `BailingMoeV2_5ForCausalLM` | Ling | `inclusionAI/Ling-2.5-1T`, `inclusionAI/Ring-2.5-1T` | | ✅︎ | | `BailingMoeV3ForCausalLM` | Ling | `inclusionAI/Ling-3.0-flash` | | ✅︎ | +| `BeamForCausalLM` | Beam | `reflection/Beam-1.0-501B-A23B-NVFP4`, etc. | | ✅︎ | | `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ | | `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ | | `CohereForCausalLM`, `Cohere2ForCausalLM` | Command-R, Command-A | `CohereLabs/c4ai-command-r-v01`, `CohereLabs/c4ai-command-r7b-12-2024`, `CohereLabs/c4ai-command-a-03-2025`, `CohereLabs/command-a-reasoning-08-2025`, etc. | ✅︎ | ✅︎ | diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 1ffb95fe377e..5abdd36cc9fd 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -434,7 +434,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_json, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string()); + expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: beam, 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] @@ -445,6 +445,6 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, hf, hy_v3, hy_v4, inkling, kimi, kimi_k2, kimi_k3, mimo, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); + expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: beam, cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, hf, hy_v3, hy_v4, inkling, kimi, kimi_k2, kimi_k3, mimo, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); } } diff --git a/rust/src/chat/src/output/default/mod.rs b/rust/src/chat/src/output/default/mod.rs index 67a5a11eae6f..dbbedfd8369c 100644 --- a/rust/src/chat/src/output/default/mod.rs +++ b/rust/src/chat/src/output/default/mod.rs @@ -23,7 +23,8 @@ use crate::output::{ChatOutputProcessor, DynChatEventStream, DynDecodedTextEvent use crate::parser::reasoning::{ReasoningParser, ReasoningParserFactory}; use crate::parser::tool::{ToolParser, ToolParserFactory}; use crate::parser::unified::{ - HfTemplateError, HfUnifiedParser, ResponseTemplate, UnifiedParserFactory, names, + BeamUnifiedParser, HfTemplateError, HfUnifiedParser, ResponseTemplate, UnifiedParserFactory, + names, }; use crate::parser::{ParserSelection, ToolStrictLevel}; use crate::request::{ChatRequest, ChatTool}; @@ -206,17 +207,19 @@ impl DefaultChatOutputProcessor { }); } - // `hf` is built from the model's template, never by its registry constructor. - let parser = if parser_name == names::HF { - let template = - response_template.as_ref().map_err(|error| Error::ParserInitialization { - kind: "unified", - name: parser_name.to_string(), - error: Box::new(error.clone()), - })?; - Box::new(HfUnifiedParser::new(template.clone(), tools, tokenizer)) - } else { - factory.create(parser_name, tools, tokenizer)? + // Parsers driven by the model's template are built from it, never by + // their registry constructors. + let template = || { + response_template.clone().map_err(|error| Error::ParserInitialization { + kind: "unified", + name: parser_name.to_string(), + error: Box::new(error), + }) + }; + let parser: Box = match parser_name { + names::HF => Box::new(HfUnifiedParser::new(template()?, tools, tokenizer)), + names::BEAM => Box::new(BeamUnifiedParser::new(template()?, tools, tokenizer)), + _ => factory.create(parser_name, tools, tokenizer)?, }; UNIFIED_PARSER_LOG_ONCE.call_once(|| info!(parser_name, "using unified parser")); @@ -589,6 +592,20 @@ mod tests { .assert_eq(&error(invalid_template())); } + #[test] + fn explicit_beam_uses_the_model_response_template() { + let beam = ParserSelection::Explicit("beam".to_string()); + crate::parser::validate_parser_overrides(&beam, &beam).unwrap(); + assert!(keeps_special_tokens("other-model", &loaded_template(), &beam, &beam).unwrap()); + + let error = + keeps_special_tokens("other-model", &Err(HfTemplateError::Missing), &beam, &beam) + .unwrap_err() + .to_report_string(); + expect_test::expect!["failed to initialize unified parser `beam`: the model's tokenizer_config.json provides no response_template"] + .assert_eq(&error); + } + #[test] fn auto_never_selects_hf() { let auto = ParserSelection::Auto; diff --git a/rust/src/chat/src/parser/unified.rs b/rust/src/chat/src/parser/unified.rs index 7db093c69684..705133bb0a71 100644 --- a/rust/src/chat/src/parser/unified.rs +++ b/rust/src/chat/src/parser/unified.rs @@ -6,8 +6,8 @@ use std::sync::LazyLock; pub use vllm_parser::unified::{ - Gemma4UnifiedParser, HfTemplateError, HfUnifiedParser, HyV3UnifiedParser, HyV4UnifiedParser, - InklingUnifiedParser, KimiK3UnifiedParser, ResponseTemplate, UnifiedParser, + BeamUnifiedParser, Gemma4UnifiedParser, HfTemplateError, HfUnifiedParser, HyV3UnifiedParser, + HyV4UnifiedParser, InklingUnifiedParser, KimiK3UnifiedParser, ResponseTemplate, UnifiedParser, }; use vllm_tokenizer::DynTokenizer; @@ -26,6 +26,9 @@ pub mod names { /// name resolution only: it is built from the model's template, never by its /// registry constructor. pub const HF: &str = "hf"; + /// Beam checkpoints: [`HF`] with Beam's tool-call grammar, built from the + /// model's template the same way. + pub const BEAM: &str = "beam"; } /// Constructor signature for one registered unified parser implementation. @@ -55,6 +58,7 @@ impl UnifiedParserFactory { factory.register_parser::(names::KIMI_K3); // Opt-in only: no model patterns map to it. factory.register_parser::(names::HF); + factory.register_parser::(names::BEAM); factory .register_pattern("gemma-4", names::GEMMA4) diff --git a/rust/src/parser/src/output_grammar.rs b/rust/src/parser/src/output_grammar.rs index ba9b4968d54b..0d05bb4ddefa 100644 --- a/rust/src/parser/src/output_grammar.rs +++ b/rust/src/parser/src/output_grammar.rs @@ -207,7 +207,8 @@ pub(crate) fn full_format_from_builder_for_test( ) } -fn tool_params(tools: &[Tool], strict_level: ToolStrictLevel) -> Vec { +/// Tools as builder inputs, strict when the tool or the server floor asks for it. +pub(crate) fn tool_params(tools: &[Tool], strict_level: ToolStrictLevel) -> Vec { tools .iter() .map(|tool| { @@ -240,7 +241,7 @@ pub(crate) fn normalize_tool_choice( } /// Whether the request asks for a tool grammar at all. -fn tool_grammar_applies(ctx: &OutputGrammarContext<'_>) -> bool { +pub(crate) fn tool_grammar_applies(ctx: &OutputGrammarContext<'_>) -> bool { if ctx.tools.is_empty() { return false; } diff --git a/rust/src/parser/src/output_grammar/arguments.rs b/rust/src/parser/src/output_grammar/arguments.rs index 0f621789a3bc..6e58c6956a73 100644 --- a/rust/src/parser/src/output_grammar/arguments.rs +++ b/rust/src/parser/src/output_grammar/arguments.rs @@ -26,7 +26,7 @@ //! - `allOf` with more than one schema accepts any value. use serde_json::{Value, json}; -use xgrammar_structural_tag::format::{Format, JsonSchemaFormat}; +use xgrammar_structural_tag::format::{EndBoundary, Format, JsonSchemaFormat, TokenFormat}; pub use crate::schema::JsonType; use crate::schema::{self, OptionSource, SchemaOption, SchemaRoot}; @@ -60,7 +60,13 @@ impl ParameterKey<'_> { /// `prefix KEY suffix content end`: a tag for a declared key, or a /// sequence matching [`Self::FREE_PATTERN`] for a free key. - pub fn tag(&self, prefix: &str, suffix: &str, content: Format, end: &str) -> Format { + pub fn tag( + &self, + prefix: &str, + suffix: &str, + content: Format, + end: impl Into, + ) -> Format { match self { Self::Declared(key) => Format::tag(format!("{prefix}{key}{suffix}"), content, end), Self::Free => Format::sequence(vec![ @@ -68,7 +74,13 @@ impl ParameterKey<'_> { Format::regex(Self::FREE_PATTERN), Format::const_string(suffix), content, - Format::const_string(end), + match end.into() { + EndBoundary::Text(end) => Format::const_string(end), + EndBoundary::Texts(ends) => { + one_of(ends.into_iter().map(Format::const_string).collect()) + } + EndBoundary::Token(end) => Format::Token(TokenFormat { token: end.token }), + }, ]), } } @@ -183,6 +195,36 @@ impl ArgumentContext<'_> { /// Build the grammar of one call's arguments under `schema`, the tool's /// parameters schema, rendering each parameter with `syntax`. pub fn arguments(schema: &Value, syntax: &dyn ArgumentSyntax, options: &ArgumentOptions) -> Format { + parameter_list( + options, + syntax.separator(), + parameters(schema, syntax, options), + ) +} + +/// The parameters of one call, each rendered by an [`ArgumentSyntax`], for a +/// syntax that lists them itself rather than through [`arguments`]. +pub struct Parameters { + /// The declared parameters, in declared order. + pub declared: Vec, + /// Any one additional parameter, if the schema allows them. + pub additional: Option, +} + +/// One declared parameter. +pub struct Parameter { + pub key: String, + pub format: Format, + pub required: bool, +} + +/// The parameters of one call under `schema`, the tool's parameters schema, +/// rendering each parameter with `syntax`. +pub fn parameters( + schema: &Value, + syntax: &dyn ArgumentSyntax, + options: &ArgumentOptions, +) -> Parameters { let cx = ArgumentContext { options, root: SchemaRoot::new(schema), @@ -190,8 +232,7 @@ pub fn arguments(schema: &Value, syntax: &dyn ArgumentSyntax, options: &Argument let parameter = |key: ParameterKey<'_>, schema: &Value| syntax.parameter(key, &cx.value_options(schema)); - let schema = cx.root.resolved(); - let (properties, additional) = match schema { + match cx.root.resolved() { Value::Object(schema) => { let additional = match schema.get("additionalProperties") { Some(Value::Bool(false)) => None, @@ -206,22 +247,33 @@ pub fn arguments(schema: &Value, syntax: &dyn ArgumentSyntax, options: &Argument .and_then(Value::as_array) .map(|required| required.iter().filter_map(Value::as_str).collect::>()) .unwrap_or_default(); - let properties = schema + let declared = schema .get("properties") .and_then(Value::as_object) .into_iter() .flatten() .filter_map(|(key, schema)| { - let parameter = parameter(ParameterKey::Declared(key), schema)?; - Some((parameter, required.contains(&key.as_str()))) + Some(Parameter { + key: key.clone(), + format: parameter(ParameterKey::Declared(key), schema)?, + required: required.contains(&key.as_str()), + }) }) .collect(); - (properties, additional) + Parameters { + declared, + additional, + } } - Value::Bool(false) => (vec![], None), - _ => (vec![], parameter(ParameterKey::Free, &ANY_SCHEMA)), - }; - parameter_list(options, syntax.separator(), properties, additional) + Value::Bool(false) => Parameters { + declared: vec![], + additional: None, + }, + _ => Parameters { + declared: vec![], + additional: parameter(ParameterKey::Free, &ANY_SCHEMA), + }, + } } /// Parameters in declared order, required ones mandatory, followed by any @@ -230,8 +282,10 @@ pub fn arguments(schema: &Value, syntax: &dyn ArgumentSyntax, options: &Argument fn parameter_list( options: &ArgumentOptions, separator: Option, - properties: Vec<(Format, bool)>, - additional: Option, + Parameters { + declared, + additional, + }: Parameters, ) -> Format { let separated = |parameter: Format| match &separator { Some(separator) => Format::sequence(vec![parameter, separator.clone()]), @@ -239,10 +293,10 @@ fn parameter_list( }; let mut elements = Vec::from_iter(separator.clone()); if options.any_order { - let any_required = properties.iter().any(|(_, required)| *required); - let parameters = properties + let any_required = declared.iter().any(|parameter| parameter.required); + let parameters = declared .into_iter() - .map(|(parameter, _)| parameter) + .map(|parameter| parameter.format) .chain(additional) .collect::>(); if !parameters.is_empty() { @@ -254,8 +308,11 @@ fn parameter_list( }); } } else { - for (parameter, required) in properties { - let parameter = separated(parameter); + for Parameter { + format, required, .. + } in declared + { + let parameter = separated(format); elements.push(if required { parameter } else { diff --git a/rust/src/parser/src/unified/beam.rs b/rust/src/parser/src/unified/beam.rs new file mode 100644 index 000000000000..fb3b904e4ea9 --- /dev/null +++ b/rust/src/parser/src/unified/beam.rs @@ -0,0 +1,543 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Unified parser for Beam checkpoints. +//! +//! Beam output is parsed by [`HfUnifiedParser`] with the checkpoint's +//! `response_template`; this parser adds Beam's output grammar. +//! A Beam message after `<|start|>assistant` is: +//! +//! ```text +//! <|channel|>analysis<|message|>REASONING<|end|><|start|>assistant +//! ( <|channel|>final<|message|>ANSWER<|return|> +//! | CALL (<|end|><|start|>assistant CALL)* <|call|> ) +//! ``` +//! +//! where reasoning may end at any closer of the template's reasoning field, +//! and `CALL` is +//! ` to=NAME<|channel|>analysis <|constrain|>xml<|message|>` followed by +//! `\nVALUE\n` parameters separated +//! by newlines. Where the template looks ahead for the next parameter, a value +//! may also end on a newline right before it. +//! +//! The grammar enforces the schema of every declared tool's arguments, in +//! declared or alphabetical order, since Beam drifts from schemas it does not +//! enforce. When the model may answer instead, calls may also name any other +//! tool with any arguments, which the parser passes on so a harness can tell +//! the model the tool does not exist. `ANSWER` follows the request's structured +//! output, if any. +//! +//! The chat template rejects a history turn without reasoning, so a message the +//! model opens holds non-empty reasoning: an empty one would come back as null +//! reasoning and fail the next request. +//! +//! Unless `VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT` turns it off, the grammar +//! applies to every message, including one the request does not constrain. + +use std::collections::BTreeMap; +use std::sync::{Arc, LazyLock}; + +use serde_json::Value; +use vllm_tokenizer::{DecodedText, DynTokenizer}; +use xgrammar_structural_tag::ToolChoice; +use xgrammar_structural_tag::format::{EndBoundary, Format}; +use xgrammar_structural_tag::tool::{ + BuilderToolChoice, NormalizedToolChoice, ToolChoiceValue, normalize_tool_choice, +}; + +use super::{ + HfPosition, HfUnifiedParser, ResponseTemplate, Result, UnifiedParser, UnifiedParserError, + UnifiedParserOutput, +}; +use crate::output_grammar::arguments::{ + self, ArgumentOptions, ArgumentSyntax, JsonType, ParameterKey, Parameters, ValueOption, one_of, +}; +use crate::output_grammar::{ + self, BuiltOutputGrammar, OutputGrammarContext, tool_grammar_applies, tool_params, +}; +use crate::tool::Tool; + +const ANALYSIS_OPEN: &str = "<|channel|>analysis<|message|>"; +const FINAL_OPEN: &str = "<|channel|>final<|message|>"; +const END: &str = "<|end|>"; +/// Closes one message and opens the next assistant message. +const NEXT_MESSAGE: &str = "<|end|><|start|>assistant"; +const RETURN: &str = "<|return|>"; +const CALL: &str = "<|call|>"; +/// Follows ` to=NAME` in a call. +const CALL_HEADER_END: &str = "<|channel|>analysis <|constrain|>xml<|message|>"; +const PARAMETER_OPEN: &str = "\n"; +const PARAMETER_CLOSE: &str = "\n"; +/// Ends a value where the template looks ahead for the next parameter. +const NEXT_PARAMETER: &str = "\n\n = LazyLock::new(|| { + std::env::var("VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT").map_or(true, |value| { + matches!(value.to_lowercase().as_str(), "true" | "1") + }) +}); + +/// Beam parser: [`HfUnifiedParser`] with Beam's output grammar. +pub struct BeamUnifiedParser { + inner: HfUnifiedParser, + /// Whether a value may end where the next parameter opens. + lenient: bool, + /// What ends the reasoning and opens the next message. + reasoning_ends: Vec, +} + +impl BeamUnifiedParser { + /// Create a parser for one request from the checkpoint's `response_template`. + pub fn new(template: Arc, tools: &[Tool], tokenizer: DynTokenizer) -> Self { + Self { + lenient: template.argument_lookahead() == Some(PARAMETER_OPEN), + reasoning_ends: match template.reasoning_closers() { + [] => vec![NEXT_MESSAGE.to_string()], + closers => closers.to_vec(), + }, + inner: HfUnifiedParser::new(template, tools, tokenizer), + } + } +} + +impl UnifiedParser for BeamUnifiedParser { + fn create(_tools: &[Tool], _tokenizer: DynTokenizer) -> Result> { + Err(UnifiedParserError::NoNamedConstructor { + parser: "beam", + built_from: "the checkpoint's response_template", + }) + } + + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> { + self.inner.initialize(prompt_token_ids) + } + + fn preserve_special_tokens(&self) -> bool { + self.inner.preserve_special_tokens() + } + + fn build_output_grammar( + &self, + ctx: &OutputGrammarContext<'_>, + ) -> output_grammar::Result> { + build_output_grammar( + self.inner.position(), + ctx, + *ALWAYS_CONSTRAIN_OUTPUT, + self.lenient, + &self.reasoning_ends, + ) + } + + fn tool_call_id(&self, tool_index: usize) -> Option<&str> { + self.inner.tool_call_id(tool_index) + } + + fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> { + self.inner.parse_into(delta, output) + } + + fn finish(&mut self) -> Result { + self.inner.finish() + } + + fn reset(&mut self) -> String { + self.inner.reset() + } +} + +/// The grammar of a Beam message from the first generated token, or `None` +/// when the parser is neither at the start of the message nor inside its +/// reasoning, or when the request constrains nothing and `always_constrain` is +/// off. Values may end where the next parameter opens if `lenient`, and +/// reasoning ends at any of `reasoning_ends`. +fn build_output_grammar( + position: HfPosition, + ctx: &OutputGrammarContext<'_>, + always_constrain: bool, + lenient: bool, + reasoning_ends: &[String], +) -> output_grammar::Result> { + let forced = !ctx.tools.is_empty() + && matches!( + ctx.tool_choice, + ToolChoice::Value(ToolChoiceValue::Required) | ToolChoice::NamedFunction(_) + ); + let answer_format = ctx.answer.filter(|_| !forced); + // The prompt may still list the tools, so only the grammar keeps the model + // from calling them. + let tools_forbidden = !ctx.tools.is_empty() + && matches!(ctx.tool_choice, ToolChoice::Value(ToolChoiceValue::None)); + if !(always_constrain + || tool_grammar_applies(ctx) + || tools_forbidden + || answer_format.is_some()) + { + return Ok(None); + } + let text = Format::any_text(); + let reasoning_end = EndBoundary::from(reasoning_ends.to_vec()); + let reasoning = match position { + HfPosition::ResponseStart => Format::tag( + ANALYSIS_OPEN, + // At least one character, which cannot start Beam's closers. + Format::sequence(vec![Format::regex("[^<]"), text]), + reasoning_end, + ), + HfPosition::Reasoning => Format::tag("", text, reasoning_end), + // Inside an answer the prompt opened, only the always-on grammar + // applies, finishing a free answer. A tool choice that forces a call + // cannot be met there, and the request's own constraint applies as is. + HfPosition::Content => { + return Ok( + (always_constrain && !forced && answer_format.is_none()).then(|| { + BuiltOutputGrammar::from_token_zero(Format::sequence(answer_rest(None))) + }), + ); + } + HfPosition::Other => return Ok(None), + }; + let NormalizedToolChoice { + function_tools, + choice, + .. + } = normalize_tool_choice( + &tool_params(ctx.tools, ctx.tool_strict_level), + ctx.tool_choice.clone(), + )?; + if function_tools.is_empty() { + return Ok(match choice { + // `none`, or no tools to call: the model can only answer. + BuilderToolChoice::Auto => { + Some(BuiltOutputGrammar::from_token_zero(Format::sequence(vec![ + reasoning, + answer(answer_format), + ]))) + } + BuilderToolChoice::Required | BuilderToolChoice::Forced => None, + }); + } + + let mut calls = function_tools + .iter() + .map(|tool| { + let header = + Format::const_string(format!(" to={}{CALL_HEADER_END}", tool.function.name)); + let schema = tool.function.parameters.clone().unwrap_or(Value::Bool(true)); + Format::sequence(vec![header, arguments(&schema, lenient)]) + }) + .collect::>(); + if matches!(choice, BuilderToolChoice::Auto) { + let names = ctx.tools.iter().map(|tool| tool.name.as_str()).collect::>(); + let header = Format::sequence(vec![ + Format::const_string(" to="), + Format::regex(other_names(&names)), + Format::const_string(CALL_HEADER_END), + ]); + calls.push(Format::sequence(vec![ + header, + arguments(&Value::Bool(true), lenient), + ])); + } + let call = one_of(calls); + let mut calls = vec![call.clone()]; + if ctx.parallel_tool_calls && !choice.is_forced() { + calls.push(Format::star(Format::sequence(vec![ + Format::const_string(NEXT_MESSAGE), + call, + ]))); + } + // The last call ends with the `<|call|>` stop token. + calls.push(Format::optional(Format::const_string(CALL))); + let calls = Format::sequence(calls); + let output = match choice { + BuilderToolChoice::Auto => Format::or(vec![calls, answer(answer_format)]), + BuilderToolChoice::Required | BuilderToolChoice::Forced => calls, + }; + Ok(Some(BuiltOutputGrammar::from_token_zero(Format::sequence( + vec![reasoning, output], + )))) +} + +/// A final answer, ending with the `<|return|>` stop token. +fn answer(format: Option<&Format>) -> Format { + Format::sequence([vec![Format::const_string(FINAL_OPEN)], answer_rest(format)].concat()) +} + +/// The rest of an open final answer, holding `format` or else free text. +fn answer_rest(format: Option<&Format>) -> Vec { + vec![ + // Without `<|end|>`, no unconstrained message can follow the answer. + format.map_or_else(|| Format::any_text_excluding(&[END, RETURN]), Format::clone), + Format::optional(Format::const_string(RETURN)), + ] +} + +/// The arguments under `schema`: its declared parameters in declared or +/// alphabetical order, required ones mandatory, then any additional ones. If +/// `lenient`, each value but the last may end by opening the next parameter. +fn arguments(schema: &Value, lenient: bool) -> Format { + let options = ArgumentOptions { + excludes: CALL_TERMINATORS.map(str::to_string).to_vec(), + ..ArgumentOptions::default() + }; + let render = |syntax: BeamArguments| arguments::parameters(schema, &syntax, &options); + if lenient { + lenient_arguments(render(BeamArguments::Joined), render(BeamArguments::Last)) + } else { + closed_arguments(render(BeamArguments::Closed)) + } +} + +/// Arguments whose values all end with a closer. +fn closed_arguments( + Parameters { + declared, + additional, + }: Parameters, +) -> Format { + let additional = additional.map(Format::star); + let lists = orders(&declared, |parameter| ¶meter.key) + .into_iter() + .map(|order| { + let listed = order + .into_iter() + .map(|parameter| listed(parameter.format.clone(), parameter.required)); + listed.chain(additional.clone()).collect::>() + }) + .filter(|list| !list.is_empty()) + .map(Format::sequence) + .collect::>(); + if lists.is_empty() { + Format::const_string("") + } else { + one_of(lists) + } +} + +/// A declared parameter as one another parameter follows, `joined`, and as the +/// last one. +struct LenientParameter { + key: String, + joined: Format, + last: Format, + required: bool, +} + +impl LenientParameter { + fn joined(&self) -> Format { + listed(self.joined.clone(), self.required) + } +} + +/// Arguments whose values end by opening the next parameter, but for the last +/// one, from the parameters as [`BeamArguments::Joined`] and as +/// [`BeamArguments::Last`]. +fn lenient_arguments(joined: Parameters, last: Parameters) -> Format { + // Both forms share their terminators, so they keep the same parameters. + let declared = joined + .declared + .into_iter() + .zip(last.declared) + .map(|(joined, last)| LenientParameter { + key: joined.key, + joined: joined.format, + last: last.format, + required: joined.required, + }) + .collect::>(); + let additional = joined.additional.zip(last.additional); + + let mut lists = vec![]; + for order in orders(&declared, |parameter| ¶meter.key) { + // The last parameter is a declared one that follows every required + // one, or an additional one. + if !order.is_empty() { + let first = order.iter().rposition(|parameter| parameter.required).unwrap_or(0); + let mut list = + order[..first].iter().map(|parameter| parameter.joined()).collect::>(); + list.push(last_present(&order[first..])); + lists.push(Format::sequence(list)); + } + if let Some((joined, last)) = &additional { + let mut list = order.iter().map(|parameter| parameter.joined()).collect::>(); + list.extend([Format::star(joined.clone()), last.clone()]); + lists.push(Format::sequence(list)); + } + } + if lists.is_empty() { + return Format::const_string(""); + } + + // Each parameter starts past the opener, as the value before writes it. + let arguments = Format::sequence(vec![ + separator(), + Format::const_string(PARAMETER_OPEN), + one_of(lists), + ]); + if declared.iter().any(|parameter| parameter.required) { + arguments + } else { + Format::optional(arguments) + } +} + +/// `parameters`, at least one, in order and the last present one in its last +/// form. Only the first may be required. +fn last_present(parameters: &[&LenientParameter]) -> Format { + if let [only] = parameters { + return only.last.clone(); + } + // The last one is in the first half, or after all of it. Halving keeps the + // grammar O(n log n) in size and O(log n) deep. + let (head, tail) = parameters.split_at(parameters.len() / 2); + let mut after_head = head.iter().map(|parameter| parameter.joined()).collect::>(); + after_head.push(last_present(tail)); + Format::or(vec![last_present(head), Format::sequence(after_head)]) +} + +/// `parameters` in declared order, and in alphabetical order by `key` if that +/// differs, since Beam often writes them sorted. +fn orders

(parameters: &[P], key: impl Fn(&P) -> &str) -> Vec> { + let declared = parameters.iter().collect::>(); + let mut sorted = declared.clone(); + sorted.sort_by_key(|parameter| key(parameter)); + if sorted + .iter() + .map(|parameter| key(parameter)) + .eq(declared.iter().map(|parameter| key(parameter))) + { + vec![declared] + } else { + vec![declared, sorted] + } +} + +/// `parameter` as listed: optional unless `required`. +fn listed(parameter: Format, required: bool) -> Format { + if required { + parameter + } else { + Format::optional(parameter) + } +} + +fn separator() -> Format { + Format::optional(Format::const_string("\n")) +} + +/// A regex for the tool names the call header can spell (no whitespace or +/// `<`) besides `names`. +fn other_names(names: &[&str]) -> String { + #[derive(Default)] + struct Node { + children: BTreeMap, + end: bool, + } + + fn escape(char: char, meta: &str) -> String { + if meta.contains(char) { + format!("\\{char}") + } else { + char.to_string() + } + } + + /// A name character other than `chars`. + fn other<'a>(chars: impl Iterator) -> String { + let chars = chars.map(|&char| escape(char, r"\]^-[")).collect::(); + format!(r"[^\s<{chars}]") + } + + /// The rest of a name past `node` that ends on none of `names`. + fn rest(node: &Node, first: bool) -> String { + let any = other([].iter()); + let mut branches = vec![format!("{}{any}*", other(node.children.keys()))]; + branches.extend(node.children.iter().map(|(&char, child)| { + format!("{}{}", escape(char, r".^$*+?{}[]\|()"), rest(child, false)) + })); + let optional = if first || node.end { "" } else { "?" }; + format!("(?:{}){optional}", branches.join("|")) + } + + let mut root = Node::default(); + for name in names { + let mut node = &mut root; + for char in name.chars() { + node = node.children.entry(char).or_default(); + } + node.end = true; + } + rest(&root, true) +} + +/// How one parameter is written: `\nVALUE` followed +/// by its end. +/// +/// The parser reads every value as text and converts it by the tool schema, so +/// string options are raw text and every other option is JSON. +enum BeamArguments { + /// The whole parameter, ending with `\n`, after an optional + /// newline. + Closed, + /// The parameter past its opener, where the template looks ahead for the + /// next one: it ends by opening the next parameter, also right after the + /// value's newline, so raw text holds no such opener. + Joined, + /// The last parameter past its opener under such a template, ending with + /// `\n`. + Last, +} + +impl ArgumentSyntax for BeamArguments { + fn separator(&self) -> Option { + None + } + + fn parameter(&self, key: ParameterKey<'_>, options: &[ValueOption<'_>]) -> Option { + let terminators: &[&str] = match self { + Self::Closed => &[""], + Self::Joined | Self::Last => &["", NEXT_PARAMETER], + }; + let values = options + .iter() + .filter_map(|option| match option.ty { + JsonType::String => option.raw_string(terminators), + _ => Some(option.json()), + }) + .collect::>(); + (!values.is_empty()).then(|| { + let value = one_of(values); + match self { + Self::Closed => Format::sequence(vec![ + separator(), + key.tag(PARAMETER_OPEN, PARAMETER_NAME_CLOSE, value, PARAMETER_CLOSE), + ]), + Self::Joined => key.tag( + "", + PARAMETER_NAME_CLOSE, + value, + NEXT_PARAMETER_ENDS.to_vec(), + ), + Self::Last => key.tag("", PARAMETER_NAME_CLOSE, value, PARAMETER_CLOSE), + } + }) + } +} + +#[cfg(test)] +mod tests; diff --git a/rust/src/parser/src/unified/beam/tests.rs b/rust/src/parser/src/unified/beam/tests.rs new file mode 100644 index 000000000000..906a01f3348a --- /dev/null +++ b/rust/src/parser/src/unified/beam/tests.rs @@ -0,0 +1,897 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::sync::Arc; + +use serde_json::{Value, json}; +use vllm_tokenizer::test_utils::TestTokenizer; +use vllm_tokenizer::{DecodedText, Tokenizer as _}; +use xgrammar_structural_tag::ToolChoice; +use xgrammar_structural_tag::format::Format; + +use super::{BeamUnifiedParser, HfPosition, build_output_grammar, other_names}; +use crate::output_grammar::test_utils::outline; +use crate::output_grammar::{GrammarCoverage, OutputGrammarContext, ToolStrictLevel}; +use crate::tool::Tool; +use crate::unified::{ResponseTemplate, UnifiedParser, UnifiedParserEvent, UnifiedParserOutput}; + +/// The `tag_pattern` of the checkpoint of [`template`]. +const TAG_PATTERN: &str = + r#"\n(?P.*?)\n"#; +/// [`TAG_PATTERN`], but a value may also end where the next parameter opens. +const LENIENT_TAG_PATTERN: &str = r#"\n(?P.*?)\n(?:|(?= Arc { + template_with(tag_pattern, json!("<|end|><|start|>assistant")) +} + +/// [`template`] with `thinking_close` as the reasoning closer. +fn template_with(tag_pattern: &str, thinking_close: Value) -> Arc { + let template = json!({ + "defaults": {"role": "assistant"}, + "start_anchor_pattern": r"<\|start\|>assistant", + "fields": { + "thinking": { + "open": "<|channel|>analysis<|message|>", + "close": thinking_close, + "content": "text" + }, + "content": { + "open": "<|channel|>final<|message|>", + "close": ["<|return|>", "<|end|>", "<|call|>"], + "content": "text" + }, + "tool_calls": { + "open_pattern": r" to=(?P[^\s<]+)<\|channel\|>analysis <\|constrain\|>xml<\|message\|>", + "close": ["<|end|><|start|>assistant", "<|call|>"], + "repeats": true, + "content": "xml-inline", + "content_args": { + "tag_pattern": tag_pattern, + "value_parser": {"name": "text", "args": {"strip": false}} + }, + "transform": { + "type": "function", + "function": {"name": "{name}", "arguments": "{content}"} + } + } + } + }); + Arc::new(ResponseTemplate::from_json(&template).unwrap()) +} + +fn tools() -> Vec { + vec![ + Tool { + name: "get_weather".to_string(), + description: None, + parameters: json!({ + "type": "object", + "properties": { + "location": {"type": "string"}, + "days": {"type": "integer"}, + "unit": {"enum": ["c", "f"]} + }, + "required": ["location"], + "additionalProperties": false + }), + strict: Some(true), + defer_loading: None, + }, + Tool { + name: "get_time".to_string(), + description: None, + parameters: json!({ + "type": "object", + "properties": {"zone": {"type": "string"}} + }), + strict: None, + defer_loading: None, + }, + ] +} + +/// A parser initialized from `prompt`, the rendered conversation. +fn parser(prompt: &str) -> BeamUnifiedParser { + parser_with(TAG_PATTERN, prompt) +} + +fn parser_with(tag_pattern: &str, prompt: &str) -> BeamUnifiedParser { + let tokenizer = Arc::new(TestTokenizer::new()); + let prompt_ids = tokenizer.encode(prompt, false).unwrap(); + let mut parser = BeamUnifiedParser::new(template(tag_pattern), &tools(), tokenizer); + parser.initialize(&prompt_ids).unwrap(); + parser +} + +const PROMPT: &str = "<|start|>user<|message|>Weather in Paris?<|end|><|start|>assistant"; + +fn grammar(position: HfPosition, tool_choice: ToolChoice, parallel: bool) -> Option { + grammar_with(position, &tools(), tool_choice, parallel, false, false) +} + +fn grammar_with( + position: HfPosition, + tools: &[Tool], + tool_choice: ToolChoice, + parallel: bool, + always_constrain: bool, + lenient: bool, +) -> Option { + let ctx = OutputGrammarContext { + tools, + tool_choice: &tool_choice, + tool_strict_level: ToolStrictLevel::Auto, + parallel_tool_calls: parallel, + answer: None, + }; + grammar_in(position, &ctx, always_constrain, lenient) +} + +fn grammar_in( + position: HfPosition, + ctx: &OutputGrammarContext<'_>, + always_constrain: bool, + lenient: bool, +) -> Option { + let grammar = build_output_grammar( + position, + ctx, + always_constrain, + lenient, + template(TAG_PATTERN).reasoning_closers(), + ) + .unwrap()?; + assert_eq!(grammar.coverage, GrammarCoverage::FromTokenZero); + Some(outline(&grammar.format)) +} + +fn constrained_grammar( + position: HfPosition, + tool_choice: ToolChoice, + answer: Format, + always_constrain: bool, +) -> Option { + let tools = tools(); + let ctx = OutputGrammarContext { + tools: &tools, + tool_choice: &tool_choice, + tool_strict_level: ToolStrictLevel::Auto, + parallel_tool_calls: true, + answer: Some(&answer), + }; + grammar_in(position, &ctx, always_constrain, false) +} + +#[test] +fn position_follows_the_prompt_and_then_the_output() { + let position = |suffix: &str| parser(&format!("{PROMPT}{suffix}")).inner.position(); + + assert_eq!(position(""), HfPosition::ResponseStart); + assert_eq!( + position("<|channel|>analysis<|message|>Let me"), + HfPosition::Reasoning + ); + // The reasoning closer ends with the anchor, so the next message starts. + assert_eq!( + position("<|channel|>analysis<|message|>Done.<|end|><|start|>assistant"), + HfPosition::ResponseStart + ); + // A partial closer stays buffered. + assert_eq!( + position("<|channel|>analysis<|message|>Done.<|end|>"), + HfPosition::Other + ); + assert_eq!( + position("<|channel|>final<|message|>It is"), + HfPosition::Content + ); + // A partial boundary stays buffered. + assert_eq!(position("<|channel|>ana"), HfPosition::Other); + // Without the anchor, the parser parses the output from the start. + assert_eq!( + parser("no anchor").inner.position(), + HfPosition::ResponseStart + ); + + let mut parser = parser(PROMPT); + parser + .parse_into( + DecodedText::unattributed("<|channel|>analysis<|message|>Hmm"), + &mut UnifiedParserOutput::default(), + ) + .unwrap(); + assert_eq!(parser.inner.position(), HfPosition::Reasoning); +} + +#[test] +fn required_calls_any_tool_after_reasoning() { + expect_test::expect![[r#" + sequence + tag `<|channel|>analysis<|message|>` .. `<|end|><|start|>assistant` + sequence + /[^<]/ + text + sequence + or + sequence + ` to=get_weather<|channel|>analysis <|constrain|>xml<|message|>` + or + sequence + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + sequence + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + sequence + ` to=get_time<|channel|>analysis <|constrain|>xml<|message|>` + sequence + optional + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + star + sequence + `<|end|><|start|>assistant` + or + sequence + ` to=get_weather<|channel|>analysis <|constrain|>xml<|message|>` + or + sequence + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + sequence + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + sequence + ` to=get_time<|channel|>analysis <|constrain|>xml<|message|>` + sequence + optional + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional `<|call|>` + "#]].assert_eq( + &grammar(HfPosition::ResponseStart, ToolChoice::required(), true).unwrap(), + ); +} + +#[test] +fn auto_with_a_strict_tool_also_allows_an_answer() { + expect_test::expect![[r#" + sequence + tag `<|channel|>analysis<|message|>` .. `<|end|><|start|>assistant` + sequence + /[^<]/ + text + or + sequence + or + sequence + ` to=get_weather<|channel|>analysis <|constrain|>xml<|message|>` + or + sequence + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + sequence + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + sequence + ` to=get_time<|channel|>analysis <|constrain|>xml<|message|>` + sequence + optional + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + sequence + sequence + ` to=` + /(?:[^\sanalysis <|constrain|>xml<|message|>` + sequence + star + sequence + optional `\n` + sequence + `\n` + or + text excluding [`<|end|>`, `<|call|>`, ``] + json(number) excluding [`<|end|>`, `<|call|>`] + json(boolean) excluding [`<|end|>`, `<|call|>`] + json(null) excluding [`<|end|>`, `<|call|>`] + json({ ... }) excluding [`<|end|>`, `<|call|>`] + json(any[]) excluding [`<|end|>`, `<|call|>`] + `\n` + star + sequence + `<|end|><|start|>assistant` + or + sequence + ` to=get_weather<|channel|>analysis <|constrain|>xml<|message|>` + or + sequence + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + sequence + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + sequence + ` to=get_time<|channel|>analysis <|constrain|>xml<|message|>` + sequence + optional + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + sequence + sequence + ` to=` + /(?:[^\sanalysis <|constrain|>xml<|message|>` + sequence + star + sequence + optional `\n` + sequence + `\n` + or + text excluding [`<|end|>`, `<|call|>`, ``] + json(number) excluding [`<|end|>`, `<|call|>`] + json(boolean) excluding [`<|end|>`, `<|call|>`] + json(null) excluding [`<|end|>`, `<|call|>`] + json({ ... }) excluding [`<|end|>`, `<|call|>`] + json(any[]) excluding [`<|end|>`, `<|call|>`] + `\n` + optional `<|call|>` + sequence + `<|channel|>final<|message|>` + text excluding [`<|end|>`, `<|return|>`] + optional `<|return|>` + "#]] + .assert_eq(&grammar(HfPosition::ResponseStart, ToolChoice::auto(), true).unwrap()); +} + +#[test] +fn named_choice_calls_that_tool_once_after_prompted_reasoning() { + expect_test::expect![[r#" + sequence + tag `` text `<|end|><|start|>assistant` + sequence + sequence + ` to=get_weather<|channel|>analysis <|constrain|>xml<|message|>` + or + sequence + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + sequence + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + optional `<|call|>` + "#]].assert_eq( + &grammar(HfPosition::Reasoning, ToolChoice::function("get_weather"), true).unwrap(), + ); +} + +#[test] +fn no_grammar_without_a_constraining_request_or_known_position() { + let non_strict = vec![tools().pop().unwrap()]; + for tools in [&[][..], &non_strict] { + assert!( + grammar_with( + HfPosition::ResponseStart, + tools, + ToolChoice::auto(), + true, + false, + false + ) + .is_none() + ); + } + assert!( + grammar_with( + HfPosition::Other, + &tools(), + ToolChoice::required(), + true, + true, + false + ) + .is_none() + ); +} + +#[test] +fn none_with_tools_only_answers() { + expect_test::expect![[r#" + sequence + tag `<|channel|>analysis<|message|>` .. `<|end|><|start|>assistant` + sequence + /[^<]/ + text + sequence + `<|channel|>final<|message|>` + text excluding [`<|end|>`, `<|return|>`] + optional `<|return|>` + "#]] + .assert_eq(&grammar(HfPosition::ResponseStart, ToolChoice::none(), true).unwrap()); +} + +#[test] +fn always_constrain_answers_without_tools() { + expect_test::expect![[r#" + sequence + tag `` text `<|end|><|start|>assistant` + sequence + `<|channel|>final<|message|>` + text excluding [`<|end|>`, `<|return|>`] + optional `<|return|>` + "#]] + .assert_eq( + &grammar_with( + HfPosition::Reasoning, + &[], + ToolChoice::auto(), + true, + true, + false, + ) + .unwrap(), + ); +} + +#[test] +fn always_constrain_allows_non_strict_calls_or_an_answer() { + let non_strict = vec![tools().pop().unwrap()]; + expect_test::expect![[r#" + sequence + tag `<|channel|>analysis<|message|>` .. `<|end|><|start|>assistant` + sequence + /[^<]/ + text + or + sequence + or + sequence + ` to=get_time<|channel|>analysis <|constrain|>xml<|message|>` + sequence + optional + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + sequence + sequence + ` to=` + /(?:[^\sanalysis <|constrain|>xml<|message|>` + sequence + star + sequence + optional `\n` + sequence + `\n` + or + text excluding [`<|end|>`, `<|call|>`, ``] + json(number) excluding [`<|end|>`, `<|call|>`] + json(boolean) excluding [`<|end|>`, `<|call|>`] + json(null) excluding [`<|end|>`, `<|call|>`] + json({ ... }) excluding [`<|end|>`, `<|call|>`] + json(any[]) excluding [`<|end|>`, `<|call|>`] + `\n` + optional `<|call|>` + sequence + `<|channel|>final<|message|>` + text excluding [`<|end|>`, `<|return|>`] + optional `<|return|>` + "#]] + .assert_eq( + &grammar_with( + HfPosition::ResponseStart, + &non_strict, + ToolChoice::auto(), + false, + true, + false, + ) + .unwrap(), + ); +} + +#[test] +fn always_constrain_finishes_an_open_answer() { + expect_test::expect![[r#" + sequence + text excluding [`<|end|>`, `<|return|>`] + optional `<|return|>` + "#]] + .assert_eq( + &grammar_with( + HfPosition::Content, + &tools(), + ToolChoice::auto(), + true, + true, + false, + ) + .unwrap(), + ); + assert!( + grammar_with( + HfPosition::Content, + &tools(), + ToolChoice::auto(), + true, + false, + false + ) + .is_none() + ); + assert!( + grammar_with( + HfPosition::Content, + &tools(), + ToolChoice::required(), + true, + true, + false + ) + .is_none() + ); +} + +#[test] +fn answer_holds_the_requested_format() { + let json = || Format::json_schema(json!({"type": "object"})); + let answer = constrained_grammar(HfPosition::ResponseStart, ToolChoice::none(), json(), false); + expect_test::expect![[r#" + sequence + tag `<|channel|>analysis<|message|>` .. `<|end|><|start|>assistant` + sequence + /[^<]/ + text + sequence + `<|channel|>final<|message|>` + json(object) + optional `<|return|>` + "#]] + .assert_eq(&answer.unwrap()); + + // Under `auto`, the answer that may replace calls holds it too. + let auto = + constrained_grammar(HfPosition::ResponseStart, ToolChoice::auto(), json(), false).unwrap(); + assert!(auto.ends_with( + "`<|channel|>final<|message|>`\n json(object)\n optional `<|return|>`\n" + )); + + // A forced call leaves no answer to constrain. + assert_eq!( + constrained_grammar( + HfPosition::ResponseStart, + ToolChoice::required(), + json(), + false + ), + grammar(HfPosition::ResponseStart, ToolChoice::required(), true) + ); +} + +#[test] +fn request_constraint_applies_alone_in_an_open_answer() { + let json = Format::json_schema(json!({"type": "object"})); + assert!(constrained_grammar(HfPosition::Content, ToolChoice::auto(), json, true).is_none()); +} + +/// Parse `output` after the prompt under `tag_pattern`, returning the calls as +/// `(name, arguments)`. +fn calls(tag_pattern: &str, output: &str) -> Vec<(String, Value)> { + let mut parser = parser_with(tag_pattern, PROMPT); + let mut events = UnifiedParserOutput::default(); + parser.parse_into(DecodedText::unattributed(output), &mut events).unwrap(); + events.append(parser.finish().unwrap()); + let mut calls: Vec<(String, String)> = Vec::new(); + for event in events.events { + if let UnifiedParserEvent::ToolCall(delta) = event { + if let Some(name) = delta.name { + calls.push((name, String::new())); + } + calls.last_mut().unwrap().1.push_str(&delta.arguments); + } + } + calls + .into_iter() + .map(|(name, arguments)| (name, serde_json::from_str(&arguments).unwrap())) + .collect() +} + +/// Output in the grammar's canonical spelling parses back to the values the +/// grammar admitted. +#[test] +fn parser_reads_back_the_grammar_spelling() { + let parameter = |key: &str, value: &str| { + format!("\n{value}\n") + }; + let call = |name: &str, parameters: &[String]| { + format!( + " to={name}<|channel|>analysis <|constrain|>xml<|message|>{}", + parameters.join("\n") + ) + }; + let output = format!( + "<|channel|>analysis<|message|>Two tools.<|end|><|start|>assistant{}<|end|><|start|>assistant{}", + call( + "get_weather", + &[ + parameter("location", "Paris"), + parameter("days", "2"), + parameter("unit", "c") + ] + ), + call("get_time", &[parameter("zone", "multi\nline")]), + ); + + for tag_pattern in [TAG_PATTERN, LENIENT_TAG_PATTERN] { + assert_eq!( + calls(tag_pattern, &output), + vec![ + ( + "get_weather".to_string(), + json!({"location": "Paris", "days": 2, "unit": "c"}) + ), + ("get_time".to_string(), json!({"zone": "multi\nline"})), + ] + ); + } +} + +#[test] +fn values_end_before_the_next_parameter_where_the_template_looks_ahead() { + let output = " to=get_weather<|channel|>analysis <|constrain|>xml<|message|>\ + \n2\n\ + \nParis\n"; + + assert!(parser_with(LENIENT_TAG_PATTERN, PROMPT).lenient); + assert_eq!( + calls(LENIENT_TAG_PATTERN, output), + vec![( + "get_weather".to_string(), + json!({"days": 2, "location": "Paris"}) + )] + ); + assert!(!parser(PROMPT).lenient); +} + +#[test] +fn reasoning_ends_at_every_closer_of_the_template() { + let closers = [ + "<|end|><|start|>assistant", + "<|constrain|><|start|>assistant", + "<|start|>assistant", + ]; + let tools = tools(); + let tokenizer = Arc::new(TestTokenizer::new()); + let mut parser = BeamUnifiedParser::new( + template_with(TAG_PATTERN, json!(closers)), + &tools, + tokenizer.clone(), + ); + parser.initialize(&tokenizer.encode(PROMPT, false).unwrap()).unwrap(); + let grammar = parser + .build_output_grammar(&OutputGrammarContext { + tools: &tools, + tool_choice: &ToolChoice::required(), + tool_strict_level: ToolStrictLevel::Auto, + parallel_tool_calls: false, + answer: None, + }) + .unwrap() + .unwrap(); + + assert_eq!( + outline(&grammar.format).lines().nth(1).unwrap(), + " tag `<|channel|>analysis<|message|>` .. `<|constrain|><|start|>assistant` | \ + `<|end|><|start|>assistant` | `<|start|>assistant`" + ); +} + +#[test] +fn lenient_arguments_come_in_declared_or_sorted_order() { + expect_test::expect![[r#" + sequence + tag `` text `<|end|><|start|>assistant` + sequence + sequence + ` to=get_weather<|channel|>analysis <|constrain|>xml<|message|>` + sequence + optional `\n` + `\n` text excluding [`<|end|>`, `<|call|>`, ``, `\n\n` text excluding [`<|end|>`, `<|call|>`, ``, `\n\n\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + sequence + optional tag `days">\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n\n` .. `\n` + or + `c` + `f` + sequence + optional tag `days">\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n\n` text excluding [`<|end|>`, `<|call|>`, ``, `\n\n` text excluding [`<|end|>`, `<|call|>`, ``, `\n\n\n` .. `\n` + or + `c` + `f` + optional `<|call|>` + "#]].assert_eq( + &grammar_with( + HfPosition::Reasoning, + &tools(), + ToolChoice::function("get_weather"), + true, + false, + true, + ) + .unwrap(), + ); +} + +#[test] +fn other_names_match_every_name_but_the_given_ones() { + let regex = regex_automata::meta::Regex::new(&format!( + "^(?:{})$", + other_names(&["get", "get_time", "a.b"]) + )) + .unwrap(); + + for name in [ + "g", + "ge", + "gets", + "get_", + "get_t", + "get_times", + "x", + "a", + "a.", + "axb", + "a.bc", + ] { + assert!(regex.is_match(name), "{name}"); + } + for name in ["get", "get_time", "a.b", "", "get time", "get<"] { + assert!(!regex.is_match(name), "{name}"); + } +} diff --git a/rust/src/parser/src/unified/hf/content.rs b/rust/src/parser/src/unified/hf/content.rs index d62ba9305e07..598c6597e1d4 100644 --- a/rust/src/parser/src/unified/hf/content.rs +++ b/rust/src/parser/src/unified/hf/content.rs @@ -6,8 +6,12 @@ //! Port of `content_parsers.py`. Each parser takes a chunk of captured text and //! parses it into a single key in the output message dictionary. +use std::borrow::Cow; + use regex_automata::meta::Regex; +use regex_automata::util::iter::Searcher; use regex_automata::util::syntax; +use regex_automata::{Input, Match}; use serde::Deserialize; use serde_json::{Map, Value}; use thiserror_ext::AsReport as _; @@ -22,6 +26,10 @@ use crate::unified::{self, parsing_failed}; const LAX_OPEN: char = '\u{1}'; const LAX_CLOSE: char = '\u{2}'; +/// Name of the group that stands for the lookahead ending an `xml-inline` tag +/// pattern. +const AHEAD_GROUP: &str = "lookahead"; + /// A configured content parser (`content` plus `content_args` of one field). #[derive(Debug, Clone)] pub(super) enum ContentParser { @@ -66,6 +74,9 @@ pub(super) struct XmlInlineArgs { key_group: usize, /// Index of the `value` group, if the pattern declares one. value_group: Option, + /// Index of the group standing for the lookahead that ends the pattern, + /// and the lookahead's own pattern, if any. + lookahead: Option<(usize, String)>, value_parser: Option>, /// Collect duplicate keys into a list. merge_duplicates: bool, @@ -143,7 +154,8 @@ impl ContentParser { } ContentKind::XmlInline => { let spec: XmlInlineSpec = deserialize(&args, &context)?; - let tag_pattern = compile_regex(scope, &spec.tag_pattern)?; + let (pattern, lookahead) = final_lookahead(&spec.tag_pattern); + let tag_pattern = compile_regex(scope, &pattern)?; let group_info = tag_pattern.group_info(); let Some(key_group) = group_info.to_index(Default::default(), "key") else { return Err(invalid!( @@ -152,10 +164,14 @@ impl ContentParser { )); }; let value_group = group_info.to_index(Default::default(), "value"); + let lookahead = group_info + .to_index(Default::default(), AHEAD_GROUP) + .zip(lookahead.map(str::to_string)); Self::XmlInline(XmlInlineArgs { tag_pattern, key_group, value_group, + lookahead, value_parser: ValueParserSpec::compile(scope, spec.value_parser)?, merge_duplicates: spec.merge_duplicates, }) @@ -193,6 +209,32 @@ impl ContentParser { Self::KvLines(args) => args.parse(text), } } + + /// The pattern of the lookahead that ends an `xml-inline` tag pattern, if + /// any. + pub fn lookahead(&self) -> Option<&str> { + match self { + Self::XmlInline(args) => args.lookahead.as_ref().map(|(_, ahead)| ahead.as_str()), + _ => None, + } + } +} + +/// `pattern` with a lookahead that ends it matched as the group +/// [`AHEAD_GROUP`] instead, as the Rust regex engine has no lookahead, and +/// the lookahead's own pattern. The engine rejects any other lookahead. +fn final_lookahead(pattern: &str) -> (Cow<'_, str>, Option<&str>) { + if let Some((head, rest)) = pattern.split_once("(?=") { + // Only closing parentheses may follow a lookahead that ends the pattern. + let ahead = rest.trim_end_matches(')'); + if ahead.len() < rest.len() && !ahead.contains(['(', ')']) { + return ( + format!("{head}(?P<{AHEAD_GROUP}>{rest}").into(), + Some(ahead), + ); + } + } + (pattern.into(), None) } impl ValueParserSpec { @@ -258,7 +300,21 @@ impl XmlInlineArgs { /// Parse shallow XML-ish tags into a dict. fn parse(&self, text: &str) -> unified::Result { let mut out = Map::new(); - for captures in self.tag_pattern.captures_iter(text) { + let mut captures = self.tag_pattern.create_captures(); + let mut tags = Searcher::new(Input::new(text)); + while tags + .advance(|input| { + self.tag_pattern.search_captures(input, &mut captures); + // The lookahead's text is not part of the tag. + let ahead = + self.lookahead.as_ref().and_then(|(group, _)| captures.get_group(*group)); + Ok(captures.get_match().map(|tag| match ahead { + Some(ahead) => Match::new(tag.pattern(), tag.start()..ahead.start), + None => tag, + })) + }) + .is_some() + { let key = captures.get_group(self.key_group).map(|span| &text[span.range()]).ok_or_else( || parsing_failed!("xml-inline: named group 'key' did not participate"), @@ -490,6 +546,25 @@ mod tests { ); } + #[test] + fn xml_inline_tags_may_end_where_the_pattern_looks_ahead() { + let xml = parser( + "xml-inline", + json!({ + "tag_pattern": r#"(?P.*?)\n(?:|(?=1\nx y\n") + .unwrap(); + assert_eq!(value, json!({"a": "1", "b": "x y"})); + assert_eq!(xml.lookahead(), Some("\w+)(?=:)(?P.*)"}); + assert!(ContentParser::new("test", kind("xml-inline"), args.as_object().unwrap()).is_err()); + } + #[test] fn kv_lines_split_and_strip() { let kv = parser("kv-lines", json!({"value_parser": {"name": "int"}})); diff --git a/rust/src/parser/src/unified/hf/mod.rs b/rust/src/parser/src/unified/hf/mod.rs index 568a8371d0ac..9d9bd3bc0664 100644 --- a/rust/src/parser/src/unified/hf/mod.rs +++ b/rust/src/parser/src/unified/hf/mod.rs @@ -25,7 +25,7 @@ mod transform; #[cfg(test)] mod tests; -pub use parser::HfUnifiedParser; +pub use parser::{HfPosition, HfUnifiedParser}; pub use template::ResponseTemplate; use thiserror::Error; use thiserror_ext::Macro; diff --git a/rust/src/parser/src/unified/hf/parser.rs b/rust/src/parser/src/unified/hf/parser.rs index 5d5065f109b6..febdd4717565 100644 --- a/rust/src/parser/src/unified/hf/parser.rs +++ b/rust/src/parser/src/unified/hf/parser.rs @@ -59,6 +59,20 @@ pub struct HfUnifiedParser { emitted_tool_count: usize, } +/// Where the parser is in the response; see [`HfUnifiedParser::position`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum HfPosition { + /// Nothing parsed since the start anchor, or since the start of the stream + /// when there is none. + ResponseStart, + /// Inside an open reasoning field. + Reasoning, + /// Inside an open content field. + Content, + /// Anywhere else. + Other, +} + /// Parser mode: the region currently receiving text. enum HfMode { /// No implicit field is declared: text outside regions is discarded. @@ -156,6 +170,36 @@ impl HfUnifiedParser { parser } + /// Where the parser is in the response. Between + /// [`UnifiedParser::initialize`] and the first generated text, this is where + /// the prompt leaves the response. + /// + /// Text outside every region, which a template without an implicit field + /// discards, leaves no state behind and so does not move the parser from + /// [`HfPosition::ResponseStart`]. + pub fn position(&self) -> HfPosition { + if !self.buffer.text.is_empty() { + return HfPosition::Other; + } + match &self.mode { + HfMode::Implicit(occurrence) | HfMode::Explicit(occurrence) if occurrence.opened => { + match &self.template.regions[occurrence.region].kind { + RegionKind::Text(TextRegion { + role: TextRole::Reasoning, + .. + }) => HfPosition::Reasoning, + RegionKind::Text(TextRegion { + role: TextRole::Content, + .. + }) => HfPosition::Content, + _ => HfPosition::Other, + } + } + _ if self.closed.contains(&true) || self.emitted_tool_count > 0 => HfPosition::Other, + _ => HfPosition::ResponseStart, + } + } + /// Clear all per-stream state. fn reset_state(&mut self) { self.buffer.clear(); diff --git a/rust/src/parser/src/unified/hf/template.rs b/rust/src/parser/src/unified/hf/template.rs index 4889f6f9b571..29c096ea3245 100644 --- a/rust/src/parser/src/unified/hf/template.rs +++ b/rust/src/parser/src/unified/hf/template.rs @@ -191,6 +191,31 @@ impl ResponseTemplate { }) } + /// The pattern of the lookahead that ends the `xml-inline` tag pattern of + /// tool-call arguments, such as for the opener of the next argument. + pub fn argument_lookahead(&self) -> Option<&str> { + self.regions.iter().find_map(|region| match ®ion.kind { + RegionKind::ToolCalls(calls) => calls.content.lookahead(), + _ => None, + }) + } + + /// The literal closers of the reasoning field, longest first, or none if it + /// has no literal close. + pub fn reasoning_closers(&self) -> &[String] { + self.regions + .iter() + .find_map(|region| match (®ion.kind, ®ion.close) { + (RegionKind::Text(text), Some(Boundary::Literals(closers))) + if text.role == TextRole::Reasoning => + { + Some(closers.as_slice()) + } + _ => None, + }) + .unwrap_or_default() + } + /// Return the start and end offsets of the last start-anchor match in /// `prompt`. pub(super) fn last_start_anchor(&self, prompt: &str) -> Option<(usize, usize)> { diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs index 74bef8e1f45f..585f28cfdebd 100644 --- a/rust/src/parser/src/unified/mod.rs +++ b/rust/src/parser/src/unified/mod.rs @@ -3,6 +3,7 @@ //! Unified parser interface for reasoning and tool-call deltas. +mod beam; mod combined; mod gemma4; mod hf; @@ -10,9 +11,10 @@ mod hy; mod inkling; mod kimi_k3; +pub use beam::BeamUnifiedParser; pub use combined::CombinedParser; pub use gemma4::Gemma4UnifiedParser; -pub use hf::{HfTemplateError, HfUnifiedParser, ResponseTemplate}; +pub use hf::{HfPosition, HfTemplateError, HfUnifiedParser, ResponseTemplate}; pub use hy::{HyV3UnifiedParser, HyV4UnifiedParser}; pub use inkling::InklingUnifiedParser; pub use kimi_k3::KimiK3UnifiedParser; diff --git a/setup.py b/setup.py index 83b59e303a6e..01d90710097c 100644 --- a/setup.py +++ b/setup.py @@ -1600,6 +1600,7 @@ def _read_requirements(filename: str) -> list[str]: "vllm": [ "py.typed", "libs/*.so*", + "models/beam/nvidia/*.cu", "model_executor/layers/fused_moe/configs/*.json", "model_executor/layers/quantization/utils/configs/*.json", "entrypoints/serve/instrumentator/static/*.js", diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index 399c898d0b01..92a44432bab4 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -1062,6 +1062,37 @@ def _tc(*, args, id=None, name=None): ) +@pytest.mark.asyncio +@pytest.mark.parametrize("after_content", [False, True]) +@pytest.mark.parametrize( + "error_type,code", [("BadRequestError", 400), ("InternalServerError", 500)] +) +async def test_stream_preserves_backend_error_and_terminates( + after_content, error_type, code +): + message = "thinking_token_budget requires reasoning_config" + + async def sse_input(): + if after_content: + yield _make_stream_chunk(delta=DeltaMessage(role="assistant")) + yield _make_stream_chunk(delta=DeltaMessage(content="Partial")) + yield "data: " + json.dumps( + {"error": {"type": error_type, "message": message, "code": code}} + ) + yield "data: [DONE]" + + converter = _make_stream_converter() + events = _parse_sse_events( + [event async for event in converter.message_stream_converter(sse_input())] + ) + assert events[-1] == ( + "error", + {"type": "error", "error": {"type": error_type, "message": message}}, + ) + assert sum(event == "error" for event, _ in events) == 1 + assert not any(event in {"message_delta", "message_stop"} for event, _ in events) + + class TestMessageStreamConverterToolUseContentBuffering: """Regression test for tool_use arguments being silently dropped. diff --git a/tests/kernels/moe/test_flashinfer_moe_ep.py b/tests/kernels/moe/test_flashinfer_moe_ep.py index 7ce3124ae144..de81a4f4b4dd 100644 --- a/tests/kernels/moe/test_flashinfer_moe_ep.py +++ b/tests/kernels/moe/test_flashinfer_moe_ep.py @@ -24,7 +24,12 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk from tests.distributed.eplb_utils import distributed_run, set_env_vars_and_device -from vllm.config import VllmConfig, set_current_vllm_config +from vllm.config import ( + KernelConfig, + ParallelConfig, + 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 @@ -32,6 +37,7 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, FusedMoEParallelConfig, + FusedMoEQuantConfig, RoutingMethodType, ) from vllm.model_executor.layers.fused_moe.experts.flashinfer_moe_ep import ( @@ -41,6 +47,7 @@ FlashInferMoeEp, FlashInferMoeEpWeights, apply_topk_in_fc1, + supports_flashinfer_moe_ep_deferred_finalize, ) from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Static from vllm.platforms import current_platform @@ -241,6 +248,7 @@ def _check_expert_views_alias_kernel_memory( assert torch.equal(output, _expected_output(geometry, topk_ids, topk_weights)) +@torch.inference_mode() def _init_warmup_forward( env: dict[str, str], world_size: int, geometry: Geometry, backend: str ) -> None: @@ -295,6 +303,216 @@ def _init_warmup_forward( adapter.destroy() +@torch.inference_mode() +def _data_parallel_boundary_forward( + env: dict[str, str], + world_size: int, + cache_path: str | None = None, +) -> None: + from vllm.distributed import graph_capture + from vllm.forward_context import set_forward_context + from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, + ) + from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput + from vllm.models.beam.nvidia.layernorm import beam_norm_residual_norm + from vllm.models.beam.nvidia.moe_finalize import ( + beam_moe_finalize_norm_residual_norm, + ) + from vllm.v1.worker.workspace import init_workspace_manager + + set_env_vars_and_device(env) + rank = torch.distributed.get_rank() + vllm_config = VllmConfig( + parallel_config=ParallelConfig( + data_parallel_size=world_size, + data_parallel_rank=rank, + enable_expert_parallel=True, + ), + kernel_config=KernelConfig(moe_backend=FLASHINFER_MOE_EP_CUTEDSL), + ) + vllm_config.scheduler_config.max_num_seqs = 32 + geometry = Geometry("dp-top8", 1024, 512, 32, 8) + with set_current_vllm_config(vllm_config): + ensure_model_parallel_initialized( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + device = torch.device(f"cuda:{rank}") + init_workspace_manager(device) + moe = _moe_config( + vllm_config, geometry, world_size, device, FLASHINFER_MOE_EP_CUTEDSL + ) + moe.defer_moe_finalize(allow_ep=True) + assert moe.use_deferred_moe_finalize + weights = _local_expert_weights(geometry, rank, moe.num_local_experts, device) + if cache_path is not None: + import os + + from flashinfer.moe_ep.kernel_src.sm100.cutedsl_megamoe.shim.knob_cache import ( # noqa: E501 + record_knobs, + ) + + os.environ["FLASHINFER_MOE_EP_KNOB_CACHE"] = cache_path + if rank == 0: + for capacity, live, cluster in ( + (32, 4, 1), + (32, 32, 4), + (MAX_TOKENS_PER_RANK, 33, 2), + ): + record_knobs( + dict( + mma_tiler_mnk=(128 if cluster == 1 else 256, 128, 256), + cluster_shape_mnk=(cluster, 1, 1), + token_back_mode="epi_warps", + load_balance_mode="atomic_counter", + group_hint=512, + flag_batch=4, + epi_flag_batch=(2, 4), + ), + dtype="nvfp4", + world_size=world_size, + hidden=geometry.hidden_size, + intermediate=2 * geometry.intermediate_size, + num_experts=geometry.num_experts, + topk=geometry.top_k, + max_tokens=capacity, + live_tokens=live, + apply_topk_in_fc1=False, + defer_topk_reduce=True, + ) + torch.distributed.barrier() + adapter = FlashInferMoeEp(moe, weights, apply_topk_in_fc1=False) + try: + adapter.warmup() + assert adapter._decode_max_num_tokens == 32 + if cache_path is not None: + assert list(adapter._live_workspaces[32]) == [4, 32] + for live, cluster in ((4, 1), (32, 4)): + workspace = adapter._live_workspaces[32][live] + assert workspace.max_tokens_per_rank == 32 + assert ( + workspace._backend_workspace._frontend.config.cluster_shape_mnk + == (cluster, 1, 1) + ) + quant_config = FusedMoEQuantConfig.make() + experts = FlashInferMoeEpExperts(moe, quant_config) + experts._adapter = adapter + prepare_finalize = maybe_make_prepare_finalize( + moe, quant_config, allow_new_interface=True, use_monolithic=False + ) + assert prepare_finalize.output_is_reduced() + kernel = mk.FusedMoEKernel(prepare_finalize, experts) + post_weight = torch.linspace( + 0.5, 1.5, geometry.hidden_size, device=device, dtype=torch.bfloat16 + ) + scales = (0.125, 0.9682458365518543) + ids, scores = _routing(geometry, rank, device) + base_input = _hidden_states(geometry, device) + reduced = (_expected_output(geometry, ids, scores) * 2.5).bfloat16() + expected = beam_norm_residual_norm( + reduced, base_input, post_weight, 1e-6, post_weight, 1e-6, scales + ) + + def run(x, batch_ids, batch_scores, sizes): + with set_forward_context( + None, + vllm_config, + num_tokens=x.shape[0], + num_tokens_across_dp=torch.tensor(sizes, dtype=torch.int32), + ): + if cache_path is not None: + selected = adapter._select_workspace(x.shape[0]) + expected_bucket = ( + 4 if max(sizes) <= 4 else (32 if max(sizes) <= 32 else 33) + ) + assert selected.live_tokens == expected_bucket + output = kernel.apply( + hidden_states=x, + w1=weights.w13, + w2=weights.w2, + topk_weights=batch_scores, + topk_ids=batch_ids, + activation=MoEActivation.SILU, + global_num_experts=geometry.num_experts, + expert_map=None, + apply_router_weight_on_input=False, + ) + assert isinstance(output, UnfinalizedMoEOutput) + return beam_moe_finalize_norm_residual_norm( + output, + x, + post_weight, + 1e-6, + post_weight, + 1e-6, + scales, + routed_scale=2.5, + ) + + graphs = {} + for sizes in ((0, 4), (3, 0), (0, 0), (31, 32), (32, 33)): + rows = torch.arange(sizes[rank], device=device) % NUM_TOKENS_PER_RANK + x, batch_ids, batch_scores = base_input[rows], ids[rows], scores[rows] + actual = run(x, batch_ids, batch_scores, sizes) + for observed, reference in zip(actual, expected): + torch.testing.assert_close( + observed, reference[rows], rtol=0, atol=0 + ) + if sizes in ((0, 4), (31, 32), (32, 33)): + graph = torch.cuda.CUDAGraph() + with ( + graph_capture(device) as capture, + torch.cuda.graph(graph, stream=capture.stream), + ): + captured = run(x, batch_ids, batch_scores, sizes) + graphs[sizes] = graph, (x, batch_ids, batch_scores), captured, rows + + for sizes in ((0, 4), (31, 32), (32, 33), (0, 4)): + graph, inputs, captured, rows = graphs[sizes] + x = inputs[0] + x.zero_() + graph.replay() + for observed in captured: + torch.testing.assert_close( + observed, torch.zeros_like(x), rtol=0, atol=0 + ) + x.copy_(base_input[rows]) + graph.replay() + for observed, reference in zip(captured, expected): + torch.testing.assert_close( + observed, reference[rows], rtol=0, atol=0 + ) + for graph, *_ in graphs.values(): + graph.reset() + finally: + adapter.destroy() + + +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="The FlashInfer MoE-EP megakernels require SM100", +) +@pytest.mark.parametrize("live_profiles", [False, True]) +def test_data_parallel_boundary_preserves_local_tokens_and_graph_profiles( + tmp_path, live_profiles +): + """DP2 keeps unequal/empty batches local across producer workspace changes.""" + moe_ep = pytest.importorskip("flashinfer.moe_ep") + if not supports_flashinfer_moe_ep_deferred_finalize(): + pytest.skip("FlashInfer does not support unfinalized MegaMoE output") + if not hasattr(moe_ep.MoEEpMegaLayer, "create_workspace"): + pytest.skip("FlashInfer does not support multiple MegaMoE workspaces") + if live_profiles and not hasattr(moe_ep.MoEEpMegaLayer, "live_token_buckets"): + pytest.skip("FlashInfer does not support live-token knob profiles") + if torch.accelerator.device_count() < 2: + pytest.skip("Need at least 2 GPUs") + distributed_run( + _data_parallel_boundary_forward, + 2, + str(tmp_path / "knobs.json") if live_profiles else None, + ) + + @pytest.mark.skipif( not current_platform.is_device_capability_family(100), reason="The FlashInfer MoE-EP megakernels require SM100", diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py index b0c3e26db99a..77eac550d904 100644 --- a/tests/kernels/moe/test_routing.py +++ b/tests/kernels/moe/test_routing.py @@ -582,6 +582,20 @@ def test_fused_topk_bias( assert_routing_results_close(topk_weights, topk_ids, baseline_weights, baseline_ids) +@pytest.mark.parametrize("scoring_func", [0, 1]) +def test_fused_grouped_topk_accepts_idle_rank(scoring_func): + """Zero-token routing must leave CUDA ready for the next batch.""" + if not current_platform.is_cuda(): + pytest.skip("CUDA required") + from vllm import _custom_ops as ops + + scores = torch.empty((0, 8), device="cuda", dtype=torch.float32) + bias = torch.zeros(8, device="cuda") + weights, ids = ops.grouped_topk(scores, 2, 1, 2, True, 1.0, bias, scoring_func) + torch.testing.assert_close(weights, scores.new_empty((0, 2))) + torch.testing.assert_close(ids, scores.new_empty((0, 2), dtype=torch.int32)) + + @pytest.mark.parametrize("m,k", MK_S) @pytest.mark.parametrize("top_k", TOP_KS) @pytest.mark.parametrize( diff --git a/tests/kernels/moe/test_trtllm_nvfp4_moe.py b/tests/kernels/moe/test_trtllm_nvfp4_moe.py index 4ee095505196..d41e433c77e7 100644 --- a/tests/kernels/moe/test_trtllm_nvfp4_moe.py +++ b/tests/kernels/moe/test_trtllm_nvfp4_moe.py @@ -38,6 +38,7 @@ from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( TrtLlmNvFp4ExpertsModular, ) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( prepare_static_weights_for_trtllm_fp4_moe, reorder_w1w3_to_w3w1, @@ -434,16 +435,22 @@ def test_trtllm_fp4_moe_reprocess_does_not_refold_swiglu_params(): torch.testing.assert_close(experts.gemm1_beta, beta * expected) -@pytest.mark.parametrize("m,chunk_size", [(1, None), (16, None), (16, 5)]) +@pytest.mark.parametrize( + "m,chunk_size,ep_size", [(1, None, 1), (16, None, 1), (16, 5, 1), (16, None, 4)] +) @torch.inference_mode() def test_trtllm_fp4_moe_deferred_finalize( m: int, chunk_size: int | None, + ep_size: int, monkeypatch: pytest.MonkeyPatch, workspace_init, ): """TRTLLM-Gen NvFP4 modular experts can leave the top-k finalize to the caller.""" e, topk, n, k = 128, 8, 1024, 1024 + local_experts = e // ep_size + ep_rank = int(ep_size > 1) + offset = ep_rank * local_experts dtype = torch.bfloat16 if chunk_size is not None: monkeypatch.setattr( @@ -456,23 +463,40 @@ def test_trtllm_fp4_moe_deferred_finalize( ): a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 w1_q, w2_q, quant_config = make_test_quant_config( - e, n, k, in_dtype=dtype, quant_dtype="nvfp4", is_scale_swizzled=False + local_experts, + n, + k, + in_dtype=dtype, + quant_dtype="nvfp4", + is_scale_swizzled=False, ) score = torch.randn((m, e), device="cuda", dtype=dtype) topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) + expert_map = None + if ep_size > 1: + topk_ids[0] = torch.arange(topk, device="cuda") + topk_ids[1] = torch.arange(topk, device="cuda") + offset + expert_map = torch.full((e,), -1, device="cuda", dtype=torch.int32) + expert_map[offset : offset + local_experts] = torch.arange( + local_experts, device="cuda", dtype=torch.int32 + ) - # One rank of a TP group, which deferral needs. + # One rank of a replicated-token TP or EP group. moe_config = FusedMoEConfig( num_experts=e, experts_per_token=topk, hidden_dim=k, - intermediate_size=2 * n, - num_local_experts=e, + intermediate_size=2 * n if ep_size == 1 else n, + num_local_experts=local_experts, num_logical_experts=e, activation=MoEActivation.SILU, device="cuda", moe_parallel_config=replace( - FusedMoEParallelConfig.make_no_parallel(), tp_size=2 + FusedMoEParallelConfig.make_no_parallel(), + tp_size=2 if ep_size == 1 else 1, + ep_size=ep_size, + ep_rank=ep_rank, + use_ep=ep_size > 1, ), in_dtype=dtype, routing_method=RoutingMethodType.TopK, @@ -480,9 +504,8 @@ def test_trtllm_fp4_moe_deferred_finalize( ) kernel = _make_trtllm_fp4_moe_kernel(moe_config, quant_config) - check_deferred_moe_finalize( - moe_config, - lambda: kernel.apply( + def run(): + output = kernel.apply( hidden_states=a, w1=w1_q, w2=w2_q, @@ -490,11 +513,24 @@ def test_trtllm_fp4_moe_deferred_finalize( topk_ids=topk_ids, activation=MoEActivation.SILU, global_num_experts=e, - expert_map=None, + expert_map=expert_map, apply_router_weight_on_input=False, - ), + ) + if ep_size > 1 and isinstance(output, UnfinalizedMoEOutput): + nonlocal_routes = (topk_ids < offset) | ( + topk_ids >= offset + local_experts + ) + torch.testing.assert_close( + output.expanded_idx_to_permuted_idx < 0, nonlocal_routes + ) + return output + + check_deferred_moe_finalize( + moe_config, + run, router_weights=topk_weights, chunked=chunk_size is not None, + allow_ep=ep_size > 1, ) diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 1536c4cec9ae..8138ec1b60c9 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -594,6 +594,7 @@ def check_deferred_moe_finalize( run: Callable[[], torch.Tensor | UnfinalizedMoEOutput], router_weights: torch.Tensor | None = None, chunked: bool = False, + allow_ep: bool = False, ) -> None: """Check a kernel that defers its finalize against the finalize it skips. @@ -608,7 +609,7 @@ def check_deferred_moe_finalize( expected_weights = None if router_weights is None else router_weights.clone() finalized = run() assert isinstance(finalized, torch.Tensor) - moe_config.defer_moe_finalize() + moe_config.defer_moe_finalize(allow_ep=allow_ep) output = run() assert moe_config.should_defer_moe_finalize(finalized.shape[0]) != chunked if chunked: @@ -621,11 +622,16 @@ def check_deferred_moe_finalize( output.expert_weights, expected_weights, atol=0, rtol=0 ) reference = torch.empty_like(finalized) + permutation = output.expanded_idx_to_permuted_idx + # moe_unpermute skips out-of-range rows, while TRTLLM-Gen marks them -1. + permutation = permutation.masked_fill( + permutation < 0, output.gemm2_permuted.shape[0] + ) moe_unpermute( reference, output.gemm2_permuted, output.expert_weights.float(), - output.expanded_idx_to_permuted_idx, + permutation, # The kernel reads its valid-row count through this pointer either way. expert_first_token_offset=output.gemm2_permuted.new_full( (1,), output.gemm2_permuted.shape[0], dtype=torch.int64 diff --git a/tests/kernels/test_fp32_router_gemm.py b/tests/kernels/test_fp32_router_gemm.py index d330abba532a..0f55d860a6c8 100644 --- a/tests/kernels/test_fp32_router_gemm.py +++ b/tests/kernels/test_fp32_router_gemm.py @@ -3,9 +3,9 @@ """Tests for fp32_router_gemm kernel: activation×weight→fp32. Supported (hidden_size, num_experts) pairs: - (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3 + (3072, 256), (6144, 128), (5120, 256) -Correctness baseline: F.linear in float32. Every M in [1, 32] is covered so +Correctness baseline: F.linear in float32. Every supported M is covered so all tuned geometries (wide-block, experts-per-block, token-group; boundaries at M=4/5, odd/even, M=15/16) are exercised on Blackwell, and the legacy 128/1 geometry everywhere else. @@ -18,8 +18,12 @@ from vllm.platforms import current_platform # (hidden_size, num_experts) -SHAPES = [(3072, 256), (6144, 128)] -ALL_M = list(range(1, 33)) +SHAPES = [(3072, 256), (6144, 128), (5120, 256)] +CASES = [ + (m, hidden, experts) + for hidden, experts in SHAPES + for m in range(1, 49 if hidden == 5120 else 33) +] # Absolute tolerance for fp32 kernel vs float64 reference ATOL_FP32 = 2e-4 ATOL_BF16 = 2e-2 # bf16 activation has lower precision @@ -40,8 +44,7 @@ def _ref(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor: return torch.nn.functional.linear(mat_a.float(), mat_b.float()) -@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES) -@pytest.mark.parametrize("num_tokens", ALL_M) +@pytest.mark.parametrize("num_tokens,hidden_dim,num_experts", CASES) def test_fp32_activation(num_tokens: int, hidden_dim: int, num_experts: int): """fp32 activation → fp32 output should match reference closely.""" _requires_sm90() @@ -58,8 +61,7 @@ def test_fp32_activation(num_tokens: int, hidden_dim: int, num_experts: int): torch.testing.assert_close(out, ref, atol=ATOL_FP32, rtol=0) -@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES) -@pytest.mark.parametrize("num_tokens", ALL_M) +@pytest.mark.parametrize("num_tokens,hidden_dim,num_experts", CASES) def test_bf16_activation(num_tokens: int, hidden_dim: int, num_experts: int): """bf16 activation → fp32 output should match reference within bf16 error.""" _requires_sm90() @@ -136,17 +138,19 @@ def test_zero_tokens_returns_empty(): assert out.dtype == torch.float32 -def test_rejects_invalid_inputs(): +@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES) +def test_rejects_invalid_inputs(hidden_dim, num_experts): """The entry must fail loudly, never compute silently wrong results.""" _requires_sm90() device = torch.device("cuda") - hidden_dim, num_experts = SHAPES[0] mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device) - # num_tokens > 32 (beyond the instantiated range) + # Beyond the instantiated range for this shape. + invalid_tokens = 49 if hidden_dim == 5120 else 33 with pytest.raises(Exception, match="num_tokens"): fp32_router_gemm( - torch.randn(33, hidden_dim, dtype=torch.float32, device=device), mat_b + torch.randn(invalid_tokens, hidden_dim, dtype=torch.float32, device=device), + mat_b, ) # unsupported (hidden_dim, num_experts) pair @@ -173,3 +177,36 @@ def test_rejects_invalid_inputs(): fp32_router_gemm( torch.randn(4, hidden_dim, dtype=torch.float16, device=device), mat_b ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + "tokens", [*range(1, 17), 17, 24, 32, 33, 40, 48, 56, 64, 65, 128] +) +@torch.inference_mode() +def test_router_retains_fp32_weights_and_logits(default_vllm_config, dist_init, tokens): + """Low-latency routing must retain FP32 weight precision and logits.""" + from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear + from vllm.platforms import current_platform + + if not current_platform.is_device_capability_family(100): + pytest.skip("SM100 required") + torch.manual_seed(42) + router = GateLinear( + 5120, 256, params_dtype=torch.float32, out_dtype=torch.float32 + ).cuda() + router.weight.normal_(std=0.01) + x = torch.randn(tokens, 5120, device="cuda", dtype=torch.bfloat16) + expected = (x.double() @ router.weight.double().T).float() + actual, bias = router(x) + assert bias is None and actual.dtype == torch.float32 + torch.testing.assert_close(actual, expected, rtol=2e-5, atol=2e-5) + torch.testing.assert_close(actual.topk(8).indices, expected.topk(8).indices) + if tokens in (1, 16, 48, 56): + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured, _ = router(x) + x.normal_() + graph.replay() + expected = (x.double() @ router.weight.double().T).float() + torch.testing.assert_close(captured, expected, rtol=2e-5, atol=2e-5) diff --git a/tests/kernels/test_ll_bf16_gemm.py b/tests/kernels/test_ll_bf16_gemm.py index 1c7a15d514c9..76ea03f60ec8 100644 --- a/tests/kernels/test_ll_bf16_gemm.py +++ b/tests/kernels/test_ll_bf16_gemm.py @@ -566,7 +566,7 @@ def test_invalid_output_dtype(): with pytest.raises(ValueError, match="output_dtype=torch.float32"): from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ll_bf16_gemm - ll_bf16_gemm(a, b, output_dtype=torch.bfloat16) + ll_bf16_gemm(a, b, output_dtype=torch.float16) def test_cache_miss_compiles_dotprod(): @@ -581,5 +581,29 @@ def test_cache_miss_compiles_dotprod(): _assert_close(out, _ref(a, b), context="cache miss dotprod") +@pytest.mark.parametrize("M", [1, 2, 8, 16]) +@pytest.mark.parametrize("N", [1792, 3584]) +@torch.inference_mode() +def test_bf16_output_and_fp32_cache_coexist(M, N): + """BF16 epilogues retain FP32 accumulation and can replay with new inputs.""" + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ll_bf16_gemm + + torch.manual_seed(42) + x = torch.randn(M, 5120, device="cuda", dtype=torch.bfloat16) + w = torch.randn(N, 5120, device="cuda", dtype=torch.bfloat16) * 0.01 + for dtype in (torch.float32, torch.bfloat16): + actual = ll_bf16_gemm(x, w, output_dtype=dtype) + assert actual.dtype == dtype + reference = (x.double() @ w.double().T).to(dtype) + torch.testing.assert_close(actual, reference, rtol=0.008, atol=1e-5) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = ll_bf16_gemm(x, w, output_dtype=torch.bfloat16) + x.normal_() + graph.replay() + reference = (x.double() @ w.double().T).bfloat16() + torch.testing.assert_close(actual, reference, rtol=0.008, atol=1e-5) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/models/registry.py b/tests/models/registry.py index f9dfafb40479..562fe18f07e9 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -216,6 +216,12 @@ def check_available_online( "BailingMoeV3ForCausalLM": _HfExamplesInfo( "inclusionAI/Ling-3.0-flash", trust_remote_code=True ), + "BeamForCausalLM": _HfExamplesInfo( + "reflection/Beam-1.0-501B-A23B-NVFP4", + is_available_online=False, + # Include MoE in the reduced one-layer initialization test. + hf_overrides={"mlp_only_layers": []}, + ), "BloomForCausalLM": _HfExamplesInfo( "bigscience/bloom-560m", {"1b": "bigscience/bloomz-1b1"} ), @@ -1490,6 +1496,12 @@ def check_available_online( max_model_len=8192, max_num_seqs=32, ), + "DFlash2BeamDraftModel": _HfExamplesInfo( + "reflection/Beam-1.0-501B-A23B-NVFP4", + speculative_model="reflection/Beam-1.0-501B-A23B-DFlash2", + is_available_online=False, + use_original_num_layers=True, + ), "DFlashLagunaForCausalLM": _HfExamplesInfo( "poolside/Laguna-XS-2.1-NVFP4", speculative_model="poolside/Laguna-XS-2.1-DFlash-NVFP4", diff --git a/tests/models/test_beam_construction.py b/tests/models/test_beam_construction.py new file mode 100644 index 000000000000..3dce0dcd558d --- /dev/null +++ b/tests/models/test_beam_construction.py @@ -0,0 +1,562 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Construct Beam from a synthetic config and generate with dummy weights. + +Exercises NoPE + SSMax, hybrid sliding-window attention, the sigmoid output +gate and its per-head bias, sandwich norm, depth-scaled residuals, grouped +sigmoid routing, and logit softcapping without downloading a checkpoint. +""" + +import json +from dataclasses import replace +from operator import attrgetter +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm import LLM, SamplingParams +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput +from vllm.models.beam.nvidia import model as beam +from vllm.models.beam.nvidia import moe as beam_moe +from vllm.transformers_utils.configs.beam import BeamConfig +from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( + set_eagle3_aux_hidden_state_layers, +) + +# `moe_intermediate_size` must be a multiple of 128 for the BF16 MoE kernel. +TINY_CONFIG = { + "architectures": ["BeamForCausalLM"], + "model_type": "beam", + "vocab_size": 1024, + "hidden_size": 128, + "intermediate_size": 256, + "num_hidden_layers": 2, + "num_attention_heads": 8, + "num_key_value_heads": 2, + "head_dim": 16, + "max_position_embeddings": 512, + "rms_norm_eps": 1e-6, + "rope_theta": 10000.0, + "sliding_window": 128, + "moe_intermediate_size": 256, + "num_experts_per_tok": 2, + "num_experts": 8, + "n_group": 2, + "topk_group": 1, + "routed_scaling_factor": 1.0, + "norm_topk_prob": True, + # Layer 1 is NoPE, so it takes the SSMax path and full attention while + # layer 0 stays RoPE + sliding. Exercises both branches in one model. + "layer_types": ["sliding_attention", "full_attention"], + "rope_parameters": { + "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, + "full_attention": None, + }, + "mlp_layer_types": ["dense", "sparse"], + "final_logit_softcapping": 30.0, + "torch_dtype": "bfloat16", + "bos_token_id": 1, + "eos_token_id": 2, +} + + +@pytest.mark.usefixtures("enable_pickle") +@pytest.mark.parametrize( + "enforce_eager,tp_size,num_experts,enable_ep", + [ + (True, 1, 8, False), + (False, 1, 8, False), + (True, 4, 2, False), + (True, 4, 8, True), + ], +) +def test_beam_builds_and_generates( + tmp_path, enforce_eager, tp_size, num_experts, enable_ep +): + """Release-style metadata must load and generate in both execution modes.""" + if torch.accelerator.device_count() < tp_size: + pytest.skip(f"Requires {tp_size} GPUs") + config = { + **TINY_CONFIG, + "num_experts": num_experts, + "moe_intermediate_size": max(256, 128 * tp_size), + "n_group": 1 if num_experts == 2 else 2, + } + (tmp_path / "config.json").write_text(json.dumps(config)) + + llm = LLM( + model=str(tmp_path), + load_format="dummy", + skip_tokenizer_init=True, + max_model_len=256, + max_num_seqs=1, + max_num_batched_tokens=256, + kv_cache_memory_bytes=16 * 1024**2, + gpu_memory_utilization=0.25, + enforce_eager=enforce_eager, + tensor_parallel_size=tp_size, + enable_expert_parallel=enable_ep, + ) + outputs = llm.generate( + {"prompt_token_ids": [5, 6, 7, 8]}, + SamplingParams(max_tokens=8, ignore_eos=True), + ) + + assert len(outputs[0].outputs[0].token_ids) == 8 + # TP shards every expert; only EP partitions the expert set across ranks. + local_experts = num_experts // tp_size if enable_ep else num_experts + assert llm.apply_model(attrgetter("num_local_physical_experts")) == ( + [local_experts] * tp_size + ) + + +def test_eplb_is_rejected_until_validated(monkeypatch): + """Expert parallelism is validated; expert rebalancing is not.""" + config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=BeamConfig()), + quant_config=None, + parallel_config=SimpleNamespace( + enable_eplb=True, use_sequence_parallel_moe=False + ), + ) + with pytest.raises(NotImplementedError, match="EPLB"): + beam.BeamForCausalLM(vllm_config=config) + + +def test_sequence_parallel_moe_is_rejected_until_validated(): + """DP attention with EP MoE takes the sequence-parallel path; unvalidated.""" + config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=BeamConfig()), + quant_config=None, + parallel_config=SimpleNamespace( + enable_eplb=False, use_sequence_parallel_moe=True + ), + ) + with pytest.raises(NotImplementedError, match="sequence-parallel"): + beam.BeamForCausalLM(vllm_config=config) + + +def test_dense_only_pipeline_stage_initializes(monkeypatch): + """A stage with dense layers and remote MoE layers needs no local experts.""" + layer = beam.BeamDecoderLayer.__new__(beam.BeamDecoderLayer) + nn.Module.__init__(layer) + layer.mlp = nn.Identity() + stage = nn.Module() + stage.layers = nn.ModuleList([layer, beam.PPMissingLayer()]) + stage.make_empty_intermediate_tensors = lambda *args, **kwargs: None + monkeypatch.setattr(beam, "BeamModel", lambda **kwargs: stage) + monkeypatch.setattr( + beam, "get_pp_group", lambda: SimpleNamespace(is_last_rank=False) + ) + monkeypatch.setattr(beam, "LogitsProcessor", lambda *args: nn.Identity()) + config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=BeamConfig()), + quant_config=None, + parallel_config=SimpleNamespace( + enable_eplb=False, use_sequence_parallel_moe=False + ), + ) + + model = beam.BeamForCausalLM(vllm_config=config) + + assert model.moe_layers == [] + assert model.num_moe_layers == 0 + assert model.num_logical_experts == 0 + assert model.num_local_physical_experts == 0 + + +class _AddOneLayer(nn.Module): + """Adds one to the residual stream and hands the next layer double that.""" + + def __init__(self): + super().__init__() + self.input_layernorm = nn.Identity() + + def forward(self, positions, hidden_states, normalized_states, next_norm_weight): + hidden_states = hidden_states + 1 + return hidden_states, hidden_states * 2 + + +@pytest.mark.cpu_test +def test_dflash_target_layers_capture_layer_outputs(monkeypatch): + """DFlash `target_layer_ids` name the layer whose output is captured.""" + model = beam.BeamModel.__new__(beam.BeamModel) + nn.Module.__init__(model) + model.start_layer, model.end_layer = 0, 3 + model.layers = nn.ModuleList([_AddOneLayer() for _ in range(3)]) + model.norm = nn.Identity() + target = beam.BeamForCausalLM.__new__(beam.BeamForCausalLM) + nn.Module.__init__(target) + target.model = model + monkeypatch.setattr( + beam, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + spec_config = SimpleNamespace( + draft_model_config=SimpleNamespace( + hf_config=SimpleNamespace(dflash_config={"target_layer_ids": [0, 1]}) + ) + ) + + set_eagle3_aux_hidden_state_layers(target, spec_config) + output, aux_hidden_states = target( + input_ids=None, + positions=torch.tensor([0]), + inputs_embeds=torch.tensor([[1.0, 2.0]]), + ) + + torch.testing.assert_close(output, torch.tensor([[8.0, 10.0]])) + assert len(aux_hidden_states) == 2 + torch.testing.assert_close(aux_hidden_states[0], torch.tensor([[2.0, 3.0]])) + torch.testing.assert_close(aux_hidden_states[1], torch.tensor([[3.0, 4.0]])) + + +@pytest.mark.parametrize("scale,expected", [(None, 1.0), (0.0, 0.0), (2.5, 2.5)]) +@pytest.mark.parametrize("sequence_parallel", [False, True]) +def test_moe_preserves_explicit_routed_scaling( + monkeypatch, scale, expected, sequence_parallel +): + """Only an absent scale defaults to one; zero must disable routed output.""" + + def gate(*args, out_dtype, **kwargs): + result = nn.Module() + result.out_dtype = out_dtype + return result + + def experts(**kwargs): + assert kwargs["routed_scaling_factor"] == expected + assert kwargs["apply_routed_scale_to_output"] is True + assert kwargs["is_sequence_parallel"] is sequence_parallel + assert kwargs["reduce_results"] is sequence_parallel + return nn.Identity() + + monkeypatch.setattr(beam_moe, "GateLinear", gate) + monkeypatch.setattr(beam_moe, "FusedMoEFactory", experts) + config = SimpleNamespace( + model_config=SimpleNamespace( + hf_config=BeamConfig(**{**TINY_CONFIG, "routed_scaling_factor": scale}) + ), + quant_config=None, + compilation_config=SimpleNamespace(static_forward_context={}), + parallel_config=SimpleNamespace( + use_sequence_parallel_moe=sequence_parallel, + enable_eplb=False, + eplb_config=SimpleNamespace(num_redundant_experts=0), + ), + ) + beam_moe.BeamSparseMoeBlock(config, reduce_results=False) + + +def _deferred_moe_config(**parallel_overrides): + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + ) + + parallel = replace(FusedMoEParallelConfig.make_no_parallel(), **parallel_overrides) + return FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=1024, + intermediate_size=1024, + num_local_experts=8 // parallel.ep_size, + num_logical_experts=8, + activation=MoEActivation.SILU, + device="cpu", + routing_method=RoutingMethodType.Renormalize, + moe_parallel_config=parallel, + in_dtype=torch.bfloat16, + skip_final_all_reduce=True, + ) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "parallel_overrides,supported,padded,expected", + [ + ({}, True, False, True), + ({"dp_size": 4, "ep_size": 4}, True, False, True), + ({"dp_size": 1, "ep_size": 1}, True, False, False), + ({"tp_size": 2}, True, False, False), + ({"pcp_size": 2}, True, False, False), + ({}, False, False, False), + ({}, True, True, False), + ], +) +def test_dp_moe_boundary_fusion_requires_supported_topology( + monkeypatch, parallel_overrides, supported, padded, expected +): + """DP owns its returned tokens and defers only with a full boundary consumer.""" + from vllm.config.kernel import FLASHINFER_MOE_EP_CUTEDSL + from vllm.model_executor.layers.fused_moe import flashinfer_moe_ep + from vllm.model_executor.layers.fused_moe.experts.flashinfer_moe_ep import ( + FlashInferMoeEpExperts, + ) + + parallel = {"dp_size": 2, "ep_size": 2, "use_ep": True, **parallel_overrides} + moe = replace( + _deferred_moe_config(**parallel), + moe_backend=FLASHINFER_MOE_EP_CUTEDSL, + skip_final_all_reduce=False, + ) + if padded: + moe.hidden_dim_unpadded -= 1 + block = beam_moe.BeamSparseMoeBlock.__new__(beam_moe.BeamSparseMoeBlock) + nn.Module.__init__(block) + block.is_sequence_parallel = False + block.gather_workspace = None + block.output_boundary_fused = False + block._boundary_layer_name = "moe.finalize" + block._compilation_config = SimpleNamespace(static_forward_context={}) + block.experts = SimpleNamespace( + moe_config=moe, + routed_experts=SimpleNamespace( + quant_method=SimpleNamespace(experts_cls=FlashInferMoeEpExperts), + apply_router_weight_on_input=False, + ), + routed_scaling_factor=2.5, + shared_experts=None, + routed_input_transform=None, + routed_output_transform=None, + ) + monkeypatch.setattr( + flashinfer_moe_ep, + "supports_flashinfer_moe_ep_deferred_finalize", + lambda: supported, + ) + + assert block.fuse_output_post_norm(None, 1e-6) is expected + assert block.output_boundary_fused is expected + assert moe.use_deferred_moe_finalize is expected + assert block.experts.routed_scaling_factor == 2.5 + assert block.gather_workspace is None + if expected: + assert block._compilation_config.static_forward_context["moe.finalize"] is block + else: + assert not block._compilation_config.static_forward_context + + moe.defer_moe_finalize(allow_ep=True) + assert moe.use_deferred_moe_finalize is expected + assert moe.should_defer_moe_finalize(1) is expected + + +@pytest.mark.cpu_test +@pytest.mark.parametrize("next_norm", [False, True]) +@pytest.mark.parametrize("sequence_parallel", [False, True]) +@pytest.mark.parametrize("rows", [0, 5]) +def test_moe_boundary_consumes_borrowed_outputs_in_opaque_op( + monkeypatch, next_norm, sequence_parallel, rows +): + """The boundary op preserves raw-output ownership and explicit residual inputs.""" + from vllm.models.beam.nvidia import moe_finalize + + block = beam_moe.BeamSparseMoeBlock.__new__(beam_moe.BeamSparseMoeBlock) + nn.Module.__init__(block) + block.is_sequence_parallel = sequence_parallel + block.output_boundary_fused = True + block.gather_workspace = ( + SimpleNamespace(world_size=2) if sequence_parallel else None + ) + local_rows = (rows + 1) // 2 if sequence_parallel else rows + partials = torch.randn(local_rows * 2, 16) + weights = torch.randn(local_rows, 2) + mapping = torch.arange(local_rows * 2, dtype=torch.int32).view(local_rows, 2) + routed_inputs = [] + calls = [] + + def routed(hidden_states, router_logits, shared_experts_input): + assert hidden_states is router_logits + assert shared_experts_input is None + routed_inputs.append(hidden_states) + return UnfinalizedMoEOutput(partials, weights, mapping) + + encoded_name = beam_moe._layer_name_type("moe") + block.experts = SimpleNamespace( + layer_name="moe", + _encode_layer_name=lambda: encoded_name, + _forward_impl=routed, + routed_scaling_factor=2.5, + ) + monkeypatch.setattr(beam_moe, "get_layer_from_name", lambda name: block.experts) + monkeypatch.setattr( + beam_moe, + "get_forward_context", + lambda: SimpleNamespace(no_compile_layers={"moe.finalize": block}), + ) + monkeypatch.setattr( + beam_moe, "sequence_parallel_chunk", lambda x: x[: (x.shape[0] + 1) // 2] + ) + + def finalize(x, residual, *args, **kwargs): + assert x.gemm2_permuted is partials and x.expert_weights is weights + assert x.expanded_idx_to_permuted_idx is mapping + calls.append((args, kwargs)) + return residual + 1, residual + 2 if args[3] is not None else None + + finalizer = ( + "beam_moe_finalize_gather_norm_residual_norm" + if sequence_parallel + else "beam_moe_finalize_norm_residual_norm" + ) + monkeypatch.setattr(moe_finalize, finalizer, finalize) + forward = block.forward_boundary + hidden_states, residual = torch.randn(rows, 16), torch.randn(rows, 16) + post_weight = torch.randn(16) + next_weight = torch.randn(16) if next_norm else None + with torch.library._scoped_library("vllm", "IMPL") as lib: + if not torch._C._dispatch_has_kernel_for_dispatch_key( + "vllm::beam_moe_forward_boundary", "CPU" + ): + lib.impl( + "beam_moe_forward_boundary", + beam_moe._beam_moe_forward_boundary, + "CPU", + ) + result, normalized = forward( + hidden_states, + residual, + post_weight, + 1e-6, + next_weight, + 1e-6 if next_norm else None, + (0.25, 0.75), + ) + torch.testing.assert_close(result, residual + 1) + if next_norm: + torch.testing.assert_close(normalized, residual + 2) + else: + assert normalized is None + torch.testing.assert_close(routed_inputs[0], hidden_states[:local_rows]) + assert len(calls) == 1 + args, kwargs = calls[0] + assert args[0] is post_weight and args[2] is next_weight + assert args[4] == (0.25, 0.75) + if sequence_parallel: + assert args[5] is block.gather_workspace + else: + assert len(args) == 5 + assert kwargs == {"routed_scale": 2.5} + + +@pytest.mark.cpu_test +@pytest.mark.parametrize("defer_output", [False, True]) +def test_sp_moe_combine_stays_reduced_when_finalization_is_deferred(defer_output): + """Deferring the local expert sum must never introduce an EP all-reduce.""" + from vllm.model_executor.layers.fused_moe.prepare_finalize.passthrough import ( + PassThroughPrepareAndFinalize, + ) + + prepare_finalize = PassThroughPrepareAndFinalize() + prepare_finalize.post_init_setup( + SimpleNamespace( + moe_config=SimpleNamespace(use_deferred_moe_finalize=defer_output) + ) + ) + assert prepare_finalize.output_is_reduced() + assert prepare_finalize.supports_deferred_moe_finalize() is defer_output + + +@pytest.mark.parametrize( + "allow_ep,parallel_overrides,padded,expected", + [ + (False, {}, False, False), + (True, {}, False, True), + (True, {"dp_size": 2}, False, False), + (True, {"pcp_size": 2}, False, False), + (True, {"sp_size": 4}, False, False), + (True, {}, True, False), + ], +) +def test_moe_finalize_deferral_requires_replicated_tokens( + allow_ep, parallel_overrides, padded, expected +): + """EP deferral requires explicit consumer support and unpadded replicated tokens.""" + config = _deferred_moe_config(ep_size=4, use_ep=True, **parallel_overrides) + if padded: + config.hidden_dim_unpadded -= 1 + + config.defer_moe_finalize(4095, allow_ep=allow_ep) + + assert config.should_defer_moe_finalize(64) is expected + assert config.should_defer_moe_finalize(4095) is expected + assert not config.should_defer_moe_finalize(4096) + assert not config.should_defer_moe_finalize(0) + + +@pytest.mark.parametrize( + "world_size,ep_size,max_tokens,finalize_limit,backend_limit,expected_limit,enabled", + [ + (2, 2, 16384, 8192, -1, 8192, True), + (4, 4, 131072, 65536, -1, 65536, True), + (8, 8, 16384, 8192, -1, 8192, True), + (4, 4, 1024, 65536, -1, 1024, True), + (4, 4, 16384, 65536, 128, 128, True), + (4, 1, 131072, 65536, -1, 131072, True), + (4, 8, 8192, 65536, -1, 8192, False), + ], +) +@pytest.mark.parametrize("sequence_parallel", [False, True]) +def test_decoder_finalizes_large_ep_batches_before_flashinfer_allreduce( + monkeypatch, + world_size, + ep_size, + max_tokens, + finalize_limit, + backend_limit, + expected_limit, + enabled, + sequence_parallel, +): + """Large EP batches return tensors; TP and smaller workspace limits are retained.""" + from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( + TrtLlmNvFp4ExpertsMonolithic, + ) + + moe_config = _deferred_moe_config( + tp_size=world_size if ep_size == 1 else 1, + ep_size=ep_size, + use_ep=ep_size > 1, + sp_size=world_size if sequence_parallel else 1, + ) + moe_config.limit_deferred_moe_finalize(backend_limit) + + def init_moe(self, *args, **kwargs): + nn.Module.__init__(self) + self.is_sequence_parallel = sequence_parallel + self.experts = SimpleNamespace( + moe_config=moe_config, + routed_scaling_factor=2.5, + routed_experts=SimpleNamespace( + quant_method=SimpleNamespace(experts_cls=TrtLlmNvFp4ExpertsMonolithic) + ), + ) + + monkeypatch.setattr(beam.BeamSparseMoeBlock, "__init__", init_moe) + monkeypatch.setattr(beam, "BeamAttention", lambda *args, **kwargs: nn.Identity()) + config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=BeamConfig(**TINY_CONFIG)), + cache_config=None, + quant_config=None, + ) + workspace = SimpleNamespace( + world_size=world_size, + max_tokens=max_tokens, + ep_finalize_max_tokens=finalize_limit, + ep_finalize_oneshot_max_tokens=64, + ) + + layer = beam.BeamDecoderLayer(config, "model.layers.1", workspace) + + assert moe_config.should_defer_moe_finalize(expected_limit) is ( + enabled and not sequence_parallel + ) + assert not moe_config.should_defer_moe_finalize(expected_limit + 1) + if sequence_parallel: + assert layer.mlp_boundary is beam.beam_norm_residual_norm + assert layer.finalize_boundary is None + assert layer.attention_boundary.func is beam.beam_allreduce_norm_residual_norm diff --git a/tests/models/test_beam_numerics.py b/tests/models/test_beam_numerics.py new file mode 100644 index 000000000000..2741a470bd04 --- /dev/null +++ b/tests/models/test_beam_numerics.py @@ -0,0 +1,1152 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regressions for Beam normalization and activation rounding.""" + +import pytest +import torch + + +def _torch_rms_norm(x, weight, eps): + y = x.float() + y = y * torch.rsqrt(y.square().mean(-1, keepdim=True) + eps) + if weight is not None: + y = y * weight.float() + return y.to(x.dtype) + + +def _torch_boundary(x, residual, post_weight, post_eps, next_weight, next_eps, scales): + if post_eps is not None: + x = _torch_rms_norm(x, post_weight, post_eps) + residual = residual.to(x.dtype) + if scales is not None: + x = x * scales[0] + residual = residual * scales[1] + hidden = x + residual + normalized = ( + _torch_rms_norm(hidden, next_weight, next_eps) if next_eps is not None else None + ) + return hidden, normalized + + +def _assert_boundary_close(actual, expected, *, relative_tolerance=1e-4): + for value, reference in zip(actual, expected): + if reference is None: + assert value is None + continue + bf16 = reference.dtype == torch.bfloat16 + torch.testing.assert_close( + value, + reference, + rtol=0.016 if bf16 else 2e-5, + atol=0.016 if bf16 else 2e-6, + ) + # Allow isolated BF16 rounding differences, but catch omitted casts. + if reference.numel(): + error = (value.float() - reference.float()).norm() + assert error <= reference.float().norm() * ( + relative_tolerance if bf16 else 2e-6 + ) + + +def _torch_megamoe_finalize(experts, weights, routed_scale, post_weight, post_eps): + # BF16 x FP32 fits in FP64; round once per fused multiply-add to FP32. + value = (experts[:, 0].double() * weights[:, 0, None].double()).float() + for k in range(1, experts.shape[1]): + value = ( + experts[:, k].double() * weights[:, k, None].double() + value.double() + ).float() + value = (value.bfloat16().float() * routed_scale).bfloat16() + if post_eps is not None: + value = _torch_rms_norm(value, post_weight, post_eps) + return value + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + "rows,hidden,top_k", + [(0, 5120, 8), (33, 5120, 8), (7, 256, 1), (3, 8192, 32)], +) +@pytest.mark.parametrize("norm", ["none", "unweighted", "weighted"]) +@torch.inference_mode() +def test_megamoe_finalize_preserves_rounding_and_ignores_capacity( + rows, hidden, top_k, norm +): + """Padded experts stay unread and the live result keeps both BF16 casts.""" + pytest.importorskip("flashinfer") + from vllm.models.beam.nvidia.moe_finalize import beam_moe_finalize + + torch.manual_seed(42) + experts = torch.randn(rows + 3, top_k, hidden, device="cuda", dtype=torch.bfloat16) + experts[:rows:2].mul_(1e-3) + experts[rows:] = float("nan") + weights = torch.rand(rows + 3, top_k, device="cuda", dtype=torch.float32) + weights[rows:] = float("nan") + post_weight = ( + torch.randn(hidden, device="cuda", dtype=torch.bfloat16) + if norm == "weighted" + else None + ) + post_eps = None if norm == "none" else 1e-6 + storage = torch.full((rows + 3, hidden), 37, device="cuda", dtype=experts.dtype) + output = storage[:rows] + actual = beam_moe_finalize(experts, weights, output, 0.137, post_weight, post_eps) + expected = _torch_megamoe_finalize( + experts[:rows], weights[:rows], 0.137, post_weight, post_eps + ) + assert actual is output + if post_eps is None: + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + else: + _assert_boundary_close((actual,), (expected,)) + torch.testing.assert_close(storage[rows:], torch.full_like(storage[rows:], 37)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + "rows,hidden,top_k", + [(0, 5120, 8), (33, 5120, 8), (3, 8192, 32)], +) +@pytest.mark.parametrize( + "post_norm,next_norm,scales,routed_scale", + [ + ("none", "none", None, 1.0), + ("unweighted", "weighted", (0.173, 0.937), 0.137), + ("weighted", "unweighted", None, 0.137), + ], +) +@torch.inference_mode() +def test_megamoe_local_boundary_preserves_rounding_and_ignores_capacity( + rows, hidden, top_k, post_norm, next_norm, scales, routed_scale +): + """DP finalization keeps every BF16 boundary and ignores workspace padding.""" + pytest.importorskip("flashinfer") + from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput + from vllm.models.beam.nvidia.moe_finalize import ( + beam_moe_finalize_norm_residual_norm, + ) + + torch.manual_seed(42) + capacity = rows + 3 + experts = torch.randn(capacity, top_k, hidden, device="cuda", dtype=torch.bfloat16) + experts[:rows:2].mul_(1e-3) + experts[rows:] = float("nan") + weights = torch.rand(capacity, top_k, device="cuda") + weights[rows:] = float("nan") + mapping = torch.arange(capacity * top_k, device="cuda", dtype=torch.int32).view( + capacity, top_k + ) + output = UnfinalizedMoEOutput(experts.view(-1, hidden), weights, mapping) + residual = torch.randn(rows, hidden, device="cuda", dtype=experts.dtype) + original_residual = residual.clone() + + def norm_args(kind): + weight = ( + torch.randn(hidden, device="cuda", dtype=experts.dtype) + if kind == "weighted" + else None + ) + return weight, None if kind == "none" else 1e-6 + + post_weight, post_eps = norm_args(post_norm) + next_weight, next_eps = norm_args(next_norm) + actual = beam_moe_finalize_norm_residual_norm( + output, + residual, + post_weight, + post_eps, + next_weight, + next_eps, + scales, + routed_scale, + ) + finalized = _torch_megamoe_finalize( + experts[:rows], weights[:rows], routed_scale, None, None + ) + expected = _torch_boundary( + finalized, residual, post_weight, post_eps, next_weight, next_eps, scales + ) + _assert_boundary_close(actual, expected) + torch.testing.assert_close(residual, original_residual, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("strategy", ["fused", "unfused"]) +@torch.inference_mode() +def test_megamoe_local_boundary_graph_replays_with_live_inputs(strategy): + """Captured DP finalization reads current routing, residual, and norm weights.""" + pytest.importorskip("flashinfer") + from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput + from vllm.models.beam.nvidia.moe_finalize import ( + beam_moe_finalize_norm_residual_norm, + ) + + torch.manual_seed(42) + experts = torch.randn(19, 8, 5120, device="cuda", dtype=torch.bfloat16) + weights = torch.rand(19, 8, device="cuda") + mapping = torch.arange(19 * 8, device="cuda", dtype=torch.int32).view(19, 8) + output = UnfinalizedMoEOutput(experts.view(-1, 5120), weights, mapping) + residual = torch.randn(17, 5120, device="cuda", dtype=experts.dtype) + post_weight = torch.randn(5120, device="cuda", dtype=experts.dtype) + next_weight = torch.randn_like(post_weight) + scales = (0.173, 0.937) + + def run(): + return beam_moe_finalize_norm_residual_norm( + output, + residual, + post_weight, + 1e-6, + next_weight, + 1e-6, + scales, + 0.137, + strategy=strategy, + ) + + run() + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = run() + for _ in range(2): + experts.add_(0.125) + weights.copy_(weights.roll(1, dims=1)) + residual.sub_(0.25) + post_weight.neg_() + next_weight.mul_(0.5) + graph.replay() + finalized = _torch_megamoe_finalize( + experts[:17], weights[:17], 0.137, None, None + ) + expected = _torch_boundary( + finalized, residual, post_weight, 1e-6, next_weight, 1e-6, scales + ) + _assert_boundary_close(actual, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("rows", [0, 1, 8, 128]) +@pytest.mark.parametrize("cap", [7.3, 30.0]) +@torch.inference_mode() +def test_softcap_rounds_into_original_logits_without_touching_padding(dtype, rows, cap): + from vllm.models.beam.nvidia.activation import beam_softcap_logits + + torch.manual_seed(42) + storage = torch.randn(rows, 264, device="cuda", dtype=dtype) * (cap * 3) + logits = storage[:, :257] + if rows: + logits[0, :6] = torch.tensor( + [0.0, -0.0, float("inf"), -float("inf"), float("nan"), 1e-38], + device=logits.device, + dtype=dtype, + ) + original = storage.clone() + expected = (torch.tanh(logits.float() / cap) * cap).to(dtype) + actual = beam_softcap_logits(logits, cap) + + assert actual is logits + assert actual.dtype == dtype + torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) + torch.testing.assert_close( + storage[:, 257:], original[:, 257:], rtol=0, atol=0, equal_nan=True + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize( + "rows,width", [(0, 2560), (1, 5120), (8, 2560), (33, 513), (2048, 2560)] +) +@torch.inference_mode() +def test_dense_activation_rounds_after_product(default_vllm_config, dtype, rows, width): + """BF16 I/O must not round the SiLU result before multiplication.""" + import vllm._custom_ops # noqa: F401 + from vllm.models.beam.nvidia.activation import beam_silu_mul + + torch.manual_seed(42) + x = torch.randn(rows, 2 * width, device="cuda", dtype=dtype) * 4 + expected = torch.empty(rows, width, device="cuda", dtype=torch.float32) + if rows: + torch.ops._C.silu_and_mul(expected, x.float()) + actual = beam_silu_mul(x) + torch.testing.assert_close( + actual, + expected.to(dtype), + rtol=0.008 if dtype == torch.bfloat16 else 1e-6, + atol=1e-7, + ) + if rows: + error = (actual.float() - expected.to(dtype).float()).norm() + assert error <= expected.norm() * 1e-5 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("has_weight", [False, True]) +@pytest.mark.parametrize("shape", [(33, 513), (17, 6, 128)]) +@torch.inference_mode() +def test_fused_norm_preserves_fp32_affine_for_strided_rows(dtype, has_weight, shape): + """QKV slices and near-epsilon values retain the reference's FP32 affine.""" + from vllm.models.beam.nvidia.layernorm import beam_rms_norm + + generator = torch.Generator(device="cuda").manual_seed(42) + storage = torch.randn( + (*shape[:-1], shape[-1] + 64), + device="cuda", + dtype=dtype, + generator=generator, + ) + x = storage[..., 1 : shape[-1] + 1] + x[::2].mul_(1e-3) + weight = ( + torch.randn(shape[-1], device="cuda", dtype=dtype, generator=generator) + if has_weight + else None + ) + expected = _torch_rms_norm(x, weight, 1e-6) + actual = beam_rms_norm(x, weight, 1e-6) + # FP32 reduction trees may differ; allow one BF16 ULP after final rounding. + torch.testing.assert_close( + actual, + expected, + rtol=0.008 if dtype == torch.bfloat16 else 1e-6, + atol=1e-7, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("shape", [(0, 5120), (33, 513)]) +@pytest.mark.parametrize( + "post,next_,scaled", + [(True, True, True), (True, True, False), (False, True, True), (True, False, True)], +) +@torch.inference_mode() +def test_norm_residual_norm_preserves_intermediate_rounding( + dtype, shape, post, next_, scaled +): + """Both residual and next input must retain BF16 casts between fused stages.""" + from vllm.models.beam.nvidia.layernorm import beam_norm_residual_norm + + torch.manual_seed(42) + x = torch.randn(shape[0], shape[1] + 32, device="cuda", dtype=dtype)[ + :, 1 : shape[1] + 1 + ] + residual = torch.randn_like(x) + x[::2].mul_(1e-3) + post_weight = torch.randn(shape[1], device="cuda", dtype=dtype) if post else None + next_weight = torch.randn(shape[1], device="cuda", dtype=dtype) if next_ else None + scales = (0.125, 0.9682458365518543) if scaled else None + original_x, original_residual = x.clone(), residual.clone() + args = ( + x, + residual, + post_weight, + 1e-6 if post else None, + next_weight, + 1e-6 if next_ else None, + scales, + ) + expected = _torch_boundary(*args) + actual = beam_norm_residual_norm(*args) + _assert_boundary_close(actual, expected) + torch.testing.assert_close(x, original_x, rtol=0, atol=0) + torch.testing.assert_close(residual, original_residual, rtol=0, atol=0) + + +def _separate_qk_norm_rope( + q, + k, + qw, + kw, + eps, + positions, + cache, + head_dim, + rotary_dim=0, + neox=True, + mode="none", + offset=1, + divisor=1.0, +): + from vllm import _custom_ops as ops + + q = _torch_rms_norm( + q.reshape(-1, q.shape[-1] // head_dim, head_dim), qw, eps + ).reshape_as(q) + k = _torch_rms_norm( + k.reshape(-1, k.shape[-1] // head_dim, head_dim), kw, eps + ).reshape_as(k) + if rotary_dim and q.shape[0]: + qf, kf = q.float(), k.float() + ops.rotary_embedding(positions, qf, kf, head_dim, cache, neox) + q, k = qf.to(q.dtype), kf.to(k.dtype) + elif mode != "none": + p = positions.float() + offset + if mode == "default": + scale = p.log() / divisor + elif mode == "root10": + scale = (0.1 * torch.log1p(p / 2048)).exp() + else: + scale = (130 * torch.expm1(0.001 * p.log())).clamp(min=1.0) + q = (q * scale[:, None]).to(q.dtype) + return q, k + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("rows", [0, 1, 17, 512, 2049]) +@pytest.mark.parametrize("bias", [False, True]) +@pytest.mark.parametrize("packed", [False, True]) +@pytest.mark.parametrize("floor,temperature", [(0.13, 0.7), (0.0, 1.0)]) +@torch.inference_mode() +def test_output_gate_preserves_activation_dtype_rounding( + dtype, rows, bias, packed, floor, temperature +): + from vllm.models.beam.nvidia.attention_ops import beam_output_gate + + torch.manual_seed(42) + x = torch.randn(rows, 768, device="cuda", dtype=dtype) + gate = torch.randn(rows, 1792 if packed else 768, device="cuda", dtype=dtype) * 8 + gate = gate[:, -768:] + weight = torch.randn(6, device="cuda", dtype=dtype) if bias else None + kwargs = dict(num_heads=6, head_dim=128, floor=floor, temperature=temperature) + # Preserve each activation-dtype operation in the parent implementation. + g = gate.view(rows, 6, 128) + if weight is not None: + g = g + weight[None, :, None] + expected = x * ( + floor + (1 - floor) * torch.sigmoid(g.reshape_as(gate) / temperature) + ) + actual = beam_output_gate(x, gate, weight, **kwargs) + torch.testing.assert_close(actual, expected, rtol=2e-6, atol=1e-7) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize( + "rotary_dim,neox,weighted", [(0, True, False), (128, True, True), (64, False, True)] +) +@pytest.mark.parametrize("head_dim,q_heads,k_heads", [(128, 6, 1), (80, 4, 2)]) +@torch.inference_mode() +def test_qk_norm_rope_keeps_norm_rounding_and_fp32_rotation( + dtype, rotary_dim, neox, weighted, head_dim, q_heads, k_heads +): + from vllm.models.beam.nvidia.attention_ops import beam_qk_norm_rope + + torch.manual_seed(42) + rows = 17 + rotary_dim = min(rotary_dim, head_dim) + sizes = (q_heads * head_dim, k_heads * head_dim, k_heads * head_dim) + qkv = torch.randn(rows, sum(sizes) + 1, device="cuda", dtype=dtype)[:, 1:] + q, k, _ = qkv.split(sizes, -1) + q[::2].mul_(1e-3) + qw = torch.randn(head_dim, device="cuda", dtype=dtype) if weighted else None + kw = torch.randn(head_dim, device="cuda", dtype=dtype) if weighted else None + positions = torch.arange(rows, device="cuda") * 3 + angle = torch.randn(max(1, rows * 3), rotary_dim // 2, device="cuda") + cache = torch.cat((angle.cos(), angle.sin()), -1) + args = (q, k, qw, kw, 1e-6, positions, cache, head_dim, rotary_dim, neox) + expected = _separate_qk_norm_rope(*args) + actual = beam_qk_norm_rope(*args) + for value, reference in zip(actual, expected): + torch.testing.assert_close( + value, reference, rtol=0.016 if dtype == torch.bfloat16 else 2e-5, atol=2e-6 + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("rows", [1, 8, 513, 2049]) +@pytest.mark.parametrize("kv_heads", [1, 2]) +@pytest.mark.parametrize("nhd", [False, True]) +@pytest.mark.parametrize("rotary_dim,neox", [(128, True), (64, False), (0, True)]) +@torch.inference_mode() +def test_qk_norm_rope_inserts_only_mapped_kv(rows, kv_heads, nhd, rotary_dim, neox): + """Cache insertion preserves Q/K and follows live slot mappings on replay.""" + from vllm.models.beam.nvidia.attention_ops import beam_qk_norm_rope + + torch.manual_seed(42) + sizes = (6 * kv_heads * 128, kv_heads * 128, kv_heads * 128) + qkv = torch.randn(rows, sum(sizes) + 1, device="cuda", dtype=torch.bfloat16)[:, 1:] + q, k, v = qkv.split(sizes, -1) + weight = torch.randn(128, device="cuda") + positions = torch.arange(rows, device="cuda") + rope = torch.randn(rows, rotary_dim, device="cuda") if rotary_dim else None + args = (q, k, weight, weight, 1e-6, positions, rope, 128, rotary_dim, neox) + pages = (rows + 127) // 128 + 1 + shape = (pages, 128, kv_heads, 256) if nhd else (pages, kv_heads, 128, 256) + cache = torch.zeros(shape, device="cuda", dtype=q.dtype) + if nhd: + cache = cache.transpose(1, 2) + live = max(1, rows - 1) + slots = torch.randperm(pages * 128, device="cuda")[:live].clone() + if live > 1: + slots[-1] = -1 + + def fused(): + return beam_qk_norm_rope( + *args, ssmax_mode="root10", value=v, kv_cache=cache, slot_mapping=slots + ) + + fused() + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = fused() + # Mutate both data and metadata after capture, including a previously skipped slot. + qkv.add_(0.125) + slots.copy_(torch.randperm(pages * 128, device="cuda")[:live]) + if live > 1: + slots[0] = -1 + cache.zero_() + graph.replay() + expected = beam_qk_norm_rope(*args, ssmax_mode="root10") + for result, reference in zip(actual, expected): + torch.testing.assert_close(result, reference, rtol=0, atol=0) + valid = slots >= 0 + mapped = slots[valid] + reference_cache = torch.zeros_like(cache) + reference_cache[mapped // 128, :, mapped % 128, :128] = expected[1][:live][ + valid + ].reshape(-1, kv_heads, 128) + reference_cache[mapped // 128, :, mapped % 128, 128:] = v[:live][valid].reshape( + -1, kv_heads, 128 + ) + torch.testing.assert_close(cache, reference_cache, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize( + "mode,offset,divisor", + [ + ("default", 1, 1.0), + ("default", 16, 2.772588722239781), + ("root10", 1, 1.0), + ("root1000", 1, 1.0), + ], +) +@torch.inference_mode() +def test_ssmax_fuses_after_norm_rounding_and_scales_only_queries( + dtype, mode, offset, divisor +): + from vllm.models.beam.nvidia.attention_ops import beam_qk_norm_rope + + torch.manual_seed(42) + positions = torch.tensor( + [0, 1, 15, 2047, 2048, 16384, 100000, 1048575], device="cuda" + ) + qkv = torch.randn(len(positions), 1025, device="cuda", dtype=dtype)[:, 1:] + q, k, _ = qkv.split((768, 128, 128), -1) + qw, kw = [torch.randn(128, device="cuda", dtype=dtype) for _ in range(2)] + args = (q, k, qw, kw, 1e-6, positions, None, 128, 0, True, mode, offset, divisor) + actual, expected = beam_qk_norm_rope(*args), _separate_qk_norm_rope(*args) + for value, reference in zip(actual, expected): + # Norm and position scaling each round to BF16; allow two rounding steps. + torch.testing.assert_close( + value, + reference, + rtol=0.016 if dtype == torch.bfloat16 else 2e-6, + atol=2e-6, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("rotary_dim", [0, 128]) +@torch.inference_mode() +def test_attention_fusions_replay_with_changed_positions_and_gates(rotary_dim): + from vllm.models.beam.nvidia.attention_ops import ( + beam_output_gate, + beam_qk_norm_rope, + ) + + torch.manual_seed(42) + qkv = torch.randn(17, 1025, device="cuda", dtype=torch.bfloat16)[:, 1:] + q, k, _ = qkv.split((768, 128, 128), -1) + positions = torch.arange(17, device="cuda") + angle = torch.randn(128, 64, device="cuda") + cache = torch.cat((angle.cos(), angle.sin()), -1) + gate = torch.randn_like(q) + + def run(): + nq, nk = beam_qk_norm_rope( + q, + k, + None, + None, + 1e-6, + positions, + cache, + 128, + rotary_dim, + ssmax_mode="root10", + ) + return beam_output_gate( + nq, gate, None, num_heads=6, head_dim=128, floor=0.13, temperature=0.7 + ), nk + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + captured = run() + graph.replay() + torch.accelerator.synchronize() + positions.add_(31) + gate.neg_() + qkv.mul_(0.5) + expected = run() + graph.replay() + torch.accelerator.synchronize() + for actual, reference in zip(captured, expected): + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.accelerator.is_available(), reason="CUDA required") +@pytest.mark.parametrize("rotary_dim", [0, 128]) +@pytest.mark.parametrize("insert_kv", [False, True]) +@torch.inference_mode() +def test_fused_kernels_reuse_specializations_across_token_counts( + monkeypatch, rotary_dim, insert_kv +): + """Warm tile buckets once, then vary batch sizes, positions offsets and scales.""" + from triton import knobs + + from vllm.models.beam.nvidia.activation import beam_softcap_logits + from vllm.models.beam.nvidia.attention_ops import ( + beam_output_gate, + beam_qk_norm_rope, + ) + from vllm.models.beam.nvidia.layernorm import beam_norm_residual_norm, beam_rms_norm + + device = "cuda" + dtype = torch.bfloat16 + rows = 8193 + qkv = torch.randn(rows, 1024, device=device, dtype=dtype) + kv_cache = torch.zeros((65, 1, 128, 256), device=device, dtype=dtype) + slots = torch.arange(rows, device=device) + positions = torch.arange(rows + 16, device=device) + cache = torch.randn(rows + 16, 128, device=device) + weight = torch.ones(128, device=device, dtype=dtype) + bias = torch.zeros(6, device=device, dtype=dtype) + x = torch.randn(rows, 5120, device=device, dtype=dtype) + residual = torch.randn_like(x) + norm_weight = torch.ones(5120, device=device, dtype=dtype) + + def run(tokens, offset, scales): + q, k = beam_qk_norm_rope( + qkv[:tokens, :768], + qkv[:tokens, 768:896], + weight, + weight, + 1e-6, + positions[offset : offset + tokens], + cache if rotary_dim else None, + 128, + rotary_dim, + ssmax_mode="none" if rotary_dim else "root10", + value=qkv[:tokens, 896:] if insert_kv else None, + kv_cache=kv_cache if insert_kv else None, + slot_mapping=slots[: tokens - 1] if insert_kv else None, + ) + beam_output_gate( + q, q, bias, num_heads=6, head_dim=128, floor=0.1, temperature=1.2 + ) + beam_rms_norm(x[:tokens], norm_weight, 1e-6) + beam_softcap_logits(x[:tokens], 30.0) + beam_norm_residual_norm( + x[:tokens], + residual[:tokens], + norm_weight, + 1e-6, + norm_weight, + 1e-6, + scales, + ) + + for tokens in (1, 512, 2049): + run(tokens, 0, (0.125, 0.9682458365518543)) + + def unexpected_compile(**kwargs): + if kwargs["fn"].module.startswith("vllm.models.beam.nvidia"): + pytest.fail(f"Unexpected shape specialization: {kwargs['fn'].name}") + + with monkeypatch.context() as patch: + patch.setattr(knobs.runtime, "jit_cache_hook", unexpected_compile) + token_counts = (2, 7, 17, 511, 513, 2048, 2050, 8193) + for i, tokens in enumerate(token_counts): + run(tokens, i % 16, (0.125 + i / 1000, 0.968 - i / 1000)) + torch.accelerator.synchronize() + + +def _beam_allreduce_worker(rank, world_size, port, hidden_size, host="127.0.0.1"): + from datetime import timedelta + + import torch.distributed as dist + + from vllm.models.beam.nvidia.allreduce import ( + BeamAllReduceWorkspace, + beam_allreduce_norm_residual_norm, + ) + + torch.accelerator.set_device_index(rank % torch.accelerator.device_count()) + dist.init_process_group( + "nccl", + init_method=f"tcp://{host}:{port}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=120), + ) + group = dist.new_group(backend="gloo") + separate_cutoff = (10 * 1024**2 * world_size + 2 * hidden_size - 1) // ( + 2 * hidden_size + ) + max_tokens = max(513, separate_cutoff + 1) + workspace = BeamAllReduceWorkspace(group, hidden_size, max_tokens) + torch.manual_seed(42) + residual = torch.randn(max_tokens, hidden_size, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16) + torch.manual_seed(43 + rank) + x = torch.randn_like(residual) + # Exercise shared Lamport generations across batch and implementation changes. + cutoff = workspace.small_batch_max_tokens + batches = [0, 1, cutoff, cutoff + 1, 513, 1] + if workspace.separate_norm_min_tokens is not None: + cutoff = workspace.separate_norm_min_tokens + batches += [cutoff - 1, cutoff, cutoff + 1, 1] + for iteration, tokens in enumerate(batches): + torch.manual_seed(42 + iteration) + residual.normal_() + weight.normal_() + torch.manual_seed(43 + iteration + rank) + x.normal_() + x[:, ::7] = -0.0 + post, next_norm = iteration % 3 != 0, iteration % 3 != 1 + scales = (0.125, 0.9682458365518543) if iteration % 2 else None + args = ( + x[:tokens], + residual[:tokens], + weight if post else None, + 1e-6 if post else None, + weight if next_norm else None, + 1e-6 if next_norm else None, + scales, + ) + + def check(result, args): + partial, r, pw, peps, nw, neps, scales = args + # The collective rounds the FP32 sum before the post-norm. + reduced = partial.float() + dist.all_reduce(reduced) + reduced = reduced.bfloat16() + expected = _torch_boundary(reduced, r, pw, peps, nw, neps, scales) + _assert_boundary_close(result, expected, relative_tolerance=1e-3) + + check(beam_allreduce_norm_residual_norm(*args, workspace), args) + torch.accelerator.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = beam_allreduce_norm_residual_norm(*args, workspace) + for _ in range(2): + x.mul_(0.5) + graph.replay() + check(captured, args) + graph.reset() + # A small configured capacity still needs enough space for forced one-shot. + cutoff = workspace.small_batch_max_tokens + small_workspace = BeamAllReduceWorkspace(group, hidden_size, cutoff) + args = (x[:cutoff], residual[:cutoff], weight, 1e-6, weight, 1e-6, None) + check(beam_allreduce_norm_residual_norm(*args, small_workspace), args) + dist.barrier() + dist.destroy_process_group() + + +@pytest.mark.parametrize( + "world_size,hidden_size", [(2, 1024), (2, 8192), (4, 5120), (8, 5120)] +) +def test_allreduce_boundary_handles_batch_changes_and_graph_replay( + world_size, hidden_size +): + """Changing token counts must not consume stale Lamport slots or deadlock.""" + from torch.multiprocessing import spawn + + from vllm.utils.network_utils import get_open_port + + if torch.accelerator.device_count() < world_size: + pytest.skip(f"Requires {world_size} GPUs") + mnnvl = pytest.importorskip("flashinfer.comm.mnnvl") + if not mnnvl.is_multicast_supported(0): + pytest.skip("NVLink multicast required") + spawn( + _beam_allreduce_worker, + args=(world_size, get_open_port(), hidden_size), + nprocs=world_size, + ) + + +def _beam_finalize_worker(rank, world_size, port, expert_parallel, hidden_size): + import torch.distributed as dist + + from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput + from vllm.models.beam.nvidia.allreduce import ( + BeamAllReduceWorkspace, + beam_allreduce_norm_residual_norm, + beam_finalize_allreduce_norm_residual_norm, + ) + + torch.accelerator.set_device_index(rank) + dist.init_process_group( + "nccl", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=world_size, + ) + group = dist.new_group(backend="gloo") + n, k = hidden_size, 8 + workspace = BeamAllReduceWorkspace(group, n, 4095) + if expert_parallel: + cutoff = workspace.ep_finalize_oneshot_max_tokens + batches = [0, 1, cutoff, cutoff + 1, 513] + if hidden_size == 5120: + batches.append(4095) + batches.append(2) + else: + batches = [0, 1, 127, 513, 2] + for iteration, m in enumerate(batches): + torch.manual_seed(100 + iteration) + residual = torch.randn(m, n, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(n, device="cuda", dtype=torch.bfloat16) + owners = torch.randint(world_size, (m, k), device="cuda") + owners[::3] = torch.arange(m, device="cuda")[::3, None] % world_size + shared_weights = torch.rand(m, k, device="cuda") / k + torch.manual_seed(200 + rank + iteration) + experts = torch.randn(m * k, n, device="cuda", dtype=torch.bfloat16) + weights = ( + shared_weights if expert_parallel else torch.rand(m, k, device="cuda") / k + ) + if iteration % 2: + weights = weights.bfloat16() + permutation = torch.randperm(m * k, device="cuda", dtype=torch.int32).view(m, k) + if expert_parallel: + permutation.masked_fill_(owners != rank, -1) + else: + permutation[::3, -1] = -1 + x = UnfinalizedMoEOutput(experts, weights, permutation) + routed_scale = 2.5 if iteration % 2 else 1.0 + post = iteration % 3 != 0 + next_norm = iteration % 3 != 1 + scales = (0.125, 0.9682458365518543) + + args = ( + x, + residual, + weight if post else None, + 1e-6 if post else None, + weight if next_norm else None, + 1e-6 if next_norm else None, + scales, + workspace, + routed_scale, + ) + + def run(args=args): + return beam_finalize_allreduce_norm_residual_norm( + *args, + oneshot_max_tokens=( + workspace.ep_finalize_oneshot_max_tokens if expert_parallel else 0 + ), + ) + + def check(actual, args=args): + x, residual, pw, peps, nw, neps, scales, _, routed_scale = args + m, n = residual.shape + local = torch.zeros(m, n, device="cuda") + for expert in range(x.expert_weights.shape[1]): + rows = x.expanded_idx_to_permuted_idx[:, expert].long() + values = x.gemm2_permuted[rows.clamp_min(0)].float() + values = torch.where(rows[:, None] >= 0, values, 0) + local = local + values * x.expert_weights[:, expert, None].float() + local = (local.bfloat16().float() * routed_scale).bfloat16().float() + dist.all_reduce(local) + expected = _torch_boundary( + local.bfloat16(), + residual, + pw, + peps, + nw, + neps, + scales, + ) + _assert_boundary_close(actual, expected) + + check(run()) + if expert_parallel and m == workspace.ep_finalize_oneshot_max_tokens: + tight_workspace = BeamAllReduceWorkspace(group, n, m) + args = (*args[:-2], tight_workspace, routed_scale) + check(run(args), args) + torch.accelerator.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run(args) + for _ in range(2): + experts.mul_(0.5) + if expert_parallel: + weights.mul_(0.75) + permutation.copy_(permutation.roll(1, dims=0)) + graph.replay() + check(captured, args) + # Alternate two-stage finalization with the one-stage decode path + # on the same Lamport buffers, including during graph replay. + small = residual[:8] + actual = beam_allreduce_norm_residual_norm( + small, small, weight, 1e-6, weight, 1e-6, scales, args[-2] + ) + reduced = small.float() + dist.all_reduce(reduced) + expected = _torch_boundary( + reduced.bfloat16(), small, weight, 1e-6, weight, 1e-6, scales + ) + _assert_boundary_close(actual, expected) + graph.reset() + dist.barrier() + dist.destroy_process_group() + + +@pytest.mark.parametrize( + "world_size,expert_parallel,hidden_size", + [ + (2, False, 5120), + (4, False, 5120), + (4, True, 5120), + (2, True, 1024), + (4, True, 2048), + (4, True, 8192), + ], +) +def test_moe_finalize_allreduce_preserves_local_rounding_and_replay( + world_size, expert_parallel, hidden_size +): + """Empty local routes and live routing metadata retain BF16 casts on replay.""" + from torch.multiprocessing import spawn + + from vllm.utils.network_utils import get_open_port + + if torch.accelerator.device_count() < world_size: + pytest.skip(f"Requires {world_size} GPUs") + mnnvl = pytest.importorskip("flashinfer.comm.mnnvl") + if not mnnvl.is_multicast_supported(0): + pytest.skip("NVLink multicast required") + spawn( + _beam_finalize_worker, + args=(world_size, get_open_port(), expert_parallel, hidden_size), + nprocs=world_size, + ) + + +def _beam_moe_gather_worker(rank, world_size, port, hidden_size): + from datetime import timedelta + + import torch.distributed as dist + + from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput + from vllm.models.beam.nvidia.allreduce import ( + BeamAllReduceWorkspace, + beam_allreduce_norm_residual_norm, + ) + from vllm.models.beam.nvidia.moe_finalize import ( + beam_moe_finalize_gather_norm_residual_norm, + ) + + torch.accelerator.set_device_index(rank) + dist.init_process_group( + "nccl", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=120), + ) + group = dist.new_group(backend="gloo") + workspace = BeamAllReduceWorkspace(group, hidden_size, 513) + + def make_case(tokens, seed, *, post, next_norm, scales): + local_rows = (tokens + world_size - 1) // world_size + valid_rows = max(0, min(local_rows, tokens - rank * local_rows)) + torch.manual_seed(seed) + residual = torch.randn(tokens, hidden_size, device="cuda", dtype=torch.bfloat16) + affine = torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16) + torch.manual_seed(seed + rank + 1) + experts = torch.randn( + local_rows + 3, 8, hidden_size, device="cuda", dtype=torch.bfloat16 + ) + weights = torch.rand(local_rows + 3, 8, device="cuda") + permutation = torch.arange( + (local_rows + 3) * 8, device="cuda", dtype=torch.int32 + ).view(local_rows + 3, 8) + experts[:valid_rows:2].mul_(1e-3) + experts[:valid_rows, :, ::97] = -0.0 + experts[:valid_rows, :, 1::97] = -1e-40 + # Both capacity padding and SP padding must stay outside live outputs. + experts[valid_rows:] = float("nan") + weights[valid_rows:] = float("nan") + return ( + UnfinalizedMoEOutput(experts.flatten(0, 1), weights, permutation), + residual, + affine if post == "weighted" else None, + None if post == "none" else 1e-6, + affine if next_norm == "weighted" else None, + None if next_norm == "none" else 1e-6, + scales, + workspace, + 0.137, + ) + + def check(actual, args): + output, residual, pw, peps, nw, neps, scales, _, scale = args + tokens = residual.shape[0] + local_rows = (tokens + world_size - 1) // world_size + valid_rows = max(0, min(local_rows, tokens - rank * local_rows)) + weights = output.expert_weights[:valid_rows] + experts = output.gemm2_permuted.view(-1, 8, hidden_size)[:valid_rows] + local = torch.zeros( + local_rows, hidden_size, dtype=residual.dtype, device=residual.device + ) + local[:valid_rows] = _torch_megamoe_finalize(experts, weights, scale, pw, peps) + gathered = torch.empty( + world_size * local_rows, + hidden_size, + dtype=local.dtype, + device=local.device, + ) + if tokens: + dist.all_gather_into_tensor(gathered, local) + expected = _torch_boundary( + gathered[:tokens], residual, None, None, nw, neps, scales + ) + _assert_boundary_close(actual, expected, relative_tolerance=1e-3) + + batches = (0, 1, world_size + 1, 513) + norms = ("none", "unweighted", "weighted") + for iteration, tokens in enumerate(batches): + args = make_case( + tokens, + 200 + iteration, + post=norms[iteration % 3], + next_norm=norms[(iteration + 1) % 3], + scales=(0.125, 0.9682458365518543) if iteration % 2 else None, + ) + for strategy in ("unfused", "oneshot", "twoshot", None): + check( + beam_moe_finalize_gather_norm_residual_norm(*args, strategy=strategy), + args, + ) + + # Independent captured graphs share the Lamport generation state. Replaying + # one then two stages must work after eager collectives and changed inputs. + graphs = [] + for strategy, tokens in ( + ("oneshot", world_size - 1), + ("oneshot", 17), + ("twoshot", world_size - 1), + ("twoshot", 129), + ): + args = make_case( + tokens, + 300 + tokens, + post="weighted", + next_norm="weighted", + scales=(0.125, 0.9682458365518543), + ) + beam_moe_finalize_gather_norm_residual_norm(*args, strategy=strategy) + torch.accelerator.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = beam_moe_finalize_gather_norm_residual_norm( + *args, strategy=strategy + ) + graphs.append((graph, actual, args)) + # Without a reference collective between replays, empty ranks may lag more + # than a Lamport generation. Padded rows must still acknowledge each rank. + for graph, actual, args in graphs: + if rank == world_size - 1: + torch.cuda._sleep(2_000_000) + for _ in range(16): + graph.replay() + check(actual, args) + for iteration in range(2): + for graph, actual, args in graphs[:: 1 if iteration % 2 else -1]: + output = args[0] + output.gemm2_permuted.mul_(0.75) + output.expert_weights.copy_(output.expert_weights.roll(1, dims=1)) + args[1].mul_(0.5) + args[2].neg_() + graph.replay() + check(actual, args) + # Attention AR and MoE gather reuse this workspace in production. + residual, weight = args[1], args[2] + x = residual * (rank + 1) + ar_args = (x, residual, weight, 1e-6, weight, 1e-6, None, workspace) + actual_ar = beam_allreduce_norm_residual_norm(*ar_args) + reduced = x.float() + dist.all_reduce(reduced) + _assert_boundary_close( + actual_ar, + _torch_boundary( + reduced.bfloat16(), residual, weight, 1e-6, weight, 1e-6, None + ), + relative_tolerance=1e-3, + ) + torch.accelerator.synchronize() + for graph, _, _ in graphs: + graph.reset() + dist.barrier() + dist.destroy_process_group() + + +@pytest.mark.parametrize("world_size,hidden_size", [(2, 8192), (4, 5120)]) +def test_megamoe_gather_restores_tokens_and_alternates_graph_strategies( + world_size, hidden_size +): + """Keep SP rank order, BF16 casts, and generation state across all strategies.""" + from torch.multiprocessing import spawn + + from vllm.utils.network_utils import get_open_port + + if torch.accelerator.device_count() < world_size: + pytest.skip(f"Requires {world_size} GPUs") + mnnvl = pytest.importorskip("flashinfer.comm.mnnvl") + if not mnnvl.is_multicast_supported(0): + pytest.skip("NVLink multicast required") + spawn( + _beam_moe_gather_worker, + args=(world_size, get_open_port(), hidden_size), + nprocs=world_size, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("rows", [1, 2, 4, 8, 12, 16, 24, 64]) +@pytest.mark.parametrize("qkv_rows,gate_rows", [(2048, 1536), (1024, 768)]) +@torch.inference_mode() +def test_packed_attention_projection(rows, qkv_rows, gate_rows): + """The shared GEMMs preserve QKV/gate ordering and BF16 output accuracy.""" + from vllm.models.beam.nvidia.projection import ( + packed_projection, + use_low_latency_gemm, + ) + + torch.manual_seed(42) + x = torch.randn(rows, 5120, device="cuda", dtype=torch.bfloat16) + w = torch.randn(qkv_rows + gate_rows, 5120, device="cuda", dtype=x.dtype) * 0.01 + packed = packed_projection(x, w, use_low_latency_gemm(w)) + for actual, weight in zip( + packed.split([qkv_rows, gate_rows], -1), + w.split([qkv_rows, gate_rows]), + ): + expected = (x.double() @ weight.double().T).to(x.dtype) + torch.testing.assert_close(actual, expected, rtol=0.008, atol=1e-5) diff --git a/tests/models/test_beam_weight_mapping.py b/tests/models/test_beam_weight_mapping.py new file mode 100644 index 000000000000..0f5feb7c480e --- /dev/null +++ b/tests/models/test_beam_weight_mapping.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""`BeamModel.hf_to_vllm_mapper` must fuse only the dense projections. + +Beam's checkpoints carry two `gate_proj` names that are *not* the SwiGLU +gate -- the per-expert `mlp.experts.N.gate_proj` (owned by +`MoERunner.load_weights`) and the attention output gate `self_attn.gate_proj`. +Folding either into `mlp.gate_up_proj` silently corrupts the weights, so pin the +mapping. The attention gate is additionally renamed to `output_gate_proj` so +that `packed_modules_mapping` consumers (LoRA) cannot match it as a +`gate_up_proj` shard either. +""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm.models.beam.nvidia import model as beam +from vllm.models.beam.nvidia.model import BeamModel +from vllm.transformers_utils.configs.beam import BeamConfig + +MAPPER = BeamModel.hf_to_vllm_mapper + + +def _map(name: str) -> tuple[str, int | str | None]: + weight = torch.empty(0) + ((mapped, out),) = MAPPER.apply([(name, weight)]) + return mapped, getattr(out, "shard_id", None) + + +@pytest.mark.parametrize( + ("checkpoint_name", "expected_name", "expected_shard"), + [ + ( + "model.layers.3.self_attn.q_proj.weight", + "model.layers.3.self_attn.qkv_proj.weight", + "q", + ), + ( + "model.layers.3.self_attn.k_proj.weight", + "model.layers.3.self_attn.qkv_proj.weight", + "k", + ), + ( + "model.layers.3.self_attn.v_proj.weight", + "model.layers.3.self_attn.qkv_proj.weight", + "v", + ), + ( + "model.layers.0.mlp.gate_proj.weight", + "model.layers.0.mlp.gate_up_proj.weight", + 0, + ), + ( + "model.layers.0.mlp.up_proj.weight", + "model.layers.0.mlp.gate_up_proj.weight", + 1, + ), + ], +) +def test_dense_projections_are_fused(checkpoint_name, expected_name, expected_shard): + assert _map(checkpoint_name) == (expected_name, expected_shard) + + +@pytest.mark.parametrize( + "checkpoint_name", + [ + # Owned by MoERunner.load_weights, which is also what keeps expert + # weights correct under EP and redundant experts. + "model.layers.1.mlp.experts.7.gate_proj.weight", + "model.layers.1.mlp.experts.7.up_proj.weight", + "model.layers.1.mlp.experts.7.down_proj.weight", + "model.layers.1.self_attn.gate_bias", + # The MoE router. + "model.layers.1.mlp.gate.weight", + "model.layers.1.mlp.gate.e_score_correction_bias", + ], +) +def test_gates_and_experts_pass_through(checkpoint_name): + assert _map(checkpoint_name) == (checkpoint_name, None) + + +def test_attention_output_gate_is_renamed_not_fused(): + """The attention gate must be renamed, and must not pick up a shard id.""" + assert _map("model.layers.1.self_attn.gate_proj.weight") == ( + "model.layers.1.self_attn.output_gate_proj.weight", + None, + ) + + +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize( + "module_path,name", + [("self_attn", "gate_bias"), ("mlp.gate", "e_score_correction_bias")], +) +def test_optional_bias_loads_only_when_present(module_path, name, enabled): + """Unknown Beam biases must fail instead of silently changing model output.""" + model = BeamModel.__new__(BeamModel) + nn.Module.__init__(model) + layer = nn.Module() + layer.self_attn = nn.Module() + layer.mlp = nn.Module() + layer.mlp.gate = nn.Module() + model.layers = nn.ModuleList([layer]) + target = layer.get_submodule(module_path) + target.register_parameter(name, nn.Parameter(torch.zeros(2)) if enabled else None) + weight_name = f"layers.0.{module_path}.{name}" + weight = torch.tensor([1.0, 2.0]) + weights = [(weight_name, weight)] + if enabled: + assert model.load_weights(weights) == {weight_name} + torch.testing.assert_close(getattr(target, name), weight) + else: + with pytest.raises(ValueError, match=name): + model.load_weights(weights) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@torch.inference_mode() +def test_checkpoint_norm_names_match_sublayer_order(default_vllm_config, monkeypatch): + """Load reference names directly and use each weight on its intended branch.""" + + class Attention(nn.Module): + def forward(self, positions, hidden_states): + return hidden_states + + monkeypatch.setattr(beam, "BeamAttention", lambda *args, **kwargs: Attention()) + monkeypatch.setattr(beam, "BeamMLP", lambda **kwargs: nn.Identity()) + config = BeamConfig( + hidden_size=8, + mlp_only_layers=[0], + enable_sandwich_norm=True, + self_attn_depth_scaled_residual=False, + mlp_depth_scaled_residual=False, + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=config), + cache_config=None, + quant_config=None, + ) + model = BeamModel.__new__(BeamModel) + nn.Module.__init__(model) + model.config = config + layer = beam.BeamDecoderLayer(vllm_config, prefix="layers.0") + model.layers = nn.ModuleList([layer]) + model.cuda() + names = [ + "input_layernorm", + "post_attention_layernorm", + "pre_feedforward_layernorm", + "post_feedforward_layernorm", + ] + weights = { + f"layers.0.{name}.weight": torch.arange(1, 9, device="cuda").float().roll(i) / 8 + for i, name in enumerate(names) + } + assert model.load_weights(weights.items()) == set(weights) + x = torch.arange(1, 17, device="cuda").float().view(2, 8) / 16 + + def norm(value, name): + return ( + value + * torch.rsqrt(value.square().mean(-1, keepdim=True) + 1e-6) + * weights[f"layers.0.{name}.weight"] + ) + + expected = x + norm(norm(x, names[0]), names[1]) + expected = expected + norm(norm(expected, names[2]), names[3]) + actual, _ = layer(torch.arange(2, device="cuda"), x, None) + torch.testing.assert_close(actual, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("tp_size", [4, 8]) +@torch.inference_mode() +def test_checkpoint_loads_into_merged_qkv_gate( + default_vllm_config, monkeypatch, tp_size +): + """All four checkpoint projections load directly into the allocated TP shard.""" + from vllm.models.beam.nvidia import attention + from vllm.models.beam.nvidia.projection import packed_projection + + for module in ( + "vllm.model_executor.layers.linear", + "vllm.model_executor.parameter", + ): + monkeypatch.setattr( + f"{module}.get_tensor_model_parallel_rank", lambda: tp_size - 1 + ) + monkeypatch.setattr( + f"{module}.get_tensor_model_parallel_world_size", lambda: tp_size + ) + monkeypatch.setattr( + attention, "get_tensor_model_parallel_world_size", lambda: tp_size + ) + attention_stub = nn.Identity() + attention_stub.attn_backend = SimpleNamespace(get_name=lambda: "TEST") + monkeypatch.setattr(attention, "Attention", lambda *args, **kwargs: attention_stub) + config = BeamConfig( + hidden_size=64, num_attention_heads=48, num_key_value_heads=8, head_dim=8 + ) + attn = ( + attention.BeamAttention(config, None, None, prefix="layers.0.self_attn") + .bfloat16() + .cuda() + ) + model = BeamModel.__new__(BeamModel) + nn.Module.__init__(model) + model.config = config + model.hf_to_vllm_mapper = model.get_weights_mapper(config, tp_size) + layer = nn.Module() + layer.self_attn = attn + model.layers = nn.ModuleList([layer]) + param = attn.qkv_gate_proj.weight + initial_pointer = param.data_ptr() + for _ in range(2): + weights = [ + torch.randn(n, 64, device="cuda", dtype=torch.bfloat16) + for n in (384, 64, 64, 384) + ] + loaded = model.load_weights( + (f"layers.0.self_attn.{name}.weight", w) + for name, w in zip(("q_proj", "k_proj", "v_proj", "gate_proj"), weights) + ) + assert loaded == {"layers.0.self_attn.qkv_gate_proj.weight"} + assert param.data_ptr() == initial_pointer + expected_weight = torch.cat([w.chunk(tp_size)[-1] for w in weights]) + torch.testing.assert_close(param, expected_weight, rtol=0, atol=0) + x = torch.randn(3, 64, device="cuda", dtype=torch.bfloat16) + torch.testing.assert_close( + packed_projection(x, param, False), + torch.nn.functional.linear(x, expected_weight), + rtol=0, + atol=0, + ) diff --git a/tests/parser/test_beam_parser.py b/tests/parser/test_beam_parser.py new file mode 100644 index 000000000000..3673b360908e --- /dev/null +++ b/tests/parser/test_beam_parser.py @@ -0,0 +1,747 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +from copy import deepcopy +from types import SimpleNamespace +from typing import Any + +import pytest +from openai.types.responses import FunctionTool, ToolChoiceFunction +from xgrammar import Grammar +from xgrammar.testing import _is_grammar_accept_string + +from vllm.config.reasoning import ReasoningConfig +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ChatCompletionToolsParam, +) +from vllm.entrypoints.openai.responses.protocol import ( + ResponsesRequest, + ResponsesResponse, +) +from vllm.exceptions import VLLMValidationError +from vllm.parser.beam import BeamParser, BeamReasoningParser, BeamToolParser +from vllm.parser.parser_manager import ParserManager +from vllm.parser.response_template import ResponseTemplateParser +from vllm.reasoning import ReasoningParserManager +from vllm.tool_parsers import ToolParserManager + +# The `response_template` of the Beam checkpoints. +BEAM_RESPONSE_TEMPLATE: dict[str, Any] = { + "defaults": {"role": "assistant"}, + "start_anchor_pattern": r"<\|start\|>assistant", + "fields": { + "thinking": { + "open": "<|channel|>analysis<|message|>", + "close": "<|end|><|start|>assistant", + "content": "text", + }, + "content": { + "open": "<|channel|>final<|message|>", + "close": ["<|return|>", "<|end|>", "<|call|>"], + "content": "text", + }, + "tool_calls": { + "open_pattern": ( + r" to=(?P[^\s<]+)" + r"<\|channel\|>analysis <\|constrain\|>xml<\|message\|>" + ), + "close": ["<|end|><|start|>assistant", "<|call|>"], + "repeats": True, + "content": "xml-inline", + "content_args": { + "tag_pattern": ( + r'' + r'\n(?P.*?)\n(?:|(?=": 200, + "<|channel|>": 201, + "<|message|>": 202, + "<|end|>": 203, + "<|call|>": 204, + "<|return|>": 205, + "<|constrain|>": 206, +} +TEXT_TOKENS = {100: "assistant", 101: "analysis", 102: "final", 103: "plan"} +PROMPT_TOKEN_IDS = [200, 100] +THOUGHT = "<|channel|>analysis<|message|>plan<|end|><|start|>assistant" +FINAL = "<|channel|>final<|message|>" +NEXT_MESSAGE = "<|end|><|start|>assistant" + + +def _tool(name, properties, *, strict=None): + return ChatCompletionToolsParam( + type="function", + function={ + "name": name, + "parameters": {"type": "object", "properties": properties}, + "strict": strict, + }, + ) + + +TOOLS = [ + _tool("get_weather", {"city": {"type": "string"}}), + _tool("set_count", {"count": {"type": "integer"}}), +] +STRICT_TOOLS = [ + _tool("get_weather", {"city": {"type": "string"}}), + _tool("set_count", {"count": {"type": "integer"}}, strict=True), +] + + +class BeamTokenizer: + def __init__(self, response_template=BEAM_RESPONSE_TEMPLATE): + self.response_template = response_template + self.init_kwargs = {} + self.name_or_path = "" + self.all_special_tokens = list(VOCAB) + self.all_special_ids = list(VOCAB.values()) + + def decode(self, token_ids, skip_special_tokens=False, **_kwargs): + inverse_vocab = {token_id: token for token, token_id in VOCAB.items()} + return "".join( + TEXT_TOKENS.get(token_id, inverse_vocab.get(token_id, "")) + for token_id in token_ids + ) + + def get_vocab(self): + return VOCAB + + def encode(self, text, add_special_tokens=False): + del add_special_tokens + return list(text.encode()) + + +def request(**overrides): + values = dict( + tool_choice="auto", + include_reasoning=True, + tools=TOOLS, + parallel_tool_calls=True, + ) + values.update(overrides) + return SimpleNamespace(**values) + + +def parameter(name, value): + return f'\n{value}\n' + + +def call(name, body, *, terminator="<|call|>"): + return ( + f" to={name}<|channel|>analysis <|constrain|>xml<|message|>{body}{terminator}" + ) + + +def make_parser(prompt_token_ids=PROMPT_TOKEN_IDS, tools=TOOLS, **kwargs): + parser_cls = ParserManager.get_parser( + tool_parser_name="beam", + reasoning_parser_name="beam", + enable_auto_tools=True, + tokenizer=BeamTokenizer(), + ) + assert parser_cls is not None + parser = parser_cls(BeamTokenizer(), tools, **kwargs) + parser.set_prompt_token_ids(prompt_token_ids) + return parser + + +def parse(text, prompt_token_ids=PROMPT_TOKEN_IDS): + reasoning, content, calls = make_parser(prompt_token_ids).parse( + text, request(), enable_auto_tools=True + ) + return reasoning, content, normalize_calls(calls) + + +def normalize_calls(calls): + return [(call.name, json.loads(call.arguments)) for call in calls or []] + + +def test_beam_parsers_are_registered(): + assert ToolParserManager.get_tool_parser("beam") is BeamToolParser + assert ReasoningParserManager.get_reasoning_parser("beam") is BeamReasoningParser + + +def test_parser_manager_selects_beam_parser(): + parser_cls = ParserManager.get_parser( + tool_parser_name="beam", + reasoning_parser_name="beam", + enable_auto_tools=True, + tokenizer=BeamTokenizer(), + ) + + assert parser_cls is not None + assert issubclass(parser_cls, BeamParser) + assert parser_cls.reasoning_parser_cls is BeamReasoningParser + assert parser_cls.tool_parser_cls is BeamToolParser + assert parser_cls._enable_auto_tools + + +@pytest.mark.parametrize("other", ["hermes", "hf"]) +def test_beam_parser_rejects_other_parsers(other): + with pytest.raises(TypeError, match="cannot be combined"): + ParserManager.get_parser( + tool_parser_name=other, + reasoning_parser_name="beam", + enable_auto_tools=True, + tokenizer=BeamTokenizer(), + ) + + +def test_beam_parser_requires_checkpoint_metadata_at_startup(): + with pytest.raises( + TypeError, match="--reasoning-parser beam requires `response_template`" + ): + ParserManager.get_parser( + reasoning_parser_name="beam", + tokenizer=BeamTokenizer(response_template=None), + ) + + template = deepcopy(BEAM_RESPONSE_TEMPLATE) + del template["fields"]["tool_calls"] + with pytest.raises(TypeError, match="--tool-call-parser beam requires"): + ParserManager.get_parser( + tool_parser_name="beam", + enable_auto_tools=True, + tokenizer=BeamTokenizer(template), + ) + + +def test_reasoning_and_final_answer(): + assert parse(THOUGHT + FINAL + "Sunny.<|return|>") == ("plan", "Sunny.", []) + + +def test_prefilled_reasoning_continues_from_prompt(): + prompt = PROMPT_TOKEN_IDS + [201, 101, 202, 103] + text = " more.<|end|><|start|>assistant" + FINAL + "Done.<|return|>" + + assert parse(text, prompt) == (" more.", "Done.", []) + + +def test_parallel_calls_are_coerced_to_their_schemas(): + text = ( + THOUGHT + + call( + "get_weather", + parameter("city", "Paris"), + terminator="<|end|><|start|>assistant", + ) + + call("set_count", parameter("count", "7")) + ) + + assert parse(text) == ( + "plan", + None, + [("get_weather", {"city": "Paris"}), ("set_count", {"count": 7})], + ) + + +def test_call_stop_token_removed_from_the_text_still_ends_the_call(): + text = THOUGHT + call("get_weather", parameter("city", "Paris"), terminator="") + + assert parse(text) == ("plan", None, [("get_weather", {"city": "Paris"})]) + + +@pytest.mark.parametrize("end", ["<|end|>", "<|end|><|start|>"]) +def test_output_cut_before_the_next_message_start_ends_its_region(end): + """Output that stops after `<|end|>`, such as at `max_tokens`, has ended + its reasoning or call and lacks only the next message's start.""" + weather = call("get_weather", parameter("city", "Paris"), terminator=end) + + assert parse("<|channel|>analysis<|message|>plan" + end) == ("plan", None, []) + assert parse(THOUGHT + weather) == ( + "plan", + None, + [("get_weather", {"city": "Paris"})], + ) + + +THOUGHT_IDS = [201, 101, 202, 103, 203, 200, 100] +FINAL_IDS = [201, 102, 202] + + +@pytest.mark.parametrize( + ("prompt_token_ids", "token_ids", "count"), + [ + (PROMPT_TOKEN_IDS, [*THOUGHT_IDS, *FINAL_IDS, 205], 1), + (PROMPT_TOKEN_IDS + [201, 101, 202], [103, 103, 203, 200, 100, *FINAL_IDS], 2), + (PROMPT_TOKEN_IDS, [201, 101, 202, 103, 103, 103], 3), + ], + ids=["reasoning_then_answer", "prefilled_reasoning", "cut_off"], +) +def test_reasoning_tokens_are_counted(prompt_token_ids, token_ids, count): + parser = make_parser(prompt_token_ids) + + assert parser.count_reasoning_tokens(token_ids) == count + + +def test_reasoning_tokens_are_counted_as_the_generation_grows(): + token_ids = [*THOUGHT_IDS, *FINAL_IDS, 205] + parser = make_parser() + + assert [ + parser.count_reasoning_tokens(token_ids[:end]) + for end in range(1, len(token_ids) + 1) + ] == [ + make_parser().count_reasoning_tokens(token_ids[:end]) + for end in range(1, len(token_ids) + 1) + ] + assert parser.count_reasoning_tokens(FINAL_IDS) == 0 + + +def test_call_stop_token_ends_the_answer(): + assert parse(FINAL + "done<|call|>") == (None, "done", []) + + +@pytest.mark.parametrize("chunk_size", [1, 7, 10_000]) +def test_streaming_matches_non_streaming(chunk_size): + text = ( + THOUGHT + + call( + "get_weather", + parameter("city", "Paris"), + terminator="<|end|><|start|>assistant", + ) + + call("set_count", parameter("count", "7")) + ) + parser = make_parser() + reasoning, calls = [], [] + for start in range(0, len(text), chunk_size): + end = min(start + chunk_size, len(text)) + delta = parser.parse_delta( + text[start:end], [], request(), finished=end == len(text) + ) + if delta is not None: + reasoning.append(delta.reasoning or "") + calls.extend(delta.tool_calls) + + assert "".join(reasoning) == "plan" + assert [ + (call.function.name, json.loads(call.function.arguments)) for call in calls + ] == parse(text)[2] + + +WEATHER_CALL = call("get_weather", parameter("city", "Paris")) +WEATHER = [("get_weather", {"city": "Paris"})] + + +def chat_request(tools=STRICT_TOOLS, **body): + return ChatCompletionRequest( + model="beam", + messages=[{"role": "user", "content": "Weather?"}], + tools=tools and [tool.model_dump() for tool in tools], + **body, + ) + + +def named(name): + return {"type": "function", "function": {"name": name}} + + +def adjust(req, prompt_token_ids=PROMPT_TOKEN_IDS, **kwargs): + make_parser(prompt_token_ids, req.tools, **kwargs).adjust_request(req) + return req + + +def accepts(req, text): + grammar = Grammar.from_structural_tag( + req.extract_structured_outputs().structural_tag + ) + return _is_grammar_accept_string(grammar, text) + + +@pytest.mark.parametrize( + ("tool_choice", "text", "calls"), + [ + ("required", THOUGHT + WEATHER_CALL, WEATHER), + ("required", WEATHER_CALL, None), + ( + "required", + THOUGHT + + call("get_weather", parameter("city", "Paris"), terminator=NEXT_MESSAGE) + + call("set_count", parameter("count", "7")), + [*WEATHER, ("set_count", {"count": 7})], + ), + ("required", THOUGHT + call("set_count", parameter("count", "x")), None), + ( + "required", + THOUGHT + call("get_weather", parameter("city", "a<|call|>b")), + None, + ), + ("required", THOUGHT + FINAL + "Paris<|return|>", None), + (named("set_count"), THOUGHT + WEATHER_CALL, None), + (named("get_weather"), THOUGHT + WEATHER_CALL, WEATHER), + ( + named("get_weather"), + THOUGHT + + call("get_weather", parameter("city", "Paris"), terminator=NEXT_MESSAGE) + + WEATHER_CALL, + None, + ), + ], +) +def test_forced_tool_calls_follow_the_template_grammar(tool_choice, text, calls): + """From the message start, the grammar opens with the reasoning, and + accepts only calls the parser reads back.""" + req = adjust(chat_request(tool_choice=tool_choice)) + + assert req._grammar_from_parser + assert accepts(req, text) is (calls is not None) + if calls is not None: + assert parse(text)[2] == calls + + +def test_reasoning_may_be_skipped_when_thinking_is_disabled(): + req = adjust( + chat_request(tool_choice="required"), + chat_template_kwargs={"enable_thinking": False}, + ) + + assert accepts(req, WEATHER_CALL) + assert accepts(req, THOUGHT + WEATHER_CALL) + + +def test_parallel_tool_calls_false_allows_a_single_call(): + req = adjust(chat_request(tool_choice="required", parallel_tool_calls=False)) + calls = call("get_weather", parameter("city", "Paris"), terminator=NEXT_MESSAGE) + + assert accepts(req, THOUGHT + WEATHER_CALL) + assert not accepts(req, THOUGHT + calls + WEATHER_CALL) + + +@pytest.mark.parametrize("strict", [False, True]) +@pytest.mark.parametrize( + ("separator", "accepted"), [("", True), ("\n", True), ("\n\n", False), (" ", False)] +) +def test_parameters_are_separated_by_at_most_a_newline(strict, separator, accepted): + """A model that skips a parameter's opener must not write whitespace until + `max_tokens`.""" + write = _tool( + "write", + {"path": {"type": "string"}, "content": {"type": "string"}}, + strict=strict, + ) + req = adjust(chat_request([write], tool_choice="required")) + body = parameter("path", "a.py") + separator + parameter("content", "x") + + assert accepts(req, THOUGHT + call("write", body)) is accepted + assert accepts(req, THOUGHT + call("write", separator + body)) is accepted + + +def test_forced_calls_in_prefilled_reasoning_start_after_it(): + req = adjust( + chat_request(tool_choice="required"), PROMPT_TOKEN_IDS + [201, 101, 202, 103] + ) + + assert not req._grammar_from_parser + assert accepts(req, WEATHER_CALL) + + +def test_forced_calls_cannot_follow_a_prefilled_answer(): + with pytest.raises(VLLMValidationError, match="required tool choice"): + adjust(chat_request(tool_choice="required"), PROMPT_TOKEN_IDS + [201, 102, 202]) + + +def test_prefilled_answer_can_only_be_finished(): + req = adjust(chat_request(TOOLS), PROMPT_TOKEN_IDS + [201, 102, 202]) + + assert accepts(req, "Sunny.<|return|>") + assert not accepts(req, "Sunny." + NEXT_MESSAGE + WEATHER_CALL) + + +def test_forced_call_with_unencodable_tool_name_is_rejected(): + tools = [_tool("get weather", {"city": {"type": "string"}})] + + with pytest.raises(VLLMValidationError, match="tool name"): + adjust(chat_request(tools, tool_choice="required")) + + +def test_string_formats_and_patterns_are_held_in_raw_text(): + """Like XGrammar, a strict string keeps its format or pattern, unless a + match could end the value early.""" + contact = _tool( + "set_contact", + { + "email": {"type": "string", "format": "email"}, + "code": {"type": "string", "pattern": "^[A-Z]{3}$"}, + "n": {"type": "integer"}, + }, + strict=True, + ) + contact.function.parameters["required"] = ["email", "code", "n"] + req = adjust(chat_request([contact], tool_choice="required")) + + def contact_call(email, code, *rest): + body = parameter("email", email) + parameter("code", code) + "".join(rest) + return THOUGHT + call("set_contact", body) + + text = contact_call("a.b@x.co", "ABC", parameter("n", "1")) + calls = make_parser(tools=[contact]).parse(text, req, enable_auto_tools=True)[2] + assert accepts(req, text) + assert normalize_calls(calls) == [ + ("set_contact", {"email": "a.b@x.co", "code": "ABC", "n": 1}) + ] + assert not accepts(req, contact_call("ab", "ABC", parameter("n", "1"))) + assert not accepts(req, contact_call("a.b@x.co", "ABCD", parameter("n", "1"))) + assert not accepts(req, contact_call("a.b@x.co", "ABC")) + # Line breaks could end the value. + dotted = _tool( + "set_code", {"code": {"type": "string", "pattern": "^a.b$"}}, strict=True + ) + with pytest.raises(VLLMValidationError, match="required tool choice"): + adjust(chat_request([dotted], tool_choice="required")) + + +@pytest.mark.parametrize( + ("text", "accepted"), + [ + (THOUGHT + WEATHER_CALL, True), + (THOUGHT + FINAL + "Sunny.<|return|>", True), + (THOUGHT + FINAL + "Sunny." + NEXT_MESSAGE + WEATHER_CALL, False), + (THOUGHT + call("set_count", parameter("count", "x")), False), + ], +) +def test_strict_tools_constrain_automatic_tool_choice(text, accepted): + assert accepts(adjust(chat_request()), text) is accepted + + +@pytest.mark.parametrize( + ("text", "accepted"), + [ + (THOUGHT + FINAL + "Sunny.<|return|>", True), + (FINAL + "Sunny.<|return|>", False), + (WEATHER_CALL, False), + (THOUGHT + WEATHER_CALL, False), + (FINAL + "Sunny." + NEXT_MESSAGE + WEATHER_CALL, False), + ], +) +def test_tool_choice_none_allows_only_an_answer(text, accepted): + req = adjust(chat_request(TOOLS, tool_choice="none")) + + assert req._grammar_from_parser + assert accepts(req, text) is accepted + + +@pytest.mark.parametrize(("tools", "tool_choice"), [(TOOLS, "auto"), (None, None)]) +def test_always_constrained_output_follows_the_template_framing(tools, tool_choice): + req = adjust(chat_request(tools, tool_choice=tool_choice)) + count_call = THOUGHT + call("set_count", parameter("count", "3")) + + assert accepts(req, THOUGHT + FINAL + "Sunny.<|return|>") + assert not accepts(req, FINAL + "Sunny.<|return|>") + assert not accepts(req, THOUGHT.replace("plan", "") + FINAL + "Sunny.<|return|>") + assert not accepts(req, THOUGHT + "Sunny.") + assert accepts(req, count_call) is (tool_choice == "auto") + + +@pytest.mark.parametrize("response_format", [None, {"type": "json_object"}]) +@pytest.mark.parametrize("strict", [None, True]) +@pytest.mark.parametrize("tool_choice", ["auto", "required"]) +def test_automatic_choice_lets_calls_name_undeclared_tools( + strict, tool_choice, response_format +): + """Unless a forced choice limits calls to the declared tools. The parser + passes such calls on, so the harness can tell the model the tool does not + exist.""" + bash = _tool("bash", {"command": {"type": "string"}}, strict=strict) + req = adjust( + chat_request([bash], tool_choice=tool_choice, response_format=response_format) + ) + undeclared = THOUGHT + call("read_file", parameter("path", "a.py")) + parser = make_parser(tools=[bash]) + + assert accepts(req, THOUGHT + call("bash", parameter("command", "ls"))) + assert accepts(req, undeclared) is (tool_choice == "auto") + calls = parser.parse(undeclared, req, enable_auto_tools=True)[2] + assert normalize_calls(calls) == [("read_file", {"path": "a.py"})] + + +@pytest.mark.parametrize( + ("name", "body", "accepted"), + [ + ("set_count", parameter("count", "1"), True), + ("set_count", parameter("count", "x"), False), + ("get_weather", parameter("town", "Paris"), False), + # The grammar cannot enforce `minProperties`. + ("find", parameter("x", "1") + parameter("x", "y"), True), + ("search_files", parameter("x", "1") + parameter("x", "y"), True), + ("set_coun", parameter("x", "1") + parameter("x", "y"), True), + ("set_counts", parameter("x", "1") + parameter("x", "y"), True), + ], +) +def test_automatic_calls_enforce_declared_arguments_beside_any_other_name( + name, body, accepted +): + find = _tool("find", {"q": {"type": "string"}}) + find.function.parameters["minProperties"] = 1 + req = adjust(chat_request([*TOOLS, find])) + + assert accepts(req, THOUGHT + call(name, body)) is accepted + + +def test_value_ends_before_an_opener_on_the_next_line(): + write = _tool("write", {"path": {"type": "string"}, "content": {"type": "string"}}) + req = adjust(chat_request([write], tool_choice="required")) + body = '\nprint(1)\n\n' + parameter("path", "a.py") + text = THOUGHT + call("write", body) + + assert accepts(req, text) + calls = make_parser(tools=[write]).parse(text, req, enable_auto_tools=True)[2] + assert normalize_calls(calls) == [ + ("write", {"content": "print(1)\n", "path": "a.py"}) + ] + + +def test_always_constrained_output_can_be_disabled(monkeypatch): + monkeypatch.setenv("VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT", "0") + + assert adjust(chat_request(TOOLS)).structured_outputs is None + + +def test_strict_tools_the_template_cannot_hold_leave_automatic_choice_unconstrained(): + dotted = _tool( + "set_code", {"code": {"type": "string", "pattern": "^a.b$"}}, strict=True + ) + + assert adjust(chat_request([dotted])).structured_outputs is None + + +def test_generic_parser_does_not_derive_the_template_grammar(): + req = chat_request(TOOLS, tool_choice="required") + parser = ResponseTemplateParser(BeamTokenizer(), req.tools, enable_auto_tools=True) + + with pytest.raises(VLLMValidationError, match="required tool choice"): + parser.adjust_request(req) + + +@pytest.mark.parametrize("tools", [None, TOOLS]) +def test_structured_output_is_framed_as_the_answer(tools): + schema = {"type": "object", "properties": {"a": {"type": "integer"}}} + req = adjust( + chat_request( + tools, + response_format={ + "type": "json_schema", + "json_schema": {"name": "answer", "schema": schema}, + }, + ) + ) + + assert req.response_format is None + assert accepts(req, THOUGHT + FINAL + '{"a": 1}<|return|>') + assert not accepts(req, THOUGHT + FINAL + "Sunny.<|return|>") + assert accepts(req, THOUGHT + WEATHER_CALL) is (tools is not None) + + +def test_responses_named_tool_choice_follows_the_template_grammar(): + req = adjust( + ResponsesRequest( + input="Weather?", + tools=[ + FunctionTool( + type="function", + name="get_weather", + parameters={ + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + strict=True, + ) + ], + tool_choice=ToolChoiceFunction(type="function", name="get_weather"), + text={"format": {"type": "json_object"}}, + ) + ) + + assert req._grammar_from_parser + assert req.text.format.type == "json_object" + assert accepts(req, THOUGHT + WEATHER_CALL) + assert not accepts(req, FINAL + "{}<|return|>") + + +@pytest.mark.parametrize("status", ["in_progress", "completed"]) +@pytest.mark.parametrize("inside_reasoning", [False, True]) +def test_responses_preserves_text_config_with_template_grammar( + status, inside_reasoning +): + text = { + "format": { + "type": "json_schema", + "name": "answer", + "strict": True, + "schema": { + "type": "object", + "properties": {"a": {"type": "integer"}}, + "required": ["a"], + "additionalProperties": False, + }, + }, + "verbosity": "low", + } + req = ResponsesRequest(input="Answer", text=text) + expected_text = req.text.model_dump(by_alias=True) + prompt = PROMPT_TOKEN_IDS + ([201, 101, 202] if inside_reasoning else []) + adjust(req, prompt_token_ids=prompt) + params = req.to_sampling_params(default_max_tokens=128) + grammar = Grammar.from_structural_tag(params.structured_outputs.structural_tag) + prefix = "" if inside_reasoning else THOUGHT + assert _is_grammar_accept_string(grammar, prefix + FINAL + '{"a": 1}<|return|>') + assert not _is_grammar_accept_string(grammar, prefix + FINAL + "Sunny.<|return|>") + response = ResponsesResponse.from_request(req, params, "beam", 0, [], status) + assert response.model_dump(by_alias=True)["text"] == expected_text + + +@pytest.mark.parametrize("reverse_alternatives", [False, True]) +def test_reasoning_budget_initializes_from_checkpoint_pattern( + monkeypatch, reverse_alternatives +): + template = deepcopy(BEAM_RESPONSE_TEMPLATE) + alternatives = [ + r"^<\|channel\|>analysis<\|message\|>", + r"<\|start\|>assistant<\|channel\|>analysis<\|message\|>", + ] + if reverse_alternatives: + alternatives.reverse() + template["fields"]["thinking"] = { + "open_pattern": "(?:" + "|".join(alternatives) + ")", + "close": "<|end|>", + "content": "text", + } + tokenizer = BeamTokenizer(template) + monkeypatch.setattr( + "vllm.config.reasoning.cached_tokenizer_from_config", lambda **_: tokenizer + ) + config = ReasoningConfig(reasoning_parser="beam") + config.initialize_token_ids(SimpleNamespace()) + + assert config.enabled + # The marker must work both after the prompt's assistant header and in + # subsequent complete assistant messages, regardless of regex branch order. + assert config.reasoning_start_token_ids == tokenizer.encode( + "<|channel|>analysis<|message|>" + ) + assert config.reasoning_end_token_ids == tokenizer.encode("<|end|>") + assert config.natural_reasoning_end_token_ids == tokenizer.encode("<|end|>") + + +@pytest.mark.parametrize("pattern", [r"analysis.*", r"(?:analysis|thinking)"]) +def test_reasoning_budget_does_not_guess_a_marker_for_other_patterns(pattern): + template = deepcopy(BEAM_RESPONSE_TEMPLATE) + thinking = template["fields"]["thinking"] + del thinking["open"] + thinking["open_pattern"] = pattern + + assert BeamReasoningParser(BeamTokenizer(template)).reasoning_start_str is None diff --git a/tests/parser/test_response_template_grammar.py b/tests/parser/test_response_template_grammar.py new file mode 100644 index 000000000000..29ae227262b5 --- /dev/null +++ b/tests/parser/test_response_template_grammar.py @@ -0,0 +1,492 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +from copy import deepcopy +from typing import Any + +import pytest +from xgrammar import Grammar +from xgrammar.testing import _is_grammar_accept_string + +from vllm.parser.chat_parsing.response_templates import load_response_template +from vllm.parser.response_template_grammar import ( + answer_rest_structural_tag, + tool_call_structural_tag, +) + +TEMPLATE: dict[str, Any] = { + "start_anchor": "", + "fields": { + "thinking": {"open": "", "close": "", "content": "text"}, + "content": {"content": "text"}, + "tool_calls": { + "open_pattern": r'', + "close": "", + "repeats": True, + "content": "xml-inline", + "content_args": { + "tag_pattern": r'(?P.*?)' + }, + "transform": { + "type": "function", + "function": {"name": "{name}", "arguments": "{content}"}, + }, + }, + }, +} +ANSWER_TEMPLATE = deepcopy(TEMPLATE) +ANSWER_TEMPLATE["fields"]["content"] = { + "open": "", + "close": "", + "content": "text", +} + + +def tool(parameters, name="f", strict=True) -> dict[str, Any]: + return {"name": name, "parameters": parameters, "strict": strict} + + +WEATHER = tool( + { + "type": "object", + "properties": {"city": {"type": "string"}, "days": {"type": "integer"}}, + "required": ["city"], + }, + name="get.weather", +) +PING = {"name": "ping"} +CALL = 'Paris' + + +def build( + template=TEMPLATE, + functions=(WEATHER, PING), + *, + parallel=True, + thinking=True, + answer=False, + answer_format=None, + sorted_arguments=False, +): + fields = load_response_template(template).fields + return tool_call_structural_tag( + fields["tool_calls"], + list(functions), + parallel_tool_calls=parallel, + thinking=fields["thinking"] if thinking else None, + answer=fields["content"] if answer else None, + answer_format=answer_format, + sorted_arguments=sorted_arguments, + ) + + +def accepts(tag, text): + return _is_grammar_accept_string(Grammar.from_structural_tag(json.dumps(tag)), text) + + +def call(*arguments): + body = "".join(f'{value}' for key, value in arguments) + return f'{body}' + + +@pytest.mark.parametrize( + ("text", "accepted"), + [ + ("plan" + CALL, True), + (CALL, True), + (CALL.removesuffix(""), True), + ( + 'Paris\n' + '3', + True, + ), + ( + 'Paris' + 'soon', + False, + ), + ('3', False), + (CALL + 'x', True), + ('', False), + ("planThe answer.", False), + ], +) +def test_forced_calls_follow_the_template_format(text, accepted): + assert accepts(build(), text) is accepted + + +def test_single_call_when_parallel_calls_are_disabled(): + tag = build(parallel=False) + + assert accepts(tag, CALL) + assert not accepts(tag, CALL + CALL) + + +def test_thinking_is_rejected_when_not_requested(): + assert not accepts(build(thinking=False), "plan" + CALL) + + +@pytest.mark.parametrize( + ("text", "accepted"), + [ + ("planSunny.", True), + ("Sunny.", True), + ("plan" + CALL, True), + ("Sunny." + CALL, False), + ('soon', False), + ], +) +def test_answer_field_lets_the_model_reply_instead(text, accepted): + assert accepts(build(ANSWER_TEMPLATE, answer=True), text) is accepted + + +@pytest.mark.parametrize("required", [False, True]) +@pytest.mark.parametrize( + ("text", "optional_only"), + [ + ("plan" + CALL, False), + ("\nplanSunny.", False), + (" \nSunny.", False), + (CALL, True), + ("Sunny.", True), + ("Sunny.", True), + ("<Sunny.", True), + ], +) +def test_required_reasoning_opens_the_message_with_nonempty_reasoning( + required, text, optional_only +): + fields = load_response_template(ANSWER_TEMPLATE).fields + tag = tool_call_structural_tag( + fields["tool_calls"], + [WEATHER], + parallel_tool_calls=True, + thinking=fields["thinking"], + answer=fields["content"], + message_start=[""], + requires_reasoning=required, + ) + + assert accepts(tag, text) is (not required or not optional_only) + + +def test_answer_rest_finishes_the_open_answer(): + answer = load_response_template(ANSWER_TEMPLATE).fields["content"] + tag = answer_rest_structural_tag(answer) + + assert accepts(tag, "Sunny.") + assert accepts(tag, "Sunny.") + assert not accepts(tag, "Sunny." + CALL) + + +def test_answer_without_literal_opener_is_not_derived(): + assert build(answer=True) is None + + +def test_answer_holds_the_answer_format(): + schema = {"type": "json_schema", "json_schema": {"type": "object"}} + with_tools = build(ANSWER_TEMPLATE, answer=True, answer_format=schema) + answer_only = build(ANSWER_TEMPLATE, (), answer=True, answer_format=schema) + + for tag in (with_tools, answer_only): + assert accepts(tag, 'plan{"a": 1}') + assert not accepts(tag, "Sunny.") + assert accepts(with_tools, CALL) + assert not accepts(answer_only, CALL) + + +def test_json_arguments_follow_the_schema(): + template = deepcopy(TEMPLATE) + template["fields"]["tool_calls"].update(content="json", content_args={}) + tag = build(template, [WEATHER]) + + assert accepts(tag, '{"city": "Paris"}') + assert not accepts(tag, '{"days": 3}') + + +END_X = {"tag_pattern": r'(?P.*?)X'} + + +@pytest.mark.parametrize( + "change", + [ + {"content": "json", "content_args": {"unquoted_keys": True}}, + { + "transform": { + "type": "function", + "function": {"name": "fixed", "arguments": "{content}"}, + } + }, + {"content": "kv-lines", "content_args": {}}, + # Framing in which argument text can spell a closer, where the parser + # would end the call. + {"close": ["}X", ""], "content_args": END_X}, + { + "close": ["}X", ""], + "content_args": {**END_X, "value_parser": {"name": "json"}}, + }, + {"close": ["7X", ""], "content_args": END_X}, + {"close": [">X", ""]}, + {"close": [">", ""]}, + {"close": [">"]}, + {"close": ["}", ""], "content": "json", "content_args": {}}, + ], +) +def test_underivable_tool_formats_are_not_derived(change): + template = deepcopy(TEMPLATE) + template["fields"]["tool_calls"].update(change) + + assert build(template) is None + + +def test_non_literal_thinking_delimiters_are_not_derived(): + template = deepcopy(TEMPLATE) + del template["fields"]["thinking"]["open"] + template["fields"]["thinking"]["open_pattern"] = r"" + + assert build(template) is None + assert build(template, thinking=False) is not None + + +LEVEL = {"type": "object", "properties": {"level": {"$ref": "#/$defs/level"}}} + + +@pytest.mark.parametrize( + ("schema", "value", "accepted"), + [ + ({"type": "integer", "nullable": True}, "null", True), + ({"type": "integer", "nullable": True}, "soon", False), + ({"type": ["integer", "null"]}, "null", True), + ({"anyOf": [{"type": "integer"}, {"type": "boolean"}]}, "true", True), + ({"anyOf": [{"type": "integer"}, {"type": "boolean"}]}, "soon", False), + ({"type": ["string", "integer"]}, "soon", True), + ({"enum": ["celsius", "fahrenheit"]}, "celsius", True), + ({"enum": ["celsius", "fahrenheit"]}, "kelvin", False), + ({"type": "string", "const": "on"}, "off", False), + (LEVEL, '{"level": 2}', True), + (LEVEL, '{"level": "high"}', False), + ], +) +def test_argument_values_follow_how_the_parser_reads_them(schema, value, accepted): + parameters = { + "type": "object", + "properties": {"v": schema}, + "required": ["v"], + "$defs": {"level": {"type": "integer"}}, + } + + assert accepts(build(functions=[tool(parameters)]), call(("v", value))) is accepted + + +@pytest.mark.parametrize( + ("schema", "strict", "derivable"), + [ + # The parser reads values without a type as text. + ({"enum": [1, 2]}, True, False), + ({"$ref": "#/$defs/level"}, True, False), + ({"$ref": "#/$defs/level"}, False, True), + # Raw text cannot hold its closer. + ({"enum": ["ab"]}, True, False), + ({"type": "object", "properties": {"x": {"$ref": "#/x"}}}, True, False), + (False, True, False), + # A match could write the start of the value's end. + ({"type": "string", "pattern": "^[a-z<]+$"}, True, False), + ({"type": "string", "pattern": "^[a-z<]+$"}, False, True), + ({"type": "string", "pattern": "^a.b$"}, True, False), + # XGrammar holds length bounds only in JSON. + ({"type": "string", "minLength": 3}, True, False), + ({"type": "string", "format": "ipv6"}, True, False), + ], +) +def test_schemas_text_values_cannot_hold_are_not_derived(schema, strict, derivable): + parameters = { + "type": "object", + "properties": {"v": schema}, + "$defs": {"level": {"type": "integer"}}, + } + + assert (build(functions=[tool(parameters, strict=strict)]) is not None) is derivable + + +@pytest.mark.parametrize( + ("schema", "value", "accepted"), + [ + ({"type": "string", "format": "email"}, "a.b@x.co", True), + ({"type": "string", "format": "email"}, "ab", False), + ({"type": "string", "format": "date-time"}, "2026-10-06T12:00:00Z", True), + ({"type": "string", "format": "date-time"}, "2026-13-06T12:00:00Z", False), + ({"type": "string", "pattern": "^[A-Z]{3}$"}, "ABC", True), + ({"type": "string", "pattern": "^[A-Z]{3}$"}, "ABCD", False), + # Like XGrammar, the format overrides the pattern and the length bounds. + ( + {"type": "string", "format": "uuid", "pattern": "^x$", "maxLength": 1}, + "123e4567-e89b-12d3-a456-426614174000", + True, + ), + ], +) +def test_text_values_hold_string_formats_and_patterns(schema, value, accepted): + parameters = {"type": "object", "properties": {"v": schema}, "required": ["v"]} + + assert accepts(build(functions=[tool(parameters)]), call(("v", value))) is accepted + + +def test_values_cannot_hold_a_closer(): + template = deepcopy(TEMPLATE) + template["fields"]["tool_calls"]["close"] = ["", "STOP"] + text = {"type": "object", "properties": {"v": {"type": "string"}}} + strings = {"type": "array", "items": {"type": "string"}} + json_text = {"type": "object", "properties": {"v": strings}} + + assert accepts(build(template, [tool(text)]), call(("v", "ab"))) + assert not accepts(build(template, [tool(text)]), call(("v", "aSTOPb"))) + assert not accepts( + build(template, [tool(None, strict=False)]), call(("v", "aSTOPb")) + ) + assert accepts(build(template, [tool(json_text)]), call(("v", '["ab"]'))) + assert not accepts(build(template, [tool(json_text)]), call(("v", '["aSTOPb"]'))) + + +def test_extra_arguments_need_additional_properties(): + parameters = {"type": "object", "properties": {"city": {"type": "string"}}} + extra = call(("city", "Paris"), ("zip", "75001")) + + assert not accepts(build(functions=[tool(parameters)]), extra) + # Next to declared keys, an extra one could repeat them unchecked. + allowed = {**parameters, "additionalProperties": True} + assert build(functions=[tool(allowed)]) is None + free_form = {"type": "object", "additionalProperties": True} + assert accepts(build(functions=[tool(free_form)]), extra) + + +def test_only_strict_tools_follow_their_schema(): + off_schema = call(("days", "soon")) + parameters = WEATHER["parameters"] + + assert accepts(build(functions=[tool(parameters, strict=False)]), off_schema) + assert not accepts(build(functions=[tool(parameters)]), off_schema) + assert accepts(build(functions=[tool(None)]), call()) + assert not accepts(build(functions=[tool(None)]), call(("a", "x"))) + + +def test_annotations_keep_strict_arguments_to_their_schema(): + parameters = { + **WEATHER["parameters"], + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "weather", + } + tag = build(functions=[tool(parameters)]) + + assert accepts(tag, call(("city", "Paris"), ("days", "3"))) + assert not accepts(tag, call(("city", "Paris"), ("days", "soon"))) + + +@pytest.mark.parametrize( + ("arguments", "accepted", "accepted_sorted"), + [ + ((("path", "a"), ("content", "b")), True, True), + ((("content", "b"), ("path", "a")), False, True), + ((("path", "a"),), True, True), + ((("content", "b"),), False, False), + ((("path", "a"), ("path", "c")), False, False), + ((("path", "a"), ("mode", "c")), False, False), + ], +) +def test_strict_arguments_follow_declared_or_sorted_order( + arguments, accepted, accepted_sorted +): + parameters = { + "properties": {"path": {"type": "string"}, "content": {"type": "string"}}, + "required": ["path"], + } + tags = [ + build(functions=[tool(parameters)], sorted_arguments=sorted_arguments) + for sorted_arguments in (False, True) + ] + + assert [accepts(tag, call(*arguments)) for tag in tags] == [ + accepted, + accepted_sorted, + ] + + +LOOKAHEAD_TEMPLATE = deepcopy(TEMPLATE) +LOOKAHEAD_TEMPLATE["fields"]["tool_calls"]["content_args"]["tag_pattern"] = ( + r'(?P.*?)\n(?:|(?=Paris\n\n3\n', True, True), + ('Paris\n3\n', True, True), + ('3\nParis\n', True, False), + ('Paris\nLyon\n', True, False), + ('Paris\n3\n', True, False), + ('Paris\nParis\na c\n', True, True), + ], +) +def test_values_end_where_a_lookahead_end_does(body, accepted_loose, accepted_strict): + parameters = WEATHER["parameters"] + tags = [ + build(LOOKAHEAD_TEMPLATE, [tool(parameters, strict=strict)]) + for strict in (False, True) + ] + text = f'{body}' + + assert [accepts(tag, text) for tag in tags] == [accepted_loose, accepted_strict] + + +@pytest.mark.parametrize("lookahead", [False, True]) +@pytest.mark.parametrize( + ("keys", "accepted"), + [ + ("b", True), + ("ab", True), + ("bc", True), + ("bd", True), + ("be", True), + ("abde", True), + ("bcde", True), + ("", False), + ("a", False), + ("acd", False), + ("bdc", False), + ("bcc", False), + ], +) +def test_last_parameter_follows_every_required_one(keys, accepted, lookahead): + parameters = { + "properties": {key: {"type": "string"} for key in "abcde"}, + "required": ["b"], + } + tag = build(LOOKAHEAD_TEMPLATE, [tool(parameters)]) + ends = ["\n" if lookahead else "\n"] * (len(keys) - 1) + ["\n"] + body = "".join(f'x{e}' for k, e in zip(keys, ends)) + + assert accepts(tag, f'{body}') is accepted + + +def test_required_argument_without_schema_is_still_required(): + parameters = {"type": "object", "properties": {}, "required": ["a"]} + tag = build(functions=[tool(parameters)]) + + assert accepts(tag, call(("a", "x"))) + assert not accepts(tag, call()) + assert build(functions=[tool({**parameters, "required": ["a", 1]})]) is None + + +@pytest.mark.parametrize( + ("function", "match"), + [ + (tool({}, name="get weather"), "tool name"), + (tool({"properties": {'a"b': {"type": "string"}}}), "parameter name"), + ], +) +def test_unencodable_names_are_rejected(function, match): + with pytest.raises(ValueError, match=match): + build(functions=[function]) diff --git a/tests/test_config.py b/tests/test_config.py index c68ecca605ac..b1f977bd6e0b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -924,7 +924,10 @@ def test_v2_model_runner_supports_custom_logits_processors(): assert config._get_v2_model_runner_unsupported_features() == [] -@pytest.mark.parametrize("architecture", ["DFlash2DraftModel", "LiLiCorrDraftModel"]) +@pytest.mark.parametrize( + "architecture", + ["DFlash2DraftModel", "DFlash2BeamDraftModel", "LiLiCorrDraftModel"], +) def test_dflash_candidate_draft_forces_v2_model_runner(architecture): """A DFlash2 draft must reach the V2 speculator, the only one that runs its candidate selector; on V1 it would draft as DFlash1 without raising.""" diff --git a/tests/transformers_utils/test_beam_config.py b/tests/transformers_utils/test_beam_config.py new file mode 100644 index 000000000000..1bed226529a7 --- /dev/null +++ b/tests/transformers_utils/test_beam_config.py @@ -0,0 +1,305 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Test Beam config aliases, hybrid attention, and serialization.""" + +from typing import Any + +import pytest + +from vllm.transformers_utils.config import get_config +from vllm.transformers_utils.configs.beam import ( + BeamConfig, + _canonical_dtype, +) + +# Synthetic configuration with hybrid attention and grouped routing. +_NUM_LAYERS = 4 +_GLOBAL_ATTENTION_LAYER_IDS = frozenset({1, 3}) + +_SYNTHETIC_CONFIG: dict[str, Any] = { + "vocab_size": 256, + "hidden_size": 64, + "intermediate_size": 128, + "num_hidden_layers": _NUM_LAYERS, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 16, + "hidden_act": "silu", + "max_position_embeddings": 256, + "rms_norm_eps": 1e-6, + "tie_word_embeddings": False, + "attention_bias": False, + "attention_output_gate": False, + "sliding_window": 32, + "layer_types": [ + "full_attention" + if layer_id in _GLOBAL_ATTENTION_LAYER_IDS + else "sliding_attention" + for layer_id in range(_NUM_LAYERS) + ], + "moe_intermediate_size": 128, + "num_experts_per_tok": 2, + "num_experts": 8, + "norm_topk_prob": True, + "scoring_func": "sigmoid", + "n_group": 2, + "topk_group": 1, + "routed_scaling_factor": 1.5, + "topk_method": "noaux_tc", + "mlp_only_layers": [0], + "rope_theta": 10000, + "rope_scaling": None, + "nope_layer_ids": sorted(_GLOBAL_ATTENTION_LAYER_IDS), + "mlp_depth_scaled_residual": True, + "self_attn_depth_scaled_residual": True, + "final_logit_softcapping": 30.0, + "moe_router_dtype": "fp32", + "ssmax": True, + "ssmax_offset": 1, + "ssmax_divide_by_log_offset": False, + "ssmax_scale_mode": "root10", + "use_paramfree_qk_norm": True, + "enable_sandwich_norm": True, +} + + +def test_synthetic_config_resolves(): + config = BeamConfig(**_SYNTHETIC_CONFIG) + assert config.model_type == "beam" + assert config.architectures == ["BeamForCausalLM"] + assert config.num_hidden_layers == 4 + assert config.hidden_size == 64 + assert config.num_attention_heads == 4 + assert config.num_key_value_heads == 2 + assert config.head_dim == 16 + assert config.num_experts == 8 + assert config.moe_intermediate_size == 128 + assert config.nope_layer_ids == [1, 3] + assert config.mlp_only_layers == [0] + assert config.final_logit_softcapping == 30.0 + assert config.moe_router_dtype == "float32" + assert config.ssmax and config.ssmax_scale_mode == "root10" + assert config.use_paramfree_qk_norm + assert config.enable_sandwich_norm + for layer_id, layer_type in enumerate(config.layer_types): + expected = ( + "full_attention" + if layer_id in _GLOBAL_ATTENTION_LAYER_IDS + else "sliding_attention" + ) + assert layer_type == expected + + +def test_config_aliases_resolve(): + """Config aliases resolve to the fields used by the model.""" + no_rope = [1] * 4 + for i in (1, 3): + no_rope[i] = 0 + config = BeamConfig( + num_hidden_layers=4, + no_rope_layers=no_rope, + num_experts=8, + router_dtype="fp32", + rope_parameters={"rope_theta": 5000.0, "rope_type": "default"}, + ) + assert config.nope_layer_ids == [1, 3] + assert config.num_experts == 8 + assert config.moe_router_dtype == "float32" + assert config.rope_theta == 5000.0 + + +def test_unsupported_dtype_keys_are_dropped(): + """Unsupported dtype options are excluded from the serialized config.""" + config = BeamConfig( + num_hidden_layers=4, residual_in_fp32=True, lm_head_dtype="fp32" + ) + for key in ("residual_in_fp32", "lm_head_dtype", "fp32_residual_connection"): + assert not hasattr(config, key) + assert key not in config.to_dict() + + +def test_layer_types_derived_when_omitted(): + """A missing `layer_types` must not leave the SWA/full split implicit. + + vLLM only skips populating the global `CacheConfig.sliding_window` when + `layer_types` says the model is partly sliding; without it the + full-attention layers pick that global value up as a fallback and silently + get a window. + """ + no_rope = [1] * 8 + for i in (3, 7): + no_rope[i] = 0 + config = BeamConfig(num_hidden_layers=8, no_rope_layers=no_rope, sliding_window=512) + assert config.layer_types == [ + "full_attention" if i in (3, 7) else "sliding_attention" for i in range(8) + ] + + +def test_explicit_args_beat_aliases(): + config = BeamConfig( + num_hidden_layers=4, + nope_layer_ids=[1], + no_rope_layers=[1, 1, 0, 1], + moe_router_dtype="bf16", + router_dtype="float32", + ) + assert config.nope_layer_ids == [1] + assert config.moe_router_dtype == "bfloat16" + + +def test_to_dict_roundtrip_keeps_canonical_fields(): + config = BeamConfig(**_SYNTHETIC_CONFIG) + serialized = config.to_dict() + assert serialized["num_experts"] == 8 + assert "num_local_experts" not in serialized + rebuilt = BeamConfig(**serialized) + assert rebuilt.num_experts == 8 + assert rebuilt.nope_layer_ids == config.nope_layer_ids + assert rebuilt.moe_router_dtype == config.moe_router_dtype + + +def test_noncanonical_expert_count_is_rejected(): + with pytest.raises(ValueError, match="requires num_experts"): + BeamConfig(num_local_experts=8) + + +def test_nested_rope_checkpoint_loads_and_infers_nope(tmp_path): + """Release metadata must survive generic loading and config round trips.""" + import json + + checkpoint = { + "model_type": "beam", + "num_hidden_layers": 2, + "layer_types": ["sliding_attention", "full_attention"], + "rope_parameters": { + "sliding_attention": {"rope_type": "default", "rope_theta": 5000.0}, + "full_attention": None, + }, + } + (tmp_path / "config.json").write_text(json.dumps(checkpoint)) + config = get_config(str(tmp_path), trust_remote_code=False) + assert config.nope_layer_ids == [1] + assert config.rope_parameters == {"rope_type": "default", "rope_theta": 5000.0} + rebuilt = BeamConfig(**config.to_dict()) + assert rebuilt.nope_layer_ids == [1] + assert rebuilt.rope_parameters == config.rope_parameters + + +def test_release_config_omitted_features_match_reference(tmp_path): + """A checkpoint config that omits the fixed-behavior switches resolves to them.""" + import json + + num_layers, nope_layers = 12, (3, 7, 11) + checkpoint = { + "model_type": "beam", + "architectures": ["BeamForCausalLM"], + "num_hidden_layers": num_layers, + "num_experts": 32, + "n_group": 4, + "topk_group": 2, + "routed_scaling_factor": 2.5, + "sliding_window": 128, + "layer_types": [ + "full_attention" if i in nope_layers else "sliding_attention" + for i in range(num_layers) + ], + "mlp_layer_types": ["dense"] * 2 + ["sparse"] * (num_layers - 2), + "rope_parameters": { + "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, + "full_attention": None, + }, + } + (tmp_path / "config.json").write_text(json.dumps(checkpoint)) + config = get_config(str(tmp_path), trust_remote_code=False) + for resolved in (config, BeamConfig(**config.to_dict())): + assert resolved.attention_output_gate + assert resolved.attention_output_gate_bias + assert resolved.enable_sandwich_norm + assert resolved.self_attn_depth_scaled_residual + assert resolved.mlp_depth_scaled_residual + assert resolved.ssmax + assert resolved.ssmax_scale_mode == "root10" + assert resolved.scoring_func == "sigmoid" + assert resolved.topk_method == "noaux_tc" + assert resolved.nope_layer_ids == list(nope_layers) + assert resolved.mlp_only_layers == [0, 1] + + +@pytest.mark.parametrize( + "explicit", [{"nope_layer_ids": []}, {"no_rope_layers": [1, 1]}] +) +def test_explicit_rope_schedule_beats_layer_types(explicit): + config = BeamConfig( + num_hidden_layers=2, + layer_types=["sliding_attention", "full_attention"], + **explicit, + ) + assert config.nope_layer_ids == [] + + +@pytest.mark.parametrize( + "overrides", + [ + {"layer_types": ["full_attention"]}, + {"layer_types": ["full_attention", "invalid"]}, + {"no_rope_layers": [1]}, + {"no_rope_layers": [1, 2]}, + {"nope_layer_ids": [2]}, + {"nope_layer_ids": [1, 1]}, + ], +) +def test_invalid_attention_schedule_is_rejected(overrides): + with pytest.raises(ValueError): + BeamConfig(num_hidden_layers=2, **overrides) + + +def test_dtype_spellings_are_canonicalized(): + """Short spellings must resolve to the keys of `STR_DTYPE_TO_TORCH_DTYPE`.""" + assert _canonical_dtype(None) is None + assert _canonical_dtype("fp32") == "float32" + assert _canonical_dtype("float32") == "float32" + assert _canonical_dtype("bf16") == "bfloat16" + assert _canonical_dtype("bfloat16") == "bfloat16" + with pytest.raises(ValueError, match="Unsupported dtype"): + _canonical_dtype("fp16") + + +def test_mlp_layer_types_derives_dense_layers(): + """Newer checkpoints list per-layer feed-forward types instead of + `mlp_only_layers`; the model keeps consuming the dense-layer list.""" + config = BeamConfig( + num_hidden_layers=4, mlp_layer_types=["dense", "sparse", "sparse", "dense"] + ) + assert config.mlp_only_layers == [0, 3] + assert config.mlp_layer_types == ["dense", "sparse", "sparse", "dense"] + + +def test_mlp_only_layers_derives_layer_types(): + config = BeamConfig(num_hidden_layers=4, mlp_only_layers=[0]) + assert config.mlp_layer_types == ["dense", "sparse", "sparse", "sparse"] + + +def test_agreeing_mlp_schemas_are_accepted_and_round_trip(): + config = BeamConfig( + num_hidden_layers=3, + mlp_only_layers=[0, 1], + mlp_layer_types=["dense", "dense", "sparse"], + ) + rebuilt = BeamConfig(**config.to_dict()) + assert rebuilt.mlp_only_layers == [0, 1] + assert rebuilt.mlp_layer_types == ["dense", "dense", "sparse"] + + +@pytest.mark.parametrize( + "overrides", + [ + {"mlp_only_layers": [0], "mlp_layer_types": ["sparse", "dense", "sparse"]}, + {"mlp_layer_types": ["dense", "sparse"]}, + {"mlp_layer_types": ["dense", "moe", "sparse"]}, + {"mlp_only_layers": [0, 0]}, + {"mlp_only_layers": [3]}, + ], +) +def test_invalid_mlp_schedules_are_rejected(overrides): + with pytest.raises(ValueError): + BeamConfig(num_hidden_layers=3, **overrides) diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index 558bb38e6bc9..efc5677c5b3a 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -39,6 +39,50 @@ def test_grouped_conv_matches_reference(block_size: int): torch.testing.assert_close(actual, expected.flatten(0, 1).flatten(-2)) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@pytest.mark.parametrize( + "rows,width,query,group,taps", + [ + (8, 5120, 8, 16, 2), + (64, 5120, 8, 16, 2), + (256, 5120, 8, 16, 2), + (15, 143, 5, 11, 3), + ], +) +@pytest.mark.parametrize("warps", [None, 4, 8, 16]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +def test_beam_conv_residual_norm(rows, width, query, group, taps, warps, dtype): + """Dtype rounding and request boundaries survive residual/norm fusion.""" + from vllm import _custom_ops as ops + from vllm.models.beam.nvidia.dflash2_conv_ops import conv_residual_norm + + torch.manual_seed(42) + hidden = torch.randn(rows, width, device="cuda", dtype=dtype) + residual = torch.randn_like(hidden) + weight = torch.randn(width, device="cuda", dtype=hidden.dtype) + base = torch.randn(taps, width, device="cuda", dtype=hidden.dtype) + delta = torch.randn( + rows, 2, taps, width // group, device="cuda", dtype=hidden.dtype + )[:, 1] + expected = _grouped_conv(hidden, delta, base, query, width // group, group, taps) + expected_residual = residual.clone() + ops.fused_add_rms_norm(expected, expected_residual, weight, 1e-6) + actual, actual_residual = conv_residual_norm( + hidden, delta, base, residual, weight, 1e-6, query, group, num_warps=warps + ) + tolerance = {torch.bfloat16: 0.015625, torch.float16: 0.002, torch.float32: 3e-6}[ + dtype + ] + residual_tolerance = 0 if dtype == torch.bfloat16 else tolerance + torch.testing.assert_close( + actual_residual, + expected_residual, + rtol=residual_tolerance, + atol=residual_tolerance, + ) + torch.testing.assert_close(actual, expected, rtol=tolerance, atol=tolerance) + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize( @@ -106,6 +150,128 @@ def test_draft_quant_exclusions_include_global_layer_indices(): assert "lilicorr.layers.88.mlp.0" not in quant_config.exclude_modules +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@pytest.mark.parametrize( + "rows,heads,scale", [(8, 1, 1.0), (64, 2, 0.25), (256, 1, 1.0)] +) +@pytest.mark.parametrize("neox", [True, False]) +def test_beam_fused_context_kv(rows, heads, scale, neox): + """All layer pointers and strided cache pages match the staged context path.""" + from vllm import _custom_ops as ops + from vllm.models.beam.nvidia.dflash2_context_ops import store_context_kv + + torch.manual_seed(42) + layers, dim, page = 5, 128, 16 + kv = torch.randn(rows, layers, 2, heads, dim, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(layers, dim, device="cuda", dtype=kv.dtype) + positions = torch.arange(rows, device="cuda", dtype=torch.int64) + angles = torch.randn(rows, dim // 2, device="cuda") + rope = torch.cat((angles.cos(), angles.sin()), -1).to(kv.dtype) + staged = kv.permute(2, 1, 0, 3, 4).contiguous() + keys = torch.empty_like(staged[0]) + ops.rms_norm(keys, staged[0], weight, 1e-6) + ops.rotary_embedding( + positions.repeat(layers), keys.view(layers * rows, -1), None, dim, rope, neox + ) + values = staged[1] * scale + storage = torch.randn( + (rows + page - 1) // page, + layers, + heads, + page, + 2 * dim, + device="cuda", + dtype=kv.dtype, + ) + expected = storage.clone() + slots = positions.flip(0)[:-1].contiguous() + slots[::3] = -1 + valid = slots >= 0 + ids = slots[valid] + for layer in range(layers): + expected[ids // page, layer, :, ids % page, :dim] = keys[layer, :-1][valid] + expected[ids // page, layer, :, ids % page, dim:] = values[layer, :-1][valid] + caches = storage.unbind(1) + pointers = torch.tensor( + [c.data_ptr() for c in caches], device="cuda", dtype=torch.int64 + ) + store_context_kv( + kv, weight, positions, rope, pointers, slots, caches[0], 1e-6, scale, neox + ) + torch.testing.assert_close(storage, expected, rtol=0.01, atol=0.015625) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@pytest.mark.parametrize("rows", [1, 8, 65, 2048]) +@pytest.mark.parametrize("cache_dtype", ["fp8_e4m3", "fp8_e5m2"]) +def test_beam_fp8_context_cache_matches_backend_and_replays(rows, cache_dtype): + """Quantize rounded context K/V with each layer's live scales and slot mask.""" + from vllm import _custom_ops as ops + from vllm.models.beam.nvidia.dflash2_context_ops import store_context_kv + + torch.manual_seed(42) + layers, heads, dim, page = 3, 2, 128, 16 + kv = torch.randn(rows, layers, 2, heads, dim, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(layers, dim, device="cuda", dtype=kv.dtype) + positions = torch.arange(rows, device="cuda") + angles = torch.randn(rows, dim // 2, device="cuda") + rope = torch.cat((angles.cos(), angles.sin()), -1).to(kv.dtype) + shape = ((rows + page - 1) // page, layers, heads, page, 2 * dim) + bf16 = torch.zeros(shape, device="cuda", dtype=kv.dtype) + actual = torch.zeros(shape, device="cuda", dtype=torch.uint8) + expected = torch.zeros_like(actual) + dtype = torch.float8_e4m3fn if cache_dtype == "fp8_e4m3" else torch.float8_e5m2 + scales = torch.tensor([[0.01, 0.02], [0.5, 1.5], [2.0, 3.0]], device="cuda") + scale_ptrs = torch.tensor( + [s.data_ptr() for pair in scales.unbind() for s in pair.unbind()], + device="cuda", + ) + + def pointers(storage): + return torch.tensor([c.data_ptr() for c in storage.unbind(1)], device="cuda") + + slots = positions.flip(0).contiguous() + slots[1::3] = -1 + args = (kv, weight, positions, rope) + bf16_ptrs, fp8_ptrs = pointers(bf16), pointers(actual) + + def run(): + store_context_kv( + *args, + fp8_ptrs, + slots, + actual[:, 0].view(dtype), + 1e-6, + 0.25, + True, + scale_ptrs, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for factor in (1.0, 0.5): + scales.mul_(factor) + kv.mul_(factor) + store_context_kv(*args, bf16_ptrs, positions, bf16[:, 0], 1e-6, 0.25, True) + for layer in range(layers): + source = bf16[:, layer].transpose(1, 2).reshape(-1, heads, 2 * dim) + cache = expected[:, layer].transpose(1, 2) + ops.reshape_and_cache_flash( + source[:rows, :, :dim].contiguous(), + source[:rows, :, dim:].contiguous(), + cache[..., :dim], + cache[..., dim:], + slots, + cache_dtype, + scales[layer, 0], + scales[layer, 1], + ) + graph.replay() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + def test_selector_edges_match_sequential_reference(): torch.manual_seed(1) batch, steps, top_k, rank = 2, 4, 3, 5 @@ -142,6 +308,78 @@ def test_selector_edges_match_sequential_reference(): torch.testing.assert_close(actual, expected) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@pytest.mark.parametrize( + "batch,steps,top_k,rank", + [(1, 1, 3, 5), (1, 7, 16, 256), (8, 7, 16, 256), (32, 7, 16, 256), (3, 5, 17, 65)], +) +@pytest.mark.parametrize("unary_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("anchor_dtype", [torch.int32, torch.int64]) +def test_beam_fused_candidate_edges( + batch, steps, top_k, rank, unary_dtype, anchor_dtype +): + """Keep the anchor, predecessor mapping, and BF16 intermediate rounding.""" + from vllm.models.beam.nvidia.dflash2_ops import candidate_edges + + torch.manual_seed(1) + vocab = 257 + predecessors = torch.randn(vocab, rank, device="cuda", dtype=torch.bfloat16) + successors = torch.randn_like(predecessors) + candidates = torch.randint(vocab, (batch, steps, top_k), device="cuda") + unary = torch.randn(batch, steps, top_k, device="cuda", dtype=unary_dtype) + hidden = torch.randn(batch, steps, rank, device="cuda", dtype=torch.bfloat16) + # Production anchors are a strided view of the query token buffer. + anchors = torch.randint(vocab, (batch * 8,), device="cuda", dtype=anchor_dtype)[::8] + args = (predecessors, successors, candidates, unary, hidden, anchors, top_k) + expected = _score_edges(*args) + actual = candidate_edges(*args) + torch.testing.assert_close(actual, expected, rtol=0.01, atol=0.015625) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@pytest.mark.parametrize("rows", [8, 64, 256]) +def test_beam_draft_qk_norm_rope_cache(rows): + """Fused query preparation preserves BF16 rounding and ignores padded slots.""" + from vllm import _custom_ops as ops + from vllm.models.beam.nvidia.attention_ops import beam_qk_norm_rope + + torch.manual_seed(42) + qkv = torch.randn(rows, 896, device="cuda", dtype=torch.bfloat16) + q, k, v = qkv.split((640, 128, 128), -1) + qw, kw = torch.randn(2, 128, device="cuda", dtype=q.dtype) + positions = torch.arange(rows, device="cuda", dtype=torch.int64) + angles = torch.randn(rows, 64, device="cuda") + rope = torch.cat((angles.cos(), angles.sin()), -1).to(q.dtype) + oq, ok = torch.empty_like(q), torch.empty_like(k) + ops.rms_norm(oq.view(-1, 5, 128), q.view(-1, 5, 128), qw, 1e-6) + ops.rms_norm(ok.view(-1, 1, 128), k.view(-1, 1, 128), kw, 1e-6) + ops.rotary_embedding(positions, oq, ok, 128, rope, True) + cache = torch.randn((rows + 15) // 16, 1, 16, 256, device="cuda", dtype=q.dtype) + expected_cache = cache.clone() + slots = positions.flip(0).contiguous() + slots[::3] = -1 + valid = slots >= 0 + slot_ids = slots[valid] + expected_cache[slot_ids // 16, 0, slot_ids % 16, :128] = ok[valid] + expected_cache[slot_ids // 16, 0, slot_ids % 16, 128:] = v[valid] + actual = beam_qk_norm_rope( + q, + k, + qw, + kw, + 1e-6, + positions, + rope, + 128, + 128, + value=v, + kv_cache=cache, + slot_mapping=slots, + ) + torch.testing.assert_close(actual, (oq, ok), rtol=0.01, atol=0.015625) + torch.testing.assert_close(cache, expected_cache, rtol=0.01, atol=0.015625) + + def _stub_base(monkeypatch, draft_logits): """A DFlashSpeculator.__init__ that allocates only what the base class would. @@ -189,20 +427,8 @@ def test_selector_asks_for_fp32_proposal_logits(): assert fill == float("-inf") -@pytest.mark.skip_global_cleanup -@pytest.mark.parametrize("variant", ["dflash2", "lilicorr", "lilicorr_plain"]) -def test_candidate_model_decoder_layer_cls(monkeypatch, variant): - from types import SimpleNamespace - - from vllm.config import set_current_vllm_config - from vllm.model_executor.models.lilicorr import LiLiCorr - from vllm.model_executor.models.qwen3_dflash import DFlashQwen3DecoderLayer - from vllm.model_executor.models.qwen3_dflash2 import ( - DFlash2Qwen3DecoderLayer, - DFlash2Qwen3Model, - ) - - # 1. Mock get_current_vllm_config and TP groups +def _mock_draft_configs(monkeypatch, hf_config): + """The current and draft vLLM configs for building a draft model on CPU.""" mock_current_vllm_config = SimpleNamespace( cache_config=SimpleNamespace( block_size=16, @@ -234,13 +460,12 @@ def test_candidate_model_decoder_layer_cls(monkeypatch, variant): dtype=torch.float32, is_mm_prefix_lm=False, rswa_window=None, + head_dtype=None, ), kernel_config=SimpleNamespace( linear_backend="auto", ), ) - from vllm.platforms import current_platform - monkeypatch.setattr( current_platform, "get_attn_backend_cls", @@ -258,7 +483,41 @@ class MockGroup: MockGroup(), ) - # 2. Mock vllm_config + vllm_config = SimpleNamespace( + speculative_config=SimpleNamespace( + draft_model_config=SimpleNamespace( + hf_config=hf_config, + quantization=None, + ), + num_speculative_tokens=4, + enable_adaptive_verification=False, + ), + model_config=SimpleNamespace( + dtype=torch.float32, + is_mm_prefix_lm=False, + get_total_num_hidden_layers=lambda: 2, + ), + load_config=SimpleNamespace( + quantization=None, + quantization_param_path=None, + ), + ) + mock_current_vllm_config.speculative_config = vllm_config.speculative_config + vllm_config.compilation_config = mock_current_vllm_config.compilation_config + return mock_current_vllm_config, vllm_config + + +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize("variant", ["dflash2", "lilicorr", "lilicorr_plain"]) +def test_candidate_model_decoder_layer_cls(monkeypatch, variant): + from vllm.config import set_current_vllm_config + from vllm.model_executor.models.lilicorr import LiLiCorr + from vllm.model_executor.models.qwen3_dflash import DFlashQwen3DecoderLayer + from vllm.model_executor.models.qwen3_dflash2 import ( + DFlash2Qwen3DecoderLayer, + DFlash2Qwen3Model, + ) + hf_config = SimpleNamespace( vocab_size=1000, hidden_size=256, @@ -287,33 +546,13 @@ class MockGroup: "use_aux_hidden_state": False, }, ) - vllm_config = SimpleNamespace( - speculative_config=SimpleNamespace( - draft_model_config=SimpleNamespace( - hf_config=hf_config, - quantization=None, - ), - num_speculative_tokens=4, - enable_adaptive_verification=False, - ), - model_config=SimpleNamespace( - dtype=torch.float32, - is_mm_prefix_lm=False, - ), - load_config=SimpleNamespace( - quantization=None, - quantization_param_path=None, - ), - ) - mock_current_vllm_config.speculative_config = vllm_config.speculative_config - vllm_config.compilation_config = mock_current_vllm_config.compilation_config + current_config, vllm_config = _mock_draft_configs(monkeypatch, hf_config) - # 3. Instantiate the model under meta device to avoid parameter allocation issues - with set_current_vllm_config(mock_current_vllm_config), torch.device("meta"): + # Instantiate the model under meta device to avoid parameter allocation issues + with set_current_vllm_config(current_config), torch.device("meta"): model_cls = DFlash2Qwen3Model if variant == "dflash2" else LiLiCorr model = model_cls(vllm_config=vllm_config) - # 4. Assert that the layers are DFlash2Qwen3DecoderLayer (the subclass) assert len(model.layers) == 2 expected = ( DFlashQwen3DecoderLayer @@ -323,6 +562,93 @@ class MockGroup: assert type(model.layers[0]) is expected +@pytest.mark.skip_global_cleanup +def test_beam_dflash2_loads_every_checkpoint_tensor(monkeypatch): + """The exported checkpoint's names load into every parameter but the + embedding and LM head, which alias the target's.""" + from vllm.config import set_current_vllm_config + from vllm.models.beam import DFlash2BeamForCausalLM + + hidden, heads, kv_heads, head_dim, ffn, vocab, rank = 64, 4, 2, 16, 128, 100, 8 + taps, groups = 2, 4 + hf_config = SimpleNamespace( + vocab_size=vocab, + hidden_size=hidden, + num_hidden_layers=2, + num_attention_heads=heads, + num_key_value_heads=kv_heads, + head_dim=head_dim, + max_position_embeddings=2048, + rms_norm_eps=1e-6, + rope_parameters={"rope_type": "default", "rope_theta": 10000.0}, + intermediate_size=ffn, + hidden_act="silu", + is_causal=False, + layer_types=["sliding_attention"] * 2, + sliding_window=64, + dflash_config={ + "block_size": 16, + "conv_group_size": hidden // groups, + "conv_kernel_size": taps, + "final_logit_softcapping": 30, + "mask_token_id": 99, + "selector_rank": rank, + "selector_top_k": 3, + "target_layer_ids": [0, 1], + }, + ) + current_config, vllm_config = _mock_draft_configs(monkeypatch, hf_config) + with set_current_vllm_config(current_config): + model = DFlash2BeamForCausalLM(vllm_config=vllm_config) + + shapes = { + "fc.weight": (hidden, 2 * hidden), + "hidden_norm.weight": (hidden,), + "norm.weight": (hidden,), + "candidate_selector.hidden_projection.weight": (rank, hidden), + "candidate_selector.predecessor_codebook": (vocab, rank), + "candidate_selector.successor_codebook": (vocab, rank), + } + for layer in range(2): + for conv in ("attention_conv", "mlp_conv"): + shapes[f"layers.{layer}.{conv}.base_kernel"] = (2, taps, hidden) + shapes[f"layers.{layer}.{conv}.kernel_projection.weight"] = ( + 2 * taps * groups, + hidden, + ) + shapes |= { + f"layers.{layer}.input_layernorm.weight": (hidden,), + f"layers.{layer}.post_attention_layernorm.weight": (hidden,), + f"layers.{layer}.self_attn.q_proj.weight": (heads * head_dim, hidden), + f"layers.{layer}.self_attn.k_proj.weight": (kv_heads * head_dim, hidden), + f"layers.{layer}.self_attn.v_proj.weight": (kv_heads * head_dim, hidden), + f"layers.{layer}.self_attn.o_proj.weight": (hidden, heads * head_dim), + f"layers.{layer}.self_attn.q_norm.weight": (head_dim,), + f"layers.{layer}.self_attn.k_norm.weight": (head_dim,), + f"layers.{layer}.mlp.gate_proj.weight": (ffn, hidden), + f"layers.{layer}.mlp.up_proj.weight": (ffn, hidden), + f"layers.{layer}.mlp.down_proj.weight": (hidden, ffn), + } + torch.manual_seed(0) + checkpoint = {name: torch.randn(shape) for name, shape in shapes.items()} + + loaded = model.load_weights(iter(checkpoint.items())) + + assert loaded == {name for name, _ in model.named_parameters()} + context_kv = [ + checkpoint[f"layers.{layer}.self_attn.{proj}_proj.weight"] + for layer in range(2) + for proj in ("k", "v") + ] + torch.testing.assert_close(model.model._context_kv_weight, torch.cat(context_kv)) + # Numbered after the target's two layers, so the KV cache keys stay distinct. + assert model.get_draft_kv_cache_layer_names() == [ + "model.layers.2.self_attn.attn", + "model.layers.3.self_attn.attn", + ] + assert model.get_draft_attn_causal() == [False, False] + + def test_conv_projections_use_draft_quant_config(monkeypatch): from torch import nn diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 70fed78573dd..d74bb6781612 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2424,6 +2424,11 @@ def grouped_topk( raise NotImplementedError( "The fused grouped_topk kernel is only available on CUDA and XPU platforms" ) + if scores.shape[0] == 0: + return ( + scores.new_empty((0, topk), dtype=torch.float32), + scores.new_empty((0, topk), dtype=torch.int32), + ) return torch.ops._moe_C.grouped_topk( scores, num_expert_group, diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 0205393f69c5..bd64e563173e 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1371,8 +1371,8 @@ def __post_init__(self): f"Unsupported speculative method: '{self.method}'" ) - if self.method in ("eagle", "eagle3"): - # EAGLE drafts share the target's positional space; a + if self.method in ("eagle", "eagle3", "dflash"): + # EAGLE and DFlash drafts share the target's positional space; a # draft checkpoint with a smaller max_position_embeddings # than the target under-sizes its rotary cache (#48894). SpeculativeConfig._maybe_override_draft_max_position_embeddings( @@ -1651,7 +1651,7 @@ def _maybe_override_draft_max_position_embeddings( draft_hf_config: PreTrainedConfig, target_max_model_len: int, ) -> None: - """Raise an EAGLE draft's max_position_embeddings up to the target's. + """Raise an EAGLE or DFlash draft's position limit up to the target's. The proposer feeds the draft positions up to the target's max_model_len, while max_position_embeddings sizes the draft's @@ -1674,7 +1674,7 @@ def _maybe_override_draft_max_position_embeddings( return logger.info( "Overriding draft model max_position_embeddings from %d to the " - "target model's max_model_len (%d); EAGLE drafts share the " + "target model's max_model_len (%d); these drafts share the " "target's positional space.", draft_max_position_embeddings, target_max_model_len, diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 53b9784dadb1..ef09c4ba87da 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -74,6 +74,7 @@ DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset( { "AliceAIForCausalLM", + "BeamForCausalLM", "DeepseekV32MTPModel", "DeepseekV32ForCausalLM", "DeepseekV4ForCausalLM", @@ -749,9 +750,11 @@ def _is_dflash_candidate_draft(self) -> bool: if draft_config is None: return False return bool( - {"DFlash2DraftModel", "LiLiCorrDraftModel"}.intersection( - draft_config.architectures or [] - ) + { + "DFlash2DraftModel", + "DFlash2BeamDraftModel", + "LiLiCorrDraftModel", + }.intersection(draft_config.architectures or []) ) def _dflash_needs_multi_kv_group(self) -> bool: diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index 4b7f2b95a59b..e66b4a62e0ce 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -1005,8 +1005,21 @@ def stop_and_flush() -> list[str]: ) yield wrap_data_with_event(data, "message_stop") else: - origin_chunk = ChatCompletionStreamResponse.model_validate_json( - data_str + payload = json.loads(data_str) + if "error" in payload: + backend_error = ErrorResponse.model_validate(payload).error + error_response = AnthropicStreamEvent( + type="error", + error=AnthropicError( + type=backend_error.type, + message=backend_error.message, + ), + ) + data = error_response.model_dump_json(exclude_unset=True) + yield wrap_data_with_event(data, "error") + return + origin_chunk = ChatCompletionStreamResponse.model_validate( + payload ) if first_item: diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 3ddf81b288aa..39d5ff9961e2 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -53,6 +53,7 @@ from pydantic import ( BeforeValidator, Field, + PrivateAttr, ValidationError, field_serializer, model_validator, @@ -469,6 +470,14 @@ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: max_output_tokens_param="max_output_tokens", ) + _grammar_from_parser: bool = PrivateAttr(default=False) + """CAUTION: Should only be set by the parser's adjust_request.""" + + _parser_structured_outputs: StructuredOutputsParams | None = PrivateAttr( + default=None + ) + """Parser-derived constraints that preserve the public request's text metadata.""" + _DEFAULT_SAMPLING_PARAMS = { "temperature": 1.0, "top_p": 1.0, @@ -478,6 +487,8 @@ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: def extract_structured_outputs(self) -> StructuredOutputsParams | None: """Normalize request constraints into ``StructuredOutputsParams``.""" + if self._parser_structured_outputs is not None: + return self._parser_structured_outputs if self.text is None or self.text.format is None: return self.structured_outputs diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index dd154f968b66..ba9d95b42873 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -526,6 +526,7 @@ async def _create_responses( trace_headers=trace_headers, session_id=session_id, reasoning_parser_kwargs=reasoning_parser_kwargs, + reasoning_ended=True if request._grammar_from_parser else None, ) generators.append(generator) @@ -647,6 +648,7 @@ async def _generate_with_builtin_tools( trace_headers: Mapping[str, str] | None = None, session_id: str | None = None, reasoning_parser_kwargs: dict[str, Any] | None = None, + reasoning_ended: bool | None = None, ): max_model_len = self.model_config.max_model_len cache_salt = cast(str | None, engine_input.get("cache_salt")) @@ -673,6 +675,7 @@ async def _generate_with_builtin_tools( priority=priority, session_id=session_id, reasoning_parser_kwargs=reasoning_parser_kwargs, + reasoning_ended=reasoning_ended, ) async for res in generator: diff --git a/vllm/envs.py b/vllm/envs.py index 19e3250e2fbe..3775496ab44f 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -246,6 +246,7 @@ VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE: int = 163840 VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS: int = 1 VLLM_ENFORCE_STRICT_TOOL_CALLING: bool = True + VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT: bool = True VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16 VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300 VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS: int = 5 @@ -1844,6 +1845,11 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_ENFORCE_STRICT_TOOL_CALLING": lambda: ( os.getenv("VLLM_ENFORCE_STRICT_TOOL_CALLING", "True").lower() in ("true", "1") ), + # Constrain every `beam` output to its response template's format, also + # without tool or structured output constraints. + "VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT": lambda: ( + os.getenv("VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT", "True").lower() in ("true", "1") + ), # Control the max chunk bytes (in MB) for the rpc message queue. # Object larger than this threshold will be broadcast to worker # processes via zmq. diff --git a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py index f06de6bd0c8c..707e2cdf454f 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py @@ -356,10 +356,12 @@ def kernel( if tidx == 0: for m in cutlass.range_constexpr(M): partials = sm[m, None].load() - gC[m, n_idx] = partials.reduce( - cute.ReductionOp.ADD, - init_val=cutlass.Float32(0.0), - reduction_profile=0, - ) + gC[m, n_idx] = cutlass.Float32( + partials.reduce( + cute.ReductionOp.ADD, + init_val=cutlass.Float32(0.0), + reduction_profile=0, + ) + ).to(gC.element_type) if const_expr(self.use_pdl): cute.arch.griddepcontrol_launch_dependents() diff --git a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py index b44a5f6e2c94..2ae3d599227d 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py @@ -47,7 +47,7 @@ class LLBf16SplitK: :param ab_dtype: Element type for A and B operands. :param acc_dtype: Accumulator type used by MMA and reductions. - :param out_dtype: Output element type. The public wrapper uses fp32. + :param out_dtype: Output element type (fp32 or bf16). :param tile_n: CTA tile size in N. M is fixed to 16 for router batches. :type tile_n: int :param tile_k: K tile size staged through shared memory. @@ -580,6 +580,6 @@ def kernel( reduction_profile=0, ) ) - gC[local_coord] = acc + gC[local_coord] = cutlass.Float32(acc).to(gC.element_type) cute.arch.sync_threads() diff --git a/vllm/model_executor/kernels/linear/cute_dsl/_ll_router_common.py b/vllm/model_executor/kernels/linear/cute_dsl/_ll_router_common.py index 23fcefdac9e5..7770c0f8de48 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/_ll_router_common.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/_ll_router_common.py @@ -76,8 +76,8 @@ def make_fake_gemm_tensors( a_dtype: torch.dtype, b_dtype: torch.dtype, divisibility: int, + output_dtype: torch.dtype = torch.float32, ): - from cutlass import Float32 from quack.compile_utils import make_fake_tensor hidden_states = make_fake_tensor( @@ -86,7 +86,7 @@ def make_fake_gemm_tensors( router_weight = make_fake_tensor( cutlass_dtype(b_dtype), (N, K), divisibility=divisibility ) - output = make_fake_tensor(Float32, (M, N), divisibility=1) + output = make_fake_tensor(cutlass_dtype(output_dtype), (M, N), divisibility=1) return hidden_states, router_weight, output @@ -97,6 +97,7 @@ def validate_common_gemm_inputs( *, op_name: str, k_multiple: int | None = None, + output_dtypes: tuple[torch.dtype, ...] = (torch.float32,), ) -> None: if hidden_states.dim() != 2 or router_weight.dim() != 2: raise ValueError("hidden_states and router_weight must be 2D tensors") @@ -106,8 +107,11 @@ def validate_common_gemm_inputs( raise ValueError( "hidden_states and router_weight must be on the same CUDA device" ) - if output_dtype != torch.float32: - raise ValueError(f"{op_name} only supports output_dtype=torch.float32") + if output_dtype not in output_dtypes: + raise ValueError( + f"{op_name} only supports output_dtype=" + + " or ".join(str(dtype) for dtype in output_dtypes) + ) if hidden_states.shape[1] != router_weight.shape[1]: raise ValueError( "hidden_states and router_weight must have matching K dimensions" diff --git a/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py b/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py index 95cb78c6a11c..9d1573bb80e4 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py @@ -47,6 +47,8 @@ def is_available() -> bool: **{M: (8, 5) for M in (5, 8)}, 9: (8, 2), }, + (5120, 1792): {M: (2, 4) for M in range(5, 17)}, + (5120, 3584): {M: (2, 4) for M in range(5, 17)}, (7168, 256): {14: (8, 2)}, (6144, 256): {M: (8, 2) for M in (9, 12, 16)}, (7168, 384): {M: (7, 5) for M in (13, 16)}, @@ -101,39 +103,52 @@ class CompileKey: bs: int = 0 split_k: int = 0 num_stages: int = 0 + output_dtype: torch.dtype = torch.float32 def __init__(self, *, prefetch_pdl_weights: bool = False) -> None: self._prefetch_pdl_weights = prefetch_pdl_weights - # Dot-prod: keyed on (M, K, bs), because M and K are Constexpr. - self._compiled_cache: dict[tuple[int, int, int], Any] = {} - # Split-K: keyed on (split_k, num_stages), fully shape-dynamic. - self._splitk_cache: dict[tuple[int, int], Any] = {} - - def dispatch(self, *, M: int, K: int, N: int) -> CompileKey: + # M and K are compile-time constants for dot-product. + self._compiled_cache: dict[tuple[int, int, int, torch.dtype], Any] = {} + # Split-K is shape-dynamic; both caches distinguish output dtype. + self._splitk_cache: dict[tuple[int, int, torch.dtype], Any] = {} + + def dispatch( + self, *, M: int, K: int, N: int, output_dtype: torch.dtype = torch.float32 + ) -> CompileKey: tuned_bs, tuned_splitk = _arch_tuned_configs() if M <= _DEFAULT_DOTPROD_MAX_M or K < 2048: bs = tuned_bs.get((K, N), {}).get(M, _DEFAULT_DOTPROD_BS) - return self.CompileKey(backend="dotprod", M=M, K=K, bs=bs) + return self.CompileKey( + backend="dotprod", M=M, K=K, bs=bs, output_dtype=output_dtype + ) split_k, num_stages = tuned_splitk.get((K, N), {}).get( M, _DEFAULT_SPLITK_CONFIG ) - return self.CompileKey(backend="splitk", split_k=split_k, num_stages=num_stages) + return self.CompileKey( + backend="splitk", + split_k=split_k, + num_stages=num_stages, + output_dtype=output_dtype, + ) def get_warmup_keys( self, *, shapes: Iterable[tuple[int, int]], m_values: Iterable[int], + output_dtype: torch.dtype = torch.float32, ) -> list[CompileKey]: return list( dict.fromkeys( - self.dispatch(M=M, K=K, N=N) for K, N in shapes for M in m_values + self.dispatch(M=M, K=K, N=N, output_dtype=output_dtype) + for K, N in shapes + for M in m_values ) ) @staticmethod - def _fake_gemm_tensors(*, M, K, N, divisibility: int): + def _fake_gemm_tensors(*, M, K, N, divisibility: int, output_dtype: torch.dtype): return make_fake_gemm_tensors( M=M, K=K, @@ -141,10 +156,13 @@ def _fake_gemm_tensors(*, M, K, N, divisibility: int): a_dtype=torch.bfloat16, b_dtype=torch.bfloat16, divisibility=divisibility, + output_dtype=output_dtype, ) def _compile_splitk(self, compile_key: CompileKey) -> None: cute, _ = cute_context() + from cutlass import BFloat16, Float32 + from ._ll_bf16_splitk import LLBf16SplitK hidden_states, router_weight, output = self._fake_gemm_tensors( @@ -152,8 +170,12 @@ def _compile_splitk(self, compile_key: CompileKey) -> None: K=cute.sym_int(), N=cute.sym_int(), divisibility=8, + output_dtype=compile_key.output_dtype, ) gemm = LLBf16SplitK( + out_dtype=BFloat16 + if compile_key.output_dtype == torch.bfloat16 + else Float32, tile_n=16, tile_k=256, num_stages=compile_key.num_stages, @@ -169,7 +191,9 @@ def _compile_splitk(self, compile_key: CompileKey) -> None: current_cuda_stream(), options="--enable-tvm-ffi", ) - self._splitk_cache[(compile_key.split_k, compile_key.num_stages)] = compiled + self._splitk_cache[ + (compile_key.split_k, compile_key.num_stages, compile_key.output_dtype) + ] = compiled logger.debug( "Compiled ll_bf16_splitk: sk=%d ns=%d", compile_key.split_k, @@ -187,6 +211,7 @@ def _compile_dotprod(self, compile_key: CompileKey) -> None: K=compile_key.K, N=N, divisibility=stride_divisibility, + output_dtype=compile_key.output_dtype, ) gemm = LLBf16Dotprod( k=compile_key.K, @@ -205,7 +230,9 @@ def _compile_dotprod(self, compile_key: CompileKey) -> None: current_cuda_stream(), options="--enable-tvm-ffi --ptxas-options -maxrregcount=64", ) - self._compiled_cache[(compile_key.M, compile_key.K, compile_key.bs)] = compiled + self._compiled_cache[ + (compile_key.M, compile_key.K, compile_key.bs, compile_key.output_dtype) + ] = compiled logger.debug( "Compiled ll_bf16_dotprod: M=%d, K=%d, bs=%d", compile_key.M, @@ -215,12 +242,21 @@ def _compile_dotprod(self, compile_key: CompileKey) -> None: def compile(self, compile_key: CompileKey) -> None: if compile_key.backend == "splitk": - splitk_cache_key = (compile_key.split_k, compile_key.num_stages) + splitk_cache_key = ( + compile_key.split_k, + compile_key.num_stages, + compile_key.output_dtype, + ) if splitk_cache_key not in self._splitk_cache: self._compile_splitk(compile_key) return - dotprod_cache_key = (compile_key.M, compile_key.K, compile_key.bs) + dotprod_cache_key = ( + compile_key.M, + compile_key.K, + compile_key.bs, + compile_key.output_dtype, + ) if dotprod_cache_key not in self._compiled_cache: self._compile_dotprod(compile_key) @@ -229,8 +265,11 @@ def warmup( *, shapes: Iterable[tuple[int, int]], m_values: Iterable[int], + output_dtype: torch.dtype = torch.float32, ) -> None: - for compile_key in self.get_warmup_keys(shapes=shapes, m_values=m_values): + for compile_key in self.get_warmup_keys( + shapes=shapes, m_values=m_values, output_dtype=output_dtype + ): self.compile(compile_key) @staticmethod @@ -251,6 +290,7 @@ def _validate_inputs( output_dtype, op_name="ll_bf16_gemm", k_multiple=8, + output_dtypes=(torch.float32, torch.bfloat16), ) def __call__( @@ -258,19 +298,28 @@ def __call__( hidden_states: torch.Tensor, # [M, K] bf16 router_weight: torch.Tensor, # [N, K] bf16 output_dtype: torch.dtype = torch.float32, - ) -> torch.Tensor: # [M, N] fp32 + ) -> torch.Tensor: # [M, N] self._validate_inputs(hidden_states, router_weight, output_dtype) M, K = hidden_states.shape N = router_weight.shape[0] - compile_key = self.dispatch(M=M, K=K, N=N) + compile_key = self.dispatch(M=M, K=K, N=N, output_dtype=output_dtype) if compile_key.backend == "splitk": - splitk_cache_key = (compile_key.split_k, compile_key.num_stages) + splitk_cache_key = ( + compile_key.split_k, + compile_key.num_stages, + compile_key.output_dtype, + ) if splitk_cache_key not in self._splitk_cache: self.compile(compile_key) kernel = self._splitk_cache[splitk_cache_key] else: - dotprod_cache_key = (compile_key.M, compile_key.K, compile_key.bs) + dotprod_cache_key = ( + compile_key.M, + compile_key.K, + compile_key.bs, + compile_key.output_dtype, + ) if dotprod_cache_key not in self._compiled_cache: self.compile(compile_key) kernel = self._compiled_cache[dotprod_cache_key] diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index e8bd822becab..557a810a34d8 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -8,7 +8,7 @@ import vllm.envs as envs from vllm.config import ParallelConfig, SchedulerConfig -from vllm.config.kernel import MoEBackend +from vllm.config.kernel import FLASHINFER_MOE_EP_CUTEDSL, MoEBackend from vllm.distributed import get_dp_group, get_pcp_group, get_tensor_model_parallel_rank from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import MoEActivation @@ -1315,6 +1315,7 @@ class FusedMoEConfig: # Requested through `defer_moe_finalize()`, read through # `should_defer_moe_finalize()`. _defer_moe_finalize: bool = field(default=False, init=False) + _defer_moe_finalize_allow_ep: bool = field(default=False, init=False) # Most tokens a deferred call covers: the consumer's capacity, lowered by # experts that would split a larger call across kernel launches. Negative # means unbounded. @@ -1474,34 +1475,58 @@ def use_deferred_moe_finalize(self) -> bool: Evaluated on read rather than in ``__post_init__`` because deferral is requested after construction, like ``skip_final_all_reduce``. """ - # The consumer fuses a TP all-reduce. Other parallel modes require a - # combine or reduce-scatter after the experts and cannot defer it, and - # the consumer has no way to strip hidden-dim padding from GEMM2 rows. + from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep import ( + supports_flashinfer_moe_ep_deferred_finalize, + ) + + if ( + self.moe_backend == FLASHINFER_MOE_EP_CUTEDSL + and self._defer_moe_finalize_allow_ep + ): + # FLASHINFER_MOE_EP_CUTEDSL supports deferred moe finalize for EP + parallelism_supported = ( + supports_flashinfer_moe_ep_deferred_finalize() + and self.tp_size == 1 + and self.ep_size > 1 + and not self.is_sequence_parallel + ) + else: + parallelism_supported = ( + self.dp_size == 1 + and not self.is_sequence_parallel + and ( + (self.tp_size > 1 and self.ep_size == 1) + or ( + self._defer_moe_finalize_allow_ep + and self.tp_size == 1 + and self.ep_size > 1 + ) + ) + ) return ( self._defer_moe_finalize - and self.tp_size > 1 - and self.dp_size == 1 - and self.ep_size == 1 + and parallelism_supported and self.pcp_size == 1 - and not self.is_sequence_parallel and self.hidden_dim == self.hidden_dim_unpadded ) - def defer_moe_finalize(self, max_num_tokens: int = -1) -> None: + def defer_moe_finalize( + self, max_num_tokens: int = -1, *, allow_ep: bool = False + ) -> None: """Ask the experts to leave the top-k reduction to the layer's consumer. A layer whose consumer can fuse the top-k reduction (e.g. into its TP all-reduce) calls this once while the model is built, and only when its - experts are TRTLLM-Gen ones that can stop after GEMM2 (the - ``do_finalize=False`` path). Other experts ignore the request and - always finalize, which ``should_defer_moe_finalize`` can't see, so the + experts support ``do_finalize=False`` (TRTLLM-Gen or MegaMoE). + Other experts ignore the request and always finalize, which + ``should_defer_moe_finalize`` can't see, so the layer checks the quant method's ``experts_cls`` first, as Kimi-K3 does. From then on: - The experts return an ``UnfinalizedMoEOutput`` instead of finalized states for every call ``should_defer_moe_finalize`` accepts. - ``should_defer_moe_finalize(num_tokens)`` is the answer for a call: - it requires a TP-only deployment without hidden-dim padding + it requires TP or explicitly supported EP, no padding (``use_deferred_moe_finalize``), a non-empty call and at most ``defer_moe_finalize_max_num_tokens`` tokens. The cap starts at ``max_num_tokens`` and only ever goes down: experts that would split a @@ -1514,12 +1539,19 @@ def defer_moe_finalize(self, max_num_tokens: int = -1) -> None: consumer then owns the top-k reduction, the shared-expert add and the all-reduce. + MegaMoE fixes deferral when its workspace is created. Its consumer must + also enter the producer on empty DP ranks to participate in communication, + then finalize only the local rows. + Args: max_num_tokens: Most tokens per call the consumer can take in deferred form. Negative means no limit of its own. + allow_ep: The consumer handles replicated-token EP reduction or + local DP finalization / SP gathering of MegaMoE outputs. """ self._defer_moe_finalize = True + self._defer_moe_finalize_allow_ep = allow_ep self.limit_deferred_moe_finalize(max_num_tokens) def should_defer_moe_finalize(self, num_tokens: int) -> bool: diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_moe_ep.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_moe_ep.py index 8fbfcd917839..b57de22c178b 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_moe_ep.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_moe_ep.py @@ -28,6 +28,7 @@ supports_current_device, validate_flashinfer_moe_ep_layer, ) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) @@ -177,12 +178,14 @@ def apply( workspace2: torch.Tensor, expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, - ) -> None: + ) -> UnfinalizedMoEOutput | None: if self._adapter is None: raise RuntimeError("FlashInfer MoE-EP weights have not been processed") - output.copy_( - self._adapter(hidden_states, topk_ids.to(torch.int32), topk_weights) - ) + result = self._adapter(hidden_states, topk_ids.to(torch.int32), topk_weights) + if isinstance(result, UnfinalizedMoEOutput): + return result + output.copy_(result) + return None def epilogue_from_quant_config( 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 63345a417438..0a0d97f1571a 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py @@ -17,7 +17,12 @@ FLASHINFER_MOE_EP_DEEP_GEMM, ) from vllm.distributed import get_ep_group +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.fused_moe.config import RoutingMethodType +from vllm.model_executor.layers.fused_moe.moe_output import ( + UnfinalizedMoEOutput, + convert_flashinfer_moe_output, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kMxfp4Static, @@ -25,6 +30,7 @@ ) from vllm.platforms import current_platform from vllm.utils.deep_gemm import _import_deep_gemm +from vllm.utils.math_utils import cdiv if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig @@ -272,6 +278,13 @@ def _load_flashinfer_moe_ep_api() -> _FlashInferMoeEpApi: return _FlashInferMoeEpApi() +def supports_flashinfer_moe_ep_deferred_finalize() -> bool: + api = _load_flashinfer_moe_ep_api() + return hasattr(api.MoEEpMegaLayer, "forward_unfinalized") and ( + "do_finalize" in api.Nvfp4CutedslMegaMoeConfig.__dataclass_fields__ + ) + + def _expose_deep_gemm_to_flashinfer() -> None: deep_gemm = _import_deep_gemm() if deep_gemm is None: @@ -294,6 +307,17 @@ def __init__( if spec.kernel == "deep_gemm": _expose_deep_gemm_to_flashinfer() api = _load_flashinfer_moe_ep_api() + # MegaMoE's workspace mode is fixed, including empty collective calls. + self._defer_finalize = moe.use_deferred_moe_finalize + if self._defer_finalize and ( + spec.kernel != "cutedsl" + or apply_topk_in_fc1 + or not supports_flashinfer_moe_ep_deferred_finalize() + ): + raise ValueError( + "Deferred finalization requires FlashInfer CuTeDSL " + "forward_unfinalized with output-side routing weights" + ) if epilogue is None: epilogue = FlashInferMoeEpEpilogue() ep_group = get_ep_group() @@ -329,6 +353,7 @@ def __init__( fc2_alpha=None, fc1_norm_const=None, knobs=None, + **({"do_finalize": False} if self._defer_finalize else {}), ) else: megakernel = api.DeepGemmMegaMoeConfig( @@ -357,6 +382,23 @@ def __init__( self._top_k = moe.experts_per_token self._epilogue: FlashInferMoeEpEpilogue | None = epilogue self._return_workspace_view = moe.routing_method is RoutingMethodType.DeepseekV4 + self._sp_size = moe.sp_size + self._decode_max_num_tokens = 0 + self._decode_workspace: Any | None = None + self._is_separate_workspace_supported = ( + spec.kernel == "cutedsl" + and moe.moe_parallel_config.pcp_size == 1 + and hasattr(self._mega_layer, "create_workspace") + ) + if self._is_separate_workspace_supported: + vllm_config = get_current_vllm_config() + decode_tokens = cdiv( + vllm_config.scheduler_config.max_num_seqs + * (1 + vllm_config.num_speculative_tokens), + self._sp_size, + ) + if decode_tokens < self._max_num_tokens: + self._decode_max_num_tokens = decode_tokens @property def can_overlap_shared_experts(self) -> bool: @@ -374,6 +416,28 @@ def topk_indices_dtype(self) -> torch.dtype: def is_monolithic(self) -> bool: return False + def _select_workspace(self, num_tokens: int) -> Any | None: + if not is_forward_context_available(): + return None + metadata = get_forward_context().dp_metadata + if metadata is None: + return None + # Every EP rank must use the same capacity: peer buffer offsets depend + # on it. Local token counts can differ, including on idle ranks. + if metadata.local_sizes is not None: + max_tokens = max(metadata.local_sizes) + else: + max_tokens = cdiv( + int(metadata.num_tokens_across_dp_cpu.max().item()), self._sp_size + ) + if ( + self._decode_workspace is not None + and max_tokens <= self._decode_max_num_tokens + ): + assert num_tokens <= self._decode_max_num_tokens + return self._decode_workspace + return None + def _tensors( self, hidden_states: torch.Tensor, @@ -395,7 +459,7 @@ def __call__( hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: if self._mega_layer is None: raise RuntimeError("FlashInfer MoE-EP layer was destroyed") if hidden_states.shape[0] > self._max_num_tokens: @@ -404,11 +468,24 @@ def __call__( f"the workspace supports at most {self._max_num_tokens}" ) tensors = self._tensors(hidden_states, topk_ids, topk_weights) + workspace = self._select_workspace(hidden_states.shape[0]) + workspace_kwargs = {"workspace": workspace} if workspace is not None else {} + if self._defer_finalize: + unfinalized = convert_flashinfer_moe_output( + list(self._mega_layer.forward_unfinalized(tensors, **workspace_kwargs)), + do_finalize=False, + num_tokens=hidden_states.shape[0], + top_k=topk_ids.shape[1], + ) + assert isinstance(unfinalized, UnfinalizedMoEOutput) + return unfinalized if self._return_workspace_view and getattr( self._mega_layer, "supports_output_view", False ): - return self._mega_layer.forward(tensors, return_workspace_view=True) - return self._mega_layer(tensors) + return self._mega_layer.forward( + tensors, return_workspace_view=True, **workspace_kwargs + ) + return self._mega_layer(tensors, **workspace_kwargs) @torch.inference_mode() def warmup(self) -> None: @@ -431,7 +508,14 @@ def warmup(self) -> None: dtype=torch.float32, device=self._device, ) - self._mega_layer.warmup(self._tensors(hidden_states, topk_ids, topk_weights)) + tensors = self._tensors(hidden_states, topk_ids, topk_weights) + self._mega_layer.warmup(tensors) + if self._is_separate_workspace_supported and self._decode_max_num_tokens: + if self._decode_workspace is None: + self._decode_workspace = self._mega_layer.create_workspace( + self._decode_max_num_tokens + ) + self._mega_layer.warmup(tensors, workspace=self._decode_workspace) def kernel_weights( self, @@ -463,6 +547,7 @@ def destroy(self) -> None: return mega_layer.destroy() self._mega_layer = None + self._decode_workspace = None self._epilogue = None def __del__(self) -> None: diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index fccbd1a9d83b..501222022129 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1116,7 +1116,7 @@ def _allocate_buffers( Returns the workspace tensors, output tensor, and whether the output uses a prepare/finalize-provided allocation. """ - assert M_full > 0 and M_chunk > 0 + assert M_full >= 0 and M_chunk >= 0 workspace_dtype = self.fused_experts.workspace_dtype(out_dtype) @@ -1158,7 +1158,8 @@ def _allocate_buffers( # externally aliased. The workspace manager profiles and locks a # process-wide high-water mark; later graph shapes may need this # storage for workspace13. - max_shape_size = max(prod(workspace13_shape), prod(fused_out_shape)) + # Idle collective ranks still need the workspace manager initialized. + max_shape_size = max(1, prod(workspace13_shape), prod(fused_out_shape)) if current_platform.is_cpu(): # Every CPU FusedMoEExpertsModular kernel reports zero-sized @@ -1314,13 +1315,9 @@ def _fused_experts( a1q, w1, w2, topk_ids ) - # This happens when none of the tokens from the all2all reach this - # EP rank. Also, note that this is only relevant for CUDAGraph - # incompatible all2all kernels like the DeepEP high-throughput - # kernels. CUDAGraph compatible all2all kernels like the DeepEP - # low-latency kernels are always batched and can never run into - # the tensor.numel() == 0 case. - if M_full == 0: + # After dispatch, an EP rank may have no work. With passthrough, + # dispatch happens inside the experts, so idle ranks must participate. + if M_full == 0 and not self.moe_parallel_config.use_passthrough_all2all: return torch.empty_like(a1q, dtype=in_dtype) workspace13, workspace2, fused_out, fused_out_is_specialized = ( diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/passthrough.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/passthrough.py index a6c541d2b9b6..09bf0c0ee343 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/passthrough.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/passthrough.py @@ -19,13 +19,14 @@ class PassThroughPrepareAndFinalize(MoEPrepareAndFinalizeNoDPEPModular): """Hands activations and routing straight to experts that own their EP communication.""" + def post_init_setup(self, fused_experts: mk.FusedMoEExperts) -> None: + self._defer_output = fused_experts.moe_config.use_deferred_moe_finalize + def output_is_reduced(self) -> bool: return True def supports_deferred_moe_finalize(self) -> bool: - # The experts apply the top-k weights and combine internally; there - # is no finalize left for a consumer to take over. - return False + return getattr(self, "_defer_output", False) def prepare( self, diff --git a/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py b/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py index 2190eea1572d..395d7d68ea93 100644 --- a/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py +++ b/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py @@ -517,6 +517,7 @@ def splitk_reduce_triton(partials: torch.Tensor, out: torch.Tensor): 256: (128, 32), 512: (128, 19), }, + (5120, 256): {64: (64, 64)}, # GB200 (6144, 128): { # MiniMax-M3 64: (32, 64), 128: (64, 64), @@ -533,7 +534,7 @@ def _pick_tile_config(N: int, K: int, M: int, num_sms: int) -> tuple[int, int]: if 32 < N <= 512: token_bucket = triton.next_power_of_2(N) overrides = _TILE_CONFIG_OVERRIDES.get((K, M)) - if overrides is not None: + if overrides is not None and token_bucket in overrides: BN, split_k = overrides[token_bucket] return BN, min(split_k, k_tiles) @@ -609,6 +610,8 @@ def _bf16x3_router_gemm(X: torch.Tensor, W: torch.Tensor) -> torch.Tensor: """Return ``X @ W.T`` using the SM100 BF16x3 router GEMM kernel.""" N, K = X.shape M, _ = W.shape + if N == 0: + return X.new_empty((N, M), dtype=torch.float32) num_sms = torch.cuda.get_device_properties(X.device).multi_processor_count BN, split_k = _pick_tile_config(N, K, M, num_sms) diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index c6b1da890bbe..7d14f105c996 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -24,8 +24,8 @@ class GateLinear(ReplicatedLinear): 1. cuteDSL ll_bf16_gemm (SM90+, M<=16, bf16 in, fp32 out, any shape w/ K divisible by 8) - 2. fp32 specialized kernel (bf16/fp32 in, fp32 out, M<=32; SM90+ - model-specific shapes, gfx950 any shape) + 2. fp32 specialized kernel (bf16/fp32 in, fp32 out; SM90+ + model-specific shapes and batch limits, gfx950 any shape with M<=32) 3. cuteDSL ll_fp32w_gemm (SM90+, M<=32, bf16/fp32 in, fp32 out, any shape) 4. bf16x3 CuteDSL kernel (SM100, bf16 in, fp32 weight) @@ -46,9 +46,8 @@ class GateLinear(ReplicatedLinear): split-K from the token count and skip the batch-invariant linear path. """ - # (hidden_size, num_experts) pairs with an instantiated fp32 kernel: - # (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3 - FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128)} + # (hidden_size, num_experts) pairs with an instantiated fp32 kernel. + FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128), (5120, 256)} FP32_MAX_TOKENS = 32 # Largest batch for which tier 1 beats the cuBLAS epilogue. Subclasses # whose router shape stays ahead further up the range may raise it. @@ -454,9 +453,9 @@ def fp32_router_gemm_dispatch_impl( weight: torch.Tensor, allow_bf16x3_router_gemm: bool, ) -> torch.Tensor: - """Dynamically run fp32 specialized gemm if num_tokens <= FP32_MAX_TOKENS, - otherwise optionally run the experimental BF16x3 kernel for medium/large - SM100 router batches, then fall back to F.linear. + """Run the FP32 kernel for small batches, with shape-specific limits. + + Larger SM100 BF16 batches use BF16x3; other inputs fall back to F.linear. This must be wrapped in a custom op because our torch.compile integration does not support runtime dispatching on num_tokens. """ @@ -473,7 +472,15 @@ def fp32_router_gemm_dispatch_impl( if weight.dtype == torch.bfloat16: return torch.mm(x, weight.T, out_dtype=torch.float32) return torch.nn.functional.linear(x.float(), weight) - if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS: + max_tokens = _FP32_ROUTER_GEMM_MAX_TOKENS + if ( + allow_bf16x3_router_gemm + and x.dtype == torch.bfloat16 + and weight.shape == (256, 5120) + ): + # GB200: four token groups win through 48 for multiples of eight. + max_tokens = 48 if x.shape[0] % 8 == 0 else 16 + if x.shape[0] <= max_tokens: return ops.fp32_router_gemm(x, weight) if allow_bf16x3_router_gemm and x.dtype == torch.bfloat16: diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 485ea952fccb..1e734fc7aa69 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -320,6 +320,8 @@ def forward( class DFlashQwen3DecoderLayer(nn.Module): + attention_cls = DFlashQwen3Attention + def __init__( self, vllm_config: VllmConfig, @@ -354,7 +356,7 @@ def __init__( # distilled the other way has to say so here. is_neox_style = getattr(config, "is_neox_style", True) - self.self_attn = DFlashQwen3Attention( + self.self_attn = self.attention_cls( hidden_size=self.hidden_size, num_heads=config.num_attention_heads, max_position=config.max_position_embeddings, diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 23e67c9f506c..b8fe347209bd 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -78,6 +78,7 @@ "BailingMoeV2ForCausalLM": ("bailing_moe", "BailingMoeV2ForCausalLM"), "BailingMoeV2_5ForCausalLM": ("bailing_moe_linear", "BailingMoeV25ForCausalLM"), "BailingMoeV3ForCausalLM": ("bailing_moe_v3", "BailingMoeV3ForCausalLM"), + "BeamForCausalLM": ("vllm.models.beam", "BeamForCausalLM"), "BloomForCausalLM": ("bloom", "BloomForCausalLM"), "ChatGLMModel": ("chatglm", "ChatGLMForCausalLM"), "ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"), @@ -636,6 +637,7 @@ "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), "LiLiCorrDraftModel": ("lilicorr", "LiLiCorrForCausalLM"), "DFlash2DraftModel": ("qwen3_dflash2", "DFlash2Qwen3ForCausalLM"), + "DFlash2BeamDraftModel": ("vllm.models.beam", "DFlash2BeamForCausalLM"), # Muse Glimmer's DFlash draft head, reusing the generic qwen3_dflash # implementation. EAGLEConfig rewrites a dflash draft's architecture to # DFlash{arch} unless it already starts or ends with "DFlash" (see diff --git a/vllm/models/beam/__init__.py b/vllm/models/beam/__init__.py new file mode 100644 index 000000000000..faed0717d1a7 --- /dev/null +++ b/vllm/models/beam/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Beam model entry point.""" + +from .nvidia.dflash2 import DFlash2BeamForCausalLM +from .nvidia.model import BeamForCausalLM + +__all__ = ["BeamForCausalLM", "DFlash2BeamForCausalLM"] diff --git a/vllm/models/beam/nvidia/__init__.py b/vllm/models/beam/nvidia/__init__.py new file mode 100644 index 000000000000..d2d40aaa7d8d --- /dev/null +++ b/vllm/models/beam/nvidia/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVIDIA-specific Beam model implementation.""" diff --git a/vllm/models/beam/nvidia/activation.py b/vllm/models/beam/nvidia/activation.py new file mode 100644 index 000000000000..9afc18a0b464 --- /dev/null +++ b/vllm/models/beam/nvidia/activation.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.triton_utils import tl, triton +from vllm.triton_utils import tldevice as libdevice + + +@triton.jit(do_not_specialize=["size"]) +def _silu_mul_kernel(X, Y, size, width: tl.constexpr, BLOCK: tl.constexpr): + i = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + src = i // width * (2 * width) + i % width + gate = tl.load(X + src, i < size, other=0).to(tl.float32) + up = tl.load(X + src + width, i < size, other=0).to(tl.float32) + act = tl.div_rn(gate, 1.0 + libdevice.exp(-gate)) + tl.store(Y + i, act * up, i < size) + + +def beam_silu_mul(x: torch.Tensor) -> torch.Tensor: + """Compute SiLU and multiplication in FP32, rounding only the product.""" + width = x.shape[-1] // 2 + output = torch.empty((*x.shape[:-1], width), dtype=x.dtype, device=x.device) + if output.numel(): + _silu_mul_kernel[(triton.cdiv(output.numel(), 1024),)]( + x, output, output.numel(), width, 1024, enable_fp_fusion=False + ) + return output + + +@triton.jit +def _softcap_logits_kernel( + X, + WIDTH: tl.constexpr, + STRIDE: tl.constexpr, + CAP: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + x = tl.load(X + row * STRIDE + col, col < WIDTH, 0).to(tl.float32) + y = libdevice.tanh(x * (1.0 / CAP)) * CAP + tl.store(X + row * STRIDE + col, y, col < WIDTH) + + +def beam_softcap_logits(logits: torch.Tensor, cap: float) -> torch.Tensor: + """Softcap in FP32 and round back into the original logits buffer.""" + assert logits.ndim == 2 and logits.stride(1) == 1 + if logits.numel(): + _softcap_logits_kernel[(logits.shape[0], triton.cdiv(logits.shape[1], 1024))]( + logits, + logits.shape[1], + logits.stride(0), + cap, + 1024, + enable_fp_fusion=False, + enable_reflect_ftz=False, + ) + return logits diff --git a/vllm/models/beam/nvidia/allreduce.cu b/vllm/models/beam/nvidia/allreduce.cu new file mode 100644 index 000000000000..e16b3496516e --- /dev/null +++ b/vllm/models/beam/nvidia/allreduce.cu @@ -0,0 +1,625 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Communication helpers and launch geometry from NVIDIA/FlashInfer's MNNVL +// all-reduce (Apache-2.0); Beam supplies the normalization epilogue. + +#include "flashinfer/comm/trtllm_mnnvl_allreduce.cuh" +#include "tvm_ffi_utils.h" + +namespace beam { +namespace ar = flashinfer::trtllm_mnnvl_allreduce; +namespace cg = cooperative_groups; +using T = __nv_bfloat16; +using Vec = ar::utils::PackedVec; + +struct Params { + ar::AllReduceKernelParams ar; + T const* nextWeight; + float nextEps; + float xScale; + float rScale; + bool postNorm; + bool nextNorm; + bool scaled; + T const* experts; + void const* expertWeights; + int const* permutation; + int topK; + bool floatWeights; + float routedScale; +}; + +template +__device__ Vec loadInput(Params const& p, int token, int packedIdx) { + Vec value; + if constexpr (!Finalize) { + value.packed = ar::utils::loadPacked( + p.ar.shardPtr + int64_t(token) * p.ar.tokenDim + packedIdx * 8); + return value; + } + float sum[8]{}; + for (int base = 0; base < p.topK; base += 8) { + int rows[8]; + float weights[8]; +#pragma unroll + for (int k = 0; k < 8; ++k) { + int index = token * p.topK + base + k; + rows[k] = base + k < p.topK ? p.permutation[index] : -1; + weights[k] = rows[k] < 0 ? 0.f + : p.floatWeights + ? static_cast(p.expertWeights)[index] + : ar::utils::toFloat( + static_cast(p.expertWeights)[index]); + } +#pragma unroll + for (int k = 0; k < 8; ++k) { + Vec expert{}; + if (rows[k] >= 0) + expert.packed = ar::utils::loadPacked( + p.experts + int64_t(rows[k]) * p.ar.tokenDim + packedIdx * 8); +#pragma unroll + for (int i = 0; i < 8; ++i) + sum[i] = __fadd_rn( + sum[i], + __fmul_rn(ar::utils::toFloat(expert.elements[i]), weights[k])); + } + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + T rounded = ar::utils::fromFloat(sum[i]); + value.elements[i] = ar::utils::fromFloat( + __fmul_rn(ar::utils::toFloat(rounded), p.routedScale)); + } + return value; +} + +template +__device__ float invRms(Vec const& value, int dim, float eps) { + float sum = 0.f; +#pragma unroll + for (int i = 0; i < 8; ++i) { + float v = ar::utils::toFloat(value.elements[i]); + sum = __fadd_rn(sum, __fmul_rn(v, v)); + } + sum = ar::utils::blockReduceSum(sum); + if constexpr (UseCGA) { + auto cluster = cg::this_cluster(); + __shared__ float partial[8]; + if (threadIdx.x < cluster.num_blocks()) { + cluster.map_shared_rank(partial, threadIdx.x)[cluster.block_rank()] = sum; + } + cluster.sync(); + sum = 0.f; + for (int i = 0; i < cluster.num_blocks(); ++i) + sum = __fadd_rn(sum, partial[i]); + } + __syncthreads(); + return rsqrtf(sum / dim + eps); +} + +template +__device__ void epilogue(Params const& p, Vec value, int packedIdx, int offset, + Vec postWeight = {}, Vec residual = {}, + Vec nextWeight = {}) { + if (p.postNorm) { + float rstd = invRms(value, p.ar.tokenDim, p.ar.epsilon); + Vec weight = postWeight; + if constexpr (!Preloaded) { + if (p.ar.gammaPtr) + weight.packed = + ar::utils::loadPacked(p.ar.gammaPtr + packedIdx * 8); + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + float v = __fmul_rn(ar::utils::toFloat(value.elements[i]), rstd); + if (p.ar.gammaPtr) + v = __fmul_rn(v, ar::utils::toFloat(weight.elements[i])); + value.elements[i] = ar::utils::fromFloat(v); + } + } + if constexpr (!Preloaded) + residual.packed = + ar::utils::loadPacked(p.ar.residualInPtr + offset); +#pragma unroll + for (int i = 0; i < 8; ++i) { + T x = value.elements[i], r = residual.elements[i]; + if (p.scaled) { + x = ar::utils::fromFloat( + __fmul_rn(ar::utils::toFloat(x), p.xScale)); + r = ar::utils::fromFloat( + __fmul_rn(ar::utils::toFloat(r), p.rScale)); + } + value.elements[i] = ar::utils::fromFloat( + __fadd_rn(ar::utils::toFloat(x), ar::utils::toFloat(r))); + } + *reinterpret_cast(p.ar.prenormedPtr + offset) = value.packed; + if (p.nextNorm) { + float rstd = invRms(value, p.ar.tokenDim, p.nextEps); + Vec weight = nextWeight; + if constexpr (!Preloaded) { + if (p.nextWeight) + weight.packed = + ar::utils::loadPacked(p.nextWeight + packedIdx * 8); + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + float v = __fmul_rn(ar::utils::toFloat(value.elements[i]), rstd); + if (p.nextWeight) + v = __fmul_rn(v, ar::utils::toFloat(weight.elements[i])); + value.elements[i] = ar::utils::fromFloat(v); + } + *reinterpret_cast(p.ar.outputPtr + offset) = value.packed; + } +} + +template +__device__ Vec reduceLocal(Vec value, T* buffer, int token, int dim, + int packedIdx) { + Vec remote[World]; + ar::utils::waitOneshotRemoteRanks( + remote, buffer, token, dim, packedIdx); + return ar::utils::reduceOneshotDeterministic( + remote, value); +} + +template +__global__ void __launch_bounds__(1024) oneshot(Params p) { + int packedIdx = threadIdx.x; + if constexpr (!SmallBatch) packedIdx = cg::this_cluster().thread_rank(); + int token = blockIdx.x; + int offset = token * p.ar.tokenDim + packedIdx * 8; + cudaGridDependencySynchronize(); + ar::utils::LamportFlags flag(p.ar.bufferFlags, 1); + auto multicast = static_cast(flag.getCurLamportBuf(p.ar.mcastPtr, 0)); + auto local = + static_cast(flag.getCurLamportBuf(p.ar.inputPtrs[p.ar.rank], 0)); + Vec value = loadInput(p, token, packedIdx); + ar::utils::sanitizeLamportPayload(value); + reinterpret_cast(multicast + token * p.ar.tokenDim * World + + p.ar.rank * p.ar.tokenDim)[packedIdx] = + value.packed; + Vec postWeight{}, residual{}, nextWeight{}; + if constexpr (SmallBatch) { + // Overlap epilogue inputs with the peer payloads in flight. + if (p.postNorm && p.ar.gammaPtr) + postWeight.packed = + ar::utils::loadPackedVolatile(p.ar.gammaPtr + packedIdx * 8) + .packed; + residual.packed = + ar::utils::loadPackedVolatile(p.ar.residualInPtr + offset) + .packed; + if (p.nextNorm && p.nextWeight) + nextWeight.packed = + ar::utils::loadPackedVolatile(p.nextWeight + packedIdx * 8) + .packed; + } + flag.ctaArrive(); + if constexpr (SmallBatch) { + // This workspace uses at most two stages. Constant indices avoid a + // dynamically indexed array spilling the Lamport state to local memory. + flag.clearDirtyLamportBuf(p.ar.inputPtrs[p.ar.rank], 0); + flag.clearDirtyLamportBuf(p.ar.inputPtrs[p.ar.rank], 1); + } else { + flag.clearDirtyLamportBuf(p.ar.inputPtrs[p.ar.rank], -1); + } +#define REDUCE(R) \ + case R: \ + if constexpr (World > R) \ + value = reduceLocal(value, local, token, p.ar.tokenDim, \ + packedIdx); \ + break + switch (p.ar.rank) { + REDUCE(0); + REDUCE(1); + REDUCE(2); + REDUCE(3); + REDUCE(4); + REDUCE(5); + REDUCE(6); + REDUCE(7); + } +#undef REDUCE + cudaTriggerProgrammaticLaunchCompletion(); + epilogue(p, value, packedIdx, offset, postWeight, + residual, nextWeight); + flag.waitAndUpdate({static_cast(p.ar.numTokens * p.ar.tokenDim * + World * sizeof(T)), + 0, 0, 0}); +} + +template +__global__ void __launch_bounds__(128) finalizeScatter(Params p) { + int packedIdx = blockIdx.y * blockDim.x + threadIdx.x; + int token = blockIdx.x; + int destRank = token % World; + int localToken = token / World; + cudaGridDependencySynchronize(); + ar::utils::LamportFlags flag(p.ar.bufferFlags, 2); + auto local = + static_cast(flag.getCurLamportBuf(p.ar.inputPtrs[p.ar.rank], 0)); + auto dest = + static_cast(flag.getCurLamportBuf(p.ar.inputPtrs[destRank], 0)); + auto broadcast = static_cast(flag.getCurLamportBuf(p.ar.mcastPtr, 1)); + if (packedIdx * 8 < p.ar.tokenDim) { + Vec value = loadInput(p, token, packedIdx); + ar::utils::sanitizeLamportPayload(value); + reinterpret_cast(dest + localToken * p.ar.tokenDim * World + + p.ar.rank * p.ar.tokenDim)[packedIdx] = + value.packed; + } + cudaTriggerProgrammaticLaunchCompletion(); + flag.clearDirtyLamportBuf(p.ar.inputPtrs[p.ar.rank], 0); + if (packedIdx * 8 < p.ar.tokenDim && destRank == p.ar.rank) { + Vec value = + ar::utils::reduceLamportRanksChunked( + local, localToken, p.ar.tokenDim, packedIdx * 8); + ar::utils::sanitizeLamportPayload(value); + reinterpret_cast(broadcast + token * p.ar.tokenDim)[packedIdx] = + value.packed; + } + flag.clearDirtyLamportBuf(p.ar.inputPtrs[p.ar.rank], 1); +} + +template +__global__ void __launch_bounds__(1024) finishTwoshot(Params p) { + int packedIdx = threadIdx.x; + if constexpr (UseCGA) packedIdx = cg::this_cluster().thread_rank(); + int offset = blockIdx.x * p.ar.tokenDim + packedIdx * 8; + cudaTriggerProgrammaticLaunchCompletion(); + ar::utils::LamportFlags flag(p.ar.bufferFlags, 2); + auto input = static_cast(flag.getCurLamportBuf(p.ar.localBufferPtr, 1)); + flag.ctaArrive(); + auto loaded = ar::utils::loadPackedVolatile(input + offset); + while (ar::utils::isLamportDirty(loaded)) { + loaded = ar::utils::loadPackedVolatile(input + offset); + } + Vec value{.packed = loaded.packed}; + epilogue(p, value, packedIdx, offset); + cudaGridDependencySynchronize(); + flag.waitAndUpdate( + {static_cast(flashinfer::round_up(p.ar.numTokens, p.ar.nRanks) * + p.ar.tokenDim * sizeof(T)), + static_cast(p.ar.numTokens * p.ar.tokenDim * sizeof(T)), 0, + 0}); +} + +__device__ Vec finalizeGatherRow(Params const& p, int token, int packedIdx) { + int64_t inputOffset = int64_t(token) * p.topK * p.ar.tokenDim + packedIdx * 8; + auto weights = static_cast(p.expertWeights); + Vec expert; + expert.packed = ar::utils::loadPacked(p.experts + inputOffset); + float weight = weights[int64_t(token) * p.topK]; + float sum[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) + sum[i] = __fmul_rn(ar::utils::toFloat(expert.elements[i]), weight); + for (int k = 1; k < p.topK; ++k) { + expert.packed = ar::utils::loadPacked(p.experts + inputOffset + + int64_t(k) * p.ar.tokenDim); + weight = weights[int64_t(token) * p.topK + k]; +#pragma unroll + for (int i = 0; i < 8; ++i) + sum[i] = + __fmaf_rn(ar::utils::toFloat(expert.elements[i]), weight, sum[i]); + } + Vec value; +#pragma unroll + for (int i = 0; i < 8; ++i) { + T reduced = ar::utils::fromFloat(sum[i]); + value.elements[i] = ar::utils::fromFloat( + __fmul_rn(ar::utils::toFloat(reduced), p.routedScale)); + } + if (p.postNorm) { + float rstd = invRms(value, p.ar.tokenDim, p.ar.epsilon); + Vec weight; + if (p.ar.gammaPtr) + weight.packed = + ar::utils::loadPacked(p.ar.gammaPtr + packedIdx * 8); +#pragma unroll + for (int i = 0; i < 8; ++i) { + float x = __fmul_rn(ar::utils::toFloat(value.elements[i]), rstd); + if (p.ar.gammaPtr) + x = __fmul_rn(x, ar::utils::toFloat(weight.elements[i])); + value.elements[i] = ar::utils::fromFloat(x); + } + } + ar::utils::sanitizeLamportPayload(value); + return value; +} + +__device__ Vec waitGatherRow(T const* buffer, int64_t offset) { + auto loaded = ar::utils::loadPackedVolatile(buffer + offset); + while (ar::utils::isLamportDirty(loaded)) + loaded = ar::utils::loadPackedVolatile(buffer + offset); + return Vec{.packed = loaded.packed}; +} + +template +__global__ void __launch_bounds__(1024) finalizeGatherOneshot(Params p) { + int localTokens = flashinfer::ceil_div(p.ar.numTokens, World); + int localToken = blockIdx.x; + int token = p.ar.rank * localTokens + localToken; + int packedIdx = threadIdx.x; + cudaGridDependencySynchronize(); + ar::utils::LamportFlags flag(p.ar.bufferFlags, 1); + auto broadcast = static_cast(flag.getCurLamportBuf(p.ar.mcastPtr)); + auto local = static_cast(flag.getCurLamportBuf(p.ar.localBufferPtr)); + Vec ownValue{}; + if (token < p.ar.numTokens) + ownValue = finalizeGatherRow(p, localToken, packedIdx); + // Empty ranks publish padding so no peer can outrun their readers. + *reinterpret_cast(broadcast + int64_t(token) * p.ar.tokenDim + + packedIdx * 8) = ownValue.packed; + flag.ctaArrive(); + flag.clearDirtyLamportBuf(p.ar.localBufferPtr, -1); + p.postNorm = false; +#pragma unroll + for (int rank = 0; rank < World; ++rank) { + token = rank * localTokens + localToken; + int offset = token * p.ar.tokenDim + packedIdx * 8; + Vec value = rank == p.ar.rank ? ownValue : waitGatherRow(local, offset); + if (token < p.ar.numTokens) { + epilogue(p, value, packedIdx, offset); + } + } + cudaTriggerProgrammaticLaunchCompletion(); + flag.waitAndUpdate({static_cast(int64_t(localTokens) * World * + p.ar.tokenDim * sizeof(T)), + 0, 0, 0}); +} + +__global__ void __launch_bounds__(1024) finalizeGatherPublish(Params p) { + int localTokens = flashinfer::ceil_div(p.ar.numTokens, p.ar.nRanks); + int localToken = blockIdx.x; + int token = p.ar.rank * localTokens + localToken; + int packedIdx = threadIdx.x; + cudaGridDependencySynchronize(); + ar::utils::LamportFlags flag(p.ar.bufferFlags, 1); + auto broadcast = static_cast(flag.getCurLamportBuf(p.ar.mcastPtr)); + Vec value{}; + if (token < p.ar.numTokens) + value = finalizeGatherRow(p, localToken, packedIdx); + *reinterpret_cast(broadcast + int64_t(token) * p.ar.tokenDim + + packedIdx * 8) = value.packed; + cudaTriggerProgrammaticLaunchCompletion(); + flag.clearDirtyLamportBuf(p.ar.localBufferPtr, -1); +} + +__global__ void __launch_bounds__(1024) finishGatherTwoshot(Params p) { + int packedIdx = threadIdx.x; + int offset = blockIdx.x * p.ar.tokenDim + packedIdx * 8; + cudaTriggerProgrammaticLaunchCompletion(); + ar::utils::LamportFlags flag(p.ar.bufferFlags, 1); + auto local = static_cast(flag.getCurLamportBuf(p.ar.localBufferPtr)); + flag.ctaArrive(); + Vec value = waitGatherRow(local, offset); + p.postNorm = false; + if (blockIdx.x < p.ar.numTokens) epilogue(p, value, packedIdx, offset); + // The producer must finish clearing the previous generation before reuse. + cudaGridDependencySynchronize(); + flag.waitAndUpdate( + {static_cast(int64_t(gridDim.x) * p.ar.tokenDim * sizeof(T)), 0, + 0, 0}); +} + +void checkCuda(cudaError_t status) { + TVM_FFI_ICHECK_EQ(status, cudaSuccess) << cudaGetErrorString(status); +} + +void launch(TensorView input, TensorView residual, + tvm::ffi::Optional postWeight, + tvm::ffi::Optional nextWeight, TensorView hidden, + tvm::ffi::Optional normalized, int64_t pointers, + TensorView flags, int64_t multicast, int64_t local, int64_t world, + int64_t rank, double postEps, double nextEps, double xScale, + double rScale, bool postNorm, bool nextNorm, bool scaled, + bool oneShot, tvm::ffi::Optional experts, + tvm::ffi::Optional expertWeights, + tvm::ffi::Optional permutation, double routedScale, + bool smallBatch) { + ffi::CUDADeviceGuard guard(input.device().device_id); + auto stream = get_stream(input.device()); + TVM_FFI_ICHECK_EQ(encode_dlpack_dtype(input.dtype()), bfloat16_code); + TVM_FFI_ICHECK_EQ(input.size(1) % 8, 0); + Params p{}; + p.ar.outputPtr = normalized.has_value() + ? static_cast(normalized.value().data_ptr()) + : nullptr; + p.ar.prenormedPtr = static_cast(hidden.data_ptr()); + p.ar.shardPtr = static_cast(input.data_ptr()); + p.ar.residualInPtr = static_cast(residual.data_ptr()); + p.ar.gammaPtr = postWeight.has_value() + ? static_cast(postWeight.value().data_ptr()) + : nullptr; + p.nextWeight = nextWeight.has_value() + ? static_cast(nextWeight.value().data_ptr()) + : nullptr; + p.ar.inputPtrs = reinterpret_cast(pointers); + p.ar.mcastPtr = reinterpret_cast(multicast); + p.ar.localBufferPtr = reinterpret_cast(local); + p.ar.bufferFlags = static_cast(flags.data_ptr()); + p.ar.nRanks = world; + p.ar.rank = rank; + p.ar.numTokens = input.size(0); + p.ar.tokenDim = input.size(1); + p.ar.epsilon = postEps; + p.nextEps = nextEps; + p.xScale = xScale; + p.rScale = rScale; + p.postNorm = postNorm; + p.nextNorm = nextNorm; + p.scaled = scaled; + if (experts.has_value()) { + TVM_FFI_ICHECK(expertWeights.has_value() && permutation.has_value()); + p.experts = static_cast(experts.value().data_ptr()); + p.expertWeights = expertWeights.value().data_ptr(); + p.permutation = static_cast(permutation.value().data_ptr()); + p.topK = permutation.value().size(1); + p.floatWeights = + encode_dlpack_dtype(expertWeights.value().dtype()) == float32_code; + p.routedScale = routedScale; + } + auto [threads, clusters, loads] = + ar::utils::adjustGridConfig(p.ar.numTokens, p.ar.tokenDim, 8, true); + TVM_FFI_ICHECK(!smallBatch || (oneShot && !p.experts)); + if (smallBatch) { + clusters = 1; + threads = p.ar.tokenDim / 8; + } + TVM_FFI_ICHECK_EQ(loads, 1); + TVM_FFI_ICHECK_EQ(threads * clusters * 8, p.ar.tokenDim); + cudaLaunchAttribute attrs[2]{}; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + attrs[1].id = cudaLaunchAttributeClusterDimension; + attrs[1].val.clusterDim = {1, static_cast(clusters), 1}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(p.ar.numTokens, clusters); + config.blockDim = dim3(threads); + config.stream = stream; + config.attrs = attrs; + config.numAttrs = smallBatch ? 1 : 2; + if (oneShot) { +#define LAUNCH(W) \ + case W: \ + if (smallBatch) \ + checkCuda(cudaLaunchKernelEx(&config, &oneshot, p)); \ + else if (p.experts) \ + checkCuda(cudaLaunchKernelEx(&config, &oneshot, p)); \ + else \ + checkCuda(cudaLaunchKernelEx(&config, &oneshot, p)); \ + break + switch (world) { + LAUNCH(2); + LAUNCH(4); + LAUNCH(8); + default: + TVM_FFI_LOG_AND_THROW(ValueError) << "Beam supports TP2/4/8"; + } +#undef LAUNCH + } else { + cudaLaunchConfig_t scatter{}; + scatter.gridDim = + dim3(p.ar.numTokens, flashinfer::ceil_div(p.ar.tokenDim, 1024)); + scatter.blockDim = dim3(128); + scatter.stream = stream; + scatter.attrs = attrs; + scatter.numAttrs = 1; +#define LAUNCH(W) \ + case W: \ + if (p.experts) \ + checkCuda(cudaLaunchKernelEx(&scatter, &finalizeScatter, p)); \ + else \ + checkCuda(cudaLaunchKernelEx( \ + &scatter, &ar::twoshotAllreduceKernel, p.ar)); \ + break + switch (world) { + LAUNCH(2); + LAUNCH(4); + LAUNCH(8); + default: + TVM_FFI_LOG_AND_THROW(ValueError) << "Beam supports TP2/4/8"; + } +#undef LAUNCH + if (clusters > 1) { + checkCuda(cudaLaunchKernelEx(&config, &finishTwoshot, p)); + } else { + config.numAttrs = 1; + checkCuda(cudaLaunchKernelEx(&config, &finishTwoshot, p)); + } + } +} + +void launchGather(TensorView experts, TensorView expertWeights, + TensorView residual, + tvm::ffi::Optional postWeight, + tvm::ffi::Optional nextWeight, TensorView hidden, + tvm::ffi::Optional normalized, int64_t pointers, + TensorView flags, int64_t multicast, int64_t local, + int64_t world, int64_t rank, double postEps, double nextEps, + double xScale, double rScale, bool postNorm, bool nextNorm, + bool scaled, bool oneShot, double routedScale) { + ffi::CUDADeviceGuard guard(residual.device().device_id); + TVM_FFI_ICHECK_EQ(encode_dlpack_dtype(experts.dtype()), bfloat16_code); + TVM_FFI_ICHECK_EQ(encode_dlpack_dtype(expertWeights.dtype()), float32_code); + TVM_FFI_ICHECK_EQ(encode_dlpack_dtype(residual.dtype()), bfloat16_code); + TVM_FFI_ICHECK_EQ(encode_dlpack_dtype(hidden.dtype()), bfloat16_code); + TVM_FFI_ICHECK(world == 2 || world == 4 || world == 8); + TVM_FFI_ICHECK(rank >= 0 && rank < world); + int64_t tokens = residual.size(0); + int64_t dim = residual.size(1); + int64_t localTokens = flashinfer::ceil_div(tokens, world); + TVM_FFI_ICHECK(dim >= 1024 && dim <= 8192 && dim % 1024 == 0); + TVM_FFI_ICHECK(experts.size(0) >= localTokens); + TVM_FFI_ICHECK(expertWeights.size(0) >= localTokens); + TVM_FFI_ICHECK_EQ(experts.size(2), dim); + TVM_FFI_ICHECK_EQ(expertWeights.size(1), experts.size(1)); + TVM_FFI_ICHECK(experts.size(1) >= 1 && experts.size(1) <= 32); + TVM_FFI_ICHECK_EQ(hidden.size(0), tokens); + TVM_FFI_ICHECK_EQ(hidden.size(1), dim); + TVM_FFI_ICHECK(!nextNorm || normalized.has_value()); + TVM_FFI_ICHECK(localTokens * world * dim <= INT32_MAX); + if (!tokens) return; + + Params p{}; + p.ar.outputPtr = normalized.has_value() + ? static_cast(normalized.value().data_ptr()) + : nullptr; + p.ar.prenormedPtr = static_cast(hidden.data_ptr()); + p.ar.residualInPtr = static_cast(residual.data_ptr()); + p.ar.gammaPtr = postWeight.has_value() + ? static_cast(postWeight.value().data_ptr()) + : nullptr; + p.nextWeight = nextWeight.has_value() + ? static_cast(nextWeight.value().data_ptr()) + : nullptr; + p.ar.inputPtrs = reinterpret_cast(pointers); + p.ar.mcastPtr = reinterpret_cast(multicast); + p.ar.localBufferPtr = reinterpret_cast(local); + p.ar.bufferFlags = static_cast(flags.data_ptr()); + p.ar.nRanks = world; + p.ar.rank = rank; + p.ar.numTokens = tokens; + p.ar.tokenDim = dim; + p.ar.epsilon = postEps; + p.nextEps = nextEps; + p.xScale = xScale; + p.rScale = rScale; + p.postNorm = postNorm; + p.nextNorm = nextNorm; + p.scaled = scaled; + p.experts = static_cast(experts.data_ptr()); + p.expertWeights = expertWeights.data_ptr(); + p.topK = experts.size(1); + p.routedScale = routedScale; + + cudaLaunchAttribute attr{}; + attr.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr.val.programmaticStreamSerializationAllowed = 1; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(localTokens); + config.blockDim = dim3(dim / 8); + config.stream = get_stream(residual.device()); + config.attrs = &attr; + config.numAttrs = 1; + if (oneShot) { + switch (world) { + case 2: + checkCuda(cudaLaunchKernelEx(&config, &finalizeGatherOneshot<2>, p)); + break; + case 4: + checkCuda(cudaLaunchKernelEx(&config, &finalizeGatherOneshot<4>, p)); + break; + case 8: + checkCuda(cudaLaunchKernelEx(&config, &finalizeGatherOneshot<8>, p)); + break; + } + } else { + checkCuda(cudaLaunchKernelEx(&config, &finalizeGatherPublish, p)); + config.gridDim = dim3(localTokens * world); + checkCuda(cudaLaunchKernelEx(&config, &finishGatherTwoshot, p)); + } +} +} // namespace beam +TVM_FFI_DLL_EXPORT_TYPED_FUNC(beam_allreduce, beam::launch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(beam_moe_gather, beam::launchGather); diff --git a/vllm/models/beam/nvidia/allreduce.py b/vllm/models/beam/nvidia/allreduce.py new file mode 100644 index 000000000000..90883b6775aa --- /dev/null +++ b/vllm/models/beam/nvidia/allreduce.py @@ -0,0 +1,274 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CUDA JIT MNNVL collective with Beam's FP32 sandwich-norm epilogue.""" + +import functools +import hashlib +from collections.abc import Callable +from pathlib import Path + +import torch + +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput +from vllm.utils.flashinfer import has_flashinfer +from vllm.utils.math_utils import cdiv + + +@functools.cache +def _module(): + from flashinfer.jit import env as jit_env + from flashinfer.jit.core import current_compilation_context, gen_jit_spec + + source = Path(__file__).with_suffix(".cu") + communication_header = ( + jit_env.FLASHINFER_INCLUDE_DIR / "flashinfer/comm/trtllm_mnnvl_allreduce.cuh" + ) + digest = hashlib.sha256( + source.read_bytes() + communication_header.read_bytes() + ).hexdigest()[:16] + flags = current_compilation_context.get_nvcc_flags_list( + supported_major_versions=[9, 10] + ) + spec = gen_jit_spec( + f"beam_allreduce_{digest}", + [source], + extra_cuda_cflags=[*flags, "--fmad=false"], + extra_include_paths=[jit_env.FLASHINFER_CSRC_DIR], + ) + spec.extra_cuda_cflags.remove("-use_fast_math") + return spec.build_and_load() + + +class BeamAllReduceWorkspace: + """Reuse FlashInfer's multicast allocation and Lamport generation protocol.""" + + @classmethod + def create( + cls, + group: torch.distributed.ProcessGroup, + hidden_size: int, + max_tokens: int, + dtype: torch.dtype, + *, + device_group: torch.distributed.ProcessGroup | None = None, + ) -> "BeamAllReduceWorkspace | None": + # The CUDA epilogue handles one packed load per thread. + if ( + group.size() not in (2, 4, 8) + or dtype != torch.bfloat16 + or hidden_size % 1024 + or not 1024 <= hidden_size <= 8192 + or not has_flashinfer() + ): + return None + from flashinfer.comm.mnnvl import ( + TorchDistBackend, + all_ranks_agree, + is_multicast_supported, + ) + + device = torch.accelerator.current_device_index() + supported = torch.cuda.get_device_capability(device)[0] in (9, 10) + supported = supported and is_multicast_supported(device) + if not all_ranks_agree(TorchDistBackend(group=group), supported): + return None + return cls(group, hidden_size, max_tokens, device_group=device_group) + + def __init__( + self, + group: torch.distributed.ProcessGroup, + hidden_size: int, + max_tokens: int, + *, + device_group: torch.distributed.ProcessGroup | None = None, + ) -> None: + from flashinfer.comm import create_allreduce_fusion_workspace + from flashinfer.comm.mnnvl import TorchDistBackend + + self.world_size = group.size() + self.rank = torch.distributed.get_rank(group) + self.hidden_size = hidden_size + self.max_tokens = max_tokens + self.device_group = device_group + self.all_gather: Callable[[torch.Tensor], torch.Tensor] | None = None + # Scale one-shot cutoffs by the BF16 payload staged across all ranks. + token_bytes = hidden_size * 2 + rank_token_bytes = self.world_size * token_bytes + self.small_batch_max_tokens = 3200 * 1024 // rank_token_bytes + self.ep_finalize_oneshot_max_tokens = 2560 * 1024 // rank_token_bytes + # GB300 TP2/4, H=5120, top-k=8: one-shot helps at one token; + # two-shot wins through 131072 tokens. Keep the unfused cutoff tunable + # for other geometries instead of imposing an unmeasured crossover. + self.moe_gather_oneshot_max_tokens = 1 + self.moe_gather_separate_norm_min_tokens: int | None = None + # Two-shot reserves two rows per token; one-shot needs one per rank. + workspace_tokens = max( + max_tokens, + min(max_tokens, self.small_batch_max_tokens) * self.world_size // 2, + ) + self.workspace = create_allreduce_fusion_workspace( + backend="mnnvl", + world_size=self.world_size, + rank=self.rank, + max_token_num=workspace_tokens, + hidden_dim=hidden_size, + dtype=torch.bfloat16, + comm_backend=TorchDistBackend(group=group), + group=group, + ) + self.module = _module() + self.allocation_bytes = ( + self.workspace.handle.mcast_device_memory.get_allocation_size() + ) + self.separate_norm_min_tokens = cdiv( + 10 * 1024**2 * self.world_size, token_bytes + ) + self.ep_finalize_max_tokens = 640 * 1024**2 // token_bytes + + +def beam_allreduce_norm_residual_norm( + x: torch.Tensor, + residual: torch.Tensor, + post_weight: torch.Tensor | None, + post_eps: float | None, + next_weight: torch.Tensor | None, + next_eps: float | None, + scales: tuple[float, float] | None, + workspace: BeamAllReduceWorkspace, +) -> tuple[torch.Tensor, torch.Tensor | None]: + assert x.shape == residual.shape + for weight in (post_weight, next_weight): + if weight is not None: + assert weight.dtype == torch.bfloat16 and weight.is_contiguous() + assert weight.shape == (workspace.hidden_size,) + assert x.dtype == residual.dtype == torch.bfloat16 + assert x.is_contiguous() and residual.is_contiguous() + assert x.shape[1] == workspace.hidden_size + assert x.shape[0] <= workspace.max_tokens + if x.shape[0] >= workspace.separate_norm_min_tokens: + from flashinfer.comm import AllReduceFusionPattern, allreduce_fusion + + from .layernorm import beam_norm_residual_norm + + reduced = allreduce_fusion( + input=x, + workspace=workspace.workspace, + pattern=AllReduceFusionPattern.kAllReduce, + launch_with_pdl=True, + ) + return beam_norm_residual_norm( + reduced, + residual, + post_weight, + post_eps, + next_weight, + next_eps, + scales, + launch_pdl=True, + ) + return _cuda_allreduce_norm_residual_norm( + x, residual, post_weight, post_eps, next_weight, next_eps, scales, workspace + ) + + +def _cuda_allreduce_norm_residual_norm( + x, + residual, + post_weight, + post_eps, + next_weight, + next_eps, + scales, + workspace, + unfinalized: UnfinalizedMoEOutput | None = None, + routed_scale: float = 1.0, + oneshot_max_tokens: int = 0, +): + from flashinfer.comm.trtllm_mnnvl_ar import MNNVLAllreduceFusionStrategy + + assert x.dtype == residual.dtype == torch.bfloat16 + assert x.is_contiguous() and residual.is_contiguous() + m, n = x.shape + assert n == workspace.hidden_size and m <= workspace.max_tokens + hidden = torch.empty_like(x) + normalized = torch.empty_like(x) if next_eps is not None else None + if m: + strategy = MNNVLAllreduceFusionStrategy.select_strategy( + workspace.world_size, m, n, x.dtype + ) + small_batch = unfinalized is None and m <= workspace.small_batch_max_tokens + w = workspace.workspace + workspace.module.beam_allreduce( + x, + residual, + post_weight, + next_weight, + hidden, + normalized, + w.uc_ptrs_dev, + w.buffer_flags, + w.mc_ptr, + w.uc_ptr_local, + workspace.world_size, + workspace.rank, + post_eps or 0.0, + next_eps or 0.0, + *(scales if scales is not None else (1.0, 1.0)), + post_eps is not None, + next_eps is not None, + scales is not None, + small_batch + or m <= oneshot_max_tokens + or strategy == MNNVLAllreduceFusionStrategy.ONESHOT, + unfinalized.gemm2_permuted if unfinalized is not None else None, + unfinalized.expert_weights if unfinalized is not None else None, + unfinalized.expanded_idx_to_permuted_idx + if unfinalized is not None + else None, + routed_scale, + small_batch, + ) + return hidden, normalized + + +def beam_finalize_allreduce_norm_residual_norm( + x: UnfinalizedMoEOutput, + residual: torch.Tensor, + post_weight: torch.Tensor | None, + post_eps: float | None, + next_weight: torch.Tensor | None, + next_eps: float | None, + scales: tuple[float, float] | None, + workspace: BeamAllReduceWorkspace, + routed_scale: float, + *, + oneshot_max_tokens: int = 0, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Finalize local experts before AR, retaining the local BF16 rounding.""" + m, n = residual.shape + assert 0 <= oneshot_max_tokens <= workspace.small_batch_max_tokens + for weight in (post_weight, next_weight): + if weight is not None: + assert weight.dtype == torch.bfloat16 and weight.is_contiguous() + assert weight.shape == (n,) + assert x.gemm2_permuted.dtype == torch.bfloat16 + assert x.gemm2_permuted.is_contiguous() and x.gemm2_permuted.shape[1] == n + assert x.expert_weights.dtype in (torch.float32, torch.bfloat16) + assert x.expanded_idx_to_permuted_idx.dtype == torch.int32 + assert x.expert_weights.is_contiguous() + assert x.expanded_idx_to_permuted_idx.is_contiguous() + assert x.expert_weights.shape == x.expanded_idx_to_permuted_idx.shape + assert x.expert_weights.shape[0] == m + return _cuda_allreduce_norm_residual_norm( + residual, + residual, + post_weight, + post_eps, + next_weight, + next_eps, + scales, + workspace, + x, + routed_scale, + oneshot_max_tokens, + ) diff --git a/vllm/models/beam/nvidia/attention.py b/vllm/models/beam/nvidia/attention.py new file mode 100644 index 000000000000..1511fb002e26 --- /dev/null +++ b/vllm/models/beam/nvidia/attention.py @@ -0,0 +1,272 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math + +import torch +from torch import nn + +from vllm.config import CacheConfig +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, + UnquantizedLinearMethod, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader +from vllm.model_executor.models.utils import extract_layer_index +from vllm.model_executor.utils import set_weight_attrs +from vllm.transformers_utils.configs.beam import BeamConfig +from vllm.v1.attention.backend import AttentionType + +from .attention_ops import beam_output_gate, beam_qk_norm_rope +from .layernorm import RMSNorm +from .projection import ( + packed_projection, + use_low_latency_gemm, + warmup_packed_projection, +) + + +def use_merged_qkv_gate(config: BeamConfig, tp_size: int) -> bool: + return ( + config.attention_output_gate + and not config.attention_bias + and config.num_key_value_heads >= tp_size + ) + + +class BeamAttention(nn.Module): + """GQA with per-head Q/K RMSNorm, then RoPE (default) or SSMax (NoPE). + + With ``attention_output_gate``, the attention output is scaled by + ``gate_floor + (1 - gate_floor) * sigmoid((gate + bias) / temperature)`` + before ``o_proj``. + """ + + def __init__( + self, + config: BeamConfig, + cache_config: CacheConfig | None, + quant_config: QuantizationConfig | None, + prefix: str = "", + reduce_results: bool = True, + ) -> None: + super().__init__() + self.layer_id = extract_layer_index(prefix) + hidden_size = config.hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = config.num_attention_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = config.num_key_value_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + else: + 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 = config.head_dim or (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 + + self.merged_qkv_gate = use_merged_qkv_gate(config, tp_size) + if self.merged_qkv_gate: + self.qkv_gate_proj = MergedColumnParallelLinear( + hidden_size, + [ + self.total_num_heads * self.head_dim, + self.total_num_kv_heads * self.head_dim, + self.total_num_kv_heads * self.head_dim, + self.total_num_heads * self.head_dim, + ], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.qkv_gate_proj", + ) + else: + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=config.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=config.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + reduce_results=reduce_results, + ) + + # Affine and parameter-free per-head Q/K norms share `rms_norm_eps`. + qk_norm_has_weight = not config.use_paramfree_qk_norm + self.rms_norm_eps = config.rms_norm_eps + self.q_norm = RMSNorm(self.head_dim, has_weight=qk_norm_has_weight) + self.k_norm = RMSNorm(self.head_dim, has_weight=qk_norm_has_weight) + + self.disable_rope = self.layer_id in config.nope_layer_ids + self.use_ssmax = bool(config.ssmax) and self.disable_rope + self.ssmax_mode = config.ssmax_scale_mode if self.use_ssmax else "none" + if self.ssmax_mode not in ("none", "default", "root10", "root1000"): + raise ValueError(f"Unknown ssmax_scale_mode: {self.ssmax_mode!r}") + self.ssmax_offset = config.ssmax_offset + self.ssmax_log_divisor = ( + math.log(config.ssmax_offset) + if self.use_ssmax + and self.ssmax_mode == "default" + and config.ssmax_divide_by_log_offset + else 1.0 + ) + # The output gate has a separate bias per head. + self.attn_output_gate = bool(config.attention_output_gate) + if self.attn_output_gate: + self.gate_floor = float(config.attention_output_gate_floor) + self.gate_temperature = float(config.attention_output_gate_temperature) + if not 0 <= self.gate_floor < 1: + raise ValueError( + "attention_output_gate_floor must be in [0, 1); " + f"got {self.gate_floor}" + ) + if self.gate_temperature <= 0: + raise ValueError( + "attention_output_gate_temperature must be positive; " + f"got {self.gate_temperature}" + ) + if not self.merged_qkv_gate: + self.output_gate_proj = ColumnParallelLinear( + hidden_size, + self.total_num_heads * self.head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.output_gate_proj", + ) + if config.attention_output_gate_bias: + self.gate_bias = nn.Parameter(torch.zeros(self.num_heads)) + set_weight_attrs( + self.gate_bias, {"weight_loader": sharded_weight_loader(0)} + ) + else: + self.register_parameter("gate_bias", None) + + if self.disable_rope: + self.rotary_emb = None + else: + self.rotary_emb = get_rope( + self.head_dim, + max_position=config.max_position_embeddings, + is_neox_style=True, + rope_parameters=config.rope_parameters, + dtype=torch.float32, + ) + + # Attention masks follow `layer_types` independently of the NoPE flags. + # Backends interpret `sliding_window` as `(q - k) <= sliding_window - 1`. + is_sliding = config.layer_types[self.layer_id] == "sliding_attention" + per_layer_sw = config.sliding_window if is_sliding else None + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + per_layer_sliding_window=per_layer_sw, + attn_type=AttentionType.DECODER, + prefix=f"{prefix}.attn", + ) + + self.fuse_kv_insert = ( + self.attn.attn_backend.get_name() in ("FLASH_ATTN", "FLASHINFER") + and self.attn.impl.dcp_world_size == 1 + and self.attn.kv_sharing_target_layer_name is None + ) + self.low_latency_projection = False + + def warmup_projection(self) -> None: + if not self.merged_qkv_gate: + return + projection = self.qkv_gate_proj + self.low_latency_projection = isinstance( + projection.quant_method, UnquantizedLinearMethod + ) and use_low_latency_gemm(projection.weight) + if self.low_latency_projection: + warmup_packed_projection(projection.weight) + + def forward( + self, positions: torch.Tensor, hidden_states: torch.Tensor + ) -> torch.Tensor: + gate = None + if self.merged_qkv_gate: + if self.low_latency_projection: + packed = packed_projection( + hidden_states, self.qkv_gate_proj.weight, True + ) + else: + packed, _ = self.qkv_gate_proj(hidden_states) + qkv, gate = packed.split( + [self.q_size + 2 * self.kv_size, self.q_size], dim=-1 + ) + else: + qkv, _ = self.qkv_proj(hidden_states) + if self.attn_output_gate: + gate, _ = self.output_gate_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + kv_cache, slot_mapping = None, None + if ( + self.fuse_kv_insert + and self.attn.kv_cache.dtype == q.dtype == torch.bfloat16 + ): + slots = get_forward_context().slot_mapping + if isinstance(slots, dict): + slot_mapping = slots.get(self.attn.layer_name) + if slot_mapping is not None: + kv_cache = self.attn.kv_cache + + rope = self.rotary_emb + q, k = beam_qk_norm_rope( + q, + k, + self.q_norm.weight, + self.k_norm.weight, + self.rms_norm_eps, + positions, + rope.cos_sin_cache if rope is not None else None, + self.head_dim, + rope.rotary_dim if rope is not None else 0, + rope.is_neox_style if rope is not None else True, + self.ssmax_mode, + self.ssmax_offset, + self.ssmax_log_divisor, + value=v, + kv_cache=kv_cache, + slot_mapping=slot_mapping, + ) + + # Both backends read the populated cache when K/V are omitted. + attn_output = ( + self.attn(q, k, v) if kv_cache is None else self.attn(q, None, None) + ) + if gate is not None: + attn_output = beam_output_gate( + attn_output, + gate, + self.gate_bias, + num_heads=self.num_heads, + head_dim=self.head_dim, + floor=self.gate_floor, + temperature=self.gate_temperature, + ) + output, _ = self.o_proj(attn_output) + return output diff --git a/vllm/models/beam/nvidia/attention_ops.py b/vllm/models/beam/nvidia/attention_ops.py new file mode 100644 index 000000000000..247b53fd74f9 --- /dev/null +++ b/vllm/models/beam/nvidia/attention_ops.py @@ -0,0 +1,289 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused Q/K normalization, position transforms, and attention output gating.""" + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.triton_utils import tldevice as libdevice + + +@triton.jit(do_not_specialize=["SIZE"]) +def _output_gate_kernel( + X, + G, + B, + Y, + SIZE, + G_STRIDE: tl.constexpr, + HEADS: tl.constexpr, + DIM: tl.constexpr, + FLOOR: tl.constexpr, + TEMPERATURE: tl.constexpr, + BIAS: tl.constexpr, + BLOCK: tl.constexpr, +): + i = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + dtype = G.dtype.element_ty + g_idx = i // (HEADS * DIM) * G_STRIDE + i % (HEADS * DIM) + g = tl.load(G + g_idx, i < SIZE, other=0).to(tl.float32) + if BIAS: + b = tl.load(B + i // DIM % HEADS).to(tl.float32) + g = (g + b).to(dtype).to(tl.float32) + g = (g / TEMPERATURE).to(dtype).to(tl.float32) + sigmoid = (1.0 / (1.0 + tl.exp(-g))).to(dtype).to(tl.float32) + multiplier = (sigmoid * (1.0 - FLOOR)).to(dtype).to(tl.float32) + multiplier = (multiplier + FLOOR).to(dtype).to(tl.float32) + x = tl.load(X + i, i < SIZE, other=0).to(tl.float32) + tl.store(Y + i, x * multiplier, i < SIZE) + + +def beam_output_gate( + attn_output: torch.Tensor, + gate: torch.Tensor, + gate_bias: torch.Tensor | None, + *, + num_heads: int, + head_dim: int, + floor: float, + temperature: float, +) -> torch.Tensor: + """Fuse gate operations while retaining activation-dtype rounding.""" + output = torch.empty_like(attn_output) + if output.numel(): + if ( + attn_output.dtype == gate.dtype == torch.bfloat16 + and (gate_bias is None or gate_bias.dtype == torch.bfloat16) + and current_platform.is_device_capability_family(100) + ): + from .attention_ops_cutedsl import run + + run( + attn_output, + gate, + gate_bias, + num_heads, + head_dim, + floor, + temperature, + output, + ) + return output + _output_gate_kernel[(triton.cdiv(output.numel(), 256),)]( + attn_output, + gate, + gate_bias, + output, + output.numel(), + gate.stride(0), + num_heads, + head_dim, + floor, + temperature, + gate_bias is not None, + 256, + enable_fp_fusion=False, + ) + return output + + +@triton.jit( + do_not_specialize=["M", "NUM_SLOTS"], do_not_specialize_on_alignment=["POS"] +) +def _qk_norm_rope_kernel( + Q, + K, + QW, + KW, + POS, + CACHE, + OQ, + OK, + M, + V, + KV_CACHE, + SLOTS, + NUM_SLOTS, + WRITE_CACHE: tl.constexpr, + V_STRIDE: tl.constexpr, + CACHE_BLOCK_STRIDE: tl.constexpr, + CACHE_HEAD_STRIDE: tl.constexpr, + CACHE_TOKEN_STRIDE: tl.constexpr, + PAGE_SIZE: tl.constexpr, + TOKENS: tl.constexpr, + Q_HEADS: tl.constexpr, + K_HEADS: tl.constexpr, + DIM: tl.constexpr, + Q_STRIDE: tl.constexpr, + K_STRIDE: tl.constexpr, + EPS: tl.constexpr, + Q_WEIGHT: tl.constexpr, + K_WEIGHT: tl.constexpr, + ROTARY: tl.constexpr, + NEOX: tl.constexpr, + SSMAX: tl.constexpr, + OFFSET: tl.constexpr, + LOG_DIVISOR: tl.constexpr, + BLOCK: tl.constexpr, +): + token = tl.program_id(0).to(tl.int64) * TOKENS + tl.arange(0, TOKENS) + head = tl.program_id(1) + c = tl.arange(0, BLOCK) + mask = (token[:, None] < M) & (c[None, :] < DIM) + if head < Q_HEADS: + x = tl.load( + Q + token[:, None] * Q_STRIDE + head * DIM + c[None, :], mask, 0 + ).to(tl.float32) + if BLOCK == 128: + squares = (x * x).reshape((TOKENS, 4, 32)) + total = tl.sum(tl.sum(squares, 2), 1) + else: + total = tl.sum(x * x, 1) + y = x * tl.rsqrt(total / DIM + EPS)[:, None] + if Q_WEIGHT: + y = y * tl.load(QW + c, c < DIM, 0).to(tl.float32)[None, :] + out = OQ + (token[:, None] * Q_HEADS + head) * DIM + else: + h = head - Q_HEADS + x = tl.load(K + token[:, None] * K_STRIDE + h * DIM + c[None, :], mask, 0).to( + tl.float32 + ) + if BLOCK == 128: + squares = (x * x).reshape((TOKENS, 4, 32)) + total = tl.sum(tl.sum(squares, 2), 1) + else: + total = tl.sum(x * x, 1) + y = x * tl.rsqrt(total / DIM + EPS)[:, None] + if K_WEIGHT: + y = y * tl.load(KW + c, c < DIM, 0).to(tl.float32)[None, :] + out = OK + (token[:, None] * K_HEADS + h) * DIM + y = y.to(Q.dtype.element_ty).to(tl.float32) + if ROTARY: + half: tl.constexpr = ROTARY // 2 + if NEOX: + partner = tl.where(c < half, c + half, c - half) + second = c >= half + offset = c % half + else: + partner = c ^ 1 + second = c % 2 != 0 + offset = c // 2 + partner = tl.minimum(tl.maximum(partner, 0), BLOCK - 1) + other = tl.gather(y, tl.broadcast_to(partner[None, :], (TOKENS, BLOCK)), axis=1) + pos = tl.load(POS + token, token < M, 0).to(tl.int64) + rope_mask = (token[:, None] < M) & (c[None, :] < ROTARY) + cos = tl.load(CACHE + pos[:, None] * ROTARY + offset[None, :], rope_mask, 0) + sin = tl.load( + CACHE + pos[:, None] * ROTARY + half + offset[None, :], rope_mask, 0 + ) + # Match the FP32 multiply/FMA ordering of the CUDA RoPE kernel. + rotated = tl.where( + second[None, :], + tl.fma(other, sin, y * cos), + tl.fma(y, cos, -other * sin), + ) + y = tl.where(c[None, :] < ROTARY, rotated, y) + elif SSMAX != "none": + if head < Q_HEADS: + p = tl.load(POS + token, token < M, 0).to(tl.float32) + OFFSET + if SSMAX == "default": + scale = libdevice.log(p) / LOG_DIVISOR + elif SSMAX == "root10": + scale = libdevice.exp(0.1 * libdevice.log1p(p / 2048.0)) + else: + scale = tl.maximum( + 130.0 * libdevice.expm1(0.001 * libdevice.log(p)), 1.0 + ) + y = y * scale[:, None] + tl.store(out + c[None, :], y, mask) + if WRITE_CACHE and head >= Q_HEADS: + slot = tl.load(SLOTS + token, token < NUM_SLOTS, -1).to(tl.int64) + cache_mask = mask & (token[:, None] < NUM_SLOTS) & (slot[:, None] >= 0) + address = ( + slot[:, None] // PAGE_SIZE * CACHE_BLOCK_STRIDE + + slot[:, None] % PAGE_SIZE * CACHE_TOKEN_STRIDE + + (head - Q_HEADS) * CACHE_HEAD_STRIDE + + c[None, :] + ) + tl.store(KV_CACHE + address, y, cache_mask) + value = tl.load( + V + token[:, None] * V_STRIDE + (head - Q_HEADS) * DIM + c[None, :], + cache_mask, + 0, + ) + tl.store(KV_CACHE + address + DIM, value, cache_mask) + + +def beam_qk_norm_rope( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor | None, + k_weight: torch.Tensor | None, + eps: float, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor | None, + head_dim: int, + rotary_dim: int = 0, + is_neox_style: bool = True, + ssmax_mode: str = "none", + ssmax_offset: int = 1, + ssmax_log_divisor: float = 1.0, + *, + value: torch.Tensor | None = None, + kv_cache: torch.Tensor | None = None, + slot_mapping: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Round norm before position transforms and optionally insert BF16 K/V.""" + if kv_cache is not None: + assert value is not None and slot_mapping is not None + assert kv_cache.dtype == q.dtype == torch.bfloat16 + assert kv_cache.ndim == 4 and kv_cache.stride(-1) == 1 + q_out, k_out = torch.empty_like(q), torch.empty_like(k) + q_heads, k_heads = q.shape[-1] // head_dim, k.shape[-1] // head_dim + if q.shape[0]: + if rotary_dim or q.shape[0] < 512: + tile, warps = 4, 4 + elif q.shape[0] <= 2048: + tile, warps = 16, 8 + else: + tile, warps = 32, 4 + _qk_norm_rope_kernel[(triton.cdiv(q.shape[0], tile), q_heads + k_heads)]( + q, + k, + q_weight, + k_weight, + positions, + cos_sin_cache, + q_out, + k_out, + q.shape[0], + value, + kv_cache, + slot_mapping, + slot_mapping.numel() if slot_mapping is not None else 0, + kv_cache is not None, + value.stride(0) if value is not None else 0, + kv_cache.stride(0) if kv_cache is not None else 0, + kv_cache.stride(1) if kv_cache is not None else 0, + kv_cache.stride(2) if kv_cache is not None else 0, + kv_cache.shape[2] if kv_cache is not None else 1, + tile, + q_heads, + k_heads, + head_dim, + q.stride(0), + k.stride(0), + eps, + q_weight is not None, + k_weight is not None, + rotary_dim, + is_neox_style, + ssmax_mode, + ssmax_offset, + ssmax_log_divisor, + triton.next_power_of_2(head_dim), + num_warps=warps, + enable_fp_fusion=False, + ) + return q_out, k_out diff --git a/vllm/models/beam/nvidia/attention_ops_cutedsl.py b/vllm/models/beam/nvidia/attention_ops_cutedsl.py new file mode 100644 index 000000000000..eb4b0678524d --- /dev/null +++ b/vllm/models/beam/nvidia/attention_ops_cutedsl.py @@ -0,0 +1,989 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Bit-exact BF16 output gate in CuTe DSL.""" + +from collections.abc import Callable + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass import Float32, Int32, Int64, Uint32, const_expr +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +from vllm.utils.torch_utils import current_stream + +_LDQ = "ld.global.nc.L1::no_allocate" + + +def _asm(res_ty, operands, text, constraints, loc=None, ip=None, side_effects=False): + return llvm.inline_asm( + res_ty, + operands, + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def ld_v4(addr: Int64, pf="", *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [Int64(addr).ir_value(loc=loc, ip=ip)], + _LDQ + pf + ".v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc, + ip, + ) + return [ + Uint32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(4) + ] + + +@dsl_user_op +def make_policy(kind: str, *, loc=None, ip=None) -> Int64: + return Int64( + _asm( + T.i64(), + [], + f"createpolicy.fractional.L2::{kind}.b64 $0, 1.0;", + "=l", + loc, + ip, + ) + ) + + +@dsl_user_op +def st_v4_hint(addr: Int64, words, pol: Int64, *, loc=None, ip=None): + ops = [Int64(addr).ir_value(loc=loc, ip=ip)] + [ + Uint32(w).ir_value(loc=loc, ip=ip) for w in words + ] + ops.append(Int64(pol).ir_value(loc=loc, ip=ip)) + _asm( + None, + ops, + "st.global.L2::cache_hint.v4.b32 [$0], {$1, $2, $3, $4}, $5;", + "l,r,r,r,r,l", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def ld_v2(addr: Int64, pf=".L2::128B", *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 2), + [Int64(addr).ir_value(loc=loc, ip=ip)], + _LDQ + pf + ".v2.b32 {$0, $1}, [$2];", + "=r,=r,l", + loc, + ip, + ) + return [ + Uint32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(2) + ] + + +@dsl_user_op +def ld_b32(addr: Int64, *, loc=None, ip=None): + return [ + Uint32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + _LDQ + ".L2::128B.b32 $0, [$1];", + "=r,l", + loc, + ip, + ) + ) + ] + + +@dsl_user_op +def ld_b16(addr: Int64, *, loc=None, ip=None) -> Uint32: + return Uint32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "{\n.reg .b16 h;\nld.global.nc.b16 h, [$1];\ncvt.u32.u16 $0, h;\n}", + "=r,l", + loc, + ip, + ) + ) + + +@dsl_user_op +def ld_b16_l2_128(addr: Int64, pf=".L2::128B", *, loc=None, ip=None) -> Uint32: + return Uint32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "{\n.reg .b16 h;\n" + _LDQ + pf + ".b16 h, [$1];\ncvt.u32.u16 $0, h;\n}", + "=r,l", + loc, + ip, + ) + ) + + +@dsl_user_op +def ld_bias_b16x2(addr: Int64, *, loc=None, ip=None) -> Uint32: + return Uint32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "{\n.reg .b16 h;\nld.global.nc.b16 h, [$1];\nmov.b32 $0, {h, h};\n}", + "=r,l", + loc, + ip, + ) + ) + + +@dsl_user_op +def st_vec(addr: Int64, words, *, loc=None, ip=None): + n = len(words) + ops = [Int64(addr).ir_value(loc=loc, ip=ip)] + [ + Uint32(w).ir_value(loc=loc, ip=ip) for w in words + ] + if n == 4: + text = "st.global.v4.b32 [$0], {$1, $2, $3, $4};" + elif n == 2: + text = "st.global.v2.b32 [$0], {$1, $2};" + else: + text = "st.global.b32 [$0], $1;" + _asm(None, ops, text, "l" + ",r" * n, loc, ip, side_effects=True) + + +@dsl_user_op +def st_b16(addr: Int64, word, *, loc=None, ip=None): + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip), Uint32(word).ir_value(loc=loc, ip=ip)], + "{\n.reg .b16 l, h;\nmov.b32 {l, h}, $1;\nst.global.b16 [$0], l;\n}", + "l,r", + loc, + ip, + side_effects=True, + ) + + +def _gate_ptx(has_bias, has_temp, has_floor): + s = [ + "{", + ".reg .b32 w, lo, hi;", + ".reg .b64 p, kl, k1;", + ".reg .f32 a0, a1;", + "mov.b64 kl, 0xBFB8AA3BBFB8AA3B;", + "mov.b64 k1, 0x3F8000003F800000;", + ] + w = "$1" + if has_bias: + s.append("add.rn.bf16x2 w, $1, $3;") + w = "w" + + def unpack(src): + s.append(f"shl.b32 lo, {src}, 16;") + s.append(f"and.b32 hi, {src}, 0xffff0000;") + + if has_temp: + unpack(w) + s += [ + "mov.b32 a0, lo;", + "mov.b32 a1, hi;", + "div.full.f32 a0, a0, $4;", + "div.full.f32 a1, a1, $4;", + "cvt.rn.bf16x2.f32 w, a1, a0;", + ] + w = "w" + unpack(w) + s += [ + "mov.b64 p, {lo, hi};", + "mul.rn.f32x2 p, p, kl;", + "mov.b64 {a0, a1}, p;", + "ex2.approx.ftz.f32 a0, a0;", + "ex2.approx.ftz.f32 a1, a1;", + "mov.b64 p, {a0, a1};", + "add.rn.f32x2 p, p, k1;", + "mov.b64 {a0, a1}, p;", + "rcp.approx.ftz.f32 a0, a0;", + "rcp.approx.ftz.f32 a1, a1;", + "cvt.rn.bf16x2.f32 w, a1, a0;", + ] + if has_floor: + s += [".reg .b64 kc, kf;", "mov.b64 kc, {$5, $5};", "mov.b64 kf, {$6, $6};"] + unpack("w") + s += [ + "mov.b64 p, {lo, hi};", + "mul.rn.f32x2 p, p, kc;", + "mov.b64 {a0, a1}, p;", + "cvt.rn.bf16x2.f32 w, a1, a0;", + ] + unpack("w") + s += [ + "mov.b64 p, {lo, hi};", + "add.rn.f32x2 p, p, kf;", + "mov.b64 {a0, a1}, p;", + "cvt.rn.bf16x2.f32 w, a1, a0;", + ] + s += ["mul.rn.bf16x2 $0, w, $2;", "}"] + return "\n".join(s) + + +@dsl_user_op +def gate_pair( + g, x, b, temp, c1, flo, has_bias, has_temp, has_floor, *, loc=None, ip=None +) -> Uint32: + ops = [ + Uint32(g).ir_value(loc=loc, ip=ip), + Uint32(x).ir_value(loc=loc, ip=ip), + Uint32(b).ir_value(loc=loc, ip=ip), + Float32(temp).ir_value(loc=loc, ip=ip), + Float32(c1).ir_value(loc=loc, ip=ip), + Float32(flo).ir_value(loc=loc, ip=ip), + ] + return Uint32( + _asm( + T.i32(), + ops, + _gate_ptx(has_bias, has_temp, has_floor), + "=r,r,r,r,f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def gate_pair_default_bias(g, x, b, *, loc=None, ip=None) -> Uint32: + ops = [ + Uint32(g).ir_value(loc=loc, ip=ip), + Uint32(x).ir_value(loc=loc, ip=ip), + Uint32(b).ir_value(loc=loc, ip=ip), + ] + return Uint32( + _asm(T.i32(), ops, _gate_ptx(True, False, False), "=r,r,r,r", loc, ip) + ) + + +@dsl_user_op +def gate_pair_default_nobias(g, x, *, loc=None, ip=None) -> Uint32: + ops = [Uint32(g).ir_value(loc=loc, ip=ip), Uint32(x).ir_value(loc=loc, ip=ip)] + return Uint32(_asm(T.i32(), ops, _gate_ptx(False, False, False), "=r,r,r", loc, ip)) + + +def gate_pair_selected(g, x, b, temp, c1, flo, has_bias, has_temp, has_floor): + if const_expr(not has_temp and not has_floor): + if const_expr(has_bias): + return gate_pair_default_bias(g, x, b) + return gate_pair_default_nobias(g, x) + return gate_pair(g, x, b, temp, c1, flo, has_bias, has_temp, has_floor) + + +def _gate_vec_exp_interleaved_default_ptx(has_bias, npairs): + """Two BF16x2 pairs with independent MUFU operations issued together.""" + assert npairs == 2 + s = [ + "{", + ".reg .b32 w0, w1, lo, hi;", + ".reg .b64 p0, p1, kl, k1;", + ".reg .f32 a0, a1, a2, a3;", + "mov.b64 kl, 0xBFB8AA3BBFB8AA3B;", + "mov.b64 k1, 0x3F8000003F800000;", + ] + g0, g1 = "$2", "$3" + if has_bias: + s += ["add.rn.bf16x2 w0, $2, $6;", "add.rn.bf16x2 w1, $3, $6;"] + g0, g1 = "w0", "w1" + s += [ + f"shl.b32 lo, {g0}, 16;", + f"and.b32 hi, {g0}, 0xffff0000;", + "mov.b64 p0, {lo, hi};", + "mul.rn.f32x2 p0, p0, kl;", + "mov.b64 {a0, a1}, p0;", + f"shl.b32 lo, {g1}, 16;", + f"and.b32 hi, {g1}, 0xffff0000;", + "mov.b64 p1, {lo, hi};", + "mul.rn.f32x2 p1, p1, kl;", + "mov.b64 {a2, a3}, p1;", + "ex2.approx.ftz.f32 a0, a0;", + "ex2.approx.ftz.f32 a1, a1;", + "ex2.approx.ftz.f32 a2, a2;", + "ex2.approx.ftz.f32 a3, a3;", + "mov.b64 p0, {a0, a1};", + "mov.b64 p1, {a2, a3};", + "add.rn.f32x2 p0, p0, k1;", + "add.rn.f32x2 p1, p1, k1;", + "mov.b64 {a0, a1}, p0;", + "mov.b64 {a2, a3}, p1;", + "rcp.approx.ftz.f32 a0, a0;", + "rcp.approx.ftz.f32 a1, a1;", + "rcp.approx.ftz.f32 a2, a2;", + "rcp.approx.ftz.f32 a3, a3;", + "cvt.rn.bf16x2.f32 w0, a1, a0;", + "cvt.rn.bf16x2.f32 w1, a3, a2;", + "mul.rn.bf16x2 $0, w0, $4;", + "mul.rn.bf16x2 $1, w1, $5;", + "}", + ] + return "\n".join(s) + + +@dsl_user_op +def gate_vec_exp_interleaved_default(gs, xs, b, *, loc=None, ip=None): + n = len(gs) + ops = [Uint32(v).ir_value(loc=loc, ip=ip) for v in gs + xs] + if b is not None: + ops.append(Uint32(b).ir_value(loc=loc, ip=ip)) + res = _asm( + llvm.StructType.get_literal([T.i32()] * n), + ops, + _gate_vec_exp_interleaved_default_ptx(b is not None, n), + ",".join(["=r"] * n + ["r"] * len(ops)), + loc, + ip, + ) + return [ + Uint32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(n) + ] + + +def _gate_vec_default_ptx(has_bias, npairs): + s = [ + "{", + ".reg .b32 w, lo, hi;", + ".reg .b64 p, kl, k1;", + ".reg .f32 a0, a1;", + "mov.b64 kl, 0xBFB8AA3BBFB8AA3B;", + "mov.b64 k1, 0x3F8000003F800000;", + ] + for j in range(npairs): + g = f"${npairs + j}" + x = f"${2 * npairs + j}" + if has_bias: + s.append(f"add.rn.bf16x2 w, {g}, ${3 * npairs};") + g = "w" + s += [ + f"shl.b32 lo, {g}, 16;", + f"and.b32 hi, {g}, 0xffff0000;", + "mov.b64 p, {lo, hi};", + "mul.rn.f32x2 p, p, kl;", + "mov.b64 {a0, a1}, p;", + "ex2.approx.ftz.f32 a0, a0;", + "ex2.approx.ftz.f32 a1, a1;", + "mov.b64 p, {a0, a1};", + "add.rn.f32x2 p, p, k1;", + "mov.b64 {a0, a1}, p;", + "rcp.approx.ftz.f32 a0, a0;", + "rcp.approx.ftz.f32 a1, a1;", + "cvt.rn.bf16x2.f32 w, a1, a0;", + f"mul.rn.bf16x2 ${j}, w, {x};", + ] + s.append("}") + return "\n".join(s) + + +@dsl_user_op +def gate_vec_default(gs, xs, b, *, loc=None, ip=None): + """Default-math (floor=0, temperature=1) gate for 2 or 4 packed BF16x2 pairs.""" + n = len(gs) + ops = [Uint32(v).ir_value(loc=loc, ip=ip) for v in gs + xs] + if b is not None: + ops.append(Uint32(b).ir_value(loc=loc, ip=ip)) + res = _asm( + llvm.StructType.get_literal([T.i32()] * n), + ops, + _gate_vec_default_ptx(b is not None, n), + ",".join(["=r"] * n + ["r"] * len(ops)), + loc, + ip, + ) + return [ + Uint32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(n) + ] + + +def _gate_ptx_scalar(has_bias, has_temp, has_floor): + s = [ + "{", + ".reg .b16 g, b, x, w;", + ".reg .f32 a, r;", + "cvt.u16.u32 g, $1;", + "cvt.u16.u32 x, $2;", + ] + w = "g" + if has_bias: + s += ["cvt.u16.u32 b, $3;", "add.rn.bf16 w, g, b;"] + w = "w" + if has_temp: + s += [ + f"cvt.f32.bf16 a, {w};", + "div.full.f32 a, a, $4;", + "cvt.rn.bf16.f32 w, a;", + ] + w = "w" + s += [ + f"cvt.f32.bf16 a, {w};", + "mul.rn.f32 a, a, 0fBFB8AA3B;", + "ex2.approx.ftz.f32 a, a;", + "add.rn.f32 a, a, 0f3F800000;", + "rcp.approx.ftz.f32 r, a;", + "cvt.rn.bf16.f32 w, r;", + ] + if has_floor: + s += [ + "cvt.f32.bf16 a, w;", + "mul.rn.f32 a, a, $5;", + "cvt.rn.bf16.f32 w, a;", + "cvt.f32.bf16 a, w;", + "add.rn.f32 a, a, $6;", + "cvt.rn.bf16.f32 w, a;", + ] + s += ["mul.rn.bf16 w, w, x;", "mov.b32 $0, {w, w};", "}"] + return "\n".join(s) + + +@dsl_user_op +def gate_one( + g, x, b, temp, c1, flo, has_bias, has_temp, has_floor, *, loc=None, ip=None +) -> Uint32: + ops = [ + Uint32(g).ir_value(loc=loc, ip=ip), + Uint32(x).ir_value(loc=loc, ip=ip), + Uint32(b).ir_value(loc=loc, ip=ip), + Float32(temp).ir_value(loc=loc, ip=ip), + Float32(c1).ir_value(loc=loc, ip=ip), + Float32(flo).ir_value(loc=loc, ip=ip), + ] + return Uint32( + _asm( + T.i32(), + ops, + _gate_ptx_scalar(has_bias, has_temp, has_floor), + "=r,r,r,r,f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def gate_one_default_bias(g, x, b, *, loc=None, ip=None) -> Uint32: + ops = [ + Uint32(g).ir_value(loc=loc, ip=ip), + Uint32(x).ir_value(loc=loc, ip=ip), + Uint32(b).ir_value(loc=loc, ip=ip), + ] + return Uint32( + _asm(T.i32(), ops, _gate_ptx_scalar(True, False, False), "=r,r,r,r", loc, ip) + ) + + +@dsl_user_op +def gate_one_default_nobias(g, x, *, loc=None, ip=None) -> Uint32: + ops = [Uint32(g).ir_value(loc=loc, ip=ip), Uint32(x).ir_value(loc=loc, ip=ip)] + return Uint32( + _asm(T.i32(), ops, _gate_ptx_scalar(False, False, False), "=r,r,r", loc, ip) + ) + + +def gate_one_selected(g, x, b, temp, c1, flo, has_bias, has_temp, has_floor): + if const_expr(not has_temp and not has_floor): + if const_expr(has_bias): + return gate_one_default_bias(g, x, b) + return gate_one_default_nobias(g, x) + return gate_one(g, x, b, temp, c1, flo, has_bias, has_temp, has_floor) + + +def _ld_words(addr, VEC, PF256): + if const_expr(VEC == 8): + return ld_v4(addr, ".L2::256B" if PF256 else "") + elif const_expr(VEC == 4): + return ld_v2(addr, ".L2::256B" if PF256 else ".L2::128B") + elif const_expr(VEC == 2): + return ld_b32(addr) + return [ld_b16_l2_128(addr, ".L2::256B" if PF256 else ".L2::128B")] + + +@cute.kernel +def _gate2d_kernel( + x_addr: Int64, + g_addr: Int64, + b_addr: Int64, + o_addr: Int64, + n_rows: Int32, + temp: Float32, + c1: Float32, + flo: Float32, + W: cutlass.Constexpr, + D: cutlass.Constexpr, + G_STRIDE: cutlass.Constexpr, + VEC: cutlass.Constexpr, + E: cutlass.Constexpr, + R: cutlass.Constexpr, + BLOCK: cutlass.Constexpr, + BIAS_MODE: cutlass.Constexpr, + HAS_TEMP: cutlass.Constexpr, + HAS_FLOOR: cutlass.Constexpr, + ROW_EXACT: cutlass.Constexpr, + ONE_ROW: cutlass.Constexpr, + ROWS_ON_X: cutlass.Constexpr, + PF256: cutlass.Constexpr, + INTERLEAVED8: cutlass.Constexpr, +): + tidx, _, _ = cute.arch.thread_idx() + bx, by, _ = cute.arch.block_idx() + if const_expr(not ROWS_ON_X): + bx, by = by, bx + cta_cols = BLOCK * VEC * E + col_exact = W % cta_cols == 0 + nwords = (VEC + 1) // 2 + has_bias = BIAS_MODE != 0 + row0 = Uint32(0) if const_expr(ONE_ROW) else Uint32(bx) * Uint32(R) + colbase = Uint32(by) * Uint32(cta_cols) + Uint32(tidx) * Uint32(VEC) + + cols = [] + colcs = [] + bws = [] + for e in cutlass.range_constexpr(E): + col = colbase + Uint32(e * BLOCK * VEC) + colc = col if const_expr(col_exact) else cutlass.min(col, Uint32(W - VEC)) + cols.append(col) + colcs.append(colc) + if const_expr(BIAS_MODE == 1): + bw = [ld_bias_b16x2(b_addr + Int64(colc // Uint32(D)) * 2)] * nwords + elif const_expr(BIAS_MODE == 2): + bw = [] + for j in cutlass.range_constexpr(nwords): + lo = ld_b16(b_addr + Int64((colc + Uint32(2 * j)) // Uint32(D)) * 2) + if const_expr(VEC > 1): + hi = ld_b16( + b_addr + Int64((colc + Uint32(2 * j + 1)) // Uint32(D)) * 2 + ) + bw.append(lo | (hi << 16)) + else: + bw.append(lo) + else: + bw = [Uint32(0)] * nwords + bws.append(bw) + + rows = [] + xs = [] + gs = [] + for r in cutlass.range_constexpr(R): + row = Uint32(0) if const_expr(ONE_ROW) else row0 + Uint32(r) + rowc = ( + row + if const_expr(ONE_ROW or ROW_EXACT) + else cutlass.min(row, Uint32(n_rows) - Uint32(1)) + ) + rows.append(row) + if const_expr(ONE_ROW): + x_row = x_addr + g_row = g_addr + else: + x_row = x_addr + Int64(rowc) * (W * 2) + g_row = g_addr + Int64(rowc) * (G_STRIDE * 2) + xr = [] + gr = [] + for e in cutlass.range_constexpr(E): + xr.append(_ld_words(x_row + Int64(colcs[e]) * 2, VEC, PF256)) + gr.append(_ld_words(g_row + Int64(colcs[e]) * 2, VEC, PF256)) + xs.append(xr) + gs.append(gr) + + outs = [] + if const_expr(VEC == 1): + flat_g = [] + flat_x = [] + flat_b = [] + for r in cutlass.range_constexpr(R): + for e in cutlass.range_constexpr(E): + flat_g.append(gs[r][e][0]) + flat_x.append(xs[r][e][0]) + flat_b.append(bws[e][0]) + n = R * E + flat_o = [] + for i in cutlass.range_constexpr(0, n, 2): + if const_expr(i + 1 < n): + rr = gate_pair_selected( + flat_g[i] | (flat_g[i + 1] << 16), + flat_x[i] | (flat_x[i + 1] << 16), + (flat_b[i] & Uint32(0xFFFF)) | (flat_b[i + 1] << 16), + temp, + c1, + flo, + has_bias, + HAS_TEMP, + HAS_FLOOR, + ) + flat_o.append(rr) + flat_o.append(rr >> 16) + else: + flat_o.append( + gate_one_selected( + flat_g[i], + flat_x[i], + flat_b[i], + temp, + c1, + flo, + has_bias, + HAS_TEMP, + HAS_FLOOR, + ) + ) + for r in cutlass.range_constexpr(R): + outs.append([[flat_o[r * E + e]] for e in range(E)]) + else: + for r in cutlass.range_constexpr(R): + orow = [] + for e in cutlass.range_constexpr(E): + if const_expr( + VEC >= 4 and not HAS_TEMP and not HAS_FLOOR and BIAS_MODE == 1 + ): + if const_expr(VEC == 4): + ow = gate_vec_exp_interleaved_default( + gs[r][e], xs[r][e], bws[e][0] + ) + elif const_expr(INTERLEAVED8): + ow = gate_vec_exp_interleaved_default( + gs[r][e][0:2], xs[r][e][0:2], bws[e][0] + ) + gate_vec_exp_interleaved_default( + gs[r][e][2:4], xs[r][e][2:4], bws[e][0] + ) + else: + ow = gate_vec_default(gs[r][e], xs[r][e], bws[e][0]) + elif const_expr( + VEC >= 4 and not HAS_TEMP and not HAS_FLOOR and BIAS_MODE == 0 + ): + if const_expr(VEC == 4): + ow = gate_vec_exp_interleaved_default(gs[r][e], xs[r][e], None) + elif const_expr(INTERLEAVED8): + ow = gate_vec_exp_interleaved_default( + gs[r][e][0:2], xs[r][e][0:2], None + ) + gate_vec_exp_interleaved_default( + gs[r][e][2:4], xs[r][e][2:4], None + ) + else: + ow = gate_vec_default(gs[r][e], xs[r][e], None) + else: + ow = [] + for j in cutlass.range_constexpr(nwords): + ow.append( + gate_pair_selected( + gs[r][e][j], + xs[r][e][j], + bws[e][j], + temp, + c1, + flo, + has_bias, + HAS_TEMP, + HAS_FLOOR, + ) + ) + orow.append(ow) + outs.append(orow) + + st_pol = make_policy("evict_last") if const_expr(VEC == 8) else None + for r in cutlass.range_constexpr(R): + o_row = o_addr if const_expr(ONE_ROW) else o_addr + Int64(rows[r]) * (W * 2) + for e in cutlass.range_constexpr(E): + addr = o_row + Int64(cols[e]) * 2 + if const_expr(ROW_EXACT and col_exact): + if const_expr(VEC == 1): + st_b16(addr, outs[r][e][0]) + elif const_expr(st_pol is not None): + st_v4_hint(addr, outs[r][e], st_pol) + else: + st_vec(addr, outs[r][e]) + else: + if const_expr(ROW_EXACT): + ok = cols[e] < Uint32(W) + elif const_expr(col_exact): + ok = rows[r] < Uint32(n_rows) + else: + ok = (rows[r] < Uint32(n_rows)) & (cols[e] < Uint32(W)) + if ok: + if const_expr(VEC == 1): + st_b16(addr, outs[r][e][0]) + elif const_expr(st_pol is not None): + st_v4_hint(addr, outs[r][e], st_pol) + else: + st_vec(addr, outs[r][e]) + + +@cute.jit +def _gate2d_launch( + x_addr: Int64, + g_addr: Int64, + b_addr: Int64, + o_addr: Int64, + n_rows: Int32, + temp: Float32, + c1: Float32, + flo: Float32, + stream: cuda.CUstream, + W: cutlass.Constexpr, + D: cutlass.Constexpr, + G_STRIDE: cutlass.Constexpr, + VEC: cutlass.Constexpr, + E: cutlass.Constexpr, + R: cutlass.Constexpr, + BLOCK: cutlass.Constexpr, + BIAS_MODE: cutlass.Constexpr, + HAS_TEMP: cutlass.Constexpr, + HAS_FLOOR: cutlass.Constexpr, + ROW_EXACT: cutlass.Constexpr, + ONE_ROW: cutlass.Constexpr, + ROWS_ON_X: cutlass.Constexpr, + CLUSTER_LAUNCH: cutlass.Constexpr, + PF256: cutlass.Constexpr, + INTERLEAVED8: cutlass.Constexpr, +): + ncb = (W + BLOCK * VEC * E - 1) // (BLOCK * VEC * E) + nrb = (n_rows + (R - 1)) // R + grid = [nrb, ncb, 1] if const_expr(ROWS_ON_X) else [ncb, nrb, 1] + kernel = _gate2d_kernel( + x_addr, + g_addr, + b_addr, + o_addr, + n_rows, + temp, + c1, + flo, + W, + D, + G_STRIDE, + VEC, + E, + R, + BLOCK, + BIAS_MODE, + HAS_TEMP, + HAS_FLOOR, + ROW_EXACT, + ONE_ROW, + ROWS_ON_X, + PF256, + INTERLEAVED8, + ) + # Small grids: launching as (trivial 1x1x1) clusters measured ~0.1 us lower end-to- + # end + # latency on B300, and adding the cooperative (co-resident) attribute on top of that + # a + # further 1-3 timer ticks; large grids lose CTA launch throughput this way. + if const_expr(CLUSTER_LAUNCH == 2): + kernel.launch( + grid=grid, + block=[BLOCK, 1, 1], + cluster=[1, 1, 1], + cooperative=True, + stream=stream, + preferred_smem_carveout=0, + min_blocks_per_mp=1, + ) + elif const_expr(CLUSTER_LAUNCH == 1): + kernel.launch( + grid=grid, + block=[BLOCK, 1, 1], + cluster=[1, 1, 1], + stream=stream, + preferred_smem_carveout=0, + ) + else: + kernel.launch( + grid=grid, block=[BLOCK, 1, 1], stream=stream, preferred_smem_carveout=0 + ) + + +_cache: dict[tuple, Callable] = {} +_no_coop = set() +_CLUSTER_MAX_CTAS = 768 +# Cooperative 1x1x1-cluster launches were measured to be refused above ~1-1.5K CTAs on +# B300 +# even when plain occupancy would allow more, so cap them here (and fall back if +# refused). +_COOP_MAX_CTAS = 1024 +_sm_count: list[int] = [] + + +def _num_sms(): + if not _sm_count: + _sm_count.append( + torch.cuda.get_device_properties( + torch.accelerator.current_device_index() + ).multi_processor_count + ) + return _sm_count[0] + + +def _max_vec(W, gs, ptrs): + for vec in (8, 4, 2): + if W % vec == 0 and gs % vec == 0 and all(p % (2 * vec) == 0 for p in ptrs): + return vec + return 1 + + +def _coop_fits(ncta, block, regs): + """Conservative co-residency bound (no shared memory; regs rounded up per warp).""" + per_warp = -(-regs * 32 // 256) * 256 + by_regs = 65536 // (per_warp * -(-block // 32)) + per_sm = min(32, 2048 // block, by_regs) + return ncta <= min(_COOP_MAX_CTAS, _num_sms() * per_sm) + + +def _choose(T_, W, maxvec, contiguous_gate): + """Return launch parameters and an optional cooperative register bound.""" + n = T_ * W + if n <= 512 * 1024: + # Decode/small-prefill sizes (any floor/temperature). Scalar lanes in 96-ish + # CTAs + # up to 64K elements, then 8-byte packets with ~1.5-3K warps in flight. + pf256 = False + if n <= 16 * 1024: + vec, e, block_cap = 1, 1, 128 + elif n <= 32 * 1024: + # A 256-thread scalar CTA spans 512 B: request 256 B per L2 miss. + vec, e, block_cap, pf256 = 1, 1, 256, True + elif n <= 64 * 1024: + # 512-thread scalar CTAs each cover 1 KB of a row: fetch 256 B per L2 miss. + vec, e, block_cap, pf256 = 1, 1, 512, True + elif n <= 128 * 1024: + vec, e, block_cap = 4, 1, 256 + elif n <= 256 * 1024: + vec, e, block_cap = 4, 1, 512 + else: + vec, e, block_cap = 4, 1, 384 + vec = min(vec, maxvec) + if vec != 1: + pf256 = False + need = -(-W // (vec * e)) + return vec, e, 1, min(block_cap, max(32, -(-need // 32) * 32)), pf256, None + if n <= 32 * 1024 * 1024: + vec, e, r = 8, 1, 1 + else: + vec, e, r = 8, 1, 1 if contiguous_gate else 2 + vec = min(vec, maxvec) + if vec == 8 and n > 1024 * 1024: + # Two rows per 192-thread CTA keep the T=512 primary grid co-resident. + ncb = -(-W // (192 * 8)) + if ncb * 192 * 8 - W < 192 * 8 // 2 and _coop_fits(ncb * -(-T_ // 2), 192, 32): + return 8, 1, 2, 192, False, 32 + need = -(-W // (vec * e)) + block_cap = 256 if vec == 8 and 512 * 1024 <= n <= 1024 * 1024 else 128 + if ( + vec == 8 + and n > 1024 * 1024 + and -(-T_ // r) * -(-W // (256 * vec * e)) <= _CLUSTER_MAX_CTAS + ): + block_cap = 256 # keeps ~1-1.5M-element grids within the cluster-launch tier + block = min(block_cap, max(32, -(-need // 32) * 32)) + return vec, e, r, block, False, None + + +def run(attn_output, gate, gate_bias, num_heads, head_dim, floor, temperature, output): + T_ = attn_output.shape[0] + H = int(num_heads) + D = int(head_dim) + W = H * D + if T_ == 0 or W == 0: + return + gs = gate.stride(0) if T_ > 1 else max(gate.stride(0), W) + has_bias = gate_bias is not None + floor = float(floor) + temperature = float(temperature) + has_temp = temperature != 1.0 + c1 = 1.0 - floor + has_floor = c1 != 1.0 or floor != 0.0 + ptrs = [attn_output.data_ptr(), gate.data_ptr(), output.data_ptr()] + b_ptr = gate_bias.data_ptr() if has_bias else 0 + + vec, e, r, block, pf256, coop_regs = _choose(T_, W, _max_vec(W, gs, ptrs), gs == W) + bias_mode = 0 if not has_bias else (1 if D % vec == 0 else 2) + nrb = -(-T_ // r) + compile_stride = W if T_ == 1 else gs + ncta = nrb * -(-W // (block * vec * e)) + launch_mode = 0 + if coop_regs is not None: + launch_mode = 2 + elif ncta <= _CLUSTER_MAX_CTAS: + # Cooperative launch needs every CTA resident at once; only claim it at <= half + # the + # SM's thread capacity so it cannot fail for any register footprint used here. + launch_mode = 2 if ncta <= _num_sms() * max(1, 1024 // block) else 1 + # T=512's two-row cooperative tile exposes the final packets' math at the one-wave + # tail. + # Interleave independent exponential sigmoid operations for these packets. + interleaved8 = ( + vec == 8 + and r == 2 + and launch_mode == 2 + and not has_temp + and not has_floor + and bias_mode != 2 + ) + key = ( + W, + D, + compile_stride, + vec, + e, + r, + block, + bias_mode, + has_temp, + has_floor, + T_ % r == 0, + T_ == 1, + nrb > 65535, + launch_mode, + pf256, + interleaved8, + ) + if key in _no_coop: + key = key[:13] + (1, pf256, interleaved8) + args = ( + Int64(ptrs[0]), + Int64(ptrs[1]), + Int64(b_ptr), + Int64(ptrs[2]), + Int32(T_), + Float32(temperature), + Float32(c1), + Float32(floor), + cuda.CUstream(current_stream().cuda_stream), + ) + fn = _cache.get(key) + if fn is None: + fn = cute.compile(_gate2d_launch, *args, *key) + _cache[key] = fn + if launch_mode == 2 and coop_regs is not None: + try: + fn(*args) + except ( + Exception + ) as exc: # cooperative grid refused by the driver: launch clustered + if "COOPERATIVE_LAUNCH_TOO_LARGE" not in str(exc): + raise + _no_coop.add(key) + key = key[:13] + (1, pf256, interleaved8) + fn = _cache.get(key) + if fn is None: + fn = cute.compile(_gate2d_launch, *args, *key) + _cache[key] = fn + fn(*args) + return + fn(*args) diff --git a/vllm/models/beam/nvidia/dflash2.py b/vllm/models/beam/nvidia/dflash2.py new file mode 100644 index 000000000000..adc45210c540 --- /dev/null +++ b/vllm/models/beam/nvidia/dflash2.py @@ -0,0 +1,520 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DFlash2 draft model for Beam targets. + +A DFlash2 drafter predicts a block of tokens at once from hidden states the +Beam target captures at `dflash_config.target_layer_ids`, then a low-rank +candidate selector rescores the top-k tokens at each block position against +the token before it. The drafter shares the target's embeddings and LM head. + +Nothing here is torch-compiled. The DFlash speculator captures the whole draft +step (decoder layers, candidate head and selector) in its own CUDA graphs, and +the context K/V precompute runs eagerly as a few fused kernels. +""" + +from collections.abc import Iterable + +import torch +import torch.nn.functional as F +from torch import nn + +from vllm import _custom_ops as ops +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Attention +from vllm.model_executor.models.qwen3_dflash2 import DFlash2Qwen3DecoderLayer +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + get_draft_quant_config, + maybe_prefix, +) + +from .dflash2_ops import candidate_edges + +logger = init_logger(__name__) + + +class DFlash2BeamAttention(DFlashQwen3Attention): + def forward( + self, positions: torch.Tensor, hidden_states: torch.Tensor + ) -> torch.Tensor: + from .attention_ops import beam_qk_norm_rope + + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + if self.v_scale is not None: + v = v * self.v_scale + cache, slots = None, None + mapping = get_forward_context().slot_mapping + assert isinstance(mapping, dict) + if mapping and self.attn.kv_cache.dtype == q.dtype == torch.bfloat16: + slots = mapping[self.attn.layer_name] + cache = self.attn.kv_cache + assert self.attn.attn_backend.get_name() in ("FLASH_ATTN", "FLASHINFER") + assert self.attn.impl.dcp_world_size == 1 + assert self.attn.kv_sharing_target_layer_name is None + assert cache.dtype == q.dtype == torch.bfloat16 + assert cache.ndim == 4 + rope = self.rotary_emb + q, k = beam_qk_norm_rope( + q, + k, + self.q_norm.weight, + self.k_norm.weight, + self.q_norm.variance_epsilon, + positions, + rope.cos_sin_cache, + self.head_dim, + rope.rotary_dim, + rope.is_neox_style, + value=v, + kv_cache=cache, + slot_mapping=slots, + ) + output = self.attn(q, k, v) if cache is None else self.attn(q, None, None) + output, _ = self.o_proj(output) + return output + + +class DFlash2BeamDecoderLayer(DFlash2Qwen3DecoderLayer): + attention_cls = DFlash2BeamAttention + + +class DFlash2BeamCandidateSelector(nn.Module): + """Scores each candidate token against every candidate one position back. + + For block position `p`, candidate `c` and predecessor `q`, the score is the + candidate's own logit plus ``, where + `h_p` is the drafter's hidden state projected to the codebook rank. The + predecessor of the first position is the anchor token. + """ + + def __init__( + self, + hidden_size: int, + vocab_size: int, + rank: int, + top_k: int, + params_dtype: torch.dtype, + prefix: str, + ) -> None: + super().__init__() + self.top_k = top_k + self.predecessor_codebook = nn.Parameter( + torch.empty(vocab_size, rank, dtype=params_dtype), requires_grad=False + ) + self.successor_codebook = nn.Parameter( + torch.empty(vocab_size, rank, dtype=params_dtype), requires_grad=False + ) + self.hidden_projection = ReplicatedLinear( + hidden_size, + rank, + bias=False, + params_dtype=params_dtype, + quant_config=None, + prefix=maybe_prefix(prefix, "hidden_projection"), + return_bias=False, + ) + + def forward( + self, + candidate_ids: torch.Tensor, + unary_logits: torch.Tensor, + hidden_states: torch.Tensor, + anchor_token_ids: torch.Tensor, + ) -> torch.Tensor: + """Edge scores `[requests, steps, top_k (predecessor), top_k]`.""" + assert hidden_states.dtype == torch.bfloat16 + return candidate_edges( + self.predecessor_codebook, + self.successor_codebook, + candidate_ids, + unary_logits, + self.hidden_projection(hidden_states), + anchor_token_ids, + self.top_k, + ) + + +class DFlash2BeamModel(nn.Module): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + }, + ) + + def __init__( + self, + *, + vllm_config: VllmConfig, + start_layer_id: int = 0, + prefix: str = "", + ) -> None: + super().__init__() + assert vllm_config.speculative_config is not None + config = vllm_config.speculative_config.draft_model_config.hf_config + self.config = config + draft_config = config.dflash_config + if getattr(config, "attention_bias", False) or get_draft_quant_config( + vllm_config + ): + raise NotImplementedError( + "The Beam DFlash2 drafter supports only unquantized drafts " + "without attention biases." + ) + params_dtype = vllm_config.model_config.dtype + self.input_embedding_scale = float( + draft_config.get("input_embedding_scale", 1.0) + ) + + # Replaced by the target's embedding when the speculator loads the draft. + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + + # Numbered after the target's layers so the attention layer names, which + # key the KV cache, do not collide with the target's. + current_vllm_config = get_current_vllm_config() + self.layers = nn.ModuleList( + [ + DFlash2BeamDecoderLayer( + current_vllm_config, + config=config, + layer_idx=layer_idx, + cache_config=current_vllm_config.cache_config, + quant_config=None, + prefix=maybe_prefix(prefix, f"layers.{layer_idx + start_layer_id}"), + ) + for layer_idx in range(config.num_hidden_layers) + ] + ) + + target_hidden_size = ( + getattr(config, "target_hidden_size", None) or config.hidden_size + ) + self.fc = ReplicatedLinear( + target_hidden_size * len(draft_config["target_layer_ids"]), + config.hidden_size, + bias=False, + params_dtype=params_dtype, + quant_config=None, + prefix=maybe_prefix(prefix, "fc"), + return_bias=False, + ) + self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self._context_cache_ptrs: torch.Tensor | None = None + self._context_scale_ptrs: torch.Tensor | None = None + self._context_cache_addresses: tuple[int, ...] | None = None + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.candidate_selector = DFlash2BeamCandidateSelector( + hidden_size=config.hidden_size, + vocab_size=config.vocab_size, + rank=int(draft_config["selector_rank"]), + top_k=int(draft_config["selector_top_k"]), + params_dtype=params_dtype, + prefix=maybe_prefix(prefix, "candidate_selector"), + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + embeds = self.embed_tokens(input_ids) + if self.input_embedding_scale != 1.0: + embeds = embeds * self.input_embedding_scale + return embeds + + def build_context_kv_weights(self) -> None: + """Stack every layer's K/V projection and K-norm weight. + + The context K/V for all layers then comes from one GEMM and one grouped + norm. Must run after the weights are loaded, since it copies them. + """ + attns = [layer.self_attn for layer in self.layers] + self._context_kv_weight = torch.cat( + [attn.qkv_proj.weight[attn.q_size :] for attn in attns] + ) + self._context_k_norm_weight = torch.stack( + [attn.k_norm.weight for attn in attns] + ).contiguous() + + def precompute_and_store_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, + ) -> None: + """Project the combined target states to each layer's K/V and cache them. + + The drafter's queries attend to these context keys, which carry the + layer's K-norm and RoPE like the queries' own keys. Profiling uses a + `None` slot mapping to run the projection without writing the cache. + """ + if not hasattr(self, "_context_kv_weight"): + # Dummy weights skip `load_weights`. + logger.warning_once( + "The Beam DFlash2 context K/V weights were built lazily; this is " + "expected only with dummy weights." + ) + self.build_context_kv_weights() + + attn0 = self.layers[0].self_attn + num_layers = len(self.layers) + num_ctx = context_states.shape[0] + num_kv_heads, head_dim = attn0.num_kv_heads, attn0.head_dim + + normed = torch.empty_like(context_states) + ops.rms_norm( + normed, + context_states, + self.hidden_norm.weight, + self.hidden_norm.variance_epsilon, + ) + kv = F.linear(normed, self._context_kv_weight).view( + num_ctx, num_layers, 2, num_kv_heads, head_dim + ) + self._store_fused_context_kv(kv, context_positions, context_slot_mapping) + + def _store_fused_context_kv( + self, + kv: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor | list[torch.Tensor | None] | None, + ) -> None: + from .dflash2_context_ops import store_context_kv + + if slot_mapping is None: + return + if isinstance(slot_mapping, (list, tuple)): + slots = slot_mapping[0] + if slots is None: + assert all(s is None for s in slot_mapping) + return + assert all( + s is not None + and s.data_ptr() == slots.data_ptr() + and s.shape == slots.shape + and s.stride() == slots.stride() + for s in slot_mapping + ) + else: + slots = slot_mapping + first = self.layers[0].self_attn + rotary = first.rotary_emb + assert kv.dtype == torch.bfloat16 + assert rotary.rotary_dim == first.head_dim + assert rotary.cos_sin_cache.dtype == kv.dtype + assert slots.stride(0) == 1 + attns = [layer.self_attn.attn for layer in self.layers] + caches = [attn.kv_cache for attn in attns] + cache = caches[0] + cache_dtype = attns[0].kv_cache_dtype + fp8_dtype = { + "fp8": torch.float8_e4m3fn, + "fp8_e4m3": torch.float8_e4m3fn, + "fp8_e5m2": torch.float8_e5m2, + }.get(cache_dtype) + for attn, c in zip(attns, caches): + assert attn.attn_backend.get_name() in ("FLASH_ATTN", "FLASHINFER") + assert attn.impl.dcp_world_size == 1 + assert attn.kv_sharing_target_layer_name is None + assert attn.kv_cache_dtype == cache_dtype + assert c.dtype == cache.dtype and c.device == kv.device + assert c.dtype == kv.dtype or fp8_dtype is not None + assert c.ndim == 4 and c.stride(-1) == 1 + assert c.shape[1] == kv.shape[3] + assert c.shape[-1] == 2 * kv.shape[-1] + assert c.shape == cache.shape and c.stride() == cache.stride() + cache_pointers = tuple(c.data_ptr() for c in caches) + scale_pointers = ( + tuple( + s.data_ptr() for attn in attns for s in (attn._k_scale, attn._v_scale) + ) + if fp8_dtype is not None + else () + ) + pointers = cache_pointers + scale_pointers + if self._context_cache_addresses != pointers: + self._context_cache_ptrs = torch.tensor( + cache_pointers, dtype=torch.int64, device=kv.device + ) + self._context_scale_ptrs = ( + torch.tensor(scale_pointers, dtype=torch.int64, device=kv.device) + if scale_pointers + else None + ) + self._context_cache_addresses = pointers + assert self._context_cache_ptrs is not None + assert self._context_cache_ptrs.device == kv.device + store_context_kv( + kv, + self._context_k_norm_weight, + positions, + rotary.cos_sin_cache, + self._context_cache_ptrs, + slots, + cache.view(fp8_dtype) if fp8_dtype is not None else cache, + first.k_norm.variance_epsilon, + first.v_scale if first.v_scale is not None else 1.0, + rotary.is_neox_style, + self._context_scale_ptrs, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + residual = inputs_embeds + hidden_states = self.layers[0].input_layernorm(inputs_embeds) + for i, layer in enumerate(self.layers): + next_norm = ( + self.layers[i + 1].input_layernorm + if i + 1 < len(self.layers) + else self.norm + ) + hidden_states, residual = _forward_normalized( + layer, positions, hidden_states, residual, next_norm + ) + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + +def _forward_normalized( + layer: DFlash2Qwen3DecoderLayer, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor, + next_norm: RMSNorm, +) -> tuple[torch.Tensor, torch.Tensor]: + """Carry normalized activations across layers to fuse both convolutions.""" + from .dflash2_conv_ops import conv_residual_norm + + hidden_states, coefficients = layer.attention_conv.prepare(hidden_states) + hidden_states = layer.self_attn(positions, hidden_states) + hidden_states, residual = conv_residual_norm( + hidden_states, + coefficients, + layer.attention_conv.base_kernel[1], + residual, + layer.post_attention_layernorm.weight, + layer.post_attention_layernorm.variance_epsilon, + layer.attention_conv.block_size, + layer.attention_conv.group_size, + ) + hidden_states, coefficients = layer.mlp_conv.prepare(hidden_states) + hidden_states = layer.mlp(hidden_states) + return conv_residual_norm( + hidden_states, + coefficients, + layer.mlp_conv.base_kernel[1], + residual, + next_norm.weight, + next_norm.variance_epsilon, + layer.mlp_conv.block_size, + layer.mlp_conv.group_size, + ) + + +class DFlash2BeamForCausalLM(nn.Module): + # The checkpoint ships neither, so the speculator always aliases the target's. + has_own_embed_tokens = False + has_own_lm_head = False + # Full vocabulary: draft ids are target ids. + draft_id_to_target_id = None + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + self.config = vllm_config.speculative_config.draft_model_config.hf_config + self.model = DFlash2BeamModel( + vllm_config=vllm_config, + start_layer_id=vllm_config.model_config.get_total_num_hidden_layers(), + prefix=maybe_prefix(prefix, "model"), + ) + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + draft_config = self.config.dflash_config + softcap = float(draft_config.get("final_logit_softcapping") or 0.0) + self.logits_processor = LogitsProcessor( + self.config.vocab_size, + scale=float(draft_config.get("output_multiplier", 1.0)), + soft_cap=softcap if softcap > 0 else None, + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.model(input_ids, positions, inputs_embeds) + + def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Project the concatenated target-layer states to the drafter's width.""" + expected = self.model.fc.input_size + if hidden_states.shape[-1] != expected: + raise ValueError( + f"The DFlash2 drafter expects {expected} concatenated target " + f"features but received {hidden_states.shape[-1]}; its " + "target_layer_ids do not match what the target captures." + ) + return self.model.fc(hidden_states) + + def precompute_and_store_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, + ) -> None: + self.model.precompute_and_store_context_kv( + context_states, context_positions, context_slot_mapping + ) + + def get_draft_kv_cache_layer_names(self) -> list[str]: + return [layer.self_attn.attn.layer_name for layer in self.model.layers] + + def get_draft_attn_causal(self) -> list[bool]: + """Per-layer causality, aligned with `get_draft_kv_cache_layer_names`.""" + return [layer.self_attn.causal for layer in self.model.layers] + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.logits_processor(self.lm_head, hidden_states) + + def compute_candidates( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """The selector's top-k candidate ids and their logits at each position.""" + return self.logits_processor.get_top_k_tokens( + self.lm_head, hidden_states, self.model.candidate_selector.top_k + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loaded = {f"model.{name}" for name in self.model.load_weights(weights)} + self.model.build_context_kv_weights() + # Aliased to the target's after loading. + return loaded | {"model.embed_tokens.weight", "lm_head.weight"} diff --git a/vllm/models/beam/nvidia/dflash2_context_ops.py b/vllm/models/beam/nvidia/dflash2_context_ops.py new file mode 100644 index 000000000000..e96d1bbd6c55 --- /dev/null +++ b/vllm/models/beam/nvidia/dflash2_context_ops.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit(do_not_specialize=["ROWS", "NUM_SLOTS"]) +def _context_kv_kernel( + KV, + WEIGHT, + POSITIONS, + ROPE, + CACHE_PTRS, + CACHE, + SCALE_PTRS, + SLOTS, + ROWS, + NUM_SLOTS, + LAYERS: tl.constexpr, + HEADS: tl.constexpr, + DIM: tl.constexpr, + BLOCK_STRIDE: tl.constexpr, + HEAD_STRIDE: tl.constexpr, + TOKEN_STRIDE: tl.constexpr, + PAGE: tl.constexpr, + EPS: tl.constexpr, + VALUE_SCALE: tl.constexpr, + NEOX: tl.constexpr, + FP8: tl.constexpr, + FP8_MAX: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) * 4 + tl.arange(0, 4) + layer_head = tl.program_id(1) + layer = layer_head // HEADS + head = layer_head % HEADS + col = tl.arange(0, BLOCK) + slot = tl.load(SLOTS + row, (row < ROWS) & (row < NUM_SLOTS), -1).to(tl.int64) + mask = (row[:, None] < ROWS) & (slot[:, None] >= 0) & (col[None, :] < DIM) + offsets = ( + row[:, None] * (LAYERS * 2 * HEADS * DIM) + + (layer * 2 * HEADS + head) * DIM + + col[None, :] + ) + key = tl.load(KV + offsets, mask, 0).to(tl.float32) + if BLOCK == 128: + squares = (key * key).reshape((4, 4, 32)) + total = tl.sum(tl.sum(squares, 2), 1) + else: + total = tl.sum(key * key, 1) + weight = tl.load(WEIGHT + layer * DIM + col, col < DIM, 0).to(tl.float32) + key = key * tl.rsqrt(total / DIM + EPS)[:, None] * weight[None, :] + key = key.to(KV.dtype.element_ty).to(tl.float32) + half: tl.constexpr = DIM // 2 + if NEOX: + partner = (col + half) % DIM + second = col >= half + rope_col = col % half + else: + partner = col ^ 1 + second = col % 2 != 0 + rope_col = col // 2 + other = tl.gather(key, tl.broadcast_to(partner[None, :], (4, BLOCK)), 1) + position = tl.load(POSITIONS + row, (row < ROWS) & (slot >= 0), 0).to(tl.int64) + cosine = tl.load(ROPE + position[:, None] * DIM + rope_col[None, :], mask, 0) + sine = tl.load(ROPE + position[:, None] * DIM + half + rope_col[None, :], mask, 0) + key = tl.where( + second[None, :], + tl.fma(other, sine, key * cosine), + tl.fma(key, cosine, -other * sine), + ) + cache = tl.load(CACHE_PTRS + layer).to(tl.pointer_type(CACHE.dtype.element_ty)) + address = ( + slot[:, None] // PAGE * BLOCK_STRIDE + + slot[:, None] % PAGE * TOKEN_STRIDE + + head * HEAD_STRIDE + + col[None, :] + ) + value = tl.load(KV + offsets + HEADS * DIM, mask, 0).to(tl.float32) + value = value * VALUE_SCALE + if FP8: + # Preserve BF16 rounding after RoPE and value scaling. + key = key.to(KV.dtype.element_ty).to(tl.float32) + value = value.to(KV.dtype.element_ty).to(tl.float32) + k_scale = tl.load(SCALE_PTRS + layer * 2).to(tl.pointer_type(tl.float32)) + v_scale = tl.load(SCALE_PTRS + layer * 2 + 1).to(tl.pointer_type(tl.float32)) + key = tl.minimum(tl.maximum(key / tl.load(k_scale), -FP8_MAX), FP8_MAX) + value = tl.minimum(tl.maximum(value / tl.load(v_scale), -FP8_MAX), FP8_MAX) + tl.store(cache + address, key, mask) + tl.store(cache + address + DIM, value, mask) + + +def store_context_kv( + kv: torch.Tensor, + weight: torch.Tensor, + positions: torch.Tensor, + rope: torch.Tensor, + cache_ptrs: torch.Tensor, + slots: torch.Tensor, + cache: torch.Tensor, + eps: float, + value_scale: float, + is_neox: bool, + scale_ptrs: torch.Tensor | None = None, +) -> None: + """Normalize, rotate, and scatter all draft layers' context K/V.""" + rows, layers, _, heads, dim = kv.shape + fp8 = cache.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) + assert cache.dtype == kv.dtype or fp8 + assert not fp8 or scale_ptrs is not None + if rows: + _context_kv_kernel[(triton.cdiv(rows, 4), layers * heads)]( + kv, + weight, + positions, + rope, + cache_ptrs, + cache, + scale_ptrs, + slots, + rows, + slots.numel(), + layers, + heads, + dim, + cache.stride(0), + cache.stride(1), + cache.stride(2), + cache.shape[2], + eps, + value_scale, + is_neox, + fp8, + torch.finfo(cache.dtype).max if fp8 else 0.0, + triton.next_power_of_2(dim), + num_warps=4, + enable_fp_fusion=False, + ) diff --git a/vllm/models/beam/nvidia/dflash2_conv_ops.py b/vllm/models/beam/nvidia/dflash2_conv_ops.py new file mode 100644 index 000000000000..4ce6508b7040 --- /dev/null +++ b/vllm/models/beam/nvidia/dflash2_conv_ops.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _conv_residual_norm_kernel( + X, + DELTA, + BASE, + RESIDUAL, + WEIGHT, + OUTPUT, + X_STRIDE: tl.constexpr, + DELTA_ROW_STRIDE: tl.constexpr, + DELTA_TAP_STRIDE: tl.constexpr, + RESIDUAL_STRIDE: tl.constexpr, + WIDTH: tl.constexpr, + GROUP: tl.constexpr, + TAPS: tl.constexpr, + QUERY: tl.constexpr, + EPS: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + col = tl.arange(0, BLOCK) + mask = col < WIDTH + group = col // GROUP + coefficient = tl.load(BASE + col, mask, 0).to(tl.float32) + tl.load( + DELTA + row * DELTA_ROW_STRIDE + group, mask, 0 + ).to(tl.float32) + value = tl.load(X + row * X_STRIDE + col, mask, 0).to(tl.float32) + result = coefficient * value + for tap in tl.static_range(1, TAPS): + valid = mask & (row % QUERY >= tap) + coefficient = tl.load(BASE + tap * WIDTH + col, valid, 0).to( + tl.float32 + ) + tl.load( + DELTA + row * DELTA_ROW_STRIDE + tap * DELTA_TAP_STRIDE + group, valid, 0 + ).to(tl.float32) + value = tl.load(X + (row - tap) * X_STRIDE + col, valid, 0).to(tl.float32) + result = tl.fma(coefficient, value, result) + result = result.to(X.dtype.element_ty).to(tl.float32) + residual = tl.load(RESIDUAL + row * RESIDUAL_STRIDE + col, mask, 0).to(tl.float32) + result = (result + residual).to(RESIDUAL.dtype.element_ty).to(tl.float32) + tl.store(RESIDUAL + row * RESIDUAL_STRIDE + col, result, mask) + variance = tl.sum(result * result, 0) / WIDTH + weight = tl.load(WEIGHT + col, mask, 0).to(tl.float32) + normalized = result * tl.rsqrt(variance + EPS) * weight + tl.store(OUTPUT + row * WIDTH + col, normalized, mask) + + +def conv_residual_norm( + hidden: torch.Tensor, + coefficients: torch.Tensor, + base: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, + query_len: int, + group_size: int, + *, + num_warps: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Finish grouped convolution and update the residual before RMSNorm.""" + rows, width = hidden.shape + if num_warps is None: + num_warps = 16 if rows <= 256 else (8 if rows < 4096 else 4) + output = torch.empty_like(hidden) + if rows: + _conv_residual_norm_kernel[(rows,)]( + hidden, + coefficients, + base, + residual, + weight, + output, + hidden.stride(0), + coefficients.stride(0), + coefficients.stride(1), + residual.stride(0), + width, + group_size, + base.shape[0], + query_len, + eps, + triton.next_power_of_2(width), + num_warps=num_warps, + enable_fp_fusion=False, + ) + return output, residual diff --git a/vllm/models/beam/nvidia/dflash2_ops.py b/vllm/models/beam/nvidia/dflash2_ops.py new file mode 100644 index 000000000000..ed1dbfce5c01 --- /dev/null +++ b/vllm/models/beam/nvidia/dflash2_ops.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _candidate_edges_kernel( + PREDECESSORS, + SUCCESSORS, + IDS, + UNARY, + HIDDEN, + ANCHORS, + OUTPUT, + ANCHOR_STRIDE: tl.constexpr, + STEPS: tl.constexpr, + TOP_K: tl.constexpr, + RANK: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_R: tl.constexpr, +): + position = tl.program_id(0) + predecessor = tl.program_id(1) + step = position % STEPS + if step == 0: + pred_id = tl.load(ANCHORS + position // STEPS * ANCHOR_STRIDE).to(tl.int64) + else: + pred_id = tl.load(IDS + (position - 1) * TOP_K + predecessor).to(tl.int64) + r = tl.arange(0, BLOCK_R) + c = tl.arange(0, BLOCK_K) + pred = tl.load(PREDECESSORS + pred_id.to(tl.int64) * RANK + r, r < RANK, 0) + hidden = tl.load(HIDDEN + position * RANK + r, r < RANK, 0) + # Preserve the materialized product and GEMM's output rounding. + weighted = (pred.to(tl.float32) * hidden.to(tl.float32)).to(hidden.dtype) + ids = tl.load(IDS + position * TOP_K + c, c < TOP_K, 0).to(tl.int64) + successors = tl.load( + SUCCESSORS + ids[:, None] * RANK + r[None, :], + (c[:, None] < TOP_K) & (r[None, :] < RANK), + 0, + ) + scores = tl.sum(successors.to(tl.float32) * weighted[None, :].to(tl.float32), 1) + scores = scores.to(hidden.dtype).to(tl.float32) + unary = tl.load(UNARY + position * TOP_K + c, c < TOP_K, 0).to(tl.float32) + tl.store( + OUTPUT + (position * TOP_K + predecessor) * TOP_K + c, scores + unary, c < TOP_K + ) + + +def candidate_edges( + predecessor_table: torch.Tensor, + successor_table: torch.Tensor, + candidate_ids: torch.Tensor, + unary_logits: torch.Tensor, + hidden: torch.Tensor, + anchor_token_ids: torch.Tensor, + top_k: int, +) -> torch.Tensor: + """Gather codebooks and score candidate edges without intermediate tensors.""" + batch, steps, _ = candidate_ids.shape + rank = hidden.shape[-1] + output = torch.empty( + (batch, steps, top_k, top_k), dtype=unary_logits.dtype, device=hidden.device + ) + if batch * steps: + _candidate_edges_kernel[(batch * steps, top_k)]( + predecessor_table, + successor_table, + candidate_ids, + unary_logits, + hidden, + anchor_token_ids, + output, + anchor_token_ids.stride(0), + steps, + top_k, + rank, + triton.next_power_of_2(top_k), + triton.next_power_of_2(rank), + num_warps=4, + enable_fp_fusion=False, + ) + return output diff --git a/vllm/models/beam/nvidia/layernorm.py b/vllm/models/beam/nvidia/layernorm.py new file mode 100644 index 000000000000..3539ff1be618 --- /dev/null +++ b/vllm/models/beam/nvidia/layernorm.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Beam kernels preserving FP32 normalization and affine arithmetic.""" + +import torch +from torch import nn + +from vllm.triton_utils import tl, triton + + +class RMSNorm(nn.Module): + """Parameters for normalization performed directly by the model kernels.""" + + def __init__(self, hidden_size: int, has_weight: bool = True) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) if has_weight else None + + def forward(self, *args, **kwargs): + raise RuntimeError("Call the model's fused normalization kernels directly") + + +@triton.jit +def _rms_norm_kernel( + X, + W, + Y, + N: tl.constexpr, + HEADS: tl.constexpr, + STRIDE_TOKEN: tl.constexpr, + STRIDE_HEAD: tl.constexpr, + EPS: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + idx = tl.arange(0, BLOCK) + token = row // HEADS + head = row % HEADS + x = tl.load( + X + token * STRIDE_TOKEN + head * STRIDE_HEAD + idx, idx < N, other=0 + ).to(tl.float32) + mean = tl.sum(x * x, 0) / N + y = x * tl.rsqrt(mean + EPS) + if HAS_WEIGHT: + w = tl.load(W + idx, idx < N, other=0).to(tl.float32) + y = y * w + tl.store(Y + row * N + idx, y, idx < N) + + +def beam_rms_norm( + x: torch.Tensor, weight: torch.Tensor | None, eps: float +) -> torch.Tensor: + """Normalize strided token/head rows, rounding only after the affine multiply.""" + assert x.ndim in (2, 3) and x.stride(-1) == 1 + n = x.shape[-1] + heads = x.shape[1] if x.ndim == 3 else 1 + out = torch.empty(x.shape, device=x.device, dtype=x.dtype) + if out.numel() == 0: + return out + _rms_norm_kernel[(x.numel() // n,)]( + x, + weight, + out, + n, + heads, + x.stride(0), + x.stride(-2) if x.ndim == 3 else n, + eps, + weight is not None, + triton.next_power_of_2(n), + num_warps=4 if n <= 1024 else 8, + enable_fp_fusion=False, + ) + return out + + +@triton.jit +def _boundary_math( + x, + r, + POST_W, + NEXT_W, + col, + dtype: tl.constexpr, + N: tl.constexpr, + POST_EPS: tl.constexpr, + NEXT_EPS: tl.constexpr, + X_SCALE, + R_SCALE, + POST_NORM: tl.constexpr, + POST_WEIGHT: tl.constexpr, + NEXT_NORM: tl.constexpr, + NEXT_WEIGHT: tl.constexpr, + SCALED: tl.constexpr, +): + if POST_NORM: + x = x * tl.rsqrt(tl.sum(x * x, 0) / N + POST_EPS) + if POST_WEIGHT: + w = tl.load(POST_W + col, col < N, other=0).to(tl.float32) + x = x * w + x = x.to(dtype).to(tl.float32) + if SCALED: + x = (x * X_SCALE).to(dtype).to(tl.float32) + r = (r * R_SCALE).to(dtype).to(tl.float32) + h = (x + r).to(dtype) + if NEXT_NORM: + h_fp32 = tl.where(col < N, h.to(tl.float32), 0) + y = h_fp32 * tl.rsqrt(tl.sum(h_fp32 * h_fp32, 0) / N + NEXT_EPS) + if NEXT_WEIGHT: + w = tl.load(NEXT_W + col, col < N, other=0).to(tl.float32) + y = y * w + else: + y = h + return h, y + + +@triton.jit +def _norm_residual_norm_kernel( + X, + R, + POST_W, + NEXT_W, + H, + Y, + N: tl.constexpr, + X_STRIDE: tl.constexpr, + R_STRIDE: tl.constexpr, + POST_EPS: tl.constexpr, + NEXT_EPS: tl.constexpr, + X_SCALE, + R_SCALE, + POST_NORM: tl.constexpr, + POST_WEIGHT: tl.constexpr, + NEXT_NORM: tl.constexpr, + NEXT_WEIGHT: tl.constexpr, + SCALED: tl.constexpr, + BLOCK: tl.constexpr, + PDL: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + col = tl.arange(0, BLOCK) + dtype = X.dtype.element_ty + if PDL: + tl.extra.cuda.gdc_wait() + x = tl.load(X + row * X_STRIDE + col, col < N, other=0).to(tl.float32) + r = tl.load(R + row * R_STRIDE + col, col < N, other=0).to(dtype).to(tl.float32) + h, y = _boundary_math( + x, + r, + POST_W, + NEXT_W, + col, + dtype, + N, + POST_EPS, + NEXT_EPS, + X_SCALE, + R_SCALE, + POST_NORM, + POST_WEIGHT, + NEXT_NORM, + NEXT_WEIGHT, + SCALED, + ) + tl.store(H + row * N + col, h, col < N) + if NEXT_NORM: + tl.store(Y + row * N + col, y, col < N) + + +def beam_norm_residual_norm( + x: torch.Tensor, + residual: torch.Tensor, + post_weight: torch.Tensor | None, + post_eps: float | None, + next_weight: torch.Tensor | None, + next_eps: float | None, + scales: tuple[float, float] | None, + *, + launch_pdl: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Fuse post-norm, scaled residual add and next pre-norm, keeping each cast.""" + assert x.ndim == residual.ndim == 2 and x.shape == residual.shape + assert x.stride(-1) == residual.stride(-1) == 1 + hidden = torch.empty(x.shape, device=x.device, dtype=x.dtype) + normalized = torch.empty_like(hidden) if next_eps is not None else None + n = x.shape[-1] + if hidden.numel(): + _norm_residual_norm_kernel[(x.shape[0],)]( + x, + residual, + post_weight, + next_weight, + hidden, + normalized, + n, + x.stride(0), + residual.stride(0), + post_eps or 0.0, + next_eps or 0.0, + *(scales if scales is not None else (1.0, 1.0)), + post_eps is not None, + post_weight is not None, + next_eps is not None, + next_weight is not None, + scales is not None, + triton.next_power_of_2(n), + launch_pdl, + num_warps=4 if launch_pdl or n <= 1024 else 8, + enable_fp_fusion=False, + launch_pdl=launch_pdl, + ) + return hidden, normalized diff --git a/vllm/models/beam/nvidia/model.py b/vllm/models/beam/nvidia/model.py new file mode 100644 index 000000000000..f1d6ba31a9bc --- /dev/null +++ b/vllm/models/beam/nvidia/model.py @@ -0,0 +1,550 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Beam model compatible with HuggingFace weights. + +Beam is a decoder-only sparse MoE causal LM. Its config-gated features +are NoPE layers with SSMax query scaling, hybrid sliding-window attention, +per-head Q/K RMSNorm, sandwich norm, depth-scaled residuals, and sigmoid +attention output gating. The MoE router uses grouped sigmoid ``noaux_tc`` +selection. +""" + +import math +from collections.abc import Iterable +from functools import partial +from typing import cast + +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.distributed import ( + get_pp_group, + get_tensor_model_parallel_world_size, + get_tp_group, +) +from vllm.model_executor.layers.fused_moe import MoERunner +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.models.interfaces import ( + EagleModelMixin, + MixtureOfExperts, + SupportsEagle3, + SupportsPP, + SupportsQuant, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + WeightsMapper, + extract_layer_index, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.beam import BeamConfig + +from .activation import beam_softcap_logits +from .allreduce import ( + BeamAllReduceWorkspace, + beam_allreduce_norm_residual_norm, + beam_finalize_allreduce_norm_residual_norm, +) +from .attention import BeamAttention, use_merged_qkv_gate +from .layernorm import RMSNorm, beam_norm_residual_norm, beam_rms_norm +from .moe import BeamMLP, BeamSparseMoeBlock + + +def _depth_scaled_coeffs(layer_id: int, is_attention: bool) -> tuple[float, float]: + """Depth-scaled residual coefficients ``(x_scale, r_scale)`` for one sublayer. + + ``num_adds`` counts residual adds up to and including this one, with a + 1-based layer index ``L = layer_id + 1``: ``2L`` for the attention add, + ``2L + 1`` for the MLP add. + """ + one_based = layer_id + 1 + num_adds = 2 * one_based if is_attention else 2 * one_based + 1 + denom = math.sqrt(num_adds) + return 1.0 / denom, math.sqrt(num_adds - 1) / denom + + +class BeamDecoderLayer(nn.Module): + """Sandwich-norm block carrying residual and next normalized input separately.""" + + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + allreduce_workspace: BeamAllReduceWorkspace | None = None, + ) -> None: + super().__init__() + config = cast(BeamConfig, vllm_config.model_config.hf_config) + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + self.layer_id = extract_layer_index(prefix) + self.rms_norm_eps = config.rms_norm_eps + fuse_attention_reduce = ( + allreduce_workspace is not None and not config.attention_bias + ) + + self.self_attn = BeamAttention( + config, + cache_config, + quant_config, + prefix=f"{prefix}.self_attn", + reduce_results=not fuse_attention_reduce, + ) + + if self.layer_id in config.mlp_only_layers: + self.mlp: nn.Module = BeamMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + reduce_results=allreduce_workspace is None, + ) + else: + self.mlp = BeamSparseMoeBlock( + vllm_config=vllm_config, + prefix=f"{prefix}.mlp", + reduce_results=allreduce_workspace is None, + ) + + fused_boundary = ( + partial(beam_allreduce_norm_residual_norm, workspace=allreduce_workspace) + if allreduce_workspace is not None + else beam_norm_residual_norm + ) + self.attention_boundary = ( + fused_boundary if fuse_attention_reduce else beam_norm_residual_norm + ) + fuse_mlp_reduce = allreduce_workspace is not None and ( + isinstance(self.mlp, BeamMLP) + or ( + not self.mlp.is_sequence_parallel + and self.mlp.experts.moe_config.skip_final_all_reduce + ) + ) + self.mlp_boundary = ( + fused_boundary if fuse_mlp_reduce else beam_norm_residual_norm + ) + self.finalize_boundary = None + if fuse_mlp_reduce and isinstance(self.mlp, BeamSparseMoeBlock): + assert allreduce_workspace is not None + moe_config = self.mlp.experts.moe_config + allow_ep = moe_config.ep_size == allreduce_workspace.world_size + max_tokens = allreduce_workspace.max_tokens + if allow_ep: + max_tokens = min(max_tokens, allreduce_workspace.ep_finalize_max_tokens) + self.mlp.defer_finalize(max_tokens, allow_ep=allow_ep) + self.finalize_boundary = partial( + beam_finalize_allreduce_norm_residual_norm, + workspace=allreduce_workspace, + routed_scale=self.mlp.experts.routed_scaling_factor, + oneshot_max_tokens=( + allreduce_workspace.ep_finalize_oneshot_max_tokens + if allow_ep + else 0 + ), + ) + + self.input_layernorm = RMSNorm(config.hidden_size) + self.pre_feedforward_layernorm = RMSNorm(config.hidden_size) + self.post_attention_layernorm = ( + RMSNorm(config.hidden_size) if config.enable_sandwich_norm else None + ) + self.post_feedforward_layernorm = ( + RMSNorm(config.hidden_size) if config.enable_sandwich_norm else None + ) + self.mlp_post_norm_fused = isinstance( + self.mlp, BeamSparseMoeBlock + ) and self.mlp.fuse_output_post_norm( + self.post_feedforward_layernorm, + self.rms_norm_eps, + workspace=allreduce_workspace, + ) + + self.attn_scales = ( + _depth_scaled_coeffs(self.layer_id, is_attention=True) + if config.self_attn_depth_scaled_residual + else None + ) + self.mlp_scales = ( + _depth_scaled_coeffs(self.layer_id, is_attention=False) + if config.mlp_depth_scaled_residual + else None + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + normalized_states: torch.Tensor | None = None, + next_norm_weight: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + x = ( + beam_rms_norm(hidden_states, self.input_layernorm.weight, self.rms_norm_eps) + if normalized_states is None + else normalized_states + ) + x = self.self_attn(positions, x) + hidden_states, x = self.attention_boundary( + x, + hidden_states, + self.post_attention_layernorm.weight + if self.post_attention_layernorm is not None + else None, + self.rms_norm_eps if self.post_attention_layernorm is not None else None, + self.pre_feedforward_layernorm.weight, + self.rms_norm_eps, + self.attn_scales, + ) + assert x is not None + if isinstance(self.mlp, BeamSparseMoeBlock) and self.mlp.output_boundary_fused: + return self.mlp.forward_boundary( + x, + hidden_states, + self.post_feedforward_layernorm.weight + if self.post_feedforward_layernorm is not None + else None, + self.rms_norm_eps + if self.post_feedforward_layernorm is not None + else None, + next_norm_weight, + self.rms_norm_eps if next_norm_weight is not None else None, + self.mlp_scales, + ) + x = self.mlp(x) + boundary = self.mlp_boundary + if isinstance(x, UnfinalizedMoEOutput): + assert self.finalize_boundary is not None + boundary = self.finalize_boundary + return boundary( + x, + hidden_states, + self.post_feedforward_layernorm.weight + if self.post_feedforward_layernorm is not None + and not self.mlp_post_norm_fused + else None, + self.rms_norm_eps + if self.post_feedforward_layernorm is not None + and not self.mlp_post_norm_fused + else None, + next_norm_weight, + self.rms_norm_eps if next_norm_weight is not None else None, + self.mlp_scales, + ) + + +class BeamModel(nn.Module, EagleModelMixin): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={ + # Renamed so it cannot be matched as a `gate_up_proj` shard by + # `packed_modules_mapping` consumers such as LoRA. + ".self_attn.gate_proj": ".self_attn.output_gate_proj", + }, + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".self_attn.q_proj": (".self_attn.qkv_proj", "q"), + ".self_attn.k_proj": (".self_attn.qkv_proj", "k"), + ".self_attn.v_proj": (".self_attn.qkv_proj", "v"), + # `.mlp.experts.N.gate_proj` does not contain `.mlp.gate_proj`, so + # expert weights fall through to `MoERunner.load_weights` (which is + # also what keeps them correct under EP and redundant experts). + ".mlp.gate_proj": (".mlp.gate_up_proj", 0), + ".mlp.up_proj": (".mlp.gate_up_proj", 1), + }, + ) + + @classmethod + def get_weights_mapper(cls, config: BeamConfig, tp_size: int) -> WeightsMapper: + mapper = cls.hf_to_vllm_mapper + if not use_merged_qkv_gate(config, tp_size): + return mapper + return WeightsMapper( + orig_to_new_substr=mapper.orig_to_new_substr, + orig_to_new_stacked={ + **mapper.orig_to_new_stacked, + **{ + f".self_attn.{name}": (".self_attn.qkv_gate_proj", shard) + for shard, name in enumerate( + ("q_proj", "k_proj", "v_proj", "output_gate_proj") + ) + }, + }, + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config = vllm_config.model_config.hf_config + self.hf_to_vllm_mapper = self.get_weights_mapper( + config, get_tensor_model_parallel_world_size() + ) + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + + # Beam drives SWA per layer off `layer_types`, passing an explicit + # `per_layer_sliding_window` (None on the full-attention layers). + # Without this, those layers would pick up `cache_config.sliding_window` + # as a fallback in `Attention.__init__` and silently get a window. + if cache_config is not None: + cache_config.sliding_window = None + + self.config = config + self.quant_config = quant_config + self.vocab_size = config.vocab_size + self.embed_tokens: VocabParallelEmbedding | PPMissingLayer + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + tp_group = get_tp_group() + self.allreduce_workspace = BeamAllReduceWorkspace.create( + tp_group.cpu_group, + config.hidden_size, + max( + vllm_config.scheduler_config.max_num_batched_tokens, + vllm_config.scheduler_config.max_num_seqs, + ), + vllm_config.model_config.dtype, + device_group=tp_group.device_group, + ) + if self.allreduce_workspace is not None: + self.allreduce_workspace.all_gather = partial(tp_group.all_gather, dim=0) + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: BeamDecoderLayer( + vllm_config=vllm_config, + prefix=prefix, + allreduce_workspace=self.allreduce_workspace, + ), + prefix=f"{prefix}.layers", + ) + self.norm = ( + RMSNorm(config.hidden_size) + if get_pp_group().is_last_rank + else PPMissingLayer() + ) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states"], config.hidden_size + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> 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 + else: + assert input_ids is not None + hidden_states = self.embed_input_ids(input_ids) + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + + aux_hidden_states = self._maybe_add_hidden_state( + [], self.start_layer, hidden_states, None + ) + normalized_states = None + for index in range(self.start_layer, self.end_layer): + next_norm = ( + self.layers[index + 1].input_layernorm + if index + 1 < self.end_layer + else self.norm + ) + hidden_states, normalized_states = self.layers[index]( + positions, + hidden_states, + normalized_states, + next_norm.weight if isinstance(next_norm, RMSNorm) else None, + ) + self._maybe_add_hidden_state( + aux_hidden_states, index + 1, hidden_states, None + ) + + if not get_pp_group().is_last_rank: + return IntermediateTensors({"hidden_states": hidden_states}) + + hidden_states = ( + beam_rms_norm(hidden_states, self.norm.weight, self.config.rms_norm_eps) + if normalized_states is None + else normalized_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]: + loader = AutoWeightsLoader( + self, + ignore_unexpected_suffixes=[ + ".weight_scale", + "_weight_scale", + ".input_scale", + "_input_scale", + ], + ) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + +class BeamForCausalLM( + nn.Module, + SupportsPP, + MixtureOfExperts, + SupportsQuant, + SupportsEagle3, +): + """Beam causal LM. LoRA support is deferred to a follow-up change.""" + + hf_to_vllm_mapper = BeamModel.hf_to_vllm_mapper + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "qkv_gate_proj": ["q_proj", "k_proj", "v_proj", "output_gate_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + fall_back_to_pt_during_load = False + + def process_weights_after_loading(self) -> None: + for module in self.model.modules(): + if isinstance(module, BeamAttention): + module.warmup_projection() + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + parallel_config = vllm_config.parallel_config + if parallel_config.enable_eplb: + # Expert parallelism is validated; expert rebalancing is not. + raise NotImplementedError( + "Beam has not been validated with EPLB; run without --enable-eplb." + ) + if parallel_config.use_sequence_parallel_moe: + raise NotImplementedError( + "Beam has not been validated with sequence-parallel MoE " + "(data parallelism with expert parallelism)." + ) + self.config = config + self.quant_config = quant_config + + self.model = BeamModel( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + if config.tie_word_embeddings: + raise NotImplementedError("Beam requires tie_word_embeddings=False.") + # Untied LM head, in the model dtype. An fp32 head is available the + # standard way, via `head_dtype` on the config (see + # `LogitsProcessor._apply_head`), rather than by pinning the parameter + # dtype here. + self.lm_head: ParallelLMHead | PPMissingLayer + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + # Apply the softcap in fp32 while keeping the head in the model dtype. + self.final_logit_softcapping = config.final_logit_softcapping + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + # MoE metadata for the MixtureOfExperts interface. + self.moe_layers: list[MoERunner] = [] + example_layer: BeamSparseMoeBlock | None = None + for layer in self.model.layers: + if isinstance(layer, PPMissingLayer): + continue + assert isinstance(layer, BeamDecoderLayer) + if isinstance(layer.mlp, BeamSparseMoeBlock): + example_layer = layer.mlp + self.moe_layers.append(layer.mlp.experts) + if example_layer is None: + # A pipeline stage may contain only the initial dense layers. + self.num_moe_layers = 0 + self.num_expert_groups = 0 + self.num_shared_experts = 0 + self.num_logical_experts = 0 + self.num_physical_experts = 0 + self.num_local_physical_experts = 0 + self.num_routed_experts = 0 + self.num_redundant_experts = 0 + return + + self.num_moe_layers = len(self.moe_layers) + self.num_expert_groups = config.n_group or 1 + self.num_shared_experts = 0 + moe_config = example_layer.experts.moe_config + self.num_logical_experts = moe_config.num_logical_experts + self.num_physical_experts = moe_config.num_experts + self.num_local_physical_experts = moe_config.num_local_experts + self.num_routed_experts = config.num_experts + self.num_redundant_experts = ( + self.num_physical_experts - self.num_logical_experts + ) + + def update_physical_experts_metadata( + self, + num_physical_experts: int, + num_local_physical_experts: int, + ) -> None: + assert self.num_local_physical_experts == num_local_physical_experts + self.num_physical_experts = num_physical_experts + self.num_local_physical_experts = num_local_physical_experts + self.num_redundant_experts = num_physical_experts - self.num_logical_experts + for layer in self.moe_layers: + layer.update_expert_map() + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + logits = self.logits_processor(self.lm_head, hidden_states) + if logits is not None and self.final_logit_softcapping is not None: + logits = beam_softcap_logits(logits, self.final_logit_softcapping) + return logits + + def _strip_lm_head_padding( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> Iterable[tuple[str, torch.Tensor]]: + # Strip padding from the LM head's input dimension. + hidden_size = self.config.hidden_size + for name, weight in weights: + if name.endswith("lm_head.weight") and weight.shape[-1] > hidden_size: + weight = weight[..., :hidden_size].contiguous() + yield name, weight + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(self._strip_lm_head_padding(weights)) diff --git a/vllm/models/beam/nvidia/moe.py b/vllm/models/beam/nvidia/moe.py new file mode 100644 index 000000000000..011aadcea6a9 --- /dev/null +++ b/vllm/models/beam/nvidia/moe.py @@ -0,0 +1,337 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import TYPE_CHECKING + +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.distributed import tensor_model_parallel_all_gather +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.fused_moe import FusedMoEFactory, GateLinear +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput +from vllm.model_executor.layers.fused_moe.runner.moe_runner import ( + _layer_name_type, + _resolve_layer_name, + get_layer_from_name, +) +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.models.utils import sequence_parallel_chunk +from vllm.utils.torch_utils import ( + STR_DTYPE_TO_TORCH_DTYPE, + direct_register_custom_op, +) + +from .activation import beam_silu_mul +from .layernorm import RMSNorm + +if TYPE_CHECKING: + from .allreduce import BeamAllReduceWorkspace + + +def _beam_moe_forward_boundary( + hidden_states: torch.Tensor, + residual: torch.Tensor, + post_weight: torch.Tensor | None, + post_eps: float | None, + next_weight: torch.Tensor | None, + next_eps: float | None, + x_scale: float | None, + r_scale: float | None, + layer_name: _layer_name_type, +) -> tuple[torch.Tensor, torch.Tensor]: + experts = get_layer_from_name(_resolve_layer_name(layer_name)) + layer = get_forward_context().no_compile_layers[f"{experts.layer_name}.finalize"] + assert isinstance(layer, BeamSparseMoeBlock) + hidden, normalized = layer._forward_boundary( + hidden_states, + residual, + post_weight, + post_eps, + next_weight, + next_eps, + (x_scale, r_scale) if x_scale is not None and r_scale is not None else None, + ) + return hidden, normalized if normalized is not None else residual.new_empty(0) + + +def _beam_moe_forward_boundary_fake( + hidden_states: torch.Tensor, + residual: torch.Tensor, + post_weight: torch.Tensor | None, + post_eps: float | None, + next_weight: torch.Tensor | None, + next_eps: float | None, + x_scale: float | None, + r_scale: float | None, + layer_name: _layer_name_type, +) -> tuple[torch.Tensor, torch.Tensor]: + return ( + torch.empty_like(residual), + torch.empty_like(residual) if next_eps is not None else residual.new_empty(0), + ) + + +direct_register_custom_op( + op_name="beam_moe_forward_boundary", + op_func=_beam_moe_forward_boundary, + mutates_args=["hidden_states"], + fake_impl=_beam_moe_forward_boundary_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) + + +class BeamMLP(nn.Module): + """Dense SwiGLU MLP (used by ``mlp_only_layers``).""" + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + reduce_results: bool = True, + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + reduce_results=reduce_results, + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported." + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + # Beam rounds once after the fp32 SiLU and gate/up multiplication. + out = beam_silu_mul(gate_up) + out, _ = self.down_proj(out) + return out + + +class BeamSparseMoeBlock(nn.Module): + """Grouped sigmoid ``noaux_tc`` MoE block.""" + + def __init__( + self, vllm_config: VllmConfig, prefix: str = "", reduce_results: bool = True + ) -> None: + super().__init__() + config = vllm_config.model_config.hf_config + parallel_config = vllm_config.parallel_config + quant_config = vllm_config.quant_config + self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe + self.gather_workspace: BeamAllReduceWorkspace | None = None + self.output_boundary_fused = False + self._compilation_config = vllm_config.compilation_config + + router_dtype = STR_DTYPE_TO_TORCH_DTYPE[config.moe_router_dtype or "float32"] + self.gate = GateLinear( + config.hidden_size, + config.num_experts, + bias=False, + params_dtype=router_dtype, + out_dtype=router_dtype, + prefix=f"{prefix}.gate", + ) + e_score_correction_bias = None + if config.topk_method == "noaux_tc": + e_score_correction_bias = nn.Parameter( + torch.zeros(config.num_experts, dtype=torch.float32) + ) + self.gate.register_parameter("e_score_correction_bias", e_score_correction_bias) + + self.experts = FusedMoEFactory( + reduce_results=reduce_results or self.is_sequence_parallel, + is_sequence_parallel=self.is_sequence_parallel, + gate=self.gate, + num_experts=config.num_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + renormalize=config.norm_topk_prob, + quant_config=quant_config, + use_grouped_topk=( + config.n_group is not None and config.topk_group is not None + ), + num_expert_group=config.n_group, + topk_group=config.topk_group, + scoring_func=config.scoring_func, + e_score_correction_bias=e_score_correction_bias, + routed_scaling_factor=( + config.routed_scaling_factor + if config.routed_scaling_factor is not None + else 1.0 + ), + apply_routed_scale_to_output=True, + router_logits_dtype=self.gate.out_dtype, + enable_eplb=parallel_config.enable_eplb, + num_redundant_experts=parallel_config.eplb_config.num_redundant_experts, + prefix=f"{prefix}.experts", + ) + self._boundary_layer_name = f"{prefix}.experts.finalize" + + def defer_finalize(self, max_tokens: int, *, allow_ep: bool = False) -> None: + from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( + TrtLlmNvFp4ExpertsModular, + TrtLlmNvFp4ExpertsMonolithic, + ) + + experts_cls = getattr( + self.experts.routed_experts.quant_method, "experts_cls", None + ) + if experts_cls in (TrtLlmNvFp4ExpertsModular, TrtLlmNvFp4ExpertsMonolithic): + self.experts.moe_config.defer_moe_finalize(max_tokens, allow_ep=allow_ep) + + def fuse_output_post_norm( + self, + post_norm: RMSNorm | None, + eps: float, + *, + workspace: "BeamAllReduceWorkspace | None" = None, + ) -> bool: + from vllm.config.kernel import FLASHINFER_MOE_EP_CUTEDSL + from vllm.model_executor.layers.fused_moe.experts.flashinfer_moe_ep import ( + FlashInferMoeEpExperts, + ) + from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep import ( + supports_flashinfer_moe_ep_deferred_finalize, + ) + + moe = self.experts.moe_config + experts = self.experts.routed_experts + if ( + (not self.is_sequence_parallel and moe.dp_size == 1) + or moe.moe_backend != FLASHINFER_MOE_EP_CUTEDSL + or getattr(experts.quant_method, "experts_cls", None) + is not FlashInferMoeEpExperts + or moe.in_dtype != torch.bfloat16 + or moe.hidden_dim != moe.hidden_dim_unpadded + or moe.hidden_dim % 256 + or not 256 <= moe.hidden_dim <= 8192 + or not 1 <= moe.experts_per_token <= 32 + or experts.apply_router_weight_on_input + or self.experts.shared_experts is not None + or self.experts.routed_input_transform is not None + or self.experts.routed_output_transform is not None + or moe.pcp_size != 1 + or moe.tp_size != 1 + or moe.ep_size <= 1 + or not supports_flashinfer_moe_ep_deferred_finalize() + ): + return False + fuse_gather = ( + workspace is not None + and moe.sp_size == workspace.world_size == moe.ep_size + and moe.dp_size == 1 + ) + if not self.is_sequence_parallel or fuse_gather: + self.gather_workspace = workspace if fuse_gather else None + self.output_boundary_fused = True + moe.defer_moe_finalize(allow_ep=True) + self._compilation_config.static_forward_context[ + self._boundary_layer_name + ] = self + return True + return False + + def forward_boundary( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + post_weight: torch.Tensor | None, + post_eps: float | None, + next_weight: torch.Tensor | None, + next_eps: float | None, + scales: tuple[float, float] | None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + hidden, normalized = torch.ops.vllm.beam_moe_forward_boundary( + hidden_states, + residual, + post_weight, + post_eps, + next_weight, + next_eps, + *(scales if scales is not None else (None, None)), + self.experts._encode_layer_name(), + ) + return hidden, normalized if next_eps is not None else None + + def _forward_boundary( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + post_weight: torch.Tensor | None, + post_eps: float | None, + next_weight: torch.Tensor | None, + next_eps: float | None, + scales: tuple[float, float] | None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + from .moe_finalize import ( + beam_moe_finalize_gather_norm_residual_norm, + beam_moe_finalize_norm_residual_norm, + ) + + local_states = ( + sequence_parallel_chunk(hidden_states) + if self.is_sequence_parallel + else hidden_states + ) + output = self.experts._forward_impl(local_states, local_states, None) + assert isinstance(output, UnfinalizedMoEOutput) + # Consume the borrowed expert workspace inside the same opaque op. + if self.gather_workspace is None: + return beam_moe_finalize_norm_residual_norm( + output, + residual, + post_weight, + post_eps, + next_weight, + next_eps, + scales, + routed_scale=self.experts.routed_scaling_factor, + ) + return beam_moe_finalize_gather_norm_residual_norm( + output, + residual, + post_weight, + post_eps, + next_weight, + next_eps, + scales, + self.gather_workspace, + routed_scale=self.experts.routed_scaling_factor, + ) + + def forward( + self, hidden_states: torch.Tensor + ) -> torch.Tensor | UnfinalizedMoEOutput: + num_tokens = hidden_states.shape[0] + if self.is_sequence_parallel: + hidden_states = sequence_parallel_chunk(hidden_states) + elif self.experts.moe_config.should_defer_moe_finalize(num_tokens): + output = self.experts._forward_impl(hidden_states, hidden_states, None) + assert isinstance(output, UnfinalizedMoEOutput) + return output + + output = self.experts(hidden_states=hidden_states, router_logits=hidden_states) + if self.is_sequence_parallel: + output = tensor_model_parallel_all_gather(output, 0)[:num_tokens] + return output diff --git a/vllm/models/beam/nvidia/moe_finalize.cu b/vllm/models/beam/nvidia/moe_finalize.cu new file mode 100644 index 000000000000..89a03c36581f --- /dev/null +++ b/vllm/models/beam/nvidia/moe_finalize.cu @@ -0,0 +1,260 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include +#include + +#include "tvm_ffi_utils.h" + +namespace beam { +using T = __nv_bfloat16; +constexpr int Packs = BEAM_MOE_PACKS_PER_THREAD; +union Vec { + float4 packed; + T elements[8]; +}; + +__device__ float blockSum(float value) { + __shared__ float partial[32]; + int lane = threadIdx.x % 32; + int warp = threadIdx.x / 32; +#pragma unroll + for (int offset = 16; offset > 0; offset /= 2) + value = __fadd_rn(value, __shfl_xor_sync(0xffffffff, value, offset)); + if (lane == 0) partial[warp] = value; + __syncthreads(); + if (warp == 0) { + value = lane < blockDim.x / 32 ? partial[lane] : 0.f; +#pragma unroll + for (int offset = 16; offset > 0; offset /= 2) + value = __fadd_rn(value, __shfl_xor_sync(0xffffffff, value, offset)); + if (lane == 0) partial[0] = value; + } + __syncthreads(); + return partial[0]; +} + +__device__ Vec finalizePack(T const* experts, float const* weights, int token, + int column, int hidden, int topK, + float routedScale) { + int64_t inputOffset = int64_t(token) * topK * hidden + column; + Vec value; + value.packed = *reinterpret_cast(experts + inputOffset); + float weight = weights[int64_t(token) * topK]; + float sum[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) + sum[i] = __fmul_rn(__bfloat162float(value.elements[i]), weight); + for (int k = 1; k < topK; ++k) { + value.packed = + *reinterpret_cast(experts + inputOffset + k * hidden); + weight = weights[int64_t(token) * topK + k]; +#pragma unroll + for (int i = 0; i < 8; ++i) + sum[i] = __fmaf_rn(__bfloat162float(value.elements[i]), weight, sum[i]); + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + T reduced = __float2bfloat16_rn(sum[i]); + value.elements[i] = + __float2bfloat16_rn(__fmul_rn(__bfloat162float(reduced), routedScale)); + } + return value; +} + +__device__ float invRms(Vec const (&values)[Packs], int hidden, float eps) { + float squareSum = 0.f; +#pragma unroll + for (int pack = 0; pack < Packs; ++pack) +#pragma unroll + for (int i = 0; i < 8; ++i) { + float x = __bfloat162float(values[pack].elements[i]); + squareSum = __fadd_rn(squareSum, __fmul_rn(x, x)); + } + return rsqrtf(blockSum(squareSum) / hidden + eps); +} + +__device__ Vec normalize(Vec value, T const* weight, int column, float rstd) { + Vec affine; + if (weight) affine.packed = *reinterpret_cast(weight + column); +#pragma unroll + for (int i = 0; i < 8; ++i) { + float x = __fmul_rn(__bfloat162float(value.elements[i]), rstd); + if (weight) x = __fmul_rn(x, __bfloat162float(affine.elements[i])); + value.elements[i] = __float2bfloat16_rn(x); + } + return value; +} + +template +__global__ void finalize(T const* experts, float const* weights, T* output, + T const* postWeight, int hidden, int topK, + float routedScale, float eps) { + int token = blockIdx.x; + Vec values[Packs]{}; +#pragma unroll + for (int pack = 0; pack < Packs; ++pack) { + int column = (threadIdx.x + pack * blockDim.x) * 8; + if (column < hidden) + values[pack] = finalizePack(experts, weights, token, column, hidden, topK, + routedScale); + } + float rstd = 1.f; + if constexpr (PostNorm) rstd = invRms(values, hidden, eps); +#pragma unroll + for (int pack = 0; pack < Packs; ++pack) { + int column = (threadIdx.x + pack * blockDim.x) * 8; + if (column >= hidden) continue; + Vec value = values[pack]; + if constexpr (PostNorm) value = normalize(value, postWeight, column, rstd); + *reinterpret_cast(output + int64_t(token) * hidden + column) = + value.packed; + } +} + +struct BoundaryParams { + T const* experts; + float const* weights; + T const* residual; + T const* postWeight; + T const* nextWeight; + T* output; + T* normalized; + int hidden; + int topK; + float routedScale; + float postEps; + float nextEps; + float xScale; + float rScale; + bool nextNorm; + bool scaled; +}; + +template +__global__ void __launch_bounds__(1024, Packs == 1 ? 2 : 1) + finalizeBoundary(BoundaryParams p) { + int token = blockIdx.x; + Vec values[Packs]{}; +#pragma unroll + for (int pack = 0; pack < Packs; ++pack) { + int column = (threadIdx.x + pack * blockDim.x) * 8; + if (column < p.hidden) + values[pack] = finalizePack(p.experts, p.weights, token, column, p.hidden, + p.topK, p.routedScale); + } + float rstd = 1.f; + if constexpr (PostNorm) rstd = invRms(values, p.hidden, p.postEps); +#pragma unroll + for (int pack = 0; pack < Packs; ++pack) { + int column = (threadIdx.x + pack * blockDim.x) * 8; + if (column >= p.hidden) continue; + int64_t offset = int64_t(token) * p.hidden + column; + Vec value = values[pack]; + if constexpr (PostNorm) + value = normalize(value, p.postWeight, column, rstd); + Vec residual; + residual.packed = *reinterpret_cast(p.residual + offset); +#pragma unroll + for (int i = 0; i < 8; ++i) { + T x = value.elements[i], r = residual.elements[i]; + if (p.scaled) { + x = __float2bfloat16_rn(__fmul_rn(__bfloat162float(x), p.xScale)); + r = __float2bfloat16_rn(__fmul_rn(__bfloat162float(r), p.rScale)); + } + value.elements[i] = __float2bfloat16_rn( + __fadd_rn(__bfloat162float(x), __bfloat162float(r))); + } + values[pack] = value; + *reinterpret_cast(p.output + offset) = value.packed; + } + if (p.nextNorm) { + // The second reduction reuses blockSum's shared storage. + if constexpr (PostNorm) __syncthreads(); + rstd = invRms(values, p.hidden, p.nextEps); +#pragma unroll + for (int pack = 0; pack < Packs; ++pack) { + int column = (threadIdx.x + pack * blockDim.x) * 8; + if (column >= p.hidden) continue; + Vec value = normalize(values[pack], p.nextWeight, column, rstd); + *reinterpret_cast(p.normalized + int64_t(token) * p.hidden + + column) = value.packed; + } + } +} + +void launch(TensorView experts, TensorView weights, TensorView output, + tvm::ffi::Optional postWeight, double routedScale, + double eps, bool postNorm) { + ffi::CUDADeviceGuard guard(output.device().device_id); + int hidden = output.size(1); + int topK = experts.size(1); + TVM_FFI_ICHECK_EQ(encode_dlpack_dtype(experts.dtype()), bfloat16_code); + TVM_FFI_ICHECK_EQ(encode_dlpack_dtype(weights.dtype()), float32_code); + TVM_FFI_ICHECK_EQ(encode_dlpack_dtype(output.dtype()), bfloat16_code); + TVM_FFI_ICHECK(hidden >= 256 && hidden <= 8192 && hidden % 256 == 0); + TVM_FFI_ICHECK(topK >= 1 && topK <= 32); + if (output.size(0) == 0) return; + auto kernel = postNorm ? finalize : finalize; + int threads = ((hidden + 8 * Packs * 32 - 1) / (8 * Packs * 32)) * 32; + kernel<<>>( + static_cast(experts.data_ptr()), + static_cast(weights.data_ptr()), + static_cast(output.data_ptr()), + postWeight.has_value() + ? static_cast(postWeight.value().data_ptr()) + : nullptr, + hidden, topK, routedScale, eps); + auto status = cudaGetLastError(); + TVM_FFI_ICHECK_EQ(status, cudaSuccess) << cudaGetErrorString(status); +} + +void launchBoundary(TensorView experts, TensorView weights, TensorView residual, + tvm::ffi::Optional postWeight, + tvm::ffi::Optional nextWeight, + TensorView output, + tvm::ffi::Optional normalized, + double routedScale, double postEps, double nextEps, + double xScale, double rScale, bool postNorm, bool nextNorm, + bool scaled) { + ffi::CUDADeviceGuard guard(output.device().device_id); + int hidden = output.size(1); + int topK = experts.size(1); + TVM_FFI_ICHECK_EQ(encode_dlpack_dtype(experts.dtype()), bfloat16_code); + TVM_FFI_ICHECK_EQ(encode_dlpack_dtype(weights.dtype()), float32_code); + TVM_FFI_ICHECK_EQ(encode_dlpack_dtype(residual.dtype()), bfloat16_code); + TVM_FFI_ICHECK_EQ(encode_dlpack_dtype(output.dtype()), bfloat16_code); + TVM_FFI_ICHECK(hidden >= 256 && hidden <= 8192 && hidden % 256 == 0); + TVM_FFI_ICHECK(topK >= 1 && topK <= 32); + if (output.size(0) == 0) return; + auto optionalData = [](tvm::ffi::Optional tensor) -> T* { + return tensor.has_value() ? static_cast(tensor.value().data_ptr()) + : nullptr; + }; + BoundaryParams p{static_cast(experts.data_ptr()), + static_cast(weights.data_ptr()), + static_cast(residual.data_ptr()), + optionalData(postWeight), + optionalData(nextWeight), + static_cast(output.data_ptr()), + optionalData(normalized), + hidden, + topK, + static_cast(routedScale), + static_cast(postEps), + static_cast(nextEps), + static_cast(xScale), + static_cast(rScale), + nextNorm, + scaled}; + int threads = ((hidden + 8 * Packs * 32 - 1) / (8 * Packs * 32)) * 32; + auto kernel = postNorm ? finalizeBoundary : finalizeBoundary; + kernel<<>>(p); + auto status = cudaGetLastError(); + TVM_FFI_ICHECK_EQ(status, cudaSuccess) << cudaGetErrorString(status); +} +} // namespace beam + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(beam_moe_finalize, beam::launch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(beam_moe_finalize_norm_residual_norm, + beam::launchBoundary); diff --git a/vllm/models/beam/nvidia/moe_finalize.py b/vllm/models/beam/nvidia/moe_finalize.py new file mode 100644 index 000000000000..271d0ec69424 --- /dev/null +++ b/vllm/models/beam/nvidia/moe_finalize.py @@ -0,0 +1,271 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Finalize live MegaMoE tokens with Beam's post normalization.""" + +import functools +import hashlib +from pathlib import Path + +import torch + +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput + +from .allreduce import BeamAllReduceWorkspace + + +@functools.cache +def _module(packs_per_thread: int = 1): + from flashinfer.jit import env as jit_env + from flashinfer.jit.core import current_compilation_context, gen_jit_spec + + source = Path(__file__).with_suffix(".cu") + digest = hashlib.sha256(source.read_bytes()).hexdigest()[:16] + flags = current_compilation_context.get_nvcc_flags_list( + supported_major_versions=[9, 10] + ) + spec = gen_jit_spec( + f"beam_moe_finalize_{digest}_{packs_per_thread}", + [source], + extra_cuda_cflags=[ + *flags, + "--fmad=false", + f"-DBEAM_MOE_PACKS_PER_THREAD={packs_per_thread}", + ], + extra_include_paths=[jit_env.FLASHINFER_CSRC_DIR], + ) + spec.extra_cuda_cflags.remove("-use_fast_math") + return spec.build_and_load() + + +def beam_moe_finalize( + expert_outputs: torch.Tensor, + expert_weights: torch.Tensor, + output: torch.Tensor, + routed_scale: float = 1.0, + post_weight: torch.Tensor | None = None, + post_eps: float | None = None, +) -> torch.Tensor: + """Reduce only output rows, preserving BF16 casts before scale and RMSNorm. + + Args: + expert_outputs: Contiguous BF16 expert results [capacity, top_k, hidden]. + expert_weights: Contiguous FP32 routing weights [capacity, top_k]. + output: Contiguous BF16 destination [live_tokens, hidden]. + routed_scale: Multiplier applied after rounding the expert reduction. + post_weight: Optional BF16 RMSNorm affine weight [hidden]. + post_eps: RMSNorm epsilon, or None to omit normalization. + + """ + assert expert_outputs.ndim == 3 and expert_weights.ndim == output.ndim == 2 + capacity, top_k, hidden = expert_outputs.shape + assert 1 <= top_k <= 32 + assert 256 <= hidden <= 8192 and hidden % 256 == 0 + assert output.shape[1] == hidden and output.shape[0] <= capacity + assert expert_weights.shape[1] == top_k + assert expert_weights.shape[0] >= output.shape[0] + assert expert_outputs.dtype == output.dtype == torch.bfloat16 + assert expert_weights.dtype == torch.float32 + assert output.is_cuda + for tensor in (expert_outputs, expert_weights, output): + assert tensor.is_contiguous() and tensor.device == output.device + if post_weight is not None: + assert post_eps is not None + assert post_weight.shape == (hidden,) and post_weight.is_contiguous() + assert post_weight.dtype == torch.bfloat16 + assert post_weight.device == output.device + if output.shape[0]: + _module().beam_moe_finalize( + expert_outputs, + expert_weights, + output, + post_weight, + routed_scale, + post_eps or 0.0, + post_eps is not None, + ) + return output + + +def beam_moe_finalize_gather_norm_residual_norm( + x: UnfinalizedMoEOutput, + residual: torch.Tensor, + post_weight: torch.Tensor | None, + post_eps: float | None, + next_weight: torch.Tensor | None, + next_eps: float | None, + scales: tuple[float, float] | None, + workspace: BeamAllReduceWorkspace, + routed_scale: float = 1.0, + *, + strategy: str | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Finalize MegaMoE output and restore the replicated token boundary. + + FlashInfer's ``forward_unfinalized`` guarantees token-major expert rows and + an identity route map, so a view recovers its [tokens, top_k, hidden] layout. + This consumer requires that producer contract. + + Each rank owns ceil(global_tokens / TP) consecutive tokens. Capacity rows + and the final rank's token padding do not contribute to the returned tensors. + ``strategy`` overrides the token-count heuristic for validation and benchmarking. + """ + assert residual.ndim == 2 and residual.is_cuda + tokens, hidden_size = residual.shape + local_tokens = (tokens + workspace.world_size - 1) // workspace.world_size + assert tokens <= workspace.max_tokens and hidden_size == workspace.hidden_size + expert_weights = x.expert_weights + assert expert_weights.ndim == 2 + assert x.gemm2_permuted.ndim == 2 and x.gemm2_permuted.is_contiguous() + assert x.gemm2_permuted.shape[1] == hidden_size + expert_outputs = x.gemm2_permuted.view(*expert_weights.shape, hidden_size) + assert expert_weights.shape[0] >= local_tokens + assert 1 <= expert_weights.shape[1] <= 32 + assert expert_outputs.dtype == residual.dtype == torch.bfloat16 + assert expert_weights.dtype == torch.float32 + for tensor in (expert_outputs, expert_weights, residual): + assert tensor.is_contiguous() and tensor.device == residual.device + for weight, eps in ((post_weight, post_eps), (next_weight, next_eps)): + if weight is not None: + assert eps is not None + assert weight.shape == (hidden_size,) and weight.is_contiguous() + assert weight.dtype == torch.bfloat16 and weight.device == residual.device + if strategy is None: + if tokens <= workspace.moe_gather_oneshot_max_tokens: + strategy = "oneshot" + elif ( + workspace.moe_gather_separate_norm_min_tokens is not None + and tokens >= workspace.moe_gather_separate_norm_min_tokens + ): + strategy = "unfused" + else: + strategy = "twoshot" + if strategy not in ("unfused", "oneshot", "twoshot"): + raise ValueError(f"Unknown MegaMoE gather strategy: {strategy}") + if not tokens: + hidden = torch.empty_like(residual) + return hidden, torch.empty_like(residual) if next_eps is not None else None + if strategy == "unfused": + from .layernorm import beam_norm_residual_norm + + local = torch.empty( + (local_tokens, hidden_size), dtype=residual.dtype, device=residual.device + ) + beam_moe_finalize( + expert_outputs, expert_weights, local, routed_scale, post_weight, post_eps + ) + if workspace.all_gather is not None: + gathered = workspace.all_gather(local) + else: + gathered = torch.empty( + (local_tokens * workspace.world_size, hidden_size), + dtype=residual.dtype, + device=residual.device, + ) + torch.distributed.all_gather_into_tensor( + gathered, local, group=workspace.device_group + ) + return beam_norm_residual_norm( + gathered[:tokens], residual, None, None, next_weight, next_eps, scales + ) + hidden = torch.empty_like(residual) + normalized = torch.empty_like(residual) if next_eps is not None else None + w = workspace.workspace + workspace.module.beam_moe_gather( + expert_outputs, + expert_weights, + residual, + post_weight, + next_weight, + hidden, + normalized, + w.uc_ptrs_dev, + w.buffer_flags, + w.mc_ptr, + w.uc_ptr_local, + workspace.world_size, + workspace.rank, + post_eps or 0.0, + next_eps or 0.0, + *(scales if scales is not None else (1.0, 1.0)), + post_eps is not None, + next_eps is not None, + scales is not None, + strategy == "oneshot", + routed_scale, + ) + return hidden, normalized + + +def beam_moe_finalize_norm_residual_norm( + x: UnfinalizedMoEOutput, + residual: torch.Tensor, + post_weight: torch.Tensor | None, + post_eps: float | None, + next_weight: torch.Tensor | None, + next_eps: float | None, + scales: tuple[float, float] | None, + routed_scale: float = 1.0, + *, + strategy: str | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Finalize local MegaMoE tokens through the next normalization boundary. + + Requires FlashInfer's token-major expert rows and identity route map. + Only the live rows in ``residual`` are read from the workspace capacity. + ``strategy`` can force the unfused reference for validation and benchmarking. + """ + assert residual.ndim == 2 and residual.is_cuda + tokens, hidden_size = residual.shape + assert 256 <= hidden_size <= 8192 and hidden_size % 256 == 0 + expert_weights = x.expert_weights + assert expert_weights.ndim == 2 and expert_weights.shape[0] >= tokens + assert 1 <= expert_weights.shape[1] <= 32 + assert x.gemm2_permuted.ndim == 2 and x.gemm2_permuted.is_contiguous() + assert x.gemm2_permuted.shape[1] == hidden_size + expert_outputs = x.gemm2_permuted.view(*expert_weights.shape, hidden_size) + assert expert_outputs.dtype == residual.dtype == torch.bfloat16 + assert expert_weights.dtype == torch.float32 + for tensor in (expert_outputs, expert_weights, residual): + assert tensor.is_contiguous() and tensor.device == residual.device + for weight, eps in ((post_weight, post_eps), (next_weight, next_eps)): + if weight is not None: + assert eps is not None + assert weight.shape == (hidden_size,) and weight.is_contiguous() + assert weight.dtype == torch.bfloat16 and weight.device == residual.device + if strategy not in (None, "fused", "unfused"): + raise ValueError(f"Unknown local MegaMoE finalize strategy: {strategy}") + if strategy == "unfused": + from .layernorm import beam_norm_residual_norm + + finalized = torch.empty_like(residual) + beam_moe_finalize( + expert_outputs, + expert_weights, + finalized, + routed_scale, + post_weight, + post_eps, + ) + return beam_norm_residual_norm( + finalized, residual, None, None, next_weight, next_eps, scales + ) + hidden = torch.empty_like(residual) + normalized = torch.empty_like(residual) if next_eps is not None else None + if tokens: + _module().beam_moe_finalize_norm_residual_norm( + expert_outputs, + expert_weights, + residual, + post_weight, + next_weight, + hidden, + normalized, + routed_scale, + post_eps or 0.0, + next_eps or 0.0, + *(scales if scales is not None else (1.0, 1.0)), + post_eps is not None, + next_eps is not None, + scales is not None, + ) + return hidden, normalized diff --git a/vllm/models/beam/nvidia/projection.py b/vllm/models/beam/nvidia/projection.py new file mode 100644 index 000000000000..ce28ccda67dd --- /dev/null +++ b/vllm/models/beam/nvidia/projection.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Packed attention projections using shared low-latency GEMMs.""" + +import functools + +import torch +import torch.nn.functional as F + +from vllm.platforms import current_platform + + +@functools.cache +def _tinygemm(): + from flashinfer.jit import gen_tinygemm2_module + + return gen_tinygemm2_module().build_and_load() + + +def use_low_latency_gemm(weight: torch.Tensor) -> bool: + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import is_available + + return ( + weight.dtype == torch.bfloat16 + and weight.device.type == "cuda" + and current_platform.is_device_capability_family(100) + and weight.shape in ((1792, 5120), (3584, 5120)) + and is_available() + ) + + +def warmup_packed_projection(weight: torch.Tensor) -> None: + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( + ll_bf16_gemm_c1_pdl_kernel, + ll_bf16_gemm_kernel, + ) + + shapes = [(weight.shape[1], weight.shape[0])] + ll_bf16_gemm_c1_pdl_kernel.warmup( + shapes=shapes, m_values=[1], output_dtype=torch.bfloat16 + ) + ll_bf16_gemm_kernel.warmup( + shapes=shapes, m_values=[2, 16], output_dtype=torch.bfloat16 + ) + _tinygemm() + + +def packed_projection( + x: torch.Tensor, weight: torch.Tensor, low_latency: bool +) -> torch.Tensor: + rows = x.shape[0] + if not low_latency or not 1 <= rows <= 16 or not x.is_contiguous(): + return F.linear(x, weight) + tiny_max = 8 if weight.shape[0] == 1792 else 4 + if 3 <= rows <= tiny_max: + output = x.new_empty((rows, weight.shape[0])) + _tinygemm().tinygemm2_nobias_op(x, weight, output, False) + return output + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ll_bf16_gemm + + return ll_bf16_gemm(x, weight, output_dtype=torch.bfloat16) diff --git a/vllm/parser/beam.py b/vllm/parser/beam.py new file mode 100644 index 000000000000..9f99e321f19f --- /dev/null +++ b/vllm/parser/beam.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Parser for Beam checkpoints.""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any, cast + +from openai.types.responses import ToolChoiceFunction + +from vllm import envs +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedToolChoiceParam, + ChatCompletionRequest, +) +from vllm.exceptions import VLLMValidationError +from vllm.logger import init_logger +from vllm.parser.abstract_parser import ( + _xgrammar_supports, + structured_outputs_to_format, +) +from vllm.parser.chat_parsing import ResponseParser +from vllm.parser.chat_parsing.response_templates import ResponseTemplate +from vllm.parser.response_template import ( + CONTENT_FIELD, + THINKING_FIELD, + TOOL_FIELD, + ResponseTemplateParser, + ResponseTemplateReasoningParser, + ResponseTemplateToolParser, + _as_tool_dicts, + _is_named_tool_choice, + _tool_parsing_enabled, +) +from vllm.parser.response_template_grammar import ( + _openers, + answer_rest_structural_tag, + message_start_anchors, + tool_call_structural_tag, +) +from vllm.sampling_params import StructuredOutputsParams +from vllm.tool_parsers.structural_tag_registry import resolve_tool_strictness + +if TYPE_CHECKING: + from vllm.entrypoints.openai.responses.protocol import ResponsesRequest + +logger = init_logger(__name__) + + +def _named_tool(request: ChatCompletionRequest | ResponsesRequest) -> str: + choice = request.tool_choice + if isinstance(choice, ToolChoiceFunction): + return choice.name + return cast(ChatCompletionNamedToolChoiceParam, choice).function.name + + +def _starts_message(template: ResponseTemplate, prefix: str) -> bool: + """Whether output after `prefix` starts the assistant message.""" + ends = [match.end() for match in template.start_anchor_re.finditer(prefix)] + return bool(ends) and ends[-1] == len(prefix) + + +def _open_field(template: ResponseTemplate, prefix: str) -> str | None: + """The field that `prefix` leaves open with its explicit opener, if any.""" + field = None + for event in ResponseParser(template, prefix=prefix).initial_events: + if event["type"] == "region_open" and event["end"] > event["start"]: + field = event["field"] + elif event["type"] in {"region_close", "region_malformed"}: + field = None + return field + + +class BeamReasoningParser(ResponseTemplateReasoningParser): + """Selects `BeamParser` as the reasoning parser.""" + + @classmethod + def unified_parser_cls(cls) -> type[ResponseTemplateParser]: + return BeamParser + + @property + def reasoning_start_str(self) -> str | None: + if literal := super().reasoning_start_str: + return literal + thinking = self.response_template.fields.get(THINKING_FIELD) + if thinking is None: + return None + openers = _openers(thinking, message_start_anchors(self.response_template)) + if not openers: + return None + # Budget tracking needs one marker for both a prompt continuation + # and a complete assistant message in the checkpoint's pattern. + marker = min(openers, key=len) + return marker if all(opener.endswith(marker) for opener in openers) else None + + +class BeamToolParser(ResponseTemplateToolParser): + """Selects `BeamParser` as the tool parser.""" + + @classmethod + def unified_parser_cls(cls) -> type[ResponseTemplateParser]: + return BeamParser + + +class BeamParser(ResponseTemplateParser): + """Parse Beam output with the checkpoint's `response_template` metadata, + and constrain it with a grammar derived from that template.""" + + reasoning_parser_cls = BeamReasoningParser + tool_parser_cls = BeamToolParser + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + chat_template_kwargs = kwargs.get("chat_template_kwargs") or {} + # From the message start, the grammar opens the message with non-empty + # reasoning, unless the request turns thinking off. + self._requires_reasoning = bool( + chat_template_kwargs.get("enable_thinking", True) + ) + + def _constrain_request( + self, + request: ChatCompletionRequest | ResponsesRequest, + ) -> ChatCompletionRequest | ResponsesRequest: + message_start = _starts_message(self.response_template, self._prefix) + tag = self._template_tag(request, message_start=message_start) + if tag is not None: + grammar = StructuredOutputsParams(structural_tag=json.dumps(tag)) + if isinstance(request, ChatCompletionRequest): + request.structured_outputs = grammar + request.response_format = None + else: + request._parser_structured_outputs = grammar + # From the message start, the tag also covers the reasoning. + request._grammar_from_parser = message_start + return request + + def _template_tag( + self, + request: ChatCompletionRequest | ResponsesRequest, + *, + message_start: bool, + ) -> dict[str, Any] | None: + """The template grammar for `request`, or `None` to leave its output + to the request's own constraint, if any.""" + fields = self.response_template.fields + named = _is_named_tool_choice(request) + tools_enabled = ( + self._tool_parser is not None + and bool(request.tools) + and _tool_parsing_enabled( + request, enable_auto_tools=self._enable_auto_tools + ) + ) + forced = tools_enabled and (request.tool_choice == "required" or named) + # The prompt may still list the tools, so only the grammar keeps the + # model from calling them. + tools_forbidden = ( + bool(request.tools) + and request.tool_choice == "none" + and TOOL_FIELD in fields + ) + constrained_tools = None + if tools_enabled and (forced or envs.VLLM_ENFORCE_STRICT_TOOL_CALLING): + assert self._tool_parser is not None + strict_level = self.tool_strict_level + if strict_level is None: + strict_level = self._tool_parser.default_tool_strict_level + constrained_tools = resolve_tool_strictness( + request.tools or [], request.tool_choice, strict_level + ) + output = request.extract_structured_outputs() + if forced and output is not None: + logger.warning_once( + "structured outputs are ignored because tool_choice forces tool call." + ) + output = None + answer = fields.get(CONTENT_FIELD) + if output is not None and ( + output.structural_tag is not None + or answer is None + or answer.open_re is None + or not _xgrammar_supports(output) + ): + if constrained_tools is not None: + logger.warning_once( + "response_template parsing cannot enforce strict tools " + "alongside this structured output." + ) + return None + # Without a tool or structured output constraint, the output still + # follows the template. + always_constrain = envs.VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT and not forced + if ( + constrained_tools is None + and output is None + and not tools_forbidden + and not always_constrain + ): + return None + + # Inside an open answer, only the always-on grammar applies, finishing a + # free answer. Inside open reasoning, the grammar needs the reasoning + # gate to defer it to the reasoning's end. + open_field = _open_field(self.response_template, self._prefix) + reasoning_gated = ( + self._reasoning_parser is not None and request.include_reasoning + ) + tag = None + if open_field == CONTENT_FIELD: + if always_constrain and output is None: + tag = answer_rest_structural_tag(fields[CONTENT_FIELD]) + elif open_field != THINKING_FIELD or reasoning_gated: + functions = [] + if tools_enabled: + tools = _as_tool_dicts(constrained_tools or request.tools) + functions = [tool.get("function", tool) for tool in tools] + if named: + name = _named_tool(request) + functions = [f for f in functions if f.get("name") == name] + answer_format = ( + None if output is None else structured_outputs_to_format(output) + ) + try: + tag = tool_call_structural_tag( + fields.get(TOOL_FIELD), + functions, + parallel_tool_calls=( + not named and request.parallel_tool_calls is not False + ), + thinking=fields.get(THINKING_FIELD), + answer=None if forced else answer, + answer_format=( + None if answer_format is None else answer_format.model_dump() + ), + message_start=( + message_start_anchors(self.response_template) + if message_start + else None + ), + requires_reasoning=self._requires_reasoning, + # Beam drifts from schemas the grammar does not enforce, + # and also writes arguments sorted by name. + sorted_arguments=True, + strict_arguments=True, + ) + except ValueError as exc: + raise VLLMValidationError(str(exc), parameter="tools") from exc + if tag is None and forced: + kind = "named" if named else "required" + raise VLLMValidationError( + f"response_template parsing cannot guarantee {kind} tool choice " + "for this request", + parameter="tool_choice", + ) + if tag is None and constrained_tools is not None: + logger.warning_once( + "response_template parsing cannot enforce strict tools for this " + "request." + ) + return tag diff --git a/vllm/parser/parser_manager.py b/vllm/parser/parser_manager.py index d6fdbce19ca8..b6f00cacd863 100644 --- a/vllm/parser/parser_manager.py +++ b/vllm/parser/parser_manager.py @@ -16,8 +16,6 @@ logger = init_logger(__name__) -HF_PARSER = "hf" - class ParserManager: """Provides a unified Parser by composing reasoning and tool parser adapters.""" @@ -122,8 +120,10 @@ def get_parser( tags (``--tool-strict-level``), or ``None`` for the tool parser's default level. tokenizer: Tokenizer the composed parser is validated against - at startup (see :meth:`validate_parser_tokenizer`). ``None`` - skips the check (e.g. ``skip_tokenizer_init``). + at startup (see :meth:`validate_parser_tokenizer`), including + its `response_template` metadata when a response-template + parser is selected. ``None`` skips the check (e.g. + ``skip_tokenizer_init``). Returns: A Parser class, or None if neither parser is specified. @@ -184,39 +184,40 @@ def _compose_parser( HarmonyParser.tool_strict_level = strict_level return HarmonyParser - if HF_PARSER in (reasoning_parser_name, tool_parser_name): - if {reasoning_parser_name, tool_parser_name} - { - HF_PARSER, - None, - "", - }: + from vllm.parser.response_template import ( + response_template_parser_cls, + validate_tokenizer_response_template, + ) + + template_parsers = { + response_template_parser_cls(parser_cls) + for parser_cls in (reasoning_parser_cls, tool_parser_cls) + if parser_cls is not None + } + template_parser_cls = next(iter(template_parsers - {None}), None) + if template_parser_cls is not None: + if len(template_parsers) > 1: raise TypeError( - "The hf parser cannot be combined with other " - "reasoning or tool call parsers" + f"--reasoning-parser {reasoning_parser_name} and " + f"--tool-call-parser {tool_parser_name} cannot be combined" ) - from vllm.parser.response_template import ( - ResponseTemplateParser, - validate_tokenizer_response_template, - ) - if tokenizer is not None: validate_tokenizer_response_template( tokenizer, - reasoning=reasoning_parser_cls is not None, - tools=tool_parser_cls is not None, + reasoning_parser=reasoning_parser_name, + # Without auto tool choice, no tool parser is selected. + tool_parser=tool_parser_name if tool_parser_cls else None, ) - - r_cls = reasoning_parser_cls - t_cls = tool_parser_cls - auto_tools = enable_auto_tools - - class _ResponseTemplateParser(ResponseTemplateParser): - reasoning_parser_cls = r_cls - tool_parser_cls = t_cls - tool_strict_level = strict_level - _enable_auto_tools = auto_tools - - return _ResponseTemplateParser + return type( + template_parser_cls.__name__, + (template_parser_cls,), + { + "reasoning_parser_cls": reasoning_parser_cls, + "tool_parser_cls": tool_parser_cls, + "tool_strict_level": strict_level, + "_enable_auto_tools": enable_auto_tools, + }, + ) if reasoning_parser_name == "kimi_k3" or tool_parser_name == "kimi_k3": from vllm.parser.kimi_k3 import KimiK3Parser diff --git a/vllm/parser/response_template.py b/vllm/parser/response_template.py index b1665f1772a5..ea9556730549 100644 --- a/vllm/parser/response_template.py +++ b/vllm/parser/response_template.py @@ -7,6 +7,7 @@ import copy import json from collections.abc import Callable, Iterable, Sequence +from math import inf from typing import TYPE_CHECKING, Any, cast from openai.types.responses import ToolChoiceCustom, ToolChoiceFunction @@ -28,6 +29,7 @@ ResponseTemplate, load_response_template, ) +from vllm.parser.response_template_grammar import message_start_anchors 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 @@ -135,27 +137,38 @@ def validate_response_template_for_serving( def validate_tokenizer_response_template( tokenizer: Any, *, - reasoning: bool, - tools: bool, + reasoning_parser: str | None, + tool_parser: str | None, ) -> None: - """Fail at startup if the tokenizer cannot serve the requested parsing.""" + """Fail at startup if the tokenizer cannot serve the selected parsers, + named as given on the command line.""" template = resolve_response_template(tokenizer) if template is None: + selected = " ".join( + f"{flag} {name}" + for flag, name in ( + ("--reasoning-parser", reasoning_parser), + ("--tool-call-parser", tool_parser), + ) + if name + ) raise TypeError( - "The hf parser requires `response_template` metadata " + f"{selected} requires `response_template` metadata " "in the tokenizer configuration" ) try: fields = validate_response_template_for_serving(template).fields except (TypeError, ValueError) as exc: raise TypeError(f"Invalid response_template metadata: {exc}") from exc - if reasoning and THINKING_FIELD not in fields: + if reasoning_parser and THINKING_FIELD not in fields: raise TypeError( - "--reasoning-parser hf requires a response_template with a thinking field" + f"--reasoning-parser {reasoning_parser} requires a response_template " + "with a thinking field" ) - if tools and TOOL_FIELD not in fields: + if tool_parser and TOOL_FIELD not in fields: raise TypeError( - "--tool-call-parser hf requires a response_template with a tool_calls field" + f"--tool-call-parser {tool_parser} requires a response_template " + "with a tool_calls field" ) @@ -292,7 +305,7 @@ def _serving_template( template = resolve_response_template(tokenizer, template) if template is None: raise ValueError( - "The hf parser requires `response_template` " + "Parsing with a response template requires `response_template` " "metadata in the tokenizer configuration" ) validate_response_template_for_serving(template) @@ -319,12 +332,17 @@ def _to_tool_calls(value: Any) -> list[tuple[str, str]]: class ResponseTemplateReasoningParser(ReasoningParser): - """Reasoning end detection for `--reasoning-parser hf`. + """Reasoning end detection for a response-template parser. `ResponseTemplateParser` parses the output; this parser tells it and the structured-output engine where reasoning ends. """ + @classmethod + def unified_parser_cls(cls) -> type[ResponseTemplateParser]: + """The unified parser that selecting this reasoning parser stands for.""" + return ResponseTemplateParser + def __init__( self, tokenizer: TokenizerLike, @@ -400,8 +418,26 @@ def extract_reasoning_streaming( class ResponseTemplateToolParser(ToolParser): - """Registers `--tool-call-parser hf`; `ResponseTemplateParser` extracts the - tool calls.""" + """Selects a response-template parser as the tool parser; + `ResponseTemplateParser` extracts the tool calls.""" + + @classmethod + def unified_parser_cls(cls) -> type[ResponseTemplateParser]: + """The unified parser that selecting this tool parser stands for.""" + return ResponseTemplateParser + + +def response_template_parser_cls( + split_parser_cls: type[ReasoningParser] | type[ToolParser], +) -> type[ResponseTemplateParser] | None: + """The response-template parser that a registered reasoning or tool parser + stands for, or `None` for any other parser.""" + if issubclass( + split_parser_cls, + (ResponseTemplateReasoningParser, ResponseTemplateToolParser), + ): + return split_parser_cls.unified_parser_cls() + return None class _DeltaBuilder: @@ -441,23 +477,50 @@ def __init__( self.include_reasoning = include_reasoning self.parse_tools = parse_tools self.new_tool_call_id = new_tool_call_id + self.template = template self.parser = ResponseParser(template, prefix=prefix, tools=tools) self._prefix_end = len(self.parser.input_text) self._next_tool_index = 0 self._finalized = False + self._open_region: dict[str, Any] | None = None + self._track(self.parser.initial_events) def feed(self, text: str, *, finished: bool) -> DeltaMessage | None: if self._finalized: return None delta = _DeltaBuilder() - self._route(self.parser.feed(text), delta) + self._route(self._track(self.parser.feed(text)), delta) if finished: self._finalized = True if len(self.parser.input_text) > self._prefix_end: - _, events = self.parser.finalize() - self._route(events, delta) + events = self.parser.feed(self._closer_rest()) + self._route(events + self.parser.finalize()[1], delta) return delta.build() + def _track(self, events: list[dict[str, Any]]) -> list[dict[str, Any]]: + for event in events: + if event["type"] == "region_open": + self._open_region = event + elif event["type"] in ("region_close", "region_malformed"): + self._open_region = None + return events + + def _closer_rest(self) -> str: + """The rest of a closer the output stopped in after ending its region, + when only the next message's start anchor is missing.""" + if self._open_region is None: + return "" + field = self.template.fields[self._open_region["field"]] + text = self.parser.input_text[self._open_region["end"] :] + for closer in field.close_literals or (): + for anchor in message_start_anchors(self.template): + if not anchor or not closer.endswith(anchor) or closer == anchor: + continue + for cut in range(len(closer) - 1, len(closer) - len(anchor) - 1, -1): + if text.endswith(closer[:cut]): + return closer[cut:] + return "" + def _route(self, events: Sequence[dict[str, Any]], delta: _DeltaBuilder) -> None: for event in events: field = event.get("field") @@ -499,6 +562,54 @@ def _route_tool_event(self, event: dict[str, Any], delta: _DeltaBuilder) -> None self._next_tool_index += 1 +class _ReasoningTokenCounter: + """Counts the reasoning tokens of one growing generation, parsing each + token once.""" + + def __init__( + self, template: ResponseTemplate, prefix: str, tokenizer: TokenizerLike + ) -> None: + self.tokenizer = tokenizer + self.parser = ResponseParser(template, prefix=prefix) + self.start = len(self.parser.input_text) + self.token_ids: list[int] = [] + self.token_ends: list[int] = [] + # Input spans between the delimiters of thinking regions. + self.spans: list[tuple[int, float]] = [] + self.open: int | None = None + self._track(self.parser.initial_events) + + def continues(self, token_ids: Sequence[int]) -> bool: + return list(token_ids[: len(self.token_ids)]) == self.token_ids + + def count(self, token_ids: Sequence[int]) -> int: + new_ids = list(token_ids[len(self.token_ids) :]) + texts = [_decode(self.tokenizer, [token_id]) for token_id in new_ids] + end = self.token_ends[-1] if self.token_ends else self.start + for text in texts: + end += len(text) + self.token_ends.append(end) + self.token_ids.extend(new_ids) + self._track(self.parser.feed("".join(texts))) + # Text held back as a possible closer counts until a closer arrives. + spans = self.spans if self.open is None else [*self.spans, (self.open, inf)] + count, start = 0, self.start + for end in self.token_ends: + count += any(begin < end and start < close for begin, close in spans) + start = end + return count + + def _track(self, events: Iterable[dict[str, Any]]) -> None: + for event in events: + if event["field"] != THINKING_FIELD or event["type"] == "region_chunk": + continue + if event["type"] == "region_open": + self.open = event["end"] + elif self.open is not None: + self.spans.append((self.open, event["start"])) + self.open = None + + class ResponseTemplateParser(Parser): """Parse reasoning, content, and tool calls with one response template.""" @@ -526,6 +637,7 @@ def __init__( self._prompt_token_ids: list[int] | None = None self._prefix = "" self._stream: _ResponseStream | None = None + self._reasoning_counter: _ReasoningTokenCounter | None = None def set_prompt_token_ids(self, prompt_token_ids: Sequence[int]) -> None: prompt_token_ids = list(prompt_token_ids) @@ -534,6 +646,7 @@ def set_prompt_token_ids(self, prompt_token_ids: Sequence[int]) -> None: self._prompt_token_ids = prompt_token_ids self._prefix = _decode(self.model_tokenizer, prompt_token_ids) self._stream = None + self._reasoning_counter = None def adjust_request( self, @@ -547,6 +660,14 @@ def adjust_request( request.skip_special_tokens = False if hasattr(request, "spaces_between_special_tokens"): request.spaces_between_special_tokens = False + return self._constrain_request(request) + + def _constrain_request( + self, + request: ChatCompletionRequest | ResponsesRequest, + ) -> ChatCompletionRequest | ResponsesRequest: + """Constrain `request` to the tool guarantees it asks for, or reject + those this parser cannot enforce.""" if ( self._tool_parser is not None and request.tools @@ -567,6 +688,21 @@ def is_reasoning_end(self, input_ids: list[int]) -> bool: or self._reasoning_parser.is_reasoning_end(input_ids) ) + def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: + """Count the generated tokens whose text the template parses as + reasoning, continuing reasoning left open by the prompt.""" + if ( + self._reasoning_parser is None + or THINKING_FIELD not in self.response_template.fields + ): + return 0 + counter = self._reasoning_counter + if counter is None or not counter.continues(token_ids): + counter = self._reasoning_counter = _ReasoningTokenCounter( + self.response_template, self._prefix, self.model_tokenizer + ) + return counter.count(token_ids) + def parse( self, model_output: str, diff --git a/vllm/parser/response_template_grammar.py b/vllm/parser/response_template_grammar.py new file mode 100644 index 000000000000..40317023f485 --- /dev/null +++ b/vllm/parser/response_template_grammar.py @@ -0,0 +1,1000 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Structural tags for tool calls and answers, derived from a response template.""" + +from __future__ import annotations + +import json +from collections.abc import Callable +from functools import lru_cache +from typing import Any + +import regex as re + +try: + from re import _parser # type: ignore[attr-defined] +except ImportError: # Python 3.10 + import sre_parse as _parser + +from vllm.parser.chat_parsing.content_parsers import parse_content +from vllm.parser.chat_parsing.response_parser import _schema_types +from vllm.parser.chat_parsing.response_templates import ( + ResponseTemplate, + ResponseTemplateField, +) + +_sre_parse = _parser +_REGEX_META = re.compile(r"[.^$*+?{}\[\]\\|()]") +_LITERAL_ESCAPES = {"n": "\n", "r": "\r", "t": "\t"} +# Stands for `^` while expanding a delimiter pattern. +_START = "\0" +# Precedes the text of a final lookahead while expanding a delimiter pattern. +_AHEAD = "\1" +# Arbitrary bound on the strings expanded from one delimiter pattern. +_MAX_LITERALS = 16 +# Templates write at most a newline between parameters. Unbounded whitespace +# would let a model that skipped a parameter's opener write it until `max_tokens`. +_SEPARATOR = {"type": "regex", "pattern": r"\n?"} +_ANY_TEXT = {"type": "any_text"} +# Character classes in XGrammar regexes take whitespace only as escapes. +_BLANKS = {" ": " ", "\t": r"\t", "\n": r"\n", "\r": r"\r"} +# Keywords that describe a schema without constraining its values. +_ANNOTATIONS = frozenset( + { + "title", + "description", + "default", + "examples", + "deprecated", + "$comment", + "$schema", + "$id", + } +) +_OBJECT_KEYWORDS = frozenset( + {"type", "properties", "required", "additionalProperties", "$defs", "definitions"} +) +# Characters JSON writes outside strings. +_JSON_SYNTAX = frozenset(' \t\n\r{}[],:"-+.0123456789eEtrufalsn') +# Pieces of the text of a call for `_spells_closer`. Text and JSON hold no +# closer, and JSON ends with a character it writes outside strings. +_TEXT_PIECE = (r"[\s\S]*", r"[\s\S]+", False) +_JSON_PIECE = ( + r"[\s\S]*", + r"[\s\S]*[" + "".join(map(re.escape, sorted(_JSON_SYNTAX))) + "]", + False, +) +_REPEATS = tuple( + getattr(_sre_parse, name) + for name in ("MAX_REPEAT", "MIN_REPEAT", "POSSESSIVE_REPEAT") + if hasattr(_sre_parse, name) +) +_CATEGORIES = { + _sre_parse.CATEGORY_DIGIT: r"\d", + _sre_parse.CATEGORY_NOT_DIGIT: r"\D", + _sre_parse.CATEGORY_SPACE: r"\s", + _sre_parse.CATEGORY_NOT_SPACE: r"\S", + _sre_parse.CATEGORY_WORD: r"\w", + _sre_parse.CATEGORY_NOT_WORD: r"\W", +} +_HOST_LABEL = r"[a-z0-9](?:[a-z0-9-]*[a-z0-9])?" +_DATE = r"[0-9]{4}-(?:0[1-9]|1[0-2])-(?:0[1-9]|[12][0-9]|3[01])" +_TIME = ( + r"(?:[01][0-9]|2[0-3]):[0-5][0-9]:(?:[0-5][0-9]|60)(?:\.[0-9]+)?" + r"(?:Z|[+-](?:[01][0-9]|2[0-3]):[0-5][0-9])" +) +_OCTET = r"(?:25[0-5]|2[0-4][0-9]|[01]?[0-9]?[0-9])" +_ATOM = r"[A-Za-z0-9_!#$%&'*+/=?^`{|}~-]+" +_URI_CHARS = r"A-Za-z0-9_.~!$&'()*+,;=" +_ESCAPED = r"%[0-9A-Fa-f]{2}" +_PATH = rf"(?:[{_URI_CHARS}:@-]|{_ESCAPED})" +_QUERY = rf"(?:[{_URI_CHARS}:@/?-]|{_ESCAPED})*" +# Raw string forms of the formats XGrammar holds in JSON strings. Emails leave +# out quoted local parts, whose quotes XGrammar writes as JSON escapes. +_STRING_FORMATS = { + "date": _DATE, + "time": _TIME, + "date-time": f"{_DATE}T{_TIME}", + "email": ( + rf"{_ATOM}(?:\.{_ATOM})*@[A-Za-z0-9](?:[A-Za-z0-9-]*[A-Za-z0-9])?" + r"(?:\.[A-Za-z0-9][A-Za-z0-9-]*[A-Za-z0-9])*" + ), + "hostname": rf"{_HOST_LABEL}(?:\.{_HOST_LABEL})*", + "ipv4": rf"(?:{_OCTET}\.){{3}}{_OCTET}", + "uuid": r"[0-9A-Fa-f]{8}(?:-[0-9A-Fa-f]{4}){3}-[0-9A-Fa-f]{12}", + "uri": ( + rf"[A-Za-z][A-Za-z+.-]*:(?://(?:(?:[{_URI_CHARS}:-]|{_ESCAPED})*@)?" + rf"(?:[{_URI_CHARS}-]|{_ESCAPED})*(?::[0-9]*)?(?:/{_PATH}*)*" + rf"|/?(?:{_PATH}+(?:/{_PATH}*)*)?)(?:\?{_QUERY})?(?:#{_QUERY})?" + ), +} + + +def _split_group(pattern: str, name: str) -> tuple[str, str, str] | None: + """Split a named regex group into its prefix, body and suffix.""" + opener = f"(?P<{name}>" + start = pattern.find(opener) + if start < 0: + return None + depth = 0 + in_class = False + index = start + while index < len(pattern): + char = pattern[index] + if char == "\\": + index += 2 + continue + if in_class: + in_class = char != "]" + elif char == "[": + in_class = True + elif char == "(": + depth += 1 + elif char == ")": + depth -= 1 + if depth == 0: + body = pattern[start + len(opener) : index] + return pattern[:start], body, pattern[index + 1 :] + index += 1 + return None + + +def _literals( + pattern: str, *, message_start: bool = False, lookahead: bool = False +) -> list[str] | None: + """Expand literal alternatives, allowing `^` only at the message start and + a final lookahead only with `lookahead`.""" + literals, index = _alternatives(pattern, 0) + if literals is None or index != len(pattern): + return None + if any(literal.count(_AHEAD) > lookahead for literal in literals): + return None + if not message_start: + return None if any(_START in literal for literal in literals) else literals + return [ + literal.removeprefix(_START) + for literal in literals + if _START not in literal[1:] and literal != _START + ] + + +def _alternatives(pattern: str, index: int) -> tuple[list[str] | None, int]: + literals = [] + while True: + branch, index = _concatenation(pattern, index) + if branch is None: + return None, index + literals.extend(branch) + if not pattern.startswith("|", index): + break + index += 1 + literals = list(dict.fromkeys(literals)) + return (literals if len(literals) <= _MAX_LITERALS else None), index + + +def _concatenation(pattern: str, index: int) -> tuple[list[str] | None, int]: + literals = [""] + while index < len(pattern) and pattern[index] not in "|)": + if pattern.startswith(("(?:", "(?="), index): + ahead = pattern[index + 2] == "=" + atom, index = _alternatives(pattern, index + 3) + if atom is None or not pattern.startswith(")", index): + return None, index + index += 1 + if ahead: + atom = [_AHEAD + literal for literal in atom] + elif pattern[index] == "\\" and index + 1 < len(pattern): + escaped = pattern[index + 1] + if escaped.isalnum() and escaped not in _LITERAL_ESCAPES: + return None, index + atom = [_LITERAL_ESCAPES.get(escaped, escaped)] + index += 2 + elif pattern[index] == "^": + atom = [_START] + index += 1 + elif _REGEX_META.fullmatch(pattern[index]): + return None, index + else: + atom = [pattern[index]] + index += 1 + if pattern.startswith("?", index): + atom = ["", *atom] + index += 1 + # A lookahead reads as a literal only at the end of the pattern. + if any(_AHEAD in prefix for prefix in literals): + return None, index + literals = [prefix + suffix for prefix in literals for suffix in atom] + if len(literals) > _MAX_LITERALS: + return None, index + return literals, index + + +def _openers( + field: ResponseTemplateField, message_start: list[str] | None = None +) -> list[str] | None: + """Derive field openers, including continuations of message-start anchors.""" + if field.open_literals or field.open_re is None: + openers = field.open_literals + else: + openers = _literals( + field.open_re.pattern, message_start=message_start is not None + ) + return _continuing(openers, message_start, field.open_re) + + +def _continuing( + literals: list[str] | None, anchors: list[str] | None, opener: re.Pattern +) -> list[str] | None: + """Keep anchor continuations only where the parser finds the opener.""" + if not literals or not anchors: + return literals + rests = [ + literal[len(anchor) :] + for literal in literals + for anchor in anchors + if len(literal) > len(anchor) and literal.startswith(anchor) + ] + rests = [ + rest + for rest in rests + if any(opener.fullmatch(rest, index) for index in range(len(rest))) + ] + return list(dict.fromkeys([*literals, *rests])) + + +def _escape(text: str) -> str: + return _REGEX_META.sub(lambda match: "\\" + match.group(), text) + + +def _call_opener( + opener: tuple[str, str, str], name: str, message_start: list[str] | None +) -> str: + before, _, after = opener + if message_start is not None and before: + literals = _literals(before, message_start=True) + # Text before a name inside a group is no pattern of its own. + if literals: + literals = _continuing(literals, message_start, re.compile(before)) + before = "(?:" + "|".join(map(_escape, literals or ())) + ")" + return before + name + after + + +def _other_names(pattern: str, names: list[str]) -> str | None: + """A regex for the names `pattern` matches besides `names`, or `None` if + `pattern` is no negatable character class repeated.""" + if not names: + return pattern + match = re.fullmatch(r"\\([wS])\+|\[\^((?:\\.|[^\]\\])+)\]\+", pattern) + if match is None: + return None + negated = {"w": r"\W", "S": r"\s"}[match[1]] if match[1] else match[2] + + def other(chars: set[str]) -> str: + return f"[^{negated}{''.join(map(_class_escape, sorted(chars)))}]" + + trie: dict[str, dict] = {} + for name in names: + node = trie + for char in name: + node = node.setdefault(char, {}) + node[""] = {} + + def rest(node: dict[str, dict], first: bool) -> str: + """The rest of a name past `node` that ends on none of `names`.""" + branches = [other(set(node) - {""}) + other(set()) + "*"] + branches += [ + _escape(char) + rest(child, False) for char, child in node.items() if char + ] + return f"(?:{'|'.join(branches)})" + ("" if first or "" in node else "?") + + return rest(trie, True) + + +def _thinking_format( + thinking: ResponseTemplateField, + message_start: list[str] | None, + nonempty: bool = False, +) -> dict[str, Any] | None: + """The reasoning, which holds at least a character if `nonempty`.""" + openers, closers = _openers(thinking, message_start), _closers(thinking) + if not openers or not closers: + return None + text = _sequence(_nonempty_start(closers), _ANY_TEXT) if nonempty else _ANY_TEXT + return _one_of([_tag(opener, text, closers) for opener in openers]) + + +def _nonempty_start(closers: list[str]) -> dict[str, Any]: + """The start of text that is not empty and does not begin with one of the + `closers`. Each branch leaves every closer's spelling before completing it, + on a character that starts no closer, so the text after it cannot hide a + closer's start and a closer written as one token cannot pass for the text.""" + closers = [closer for closer in closers if closer] + starts = {closer[0] for closer in closers} + prefixes = dict.fromkeys( + [""] + [closer[:end] for closer in closers for end in range(1, len(closer))] + ) + branches = [] + for prefix in prefixes: + # A prefix holding a whole closer has already closed the text. + if any(closer in prefix for closer in closers): + continue + nexts = { + closer[len(prefix)] + for closer in closers + if len(closer) > len(prefix) and closer.startswith(prefix) + } + excluded = "".join(map(_class_escape, sorted(starts | nexts))) + branches.append(f"{_escape(prefix)}[^{excluded}]") + return {"type": "regex", "pattern": "|".join(branches)} + + +def _class_escape(char: str) -> str: + if char in _BLANKS: + return _BLANKS[char] + return "\\" + char if char in "\\]^-[" else char + + +def message_start_anchors(template: ResponseTemplate) -> list[str]: + """The strings the assistant message starts after, if literal.""" + if template.start_anchor_literals: + return template.start_anchor_literals + return _literals(template.start_anchor_re.pattern) or [] + + +def _closers(field: ResponseTemplateField) -> list[str] | None: + """Return literal closers, an empty list if absent, or `None` if underivable.""" + if field.close_re is None: + return [] + return field.close_literals or _literals(field.close_re.pattern) + + +# Every request derives the tags for its tools anew. +@lru_cache(maxsize=512) +def _endings(literals: tuple[str, ...]) -> str: + return "|".join( + re.escape(literal[index:]) + for literal in literals + for index in range(len(literal)) + ) + + +def _literal_piece(literals: list[str], before: str = "") -> tuple[str, str, bool]: + either = before + "(?:" + "|".join(map(re.escape, literals)) + ")" + return either, f"{either}|{_endings(tuple(literals))}", True + + +@lru_cache(maxsize=512) +def _spells_closer( + closers: tuple[str, ...], + pieces: tuple[tuple[str, str, bool], ...], + *, + repeats: bool = False, +) -> bool: + """Check whether argument framing can spell a call closer across pieces.""" + texts = [text for text, _, _ in pieces] + loop = f"(?:{''.join(texts)})*" if repeats else "" + # Another call or an answer may follow the call's own closer. + closed = "(?:" + "|".join(map(re.escape, closers)) + r")[\s\S]*" + for index, (_, endings, holds) in enumerate(pieces): + ends = re.compile(endings) + rest = re.compile("".join(texts[index + 1 :]) + loop + closed, re.DOTALL) + for closer in closers: + if holds and ends.fullmatch(closer, partial=True): + return True + # A closer that starts in this piece and ends in the text after it. + for split in range(1, len(closer)): + if ends.fullmatch(closer[:split]) and rest.fullmatch( + closer[split:], partial=True + ): + return True + return False + + +def _const_strings(values: list[str]) -> dict[str, Any]: + return _one_of([{"type": "const_string", "value": value} for value in values]) + + +def _sequence(*elements: dict[str, Any]) -> dict[str, Any]: + return {"type": "sequence", "elements": list(elements)} + + +def _one_of(formats: list[dict[str, Any]]) -> dict[str, Any]: + return formats[0] if len(formats) == 1 else {"type": "or", "elements": formats} + + +def _optional(content: dict[str, Any]) -> dict[str, Any]: + return {"type": "optional", "content": content} + + +def _tag(begin: str, content: dict[str, Any], end: str | list[str]) -> dict[str, Any]: + return {"type": "tag", "begin": begin, "content": content, "end": end} + + +def _json_value_schema(schema: Any, parameters: dict[str, Any]) -> dict[str, Any]: + if not isinstance(schema, dict): + return {} + if schema.get("nullable"): + # XGrammar ignores `nullable`. + schema = {"anyOf": [schema, {"type": "null"}]} + definitions = { + key: parameters[key] for key in ("$defs", "definitions") if key in parameters + } + return {**definitions, **schema} + + +def _string_choices(schema: Any) -> list[str] | None: + if not isinstance(schema, dict): + return None + choices = [schema["const"]] if "const" in schema else schema.get("enum") + if not isinstance(choices, list) or not choices: + return None + return choices if all(isinstance(choice, str) for choice in choices) else None + + +def _typed_by_parser(schema: Any) -> bool: + """Whether the parser reads schema values with their JSON type.""" + if schema is True or _schema_types(schema): + return True + return isinstance(schema, dict) and ( + not set(schema) - _ANNOTATIONS or _string_choices(schema) is not None + ) + + +def _text_format(schema: Any, excludes: list[str]) -> dict[str, Any] | None: + """Constrain raw text, which the `excludes` end, without changing its + parsed value.""" + if schema is True: + return _ANY_TEXT + if isinstance(schema, dict) and {"format", "pattern"} & set(schema): + return _string_regex(schema, excludes) + if not isinstance(schema, dict) or set(schema) - _ANNOTATIONS - { + "type", + "nullable", + "enum", + "const", + "anyOf", + "oneOf", + }: + return None + branches = [*(schema.get("anyOf") or []), *(schema.get("oneOf") or [])] + if "enum" in schema or "const" in schema: + choices = _string_choices(schema) + if choices is None or branches or set(_schema_types(schema)) - {"string"}: + return None + return _const_strings(choices) + if any(_text_format(branch, excludes) is not _ANY_TEXT for branch in branches): + return None + return _ANY_TEXT + + +def _string_regex(schema: dict[str, Any], excludes: list[str]) -> dict[str, Any] | None: + """Hold the `format` or `pattern` of a raw string. Like XGrammar, a format + overrides a pattern, either one drops length bounds, and the pattern + matches the whole string.""" + if set(schema) - _ANNOTATIONS - { + "type", + "format", + "pattern", + "minLength", + "maxLength", + } or _schema_types(schema) != ("string",): + return None + pattern = ( + _STRING_FORMATS.get(schema["format"]) + if "format" in schema + else schema["pattern"] + ) + if ( + not isinstance(pattern, str) + or _may_write(pattern, {text[0] for text in excludes if text}) + or not _regex_compiles(pattern) + ): + return None + return {"type": "regex", "pattern": pattern} + + +def _may_write(pattern: str, chars: set[str]) -> bool: + """Whether a match of `pattern` may hold one of `chars`, as any pattern may + that Python cannot parse.""" + try: + parsed = _sre_parse.parse(pattern) + except (_sre_parse.error, OverflowError, RecursionError): + return True + if parsed.state.flags & _sre_parse.SRE_FLAG_IGNORECASE: + return True + + def may(items: Any) -> bool: + for op, arg in items: + if op is _sre_parse.LITERAL: + written = {chr(arg)} + elif op is _sre_parse.IN: + written = {char for char in chars if _in_class(arg, char)} + elif op in _REPEATS: + written = chars if arg[1] and may(arg[2]) else set() + elif op is _sre_parse.SUBPATTERN: + if arg[1] & _sre_parse.SRE_FLAG_IGNORECASE: + return True + written = chars if may(arg[3]) else set() + elif op is _sre_parse.BRANCH: + written = chars if any(map(may, arg[1])) else set() + elif op is _sre_parse.AT: + continue + else: + # Such as `.`, which XGrammar lets match line breaks. + return True + if written & chars: + return True + return False + + return may(parsed) + + +def _in_class(items: list[tuple[Any, Any]], char: str) -> bool: + negated = bool(items) and items[0][0] is _sre_parse.NEGATE + matched = False + for op, arg in items[negated:]: + if op is _sre_parse.LITERAL: + matched |= ord(char) == arg + elif op is _sre_parse.RANGE: + matched |= arg[0] <= ord(char) <= arg[1] + elif op is _sre_parse.CATEGORY and arg in _CATEGORIES: + matched |= bool(re.fullmatch(_CATEGORIES[arg], char)) + else: + return True + return matched != negated + + +# Every request derives the tags for its tools anew. +@lru_cache(maxsize=512) +def _regex_compiles(pattern: str) -> bool: + from xgrammar import Grammar + + try: + Grammar.from_regex(pattern) + except RuntimeError: + return False + return True + + +def _value_format( + schema: Any, parameters: dict[str, Any], json_values: bool, closers: list[str] +) -> dict[str, Any] | None: + """Use JSON for decoded values and raw text for string values.""" + if schema is False: + return None + types = _schema_types(schema) + if json_values or (types and "string" not in types): + return _json_format(_json_value_schema(schema, parameters), closers) + return _text_format(schema, closers) + + +def _json_format(schema: Any, closers: list[str]) -> dict[str, Any] | None: + """Use JSON unless a closer is ambiguous with its syntax or the schema is + unsupported. Strings hold no closer, as the parser ends the call there.""" + if closers and ( + any(set(closer) <= _JSON_SYNTAX for closer in closers) + or not _json_compiles(json.dumps(schema)) + ): + return None + return {"type": "json_schema", "json_schema": schema, "excludes": closers} + + +# Every request derives the tags for its tools anew. +@lru_cache(maxsize=512) +def _json_compiles(schema: str) -> bool: + from xgrammar import Grammar + + try: + Grammar.from_json_schema(schema) + except RuntimeError: + return False + return True + + +def _xml_arguments_format( + content_args: dict[str, Any], + parameters: dict[str, Any], + *, + strict: bool, + closers: list[str], + sorted_arguments: bool, +) -> dict[str, Any] | None: + if set(content_args) - {"tag_pattern", "value_parser", "merge_duplicates"}: + return None + value_parser = content_args.get("value_parser") or {"name": "text"} + if value_parser.get("name") not in ("text", "json") or ( + value_parser.get("name") == "json" and value_parser.get("args") + ): + return None + key_split = _split_group(content_args["tag_pattern"], "key") + value_split = key_split and _split_group(key_split[2], "value") + if not key_split or not value_split: + return None + key_pattern = key_split[1] + begins, middles = _literals(key_split[0]), _literals(value_split[0]) + ends = _literals(value_split[2], lookahead=True) + if not begins or not middles or not ends: + return None + aheads = [end.split(_AHEAD) for end in ends if _AHEAD in end] + ends = [end for end in ends if _AHEAD not in end] + # A value may also end where the next parameter opens, if the template + # looks ahead for its only opener. + if not ends or ( + aheads and (len(begins) > 1 or {a for _, a in aheads} != {*begins}) + ): + return None + follows = [before + ahead for before, ahead in aheads] + properties = parameters.get("properties") + json_values = value_parser.get("name") == "json" + typed = properties is None or ( + isinstance(properties, dict) + and (json_values or all(map(_typed_by_parser, properties.values()))) + ) + # Other schema shapes cannot be typed by the XML parser. + if strict and (set(parameters) - _ANNOTATIONS - _OBJECT_KEYWORDS or not typed): + return None + required = parameters.get("required") or [] + names: tuple[str, ...] = () + if strict and properties is not None: + # Required entries that are no names leave the tool to the JSON schema. + if not isinstance(required, list) or not all( + isinstance(key, str) for key in required + ): + return None + names = (*properties, *(key for key in required if key not in properties)) + pieces = ( + _literal_piece(begins, before=_SEPARATOR["pattern"]), + (f"(?:{key_pattern})", _endings(names), False), + _literal_piece(middles), + # Values of every format are some text. + _TEXT_PIECE, + _literal_piece(ends), + ) + if _spells_closer(tuple(closers), pieces, repeats=True): + return None + + # Raw text has no escaping for a call closer, nor, where a value may end + # before the next opener, for any value end. + excludes = closers + follows + ends if follows else closers + text = {"type": "any_text", "excludes": excludes} if excludes else _ANY_TEXT + # Each value but the last then ends by opening the next parameter, whose + # tag starts after the opener. + starts = [""] if follows else begins + + def tag( + begin: str, value: dict[str, Any], fits: list[str], last: bool + ) -> dict[str, Any]: + if follows and not last: + fits = [e + sep + begins[0] for e in fits for sep in ("", "\n")] + follows + return _tag(begin, value, fits) + + def separated(tags: list[dict[str, Any]]) -> dict[str, Any]: + return _one_of(tags) if follows else _sequence(_SEPARATOR, _one_of(tags)) + + def arguments(body: dict[str, Any], optional: bool) -> dict[str, Any]: + """`body` after the opener of its first parameter.""" + if not follows: + return body + body = _sequence(_SEPARATOR, _const_strings(begins), body) + return _optional(body) if optional else body + + def any_parameter(last: bool) -> dict[str, Any]: + name = {"type": "regex", "pattern": key_pattern} + content = _sequence(name, _const_strings(middles), text) + return separated([tag(start, content, ends, last) for start in starts]) + + def any_parameters() -> dict[str, Any]: + more = {"type": "star", "content": any_parameter(False)} + return ( + arguments(_sequence(more, any_parameter(True)), True) if follows else more + ) + + if not strict: + return any_parameters() + if properties is None: + extra = parameters.get("additionalProperties", True) + if extra is False: + return {"type": "const_string", "value": ""} + return any_parameters() if extra is True or extra == {} else None + # A declared key repeated as an extra one would bypass its schema. + if parameters.get("additionalProperties", False) is not False: + return None + + def read_back(written: str) -> dict[str, Any] | None: + """The arguments the parser reads from `written`, unless a closer ends + the call in it.""" + if any(closer in written for closer in closers): + return None + try: + return parse_content(written, "xml-inline", content_args) + except ValueError: + return None + + def parameter( + key: str, value: dict[str, Any], choices: list[str] | None, last: bool + ) -> dict[str, Any] | None: + if not re.fullmatch(key_pattern, key) or any( + list(read_back(b + key + m + "0" + e) or ()) != [key] + for b in begins + for m in middles + for e in ends + ): + raise ValueError(f"parameter name {key!r} cannot be encoded in tool calls") + if choices is None: + return separated( + [tag(s + key + m, value, ends, last) for s in starts for m in middles] + ) + # Optional framing whitespace must not change a text choice. + tags, framed = [], set() + for b, s in zip(begins, starts): + for m in middles: + by_ends: dict[tuple[str, ...], list[str]] = {} + for choice in choices: + fits = tuple( + e + for e in ends + if read_back(b + key + m + choice + e) == {key: choice} + ) + by_ends.setdefault(fits, []).append(choice) + for fits, group in by_ends.items(): + if fits: + tags.append( + tag(s + key + m, _const_strings(group), list(fits), last) + ) + framed.update(group) + if framed != set(choices): + return None + return separated(tags) + + if not names: + return {"type": "const_string", "value": ""} + more: dict[str, dict[str, Any]] = {} + final: dict[str, dict[str, Any]] = {} + for key in names: + # A required name without a schema takes any value. + schema = properties.get(key, True) + value = _value_format(schema, parameters, json_values, closers + follows + ends) + if value is None: + return None + choices = None + if value["type"] in ("const_string", "or"): + choices = _string_choices(schema) + value = text if value == _ANY_TEXT else value + joined = parameter(key, value, choices, last=False) + ending = parameter(key, value, choices, last=True) + if joined is None or ending is None: + return None + more[key], final[key] = joined, ending + + def ordered(keys: tuple[str, ...]) -> dict[str, Any]: + before = [more[k] if k in required else _optional(more[k]) for k in keys] + if not follows: + return _sequence(*before) + + def last(start: int, stop: int) -> dict[str, Any]: + """Parameters of `keys[start:stop]` that end at one of them.""" + if stop - start == 1: + return final[keys[start]] + # Halving keeps the tag O(n log n) in size and O(log n) deep. + middle = (start + stop) // 2 + return _one_of( + [ + last(start, middle), + _sequence(*before[start:middle], last(middle, stop)), + ] + ) + + # The last parameter follows every required one. + first = max((i for i, k in enumerate(keys) if k in required), default=0) + return _sequence(*before[:first], last(first, len(keys))) + + orders = [names] + if sorted_arguments and sorted(names) != list(names): + orders.append(tuple(sorted(names))) + return arguments(_one_of([ordered(keys) for keys in orders]), not required) + + +def _arguments_format( + field: ResponseTemplateField, + parameters: dict[str, Any], + *, + strict: bool, + sorted_arguments: bool = False, +) -> dict[str, Any] | None: + closers = _closers(field) or [] + if field.content == "json" and not field.content_args: + if _spells_closer(tuple(closers), (_JSON_PIECE,)): + return None + schema = parameters if strict else {"type": "object"} + return _json_format(schema, closers) + if field.content == "xml-inline": + return _xml_arguments_format( + field.content_args, + parameters, + strict=strict, + closers=closers, + sorted_arguments=sorted_arguments, + ) + return None + + +def _names_function_by_opener(field: ResponseTemplateField) -> bool: + function = field.transform.get("function") if field.transform else None + return ( + field.open_re is not None + and "name" in field.open_re.groupindex + and not field.transform_each + and isinstance(function, dict) + and function.get("name") == "{name}" + and function.get("arguments") == "{content}" + ) + + +def _reply_format( + answer: ResponseTemplateField, + message_start: list[str] | None, + content: dict[str, Any] | None, +) -> dict[str, Any] | None: + """The `answer` field holding `content`, or free text if not given, or + `None` if it cannot be derived.""" + openers, closers = _openers(answer, message_start), _closers(answer) + if not openers or closers is None: + return None + return _sequence(_const_strings(openers), *_answer_text(closers, content)) + + +def _answer_text( + closers: list[str], content: dict[str, Any] | None = None +) -> list[dict[str, Any]]: + # Text cannot close the answer early, so no unconstrained call can follow + # it. Like a call, the answer may end at a stop token instead of a closer. + elements = [content or {"type": "any_text", "excludes": closers}] + if closers: + elements.append(_optional(_const_strings(closers))) + return elements + + +def answer_rest_structural_tag(answer: ResponseTemplateField) -> dict[str, Any] | None: + """The rest of an `answer` that is already open, or `None` if underivable.""" + closers = _closers(answer) + if closers is None: + return None + return {"type": "structural_tag", "format": _sequence(*_answer_text(closers))} + + +def tool_call_structural_tag( + field: ResponseTemplateField | None, + functions: list[dict[str, Any]], + *, + parallel_tool_calls: bool, + thinking: ResponseTemplateField | None = None, + answer: ResponseTemplateField | None = None, + answer_format: dict[str, Any] | None = None, + message_start: list[str] | None = None, + requires_reasoning: bool = False, + sorted_arguments: bool = False, + strict_arguments: bool = False, +) -> dict[str, Any] | None: + """Derive native calls to `functions` in the format of the tool call + `field`, respecting strict arguments, count and reasoning. + + `functions` are OpenAI function definitions. `strict_arguments` enforces + the schema of every function's arguments where the template can. Strict + `xml-inline` arguments come in declared order, or also alphabetical with + `sorted_arguments`. An `answer` field lets the model answer instead, in + `answer_format` if given; without `functions`, only the answer is derived. + When the model may answer, calls may also name any other function the + template can spell, with any arguments, so a harness can tell the model + the function does not exist. At the `message_start`, + the tag covers the output from its first token, which opens with non-empty + reasoning if `requires_reasoning`. Returns `None` if the template format + cannot be derived, and raises `ValueError` for names it cannot encode.""" + if not functions and answer is not None: + + def reply(start: list[str] | None) -> list[dict[str, Any]] | None: + format = _reply_format(answer, start, answer_format) + return None if format is None else [format] + + return _structural_tag(reply, thinking, message_start, requires_reasoning) + if not functions or field is None or not _names_function_by_opener(field): + return None + opener = _split_group(field.open_re.pattern, "name") + closer_literals = _closers(field) + if opener is None or closer_literals is None: + return None + arguments, enforced = [], [] + for function in functions: + # Like the generic JSON schema, a function without parameters takes none. + parameters = function.get("parameters") or {"type": "object", "properties": {}} + strict = bool(function.get("strict")) + function_arguments = None + if strict or strict_arguments: + try: + function_arguments = _arguments_format( + field, parameters, strict=True, sorted_arguments=sorted_arguments + ) + except ValueError: + if strict: + raise + enforced.append(function_arguments is not None) + if function_arguments is None and not strict: + function_arguments = _arguments_format(field, parameters, strict=False) + if function_arguments is None: + return None + arguments.append(function_arguments) + names = [] + for function in functions: + name = function["name"] + if not re.fullmatch(opener[1], name): + raise ValueError(f"tool name {name!r} cannot be encoded in tool calls") + names.append(_escape(name)) + if answer is not None: + others = _other_names( + opener[1], [f["name"] for f, e in zip(functions, enforced) if e] + ) + loose = _arguments_format(field, {}, strict=False) + if others is not None and loose is not None: + # Calls to any function whose schema is not enforced take any arguments. + names = [name for name, e in zip(names, enforced) if e] + [f"(?:{others})"] + arguments = [args for args, e in zip(arguments, enforced) if e] + [loose] + closers = _const_strings(closer_literals) if closer_literals else None + + def call(start: list[str] | None) -> dict[str, Any]: + return _one_of( + [ + _sequence( + {"type": "regex", "pattern": _call_opener(opener, name, start)}, + function_arguments, + ) + for name, function_arguments in zip(names, arguments) + ] + ) + + def output(start: list[str] | None) -> list[dict[str, Any]] | None: + elements = [call(start)] + if parallel_tool_calls and field.repeats: + repeated = _sequence(closers, call(None)) if closers else call(None) + elements.append({"type": "star", "content": repeated}) + if closers: + # EOS can terminate a call without writing its closer. + elements.append(_optional(closers)) + if answer is not None: + reply = _reply_format(answer, start, answer_format) + if reply is None: + return None + elements = [_one_of([_sequence(*elements), reply])] + return elements + + return _structural_tag(output, thinking, message_start, requires_reasoning) + + +def _structural_tag( + output: Callable[[list[str] | None], list[dict[str, Any]] | None], + thinking: ResponseTemplateField | None, + message_start: list[str] | None, + requires_reasoning: bool = False, +) -> dict[str, Any] | None: + """Frame output after reasoning, from the first token at the + `message_start`. The reasoning is optional unless `requires_reasoning` + makes the message open with non-empty reasoning.""" + rest = output(None) + if rest is None: + return None + think = None + if thinking is not None: + think = _thinking_format(thinking, message_start, requires_reasoning) + if think is None: + return None + if message_start is None: + if think is not None: + rest.insert(0, _optional(think)) + return {"type": "structural_tag", "format": _sequence(*rest)} + if think is not None and requires_reasoning: + return {"type": "structural_tag", "format": _sequence(think, *rest)} + start = output(message_start) + if start is None: + return None + format = _sequence(*start) + if think is not None: + format = _one_of([_sequence(think, *rest), format]) + return {"type": "structural_tag", "format": format} diff --git a/vllm/reasoning/__init__.py b/vllm/reasoning/__init__.py index 9e75e9b77ef4..796fc3ba3f37 100644 --- a/vllm/reasoning/__init__.py +++ b/vllm/reasoning/__init__.py @@ -84,6 +84,10 @@ "response_template_reasoning_parser", "ResponseTemplateReasoningParser", ), + "beam": ( + "beam_reasoning_parser", + "BeamReasoningParser", + ), "holo2": ( "deepseek_v3_reasoning_parser", "DeepSeekV3ReasoningWithThinkingParser", diff --git a/vllm/reasoning/beam_reasoning_parser.py b/vllm/reasoning/beam_reasoning_parser.py new file mode 100644 index 000000000000..b5d87170f0ba --- /dev/null +++ b/vllm/reasoning/beam_reasoning_parser.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.parser.beam import BeamReasoningParser + +__all__ = ["BeamReasoningParser"] diff --git a/vllm/renderers/online_renderer.py b/vllm/renderers/online_renderer.py index babd0ff6e5fc..65771211d0a5 100644 --- a/vllm/renderers/online_renderer.py +++ b/vllm/renderers/online_renderer.py @@ -60,6 +60,7 @@ has_non_text_content, ) from vllm.renderers.inputs.preprocess import ( + extract_prompt_components, parse_model_prompt, prompt_to_seq, ) @@ -844,13 +845,17 @@ async def preprocess_chat( f"but got {type(request).__name__}" ) raise NotImplementedError(msg) - request = parser( + request_parser = parser( tokenizer, request.tools, model_config=self.model_config, chat_template_kwargs=chat_params.chat_template_kwargs, - ).adjust_request( - request=request, ) + prompt_token_ids = extract_prompt_components( + self.model_config, engine_input + ).token_ids + if prompt_token_ids is not None: + request_parser.set_prompt_token_ids(prompt_token_ids) + request = request_parser.adjust_request(request=request) return conversation, [engine_input] diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py index 4b9cb33502e5..9e2e4a6d9fd8 100644 --- a/vllm/tool_parsers/__init__.py +++ b/vllm/tool_parsers/__init__.py @@ -90,6 +90,10 @@ "response_template_tool_parser", "ResponseTemplateToolParser", ), + "beam": ( + "beam_tool_parser", + "BeamToolParser", + ), "poolside_v1": ( "poolside_v1_tool_parser", "PoolsideV1ToolParser", diff --git a/vllm/tool_parsers/beam_tool_parser.py b/vllm/tool_parsers/beam_tool_parser.py new file mode 100644 index 000000000000..dbfd820688a2 --- /dev/null +++ b/vllm/tool_parsers/beam_tool_parser.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.parser.beam import BeamToolParser + +__all__ = ["BeamToolParser"] diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index c67d468a13b3..9867dbaba2da 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -118,6 +118,7 @@ def __getitem__(self, key): _CONFIG_REGISTRY: dict[str, type[PreTrainedConfig]] = LazyConfigDict( bagel="BagelConfig", bailing_moe_v3_vl="BailingMoeV3VLConfig", + beam="BeamConfig", chatglm="ChatGLMConfig", modernvbert="ColModernVBertConfig", colpali="ColPaliConfig", diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 68f82ead6b39..e674dd5a2e85 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -18,6 +18,7 @@ "BailingMoeV3TextConfig": "vllm.transformers_utils.configs.bailing_moe_v3_vl", "BailingMoeV3VisionConfig": "vllm.transformers_utils.configs.bailing_moe_v3_vl", "BailingMoeV3VLConfig": "vllm.transformers_utils.configs.bailing_moe_v3_vl", + "BeamConfig": "vllm.transformers_utils.configs.beam", "ChatGLMConfig": "vllm.transformers_utils.configs.chatglm", "ColModernVBertConfig": "vllm.transformers_utils.configs.colmodernvbert", "ColPaliConfig": "vllm.transformers_utils.configs.colpali", @@ -90,6 +91,7 @@ "BailingMoeV3TextConfig", "BailingMoeV3VisionConfig", "BailingMoeV3VLConfig", + "BeamConfig", "ChatGLMConfig", "ColModernVBertConfig", "ColPaliConfig", diff --git a/vllm/transformers_utils/configs/beam.py b/vllm/transformers_utils/configs/beam.py new file mode 100644 index 000000000000..7b7fee7c81d7 --- /dev/null +++ b/vllm/transformers_utils/configs/beam.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Beam model configuration.""" + +from typing import Any + +from transformers.configuration_utils import PreTrainedConfig + +# Normalize dtype aliases for ``STR_DTYPE_TO_TORCH_DTYPE``. +_DTYPE_ALIASES = {"fp32": "float32", "bf16": "bfloat16"} + + +def _canonical_dtype(name: str | None) -> str | None: + if name is None: + return None + canonical = _DTYPE_ALIASES.get(name, name) + if canonical not in ("float32", "bfloat16"): + raise ValueError( + f"Unsupported dtype {name!r}; expected one of " + f"{sorted(set(_DTYPE_ALIASES) | {'float32', 'bfloat16'})}." + ) + return canonical + + +class BeamConfig(PreTrainedConfig): + """Configuration for the Beam decoder-only MoE causal LM. + + Defaults match the Beam reference, including NoPE layers with + SSMax query scaling, hybrid sliding-window attention, per-head Q/K RMSNorm, + sandwich norm, depth-scaled residuals, sigmoid attention output gating, and + grouped sigmoid ``noaux_tc`` routing. + """ + + model_type = "beam" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size: int = 151936, + hidden_size: int = 2048, + intermediate_size: int = 6144, + num_hidden_layers: int = 24, + num_attention_heads: int = 32, + num_key_value_heads: int = 4, + head_dim: int | None = None, + hidden_act: str = "silu", + max_position_embeddings: int = 32768, + initializer_range: float = 0.02, + rms_norm_eps: float = 1e-6, + rope_theta: float = 10000.0, + rope_scaling: dict[str, Any] | None = None, + ssmax: bool = True, + ssmax_offset: int = 1, + ssmax_divide_by_log_offset: bool = False, + ssmax_scale_mode: str = "root10", + use_paramfree_qk_norm: bool = False, + use_cache: bool = True, + tie_word_embeddings: bool = False, + attention_bias: bool = False, + attention_output_gate: bool = True, + attention_output_gate_bias: bool = True, + attention_output_gate_floor: float = 0.0, + attention_output_gate_temperature: float = 1.0, + attention_dropout: float = 0.0, + sliding_window: int = 4096, + layer_types: list[str] | None = None, + nope_layer_ids: list[int] | None = None, + moe_intermediate_size: int = 768, + num_experts_per_tok: int = 8, + num_experts: int = 128, + norm_topk_prob: bool = True, + scoring_func: str = "sigmoid", + n_group: int | None = None, + topk_group: int | None = None, + routed_scaling_factor: float | None = None, + topk_method: str | None = "noaux_tc", + output_router_logits: bool = False, + router_aux_loss_coef: float = 0.001, + enable_sandwich_norm: bool = True, + self_attn_depth_scaled_residual: bool = True, + mlp_depth_scaled_residual: bool = True, + mlp_only_layers: list[int] | None = None, + mlp_layer_types: list[str] | None = None, + final_logit_softcapping: float | None = None, + moe_router_dtype: str | None = None, + pad_token_id: int | None = None, + bos_token_id: int | None = None, + eos_token_id: int | None = None, + **kwargs, + ): + if "num_local_experts" in kwargs: + raise ValueError("Beam requires num_experts instead of num_local_experts.") + if "architectures" not in kwargs: + kwargs["architectures"] = ["BeamForCausalLM"] + + no_rope_layers = kwargs.get("no_rope_layers") + if no_rope_layers is not None and ( + len(no_rope_layers) != num_hidden_layers + or any(flag not in (0, 1) for flag in no_rope_layers) + ): + raise ValueError("no_rope_layers must contain one 0/1 flag per layer.") + if layer_types is not None and ( + len(layer_types) != num_hidden_layers + or any( + t not in ("full_attention", "sliding_attention") for t in layer_types + ) + ): + raise ValueError( + "layer_types must contain one valid attention type per layer." + ) + if nope_layer_ids is None and no_rope_layers is not None: + # no_rope_layers[i] == 0 marks layer i as NoPE. + nope_layer_ids = [ + i + for i, uses_rope in enumerate(kwargs["no_rope_layers"]) + if not uses_rope + ] + elif nope_layer_ids is None and layer_types is not None: + nope_layer_ids = [ + i for i, kind in enumerate(layer_types) if kind == "full_attention" + ] + if nope_layer_ids is not None and ( + len(set(nope_layer_ids)) != len(nope_layer_ids) + or any(i < 0 or i >= num_hidden_layers for i in nope_layer_ids) + ): + raise ValueError("nope_layer_ids must contain unique valid layer indices.") + # These dtype options are not supported by Beam. + kwargs.pop("residual_in_fp32", None) + kwargs.pop("lm_head_dtype", None) + if moe_router_dtype is None: + moe_router_dtype = kwargs.get("router_dtype") + moe_router_dtype = _canonical_dtype(moe_router_dtype) + rope_parameters = kwargs.get("rope_parameters") + if rope_parameters and set(rope_parameters).issubset( + {"sliding_attention", "full_attention"} + ): + # NoPE layers have no rotary parameters. Canonicalize before HF's + # generic RoPE validation, which cannot process null entries. + rotary_parameters = [p for p in rope_parameters.values() if p is not None] + if rotary_parameters and any( + p != rotary_parameters[0] for p in rotary_parameters[1:] + ): + raise ValueError( + "Beam requires identical RoPE parameters across layers." + ) + rope_parameters = dict(rotary_parameters[0]) if rotary_parameters else {} + rope_parameters.setdefault("rope_type", "default") + rope_parameters.setdefault("rope_theta", rope_theta) + kwargs["rope_parameters"] = rope_parameters + if rope_parameters: + rope_theta = rope_parameters.get("rope_theta", rope_theta) + if layer_types is None: + # Explicit layer types keep full-attention layers from inheriting + # the global sliding window. + nope_ids = set(nope_layer_ids or ()) + layer_types = [ + "full_attention" + if not sliding_window or layer_id in nope_ids + else "sliding_attention" + for layer_id in range(num_hidden_layers) + ] + + # Newer checkpoints carry per-layer feed-forward types; older ones list + # the dense layers. Accept either, derive the other, reject disagreement. + if mlp_layer_types is not None: + if len(mlp_layer_types) != num_hidden_layers or any( + t not in ("dense", "sparse") for t in mlp_layer_types + ): + raise ValueError( + "mlp_layer_types must contain one 'dense' or 'sparse' entry " + "per layer." + ) + dense_layers = [i for i, t in enumerate(mlp_layer_types) if t == "dense"] + if mlp_only_layers is not None and sorted(mlp_only_layers) != dense_layers: + raise ValueError( + "mlp_only_layers disagrees with mlp_layer_types: " + f"{sorted(mlp_only_layers)} vs {dense_layers}." + ) + mlp_only_layers = dense_layers + else: + mlp_only_layers = [] if mlp_only_layers is None else list(mlp_only_layers) + if len(set(mlp_only_layers)) != len(mlp_only_layers) or any( + i < 0 or i >= num_hidden_layers for i in mlp_only_layers + ): + raise ValueError( + "mlp_only_layers must contain unique valid layer indices." + ) + mlp_layer_types = [ + "dense" if i in set(mlp_only_layers) else "sparse" + for i in range(num_hidden_layers) + ] + + 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.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.ssmax = ssmax + self.ssmax_offset = ssmax_offset + self.ssmax_divide_by_log_offset = ssmax_divide_by_log_offset + self.ssmax_scale_mode = ssmax_scale_mode + self.use_paramfree_qk_norm = use_paramfree_qk_norm + self.use_cache = use_cache + self.attention_bias = attention_bias + self.attention_output_gate = attention_output_gate + self.attention_output_gate_bias = attention_output_gate_bias + self.attention_output_gate_floor = attention_output_gate_floor + self.attention_output_gate_temperature = attention_output_gate_temperature + self.attention_dropout = attention_dropout + self.sliding_window = sliding_window + self.layer_types = layer_types + self.nope_layer_ids = [] if nope_layer_ids is None else nope_layer_ids + self.moe_intermediate_size = moe_intermediate_size + self.num_experts_per_tok = num_experts_per_tok + self.num_experts = num_experts + self.norm_topk_prob = norm_topk_prob + self.scoring_func = scoring_func + self.n_group = n_group + self.topk_group = topk_group + self.routed_scaling_factor = routed_scaling_factor + self.topk_method = topk_method + self.output_router_logits = output_router_logits + self.router_aux_loss_coef = router_aux_loss_coef + self.enable_sandwich_norm = enable_sandwich_norm + self.self_attn_depth_scaled_residual = self_attn_depth_scaled_residual + self.mlp_depth_scaled_residual = mlp_depth_scaled_residual + self.mlp_only_layers = mlp_only_layers + self.mlp_layer_types = mlp_layer_types + self.final_logit_softcapping = final_logit_softcapping + self.moe_router_dtype = moe_router_dtype + + kwargs.update( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + ) + super().__init__(**kwargs) + + +__all__ = ["BeamConfig"] diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index fc028941cf0d..e1f97d2affcf 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/vllm/v1/worker/gpu/spec_decode/__init__.py @@ -31,7 +31,9 @@ def init_speculator( ) return LiLiCorrSpeculator(vllm_config, device) - if "DFlash2DraftModel" in speculative_config.draft_model_config.architectures: + if {"DFlash2DraftModel", "DFlash2BeamDraftModel"}.intersection( + speculative_config.draft_model_config.architectures + ): from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import ( DFlash2Speculator, )