@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6570c9800c

[Model] Add Inkling model support [1/N] (#48799)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Signed-off-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <mozf@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Woosuk Kwon committed Jul 15, 2026 at 23:40 UTC 6570c9800cda6cc510525f485141d527f105b46c
95 files changed +12137 -62
.gitignore
+3
@@ -18,6 +18,9 @@ vllm/third_party/deep_gemm/
18 # fmha_sm100 vendored package built from source
19 vllm/third_party/fmha_sm100/
20
21 +# tml-fa4 vendored package built from source
22 +vllm/third_party/tml_fa4/
23 +
24 # triton jit
25 .triton
26
CMakeLists.txt
+1
@@ -1408,6 +1408,7 @@ if (VLLM_GPU_LANG STREQUAL "CUDA")
1408 include(cmake/external_projects/fmha_sm100.cmake)
1409 include(cmake/external_projects/flashmla.cmake)
1410 include(cmake/external_projects/qutlass.cmake)
1411 + include(cmake/external_projects/tml_fa4.cmake)
1412
1413 # vllm-flash-attn should be last as it overwrites some CMake functions
1414 include(cmake/external_projects/vllm_flash_attn.cmake)
cmake/external_projects/tml_fa4.cmake new
+50
@@ -0,0 +1,50 @@
1 +include(FetchContent)
2 +
3 +if(DEFINED ENV{TML_FA4_SRC_DIR})
4 + set(TML_FA4_SRC_DIR $ENV{TML_FA4_SRC_DIR})
5 +endif()
6 +
7 +if(TML_FA4_SRC_DIR)
8 + FetchContent_Declare(
9 + tml_fa4
10 + SOURCE_DIR ${TML_FA4_SRC_DIR}
11 + CONFIGURE_COMMAND ""
12 + BUILD_COMMAND "")
13 +else()
14 + FetchContent_Declare(
15 + tml_fa4
16 + GIT_REPOSITORY https://github.com/vllm-project/tml-fa4.git
17 + GIT_TAG 13374f0c855acc1add1bf30444bd67aebbc24a8e
18 + GIT_PROGRESS TRUE
19 + CONFIGURE_COMMAND ""
20 + BUILD_COMMAND "")
21 +endif()
22 +
23 +FetchContent_GetProperties(tml_fa4)
24 +if(NOT tml_fa4_POPULATED)
25 + FetchContent_Populate(tml_fa4)
26 +endif()
27 +message(STATUS "tml-fa4 is available at ${tml_fa4_SOURCE_DIR}")
28 +
29 +add_custom_target(tml_fa4)
30 +
31 +# Install into a private namespace so this implementation cannot shadow the
32 +# flash_attn package used by vLLM's standard attention backends.
33 +install(CODE "
34 + file(GLOB_RECURSE TML_FA4_PY_FILES
35 + \"${tml_fa4_SOURCE_DIR}/flash_attn/cute/*.py\")
36 + foreach(SRC_FILE \${TML_FA4_PY_FILES})
37 + file(RELATIVE_PATH REL_PATH
38 + \"${tml_fa4_SOURCE_DIR}/flash_attn/cute\" \${SRC_FILE})
39 + set(DST_FILE
40 + \"\${CMAKE_INSTALL_PREFIX}/vllm/third_party/tml_fa4/\${REL_PATH}\")
41 + get_filename_component(DST_DIR \${DST_FILE} DIRECTORY)
42 + file(MAKE_DIRECTORY \${DST_DIR})
43 + file(READ \${SRC_FILE} FILE_CONTENTS)
44 + string(REPLACE
45 + \"flash_attn.cute\"
46 + \"vllm.third_party.tml_fa4\"
47 + FILE_CONTENTS \"\${FILE_CONTENTS}\")
48 + file(WRITE \${DST_FILE} \"\${FILE_CONTENTS}\")
49 + endforeach()
50 +" COMPONENT tml_fa4)
requirements/cuda.txt
+1 -1
@@ -26,7 +26,7 @@ fastsafetensors >= 0.3.2
26
27 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
28 nvidia-cutlass-dsl[cu13]==4.5.2
29 -quack-kernels>=0.3.3
29 +quack-kernels>=0.4.0 # Required for tml-fa4
30
31 # Tokenspeed_MLA for faster mla with spec decode
32 tokenspeed-mla==0.1.8; platform_system == "Linux"
rust/Cargo.lock
+2 -1
@@ -2220,7 +2220,7 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092"
2220 [[package]]
2221 name = "llm-multimodal"
2222 version = "1.7.1"
2223 -source = "git+https://github.com/smg-project/llm-multimodal?rev=7df38e53f99aefaebe86934f010aa8084ec99b2f#7df38e53f99aefaebe86934f010aa8084ec99b2f"
2223 +source = "git+https://github.com/smg-project/llm-multimodal?rev=5390032d6dc8a3e6fdc83acd320260367eb4b9b5#5390032d6dc8a3e6fdc83acd320260367eb4b9b5"
2224 dependencies = [
2225 "anyhow",
2226 "base64 0.22.1",
@@ -2235,6 +2235,7 @@ dependencies = [
2235 "once_cell",
2236 "pkg-config",
2237 "rayon",
2238 + "realfft",
2239 "reqwest 0.13.4",
2240 "rustfft",
2241 "serde",
rust/Cargo.toml
+1 -1
@@ -53,7 +53,7 @@ hyper-util = { version = "0.1.20", features = [
53 indexmap = "2.13.0"
54 itertools = "0.14.0"
55 libc = "0.2.177"
56 -llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "7df38e53f99aefaebe86934f010aa8084ec99b2f", default-features = false, features = ["native-tls"] }
56 +llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "5390032d6dc8a3e6fdc83acd320260367eb4b9b5", default-features = false, features = ["native-tls"] }
57 mimalloc = "0.1.52"
58 minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] }
59 minijinja-contrib = { version = "2.0", features = ["pycompat"] }
rust/src/chat/src/backend/hf.rs
+2
@@ -20,6 +20,7 @@ use crate::output::{
20 use crate::renderer::hf::{HfChatRenderer, MultimodalRenderInfo};
21 use crate::renderer::{
22 DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, HarmonyChatRenderer,
23 + InklingChatRenderer,
24 };
25 use crate::request::ChatRequest;
26 use crate::{DynChatOutputProcessor, RendererSelection};
@@ -71,6 +72,7 @@ impl HfChatBackend {
72 RendererSelection::DeepSeekV32 => Arc::new(DeepSeekV32ChatRenderer::new()),
73 RendererSelection::DeepSeekV4 => Arc::new(DeepSeekV4ChatRenderer::new()),
74 RendererSelection::Harmony => Arc::new(HarmonyChatRenderer::new()?),
75 + RendererSelection::Inkling => Arc::new(InklingChatRenderer::new(tokenizer.clone())?),
76 };
77
78 info!(
rust/src/chat/src/lib.rs
+3 -3
@@ -33,7 +33,7 @@ pub use parser::tool::{ToolParser, ToolParserError, ToolParserFactory};
33 pub use renderer::hf::ChatTemplateContentFormatOption;
34 pub use renderer::{
35 ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer,
36 - HarmonyChatRenderer, RenderedPrompt, RendererSelection,
36 + HarmonyChatRenderer, InklingChatRenderer, RenderedPrompt, RendererSelection,
37 };
38 pub use request::{
39 ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool,
@@ -294,7 +294,7 @@ mod tests {
294 )
295 .unwrap_err();
296
297 - expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
297 + expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, inkling, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
298 }
299
300 #[test]
@@ -305,6 +305,6 @@ mod tests {
305 )
306 .unwrap_err();
307
308 - expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, kimi, kimi_k2, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
308 + expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, inkling, kimi, kimi_k2, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
309 }
310 }
rust/src/chat/src/multimodal/audio.rs
+112
@@ -79,6 +79,38 @@ mod tests {
79 use super::*;
80
81 const AUDIO_PAD_ID: u32 = 151_676;
82 + const INKLING_AUDIO_MARKER_ID: u32 = 200_020;
83 + const INKLING_AUDIO_EMBED_ID: i32 = 200_053;
84 +
85 + fn inkling_info(decoder_dmodel: serde_json::Value) -> MultimodalModelInfo {
86 + let config = serde_json::json!({
87 + "model_type": "inkling_mm_model",
88 + "audio_config": {
89 + "decoder_dmodel": decoder_dmodel,
90 + "n_mel_bins": 80,
91 + "mel_vocab_size": 16,
92 + "dmel_min_value": -7.0,
93 + "dmel_max_value": 2.0
94 + }
95 + });
96 + let tokenizer = TestTokenizer::new()
97 + .with_regular_token("<|content_image|>", 200_005)
98 + .with_regular_token("<|content_audio_input|>", INKLING_AUDIO_MARKER_ID);
99 + let context = MultimodalModelContext {
100 + model_id: "inkling-test".to_string(),
101 + model_type: Some("inkling_mm_model".to_string()),
102 + config,
103 + tokenizer: TokenizerResolver(Arc::new(tokenizer)),
104 + };
105 +
106 + MultimodalModelInfo::from_loaded(
107 + context,
108 + PreProcessorConfig::default(),
109 + PreProcessorConfig::default(),
110 + )
111 + .unwrap()
112 + .expect("Inkling multimodal support")
113 + }
114
115 fn qwen3_asr_info() -> MultimodalModelInfo {
116 let context = MultimodalModelContext {
@@ -120,6 +152,40 @@ mod tests {
152 bytes
153 }
154
155 + #[test]
156 + fn resolves_inkling_audio_from_model_spec() {
157 + let info = inkling_info(serde_json::json!(1024));
158 + let support = info.audio.as_ref().expect("audio support");
159 +
160 + assert_eq!(
161 + info.placeholder_token(Modality::Audio),
162 + Some("<|content_audio_input|>")
163 + );
164 + assert_eq!(support.placeholder.marker_token_id, INKLING_AUDIO_MARKER_ID);
165 + assert_eq!(
166 + support.placeholder.embed_token_id,
167 + INKLING_AUDIO_EMBED_ID as u32
168 + );
169 + assert_eq!(support.spec.primary_key(), AUDIO_PRIMARY_KEY);
170 + assert!(matches!(
171 + &support.spec.field_layouts.encoder_input,
172 + llm_multimodal::FieldLayout::Flat { sizes_key }
173 + if sizes_key == "num_audio_tokens"
174 + ));
175 + assert!(matches!(
176 + support.spec.field_layouts.model_specific.get("num_audio_tokens"),
177 + Some(llm_multimodal::FieldLayout::Batched)
178 + ));
179 + }
180 +
181 + #[test]
182 + fn inkling_decoder_config_gates_audio_capability() {
183 + let info = inkling_info(serde_json::Value::Null);
184 +
185 + assert!(info.audio.is_none());
186 + assert_eq!(info.placeholder_token(Modality::Audio), None);
187 + }
188 +
189 #[test]
190 fn resolves_qwen_audio_from_model_spec() {
191 let info = qwen3_asr_info();
@@ -210,4 +276,50 @@ mod tests {
276 if tensor.dtype == "int64" && tensor.shape.is_empty()
277 ));
278 }
279 +
280 + #[tokio::test]
281 + async fn inkling_tracker_and_processor_use_standard_audio_key() {
282 + let info = inkling_info(serde_json::json!(1024));
283 + let wav = wav_i16_mono(16_000, &[0; 1_600]);
284 + let expected_hash = llm_multimodal::hasher::hash_audio(&wav);
285 + let fetched = info
286 + .fetch_media(vec![MediaContentPart::AudioData {
287 + data: wav,
288 + mime_type: Some("audio/wav".to_string()),
289 + uuid: Some("audio-1".to_string()),
290 + }])
291 + .await
292 + .unwrap();
293 +
294 + let prepared = info.prepare_audios(fetched.audios, fetched.audio_uuids).await.unwrap();
295 +
296 + assert_eq!(prepared.replacements.len(), 1);
297 + assert_eq!(
298 + prepared.replacements[0].tokens[0],
299 + INKLING_AUDIO_MARKER_ID as i32
300 + );
301 + assert!(
302 + prepared.replacements[0].tokens[1..]
303 + .iter()
304 + .all(|token| *token == INKLING_AUDIO_EMBED_ID)
305 + );
306 + let item = &prepared.items[0];
307 + assert_eq!(item.hash, expected_hash);
308 + assert_eq!(item.uuid.as_deref(), Some("audio-1"));
309 +
310 + let features = &item.data[AUDIO_PRIMARY_KEY];
311 + assert!(matches!(&features.field, MmField::Flat(_)));
312 + assert!(matches!(
313 + features.data.as_ref(),
314 + Some(MmKwargValue::Tensor(tensor))
315 + if tensor.dtype == "float32" && tensor.shape.get(1) == Some(&80)
316 + ));
317 + let count = &item.data["num_audio_tokens"];
318 + assert!(matches!(&count.field, MmField::Batched(_)));
319 + assert!(matches!(
320 + count.data.as_ref(),
321 + Some(MmKwargValue::Tensor(tensor))
322 + if tensor.dtype == "int64" && tensor.shape.is_empty()
323 + ));
324 + }
325 }
rust/src/chat/src/parser/reasoning/mod.rs
+2
@@ -23,6 +23,7 @@ pub mod names {
23 pub const DEEPSEEK_V3: &str = "deepseek_v3";
24 pub const DEEPSEEK_V4: &str = "deepseek_v4";
25 pub const GEMMA4: &str = "gemma4";
26 + pub const INKLING: &str = "inkling";
27 pub const GLM45: &str = "glm45";
28 pub const KIMI: &str = "kimi";
29 pub const KIMI_K2: &str = "kimi_k2";
@@ -63,6 +64,7 @@ impl ReasoningParserFactory {
64 .register_parser::<DeepSeekV3ReasoningParser>(names::DEEPSEEK_V3)
65 .register_parser::<DeepSeekV4ReasoningParser>(names::DEEPSEEK_V4)
66 .register_unified_dummy(names::GEMMA4)
67 + .register_unified_dummy(names::INKLING)
68 .register_parser::<Glm45ReasoningParser>(names::GLM45)
69 .register_parser::<KimiReasoningParser>(names::KIMI)
70 .register_parser::<KimiK2ReasoningParser>(names::KIMI_K2)
rust/src/chat/src/parser/tool/mod.rs
+2
@@ -25,6 +25,7 @@ pub mod names {
25 pub const GLM45: &str = "glm45";
26 pub const GLM47: &str = "glm47";
27 pub const GEMMA4: &str = "gemma4";
28 + pub const INKLING: &str = "inkling";
29 pub const GRANITE4: &str = "granite4";
30 pub const HERMES: &str = "hermes";
31 pub const HY_V3: &str = "hy_v3";
@@ -70,6 +71,7 @@ impl ToolParserFactory {
71 .register_parser::<Glm45MoeToolParser>(names::GLM45)
72 .register_parser::<Glm47MoeToolParser>(names::GLM47)
73 .register_unified_dummy(names::GEMMA4)
74 + .register_unified_dummy(names::INKLING)
75 .register_parser::<Granite4ToolParser>(names::GRANITE4)
76 .register_parser::<HermesToolParser>(names::HERMES)
77 .register_parser::<HyV3ToolParser>(names::HY_V3)
rust/src/chat/src/parser/unified.rs
+24 -2
@@ -5,7 +5,7 @@
5
6 use std::sync::LazyLock;
7
8 -pub use vllm_parser::unified::{Gemma4UnifiedParser, UnifiedParser};
8 +pub use vllm_parser::unified::{Gemma4UnifiedParser, InklingUnifiedParser, UnifiedParser};
9 use vllm_tokenizer::DynTokenizer;
10
11 use crate::parser::ParserFactory;
@@ -14,6 +14,7 @@ use crate::request::ChatTool;
14 /// Canonical public names for registered unified parsers.
15 pub mod names {
16 pub const GEMMA4: &str = "gemma4";
17 + pub const INKLING: &str = "inkling";
18 }
19
20 /// Constructor signature for one registered unified parser implementation.
@@ -37,10 +38,12 @@ impl UnifiedParserFactory {
38 let mut factory = Self::default();
39
40 factory.register_parser::<Gemma4UnifiedParser>(names::GEMMA4);
41 + factory.register_parser::<InklingUnifiedParser>(names::INKLING);
42
43 factory
44 .register_pattern("gemma-4", names::GEMMA4)
43 - .register_pattern("gemma4", names::GEMMA4);
45 + .register_pattern("gemma4", names::GEMMA4)
46 + .register_pattern("inkling", names::INKLING);
47
48 factory
49 }
@@ -88,6 +91,13 @@ mod tests {
91 .with_regular_token("<channel|>", 257)
92 }
93
94 + fn inkling_tokenizer() -> TestTokenizer {
95 + TestTokenizer::new()
96 + .with_regular_token("<|message_model|>", 200001)
97 + .with_regular_token("<|content_text|>", 200004)
98 + .with_regular_token("<|content_thinking|>", 200008)
99 + }
100 +
101 #[test]
102 fn factory_registers_gemma4() {
103 let factory = UnifiedParserFactory::new();
@@ -99,4 +109,16 @@ mod tests {
109 );
110 factory.create(names::GEMMA4, &[], Arc::new(tokenizer())).unwrap();
111 }
112 +
113 + #[test]
114 + fn factory_registers_inkling() {
115 + let factory = UnifiedParserFactory::new();
116 +
117 + assert!(factory.contains(names::INKLING));
118 + assert_eq!(
119 + factory.resolve_name_for_model("thinkingmachines/Inkling"),
120 + Some(names::INKLING)
121 + );
122 + factory.create(names::INKLING, &[], Arc::new(inkling_tokenizer())).unwrap();
123 + }
124 }
rust/src/chat/src/renderer/inkling/fixtures/text_audio_input.json new
+21
@@ -0,0 +1,21 @@
1 +{
2 + "messages": [
3 + {
4 + "role": "user",
5 + "content": [
6 + {
7 + "type": "text",
8 + "text": "transcribe"
9 + },
10 + {
11 + "type": "input_audio",
12 + "data": ""
13 + },
14 + {
15 + "type": "audio_url",
16 + "audio_url": "data:audio/wav;base64,"
17 + }
18 + ]
19 + }
20 + ]
21 +}
rust/src/chat/src/renderer/inkling/fixtures/text_audio_output.txt new
+1
@@ -0,0 +1 @@
1 +<|message_system|><|content_text|>Thinking effort level: 0.9<|end_message|><|message_user|><|content_text|>transcribe<|end_message|><|message_user|><|content_audio_input|><|audio_end|><|end_message|><|message_user|><|content_audio_input|><|audio_end|><|end_message|><|message_model|>
rust/src/chat/src/renderer/inkling/fixtures/text_image_input.json new
+17
@@ -0,0 +1,17 @@
1 +{
2 + "messages": [
3 + {
4 + "role": "user",
5 + "content": [
6 + {
7 + "type": "text",
8 + "text": "look"
9 + },
10 + {
11 + "type": "image_url",
12 + "image_url": "data:image/png;base64,"
13 + }
14 + ]
15 + }
16 + ]
17 +}
rust/src/chat/src/renderer/inkling/fixtures/text_image_output.txt new
+1
@@ -0,0 +1 @@
1 +<|message_system|><|content_text|>Thinking effort level: 0.9<|end_message|><|message_user|><|content_text|>look<|end_message|><|message_user|><|content_image|><|end_message|><|message_model|>
rust/src/chat/src/renderer/inkling/fixtures/tool_declare_input.json new
+46
@@ -0,0 +1,46 @@
1 +{
2 + "tools": [
3 + {
4 + "type": "function",
5 + "function": {
6 + "name": "get_weather",
7 + "description": "Get weather information",
8 + "parameters": {
9 + "type": "object",
10 + "required": [
11 + "city"
12 + ],
13 + "properties": {
14 + "city": {
15 + "type": "string"
16 + }
17 + }
18 + }
19 + }
20 + }
21 + ],
22 + "messages": [
23 + {
24 + "role": "developer",
25 + "content": "rules",
26 + "tools": [
27 + {
28 + "type": "function",
29 + "function": {
30 + "name": "local_tool",
31 + "parameters": {
32 + "z": 1,
33 + "a": {
34 + "b": 2
35 + }
36 + }
37 + }
38 + }
39 + ]
40 + },
41 + {
42 + "role": "user",
43 + "content": "hi"
44 + }
45 + ]
46 +}
rust/src/chat/src/renderer/inkling/fixtures/tool_declare_output.txt new
+1
@@ -0,0 +1 @@
1 +<|message_system|>tool_declare<|content_xml|>[{"description":"Get weather information","name":"get_weather","parameters":{"properties":{"city":{"type":"string"}},"required":["city"],"type":"object"},"type":"function"},{"description":"","name":"local_tool","parameters":{"a":{"b":2},"z":1},"type":"function"}]<|end_message|><|message_system|><|content_text|>rules<|end_message|><|message_system|><|content_text|>Thinking effort level: 0.9<|end_message|><|message_user|><|content_text|>hi<|end_message|><|message_model|>
rust/src/chat/src/renderer/inkling/fixtures/tool_round_trip_input.json new
+24
@@ -0,0 +1,24 @@
1 +{
2 + "add_generation_prompt": false,
3 + "messages": [
4 + {
5 + "role": "assistant",
6 + "reasoning_content": "think",
7 + "content": "answer",
8 + "tool_calls": [
9 + {
10 + "id": "call_1",
11 + "function": {
12 + "name": "get_weather",
13 + "arguments": "{\"city\":\"SF\"}"
14 + }
15 + }
16 + ]
17 + },
18 + {
19 + "role": "tool",
20 + "tool_call_id": "call_1",
21 + "content": "sunny"
22 + }
23 + ]
24 +}
rust/src/chat/src/renderer/inkling/fixtures/tool_round_trip_output.txt new
+1
@@ -0,0 +1 @@
1 +<|message_system|><|content_text|>Thinking effort level: 0.9<|end_message|><|message_model|><|content_thinking|>think<|end_message|><|message_model|><|content_text|>answer<|end_message|><|message_model|>get_weather<|content_invoke_tool_json|>{"name":"get_weather","args":{"city":"SF"}}<|end_message|><|content_model_end_sampling|><|message_tool|>get_weather<|content_text|>sunny<|end_message|>
rust/src/chat/src/renderer/inkling/mod.rs new
+442
@@ -0,0 +1,442 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::collections::HashMap;
5 +
6 +use serde_json::{Map, Value, json};
7 +use thiserror_ext::AsReport as _;
8 +use vllm_text::Prompt;
9 +use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
10 +
11 +use super::{ChatRenderer, RenderedPrompt, request_template_kwargs};
12 +use crate::error::{Error, Result};
13 +use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool};
14 +use crate::{AssistantContentBlock, AssistantToolCall};
15 +
16 +const MESSAGE_USER: &str = "<|message_user|>";
17 +const MESSAGE_MODEL: &str = "<|message_model|>";
18 +const MESSAGE_SYSTEM: &str = "<|message_system|>";
19 +const MESSAGE_TOOL: &str = "<|message_tool|>";
20 +const CONTENT_TEXT: &str = "<|content_text|>";
21 +const CONTENT_IMAGE: &str = "<|content_image|>";
22 +const CONTENT_MODEL_END_SAMPLING: &str = "<|content_model_end_sampling|>";
23 +const CONTENT_AUDIO_INPUT: &str = "<|content_audio_input|>";
24 +const CONTENT_THINKING: &str = "<|content_thinking|>";
25 +// Inkling renderer semantics use this slot for structured payloads such as tool
26 +// declarations.
27 +const CONTENT_XML: &str = "<|content_xml|>";
28 +const CONTENT_INVOKE_TOOL_JSON: &str = "<|content_invoke_tool_json|>";
29 +const END_MESSAGE: &str = "<|end_message|>";
30 +const AUDIO_END: &str = "<|audio_end|>";
31 +const MAX_REASONING_EFFORT: f64 = 0.99;
32 +
33 +/// Native Inkling renderer that emits token IDs directly.
34 +#[derive(Clone)]
35 +pub struct InklingChatRenderer {
36 + tokenizer: DynTokenizer,
37 + special: InklingSpecialTokenIds,
38 +}
39 +
40 +#[derive(Debug, Clone, Copy)]
41 +struct InklingSpecialTokenIds {
42 + message_user: u32,
43 + message_model: u32,
44 + message_system: u32,
45 + message_tool: u32,
46 + content_text: u32,
47 + content_image: u32,
48 + content_model_end_sampling: u32,
49 + content_audio_input: u32,
50 + content_thinking: u32,
51 + content_xml: u32,
52 + content_invoke_tool_json: u32,
53 + end_message: u32,
54 + audio_end: u32,
55 +}
56 +
57 +impl InklingSpecialTokenIds {
58 + fn resolve(tokenizer: &dyn Tokenizer) -> Result<Self> {
59 + Ok(Self {
60 + message_user: resolve_special_token(tokenizer, MESSAGE_USER)?,
61 + message_model: resolve_special_token(tokenizer, MESSAGE_MODEL)?,
62 + message_system: resolve_special_token(tokenizer, MESSAGE_SYSTEM)?,
63 + message_tool: resolve_special_token(tokenizer, MESSAGE_TOOL)?,
64 + content_text: resolve_special_token(tokenizer, CONTENT_TEXT)?,
65 + content_image: resolve_special_token(tokenizer, CONTENT_IMAGE)?,
66 + content_model_end_sampling: resolve_special_token(
67 + tokenizer,
68 + CONTENT_MODEL_END_SAMPLING,
69 + )?,
70 + content_audio_input: resolve_special_token(tokenizer, CONTENT_AUDIO_INPUT)?,
71 + content_thinking: resolve_special_token(tokenizer, CONTENT_THINKING)?,
72 + content_xml: resolve_special_token(tokenizer, CONTENT_XML)?,
73 + content_invoke_tool_json: resolve_special_token(tokenizer, CONTENT_INVOKE_TOOL_JSON)?,
74 + end_message: resolve_special_token(tokenizer, END_MESSAGE)?,
75 + audio_end: resolve_special_token(tokenizer, AUDIO_END)?,
76 + })
77 + }
78 +}
79 +
80 +impl InklingChatRenderer {
81 + pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
82 + let special = InklingSpecialTokenIds::resolve(tokenizer.as_ref())?;
83 + Ok(Self { tokenizer, special })
84 + }
85 +
86 + fn write_text_tokens(&self, out: &mut Vec<u32>, text: &str) -> Result<()> {
87 + out.extend(self.tokenizer.encode(text, false)?);
88 + Ok(())
89 + }
90 +
91 + fn write_message_start(
92 + &self,
93 + out: &mut Vec<u32>,
94 + role_token_id: u32,
95 + author_name: Option<&str>,
96 + ) -> Result<()> {
97 + out.push(role_token_id);
98 + if let Some(author_name) = author_name
99 + && !author_name.is_empty()
100 + {
101 + self.write_text_tokens(out, author_name)?;
102 + }
103 + Ok(())
104 + }
105 +
106 + fn write_text_block(
107 + &self,
108 + out: &mut Vec<u32>,
109 + role_token_id: u32,
110 + author_name: Option<&str>,
111 + text: &str,
112 + ) -> Result<()> {
113 + self.write_message_start(out, role_token_id, author_name)?;
114 + out.push(self.special.content_text);
115 + self.write_text_tokens(out, text)?;
116 + out.push(self.special.end_message);
117 + Ok(())
118 + }
119 +
120 + /// Write one image block holding only the `<|content_image|>` marker.
121 + /// Multimodal preprocessing later expands the marker into per-patch
122 + /// image placeholder tokens once the patch count is known, mirroring the
123 + /// Python `InklingMultiModalProcessor` marker-anchored prompt updates.
124 + fn write_image_block(&self, out: &mut Vec<u32>, role_token_id: u32) {
125 + out.push(role_token_id);
126 + out.push(self.special.content_image);
127 + out.push(self.special.end_message);
128 + }
129 +
130 + /// Write one audio block holding the `<|content_audio_input|>` marker
131 + /// followed by the `<|audio_end|>` terminator. Multimodal preprocessing
132 + /// later expands the marker into per-frame audio placeholder tokens
133 + /// (landing before `<|audio_end|>`) once the clip length is known.
134 + fn write_audio_block(&self, out: &mut Vec<u32>, role_token_id: u32) {
135 + out.push(role_token_id);
136 + out.push(self.special.content_audio_input);
137 + out.push(self.special.audio_end);
138 + out.push(self.special.end_message);
139 + }
140 +
141 + fn write_reasoning_block(&self, out: &mut Vec<u32>, text: &str) -> Result<()> {
142 + if text.is_empty() {
143 + return Ok(());
144 + }
145 + out.push(self.special.message_model);
146 + out.push(self.special.content_thinking);
147 + self.write_text_tokens(out, text)?;
148 + out.push(self.special.end_message);
149 + Ok(())
150 + }
151 +
152 + fn write_reasoning_effort(&self, out: &mut Vec<u32>, effort: f64) -> Result<()> {
153 + if !(0.0..=MAX_REASONING_EFFORT).contains(&effort) {
154 + return Err(Error::ChatTemplate(format!(
155 + "Inkling reasoning_effort must be in [0.0, 0.99], got {effort}"
156 + )));
157 + }
158 + let formatted = format!("{effort:.2}");
159 + let effort = formatted.trim_end_matches('0').trim_end_matches('.');
160 + let effort = if matches!(effort, "0" | "-0") {
161 + "0.0"
162 + } else {
163 + effort
164 + };
165 + self.write_text_block(
166 + out,
167 + self.special.message_system,
168 + None,
169 + &format!("Thinking effort level: {effort}"),
170 + )
171 + }
172 +
173 + fn write_tool_declarations(&self, out: &mut Vec<u32>, tools: &[&ChatTool]) -> Result<()> {
174 + if tools.is_empty() {
175 + return Ok(());
176 + }
177 +
178 + let mut specs = Vec::with_capacity(tools.len());
179 + for tool in tools {
180 + specs.push(json!({
181 + "description": tool.description.as_deref().unwrap_or(""),
182 + "name": tool.name,
183 + "parameters": sort_json(&tool.parameters),
184 + "type": "function",
185 + }));
186 + }
187 + let payload = compact_json(&sort_json(&Value::Array(specs)))?;
188 +
189 + self.write_message_start(out, self.special.message_system, Some("tool_declare"))?;
190 + out.push(self.special.content_xml);
191 + self.write_text_tokens(out, &payload)?;
192 + out.push(self.special.end_message);
193 + Ok(())
194 + }
195 +
196 + fn write_chat_content(
197 + &self,
198 + out: &mut Vec<u32>,
199 + role_token_id: u32,
200 + content: &ChatContent,
201 + ) -> Result<()> {
202 + match content {
203 + ChatContent::Text(text) => {
204 + if !text.is_empty() {
205 + self.write_text_block(out, role_token_id, None, text)?;
206 + }
207 + }
208 + ChatContent::Parts(parts) => {
209 + for part in parts {
210 + match part {
211 + ChatContentPart::Text { text } => {
212 + self.write_text_block(out, role_token_id, None, text)?;
213 + }
214 + ChatContentPart::ImageUrl { .. } => {
215 + self.write_image_block(out, role_token_id);
216 + }
217 + ChatContentPart::InputAudio { .. } | ChatContentPart::AudioUrl { .. } => {
218 + self.write_audio_block(out, role_token_id);
219 + }
220 + // Inkling has no video modality.
221 + ChatContentPart::VideoUrl { .. } => {
222 + return Err(Error::UnsupportedMultimodalContent("video_url"));
223 + }
224 + }
225 + }
226 + }
227 + }
228 + Ok(())
229 + }
230 +
231 + fn write_assistant_tool_call(
232 + &self,
233 + out: &mut Vec<u32>,
234 + tool_call: &AssistantToolCall,
235 + ) -> Result<()> {
236 + let payload = tool_call_json(tool_call)?;
237 +
238 + self.write_message_start(out, self.special.message_model, Some(&tool_call.name))?;
239 + out.push(self.special.content_invoke_tool_json);
240 + self.write_text_tokens(out, &payload)?;
241 + out.push(self.special.end_message);
242 + Ok(())
243 + }
244 +
245 + fn write_assistant_content(
246 + &self,
247 + out: &mut Vec<u32>,
248 + content: &[AssistantContentBlock],
249 + tool_call_id_to_name: &mut HashMap<String, String>,
250 + ) -> Result<()> {
251 + for block in content {
252 + match block {
253 + AssistantContentBlock::Reasoning { text } => {
254 + self.write_reasoning_block(out, text)?;
255 + }
256 + AssistantContentBlock::Text { text } => {
257 + if !text.is_empty() {
258 + self.write_text_block(out, self.special.message_model, None, text)?;
259 + }
260 + }
261 + AssistantContentBlock::ToolCall(tool_call) => {
262 + if !tool_call.id.is_empty() {
263 + tool_call_id_to_name.insert(tool_call.id.clone(), tool_call.name.clone());
264 + }
265 + self.write_assistant_tool_call(out, tool_call)?;
266 + }
267 + }
268 + }
269 + out.push(self.special.content_model_end_sampling);
270 + Ok(())
271 + }
272 +
273 + fn write_tool_response(
274 + &self,
275 + out: &mut Vec<u32>,
276 + content: &ChatContent,
277 + tool_call_id: &str,
278 + tool_call_id_to_name: &HashMap<String, String>,
279 + ) -> Result<()> {
280 + let text = content.try_flatten_to_text()?;
281 + let tool_name = tool_call_id_to_name.get(tool_call_id).map(String::as_str).unwrap_or("");
282 + self.write_text_block(out, self.special.message_tool, Some(tool_name), &text)
283 + }
284 +}
285 +
286 +impl ChatRenderer for InklingChatRenderer {
287 + fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
288 + request.validate()?;
289 + if request.chat_options.continue_final_message() {
290 + return Err(Error::ChatTemplate(
291 + "Inkling renderer does not support continue_final_message".to_string(),
292 + ));
293 + }
294 +
295 + let mut out = Vec::new();
296 + let mut tool_call_id_to_name = HashMap::new();
297 + let effective_template_kwargs = request_template_kwargs(request);
298 + let tools = rendered_tools(request);
299 + self.write_tool_declarations(&mut out, &tools)?;
300 + let mut reasoning_effort =
301 + resolve_reasoning_effort(effective_template_kwargs.get("reasoning_effort"));
302 +
303 + for message in &request.messages {
304 + if !matches!(
305 + message,
306 + ChatMessage::System { .. } | ChatMessage::Developer { .. }
307 + ) && let Some(effort) = reasoning_effort.take()
308 + {
309 + self.write_reasoning_effort(&mut out, effort)?;
310 + }
311 +
312 + match message {
313 + ChatMessage::System { content } => {
314 + self.write_chat_content(&mut out, self.special.message_system, content)?;
315 + }
316 + ChatMessage::Developer { content, .. } => {
317 + self.write_chat_content(&mut out, self.special.message_system, content)?;
318 + }
319 + ChatMessage::User { content } => {
320 + self.write_chat_content(&mut out, self.special.message_user, content)?;
321 + }
322 + ChatMessage::Assistant { content } => {
323 + self.write_assistant_content(&mut out, content, &mut tool_call_id_to_name)?;
324 + }
325 + ChatMessage::ToolResponse {
326 + content,
327 + tool_call_id,
328 + } => {
329 + self.write_tool_response(
330 + &mut out,
331 + content,
332 + tool_call_id,
333 + &tool_call_id_to_name,
334 + )?;
335 + }
336 + }
337 + }
338 +
339 + if let Some(effort) = reasoning_effort {
340 + self.write_reasoning_effort(&mut out, effort)?;
341 + }
342 +
343 + if request.chat_options.add_generation_prompt() {
344 + out.push(self.special.message_model);
345 + }
346 +
347 + Ok(RenderedPrompt {
348 + prompt: Prompt::TokenIds(out),
349 + effective_template_kwargs,
350 + })
351 + }
352 +}
353 +
354 +fn resolve_reasoning_effort(value: Option<&Value>) -> Option<f64> {
355 + let Some(value) = value else {
356 + return Some(0.9);
357 + };
358 + match value {
359 + Value::String(name) => match name.as_str() {
360 + "none" => Some(0.0),
361 + "minimal" => Some(0.1),
362 + "low" => Some(0.2),
363 + "medium" => Some(0.7),
364 + "high" => Some(0.9),
365 + "xhigh" | "max" => Some(0.99),
366 + _ => None,
367 + },
368 + Value::Number(number) => number.as_f64(),
369 + _ => None,
370 + }
371 +}
372 +
373 +fn resolve_special_token(tokenizer: &dyn Tokenizer, token: &str) -> Result<u32> {
374 + tokenizer.token_to_id(token).ok_or_else(|| {
375 + Error::ChatTemplate(format!(
376 + "Inkling tokenizer is missing special token `{token}`"
377 + ))
378 + })
379 +}
380 +
381 +fn rendered_tools(request: &ChatRequest) -> Vec<&ChatTool> {
382 + if !request.tool_parsing_enabled() {
383 + return Vec::new();
384 + }
385 +
386 + let mut tools = Vec::with_capacity(request.tools.len());
387 + tools.extend(request.tools.iter());
388 + for message in &request.messages {
389 + if let ChatMessage::Developer {
390 + tools: Some(local_tools),
391 + ..
392 + } = message
393 + {
394 + tools.extend(local_tools.iter());
395 + }
396 + }
397 + tools
398 +}
399 +
400 +fn tool_call_json(tool_call: &AssistantToolCall) -> Result<String> {
401 + let name_json = serde_json::to_string(&tool_call.name)
402 + .map_err(|error| Error::ChatTemplate(error.as_report().to_string()))?;
403 + let arguments = if tool_call.arguments.trim().is_empty() {
404 + Value::Object(Map::new())
405 + } else {
406 + serde_json::from_str(&tool_call.arguments).map_err(|error| {
407 + Error::ChatTemplate(format!(
408 + "Inkling tool call arguments must decode to a JSON object: {error}"
409 + ))
410 + })?
411 + };
412 + let Value::Object(_) = arguments else {
413 + return Err(Error::ChatTemplate(
414 + "Inkling tool call arguments must decode to a JSON object".to_string(),
415 + ));
416 + };
417 + let args_json = compact_json(&sort_json(&arguments))?;
418 + Ok(format!("{{\"name\":{name_json},\"args\":{args_json}}}"))
419 +}
420 +
421 +fn compact_json(value: &Value) -> Result<String> {
422 + serde_json::to_string(value).map_err(|error| Error::ChatTemplate(error.as_report().to_string()))
423 +}
424 +
425 +fn sort_json(value: &Value) -> Value {
426 + match value {
427 + Value::Array(items) => Value::Array(items.iter().map(sort_json).collect()),
428 + Value::Object(map) => {
429 + let mut sorted = Map::new();
430 + let mut keys = map.keys().collect::<Vec<_>>();
431 + keys.sort();
432 + for key in keys {
433 + sorted.insert(key.clone(), sort_json(&map[key]));
434 + }
435 + Value::Object(sorted)
436 + }
437 + _ => value.clone(),
438 + }
439 +}
440 +
441 +#[cfg(test)]
442 +mod tests;
rust/src/chat/src/renderer/inkling/tests.rs new
+375
@@ -0,0 +1,375 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::path::PathBuf;
5 +use std::sync::Arc;
6 +
7 +use expect_test::{ExpectFile, expect_file};
8 +use serde_json::json;
9 +use thiserror_ext::AsReport;
10 +use vllm_text::tokenizer::Tokenizer;
11 +
12 +use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request};
13 +
14 +use super::{
15 + AUDIO_END, CONTENT_AUDIO_INPUT, CONTENT_IMAGE, CONTENT_INVOKE_TOOL_JSON, CONTENT_TEXT,
16 + CONTENT_THINKING, CONTENT_XML, END_MESSAGE, InklingChatRenderer, MESSAGE_MODEL, MESSAGE_SYSTEM,
17 + MESSAGE_TOOL, MESSAGE_USER,
18 +};
19 +use crate::event::{AssistantContentBlock, AssistantToolCall};
20 +use crate::request::{ChatMessage, ChatRequest, GenerationPromptMode, ReasoningEffort};
21 +use crate::{ChatRenderer, Error};
22 +
23 +struct FixtureTokenizer;
24 +
25 +impl Tokenizer for FixtureTokenizer {
26 + fn encode(
27 + &self,
28 + text: &str,
29 + _add_special_tokens: bool,
30 + ) -> vllm_text::tokenizer::Result<Vec<u32>> {
31 + Ok(text.bytes().map(u32::from).collect())
32 + }
33 +
34 + fn decode(
35 + &self,
36 + token_ids: &[u32],
37 + _skip_special_tokens: bool,
38 + ) -> vllm_text::tokenizer::Result<String> {
39 + Ok(token_ids
40 + .iter()
41 + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
42 + .collect())
43 + }
44 +
45 + fn token_to_id(&self, token: &str) -> Option<u32> {
46 + match token {
47 + "<|message_user|>" => Some(200000),
48 + "<|message_model|>" => Some(200001),
49 + "<|message_system|>" => Some(200002),
50 + "<|message_tool|>" => Some(200003),
51 + "<|content_text|>" => Some(200004),
52 + "<|content_image|>" => Some(200005),
53 + "<|content_model_end_sampling|>" => Some(200006),
54 + "<|content_thinking|>" => Some(200008),
55 + "<|end_message|>" => Some(200010),
56 + "<|content_audio_input|>" => Some(200020),
57 + CONTENT_XML => Some(200024),
58 + "<|audio_end|>" => Some(200043),
59 + "<|content_invoke_tool_json|>" => Some(200049),
60 + _ => None,
61 + }
62 + }
63 +
64 + fn id_to_token(&self, id: u32) -> Option<String> {
65 + let token = match id {
66 + 200000 => "<|message_user|>",
67 + 200001 => "<|message_model|>",
68 + 200002 => "<|message_system|>",
69 + 200003 => "<|message_tool|>",
70 + 200004 => "<|content_text|>",
71 + 200005 => "<|content_image|>",
72 + 200006 => "<|content_model_end_sampling|>",
73 + 200008 => "<|content_thinking|>",
74 + 200010 => "<|end_message|>",
75 + 200020 => "<|content_audio_input|>",
76 + 200024 => CONTENT_XML,
77 + 200043 => "<|audio_end|>",
78 + 200049 => "<|content_invoke_tool_json|>",
79 + _ => return None,
80 + };
81 + Some(token.to_string())
82 + }
83 +}
84 +
85 +fn renderer() -> InklingChatRenderer {
86 + InklingChatRenderer::new(Arc::new(FixtureTokenizer)).unwrap()
87 +}
88 +
89 +fn render_token_ids(request: &ChatRequest) -> Vec<u32> {
90 + renderer()
91 + .render(request)
92 + .unwrap()
93 + .prompt
94 + .into_token_ids()
95 + .expect("Inkling renderer returns token IDs")
96 +}
97 +
98 +fn fixture_request(name: &str) -> ChatRequest {
99 + fixture_chat_request(&fixture_path(name), inkling_fixture_options())
100 +}
101 +
102 +fn inkling_fixture_options() -> FixtureRequestOptions {
103 + FixtureRequestOptions {
104 + enable_thinking: false,
105 + no_generation_prompt_when_last_assistant: false,
106 + }
107 +}
108 +
109 +fn fixture_path(name: &str) -> PathBuf {
110 + PathBuf::from(env!("CARGO_MANIFEST_DIR"))
111 + .join("src/renderer/inkling")
112 + .join("fixtures")
113 + .join(name)
114 +}
115 +
116 +fn assert_fixture(input_name: &str, expected: ExpectFile) {
117 + let request = fixture_request(input_name);
118 + let rendered = format!("{}\n", render_symbolic_tokens(&render_token_ids(&request)));
119 +
120 + expected.assert_eq(&rendered);
121 +}
122 +
123 +fn render_symbolic_tokens(token_ids: &[u32]) -> String {
124 + let mut out = String::new();
125 + let mut text_bytes = Vec::new();
126 + for token_id in token_ids {
127 + if let Some(marker) = symbolic_tokens()
128 + .iter()
129 + .find_map(|(marker, marker_id)| (token_id == marker_id).then_some(*marker))
130 + {
131 + flush_text_bytes(&mut out, &mut text_bytes);
132 + out.push_str(marker);
133 + continue;
134 + }
135 +
136 + let byte = u8::try_from(*token_id)
137 + .unwrap_or_else(|_| panic!("unexpected non-byte fixture text token {token_id}"));
138 + text_bytes.push(byte);
139 + }
140 + flush_text_bytes(&mut out, &mut text_bytes);
141 + out
142 +}
143 +
144 +fn flush_text_bytes(out: &mut String, text_bytes: &mut Vec<u8>) {
145 + if text_bytes.is_empty() {
146 + return;
147 + }
148 + out.push_str(std::str::from_utf8(text_bytes).unwrap());
149 + text_bytes.clear();
150 +}
151 +
152 +fn symbolic_tokens() -> [(&'static str, u32); 13] {
153 + [
154 + (MESSAGE_USER, 200000),
155 + (MESSAGE_MODEL, 200001),
156 + (MESSAGE_SYSTEM, 200002),
157 + (MESSAGE_TOOL, 200003),
158 + (CONTENT_TEXT, 200004),
159 + (CONTENT_IMAGE, 200005),
160 + (super::CONTENT_MODEL_END_SAMPLING, 200006),
161 + (CONTENT_THINKING, 200008),
162 + (END_MESSAGE, 200010),
163 + (CONTENT_AUDIO_INPUT, 200020),
164 + (CONTENT_XML, 200024),
165 + (AUDIO_END, 200043),
166 + (CONTENT_INVOKE_TOOL_JSON, 200049),
167 + ]
168 +}
169 +
170 +#[test]
171 +fn renders_text_and_image_as_inkling_tokens() {
172 + assert_fixture(
173 + "text_image_input.json",
174 + expect_file!["fixtures/text_image_output.txt"],
175 + );
176 +}
177 +
178 +#[test]
179 +fn renders_audio_marker_only_blocks() {
180 + assert_fixture(
181 + "text_audio_input.json",
182 + expect_file!["fixtures/text_audio_output.txt"],
183 + );
184 +}
185 +
186 +#[test]
187 +fn renders_reasoning_text_tool_call_and_tool_response() {
188 + assert_fixture(
189 + "tool_round_trip_input.json",
190 + expect_file!["fixtures/tool_round_trip_output.txt"],
191 + );
192 +}
193 +
194 +#[test]
195 +fn renders_request_and_developer_tools_as_tool_declare() {
196 + assert_fixture(
197 + "tool_declare_input.json",
198 + expect_file!["fixtures/tool_declare_output.txt"],
199 + );
200 +}
201 +
202 +#[test]
203 +fn renders_named_reasoning_effort_after_tool_declarations() {
204 + for (effort, expected) in [
205 + (ReasoningEffort::None, "0.0"),
206 + (ReasoningEffort::Minimal, "0.1"),
207 + (ReasoningEffort::Low, "0.2"),
208 + (ReasoningEffort::Medium, "0.7"),
209 + (ReasoningEffort::High, "0.9"),
210 + (ReasoningEffort::XHigh, "0.99"),
211 + (ReasoningEffort::Max, "0.99"),
212 + ] {
213 + let mut request = fixture_request("tool_declare_input.json");
214 + request.chat_options.reasoning_effort = Some(effort);
215 +
216 + let rendered = render_symbolic_tokens(&render_token_ids(&request));
217 + let tool_end = rendered.find("<|end_message|>").unwrap();
218 + let effort_block = format!(
219 + "<|message_system|><|content_text|>Thinking effort level: \
220 + {expected}<|end_message|>"
221 + );
222 + let effort_start = rendered.find(&effort_block).unwrap();
223 + let system_end = rendered.find("rules<|end_message|>").unwrap();
224 +
225 + assert!(effort_start > tool_end);
226 + assert!(effort_start > system_end);
227 + }
228 +}
229 +
230 +#[test]
231 +fn emits_one_reasoning_effort_for_multi_turn_conversation() {
232 + let request = ChatRequest {
233 + messages: vec![
234 + ChatMessage::system("rules"),
235 + ChatMessage::user("user1"),
236 + ChatMessage::assistant_text("assistant1"),
237 + ChatMessage::user("user2"),
238 + ],
239 + ..ChatRequest::for_test()
240 + };
241 +
242 + let rendered = render_symbolic_tokens(&render_token_ids(&request));
243 + let effort = "<|message_system|><|content_text|>Thinking effort level: 0.9\
244 + <|end_message|>";
245 +
246 + assert_eq!(rendered.matches(effort).count(), 1);
247 + assert!(rendered.find("rules").unwrap() < rendered.find(effort).unwrap());
248 + assert!(rendered.find(effort).unwrap() < rendered.find("user1").unwrap());
249 + assert!(rendered.find("user1").unwrap() < rendered.find("assistant1").unwrap());
250 + assert!(rendered.find("assistant1").unwrap() < rendered.find("user2").unwrap());
251 +}
252 +
253 +#[test]
254 +fn canonicalizes_numeric_zero_reasoning_effort() {
255 + for value in [0.0, -0.0] {
256 + let mut request = ChatRequest::for_test();
257 + request
258 + .chat_options
259 + .template_kwargs
260 + .insert("reasoning_effort".to_string(), json!(value));
261 +
262 + assert!(
263 + render_symbolic_tokens(&render_token_ids(&request)).starts_with(
264 + "<|message_system|><|content_text|>Thinking effort level: 0.0\
265 + <|end_message|>"
266 + )
267 + );
268 + }
269 +}
270 +
271 +#[test]
272 +fn defaults_reasoning_effort_to_high() {
273 + let request = ChatRequest::for_test();
274 +
275 + assert!(
276 + render_symbolic_tokens(&render_token_ids(&request)).starts_with(
277 + "<|message_system|><|content_text|>Thinking effort level: 0.9\
278 + <|end_message|>"
279 + )
280 + );
281 +}
282 +
283 +#[test]
284 +fn renders_numeric_reasoning_effort_template_kwarg() {
285 + let mut request = ChatRequest::for_test();
286 + request
287 + .chat_options
288 + .template_kwargs
289 + .insert("reasoning_effort".to_string(), json!(0.8));
290 +
291 + assert_eq!(
292 + render_symbolic_tokens(&render_token_ids(&request)),
293 + "<|message_system|><|content_text|>Thinking effort level: 0.8\
294 + <|end_message|><|message_user|><|content_text|>test\
295 + <|end_message|><|message_model|>"
296 + );
297 +}
298 +
299 +#[test]
300 +fn ignores_unsupported_reasoning_effort_values() {
301 + for value in [json!(true), json!("invalid"), json!(null)] {
302 + let mut request = ChatRequest::for_test();
303 + request
304 + .chat_options
305 + .template_kwargs
306 + .insert("reasoning_effort".to_string(), value);
307 +
308 + assert_eq!(
309 + render_symbolic_tokens(&render_token_ids(&request)),
310 + "<|message_user|><|content_text|>test<|end_message|><|message_model|>"
311 + );
312 + }
313 +}
314 +
315 +#[test]
316 +fn rejects_out_of_range_reasoning_effort() {
317 + for value in [0.990_000_1, 1.0, 1.5, -0.1] {
318 + let mut request = ChatRequest::for_test();
319 + request
320 + .chat_options
321 + .template_kwargs
322 + .insert("reasoning_effort".to_string(), json!(value));
323 +
324 + let error = renderer().render(&request).unwrap_err();
325 + assert!(error.as_report().to_string().contains("must be in [0.0, 0.99]"));
326 + }
327 +}
328 +
329 +#[test]
330 +fn rejects_continue_final_message() {
331 + let request = ChatRequest {
332 + messages: vec![ChatMessage::assistant_text("partial")],
333 + chat_options: crate::ChatOptions {
334 + generation_prompt_mode: GenerationPromptMode::ContinueFinalAssistant,
335 + ..Default::default()
336 + },
337 + ..ChatRequest::for_test()
338 + };
339 +
340 + let error = renderer().render(&request).unwrap_err();
341 + assert!(matches!(error, Error::ChatTemplate(_)));
342 + assert!(error.as_report().to_string().contains("continue_final_message"));
343 +}
344 +
345 +#[test]
346 +fn renders_developer_messages_as_system() {
347 + let request = ChatRequest {
348 + messages: vec![ChatMessage::developer("rules", None)],
349 + ..ChatRequest::for_test()
350 + };
351 +
352 + assert_eq!(
353 + render_symbolic_tokens(&render_token_ids(&request)),
354 + "<|message_system|><|content_text|>rules<|end_message|>\
355 + <|message_system|><|content_text|>Thinking effort level: 0.9\
356 + <|end_message|><|message_model|>"
357 + );
358 +}
359 +
360 +#[test]
361 +fn rejects_non_object_tool_call_arguments() {
362 + let request = ChatRequest {
363 + messages: vec![ChatMessage::assistant_blocks(vec![
364 + AssistantContentBlock::ToolCall(AssistantToolCall {
365 + id: "call_1".to_string(),
366 + name: "get_weather".to_string(),
367 + arguments: "[]".to_string(),
368 + }),
369 + ])],
370 + ..ChatRequest::for_test()
371 + };
372 +
373 + let error = renderer().render(&request).unwrap_err();
374 + assert!(error.as_report().to_string().contains("JSON object"));
375 +}
rust/src/chat/src/renderer/mod.rs
+2
@@ -14,6 +14,7 @@ pub mod deepseek_v32;
14 pub mod deepseek_v4;
15 pub mod harmony;
16 pub mod hf;
17 +mod inkling;
18 mod selection;
19 #[cfg(test)]
20 mod test_utils;
@@ -21,6 +22,7 @@ mod test_utils;
22 pub use deepseek_v4::DeepSeekV4ChatRenderer;
23 pub use deepseek_v32::DeepSeekV32ChatRenderer;
24 pub use harmony::HarmonyChatRenderer;
25 +pub use inkling::InklingChatRenderer;
26 pub use selection::RendererSelection;
27
28 /// Rendered chat prompt submitted to the text backend.
rust/src/chat/src/renderer/selection.rs
+9 -1
@@ -24,6 +24,8 @@ pub enum RendererSelection {
24 DeepSeekV4,
25 /// Force the GPT-OSS Harmony renderer.
26 Harmony,
27 + /// Force the Inkling native token renderer.
28 + Inkling,
29 }
30
31 impl RendererSelection {
@@ -33,6 +35,8 @@ impl RendererSelection {
35 pub const GPT_OSS_MODEL_TYPE: &str = "gpt_oss";
36 pub const HARMONY_LITERAL: &str = "harmony";
37 pub const HF_LITERAL: &str = "hf";
38 + pub const INKLING_LITERAL: &str = "inkling";
39 + pub const INKLING_MODEL_TYPE: &str = "inkling_mm_model";
40
41 /// Resolve the renderer selection using the given model type string, if
42 /// it's `Auto`.
@@ -42,6 +46,7 @@ impl RendererSelection {
46 Self::DEEPSEEK_V32_LITERAL => Self::DeepSeekV32,
47 Self::DEEPSEEK_V4_LITERAL => Self::DeepSeekV4,
48 Self::GPT_OSS_MODEL_TYPE => Self::Harmony,
49 + Self::INKLING_MODEL_TYPE => Self::Inkling,
50 _ => Self::Hf,
51 },
52 selection => selection,
@@ -63,6 +68,8 @@ impl FromStr for RendererSelection {
68 Ok(Self::DeepSeekV4)
69 } else if value.eq_ignore_ascii_case(Self::HARMONY_LITERAL) {
70 Ok(Self::Harmony)
71 + } else if value.eq_ignore_ascii_case(Self::INKLING_LITERAL) {
72 + Ok(Self::Inkling)
73 } else {
74 Err(format!(
75 "unknown renderer `{value}` (expected one of: {})",
@@ -80,6 +87,7 @@ impl fmt::Display for RendererSelection {
87 Self::DeepSeekV32 => f.write_str(Self::DEEPSEEK_V32_LITERAL),
88 Self::DeepSeekV4 => f.write_str(Self::DEEPSEEK_V4_LITERAL),
89 Self::Harmony => f.write_str(Self::HARMONY_LITERAL),
90 + Self::Inkling => f.write_str(Self::INKLING_LITERAL),
91 }
92 }
93 }
@@ -106,7 +114,7 @@ mod tests {
114 fn renderer_selection_expected_error_message() {
115 let err = RendererSelection::from_str("unknown").unwrap_err();
116 expect_test::expect![
109 - "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony)"
117 + "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling)"
118 ]
119 .assert_eq(&err);
120 }
rust/src/chat/src/renderer/test_utils.rs
+11
@@ -103,6 +103,8 @@ pub(crate) enum FixtureContent {
103 pub(crate) enum FixtureContentPart {
104 Text { text: String },
105 ImageUrl { image_url: String },
106 + InputAudio { data: String },
107 + AudioUrl { audio_url: String },
108 }
109
110 #[derive(Debug, Deserialize)]
@@ -226,6 +228,15 @@ fn to_chat_content(content: FixtureContent) -> ChatContent {
228 FixtureContentPart::ImageUrl { image_url } => {
229 ChatContentPart::image_url(image_url)
230 }
231 + FixtureContentPart::InputAudio { data } => ChatContentPart::InputAudio {
232 + data,
233 + format: None,
234 + uuid: None,
235 + },
236 + FixtureContentPart::AudioUrl { audio_url } => ChatContentPart::AudioUrl {
237 + audio_url,
238 + uuid: None,
239 + },
240 })
241 .collect(),
242 ),
rust/src/chat/tests/roundtrip.rs
+16 -2
@@ -263,19 +263,32 @@ impl RoundtripCase {
263 sort_json_keys: false,
264 }
265 }
266 +
267 + /// Inkling typed content blocks with native token-id rendering.
268 + fn inkling() -> Self {
269 + Self {
270 + model_id: "thinkingmachines/Inkling",
271 + assistant_stop_suffix: "",
272 + tool_call_parser: ParserSelection::Auto,
273 + reasoning_parser: ParserSelection::Auto,
274 + thinking_behavior: ThinkingBehavior::Always { value: true },
275 + json_fmt: compact_json_fmt(),
276 + sort_json_keys: true,
277 + }
278 + }
279 }
280
281 macro_rules! roundtrip_tests {
269 - ($($case:ident => [$($(#[$fixture_attr:meta])* $fixture:ident),* $(,)?]),+ $(,)?) => {
282 + ($($case:ident => $(#[$case_attr:meta])* [$($fixture:ident),* $(,)?]),+ $(,)?) => {
283 paste::paste! {
284 $(
285 #[tokio::test]
286 #[file_serial([<hf_ $case>])]
287 + $(#[$case_attr])*
288 async fn [<roundtrip_ $case>]() -> Result<()> {
289 let case = RoundtripCase::$case();
290 let backends = load_roundtrip_backends(&case).await?;
291 $(
278 - $(#[$fixture_attr])*
292 [<run_roundtrip_ $fixture>](&case, &backends).await?;
293 )*
294 Ok(())
@@ -300,6 +313,7 @@ roundtrip_tests! {
313 gemma4 => [tool_call_mix], // Gemma4 strips reasoning in history if there's no tool call
314 kimi_k25 => [tool_call_mix], // Kimi K2.5 strips reasoning in history
315 gpt_oss => [tool_call_mix], // Harmony strips reasoning in history if there's no tool call
316 + inkling => [reasoning_and_content, tool_call_mix],
317 }
318
319 /// Run the fixed reasoning+content fixture for one model/parser case.
rust/src/cmd/src/cli/tests.rs
+1 -1
@@ -646,7 +646,7 @@ fn serve_args_reject_unknown_renderer_value() {
646 .unwrap_err();
647
648 expect![[r#"
649 - error: invalid value 'definitely_missing' for '--tokenizer-mode <RENDERER>': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony)
649 + error: invalid value 'definitely_missing' for '--tokenizer-mode <RENDERER>': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling)
650
651 For more information, try '--help'.
652 "#]]
rust/src/parser/src/tool/json/mod.rs
+12 -12
@@ -31,24 +31,24 @@ use super::utils::{
31 };
32 use super::{Result, ToolCallDelta, ToolParserOutput};
33
34 -type JsonToolInput<'i> = Partial<&'i str>;
34 +pub(crate) type JsonToolInput<'i> = Partial<&'i str>;
35
36 #[derive(Debug, Clone, Copy)]
37 -struct JsonToolCallConfig {
38 - parser_name: &'static str,
39 - start_marker: &'static str,
40 - end_marker: &'static str,
41 - marker_whitespace: JsonToolCallWhitespace,
42 - delimiter: Option<&'static str>,
43 - name_key: &'static str,
37 +pub(crate) struct JsonToolCallConfig {
38 + pub parser_name: &'static str,
39 + pub start_marker: &'static str,
40 + pub end_marker: &'static str,
41 + pub marker_whitespace: JsonToolCallWhitespace,
42 + pub delimiter: Option<&'static str>,
43 + pub name_key: &'static str,
44 /// Candidate JSON keys naming the arguments payload, tried in order.
45 /// Most parsers use a single key like `["arguments"]`, but some accept
46 /// multiple (e.g. InternLM2 accepts `parameters` or `arguments`).
47 - arguments_key: &'static [&'static str],
47 + pub arguments_key: &'static [&'static str],
48 }
49
50 #[derive(Debug, Clone, Copy)]
51 -enum JsonToolCallWhitespace {
51 +pub(crate) enum JsonToolCallWhitespace {
52 Optional,
53 Exact(&'static str),
54 }
@@ -61,7 +61,7 @@ enum JsonToolCallMode {
61 }
62
63 #[derive(Debug, Clone, PartialEq, Eq)]
64 -enum JsonToolCallEvent {
64 +pub(crate) enum JsonToolCallEvent {
65 Text { len: usize },
66 ToolCallStart,
67 ToolCallHeader { function_name: String },
@@ -220,7 +220,7 @@ fn tool_call_start_event(
220
221 /// Parse a marker-wrapped JSON tool-call header before the raw arguments
222 /// payload.
223 -fn tool_call_header_event(
223 +pub(crate) fn tool_call_header_event(
224 input: &mut JsonToolInput<'_>,
225 config: JsonToolCallConfig,
226 ) -> ModalResult<JsonToolCallEvent> {
rust/src/parser/src/tool/mod.rs
+1 -1
@@ -9,7 +9,7 @@ mod deepseek_dsml;
9 pub(crate) mod deepseek_json;
10 mod glm_xml;
11 mod hy_v3;
12 -mod json;
12 +pub(crate) mod json;
13 mod kimi_k2;
14 mod minimax_m2;
15 mod minimax_m3;
rust/src/parser/src/unified/gemma4.rs
+5 -14
@@ -10,7 +10,7 @@ use winnow::prelude::*;
10 use winnow::stream::{Partial, Stream};
11 use winnow::token::{literal, take_till, take_until};
12
13 -use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput};
13 +use super::{Result, UnifiedParser, UnifiedParserOutput, token_id};
14 use crate::reasoning::last_reasoning_boundary;
15 use crate::tool::{Tool, ToolCallDelta};
16 use crate::unified::parsing_failed;
@@ -79,17 +79,8 @@ pub struct Gemma4UnifiedParser {
79 impl Gemma4UnifiedParser {
80 /// Create a Gemma4 parser.
81 pub fn new(_tools: &[Tool], tokenizer: DynTokenizer) -> Result<Self> {
82 - let channel_start_token_id = tokenizer.token_to_id(CHANNEL_START).ok_or_else(|| {
83 - UnifiedParserError::MissingToken {
84 - token: CHANNEL_START.to_string(),
85 - }
86 - })?;
87 - let channel_end_token_id =
88 - tokenizer
89 - .token_to_id(CHANNEL_END)
90 - .ok_or_else(|| UnifiedParserError::MissingToken {
91 - token: CHANNEL_END.to_string(),
92 - })?;
82 + let channel_start_token_id = token_id(tokenizer.as_ref(), CHANNEL_START)?;
83 + let channel_end_token_id = token_id(tokenizer.as_ref(), CHANNEL_END)?;
84
85 Ok(Self {
86 buffer: String::new(),
@@ -524,10 +515,10 @@ mod tests {
515
516 use super::{
517 CHANNEL_END, CHANNEL_START, Gemma4UnifiedParser, ToolCallDelta, UnifiedParser,
527 - UnifiedParserError, UnifiedParserOutput, gemma4_array_content, parse_gemma4_args,
518 + UnifiedParserOutput, gemma4_array_content, parse_gemma4_args,
519 };
520 use crate::tool::Tool;
530 - use crate::unified::{UnifiedParserEvent, parsing_failed};
521 + use crate::unified::{UnifiedParserError, UnifiedParserEvent, parsing_failed};
522
523 const CHANNEL_START_ID: u32 = 256;
524 const CHANNEL_END_ID: u32 = 257;
rust/src/parser/src/unified/inkling.rs new
+768
@@ -0,0 +1,768 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use winnow::ascii::multispace0 as ws0;
5 +use winnow::combinator::{alt, seq};
6 +use winnow::error::ModalResult;
7 +use winnow::prelude::*;
8 +use winnow::stream::Partial;
9 +use winnow::token::literal;
10 +
11 +use vllm_tokenizer::DynTokenizer;
12 +
13 +use super::{Result, UnifiedParser, UnifiedParserOutput, token_id};
14 +use crate::tool::json::{
15 + JsonToolCallConfig, JsonToolCallEvent, JsonToolCallWhitespace, JsonToolInput,
16 + tool_call_header_event,
17 +};
18 +use crate::tool::{Tool, ToolCallDelta};
19 +use crate::unified::parsing_failed;
20 +use crate::utils::{
21 + JsonObjectScanState, parse_buffered_event, safe_text_len_mul, take_json_object,
22 +};
23 +
24 +const CONTENT_TEXT: &str = "<|content_text|>";
25 +const CONTENT_THINKING: &str = "<|content_thinking|>";
26 +const CONTENT_INVOKE_TOOL_JSON: &str = "<|content_invoke_tool_json|>";
27 +const CONTENT_INVOKE_TOOL_TEXT: &str = "<|content_invoke_tool_text|>";
28 +const CONTENT_MODEL_END_SAMPLING: &str = "<|content_model_end_sampling|>";
29 +const CONTENT_TOOL_ERROR: &str = "<|content_tool_error|>";
30 +const MESSAGE_MODEL: &str = "<|message_model|>";
31 +const END_MESSAGE: &str = "<|end_message|>";
32 +
33 +const IDLE_MARKERS: &[&str] = &[
34 + MESSAGE_MODEL,
35 + CONTENT_TEXT,
36 + CONTENT_THINKING,
37 + CONTENT_INVOKE_TOOL_JSON,
38 + CONTENT_INVOKE_TOOL_TEXT,
39 + CONTENT_MODEL_END_SAMPLING,
40 + CONTENT_TOOL_ERROR,
41 + END_MESSAGE,
42 +];
43 +const BLOCK_END_MARKERS: &[&str] = &[END_MESSAGE, CONTENT_MODEL_END_SAMPLING];
44 +
45 +const INKLING_TOOL_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
46 + parser_name: "Inkling",
47 + start_marker: CONTENT_INVOKE_TOOL_JSON,
48 + end_marker: END_MESSAGE,
49 + marker_whitespace: JsonToolCallWhitespace::Optional,
50 + delimiter: None,
51 + name_key: "name",
52 + arguments_key: &["args"],
53 +};
54 +
55 +type InklingInput<'i> = Partial<&'i str>;
56 +
57 +#[derive(Debug, Clone, PartialEq, Eq)]
58 +enum InklingEvent {
59 + Text { len: usize },
60 + Reasoning { len: usize },
61 + TextStart,
62 + ReasoningStart,
63 + MessageStart,
64 + Header,
65 + ToolJsonStart,
66 + ToolJsonHeader { name: String },
67 + ToolJsonArgs { len: usize, complete: bool },
68 + BlockEnd,
69 +}
70 +
71 +#[derive(Debug, Clone, Default, PartialEq, Eq)]
72 +enum InklingMode {
73 + #[default]
74 + Idle,
75 + MessageHeader,
76 + Text,
77 + Reasoning,
78 + ToolJsonHeader,
79 + ToolJsonArgs {
80 + json_scan: JsonObjectScanState,
81 + },
82 + ToolJsonClose,
83 +}
84 +
85 +/// Unified parser for Inkling typed content blocks.
86 +pub struct InklingUnifiedParser {
87 + buffer: String,
88 + mode: InklingMode,
89 + emitted_tool_count: usize,
90 + active_tool_index: Option<usize>,
91 + tokenizer: DynTokenizer,
92 + message_model_token_id: u32,
93 + content_text_token_id: u32,
94 + content_thinking_token_id: u32,
95 +}
96 +
97 +impl InklingUnifiedParser {
98 + /// Create a Inkling parser.
99 + pub fn new(_tools: &[Tool], tokenizer: DynTokenizer) -> Result<Self> {
100 + let message_model_token_id = token_id(tokenizer.as_ref(), MESSAGE_MODEL)?;
101 + let content_text_token_id = token_id(tokenizer.as_ref(), CONTENT_TEXT)?;
102 + let content_thinking_token_id = token_id(tokenizer.as_ref(), CONTENT_THINKING)?;
103 +
104 + Ok(Self {
105 + buffer: String::new(),
106 + mode: InklingMode::Idle,
107 + emitted_tool_count: 0,
108 + active_tool_index: None,
109 + tokenizer,
110 + message_model_token_id,
111 + content_text_token_id,
112 + content_thinking_token_id,
113 + })
114 + }
115 +
116 + fn initialize_mode(&mut self, prompt_token_ids: &[u32]) {
117 + self.mode = InklingMode::Idle;
118 + for token_id in prompt_token_ids.iter().rev().copied() {
119 + if token_id == self.message_model_token_id {
120 + self.mode = InklingMode::MessageHeader;
121 + return;
122 + }
123 + if token_id == self.content_thinking_token_id {
124 + self.mode = InklingMode::Reasoning;
125 + return;
126 + }
127 + if token_id == self.content_text_token_id {
128 + self.mode = InklingMode::Text;
129 + return;
130 + }
131 + if self.tokenizer.is_special_id(token_id) {
132 + return;
133 + }
134 + }
135 + }
136 +
137 + fn apply_event(&mut self, event: InklingEvent, output: &mut UnifiedParserOutput) -> Result<()> {
138 + match event {
139 + InklingEvent::Text { len } => output.push_text(self.buffer[..len].to_string()),
140 + InklingEvent::Reasoning { len } => {
141 + output.push_reasoning(self.buffer[..len].to_string());
142 + }
143 + InklingEvent::MessageStart => self.mode = InklingMode::MessageHeader,
144 + InklingEvent::Header => {}
145 + InklingEvent::TextStart => self.mode = InklingMode::Text,
146 + InklingEvent::ReasoningStart => self.mode = InklingMode::Reasoning,
147 + InklingEvent::ToolJsonStart => self.mode = InklingMode::ToolJsonHeader,
148 + InklingEvent::ToolJsonHeader { name } => {
149 + let tool_index = self.emitted_tool_count;
150 + self.emitted_tool_count += 1;
151 + self.active_tool_index = Some(tool_index);
152 + self.mode = InklingMode::ToolJsonArgs {
153 + json_scan: JsonObjectScanState::default(),
154 + };
155 + output.push_call(ToolCallDelta {
156 + tool_index,
157 + name: Some(name),
158 + arguments: String::new(),
159 + });
160 + }
161 + InklingEvent::ToolJsonArgs { len, complete } => {
162 + let Some(tool_index) = self.active_tool_index else {
163 + return Err(parsing_failed!(
164 + "Inkling arguments without an active tool call"
165 + ));
166 + };
167 + output.push_call(ToolCallDelta {
168 + tool_index,
169 + name: None,
170 + arguments: self.buffer[..len].to_string(),
171 + });
172 + if complete {
173 + self.mode = InklingMode::ToolJsonClose;
174 + }
175 + }
176 + InklingEvent::BlockEnd => {
177 + self.mode = InklingMode::Idle;
178 + self.active_tool_index = None;
179 + }
180 + }
181 + Ok(())
182 + }
183 +
184 + fn reset(&mut self) -> String {
185 + self.mode = InklingMode::Idle;
186 + self.active_tool_index = None;
187 + self.emitted_tool_count = 0;
188 + std::mem::take(&mut self.buffer)
189 + }
190 +}
191 +
192 +impl UnifiedParser for InklingUnifiedParser {
193 + fn create(tools: &[Tool], tokenizer: DynTokenizer) -> Result<Box<dyn UnifiedParser>>
194 + where
195 + Self: Sized + 'static,
196 + {
197 + Self::new(tools, tokenizer).map(|parser| Box::new(parser) as Box<dyn UnifiedParser>)
198 + }
199 +
200 + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
201 + self.buffer.clear();
202 + self.emitted_tool_count = 0;
203 + self.active_tool_index = None;
204 + self.initialize_mode(prompt_token_ids);
205 + Ok(())
206 + }
207 +
208 + fn preserve_special_tokens(&self) -> bool {
209 + true
210 + }
211 +
212 + fn parse_into(&mut self, chunk: &str, output: &mut UnifiedParserOutput) -> Result<()> {
213 + self.buffer.push_str(chunk);
214 +
215 + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
216 + parse_next_inkling_event(input, &mut self.mode)
217 + })? {
218 + self.apply_event(event, output)?;
219 + self.buffer.drain(..consumed_len);
220 + }
221 +
222 + Ok(())
223 + }
224 +
225 + fn finish(&mut self) -> Result<UnifiedParserOutput> {
226 + let mut output = UnifiedParserOutput::default();
227 +
228 + match &self.mode {
229 + InklingMode::Idle | InklingMode::Text => {
230 + output.push_text(std::mem::take(&mut self.buffer))
231 + }
232 + InklingMode::MessageHeader => self.buffer.clear(),
233 + InklingMode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)),
234 + InklingMode::ToolJsonHeader
235 + | InklingMode::ToolJsonArgs { .. }
236 + | InklingMode::ToolJsonClose => {
237 + return Err(parsing_failed!("incomplete Inkling tool call"));
238 + }
239 + }
240 +
241 + let _ = self.reset();
242 + Ok(output)
243 + }
244 +
245 + fn reset(&mut self) -> String {
246 + InklingUnifiedParser::reset(self)
247 + }
248 +}
249 +
250 +/// Parse one Inkling event from buffered streaming input.
251 +fn parse_next_inkling_event(
252 + input: &mut InklingInput<'_>,
253 + mode: &mut InklingMode,
254 +) -> ModalResult<InklingEvent> {
255 + match mode {
256 + InklingMode::Idle => parse_idle_event(input),
257 + InklingMode::MessageHeader => parse_message_header_event(input),
258 + InklingMode::Text => parse_text_event(input),
259 + InklingMode::Reasoning => parse_reasoning_event(input),
260 + InklingMode::ToolJsonHeader => parse_tool_json_header_event(input),
261 + InklingMode::ToolJsonArgs { json_scan } => parse_tool_json_args_event(input, json_scan),
262 + InklingMode::ToolJsonClose => parse_tool_json_close_event(input),
263 + }
264 +}
265 +
266 +/// Parse an event while waiting for a Inkling content kind.
267 +fn parse_idle_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {
268 + alt((
269 + message_start_event,
270 + reasoning_start_event,
271 + text_start_event,
272 + tool_json_start_event,
273 + raw_text_start_event,
274 + block_end_event,
275 + safe_idle_text_event,
276 + ))
277 + .parse_next(input)
278 +}
279 +
280 +/// Parse a Inkling model-authored message start marker.
281 +fn message_start_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {
282 + literal(MESSAGE_MODEL).value(InklingEvent::MessageStart).parse_next(input)
283 +}
284 +
285 +/// Parse an event while waiting for an Inkling message content kind.
286 +fn parse_message_header_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {
287 + alt((
288 + reasoning_start_event,
289 + text_start_event,
290 + tool_json_start_event,
291 + raw_text_start_event,
292 + block_end_event,
293 + safe_header_event,
294 + ))
295 + .parse_next(input)
296 +}
297 +
298 +/// Parse an event inside a Inkling text block.
299 +fn parse_text_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {
300 + alt((block_end_event, safe_text_event)).parse_next(input)
301 +}
302 +
303 +/// Parse an event inside a Inkling reasoning block.
304 +fn parse_reasoning_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {
305 + alt((block_end_event, safe_reasoning_event)).parse_next(input)
306 +}
307 +
308 +/// Parse a Inkling text start marker.
309 +fn text_start_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {
310 + literal(CONTENT_TEXT).value(InklingEvent::TextStart).parse_next(input)
311 +}
312 +
313 +/// Parse a Inkling reasoning start marker.
314 +fn reasoning_start_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {
315 + literal(CONTENT_THINKING).value(InklingEvent::ReasoningStart).parse_next(input)
316 +}
317 +
318 +/// Parse a Inkling JSON tool-call start marker.
319 +fn tool_json_start_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {
320 + literal(CONTENT_INVOKE_TOOL_JSON)
321 + .value(InklingEvent::ToolJsonStart)
322 + .parse_next(input)
323 +}
324 +
325 +/// Parse a Inkling content kind treated as visible text.
326 +fn raw_text_start_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {
327 + alt((
328 + literal(CONTENT_INVOKE_TOOL_TEXT),
329 + literal(CONTENT_TOOL_ERROR),
330 + ))
331 + .value(InklingEvent::TextStart)
332 + .parse_next(input)
333 +}
334 +
335 +/// Parse a Inkling block end marker.
336 +fn block_end_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {
337 + alt((literal(END_MESSAGE), literal(CONTENT_MODEL_END_SAMPLING)))
338 + .value(InklingEvent::BlockEnd)
339 + .parse_next(input)
340 +}
341 +
342 +/// Parse safe text while waiting for the next Inkling marker.
343 +fn safe_idle_text_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {
344 + safe_text_len_mul(input, IDLE_MARKERS).map(|len| InklingEvent::Text { len })
345 +}
346 +
347 +/// Parse safe header text before the next Inkling marker.
348 +fn safe_header_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {
349 + safe_text_len_mul(input, IDLE_MARKERS).map(|_| InklingEvent::Header)
350 +}
351 +
352 +/// Parse safe text before the end of a Inkling text block.
353 +fn safe_text_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {
354 + safe_text_len_mul(input, BLOCK_END_MARKERS).map(|len| InklingEvent::Text { len })
355 +}
356 +
357 +/// Parse safe reasoning before the end of a Inkling reasoning block.
358 +fn safe_reasoning_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {
359 + safe_text_len_mul(input, BLOCK_END_MARKERS).map(|len| InklingEvent::Reasoning { len })
360 +}
361 +
362 +/// Parse a Inkling JSON tool-call header.
363 +fn parse_tool_json_header_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {
364 + match tool_call_header_event(input, INKLING_TOOL_CONFIG)? {
365 + JsonToolCallEvent::ToolCallHeader { function_name } => Ok(InklingEvent::ToolJsonHeader {
366 + name: function_name,
367 + }),
368 + _ => unreachable!("tool_call_header_event only emits ToolCallHeader"),
369 + }
370 +}
371 +
372 +/// Parse raw Inkling JSON tool-call argument bytes.
373 +fn parse_tool_json_args_event(
374 + input: &mut JsonToolInput<'_>,
375 + json_scan: &mut JsonObjectScanState,
376 +) -> ModalResult<InklingEvent> {
377 + let len = take_json_object(input, json_scan)?;
378 + Ok(InklingEvent::ToolJsonArgs {
379 + len,
380 + complete: json_scan.complete(),
381 + })
382 +}
383 +
384 +/// Parse the close of a Inkling JSON tool-call block.
385 +fn parse_tool_json_close_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {
386 + seq!(
387 + _: ws0,
388 + _: literal("}"),
389 + _: ws0,
390 + _: literal(END_MESSAGE),
391 + )
392 + .value(InklingEvent::BlockEnd)
393 + .parse_next(input)
394 +}
395 +
396 +#[cfg(test)]
397 +mod tests {
398 + use std::sync::Arc;
399 +
400 + use super::{CONTENT_TEXT, CONTENT_THINKING, InklingUnifiedParser, MESSAGE_MODEL};
401 + use crate::tool::Tool;
402 + use crate::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput};
403 + use thiserror_ext::AsReport;
404 + use vllm_tokenizer::Tokenizer;
405 +
406 + struct FakeTokenizer;
407 +
408 + impl Tokenizer for FakeTokenizer {
409 + fn encode(
410 + &self,
411 + text: &str,
412 + _add_special_tokens: bool,
413 + ) -> vllm_tokenizer::Result<Vec<u32>> {
414 + Ok(text.chars().map(u32::from).collect())
415 + }
416 +
417 + fn decode(
418 + &self,
419 + token_ids: &[u32],
420 + _skip_special_tokens: bool,
421 + ) -> vllm_tokenizer::Result<String> {
422 + Ok(token_ids
423 + .iter()
424 + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
425 + .collect())
426 + }
427 +
428 + fn token_to_id(&self, token: &str) -> Option<u32> {
429 + match token {
430 + MESSAGE_MODEL => Some(200001),
431 + "<|content_text|>" => Some(200004),
432 + "<|content_model_end_sampling|>" => Some(200006),
433 + "<|content_thinking|>" => Some(200008),
434 + "<|end_message|>" => Some(200010),
435 + "<|content_tool_error|>" => Some(200022),
436 + "<|content_invoke_tool_json|>" => Some(200049),
437 + "<|content_invoke_tool_text|>" => Some(200057),
438 + _ => None,
439 + }
440 + }
441 +
442 + fn id_to_token(&self, id: u32) -> Option<String> {
443 + let token = match id {
444 + 200001 => MESSAGE_MODEL,
445 + 200004 => "<|content_text|>",
446 + 200006 => "<|content_model_end_sampling|>",
447 + 200008 => "<|content_thinking|>",
448 + 200010 => "<|end_message|>",
449 + 200022 => "<|content_tool_error|>",
450 + 200049 => "<|content_invoke_tool_json|>",
451 + 200057 => "<|content_invoke_tool_text|>",
452 + _ => return None,
453 + };
454 + Some(token.to_string())
455 + }
456 +
457 + fn is_special_id(&self, token_id: u32) -> bool {
458 + (199999..=200057).contains(&token_id)
459 + }
460 + }
461 +
462 + trait UnifiedParserTestExt {
463 + fn parse_chunk(&mut self, chunk: &str) -> super::Result<UnifiedParserOutput>;
464 + fn parse_complete(&mut self, text: &str) -> super::Result<UnifiedParserOutput>;
465 + }
466 +
467 + impl<T: UnifiedParser + ?Sized> UnifiedParserTestExt for T {
468 + fn parse_chunk(&mut self, chunk: &str) -> super::Result<UnifiedParserOutput> {
469 + let mut output = UnifiedParserOutput::default();
470 + self.parse_into(chunk, &mut output)?;
471 + Ok(output)
472 + }
473 +
474 + fn parse_complete(&mut self, text: &str) -> super::Result<UnifiedParserOutput> {
475 + let mut output = self.parse_chunk(text)?;
476 + output.append(self.finish()?);
477 + Ok(output)
478 + }
479 + }
480 +
481 + trait UnifiedParserOutputTestExt {
482 + fn normal_text(&self) -> String;
483 + fn reasoning_text(&self) -> String;
484 + fn calls(&self) -> Vec<crate::tool::ToolCallDelta>;
485 + }
486 +
487 + impl UnifiedParserOutputTestExt for UnifiedParserOutput {
488 + fn normal_text(&self) -> String {
489 + self.events
490 + .iter()
491 + .filter_map(|event| match event {
492 + UnifiedParserEvent::Text(text) => Some(text.as_str()),
493 + _ => None,
494 + })
495 + .collect()
496 + }
497 +
498 + fn reasoning_text(&self) -> String {
499 + self.events
500 + .iter()
501 + .filter_map(|event| match event {
502 + UnifiedParserEvent::Reasoning(text) => Some(text.as_str()),
503 + _ => None,
504 + })
505 + .collect()
506 + }
507 +
508 + fn calls(&self) -> Vec<crate::tool::ToolCallDelta> {
509 + self.events
510 + .iter()
511 + .filter_map(|event| match event {
512 + UnifiedParserEvent::ToolCall(call) => Some(call.clone()),
513 + _ => None,
514 + })
515 + .collect()
516 + }
517 + }
518 +
519 + fn test_tools() -> Vec<Tool> {
520 + vec![Tool {
521 + name: "get_weather".to_string(),
522 + description: Some("Get weather".to_string()),
523 + parameters: serde_json::json!({
524 + "type": "object",
525 + "properties": {
526 + "city": {"type": "string"}
527 + },
528 + }),
529 + strict: None,
530 + }]
531 + }
532 +
533 + fn test_parser() -> InklingUnifiedParser {
534 + InklingUnifiedParser::new(&test_tools(), Arc::new(FakeTokenizer)).unwrap()
535 + }
536 +
537 + fn collect_stream(chunks: &[&str]) -> UnifiedParserOutput {
538 + let mut parser = test_parser();
539 + let mut output = UnifiedParserOutput::default();
540 + for chunk in chunks {
541 + output.append(parser.parse_chunk(chunk).unwrap());
542 + }
543 + output.append(parser.finish().unwrap());
544 + output
545 + }
546 +
547 + #[test]
548 + fn inkling_streaming_emits_reasoning_then_text() {
549 + let output = collect_stream(&[concat!(
550 + "<|content_thinking|>reason<|end_message|>",
551 + "<|message_model|><|content_text|>answer<|end_message|>",
552 + "<|content_model_end_sampling|>"
553 + )]);
554 +
555 + assert_eq!(output.reasoning_text(), "reason");
556 + assert_eq!(output.normal_text(), "answer");
557 + assert!(output.calls().is_empty());
558 + }
559 +
560 + #[test]
561 + fn inkling_streaming_holds_split_markers() {
562 + let output = collect_stream(&[
563 + "<|content_thin",
564 + "king|>rea",
565 + "son<|end_mes",
566 + "sage|><|message_",
567 + "model|><|content_text|>answer",
568 + "<|end_message|>",
569 + ]);
570 +
571 + assert_eq!(output.reasoning_text(), "reason");
572 + assert_eq!(output.normal_text(), "answer");
573 + }
574 +
575 + #[test]
576 + fn inkling_streaming_accepts_content_blocks_without_message_start() {
577 + let output = collect_stream(&[concat!(
578 + "<|content_thinking|>reason<|end_message|>",
579 + "<|content_text|>answer<|end_message|>",
580 + )]);
581 +
582 + assert_eq!(output.reasoning_text(), "reason");
583 + assert_eq!(output.normal_text(), "answer");
584 + }
585 +
586 + #[test]
587 + fn inkling_tool_json_streams_name_then_argument_deltas() {
588 + let mut parser = test_parser();
589 + let chunks = [
590 + "<|message_model|><|content_invoke_tool_json|>{\"name\":\"get_weather\",\"args\":",
591 + "{\"city\":",
592 + "\"SF\"",
593 + "}}<|end_message|>",
594 + ];
595 +
596 + let mut output = UnifiedParserOutput::default();
597 + let mut observed_args = Vec::new();
598 + for chunk in chunks {
599 + let next = parser.parse_chunk(chunk).unwrap();
600 + observed_args.extend(
601 + next.calls()
602 + .into_iter()
603 + .filter(|call| call.name.is_none())
604 + .map(|call| call.arguments),
605 + );
606 + output.append(next);
607 + }
608 + output.append(parser.finish().unwrap());
609 +
610 + let calls = output.calls();
611 + assert_eq!(calls[0].name.as_deref(), Some("get_weather"));
612 + assert_eq!(observed_args, ["{\"city\":", "\"SF\"", "}"]);
613 + assert_eq!(
614 + calls.iter().map(|call| call.arguments.as_str()).collect::<String>(),
615 + "{\"city\":\"SF\"}"
616 + );
617 + assert!(output.normal_text().is_empty());
618 + }
619 +
620 + #[test]
621 + fn inkling_discards_tool_name_from_message_header() {
622 + let output = collect_stream(&[concat!(
623 + "<|message_model|>get_weather<|content_invoke_tool_json|>",
624 + "{\"name\":\"get_weather\",\"args\":{}}<|end_message|>"
625 + )]);
626 +
627 + assert!(output.normal_text().is_empty());
628 + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
629 + }
630 +
631 + #[test]
632 + fn inkling_streaming_handles_multiple_tool_blocks() {
633 + let output = collect_stream(&[concat!(
634 + "<|content_invoke_tool_json|>{\"name\":\"get_weather\",\"args\":{\"city\":\"SF\"}}<|end_message|>",
635 + "<|message_model|><|content_invoke_tool_json|>{\"name\":\"get_weather\",\"args\":{\"city\":\"NYC\"}}<|end_message|>"
636 + )]);
637 +
638 + let calls = output.calls();
639 + assert_eq!(calls.iter().filter(|call| call.name.is_some()).count(), 2);
640 + assert_eq!(calls[0].tool_index, 0);
641 + assert_eq!(calls[2].tool_index, 1);
642 + }
643 +
644 + #[test]
645 + fn inkling_initialize_open_text_prompt_starts_in_text() {
646 + let mut parser = test_parser();
647 + parser.initialize(&[200004]).unwrap();
648 +
649 + let output = parser.parse_complete("answer<|end_message|>").unwrap();
650 +
651 + assert_eq!(output.normal_text(), "answer");
652 + assert!(output.reasoning_text().is_empty());
653 + }
654 +
655 + #[test]
656 + fn inkling_initialize_open_reasoning_prompt_starts_in_reasoning() {
657 + let mut parser = test_parser();
658 + parser.initialize(&[200008]).unwrap();
659 +
660 + let output = parser.parse_complete("reason<|end_message|>").unwrap();
661 +
662 + assert_eq!(output.reasoning_text(), "reason");
663 + assert!(output.normal_text().is_empty());
664 + }
665 +
666 + #[test]
667 + fn inkling_initialize_model_opener_starts_in_message_header() {
668 + let mut parser = test_parser();
669 + parser.initialize(&[200001]).unwrap();
670 +
671 + let output = parser
672 + .parse_complete(concat!(
673 + "get_weather<|content_invoke_tool_json|>",
674 + "{\"name\":\"get_weather\",\"args\":{}}<|end_message|>"
675 + ))
676 + .unwrap();
677 +
678 + assert!(output.normal_text().is_empty());
679 + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
680 + }
681 +
682 + #[test]
683 + fn inkling_plain_text_falls_through_as_text() {
684 + let output = collect_stream(&["plain ", "answer"]);
685 +
686 + assert_eq!(output.normal_text(), "plain answer");
687 + assert!(output.reasoning_text().is_empty());
688 + }
689 +
690 + #[test]
691 + fn inkling_raw_tool_text_and_tool_error_are_visible_text() {
692 + let output = collect_stream(&[concat!(
693 + "<|content_invoke_tool_text|>search SF<|end_message|>",
694 + "<|content_tool_error|>failed<|end_message|>"
695 + )]);
696 +
697 + assert_eq!(output.normal_text(), "search SFfailed");
698 + assert!(output.calls().is_empty());
699 + }
700 +
701 + #[test]
702 + fn inkling_finish_fails_incomplete_tool_call() {
703 + let mut parser = test_parser();
704 + parser
705 + .parse_chunk("<|content_invoke_tool_json|>{\"name\":\"get_weather\",\"args\":{\"city\"")
706 + .unwrap();
707 +
708 + let error = parser.finish().unwrap_err();
709 +
710 + assert!(error.as_report().to_string().contains("incomplete Inkling tool call"));
711 + }
712 +
713 + #[test]
714 + fn inkling_rejects_non_object_args() {
715 + let mut parser = test_parser();
716 + let error = parser
717 + .parse_chunk("<|content_invoke_tool_json|>{\"name\":\"get_weather\",\"args\":42")
718 + .unwrap_err();
719 +
720 + assert!(error.as_report().to_string().contains("JSON object argument"));
721 + }
722 +
723 + #[test]
724 + fn inkling_missing_token_fails_create() {
725 + struct MissingTokenizer;
726 +
727 + impl Tokenizer for MissingTokenizer {
728 + fn encode(
729 + &self,
730 + _text: &str,
731 + _add_special_tokens: bool,
732 + ) -> vllm_tokenizer::Result<Vec<u32>> {
733 + Ok(vec![])
734 + }
735 +
736 + fn decode(
737 + &self,
738 + _token_ids: &[u32],
739 + _skip_special_tokens: bool,
740 + ) -> vllm_tokenizer::Result<String> {
741 + Ok(String::new())
742 + }
743 +
744 + fn token_to_id(&self, token: &str) -> Option<u32> {
745 + match token {
746 + MESSAGE_MODEL => Some(200001),
747 + CONTENT_TEXT => Some(200004),
748 + _ => None,
749 + }
750 + }
751 +
752 + fn id_to_token(&self, id: u32) -> Option<String> {
753 + match id {
754 + 200001 => Some(MESSAGE_MODEL.to_string()),
755 + 200004 => Some(CONTENT_TEXT.to_string()),
756 + _ => None,
757 + }
758 + }
759 + }
760 +
761 + let error = match InklingUnifiedParser::new(&[], Arc::new(MissingTokenizer)) {
762 + Ok(_) => panic!("expected parser creation to fail"),
763 + Err(error) => error,
764 + };
765 +
766 + assert!(error.as_report().to_string().contains(CONTENT_THINKING));
767 + }
768 +}
rust/src/parser/src/unified/mod.rs
+9
@@ -5,9 +5,11 @@
5
6 mod combined;
7 mod gemma4;
8 +mod inkling;
9
10 pub use combined::CombinedParser;
11 pub use gemma4::Gemma4UnifiedParser;
12 +pub use inkling::InklingUnifiedParser;
13 use thiserror::Error;
14 use thiserror_ext::Macro;
15 use vllm_tokenizer::DynTokenizer;
@@ -208,3 +210,10 @@ pub enum UnifiedParserError {
210 #[error(transparent)]
211 Tool(#[from] ToolParserError),
212 }
213 +
214 +/// Returns the ID for the given token, or an error if it's not found.
215 +fn token_id(tokenizer: &dyn vllm_tokenizer::Tokenizer, token: &str) -> Result<u32> {
216 + tokenizer.token_to_id(token).ok_or_else(|| UnifiedParserError::MissingToken {
217 + token: token.to_string(),
218 + })
219 +}
rust/src/text/src/backend/hf/config.rs
+46 -3
@@ -95,6 +95,7 @@ impl HfSpecialTokens {
95 pub struct ModelConfig {
96 model_type: Option<String>,
97 vocab_size: Option<u32>,
98 + eos_token_id: Option<OneOrManyTokenIds>,
99 num_experts: Option<OneOrManyExpertCount>,
100 moe_num_experts: Option<OneOrManyExpertCount>,
101 n_routed_experts: Option<OneOrManyExpertCount>,
@@ -125,12 +126,16 @@ pub(super) enum OneOrManyTokenIds {
126 }
127
128 impl OneOrManyTokenIds {
128 - pub(super) fn into_set(self) -> BTreeSet<u32> {
129 + pub(super) fn as_slice(&self) -> &[u32] {
130 match self {
130 - Self::One(id) => BTreeSet::from([id]),
131 - Self::Many(ids) => ids.into_iter().collect(),
131 + Self::One(id) => std::slice::from_ref(id),
132 + Self::Many(ids) => ids.as_slice(),
133 }
134 }
135 +
136 + pub(super) fn into_set(self) -> BTreeSet<u32> {
137 + self.as_slice().iter().copied().collect()
138 + }
139 }
140
141 /// Hugging Face configs may expose the expert count either as one integer or
@@ -195,6 +200,18 @@ impl ModelConfig {
200 }
201 }
202
203 + /// Return the effective model-side EOS token ids, following the same
204 + /// simplified text-config selection as `vocab_size`.
205 + pub(super) fn eos_token_ids(&self) -> &[u32] {
206 + if let Some(eos_token_id) = self.eos_token_id.as_ref() {
207 + eos_token_id.as_slice()
208 + } else if let Some(text_config) = self.text_config.as_deref() {
209 + text_config.eos_token_ids()
210 + } else {
211 + &[]
212 + }
213 + }
214 +
215 /// Match Python's current expert-count priority on the selected text
216 /// config.
217 ///
@@ -354,6 +371,32 @@ mod tests {
371 assert_eq!(config.vocab_size().unwrap(), 151936);
372 }
373
374 + #[test]
375 + fn model_config_reads_top_level_eos_token_ids() {
376 + let single: ModelConfig = serde_json::from_str(r#"{"eos_token_id":151645}"#).unwrap();
377 + let many: ModelConfig =
378 + serde_json::from_str(r#"{"eos_token_id":[128001,128008,128009]}"#).unwrap();
379 + let null: ModelConfig = serde_json::from_str(r#"{"eos_token_id":null}"#).unwrap();
380 +
381 + assert_eq!(single.eos_token_ids(), &[151645]);
382 + assert_eq!(many.eos_token_ids(), &[128001, 128008, 128009]);
383 + assert!(null.eos_token_ids().is_empty());
384 + }
385 +
386 + #[test]
387 + fn model_config_uses_nested_eos_token_ids_when_top_level_is_absent() {
388 + let config: ModelConfig = serde_json::from_str(
389 + r#"{
390 + "text_config": {
391 + "eos_token_id": [59246, 59253, 59255]
392 + }
393 + }"#,
394 + )
395 + .unwrap();
396 +
397 + assert_eq!(config.eos_token_ids(), &[59246, 59253, 59255]);
398 + }
399 +
400 #[test]
401 fn model_config_rejects_missing_vocab_size() {
402 let config: ModelConfig = serde_json::from_str(r#"{}"#).unwrap();
rust/src/text/src/backend/hf/mod.rs
+119 -15
@@ -8,7 +8,9 @@ use std::collections::BTreeSet;
8 use std::sync::Arc;
9
10 use tracing::info;
11 -use vllm_tokenizer::{DynTokenizer, HuggingFaceTokenizer, TekkenTokenizer, TiktokenTokenizer};
11 +use vllm_tokenizer::{
12 + DynTokenizer, HuggingFaceTokenizer, TekkenTokenizer, TiktokenTokenizer, Tokenizer,
13 +};
14
15 use self::config::{GenerationConfig, load_generation_config};
16 pub use self::config::{
@@ -56,23 +58,15 @@ impl HfTextBackend {
58 pub fn from_resolved_model_files(files: ResolvedModelFiles, model_id: String) -> Result<Self> {
59 let tokenizer_config = load_tokenizer_config(files.tokenizer_config_path.as_deref())?;
60 let tokenizer = load_tokenizer(&files.tokenizer)?;
59 - let primary_eos_token_id = tokenizer_config
60 - .special_tokens
61 - .eos_token
62 - .as_ref()
63 - .and_then(|token| tokenizer.token_to_id(token.as_str()));
64 -
61 let model_config = load_model_config(files.config_path.as_deref())?;
62 let model_vocab_size = model_config.vocab_size()? as usize;
63 let generation_config = load_generation_config(files.generation_config_path.as_deref())?;
68 - let mut extra_eos_token_ids = generation_config
69 - .eos_token_id
70 - .clone()
71 - .map(|value| value.into_set())
72 - .unwrap_or_default();
73 - if let Some(primary_eos_token_id) = primary_eos_token_id {
74 - extra_eos_token_ids.remove(&primary_eos_token_id);
75 - }
64 + let (primary_eos_token_id, extra_eos_token_ids) = resolve_eos_token_ids(
65 + &tokenizer_config,
66 + &model_config,
67 + &generation_config,
68 + tokenizer.as_ref(),
69 + );
70
71 info!(
72 model_id,
@@ -98,6 +92,42 @@ impl HfTextBackend {
92 }
93 }
94
95 +/// Resolve EOS hints from tokenizer, model, and generation configs.
96 +///
97 +/// Resolution rules:
98 +/// 1. Use the tokenizer-side EOS token as the primary EOS when it resolves to a
99 +/// token id.
100 +/// 2. Fall back to the first model-config EOS id when the tokenizer config does
101 +/// not provide a primary EOS.
102 +/// 3. Keep any remaining model/generation EOS ids as extra stop-token ids so
103 +/// they still participate in stopping and min-token handling.
104 +fn resolve_eos_token_ids(
105 + tokenizer_config: &HfTokenizerConfig,
106 + model_config: &ModelConfig,
107 + generation_config: &GenerationConfig,
108 + tokenizer: &dyn Tokenizer,
109 +) -> (Option<u32>, BTreeSet<u32>) {
110 + let model_config_eos_token_ids = model_config.eos_token_ids();
111 + let primary_eos_token_id = tokenizer_config
112 + .special_tokens
113 + .eos_token
114 + .as_ref()
115 + .and_then(|token| tokenizer.token_to_id(token.as_str()))
116 + .or_else(|| model_config_eos_token_ids.first().copied());
117 +
118 + let mut extra_eos_token_ids = generation_config
119 + .eos_token_id
120 + .clone()
121 + .map(|value| value.into_set())
122 + .unwrap_or_default();
123 + extra_eos_token_ids.extend(model_config_eos_token_ids.iter().copied());
124 + if let Some(primary_eos_token_id) = primary_eos_token_id {
125 + extra_eos_token_ids.remove(&primary_eos_token_id);
126 + }
127 +
128 + (primary_eos_token_id, extra_eos_token_ids)
129 +}
130 +
131 impl TextBackend for HfTextBackend {
132 fn tokenizer(&self) -> DynTokenizer {
133 self.tokenizer.clone()
@@ -128,3 +158,77 @@ impl TextBackend for HfTextBackend {
158 })
159 }
160 }
161 +
162 +#[cfg(test)]
163 +mod tests {
164 + use std::collections::BTreeSet;
165 +
166 + use super::{GenerationConfig, HfTokenizerConfig, ModelConfig, resolve_eos_token_ids};
167 + use vllm_tokenizer::Tokenizer;
168 +
169 + struct FakeTokenizer;
170 +
171 + impl Tokenizer for FakeTokenizer {
172 + fn encode(
173 + &self,
174 + _text: &str,
175 + _add_special_tokens: bool,
176 + ) -> vllm_tokenizer::Result<Vec<u32>> {
177 + Ok(vec![])
178 + }
179 +
180 + fn decode(
181 + &self,
182 + _token_ids: &[u32],
183 + _skip_special_tokens: bool,
184 + ) -> vllm_tokenizer::Result<String> {
185 + Ok(String::new())
186 + }
187 +
188 + fn token_to_id(&self, token: &str) -> Option<u32> {
189 + (token == "</s>").then_some(2)
190 + }
191 +
192 + fn id_to_token(&self, id: u32) -> Option<String> {
193 + (id == 2).then(|| "</s>".to_string())
194 + }
195 + }
196 +
197 + #[test]
198 + fn eos_resolution_uses_model_config_when_tokenizer_has_no_eos() {
199 + let tokenizer_config: HfTokenizerConfig = serde_json::from_str("{}").unwrap();
200 + let model_config: ModelConfig =
201 + serde_json::from_str(r#"{"eos_token_id":[200006,200010]}"#).unwrap();
202 + let generation_config: GenerationConfig = serde_json::from_str("{}").unwrap();
203 +
204 + let (primary, extra) = resolve_eos_token_ids(
205 + &tokenizer_config,
206 + &model_config,
207 + &generation_config,
208 + &FakeTokenizer,
209 + );
210 +
211 + assert_eq!(primary, Some(200006));
212 + assert_eq!(extra, BTreeSet::from([200010]));
213 + }
214 +
215 + #[test]
216 + fn eos_resolution_keeps_tokenizer_eos_primary() {
217 + let tokenizer_config: HfTokenizerConfig =
218 + serde_json::from_str(r#"{"eos_token":"</s>"}"#).unwrap();
219 + let model_config: ModelConfig =
220 + serde_json::from_str(r#"{"eos_token_id":[2,200006]}"#).unwrap();
221 + let generation_config: GenerationConfig =
222 + serde_json::from_str(r#"{"eos_token_id":[2,200010]}"#).unwrap();
223 +
224 + let (primary, extra) = resolve_eos_token_ids(
225 + &tokenizer_config,
226 + &model_config,
227 + &generation_config,
228 + &FakeTokenizer,
229 + );
230 +
231 + assert_eq!(primary, Some(2));
232 + assert_eq!(extra, BTreeSet::from([200006, 200010]));
233 + }
234 +}
setup.py
+15
@@ -445,6 +445,17 @@ class cmake_build_ext(build_ext):
445 dirs_exist_ok=True,
446 )
447
448 + tml_fa4_build = os.path.join(
449 + self.build_lib, "vllm", "third_party", "tml_fa4"
450 + )
451 + if os.path.exists(tml_fa4_build):
452 + print(f"Copying {tml_fa4_build} to vllm/third_party/tml_fa4")
453 + shutil.copytree(
454 + tml_fa4_build,
455 + "vllm/third_party/tml_fa4",
456 + dirs_exist_ok=True,
457 + )
458 +
459
460 class precompiled_build_ext(build_ext):
461 """Disables extension building when using precompiled binaries."""
@@ -803,6 +814,7 @@ class precompiled_wheel_utils:
814 # DeepGEMM: extract all files (.py, .so, .cuh, .h, .hpp, etc.)
815 deep_gemm_regex = re.compile(r"vllm/third_party/deep_gemm/.*")
816 fmha_sm100_regex = re.compile(r"vllm/third_party/fmha_sm100/.*")
817 + tml_fa4_regex = re.compile(r"vllm/third_party/tml_fa4/.*")
818 file_members = []
819 for member in wheel.filelist:
820 if member.filename in exact_members:
@@ -828,6 +840,7 @@ class precompiled_wheel_utils:
840 or triton_kernels_regex.match(member.filename)
841 or flashmla_regex.match(member.filename)
842 or deep_gemm_regex.match(member.filename)
843 + or tml_fa4_regex.match(member.filename)
844 or fmha_sm100_regex.match(member.filename)
845 ):
846 file_members.append(member)
@@ -1146,6 +1159,8 @@ if _is_cuda():
1159 ext_modules.append(CMakeExtension(name="vllm._qutlass_C", optional=True))
1160 # fmha_sm100 is a Python/CuTe-DSL package installed into vllm.third_party.
1161 ext_modules.append(CMakeExtension(name="vllm.fmha_sm100", optional=True))
1162 + # tml-fa4 is copied into an isolated vllm.third_party package.
1163 + ext_modules.append(CMakeExtension(name="vllm.tml_fa4", optional=True))
1164
1165 if _is_cpu():
1166 import platform
tests/kernels/test_ll_bf16_gemm.py
+11 -2
@@ -190,8 +190,8 @@ def test_arbitrary_N_splitk(N):
190
191 @pytest.mark.parametrize(
192 "N,K",
193 - [(256, 7168), (256, 14400), (8, 4096), (384, 7168)],
194 - ids=["DSV3", "DSV4-Flash", "Mixtral", "DSV4-Pro"],
193 + [(256, 7168), (256, 14400), (8, 4096), (384, 7168), (264, 6144)],
194 + ids=["DSV3", "DSV4-Flash", "Mixtral", "DSV4-Pro", "Inkling"],
195 )
196 def test_single_token(N, K):
197 torch.manual_seed(42)
@@ -202,6 +202,15 @@ def test_single_token(N, K):
202 _assert_close(out, _ref(a, b), context=f"M=1 {N}x{K}")
203
204
205 +def test_inkling_max_tokens():
206 + torch.manual_seed(42)
207 + a = torch.randn(64, 6144, dtype=torch.bfloat16, device="cuda")
208 + b = torch.randn(264, 6144, dtype=torch.bfloat16, device="cuda")
209 + out = _gemm(a, b)
210 + assert out.shape == (64, 264)
211 + _assert_close(out, _ref(a, b), context="Inkling M=64")
212 +
213 +
214 # =================================================================
215 # Numerical robustness
216 # =================================================================
tests/models/inkling/test_contract_validation.py new
+50
@@ -0,0 +1,50 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import numpy as np
5 +import pytest
6 +
7 +from vllm.config.compilation import CompilationConfig, CUDAGraphMode
8 +from vllm.models.inkling.common.mm_preprocess import InklingMultiModalDataParser
9 +from vllm.models.inkling.configs import (
10 + InklingAudioConfig,
11 + InklingVisionConfig,
12 +)
13 +from vllm.models.inkling.nvidia.sconv_swa_attn import (
14 + InklingSconvMetadataBuilder,
15 +)
16 +from vllm.v1.attention.backend import AttentionCGSupport
17 +
18 +
19 +@pytest.mark.parametrize(
20 + ("config_cls", "kwargs", "missing"),
21 + [
22 + (InklingAudioConfig, {"decoder_dmodel": 16}, "n_mel_bins"),
23 + (InklingVisionConfig, {"decoder_dmodel": 16}, "vision_encoder_type"),
24 + ],
25 +)
26 +def test_enabled_tower_requires_architecture_fields(config_cls, kwargs, missing):
27 + with pytest.raises(ValueError, match=missing):
28 + config_cls(**kwargs)
29 +
30 +
31 +def test_inkling_raw_2d_audio_is_rejected_as_ambiguous():
32 + parser = InklingMultiModalDataParser(target_sr=16_000, target_channels=1)
33 + with pytest.raises(ValueError, match="ambiguous channel layout"):
34 + parser._parse_audio_data(np.zeros((2, 100), dtype=np.float32))
35 +
36 +
37 +def test_inkling_supports_full_decode_only_cudagraphs():
38 + support = InklingSconvMetadataBuilder.get_cudagraph_support
39 + assert support(None, None) == AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
40 +
41 + compilation_config = CompilationConfig(
42 + cudagraph_mode=CUDAGraphMode.FULL,
43 + splitting_ops=[],
44 + )
45 + resolved_mode = compilation_config.resolve_cudagraph_mode_and_sizes(
46 + AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE,
47 + "InklingSconvBackend",
48 + )
49 +
50 + assert resolved_mode == CUDAGraphMode.FULL_DECODE_ONLY
tests/models/inkling/test_fa4_rel_attention.py new
+381
@@ -0,0 +1,381 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Correctness test for the Inkling FA4 relative-attention score-mod kernel.
4 +
5 +Checks ``inkling_fa4_rel_attention`` against a pure-PyTorch reference that
6 +implements the relative bias exactly as documented in the Inkling architecture
7 +guide::
8 +
9 + logit(i, j, h) = (1 / head_dim) * dot(q[i, h], k[j, h]) + rel_bias(i, j, h)
10 + rel_bias(i, j, h) = rel_logits[i, h, i - j] if 0 <= i - j < rel_extent
11 + = 0 otherwise
12 +
13 +with causal (and optionally sliding-window) masking handled by the backend.
14 +"""
15 +
16 +import pytest
17 +import torch
18 +
19 +from vllm.models.inkling.nvidia.attention import (
20 + InklingAttention,
21 + compute_log_scaling_tau,
22 +)
23 +from vllm.models.inkling.nvidia.ops.fa4_rel_attention import (
24 + inkling_fa4_num_splits,
25 + inkling_fa4_rel_attention,
26 +)
27 +from vllm.platforms import current_platform
28 +from vllm.platforms.interface import DeviceCapability
29 +
30 +_cap = current_platform.get_device_capability() if current_platform.is_cuda() else None
31 +
32 +NUM_HEADS = [(4, 4), (8, 2)] # (num_heads, num_kv_heads)
33 +GLOBAL_REL_EXTENTS = [128, 1024]
34 +LOCAL_REL_EXTENTS = [128, 256]
35 +HEAD_DIM = 128
36 +BLOCK_SIZE = 16
37 +DTYPE = torch.bfloat16
38 +
39 +
40 +def test_log_scaling_tau_matches_reference():
41 + positions = torch.tensor([0, 127999, 128000, 999999], dtype=torch.int64)
42 + actual = compute_log_scaling_tau(positions, 128000, 0.1)
43 + expected = 1.0 + 0.1 * torch.log(
44 + torch.clamp((positions + 1).float() / 128000.0, min=1.0)
45 + )
46 + torch.testing.assert_close(actual, expected, rtol=0, atol=0)
47 +
48 +
49 +def test_split_packed_kv_cache():
50 + attention = InklingAttention.__new__(InklingAttention)
51 + torch.nn.Module.__init__(attention)
52 + attention.head_dim = 8
53 + attention.kv_cache = torch.arange(2 * 3 * 4 * 16).reshape(2, 3, 4, 16)
54 +
55 + key_cache, value_cache = attention._split_kv_cache()
56 +
57 + assert key_cache.shape == value_cache.shape == (2, 4, 3, 8)
58 + torch.testing.assert_close(key_cache, attention.kv_cache[..., :8].transpose(1, 2))
59 + torch.testing.assert_close(value_cache, attention.kv_cache[..., 8:].transpose(1, 2))
60 +
61 +
62 +def test_num_splits_hopper_is_unsplit(monkeypatch):
63 + monkeypatch.setattr(
64 + current_platform,
65 + "get_device_capability",
66 + lambda: DeviceCapability(major=9, minor=0),
67 + )
68 + assert (
69 + inkling_fa4_num_splits(
70 + is_local=False,
71 + batch_size=1,
72 + max_query_len=1,
73 + num_heads=16,
74 + num_kv_heads=2,
75 + max_kv_len=1_048_576,
76 + )
77 + == 1
78 + )
79 +
80 +
81 +@pytest.fixture
82 +def blackwell_platform(monkeypatch):
83 + monkeypatch.setattr(
84 + current_platform,
85 + "get_device_capability",
86 + lambda: DeviceCapability(major=10, minor=0),
87 + )
88 +
89 +
90 +@pytest.mark.parametrize(
91 + ("batch_size", "max_query_len", "expected"),
92 + [
93 + (1, 1, (16, 32, 128, 128)),
94 + (8, 1, (2, 4, 8, 16)),
95 + (32, 1, (1, 1, 2, 4)),
96 + (1, 128, (2, 4, 8, 16)),
97 + (1, 2048, (1, 1, 1, 1)),
98 + ],
99 +)
100 +def test_num_splits_all_tp(blackwell_platform, batch_size, max_query_len, expected):
101 + actual = tuple(
102 + inkling_fa4_num_splits(
103 + is_local=False,
104 + batch_size=batch_size,
105 + max_query_len=max_query_len,
106 + num_heads=64 // tp,
107 + num_kv_heads=8 // tp,
108 + max_kv_len=131072,
109 + )
110 + for tp in (1, 2, 4, 8)
111 + )
112 + assert actual == expected
113 +
114 +
115 +@pytest.mark.parametrize("tp", [1, 2, 4, 8])
116 +def test_num_splits_local_is_unsplit(tp):
117 + assert (
118 + inkling_fa4_num_splits(
119 + is_local=True,
120 + batch_size=1,
121 + max_query_len=1,
122 + num_heads=64 // tp,
123 + num_kv_heads=16 // tp,
124 + max_kv_len=512,
125 + )
126 + == 1
127 + )
128 +
129 +
130 +@pytest.mark.parametrize(
131 + ("max_kv_len", "expected"),
132 + [(8192, 32), (65536, 64), (1048576, 128)],
133 +)
134 +@pytest.mark.parametrize("tp", [4, 8])
135 +def test_num_splits_long_context_bound(blackwell_platform, tp, max_kv_len, expected):
136 + assert (
137 + inkling_fa4_num_splits(
138 + is_local=False,
139 + batch_size=1,
140 + max_query_len=1,
141 + num_heads=64 // tp,
142 + num_kv_heads=8 // tp,
143 + max_kv_len=max_kv_len,
144 + )
145 + == expected
146 + )
147 +
148 +
149 +def _ref_rel_attn(
150 + q: torch.Tensor, # [total_q, H, D]
151 + key_cache: torch.Tensor, # [num_blocks, block, Hkv, D]
152 + value_cache: torch.Tensor,
153 + rel_logits: torch.Tensor, # [total_q, H, rel_extent]
154 + *,
155 + q_lens: list[int],
156 + kv_lens: list[int],
157 + block_table: torch.Tensor,
158 + scale: float,
159 + rel_extent: int,
160 + window_left: int | None,
161 +) -> torch.Tensor:
162 + num_kv_heads = key_cache.shape[2]
163 + num_heads = q.shape[1]
164 + g = num_heads // num_kv_heads
165 + bt = block_table.cpu().numpy()
166 + out = torch.empty_like(q)
167 +
168 + start = 0
169 + for i, (ql, kl) in enumerate(zip(q_lens, kv_lens)):
170 + qi = q[start : start + ql].float() # [ql, H, D]
171 + rl = rel_logits[start : start + ql].float() # [ql, H, rel_extent]
172 +
173 + nblk = (kl + BLOCK_SIZE - 1) // BLOCK_SIZE
174 + blk = bt[i, :nblk]
175 + k = key_cache[blk].reshape(-1, num_kv_heads, HEAD_DIM)[:kl].float()
176 + v = value_cache[blk].reshape(-1, num_kv_heads, HEAD_DIM)[:kl].float()
177 + k = k.repeat_interleave(g, dim=1) # [kl, H, D]
178 + v = v.repeat_interleave(g, dim=1)
179 +
180 + # [H, ql, kl]
181 + scores = torch.einsum("qhd,khd->hqk", qi, k) * scale
182 +
183 + dev = q.device
184 + qpos = torch.arange(ql, device=dev).view(ql, 1) + (kl - ql) # query pos
185 + kpos = torch.arange(kl, device=dev).view(1, kl)
186 + dist = qpos - kpos # [ql, kl] = i - j
187 +
188 + # Relative bias: rel_logits[i, h, dist] when 0 <= dist < rel_extent.
189 + in_rng = (dist >= 0) & (dist < rel_extent) # [ql, kl]
190 + idx = dist.clamp(0, rel_extent - 1)
191 + # gather per head: bias[h, i, j] = rl[i, h, idx[i, j]]
192 + bias = rl.permute(1, 0, 2).gather( # [H, ql, rel_extent]
193 + 2, idx.unsqueeze(0).expand(num_heads, -1, -1)
194 + ) # [H, ql, kl]
195 + bias = torch.where(in_rng.unsqueeze(0), bias, torch.zeros_like(bias))
196 + scores = scores + bias
197 +
198 + mask = dist < 0 # causal
199 + if window_left is not None:
200 + mask = mask | (dist > window_left)
201 + scores.masked_fill_(mask.unsqueeze(0), float("-inf"))
202 +
203 + probs = torch.softmax(scores, dim=-1)
204 + out[start : start + ql] = torch.einsum("hqk,khd->qhd", probs, v).to(q.dtype)
205 + start += ql
206 + return out
207 +
208 +
209 +def _run_case(seq_lens, num_heads, num_kv_heads, rel_extent, window_left, seed=0):
210 + torch.manual_seed(seed)
211 + device = "cuda"
212 + q_lens = [s[0] for s in seq_lens]
213 + kv_lens = [s[1] for s in seq_lens]
214 + total_q = sum(q_lens)
215 + num_seqs = len(seq_lens)
216 +
217 + scale = 1.0 / HEAD_DIM
218 +
219 + # q/k are RMS-normed in the model (unit-ish norm); normalize here so the
220 + # logit magnitudes are realistic and the bias is not numerically dwarfed.
221 + q = torch.randn(total_q, num_heads, HEAD_DIM, device=device, dtype=DTYPE)
222 + q = torch.nn.functional.normalize(q.float(), dim=-1).to(DTYPE)
223 +
224 + # Paged KV cache.
225 + max_blocks = (max(kv_lens) + BLOCK_SIZE - 1) // BLOCK_SIZE
226 + num_blocks = num_seqs * max_blocks + 1
227 + key_cache = torch.randn(
228 + num_blocks, BLOCK_SIZE, num_kv_heads, HEAD_DIM, device=device, dtype=DTYPE
229 + )
230 + key_cache = torch.nn.functional.normalize(key_cache.float(), dim=-1).to(DTYPE)
231 + value_cache = torch.randn(
232 + num_blocks, BLOCK_SIZE, num_kv_heads, HEAD_DIM, device=device, dtype=DTYPE
233 + )
234 +
235 + # Distinct blocks per sequence (block 0 left as a never-referenced pad).
236 + block_table = torch.zeros(num_seqs, max_blocks, dtype=torch.int32, device=device)
237 + for i in range(num_seqs):
238 + block_table[i] = torch.arange(
239 + 1 + i * max_blocks, 1 + (i + 1) * max_blocks, dtype=torch.int32
240 + )
241 +
242 + cu_seqlens_q = torch.tensor(
243 + [0, *torch.cumsum(torch.tensor(q_lens), 0).tolist()],
244 + dtype=torch.int32,
245 + device=device,
246 + )
247 + cache_seqlens = torch.tensor(kv_lens, dtype=torch.int32, device=device)
248 +
249 + rel_logits = torch.randn(total_q, num_heads, rel_extent, device=device, dtype=DTYPE)
250 +
251 + window_size = (-1, -1) if window_left is None else (window_left, 0)
252 +
253 + preallocated_out = torch.empty_like(q)
254 + out = inkling_fa4_rel_attention(
255 + q,
256 + key_cache,
257 + value_cache,
258 + block_table=block_table,
259 + cache_seqlens=cache_seqlens,
260 + cu_seqlens_q=cu_seqlens_q,
261 + max_seqlen_q=max(q_lens),
262 + softmax_scale=scale,
263 + causal=True,
264 + window_size=window_size,
265 + rel_extent=rel_extent,
266 + rel_logits=rel_logits,
267 + out=preallocated_out,
268 + )
269 + assert out.data_ptr() == preallocated_out.data_ptr()
270 + out = out.view(total_q, num_heads, HEAD_DIM)
271 +
272 + ref = _ref_rel_attn(
273 + q,
274 + key_cache,
275 + value_cache,
276 + rel_logits,
277 + q_lens=q_lens,
278 + kv_lens=kv_lens,
279 + block_table=block_table,
280 + scale=scale,
281 + rel_extent=rel_extent,
282 + window_left=window_left,
283 + )
284 +
285 + torch.testing.assert_close(out.float(), ref.float(), atol=2e-2, rtol=2e-2)
286 +
287 +
288 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA")
289 +@pytest.mark.skipif(
290 + _cap is None or _cap.major < 9,
291 + reason="FA4 score-mod requires Hopper+ (SM90+)",
292 +)
293 +@pytest.mark.parametrize("num_heads", NUM_HEADS)
294 +@pytest.mark.parametrize(
295 + "seq_lens",
296 + [
297 + [(64, 64)], # single full prefill
298 + [(64, 64), (33, 33), (17, 17)], # ragged prefill batch
299 + [(512, 512)], # seq_len >> rel_extent (most keys get zero bias)
300 + [(300, 300), (512, 512), (129, 129)], # large ragged batch
301 + ],
302 +)
303 +@pytest.mark.parametrize("rel_extent", GLOBAL_REL_EXTENTS)
304 +@torch.inference_mode()
305 +def test_full_attention(seq_lens, num_heads, rel_extent):
306 + # rel_extent=128 exercises the out-of-range (zero bias) path; 1024 covers all.
307 + # With the 512-token cases and rel_extent=128, query/seq lengths are far
308 + # larger than rel_extent so the vast majority of (i, j) pairs are out of
309 + # range and must contribute zero bias.
310 + _run_case(seq_lens, num_heads[0], num_heads[1], rel_extent, window_left=None)
311 +
312 +
313 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA")
314 +@pytest.mark.skipif(
315 + _cap is None or _cap.major < 9,
316 + reason="FA4 score-mod requires Hopper+ (SM90+)",
317 +)
318 +@pytest.mark.parametrize("num_heads", NUM_HEADS)
319 +@pytest.mark.parametrize(
320 + "seq_lens",
321 + [
322 + [(200, 512)], # chunked prefill: q_len=200 (> rel_extent), 312 cached
323 + [(200, 512), (50, 300), (1, 400)], # mixed chunked + decode
324 + ],
325 +)
326 +@pytest.mark.parametrize("rel_extent", GLOBAL_REL_EXTENTS)
327 +@torch.inference_mode()
328 +def test_chunked_prefill(seq_lens, num_heads, rel_extent):
329 + # q_len < kv_len with q_len itself larger than rel_extent (for the 128 case):
330 + # exercises the seqlen_k - seqlen_q offset together with the out-of-range path.
331 + _run_case(seq_lens, num_heads[0], num_heads[1], rel_extent, window_left=None)
332 +
333 +
334 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA")
335 +@pytest.mark.skipif(
336 + _cap is None or _cap.major < 9,
337 + reason="FA4 score-mod requires Hopper+ (SM90+)",
338 +)
339 +@pytest.mark.parametrize("num_heads", NUM_HEADS)
340 +@pytest.mark.parametrize(
341 + "seq_lens",
342 + [
343 + [(64, 64), (40, 40)], # seq_len > window
344 + [(512, 512), (300, 300)], # seq_len/query_len >> window
345 + [(1, 512)], # decode with kv_len >> window
346 + ],
347 +)
348 +@pytest.mark.parametrize("local_extent", LOCAL_REL_EXTENTS)
349 +@torch.inference_mode()
350 +def test_sliding_window(seq_lens, num_heads, local_extent):
351 + # Local layers use window_size=(local_extent-1, 0) and rel_extent==local_extent.
352 + # With the 512-token cases, query/seq lengths far exceed the window so most
353 + # keys are masked out by the sliding window.
354 + _run_case(
355 + seq_lens,
356 + num_heads[0],
357 + num_heads[1],
358 + rel_extent=local_extent,
359 + window_left=local_extent - 1,
360 + )
361 +
362 +
363 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA")
364 +@pytest.mark.skipif(
365 + _cap is None or _cap.major < 9,
366 + reason="FA4 score-mod requires Hopper+ (SM90+)",
367 +)
368 +@pytest.mark.parametrize("num_heads", NUM_HEADS)
369 +@pytest.mark.parametrize(
370 + "seq_lens",
371 + [
372 + [(1, 50)],
373 + [(1, 50), (1, 7), (1, 200)],
374 + [(1, 512), (1, 333)], # kv_len >> rel_extent
375 + ],
376 +)
377 +@pytest.mark.parametrize("rel_extent", GLOBAL_REL_EXTENTS)
378 +@torch.inference_mode()
379 +def test_decode(seq_lens, num_heads, rel_extent):
380 + # q_len=1 with kv_len>q_len: the score-mod's seqlen_k - seqlen_q offset path.
381 + _run_case(seq_lens, num_heads[0], num_heads[1], rel_extent, window_left=None)
tests/models/inkling/test_fa4_warmup.py new
+69
@@ -0,0 +1,69 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm.models.inkling.nvidia.ops.fa4_rel_attention import (
7 + bucket_max_seqlen_q,
8 + inkling_fa4_num_splits,
9 +)
10 +from vllm.models.inkling.nvidia.ops.fa4_warmup import (
11 + InklingFA4WarmupConfig,
12 + _iter_compile_units,
13 + _num_warps_bucket,
14 +)
15 +
16 +
17 +def test_bucket_max_seqlen_q():
18 + assert [bucket_max_seqlen_q(n) for n in range(1, 10)] == [
19 + 1,
20 + 2,
21 + 4,
22 + 4,
23 + 8,
24 + 8,
25 + 8,
26 + 8,
27 + 16,
28 + ]
29 +
30 +
31 +def test_warmup_enumerates_every_runtime_compile_class():
32 + config = InklingFA4WarmupConfig(
33 + num_heads=16,
34 + num_kv_heads=2,
35 + head_dim=128,
36 + rel_extent=1024,
37 + window_size=(-1, -1),
38 + is_local=False,
39 + max_kv_len=65536,
40 + dtype=torch.bfloat16,
41 + kv_dtype=torch.bfloat16,
42 + block_size=16,
43 + max_num_reqs=64,
44 + max_num_batched_tokens=192,
45 + )
46 + warmed_keys = {unit.key[-1] for unit in _iter_compile_units(config)}
47 +
48 + for query_len in range(1, config.max_num_batched_tokens + 1):
49 + max_seqlen_q = bucket_max_seqlen_q(query_len)
50 + max_num_reqs = min(
51 + config.max_num_reqs,
52 + config.max_num_batched_tokens - query_len + 1,
53 + )
54 + for num_reqs in range(1, max_num_reqs + 1):
55 + num_splits = inkling_fa4_num_splits(
56 + is_local=config.is_local,
57 + batch_size=num_reqs,
58 + max_query_len=max_seqlen_q,
59 + num_heads=config.num_heads,
60 + num_kv_heads=config.num_kv_heads,
61 + max_kv_len=config.max_kv_len,
62 + )
63 + runtime_key = (
64 + max_seqlen_q,
65 + num_splits,
66 + _num_warps_bucket(num_reqs) if num_splits > 1 else None,
67 + num_reqs > 1024,
68 + )
69 + assert runtime_key in warmed_keys
tests/models/inkling/test_mm_towers.py new
+138
@@ -0,0 +1,138 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Correctness tests for the fused Inkling vision/audio tower kernels.
4 +
5 +Each kernel is checked against a pure-PyTorch reference implementing the
6 +towers' original op sequence (native ``rms_norm`` semantics: fp32 variance,
7 +bf16-rounded weight multiply; exact-erf GELU; fp32-accumulated embedding sum).
8 +The fused kernels keep the same accumulation dtype and rounding points, so
9 +outputs differ from the reference only by reduction order — a few bf16 ulps.
10 +"""
11 +
12 +import pytest
13 +import torch
14 +import torch.nn.functional as F
15 +
16 +from vllm.platforms import current_platform
17 +
18 +if not current_platform.is_cuda():
19 + pytest.skip("requires CUDA", allow_module_level=True)
20 +
21 +from vllm.models.inkling.common.towers import fold_timespace_to_depth
22 +from vllm.models.inkling.nvidia.ops.mm_towers import dmel_embed_sum_norm, rmsnorm_gelu
23 +
24 +DTYPE = torch.bfloat16
25 +
26 +
27 +def _bf16_spacing(x: torch.Tensor) -> torch.Tensor:
28 + return torch.exp2(torch.floor(torch.log2(x.float().abs().clamp(min=1e-30)))) * (
29 + 2**-7
30 + )
31 +
32 +
33 +def _assert_close_ulps(
34 + got: torch.Tensor,
35 + ref: torch.Tensor,
36 + max_ulps: float = 4.0,
37 + atol: float = 1e-3,
38 + pre_act: torch.Tensor | None = None,
39 +) -> None:
40 + """Fused vs reference differ only by fp32 reduction order, which shows up
41 + as a few bf16 ulps at any magnitude — a fixed rtol misrepresents that, so
42 + compare against the reference's local ulp spacing. ``pre_act`` (the
43 + reference pre-activation) adds derivative-propagated slack: near
44 + ``gelu(x) ~ 0`` an ulp-level input flip legitimately moves the output by
45 + many of ITS (tiny) ulps, bounded by |gelu'| <= 1.13 times the input ulps."""
46 + g, r = got.float(), ref.float()
47 + tol = torch.clamp(max_ulps * _bf16_spacing(ref), min=atol)
48 + if pre_act is not None:
49 + tol = tol + 2.5 * _bf16_spacing(pre_act)
50 + bad = (g - r).abs() > tol
51 + assert not bad.any(), (
52 + f"{int(bad.sum())}/{ref.numel()} elements beyond tolerance; "
53 + f"max abs diff {(g - r).abs().max().item():.3e}"
54 + )
55 +
56 +
57 +def _ref_rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
58 + # Matches ir.ops.rms_norm: fp32 normalize, bf16 weight multiply.
59 + x32 = x.float()
60 + var = x32.pow(2).mean(dim=-1, keepdim=True)
61 + xn = x32 * torch.rsqrt(var + eps)
62 + return (xn.to(weight.dtype) * weight).to(x.dtype)
63 +
64 +
65 +@pytest.mark.parametrize("rows", [1, 7, 1200, 38400])
66 +@pytest.mark.parametrize("dim", [128, 320, 4800, 6144])
67 +@pytest.mark.parametrize("gelu", [True, False])
68 +def test_rmsnorm_gelu(rows: int, dim: int, gelu: bool) -> None:
69 + torch.manual_seed(rows * dim)
70 + x = torch.randn(rows, dim, device="cuda", dtype=DTYPE)
71 + w = torch.randn(dim, device="cuda", dtype=DTYPE)
72 +
73 + h = _ref_rms_norm(x, w, 1e-5)
74 + ref = F.gelu(h) if gelu else h
75 + got = rmsnorm_gelu(x, w, 1e-5, gelu=gelu)
76 +
77 + _assert_close_ulps(got, ref, pre_act=h if gelu else None)
78 +
79 +
80 +@pytest.mark.parametrize("n", [1, 5, 64])
81 +@pytest.mark.parametrize(
82 + "shape,fold",
83 + [
84 + ((2, 8, 8, 128), (1, 2)), # the real L0 -> L1 vision transition
85 + ((2, 4, 4, 320), (1, 2)),
86 + ((2, 8, 8, 128), (2, 2)), # temporal + spatial fold
87 + ],
88 +)
89 +def test_rmsnorm_gelu_folded_store(
90 + n: int, shape: tuple[int, ...], fold: tuple[int, int]
91 +) -> None:
92 + torch.manual_seed(n)
93 + x = torch.randn(n, *shape, device="cuda", dtype=DTYPE)
94 + w = torch.randn(shape[-1], device="cuda", dtype=DTYPE)
95 +
96 + plain = rmsnorm_gelu(x, w, 1e-5, gelu=True)
97 + ref = fold_timespace_to_depth(plain, *fold)
98 + got = rmsnorm_gelu(x, w, 1e-5, gelu=True, fold=fold)
99 +
100 + # The folded store is a pure permutation of the unfolded output.
101 + assert got.shape == ref.shape
102 + torch.testing.assert_close(got, ref, rtol=0, atol=0)
103 +
104 +
105 +@pytest.mark.parametrize("num_frames", [1, 3, 100, 4097])
106 +@pytest.mark.parametrize("with_norm", [True, False])
107 +def test_dmel_embed_sum_norm(num_frames: int, with_norm: bool) -> None:
108 + torch.manual_seed(num_frames)
109 + n_bins, vocab, dim = 80, 16, 6144
110 + idx = torch.randint(
111 + 0, vocab, (num_frames, n_bins), device="cuda", dtype=torch.int32
112 + )
113 + table = torch.randn(n_bins * vocab, dim, device="cuda", dtype=DTYPE)
114 + norm_w = torch.randn(dim, device="cuda", dtype=DTYPE)
115 +
116 + flat = (torch.arange(n_bins, device="cuda", dtype=torch.int32) * vocab).unsqueeze(
117 + 0
118 + ) + idx
119 + ref = (
120 + F.embedding(flat.reshape(-1).long(), table)
121 + .reshape(num_frames, n_bins, dim)
122 + .sum(dim=1)
123 + )
124 + if with_norm:
125 + ref = _ref_rms_norm(ref, norm_w, 1e-6)
126 +
127 + got = dmel_embed_sum_norm(idx, table, norm_w if with_norm else None, 1e-6)
128 + _assert_close_ulps(got, ref)
129 +
130 +
131 +def test_empty_inputs() -> None:
132 + x = torch.empty(0, 320, device="cuda", dtype=DTYPE)
133 + w = torch.randn(320, device="cuda", dtype=DTYPE)
134 + assert rmsnorm_gelu(x, w, 1e-5).shape == (0, 320)
135 +
136 + idx = torch.empty(0, 80, device="cuda", dtype=torch.int32)
137 + table = torch.randn(1280, 6144, device="cuda", dtype=DTYPE)
138 + assert dmel_embed_sum_norm(idx, table, None, 0.0).shape == (0, 6144)
tests/models/inkling/test_moe_weight_layout.py new
+118
@@ -0,0 +1,118 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from types import SimpleNamespace
5 +
6 +import pytest
7 +import torch
8 +
9 +from vllm.models.inkling.nvidia import moe
10 +from vllm.platforms import current_platform
11 +
12 +
13 +def test_gate_loads_directly_into_padded_runtime_weight() -> None:
14 + gate = moe.InklingGate(
15 + d_model=4,
16 + n_routed_experts=5,
17 + n_shared_experts=2,
18 + experts_per_token=2,
19 + route_scale=1.0,
20 + )
21 + loaded = torch.arange(28, dtype=gate.weight.dtype).reshape(7, 4)
22 +
23 + gate.weight.weight_loader(gate.weight, loaded)
24 +
25 + assert gate.weight.shape == (8, 4)
26 + torch.testing.assert_close(gate.weight[:7], loaded)
27 + torch.testing.assert_close(gate.weight[7], torch.zeros(4))
28 +
29 +
30 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA")
31 +@pytest.mark.parametrize(("num_tokens", "expected_calls"), [(1, 1), (64, 1), (65, 0)])
32 +def test_gate_uses_ll_bf16_gemm_through_token_limit(
33 + monkeypatch, num_tokens, expected_calls
34 +) -> None:
35 + gate = moe.InklingGate(
36 + d_model=8,
37 + n_routed_experts=5,
38 + n_shared_experts=2,
39 + experts_per_token=2,
40 + route_scale=1.0,
41 + ).to(device="cuda", dtype=torch.bfloat16)
42 + hidden_states = torch.randn(num_tokens, 8, device="cuda", dtype=torch.bfloat16)
43 + calls = []
44 +
45 + def fake_ll_bf16_gemm(x, weight):
46 + calls.append((x, weight))
47 + return torch.ones(num_tokens, 8, device="cuda", dtype=torch.float32)
48 +
49 + monkeypatch.setattr(
50 + moe.current_platform, "has_device_capability", lambda capability: True
51 + )
52 + monkeypatch.setattr(moe.ll_bf16, "is_available", lambda: True)
53 + monkeypatch.setattr(moe.ll_bf16, "ll_bf16_gemm", fake_ll_bf16_gemm)
54 +
55 + logits = gate.compute_logits(hidden_states)
56 +
57 + assert len(calls) == expected_calls
58 + if calls:
59 + assert calls[0][0] is hidden_states
60 + assert calls[0][1] is gate.weight
61 + assert logits.shape == (num_tokens, 8)
62 + assert logits.dtype == torch.float32
63 +
64 +
65 +@pytest.mark.parametrize(("projection", "amax"), [("w13", 4.375), ("w2", 2960.0)])
66 +def test_moe_loads_calibrated_input_scale(projection: str, amax: float) -> None:
67 + experts = SimpleNamespace(
68 + w13_input_scale=torch.nn.Parameter(torch.empty(3, 2)),
69 + w2_input_scale=torch.nn.Parameter(torch.empty(3)),
70 + )
71 + layer = SimpleNamespace(experts=SimpleNamespace(routed_experts=experts))
72 +
73 + loaded = moe.InklingMoE.load_expert_weight(
74 + layer,
75 + f"experts.{projection}_weight.input_amax",
76 + torch.tensor([amax]),
77 + )
78 +
79 + scale = getattr(experts, f"{projection}_input_scale")
80 + expected = torch.full_like(scale, amax / (448.0 * 6.0))
81 + torch.testing.assert_close(scale, expected)
82 + assert loaded == [f"experts.routed_experts.{projection}_input_scale"]
83 +
84 +
85 +def test_sink_down_projection_is_packed_during_load(monkeypatch) -> None:
86 + monkeypatch.setattr(moe, "get_tensor_model_parallel_world_size", lambda: 2)
87 + monkeypatch.setattr(moe, "get_tensor_model_parallel_rank", lambda: 1)
88 + sink = moe.InklingSinkExperts(n_experts=2, d_model=3, d_mlp=8)
89 + loaded = torch.arange(48, dtype=sink.w2_weight.dtype).reshape(2, 3, 8)
90 +
91 + sink.load_weight("w2_weight", loaded)
92 +
93 + expected = loaded[:, :, 4:].permute(1, 0, 2).reshape(3, 8)
94 + torch.testing.assert_close(sink.w2_weight, expected)
95 +
96 +
97 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA")
98 +def test_sink_packed_weight_forward_matches_expert_sum(monkeypatch) -> None:
99 + monkeypatch.setattr(moe, "get_tensor_model_parallel_world_size", lambda: 2)
100 + monkeypatch.setattr(moe, "get_tensor_model_parallel_rank", lambda: 1)
101 + sink = moe.InklingSinkExperts(n_experts=2, d_model=3, d_mlp=8).to(
102 + device="cuda", dtype=torch.bfloat16
103 + )
104 + torch.manual_seed(1)
105 + w13 = torch.randn(2, 16, 3, dtype=torch.bfloat16)
106 + w2 = torch.randn(2, 3, 8, dtype=torch.bfloat16)
107 + sink.load_weight("w13_weight", w13)
108 + sink.load_weight("w2_weight", w2)
109 + x = torch.randn(5, 3, device="cuda", dtype=torch.bfloat16)
110 + gammas = torch.randn(5, 2, device="cuda")
111 +
112 + output = sink(x, gammas)
113 +
114 + raw = torch.einsum("td,efd->tef", x, w13[:, 8:].to("cuda"))
115 + hidden = torch.nn.functional.silu(raw[:, :, 0::2].float())
116 + hidden = (hidden * raw[:, :, 1::2] * gammas[:, :, None]).to(torch.bfloat16)
117 + expected = torch.einsum("tef,edf->td", hidden, w2[:, :, 4:].to("cuda"))
118 + torch.testing.assert_close(output, expected, rtol=0, atol=0)
tests/models/inkling/test_qkvr_prep.py new
+258
@@ -0,0 +1,258 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +import vllm.models.inkling.nvidia.ops.sconv as sconv
8 +from vllm.models.inkling.nvidia.ops import qkvr_prep
9 +from vllm.platforms import current_platform
10 +
11 +_cap = current_platform.get_device_capability() if current_platform.is_cuda() else None
12 +
13 +
14 +def _make_inputs(*, is_local: bool, tokens: int = 33, tp_size: int = 4):
15 + torch.manual_seed(0)
16 + heads = 64 // tp_size
17 + kv_heads = (16 if is_local else 8) // tp_size
18 + head_dim = 128
19 + d_rel = 16
20 + rel_extent = 512 if is_local else 1024
21 + page_size = 16
22 + num_blocks = (tokens + page_size - 1) // page_size
23 + q_width = heads * head_dim
24 + kv_width = kv_heads * head_dim
25 + r_width = heads * d_rel
26 + device = "cuda"
27 +
28 + qkvr = torch.randn(
29 + tokens,
30 + q_width + 2 * kv_width + r_width,
31 + device=device,
32 + dtype=torch.bfloat16,
33 + )
34 + k_weight = torch.randn(kv_width, 4, device=device, dtype=torch.bfloat16)
35 + v_weight = torch.randn_like(k_weight)
36 + q_norm_weight = torch.randn(head_dim, device=device, dtype=torch.bfloat16)
37 + k_norm_weight = torch.randn_like(q_norm_weight)
38 + rel_proj = torch.randn(d_rel, rel_extent, device=device, dtype=torch.bfloat16)
39 + conv_cache = torch.zeros(
40 + num_blocks,
41 + kv_heads,
42 + page_size,
43 + 2 * head_dim,
44 + device=device,
45 + dtype=torch.bfloat16,
46 + )
47 + key_cache = torch.empty(
48 + num_blocks,
49 + page_size,
50 + kv_heads,
51 + head_dim,
52 + device=device,
53 + dtype=torch.bfloat16,
54 + )
55 + value_cache = torch.empty_like(key_cache)
56 + positions = torch.arange(tokens, device=device, dtype=torch.int64)
57 + block_table = torch.arange(num_blocks, device=device, dtype=torch.int32)[None]
58 + seq_idx = torch.zeros(tokens, device=device, dtype=torch.int32)
59 + slots = torch.arange(tokens, device=device, dtype=torch.int64)
60 + query_start = torch.zeros(tokens, device=device, dtype=torch.int32)
61 + log_scaling_n_floor = None if is_local else 128000
62 + log_scaling = None
63 + if log_scaling_n_floor is not None:
64 + log_scaling = torch.linspace(
65 + 1.0,
66 + 1.1,
67 + tokens,
68 + device=device,
69 + dtype=torch.float32,
70 + )
71 + return (
72 + qkvr,
73 + k_weight,
74 + v_weight,
75 + q_norm_weight,
76 + k_norm_weight,
77 + rel_proj,
78 + 1e-6,
79 + heads,
80 + kv_heads,
81 + head_dim,
82 + d_rel,
83 + conv_cache,
84 + key_cache,
85 + value_cache,
86 + positions,
87 + block_table,
88 + seq_idx,
89 + slots,
90 + query_start,
91 + slots,
92 + 0,
93 + head_dim,
94 + page_size,
95 + log_scaling,
96 + )
97 +
98 +
99 +def _reference(args):
100 + qkvr = args[0]
101 + heads, kv_heads, head_dim, d_rel = args[7:11]
102 + q_width = heads * head_dim
103 + kv_width = kv_heads * head_dim
104 + q, k, v, r = qkvr.split((q_width, kv_width, kv_width, heads * d_rel), dim=1)
105 + conv_cache = args[11].clone()
106 + k = sconv.fused_sconv(
107 + k.contiguous(),
108 + args[1],
109 + conv_cache,
110 + args[14],
111 + args[15],
112 + args[16],
113 + args[17],
114 + args[18],
115 + args[20],
116 + head_dim,
117 + args[22],
118 + )
119 + v = sconv.fused_sconv(
120 + v.contiguous(),
121 + args[2],
122 + conv_cache,
123 + args[14],
124 + args[15],
125 + args[16],
126 + args[17],
127 + args[18],
128 + args[21],
129 + head_dim,
130 + args[22],
131 + )
132 +
133 + def rms_norm(x, weight):
134 + x = x.reshape(-1, head_dim).float()
135 + rstd = torch.rsqrt(x.square().mean(1, keepdim=True) + args[6])
136 + return (x * rstd * weight.float()).to(qkvr.dtype)
137 +
138 + q = rms_norm(q, args[3]).view(qkvr.shape[0], heads, head_dim)
139 + k = rms_norm(k, args[4]).view(qkvr.shape[0], kv_heads, head_dim)
140 + v = v.view(qkvr.shape[0], kv_heads, head_dim)
141 + rel = torch.mm(r.reshape(-1, d_rel), args[5]).view(qkvr.shape[0], heads, -1)
142 + if args[23] is not None:
143 + q = (q.float() * args[23][:, None, None]).to(q.dtype)
144 + rel = (rel.float() * args[23][:, None, None]).to(rel.dtype)
145 +
146 + key_cache = args[12].clone()
147 + value_cache = args[13].clone()
148 + slots = args[19]
149 + valid = slots >= 0
150 + key_cache.view(-1, kv_heads, head_dim)[slots[valid]] = k[valid]
151 + value_cache.view(-1, kv_heads, head_dim)[slots[valid]] = v[valid]
152 + return q.flatten(1), rel, key_cache, value_cache
153 +
154 +
155 +@pytest.mark.skipif(
156 + _cap is None,
157 + reason="Inkling QKVR prep kernels require CUDA",
158 +)
159 +@pytest.mark.parametrize(
160 + ("is_local", "tokens", "tp_size"),
161 + [
162 + (True, 9, 4),
163 + (False, 33, 4),
164 + (True, 33, 8),
165 + (True, 128, 4),
166 + (False, 128, 8),
167 + (True, 512, 4),
168 + (False, 640, 8),
169 + ],
170 +)
171 +@torch.inference_mode()
172 +def test_qkvr_prep_matches_reference(is_local, tokens, tp_size):
173 + args = _make_inputs(is_local=is_local, tokens=tokens, tp_size=tp_size)
174 + ref_q, ref_rel, ref_key, ref_value = _reference(args)
175 +
176 + q, rel = qkvr_prep.fused_qkvr_prep(*args)
177 +
178 + torch.testing.assert_close(q, ref_q, rtol=0.01, atol=0.02)
179 + torch.testing.assert_close(rel, ref_rel, rtol=0.02, atol=0.125)
180 + torch.testing.assert_close(args[12], ref_key, rtol=0.01, atol=0.01)
181 + torch.testing.assert_close(args[13], ref_value, rtol=0, atol=0)
182 +
183 +
184 +@pytest.mark.skipif(
185 + _cap is None,
186 + reason="Inkling QKVR prep kernels require CUDA",
187 +)
188 +@pytest.mark.parametrize("tokens", [9, 128])
189 +@torch.inference_mode()
190 +def test_qkvr_log_scaling_preserves_bf16_norm_output(tokens):
191 + args = list(_make_inputs(is_local=False, tokens=tokens, tp_size=8))
192 + tau = args[23]
193 + assert tau is not None
194 +
195 + scaled_q, scaled_rel = qkvr_prep.fused_qkvr_prep(*args)
196 + args[23] = None
197 + unscaled_q, unscaled_rel = qkvr_prep.fused_qkvr_prep(*args)
198 +
199 + expected_q = (unscaled_q.float() * tau[:, None]).to(unscaled_q.dtype)
200 + expected_rel = (unscaled_rel.float() * tau[:, None, None]).to(unscaled_rel.dtype)
201 + torch.testing.assert_close(scaled_q, expected_q, rtol=0, atol=0)
202 + torch.testing.assert_close(scaled_rel, expected_rel, rtol=0, atol=0)
203 +
204 +
205 +@pytest.mark.skipif(
206 + _cap is None,
207 + reason="Inkling QKVR prep kernels require CUDA",
208 +)
209 +@pytest.mark.parametrize("tokens", [9, 128])
210 +@torch.inference_mode()
211 +def test_qkvr_prep_negative_slots_skip_cache_writes(tokens):
212 + args = list(_make_inputs(is_local=True, tokens=tokens))
213 + args[17].fill_(-1)
214 + args[19].fill_(-1)
215 + conv_cache = args[11].clone()
216 + key_cache = args[12].clone()
217 + value_cache = args[13].clone()
218 +
219 + qkvr_prep.fused_qkvr_prep(*args)
220 +
221 + torch.testing.assert_close(args[11], conv_cache, rtol=0, atol=0)
222 + torch.testing.assert_close(args[12], key_cache, rtol=0, atol=0)
223 + torch.testing.assert_close(args[13], value_cache, rtol=0, atol=0)
224 +
225 +
226 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
227 +@torch.inference_mode()
228 +def test_fused_sconv_negative_slots_skip_cache_writes():
229 + args = list(_make_inputs(is_local=True, tokens=9))
230 + args[17].fill_(-1)
231 + cache = args[11].clone()
232 +
233 + sconv.fused_sconv(
234 + args[0][:, 16 * 128 : 16 * 128 + 4 * 128].contiguous(),
235 + args[1],
236 + args[11],
237 + args[14],
238 + args[15],
239 + args[16],
240 + args[17],
241 + args[18],
242 + 0,
243 + 128,
244 + args[22],
245 + )
246 +
247 + torch.testing.assert_close(args[11], cache, rtol=0, atol=0)
248 +
249 +
250 +@pytest.mark.parametrize(
251 + ("rel_extent", "last_latency_rows", "first_throughput_rows"),
252 + [(512, 8191, 8192), (1024, 2047, 2048)],
253 +)
254 +def test_rel_projection_schedule_crossover(
255 + rel_extent, last_latency_rows, first_throughput_rows
256 +):
257 + assert not qkvr_prep.use_rel_proj_throughput(last_latency_rows, rel_extent)
258 + assert qkvr_prep.use_rel_proj_throughput(first_throughput_rows, rel_extent)
tests/models/inkling/test_sconv_metadata.py new
+119
@@ -0,0 +1,119 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Parity test for the fused sconv seq-metadata kernel.
4 +
5 +``sconv_seq_metadata`` fills the per-token ``seq_idx`` (owning request) and
6 +``query_start`` (first x-row of that request) buffers in a single launch. Actual
7 +tokens must match the searchsorted-based reference, while CUDA graph padding
8 +must be initialized to safe zero values.
9 +"""
10 +
11 +import pytest
12 +import torch
13 +
14 +from vllm.models.inkling.nvidia.ops.sconv import sconv_seq_metadata
15 +from vllm.models.inkling.nvidia.sconv_swa_attn import InklingSconvMetadataBuilder
16 +
17 +CASES = [
18 + # (query_lens, extra_pad_tokens)
19 + ([1], 0), # bsz1 decode
20 + ([1] * 8, 0), # uniform decode
21 + ([1] * 8, 3), # uniform decode, padded tokens past the last request
22 + ([2] * 4, 0), # uniform spec-decode
23 + ([2048], 0), # single prefill
24 + ([517, 1, 1, 33, 1, 128], 0), # mixed prefill/decode
25 + ([517, 1, 1, 33, 1, 128], 5), # mixed, padded
26 + ([1] * 500, 0), # many requests (deep binary search)
27 +]
28 +
29 +
30 +def _ref(query_start_loc: torch.Tensor, num_reqs: int, num_tokens: int):
31 + cu_seqlens = query_start_loc[: num_reqs + 1].to(torch.int64)
32 + token_idx = torch.arange(num_tokens, device=cu_seqlens.device, dtype=torch.int64)
33 + seq_idx = (torch.searchsorted(cu_seqlens, token_idx, right=True) - 1).clamp(
34 + max=num_reqs - 1
35 + )
36 + return seq_idx.to(torch.int32), cu_seqlens[seq_idx].to(torch.int32)
37 +
38 +
39 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
40 +@pytest.mark.parametrize("query_lens,extra_pad", CASES)
41 +def test_sconv_seq_metadata_matches_searchsorted(query_lens, extra_pad):
42 + device = "cuda"
43 + num_reqs = len(query_lens)
44 + query_start_loc = torch.tensor(
45 + [0] + list(torch.tensor(query_lens).cumsum(0)), dtype=torch.int32
46 + ).to(device)
47 + num_actual_tokens = int(query_start_loc[-1])
48 + num_padded_tokens = num_actual_tokens + extra_pad
49 +
50 + ref_seq, ref_qs = _ref(query_start_loc, num_reqs, num_actual_tokens)
51 +
52 + seq_idx = torch.full((num_padded_tokens,), -1, dtype=torch.int32, device=device)
53 + query_start = torch.full_like(seq_idx, -1)
54 + sconv_seq_metadata(
55 + query_start_loc,
56 + num_reqs,
57 + num_actual_tokens,
58 + seq_idx,
59 + query_start,
60 + num_padded_tokens,
61 + )
62 +
63 + torch.testing.assert_close(seq_idx[:num_actual_tokens], ref_seq, rtol=0, atol=0)
64 + torch.testing.assert_close(query_start[:num_actual_tokens], ref_qs, rtol=0, atol=0)
65 + assert torch.count_nonzero(seq_idx[num_actual_tokens:]) == 0
66 + assert torch.count_nonzero(query_start[num_actual_tokens:]) == 0
67 +
68 +
69 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
70 +def test_sconv_metadata_reuses_padded_static_buffers():
71 + device = torch.device("cuda")
72 + builder = object.__new__(InklingSconvMetadataBuilder)
73 + builder.seq_idx_buffer = torch.empty(8, dtype=torch.int32, device=device)
74 + builder.query_start_buffer = torch.empty(8, dtype=torch.int32, device=device)
75 +
76 + class CommonMetadata:
77 + num_reqs = 1
78 + num_actual_tokens = 8
79 + query_start_loc = torch.tensor([0, 3], dtype=torch.int32, device=device)
80 + query_start_loc_cpu = torch.tensor([0, 3], dtype=torch.int32)
81 + block_table_tensor = torch.zeros((1, 1), dtype=torch.int32, device=device)
82 + slot_mapping = torch.tensor(
83 + [0, 1, 2, -1, -1, -1, -1, -1],
84 + dtype=torch.int64,
85 + device=device,
86 + )
87 +
88 + common = CommonMetadata()
89 + first = builder.build(0, common)
90 + pointers = tuple(
91 + tensor.data_ptr()
92 + for tensor in (
93 + first.block_table,
94 + first.slot_mapping,
95 + first.seq_idx,
96 + first.query_start,
97 + )
98 + )
99 +
100 + common.query_start_loc[1] = 5
101 + common.query_start_loc_cpu[1] = 5
102 + common.slot_mapping[:5] = torch.arange(5, dtype=torch.int64, device=device)
103 + second = builder.build(0, common)
104 +
105 + assert second.slot_mapping.shape == (8,)
106 + assert second.seq_idx.shape == (8,)
107 + assert second.query_start.shape == (8,)
108 + assert pointers == tuple(
109 + tensor.data_ptr()
110 + for tensor in (
111 + second.block_table,
112 + second.slot_mapping,
113 + second.seq_idx,
114 + second.query_start,
115 + )
116 + )
117 + assert torch.count_nonzero(second.seq_idx[5:]) == 0
118 + assert torch.count_nonzero(second.query_start[5:]) == 0
119 + assert torch.all(second.slot_mapping[5:] == -1)
tests/models/registry.py
+12
@@ -330,6 +330,12 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
330 "naver-hyperclovax/HyperCLOVAX-SEED-Think-14B",
331 min_transformers_version="5.9.0",
332 ),
333 + "InklingForCausalLM": _HfExamplesInfo(
334 + "thinkingmachines/Inkling-NVFP4",
335 + tokenizer_mode="inkling",
336 + trust_remote_code=True,
337 + max_model_len=4096,
338 + ),
339 "InternLM2ForCausalLM": _HfExamplesInfo(
340 "internlm/internlm2-chat-7b", trust_remote_code=True
341 ),
@@ -956,6 +962,12 @@ _MULTIMODAL_EXAMPLE_MODELS = {
962 "HuggingFaceM4/Idefics3-8B-Llama3",
963 extras={"tiny": "HuggingFaceTB/SmolVLM-256M-Instruct"},
964 ),
965 + "InklingForConditionalGeneration": _HfExamplesInfo(
966 + "thinkingmachines/Inkling-NVFP4",
967 + tokenizer_mode="inkling",
968 + trust_remote_code=True,
969 + max_model_len=4096,
970 + ),
971 "IsaacForConditionalGeneration": _HfExamplesInfo(
972 "PerceptronAI/Isaac-0.1",
973 trust_remote_code=True,
tests/parser/engine/test_delegating_replay.py
+5
@@ -88,6 +88,11 @@ def _discover_pairings() -> list[_PairingInfo]:
88 if cfg.name not in _BUILDERS:
89 missing_builders.append(f"{engine_cls.__name__} (config.name={cfg.name!r})")
90 continue
91 + if cfg.name == "inkling":
92 + # Inkling uses typed structural blocks and opts out of token-id
93 + # terminal matching; combined-parser replay coverage lives in
94 + # test_inkling.py.
95 + continue
96
97 parser_cls = type(
98 f"_Delegating{engine_cls.__name__}",
tests/parser/engine/test_inkling.py new
+506
@@ -0,0 +1,506 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for the engine-based Inkling parser.
4 +
5 +Inkling output is a sequence of typed content blocks delimited by dedicated
6 +special tokens; the tool-call payload is ``{"name":...,"args":{...}}``
7 +between ``<|content_invoke_tool_json|>`` and ``<|end_message|>``. The
8 +cases mirror the Rust unified parser's tests
9 +(``rust/src/parser/src/unified/inkling.rs``) where applicable.
10 +"""
11 +
12 +import json
13 +
14 +import pytest
15 +
16 +from tests.parser.engine.conftest import make_mock_tokenizer
17 +from tests.parser.engine.streaming_helpers import (
18 + collect_content,
19 + collect_function_name,
20 + collect_tool_arguments,
21 +)
22 +from vllm.parser.engine.parser_engine_config import ParserState
23 +from vllm.parser.inkling import InklingParser, _inkling_arg_converter
24 +from vllm.parser.parser_manager import ParserManager
25 +
26 +MSG_MODEL = "<|message_model|>"
27 +TEXT_START = "<|content_text|>"
28 +THINK_START = "<|content_thinking|>"
29 +TOOL_JSON = "<|content_invoke_tool_json|>"
30 +TOOL_TEXT = "<|content_invoke_tool_text|>"
31 +TOOL_ERROR = "<|content_tool_error|>"
32 +END_MESSAGE = "<|end_message|>"
33 +END_SAMPLING = "<|content_model_end_sampling|>"
34 +
35 +_TML_VOCAB = {
36 + MSG_MODEL: 200001,
37 + TEXT_START: 200004,
38 + END_SAMPLING: 200006,
39 + THINK_START: 200008,
40 + END_MESSAGE: 200010,
41 + TOOL_ERROR: 200022,
42 + TOOL_JSON: 200049,
43 + TOOL_TEXT: 200057,
44 +}
45 +
46 +
47 +@pytest.fixture
48 +def mock_tokenizer():
49 + return make_mock_tokenizer(_TML_VOCAB)
50 +
51 +
52 +@pytest.fixture
53 +def parser(mock_tokenizer):
54 + return InklingParser(mock_tokenizer)
55 +
56 +
57 +def _tool_block(name: str, args: str) -> str:
58 + return f'{TOOL_JSON}{{"name":"{name}","args":{args}}}{END_MESSAGE}'
59 +
60 +
61 +_MARKERS = sorted(_TML_VOCAB, key=len, reverse=True)
62 +
63 +
64 +def _tokenize(text: str) -> list[tuple[int, str]]:
65 + """Tokenize like the real stream: markers are atomic special tokens,
66 + plain text becomes one token per character (matching the mock
67 + tokenizer's ``chr``-based decode)."""
68 + tokens: list[tuple[int, str]] = []
69 + i = 0
70 + while i < len(text):
71 + for marker in _MARKERS:
72 + if text.startswith(marker, i):
73 + tokens.append((_TML_VOCAB[marker], marker))
74 + i += len(marker)
75 + break
76 + else:
77 + tokens.append((ord(text[i]), text[i]))
78 + i += 1
79 + return tokens
80 +
81 +
82 +def _stream(parser, request, text: str, chunk_size: int):
83 + """Stream production-shaped deltas: ``chunk_size`` tokens per delta,
84 + with delta_token_ids covering every token (specials and text)."""
85 + tokens = _tokenize(text)
86 + results = []
87 + previous_text = ""
88 + previous_token_ids: list[int] = []
89 + for start in range(0, len(tokens), chunk_size):
90 + batch = tokens[start : start + chunk_size]
91 + delta_text = "".join(t for _, t in batch)
92 + delta_token_ids = [tid for tid, _ in batch]
93 + current_text = previous_text + delta_text
94 + current_token_ids = previous_token_ids + delta_token_ids
95 + delta = parser.extract_tool_calls_streaming(
96 + previous_text=previous_text,
97 + current_text=current_text,
98 + delta_text=delta_text,
99 + previous_token_ids=tuple(previous_token_ids),
100 + current_token_ids=tuple(current_token_ids),
101 + delta_token_ids=tuple(delta_token_ids),
102 + request=request,
103 + )
104 + results.append((delta, current_text))
105 + previous_text = current_text
106 + previous_token_ids = current_token_ids
107 + finish = parser.finish_streaming()
108 + if finish is not None:
109 + results.append((finish, text))
110 + return results
111 +
112 +
113 +def _stream_text_only(parser, request, text: str, chunk_size: int):
114 + """Stream text-only deltas (no token ids), chunked at arbitrary
115 + character boundaries — exercises the text-lexing fallback path,
116 + including markers split across chunks."""
117 + results = []
118 + previous_text = ""
119 + for start in range(0, len(text), chunk_size):
120 + delta_text = text[start : start + chunk_size]
121 + current_text = previous_text + delta_text
122 + delta = parser.extract_tool_calls_streaming(
123 + previous_text=previous_text,
124 + current_text=current_text,
125 + delta_text=delta_text,
126 + previous_token_ids=(),
127 + current_token_ids=(),
128 + delta_token_ids=(),
129 + request=request,
130 + )
131 + results.append((delta, current_text))
132 + previous_text = current_text
133 + finish = parser.finish_streaming()
134 + if finish is not None:
135 + results.append((finish, text))
136 + return results
137 +
138 +
139 +def _collect_reasoning(results) -> str:
140 + return "".join(d.reasoning for d, _ in results if d and d.reasoning)
141 +
142 +
143 +class TestArgConverter:
144 + def test_complete_wrapper(self):
145 + raw = '{"name":"get_weather","args":{"city":"SF"}}'
146 + assert _inkling_arg_converter(raw, False) == '{"city":"SF"}'
147 +
148 + def test_partial_before_args(self):
149 + assert _inkling_arg_converter('{"name":"get_w', True) == ""
150 +
151 + def test_partial_inside_args(self):
152 + raw = '{"name":"x","args":{"a":1'
153 + assert _inkling_arg_converter(raw, True) == '{"a":1'
154 +
155 + def test_prefix_stability(self):
156 + full = '{"name":"x","args":{"a":{"b":[1,2]},"c":"d"}}'
157 + prev = ""
158 + for end in range(len(full)):
159 + out = _inkling_arg_converter(full[:end], True)
160 + assert out.startswith(prev) or prev.startswith(out) or not prev
161 + if out.startswith(prev):
162 + prev = out
163 +
164 + def test_args_value_appearing_in_name(self):
165 + raw = '{"name":"args","args":{"k":1}}'
166 + assert _inkling_arg_converter(raw, False) == '{"k":1}'
167 +
168 + def test_whitespace_tolerated(self):
169 + raw = '{ "name" : "x" , "args" : {"a": 1} }'
170 + assert _inkling_arg_converter(raw, False) == '{"a": 1}'
171 +
172 + def test_missing_args_defaults_empty(self):
173 + assert _inkling_arg_converter('{"name":"x"}', False) == "{}"
174 +
175 + def test_non_object_args_rejected(self):
176 + with pytest.raises(ValueError, match="JSON object"):
177 + _inkling_arg_converter('{"name":"x","args":[1]}', False)
178 +
179 +
180 +class TestNonStreaming:
181 + def test_plain_text(self, parser, mock_request):
182 + reasoning, content, tools = parser.parse(
183 + f"{TEXT_START}hello world{END_MESSAGE}", mock_request
184 + )
185 + assert reasoning is None
186 + assert content == "hello world"
187 + assert tools is None
188 +
189 + def test_reasoning_text_tool(self, parser, mock_request):
190 + text = (
191 + f"{THINK_START}I should check the weather.{END_MESSAGE}"
192 + f"{MSG_MODEL}{TEXT_START}Let me check.{END_MESSAGE}"
193 + f"{MSG_MODEL}" + _tool_block("get_weather", '{"city":"SF"}')
194 + )
195 + reasoning, content, tools = parser.parse(text, mock_request)
196 + assert reasoning == "I should check the weather."
197 + assert content == "Let me check."
198 + assert [t.name for t in tools] == ["get_weather"]
199 + assert json.loads(tools[0].arguments) == {"city": "SF"}
200 +
201 + def test_tool_header_name_is_not_visible_content(self, parser, mock_request):
202 + text = "get_weather" + _tool_block("get_weather", '{"city":"SF"}')
203 + _, content, tools = parser.parse(text, mock_request)
204 + assert content is None
205 + assert [tool.name for tool in tools] == ["get_weather"]
206 +
207 + def test_parallel_tool_calls(self, parser, mock_request):
208 + text = _tool_block("a", "{}") + MSG_MODEL + _tool_block("b", '{"x":[1,2]}')
209 + _, _, tools = parser.parse(text, mock_request)
210 + assert [t.name for t in tools] == ["a", "b"]
211 + assert json.loads(tools[0].arguments) == {}
212 + assert json.loads(tools[1].arguments) == {"x": [1, 2]}
213 +
214 + def test_nested_args(self, parser, mock_request):
215 + args = '{"q":{"deep":{"list":[{"k":"v"}]}},"s":"a}b"}'
216 + _, _, tools = parser.parse(_tool_block("f", args), mock_request)
217 + assert json.loads(tools[0].arguments) == json.loads(args)
218 +
219 + def test_invoke_tool_text_is_visible_text(self, parser, mock_request):
220 + reasoning, content, tools = parser.parse(
221 + f"{TOOL_TEXT}do something{END_MESSAGE}", mock_request
222 + )
223 + assert content == "do something"
224 + assert tools is None
225 +
226 + def test_tool_error_is_visible_text(self, parser, mock_request):
227 + _, content, tools = parser.parse(f"{TOOL_ERROR}boom{END_MESSAGE}", mock_request)
228 + assert content == "boom"
229 + assert tools is None
230 +
231 + def test_end_sampling_closes_blocks(self, parser, mock_request):
232 + reasoning, content, _ = parser.parse(
233 + f"{THINK_START}hm{END_MESSAGE}{MSG_MODEL}{TEXT_START}hi{END_SAMPLING}",
234 + mock_request,
235 + )
236 + assert reasoning == "hm"
237 + assert content == "hi"
238 +
239 + def test_multiple_reasoning_blocks_concatenate(self, parser, mock_request):
240 + text = (
241 + f"{THINK_START}one{END_MESSAGE}"
242 + f"{MSG_MODEL}{TEXT_START}mid{END_MESSAGE}"
243 + f"{MSG_MODEL}{THINK_START}two{END_MESSAGE}"
244 + )
245 + reasoning, content, _ = parser.parse(text, mock_request)
246 + assert reasoning == "onetwo"
247 + assert content == "mid"
248 +
249 + def test_text_after_tool_call(self, parser, mock_request):
250 + text = _tool_block("f", "{}") + f"{MSG_MODEL}{TEXT_START}done{END_MESSAGE}"
251 + _, content, tools = parser.parse(text, mock_request)
252 + assert [t.name for t in tools] == ["f"]
253 + assert content == "done"
254 +
255 + def test_incomplete_tool_call_at_eos(self, parser, mock_request):
256 + # Engine convention: best-effort with what arrived. (The Rust
257 + # parser instead errors with "incomplete Inkling tool call".)
258 + _, _, tools = parser.parse(
259 + f'{TOOL_JSON}{{"name":"d","args":{{"k":"v"', mock_request
260 + )
261 + assert [t.name for t in tools] == ["d"]
262 +
263 + def test_prose_marker_without_token_ids_is_structural(self, parser, mock_request):
264 + # Inkling opts into text-lexer terminal recognition so held-back
265 + # structural marker text from the detokenizer is still parsed.
266 + _, content, _ = parser.parse(
267 + f"{TEXT_START}see {TEXT_START} token{END_MESSAGE}", mock_request
268 + )
269 + assert content == "see token"
270 +
271 +
272 +class TestStreaming:
273 + @pytest.mark.parametrize("chunk_size", [1, 3, 7, 64, 4096])
274 + def test_chunk_invariance_tool_call(self, mock_tokenizer, mock_request, chunk_size):
275 + parser = InklingParser(mock_tokenizer)
276 + text = f"{TEXT_START}Check this.{END_MESSAGE}{MSG_MODEL}" + _tool_block(
277 + "get_weather", '{"city":"San Francisco"}'
278 + )
279 + results = _stream(parser, mock_request, text, chunk_size)
280 + assert collect_content(results) == "Check this."
281 + assert collect_function_name(results) == "get_weather"
282 + assert json.loads(collect_tool_arguments(results)) == {"city": "San Francisco"}
283 +
284 + @pytest.mark.parametrize("chunk_size", [1, 3, 7, 64])
285 + def test_chunk_invariance_tool_call_text_only(
286 + self, mock_tokenizer, mock_request, chunk_size
287 + ):
288 + # Same case through the text-lexing fallback (no token ids),
289 + # with markers split at arbitrary character boundaries.
290 + parser = InklingParser(mock_tokenizer)
291 + text = f"{TEXT_START}Check this.{END_MESSAGE}{MSG_MODEL}" + _tool_block(
292 + "get_weather", '{"city":"San Francisco"}'
293 + )
294 + results = _stream_text_only(parser, mock_request, text, chunk_size)
295 + assert collect_content(results) == "Check this."
296 + assert collect_function_name(results) == "get_weather"
297 + assert json.loads(collect_tool_arguments(results)) == {"city": "San Francisco"}
298 +
299 + @pytest.mark.parametrize("chunk_size", [1, 5, 11])
300 + def test_chunk_invariance_reasoning(self, mock_tokenizer, mock_request, chunk_size):
301 + parser = InklingParser(mock_tokenizer)
302 + text = (
303 + f"{THINK_START}thinking...{END_MESSAGE}"
304 + f"{MSG_MODEL}{TEXT_START}answer{END_MESSAGE}"
305 + )
306 + results = _stream(parser, mock_request, text, chunk_size)
307 + assert _collect_reasoning(results) == "thinking..."
308 + assert collect_content(results) == "answer"
309 +
310 + def test_split_marker_held_across_chunks(self, parser, mock_request):
311 + # Mirrors Rust `inkling_streaming_holds_split_markers`.
312 + text = f"{TEXT_START}hello{END_MESSAGE}"
313 + results = _stream_text_only(parser, mock_request, text, 9)
314 + assert collect_content(results) == "hello"
315 +
316 + def test_name_streams_before_args_complete(self, parser, mock_request):
317 + # Feed only up to the name's closing quote — the name delta must
318 + # already be emitted before any args arrive.
319 + prefix = f'{TOOL_JSON}{{"name":"get_weather",'
320 + results = _stream(parser, mock_request, prefix, 4096)
321 + assert collect_function_name(results) == "get_weather"
322 +
323 + def test_combined_parser_reasoning_to_tool_handoff_uses_text_markers(
324 + self, mock_tokenizer, mock_request
325 + ):
326 + parser_cls = ParserManager.get_parser(
327 + tool_parser_name="inkling",
328 + reasoning_parser_name="inkling",
329 + enable_auto_tools=True,
330 + )
331 + parser = parser_cls(mock_tokenizer, [])
332 +
333 + first = parser.parse_delta(
334 + THINK_START,
335 + [_TML_VOCAB[THINK_START]],
336 + mock_request,
337 + prompt_token_ids=[_TML_VOCAB[MSG_MODEL]],
338 + finished=False,
339 + )
340 + assert first is None
341 +
342 + second = parser.parse_delta(
343 + "thinking",
344 + [ord(c) for c in "thinking"],
345 + mock_request,
346 + finished=False,
347 + )
348 + assert second is not None
349 + assert second.reasoning == "thinking"
350 +
351 + # Mirrors the DelegatingParser handoff after reasoning closes: the
352 + # tool pass receives reconstructed text that starts at the Inkling
353 + # tool marker, while the token-id slice has already moved past it.
354 + body = (
355 + "get_weather"
356 + f'{TOOL_JSON}{{"name":"get_weather","args":{{"city":"Seattle"}}}}'
357 + f"{END_MESSAGE}"
358 + )
359 + third = parser.parse_delta(
360 + body,
361 + [_TML_VOCAB[END_MESSAGE], _TML_VOCAB[END_SAMPLING]],
362 + mock_request,
363 + finished=True,
364 + )
365 + assert third is not None
366 + assert third.tool_calls
367 + assert third.tool_calls[0].function.name == "get_weather"
368 + assert third.tool_calls[0].function.arguments == '{"city":"Seattle"}'
369 + assert TOOL_JSON not in ((third.content or "") + (third.reasoning or ""))
370 +
371 + def test_streamed_args_are_object_only(self, parser, mock_request):
372 + # The streamed `arguments` must be the bare args object, never
373 + # the `{"name":...}` wrapper.
374 + text = _tool_block("f", '{"a":1}')
375 + results = _stream(parser, mock_request, text, 3)
376 + args = collect_tool_arguments(results)
377 + assert json.loads(args) == {"a": 1}
378 + assert "name" not in args
379 +
380 + @pytest.mark.parametrize("chunk_size", [1, 9])
381 + def test_parallel_calls_streaming(self, mock_tokenizer, mock_request, chunk_size):
382 + parser = InklingParser(mock_tokenizer)
383 + text = _tool_block("a", '{"i":1}') + MSG_MODEL + _tool_block("b", '{"i":2}')
384 + results = _stream(parser, mock_request, text, chunk_size)
385 + indexed: dict[int, dict[str, str]] = {}
386 + for delta, _ in results:
387 + if not (delta and delta.tool_calls):
388 + continue
389 + for tc in delta.tool_calls:
390 + slot = indexed.setdefault(tc.index, {"name": "", "args": ""})
391 + if tc.function and tc.function.name:
392 + slot["name"] = tc.function.name
393 + if tc.function and tc.function.arguments:
394 + slot["args"] += tc.function.arguments
395 + assert indexed[0]["name"] == "a"
396 + assert indexed[1]["name"] == "b"
397 + assert json.loads(indexed[0]["args"]) == {"i": 1}
398 + assert json.loads(indexed[1]["args"]) == {"i": 2}
399 +
400 +
401 +class TestPromptSeededState:
402 + def test_prompt_ending_in_thinking_starts_reasoning(self, parser, mock_request):
403 + parser.adjust_initial_state_from_prompt([200001, _TML_VOCAB[THINK_START]])
404 + assert parser._engine.state == ParserState.REASONING
405 +
406 + def test_prompt_ending_in_text_starts_content(self, parser):
407 + parser.adjust_initial_state_from_prompt([200001, _TML_VOCAB[TEXT_START]])
408 + assert parser._engine.state == ParserState.CONTENT
409 +
410 + def test_generation_prompt_tail_starts_message_header(self, parser):
411 + parser.adjust_initial_state_from_prompt(
412 + [_TML_VOCAB[END_MESSAGE], _TML_VOCAB[MSG_MODEL]]
413 + )
414 + assert parser._engine.state == ParserState.MESSAGE_HEADER
415 +
416 + def test_generation_prompt_header_hides_tool_name(self, parser, mock_request):
417 + text = "get_weather" + _tool_block("get_weather", '{"city":"SF"}')
418 + delta = parser.parse_delta(
419 + text,
420 + [token_id for token_id, _ in _tokenize(text)],
421 + mock_request,
422 + prompt_token_ids=[_TML_VOCAB[END_MESSAGE], _TML_VOCAB[MSG_MODEL]],
423 + finished=True,
424 + )
425 + assert delta is not None
426 + assert delta.content is None
427 + assert delta.tool_calls[0].function.name == "get_weather"
428 +
429 +
430 +class TestToolCallFiltering:
431 + """Inkling equivalents of the generic tool-call-filtering replay tests
432 + (Inkling is excluded from those in test_replay.py: its structural
433 + role/kind tokens and shared block-end token don't fit the generic
434 + reasoning/tool split model)."""
435 +
436 + def test_skip_tool_parsing_round_trip(self, mock_tokenizer, mock_request):
437 + # First pass (reasoning adapter, skip_tool_parsing): reasoning is
438 + # classified as reasoning while tool markup survives in content;
439 + # second pass (tool adapter) re-extracts the calls from it.
440 + text = (
441 + f"{THINK_START}plan{END_MESSAGE}{MSG_MODEL}"
442 + + _tool_block("f", '{"a":1}')
443 + + MSG_MODEL
444 + + _tool_block("g", '{"b":[2]}')
445 + )
446 + first = InklingParser(mock_tokenizer)
447 + first.skip_tool_parsing = True
448 + reasoning, content = first.extract_reasoning(text, mock_request)
449 + assert reasoning == "plan"
450 + assert content.count(TOOL_JSON) == 2
451 +
452 + second = InklingParser(mock_tokenizer)
453 + result = second.extract_tool_calls_from_content(content, mock_request)
454 + assert result.tools_called
455 + assert [tc.function.name for tc in result.tool_calls] == ["f", "g"]
456 + assert json.loads(result.tool_calls[0].function.arguments) == {"a": 1}
457 + assert json.loads(result.tool_calls[1].function.arguments) == {"b": [2]}
458 +
459 + @pytest.fixture
460 + def none_request(self, mock_request):
461 + mock_request.tools = [{"type": "function", "function": {"name": "f"}}]
462 + mock_request.tool_choice = "none"
463 + return mock_request
464 +
465 + def test_tool_choice_none_non_streaming(self, mock_tokenizer, none_request):
466 + parser = InklingParser(mock_tokenizer)
467 + text = (
468 + f"{THINK_START}plan{END_MESSAGE}"
469 + f"{MSG_MODEL}{TEXT_START}visible{END_MESSAGE}"
470 + f"{MSG_MODEL}" + _tool_block("f", '{"a":1}')
471 + )
472 + reasoning, content, tools = parser.parse(text, none_request)
473 + assert reasoning == "plan"
474 + assert content == "visible"
475 + assert not tools
476 +
477 + def test_tool_choice_none_streaming(self, mock_tokenizer, none_request):
478 + parser = InklingParser(mock_tokenizer)
479 + text = f"{TEXT_START}visible{END_MESSAGE}{MSG_MODEL}" + _tool_block(
480 + "f", '{"a":1}'
481 + )
482 + results = _stream(parser, none_request, text, 3)
483 + assert collect_content(results) == "visible"
484 + assert all(not (d and d.tool_calls) for d, _ in results)
485 +
486 +
487 +class TestRegisteredAdapters:
488 + def test_adapters_resolve(self):
489 + from vllm.reasoning import ReasoningParserManager
490 + from vllm.tool_parsers import ToolParserManager
491 +
492 + reasoning_cls = ReasoningParserManager.get_reasoning_parser("inkling")
493 + tool_cls = ToolParserManager.get_tool_parser("inkling")
494 + assert reasoning_cls._parser_engine_cls is InklingParser
495 + assert tool_cls._parser_engine_cls is InklingParser
496 + assert tool_cls.supports_required_and_named is False
497 +
498 + def test_adapter_round_trip(self, mock_tokenizer, mock_request):
499 + from vllm.tool_parsers import ToolParserManager
500 +
501 + tool_cls = ToolParserManager.get_tool_parser("inkling")
502 + adapter = tool_cls(mock_tokenizer)
503 + result = adapter.extract_tool_calls(_tool_block("f", '{"a":1}'), mock_request)
504 + assert result.tools_called
505 + assert result.tool_calls[0].function.name == "f"
506 + assert json.loads(result.tool_calls[0].function.arguments) == {"a": 1}
tests/parser/engine/test_replay.py
+10 -1
@@ -68,6 +68,10 @@ def _discover_parsers() -> list[_ParserInfo]:
68 if cfg.name not in _BUILDERS:
69 missing_builders.append(f"{obj.__name__} (config.name={cfg.name!r})")
70 continue
71 + if cfg.name == "inkling":
72 + # Inkling opts out of token-id terminal matching and has typed
73 + # structural blocks; its replay coverage lives in test_inkling.py.
74 + continue
75 tool_end = cfg.token_id_terminals.get("TOOL_END")
76 if not tool_end:
77 raise RuntimeError(
@@ -299,7 +303,12 @@ _TOOL_CALL_SAMPLES = [
303 (p.parser_cls, s, p.think_end, p.tool_start)
304 for p in _PARSERS
305 for s in p.samples
302 - if s.expected_tool_calls and s.expected_reasoning
306 + if s.expected_tool_calls
307 + and s.expected_reasoning
308 + # Inkling's typed-block format (structural role/kind tokens, shared
309 + # block-end token) doesn't fit the generic reasoning/tool split
310 + # below; its filtering modes are covered in test_inkling.py instead.
311 + and p.name != "inkling"
312 ]
313
314
tests/parser/engine/trace_builder.py
+73
@@ -33,6 +33,7 @@ from vllm.parser.engine.registered_adapters import (
33 DeepSeekV32Parser,
34 Gemma4Parser,
35 Glm47MoeParser,
36 + InklingParser,
37 KimiK2Parser,
38 MinimaxM2Parser,
39 NemotronV3Parser,
@@ -948,6 +949,77 @@ _KIMI_K2_SCENARIOS = [
949 ]
950
951
952 +# ── Inkling (typed content blocks, JSON tool payloads) ───────────────────
953 +
954 +_TML_VOCAB: dict[str, int] = {
955 + "<|message_model|>": 200001,
956 + "<|content_text|>": 200004,
957 + "<|content_model_end_sampling|>": 200006,
958 + "<|content_thinking|>": 200008,
959 + "<|end_message|>": 200010,
960 + "<|content_tool_error|>": 200022,
961 + "<|content_invoke_tool_json|>": 200049,
962 + "<|content_invoke_tool_text|>": 200057,
963 +}
964 +
965 +
966 +def _inkling_block(
967 + segs: list[tuple[str, bool]],
968 + kind_token: str,
969 + body: str,
970 +) -> None:
971 + """Append one Inkling content block; blocks after the first start with
972 + the ``<|message_model|>`` role token (the first block continues the
973 + generation prompt directly)."""
974 + if segs:
975 + segs.append(("<|message_model|>", True))
976 + segs.append((kind_token, True))
977 + if body:
978 + segs.append((body, False))
979 + segs.append(("<|end_message|>", True))
980 +
981 +
982 +def _inkling_segments(scenario: Scenario) -> list[tuple[str, bool]]:
983 + segs: list[tuple[str, bool]] = []
984 + if scenario.reasoning is not None:
985 + _inkling_block(segs, "<|content_thinking|>", scenario.reasoning)
986 + if scenario.tool_calls is not None and not scenario.tool_calls:
987 + _inkling_block(segs, "<|content_invoke_tool_json|>", "")
988 + if scenario.content is not None:
989 + _inkling_block(segs, "<|content_text|>", scenario.content)
990 + if scenario.tool_calls:
991 + for tc in scenario.tool_calls:
992 + args = json.dumps(tc.arguments, ensure_ascii=False, separators=(",", ":"))
993 + payload = f'{{"name":"{tc.name}","args":{args}}}'
994 + _inkling_block(segs, "<|content_invoke_tool_json|>", payload)
995 + return segs
996 +
997 +
998 +def _build_inkling(scenario: Scenario, validate: bool = True) -> Sample:
999 + prompt_token_ids = None
1000 + if scenario.after_tool_response:
1001 + # Prompt ends with a closed tool-response block and the
1002 + # generation-prompt role token.
1003 + prompt_token_ids = [
1004 + _TML_VOCAB["<|end_message|>"],
1005 + _TML_VOCAB["<|message_model|>"],
1006 + ]
1007 + sample = _make_sample(
1008 + sample_id=f"inkling-{scenario.id}",
1009 + description=scenario.description,
1010 + vocab=_TML_VOCAB,
1011 + segments=_inkling_segments(scenario),
1012 + expected_reasoning=scenario.reasoning,
1013 + expected_content=_qwen3_expected_content(scenario),
1014 + expected_tool_calls=_expected_tc(scenario),
1015 + tools=_expected_tools(scenario),
1016 + prompt_token_ids=prompt_token_ids,
1017 + )
1018 + if validate:
1019 + _validate_sample(sample, InklingParser)
1020 + return sample
1021 +
1022 +
1023 # ── Registry and public API ──────────────────────────────────────────
1024
1025 _BUILDERS: dict[str, Any] = {
@@ -960,6 +1032,7 @@ _BUILDERS: dict[str, Any] = {
1032 "glm47_moe": _build_glm47_moe,
1033 "kimi_k2": _build_kimi_k2,
1034 "qwen3": _build_qwen3,
1035 + "inkling": _build_inkling,
1036 }
1037
1038
tests/renderers/test_inkling.py new
+484
@@ -0,0 +1,484 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Golden tests for the native Inkling renderer encoding.
4 +
5 +The message fixtures mirror the Rust renderer's fixture tests
6 +(``rust/src/chat/src/renderer/inkling/tests.rs``) so the two frontends stay
7 +in token-level parity. Image blocks emit the bare ``<|content_image|>`` marker;
8 +``InklingMultiModalProcessor`` inserts the per-patch placeholder run.
9 +"""
10 +
11 +import pytest
12 +
13 +from vllm.renderers.inkling import (
14 + InklingRenderer,
15 + _HfBackedTmlTokenizer,
16 + _resolve_reasoning_effort,
17 +)
18 +from vllm.renderers.inkling_encoding import (
19 + SPECIAL_TOKEN_SPELLINGS,
20 + render_inkling_messages,
21 +)
22 +from vllm.renderers.params import ChatParams
23 +
24 +
25 +@pytest.fixture()
26 +def should_do_global_cleanup_after_test() -> bool:
27 + # These tests touch no distributed or device state; the global
28 + # cleanup fixture is unnecessary (and trips a torch MPS allocator
29 + # assert on macOS dev machines).
30 + return False
31 +
32 +
33 +# One id per special token, mirroring the real Inkling vocab layout; plain
34 +# text encodes one token per character so decoded output is exact.
35 +_SPECIAL_VOCAB = {
36 + "<|message_user|>": 200000,
37 + "<|message_model|>": 200001,
38 + "<|message_system|>": 200002,
39 + "<|message_tool|>": 200003,
40 + "<|content_text|>": 200004,
41 + "<|content_image|>": 200005,
42 + "<|content_model_end_sampling|>": 200006,
43 + "<|content_thinking|>": 200008,
44 + "<|end_message|>": 200010,
45 + "<|content_audio_input|>": 200020,
46 + # The HF vocab spells CONTENT_XML as an unused slot.
47 + "<|unused_200024|>": 200024,
48 + "<|audio_end|>": 200043,
49 + "<|content_invoke_tool_json|>": 200049,
50 +}
51 +
52 +_ID_TO_SPECIAL = {v: k for k, v in _SPECIAL_VOCAB.items()}
53 +
54 +
55 +class FakeHfTokenizer:
56 + def get_vocab(self):
57 + return dict(_SPECIAL_VOCAB)
58 +
59 + def encode(self, text, add_special_tokens=False):
60 + assert not add_special_tokens
61 + return [ord(ch) for ch in text]
62 +
63 +
64 +def decode(token_ids):
65 + return "".join(
66 + _ID_TO_SPECIAL.get(tid, chr(tid) if tid < 200000 else f"<{tid}>")
67 + for tid in token_ids
68 + )
69 +
70 +
71 +@pytest.fixture
72 +def inkling_tokenizer():
73 + return _HfBackedTmlTokenizer(FakeHfTokenizer())
74 +
75 +
76 +def render_text(inkling_tokenizer, messages, **kwargs):
77 + return decode(render_inkling_messages(messages, inkling_tokenizer, **kwargs))
78 +
79 +
80 +class TestRustFixtureParity:
81 + def test_tool_round_trip(self, inkling_tokenizer):
82 + messages = [
83 + {
84 + "role": "assistant",
85 + "reasoning_content": "think",
86 + "content": "answer",
87 + "tool_calls": [
88 + {
89 + "id": "call_1",
90 + "function": {
91 + "name": "get_weather",
92 + "arguments": '{"city":"SF"}',
93 + },
94 + }
95 + ],
96 + },
97 + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"},
98 + ]
99 + assert render_text(
100 + inkling_tokenizer, messages, add_generation_prompt=False
101 + ) == (
102 + "<|message_model|><|content_thinking|>think<|end_message|>"
103 + "<|message_model|><|content_text|>answer<|end_message|>"
104 + "<|message_model|>get_weather<|content_invoke_tool_json|>"
105 + '{"name":"get_weather","args":{"city":"SF"}}<|end_message|>'
106 + "<|content_model_end_sampling|>"
107 + "<|message_tool|>get_weather<|content_text|>sunny<|end_message|>"
108 + )
109 +
110 + def test_tool_declare(self, inkling_tokenizer):
111 + tools = [
112 + {
113 + "type": "function",
114 + "function": {
115 + "name": "get_weather",
116 + "description": "Get weather information",
117 + "parameters": {
118 + "type": "object",
119 + "required": ["city"],
120 + "properties": {"city": {"type": "string"}},
121 + },
122 + },
123 + }
124 + ]
125 + messages = [
126 + {
127 + "role": "developer",
128 + "content": "rules",
129 + "tools": [
130 + {
131 + "type": "function",
132 + "function": {
133 + "name": "local_tool",
134 + "parameters": {"z": 1, "a": {"b": 2}},
135 + },
136 + }
137 + ],
138 + },
139 + {"role": "user", "content": "hi"},
140 + ]
141 + assert render_text(inkling_tokenizer, messages, tools=tools) == (
142 + "<|message_system|>tool_declare<|unused_200024|>"
143 + '[{"description":"Get weather information","name":"get_weather",'
144 + '"parameters":{"properties":{"city":{"type":"string"}},'
145 + '"required":["city"],"type":"object"},"type":"function"},'
146 + '{"description":"","name":"local_tool",'
147 + '"parameters":{"a":{"b":2},"z":1},"type":"function"}]'
148 + "<|end_message|>"
149 + "<|message_system|><|content_text|>rules<|end_message|>"
150 + "<|message_user|><|content_text|>hi<|end_message|>"
151 + "<|message_model|>"
152 + )
153 +
154 + def test_text_image(self, inkling_tokenizer):
155 + messages = [
156 + {
157 + "role": "user",
158 + "content": [
159 + {"type": "text", "text": "look"},
160 + {"type": "image_url", "image_url": "data:image/png;base64,"},
161 + ],
162 + }
163 + ]
164 + # Multimodal preprocessing expands this bare marker after rendering.
165 + assert render_text(inkling_tokenizer, messages) == (
166 + "<|message_user|><|content_text|>look<|end_message|>"
167 + "<|message_user|><|content_image|><|end_message|>"
168 + "<|message_model|>"
169 + )
170 +
171 +
172 +class TestRenderingSemantics:
173 + @pytest.mark.parametrize(
174 + ("value", "expected"),
175 + [
176 + ("none", 0.0),
177 + ("minimal", 0.1),
178 + ("low", 0.2),
179 + ("medium", 0.7),
180 + ("high", 0.9),
181 + ("xhigh", 0.99),
182 + ("max", 0.99),
183 + (None, 0.9),
184 + (0.8, 0.8),
185 + (True, None),
186 + ("invalid", None),
187 + ],
188 + )
189 + def test_resolve_reasoning_effort(self, value, expected):
190 + assert _resolve_reasoning_effort(value) == expected
191 +
192 + def test_generation_prompt_default(self, inkling_tokenizer):
193 + text = render_text(inkling_tokenizer, [{"role": "user", "content": "hi"}])
194 + assert text.endswith("<|end_message|><|message_model|>")
195 +
196 + def test_developer_folds_into_system(self, inkling_tokenizer):
197 + assert (
198 + render_text(
199 + inkling_tokenizer,
200 + [{"role": "developer", "content": "be nice"}],
201 + add_generation_prompt=False,
202 + )
203 + == "<|message_system|><|content_text|>be nice<|end_message|>"
204 + )
205 +
206 + def test_empty_string_content_skipped(self, inkling_tokenizer):
207 + assert (
208 + render_text(
209 + inkling_tokenizer,
210 + [{"role": "user", "content": ""}],
211 + add_generation_prompt=False,
212 + )
213 + == ""
214 + )
215 +
216 + def test_empty_reasoning_skipped(self, inkling_tokenizer):
217 + assert (
218 + render_text(
219 + inkling_tokenizer,
220 + [{"role": "assistant", "reasoning_content": "", "content": "hi"}],
221 + add_generation_prompt=False,
222 + )
223 + == "<|message_model|><|content_text|>hi<|end_message|>"
224 + "<|content_model_end_sampling|>"
225 + )
226 +
227 + def test_reasoning_field(self, inkling_tokenizer):
228 + messages = [
229 + {
230 + "role": "assistant",
231 + "reasoning": "think",
232 + "content": "answer",
233 + }
234 + ]
235 + assert render_text(
236 + inkling_tokenizer, messages, add_generation_prompt=False
237 + ) == (
238 + "<|message_model|><|content_thinking|>think<|end_message|>"
239 + "<|message_model|><|content_text|>answer<|end_message|>"
240 + "<|content_model_end_sampling|>"
241 + )
242 +
243 + def test_audio_part(self, inkling_tokenizer):
244 + messages = [
245 + {
246 + "role": "user",
247 + "content": [{"type": "input_audio", "input_audio": {}}],
248 + }
249 + ]
250 + assert render_text(
251 + inkling_tokenizer, messages, add_generation_prompt=False
252 + ) == ("<|message_user|><|content_audio_input|><|audio_end|><|end_message|>")
253 +
254 + def test_tool_response_name_from_message(self, inkling_tokenizer):
255 + assert (
256 + render_text(
257 + inkling_tokenizer,
258 + [{"role": "tool", "name": "my_tool", "content": "ok"}],
259 + add_generation_prompt=False,
260 + )
261 + == "<|message_tool|>my_tool<|content_text|>ok<|end_message|>"
262 + )
263 +
264 + def test_tool_call_args_object_form(self, inkling_tokenizer):
265 + messages = [
266 + {
267 + "role": "assistant",
268 + "tool_calls": [
269 + {
270 + "id": "x",
271 + "function": {
272 + "name": "f",
273 + # dict-form arguments, unsorted keys
274 + "arguments": {"z": 1, "a": 2},
275 + },
276 + }
277 + ],
278 + }
279 + ]
280 + assert render_text(
281 + inkling_tokenizer, messages, add_generation_prompt=False
282 + ) == (
283 + "<|message_model|>f<|content_invoke_tool_json|>"
284 + '{"name":"f","args":{"a":2,"z":1}}<|end_message|>'
285 + "<|content_model_end_sampling|>"
286 + )
287 +
288 + def test_tool_call_empty_args(self, inkling_tokenizer):
289 + messages = [
290 + {
291 + "role": "assistant",
292 + "tool_calls": [{"id": "x", "function": {"name": "f", "arguments": ""}}],
293 + }
294 + ]
295 + assert render_text(
296 + inkling_tokenizer, messages, add_generation_prompt=False
297 + ) == (
298 + "<|message_model|>f<|content_invoke_tool_json|>"
299 + '{"name":"f","args":{}}<|end_message|>'
300 + "<|content_model_end_sampling|>"
301 + )
302 +
303 + def test_tool_call_non_object_args_rejected(self, inkling_tokenizer):
304 + messages = [
305 + {
306 + "role": "assistant",
307 + "tool_calls": [
308 + {"id": "x", "function": {"name": "f", "arguments": "[1]"}}
309 + ],
310 + }
311 + ]
312 + with pytest.raises(TypeError, match="decode to an object"):
313 + render_inkling_messages(
314 + messages, inkling_tokenizer, add_generation_prompt=False
315 + )
316 +
317 + def test_unsupported_role_rejected(self, inkling_tokenizer):
318 + with pytest.raises(ValueError, match="unsupported Inkling message role"):
319 + render_inkling_messages(
320 + [{"role": "narrator", "content": "hi"}], inkling_tokenizer
321 + )
322 +
323 +
324 +class TestReasoningEffort:
325 + def test_frontend_defaults_to_high(self, inkling_tokenizer):
326 + renderer = InklingRenderer.__new__(InklingRenderer)
327 + renderer._inkling_tokenizer = inkling_tokenizer
328 +
329 + text = decode(
330 + renderer._render(
331 + [
332 + {"role": "system", "content": "rules"},
333 + {"role": "user", "content": "hi"},
334 + ],
335 + ChatParams(),
336 + )
337 + )
338 +
339 + assert text.startswith(
340 + "<|message_system|><|content_text|>rules<|end_message|>"
341 + "<|message_system|><|content_text|>Thinking effort level: 0.9"
342 + "<|end_message|>"
343 + )
344 +
345 + @pytest.mark.parametrize("value", ["none", 0, 0.0, -0.0])
346 + def test_zero_effort_has_one_canonical_spelling(self, inkling_tokenizer, value):
347 + renderer = InklingRenderer.__new__(InklingRenderer)
348 + renderer._inkling_tokenizer = inkling_tokenizer
349 +
350 + text = decode(
351 + renderer._render(
352 + [{"role": "user", "content": "hi"}],
353 + ChatParams(chat_template_kwargs={"reasoning_effort": value}),
354 + )
355 + )
356 +
357 + assert text.startswith(
358 + "<|message_system|><|content_text|>Thinking effort level: 0.0"
359 + "<|end_message|>"
360 + )
361 +
362 + def test_emits_one_effort_after_initial_prefix(self, inkling_tokenizer):
363 + effort_block = (
364 + "<|message_system|><|content_text|>Thinking effort level: 0.7"
365 + "<|end_message|>"
366 + )
367 + text = render_text(
368 + inkling_tokenizer,
369 + [
370 + {"role": "system", "content": "rules"},
371 + {"role": "developer", "content": "policy"},
372 + {"role": "user", "content": "user1"},
373 + {"role": "assistant", "content": "assistant1"},
374 + {"role": "user", "content": "user2"},
375 + ],
376 + tools=[{"type": "function", "function": {"name": "f"}}],
377 + reasoning_effort=0.7,
378 + )
379 +
380 + assert text.count(effort_block) == 1
381 + assert (
382 + text.index("tool_declare")
383 + < text.index("rules")
384 + < text.index("policy")
385 + < text.index(effort_block)
386 + < text.index("user1")
387 + < text.index("assistant1")
388 + < text.index("user2")
389 + )
390 +
391 + def test_renders_after_tool_declare(self, inkling_tokenizer):
392 + tools = [{"type": "function", "function": {"name": "f"}}]
393 + text = render_text(
394 + inkling_tokenizer,
395 + [{"role": "user", "content": "hi"}],
396 + tools=tools,
397 + reasoning_effort=0.8,
398 + )
399 + declare_end = text.index("<|end_message|>") + len("<|end_message|>")
400 + assert text[declare_end:].startswith(
401 + "<|message_system|><|content_text|>Thinking effort level: 0.8"
402 + "<|end_message|>"
403 + )
404 +
405 + @pytest.mark.parametrize(
406 + ("value", "expected"),
407 + [(0.2, "0.2"), (0.7, "0.7"), (0.9, "0.9"), (0.99, "0.99")],
408 + )
409 + def test_formats_at_most_two_decimals(self, inkling_tokenizer, value, expected):
410 + text = render_text(
411 + inkling_tokenizer,
412 + [{"role": "user", "content": "hi"}],
413 + reasoning_effort=value,
414 + )
415 + assert f"Thinking effort level: {expected}<|end_message|>" in text
416 +
417 + def test_absent_emits_no_block(self, inkling_tokenizer):
418 + text = render_text(inkling_tokenizer, [{"role": "user", "content": "hi"}])
419 + assert "Thinking effort" not in text
420 +
421 + @pytest.mark.parametrize("value", [0.9900001, 1, 1.5, -0.1])
422 + def test_out_of_range_rejected(self, inkling_tokenizer, value):
423 + with pytest.raises(ValueError, match="must be in"):
424 + render_inkling_messages(
425 + [{"role": "user", "content": "hi"}],
426 + inkling_tokenizer,
427 + reasoning_effort=value,
428 + )
429 +
430 +
431 +class TestSpecialTokenResolution:
432 + def test_missing_special_token_raises(self):
433 + class IncompleteTokenizer(FakeHfTokenizer):
434 + def get_vocab(self):
435 + vocab = dict(_SPECIAL_VOCAB)
436 + del vocab["<|content_invoke_tool_json|>"]
437 + return vocab
438 +
439 + with pytest.raises(ValueError, match="missing special tokens"):
440 + _HfBackedTmlTokenizer(IncompleteTokenizer())
441 +
442 + def test_semantic_spelling_preferred(self):
443 + class SemanticSpellingTokenizer(FakeHfTokenizer):
444 + def get_vocab(self):
445 + vocab = dict(_SPECIAL_VOCAB)
446 + del vocab["<|unused_200024|>"]
447 + vocab["<|content_xml|>"] = 200024
448 + return vocab
449 +
450 + inkling_tokenizer = _HfBackedTmlTokenizer(SemanticSpellingTokenizer())
451 + tools = [{"type": "function", "function": {"name": "f"}}]
452 + ids = render_inkling_messages(
453 + [{"role": "user", "content": "hi"}], inkling_tokenizer, tools=tools
454 + )
455 + assert 200024 in ids
456 +
457 + def test_all_spellings_covered(self):
458 + # Every semantic token must resolve from the reference vocab.
459 + for token, spellings in SPECIAL_TOKEN_SPELLINGS.items():
460 + assert any(s in _SPECIAL_VOCAB for s in spellings), token
461 +
462 +
463 +def test_render_inkling_messages_direct_protocol():
464 + """The encoding core only needs the structural tokenizer protocol."""
465 +
466 + class ProtocolTokenizer:
467 + def encode_text(self, text):
468 + return [ord(c) for c in text]
469 +
470 + def encode_special(self, token):
471 + return {
472 + "<|message_user|>": 200000,
473 + "<|message_model|>": 200001,
474 + "<|content_text|>": 200004,
475 + "<|content_model_end_sampling|>": 200006,
476 + "<|end_message|>": 200010,
477 + }[token]
478 +
479 + ids = render_inkling_messages(
480 + [{"role": "user", "content": "hi"}],
481 + ProtocolTokenizer(),
482 + add_generation_prompt=True,
483 + )
484 + assert ids == [200000, 200004, ord("h"), ord("i"), 200010, 200001]
tests/test_config.py
+20
@@ -211,6 +211,26 @@ def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1(
211 ),
212 True,
213 ),
214 + (
215 + SimpleNamespace(
216 + model="thinkingmachines/Inkling",
217 + architectures=["InklingForCausalLM"],
218 + runner_type="generate",
219 + is_moe=True,
220 + is_quantized=False,
221 + ),
222 + True,
223 + ),
224 + (
225 + SimpleNamespace(
226 + model="thinkingmachines/Inkling",
227 + architectures=["InklingForConditionalGeneration"],
228 + runner_type="generate",
229 + is_moe=True,
230 + is_quantized=False,
231 + ),
232 + True,
233 + ),
234 (
235 SimpleNamespace(
236 model="mistralai/Mixtral-8x7B-Instruct-v0.1",
vllm/config/model.py
+5 -1
@@ -83,7 +83,9 @@ _REGISTERED_MODEL_CLASS_OVERRIDES: set[tuple[str, str]] = set()
83 RunnerOption = Literal["auto", RunnerType]
84 ConvertType = Literal["none", "embed", "classify"]
85 ConvertOption = Literal["auto", ConvertType]
86 -TokenizerMode = Literal["auto", "hf", "slow", "mistral", "deepseek_v32", "deepseek_v4"]
86 +TokenizerMode = Literal[
87 + "auto", "hf", "slow", "mistral", "deepseek_v32", "deepseek_v4", "inkling"
88 +]
89 ModelDType = Literal["auto", "half", "float16", "bfloat16", "float", "float32"]
90 LogprobsMode = Literal[
91 "raw_logits", "raw_logprobs", "processed_logits", "processed_logprobs"
@@ -628,6 +630,8 @@ class ModelConfig:
630 self.tokenizer_mode = "deepseek_v32"
631 elif arch == "DeepseekV4ForCausalLM":
632 self.tokenizer_mode = "deepseek_v4"
633 + elif arch in ("InklingForCausalLM", "InklingForConditionalGeneration"):
634 + self.tokenizer_mode = "inkling"
635
636 if self.tokenizer_mode != "auto":
637 logger.info(
vllm/config/vllm.py
+3 -1
@@ -68,9 +68,11 @@ logger = init_logger(__name__)
68 DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset(
69 {
70 "DeepseekV2ForCausalLM",
71 - "Qwen2MoeForCausalLM",
71 "GraniteMoeForCausalLM",
72 + "InklingForCausalLM",
73 + "InklingForConditionalGeneration",
74 "LongcatFlashNgramForCausalLM",
75 + "Qwen2MoeForCausalLM",
76 }
77 )
78
vllm/model_executor/models/registry.py
+5
@@ -159,6 +159,11 @@ _TEXT_GENERATION_MODELS = {
159 "vllm.models.minimax_m3",
160 "MiniMaxM3SparseForCausalLM",
161 ),
162 + "InklingForCausalLM": ("vllm.models.inkling", "InklingForCausalLM"),
163 + "InklingForConditionalGeneration": (
164 + "vllm.models.inkling",
165 + "InklingForConditionalGeneration",
166 + ),
167 "Ministral3ForCausalLM": ("mistral", "MistralForCausalLM"),
168 "MistralForCausalLM": ("mistral", "MistralForCausalLM"),
169 "MistralLarge3ForCausalLM": ("mistral_large_3", "MistralLarge3ForCausalLM"),
vllm/model_executor/warmup/kernel_warmup.py
+1
@@ -43,6 +43,7 @@ if TYPE_CHECKING:
43 logger = init_logger(__name__)
44
45 _LL_BF16_WARMUP_MODEL_SHAPES: tuple[tuple[int, int], ...] = (
46 + (6144, 264), # Inkling
47 (7168, 256), # DSV3
48 (7168, 384), # DSV4-Pro
49 (14400, 256), # DSV4-Flash
vllm/models/inkling/__init__.py new
+22
@@ -0,0 +1,22 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import TYPE_CHECKING
4 +
5 +if TYPE_CHECKING:
6 + from .nvidia.model import (
7 + InklingForCausalLM,
8 + InklingForConditionalGeneration,
9 + )
10 +
11 +__all__ = [
12 + "InklingForConditionalGeneration",
13 + "InklingForCausalLM",
14 +]
15 +
16 +
17 +def __getattr__(name: str):
18 + if name in __all__:
19 + from .nvidia import model
20 +
21 + return getattr(model, name)
22 + raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
vllm/models/inkling/common/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/inkling/common/mm_preprocess.py new
+361
@@ -0,0 +1,361 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inkling multimodal preprocessing."""
4 +
5 +from __future__ import annotations
6 +
7 +from collections.abc import Iterable, Mapping, Sequence
8 +from typing import Any, cast
9 +
10 +import numpy as np
11 +import regex as re
12 +import torch
13 +from transformers.feature_extraction_utils import BatchFeature
14 +
15 +from vllm.config.multimodal import (
16 + AudioDummyOptions,
17 + BaseDummyOptions,
18 + ImageDummyOptions,
19 +)
20 +from vllm.inputs import MultiModalDataDict
21 +from vllm.multimodal.inputs import (
22 + MultiModalFieldConfig,
23 + MultiModalKwargsItems,
24 +)
25 +from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
26 +from vllm.multimodal.processing import (
27 + BaseDummyInputsBuilder,
28 + BaseMultiModalProcessor,
29 + BaseProcessingInfo,
30 + PromptReplacement,
31 + PromptUpdate,
32 + PromptUpdateDetails,
33 +)
34 +from vllm.transformers_utils.processors.inkling import (
35 + AUDIO_MARKER_ID,
36 + AUDIO_TOKEN_ID,
37 + IMAGE_MARKER_ID,
38 + IMAGE_TOKEN_ID,
39 + InklingAudioFeatureExtractor,
40 + InklingImageProcessor,
41 + InklingProcessor,
42 +)
43 +
44 +from ..configs import InklingMMConfig
45 +
46 +# Maximum audio tokens accepted per clip. At the dMel rate of 20 tokens/s
47 +# (50 ms hop) this is ~10 minutes of audio. It bounds the persistent per-request
48 +# buffers and the encoder/memory budget; longer clips are rejected up front.
49 +MAX_AUDIO_TOKENS = 12_000
50 +
51 +
52 +class InklingMultiModalDataParser(MultiModalDataParser):
53 + def _parse_audio_data(self, data: Any) -> Any:
54 + if isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 2:
55 + raise ValueError(
56 + "Inkling raw 2-D audio has an ambiguous channel layout. "
57 + "Provide encoded audio or a list of mono waveforms."
58 + )
59 + return super()._parse_audio_data(data)
60 +
61 +
62 +def inkling_vision_enabled(config: InklingMMConfig) -> bool:
63 + return getattr(config.vision_config, "decoder_dmodel", None) is not None
64 +
65 +
66 +def inkling_audio_enabled(config: InklingMMConfig) -> bool:
67 + return getattr(config.audio_config, "decoder_dmodel", None) is not None
68 +
69 +
70 +class InklingProcessingInfo(BaseProcessingInfo):
71 + def get_hf_config(self) -> InklingMMConfig:
72 + return self.ctx.get_hf_config(InklingMMConfig)
73 +
74 + def get_hf_processor(self, **kwargs: object) -> InklingProcessor:
75 + config = self.get_hf_config()
76 + vision_config = config.vision_config
77 + audio_config = config.audio_config
78 +
79 + image_processor = InklingImageProcessor(
80 + patch_size=getattr(vision_config, "patch_size", None) or 40,
81 + )
82 +
83 + if inkling_audio_enabled(config):
84 + audio_params = {
85 + "n_mels": audio_config.n_mel_bins,
86 + "num_dmel_bins": audio_config.mel_vocab_size,
87 + "dmel_min_value": audio_config.dmel_min_value,
88 + "dmel_max_value": audio_config.dmel_max_value,
89 + }
90 + else:
91 + audio_params = {}
92 + audio_extractor = InklingAudioFeatureExtractor(params=audio_params)
93 +
94 + return InklingProcessor(
95 + image_processor=image_processor,
96 + audio_feature_extractor=audio_extractor,
97 + tokenizer=self.get_tokenizer(),
98 + )
99 +
100 + def get_data_parser(self) -> MultiModalDataParser:
101 + # Audio inputs must be resampled to the dMel feature extractor's rate
102 + # before `process_audios` (see InklingAudioFeatureExtractor._decode_one).
103 + # Without a target_sr the default parser raises on any audio input.
104 + extractor = self.get_hf_processor().audio_feature_extractor
105 + return InklingMultiModalDataParser(
106 + target_sr=extractor.params.sample_rate,
107 + target_channels=1,
108 + expected_hidden_size=self._get_expected_hidden_size(),
109 + )
110 +
111 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:
112 + config = self.get_hf_config()
113 + limits: dict[str, int | None] = {}
114 + if inkling_vision_enabled(config):
115 + limits["image"] = None
116 + if inkling_audio_enabled(config):
117 + limits["audio"] = None
118 + return limits
119 +
120 + def get_mm_max_tokens_per_item(
121 + self,
122 + seq_len: int,
123 + mm_counts: Mapping[str, int],
124 + ) -> Mapping[str, int] | None:
125 + # Let vLLM profile dummy inputs to determine the max token counts; the
126 + # image patch count is data-dependent, and the dummy audio is sized to
127 + # MAX_AUDIO_TOKENS so audio is profiled/budgeted at its allowed maximum.
128 + return None
129 +
130 +
131 +class InklingDummyInputsBuilder(BaseDummyInputsBuilder[InklingProcessingInfo]):
132 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
133 + # One placeholder per media item; the processor expands each into N
134 + # copies once feature row counts are known.
135 + num_images = mm_counts.get("image", 0)
136 + num_audios = mm_counts.get("audio", 0)
137 + # Use spellings the renderer would emit; tokenization is bypassed in
138 + # _call_hf_processor (we build input_ids directly), so the exact text
139 + # only needs to be a stable per-item marker.
140 + return ("<|content_image|>" * num_images) + (
141 + "<|content_audio_input|>" * num_audios
142 + )
143 +
144 + def get_dummy_mm_data(
145 + self,
146 + seq_len: int,
147 + mm_counts: Mapping[str, int],
148 + mm_options: Mapping[str, BaseDummyOptions],
149 + ) -> MultiModalDataDict:
150 + config = self.info.get_hf_config()
151 + num_images = mm_counts.get("image", 0)
152 + num_audios = mm_counts.get("audio", 0)
153 +
154 + mm_data: dict[str, Any] = {}
155 + if num_images:
156 + patch_size = getattr(config.vision_config, "patch_size", 40)
157 + # A square image ~4 patches wide so the dummy emits several patches.
158 + side = patch_size * 4
159 + image_overrides = mm_options.get("image")
160 + mm_data["image"] = self._get_dummy_images(
161 + width=side,
162 + height=side,
163 + num_images=num_images,
164 + overrides=cast(ImageDummyOptions | None, image_overrides),
165 + )
166 + if num_audios:
167 + # Size the dummy at the maximum allowed audio so memory/encoder
168 + # budgeting reflects the largest clip we accept (MAX_AUDIO_TOKENS).
169 + params = self.info.get_hf_processor().audio_feature_extractor.params
170 + hop = int(round(params.audio_token_duration_s * params.sample_rate))
171 + audio_len = MAX_AUDIO_TOKENS * hop
172 + audio_overrides = mm_options.get("audio")
173 + mm_data["audio"] = self._get_dummy_audios(
174 + length=audio_len,
175 + num_audios=num_audios,
176 + overrides=cast(AudioDummyOptions | None, audio_overrides),
177 + )
178 + return mm_data
179 +
180 +
181 +class InklingMultiModalProcessor(BaseMultiModalProcessor[InklingProcessingInfo]):
182 + def _hf_processor_applies_updates(
183 + self,
184 + prompt_text: str,
185 + mm_items: MultiModalDataItems,
186 + hf_processor_mm_kwargs: Mapping[str, object],
187 + tokenization_kwargs: Mapping[str, object],
188 + ) -> bool:
189 + return False
190 +
191 + def _call_hf_processor(
192 + self,
193 + prompt: str,
194 + mm_data: Mapping[str, object],
195 + mm_kwargs: Mapping[str, object],
196 + tok_kwargs: Mapping[str, object],
197 + ) -> BatchFeature:
198 + # Inkling is not a standard HF processor (no fused text+mm call), so we run
199 + # the vendored extractors ourselves and tokenize the text separately.
200 + # The MM placeholders in `prompt` are expanded later by the prompt
201 + # updates, so here we emit ONE placeholder id per media item.
202 + processor = self.info.get_hf_processor(**mm_kwargs)
203 + tokenizer = self.info.get_tokenizer()
204 +
205 + images = mm_data.get("images") or []
206 + audios = mm_data.get("audios") or []
207 + if not isinstance(images, list):
208 + images = list(cast(Iterable[Any], images))
209 + if not isinstance(audios, list):
210 + audios = list(cast(Iterable[Any], audios))
211 +
212 + prompt_ids = self._tokenize_with_placeholders(
213 + prompt, tokenizer, len(images), len(audios)
214 + )
215 +
216 + data: dict[str, Any] = {"input_ids": [prompt_ids]}
217 +
218 + if images:
219 + img_feat = processor.process_images(images)
220 + data["pixel_values"] = img_feat["vision_patches_bthwc"]
221 + data["num_patches"] = torch.tensor(
222 + img_feat["num_patches"], dtype=torch.int64
223 + )
224 +
225 + if audios:
226 + aud_feat = processor.process_audios(audios)
227 + per_clip = aud_feat["dmel_bins"]
228 + num_audio_tokens = aud_feat["num_audio_tokens"]
229 + for i, n in enumerate(num_audio_tokens):
230 + if int(n) > MAX_AUDIO_TOKENS:
231 + raise ValueError(
232 + f"Audio clip {i} produces {int(n)} tokens, exceeding the "
233 + f"maximum of {MAX_AUDIO_TOKENS} (~10 min at 20 tokens/s). "
234 + "Provide a shorter clip."
235 + )
236 + if per_clip:
237 + input_audio_features = torch.cat(
238 + [torch.as_tensor(c) for c in per_clip], dim=0
239 + )
240 + else:
241 + input_audio_features = torch.empty(0)
242 + data["input_audio_features"] = input_audio_features
243 + data["num_audio_tokens"] = torch.tensor(num_audio_tokens, dtype=torch.int64)
244 +
245 + return BatchFeature(data=data, tensor_type=None)
246 +
247 + def _tokenize_with_placeholders(
248 + self,
249 + prompt: str,
250 + tokenizer: Any,
251 + num_images: int,
252 + num_audios: int,
253 + ) -> list[int]:
254 + """Tokenize `prompt`, emitting the block-start marker id per media item.
255 +
256 + Each marker (kept verbatim) is later expanded by ``_get_prompt_updates``
257 + into ``<marker> + <placeholder> * N``.
258 + """
259 + image_marker = "<|content_image|>"
260 + audio_marker = "<|content_audio_input|>"
261 +
262 + pattern = f"({re.escape(image_marker)}|{re.escape(audio_marker)})"
263 + chunks = re.split(pattern, prompt)
264 +
265 + ids: list[int] = []
266 + seen_img = seen_aud = 0
267 + for chunk in chunks:
268 + if chunk == image_marker:
269 + ids.append(IMAGE_MARKER_ID)
270 + seen_img += 1
271 + elif chunk == audio_marker:
272 + ids.append(AUDIO_MARKER_ID)
273 + seen_aud += 1
274 + elif chunk:
275 + ids.extend(tokenizer.encode(chunk, add_special_tokens=False))
276 +
277 + # Reconcile against the declared media counts only when media is
278 + # present. With no media items -- e.g. the base text-only tokenization
279 + # probe (``_apply_hf_processor_text_only``), which calls this via
280 + # ``_call_hf_processor`` with empty ``mm_data`` -- emit the markers
281 + # verbatim; the marker<->item correspondence is enforced later by
282 + # ``_get_prompt_updates`` once the media features are available.
283 + if num_images or num_audios:
284 + # Fail clearly on a placeholder/media-count mismatch instead of
285 + # crashing with an IndexError deep in the per-item replacement logic.
286 + if num_images and seen_img != num_images:
287 + raise ValueError(
288 + f"Prompt contains {seen_img} image placeholder(s), but only "
289 + f"{num_images} image(s) were provided."
290 + )
291 + if num_audios and seen_aud != num_audios:
292 + raise ValueError(
293 + f"Prompt contains {seen_aud} audio placeholder(s), but only "
294 + f"{num_audios} audio input(s) were provided."
295 + )
296 + return ids
297 +
298 + def _get_mm_fields_config(
299 + self,
300 + hf_inputs: BatchFeature,
301 + hf_processor_mm_kwargs: Mapping[str, object],
302 + ) -> Mapping[str, MultiModalFieldConfig]:
303 + num_patches = hf_inputs.get("num_patches", torch.empty(0, dtype=torch.int64))
304 + num_audio_tokens = hf_inputs.get(
305 + "num_audio_tokens", torch.empty(0, dtype=torch.int64)
306 + )
307 + return dict(
308 + # Ragged per-image patches, grouped by num_patches.
309 + pixel_values=MultiModalFieldConfig.flat_from_sizes("image", num_patches),
310 + num_patches=MultiModalFieldConfig.batched("image"),
311 + # Ragged per-audio frames, grouped by num_audio_tokens.
312 + input_audio_features=MultiModalFieldConfig.flat_from_sizes(
313 + "audio", num_audio_tokens
314 + ),
315 + num_audio_tokens=MultiModalFieldConfig.batched("audio"),
316 + )
317 +
318 + def _get_prompt_updates(
319 + self,
320 + mm_items: Any,
321 + hf_processor_mm_kwargs: Mapping[str, Any],
322 + out_mm_kwargs: MultiModalKwargsItems,
323 + ) -> Sequence[PromptUpdate]:
324 + out_mm_data = out_mm_kwargs.get_data()
325 + num_patches: Any = out_mm_data.get("num_patches")
326 + num_audio_tokens: Any = out_mm_data.get("num_audio_tokens")
327 +
328 + # Keep the block-start marker and append N placeholder tokens after it;
329 + # only the placeholder positions are flagged as embeddings (is_embed), so
330 + # the marker stays a normal text token while the tower features scatter
331 + # into the placeholders.
332 + def image_replacement(item_idx: int) -> PromptUpdateDetails:
333 + n = int(num_patches[item_idx])
334 + return PromptUpdateDetails.select_token_id(
335 + [IMAGE_MARKER_ID] + [IMAGE_TOKEN_ID] * n, IMAGE_TOKEN_ID
336 + )
337 +
338 + def audio_replacement(item_idx: int) -> PromptUpdateDetails:
339 + n = int(num_audio_tokens[item_idx])
340 + return PromptUpdateDetails.select_token_id(
341 + [AUDIO_MARKER_ID] + [AUDIO_TOKEN_ID] * n, AUDIO_TOKEN_ID
342 + )
343 +
344 + updates: list[PromptUpdate] = []
345 + if num_patches is not None and len(num_patches) > 0:
346 + updates.append(
347 + PromptReplacement(
348 + modality="image",
349 + target=[IMAGE_MARKER_ID],
350 + replacement=image_replacement,
351 + )
352 + )
353 + if num_audio_tokens is not None and len(num_audio_tokens) > 0:
354 + updates.append(
355 + PromptReplacement(
356 + modality="audio",
357 + target=[AUDIO_MARKER_ID],
358 + replacement=audio_replacement,
359 + )
360 + )
361 + return updates
vllm/models/inkling/common/towers.py new
+321
@@ -0,0 +1,321 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inkling Titan vision + audio towers.
4 +
5 +The vision tower (``InklingVision`` / ``HMLPPatchEncoder``) emits one token per
6 +image patch; the audio tower (``InklingAudio``) emits one token per audio frame.
7 +Both use vLLM's standard ``RMSNorm`` (CPU-friendly, with a native fallback).
8 +"""
9 +
10 +from __future__ import annotations
11 +
12 +from typing import cast
13 +
14 +import numpy as np
15 +import torch
16 +import torch.nn.functional as F
17 +from torch import nn
18 +
19 +from vllm.model_executor.layers.layernorm import RMSNorm
20 +
21 +from ..configs import InklingAudioConfig, InklingVisionConfig
22 +
23 +# ===========================================================================
24 +# Vision tower (HMLPPatchEncoder / InklingVision)
25 +# ===========================================================================
26 +
27 +
28 +def _prime_factors(n: int) -> list[int]:
29 + """Return the prime factors of *n* in ascending order."""
30 + if n < 1:
31 + raise ValueError("n must be a positive integer")
32 +
33 + factors: list[int] = []
34 + while n % 2 == 0:
35 + factors.append(2)
36 + n //= 2
37 + p = 3
38 + while p * p <= n:
39 + while n % p == 0:
40 + factors.append(p)
41 + n //= p
42 + p += 2
43 + if n > 1:
44 + factors.append(n)
45 + return factors
46 +
47 +
48 +def plan_out_scales(
49 + temporal_patch_size: int, patch_size: int, n_layers: int, n_channels: int = 3
50 +) -> list[tuple[int, int, int, int]]:
51 + """Plan the (t, h, w, c) dimensions for each HMLP layer.
52 +
53 + Spatial dims expand first, then temporal; channel counts round up to
54 + multiples of 64.
55 + """
56 + if patch_size <= 1:
57 + raise ValueError(
58 + "patch_size must be greater than 1, otherwise this doesn't make sense"
59 + )
60 +
61 + def _round_up(x: int) -> int:
62 + return int(np.ceil(x / 64)) * 64
63 +
64 + last_h_scale = 1
65 + scales: list[tuple[int, int, int, int]] = [(1, 1, 1, n_channels)]
66 + for pscale in _prime_factors(patch_size)[::-1]:
67 + last_h_scale *= pscale
68 + scales.append(
69 + (
70 + 1,
71 + last_h_scale,
72 + last_h_scale,
73 + _round_up((last_h_scale**2) * n_channels),
74 + )
75 + )
76 + last_t_scale = 1
77 + for tscale in _prime_factors(temporal_patch_size)[::-1]:
78 + last_t_scale *= tscale
79 + scales.append(
80 + (
81 + last_t_scale,
82 + last_h_scale,
83 + last_h_scale,
84 + _round_up((last_h_scale**2) * n_channels * last_t_scale),
85 + )
86 + )
87 +
88 + size_reduction = np.prod(np.array(scales)[:, :-1], 1)
89 +
90 + log_ideal_scales = np.linspace(
91 + 0,
92 + np.log(patch_size * patch_size * temporal_patch_size * n_channels),
93 + n_layers + 1,
94 + )
95 + cost_matrix = np.abs(log_ideal_scales[:, None] - np.log(size_reduction)[None])
96 +
97 + if n_layers >= len(scales):
98 + idxs = np.argmin(cost_matrix, axis=1)
99 + else:
100 + from scipy.optimize import linear_sum_assignment
101 +
102 + idxs = linear_sum_assignment(cost_matrix)[1]
103 +
104 + assert len(idxs) >= 2
105 + idxs[0] = 0
106 + idxs[-1] = len(scales) - 1
107 +
108 + return [scales[i] for i in idxs]
109 +
110 +
111 +def fold_timespace_to_depth(
112 + vision_patches_bthwc: torch.Tensor, t_fold: int, hw_fold: int
113 +) -> torch.Tensor:
114 + """(B, T, H, W, C) -> (B, T//t, H//hw, W//hw, C*(t*hw**2))."""
115 + B, T, H, W, C = vision_patches_bthwc.shape
116 +
117 + assert T % t_fold == 0, f"Temporal dimension {T} must be divisible by {t_fold}"
118 + assert H % hw_fold == 0, f"Height dimension {H} must be divisible by {hw_fold}"
119 + assert W % hw_fold == 0, f"Width dimension {W} must be divisible by {hw_fold}"
120 +
121 + t_new = T // t_fold
122 + h_new = H // hw_fold
123 + w_new = W // hw_fold
124 +
125 + x = vision_patches_bthwc.reshape(
126 + B, t_new, t_fold, h_new, hw_fold, w_new, hw_fold, C
127 + )
128 + x = x.permute(0, 1, 3, 5, 2, 4, 6, 7)
129 + x = x.reshape(B, t_new, h_new, w_new, t_fold * hw_fold * hw_fold * C)
130 + return x
131 +
132 +
133 +class HMLPPatchEncoder(nn.Module):
134 + def __init__(self, config: InklingVisionConfig):
135 + super().__init__()
136 + self.decoder_dmodel = config.decoder_dmodel
137 + self.patch_size = config.patch_size
138 + self.temporal_patch_size = config.temporal_patch_size
139 + self.n_channels = config.n_channels
140 + self.n_layers = config.n_layers
141 + self.use_vision_norm = config.use_vision_norm
142 +
143 + self.scales: list[tuple[int, int, int, int]] = plan_out_scales(
144 + self.temporal_patch_size, self.patch_size, self.n_layers, self.n_channels
145 + )
146 + self.layers: nn.ModuleDict = nn.ModuleDict()
147 + for i, (start_scale, end_scale) in enumerate(
148 + zip(self.scales[:-1], self.scales[1:])
149 + ):
150 + shuffle_mult = (
151 + (end_scale[0] // start_scale[0])
152 + * (end_scale[1] // start_scale[1])
153 + * (end_scale[2] // start_scale[2])
154 + )
155 + if i == self.n_layers - 1:
156 + self.layers[f"linear_{i}"] = nn.Linear(
157 + start_scale[3] * shuffle_mult, self.decoder_dmodel, bias=False
158 + )
159 + else:
160 + self.layers[f"linear_{i}"] = nn.Linear(
161 + start_scale[3] * shuffle_mult, end_scale[3], bias=False
162 + )
163 + self.layers[f"norm_{i}"] = RMSNorm(end_scale[3])
164 +
165 + self.final_norm: RMSNorm | None = None
166 + if self.use_vision_norm:
167 + assert self.decoder_dmodel is not None
168 + self.final_norm = RMSNorm(self.decoder_dmodel)
169 +
170 + def forward(self, x: torch.Tensor) -> torch.Tensor:
171 + # Fused norm+gelu on CUDA (same fp32-accum/bf16-rounding structure as
172 + # the generic path below; differs only by reduction order).
173 + fused = None
174 + if x.is_cuda and x.dtype == torch.bfloat16:
175 + from vllm.models.inkling.nvidia.ops.mm_towers import rmsnorm_gelu
176 +
177 + fused = rmsnorm_gelu
178 +
179 + num_patches, T, H, W, C = x.shape
180 + prefolded = False
181 + for i, (start_scale, end_scale) in enumerate(
182 + zip(self.scales[:-1], self.scales[1:])
183 + ):
184 + t_fold = end_scale[0] // start_scale[0]
185 + hw_fold = end_scale[1] // start_scale[1]
186 + if (hw_fold > 1 or t_fold > 1) and not prefolded:
187 + x = fold_timespace_to_depth(x, t_fold, hw_fold)
188 + prefolded = False
189 + assert x.shape[1:-1] == (
190 + T // end_scale[0],
191 + H // end_scale[1],
192 + W // end_scale[2],
193 + )
194 + x = self.layers[f"linear_{i}"](x)
195 + if i < self.n_layers - 1:
196 + norm = cast(RMSNorm, self.layers[f"norm_{i}"])
197 + if fused is not None:
198 + # If the NEXT layer starts with a copying fold (spatial
199 + # dims still > 1 after folding), store this layer's
200 + # output directly in the folded layout instead.
201 + nxt = self.scales[i + 2]
202 + ntf = nxt[0] // end_scale[0]
203 + nhf = nxt[1] // end_scale[1]
204 + copy_fold = (ntf > 1 or nhf > 1) and (
205 + x.shape[2] // nhf > 1 or x.shape[3] // nhf > 1
206 + )
207 + x = fused(
208 + x,
209 + norm.weight,
210 + norm.variance_epsilon,
211 + gelu=True,
212 + fold=(ntf, nhf) if copy_fold else None,
213 + )
214 + prefolded = copy_fold
215 + else:
216 + # rms_norm kernel only supports rank 2-4; x is 5-D here.
217 + orig = x.shape
218 + x = norm(x.reshape(-1, x.shape[-1])).reshape(orig)
219 + x = F.gelu(x)
220 +
221 + if self.final_norm is not None:
222 + if fused is not None:
223 + x = fused(
224 + x,
225 + self.final_norm.weight,
226 + self.final_norm.variance_epsilon,
227 + gelu=False,
228 + )
229 + else:
230 + orig = x.shape
231 + x = self.final_norm(x.reshape(-1, x.shape[-1])).reshape(orig)
232 +
233 + x = x.reshape(num_patches, -1)
234 + return x
235 +
236 +
237 +class InklingVision(nn.Module):
238 + def __init__(self, config: InklingVisionConfig, prefix: str = ""):
239 + del prefix
240 + super().__init__()
241 + assert config.vision_encoder_type == "hmlp"
242 + self.vision_encoder = HMLPPatchEncoder(config)
243 +
244 + @property
245 + def dtype(self) -> torch.dtype:
246 + return next(self.parameters()).dtype
247 +
248 + @property
249 + def device(self) -> torch.device:
250 + return next(self.parameters()).device
251 +
252 + def forward(self, vision_features: torch.Tensor) -> torch.Tensor:
253 + return self.vision_encoder(vision_features)
254 +
255 +
256 +# ===========================================================================
257 +# Audio tower (InklingAudio)
258 +# ===========================================================================
259 +
260 +
261 +class InklingAudio(nn.Module):
262 + def __init__(self, config: InklingAudioConfig, prefix: str = ""):
263 + del prefix
264 + super().__init__()
265 + assert config.audio_mode == "dmel"
266 + self.n_mel_bins = config.n_mel_bins
267 + self.mel_vocab_size = config.mel_vocab_size
268 + self.use_audio_norm = config.use_audio_norm
269 + self.encoder = nn.Embedding(
270 + config.n_mel_bins * config.mel_vocab_size, config.decoder_dmodel
271 + )
272 + self.final_norm: RMSNorm | None = None
273 + if self.use_audio_norm:
274 + assert config.decoder_dmodel is not None
275 + self.final_norm = RMSNorm(config.decoder_dmodel, eps=1e-6)
276 +
277 + @property
278 + def dtype(self) -> torch.dtype:
279 + return self.encoder.weight.dtype
280 +
281 + @property
282 + def device(self) -> torch.device:
283 + return self.encoder.weight.device
284 +
285 + def forward(self, audio_features: torch.Tensor) -> torch.Tensor:
286 + assert audio_features.shape[1] == self.n_mel_bins
287 +
288 + # dMel bins are integer indices; cast once to int32 on the right device
289 + # (no float round-trip).
290 + audio_features = audio_features.to(
291 + device=self.encoder.weight.device, dtype=torch.int32
292 + )
293 +
294 + weight = self.encoder.weight
295 + if audio_features.is_cuda and weight.dtype == torch.bfloat16:
296 + # One kernel: per-bin offset + embedding gather + fp32 sum + norm.
297 + # Skips the [T, n_mel_bins, D] intermediate entirely (bit-exact).
298 + from vllm.models.inkling.nvidia.ops.mm_towers import dmel_embed_sum_norm
299 +
300 + return dmel_embed_sum_norm(
301 + audio_features.contiguous(),
302 + weight,
303 + self.final_norm.weight if self.final_norm is not None else None,
304 + self.final_norm.variance_epsilon if self.final_norm else 0.0,
305 + )
306 +
307 + embedding_indices = (
308 + torch.arange(self.n_mel_bins, device=audio_features.device)
309 + * self.mel_vocab_size
310 + ).unsqueeze(0) + audio_features
311 +
312 + hidden_states = (
313 + self.encoder(embedding_indices.reshape(-1))
314 + .reshape(audio_features.shape[0], audio_features.shape[1], -1)
315 + .sum(axis=1)
316 + )
317 +
318 + if self.final_norm is not None:
319 + hidden_states = self.final_norm(hidden_states)
320 +
321 + return hidden_states
vllm/models/inkling/configs.py new
+373
@@ -0,0 +1,373 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inkling model configs for the text backbone and audio/vision towers."""
4 +
5 +from __future__ import annotations
6 +
7 +from dataclasses import dataclass, field
8 +from typing import Any, ClassVar, Literal, cast
9 +
10 +import torch
11 +from transformers.configuration_utils import PretrainedConfig
12 +
13 +
14 +class InklingModelConfig(PretrainedConfig):
15 + model_type = "inkling_model"
16 + keys_to_ignore_at_inference = ["past_key_values"]
17 +
18 + def __init__(
19 + self,
20 + *,
21 + vocab_size: int = 201024,
22 + hidden_size: int = 1536,
23 + intermediate_size: int = 768,
24 + dense_intermediate_size: int | None = None,
25 + num_hidden_layers: int = 16,
26 + num_attention_heads: int = 12,
27 + num_key_value_heads: int = 4,
28 + head_dim: int | None = None,
29 + v_head_dim: int | None = None,
30 + d_rel: int = 16,
31 + rel_extent: int = 1024,
32 + local_layer_ids: list[int] | None = None,
33 + sliding_window_size: int = 512,
34 + swa_num_attention_heads: int | None = None,
35 + swa_num_key_value_heads: int | None = None,
36 + swa_head_dim: int | None = None,
37 + swa_v_head_dim: int | None = None,
38 + rms_norm_eps: float = 1e-6,
39 + hidden_act: str = "silu",
40 + q_bias: bool = False,
41 + o_bias: bool = False,
42 + use_embed_norm: bool = False,
43 + use_sconv: bool = False,
44 + sconv_kernel_size: int = 4,
45 + dense_mlp_idx: int = 0,
46 + n_routed_experts: int = 0,
47 + n_shared_experts: int = 0,
48 + num_experts_per_tok: int = 1,
49 + route_scale: float = 1.0,
50 + use_gate_bias: bool = False,
51 + use_global_scale: bool = False,
52 + norm_after_topk: bool = True,
53 + gate_activation: Literal["sigmoid", "softmax"] = "sigmoid",
54 + shared_expert_sink: bool = False,
55 + shared_experts_size: int = 1,
56 + inference_moe_w13_interleaved: bool = True,
57 + log_scaling_n_floor: int | None = None,
58 + log_scaling_alpha: float = 0.1,
59 + unpadded_vocab_size: int | None = None,
60 + padded_vocab_size: int | None = None,
61 + logits_mup_width_multiplier: float | None = None,
62 + final_logit_softcapping: float | None = None,
63 + tie_word_embeddings: bool = False,
64 + **kwargs: Any,
65 + ) -> None:
66 + if head_dim is None:
67 + head_dim = hidden_size // num_attention_heads
68 + if v_head_dim is None:
69 + v_head_dim = head_dim
70 + if swa_num_attention_heads is None:
71 + swa_num_attention_heads = num_attention_heads
72 + if swa_num_key_value_heads is None:
73 + swa_num_key_value_heads = num_key_value_heads
74 + if swa_head_dim is None:
75 + swa_head_dim = head_dim
76 + if swa_v_head_dim is None:
77 + swa_v_head_dim = swa_head_dim
78 + if dense_intermediate_size is None:
79 + dense_intermediate_size = intermediate_size
80 + if local_layer_ids is None:
81 + local_layer_ids = []
82 +
83 + if padded_vocab_size is None:
84 + padded_vocab_size = vocab_size
85 + vocab_size = (
86 + unpadded_vocab_size
87 + if (
88 + unpadded_vocab_size is not None
89 + and unpadded_vocab_size < padded_vocab_size
90 + )
91 + else vocab_size
92 + )
93 +
94 + self.vocab_size = vocab_size
95 + self.padded_vocab_size = padded_vocab_size
96 + self.hidden_size = hidden_size
97 + self.intermediate_size = intermediate_size
98 + self.dense_intermediate_size = dense_intermediate_size
99 + self.num_hidden_layers = num_hidden_layers
100 + self.num_attention_heads = num_attention_heads
101 + self.num_key_value_heads = num_key_value_heads
102 + self.head_dim = head_dim
103 + self.v_head_dim = v_head_dim
104 + self.d_rel = d_rel
105 + self.rel_extent = rel_extent
106 + self.local_layer_ids = local_layer_ids
107 + self.sliding_window_size = sliding_window_size
108 + self.swa_num_attention_heads = swa_num_attention_heads
109 + self.swa_num_key_value_heads = swa_num_key_value_heads
110 + self.swa_head_dim = swa_head_dim
111 + self.swa_v_head_dim = swa_v_head_dim
112 + self.rms_norm_eps = rms_norm_eps
113 + self.hidden_act = hidden_act
114 + self.q_bias = q_bias
115 + self.o_bias = o_bias
116 + self.use_embed_norm = use_embed_norm
117 + self.use_sconv = use_sconv
118 + self.sconv_kernel_size = sconv_kernel_size
119 + self.dense_mlp_idx = dense_mlp_idx
120 + self.n_routed_experts = n_routed_experts
121 + self.num_experts = n_routed_experts
122 + self.n_shared_experts = n_shared_experts
123 + self.num_shared_experts = n_shared_experts
124 + self.num_experts_per_tok = num_experts_per_tok
125 + self.route_scale = route_scale
126 + self.use_gate_bias = use_gate_bias
127 + self.use_global_scale = use_global_scale
128 + self.norm_after_topk = norm_after_topk
129 + self.gate_activation = gate_activation
130 + self.shared_expert_sink = shared_expert_sink
131 + self.shared_experts_size = shared_experts_size
132 + self.inference_moe_w13_interleaved = inference_moe_w13_interleaved
133 + self.log_scaling_n_floor = log_scaling_n_floor
134 + self.log_scaling_alpha = log_scaling_alpha
135 + self.unpadded_vocab_size = self.vocab_size
136 + self.logits_mup_width_multiplier = logits_mup_width_multiplier
137 + self.final_logit_softcapping = final_logit_softcapping
138 +
139 + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
140 +
141 + @property
142 + def conv_layer_ids(self) -> list[int]:
143 + return list(range(self.num_hidden_layers))
144 +
145 + @property
146 + def linear_layer_ids(self) -> list[int]:
147 + return self.conv_layer_ids
148 +
149 + @property
150 + def full_attention_layer_ids(self) -> list[int]:
151 + return list(range(self.num_hidden_layers))
152 +
153 + @property
154 + def mamba_chunk_size(self) -> int:
155 + # Floor at 64: mamba_cache_chunk_size = max(mamba_chunk_size, page_size),
156 + # and a floor of 1 lets the radix tree adopt another request's KV at
157 + # tiny shared prefixes, whose different kernel-rounding perturbs decode
158 + # logits.
159 + return 64
160 +
161 + @property
162 + def mamba2_cache_params(self) -> TMLConvCacheParams | None:
163 + try:
164 + from vllm.distributed import get_tensor_model_parallel_world_size
165 +
166 + tp_size = get_tensor_model_parallel_world_size()
167 + except (AssertionError, RuntimeError):
168 + tp_size = 1
169 +
170 + def tp_local_kv_conv_dim(num_kv_heads: int, head_dim: int) -> int:
171 + return max(1, num_kv_heads // tp_size) * head_dim
172 +
173 + full_kv_conv_dim = tp_local_kv_conv_dim(self.num_key_value_heads, self.head_dim)
174 + local_kv_conv_dim = tp_local_kv_conv_dim(
175 + self.swa_num_key_value_heads, self.swa_head_dim
176 + )
177 + stream_dim = self.hidden_size
178 + conv_len = self.sconv_kernel_size - 1
179 + shape = TMLConvStateShape(
180 + conv=[
181 + (conv_len, full_kv_conv_dim),
182 + (conv_len, full_kv_conv_dim),
183 + (conv_len, local_kv_conv_dim),
184 + (conv_len, local_kv_conv_dim),
185 + (conv_len, stream_dim),
186 + (conv_len, stream_dim),
187 + ],
188 + temporal=(0, 0, 0),
189 + )
190 + dtype = TMLStateDType(conv=torch.bfloat16, temporal=torch.bfloat16)
191 + return TMLConvCacheParams(shape=shape, layers=self.conv_layer_ids, dtype=dtype)
192 +
193 +
194 +class InklingAudioConfig(PretrainedConfig):
195 + model_type = "inkling_audio_model"
196 +
197 + def __init__(
198 + self,
199 + *,
200 + decoder_dmodel: int | None = None,
201 + n_mel_bins: int | None = None,
202 + mel_vocab_size: int | None = None,
203 + dmel_min_value: float | None = None,
204 + dmel_max_value: float | None = None,
205 + use_audio_norm: bool | None = None,
206 + audio_mode: Literal["dmel", "flow"] | None = None,
207 + **kwargs: Any,
208 + ) -> None:
209 + values = {
210 + "n_mel_bins": n_mel_bins,
211 + "mel_vocab_size": mel_vocab_size,
212 + "dmel_min_value": dmel_min_value,
213 + "dmel_max_value": dmel_max_value,
214 + "use_audio_norm": use_audio_norm,
215 + "audio_mode": audio_mode,
216 + }
217 + if decoder_dmodel is not None and (
218 + missing := [name for name, value in values.items() if value is None]
219 + ):
220 + raise ValueError(
221 + "Enabled Inkling audio tower is missing config fields: "
222 + + ", ".join(missing)
223 + )
224 + self.decoder_dmodel = decoder_dmodel
225 + self.n_mel_bins = cast(int, n_mel_bins)
226 + self.mel_vocab_size = cast(int, mel_vocab_size)
227 + self.dmel_min_value = cast(float, dmel_min_value)
228 + self.dmel_max_value = cast(float, dmel_max_value)
229 + self.use_audio_norm = cast(bool, use_audio_norm)
230 + self.audio_mode = cast(Literal["dmel", "flow"], audio_mode)
231 + super().__init__(**kwargs)
232 +
233 +
234 +class InklingVisionConfig(PretrainedConfig):
235 + model_type = "inkling_vision_model"
236 +
237 + def __init__(
238 + self,
239 + *,
240 + vision_encoder_type: Literal["linear", "hmlp"] | None = None,
241 + decoder_dmodel: int | None = None,
242 + patch_size: int | None = None,
243 + temporal_patch_size: int | None = None,
244 + n_channels: int | None = None,
245 + n_layers: int | None = None,
246 + use_vision_norm: bool | None = None,
247 + **kwargs: Any,
248 + ) -> None:
249 + values = {
250 + "vision_encoder_type": vision_encoder_type,
251 + "patch_size": patch_size,
252 + "temporal_patch_size": temporal_patch_size,
253 + "n_channels": n_channels,
254 + "n_layers": n_layers,
255 + "use_vision_norm": use_vision_norm,
256 + }
257 + if decoder_dmodel is not None and (
258 + missing := [name for name, value in values.items() if value is None]
259 + ):
260 + raise ValueError(
261 + "Enabled Inkling vision tower is missing config fields: "
262 + + ", ".join(missing)
263 + )
264 + self.vision_encoder_type = cast(Literal["linear", "hmlp"], vision_encoder_type)
265 + self.decoder_dmodel = decoder_dmodel
266 + self.patch_size = cast(int, patch_size)
267 + self.temporal_patch_size = cast(int, temporal_patch_size)
268 + self.n_channels = cast(int, n_channels)
269 + self.n_layers = cast(int, n_layers)
270 + self.use_vision_norm = cast(bool, use_vision_norm)
271 + super().__init__(**kwargs)
272 +
273 +
274 +class InklingMMConfig(PretrainedConfig):
275 + model_type = "inkling_mm_model"
276 + keys_to_ignore_at_inference = ["past_key_values"]
277 + sub_configs: ClassVar[dict[str, type[PretrainedConfig]]] = {
278 + "text_config": InklingModelConfig,
279 + "audio_config": InklingAudioConfig,
280 + "vision_config": InklingVisionConfig,
281 + }
282 +
283 + def __init__(
284 + self,
285 + *,
286 + text_config: dict[str, Any] | InklingModelConfig | None = None,
287 + audio_config: dict[str, Any] | InklingAudioConfig | None = None,
288 + vision_config: dict[str, Any] | InklingVisionConfig | None = None,
289 + tie_word_embeddings: bool = False,
290 + **kwargs: Any,
291 + ) -> None:
292 + self.text_config = (
293 + text_config
294 + if isinstance(text_config, InklingModelConfig)
295 + else InklingModelConfig(**(text_config or {}))
296 + )
297 + self.audio_config = (
298 + audio_config
299 + if isinstance(audio_config, InklingAudioConfig)
300 + else InklingAudioConfig(**(audio_config or {}))
301 + )
302 + self.vision_config = (
303 + vision_config
304 + if isinstance(vision_config, InklingVisionConfig)
305 + else InklingVisionConfig(**(vision_config or {}))
306 + )
307 + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
308 +
309 + def get_text_config(self, *args: Any, **kwargs: Any) -> InklingModelConfig:
310 + return self.text_config
311 +
312 + @property
313 + def vocab_size(self) -> int:
314 + return self.text_config.vocab_size
315 +
316 + @property
317 + def hidden_size(self) -> int:
318 + return self.text_config.hidden_size
319 +
320 + @property
321 + def num_hidden_layers(self) -> int:
322 + return self.text_config.num_hidden_layers
323 +
324 + @property
325 + def num_attention_heads(self) -> int:
326 + return self.text_config.num_attention_heads
327 +
328 + @property
329 + def num_key_value_heads(self) -> int:
330 + return self.text_config.num_key_value_heads
331 +
332 + @property
333 + def head_dim(self) -> int:
334 + return self.text_config.head_dim
335 +
336 + @property
337 + def full_attention_layer_ids(self) -> list[int]:
338 + return self.text_config.full_attention_layer_ids
339 +
340 + @property
341 + def linear_layer_ids(self) -> list[int]:
342 + return self.text_config.linear_layer_ids
343 +
344 + @property
345 + def conv_layer_ids(self) -> list[int]:
346 + return self.text_config.conv_layer_ids
347 +
348 + @property
349 + def mamba_chunk_size(self) -> int:
350 + return self.text_config.mamba_chunk_size
351 +
352 + @property
353 + def mamba2_cache_params(self) -> TMLConvCacheParams | None:
354 + return self.text_config.mamba2_cache_params
355 +
356 +
357 +@dataclass(kw_only=True, frozen=True)
358 +class TMLConvStateShape:
359 + conv: list[tuple[int, int]]
360 + temporal: tuple[int, int, int]
361 +
362 +
363 +@dataclass(kw_only=True, frozen=True)
364 +class TMLStateDType:
365 + conv: torch.dtype = torch.bfloat16
366 + temporal: torch.dtype = torch.bfloat16
367 +
368 +
369 +@dataclass(kw_only=True, frozen=True)
370 +class TMLConvCacheParams:
371 + shape: TMLConvStateShape
372 + layers: list[int]
373 + dtype: TMLStateDType = field(default_factory=TMLStateDType)
vllm/models/inkling/nvfp4.py new
+64
@@ -0,0 +1,64 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""NVFP4 (ModelOpt) support for the Inkling mixture-of-experts.
4 +
5 +Only the routed MoE experts are quantized in the Inkling checkpoint;
6 +attention, the dense MLP, and the shared "sink" experts stay bf16 (they are
7 +in the checkpoint ``exclude_modules``). The routed experts are served by
8 +vLLM's standard ModelOpt NVFP4 fused-MoE stack (see ``moe.py``); this module
9 +keeps the checkpoint detection.
10 +"""
11 +
12 +from __future__ import annotations
13 +
14 +FLOAT8_E4M3_MAX = 448.0
15 +FLOAT4_E2M1_MAX = 6.0
16 +
17 +
18 +class InklingNvfp4Config:
19 + """Lightweight NVFP4 descriptor parsed from the checkpoint quant config.
20 +
21 + Holds the (mapped) ``exclude_modules`` so the model can decide, per MoE
22 + layer and per expert group, whether the weights are NVFP4 or plain bf16.
23 + """
24 +
25 + def __init__(self, group_size: int, exclude_modules: list[str]) -> None:
26 + self.group_size = group_size
27 + self.exclude_modules = set(exclude_modules)
28 +
29 + @staticmethod
30 + def _is_nvfp4(quant_cfg: dict) -> bool:
31 + wq = quant_cfg["modelopt_quant_config"]["quant_cfg"]["*weight_quantizer"]
32 + return tuple(wq["num_bits"]) == (2, 1) and tuple(
33 + wq["block_sizes"].get("scale_bits", [])
34 + ) == (4, 3)
35 +
36 + @classmethod
37 + def from_hf_config(cls, hf_config) -> InklingNvfp4Config | None:
38 + quant_cfg = getattr(hf_config, "quantization_config", None)
39 + text_config = getattr(hf_config, "text_config", None)
40 + if quant_cfg is None and text_config is not None:
41 + quant_cfg = getattr(text_config, "quantization_config", None)
42 + if quant_cfg is None:
43 + return None
44 + # ModelOpt <=0.29 nests everything under "quantization".
45 + if "quantization" in quant_cfg:
46 + quant_cfg = quant_cfg["quantization"]
47 + if not cls._is_nvfp4(quant_cfg):
48 + return None
49 + group_size = quant_cfg.get("group_size", 16)
50 + if group_size != 16:
51 + raise ValueError("Inkling NVFP4 only supports group size 16")
52 + exclude = list(quant_cfg.get("exclude_modules", []) or [])
53 + return cls(group_size=group_size, exclude_modules=exclude)
54 +
55 + def experts_quantized(self, layer_id: int) -> bool:
56 + """Whether the routed experts of ``layer_id`` are NVFP4 (vs excluded)."""
57 + return f"model.llm.layers.{layer_id}.mlp.experts" not in self.exclude_modules
58 +
59 + def shared_experts_quantized(self, layer_id: int) -> bool:
60 + """Whether the shared sink experts of ``layer_id`` are NVFP4."""
61 + return (
62 + f"model.llm.layers.{layer_id}.mlp.shared_experts"
63 + not in self.exclude_modules
64 + )
vllm/models/inkling/nvidia/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/inkling/nvidia/attention.py new
+325
@@ -0,0 +1,325 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from __future__ import annotations
4 +
5 +from typing import cast
6 +
7 +import torch
8 +from torch import nn
9 +
10 +from vllm.config import VllmConfig, get_current_vllm_config
11 +from vllm.distributed import get_tensor_model_parallel_world_size
12 +from vllm.forward_context import get_forward_context
13 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
14 +from vllm.model_executor.layers.linear import (
15 + MergedColumnParallelLinear,
16 + RowParallelLinear,
17 +)
18 +from vllm.model_executor.layers.quantization import QuantizationConfig
19 +from vllm.utils.torch_utils import (
20 + canonicalize_singleton_dim_strides,
21 + kv_cache_dtype_str_to_dtype,
22 +)
23 +from vllm.v1.attention.backend import AttentionBackend
24 +from vllm.v1.attention.backends.flash_attn import (
25 + FlashAttentionBackend,
26 + FlashAttentionMetadata,
27 +)
28 +from vllm.v1.kv_cache_interface import (
29 + FullAttentionSpec,
30 + KVCacheSpec,
31 + SlidingWindowSpec,
32 +)
33 +
34 +from ..configs import InklingModelConfig
35 +from .layernorm import InklingRMSNorm
36 +from .ops.fa4_rel_attention import (
37 + bucket_max_seqlen_q,
38 + inkling_fa4_num_splits,
39 + inkling_fa4_rel_attention,
40 +)
41 +from .ops.fa4_warmup import InklingFA4WarmupConfig, register_fa4_warmup
42 +from .ops.qkvr_prep import fused_qkvr_prep
43 +from .sconv_swa_attn import _K, _V, InklingConvState, InklingSconvMetadata
44 +from .short_conv import InklingShortConv
45 +
46 +
47 +def compute_log_scaling_tau(
48 + positions: torch.Tensor, n_floor: int, alpha: float
49 +) -> torch.Tensor:
50 + effective_n = (positions + 1).to(torch.float32)
51 + return 1.0 + alpha * torch.log(torch.clamp(effective_n / float(n_floor), min=1.0))
52 +
53 +
54 +class RelLogitsProj(nn.Module):
55 + """Project the per-head relative branch ``r`` to per-distance logits."""
56 +
57 + def __init__(self, d_rel: int, rel_extent: int) -> None:
58 + super().__init__()
59 + self.d_rel = d_rel
60 + self.rel_extent = rel_extent
61 + self.proj = nn.Parameter(torch.empty(d_rel, rel_extent), requires_grad=False)
62 +
63 + def forward(self, r_out: torch.Tensor) -> torch.Tensor:
64 + # r_out: (T, num_heads, d_rel) -> (T, num_heads, rel_extent)
65 + return torch.einsum("thd,de->the", r_out, self.proj)
66 +
67 +
68 +class InklingAttention(nn.Module, AttentionLayerBase):
69 + def __init__(
70 + self,
71 + config: InklingModelConfig,
72 + *,
73 + num_heads: int,
74 + num_kv_heads: int,
75 + head_dim: int,
76 + rel_extent: int,
77 + local_extent: int,
78 + is_local: bool,
79 + prefix: str,
80 + quant_config: QuantizationConfig | None = None,
81 + conv_owner: InklingConvState,
82 + ) -> None:
83 + super().__init__()
84 + self.prefix = prefix
85 + self.is_local = is_local
86 + self.hidden_size = config.hidden_size
87 + self.head_dim = head_dim
88 + self.d_rel = config.d_rel
89 + self.log_scaling_n_floor = config.log_scaling_n_floor
90 + self.log_scaling_alpha = config.log_scaling_alpha
91 + # q/k are per-head RMS-normed (unit norm), so Inkling scales by 1/head_dim.
92 + self.scaling = 1.0 / head_dim
93 +
94 + tp_size = get_tensor_model_parallel_world_size()
95 + self.num_total_heads = num_heads
96 + self.num_total_kv_heads = num_kv_heads
97 + assert self.num_total_heads % tp_size == 0
98 + self.num_heads = self.num_total_heads // tp_size
99 + if self.num_total_kv_heads >= tp_size:
100 + assert self.num_total_kv_heads % tp_size == 0
101 + else:
102 + assert tp_size % self.num_total_kv_heads == 0
103 + self.num_kv_heads = max(1, self.num_total_kv_heads // tp_size)
104 + # When tp_size > num_kv_heads the K/V projections are padded up to
105 + # tp_size heads so each rank gets at least one (GQA replication).
106 + kv_total_for_sizing = max(self.num_total_kv_heads, tp_size)
107 +
108 + self.qkvr = MergedColumnParallelLinear(
109 + input_size=config.hidden_size,
110 + output_sizes=[
111 + head_dim * self.num_total_heads,
112 + head_dim * kv_total_for_sizing,
113 + head_dim * kv_total_for_sizing,
114 + self.d_rel * self.num_total_heads,
115 + ],
116 + bias=config.q_bias,
117 + quant_config=quant_config,
118 + prefix=f"{prefix}.qkvr",
119 + )
120 + self.wo_ud = RowParallelLinear(
121 + input_size=head_dim * self.num_total_heads,
122 + output_size=config.hidden_size,
123 + bias=config.o_bias,
124 + quant_config=quant_config,
125 + # reduce_results=False: the partial output is all-reduced below
126 + # (one-shot custom AR) so the attention-output sconv can run on the
127 + # full hidden width fused with the residual add + rmsnorm.
128 + reduce_results=False,
129 + prefix=f"{prefix}.wo_ud",
130 + )
131 + self.rel_extent = local_extent if is_local else rel_extent
132 + self.local_extent = local_extent if is_local else None
133 + self.rel_logits_proj = RelLogitsProj(self.d_rel, self.rel_extent)
134 + self.q_norm = InklingRMSNorm(head_dim, eps=config.rms_norm_eps)
135 + self.k_norm = InklingRMSNorm(head_dim, eps=config.rms_norm_eps)
136 +
137 + # Short convolution on the K/V streams (per-head-width, TP sharded),
138 + # applied after the qkvr projection and before q/k norm.
139 + kv_conv_dim = self.num_kv_heads * head_dim
140 + self.conv_owner = conv_owner
141 + self.k_sconv = InklingShortConv(
142 + kv_conv_dim, config.sconv_kernel_size, owner=conv_owner, stream_idx=_K
143 + )
144 + self.v_sconv = InklingShortConv(
145 + kv_conv_dim, config.sconv_kernel_size, owner=conv_owner, stream_idx=_V
146 + )
147 +
148 + # FA4 left/right window; right=0 keeps it causal. local_extent-1 mirrors
149 + # the source (sliding_window_size - 1).
150 + self.window_size: tuple[int, int] = (
151 + (local_extent - 1, 0) if is_local else (-1, -1)
152 + )
153 + # Static per-layer-type KV length bound for the split heuristic: local
154 + # layers never see more than the sliding window.
155 + vllm_config = get_current_vllm_config()
156 + self._max_kv_len = (
157 + local_extent if is_local else vllm_config.model_config.max_model_len
158 + )
159 +
160 + # ---- KV-cache wiring (reuse FlashAttentionBackend for metadata) ----
161 + cache_config = vllm_config.cache_config
162 + self.kv_cache_dtype = (
163 + cache_config.cache_dtype if cache_config is not None else "auto"
164 + )
165 + self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
166 + self.kv_cache_dtype, vllm_config.model_config
167 + )
168 + self.register_buffer("k_scale", torch.ones((), dtype=torch.float32))
169 + self.register_buffer("v_scale", torch.ones((), dtype=torch.float32))
170 +
171 + compilation_config = vllm_config.compilation_config
172 + if prefix in compilation_config.static_forward_context:
173 + raise ValueError(f"Duplicate layer name: {prefix}")
174 + compilation_config.static_forward_context[prefix] = self
175 + self.kv_cache = torch.tensor([]) # replaced by bind_kv_cache
176 +
177 + register_fa4_warmup(
178 + InklingFA4WarmupConfig(
179 + num_heads=self.num_heads,
180 + num_kv_heads=self.num_kv_heads,
181 + head_dim=self.head_dim,
182 + rel_extent=self.rel_extent,
183 + window_size=self.window_size,
184 + is_local=self.is_local,
185 + max_kv_len=self._max_kv_len,
186 + dtype=vllm_config.model_config.dtype,
187 + kv_dtype=self.kv_cache_torch_dtype,
188 + block_size=vllm_config.cache_config.block_size,
189 + max_num_reqs=vllm_config.scheduler_config.max_num_seqs,
190 + max_num_batched_tokens=(
191 + vllm_config.scheduler_config.max_num_batched_tokens
192 + ),
193 + )
194 + )
195 +
196 + def get_attn_backend(self) -> type[AttentionBackend]:
197 + return FlashAttentionBackend
198 +
199 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
200 + block_size = vllm_config.cache_config.block_size
201 + if self.is_local:
202 + assert self.local_extent is not None
203 + return SlidingWindowSpec(
204 + block_size=block_size,
205 + num_kv_heads=self.num_kv_heads,
206 + head_size=self.head_dim,
207 + dtype=self.kv_cache_torch_dtype,
208 + sliding_window=self.local_extent,
209 + )
210 + return FullAttentionSpec(
211 + block_size=block_size,
212 + num_kv_heads=self.num_kv_heads,
213 + head_size=self.head_dim,
214 + dtype=self.kv_cache_torch_dtype,
215 + )
216 +
217 + def _split_kv_cache(self) -> tuple[torch.Tensor, torch.Tensor]:
218 + key_cache, value_cache = self.kv_cache.transpose(1, 2).split(
219 + self.head_dim, dim=-1
220 + )
221 + return (
222 + canonicalize_singleton_dim_strides(key_cache),
223 + canonicalize_singleton_dim_strides(value_cache),
224 + )
225 +
226 + def forward(
227 + self,
228 + positions: torch.Tensor,
229 + hidden_states: torch.Tensor,
230 + log_scaling: torch.Tensor | None = None,
231 + ) -> torch.Tensor:
232 + num_tokens = hidden_states.shape[0]
233 + qkvr, _ = self.qkvr(hidden_states)
234 +
235 + attn_metadata = get_forward_context().attn_metadata
236 + attn_output = torch.empty(
237 + (num_tokens, self.num_heads, self.head_dim),
238 + dtype=qkvr.dtype,
239 + device=qkvr.device,
240 + )
241 + if not isinstance(attn_metadata, dict):
242 + attn_output.zero_()
243 + else:
244 + conv_meta = attn_metadata[self.conv_owner.prefix]
245 + md = attn_metadata[self.prefix]
246 + assert isinstance(conv_meta, InklingSconvMetadata)
247 + fa_md = cast(FlashAttentionMetadata, md)
248 + assert self.kv_cache.numel() > 0
249 + assert self.conv_owner.kv_cache.numel() > 0
250 + # One launch: K/V sconv (conv-cache insert + conv + residual),
251 + # Q/K per-head rmsnorm, and the attention KV-cache write. K/V are
252 + # consumed via the KV cache; only normed q is materialized.
253 + key_cache, value_cache = self._split_kv_cache()
254 + off_k, _ = self.conv_owner.stream_ranges[_K]
255 + off_v, _ = self.conv_owner.stream_ranges[_V]
256 + q, rel_logits = fused_qkvr_prep(
257 + qkvr,
258 + self.k_sconv.weight.squeeze(1),
259 + self.v_sconv.weight.squeeze(1),
260 + self.q_norm.weight,
261 + self.k_norm.weight,
262 + self.rel_logits_proj.proj,
263 + self.q_norm.variance_epsilon,
264 + self.num_heads,
265 + self.num_kv_heads,
266 + self.head_dim,
267 + self.d_rel,
268 + self.conv_owner.kv_cache,
269 + key_cache,
270 + value_cache,
271 + positions,
272 + conv_meta.block_table,
273 + conv_meta.seq_idx,
274 + conv_meta.slot_mapping,
275 + conv_meta.query_start,
276 + fa_md.slot_mapping,
277 + off_k,
278 + off_v,
279 + self.conv_owner.block_size,
280 + log_scaling if not self.is_local else None,
281 + )
282 + q = q.view(num_tokens, self.num_heads, self.head_dim)
283 + self._attention(q, rel_logits, attn_output)
284 +
285 + flat = attn_output.view(num_tokens, -1)
286 + output, _ = self.wo_ud(flat)
287 + return output
288 +
289 + def _attention(
290 + self,
291 + q: torch.Tensor,
292 + rel_logits: torch.Tensor,
293 + output: torch.Tensor,
294 + ) -> None:
295 + attn_metadata = get_forward_context().attn_metadata
296 + assert isinstance(attn_metadata, dict)
297 + md = cast(FlashAttentionMetadata, attn_metadata[self.prefix])
298 +
299 + nt = md.num_actual_tokens
300 + key_cache, value_cache = self._split_kv_cache()
301 + max_seqlen_q = bucket_max_seqlen_q(md.max_query_len)
302 + num_splits = inkling_fa4_num_splits(
303 + is_local=self.is_local,
304 + batch_size=md.seq_lens.shape[0],
305 + max_query_len=max_seqlen_q,
306 + num_heads=self.num_heads,
307 + num_kv_heads=self.num_kv_heads,
308 + max_kv_len=self._max_kv_len,
309 + )
310 + inkling_fa4_rel_attention(
311 + q[:nt],
312 + key_cache,
313 + value_cache,
314 + block_table=md.block_table,
315 + cache_seqlens=md.seq_lens,
316 + cu_seqlens_q=md.query_start_loc,
317 + max_seqlen_q=max_seqlen_q,
318 + softmax_scale=self.scaling,
319 + causal=True,
320 + window_size=self.window_size,
321 + rel_extent=self.rel_extent,
322 + rel_logits=rel_logits[:nt],
323 + num_splits=num_splits,
324 + out=output[:nt],
325 + )
vllm/models/inkling/nvidia/layernorm.py new
+26
@@ -0,0 +1,26 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inkling RMSNorm (no bias, weight-scaled), backed by the vendored Triton kernel."""
4 +
5 +from __future__ import annotations
6 +
7 +import torch
8 +from torch import nn
9 +
10 +from .ops import rmsnorm
11 +
12 +
13 +class InklingRMSNorm(nn.Module):
14 + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
15 + super().__init__()
16 + self.weight = nn.Parameter(torch.ones(hidden_size))
17 + self.variance_epsilon = eps
18 + self.hidden_size = hidden_size
19 +
20 + def forward(self, x: torch.Tensor) -> torch.Tensor:
21 + if x.numel() == 0:
22 + return x
23 + original_shape = x.shape
24 + x_2d = x.contiguous().view(-1, self.hidden_size)
25 + y = rmsnorm(x_2d, self.weight, self.variance_epsilon)
26 + return y.view(original_shape)
vllm/models/inkling/nvidia/logits_processor.py new
+73
@@ -0,0 +1,73 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inkling logits processor with muP output scaling."""
4 +
5 +from __future__ import annotations
6 +
7 +import torch
8 +
9 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
10 +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
11 +
12 +
13 +class InklingLogitsProcessor(LogitsProcessor):
14 + """``LogitsProcessor`` that applies Inkling's muP logits width multiplier.
15 +
16 + Args:
17 + vocab_size: Padded vocabulary size.
18 + org_vocab_size: Unpadded vocabulary size (defaults to ``vocab_size``).
19 + scale: Base logits scale (kept ``1.0`` for the served checkpoint).
20 + logits_as_input: Whether the input is already logits.
21 + soft_cap: Optional logit soft cap (``None`` for the served checkpoint).
22 + logits_mup_width_multiplier: muP width divisor for the final logits;
23 + ``None`` or ``0`` disables it.
24 + """
25 +
26 + def __init__(
27 + self,
28 + vocab_size: int,
29 + org_vocab_size: int | None = None,
30 + scale: float = 1.0,
31 + logits_as_input: bool = False,
32 + soft_cap: float | None = None,
33 + logits_mup_width_multiplier: float | None = None,
34 + ) -> None:
35 + super().__init__(
36 + vocab_size=vocab_size,
37 + org_vocab_size=org_vocab_size,
38 + scale=scale,
39 + logits_as_input=logits_as_input,
40 + soft_cap=soft_cap,
41 + )
42 + self.logits_mup_width_multiplier = logits_mup_width_multiplier
43 + self._logits_zero: torch.Tensor | None = None
44 +
45 + def forward(
46 + self,
47 + lm_head: VocabParallelEmbedding,
48 + hidden_states: torch.Tensor,
49 + embedding_bias: torch.Tensor | None = None,
50 + ) -> torch.Tensor | None:
51 + mup = self.logits_mup_width_multiplier
52 + if not mup:
53 + return super().forward(lm_head, hidden_states, embedding_bias)
54 + # Fold the muP width divisor into the lm_head GEMM alpha (fp32 epilogue):
55 + # no separate elementwise kernel, no bf16 rounding of scaled logits, and
56 + # no weight mutation. Overfit to the served checkpoint: bf16 lm_head, no
57 + # soft cap, unit logits scale.
58 + assert self.soft_cap is None
59 + assert self.scale == 1.0
60 + w = lm_head.weight
61 + if self._logits_zero is None:
62 + self._logits_zero = w.new_zeros(1)
63 + logits = torch.addmm(
64 + self._logits_zero,
65 + hidden_states,
66 + w.t(),
67 + beta=0.0,
68 + alpha=1.0 / mup,
69 + )
70 + logits = self._gather_logits(logits)
71 + if logits is not None:
72 + logits = logits[..., : self.org_vocab_size]
73 + return logits
vllm/models/inkling/nvidia/mlp.py new
+63
@@ -0,0 +1,63 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inkling dense SwiGLU MLP (also used as the MoE shared expert).
4 +
5 +The checkpoint stores the gate/up projection as a single fused, *interleaved*
6 +weight (``[gate0, up0, gate1, up1, ...]``), so we use a plain
7 +``ColumnParallelLinear`` whose contiguous row-sharding keeps each gate/up pair
8 +together, and an interleaved SwiGLU activation.
9 +"""
10 +
11 +from __future__ import annotations
12 +
13 +import torch
14 +from torch import nn
15 +
16 +from vllm.model_executor.layers.linear import (
17 + ColumnParallelLinear,
18 + RowParallelLinear,
19 +)
20 +from vllm.model_executor.layers.quantization import QuantizationConfig
21 +
22 +
23 +class InklingDenseMLP(nn.Module):
24 + def __init__(
25 + self,
26 + hidden_size: int,
27 + intermediate_size: int,
28 + *,
29 + use_global_scale: bool = False,
30 + quant_config: QuantizationConfig | None = None,
31 + prefix: str = "",
32 + ) -> None:
33 + super().__init__()
34 + self.gate_up_proj = ColumnParallelLinear(
35 + hidden_size,
36 + 2 * intermediate_size,
37 + bias=False,
38 + quant_config=quant_config,
39 + prefix=f"{prefix}.gate_up_proj",
40 + )
41 + self.down_proj = RowParallelLinear(
42 + intermediate_size,
43 + hidden_size,
44 + bias=False,
45 + quant_config=quant_config,
46 + reduce_results=False,
47 + prefix=f"{prefix}.down_proj",
48 + )
49 + if use_global_scale:
50 + self.global_scale = nn.Parameter(torch.empty(1), requires_grad=False)
51 + else:
52 + self.global_scale = None
53 +
54 + def forward(self, x: torch.Tensor) -> torch.Tensor:
55 + from .ops import silu_and_mul_triton
56 +
57 + gate_up, _ = self.gate_up_proj(x)
58 + x = silu_and_mul_triton(gate_up)
59 + x, _ = self.down_proj(x)
60 + if self.global_scale is not None:
61 + x = x * self.global_scale
62 + # TP-partial output: the layer's reduce-scatter fallback consumes it.
63 + return x
vllm/models/inkling/nvidia/model.py new
+678
@@ -0,0 +1,678 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inkling model implementation for NVIDIA GPUs."""
4 +
5 +from __future__ import annotations
6 +
7 +from collections.abc import Iterable
8 +from typing import Any
9 +
10 +import regex as re
11 +import torch
12 +from torch import nn
13 +
14 +from vllm.config import VllmConfig
15 +from vllm.distributed import (
16 + get_pp_group,
17 + get_tensor_model_parallel_rank,
18 + get_tensor_model_parallel_world_size,
19 + tensor_model_parallel_all_gather,
20 + tensor_model_parallel_reduce_scatter,
21 +)
22 +from vllm.forward_context import get_forward_context
23 +from vllm.model_executor.layers.quantization import QuantizationConfig
24 +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
25 +from vllm.model_executor.models.interfaces import (
26 + MultiModalEmbeddings,
27 + SupportsMultiModal,
28 + SupportsPP,
29 +)
30 +from vllm.model_executor.models.utils import (
31 + AutoWeightsLoader,
32 + WeightsMapper,
33 + make_empty_intermediate_tensors_factory,
34 + make_layers,
35 + maybe_prefix,
36 +)
37 +from vllm.models.inkling.common.mm_preprocess import (
38 + InklingDummyInputsBuilder,
39 + InklingMultiModalProcessor,
40 + InklingProcessingInfo,
41 + inkling_audio_enabled,
42 + inkling_vision_enabled,
43 +)
44 +from vllm.models.inkling.common.towers import InklingAudio, InklingVision
45 +from vllm.multimodal import MULTIMODAL_REGISTRY
46 +from vllm.sequence import IntermediateTensors
47 +
48 +from ..configs import InklingMMConfig, InklingModelConfig
49 +from ..nvfp4 import InklingNvfp4Config
50 +from .attention import InklingAttention, compute_log_scaling_tau
51 +from .layernorm import InklingRMSNorm
52 +from .logits_processor import InklingLogitsProcessor
53 +from .mlp import InklingDenseMLP
54 +from .moe import InklingMoE
55 +from .ops.lamport import get_lamport_rs_conv, initialize_lamport_rs_conv
56 +from .ops.norm import add_rmsnorm, embed_rmsnorm
57 +from .sconv_swa_attn import _ATTN, _MLP, InklingConvState, InklingSconvMetadata
58 +from .short_conv import InklingShortConv
59 +
60 +
61 +def _layer_id(name: str) -> int | None:
62 + m = re.search(r"\.layers\.(\d+)\.", name)
63 + return int(m.group(1)) if m else None
64 +
65 +
66 +def _sconv_add_norm(
67 + delta: torch.Tensor,
68 + hidden: torch.Tensor,
69 + sconv: InklingShortConv,
70 + norm: InklingRMSNorm | None,
71 + positions: torch.Tensor,
72 +) -> tuple[torch.Tensor | None, torch.Tensor]:
73 + """``h = hidden + sconv(TP-sum(delta)); y = rmsnorm(h)``.
74 +
75 + The Lamport path performs reduce-scatter + shard sconv + all-gather +
76 + residual add + norm. The NCCL path handles unsupported configurations."""
77 + attn_metadata = get_forward_context().attn_metadata
78 + m = (
79 + attn_metadata.get(sconv.owner.prefix)
80 + if isinstance(attn_metadata, dict)
81 + else None
82 + )
83 + cache = sconv.owner.kv_cache
84 + off_s, ws = sconv.owner.stream_ranges[sconv.stream_idx]
85 + norm_w = norm.weight if norm is not None else None
86 + eps = norm.variance_epsilon if norm is not None else 0.0
87 +
88 + mm = get_lamport_rs_conv(hidden.shape[-1], sconv.kernel_size)
89 + if mm is not None and mm.usable(delta.shape[0]) and m is not None:
90 + assert cache.numel() > 0
91 + assert isinstance(m, InklingSconvMetadata)
92 + return mm.rs_sconv_ag_add_norm(
93 + delta,
94 + hidden,
95 + sconv.weight.squeeze(1),
96 + norm_w,
97 + eps,
98 + cache,
99 + positions,
100 + m.block_table,
101 + m.seq_idx,
102 + m.slot_mapping,
103 + off_s,
104 + ws,
105 + sconv.owner.block_size,
106 + )
107 +
108 + # Fallback: NCCL RS -> shard sconv -> AG -> fused add(+rmsnorm).
109 + shard = tensor_model_parallel_reduce_scatter(delta, dim=-1)
110 + shard = sconv(shard.contiguous(), positions)
111 + full = tensor_model_parallel_all_gather(shard, dim=-1)
112 + if norm is None:
113 + return None, hidden + full
114 + return add_rmsnorm(hidden, full, norm_w, eps)
115 +
116 +
117 +class InklingDecoderLayer(nn.Module):
118 + def __init__(
119 + self,
120 + config: InklingModelConfig,
121 + layer_id: int,
122 + is_local: bool,
123 + quant_config: QuantizationConfig | None,
124 + prefix: str,
125 + nvfp4_config: InklingNvfp4Config | None = None,
126 + ) -> None:
127 + super().__init__()
128 + # Per-layer owner of the conv state as a paged SWA cache. The 4 sconv
129 + # streams (K/V/attn/mlp) are packed head-major into one block and share
130 + # it. Built first so the attention layer can wire its K/V sconv to it.
131 + self.conv_state = InklingConvState(
132 + num_kv_heads=(
133 + config.swa_num_key_value_heads
134 + if is_local
135 + else config.num_key_value_heads
136 + ),
137 + head_dim=config.swa_head_dim if is_local else config.head_dim,
138 + hidden_size=config.hidden_size,
139 + kernel_size=config.sconv_kernel_size,
140 + prefix=f"{prefix}.conv_state",
141 + )
142 + self.attn_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
143 + self.attn = InklingAttention(
144 + config,
145 + num_heads=(
146 + config.swa_num_attention_heads
147 + if is_local
148 + else config.num_attention_heads
149 + ),
150 + num_kv_heads=(
151 + config.swa_num_key_value_heads
152 + if is_local
153 + else config.num_key_value_heads
154 + ),
155 + head_dim=config.swa_head_dim if is_local else config.head_dim,
156 + rel_extent=config.rel_extent,
157 + local_extent=config.sliding_window_size,
158 + is_local=is_local,
159 + prefix=f"{prefix}.attn",
160 + quant_config=quant_config,
161 + conv_owner=self.conv_state,
162 + )
163 + self.mlp_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
164 + if layer_id < config.dense_mlp_idx:
165 + self.mlp: nn.Module = InklingDenseMLP(
166 + hidden_size=config.hidden_size,
167 + intermediate_size=config.dense_intermediate_size,
168 + use_global_scale=config.use_global_scale,
169 + quant_config=quant_config,
170 + prefix=f"{prefix}.mlp",
171 + )
172 + else:
173 + # InklingMoE decides per layer (from the checkpoint exclude list)
174 + # whether the routed experts are NVFP4 or bf16; the shared sink
175 + # experts are always bf16.
176 + self.mlp = InklingMoE(
177 + config,
178 + layer_id,
179 + prefix=f"{prefix}.mlp",
180 + nvfp4_config=nvfp4_config,
181 + )
182 +
183 + # Short convolution on the attention-output and MLP-output residual
184 + # streams, hidden-sharded: the sublayer outputs are reduce-scattered
185 + # to [T, H/tp], the sconv runs on the shard, and an all-gather
186 + # restores the full residual — all fused with the residual add + next
187 + # rmsnorm via the Lamport P2P kernels for decode-sized batches.
188 + tp_size = get_tensor_model_parallel_world_size()
189 + sconv_dim = config.hidden_size // tp_size
190 + self.attn_sconv = InklingShortConv(
191 + sconv_dim, config.sconv_kernel_size, owner=self.conv_state, stream_idx=_ATTN
192 + )
193 + self.mlp_sconv = InklingShortConv(
194 + sconv_dim, config.sconv_kernel_size, owner=self.conv_state, stream_idx=_MLP
195 + )
196 +
197 + def forward(
198 + self,
199 + positions: torch.Tensor,
200 + hidden_states: torch.Tensor,
201 + pending: tuple[torch.Tensor | None, InklingShortConv] | None = None,
202 + attn_in: torch.Tensor | None = None,
203 + log_scaling: torch.Tensor | None = None,
204 + ) -> tuple[torch.Tensor, tuple[torch.Tensor | None, InklingShortConv]]:
205 + # The previous sublayer's (pre-reduce, pre-sconv) delta is folded in
206 + # fused with its RS/sconv/AG and this layer's pre-attention rmsnorm.
207 + # A None delta means the partials sit in the NVLS symm buffer.
208 + if pending is None:
209 + if attn_in is None:
210 + # First layer; on the text path attn_norm comes fused with
211 + # the embedding gather (chain_weight in embed_rmsnorm).
212 + attn_in = self.attn_norm(hidden_states)
213 + else:
214 + attn_in, hidden_states = _sconv_add_norm(
215 + pending[0], hidden_states, pending[1], self.attn_norm, positions
216 + )
217 + attn_output = self.attn(positions, attn_in, log_scaling)
218 + mlp_in, hidden_states = _sconv_add_norm(
219 + attn_output, hidden_states, self.attn_sconv, self.mlp_norm, positions
220 + )
221 + mlp_output = self.mlp(mlp_in)
222 + # The caller folds mlp_output (pre-reduce, pre-sconv) into the next
223 + # fused sconv+add+rmsnorm.
224 + return hidden_states, (mlp_output, self.mlp_sconv)
225 +
226 +
227 +class InklingReplicatedEmbedding(nn.Module):
228 + """Full-vocab embedding table replicated on every TP rank.
229 +
230 + Trades the full table per rank (~2.3 GiB at V=201k / H=6144 bf16, vs a
231 + 1/tp shard) for no masked lookup or per-lookup TP all-reduce, and keeps the
232 + full table on-rank for the fused gather+norm kernel. Bit-exact vs
233 + vocab-parallel: the all-reduce there only ever summed one real row against
234 + exact zeros. The LM head stays vocab-sharded.
235 + """
236 +
237 + def __init__(self, num_embeddings: int, embedding_dim: int) -> None:
238 + super().__init__()
239 + self.weight = nn.Parameter(
240 + torch.empty(num_embeddings, embedding_dim, dtype=torch.get_default_dtype()),
241 + requires_grad=False,
242 + )
243 +
244 + def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
245 + return embed_rmsnorm(input_ids, self.weight, None, 0.0)
246 +
247 +
248 +class InklingModel(nn.Module):
249 + def __init__(
250 + self,
251 + *,
252 + config: InklingModelConfig,
253 + quant_config: QuantizationConfig | None,
254 + prefix: str,
255 + nvfp4_config: InklingNvfp4Config | None = None,
256 + ) -> None:
257 + super().__init__()
258 + self.config = config
259 + self.embed_tokens = InklingReplicatedEmbedding(
260 + config.padded_vocab_size, config.hidden_size
261 + )
262 + self.embed_norm = (
263 + InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
264 + if config.use_embed_norm
265 + else None
266 + )
267 + local_ids = set(config.local_layer_ids)
268 +
269 + def get_layer(prefix: str) -> InklingDecoderLayer:
270 + idx = _layer_id(prefix + ".") or int(prefix.split(".")[-1])
271 + return InklingDecoderLayer(
272 + config, idx, idx in local_ids, quant_config, prefix, nvfp4_config
273 + )
274 +
275 + self.start_layer, self.end_layer, self.layers = make_layers(
276 + config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
277 + )
278 + self.norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
279 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
280 + ["hidden_states"], config.hidden_size
281 + )
282 +
283 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
284 + # Row gather + embed_norm in one launch.
285 + norm = self.embed_norm
286 + return embed_rmsnorm(
287 + input_ids,
288 + self.embed_tokens.weight,
289 + norm.weight if norm is not None else None,
290 + norm.variance_epsilon if norm is not None else 0.0,
291 + )
292 +
293 + def forward(
294 + self,
295 + input_ids: torch.Tensor | None,
296 + positions: torch.Tensor,
297 + intermediate_tensors: IntermediateTensors | None,
298 + inputs_embeds: torch.Tensor | None = None,
299 + ) -> torch.Tensor | IntermediateTensors:
300 + attn_in0: torch.Tensor | None = None
301 + if get_pp_group().is_first_rank:
302 + if inputs_embeds is not None:
303 + # embed_norm was already applied when producing inputs_embeds.
304 + hidden_states = inputs_embeds
305 + else:
306 + # Gather + embed_norm + the first layer's attn_norm, one launch.
307 + norm = self.embed_norm
308 + hidden_states, attn_in0 = embed_rmsnorm(
309 + input_ids,
310 + self.embed_tokens.weight,
311 + norm.weight if norm is not None else None,
312 + self.config.rms_norm_eps,
313 + chain_weight=self.layers[self.start_layer].attn_norm.weight,
314 + )
315 + else:
316 + assert intermediate_tensors is not None
317 + hidden_states = intermediate_tensors["hidden_states"]
318 + hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
319 + log_scaling = None
320 + if self.config.log_scaling_n_floor is not None:
321 + log_scaling = compute_log_scaling_tau(
322 + positions,
323 + self.config.log_scaling_n_floor,
324 + self.config.log_scaling_alpha,
325 + )
326 +
327 + pending: tuple[torch.Tensor | None, InklingShortConv] | None = None
328 + for layer in self.layers[self.start_layer : self.end_layer]:
329 + hidden_states, pending = layer(
330 + positions,
331 + hidden_states,
332 + pending=pending,
333 + attn_in=attn_in0,
334 + log_scaling=log_scaling,
335 + )
336 + attn_in0 = None
337 +
338 + if not get_pp_group().is_last_rank:
339 + if pending is not None:
340 + hidden_states = _sconv_add_norm(
341 + pending[0], hidden_states, pending[1], None, positions
342 + )[1]
343 + return IntermediateTensors({"hidden_states": hidden_states})
344 + if pending is not None:
345 + # Final RS/sconv/AG + residual add fused with the final rmsnorm.
346 + norm_out = _sconv_add_norm(
347 + pending[0], hidden_states, pending[1], self.norm, positions
348 + )[0]
349 + assert norm_out is not None
350 + return norm_out
351 + return self.norm(hidden_states)
352 +
353 +
354 +class _TmlForCausalLMBase(nn.Module, SupportsPP):
355 + """Shared text-backbone causal-LM scaffolding for both entry classes."""
356 +
357 + hf_to_vllm_mapper = WeightsMapper(
358 + orig_to_new_substr={
359 + ".w13_dn": ".gate_up_proj",
360 + ".w2_md": ".down_proj",
361 + },
362 + orig_to_new_stacked={
363 + ".attn.wq_du.": (".attn.qkvr.", 0),
364 + ".attn.wk_dv.": (".attn.qkvr.", 1),
365 + ".attn.wv_dv.": (".attn.qkvr.", 2),
366 + ".attn.wr_du.": (".attn.qkvr.", 3),
367 + },
368 + orig_to_new_prefix={
369 + "model.llm.layers.": "model.layers.",
370 + "model.llm.embed_norm": "model.embed_norm",
371 + "model.llm.embed": "model.embed_tokens",
372 + "model.llm.norm": "model.norm",
373 + "model.llm.unembed": "lm_head",
374 + },
375 + orig_to_new_suffix={
376 + # NVFP4 scale
377 + ".w13_weight.scale": ".w13_weight_scale",
378 + ".w13_weight.scale2": ".w13_weight_scale_2",
379 + ".w2_weight.scale": ".w2_weight_scale",
380 + ".w2_weight.scale2": ".w2_weight_scale_2",
381 + },
382 + )
383 +
384 + def _build(
385 + self,
386 + vllm_config: VllmConfig,
387 + text_config: InklingModelConfig,
388 + prefix: str,
389 + ) -> None:
390 + quant_config = vllm_config.quant_config
391 + self.config = text_config
392 + # NVFP4 experts are detected directly from the checkpoint quant config;
393 + # only the MoE experts are quantized (attention/dense MLP stay bf16).
394 + self.nvfp4_config = InklingNvfp4Config.from_hf_config(
395 + vllm_config.model_config.hf_config
396 + )
397 + # Read by the MRV2 runner to publish per-request short-conv metadata.
398 + # Short convolution is intrinsic to Inkling, so this is always set.
399 + self.uses_sconv = True
400 + self.model = InklingModel(
401 + config=text_config,
402 + quant_config=quant_config,
403 + prefix=maybe_prefix(prefix, "model"),
404 + nvfp4_config=self.nvfp4_config,
405 + )
406 + initialize_lamport_rs_conv(
407 + text_config.hidden_size,
408 + text_config.sconv_kernel_size,
409 + vllm_config.scheduler_config.max_num_batched_tokens,
410 + )
411 + self.lm_head = ParallelLMHead(
412 + text_config.padded_vocab_size,
413 + text_config.hidden_size,
414 + org_num_embeddings=text_config.padded_vocab_size,
415 + quant_config=quant_config,
416 + prefix=maybe_prefix(prefix, "lm_head"),
417 + )
418 + self.logits_processor = InklingLogitsProcessor(
419 + text_config.padded_vocab_size,
420 + org_vocab_size=text_config.vocab_size,
421 + soft_cap=text_config.final_logit_softcapping,
422 + logits_mup_width_multiplier=text_config.logits_mup_width_multiplier,
423 + )
424 + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
425 + self.model.make_empty_intermediate_tensors
426 + )
427 +
428 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
429 + return self.model.embed_input_ids(input_ids)
430 +
431 + def forward(
432 + self,
433 + input_ids: torch.Tensor | None,
434 + positions: torch.Tensor,
435 + intermediate_tensors: IntermediateTensors | None = None,
436 + inputs_embeds: torch.Tensor | None = None,
437 + **kwargs: object,
438 + ) -> torch.Tensor | IntermediateTensors:
439 + return self.model(
440 + input_ids,
441 + positions,
442 + intermediate_tensors,
443 + inputs_embeds,
444 + )
445 +
446 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
447 + return self.logits_processor(self.lm_head, hidden_states)
448 +
449 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
450 + return _load_inkling_weights(self, weights, self.config)
451 +
452 +
453 +class InklingForCausalLM(_TmlForCausalLMBase):
454 + """Text-only entry point (``inkling_model`` checkpoints)."""
455 +
456 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
457 + super().__init__()
458 + self._build(vllm_config, vllm_config.model_config.hf_config, prefix)
459 +
460 +
461 +@MULTIMODAL_REGISTRY.register_processor(
462 + InklingMultiModalProcessor,
463 + info=InklingProcessingInfo,
464 + dummy_inputs=InklingDummyInputsBuilder,
465 +)
466 +class InklingForConditionalGeneration(_TmlForCausalLMBase, SupportsMultiModal):
467 + """Top-level (multimodal) entry point.
468 +
469 + Builds the vision + audio towers on top of the shared text backbone. Inkling has
470 + NO cross-modal fusion (the vision tower emits one token per patch, the audio
471 + tower one token per frame), so generation reuses the inherited backbone
472 + ``forward`` / ``compute_logits`` (the latter already applies muP) and this
473 + class only adds multimodal embedding + merge.
474 + """
475 +
476 + hf_to_vllm_mapper = _TmlForCausalLMBase.hf_to_vllm_mapper | WeightsMapper(
477 + orig_to_new_prefix={
478 + "model.audio.": "audio.",
479 + "model.visual.": "visual.vision_encoder.",
480 + },
481 + )
482 +
483 + @classmethod
484 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:
485 + if modality.startswith("image"):
486 + return "<|content_image|>"
487 + if modality.startswith("audio"):
488 + return "<|content_audio_input|>"
489 + raise ValueError("Only image or audio modality is supported")
490 +
491 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
492 + super().__init__()
493 + config: InklingMMConfig = vllm_config.model_config.hf_config
494 +
495 + self.visual = (
496 + InklingVision(config.vision_config, prefix=maybe_prefix(prefix, "visual"))
497 + if inkling_vision_enabled(config)
498 + else None
499 + )
500 + self.audio = (
501 + InklingAudio(config.audio_config, prefix=maybe_prefix(prefix, "audio"))
502 + if inkling_audio_enabled(config)
503 + else None
504 + )
505 +
506 + self._build(vllm_config, config.text_config, prefix)
507 +
508 + # -- multimodal embedding -------------------------------------------
509 +
510 + def _process_image_input(
511 + self, pixel_values: Any, num_patches: Any
512 + ) -> tuple[torch.Tensor, ...]:
513 + assert self.visual is not None
514 + # pixel_values is a list (per item) of [P_i, 2, P, P, 3] tensors,
515 + # or a single concatenated tensor. Normalize to a flat batch, run the
516 + # tower once, then split back per item.
517 + if isinstance(pixel_values, (list, tuple)):
518 + if not pixel_values:
519 + return ()
520 + sizes = [int(p.shape[0]) for p in pixel_values]
521 + patches = torch.cat(list(pixel_values), dim=0)
522 + else:
523 + patches = pixel_values
524 + sizes = self._sizes_from(num_patches, patches.shape[0])
525 +
526 + patches = patches.to(device=self.visual.device, dtype=self.visual.dtype)
527 + embeds = self.visual(patches) # [total_patches, D]
528 + return tuple(embeds.split(sizes))
529 +
530 + def _process_audio_input(
531 + self, input_audio_features: Any, num_audio_tokens: Any
532 + ) -> tuple[torch.Tensor, ...]:
533 + assert self.audio is not None
534 + if isinstance(input_audio_features, (list, tuple)):
535 + if not input_audio_features:
536 + return ()
537 + sizes = [int(d.shape[0]) for d in input_audio_features]
538 + dmel = torch.cat(list(input_audio_features), dim=0)
539 + else:
540 + dmel = input_audio_features
541 + sizes = self._sizes_from(num_audio_tokens, dmel.shape[0])
542 +
543 + dmel = dmel.to(device=self.audio.device)
544 + embeds = self.audio(dmel) # [total_frames, D]
545 + return tuple(embeds.split(sizes))
546 +
547 + @staticmethod
548 + def _sizes_from(counts: Any, total: int) -> list[int]:
549 + if counts is None:
550 + return [total]
551 + if isinstance(counts, torch.Tensor):
552 + return [int(c) for c in counts.flatten().tolist()]
553 + if isinstance(counts, (list, tuple)):
554 + flat: list[int] = []
555 + for c in counts:
556 + flat.append(int(c.item()) if isinstance(c, torch.Tensor) else int(c))
557 + return flat
558 + return [int(counts)]
559 +
560 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
561 + # Iterate modalities in a stable order so the returned per-item tensors
562 + # line up with their appearance order; the positional merge in
563 + # embed_input_ids handles actual placement.
564 + pixel_values = kwargs.get("pixel_values")
565 + num_patches = kwargs.get("num_patches")
566 + input_audio_features = kwargs.get("input_audio_features")
567 + num_audio_tokens = kwargs.get("num_audio_tokens")
568 +
569 + embeddings: tuple[torch.Tensor, ...] = ()
570 + if pixel_values is not None and self.visual is not None:
571 + embeddings += self._process_image_input(pixel_values, num_patches)
572 + if input_audio_features is not None and self.audio is not None:
573 + embeddings += self._process_audio_input(
574 + input_audio_features, num_audio_tokens
575 + )
576 + return embeddings
577 +
578 + def embed_input_ids(
579 + self,
580 + input_ids: torch.Tensor,
581 + multimodal_embeddings: MultiModalEmbeddings | None = None,
582 + *,
583 + is_multimodal: torch.Tensor | None = None,
584 + ) -> torch.Tensor:
585 + # Override the base's 1-arg embed_input_ids: the runner calls this 3-arg
586 + # signature for multimodal models. Text embeddings come from the shared
587 + # backbone (which applies embed_norm); MM embeddings are scattered in.
588 + from vllm.model_executor.models.utils import _merge_multimodal_embeddings
589 +
590 + # Placeholder ids use unused vocabulary slots and these positions are
591 + # overwritten by MM embeds below.
592 + inputs_embeds = self.model.embed_input_ids(input_ids)
593 + if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
594 + return inputs_embeds
595 + assert is_multimodal is not None
596 + return _merge_multimodal_embeddings(
597 + inputs_embeds=inputs_embeds,
598 + multimodal_embeddings=multimodal_embeddings,
599 + is_multimodal=is_multimodal,
600 + )
601 +
602 + def get_language_model(self) -> nn.Module:
603 + # This class IS the causal LM (the towers are side branches), so the
604 + # language model is self — callers expect a module exposing ``.model``
605 + # and ``.lm_head``.
606 + return self
607 +
608 +
609 +# ===========================================================================
610 +# Weight loading
611 +# ===========================================================================
612 +
613 +
614 +_MOE_EXPERT_WEIGHT_RE = re.compile(
615 + r"^(?P<mlp>.*\.mlp)\.(?P<rest>(?:shared_)?experts\..+)$"
616 +)
617 +
618 +
619 +def _load_inkling_weights(
620 + module: nn.Module,
621 + weights: Iterable[tuple[str, torch.Tensor]],
622 + config: InklingModelConfig,
623 +) -> set[str]:
624 + moe_modules = {
625 + name: mod for name, mod in module.named_modules() if isinstance(mod, InklingMoE)
626 + }
627 + loaded: set[str] = set()
628 + tp_size = get_tensor_model_parallel_world_size()
629 + tp_rank = get_tensor_model_parallel_rank()
630 + local_ids = set(config.local_layer_ids)
631 +
632 + def _iter_loadable_weights() -> Iterable[tuple[str, torch.Tensor]]:
633 + for name, weight in module.hf_to_vllm_mapper.apply(weights):
634 + shard_id = getattr(weight, "shard_id", None)
635 + # Replicate K/V conv-free GQA heads when tp_size > num_kv_heads.
636 + if (
637 + shard_id in (1, 2)
638 + and name.endswith(".attn.qkvr.weight")
639 + and weight.shape[0] > 0
640 + ):
641 + lid = _layer_id(name)
642 + if lid is not None:
643 + is_local = lid in local_ids
644 + n_kv = (
645 + config.swa_num_key_value_heads
646 + if is_local
647 + else config.num_key_value_heads
648 + )
649 + head_dim = config.swa_head_dim if is_local else config.head_dim
650 + if tp_size > n_kv and weight.shape[0] == n_kv * head_dim:
651 + kv_idx = (tp_rank * n_kv) // tp_size
652 + weight = weight.narrow(0, kv_idx * head_dim, head_dim)
653 + weight.shard_id = shard_id
654 +
655 + # MoE expert tensors (fused stacked, routed + shared sink): translate
656 + # the checkpoint layout to per-expert FusedMoE loads.
657 + moe_match = _MOE_EXPERT_WEIGHT_RE.match(name)
658 + if moe_match is not None and moe_match.group("mlp") in moe_modules:
659 + moe = moe_modules[moe_match.group("mlp")]
660 + for rel in moe.load_expert_weight(moe_match.group("rest"), weight):
661 + loaded.add(f"{moe_match.group('mlp')}.{rel}")
662 + continue
663 +
664 + yield name, weight
665 +
666 + # The release checkpoint also carries auxiliary prediction-head weights;
667 + # they are not part of the causal LM served by this implementation.
668 + loader = AutoWeightsLoader(module, skip_prefixes=["model.mtp."])
669 + loaded |= loader.load_weights(_iter_loadable_weights())
670 +
671 + # Post-load MoE fixups (default input scales, zeroed EP-padding experts).
672 + for moe_name, moe in moe_modules.items():
673 + for rel in moe.finalize_load():
674 + loaded.add(f"{moe_name}.{rel}")
675 + return loaded
676 +
677 +
678 +EntryClass = [InklingForCausalLM, InklingForConditionalGeneration]
vllm/models/inkling/nvidia/moe.py new
+578
@@ -0,0 +1,578 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inkling mixture-of-experts on vLLM's FusedMoE abstraction.
4 +
5 +Overfit to the served checkpoint: sigmoid gate (+ selection bias) top-k over
6 +the routed experts, log-sigmoid renormalization over the k routed + S shared
7 +"sink" logits, scaled by route_scale * global_scale. The routed top-k goes
8 +through vLLM's FusedMoE (which handles TP/EP); the sink experts run in
9 +:class:`InklingSinkExperts` -- replicated across EP ranks (every token
10 +activates every sink) and always bf16 (the checkpoint excludes every
11 +``shared_experts`` from quantization).
12 +
13 +NVFP4 routed experts reuse vLLM's ModelOpt NVFP4 fused-MoE method; excluded
14 +(bf16) layers fall back to the unquantized method. The checkpoint's fused
15 +stacked tensors (interleaved gate/up rows, ``.scale`` / ``.scale2`` /
16 +``.input_amax`` aux tensors) are translated to the standard per-expert loads
17 +in :meth:`InklingMoE.load_expert_weight`.
18 +"""
19 +
20 +from __future__ import annotations
21 +
22 +import math
23 +from typing import TYPE_CHECKING
24 +
25 +import torch
26 +from torch import nn
27 +from torch.nn.parameter import Parameter
28 +
29 +import vllm.envs as envs
30 +from vllm.config import get_current_vllm_config
31 +from vllm.distributed import (
32 + get_dp_group,
33 + get_pcp_group,
34 + get_tensor_model_parallel_rank,
35 + get_tensor_model_parallel_world_size,
36 +)
37 +from vllm.model_executor.kernels.linear.cute_dsl import ll_bf16
38 +from vllm.model_executor.layers.fused_moe import FusedMoE
39 +from vllm.model_executor.utils import set_weight_attrs
40 +from vllm.platforms import current_platform
41 +from vllm.triton_utils import tl, tldevice, triton
42 +from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
43 +from vllm.utils.torch_utils import aux_stream
44 +
45 +from ..configs import InklingModelConfig
46 +from ..nvfp4 import FLOAT4_E2M1_MAX, FLOAT8_E4M3_MAX
47 +
48 +if TYPE_CHECKING:
49 + from vllm.model_executor.layers.fused_moe.routed_experts import (
50 + RoutedExperts,
51 + )
52 +
53 + from ..nvfp4 import InklingNvfp4Config
54 +
55 +# ---------------------------------------------------------------------------
56 +# Gate / expert selection
57 +# ---------------------------------------------------------------------------
58 +
59 +_INKLING_LL_BF16_MAX_TOKENS = 64
60 +
61 +
62 +def _linear_with_fp32_out(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
63 + leading = list(x.shape[:-1])
64 + flat = x.flatten(0, -2)
65 + if (
66 + flat.shape[0] <= _INKLING_LL_BF16_MAX_TOKENS
67 + and flat.dtype == torch.bfloat16
68 + and weight.dtype == torch.bfloat16
69 + and flat.is_cuda
70 + and flat.is_contiguous()
71 + and weight.is_contiguous()
72 + and flat.shape[1] % 8 == 0
73 + and current_platform.has_device_capability(90)
74 + and ll_bf16.is_available()
75 + ):
76 + out = ll_bf16.ll_bf16_gemm(flat, weight)
77 + else:
78 + out = torch.mm(flat, weight.T, out_dtype=torch.float32)
79 + return out.view(*leading, weight.shape[0])
80 +
81 +
82 +@triton.jit(do_not_specialize=["T", "route_scale"])
83 +def _inkling_gate_select_kernel(
84 + logits_ptr, # [T, G] fp32 gate logits (stride_logits_0 may include pad)
85 + bias_ptr, # [R] fp32 selection bias (or 0 ptr if HAS_BIAS=False)
86 + global_scale_ptr, # [1] fp32 (or unused if HAS_GSCALE=False)
87 + ids_ptr, # [T, K + S] int32 out: selected expert ids
88 + weights_ptr, # [T, K + S] fp32 out: renormalized weights
89 + route_scale,
90 + T,
91 + G: tl.constexpr, # total gate experts (routed + shared)
92 + stride_logits_0,
93 + R: tl.constexpr, # routed experts
94 + K: tl.constexpr, # top-k routed
95 + S: tl.constexpr, # shared (sink) experts
96 + HAS_BIAS: tl.constexpr,
97 + HAS_GSCALE: tl.constexpr,
98 + BLOCK_G: tl.constexpr,
99 +):
100 + pid = tl.program_id(0).to(tl.int64)
101 + if pid >= T:
102 + return
103 + offs = tl.arange(0, BLOCK_G)
104 + mask_r = offs < R
105 + logits = tl.load(
106 + logits_ptr + pid * stride_logits_0 + offs,
107 + mask=offs < G,
108 + other=float("-inf"),
109 + ).to(tl.float32)
110 +
111 + # Selection scores: sigmoid(routed logits) (+ bias), non-routed lanes -inf.
112 + sel = tl.where(mask_r, tl.sigmoid(logits), float("-inf"))
113 + if HAS_BIAS:
114 + bias = tl.load(bias_ptr + offs, mask=mask_r, other=0.0).to(tl.float32)
115 + sel = tl.where(mask_r, sel + bias, float("-inf"))
116 +
117 + scale = route_scale
118 + if HAS_GSCALE:
119 + scale = scale * tl.load(global_scale_ptr).to(tl.float32)
120 +
121 + # Iterative top-K (K is small); argmax tie-breaks to the lowest index
122 + # (stable ordering).
123 + A: tl.constexpr = K + S
124 + offs_a = tl.arange(0, A)
125 + top_ids = tl.zeros([A], dtype=tl.int32)
126 + active = tl.zeros([A], dtype=tl.float32)
127 + for kk in tl.static_range(K):
128 + idx = tl.argmax(sel, axis=0).to(tl.int32)
129 + raw = tl.max(tl.where(offs == idx, logits, float("-inf")), axis=0)
130 + top_ids = tl.where(offs_a == kk, idx, top_ids)
131 + active = tl.where(offs_a == kk, raw, active)
132 + sel = tl.where(offs == idx, float("-inf"), sel)
133 + if S > 0:
134 + # Shared sink logits sit at the tail of the gate output; their expert
135 + # ids continue after the routed range (R + j).
136 + for jj in tl.static_range(S):
137 + raw = tl.max(tl.where(offs == R + jj, logits, float("-inf")), axis=0)
138 + top_ids = tl.where(offs_a == K + jj, tl.full([], R + jj, tl.int32), top_ids)
139 + active = tl.where(offs_a == K + jj, raw, active)
140 +
141 + # Log-sigmoid renormalization over the K + S active logits.
142 + abs_l = tl.abs(active)
143 + min_l = tl.minimum(active, 0.0)
144 + log_probs = min_l - tldevice.log1p(tldevice.exp(-abs_l))
145 + max_lp = tl.max(log_probs, axis=0)
146 + exp_shifted = tldevice.exp(log_probs - max_lp)
147 + sum_exp = tl.sum(exp_shifted, axis=0)
148 + weights = exp_shifted / sum_exp * scale
149 +
150 + tl.store(ids_ptr + pid * A + offs_a, top_ids)
151 + tl.store(weights_ptr + pid * A + offs_a, weights)
152 +
153 +
154 +def inkling_gate_select(
155 + logits: torch.Tensor, # [T, >=G] fp32 (rows may carry GEMM padding)
156 + n_gate_experts: int,
157 + n_routed_experts: int,
158 + topk: int,
159 + n_shared_experts: int,
160 + bias: torch.Tensor | None,
161 + route_scale: float,
162 + global_scale: torch.Tensor | None,
163 +) -> tuple[torch.Tensor, torch.Tensor]:
164 + """Sigmoid + bias + top-k + log-sigmoid renorm; returns (weights, ids)."""
165 + assert logits.dtype == torch.float32
166 + tokens = logits.shape[0]
167 + active = topk + n_shared_experts
168 + topk_ids = torch.empty((tokens, active), dtype=torch.int32, device=logits.device)
169 + topk_weights = torch.empty(
170 + (tokens, active), dtype=torch.float32, device=logits.device
171 + )
172 + if tokens == 0:
173 + return topk_weights, topk_ids
174 + _inkling_gate_select_kernel[(tokens,)](
175 + logits,
176 + bias if bias is not None else logits,
177 + global_scale if global_scale is not None else logits,
178 + topk_ids,
179 + topk_weights,
180 + route_scale,
181 + tokens,
182 + n_gate_experts,
183 + logits.stride(0),
184 + n_routed_experts,
185 + topk,
186 + n_shared_experts,
187 + HAS_BIAS=bias is not None,
188 + HAS_GSCALE=global_scale is not None,
189 + BLOCK_G=triton.next_power_of_2(n_gate_experts),
190 + )
191 + return topk_weights, topk_ids
192 +
193 +
194 +class InklingGate(nn.Module):
195 + """Sigmoid gate with selection bias, log-sigmoid renorm after top-k, and
196 + global scale (the served checkpoint's only configuration)."""
197 +
198 + def __init__(
199 + self,
200 + d_model: int,
201 + n_routed_experts: int,
202 + n_shared_experts: int,
203 + experts_per_token: int,
204 + route_scale: float,
205 + *,
206 + use_global_scale: bool = False,
207 + use_gate_bias: bool = False,
208 + ) -> None:
209 + super().__init__()
210 + self.n_routed_experts = n_routed_experts
211 + self.n_shared_experts = n_shared_experts
212 + self.n_total_experts = n_routed_experts + n_shared_experts
213 + self.topk = experts_per_token
214 + self.route_scale = route_scale
215 +
216 + padded_experts = self.n_total_experts + (-self.n_total_experts) % 8
217 + self.weight = Parameter(
218 + torch.empty(padded_experts, d_model), requires_grad=False
219 + )
220 + set_weight_attrs(self.weight, {"weight_loader": self._load_weight})
221 + if use_global_scale:
222 + self.global_scale = Parameter(
223 + torch.empty(1, dtype=torch.float32), requires_grad=False
224 + )
225 + else:
226 + self.global_scale = None
227 + if use_gate_bias:
228 + self.bias = Parameter(
229 + torch.empty(n_routed_experts, dtype=torch.float32),
230 + requires_grad=False,
231 + )
232 + else:
233 + self.bias = None
234 +
235 + @staticmethod
236 + def _load_weight(param: Parameter, loaded_weight: torch.Tensor) -> None:
237 + param.data.zero_()
238 + param.data[: loaded_weight.shape[0]].copy_(loaded_weight)
239 +
240 + def compute_logits(self, x: torch.Tensor) -> torch.Tensor:
241 + """fp32 gate logits [T, n_total_experts + pad] (pad columns are junk)."""
242 + return _linear_with_fp32_out(x, self.weight)
243 +
244 + def select_experts(
245 + self, gating_output: torch.Tensor
246 + ) -> tuple[torch.Tensor, torch.Tensor]:
247 + """Full selection: (weights, ids) of [T, K + S]. The first K entries
248 + are the routed top-k; the S trailing entries are the sink gammas."""
249 + return inkling_gate_select(
250 + gating_output,
251 + self.n_total_experts,
252 + self.n_routed_experts,
253 + self.topk,
254 + self.n_shared_experts,
255 + self.bias,
256 + self.route_scale,
257 + self.global_scale,
258 + )
259 +
260 +
261 +# ---------------------------------------------------------------------------
262 +# MoE layer
263 +# ---------------------------------------------------------------------------
264 +
265 +
266 +def _inkling_moe_ep_size() -> int:
267 + """EP size the FusedMoE layer will run with (mirrors
268 + FusedMoEParallelConfig.make: experts shard over tp * dp * pcp when
269 + expert parallelism is enabled)."""
270 + parallel_config = get_current_vllm_config().parallel_config
271 + if not parallel_config.enable_expert_parallel:
272 + return 1
273 + world = (
274 + get_tensor_model_parallel_world_size()
275 + * get_dp_group().world_size
276 + * get_pcp_group().world_size
277 + )
278 + return world if world > 1 else 1
279 +
280 +
281 +class InklingSinkExperts(nn.Module):
282 + """Shared "sink" experts with per-token gammas, in bf16.
283 +
284 + Replicated across EP ranks (every token activates every sink, so
285 + EP-sharding them would hotspot the owning rank) and TP-sharded on the
286 + intermediate dim so the output remains a TP-partial sum like the routed
287 + output. The sinks are always bf16 (the checkpoint excludes every
288 + ``shared_experts`` from quantization): the experts concatenate into two
289 + plain dense GEMMs with the fused sink epilogue between them.
290 + """
291 +
292 + def __init__(
293 + self, n_experts: int, d_model: int, d_mlp: int, *, prefix: str = ""
294 + ) -> None:
295 + super().__init__()
296 + self.n_experts = n_experts
297 + tp_size = get_tensor_model_parallel_world_size()
298 + self.tp_rank = get_tensor_model_parallel_rank()
299 + intermediate_pp = d_mlp // tp_size
300 + self.w13_weight = Parameter(
301 + torch.empty(n_experts, 2 * intermediate_pp, d_model),
302 + requires_grad=False,
303 + )
304 + self.w2_weight = Parameter(
305 + torch.empty(d_model, n_experts * intermediate_pp),
306 + requires_grad=False,
307 + )
308 + self._unit: torch.Tensor | None = None
309 +
310 + def load_weight(self, key: str, weight: torch.Tensor) -> list[str]:
311 + """Load one checkpoint sink tensor (stacked over the S experts)."""
312 + if key == "w13_weight":
313 + if weight.shape != self.w13_weight.shape:
314 + shard = self.w13_weight.shape[1]
315 + weight = weight.narrow(1, self.tp_rank * shard, shard)
316 + self.w13_weight.data.copy_(weight)
317 + return [key]
318 +
319 + assert key == "w2_weight"
320 + shard = self.w2_weight.shape[1] // self.n_experts
321 + shard_start = 0 if weight.shape[2] == shard else self.tp_rank * shard
322 + for expert_idx, expert_weight in enumerate(weight):
323 + local_weight = expert_weight.narrow(1, shard_start, shard)
324 + start = expert_idx * shard
325 + self.w2_weight.data[:, start : start + shard].copy_(local_weight)
326 + return [key]
327 +
328 + def forward(self, x: torch.Tensor, gammas: torch.Tensor) -> torch.Tensor:
329 + """``sum_e gammas[:, e] * MLP_e(x)`` (TP-partial along d_mlp)."""
330 + from .ops import sink_silu_mul_epilogue
331 +
332 + # One GEMM over the experts' stacked w13 (a view), fused epilogue,
333 + # then one GEMM whose K-reduction over the K-concatenated w2 performs
334 + # the expert sum.
335 + if self._unit is None or self._unit.device != x.device:
336 + self._unit = torch.ones(
337 + self.n_experts, dtype=torch.float32, device=x.device
338 + )
339 + raw = x @ self.w13_weight.view(-1, x.shape[-1]).T # (T, S*2F)
340 + h = sink_silu_mul_epilogue(
341 + raw, self._unit, gammas, self._unit, self.n_experts, x.dtype
342 + )
343 + return h @ self.w2_weight.T # (T, D)
344 +
345 +
346 +class InklingMoE(nn.Module):
347 + def __init__(
348 + self,
349 + config: InklingModelConfig,
350 + layer_id: int,
351 + *,
352 + prefix: str = "",
353 + nvfp4_config: InklingNvfp4Config | None = None,
354 + ) -> None:
355 + super().__init__()
356 + # Overfit to the served checkpoint: sigmoid gate renormalized after
357 + # top-k, shared sink experts, interleaved gate/up checkpoint rows.
358 + assert config.gate_activation == "sigmoid" and config.norm_after_topk
359 + assert config.n_shared_experts > 0 and config.shared_expert_sink
360 + assert config.inference_moe_w13_interleaved
361 + n_routed = config.n_routed_experts
362 + n_shared = config.n_shared_experts
363 + self.n_routed_experts = n_routed
364 + self.gate = InklingGate(
365 + d_model=config.hidden_size,
366 + n_routed_experts=n_routed,
367 + n_shared_experts=n_shared,
368 + experts_per_token=config.num_experts_per_tok,
369 + route_scale=config.route_scale,
370 + use_global_scale=config.use_global_scale,
371 + use_gate_bias=config.use_gate_bias,
372 + )
373 +
374 + moe_quant_config = None
375 + if nvfp4_config is not None and nvfp4_config.experts_quantized(layer_id):
376 + from vllm.model_executor.layers.quantization.modelopt import (
377 + ModelOptNvFp4Config,
378 + )
379 +
380 + # The Inkling checkpoint is ModelOpt NVFP4; exclusion is decided per
381 + # layer right here, so no exclude list is needed.
382 + moe_quant_config = ModelOptNvFp4Config(
383 + quant_method="NVFP4",
384 + is_checkpoint_nvfp4_serialized=True,
385 + kv_cache_quant_algo=None,
386 + exclude_modules=[],
387 + group_size=nvfp4_config.group_size,
388 + )
389 +
390 + # TRTLLM MoE kernels assume equal, contiguous per-rank expert slabs
391 + # (local_expert_offset = ep_rank * local_num_experts), so pad the
392 + # expert count to a multiple of the EP size. A no-op for the usual
393 + # power-of-two EP sizes (n_routed is a power of two).
394 + num_experts = n_routed + (-n_routed) % _inkling_moe_ep_size()
395 +
396 + self.experts = FusedMoE(
397 + num_experts=num_experts,
398 + top_k=config.num_experts_per_tok,
399 + hidden_size=config.hidden_size,
400 + intermediate_size=config.intermediate_size,
401 + renormalize=False,
402 + quant_config=moe_quant_config,
403 + prefix=f"{prefix}.experts",
404 + custom_routing_function=self._select_routed,
405 + router_logits_dtype=torch.float32,
406 + activation="silu",
407 + )
408 + # The decoder layer reduce-scatters the MoE delta into the sconv
409 + # stream itself (RS -> shard sconv -> AG); the runner must return the
410 + # per-rank partial sum instead of all-reducing.
411 + self.experts.moe_config.skip_final_all_reduce = True
412 +
413 + self._routed_sel: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None
414 + # The sinks are always bf16; fail loudly on a checkpoint that
415 + # quantizes them instead of silently misloading.
416 + assert nvfp4_config is None or not nvfp4_config.shared_experts_quantized(
417 + layer_id
418 + ), f"layer {layer_id}: NVFP4 shared experts are not supported"
419 +
420 + self.sink_experts = InklingSinkExperts(
421 + n_experts=n_shared,
422 + d_model=config.hidden_size,
423 + d_mlp=config.intermediate_size,
424 + prefix=f"{prefix}.shared_experts",
425 + )
426 +
427 + # Sink chain overlaps the routed MoE call on the aux stream for
428 + # decode-sized batches (same pattern as the runner's SharedExperts
429 + # multi-stream overlap). The routed GEMM runs on the default stream and
430 + # the sink chain on the aux stream, joined via these two events by
431 + # ``maybe_execute_in_parallel``.
432 + self._sink_stream: torch.cuda.Stream | None = aux_stream()
433 + self._sink_events = (torch.cuda.Event(), torch.cuda.Event())
434 +
435 + def _select_routed(
436 + self,
437 + hidden_states: torch.Tensor,
438 + gating_output: torch.Tensor,
439 + topk: int,
440 + renormalize: bool,
441 + ) -> tuple[torch.Tensor, torch.Tensor]:
442 + """FusedMoE ``custom_routing_function``: the routed top-k slice of the
443 + full (routed + sink) selection.
444 +
445 + forward() stashes its selection (keyed by logits identity) so the
446 + gate select runs once per layer; the fallback covers paths where the
447 + runner re-derives the logits (e.g. naive DP dispatch).
448 + """
449 + del hidden_states, renormalize
450 + assert topk == self.gate.topk
451 + cached = self._routed_sel
452 + self._routed_sel = None
453 + if cached is not None and cached[0] is gating_output:
454 + return cached[1], cached[2]
455 + weights, ids = self.gate.select_experts(gating_output)
456 + return weights[:, :topk].contiguous(), ids[:, :topk].contiguous()
457 +
458 + def forward(self, x: torch.Tensor) -> torch.Tensor | None:
459 + router_logits = self.gate.compute_logits(x)
460 + num_tokens = x.shape[0]
461 + # One gate select per layer: the routed slice is stashed for the
462 + # routing function inside the FusedMoE op; the sink gammas are the
463 + # trailing columns.
464 + k = self.gate.topk
465 + weights, ids = self.gate.select_experts(router_logits)
466 + self._routed_sel = (
467 + router_logits,
468 + weights[:, :k].contiguous(),
469 + ids[:, :k].contiguous(),
470 + )
471 + gammas = weights[:, k:]
472 +
473 + out, sink_out = maybe_execute_in_parallel(
474 + lambda: self.experts(hidden_states=x, router_logits=router_logits),
475 + lambda: self.sink_experts(x, gammas),
476 + self._sink_events[0],
477 + self._sink_events[1],
478 + self._sink_stream
479 + if num_tokens <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
480 + else None,
481 + )
482 + self._routed_sel = None
483 +
484 + return out + sink_out
485 +
486 + # -- weight loading ----------------------------------------------------
487 +
488 + def _local_expert_slots(self) -> dict[int, int]:
489 + """Global expert id -> local slot for this rank's expert partition."""
490 + manager = self.experts.routed_experts.expert_map_manager
491 + if manager.expert_map is None:
492 + return {g: g for g in range(manager.global_num_experts)}
493 + emap = manager.expert_map.tolist()
494 + return {g: slot for g, slot in enumerate(emap) if slot >= 0}
495 +
496 + def load_expert_weight(self, name: str, weight: torch.Tensor) -> list[str]:
497 + """Load one checkpoint expert tensor.
498 +
499 + ``name`` is relative to the mlp module: ``experts.<t>`` (routed
500 + stack) or ``shared_experts.shared_<t>`` (sink experts). Returns the
501 + loaded param names (relative to this module).
502 + """
503 + if name.startswith("shared_experts."):
504 + key = name.split(".", 1)[1].replace("shared_", "", 1)
505 + return [
506 + f"sink_experts.{p}" for p in self.sink_experts.load_weight(key, weight)
507 + ]
508 +
509 + experts: RoutedExperts = self.experts.routed_experts
510 + key = name.split(".", 1)[1]
511 +
512 + # original_shape is unused by the vLLM serving layout.
513 + if key.endswith(".original_shape"):
514 + return []
515 + if key.endswith(".input_amax"):
516 + projection = "w13" if key.startswith("w13") else "w2"
517 + amax = float(weight.max())
518 + assert math.isfinite(amax) and amax > 0, (
519 + f"bad {projection} input_amax: {amax}"
520 + )
521 + input_scale = getattr(experts, f"{projection}_input_scale")
522 + input_scale.data.fill_(amax / (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX))
523 + return [f"experts.routed_experts.{projection}_input_scale"]
524 +
525 + param = getattr(experts, key)
526 + slots = self._local_expert_slots()
527 + gids = sorted(slots)
528 + lids = [slots[g] for g in gids]
529 + tp_rank = experts.moe_config.moe_parallel_config.tp_rank
530 +
531 + if key.endswith("_scale_2"):
532 + # Per-expert scalars, vectorized over the local experts. The
533 + # fused w13 param carries one slot per gate/up half.
534 + vals = weight[gids].float().to(param.device)
535 + param.data[lids] = vals[:, None] if param.data.ndim == 2 else vals
536 + elif key.startswith("w13"):
537 + # Checkpoint w13 rows are interleaved [g0, u0, g1, u1, ...]; the
538 + # fused param layout is [w1(gate); w3(up)]. The TP-local rows form
539 + # one contiguous slab of the interleaved tensor, so upload just
540 + # that slab (a single bounded synchronous H2D; pre-uploading whole
541 + # untrimmed tensors pins the mmap pages of the entire checkpoint
542 + # and OOMs the host) and de-interleave on device.
543 + half = param.shape[1] // 2
544 + for gid, lid in slots.items():
545 + slab = weight[gid].narrow(0, tp_rank * 2 * half, 2 * half)
546 + slab = slab.to(param.device)
547 + param.data[lid, :half].copy_(slab[0::2])
548 + param.data[lid, half:].copy_(slab[1::2])
549 + else:
550 + # w2: shard the packed intermediate (last) dim.
551 + shard = param.shape[2]
552 + for gid, lid in slots.items():
553 + param.data[lid].copy_(weight[gid].narrow(1, tp_rank * shard, shard))
554 + return [f"experts.routed_experts.{key}"]
555 +
556 + def finalize_load(self) -> list[str]:
557 + """Post-load fixups for zeroed padding experts."""
558 + experts = self.experts.routed_experts
559 + out: list[str] = []
560 + # Zero the EP-alignment padding experts (if any) so their
561 + # (never-routed) slots hold defined values.
562 + slots = self._local_expert_slots()
563 + for gid in range(self.n_routed_experts, experts.global_num_experts):
564 + lid = slots.get(gid)
565 + if lid is None:
566 + continue
567 + for pname in (
568 + "w13_weight",
569 + "w2_weight",
570 + "w13_weight_scale",
571 + "w2_weight_scale",
572 + "w13_weight_scale_2",
573 + "w2_weight_scale_2",
574 + ):
575 + p = getattr(experts, pname, None)
576 + if p is not None:
577 + p.data[lid].zero_()
578 + return out
vllm/models/inkling/nvidia/ops/__init__.py new
+43
@@ -0,0 +1,43 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inkling kernels (NVIDIA).
4 +
5 +``rmsnorm`` / ``sconv`` import eagerly. The SwiGLU kernels and the FA4
6 +relative-attention wrapper are exposed lazily to keep this package's
7 +import path lightweight.
8 +"""
9 +
10 +from typing import TYPE_CHECKING
11 +
12 +from .norm import add_rmsnorm, rmsnorm
13 +from .sconv import fused_sconv
14 +
15 +_LAZY_EXPORTS = {
16 + "silu_and_mul_triton": "silu_and_mul",
17 + "sink_silu_mul_epilogue": "silu_and_mul",
18 + "inkling_fa4_rel_attention": "fa4_rel_attention",
19 +}
20 +
21 +if TYPE_CHECKING:
22 + from .fa4_rel_attention import inkling_fa4_rel_attention # noqa: F401
23 + from .silu_and_mul import ( # noqa: F401
24 + silu_and_mul_triton,
25 + sink_silu_mul_epilogue,
26 + )
27 +
28 +__all__ = [
29 + "add_rmsnorm",
30 + "rmsnorm",
31 + "fused_sconv",
32 + *sorted(_LAZY_EXPORTS),
33 +]
34 +
35 +
36 +def __getattr__(name: str):
37 + module = _LAZY_EXPORTS.get(name)
38 + if module is not None:
39 + import importlib
40 +
41 + mod = importlib.import_module(f".{module}", __name__)
42 + return getattr(mod, name)
43 + raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
vllm/models/inkling/nvidia/ops/fa4_rel_attention.py new
+106
@@ -0,0 +1,106 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from __future__ import annotations
4 +
5 +import torch
6 +
7 +from vllm.platforms import current_platform
8 +
9 +
10 +def bucket_max_seqlen_q(max_seqlen_q: int) -> int:
11 + """Round the FA4 scheduling bound up to a power of two."""
12 + return 1 << max(0, max_seqlen_q - 1).bit_length()
13 +
14 +
15 +def inkling_fa4_num_splits(
16 + *,
17 + is_local: bool,
18 + batch_size: int,
19 + max_query_len: int,
20 + num_heads: int,
21 + num_kv_heads: int,
22 + max_kv_len: int,
23 +) -> int:
24 + """Return the split-KV cap for FA4 with sheared relative bias."""
25 + capability = current_platform.get_device_capability()
26 + if capability is not None and capability.major == 9:
27 + return 1
28 + if is_local:
29 + return 1
30 +
31 + q_rows = max_query_len * (num_heads // num_kv_heads)
32 + q_tiles = (q_rows + 255) // 256
33 + base_ctas = batch_size * num_kv_heads * q_tiles
34 + # Shearing makes split/combine overhead more visible. Multi-tile causal
35 + # prefill saturates around 64 CTAs. Batch-1 decode at very long context is
36 + # memory-bound and uses a TP-specific cap measured through 1M KV tokens.
37 + target_ctas = (
38 + 256 if q_tiles == 1 and batch_size == 1 else (128 if q_tiles == 1 else 64)
39 + )
40 + max_splits = 128
41 + if q_tiles == 1 and batch_size == 1:
42 + if num_kv_heads == 8:
43 + max_splits = 16
44 + elif num_kv_heads == 4 or max_kv_len <= 8192:
45 + max_splits = 32
46 + elif max_kv_len <= 65536:
47 + max_splits = 64
48 + else:
49 + max_splits = 128
50 + return max(
51 + 1,
52 + min(target_ctas // base_ctas, max_splits, (max_kv_len + 127) // 128),
53 + )
54 +
55 +
56 +def inkling_fa4_rel_attention(
57 + q: torch.Tensor,
58 + key_cache: torch.Tensor,
59 + value_cache: torch.Tensor,
60 + *,
61 + block_table: torch.Tensor,
62 + cache_seqlens: torch.Tensor,
63 + cu_seqlens_q: torch.Tensor,
64 + max_seqlen_q: int,
65 + softmax_scale: float,
66 + causal: bool,
67 + window_size: tuple[int, int],
68 + rel_extent: int,
69 + rel_logits: torch.Tensor,
70 + num_splits: int = 32,
71 + out: torch.Tensor | None = None,
72 +) -> torch.Tensor:
73 + """Paged varlen FA4 over the bound K/V cache with the Inkling relative bias.
74 +
75 + ``q`` is ``(num_tokens, num_heads, head_dim)``; ``key_cache`` / ``value_cache``
76 + are the paged caches ``(num_blocks, block_size, num_kv_heads, head_dim)``;
77 + ``block_table`` is the per-request page table and ``cache_seqlens`` the
78 + per-request KV lengths (``seqused_k``). ``rel_logits`` is
79 + ``(num_tokens, num_heads, rel_extent)``.
80 +
81 + The bias uses tml-fa4's sheared relative-bias layout.
82 + """
83 + from vllm.third_party.tml_fa4 import flash_attn_varlen_func
84 +
85 + # cute uses (None, None) to mean "no window".
86 + cute_window = (None, None) if window_size == (-1, -1) else window_size
87 +
88 + ret = flash_attn_varlen_func(
89 + q=q,
90 + k=key_cache,
91 + v=value_cache,
92 + cu_seqlens_q=cu_seqlens_q,
93 + seqused_k=cache_seqlens,
94 + max_seqlen_q=max_seqlen_q,
95 + page_table=block_table,
96 + softmax_scale=softmax_scale,
97 + causal=causal,
98 + window_size=cute_window,
99 + num_splits=num_splits,
100 + return_lse=False,
101 + out=out,
102 + rel_bias=rel_logits.contiguous(),
103 + )
104 + if isinstance(ret, tuple):
105 + return ret[0]
106 + return ret
vllm/models/inkling/nvidia/ops/fa4_warmup.py new
+154
@@ -0,0 +1,154 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Startup compilation of the FA4 kernels used by Inkling."""
4 +
5 +from __future__ import annotations
6 +
7 +from collections.abc import Iterator
8 +from dataclasses import dataclass
9 +from functools import partial
10 +
11 +import torch
12 +
13 +from vllm.model_executor.warmup.cutedsl_warmup import (
14 + CuTeDSLCompileUnit,
15 + register_cutedsl_warmup_provider,
16 +)
17 +
18 +from .fa4_rel_attention import (
19 + bucket_max_seqlen_q,
20 + inkling_fa4_num_splits,
21 + inkling_fa4_rel_attention,
22 +)
23 +
24 +
25 +@dataclass(frozen=True)
26 +class InklingFA4WarmupConfig:
27 + num_heads: int
28 + num_kv_heads: int
29 + head_dim: int
30 + rel_extent: int
31 + window_size: tuple[int, int]
32 + is_local: bool
33 + max_kv_len: int
34 + dtype: torch.dtype
35 + kv_dtype: torch.dtype
36 + block_size: int
37 + max_num_reqs: int
38 + max_num_batched_tokens: int
39 +
40 +
41 +def _num_warps_bucket(num_reqs: int) -> int:
42 + num_warps = min((num_reqs + 30) // 31, 32)
43 + return 1 << (num_warps - 1).bit_length()
44 +
45 +
46 +def _compile(config: InklingFA4WarmupConfig, max_seqlen_q: int, num_reqs: int) -> None:
47 + from torch._subclasses.fake_tensor import FakeTensorMode
48 +
49 + num_splits = inkling_fa4_num_splits(
50 + is_local=config.is_local,
51 + batch_size=num_reqs,
52 + max_query_len=max_seqlen_q,
53 + num_heads=config.num_heads,
54 + num_kv_heads=config.num_kv_heads,
55 + max_kv_len=config.max_kv_len,
56 + )
57 + with FakeTensorMode():
58 + device = torch.accelerator.current_accelerator()
59 + total_q = max_seqlen_q + num_reqs - 1
60 + q = torch.empty(
61 + total_q,
62 + config.num_heads,
63 + config.head_dim,
64 + dtype=config.dtype,
65 + device=device,
66 + )
67 + kv = torch.empty(
68 + 1,
69 + 2,
70 + config.block_size,
71 + config.num_kv_heads,
72 + config.head_dim,
73 + dtype=config.kv_dtype,
74 + device=device,
75 + )
76 + key_cache, value_cache = kv.unbind(1)
77 + inkling_fa4_rel_attention(
78 + q,
79 + key_cache,
80 + value_cache,
81 + block_table=torch.empty(num_reqs, 1, dtype=torch.int32, device=device),
82 + cache_seqlens=torch.empty(num_reqs, dtype=torch.int32, device=device),
83 + cu_seqlens_q=torch.empty(num_reqs + 1, dtype=torch.int32, device=device),
84 + max_seqlen_q=max_seqlen_q,
85 + softmax_scale=1.0 / config.head_dim,
86 + causal=True,
87 + window_size=config.window_size,
88 + rel_extent=config.rel_extent,
89 + rel_logits=torch.empty(
90 + total_q,
91 + config.num_heads,
92 + config.rel_extent,
93 + dtype=config.dtype,
94 + device=device,
95 + ),
96 + num_splits=num_splits,
97 + out=torch.empty_like(q),
98 + )
99 +
100 +
101 +def _iter_compile_units(
102 + config: InklingFA4WarmupConfig,
103 +) -> Iterator[CuTeDSLCompileUnit]:
104 + max_bucket = bucket_max_seqlen_q(config.max_num_batched_tokens)
105 + max_seqlen_q = 1
106 + while max_seqlen_q <= max_bucket:
107 + min_query_len = max_seqlen_q // 2 + 1
108 + max_num_reqs = min(
109 + config.max_num_reqs,
110 + config.max_num_batched_tokens - min_query_len + 1,
111 + )
112 + seen: set[tuple[int, int, int | None, bool]] = set()
113 + for num_reqs in range(1, max_num_reqs + 1):
114 + num_splits = inkling_fa4_num_splits(
115 + is_local=config.is_local,
116 + batch_size=num_reqs,
117 + max_query_len=max_seqlen_q,
118 + num_heads=config.num_heads,
119 + num_kv_heads=config.num_kv_heads,
120 + max_kv_len=config.max_kv_len,
121 + )
122 + key = (
123 + max_seqlen_q,
124 + num_splits,
125 + _num_warps_bucket(num_reqs) if num_splits > 1 else None,
126 + num_reqs > 1024,
127 + )
128 + if key in seen:
129 + continue
130 + seen.add(key)
131 + yield CuTeDSLCompileUnit(
132 + name="inkling_fa4",
133 + key=("inkling_fa4", config, key),
134 + compile=partial(_compile, config, max_seqlen_q, num_reqs),
135 + )
136 + max_seqlen_q *= 2
137 +
138 +
139 +class _WarmupProvider:
140 + def __init__(self) -> None:
141 + self.configs: set[InklingFA4WarmupConfig] = set()
142 +
143 + def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]:
144 + return tuple(
145 + unit for config in self.configs for unit in _iter_compile_units(config)
146 + )
147 +
148 +
149 +_PROVIDER = _WarmupProvider()
150 +
151 +
152 +def register_fa4_warmup(config: InklingFA4WarmupConfig) -> None:
153 + _PROVIDER.configs.add(config)
154 + register_cutedsl_warmup_provider(_PROVIDER)
vllm/models/inkling/nvidia/ops/lamport.py new
+766
@@ -0,0 +1,766 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Deadlock-free fused RS + short-conv + AG + residual + RMSNorm.
4 +
5 +The public integration surface is ``LamportRSConv.rs_sconv_ag_add_norm``.
6 +
7 +Liveness
8 +--------
9 +Large grids are deliberately split at the two communication dependencies:
10 +
11 + 1. ``_publish_input_kernel`` only publishes rank partials.
12 + 2. ``_reduce_insert_kernel`` waits, reduces, and inserts the local shard.
13 + 3. ``_sconv_publish_kernel`` only computes and publishes the local result.
14 + 4. ``_gather_norm_kernel`` waits, gathers, and normalizes.
15 +
16 +Without PDL, CUDA stream order completes a producer before its consumer. With
17 +PDL, every producer CTA posts all peer stores before triggering its dependent,
18 +and the consumer executes ``gdc_wait`` before polling. A consumer therefore
19 +waits only for stores from a producer that is already running (or complete) on
20 +another GPU. It never waits for another block in its own grid. Consequently
21 +spinning consumers cannot occupy resources needed by any producer, and the
22 +wait-for graph has no cycle. The proof is independent of grid size, block
23 +dispatch order, and occupancy.
24 +
25 +For one token, the first three phases use eight independent channel slices to
26 +expose enough CTA parallelism for decode latency. Each slice has exclusive
27 +ownership of its cache and Lamport columns; the gather/RMSNorm phase retains
28 +one CTA per token so no cross-CTA reduction or completion counter is needed.
29 +
30 +Immediate buffer reuse (including replay of a captured CUDA graph) is also
31 +safe. Rank R cannot republish input for call n+1 until its gather for n has
32 +finished; that gather waited for owner O's output, which O publishes only
33 +after consuming R's input for n. Likewise, R cannot republish output for n+1
34 +until its reduction for n+1 has observed destination D's input; D publishes
35 +that input only after its gather consumed R's output for n. Thus every prior
36 +read happens-before a same-slot rewrite. Three generations reduce incidental
37 +coupling for ordinary launches, but correctness does not depend on rotation.
38 +
39 +The payload itself is the Lamport flag. A 32-bit store publishes two bf16s
40 +atomically; 0x80008000 (two negative zeroes) denotes an empty pair. Real
41 +negative zeroes are changed to positive zero before publication. Consumers
42 +use volatile 32-bit loads and restore the sentinel after consuming a slot.
43 +"""
44 +
45 +from __future__ import annotations
46 +
47 +import os
48 +
49 +import torch
50 +
51 +from vllm.distributed import get_tp_group
52 +from vllm.distributed.parallel_state import in_the_same_node_as
53 +from vllm.logger import init_logger
54 +from vllm.triton_utils import HAS_TRITON, tl, triton
55 +
56 +logger = init_logger(__name__)
57 +
58 +_MAX_TOKENS = 16384
59 +_EMPTY_PAIR = tl.constexpr(0x80008000) if HAS_TRITON else 0x80008000
60 +
61 +
62 +@triton.jit
63 +def _pack_bf16_pairs(values):
64 + """Pack bf16 values into atomic u32 pairs and reserve negative zero."""
65 + lo, hi = tl.split(values.reshape([values.shape[0] // 2, 2]))
66 + lo = lo.to(tl.uint16, bitcast=True)
67 + hi = hi.to(tl.uint16, bitcast=True)
68 + lo = tl.where(lo == 0x8000, 0, lo).to(tl.uint32)
69 + hi = tl.where(hi == 0x8000, 0, hi).to(tl.uint32)
70 + return lo | (hi << 16)
71 +
72 +
73 +@triton.jit
74 +def _unpack_bf16_pairs(values):
75 + lo = (values & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True)
76 + hi = (values >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True)
77 + return tl.interleave(lo, hi)
78 +
79 +
80 +@triton.jit
81 +def _wait_pairs(ptr, offsets, mask):
82 + values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True)
83 + while tl.max(tl.where(mask & (values == _EMPTY_PAIR), 1, 0)) != 0:
84 + values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True)
85 + return values
86 +
87 +
88 +@triton.jit
89 +def _publish_input_kernel(
90 + stage_ptr,
91 + peer_ptrs,
92 + peer_offset_u32,
93 + stride_stage_t,
94 + C: tl.constexpr,
95 + CS: tl.constexpr,
96 + CS_P2: tl.constexpr,
97 + SPLITS: tl.constexpr,
98 + RANK: tl.constexpr,
99 + WORLD: tl.constexpr,
100 + USE_PDL: tl.constexpr,
101 + launch_pdl: tl.constexpr,
102 +):
103 + """Publish this rank's full partial row into every shard owner's slots."""
104 + token = tl.program_id(0).to(tl.int64)
105 + split = tl.program_id(1).to(tl.int64)
106 + CSS: tl.constexpr = CS // SPLITS
107 + pair = tl.arange(0, CS_P2 // 2)
108 + pair_mask = pair < CSS // 2
109 + elem = tl.arange(0, CS_P2)
110 + elem_mask = elem < CSS
111 + ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64))
112 + # Address generation is independent of the preceding stream kernel. The
113 + # acquire remains before stage/peer loads, which may consume its writes.
114 + if USE_PDL:
115 + tl.extra.cuda.gdc_wait()
116 +
117 + for owner in tl.static_range(WORLD):
118 + values = tl.load(
119 + stage_ptr + token * stride_stage_t + owner * CS + split * CSS + elem,
120 + mask=elem_mask,
121 + other=0.0,
122 + )
123 + packed = _pack_bf16_pairs(values)
124 + base = tl.load(ptrs + owner).to(tl.pointer_type(tl.uint32))
125 + dst = (token * WORLD + RANK) * (CS // 2) + split * (CSS // 2) + pair
126 + tl.store(base + peer_offset_u32 + dst, packed, mask=pair_mask)
127 + if USE_PDL:
128 + tl.extra.cuda.gdc_launch_dependents()
129 +
130 +
131 +@triton.jit
132 +def _reduce_insert_kernel(
133 + input_peer_ptrs,
134 + input_peer_offset_u32,
135 + cache_ptr,
136 + slot_ptr,
137 + stride_cache_block,
138 + stride_cache_head,
139 + stride_cache_token,
140 + stride_cache_dim,
141 + block_size,
142 + C: tl.constexpr,
143 + CS: tl.constexpr,
144 + CS_P2: tl.constexpr,
145 + SPLITS: tl.constexpr,
146 + HEAD_SIZE: tl.constexpr,
147 + CACHE_OFFSET: tl.constexpr,
148 + RANK: tl.constexpr,
149 + WORLD: tl.constexpr,
150 + USE_PDL: tl.constexpr,
151 + launch_pdl: tl.constexpr,
152 +):
153 + """Consume all partials for this rank and insert the reduced cache row."""
154 + token = tl.program_id(0).to(tl.int64)
155 + split = tl.program_id(1).to(tl.int64)
156 + CSS: tl.constexpr = CS // SPLITS
157 + source = tl.arange(0, WORLD)
158 + pair = tl.arange(0, CS_P2 // 2)
159 + pair_mask = pair < CSS // 2
160 + offsets = (
161 + (token * WORLD + source)[:, None] * (CS // 2)
162 + + split * (CSS // 2)
163 + + pair[None, :]
164 + )
165 + mask = tl.full([WORLD], True, tl.int1)[:, None] & pair_mask[None, :]
166 + input_ptrs = input_peer_ptrs.to(tl.pointer_type(tl.uint64))
167 + input_u32 = tl.load(input_ptrs + RANK).to(tl.pointer_type(tl.uint32))
168 + input_u32 += input_peer_offset_u32
169 + # Slot metadata and the cache destination do not depend on input publish.
170 + slot = tl.load(slot_ptr + token)
171 + valid = slot >= 0
172 + channel = tl.arange(0, CS_P2)
173 + channel_mask = channel < CSS
174 + global_channel = split * CSS + channel
175 + head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1)
176 + dim = CACHE_OFFSET + global_channel % HEAD_SIZE
177 + safe_slot = tl.maximum(slot, 0).to(tl.int64)
178 + dst = (
179 + cache_ptr
180 + + (safe_slot // block_size) * stride_cache_block
181 + + head * stride_cache_head
182 + + (safe_slot % block_size) * stride_cache_token
183 + + dim * stride_cache_dim
184 + )
185 + if USE_PDL:
186 + tl.extra.cuda.gdc_wait()
187 + packed = _wait_pairs(input_u32, offsets, mask)
188 +
189 + lo = (packed & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True)
190 + hi = (packed >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True)
191 + reduced = tl.interleave(
192 + tl.sum(lo.to(tl.float32), axis=0).to(tl.bfloat16),
193 + tl.sum(hi.to(tl.float32), axis=0).to(tl.bfloat16),
194 + )
195 + tl.store(
196 + input_u32 + offsets,
197 + tl.full([WORLD, CS_P2 // 2], _EMPTY_PAIR, tl.uint32),
198 + mask=mask,
199 + )
200 +
201 + tl.store(dst, reduced, mask=valid & channel_mask)
202 + if USE_PDL:
203 + tl.extra.cuda.gdc_launch_dependents()
204 +
205 +
206 +@triton.jit
207 +def _sconv_publish_kernel(
208 + peer_ptrs,
209 + peer_offset_u32,
210 + residual_ptr,
211 + weight_ptr,
212 + cache_ptr,
213 + position_ptr,
214 + sequence_ptr,
215 + slot_ptr,
216 + block_table_ptr,
217 + stride_residual_t,
218 + stride_cache_block,
219 + stride_cache_head,
220 + stride_cache_token,
221 + stride_cache_dim,
222 + stride_block_table_r,
223 + max_blocks,
224 + block_size,
225 + C: tl.constexpr,
226 + CS: tl.constexpr,
227 + CS_P2: tl.constexpr,
228 + SPLITS: tl.constexpr,
229 + HEAD_SIZE: tl.constexpr,
230 + CACHE_OFFSET: tl.constexpr,
231 + RANK: tl.constexpr,
232 + WORLD: tl.constexpr,
233 + WINDOW: tl.constexpr,
234 + USE_PDL: tl.constexpr,
235 + launch_pdl: tl.constexpr,
236 +):
237 + """Compute this rank's shard, then publish it to every rank."""
238 + tl.static_assert(WINDOW == 4)
239 + token = tl.program_id(0).to(tl.int64)
240 + split = tl.program_id(1).to(tl.int64)
241 + CSS: tl.constexpr = CS // SPLITS
242 + channel = tl.arange(0, CS_P2)
243 + channel_mask = channel < CSS
244 + global_channel = split * CSS + channel
245 + head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1)
246 + dim = CACHE_OFFSET + global_channel % HEAD_SIZE
247 + slot = tl.load(slot_ptr + token)
248 + valid = slot >= 0
249 + position = tl.load(position_ptr + token)
250 + sequence = tl.load(sequence_ptr + token)
251 + safe_slot = tl.maximum(slot, 0).to(tl.int64)
252 + own_ptr = (
253 + cache_ptr
254 + + (safe_slot // block_size) * stride_cache_block
255 + + head * stride_cache_head
256 + + (safe_slot % block_size) * stride_cache_token
257 + + dim * stride_cache_dim
258 + )
259 + # These loads were made visible before reduce/insert could signal us and
260 + # are independent of its cache store. Hoisting them gives PDL useful work
261 + # to overlap while retaining the acquire before every cache read.
262 + residual = tl.load(
263 + residual_ptr + token * stride_residual_t + RANK * CS + global_channel,
264 + mask=channel_mask,
265 + other=0.0,
266 + )
267 + weight0 = tl.load(
268 + weight_ptr + global_channel * WINDOW,
269 + mask=channel_mask,
270 + other=0.0,
271 + ).to(tl.float32)
272 + weight1 = tl.load(
273 + weight_ptr + global_channel * WINDOW + 1,
274 + mask=channel_mask,
275 + other=0.0,
276 + ).to(tl.float32)
277 + weight2 = tl.load(
278 + weight_ptr + global_channel * WINDOW + 2,
279 + mask=channel_mask,
280 + other=0.0,
281 + ).to(tl.float32)
282 + weight3 = tl.load(
283 + weight_ptr + global_channel * WINDOW + 3,
284 + mask=channel_mask,
285 + other=0.0,
286 + ).to(tl.float32)
287 + ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64))
288 + if USE_PDL:
289 + tl.extra.cuda.gdc_wait()
290 + current = tl.load(own_ptr, mask=valid & channel_mask, other=0.0)
291 +
292 + conv = tl.zeros([CS_P2], tl.float32)
293 + for tap_idx in tl.static_range(WINDOW):
294 + source_position = position - (WINDOW - 1) + tap_idx
295 + take = valid & (source_position >= 0)
296 + if tap_idx == WINDOW - 1:
297 + value = tl.where(take, current.to(tl.float32), 0.0)
298 + else:
299 + safe_position = tl.maximum(source_position, 0)
300 + logical_block = tl.minimum(safe_position // block_size, max_blocks - 1)
301 + physical_block = tl.load(
302 + block_table_ptr + sequence * stride_block_table_r + logical_block,
303 + mask=take,
304 + other=0,
305 + ).to(tl.int64)
306 + source_ptr = (
307 + cache_ptr
308 + + physical_block * stride_cache_block
309 + + head * stride_cache_head
310 + + (safe_position % block_size) * stride_cache_token
311 + + dim * stride_cache_dim
312 + )
313 + cached = tl.load(source_ptr, mask=take & channel_mask, other=0.0)
314 + value = tl.where(take, cached.to(tl.float32), 0.0)
315 + if tap_idx == 0:
316 + weight = weight0
317 + elif tap_idx == 1:
318 + weight = weight1
319 + elif tap_idx == 2:
320 + weight = weight2
321 + else:
322 + weight = weight3
323 + conv += value * weight
324 +
325 + # Preserve both bf16 rounding points of the original sublayer.
326 + short_conv_with_skip = (conv + current.to(tl.float32)).to(tl.bfloat16)
327 + output = (residual.to(tl.float32) + short_conv_with_skip.to(tl.float32)).to(
328 + tl.bfloat16
329 + )
330 +
331 + pair = tl.arange(0, CS_P2 // 2)
332 + pair_mask = pair < CSS // 2
333 + packed = _pack_bf16_pairs(output)
334 + row_offset = token * (C // 2) + RANK * (CS // 2) + split * (CSS // 2) + pair
335 +
336 + for destination in tl.static_range(WORLD):
337 + base = tl.load(ptrs + destination).to(tl.pointer_type(tl.uint32))
338 + tl.store(base + peer_offset_u32 + row_offset, packed, mask=pair_mask)
339 + if USE_PDL:
340 + tl.extra.cuda.gdc_launch_dependents()
341 +
342 +
343 +@triton.jit
344 +def _gather_norm_kernel(
345 + output_peer_ptrs,
346 + output_peer_offset_u32,
347 + norm_weight_ptr,
348 + normed_ptr,
349 + residual_out_ptr,
350 + eps,
351 + stride_output_t,
352 + C: tl.constexpr,
353 + C_P2: tl.constexpr,
354 + RANK: tl.constexpr,
355 + HAS_NORM: tl.constexpr,
356 + USE_PDL: tl.constexpr,
357 + launch_pdl: tl.constexpr,
358 +):
359 + """Consume a complete row; one CTA owns all outputs for one token."""
360 + token = tl.program_id(0).to(tl.int64)
361 + pair = tl.arange(0, C_P2 // 2)
362 + pair_mask = pair < C // 2
363 + output_ptrs = output_peer_ptrs.to(tl.pointer_type(tl.uint64))
364 + output_u32 = tl.load(output_ptrs + RANK).to(tl.pointer_type(tl.uint32))
365 + output_u32 += output_peer_offset_u32
366 + offsets = token * (C // 2) + pair
367 + channel = tl.arange(0, C_P2)
368 + channel_mask = channel < C
369 + # Gamma is independent of the preceding sconv publication. Keep the
370 + # acquire immediately before polling the Lamport output slots.
371 + if HAS_NORM:
372 + weight = tl.load(norm_weight_ptr + channel, mask=channel_mask, other=0.0)
373 + if USE_PDL:
374 + tl.extra.cuda.gdc_wait()
375 + packed = _wait_pairs(output_u32, offsets, pair_mask)
376 + row = _unpack_bf16_pairs(packed)
377 + tl.store(
378 + residual_out_ptr + token * stride_output_t + channel, row, mask=channel_mask
379 + )
380 + if HAS_NORM:
381 + row_f32 = tl.where(channel_mask, row.to(tl.float32), 0.0)
382 + inv_rms = tl.rsqrt(tl.sum(row_f32 * row_f32, axis=0) / C + eps)
383 + tl.store(
384 + normed_ptr + token * stride_output_t + channel,
385 + (row_f32 * inv_rms * weight.to(tl.float32)).to(tl.bfloat16),
386 + mask=channel_mask,
387 + )
388 + tl.store(
389 + output_u32 + offsets,
390 + tl.full([C_P2 // 2], _EMPTY_PAIR, tl.uint32),
391 + mask=pair_mask,
392 + )
393 + if USE_PDL:
394 + tl.extra.cuda.gdc_launch_dependents()
395 +
396 +
397 +@triton.jit
398 +def _validate_lamport_init_kernel(
399 + input_peer_ptrs,
400 + output_peer_ptrs,
401 + bad_ptr,
402 + num_pairs,
403 + RANK: tl.constexpr,
404 + BLOCK: tl.constexpr,
405 +):
406 + """Validate both complete local allocations through their fabric pointers."""
407 + offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
408 + mask = offsets < num_pairs
409 + ptrs_in = input_peer_ptrs.to(tl.pointer_type(tl.uint64))
410 + ptrs_out = output_peer_ptrs.to(tl.pointer_type(tl.uint64))
411 + local_in = tl.load(ptrs_in + RANK).to(tl.pointer_type(tl.uint32))
412 + local_out = tl.load(ptrs_out + RANK).to(tl.pointer_type(tl.uint32))
413 + value_in = tl.load(local_in + offsets, mask=mask, other=_EMPTY_PAIR)
414 + value_out = tl.load(local_out + offsets, mask=mask, other=_EMPTY_PAIR)
415 + bad = tl.max(
416 + tl.where(mask & ((value_in != _EMPTY_PAIR) | (value_out != _EMPTY_PAIR)), 1, 0)
417 + )
418 + if bad != 0:
419 + tl.atomic_max(bad_ptr, 1)
420 +
421 +
422 +class LamportRSConv:
423 + """Persistent symmetric buffers for one TP group."""
424 +
425 + def _initialize_lamport_buffers(self) -> None:
426 + """Arm every slot and collectively verify the exact sentinel bits.
427 +
428 + Do not replace this with a zero-fill: Lamport readiness distinguishes
429 + the bf16 bit pattern for -0.0 (0x8000) from every published payload.
430 + The validation is deliberately collective so one rank cannot enter the
431 + first polling kernel while another rank still has an unarmed buffer.
432 + """
433 + if not self.buf_in.is_contiguous() or not self.buf_out.is_contiguous():
434 + raise RuntimeError("Lamport symmetric buffers must be contiguous")
435 + if self.buf_in.numel() % 2 or self.buf_out.numel() % 2:
436 + raise RuntimeError("Lamport buffers must contain whole uint32 pairs")
437 +
438 + # Fill through int16 so each bf16 lane receives the exact -0.0 bits.
439 + # This covers all three generations and all max-token slots, including
440 + # slots that a smaller first invocation does not touch.
441 + self.buf_in.view(torch.int16).fill_(-0x8000)
442 + self.buf_out.view(torch.int16).fill_(-0x8000)
443 + torch.accelerator.synchronize(self.device)
444 + self.tp.barrier()
445 +
446 + # Scan the complete local allocations. Since every rank performs the
447 + # scan and participates in MAX, success proves every symmetric backing
448 + # allocation was armed before any rank is allowed to use generation 0.
449 + bad = torch.logical_or(
450 + self.buf_in.view(torch.int16).ne(-0x8000).any(),
451 + self.buf_out.view(torch.int16).ne(-0x8000).any(),
452 + ).to(dtype=torch.float32)
453 + bad = self.tp.all_reduce(bad)
454 + if int(bad.item()) != 0:
455 + raise RuntimeError("Lamport sentinel initialization failed on a TP rank")
456 + self.tp.barrier()
457 +
458 + def _initialize_mnnvl_buffers(self) -> None:
459 + """Initialize and validate FlashInfer fabric-mapped Lamport storage."""
460 + self._mnnvl_input_handle.lamport_initialize(self.rank, torch.bfloat16)
461 + self._mnnvl_output_handle.lamport_initialize(self.rank, torch.bfloat16)
462 + torch.accelerator.synchronize(self.device)
463 + self.tp.barrier()
464 +
465 + bad = torch.zeros((), dtype=torch.int32, device=self.device)
466 + num_pairs = self.num_buffers * self.max_tokens * self.hidden_size // 2
467 + _validate_lamport_init_kernel[(triton.cdiv(num_pairs, 256),)](
468 + self.input_peer_ptrs,
469 + self.output_peer_ptrs,
470 + bad,
471 + num_pairs,
472 + RANK=self.rank,
473 + BLOCK=256,
474 + num_warps=4,
475 + )
476 + bad = self.tp.all_reduce(bad.to(torch.float32))
477 + if int(bad.item()) != 0:
478 + raise RuntimeError("MNNVL Lamport sentinel initialization failed")
479 + self.tp.barrier()
480 +
481 + def __init__(
482 + self, hidden_size: int, window_size: int, max_tokens: int = _MAX_TOKENS
483 + ) -> None:
484 + import torch.distributed._symmetric_memory as symm_mem
485 +
486 + tp = get_tp_group()
487 + self.tp = tp
488 + self.group = tp.device_group
489 + self.world_size = tp.world_size
490 + self.rank = tp.rank_in_group
491 + self.device = torch.device(tp.device)
492 + if self.world_size not in (2, 4, 8):
493 + raise ValueError(f"TP world size must be 2, 4, or 8, got {self.world_size}")
494 + if hidden_size % (2 * self.world_size) != 0:
495 + raise ValueError("hidden size must produce an even shard on every rank")
496 + if window_size != 4:
497 + raise ValueError(f"short-conv window size must be 4, got {window_size}")
498 + if max_tokens < 1 or max_tokens > _MAX_TOKENS:
499 + raise ValueError(f"max_tokens must be in [1, {_MAX_TOKENS}]")
500 +
501 + is_cross_node = not all(in_the_same_node_as(tp.cpu_group))
502 + self.hidden_size = hidden_size
503 + self.window_size = window_size
504 + self.max_tokens = max_tokens
505 + self.shard_size = hidden_size // self.world_size
506 + # Three generations follow FlashInfer's Lamport layout. A generation
507 + # is reused only after two intervening collective calls.
508 + self.num_buffers = 3
509 + self.input_generation_bytes = max_tokens * hidden_size * 2
510 + self.output_generation_bytes = max_tokens * hidden_size * 2
511 + self.use_pdl = torch.cuda.get_device_capability(self.device)[0] >= 9
512 + if is_cross_node:
513 + try:
514 + from flashinfer.comm.mnnvl import (
515 + McastGPUBuffer,
516 + TorchDistBackend,
517 + is_mnnvl_fabric_supported,
518 + )
519 + except ImportError as error:
520 + raise RuntimeError(
521 + "cross-node TP requires FlashInfer MNNVL support"
522 + ) from error
523 +
524 + local_supported = int(
525 + is_mnnvl_fabric_supported(torch.accelerator.current_device_index())
526 + )
527 + unsupported = torch.tensor(
528 + 1 - local_supported, dtype=torch.float32, device=self.device
529 + )
530 + unsupported = tp.all_reduce(unsupported)
531 + if int(unsupported.item()) != 0:
532 + raise RuntimeError("cross-node TP is supported only on MNNVL fabric")
533 +
534 + comm_backend = TorchDistBackend(self.group)
535 + allocation_bytes = self.num_buffers * max_tokens * hidden_size * 2
536 + self._mnnvl_input_handle = McastGPUBuffer(
537 + allocation_bytes,
538 + self.world_size,
539 + self.rank,
540 + self.device,
541 + comm_backend,
542 + )
543 + self._mnnvl_output_handle = McastGPUBuffer(
544 + allocation_bytes,
545 + self.world_size,
546 + self.rank,
547 + self.device,
548 + comm_backend,
549 + )
550 + self.input_peer_ptrs = self._mnnvl_input_handle.get_buffer_ptrs_dev()
551 + self.output_peer_ptrs = self._mnnvl_output_handle.get_buffer_ptrs_dev()
552 + self._initialize_mnnvl_buffers()
553 + logger.info("using FlashInfer fabric-mapped MNNVL Lamport buffers")
554 + else:
555 + self.buf_in = symm_mem.empty(
556 + self.num_buffers,
557 + max_tokens,
558 + self.world_size,
559 + self.shard_size,
560 + dtype=torch.bfloat16,
561 + device=self.device,
562 + )
563 + self.buf_out = symm_mem.empty(
564 + self.num_buffers,
565 + max_tokens,
566 + hidden_size,
567 + dtype=torch.bfloat16,
568 + device=self.device,
569 + )
570 + group_name = self.group.group_name
571 + input_handle = symm_mem.rendezvous(self.buf_in, group_name)
572 + output_handle = symm_mem.rendezvous(self.buf_out, group_name)
573 + self.input_peer_ptrs = input_handle.buffer_ptrs_dev
574 + self.output_peer_ptrs = output_handle.buffer_ptrs_dev
575 + self._input_handle = input_handle
576 + self._output_handle = output_handle
577 + self._initialize_lamport_buffers()
578 + self.generation = 0
579 +
580 + def usable(self, num_tokens: int) -> bool:
581 + return 0 < num_tokens <= self.max_tokens
582 +
583 + def rs_sconv_ag_add_norm(
584 + self,
585 + input_tensor: torch.Tensor,
586 + residual: torch.Tensor,
587 + conv_weight: torch.Tensor,
588 + norm_weight: torch.Tensor | None,
589 + eps: float,
590 + cache: torch.Tensor,
591 + positions: torch.Tensor,
592 + block_table: torch.Tensor,
593 + seq_idx: torch.Tensor,
594 + slot_mapping: torch.Tensor,
595 + off_s: int,
596 + ws: int,
597 + block_size: int,
598 + ) -> tuple[torch.Tensor | None, torch.Tensor]:
599 + """Return ``(normed | None, new_residual)``, both shaped ``[T, 6144]``."""
600 + tokens, hidden_size = residual.shape
601 + if not self.usable(tokens):
602 + raise ValueError(f"num_tokens must be in [1, {self.max_tokens}]")
603 + if hidden_size != self.hidden_size or residual.dtype != torch.bfloat16:
604 + raise ValueError("residual must be bf16 [T, 6144]")
605 + if (
606 + input_tensor.shape != residual.shape
607 + or input_tensor.dtype != torch.bfloat16
608 + or input_tensor.stride(1) != 1
609 + ):
610 + raise ValueError("input_tensor must be channel-contiguous bf16 [T, 6144]")
611 + shard_size = hidden_size // self.world_size
612 + if conv_weight.shape != (shard_size, self.window_size):
613 + raise ValueError(
614 + f"conv_weight must have shape [{shard_size}, {self.window_size}]"
615 + )
616 + if (
617 + conv_weight.dtype != torch.bfloat16
618 + or conv_weight.stride(0) != self.window_size
619 + ):
620 + raise ValueError("conv_weight must be contiguous bf16")
621 + if norm_weight is not None and (
622 + norm_weight.shape != (hidden_size,) or norm_weight.dtype != torch.bfloat16
623 + ):
624 + raise ValueError("norm_weight must be bf16 [6144] or None")
625 + if cache.dtype != torch.bfloat16 or cache.ndim != 4:
626 + raise ValueError("cache must be a 4-D bf16 tensor")
627 + if shard_size % ws != 0 or cache.shape[1] != shard_size // ws:
628 + raise ValueError("cache head layout is inconsistent with ws")
629 + if off_s < 0 or off_s + ws > cache.shape[3]:
630 + raise ValueError("cache channel offset is out of bounds")
631 +
632 + index = self.generation
633 + input_offset = index * self.input_generation_bytes // 4
634 + output_offset = index * self.output_generation_bytes // 4
635 + normed = torch.empty_like(residual) if norm_weight is not None else None
636 + residual_out = torch.empty_like(residual)
637 + phase_splits = 8 if tokens == 1 and shard_size % 8 == 0 else 1
638 + phase_tile_p2 = triton.next_power_of_2(shard_size // phase_splits)
639 + phase_grid = (tokens, phase_splits)
640 + # Wide CTAs win before the grid saturates; smaller CTAs reduce register
641 + # pressure once high-throughput batches provide enough parallelism.
642 + if 128 <= tokens <= 2048:
643 + phase_warps = 16
644 + elif tokens > 2048:
645 + phase_warps = 8
646 + else:
647 + phase_warps = 4
648 + gather_warps = 4 if tokens >= 256 else 8
649 + _publish_input_kernel[phase_grid](
650 + input_tensor,
651 + self.input_peer_ptrs,
652 + input_offset,
653 + input_tensor.stride(0),
654 + C=hidden_size,
655 + CS=shard_size,
656 + CS_P2=phase_tile_p2,
657 + SPLITS=phase_splits,
658 + RANK=self.rank,
659 + WORLD=self.world_size,
660 + USE_PDL=self.use_pdl,
661 + launch_pdl=self.use_pdl,
662 + num_warps=phase_warps,
663 + )
664 + _reduce_insert_kernel[phase_grid](
665 + self.input_peer_ptrs,
666 + input_offset,
667 + cache,
668 + slot_mapping,
669 + cache.stride(0),
670 + cache.stride(1),
671 + cache.stride(2),
672 + cache.stride(3),
673 + block_size,
674 + C=hidden_size,
675 + CS=shard_size,
676 + CS_P2=phase_tile_p2,
677 + SPLITS=phase_splits,
678 + HEAD_SIZE=ws,
679 + CACHE_OFFSET=off_s,
680 + RANK=self.rank,
681 + WORLD=self.world_size,
682 + USE_PDL=self.use_pdl,
683 + launch_pdl=self.use_pdl,
684 + num_warps=phase_warps,
685 + )
686 + _sconv_publish_kernel[phase_grid](
687 + self.output_peer_ptrs,
688 + output_offset,
689 + residual,
690 + conv_weight,
691 + cache,
692 + positions,
693 + seq_idx,
694 + slot_mapping,
695 + block_table,
696 + residual.stride(0),
697 + cache.stride(0),
698 + cache.stride(1),
699 + cache.stride(2),
700 + cache.stride(3),
701 + block_table.stride(0),
702 + block_table.shape[1],
703 + block_size,
704 + C=hidden_size,
705 + CS=shard_size,
706 + CS_P2=phase_tile_p2,
707 + SPLITS=phase_splits,
708 + HEAD_SIZE=ws,
709 + CACHE_OFFSET=off_s,
710 + RANK=self.rank,
711 + WORLD=self.world_size,
712 + WINDOW=self.window_size,
713 + USE_PDL=self.use_pdl,
714 + launch_pdl=self.use_pdl,
715 + num_warps=phase_warps,
716 + )
717 + _gather_norm_kernel[(tokens,)](
718 + self.output_peer_ptrs,
719 + output_offset,
720 + norm_weight if norm_weight is not None else residual,
721 + normed if normed is not None else residual_out,
722 + residual_out,
723 + eps,
724 + residual.stride(0),
725 + C=hidden_size,
726 + C_P2=triton.next_power_of_2(hidden_size),
727 + RANK=self.rank,
728 + HAS_NORM=norm_weight is not None,
729 + USE_PDL=self.use_pdl,
730 + launch_pdl=self.use_pdl,
731 + num_warps=gather_warps,
732 + )
733 + self.generation = (index + 1) % self.num_buffers
734 + return normed, residual_out
735 +
736 +
737 +_STATE: LamportRSConv | None = None
738 +_STATE_FAILED = False
739 +
740 +
741 +def initialize_lamport_rs_conv(
742 + hidden_size: int, window_size: int, max_num_batched_tokens: int
743 +) -> None:
744 + """Collectively initialize the TP-group state during model construction."""
745 + global _STATE, _STATE_FAILED
746 + if _STATE is not None:
747 + if _STATE.hidden_size != hidden_size or _STATE.window_size != window_size:
748 + raise RuntimeError("all Lamport users must share hidden and window sizes")
749 + return
750 + if _STATE_FAILED or os.environ.get("LAMPORT_RS_SCONV", "1") == "0":
751 + return
752 + try:
753 + max_tokens = min(_MAX_TOKENS, max_num_batched_tokens)
754 + _STATE = LamportRSConv(hidden_size, window_size, max_tokens=max_tokens)
755 + except Exception:
756 + _STATE_FAILED = True
757 + logger.exception("fused collective unavailable; use the NCCL fallback")
758 +
759 +
760 +def get_lamport_rs_conv(hidden_size: int, window_size: int) -> LamportRSConv | None:
761 + """Return the state initialized with the model, or ``None`` for fallback."""
762 + if _STATE is not None and (
763 + _STATE.hidden_size != hidden_size or _STATE.window_size != window_size
764 + ):
765 + raise RuntimeError("all Lamport users must share hidden and window sizes")
766 + return _STATE
vllm/models/inkling/nvidia/ops/mm_towers.py new
+190
@@ -0,0 +1,190 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Fused CUDA kernels for the Inkling vision/audio towers.
4 +
5 +Both kernels keep the reference paths' fp32 accumulation and per-op bf16
6 +rounding points (native ``rms_norm`` / ``F.gelu``); outputs are frequently
7 +bit-identical and otherwise differ by 1-2 bf16 ulps from reduction-order
8 +(real-checkpoint-weight cosine vs reference > 0.9999998).
9 +"""
10 +
11 +from __future__ import annotations
12 +
13 +import torch
14 +
15 +from vllm.triton_utils import tl, tldevice, triton
16 +
17 +from .norm import _get_num_warps_from_block_size
18 +
19 +
20 +@triton.jit
21 +def _dmel_embed_sum_norm_kernel(
22 + idx_ptr, # [T, NB] int32 dMel bin indices (values in [0, VOCAB))
23 + w_ptr, # [NB * VOCAB, D] bf16 embedding table
24 + norm_w_ptr, # [D] (unused if HAS_NORM=False)
25 + out_ptr, # [T, D] bf16
26 + eps,
27 + D,
28 + stride_idx_t,
29 + NB: tl.constexpr,
30 + VOCAB: tl.constexpr,
31 + D_P2: tl.constexpr,
32 + HAS_NORM: tl.constexpr,
33 +):
34 + t = tl.program_id(0).to(tl.int64)
35 + offs = tl.arange(0, D_P2)
36 + mask = offs < D
37 + # One embedding row per mel bin (bin b uses table rows [b*VOCAB, (b+1)*VOCAB)),
38 + # summed in fp32 (matches torch's fp32-accumulated bf16 .sum()).
39 + acc = tl.zeros([D_P2], dtype=tl.float32)
40 + for b in tl.static_range(NB):
41 + v = tl.load(idx_ptr + t * stride_idx_t + b)
42 + row = (b * VOCAB + v).to(tl.int64)
43 + acc += tl.load(w_ptr + row * D + offs, mask=mask, other=0.0).to(tl.float32)
44 + h = acc.to(tl.bfloat16)
45 + if HAS_NORM:
46 + # Match ir.ops.rms_norm: fp32 variance/normalize, then a single-rounded
47 + # bf16 multiply with the bf16 weight.
48 + x32 = h.to(tl.float32)
49 + var = tl.sum(tl.where(mask, x32 * x32, 0.0), axis=0) / D
50 + xn = (x32 * tl.math.rsqrt(var + eps)).to(tl.bfloat16)
51 + w = tl.load(norm_w_ptr + offs, mask=mask, other=0.0)
52 + h = xn * w
53 + tl.store(out_ptr + t * D + offs, h, mask=mask)
54 +
55 +
56 +def dmel_embed_sum_norm(
57 + dmel_idx: torch.Tensor, # [T, NB] int32
58 + weight: torch.Tensor, # [NB * VOCAB, D] bf16
59 + norm_weight: torch.Tensor | None,
60 + eps: float,
61 +) -> torch.Tensor:
62 + """``rmsnorm(sum_b weight[b * VOCAB + idx[:, b]])`` in one launch (no
63 + [T, NB, D] intermediate)."""
64 + T, nb = dmel_idx.shape
65 + D = weight.shape[1]
66 + assert weight.shape[0] % nb == 0
67 + vocab = weight.shape[0] // nb
68 + out = torch.empty((T, D), dtype=weight.dtype, device=weight.device)
69 + if T == 0:
70 + return out
71 + d_p2 = triton.next_power_of_2(D)
72 + _dmel_embed_sum_norm_kernel[(T,)](
73 + dmel_idx,
74 + weight,
75 + norm_weight if norm_weight is not None else weight,
76 + out,
77 + eps,
78 + D,
79 + dmel_idx.stride(0),
80 + NB=nb,
81 + VOCAB=vocab,
82 + D_P2=d_p2,
83 + HAS_NORM=norm_weight is not None,
84 + # Swept on GB200: 8 warps beats the block-size heuristic's 16 by ~4%
85 + # at large T (the 80-row gather chain is latency- not lane-bound).
86 + num_warps=8,
87 + )
88 + return out
89 +
90 +
91 +@triton.jit
92 +def _rmsnorm_gelu_kernel(
93 + x_ptr, # [R, D] bf16
94 + w_ptr, # [D]
95 + out_ptr, # [R, D] bf16 (or the folded layout when FOLD)
96 + eps,
97 + R,
98 + D,
99 + D_P2: tl.constexpr,
100 + BLOCK_M: tl.constexpr, # rows per block (>1 for small D)
101 + HAS_GELU: tl.constexpr,
102 + FOLD: tl.constexpr,
103 + # fold geometry: input rows index [N, T, H, W]; the store scatters each
104 + # row to (out_row, slot) of fold_timespace_to_depth's output layout.
105 + FT: tl.constexpr,
106 + FH: tl.constexpr,
107 + FW: tl.constexpr,
108 + TF: tl.constexpr,
109 + HF: tl.constexpr,
110 +):
111 + pid = tl.program_id(0).to(tl.int64)
112 + rows = pid * BLOCK_M + tl.arange(0, BLOCK_M)
113 + rmask = rows < R
114 + offs = tl.arange(0, D_P2)
115 + mask = rmask[:, None] & (offs < D)[None, :]
116 + x32 = tl.load(x_ptr + rows[:, None] * D + offs[None, :], mask=mask, other=0.0).to(
117 + tl.float32
118 + )
119 + var = tl.sum(x32 * x32, axis=1) / D
120 + xn = (x32 * tl.math.rsqrt(var + eps)[:, None]).to(tl.bfloat16)
121 + w = tl.load(w_ptr + offs, mask=offs < D, other=0.0)
122 + h = xn * w[None, :] # bf16 multiply, matching ir.ops.rms_norm
123 + if HAS_GELU:
124 + # Exact (erf) GELU on the bf16-rounded norm output, fp32 math,
125 + # matching F.gelu's opmath on a bf16 tensor.
126 + g32 = h.to(tl.float32)
127 + h = (0.5 * g32 * (1.0 + tldevice.erf(g32 * 0.7071067811865476))).to(tl.bfloat16)
128 + if FOLD:
129 + # Store directly into the next layer's folded layout (a pure
130 + # permutation — replaces the separate fold copy pass).
131 + t = (rows // (FH * FW)) % FT
132 + hh = (rows // FW) % FH
133 + ww = rows % FW
134 + n = rows // (FT * FH * FW)
135 + slot = ((t % TF) * HF + hh % HF) * HF + ww % HF
136 + out_row = ((n * (FT // TF) + t // TF) * (FH // HF) + hh // HF) * (
137 + FW // HF
138 + ) + ww // HF
139 + base = (out_row * (TF * HF * HF) + slot) * D
140 + tl.store(out_ptr + base[:, None] + offs[None, :], h, mask=mask)
141 + else:
142 + tl.store(out_ptr + rows[:, None] * D + offs[None, :], h, mask=mask)
143 +
144 +
145 +def rmsnorm_gelu(
146 + x: torch.Tensor, # [..., D] bf16 contiguous
147 + weight: torch.Tensor,
148 + eps: float,
149 + gelu: bool = True,
150 + fold: tuple[int, int] | None = None, # (t_fold, hw_fold) of the NEXT fold
151 +) -> torch.Tensor:
152 + """Fused ``gelu(rmsnorm(x))`` (or plain rmsnorm); multiple rows per block
153 + when D is small. With ``fold``, x must be [N, T, H, W, D] and the output
154 + comes back as ``fold_timespace_to_depth(result, *fold)``."""
155 + D = x.shape[-1]
156 + flat = x.reshape(-1, D)
157 + assert flat.stride(1) == 1 and flat.stride(0) == D
158 + R = flat.shape[0]
159 + if fold is None:
160 + out = torch.empty_like(flat)
161 + ft = fh = fw = tf = hf = 1
162 + out_shape = x.shape
163 + else:
164 + tf, hf = fold
165 + N, ft, fh, fw, _ = x.shape
166 + out_shape = (N, ft // tf, fh // hf, fw // hf, tf * hf * hf * D)
167 + out = torch.empty(out_shape, dtype=x.dtype, device=x.device)
168 + if R == 0:
169 + return out.reshape(out_shape)
170 + d_p2 = triton.next_power_of_2(D)
171 + block_m = max(1, 4096 // d_p2)
172 + _rmsnorm_gelu_kernel[(triton.cdiv(R, block_m),)](
173 + flat,
174 + weight,
175 + out,
176 + eps,
177 + R,
178 + D,
179 + D_P2=d_p2,
180 + BLOCK_M=block_m,
181 + HAS_GELU=gelu,
182 + FOLD=fold is not None,
183 + FT=ft,
184 + FH=fh,
185 + FW=fw,
186 + TF=tf,
187 + HF=hf,
188 + num_warps=_get_num_warps_from_block_size(d_p2 * block_m),
189 + )
190 + return out.reshape(out_shape)
vllm/models/inkling/nvidia/ops/norm.py new
+315
@@ -0,0 +1,315 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from __future__ import annotations
4 +
5 +from functools import lru_cache
6 +
7 +import torch
8 +
9 +from vllm.triton_utils import tl, triton
10 +
11 +_MAX_FUSED_SIZE = 65536
12 +
13 +
14 +def _get_num_warps_from_block_size(block_size: int) -> int:
15 + if block_size >= 32768:
16 + return 32
17 + if block_size >= 8192:
18 + return 16
19 + if block_size >= 2048:
20 + return 8
21 + return 4
22 +
23 +
24 +def _largest_power_of_2(n: int) -> int:
25 + assert n > 0, f"{n=}"
26 + return 1 << (n.bit_length() - 1)
27 +
28 +
29 +@lru_cache(maxsize=128)
30 +def _get_grid_size_for_mem_bw_kernel(device: torch.device, factor: int = 8) -> int:
31 + num_sms = torch.cuda.get_device_properties(device).multi_processor_count
32 + return _largest_power_of_2(num_sms) * factor
33 +
34 +
35 +@triton.jit
36 +def _rmsnorm_fwd_kernel(
37 + x_ptr,
38 + weight_ptr,
39 + y_ptr,
40 + rstd_ptr,
41 + eps,
42 + x_stride_0,
43 + y_stride_0,
44 + n_cols,
45 + block_size_n: tl.constexpr,
46 +):
47 + pid_m = tl.program_id(0).to(tl.int64)
48 + offs_n = tl.arange(0, block_size_n)
49 + mask_n = offs_n < n_cols
50 +
51 + weight = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
52 + x = tl.load(x_ptr + pid_m * x_stride_0 + offs_n, mask=mask_n, other=0.0).to(
53 + tl.float32
54 + )
55 + row_var = tl.sum(x * x, axis=0) / n_cols
56 + rstd = tl.math.rsqrt(row_var + eps)
57 + tl.store(rstd_ptr + pid_m, rstd)
58 +
59 + y = x * rstd * weight
60 + tl.store(y_ptr + pid_m * y_stride_0 + offs_n, y, mask=mask_n)
61 +
62 +
63 +@triton.jit(do_not_specialize=["n_rows"])
64 +def _rmsnorm_fwd_kernel_block_m(
65 + x_ptr,
66 + weight_ptr,
67 + y_ptr,
68 + rstd_ptr,
69 + eps,
70 + x_stride_0,
71 + y_stride_0,
72 + n_rows,
73 + n_cols,
74 + block_size_m: tl.constexpr,
75 + block_size_n: tl.constexpr,
76 +):
77 + pid_m = tl.program_id(0).to(tl.int64)
78 + offs_n = tl.arange(0, block_size_n)
79 + mask_n = offs_n < n_cols
80 +
81 + num_blocks_m = tl.cdiv(n_rows, block_size_m)
82 + blocks_per_pid = tl.cdiv(num_blocks_m, tl.num_programs(0))
83 + block_id_start = pid_m * blocks_per_pid
84 + block_id_end = min(block_id_start + blocks_per_pid, num_blocks_m)
85 +
86 + weight = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
87 +
88 + for block_id in range(block_id_start, block_id_end):
89 + offs_m = block_id * block_size_m + tl.arange(0, block_size_m)
90 + mask_m = offs_m < n_rows
91 + mask_mn = mask_m[:, None] & mask_n[None, :]
92 + x = tl.load(
93 + x_ptr + offs_m[:, None] * x_stride_0 + offs_n[None, :],
94 + mask=mask_mn,
95 + other=0.0,
96 + ).to(tl.float32)
97 + row_var = tl.sum(x * x, axis=1) / n_cols
98 + rstd = tl.math.rsqrt(row_var + eps)
99 + tl.store(rstd_ptr + offs_m, rstd, mask=mask_m)
100 +
101 + y = x * rstd[:, None] * weight
102 + tl.store(
103 + y_ptr + offs_m[:, None] * y_stride_0 + offs_n[None, :],
104 + y,
105 + mask=mask_mn,
106 + )
107 +
108 +
109 +@triton.jit
110 +def _add_rmsnorm_fwd_kernel(
111 + res_ptr, # [T, N] residual (read)
112 + delta_ptr, # [T, N] delta to add (read)
113 + weight_ptr,
114 + y_ptr, # [T, N] normed output
115 + res_out_ptr, # [T, N] updated residual output
116 + eps,
117 + res_stride_0,
118 + delta_stride_0,
119 + y_stride_0,
120 + res_out_stride_0,
121 + n_cols,
122 + block_size_n: tl.constexpr,
123 +):
124 + pid_m = tl.program_id(0).to(tl.int64)
125 + offs_n = tl.arange(0, block_size_n)
126 + mask_n = offs_n < n_cols
127 +
128 + weight = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
129 + r = tl.load(res_ptr + pid_m * res_stride_0 + offs_n, mask=mask_n, other=0.0).to(
130 + tl.float32
131 + )
132 + d = tl.load(delta_ptr + pid_m * delta_stride_0 + offs_n, mask=mask_n, other=0.0).to(
133 + tl.float32
134 + )
135 + # Round the sum to the residual dtype first (matches the eager
136 + # `residual + delta` then rmsnorm-on-bf16 sequence bit-for-bit).
137 + s = (r + d).to(res_out_ptr.dtype.element_ty)
138 + tl.store(res_out_ptr + pid_m * res_out_stride_0 + offs_n, s, mask=mask_n)
139 + x = s.to(tl.float32)
140 + row_var = tl.sum(x * x, axis=0) / n_cols
141 + rstd = tl.math.rsqrt(row_var + eps)
142 + y = x * rstd * weight
143 + tl.store(y_ptr + pid_m * y_stride_0 + offs_n, y, mask=mask_n)
144 +
145 +
146 +def add_rmsnorm(
147 + residual: torch.Tensor,
148 + delta: torch.Tensor,
149 + weight: torch.Tensor,
150 + eps: float,
151 +) -> tuple[torch.Tensor, torch.Tensor]:
152 + """Fused ``res = residual + delta; y = rmsnorm(res)``.
153 +
154 + Returns ``(y, res)``; both are fresh tensors (cudagraph-friendly, no
155 + in-place update of the inputs).
156 + """
157 + assert residual.ndim == 2 and delta.ndim == 2, (residual.shape, delta.shape)
158 + n_rows, n_cols = residual.shape
159 + assert weight.shape[0] == n_cols
160 + y = torch.empty_like(residual)
161 + res_out = torch.empty_like(residual)
162 + if n_rows == 0:
163 + return y, res_out
164 +
165 + block_size_n = triton.next_power_of_2(n_cols)
166 + max_block_size_n = _MAX_FUSED_SIZE // residual.element_size()
167 + if max_block_size_n < block_size_n:
168 + raise RuntimeError(f"Large {n_cols=} is not supported")
169 + num_warps = _get_num_warps_from_block_size(block_size_n)
170 + _add_rmsnorm_fwd_kernel[(n_rows,)](
171 + residual,
172 + delta,
173 + weight,
174 + y,
175 + res_out,
176 + eps,
177 + residual.stride(0),
178 + delta.stride(0),
179 + y.stride(0),
180 + res_out.stride(0),
181 + n_cols,
182 + block_size_n,
183 + num_warps=num_warps,
184 + )
185 + return y, res_out
186 +
187 +
188 +@triton.jit
189 +def _embed_rmsnorm_kernel(
190 + ids_ptr, # [T] token ids
191 + table_ptr, # [V, N] embedding table
192 + weight_ptr, # [N] (HAS_NORM only)
193 + chain_weight_ptr, # [N] (HAS_CHAIN only)
194 + out_ptr, # [T, N] rmsnorm(table[ids], weight)
195 + chain_out_ptr, # [T, N] rmsnorm(out, chain_weight) (HAS_CHAIN only)
196 + eps,
197 + table_stride_0,
198 + n_cols,
199 + block_size_n: tl.constexpr,
200 + HAS_NORM: tl.constexpr,
201 + HAS_CHAIN: tl.constexpr,
202 +):
203 + pid_m = tl.program_id(0).to(tl.int64)
204 + offs_n = tl.arange(0, block_size_n)
205 + mask_n = offs_n < n_cols
206 + row = tl.load(ids_ptr + pid_m).to(tl.int64)
207 + x = tl.load(table_ptr + row * table_stride_0 + offs_n, mask=mask_n, other=0.0)
208 + if HAS_NORM:
209 + xf = x.to(tl.float32)
210 + rstd = tl.math.rsqrt(tl.sum(xf * xf, axis=0) / n_cols + eps)
211 + w = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
212 + # Round to the output dtype so the chained norm is bit-exact vs the
213 + # unfused pair (which stores bf16 in between).
214 + x = (xf * rstd * w).to(out_ptr.dtype.element_ty)
215 + tl.store(out_ptr + pid_m * n_cols + offs_n, x, mask=mask_n)
216 + if HAS_CHAIN:
217 + xf = x.to(tl.float32)
218 + rstd = tl.math.rsqrt(tl.sum(xf * xf, axis=0) / n_cols + eps)
219 + w = tl.load(chain_weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
220 + tl.store(
221 + chain_out_ptr + pid_m * n_cols + offs_n,
222 + (xf * rstd * w).to(chain_out_ptr.dtype.element_ty),
223 + mask=mask_n,
224 + )
225 +
226 +
227 +def embed_rmsnorm(
228 + input_ids: torch.Tensor,
229 + embed_table: torch.Tensor,
230 + weight: torch.Tensor | None,
231 + eps: float,
232 + chain_weight: torch.Tensor | None = None,
233 +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
234 + """Fused ``rmsnorm(embed_table[input_ids], weight)`` row gather + norm.
235 +
236 + Requires the full vocab on-rank (replicated or tp_size == 1).
237 + ``weight=None`` skips the norm (``use_embed_norm=False``), leaving a pure
238 + embedding-table gather.
239 + ``chain_weight`` additionally emits ``rmsnorm(out, chain_weight)`` (the
240 + first decoder layer's pre-attention norm) as a second output, still one
241 + launch. Bit-exact vs the unfused module sequence."""
242 + ids = input_ids.view(-1)
243 + (T,) = ids.shape
244 + n = embed_table.shape[1]
245 + out = torch.empty(
246 + (*input_ids.shape, n), dtype=embed_table.dtype, device=embed_table.device
247 + )
248 + chain_out = torch.empty_like(out) if chain_weight is not None else None
249 + if T > 0:
250 + block_size_n = triton.next_power_of_2(n)
251 + _embed_rmsnorm_kernel[(T,)](
252 + ids,
253 + embed_table,
254 + weight if weight is not None else embed_table,
255 + chain_weight if chain_weight is not None else embed_table,
256 + out,
257 + chain_out if chain_out is not None else out,
258 + eps,
259 + embed_table.stride(0),
260 + n,
261 + block_size_n,
262 + HAS_NORM=weight is not None,
263 + HAS_CHAIN=chain_weight is not None,
264 + num_warps=_get_num_warps_from_block_size(block_size_n),
265 + )
266 + if chain_out is not None:
267 + return out, chain_out
268 + return out
269 +
270 +
271 +def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
272 + assert x.ndim == 2, f"{x.shape=}"
273 + assert weight.ndim == 1, f"{weight.shape=}"
274 + n_rows, n_cols = x.shape
275 + assert weight.shape[0] == n_cols, f"{weight.shape=} {x.shape=}"
276 + y = torch.empty_like(x)
277 + rstd = torch.empty((n_rows,), dtype=torch.float32, device=x.device)
278 +
279 + block_size_n = triton.next_power_of_2(n_cols)
280 + max_block_size_n = _MAX_FUSED_SIZE // x.element_size()
281 + if max_block_size_n < block_size_n:
282 + raise RuntimeError(f"Large {n_cols=} is not supported")
283 + block_size_m = max(1, 4096 // block_size_n)
284 + num_warps = _get_num_warps_from_block_size(block_size_n)
285 +
286 + if block_size_m == 1:
287 + _rmsnorm_fwd_kernel[(n_rows,)](
288 + x,
289 + weight,
290 + y,
291 + rstd,
292 + eps,
293 + x.stride(0),
294 + y.stride(0),
295 + n_cols,
296 + block_size_n,
297 + num_warps=num_warps,
298 + )
299 + else:
300 + grid_size = _get_grid_size_for_mem_bw_kernel(x.device)
301 + _rmsnorm_fwd_kernel_block_m[(grid_size,)](
302 + x,
303 + weight,
304 + y,
305 + rstd,
306 + eps,
307 + x.stride(0),
308 + y.stride(0),
309 + n_rows,
310 + n_cols,
311 + block_size_m,
312 + block_size_n,
313 + num_warps=num_warps,
314 + )
315 + return y
vllm/models/inkling/nvidia/ops/qkvr_prep.py new
+918
@@ -0,0 +1,918 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm.triton_utils import tl, triton
7 +from vllm.utils.torch_utils import aux_stream
8 +
9 +LOW_BLOCK_M = 32
10 +LOW_BLOCK_N = 64
11 +LOW_NUM_WARPS = 4
12 +THROUGHPUT_BLOCK_M = 32
13 +THROUGHPUT_BLOCK_N = 128
14 +THROUGHPUT_GROUP_M = 2
15 +THROUGHPUT_NUM_WARPS = 4
16 +SMALL_TOKEN_THRESHOLD = 128
17 +SMALL_NUM_WARPS = 2
18 +Q_BLOCK_ROWS = 8
19 +Q_NUM_WARPS = 2
20 +KV_BLOCK_ROWS = 4
21 +KV_NUM_WARPS = 2
22 +
23 +
24 +@triton.jit(do_not_specialize=["rows"])
25 +def _rel_proj_low_latency_kernel(
26 + qkvr_ptr,
27 + rel_proj_ptr,
28 + rel_out_ptr,
29 + log_scaling_ptr,
30 + rows,
31 + stride_x_t,
32 + R_OFFSET: tl.constexpr,
33 + NUM_Q_HEADS: tl.constexpr,
34 + REL_EXTENT: tl.constexpr,
35 + BLOCK_M: tl.constexpr,
36 + BLOCK_N: tl.constexpr,
37 + APPLY_LOG_SCALING: tl.constexpr,
38 +):
39 + row = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
40 + col = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
41 + inner = tl.arange(0, 16)
42 + token = row // NUM_Q_HEADS
43 + head = row % NUM_Q_HEADS
44 + relative = tl.load(
45 + qkvr_ptr
46 + + token[:, None] * stride_x_t
47 + + R_OFFSET
48 + + head[:, None] * 16
49 + + inner[None, :],
50 + mask=row[:, None] < rows,
51 + other=0.0,
52 + )
53 + projection = tl.load(
54 + rel_proj_ptr + inner[:, None] * REL_EXTENT + col[None, :],
55 + mask=col[None, :] < REL_EXTENT,
56 + other=0.0,
57 + )
58 + values = tl.dot(relative, projection, out_dtype=tl.float32).to(
59 + rel_out_ptr.dtype.element_ty
60 + )
61 + values = values.to(tl.float32)
62 + if APPLY_LOG_SCALING:
63 + values *= tl.load(log_scaling_ptr + token, mask=row < rows, other=1.0)[:, None]
64 + tl.store(
65 + rel_out_ptr + row[:, None] * REL_EXTENT + col[None, :],
66 + values.to(rel_out_ptr.dtype.element_ty),
67 + mask=(row[:, None] < rows) & (col[None, :] < REL_EXTENT),
68 + )
69 +
70 +
71 +@triton.jit(do_not_specialize=["rows"])
72 +def _rel_proj_throughput_kernel(
73 + qkvr_ptr,
74 + rel_proj_ptr,
75 + rel_out_ptr,
76 + log_scaling_ptr,
77 + rows,
78 + stride_x_t,
79 + R_OFFSET: tl.constexpr,
80 + NUM_Q_HEADS: tl.constexpr,
81 + REL_EXTENT: tl.constexpr,
82 + BLOCK_M: tl.constexpr,
83 + BLOCK_N: tl.constexpr,
84 + GROUP_M: tl.constexpr,
85 + APPLY_LOG_SCALING: tl.constexpr,
86 +):
87 + row_group = tl.program_id(0)
88 + col = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
89 + inner = tl.arange(0, 16)
90 + projection = tl.load(
91 + rel_proj_ptr + inner[:, None] * REL_EXTENT + col[None, :],
92 + mask=col[None, :] < REL_EXTENT,
93 + other=0.0,
94 + )
95 + row_offsets = tl.arange(0, BLOCK_M)
96 + for group_offset in tl.static_range(GROUP_M):
97 + row = (row_group * GROUP_M + group_offset) * BLOCK_M + row_offsets
98 + token = row // NUM_Q_HEADS
99 + head = row % NUM_Q_HEADS
100 + relative = tl.load(
101 + qkvr_ptr
102 + + token[:, None] * stride_x_t
103 + + R_OFFSET
104 + + head[:, None] * 16
105 + + inner[None, :],
106 + mask=row[:, None] < rows,
107 + other=0.0,
108 + )
109 + values = tl.dot(relative, projection, out_dtype=tl.float32).to(
110 + rel_out_ptr.dtype.element_ty
111 + )
112 + values = values.to(tl.float32)
113 + if APPLY_LOG_SCALING:
114 + values *= tl.load(log_scaling_ptr + token, mask=row < rows, other=1.0)[
115 + :, None
116 + ]
117 + tl.store(
118 + rel_out_ptr + row[:, None] * REL_EXTENT + col[None, :],
119 + values.to(rel_out_ptr.dtype.element_ty),
120 + mask=(row[:, None] < rows) & (col[None, :] < REL_EXTENT),
121 + )
122 +
123 +
124 +def use_rel_proj_throughput(rows: int, rel_extent: int) -> bool:
125 + min_rows = 8192 if rel_extent == 512 else 2048
126 + return rows >= min_rows
127 +
128 +
129 +def qkvr_rel_proj(
130 + qkvr: torch.Tensor,
131 + rel_proj: torch.Tensor,
132 + rel_out: torch.Tensor,
133 + log_scaling: torch.Tensor | None,
134 + *,
135 + num_q_heads: int,
136 + num_kv_heads: int,
137 + head_dim: int,
138 + d_rel: int,
139 +) -> None:
140 + rows = qkvr.shape[0] * num_q_heads
141 + rel_extent = rel_proj.shape[1]
142 + assert d_rel == 16 and rel_proj.shape[0] == 16
143 + r_offset = num_q_heads * head_dim + 2 * num_kv_heads * head_dim
144 + log_scaling_ptr = log_scaling if log_scaling is not None else qkvr
145 + common = dict(
146 + R_OFFSET=r_offset,
147 + NUM_Q_HEADS=num_q_heads,
148 + REL_EXTENT=rel_extent,
149 + APPLY_LOG_SCALING=log_scaling is not None,
150 + )
151 +
152 + if use_rel_proj_throughput(rows, rel_extent):
153 + grid = (
154 + triton.cdiv(rows, THROUGHPUT_BLOCK_M * THROUGHPUT_GROUP_M),
155 + triton.cdiv(rel_extent, THROUGHPUT_BLOCK_N),
156 + )
157 + _rel_proj_throughput_kernel[grid](
158 + qkvr,
159 + rel_proj,
160 + rel_out,
161 + log_scaling_ptr,
162 + rows,
163 + qkvr.stride(0),
164 + BLOCK_M=THROUGHPUT_BLOCK_M,
165 + BLOCK_N=THROUGHPUT_BLOCK_N,
166 + GROUP_M=THROUGHPUT_GROUP_M,
167 + num_warps=THROUGHPUT_NUM_WARPS,
168 + **common,
169 + )
170 + return
171 +
172 + grid = (
173 + triton.cdiv(rows, LOW_BLOCK_M),
174 + triton.cdiv(rel_extent, LOW_BLOCK_N),
175 + )
176 + _rel_proj_low_latency_kernel[grid](
177 + qkvr,
178 + rel_proj,
179 + rel_out,
180 + log_scaling_ptr,
181 + rows,
182 + qkvr.stride(0),
183 + BLOCK_M=LOW_BLOCK_M,
184 + BLOCK_N=LOW_BLOCK_N,
185 + num_warps=LOW_NUM_WARPS,
186 + **common,
187 + )
188 +
189 +
190 +@triton.jit(do_not_specialize=["tokens", "stride_block_table_req", "max_blocks"])
191 +def _qkvr_qkv_kernel(
192 + qkvr_ptr,
193 + q_norm_weight_ptr,
194 + q_out_ptr,
195 + rel_proj_ptr,
196 + rel_out_ptr,
197 + k_weight_ptr,
198 + v_weight_ptr,
199 + k_norm_weight_ptr,
200 + conv_cache_ptr,
201 + key_cache_ptr,
202 + value_cache_ptr,
203 + positions_ptr,
204 + seq_idx_ptr,
205 + conv_slot_mapping_ptr,
206 + conv_block_table_ptr,
207 + query_start_ptr,
208 + attention_slot_mapping_ptr,
209 + log_scaling_ptr,
210 + tokens,
211 + eps,
212 + stride_x_t,
213 + stride_cc_block,
214 + stride_cc_head,
215 + stride_cc_token,
216 + stride_cc_dim,
217 + stride_kc_block,
218 + stride_kc_token,
219 + stride_kc_head,
220 + stride_vc_block,
221 + stride_vc_token,
222 + stride_vc_head,
223 + stride_block_table_req,
224 + max_blocks,
225 + conv_block_size,
226 + attention_page_size,
227 + Q_WIDTH: tl.constexpr,
228 + KV_WIDTH: tl.constexpr,
229 + NUM_Q_HEADS: tl.constexpr,
230 + NUM_KV_HEADS: tl.constexpr,
231 + HEAD_DIM: tl.constexpr,
232 + WINDOW_SIZE: tl.constexpr,
233 + OFF_K: tl.constexpr,
234 + OFF_V: tl.constexpr,
235 + APPLY_LOG_SCALING: tl.constexpr,
236 + D_REL: tl.constexpr,
237 + REL_EXTENT: tl.constexpr,
238 + REL_EXTENT_PADDED: tl.constexpr,
239 +):
240 + block = tl.program_id(0)
241 + num_q_rows = tokens * NUM_Q_HEADS
242 + dims = tl.arange(0, HEAD_DIM)
243 +
244 + if block < num_q_rows:
245 + row = block
246 + token = row // NUM_Q_HEADS
247 + head = row % NUM_Q_HEADS
248 + values = tl.load(
249 + qkvr_ptr + token * stride_x_t + head * HEAD_DIM + dims,
250 + ).to(tl.float32)
251 + weight = tl.load(q_norm_weight_ptr + dims).to(tl.float32)
252 + rstd = tl.rsqrt(tl.sum(values * values, axis=0) / HEAD_DIM + eps)
253 + normalized = values * rstd * weight
254 + if APPLY_LOG_SCALING:
255 + normalized = normalized.to(q_out_ptr.dtype.element_ty).to(tl.float32)
256 + normalized *= tl.load(log_scaling_ptr + token)
257 + tl.store(
258 + q_out_ptr + row * HEAD_DIM + dims,
259 + normalized.to(q_out_ptr.dtype.element_ty),
260 + )
261 + rel_cols = tl.arange(0, REL_EXTENT_PADDED)
262 + rel_mask = rel_cols < REL_EXTENT
263 + projected = tl.zeros([REL_EXTENT_PADDED], dtype=tl.float32)
264 + rel_offset = Q_WIDTH + 2 * KV_WIDTH + head * D_REL
265 + for rel_dim in tl.static_range(D_REL):
266 + rel_value = tl.load(
267 + qkvr_ptr + token * stride_x_t + rel_offset + rel_dim
268 + ).to(tl.float32)
269 + proj = tl.load(
270 + rel_proj_ptr + rel_dim * REL_EXTENT + rel_cols,
271 + mask=rel_mask,
272 + other=0.0,
273 + ).to(tl.float32)
274 + projected += rel_value * proj
275 + projected = projected.to(rel_out_ptr.dtype.element_ty).to(tl.float32)
276 + if APPLY_LOG_SCALING:
277 + projected *= tl.load(log_scaling_ptr + token)
278 + tl.store(
279 + rel_out_ptr + row * REL_EXTENT + rel_cols,
280 + projected.to(rel_out_ptr.dtype.element_ty),
281 + mask=rel_mask,
282 + )
283 + else:
284 + row = block - num_q_rows
285 + if row < tokens * NUM_KV_HEADS:
286 + token = row // NUM_KV_HEADS
287 + head = row % NUM_KV_HEADS
288 + position = tl.load(positions_ptr + token)
289 + request = tl.load(seq_idx_ptr + token)
290 + conv_slot = tl.load(conv_slot_mapping_ptr + token)
291 + query_start = tl.load(query_start_ptr + token)
292 + attention_slot = tl.load(attention_slot_mapping_ptr + token)
293 + valid = conv_slot >= 0
294 +
295 + k_col = Q_WIDTH + head * HEAD_DIM
296 + v_col = Q_WIDTH + KV_WIDTH + head * HEAD_DIM
297 + k_value = tl.load(qkvr_ptr + token * stride_x_t + k_col + dims)
298 + v_value = tl.load(qkvr_ptr + token * stride_x_t + v_col + dims)
299 +
300 + safe_slot = tl.maximum(conv_slot, 0)
301 + cache_base = (
302 + conv_cache_ptr
303 + + (safe_slot // conv_block_size) * stride_cc_block
304 + + head * stride_cc_head
305 + + (safe_slot % conv_block_size) * stride_cc_token
306 + )
307 + tl.store(
308 + cache_base + (OFF_K + dims) * stride_cc_dim,
309 + k_value,
310 + mask=valid,
311 + )
312 + tl.store(
313 + cache_base + (OFF_V + dims) * stride_cc_dim,
314 + v_value,
315 + mask=valid,
316 + )
317 +
318 + acc_k = tl.zeros([HEAD_DIM], dtype=tl.float32)
319 + acc_v = tl.zeros([HEAD_DIM], dtype=tl.float32)
320 + for tap in tl.static_range(WINDOW_SIZE):
321 + source_position = position - (WINDOW_SIZE - 1) + tap
322 + source_row = token - (WINDOW_SIZE - 1) + tap
323 + in_window = valid & (source_position >= 0)
324 + intra = in_window & (source_row >= query_start)
325 + cached = in_window & (source_row < query_start)
326 + safe_row = tl.maximum(source_row, 0)
327 + source_k = tl.load(
328 + qkvr_ptr + safe_row * stride_x_t + k_col + dims,
329 + mask=intra,
330 + other=0.0,
331 + ).to(tl.float32)
332 + source_v = tl.load(
333 + qkvr_ptr + safe_row * stride_x_t + v_col + dims,
334 + mask=intra,
335 + other=0.0,
336 + ).to(tl.float32)
337 + safe_position = tl.maximum(source_position, 0)
338 + logical_block = tl.minimum(
339 + safe_position // conv_block_size, max_blocks - 1
340 + )
341 + physical_block = tl.load(
342 + conv_block_table_ptr
343 + + request * stride_block_table_req
344 + + logical_block,
345 + mask=cached,
346 + other=0,
347 + ).to(tl.int64)
348 + tap_base = (
349 + conv_cache_ptr
350 + + physical_block * stride_cc_block
351 + + head * stride_cc_head
352 + + (safe_position % conv_block_size) * stride_cc_token
353 + )
354 + source_k += tl.load(
355 + tap_base + (OFF_K + dims) * stride_cc_dim,
356 + mask=cached,
357 + other=0.0,
358 + ).to(tl.float32)
359 + source_v += tl.load(
360 + tap_base + (OFF_V + dims) * stride_cc_dim,
361 + mask=cached,
362 + other=0.0,
363 + ).to(tl.float32)
364 + k_weight = tl.load(
365 + k_weight_ptr + (head * HEAD_DIM + dims) * WINDOW_SIZE + tap
366 + ).to(tl.float32)
367 + v_weight = tl.load(
368 + v_weight_ptr + (head * HEAD_DIM + dims) * WINDOW_SIZE + tap
369 + ).to(tl.float32)
370 + acc_k += source_k * k_weight
371 + acc_v += source_v * v_weight
372 +
373 + k_rounded = (acc_k + k_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty)
374 + v_rounded = (acc_v + v_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty)
375 + k_float = k_rounded.to(tl.float32)
376 + k_norm_weight = tl.load(k_norm_weight_ptr + dims).to(tl.float32)
377 + rstd = tl.rsqrt(tl.sum(k_float * k_float, axis=0) / HEAD_DIM + eps)
378 + k_normalized = (k_float * rstd * k_norm_weight).to(
379 + qkvr_ptr.dtype.element_ty
380 + )
381 +
382 + safe_attention_slot = tl.maximum(attention_slot, 0)
383 + attention_block = safe_attention_slot // attention_page_size
384 + attention_offset = safe_attention_slot % attention_page_size
385 + tl.store(
386 + key_cache_ptr
387 + + attention_block * stride_kc_block
388 + + attention_offset * stride_kc_token
389 + + head * stride_kc_head
390 + + dims,
391 + k_normalized,
392 + mask=attention_slot >= 0,
393 + )
394 + tl.store(
395 + value_cache_ptr
396 + + attention_block * stride_vc_block
397 + + attention_offset * stride_vc_token
398 + + head * stride_vc_head
399 + + dims,
400 + v_rounded,
401 + mask=attention_slot >= 0,
402 + )
403 +
404 +
405 +@triton.jit(do_not_specialize=["num_rows"])
406 +def _q_kernel(
407 + qkvr_ptr,
408 + q_norm_weight_ptr,
409 + q_out_ptr,
410 + log_scaling_ptr,
411 + num_rows,
412 + stride_x_t,
413 + eps,
414 + NUM_Q_HEADS: tl.constexpr,
415 + HEAD_DIM: tl.constexpr,
416 + BLOCK_ROWS: tl.constexpr,
417 + APPLY_LOG_SCALING: tl.constexpr,
418 +):
419 + rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
420 + dims = tl.arange(0, HEAD_DIM)
421 + row_mask = rows < num_rows
422 + tokens = rows // NUM_Q_HEADS
423 + heads = rows % NUM_Q_HEADS
424 + values = tl.load(
425 + qkvr_ptr
426 + + tokens[:, None] * stride_x_t
427 + + heads[:, None] * HEAD_DIM
428 + + dims[None, :],
429 + mask=row_mask[:, None],
430 + other=0.0,
431 + ).to(tl.float32)
432 + weight = tl.load(q_norm_weight_ptr + dims).to(tl.float32)
433 + rstd = tl.rsqrt(tl.sum(values * values, axis=1) / HEAD_DIM + eps)
434 + normalized = values * rstd[:, None] * weight[None, :]
435 + if APPLY_LOG_SCALING:
436 + normalized = normalized.to(q_out_ptr.dtype.element_ty).to(tl.float32)
437 + tau = tl.load(log_scaling_ptr + tokens, mask=row_mask, other=1.0)
438 + normalized *= tau[:, None]
439 + tl.store(
440 + q_out_ptr + rows[:, None] * HEAD_DIM + dims[None, :],
441 + normalized.to(q_out_ptr.dtype.element_ty),
442 + mask=row_mask[:, None],
443 + )
444 +
445 +
446 +@triton.jit(do_not_specialize=["tokens", "stride_block_table_req", "max_blocks"])
447 +def _kv_kernel(
448 + qkvr_ptr,
449 + k_weight_ptr,
450 + v_weight_ptr,
451 + k_norm_weight_ptr,
452 + conv_cache_ptr,
453 + key_cache_ptr,
454 + value_cache_ptr,
455 + positions_ptr,
456 + seq_idx_ptr,
457 + conv_slot_mapping_ptr,
458 + conv_block_table_ptr,
459 + query_start_ptr,
460 + attention_slot_mapping_ptr,
461 + tokens,
462 + eps,
463 + stride_x_t,
464 + stride_cc_block,
465 + stride_cc_head,
466 + stride_cc_token,
467 + stride_cc_dim,
468 + stride_kc_block,
469 + stride_kc_token,
470 + stride_kc_head,
471 + stride_vc_block,
472 + stride_vc_token,
473 + stride_vc_head,
474 + stride_block_table_req,
475 + max_blocks,
476 + conv_block_size,
477 + attention_page_size,
478 + Q_WIDTH: tl.constexpr,
479 + KV_WIDTH: tl.constexpr,
480 + NUM_KV_HEADS: tl.constexpr,
481 + HEAD_DIM: tl.constexpr,
482 + WINDOW_SIZE: tl.constexpr,
483 + OFF_K: tl.constexpr,
484 + OFF_V: tl.constexpr,
485 + BLOCK_ROWS: tl.constexpr,
486 +):
487 + token = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
488 + dims = tl.arange(0, HEAD_DIM)
489 + row_mask = token < tokens
490 + head_id = tl.program_id(1)
491 + head = tl.full([BLOCK_ROWS], head_id, tl.int64)
492 + position = tl.load(positions_ptr + token, mask=row_mask, other=0)
493 + request = tl.load(seq_idx_ptr + token, mask=row_mask, other=0)
494 + conv_slot = tl.load(conv_slot_mapping_ptr + token, mask=row_mask, other=-1)
495 + query_start = tl.load(query_start_ptr + token, mask=row_mask, other=0)
496 + attention_slot = tl.load(
497 + attention_slot_mapping_ptr + token, mask=row_mask, other=-1
498 + )
499 + valid = row_mask & (conv_slot >= 0)
500 +
501 + k_col = Q_WIDTH + head * HEAD_DIM
502 + v_col = Q_WIDTH + KV_WIDTH + head * HEAD_DIM
503 + k_value = tl.load(
504 + qkvr_ptr + token[:, None] * stride_x_t + k_col[:, None] + dims[None, :],
505 + mask=row_mask[:, None],
506 + other=0.0,
507 + )
508 + v_value = tl.load(
509 + qkvr_ptr + token[:, None] * stride_x_t + v_col[:, None] + dims[None, :],
510 + mask=row_mask[:, None],
511 + other=0.0,
512 + )
513 + safe_slot = tl.maximum(conv_slot, 0)
514 + cache_base = (
515 + conv_cache_ptr
516 + + (safe_slot // conv_block_size) * stride_cc_block
517 + + head * stride_cc_head
518 + + (safe_slot % conv_block_size) * stride_cc_token
519 + )
520 + tl.store(
521 + cache_base[:, None] + (OFF_K + dims[None, :]) * stride_cc_dim,
522 + k_value,
523 + mask=valid[:, None],
524 + )
525 + tl.store(
526 + cache_base[:, None] + (OFF_V + dims[None, :]) * stride_cc_dim,
527 + v_value,
528 + mask=valid[:, None],
529 + )
530 +
531 + acc_k = tl.zeros([BLOCK_ROWS, HEAD_DIM], dtype=tl.float32)
532 + acc_v = tl.zeros([BLOCK_ROWS, HEAD_DIM], dtype=tl.float32)
533 + for tap in tl.static_range(WINDOW_SIZE):
534 + source_position = position - (WINDOW_SIZE - 1) + tap
535 + source_row = token - (WINDOW_SIZE - 1) + tap
536 + in_window = valid & (source_position >= 0)
537 + intra = in_window & (source_row >= query_start)
538 + cached = in_window & (source_row < query_start)
539 + safe_row = tl.maximum(source_row, 0)
540 + source_k = tl.load(
541 + qkvr_ptr + safe_row[:, None] * stride_x_t + k_col[:, None] + dims[None, :],
542 + mask=intra[:, None],
543 + other=0.0,
544 + ).to(tl.float32)
545 + source_v = tl.load(
546 + qkvr_ptr + safe_row[:, None] * stride_x_t + v_col[:, None] + dims[None, :],
547 + mask=intra[:, None],
548 + other=0.0,
549 + ).to(tl.float32)
550 + safe_position = tl.maximum(source_position, 0)
551 + logical_block = tl.minimum(safe_position // conv_block_size, max_blocks - 1)
552 + physical_block = tl.load(
553 + conv_block_table_ptr + request * stride_block_table_req + logical_block,
554 + mask=cached,
555 + other=0,
556 + ).to(tl.int64)
557 + tap_base = (
558 + conv_cache_ptr
559 + + physical_block * stride_cc_block
560 + + head * stride_cc_head
561 + + (safe_position % conv_block_size) * stride_cc_token
562 + )
563 + source_k += tl.load(
564 + tap_base[:, None] + (OFF_K + dims[None, :]) * stride_cc_dim,
565 + mask=cached[:, None],
566 + other=0.0,
567 + ).to(tl.float32)
568 + source_v += tl.load(
569 + tap_base[:, None] + (OFF_V + dims[None, :]) * stride_cc_dim,
570 + mask=cached[:, None],
571 + other=0.0,
572 + ).to(tl.float32)
573 + k_weight = tl.load(
574 + k_weight_ptr + (head_id * HEAD_DIM + dims) * WINDOW_SIZE + tap
575 + ).to(tl.float32)
576 + v_weight = tl.load(
577 + v_weight_ptr + (head_id * HEAD_DIM + dims) * WINDOW_SIZE + tap
578 + ).to(tl.float32)
579 + acc_k += source_k * k_weight[None, :]
580 + acc_v += source_v * v_weight[None, :]
581 +
582 + k_rounded = (acc_k + k_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty)
583 + v_rounded = (acc_v + v_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty)
584 + k_float = k_rounded.to(tl.float32)
585 + k_norm_weight = tl.load(k_norm_weight_ptr + dims).to(tl.float32)
586 + rstd = tl.rsqrt(tl.sum(k_float * k_float, axis=1) / HEAD_DIM + eps)
587 + k_normalized = (k_float * rstd[:, None] * k_norm_weight[None, :]).to(
588 + qkvr_ptr.dtype.element_ty
589 + )
590 +
591 + safe_attention_slot = tl.maximum(attention_slot, 0)
592 + attention_block = safe_attention_slot // attention_page_size
593 + attention_offset = safe_attention_slot % attention_page_size
594 + attention_mask = row_mask & (attention_slot >= 0)
595 + tl.store(
596 + key_cache_ptr
597 + + attention_block[:, None] * stride_kc_block
598 + + attention_offset[:, None] * stride_kc_token
599 + + head[:, None] * stride_kc_head
600 + + dims[None, :],
601 + k_normalized,
602 + mask=attention_mask[:, None],
603 + )
604 + tl.store(
605 + value_cache_ptr
606 + + attention_block[:, None] * stride_vc_block
607 + + attention_offset[:, None] * stride_vc_token
608 + + head[:, None] * stride_vc_head
609 + + dims[None, :],
610 + v_rounded,
611 + mask=attention_mask[:, None],
612 + )
613 +
614 +
615 +def _run_tiled_q(
616 + qkvr: torch.Tensor,
617 + q_norm_weight: torch.Tensor,
618 + q_out: torch.Tensor,
619 + positions: torch.Tensor,
620 + *,
621 + eps: float,
622 + num_q_heads: int,
623 + head_dim: int,
624 + log_scaling: torch.Tensor | None,
625 +) -> None:
626 + num_rows = qkvr.shape[0] * num_q_heads
627 + _q_kernel[(triton.cdiv(num_rows, Q_BLOCK_ROWS),)](
628 + qkvr,
629 + q_norm_weight,
630 + q_out,
631 + log_scaling if log_scaling is not None else positions,
632 + num_rows,
633 + qkvr.stride(0),
634 + eps,
635 + NUM_Q_HEADS=num_q_heads,
636 + HEAD_DIM=head_dim,
637 + BLOCK_ROWS=Q_BLOCK_ROWS,
638 + APPLY_LOG_SCALING=log_scaling is not None,
639 + num_warps=Q_NUM_WARPS,
640 + )
641 +
642 +
643 +def _run_tiled_kv(
644 + qkvr: torch.Tensor,
645 + k_weight: torch.Tensor,
646 + v_weight: torch.Tensor,
647 + k_norm_weight: torch.Tensor,
648 + conv_cache: torch.Tensor,
649 + key_cache: torch.Tensor,
650 + value_cache: torch.Tensor,
651 + positions: torch.Tensor,
652 + seq_idx: torch.Tensor,
653 + conv_slot_mapping: torch.Tensor,
654 + conv_block_table: torch.Tensor,
655 + query_start: torch.Tensor,
656 + attention_slot_mapping: torch.Tensor,
657 + *,
658 + eps: float,
659 + num_q_heads: int,
660 + num_kv_heads: int,
661 + head_dim: int,
662 + off_k: int,
663 + off_v: int,
664 + conv_block_size: int,
665 +) -> None:
666 + tokens = qkvr.shape[0]
667 + _kv_kernel[(triton.cdiv(tokens, KV_BLOCK_ROWS), num_kv_heads)](
668 + qkvr,
669 + k_weight,
670 + v_weight,
671 + k_norm_weight,
672 + conv_cache,
673 + key_cache,
674 + value_cache,
675 + positions,
676 + seq_idx,
677 + conv_slot_mapping,
678 + conv_block_table,
679 + query_start,
680 + attention_slot_mapping,
681 + tokens,
682 + eps,
683 + qkvr.stride(0),
684 + conv_cache.stride(0),
685 + conv_cache.stride(1),
686 + conv_cache.stride(2),
687 + conv_cache.stride(3),
688 + key_cache.stride(0),
689 + key_cache.stride(1),
690 + key_cache.stride(2),
691 + value_cache.stride(0),
692 + value_cache.stride(1),
693 + value_cache.stride(2),
694 + conv_block_table.stride(0),
695 + conv_block_table.shape[1],
696 + conv_block_size,
697 + key_cache.shape[1],
698 + Q_WIDTH=num_q_heads * head_dim,
699 + KV_WIDTH=num_kv_heads * head_dim,
700 + NUM_KV_HEADS=num_kv_heads,
701 + HEAD_DIM=head_dim,
702 + WINDOW_SIZE=k_weight.shape[1],
703 + OFF_K=off_k,
704 + OFF_V=off_v,
705 + BLOCK_ROWS=KV_BLOCK_ROWS,
706 + num_warps=KV_NUM_WARPS,
707 + )
708 +
709 +
710 +def _run_fused_small(
711 + qkvr: torch.Tensor,
712 + q_norm_weight: torch.Tensor,
713 + q_out: torch.Tensor,
714 + rel_proj: torch.Tensor,
715 + rel_out: torch.Tensor,
716 + k_weight: torch.Tensor,
717 + v_weight: torch.Tensor,
718 + k_norm_weight: torch.Tensor,
719 + conv_cache: torch.Tensor,
720 + key_cache: torch.Tensor,
721 + value_cache: torch.Tensor,
722 + positions: torch.Tensor,
723 + seq_idx: torch.Tensor,
724 + conv_slot_mapping: torch.Tensor,
725 + conv_block_table: torch.Tensor,
726 + query_start: torch.Tensor,
727 + attention_slot_mapping: torch.Tensor,
728 + *,
729 + eps: float,
730 + num_q_heads: int,
731 + num_kv_heads: int,
732 + head_dim: int,
733 + off_k: int,
734 + off_v: int,
735 + conv_block_size: int,
736 + log_scaling: torch.Tensor | None,
737 +) -> None:
738 + tokens = qkvr.shape[0]
739 +
740 + num_q_rows = tokens * num_q_heads
741 + grid = (num_q_rows + tokens * num_kv_heads,)
742 + _qkvr_qkv_kernel[grid](
743 + qkvr,
744 + q_norm_weight,
745 + q_out,
746 + rel_proj,
747 + rel_out,
748 + k_weight,
749 + v_weight,
750 + k_norm_weight,
751 + conv_cache,
752 + key_cache,
753 + value_cache,
754 + positions,
755 + seq_idx,
756 + conv_slot_mapping,
757 + conv_block_table,
758 + query_start,
759 + attention_slot_mapping,
760 + log_scaling if log_scaling is not None else positions,
761 + tokens,
762 + eps,
763 + qkvr.stride(0),
764 + conv_cache.stride(0),
765 + conv_cache.stride(1),
766 + conv_cache.stride(2),
767 + conv_cache.stride(3),
768 + key_cache.stride(0),
769 + key_cache.stride(1),
770 + key_cache.stride(2),
771 + value_cache.stride(0),
772 + value_cache.stride(1),
773 + value_cache.stride(2),
774 + conv_block_table.stride(0),
775 + conv_block_table.shape[1],
776 + conv_block_size,
777 + key_cache.shape[1],
778 + Q_WIDTH=num_q_heads * head_dim,
779 + KV_WIDTH=num_kv_heads * head_dim,
780 + NUM_Q_HEADS=num_q_heads,
781 + NUM_KV_HEADS=num_kv_heads,
782 + HEAD_DIM=head_dim,
783 + WINDOW_SIZE=k_weight.shape[1],
784 + OFF_K=off_k,
785 + OFF_V=off_v,
786 + APPLY_LOG_SCALING=log_scaling is not None,
787 + D_REL=16,
788 + REL_EXTENT=rel_proj.shape[1],
789 + REL_EXTENT_PADDED=triton.next_power_of_2(rel_proj.shape[1]),
790 + num_warps=SMALL_NUM_WARPS,
791 + )
792 +
793 +
794 +def fused_qkvr_prep(
795 + qkvr: torch.Tensor,
796 + k_weight: torch.Tensor,
797 + v_weight: torch.Tensor,
798 + q_norm_weight: torch.Tensor,
799 + k_norm_weight: torch.Tensor,
800 + rel_proj: torch.Tensor,
801 + eps: float,
802 + num_q_heads: int,
803 + num_kv_heads: int,
804 + head_dim: int,
805 + d_rel: int,
806 + conv_cache: torch.Tensor,
807 + key_cache: torch.Tensor,
808 + value_cache: torch.Tensor,
809 + positions: torch.Tensor,
810 + conv_block_table: torch.Tensor,
811 + seq_idx: torch.Tensor,
812 + conv_slot_mapping: torch.Tensor,
813 + query_start: torch.Tensor,
814 + attention_slot_mapping: torch.Tensor,
815 + off_k: int,
816 + off_v: int,
817 + conv_block_size: int,
818 + log_scaling: torch.Tensor | None = None,
819 +) -> tuple[torch.Tensor, torch.Tensor]:
820 + assert d_rel == 16 and rel_proj.shape[0] == 16
821 + assert head_dim == 128
822 + assert qkvr.is_contiguous()
823 + assert k_weight.stride() == (k_weight.shape[1], 1)
824 + assert v_weight.stride() == (v_weight.shape[1], 1)
825 + assert rel_proj.stride() == (rel_proj.shape[1], 1)
826 + assert conv_cache.stride(3) == 1
827 + assert key_cache.stride(3) == 1 and value_cache.stride(3) == 1
828 + tokens = qkvr.shape[0]
829 + q_out = torch.empty(
830 + (tokens, num_q_heads * head_dim), dtype=qkvr.dtype, device=qkvr.device
831 + )
832 + rel_out = torch.empty(
833 + (tokens, num_q_heads, rel_proj.shape[1]),
834 + dtype=qkvr.dtype,
835 + device=qkvr.device,
836 + )
837 + if tokens == 0:
838 + return q_out, rel_out
839 +
840 + if tokens < SMALL_TOKEN_THRESHOLD:
841 + _run_fused_small(
842 + qkvr,
843 + q_norm_weight,
844 + q_out,
845 + rel_proj,
846 + rel_out,
847 + k_weight,
848 + v_weight,
849 + k_norm_weight,
850 + conv_cache,
851 + key_cache,
852 + value_cache,
853 + positions,
854 + seq_idx,
855 + conv_slot_mapping,
856 + conv_block_table,
857 + query_start,
858 + attention_slot_mapping,
859 + eps=eps,
860 + num_q_heads=num_q_heads,
861 + num_kv_heads=num_kv_heads,
862 + head_dim=head_dim,
863 + off_k=off_k,
864 + off_v=off_v,
865 + conv_block_size=conv_block_size,
866 + log_scaling=log_scaling,
867 + )
868 + return q_out, rel_out
869 +
870 + kv_stream = aux_stream()
871 + assert kv_stream is not None
872 + current_stream = torch.cuda.current_stream()
873 + kv_stream.wait_stream(current_stream)
874 + with torch.cuda.stream(kv_stream):
875 + _run_tiled_kv(
876 + qkvr,
877 + k_weight,
878 + v_weight,
879 + k_norm_weight,
880 + conv_cache,
881 + key_cache,
882 + value_cache,
883 + positions,
884 + seq_idx,
885 + conv_slot_mapping,
886 + conv_block_table,
887 + query_start,
888 + attention_slot_mapping,
889 + eps=eps,
890 + num_q_heads=num_q_heads,
891 + num_kv_heads=num_kv_heads,
892 + head_dim=head_dim,
893 + off_k=off_k,
894 + off_v=off_v,
895 + conv_block_size=conv_block_size,
896 + )
897 + _run_tiled_q(
898 + qkvr,
899 + q_norm_weight,
900 + q_out,
901 + positions,
902 + eps=eps,
903 + num_q_heads=num_q_heads,
904 + head_dim=head_dim,
905 + log_scaling=log_scaling,
906 + )
907 + qkvr_rel_proj(
908 + qkvr,
909 + rel_proj,
910 + rel_out,
911 + log_scaling,
912 + num_q_heads=num_q_heads,
913 + num_kv_heads=num_kv_heads,
914 + head_dim=head_dim,
915 + d_rel=d_rel,
916 + )
917 + current_stream.wait_stream(kv_stream)
918 + return q_out, rel_out
vllm/models/inkling/nvidia/ops/sconv.py new
+290
@@ -0,0 +1,290 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inkling short-convolution kernels backed by a paged sliding-window state cache.
4 +
5 +Each layer's 4 conv streams (K, V, attn-output, mlp-output) share one paged KV
6 +cache ``[num_blocks, H, N, D]`` (head-major; see ``sconv_swa_attn.py``). A stream
7 +occupies the contiguous D-sub-range ``[off_s, off_s + ws)`` across all ``H``
8 +heads, so its flat per-token width is ``H * ws`` and that is the conv channel
9 +dim. The cache stores the conv *input* at every absolute position.
10 +
11 +``fused_sconv`` is the single-launch path used by the model: per token it writes
12 +the current input to its slot and convolves the ``W`` taps ending at its
13 +absolute position. A tap landing inside the current forward is read from the
14 +immutable input ``x`` (row ``src - pos + pid_t``); only pre-forward taps are read
15 +from the paged cache (window position ``src`` -> physical block via
16 +``block_table[req, src // N]``). The just-written slot is never read back this
17 +step, so there is no write/read hazard within or across programs -- which is
18 +why this needs no decode-vs-prefill split and is valid for prefill / decode /
19 +spec alike.
20 +
21 +All kernels address the cache purely by ``(slot, absolute_position)`` and
22 +allocate nothing inside the captured region; their grids depend only on the
23 +token count (``fused_sconv`` on a fixed token/channel tiling), so the same
24 +decode path replays correctly under a full CUDA graph without any
25 +data-dependent shape or branch.
26 +"""
27 +
28 +from __future__ import annotations
29 +
30 +import torch
31 +
32 +from vllm.triton_utils import tl, triton
33 +
34 +
35 +@triton.jit
36 +def _fused_sconv_kernel(
37 + x_ptr, # [T, H*WS] head-major current-token inputs (also residual)
38 + cache_ptr, # [num_blocks, H, N, D] paged (page-strided view)
39 + weight_ptr, # [H*WS, W]
40 + out_ptr, # [T, H*WS]
41 + pos_ptr, # [T] int64 absolute position per token
42 + seq_idx_ptr, # [T] int32 token -> batch request
43 + slot_ptr, # [T] int64 flat slot (block*N + blk_off); < 0 => PAD
44 + block_table_ptr, # [num_reqs, max_blocks] int32 block_table
45 + qstart_ptr, # [T] int32 first x-row of the token's request
46 + T, # num tokens
47 + stride_x_t,
48 + stride_c_blk,
49 + stride_c_h,
50 + stride_c_n,
51 + stride_c_d,
52 + stride_w_d,
53 + stride_w_w,
54 + stride_bt_r,
55 + MAX_BLOCKS,
56 + N, # block_size
57 + W: tl.constexpr,
58 + USE_SILU: tl.constexpr,
59 + USE_RESIDUAL: tl.constexpr,
60 + OFF_S: tl.constexpr,
61 + WS: tl.constexpr,
62 + H: tl.constexpr,
63 + BT: tl.constexpr,
64 + BLOCK_C: tl.constexpr,
65 +):
66 + # Each program owns a [BT tokens, BLOCK_C channels] tile. The flat channel
67 + # index packs all H heads head-major (head = c // WS, in-stream offset =
68 + # c % WS), so one program spans heads -- no per-head launch and no
69 + # next_power_of_2(WS) lane waste.
70 + pid_t = tl.program_id(0)
71 + pid_c = tl.program_id(1)
72 + toff = pid_t * BT + tl.arange(0, BT) # [BT] token rows
73 + coff = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) # [BLOCK_C] flat channels
74 + C = H * WS
75 + t_mask = toff < T
76 + c_mask = coff < C
77 + head = tl.minimum(coff // WS, H - 1) # clamp keeps masked lanes in-buffer
78 + cd = OFF_S + coff % WS # cache D-index of the channel's stream slot
79 +
80 + slot = tl.load(slot_ptr + toff, mask=t_mask, other=-1) # [BT]
81 + valid = slot >= 0
82 + pos = tl.load(pos_ptr + toff, mask=t_mask, other=0)
83 + req = tl.load(seq_idx_ptr + toff, mask=t_mask, other=0)
84 + qstart = tl.load(qstart_ptr + toff, mask=t_mask, other=0)
85 +
86 + tc_mask = t_mask[:, None] & c_mask[None, :]
87 +
88 + # 1) Insert each token's input into its paged slot (skip PAD rows).
89 + xv = tl.load(x_ptr + toff[:, None] * stride_x_t + coff[None, :], mask=tc_mask)
90 + safe_slot = tl.maximum(slot, 0)
91 + dst = (
92 + cache_ptr
93 + + (safe_slot // N)[:, None] * stride_c_blk
94 + + head[None, :] * stride_c_h
95 + + (safe_slot % N)[:, None] * stride_c_n
96 + + cd[None, :] * stride_c_d
97 + )
98 + tl.store(dst, xv, mask=tc_mask & valid[:, None])
99 +
100 + # 2) Convolve the W taps ending at each token's `pos`. Each tap is read from
101 + # `x` if it falls inside this forward (row >= the request's first row), else
102 + # from the paged cache. Exactly one source is unmasked per tap, so we sum both.
103 + acc = tl.zeros([BT, BLOCK_C], dtype=tl.float32)
104 + for iw in tl.static_range(W):
105 + src = pos - (W - 1) + iw # [BT] absolute window position
106 + row = toff - (W - 1) + iw # [BT] x-row of `src` (== src - pos + token)
107 + in_win = valid & (src >= 0)
108 + intra = in_win & (row >= qstart)
109 + cached = in_win & (row < qstart)
110 + # intra-forward tap: read the immutable input x (never the slot we just
111 + # wrote), so there is no write/read hazard.
112 + safe_row = tl.maximum(row, 0)
113 + xt = tl.load(
114 + x_ptr + safe_row[:, None] * stride_x_t + coff[None, :],
115 + mask=c_mask[None, :] & intra[:, None],
116 + other=0.0,
117 + ).to(tl.float32)
118 + # pre-forward tap: read from the paged cache via the block table. Clamp
119 + # addressing terms; the load is masked off when out of window.
120 + safe_src = tl.maximum(src, 0)
121 + safe_lblk = tl.minimum(safe_src // N, MAX_BLOCKS - 1)
122 + blk = tl.load(
123 + block_table_ptr + req * stride_bt_r + safe_lblk, mask=cached, other=0
124 + ).to(tl.int64)
125 + cbase = (
126 + cache_ptr
127 + + blk[:, None] * stride_c_blk
128 + + head[None, :] * stride_c_h
129 + + (safe_src % N)[:, None] * stride_c_n
130 + + cd[None, :] * stride_c_d
131 + )
132 + cv = tl.load(cbase, mask=c_mask[None, :] & cached[:, None], other=0.0).to(
133 + tl.float32
134 + )
135 + wv = tl.load(
136 + weight_ptr + coff * stride_w_d + iw * stride_w_w, mask=c_mask, other=0.0
137 + ).to(tl.float32)
138 + acc += (xt + cv) * wv[None, :]
139 +
140 + if USE_SILU:
141 + acc = acc * tl.sigmoid(acc)
142 + if USE_RESIDUAL:
143 + acc += xv.to(tl.float32)
144 +
145 + tl.store(
146 + out_ptr + toff[:, None] * stride_x_t + coff[None, :],
147 + acc.to(out_ptr.dtype.element_ty),
148 + mask=tc_mask,
149 + )
150 +
151 +
152 +def fused_sconv(
153 + x: torch.Tensor, # [T, H*ws] head-major current-token inputs
154 + weight: torch.Tensor, # [H*ws, W]
155 + cache: torch.Tensor, # [num_blocks, H, N, D] paged
156 + positions: torch.Tensor, # [T] int64 absolute position per token
157 + block_table: torch.Tensor, # [num_reqs, max_blocks] int32
158 + seq_idx: torch.Tensor, # [T] int32 token -> batch request
159 + slot_mapping: torch.Tensor, # [T] int64 flat slot (PAD = -1 => skip)
160 + query_start: torch.Tensor, # [T] int32 first x-row of the token's request
161 + off_s: int,
162 + ws: int,
163 + block_size: int,
164 + activation: str | None = None,
165 + use_residual: bool = True,
166 +) -> torch.Tensor:
167 + """Single-launch insert + depthwise causal conv1d over the paged cache.
168 +
169 + Reads same-forward taps from ``x`` and pre-forward taps from the cache, so
170 + it is race-free in one launch for prefill / decode / spec and supports full
171 + CUDA-graph capture for decode.
172 + """
173 + T = x.shape[0]
174 + out = torch.empty_like(x)
175 + if T == 0:
176 + return out
177 + assert x.is_contiguous()
178 + assert cache.stride(3) == 1, "cache D-dim must be contiguous"
179 + H = cache.shape[1]
180 + W = weight.shape[1]
181 + C = H * ws # flat conv channel dim (all heads, head-major)
182 + # Tile BT tokens x BLOCK_C channels per program: enough work per CTA to
183 + # amortize the per-token addressing, while keeping the grid large for
184 + # prefill. BLOCK_C spans heads so there is no per-head launch. A ~2K-element
185 + # tile at 4 warps measured best on Blackwell; larger tiles spill registers.
186 + BLOCK_C = min(triton.next_power_of_2(C), 256)
187 + BT = 8
188 + grid = (triton.cdiv(T, BT), triton.cdiv(C, BLOCK_C))
189 + _fused_sconv_kernel[grid](
190 + x,
191 + cache,
192 + weight,
193 + out,
194 + positions,
195 + seq_idx,
196 + slot_mapping,
197 + block_table,
198 + query_start,
199 + T,
200 + x.stride(0),
201 + cache.stride(0),
202 + cache.stride(1),
203 + cache.stride(2),
204 + cache.stride(3),
205 + weight.stride(0),
206 + weight.stride(1),
207 + block_table.stride(0),
208 + block_table.shape[1],
209 + block_size,
210 + W=W,
211 + USE_SILU=activation in ("silu", "swish"),
212 + USE_RESIDUAL=use_residual,
213 + OFF_S=off_s,
214 + WS=ws,
215 + H=H,
216 + BT=BT,
217 + BLOCK_C=BLOCK_C,
218 + num_warps=4,
219 + )
220 + return out
221 +
222 +
223 +@triton.jit
224 +def _seq_metadata_kernel(
225 + qsl_ptr, # [num_reqs + 1] int32 cumulative query start rows
226 + seq_idx_ptr, # [T] int32 out: token -> owning request
227 + query_start_ptr, # [T] int32 out: first x-row of the token's request
228 + num_reqs,
229 + num_actual_tokens,
230 + num_padded_tokens,
231 + n_iters, # ceil(log2(num_reqs)): binary-search depth
232 + BLOCK: tl.constexpr,
233 +):
234 + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
235 + tok = offs.to(tl.int32)
236 + # Largest j in [0, num_reqs) with qsl[j] <= tok.
237 + lo = tl.zeros([BLOCK], tl.int32)
238 + hi = tl.full([BLOCK], num_reqs - 1, tl.int32)
239 + for _ in range(n_iters):
240 + mid = (lo + hi + 1) // 2
241 + below = tl.load(qsl_ptr + mid) <= tok
242 + lo = tl.where(below, mid, lo)
243 + hi = tl.where(below, hi, mid - 1)
244 + actual = offs < num_actual_tokens
245 + padded = offs < num_padded_tokens
246 + query_start = tl.load(qsl_ptr + lo)
247 + tl.store(seq_idx_ptr + offs, tl.where(actual, lo, 0), mask=padded)
248 + tl.store(
249 + query_start_ptr + offs,
250 + tl.where(actual, query_start, 0),
251 + mask=padded,
252 + )
253 +
254 +
255 +def sconv_seq_metadata(
256 + query_start_loc: torch.Tensor,
257 + num_reqs: int,
258 + num_actual_tokens: int,
259 + seq_idx_out: torch.Tensor,
260 + query_start_out: torch.Tensor,
261 + num_padded_tokens: int | None = None,
262 +) -> None:
263 + """Fill static per-token seq_idx / query_start buffers in one launch.
264 +
265 + Replaces the arange + searchsorted + clamp + gather + 2x copy chain of the
266 + sconv metadata build with a single kernel writing both persistent buffers.
267 + Padded rows are filled with zero and must have ``slot_mapping == -1``.
268 + """
269 + if num_padded_tokens is None:
270 + num_padded_tokens = num_actual_tokens
271 + if num_padded_tokens < num_actual_tokens:
272 + raise ValueError("num_padded_tokens must cover all actual tokens")
273 + if num_padded_tokens > seq_idx_out.shape[0]:
274 + raise ValueError("seq_idx_out is too small for the padded token count")
275 + if num_padded_tokens > query_start_out.shape[0]:
276 + raise ValueError("query_start_out is too small for the padded token count")
277 +
278 + BLOCK = 256
279 + n_iters = (num_reqs - 1).bit_length()
280 + grid = (triton.cdiv(num_padded_tokens, BLOCK),)
281 + _seq_metadata_kernel[grid](
282 + query_start_loc,
283 + seq_idx_out,
284 + query_start_out,
285 + num_reqs,
286 + num_actual_tokens,
287 + num_padded_tokens,
288 + n_iters,
289 + BLOCK=BLOCK,
290 + )
vllm/models/inkling/nvidia/ops/silu_and_mul.py new
+197
@@ -0,0 +1,197 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""SwiGLU kernels for the Inkling MLP layers.
4 +
5 +``silu_and_mul_triton``: SiLU-and-mul over the checkpoint's interleaved
6 +fused gate/up layout (dense MLP). ``sink_silu_mul_epilogue``: the sink-expert
7 +variant with the per-expert dequant scale and per-token gamma fused in.
8 +"""
9 +
10 +import torch
11 +
12 +from vllm.triton_utils import tl, triton
13 +
14 +
15 +@triton.jit(do_not_specialize=["M"])
16 +def _silu_and_mul_triton_kernel(
17 + gateup_out_ptr,
18 + down_inp_ptr,
19 + M,
20 + N: tl.constexpr,
21 + GRID_SIZE: tl.constexpr,
22 + NUM_STAGES: tl.constexpr,
23 + BLOCK_SIZE_M: tl.constexpr,
24 + BLOCK_SIZE_N: tl.constexpr,
25 + EVEN_N: tl.constexpr,
26 + INT64_INDEX: tl.constexpr,
27 +):
28 + start_pid = tl.program_id(0)
29 + if INT64_INDEX:
30 + start_pid = start_pid.to(tl.int64)
31 + M = M.to(tl.int64)
32 +
33 + NUM_BLOCKS_N: tl.constexpr = tl.cdiv(N, BLOCK_SIZE_N)
34 + num_blocks_mn = tl.cdiv(M, BLOCK_SIZE_M) * NUM_BLOCKS_N
35 +
36 + for pid in tl.range(start_pid, num_blocks_mn, GRID_SIZE, num_stages=NUM_STAGES):
37 + pid_m = pid // NUM_BLOCKS_N
38 + pid_n = pid % NUM_BLOCKS_N
39 +
40 + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
41 + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
42 + mask_m = offs_m < M
43 + mask_n = offs_n < N
44 +
45 + # Interleaved fused gate/up: [g0, u0, g1, u1, ...].
46 + mask_offs_2n = pid_n * BLOCK_SIZE_N + tl.arange(0, 2 * BLOCK_SIZE_N) // 2
47 + tl.static_assert(BLOCK_SIZE_N % 8 == 0, f"{BLOCK_SIZE_N=}")
48 + mask_2n = mask_offs_2n < N
49 + mask_2n = tl.max_constancy(mask_2n, [16])
50 +
51 + offs_2n = pid_n * 2 * BLOCK_SIZE_N + tl.arange(0, 2 * BLOCK_SIZE_N)
52 + offs_m2n = offs_m[:, None] * N * 2 + offs_2n[None, :]
53 +
54 + if EVEN_N or pid_n * BLOCK_SIZE_N + BLOCK_SIZE_N <= N:
55 + gateup_out = tl.load(
56 + gateup_out_ptr + offs_m2n, mask=mask_m[:, None], other=0.0
57 + )
58 + else:
59 + mask_m2n = mask_m[:, None] & mask_2n[None, :]
60 + gateup_out = tl.load(gateup_out_ptr + offs_m2n, mask=mask_m2n, other=0.0)
61 +
62 + gate_out, up_out = tl.split(
63 + tl.reshape(gateup_out, (BLOCK_SIZE_M, BLOCK_SIZE_N, 2))
64 + )
65 + gate_out = gate_out.to(tl.float32)
66 + up_out = up_out.to(tl.float32)
67 +
68 + down_inp = gate_out * tl.sigmoid(gate_out) * up_out
69 +
70 + mask_mn = mask_m[:, None] if EVEN_N else mask_m[:, None] & mask_n[None, :]
71 + offs_mn = offs_m[:, None] * N + offs_n[None, :]
72 + tl.store(down_inp_ptr + offs_mn, down_inp, mask=mask_mn)
73 +
74 +
75 +def silu_and_mul_triton(gateup_output: torch.Tensor) -> torch.Tensor:
76 + """SiLU-and-mul for the interleaved fused gate/up layout.
77 +
78 + Adapted from ``inkling_kernels.activation.silu_and_mul_fwd`` (without MXFP).
79 + """
80 + assert gateup_output.is_contiguous(), (
81 + f"{gateup_output.shape=} {gateup_output.stride()=}"
82 + )
83 + assert gateup_output.ndim == 2, f"{gateup_output.shape=}"
84 +
85 + M = gateup_output.shape[0]
86 + hidden_size = gateup_output.shape[1]
87 + assert hidden_size % 2 == 0, f"{hidden_size=}"
88 + N = hidden_size // 2
89 +
90 + down_input = torch.empty(
91 + (M, N), device=gateup_output.device, dtype=gateup_output.dtype
92 + )
93 + if M == 0:
94 + return down_input
95 +
96 + BLOCK_SIZE_N = max(8, min(256, triton.next_power_of_2(N)))
97 + if M <= 1:
98 + BLOCK_SIZE_M = 4
99 + elif M <= 256:
100 + BLOCK_SIZE_M = 2
101 + elif M < 4096:
102 + BLOCK_SIZE_M = 4
103 + else:
104 + BLOCK_SIZE_M = 16
105 + BLOCK_SIZE_N = max(8, min(128, triton.next_power_of_2(N)))
106 + max_grid_size = triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N)
107 + num_sms = torch.cuda.get_device_properties(
108 + gateup_output.device
109 + ).multi_processor_count
110 + grid_size = min(num_sms * 4, max_grid_size)
111 +
112 + _silu_and_mul_triton_kernel[(grid_size,)](
113 + gateup_out_ptr=gateup_output,
114 + down_inp_ptr=down_input,
115 + M=M,
116 + N=N,
117 + GRID_SIZE=grid_size,
118 + NUM_STAGES=1,
119 + BLOCK_SIZE_M=BLOCK_SIZE_M,
120 + BLOCK_SIZE_N=BLOCK_SIZE_N,
121 + EVEN_N=N % BLOCK_SIZE_N == 0,
122 + INT64_INDEX=gateup_output.nbytes >= 2**31,
123 + num_warps=8,
124 + )
125 +
126 + return down_input
127 +
128 +
129 +@triton.jit(do_not_specialize=["T"])
130 +def _sink_epilogue_kernel(
131 + raw_ptr, # [T, S * 2F] gemm1 output, interleaved g/u pairs per expert block
132 + alpha_ptr, # [S] fp32 per-expert pre-SiLU dequant scale
133 + gamma_ptr, # [T, S] fp32 per-token sink weights (may be strided)
134 + ratio_ptr, # [S] fp32 per-expert post-SiLU scale (gemm2 alpha ratio)
135 + out_ptr, # [T, S * F] output
136 + T,
137 + stride_raw_0,
138 + stride_gamma_0,
139 + F: tl.constexpr,
140 + S: tl.constexpr,
141 + BLOCK_F: tl.constexpr,
142 +):
143 + pid_t = tl.program_id(0).to(tl.int64)
144 + pid_sf = tl.program_id(1)
145 + if pid_t >= T:
146 + return
147 + s = pid_sf // (F // BLOCK_F)
148 + offs_f = (pid_sf % (F // BLOCK_F)) * BLOCK_F + tl.arange(0, BLOCK_F)
149 +
150 + base = pid_t * stride_raw_0 + s * 2 * F
151 + gate = tl.load(raw_ptr + base + 2 * offs_f).to(tl.float32)
152 + up = tl.load(raw_ptr + base + 2 * offs_f + 1).to(tl.float32)
153 + alpha = tl.load(alpha_ptr + s)
154 + weight = tl.load(gamma_ptr + pid_t * stride_gamma_0 + s) * tl.load(ratio_ptr + s)
155 +
156 + gate *= alpha
157 + up *= alpha
158 + h = gate * tl.sigmoid(gate) * up * weight
159 + tl.store(out_ptr + pid_t * (S * F) + s * F + offs_f, h)
160 +
161 +
162 +def sink_silu_mul_epilogue(
163 + raw: torch.Tensor, # [T, S * 2F] gemm1 output (interleaved gate/up rows)
164 + alphas: torch.Tensor, # [S] fp32
165 + gammas: torch.Tensor, # [T, S] fp32
166 + ratios: torch.Tensor, # [S] fp32
167 + n_experts: int,
168 + out_dtype: torch.dtype,
169 +) -> torch.Tensor:
170 + """Fused sink-expert epilogue: silu(g * a_e) * (u * a_e) * (gamma * r_e).
171 +
172 + One kernel replaces the per-expert dequant column scale, the SwiGLU, and
173 + the per-token gamma multiply between the two sink GEMMs.
174 + """
175 + tokens = raw.shape[0]
176 + f = raw.shape[1] // (2 * n_experts)
177 + out = torch.empty((tokens, n_experts * f), device=raw.device, dtype=out_dtype)
178 + if tokens == 0:
179 + return out
180 + # raw may be a column-slice of a padded GEMM output (rows strided).
181 + assert raw.stride(1) == 1 and gammas.stride(1) == 1
182 + # Largest power-of-two divisor of f (f = 768 -> 256), capped at 512.
183 + block_f = min(512, f & (-f))
184 + _sink_epilogue_kernel[(tokens, n_experts * (f // block_f))](
185 + raw,
186 + alphas,
187 + gammas,
188 + ratios,
189 + out,
190 + tokens,
191 + raw.stride(0),
192 + gammas.stride(0),
193 + F=f,
194 + S=n_experts,
195 + BLOCK_F=block_f,
196 + )
197 + return out
vllm/models/inkling/nvidia/sconv_swa_attn.py new
+221
@@ -0,0 +1,221 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inkling short-conv state managed as a sliding-window KV cache.
4 +
5 +Each decoder layer owns one ``InklingConvState`` (an ``AttentionLayerBase``) that
6 +emits a single ``SlidingWindowSpec`` for the layer's 4 sconv streams (K, V,
7 +attn-output, mlp-output), packed head-major into one block:
8 +
9 + H = num_kv_heads (per-rank), N = block_size = sconv_kernel_size,
10 + D = head_dim(K) + head_dim(V) + hidden/H(attn) + hidden/H(mlp)
11 +
12 +``D`` is TP-invariant; per rank we store ``H/TP`` heads of width ``D``. The conv
13 +reads/writes this cache out-of-band via a custom backend; the (smaller) conv page
14 +is padded up to the uniform attention page by ``unify_kv_cache_spec_page_size``.
15 +"""
16 +
17 +from __future__ import annotations
18 +
19 +from dataclasses import dataclass
20 +from typing import ClassVar
21 +
22 +import torch
23 +from torch import nn
24 +
25 +from vllm.config import VllmConfig, get_current_vllm_config
26 +from vllm.distributed import get_tensor_model_parallel_world_size
27 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
28 +from vllm.v1.attention.backend import (
29 + AttentionBackend,
30 + AttentionCGSupport,
31 + AttentionMetadata,
32 + AttentionMetadataBuilder,
33 + CommonAttentionMetadata,
34 +)
35 +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec, SlidingWindowSpec
36 +
37 +from .ops.sconv import sconv_seq_metadata
38 +
39 +# Stream order within the per-head packed D (== contiguous sub-ranges).
40 +_K, _V, _ATTN, _MLP = 0, 1, 2, 3
41 +
42 +
43 +@dataclass
44 +class InklingSconvMetadata(AttentionMetadata):
45 + block_table: torch.Tensor # [num_reqs, max_blocks] physical blocks per req
46 + slot_mapping: torch.Tensor # [T] int64 flat slot of each token (-1 => skip)
47 + seq_idx: torch.Tensor # [T] int32 token -> batch request
48 + query_start: torch.Tensor # [T] int32 first x-row of each token's request
49 +
50 +
51 +class InklingSconvMetadataBuilder(AttentionMetadataBuilder[InklingSconvMetadata]):
52 + _cudagraph_support: ClassVar[AttentionCGSupport] = (
53 + AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
54 + )
55 +
56 + def __init__(
57 + self,
58 + kv_cache_spec: AttentionSpec,
59 + layer_names: list[str],
60 + vllm_config: VllmConfig,
61 + device: torch.device,
62 + ) -> None:
63 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
64 + assert isinstance(kv_cache_spec, SlidingWindowSpec)
65 + # Persistent per-token buffers for CUDA graph capture.
66 + max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
67 + self.seq_idx_buffer = torch.empty(
68 + max_num_tokens, dtype=torch.int32, device=device
69 + )
70 + self.query_start_buffer = torch.empty(
71 + max_num_tokens, dtype=torch.int32, device=device
72 + )
73 +
74 + def build(
75 + self,
76 + common_prefix_len: int,
77 + common_attn_metadata: CommonAttentionMetadata,
78 + fast_build: bool = False,
79 + ) -> InklingSconvMetadata:
80 + num_reqs = common_attn_metadata.num_reqs
81 + num_actual_tokens = int(common_attn_metadata.query_start_loc_cpu[-1])
82 + num_padded_tokens = common_attn_metadata.slot_mapping.shape[0]
83 + assert num_padded_tokens >= num_actual_tokens
84 +
85 + # Per-token seq_idx (owning request) and query_start (first x-row of
86 + # that request; the fused kernel uses it to tell same-forward taps,
87 + # read from x, from pre-forward taps, read from cache) in one launch.
88 + sconv_seq_metadata(
89 + common_attn_metadata.query_start_loc,
90 + num_reqs,
91 + num_actual_tokens,
92 + self.seq_idx_buffer,
93 + self.query_start_buffer,
94 + num_padded_tokens,
95 + )
96 +
97 + return InklingSconvMetadata(
98 + block_table=common_attn_metadata.block_table_tensor,
99 + slot_mapping=common_attn_metadata.slot_mapping[:num_padded_tokens],
100 + seq_idx=self.seq_idx_buffer[:num_padded_tokens],
101 + query_start=self.query_start_buffer[:num_padded_tokens],
102 + )
103 +
104 +
105 +class InklingSconvBackend(AttentionBackend):
106 + """Custom dummy backend for the sconv sliding-window cache management."""
107 +
108 + @staticmethod
109 + def get_name() -> str:
110 + return "INKLING_SCONV_SWA"
111 +
112 + @classmethod
113 + def indexes_kv_by_block_stride(cls) -> bool:
114 + # num_blocks is the outermost dim (HND, see get_kv_cache_shape), so the
115 + # padded conv page is read through a strided view.
116 + return True
117 +
118 + @staticmethod
119 + def get_kv_cache_shape(
120 + num_blocks: int,
121 + block_size: int,
122 + num_kv_heads: int,
123 + head_size: int,
124 + cache_dtype_str: str = "auto",
125 + ) -> tuple[int, ...]:
126 + # HND, num-blocks-first, head-major: [num_blocks, H, N, D].
127 + return (num_blocks, num_kv_heads, block_size, head_size)
128 +
129 + @staticmethod
130 + def get_kv_cache_stride_order(
131 + include_num_layers_dimension: bool = False,
132 + ) -> tuple[int, ...]:
133 + # Identity: physical layout == logical [num_blocks, H, N, D].
134 + if include_num_layers_dimension:
135 + return (0, 1, 2, 3, 4)
136 + return (0, 1, 2, 3)
137 +
138 + @staticmethod
139 + def get_impl_cls():
140 + raise NotImplementedError(
141 + "InklingSconvBackend has no attention impl; the conv runs out-of-band."
142 + )
143 +
144 + @staticmethod
145 + def get_builder_cls() -> type[InklingSconvMetadataBuilder]:
146 + return InklingSconvMetadataBuilder
147 +
148 +
149 +class InklingConvState(nn.Module, AttentionLayerBase):
150 + """Per-decoder-layer owner emitting one sliding-window conv-state spec."""
151 +
152 + def __init__(
153 + self,
154 + *,
155 + num_kv_heads: int,
156 + head_dim: int,
157 + hidden_size: int,
158 + kernel_size: int,
159 + prefix: str,
160 + ) -> None:
161 + super().__init__()
162 + self.prefix = prefix
163 + # Bound to the manager-allocated paged cache by bind_kv_cache; a
164 + # placeholder until then. Read out-of-band by InklingShortConv.
165 + self.kv_cache = torch.tensor([])
166 + tp_size = get_tensor_model_parallel_world_size()
167 + # Guardrails for the conv-state layout below; only these are exercised.
168 + # tp_size <= num_kv_heads keeps >=1 whole KV head per rank (no
169 + # replication/clamping), so the per-head width stays TP-invariant.
170 + assert tp_size <= num_kv_heads, (
171 + f"sconv SWA cache supports tp_size <= num_kv_heads ({num_kv_heads}), "
172 + f"got {tp_size}"
173 + )
174 + # Per-rank head count; D is TP-invariant (K/V heads and the hidden
175 + # chunk both scale 1/TP together). The attn-/mlp-output sconv streams
176 + # are hidden-sharded: each rank owns its H/tp chunk (the sublayer
177 + # outputs are reduce-scattered / all-gathered around the conv).
178 + self.num_kv_heads = num_kv_heads // tp_size
179 + hidden_per_head = hidden_size // num_kv_heads
180 + # Packed per-head width: K + V + attn-output chunk + mlp-output chunk,
181 + # padded to a power of two so every layer's conv page is the same size
182 + # and an exact multiple of the attention page (the page unifier then
183 + # scales attention block sizes instead of padding).
184 + raw_head_size = 2 * head_dim + 2 * hidden_per_head
185 + self.head_size = 1 << (raw_head_size - 1).bit_length()
186 + self.sliding_window = kernel_size
187 + self.block_size = kernel_size
188 + # Per-head D-sub-range (offset, width) for each stream. Streams share
189 + # the cache; each writes/reads its own width across all H heads.
190 + self.stream_ranges: tuple[tuple[int, int], ...] = (
191 + (0, head_dim), # _K
192 + (head_dim, head_dim), # _V
193 + (2 * head_dim, hidden_per_head), # _ATTN
194 + (2 * head_dim + hidden_per_head, hidden_per_head), # _MLP
195 + )
196 + vllm_config = get_current_vllm_config()
197 + self._dtype = vllm_config.model_config.dtype
198 + assert self._dtype == torch.bfloat16, (
199 + f"sconv SWA cache supports bfloat16 only, got {self._dtype}"
200 + )
201 + # Register in the forward context so the runner enumerates this owner as
202 + # an attention-like layer (get_kv_cache_spec / get_attn_backend).
203 + compilation_config = vllm_config.compilation_config
204 + if prefix in compilation_config.static_forward_context:
205 + raise ValueError(f"Duplicate layer name: {prefix}")
206 + compilation_config.static_forward_context[prefix] = self
207 +
208 + def forward(self): ...
209 +
210 + def get_attn_backend(self) -> type[AttentionBackend]:
211 + return InklingSconvBackend
212 +
213 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
214 + return SlidingWindowSpec(
215 + block_size=self.block_size,
216 + num_kv_heads=self.num_kv_heads,
217 + head_size=self.head_size,
218 + head_size_v=0, # all 4 streams packed into head_size
219 + dtype=self._dtype,
220 + sliding_window=self.sliding_window,
221 + )
vllm/models/inkling/nvidia/short_conv.py new
+93
@@ -0,0 +1,93 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inkling short convolution: depthwise causal conv1d (+ residual) over a paged
4 +sliding-window conv-state cache.
5 +
6 +Each decoder layer owns one ``InklingConvState`` (``sconv_swa_attn.py``) holding the
7 +manager-allocated paged cache for the layer's 4 sconv streams (K, V, attn-output,
8 +mlp-output), packed head-major into one block. Each ``InklingShortConv`` is a
9 +stateless weight + kernel launcher that, per forward (positions-addressed, the
10 +same path for prefill / decode / mixed), inserts the current tokens' inputs
11 +into their paged slot and convolves each token against the ``W`` taps ending
12 +at its absolute position, reading pre-forward window positions out of the
13 +paged cache via the block table.
14 +
15 +Per-forward metadata (``block_table`` / ``slot_mapping`` / ``seq_idx`` /
16 +``query_start``) is built once by ``InklingSconvMetadataBuilder`` and published under
17 +the owner's prefix in the forward context; the absolute ``positions`` are
18 +threaded in from the model. The insert + conv run in a single ``fused_sconv``
19 +launch (same path for prefill / decode / mixed / spec). All inputs are
20 +fixed-address persistent buffers and the grid is fixed, so decode can replay
21 +under a full CUDA graph.
22 +"""
23 +
24 +from __future__ import annotations
25 +
26 +import torch
27 +from torch import nn
28 +from torch.nn.parameter import Parameter
29 +
30 +from vllm.distributed import get_tensor_model_parallel_rank
31 +from vllm.forward_context import get_forward_context
32 +from vllm.model_executor.utils import set_weight_attrs
33 +
34 +from .ops import fused_sconv
35 +from .sconv_swa_attn import InklingConvState, InklingSconvMetadata
36 +
37 +
38 +class InklingShortConv(nn.Module):
39 + def __init__(
40 + self, dim: int, kernel_size: int, owner: InklingConvState, stream_idx: int
41 + ) -> None:
42 + super().__init__()
43 + self.dim = dim
44 + self.kernel_size = kernel_size
45 + self.owner = owner
46 + self.stream_idx = stream_idx
47 + self.tp_rank = get_tensor_model_parallel_rank()
48 +
49 + # Depthwise conv weight; checkpoint stores (dim, 1, W).
50 + self.weight = Parameter(torch.empty(dim, 1, kernel_size), requires_grad=False)
51 + set_weight_attrs(self.weight, {"weight_loader": self.weight_loader})
52 +
53 + def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor) -> None:
54 + if loaded_weight.shape[0] != param.shape[0]:
55 + shard = param.shape[0]
56 + loaded_weight = loaded_weight.narrow(0, self.tp_rank * shard, shard)
57 + param.data.copy_(loaded_weight)
58 +
59 + def forward(self, x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
60 + # x: (num_tokens, dim); positions: (num_tokens,) absolute positions.
61 + attn_metadata = get_forward_context().attn_metadata
62 + if not isinstance(attn_metadata, dict):
63 + # Memory-profiling / no metadata: identity (residual).
64 + return x
65 + m = attn_metadata.get(self.owner.prefix)
66 + if m is None:
67 + return x
68 + assert isinstance(m, InklingSconvMetadata)
69 + cache = self.owner.kv_cache
70 + if cache.numel() == 0:
71 + # Cache not yet bound (profiling before KV alloc): identity.
72 + return x
73 +
74 + off_s, ws = self.owner.stream_ranges[self.stream_idx]
75 + block_size = self.owner.block_size
76 + x = x.contiguous()
77 + weight = self.weight.squeeze(1) # (dim, W)
78 +
79 + return fused_sconv(
80 + x,
81 + weight,
82 + cache,
83 + positions,
84 + m.block_table,
85 + m.seq_idx,
86 + m.slot_mapping,
87 + m.query_start,
88 + off_s,
89 + ws,
90 + block_size,
91 + activation=None,
92 + use_residual=True,
93 + )
vllm/parser/engine/parser_engine.py
+2
@@ -110,6 +110,8 @@ class ParserEngine(Parser):
110
111 self._has_reasoning = (
112 "THINK_END" in parser_engine_config.token_id_terminals
113 + or "THINK_START" in parser_engine_config.terminals
114 + or "THINK_END" in parser_engine_config.terminals
115 or parser_engine_config.initial_state == ParserState.REASONING
116 )
117 self._reasoning_ended: bool = not self._has_reasoning
vllm/parser/engine/parser_engine_config.py
+1
@@ -28,6 +28,7 @@ from vllm.parser.engine.events import EventType
28 class ParserState(Enum):
29 CONTENT = auto()
30 REASONING = auto()
31 + MESSAGE_HEADER = auto()
32 TOOL_PREAMBLE = auto()
33 TOOL_NAME = auto()
34 TOOL_ARGS = auto()
vllm/parser/engine/registered_adapters.py
+6
@@ -12,6 +12,7 @@ from vllm.parser.deepseek_v32 import DeepSeekV32Parser
12 from vllm.parser.engine.adapters import make_adapters
13 from vllm.parser.gemma4 import Gemma4Parser
14 from vllm.parser.glm47_moe import Glm47MoeParser
15 +from vllm.parser.inkling import InklingParser
16 from vllm.parser.kimi_k2 import KimiK2Parser
17 from vllm.parser.minimax_m2 import MinimaxM2Parser
18 from vllm.parser.nemotron_v3 import NemotronV3Parser
@@ -62,3 +63,8 @@ from vllm.parser.seed_oss import SeedOssParser
63 KimiK2ParserReasoningAdapter,
64 KimiK2ParserToolAdapter,
65 ) = make_adapters(KimiK2Parser)
66 +
67 +(
68 + InklingParserReasoningAdapter,
69 + InklingParserToolAdapter,
70 +) = make_adapters(InklingParser)
vllm/parser/engine/streaming_parser_engine.py
+11
@@ -269,6 +269,8 @@ class StreamingParserEngine:
269 SemanticEvent(EventType.REASONING_END, tool_index=self.tool_index)
270 )
271 self.state = ParserState.CONTENT
272 + elif self.state == ParserState.MESSAGE_HEADER:
273 + self.state = ParserState.CONTENT
274
275 return events
276
@@ -314,6 +316,15 @@ class StreamingParserEngine:
316 return self._emit_for_state(value)
317
318 if self.skip_tool_parsing and terminal in self._tool_terminals:
319 + if self.state == ParserState.MESSAGE_HEADER:
320 + self.state = ParserState.CONTENT
321 + return [
322 + SemanticEvent(
323 + EventType.TEXT_CHUNK,
324 + value=value,
325 + tool_index=self.tool_index,
326 + )
327 + ]
328 if EventType.REASONING_END in transition.events:
329 self.state = ParserState.CONTENT
330 return [
vllm/parser/inkling.py new
+411
@@ -0,0 +1,411 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inkling parser: typed content blocks parsed by a single state machine.
4 +
5 +Inkling output format (every marker is a dedicated special token)::
6 +
7 + <|message_model|><|content_thinking|>...reasoning...<|end_message|>
8 + <|message_model|><|content_text|>...visible text...<|end_message|>
9 + <|message_model|><|content_invoke_tool_json|>
10 + {"name":"get_weather","args":{"city":"SF"}}<|end_message|>
11 +
12 +Blocks are self-describing and may repeat in any order; sampling may
13 +also end a block with the standalone ``<|content_model_end_sampling|>``
14 +token. The tool-call payload is a single JSON object whose ``name`` is
15 +extracted by the engine's name-from-args path and whose ``args`` object
16 +is carved out of the wrapper by :func:`_inkling_arg_converter`.
17 +
18 +Note the terminal *labels*: ``THINK_START``/``THINK_END`` are what the
19 +engine keys its reasoning plumbing on (``is_reasoning_end``,
20 +``count_reasoning_tokens``, initial-state seeding), so ``<|end_message|>``
21 +is labelled ``THINK_END`` here even though it ends every block kind —
22 +the transition table, not the label, carries the semantics.
23 +"""
24 +
25 +from __future__ import annotations
26 +
27 +import functools
28 +import json
29 +from collections.abc import Sequence
30 +from typing import TYPE_CHECKING
31 +
32 +from vllm.entrypoints.openai.engine.protocol import (
33 + ExtractedToolCallInformation,
34 +)
35 +from vllm.parser.engine.events import EventType
36 +from vllm.parser.engine.parser_engine import ParserEngine
37 +from vllm.parser.engine.parser_engine_config import (
38 + ParserEngineConfig,
39 + ParserState,
40 + Transition,
41 +)
42 +
43 +if TYPE_CHECKING:
44 + from vllm.tokenizers import TokenizerLike
45 + from vllm.tool_parsers.abstract_tool_parser import Tool
46 +
47 +MESSAGE_MODEL = "<|message_model|>"
48 +CONTENT_TEXT = "<|content_text|>"
49 +CONTENT_THINKING = "<|content_thinking|>"
50 +CONTENT_INVOKE_TOOL_JSON = "<|content_invoke_tool_json|>"
51 +CONTENT_INVOKE_TOOL_TEXT = "<|content_invoke_tool_text|>"
52 +CONTENT_TOOL_ERROR = "<|content_tool_error|>"
53 +CONTENT_MODEL_END_SAMPLING = "<|content_model_end_sampling|>"
54 +END_MESSAGE = "<|end_message|>"
55 +
56 +INKLING_SPECIAL_TOKENS = (
57 + MESSAGE_MODEL,
58 + CONTENT_TEXT,
59 + CONTENT_THINKING,
60 + CONTENT_INVOKE_TOOL_JSON,
61 + CONTENT_INVOKE_TOOL_TEXT,
62 + CONTENT_TOOL_ERROR,
63 + CONTENT_MODEL_END_SAMPLING,
64 + END_MESSAGE,
65 +)
66 +
67 +_WS = " \t\r\n"
68 +
69 +
70 +def _scan_json_value(raw: str, start: int) -> int | None:
71 + """Return the end index (exclusive) of the JSON object starting at
72 + ``raw[start]``, or ``None`` when the object is still unterminated."""
73 + depth = 0
74 + in_string = False
75 + escape = False
76 + for i in range(start, len(raw)):
77 + ch = raw[i]
78 + if escape:
79 + escape = False
80 + continue
81 + if in_string:
82 + if ch == "\\":
83 + escape = True
84 + elif ch == '"':
85 + in_string = False
86 + continue
87 + if ch == '"':
88 + in_string = True
89 + elif ch == "{":
90 + depth += 1
91 + elif ch == "}":
92 + depth -= 1
93 + if depth == 0:
94 + return i + 1
95 + return None
96 +
97 +
98 +def _args_value_span(raw: str) -> str | None:
99 + """Extract the raw text span of the top-level ``"args"`` value from a
100 + (possibly incomplete) ``{"name":...,"args":{...}}`` wrapper.
101 +
102 + Returns the verbatim substring (prefix-stable across growing input,
103 + which the engine's argument-delta diffing relies on), possibly an
104 + unterminated object prefix; ``None`` when the value has not started.
105 + Raises ``ValueError`` when the value is not a JSON object.
106 + """
107 + depth = 0
108 + in_string = False
109 + escape = False
110 + string_start = -1
111 + last_string: str | None = None
112 + for i, ch in enumerate(raw):
113 + if escape:
114 + escape = False
115 + continue
116 + if in_string:
117 + if ch == "\\":
118 + escape = True
119 + elif ch == '"':
120 + in_string = False
121 + if depth == 1:
122 + last_string = raw[string_start + 1 : i]
123 + continue
124 + if ch == '"':
125 + in_string = True
126 + string_start = i
127 + elif ch == ":" and depth == 1 and last_string == "args":
128 + value_start = i + 1
129 + while value_start < len(raw) and raw[value_start] in _WS:
130 + value_start += 1
131 + if value_start >= len(raw):
132 + return None
133 + if raw[value_start] != "{":
134 + raise ValueError("Inkling tool call args must be a JSON object")
135 + value_end = _scan_json_value(raw, value_start)
136 + if value_end is None:
137 + return raw[value_start:]
138 + return raw[value_start:value_end]
139 + elif ch in "{[":
140 + depth += 1
141 + elif ch in "}]":
142 + depth -= 1
143 + return None
144 +
145 +
146 +def _inkling_arg_converter(raw_args: str, partial: bool) -> str:
147 + """Carve the ``args`` object out of the tool-call JSON wrapper.
148 +
149 + Why a converter at all: the engine's ``tool_args_json`` machinery
150 + treats the *entire* TOOL_ARGS text as the tool arguments, but Inkling's
151 + payload is the ``{"name":...,"args":{...}}`` wrapper — without a
152 + converter, ``_compute_arg_delta`` streams the wrapper verbatim into
153 + the OpenAI ``arguments`` field (``converter is None -> raw delta``,
154 + unconditionally; ``stream_arg_deltas=False`` does not stop it).
155 +
156 + Why a hand-rolled scanner instead of (partial) ``json.loads`` +
157 + ``json.dumps``: the engine diffs successive converter outputs and
158 + requires each to extend the previous one (``startswith``); a
159 + violation silently drops argument deltas. Re-serialization changes
160 + whitespace and closes unterminated structures differently across
161 + ticks, so the only prefix-stable output is a verbatim substring of
162 + the input. The scanner is also string/escape-aware so an ``"args"``
163 + literal inside the name or a string value cannot mislead it, and it
164 + still recovers a partial span when EOS truncates the wrapper (where
165 + ``json.loads`` would fail).
166 + """
167 + span = _args_value_span(raw_args)
168 + if span is None:
169 + # No args value yet (streaming) or none at all (treat as empty).
170 + return "" if partial else "{}"
171 + return span
172 +
173 +
174 +@functools.cache
175 +def inkling_config() -> ParserEngineConfig:
176 + terminals = {
177 + "MSG_MODEL": MESSAGE_MODEL,
178 + "TEXT_START": CONTENT_TEXT,
179 + "THINK_START": CONTENT_THINKING,
180 + "THINK_END": END_MESSAGE,
181 + "END_SAMPLING": CONTENT_MODEL_END_SAMPLING,
182 + "TOOL_START": CONTENT_INVOKE_TOOL_JSON,
183 + "TOOL_TEXT": CONTENT_INVOKE_TOOL_TEXT,
184 + "TOOL_ERROR": CONTENT_TOOL_ERROR,
185 + }
186 + transitions: dict[tuple[ParserState, str], Transition] = {
187 + # ── Between blocks / inside a text block ──────────────────────
188 + (ParserState.CONTENT, "MSG_MODEL"): Transition(
189 + ParserState.MESSAGE_HEADER,
190 + (),
191 + ),
192 + (ParserState.CONTENT, "TEXT_START"): Transition(
193 + ParserState.CONTENT,
194 + (),
195 + ),
196 + (ParserState.CONTENT, "THINK_START"): Transition(
197 + ParserState.REASONING,
198 + (EventType.REASONING_START,),
199 + ),
200 + (ParserState.CONTENT, "TOOL_START"): Transition(
201 + ParserState.TOOL_ARGS,
202 + (EventType.TOOL_CALL_START,),
203 + ),
204 + # Raw / error tool blocks render as visible text.
205 + (ParserState.CONTENT, "TOOL_TEXT"): Transition(
206 + ParserState.CONTENT,
207 + (),
208 + ),
209 + (ParserState.CONTENT, "TOOL_ERROR"): Transition(
210 + ParserState.CONTENT,
211 + (),
212 + ),
213 + # The optional function name between the model-role and content-kind
214 + # markers is metadata, not visible assistant content.
215 + (ParserState.MESSAGE_HEADER, "MSG_MODEL"): Transition(
216 + ParserState.MESSAGE_HEADER,
217 + (),
218 + ),
219 + (ParserState.MESSAGE_HEADER, "TEXT_START"): Transition(
220 + ParserState.CONTENT,
221 + (),
222 + ),
223 + (ParserState.MESSAGE_HEADER, "THINK_START"): Transition(
224 + ParserState.REASONING,
225 + (EventType.REASONING_START,),
226 + ),
227 + (ParserState.MESSAGE_HEADER, "TOOL_START"): Transition(
228 + ParserState.TOOL_ARGS,
229 + (EventType.TOOL_CALL_START,),
230 + ),
231 + (ParserState.MESSAGE_HEADER, "TOOL_TEXT"): Transition(
232 + ParserState.CONTENT,
233 + (),
234 + ),
235 + (ParserState.MESSAGE_HEADER, "TOOL_ERROR"): Transition(
236 + ParserState.CONTENT,
237 + (),
238 + ),
239 + # ── Inside a thinking block ───────────────────────────────────
240 + (ParserState.REASONING, "THINK_START"): Transition(
241 + ParserState.REASONING,
242 + (),
243 + ),
244 + # Defensive: tool call opening while a thinking block is still
245 + # unclosed.
246 + (ParserState.REASONING, "TOOL_START"): Transition(
247 + ParserState.TOOL_ARGS,
248 + (EventType.REASONING_END, EventType.TOOL_CALL_START),
249 + ),
250 + }
251 + # Block-end terminals behave identically regardless of label. A
252 + # closed tool block returns to CONTENT (Inkling has no section wrapper;
253 + # blocks of any kind may follow), which also keeps the block-kind
254 + # and role tokens out of the engine's tool-terminal set so the
255 + # skip_tool_parsing reasoning pass still classifies reasoning.
256 + for end in ("THINK_END", "END_SAMPLING"):
257 + transitions[(ParserState.CONTENT, end)] = Transition(
258 + ParserState.CONTENT,
259 + (),
260 + )
261 + transitions[(ParserState.REASONING, end)] = Transition(
262 + ParserState.CONTENT,
263 + (EventType.REASONING_END,),
264 + )
265 + transitions[(ParserState.TOOL_ARGS, end)] = Transition(
266 + ParserState.CONTENT,
267 + (EventType.TOOL_CALL_END,),
268 + )
269 + transitions[(ParserState.MESSAGE_HEADER, end)] = Transition(
270 + ParserState.CONTENT,
271 + (),
272 + )
273 +
274 + return ParserEngineConfig(
275 + name="inkling",
276 + # Normal generation continues after a prompt-prefilled
277 + # `<|message_model|>`. Non-streaming parsing receives only the generated
278 + # suffix, so begin in the corresponding message-header state as well.
279 + initial_state=ParserState.MESSAGE_HEADER,
280 + terminals=terminals,
281 + # Inkling content-kind markers are the grammar. When the engine is
282 + # used through DelegatingParser, the reasoning pass can hand the tool
283 + # pass reconstructed text whose token-id slice no longer contains the
284 + # content-kind marker that starts the tool block, so keep Inkling on
285 + # the text grammar instead of token-id-only terminal matching.
286 + token_id_terminals={},
287 + transitions=transitions,
288 + arg_converter=_inkling_arg_converter,
289 + stream_arg_deltas=True,
290 + tool_args_json=True,
291 + strip_trailing_reasoning_whitespace=True,
292 + drop_whitespace_only_content_before_tools=True,
293 + strip_content_whitespace_with_tools=False,
294 + validate_tool_names=False,
295 + )
296 +
297 +
298 +class InklingParser(ParserEngine):
299 + CONFIG_NAME = "inkling"
300 +
301 + def __init__(
302 + self,
303 + tokenizer: TokenizerLike,
304 + tools: list[Tool] | None = None,
305 + **kwargs,
306 + ) -> None:
307 + kwargs.setdefault("parser_engine_config", inkling_config())
308 + super().__init__(tokenizer, tools, **kwargs)
309 +
310 + def adjust_initial_state_from_prompt(self, prompt_token_ids: Sequence[int]) -> None:
311 + """Seed the initial parsing state from the prompt tail.
312 +
313 + Mirrors the Rust parser's ``initialize()``: scanning the prompt
314 + backwards, the last relevant special token decides whether generation
315 + continues inside a thinking block, a text block, a model message
316 + header, or between blocks.
317 + """
318 + vocab = self.vocab
319 + thinking_id = vocab.get(CONTENT_THINKING)
320 + text_id = vocab.get(CONTENT_TEXT)
321 + model_id = vocab.get(MESSAGE_MODEL)
322 + special_ids = {vocab[text] for text in INKLING_SPECIAL_TOKENS if text in vocab}
323 + for token_id in reversed(prompt_token_ids):
324 + if token_id == thinking_id:
325 + self._engine.reset(initial_state=ParserState.REASONING)
326 + self._streaming_initialized = True
327 + return
328 + if token_id == text_id:
329 + self._engine.reset(initial_state=ParserState.CONTENT)
330 + self._streaming_initialized = True
331 + return
332 + if token_id == model_id:
333 + self._engine.reset(initial_state=ParserState.MESSAGE_HEADER)
334 + self._streaming_initialized = True
335 + return
336 + if token_id in special_ids:
337 + break
338 +
339 + def is_reasoning_end(self, input_ids: list[int]) -> bool:
340 + vocab = self.vocab
341 + thinking_id = vocab.get(CONTENT_THINKING)
342 + text_id = vocab.get(CONTENT_TEXT)
343 + model_id = vocab.get(MESSAGE_MODEL)
344 + end_sampling_id = vocab.get(CONTENT_MODEL_END_SAMPLING)
345 + for token_id in reversed(input_ids):
346 + if token_id in (thinking_id, model_id):
347 + return False
348 + if token_id in (text_id, end_sampling_id):
349 + return True
350 + return False
351 +
352 + def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
353 + vocab = self.vocab
354 + thinking_id = vocab.get(CONTENT_THINKING)
355 + end_ids = {
356 + token_id
357 + for token_id in (
358 + vocab.get(END_MESSAGE),
359 + vocab.get(CONTENT_MODEL_END_SAMPLING),
360 + )
361 + if token_id is not None
362 + }
363 + in_reasoning = False
364 + count = 0
365 + for token_id in token_ids:
366 + if token_id == thinking_id:
367 + in_reasoning = True
368 + continue
369 + if token_id in end_ids:
370 + in_reasoning = False
371 + continue
372 + if in_reasoning:
373 + count += 1
374 + return count
375 +
376 + def _single_pass_parse(
377 + self,
378 + text: str,
379 + token_ids: Sequence[int],
380 + initial_state: ParserState | None = None,
381 + ) -> tuple[str | None, str | None, ExtractedToolCallInformation]:
382 + reasoning, content, tool_call_info = super()._single_pass_parse(
383 + text, token_ids, initial_state=initial_state
384 + )
385 + # The engine defers content that follows tool-call events within a
386 + # single pass; Inkling allows text blocks after tool-call blocks, so
387 + # flush the trailing text (matching the Rust unified parser).
388 + if self._deferred_content:
389 + trailing = self._deferred_content
390 + self._deferred_content = ""
391 + content = self._strip_content_whitespace(
392 + (content or "") + trailing,
393 + tool_call_info.tools_called,
394 + )
395 + tool_call_info = ExtractedToolCallInformation(
396 + tools_called=tool_call_info.tools_called,
397 + tool_calls=tool_call_info.tool_calls,
398 + content=content,
399 + )
400 + return reasoning, content, tool_call_info
401 +
402 + @staticmethod
403 + def _extract_args_value(parsed: dict) -> str | None:
404 + # Inkling wraps arguments under "args" rather than "arguments".
405 + for key in ("args", "arguments", "parameters"):
406 + if key in parsed:
407 + val = parsed[key]
408 + if isinstance(val, str):
409 + return val
410 + return json.dumps(val, ensure_ascii=False)
411 + return None
vllm/reasoning/__init__.py
+4
@@ -128,6 +128,10 @@ _REASONING_PARSERS_TO_REGISTER = {
128 "step3p5_reasoning_parser",
129 "Step3p5ReasoningParser",
130 ),
131 + "inkling": (
132 + "inkling_reasoning_parser",
133 + "InklingParserReasoningAdapter",
134 + ),
135 }
136
137
vllm/reasoning/inkling_reasoning_parser.py new
+6
@@ -0,0 +1,6 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from vllm.parser.engine.registered_adapters import InklingParserReasoningAdapter
5 +
6 +__all__ = ["InklingParserReasoningAdapter"]
vllm/renderers/inkling.py new
+178
@@ -0,0 +1,178 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Native Inkling chat renderer for the Python frontend.
4 +
5 +Mirrors the Rust frontend's native Inkling renderer
6 +(``rust/src/chat/src/renderer/inkling/mod.rs``): chat messages are rendered
7 +directly to token ids — Inkling has no Jinja chat template and no faithful
8 +text form.
9 +
10 +The encoding logic lives in ``inkling_encoding.py`` behind a narrow
11 +"OpenAI messages + tools -> token ids" call; see the swap-point comment
12 +in :meth:`InklingRenderer._render` for adopting a standalone Inkling
13 +input-processing library (mistral-common style) later.
14 +"""
15 +
16 +from vllm.config import VllmConfig
17 +from vllm.entrypoints.chat_utils import (
18 + ChatCompletionMessageParam,
19 + ConversationMessage,
20 + parse_chat_messages,
21 + parse_chat_messages_async,
22 +)
23 +from vllm.logger import init_logger
24 +from vllm.tokenizers.hf import HfTokenizer
25 +from vllm.utils.async_utils import make_async
26 +
27 +from .base import BaseRenderer
28 +from .inkling_encoding import SPECIAL_TOKEN_SPELLINGS, render_inkling_messages
29 +from .inputs import DictPrompt
30 +from .inputs.preprocess import parse_dec_only_prompt
31 +from .params import ChatParams
32 +
33 +logger = init_logger(__name__)
34 +
35 +_NAMED_REASONING_EFFORT = {
36 + "none": 0.0,
37 + "minimal": 0.1,
38 + "low": 0.2,
39 + "medium": 0.7,
40 + "high": 0.9,
41 + "xhigh": 0.99,
42 + "max": 0.99,
43 +}
44 +
45 +_DEFAULT_REASONING_EFFORT = 0.9
46 +
47 +
48 +def _resolve_reasoning_effort(value: object) -> float | int | None:
49 + if value is None:
50 + return _DEFAULT_REASONING_EFFORT
51 + if isinstance(value, str):
52 + return _NAMED_REASONING_EFFORT.get(value)
53 + if isinstance(value, bool) or not isinstance(value, (int, float)):
54 + return None
55 + return value
56 +
57 +
58 +class _HfBackedTmlTokenizer:
59 + """Adapts an HF tokenizer to the encoding core's tokenizer protocol.
60 +
61 + Special-token ids are resolved from the tokenizer vocab at
62 + construction (never hardcoded), trying each known spelling — the Inkling
63 + HF vocab exposes some semantic slots as ``<|unused_NNNNNN|>`` tokens.
64 + """
65 +
66 + def __init__(self, tokenizer: HfTokenizer) -> None:
67 + self._tokenizer = tokenizer
68 +
69 + vocab = tokenizer.get_vocab()
70 + special_ids: dict[str, int] = {}
71 + missing: list[str] = []
72 + for token, spellings in SPECIAL_TOKEN_SPELLINGS.items():
73 + for spelling in spellings:
74 + token_id = vocab.get(spelling)
75 + if token_id is not None:
76 + special_ids[token] = token_id
77 + break
78 + else:
79 + missing.append(token)
80 + if missing:
81 + raise ValueError(f"Inkling tokenizer is missing special tokens: {missing}")
82 + self._special_ids = special_ids
83 +
84 + def encode_text(self, text: str) -> list[int]:
85 + return self._tokenizer.encode(text, add_special_tokens=False)
86 +
87 + def encode_special(self, token: str) -> int:
88 + return self._special_ids[token]
89 +
90 +
91 +class InklingRenderer(BaseRenderer[HfTokenizer]):
92 + def __init__(
93 + self,
94 + config: VllmConfig,
95 + tokenizer: HfTokenizer | None,
96 + ) -> None:
97 + super().__init__(config, tokenizer)
98 +
99 + self._inkling_tokenizer = _HfBackedTmlTokenizer(self.get_tokenizer())
100 + self._render_async = make_async(self._render, executor=self._executor)
101 +
102 + def _render(
103 + self,
104 + messages: list[ChatCompletionMessageParam],
105 + params: ChatParams,
106 + ) -> list[int]:
107 + kwargs = params.chat_template_kwargs or {}
108 +
109 + if kwargs.get("continue_final_message"):
110 + raise ValueError("Inkling renderer does not support continue_final_message")
111 +
112 + reasoning_effort = _resolve_reasoning_effort(kwargs.get("reasoning_effort"))
113 +
114 + try:
115 + # Swap point: to adopt a standalone Inkling input-processing
116 + # library, replace this call (and the _HfBackedTmlTokenizer
117 + # adapter above) with the library's renderer.
118 + return render_inkling_messages(
119 + messages,
120 + self._inkling_tokenizer,
121 + add_generation_prompt=kwargs.get("add_generation_prompt", True),
122 + tools=kwargs.get("tools"),
123 + reasoning_effort=reasoning_effort,
124 + )
125 + except ValueError:
126 + raise
127 + except (TypeError, KeyError) as e:
128 + # Malformed request content; surface as a request error.
129 + raise ValueError(str(e)) from e
130 + except Exception as e:
131 + logger.exception("Error while rendering Inkling chat messages")
132 + raise ValueError(str(e)) from e
133 +
134 + def render_messages(
135 + self,
136 + messages: list[ChatCompletionMessageParam],
137 + params: ChatParams,
138 + ) -> tuple[list[ConversationMessage], DictPrompt]:
139 + conversation, mm_data, mm_uuids = parse_chat_messages(
140 + messages,
141 + self.model_config,
142 + content_format="string",
143 + media_io_kwargs=params.media_io_kwargs,
144 + mm_processor_kwargs=params.mm_processor_kwargs,
145 + )
146 +
147 + token_ids = self._render(messages, params)
148 +
149 + prompt = parse_dec_only_prompt(token_ids)
150 + if mm_data is not None:
151 + prompt["multi_modal_data"] = mm_data
152 + if mm_uuids is not None:
153 + prompt["multi_modal_uuids"] = mm_uuids
154 +
155 + return conversation, prompt
156 +
157 + async def render_messages_async(
158 + self,
159 + messages: list[ChatCompletionMessageParam],
160 + params: ChatParams,
161 + ) -> tuple[list[ConversationMessage], DictPrompt]:
162 + conversation, mm_data, mm_uuids = await parse_chat_messages_async(
163 + messages,
164 + self.model_config,
165 + content_format="string",
166 + media_io_kwargs=params.media_io_kwargs,
167 + mm_processor_kwargs=params.mm_processor_kwargs,
168 + )
169 +
170 + token_ids = await self._render_async(messages, params)
171 +
172 + prompt = parse_dec_only_prompt(token_ids)
173 + if mm_data is not None:
174 + prompt["multi_modal_data"] = mm_data
175 + if mm_uuids is not None:
176 + prompt["multi_modal_uuids"] = mm_uuids
177 +
178 + return conversation, prompt
vllm/renderers/inkling_encoding.py new
+383
@@ -0,0 +1,383 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inkling chat-encoding core.
4 +
5 +Pure implementation of Inkling chat rendering, kept deliberately free of
6 +vLLM imports: it depends only on the :class:`InklingTextTokenizer` protocol
7 +and speaks OpenAI-style message dicts. If a standalone Inkling
8 +input-processing library becomes available, this module is the unit to
9 +swap out — the swap point is marked in ``vllm/renderers/inkling.py``.
10 +
11 +Multimodal parts follow the contract of vLLM's Inkling multimodal processor
12 +(``InklingMultiModalProcessor`` anchors prompt updates on the bare
13 +content-kind marker and inserts the per-patch placeholder run itself):
14 +
15 +* image parts emit only ``<|content_image|>`` — no seed placeholder id;
16 +* audio parts emit ``<|content_audio_input|><|audio_end|>`` — no seed
17 + placeholder id between them.
18 +
19 +``reasoning_effort`` (a float in [0, 0.99], sourced from
20 +``chat_template_kwargs`` only) renders the ``Thinking effort level:``
21 +system control block after the tool declarations and initial system messages.
22 +"""
23 +
24 +from __future__ import annotations
25 +
26 +import json
27 +from collections.abc import Iterator, Mapping, Sequence
28 +from typing import Any, Protocol
29 +
30 +END_OF_TEXT = "<|endoftext|>"
31 +MESSAGE_USER = "<|message_user|>"
32 +MESSAGE_MODEL = "<|message_model|>"
33 +MESSAGE_SYSTEM = "<|message_system|>"
34 +MESSAGE_TOOL = "<|message_tool|>"
35 +CONTENT_TEXT = "<|content_text|>"
36 +CONTENT_IMAGE = "<|content_image|>"
37 +CONTENT_MODEL_END_SAMPLING = "<|content_model_end_sampling|>"
38 +CONTENT_THINKING = "<|content_thinking|>"
39 +CONTENT_AUDIO_INPUT = "<|content_audio_input|>"
40 +CONTENT_TOOL_ERROR = "<|content_tool_error|>"
41 +CONTENT_XML = "<|content_xml|>"
42 +CONTENT_INVOKE_TOOL_JSON = "<|content_invoke_tool_json|>"
43 +CONTENT_INVOKE_TOOL_TEXT = "<|content_invoke_tool_text|>"
44 +END_MESSAGE = "<|end_message|>"
45 +AUDIO_END = "<|audio_end|>"
46 +
47 +ROLE_MESSAGE_TOKENS: dict[str, str] = {
48 + "user": MESSAGE_USER,
49 + "assistant": MESSAGE_MODEL,
50 + "system": MESSAGE_SYSTEM,
51 + # The developer role folds into the system role on the Inkling wire.
52 + "developer": MESSAGE_SYSTEM,
53 + "tool": MESSAGE_TOOL,
54 +}
55 +
56 +# Alternate vocab spellings per semantic token. The Inkling HF tokenizer
57 +# exposes some semantic slots as ``<|unused_NNNNNN|>`` tokens (notably
58 +# CONTENT_XML = <|unused_200024|>), so ID resolution must try these
59 +# spellings in order.
60 +SPECIAL_TOKEN_SPELLINGS: dict[str, tuple[str, ...]] = {
61 + MESSAGE_USER: (MESSAGE_USER,),
62 + MESSAGE_MODEL: (MESSAGE_MODEL,),
63 + MESSAGE_SYSTEM: (MESSAGE_SYSTEM,),
64 + MESSAGE_TOOL: (MESSAGE_TOOL,),
65 + CONTENT_TEXT: (CONTENT_TEXT,),
66 + CONTENT_IMAGE: (CONTENT_IMAGE,),
67 + CONTENT_MODEL_END_SAMPLING: (CONTENT_MODEL_END_SAMPLING,),
68 + CONTENT_THINKING: (CONTENT_THINKING,),
69 + CONTENT_AUDIO_INPUT: (CONTENT_AUDIO_INPUT,),
70 + CONTENT_XML: (CONTENT_XML, "<|unused_200024|>"),
71 + CONTENT_INVOKE_TOOL_JSON: (CONTENT_INVOKE_TOOL_JSON,),
72 + END_MESSAGE: (END_MESSAGE,),
73 + AUDIO_END: (AUDIO_END,),
74 +}
75 +
76 +
77 +class InklingTextTokenizer(Protocol):
78 + """Structural tokenizer contract required by the renderer."""
79 +
80 + def encode_text(self, text: str) -> list[int]: ...
81 +
82 + def encode_special(self, token: str) -> int: ...
83 +
84 +
85 +# OpenAI content-part type spellings that mean image / audio (rendering only
86 +# needs the kind, not the bytes — the bytes are handled by the MM processor).
87 +_IMAGE_PART_TYPES = frozenset({"image", "input_image", "image_url"})
88 +_AUDIO_PART_TYPES = frozenset({"audio", "input_audio", "audio_url"})
89 +
90 +_MAX_REASONING_EFFORT = 0.99
91 +
92 +
93 +def render_inkling_messages(
94 + messages: Sequence[Mapping[str, Any]],
95 + tokenizer: InklingTextTokenizer,
96 + *,
97 + add_generation_prompt: bool = True,
98 + tools: Sequence[Mapping[str, Any]] | None = None,
99 + reasoning_effort: float | None = None,
100 +) -> list[int]:
101 + """Render chat messages to Inkling input ids.
102 +
103 + PURE renderer: emits Inkling framing plus bare media markers; media
104 + encoding and placeholder expansion happen later in the MM processor.
105 + ``add_generation_prompt`` appends the assistant turn opener so the
106 + model continues into the response.
107 + """
108 + input_ids: list[int] = []
109 + tool_call_id_to_name: dict[str, str] = {}
110 +
111 + # Request-level tools plus per-developer-message tools (Rust renderer
112 + # semantics) are declared in a single leading system block.
113 + all_tools = list(tools or [])
114 + for message in messages:
115 + if message.get("role") == "developer":
116 + all_tools.extend(message.get("tools") or [])
117 +
118 + if all_tools:
119 + _append_message(
120 + input_ids,
121 + tokenizer,
122 + "system",
123 + "xml",
124 + _tool_declare_json(all_tools),
125 + author_name="tool_declare",
126 + )
127 +
128 + for message in messages:
129 + role = _expect_role(message)
130 + if role not in {"system", "developer"} and reasoning_effort is not None:
131 + _append_reasoning_effort(input_ids, tokenizer, reasoning_effort)
132 + reasoning_effort = None
133 +
134 + if role == "tool":
135 + tool_name = message.get("name") or tool_call_id_to_name.get(
136 + str(message.get("tool_call_id") or ""), ""
137 + )
138 + _append_message(
139 + input_ids,
140 + tokenizer,
141 + "tool",
142 + "text",
143 + _flatten_text_content(message.get("content")),
144 + author_name=str(tool_name),
145 + )
146 + continue
147 +
148 + if role == "assistant":
149 + reasoning_content = message.get("reasoning")
150 + if reasoning_content is None:
151 + reasoning_content = message.get("reasoning_content")
152 + if reasoning_content:
153 + if not isinstance(reasoning_content, str):
154 + raise TypeError(
155 + "assistant reasoning_content must be a string for "
156 + "Inkling rendering"
157 + )
158 + _append_message(
159 + input_ids,
160 + tokenizer,
161 + "assistant",
162 + "thinking",
163 + reasoning_content,
164 + )
165 +
166 + for kind, text in _iter_render_parts(message.get("content")):
167 + _append_message(input_ids, tokenizer, role, kind, text)
168 +
169 + if role == "assistant":
170 + for tool_call in message.get("tool_calls") or []:
171 + name, args = _tool_call_name_and_args(tool_call)
172 + tool_call_id = _as_mapping(tool_call).get("id")
173 + if tool_call_id:
174 + tool_call_id_to_name[str(tool_call_id)] = name
175 + _append_message(
176 + input_ids,
177 + tokenizer,
178 + "assistant",
179 + "invoke_tool_json",
180 + _tool_call_json(name, args),
181 + author_name=name,
182 + )
183 +
184 + input_ids.append(tokenizer.encode_special(CONTENT_MODEL_END_SAMPLING))
185 +
186 + if reasoning_effort is not None:
187 + _append_reasoning_effort(input_ids, tokenizer, reasoning_effort)
188 +
189 + if add_generation_prompt:
190 + input_ids.append(tokenizer.encode_special(MESSAGE_MODEL))
191 + return input_ids
192 +
193 +
194 +def _append_reasoning_effort(
195 + input_ids: list[int],
196 + tokenizer: InklingTextTokenizer,
197 + reasoning_effort: float,
198 +) -> None:
199 + _append_message(
200 + input_ids,
201 + tokenizer,
202 + "system",
203 + "text",
204 + _thinking_effort_text(reasoning_effort),
205 + )
206 +
207 +
208 +def _thinking_effort_text(reasoning_effort: Any) -> str:
209 + if isinstance(reasoning_effort, bool) or not isinstance(
210 + reasoning_effort, (int, float)
211 + ):
212 + raise TypeError(
213 + "Inkling reasoning_effort must be a number in [0.0, 0.99], "
214 + f"got {type(reasoning_effort).__name__}"
215 + )
216 + value = float(reasoning_effort)
217 + if not 0.0 <= value <= _MAX_REASONING_EFFORT:
218 + raise ValueError(
219 + f"Inkling reasoning_effort must be in [0.0, 0.99], got {value}"
220 + )
221 +
222 + effort_text = f"{value:.2f}".rstrip("0").rstrip(".")
223 + if effort_text in {"0", "-0"}:
224 + effort_text = "0.0"
225 + return f"Thinking effort level: {effort_text}"
226 +
227 +
228 +def _append_message(
229 + input_ids: list[int],
230 + tokenizer: InklingTextTokenizer,
231 + role: str,
232 + kind: str,
233 + text: str,
234 + *,
235 + author_name: str | None = None,
236 +) -> None:
237 + input_ids.append(tokenizer.encode_special(ROLE_MESSAGE_TOKENS[role]))
238 + if author_name:
239 + input_ids.extend(tokenizer.encode_text(author_name))
240 +
241 + if kind == "text":
242 + input_ids.append(tokenizer.encode_special(CONTENT_TEXT))
243 + input_ids.extend(tokenizer.encode_text(text))
244 + elif kind == "image":
245 + # Bare marker only: the MM processor replaces the marker with
246 + # `marker + placeholder * num_patches` (see mm_preprocess.py).
247 + input_ids.append(tokenizer.encode_special(CONTENT_IMAGE))
248 + elif kind == "audio":
249 + input_ids.append(tokenizer.encode_special(CONTENT_AUDIO_INPUT))
250 + input_ids.append(tokenizer.encode_special(AUDIO_END))
251 + elif kind == "thinking":
252 + input_ids.append(tokenizer.encode_special(CONTENT_THINKING))
253 + input_ids.extend(tokenizer.encode_text(text))
254 + elif kind == "xml":
255 + input_ids.append(tokenizer.encode_special(CONTENT_XML))
256 + input_ids.extend(tokenizer.encode_text(text))
257 + elif kind == "invoke_tool_json":
258 + input_ids.append(tokenizer.encode_special(CONTENT_INVOKE_TOOL_JSON))
259 + input_ids.extend(tokenizer.encode_text(text))
260 + else:
261 + raise ValueError(f"unsupported Inkling render part kind: {kind!r}")
262 +
263 + input_ids.append(tokenizer.encode_special(END_MESSAGE))
264 +
265 +
266 +def _iter_render_parts(content: Any) -> Iterator[tuple[str, str]]:
267 + """Yield (kind, text) per content part: kind in {text, image, audio}."""
268 + if content is None:
269 + return
270 + if isinstance(content, str):
271 + if content:
272 + yield ("text", content)
273 + return
274 + if not isinstance(content, Sequence) or isinstance(content, (bytes, bytearray)):
275 + raise TypeError("message content must be a string or a sequence of parts")
276 + for part in content:
277 + if isinstance(part, str):
278 + yield ("text", part)
279 + continue
280 + if not isinstance(part, Mapping):
281 + raise TypeError(f"content part must be mapping, got {type(part).__name__}")
282 + ptype = part.get("type")
283 + if ptype in (None, "text", "input_text"):
284 + text = part.get("text", "")
285 + yield ("text", text if isinstance(text, str) else "")
286 + elif ptype in _IMAGE_PART_TYPES:
287 + yield ("image", "")
288 + elif ptype in _AUDIO_PART_TYPES:
289 + yield ("audio", "")
290 + else:
291 + raise ValueError(f"unsupported content part type: {ptype!r}")
292 +
293 +
294 +def _flatten_text_content(content: Any) -> str:
295 + """Flatten a tool-response content (string or text parts) to text."""
296 + if content is None:
297 + return ""
298 + if isinstance(content, str):
299 + return content
300 + parts: list[str] = []
301 + for kind, text in _iter_render_parts(content):
302 + if kind != "text":
303 + raise ValueError(
304 + "Inkling tool response content must be text, "
305 + f"got a part of kind {kind!r}"
306 + )
307 + parts.append(text)
308 + return "".join(parts)
309 +
310 +
311 +def _expect_role(message: Mapping[str, Any]) -> str:
312 + role = message.get("role")
313 + if role not in ROLE_MESSAGE_TOKENS:
314 + raise ValueError(
315 + f"unsupported Inkling message role {role!r}; "
316 + f"expected one of {sorted(ROLE_MESSAGE_TOKENS)}"
317 + )
318 + return str(role)
319 +
320 +
321 +def _as_mapping(value: Any) -> Mapping[str, Any]:
322 + if isinstance(value, Mapping):
323 + return value
324 + if hasattr(value, "model_dump"):
325 + dumped = value.model_dump()
326 + if isinstance(dumped, Mapping):
327 + return dumped
328 + raise TypeError(f"expected mapping, got {type(value).__name__}")
329 +
330 +
331 +def _canonical_json(value: Any) -> str:
332 + return json.dumps(
333 + _sort_json(value),
334 + ensure_ascii=False,
335 + allow_nan=False,
336 + separators=(",", ":"),
337 + )
338 +
339 +
340 +def _sort_json(value: Any) -> Any:
341 + if isinstance(value, Mapping):
342 + return {str(key): _sort_json(value[key]) for key in sorted(value)}
343 + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
344 + return [_sort_json(item) for item in value]
345 + return value
346 +
347 +
348 +def _tool_declare_json(tools: Sequence[Mapping[str, Any]]) -> str:
349 + tool_specs = []
350 + for tool_value in tools:
351 + tool = _as_mapping(tool_value)
352 + function = _as_mapping(tool.get("function", {}))
353 + tool_specs.append(
354 + {
355 + "description": function.get("description") or "",
356 + "name": function["name"],
357 + "parameters": function.get("parameters") or {},
358 + "type": tool.get("type", "function"),
359 + }
360 + )
361 + return _canonical_json(tool_specs)
362 +
363 +
364 +def _tool_call_name_and_args(tool_call_value: Any) -> tuple[str, Mapping[str, Any]]:
365 + tool_call = _as_mapping(tool_call_value)
366 + function = _as_mapping(tool_call.get("function", {}))
367 + name = function.get("name")
368 + if not isinstance(name, str):
369 + raise TypeError("tool call function name must be a string")
370 +
371 + raw_args = function.get("arguments") or {}
372 + if isinstance(raw_args, str):
373 + args = json.loads(raw_args) if raw_args.strip() else {}
374 + else:
375 + args = raw_args
376 + if not isinstance(args, Mapping):
377 + raise TypeError("tool call function arguments must decode to an object")
378 + return name, args
379 +
380 +
381 +def _tool_call_json(name: str, args: Mapping[str, Any]) -> str:
382 + name_json = json.dumps(name, ensure_ascii=False, allow_nan=False)
383 + return f'{{"name":{name_json},"args":{_canonical_json(args)}}}'
vllm/renderers/registry.py
+1
@@ -26,6 +26,7 @@ _VLLM_RENDERERS = {
26 "kimi_audio": ("hf", "HfRenderer"),
27 "mistral": ("mistral", "MistralRenderer"),
28 "terratorch": ("terratorch", "TerratorchRenderer"),
29 + "inkling": ("inkling", "InklingRenderer"),
30 }
31
32
vllm/tokenizers/registry.py
+4
@@ -44,6 +44,10 @@ _VLLM_TOKENIZERS = {
44 "hf": ("hf", "CachedHfTokenizer"),
45 "kimi_audio": ("kimi_audio", "KimiAudioTokenizer"),
46 "mistral": ("mistral", "MistralTokenizer"),
47 + # Inkling uses the plain HF tokenizer for token operations; the "inkling"
48 + # mode exists to select the InklingRenderer, which renders chat to
49 + # token ids natively (Inkling has no Jinja chat template).
50 + "inkling": ("hf", "CachedHfTokenizer"),
51 }
52
53
vllm/tool_parsers/__init__.py
+4
@@ -174,6 +174,10 @@ _TOOL_PARSERS_TO_REGISTER = {
174 "step3p5_tool_parser",
175 "Step3p5ToolParser",
176 ),
177 + "inkling": (
178 + "inkling_tool_parser",
179 + "InklingEngineToolParser",
180 + ),
181 "xlam": (
182 "xlam_tool_parser",
183 "xLAMToolParser",
vllm/tool_parsers/inkling_tool_parser.py new
+11
@@ -0,0 +1,11 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from vllm.parser.engine.registered_adapters import InklingParserToolAdapter
5 +
6 +
7 +class InklingEngineToolParser(InklingParserToolAdapter): # type: ignore[valid-type, misc]
8 + # No Inkling structural-tag grammar is wired up yet; fall back to auto
9 + # parsing for named/required tool choice.
10 + structural_tag_model = None
11 + supports_required_and_named = False
vllm/transformers_utils/config.py
+2
@@ -127,6 +127,8 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
127 laguna="LagunaConfig",
128 lfm2_moe="Lfm2MoeConfig",
129 **{"unlimited-ocr": "UnlimitedOCRConfig"},
130 + inkling_mm_model="InklingMMConfig",
131 + inkling_model="InklingModelConfig",
132 )
133
134 _SPECULATIVE_DECODING_CONFIGS: set[str] = {"eagle", "speculators", "medusa"}
vllm/transformers_utils/configs/__init__.py
+8
@@ -94,6 +94,10 @@ _CLASS_TO_MODULE: dict[str, str] = {
94 "Qwen3_5TextConfig": "vllm.transformers_utils.configs.qwen3_5",
95 "Qwen3_5MoeConfig": "vllm.transformers_utils.configs.qwen3_5_moe",
96 "Qwen3_5MoeTextConfig": "vllm.transformers_utils.configs.qwen3_5_moe",
97 + "InklingModelConfig": "vllm.models.inkling.configs",
98 + "InklingAudioConfig": "vllm.models.inkling.configs",
99 + "InklingVisionConfig": "vllm.models.inkling.configs",
100 + "InklingMMConfig": "vllm.models.inkling.configs",
101 # Special case: DeepseekV3Config is from HuggingFace Transformers
102 "DeepseekV3Config": "transformers",
103 }
@@ -174,6 +178,10 @@ __all__ = [
178 "Qwen3_5TextConfig",
179 "Qwen3_5MoeConfig",
180 "Qwen3_5MoeTextConfig",
181 + "InklingModelConfig",
182 + "InklingAudioConfig",
183 + "InklingVisionConfig",
184 + "InklingMMConfig",
185 ]
186
187
vllm/transformers_utils/processors/__init__.py
+6
@@ -44,6 +44,9 @@ __all__ = [
44 "Ovis2_5Processor",
45 "Qwen3ASRProcessor",
46 "Step3VLProcessor",
47 + "InklingProcessor",
48 + "InklingImageProcessor",
49 + "InklingAudioFeatureExtractor",
50 ]
51
52 _CLASS_TO_MODULE: dict[str, str] = {
@@ -80,6 +83,9 @@ _CLASS_TO_MODULE: dict[str, str] = {
83 "Ovis2_5Processor": "vllm.transformers_utils.processors.ovis2_5",
84 "Qwen3ASRProcessor": "vllm.transformers_utils.processors.qwen3_asr",
85 "Step3VLProcessor": "vllm.transformers_utils.processors.step3_vl",
86 + "InklingProcessor": "vllm.transformers_utils.processors.inkling",
87 + "InklingImageProcessor": "vllm.transformers_utils.processors.inkling",
88 + "InklingAudioFeatureExtractor": "vllm.transformers_utils.processors.inkling",
89 }
90
91
vllm/transformers_utils/processors/inkling.py new
+504
@@ -0,0 +1,504 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Vendored HuggingFace-convention processors for the Inkling Titan model.
4 +
5 +Implements the image processor (numba patchifier), the audio feature extractor
6 +(STFT/dMel path), the composite processor, and the MM token-id constants.
7 +
8 +Besides raw bytes / file paths, the extractors also accept the dummy inputs
9 +vLLM generates during profiling (PIL images / numpy audio arrays).
10 +"""
11 +
12 +from __future__ import annotations
13 +
14 +import io
15 +import math
16 +from collections.abc import Sequence
17 +from dataclasses import dataclass
18 +
19 +import numpy as np
20 +import torch
21 +import torch.nn.functional as F
22 +from numba import njit
23 +from transformers.feature_extraction_utils import (
24 + BatchFeature,
25 + FeatureExtractionMixin,
26 +)
27 +from transformers.image_processing_utils import BaseImageProcessor
28 +from transformers.image_utils import ImageInput
29 +
30 +# ---------------------------------------------------------------------------
31 +# MM token-id constants
32 +# ---------------------------------------------------------------------------
33 +
34 +# Block-start marker token ids. These are real tokens the model was trained on
35 +# (``<|content_image|>`` / ``<|content_audio_input|>``) that mark the start of an
36 +# image/audio embedding block; they are kept verbatim in ``input_ids``.
37 +IMAGE_MARKER_ID = 200005 # <|content_image|>
38 +AUDIO_MARKER_ID = 200020 # <|content_audio_input|>
39 +
40 +# Per-patch / per-frame placeholder ids marking where tower embeddings are
41 +# scattered in. These are unused slots in the padded vocabulary and the
42 +# corresponding positions are always overwritten by tower embeddings.
43 +IMAGE_TOKEN_ID = 200054 # <|unused_200054|>
44 +AUDIO_TOKEN_ID = 200053 # <|unused_200053|>
45 +
46 +# ---------------------------------------------------------------------------
47 +# Image processing
48 +# ---------------------------------------------------------------------------
49 +
50 +IMAGE_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)
51 +IMAGE_STD = np.array([0.26862954, 0.2613026, 0.2757771], dtype=np.float32)
52 +PAD_RAW_VALUE = np.float32(-1.0 / 255.0)
53 +PAD_NORM = (np.full((3,), PAD_RAW_VALUE, dtype=np.float32) - IMAGE_MEAN) / IMAGE_STD
54 +
55 +
56 +def _validate_image_rescale(
57 + rescale_image_frac: float | None,
58 + rescale_image_max_upscaled_long_edge: int | None,
59 +) -> None:
60 + if rescale_image_frac is not None and (
61 + not math.isfinite(rescale_image_frac) or rescale_image_frac <= 0
62 + ):
63 + raise ValueError(
64 + "rescale_image_frac must be positive and finite or None, "
65 + f"got {rescale_image_frac}"
66 + )
67 + if rescale_image_max_upscaled_long_edge is None:
68 + return
69 + if rescale_image_max_upscaled_long_edge <= 0:
70 + raise ValueError(
71 + "rescale_image_max_upscaled_long_edge must be positive or None, "
72 + f"got {rescale_image_max_upscaled_long_edge}"
73 + )
74 + if rescale_image_frac is None or rescale_image_frac <= 1.0:
75 + raise ValueError(
76 + "rescale_image_max_upscaled_long_edge requires rescale_image_frac > 1, "
77 + f"got {rescale_image_frac}"
78 + )
79 +
80 +
81 +def _scaled_image_dimensions(
82 + width: int,
83 + height: int,
84 + rescale_image_frac: float | None,
85 + rescale_image_max_upscaled_long_edge: int | None,
86 +) -> tuple[int, int]:
87 + """Return the long-edge-scaled ``(width, height)``."""
88 + if rescale_image_frac is None:
89 + return width, height
90 +
91 + long_edge = max(width, height)
92 + if long_edge == 0:
93 + return width, height
94 +
95 + target_long_edge = float(long_edge) * rescale_image_frac
96 + if rescale_image_max_upscaled_long_edge is not None:
97 + effective_cap = max(rescale_image_max_upscaled_long_edge, long_edge)
98 + target_long_edge = min(target_long_edge, float(effective_cap))
99 +
100 + ratio = target_long_edge / float(long_edge)
101 + if ratio == 1.0:
102 + return width, height
103 +
104 + def scale(value: int) -> int:
105 + return max(1, math.floor(float(value) * ratio + 0.5))
106 +
107 + return scale(width), scale(height)
108 +
109 +
110 +def _load_image_bytes(image) -> bytes:
111 + """Encode a PIL image as raw PNG bytes for preprocessing.
112 +
113 + The HF processor is always handed ``PIL.Image`` instances by vLLM, so no
114 + other input types need to be supported here.
115 + """
116 + if image.mode != "RGB":
117 + image = image.convert("RGB")
118 + buf = io.BytesIO()
119 + image.save(buf, format="PNG")
120 + return buf.getvalue()
121 +
122 +
123 +@njit(cache=True)
124 +def _fill_patches_numba(
125 + arr: np.ndarray,
126 + patch_size: int,
127 + patches: np.ndarray,
128 + mean: np.ndarray,
129 + std: np.ndarray,
130 + pad_norm: np.ndarray,
131 +) -> None:
132 + h = arr.shape[0]
133 + w = arr.shape[1]
134 + nph = (h + patch_size - 1) // patch_size
135 + npw = w // patch_size + 1
136 + inv255 = np.float32(1.0 / 255.0)
137 +
138 + for k in range(nph * npw):
139 + i = k // npw
140 + j = k - i * npw
141 + y_base = i * patch_size
142 + x_base = j * patch_size
143 +
144 + for y in range(patch_size):
145 + iy = y_base + y
146 + for x in range(patch_size):
147 + ix = x_base + x
148 + if iy < h and ix < w:
149 + for c in range(3):
150 + raw = np.float32(arr[iy, ix, c]) * inv255
151 + patches[k, y, x, c] = (raw - mean[c]) / std[c]
152 + else:
153 + for c in range(3):
154 + patches[k, y, x, c] = pad_norm[c]
155 +
156 +
157 +def _encode_image_bytes(
158 + image_bytes: bytes,
159 + *,
160 + patch_size: int,
161 + rescale_image_frac: float | None,
162 + rescale_image_max_upscaled_long_edge: int | None,
163 +) -> torch.Tensor:
164 + if patch_size <= 0:
165 + raise ValueError("patch_size must be greater than zero")
166 + _validate_image_rescale(
167 + rescale_image_frac,
168 + rescale_image_max_upscaled_long_edge,
169 + )
170 +
171 + from PIL import Image
172 +
173 + image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
174 + scaled_size = _scaled_image_dimensions(
175 + image.width,
176 + image.height,
177 + rescale_image_frac=rescale_image_frac,
178 + rescale_image_max_upscaled_long_edge=rescale_image_max_upscaled_long_edge,
179 + )
180 + if scaled_size != image.size:
181 + image = image.resize(scaled_size, resample=Image.Resampling.LANCZOS)
182 + arr = np.array(image, dtype=np.uint8, copy=True)
183 + height, width, _ = arr.shape
184 +
185 + nph = (height + patch_size - 1) // patch_size
186 + npw = width // patch_size + 1
187 + num_patches = nph * npw
188 +
189 + patches = np.empty((num_patches, patch_size, patch_size, 3), dtype=np.float32)
190 + _fill_patches_numba(arr, patch_size, patches, IMAGE_MEAN, IMAGE_STD, PAD_NORM)
191 +
192 + return (
193 + torch.from_numpy(patches)
194 + .to(torch.bfloat16)
195 + .view(num_patches, 1, patch_size, patch_size, 3)
196 + .expand(num_patches, 2, patch_size, patch_size, 3)
197 + )
198 +
199 +
200 +class InklingImageProcessor(BaseImageProcessor):
201 + r"""Turn raw images into ``vision_patches_bthwc`` for Inkling hMLP.
202 +
203 + ``rescale_image_frac`` scales the long edge while preserving aspect ratio.
204 + ``rescale_image_max_upscaled_long_edge`` optionally caps only upscaling and
205 + therefore requires a scale factor greater than one. The defaults, ``2.0`` and
206 + ``2048``, grow images toward a 2048-pixel long edge by at most 2x, while leaving
207 + images already at or above 2048 unchanged.
208 + """
209 +
210 + model_input_names = ["vision_patches_bthwc"]
211 +
212 + def __init__(
213 + self,
214 + patch_size: int = 40,
215 + rescale_image_frac: float | None = 2.0,
216 + rescale_image_max_upscaled_long_edge: int | None = 2048,
217 + **kwargs,
218 + ):
219 + if patch_size <= 0:
220 + raise ValueError("patch_size must be greater than zero")
221 + _validate_image_rescale(
222 + rescale_image_frac,
223 + rescale_image_max_upscaled_long_edge,
224 + )
225 + super().__init__(**kwargs)
226 + self.patch_size = patch_size
227 + self.rescale_image_frac = rescale_image_frac
228 + self.rescale_image_max_upscaled_long_edge = rescale_image_max_upscaled_long_edge
229 +
230 + def _encode_one(self, image) -> torch.Tensor:
231 + return _encode_image_bytes(
232 + _load_image_bytes(image),
233 + patch_size=self.patch_size,
234 + rescale_image_frac=self.rescale_image_frac,
235 + rescale_image_max_upscaled_long_edge=self.rescale_image_max_upscaled_long_edge,
236 + )
237 +
238 + def preprocess(
239 + self,
240 + images: ImageInput | list,
241 + return_tensors: str | None = "pt",
242 + **kwargs,
243 + ) -> BatchFeature:
244 + del return_tensors, kwargs
245 + if not isinstance(images, (list, tuple)):
246 + images = [images]
247 +
248 + per_image_patches: list[torch.Tensor] = []
249 + num_patches: list[int] = []
250 + num_tokens: list[int] = []
251 + for img in images:
252 + vp = self._encode_one(img)
253 + n_patches = int(vp.shape[0])
254 + per_image_patches.append(vp)
255 + num_patches.append(n_patches)
256 + num_tokens.append(n_patches)
257 +
258 + if len(per_image_patches) == 1:
259 + vision_patches_bthwc = per_image_patches[0]
260 + elif per_image_patches:
261 + vision_patches_bthwc = torch.cat(per_image_patches, dim=0)
262 + else:
263 + vision_patches_bthwc = torch.empty(0)
264 +
265 + data = {
266 + "vision_patches_bthwc": vision_patches_bthwc,
267 + "num_patches": num_patches,
268 + "num_tokens": num_tokens,
269 + }
270 + return BatchFeature(data=data, tensor_type=None)
271 +
272 +
273 +# ---------------------------------------------------------------------------
274 +# Audio feature extraction
275 +# ---------------------------------------------------------------------------
276 +
277 +
278 +@dataclass
279 +class InklingAudioEncoderParams:
280 + """Audio preprocessing parameters used to convert raw audio into dMel bins."""
281 +
282 + sample_rate: int = 16_000
283 + window_size_multiplier: float = 2.0
284 + n_fft: int | None = None
285 + n_mels: int = 80
286 + num_dmel_bins: int = 16
287 + dmel_min_value: float = -7.0
288 + dmel_max_value: float = 2.0
289 + audio_token_duration_s: float = 0.05
290 +
291 +
292 +def _to_exact_int(value: float, name: str, tolerance: float = 1e-6) -> int:
293 + rounded = round(value)
294 + if abs(value - rounded) > tolerance:
295 + raise ValueError(f"{name} must resolve to an integer sample count, got {value}")
296 + return int(rounded)
297 +
298 +
299 +def _hz_to_mel(frequencies: np.ndarray) -> np.ndarray:
300 + """Slaney mel scale, matching the librosa/torchaudio convention."""
301 + frequencies = np.asarray(frequencies, dtype=np.float64)
302 + f_sp = 200.0 / 3.0
303 + min_log_hz = 1000.0
304 + min_log_mel = min_log_hz / f_sp
305 + logstep = np.log(6.4) / 27.0
306 + linear = frequencies / f_sp
307 + log = (
308 + min_log_mel + np.log(np.maximum(frequencies, min_log_hz) / min_log_hz) / logstep
309 + )
310 + return np.where(frequencies >= min_log_hz, log, linear)
311 +
312 +
313 +def _mel_to_hz(mels: np.ndarray) -> np.ndarray:
314 + mels = np.asarray(mels, dtype=np.float64)
315 + f_sp = 200.0 / 3.0
316 + min_log_hz = 1000.0
317 + min_log_mel = min_log_hz / f_sp
318 + logstep = np.log(6.4) / 27.0
319 + linear = mels * f_sp
320 + log = min_log_hz * np.exp(logstep * (mels - min_log_mel))
321 + return np.where(mels >= min_log_mel, log, linear)
322 +
323 +
324 +_MEL_BASIS_CACHE: dict[tuple[int, int, int], torch.Tensor] = {}
325 +
326 +
327 +def _mel_basis(sample_rate: int, n_fft: int, n_mels: int) -> torch.Tensor:
328 + key = (sample_rate, n_fft, n_mels)
329 + cached = _MEL_BASIS_CACHE.get(key)
330 + if cached is not None:
331 + return cached
332 +
333 + fft_bins = n_fft // 2 + 1
334 + fft_freqs = np.arange(fft_bins, dtype=np.float64) * sample_rate / n_fft
335 + mel_edges = _mel_to_hz(
336 + np.linspace(
337 + _hz_to_mel(np.array([0.0]))[0],
338 + _hz_to_mel(np.array([sample_rate / 2.0]))[0],
339 + n_mels + 2,
340 + dtype=np.float64,
341 + )
342 + )
343 + mel_widths = np.diff(mel_edges)
344 + lower = (fft_freqs[None, :] - mel_edges[:-2, None]) / mel_widths[:-1, None]
345 + upper = (mel_edges[2:, None] - fft_freqs[None, :]) / mel_widths[1:, None]
346 + weights = np.maximum(0.0, np.minimum(lower, upper))
347 +
348 + # Slaney area normalization.
349 + weights *= (2.0 / (mel_edges[2:] - mel_edges[:-2]))[:, None]
350 + basis = torch.from_numpy(weights.astype(np.float32, copy=False)).contiguous()
351 + _MEL_BASIS_CACHE[key] = basis
352 + return basis
353 +
354 +
355 +def _dmel_bins(audio: torch.Tensor, params: InklingAudioEncoderParams) -> torch.Tensor:
356 + hop_length = _to_exact_int(
357 + params.audio_token_duration_s * params.sample_rate,
358 + "audio_token_duration_s * sample_rate",
359 + )
360 + window_size = _to_exact_int(
361 + params.audio_token_duration_s
362 + * params.window_size_multiplier
363 + * params.sample_rate,
364 + "audio_token_duration_s * window_size_multiplier * sample_rate",
365 + )
366 + n_fft = params.n_fft or window_size
367 + if hop_length <= 0 or window_size <= 0 or n_fft <= 0:
368 + raise ValueError("audio hop length, window size, and n_fft must be positive")
369 + if audio.numel() == 0:
370 + return torch.empty((0, params.n_mels), dtype=torch.int32)
371 +
372 + right_pad = math.ceil(audio.numel() / hop_length) * hop_length - audio.numel()
373 + left_pad = max(n_fft - hop_length, 0)
374 + audio = F.pad(audio, (left_pad, right_pad))
375 +
376 + window = torch.hann_window(window_size, periodic=True, dtype=torch.float32)
377 + spec = torch.stft(
378 + audio.unsqueeze(0),
379 + n_fft=n_fft,
380 + hop_length=hop_length,
381 + win_length=window_size,
382 + window=window,
383 + center=False,
384 + normalized=False,
385 + onesided=True,
386 + return_complex=True,
387 + )
388 + spec_ri = torch.view_as_real(spec)
389 + magnitude = (
390 + (spec_ri[..., 0].square() + spec_ri[..., 1].square())
391 + .clamp_min(1e-10)
392 + .sqrt()
393 + .squeeze(0)
394 + )
395 +
396 + mel = (
397 + _mel_basis(params.sample_rate, n_fft, params.n_mels)
398 + .matmul(magnitude)
399 + .clamp_min(1e-10)
400 + .log10()
401 + )
402 + mel = mel.to(torch.float64).clamp(
403 + min=params.dmel_min_value, max=params.dmel_max_value
404 + )
405 + bin_centers = torch.linspace(
406 + params.dmel_min_value,
407 + params.dmel_max_value,
408 + params.num_dmel_bins,
409 + dtype=torch.float64,
410 + )
411 + dmel_bins = (mel.unsqueeze(-1) - bin_centers).abs().argmin(dim=-1)
412 + return dmel_bins.to(torch.int32).T.contiguous()
413 +
414 +
415 +class InklingAudioFeatureExtractor(FeatureExtractionMixin):
416 + """Convert raw audio into Inkling dMel bins in the HF feature-extractor API."""
417 +
418 + model_input_names = ["dmel_bins"]
419 +
420 + def __init__(self, params: dict | None = None, **kwargs):
421 + super().__init__(**kwargs)
422 + merged = InklingAudioEncoderParams()
423 + if params:
424 + for k, v in params.items():
425 + if hasattr(merged, k):
426 + setattr(merged, k, v)
427 + # also accept flat kwargs (HF config style)
428 + for k in list(kwargs.keys()):
429 + if hasattr(merged, k):
430 + setattr(merged, k, kwargs[k])
431 + self.params = merged
432 +
433 + def _decode_one(self, audio) -> torch.Tensor:
434 + # vLLM hands the feature extractor numpy arrays (the dummy-input builder
435 + # during profiling, and MultiModalDataParser after resampling to the
436 + # target sample rate), so no other input types need to be supported.
437 + return torch.from_numpy(
438 + np.ascontiguousarray(audio.astype(np.float32, copy=False))
439 + ).flatten()
440 +
441 + def _encode_one(self, audio) -> torch.Tensor:
442 + return _dmel_bins(self._decode_one(audio), self.params)
443 +
444 + def __call__(
445 + self,
446 + audios: Sequence | None,
447 + return_tensors: str | None = None,
448 + **kwargs,
449 + ) -> BatchFeature:
450 + del return_tensors, kwargs
451 + if audios is None:
452 + audios = []
453 + if not isinstance(audios, (list, tuple)):
454 + audios = [audios]
455 +
456 + dmel_bins = [self._encode_one(a) for a in audios]
457 + data = {
458 + # per-clip feature: dmel bins as float32 [T, n_mels]
459 + "dmel_bins": [bins.to(torch.float32) for bins in dmel_bins],
460 + "num_audio_tokens": [int(bins.shape[0]) for bins in dmel_bins],
461 + }
462 + # return_tensors intentionally ignored: per-clip features have ragged T.
463 + return BatchFeature(data=data, tensor_type=None)
464 +
465 +
466 +# ---------------------------------------------------------------------------
467 +# Composite processor
468 +# ---------------------------------------------------------------------------
469 +
470 +
471 +class InklingProcessor:
472 + """Bundle Inkling image + audio preprocessing with the MM token ids."""
473 +
474 + def __init__(
475 + self,
476 + image_processor: InklingImageProcessor | None = None,
477 + audio_feature_extractor: InklingAudioFeatureExtractor | None = None,
478 + tokenizer=None,
479 + ):
480 + self.image_processor = image_processor or InklingImageProcessor()
481 + self.audio_feature_extractor = (
482 + audio_feature_extractor or InklingAudioFeatureExtractor()
483 + )
484 + self.tokenizer = tokenizer
485 +
486 + def process_images(self, images: list):
487 + """Raw images -> BatchFeature(vision_patches_bthwc, num_patches, num_tokens)."""
488 + return self.image_processor.preprocess(images, return_tensors="pt")
489 +
490 + def process_audios(self, audios: list):
491 + """Raw audios -> BatchFeature(dmel_bins, num_audio_tokens)."""
492 + return self.audio_feature_extractor(audios)
493 +
494 +
495 +__all__ = [
496 + "InklingImageProcessor",
497 + "InklingAudioFeatureExtractor",
498 + "InklingAudioEncoderParams",
499 + "InklingProcessor",
500 + "IMAGE_MARKER_ID",
501 + "AUDIO_MARKER_ID",
502 + "IMAGE_TOKEN_ID",
503 + "AUDIO_TOKEN_ID",
504 +]