@@ -302,3 +302,135 @@ def test_attach_lumid_cfg_not_injected_for_non_s3_list_inputs() -> None:
302302 ds = llm_node .data_spec
303303 assert "lumid_cfg" not in ds
304304 assert "s3_cfg" not in ds
305+
306+
307+ def test_runtime_graph_propagates_extended_sampler_fields () -> None :
308+ stock = input_placeholder ("Stock" )
309+ llm = LLMChatOp (
310+ [OpMessage (role = "user" , content = stock )],
311+ config = GenerationConfig (
312+ model = "meta-llama/Llama-3.1-8B-Instruct" ,
313+ max_tokens = 256 ,
314+ min_tokens = 4 ,
315+ temperature = 0.7 ,
316+ top_p = 0.8 ,
317+ top_k = 20 ,
318+ min_p = 0.0 ,
319+ presence_penalty = 1.5 ,
320+ frequency_penalty = 0.5 ,
321+ repetition_penalty = 1.0 ,
322+ seed = 42 ,
323+ chat_template_kwargs = {"enable_thinking" : False },
324+ extra_sampling_params = {"length_penalty" : 1.1 },
325+ ),
326+ )
327+ output = as_output ("result" , llm )
328+ compiled = Graph .from_ops ([output ]).compile (Stock = ["NVDA" ])
329+ runtime_graph = RuntimeGraphBuilder ().build (compiled )
330+
331+ spec = runtime_graph .nodes [llm .id ].inference_spec
332+ assert spec ["max_tokens" ] == 256
333+ assert spec ["min_tokens" ] == 4
334+ assert spec ["temperature" ] == 0.7
335+ assert spec ["top_p" ] == 0.8
336+ assert spec ["top_k" ] == 20
337+ assert spec ["min_p" ] == 0.0
338+ assert spec ["presence_penalty" ] == 1.5
339+ assert spec ["frequency_penalty" ] == 0.5
340+ assert spec ["repetition_penalty" ] == 1.0
341+ assert spec ["seed" ] == 42
342+ assert spec ["chat_template_kwargs" ] == {"enable_thinking" : False }
343+ # extra_sampling_params keys are merged into the spec as flat keys.
344+ assert spec ["length_penalty" ] == 1.1
345+ # model never appears in inference_spec — it's carried separately.
346+ assert "model" not in spec
347+ # The escape-hatch wrapper key itself is not leaked into the spec.
348+ assert "extra_sampling_params" not in spec
349+
350+
351+ def _build_llm_graph (** config_kwargs ):
352+ stock = input_placeholder ("Stock" )
353+ cfg = GenerationConfig (model = "meta-llama/Llama-3.1-8B-Instruct" , ** config_kwargs )
354+ llm = LLMChatOp ([OpMessage (role = "user" , content = stock )], config = cfg )
355+ output = as_output ("result" , llm )
356+ compiled = Graph .from_ops ([output ]).compile (Stock = ["NVDA" ])
357+ return RuntimeGraphBuilder ().build (compiled ), llm .id
358+
359+
360+ _ENGINE_FIELDS = (
361+ "max_model_len" ,
362+ "gpu_memory_utilization" ,
363+ "tensor_parallel_size" ,
364+ "dtype" ,
365+ )
366+
367+
368+ def test_engine_overlay_lands_in_vllm_backend () -> None :
369+ runtime_graph , llm_id = _build_llm_graph (
370+ max_model_len = 8192 ,
371+ gpu_memory_utilization = 0.75 ,
372+ tensor_parallel_size = 2 ,
373+ dtype = "bf16" ,
374+ extra_engine_kwargs = {"quantization" : "fp8" , "kv_cache_dtype" : "fp8" },
375+ )
376+ node = runtime_graph .nodes [llm_id ]
377+ vllm_cfg = node .model_spec ["vllm" ]
378+ assert vllm_cfg ["max_model_len" ] == 8192
379+ assert vllm_cfg ["gpu_memory_utilization" ] == 0.75
380+ assert vllm_cfg ["tensor_parallel_size" ] == 2
381+ assert vllm_cfg ["dtype" ] == "bf16"
382+ assert vllm_cfg ["quantization" ] == "fp8"
383+ assert vllm_cfg ["kv_cache_dtype" ] == "fp8"
384+ # Engine fields must NOT leak into inference_spec.
385+ for k in _ENGINE_FIELDS + ("quantization" , "kv_cache_dtype" ):
386+ assert k not in node .inference_spec
387+
388+
389+ def test_engine_overlay_omitted_when_unset () -> None :
390+ # gpu_memory_utilization has an env-var default in the baseline vLLM
391+ # config, so unset means "stays at env default" not "absent". The other
392+ # three fields have no default and should not appear when unset.
393+ runtime_graph , llm_id = _build_llm_graph ()
394+ vllm_cfg = runtime_graph .nodes [llm_id ].model_spec ["vllm" ]
395+ for k in ("max_model_len" , "tensor_parallel_size" , "dtype" ):
396+ assert k not in vllm_cfg
397+
398+
399+ def test_engine_overlay_overrides_baseline_default () -> None :
400+ # gpu_memory_utilization is set in the baseline config; an explicit
401+ # value on GenerationConfig must override it.
402+ runtime_graph , llm_id = _build_llm_graph (gpu_memory_utilization = 0.6 )
403+ vllm_cfg = runtime_graph .nodes [llm_id ].model_spec ["vllm" ]
404+ assert vllm_cfg ["gpu_memory_utilization" ] == 0.6
405+
406+
407+ def test_extra_sampling_params_conflict_rejected () -> None :
408+ stock = input_placeholder ("Stock" )
409+ llm = LLMChatOp (
410+ [OpMessage (role = "user" , content = stock )],
411+ config = GenerationConfig (
412+ model = "meta-llama/Llama-3.1-8B-Instruct" ,
413+ temperature = 0.7 ,
414+ extra_sampling_params = {"temperature" : 0.0 },
415+ ),
416+ )
417+ output = as_output ("result" , llm )
418+ compiled = Graph .from_ops ([output ]).compile (Stock = ["NVDA" ])
419+ with pytest .raises (ValueError , match = "extra_sampling_params conflict" ):
420+ RuntimeGraphBuilder ().build (compiled )
421+
422+
423+ def test_extra_engine_kwargs_conflict_rejected () -> None :
424+ stock = input_placeholder ("Stock" )
425+ llm = LLMChatOp (
426+ [OpMessage (role = "user" , content = stock )],
427+ config = GenerationConfig (
428+ model = "meta-llama/Llama-3.1-8B-Instruct" ,
429+ max_model_len = 8192 ,
430+ extra_engine_kwargs = {"max_model_len" : 4096 },
431+ ),
432+ )
433+ output = as_output ("result" , llm )
434+ compiled = Graph .from_ops ([output ]).compile (Stock = ["NVDA" ])
435+ with pytest .raises (ValueError , match = "extra_engine_kwargs conflict" ):
436+ RuntimeGraphBuilder ().build (compiled )
0 commit comments