From 8c1557a79c539ffe82d004d2a0c8d7b5e71159ce Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 16 Sep 2026 17:44:29 +0100 Subject: [PATCH] Add `pydocstyle` to the `ruff` rules (#52136) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Harry Mellor Co-authored-by: Claude Opus 5 (1M context) --- .../test_lm_eval_correctness.py | 3 +- .../scripts/compare-json-results.py | 6 +- .../convert-results-json-to-markdown.py | 3 +- .buildkite/scripts/generate-nightly-index.py | 21 +- benchmarks/attention_benchmarks/batch_spec.py | 28 +- benchmarks/attention_benchmarks/benchmark.py | 20 +- benchmarks/attention_benchmarks/common.py | 15 +- benchmarks/attention_benchmarks/mla_runner.py | 51 +- benchmarks/attention_benchmarks/runner.py | 12 +- benchmarks/benchmark_batch_invariance.py | 6 +- benchmarks/benchmark_hash.py | 3 +- .../benchmark_hidden_state_extraction.py | 6 +- .../benchmark_long_document_qa_throughput.py | 11 +- benchmarks/benchmark_prefix_block_hash.py | 4 +- benchmarks/benchmark_prefix_caching.py | 3 +- .../benchmark_serving_structured_output.py | 5 +- benchmarks/benchmark_topk_topp.py | 9 +- .../merge_attn_states_benchmarks.py | 6 +- benchmarks/kernels/bench_concat_mla_q.py | 6 +- .../benchmark_2d_silu_mul_fp8_quant.py | 3 +- .../kernels/benchmark_cutlass_moe_fp8.py | 5 +- .../kernels/benchmark_cutlass_moe_nvfp4.py | 3 +- .../kernels/benchmark_device_communicators.py | 5 +- .../kernels/benchmark_fused_collective.py | 9 +- .../kernels/benchmark_kimi_k3_attn_res.py | 1 + .../kernels/benchmark_kimi_k3_kda_decode.py | 1 + .../benchmark_kimi_k3_latent_moe_tail.py | 1 + benchmarks/kernels/benchmark_lora.py | 55 +- benchmarks/kernels/benchmark_mla_k_concat.py | 3 +- benchmarks/kernels/benchmark_moe_defaults.py | 6 +- .../benchmark_per_token_group_quant.py | 2 +- .../kernels/benchmark_per_token_quant_fp8.py | 6 +- .../benchmark_rdna_hybrid_w4a16_gemm.py | 3 +- .../kernels/benchmark_reshape_and_cache.py | 1 - .../benchmark_reshape_and_cache_flash.py | 1 - .../benchmark_selective_state_update.py | 9 +- .../kernels/benchmark_silu_mul_fp8_quant.py | 7 +- .../kernels/benchmark_vit_aiter_fp8_attn.py | 1 + .../benchmark_vocab_parallel_embedding.py | 3 +- .../kernels/benchmark_w8a8_block_fp8.py | 2 + benchmarks/kernels/ir/bench_ir_ops.py | 3 +- benchmarks/kernels/ir/shapes.py | 4 +- benchmarks/kernels/utils.py | 3 +- benchmarks/multi_turn/bench_dataset.py | 3 +- .../multi_turn/convert_sharegpt_to_openai.py | 3 +- benchmarks/replayssm/e2e_decode_speedup.py | 1 + csrc/cpu/generate_cpu_attn_dispatch.py | 4 +- docs/mkdocs/gen_files/generate_argparse.py | 6 +- .../gen_files/generate_attention_backends.py | 6 +- docs/mkdocs/gen_files/generate_examples.py | 14 +- docs/mkdocs/gen_files/generated_content.py | 1 + docs/mkdocs/hooks/autoref_code.py | 3 +- docs/mkdocs/hooks/url_schemes.py | 11 +- examples/applications/api_server/server.py | 3 +- .../streamlit_openai_chatbot_webserver.py | 7 +- ...val_augmented_generation_with_langchain.py | 37 +- ...al_augmented_generation_with_llamaindex.py | 18 +- examples/deployment/async_llm_streaming.py | 6 +- examples/deployment/llm_engine_example.py | 3 +- .../disaggregated_encoder/disagg_epd_proxy.py | 20 +- .../disagg_proxy_demo.py | 3 +- .../disagg_proxy_multiturn.py | 6 +- .../disagg_proxy_pushconnector_demo.py | 3 +- .../example_connector/decode_example.py | 2 +- .../flexkv_connector/prefix_caching_flexkv.py | 3 +- .../decode_example.py | 2 +- .../load_recovery_example_connector.py | 3 +- .../lmcache/cpu_offload_lmcache.py | 3 +- .../disagg_proxy_server.py | 12 +- .../lmcache/kv_cache_sharing_lmcache_v1.py | 3 +- .../mooncake_connector_proxy.py | 18 +- .../automatic_prefix_caching_offline.py | 3 +- .../reproducibility_offline.py | 3 +- .../context_extension_offline.py | 3 +- .../data_parallel/data_parallel_offline.py | 3 +- .../kv_events/kv_events_subscriber.py | 2 +- examples/features/logits_processor/custom.py | 2 +- .../features/logits_processor/custom_req.py | 9 +- .../logits_processor/custom_req_init.py | 9 +- .../lora/lora_with_quantization_offline.py | 5 +- examples/features/lora/multilora_offline.py | 3 +- .../data_parallel_pause_resume.py | 3 +- .../prompt_embed/prompt_embed_offline.py | 3 +- .../features/reset_kv/reset_kv_offline.py | 3 +- .../load_sharded_state_offline.py | 3 +- .../save_sharded_state_offline.py | 3 +- .../mlpspeculator_offline.py | 3 +- .../structured_outputs_offline.py | 3 +- .../torchrun/torchrun_dp_example_offline.py | 3 +- .../torchrun/torchrun_example_offline.py | 3 +- .../multimodal/audio_language_offline.py | 6 +- .../encoder_decoder_multimodal_offline.py | 6 +- ...i_chat_completion_client_for_multimodal.py | 2 - .../multimodal/qwen2_5_omni/only_thinker.py | 3 +- .../multimodal/qwen3_omni/only_thinker.py | 3 +- .../vision_language_multi_image_offline.py | 7 +- .../multimodal/vision_language_offline.py | 9 +- examples/generate/trace_replay_offline.py | 3 +- .../openai_embedding_long_text/client.py | 3 +- .../pooling/embed/vision_embedding_offline.py | 3 +- .../pooling/embed/vision_embedding_online.py | 22 +- .../plugin/prithvi_geospatial_mae_offline.py | 13 +- .../pooling/reward/sequence_reward_offline.py | 3 +- .../pooling/reward/sequence_reward_online.py | 3 +- .../pooling/reward/token_reward_offline.py | 3 +- .../pooling/reward/token_reward_online.py | 3 +- .../pooling/score/cohere_rerank_client.py | 3 +- .../pooling/score/colbert_rerank_online.py | 3 +- .../score/colmodernvbert_rerank_online.py | 3 +- .../pooling/score/colqwen3_5_rerank_online.py | 3 +- .../pooling/score/colqwen3_rerank_online.py | 3 +- .../pooling/score/convert_model_to_seq_cls.py | 15 +- .../pooling/score/qwen3_reranker_offline.py | 7 +- .../pooling/score/qwen3_reranker_online.py | 3 +- examples/pooling/score/rerank_api_online.py | 3 +- examples/pooling/score/score_api_online.py | 3 +- .../pooling/score/using_template_offline.py | 5 +- .../pooling/score/using_template_online.py | 6 +- .../pooling/score/vision_rerank_api_online.py | 3 +- .../pooling/score/vision_reranker_offline.py | 3 +- .../pooling/score/vision_score_api_online.py | 3 +- .../forced_alignment_offline.py | 3 +- .../token_classify/forced_alignment_online.py | 3 +- examples/pooling/token_classify/ner_online.py | 3 +- .../colqwen3_token_embed_online.py | 3 +- .../token_embed/jina_reranker_v3_online.py | 3 +- .../multi_vector_retrieval_online.py | 3 +- examples/ray_serving/batch_llm_inference.py | 3 +- examples/ray_serving/ray_serve_deepseek.py | 3 +- ...at_completion_tool_calls_with_reasoning.py | 3 +- .../openai_chat_completion_with_reasoning.py | 3 +- ...hat_completion_with_reasoning_streaming.py | 3 +- examples/reasoning/openai_responses_client.py | 3 +- examples/rl/rlhf_async_new_apis.py | 3 +- examples/rl/rlhf_http_ipc.py | 3 +- examples/rl/rlhf_http_nccl.py | 3 +- examples/rl/rlhf_ipc_fsdp_ep.py | 3 +- examples/rl/rlhf_nccl_fsdp_ep.py | 6 +- examples/rl/rlhf_sharded_rdt_small_ep.py | 3 +- examples/rl/routed_experts_e2e.py | 3 +- .../openai/openai_transcription_client.py | 14 +- .../realtime/openai_realtime_client.py | 11 +- .../openai_realtime_microphone_client.py | 3 +- ...penai_chat_completion_client_with_tools.py | 3 +- ...t_completion_client_with_tools_required.py | 3 +- ..._chat_completion_client_with_tools_xlam.py | 8 +- ...letion_client_with_tools_xlam_streaming.py | 7 +- .../openai_responses_client_with_mcp_tools.py | 3 +- .../openai_responses_client_with_tools.py | 6 +- pyproject.toml | 11 + scripts/autotune_helion_kernels.py | 3 +- scripts/benchmark_helion_kernels.py | 4 +- setup.py | 10 +- .../test_basic_correctness.py | 2 +- tests/basic_correctness/test_mem.py | 6 +- tests/benchmarks/test_bfcl_dataset.py | 18 +- ...est_custom_dataset_chat_template_kwargs.py | 3 +- .../test_endpoint_request_func_timing.py | 2 +- tests/benchmarks/test_sampling_params.py | 9 +- tests/benchmarks/test_skip_tokenizer_init.py | 3 +- tests/ci_envs.py | 4 +- tests/compile/backend.py | 6 +- tests/compile/conftest.py | 6 +- .../correctness_e2e/test_sequence_parallel.py | 11 +- .../compile/fullgraph/test_full_cudagraph.py | 10 +- .../compile/fullgraph/test_multiple_graphs.py | 3 +- tests/compile/fullgraph/test_simple.py | 6 +- tests/compile/fullgraph/test_toy_llama.py | 3 +- .../passes/distributed/test_async_tp.py | 24 +- tests/compile/passes/ir/test_clone_cleanup.py | 3 +- .../ir/test_inplace_functionalization.py | 9 +- .../test_double_aiter_rms_quant_fusion.py | 6 +- .../passes/test_fuse_mla_dual_rms_norm.py | 29 +- tests/compile/passes/test_fusion.py | 3 +- tests/compile/passes/test_fusion_attn.py | 3 +- .../passes/test_mla_attn_quant_fusion.py | 3 +- .../test_silu_mul_quant_manual_fusion.py | 3 +- tests/compile/passes/test_split_coalescing.py | 3 +- tests/compile/silly_attention.py | 12 +- tests/compile/test_aot_compile.py | 19 +- tests/compile/test_codegen.py | 30 +- tests/compile/test_compile_ranges.py | 6 +- tests/compile/test_config.py | 9 +- .../test_dynamic_shapes_compilation.py | 3 +- tests/compile/test_graph_partition.py | 43 +- ...test_inductor_fallback_allow_list_patch.py | 3 +- tests/compile/test_structured_logging.py | 1 + tests/config/test_config_utils.py | 2 +- tests/config/test_model_arch_config.py | 6 +- tests/config/test_mp_reducer.py | 4 +- tests/config/test_multimodal_config.py | 15 +- .../test_speculative_draft_hf_overrides.py | 12 +- ...eculative_draft_max_position_embeddings.py | 3 +- tests/cuda/test_cuda_context.py | 4 +- tests/cuda/test_platform_no_cuda_init.py | 3 +- ...stop_string_while_stop_model_terminates.py | 4 +- tests/distributed/conftest.py | 21 +- tests/distributed/test_context_parallel.py | 11 +- tests/distributed/test_eplb_algo.py | 24 +- tests/distributed/test_eplb_events.py | 5 +- tests/distributed/test_eplb_execute.py | 19 +- .../test_eplb_quant_scale_consistency.py | 12 +- tests/distributed/test_eplb_spec_decode.py | 8 +- tests/distributed/test_eplb_utils.py | 6 +- tests/distributed/test_events.py | 30 +- tests/distributed/test_expert_placement.py | 3 - tests/distributed/test_mnnvl_alltoall.py | 11 +- tests/distributed/test_mq_connect_ip.py | 12 +- tests/distributed/test_multiproc_executor.py | 6 +- tests/distributed/test_pipeline_parallel.py | 11 +- tests/distributed/test_pynccl.py | 3 +- tests/distributed/test_quick_all_reduce.py | 3 +- tests/distributed/test_ray_v2_executor.py | 4 +- tests/distributed/test_ray_v2_executor_e2e.py | 10 +- tests/distributed/test_sharded_rdt_plan.py | 189 +- .../distributed/test_sharded_rdt_producer.py | 119 +- tests/distributed/test_sharded_rdt_trainer.py | 36 +- tests/distributed/test_shm_broadcast.py | 14 +- tests/distributed/test_shm_buffer.py | 18 +- tests/distributed/test_utils.py | 3 +- tests/distributed/test_weight_transfer.py | 66 +- .../test_weight_transfer_nccl_uid.py | 9 +- tests/engine/test_arg_utils.py | 22 +- .../test_anthropic_messages_conversion.py | 34 +- .../cohere/test_serving_streaming.py | 6 +- .../test_api_server_process_manager.py | 12 +- tests/entrypoints/launchers/test_cli_args.py | 32 +- .../launchers/test_dp_supervisor.py | 24 +- tests/entrypoints/launchers/test_run_batch.py | 11 +- tests/entrypoints/launchers/test_shutdown.py | 1 - .../llm/offline_mode/test_offline_mode.py | 2 +- tests/entrypoints/llm/test_accuracy.py | 6 +- tests/entrypoints/llm/test_chat.py | 9 +- .../llm/test_mm_cache_external_injection.py | 3 +- .../chat_completion/test_audio_in_video.py | 9 +- ...chat_completion_with_mixed_audio_embeds.py | 3 +- ...chat_completion_with_mixed_image_embeds.py | 3 +- .../openai/chat_completion/test_vision.py | 9 +- .../test_batched_chat_completions.py | 2 +- .../openai/chat_completion/test_chat_error.py | 10 +- .../test_chat_logit_bias_validation.py | 9 +- .../test_completion_with_function_calling.py | 2 +- .../test_extra_content_fields.py | 10 +- .../chat_completion/test_logprob_token_ids.py | 3 +- .../chat_completion/test_serving_chat.py | 17 +- .../test_thinking_token_budget.py | 7 +- .../openai/completion/test_completion.py | 2 - .../completion/test_completion_error.py | 13 +- .../completion/test_token_in_token_out.py | 3 +- .../openai/correctness/test_lmeval.py | 5 +- .../parser/test_harmony_render_parity.py | 12 +- .../openai/parser/test_harmony_utils.py | 11 +- .../openai/responses/test_errors.py | 4 +- .../openai/responses/test_function_call.py | 3 +- .../responses/test_function_call_parsing.py | 12 +- .../openai/responses/test_harmony_utils.py | 3 +- .../test_response_input_to_harmony.py | 12 +- .../openai/responses/test_responses_utils.py | 9 +- .../responses/test_serving_responses.py | 50 +- .../openai/responses/test_simple.py | 3 +- .../openai/responses/test_streaming_events.py | 6 +- .../entrypoints/openai/test_openai_schema.py | 3 +- .../openai/test_render_token_offsets.py | 9 +- .../openai/test_return_token_ids.py | 3 +- .../entrypoints/openai/test_stop_token_ids.py | 3 +- tests/entrypoints/openai/utils.py | 3 +- .../pooling/basic/test_tiling_engine.py | 12 +- .../pooling/embed/test_cohere_online.py | 6 +- .../embed/test_cohere_openai_parity.py | 3 +- .../pooling/embed/test_io_processor.py | 3 +- .../pooling/embed/test_online_dimensions.py | 4 +- .../pooling/embed/test_online_long_text.py | 17 +- .../pooling/embed/test_protocol.py | 3 +- .../scale_out/derender/test_derender.py | 5 +- .../derender/test_derender_parity.py | 9 +- .../derender/test_derender_stream.py | 3 - .../scale_out/render/test_render.py | 13 +- .../render/test_render_multimodal.py | 3 - .../token_in_token_out/test_mm_serde.py | 3 +- .../token_in_token_out/test_protocol.py | 9 +- tests/entrypoints/serve/dev/rlhf/conftest.py | 4 +- .../serve/dev/rpc/test_collective_rpc.py | 8 +- tests/entrypoints/serve/dev/test_sleep.py | 3 +- .../test_error_sanitization.py | 6 +- .../test_validation_exception_handler.py | 15 +- .../serve/instrumentator/test_basic.py | 2 +- .../serve/instrumentator/test_metrics.py | 3 +- .../serve/lora/test_lora_adapters.py | 3 +- .../test_authentication_middleware.py | 9 +- .../middleware/test_optional_middleware.py | 3 +- .../test_sagemaker_handler_overrides.py | 18 +- .../serve/tokenize/test_tokenization_vlm.py | 3 +- .../tokenize/test_tokenize_then_chat_vlm.py | 7 +- .../entrypoints/serve/utils/test_api_utils.py | 12 +- .../serve/utils/test_request_logger.py | 3 +- .../serve/utils/test_sse_keep_alive.py | 24 +- tests/entrypoints/speech_to_text/conftest.py | 1 + .../test_transcription_api_correctness.py | 3 +- .../realtime/test_realtime_validation.py | 3 +- .../test_chunk_timestamp_offset.py | 5 +- .../test_transcription_validation_whisper.py | 3 +- .../tool_parsers/test_hermes_tool_parser.py | 4 - .../tool_parsers/test_openai_tool_parser.py | 11 +- .../entrypoints/unit_tests/test_chat_utils.py | 3 - tests/entrypoints/unit_tests/test_context.py | 22 +- .../test_weight_transfer_llm.py | 3 +- tests/evals/gpt_oss/conftest.py | 4 +- tests/evals/gpt_oss/test_gpqa_correctness.py | 4 +- tests/evals/gsm8k/gsm8k_eval.py | 10 +- tests/evals/gsm8k/test_gsm8k_correctness.py | 3 +- tests/evals/gsm8k/test_gsm8k_offloading.py | 3 +- tests/evals/mrcr/test_mrcr_correctness.py | 3 +- .../fusion/test_quant_activation_contract.py | 3 +- tests/ir/ir_test_utils.py | 4 +- tests/ir/test_op.py | 4 +- tests/kernels/attention/test_amx_mla.py | 9 +- .../attention/test_attention_selector.py | 6 +- .../attention/test_cutlass_mla_decode.py | 3 +- tests/kernels/attention/test_flashinfer.py | 4 +- .../attention/test_flashinfer_mla_decode.py | 2 +- .../test_flashinfer_trtllm_attention.py | 3 +- .../test_kimi_k3_mla_fused_epilogue.py | 3 +- .../kernels/attention/test_lightning_attn.py | 4 +- .../attention/test_merge_attn_states.py | 3 +- tests/kernels/attention/test_mha_attn.py | 10 +- tests/kernels/attention/test_minimax_m3.py | 42 +- ...test_mla_cross_layer_kernel_equivalence.py | 24 +- tests/kernels/attention/test_rocm_aiter_fa.py | 21 +- .../test_rocm_aiter_mla_op_registration.py | 2 +- ...est_rocm_aiter_mla_sparse_metadata_sync.py | 6 +- .../attention/test_triton_decode_attention.py | 6 +- .../test_triton_unified_attention_diffkv.py | 4 +- .../attention/test_trtllm_kvfp8_dequant.py | 9 +- .../kernels/attention/test_xpu_mla_sparse.py | 4 +- tests/kernels/core/test_apply_rotary_emb.py | 6 +- tests/kernels/core/test_fused_embed_norm.py | 6 +- tests/kernels/core/test_fused_q_kv_rmsnorm.py | 6 +- .../kernels/core/test_fused_rms_norm_gated.py | 3 +- tests/kernels/core/test_minimax_reduce_rms.py | 1 - tests/kernels/core/test_mrope.py | 2 +- tests/kernels/core/test_opcheck.py | 4 +- tests/kernels/core/test_rocm_aiter_ops.py | 3 +- tests/kernels/core/test_rotary_embedding.py | 4 +- .../test_rotary_embedding_mla_cache_fused.py | 4 +- .../core/test_vocab_parallel_embedding.py | 3 +- tests/kernels/helion/test_autotune.py | 3 +- tests/kernels/helion/test_config_manager.py | 3 +- ...test_dynamic_per_token_scaled_fp8_quant.py | 2 +- .../kernels/helion/test_fused_qk_norm_rope.py | 2 +- tests/kernels/helion/test_helion_available.py | 3 +- .../helion/test_per_token_group_fp8_quant.py | 2 +- tests/kernels/helion/test_register.py | 13 +- .../test_rms_norm_dynamic_per_token_quant.py | 2 +- .../helion/test_rms_norm_per_block_quant.py | 2 +- tests/kernels/helion/utils.py | 2 +- tests/kernels/ir/test_aiter_norm_dispatch.py | 2 +- tests/kernels/ir/test_ir_ops.py | 3 +- tests/kernels/ir/test_layernorm.py | 3 +- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py | 9 +- tests/kernels/mamba/test_causal_conv1d.py | 9 +- tests/kernels/mamba/test_mamba_ssm.py | 5 +- tests/kernels/mamba/test_mamba_ssm_configs.py | 12 +- tests/kernels/mamba/test_mamba_ssm_ssd.py | 5 +- ...ayssm_prefill_decode_equivalence_mamba2.py | 3 +- .../test_replayssm_standard_decode_mamba2.py | 6 +- tests/kernels/mamba/utils.py | 5 +- .../moe/modular_kernel_tools/common.py | 11 +- tests/kernels/moe/parallel_utils.py | 4 +- tests/kernels/moe/test_b12x.py | 3 +- tests/kernels/moe/test_batched_deepgemm.py | 1 - tests/kernels/moe/test_batched_moe.py | 9 +- tests/kernels/moe/test_block_int8.py | 6 +- .../moe/test_count_expert_num_tokens.py | 4 +- tests/kernels/moe/test_cpu_fused_moe.py | 7 +- tests/kernels/moe/test_cpu_int4_moe.py | 1 + tests/kernels/moe/test_cpu_quant_fused_moe.py | 1 + tests/kernels/moe/test_cutedsl_moe.py | 8 +- tests/kernels/moe/test_deepep_deepgemm_moe.py | 16 +- tests/kernels/moe/test_deepep_moe.py | 8 +- tests/kernels/moe/test_deepep_v2_moe.py | 3 +- tests/kernels/moe/test_deepgemm.py | 15 +- tests/kernels/moe/test_flashinfer.py | 3 +- .../moe/test_flashinfer_cutedsl_nvfp4_moe.py | 2 +- tests/kernels/moe/test_fused_topk.py | 2 +- .../moe/test_gpt_oss_triton_kernels.py | 3 +- tests/kernels/moe/test_grouped_topk.py | 2 +- .../moe/test_marlin_vs_trtllm_mxint4.py | 6 +- .../moe/test_modular_kernel_combinations.py | 3 +- .../moe/test_modular_oai_triton_moe.py | 4 +- tests/kernels/moe/test_moe.py | 11 +- .../kernels/moe/test_moe_align_block_size.py | 10 +- tests/kernels/moe/test_moe_kernel_oracle.py | 3 +- tests/kernels/moe/test_moe_layer.py | 10 +- .../kernels/moe/test_moe_permute_unpermute.py | 2 +- tests/kernels/moe/test_mxfp4_moe.py | 13 +- .../moe/test_mxfp8_aiter_backend_selection.py | 8 +- tests/kernels/moe/test_ocp_mx_moe.py | 10 +- tests/kernels/moe/test_rocm_aiter_moe.py | 33 +- .../test_routed_experts_capture_monolithic.py | 24 +- tests/kernels/moe/test_routing.py | 18 +- tests/kernels/moe/test_routing_simulator.py | 9 +- .../test_shared_fused_moe_routed_transform.py | 3 +- .../moe/test_silu_mul_fp8_quant_deep_gemm.py | 10 +- ..._mul_per_token_group_quant_fp8_colmajor.py | 7 +- tests/kernels/moe/test_triton_moe_ptpc_fp8.py | 7 +- tests/kernels/moe/test_trtllm_bf16_moe.py | 3 +- tests/kernels/moe/test_trtllm_nvfp4_moe.py | 3 +- .../moe/test_unquantized_backend_selection.py | 6 +- tests/kernels/moe/test_zero_expert_moe.py | 3 +- tests/kernels/moe/utils.py | 3 +- tests/kernels/quant_utils.py | 3 +- tests/kernels/quantization/nvfp4_utils.py | 1 + .../test_aiter_hipb_mm_linear_kernel.py | 3 +- .../quantization/test_cpu_fp8_scaled_mm.py | 1 + .../quantization/test_cutlass_scaled_mm.py | 2 +- .../kernels/quantization/test_cutlass_w4a8.py | 3 +- .../quantization/test_cutlass_w4a8_moe.py | 16 +- .../quantization/test_fp8_min_max_helper.py | 3 +- .../kernels/quantization/test_int8_kernel.py | 7 +- .../quantization/test_marlin_tile_padding.py | 15 +- .../test_mxfp4_kernel_selection.py | 20 +- .../quantization/test_mxfp4_triton_ep.py | 6 +- .../test_mxfp6_kernel_selection.py | 11 +- .../quantization/test_nvfp4_emulation.py | 3 +- .../kernels/quantization/test_nvfp4_quant.py | 3 +- .../test_per_token_group_quant.py | 14 +- .../quantization/test_quantized_embedding.py | 3 +- .../test_rocm_aiter_grouped_quant.py | 6 +- tests/kernels/quantization/test_rocm_mxfp4.py | 5 +- .../test_scaled_mm_kernel_selection.py | 2 +- .../quantization/test_triton_scaled_mm.py | 2 +- .../test_apply_repetition_penalties.py | 6 +- tests/kernels/test_awq_int4_to_int8.py | 4 +- tests/kernels/test_cache_kernels.py | 4 +- tests/kernels/test_compressor_kv_cache.py | 18 +- tests/kernels/test_cp_gather_fp8.py | 4 +- tests/kernels/test_deepseek_v4_cpu_kernels.py | 57 +- tests/kernels/test_engram.py | 9 +- tests/kernels/test_fla_layernorm_guard.py | 3 +- tests/kernels/test_flex_attention.py | 2 +- tests/kernels/test_fp32_router_gemm.py | 3 +- .../test_fused_deepseek_v32_norm_rope.py | 3 +- ..._fused_deepseek_v4_qnorm_rope_kv_insert.py | 9 +- .../test_fused_indexer_q_rope_quant.py | 4 +- .../kernels/test_fused_inv_rope_fp8_quant.py | 17 +- tests/kernels/test_mhc_kernels.py | 10 +- tests/kernels/test_minimax_m3_amd_ops.py | 12 +- tests/kernels/test_shuffle_rows.py | 8 +- tests/kernels/test_top_k_per_row.py | 51 +- tests/kernels/utils.py | 262 +- tests/lora/conftest.py | 4 +- tests/lora/test_add_lora.py | 3 +- tests/lora/test_default_mm_loras.py | 7 +- tests/lora/test_fused_moe_lora_kernel.py | 22 +- tests/lora/test_layers.py | 9 +- tests/lora/test_lora_functions.py | 4 +- tests/lora/test_mixtral.py | 2 +- tests/lora/test_moe_lora_align_sum.py | 6 +- tests/lora/test_moe_lora_ep_load.py | 9 +- tests/lora/test_punica_ops.py | 23 +- tests/lora/test_punica_ops_fp8.py | 6 +- tests/lora/test_punica_xpu_ops.py | 12 +- tests/lora/test_qwen3_unembed.py | 3 +- tests/lora/test_qwen3_with_multi_loras.py | 12 +- tests/lora/test_qwenvl.py | 9 +- tests/lora/test_whisper.py | 3 +- tests/lora/utils.py | 3 +- .../layers/test_fused_shared_expert.py | 7 +- .../layers/test_mla_short_prefill_indexer.py | 36 +- .../model_loader/test_ep_weight_filter.py | 9 +- .../model_loader/test_reload.py | 9 +- tests/model_executor/test_deep_gemm_warmup.py | 8 +- .../test_model_load_with_params.py | 12 +- tests/model_executor/test_oink_integration.py | 1 - tests/model_executor/test_qwen3_omni.py | 4 +- tests/model_executor/test_qwen3_vl_mrope.py | 4 +- .../test_routed_experts_capture.py | 7 +- tests/model_executor/test_weight_utils.py | 32 +- tests/models/inkling/test_mm_towers.py | 3 +- .../kimi_k3/test_amd_latent_moe_runner.py | 3 +- .../kimi_k3/test_aux_attn_res_stream.py | 6 +- tests/models/kimi_k3/test_kda_metadata.py | 3 +- .../models/language/generation/test_gemma.py | 6 +- .../models/language/generation/test_hybrid.py | 12 +- .../language/generation/test_mistral.py | 3 +- .../pooling/test_max_tokens_per_doc.py | 4 +- .../pooling/test_splade_sparse_pooler.py | 3 +- .../multimodal/generation/test_interleaved.py | 4 +- .../multimodal/generation/test_maverick.py | 24 +- .../generation/test_moondream3_moe.py | 3 +- .../generation/test_nemotron_parse.py | 3 +- .../multimodal/generation/test_qwen2_vl.py | 6 +- .../generation/test_transformers_audio.py | 3 +- .../multimodal/generation/test_ultravox.py | 1 - .../multimodal/generation/test_voxtral.py | 1 - .../generation/vlm_utils/builders.py | 3 +- .../generation/vlm_utils/case_filtering.py | 1 + .../multimodal/generation/vlm_utils/core.py | 2 +- .../generation/vlm_utils/custom_inputs.py | 2 + .../generation/vlm_utils/model_utils.py | 10 +- .../multimodal/pooling/test_dse_qwen2_vl.py | 2 +- .../pooling/test_jinavl_reranker.py | 13 +- .../pooling/test_llama_nemotron_vl.py | 3 +- .../processing/test_audio_in_video.py | 9 +- .../multimodal/processing/test_common.py | 7 +- .../processing/test_deepseek_ocr.py | 9 +- .../multimodal/processing/test_glm4_1v.py | 3 +- .../multimodal/processing/test_llava_next.py | 3 +- .../processing/test_llava_next_video.py | 6 +- .../processing/test_llava_onevision.py | 3 +- .../multimodal/processing/test_molmo2.py | 3 +- .../processing/test_qwen2_5_omni_embed.py | 29 +- .../multimodal/processing/test_qwen2_vl.py | 1 - .../multimodal/processing/test_qwen3_omni.py | 6 +- .../multimodal/processing/test_qwen3_vl.py | 12 +- .../processing/test_transformers_audio.py | 3 +- .../processing/test_transformers_image.py | 24 +- .../processing/transformers_backend.py | 6 +- tests/models/multimodal/test_cohere_asr.py | 3 +- tests/models/multimodal/test_mapping.py | 13 +- .../multimodal/test_nano_nemotron_vl.py | 9 +- tests/models/quantization/test_gpt_oss.py | 6 +- tests/models/qwen4_exp/test_config.py | 3 +- tests/models/registry.py | 7 +- tests/models/test_deepseek_v4_fi_moe_ep.py | 17 +- tests/models/test_deepseek_v4_mega_moe.py | 6 +- tests/models/test_dspark_mla.py | 3 +- tests/models/test_initialization.py | 5 +- tests/models/test_qwen3_5_mtp_config.py | 9 +- tests/models/test_utils.py | 3 +- tests/models/test_vision.py | 14 +- .../models/transformers/fusers/test_linear.py | 69 +- tests/models/transformers/fusers/test_mla.py | 15 +- tests/models/transformers/fusers/test_moe.py | 18 +- .../transformers/fusers/test_rms_norm.py | 27 +- .../transformers/test_layer_registry.py | 6 +- tests/models/utils.py | 8 +- tests/multimodal/media/test_audio.py | 22 +- tests/multimodal/media/test_connector.py | 3 +- tests/multimodal/media/test_image.py | 1 - tests/multimodal/media/test_video.py | 28 +- tests/multimodal/test_audio.py | 6 - .../test_embedding_shape_validation.py | 12 +- .../test_embedding_shape_validation_unit.py | 13 +- tests/multimodal/test_hasher.py | 3 +- tests/multimodal/test_image.py | 3 +- tests/multimodal/test_parse.py | 6 +- tests/multimodal/test_processing.py | 21 +- tests/multimodal/test_registry.py | 6 +- .../test_sparse_tensor_validation_unit.py | 6 +- tests/multimodal/test_utils.py | 3 +- tests/multimodal/test_video.py | 45 +- tests/parser/cohere/test_reasoning.py | 3 +- tests/parser/cohere/utils.py | 9 +- tests/parser/engine/conftest.py | 1 + tests/parser/engine/replay_harness.py | 1 + tests/parser/engine/test_deepseek_v32.py | 3 +- tests/parser/engine/test_deepseek_v4.py | 18 +- tests/parser/engine/test_delegating_replay.py | 3 +- tests/parser/engine/test_engine.py | 57 +- .../engine/test_gemma4_streaming_reasoning.py | 27 +- tests/parser/engine/test_glm47_moe.py | 6 +- tests/parser/engine/test_inkling.py | 30 +- tests/parser/engine/test_nemotron_v3.py | 18 +- tests/parser/engine/test_parser_engine.py | 69 +- tests/parser/engine/test_qwen3.py | 15 +- tests/parser/engine/test_qwen3_reasoning.py | 18 +- tests/parser/engine/test_replay.py | 12 +- tests/parser/engine/test_seed_oss.py | 6 +- .../engine/test_ufffd_reasoning_transition.py | 3 +- tests/parser/engine/trace_builder.py | 3 +- tests/parser/mistral/test_tool_calls.py | 18 +- tests/parser/test_streaming.py | 36 +- .../query_embedding_processor.py | 2 +- .../prithvi_io_processor/prithvi_processor.py | 5 +- .../prithvi_io_processor/types.py | 4 +- .../dummy_stat_logger/dummy_stat_logger.py | 3 +- .../gguf/test_gguf_plugin_generate.py | 4 +- tests/plugins_tests/test_endpoint_plugins.py | 12 +- .../test_io_processor_plugins.py | 3 +- tests/quantization/test_auto_round.py | 15 +- tests/quantization/test_compressed_tensors.py | 4 +- tests/quantization/test_fp8.py | 3 +- tests/quantization/test_gfx950_moe.py | 6 +- .../test_humming_mxfp4_block_fp8.py | 6 +- tests/quantization/test_int8_moe_oracle.py | 6 +- tests/quantization/test_modelopt.py | 12 +- tests/quantization/test_online.py | 5 +- tests/quantization/test_online_mxfp4.py | 12 +- tests/quantization/test_per_token_kv_cache.py | 3 +- .../test_register_quantization_config.py | 2 - .../test_base_thinking_reasoning_parser.py | 12 +- .../test_kimi_k3_reasoning_parser.py | 12 +- .../test_chat_utils_prompt_embeds.py | 9 +- tests/renderers/test_completions.py | 2 +- tests/renderers/test_gemma4_chat_template.py | 31 +- tests/renderers/test_hf.py | 13 +- .../test_process_multi_modal_uuids.py | 3 +- .../test_prompt_embeds_payload_validation.py | 5 +- .../test_sparse_tensor_concurrent_race.py | 3 +- .../test_sparse_tensor_validation.py | 48 +- tests/renderers/test_token_offsets.py | 15 +- tests/spec_decode/test_custom_proposer.py | 2 + tests/test_access_log_filter.py | 4 +- tests/test_cmake_utils.py | 5 +- tests/test_config.py | 36 +- tests/test_logger.py | 24 +- tests/test_ray_env.py | 6 +- tests/tokenizers_/test_deepseek_v4.py | 3 +- tests/tokenizers_/test_hf.py | 3 +- tests/tool_parsers/common_tests.py | 2 + .../tool_parsers/test_apertus_tool_parser.py | 12 +- .../test_deepseekv32_tool_parser.py | 30 +- tests/tool_parsers/test_gemma4_tool_parser.py | 3 +- .../test_gigachat3_tool_parser.py | 7 +- .../test_internlm2_tool_parser.py | 1 - .../tool_parsers/test_kimi_k2_tool_parser.py | 1 + .../test_kimi_k3_named_tool_choice.py | 3 +- tests/tool_parsers/test_lfm2_tool_parser.py | 27 +- .../test_llama3_json_tool_parser.py | 2 +- .../test_llama4_pythonic_tool_parser.py | 2 +- tests/tool_parsers/test_olmo3_tool_parser.py | 2 +- .../tool_parsers/test_phi4mini_tool_parser.py | 1 - .../tool_parsers/test_pythonic_tool_parser.py | 2 +- .../test_qwen3coder_tool_parser.py | 14 +- .../tool_parsers/test_step3p5_tool_parser.py | 14 +- tests/tool_use/test_chat_completions.py | 3 +- .../test_gemma4_responses_adjust_request.py | 10 +- tests/tool_use/test_muse_glimmer.py | 3 +- tests/tool_use/test_parallel_tool_calls.py | 5 +- tests/tool_use/test_tool_choice_required.py | 3 +- .../processors/test_glm5next.py | 6 +- tests/transformers_utils/test_config.py | 9 +- .../test_hf_overrides_model_type.py | 4 +- tests/utils.py | 61 +- tests/utils_/__init__.py | 3 +- tests/utils_/test_gpu_sync_debug.py | 15 +- tests/utils_/test_jsontree.py | 1 - tests/utils_/test_network_utils.py | 15 +- tests/utils_/test_numa_utils.py | 18 +- tests/v1/attention/test_attention_backends.py | 8 +- .../test_attention_backends_selection.py | 3 +- .../v1/attention/test_attention_splitting.py | 31 +- tests/v1/attention/test_dcp_a2a_pack_mask.py | 3 +- .../test_dspark_noncausal_sparse_mla.py | 3 +- .../test_flashinfer_mla_sparse_sm90.py | 3 +- .../v1/attention/test_gdn_metadata_builder.py | 3 +- .../v1/attention/test_indexer_dcp_localize.py | 18 +- .../attention/test_kpool_tail_slot_mapping.py | 21 +- .../test_mamba_update_block_table.py | 27 +- tests/v1/attention/test_mla_backends.py | 9 +- .../test_mla_prefill_quant_output.py | 9 +- .../v1/attention/test_mla_prefill_selector.py | 3 +- .../test_replayssm_metadata_builder.py | 3 +- .../test_rocm_aiter_mla_mtp_split.py | 3 - .../test_rocm_attention_backends_selection.py | 3 +- .../test_sparse_indexer_decode_seq_lens.py | 12 +- .../v1/attention/test_sparse_mla_backends.py | 31 +- .../test_trtllm_attention_integration.py | 10 +- tests/v1/attention/utils.py | 1 - .../test_partial_prefix_cache_hits.py | 51 +- tests/v1/core/test_async_scheduler.py | 6 +- tests/v1/core/test_contiguous_kv_packing.py | 18 +- tests/v1/core/test_encoder_cache_manager.py | 9 +- tests/v1/core/test_kv_cache_utils.py | 43 +- tests/v1/core/test_kv_sharing.py | 9 +- tests/v1/core/test_output.py | 9 +- tests/v1/core/test_prefix_caching.py | 120 +- tests/v1/core/test_repetition_detection.py | 42 +- tests/v1/core/test_scheduler.py | 120 +- .../v1/core/test_swa_inflight_window_free.py | 12 +- tests/v1/core/utils.py | 1 + .../v1/cudagraph/test_breakable_cudagraph.py | 13 +- tests/v1/cudagraph/test_cudagraph_dispatch.py | 1 - tests/v1/determinism/test_batch_invariance.py | 16 +- .../test_matmul_batch_invariant.py | 11 +- .../test_online_batch_invariance.py | 3 +- .../test_rms_norm_batch_invariant.py | 25 +- tests/v1/distributed/test_async_llm_dp.py | 16 +- tests/v1/distributed/test_dbo.py | 7 +- tests/v1/distributed/test_external_lb_dp.py | 3 +- tests/v1/distributed/test_hybrid_lb_dp.py | 3 +- tests/v1/distributed/test_internal_lb_dp.py | 11 +- tests/v1/distributed/test_pp_dp_v2.py | 14 +- tests/v1/e2e/general/test_async_scheduling.py | 9 +- .../test_attention_backend_per_kind.py | 3 +- tests/v1/e2e/general/test_context_length.py | 11 +- .../test_correctness_sliding_window.py | 7 +- .../general/test_kv_sharing_fast_prefill.py | 4 +- tests/v1/e2e/general/test_min_tokens.py | 33 +- .../general/test_pooling_chunked_prefill.py | 2 - tests/v1/e2e/general/test_sharded_sampling.py | 6 +- tests/v1/e2e/general/test_streaming_input.py | 3 +- .../mtp_other/test_synthetic.py | 3 +- .../e2e/spec_decode/acceptance_rates/utils.py | 3 +- .../draft_model/test_draft_model.py | 22 +- .../e2e/spec_decode/draft_model/test_lora.py | 6 +- tests/v1/e2e/spec_decode/eagle/utils.py | 3 +- tests/v1/e2e/spec_decode/mtp/test_mtp.py | 3 +- .../ngram_suffix/test_ngram_suffix.py | 6 +- .../speculators/test_speculators.py | 3 +- .../e2e/spec_decode/test_sharded_sampling.py | 3 +- tests/v1/e2e/test_hybrid_chunked_prefill.py | 4 +- .../integration/test_epd_correctness.py | 5 +- .../cpu/scheduler/test_embedding_cache.py | 3 +- .../unit/cpu/scheduler/test_scheduler.py | 36 +- .../unit/cpu/test_ec_shared_region.py | 12 +- .../unit/cpu/worker/test_worker.py | 18 +- .../unit/test_ec_cpu_connector.py | 3 +- .../unit/test_ec_example_connector.py | 4 +- .../unit/test_epd_proxy_round_robin.py | 3 +- tests/v1/ec_connector/unit/test_metadata.py | 5 +- .../unit/test_scheduler_nixl_producer.py | 3 +- tests/v1/ec_connector/unit/test_session.py | 7 +- tests/v1/ec_connector/unit/utils.py | 1 + tests/v1/engine/conftest.py | 7 +- tests/v1/engine/test_abort_final_step.py | 10 +- tests/v1/engine/test_async_llm.py | 2 - .../engine/test_core_engine_actor_manager.py | 3 +- tests/v1/engine/test_engine_core.py | 28 +- tests/v1/engine/test_engine_core_client.py | 13 +- .../v1/engine/test_fast_incdec_prefix_err.py | 3 +- tests/v1/engine/test_llm_engine.py | 3 +- tests/v1/engine/test_logprobs_processor.py | 3 +- tests/v1/engine/test_output_processor.py | 5 +- .../engine/test_preprocess_error_handling.py | 1 - tests/v1/engine/utils.py | 22 +- .../predictable_llama.py | 1 + .../nixl_integration/test_disagg_accuracy.py | 7 +- .../test_multi_connector_edge_cases.py | 3 +- .../nixl_integration/toy_proxy_server.py | 16 +- .../test_canonical_mapping.py | 3 +- .../unit/offloading_connector/test_config.py | 6 +- .../unit/offloading_connector/test_metrics.py | 6 +- .../offloading_connector/test_scheduler.py | 48 +- .../unit/offloading_connector/test_worker.py | 7 +- .../unit/offloading_connector/utils.py | 10 +- .../unit/test_bidirectional_kv_transfer.py | 57 +- .../unit/test_cache_pollution_prevention.py | 8 +- tests/v1/kv_connector/unit/test_config.py | 18 +- .../unit/test_decode_bench_connector.py | 3 +- .../unit/test_error_propagation.py | 6 +- .../unit/test_example_connector.py | 14 +- .../unit/test_flexkv_connector.py | 9 +- .../unit/test_handshake_pp_aggregation.py | 9 +- .../v1/kv_connector/unit/test_hf3fs_client.py | 9 +- .../kv_connector/unit/test_hf3fs_connector.py | 7 +- .../unit/test_hf3fs_metadata_server.py | 9 +- .../unit/test_invalid_blocks_correctness.py | 16 +- .../unit/test_lmcache_connector.py | 6 +- .../unit/test_mooncake_connector.py | 32 +- .../unit/test_mooncake_connector_hma.py | 2 - .../kv_connector/unit/test_mooncake_stats.py | 3 +- .../unit/test_mooncake_store_connector.py | 9 +- .../unit/test_mooncake_store_coordinator.py | 54 +- .../unit/test_mooncake_store_hma_e2e.py | 24 +- .../unit/test_mooncake_store_scheduler.py | 6 +- .../unit/test_mooncake_store_worker.py | 62 +- .../unit/test_moriio_connector.py | 24 +- .../unit/test_moriio_routing_fairness.py | 6 +- .../kv_connector/unit/test_multi_connector.py | 22 +- .../kv_connector/unit/test_nixl_connector.py | 66 +- .../unit/test_nixl_connector_hma.py | 54 +- .../unit/test_nixl_desc_geometry.py | 30 +- .../unit/test_nixl_push_connector.py | 105 +- .../unit/test_nixl_rocm_gpu_mem_diag.py | 3 +- .../unit/test_nixl_simple_cpu_offload.py | 9 +- .../unit/test_offloading_connector.py | 14 +- .../unit/test_remote_decode_lifecycle.py | 4 - .../unit/test_remote_prefill_lifecycle.py | 18 +- .../unit/test_simple_cpu_offload_connector.py | 3 +- tests/v1/kv_connector/unit/test_tp_mapping.py | 6 +- tests/v1/kv_connector/unit/utils.py | 1 - .../kv_offload/cpu/policies/test_factory.py | 15 +- .../kv_offload/cpu/test_canonical_layout.py | 9 +- tests/v1/kv_offload/cpu/test_manager.py | 41 +- .../cpu/test_shared_offload_region.py | 42 +- tests/v1/kv_offload/test_file_mapper.py | 3 +- .../tiering/p2p/p2p_connector_proxy.py | 3 +- .../tiering/p2p/test_data_transport.py | 8 +- .../v1/kv_offload/tiering/p2p/test_manager.py | 93 +- .../kv_offload/tiering/p2p/test_sessions.py | 122 +- .../kv_offload/tiering/test_async_lookup.py | 6 +- tests/v1/kv_offload/tiering/test_factory.py | 3 +- tests/v1/kv_offload/tiering/test_fs_tier.py | 27 +- tests/v1/kv_offload/tiering/test_obj_tier.py | 9 +- .../tiering/test_tiering_offloading.py | 36 +- .../v1/logits_processors/test_correctness.py | 30 +- .../logits_processors/test_custom_offline.py | 11 +- .../logits_processors/test_custom_online.py | 7 +- tests/v1/logits_processors/utils.py | 19 +- tests/v1/metrics/test_engine_logger_apis.py | 11 +- tests/v1/metrics/test_perf_metrics.py | 18 +- tests/v1/metrics/test_ray_metrics.py | 22 +- .../sample/test_batched_count_greater_than.py | 3 +- tests/v1/sample/test_logprobs.py | 20 +- tests/v1/sample/test_rejection_sampler.py | 26 +- tests/v1/sample/test_sampler.py | 19 +- tests/v1/sample/test_sampling_params_e2e.py | 8 - tests/v1/sample/test_topk_topp_sampler.py | 9 +- tests/v1/sample/utils.py | 18 +- tests/v1/shutdown/test_delete.py | 2 + tests/v1/shutdown/test_startup_error.py | 1 - tests/v1/shutdown/utils.py | 2 +- tests/v1/simple_kv_offload/test_kv_events.py | 9 +- tests/v1/simple_kv_offload/test_scheduler.py | 19 +- .../v1/spec_decode/test_acceptance_length.py | 3 +- .../test_backup_token_async_spec.py | 3 +- tests/v1/spec_decode/test_dynamic_sd_cug.py | 6 - tests/v1/spec_decode/test_eagle.py | 22 +- .../test_eagle_draft_attn_metadata.py | 9 +- .../v1/spec_decode/test_eagle_step_kernel.py | 3 +- .../spec_decode/test_extract_hidden_states.py | 6 +- .../v1/spec_decode/test_llm_base_proposer.py | 6 +- tests/v1/spec_decode/test_max_len.py | 3 +- tests/v1/spec_decode/test_mtp.py | 4 +- .../test_rejection_sampler_utils.py | 22 +- .../test_speculators_correctness.py | 6 +- .../v1/spec_decode/test_speculators_eagle3.py | 3 +- .../test_gpu_model_runner_streaming.py | 1 - .../test_gpu_model_runner_v2_streaming.py | 1 - .../test_scheduler_streaming.py | 3 +- tests/v1/structured_output/test_validation.py | 12 +- tests/v1/test_kv_cache_spec_registry.py | 16 +- tests/v1/test_outputs.py | 24 +- tests/v1/test_serial_utils.py | 5 +- tests/v1/utils.py | 3 + .../test_dsv4_packed_zeroer_geometry.py | 3 +- tests/v1/worker/test_encoder_runner.py | 15 +- tests/v1/worker/test_gpu_bad_words.py | 9 +- tests/v1/worker/test_gpu_batch_shard.py | 27 +- tests/v1/worker/test_gpu_block_table.py | 6 +- tests/v1/worker/test_gpu_gumbel_sample.py | 2 +- tests/v1/worker/test_gpu_input_batch.py | 13 +- tests/v1/worker/test_gpu_logit_bias.py | 3 +- tests/v1/worker/test_gpu_model_runner.py | 14 +- .../worker/test_gpu_model_runner_mm_gather.py | 9 +- tests/v1/worker/test_gpu_model_runner_v2.py | 13 +- ...gpu_model_runner_v2_cudagraph_profiling.py | 3 +- tests/v1/worker/test_gpu_profiler.py | 10 +- .../worker/test_gpu_rejection_sampler_i64.py | 3 +- tests/v1/worker/test_gpu_thinking_budget.py | 3 +- tests/v1/worker/test_gpu_ubatch_slicing.py | 3 +- tests/v1/worker/test_gpu_worker.py | 15 +- tests/v1/worker/test_kv_block_zeroer.py | 3 +- tests/v1/worker/test_mamba_utils.py | 45 +- tests/v1/worker/test_mixed_warmup_gate.py | 3 +- tests/v1/worker/test_prompt_embeds_state.py | 12 +- .../v1/worker/test_worker_memory_snapshot.py | 1 + tests/v1/worker/test_workspace.py | 3 +- .../vllm_test_utils/__init__.py | 3 +- .../vllm_test_utils/vllm_test_utils/blame.py | 3 +- .../vllm_test_utils/monitor.py | 3 +- tests/weight_loading/test_weight_loading.py | 5 +- tools/generate_cmake_presets.py | 1 - tools/generate_versions_json.py | 3 +- tools/install_nixl_from_source_ubuntu.py | 3 +- tools/pre_commit/check_spdx_header.py | 8 +- .../pre_commit/generate_nightly_torch_test.py | 3 +- tools/pre_commit/mypy.py | 12 +- tools/pre_commit/validate_config.py | 7 +- .../nsys_profile_tools/gputrc2graph.py | 27 +- tools/recipes/recipe_json_to_vllm_config.py | 7 +- tools/report_build_time_ninja.py | 10 +- tools/vllm-rocm/pin_rocm_dependencies.py | 14 +- vllm/__init__.py | 2 +- vllm/_aiter_ops.py | 43 +- vllm/_custom_ops.py | 110 +- vllm/_xpu_ops.py | 3 + vllm/assets/base.py | 3 +- vllm/assets/image.py | 8 +- vllm/assets/video.py | 6 +- .../datasets/create_txt_slices_dataset.py | 7 +- vllm/benchmarks/datasets/datasets.py | 163 +- vllm/benchmarks/datasets/utils.py | 8 +- vllm/benchmarks/lib/endpoint_request_func.py | 8 +- vllm/benchmarks/lib/ready_checker.py | 5 +- vllm/benchmarks/lib/utils.py | 7 +- vllm/benchmarks/mm_processor.py | 20 +- vllm/benchmarks/plot.py | 13 +- vllm/benchmarks/serve.py | 10 +- vllm/benchmarks/startup.py | 10 +- vllm/benchmarks/sweep/param_sweep.py | 10 +- vllm/benchmarks/sweep/plot.py | 3 +- vllm/benchmarks/throughput.py | 9 +- vllm/compilation/backends.py | 21 +- vllm/compilation/base_static_graph.py | 12 +- vllm/compilation/caching.py | 9 +- vllm/compilation/codegen.py | 2 + vllm/compilation/compiler_interface.py | 32 +- vllm/compilation/cuda_graph.py | 2 +- vllm/compilation/decorators.py | 14 +- vllm/compilation/partition_rules.py | 5 +- .../passes/fusion/act_quant_fusion.py | 17 +- .../passes/fusion/allreduce_rms_fusion.py | 21 +- .../passes/fusion/attn_quant_fusion.py | 9 +- .../passes/fusion/matcher_utils.py | 2 +- .../passes/fusion/mla_attn_quant_fusion.py | 12 +- .../passes/fusion/qk_norm_rope_fusion.py | 3 +- .../fusion/qk_norm_rope_kvcache_fusion.py | 6 +- .../passes/fusion/rms_quant_fusion.py | 6 +- .../passes/fusion/rocm_aiter_fusion.py | 39 +- .../passes/fusion/rope_kvcache_fusion.py | 12 +- .../passes/fusion/sequence_parallelism.py | 9 +- vllm/compilation/passes/inductor_pass.py | 20 +- .../passes/ir/clone_elimination.py | 3 +- .../passes/ir/inplace_functionalization.py | 3 +- vllm/compilation/passes/ir/lowering_pass.py | 7 +- vllm/compilation/passes/pass_manager.py | 9 +- .../passes/utility/fix_functionalization.py | 21 +- .../passes/utility/noop_elimination.py | 7 +- .../passes/utility/post_cleanup.py | 3 +- .../passes/utility/scatter_split_replace.py | 3 +- .../passes/utility/split_coalescing.py | 9 +- vllm/compilation/passes/vllm_inductor_pass.py | 20 +- vllm/compilation/piecewise_backend.py | 3 +- vllm/compilation/wrapper.py | 10 +- vllm/config/attention.py | 3 +- vllm/config/cache.py | 3 +- vllm/config/compilation.py | 39 +- vllm/config/device.py | 3 +- vllm/config/ec_manager_config.py | 3 +- vllm/config/ec_transfer.py | 3 +- vllm/config/kernel.py | 22 +- vllm/config/kv_transfer.py | 3 +- vllm/config/load.py | 3 +- vllm/config/lora.py | 3 +- vllm/config/model.py | 36 +- vllm/config/model_arch.py | 4 +- vllm/config/multimodal.py | 15 +- vllm/config/observability.py | 6 +- vllm/config/offload.py | 3 +- vllm/config/parallel.py | 10 +- vllm/config/pooler.py | 3 +- vllm/config/profiler.py | 3 +- vllm/config/reasoning.py | 6 +- vllm/config/scheduler.py | 7 +- vllm/config/speculative.py | 17 +- vllm/config/structured_outputs.py | 3 +- vllm/config/utils.py | 28 +- vllm/config/vllm.py | 47 +- vllm/device_allocator/cumem.py | 30 +- vllm/device_allocator/sleep_mode_backend.py | 15 +- vllm/device_allocator/xpumem.py | 1 + .../aiter_custom_all_reduce.py | 4 +- .../device_communicators/all2all.py | 52 +- .../device_communicators/all_reduce_utils.py | 7 +- .../base_device_communicator.py | 18 +- .../device_communicators/cpu_communicator.py | 16 +- .../device_communicators/cuda_communicator.py | 12 +- .../device_communicators/custom_all_reduce.py | 7 +- .../flashinfer_all_reduce.py | 7 +- .../device_communicators/pynccl.py | 4 +- .../device_communicators/quick_all_reduce.py | 11 +- .../device_communicators/ray_communicator.py | 30 +- .../device_communicators/shm_broadcast.py | 20 +- .../shm_object_storage.py | 41 +- .../device_communicators/xpu_communicator.py | 10 +- .../ec_transfer/ec_connector/base.py | 65 +- .../ec_connector/cpu/ec_shared_region.py | 3 +- .../ec_connector/cpu/scheduler/__init__.py | 3 + .../cpu/scheduler/embedding_cache.py | 3 +- .../ec_connector/example_connector.py | 22 +- .../ec_transfer/ec_transfer_state.py | 5 +- vllm/distributed/eplb/async_worker.py | 4 +- vllm/distributed/eplb/eplb_communicator.py | 12 +- vllm/distributed/eplb/eplb_state.py | 51 +- vllm/distributed/eplb/eplb_utils.py | 13 +- vllm/distributed/eplb/policy/abstract.py | 7 +- vllm/distributed/eplb/policy/default.py | 30 +- vllm/distributed/eplb/rebalance_execute.py | 33 +- vllm/distributed/kv_events.py | 38 +- .../kv_transfer/kv_connector/factory.py | 1 + .../kv_transfer/kv_connector/utils.py | 23 +- .../kv_transfer/kv_connector/v1/base.py | 140 +- .../kv_connector/v1/decode_bench_connector.py | 32 +- .../kv_connector/v1/example_connector.py | 14 +- .../v1/example_hidden_states_connector.py | 21 +- .../kv_connector/v1/flexkv_connector.py | 10 + .../kv_connector/v1/hf3fs/hf3fs_client.py | 8 +- .../kv_connector/v1/hf3fs/hf3fs_connector.py | 16 +- .../v1/hf3fs/hf3fs_metadata_server.py | 8 +- .../v1/hf3fs/utils/gather_scatter_helper.py | 2 + .../kv_connector/v1/lmcache_connector.py | 66 +- .../multi_process_adapter.py | 93 +- .../v1/lmcache_integration/utils.py | 11 +- .../v1/lmcache_integration/vllm_v1_adapter.py | 21 +- .../kv_connector/v1/lmcache_mp_connector.py | 103 +- .../kv_transfer/kv_connector/v1/metrics.py | 21 +- .../v1/mooncake/mooncake_connector.py | 39 +- .../v1/mooncake/mooncake_utils.py | 3 +- .../v1/mooncake/store/coordinator.py | 3 +- .../kv_connector/v1/mooncake/store/data.py | 3 + .../kv_connector/v1/mooncake/store/worker.py | 4 +- .../kv_connector/v1/moriio/moriio_common.py | 3 +- .../v1/moriio/moriio_connector.py | 33 +- .../kv_connector/v1/moriio/moriio_engine.py | 10 +- .../kv_connector/v1/multi_connector.py | 20 +- .../kv_connector/v1/nixl/base_scheduler.py | 10 +- .../kv_connector/v1/nixl/base_worker.py | 55 +- .../kv_connector/v1/nixl/connector.py | 8 +- .../kv_connector/v1/nixl/metadata.py | 7 +- .../kv_connector/v1/nixl/pull_scheduler.py | 8 +- .../kv_connector/v1/nixl/pull_worker.py | 9 +- .../kv_connector/v1/nixl/push_worker.py | 6 +- .../kv_transfer/kv_connector/v1/nixl/stats.py | 2 +- .../kv_transfer/kv_connector/v1/nixl/utils.py | 3 +- .../v1/offloading/canonical_mapping.py | 12 +- .../kv_connector/v1/offloading/events.py | 10 +- .../kv_connector/v1/offloading/metrics.py | 6 +- .../kv_connector/v1/offloading/scheduler.py | 30 +- .../kv_connector/v1/offloading/worker.py | 18 +- .../v1/ssm_conv_transfer_utils.py | 3 + .../kv_transfer/kv_transfer_state.py | 6 +- vllm/distributed/parallel_state.py | 51 +- vllm/distributed/stateless_coordinator.py | 3 +- vllm/distributed/utils.py | 18 +- vllm/distributed/weight_transfer/__init__.py | 3 +- vllm/distributed/weight_transfer/base.py | 37 +- vllm/distributed/weight_transfer/factory.py | 6 +- .../distributed/weight_transfer/ipc_engine.py | 17 +- .../weight_transfer/nccl_common.py | 10 +- .../weight_transfer/nccl_engine.py | 17 +- .../weight_transfer/packed_tensor.py | 4 + .../weight_transfer/sharded_rdt_common.py | 12 +- .../weight_transfer/sharded_rdt_engine.py | 48 +- .../weight_transfer/sharded_rdt_fake.py | 6 +- .../weight_transfer/sharded_rdt_trainer.py | 45 +- .../weight_transfer/sparse_nccl_engine.py | 9 +- vllm/engine/arg_utils.py | 4 +- vllm/engine/protocol.py | 29 +- vllm/entrypoints/anthropic/protocol.py | 30 +- vllm/entrypoints/anthropic/serving.py | 33 +- vllm/entrypoints/chat_utils.py | 26 +- vllm/entrypoints/cli/benchmark/base.py | 1 + vllm/entrypoints/cli/main.py | 5 +- vllm/entrypoints/generate/base/serving.py | 6 +- .../generate/beam_search/offline.py | 6 +- .../generate/beam_search/online.py | 2 +- .../generate/generative_scoring/serving.py | 6 + .../entrypoints/launchers/api_server/entry.py | 7 +- vllm/entrypoints/launchers/cli_args.py | 3 +- vllm/entrypoints/launchers/dp_supervisor.py | 27 +- vllm/entrypoints/launchers/grpc_server.py | 8 +- vllm/entrypoints/launchers/launcher.py | 10 +- .../entrypoints/launchers/render/app_state.py | 1 - vllm/entrypoints/launchers/render/entry.py | 1 - vllm/entrypoints/launchers/run_batch.py | 38 +- .../launchers/utils/server_utils.py | 3 +- vllm/entrypoints/llm.py | 38 +- vllm/entrypoints/mcp/tool.py | 6 +- vllm/entrypoints/mcp/tool_server.py | 11 +- vllm/entrypoints/offline_utils.py | 10 +- .../openai/chat_completion/batch_serving.py | 1 + .../openai/chat_completion/serving.py | 7 +- vllm/entrypoints/openai/completion/serving.py | 4 +- vllm/entrypoints/openai/models/serving.py | 7 +- .../openai/parser/harmony_utils.py | 24 +- .../openai/responses/api_router.py | 2 +- vllm/entrypoints/openai/responses/context.py | 21 +- vllm/entrypoints/openai/responses/harmony.py | 6 +- vllm/entrypoints/openai/responses/protocol.py | 11 +- .../openai/responses/streaming_events.py | 15 +- vllm/entrypoints/openai/responses/utils.py | 14 +- vllm/entrypoints/pooling/offline.py | 17 +- .../pooling/pooling/io_processor.py | 3 +- vllm/entrypoints/pooling/scoring/typing.py | 3 +- vllm/entrypoints/pooling/scoring/utils.py | 7 +- .../token_in_token_out/api_router.py | 3 +- .../scale_out/token_in_token_out/mm_serde.py | 3 +- .../entrypoints/serve/dev/cache/api_router.py | 10 +- vllm/entrypoints/serve/dev/rlhf/api_router.py | 10 +- .../serve/elastic_ep/middleware.py | 3 +- vllm/entrypoints/serve/engine/serving.py | 2 +- .../exception_handling/handlers/validation.py | 3 +- .../exception_handling/handlers/vllm_error.py | 4 +- .../serve/exception_handling/register.py | 3 +- .../serve/instrumentator/metrics.py | 1 - vllm/entrypoints/serve/lora/api_router.py | 6 +- .../serve/middleware/authenticate.py | 4 +- .../serve/middleware/log_response.py | 1 - .../serve/middleware/x_request_id.py | 6 +- .../entrypoints/serve/sagemaker/api_router.py | 8 +- vllm/entrypoints/serve/tokenize/protocol.py | 3 +- vllm/entrypoints/serve/utils/api_utils.py | 2 +- vllm/entrypoints/serve/utils/fingerprint.py | 6 +- vllm/entrypoints/serve/utils/orca_metrics.py | 35 +- .../serve/utils/tool_calls_utils.py | 1 - .../speech_to_text/base/serving.py | 9 +- vllm/entrypoints/speech_to_text/base/utils.py | 1 + .../speech_to_text/realtime/protocol.py | 14 +- .../speech_to_text/realtime/serving.py | 1 + vllm/env_override.py | 15 +- vllm/envs.py | 30 +- vllm/exceptions.py | 5 +- vllm/forward_context.py | 9 +- vllm/inputs/engine.py | 19 +- vllm/inputs/llm.py | 10 +- vllm/ir/op.py | 53 +- vllm/ir/ops/layernorm.py | 4 +- vllm/ir/util.py | 10 +- vllm/kernels/helion/case_key.py | 4 +- vllm/kernels/helion/config_manager.py | 3 +- vllm/kernels/helion/ops/__init__.py | 1 + .../ops/dynamic_per_token_scaled_fp8_quant.py | 1 - vllm/kernels/helion/ops/fused_qk_norm_rope.py | 1 - .../helion/ops/per_token_group_fp8_quant.py | 1 - .../ops/rms_norm_dynamic_per_token_quant.py | 1 - .../helion/ops/rms_norm_per_block_quant.py | 1 - .../ops/silu_and_mul_per_block_quant.py | 1 - vllm/kernels/helion/register.py | 11 +- vllm/kernels/helion/utils.py | 3 +- vllm/logger.py | 33 +- vllm/logging_utils/access_log_filter.py | 15 +- vllm/logging_utils/formatter.py | 5 +- vllm/logging_utils/log_time.py | 7 +- vllm/logprobs.py | 24 +- vllm/lora/layers/base_linear.py | 3 +- vllm/lora/layers/column_parallel_linear.py | 31 +- vllm/lora/layers/fused_moe.py | 31 +- vllm/lora/layers/logits_processor.py | 4 +- vllm/lora/layers/replicated_linear.py | 3 +- vllm/lora/layers/row_parallel_linear.py | 6 +- vllm/lora/layers/utils.py | 6 +- vllm/lora/lora_model.py | 35 +- vllm/lora/lora_weights.py | 1 - vllm/lora/model_manager.py | 12 +- vllm/lora/ops/triton_ops/fp8_kernel_utils.py | 21 +- .../ops/triton_ops/fused_moe_lora_fp8_op.py | 13 +- vllm/lora/ops/triton_ops/fused_moe_lora_op.py | 13 +- vllm/lora/ops/triton_ops/kernel_utils.py | 34 +- .../lora/ops/triton_ops/lora_expand_fp8_op.py | 16 +- vllm/lora/ops/triton_ops/lora_expand_op.py | 52 +- .../ops/triton_ops/lora_kernel_metadata.py | 16 +- .../lora/ops/triton_ops/lora_shrink_fp8_op.py | 40 +- vllm/lora/ops/triton_ops/lora_shrink_op.py | 52 +- vllm/lora/ops/triton_ops/utils.py | 11 +- vllm/lora/peft_helper.py | 10 +- vllm/lora/punica_wrapper/punica_base.py | 72 +- vllm/lora/punica_wrapper/punica_cpu.py | 50 +- vllm/lora/punica_wrapper/punica_gpu.py | 56 +- vllm/lora/punica_wrapper/punica_xpu.py | 59 +- vllm/lora/punica_wrapper/utils.py | 6 +- vllm/lora/request.py | 9 +- vllm/lora/resolver.py | 7 + vllm/lora/utils.py | 27 +- vllm/lora/worker_manager.py | 9 +- vllm/model_executor/custom_op.py | 12 +- .../determinism/batch_invariant.py | 25 +- vllm/model_executor/hw_agnostic/custom_op.py | 15 +- .../hw_agnostic/layers/layernorm.py | 3 +- .../attention/dsa/sparse_mqa_logits.py | 18 + .../model_executor/kernels/linear/__init__.py | 28 +- vllm/model_executor/kernels/linear/base.py | 14 +- .../mixed_precision/rdna_hybrid_w4a16.py | 10 +- .../linear/mixed_precision/triton_w4a16.py | 16 +- .../kernels/linear/mxfp4/base.py | 1 + .../kernels/linear/mxfp6/base.py | 1 + .../kernels/linear/nvfp4/flashinfer.py | 3 +- .../scaled_mm/BlockScaledMMLinearKernel.py | 3 +- .../kernels/linear/scaled_mm/aiter.py | 3 +- .../kernels/linear/scaled_mm/cutlass.py | 6 +- .../kernels/linear/scaled_mm/flashinfer.py | 15 +- .../kernels/linear/scaled_mm/marlin.py | 3 +- .../kernels/linear/scaled_mm/pytorch.py | 3 +- vllm/model_executor/kernels/mhc/aiter.py | 8 +- vllm/model_executor/kernels/mhc/tilelang.py | 11 +- .../kernels/mhc/tilelang_kernels.py | 2 +- vllm/model_executor/kernels/mhc/torch.py | 5 +- vllm/model_executor/layers/activation.py | 13 +- .../layers/attention/attention.py | 12 +- .../layers/attention/cross_attention.py | 4 +- .../attention/encoder_only_attention.py | 4 +- .../layers/attention/mla_attention.py | 45 +- .../layers/attention/mm_encoder_attention.py | 19 +- .../layers/attention/static_sink_attention.py | 4 +- .../layers/attention_layer_base.py | 6 +- vllm/model_executor/layers/conv.py | 4 +- .../layers/fused_allreduce_gemma_rms_norm.py | 3 +- .../model_executor/layers/fused_embed_norm.py | 3 +- .../layers/fused_moe/activation.py | 4 +- .../layers/fused_moe/all2all_utils.py | 4 +- .../model_executor/layers/fused_moe/config.py | 88 +- .../layers/fused_moe/deep_gemm_utils.py | 3 +- .../layers/fused_moe/expert_map_manager.py | 40 +- .../experts/aiter_mxfp4_w4a16_moe.py | 4 +- .../fused_moe/experts/aiter_mxfp4_w4a8_moe.py | 6 +- .../fused_moe/experts/aiter_mxfp8_moe.py | 3 +- .../experts/batched_deep_gemm_moe.py | 6 +- .../layers/fused_moe/experts/cpu_int4_moe.py | 2 + .../layers/fused_moe/experts/cpu_moe.py | 11 +- .../layers/fused_moe/experts/cutlass_moe.py | 3 +- .../layers/fused_moe/experts/deep_gemm_moe.py | 3 +- .../layers/fused_moe/experts/fallback.py | 3 +- .../experts/flashinfer_cutedsl_batched_moe.py | 10 +- .../experts/flashinfer_cutedsl_moe.py | 3 +- .../experts/flashinfer_cutlass_moe.py | 6 +- .../fused_moe/experts/fused_batched_moe.py | 6 +- .../fused_moe/experts/fused_humming_moe.py | 15 +- .../experts/gpt_oss_triton_kernels_moe.py | 9 +- .../layers/fused_moe/experts/lora_context.py | 3 +- .../fused_moe/experts/lora_experts_mixin.py | 3 +- .../layers/fused_moe/experts/marlin_moe.py | 64 +- .../fused_moe/experts/mxfp8_native_moe.py | 3 +- .../fused_moe/experts/nvfp4_emulation_moe.py | 12 +- .../fused_moe/experts/ocp_mx_emulation_moe.py | 10 +- .../fused_moe/experts/trtllm_bf16_moe.py | 11 +- .../fused_moe/experts/trtllm_fp8_moe.py | 11 +- .../fused_moe/experts/trtllm_lora_moe.py | 5 +- .../fused_moe/experts/trtllm_mxfp4_moe.py | 10 +- .../fused_moe/experts/trtllm_mxint4_moe.py | 3 +- .../fused_moe/experts/trtllm_nvfp4_moe.py | 12 +- .../layers/fused_moe/fused_moe.py | 14 +- .../layers/fused_moe/fused_moe_method_base.py | 20 +- .../layers/fused_moe/hpc_moe.py | 3 +- vllm/model_executor/layers/fused_moe/layer.py | 6 + .../layers/fused_moe/modular_kernel.py | 273 +- .../layers/fused_moe/moe_align_block_size.py | 56 +- .../layers/fused_moe/moe_fused_mul_sum.py | 4 +- .../layers/fused_moe/moe_output.py | 2 +- .../layers/fused_moe/moe_permute_unpermute.py | 51 +- .../layers/fused_moe/oracle/base.py | 15 +- .../layers/fused_moe/oracle/fp8.py | 18 +- .../layers/fused_moe/oracle/int8.py | 11 +- .../layers/fused_moe/oracle/int_wna16.py | 25 +- .../layers/fused_moe/oracle/mxfp4.py | 20 +- .../layers/fused_moe/oracle/mxfp8.py | 2 +- .../layers/fused_moe/oracle/nvfp4.py | 4 +- .../layers/fused_moe/oracle/unquantized.py | 7 +- .../layers/fused_moe/oracle/w4a8_int8.py | 24 +- .../fused_moe/prepare_finalize/batched.py | 3 +- .../fused_moe/prepare_finalize/deepep_ht.py | 4 +- .../fused_moe/prepare_finalize/deepep_ll.py | 8 +- .../fused_moe/prepare_finalize/deepep_v2.py | 6 +- .../layers/fused_moe/prepare_finalize/mori.py | 7 +- .../fused_moe/prepare_finalize/naive_dp_ep.py | 8 +- .../fused_moe/prepare_finalize/nixl_ep.py | 8 +- .../layers/fused_moe/routed_experts.py | 30 +- .../fused_moe/routed_experts_capturer.py | 4 +- .../aiter_shared_routed_fused_moe_router.py | 3 +- .../layers/fused_moe/router/base_router.py | 22 +- .../fused_moe/router/fused_moe_router.py | 8 +- .../layers/fused_moe/router/gate_linear.py | 3 +- .../layers/fused_moe/router/router_factory.py | 5 +- .../router/routing_simulator_router.py | 31 +- .../fused_moe/router/zero_expert_router.py | 3 +- .../layers/fused_moe/runner/moe_runner.py | 7 +- .../fused_moe/runner/moe_runner_interface.py | 3 +- .../fused_moe/topk_weight_and_reduce.py | 12 +- vllm/model_executor/layers/fused_moe/utils.py | 28 +- .../layers/fused_qk_norm_rope.py | 1 + .../layers/fusion/fused_act_quant.py | 3 +- .../layers/fusion/quant_activation.py | 3 +- vllm/model_executor/layers/hpc/gated_mla.py | 1 + vllm/model_executor/layers/hpc/hpc_ihc.py | 1 + vllm/model_executor/layers/indexer_topk.py | 10 +- vllm/model_executor/layers/layernorm.py | 5 +- vllm/model_executor/layers/lightning_attn.py | 11 +- vllm/model_executor/layers/linear.py | 18 +- .../model_executor/layers/logits_processor.py | 6 +- vllm/model_executor/layers/mamba/abstract.py | 6 +- .../layers/mamba/gdn/olmo_gdn_linear_attn.py | 7 +- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 25 +- .../mamba/linear/bailing_linear_attn.py | 2 +- .../layers/mamba/mamba_mixer.py | 7 +- .../layers/mamba/mamba_mixer2.py | 4 +- .../layers/mamba/mamba_utils.py | 8 +- .../layers/mamba/ops/causal_conv1d.py | 5 +- .../layers/mamba/ops/cpu/causal_conv1d.py | 3 +- .../mamba/ops/gdn_chunk_cutedsl/__init__.py | 1 + .../mamba/ops/gdn_chunk_cutedsl/kernel_h.py | 1 - .../layers/mamba/ops/mamba_ssm.py | 64 +- .../layers/mamba/ops/ssd_bmm.py | 5 +- .../layers/mamba/ops/ssd_combined.py | 4 +- .../layers/mamba/ops/ssu_dispatch.py | 3 +- .../minimax_rms_norm/lamport_workspace.py | 20 +- vllm/model_executor/layers/pooler/abstract.py | 4 +- .../layers/pooler/seqwise/poolers.py | 3 +- .../layers/pooler/tokwise/poolers.py | 3 +- .../layers/quantization/__init__.py | 1 + .../layers/quantization/auto_awq.py | 3 +- .../layers/quantization/auto_gptq.py | 2 +- .../layers/quantization/base_config.py | 28 +- .../compressed_tensors/compressed_tensors.py | 47 +- .../compressed_tensors_embedding.py | 3 +- .../compressed_tensors_moe_w4a4_nvfp4.py | 4 +- .../compressed_tensors_moe_w4a8_int8.py | 3 +- .../compressed_tensors_moe_wna16.py | 3 +- .../schemes/compressed_tensors_scheme.py | 19 +- .../schemes/compressed_tensors_w4a4_mxfp4.py | 3 +- .../schemes/compressed_tensors_w8a8_mxfp8.py | 3 +- .../schemes/compressed_tensors_wNa8o8.py | 3 +- .../compressed_tensors/transform/linear.py | 10 +- .../compressed_tensors/transform/module.py | 3 +- .../transform/schemes/linear_qutlass_nvfp4.py | 3 +- .../quantization/compressed_tensors/utils.py | 14 +- .../model_executor/layers/quantization/fp8.py | 6 +- .../layers/quantization/fp_quant.py | 1 + .../layers/quantization/humming.py | 3 +- .../layers/quantization/inc/config_parser.py | 3 +- .../quantization/inc/schemes/inc_scheme.py | 3 +- .../layers/quantization/input_quant_fp8.py | 27 +- .../layers/quantization/kv_cache.py | 7 +- .../layers/quantization/modelopt.py | 42 +- .../layers/quantization/moe_wna16.py | 1 + .../layers/quantization/mxfp4.py | 3 +- .../layers/quantization/online/base.py | 6 +- .../layers/quantization/online/fp8.py | 18 +- .../layers/quantization/quark/quark.py | 22 +- .../layers/quantization/quark/quark_moe.py | 5 +- .../quantization/quark/schemes/quark_nvfp4.py | 3 +- .../quark/schemes/quark_scheme.py | 19 +- .../quark/schemes/quark_w4a8_mxfp4_fp8.py | 3 +- .../layers/quantization/qutlass_utils.py | 12 +- .../layers/quantization/torchao.py | 11 +- .../quantization/turboquant/centroids.py | 1 + .../layers/quantization/turboquant/config.py | 1 + .../layers/quantization/utils/config_utils.py | 5 +- .../quantization/utils/flashinfer_fp4_moe.py | 2 +- .../utils/flashinfer_mxint4_moe.py | 12 +- .../quantization/utils/flashinfer_utils.py | 8 +- .../layers/quantization/utils/fp8_utils.py | 36 +- .../quantization/utils/humming_utils.py | 30 +- .../layers/quantization/utils/int8_utils.py | 1 + .../quantization/utils/machete_utils.py | 4 +- .../layers/quantization/utils/marlin_utils.py | 12 +- .../quantization/utils/marlin_utils_fp4.py | 13 +- .../quantization/utils/marlin_utils_fp8.py | 12 +- .../quantization/utils/marlin_utils_test.py | 2 +- .../layers/quantization/utils/mxfp4_utils.py | 5 +- .../utils/nvfp4_emulation_utils.py | 3 +- .../layers/quantization/utils/nvfp4_utils.py | 28 +- .../layers/quantization/utils/quant_utils.py | 38 +- vllm/model_executor/layers/resampler.py | 12 +- .../layers/rotary_embedding/common.py | 16 +- .../rotary_embedding/deepseek_scaling_rope.py | 1 - .../rotary_embedding/ernie45_vl_rope.py | 2 +- .../layers/rotary_embedding/fope.py | 2 +- .../layers/rotary_embedding/mrope.py | 6 + .../rotary_embedding/mrope_interleaved.py | 6 +- .../rotary_embedding/ntk_scaling_rope.py | 3 +- .../layers/sparse_attn_indexer.py | 7 +- .../layers/sparse_mqa_indexer.py | 6 +- vllm/model_executor/layers/utils.py | 4 +- .../layers/vocab_parallel_embedding.py | 5 +- vllm/model_executor/model_loader/__init__.py | 1 + .../model_loader/base_loader.py | 3 +- .../model_loader/checkpoint_weight_patch.py | 3 +- .../model_loader/default_loader.py | 3 +- .../model_loader/ep_weight_filter.py | 10 +- .../model_loader/reload/__init__.py | 3 +- .../model_loader/reload/layerwise.py | 28 +- .../model_loader/reload/meta.py | 12 +- .../model_loader/reload/sanitize.py | 8 +- .../model_loader/reload/torchao_decorator.py | 3 +- .../model_loader/reload/utils.py | 8 +- .../model_loader/runai_streamer_loader.py | 9 +- .../model_loader/sharded_state_loader.py | 6 +- .../model_executor/model_loader/tensorizer.py | 15 +- .../model_loader/tensorizer_loader.py | 3 +- vllm/model_executor/model_loader/utils.py | 3 +- .../model_loader/weight_cache/daemon.py | 2 + .../model_loader/weight_cache/ipc_loader.py | 3 +- .../model_loader/weight_cache/protocol.py | 2 + .../model_loader/weight_tying.py | 3 +- .../model_loader/weight_utils.py | 29 +- vllm/model_executor/models/AXK1.py | 3 +- vllm/model_executor/models/adapters.py | 15 +- vllm/model_executor/models/arcee.py | 15 +- vllm/model_executor/models/aria.py | 34 +- vllm/model_executor/models/audioflamingo3.py | 27 +- vllm/model_executor/models/bagel.py | 21 +- .../models/bailing_moe_linear.py | 5 +- vllm/model_executor/models/bert.py | 8 +- vllm/model_executor/models/blip.py | 12 +- vllm/model_executor/models/blip2.py | 26 +- vllm/model_executor/models/clip.py | 18 +- vllm/model_executor/models/cohere2_moe.py | 3 +- vllm/model_executor/models/cohere2_vision.py | 23 +- vllm/model_executor/models/cohere_asr.py | 54 +- vllm/model_executor/models/cohere_compass.py | 21 +- vllm/model_executor/models/colbert.py | 5 +- vllm/model_executor/models/colpali.py | 3 +- vllm/model_executor/models/colqwen3.py | 4 +- vllm/model_executor/models/colqwen3_5.py | 4 +- vllm/model_executor/models/commandr.py | 3 +- vllm/model_executor/models/config.py | 8 +- .../models/conformer_encoder.py | 7 +- vllm/model_executor/models/cosmos3_edge.py | 6 +- vllm/model_executor/models/deepencoder.py | 122 +- vllm/model_executor/models/deepencoder2.py | 18 +- vllm/model_executor/models/deepseek_eagle3.py | 3 +- vllm/model_executor/models/deepseek_mtp.py | 3 +- vllm/model_executor/models/deepseek_ocr.py | 32 +- vllm/model_executor/models/deepseek_ocr2.py | 4 +- vllm/model_executor/models/deepseek_v2.py | 9 +- vllm/model_executor/models/deepseek_vl2.py | 22 +- vllm/model_executor/models/diffusion_gemma.py | 4 +- vllm/model_executor/models/dots_ocr.py | 24 +- vllm/model_executor/models/eagle2_5_vl.py | 33 +- vllm/model_executor/models/ernie45_vl.py | 27 +- .../models/extract_hidden_states.py | 6 +- vllm/model_executor/models/falcon_h1.py | 7 +- vllm/model_executor/models/fireredasr2.py | 9 +- vllm/model_executor/models/funasr.py | 9 +- vllm/model_executor/models/gemma3_mm.py | 19 +- vllm/model_executor/models/gemma3n.py | 13 +- .../models/gemma3n_audio_utils.py | 4 +- vllm/model_executor/models/gemma3n_mm.py | 41 +- vllm/model_executor/models/gemma4.py | 2 + vllm/model_executor/models/gemma4_mm.py | 19 +- vllm/model_executor/models/gemma4_mtp.py | 1 + vllm/model_executor/models/glm4_1v.py | 60 +- vllm/model_executor/models/glm4_moe.py | 3 +- .../models/glm4_moe_lite_mtp.py | 3 +- vllm/model_executor/models/glm4_moe_mtp.py | 6 +- vllm/model_executor/models/glm4v.py | 32 +- vllm/model_executor/models/glm_ocr_mtp.py | 3 +- vllm/model_executor/models/glmasr.py | 78 +- vllm/model_executor/models/glmasr_utils.py | 4 +- vllm/model_executor/models/gpt2.py | 1 + vllm/model_executor/models/gpt_oss.py | 1 + vllm/model_executor/models/granite.py | 1 + vllm/model_executor/models/granite4_vision.py | 1 + vllm/model_executor/models/granite_speech.py | 12 +- .../model_executor/models/granitemoehybrid.py | 1 + vllm/model_executor/models/hrm_text.py | 3 +- .../models/hyperclovax_vision_v2.py | 18 +- .../models/idefics2_vision_model.py | 32 +- vllm/model_executor/models/idefics3.py | 26 +- vllm/model_executor/models/interfaces.py | 107 +- vllm/model_executor/models/intern_vit.py | 2 +- vllm/model_executor/models/interns1.py | 48 +- vllm/model_executor/models/interns1_pro.py | 2 +- vllm/model_executor/models/interns1_vit.py | 9 +- vllm/model_executor/models/internvl.py | 54 +- .../model_executor/models/iquest_loopcoder.py | 1 + vllm/model_executor/models/isaac.py | 24 +- vllm/model_executor/models/kanana_v.py | 9 +- vllm/model_executor/models/keye.py | 61 +- vllm/model_executor/models/keye_vl1_5.py | 69 +- vllm/model_executor/models/kimi_k25.py | 6 +- vllm/model_executor/models/kimi_k25_vit.py | 15 +- vllm/model_executor/models/kimi_vl.py | 11 +- vllm/model_executor/models/lfm2.py | 1 + vllm/model_executor/models/lfm2_moe.py | 1 + vllm/model_executor/models/lfm2_siglip2.py | 29 +- vllm/model_executor/models/lfm2_vl.py | 17 +- vllm/model_executor/models/llama4.py | 5 +- vllm/model_executor/models/llava.py | 23 +- vllm/model_executor/models/llava_next.py | 19 +- .../model_executor/models/llava_next_video.py | 19 +- vllm/model_executor/models/llava_onevision.py | 57 +- .../model_executor/models/llava_onevision2.py | 3 +- .../models/longcat_flash_mtp.py | 3 +- .../models/longcat_flash_ngram.py | 1 + vllm/model_executor/models/mamba2.py | 1 + vllm/model_executor/models/medusa.py | 3 +- vllm/model_executor/models/mellum.py | 12 +- vllm/model_executor/models/midashenglm.py | 8 +- vllm/model_executor/models/mimo_audio.py | 2 + vllm/model_executor/models/mimo_v2.py | 4 +- vllm/model_executor/models/mimo_v2_mtp.py | 2 +- vllm/model_executor/models/mimo_v2_omni.py | 25 +- vllm/model_executor/models/minicpmo.py | 19 +- vllm/model_executor/models/minicpmv.py | 36 +- vllm/model_executor/models/minicpmv4_6.py | 16 +- vllm/model_executor/models/mistral3.py | 23 +- vllm/model_executor/models/mllama4.py | 35 +- vllm/model_executor/models/mlp_speculator.py | 10 +- vllm/model_executor/models/molmo.py | 25 +- vllm/model_executor/models/molmo2.py | 64 +- vllm/model_executor/models/moondream3.py | 4 +- vllm/model_executor/models/moonvit.py | 71 +- vllm/model_executor/models/moss_audio.py | 11 +- .../models/moss_transcribe_diarize.py | 20 +- vllm/model_executor/models/muse_glimmer.py | 13 +- .../model_executor/models/nano_nemotron_vl.py | 68 +- vllm/model_executor/models/nemotron_h.py | 1 + vllm/model_executor/models/nemotron_parse.py | 33 +- vllm/model_executor/models/nemotron_vl.py | 7 +- vllm/model_executor/models/openpangu_mtp.py | 3 +- vllm/model_executor/models/openpangu_vl.py | 8 +- vllm/model_executor/models/openvla.py | 11 +- vllm/model_executor/models/ovis.py | 18 +- vllm/model_executor/models/ovis2_5.py | 29 +- vllm/model_executor/models/paddleocr_vl.py | 5 +- vllm/model_executor/models/paligemma.py | 20 +- vllm/model_executor/models/parakeet.py | 4 +- vllm/model_executor/models/param2moe.py | 12 +- vllm/model_executor/models/phi3.py | 2 +- vllm/model_executor/models/phi3v.py | 37 +- vllm/model_executor/models/phi4mm.py | 74 +- vllm/model_executor/models/phi4mm_audio.py | 27 +- vllm/model_executor/models/phi4mm_utils.py | 83 +- vllm/model_executor/models/phi4siglip.py | 9 +- vllm/model_executor/models/pixtral.py | 29 +- .../models/qwen2_5_omni_thinker.py | 47 +- vllm/model_executor/models/qwen2_5_vl.py | 38 +- vllm/model_executor/models/qwen2_audio.py | 18 +- vllm/model_executor/models/qwen2_vl.py | 24 +- vllm/model_executor/models/qwen3_5.py | 2 +- vllm/model_executor/models/qwen3_asr.py | 7 +- vllm/model_executor/models/qwen3_dflash.py | 12 +- .../models/qwen3_omni_moe_thinker.py | 35 +- vllm/model_executor/models/qwen3_vl.py | 27 +- vllm/model_executor/models/radio.py | 7 +- vllm/model_executor/models/registry.py | 16 +- vllm/model_executor/models/roberta.py | 1 + vllm/model_executor/models/sarvam.py | 1 + vllm/model_executor/models/siglip.py | 16 +- vllm/model_executor/models/siglip2navit.py | 66 +- vllm/model_executor/models/skyworkr1v.py | 24 +- vllm/model_executor/models/stablelm.py | 3 +- vllm/model_executor/models/step3_vl.py | 30 +- vllm/model_executor/models/step3p5.py | 4 +- vllm/model_executor/models/step3p5_mtp.py | 3 +- vllm/model_executor/models/step_vl.py | 2 +- vllm/model_executor/models/terratorch.py | 2 +- .../models/transformers/__init__.py | 2 +- .../models/transformers/base.py | 27 +- .../models/transformers/causal.py | 1 - .../models/transformers/fusers/base.py | 27 +- .../models/transformers/fusers/mla.py | 3 +- .../models/transformers/fusers/moe.py | 12 +- .../models/transformers/fusers/rms_norm.py | 2 +- .../models/transformers/fx_utils.py | 37 +- .../models/transformers/layers.py | 3 +- .../model_executor/models/transformers/moe.py | 6 +- .../models/transformers/multimodal.py | 37 +- .../models/transformers/pooling.py | 3 +- .../models/transformers/utils.py | 27 +- vllm/model_executor/models/ultravox.py | 23 +- vllm/model_executor/models/utils.py | 62 +- vllm/model_executor/models/vision.py | 27 +- vllm/model_executor/models/voxtral.py | 6 +- .../model_executor/models/voxtral_realtime.py | 4 +- vllm/model_executor/models/voyage.py | 3 +- vllm/model_executor/models/whisper.py | 18 +- vllm/model_executor/models/zamba2.py | 33 +- vllm/model_executor/offloader/base.py | 1 + vllm/model_executor/offloader/prefetch.py | 1 + vllm/model_executor/offloader/prefetch_ops.py | 2 + vllm/model_executor/offloader/uva.py | 1 + vllm/model_executor/parameter.py | 50 +- vllm/model_executor/utils.py | 10 +- .../model_executor/warmup/deep_gemm_warmup.py | 15 +- .../warmup/jit_warmup_tilelang_helper.py | 1 - vllm/model_executor/warmup/kernel_warmup.py | 6 +- .../warmup/qwen4_exp_qsa_warmup.py | 4 +- vllm/models/deepseek_v32/common/kernels.py | 3 +- vllm/models/deepseek_v4/amd/model.py | 3 +- vllm/models/deepseek_v4/amd/mtp.py | 3 +- vllm/models/deepseek_v4/attention.py | 4 +- .../deepseek_v4/common/mm_preprocess.py | 9 +- .../deepseek_v4/common/ops/cache_utils.py | 9 +- .../common/ops/fused_compress_quant_cache.py | 7 +- .../deepseek_v4/common/ops/fused_indexer_q.py | 3 +- .../common/ops/fused_inv_rope_fp8_quant.py | 4 +- vllm/models/deepseek_v4/common/vision.py | 1 + vllm/models/deepseek_v4/cpu/cpu_compressor.py | 3 +- vllm/models/deepseek_v4/cpu/cpu_mla.py | 3 +- vllm/models/deepseek_v4/cpu/cpu_sparse.py | 24 +- vllm/models/deepseek_v4/cpu/cpu_utils.py | 1 + vllm/models/deepseek_v4/cpu/model.py | 3 +- vllm/models/deepseek_v4/nvidia/fi_moe.py | 3 +- .../deepseek_v4/nvidia/flashinfer_sparse.py | 3 +- vllm/models/deepseek_v4/nvidia/model.py | 3 +- vllm/models/deepseek_v4/nvidia/mtp.py | 3 +- vllm/models/deepseek_v4/xpu/model.py | 3 +- vllm/models/deepseek_v4/xpu/mtp.py | 3 +- vllm/models/deepseek_v41/amd/model.py | 6 +- vllm/models/deepseek_v41/amd/vl_model.py | 3 +- vllm/models/deepseek_v41/attention.py | 7 +- vllm/models/deepseek_v41/common/engram.py | 9 +- .../deepseek_v41/common/mm_preprocess.py | 9 +- .../deepseek_v41/common/ops/cache_utils.py | 9 +- .../common/ops/fused_compress_quant_cache.py | 1 + .../common/ops/indexer_k_store.py | 2 + .../deepseek_v41/nvidia/flashinfer_sparse.py | 3 +- vllm/models/deepseek_v41/nvidia/model.py | 6 +- vllm/models/deepseek_v41/nvidia/vl_model.py | 3 +- .../models/dots3_note/nvidia/audio_encoder.py | 10 +- .../models/glm5next/amd/ops/kpool_compress.py | 14 +- .../ops/third_party/kda/fused_recurrent.py | 4 +- .../amd/ops/third_party/kda/kernels.py | 29 +- vllm/models/glm5next/common/kda.py | 6 +- vllm/models/glm5next/common/multimodal.py | 3 +- vllm/models/glm5next/common/sparse_indexer.py | 6 +- .../glm5next/nvidia/ops/fused_eh_norm.py | 3 +- .../glm5next/nvidia/ops/kpool_compress.py | 15 +- .../ops/third_party/kda/fused_recurrent.py | 4 +- .../nvidia/ops/third_party/kda/kernels.py | 29 +- vllm/models/hy_v4/nvidia/attention.py | 3 + vllm/models/hy_v4/nvidia/flashmla_sparse.py | 2 + vllm/models/hy_v4/nvidia/hc.py | 4 + vllm/models/hy_v4/nvidia/model.py | 2 +- vllm/models/hy_v4/nvidia/mtp.py | 4 + vllm/models/inkling/amd/logits_processor.py | 1 + vllm/models/inkling/amd/model.py | 3 +- vllm/models/inkling/amd/moe.py | 9 +- vllm/models/inkling/amd/mtp.py | 9 +- vllm/models/inkling/amd/ops/norm.py | 6 +- .../models/inkling/nvidia/logits_processor.py | 1 + vllm/models/inkling/nvidia/model.py | 3 +- vllm/models/inkling/nvidia/moe.py | 9 +- vllm/models/inkling/nvidia/mtp.py | 9 +- vllm/models/inkling/nvidia/ops/mm_towers.py | 6 +- vllm/models/inkling/nvidia/ops/norm.py | 6 +- vllm/models/kimi_k3/amd/latent_moe_runner.py | 4 +- vllm/models/kimi_k3/amd/linear.py | 4 +- vllm/models/kimi_k3/amd/ops/kda_chunk.py | 9 + vllm/models/kimi_k3/amd/ops/kda_prefill.py | 17 + .../kimi_k3/amd/ops/third_party/kda/chunk.py | 13 +- .../amd/ops/third_party/kda/chunk_intra.py | 3 +- .../kda/chunk_intra_token_parallel.py | 5 +- .../kimi_k3/nvidia/latent_moe_runner.py | 4 +- .../models/kimi_k3/nvidia/low_latency_gemm.py | 1 + vllm/models/kimi_k3/nvidia/mla.py | 3 +- vllm/models/kimi_k3/nvidia/model.py | 2 + ...educe_rmsnorm_reduce_scatter_early_exit.py | 3 - .../fused_add_multicast_gemm.py | 1 - .../cute_dsl/latent_moe_tail/primitives.py | 13 - .../nvidia/ops/third_party/kda/chunk.py | 17 +- .../nvidia/ops/third_party/kda/chunk_intra.py | 298 +- .../kda/chunk_intra_token_parallel.py | 56 +- .../ops/third_party/kda/fused_recurrent.py | 4 +- vllm/models/minimax_m3/amd/model.py | 3 +- vllm/models/minimax_m3/common/indexer.py | 6 +- .../minimax_m3/common/sparse_attention.py | 3 +- vllm/models/minimax_m3/nvidia/indexer_msa.py | 6 +- vllm/models/minimax_m3/nvidia/model.py | 6 +- vllm/models/qwen4_exp/amd/indexer_qsa.py | 5 - vllm/models/qwen4_exp/amd/low_latency_gemm.py | 1 - vllm/models/qwen4_exp/amd/model.py | 2 - vllm/models/qwen4_exp/amd/mtp.py | 4 - vllm/models/qwen4_exp/amd/ops/qsa.py | 6 - vllm/models/qwen4_exp/amd/ple_layer.py | 1 - vllm/models/qwen4_exp/amd/qsa.py | 1 - vllm/models/qwen4_exp/common/ple.py | 2 - vllm/models/qwen4_exp/common/qsa_cache.py | 3 - vllm/models/qwen4_exp/nvidia/indexer_qsa.py | 2 - vllm/models/qwen4_exp/nvidia/model.py | 2 - vllm/models/qwen4_exp/nvidia/mtp.py | 4 - .../qwen4_exp/nvidia/ngram_embedding.py | 1 - vllm/models/qwen4_exp/nvidia/ops/qsa.py | 3 - .../qwen4_exp/nvidia/ops/qsa_indexer.py | 6 +- vllm/models/qwen4_exp/nvidia/ple_layer.py | 1 - vllm/multimodal/audio.py | 7 + vllm/multimodal/cache.py | 68 +- vllm/multimodal/encoder_budget.py | 3 +- vllm/multimodal/gpu_ipc_memory.py | 1 + vllm/multimodal/hasher.py | 5 +- vllm/multimodal/image.py | 3 +- vllm/multimodal/inputs.py | 77 +- vllm/multimodal/media/audio.py | 15 +- vllm/multimodal/media/base.py | 6 +- vllm/multimodal/media/connector.py | 66 +- vllm/multimodal/parse.py | 50 +- vllm/multimodal/processing/context.py | 48 +- vllm/multimodal/processing/dummy_inputs.py | 15 +- vllm/multimodal/processing/inputs.py | 3 +- vllm/multimodal/processing/processor.py | 85 +- vllm/multimodal/registry.py | 26 +- vllm/multimodal/utils.py | 33 +- vllm/multimodal/video.py | 13 +- vllm/multimodal/video_decoders/opencv.py | 4 +- vllm/multimodal/video_prune/evs.py | 17 +- vllm/multimodal/video_prune/vidcom2.py | 1 + vllm/outputs.py | 11 +- vllm/parser/abstract_parser.py | 43 +- vllm/parser/cohere_command.py | 12 +- vllm/parser/engine/__init__.py | 3 +- vllm/parser/engine/incremental_lexer.py | 3 +- vllm/parser/engine/streaming_parser_engine.py | 3 +- vllm/parser/engine/token_id_scanner.py | 3 +- vllm/parser/gemma4.py | 1 + vllm/parser/inkling.py | 3 +- vllm/parser/mistral.py | 2 + vllm/parser/parser_manager.py | 8 +- vllm/platforms/__init__.py | 4 +- vllm/platforms/cpu.py | 15 +- vllm/platforms/cuda.py | 11 +- vllm/platforms/interface.py | 186 +- vllm/platforms/rocm.py | 21 +- vllm/platforms/xpu.py | 8 +- vllm/plugins/__init__.py | 1 + .../lora_resolvers/filesystem_resolver.py | 3 +- .../plugins/lora_resolvers/hf_hub_resolver.py | 6 +- vllm/pooling_params.py | 2 + vllm/profiler/layerwise_profile.py | 6 +- vllm/profiler/wrapper.py | 5 +- vllm/ray/lazy_utils.py | 1 - vllm/ray/ray_env.py | 1 + vllm/reasoning/abs_reasoning_parsers.py | 64 +- vllm/reasoning/basic_parsers.py | 13 +- .../reasoning/deepseek_r1_reasoning_parser.py | 3 +- .../reasoning/deepseek_v3_reasoning_parser.py | 7 +- vllm/reasoning/ernie45_reasoning_parser.py | 9 +- vllm/reasoning/gemma4_utils.py | 1 + vllm/reasoning/gptoss_reasoning_parser.py | 3 +- vllm/reasoning/granite_reasoning_parser.py | 11 +- .../hunyuan_a13b_reasoning_parser.py | 7 +- vllm/reasoning/hy_v3_reasoning_parser.py | 3 +- vllm/reasoning/hy_v4_reasoning_parser.py | 2 +- vllm/reasoning/identity_reasoning_parser.py | 3 +- vllm/reasoning/minimax_m2_reasoning_parser.py | 7 +- vllm/reasoning/olmo3_reasoning_parser.py | 13 +- .../reasoning/poolside_v1_reasoning_parser.py | 3 +- vllm/reasoning/step3_reasoning_parser.py | 6 +- vllm/reasoning/step3p5_reasoning_parser.py | 3 +- vllm/renderers/base.py | 9 +- vllm/renderers/hf.py | 8 +- vllm/renderers/inputs/preprocess.py | 15 +- vllm/renderers/inputs/tokenize.py | 7 +- vllm/renderers/online_derenderer.py | 6 +- vllm/renderers/params.py | 6 +- vllm/sampling_params.py | 13 +- vllm/scalar_type.py | 37 +- .../ops/fused_norm_gate.py | 1 + .../flash_linear_attention/ops/kda.py | 1 + .../ops/layernorm_guard.py | 4 +- vllm/third_party/pynvml.py | 4446 ++++++++++------- vllm/tilelang_utils/__init__.py | 1 + vllm/tokenizers/deepseek_v32.py | 4 +- vllm/tokenizers/deepseek_v4.py | 4 +- vllm/tokenizers/detokenizer_utils.py | 1 + vllm/tokenizers/hf.py | 6 +- vllm/tool_parsers/abstract_tool_parser.py | 22 +- vllm/tool_parsers/apertus_tool_parser.py | 46 +- vllm/tool_parsers/ernie45_tool_parser.py | 3 +- .../tool_parsers/functiongemma_tool_parser.py | 3 +- vllm/tool_parsers/gemma4_utils.py | 3 + vllm/tool_parsers/gptoss_tool_parser.py | 3 +- .../granite_20b_fc_tool_parser.py | 3 +- vllm/tool_parsers/granite_tool_parser.py | 3 +- vllm/tool_parsers/hermes_tool_parser.py | 3 +- vllm/tool_parsers/hunyuan_a13b_tool_parser.py | 9 +- vllm/tool_parsers/hy_v3_tool_parser.py | 1 + vllm/tool_parsers/hy_v4_tool_parser.py | 5 +- vllm/tool_parsers/lfm2_tool_parser.py | 6 +- .../llama4_pythonic_tool_parser.py | 8 +- vllm/tool_parsers/llama_tool_parser.py | 6 +- vllm/tool_parsers/olmo3_tool_parser.py | 7 +- vllm/tool_parsers/phi4mini_tool_parser.py | 7 +- vllm/tool_parsers/poolside_v1_tool_parser.py | 3 +- vllm/tool_parsers/pythonic_tool_parser.py | 7 +- vllm/tool_parsers/step3_tool_parser.py | 3 +- vllm/tool_parsers/step3p5_tool_parser.py | 75 +- vllm/tool_parsers/structural_tag_registry.py | 4 +- vllm/tool_parsers/utils.py | 18 +- vllm/tracing/__init__.py | 8 +- vllm/tracing/otel.py | 11 +- vllm/tracing/utils.py | 6 +- vllm/transformers_utils/config.py | 45 +- vllm/transformers_utils/configs/AXK1.py | 5 +- vllm/transformers_utils/configs/__init__.py | 3 +- vllm/transformers_utils/configs/colpali.py | 3 +- vllm/transformers_utils/configs/colqwen3.py | 3 +- .../configs/extract_hidden_states.py | 3 +- vllm/transformers_utils/configs/falcon.py | 2 +- vllm/transformers_utils/configs/hy_v3.py | 4 +- vllm/transformers_utils/configs/hy_v4.py | 1 + .../transformers_utils/configs/hyperclovax.py | 4 +- vllm/transformers_utils/configs/kimi_k25.py | 4 +- vllm/transformers_utils/configs/lfm2_moe.py | 7 +- .../configs/mlp_speculator.py | 4 +- vllm/transformers_utils/configs/moondream3.py | 1 + vllm/transformers_utils/configs/nemotron.py | 13 +- vllm/transformers_utils/configs/nemotron_h.py | 7 +- .../transformers_utils/configs/olmo_hybrid.py | 6 +- vllm/transformers_utils/configs/ovis.py | 2 + vllm/transformers_utils/configs/qwen3_5.py | 2 +- .../transformers_utils/configs/qwen3_5_moe.py | 2 +- vllm/transformers_utils/configs/qwen3_next.py | 6 +- vllm/transformers_utils/configs/radio.py | 6 +- .../configs/speculators/algos.py | 16 +- .../configs/speculators/base.py | 8 +- vllm/transformers_utils/configs/ultravox.py | 4 +- vllm/transformers_utils/dynamic_module.py | 3 +- vllm/transformers_utils/processor.py | 17 +- .../transformers_utils/processors/__init__.py | 3 +- vllm/transformers_utils/processors/bagel.py | 15 +- .../processors/cohere_asr.py | 5 +- .../processors/cosmos3_edge.py | 4 +- .../processors/deepseek_ocr.py | 7 +- .../processors/deepseek_vl2.py | 12 +- .../processors/fireredasr2.py | 13 +- vllm/transformers_utils/processors/funasr.py | 16 +- vllm/transformers_utils/processors/glm4v.py | 3 +- .../transformers_utils/processors/glm5next.py | 6 +- .../transformers_utils/processors/internvl.py | 3 +- vllm/transformers_utils/processors/isaac.py | 5 + .../transformers_utils/processors/kimi_k25.py | 4 +- .../processors/mimo_v2_omni.py | 3 + .../transformers_utils/processors/minicpmo.py | 22 +- .../transformers_utils/processors/minicpmv.py | 14 +- .../processors/moondream3.py | 15 +- .../processors/nano_nemotron_vl.py | 15 +- .../processors/nemotron_vl.py | 6 +- vllm/transformers_utils/processors/ovis.py | 28 +- vllm/transformers_utils/processors/ovis2_5.py | 64 +- vllm/transformers_utils/processors/pixtral.py | 3 +- vllm/transformers_utils/processors/voxtral.py | 3 +- vllm/transformers_utils/repo_utils.py | 18 +- vllm/transformers_utils/runai_utils.py | 11 +- vllm/transformers_utils/s3_utils.py | 8 +- vllm/transformers_utils/utils.py | 8 +- vllm/utils/argparse_utils.py | 8 +- vllm/utils/async_utils.py | 10 +- vllm/utils/cache.py | 10 +- vllm/utils/collection_utils.py | 9 +- vllm/utils/counter.py | 8 +- vllm/utils/cpu_resource_utils.py | 6 +- vllm/utils/cpu_triton_utils.py | 4 +- vllm/utils/deep_gemm.py | 20 +- vllm/utils/flashinfer.py | 5 +- vllm/utils/func_utils.py | 8 +- vllm/utils/gc_utils.py | 24 +- vllm/utils/gpu_sync_debug.py | 6 +- vllm/utils/hashing.py | 4 + vllm/utils/import_utils.py | 32 +- vllm/utils/jit_monitor.py | 3 +- vllm/utils/jsontree.py | 3 +- vllm/utils/math_utils.py | 2 +- vllm/utils/mem_utils.py | 4 +- vllm/utils/multi_stream_utils.py | 2 + vllm/utils/network_utils.py | 8 +- vllm/utils/nvtx_pytorch_hooks.py | 5 +- vllm/utils/ompmultiprocessing.py | 13 +- vllm/utils/platform_utils.py | 3 +- vllm/utils/registry.py | 15 +- vllm/utils/sparse_utils.py | 3 +- vllm/utils/system_utils.py | 9 +- vllm/utils/tensor_schema.py | 1 - vllm/utils/torch_utils.py | 44 +- vllm/v1/attention/backend.py | 49 +- vllm/v1/attention/backends/cpu_attn.py | 11 +- vllm/v1/attention/backends/fa_utils.py | 1 + vllm/v1/attention/backends/flash_attn.py | 12 +- .../attention/backends/flash_attn_diffkv.py | 8 + vllm/v1/attention/backends/flashinfer.py | 18 +- vllm/v1/attention/backends/flex_attention.py | 21 +- vllm/v1/attention/backends/gdn_attn.py | 3 +- vllm/v1/attention/backends/mamba2_attn.py | 3 +- vllm/v1/attention/backends/mamba_attn.py | 19 +- vllm/v1/attention/backends/mla/amx_mla.py | 3 +- .../attention/backends/mla/flashmla_sparse.py | 4 +- vllm/v1/attention/backends/mla/indexer.py | 3 +- .../backends/mla/prefill/aiter_flash_attn.py | 2 +- .../v1/attention/backends/mla/prefill/base.py | 3 +- .../backends/mla/prefill/registry.py | 3 + .../backends/mla/prefill/selector.py | 3 + .../attention/backends/mla/rocm_aiter_mla.py | 4 +- .../backends/mla/rocm_aiter_mla_sparse.py | 3 +- .../v1/attention/backends/mla/sparse_utils.py | 4 +- vllm/v1/attention/backends/registry.py | 10 +- vllm/v1/attention/backends/rocm_aiter_fa.py | 11 +- .../backends/rocm_aiter_unified_attn.py | 10 +- vllm/v1/attention/backends/rocm_attn.py | 11 +- vllm/v1/attention/backends/triton_attn.py | 11 +- vllm/v1/attention/backends/utils.py | 37 +- vllm/v1/attention/ops/common.py | 6 +- vllm/v1/attention/ops/dcp.py | 33 +- vllm/v1/attention/ops/flashmla.py | 8 +- vllm/v1/attention/ops/merge_attn_states.py | 2 +- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 4 +- .../attention/ops/triton_decode_attention.py | 3 +- .../attention/ops/triton_prefill_attention.py | 6 +- .../turboquant_soa/triton_turboquant_store.py | 3 +- vllm/v1/attention/ops/vit_attn_wrappers.py | 6 +- vllm/v1/attention/ops/xpu_mla_sparse.py | 9 +- vllm/v1/attention/selector.py | 2 +- vllm/v1/core/block_pool.py | 39 +- vllm/v1/core/encoder_cache_manager.py | 10 +- vllm/v1/core/kv_cache_coordinator.py | 62 +- vllm/v1/core/kv_cache_manager.py | 29 +- vllm/v1/core/kv_cache_utils.py | 117 +- vllm/v1/core/sched/interface.py | 14 +- vllm/v1/core/sched/request_queue.py | 12 +- vllm/v1/core/sched/scheduler.py | 31 +- vllm/v1/core/sched/utils.py | 4 + vllm/v1/core/single_type_kv_cache_manager.py | 92 +- vllm/v1/cudagraph_dispatcher.py | 14 +- vllm/v1/engine/__init__.py | 10 +- vllm/v1/engine/async_llm.py | 46 +- vllm/v1/engine/core.py | 43 +- vllm/v1/engine/core_client.py | 46 +- vllm/v1/engine/detokenizer.py | 7 +- vllm/v1/engine/input_processor.py | 3 +- vllm/v1/engine/llm_engine.py | 2 - vllm/v1/engine/logprobs.py | 8 +- vllm/v1/engine/output_processor.py | 8 +- vllm/v1/engine/parallel_sampling.py | 4 +- vllm/v1/engine/tensor_ipc.py | 4 +- vllm/v1/engine/utils.py | 26 +- vllm/v1/executor/abstract.py | 57 +- vllm/v1/executor/multiproc_executor.py | 19 +- vllm/v1/executor/ray_executor.py | 9 +- vllm/v1/executor/ray_executor_v2.py | 5 +- vllm/v1/executor/ray_utils.py | 13 +- vllm/v1/kv_cache_interface.py | 63 +- vllm/v1/kv_cache_layout.py | 3 +- vllm/v1/kv_cache_spec_registry.py | 27 +- vllm/v1/kv_offload/base.py | 71 +- vllm/v1/kv_offload/cpu/gpu_worker.py | 33 +- vllm/v1/kv_offload/cpu/manager.py | 3 +- vllm/v1/kv_offload/cpu/policies/arc.py | 3 +- vllm/v1/kv_offload/cpu/policies/base.py | 20 +- vllm/v1/kv_offload/cpu/policies/factory.py | 1 + vllm/v1/kv_offload/cpu/policies/lru.py | 3 +- .../kv_offload/cpu/shared_offload_region.py | 5 +- vllm/v1/kv_offload/cpu/spec.py | 6 +- vllm/v1/kv_offload/file_mapper.py | 10 +- vllm/v1/kv_offload/tiering/async_lookup.py | 16 +- vllm/v1/kv_offload/tiering/base.py | 48 +- vllm/v1/kv_offload/tiering/example/manager.py | 32 +- vllm/v1/kv_offload/tiering/factory.py | 1 + vllm/v1/kv_offload/tiering/fs/io.py | 13 +- vllm/v1/kv_offload/tiering/fs/manager.py | 40 +- vllm/v1/kv_offload/tiering/fs/thread_pool.py | 19 +- vllm/v1/kv_offload/tiering/manager.py | 55 +- vllm/v1/kv_offload/tiering/obj/manager.py | 32 +- .../v1/kv_offload/tiering/p2p/control/base.py | 2 + vllm/v1/kv_offload/tiering/p2p/control/zmq.py | 3 +- vllm/v1/kv_offload/tiering/p2p/data/base.py | 5 + vllm/v1/kv_offload/tiering/p2p/data/nixl.py | 4 +- vllm/v1/kv_offload/tiering/p2p/manager.py | 10 +- .../tiering/p2p/session/protocol.py | 3 +- .../kv_offload/tiering/p2p/session/session.py | 3 +- vllm/v1/kv_offload/tiering/spec.py | 13 +- vllm/v1/metrics/buckets.py | 3 + vllm/v1/metrics/loggers.py | 20 +- vllm/v1/metrics/perf.py | 81 +- vllm/v1/metrics/prometheus.py | 2 +- vllm/v1/metrics/ray_wrappers.py | 22 +- vllm/v1/metrics/reader.py | 1 + vllm/v1/metrics/stats.py | 12 +- vllm/v1/outputs.py | 4 +- vllm/v1/request.py | 3 +- vllm/v1/sample/logits_processor/__init__.py | 9 +- vllm/v1/sample/logits_processor/builtin.py | 9 +- vllm/v1/sample/logits_processor/interface.py | 3 +- vllm/v1/sample/logits_processor/state.py | 9 +- vllm/v1/sample/ops/logprobs.py | 4 +- vllm/v1/sample/ops/penalties.py | 8 +- vllm/v1/sample/ops/topk_topp_sampler.py | 12 +- vllm/v1/sample/ops/topk_topp_triton.py | 7 +- vllm/v1/sample/rejection_sampler.py | 19 +- vllm/v1/sample/sampler.py | 9 +- vllm/v1/serial_utils.py | 15 +- vllm/v1/simple_kv_offload/manager.py | 13 +- vllm/v1/simple_kv_offload/metadata.py | 3 +- vllm/v1/simple_kv_offload/worker.py | 2 + vllm/v1/spec_decode/dflash.py | 3 +- vllm/v1/spec_decode/extract_hidden_states.py | 10 +- vllm/v1/spec_decode/gemma4.py | 6 +- vllm/v1/spec_decode/llm_base_proposer.py | 33 +- vllm/v1/spec_decode/medusa.py | 4 +- vllm/v1/spec_decode/ngram_proposer.py | 6 +- vllm/v1/spec_decode/ngram_proposer_gpu.py | 28 +- vllm/v1/spec_decode/step3p5.py | 1 - vllm/v1/spec_decode/suffix_decoding.py | 3 +- vllm/v1/spec_decode/utils.py | 25 +- vllm/v1/structured_output/backend_guidance.py | 1 - vllm/v1/structured_output/backend_outlines.py | 5 +- vllm/v1/structured_output/backend_types.py | 36 +- vllm/v1/structured_output/utils.py | 23 +- vllm/v1/utils.py | 26 +- vllm/v1/worker/block_table.py | 29 +- vllm/v1/worker/dp_utils.py | 9 +- .../worker/ec_connector_model_runner_mixin.py | 4 +- vllm/v1/worker/encoder_cudagraph.py | 9 +- vllm/v1/worker/gpu/cudagraph_utils.py | 15 +- vllm/v1/worker/gpu/dp_utils.py | 10 +- vllm/v1/worker/gpu/lora_utils.py | 3 +- vllm/v1/worker/gpu/mm/encoder_cache.py | 3 +- vllm/v1/worker/gpu/model_runner.py | 6 +- vllm/v1/worker/gpu/model_states/interface.py | 3 +- vllm/v1/worker/gpu/pp_utils.py | 7 +- vllm/v1/worker/gpu/sample/output.py | 3 +- vllm/v1/worker/gpu/sample/sampler.py | 3 +- vllm/v1/worker/gpu/sample/watermark.py | 1 + .../gpu/spec_decode/adaptive_verification.py | 11 +- .../spec_decode/autoregressive/speculator.py | 3 +- .../gpu/spec_decode/dflash/cudagraph.py | 3 +- vllm/v1/worker/gpu/spec_decode/eagle/utils.py | 4 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 3 +- vllm/v1/worker/gpu_input_batch.py | 27 +- vllm/v1/worker/gpu_model_runner.py | 151 +- vllm/v1/worker/gpu_ubatch_wrapper.py | 3 +- vllm/v1/worker/gpu_worker.py | 20 +- .../worker/kv_connector_model_runner_mixin.py | 4 +- vllm/v1/worker/lora_model_runner_mixin.py | 14 +- vllm/v1/worker/mamba_utils.py | 20 +- vllm/v1/worker/startup_plan.py | 9 +- vllm/v1/worker/ubatch_utils.py | 15 +- vllm/v1/worker/ubatching.py | 4 +- vllm/v1/worker/utils.py | 27 +- vllm/v1/worker/worker_base.py | 17 +- vllm/v1/worker/workspace.py | 6 + vllm/vllm_flash_attn/flash_attn_interface.py | 39 + 2002 files changed, 15024 insertions(+), 14269 deletions(-) diff --git a/.buildkite/lm-eval-harness/test_lm_eval_correctness.py b/.buildkite/lm-eval-harness/test_lm_eval_correctness.py index fe9e332bdf..5576c84b05 100644 --- a/.buildkite/lm-eval-harness/test_lm_eval_correctness.py +++ b/.buildkite/lm-eval-harness/test_lm_eval_correctness.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -LM eval harness on model to compare vs HF baseline computed offline. +"""LM eval harness on model to compare vs HF baseline computed offline. Configs are found in configs/$MODEL.yaml pytest -s -v test_lm_eval_correctness.py \ diff --git a/.buildkite/performance-benchmarks/scripts/compare-json-results.py b/.buildkite/performance-benchmarks/scripts/compare-json-results.py index c9f8139fe6..823c7488ec 100644 --- a/.buildkite/performance-benchmarks/scripts/compare-json-results.py +++ b/.buildkite/performance-benchmarks/scripts/compare-json-results.py @@ -83,8 +83,7 @@ def compare_data_columns( drop_column: str, debug: bool = False, ): - """ - Align concatenation by keys derived from info_cols instead of row order. + """Align concatenation by keys derived from info_cols instead of row order. - Pick one canonical key list: subset of info_cols present in ALL files. - For each file: set index to those keys, aggregate duplicates (mean for metric, first for names). @@ -363,8 +362,7 @@ def _apply_two_decimals( # Export helpers (Excel + CSV) # ----------------------------- def _sanitize_sheet_name(name: str) -> str: - """ - Excel sheet constraints: + """Excel sheet constraints: - max 31 chars - cannot contain: : \ / ? * [ ] - cannot be empty diff --git a/.buildkite/performance-benchmarks/scripts/convert-results-json-to-markdown.py b/.buildkite/performance-benchmarks/scripts/convert-results-json-to-markdown.py index e694c38780..deb14a65e4 100644 --- a/.buildkite/performance-benchmarks/scripts/convert-results-json-to-markdown.py +++ b/.buildkite/performance-benchmarks/scripts/convert-results-json-to-markdown.py @@ -93,8 +93,7 @@ def results_to_json(latency, throughput, serving): def get_size_with_unit(bytes, suffix="B"): - """ - Scale bytes to its proper format + """Scale bytes to its proper format e.g: 1253656 => '1.20MB' 1253656678 => '1.17GB' diff --git a/.buildkite/scripts/generate-nightly-index.py b/.buildkite/scripts/generate-nightly-index.py index 67adf9c92b..8f92996ff9 100644 --- a/.buildkite/scripts/generate-nightly-index.py +++ b/.buildkite/scripts/generate-nightly-index.py @@ -18,8 +18,7 @@ import regex as re def normalize_package_name(name: str) -> str: - """ - Normalize package name according to PEP 503. + """Normalize package name according to PEP 503. https://peps.python.org/pep-0503/#normalized-names Replace runs of underscores, hyphens, and periods with a single hyphen, @@ -55,17 +54,18 @@ class WheelFileInfo: def parse_from_filename(file: str) -> WheelFileInfo: - """ - Parse wheel file name to extract metadata. + """Parse wheel file name to extract metadata. The format of wheel names: {package_name}-{version}(-{build_tag})?-{python_tag}-{abi_tag}-{platform_tag}.whl All versions could contain a variant like '+cu129' or '.cpu' or `.rocm` (or not). + Example: vllm-0.11.0-cp38-abi3-manylinux1_x86_64.whl vllm-0.10.2rc2+cu129-cp38-abi3-manylinux2014_aarch64.whl vllm-0.11.1rc8.dev14+gaa384b3c0-cp38-abi3-manylinux2014_aarch64.whl vllm-0.11.1rc8.dev14+gaa384b3c0.cu130-cp38-abi3-manylinux1_x86_64.whl + """ wheel_file_re = re.compile( r"^(?P.+)-(?P[^-]+?)(-(?P[^-]+))?-(?P[^-]+)-(?P[^-]+)-(?P[^-]+)\.whl$" @@ -111,9 +111,7 @@ def parse_from_filename(file: str) -> WheelFileInfo: def generate_project_list(subdir_names: list[str], comment: str = "") -> str: - """ - Generate project list HTML content linking to each project & variant subdirectory. - """ + """Generate project list HTML content linking to each project & variant subdirectory.""" href_tags = [] for name in sorted(subdir_names): name = name.strip("/").strip(".") @@ -127,9 +125,7 @@ def generate_package_index_and_metadata( index_base_dir: Path, comment: str = "", ) -> tuple[str, str]: - """ - Generate package index HTML content for a specific package, linking to actual wheel files. - """ + """Generate package index HTML content for a specific package, linking to actual wheel files.""" href_tags = [] metadata = [] for file in sorted(wheel_files, key=lambda x: x.filename): @@ -156,8 +152,7 @@ def generate_index_and_metadata( alias_to_default: str | None = None, comment: str = "", ): - """ - Generate index for all wheel files. + """Generate index for all wheel files. Args: whl_files (list[str]): List of wheel files (must be directly under `wheel_base_dir`). @@ -216,8 +211,8 @@ def generate_index_and_metadata( }, ... ] - """ + """ parsed_files = [parse_from_filename(f) for f in whl_files] if not parsed_files: diff --git a/benchmarks/attention_benchmarks/batch_spec.py b/benchmarks/attention_benchmarks/batch_spec.py index 9f15f1d809..85c689b489 100644 --- a/benchmarks/attention_benchmarks/batch_spec.py +++ b/benchmarks/attention_benchmarks/batch_spec.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Simplified batch specification grammar for attention benchmarks. +"""Simplified batch specification grammar for attention benchmarks. Grammar (underscore-separated segments): Format: (?) q(k?) (s(k?))? @@ -24,6 +23,7 @@ Examples: q4s1k -> [(4, 1024)] # 4-token extend (spec decode) 2q1k_32q1s1k -> [(1024, 1024)] * 2 + [(1, 1024)] * 32 # Mixed batch 16q4s1k -> [(4, 1024)] * 16 # 16 spec decode requests + """ from collections import Counter @@ -71,8 +71,7 @@ def _parse_size(size_str: str, k_suffix: str) -> int: def parse_batch_spec(spec: str) -> list[BatchRequest]: - """ - Parse batch specification string into list of BatchRequest objects. + """Parse batch specification string into list of BatchRequest objects. Grammar: (?) q(k?) (s(k?))? @@ -84,6 +83,7 @@ def parse_batch_spec(spec: str) -> list[BatchRequest]: Raises: ValueError: If spec format is invalid + """ requests = [] @@ -103,8 +103,7 @@ def parse_batch_spec(spec: str) -> list[BatchRequest]: def format_batch_spec(requests: list[BatchRequest]) -> str: - """ - Format list of BatchRequest into human-readable string. + """Format list of BatchRequest into human-readable string. Groups requests by type and provides counts and sizes. @@ -113,6 +112,7 @@ def format_batch_spec(requests: list[BatchRequest]) -> str: Returns: Formatted string describing the batch + """ kinds = { "prefill": [], @@ -157,8 +157,7 @@ def format_batch_spec(requests: list[BatchRequest]) -> str: def reorder_for_flashinfer(requests: list[BatchRequest]) -> list[BatchRequest]: - """ - Reorder requests for FlashInfer: decode first, then prefill. + """Reorder requests for FlashInfer: decode first, then prefill. FlashInfer expects decode requests before prefill requests for optimal performance. @@ -168,6 +167,7 @@ def reorder_for_flashinfer(requests: list[BatchRequest]) -> list[BatchRequest]: Returns: Reordered list with decode requests first + """ decodes = [r for r in requests if r.is_decode] non_decodes = [r for r in requests if not r.is_decode] @@ -177,14 +177,14 @@ def reorder_for_flashinfer(requests: list[BatchRequest]) -> list[BatchRequest]: def split_by_type( requests: list[BatchRequest], ) -> dict[str, list[BatchRequest]]: - """ - Split requests by type for analysis. + """Split requests by type for analysis. Args: requests: List of BatchRequest Returns: Dict with keys: 'decode', 'prefill', 'extend' + """ result = { "decode": [], @@ -204,14 +204,14 @@ def split_by_type( def get_batch_stats(requests: list[BatchRequest]) -> dict: - """ - Compute statistics about a batch. + """Compute statistics about a batch. Args: requests: List of BatchRequest Returns: Dict with batch statistics + """ by_type = split_by_type(requests) @@ -232,8 +232,7 @@ def get_batch_stats(requests: list[BatchRequest]) -> dict: def get_batch_type(batch_spec: str, spec_decode_threshold: int = 8) -> str: - """ - Classify a batch spec into a type string. + """Classify a batch spec into a type string. Args: batch_spec: Batch specification string (e.g., "q2k", "8q1s1k", "2q2k_8q1s1k") @@ -241,6 +240,7 @@ def get_batch_type(batch_spec: str, spec_decode_threshold: int = 8) -> str: Returns: Type string: "prefill", "decode", "spec-decode", "extend", or "mixed (types...)" + """ requests = parse_batch_spec(batch_spec) diff --git a/benchmarks/attention_benchmarks/benchmark.py b/benchmarks/attention_benchmarks/benchmark.py index 9c1ea1add7..943133252b 100644 --- a/benchmarks/attention_benchmarks/benchmark.py +++ b/benchmarks/attention_benchmarks/benchmark.py @@ -2,8 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Universal vLLM Attention Benchmark +"""Universal vLLM Attention Benchmark. Benchmark any attention backend with the extended grammar. Supports standard attention (Flash/Triton/FlashInfer) and MLA backends. @@ -23,6 +22,7 @@ Examples: # Parameter sweep (YAML config - recommended) python benchmark.py --config configs/cutlass_numsplits.yaml + """ import argparse @@ -84,8 +84,7 @@ def run_mla_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult: def run_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult: - """ - Run a single benchmark with proper backend selection. + """Run a single benchmark with proper backend selection. Args: config: BenchmarkConfig with backend, batch_spec, and model params @@ -93,6 +92,7 @@ def run_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult: Returns: BenchmarkResult (may have error field set on failure) + """ try: if is_mla_backend(config.backend): @@ -119,8 +119,7 @@ def run_model_parameter_sweep( sweep: ModelParameterSweep, console: Console, ) -> list[BenchmarkResult]: - """ - Run model parameter sweep for given backends and batch specs. + """Run model parameter sweep for given backends and batch specs. Args: backends: List of backend names @@ -131,6 +130,7 @@ def run_model_parameter_sweep( Returns: List of BenchmarkResult objects + """ all_results = [] @@ -304,8 +304,7 @@ def run_parameter_sweep( sweep: ParameterSweep, console: Console, ) -> list[BenchmarkResult]: - """ - Run parameter sweep for given backends and batch specs. + """Run parameter sweep for given backends and batch specs. Args: backends: List of backend names @@ -316,6 +315,7 @@ def run_parameter_sweep( Returns: List of BenchmarkResult objects + """ all_results = [] @@ -395,8 +395,7 @@ def load_config_from_yaml(config_path: str) -> dict: def generate_batch_specs_from_ranges(ranges: list[dict]) -> list[str]: - """ - Generate batch specs from range specifications. + """Generate batch specs from range specifications. Args: ranges: List of range specifications, each containing: @@ -420,6 +419,7 @@ def generate_batch_specs_from_ranges(ranges: list[dict]) -> list[str]: } ] Returns: ["q1kv1k", "q2kv1k", ..., "q16kv1k"] + """ all_specs = [] diff --git a/benchmarks/attention_benchmarks/common.py b/benchmarks/attention_benchmarks/common.py index 45858f9148..b6fd26c803 100644 --- a/benchmarks/attention_benchmarks/common.py +++ b/benchmarks/attention_benchmarks/common.py @@ -21,8 +21,7 @@ from vllm.triton_utils import triton def batch_spec_sort_key(spec: str) -> tuple[int, int, int]: - """ - Extract sorting key from batch spec: (batch_size, max_q_len, max_kv_len). + """Extract sorting key from batch spec: (batch_size, max_q_len, max_kv_len). This ensures results are sorted by batch size first, then query length, then sequence length, rather than alphabetically. @@ -354,13 +353,13 @@ class ResultsFormatter: backends: list[str], compare_to_fastest: bool = True, ): - """ - Print results as a rich table. + """Print results as a rich table. Args: results: List of BenchmarkResult backends: List of backend names being compared compare_to_fastest: Show percentage comparison to fastest + """ # Group by batch spec, preserving first-occurrence order by_spec = {} @@ -489,14 +488,14 @@ class ResultsFormatter: def setup_mla_dims(model_name: str = "deepseek-v3") -> dict: - """ - Get MLA dimensions for known models. + """Get MLA dimensions for known models. Args: model_name: Model identifier Returns: Dict with MLA dimension configuration + """ configs = { "deepseek-v2": { @@ -542,8 +541,7 @@ def get_attention_scale(head_dim: int) -> float: def is_mla_backend(backend: str) -> bool: - """ - Check if backend is an MLA backend using the AttentionBackendEnum. + """Check if backend is an MLA backend using the AttentionBackendEnum. Args: backend: Backend name matching AttentionBackendEnum exactly @@ -551,6 +549,7 @@ def is_mla_backend(backend: str) -> bool: Returns: True if the backend is an MLA backend, False otherwise + """ from vllm.v1.attention.backends.registry import AttentionBackendEnum diff --git a/benchmarks/attention_benchmarks/mla_runner.py b/benchmarks/attention_benchmarks/mla_runner.py index 8a6ae3f87d..04dcd73719 100644 --- a/benchmarks/attention_benchmarks/mla_runner.py +++ b/benchmarks/attention_benchmarks/mla_runner.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -MLA benchmark runner - shared utilities for MLA benchmarks. +"""MLA benchmark runner - shared utilities for MLA benchmarks. This module provides helpers for running MLA backends without needing full VllmConfig integration. @@ -93,8 +92,7 @@ def _ensure_single_rank_model_parallel() -> None: def _add_mock_methods_to_model_config(model_config: ModelConfig) -> None: - """ - Add mock methods for layer-specific queries to ModelConfig. + """Add mock methods for layer-specific queries to ModelConfig. These methods are needed by metadata builders but aren't normally present on ModelConfig when used in benchmark contexts. @@ -125,14 +123,15 @@ def create_minimal_vllm_config( kv_cache_dtype: str = "auto", sparse_mla_force_mqa: bool = False, ) -> VllmConfig: - """ - Create minimal VllmConfig for MLA benchmarks. + """Create minimal VllmConfig for MLA benchmarks. Args: model_name: Model name (deepseek-v2, deepseek-v3, etc.) - used if mla_dims not provided block_size: KV cache block size max_num_seqs: Maximum number of sequences + max_num_batched_tokens: Maximum number of batched tokens + max_model_len: Maximum model context length mla_dims: Optional custom MLA dimensions dict. If not provided, uses setup_mla_dims(model_name) index_topk: Optional topk value for sparse MLA backends. If provided, @@ -140,11 +139,13 @@ def create_minimal_vllm_config( prefill_backend: Prefill backend name (e.g., "fa3", "fa4", "flashinfer", "trtllm"). Configures the attention config to force the specified prefill backend. + kv_cache_dtype: KV cache dtype, e.g. "auto" or "fp8" sparse_mla_force_mqa: If True, forces all sparse MLA tokens through forward_mqa (even prefill tokens). Returns: VllmConfig for benchmarking + """ # Get MLA dimensions - use provided or load from model name if mla_dims is None: @@ -320,8 +321,7 @@ _BACKEND_PROPERTIES = { def _get_backend_config(backend: str) -> dict: - """ - Get backend configuration from AttentionBackendEnum. + """Get backend configuration from AttentionBackendEnum. Uses the registry to get the backend class and extract configuration from its methods (get_impl_cls, get_builder_cls, is_sparse, etc.). @@ -332,6 +332,7 @@ def _get_backend_config(backend: str) -> dict: Returns: Dict with backend configuration + """ from vllm.v1.attention.backend import MultipleOf from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -381,8 +382,7 @@ def _build_attention_metadata( device: torch.device, builder_instance, ) -> tuple: - """ - Build attention metadata from batch requests. + """Build attention metadata from batch requests. Args: requests: List of BatchRequest objects @@ -392,6 +392,7 @@ def _build_attention_metadata( Returns: Tuple of (metadata, kv_cache_num_blocks) + """ q_lens = [r.q_len for r in requests] kv_lens = [r.kv_len for r in requests] @@ -472,8 +473,7 @@ def _create_input_tensors( device: torch.device, dtype: torch.dtype, ): - """ - Create input tensors for both decode and prefill modes. + """Create input tensors for both decode and prefill modes. MLA requires different tensor formats for decode vs prefill: - Decode: Uses kv_lora_rank (512) dimension @@ -490,6 +490,7 @@ def _create_input_tensors( Tuple of (decode_inputs, prefill_inputs) - decode_inputs: Query tensor(s) for decode mode - prefill_inputs: Dict with 'q', 'k_c_normed', 'k_pe', 'k_scale' for prefill + """ if query_format == "tuple": # Decode mode format: (q_nope, q_pe) where q_nope has kv_lora_rank dim @@ -583,8 +584,7 @@ def _create_backend_impl( index_topk: int | None = None, kv_cache_dtype: str = "auto", ): - """ - Create backend implementation instance. + """Create backend implementation instance. Args: backend_cfg: Backend configuration dict from _get_backend_config() @@ -593,9 +593,11 @@ def _create_backend_impl( device: Target device max_num_tokens: Maximum number of tokens for sparse indexer buffer index_topk: Topk value for sparse MLA backends + kv_cache_dtype: KV cache dtype, e.g. "auto" or "fp8" Returns: Tuple of (impl, layer, builder_instance, indexer) + """ # Get classes from backend config (already resolved by _get_backend_config) impl_class = backend_cfg["impl_class"] @@ -701,14 +703,14 @@ def _create_backend_impl( def _extract_mla_dims_from_config(config) -> dict | None: - """ - Extract MLA dimensions from BenchmarkConfig if all required fields are present. + """Extract MLA dimensions from BenchmarkConfig if all required fields are present. Args: config: BenchmarkConfig instance Returns: Dict with MLA dimensions if all fields are provided, None otherwise + """ # Check if all MLA-specific fields are provided if all( @@ -762,8 +764,7 @@ def _run_single_benchmark( output_scale: float | None = None, fuse_quant_op: bool = False, ) -> BenchmarkResult: - """ - Run a single benchmark iteration. + """Run a single benchmark iteration. Args: config: BenchmarkConfig instance @@ -774,6 +775,7 @@ def _run_single_benchmark( mla_dims: MLA dimension configuration device: Target device indexer: Optional MockIndexer for sparse backends + kv_cache_dtype: KV cache dtype. None keeps the backend default. output_scale: Static per-tensor FP8 scale for prefill output. None keeps the plain bf16 output (no quantization). fuse_quant_op: With output_scale set, True lets the prefill kernel write @@ -782,6 +784,7 @@ def _run_single_benchmark( Returns: BenchmarkResult with timing statistics + """ # Parse batch spec requests = parse_batch_spec(config.batch_spec) @@ -1044,8 +1047,7 @@ def _run_mla_benchmark_batched( output_scale: float | None = None, fuse_quant_op: bool = False, ) -> list[BenchmarkResult]: - """ - Unified batched MLA benchmark runner for all backends. + """Unified batched MLA benchmark runner for all backends. Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla, flashinfer_mla_sparse, flashmla_sparse @@ -1061,11 +1063,16 @@ def _run_mla_benchmark_batched( index_topk: Topk value for sparse MLA backends (default 2048) prefill_backend: Prefill backend name (e.g., "fa3", "fa4"). When set, forces the specified FlashAttention version for prefill. + output_scale: Static per-tensor FP8 scale for prefill output. None + keeps the plain bf16 output (no quantization). + fuse_quant_op: With output_scale set, True lets the prefill kernel write + FP8 directly; False runs bf16 attention then a standalone quant. sparse_mla_force_mqa: If True, forces all sparse MLA tokens through forward_mqa (even prefill tokens). Returns: List of BenchmarkResult objects + """ if not configs_with_params: return [] @@ -1277,8 +1284,7 @@ def run_mla_benchmark( output_scale: float | None = None, fuse_quant_op: bool = False, ) -> BenchmarkResult | list[BenchmarkResult]: - """ - Unified MLA benchmark runner for all backends. + """Unified MLA benchmark runner for all backends. Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla, flashinfer_mla_sparse, flashmla_sparse @@ -1303,6 +1309,7 @@ def run_mla_benchmark( Returns: BenchmarkResult (single mode) or list of BenchmarkResult (batched mode) + """ # Normalize to batched mode: (config, threshold, num_splits) if isinstance(config, list): diff --git a/benchmarks/attention_benchmarks/runner.py b/benchmarks/attention_benchmarks/runner.py index 01fd558043..a1b7d04505 100644 --- a/benchmarks/attention_benchmarks/runner.py +++ b/benchmarks/attention_benchmarks/runner.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Standard attention benchmark runner - shared utilities for non-MLA benchmarks. +"""Standard attention benchmark runner - shared utilities for non-MLA benchmarks. This module provides helpers for running standard attention backends (FlashAttention, Triton, FlashInfer) with real vLLM integration. @@ -57,8 +56,7 @@ from vllm.v1.kv_cache_interface import ( def _get_backend_config(backend: str) -> dict: - """ - Get backend configuration from AttentionBackendEnum. + """Get backend configuration from AttentionBackendEnum. Args: backend: Backend name matching AttentionBackendEnum exactly @@ -66,6 +64,7 @@ def _get_backend_config(backend: str) -> dict: Returns: Dict with backend_class + """ from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -409,6 +408,7 @@ def _run_single_benchmark( Returns: (timing_stats, mem_stats) where timing_stats is a dict with mean/std/min/max in seconds per layer. + """ total_q = q_list[0].shape[0] v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim @@ -458,8 +458,7 @@ def _run_single_benchmark( def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult: - """ - Run standard attention benchmark with real kernels. + """Run standard attention benchmark with real kernels. Supports: FLASH_ATTN, TRITON_ATTN, FLASHINFER @@ -468,6 +467,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult: Returns: BenchmarkResult with timing and memory statistics + """ device = torch.device(config.device) torch.accelerator.set_device_index(device) diff --git a/benchmarks/benchmark_batch_invariance.py b/benchmarks/benchmark_batch_invariance.py index 7473a41e51..dd22523ce2 100755 --- a/benchmarks/benchmark_batch_invariance.py +++ b/benchmarks/benchmark_batch_invariance.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Benchmark to measure the performance overhead of VLLM_BATCH_INVARIANT mode. +"""Benchmark to measure the performance overhead of VLLM_BATCH_INVARIANT mode. This benchmark runs the same workload twice: 1. With VLLM_BATCH_INVARIANT=0 (baseline) @@ -96,8 +95,7 @@ def run_benchmark_with_batch_invariant( batch_invariant: bool, seed: int = 12345, ) -> dict: - """ - Run the benchmark with the specified configuration. + """Run the benchmark with the specified configuration. Returns a dict with timing and throughput metrics. """ diff --git a/benchmarks/benchmark_hash.py b/benchmarks/benchmark_hash.py index 08cdc012d6..c5af4da999 100644 --- a/benchmarks/benchmark_hash.py +++ b/benchmarks/benchmark_hash.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Micro benchmark comparing built-in hash(), SHA-256, and xxHash. +"""Micro benchmark comparing built-in hash(), SHA-256, and xxHash. This focuses on a single test payload shaped like the prefix-cache hash input: (32-byte bytes object, 32-int tuple) diff --git a/benchmarks/benchmark_hidden_state_extraction.py b/benchmarks/benchmark_hidden_state_extraction.py index f0a35a0cf1..1f5da88c2c 100644 --- a/benchmarks/benchmark_hidden_state_extraction.py +++ b/benchmarks/benchmark_hidden_state_extraction.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Benchmark hidden state extraction throughput. +"""Benchmark hidden state extraction throughput. Measures two modes: 1. Baseline: bulk inference with max_tokens=1, no extraction. @@ -137,7 +136,8 @@ async def _client_loop( client_id: int, ): """A single async client: pulls prompts, submits to engine, consumes - hidden states as soon as each request finishes.""" + hidden states as soon as each request finishes. + """ loop = asyncio.get_event_loop() while True: item = await prompt_queue.get() diff --git a/benchmarks/benchmark_long_document_qa_throughput.py b/benchmarks/benchmark_long_document_qa_throughput.py index 6106b90149..1f5b5f5504 100644 --- a/benchmarks/benchmark_long_document_qa_throughput.py +++ b/benchmarks/benchmark_long_document_qa_throughput.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Offline benchmark to test the long document QA throughput. +"""Offline benchmark to test the long document QA throughput. Example usage: # This workload samples 8 different prompts with a default input @@ -49,14 +48,14 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser def test_long_document_qa(llm=None, sampling_params=None, prompts=None): - """ - Test long document QA with the given prompts and sampling parameters. + """Test long document QA with the given prompts and sampling parameters. Print the time spent in processing all the prompts. Args: llm: The language model used for generating responses. sampling_params: Sampling parameter used to generate the response. prompts: A list of prompt strings to be processed by the LLM. + """ start_time = time.time() llm.generate(prompts, sampling_params=sampling_params) @@ -65,8 +64,7 @@ def test_long_document_qa(llm=None, sampling_params=None, prompts=None): def repeat_prompts(prompts, repeat_count, mode: str): - """ - Repeat each prompt in the list for a specified number of times. + """Repeat each prompt in the list for a specified number of times. The order of prompts in the output list depends on the mode. Args: @@ -84,6 +82,7 @@ def repeat_prompts(prompts, repeat_count, mode: str): Raises: ValueError: If an invalid mode is provided. + """ print("Repeat mode: ", mode) if mode == "random": diff --git a/benchmarks/benchmark_prefix_block_hash.py b/benchmarks/benchmark_prefix_block_hash.py index 8bcd8af0d3..522bf9d1ab 100644 --- a/benchmarks/benchmark_prefix_block_hash.py +++ b/benchmarks/benchmark_prefix_block_hash.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Simple benchmark to compare prefix-cache block hashing algorithms. +"""Simple benchmark to compare prefix-cache block hashing algorithms. Example: python benchmark_prefix_block_hash.py --num-blocks 20000 --block-size 32 + """ from __future__ import annotations diff --git a/benchmarks/benchmark_prefix_caching.py b/benchmarks/benchmark_prefix_caching.py index 0145f6b7c8..5abdd65b2a 100644 --- a/benchmarks/benchmark_prefix_caching.py +++ b/benchmarks/benchmark_prefix_caching.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Benchmark the efficiency of prefix caching. +"""Benchmark the efficiency of prefix caching. This script allows you to benchmark the performance of a model with and without prefix caching using either fixed prompts diff --git a/benchmarks/benchmark_serving_structured_output.py b/benchmarks/benchmark_serving_structured_output.py index 664fa58dd4..dc8e30268e 100644 --- a/benchmarks/benchmark_serving_structured_output.py +++ b/benchmarks/benchmark_serving_structured_output.py @@ -102,6 +102,7 @@ class SampleRequest: images). prompt_len: The length of the prompt in tokens. expected_output_len: The expected length of the output in tokens. + """ prompt: str @@ -305,8 +306,7 @@ async def get_request( request_rate: float, burstiness: float = 1.0, ) -> AsyncGenerator[tuple[int, SampleRequest], None]: - """ - Asynchronously generates requests at a specified rate + """Asynchronously generates requests at a specified rate with OPTIONAL burstiness. Args: @@ -322,6 +322,7 @@ async def get_request( A lower burstiness value (0 < burstiness < 1) results in more bursty requests, while a higher burstiness value (burstiness > 1) results in a more uniform arrival of requests. + """ input_requests = iter(input_requests) diff --git a/benchmarks/benchmark_topk_topp.py b/benchmarks/benchmark_topk_topp.py index f727f16ea2..14111f733b 100644 --- a/benchmarks/benchmark_topk_topp.py +++ b/benchmarks/benchmark_topk_topp.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Benchmark comparing Triton vs PyTorch sort-based top-k/top-p implementations. +"""Benchmark comparing Triton vs PyTorch sort-based top-k/top-p implementations. Compares: - apply_top_k_top_p_triton (Triton binary search) @@ -47,8 +46,7 @@ def calculate_ops_pct( vocab_size: int, batch_size: int, ) -> float: - """ - Calculate the percentage of active top-k and top-p operations. + """Calculate the percentage of active top-k and top-p operations. Returns percentage where 100% = batch_size ops. E.g., if all rows have both top-k and top-p active, returns 200%. @@ -117,8 +115,7 @@ def benchmark_function( warmup_iters: int = 5, benchmark_iters: int = 20, ) -> tuple[float, int]: - """ - Benchmark a function and return (avg_time_ms, peak_memory_bytes). + """Benchmark a function and return (avg_time_ms, peak_memory_bytes). Returns average time in milliseconds and peak memory usage. """ diff --git a/benchmarks/fused_kernels/merge_attn_states_benchmarks.py b/benchmarks/fused_kernels/merge_attn_states_benchmarks.py index 26b04299b3..a4b4533fa3 100644 --- a/benchmarks/fused_kernels/merge_attn_states_benchmarks.py +++ b/benchmarks/fused_kernels/merge_attn_states_benchmarks.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Benchmark: Fused FP8 output quantization in merge_attn_states +"""Benchmark: Fused FP8 output quantization in merge_attn_states. Compares fused vs unfused approaches for producing FP8-quantized merged attention output: @@ -85,7 +84,8 @@ def make_inputs( def build_configs(head_configs, num_tokens_list, input_dtypes, tp_sizes): """Build (num_tokens, num_heads, head_size, dtype_str) config tuples, - applying TP division to num_heads and skipping invalid combos.""" + applying TP division to num_heads and skipping invalid combos. + """ configs = [] for (_, nh, hs), nt, dtype, tp in itertools.product( head_configs, num_tokens_list, input_dtypes, tp_sizes diff --git a/benchmarks/kernels/bench_concat_mla_q.py b/benchmarks/kernels/bench_concat_mla_q.py index 8d940484d6..c528ab86b2 100644 --- a/benchmarks/kernels/bench_concat_mla_q.py +++ b/benchmarks/kernels/bench_concat_mla_q.py @@ -24,9 +24,9 @@ def make_inputs(num_tokens, dtype): """Create inputs matching the real code path. Args: - contiguous_nope: If False, simulate the transposed BMM output - (non-contiguous nope with stride pattern from - [N,B,L].transpose(0,1)). + num_tokens: Number of tokens to generate inputs for. + dtype: dtype of the generated tensors. + """ # Simulate: bmm output [N, B, L].transpose(0, 1) -> [B, N, L] raw = torch.randn(NUM_HEADS, num_tokens, NOPE_DIM, dtype=dtype, device="cuda") diff --git a/benchmarks/kernels/benchmark_2d_silu_mul_fp8_quant.py b/benchmarks/kernels/benchmark_2d_silu_mul_fp8_quant.py index 0dd5c6d848..d0880b3f03 100644 --- a/benchmarks/kernels/benchmark_2d_silu_mul_fp8_quant.py +++ b/benchmarks/kernels/benchmark_2d_silu_mul_fp8_quant.py @@ -106,8 +106,7 @@ class BenchmarkTensors: def reference_quant(x: torch.Tensor, quant_out: torch.Tensor, use_ue8m0: bool): - """ - Reference triton quant kernel from, + """Reference triton quant kernel from, vllm.model_executor.layers.quantization.utils.fp8_utils """ assert quant_out.size() == x.size() diff --git a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py index 03d7fb386f..bff4f73e2c 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Benchmark the performance of the cutlass_moe_fp8 kernel vs the triton_moe +"""Benchmark the performance of the cutlass_moe_fp8 kernel vs the triton_moe kernel. Both kernels take in fp8 quantized weights and 16-bit activations, but use different quantization strategies and backends. """ @@ -190,7 +189,7 @@ def bench_run( torch.accelerator.synchronize() def bench_cuda_graph(graph, num_warmup=5, num_iters=100): - """Benchmark CUDA graph using events like benchmark_moe.py""" + """Benchmark CUDA graph using events like benchmark_moe.py.""" # Warmup for _ in range(num_warmup): graph.replay() diff --git a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py index 7379bf8588..a14242732a 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Benchmark the performance of the cutlass_moe_fp4 kernel vs the triton_moe +"""Benchmark the performance of the cutlass_moe_fp4 kernel vs the triton_moe kernel. The cutlass_moe_fp4 kernel takes in fp4 quantized weights and 16-bit activations. The triton_moe kernel takes in fp8 weights(tensor scaled to fp8) and 16-bit activations. diff --git a/benchmarks/kernels/benchmark_device_communicators.py b/benchmarks/kernels/benchmark_device_communicators.py index 6c3a726a97..9f9a150471 100644 --- a/benchmarks/kernels/benchmark_device_communicators.py +++ b/benchmarks/kernels/benchmark_device_communicators.py @@ -2,8 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Benchmark script for device communicators: +"""Benchmark script for device communicators: CustomAllreduce (oneshot, twoshot), PyNcclCommunicator, and SymmMemCommunicator (multimem, two-shot). @@ -17,6 +16,7 @@ Usage: Example: torchrun --nproc_per_node=2 benchmark_device_communicators.py --sequence-lengths 512 1024 2048 --num-warmup 10 --num-trials 100 + """ import json @@ -241,7 +241,6 @@ class CommunicatorBenchmark: self, sequence_length: int, num_warmup: int, num_trials: int ) -> dict[str, float]: """Benchmark allreduce operations for all available communicators.""" - results = {} # Define communicators with their benchmark functions diff --git a/benchmarks/kernels/benchmark_fused_collective.py b/benchmarks/kernels/benchmark_fused_collective.py index c999c16021..aaa39fd9fe 100644 --- a/benchmarks/kernels/benchmark_fused_collective.py +++ b/benchmarks/kernels/benchmark_fused_collective.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Benchmark for FlashInfer fused collective operations vs standard operations. +"""Benchmark for FlashInfer fused collective operations vs standard operations. This benchmark compares: 1. FlashInfer's allreduce_fusion with trtllm backend @@ -446,9 +445,15 @@ def run_benchmarks( """Run all benchmarks for given configuration. Args: + num_tokens: Number of tokens in the benchmark input. + hidden_dim: Hidden dimension of the benchmark input. + dtype: dtype of the benchmark input. + use_residual: Whether to include a residual add in the fused op. + no_oneshot: Skip the one-shot allreduce variants. allreduce_params: Shared parameters for FlashInfer fused allreduce. workspaces: Dict mapping backend name ("trtllm", "mnnvl") to workspace. quant_modes: Set of quantization modes: "none", "fp8", "fp4". + """ ( input_tensor, diff --git a/benchmarks/kernels/benchmark_kimi_k3_attn_res.py b/benchmarks/kernels/benchmark_kimi_k3_attn_res.py index c8cf6f9acc..dba82a8722 100644 --- a/benchmarks/kernels/benchmark_kimi_k3_attn_res.py +++ b/benchmarks/kernels/benchmark_kimi_k3_attn_res.py @@ -10,6 +10,7 @@ Example: .venv/bin/python benchmarks/kernels/benchmark_kimi_k3_attn_res.py \ --output /tmp/attn_res.json + """ from __future__ import annotations diff --git a/benchmarks/kernels/benchmark_kimi_k3_kda_decode.py b/benchmarks/kernels/benchmark_kimi_k3_kda_decode.py index fcce640152..6b6d944455 100644 --- a/benchmarks/kernels/benchmark_kimi_k3_kda_decode.py +++ b/benchmarks/kernels/benchmark_kimi_k3_kda_decode.py @@ -16,6 +16,7 @@ state traffic itself. Example: python benchmarks/kernels/benchmark_kimi_k3_kda_decode.py \ --tokens 1 8 32 64 128 --heads 12 + """ import argparse diff --git a/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py b/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py index 2669dadb34..855769263d 100644 --- a/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py +++ b/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py @@ -27,6 +27,7 @@ Examples: For multi-node runs, launch one ``torchrun`` agent per node and use a shared rendezvous endpoint. + """ from __future__ import annotations diff --git a/benchmarks/kernels/benchmark_lora.py b/benchmarks/kernels/benchmark_lora.py index ab930c59d2..606741d047 100644 --- a/benchmarks/kernels/benchmark_lora.py +++ b/benchmarks/kernels/benchmark_lora.py @@ -102,9 +102,7 @@ def make_rand_tensors( num_slices: int, device: str = "cuda", ) -> tuple[torch.Tensor, list[torch.Tensor], torch.Tensor]: - """ - Make LoRA input/output matrices. - """ + """Make LoRA input/output matrices.""" A = torch.rand(a_shape, dtype=a_dtype).to(device) # LoRA weights column major @@ -117,8 +115,7 @@ def make_rand_tensors( def make_prompt_lora_mapping( num_prompts: int, num_active_loras: int, sort_by_lora_id: bool, device: str ) -> torch.Tensor: - """ - All prompts are mapped to a LoRA ID in range [0, num_active_loras). + """All prompts are mapped to a LoRA ID in range [0, num_active_loras). where 0 refers to first lora, 1 refers to second lora and so on. """ assert num_active_loras > 0 @@ -147,9 +144,7 @@ def make_token_lora_mapping( seq_len_tensor: torch.Tensor, device: str, ): - """ - Make token_lora_mapping from prompt_lora_mapping and seq_lens_tensor - """ + """Make token_lora_mapping from prompt_lora_mapping and seq_lens_tensor.""" assert prompt_lora_mapping.shape[0] == num_prompts # token to lora index mapping @@ -174,8 +169,7 @@ def ref_group_gemm( scaling: float, add_inputs: bool | None, ): - """ - Torch group gemm reference implementation to test correctness of + """Torch group gemm reference implementation to test correctness of benchmarking operations. """ batches = seq_lens_cpu.size(0) @@ -198,9 +192,7 @@ def ref_group_gemm( class OpType(Enum): - """ - LoRA Ops to benchmark and its properties. - """ + """LoRA Ops to benchmark and its properties.""" LORA_SHRINK = auto() LORA_EXPAND = auto() @@ -291,9 +283,7 @@ class OpType(Enum): def matmul_dtypes( self, op_dtype: torch.dtype ) -> tuple[torch.dtype, torch.dtype, torch.dtype]: - """ - return a type, b type and c type for A x B = C - """ + """Return a type, b type and c type for A x B = C.""" if self.is_shrink_fn(): return op_dtype, op_dtype, torch.float32 elif self.is_expand_fn(): @@ -338,8 +328,7 @@ class OpType(Enum): top_k_num: int | None = None, num_experts: int | None = None, ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: - """ - Given num_slices, return the shapes of the A, B, and C matrices + """Given num_slices, return the shapes of the A, B, and C matrices in A x B = C, for the op_type """ m, k, n = self.mkn(batch_size, seq_length, hidden_size, lora_rank) @@ -419,9 +408,7 @@ class OpType(Enum): @dataclass class BenchmarkContext: - """ - LoRA benchmark context - """ + """LoRA benchmark context.""" batch_size: int hidden_size: int @@ -467,9 +454,7 @@ class BenchmarkContext: @dataclass class BenchmarkTensors: - """ - Input/Output tensors used for benchmarks - """ + """Input/Output tensors used for benchmarks.""" # matmul tensors input: torch.Tensor @@ -555,9 +540,7 @@ class BenchmarkTensors: ) def sanity_check(self, ctx: BenchmarkContext, op_type: OpType) -> None: - """ - Fails asserts when non-conformality is detected. - """ + """Fails asserts when non-conformality is detected.""" num_tokens = ( self.input.shape[1] if op_type.is_fused_moe_lora_expand_fn() @@ -577,9 +560,7 @@ class BenchmarkTensors: ) def to_device(self, device: str): - """ - Transfer tensors to device if the tensors aren't already on the device - """ + """Transfer tensors to device if the tensors aren't already on the device.""" def to_device(tensor: torch.Tensor): if tensor.device != device: @@ -604,9 +585,7 @@ class BenchmarkTensors: ) def metadata(self, ctx: BenchmarkContext, op_type: OpType) -> tuple[int, int, int]: - """ - Return num_seqs, num_tokens and max_seq_len - """ + """Return num_seqs, num_tokens and max_seq_len.""" num_seqs = self.seq_lens.shape[0] num_tokens = self.get_num_tokens( self.lora_kernel_meta.token_lora_mapping.shape[0], ctx.top_k_num, op_type @@ -630,8 +609,7 @@ class BenchmarkTensors: expert_map: torch.Tensor | None = None, pad_sorted_ids: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Aligns tokens and experts into block-sized chunks for LoRA-based + """Aligns tokens and experts into block-sized chunks for LoRA-based mixture-of-experts (MoE) execution. """ max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1) @@ -959,8 +937,7 @@ class BenchmarkTensors: def test_correctness( self, op_type: OpType, expand_fn_add_inputs: bool | None ) -> bool: - """ - Test correctness of op_type implementation against a grouped gemm + """Test correctness of op_type implementation against a grouped gemm reference implementation. """ seq_lens_cpu = self.seq_lens.to(device="cpu") @@ -1070,15 +1047,13 @@ def bench_torch_mm( op_type: OpType, cuda_graph_nops: int | None = None, ) -> TMeasurement: - """ - Benchmark basic torch.mm as a roofline. + """Benchmark basic torch.mm as a roofline. When all the input tokens have the same LoRA ID, the LoRA kernels are just a matmul. This torch.mm benchmark serves as a roofline for that case. input op_type is used in determining the m, k, n dimensions for the matmul. """ - batch_size, hidden_size, lora_rank, seq_length, dtype = ( ctx.batch_size, ctx.hidden_size, diff --git a/benchmarks/kernels/benchmark_mla_k_concat.py b/benchmarks/kernels/benchmark_mla_k_concat.py index 7debf36348..3cae613e82 100644 --- a/benchmarks/kernels/benchmark_mla_k_concat.py +++ b/benchmarks/kernels/benchmark_mla_k_concat.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Benchmark script comparing torch.cat vs direct copy for k_nope/k_pe concatenation +"""Benchmark script comparing torch.cat vs direct copy for k_nope/k_pe concatenation in MLA (Multi-head Latent Attention) prefill. This validates that the optimization from commit 8d4142bd is beneficial across diff --git a/benchmarks/kernels/benchmark_moe_defaults.py b/benchmarks/kernels/benchmark_moe_defaults.py index f6ad59366d..f21fbaa703 100644 --- a/benchmarks/kernels/benchmark_moe_defaults.py +++ b/benchmarks/kernels/benchmark_moe_defaults.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Benchmark comparing old vs new default fused MoE configs. +"""Benchmark comparing old vs new default fused MoE configs. Runs the triton fused_moe kernel with three configurations for each scenario: 1. Tuned config (from JSON file, if available) — the target to match @@ -33,7 +32,8 @@ FP8_DTYPE = current_platform.fp8_dtype() def old_default_config(M, E, N, K, topk, dtype=None, block_shape=None): """The original defaults before https://github.com/vllm-project/vllm/pull/34846, - for comparison.""" + for comparison. + """ if dtype == "fp8_w8a8" and block_shape is not None: return { "BLOCK_SIZE_M": 64, diff --git a/benchmarks/kernels/benchmark_per_token_group_quant.py b/benchmarks/kernels/benchmark_per_token_group_quant.py index f2195a6d78..fda3d2a652 100644 --- a/benchmarks/kernels/benchmark_per_token_group_quant.py +++ b/benchmarks/kernels/benchmark_per_token_group_quant.py @@ -15,7 +15,7 @@ from vllm.platforms import current_platform @contextmanager def _triton_mode(): - """Temporarily force the Triton fallback path""" + """Temporarily force the Triton fallback path.""" with patch("vllm.platforms.current_platform.is_cuda", return_value=False): yield diff --git a/benchmarks/kernels/benchmark_per_token_quant_fp8.py b/benchmarks/kernels/benchmark_per_token_quant_fp8.py index 6ce97e3036..0069b06c64 100644 --- a/benchmarks/kernels/benchmark_per_token_quant_fp8.py +++ b/benchmarks/kernels/benchmark_per_token_quant_fp8.py @@ -16,7 +16,7 @@ from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE def with_triton_mode(fn): - """Temporarily force the Triton fallback path""" + """Temporarily force the Triton fallback path.""" def wrapped(*args, **kwargs): with patch("vllm.platforms.current_platform.is_cuda", return_value=False): @@ -124,8 +124,7 @@ def compute_geomean_speedups( speedup_cols: list[str], groupby_cols: list[str] | None = None, ) -> pd.DataFrame: - """ - Compute geometric mean speedups over a baseline column. + """Compute geometric mean speedups over a baseline column. Args: df: Input dataframe @@ -135,6 +134,7 @@ def compute_geomean_speedups( Returns: pd.DataFrame with geometric mean speedups + """ from scipy.stats import gmean diff --git a/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py b/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py index 4ecd917f31..2d5ec80c6a 100644 --- a/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py +++ b/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Benchmark the RDNAHybridW4A16LinearKernel across decode and prefill shapes. +"""Benchmark the RDNAHybridW4A16LinearKernel across decode and prefill shapes. Usage: python benchmark_int4_gemm.py diff --git a/benchmarks/kernels/benchmark_reshape_and_cache.py b/benchmarks/kernels/benchmark_reshape_and_cache.py index 97af4ac976..f49bf61036 100644 --- a/benchmarks/kernels/benchmark_reshape_and_cache.py +++ b/benchmarks/kernels/benchmark_reshape_and_cache.py @@ -32,7 +32,6 @@ def run_benchmark( device: str = "cuda", ) -> float: """Return latency (seconds) for given num_tokens.""" - if kv_cache_dtype == "fp8" and head_size % 16: raise ValueError("fp8 kv-cache requires head_size to be a multiple of 16.") diff --git a/benchmarks/kernels/benchmark_reshape_and_cache_flash.py b/benchmarks/kernels/benchmark_reshape_and_cache_flash.py index 55c2037251..84d9cc707c 100644 --- a/benchmarks/kernels/benchmark_reshape_and_cache_flash.py +++ b/benchmarks/kernels/benchmark_reshape_and_cache_flash.py @@ -37,7 +37,6 @@ def run_benchmark( device: str = "cuda", ) -> float: """Return latency (seconds) for given num_tokens.""" - if kv_cache_dtype == "fp8" and head_size % 16: raise ValueError("fp8 kv-cache requires head_size to be a multiple of 16.") diff --git a/benchmarks/kernels/benchmark_selective_state_update.py b/benchmarks/kernels/benchmark_selective_state_update.py index ddf1efd018..476dce7dff 100644 --- a/benchmarks/kernels/benchmark_selective_state_update.py +++ b/benchmarks/kernels/benchmark_selective_state_update.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Benchmark and tuning script for the Mamba selective_state_update kernel. +"""Benchmark and tuning script for the Mamba selective_state_update kernel. Mirrors the fused MoE tuning workflow: sweeps (BLOCK_SIZE_M, num_warps) across an effective_batch grid for a given (headdim, dstate, ngroups, cache_dtype) and @@ -124,8 +123,7 @@ def benchmark_config( num_warmup: int = 20, graph_batch_size: int = 10, ) -> float | None: - """ - Time one (BLOCK_SIZE_M, num_warps) config for selective_state_update. + """Time one (BLOCK_SIZE_M, num_warps) config for selective_state_update. Returns elapsed time in microseconds, or None on error. Uses accelerator graph capture-and-replay to isolate kernel time from @@ -356,8 +354,7 @@ def validate_configs( rtol: float = 1e-2, state_dtype: torch.dtype | None = None, ) -> dict[int, bool]: - """ - For every (effective_batch, batch, nheads) in *active* that has a tuned + """For every (effective_batch, batch, nheads) in *active* that has a tuned config, run the kernel with that config and compare against the reference. Returns {effective_batch: passed}. """ diff --git a/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py b/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py index 9fcf278f2e..22ef378b07 100644 --- a/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py +++ b/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Comprehensive 3-way SiLU Benchmark Suite +"""Comprehensive 3-way SiLU Benchmark Suite. This benchmark compares three SiLU implementations: 1. SiLU V2 (CUDA) - Optimized CUDA kernel implementation @@ -110,7 +109,7 @@ def silu_mul_fp8_quant_deep_gemm_triton( eps: float = 1e-10, expert_offsets: torch.Tensor = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """Quantize silu(y[..., :H]) * y[..., H:] to FP8 with group per-token scales + """Quantize silu(y[..., :H]) * y[..., H:] to FP8 with group per-token scales. y has shape (E, T, 2*H). The first half of the last dimension is silu-activated, multiplied by the second half, then quantized into FP8. @@ -206,7 +205,7 @@ def benchmark( iterations_per_run: int = 20, ): def generate_data(seed_offset=0): - """Generate input data with given seed offset""" + """Generate input data with given seed offset.""" set_random_seed(42 + seed_offset) y = torch.rand((E, T, 2 * H), dtype=torch.bfloat16, device="cuda").contiguous() diff --git a/benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py b/benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py index 6bb504b7f0..29ff156b60 100644 --- a/benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py +++ b/benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py @@ -10,6 +10,7 @@ with a static scale file. Example: python benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py \ --seq-lens 2304 4096 8192 16384 --head-dim 72 + """ from types import SimpleNamespace diff --git a/benchmarks/kernels/benchmark_vocab_parallel_embedding.py b/benchmarks/kernels/benchmark_vocab_parallel_embedding.py index 4925e6c024..2922c449ea 100644 --- a/benchmarks/kernels/benchmark_vocab_parallel_embedding.py +++ b/benchmarks/kernels/benchmark_vocab_parallel_embedding.py @@ -97,7 +97,8 @@ def graph_replay(fn): def calculate_diff(vocab_size, hidden_size, tp_size, num_tokens): """The fused kernel must be bit-exact with the path it replaces, on every - rank of the shard.""" + rank of the shard. + """ dtype = torch.bfloat16 for tp_rank in range(tp_size): input_ids, weight, shard_indices = make_inputs( diff --git a/benchmarks/kernels/benchmark_w8a8_block_fp8.py b/benchmarks/kernels/benchmark_w8a8_block_fp8.py index 590d4cfdc6..3ad7ec8887 100644 --- a/benchmarks/kernels/benchmark_w8a8_block_fp8.py +++ b/benchmarks/kernels/benchmark_w8a8_block_fp8.py @@ -57,10 +57,12 @@ def w8a8_block_matmul( Bs: The per-block quantization scale for `B`. block_size: The block size for per-block quantization. It should be 2-dim, e.g., [128, 128]. + config: Triton kernel launch config, e.g. tile sizes and num_warps. output_dtype: The dtype of the returned tensor. Returns: torch.Tensor: The result of matmul. + """ assert len(block_size) == 2 block_n, block_k = block_size[0], block_size[1] diff --git a/benchmarks/kernels/ir/bench_ir_ops.py b/benchmarks/kernels/ir/bench_ir_ops.py index b23c4e8ae3..80fdd87acb 100644 --- a/benchmarks/kernels/ir/bench_ir_ops.py +++ b/benchmarks/kernels/ir/bench_ir_ops.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Generic benchmark harness for vLLM IR ops. +"""Generic benchmark harness for vLLM IR ops. Usage: python benchmarks/kernels/ir/bench_ir_ops.py diff --git a/benchmarks/kernels/ir/shapes.py b/benchmarks/kernels/ir/shapes.py index 6cc44cf6ce..bac2920103 100644 --- a/benchmarks/kernels/ir/shapes.py +++ b/benchmarks/kernels/ir/shapes.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Shape configurations for IR op benchmarks. -""" +"""Shape configurations for IR op benchmarks.""" import torch diff --git a/benchmarks/kernels/utils.py b/benchmarks/kernels/utils.py index a9af811bbe..64d23691bd 100644 --- a/benchmarks/kernels/utils.py +++ b/benchmarks/kernels/utils.py @@ -17,8 +17,7 @@ class CudaGraphBenchParams: @dataclasses.dataclass class ArgPool: - """ - When some argument of the benchmarking function is annotated with this type, + """When some argument of the benchmarking function is annotated with this type, the benchmarking class (BenchMM) will collapse the argument to a pick a single value from the given list of values, during function invocation. For every invocation during a benchmarking run, it will choose a diff --git a/benchmarks/multi_turn/bench_dataset.py b/benchmarks/multi_turn/bench_dataset.py index 8cb8a2f386..9b159a9ee9 100644 --- a/benchmarks/multi_turn/bench_dataset.py +++ b/benchmarks/multi_turn/bench_dataset.py @@ -148,8 +148,7 @@ class LognormalDistribution(Distribution): def _generate_lognormal_by_median( target_average: int, median_ratio: float ) -> tuple[float, float]: - """ - Compute (mu, sigma) for a lognormal distribution given: + """Compute (mu, sigma) for a lognormal distribution given: - a target average (mean of the distribution) - a ratio of median / mean (controls skewness), assume mean > median diff --git a/benchmarks/multi_turn/convert_sharegpt_to_openai.py b/benchmarks/multi_turn/convert_sharegpt_to_openai.py index fccab4d0ce..b415744873 100644 --- a/benchmarks/multi_turn/convert_sharegpt_to_openai.py +++ b/benchmarks/multi_turn/convert_sharegpt_to_openai.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Download dataset from: +"""Download dataset from: https://huggingface.co/datasets/philschmid/sharegpt-raw/blob/main/sharegpt_20230401_clean_lang_split.json Convert to OpenAI API: diff --git a/benchmarks/replayssm/e2e_decode_speedup.py b/benchmarks/replayssm/e2e_decode_speedup.py index b7084a2131..ff7c30669a 100644 --- a/benchmarks/replayssm/e2e_decode_speedup.py +++ b/benchmarks/replayssm/e2e_decode_speedup.py @@ -15,6 +15,7 @@ Examples: python e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 python e2e_decode_speedup.py --dtype auto --buffer-len 16 \ --model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 # B300 NVFP4 + """ import argparse diff --git a/csrc/cpu/generate_cpu_attn_dispatch.py b/csrc/cpu/generate_cpu_attn_dispatch.py index a82a27b844..12a14ef587 100644 --- a/csrc/cpu/generate_cpu_attn_dispatch.py +++ b/csrc/cpu/generate_cpu_attn_dispatch.py @@ -1,9 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Generate CPU attention dispatch switch cases and kernel instantiations. -""" +"""Generate CPU attention dispatch switch cases and kernel instantiations.""" import os diff --git a/docs/mkdocs/gen_files/generate_argparse.py b/docs/mkdocs/gen_files/generate_argparse.py index 1c7025366f..ecdce7761f 100644 --- a/docs/mkdocs/gen_files/generate_argparse.py +++ b/docs/mkdocs/gen_files/generate_argparse.py @@ -98,8 +98,7 @@ mock_if_no_torch("torch", PydanticMagicMock(name="torch", nn=mock_nn)) # Mock torch.library.infer_schema for vllm.ir.ops.IrOpInplaceOverload.__init__ # We need to return the corresponding number of inputs, as IR infra will assert it def get_outputs(native_fn: Callable) -> str: - """ - Extract output schema from function's return type annotation, + """Extract output schema from function's return type annotation, e.g. 'Tensor' or 'Tensor, Tensor'. """ import typing @@ -249,11 +248,12 @@ def create_parser(add_cli_args, **kwargs) -> FlexibleArgumentParser: """Create a parser for the given class with markdown formatting. Args: - cls: The class to create a parser for + add_cli_args: The `add_cli_args` classmethod to build the parser from. **kwargs: Additional keyword arguments to pass to `cls.add_cli_args`. Returns: FlexibleArgumentParser: A parser with markdown formatting for the class. + """ try: parser = FlexibleArgumentParser(add_json_tip=False) diff --git a/docs/mkdocs/gen_files/generate_attention_backends.py b/docs/mkdocs/gen_files/generate_attention_backends.py index 942165eaa9..725d7072ca 100644 --- a/docs/mkdocs/gen_files/generate_attention_backends.py +++ b/docs/mkdocs/gen_files/generate_attention_backends.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Generates documentation table for attention backends showing feature support. +"""Generates documentation table for attention backends showing feature support. This script parses all registered attention backends using AST (no imports needed) and generates a markdown table showing what features each backend supports, @@ -304,6 +303,7 @@ def parse_mla_prefill_registry() -> dict[str, str]: Returns: A dict mapping backend names to their class paths. + """ if not MLA_PREFILL_REGISTRY_FILE.exists(): return {} @@ -325,6 +325,7 @@ def parse_mla_prefill_priorities() -> dict[str, list[str]]: Returns: A dict with keys like 'blackwell' and 'default' containing lists of backend enum names in priority order. + """ if not MLA_PREFILL_SELECTOR_FILE.exists(): return {} @@ -549,6 +550,7 @@ def parse_mla_prefill_backend_file(class_path: str) -> dict[str, Any] | None: Returns: A dict with backend properties, or None if parsing fails. + """ file_path = get_file_from_class_path(class_path) if file_path is None: diff --git a/docs/mkdocs/gen_files/generate_examples.py b/docs/mkdocs/gen_files/generate_examples.py index c9c10018ec..c051653cf9 100644 --- a/docs/mkdocs/gen_files/generate_examples.py +++ b/docs/mkdocs/gen_files/generate_examples.py @@ -49,8 +49,7 @@ def title(text: str) -> str: @dataclass class Example: - """ - Example class for generating documentation content from a given path. + """Example class for generating documentation content from a given path. Attributes: path (Path): The path to the main directory or file. @@ -64,6 +63,7 @@ class Example: Methods: generate() -> str: Generates the documentation content. + """ path: Path @@ -75,7 +75,8 @@ class Example: If path is a file, it returns the path itself. If path is a directory, it searches for Markdown files (*.md) in the directory and returns the first one - found. If no Markdown files are found, it returns None.""" + found. If no Markdown files are found, it returns None. + """ # Single file example if self.path.is_file(): return self.path @@ -90,7 +91,8 @@ class Example: """Determine other files in the directory excluding the main file. If path is a file, it returns an empty list. Otherwise, it returns every file - in the directory except the main file in a list.""" + in the directory except the main file in a list. + """ # Single file example if self.path.is_file(): return [] @@ -116,8 +118,7 @@ class Example: raise ValueError(f"Title not found in {self.main_file}") def fix_relative_links(self, content: str) -> str: - """ - Fix relative links in markdown content by converting them to gh-file + """Fix relative links in markdown content by converting them to gh-file format. Args: @@ -125,6 +126,7 @@ class Example: Returns: str: Content with relative links converted to gh-file format + """ # Regex to match markdown links [text](relative_path) # This matches links that don't start with http, https, ftp, or # diff --git a/docs/mkdocs/gen_files/generated_content.py b/docs/mkdocs/gen_files/generated_content.py index fa3f31b294..d669194424 100644 --- a/docs/mkdocs/gen_files/generated_content.py +++ b/docs/mkdocs/gen_files/generated_content.py @@ -35,6 +35,7 @@ def fill_markers(doc_path: str, blocks: dict[str, str]) -> None: FileNotFoundError: If the source page does not exist. ValueError: If an expected marker is missing, or any `gen:` marker is left unsubstituted after filling. + """ source = DOCS_DIR / doc_path if not source.exists(): diff --git a/docs/mkdocs/hooks/autoref_code.py b/docs/mkdocs/hooks/autoref_code.py index 647f74f202..58cec14d45 100644 --- a/docs/mkdocs/hooks/autoref_code.py +++ b/docs/mkdocs/hooks/autoref_code.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -MkDocs hook to automatically convert inline code references to API doc links. +"""MkDocs hook to automatically convert inline code references to API doc links. For example, `WeightTransferConfig` becomes [`WeightTransferConfig`][vllm.config.WeightTransferConfig] diff --git a/docs/mkdocs/hooks/url_schemes.py b/docs/mkdocs/hooks/url_schemes.py index d208b9d677..cee418936f 100644 --- a/docs/mkdocs/hooks/url_schemes.py +++ b/docs/mkdocs/hooks/url_schemes.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -MkDocs hook + markdown extension to enable the following links to render correctly, +"""MkDocs hook + markdown extension to enable the following links to render correctly, including inside content included via pymdownx.snippets: - Relative file links outside of the `docs/` directory, e.g.: @@ -71,9 +70,7 @@ class UrlSchemesPreprocessor(Preprocessor): return lines def replace_relative_link(match: re.Match) -> str: - """ - Replace relative file links with URLs if they point outside the docs dir. - """ + """Replace relative file links pointing outside the docs dir with URLs.""" title = match.group("title") path = match.group("path") path = ((DOC_DIR / page.file.src_uri).parent / path).resolve() @@ -91,9 +88,7 @@ class UrlSchemesPreprocessor(Preprocessor): return f"[{gh_icon} {title}]({url})" def replace_github_link(match: re.Match) -> str: - """ - Replace GitHub issue, PR, and project links with enhanced Markdown links. - """ + """Replace GitHub issue, PR, and project links with rich Markdown links.""" repo = match.group("repo") type = match.group("type") number = match.group("number") diff --git a/examples/applications/api_server/server.py b/examples/applications/api_server/server.py index e56a7931f6..36474e2101 100644 --- a/examples/applications/api_server/server.py +++ b/examples/applications/api_server/server.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -NOTE: This API server is used only for demonstrating usage of AsyncEngine +"""NOTE: This API server is used only for demonstrating usage of AsyncEngine and simple performance benchmarks. It is not intended for production use. For production use, we recommend using our OpenAI compatible server. We are also not going to accept PRs modifying this file. diff --git a/examples/applications/chatbot/streamlit_openai_chatbot_webserver.py b/examples/applications/chatbot/streamlit_openai_chatbot_webserver.py index d60dbf4d75..eb18b41a5a 100644 --- a/examples/applications/chatbot/streamlit_openai_chatbot_webserver.py +++ b/examples/applications/chatbot/streamlit_openai_chatbot_webserver.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -vLLM Chat Assistant - A Streamlit Web Interface +"""vLLM Chat Assistant - A Streamlit Web Interface. A streamlined chat interface that quickly integrates with vLLM API server. @@ -78,6 +77,7 @@ def create_new_chat_session(): - current_session: Sets to new session ID - active_session: Sets to new session ID - messages: Resets to empty list + """ session_id = datetime.now().strftime("%Y-%m-%d %H:%M:%S") st.session_state.sessions[session_id] = [] @@ -101,6 +101,7 @@ def switch_to_chat_session(session_id): - current_session: Updated to specified session_id - active_session: Updated to specified session_id - messages: Loaded from sessions[session_id] + """ st.session_state.current_session = session_id st.session_state.active_session = session_id @@ -134,6 +135,7 @@ def get_llm_response(messages, model, reason, content_ph=None, reasoning_ph=None Note: The function uses streamlit placeholders for live updates. When reason=True, the reasoning process appears above the response. + """ full_text = "" think_text = "" @@ -256,6 +258,7 @@ def server_supports_reasoning(): Returns: bool: True if the model supports reasoning, False otherwise + """ resp = client.chat.completions.create( model=model, diff --git a/examples/applications/rag/retrieval_augmented_generation_with_langchain.py b/examples/applications/rag/retrieval_augmented_generation_with_langchain.py index d9a4cadb03..1681ea1466 100644 --- a/examples/applications/rag/retrieval_augmented_generation_with_langchain.py +++ b/examples/applications/rag/retrieval_augmented_generation_with_langchain.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Retrieval Augmented Generation (RAG) Implementation with Langchain +"""Retrieval Augmented Generation (RAG) Implementation with Langchain ================================================================== This script demonstrates a RAG implementation using LangChain, Milvus @@ -35,6 +34,7 @@ Notes: - Ensure both vLLM services are running before executing - Default ports: 8000 (embedding), 8001 (chat) - First run may take time to download models + """ import argparse @@ -52,9 +52,7 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter def load_and_split_documents(config: dict[str, Any]): - """ - Load and split documents from web URL - """ + """Load and split documents from web URL.""" try: loader = WebBaseLoader(web_paths=(config["url"],)) docs = loader.load() @@ -70,9 +68,7 @@ def load_and_split_documents(config: dict[str, Any]): def init_vectorstore(config: dict[str, Any], documents: list[Document]): - """ - Initialize vector store with documents - """ + """Initialize vector store with documents.""" return Milvus.from_documents( documents=documents, embedding=OpenAIEmbeddings( @@ -86,9 +82,7 @@ def init_vectorstore(config: dict[str, Any], documents: list[Document]): def init_llm(config: dict[str, Any]): - """ - Initialize llm - """ + """Initialize llm.""" return ChatOpenAI( model=config["chat_model"], openai_api_key=config["vllm_api_key"], @@ -97,9 +91,7 @@ def init_llm(config: dict[str, Any]): def get_qa_prompt(): - """ - Get question answering prompt template - """ + """Get question answering prompt template.""" template = """You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. @@ -112,16 +104,12 @@ Answer: def format_docs(docs: list[Document]): - """ - Format documents for prompt - """ + """Format documents for prompt.""" return "\n\n".join(doc.page_content for doc in docs) def create_qa_chain(retriever: Any, llm: ChatOpenAI, prompt: PromptTemplate): - """ - Set up question answering chain - """ + """Set up question answering chain.""" return ( { "context": retriever | format_docs, @@ -134,9 +122,7 @@ def create_qa_chain(retriever: Any, llm: ChatOpenAI, prompt: PromptTemplate): def get_parser() -> argparse.ArgumentParser: - """ - Parse command line arguments - """ + """Parse command line arguments.""" parser = argparse.ArgumentParser(description="RAG with vLLM and langchain") # Add command line arguments @@ -192,10 +178,7 @@ def get_parser() -> argparse.ArgumentParser: def init_config(args: Namespace): - """ - Initialize configuration settings from command line arguments - """ - + """Initialize configuration settings from command line arguments.""" return { "vllm_api_key": args.vllm_api_key, "vllm_embedding_endpoint": args.vllm_embedding_endpoint, diff --git a/examples/applications/rag/retrieval_augmented_generation_with_llamaindex.py b/examples/applications/rag/retrieval_augmented_generation_with_llamaindex.py index be4796acd1..0df32206bf 100644 --- a/examples/applications/rag/retrieval_augmented_generation_with_llamaindex.py +++ b/examples/applications/rag/retrieval_augmented_generation_with_llamaindex.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -RAG (Retrieval Augmented Generation) Implementation with LlamaIndex +"""RAG (Retrieval Augmented Generation) Implementation with LlamaIndex ================================================================ This script demonstrates a RAG system using: @@ -35,6 +34,7 @@ Notes: - Ensure both vLLM services are running before executing - Default ports: 8000 (embedding), 8001 (chat) - First run may take time to download models + """ import argparse @@ -50,7 +50,7 @@ from llama_index.vector_stores.milvus import MilvusVectorStore def init_config(args: Namespace): - """Initialize configuration with command line arguments""" + """Initialize configuration with command line arguments.""" return { "url": args.url, "embedding_model": args.embedding_model, @@ -66,12 +66,12 @@ def init_config(args: Namespace): def load_documents(url: str) -> list: - """Load and process web documents""" + """Load and process web documents.""" return SimpleWebPageReader(html_to_text=True).load_data([url]) def setup_models(config: dict[str, Any]): - """Configure embedding and chat models""" + """Configure embedding and chat models.""" Settings.embed_model = OpenAILikeEmbedding( api_base=config["embedding_endpoint"], api_key=config["vllm_api_key"], @@ -96,14 +96,14 @@ def setup_models(config: dict[str, Any]): def setup_vector_store(db_path: str) -> MilvusVectorStore: - """Initialize vector store""" + """Initialize vector store.""" sample_emb = Settings.embed_model.get_text_embedding("test") print(f"Embedding dimension: {len(sample_emb)}") return MilvusVectorStore(uri=db_path, dim=len(sample_emb), overwrite=True) def create_index(documents: list, vector_store: MilvusVectorStore): - """Create document index""" + """Create document index.""" storage_context = StorageContext.from_defaults(vector_store=vector_store) return VectorStoreIndex.from_documents( documents, @@ -112,13 +112,13 @@ def create_index(documents: list, vector_store: MilvusVectorStore): def query_document(index: VectorStoreIndex, question: str, top_k: int): - """Query document with given question""" + """Query document with given question.""" query_engine = index.as_query_engine(similarity_top_k=top_k) return query_engine.query(question) def get_parser() -> argparse.ArgumentParser: - """Parse command line arguments""" + """Parse command line arguments.""" parser = argparse.ArgumentParser(description="RAG with vLLM and LlamaIndex") # Add command line arguments diff --git a/examples/deployment/async_llm_streaming.py b/examples/deployment/async_llm_streaming.py index ef69089a04..bf9b36d05c 100644 --- a/examples/deployment/async_llm_streaming.py +++ b/examples/deployment/async_llm_streaming.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Simple example demonstrating streaming offline inference with AsyncLLM (V1 engine). +"""Simple example demonstrating streaming offline inference with AsyncLLM (V1 engine). This script shows the core functionality of vLLM's AsyncLLM engine for streaming token-by-token output in offline inference scenarios. It demonstrates DELTA mode @@ -20,8 +19,7 @@ from vllm.v1.engine.async_llm import AsyncLLM async def stream_response(engine: AsyncLLM, prompt: str, request_id: str) -> None: - """ - Stream response from AsyncLLM and display tokens as they arrive. + """Stream response from AsyncLLM and display tokens as they arrive. This function demonstrates the core streaming pattern: 1. Create SamplingParams with DELTA output kind diff --git a/examples/deployment/llm_engine_example.py b/examples/deployment/llm_engine_example.py index d9215255a8..0d25a157b4 100644 --- a/examples/deployment/llm_engine_example.py +++ b/examples/deployment/llm_engine_example.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This file demonstrates using the `LLMEngine` +"""This file demonstrates using the `LLMEngine` for processing prompts with various sampling parameters. """ diff --git a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py index 830da9b5d6..c44c902690 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py +++ b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -disagg_encoder_proxy.py +"""disagg_encoder_proxy.py Proxy that routes OpenAI-compatible “/v1/chat/completions” requests to two clusters: @@ -200,8 +199,7 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict: def extract_mm_items(request_data: dict) -> list[dict]: - """ - Return *all* image/audio/video items that appear anywhere in `messages`. + """Return *all* image/audio/video items that appear anywhere in `messages`. Each returned dict looks like: { "type": "image_url", "image_url": {...} } @@ -224,8 +222,7 @@ async def fanout_encoder_primer( req_id: str, consumer_zmq: str | None = None, ) -> tuple[dict[int, dict], dict[str, Any]]: - """ - 1. Build one request *per MM item* with all text removed. + """1. Build one request *per MM item* with all text removed. 2. Send them concurrently to the encode cluster. 3. Raise if any of them fails. @@ -402,8 +399,7 @@ async def maybe_prefill( p_url: str, req_id: str, ) -> dict: - """ - - Do prefill-only task if p_url exist; + """- Do prefill-only task if p_url exist; - Return a new body carrying kv transfer params (for nixl connector) - Else, skip and return the original request data for decode @@ -849,14 +845,14 @@ async def _post_if_available( payload: dict, headers: dict, ) -> dict | None: - """ - POST `payload` to `url`. + """POST `payload` to `url`. Returns ------- • The decoded JSON body on success (2xx) • None if the endpoint does not exist (404) • Raises for anything else. + """ try: resp = await session.post(url, json=payload, headers=headers) @@ -877,9 +873,7 @@ async def _post_if_available( async def _profile_cmd(cmd: str, payload: dict, e_url: str, p_url: str, d_url: str): - """ - Fire & forget to both clusters, tolerate 404. - """ + """Fire & forget to both clusters, tolerate 404.""" headers = {"Authorization": f"Bearer {os.getenv('OPENAI_API_KEY', '')}"} encode_task = _post_if_available( diff --git a/examples/disaggregated/disaggregated_serving/disagg_proxy_demo.py b/examples/disaggregated/disaggregated_serving/disagg_proxy_demo.py index 57deef6a15..ae35ec6d76 100644 --- a/examples/disaggregated/disaggregated_serving/disagg_proxy_demo.py +++ b/examples/disaggregated/disaggregated_serving/disagg_proxy_demo.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This file provides a disaggregated prefilling proxy demo to demonstrate an +"""This file provides a disaggregated prefilling proxy demo to demonstrate an example usage of XpYd disaggregated prefilling. We can launch multiple vllm instances (2 for prefill and 2 for decode), and launch this proxy demo through: diff --git a/examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py b/examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py index cc1cc402d2..3102f344b0 100644 --- a/examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py +++ b/examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Disaggregated Prefill/Decode Proxy with Bidirectional KV Transfer +"""Disaggregated Prefill/Decode Proxy with Bidirectional KV Transfer. This proxy sits between clients and a vLLM Prefill/Decode (P/D) deployment, routing multi-turn chat requests so that each turn reuses KV cache blocks @@ -54,7 +53,7 @@ Benchmarking: it the benchmark sends OpenAI-schema-compliant payloads and every turn lands as a cache MISS in this proxy. - Example: +Example: python benchmarks/multi_turn/benchmark_serving_multi_turn.py \\ --model --served-model-name \\ --url http://:8000 \\ @@ -67,6 +66,7 @@ Benchmarking: Dependencies: pip install fastapi uvicorn httpx + """ from __future__ import annotations diff --git a/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py b/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py index a35a44fbba..53d8b210d6 100644 --- a/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py +++ b/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Push-mode disaggregated prefilling proxy demo. +"""Push-mode disaggregated prefilling proxy demo. Companion to ``disagg_proxy_demo.py`` (pull mode). The client-facing API is the same; the difference is in how P and D coordinate the KV transfer: diff --git a/examples/disaggregated/example_connector/decode_example.py b/examples/disaggregated/example_connector/decode_example.py index 2d575840e6..9b76e2f001 100644 --- a/examples/disaggregated/example_connector/decode_example.py +++ b/examples/disaggregated/example_connector/decode_example.py @@ -6,7 +6,7 @@ from vllm.config import KVTransferConfig def read_prompts(): - """Read prompts from output.txt""" + """Read prompts from output.txt.""" prompts = [] try: with open("output.txt") as f: diff --git a/examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py b/examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py index b67c245931..55f4fe5200 100644 --- a/examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py +++ b/examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This example shows how to use FlexKV with vLLM for prefix caching. +"""This example shows how to use FlexKV with vLLM for prefix caching. FlexKV is a distributed KV Store and multi-level cache management system for ultra-large-scale LLM inference. diff --git a/examples/disaggregated/kv_load_failure_recovery_offline/decode_example.py b/examples/disaggregated/kv_load_failure_recovery_offline/decode_example.py index db9c5a85f7..3c3d5dbff9 100644 --- a/examples/disaggregated/kv_load_failure_recovery_offline/decode_example.py +++ b/examples/disaggregated/kv_load_failure_recovery_offline/decode_example.py @@ -7,7 +7,7 @@ from vllm.config import KVTransferConfig def read_prompts(): - """Read prompts from prefill_output.txt""" + """Read prompts from prefill_output.txt.""" prompts = [] try: with open("prefill_output.txt") as f: diff --git a/examples/disaggregated/kv_load_failure_recovery_offline/load_recovery_example_connector.py b/examples/disaggregated/kv_load_failure_recovery_offline/load_recovery_example_connector.py index 5bae044267..2b2e94aa74 100644 --- a/examples/disaggregated/kv_load_failure_recovery_offline/load_recovery_example_connector.py +++ b/examples/disaggregated/kv_load_failure_recovery_offline/load_recovery_example_connector.py @@ -119,8 +119,7 @@ class LoadRecoveryExampleConnector(ExampleConnector): def update_state_after_alloc( self, request: Request, blocks: KVCacheBlocks, num_external_tokens: int ): - """ - Update KVConnector state after block allocation. + """Update KVConnector state after block allocation. If blocks were allocated, add to _requests_need_load, such that we load the KVs in the next forward pass. diff --git a/examples/disaggregated/lmcache/cpu_offload_lmcache.py b/examples/disaggregated/lmcache/cpu_offload_lmcache.py index b67a929e5d..8a84675952 100644 --- a/examples/disaggregated/lmcache/cpu_offload_lmcache.py +++ b/examples/disaggregated/lmcache/cpu_offload_lmcache.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This file demonstrates the example usage of CPU offloading +"""This file demonstrates the example usage of CPU offloading with LMCache in vLLM v1. Note that `lmcache` is needed to run this example. diff --git a/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_proxy_server.py b/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_proxy_server.py index e048aecff9..3a64139c63 100644 --- a/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_proxy_server.py +++ b/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_proxy_server.py @@ -14,9 +14,7 @@ from fastapi.responses import StreamingResponse @asynccontextmanager async def lifespan(app: FastAPI): - """ - Lifespan context manager to handle startup and shutdown events. - """ + """Lifespan context manager to handle startup and shutdown events.""" # Startup: Initialize clients prefiller_base_url = ( f"http://{global_args.prefiller_host}:{global_args.prefiller_port}/v1" @@ -106,9 +104,7 @@ app.state.decode_client = None async def send_request_to_service( client: httpx.AsyncClient, endpoint: str, req_data: dict ): - """ - Send a request to a service using a persistent client. - """ + """Send a request to a service using a persistent client.""" req_data = req_data.copy() req_data["max_tokens"] = 1 if "max_completion_tokens" in req_data: @@ -128,9 +124,7 @@ async def send_request_to_service( async def stream_service_response( client: httpx.AsyncClient, endpoint: str, req_data: dict ): - """ - Asynchronously stream the response from a service using a persistent client. - """ + """Asynchronously stream the response from a service using a persistent client.""" headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"} async with client.stream( "POST", endpoint, json=req_data, headers=headers diff --git a/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py b/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py index 489ff13212..a50d6cbbc3 100644 --- a/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py +++ b/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This file demonstrates the example usage of remote KV cache sharing +"""This file demonstrates the example usage of remote KV cache sharing with LMCache. We will launch 2 vllm instances, and launch an additional LMCache server. KV cache is transferred in the following manner: diff --git a/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py b/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py index 09880a32aa..6815f42061 100644 --- a/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py +++ b/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py @@ -65,9 +65,7 @@ async def get_prefiller_info(prefill_clients: list, ready: asyncio.Event): @asynccontextmanager async def lifespan(app: FastAPI): - """ - Lifespan context manager to handle startup and shutdown events. - """ + """Lifespan context manager to handle startup and shutdown events.""" # Startup: Initialize client pools for prefiller and decoder services app.state.prefill_clients = [] app.state.decode_clients = [] @@ -180,10 +178,12 @@ def _parse_prefill_urls(prefill_list): """Parse prefill URLs from --prefill arguments. Format: --prefill URL [BOOTSTRAP_PORT] + Example: --prefill http://prefill1:8080 9000 # With bootstrap port --prefill http://prefill2:8080 none # Explicitly no bootstrap port --prefill http://prefill3:8080 # Defaults to no bootstrap port + """ if not prefill_list: return [] @@ -228,8 +228,7 @@ def _parse_decode_urls(decode_list): def get_next_client(app, service_type: str): - """ - Get the next client in round-robin fashion. + """Get the next client in round-robin fashion. Args: app: The FastAPI app instance @@ -237,6 +236,7 @@ def get_next_client(app, service_type: str): Returns: The next client to use + """ if service_type == "prefill": return next(app.state.prefill_iterator) @@ -250,9 +250,7 @@ def get_next_client(app, service_type: str): async def send_request_to_service( client_info: dict, dp_rank: int, endpoint: str, req_data: dict, request_id: str ): - """ - Send a request to a service using a client from the pool. - """ + """Send a request to a service using a client from the pool.""" req_data = req_data.copy() req_data["kv_transfer_params"] = { "do_remote_decode": True, @@ -288,9 +286,7 @@ async def stream_service_response( req_data: dict, request_id: str, ): - """ - Asynchronously stream response from a service using a client from the pool. - """ + """Asynchronously stream response from a service using a client from the pool.""" headers = { "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", "X-Request-Id": request_id, diff --git a/examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py b/examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py index 801b4b7697..db033e1861 100644 --- a/examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py +++ b/examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Demonstration script for Automatic Prefix Caching (APC) in vLLM. +"""Demonstration script for Automatic Prefix Caching (APC) in vLLM. Automatic Prefix Caching (APC) allows the vLLM engine to reuse cached KV (key-value) pairs from previous prompts if a new query shares the same diff --git a/examples/features/batch_invariance/reproducibility_offline.py b/examples/features/batch_invariance/reproducibility_offline.py index 72c1e841dc..54afcdee9f 100644 --- a/examples/features/batch_invariance/reproducibility_offline.py +++ b/examples/features/batch_invariance/reproducibility_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Demonstrates how to achieve reproducibility in vLLM. +"""Demonstrates how to achieve reproducibility in vLLM. Main article: https://docs.vllm.ai/en/latest/usage/reproducibility.html """ diff --git a/examples/features/context_extension/context_extension_offline.py b/examples/features/context_extension/context_extension_offline.py index 3874288b5e..cf270a2230 100644 --- a/examples/features/context_extension/context_extension_offline.py +++ b/examples/features/context_extension/context_extension_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This script demonstrates how to extend the context length +"""This script demonstrates how to extend the context length of a Qwen model using the YARN method (rope_parameters) and run a simple chat example. diff --git a/examples/features/data_parallel/data_parallel_offline.py b/examples/features/data_parallel/data_parallel_offline.py index c38ff7297a..86e61e04f0 100644 --- a/examples/features/data_parallel/data_parallel_offline.py +++ b/examples/features/data_parallel/data_parallel_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Usage: +"""Usage: Single node: python examples/features/data_parallel/data_parallel_offline.py \ --model="ibm-research/PowerMoE-3b" \ diff --git a/examples/features/kv_events/kv_events_subscriber.py b/examples/features/kv_events/kv_events_subscriber.py index 9c1b5138bc..91d3829c54 100644 --- a/examples/features/kv_events/kv_events_subscriber.py +++ b/examples/features/kv_events/kv_events_subscriber.py @@ -18,7 +18,7 @@ class EventBatch(msgspec.Struct, array_like=True, omit_defaults=True, gc=False): class KVCacheEvent(msgspec.Struct, omit_defaults=True, gc=False, tag=True): - """Base class for all KV cache-related events""" + """Base class for all KV cache-related events.""" class BlockStored(KVCacheEvent): diff --git a/examples/features/logits_processor/custom.py b/examples/features/logits_processor/custom.py index ce000872dc..c5028aa797 100644 --- a/examples/features/logits_processor/custom.py +++ b/examples/features/logits_processor/custom.py @@ -48,7 +48,7 @@ from vllm.v1.sample.logits_processor.builtin import process_dict_updates # Hypothetical custom logits processor class DummyLogitsProcessor(LogitsProcessor): - """Fake logit processor to support unit testing and examples""" + """Fake logit processor to support unit testing and examples.""" @classmethod def validate_params(cls, params: SamplingParams): diff --git a/examples/features/logits_processor/custom_req.py b/examples/features/logits_processor/custom_req.py index 5763fff541..36b83ec958 100644 --- a/examples/features/logits_processor/custom_req.py +++ b/examples/features/logits_processor/custom_req.py @@ -55,10 +55,11 @@ logger = init_logger(__name__) class DummyPerReqLogitsProcessor: """The request-level logits processor masks out all logits except the - token id identified by `target_token`""" + token id identified by `target_token` + """ def __init__(self, target_token: int) -> None: - """Specify `target_token`""" + """Specify `target_token`.""" self.target_token = target_token def __call__( @@ -74,7 +75,8 @@ class DummyPerReqLogitsProcessor: class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor): """Example of wrapping a fake request-level logit processor to create a - batch-level logits processor""" + batch-level logits processor + """ @classmethod def validate_params(cls, params: SamplingParams): @@ -103,6 +105,7 @@ class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor): Returns: `Callable` request logits processor, or None + """ target_token: Any | None = params.extra_args and params.extra_args.get( "target_token" diff --git a/examples/features/logits_processor/custom_req_init.py b/examples/features/logits_processor/custom_req_init.py index acd2c47f23..4f29bd2714 100644 --- a/examples/features/logits_processor/custom_req_init.py +++ b/examples/features/logits_processor/custom_req_init.py @@ -56,10 +56,11 @@ logger = init_logger(__name__) class DummyPerReqLogitsProcessor: """The request-level logits processor masks out all logits except the - token id identified by `target_token`""" + token id identified by `target_token` + """ def __init__(self, target_token: int) -> None: - """Specify `target_token`""" + """Specify `target_token`.""" self.target_token = target_token def __call__( @@ -75,7 +76,8 @@ class DummyPerReqLogitsProcessor: class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor): """Example of overriding the wrapper class `__init__()` in order to utilize - info about the device type""" + info about the device type + """ @classmethod def validate_params(cls, params: SamplingParams): @@ -111,6 +113,7 @@ class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor): Returns: `Callable` request logits processor, or None + """ if ( not self.is_cuda diff --git a/examples/features/lora/lora_with_quantization_offline.py b/examples/features/lora/lora_with_quantization_offline.py index ef6dbd4534..23f165cc5f 100644 --- a/examples/features/lora/lora_with_quantization_offline.py +++ b/examples/features/lora/lora_with_quantization_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This example shows how to use LoRA with different quantization techniques +"""This example shows how to use LoRA with different quantization techniques for offline inference. Requires HuggingFace credentials for access. @@ -73,7 +72,6 @@ def initialize_engine( model: str, quantization: str, lora_repo: str | None ) -> LLMEngine: """Initialize the LLMEngine.""" - engine_args = EngineArgs( model=model, quantization=quantization, @@ -86,7 +84,6 @@ def initialize_engine( def main(): """Main function that sets up and runs the prompt processing.""" - test_configs = [ # QLoRA (https://arxiv.org/abs/2305.14314) { diff --git a/examples/features/lora/multilora_offline.py b/examples/features/lora/multilora_offline.py index 92021f9fb2..d870ca7562 100644 --- a/examples/features/lora/multilora_offline.py +++ b/examples/features/lora/multilora_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This example shows how to use the multi-LoRA functionality +"""This example shows how to use the multi-LoRA functionality for offline inference. Requires HuggingFace credentials for access to Llama2. diff --git a/examples/features/pause_resume/data_parallel_pause_resume.py b/examples/features/pause_resume/data_parallel_pause_resume.py index 1f11536e53..fa26d0aa3c 100644 --- a/examples/features/pause_resume/data_parallel_pause_resume.py +++ b/examples/features/pause_resume/data_parallel_pause_resume.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test pause/resume with Data Parallel (DP) via HTTP API. +"""Test pause/resume with Data Parallel (DP) via HTTP API. This example demonstrates coordinated pause/resume across multiple DP ranks. The pause synchronizes across all DP engines via all-reduce. diff --git a/examples/features/prompt_embed/prompt_embed_offline.py b/examples/features/prompt_embed/prompt_embed_offline.py index 9e90aa46b5..2b5d3121c8 100644 --- a/examples/features/prompt_embed/prompt_embed_offline.py +++ b/examples/features/prompt_embed/prompt_embed_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Demonstrates how to generate prompt embeddings using +"""Demonstrates how to generate prompt embeddings using Hugging Face Transformers and use them as input to vLLM for both single and batch inference. diff --git a/examples/features/reset_kv/reset_kv_offline.py b/examples/features/reset_kv/reset_kv_offline.py index 3fbe7fa754..c19d88a35f 100644 --- a/examples/features/reset_kv/reset_kv_offline.py +++ b/examples/features/reset_kv/reset_kv_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This file demonstrates preempt requests when using the `LLMEngine` +"""This file demonstrates preempt requests when using the `LLMEngine` for processing prompts with various sampling parameters. """ diff --git a/examples/features/sharded_state/load_sharded_state_offline.py b/examples/features/sharded_state/load_sharded_state_offline.py index e867db5d12..66d8a75107 100644 --- a/examples/features/sharded_state/load_sharded_state_offline.py +++ b/examples/features/sharded_state/load_sharded_state_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Validates the loading of a model saved with the sharded_state format. +"""Validates the loading of a model saved with the sharded_state format. This script demonstrates how to load a model that was previously saved using save_sharded_state_offline.py and validates it by running inference. Example usage: diff --git a/examples/features/sharded_state/save_sharded_state_offline.py b/examples/features/sharded_state/save_sharded_state_offline.py index 675f2e35a5..7ed0de1f10 100644 --- a/examples/features/sharded_state/save_sharded_state_offline.py +++ b/examples/features/sharded_state/save_sharded_state_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Saves each worker's model state dict directly to a checkpoint, which enables a +"""Saves each worker's model state dict directly to a checkpoint, which enables a fast load path for large tensor-parallel models where each worker only needs to read its own shard rather than the entire checkpoint. diff --git a/examples/features/speculative_decoding/mlpspeculator_offline.py b/examples/features/speculative_decoding/mlpspeculator_offline.py index 6a533eb5c9..b863b28222 100644 --- a/examples/features/speculative_decoding/mlpspeculator_offline.py +++ b/examples/features/speculative_decoding/mlpspeculator_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This file demonstrates the usage of text generation with an LLM model, +"""This file demonstrates the usage of text generation with an LLM model, comparing the performance with and without speculative decoding. Note that this example is out of date and not supported in vLLM v1. diff --git a/examples/features/structured_outputs/structured_outputs_offline.py b/examples/features/structured_outputs/structured_outputs_offline.py index 6b6099f71b..e716935755 100644 --- a/examples/features/structured_outputs/structured_outputs_offline.py +++ b/examples/features/structured_outputs/structured_outputs_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This file demonstrates the example usage of structured outputs +"""This file demonstrates the example usage of structured outputs in vLLM. It shows how to apply different constraints such as choice, regex, json schema, and grammar to produce structured and formatted results based on specific prompts. diff --git a/examples/features/torchrun/torchrun_dp_example_offline.py b/examples/features/torchrun/torchrun_dp_example_offline.py index f18f6042e9..1f8ac60377 100644 --- a/examples/features/torchrun/torchrun_dp_example_offline.py +++ b/examples/features/torchrun/torchrun_dp_example_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -experimental support for data-parallel inference with torchrun +"""experimental support for data-parallel inference with torchrun Note the data load balancing and distribution is done out of the vllm engine, no internal lb supported in external_launcher mode. diff --git a/examples/features/torchrun/torchrun_example_offline.py b/examples/features/torchrun/torchrun_example_offline.py index e41bcd420c..b4b7cbd6a9 100644 --- a/examples/features/torchrun/torchrun_example_offline.py +++ b/examples/features/torchrun/torchrun_example_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -experimental support for tensor-parallel inference with torchrun, +"""experimental support for tensor-parallel inference with torchrun, see https://github.com/vllm-project/vllm/issues/11400 for the motivation and use case for this example. run the script with `torchrun --nproc-per-node=4 torchrun_example_offline.py`, diff --git a/examples/generate/multimodal/audio_language_offline.py b/examples/generate/multimodal/audio_language_offline.py index a2e9b19756..8b2215d8c1 100644 --- a/examples/generate/multimodal/audio_language_offline.py +++ b/examples/generate/multimodal/audio_language_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This example shows how to use vLLM for running offline inference +"""This example shows how to use vLLM for running offline inference with the correct prompt format on audio language models. For most models, the prompt format should follow corresponding examples @@ -286,8 +285,7 @@ def run_minicpmo(question: str, audio_count: int) -> ModelRequestData: # Phi-4-multimodal-instruct def run_phi4mm(question: str, audio_count: int) -> ModelRequestData: - """ - Phi-4-multimodal-instruct supports both image and audio inputs. Here, we + """Phi-4-multimodal-instruct supports both image and audio inputs. Here, we show how to process audio inputs. """ model_path = snapshot_download("microsoft/Phi-4-multimodal-instruct") diff --git a/examples/generate/multimodal/encoder_decoder_multimodal_offline.py b/examples/generate/multimodal/encoder_decoder_multimodal_offline.py index 1ab087de3f..b4ea893fc0 100644 --- a/examples/generate/multimodal/encoder_decoder_multimodal_offline.py +++ b/examples/generate/multimodal/encoder_decoder_multimodal_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This example shows how to use vLLM for running offline inference with +"""This example shows how to use vLLM for running offline inference with the explicit/implicit prompt format on enc-dec LMMs for text generation. """ @@ -56,8 +55,7 @@ def run_whisper(): def run_fireredasr2(): - """ - FireRedASR2 – Automatic Speech Recognition model. + """FireRedASR2 – Automatic Speech Recognition model. This model uses a Conformer encoder + Qwen2 LLM decoder architecture for speech-to-text transcription. Audio is passed via the implicit diff --git a/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py b/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py index e1c4fd76d7..60aba7abd0 100644 --- a/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py +++ b/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py @@ -43,7 +43,6 @@ headers = {"User-Agent": "vLLM Example Client"} def encode_base64_content_from_url(content_url: str) -> str: """Encode a content retrieved from a remote url to base64 format.""" - with requests.get(content_url, headers=headers) as response: response.raise_for_status() result = base64.b64encode(response.content).decode("utf-8") @@ -53,7 +52,6 @@ def encode_base64_content_from_url(content_url: str) -> str: def encode_base64_content_from_file(file_path: str) -> str: """Encode a local file content to base64 format.""" - with open(file_path, "rb") as file: file_content = file.read() result = base64.b64encode(file_content).decode("utf-8") diff --git a/examples/generate/multimodal/qwen2_5_omni/only_thinker.py b/examples/generate/multimodal/qwen2_5_omni/only_thinker.py index 476ca91d9f..66bfd49003 100644 --- a/examples/generate/multimodal/qwen2_5_omni/only_thinker.py +++ b/examples/generate/multimodal/qwen2_5_omni/only_thinker.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This example shows how to use vLLM for running offline inference +"""This example shows how to use vLLM for running offline inference with the correct prompt format on Qwen2.5-Omni (thinker only). """ diff --git a/examples/generate/multimodal/qwen3_omni/only_thinker.py b/examples/generate/multimodal/qwen3_omni/only_thinker.py index 84dd5228e5..57e845ac4f 100644 --- a/examples/generate/multimodal/qwen3_omni/only_thinker.py +++ b/examples/generate/multimodal/qwen3_omni/only_thinker.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This example shows how to use vLLM for running offline inference +"""This example shows how to use vLLM for running offline inference with the correct prompt format on Qwen3-Omni (thinker only). """ diff --git a/examples/generate/multimodal/vision_language_multi_image_offline.py b/examples/generate/multimodal/vision_language_multi_image_offline.py index 8b12df3991..cccb75d838 100644 --- a/examples/generate/multimodal/vision_language_multi_image_offline.py +++ b/examples/generate/multimodal/vision_language_multi_image_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This example shows how to use vLLM for running offline inference with +"""This example shows how to use vLLM for running offline inference with multi-image input on vision language models for text generation, using the chat template defined by the model. """ @@ -913,11 +912,9 @@ def load_phi3v(question: str, image_urls: list[str]) -> ModelRequestData: def load_phi4mm(question: str, image_urls: list[str]) -> ModelRequestData: - """ - Phi-4-multimodal-instruct supports both image and audio inputs. Here, we + """Phi-4-multimodal-instruct supports both image and audio inputs. Here, we show how to process multi images inputs. """ - model_path = snapshot_download("microsoft/Phi-4-multimodal-instruct") # Since the vision-lora and speech-lora co-exist with the base model, # we have to manually specify the path of the lora weights. diff --git a/examples/generate/multimodal/vision_language_offline.py b/examples/generate/multimodal/vision_language_offline.py index 33e1a8a36b..7cd471d266 100644 --- a/examples/generate/multimodal/vision_language_offline.py +++ b/examples/generate/multimodal/vision_language_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This example shows how to use vLLM for running offline inference with +"""This example shows how to use vLLM for running offline inference with the correct prompt format on vision language models for text generation. For most models, the prompt format should follow corresponding examples @@ -1660,8 +1659,7 @@ def run_phi3v(questions: list[str], modality: str) -> ModelRequestData: # Phi-4-multimodal-instruct def run_phi4mm(questions: list[str], modality: str) -> ModelRequestData: - """ - Phi-4-multimodal-instruct supports both image and audio inputs. Here, we + """Phi-4-multimodal-instruct supports both image and audio inputs. Here, we show how to process image inputs. """ assert modality == "image" @@ -2283,8 +2281,7 @@ MODELS_SUPPORT_VIT_CUDA_GRAPH = [ def get_multi_modal_input(args): - """ - return { + """Return { "data": image or video, "question": question, } diff --git a/examples/generate/trace_replay_offline.py b/examples/generate/trace_replay_offline.py index 2baa03307d..498a46cbd4 100644 --- a/examples/generate/trace_replay_offline.py +++ b/examples/generate/trace_replay_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Trace-replay with vLLM offline inference. +"""Trace-replay with vLLM offline inference. Trace-replay lets you supply a known sequence of decode token IDs alongside the prompt. Instead of sampling from the model distribution, the engine diff --git a/examples/pooling/embed/openai_embedding_long_text/client.py b/examples/pooling/embed/openai_embedding_long_text/client.py index 7cc33b1f24..44e415f670 100644 --- a/examples/pooling/embed/openai_embedding_long_text/client.py +++ b/examples/pooling/embed/openai_embedding_long_text/client.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example script demonstrating long text embedding with chunked processing in vLLM. +"""Example script demonstrating long text embedding with chunked processing in vLLM. This example shows how to use vLLM's chunked processing feature to handle text inputs that exceed the model's maximum token length. The feature automatically diff --git a/examples/pooling/embed/vision_embedding_offline.py b/examples/pooling/embed/vision_embedding_offline.py index 2c82d0374d..5175a6cd46 100644 --- a/examples/pooling/embed/vision_embedding_offline.py +++ b/examples/pooling/embed/vision_embedding_offline.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -This example shows how to use vLLM for running offline inference with +"""This example shows how to use vLLM for running offline inference with the correct prompt format on vision language models for multimodal embedding. For most models, the prompt format should follow corresponding examples diff --git a/examples/pooling/embed/vision_embedding_online.py b/examples/pooling/embed/vision_embedding_online.py index 79d6e397c8..765528978b 100644 --- a/examples/pooling/embed/vision_embedding_online.py +++ b/examples/pooling/embed/vision_embedding_online.py @@ -37,8 +37,7 @@ def create_chat_embeddings( add_special_tokens: bool = False, padding: str | NotGiven = NOT_GIVEN, ) -> CreateEmbeddingResponse: - """ - Convenience function for accessing vLLM's Chat Embeddings API, + """Convenience function for accessing vLLM's Chat Embeddings API, which is an extension of OpenAI's existing Embeddings API. """ return client.post( @@ -56,13 +55,11 @@ def create_chat_embeddings( def run_clip(client: OpenAI, model: str): - """ - Start the server using: + """Start the server using: vllm serve openai/clip-vit-base-patch32 \ --runner pooling """ - response = create_chat_embeddings( client, messages=[ @@ -97,8 +94,7 @@ def run_clip(client: OpenAI, model: str): def run_dse_qwen2_vl(client: OpenAI, model: str): - """ - Start the server using: + """Start the server using: vllm serve MrLight/dse-qwen2-2b-mrl-v1 \ --runner pooling \ @@ -159,14 +155,12 @@ def run_dse_qwen2_vl(client: OpenAI, model: str): def run_qwen3_vl(client: OpenAI, model: str): - """ - Start the server using: + """Start the server using: vllm serve Qwen/Qwen3-VL-Embedding-2B \ --runner pooling \ --max-model-len 8192 """ - default_instruction = "Represent the user's input." print("Text embedding output:") @@ -266,14 +260,12 @@ def run_qwen3_vl(client: OpenAI, model: str): def run_siglip(client: OpenAI, model: str): - """ - Start the server using: + """Start the server using: vllm serve google/siglip-base-patch16-224 \ --runner pooling \ --chat-template template_basic.jinja """ - response = create_chat_embeddings( client, messages=[ @@ -311,8 +303,7 @@ def run_siglip(client: OpenAI, model: str): def run_vlm2vec(client: OpenAI, model: str): - """ - Start the server using: + """Start the server using: vllm serve TIGER-Lab/VLM2Vec-Full \ --runner pooling \ @@ -320,7 +311,6 @@ def run_vlm2vec(client: OpenAI, model: str): --max-model-len 4096 \ --chat-template examples/pooling/embed/template/vlm2vec_phi3v.jinja """ - response = create_chat_embeddings( client, messages=[ diff --git a/examples/pooling/plugin/prithvi_geospatial_mae_offline.py b/examples/pooling/plugin/prithvi_geospatial_mae_offline.py index f7b30d9313..6bbde74264 100644 --- a/examples/pooling/plugin/prithvi_geospatial_mae_offline.py +++ b/examples/pooling/plugin/prithvi_geospatial_mae_offline.py @@ -85,8 +85,7 @@ def generate_datamodule(): def process_channel_group(orig_img, channels): - """ - Args: + """Args: orig_img: torch.Tensor representing original image (reference) with shape = (bands, H, W). channels: list of indices representing RGB channels. @@ -94,8 +93,8 @@ def process_channel_group(orig_img, channels): Returns: torch.Tensor with shape (num_channels, height, width) for original image - """ + """ orig_img = orig_img[channels, ...] valid_mask = torch.ones_like(orig_img, dtype=torch.bool) valid_mask[orig_img == NO_DATA_FLOAT] = False @@ -121,8 +120,8 @@ def read_geotiff(file_path: str): Returns: np.ndarray with shape (bands, height, width) meta info dict - """ + """ with rasterio.open(file_path) as src: img = src.read() meta = src.meta @@ -142,8 +141,8 @@ def save_geotiff(image, output_path: str, meta: dict): image: np.ndarray with shape (bands, height, width) output_path: path where to save the image meta: dict with meta info. - """ + """ with rasterio.open(output_path, "w", **meta) as dest: for i in range(image.shape[0]): dest.write(image[i, :, :], i + 1) @@ -172,12 +171,14 @@ def load_example( images in *file_paths*. std: list containing std values for each band in the images in *file_paths*. + indices: optional band indices to select from each image. All bands + are kept when None. Returns: np.array containing created example list of meta info for each image in *file_paths* - """ + """ imgs = [] metas = [] temporal_coords = [] diff --git a/examples/pooling/reward/sequence_reward_offline.py b/examples/pooling/reward/sequence_reward_offline.py index 0727bceee1..4453f4a390 100644 --- a/examples/pooling/reward/sequence_reward_offline.py +++ b/examples/pooling/reward/sequence_reward_offline.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example offline usage of sequence reward models. +"""Example offline usage of sequence reward models. The key distinction between sequence classification and token classification lies in their output granularity: sequence classification produces a single diff --git a/examples/pooling/reward/sequence_reward_online.py b/examples/pooling/reward/sequence_reward_online.py index 40d8d28e39..c8362265dd 100644 --- a/examples/pooling/reward/sequence_reward_online.py +++ b/examples/pooling/reward/sequence_reward_online.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example online usage of sequence reward models. +"""Example online usage of sequence reward models. Run `vllm serve --runner pooling` to start up the server in vLLM. e.g. diff --git a/examples/pooling/reward/token_reward_offline.py b/examples/pooling/reward/token_reward_offline.py index 4705c04912..1820c390f8 100644 --- a/examples/pooling/reward/token_reward_offline.py +++ b/examples/pooling/reward/token_reward_offline.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example offline usage of token reward models. +"""Example offline usage of token reward models. The key distinction between sequence classification and token classification lies in their output granularity: sequence classification produces a single diff --git a/examples/pooling/reward/token_reward_online.py b/examples/pooling/reward/token_reward_online.py index 64ee0c9dfd..38f2394451 100644 --- a/examples/pooling/reward/token_reward_online.py +++ b/examples/pooling/reward/token_reward_online.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example online usage of token reward models. +"""Example online usage of token reward models. Run `vllm serve --runner pooling` to start up the server in vLLM. e.g. diff --git a/examples/pooling/score/cohere_rerank_client.py b/examples/pooling/score/cohere_rerank_client.py index b32209967b..e7bf720cac 100644 --- a/examples/pooling/score/cohere_rerank_client.py +++ b/examples/pooling/score/cohere_rerank_client.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example of using the OpenAI entrypoint's rerank API which is compatible with +"""Example of using the OpenAI entrypoint's rerank API which is compatible with the Cohere SDK: https://github.com/cohere-ai/cohere-python Note that `pip install cohere` is needed to run this example. diff --git a/examples/pooling/score/colbert_rerank_online.py b/examples/pooling/score/colbert_rerank_online.py index 4cc509b95a..b8bc05b9e0 100644 --- a/examples/pooling/score/colbert_rerank_online.py +++ b/examples/pooling/score/colbert_rerank_online.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example of using ColBERT late interaction models for reranking and scoring. +"""Example of using ColBERT late interaction models for reranking and scoring. ColBERT (Contextualized Late Interaction over BERT) uses per-token embeddings and MaxSim scoring for document reranking, providing better accuracy than diff --git a/examples/pooling/score/colmodernvbert_rerank_online.py b/examples/pooling/score/colmodernvbert_rerank_online.py index de827ae062..95ca36ed71 100644 --- a/examples/pooling/score/colmodernvbert_rerank_online.py +++ b/examples/pooling/score/colmodernvbert_rerank_online.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example of using ColModernVBERT late interaction model for reranking. +"""Example of using ColModernVBERT late interaction model for reranking. ColModernVBERT is a multi-modal ColBERT-style model combining a SigLIP vision encoder with a ModernBERT text encoder. It produces per-token diff --git a/examples/pooling/score/colqwen3_5_rerank_online.py b/examples/pooling/score/colqwen3_5_rerank_online.py index 00746634d5..a2cd635361 100644 --- a/examples/pooling/score/colqwen3_5_rerank_online.py +++ b/examples/pooling/score/colqwen3_5_rerank_online.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example of using ColQwen3.5 late interaction model for reranking. +"""Example of using ColQwen3.5 late interaction model for reranking. ColQwen3.5 is a multi-modal ColBERT-style model based on Qwen3.5. It produces per-token embeddings and uses MaxSim scoring for retrieval diff --git a/examples/pooling/score/colqwen3_rerank_online.py b/examples/pooling/score/colqwen3_rerank_online.py index 0e61531bfd..22d5fee079 100644 --- a/examples/pooling/score/colqwen3_rerank_online.py +++ b/examples/pooling/score/colqwen3_rerank_online.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -Example of using ColQwen3 late interaction model for reranking and scoring. +"""Example of using ColQwen3 late interaction model for reranking and scoring. ColQwen3 is a multi-modal ColBERT-style model based on Qwen3-VL. It produces per-token embeddings and uses MaxSim scoring for retrieval diff --git a/examples/pooling/score/convert_model_to_seq_cls.py b/examples/pooling/score/convert_model_to_seq_cls.py index 21e2941664..494f683ae7 100644 --- a/examples/pooling/score/convert_model_to_seq_cls.py +++ b/examples/pooling/score/convert_model_to_seq_cls.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -Script to convert Large Language Models (LLMs) to Sequence Classification models. +"""Script to convert Large Language Models (LLMs) to Sequence Classification models. This is particularly useful for converting reranker models that use next-token prediction to a sequence classification format for compatibility with standard classification and rerank pipelines. @@ -35,8 +34,7 @@ import transformers def from_2_way_softmax(causal_lm, seq_cls_model, tokenizer, tokens, device): - """ - This method extracts the difference between weights for 'true' and 'false' tokens + """This method extracts the difference between weights for 'true' and 'false' tokens from the language model head to create a single classification weight vector. Args: @@ -47,6 +45,7 @@ def from_2_way_softmax(causal_lm, seq_cls_model, tokenizer, tokens, device): device: Target device (cpu/cuda) Reference: https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3 + """ assert len(tokens) == 2, ( "Method requires exactly two tokens for binary classification" @@ -73,8 +72,7 @@ def from_2_way_softmax(causal_lm, seq_cls_model, tokenizer, tokens, device): def no_post_processing(causal_lm, seq_cls_model, tokenizer, tokens, device): - """ - Directly use token weights from the language model head for classification. + """Directly use token weights from the language model head for classification. This method maps each classification label directly to a corresponding token in the vocabulary without additional transformation. @@ -85,6 +83,7 @@ def no_post_processing(causal_lm, seq_cls_model, tokenizer, tokens, device): tokenizer: Model tokenizer tokens: List of tokens representing class labels device: Target device (cpu/cuda) + """ # Get the language model head weights (vocabulary_size x hidden_size) lm_head_weights = causal_lm.lm_head.weight @@ -110,8 +109,7 @@ method_map = { def converting( model_name, classifier_from_tokens, path, method, use_sep_token=False, device="cpu" ): - """ - Main conversion function to transform a CausalLM model to SequenceClassification. + """Main conversion function to transform a CausalLM model to SequenceClassification. Args: model_name: Name or path of the pretrained model @@ -120,6 +118,7 @@ def converting( method: Conversion method ('from_2_way_softmax' or 'no_post_processing') use_sep_token: Whether to use separating token in the sequence classification model device: Device to load the model on ('cpu' or 'cuda') + """ assert method in method_map, f"Unknown method: {method}" diff --git a/examples/pooling/score/qwen3_reranker_offline.py b/examples/pooling/score/qwen3_reranker_offline.py index c79ebf97fe..dbe8d47522 100644 --- a/examples/pooling/score/qwen3_reranker_offline.py +++ b/examples/pooling/score/qwen3_reranker_offline.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -What is the difference between the official original version and one +"""What is the difference between the official original version and one that has been converted into a sequence classification model? Qwen3-Reranker is a language model that doing reranker by using the @@ -34,8 +33,7 @@ model_name = "Qwen/Qwen3-Reranker-0.6B" def get_llm() -> LLM: - """ - Initializes and returns the LLM model for Qwen3-Reranker. + """Initializes and returns the LLM model for Qwen3-Reranker. Returns: LLM: Configured vLLM instance for reranking tasks. @@ -43,6 +41,7 @@ def get_llm() -> LLM: Note: This function loads the ORIGINAL Qwen3-Reranker model with specific overrides to make it compatible with vLLM's score API. + """ return LLM( # Specify the original model from HuggingFace diff --git a/examples/pooling/score/qwen3_reranker_online.py b/examples/pooling/score/qwen3_reranker_online.py index f117b3b5b8..324e3f5bc3 100644 --- a/examples/pooling/score/qwen3_reranker_online.py +++ b/examples/pooling/score/qwen3_reranker_online.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -What is the difference between the official original version and one +"""What is the difference between the official original version and one that has been converted into a sequence classification model? Qwen3-Reranker is a language model that doing reranker by using the diff --git a/examples/pooling/score/rerank_api_online.py b/examples/pooling/score/rerank_api_online.py index 908d6a9240..caf0586282 100644 --- a/examples/pooling/score/rerank_api_online.py +++ b/examples/pooling/score/rerank_api_online.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example of using the OpenAI entrypoint's rerank API which is compatible with +"""Example of using the OpenAI entrypoint's rerank API which is compatible with Jina and Cohere https://jina.ai/reranker run: vllm serve BAAI/bge-reranker-base diff --git a/examples/pooling/score/score_api_online.py b/examples/pooling/score/score_api_online.py index af2886f136..a60cf04bf8 100644 --- a/examples/pooling/score/score_api_online.py +++ b/examples/pooling/score/score_api_online.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example online usage of Score API. +"""Example online usage of Score API. Run `vllm serve --runner pooling` to start up the server in vLLM. """ diff --git a/examples/pooling/score/using_template_offline.py b/examples/pooling/score/using_template_offline.py index f434e699ff..abe8157d4c 100644 --- a/examples/pooling/score/using_template_offline.py +++ b/examples/pooling/score/using_template_offline.py @@ -68,13 +68,13 @@ def get_chat_template(model: str) -> str: def get_hf_overrides(model: str) -> dict[str, Any]: """Convert Large Language Models (LLMs) to Sequence Classification models. - note: + Note: Some reranking models require special configuration overrides to work correctly with vLLM's score API. Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/qwen3_reranker_offline.py Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/convert_model_to_seq_cls.py - """ + """ model_name_to_hf_overrides_map = { "BAAI/bge-reranker-v2-gemma": { "architectures": ["GemmaForSequenceClassification"], @@ -122,7 +122,6 @@ def get_hf_overrides(model: str) -> dict[str, Any]: def main(args: Namespace): """Main execution function for the reranking example.""" - # Get the overrides for the specified model args.hf_overrides = get_hf_overrides(args.model) diff --git a/examples/pooling/score/using_template_online.py b/examples/pooling/score/using_template_online.py index f0bfa7d157..f7dc109dd8 100644 --- a/examples/pooling/score/using_template_online.py +++ b/examples/pooling/score/using_template_online.py @@ -1,15 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -Example of using the rerank API with template. +"""Example of using the rerank API with template. This script demonstrates how to interact with a vLLM server running a reranking model via the REST API. Before running this script, start the vLLM server with one of the supported reranking models using the commands below. -note: +Note: Some reranking models require special configuration overrides to work correctly with vLLM's score API. Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/qwen3_reranker_online.py @@ -21,6 +20,7 @@ run: vllm serve mixedbread-ai/mxbai-rerank-base-v2 --hf_overrides '{"architectures": ["Qwen2ForSequenceClassification"],"classifier_from_token": ["0", "1"], "method": "from_2_way_softmax"}' --chat-template examples/pooling/score/template/mxbai_rerank_v2.jinja vllm serve nvidia/llama-nemotron-rerank-1b-v2 --runner pooling --trust-remote-code --chat-template examples/pooling/score/template/nemotron-rerank.jinja vllm serve Qwen/Qwen3-Reranker-0.6B --runner pooling --hf_overrides '{"architectures": ["Qwen3ForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}' --chat-template examples/pooling/score/template/qwen3_reranker.jinja + """ import json diff --git a/examples/pooling/score/vision_rerank_api_online.py b/examples/pooling/score/vision_rerank_api_online.py index dce2efd1dc..cc1f2e97bf 100644 --- a/examples/pooling/score/vision_rerank_api_online.py +++ b/examples/pooling/score/vision_rerank_api_online.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -Example Python client for multimodal rerank API which is compatible with +"""Example Python client for multimodal rerank API which is compatible with Jina and Cohere https://jina.ai/reranker Run `vllm serve --runner pooling` to start up the server in vLLM. diff --git a/examples/pooling/score/vision_reranker_offline.py b/examples/pooling/score/vision_reranker_offline.py index ef2b354877..24c8b9068f 100644 --- a/examples/pooling/score/vision_reranker_offline.py +++ b/examples/pooling/score/vision_reranker_offline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This example shows how to use vLLM for running offline inference with +"""This example shows how to use vLLM for running offline inference with vision language reranker models for multimodal scoring tasks. Vision language rerankers score the relevance between a text query and diff --git a/examples/pooling/score/vision_score_api_online.py b/examples/pooling/score/vision_score_api_online.py index 543d4bfa21..3a3ac62cf7 100644 --- a/examples/pooling/score/vision_score_api_online.py +++ b/examples/pooling/score/vision_score_api_online.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -Example online usage of Score API. +"""Example online usage of Score API. Run `vllm serve --runner pooling` to start up the server in vLLM. e.g. diff --git a/examples/pooling/token_classify/forced_alignment_offline.py b/examples/pooling/token_classify/forced_alignment_offline.py index e97a373576..13facab480 100644 --- a/examples/pooling/token_classify/forced_alignment_offline.py +++ b/examples/pooling/token_classify/forced_alignment_offline.py @@ -3,8 +3,7 @@ # Adapted from Qwen3-ForcedAligner inference: # https://github.com/QwenLM/Qwen3-ASR -""" -Offline forced alignment example using Qwen3-ForcedAligner-0.6B. +"""Offline forced alignment example using Qwen3-ForcedAligner-0.6B. Forced alignment takes audio and reference text as input and produces word-level timestamps. The model predicts a time bin at each diff --git a/examples/pooling/token_classify/forced_alignment_online.py b/examples/pooling/token_classify/forced_alignment_online.py index e967a8ba70..ab3026160d 100644 --- a/examples/pooling/token_classify/forced_alignment_online.py +++ b/examples/pooling/token_classify/forced_alignment_online.py @@ -3,8 +3,7 @@ # Adapted from Qwen3-ForcedAligner inference: # https://github.com/QwenLM/Qwen3-ASR -""" -Online forced alignment example using Qwen3-ForcedAligner-0.6B-hf. +"""Online forced alignment example using Qwen3-ForcedAligner-0.6B-hf. Forced alignment takes audio and reference text as input and produces word-level timestamps. The model predicts a time bin at each diff --git a/examples/pooling/token_classify/ner_online.py b/examples/pooling/token_classify/ner_online.py index 9ec2bd45a0..bd60d68626 100644 --- a/examples/pooling/token_classify/ner_online.py +++ b/examples/pooling/token_classify/ner_online.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Adapted from https://huggingface.co/boltuix/NeuroBERT-NER -""" -Example online usage of Pooling API for Named Entity Recognition (NER). +"""Example online usage of Pooling API for Named Entity Recognition (NER). Run `vllm serve --runner pooling` to start up the server in vLLM. e.g. diff --git a/examples/pooling/token_embed/colqwen3_token_embed_online.py b/examples/pooling/token_embed/colqwen3_token_embed_online.py index cac11188e8..c1a8f3a56e 100644 --- a/examples/pooling/token_embed/colqwen3_token_embed_online.py +++ b/examples/pooling/token_embed/colqwen3_token_embed_online.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -Example online usage of Pooling API for ColQwen3 multi-vector retrieval. +"""Example online usage of Pooling API for ColQwen3 multi-vector retrieval. ColQwen3 is a multi-modal late interaction model based on Qwen3-VL that produces per-token embeddings (320-dim, L2-normalized) for both text and diff --git a/examples/pooling/token_embed/jina_reranker_v3_online.py b/examples/pooling/token_embed/jina_reranker_v3_online.py index 8350aee2f1..b9018f0f23 100644 --- a/examples/pooling/token_embed/jina_reranker_v3_online.py +++ b/examples/pooling/token_embed/jina_reranker_v3_online.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -Example online usage of the Jina Reranker v3 score and rerank APIs with a task +"""Example online usage of the Jina Reranker v3 score and rerank APIs with a task instruction. Run `vllm serve jinaai/jina-reranker-v3 --runner pooling` to start up the diff --git a/examples/pooling/token_embed/multi_vector_retrieval_online.py b/examples/pooling/token_embed/multi_vector_retrieval_online.py index c2aae88294..f36e3dc31d 100644 --- a/examples/pooling/token_embed/multi_vector_retrieval_online.py +++ b/examples/pooling/token_embed/multi_vector_retrieval_online.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example online usage of Pooling API for multi vector retrieval. +"""Example online usage of Pooling API for multi vector retrieval. Run `vllm serve --runner pooling` to start up the server in vLLM. e.g. diff --git a/examples/ray_serving/batch_llm_inference.py b/examples/ray_serving/batch_llm_inference.py index 22408dc950..1b265efb17 100644 --- a/examples/ray_serving/batch_llm_inference.py +++ b/examples/ray_serving/batch_llm_inference.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This example shows how to use Ray Data for data parallel batch inference. +"""This example shows how to use Ray Data for data parallel batch inference. Ray Data is a data processing framework that can process very large datasets with first-class support for vLLM. diff --git a/examples/ray_serving/ray_serve_deepseek.py b/examples/ray_serving/ray_serve_deepseek.py index af53443b91..a1e9348115 100644 --- a/examples/ray_serving/ray_serve_deepseek.py +++ b/examples/ray_serving/ray_serve_deepseek.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Deploy DeepSeek R1 or V3 with Ray Serve LLM. +"""Deploy DeepSeek R1 or V3 with Ray Serve LLM. Ray Serve LLM is a scalable and production-grade model serving library built on the Ray distributed computing framework and first-class support for the vLLM engine. diff --git a/examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py b/examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py index 1dfc308464..53b30f534d 100644 --- a/examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py +++ b/examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -An example demonstrates how to use tool calling with reasoning models +"""An example demonstrates how to use tool calling with reasoning models like QwQ-32B. The reasoning will not be parsed by the tool calling process; only the final output will be parsed. diff --git a/examples/reasoning/openai_chat_completion_with_reasoning.py b/examples/reasoning/openai_chat_completion_with_reasoning.py index 87043897b0..a72606108f 100644 --- a/examples/reasoning/openai_chat_completion_with_reasoning.py +++ b/examples/reasoning/openai_chat_completion_with_reasoning.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -An example shows how to generate chat completions from reasoning models +"""An example shows how to generate chat completions from reasoning models like DeepSeekR1. To run this example, you need to start the vLLM server diff --git a/examples/reasoning/openai_chat_completion_with_reasoning_streaming.py b/examples/reasoning/openai_chat_completion_with_reasoning_streaming.py index 8e262701b7..36f4c1f8e6 100644 --- a/examples/reasoning/openai_chat_completion_with_reasoning_streaming.py +++ b/examples/reasoning/openai_chat_completion_with_reasoning_streaming.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -An example shows how to generate chat completions from reasoning models +"""An example shows how to generate chat completions from reasoning models like DeepSeekR1. To run this example, you need to start the vLLM server with the reasoning diff --git a/examples/reasoning/openai_responses_client.py b/examples/reasoning/openai_responses_client.py index b4eb246715..94ecd84468 100644 --- a/examples/reasoning/openai_responses_client.py +++ b/examples/reasoning/openai_responses_client.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Set up this example by starting a vLLM OpenAI-compatible server. +"""Set up this example by starting a vLLM OpenAI-compatible server. Reasoning models can be used through the Responses API as seen here https://platform.openai.com/docs/api-reference/responses For example: diff --git a/examples/rl/rlhf_async_new_apis.py b/examples/rl/rlhf_async_new_apis.py index ea7099962f..081023156e 100644 --- a/examples/rl/rlhf_async_new_apis.py +++ b/examples/rl/rlhf_async_new_apis.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Demonstrates async reinforcement learning using vLLM and Ray, +"""Demonstrates async reinforcement learning using vLLM and Ray, with native weight syncing APIs and batch-invariant generation. The script separates training and inference workloads onto distinct GPUs diff --git a/examples/rl/rlhf_http_ipc.py b/examples/rl/rlhf_http_ipc.py index c688707df9..0ba9350216 100644 --- a/examples/rl/rlhf_http_ipc.py +++ b/examples/rl/rlhf_http_ipc.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -RLHF weight syncing against a `vllm serve` HTTP server, using CUDA IPC for the +"""RLHF weight syncing against a `vllm serve` HTTP server, using CUDA IPC for the data plane. * OpenAI-compatible API for inference requests diff --git a/examples/rl/rlhf_http_nccl.py b/examples/rl/rlhf_http_nccl.py index 97b0db0b0b..d095b647aa 100644 --- a/examples/rl/rlhf_http_nccl.py +++ b/examples/rl/rlhf_http_nccl.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -RLHF weight syncing against a `vllm serve` HTTP server, using NCCL for the +"""RLHF weight syncing against a `vllm serve` HTTP server, using NCCL for the data plane. * OpenAI-compatible API for inference requests diff --git a/examples/rl/rlhf_ipc_fsdp_ep.py b/examples/rl/rlhf_ipc_fsdp_ep.py index e3fa1c44f1..4eb3fbcbbb 100644 --- a/examples/rl/rlhf_ipc_fsdp_ep.py +++ b/examples/rl/rlhf_ipc_fsdp_ep.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -RLHF with FSDP2 training and vLLM expert-parallel inference using **CUDA IPC** +"""RLHF with FSDP2 training and vLLM expert-parallel inference using **CUDA IPC** weight transfer and **packed** tensors. Multi-rank version of `rlhf_http_ipc.py`: the trainer is 4 FSDP2 Ray actors diff --git a/examples/rl/rlhf_nccl_fsdp_ep.py b/examples/rl/rlhf_nccl_fsdp_ep.py index e4380bea94..b9b72c012e 100644 --- a/examples/rl/rlhf_nccl_fsdp_ep.py +++ b/examples/rl/rlhf_nccl_fsdp_ep.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -RLHF with FSDP2 training (4 GPUs) and vLLM expert-parallel inference (4 GPUs). +"""RLHF with FSDP2 training (4 GPUs) and vLLM expert-parallel inference (4 GPUs). 8-GPU layout: Training — 4 GPUs, PyTorch FSDP2 (fully_shard), as Ray actors @@ -74,8 +73,7 @@ BASE_URL = f"http://localhost:{SERVER_PORT}" @ray.remote(num_gpus=1) class FSDPTrainWorker: - """ - One FSDP2 training worker per GPU. Four of these form the FSDP group. + """One FSDP2 training worker per GPU. Four of these form the FSDP group. Rank 0 additionally drives weight transfer to the vLLM server. """ diff --git a/examples/rl/rlhf_sharded_rdt_small_ep.py b/examples/rl/rlhf_sharded_rdt_small_ep.py index 213eae55ef..ed31a32ad8 100644 --- a/examples/rl/rlhf_sharded_rdt_small_ep.py +++ b/examples/rl/rlhf_sharded_rdt_small_ep.py @@ -101,7 +101,8 @@ PROMPTS = ["The capital of France is", "The future of AI is"] class FSDPTrainWorker: """One FSDP2 rank per GPU. Identical to the 8-GPU example's worker: the sharded-RDT engine owns the NIXL serve surface, so this stays a plain Ray - actor with no producer mixin.""" + actor with no producer mixin. + """ def __init__(self, model_name, rank, world_size, master_addr, master_port): self.rank = rank diff --git a/examples/rl/routed_experts_e2e.py b/examples/rl/routed_experts_e2e.py index 1666bc3ffe..04b5871fd1 100644 --- a/examples/rl/routed_experts_e2e.py +++ b/examples/rl/routed_experts_e2e.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -End-to-end example for routed experts capture with hybrid models. +"""End-to-end example for routed experts capture with hybrid models. Validates that: 1. routed_experts is returned in CompletionOutput for MoE models. diff --git a/examples/speech_to_text/openai/openai_transcription_client.py b/examples/speech_to_text/openai/openai_transcription_client.py index f928c06d45..cc6757bd96 100644 --- a/examples/speech_to_text/openai/openai_transcription_client.py +++ b/examples/speech_to_text/openai/openai_transcription_client.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This script demonstrates how to use the vLLM API server to perform audio +"""This script demonstrates how to use the vLLM API server to perform audio transcription with the `openai/whisper-large-v3` model. Before running this script, you must start the vLLM server with the following command: @@ -35,8 +34,7 @@ def sync_openai( hotwords: str = None, prompt: str | None = None, ): - """ - Perform synchronous transcription using OpenAI-compatible API. + """Perform synchronous transcription using OpenAI-compatible API. The optional ``prompt`` is the OpenAI-API ``prompt`` field (style / vocabulary hint). It is wired through model-by-model: Whisper uses it @@ -69,9 +67,7 @@ async def stream_openai_response( hotwords: str = None, prompt: str | None = None, ): - """ - Perform asynchronous transcription using OpenAI-compatible API. - """ + """Perform asynchronous transcription using OpenAI-compatible API.""" print("\ntranscription result [stream]:", end=" ") with open(audio_path, "rb") as f: transcription = await client.audio.transcriptions.create( @@ -98,9 +94,7 @@ async def stream_openai_response( def stream_api_response(audio_path: str, model: str, openai_api_base: str): - """ - Perform streaming transcription using raw HTTP requests to the vLLM API server. - """ + """Stream a transcription using raw HTTP requests to the vLLM API server.""" import json import os diff --git a/examples/speech_to_text/realtime/openai_realtime_client.py b/examples/speech_to_text/realtime/openai_realtime_client.py index fda3d7cb45..a290b63e83 100644 --- a/examples/speech_to_text/realtime/openai_realtime_client.py +++ b/examples/speech_to_text/realtime/openai_realtime_client.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This script demonstrates how to use the vLLM Realtime WebSocket API to perform +"""This script demonstrates how to use the vLLM Realtime WebSocket API to perform audio transcription by uploading an audio file. Before running this script, you must start the vLLM server with a realtime-capable @@ -34,9 +33,7 @@ from vllm.multimodal.media.audio import load_audio def audio_to_pcm16_base64(audio_path: str) -> str: - """ - Load an audio file and convert it to base64-encoded PCM16 @ 16kHz. - """ + """Load an audio file and convert it to base64-encoded PCM16 @ 16kHz.""" # Load audio and resample to 16kHz mono audio, _ = load_audio(audio_path, sr=16000, mono=True) # Convert to PCM16 @@ -46,9 +43,7 @@ def audio_to_pcm16_base64(audio_path: str) -> str: async def realtime_transcribe(audio_path: str, host: str, port: int, model: str): - """ - Connect to the Realtime API and transcribe an audio file. - """ + """Connect to the Realtime API and transcribe an audio file.""" uri = f"ws://{host}:{port}/v1/realtime" async with websockets.connect(uri) as ws: diff --git a/examples/speech_to_text/realtime/openai_realtime_microphone_client.py b/examples/speech_to_text/realtime/openai_realtime_microphone_client.py index a3c07673ff..c29486120c 100644 --- a/examples/speech_to_text/realtime/openai_realtime_microphone_client.py +++ b/examples/speech_to_text/realtime/openai_realtime_microphone_client.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Minimal Gradio demo for real-time speech transcription using the vLLM Realtime API. +"""Minimal Gradio demo for real-time speech transcription using the vLLM Realtime API. Start the vLLM server first: diff --git a/examples/tool_calling/openai_chat_completion_client_with_tools.py b/examples/tool_calling/openai_chat_completion_client_with_tools.py index 0bd1d05322..508c6900b6 100644 --- a/examples/tool_calling/openai_chat_completion_client_with_tools.py +++ b/examples/tool_calling/openai_chat_completion_client_with_tools.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Set up this example by starting a vLLM OpenAI-compatible server with tool call +"""Set up this example by starting a vLLM OpenAI-compatible server with tool call options enabled. For example: IMPORTANT: for mistral, you must use one of the provided mistral tool call diff --git a/examples/tool_calling/openai_chat_completion_client_with_tools_required.py b/examples/tool_calling/openai_chat_completion_client_with_tools_required.py index c00d712b35..f38011e10d 100644 --- a/examples/tool_calling/openai_chat_completion_client_with_tools_required.py +++ b/examples/tool_calling/openai_chat_completion_client_with_tools_required.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -To run this example, you can start the vLLM server +"""To run this example, you can start the vLLM server without any specific flags: ```bash diff --git a/examples/tool_calling/openai_chat_completion_client_with_tools_xlam.py b/examples/tool_calling/openai_chat_completion_client_with_tools_xlam.py index f0b0a2db44..004b5fc7fc 100644 --- a/examples/tool_calling/openai_chat_completion_client_with_tools_xlam.py +++ b/examples/tool_calling/openai_chat_completion_client_with_tools_xlam.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -Set up this example by starting a vLLM OpenAI-compatible server with tool call +"""Set up this example by starting a vLLM OpenAI-compatible server with tool call options enabled for xLAM-2 models: vllm serve --model Salesforce/Llama-xLAM-2-8b-fc-r --enable-auto-tool-choice --tool-call-parser xlam @@ -105,8 +104,7 @@ tool_functions = { def process_response(response, tool_functions, original_query): - """Process a non-streaming response with possible tool calls""" - + """Process a non-streaming response with possible tool calls.""" print("\n--- Response Output ---") # Check if the response has content @@ -195,7 +193,7 @@ def process_response(response, tool_functions, original_query): def run_test_case(query, test_name): - """Run a single test case with the given query""" + """Run a single test case with the given query.""" print(f"\n{'=' * 50}\nTEST CASE: {test_name}\n{'=' * 50}") print(f"Query: '{query}'") diff --git a/examples/tool_calling/openai_chat_completion_client_with_tools_xlam_streaming.py b/examples/tool_calling/openai_chat_completion_client_with_tools_xlam_streaming.py index 94e664c9ec..43e4ed21a0 100644 --- a/examples/tool_calling/openai_chat_completion_client_with_tools_xlam_streaming.py +++ b/examples/tool_calling/openai_chat_completion_client_with_tools_xlam_streaming.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -Set up this example by starting a vLLM OpenAI-compatible server with tool call +"""Set up this example by starting a vLLM OpenAI-compatible server with tool call options enabled for xLAM-2 models: vllm serve --model Salesforce/Llama-xLAM-2-8b-fc-r --enable-auto-tool-choice --tool-call-parser xlam @@ -107,7 +106,7 @@ tool_functions = { def process_stream(response, tool_functions, original_query): - """Process a streaming response with possible tool calls""" + """Process a streaming response with possible tool calls.""" # Track multiple tool calls tool_calls = {} # Dictionary to store tool calls by ID @@ -223,7 +222,7 @@ def process_stream(response, tool_functions, original_query): def run_test_case(query, test_name): - """Run a single test case with the given query""" + """Run a single test case with the given query.""" print(f"\n{'=' * 50}\nTEST CASE: {test_name}\n{'=' * 50}") print(f"Query: '{query}'") diff --git a/examples/tool_calling/openai_responses_client_with_mcp_tools.py b/examples/tool_calling/openai_responses_client_with_mcp_tools.py index 1e685a4f34..0b6ea412bf 100644 --- a/examples/tool_calling/openai_responses_client_with_mcp_tools.py +++ b/examples/tool_calling/openai_responses_client_with_mcp_tools.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example demonstrating MCP (Model Context Protocol) tools with the Responses API. +"""Example demonstrating MCP (Model Context Protocol) tools with the Responses API. This example shows how to use MCP tools with different allowed_tools configurations: 1. No filter (allows all tools from the MCP server) diff --git a/examples/tool_calling/openai_responses_client_with_tools.py b/examples/tool_calling/openai_responses_client_with_tools.py index a47210f600..6530f331e4 100644 --- a/examples/tool_calling/openai_responses_client_with_tools.py +++ b/examples/tool_calling/openai_responses_client_with_tools.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Set up this example by starting a vLLM OpenAI-compatible server with tool call +"""Set up this example by starting a vLLM OpenAI-compatible server with tool call options enabled. Reasoning models can be used through the Responses API as seen here https://platform.openai.com/docs/api-reference/responses @@ -17,8 +16,7 @@ from openai import OpenAI def get_weather(latitude: float, longitude: float) -> str: - """ - Mock function to simulate getting weather data. + """Mock function to simulate getting weather data. In a real application, this would call an external weather API. """ return f"Current temperature at ({latitude}, {longitude}) is 20°C." diff --git a/pyproject.toml b/pyproject.toml index 1a77433f9c..83949db8fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,9 +60,13 @@ include = ["vllm*"] "vllm/third_party/**" = ["ALL"] "vllm/version.py" = ["F401"] "vllm/_version.py" = ["ALL"] +# test args are mostly pytest fixtures, which don't benefit from descriptions +"tests/**" = ["D417"] [tool.ruff.lint] select = [ + # pydocstyle + "D", # pycodestyle "E", # Pyflakes @@ -83,6 +87,10 @@ select = [ "INP", ] ignore = [ + # don't enforce that everything has a docstring + "D100", "D101", "D102", "D103", "D104", "D105", "D106", "D107", + # unwanted pydocstyle rules + "D205", "D301", "D400", "D401", "D404", "D415", # star imports "F405", "F403", # lambda expression assignment @@ -95,6 +103,9 @@ ignore = [ "UP032", ] +[tool.ruff.lint.pydocstyle] +ignore-var-parameters = true + [tool.ruff.format] docstring-code-format = true diff --git a/scripts/autotune_helion_kernels.py b/scripts/autotune_helion_kernels.py index 864aabfc09..e6352a40af 100644 --- a/scripts/autotune_helion_kernels.py +++ b/scripts/autotune_helion_kernels.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Autotune registered Helion kernels for optimal configurations. +"""Autotune registered Helion kernels for optimal configurations. Usage: # Autotune all registered kernels diff --git a/scripts/benchmark_helion_kernels.py b/scripts/benchmark_helion_kernels.py index fea30d6f9c..04350896af 100644 --- a/scripts/benchmark_helion_kernels.py +++ b/scripts/benchmark_helion_kernels.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Benchmark a registered Helion kernel against a baseline. +"""Benchmark a registered Helion kernel against a baseline. For each input case produced by the kernel's registered input generator, this checks the Helion kernel's numerics once against its eager reference, measures @@ -394,6 +393,7 @@ def check_kernel_correctness( ``passed=False`` with the exception text in ``error``; iteration continues regardless. An empty input mapping returns an empty list, which the CLI reports as a skipped check. + """ if inputs_dict is None: inputs_dict = kernel.get_inputs() diff --git a/setup.py b/setup.py index 4012000e63..48cbdcf763 100644 --- a/setup.py +++ b/setup.py @@ -512,8 +512,7 @@ class precompiled_wheel_utils: *, rocm: bool = False, ) -> tuple[list[dict], str]: - """ - Fetches metadata for a specific variant of the precompiled wheel. + """Fetches metadata for a specific variant of the precompiled wheel. For non-ROCm, fetches vllm metadata. @@ -591,7 +590,6 @@ class precompiled_wheel_utils: @staticmethod def detect_system_cuda_variant() -> str: """Auto-detect CUDA variant from torch, nvidia-smi, or env default.""" - # Map CUDA major version to hosted wheel variants on wheels.vllm.ai supported = {12: "cu129", 13: "cu130"} @@ -726,7 +724,8 @@ class precompiled_wheel_utils: wheels: list[dict], repo_url: str, arch: str ) -> None: """Warn if installed torch differs from the custom ROCm build on - wheels.vllm.ai and suggest the correct install command.""" + wheels.vllm.ai and suggest the correct install command. + """ try: installed = torch.__version__ except Exception: @@ -886,8 +885,7 @@ class precompiled_wheel_utils: @staticmethod def determine_wheel_url() -> tuple[str, str | None]: - """ - Try to determine the precompiled wheel URL or path to use. + """Try to determine the precompiled wheel URL or path to use. The order of preference is: 1. user-specified wheel location (can be either local or remote, via VLLM_PRECOMPILED_WHEEL_LOCATION) diff --git a/tests/basic_correctness/test_basic_correctness.py b/tests/basic_correctness/test_basic_correctness.py index 3412b30814..1b3c92ee5b 100644 --- a/tests/basic_correctness/test_basic_correctness.py +++ b/tests/basic_correctness/test_basic_correctness.py @@ -89,7 +89,7 @@ TARGET_TEST_SUITE = _resolve_target_test_suite() condition=current_platform.is_rocm(), ) def test_vllm_gc_ed(): - """Verify vllm instance is GC'ed when it is deleted""" + """Verify vllm instance is GC'ed when it is deleted.""" llm = LLM("hmellor/tiny-random-LlamaForCausalLM") weak_llm = weakref.ref(llm) del llm diff --git a/tests/basic_correctness/test_mem.py b/tests/basic_correctness/test_mem.py index c3f66e24ab..f04bb397f0 100644 --- a/tests/basic_correctness/test_mem.py +++ b/tests/basic_correctness/test_mem.py @@ -51,8 +51,7 @@ def _wake_up_with_poisoned_mappings(allocator, byte_value: int = 0xA5) -> None: @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") def test_python_error(): - """ - Test if Python error occurs when there's low-level + """Test if Python error occurs when there's low-level error happening from the C++ side. """ allocator = get_mem_allocator_instance() @@ -113,7 +112,8 @@ def test_basic_cumem(): @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") def test_discard_tags(): """Test that discard(tags) selectively frees GPU memory for specific - tags while keeping other tags mapped and usable.""" + tags while keeping other tags mapped and usable. + """ allocator = get_mem_allocator_instance() with allocator.use_memory_pool("weights"): diff --git a/tests/benchmarks/test_bfcl_dataset.py b/tests/benchmarks/test_bfcl_dataset.py index d191922319..982b93ad8e 100644 --- a/tests/benchmarks/test_bfcl_dataset.py +++ b/tests/benchmarks/test_bfcl_dataset.py @@ -13,7 +13,8 @@ from vllm.benchmarks.datasets import BFCLDataset, get_samples def _patch_hf_api(side_effect): """Return a patch context that swaps `hf_api()` to a stub whose - `.hf_hub_download` attribute uses `side_effect`.""" + `.hf_hub_download` attribute uses `side_effect`. + """ fake_api = MagicMock() fake_api.hf_hub_download.side_effect = side_effect return patch("vllm.benchmarks.datasets.datasets.hf_api", return_value=fake_api) @@ -115,7 +116,8 @@ def test_bfcl_dataset_translates_schema_and_attaches_tools( ) -> None: """BFCLDataset should translate schemas to OpenAI tool format, set `messages` directly on SampleRequest, and attach tools/tool_choice via - request_overrides.""" + request_overrides. + """ paths = _write_fake_files(tmp_path) def fake_download(_repo, filename, **_kwargs): @@ -166,7 +168,8 @@ def test_bfcl_dataset_missing_category_raises_clear_error( hf_tokenizer: PreTrainedTokenizerBase, ) -> None: """A typo'd category should produce an actionable ValueError, not an - opaque huggingface_hub exception.""" + opaque huggingface_hub exception. + """ from huggingface_hub.errors import EntryNotFoundError args = _args_for_bfcl(categories=["simpl"]) # typo @@ -184,7 +187,8 @@ def test_bfcl_dataset_missing_category_raises_clear_error( @pytest.mark.benchmark def test_chat_backend_uses_messages_field_when_set() -> None: """When RequestFuncInput.chat_messages is set, the chat backend must use - it verbatim and skip default content construction from `prompt`.""" + it verbatim and skip default content construction from `prompt`. + """ import asyncio from vllm.benchmarks.lib.endpoint_request_func import ( @@ -244,7 +248,8 @@ def test_chat_backend_uses_messages_field_when_set() -> None: def test_bfcl_prompt_len_includes_tools(tmp_path: Path) -> None: """prompt_len must reflect tokens from both messages *and* tool schemas, so percentile buckets and input-distribution summaries aren't biased - low for tool-heavy traffic.""" + low for tool-heavy traffic. + """ paths = _write_fake_files(tmp_path) def fake_download(_repo, filename, **_kwargs): @@ -291,7 +296,8 @@ def test_bfcl_prompt_len_falls_back_when_tokenizer_rejects_tools( tmp_path: Path, ) -> None: """Older tokenizers don't accept tools=; fallback must still produce a - non-zero prompt_len without crashing.""" + non-zero prompt_len without crashing. + """ paths = _write_fake_files(tmp_path) def fake_download(_repo, filename, **_kwargs): diff --git a/tests/benchmarks/test_custom_dataset_chat_template_kwargs.py b/tests/benchmarks/test_custom_dataset_chat_template_kwargs.py index 149ea6625a..28ecb6e443 100644 --- a/tests/benchmarks/test_custom_dataset_chat_template_kwargs.py +++ b/tests/benchmarks/test_custom_dataset_chat_template_kwargs.py @@ -12,7 +12,8 @@ from vllm.benchmarks.datasets import get_samples class _RecordingTokenizer: """Minimal tokenizer stub that records the kwargs forwarded to apply_chat_template, so we can assert chat_template_kwargs propagation - without loading a real model/template.""" + without loading a real model/template. + """ def __init__(self) -> None: self.captured_kwargs: dict | None = None diff --git a/tests/benchmarks/test_endpoint_request_func_timing.py b/tests/benchmarks/test_endpoint_request_func_timing.py index 34394e3b01..57ac5824ab 100644 --- a/tests/benchmarks/test_endpoint_request_func_timing.py +++ b/tests/benchmarks/test_endpoint_request_func_timing.py @@ -228,7 +228,7 @@ def test_decode_span_identity_is_exact( api_url: str, chat: bool, ) -> None: - """latency - ttft must equal sum(itl) exactly. + """Latency - ttft must equal sum(itl) exactly. On a scripted clock this can only drift if a code path reads the clock an extra time or attributes a read to the wrong event. diff --git a/tests/benchmarks/test_sampling_params.py b/tests/benchmarks/test_sampling_params.py index 3bc34a84b3..13dcf8ad20 100644 --- a/tests/benchmarks/test_sampling_params.py +++ b/tests/benchmarks/test_sampling_params.py @@ -10,7 +10,8 @@ from vllm.tokenizers import TokenizerLike class _FakeTokenizer(TokenizerLike): """Minimal tokenizer implementing the TokenizerLike protocol - for testing get_sampling_params.""" + for testing get_sampling_params. + """ def __init__(self, vocab_size: int = 1000, num_special_tokens: int = 0) -> None: self._vocab_size = vocab_size @@ -171,7 +172,8 @@ class TestGetSamplingParams: def test_output_low_clamped_to_one(self): """Even with a high ratio that would push output_low to 0, - the function clamps it to 1.""" + the function clamps it to 1. + """ rng = np.random.default_rng(0) # output_len=1, ratio=0.99 → floor(1*0.01)=0, should clamp to 1 _, output_lens, _ = get_sampling_params( @@ -233,7 +235,8 @@ class TestGetSamplingParams: def test_input_len_zero_with_special_tokens(self): """input_len < num_special_tokens → real_input_len = 0, which is fine - (range [0, 0]).""" + (range [0, 0]). + """ rng = np.random.default_rng(0) input_lens, _, _ = get_sampling_params( rng, 5, 0.0, 5, 50, self._tok(num_special=10) diff --git a/tests/benchmarks/test_skip_tokenizer_init.py b/tests/benchmarks/test_skip_tokenizer_init.py index 741db2e0d6..1eb9a2536e 100644 --- a/tests/benchmarks/test_skip_tokenizer_init.py +++ b/tests/benchmarks/test_skip_tokenizer_init.py @@ -52,7 +52,8 @@ def _write_dataset(path: Path) -> None: def _args(dataset_path: str) -> argparse.Namespace: """Reproduce the argparse.Namespace that serve.py builds from the - failing command, including skip_tokenizer_init=True.""" + failing command, including skip_tokenizer_init=True. + """ return argparse.Namespace( # dataset dataset_name="custom", diff --git a/tests/ci_envs.py b/tests/ci_envs.py index f3a54f308c..171fac996a 100644 --- a/tests/ci_envs.py +++ b/tests/ci_envs.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -These envs only work for a small part of the tests, fix what you need! -""" +"""These envs only work for a small part of the tests, fix what you need!""" import os from collections.abc import Callable diff --git a/tests/compile/backend.py b/tests/compile/backend.py index 87f98946a8..167de6ad68 100644 --- a/tests/compile/backend.py +++ b/tests/compile/backend.py @@ -29,8 +29,7 @@ logger = init_logger("vllm.tests.compile.backend") class LazyInitPass(InductorPass): - """ - If there's a pass that we want to initialize lazily in a test, + """If there's a pass that we want to initialize lazily in a test, we can wrap it in LazyInitPass, which will initialize the pass when invoked and then immediately invoke it. """ @@ -45,8 +44,7 @@ class LazyInitPass(InductorPass): class TestBackend: - """ - This class provides a simple Inductor backend that can be used for testing. + """This class provides a simple Inductor backend that can be used for testing. It takes a list of custom passes and runs them after Inductor's passes. It also saves the graph before and after the custom passes for inspection. diff --git a/tests/compile/conftest.py b/tests/compile/conftest.py index 7d15b5c47e..972ec3fef5 100644 --- a/tests/compile/conftest.py +++ b/tests/compile/conftest.py @@ -11,8 +11,7 @@ from vllm.platforms.interface import DeviceCapability @pytest.fixture def mock_cuda_platform(): - """ - Fixture that returns a factory for creating mocked CUDA platforms. + """Fixture that returns a factory for creating mocked CUDA platforms. Usage: def test_something(mock_cuda_platform): @@ -51,8 +50,7 @@ def mock_cuda_platform(): @pytest.fixture def mock_xpu_platform(): - """ - Fixture that returns a factory for creating mocked XPU platforms. + """Fixture that returns a factory for creating mocked XPU platforms. Usage: def test_something(mock_xpu_platform): diff --git a/tests/compile/correctness_e2e/test_sequence_parallel.py b/tests/compile/correctness_e2e/test_sequence_parallel.py index eed5efb7b3..e5f8d1870e 100644 --- a/tests/compile/correctness_e2e/test_sequence_parallel.py +++ b/tests/compile/correctness_e2e/test_sequence_parallel.py @@ -1,11 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -WARNING: This test runs in both single-node (4 GPUs) and multi-node - (2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is - important to set the distributed backend to "mp" to avoid Ray scheduling - all workers in a node other than the head node, which can cause the test - to fail. +"""WARNING: This test runs in both single-node (4 GPUs) and multi-node +(2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is +important to set the distributed backend to "mp" to avoid Ray scheduling +all workers in a node other than the head node, which can cause the test +to fail. """ import json diff --git a/tests/compile/fullgraph/test_full_cudagraph.py b/tests/compile/fullgraph/test_full_cudagraph.py index c2af7be048..d91cd84816 100644 --- a/tests/compile/fullgraph/test_full_cudagraph.py +++ b/tests/compile/fullgraph/test_full_cudagraph.py @@ -17,8 +17,7 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum @contextlib.contextmanager def temporary_environ(env_vars): - """ - Temporarily set environment variables and restore them afterward. + """Temporarily set environment variables and restore them afterward. We have to do this vs monkeypatch because monkeypatch doesn't work with "module" scoped fixtures. """ @@ -125,8 +124,7 @@ def llm_pair(request): indirect=True, ) class TestFullCUDAGraph: - """ - Use a class such that an llm pair is constructed once for all + """Use a class such that an llm pair is constructed once for all batch_size/max_tokens combinations and released immediately after. Module-scope fixtures would stick around the whole time, @@ -149,11 +147,9 @@ class TestFullCUDAGraph: ], ) def test_full_cudagraph(self, batch_size, max_tokens, llm_pair: tuple[LLM, LLM]): - """ - Test various batch sizes and max_tokens to ensure that the + """Test various batch sizes and max_tokens to ensure that the full cudagraph compilation works for padded cases too. """ - full_cudagraph_llm, piecewise_llm = llm_pair prompts = ["the quick brown fox"] * batch_size diff --git a/tests/compile/fullgraph/test_multiple_graphs.py b/tests/compile/fullgraph/test_multiple_graphs.py index 6d3788af9d..e265663334 100644 --- a/tests/compile/fullgraph/test_multiple_graphs.py +++ b/tests/compile/fullgraph/test_multiple_graphs.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test (piecewise) compilation with a simple model where multiple submodules +"""Test (piecewise) compilation with a simple model where multiple submodules are compiled and graph captured separately. """ diff --git a/tests/compile/fullgraph/test_simple.py b/tests/compile/fullgraph/test_simple.py index f1ea0e414d..4d10a92a07 100644 --- a/tests/compile/fullgraph/test_simple.py +++ b/tests/compile/fullgraph/test_simple.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test the piecewise compilation with a simple model so that we +"""Test the piecewise compilation with a simple model so that we can exactly calculate the expected output and side effects. """ @@ -52,8 +51,7 @@ class SillyModel(nn.Module): self.intermediate_unbacked = intermediate_unbacked def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Overall effect: + """Overall effect: x = 3 * x + 19 global_counter += 2 """ diff --git a/tests/compile/fullgraph/test_toy_llama.py b/tests/compile/fullgraph/test_toy_llama.py index 69c758702e..7823c362f6 100644 --- a/tests/compile/fullgraph/test_toy_llama.py +++ b/tests/compile/fullgraph/test_toy_llama.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test the piecewise compilation with a simple model, comparing the output +"""Test the piecewise compilation with a simple model, comparing the output with and without the piecewise compilation. This is a tractable model, the weights and computation are specially designed diff --git a/tests/compile/passes/distributed/test_async_tp.py b/tests/compile/passes/distributed/test_async_tp.py index 7b547c9720..ff0b048962 100644 --- a/tests/compile/passes/distributed/test_async_tp.py +++ b/tests/compile/passes/distributed/test_async_tp.py @@ -56,10 +56,7 @@ class TestMMRSModel(torch.nn.Module): torch.nn.init.normal_(self.gate_proj, std=0.02) def forward(self, hidden_states): - """ - Forward pass implementing the mm + reduce scatter in the FX graph - - """ + """Forward pass implementing the mm + reduce scatter in the FX graph.""" # Reshape input view = hidden_states.reshape(-1, self.hidden_size) @@ -88,9 +85,7 @@ class TestAGMMModel(torch.nn.Module): torch.nn.init.normal_(self.weight, std=0.02) def forward(self, hidden_states): - """ - Forward pass implementing the mm + all gather in the FX graph - """ + """Forward pass implementing the mm + all gather in the FX graph.""" # Reshape input view = hidden_states.reshape(-1, self.hidden_size) all_gather = tensor_model_parallel_all_gather(view, dim=0) @@ -122,10 +117,7 @@ class _BaseScaledMMModel(torch.nn.Module): class TestScaledMMRSModel(_BaseScaledMMModel): def forward(self, input: torch.Tensor): - """ - Forward pass implementing the scaled_mm + reduce scatter in the FX graph - - """ + """Forward pass implementing the scaled_mm + reduce scatter in the FX graph.""" fp8_input = input.to(FP8_DTYPE) scale_a = torch.ones(input.shape[0], 1, dtype=torch.float32) scaled_mm = torch._scaled_mm( @@ -147,9 +139,7 @@ class TestScaledMMRSModel(_BaseScaledMMModel): class TestAGScaledMMModel(_BaseScaledMMModel): def forward(self, input: torch.Tensor): - """ - Forward pass implementing the all gather + scaled_mm in the FX graph - """ + """Forward pass implementing the all gather + scaled_mm in the FX graph.""" # Reshape input fp8_input = input.to(FP8_DTYPE) all_gather = tensor_model_parallel_all_gather(fp8_input, dim=0) @@ -173,8 +163,7 @@ class TestAGScaledMMModel(_BaseScaledMMModel): class TestCutlassScaledMMRSModel(_BaseScaledMMModel): def forward(self, input: torch.Tensor): - """ - Forward pass implementing the cutlass_scaled_mm + reduce scatter + """Forward pass implementing the cutlass_scaled_mm + reduce scatter in the FX graph """ @@ -200,8 +189,7 @@ class TestCutlassScaledMMRSModel(_BaseScaledMMModel): class TestAGCutlassScaledMMModel(_BaseScaledMMModel): def forward(self, input: torch.Tensor): - """ - Forward pass implementing the all gather + cutlass_scaled_mm + """Forward pass implementing the all gather + cutlass_scaled_mm in the FX graph """ # Reshape input diff --git a/tests/compile/passes/ir/test_clone_cleanup.py b/tests/compile/passes/ir/test_clone_cleanup.py index b6626a5cb7..8a51e6d539 100644 --- a/tests/compile/passes/ir/test_clone_cleanup.py +++ b/tests/compile/passes/ir/test_clone_cleanup.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Comprehensive tests for UnsafeCloneEliminationPass. +"""Comprehensive tests for UnsafeCloneEliminationPass. This test suite exercises all possible valid FX graph patterns involving clones: 1. Clone with no users (dead code) diff --git a/tests/compile/passes/ir/test_inplace_functionalization.py b/tests/compile/passes/ir/test_inplace_functionalization.py index 1e8d566216..773ea93ce6 100644 --- a/tests/compile/passes/ir/test_inplace_functionalization.py +++ b/tests/compile/passes/ir/test_inplace_functionalization.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for IR inplace functionalization pass integration. +"""Tests for IR inplace functionalization pass integration. This test suite verifies that the inplace functionalization pass, lowering pass, and clone cleanup pass work together correctly with donated buffer tracking. @@ -108,8 +107,7 @@ class MixedModel(nn.Module): class ModelWithTritonAfterMaybeInplace(nn.Module): - """ - Model using maybe_inplace followed by a Triton kernel. + """Model using maybe_inplace followed by a Triton kernel. Test clone elimination can handle Triton in the graph """ @@ -410,8 +408,7 @@ def with_dyn_arg(fn: Callable, arg_index: int, dim_index: int): reason="Only test on cuda and rocm platform", ) def test_piecewise_compilation_with_donated_buffers(monkeypatch, fresh_vllm_cache): - """ - Test piecewise compilation with donated buffers across graph splits. + """Test piecewise compilation with donated buffers across graph splits. Utilizes a custom splitting op. Uses fresh cache to avoid compilation caching. """ torch.set_default_device(current_platform.device_type) diff --git a/tests/compile/passes/test_double_aiter_rms_quant_fusion.py b/tests/compile/passes/test_double_aiter_rms_quant_fusion.py index 6a620d11a4..f05f821852 100644 --- a/tests/compile/passes/test_double_aiter_rms_quant_fusion.py +++ b/tests/compile/passes/test_double_aiter_rms_quant_fusion.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for the DoubleQuant fan-out variants registered by +"""Unit tests for the DoubleQuant fan-out variants registered by ``RocmAiterRMSNormQuantFusionPass``. Both variants target a 1-to-2 fan-out where one ``rms_norm`` output feeds @@ -90,8 +89,7 @@ def test_double_aiter_rms_fp8_group_quant_fusion( model_cls: type[torch.nn.Module], monkeypatch: pytest.MonkeyPatch, ) -> None: - """ - Both fan-out shapes (with and without an intermediate view) must fuse + """Both fan-out shapes (with and without an intermediate view) must fuse into ``rocm_aiter_rmsnorm_fp8_group_quant``: the no-view shape via ``DoubleAiterRMSFp8GroupQuantPattern`` and the viewed shape via the new ``DoubleAiterRMSFp8GroupQuantViewPattern`` sibling. diff --git a/tests/compile/passes/test_fuse_mla_dual_rms_norm.py b/tests/compile/passes/test_fuse_mla_dual_rms_norm.py index 6f20d4e158..0f631b9472 100644 --- a/tests/compile/passes/test_fuse_mla_dual_rms_norm.py +++ b/tests/compile/passes/test_fuse_mla_dual_rms_norm.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit test for the MLADualRMSNormFusionPass. +"""Unit test for the MLADualRMSNormFusionPass. The pass fuses paired q/kv RMS norms in MLA attention into a single fused_mla_dual_rms_norm op backed by AITER's fused_qk_rmsnorm kernel. @@ -38,13 +37,12 @@ FP8_DTYPE = current_platform.fp8_dtype() class MLADualRMSNormTestModel(torch.nn.Module): - """ - Minimal model reproducing the MLA dual RMS norm pattern: - linear -> split([q_dim, kv_dim]) - +-- q_c (getitem 0) -> rms_norm(q_w, eps) -> linear - +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim]) - +-- kv_c (getitem 0) -> rms_norm(kv_w, eps) - +-- k_pe + """Minimal model reproducing the MLA dual RMS norm pattern: + linear -> split([q_dim, kv_dim]) + +-- q_c (getitem 0) -> rms_norm(q_w, eps) -> linear + +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim]) + +-- kv_c (getitem 0) -> rms_norm(kv_w, eps) + +-- k_pe """ def __init__( @@ -155,13 +153,12 @@ def test_fuse_mla_dual_rms_norm( class MLADualRMSNormFp8PerTokenTestModel(torch.nn.Module): - """ - Minimal model reproducing the FP8 MLA attention path with *per-token* quant: - linear -> split([q_dim, kv_dim]) - +-- q_c (getitem 0) -> rocm_aiter_rmsnorm_fused_dynamic_quant -> dequant - +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim]) - +-- kv_c (getitem 0) -> rms_norm (bf16) - +-- k_pe + """Minimal model reproducing the FP8 MLA attention path with *per-token* quant: + linear -> split([q_dim, kv_dim]) + +-- q_c (getitem 0) -> rocm_aiter_rmsnorm_fused_dynamic_quant -> dequant + +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim]) + +-- kv_c (getitem 0) -> rms_norm (bf16) + +-- k_pe """ def __init__( diff --git a/tests/compile/passes/test_fusion.py b/tests/compile/passes/test_fusion.py index 843052f4af..0945bd622a 100644 --- a/tests/compile/passes/test_fusion.py +++ b/tests/compile/passes/test_fusion.py @@ -732,7 +732,8 @@ def test_aiter_fusion_rmsnorm_gated_quant_no_gdn_layers( monkeypatch: pytest.MonkeyPatch, ): """Verify that without GDN layers in static_forward_context, - the gated pattern is not registered and no matches occur.""" + the gated pattern is not registered and no matches occur. + """ group_shape = GroupShape(1, 128) vllm_config = VllmConfig( model_config=ModelConfig(dtype=dtype), diff --git a/tests/compile/passes/test_fusion_attn.py b/tests/compile/passes/test_fusion_attn.py index 4df02cab95..37f0d4126c 100644 --- a/tests/compile/passes/test_fusion_attn.py +++ b/tests/compile/passes/test_fusion_attn.py @@ -107,7 +107,6 @@ class AttentionQuantPatternModel(torch.nn.Module): def build_attn_metadata(self, batch_size: int) -> AttentionMetadata: """Initialize attention metadata.""" - # TODO (Rohan138) reuse utils from vllm/v1/worker/gpu/attn_utils.py # Create common attn metadata @@ -304,7 +303,7 @@ def test_attention_quant_pattern( dist_init, disable_vllm_compile_cache, ): - """Test AttentionStaticQuantPattern fusion pass""" + """Test AttentionStaticQuantPattern fusion pass.""" if backend == AttentionBackendEnum.FLASHINFER and ( not current_platform.is_device_capability((10, 0)) or not has_flashinfer() ): diff --git a/tests/compile/passes/test_mla_attn_quant_fusion.py b/tests/compile/passes/test_mla_attn_quant_fusion.py index 11c455bfa2..f1de3dde0c 100644 --- a/tests/compile/passes/test_mla_attn_quant_fusion.py +++ b/tests/compile/passes/test_mla_attn_quant_fusion.py @@ -141,7 +141,6 @@ class MLAAttentionQuantPatternModel(torch.nn.Module): the full output[:num_actual_toks] buffer after both forward_mha and forward_mqa have written their results. """ - batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size) common_attn_metadata = create_common_attn_metadata( batch_spec, self.block_size, self.device, arange_block_indices=True @@ -408,7 +407,7 @@ def test_mla_attention_quant_pattern( dist_init, disable_vllm_compile_cache, ): - """Test MLA AttentionQuantPattern fusion pass""" + """Test MLA AttentionQuantPattern fusion pass.""" if ( model_class is TestMLAAttentionNvfp4QuantPatternModel and not is_nvfp4_supported() diff --git a/tests/compile/passes/test_silu_mul_quant_manual_fusion.py b/tests/compile/passes/test_silu_mul_quant_manual_fusion.py index ee94960844..10dc401a26 100644 --- a/tests/compile/passes/test_silu_mul_quant_manual_fusion.py +++ b/tests/compile/passes/test_silu_mul_quant_manual_fusion.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for manual fusion via maybe_fused_act_quant. +"""Tests for manual fusion via maybe_fused_act_quant. Tests all fusion paths in _FUSED_ACT_QUANT: - kFp8StaticTensorSym: all platforms diff --git a/tests/compile/passes/test_split_coalescing.py b/tests/compile/passes/test_split_coalescing.py index ab7a0be1a2..52e531e127 100644 --- a/tests/compile/passes/test_split_coalescing.py +++ b/tests/compile/passes/test_split_coalescing.py @@ -15,7 +15,8 @@ DEVICE_TYPE = current_platform.device_type class SplitCoalescingModel(torch.nn.Module): """Model with 3 separate split_with_sizes calls on the same input, - simulating the B200+FP8 graph where CSE fails to merge them.""" + simulating the B200+FP8 graph where CSE fails to merge them. + """ def __init__(self, q_size: int, kv_size: int) -> None: super().__init__() diff --git a/tests/compile/silly_attention.py b/tests/compile/silly_attention.py index 29c02f6e6a..0d33135129 100644 --- a/tests/compile/silly_attention.py +++ b/tests/compile/silly_attention.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Shared PyTorch custom silly attention for compilation tests. +"""Shared PyTorch custom silly attention for compilation tests. Centralizes custom operation definitions to avoid duplicate registrations. """ @@ -21,12 +20,12 @@ _global_counter = 0 def get_global_counter(): - """Get the current global counter value""" + """Get the current global counter value.""" return _global_counter def reset_global_counter(): - """Reset the global counter to 0""" + """Reset the global counter to 0.""" global _global_counter _global_counter = 0 @@ -34,8 +33,7 @@ def reset_global_counter(): def silly_attention( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, out: torch.Tensor ) -> None: - """ - Unified attention implementation that depends on + """Unified attention implementation that depends on all inputs and affects the output. Always increments a global counter that tests can use or ignore. """ @@ -51,7 +49,7 @@ def silly_attention( def silly_attention_fake( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, out: torch.Tensor ) -> None: - """Fake implementation for testing""" + """Fake implementation for testing.""" return diff --git a/tests/compile/test_aot_compile.py b/tests/compile/test_aot_compile.py index 90c132bee1..e5c11b8acc 100644 --- a/tests/compile/test_aot_compile.py +++ b/tests/compile/test_aot_compile.py @@ -268,8 +268,7 @@ def test_save_and_load_slice(monkeypatch: pytest.MonkeyPatch): def test_cache_load_returns_tuple_consistency_tuple_output( monkeypatch: pytest.MonkeyPatch, ): - """ - Test that cache loading correctly handles models that return tuples. + """Test that cache loading correctly handles models that return tuples. This verifies that when a model returns a tuple of tensors, the output type is preserved as a tuple between fresh compilation and cache load. @@ -335,8 +334,7 @@ def test_cache_load_returns_tuple_consistency_tuple_output( @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") def test_shape_env(monkeypatch: pytest.MonkeyPatch): - """ - Test that the shape environment is correctly serialized and preserved + """Test that the shape environment is correctly serialized and preserved when loading from cache. """ with monkeypatch.context() as m: @@ -374,8 +372,7 @@ def test_shape_env(monkeypatch: pytest.MonkeyPatch): def test_partition_wrapper_applied_on_aot_load( monkeypatch: pytest.MonkeyPatch, vllm_tmp_cache: Path, mocker ): - """ - Test that partition wrappers are applied when loading AOT cached functions. + """Test that partition wrappers are applied when loading AOT cached functions. This test verifies the fix for GitHub issue #31439 where AOT compile caused 2x latency regression when use_inductor_graph_partition=True. @@ -498,14 +495,12 @@ def test_standalone_compile_correctness(): @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") @create_new_process_for_each_test("spawn") def test_gpt2_cache_hit(monkeypatch: pytest.MonkeyPatch): - """ - Test that compiling gpt2 twice results in a cache hit. + """Test that compiling gpt2 twice results in a cache hit. Counter values are read from the EngineCore subprocess via ``LLM.collective_rpc`` so the test works under default V1 multiprocessing (no shared memory between test and engine). """ - from vllm import LLM def _snap(self): @@ -641,7 +636,7 @@ class TestStandaloneCompiledArtifacts: @patch("torch._inductor.standalone_compile.AOTCompiledArtifact.deserialize") def test_load_all_success(self, mock_deserialize): - """Test successful loading of all artifacts""" + """Test successful loading of all artifacts.""" cache = StandaloneCompiledArtifacts() mock_artifact1 = Mock() @@ -658,7 +653,7 @@ class TestStandaloneCompiledArtifacts: @patch("torch._inductor.standalone_compile.AOTCompiledArtifact.deserialize") def test_load_all_already_loaded(self, mock_deserialize): - """Test that load_all skips if already loaded""" + """Test that load_all skips if already loaded.""" cache = StandaloneCompiledArtifacts() mock_artifact = Mock() @@ -671,7 +666,7 @@ class TestStandaloneCompiledArtifacts: @patch("torch._inductor.standalone_compile.AOTCompiledArtifact.deserialize") def test_get_loaded_artifact(self, mock_deserialize): - """Test retrieving loaded artifacts""" + """Test retrieving loaded artifacts.""" cache = StandaloneCompiledArtifacts() mock_artifact = Mock() diff --git a/tests/compile/test_codegen.py b/tests/compile/test_codegen.py index 21db287a24..a0086cb702 100644 --- a/tests/compile/test_codegen.py +++ b/tests/compile/test_codegen.py @@ -77,7 +77,8 @@ def test_non_primitive_kwargs_lifted_to_consts( into the namespace, so ``device`` is unbound. The fix collects such objects into ``__vllm_consts__`` and references them by index. The unqualified ``device(type=...)`` form must never appear in the - generated source.""" + generated source. + """ split_gm = _trace_and_split(model_fn, (x,), split_ops) code, submod_names, consts = generate_execution_code(split_gm) @@ -98,7 +99,8 @@ def test_dtype_singleton_deduped(x: torch.Tensor) -> None: """``torch.float16`` is a process-wide singleton, so two ops referring to it in the traced graph share a single consts slot via ``id()``-based dedup. Distinct expressions (``x.to(...)`` vs ``(x*2).to(...)``) ensure - the tracer can't CSE the two ops into a single node.""" + the tracer can't CSE the two ops into a single node. + """ def model_fn(x: torch.Tensor) -> torch.Tensor: return x.to(dtype=torch.float16) + (x * 2).to(dtype=torch.float16) @@ -147,7 +149,8 @@ def test_distinct_dtypes_get_distinct_slots(x: torch.Tensor) -> None: def test_consts_ordering_deterministic(x: torch.Tensor) -> None: """Two independent traces of the same model must produce equal consts lists *in the same order*. Cache artifacts identify const slots by - index, so a non-deterministic order would invalidate cached code.""" + index, so a non-deterministic order would invalidate cached code. + """ def model_fn(x: torch.Tensor) -> torch.Tensor: # Multiple distinct non-primitives encountered in a fixed graph order. @@ -185,7 +188,8 @@ def test_consts_shared_across_split_submods(x: torch.Tensor) -> None: The function below splits into three inlined submods, two of which independently reference ``torch.float16``. The shared ``const_index`` threaded through recursive ``generate_execution_code_with_name`` calls - must collapse the dtype to a single slot used from both submods.""" + must collapse the dtype to a single slot used from both submods. + """ def model_fn(x: torch.Tensor) -> torch.Tensor: a = x.to(dtype=torch.float16) # submod_0: _to_copy(fp16) @@ -230,7 +234,8 @@ def test_non_graphmodule_submod_uses_indexed_callable(x: torch.Tensor) -> None: """When a child of split_gm is *not* a ``torch.fx.GraphModule`` — as happens in production once ``PiecewiseBackend`` replaces submods — codegen emits ``__vllm_submods__[idx](...)`` instead of inlining, and - the runtime callable is bound from ``submod_callables``.""" + the runtime callable is bound from ``submod_callables``. + """ def model_fn(x: torch.Tensor) -> torch.Tensor: return x.relu().sigmoid() @@ -280,7 +285,8 @@ def test_non_graphmodule_submod_uses_indexed_callable(x: torch.Tensor) -> None: def test_getitem_in_stitching_graph(x: torch.Tensor) -> None: """``operator.getitem`` on submod tuple returns is the ``call_function`` special case at codegen.py — emitted as ``name = source[index]`` - rather than a function call.""" + rather than a function call. + """ def model_fn(x: torch.Tensor) -> torch.Tensor: return x.relu().sigmoid() @@ -298,7 +304,8 @@ def test_getitem_in_stitching_graph(x: torch.Tensor) -> None: def test_del_emitted_for_intermediate_values(x: torch.Tensor) -> None: """The codegen schedules ``del`` after a value's last use to free memory early. Multi-submod splits naturally have intermediates whose - last use is not the output node.""" + last use is not the output node. + """ def model_fn(x: torch.Tensor) -> torch.Tensor: return x.relu().sigmoid().tanh() @@ -318,7 +325,8 @@ def test_with_submod_false_rejects_call_module() -> None: """``generate_execution_code_with_name(with_submod=False)`` is the recursive entry for inlining a GraphModule into its parent. It must refuse a graph that itself contains ``call_module`` nodes — the parent - is responsible for handling those.""" + is responsible for handling those. + """ g = fx.Graph() x_node = g.placeholder("x") root = torch.nn.Module() @@ -335,7 +343,8 @@ def test_node_ref_recurses_through_containers() -> None: """``_node_ref`` is the recursive walker that lifts non-primitives nested inside list/tuple/dict args. Real aten ops rarely produce such structures, but the path is needed for DTensor placement lists and - other future cases — unit-test the walker directly.""" + other future cases — unit-test the walker directly. + """ consts: list = [] const_index: dict[int, int] = {} cpu = torch.device("cpu") @@ -359,7 +368,8 @@ def test_node_ref_recurses_through_containers() -> None: def test_legacy_code_without_consts() -> None: """``compile_execution_fn(consts=None)`` must still load code that has no ``__vllm_consts__`` reference, so older serialized cache artifacts - keep working.""" + keep working. + """ # Pre-consts codegen: no __vllm_consts__ reference, only torch/operator. legacy_code = ( "import torch\n" diff --git a/tests/compile/test_compile_ranges.py b/tests/compile/test_compile_ranges.py index 4dfea42a6b..b3e2ec4731 100644 --- a/tests/compile/test_compile_ranges.py +++ b/tests/compile/test_compile_ranges.py @@ -129,7 +129,8 @@ def test_compile_config_get_compile_ranges(): class PostGradStaticShapeChecker(InductorPass): """Asserts that compile_sizes entries produce graphs with fully concrete - (non-symbolic) shapes, and compile_ranges entries have symbolic shapes.""" + (non-symbolic) shapes, and compile_ranges entries have symbolic shapes. + """ def __init__(self): self.num_static_calls = 0 @@ -170,7 +171,8 @@ class PostGradStaticShapeChecker(InductorPass): def test_compile_sizes_produce_static_shapes(disable_vllm_compile_cache): """Verify that compile_sizes entries are compiled with fully concrete - shapes (no SymInts), while compile_ranges entries retain dynamic shapes.""" + shapes (no SymInts), while compile_ranges entries retain dynamic shapes. + """ checker = PostGradStaticShapeChecker() torch.set_default_device("cuda") vllm_config = VllmConfig( diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index 5c9132ccbf..fc32956218 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -1229,7 +1229,8 @@ def test_compile_sizes_padding_validation(): def test_inductor_asserts_default_disabled(monkeypatch): """Test that inductor runtime asserts are disabled by default - (INFO logging level) on torch < 2.12.""" + (INFO logging level) on torch < 2.12. + """ monkeypatch.setenv("VLLM_LOGGING_LEVEL", "INFO") import importlib @@ -1247,7 +1248,8 @@ def test_inductor_asserts_default_disabled(monkeypatch): def test_inductor_asserts_enabled_in_debug(monkeypatch): """Test that VLLM_LOGGING_LEVEL=DEBUG enables inductor runtime asserts - on torch < 2.12.""" + on torch < 2.12. + """ monkeypatch.setenv("VLLM_LOGGING_LEVEL", "DEBUG") import importlib @@ -1283,7 +1285,8 @@ def test_get_inductor_factors_includes_configs(): def test_inductor_asserts_user_override(monkeypatch): """Test that explicit inductor_compile_config overrides the - debug-logging default.""" + debug-logging default. + """ monkeypatch.setenv("VLLM_LOGGING_LEVEL", "INFO") import importlib diff --git a/tests/compile/test_dynamic_shapes_compilation.py b/tests/compile/test_dynamic_shapes_compilation.py index 71f66fce40..ea2ccafbe7 100644 --- a/tests/compile/test_dynamic_shapes_compilation.py +++ b/tests/compile/test_dynamic_shapes_compilation.py @@ -22,7 +22,7 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer def get_test_models(): - """Get list of models to test based on PyTorch version""" + """Get list of models to test based on PyTorch version.""" models = [ "Qwen/Qwen3-0.6B", "openai-community/gpt2", @@ -172,7 +172,6 @@ def test_model_specialization_with_evaluate_guards( """Test that evaluate_guards correctly detects shape specialization violations. """ - if ( use_aot_compile and dynamic_shapes_type == DynamicShapesType.BACKED diff --git a/tests/compile/test_graph_partition.py b/tests/compile/test_graph_partition.py index bb2a6f2aee..f5aaffe00e 100644 --- a/tests/compile/test_graph_partition.py +++ b/tests/compile/test_graph_partition.py @@ -24,8 +24,7 @@ DEVICE_TYPE = current_platform.device_type def test_getitem_moved_to_producer_subgraph(): - """ - Test that getitem operations are moved to the same subgraph as their input, + """Test that getitem operations are moved to the same subgraph as their input, preventing tuple inputs to submodules. """ @@ -79,8 +78,7 @@ def test_getitem_moved_to_producer_subgraph(): def test_no_tuple_inputs_with_multiple_consumers(): - """ - Test that when a tuple is consumed by multiple split operations, + """Test that when a tuple is consumed by multiple split operations, getitem operations are properly moved to avoid tuple inputs. """ @@ -138,13 +136,10 @@ def test_no_tuple_inputs_with_multiple_consumers(): def test_consecutive_ops_in_split(): - """ - Test that consecutive splitting operations are grouped into the same subgraph - """ + """Test that consecutive splitting operations are grouped into the same subgraph.""" def model_fn(x: torch.Tensor) -> torch.Tensor: - """ - Define a simple model where consecutive operations create opportunities + """Define a simple model where consecutive operations create opportunities for splitting subgraphs. """ # Apply silly attention followed by consecutive operations @@ -211,8 +206,7 @@ def _subgraphs_with_empty_nodes(split_items, *, is_splitting_graph): def test_empty_only_partition_stays_separate_after_splitting_predecessor(): - """ - Empty-only subgraphs should not be merged when the only predecessor is + """Empty-only subgraphs should not be merged when the only predecessor is a splitting-op subgraph. """ @@ -248,8 +242,7 @@ def test_empty_only_partition_stays_separate_after_splitting_predecessor(): def test_empty_only_partition_is_merged(): - """ - Empty-only subgraphs should still be merged when a non-splitting predecessor + """Empty-only subgraphs should still be merged when a non-splitting predecessor exists. The merged empty node must remain outside splitting-op subgraphs. """ @@ -294,8 +287,7 @@ def test_empty_only_partition_is_merged(): def test_builtin_empty_only_partition_is_merged(): - """ - In Dynamo graphs, torch.empty/empty_like may appear as builtin call targets + """In Dynamo graphs, torch.empty/empty_like may appear as builtin call targets (not aten OpOverload). Ensure empty-only partitions are still merged. """ @@ -340,8 +332,7 @@ def test_builtin_empty_only_partition_is_merged(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_sym_size_whole_shape_boundary(): - """ - Test that using x.size() (whole shape) across a split boundary can be + """Test that using x.size() (whole shape) across a split boundary can be compiled by standalone_compile. The dynamo graph looks like: @@ -400,8 +391,7 @@ def test_sym_size_whole_shape_boundary(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_symint_crosses_split_boundary(): - """ - Test that SymInt placeholders from torch.compile + mark_dynamic + """Test that SymInt placeholders from torch.compile + mark_dynamic cross split boundaries safely via split_module's natural threading. SymInt values are threaded through subgraphs by split_module and @@ -459,8 +449,7 @@ def test_symint_crosses_split_boundary(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_shape_boundary_standalone_compile(): - """ - Repro for the original production bug: + """Repro for the original production bug: AssertionError: out_spec mismatch TreeSpec(tuple, None, [*, *, TreeSpec(Size, None, [*, *]), *]) @@ -529,8 +518,7 @@ def test_shape_boundary_standalone_compile(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_size_used_in_multiple_consumer_subgraphs(): - """ - Validates that x.size() (whole shape) used by multiple downstream subgraphs + """Validates that x.size() (whole shape) used by multiple downstream subgraphs does not cause torch.Size to cross split boundaries. Model: @@ -584,8 +572,7 @@ def test_size_used_in_multiple_consumer_subgraphs(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_sym_size_metadata_propagated(): - """ - Validates that new sym_size.int nodes created by the pre-pass have + """Validates that new sym_size.int nodes created by the pre-pass have example_value metadata set. Without it, placeholder metadata in consumer subgraphs would be None, breaking any code that dynamically builds example inputs from metadata (e.g. standalone_compile per-submodule). @@ -634,8 +621,7 @@ def test_sym_size_metadata_propagated(): def test_decompose_size_with_getitem_user(): - """ - Regression test: _decompose_size_nodes must handle getitem users of size() + """Regression test: _decompose_size_nodes must handle getitem users of size() correctly. When a graph contains x.shape[i], it can appear as: @@ -704,8 +690,7 @@ def test_decompose_size_with_getitem_user(): def test_decompose_size_leaves_scalar_size_with_dim(): - """ - Regression test: _decompose_size_nodes must leave x.size(dim) alone. + """Regression test: _decompose_size_nodes must leave x.size(dim) alone. x.size() returns a torch.Size tuple that can't cross split boundaries and must be decomposed. x.size(dim), however, already returns a scalar diff --git a/tests/compile/test_inductor_fallback_allow_list_patch.py b/tests/compile/test_inductor_fallback_allow_list_patch.py index 29fe9962e3..c5705da211 100644 --- a/tests/compile/test_inductor_fallback_allow_list_patch.py +++ b/tests/compile/test_inductor_fallback_allow_list_patch.py @@ -236,7 +236,8 @@ class TestInductorFallbackFastPath: def test_inner_set_membership_still_works_for_standard_ops(self): """The patch must not break Inductor's existing fallback decisions - for non-vllm ops such as ``torchvision::roi_align``.""" + for non-vllm ops such as ``torchvision::roi_align``. + """ import torch._inductor.lowering as _lowering allow_list = _lowering.FALLBACK_ALLOW_LIST diff --git a/tests/compile/test_structured_logging.py b/tests/compile/test_structured_logging.py index 10b0ed139c..98f092bfce 100644 --- a/tests/compile/test_structured_logging.py +++ b/tests/compile/test_structured_logging.py @@ -60,6 +60,7 @@ class TraceStructuredCapture: Args: event_type: The event type to filter by (e.g., "artifact", "graph_dump") name_pattern: Regex pattern to match against the artifact name + """ regex = re.compile(name_pattern) return [ diff --git a/tests/config/test_config_utils.py b/tests/config/test_config_utils.py index 590e97c1eb..1d0f4d915b 100644 --- a/tests/config/test_config_utils.py +++ b/tests/config/test_config_utils.py @@ -40,7 +40,7 @@ class DummyLogprobsMode(Enum): def test_hash_factors_deterministic(): - """Test that hash_factors produces consistent SHA-256 hashes""" + """Test that hash_factors produces consistent SHA-256 hashes.""" factors = {"a": 1, "b": "test"} hash1 = hash_factors(factors) hash2 = hash_factors(factors) diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index b68820c0d5..f71d4875d7 100644 --- a/tests/config/test_model_arch_config.py +++ b/tests/config/test_model_arch_config.py @@ -372,7 +372,8 @@ class _HeterogeneousConfig(PretrainedConfig): """A heterogeneous config with no convertor of its own. Mirrors the parts vLLM uses: per-layer configs are shallow copies with the - varying attributes applied and heterogeneity stripped, so they do not recurse.""" + varying attributes applied and heterogeneity stripped, so they do not recurse. + """ is_heterogeneous = True @@ -471,7 +472,8 @@ def test_deepseek_v4_convertor_splits_vision_architecture(): """The vision checkpoint shares model_type/architectures with the text-only DeepSeek-V4; the convertor routes it to the VL wrapper class by rewriting hf_config.architectures (model-class resolution reads the - raw attribute).""" + raw attribute). + """ from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config from vllm.transformers_utils.model_arch_config_convertor import ( DeepseekV4ModelArchConfigConvertor, diff --git a/tests/config/test_mp_reducer.py b/tests/config/test_mp_reducer.py index 56dc542f1c..56e4e47474 100644 --- a/tests/config/test_mp_reducer.py +++ b/tests/config/test_mp_reducer.py @@ -9,12 +9,10 @@ from vllm.v1.engine.async_llm import AsyncLLM def test_mp_reducer(): - """ - Test that _reduce_config reducer is registered when AsyncLLM is instantiated + """Test that _reduce_config reducer is registered when AsyncLLM is instantiated without transformers_modules. This is a regression test for https://github.com/vllm-project/vllm/pull/18640. """ - # Ensure transformers_modules is not in sys.modules if "transformers_modules" in sys.modules: del sys.modules["transformers_modules"] diff --git a/tests/config/test_multimodal_config.py b/tests/config/test_multimodal_config.py index 84cf092c21..abf834608a 100644 --- a/tests/config/test_multimodal_config.py +++ b/tests/config/test_multimodal_config.py @@ -42,7 +42,8 @@ def test_mm_encoder_attn_backend_hash_updates(): def test_language_model_only_does_not_affect_mm_hash(): """language_model_only does not affect the ViT computation graph, - so it should not change the multimodal config hash.""" + so it should not change the multimodal config hash. + """ base_hash = MultiModalConfig().compute_hash() lm_only_hash = MultiModalConfig(language_model_only=True).compute_hash() assert base_hash == lm_only_hash @@ -50,7 +51,8 @@ def test_language_model_only_does_not_affect_mm_hash(): def test_language_model_only_affects_model_hash(): """language_model_only affects the LM computation graph, - so it should change the model config hash.""" + so it should change the model config hash. + """ model = "llava-hf/llava-1.5-7b-hf" base_hash = ModelConfig(model).compute_hash() lm_only_hash = ModelConfig(model, language_model_only=True).compute_hash() @@ -386,7 +388,8 @@ def _resolve_mm_video_decode_device( torchcodec_available: bool = True, ) -> dict: """Run processor-device then video-decode-device resolution and report - the resulting video media IO kwargs.""" + the resulting video media IO kwargs. + """ mm_config = MultiModalConfig( mm_processor_kwargs={} if device is None else {"device": device}, mm_tensor_ipc=mm_tensor_ipc, # type: ignore[arg-type] @@ -416,7 +419,8 @@ def _resolve_mm_video_decode_device( def test_auto_video_decode_uses_nvdec_on_encoder_instance(): """The processor runs on the accelerator there, so decoded frames should - stay on-device too.""" + stay on-device too. + """ assert _resolve_mm_video_decode_device(ec_role="ec_producer") == { "backend": "torchcodec", "device": "cuda", @@ -428,7 +432,8 @@ def test_auto_video_decode_stays_on_cpu_off_encoder_instance( ec_role: ECRole | None, ): """The processor stays on CPU off encode-only instances, so video - decoding should too.""" + decoding should too. + """ assert _resolve_mm_video_decode_device(ec_role=ec_role) == {} diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py index 8505172da3..10f8c994f4 100644 --- a/tests/config/test_speculative_draft_hf_overrides.py +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -34,7 +34,8 @@ def _make_hf_config(**kwargs) -> PretrainedConfig: @pytest.mark.cpu_test def test_dict_overrides_are_not_forwarded_to_draft(): """Dict overrides are target-specific key patches; the draft must get - only the architecture-mapping override.""" + only the architecture-mapping override. + """ composed = SpeculativeConfig.compose_draft_hf_overrides( {"max_position_embeddings": 1234} ) @@ -50,7 +51,8 @@ def test_none_overrides_fall_back_to_arch_mapping(): @pytest.mark.cpu_test def test_callable_overrides_reach_the_draft_config(): """A callable override (config-to-config transform) composes with the - architecture-mapping override and is applied to the draft config.""" + architecture-mapping override and is applied to the draft config. + """ def shrink(hf_config: PretrainedConfig) -> PretrainedConfig: hf_config.num_hidden_layers = 1 @@ -67,7 +69,8 @@ def test_callable_overrides_reach_the_draft_config(): @pytest.mark.cpu_test def test_arch_mapping_applies_before_callable_override(): """The static arch-mapping override runs first, so the user callable - observes (and may adjust) the post-mapping config.""" + observes (and may adjust) the post-mapping config. + """ seen_architectures: list[str] = [] def record(hf_config: PretrainedConfig) -> PretrainedConfig: @@ -128,7 +131,8 @@ def test_composed_override_is_picklable(): the composed override must be picklable. A nested local closure is not (it raised ``Can't get local object`` on DFlashDraftModel); a ``functools.partial`` over a module-referenceable static method is. - Guard against regressing to a closure.""" + Guard against regressing to a closure. + """ composed = SpeculativeConfig.compose_draft_hf_overrides(_module_level_shrink) assert isinstance(composed, functools.partial) diff --git a/tests/config/test_speculative_draft_max_position_embeddings.py b/tests/config/test_speculative_draft_max_position_embeddings.py index 117ae5a66d..711aaf9b03 100644 --- a/tests/config/test_speculative_draft_max_position_embeddings.py +++ b/tests/config/test_speculative_draft_max_position_embeddings.py @@ -99,7 +99,8 @@ def test_independent_draft_model_keeps_its_own_limit( vllm_caplog: pytest.LogCaptureFixture, ): """An independent AR draft may genuinely have a smaller context than the - target; its max_position_embeddings must not be resized.""" + target; its max_position_embeddings must not be resized. + """ target_model_config = ModelConfig( AR_MODEL, hf_overrides={"max_position_embeddings": 8192} ) diff --git a/tests/cuda/test_cuda_context.py b/tests/cuda/test_cuda_context.py index aca490d156..cb2ba1b199 100644 --- a/tests/cuda/test_cuda_context.py +++ b/tests/cuda/test_cuda_context.py @@ -11,7 +11,7 @@ from vllm.platforms import current_platform def check_cuda_context(): - """Check CUDA driver context status""" + """Check CUDA driver context status.""" try: cuda = ctypes.CDLL("libcuda.so") device = ctypes.c_int() @@ -22,7 +22,7 @@ def check_cuda_context(): def run_cuda_test_in_thread(device_input, expected_device_id): - """Run CUDA context test in separate thread for isolation""" + """Run CUDA context test in separate thread for isolation.""" try: # New thread should have no CUDA context initially valid_before, device_before = check_cuda_context() diff --git a/tests/cuda/test_platform_no_cuda_init.py b/tests/cuda/test_platform_no_cuda_init.py index 697d2cfec8..6f50ce606b 100644 --- a/tests/cuda/test_platform_no_cuda_init.py +++ b/tests/cuda/test_platform_no_cuda_init.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test that platform imports do not prematurely initialize CUDA. +"""Test that platform imports do not prematurely initialize CUDA. This is critical for Ray-based multi-GPU setups where workers need to set CUDA_VISIBLE_DEVICES after importing vLLM but before CUDA is initialized. diff --git a/tests/detokenizer/test_stop_string_while_stop_model_terminates.py b/tests/detokenizer/test_stop_string_while_stop_model_terminates.py index 44215cb72a..3b72d612dd 100644 --- a/tests/detokenizer/test_stop_string_while_stop_model_terminates.py +++ b/tests/detokenizer/test_stop_string_while_stop_model_terminates.py @@ -44,8 +44,7 @@ def _make_request(stop, include_stop_str_in_output: bool, min_tokens: int = 0): def test_stop_string_while_stop_token_terminates(include_stop_str_in_output: bool): - """ - This test verifies that the detokenizer correctly handles the case where + """This test verifies that the detokenizer correctly handles the case where the generated token sequence contains both: - a stop token - an token @@ -67,7 +66,6 @@ def test_stop_string_while_stop_token_terminates(include_stop_str_in_output: boo where stop-string evaluation occurs before the early-return on stop_terminated. """ - # Generate text "abcdeZ" and tokenize it. generated_text = "abcde" eos_token = "Z" diff --git a/tests/distributed/conftest.py b/tests/distributed/conftest.py index 48df856f5e..876ee9dcf5 100644 --- a/tests/distributed/conftest.py +++ b/tests/distributed/conftest.py @@ -18,13 +18,13 @@ DP_RANK = 0 @pytest.fixture def random_port(): - """Generate a random port number for testing""" + """Generate a random port number for testing.""" return random.randint(10000, 59900) @pytest.fixture def publisher_config(random_port, request): - """Create a publisher config with inproc transport""" + """Create a publisher config with inproc transport.""" how = request.param if hasattr(request, "param") else "inproc" if how == "inproc": @@ -47,7 +47,7 @@ def publisher_config(random_port, request): @pytest.fixture def publisher(publisher_config): - """Create and return a publisher instance""" + """Create and return a publisher instance.""" pub = EventPublisherFactory.create(publisher_config, DP_RANK) yield pub pub.shutdown() @@ -55,7 +55,7 @@ def publisher(publisher_config): @pytest.fixture def subscriber(publisher_config): - """Create and return a subscriber for testing""" + """Create and return a subscriber for testing.""" endpoint = publisher_config.endpoint replay_endpoint = publisher_config.replay_endpoint @@ -74,7 +74,7 @@ def subscriber(publisher_config): class MockSubscriber: - """Helper class to receive and verify published events""" + """Helper class to receive and verify published events.""" def __init__( self, @@ -113,7 +113,7 @@ class MockSubscriber: self.decoder = msgspec.msgpack.Decoder(type=decode_type) def receive_one(self, timeout=1000) -> tuple[int, SampleBatch] | None: - """Receive a single message with timeout""" + """Receive a single message with timeout.""" if not self.sub.poll(timeout): return None @@ -127,7 +127,7 @@ class MockSubscriber: return seq, data def request_replay(self, start_seq: int, socket_idx: int = 0) -> None: - """Request replay of messages starting from start_seq""" + """Request replay of messages starting from start_seq.""" if not self.replay_sockets: raise ValueError("Replay sockets not initialized") if socket_idx >= len(self.replay_sockets): @@ -138,7 +138,7 @@ class MockSubscriber: ) def receive_replay(self, socket_idx: int = 0) -> list[tuple[int, SampleBatch]]: - """Receive replayed messages from a specific replay socket""" + """Receive replayed messages from a specific replay socket.""" if not self.replay_sockets: raise ValueError("Replay sockets not initialized") if socket_idx >= len(self.replay_sockets): @@ -170,7 +170,7 @@ class MockSubscriber: return replayed def close(self): - """Clean up resources""" + """Clean up resources.""" self.sub.close() for replay in self.replay_sockets: replay.close() @@ -179,7 +179,8 @@ class MockSubscriber: @pytest.fixture def enable_ray_v2_backend(): """Set env vars for the Ray V2 executor backend and shut down Ray - between tests.""" + between tests. + """ import ray saved = { diff --git a/tests/distributed/test_context_parallel.py b/tests/distributed/test_context_parallel.py index 007948aeba..6d5d454b4a 100644 --- a/tests/distributed/test_context_parallel.py +++ b/tests/distributed/test_context_parallel.py @@ -1,11 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -WARNING: This test runs in both single-node (4 GPUs) and multi-node - (2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is - important to set the distributed backend to "mp" to avoid Ray scheduling - all workers in a node other than the head node, which can cause the test - to fail. +"""WARNING: This test runs in both single-node (4 GPUs) and multi-node +(2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is +important to set the distributed backend to "mp" to avoid Ray scheduling +all workers in a node other than the head node, which can cause the test +to fail. """ import json diff --git a/tests/distributed/test_eplb_algo.py b/tests/distributed/test_eplb_algo.py index 20360c8c61..aaa987b9fc 100644 --- a/tests/distributed/test_eplb_algo.py +++ b/tests/distributed/test_eplb_algo.py @@ -29,7 +29,7 @@ def test_eplb_load_stats_reduce_across_ranks(): def test_basic_rebalance(): - """Test basic rebalancing functionality""" + """Test basic rebalancing functionality.""" # Example from https://github.com/deepseek-ai/eplb weight = torch.tensor( [ @@ -92,7 +92,7 @@ def test_basic_rebalance(): def test_single_gpu_case(): - """Test single GPU case""" + """Test single GPU case.""" weight = torch.tensor([[10, 20, 30, 40]]) num_replicas = 4 num_groups = 1 @@ -115,7 +115,7 @@ def test_single_gpu_case(): def test_equal_weights(): - """Test case with equal weights""" + """Test case with equal weights.""" weight = torch.tensor([[50, 50, 50, 50, 50, 50, 50, 50]]) num_replicas = 8 num_groups = 2 @@ -139,7 +139,7 @@ def test_equal_weights(): def test_extreme_weight_imbalance(): - """Test extreme weight imbalance case""" + """Test extreme weight imbalance case.""" weight = torch.tensor([[1000, 1, 1, 1, 1, 1, 1, 1]]) num_replicas = 12 num_groups = 2 @@ -162,7 +162,7 @@ def test_extreme_weight_imbalance(): def test_multiple_layers(): - """Test multiple layers case""" + """Test multiple layers case.""" weight = torch.tensor( [ [10, 20, 30, 40, 50, 60], # First layer @@ -195,7 +195,7 @@ def test_multiple_layers(): def test_parameter_validation(): - """Test parameter validation""" + """Test parameter validation.""" weight = torch.tensor([[10, 20, 30, 40]]) # Test non-divisible case - this should handle normally without throwing @@ -213,7 +213,7 @@ def test_parameter_validation(): def test_small_scale_hierarchical(): - """Test small-scale hierarchical load balancing""" + """Test small-scale hierarchical load balancing.""" weight = torch.tensor( [ [100, 50, 200, 75, 150, 25, 300, 80], # 8 experts @@ -243,7 +243,7 @@ def test_small_scale_hierarchical(): def test_global_load_balance_fallback(): - """Test global load balancing fallback case""" + """Test global load balancing fallback case.""" # When num_groups % num_nodes != 0, should fall back to global load # balancing weight = torch.tensor([[10, 20, 30, 40, 50, 60]]) @@ -265,7 +265,7 @@ def test_global_load_balance_fallback(): @pytest.mark.parametrize("device", ["cpu", "cuda"]) def test_device_compatibility(device): - """Test device compatibility""" + """Test device compatibility.""" if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA not available") @@ -287,8 +287,7 @@ def test_device_compatibility(device): def test_additional_cases(): - """Test more edge cases and different parameter combinations""" - + """Test more edge cases and different parameter combinations.""" # Test case 1: Large-scale distributed setup weight1 = torch.tensor( [[50, 100, 75, 120, 90, 60, 80, 110, 40, 70, 95, 85, 65, 55, 45, 35]] @@ -320,8 +319,7 @@ def test_additional_cases(): def test_compute_logical_maps_with_negative_indices(): - """ - Test that compute_logical_maps correctly handles physical slots containing + """Test that compute_logical_maps correctly handles physical slots containing -1 (unused slots). """ # 2 layers, 6 physical slots, 4 logical experts. diff --git a/tests/distributed/test_eplb_events.py b/tests/distributed/test_eplb_events.py index 40323dc898..963e3cbb77 100644 --- a/tests/distributed/test_eplb_events.py +++ b/tests/distributed/test_eplb_events.py @@ -57,9 +57,8 @@ def test_reuse_across_multiple_cycles(): def test_producer_consumer(): - """ - This test uses the CpuGpuEvent to synchronize reads and writes to/from a shared GPU - tensor on multiple CPU threads. + """This test uses the CpuGpuEvent to synchronize reads and writes to/from a shared + GPU tensor on multiple CPU threads. """ worker_stream = torch.cuda.Stream() # Create a single element counter that will be shared between two threads diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py index 1ff41f5df2..aed0decc6d 100644 --- a/tests/distributed/test_eplb_execute.py +++ b/tests/distributed/test_eplb_execute.py @@ -56,8 +56,7 @@ def create_expert_indices_with_redundancy( total_physical_experts: int, redundancy_config: list[int], # redundancy for each logical expert ) -> torch.Tensor: - """ - Create expert indices with redundancy. + """Create expert indices with redundancy. Args: num_layers: number of layers @@ -67,6 +66,7 @@ def create_expert_indices_with_redundancy( Returns: indices: Shape (num_layers, total_physical_experts) + """ assert sum(redundancy_config) == total_physical_experts assert len(redundancy_config) == num_logical_experts @@ -95,8 +95,7 @@ def create_expert_weights( device: torch.device, physical_to_logical_mapping: torch.Tensor, ) -> list[list[torch.Tensor]]: - """ - Create fake expert weights tensor for testing. + """Create fake expert weights tensor for testing. Use `arange` to generate predictable weights values, based on logical expert ID. @@ -105,6 +104,7 @@ def create_expert_weights( Args: physical_to_logical_mapping: Shape (num_layers, num_local_experts) mapping[layer, physical_pos] = logical_expert_id + """ expert_weights = [] @@ -209,9 +209,7 @@ def verify_redundant_experts_have_same_weights( world_size: int, num_local_experts: int, ) -> bool: - """ - Verify that all replicas of the same logical expert have the same weights. - """ + """Verify that all replicas of the same logical expert have the same weights.""" num_layers = len(expert_weights) total_physical_experts = world_size * num_local_experts @@ -568,7 +566,6 @@ def test_rearrange_expert_weights_with_redundancy( eplb_communicator, ): """Test the functionality of rearranging expert weights with redundancy.""" - if eplb_communicator == "nixl" and not has_nixl(): pytest.skip("NIXL is not available") if torch.accelerator.device_count() < world_size: @@ -678,7 +675,6 @@ def test_async_transfer_layer_without_mtp( eplb_communicator: str, ): """Exercise async EPLB transfer path without MTP/spec decode.""" - if eplb_communicator == "nixl" and not has_nixl(): pytest.skip("NIXL is not available") if torch.accelerator.device_count() < world_size: @@ -696,11 +692,9 @@ def test_async_transfer_layer_without_mtp( @pytest.mark.parametrize("world_size", [2, 4]) def test_rearrange_expert_weights_no_change(world_size): - """ - Test that when the indices do not change, the weights should remain + """Test that when the indices do not change, the weights should remain unchanged. """ - if torch.accelerator.device_count() < world_size: pytest.skip(f"Need at least {world_size} GPUs to run the test") distributed_run( @@ -800,7 +794,6 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None: @pytest.mark.parametrize("world_size", [2, 4]) def test_rearrange_expert_weights_profile_mode(world_size): """Test profile mode (should not copy actual weights)""" - if torch.accelerator.device_count() < world_size: pytest.skip(f"Need at least {world_size} GPUs to run the test") distributed_run( diff --git a/tests/distributed/test_eplb_quant_scale_consistency.py b/tests/distributed/test_eplb_quant_scale_consistency.py index 3b50b315f0..49dc8afbf3 100644 --- a/tests/distributed/test_eplb_quant_scale_consistency.py +++ b/tests/distributed/test_eplb_quant_scale_consistency.py @@ -124,7 +124,8 @@ def _build_processed_layer( backend: str, raw: dict[str, torch.Tensor], device: torch.device ) -> tuple[CompressedTensorsW4A4Nvfp4MoEMethod, _RoutedExpertsStub]: """Create the method + layer, load ``raw`` and run the real - process_weights_after_loading (which builds the flashinfer kernel).""" + process_weights_after_loading (which builds the flashinfer kernel). + """ moe_config = _make_moe_config(backend) method = CompressedTensorsW4A4Nvfp4MoEMethod(moe_config, "layer.0", use_a16=False) layer = _RoutedExpertsStub(moe_config).to(device) @@ -157,7 +158,8 @@ def _assert_tensors_equal( def _simulate_eplb_rearrangement(layer: torch.nn.Module, perm: torch.Tensor) -> None: """Permute experts of every registered Parameter in place, the way - EPLB's in-place rearrangement moves expert slices along dim 0.""" + EPLB's in-place rearrangement moves expert slices along dim 0. + """ with torch.no_grad(): for _, param in layer.named_parameters(): param.copy_(param[perm]) @@ -165,7 +167,8 @@ def _simulate_eplb_rearrangement(layer: torch.nn.Module, perm: torch.Tensor) -> def _ensure_world1_distributed() -> None: """The enable_eplb config makes weight processing all-reduce activation - scale amaxes over the EP group, so a (single-rank) EP group must exist.""" + scale amaxes over the EP group, so a (single-rank) EP group must exist. + """ from vllm.distributed.parallel_state import ( ensure_model_parallel_initialized, init_distributed_environment, @@ -253,7 +256,8 @@ def test_nvfp4_eplb_rearrangement_matches_reload(backend: str) -> None: def test_fp8_per_tensor_alphas_registered_and_aliased() -> None: """CPU-only: the fp8 per-tensor oracle must register the fused (w_scale * a_scale) products as layer Parameters aliased by the quant - config, so EPLB rearrangement keeps kernels consistent.""" + config, so EPLB rearrangement keeps kernels consistent. + """ from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( Fp8MoeBackend, make_fp8_moe_quant_config, diff --git a/tests/distributed/test_eplb_spec_decode.py b/tests/distributed/test_eplb_spec_decode.py index b4211c74c1..ced206a305 100644 --- a/tests/distributed/test_eplb_spec_decode.py +++ b/tests/distributed/test_eplb_spec_decode.py @@ -75,8 +75,7 @@ def test_eplb_spec_decode( monkeypatch: pytest.MonkeyPatch, model_setup: tuple[str, str, str, int, float], ): - """ - Test the correctness of EPLB speculative decoding with GSM8K dataset. + """Test the correctness of EPLB speculative decoding with GSM8K dataset. Applicable to MoE models with mtp or eagle spec decode. """ method, model_name, spec_model_name, tp_size, expected_gsm8k_value = model_setup @@ -109,10 +108,7 @@ def test_eplb_spec_decode( @large_gpu_mark(min_gb=80) def test_eplb_spec_decode_qwen3_next_mtp_async() -> None: - """ - Ensure async EPLB works with MTP speculative decoding for Qwen3-Next. - """ - + """Ensure async EPLB works with MTP speculative decoding for Qwen3-Next.""" TASK = "gsm8k" FILTER = "exact_match,strict-match" RTOL = 0.03 diff --git a/tests/distributed/test_eplb_utils.py b/tests/distributed/test_eplb_utils.py index a6cfb0d29d..dd2bb2760b 100644 --- a/tests/distributed/test_eplb_utils.py +++ b/tests/distributed/test_eplb_utils.py @@ -30,8 +30,7 @@ def _make_model_state( def test_commit_eplb_maps_shape_change(): - """ - When the number of physical experts changes, resize the active map within its + """When the number of physical experts changes, resize the active map within its preallocated storage. """ num_layers, num_logical, num_physical = 2, 4, 6 @@ -65,8 +64,7 @@ def test_commit_eplb_maps_shape_change(): def test_commit_eplb_maps_for_layer_logical_padding(): - """ - Test that logical_to_physical_map is padded with -1 to fill the + """Test that logical_to_physical_map is padded with -1 to fill the pre-allocated slots when the new map has fewer replicas than the max. """ num_layers, num_logical, num_physical = 2, 4, 6 diff --git a/tests/distributed/test_events.py b/tests/distributed/test_events.py index 894f3f7f8b..aae80d9cce 100644 --- a/tests/distributed/test_events.py +++ b/tests/distributed/test_events.py @@ -20,27 +20,26 @@ class EventSample( tag=True, # type: ignore array_like=True, # type: ignore ): - """Test event for publisher testing""" + """Test event for publisher testing.""" id: int value: str class SampleBatch(EventBatch): - """Test event batch for publisher testing""" + """Test event batch for publisher testing.""" events: list[EventSample] def create_test_events(count: int) -> SampleBatch: - """Create a batch of test events""" + """Create a batch of test events.""" events = [EventSample(id=i, value=f"test-{i}") for i in range(count)] return SampleBatch(ts=time.time(), events=events) def test_basic_publishing(publisher, subscriber): - """Test basic event publishing works""" - + """Test basic event publishing works.""" test_batch = create_test_events(5) publisher.publish(test_batch) @@ -58,7 +57,7 @@ def test_basic_publishing(publisher, subscriber): def test_multiple_events(publisher, subscriber): - """Test publishing and receiving multiple event batches""" + """Test publishing and receiving multiple event batches.""" for _ in range(10): batch = create_test_events(2) publisher.publish(batch) @@ -75,7 +74,7 @@ def test_multiple_events(publisher, subscriber): def test_replay_mechanism(publisher, subscriber): - """Test the replay mechanism works correctly""" + """Test the replay mechanism works correctly.""" for _ in range(19): batch = create_test_events(1) publisher.publish(batch) @@ -96,7 +95,7 @@ def test_replay_mechanism(publisher, subscriber): def test_replay_includes_topic(publisher, subscriber, publisher_config): - """Test that replay responses include the topic, matching PUB format""" + """Test that replay responses include the topic, matching PUB format.""" for _ in range(5): publisher.publish(create_test_events(1)) @@ -115,7 +114,7 @@ def test_replay_includes_topic(publisher, subscriber, publisher_config): def test_buffer_limit(publisher, subscriber, publisher_config): - """Test buffer limit behavior""" + """Test buffer limit behavior.""" buffer_size = publisher_config.buffer_steps # Publish more events than the buffer can hold @@ -139,9 +138,7 @@ def test_buffer_limit(publisher, subscriber, publisher_config): def test_topic_filtering(publisher_config): - """ - Test that a subscriber only receives messages matching its topic filter - """ + """Test that a subscriber only receives messages matching its topic filter.""" publisher_config.replay_endpoint = None publisher_config.topic = "foo" @@ -174,7 +171,7 @@ def test_topic_filtering(publisher_config): def test_high_volume(publisher, subscriber): - """Test publishing and receiving a high volume of events""" + """Test publishing and receiving a high volume of events.""" num_batches = 10_000 events_per_batch = 100 @@ -210,7 +207,7 @@ def test_high_volume(publisher, subscriber): def test_null_publisher(): - """Test that NullEventPublisher can be used without errors""" + """Test that NullEventPublisher can be used without errors.""" publisher = NullEventPublisher(DP_RANK) # This should not raise any errors @@ -220,8 +217,7 @@ def test_null_publisher(): def test_data_parallel_rank_tagging(publisher_config): - """Test that events are properly tagged with their data parallel rank""" - + """Test that events are properly tagged with their data parallel rank.""" publisher_config.topic = "foo" pub_0 = EventPublisherFactory.create(publisher_config, DP_RANK) pub_1 = EventPublisherFactory.create(publisher_config, DP_RANK + 1) @@ -285,7 +281,7 @@ def test_data_parallel_rank_tagging(publisher_config): def test_event_publisher_factory(random_port): - """Test event publisher factory creation behavior under different configurations""" + """Test event publisher factory creation behavior under different configurations.""" from vllm.config.kv_events import KVEventsConfig from vllm.distributed.kv_events import ZmqEventPublisher diff --git a/tests/distributed/test_expert_placement.py b/tests/distributed/test_expert_placement.py index 46f63408f4..4174b5c6cc 100644 --- a/tests/distributed/test_expert_placement.py +++ b/tests/distributed/test_expert_placement.py @@ -49,7 +49,6 @@ def verify_round_robin_pattern(expert_map, ep_rank, ep_size, global_num_experts) @pytest.mark.parametrize("world_size", [2, 4]) def test_expert_placement_various_sizes(expert_placement_strategy, world_size): """Test round_robin expert placement with various expert counts.""" - # Test with different global_num_experts values # Include both divisible and non-divisible cases if world_size == 2: @@ -116,7 +115,6 @@ def test_expert_placement_various_sizes(expert_placement_strategy, world_size): @pytest.mark.parametrize("world_size", [2, 4]) def test_expert_placement_edge_cases(expert_placement_strategy, world_size): """Test edge cases for round_robin expert placement.""" - # Test case 1: ep_size = 1 (should return None for expert_map) local_num_experts, expert_map, _ = determine_expert_map( ep_size=1, @@ -139,7 +137,6 @@ def test_expert_placement_edge_cases(expert_placement_strategy, world_size): def test_determine_expert_map_comprehensive(): """Test of determine_expert_map function with various configurations.""" - # Test cases: (ep_size, ep_rank, global_num_experts, # expert_placement_strategy, expected_local, expected_map_pattern) test_cases = [ diff --git a/tests/distributed/test_mnnvl_alltoall.py b/tests/distributed/test_mnnvl_alltoall.py index 636f7cb525..0258c781c7 100644 --- a/tests/distributed/test_mnnvl_alltoall.py +++ b/tests/distributed/test_mnnvl_alltoall.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for MNNVL AllToAll operations. +"""Tests for MNNVL AllToAll operations. Requires: docker run ... --cap-add=SYS_PTRACE ... Run: pytest tests/distributed/test_mnnvl_alltoall.py -v @@ -87,6 +86,7 @@ def _run_worker(rank, world_size, port, worker_fn, dp_size, dp_port, err_queue): Otherwise use tp=world_size (default for EP-based tests). dp_port: Separate port for the DP master (only used when dp_size is set). err_queue: Queue for propagating tracebacks to the parent process. + """ try: os.environ.pop("CUDA_VISIBLE_DEVICES", None) @@ -115,6 +115,7 @@ def _init_dp_environment(world_size, rank, port, dp_size, dp_port): Args: port: Port for torch.distributed init. dp_port: Separate port for the DP master group init. + """ from vllm.config import VllmConfig, set_current_vllm_config from vllm.config.parallel import ParallelConfig @@ -159,7 +160,8 @@ def _make_forward_context(rank, world_size, num_tokens_per_rank): class _AttnMeta: """Minimal placeholder so set_forward_context's ``attn_metadata is not None`` guard (forward_context.py:334) - is satisfied. The real DPMetadata is built from num_tokens_across_dp.""" + is satisfied. The real DPMetadata is built from num_tokens_across_dp. + """ dp_metadata = None @@ -458,7 +460,8 @@ def _one_sided_workspace_grow_worker(rank, world_size): def test_one_sided_manager_workspace_grow(world_size): """A later initialize() with a larger per-token payload must grow the workspace and rebuild MoeAlltoAll; a later initialize() with a smaller - payload must no-op.""" + payload must no-op. + """ _spawn_workers( _one_sided_workspace_grow_worker, world_size, diff --git a/tests/distributed/test_mq_connect_ip.py b/tests/distributed/test_mq_connect_ip.py index 4b0cdda3ad..701dc8f248 100644 --- a/tests/distributed/test_mq_connect_ip.py +++ b/tests/distributed/test_mq_connect_ip.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test that MessageQueue uses the local node's IP for binding, +"""Test that MessageQueue uses the local node's IP for binding, not a remote master_addr. This validates the fix for cross-node data-parallel where each DP group leader must bind to its own IP. @@ -22,7 +21,8 @@ from vllm.utils.network_utils import get_ip def test_mq_bind_with_local_ip(): """MessageQueue with remote readers should successfully bind - when connect_ip is the local node's IP.""" + when connect_ip is the local node's IP. + """ # n_reader=2, n_local_reader=1 means 1 remote reader, # which triggers the remote ZMQ socket bind. mq = MessageQueue( @@ -44,7 +44,8 @@ def test_mq_bind_with_local_ip(): def test_mq_bind_with_non_local_ip_fails(): """MessageQueue should fail to bind when connect_ip is a non-local IP address (simulating the bug where master_addr - from a different node was used).""" + from a different node was used). + """ # Use a non-local IP that we definitely can't bind to. # 198.51.100.1 is from TEST-NET-2 (RFC 5737), never locally assigned. non_local_ip = "198.51.100.1" @@ -58,7 +59,8 @@ def test_mq_bind_with_non_local_ip_fails(): def test_mq_bind_defaults_to_local_ip(): """When connect_ip is None, MessageQueue should auto-detect - the local IP and bind successfully.""" + the local IP and bind successfully. + """ mq = MessageQueue( n_reader=2, n_local_reader=1, diff --git a/tests/distributed/test_multiproc_executor.py b/tests/distributed/test_multiproc_executor.py index 599025d1c8..d0a417a788 100644 --- a/tests/distributed/test_multiproc_executor.py +++ b/tests/distributed/test_multiproc_executor.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Integration tests for MultiprocExecutor at the executor level. +"""Integration tests for MultiprocExecutor at the executor level. This test directly tests the executor without going through the LLM interface, focusing on executor initialization, RPC calls, and distributed execution. """ @@ -377,8 +376,7 @@ def test_multiproc_executor_properties(): @multi_gpu_test(num_gpus=4) def test_multiproc_executor_multi_node(): - """ - Test MultiprocExecutor with multi-node configuration. + """Test MultiprocExecutor with multi-node configuration. This simulates 2 nodes with TP=4: - Node 0 (rank 0): Uses GPUs 0,1 (CUDA_VISIBLE_DEVICES=0,1) with TP=2 - Node 1 (rank 1): Uses GPUs 2,3 (CUDA_VISIBLE_DEVICES=2,3) with TP=2 diff --git a/tests/distributed/test_pipeline_parallel.py b/tests/distributed/test_pipeline_parallel.py index a464a61683..8ba93df4f5 100644 --- a/tests/distributed/test_pipeline_parallel.py +++ b/tests/distributed/test_pipeline_parallel.py @@ -1,11 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -WARNING: This test runs in both single-node (4 GPUs) and multi-node - (2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is - important to set the distributed backend to "mp" to avoid Ray scheduling - all workers in a node other than the head node, which can cause the test - to fail. +"""WARNING: This test runs in both single-node (4 GPUs) and multi-node +(2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is +important to set the distributed backend to "mp" to avoid Ray scheduling +all workers in a node other than the head node, which can cause the test +to fail. """ import json diff --git a/tests/distributed/test_pynccl.py b/tests/distributed/test_pynccl.py index 91dccd5798..0eda9a14e0 100644 --- a/tests/distributed/test_pynccl.py +++ b/tests/distributed/test_pynccl.py @@ -495,7 +495,8 @@ def test_pynccl_suspend_resume_idempotent(): def test_pynccl_suspend_noop_without_symbol(): """RCCL / NCCL < 2.29.7 lack ncclCommSuspend; suspend() must no-op, not - crash, and resume() must then no-op too since nothing was suspended.""" + crash, and resume() must then no-op too since nothing was suspended. + """ from unittest.mock import Mock comm = object.__new__(PyNcclCommunicator) diff --git a/tests/distributed/test_quick_all_reduce.py b/tests/distributed/test_quick_all_reduce.py index fb2b04abf0..f9228fbd68 100644 --- a/tests/distributed/test_quick_all_reduce.py +++ b/tests/distributed/test_quick_all_reduce.py @@ -396,8 +396,7 @@ def test_custom_quick_allreduce_bf16_cast(monkeypatch: pytest.MonkeyPatch): def qr_variable_input(rank, world_size): - """ - When the tensor parallelism is set to 4 or 8, frequent changes + """When the tensor parallelism is set to 4 or 8, frequent changes in the input shape can cause QuickReduce to hang (this issue has been observed with the gpt_oss model). """ diff --git a/tests/distributed/test_ray_v2_executor.py b/tests/distributed/test_ray_v2_executor.py index bd4a253b96..f9b9e57fc5 100644 --- a/tests/distributed/test_ray_v2_executor.py +++ b/tests/distributed/test_ray_v2_executor.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Integration tests for RayExecutorV2 at the executor level. +"""Integration tests for RayExecutorV2 at the executor level. Validates executor initialization, placement group support, RPC calls, and distributed execution with various TP/PP configurations. """ @@ -299,7 +298,6 @@ def test_ray_v2_run_refs_stored_for_monitoring(executor): @pytest.mark.parametrize("tp_size, pp_size", [(2, 1), (2, 2)]) def test_ray_v2_single_node_generation(tp_size, pp_size): """End-to-end LLM generation with RayExecutorV2.""" - llm = LLM( model=MODEL, tensor_parallel_size=tp_size, diff --git a/tests/distributed/test_ray_v2_executor_e2e.py b/tests/distributed/test_ray_v2_executor_e2e.py index fb58301326..27de410e38 100644 --- a/tests/distributed/test_ray_v2_executor_e2e.py +++ b/tests/distributed/test_ray_v2_executor_e2e.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Orchestration-level integration tests for RayExecutorV2. -""" +"""Orchestration-level integration tests for RayExecutorV2.""" import gc import os @@ -25,7 +23,8 @@ def _ray_init(): Without this, workers cannot unpickle actor classes defined in the ``tests`` package, causing FunctionActorManager to fall back to - TemporaryActor which drops async method signatures.""" + TemporaryActor which drops async method signatures. + """ project_root = str(pathlib.Path(__file__).resolve().parents[2]) ray.init( ignore_reinit_error=True, @@ -158,8 +157,7 @@ def test_multi_replicas_with_bundle_indices(ray_init): def test_env_var_and_runtime_env_propagation(): - """ - Verify env vars (NCCL_, HF_) and parallel_config.ray_runtime_env + """Verify env vars (NCCL_, HF_) and parallel_config.ray_runtime_env propagate to RayWorkerProc actors. """ sentinel_vars = { diff --git a/tests/distributed/test_sharded_rdt_plan.py b/tests/distributed/test_sharded_rdt_plan.py index 7a75f0b7fc..5cf0a82203 100644 --- a/tests/distributed/test_sharded_rdt_plan.py +++ b/tests/distributed/test_sharded_rdt_plan.py @@ -53,7 +53,8 @@ META = torch.device("meta") class _FakeLayer: """Stand-in for a leaf module. The planner only stores it and keys on - ``id()``; nothing in the plan path touches its attributes.""" + ``id()``; nothing in the plan path touches its attributes. + """ def __init__(self, tag: str) -> None: self.tag = tag @@ -64,7 +65,8 @@ class _FakeLayer: def _module(layer, copies): """What the bake records for one leaf module: a scatter list whose entries - all carry that module as their ``layer``.""" + all carry that module as their ``layer``. + """ for c in copies: c.layer = layer return copies @@ -81,7 +83,8 @@ def _copy( dtype=torch.bfloat16, ): """A `_Scatter` as the bake would record it: source key, owning layer, and - the meta destination region.""" + the meta destination region. + """ stride = ( (1,) if len(shape) == 1 @@ -140,7 +143,8 @@ def _planner(baked, *, name_meta, live=None, held_by=None, group_lens=None): def _one_module_per_layer(n_layers, *, dtype="bfloat16", numel=4): """`(baked, name_meta, names, group_lens)` for a pre / N-layer / post model - where every gather group is one fully-baked leaf module.""" + where every gather group is one fully-baked leaf module. + """ names = ["embed.weight"] names += [f"model.layers.{i}.w" for i in range(n_layers)] names += ["norm.weight"] @@ -161,7 +165,8 @@ def _one_module_per_layer(n_layers, *, dtype="bfloat16", numel=4): class TestFakeOpChains: """Every allowlisted op must append itself to the chain and hand back a child whose shape/dtype PyTorch itself computed. The chain is the wire - format the producer replays, so its exact contents are load-bearing.""" + format the producer replays, so its exact contents are load-bearing. + """ def _fake(self, shape=(4, 6), dtype=torch.bfloat16, sink=None): return FakeRDTTensor( @@ -222,7 +227,8 @@ class TestFakeOpChains: @pytest.mark.parametrize("op,n", [("chunk", 2), ("unbind", 4)]) def test_multi_return_op_emits_one_child_per_output(self, op, n): """chunk/unbind hand back a tuple; each child carries the base op plus a - trailing __getitem__(i) so the producer can index the replayed result.""" + trailing __getitem__(i) so the producer can index the replayed result. + """ parts = getattr(self._fake(), op)(*((n,) if op == "chunk" else ()), 0) assert isinstance(parts, tuple) assert len(parts) == n @@ -238,7 +244,8 @@ class TestFakeOpChains: class TestFakeUnsupportedOps: """Anything that needs real data must fail loudly at bake time rather than - silently transferring the wrong bytes.""" + silently transferring the wrong bytes. + """ def _fake(self): return FakeRDTTensor( @@ -278,7 +285,8 @@ class TestFakeUnsupportedOps: class TestBakeRecording: """During the dry run the fake's ``copy_`` is the data sink: it records the - source chain plus the meta destination's strided region and moves nothing.""" + source chain plus the meta destination's strided region and moves nothing. + """ def _recorder_and_fake(self, shape=(4, 6)): rec = BakeSink() @@ -314,7 +322,8 @@ class TestBakeRecording: def test_unattributed_copy_is_live_but_unrecorded(self): """A copy_ with no loader stamp cannot be attributed to a param, so its - module must fall back to the plain load — but the name still moved data.""" + module must fall back to the plain load — but the name still moved data. + """ rec, fake = self._recorder_and_fake() rec.current = None torch.empty((4, 6), dtype=torch.bfloat16, device=META).copy_(fake) @@ -354,7 +363,8 @@ class TestBakeRecording: producer packs what the replay yields, so the record must carry the POST-chain dtype: taking it from the source name's metadata instead sizes the slice with the wrong itemsize and shifts every later slice in the - chunk, carving the packed blob differently on the two sides.""" + chunk, carving the packed blob differently on the two sides. + """ rec, fake = self._recorder_and_fake(shape=(4,)) layer = _FakeLayer("reinterpreted") rec.current = (layer, "weight") @@ -374,7 +384,8 @@ class TestBakeRecording: the producer packs what the chain yields, and nothing downstream can catch a mismatch — the consumer's buffer view is exactly prod(dest.shape) elements, so it reshapes cleanly over bytes laid out at other offsets. - So the bake refuses rather than recording a slice it cannot carve.""" + So the bake refuses rather than recording a slice it cannot carve. + """ rec, fake = self._recorder_and_fake(shape=(1, 6)) rec.current = (_FakeLayer("broadcast"), "weight") param = torch.empty((4, 6), dtype=torch.bfloat16, device=META) @@ -390,7 +401,8 @@ class TestBakeRecording: class TestConsumerIdentity: """Every worker in the fleet needs a DISTINCT id in 0..C-1: it selects the worker's producer block and keys the producer's per-consumer serve ring, so - a collision silently serves two workers out of one ring.""" + a collision silently serves two workers out of one ring. + """ def _engine(self, *, dp_index, rank, world_size): eng = object.__new__(ShardedRDTWeightTransferEngine) @@ -427,7 +439,8 @@ class TestConsumerIdentity: def test_independent_engines_offset_into_distinct_ranges(self): """A fleet of separate engines restarts _global_worker_index at 0 in each, so without the replica offset every engine would claim 0..w-1. - The driver sets replica_rank on the payload; these are the ids it buys.""" + The driver sets replica_rank on the payload; these are the ids it buys. + """ ids = [ self._ids(num_consumers=8, workers=range(4), replica_rank=r, num_replicas=2) for r in (0, 1) @@ -453,7 +466,8 @@ def _router( group_lens=None, ): """A router over ``n`` single-name groups unless told otherwise, so a test - can state expectations in group indices.""" + can state expectations in group indices. + """ if names is None: n = len(name_owner_class or []) or (len(group_lens or []) or 1) names = [f"g{i}" for i in range(n)] @@ -482,7 +496,8 @@ class TestRdtRouter: def test_every_producer_nic_carries_traffic(self): """16 producers / 8 consumers: the block rule spreads each consumer's pulls over its block, alternating by group, so no producer sits idle. - Plain `consumer_id % len(owners)` would idle half of them.""" + Plain `consumer_id % len(owners)` would idle half of them. + """ r = _router(16, 8, group_lens=[1] * 95) names = [f"g{i}" for i in range(95)] for c in range(8): @@ -500,7 +515,8 @@ class TestRdtRouter: def test_pipeline_stages_route_to_the_owning_stage(self): """2 stages x 8 ranks: groups 0-2 on stage 0, 3-5 on stage 1. Two owner - sets, one per stage; each consumer reaches both.""" + sets, one per stage; each consumer reaches both. + """ r = _router( 16, 8, @@ -516,7 +532,8 @@ class TestRdtRouter: def test_expert_names_route_to_the_holding_ranks(self): """A name held by only some ranks routes to those; a replicated name - keeps the full set.""" + keeps the full set. + """ r = RdtRouter( 8, 4, @@ -535,7 +552,8 @@ class TestRdtRouter: def test_consumers_spread_over_an_owner_set(self): """Several ranks hold a name (its TP peers); the block rule must spread - consumers across them, not funnel through one.""" + consumers across them, not funnel through one. + """ r = _router( 4, 4, owner_sets=[[0, 1]], name_owner_class=[0] * 6, group_lens=[1] * 6 ) @@ -544,7 +562,8 @@ class TestRdtRouter: def test_a_group_frees_at_every_owner_of_any_of_its_names(self): """The free barrier is per group, so its fan-out is the UNION over the - group's names: a rank holding none of them is not signalled.""" + group's names: a rank holding none of them is not signalled. + """ r = RdtRouter(4, 2, [[0, 1], [2, 3]], [0, 1, 0], ["a", "b", "c"], [2, 1]) assert r.group_owners(0) == [0, 1, 2, 3] assert r.group_owners(1) == [0, 1] @@ -614,7 +633,8 @@ class TestReplicaOverlay: def test_adding_a_deployment_does_not_move_the_first_one(self): """The strongest form: an existing consumer's routes are byte-identical - before and after another deployment is provisioned.""" + before and after another deployment is provisioned. + """ one, names = self._fleet(16, 8, 1) two, _ = self._fleet(16, 8, 2) for consumer_id in range(8): @@ -624,7 +644,8 @@ class TestReplicaOverlay: def test_one_deployment_is_the_historical_carve(self): """Default (no width) and width == the whole fleet must agree with the - plain block rule, so a single-deployment fleet is untouched.""" + plain block rule, so a single-deployment fleet is untouched. + """ plain = _router(16, 8, group_lens=[1] * 6) explicit, names = self._fleet(16, 8, 1) for consumer_id in range(8): @@ -638,7 +659,8 @@ class TestReplicaOverlay: def test_the_overlay_holds_for_partially_owned_names(self): """Pipeline stages and experts too: two owner sets, one narrow. The narrow one already forced agreement (one owner, no block to carve); the - wide one is what the carve decides.""" + wide one is what the carve decides. + """ r, names = self._fleet( 8, 4, @@ -654,14 +676,16 @@ class TestReplicaOverlay: def test_a_width_that_does_not_divide_the_fleet_raises(self): """A uniform fleet is asserted upstream in ``get_world_size``; if it ever is not, the overlay would map two workers of one deployment onto the same - block index, so refuse rather than serve the wrong bytes.""" + block index, so refuse rather than serve the wrong bytes. + """ with pytest.raises(ValueError, match="does not divide"): RdtRouter(4, 6, None, None, ["a"], [1], workers_per_replica=4) def test_pulls_still_carry_the_global_consumer_id(self): """The carve uses the index WITHIN a deployment, but the wire must keep the fleet-global id: it is what lets the producer tell the sharers of a - slot apart, and count their arrivals separately.""" + slot apart, and count their arrivals separately. + """ class _Method: def __init__(self): @@ -704,7 +728,8 @@ def _producer_pack_offsets(slices): class TestPackedLayout: def test_consumer_layout_matches_the_producer_rule(self): """Mixed dtypes, sizes that do not land on 16B — i.e. the Kimi-style - group (fp8 weights + fp32 scales + bf16 norms) that exposes alignment.""" + group (fp8 weights + fp32 scales + bf16 norms) that exposes alignment. + """ specs = [ ("w.fp8", "float8_e4m3fn", [17]), ("w.scale", "float32", [3]), @@ -779,7 +804,8 @@ class TestPackedLayout: class TestChunkModuleScatters: """The chunk cut: one chunk per distinct owner class present in the copies, ascending by class index. Derived purely from the bake plus the ownership - table, so vLLM's expert placement needs no cases.""" + table, so vLLM's expert placement needs no cases. + """ def test_unstamped_copies_form_one_chunk(self): layer = _FakeLayer("l") @@ -793,7 +819,8 @@ class TestChunkModuleScatters: """Classes are numbered by first appearance in metadata order, so the replicated class (0) leads and the rest follow — the ordering the old "-1 first, then ascending stamp" rule produced, now falling out of the - table instead of a special case.""" + table instead of a special case. + """ layer = _FakeLayer("moe") names = ["norm", "e5", "e0", "e2"] copies = [_copy(n, shape=(4,)) for n in names] @@ -813,7 +840,8 @@ class TestChunkModuleScatters: def test_copy_order_within_a_chunk_is_bake_order(self): """Scatter order within a chunk must follow the bake, not the stamp - sort: the packed layout and the producer's replay agree on keys order.""" + sort: the packed layout and the producer's replay agree on keys order. + """ layer = _FakeLayer("moe") names = ["e3", "e1", "e2"] # one stamp, deliberately unsorted names copies = [_copy(n, shape=(4,)) for n in names] @@ -827,7 +855,8 @@ class TestChunkModuleScatters: def test_stamps_interleave_across_modules_without_reordering_within(self): """Two modules' copies bucket by stamp; within each bucket the order is - module-then-bake order.""" + module-then-bake order. + """ a, b = _FakeLayer("a"), _FakeLayer("b") mods = [ _module(a, [_copy("a0", shape=(4,)), _copy("a1", shape=(4,))]), @@ -845,9 +874,10 @@ class TestChunkModuleScatters: ] def test_scatters_carry_their_own_dtype(self): - """dtype rides the record from the bake, where it is the fake's dtype + """Dtype rides the record from the bake, where it is the fake's dtype AFTER its op chain — not a plan-time lookup of the source name, which - would be wrong for any chain that reinterprets dtype.""" + would be wrong for any chain that reinterprets dtype. + """ layer = _FakeLayer("l") eng = _planner({}, name_meta={"w": ("float32", [10])}) ((_er, scatters),) = eng._chunk_module_scatters( @@ -886,7 +916,8 @@ class TestBuildCallPlan: def test_materialize_fires_on_a_modules_first_chunk_only(self): """Empty HF params are allocated once per module, by construction — - including a FusedMoE-like module whose copies span owner-class chunks.""" + including a FusedMoE-like module whose copies span owner-class chunks. + """ layer = _FakeLayer("spanning") names = [f"w{i}" for i in range(4)] copies = [_copy(n, shape=(4,)) for n in names] @@ -921,7 +952,8 @@ class TestBuildCallPlan: def test_free_signal_waits_for_a_groups_last_chunk_when_it_is_split(self): """The load-bearing case: with the group cut across owner-class chunks the signal must hang off the LAST one. Signaling earlier lets the producers - drop the gather buffers while a later chunk's RDMA is still reading.""" + drop the gather buffers while a later chunk's RDMA is still reading. + """ layer = _FakeLayer("l") names = [f"w{i}" for i in range(4)] copies = [_copy(n, shape=(4,)) for n in names] @@ -937,7 +969,8 @@ class TestBuildCallPlan: def test_free_signal_timing_across_two_split_groups(self): """Two groups x 2 owner-class chunks each: signals land on each group's - last chunk only.""" + last chunk only. + """ names_a = [f"model.layers.0.w{i}" for i in range(4)] names_b = [f"model.layers.1.w{i}" for i in range(4)] baked, name_meta, stamps = {}, {}, {} @@ -955,7 +988,8 @@ class TestBuildCallPlan: def test_a_group_with_nothing_local_is_signaled_at_sync_start(self): """The owners still published it, so this consumer must signal it — immediately, before the pipeline (signal-before-publish is tolerated), - never hung off another group's chunk.""" + never hung off another group's chunk. + """ baked, name_meta, names, group_lens = _one_module_per_layer(2) del baked["model.layers.1.w"] # nothing baked, and not live => no pull eng = _planner( @@ -976,7 +1010,8 @@ class TestBuildCallPlan: def test_modules_are_deduped_within_a_group(self): """Several names of one fused module map to the same scatter list; the - plan must not chunk it twice.""" + plan must not chunk it twice. + """ layer = _FakeLayer("fused") names = ["qkv.q", "qkv.k", "qkv.v"] copies = [_copy(n, shape=(4,)) for n in names] @@ -992,7 +1027,8 @@ class TestBuildCallPlan: def test_plan_is_a_pure_function_of_its_inputs(self): """Called twice with the same arguments it must produce the same plan — - that is what lets the engine build it once and reuse it every sync.""" + that is what lets the engine build it once and reuse it every sync. + """ baked, name_meta, names, group_lens = _one_module_per_layer(2) eng = _planner(baked, name_meta=name_meta) a = eng._build_call_plan(names, group_lens) @@ -1012,7 +1048,8 @@ class TestBuildCallPlan: class TestUnbakedNamesGuard: """A live name with no baked plan has no way to load — there is no fallback, and its pull would target groups the pipeline already freed — so the plan - build fails at init, naming the names.""" + build fails at init, naming the names. + """ def test_a_live_unbaked_name_fails_the_plan_build(self): baked, name_meta, names, group_lens = _one_module_per_layer(2) @@ -1023,7 +1060,8 @@ class TestUnbakedNamesGuard: def test_never_copied_names_are_dropped_entirely(self): """Experts owned by another EP rank no-op in their loader: not an error, - and not a chunk — their group is pre-freed if nothing else fills it.""" + and not a chunk — their group is pre-freed if nothing else fills it. + """ baked, name_meta, names, group_lens = _one_module_per_layer(2) del baked["norm.weight"] eng = _planner(baked, name_meta=name_meta, live=set(names) - {"norm.weight"}) @@ -1035,13 +1073,15 @@ class TestUnbakedNamesGuard: class TestCallPlanRouting: """Under partial ownership each chunk must be pulled from a producer that actually holds every name in it. Consumers bind EVERY producer (signals fan - out to all owners), so local index == trainer rank.""" + out to all owners), so local index == trainer rank. + """ def _routed_planner( self, held_by_group, num_consumers, consumer_id, *, held_by=None ): """``held_by_group[g]`` = the ranks holding group g's names, so a test - can still speak in groups; ``held_by`` overrides individual names.""" + can still speak in groups; ``held_by`` overrides individual names. + """ baked, name_meta, names, group_lens = _one_module_per_layer( len(held_by_group) - 2 ) @@ -1068,14 +1108,16 @@ class TestCallPlanRouting: def test_two_stage_ownership_splits_chunks_between_owners(self): """Two PP stages, 4 groups: stage 0 holds the first half, stage 1 the - second. Every chunk goes to a holder of its names.""" + second. Every chunk goes to a holder of its names. + """ eng, names, group_lens = self._routed_planner([[0], [0], [1], [1]], 1, 0) plan = eng._build_call_plan(names, group_lens) assert [c.owner for c in plan.chunks] == [0, 0, 1, 1] def test_a_multi_owner_group_still_resolves_to_one_producer(self): """Every chunk is served by exactly ONE producer: splitting a pull only - multiplies produce calls, since the consumer's NIC bounds it.""" + multiplies produce calls, since the consumer's NIC bounds it. + """ eng, names, group_lens = self._routed_planner([[0, 1]] * 4, 1, 0) plan = eng._build_call_plan(names, group_lens) assert all(0 <= c.owner < 2 for c in plan.chunks) @@ -1083,7 +1125,8 @@ class TestCallPlanRouting: def test_an_expert_name_routes_to_the_rank_holding_it(self): """PP ∩ EP: stage 0 = ranks 0-1, stage 1 = ranks 2-3, and one expert name inside a stage-1 group is held by rank 3 alone. The name's own owner - set is what routes it.""" + set is what routes it. + """ eng, names, group_lens = self._routed_planner( [[0, 1], [0, 1], [2, 3], [2, 3]], 1, @@ -1098,7 +1141,8 @@ class TestCallPlanRouting: def test_an_unheld_name_is_rejected_at_init(self): """A name no rank holds can never be served. The trainer raises when it transposes the holdings; the router's validate() is the consumer-side - backstop, and either way it fails at init rather than at first pull.""" + backstop, and either way it fails at init rather than at first pull. + """ with pytest.raises(ValueError, match="empty"): RdtRouter(2, 1, [[]], [0], ["w"], [1]).validate() @@ -1107,14 +1151,16 @@ class TestSignalCompleteness: """Every gather group is signaled exactly once — after its last chunk when the worker pulls from it, at sync start otherwise — for ANY expert placement. A group signaled twice over-credits the barrier; a group never - signaled parks a producer credit and hangs end_sync.""" + signaled parks a producer credit and hangs end_sync. + """ def _moe_planner(self, worker_experts, *, ep_size=4, n_experts=8, n_layers=2): - """pre / n_layers MoE layers / post. Each layer: one norm (its own + """Pre / n_layers MoE layers / post. Each layer: one norm (its own module, stamp -1) + n_experts expert names stamped ``e // n_local``. The worker's bake covers only ``worker_experts`` (its placement); foreign expert names never copied => dropped entirely, as in the real - bake.""" + bake. + """ n_local = n_experts // ep_size names = ["embed.weight"] group_lens = [1] @@ -1173,8 +1219,9 @@ class TestSignalCompleteness: assert self._signals(plan) == list(range(len(group_lens))) def test_placement_changes_only_the_chunk_count_never_the_signals(self): - """linear experts 0-3 hit 2 owner classes; round_robin 0,2,4,6 - hits all 4. More chunks per group, identical signal set.""" + """Linear experts 0-3 hit 2 owner classes; round_robin 0,2,4,6 + hits all 4. More chunks per group, identical signal set. + """ plans = {} for label, experts in (("linear", [0, 1, 2, 3]), ("round_robin", [0, 2, 4, 6])): eng, names, group_lens = self._moe_planner(worker_experts=experts) @@ -1188,7 +1235,8 @@ class TestSignalCompleteness: def test_the_fused_module_materializes_first_and_quants_last_across_chunks(self): """The FusedMoE shape: one module's copies span every owner-class chunk of - its group; materialize on its first, quant on its last.""" + its group; materialize on its first, quant on its last. + """ eng, names, group_lens = self._moe_planner( worker_experts=[0, 2, 4, 6], n_layers=1 ) @@ -1246,7 +1294,8 @@ class TestLayerwiseGroups: def test_flattening_a_partition_returns_the_input_order(self): """Un-indexed names group by POSITION, not by name class — so one that looks like a pre-block name still lands after the layers it follows. - Group index therefore means the same thing on every rank.""" + Group index therefore means the same thing on every rank. + """ names = [ "embed.weight", "model.layers.0.a", @@ -1264,7 +1313,8 @@ class TestLayerwiseGroups: def test_stacks_keep_prefix_appearance_order_and_sort_within(self): """Two stacks stay in the order they appear — the vision tower before the - text stack — while each sorts by index internally.""" + text stack — while each sorts by index internally. + """ names = [ "visual.blocks.1.w", "visual.blocks.0.w", @@ -1281,7 +1331,8 @@ class TestLayerwiseGroups: def test_the_post_block_lands_last_however_early_it_arrives(self): """Megatron-Bridge streams the last pipeline stage's output block before its layers. Sweeping un-indexed names after the first layer into a - trailing group is what lets the gather loop walk groups ascending.""" + trailing group is what lets the gather loop walk groups ascending. + """ names = ["model.layers.0.w", "model.norm.weight", "model.layers.1.w"] assert layerwise_groups(names) == [ ["model.layers.0.w"], @@ -1291,7 +1342,8 @@ class TestLayerwiseGroups: def test_names_sharing_a_layer_coalesce_when_the_source_interleaves(self): """A raw-checkpoint source may yield in shard-packing order. Coalescing on - the key keeps that from shattering one layer into many groups.""" + the key keeps that from shattering one layer into many groups. + """ names = [ "model.layers.0.a", "model.layers.1.a", @@ -1333,13 +1385,15 @@ class TestLayerwiseGroups: def test_every_supported_naming_convention_partitions_per_layer(self, prefix): """One decoder layer per group under any prefix, with no per-architecture table. A literal `model.layers.` match silently yields one whole-model - group for most of these.""" + group for most of these. + """ names = [f"{prefix}{i}.self_attn.q_proj.weight" for i in range(3)] assert layerwise_groups(names) == [[n] for n in names] def test_a_moe_layer_stays_whole(self): """The OUTERMOST index wins, so per-expert names group by layer. Splitting - on the expert would cut FusedMoE — one leaf module — across groups.""" + on the expert would cut FusedMoE — one leaf module — across groups. + """ names = [ "model.layers.0.mlp.experts.0.w1.weight", "model.layers.0.mlp.experts.1.w1.weight", @@ -1356,7 +1410,8 @@ class TestLayerwiseGroups: def test_a_vlm_partitions_the_vision_tower_and_the_text_stack(self): """Qwen-VL layout. Today's literal `model.layers.` match puts every one of these in ONE group — the whole model — because the text stack ships as - `model.language_model.layers.`.""" + `model.language_model.layers.`. + """ names = [ "visual.patch_embed.proj.weight", "visual.blocks.0.attn.qkv.weight", @@ -1375,7 +1430,8 @@ class TestLayerwiseGroups: def test_two_stacks_sharing_an_index_do_not_merge(self): """The key carries the prefix, so vision block 0 and text layer 0 are - different groups even though both are index 0.""" + different groups even though both are index 0. + """ names = [ "visual.blocks.0.attn.qkv.weight", "model.language_model.layers.0.self_attn.q_proj.weight", @@ -1412,20 +1468,23 @@ class TestOpAllowlistAgreement: def test_every_bake_recordable_op_is_replayable(self): """Anything the bake can record must be serveable, or the failure lands a - whole sync later than the mistake.""" + whole sync later than the mistake. + """ assert not set(SUPPORTED_OPS.values()) - set(ALLOWED_OPS) def test_the_producer_allows_nothing_the_consumer_cannot_emit(self): """The allowlist is also a guard against a misbehaving or spoofed consumer invoking arbitrary methods on trainer tensors, so it must not be wider than the ops the bake can actually produce. ``to`` in particular - would let a replay change dtype or device.""" + would let a replay change dtype or device. + """ assert not set(ALLOWED_OPS) - set(SUPPORTED_OPS.values()) assert "to" not in ALLOWED_OPS def test_transpose_via_t_is_serveable(self): """The specific regression: a chain recorded from ``.t()`` must pass the - producer's guard.""" + producer's guard. + """ fake = FakeRDTTensor( name="w", shape=torch.Size((4, 6)), @@ -1470,7 +1529,8 @@ BAKE_MODELS = [ class TestRequiresTheRayExecutor: """The engine's data plane is Ray's, so a non-Ray executor cannot work. The check is at construction because the alternative is an opaque failure during - the first handshake, long after the misconfiguration.""" + the first handshake, long after the misconfiguration. + """ def _construct(self, backend): cfg = SimpleNamespace( @@ -1497,7 +1557,8 @@ class TestRequiresTheRayExecutor: def test_a_custom_executor_class_is_left_alone(self): """A `type[Executor]` override is deliberate and unjudgeable here, so it - must not be rejected for merely not being the string "ray".""" + must not be rejected for merely not being the string "ray". + """ class _CustomExecutor: pass diff --git a/tests/distributed/test_sharded_rdt_producer.py b/tests/distributed/test_sharded_rdt_producer.py index 5cfab1957d..7036d10289 100644 --- a/tests/distributed/test_sharded_rdt_producer.py +++ b/tests/distributed/test_sharded_rdt_producer.py @@ -236,7 +236,8 @@ class TestPublishAndRebuild: def test_views_of_one_storage_do_not_overlap(self, server_factory): """One IPC export per storage, one as_strided view per name — the whole point of the storage/view split. Overlapping views would serve the same - bytes for two different weights.""" + bytes for two different weights. + """ server = server_factory() server.begin_sync(1) _publish(server, GI_A, GROUP_A, nbytes=64) @@ -259,7 +260,8 @@ class TestPublishAndRebuild: def test_the_default_lookahead_is_one(self): """1 is the sweet spot: group N+1 is gathered AND pullable while N is being pulled, with resident memory at its floor of 2 groups — the bound - the larger-model runs size against.""" + the larger-model runs size against. + """ assert DEFAULT_GATHER_LOOKAHEAD == 1 @@ -274,7 +276,8 @@ class TestFreeBarrier: def test_group_is_held_until_the_last_live_consumer_signals(self, server_factory): """Every live consumer signals every owner; releasing on the first - signal would drop storage other consumers are still reading.""" + signal would drop storage other consumers are still reading. + """ server = server_factory() server.begin_sync(3) _publish(server, GI_A, GROUP_A) @@ -301,7 +304,8 @@ class TestFreeBarrier: ): """A consumer with nothing to pull from a group signals it at sync start, which can precede the publish. The publish must then release the - group rather than wait for a signal that will never come again.""" + group rather than wait for a signal that will never come again. + """ server = server_factory() server.begin_sync(1) server.free_group(GI_A) @@ -324,7 +328,8 @@ class TestFreeBarrier: the server says the group is gone — and it hears that ONLY through ``wait_freed`` / ``end_sync``, never a publish return (a freed notice riding an unharvested async publish while the engine blocks in - ``wait_freed`` would wedge the loop).""" + ``wait_freed`` would wedge the loop). + """ server = server_factory() server.begin_sync(1) _publish(server, GI_A, GROUP_A) @@ -353,7 +358,8 @@ class TestFreeBarrier: class _CountingSource: """`iter_groups` yielding one CPU-tensor group per call, built lazily so a - group's tensors exist only once the engine's credit gate lets it gather.""" + group's tensors exist only once the engine's credit gate lets it gather. + """ def __init__(self, groups): self._groups = groups @@ -366,7 +372,8 @@ class _CountingSource: class _ResidencyDict(dict): """The engine's `_inflight`, instrumented: `max_len` is the high-water mark of groups whose CUDA-IPC export refs were alive at once — the trainer's - resident-memory bound in groups.""" + resident-memory bound in groups. + """ def __init__(self): super().__init__() @@ -380,7 +387,8 @@ class _ResidencyDict(dict): def _loop_engine(server, n_groups, *, lookahead): """A real engine wired straight to a real (non-Ray) server, with only the state `_run_gather_loop` reads. The `_rpc` seam dispatches inline; - `_publish_async` sees no `.remote` and runs inline too.""" + `_publish_async` sees no `.remote` and runs inline too. + """ groups = [(f"model.layers.{gi}.w",) for gi in range(n_groups)] e = ShardedRDTTrainerWeightTransferEngine.__new__( ShardedRDTTrainerWeightTransferEngine @@ -403,14 +411,16 @@ class TestGatherCredit: serveable immediately, and the ENGINE's loop stops gathering while more than `gather_lookahead` groups are unfreed. So at most `lookahead + 1` groups are resident — two at the default — while group N+1 is pulled the instant N's - pulls finish.""" + pulls finish. + """ @pytest.fixture def gather_engine(self, monkeypatch): """Neutralize the two CUDA touches in the export path so the REAL gather loop runs on CPU: `.cuda()` becomes identity, and the storage export ships the `(nbytes, ..., device, ...)` tuple the fixture's - stubbed `rebuild_cuda_tensor` expects.""" + stubbed `rebuild_cuda_tensor` expects. + """ monkeypatch.setattr(torch.Tensor, "cuda", lambda self: self) monkeypatch.setattr( trainer_mod, @@ -422,7 +432,8 @@ class TestGatherCredit: def test_publish_never_blocks(self, server_factory): """Publishing must not park on a credit: the whole point is that a gathered group becomes serveable immediately. Three publishes with no - frees, inline on this thread — a block here hangs the test.""" + frees, inline on this thread — a block here hangs the test. + """ server = server_factory(gather_lookahead=1) server.begin_sync(1) for gi in range(3): @@ -436,7 +447,8 @@ class TestGatherCredit: than 2 groups of gathered weights, no matter how the consumer paces — and the pipeline still overlaps, because the consumer here refuses to free group N until N+1 is already published (pullable). A gate an - off-by-one too tight deadlocks this test; too loose fails max_len.""" + off-by-one too tight deadlocks this test; too loose fails max_len. + """ server = server_factory(gather_lookahead=1) n_groups = 8 engine = gather_engine(server, n_groups, lookahead=1) @@ -478,7 +490,8 @@ class TestGatherCredit: ): """0 is the serialized baseline: the loop waits for group N to be freed before gathering N+1, so exactly ONE group is ever resident and nothing - overlaps. A floor of 1 on the bound would silently pipeline instead.""" + overlaps. A floor of 1 on the bound would silently pipeline instead. + """ server = server_factory(gather_lookahead=0) engine = gather_engine(server, 4, lookahead=0) @@ -503,8 +516,9 @@ class TestGatherCredit: def test_the_residency_bound_scales_with_the_lookahead( self, server_factory, gather_engine ): - """lookahead + 1, not a hardcoded 2: at lookahead=2 a free-nothing - consumer sees exactly 3 groups gathered before the loop parks.""" + """Lookahead + 1, not a hardcoded 2: at lookahead=2 a free-nothing + consumer sees exactly 3 groups gathered before the loop parks. + """ server = server_factory(gather_lookahead=2) engine = gather_engine(server, 8, lookahead=2) @@ -540,7 +554,8 @@ class TestGatherCredit: def test_a_gather_error_releases_a_blocked_wait_freed(self, server_factory): """Otherwise a trainer-side failure on another rank deadlocks this one inside its credit gate. It must RAISE, not return empty — an empty - return would spin the engine straight back into the wait.""" + return would spin the engine straight back into the wait. + """ server = server_factory() server.begin_sync(1) _publish(server, GI_A, GROUP_A) @@ -621,7 +636,8 @@ class TestBeginSync: def test_the_target_is_floored_at_one(self, server_factory): """A bare/zero call must not make every group free instantly (or divide - the barrier by zero).""" + the barrier by zero). + """ server = server_factory() server.begin_sync(0) assert server._live_count == 1 @@ -629,7 +645,8 @@ class TestBeginSync: def test_a_degraded_sync_lowers_the_target(self, server_factory): """FT degraded sync: the live count is the WHOLE degraded-sync mechanism on this side — the group releases after the live consumers - alone.""" + alone. + """ server = server_factory() server.begin_sync(2) # 4 provisioned, 2 alive _publish(server, GI_A, GROUP_A) @@ -653,7 +670,8 @@ class TestBeginSync: def test_a_straggler_signal_cannot_credit_the_next_sync(self, server_factory): """begin_sync resets the counts, so a signal that leaked past the sync boundary would over-credit — the reason the consumer drains its fired - signals before finishing. Pin the reset side of that contract.""" + signals before finishing. Pin the reset side of that contract. + """ server = server_factory() server.begin_sync(2) _publish(server, GI_A, GROUP_A) @@ -675,7 +693,8 @@ class TestBeginSync: self, server_factory ): """The packed layout repeats every sync — that is what makes caching the - destination views worth ~5ms per 384-spec group.""" + destination views worth ~5ms per 384-spec group. + """ server = server_factory() server._pack_dsts[("sentinel",)] = (0, []) server.begin_sync(1) @@ -686,7 +705,8 @@ class TestServedNamesGuard: """A producer serves only the names it holds: its stage's groups, and within them its own EP coordinate's experts plus the replicated names. Without this guard a misrouted pull would block forever in the cache wait - for a name this rank never publishes.""" + for a name this rank never publishes. + """ def test_a_pull_for_an_unserved_name_fails_loudly(self, server_factory): server = server_factory(served_names=list(GROUP_A)) @@ -702,7 +722,8 @@ class TestServedNamesGuard: def test_served_names_none_accepts_anything(self, server_factory): """Gather-to-all: every producer holds every group, so there is nothing - to guard. The call proceeds to the cache wait (not exercised here).""" + to guard. The call proceeds to the cache wait (not exercised here). + """ server = server_factory(served_names=None) assert server._served_names is None @@ -749,7 +770,8 @@ class TestStallWatchdog: def test_end_sync_fails_when_a_consumer_never_signals(self, server_factory): """The exact mid-sync death signature: the group is published, one of - the two live consumers that owed a signal is gone.""" + the two live consumers that owed a signal is gone. + """ server = server_factory(stall_timeout_s=0.3) server.begin_sync(2) _publish(server, GI_A, GROUP_A) @@ -763,7 +785,8 @@ class TestStallWatchdog: """The watchdog fires on the existing `set_gather_error` channel, so one stall unwinds the whole rank through one path rather than each waiter timing out separately — here both a consumer parked in the serve cache - wait and the engine parked in its credit gate.""" + wait and the engine parked in its credit gate. + """ server = server_factory(gather_lookahead=1, stall_timeout_s=0.3) server.begin_sync(1) _publish(server, GI_A, GROUP_A) @@ -796,7 +819,8 @@ class TestStallWatchdog: """A consumer that is slow but signaling must never trip it: the stamp is global to the producer, so a steady trickle of signals holds the whole rank open. Nine sequential publishes with a 0.3s timeout each take - longer in total than the timeout, and none may fire.""" + longer in total than the timeout, and none may fire. + """ server = server_factory(gather_lookahead=1, stall_timeout_s=0.3) server.begin_sync(1) for gi in range(9): @@ -809,7 +833,8 @@ class TestStallWatchdog: def test_begin_sync_resets_the_progress_stamp(self, server_factory): """Syncs are minutes apart. Without the reset the first publish of sync N+1 - would measure its stall from somewhere inside sync N and fire immediately.""" + would measure its stall from somewhere inside sync N and fire immediately. + """ server = server_factory(gather_lookahead=1, stall_timeout_s=0.3) server.begin_sync(1) _publish(server, GI_A, GROUP_A) @@ -831,7 +856,8 @@ class TestConcurrentPublishAndFree: self, server_factory ): """Every group published must end up freed and out of the cache, with - end_sync returning cleanly — the property whose violation is a hang.""" + end_sync returning cleanly — the property whose violation is a hang. + """ server = server_factory(gather_lookahead=2) server.begin_sync(1) groups = [(gi, (f"model.layers.{gi}.w",)) for gi in range(12)] @@ -866,7 +892,8 @@ class TestConcurrentPublishAndFree: ): """All live consumers signaling at once: the release must happen on the last signal and only once, or the engine drops IPC refs twice (or - never).""" + never). + """ server = server_factory() server.begin_sync(8) _publish(server, GI_A, GROUP_A) @@ -893,7 +920,8 @@ class TestFakeServerAgreesWithTheRealOne: """`test_sharded_rdt_trainer.py`'s `_FakeProducerServer` is a second, independent model of this protocol that the engine-side tests assert on. Pin the two against each other so the fake cannot silently drift from the - semantics the real server enforces.""" + semantics the real server enforces. + """ @staticmethod def _fake(): @@ -924,7 +952,8 @@ class TestExportRing: """The packed export path. Its one hard invariant is that the trainer packs a group exactly the way ``publish_group`` rebuilds it — a layout disagreement would serve the wrong bytes with nothing downstream to catch - it.""" + it. + """ @staticmethod def _engine(): @@ -1016,7 +1045,8 @@ class TestExportRingSlotSafety: ): """Free the later group of each pair first and hold the earlier one live until the NEXT group is gathered -- the window a modular slot counter got - wrong (at lookahead=1 there are 2 slots, so group 2 took group 0's).""" + wrong (at lookahead=1 there are 2 slots, so group 2 took group 0's). + """ server = server_factory(gather_lookahead=1) n_groups = 8 engine = gather_engine(server, n_groups, lookahead=1) @@ -1140,7 +1170,8 @@ class TestSharedServeSlots: def test_one_deployment_never_shares(self, sharing_server): """With the width equal to the live count every group is a singleton, so - each call packs for itself out of its own ring.""" + each call packs for itself out of its own ring. + """ server = sharing_server(workers_per_replica=2) server.begin_sync(2, [0, 1]) _publish(server, GI_A, GROUP_A) @@ -1220,7 +1251,8 @@ class TestSharedServeSlots: def test_a_seqless_call_is_refused(self, sharing_server): """The slot comes from ``seq``, so a caller that sends none cannot be - served safely.""" + served safely. + """ server = sharing_server(workers_per_replica=2) server.begin_sync(1, [0]) _publish(server, GI_A, GROUP_A) @@ -1230,7 +1262,8 @@ class TestSharedServeSlots: def test_a_failed_pack_reaches_every_sharer(self, sharing_server): """Only the packer touches the buffer, so its failure has to be - republished: a waiter must raise rather than read a half-written slot.""" + republished: a waiter must raise rather than read a half-written slot. + """ server = sharing_server(workers_per_replica=2) server.begin_sync(4, [0, 1, 2, 3]) _publish(server, GI_A, GROUP_A) @@ -1254,7 +1287,8 @@ class TestSharedServeSlots: def test_a_degraded_sync_narrows_the_rendezvous(self, sharing_server): """A dead deployment must not be waited for: a survivor left alone in - its group packs for itself instead of blocking on the corpse.""" + its group packs for itself instead of blocking on the corpse. + """ server = sharing_server(workers_per_replica=2, stall_timeout_s=30.0) server.begin_sync(3, [0, 1, 3]) # consumer 2 died _publish(server, GI_A, GROUP_A) @@ -1268,7 +1302,8 @@ class TestSharedServeSlots: def test_a_sharer_that_never_arrives_trips_the_watchdog(self, sharing_server): """A consumer that dies INSIDE the sync window has no detector, exactly - as for the free barrier; the rendezvous waits on the same channel.""" + as for the free barrier; the rendezvous waits on the same channel. + """ server = sharing_server(workers_per_replica=2, stall_timeout_s=0.3) server.begin_sync(4, [0, 1, 2, 3]) _publish(server, GI_A, GROUP_A) @@ -1280,7 +1315,8 @@ class TestSharedServeSlots: class TestSharedReservation: """``reserve_serve_buffer`` under sharing: one ring per GROUP, and the - init-time check that the group's members pull the same chunks.""" + init-time check that the group's members pull the same chunks. + """ def test_matching_plans_reserve_one_ring_between_them(self, sharing_server): server = sharing_server(workers_per_replica=2) @@ -1302,7 +1338,8 @@ class TestSharedReservation: def test_mismatched_plans_fail_at_reservation(self, sharing_server): """Sharing rests on the sharers pulling the same chunks. Checked here, where it names both consumers, rather than at serve time, where it is a - rendezvous nobody completes.""" + rendezvous nobody completes. + """ server = sharing_server(workers_per_replica=2) server.reserve_serve_buffer(0, 1024, "plan-a") @@ -1311,7 +1348,8 @@ class TestSharedReservation: def test_different_groups_may_differ(self, sharing_server): """Only the members of ONE group have to agree; worker 0 and worker 1 of - the same deployment hold different slices and never share a slot.""" + the same deployment hold different slices and never share a slot. + """ server = sharing_server(workers_per_replica=2) server.reserve_serve_buffer(0, 1024, "plan-a") server.reserve_serve_buffer(1, 1024, "plan-b") @@ -1320,7 +1358,8 @@ class TestSharedReservation: def test_a_missing_digest_skips_the_check(self, sharing_server): """The digest is optional on the wire, so an engine that does not send - one still reserves.""" + one still reserves. + """ server = sharing_server(workers_per_replica=2) server.reserve_serve_buffer(0, 1024, None) server.reserve_serve_buffer(2, 1024, None) diff --git a/tests/distributed/test_sharded_rdt_trainer.py b/tests/distributed/test_sharded_rdt_trainer.py index 9339e503c4..df9628cc85 100644 --- a/tests/distributed/test_sharded_rdt_trainer.py +++ b/tests/distributed/test_sharded_rdt_trainer.py @@ -34,7 +34,8 @@ from vllm.distributed.weight_transfer.sharded_rdt_trainer import ( class _ListSource(WeightSource): """A WeightSource over an explicit ordered (name, cpu-tensor) list, so the - sharded-RDT group/order logic can be tested without a real model.""" + sharded-RDT group/order logic can be tested without a real model. + """ def __init__(self, pairs): self._pairs = list(pairs) @@ -54,7 +55,8 @@ class _FakeProducerServer: per-group barrier: publish/free are keyed by GROUP INDEX, signals count to the ``begin_sync`` live total, publish_group returns nothing, and freed groups flow back only through wait_freed / end_sync (see - test_sharded_rdt_producer.TestFakeServerAgreesWithTheRealOne).""" + test_sharded_rdt_producer.TestFakeServerAgreesWithTheRealOne). + """ def __init__(self, auto_free=True): self.order: list[str] = [] @@ -98,7 +100,8 @@ class _FakeProducerServer: must already have a banked credit when the gate asks (auto_free, or a test's own free_group) — anything else is the deadlock the real watchdog would kill, so fail loudly. Not appended to ``order``: it is - pacing, not a lifecycle milestone.""" + pacing, not a lifecycle milestone. + """ assert self._pending_freed, ( "wait_freed with nothing freed: the gather loop would deadlock" ) @@ -121,7 +124,8 @@ def _rdt_engine_with_fake_server( ): """Build a ShardedRDTTrainerWeightTransferEngine wired to an in-process fake server (no Ray, no CUDA IPC): bypass trainer_init's spawn, set the - group-major metadata, and route _rpc to the fake.""" + group-major metadata, and route _rpc to the fake. + """ import vllm.distributed.weight_transfer.sharded_rdt_trainer as mod # reduce_tensor needs CUDA; the fake server never rebuilds, so stub it. @@ -321,7 +325,8 @@ def test_sharded_rdt_send_weights_surfaces_update_error(monkeypatch): class _OwnedSource(_ListSource): """A source holding only some groups' names, like a pipeline-parallel rank. - Takes group indices for convenience and declares the names inside them.""" + Takes group indices for convenience and declares the names inside them. + """ def __init__(self, pairs, owned_group_idx): super().__init__(pairs) @@ -387,7 +392,8 @@ def test_sharded_rdt_owned_group_order_mismatch_raises(monkeypatch): ) def test_sharded_rdt_begin_sync_carries_the_live_count(monkeypatch): """The free barrier's target is one integer per sync: the live consumer - count, defaulting to the whole provisioned fleet.""" + count, defaulting to the whole provisioned fleet. + """ server = _FakeProducerServer(auto_free=True) engine = _rdt_engine_with_fake_server( _rdt_source_two_layers(), @@ -405,7 +411,8 @@ def test_sharded_rdt_begin_sync_carries_the_live_count(monkeypatch): class TestLiveCountPlumbing: """``send_weights(live_consumer_ids)`` -> the barrier target. The provisioned geometry is frozen; a degraded sync only lowers the target and - narrows each slot-sharing group's rendezvous.""" + narrows each slot-sharing group's rendezvous. + """ @staticmethod def _engine(num_consumers, world=2, rank=0): @@ -438,7 +445,8 @@ class TestLiveCountPlumbing: def test_which_consumers_died_does_not_matter_only_how_many(self): """The whole point of the barrier: no routed per-producer targets, so - the identity of the dead consumer is irrelevant to the producers.""" + the identity of the dead consumer is irrelevant to the producers. + """ counts = [] for live in ([0, 1, 2, 3], [4, 5, 6, 7], [0, 2, 4, 6]): engine, got = self._engine(num_consumers=8) @@ -449,7 +457,8 @@ class TestLiveCountPlumbing: def test_the_live_ids_travel_with_the_count(self): """The count sizes the free barrier, the ids size the slot-sharing rendezvous, so a producer that shares slots can tell WHICH consumers it - is still waiting for. They must describe the same set.""" + is still waiting for. They must describe the same set. + """ engine, got = self._engine(num_consumers=8) engine.send_weights([6, 0, 2, 2]) ((count, ids),) = got @@ -517,7 +526,8 @@ def test_sharded_rdt_rejects_a_rank_holding_nothing(monkeypatch): def test_sharded_rdt_rejects_a_name_no_rank_holds(monkeypatch): """Every name must be held somewhere or it can never be served — caught - when the holdings are transposed, naming the orphan.""" + when the holdings are transposed, naming the orphan. + """ with pytest.raises(ValueError, match="no trainer rank holds"): _rdt_engine_with_fake_server( _OwnedSource(_rdt_source_two_layers()._pairs, [0, 1]), @@ -531,7 +541,8 @@ def test_sharded_rdt_rejects_a_name_no_rank_holds(monkeypatch): def test_sharded_rdt_rejects_metadata_disagreement_across_ranks(monkeypatch): """Only the sender's metadata reaches the consumers, so a rank describing - just its own share must fail loudly rather than silently drop the rest.""" + just its own share must fail loudly rather than silently drop the rest. + """ engine = _rdt_engine_with_fake_server( _OwnedSource(_rdt_source_two_layers()._pairs, [1, 2]), is_sender=False, @@ -554,7 +565,8 @@ def _serve_ring_server(src_name, src): """A producer server with one cached tensor and a pre-seeded serve ring, so a pull needs no Ray and no NIXL registration. Returns (server, serve) where ``serve(chain)`` packs one spec into the SAME ring slot every time — which is - what puts the destination-view cache, and only it, under test.""" + what puts the destination-view cache, and only it, under test. + """ from vllm.distributed.weight_transfer.sharded_rdt_trainer import ( _RDTProducerServer, ) diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py index ef7f523b78..2274f15d97 100644 --- a/tests/distributed/test_shm_broadcast.py +++ b/tests/distributed/test_shm_broadcast.py @@ -43,6 +43,7 @@ def distributed_run(fn, world_size, timeout=60): fn: Function to run in each process world_size: Number of processes to spawn timeout: Maximum time in seconds to wait for processes (default: 60) + """ number_of_processes = world_size processes = [] @@ -418,7 +419,8 @@ def test_tensor_broadcast(): def _dumps_oob(obj) -> tuple[bytes, list]: """Pickle `obj` the same way `MessageQueue.enqueue` does: tensor - dispatch table + out-of-band buffers >= 1MiB.""" + dispatch table + out-of-band buffers >= 1MiB. + """ buffers = [] def callback(buf: pickle.PickleBuffer) -> bool: @@ -489,7 +491,8 @@ def test_tensor_pickle_roundtrip(case: str): @pytest.mark.parametrize("case", ["cuda", "requires_grad", "conj"]) def test_reduce_tensor_fallback(case: str): """Tensors the zero-copy reducer can't safely alias must fall back to - torch's default reduction.""" + torch's default reduction. + """ if case == "cuda": if not torch.cuda.is_available(): pytest.skip("requires CUDA") @@ -634,11 +637,9 @@ def test_acquire_read_releases_slot_when_reader_raises(): def test_warning_logs(caplog_vllm): - """ - Test that warning logs are emitted at VLLM_RINGBUFFER_WARNING_INTERVAL intervals + """Test that warning logs are emitted at VLLM_RINGBUFFER_WARNING_INTERVAL intervals when indefinite=False, and are not emitted when indefinite=True. """ - # Patch the warning log interval to every 1 ms during reads with mock.patch( "vllm.distributed.device_communicators.shm_broadcast.VLLM_RINGBUFFER_WARNING_INTERVAL", @@ -754,7 +755,8 @@ def test_remote_subscribe_addr_unique_concurrent_writers( Pre-fix, the writer probed a port with get_open_port() and bound it afterwards; pinning the probe to one free port makes every writer bind the same port and fail deterministically on that code path, - while the late-binding implementation never consults the probe.""" + while the late-binding implementation never consults the probe. + """ from vllm.distributed.device_communicators import shm_broadcast colliding_port = get_open_port() diff --git a/tests/distributed/test_shm_buffer.py b/tests/distributed/test_shm_buffer.py index e56c71b1c2..6f61c28322 100644 --- a/tests/distributed/test_shm_buffer.py +++ b/tests/distributed/test_shm_buffer.py @@ -12,20 +12,20 @@ from vllm.distributed.device_communicators.shm_object_storage import ( class TestSingleWriterShmRingBuffer(unittest.TestCase): - """Test suite for the ring buffer implementation""" + """Test suite for the ring buffer implementation.""" def setUp(self): - """Set up test fixtures""" + """Set up test fixtures.""" self.buffer_size = 4096 self.ring_buffer = None def tearDown(self): - """Clean up after tests""" + """Clean up after tests.""" if self.ring_buffer: self.ring_buffer.close() def test_buffer_opening(self): - """Test opening an existing buffer""" + """Test opening an existing buffer.""" # First create a buffer self.ring_buffer = SingleWriterShmRingBuffer( data_buffer_size=self.buffer_size, create=True @@ -39,7 +39,7 @@ class TestSingleWriterShmRingBuffer(unittest.TestCase): ) def test_buffer_access(self): - """Test accessing allocated buffers""" + """Test accessing allocated buffers.""" self.ring_buffer = SingleWriterShmRingBuffer( data_buffer_size=self.buffer_size, create=True ) @@ -61,7 +61,7 @@ class TestSingleWriterShmRingBuffer(unittest.TestCase): self.assertEqual(read_id, monotonic_id) def test_memory_error_on_full_buffer(self): - """Test that MemoryError is raised when buffer is full""" + """Test that MemoryError is raised when buffer is full.""" small_buffer_size = 200 self.ring_buffer = SingleWriterShmRingBuffer( data_buffer_size=small_buffer_size, create=True @@ -76,7 +76,7 @@ class TestSingleWriterShmRingBuffer(unittest.TestCase): self.ring_buffer.allocate_buf(1) # Would exceed buffer capacity def test_allocation_and_free(self): - """Test allocation and freeing of buffers""" + """Test allocation and freeing of buffers.""" small_buffer_size = 200 self.ring_buffer = SingleWriterShmRingBuffer( data_buffer_size=small_buffer_size, create=True @@ -96,7 +96,7 @@ class TestSingleWriterShmRingBuffer(unittest.TestCase): self.assertEqual(freed_ids[0], i) def test_clear_buffer(self): - """Test clearing the buffer""" + """Test clearing the buffer.""" self.ring_buffer = SingleWriterShmRingBuffer( data_buffer_size=self.buffer_size, create=True ) @@ -180,7 +180,7 @@ class TestSingleWriterShmRingBuffer(unittest.TestCase): def main(): - """Main function demonstrating usage and running tests""" + """Main function demonstrating usage and running tests.""" print("=== SingleWriterShmRingBuffer Test Suite ===\n") # Run unit tests diff --git a/tests/distributed/test_utils.py b/tests/distributed/test_utils.py index 031937d7f5..fad1d2a006 100644 --- a/tests/distributed/test_utils.py +++ b/tests/distributed/test_utils.py @@ -31,7 +31,8 @@ class _CUDADeviceCountStatelessTestActor: def test_cuda_device_count_stateless(): """Test that cuda_device_count_stateless changes return value if - CUDA_VISIBLE_DEVICES is changed.""" + CUDA_VISIBLE_DEVICES is changed. + """ if current_platform.is_rocm(): pytest.skip("Skip for ROCm because Ray uses HIP_VISIBLE_DEVICES.") actor = _CUDADeviceCountStatelessTestActor.options( # type: ignore diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index 1353def738..ced71ead76 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -1166,7 +1166,8 @@ class RecordingClient: def _module_with(*pairs): """A tiny nn.Module exposing the given (name, tensor) pairs as parameters, - so trainer tests can build a ModuleSource without a real model.""" + so trainer tests can build a ModuleSource without a real model. + """ module = torch.nn.Module() for name, tensor in pairs: module.register_parameter(name, torch.nn.Parameter(tensor, requires_grad=False)) @@ -1287,7 +1288,8 @@ class TestModuleSource: def test_metadata_agrees_with_iteration(self): """The two channels must line up element-for-element: engines declare - the round from `metadata()` and then send what iteration yields.""" + the round from `metadata()` and then send what iteration yields. + """ source = ModuleSource( _module_with(("w", torch.zeros(2, 3)), ("b", torch.zeros(3))) ) @@ -1301,11 +1303,13 @@ class TestModuleSource: class TestWeightSourceGroupContract: """`groups()` / `iter_groups()` on the WeightSource ABC. Groups are what backends gather and free by, so the default must agree with - `layerwise_groups` over `metadata()`, restricted to what this rank holds.""" + `layerwise_groups` over `metadata()`, restricted to what this rank holds. + """ class _Source(WeightSource): """Minimal source over an ordered (name, tensor) list, optionally holding - only some names (in which case it iterates only their groups).""" + only some names (in which case it iterates only their groups). + """ def __init__(self, names, held=None, reverse=False): self._pairs = [(n, torch.full((2,), float(i))) for i, n in enumerate(names)] @@ -1360,7 +1364,8 @@ class TestWeightSourceGroupContract: def test_groups_order_follows_metadata_not_the_declaration(self): """The declaration is a SET of names; the groups it selects still come out in metadata order, so each pairs with the right ``iter_groups()`` - batch however the source listed them.""" + batch however the source listed them. + """ names = ["embed.w", "model.layers.0.a", "model.layers.1.a", "norm.w"] source = self._source( names, held=["model.layers.1.a", "model.layers.0.a", "model.layers.1.a"] @@ -1394,14 +1399,16 @@ class TestWeightSourceGroupContract: def test_out_of_order_iteration_raises(self): """Materializing is usually a collective, so a rank that iterates out of - order deadlocks its peers -- fail loudly instead.""" + order deadlocks its peers -- fail loudly instead. + """ source = self._source(["model.layers.0.a", "model.layers.0.b"], reverse=True) with pytest.raises(RuntimeError, match="iteration order must match"): list(source.iter_groups()) def test_a_source_may_override_iter_groups(self): """The extension point: materialize a whole group in one step instead of - one generator resume per tensor.""" + one generator resume per tensor. + """ calls = [] class _Batched(TestWeightSourceGroupContract._Source): @@ -1422,7 +1429,8 @@ class TestDeferredProcessingContract: caller that takes over the update tail (running its own `finalize_layerwise_reload` instead of going through `finish_weight_update`) reads the flag and calls the method. Both must be answerable on any engine, or - that caller ends up reaching through a getattr.""" + that caller ends up reaching through a getattr. + """ def _engines(self): return { @@ -1436,7 +1444,8 @@ class TestDeferredProcessingContract: def test_every_engine_can_be_drained(self): """The default is a no-op, so a caller never has to check whether the - method exists before calling it.""" + method exists before calling it. + """ for name, cls in self._engines().items(): assert callable(cls.drain_pending), name @@ -1530,7 +1539,8 @@ class TestTrainerEngineBase: ) def test_ipc_trainer_send_weights_drives_client_in_order(): """send_weights issues start -> update -> finish and ships per-round metadata; - the packed wire param rides the init info, not the per-round update_info.""" + the packed wire param rides the init info, not the per-round update_info. + """ client = RecordingClient() engine = IPCTrainerWeightTransferEngine( client=client, @@ -1549,7 +1559,8 @@ def test_ipc_trainer_send_weights_drives_client_in_order(): def test_ipc_trainer_init_ships_packed_to_worker(): """trainer_init drives the inference-side init handshake and propagates the - must-agree `packed` flag to the worker.""" + must-agree `packed` flag to the worker. + """ if torch.accelerator.device_count() < 1: pytest.skip("Need at least 1 GPU (CUDA IPC handles).") @@ -1570,7 +1581,8 @@ def test_ipc_trainer_init_ships_packed_to_worker(): def test_nccl_trainer_init_ships_worker_init_info(monkeypatch): """The sender's trainer_init drives the inference-side init handshake with the worker-shaped init info (rank_offset=1) while opening its own endpoint, - and propagates the must-agree wire params to the worker.""" + and propagates the must-agree wire params to the worker. + """ import vllm.distributed.weight_transfer.nccl_engine as nccl_engine_mod # Bypass the real NCCL rendezvous. @@ -1611,7 +1623,8 @@ def test_nccl_trainer_init_ships_worker_init_info(monkeypatch): def test_nccl_worker_learns_wire_params_from_init_handshake(monkeypatch): """The worker engine reads packed + buffer geometry from the trainer-supplied init info at the handshake, not from the config or the - per-round update info.""" + per-round update info. + """ import vllm.distributed.weight_transfer.nccl_engine as nccl_engine_mod monkeypatch.setattr( @@ -1644,7 +1657,8 @@ def test_nccl_worker_learns_wire_params_from_init_handshake(monkeypatch): def test_nccl_trainer_init_non_sender_skips_rendezvous_and_client(): """Non-sender trainer ranks build an engine without opening an endpoint or - touching the client; they only join the collectives in send_weights.""" + touching the client; they only join the collectives in send_weights. + """ client = RecordingClient() engine = WeightTransferTrainerFactory.trainer_init( init_info=NCCLTrainerInitInfo( @@ -1674,7 +1688,8 @@ def test_nccl_trainer_init_non_sender_skips_rendezvous_and_client(): def test_nccl_trainer_send_weights_drives_client_in_order(): """send_weights issues start -> update -> finish and ships per-round metadata; the packed wire params ride the init handshake, not the - per-round update_info.""" + per-round update_info. + """ client = RecordingClient() engine = NCCLTrainerWeightTransferEngine( client=client, @@ -1695,7 +1710,8 @@ def test_nccl_trainer_send_weights_drives_client_in_order(): class _ScriptedSource(WeightSource): """Declares `meta` but yields whatever `pairs` says — used to drive the - metadata/iteration agreement checks.""" + metadata/iteration agreement checks. + """ def __init__(self, meta, pairs): self._meta = meta @@ -1731,7 +1747,8 @@ def test_nccl_trainer_init_requires_source(): def test_nccl_trainer_send_weights_rejects_reordered_source(monkeypatch): """The worker sizes its buffers (and cuts packed chunks) from metadata(), so - iteration disagreeing with it must raise rather than corrupt the stream.""" + iteration disagreeing with it must raise rather than corrupt the stream. + """ meta = [ ParamMeta("w", torch.float32, (4,)), ParamMeta("b", torch.float32, (2,)), @@ -1745,7 +1762,8 @@ def test_nccl_trainer_send_weights_rejects_reordered_source(monkeypatch): def test_nccl_trainer_send_weights_rejects_dtype_disagreement(monkeypatch): """A source that declares one wire dtype and materializes another would make - the two sides disagree on every byte offset.""" + the two sides disagree on every byte offset. + """ meta = [ParamMeta("w", torch.float32, (4,))] engine = _mock_group_engine( _ScriptedSource(meta, [("w", torch.zeros(4, dtype=torch.bfloat16))]), @@ -1772,7 +1790,8 @@ def test_nccl_trainer_send_weights_rejects_truncated_source(monkeypatch): def test_nccl_trainer_send_weights_broadcasts_contiguous(monkeypatch): """NCCL sends numel elements from data_ptr(), so a non-contiguous view must - be linearized first or the worker receives unrelated memory.""" + be linearized first or the worker receives unrelated memory. + """ base = torch.arange(6, dtype=torch.float32).reshape(2, 3) view = base.t() # non-contiguous meta = [ParamMeta("w", torch.float32, tuple(view.shape))] @@ -1842,7 +1861,8 @@ def _sparse_patch(device: str = "cpu") -> SparseWeightPatch: def test_sparse_nccl_trainer_init_ships_worker_init_info(monkeypatch): """The sender's trainer_init drives the init handshake with the worker-shaped init info; sparse ships no packed wire params, so the worker - keeps its unpacked defaults. Sparse takes no `source`.""" + keeps its unpacked defaults. Sparse takes no `source`. + """ import vllm.distributed.weight_transfer.sparse_nccl_engine as sparse_mod monkeypatch.setattr(sparse_mod, "open_trainer_endpoint", lambda info: MagicMock()) @@ -1895,7 +1915,8 @@ def test_sparse_nccl_trainer_send_weights_drives_client_in_order(monkeypatch): def test_sparse_nccl_trainer_send_weights_empty_round_is_noop(): """A round with no patches must not touch the client (an empty sparse - update info is invalid by construction).""" + update info is invalid by construction). + """ client = RecordingClient() engine = SparseNCCLTrainerWeightTransferEngine(client=client) engine.model_update_group = MagicMock() @@ -1931,7 +1952,8 @@ def test_sparse_nccl_trainer_rejects_source(): def test_sparse_nccl_trainer_validates_patch_before_any_rpc(): """Malformed patches must fail on the trainer, before start_weight_update: the worker's own checks only run once the broadcasts are already in flight, - where a size mismatch wedges both sides instead of raising.""" + where a size mismatch wedges both sides instead of raising. + """ client = RecordingClient() engine = SparseNCCLTrainerWeightTransferEngine(client=client) engine.model_update_group = MagicMock() diff --git a/tests/distributed/test_weight_transfer_nccl_uid.py b/tests/distributed/test_weight_transfer_nccl_uid.py index 18a6442a2a..5820137604 100644 --- a/tests/distributed/test_weight_transfer_nccl_uid.py +++ b/tests/distributed/test_weight_transfer_nccl_uid.py @@ -403,7 +403,8 @@ def trainer_broadcast_tensor_torch_free( def _run_uid_transfer(trainer_task) -> None: """Mint a unique id, run the worker and the given trainer peer concurrently, and assert the tensor round-trips. Both ranks must enter init together -- - there is no store barrier on the UID path.""" + there is no store barrier on the UID path. + """ from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary _init_ray_for_weight_transfer() @@ -450,7 +451,8 @@ def _run_uid_transfer(trainer_task) -> None: ) def test_nccl_weight_transfer_between_processes_uid(): """Weight transfer over a pre-shared ncclUniqueId (no TCPStore), with both - ranks on vLLM's ``PyNcclCommunicator``.""" + ranks on vLLM's ``PyNcclCommunicator``. + """ _run_uid_transfer(trainer_broadcast_tensor_uid) @@ -460,6 +462,7 @@ def test_nccl_weight_transfer_between_processes_uid(): ) def test_nccl_weight_transfer_torch_free_trainer(): """UID weight transfer against a foreign peer that uses no torch and no vLLM - communicator, only a cupy buffer and raw ``ncclBroadcast``.""" + communicator, only a cupy buffer and raw ``ncclBroadcast``. + """ pytest.importorskip("cupy") _run_uid_transfer(trainer_broadcast_tensor_torch_free) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index e583798746..9adfedc9f4 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -394,8 +394,7 @@ def test_media_io_kwargs_parser(arg, expected): ], ) def test_optimization_level(args, expected): - """ - Test space-separated optimization levels (-O 1, -O 2, -O 3) map to + """Test space-separated optimization levels (-O 1, -O 2, -O 3) map to optimization_level. """ parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) @@ -414,9 +413,7 @@ def test_optimization_level(args, expected): ], ) def test_mode_parser(args, expected): - """ - Test compilation config modes (-cc.mode=int) map to compilation_config. - """ + """Test compilation config modes (-cc.mode=int) map to compilation_config.""" parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) parsed_args = parser.parse_args(args) assert parsed_args.compilation_config.mode == expected @@ -573,7 +570,8 @@ def test_attention_config(): def test_multi_node_world_size_includes_pcp(monkeypatch): """PCP expands the process world size, so the --nnodes divisibility check must include it. Without this, TP=1/PCP=2 over 2 nodes computes a world - size of 1 and the launch is rejected before the engine starts.""" + size of 1 and the launch is rejected before the engine starts. + """ import vllm.config.vllm # PCP requires the V2 model runner, which is gated on Triton. @@ -799,7 +797,8 @@ def test_expand_json_human_readable_numbers(input_json, expected_json): ) def test_cloud_storage_uri_skips_get_model_path(uri, monkeypatch): """Cloud storage URIs should not be passed to get_model_path() - when HF_HUB_OFFLINE=1, as they are not valid HF repo IDs.""" + when HF_HUB_OFFLINE=1, as they are not valid HF repo IDs. + """ import huggingface_hub monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True) @@ -811,7 +810,8 @@ def test_cloud_storage_uri_skips_get_model_path(uri, monkeypatch): def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch): """Cloud storage tokenizer URI should not be passed to - get_model_path() when HF_HUB_OFFLINE=1.""" + get_model_path() when HF_HUB_OFFLINE=1. + """ import huggingface_hub monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True) @@ -973,7 +973,8 @@ class TestDpDeviceIdSharding: def test_dp_supervisor_device_ids_stay_env_relative(self): """Regression test: the DP supervisor must pass env-relative indices, not physical IDs, because each child re-resolves --device-ids - against its inherited device-control env var.""" + against its inherited device-control env var. + """ import argparse from vllm.entrypoints.launchers.dp_supervisor import _build_device_ids @@ -1000,7 +1001,8 @@ class TestDpDeviceIdSharding: def test_dp_rank_shards_user_assigned_gpu_ids(self): """get_physical_gpu_ids_for_local_dp_rank slices the user-provided - --device-ids list instead of recomputing from the env var.""" + --device-ids list instead of recomputing from the env var. + """ from vllm.platforms import current_platform from vllm.v1.engine.utils import get_physical_gpu_ids_for_local_dp_rank diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index 0c986f1291..63e44eaceb 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -241,7 +241,8 @@ class TestToolResultContent: self, tool_result_content ) -> AnthropicMessagesRequest: """Build a request with assistant tool_use followed by user - tool_result.""" + tool_result. + """ return _make_request( [ { @@ -326,7 +327,8 @@ class TestToolResultContent: def test_tool_result_with_text_and_image(self): """Mixed text+image tool_result: text in tool msg, image in user - msg.""" + msg. + """ request = self._make_tool_result_request( [ {"type": "text", "text": "Here is the screenshot"}, @@ -423,7 +425,8 @@ class TestToolResultContent: class TestAttributionHeaderStripping: def test_billing_header_stripped_from_system(self): """Claude Code's x-anthropic-billing-header block should be - stripped to preserve prefix caching.""" + stripped to preserve prefix caching. + """ request = _make_request( [{"role": "user", "content": "Hello"}], system=[ @@ -480,7 +483,7 @@ class TestThinkingBlockConversion: """ def test_thinking_plus_text_in_assistant_message(self): - """thinking + text → reasoning field + plain-string content.""" + """Thinking + text → reasoning field + plain-string content.""" request = _make_request( [ {"role": "user", "content": "Write me some code."}, @@ -542,7 +545,7 @@ class TestThinkingBlockConversion: assert asst.get("content") is None def test_thinking_plus_tool_use_in_assistant_message(self): - """thinking + tool_use: reasoning field set, tool_calls populated.""" + """Thinking + tool_use: reasoning field set, tool_calls populated.""" request = _make_request( [ {"role": "user", "content": "What is 2+2?"}, @@ -1147,7 +1150,8 @@ class TestMessageStreamConverterToolUseContentBuffering: async def test_buffered_content_flushed_on_done_without_usage_chunk(self): """Content buffered during tool_use must be emitted even if the stream jumps straight from finish_reason to [DONE], skipping the - empty-choices usage chunk.""" + empty-choices usage chunk. + """ async def sse_input(): yield _make_stream_chunk( @@ -1251,7 +1255,8 @@ class TestStreamingCacheUsageSemantics: async def test_streaming_cache_fields_absent_then_populated(self): """First chunk lacks prompt_tokens_details (vLLM contract); message_start omits cache fields. The final chunk carries - prompt_tokens_details, so message_delta carries resolved values.""" + prompt_tokens_details, so message_delta carries resolved values. + """ async def sse_input(): yield _make_stream_chunk( @@ -1296,7 +1301,8 @@ class TestStreamingCacheUsageSemantics: @pytest.mark.asyncio async def test_streaming_no_cache_hit(self): """When the final chunk reports cached_tokens=0, message_delta carries - cache fields = 0 (cache miss); message_start still omits them.""" + cache fields = 0 (cache miss); message_start still omits them. + """ async def sse_input(): yield _make_stream_chunk( @@ -1337,7 +1343,8 @@ class TestStreamingCacheUsageSemantics: @pytest.mark.asyncio async def test_streaming_no_prompt_tokens_details_at_all(self): """If --enable-prompt-tokens-details is off, no chunk carries cache - info; both message_start and message_delta omit cache fields.""" + info; both message_start and message_delta omit cache fields. + """ async def sse_input(): yield _make_stream_chunk( @@ -1458,7 +1465,8 @@ class TestMessagesFullConverter: class TestCacheSalt: def test_cache_salt_passed_through(self): """cache_salt on the Anthropic request reaches the converted - ChatCompletionRequest so prefix-cache isolation works via /v1/messages.""" + ChatCompletionRequest so prefix-cache isolation works via /v1/messages. + """ request = _make_request( [{"role": "user", "content": "Hello"}], cache_salt="tenant-abc-secret-salt", @@ -1663,7 +1671,8 @@ class TestClientErrorResponses: @staticmethod def _conversion_error() -> ValidationError: """A real pydantic ValidationError like the one ChatCompletionRequest - construction raises when Anthropic input violates the OpenAI schema.""" + construction raises when Anthropic input violates the OpenAI schema. + """ class _StubRequest(BaseModel): stop: Annotated[list[str], Field(max_length=4)] | None = None @@ -1674,7 +1683,8 @@ class TestClientErrorResponses: def test_validation_error_returns_bad_request(self): """A pydantic ValidationError during Anthropic->OpenAI conversion is - surfaced as a 400 BadRequestError, not a 500.""" + surfaced as a 400 BadRequestError, not a 500. + """ handler = MagicMock(spec=AnthropicServingMessages) handler.create_messages.side_effect = self._conversion_error() diff --git a/tests/entrypoints/cohere/test_serving_streaming.py b/tests/entrypoints/cohere/test_serving_streaming.py index 8c2e8623ad..8723c6ecc7 100644 --- a/tests/entrypoints/cohere/test_serving_streaming.py +++ b/tests/entrypoints/cohere/test_serving_streaming.py @@ -87,7 +87,8 @@ def _make_chunk( citations: list[Any] | None = None, ) -> str: """Build the ``data: {...}\\n\\n`` SSE frame the production code - consumes.""" + consumes. + """ delta: dict[str, Any] = {} if role is not None: delta["role"] = role @@ -135,7 +136,8 @@ async def _stream_from(items: list[str]) -> AsyncGenerator[str, None]: async def _drain(serving: CohereServingChatV2, items: list[str]) -> list[str]: """Drive ``_chat_completion_stream_to_v2`` over ``items`` and collect - the emitted SSE frames.""" + the emitted SSE frames. + """ request = CohereChatV2Request( model="m", messages=[{"role": "user", "content": "hi"}], stream=True ) diff --git a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py index 668f115448..7f2f771107 100644 --- a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py +++ b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py @@ -50,7 +50,8 @@ def update_admission_stats_worker(listen_address, sock, args, client_config): # importable by `multiprocessing.spawn` (no closures, no nesting). def defer_addresses_stub_worker(listen_address, sock, args, client_config): """Bind ROUTER/PULL with a kernel-assigned port, report the actual - endpoints back via the pipe, then exit.""" + endpoints back via the pipe, then exit. + """ ctx = zmq.Context() try: in_sock = make_zmq_socket( @@ -357,7 +358,8 @@ def test_external_process_monitoring(api_server_args): def test_gather_actual_addresses_end_to_end(): """Each child binds ROUTER/PULL with a kernel-picked port and reports the bound endpoints back via its per-child pipe; the manager surfaces - them via :py:meth:`gather_actual_addresses`.""" + them via :py:meth:`gather_actual_addresses`. + """ host = "127.0.0.1" num_servers = 4 @@ -407,7 +409,8 @@ def test_gather_actual_addresses_end_to_end(): @pytest.mark.timeout(30) def test_gather_actual_addresses_child_crash_before_report(): """A child that exits before sending its endpoints must surface a - clear ``RuntimeError`` rather than hang or return ``None`` slots.""" + clear ``RuntimeError`` rather than hang or return ``None`` slots. + """ host = "127.0.0.1" num_servers = 2 placeholder_inputs = [ @@ -444,7 +447,8 @@ def test_gather_actual_addresses_child_crash_before_report(): def test_rust_frontend_launch_log_redacts_credentials(monkeypatch, caplog): """The Rust frontend command carries every non-default arg as JSON. - Credentials must not reach the log line.""" + Credentials must not reach the log line. + """ import subprocess as subprocess_mod from vllm.entrypoints.launchers.cli_args import make_arg_parser diff --git a/tests/entrypoints/launchers/test_cli_args.py b/tests/entrypoints/launchers/test_cli_args.py index e84abe5d26..cf2f678e40 100644 --- a/tests/entrypoints/launchers/test_cli_args.py +++ b/tests/entrypoints/launchers/test_cli_args.py @@ -152,7 +152,7 @@ def test_multiple_valid_inputs(serve_parser): ### Tests for serve argument validation that run prior to loading def test_enable_auto_choice_passes_without_tool_call_parser(serve_parser): - """Ensure validation fails if tool choice is enabled with no call parser""" + """Ensure validation fails if tool choice is enabled with no call parser.""" # If we enable-auto-tool-choice, explode with no tool-call-parser args = serve_parser.parse_args(args=["--enable-auto-tool-choice"]) with pytest.raises(TypeError): @@ -160,7 +160,7 @@ def test_enable_auto_choice_passes_without_tool_call_parser(serve_parser): def test_enable_auto_choice_passes_with_tool_call_parser(serve_parser): - """Ensure validation passes with tool choice enabled with a call parser""" + """Ensure validation passes with tool choice enabled with a call parser.""" args = serve_parser.parse_args( args=[ "--enable-auto-tool-choice", @@ -172,7 +172,7 @@ def test_enable_auto_choice_passes_with_tool_call_parser(serve_parser): def test_enable_auto_choice_fails_with_enable_reasoning(serve_parser): - """Ensure validation fails if reasoning is enabled with auto tool choice""" + """Ensure validation fails if reasoning is enabled with auto tool choice.""" args = serve_parser.parse_args( args=[ "--enable-auto-tool-choice", @@ -186,7 +186,8 @@ def test_enable_auto_choice_fails_with_enable_reasoning(serve_parser): def test_passes_with_reasoning_parser(serve_parser): """Ensure validation passes if reasoning is enabled - with a reasoning parser""" + with a reasoning parser + """ args = serve_parser.parse_args( args=[ "--reasoning-parser", @@ -197,7 +198,7 @@ def test_passes_with_reasoning_parser(serve_parser): def test_chat_template_validation_for_happy_paths(serve_parser): - """Ensure validation passes if the chat template exists""" + """Ensure validation passes if the chat template exists.""" args = serve_parser.parse_args( args=["--chat-template", CHATML_JINJA_PATH.absolute().as_posix()] ) @@ -205,7 +206,7 @@ def test_chat_template_validation_for_happy_paths(serve_parser): def test_chat_template_validation_for_sad_paths(serve_parser): - """Ensure validation fails if the chat template doesn't exist""" + """Ensure validation fails if the chat template doesn't exist.""" args = serve_parser.parse_args(args=["--chat-template", "does/not/exist"]) with pytest.raises(VLLMValidationError): validate_parsed_serve_args(args) @@ -304,13 +305,13 @@ def test_sse_keep_alive_interval_non_integer(serve_parser, value): ], ) def test_middleware(serve_parser, cli_args, expected_middleware): - """Ensure multiple middleware args are parsed properly""" + """Ensure multiple middleware args are parsed properly.""" args = serve_parser.parse_args(args=cli_args) assert args.middleware == expected_middleware def test_default_chat_template_kwargs_parsing(serve_parser): - """Ensure default_chat_template_kwargs JSON is parsed correctly""" + """Ensure default_chat_template_kwargs JSON is parsed correctly.""" args = serve_parser.parse_args( args=["--default-chat-template-kwargs", '{"enable_thinking": false}'] ) @@ -318,7 +319,7 @@ def test_default_chat_template_kwargs_parsing(serve_parser): def test_default_chat_template_kwargs_complex(serve_parser): - """Ensure complex default_chat_template_kwargs JSON is parsed correctly""" + """Ensure complex default_chat_template_kwargs JSON is parsed correctly.""" kwargs_json = '{"enable_thinking": false, "custom_param": "value", "num": 42}' args = serve_parser.parse_args(args=["--default-chat-template-kwargs", kwargs_json]) assert args.default_chat_template_kwargs == { @@ -329,7 +330,7 @@ def test_default_chat_template_kwargs_complex(serve_parser): def test_default_chat_template_kwargs_default_none(serve_parser): - """Ensure default_chat_template_kwargs defaults to None""" + """Ensure default_chat_template_kwargs defaults to None.""" args = serve_parser.parse_args(args=[]) assert args.default_chat_template_kwargs is None @@ -355,7 +356,7 @@ def test_resolve_default_chat_template_kwargs(default_kwargs, cohere_format, exp def test_default_chat_template_kwargs_invalid_json(serve_parser): - """Ensure invalid JSON raises an error""" + """Ensure invalid JSON raises an error.""" with pytest.raises(SystemExit): serve_parser.parse_args( args=["--default-chat-template-kwargs", "not valid json"] @@ -381,7 +382,8 @@ def test_default_chat_template_kwargs_invalid_json(serve_parser): ) def test_served_model_name_parsing(tmp_path, vllm_parser, args, raises): """Ensure that users don't misuse --served-model-name and end up with the default - model tag instead of the one they intended to serve.""" + model tag instead of the one they intended to serve. + """ # Call the serve subparser args.insert(0, "serve") # Create a dummy config file if the test case includes it @@ -402,7 +404,7 @@ def test_served_model_name_parsing(tmp_path, vllm_parser, args, raises): ### Tests for LoRA target modules parsing def test_lora_target_modules_single(serve_parser): - """Test parsing single lora-target-modules argument""" + """Test parsing single lora-target-modules argument.""" args = serve_parser.parse_args( args=["--enable-lora", "--lora-target-modules", "o_proj"] ) @@ -410,7 +412,7 @@ def test_lora_target_modules_single(serve_parser): def test_lora_target_modules_multiple(serve_parser): - """Test parsing multiple lora-target-modules arguments""" + """Test parsing multiple lora-target-modules arguments.""" args = serve_parser.parse_args( args=[ "--enable-lora", @@ -424,6 +426,6 @@ def test_lora_target_modules_multiple(serve_parser): def test_lora_target_modules_default_none(serve_parser): - """Test that lora-target-modules defaults to None""" + """Test that lora-target-modules defaults to None.""" args = serve_parser.parse_args(args=[]) assert args.lora_target_modules is None diff --git a/tests/entrypoints/launchers/test_dp_supervisor.py b/tests/entrypoints/launchers/test_dp_supervisor.py index eb5bf0a088..f96b020d10 100644 --- a/tests/entrypoints/launchers/test_dp_supervisor.py +++ b/tests/entrypoints/launchers/test_dp_supervisor.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for DPSupervisor: unit tests and lifecycle integration tests. +"""Tests for DPSupervisor: unit tests and lifecycle integration tests. Lifecycle integration tests replace child vLLM servers with lightweight aiohttp "fake" servers controlled by the test, so the suite runs without GPUs. @@ -463,8 +462,7 @@ async def test_start_server_no_log_config_when_no_filter(monkeypatch): class MockVLLMServer: - """ - Minimal FastAPI server that mimics one vLLM replica. + """Minimal FastAPI server that mimics one vLLM replica. GET /health returns 200 when healthy, 503 otherwise. Health state is toggled by the test via set_healthy(). """ @@ -582,8 +580,7 @@ def launch_mock_vllm_with_drain( async def _poll_supervisor_health(expected_status: int, use_ssl: bool = False) -> bool: - """ - GET /health on the supervisor once and check for expected_status. + """GET /health on the supervisor once and check for expected_status. Pass expected_status=-1 to assert the supervisor is not listening yet (a connection error is expected). The supervisor only starts its HTTP @@ -705,8 +702,7 @@ async def _run_supervisor( @pytest.mark.asyncio async def test_basic_lifecycle(monkeypatch): - """ - A) Supervisor is not listening while children are unhealthy. + """A) Supervisor is not listening while children are unhealthy. B) /health returns 200 once every child reports healthy. C) SIGTERM and shutdown """ @@ -777,8 +773,7 @@ async def test_basic_lifecycle_with_ssl(monkeypatch): @pytest.mark.asyncio async def test_failed_startup(monkeypatch): - """ - A) One of the vLLM servers crashes during startup. + """A) One of the vLLM servers crashes during startup. B) DPSupervisor detects this, and cleans up resources. """ args = _make_args() @@ -801,8 +796,7 @@ async def test_failed_startup(monkeypatch): @pytest.mark.asyncio async def test_becomes_unhealthy(monkeypatch): - """ - A) Supervisor is not listening while children are unhealthy. + """A) Supervisor is not listening while children are unhealthy. B) /health returns 200 once every child reports healthy. C) Child process becomes unhealthy. D) Detected and shutdown. @@ -842,8 +836,7 @@ async def test_becomes_unhealthy(monkeypatch): @pytest.mark.asyncio async def test_dp_server_fails(monkeypatch): - """ - A) Supervisor is not listening while children are unhealthy. + """A) Supervisor is not listening while children are unhealthy. B) /health returns 200 once every child reports healthy. C) Child process fails. D) Detected and shutdown. @@ -886,8 +879,7 @@ async def test_dp_server_fails(monkeypatch): @pytest.mark.asyncio async def test_shutdown_timeout(monkeypatch: pytest.MonkeyPatch): - """ - Child mock servers delay shutdown by 10s on SIGTERM (simulating in-flight + """Child mock servers delay shutdown by 10s on SIGTERM (simulating in-flight request drain). The supervisor is configured with shutdown_timeout=10, so its total wait budget is 10 + CHILD_EXIT_GRACE_S seconds. The children exit naturally within that window, so no force-kill should occur diff --git a/tests/entrypoints/launchers/test_run_batch.py b/tests/entrypoints/launchers/test_run_batch.py index 7e6ba3d78d..f7a2505807 100644 --- a/tests/entrypoints/launchers/test_run_batch.py +++ b/tests/entrypoints/launchers/test_run_batch.py @@ -509,9 +509,7 @@ def test_completions(): def test_completions_invalid_input(): - """ - Ensure that we fail when the input doesn't conform to the openai api. - """ + """Ensure that we fail when the input doesn't conform to the openai api.""" with ( tempfile.NamedTemporaryFile("w") as input_file, tempfile.NamedTemporaryFile("r") as output_file, @@ -602,9 +600,7 @@ def test_score(input_batch): def test_reasoning_parser(): - """ - Test that reasoning_parser parameter works correctly in run_batch. - """ + """Test that reasoning_parser parameter works correctly in run_batch.""" with ( tempfile.NamedTemporaryFile("w") as input_file, tempfile.NamedTemporaryFile("r") as output_file, @@ -767,8 +763,7 @@ def test_translation(): def test_tool_calling(): - """ - Test that tool calling works correctly in run_batch. + """Test that tool calling works correctly in run_batch. Verifies that requests with tools return tool_calls in the response. """ with ( diff --git a/tests/entrypoints/launchers/test_shutdown.py b/tests/entrypoints/launchers/test_shutdown.py index 643eff7594..6d18d825eb 100644 --- a/tests/entrypoints/launchers/test_shutdown.py +++ b/tests/entrypoints/launchers/test_shutdown.py @@ -134,7 +134,6 @@ async def test_shutdown_on_engine_failure(tmp_path: Path): Starts a vLLM server, kills it to simulate a crash, then verifies that subsequent API calls fail appropriately. """ - port = get_open_port() # Redirect server output to a file rather than a pipe: nothing drains the diff --git a/tests/entrypoints/llm/offline_mode/test_offline_mode.py b/tests/entrypoints/llm/offline_mode/test_offline_mode.py index 555e8b9b7d..ab7edd019b 100644 --- a/tests/entrypoints/llm/offline_mode/test_offline_mode.py +++ b/tests/entrypoints/llm/offline_mode/test_offline_mode.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for HF_HUB_OFFLINE mode""" +"""Tests for HF_HUB_OFFLINE mode.""" import importlib import sys diff --git a/tests/entrypoints/llm/test_accuracy.py b/tests/entrypoints/llm/test_accuracy.py index af607720c8..49d83b9cf6 100644 --- a/tests/entrypoints/llm/test_accuracy.py +++ b/tests/entrypoints/llm/test_accuracy.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This file test accuracy of the vLLM server via LMEval. +"""This file test accuracy of the vLLM server via LMEval. It uses local-completions, which interacts with vLLM through the OAI API with N concurrent connections. This simulates real work usage of the API and makes @@ -33,7 +32,6 @@ EXPECTED_VALUES = { def run_test(model_name, more_args=None): """Run the end to end accuracy test.""" - model_args = f"pretrained={model_name},max_model_len=4096" if more_args is not None: @@ -64,7 +62,6 @@ TPU_TP_TEST_STR = "" # "tensor_parallel_size=4" @pytest.mark.parametrize("model", MODEL_NAMES) def test_lm_eval_accuracy_v1_engine(model): """Run with the V1 Engine.""" - more_args = None if current_platform.is_tpu(): # Limit compilation time for TPU V1 @@ -81,7 +78,6 @@ def test_lm_eval_accuracy_v1_engine(model): @pytest.mark.parametrize("model", FP8_KV_MODEL_NAMES) def test_lm_eval_accuracy_v1_engine_fp8_kv_cache(model): """Run with the V1 Engine.""" - more_args = None if current_platform.is_tpu(): # Limit compilation time for TPU V1 diff --git a/tests/entrypoints/llm/test_chat.py b/tests/entrypoints/llm/test_chat.py index 6ee8d34ce1..276426e692 100644 --- a/tests/entrypoints/llm/test_chat.py +++ b/tests/entrypoints/llm/test_chat.py @@ -20,8 +20,7 @@ def text_llm(vllm_runner): @pytest.fixture(scope="function") def llm_for_failure_test(vllm_runner): - """ - Fixture for testing issue #26081. + """Fixture for testing issue #26081. Uses a small max_model_len to easily trigger length errors. """ with vllm_runner( @@ -67,8 +66,7 @@ def test_multi_chat(text_llm): def test_llm_chat_tokenization_no_double_bos(text_llm): - """ - LLM.chat() should not add special tokens when using chat templates. + """LLM.chat() should not add special tokens when using chat templates. Check we get a single BOS token for llama chat. """ messages = [ @@ -127,8 +125,7 @@ def test_chat_extra_kwargs(thinking_llm, enable_thinking): def test_chat_batch_failure_cleanup(llm_for_failure_test): - """ - Tests that if a batch call to llm.chat() fails mid-way + """Tests that if a batch call to llm.chat() fails mid-way (e.g., due to one invalid prompt), the requests that were already enqueued are properly aborted and do not pollute the queue for subsequent calls. diff --git a/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py b/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py index 5c378de35d..a3511b9528 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py +++ b/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test that ``InputProcessor.inject_into_mm_cache()`` correctly injects +"""Test that ``InputProcessor.inject_into_mm_cache()`` correctly injects pre-processed mm_kwargs into the processor cache and reports MM cache hit rate metrics accurately. diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py b/tests/entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py index 14ac7ce043..5031cbea1a 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py @@ -47,8 +47,7 @@ async def client(server): async def test_online_audio_in_video( client: openai.AsyncOpenAI, video_assets: VideoTestAssets ): - """Test video input with `audio_in_video=True`""" - + """Test video input with `audio_in_video=True`.""" # we don't use video_urls above because they missed audio stream. video_path = video_assets[0].video_path with open(video_path, "rb") as f: @@ -97,8 +96,7 @@ async def test_online_audio_in_video( async def test_online_audio_in_video_multi_videos( client: openai.AsyncOpenAI, video_assets: VideoTestAssets ): - """Test multi-video input with `audio_in_video=True`""" - + """Test multi-video input with `audio_in_video=True`.""" # we don't use video_urls above because they missed audio stream. video_path = video_assets[0].video_path with open(video_path, "rb") as f: @@ -151,8 +149,7 @@ async def test_online_audio_in_video_multi_videos( async def test_online_audio_in_video_interleaved( client: openai.AsyncOpenAI, video_assets: VideoTestAssets ): - """Test interleaved video/audio input with `audio_in_video=True`""" - + """Test interleaved video/audio input with `audio_in_video=True`.""" # we don't use video_urls above because they missed audio stream. video_path = video_assets[0].video_path with open(video_path, "rb") as f: diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py index aac68334bf..b12bf3a60b 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py @@ -2,7 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """E2E test for mixing `prompt_embeds` with `audio_embeds` in a single -Chat Completions request.""" +Chat Completions request. +""" import json diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py index c6ce165cf3..b72d940714 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py @@ -68,7 +68,8 @@ async def client(server): def image_url() -> str: """Stable real image as a data URL, kept identical across both the text and prompt_embeds requests so any output difference must come from - how the text content is delivered.""" + how the text content is delivered. + """ return encode_image_url(ImageAsset("stop_sign").pil_image) diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py b/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py index e6b8dc23e9..6d1c2f24e9 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py @@ -34,8 +34,7 @@ REQUIRED_BEAM_SEARCH_TERMS = [ def check_output_matches_terms(content: str, term_groups: list[list[str]]) -> bool: - """ - Check if content matches all required term groups. + """Check if content matches all required term groups. Each term group requires at least one of its terms to be present. All term groups must be satisfied. """ @@ -146,7 +145,8 @@ def describe_image_messages( ) -> list[dict]: """Build the system + user messages used by the completions-with-image family of tests. *extra_image_fields* is merged into the top-level - image content block (for uuid / bad-key tests).""" + image content block (for uuid / bad-key tests). + """ image_block: dict = { "type": "image_url", "image_url": {"url": image_url}, @@ -176,7 +176,8 @@ async def complete_and_check( temperature: float = 0.0, ) -> str: """Run a chat completion and assert the output is non-empty. - Returns the content string.""" + Returns the content string. + """ chat_completion = await client.chat.completions.create( model=model_name, messages=messages, diff --git a/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py b/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py index 20218ac6c8..467e23e7d3 100644 --- a/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py +++ b/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py @@ -259,7 +259,7 @@ async def _generator(prompt_idx: int, text: str) -> AsyncGenerator[RequestOutput @pytest.mark.asyncio @pytest.mark.skip_global_cleanup async def test_batched_echo_does_not_prepend_user_prompt() -> None: - """echo must only echo the assistant turn, never the user's prompt. + """Echo must only echo the assistant turn, never the user's prompt. With `add_generation_prompt` (the default) the response role is `assistant` while the last conversation message is the user's, so there is nothing to diff --git a/tests/entrypoints/openai/chat_completion/test_chat_error.py b/tests/entrypoints/openai/chat_completion/test_chat_error.py index aafb899f00..e0f7128744 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_error.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_error.py @@ -121,7 +121,7 @@ def _build_serving_chat(engine: AsyncLLM) -> OpenAIServingChat: @pytest.mark.asyncio async def test_chat_error_non_stream(): - """test finish_reason='error' returns 500 InternalServerError (non-streaming)""" + """Test finish_reason='error' returns 500 InternalServerError (non-streaming).""" mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig() @@ -246,7 +246,7 @@ async def test_renderer_only_chat_request_skips_mm_cache(): @pytest.mark.asyncio async def test_chat_error_stream(): - """test finish_reason='error' returns 500 InternalServerError (streaming)""" + """Test finish_reason='error' returns 500 InternalServerError (streaming).""" mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig() @@ -465,7 +465,8 @@ def test_system_message_warns_on_video(video_content): def test_json_schema_response_format_missing_schema(): """When response_format type is 'json_schema' but the json_schema field is not provided, request construction should raise a validation error - so the API returns 400 instead of 500.""" + so the API returns 400 instead of 500. + """ with pytest.raises(Exception, match="json_schema.*must be provided"): ChatCompletionRequest( model=MODEL_NAME, @@ -519,7 +520,8 @@ def test_structured_outputs_structural_tag_invalid(structural_tag): @pytest.mark.parametrize("field_name", ["prompt_logprobs", "top_logprobs"]) def test_non_numeric_logprobs_rejected(field_name): """A non-numeric logprobs value must be a clean 400 validation error, not a - TypeError from the mode='before' comparison (which surfaces as HTTP 500).""" + TypeError from the mode='before' comparison (which surfaces as HTTP 500). + """ with pytest.raises(VLLMValidationError, match=f"`{field_name}` must be an integer"): ChatCompletionRequest( model=MODEL_NAME, diff --git a/tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py b/tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py index b415fa116d..f1801701e3 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py @@ -81,7 +81,8 @@ async def test_chat_logit_bias_invalid(client): @pytest.mark.asyncio async def test_chat_logit_bias_non_integer_key(client): """Test that a non-integer logit_bias key is rejected with a clean, - informative error instead of a raw 'invalid literal for int()' message.""" + informative error instead of a raw 'invalid literal for int()' message. + """ with pytest.raises(openai.BadRequestError) as excinfo: await client.chat.completions.create( model=MODEL_NAME, @@ -101,7 +102,8 @@ async def test_chat_logit_bias_non_integer_key(client): @pytest.mark.asyncio async def test_chat_logit_bias_non_numeric_value(client): """Test that a non-numeric logit_bias value is rejected with a message - that names the specific offending token, not just a generic TypeError.""" + that names the specific offending token, not just a generic TypeError. + """ with pytest.raises(openai.BadRequestError) as excinfo: await client.chat.completions.create( model=MODEL_NAME, @@ -120,7 +122,8 @@ async def test_chat_logit_bias_non_numeric_value(client): @pytest.mark.asyncio async def test_chat_logit_bias_multiple_non_integer_keys(client): """Test that ALL invalid logit_bias keys are reported together, - not just the first one encountered.""" + not just the first one encountered. + """ with pytest.raises(openai.BadRequestError) as excinfo: await client.chat.completions.create( model=MODEL_NAME, diff --git a/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py b/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py index 3e59e7402e..61fd2f1576 100644 --- a/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py +++ b/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py @@ -451,7 +451,7 @@ async def test_inconsistent_tool_choice_and_tools( async def test_max_tokens_with_tool_choice_required( client: openai.AsyncOpenAI, tool_choice ): - """ """ + """Requesting a single token with a required tool choice must not crash.""" models = await client.models.list() model_name: str = models.data[0].id diff --git a/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py b/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py index 8c03fc1daa..9c0a4db68b 100644 --- a/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py +++ b/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py @@ -52,9 +52,10 @@ def stop_sign_image_url(): @pytest.mark.asyncio async def test_translategemma_extra_lang_code_fields(client: openai.AsyncOpenAI): - """en -> es translation through TranslateGemma's bundled chat template's + """En -> es translation through TranslateGemma's bundled chat template's ``text`` content branch, which depends on ``source_lang_code`` / - ``target_lang_code`` extra fields being preserved on the content part.""" + ``target_lang_code`` extra fields being preserved on the content part. + """ completion = await client.chat.completions.create( model=MODEL_NAME, messages=[ @@ -92,13 +93,14 @@ async def test_translategemma_image_extra_lang_code_fields( client: openai.AsyncOpenAI, stop_sign_image_url: str, ): - """en -> es OCR-translation through TranslateGemma's bundled chat + """En -> es OCR-translation through TranslateGemma's bundled chat template's ``image`` content branch. Exercises the multimodal branch of ``_collect_extra_fields`` in ``_parse_chat_message_content_part``: the ``source_lang_code`` / ``target_lang_code`` extras must survive parsing and end up alongside the ``{"type": "image"}`` placeholder so the template's image branch (``content["type"] == 'image'``) renders the - right OCR-translation prompt.""" + right OCR-translation prompt. + """ completion = await client.chat.completions.create( model=MODEL_NAME, messages=[ diff --git a/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py b/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py index 1eea6db48a..62f915b054 100644 --- a/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py +++ b/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py @@ -191,7 +191,8 @@ async def test_logprob_token_ids_stream_with_no_top_k(server): async def test_logprob_token_ids_default_behavior_unchanged(server): """Without `logprob_token_ids`, the response carries the natural top-k most-likely tokens. This guards against the new field accidentally - changing the default-path output.""" + changing the default-path output. + """ async with server.get_async_client() as client: resp = await client.chat.completions.create( model=MODEL_NAME, diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index e7ea15877c..73cffa76cf 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -1232,7 +1232,6 @@ async def test_serving_chat_mistral_token_ids_prompt_is_validated(): """Regression test: when the Mistral tokenizer path returns token IDs directly, we must still apply input length + max_tokens validation. """ - mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig(skip_tokenizer_init=True) @@ -1271,7 +1270,6 @@ async def test_serving_chat_mistral_token_ids_prompt_too_long_is_rejected(): """Regression test: MistralTokenizer token-id prompts must still enforce the max context length for the input itself (token_num >= max_model_len). """ - mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig(skip_tokenizer_init=True) @@ -1408,7 +1406,8 @@ async def test_serving_chat_did_set_correct_cache_salt(model_type): @pytest.mark.asyncio async def test_serving_chat_data_parallel_rank_extraction(): """Test that data_parallel_rank is properly extracted from header and - passed to engine.""" + passed to engine. + """ mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig() @@ -1512,8 +1511,7 @@ async def _render_chat_prompt_token_ids( class TestServingChatWithHarmony: - """ - These tests ensure Chat Completion requests are being properly converted into + """These tests ensure Chat Completion requests are being properly converted into Harmony messages and Harmony response messages back into Chat Completion responses. These tests are not exhaustive, but each one was created to cover a specific case that we got wrong but is now fixed. @@ -1731,7 +1729,8 @@ class TestServingChatWithHarmony: @pytest.mark.asyncio async def test_system_message_without_tools(self, serving_chat, stream): """Leading system message produces a developer message with - DeveloperContent (# Instructions header).""" + DeveloperContent (# Instructions header). + """ messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Hello"}, @@ -1753,7 +1752,8 @@ class TestServingChatWithHarmony: @pytest.mark.asyncio async def test_system_message_with_tools(self, serving_chat, stream, weather_tools): """Leading system message is folded into the developer message - alongside tool definitions.""" + alongside tool definitions. + """ messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What's the weather?"}, @@ -2168,7 +2168,8 @@ class TestServingChatWithHarmony: @pytest.mark.asyncio async def test_tool_choice_validation_without_parser(): """Test that tool_choice='required' or named tool without tool_parser - returns an appropriate error message.""" + returns an appropriate error message. + """ mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig() diff --git a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py index ec7a977f1a..b7ed422223 100644 --- a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py +++ b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py @@ -33,7 +33,6 @@ def _count_reasoning_decode_token_ids_between_markers( reasoning_end_ids: list[int], ) -> int | None: """Count decode tokens in the thinking span (after last start, before first end).""" - if not reasoning_start_ids or not reasoning_end_ids: raise ValueError("reasoning marker token id lists must be non-empty") @@ -152,8 +151,8 @@ async def client(request, server, server_with_auto_reasoning_config): @pytest.mark.parametrize("client", ["default", "auto_config"], indirect=True) async def test_thinking_token_budget_mixed_requests(client: openai.AsyncOpenAI): """Test that mixed requests (some with thinking_token_budget, some without) - complete successfully without errors.""" - + complete successfully without errors. + """ response_with_budget = await client.chat.completions.create( model=MODEL_NAME, messages=MESSAGES, @@ -182,7 +181,6 @@ async def test_thinking_token_budget_limits_reasoning(client: openai.AsyncOpenAI grouped into streamed chunks (a single chunk can carry several tokens under async scheduling / stream_interval > 1). Counting chunks under-counts. """ - tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME) start_ids = list(tokenizer.encode(REASONING_START_STR, add_special_tokens=False)) end_ids = list(tokenizer.encode(REASONING_END_STR, add_special_tokens=False)) @@ -227,7 +225,6 @@ async def test_thinking_token_budget_qwen35_fp8_mtp_concurrent_mixed_budget_and_ Qwen3.5 FP8 + MTP (TP=2) server. Budgeted calls are checked with ``_count_reasoning_decode_token_ids_between_markers`` on full token ids. """ - _batch_spec: list[tuple[Literal["budget"], int] | tuple[Literal["plain"], None]] = [ ("budget", 1), ("budget", 12), diff --git a/tests/entrypoints/openai/completion/test_completion.py b/tests/entrypoints/openai/completion/test_completion.py index 1bc6260cde..2876642241 100644 --- a/tests/entrypoints/openai/completion/test_completion.py +++ b/tests/entrypoints/openai/completion/test_completion.py @@ -350,7 +350,6 @@ async def test_parallel_no_streaming(client: openai.AsyncOpenAI, model_name: str """Parallel sampling without streaming. A single request output contains a list of completions. """ - prompt = "What is an LLM?" n = 3 max_tokens = 50 @@ -396,7 +395,6 @@ async def test_parallel_streaming(client: openai.AsyncOpenAI, model_name: str): The tokens from multiple samples, are flattened into a single stream, with an index to indicate which sample the token belongs to. """ - prompt = "What is an LLM?" n = 3 max_tokens = 50 diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index 513bde465d..1fee99be57 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -303,7 +303,7 @@ def test_completion_metrics_carries_both_timing_and_spec_decode(): @pytest.mark.asyncio async def test_completion_error_non_stream(): - """test finish_reason='error' returns 500 InternalServerError (non-streaming)""" + """Test finish_reason='error' returns 500 InternalServerError (non-streaming).""" mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig() @@ -439,7 +439,7 @@ async def test_renderer_only_completion_request_skips_mm_cache(): @pytest.mark.asyncio async def test_completion_error_stream(): - """test finish_reason='error' returns 500 InternalServerError (streaming)""" + """Test finish_reason='error' returns 500 InternalServerError (streaming).""" mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig() @@ -521,7 +521,8 @@ async def test_completion_error_stream(): def test_json_schema_response_format_missing_schema(): """When response_format type is 'json_schema' but the json_schema field is not provided, request construction should raise a validation error - so the API returns 400 instead of 500.""" + so the API returns 400 instead of 500. + """ with pytest.raises(Exception, match="json_schema.*must be provided"): CompletionRequest( model=MODEL_NAME, @@ -584,7 +585,8 @@ def test_negative_prompt_token_ids_flat(): def test_logprobs_minus_one_allowed(): """logprobs=-1 means "return all logprobs". The sampling layer and the chat top_logprobs / prompt_logprobs validators all accept -1, so the completion - logprobs validator must accept it too instead of rejecting it as negative.""" + logprobs validator must accept it too instead of rejecting it as negative. + """ request = CompletionRequest( model=MODEL_NAME, prompt="Test prompt", @@ -744,7 +746,8 @@ class TestCompletionPromptListLimit: @pytest.mark.parametrize("field_name", ["prompt_logprobs", "logprobs"]) def test_non_numeric_logprobs_rejected(field_name): """A non-numeric logprobs value must be a clean 400 validation error, not a - TypeError from the mode='before' comparison (which surfaces as HTTP 500).""" + TypeError from the mode='before' comparison (which surfaces as HTTP 500). + """ with pytest.raises(VLLMValidationError, match=f"`{field_name}` must be an integer"): CompletionRequest( model=MODEL_NAME, diff --git a/tests/entrypoints/openai/completion/test_token_in_token_out.py b/tests/entrypoints/openai/completion/test_token_in_token_out.py index 8882ae6244..45024d8277 100644 --- a/tests/entrypoints/openai/completion/test_token_in_token_out.py +++ b/tests/entrypoints/openai/completion/test_token_in_token_out.py @@ -39,8 +39,7 @@ def server(): @pytest.mark.asyncio async def test_token_in_token_out_and_logprobs(server): - """ - Test token-in-token-out and token_ids align with prompt_logprobs + """Test token-in-token-out and token_ids align with prompt_logprobs & logprobs when return_tokens_as_token_ids is enabled. """ tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME) diff --git a/tests/entrypoints/openai/correctness/test_lmeval.py b/tests/entrypoints/openai/correctness/test_lmeval.py index aad1b5e062..c6ce3ebc2b 100644 --- a/tests/entrypoints/openai/correctness/test_lmeval.py +++ b/tests/entrypoints/openai/correctness/test_lmeval.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This file test accuracy of the vLLM server via LMEval. +"""This file test accuracy of the vLLM server via LMEval. It uses local-completions, which interacts with vLLM through the OAI API with N concurrent connections. This simulates real work usage of the API and makes @@ -37,7 +36,6 @@ if current_platform.is_tpu(): def run_test(more_args): """Run the end to end accuracy test.""" - args = list(DEFAULT_ARGS) args.extend(more_args) print(f"Running with: {args}") @@ -68,7 +66,6 @@ def run_test(more_args): def test_lm_eval_accuracy_v1_engine(): """Run with the V1 Engine.""" - more_args = [] # Limit compilation time for V1 on TPU diff --git a/tests/entrypoints/openai/parser/test_harmony_render_parity.py b/tests/entrypoints/openai/parser/test_harmony_render_parity.py index 5cb4461223..ce6e4c80f8 100644 --- a/tests/entrypoints/openai/parser/test_harmony_render_parity.py +++ b/tests/entrypoints/openai/parser/test_harmony_render_parity.py @@ -43,7 +43,8 @@ def _system(): class TestResponseInputToHarmonyRenderParity: """Each test drives the same conversation through both APIs and asserts - identical Harmony messages and rendered token sequences.""" + identical Harmony messages and rendered token sequences. + """ # ----------------------------------------------------------------------- # Single-message cases @@ -51,7 +52,8 @@ class TestResponseInputToHarmonyRenderParity: def test_developer_message(self): """Both APIs must render developer messages identically using - DeveloperContent (with the '# Instructions' header).""" + DeveloperContent (with the '# Instructions' header). + """ chat_msgs = parse_chat_input_to_harmony_message( {"role": "developer", "content": "Be concise."} ) @@ -246,7 +248,8 @@ class TestResponseInputToHarmonyRenderParity: def test_reasoning_combined_with_function_call(self): """Chat API packs reasoning + tool_calls into one dict; responses API represents them as two separate items. Both must produce the same two - Harmony messages in the same order: analysis then commentary.""" + Harmony messages in the same order: analysis then commentary. + """ chat_msgs = parse_chat_input_to_harmony_message( { "role": "assistant", @@ -529,7 +532,8 @@ class TestResponseInputToHarmonyRenderParity: def test_completed_turns_drop_reasoning(self): """Validates that reasoning from completed turns is dropped, while reasoning from the current in-progress tool-call turn is preserved - in both chat completions and responses previous_input_messages.""" + in both chat completions and responses previous_input_messages. + """ first_turn_reasoning = "FIRST_TURN_REASONING" second_turn_reasoning = "SECOND_TURN_REASONING" diff --git a/tests/entrypoints/openai/parser/test_harmony_utils.py b/tests/entrypoints/openai/parser/test_harmony_utils.py index 1b19714ac2..816e879342 100644 --- a/tests/entrypoints/openai/parser/test_harmony_utils.py +++ b/tests/entrypoints/openai/parser/test_harmony_utils.py @@ -170,7 +170,7 @@ class TestIsFunctionRecipientWithAllowedNames: """Tests for is_function_recipient with allowed_function_tool_names.""" def test_prefixed_always_accepted(self): - """functions. prefix is always accepted regardless of allowed names.""" + """Functions. prefix is always accepted regardless of allowed names.""" fn_names = frozenset({"other_tool"}) assert is_function_recipient("functions.get_weather", fn_names) is True @@ -252,8 +252,7 @@ class TestExtractFunctionFromRecipient: class TestCommonParseInputToHarmonyMessage: - """ - Tests for scenarios that are common to both Chat Completion + """Tests for scenarios that are common to both Chat Completion parse_chat_input_to_harmony_message and Responses API response_previous_input_to_harmony functions. """ @@ -324,7 +323,8 @@ class TestCommonParseInputToHarmonyMessage: def test_system_message(self, parse_function): """Test parsing system messages, which are parsed into developer messages - with DeveloperContent.""" + with DeveloperContent. + """ chat_msg = { "role": "system", "content": "You are a helpful assistant", @@ -448,8 +448,7 @@ class TestCommonParseInputToHarmonyMessage: class TestParseChatInputToHarmonyMessage: - """ - Tests for scenarios that are specific to the Chat Completion API + """Tests for scenarios that are specific to the Chat Completion API parse_chat_input_to_harmony_message function. """ diff --git a/tests/entrypoints/openai/responses/test_errors.py b/tests/entrypoints/openai/responses/test_errors.py index 0ea7833da9..1e0182ba6b 100644 --- a/tests/entrypoints/openai/responses/test_errors.py +++ b/tests/entrypoints/openai/responses/test_errors.py @@ -14,7 +14,7 @@ from vllm.exceptions import GenerationError @pytest.mark.asyncio async def test_raise_if_error_raises_generation_error(): - """test _raise_if_error raises GenerationError""" + """Test _raise_if_error raises GenerationError.""" # create a minimal GenerateBaseServing instance mock_engine = MagicMock() mock_engine.model_config = MagicMock() @@ -42,7 +42,7 @@ async def test_raise_if_error_raises_generation_error(): @pytest.mark.asyncio async def test_convert_generation_error_to_streaming_response(): - """test _convert_generation_error_to_streaming_response output""" + """Test _convert_generation_error_to_streaming_response output.""" mock_engine = MagicMock() mock_engine.model_config = MagicMock() mock_engine.model_config.max_model_len = 100 diff --git a/tests/entrypoints/openai/responses/test_function_call.py b/tests/entrypoints/openai/responses/test_function_call.py index 9dcbd74c89..006ee0f6f2 100644 --- a/tests/entrypoints/openai/responses/test_function_call.py +++ b/tests/entrypoints/openai/responses/test_function_call.py @@ -168,8 +168,7 @@ async def test_max_tokens_with_tool_choice_required( @pytest.mark.asyncio async def test_named_tool_use(client: openai.AsyncOpenAI): def get_weather(latitude: float, longitude: float) -> str: - """ - Mock function to simulate getting weather data. + """Mock function to simulate getting weather data. In a real application, this would call an external weather API. """ return f"Current temperature at ({latitude}, {longitude}) is 20°C." diff --git a/tests/entrypoints/openai/responses/test_function_call_parsing.py b/tests/entrypoints/openai/responses/test_function_call_parsing.py index f90a641db7..c83fe0f61c 100644 --- a/tests/entrypoints/openai/responses/test_function_call_parsing.py +++ b/tests/entrypoints/openai/responses/test_function_call_parsing.py @@ -12,7 +12,8 @@ from vllm.entrypoints.openai.responses.protocol import ResponsesRequest def test_function_call_dict_converted_to_object(): """Test that function_call dictionaries are correctly parsed into - ResponseFunctionToolCall objects.""" + ResponseFunctionToolCall objects. + """ # Create a request with function_call as dict request_data = { "model": "gpt-oss", @@ -57,7 +58,6 @@ def test_direct_function_call_object_preservation(): def test_mixed_input_types_with_function_calls(): """Test parsing with mixed input types including function calls.""" - request_data = { "model": "gpt-oss", "input": [ @@ -177,7 +177,6 @@ def test_empty_list_input(): def test_function_call_output_not_affected(): """Test that FunctionCallOutput is not affected by the function_call parsing.""" - # Test with FunctionCallOutput as dict (should not be parsed) request_data = { "model": "gpt-oss", @@ -265,7 +264,6 @@ def test_function_call_validation_failure_logs_debug(caplog): def test_validator_handles_iterator_input(): """Test that validator can handle ValidatorIterator input (Pydantic internal).""" - # This test simulates when Pydantic passes a ValidatorIterator instead of a list # This happened with complex nested structures containing reasoning + function_call @@ -332,7 +330,8 @@ def test_validator_handles_empty_iterator(): def test_assistant_string_content_stays_easyinput(): """EasyInput assistant message with plain string content is not - coerced into a ResponseOutputMessage.""" + coerced into a ResponseOutputMessage. + """ request_data = { "model": "test-model", "input": [ @@ -354,7 +353,8 @@ def test_assistant_string_content_stays_easyinput(): def test_assistant_output_style_content_coerced(): """Assistant message whose content is output-message-shaped (list of - output_text items) should be coerced to ResponseOutputMessage.""" + output_text items) should be coerced to ResponseOutputMessage. + """ request_data = { "model": "test-model", "input": [ diff --git a/tests/entrypoints/openai/responses/test_harmony_utils.py b/tests/entrypoints/openai/responses/test_harmony_utils.py index 54c034de44..b524842467 100644 --- a/tests/entrypoints/openai/responses/test_harmony_utils.py +++ b/tests/entrypoints/openai/responses/test_harmony_utils.py @@ -20,8 +20,7 @@ from vllm.exceptions import VLLMValidationError class TestResponsePreviousInputToHarmony: - """ - Tests for scenarios that are specific to the Responses API + """Tests for scenarios that are specific to the Responses API response_previous_input_to_harmony function. """ diff --git a/tests/entrypoints/openai/responses/test_response_input_to_harmony.py b/tests/entrypoints/openai/responses/test_response_input_to_harmony.py index e2be8477e4..2485983bf3 100644 --- a/tests/entrypoints/openai/responses/test_response_input_to_harmony.py +++ b/tests/entrypoints/openai/responses/test_response_input_to_harmony.py @@ -67,7 +67,8 @@ class TestResponseInputToHarmonyMessage: def test_system_message(self): """System messages carry developer instructions and must be rendered - as developer messages with DeveloperContent.""" + as developer messages with DeveloperContent. + """ msg = response_input_to_harmony( {"type": "message", "role": "system", "content": "Be helpful."}, prev_responses=[], @@ -89,7 +90,8 @@ class TestResponseInputToHarmonyMessage: def test_developer_message_gets_instructions_prefix(self): """Developer messages must use DeveloperContent which adds the - '# Instructions' header the model was trained on.""" + '# Instructions' header the model was trained on. + """ msg = response_input_to_harmony( {"type": "message", "role": "developer", "content": "Be concise."}, prev_responses=[], @@ -119,7 +121,8 @@ class TestResponseInputToHarmonyMessage: def test_developer_message_array_content_concatenated(self): """Array content in developer messages is flattened and rendered - via DeveloperContent with the '# Instructions' header.""" + via DeveloperContent with the '# Instructions' header. + """ msg = response_input_to_harmony( { "type": "message", @@ -210,7 +213,8 @@ class TestResponseInputToHarmonyMessage: def test_function_call_output_uses_most_recent_matching_call(self): """When multiple prev_responses share a call_id, the last one wins - because the search is reversed.""" + because the search is reversed. + """ earlier = ResponseFunctionToolCall( id="fc_old", call_id="call_test", diff --git a/tests/entrypoints/openai/responses/test_responses_utils.py b/tests/entrypoints/openai/responses/test_responses_utils.py index 967ddacda6..1e59c82244 100644 --- a/tests/entrypoints/openai/responses/test_responses_utils.py +++ b/tests/entrypoints/openai/responses/test_responses_utils.py @@ -839,11 +839,13 @@ class TestConstructChatMessagesCombinePolicy: class TestConstructInputMessagesInstructionsLeak: """Regression tests for #37697: instructions from a prior response - should NOT leak through previous_response_id.""" + should NOT leak through previous_response_id. + """ def test_old_instructions_stripped_from_prev_msg(self): """System message in prev_msg must be dropped so the new request's - instructions are the only system message in the conversation.""" + instructions are the only system message in the conversation. + """ prev = [ {"role": "system", "content": "old instructions"}, {"role": "user", "content": "What is 2+2?"}, @@ -860,7 +862,8 @@ class TestConstructInputMessagesInstructionsLeak: def test_no_instructions_in_new_request(self): """If the new request has no instructions, old ones should still - be stripped -- they must not carry over.""" + be stripped -- they must not carry over. + """ prev = [ {"role": "system", "content": "old instructions"}, {"role": "user", "content": "Hi"}, diff --git a/tests/entrypoints/openai/responses/test_serving_responses.py b/tests/entrypoints/openai/responses/test_serving_responses.py index bd41b8cda2..1f9d95ea2e 100644 --- a/tests/entrypoints/openai/responses/test_serving_responses.py +++ b/tests/entrypoints/openai/responses/test_serving_responses.py @@ -81,7 +81,7 @@ def _new_online_renderer() -> OnlineRenderer: class MockConversationContext(ConversationContext): - """Mock conversation context for testing""" + """Mock conversation context for testing.""" def __init__(self): self.init_tool_sessions_called = False @@ -123,7 +123,7 @@ def test_serialize_message_pydantic_model_returns_dict() -> None: @pytest.fixture def mock_serving_responses(): - """Create a mock OpenAIServingResponses instance""" + """Create a mock OpenAIServingResponses instance.""" serving_responses = MagicMock(spec=OpenAIServingResponses) serving_responses.tool_server = MagicMock(spec=ToolServer) return serving_responses @@ -131,13 +131,13 @@ def mock_serving_responses(): @pytest.fixture def mock_context(): - """Create a mock conversation context""" + """Create a mock conversation context.""" return MockConversationContext() @pytest.fixture def mock_exit_stack(): - """Create a mock async exit stack""" + """Create a mock async exit stack.""" return MagicMock(spec=AsyncExitStack) @@ -550,7 +550,7 @@ def test_responses_render_result_rejects_non_single_prompt(engine_inputs): class TestInitializeToolSessions: - """Test class for _initialize_tool_sessions method""" + """Test class for _initialize_tool_sessions method.""" @pytest.fixture(params=[ParsableContext, HarmonyContext]) def tool_context(self, request): @@ -635,7 +635,7 @@ class TestInitializeToolSessions: @pytest_asyncio.fixture async def serving_responses_instance(self): - """Create a real OpenAIServingResponses instance for testing""" + """Create a real OpenAIServingResponses instance for testing.""" # Create minimal mocks for required dependencies engine_client = MagicMock() @@ -802,8 +802,7 @@ class TestInitializeToolSessions: async def test_initialize_tool_sessions( self, serving_responses_instance, mock_context, mock_exit_stack ): - """Test that method works correctly with only MCP tools""" - + """Test that method works correctly with only MCP tools.""" request = ResponsesRequest(input="test input", tools=[]) # Call the method @@ -852,11 +851,11 @@ class TestInitializeToolSessions: class TestValidateGeneratorInput: - """Test class for _validate_generator_input method""" + """Test class for _validate_generator_input method.""" @pytest_asyncio.fixture async def serving_responses_instance(self): - """Create a real OpenAIServingResponses instance for testing""" + """Create a real OpenAIServingResponses instance for testing.""" # Create minimal mocks for required dependencies engine_client = MagicMock() @@ -884,7 +883,7 @@ class TestValidateGeneratorInput: return instance def test_validate_generator_input(self, serving_responses_instance): - """Test _validate_generator_input with valid prompt length""" + """Test _validate_generator_input with valid prompt length.""" # Create an engine prompt with valid length (less than max_model_len) valid_prompt_token_ids = list(range(5)) # 5 tokens < 100 max_model_len engine_input = tokens_input(valid_prompt_token_ids) @@ -999,7 +998,7 @@ async def test_reasoning_tokens_counted_for_text_reasoning_model(monkeypatch): class TestExtractAllowedToolsFromMcpRequests: - """Test class for _extract_allowed_tools_from_mcp_requests function""" + """Test class for _extract_allowed_tools_from_mcp_requests function.""" def test_extract_allowed_tools_basic_formats(self): """Test extraction with list format, object format, and None.""" @@ -1113,7 +1112,7 @@ class TestHarmonyPreambleStreaming: return item def test_preamble_delta_emits_text_events(self) -> None: - """commentary + recipient=None should emit output_text.delta events.""" + """Commentary + recipient=None should emit output_text.delta events.""" from vllm.entrypoints.openai.responses.streaming_events import ( emit_content_delta_events, ) @@ -1146,7 +1145,7 @@ class TestHarmonyPreambleStreaming: assert "response.output_item.added" not in type_names def test_commentary_with_function_recipient_not_preamble(self) -> None: - """commentary + recipient='functions.X' must NOT use preamble path.""" + """Commentary + recipient='functions.X' must NOT use preamble path.""" from vllm.entrypoints.openai.responses.streaming_events import ( emit_content_delta_events, ) @@ -1164,7 +1163,8 @@ class TestHarmonyPreambleStreaming: def test_preamble_done_emits_text_done_events(self) -> None: """Completed preamble should emit text done + content_part done + - output_item done, same shape as final channel.""" + output_item done, same shape as final channel. + """ from vllm.entrypoints.openai.responses.streaming_events import ( emit_previous_item_done_events, ) @@ -1184,8 +1184,9 @@ class TestHarmonyPreambleStreaming: assert "response.output_item.done" in type_names def test_commentary_with_recipient_no_preamble_done(self) -> None: - """commentary + recipient='functions.X' should route to function call - done, not preamble done.""" + """Commentary + recipient='functions.X' should route to function call + done, not preamble done. + """ from vllm.entrypoints.openai.responses.streaming_events import ( emit_previous_item_done_events, ) @@ -1331,7 +1332,8 @@ def _mock_parser_with_reasoning(serving, delta_sequence: list[DeltaMessage]): class TestStreamingReasoningToContentTransition: """Tests for _process_simple_streaming_events reasoning-to-content transition, specifically the fix for mixed deltas that carry both - reasoning and content simultaneously.""" + reasoning and content simultaneously. + """ @pytest.mark.asyncio async def test_mixed_delta_reasoning_and_content_emits_reasoning_delta( @@ -1341,8 +1343,8 @@ class TestStreamingReasoningToContentTransition: and content set (e.g. reasoning end and content start in the same chunk), the trailing reasoning text must be emitted as a ResponseReasoningTextDeltaEvent and included in the - ResponseReasoningTextDoneEvent text.""" - + ResponseReasoningTextDoneEvent text. + """ monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False) serving = _make_serving_instance_with_reasoning() @@ -1411,8 +1413,8 @@ class TestStreamingReasoningToContentTransition: self, monkeypatch ): """When the transition from reasoning to content is clean (no mixed - delta), no extra reasoning delta event should be emitted.""" - + delta), no extra reasoning delta event should be emitted. + """ monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False) serving = _make_serving_instance_with_reasoning() @@ -1473,8 +1475,8 @@ class TestStreamingReasoningToContentTransition: async def test_reasoning_only_stream_no_content(self, monkeypatch): """When the stream has only reasoning deltas and no content, the reasoning done event should be emitted at finalization with the - full accumulated text, and no text delta events should appear.""" - + full accumulated text, and no text delta events should appear. + """ monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False) serving = _make_serving_instance_with_reasoning() diff --git a/tests/entrypoints/openai/responses/test_simple.py b/tests/entrypoints/openai/responses/test_simple.py index ba55a58e7d..1880c06392 100644 --- a/tests/entrypoints/openai/responses/test_simple.py +++ b/tests/entrypoints/openai/responses/test_simple.py @@ -143,7 +143,8 @@ async def test_streaming_output_consistency(client: OpenAI, model_name: str): @pytest.mark.parametrize("model_name", [MODEL_NAME]) async def test_streaming_logprobs(client: OpenAI, model_name: str): """Test that streaming with logprobs returns valid logprob data on - output_text.delta events and that top_logprobs has the requested count.""" + output_text.delta events and that top_logprobs has the requested count. + """ response = await client.responses.create( model=model_name, input="Say hello.", diff --git a/tests/entrypoints/openai/responses/test_streaming_events.py b/tests/entrypoints/openai/responses/test_streaming_events.py index e1a8bf46f6..2e9b47133d 100644 --- a/tests/entrypoints/openai/responses/test_streaming_events.py +++ b/tests/entrypoints/openai/responses/test_streaming_events.py @@ -124,7 +124,8 @@ class TestProcessorCompoundDeltas: def test_split_name_and_args_same_index(self): """Regression: parsers like KimiK2 emit name and args as separate - DeltaToolCalls at the same index within one DeltaMessage.""" + DeltaToolCalls at the same index within one DeltaMessage. + """ tc_name = _make_tool_call(0, name="get_weather") tc_args = _make_tool_call(0, arguments='{"city":"SF"}') delta = DeltaMessage(tool_calls=[tc_name, tc_args]) @@ -140,7 +141,8 @@ class TestProcessorCompoundDeltas: def test_reasoning_to_content_transition(self): """Regression: the old special case in emit_delta handled this; - now split_delta handles it generically.""" + now split_delta handles it generically. + """ processor = SimpleStreamingEventProcessor() _run_through_processor(processor, DeltaMessage(reasoning="think")) assert processor.state.current_state == _StateType.REASONING diff --git a/tests/entrypoints/openai/test_openai_schema.py b/tests/entrypoints/openai/test_openai_schema.py index 6d3fc2f447..18b01b85a8 100644 --- a/tests/entrypoints/openai/test_openai_schema.py +++ b/tests/entrypoints/openai/test_openai_schema.py @@ -83,8 +83,7 @@ def before_generate_case(context: schemathesis.HookContext, strategy): assert op is not None def no_invalid_types(case: schemathesis.Case): - """ - Skips tool_calls with `"type": "custom"` which schemathesis incorrectly + """Skips tool_calls with `"type": "custom"` which schemathesis incorrectly generates instead of the valid `"type": "function"`. Example test case that is skipped: diff --git a/tests/entrypoints/openai/test_render_token_offsets.py b/tests/entrypoints/openai/test_render_token_offsets.py index a2e66b7bd6..6ef66ee37b 100644 --- a/tests/entrypoints/openai/test_render_token_offsets.py +++ b/tests/entrypoints/openai/test_render_token_offsets.py @@ -25,7 +25,8 @@ def _model_config() -> Mock: def test_completion_flag_forwarded_to_tok_params(): """build_tok_params must forward return_token_offsets, defaulting to False (zero behavioral change for existing callers) and coercing JSON - null to False via the bool() guard.""" + null to False via the bool() guard. + """ cfg = _model_config() default = CompletionRequest(model="m", prompt="hi") @@ -40,7 +41,8 @@ def test_completion_flag_forwarded_to_tok_params(): def test_chat_flag_forwarded_to_tok_params(): """Chat build_tok_params has its own (max_completion_tokens) branch, so - its return_token_offsets forwarding is verified independently.""" + its return_token_offsets forwarding is verified independently. + """ cfg = _model_config() messages = [{"role": "user", "content": "hi"}] @@ -64,7 +66,8 @@ def test_generate_request_token_offsets_default_none(): def test_generate_request_token_offsets_survive_json_round_trip(): """GenerateRequest crosses the disagg serialization boundary; the - tuple[int, int] offsets must survive model_dump and re-validate.""" + tuple[int, int] offsets must survive model_dump and re-validate. + """ req = GenerateRequest( token_ids=[10, 20], sampling_params=SamplingParams(), diff --git a/tests/entrypoints/openai/test_return_token_ids.py b/tests/entrypoints/openai/test_return_token_ids.py index 8537082e3f..9fc8c102d4 100644 --- a/tests/entrypoints/openai/test_return_token_ids.py +++ b/tests/entrypoints/openai/test_return_token_ids.py @@ -189,8 +189,7 @@ async def test_chat_completion_with_tool_use(server): @pytest.mark.asyncio async def test_comparison_with_prompt_logprobs_and_logprobs(server): - """ - Test that token_ids align with prompt_logprobs and + """Test that token_ids align with prompt_logprobs and logprobs when return_tokens_as_token_ids is enabled. """ async with server.get_async_client() as client: diff --git a/tests/entrypoints/openai/test_stop_token_ids.py b/tests/entrypoints/openai/test_stop_token_ids.py index 5e9d79f686..d6ac2cf190 100644 --- a/tests/entrypoints/openai/test_stop_token_ids.py +++ b/tests/entrypoints/openai/test_stop_token_ids.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for stop_token_ids propagation from default_sampling_params +"""Unit tests for stop_token_ids propagation from default_sampling_params to SamplingParams in ChatCompletionRequest and CompletionRequest. Regression test for https://github.com/vllm-project/vllm/issues/22519 diff --git a/tests/entrypoints/openai/utils.py b/tests/entrypoints/openai/utils.py index 4cc4800bf8..02a4c51f94 100644 --- a/tests/entrypoints/openai/utils.py +++ b/tests/entrypoints/openai/utils.py @@ -16,8 +16,7 @@ from vllm.entrypoints.serve.engine.protocol import UsageInfo async def accumulate_streaming_response( stream_generator: AsyncGenerator[str, None], ) -> ChatCompletionResponse: - """ - Accumulate streaming SSE chunks into a complete ChatCompletionResponse. + """Accumulate streaming SSE chunks into a complete ChatCompletionResponse. This helper parses the SSE format and builds up the complete response by combining all the delta chunks. diff --git a/tests/entrypoints/pooling/basic/test_tiling_engine.py b/tests/entrypoints/pooling/basic/test_tiling_engine.py index 88c9b434b3..c620816dc5 100644 --- a/tests/entrypoints/pooling/basic/test_tiling_engine.py +++ b/tests/entrypoints/pooling/basic/test_tiling_engine.py @@ -30,8 +30,7 @@ def llm(vllm_runner): @pytest.mark.skip_global_cleanup def test_tiling_engine_basic(llm): - """ - Basic test with a small number of prompts (less than max_num_seqs). + """Basic test with a small number of prompts (less than max_num_seqs). No tiling should be triggered, but the engine still processes correctly. """ prompts = ["Hello", "World"] @@ -41,8 +40,7 @@ def test_tiling_engine_basic(llm): @pytest.mark.skip_global_cleanup def test_tiling_engine_many_requests(llm): - """ - Test with a large number of prompts that exceeds max_num_seqs. + """Test with a large number of prompts that exceeds max_num_seqs. This verifies that _run_tiling_engine correctly chunks requests, processes all of them, and returns outputs in the correct order. """ @@ -54,8 +52,7 @@ def test_tiling_engine_many_requests(llm): @pytest.mark.skip_global_cleanup def test_tiling_engine_with_pooling_params(llm): - """ - Test the tiling engine when different PoolingParams are provided. + """Test the tiling engine when different PoolingParams are provided. The engine must handle a list of params that matches the number of prompts. """ num_prompts = 10 @@ -77,8 +74,7 @@ def test_tiling_engine_with_pooling_params(llm): @pytest.mark.skip_global_cleanup def test_tiling_engine_abort_on_exception(llm): - """ - Test that abort_request IS called with the correct arguments when an + """Test that abort_request IS called with the correct arguments when an exception occurs inside the engine's step() loop. """ prompts = ["Prompt 0", "Prompt 1", "Prompt 2"] diff --git a/tests/entrypoints/pooling/embed/test_cohere_online.py b/tests/entrypoints/pooling/embed/test_cohere_online.py index 4964d99e0c..29276cb050 100644 --- a/tests/entrypoints/pooling/embed/test_cohere_online.py +++ b/tests/entrypoints/pooling/embed/test_cohere_online.py @@ -111,7 +111,8 @@ def test_basic_embed(server: RemoteOpenAIServer, model_name: str): def test_unsupported_input_type_rejected(server: RemoteOpenAIServer, model_name: str): """An input_type not defined in the model's prompt config should be - rejected with a 400 error.""" + rejected with a 400 error. + """ body = { "model": model_name, "input_type": "nonexistent_type", @@ -138,7 +139,8 @@ def test_omitted_input_type_accepted(server: RemoteOpenAIServer, model_name: str def test_v1_v2_parity(server: RemoteOpenAIServer, model_name: str): """v1 (OpenAI) and v2 (Cohere) endpoints should produce the same - float embeddings for a generic model.""" + float embeddings for a generic model. + """ texts = ["hello world"] v2 = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"]) v1 = _openai_embed(server, model_name, texts) diff --git a/tests/entrypoints/pooling/embed/test_cohere_openai_parity.py b/tests/entrypoints/pooling/embed/test_cohere_openai_parity.py index 3e168adbdb..5608cbcb4e 100644 --- a/tests/entrypoints/pooling/embed/test_cohere_openai_parity.py +++ b/tests/entrypoints/pooling/embed/test_cohere_openai_parity.py @@ -76,7 +76,8 @@ def test_single_text_parity(server: RemoteOpenAIServer): def test_batch_parity(server: RemoteOpenAIServer): """A batch of texts should produce equivalent embeddings via both APIs, - in the same order.""" + in the same order. + """ texts = [ "machine learning", "deep learning", diff --git a/tests/entrypoints/pooling/embed/test_io_processor.py b/tests/entrypoints/pooling/embed/test_io_processor.py index f06e43d143..b002aae8db 100644 --- a/tests/entrypoints/pooling/embed/test_io_processor.py +++ b/tests/entrypoints/pooling/embed/test_io_processor.py @@ -241,7 +241,8 @@ class TestResolveTruncation: def test_truncate_none_with_max_tokens(self): """truncate=NONE should NOT set truncate_prompt_tokens; the - max_tokens limit is enforced separately via _check_max_tokens.""" + max_tokens limit is enforced separately via _check_max_tokens. + """ req = self._make_request(truncate="NONE", max_tokens=10) tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req) assert tokens is None diff --git a/tests/entrypoints/pooling/embed/test_online_dimensions.py b/tests/entrypoints/pooling/embed/test_online_dimensions.py index 7abec99ec3..67c6d7b134 100644 --- a/tests/entrypoints/pooling/embed/test_online_dimensions.py +++ b/tests/entrypoints/pooling/embed/test_online_dimensions.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Run `pytest tests/entrypoints/openai/test_embedding_dimensions.py`. -""" +"""Run `pytest tests/entrypoints/openai/test_embedding_dimensions.py`.""" import openai import pytest diff --git a/tests/entrypoints/pooling/embed/test_online_long_text.py b/tests/entrypoints/pooling/embed/test_online_long_text.py index 60cacb17f8..aeb82336fa 100644 --- a/tests/entrypoints/pooling/embed/test_online_long_text.py +++ b/tests/entrypoints/pooling/embed/test_online_long_text.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test cases for long text embedding with automatic chunking mechanism. +"""Test cases for long text embedding with automatic chunking mechanism. This test suite validates vLLM's automatic chunking functionality for handling text inputs that exceed the model's maximum token length, specifically targeting @@ -241,8 +240,8 @@ async def test_long_text_embedding_1500_chars( client_with_chunked_processing: openai.AsyncOpenAI, model_name: str ): """Test embedding processing for ~1500 character long text - (~1028 tokens, exceeding 512 token limit).""" - + (~1028 tokens, exceeding 512 token limit). + """ # Verify text length # Verify text has sufficient word count (approximately 1500 words) word_count = len(LONG_TEXT_1500_WORDS.split()) @@ -290,8 +289,8 @@ async def test_long_text_embedding_2500_chars( client_with_chunked_processing: openai.AsyncOpenAI, model_name: str ): """Test embedding processing for ~2500 character long text - (~2048 tokens, requiring multiple chunks).""" - + (~2048 tokens, requiring multiple chunks). + """ # Verify text length # Verify text has sufficient word count (approximately 2500 words) word_count = len(LONG_TEXT_2500_WORDS.split()) @@ -339,7 +338,6 @@ async def test_batch_long_text_embedding( client_with_chunked_processing: openai.AsyncOpenAI, model_name: str ): """Test batch long text embedding processing.""" - input_texts = [ LONG_TEXT_1500_WORDS, LONG_TEXT_2500_WORDS, @@ -384,8 +382,8 @@ async def test_chunked_vs_normal_consistency( client_with_chunked_processing: openai.AsyncOpenAI, model_name: str ): """Test consistency between chunked and - normal processing (using short text).""" - + normal processing (using short text). + """ # Use a short text within the 512 token limit short_text = ( "Artificial intelligence technology is changing our world, " @@ -424,7 +422,6 @@ async def test_chunked_processing_response_format( client_with_chunked_processing: openai.AsyncOpenAI, model_name: str ): """Test response format and structure during chunked processing.""" - # Test with long text to trigger chunking embedding_response = await client_with_chunked_processing.embeddings.create( model=model_name, diff --git a/tests/entrypoints/pooling/embed/test_protocol.py b/tests/entrypoints/pooling/embed/test_protocol.py index 91c853f910..24d4f9e6d7 100644 --- a/tests/entrypoints/pooling/embed/test_protocol.py +++ b/tests/entrypoints/pooling/embed/test_protocol.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for Cohere embed protocol: build_typed_embeddings and its -underlying packing helpers, plus Cohere-specific serving helpers.""" +underlying packing helpers, plus Cohere-specific serving helpers. +""" import struct diff --git a/tests/entrypoints/scale_out/derender/test_derender.py b/tests/entrypoints/scale_out/derender/test_derender.py index a30e487450..db702b4b07 100644 --- a/tests/entrypoints/scale_out/derender/test_derender.py +++ b/tests/entrypoints/scale_out/derender/test_derender.py @@ -454,7 +454,8 @@ async def test_derender_completion_empty_generate_responses(client): @pytest.mark.asyncio async def test_derender_completion_logprobs(client): """token_id:N placeholders in logprobs are resolved; CompletionLogProbs - flat-list structure is returned with non-empty tokens and text_offsets.""" + flat-list structure is returned with non-empty tokens and text_offsets. + """ gr1 = await _render_completion(client, "Hello world") ids1 = gr1["token_ids"][:3] token_id = ids1[0] @@ -555,7 +556,7 @@ async def test_derender_chat_oversized_token_ids_rejected(client): @pytest.mark.asyncio async def test_derender_chat_too_many_choices_rejected(client): - """choices count exceeding VLLM_MAX_N_SEQUENCES returns 400.""" + """Choices count exceeding VLLM_MAX_N_SEQUENCES returns 400.""" # Default VLLM_MAX_N_SEQUENCES is 16384; use a larger count. oversized_choices = [ {"index": i, "token_ids": [42], "finish_reason": "stop"} for i in range(20_000) diff --git a/tests/entrypoints/scale_out/derender/test_derender_parity.py b/tests/entrypoints/scale_out/derender/test_derender_parity.py index adece09c0a..6bcede3199 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_parity.py +++ b/tests/entrypoints/scale_out/derender/test_derender_parity.py @@ -121,7 +121,8 @@ async def _disagg( def _tool_sig(response_choice: dict) -> list[tuple[str, dict]]: """[(name, json normalized args)] so key ordering / whitespace don't - cause false negatives.""" + cause false negatives. + """ return [ (tc["function"]["name"], json.loads(tc["function"]["arguments"])) for tc in (response_choice["message"].get("tool_calls") or []) @@ -143,7 +144,8 @@ async def _run_parity_case( client: httpx.AsyncClient, messages: list[dict], **extra ) -> tuple[dict, dict]: """Run the coupled request then feed its generated tokens into the - disaggregated derender endpoint. Returns (coupled, disagg).""" + disaggregated derender endpoint. Returns (coupled, disagg). + """ coupled = await _coupled(client, messages, **extra) ch = coupled["choices"][0] chat_request = {"model": MODEL, "messages": messages, **extra} @@ -203,7 +205,8 @@ async def test_parity_tool_call(client): @pytest.mark.asyncio async def test_parity_reasoning_and_tool_call(client): """Combined reasoning + tool call parity means the highest drift risk - since it exercises both parser branches on the same output.""" + since it exercises both parser branches on the same output. + """ messages = [{"role": "user", "content": "What's the weather in Paris?"}] coupled, disagg = await _run_parity_case( client, diff --git a/tests/entrypoints/scale_out/derender/test_derender_stream.py b/tests/entrypoints/scale_out/derender/test_derender_stream.py index b783f40ae3..df10b31d28 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_stream.py +++ b/tests/entrypoints/scale_out/derender/test_derender_stream.py @@ -41,7 +41,6 @@ def _make_stream_chunk( usage: dict | None = None, ) -> GenerateStreamResponse: """Build a GenerateStreamResponse SSE chunk.""" - return GenerateStreamResponse( request_id=request_id, choices=[ @@ -61,7 +60,6 @@ def _make_usage_chunk( request_id: str = "test-req", ) -> GenerateStreamResponse: """Build a usage only final SSE chunk (empty choices).""" - return GenerateStreamResponse( request_id=request_id, choices=[], @@ -531,7 +529,6 @@ async def client(server): async def _render_chat(client) -> dict: """Render a minimal chat request and return the GenerateRequest dict.""" - resp = await client.post( "/v1/chat/completions/render", json={ diff --git a/tests/entrypoints/scale_out/render/test_render.py b/tests/entrypoints/scale_out/render/test_render.py index 6aaf119f16..bb6a521ef1 100644 --- a/tests/entrypoints/scale_out/render/test_render.py +++ b/tests/entrypoints/scale_out/render/test_render.py @@ -337,7 +337,6 @@ async def test_chat_completion_render_multi_turn(client): @pytest.mark.asyncio async def test_chat_completion_render_with_stream_true(client): """Render accepts stream params but still returns JSON (non-streamed).""" - response = await client.post( "/v1/chat/completions/render", json={ @@ -456,7 +455,8 @@ async def test_chat_completion_render_with_sampling_params(client): @pytest.mark.asyncio async def test_completion_render_emits_token_offsets(client): """With return_token_offsets, /v1/completions/render returns per-token - (start, end) char offsets aligned with token_ids.""" + (start, end) char offsets aligned with token_ids. + """ prompt = "Hello, world." response = await client.post( "/v1/completions/render", @@ -481,7 +481,8 @@ async def test_completion_render_emits_token_offsets(client): @pytest.mark.asyncio async def test_completion_render_default_no_token_offsets(client): """Without the flag, token_offsets must be null (existing responses - unchanged).""" + unchanged). + """ response = await client.post( "/v1/completions/render", json={ @@ -498,7 +499,8 @@ async def test_completion_render_default_no_token_offsets(client): @pytest.mark.asyncio async def test_chat_render_emits_token_offsets(client): """With return_token_offsets, /v1/chat/completions/render returns - per-token offsets relative to the templated prompt string.""" + per-token offsets relative to the templated prompt string. + """ response = await client.post( "/v1/chat/completions/render", json={ @@ -672,7 +674,8 @@ async def test_chat_completion_render_assistant_tokens_mask_with_generation_tags client, ): """With a ``{% generation %}``-enabled template, the mask marks assistant - tokens and the masked tokens decode to the assistant content.""" + tokens and the masked tokens decode to the assistant content. + """ response = await client.post( "/v1/chat/completions/render", json={ diff --git a/tests/entrypoints/scale_out/render/test_render_multimodal.py b/tests/entrypoints/scale_out/render/test_render_multimodal.py index c6a6931948..eb45ce68ea 100644 --- a/tests/entrypoints/scale_out/render/test_render_multimodal.py +++ b/tests/entrypoints/scale_out/render/test_render_multimodal.py @@ -16,7 +16,6 @@ VISION_MODEL_NAME = "Qwen/Qwen3-VL-2B-Instruct" @pytest.fixture(scope="module") def vision_server(): """Vision-capable server used for multimodal /render tests.""" - args = [ "--enforce-eager", "--max-model-len", @@ -48,7 +47,6 @@ async def test_chat_completion_render_with_base64_image_url( local_asset_server, ): """Render a multimodal chat request and verify tokens are returned.""" - image = local_asset_server.get_image_asset("RGBA_comp.png") data_url = encode_image_url(image, format="PNG") @@ -112,7 +110,6 @@ async def test_tokenize_matches_render_for_multimodal_input( local_asset_server, ): """`/tokenize` should match `/v1/chat/completions/render` token output.""" - image = local_asset_server.get_image_asset("RGBA_comp.png") data_url = encode_image_url(image, format="PNG") diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py b/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py index 764d684445..5cd60a959f 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Roundtrip tests for multimodal serde used by the +"""Roundtrip tests for multimodal serde used by the token_in_token_out generate endpoint. """ diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py b/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py index e51530ce45..1dcf93330e 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py @@ -25,7 +25,8 @@ def _base_payload() -> dict: def test_omitted_max_tokens_is_not_provided(): """Body without ``max_tokens`` must surface as 'not provided' so the - server can apply its own default instead of the dataclass 16.""" + server can apply its own default instead of the dataclass 16. + """ req = GenerateRequest.model_validate(_base_payload()) # SamplingParams' dataclass default leaks through the parsed instance — # this is exactly the bug the server-side defaulting works around. @@ -35,7 +36,8 @@ def test_omitted_max_tokens_is_not_provided(): def test_explicit_max_tokens_is_provided(): """Even when the client picks the same value as the dataclass default, - it must register as explicitly set so the server won't override it.""" + it must register as explicitly set so the server won't override it. + """ payload = _base_payload() payload["sampling_params"] = {"max_tokens": 16} req = GenerateRequest.model_validate(payload) @@ -67,7 +69,8 @@ def test_json_roundtrip_preserves_provided_keys(): def test_internal_instance_construction_treats_all_as_provided(): """When internal callers build ``GenerateRequest`` from a pre-resolved ``SamplingParams`` instance, every field is considered explicitly set - so server-side defaulting can't clobber values resolved upstream.""" + so server-side defaulting can't clobber values resolved upstream. + """ sp = SamplingParams(max_tokens=500, temperature=0.0) req = GenerateRequest(token_ids=[1, 2, 3], sampling_params=sp) assert req.is_sampling_param_provided("max_tokens") diff --git a/tests/entrypoints/serve/dev/rlhf/conftest.py b/tests/entrypoints/serve/dev/rlhf/conftest.py index 1d2c598b19..c636b51572 100644 --- a/tests/entrypoints/serve/dev/rlhf/conftest.py +++ b/tests/entrypoints/serve/dev/rlhf/conftest.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Shared fixtures and helpers for the RL lifecycle test suite. +"""Shared fixtures and helpers for the RL lifecycle test suite. All test modules under this directory import from here to avoid duplication. @@ -92,6 +91,7 @@ def server( port: HTTP port to bind (caller is responsible for uniqueness). timeout: Seconds to wait for /health before giving up. dummy_weights: If True, use --load-format dummy (fast, no real weights). + """ env = {**os.environ, "VLLM_SERVER_DEV_MODE": "1"} base = _DUMMY_ARGS if dummy_weights else _BASE_ARGS diff --git a/tests/entrypoints/serve/dev/rpc/test_collective_rpc.py b/tests/entrypoints/serve/dev/rpc/test_collective_rpc.py index eb9aa7663c..7cf6f4490c 100644 --- a/tests/entrypoints/serve/dev/rpc/test_collective_rpc.py +++ b/tests/entrypoints/serve/dev/rpc/test_collective_rpc.py @@ -25,7 +25,7 @@ class TestWorkerExtension: ) def return_none(self, *args, **kwargs) -> None: - """Test method that does not return anything""" + """Test method that does not return anything.""" return @@ -48,7 +48,7 @@ def server(): def test_get_model_name(server): - """Test basic response""" + """Test basic response.""" response = requests.post( server.url_for("collective_rpc"), json={"method": "get_model_name"} ) @@ -59,7 +59,7 @@ def test_get_model_name(server): def test_return_none(server): - """Test return none""" + """Test return none.""" response = requests.post( server.url_for("collective_rpc"), json={"method": "return_none"} ) @@ -69,7 +69,7 @@ def test_return_none(server): def test_echo_args_kwargs(server): - """Test args, kwargs, and dict response""" + """Test args, kwargs, and dict response.""" args = ["arg1", "arg2"] kwargs = {"key1": "value1", "key2": "value2"} response = requests.post( diff --git a/tests/entrypoints/serve/dev/test_sleep.py b/tests/entrypoints/serve/dev/test_sleep.py index 260dcd00ba..5ebee50600 100644 --- a/tests/entrypoints/serve/dev/test_sleep.py +++ b/tests/entrypoints/serve/dev/test_sleep.py @@ -87,8 +87,7 @@ def test_sleep_mode(): def _get_sleep_metrics_from_api(response: requests.Response): - """Return (awake, weights_offloaded, discard_all)""" - + """Return (awake, weights_offloaded, discard_all).""" awake, weights_offloaded, discard_all = None, None, None for family in text_string_to_metric_families(response.text): diff --git a/tests/entrypoints/serve/exception_handling/test_error_sanitization.py b/tests/entrypoints/serve/exception_handling/test_error_sanitization.py index 426d1ede74..cb14260ab8 100644 --- a/tests/entrypoints/serve/exception_handling/test_error_sanitization.py +++ b/tests/entrypoints/serve/exception_handling/test_error_sanitization.py @@ -22,7 +22,8 @@ def test_sanitize_message(): class TestSanitizeMessageFilePaths: """sanitize_message should also strip file paths and traceback - frames, not just memory addresses - see #31683.""" + frames, not just memory addresses - see #31683. + """ def test_strips_traceback_style_frame(self): msg = ( @@ -44,7 +45,8 @@ class TestSanitizeMessageFilePaths: def test_strips_single_parent_container_path(self): """Regression: /app/server.py and /workspace/server.py (common in - container deployments) were missed by the original {2,} quantifier.""" + container deployments) were missed by the original {2,} quantifier. + """ assert "/app/" not in sanitize_message("Error in /app/server.py") assert "/workspace/" not in sanitize_message("Error in /workspace/server.py") diff --git a/tests/entrypoints/serve/exception_handling/test_validation_exception_handler.py b/tests/entrypoints/serve/exception_handling/test_validation_exception_handler.py index 942dd95c18..bd4e87a155 100644 --- a/tests/entrypoints/serve/exception_handling/test_validation_exception_handler.py +++ b/tests/entrypoints/serve/exception_handling/test_validation_exception_handler.py @@ -24,7 +24,8 @@ from vllm.entrypoints.serve.exception_handling.handlers.validation import ( def _fake_request(log_error_stack: bool = False) -> SimpleNamespace: """Minimal stand-in for a FastAPI Request - just enough for the - handler to read req.app.state.args.log_error_stack.""" + handler to read req.app.state.args.log_error_stack. + """ return SimpleNamespace( app=SimpleNamespace( state=SimpleNamespace(args=SimpleNamespace(log_error_stack=log_error_stack)) @@ -35,7 +36,8 @@ def _fake_request(log_error_stack: bool = False) -> SimpleNamespace: class TestValidationErrorParamFallback: """Ensure `param` falls back to the Pydantic error's `loc` when no - custom VLLMValidationError context is present.""" + custom VLLMValidationError context is present. + """ @pytest.mark.parametrize( ("error_type", "msg"), @@ -59,7 +61,8 @@ class TestValidationErrorParamFallback: async def test_param_fallback_does_not_crash_on_non_dict_error(self): """Schemathesis fuzzing found that errors[0] isn't always a dict. The fallback must not crash in that case - it should just leave - param as None instead of raising.""" + param as None instead of raising. + """ exc = RequestValidationError(["some unexpected non-dict error"]) response = await validation_exception_handler(_fake_request(), exc) @@ -97,7 +100,8 @@ class TestCleanLocForParam: def test_all_internal_falls_back_to_raw_join(self): """If every loc segment looks internal, fall back to the raw - dot-join rather than returning an empty string.""" + dot-join rather than returning an empty string. + """ loc = ("function-wrap[__log_extra_fields__()]",) assert clean_loc_for_param(loc) == "function-wrap[__log_extra_fields__()]" @@ -105,7 +109,8 @@ class TestCleanLocForParam: class TestValidationErrorDoesNotLeakServerPaths: """FastAPI sets endpoint_file/line/function/path on the exception during routing, not in the constructor - so we set them manually - here to reproduce the real leak. See #31683.""" + here to reproduce the real leak. See #31683. + """ @pytest.mark.asyncio async def test_handler_strips_endpoint_file_context(self): diff --git a/tests/entrypoints/serve/instrumentator/test_basic.py b/tests/entrypoints/serve/instrumentator/test_basic.py index 73a97c4fa2..0711cbee95 100644 --- a/tests/entrypoints/serve/instrumentator/test_basic.py +++ b/tests/entrypoints/serve/instrumentator/test_basic.py @@ -21,7 +21,7 @@ MODEL_NAME = "Qwen/Qwen3-0.6B" @pytest.fixture(scope="module") def server_args(request: pytest.FixtureRequest) -> list[str]: - """Provide extra arguments to the server via indirect parametrization + """Provide extra arguments to the server via indirect parametrization. Usage: diff --git a/tests/entrypoints/serve/instrumentator/test_metrics.py b/tests/entrypoints/serve/instrumentator/test_metrics.py index 297fe233b5..ac4576fbb9 100644 --- a/tests/entrypoints/serve/instrumentator/test_metrics.py +++ b/tests/entrypoints/serve/instrumentator/test_metrics.py @@ -403,8 +403,7 @@ async def _poll_until( def _get_running_metrics_from_api(server: RemoteOpenAIServer): - """Return (running_count, waiting_count, kv_cache_usage)""" - + """Return (running_count, waiting_count, kv_cache_usage).""" response = requests.get(server.url_for("metrics")) assert response.status_code == HTTPStatus.OK diff --git a/tests/entrypoints/serve/lora/test_lora_adapters.py b/tests/entrypoints/serve/lora/test_lora_adapters.py index 8a860dc625..79356fc951 100644 --- a/tests/entrypoints/serve/lora/test_lora_adapters.py +++ b/tests/entrypoints/serve/lora/test_lora_adapters.py @@ -245,7 +245,8 @@ async def test_multiple_lora_adapters( client: openai.AsyncOpenAI, tmp_path, qwen3_lora_files ): """Validate that many loras can be dynamically registered and inferenced - with concurrently""" + with concurrently + """ # This test file configures the server with --max-cpu-loras=2 and this test # will concurrently load 10 adapters, so it should flex the LRU cache diff --git a/tests/entrypoints/serve/middleware/test_authentication_middleware.py b/tests/entrypoints/serve/middleware/test_authentication_middleware.py index 528261e644..ee22407fdd 100644 --- a/tests/entrypoints/serve/middleware/test_authentication_middleware.py +++ b/tests/entrypoints/serve/middleware/test_authentication_middleware.py @@ -72,7 +72,8 @@ class MockModelConfig: @pytest.fixture(params=get_args(SupportedTask)) def task_routes(request, monkeypatch) -> tuple[str, list[tuple[str, list[str]]]]: """For each supported task, build an app with only that task's routers, - extract all routes, and return the task name and routes.""" + extract all routes, and return the task name and routes. + """ task = request.param # Enable development mode to register all routes (including dev-only routes). monkeypatch.setenv("VLLM_SERVER_DEV_MODE", "1") @@ -98,7 +99,8 @@ def task_routes(request, monkeypatch) -> tuple[str, list[tuple[str, list[str]]]] def test_auto_discovered_protected_routes_require_auth(task_routes): """For every auto-discovered route that starts with a guarded prefix, - verify that authentication is enforced.""" + verify that authentication is enforced. + """ task, routes = task_routes app = _create_app_with_mock_routes(routes) client = TestClient(app) @@ -132,7 +134,8 @@ def test_auto_discovered_protected_routes_require_auth(task_routes): def test_auto_discovered_unprotected_routes_no_auth(task_routes): """For every auto-discovered route that does NOT start with a guarded - prefix, verify that no authentication is required.""" + prefix, verify that no authentication is required. + """ task, routes = task_routes app = _create_app_with_mock_routes(routes) client = TestClient(app) diff --git a/tests/entrypoints/serve/middleware/test_optional_middleware.py b/tests/entrypoints/serve/middleware/test_optional_middleware.py index 865bd445c8..160f1e248d 100644 --- a/tests/entrypoints/serve/middleware/test_optional_middleware.py +++ b/tests/entrypoints/serve/middleware/test_optional_middleware.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for middleware that's off by default and can be toggled through +"""Tests for middleware that's off by default and can be toggled through server arguments, mainly --api-key and --enable-request-id-headers. """ diff --git a/tests/entrypoints/serve/sagemaker/test_sagemaker_handler_overrides.py b/tests/entrypoints/serve/sagemaker/test_sagemaker_handler_overrides.py index 20b917400b..b6137dae93 100644 --- a/tests/entrypoints/serve/sagemaker/test_sagemaker_handler_overrides.py +++ b/tests/entrypoints/serve/sagemaker/test_sagemaker_handler_overrides.py @@ -114,7 +114,8 @@ class TestHandlerOverrideIntegration: def test_customer_script_functions_auto_loaded(self, monkeypatch, tmp_path): """Test customer scenario: script functions automatically override - framework defaults.""" + framework defaults. + """ try: from model_hosting_container_standards.sagemaker.config import ( SageMakerEnvVars, @@ -175,7 +176,8 @@ async def custom_sagemaker_invocation_handler(request: Request): def test_customer_decorator_usage(self, monkeypatch, tmp_path): """Test customer scenario: using @custom_ping_handler and - @custom_invocation_handler decorators.""" + @custom_invocation_handler decorators. + """ try: from model_hosting_container_standards.sagemaker.config import ( SageMakerEnvVars, @@ -231,7 +233,8 @@ async def my_invoke(request: Request): def test_handler_priority_order(self, monkeypatch, tmp_path): """Test priority: @custom_ping_handler/@custom_invocation_handler - decorators vs script functions.""" + decorators vs script functions. + """ try: from model_hosting_container_standards.sagemaker.config import ( SageMakerEnvVars, @@ -304,7 +307,8 @@ async def custom_sagemaker_invocation_handler(request: Request): def test_environment_variable_script_loading(self, monkeypatch, tmp_path): """Test that environment variables correctly specify script location - and loading.""" + and loading. + """ try: from model_hosting_container_standards.sagemaker.config import ( SageMakerEnvVars, @@ -416,7 +420,8 @@ async def custom_sagemaker_invocation_handler(request: Request): def test_handler_env_var_override(self, monkeypatch, tmp_path): """Test CUSTOM_FASTAPI_PING_HANDLER and CUSTOM_FASTAPI_INVOCATION_HANDLER - environment variable overrides.""" + environment variable overrides. + """ try: from model_hosting_container_standards.common.fastapi.config import ( FastAPIEnvVars, @@ -510,7 +515,8 @@ async def custom_sagemaker_invocation_handler(request: Request): def test_env_var_priority_over_decorator_and_script(self, monkeypatch, tmp_path): """Test that environment variables have highest priority over decorators - and script functions for both ping and invocation handlers.""" + and script functions for both ping and invocation handlers. + """ try: from model_hosting_container_standards.common.fastapi.config import ( FastAPIEnvVars, diff --git a/tests/entrypoints/serve/tokenize/test_tokenization_vlm.py b/tests/entrypoints/serve/tokenize/test_tokenization_vlm.py index 6b226c6999..c4317973cb 100644 --- a/tests/entrypoints/serve/tokenize/test_tokenization_vlm.py +++ b/tests/entrypoints/serve/tokenize/test_tokenization_vlm.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Regression test: ``/tokenize`` must expand image placeholders for VLM models. +"""Regression test: ``/tokenize`` must expand image placeholders for VLM models. Fixed by PR #34560 ("Move InputPreprocessor into Renderer (2/2)"). Before that change, ``/tokenize`` returned ~26 tokens for a message with an diff --git a/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py b/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py index 50083e3455..49e181e36e 100644 --- a/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py +++ b/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Regression test: calling ``/tokenize`` with multimodal data followed by +"""Regression test: calling ``/tokenize`` with multimodal data followed by ``/v1/chat/completions`` with the same data must not cause an error. Ensures that the ``/tokenize`` endpoint does not pollute internal caches @@ -52,8 +51,8 @@ async def test_tokenize_then_chat_completion_with_image( local_asset_server, ): """Tokenize a multimodal message, then send the same message to chat - completions. The chat completion must succeed (not 500).""" - + completions. The chat completion must succeed (not 500). + """ image_url = local_asset_server.url_for("stop_sign.jpg") messages = [ { diff --git a/tests/entrypoints/serve/utils/test_api_utils.py b/tests/entrypoints/serve/utils/test_api_utils.py index a41f39187f..051978b60c 100644 --- a/tests/entrypoints/serve/utils/test_api_utils.py +++ b/tests/entrypoints/serve/utils/test_api_utils.py @@ -36,11 +36,13 @@ def test_should_include_usage_force_enables_continuous_usage(stream_options, exp class TestGetMaxTokens: """Tests for get_max_tokens() to ensure generation_config's max_tokens acts as a default when from model author, and as a ceiling when - explicitly set by the user.""" + explicitly set by the user. + """ def test_default_sampling_params_used_when_no_request_max_tokens(self): """When user doesn't specify max_tokens, generation_config default - should apply.""" + should apply. + """ result = get_max_tokens( max_model_len=24000, max_tokens=None, @@ -51,7 +53,8 @@ class TestGetMaxTokens: def test_request_max_tokens_not_capped_by_default_sampling_params(self): """When user specifies max_tokens in request, model author's - generation_config max_tokens must NOT cap it (fixes #34005).""" + generation_config max_tokens must NOT cap it (fixes #34005). + """ result = get_max_tokens( max_model_len=24000, max_tokens=5000, @@ -94,7 +97,8 @@ class TestGetMaxTokens: def test_request_max_tokens_smaller_than_default(self): """When user explicitly requests fewer tokens than gen_config default, - that should be respected.""" + that should be respected. + """ result = get_max_tokens( max_model_len=24000, max_tokens=512, diff --git a/tests/entrypoints/serve/utils/test_request_logger.py b/tests/entrypoints/serve/utils/test_request_logger.py index 717ff190b3..5df3c3cd84 100644 --- a/tests/entrypoints/serve/utils/test_request_logger.py +++ b/tests/entrypoints/serve/utils/test_request_logger.py @@ -239,7 +239,8 @@ def test_request_logger_log_outputs_integration(): def test_streaming_complete_logs_full_text_content(): """Test that streaming complete logging includes - full accumulated text, not just token count.""" + full accumulated text, not just token count. + """ mock_logger = MagicMock() with patch("vllm.entrypoints.serve.utils.request_logger.logger", mock_logger): diff --git a/tests/entrypoints/serve/utils/test_sse_keep_alive.py b/tests/entrypoints/serve/utils/test_sse_keep_alive.py index 519b9bb4a1..6c61080f8f 100644 --- a/tests/entrypoints/serve/utils/test_sse_keep_alive.py +++ b/tests/entrypoints/serve/utils/test_sse_keep_alive.py @@ -17,7 +17,8 @@ from vllm.entrypoints.serve.utils.sse_keep_alive import ( def test_disabled_returns_same_generator(): """For 0, negative, and NaN intervals the exact same generator object must be returned (identity), proving no wrapper or per-chunk overhead when - disabled.""" + disabled. + """ async def gen(): yield "x" @@ -33,7 +34,8 @@ def test_disabled_returns_same_generator(): @pytest.mark.asyncio async def test_keep_alive_during_silent_start(): """A silent start (queue wait / prefill) emits keep-alive comments before - the first real chunk.""" + the first real chunk. + """ release = asyncio.Event() async def gen(): @@ -55,7 +57,8 @@ async def test_keep_alive_during_silent_start(): @pytest.mark.asyncio async def test_no_keep_alive_when_streaming_is_fast(): """If chunks arrive faster than the interval, no comment is emitted and all - data is forwarded.""" + data is forwarded. + """ async def gen(): for chunk in ["a", "b", "c"]: @@ -69,7 +72,8 @@ async def test_no_keep_alive_when_streaming_is_fast(): @pytest.mark.asyncio async def test_keep_alive_between_chunks_does_not_drop_data(): """Keep-alives interleaved with real chunks must never drop or reorder - data.""" + data. + """ release_b = asyncio.Event() release_c = asyncio.Event() @@ -135,7 +139,8 @@ async def test_upstream_cancelled_error_propagates(): @pytest.mark.asyncio async def test_aclose_closes_upstream(): """Closing a wrapper that is mid-stream must complete promptly and close - the upstream generator (regression for the earlier leaked-task hang).""" + the upstream generator (regression for the earlier leaked-task hang). + """ finalized = asyncio.Event() async def gen(): @@ -154,7 +159,8 @@ async def test_aclose_closes_upstream(): @pytest.mark.asyncio async def test_client_disconnect_cancels_upstream(): """Cancelling the consumer mid-stream must cancel the pending anext and - close the upstream generator so its cleanup runs.""" + close the upstream generator so its cleanup runs. + """ finalized = asyncio.Event() async def gen(): @@ -181,7 +187,8 @@ async def test_client_disconnect_cancels_upstream(): @pytest.mark.asyncio async def test_streaming_response_emits_comment(): """A keep-alive comment must reach the client through the real - StreamingResponse used by the routers.""" + StreamingResponse used by the routers. + """ release = asyncio.Event() async def gen(): @@ -220,7 +227,8 @@ async def test_streaming_response_emits_comment(): async def test_streaming_response_disconnect_closes_upstream(): """On the ASGI 2.3 disconnect path used by the pinned Uvicorn, Starlette cancels streaming and the wrapper's finally must close the upstream - generator.""" + generator. + """ finalized = asyncio.Event() async def gen(): diff --git a/tests/entrypoints/speech_to_text/conftest.py b/tests/entrypoints/speech_to_text/conftest.py index 94505bd054..bf65a72ab5 100644 --- a/tests/entrypoints/speech_to_text/conftest.py +++ b/tests/entrypoints/speech_to_text/conftest.py @@ -11,6 +11,7 @@ def add_attention_backend(server_args, attention_config): Args: server_args: List of server arguments to extend in-place. attention_config: Dict with 'backend' key, or None. + """ if attention_config and "backend" in attention_config: server_args.extend(["--attention-backend", attention_config["backend"]]) diff --git a/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py b/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py index 713fa48eee..855dafaaba 100644 --- a/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py +++ b/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Evaluate Transcription API correctness by computing Word Error Rate (WER) +"""Evaluate Transcription API correctness by computing Word Error Rate (WER) on a given ASR dataset. When provided, it will also compare the WER against a baseline. This simulates real work usage of the API and makes sure that the frontend and diff --git a/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py b/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py index 7fb7a5522d..b9ea18bbae 100644 --- a/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py +++ b/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py @@ -295,7 +295,8 @@ async def test_session_update_invalid_model_returns_error(model_name): @pytest.mark.parametrize("model_name", [MODEL_NAME]) async def test_commit_without_session_update_returns_error(model_name): """Test that committing before validating the model returns an error - and does not fall through to processing.""" + and does not fall through to processing. + """ server_args = ["--enforce-eager", "--max-model-len", "2048"] if model_name.startswith("mistralai"): diff --git a/tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py b/tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py index cd61507728..5534845326 100644 --- a/tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py +++ b/tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py @@ -18,9 +18,8 @@ _PATCH = "vllm.entrypoints.speech_to_text.base.serving" @pytest.mark.asyncio async def test_chunk_offsets_are_cumulative_not_nominal(): - """ - chunk_start_offsets must be cumulative actual chunk lengths, not the old approach of - 'idx * max_audio_clip_s'. When split_audio places a boundary before the + """chunk_start_offsets must be cumulative actual chunk lengths, not the old approach + of 'idx * max_audio_clip_s'. When split_audio places a boundary before the nominal 30 s mark, the old formula drifts; the fixed formula stays exact. """ # Chunks shorter than exactly 30 s, as split_audio produces when a quiet diff --git a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py index 51599d359e..d5658e7cf0 100644 --- a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py +++ b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py @@ -238,8 +238,7 @@ async def test_stream_options(winning_call, whisper_client): @pytest.mark.asyncio async def test_sampling_params(mary_had_lamb, whisper_client): - """ - Compare sampling with params and greedy sampling to assert results + """Compare sampling with params and greedy sampling to assert results are different when extreme sampling parameters values are picked. """ transcription = await whisper_client.audio.transcriptions.create( diff --git a/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py b/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py index f03b6e8a38..79b46bea51 100644 --- a/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py +++ b/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py @@ -150,7 +150,6 @@ async def test_non_streaming_tool_call( client: openai.AsyncOpenAI, server_config: ServerConfig ): """Test tool call in non-streaming mode.""" - response = await client.chat.completions.create( model=server_config["model_arg"], messages=MESSAGES, @@ -183,7 +182,6 @@ async def test_streaming_tool_call( client: openai.AsyncOpenAI, server_config: ServerConfig ): """Test tool call in streaming mode.""" - stream = await client.chat.completions.create( model=server_config["model_arg"], messages=MESSAGES, @@ -230,7 +228,6 @@ async def test_non_streaming_product_tool_call( client: openai.AsyncOpenAI, server_config: ServerConfig ): """Test tool call integer and boolean parameters in non-streaming mode.""" - response = await client.chat.completions.create( model=server_config["model_arg"], messages=PRODUCT_MESSAGES, @@ -272,7 +269,6 @@ async def test_streaming_product_tool_call( client: openai.AsyncOpenAI, server_config: ServerConfig ): """Test tool call integer and boolean parameters in streaming mode.""" - stream = await client.chat.completions.create( model=server_config["model_arg"], messages=PRODUCT_MESSAGES, diff --git a/tests/entrypoints/tool_parsers/test_openai_tool_parser.py b/tests/entrypoints/tool_parsers/test_openai_tool_parser.py index d99b66d9ac..de402b435c 100644 --- a/tests/entrypoints/tool_parsers/test_openai_tool_parser.py +++ b/tests/entrypoints/tool_parsers/test_openai_tool_parser.py @@ -127,8 +127,7 @@ FUNC_ARGS_TIME = '{"city": "New York"}' # Utility to extract reasoning and tool calls # ========================================================== def extract_reasoning_and_calls(chunks: list) -> tuple[str, list[str], list[str]]: - """ - Extract accumulated reasoning text and tool call arguments + """Extract accumulated reasoning text and tool call arguments from streaming chunks. """ reasoning: str = "" @@ -165,7 +164,6 @@ def extract_reasoning_and_calls(chunks: list) -> tuple[str, list[str], list[str] @pytest.mark.asyncio async def test_calculator_tool_call_and_argument_accuracy(client: openai.AsyncOpenAI): """Verify calculator tool call is made and arguments are accurate.""" - response = await client.chat.completions.create( model=MODEL_NAME, messages=MESSAGES_CALC, @@ -205,7 +203,6 @@ async def test_calculator_tool_call_and_argument_accuracy(client: openai.AsyncOp @pytest.mark.asyncio async def test_streaming_tool_call_get_time_with_reasoning(client: openai.AsyncOpenAI): """Verify streamed reasoning and tool call behavior for get_time.""" - stream = await client.chat.completions.create( model=MODEL_NAME, messages=MESSAGES_GET_TIME, @@ -258,8 +255,7 @@ async def test_streaming_multiple_tools(client: openai.AsyncOpenAI): @pytest.mark.asyncio async def test_invalid_tool_call(client: openai.AsyncOpenAI): - """ - Verify that ambiguous instructions that should not trigger a tool + """Verify that ambiguous instructions that should not trigger a tool do not produce any tool calls. """ response = await client.chat.completions.create( @@ -283,8 +279,7 @@ async def test_invalid_tool_call(client: openai.AsyncOpenAI): @pytest.mark.asyncio async def test_tool_call_with_temperature(client: openai.AsyncOpenAI): - """ - Verify model produces valid tool or text output + """Verify model produces valid tool or text output under non-deterministic sampling. """ response = await client.chat.completions.create( diff --git a/tests/entrypoints/unit_tests/test_chat_utils.py b/tests/entrypoints/unit_tests/test_chat_utils.py index 630f641df5..ab57140c2b 100644 --- a/tests/entrypoints/unit_tests/test_chat_utils.py +++ b/tests/entrypoints/unit_tests/test_chat_utils.py @@ -940,7 +940,6 @@ def test_parse_chat_messages_audio_embeds_with_string( audio_embeds_model_config, ): """Test audio_embeds with base64 string embedding data.""" - import torch # Create a sample audio embedding tensor @@ -983,7 +982,6 @@ async def test_parse_chat_messages_audio_embeds_async( audio_embeds_model_config, ): """Test audio_embeds with async futures.""" - import torch # Create a sample audio embedding tensor @@ -2986,7 +2984,6 @@ def test_postprocess_messages_null_arguments_string(): @pytest.mark.asyncio async def test_resolve_items_runs_modalities_concurrently_and_preserves_order(): """Media fetches overlap while modality and item order are preserved.""" - active_fetches = 0 max_active_fetches = 0 diff --git a/tests/entrypoints/unit_tests/test_context.py b/tests/entrypoints/unit_tests/test_context.py index 1c1f6ed235..e09687425b 100644 --- a/tests/entrypoints/unit_tests/test_context.py +++ b/tests/entrypoints/unit_tests/test_context.py @@ -51,7 +51,8 @@ async def generate_mock_outputs( num_turns, prompt_token_counts, output_token_counts, cached_token_counts=None ): """Generate a sequence of mock RequestOutput objects to simulate multiple - turns.""" + turns. + """ if cached_token_counts is None: cached_token_counts = [0] * num_turns @@ -270,7 +271,8 @@ def test_reasoning_tokens_counting(): def test_preamble_tokens_not_counted_as_reasoning(): """Preambles (commentary with no recipient) are visible user text, - not hidden reasoning. They must NOT inflate num_reasoning_tokens.""" + not hidden reasoning. They must NOT inflate num_reasoning_tokens. + """ context, parser = make_harmony_context() parser.enqueue_chunk_result(reasoning_token_count=0) @@ -287,7 +289,8 @@ def test_preamble_tokens_not_counted_as_reasoning(): def test_commentary_with_recipient_counted_as_reasoning(): """Commentary directed at a tool (recipient != None) is hidden from - the user, so it should still count as reasoning tokens.""" + the user, so it should still count as reasoning tokens. + """ context, parser = make_harmony_context() parser.enqueue_chunk_result(reasoning_token_count=3) @@ -346,7 +349,8 @@ async def test_single_turn_no_tool_output(): @pytest.mark.asyncio async def test_negative_tool_tokens_edge_case(): """Test edge case where calculation could result in negative tool - tokens. We should log an error and clamp the value to 0.""" + tokens. We should log an error and clamp the value to 0. + """ # Use patch to check if logger.error was called with patch("vllm.entrypoints.openai.responses.context.logger.error") as mock_log: context, _ = make_harmony_context(available_tools=["browser"]) @@ -544,7 +548,6 @@ async def test_streaming_multi_turn_token_counting(): @pytest.mark.asyncio async def test_streaming_message_synchronization(): """Completed messages from append-local and flush segments sync into context.""" - # Create a streaming context with some initial messages initial_messages = [ Message( @@ -687,7 +690,8 @@ def create_simple_context_output( finished=True, ): """Helper to create a RequestOutput with customizable text for - SimpleContext tests.""" + SimpleContext tests. + """ if token_ids is None: token_ids = [] return RequestOutput( @@ -871,7 +875,8 @@ def test_simple_context_final_output(): def test_simple_context_output_messages_empty_text_with_tokens(): """output_messages should be returned when tokens exist even if text is - empty (e.g. special tokens).""" + empty (e.g. special tokens). + """ context = SimpleContext() context.append_output( create_simple_context_output( @@ -889,7 +894,8 @@ def test_simple_context_output_messages_empty_text_with_tokens(): def test_simple_context_output_messages_no_mutation(): """Each call to output_messages returns a fresh list; callers can't - corrupt internal state.""" + corrupt internal state. + """ context = SimpleContext() context.append_output( create_simple_context_output( diff --git a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py index e5c3616d6e..9b36ab398d 100644 --- a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py +++ b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py @@ -121,7 +121,8 @@ def test_get_world_size_tp1(vllm_runner): @create_new_process_for_each_test() def test_init_weight_transfer_engine_calls_engine(vllm_runner): """Test that init_weight_transfer_engine calls the engine's - init_transfer_engine method.""" + init_transfer_engine method. + """ if torch.accelerator.device_count() < 1: pytest.skip("Need at least 1 GPU for this test") diff --git a/tests/evals/gpt_oss/conftest.py b/tests/evals/gpt_oss/conftest.py index d35dec4831..28fbb88d52 100644 --- a/tests/evals/gpt_oss/conftest.py +++ b/tests/evals/gpt_oss/conftest.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Pytest configuration for GPT-OSS evaluation tests. -""" +"""Pytest configuration for GPT-OSS evaluation tests.""" from pathlib import Path diff --git a/tests/evals/gpt_oss/test_gpqa_correctness.py b/tests/evals/gpt_oss/test_gpqa_correctness.py index 63188ec407..99d799698a 100644 --- a/tests/evals/gpt_oss/test_gpqa_correctness.py +++ b/tests/evals/gpt_oss/test_gpqa_correctness.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -GPQA evaluation using vLLM server and GPT-OSS evaluation package. +"""GPQA evaluation using vLLM server and GPT-OSS evaluation package. Usage: pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py \ @@ -48,7 +47,6 @@ def ensure_tiktoken_files(): def run_gpqa_eval(model_name: str, base_url: str, reasoning_effort: str) -> float: """Run GPQA evaluation using the gpt-oss evaluation package.""" - # Build the command to run the evaluation cmd = [ sys.executable, diff --git a/tests/evals/gsm8k/gsm8k_eval.py b/tests/evals/gsm8k/gsm8k_eval.py index 83e54afd9e..7e2534b55a 100644 --- a/tests/evals/gsm8k/gsm8k_eval.py +++ b/tests/evals/gsm8k/gsm8k_eval.py @@ -1,9 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Isolated GSM8K evaluation script for vLLM serve endpoint. -""" +"""Isolated GSM8K evaluation script for vLLM serve endpoint.""" import argparse import ast @@ -45,7 +43,7 @@ def download_and_cache_file(url: str, filename: str | None = None) -> str: def load_gsm8k_data() -> tuple[list[dict], list[dict]]: - """Load GSM8K train and test data""" + """Load GSM8K train and test data.""" train_url = f"{VLLM_S3_BUCKET_URL}/ci-datasets/gsm8k/train.jsonl" test_url = f"{VLLM_S3_BUCKET_URL}/ci-datasets/gsm8k/test.jsonl" @@ -91,6 +89,7 @@ async def call_vllm_api( Returns: Tuple of (response_text, completion_tokens) + """ data = { "prompt": prompt, @@ -221,8 +220,7 @@ def evaluate_gsm8k( gen_prefix: str = "", max_concurrency: int | None = None, ) -> dict[str, float | int]: - """ - Evaluate GSM8K accuracy using vLLM serve endpoint. + """Evaluate GSM8K accuracy using vLLM serve endpoint. Returns dict with accuracy, invalid_rate, latency, etc. """ diff --git a/tests/evals/gsm8k/test_gsm8k_correctness.py b/tests/evals/gsm8k/test_gsm8k_correctness.py index 5a237cf947..477a69a16c 100644 --- a/tests/evals/gsm8k/test_gsm8k_correctness.py +++ b/tests/evals/gsm8k/test_gsm8k_correctness.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -GSM8K evaluation using vLLM server and isolated GSM8K script. +"""GSM8K evaluation using vLLM server and isolated GSM8K script. Replacement for lm-eval-harness with better performance and control. Usage: diff --git a/tests/evals/gsm8k/test_gsm8k_offloading.py b/tests/evals/gsm8k/test_gsm8k_offloading.py index c140a41e26..bcc001e2bd 100644 --- a/tests/evals/gsm8k/test_gsm8k_offloading.py +++ b/tests/evals/gsm8k/test_gsm8k_offloading.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -GSM8K correctness test for CPU KV offloading connectors. +"""GSM8K correctness test for CPU KV offloading connectors. Regression guard for stride computation bugs in the offloading worker (e.g. https://github.com/vllm-project/vllm/pull/46888) and silent KV diff --git a/tests/evals/mrcr/test_mrcr_correctness.py b/tests/evals/mrcr/test_mrcr_correctness.py index 3adfd4dc87..f381a592dd 100644 --- a/tests/evals/mrcr/test_mrcr_correctness.py +++ b/tests/evals/mrcr/test_mrcr_correctness.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -MRCR long-context accuracy test. +"""MRCR long-context accuracy test. Usage: pytest -s -v tests/evals/mrcr/test_mrcr_correctness.py \ diff --git a/tests/fusion/test_quant_activation_contract.py b/tests/fusion/test_quant_activation_contract.py index ab12393692..ec9bb6110d 100644 --- a/tests/fusion/test_quant_activation_contract.py +++ b/tests/fusion/test_quant_activation_contract.py @@ -78,7 +78,8 @@ def _all_kernel_classes() -> list[type]: def _probe(cls: type): """A bare kernel instance with a plausible config, so input_quant_key() - can be queried without the hardware-gated constructor.""" + can be queried without the hardware-gated constructor. + """ obj = cls.__new__(cls) # type: ignore[call-overload] if issubclass(cls, NvFp4LinearKernel): obj.config = NvFp4LinearLayerConfig() diff --git a/tests/ir/ir_test_utils.py b/tests/ir/ir_test_utils.py index a82206b6f7..5cd05b9bcf 100644 --- a/tests/ir/ir_test_utils.py +++ b/tests/ir/ir_test_utils.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Shared test utilities for vLLM IR op correctness tests. -""" +"""Shared test utilities for vLLM IR op correctness tests.""" import torch diff --git a/tests/ir/test_op.py b/tests/ir/test_op.py index ca8185b2d3..5033b0ec10 100644 --- a/tests/ir/test_op.py +++ b/tests/ir/test_op.py @@ -34,7 +34,7 @@ def custom_add_op(fake_vllm_ir): @_custom_add.register_impl("impl_b", inplace=True) def impl_b(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - """Computes x+y+20""" + """Computes x+y+20.""" x.add_(y) x.add_(20) return x @@ -101,7 +101,7 @@ def test_no_kw_only_args(fake_vllm_ir): class TestIrOpCustomAdd: # Registration invariants def test_decorated_object(self, custom_add_op): - """Make sure that referring directly to an op is correct""" + """Make sure that referring directly to an op is correct.""" _custom_add = custom_add_op assert isinstance(_custom_add, IrOp) assert "_custom_add" in IrOp.registry diff --git a/tests/kernels/attention/test_amx_mla.py b/tests/kernels/attention/test_amx_mla.py index f5fa2fad16..7dab61f1c4 100644 --- a/tests/kernels/attention/test_amx_mla.py +++ b/tests/kernels/attention/test_amx_mla.py @@ -247,7 +247,8 @@ def test_extend_attention_cpu_matches_reference(prefix_lens, extend_lens): class _FakeLinear: """Minimal duck-typed stand-in for kv_b_proj: only .weight and - .quant_method are read by get_and_maybe_dequant_weights().""" + .quant_method are read by get_and_maybe_dequant_weights(). + """ def __init__(self, weight: torch.Tensor): self.weight = weight @@ -318,7 +319,8 @@ class _FakeAttnMetadata: def test_amx_mla_impl_forward_mqa_matches_reference(default_vllm_config): """forward_mqa receives already-absorbed Q (as MLAAttention.forward_impl would produce via layer.W_UK_T) and must reproduce plain per-head MLA - decode attention through W_UK/W_UV.""" + decode attention through W_UK/W_UV. + """ set_random_seed(1) num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank = 4, 32, 32, 64 impl, kv_b_weight = _make_amx_mla_impl( @@ -384,7 +386,8 @@ def test_amx_mla_impl_forward_mqa_matches_reference(default_vllm_config): def test_amx_mla_impl_forward_mha_matches_reference(default_vllm_config): """forward_mha receives raw unabsorbed Q and must attend correctly for a prefill batch with a mix of fresh and continued (cached-prefix) - sequences.""" + sequences. + """ set_random_seed(2) num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank = 4, 32, 32, 64 impl, kv_b_weight = _make_amx_mla_impl( diff --git a/tests/kernels/attention/test_attention_selector.py b/tests/kernels/attention/test_attention_selector.py index 53867a5edf..0511fe7b07 100644 --- a/tests/kernels/attention/test_attention_selector.py +++ b/tests/kernels/attention/test_attention_selector.py @@ -616,7 +616,8 @@ def test_non_causal_autoselect_backend(): ) def test_flash_attn_rejects_unhandled_kv_cache_dtypes(kv_cache_dtype: str): """FlashAttentionBackend must not claim support for kv_cache dtypes - that it cannot handle.""" + that it cannot handle. + """ from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend assert not FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype) @@ -627,7 +628,8 @@ def test_flash_attn_accepts_handled_fp8_variants( kv_cache_dtype: str, monkeypatch: pytest.MonkeyPatch ): """FlashAttentionBackend must accept the two fp8 dtypes it can actually - handle: 'fp8' (alias for fp8_e4m3fn) and 'fp8_e4m3'.""" + handle: 'fp8' (alias for fp8_e4m3fn) and 'fp8_e4m3'. + """ import vllm.v1.attention.backends.fa_utils as fa_utils_mod from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend diff --git a/tests/kernels/attention/test_cutlass_mla_decode.py b/tests/kernels/attention/test_cutlass_mla_decode.py index 7a5011b5ca..d56d72adba 100644 --- a/tests/kernels/attention/test_cutlass_mla_decode.py +++ b/tests/kernels/attention/test_cutlass_mla_decode.py @@ -225,7 +225,8 @@ def test_cutlass_mla_decode_cross_layer_view(): """The kernel must read the cache's page-dim stride instead of assuming pages are packed back-to-back. A per-layer view into a cross-layer (block-major) cache has stride(0) inflated by num_layers; outputs must - match a contiguous cache holding the same data exactly.""" + match a contiguous cache holding the same data exactly. + """ device = torch.device("cuda:0") torch.set_default_dtype(torch.bfloat16) torch.set_default_device(device) diff --git a/tests/kernels/attention/test_flashinfer.py b/tests/kernels/attention/test_flashinfer.py index 109f59bc85..b786fac095 100644 --- a/tests/kernels/attention/test_flashinfer.py +++ b/tests/kernels/attention/test_flashinfer.py @@ -96,6 +96,7 @@ def _make_paged_kv_metadata( kv_indices – CUDA int32, shape [total_blocks] kv_last_page_lens – CPU int32, shape [num_seqs] block_tables – CUDA int32, shape [num_seqs, max_blocks_per_seq] + """ num_seqs = len(kv_lens) max_blocks = (max(kv_lens) + block_size - 1) // block_size @@ -161,7 +162,8 @@ def test_fast_decode_plan_importable() -> None: @torch.inference_mode def test_fast_plan_decode_warmup_uses_full_plan(dtype: torch.dtype) -> None: """On the first call fast_plan_decode must route through self.plan() and - flip vllm_first_call to False on the wrapper object.""" + flip vllm_first_call to False on the wrapper object. + """ from unittest.mock import patch from vllm.v1.attention.backends.flashinfer import fast_plan_decode diff --git a/tests/kernels/attention/test_flashinfer_mla_decode.py b/tests/kernels/attention/test_flashinfer_mla_decode.py index a49ae3ae1a..0d37919643 100644 --- a/tests/kernels/attention/test_flashinfer_mla_decode.py +++ b/tests/kernels/attention/test_flashinfer_mla_decode.py @@ -286,7 +286,7 @@ def test_flashinfer_sm90_fp8_mla_decode_without_rope(): @requires_sm10x def test_flashinfer_mla_decode_workspace_supports_autotune(): - """vLLM's FlashInfer MLA decode workspace must be int8 for autotuning. + """VLLM's FlashInfer MLA decode workspace must be int8 for autotuning. Model Runner V2's warmup autotunes ``trtllm_batch_decode_mla``, which makes the FlashInfer autotuner enumerate the CuteDSL tactic. That tactic asserts diff --git a/tests/kernels/attention/test_flashinfer_trtllm_attention.py b/tests/kernels/attention/test_flashinfer_trtllm_attention.py index 48eb58b5e1..b35005d361 100644 --- a/tests/kernels/attention/test_flashinfer_trtllm_attention.py +++ b/tests/kernels/attention/test_flashinfer_trtllm_attention.py @@ -134,7 +134,8 @@ def make_quantized_kv_cache( head_size: int, ) -> tuple: """Quantize kv_cache based on dtype. Returns (kv_cache, kv_cache_sf, - kv_scale, ref_kv_cache, is_nvfp4_kv).""" + kv_scale, ref_kv_cache, is_nvfp4_kv). + """ is_nvfp4_kv = kv_quant_dtype == FP4_DTYPE if is_nvfp4_kv: data, scales, kv_scale, ref = make_nvfp4_kv_cache( diff --git a/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py b/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py index aa7c836666..e86686ee86 100644 --- a/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py +++ b/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py @@ -73,7 +73,8 @@ def _strided_context_inputs( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """The real prefill-context layouts: ``k_nope``/``v`` as the two strided halves of one kv_b_proj output and ``k_pe`` as a column slice of the gather - workspace (left in the fp8 cache layout for a plain fp8 cache).""" + workspace (left in the fp8 cache layout for a plain fp8 cache). + """ kv_nope = _randn(num_tokens, num_heads, 256, dtype=dtype) k_nope, v = kv_nope.split((128, 128), dim=-1) workspace = _randn(num_tokens, 576, dtype=dtype) diff --git a/tests/kernels/attention/test_lightning_attn.py b/tests/kernels/attention/test_lightning_attn.py index 61e1316680..d4b1a359d2 100644 --- a/tests/kernels/attention/test_lightning_attn.py +++ b/tests/kernels/attention/test_lightning_attn.py @@ -23,7 +23,7 @@ DTYPES = [torch.float32] def reference_lightning_attention(q, k, v, ed, block_size, kv_history): - """Reference implementation of lightning attention core algorithm + """Reference implementation of lightning attention core algorithm. The difference from the main implementation is that this processes each step sequentially, instead of using parallelized triton kernels @@ -72,7 +72,7 @@ def reference_lightning_attention(q, k, v, ed, block_size, kv_history): def reference_linear_decode(q, k, v, kv_caches, slope_rate, slot_idx): - """Reference implementation: linear attention decode function""" + """Reference implementation: linear attention decode function.""" B, H, _, D = q.shape output = torch.zeros(B, H * D, dtype=q.dtype, device=q.device) diff --git a/tests/kernels/attention/test_merge_attn_states.py b/tests/kernels/attention/test_merge_attn_states.py index 1b89e4db66..07295ed0b8 100644 --- a/tests/kernels/attention/test_merge_attn_states.py +++ b/tests/kernels/attention/test_merge_attn_states.py @@ -89,7 +89,8 @@ all_case_info: list[tuple] = [] @pytest.mark.parametrize("output_dtype", [torch.float32, torch.half, torch.bfloat16]) def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None: """When a token is empty on both sides (both LSE -inf), the 0/0 softmax - scales must not surface as NaN in the merged output.""" + scales must not surface as NaN in the merged output. + """ num_tokens, num_heads, head_size = 6, 8, 128 prefix_output = torch.zeros( num_tokens, num_heads, head_size, device=DEVICE, dtype=output_dtype diff --git a/tests/kernels/attention/test_mha_attn.py b/tests/kernels/attention/test_mha_attn.py index b6578fa2f9..8fa4963a6a 100644 --- a/tests/kernels/attention/test_mha_attn.py +++ b/tests/kernels/attention/test_mha_attn.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test: +"""Test: * Tests for MMEncoderAttention layer """ @@ -49,9 +48,7 @@ if current_platform.is_rocm(): @pytest.mark.parametrize("device", devices) def test_mha_attn_platform(default_vllm_config, device: str): - """ - Test the attention selector between different platform and device. - """ + """Test the attention selector between different platform and device.""" torch.set_default_dtype(torch.float16) if device == "cpu": @@ -113,8 +110,7 @@ def ref_attention( value: torch.Tensor, scale: float, ) -> torch.Tensor: - """ - Native implementation of scaled dot product attention without mask: + """Native implementation of scaled dot product attention without mask: - query, key, value: [batch_size, seq_len, num_heads, head_size] - attn_mask: [batch_size, seq_len, seq_len] """ diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index a7e6d499c1..c7b763e7c7 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -45,7 +45,8 @@ def kv_layout(request) -> KVCacheLayout: def _layer_stride_order(layout: KVCacheLayout, ndim: int) -> tuple[int, ...]: """Per-layer physical stride order for the given layout; the 3-dim - indexer side cache (H=1) is contiguous, so identity.""" + indexer side cache (H=1) is contiguous, so identity. + """ if ndim == 3: return (0, 1, 2) stride_order = layout.layer_view_order @@ -65,7 +66,8 @@ def _main_spec() -> FullAttentionSpec: def _main_kv_logical_shape(num_pages: int) -> tuple[int, ...]: """Standardized per-layer logical shape [B, H, N, C] for the main cache, - derived the same way the production allocator does.""" + derived the same way the production allocator does. + """ shape_bytes = compute_layer_kv_cache_shape_bytes(_main_spec(), num_pages) return (*shape_bytes[:-1], shape_bytes[-1] // DTYPE.itemsize) @@ -75,7 +77,8 @@ def _allocate_main_kv_via_contract( ) -> torch.Tensor: """Build the main KV cache exactly as the production allocator does for the given layout: allocate the physical (permuted) tensor, then expose the - inverse-permuted logical [B, H, N, C] view the backend sees.""" + inverse-permuted logical [B, H, N, C] view the backend sees. + """ logical_shape = _main_kv_logical_shape(num_pages) stride_order = _layer_stride_order(layout, len(logical_shape)) physical_shape = tuple(logical_shape[i] for i in stride_order) @@ -1376,7 +1379,8 @@ def test_amd_decode_index_topk_end_to_end( @pytest.mark.parametrize("num_idx_heads", [1, 4]) def test_decode_index_topk_fp8(num_idx_heads: int): """The standalone Triton path must score FP8 inputs in FP32 so its top-k - matches a reference computed from the dequantized FP8 values.""" + matches a reference computed from the dequantized FP8 values. + """ torch.manual_seed(0) topk, init_blocks, local_blocks, head_dim = 8, 0, 1, 128 decode_query_len = 1 @@ -1649,7 +1653,8 @@ def test_prefill_sparse_attention_correctness( def test_main_cache_layout_contract(): """The standardized per-layer logical shape is [B, H, N, C] with packed K/V content, and the legacy layout aliases resolve to the expected - per-layer stride orders.""" + per-layer stride orders. + """ nb, bs, h, d = 7, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM logical = _main_kv_logical_shape(nb) assert logical == (nb, h, bs, 2 * d) @@ -1773,7 +1778,8 @@ def test_indexer_cache_squeezes_to_contiguous_3d(): """The indexer side cache is standardized 4D with H=1: under both layouts the allocator's logical view stays contiguous and squeezes (as `MiniMaxM3IndexerCache.bind_kv_cache` does) to the 3-dim - [num_blocks, block_size, head_dim] cache the kernels consume.""" + [num_blocks, block_size, head_dim] cache the kernels consume. + """ nb = 5 ispec = MLAAttentionSpec( block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE @@ -1793,7 +1799,8 @@ def test_indexer_cache_squeezes_to_contiguous_3d(): def test_hnd_allocation_is_packed_head_major(): """Under HND the backend-visible logical view is the packed head-major - physical allocation.""" + physical allocation. + """ nb, bs, h, d = 4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM logical = _main_kv_logical_shape(nb) stride_order = KVCacheLayout.LBHNC.layer_view_order @@ -1821,7 +1828,8 @@ def test_hnd_allocation_is_packed_head_major(): def test_main_cache_is_block_first_and_unpadded(): """The allocator's contiguous-view branch (not the padded-strided branch) is used for the main GQA cache: its spec is unpadded and the physical - layout keeps num_blocks as the first dimension under both layouts.""" + layout keeps num_blocks as the first dimension under both layouts. + """ spec = _main_spec() # Unpadded -> allocator uses kv_tensor.view(...) rather than as_strided(). assert spec.page_size_padded is None @@ -1843,7 +1851,8 @@ def _build_decode_inputs( ): """Shared decode setup: uniform query tokens per request, a non-identity block table, and topk indices selecting the current block plus older causal - blocks for each query token.""" + blocks for each query token. + """ active_batch = len(seq_lens_list) batch = active_batch + num_padded_reqs pages_per_req = [(s + BLOCK_SIZE - 1) // BLOCK_SIZE for s in seq_lens_list] @@ -2179,7 +2188,8 @@ def test_decode_sparse_attention_correctness( ): """Decode (split-K) parity under both layouts: this is the only coverage of the decode-site cache feed, and the strided HND case fails if the kernel - ignores the cache strides.""" + ignores the cache strides. + """ torch.manual_seed(0) q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs( seq_lens_list, decode_query_len, num_padded_reqs @@ -2229,7 +2239,8 @@ def test_decode_wrong_layout_breaks_parity(): already contiguous-NHD (i.e. skipping the allocator's inverse permute) reorders the K/V content, so the decode output no longer matches the reference computed on the correct logical view. The mislabeled tensor keeps - the same shape as the correct view, so the kernel stays in bounds.""" + the same shape as the correct view, so the kernel stays in bounds. + """ torch.manual_seed(0) seq_lens_list = (130, 257) q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs(seq_lens_list) @@ -2264,7 +2275,8 @@ def test_decode_wrong_layout_breaks_parity(): def test_main_cache_byte_identical_through_production_allocator(): """AC-2: drive the real allocator (`create_kv_cache_views`) for the M3 main `FullAttentionSpec` under HND and assert the backend-visible view has the - same shape, stride, and storage offset as the packed-HND allocation.""" + same shape, stride, and storage offset as the packed-HND allocation. + """ nb = 4 spec = _main_spec() raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8) @@ -2279,7 +2291,8 @@ def test_main_cache_byte_identical_through_production_allocator(): def test_padded_main_cache_is_flagged(): """AC-2.1 negative: the M3 main cache relies on the allocator's contiguous-view branch (`page_size_padded is None`). A spec that sets - `page_size_padded` is explicitly flagged rather than silently wrong-strided.""" + `page_size_padded` is explicitly flagged rather than silently wrong-strided. + """ def _require_unpadded_block_first(spec, stride_order): inv_order = [stride_order.index(i) for i in range(len(stride_order))] @@ -2316,7 +2329,8 @@ def test_padded_main_cache_is_flagged(): def test_reshape_and_cache_flash_write_persists(kv_layout: KVCacheLayout): """AC-5 write path: the `reshape_and_cache_flash` write site now consumes packed-content K/V split views. Writing through those views must persist - into the bound storage under both layouts.""" + into the bound storage under both layouts. + """ torch.manual_seed(0) num_pages = 4 kv_cache = _allocate_main_kv_via_contract(num_pages, kv_layout) diff --git a/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py b/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py index ffdb76e176..8106ba37a6 100644 --- a/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py +++ b/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py @@ -23,7 +23,8 @@ pytestmark = pytest.mark.skipif( def test_concat_and_cache_mla_into_unified_slot_view(): """concat_and_cache_mla must write correctly into a per-layer view whose block stride is the full unified slot (block-major), with zero bleed into - the other layers' segments of the same slot.""" + the other layers' segments of the same slot. + """ from vllm import _custom_ops as ops torch.manual_seed(0) @@ -85,7 +86,8 @@ def test_concat_and_cache_mla_into_unified_slot_view(): def test_flashmla_dense_decode_unified_slot_view(): """FlashMLA dense decode (FLASHMLA backend, e.g. Kimi-K2-style dense MLA on Hopper) must read a unified-slot block-major view bit-identically to a - contiguous per-layer cache.""" + contiguous per-layer cache. + """ import vllm.v1.attention.ops.flashmla as fm ok, reason = fm.is_flashmla_dense_supported() @@ -150,7 +152,8 @@ def test_flashmla_dense_decode_unified_slot_view(): def test_flashinfer_mla_dense_decode_unified_slot_view(): """FlashInfer MLA dense decode must read a unified-slot block-major view (inflated stride(0), non-zero storage offset) bit-identically to a - contiguous per-layer cache.""" + contiguous per-layer cache. + """ try: from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla except ImportError: @@ -221,7 +224,8 @@ def test_flashinfer_mla_dense_decode_unified_slot_view(): def test_flashmla_fp8_sparse_decode_unified_slot_view(next_n: int): """FlashMLA fp8 sparse decode (DeepSeek V3.2/V4 DSA path) must read a unified-slot block-major view bit-identically to a contiguous fp8_ds_mla - cache, with finite nonzero output.""" + cache, with finite nonzero output. + """ import vllm.v1.attention.ops.flashmla as fm ok, reason = fm.is_flashmla_sparse_supported() @@ -310,7 +314,8 @@ def test_flashmla_fp8_sparse_decode_unified_slot_view(next_n: int): def test_indexer_k_quant_and_cache_into_unified_slot_view(): """indexer_k_quant_and_cache (DeepSeek V3.2/V4 DSA indexer K write) must write correctly into a per-layer view whose block stride is the full - unified slot, with zero bleed into the other layers' segments.""" + unified slot, with zero bleed into the other layers' segments. + """ from vllm import _custom_ops as ops torch.manual_seed(0) @@ -355,7 +360,8 @@ def test_indexer_k_quant_and_cache_into_unified_slot_view(): def test_flashattn_mla_dense_decode_unified_slot_view(): """FA3 decode (FLASH_ATTN_MLA backend) must read a unified-slot - block-major view bit-identically to a contiguous per-layer cache.""" + block-major view bit-identically to a contiguous per-layer cache. + """ try: from vllm.vllm_flash_attn import flash_attn_varlen_func except ImportError: @@ -428,7 +434,8 @@ def test_flashattn_mla_dense_decode_unified_slot_view(): def test_flashmla_dense_fp8_decode_unified_slot_view(): """FlashMLA dense fp8 decode (FLASHMLA backend with quantized KV cache) must read a unified-slot block-major view bit-identically to a contiguous - per-layer fp8 cache.""" + per-layer fp8 cache. + """ import vllm.v1.attention.ops.flashmla as fm ok, reason = fm.is_flashmla_dense_supported() @@ -499,7 +506,8 @@ def test_flashmla_dense_fp8_decode_unified_slot_view(): def test_flashinfer_mla_dense_fp8_decode_unified_slot_view(): """FlashInfer MLA dense decode with an fp8 KV cache must read a unified-slot block-major view bit-identically to a contiguous per-layer - cache.""" + cache. + """ try: from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla except ImportError: diff --git a/tests/kernels/attention/test_rocm_aiter_fa.py b/tests/kernels/attention/test_rocm_aiter_fa.py index dd2870ccf1..0fbe7bced2 100644 --- a/tests/kernels/attention/test_rocm_aiter_fa.py +++ b/tests/kernels/attention/test_rocm_aiter_fa.py @@ -283,7 +283,8 @@ def _run_direct_flash_attn_case( # Backend contract tests -------------------------------------------------- def test_aiter_mha_backend_contract(): """The ROCm backend advertises the dtypes, shapes, and attention types it - is designed to handle.""" + is designed to handle. + """ from vllm.v1.attention.backend import AttentionType from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend @@ -311,7 +312,8 @@ def test_aiter_mha_backend_contract(): def test_aiter_mha_backend_validates_kv_cache_block_size(): """The backend should reject KV cache specs that cannot be gathered - correctly by the ROCm kernel.""" + correctly by the ROCm kernel. + """ from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend from vllm.v1.kv_cache_interface import FullAttentionSpec @@ -331,7 +333,8 @@ def test_aiter_mha_backend_validates_kv_cache_block_size(): def test_aiter_mha_backend_supports_compute_capability_matches_mi3xx_probe(): """The backend should trust the ROCm MI3xx probe instead of the raw torch - capability tuple.""" + capability tuple. + """ from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend @@ -356,7 +359,8 @@ def test_aiter_mha_env_flags_control_enablement( use_aiter, use_mha, expected, monkeypatch ): """Both the global AITER flag and the MHA-specific flag must be enabled - before the ROCm MHA path is considered active.""" + before the ROCm MHA path is considered active. + """ from vllm._aiter_ops import rocm_aiter_ops _assert_aiter_supported() @@ -374,7 +378,8 @@ def test_aiter_mha_env_flags_control_enablement( def test_aiter_mha_platform_gate_matches_install_and_arch(): """The global AITER availability check should only open on ROCm MI3xx when - the aiter package is installed.""" + the aiter package is installed. + """ from vllm._aiter_ops import IS_AITER_FOUND, is_aiter_found_and_supported assert is_aiter_found_and_supported() is ( @@ -645,7 +650,8 @@ def test_aiter_mha_multi_batch(num_heads, head_size, dtype): @pytest.mark.parametrize("num_blocks", DIRECT_NUM_BLOCKS) def test_aiter_fa_large_block_table_matches_reference(num_blocks): """The direct paged-KV path should stay stable for both normal and very - large block tables.""" + large block tables. + """ atol = 2e-2 rtol = 2e-2 _assert_aiter_supported() @@ -667,7 +673,8 @@ def test_aiter_fa_large_block_table_matches_reference(num_blocks): @pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only") def test_aiter_fa_sliding_window_matches_reference(): """The direct kernel should respect the same sliding-window causal mask as - the naive reference implementation.""" + the naive reference implementation. + """ atol = 2e-2 rtol = 2e-2 _assert_aiter_supported() diff --git a/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py b/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py index 659211521a..0944be8937 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py @@ -35,7 +35,7 @@ def _require_aiter(): @torch.inference_mode() def test_mla_decode_fwd_op_schema() -> None: - """opcheck validates registration, schema, fake-tensor, and ``o`` aliasing. + """Opcheck validates registration, schema, fake-tensor, and ``o`` aliasing. A single opcheck call covers that the op is registered/callable, that its fake implementation matches the real op (torch.compile tracing), and that diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py index aabce97fac..17b59c0743 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py @@ -111,7 +111,8 @@ def _make_mixed_common_metadata(): def _patch_build_deps(monkeypatch, events=None): """Stub the aiter kernel, triton helper and CUDA sync so ``build()`` runs - on CPU.""" + on CPU. + """ def fake_generate_sparse_seqlen_triton( query_lens, seq_lens, cu_query_lens, topk_token, num_tokens, max_query_len @@ -152,7 +153,8 @@ def test_build_populates_decode_only_split_fields(monkeypatch): def test_build_populates_mixed_split_fields(monkeypatch): """Mixed decode+prefill batch: split is reported, prefill fields stay - default because this impl always runs the MQA path.""" + default because this impl always runs the MQA path. + """ builder = _make_builder() _patch_build_deps(monkeypatch) diff --git a/tests/kernels/attention/test_triton_decode_attention.py b/tests/kernels/attention/test_triton_decode_attention.py index 5ac1d6cdcb..a46916e14e 100644 --- a/tests/kernels/attention/test_triton_decode_attention.py +++ b/tests/kernels/attention/test_triton_decode_attention.py @@ -125,7 +125,8 @@ def _quantize_to_fp8(tensor: torch.Tensor): @pytest.mark.parametrize("PAGE_SIZE", [1, 16]) def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE): """Test FP8 KV cache path: quantize K/V to FP8, run kernel with scales, - and compare against BF16 reference output.""" + and compare against BF16 reference output. + """ assert CACHE_SIZE % PAGE_SIZE == 0 dtype = torch.bfloat16 seq_len = L @@ -248,7 +249,8 @@ def test_decode_attention_cross_layer_view(H_Q, H_KV, D_QK, D_V, is_mla, PAGE_SI """The kernel must honor the cache's page-dim stride, not assume pages are packed back-to-back. A per-layer view into a cross-layer (block-major) cache has stride(0) inflated by num_layers; outputs must match a - contiguous cache holding the same data exactly.""" + contiguous cache holding the same data exactly. + """ B = 3 seq_len = 1027 CACHE_SIZE = 16384 diff --git a/tests/kernels/attention/test_triton_unified_attention_diffkv.py b/tests/kernels/attention/test_triton_unified_attention_diffkv.py index 513811a09c..04b7755853 100644 --- a/tests/kernels/attention/test_triton_unified_attention_diffkv.py +++ b/tests/kernels/attention/test_triton_unified_attention_diffkv.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for the Triton DiffKV unified-attention kernel. -""" +"""Unit tests for the Triton DiffKV unified-attention kernel.""" import pytest import torch diff --git a/tests/kernels/attention/test_trtllm_kvfp8_dequant.py b/tests/kernels/attention/test_trtllm_kvfp8_dequant.py index bd267c43be..8368869b65 100644 --- a/tests/kernels/attention/test_trtllm_kvfp8_dequant.py +++ b/tests/kernels/attention/test_trtllm_kvfp8_dequant.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Standalone unit tests for trtllm_prefill_attn_kvfp8_dequant. +"""Standalone unit tests for trtllm_prefill_attn_kvfp8_dequant. Tests both contiguous and non-contiguous (cross-layer unified) KV cache layouts against a pure-PyTorch reference implementation. @@ -43,8 +42,7 @@ def make_random_kv_cache( layout=KVCacheLayout.LBHNC, num_layers=None, ): - """ - Create a random fp8 KV cache mimicking the production allocation. + """Create a random fp8 KV cache mimicking the production allocation. Physical layout: [L, B, H, N, 2 * hs] permuted by ``layout.stride_order`` Returned view: (num_blocks, 2, num_kv_heads, block_size, head_size) @@ -382,8 +380,7 @@ def test_large_block_size(): @torch.inference_mode() def test_cross_layer_many_layers(): - """ - Non-contiguous with 36 layers -- matches real gpt-oss-120b. + """Non-contiguous with 36 layers -- matches real gpt-oss-120b. Cross-layer (BLHNC) strides are far from contiguous (factor of 36 in the gaps). """ from vllm.v1.attention.backends.flashinfer import ( diff --git a/tests/kernels/attention/test_xpu_mla_sparse.py b/tests/kernels/attention/test_xpu_mla_sparse.py index 965c6f735f..47139fc465 100644 --- a/tests/kernels/attention/test_xpu_mla_sparse.py +++ b/tests/kernels/attention/test_xpu_mla_sparse.py @@ -30,12 +30,12 @@ def reference_mla_sparse_prefill( topk_length: torch.Tensor | None = None, attn_sink: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Returns: + """Returns: - o: [s_q, h_q, dv] - o_fp32: [s_q, h_q, dv] - max_logits: [s_q, h_q] - lse: [s_q, h_q] + """ s_q, h_q, d_qk = q.shape s_kv, _, _ = kv.shape diff --git a/tests/kernels/core/test_apply_rotary_emb.py b/tests/kernels/core/test_apply_rotary_emb.py index 5c765bbb4a..fd92de16bb 100644 --- a/tests/kernels/core/test_apply_rotary_emb.py +++ b/tests/kernels/core/test_apply_rotary_emb.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for ApplyRotaryEmb CustomOp dispatch behavior. +"""Tests for ApplyRotaryEmb CustomOp dispatch behavior. This test ensures that RotaryEmbedding classes correctly call the appropriate ApplyRotaryEmb methods based on the calling context: @@ -179,8 +178,7 @@ def test_rotary_embedding_dispatch( test_case: RotaryEmbeddingTestCase, device: str, ): - """ - Test that RotaryEmbedding classes dispatch to the correct ApplyRotaryEmb method. + """Test that RotaryEmbedding classes dispatch to the correct ApplyRotaryEmb method. - forward_native methods should call ApplyRotaryEmb.forward_native() - forward_cuda/forward methods should call ApplyRotaryEmb.forward() diff --git a/tests/kernels/core/test_fused_embed_norm.py b/tests/kernels/core/test_fused_embed_norm.py index 9da43dba83..fc30ab3aca 100644 --- a/tests/kernels/core/test_fused_embed_norm.py +++ b/tests/kernels/core/test_fused_embed_norm.py @@ -41,7 +41,8 @@ def _rmsnorm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor: def test_fused_embed_norm_matches_reference(): """Main-model fusion: the residual is the exact gather and the second output is a correct RMSNorm. The norm matches a full-precision reference to ~2 bf16 - ulp (rtol 1e-2) -- that gap is bf16 rounding, not the kernel.""" + ulp (rtol 1e-2) -- that gap is bf16 rounding, not the kernel. + """ set_random_seed(13) table = torch.randn(VOCAB, HIDDEN, dtype=DTYPE, device="cuda") ids = torch.randint(0, VOCAB, (NUM_TOKENS,), dtype=torch.int32, device="cuda") @@ -60,7 +61,8 @@ def test_fused_embed_norm_matches_reference(): @torch.inference_mode() def test_fused_embed_eh_norm_matches_reference(): """MTP fusion (folded gather) is bit-exact vs gathering the embeds and - feeding the untouched model-local ``fused_eh_norm``.""" + feeding the untouched model-local ``fused_eh_norm``. + """ set_random_seed(13) table = torch.randn(VOCAB, HIDDEN, dtype=DTYPE, device="cuda") ids = torch.randint(0, VOCAB, (NUM_TOKENS,), dtype=torch.int32, device="cuda") diff --git a/tests/kernels/core/test_fused_q_kv_rmsnorm.py b/tests/kernels/core/test_fused_q_kv_rmsnorm.py index 3d0e2ac11e..81382bdda9 100644 --- a/tests/kernels/core/test_fused_q_kv_rmsnorm.py +++ b/tests/kernels/core/test_fused_q_kv_rmsnorm.py @@ -62,7 +62,8 @@ def test_fused_q_kv_rmsnorm_outputs_are_packed(num_tokens: int): a [1, q_size] slice of a [1, q_size + kv_size] buffer used to produce a qr_out with row stride q_size + kv_size. Downstream dispatchers that require packed row-major inputs then reject the tensor and silently fall - back to a slower GEMM path on every decode step.""" + back to a slower GEMM path on every decode step. + """ device = "cuda" dtype = torch.bfloat16 q_size, kv_size = 192, 576 @@ -88,7 +89,8 @@ def test_fused_q_kv_rmsnorm_outputs_are_packed(num_tokens: int): def test_fused_q_kv_rmsnorm_launches_past_grid_y_cap(num_tokens: int): """Regression guard: grid used to be (2, num_tokens), hitting CUDA's 65535 grid-y cap at num_tokens >= 65536. The new grid (num_tokens, 2) - lifts that bound to 2**31-1.""" + lifts that bound to 2**31-1. + """ device = "cuda" dtype = torch.bfloat16 q_size, kv_size = 192, 576 diff --git a/tests/kernels/core/test_fused_rms_norm_gated.py b/tests/kernels/core/test_fused_rms_norm_gated.py index 86b8e79656..74fdb00c30 100644 --- a/tests/kernels/core/test_fused_rms_norm_gated.py +++ b/tests/kernels/core/test_fused_rms_norm_gated.py @@ -2,7 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests that FusedRMSNormGated decomposes correctly under torch.compile, -matching the eager triton kernel output.""" +matching the eager triton kernel output. +""" import pytest import torch diff --git a/tests/kernels/core/test_minimax_reduce_rms.py b/tests/kernels/core/test_minimax_reduce_rms.py index cd8efc14a3..b92a08a317 100644 --- a/tests/kernels/core/test_minimax_reduce_rms.py +++ b/tests/kernels/core/test_minimax_reduce_rms.py @@ -36,7 +36,6 @@ def _worker_forward_qk( eps, ): """Per-rank worker: compare NCCL allreduce path vs Lamport fused kernel.""" - if not hasattr(torch.ops._C, "minimax_allreduce_rms_qk"): cleanup_dist_env_and_memory() return diff --git a/tests/kernels/core/test_mrope.py b/tests/kernels/core/test_mrope.py index 81e0a916a3..0e355d7cd3 100644 --- a/tests/kernels/core/test_mrope.py +++ b/tests/kernels/core/test_mrope.py @@ -229,7 +229,7 @@ def test_mrope_torch_compile_tracing( # Create a wrapper that makes the in-place function appear functional def functional_forward_cuda(pos, q, k): - """Wrapper that converts in-place operation to functional style + """Wrapper that converts in-place operation to functional style. CUDA Graph does not support in-place operations. This wrapper creates working copies of the diff --git a/tests/kernels/core/test_opcheck.py b/tests/kernels/core/test_opcheck.py index 40ced08b93..45dd01d44e 100644 --- a/tests/kernels/core/test_opcheck.py +++ b/tests/kernels/core/test_opcheck.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for miscellaneous utilities -""" +"""Tests for miscellaneous utilities.""" import torch diff --git a/tests/kernels/core/test_rocm_aiter_ops.py b/tests/kernels/core/test_rocm_aiter_ops.py index c15afd72d7..b2bed636e9 100644 --- a/tests/kernels/core/test_rocm_aiter_ops.py +++ b/tests/kernels/core/test_rocm_aiter_ops.py @@ -579,7 +579,8 @@ def test_rocm_aiter_mla_dual_rmsnorm_group_quant_vs_sequential(m): def test_rocm_aiter_mla_dual_rmsnorm_group_quant_strided_split_view(): """Split-view inputs, as produced by the DeepSeek-V4 q/kv split, match - contiguous inputs.""" + contiguous inputs. + """ require_aiter() require_fp8() from vllm._aiter_ops import rocm_aiter_ops diff --git a/tests/kernels/core/test_rotary_embedding.py b/tests/kernels/core/test_rotary_embedding.py index 1cbb5dbd18..353aefa11a 100644 --- a/tests/kernels/core/test_rotary_embedding.py +++ b/tests/kernels/core/test_rotary_embedding.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for miscellaneous utilities -""" +"""Tests for miscellaneous utilities.""" import pytest import torch diff --git a/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py b/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py index 1ac17335b6..3f192ba2e7 100644 --- a/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py +++ b/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for fused MLA KV-cache write and RoPE fused kernel -""" +"""Tests for fused MLA KV-cache write and RoPE fused kernel.""" import random diff --git a/tests/kernels/core/test_vocab_parallel_embedding.py b/tests/kernels/core/test_vocab_parallel_embedding.py index b0038c084e..fa50d2b3eb 100644 --- a/tests/kernels/core/test_vocab_parallel_embedding.py +++ b/tests/kernels/core/test_vocab_parallel_embedding.py @@ -87,7 +87,8 @@ def _shard(full_table, si, rows_per_rank, dtype): @pytest.mark.parametrize("id_dtype", [torch.int32, torch.int64]) def test_matches_masked_gather(tp_size, vocab_size, org_vocab_size, hidden, id_dtype): """Every rank's partial output is bit-exact with the eager path, including - the added-vocab (LoRA) rows and the ids owned by another rank.""" + the added-vocab (LoRA) rows and the ids owned by another rank. + """ set_random_seed(7) dtype = torch.bfloat16 rows = pad_vocab_size(vocab_size, 64) // tp_size diff --git a/tests/kernels/helion/test_autotune.py b/tests/kernels/helion/test_autotune.py index 8b42e145d4..f28e892109 100644 --- a/tests/kernels/helion/test_autotune.py +++ b/tests/kernels/helion/test_autotune.py @@ -59,7 +59,8 @@ class TestAutotuneDisabledKernel: def test_autotune_disabled_kernel_produces_valid_config(self): """Register a kernel with no configs (disabled), run autotune, - verify it produces a valid helion.Config.""" + verify it produces a valid helion.Config. + """ with dummy_kernel_registry(configs={}) as register: wrapper = register( "autotune_test_kernel", diff --git a/tests/kernels/helion/test_config_manager.py b/tests/kernels/helion/test_config_manager.py index f8e5eae6f1..63f8130849 100644 --- a/tests/kernels/helion/test_config_manager.py +++ b/tests/kernels/helion/test_config_manager.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for Helion ConfigManager and ConfigSet. +"""Unit tests for Helion ConfigManager and ConfigSet. Tests the simplified configuration management system for Helion custom kernels. """ diff --git a/tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py b/tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py index 6f5a40fad8..14a4356e14 100644 --- a/tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py +++ b/tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the dynamic_per_token_scaled_fp8_quant helion kernel +"""Tests for the dynamic_per_token_scaled_fp8_quant helion kernel. Run `pytest tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py`. """ diff --git a/tests/kernels/helion/test_fused_qk_norm_rope.py b/tests/kernels/helion/test_fused_qk_norm_rope.py index 4ced522d91..a973533169 100644 --- a/tests/kernels/helion/test_fused_qk_norm_rope.py +++ b/tests/kernels/helion/test_fused_qk_norm_rope.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the fused_qk_norm_rope helion kernel +"""Tests for the fused_qk_norm_rope helion kernel. Run `pytest tests/kernels/helion/test_fused_qk_norm_rope.py`. """ diff --git a/tests/kernels/helion/test_helion_available.py b/tests/kernels/helion/test_helion_available.py index 560a671f14..a4de68bd67 100644 --- a/tests/kernels/helion/test_helion_available.py +++ b/tests/kernels/helion/test_helion_available.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for Helion kernel availability and basic functionality. +"""Tests for Helion kernel availability and basic functionality. This module demonstrates the pattern for testing optional Helion kernels. Tests in this directory will be skipped if Helion is not installed. diff --git a/tests/kernels/helion/test_per_token_group_fp8_quant.py b/tests/kernels/helion/test_per_token_group_fp8_quant.py index d9d96f9287..d9dd8a8817 100644 --- a/tests/kernels/helion/test_per_token_group_fp8_quant.py +++ b/tests/kernels/helion/test_per_token_group_fp8_quant.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the per_token_group_fp8_quant helion kernel +"""Tests for the per_token_group_fp8_quant helion kernel. Run `pytest tests/kernels/helion/test_per_token_group_fp8_quant.py`. """ diff --git a/tests/kernels/helion/test_register.py b/tests/kernels/helion/test_register.py index 9876135056..572a4126cf 100644 --- a/tests/kernels/helion/test_register.py +++ b/tests/kernels/helion/test_register.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for Helion kernel registration. +"""Unit tests for Helion kernel registration. Tests ConfiguredHelionKernel, HelionKernelWrapper, and PresetConfigSearch including config picker registration and custom autotuner integration. @@ -599,7 +598,8 @@ class TestHelionKernelWrapper: ) def test_init_eagerly_initializes_hop_path(self): """Test that register_kernel eagerly builds the configured kernel - on the HOP path (no custom op registration needed).""" + on the HOP path (no custom op registration needed). + """ from vllm.kernels.helion.utils import get_canonical_gpu_name configs: dict[CaseKey, helion.Config] = { @@ -634,7 +634,8 @@ class TestHelionKernelWrapper: ) def test_init_eagerly_initializes(self): """Test that register_kernel eagerly loads configs and detects GPU - during construction so __call__ needs no further initialization.""" + during construction so __call__ needs no further initialization. + """ from vllm.kernels.helion.utils import get_canonical_gpu_name with ( @@ -962,7 +963,8 @@ class TestTorchCompileHOP: def test_compiled_graph_contains_helion_hop(self): """Verify torch.compile on a HelionKernelWrapper emits a - helion_kernel_wrapper_mutation HOP node in the FX graph.""" + helion_kernel_wrapper_mutation HOP node in the FX graph. + """ configs: dict[CaseKey, helion.Config] = { CaseKey.default(): helion.Config(block_sizes=[4, 4]) } @@ -1020,7 +1022,6 @@ class TestTorchCompileHOP: ) def test_inductor_backend_compiles_helion_hop(self): """Test torch.compile with inductor backend and Helion fusion enabled.""" - configs: dict[CaseKey, helion.Config] = { CaseKey.default(): helion.Config(block_sizes=[4, 4]) } diff --git a/tests/kernels/helion/test_rms_norm_dynamic_per_token_quant.py b/tests/kernels/helion/test_rms_norm_dynamic_per_token_quant.py index 50591ae34a..d0650dfc7a 100644 --- a/tests/kernels/helion/test_rms_norm_dynamic_per_token_quant.py +++ b/tests/kernels/helion/test_rms_norm_dynamic_per_token_quant.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the rms_norm_dynamic_per_token_quant helion kernel +"""Tests for the rms_norm_dynamic_per_token_quant helion kernel. Run `pytest tests/kernels/helion/test_rms_norm_dynamic_per_token_quant.py`. """ diff --git a/tests/kernels/helion/test_rms_norm_per_block_quant.py b/tests/kernels/helion/test_rms_norm_per_block_quant.py index 4ffc759c63..e9c003f4e3 100644 --- a/tests/kernels/helion/test_rms_norm_per_block_quant.py +++ b/tests/kernels/helion/test_rms_norm_per_block_quant.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the rms_norm_per_block_quant helion kernel +"""Tests for the rms_norm_per_block_quant helion kernel. Run `pytest tests/kernels/helion/test_rms_norm_per_block_quant.py`. """ diff --git a/tests/kernels/helion/utils.py b/tests/kernels/helion/utils.py index 38893fc8fe..4a282f275d 100644 --- a/tests/kernels/helion/utils.py +++ b/tests/kernels/helion/utils.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Helion Kernel test utils""" +"""Helion Kernel test utils.""" import pytest import torch diff --git a/tests/kernels/ir/test_aiter_norm_dispatch.py b/tests/kernels/ir/test_aiter_norm_dispatch.py index 19cf5d3112..f84a7615b1 100644 --- a/tests/kernels/ir/test_aiter_norm_dispatch.py +++ b/tests/kernels/ir/test_aiter_norm_dispatch.py @@ -21,7 +21,7 @@ pytestmark = pytest.mark.skipif( def _qkv_slice_by_head(num_tokens, num_q_heads, num_kv_heads, head_dim): - """q viewed per-head from a fused QKV projection, as Qwen3-style QK-norm does.""" + """Q viewed per-head from a fused QKV projection, as Qwen3-style QK-norm does.""" q_size, kv_size = num_q_heads * head_dim, num_kv_heads * head_dim qkv = torch.empty(num_tokens, q_size + 2 * kv_size) q = qkv.split([q_size, kv_size, kv_size], dim=-1)[0] diff --git a/tests/kernels/ir/test_ir_ops.py b/tests/kernels/ir/test_ir_ops.py index 730db9833b..560d047c0f 100644 --- a/tests/kernels/ir/test_ir_ops.py +++ b/tests/kernels/ir/test_ir_ops.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Meta-tests for vLLM IR op infrastructure. +"""Meta-tests for vLLM IR op infrastructure. Ensures all registered ops have input generators defined. Per-op correctness tests live alongside their op definitions diff --git a/tests/kernels/ir/test_layernorm.py b/tests/kernels/ir/test_layernorm.py index c324b189b9..e41b7c09d3 100644 --- a/tests/kernels/ir/test_layernorm.py +++ b/tests/kernels/ir/test_layernorm.py @@ -368,7 +368,8 @@ class TestFusedAddRMSNorm: @pytest.mark.parametrize("provider", ["vllm_c"]) def test_inplace_semantics(self, dtype, n_tokens, hidden_size, epsilon, provider): """Test that inplace implementations reuse inputs, - for maybe_inplace overload but not for default overload.""" + for maybe_inplace overload but not for default overload. + """ impl = ir.ops.fused_add_rms_norm.impls[provider] if not impl.supported: pytest.skip(f"{provider} impl not supported on this platform") diff --git a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py index c239f831c6..9495465b9e 100644 --- a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py +++ b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py @@ -1055,7 +1055,8 @@ CONV_EQUIV_SEQ_LENS = [[1], [7], [64], [65], [1, 2, 3], [63, 64, 65], [128, 129] def _conv_fp32_oracle(x, weight, bias, seq_lens, activation="silu"): """High-precision conv reference: everything in float32, no bf16 rounding. - x: [total_tokens, dim] (no initial state). Returns [total_tokens, dim].""" + x: [total_tokens, dim] (no initial state). Returns [total_tokens, dim]. + """ xf = x.float() wf = weight.float().unsqueeze(1) bf = bias.float() @@ -1164,7 +1165,8 @@ def test_conv_cpp_no_worse_than_torch_vs_fp32(seq_lens): def test_conv_cpp_vnni_packed_matches_torch(seq_lens): """The exact runtime prefill sequence (VNNI-packed weight + SD-layout conv_state view + is_vnni=True) must match the torch fallback. Validates - the packing + layout handoff on any AVX-512BF16 CPU.""" + the packing + layout handoff on any AVX-512BF16 CPU. + """ x, weight, bias = _conv_inputs(sum(seq_lens)) out_torch, _ = _run_prefill_torch(x, weight, bias, seq_lens) out_vnni, _ = _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=True) @@ -1178,7 +1180,8 @@ def test_conv_cpp_vnni_packed_matches_torch(seq_lens): @torch.inference_mode() def test_conv_update_cpp_matches_torch(batch): """Decode conv: causal_conv1d_update_cpu matches causal_conv1d_update_torch, - including the in-place conv_state update (the next-step handoff).""" + including the in-place conv_state update (the next-step handoff). + """ from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( causal_conv1d_update_torch, ) diff --git a/tests/kernels/mamba/test_causal_conv1d.py b/tests/kernels/mamba/test_causal_conv1d.py index d2a2981edf..d4f5fac59c 100644 --- a/tests/kernels/mamba/test_causal_conv1d.py +++ b/tests/kernels/mamba/test_causal_conv1d.py @@ -36,8 +36,7 @@ def causal_conv1d_ref( final_states_out: torch.Tensor | None = None, activation: str | None = "silu", ): - """ - x: (batch, dim, seqlen) + """x: (batch, dim, seqlen) weight: (dim, width) bias: (dim,) initial_states: (batch, dim, width - 1) @@ -72,8 +71,7 @@ def causal_conv1d_ref( def causal_conv1d_update_ref( x, conv_state, weight, bias=None, activation=None, cache_seqlens=None ): - """ - x: (batch, dim) or (batch, dim, seqlen) + """x: (batch, dim) or (batch, dim, seqlen) conv_state: (batch, dim, state_len), where state_len >= width - 1 weight: (dim, width) bias: (dim,) @@ -136,8 +134,7 @@ def causal_conv1d_opcheck_fn( activation: str | None = "silu", null_block_id: int = NULL_BLOCK_ID, ): - """ - x: (batch, dim, seqlen) + """x: (batch, dim, seqlen) weight: (dim, width) bias: (dim,) seq_idx: (batch, seqlen) diff --git a/tests/kernels/mamba/test_mamba_ssm.py b/tests/kernels/mamba/test_mamba_ssm.py index f8af3db56b..c24f4c0f87 100644 --- a/tests/kernels/mamba/test_mamba_ssm.py +++ b/tests/kernels/mamba/test_mamba_ssm.py @@ -51,8 +51,7 @@ def selective_scan_ref( prev_state=None, final_state_out=None, ): - """ - u: r(B D L) + """u: r(B D L) delta: r(B D L) A: c(D N) or r(D N) B: c(D N) or r(B N L) or r(B N 2L) or r(B G N L) or (B G N L) @@ -135,7 +134,7 @@ def selective_scan_opcheck_fn( cu_chunk_seqlen=None, last_chunk_indices=None, ): - """if return_last_state is True, returns (out, last_state) + """If return_last_state is True, returns (out, last_state) last_state has shape (batch, dim, dstate). """ if u.stride(-1) != 1: diff --git a/tests/kernels/mamba/test_mamba_ssm_configs.py b/tests/kernels/mamba/test_mamba_ssm_configs.py index d35ce95746..4330caa611 100644 --- a/tests/kernels/mamba/test_mamba_ssm_configs.py +++ b/tests/kernels/mamba/test_mamba_ssm_configs.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for the JSON-based config loader added to selective_state_update. +"""Unit tests for the JSON-based config loader added to selective_state_update. Tests cover: - Flat MoE-style filename generation @@ -123,7 +122,8 @@ def test_fallback_when_no_config(monkeypatch, tmp_path): def test_nearest_effective_batch_interpolation(monkeypatch, tmp_path): """When effective_batch = batch*nheads is not an exact key, the closest - key should be selected.""" + key should be selected. + """ _write_config( tmp_path, dstate=32, @@ -168,7 +168,8 @@ def test_nearest_effective_batch_interpolation(monkeypatch, tmp_path): def test_non_dict_json_returns_none(monkeypatch, tmp_path): """A valid JSON file that is not a dict (e.g. a list) must be ignored - and return None rather than raising AttributeError.""" + and return None rather than raising AttributeError. + """ device_name = get_ssm_device_name() config_path = tmp_path / get_ssm_config_file_name( _HEADDIM, 16, _CACHE_DTYPE, device_name @@ -190,7 +191,8 @@ def test_non_dict_json_returns_none(monkeypatch, tmp_path): def test_empty_config_falls_back_to_heuristic(monkeypatch, tmp_path): """An empty JSON object {} must not crash min() — should fall back - to the hard-coded heuristic.""" + to the hard-coded heuristic. + """ _write_config(tmp_path, dstate=64, payload={}) monkeypatch.setenv("VLLM_TUNED_CONFIG_FOLDER", str(tmp_path)) diff --git a/tests/kernels/mamba/test_mamba_ssm_ssd.py b/tests/kernels/mamba/test_mamba_ssm_ssd.py index 1de25780ea..22abeab1ed 100644 --- a/tests/kernels/mamba/test_mamba_ssm_ssd.py +++ b/tests/kernels/mamba/test_mamba_ssm_ssd.py @@ -41,14 +41,15 @@ def segsum(x): def ssd_minimal_discrete(X, A, B, C, block_len, initial_states=None): - """ - Arguments: + """Arguments: X: (batch, length, n_heads, d_head) A: (batch, length, n_heads) B: (batch, length, n_heads, d_state) C: (batch, length, n_heads, d_state) + Return: Y: (batch, length, n_heads, d_head) + """ assert X.dtype == A.dtype == B.dtype == C.dtype assert X.shape[1] % block_len == 0 diff --git a/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py b/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py index 063ff17dd3..fbc7b69187 100644 --- a/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py +++ b/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py @@ -68,7 +68,8 @@ def _run_prefill_decode_equivalence( the exact fp32 recurrence (and each other). All paths use the production dt flow (raw dt + softplus + a per-head dt_bias), so this also checks prefill and decode apply the softplus/bias preprocessing consistently. ``state_dtype`` - is the recurrent-state precision; ``act_dtype`` the activation/buffer one.""" + is the recurrent-state precision; ``act_dtype`` the activation/buffer one. + """ device = current_platform.device_type rtol, atol = _prefill_tolerances(act_dtype) set_random_seed(seed) diff --git a/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py b/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py index 588ce52dbe..6e5ada19bc 100644 --- a/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py +++ b/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py @@ -95,7 +95,8 @@ def _run_standard_decode( ) -> None: """Drive ``num_steps`` decode steps and check every path against the trusted recurrence. ``state_dtype`` is the recurrent-state precision; - ``act_dtype`` is the activation/buffer precision.""" + ``act_dtype`` is the activation/buffer precision. + """ device = "cuda" both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 rtol, atol = _tolerances(torch.float32 if both_fp32 else torch.bfloat16) @@ -539,7 +540,8 @@ def test_replayssm_standard_decode_tp_head_shard_equivalence( separable: one run over all heads must equal concatenating ``tp`` independent per-rank runs on ``nheads // tp`` heads and ``ngroups // tp`` groups. This guards the per-rank divisors and group->head mapping the state-shape wiring - relies on. The real TP1==TP2 engine check lives in the v1/e2e suite.""" + relies on. The real TP1==TP2 engine check lives in the v1/e2e suite. + """ state_dtype, act_dtype = precision nheads, headdim, dstate, ngroups = geometry assert nheads % tp == 0 and ngroups % tp == 0 diff --git a/tests/kernels/mamba/utils.py b/tests/kernels/mamba/utils.py index b82c4d82c3..c7cf473427 100644 --- a/tests/kernels/mamba/utils.py +++ b/tests/kernels/mamba/utils.py @@ -9,8 +9,7 @@ from einops import rearrange, repeat def selective_state_update_ref( state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False ): - """ - Argument: + """Argument: state: (batch, dim, dstate) or (batch, nheads, dim, dstate) x: (batch, dim) or (batch, nheads, dim) dt: (batch, dim) or (batch, nheads, dim) @@ -20,8 +19,10 @@ def selective_state_update_ref( D: (dim,) or (nheads, dim) z: (batch, dim) or (batch, nheads, dim) dt_bias: (dim,) or (nheads, dim) + Return: out: (batch, dim) or (batch, nheads, dim) + """ has_heads = state.dim() > 3 if state.dim() == 3: diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 8041db68d7..dd17298cb3 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -148,9 +148,7 @@ class Config: return self.E // self.world_size def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]: - """ - make env data for vllm launch. - """ + """Make env data for vllm launch.""" vllm_config = VllmConfig() vllm_config.model_config = SimpleNamespace( enforce_eager=True, @@ -169,7 +167,8 @@ class Config: def fe_supports_quant_scheme(self) -> bool: """Check if the fused experts class supports this quant config. - See https://github.com/ROCm/aiter/issues/2419 for AITER gaps.""" + See https://github.com/ROCm/aiter/issues/2419 for AITER gaps. + """ if self.quant_config is None or self.quant_dtype is None: return True if self.quant_dtype != torch.float8_e4m3fn: @@ -444,9 +443,7 @@ class RankTensors: def make_hidden_states( config: Config, ) -> tuple[torch.Tensor, torch.Tensor | None]: - """ - Return hidden_states - """ + """Return hidden_states.""" m, k, dtype = (config.M, config.K, config.dtype) device = torch.accelerator.current_device_index() a = torch.randn((m, k), device=device, dtype=dtype) / 15.0 diff --git a/tests/kernels/moe/parallel_utils.py b/tests/kernels/moe/parallel_utils.py index f926b034b3..972ebe6512 100644 --- a/tests/kernels/moe/parallel_utils.py +++ b/tests/kernels/moe/parallel_utils.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -DeepEP test utilities -""" +"""DeepEP test utilities.""" import dataclasses import os diff --git a/tests/kernels/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py index d0b1ce7cae..de8f75d0e9 100644 --- a/tests/kernels/moe/test_b12x.py +++ b/tests/kernels/moe/test_b12x.py @@ -468,7 +468,8 @@ def test_deepseek_v4_flashinfer_cutlass_falls_through_to_w4a8( monkeypatch: pytest.MonkeyPatch, ) -> None: """Explicit flashinfer_cutlass must try the W4A8 variant when the BF16 - variant is unsupported (the BF16 variant is gated to SM90).""" + variant is unsupported (the BF16 variant is gated to SM90). + """ from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( FlashInferExperts, ) diff --git a/tests/kernels/moe/test_batched_deepgemm.py b/tests/kernels/moe/test_batched_deepgemm.py index 3fca9190dc..4c4c82d54e 100644 --- a/tests/kernels/moe/test_batched_deepgemm.py +++ b/tests/kernels/moe/test_batched_deepgemm.py @@ -34,7 +34,6 @@ def test_batched_deepgemm_vs_triton( E: int, T: int, K: int, N: int, topk: int, monkeypatch, workspace_init ): """Compare BatchedDeepGemmExperts to BatchedTritonExperts.""" - monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1") device = "cuda" diff --git a/tests/kernels/moe/test_batched_moe.py b/tests/kernels/moe/test_batched_moe.py index 9be58e6337..c4ebc5ff40 100644 --- a/tests/kernels/moe/test_batched_moe.py +++ b/tests/kernels/moe/test_batched_moe.py @@ -129,7 +129,8 @@ def test_batched_mm( per_act_token_quant: bool, ): """Note: float8_e4m3fn is not supported on CUDA architecture < 89, - and those tests will be skipped on unsupported hardware.""" + and those tests will be skipped on unsupported hardware. + """ set_random_seed(7) use_fp8_w8a8 = dtype == torch.float8_e4m3fn @@ -268,7 +269,8 @@ def test_fused_moe_batched_experts( workspace_init, ): """Note: float8_e4m3fn is not supported on CUDA architecture < 89, - and those tests will be skipped on unsupported hardware.""" + and those tests will be skipped on unsupported hardware. + """ set_random_seed(7) use_fp8_w8a8 = dtype == torch.float8_e4m3fn @@ -649,7 +651,8 @@ def test_batched_experts_masked_activation_end_to_end( def test_batched_experts_end_to_end(m, n, k, e, topk): """End-to-end BatchedTritonExperts via the reference (no-comms) BatchedPrepareAndFinalize, validated against a torch reference. Exercises - the device-enablement path.""" + the device-enablement path. + """ from vllm.v1.worker.workspace import init_workspace_manager set_random_seed(7) diff --git a/tests/kernels/moe/test_block_int8.py b/tests/kernels/moe/test_block_int8.py index 415faafb3a..84a4a0d2d1 100644 --- a/tests/kernels/moe/test_block_int8.py +++ b/tests/kernels/moe/test_block_int8.py @@ -54,7 +54,8 @@ SEEDS = [0] # For test def torch_w8a8_block_int8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape): """This function performs fused moe with block-wise quantization using - native torch.""" + native torch. + """ B, D = a.shape a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D) out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device) @@ -97,7 +98,8 @@ def setup_cuda(): @torch.inference_mode() def test_w8a8_block_int8_fused_moe(M, N, K, E, topk, block_size, dtype, seed): """Tests the fused_moe kernel with W8A8 INT8 block quantization against a - native torch reference.""" + native torch reference. + """ torch.manual_seed(seed) a = torch.randn((M, K), dtype=dtype) / 10 diff --git a/tests/kernels/moe/test_count_expert_num_tokens.py b/tests/kernels/moe/test_count_expert_num_tokens.py index 8014076192..b2b6abc5b4 100644 --- a/tests/kernels/moe/test_count_expert_num_tokens.py +++ b/tests/kernels/moe/test_count_expert_num_tokens.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests compute_expert_num_tokens kernels -""" +"""Tests compute_expert_num_tokens kernels.""" import dataclasses diff --git a/tests/kernels/moe/test_cpu_fused_moe.py b/tests/kernels/moe/test_cpu_fused_moe.py index dccda70c0b..e7af163b04 100644 --- a/tests/kernels/moe/test_cpu_fused_moe.py +++ b/tests/kernels/moe/test_cpu_fused_moe.py @@ -397,7 +397,8 @@ UNALIGNED_INTERMEDIATE_DIM = 176 class _StubMoELayer(torch.nn.Module): """Minimal stand-in for the real MoE layer module, exposing just what the unquantized CPU experts read/replace (w13_weight, w2_weight, the - router configuration, and optionally w13_bias/w2_bias).""" + router configuration, and optionally w13_bias/w2_bias). + """ def __init__( self, @@ -500,8 +501,8 @@ def test_cpu_fused_moe_unaligned_intermediate_size( act: MoEActivation, ): """CPU kernels handle unaligned intermediate sizes by zero-padding the - weights before prepacking.""" - + weights before prepacking. + """ set_random_seed(0) batch_size = 64 intermediate_size = UNALIGNED_INTERMEDIATE_DIM diff --git a/tests/kernels/moe/test_cpu_int4_moe.py b/tests/kernels/moe/test_cpu_int4_moe.py index 04931e386f..2d5534ba08 100644 --- a/tests/kernels/moe/test_cpu_int4_moe.py +++ b/tests/kernels/moe/test_cpu_int4_moe.py @@ -62,6 +62,7 @@ def _make_int4_moe_weights( Returns: (w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias) where *_ref are the dequantized float reference weights + """ # Generate INT4 weights as int8 values in [-8, 7] w13_int4 = torch.randint(-8, 8, (E, 2 * N, K), dtype=torch.int8) diff --git a/tests/kernels/moe/test_cpu_quant_fused_moe.py b/tests/kernels/moe/test_cpu_quant_fused_moe.py index 84ff1ae266..99cae4534b 100644 --- a/tests/kernels/moe/test_cpu_quant_fused_moe.py +++ b/tests/kernels/moe/test_cpu_quant_fused_moe.py @@ -826,6 +826,7 @@ def _make_int4_moe_weights(E, N, K, group_size, quant_algo): w1_zeros, w2_zeros, w1_zeros_packed, w2_zeros_packed, w1_s, w2_s + """ w1_int4 = torch.randint(0, 16, (E, K, 2 * N), dtype=torch.int32) w2_int4 = torch.randint(0, 16, (E, N, K), dtype=torch.int32) diff --git a/tests/kernels/moe/test_cutedsl_moe.py b/tests/kernels/moe/test_cutedsl_moe.py index 3a371968fc..09a7256a5c 100644 --- a/tests/kernels/moe/test_cutedsl_moe.py +++ b/tests/kernels/moe/test_cutedsl_moe.py @@ -91,11 +91,9 @@ def break_fp4_bytes(a, dtype): def generate_balanced_routing( hidden_states: torch.Tensor, num_experts: int, top_k: int ): - """ - Generate routing weights and topk indices such that every expert is active. + """Generate routing weights and topk indices such that every expert is active. Returns routing_weights, topk_idx """ - num_tokens, hidden_dim = hidden_states.shape # num_tokens = batch_size * seq_len @@ -232,13 +230,13 @@ def grouped_gemm_ref( *, block_size: int = 16, ) -> torch.Tensor: - """ - Computes the reference grouped GEMM (fp4 quantized per-expert loop), + """Computes the reference grouped GEMM (fp4 quantized per-expert loop), computes flashinfer grouped GEMM (for scale consistency), and returns ONLY the repacked reference output: out_ref. Returns: out_ref: Tensor [num_experts, max_m, n_out] + """ device_hs = hidden_states_expanded.device device_w = weights.device diff --git a/tests/kernels/moe/test_deepep_deepgemm_moe.py b/tests/kernels/moe/test_deepep_deepgemm_moe.py index efb1e2f296..b755741380 100644 --- a/tests/kernels/moe/test_deepep_deepgemm_moe.py +++ b/tests/kernels/moe/test_deepep_deepgemm_moe.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test DeepEP + DeepGEMM integration +"""Test DeepEP + DeepGEMM integration DeepGEMM are gemm kernels specialized for the fp8 block-quantized case. """ @@ -92,9 +91,7 @@ def make_block_quant_fp8_weights( k: int, block_size: list[int], ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Return weights w1q, w2q, w1_scale, w2_scale - """ + """Return weights w1q, w2q, w1_scale, w2_scale.""" (_, w1q, w1_scale, _), (_, w2q, w2_scale, _) = make_test_weights( e, n, k, torch.bfloat16, torch.float8_e4m3fn, block_shape=block_size ) @@ -459,10 +456,7 @@ def test_ht_deepep_deepgemm_moe( disable_deepgemm_ue8m0, workspace_init, ): - """ - Tests for High-Throughput DeepEP + DeepGemm integration. - """ - + """Tests for High-Throughput DeepEP + DeepGemm integration.""" m, n, k = mnk set_random_seed(7) @@ -533,9 +527,7 @@ def test_ll_deepep_deepgemm_moe( disable_deepgemm_ue8m0, workspace_init, ): - """ - Tests for Low-Latency DeepEP + DeepGemm integration. - """ + """Tests for Low-Latency DeepEP + DeepGemm integration.""" assert not is_deep_gemm_e8m0_used() m, n, k = mnk diff --git a/tests/kernels/moe/test_deepep_moe.py b/tests/kernels/moe/test_deepep_moe.py index ce9e4f343d..ea2113d64b 100644 --- a/tests/kernels/moe/test_deepep_moe.py +++ b/tests/kernels/moe/test_deepep_moe.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test deepep dispatch-combine logic -""" +"""Test deepep dispatch-combine logic.""" import dataclasses @@ -56,9 +54,7 @@ MAX_TOKENS_PER_RANK = 64 def make_weights( e, n, k, dtype ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Return weights w1, w2, w1_scale, w2_scale - """ + """Return weights w1, w2, w1_scale, w2_scale.""" if dtype in [torch.float16, torch.bfloat16]: w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10 w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10 diff --git a/tests/kernels/moe/test_deepep_v2_moe.py b/tests/kernels/moe/test_deepep_v2_moe.py index f1e34d5bd7..2fbde2f86e 100644 --- a/tests/kernels/moe/test_deepep_v2_moe.py +++ b/tests/kernels/moe/test_deepep_v2_moe.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test DeepEP v2 (ElasticBuffer) dispatch-combine logic. +"""Test DeepEP v2 (ElasticBuffer) dispatch-combine logic. Compares against a pure-PyTorch reference MoE implementation. """ diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index ae6c9685a4..f6df64c51b 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit-test DeepGEMM FP8 and FP4 kernels (no DeepEP). +"""Unit-test DeepGEMM FP8 and FP4 kernels (no DeepEP). Compare DeepGEMM path against the Triton fallback inside vLLM's fused_experts. """ @@ -79,8 +78,7 @@ def make_block_quant_fp8_weights( k: int, block_size: list[int], ): - """ - Generate (w1, w2) expert weights and their per-block scale tensors + """Generate (w1, w2) expert weights and their per-block scale tensors in FP8 block-quantized format. w1 shape: (E, 2N, K) @@ -121,8 +119,7 @@ def make_block_quant_fp8_weights( def run_single_case(m, n, k, topk, num_experts, block_size): - """ - Run one (M,N,K) configuration on a single GPU and assert DeepGEMM == + """Run one (M,N,K) configuration on a single GPU and assert DeepGEMM == Triton baseline within tolerance. """ tokens_bf16 = ( @@ -248,8 +245,7 @@ def make_mxfp4_weights( n: int, k: int, ): - """ - Generate (w1, w2) expert weights in MXFP4 packed format with float32 scales, + """Generate (w1, w2) expert weights in MXFP4 packed format with float32 scales, plus BF16 reference weights for validation. w1 shape: (E, 2N, K//2) uint8 — packed FP4 @@ -310,8 +306,7 @@ def _bf16_moe_reference(x, w1, w2, topk_weights, topk_ids): def run_single_fp4_case(m, n, k, topk, num_experts): - """ - Run one (M,N,K) configuration with FP4 weights on DeepGEMM and assert + """Run one (M,N,K) configuration with FP4 weights on DeepGEMM and assert DeepGEMM FP4 == BF16 reference within tolerance. """ tokens_bf16 = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * (k**-0.5) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 2e022ede36..9627e2c67d 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -464,7 +464,8 @@ def test_trtllm_fp8_swiglu_clamp_support( ): """FlashInfer >= 0.6.18 applies the SwiGLU clamp for both block-scaled kernels with a SwiGLU activation (DeepSeek-V4 sets swiglu_limit); the - per-tensor kernel has no clamp, and Relu2 rejects the parameters.""" + per-tensor kernel has no clamp, and Relu2 rejects the parameters. + """ class _Experts(TrtLlmFp8ExpertsModular): @staticmethod diff --git a/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py b/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py index d8e449ec0a..3d637258d3 100644 --- a/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py +++ b/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py @@ -83,7 +83,7 @@ def _reference_activation( beta: float | None, limit: float | None, ): - """vLLM's own op for this activation, so the reference is not re-derived.""" + """VLLM's own op for this activation, so the reference is not re-derived.""" if activation == MoEActivation.RELU2_NO_MUL: return lambda x: torch.square(torch.relu(x)) if activation == MoEActivation.SWIGLUOAI: diff --git a/tests/kernels/moe/test_fused_topk.py b/tests/kernels/moe/test_fused_topk.py index 9c196a884d..18d1d9b4df 100644 --- a/tests/kernels/moe/test_fused_topk.py +++ b/tests/kernels/moe/test_fused_topk.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the MoE fused topk kernel +"""Tests for the MoE fused topk kernel. Run `pytest tests/kernels/moe/test_fused_topk.py`. """ diff --git a/tests/kernels/moe/test_gpt_oss_triton_kernels.py b/tests/kernels/moe/test_gpt_oss_triton_kernels.py index 4afcba706c..0a16a31797 100644 --- a/tests/kernels/moe/test_gpt_oss_triton_kernels.py +++ b/tests/kernels/moe/test_gpt_oss_triton_kernels.py @@ -377,7 +377,8 @@ def test_unit_shuffle(): @pytest.mark.parametrize("n_experts,topk", [(32, 4), (128, 4)]) def test_routing_data_from_sparse_topk_parity(n_tokens, n_experts, topk): """routing_data_from_sparse_topk must produce routing structures - identical to make_routing_data for the same topk result.""" + identical to make_routing_data for the same topk result. + """ from vllm.model_executor.layers.fused_moe.experts import ( gpt_oss_triton_kernels_moe as gptoss_moe, ) diff --git a/tests/kernels/moe/test_grouped_topk.py b/tests/kernels/moe/test_grouped_topk.py index 1ab8550d64..c65cd381db 100644 --- a/tests/kernels/moe/test_grouped_topk.py +++ b/tests/kernels/moe/test_grouped_topk.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the MoE grouped topk kernel +"""Tests for the MoE grouped topk kernel. Run `pytest tests/kernels/moe/test_grouped_topk.py`. """ diff --git a/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py b/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py index 6756d479fe..b647439c49 100644 --- a/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py +++ b/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py @@ -31,6 +31,7 @@ def mxint4_quantize( Returns: - uint8 packed (2 INT4/byte): [..., k//2] - stores SIGNED INT4 [-8, 7] - scales in BF16: [..., k//sf_vec_size] + """ x_reshaped = x.reshape(-1, sf_vec_size) x_max = x_reshaped.max(dim=-1, keepdim=True)[0].to(torch.float32) @@ -61,6 +62,7 @@ def mxint4_quantize_moe_weights( Returns: - weights_mxint4: Quantized weights [e, n, k//2] uint8 - scales_mxint4: Quantization scales [e, n, k//group_size] bf16 + """ e = weights_bf16.shape[0] weight_list = [] @@ -101,6 +103,7 @@ def marlin_quantize_moe_weights( Returns: - weights_marlin: Marlin quantized weights [e, k//8, n] int32 - scales_marlin: Marlin quantization scales [e, k//group_size, n] bf16 + """ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import ( marlin_quantize, @@ -286,7 +289,8 @@ def test_marlin_vs_trtllm_mxint4_moe_kimik2(monkeypatch, m, n, k, e, topk, group @torch.inference_mode() def test_flashinfer_trtllm_mxint4_moe_wrapper(m, n, k, e, topk): """Test that the flashinfer_trtllm_mxint4_moe wrapper matches the raw - trtllm_mxint4_block_scale_moe kernel call.""" + trtllm_mxint4_block_scale_moe kernel call. + """ pytest.importorskip("flashinfer") from flashinfer import RoutingMethodType from flashinfer.fused_moe import trtllm_mxint4_block_scale_moe diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index 0c0e1d61f9..6c21eb0d9a 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -352,7 +352,8 @@ def test_modular_kernel_combinations_singlegpu( workspace_init, ): """Note: float8_e4m3fn is not supported on CUDA architecture < 89, - and those tests will be skipped on unsupported hardware.""" + and those tests will be skipped on unsupported hardware. + """ config = Config( Ms=Ms, K=k, diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index d0fd4df23c..47be82dd2a 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test modular OAI Triton MoE -""" +"""Test modular OAI Triton MoE.""" from __future__ import annotations diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index ea732d6887..8b4efaf913 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -119,8 +119,7 @@ def iterative_moe( expert_map: torch.Tensor = None, renormalize: bool = False, ) -> torch.Tensor: - """ - Baseline implementation of fused moe. + """Baseline implementation of fused moe. Args: hidden_states: [*, hidden_size] @@ -128,6 +127,7 @@ def iterative_moe( w2: [num_experts, hidden_size, intermediate_size] gating_output: [*, num_experts] expert_map: [num_experts] + """ orig_shape = hidden_states.shape hidden_size = hidden_states.shape[-1] @@ -1956,8 +1956,7 @@ def test_batched_fused_marlin_moe( ) def is_valid(self): - """ - Return True only if the input can be represented in a Batched + """Return True only if the input can be represented in a Batched format. """ return torch.all(self.expert_num_tokens_cpu <= self.max_tokens_per_batch) @@ -2078,9 +2077,7 @@ def test_unquantized_bf16_flashinfer_trtllm_backend( dtype: torch.dtype, workspace_init, ): - """ - Test BF16 unquantized MoE with FlashInfer TRTLLM backend. - """ + """Test BF16 unquantized MoE with FlashInfer TRTLLM backend.""" set_random_seed(7) from vllm.model_executor.layers.fused_moe.config import ( diff --git a/tests/kernels/moe/test_moe_align_block_size.py b/tests/kernels/moe/test_moe_align_block_size.py index 8c16690785..a1599f5ad3 100644 --- a/tests/kernels/moe/test_moe_align_block_size.py +++ b/tests/kernels/moe/test_moe_align_block_size.py @@ -54,8 +54,7 @@ def _verify_expert_level_sorting( valid_length: int, total_tokens: int, ): - """ - Verify that actual_sorted_ids follows the correct expert-level sorting. + """Verify that actual_sorted_ids follows the correct expert-level sorting. The kerne limplementation may or may not preserve original token order in topk_ids in the final sorted_ids however this does not impact quality. """ @@ -96,8 +95,7 @@ def torch_moe_align_block_size( expert_map: torch.Tensor | None = None, pad_sorted_ids: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Golden torch implementation of moe_align_block_size. + """Golden torch implementation of moe_align_block_size. This function aligns the token distribution across experts to be compatible with block size for matrix multiplication by sorting tokens by expert and @@ -189,7 +187,7 @@ def torch_moe_align_block_size( def test_moe_align_block_size( m: int, topk: int, num_experts: int, block_size: int, pad_sorted_ids: bool ): - """Test moe_align_block_size without expert mapping""" + """Test moe_align_block_size without expert mapping.""" topk_ids = torch.zeros((m, topk), device="cuda", dtype=torch.int32) for i in range(m): experts = torch.randperm(num_experts, device="cuda")[:topk] @@ -254,7 +252,7 @@ def test_moe_align_block_size_with_expert_map( block_size: int, mask_inactive_experts: bool, ): - """Test moe_align_block_size with expert mapping (EP scenario)""" + """Test moe_align_block_size with expert mapping (EP scenario).""" expert_map = torch.full((num_experts,), -1, device="cuda", dtype=torch.int32) local_experts = list(range(0, num_experts, 2)) for i, expert_id in enumerate(local_experts): diff --git a/tests/kernels/moe/test_moe_kernel_oracle.py b/tests/kernels/moe/test_moe_kernel_oracle.py index fbdf804a3b..6ba5753f8d 100644 --- a/tests/kernels/moe/test_moe_kernel_oracle.py +++ b/tests/kernels/moe/test_moe_kernel_oracle.py @@ -21,7 +21,8 @@ from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( class TestUnquantizedDelegation: """UnquantizedMoEKernelOracle methods must delegate to the existing - module-level functions; behaviour is bit-identical.""" + module-level functions; behaviour is bit-identical. + """ def test_make_kernel_delegates(self) -> None: quant_config = object() diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index e1a152a233..38b9dd0252 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -240,8 +240,7 @@ def maybe_roundup_layer_hidden_size( act_dtype: torch.dtype, backend: str | None, ) -> int: - """ - Given layer hidden size and MoE configurations, round up hidden_size + """Given layer hidden size and MoE configurations, round up hidden_size if necessary. Args: @@ -253,6 +252,7 @@ def maybe_roundup_layer_hidden_size( Rounded up hidden_size if rounding up is required based on the configs and all2all backend. Original hidden size otherwise. + """ if backend == "deepep_high_throughput": from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import ( @@ -315,7 +315,8 @@ def tp_chunk_gate_up( device: torch.device | int | None = None, ) -> torch.Tensor: """TP-chunk a combined [gate; up] weight, splitting each half separately - so every rank gets a portion of both gate and up.""" + so every rank gets a portion of both gate and up. + """ half = w.shape[dim] // 2 gate = chunk_by_rank( w.narrow(dim, 0, half), tp_rank, tp_size, dim=dim, device=device @@ -945,6 +946,7 @@ def create_shared_experts_from_config( Returns: TestMLP instance or None if config is None + """ if shared_experts_config is None: return None @@ -991,6 +993,7 @@ def setup_moe_test_data( Returns: MoETestData containing all test data and transforms + """ # For latent MoE: latent_size = k // 2 latent_size = k // 2 @@ -1710,7 +1713,6 @@ def test_moe_layer_no_parallel( monkeypatch, ): """Test MoE layer without parallelism (dp_size=1, tp_size=1, use_ep=False).""" - if os.environ.get("VLLM_LOGGING_LEVEL") is None: monkeypatch.setenv("VLLM_LOGGING_LEVEL", "ERROR") diff --git a/tests/kernels/moe/test_moe_permute_unpermute.py b/tests/kernels/moe/test_moe_permute_unpermute.py index 4089bbfe84..76d1e57f58 100644 --- a/tests/kernels/moe/test_moe_permute_unpermute.py +++ b/tests/kernels/moe/test_moe_permute_unpermute.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the MOE permute/unpermute kernel +"""Tests for the MOE permute/unpermute kernel. Run `pytest tests/kernels/test_moe_permute_unpermute.py`. """ diff --git a/tests/kernels/moe/test_mxfp4_moe.py b/tests/kernels/moe/test_mxfp4_moe.py index 16b233b935..7c7ca10885 100644 --- a/tests/kernels/moe/test_mxfp4_moe.py +++ b/tests/kernels/moe/test_mxfp4_moe.py @@ -67,8 +67,7 @@ def compute_ref_output( @pytest.mark.parametrize("num_experts", [8, 16, 32]) @pytest.mark.parametrize("out_dtype", [torch.bfloat16]) def test_cutlass_mxfp4_grouped_mm(num_experts, out_dtype): - """ - Test the MXFP4 grouped GEMM kernel by: + """Test the MXFP4 grouped GEMM kernel by: 1. Creating random per-expert inputs and weights 2. Quantizing both to MXFP4 using the CUDA kernel 3. Running the CUTLASS grouped GEMM @@ -205,9 +204,7 @@ def test_cutlass_mxfp4_grouped_mm(num_experts, out_dtype): reason="mxfp4_experts_quant requires CUDA SM100", ) def test_mxfp4_experts_quant_basic(): - """ - Basic smoke test for the MXFP4 experts quantization kernel. - """ + """Basic smoke test for the MXFP4 experts quantization kernel.""" device = "cuda" num_experts = 4 k = 256 @@ -295,8 +292,7 @@ def compute_reference_e8m0_scale(block_max: float) -> int: @pytest.mark.parametrize("k", [256, 7168]) @pytest.mark.parametrize("m", [16, 64]) def test_mxfp4_experts_quant_e8m0_scale_correctness(m, k): - """ - Test that mxfp4_experts_quant computes E8M0 block scales correctly + """Test that mxfp4_experts_quant computes E8M0 block scales correctly per OCP MX spec (not the NVFP4 formula). The old buggy kernel used: floor(log2(max/6)) + 127 @@ -416,8 +412,7 @@ def test_mxfp4_experts_quant_e8m0_scale_correctness(m, k): reason="mxfp4_experts_quant requires CUDA SM100", ) def test_mxfp4_experts_quant_no_saturation(): - """ - Test that the E8M0 scale is large enough to avoid unexpected saturation. + """Test that the E8M0 scale is large enough to avoid unexpected saturation. With the buggy NVFP4 formula, the scale was too small causing most values to saturate to ±6 in FP4. The fixed OCP MX formula should ensure that diff --git a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py index 830ff48734..e94882a310 100644 --- a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py +++ b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py @@ -130,7 +130,8 @@ def test_is_supported_config(present, ep_size, supported, reason_substr): def test_explicit_moe_backend_aiter(): """--moe-backend aiter: returns FlyDSL when usable (TP or EP), else a clear - ValueError when the flydsl package is missing.""" + ValueError when the flydsl package is missing. + """ with _gfx950(), _flydsl_installed(True): assert ( _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1)) @@ -161,8 +162,9 @@ def test_gfx950_picks_aiter(): def test_gfx942_picks_emulation(): - """flydsl unusable (e.g. gfx942, no FlyDSL support) -> native Triton - dot_scaled backend wins instead.""" + """Flydsl unusable (e.g. gfx942, no FlyDSL support) -> native Triton + dot_scaled backend wins instead. + """ with patch(f"{_AITER_MOD}.current_platform.supports_mx", return_value=False): backend, experts_cls = select_mxfp8_moe_backend(_config()) assert backend is Fp8MoeBackend.EMULATION diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 9157d4eaa2..d513a14f77 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -214,8 +214,7 @@ def reference_moe( activation: str = "swiglu", use_interleaved_layout: bool = False, ): - """ - Reference MoE implementation for accuracy testing. + """Reference MoE implementation for accuracy testing. Args: activation: One of "swiglu", "silu", "relu2". Controls the activation @@ -224,6 +223,7 @@ def reference_moe( (gate=x[..., ::2], up=x[..., 1::2]) as used by SWIGLUOAI. If False, uses chunked layout (gate, up = chunk(x, 2)) as used by standard swiglu/silu. + """ # renormalize routing experts = torch.topk(roouting_logits, k=topk, dim=-1, sorted=True) @@ -1441,8 +1441,7 @@ def test_rocm_mxfp4_moe_oracle( intermediate_size: int, monkeypatch: pytest.MonkeyPatch, ): - """ - Test ROCm MXFP4 MoE using oracle functions. + """Test ROCm MXFP4 MoE using oracle functions. This test validates that the oracle functions work end-to-end: - select_mxfp4_moe_backend() selects a valid backend @@ -1825,8 +1824,7 @@ def test_mxfp4_emulation_rounds_up_to_block_size( def test_select_mxfp4_moe_backend_raises_with_unsupported_reasons( monkeypatch: pytest.MonkeyPatch, ): - """ - select_mxfp4_moe_backend() must raise NotImplementedError, with the + """select_mxfp4_moe_backend() must raise NotImplementedError, with the collected per-backend unsupported reasons in the message, when no backend supports the requested deployment configuration. """ diff --git a/tests/kernels/moe/test_rocm_aiter_moe.py b/tests/kernels/moe/test_rocm_aiter_moe.py index 0967f9499e..de4ae5ae31 100644 --- a/tests/kernels/moe/test_rocm_aiter_moe.py +++ b/tests/kernels/moe/test_rocm_aiter_moe.py @@ -249,7 +249,8 @@ def _make_topk_ids( device: str = "cuda", ) -> torch.Tensor: """Generate distinct expert IDs per token with the same top-k shape as - production routing.""" + production routing. + """ router_logits = torch.randn(num_tokens, num_experts, device=device) _, topk_ids = torch.topk(torch.softmax(router_logits, dim=-1), k=topk, dim=-1) return topk_ids.to(torch.int32) @@ -443,7 +444,8 @@ def test_aiter_fused_moe_custom_op_registered(): def test_aiter_asm_moe_tkw1_custom_op_registered(): """The tkw1 custom op should stay registered for FP8 apply-router-weight - paths.""" + paths. + """ _assert_aiter_supported() import vllm._aiter_ops as aiter_ops # noqa: F401 @@ -453,7 +455,8 @@ def test_aiter_asm_moe_tkw1_custom_op_registered(): def test_aiter_fused_moe_fake_tensor_support(): """The fused-MoE op should preserve fake-tensor compatibility for - torch.compile-style tracing.""" + torch.compile-style tracing. + """ _assert_aiter_supported() import vllm._aiter_ops # noqa: F401 @@ -516,7 +519,8 @@ def test_aiter_moe_enablement_follows_env( monkeypatch: pytest.MonkeyPatch, ): """The fused-MoE gate should depend only on the main AITER toggle and the - MoE-specific toggle.""" + MoE-specific toggle. + """ from vllm._aiter_ops import rocm_aiter_ops _assert_aiter_supported() @@ -547,7 +551,8 @@ def test_aiter_moe_shared_experts_enablement_follows_env( monkeypatch: pytest.MonkeyPatch, ): """Shared-expert fusion should only be enabled when the fused-MoE path is - enabled too.""" + enabled too. + """ from vllm._aiter_ops import rocm_aiter_ops _assert_aiter_supported() @@ -688,7 +693,8 @@ def test_aiter_mxfp4_quant_scheme_support_matches_gfx950(): @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") def test_aiter_fused_moe_mi350_mxfp4_w4a16_accuracy(): """The gfx950 AITER MXFP4 W4A16 MoE path should match the dequantized - MXFP4 reference.""" + MXFP4 reference. + """ from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( rocm_aiter_fused_experts, @@ -737,7 +743,8 @@ def test_aiter_fused_moe_mi350_mxfp4_w4a16_accuracy(): @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") def test_aiter_fused_moe_mi350_mxfp4_w4a16_determinism(): """The gfx950 AITER MXFP4 W4A16 MoE path should stay bitwise - deterministic.""" + deterministic. + """ from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( rocm_aiter_fused_experts, @@ -830,7 +837,8 @@ def test_aiter_fused_moe_bf16_accuracy( intermediate_dim: int, ): """The ROCm AITER fused-MoE BF16 path should match the float32 reference - on representative shapes.""" + on representative shapes. + """ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( ActivationMethod, QuantMethod, @@ -875,7 +883,8 @@ def test_aiter_fused_moe_bf16_accuracy( def test_aiter_fused_moe_gelu_accuracy(): """The GELU activation variant should stay aligned with the float32 - reference.""" + reference. + """ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( ActivationMethod, QuantMethod, @@ -920,7 +929,8 @@ def test_aiter_fused_moe_gelu_accuracy(): def test_aiter_fused_moe_determinism(): """The BF16 fused-MoE kernel should stay bitwise deterministic for the - same inputs.""" + same inputs. + """ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( ActivationMethod, QuantMethod, @@ -1075,7 +1085,8 @@ def test_aiter_fused_moe_mi3xx_bf16_accuracy(): ) def test_aiter_fused_moe_mi3xx_fp8_accuracy(): """The MI3xx FP8 per-tensor MoE path should stay within the measured FP8 - error budget.""" + error budget. + """ from tests.kernels.moe.utils import make_test_weights from vllm._aiter_ops import rocm_aiter_ops from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( diff --git a/tests/kernels/moe/test_routed_experts_capture_monolithic.py b/tests/kernels/moe/test_routed_experts_capture_monolithic.py index eec6180550..ba5eef6624 100644 --- a/tests/kernels/moe/test_routed_experts_capture_monolithic.py +++ b/tests/kernels/moe/test_routed_experts_capture_monolithic.py @@ -185,7 +185,8 @@ def test_trtllm_bf16_monolithic_routing_replay_records_valid_experts( ) -> None: """The capture callback should receive the int16 routed-expert IDs the kernel actually used, the values should be valid expert indices, and - each token should pick ``top_k`` distinct experts.""" + each token should pick ``top_k`` distinct experts. + """ if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP: pytest.skip( f"DSV3 requires top_k <= n_group * topk_group " @@ -264,7 +265,8 @@ def test_trtllm_bf16_monolithic_routing_replay_non_dsv3( routing_method: RoutingMethodType, ) -> None: """Routing replay works for non-DeepSeekV3 routing methods too. - FlashInfer's ``routing_replay_out`` is routing-method-agnostic.""" + FlashInfer's ``routing_replay_out`` is routing-method-agnostic. + """ torch.manual_seed(0) device = torch.device("cuda:0") @@ -317,7 +319,8 @@ def test_trtllm_bf16_monolithic_routing_replay_non_dsv3( def test_trtllm_bf16_monolithic_capture_disabled_skips_buffer_alloc() -> None: """With no callback installed the kernel should not see a - ``routing_replay_out`` tensor — verify the helper short-circuits.""" + ``routing_replay_out`` tensor — verify the helper short-circuits. + """ torch.manual_seed(0) device = torch.device("cuda:0") experts, _, _ = _make_bf16_monolithic_experts( @@ -339,7 +342,8 @@ def test_trtllm_bf16_monolithic_capture_disabled_skips_buffer_alloc() -> None: def test_trtllm_bf16_monolithic_supports_capture_for_all_routing() -> None: """FlashInfer's ``routing_replay_out`` is supported by all routing methods, so ``supports_routing_replay_capture`` should be True - regardless of routing method.""" + regardless of routing method. + """ device = torch.device("cuda:0") for routing_method in ( RoutingMethodType.DeepSeekV3, @@ -361,7 +365,8 @@ def test_trtllm_bf16_monolithic_supports_capture_for_all_routing() -> None: def test_trtllm_bf16_monolithic_capture_buffer_shape_and_dtype() -> None: """When capture is installed, the allocated buffer is int16 and shaped - ``(num_tokens, experts_per_token)``.""" + ``(num_tokens, experts_per_token)``. + """ device = torch.device("cuda:0") experts, _, _ = _make_bf16_monolithic_experts( num_experts=_DSV3_NUM_EXPERTS, @@ -592,7 +597,8 @@ def test_trtllm_fp8_block_scale_monolithic_routing_replay_records_valid_experts( top_k: int, ) -> None: """End-to-end: ``TrtLlmFp8ExpertsMonolithic`` (DeepSeekFp8 block-scale - path, DSV3 routing) captures valid expert IDs.""" + path, DSV3 routing) captures valid expert IDs. + """ if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP: pytest.skip( f"DSV3 requires top_k <= n_group * topk_group " @@ -778,7 +784,8 @@ def _run_nvfp4_monolithic( routing_bias: torch.Tensor, ) -> torch.Tensor: """The monolithic NVFP4 apply expects packed fp4 hidden states + the - matching fp8 per-block scale stored in the ``a1q_scale`` slot.""" + matching fp8 per-block scale stored in the ``a1q_scale`` slot. + """ # Stash the weight tensors on the experts in the locations the apply() # implementation reads from (it pulls them from quant_config / scales # already; w1/w2 come in as args). @@ -806,7 +813,8 @@ def test_trtllm_nvfp4_monolithic_routing_replay_records_valid_experts( top_k: int, ) -> None: """End-to-end: ``TrtLlmNvFp4ExpertsMonolithic`` captures valid expert IDs - on the DSV3 routing path.""" + on the DSV3 routing path. + """ if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP: pytest.skip( f"DSV3 requires top_k <= n_group * topk_group " diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py index 0eded09b1a..f922a3ebd3 100644 --- a/tests/kernels/moe/test_routing.py +++ b/tests/kernels/moe/test_routing.py @@ -214,8 +214,7 @@ def assert_routing_results_close( rtol: float = 1e-3, atol: float = 1e-3, ): - """ - Compare routing results, sorting by expert ID first to handle non-deterministic + """Compare routing results, sorting by expert ID first to handle non-deterministic ordering from sorted=False in topk. """ # Sort both results by expert IDs for consistent comparison @@ -251,7 +250,8 @@ def assert_aiter_routing_valid( the Python baseline (different group selection, scoring internals), so numerical comparison is not meaningful. Instead we verify the outputs satisfy the routing contract: correct shapes, valid expert - IDs, non-negative weights, and proper normalization.""" + IDs, non-negative weights, and proper normalization. + """ n_tokens = topk_weights.shape[0] # Shape @@ -294,8 +294,7 @@ def assert_aiter_routing_valid( def baseline_fused_topk( router_logits: torch.Tensor, top_k: int, renormalize: bool ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Baseline for standard fused top-k routing. + """Baseline for standard fused top-k routing. Algorithm: 1. Apply softmax to router logits @@ -320,8 +319,7 @@ def baseline_fused_topk_bias( e_score_correction_bias: torch.Tensor, routed_scaling_factor: float, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Baseline for fused top-k with bias correction. + """Baseline for fused top-k with bias correction. Algorithm: 1. Apply softmax to router logits @@ -364,8 +362,7 @@ def baseline_grouped_topk( e_score_correction_bias: torch.Tensor | None, routed_scaling_factor: float, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Baseline for grouped top-k routing (e.g., DeepSeek). + """Baseline for grouped top-k routing (e.g., DeepSeek). Algorithm: 1. Apply scoring function (softmax or sigmoid) @@ -436,8 +433,7 @@ def baseline_grouped_topk( def baseline_custom_llama4( router_logits: torch.Tensor, top_k: int ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Baseline for Llama4 custom routing. + """Baseline for Llama4 custom routing. Algorithm: 1. Select top-k expert indices (without softmax) diff --git a/tests/kernels/moe/test_routing_simulator.py b/tests/kernels/moe/test_routing_simulator.py index 35b1af618c..b1adcd53d9 100644 --- a/tests/kernels/moe/test_routing_simulator.py +++ b/tests/kernels/moe/test_routing_simulator.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test script for the token-to-expert routing simulator. +"""Test script for the token-to-expert routing simulator. This script demonstrates how to use the routing simulator to test different routing strategies and analyze their performance, including @@ -77,7 +76,8 @@ def test_basic_functionality( def test_routing_strategy_integration(monkeypatch, device): """Test that the routing strategy environment variable works with - FusedMoEFactory.""" + FusedMoEFactory. + """ pytest.importorskip("vllm.model_executor.layers.fused_moe.layer") import vllm.envs as envs @@ -158,7 +158,8 @@ def test_routing_strategy_integration(monkeypatch, device): def test_distribution_based_routing_with_custom_strategy(): """Test registering and using DistributionBasedRouting with custom - parameters.""" + parameters. + """ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Register custom distribution-based strategy diff --git a/tests/kernels/moe/test_shared_fused_moe_routed_transform.py b/tests/kernels/moe/test_shared_fused_moe_routed_transform.py index af9455b960..895e168b2d 100644 --- a/tests/kernels/moe/test_shared_fused_moe_routed_transform.py +++ b/tests/kernels/moe/test_shared_fused_moe_routed_transform.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for FusedMoEFactory with routed_input_transform. +"""Tests for FusedMoEFactory with routed_input_transform. Verifies that applying routed_input_transform inside FusedMoEFactory produces the same results as applying the transform manually outside. diff --git a/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py b/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py index 85b992ad65..0ca08c50b3 100644 --- a/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py +++ b/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py @@ -142,9 +142,7 @@ def silu_mul_quant( def pack_scales(x: torch.Tensor, tokens_per_expert: torch.Tensor) -> torch.Tensor: - """ - pack float32 scales into a int32 tensor - """ + """Pack float32 scales into a int32 tensor.""" assert x.dtype == torch.float32 E, T, G = x.size() @@ -170,8 +168,7 @@ def ref_with_scale_fmt( up: torch.Tensor, scale_fmt: DeepGemmQuantScaleFMT, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - The precision types of the operations triggered by this function + """The precision types of the operations triggered by this function match closely with the kernel implementation so we compare more accurately. """ @@ -204,8 +201,7 @@ def ref_with_scale_fmt( def token_random(E, T, H2, tokens_per_expert): - """ - Initialize each token in a random range so we test a range of + """Initialize each token in a random range so we test a range of scale values. """ y = torch.empty((E, T, H2), dtype=torch.bfloat16, device=DEVICE) diff --git a/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py b/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py index bfe763b44a..76744da8f4 100644 --- a/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py +++ b/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py @@ -19,11 +19,9 @@ GROUP_SIZE = 128 def reference_quant(x: torch.Tensor, use_ue8m0: bool): - """ - Reference triton quant kernel from, + """Reference triton quant kernel from, vllm.model_executor.layers.quantization.utils.fp8_utils """ - x_q = torch.empty_like(x, device=x.device, dtype=FLOAT8_DTYPE) # Allocate the scale tensor in column-major format. @@ -72,7 +70,8 @@ def reference_with_clamp( ) -> tuple[torch.Tensor, torch.Tensor]: """Pre-clamp inputs (gate from above, up symmetric) at the input dtype to match the C++ compute() template, then run the standard silu_and_mul + - quant reference.""" + quant reference. + """ N_2 = x.size(1) // 2 dtype = x.dtype gate = x[..., :N_2].to(torch.float32).clamp(max=clamp_limit).to(dtype) diff --git a/tests/kernels/moe/test_triton_moe_ptpc_fp8.py b/tests/kernels/moe/test_triton_moe_ptpc_fp8.py index 01246325c8..15cef0df38 100644 --- a/tests/kernels/moe/test_triton_moe_ptpc_fp8.py +++ b/tests/kernels/moe/test_triton_moe_ptpc_fp8.py @@ -28,7 +28,8 @@ if current_platform.is_fp8_fnuz(): def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16): """Matrix multiplication function that supports per-token input - quantization and per-column weight quantization""" + quantization and per-column weight quantization + """ A = A.to(torch.float32) B = B.to(torch.float32) @@ -56,8 +57,8 @@ def fp8_mask(a, mask): def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk): """This function performs fused moe with per-column int8 - quantization using native torch.""" - + quantization using native torch. + """ B, D = a.shape # Perform per-token quantization a_q, a_s = ops.scaled_fp8_quant(a, use_per_token_if_dynamic=True) diff --git a/tests/kernels/moe/test_trtllm_bf16_moe.py b/tests/kernels/moe/test_trtllm_bf16_moe.py index 7b6b7d81ff..9e50b4fc3b 100644 --- a/tests/kernels/moe/test_trtllm_bf16_moe.py +++ b/tests/kernels/moe/test_trtllm_bf16_moe.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for the FlashInfer TRTLLM BF16 MoE backend +"""Tests for the FlashInfer TRTLLM BF16 MoE backend (`TrtLlmBf16ExpertsModular`). This mirrors the TRTLLM NvFP4 modular test shape: construct the modular diff --git a/tests/kernels/moe/test_trtllm_nvfp4_moe.py b/tests/kernels/moe/test_trtllm_nvfp4_moe.py index 59eaa4c139..b9684396b2 100644 --- a/tests/kernels/moe/test_trtllm_nvfp4_moe.py +++ b/tests/kernels/moe/test_trtllm_nvfp4_moe.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for the FlashInfer TRTLLM NvFP4 MoE backend +"""Tests for the FlashInfer TRTLLM NvFP4 MoE backend (`TrtLlmNvFp4ExpertsModular`). Covers the activations the wrapper claims to support — SiLU, RELU^2 (non-gated), diff --git a/tests/kernels/moe/test_unquantized_backend_selection.py b/tests/kernels/moe/test_unquantized_backend_selection.py index d4c037d978..48e98f9c14 100644 --- a/tests/kernels/moe/test_unquantized_backend_selection.py +++ b/tests/kernels/moe/test_unquantized_backend_selection.py @@ -150,7 +150,8 @@ def test_select_default_backend_by_platform( expected_backend, ): """Test default backend selection per platform with all optional - accelerators (FlashInfer, AITER) disabled.""" + accelerators (FlashInfer, AITER) disabled. + """ with patch( "vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform" ) as mock_platform: @@ -444,7 +445,8 @@ def test_select_cuda_flashinfer_cutlass_backend( mock_has_flashinfer, ): """Test CUDA backend selection when FlashInfer TRTLLM is not available - and FlashInfer CUTLASS is available.""" + and FlashInfer CUTLASS is available. + """ with ( patch.object(current_platform, "is_cuda", return_value=True), patch.object(current_platform, "is_rocm", return_value=False), diff --git a/tests/kernels/moe/test_zero_expert_moe.py b/tests/kernels/moe/test_zero_expert_moe.py index c2084800f9..0202b380c5 100644 --- a/tests/kernels/moe/test_zero_expert_moe.py +++ b/tests/kernels/moe/test_zero_expert_moe.py @@ -83,8 +83,7 @@ def test_zero_expert_moe_router_is_zero_expert_router(zero_expert_moe, num_token @pytest.mark.parametrize("num_tokens", [1, 32]) def test_zero_expert_moe_forward(zero_expert_moe, num_tokens): - """ - Run a forward pass through FusedMoEFactory with zero experts + """Run a forward pass through FusedMoEFactory with zero experts and verify output shape. """ layer, vllm_config = zero_expert_moe diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index e91fa58f3e..e762472950 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -64,8 +64,7 @@ def make_dummy_moe_config( max_num_tokens: int = 512, activation: MoEActivation = MoEActivation.SILU, ) -> FusedMoEConfig: - """ - This is a dummy config for the mk constructor interface + """This is a dummy config for the mk constructor interface as most kernels like DeepGEMM, CUTLASSFp4, Triton, MARLIN do not actually use this config. diff --git a/tests/kernels/quant_utils.py b/tests/kernels/quant_utils.py index 7e14ec95d3..93a21596ca 100644 --- a/tests/kernels/quant_utils.py +++ b/tests/kernels/quant_utils.py @@ -159,7 +159,8 @@ def native_per_token_group_quant_fp8( x, group_size, eps=1e-10, dtype=torch.float8_e4m3fn ): """Function to perform per-token-group quantization on an input tensor - `x` using native torch.""" + `x` using native torch. + """ assert x.shape[-1] % group_size == 0, ( "the last dimension of `x` must be divisible by `group_size`" ) diff --git a/tests/kernels/quantization/nvfp4_utils.py b/tests/kernels/quantization/nvfp4_utils.py index e0120301f7..04f7943ae6 100644 --- a/tests/kernels/quantization/nvfp4_utils.py +++ b/tests/kernels/quantization/nvfp4_utils.py @@ -118,6 +118,7 @@ def dequant_nvfp4_kv_cache( Returns: [..., num_heads, block_size, head_size] float32. + """ data_dim = head_size // 2 scale_dim = head_size // 16 diff --git a/tests/kernels/quantization/test_aiter_hipb_mm_linear_kernel.py b/tests/kernels/quantization/test_aiter_hipb_mm_linear_kernel.py index c855d0e819..70abcbfc8e 100644 --- a/tests/kernels/quantization/test_aiter_hipb_mm_linear_kernel.py +++ b/tests/kernels/quantization/test_aiter_hipb_mm_linear_kernel.py @@ -304,7 +304,8 @@ def test_hipb_mm_kernel_forward_matches_raw_aiter_hipb_mm(enable_hipb_mm_kernel) def test_hipb_mm_kernel_forward_accuracy(enable_hipb_mm_kernel): """Kernel output should match a dequantized fp32 reference within - fp8 per-token / per-channel quantization noise.""" + fp8 per-token / per-channel quantization noise. + """ weight_shape = (512, 4096) # (N, K) num_tokens = 32 _check_bpreshuffle_runtime_support(weight_shape, num_tokens=num_tokens) diff --git a/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py b/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py index b70a1a8543..e4dadb53bd 100644 --- a/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py +++ b/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py @@ -45,6 +45,7 @@ def quantize_weight_block_fp8( Returns: fp8_weight: [N, K] float8_e4m3fn scales: [n_tiles, k_tiles] float32 + """ N, K = weight.shape block_n, block_k = block_size diff --git a/tests/kernels/quantization/test_cutlass_scaled_mm.py b/tests/kernels/quantization/test_cutlass_scaled_mm.py index a522d31714..9e1d6d2cc7 100644 --- a/tests/kernels/quantization/test_cutlass_scaled_mm.py +++ b/tests/kernels/quantization/test_cutlass_scaled_mm.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for cutlass kernels +"""Tests for cutlass kernels. Run `pytest tests/kernels/quantization/test_cutlass_scaled_mm.py`. """ diff --git a/tests/kernels/quantization/test_cutlass_w4a8.py b/tests/kernels/quantization/test_cutlass_w4a8.py index 8cfc993fe8..33bb7bccc7 100644 --- a/tests/kernels/quantization/test_cutlass_w4a8.py +++ b/tests/kernels/quantization/test_cutlass_w4a8.py @@ -300,8 +300,7 @@ def test_w4a8_cuda_graph(): ) @pytest.mark.parametrize("shape", MNK_SHAPES) def test_convert_packed_uint4b8_to_signed_int4_inplace(shape): - """ - The W4A16 checkpoints encode the weights as int4b8 packed to int32. + """The W4A16 checkpoints encode the weights as int4b8 packed to int32. The CUTLASS kernels expect signed int4 packed to int32. This tests checks that the runtime int4b8 -> signed int4 conversion matches the offline conversion step exactly. diff --git a/tests/kernels/quantization/test_cutlass_w4a8_moe.py b/tests/kernels/quantization/test_cutlass_w4a8_moe.py index 5e6c170db6..10a76e553e 100644 --- a/tests/kernels/quantization/test_cutlass_w4a8_moe.py +++ b/tests/kernels/quantization/test_cutlass_w4a8_moe.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for the CUTLASS-based W4A8 grouped GEMM kernel and the full MoE layer. -""" +"""Tests for the CUTLASS-based W4A8 grouped GEMM kernel and the full MoE layer.""" import random from dataclasses import dataclass @@ -37,8 +35,7 @@ def cutlass_quantize( group_size: int | None, zero_points: bool = False, ): - """ - Quantize weights into W4 and compute reference dequantized weights. + """Quantize weights into W4 and compute reference dequantized weights. Encoding/reordering of weights and packing of scales is deferred until after all experts are combined. @@ -66,13 +63,13 @@ def cutlass_quantize( def cutlass_preprocess( w_q_experts: list[torch.Tensor], w_s_experts: list[torch.Tensor] ): - """ - Reorder/encode expert weights and pack scales. + """Reorder/encode expert weights and pack scales. Returns: w_q_packed: Packed/encoded int4 weights for all experts. w_s_packed: Packed fp8 scales for all experts. packed_layout: Layout/stride metadata for grouped GEMM. + """ w_s_packed = ops.cutlass_pack_scale_fp8(torch.stack(w_s_experts)) w_q_packed, packed_layout = ops.cutlass_encode_and_reorder_int4b_grouped( @@ -130,7 +127,6 @@ def make_moe_test_setup( random_zero: bool = False, ) -> MoETestSetup: """Create a full set of tensors for testing cutlass_w4a8_moe_mm.""" - assert K % GROUP_SIZE == 0 # Token counts per expert (multiples of `alignment`). Ms = [alignment * random.randint(1, max_blocks) for _ in range(num_experts)] @@ -276,9 +272,7 @@ def test_cutlass_w4a8_moe_mm_end_to_end(shape, random_zero): class W4A8MoELayer(torch.nn.Module): - """ - Minimal wrapper module to test cuda graphs - """ + """Minimal wrapper module to test cuda graphs.""" def __init__(self, setup: MoETestSetup): super().__init__() diff --git a/tests/kernels/quantization/test_fp8_min_max_helper.py b/tests/kernels/quantization/test_fp8_min_max_helper.py index 8cd68a3fef..4f1beb8576 100644 --- a/tests/kernels/quantization/test_fp8_min_max_helper.py +++ b/tests/kernels/quantization/test_fp8_min_max_helper.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for the get_fp8_min_max() helper function. +"""Unit tests for the get_fp8_min_max() helper function. These tests verify the FP8 min/max value logic for both standard and fnuz (ROCm MI300) dtype handling. diff --git a/tests/kernels/quantization/test_int8_kernel.py b/tests/kernels/quantization/test_int8_kernel.py index 704bb90f84..c01929c5d1 100644 --- a/tests/kernels/quantization/test_int8_kernel.py +++ b/tests/kernels/quantization/test_int8_kernel.py @@ -30,7 +30,8 @@ if current_platform.is_cuda_alike() and not current_platform.has_device_capabili def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16): """Matrix multiplication function that supports per-token input - quantization and per-column weight quantization""" + quantization and per-column weight quantization + """ A = A.to(torch.float32) B = B.to(torch.float32) @@ -53,8 +54,8 @@ def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16): def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, topk, topk_weight, topk_ids): """This function performs fused moe with per-column int8 quantization - using native torch.""" - + using native torch. + """ B, D = a.shape # Perform per-token quantization a_q, a_s = per_token_quant_int8(a) diff --git a/tests/kernels/quantization/test_marlin_tile_padding.py b/tests/kernels/quantization/test_marlin_tile_padding.py index 027ffbf10c..f160125809 100644 --- a/tests/kernels/quantization/test_marlin_tile_padding.py +++ b/tests/kernels/quantization/test_marlin_tile_padding.py @@ -360,7 +360,8 @@ def test_gptq_marlin_padded_round_trip(shape, group_size): def test_fp8_block_marlin_padded_round_trip(shape): """Block-quantized FP8 (e.g. Nemotron NVFP4 checkpoints' FP8 layers): group_size=128 exercises the lcm K-alignment in marlin_padded_nk and the - weight_scale_inv group-wise scale padding.""" + weight_scale_inv group-wise scale padding. + """ size_n, size_k = shape block = 128 dtype = torch.float16 @@ -410,7 +411,8 @@ def test_fp8_block_marlin_padded_round_trip(shape): @pytest.mark.parametrize("shape", [(200, 288), (4640, 512)]) def test_mxfp8_marlin_padded_round_trip(shape): """MXFP8 exercises the e8m0 scale path, where padded 0.0 scales clamp to - 2^-127 instead of zero and must still contribute nothing.""" + 2^-127 instead of zero and must still contribute nothing. + """ size_n, size_k = shape group_size = 32 # The e8m0-scale Marlin kernels are only instantiated for bf16 activations. @@ -457,7 +459,8 @@ def test_mxfp8_marlin_padded_round_trip(shape): @pytest.mark.parametrize("shape", [(200, 512), (4640, 512)]) def test_awq_zp_marlin_padded_round_trip(shape): """AWQ-style uint4 with runtime zero-points, padded the way - MarlinLinearKernel does: padded columns rely on (q=0 - zp=0) * scale=0.""" + MarlinLinearKernel does: padded columns rely on (q=0 - zp=0) * scale=0. + """ size_n, size_k = shape group_size = 128 dtype = torch.float16 @@ -692,7 +695,8 @@ def test_check_moe_marlin_supports_layer_padding(): @pytest.mark.parametrize("shape", [(96, 256, 8), (160, 512, 4)]) def test_fp8_marlin_moe_padded_round_trip(shape, quant): """FP8 weight-only MoE: pad a tile-misaligned intermediate and check the - real prepare + fused_marlin_moe against the dequantized reference.""" + real prepare + fused_marlin_moe against the dequantized reference. + """ from tests.kernels.utils import torch_experts from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe import fused_topk @@ -777,7 +781,8 @@ def test_fp8_marlin_moe_padded_round_trip(shape, quant): @pytest.mark.parametrize("shape", [(96, 256, 8), (160, 512, 4)]) def test_mxfp8_marlin_moe_padded_round_trip(shape): """MXFP8 weight-only MoE round-trip at a tile-misaligned intermediate, with - unit e8m0 scales so the reference is the exact fp8 dequant.""" + unit e8m0 scales so the reference is the exact fp8 dequant. + """ from tests.kernels.utils import torch_experts from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe import fused_topk diff --git a/tests/kernels/quantization/test_mxfp4_kernel_selection.py b/tests/kernels/quantization/test_mxfp4_kernel_selection.py index 28ce615ce6..de9feb7dfa 100644 --- a/tests/kernels/quantization/test_mxfp4_kernel_selection.py +++ b/tests/kernels/quantization/test_mxfp4_kernel_selection.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for MXFP4 linear kernel selection logic (CPU-only) +"""Tests for MXFP4 linear kernel selection logic (CPU-only). Run `pytest tests/kernels/quantization/test_mxfp4_kernel_selection.py`. """ @@ -72,7 +72,8 @@ def test_true_w4a4_kernels_reject_unset_activation(kernel_cls): @pytest.mark.parametrize("kernel_cls", _TRUE_W4A4_KERNELS) def test_true_w4a4_kernels_reject_explicit_non_mxfp4_activation(kernel_cls): """FlashInfer/XPU/Aiter quantize activations to MXFP4 internally, so an - explicit request for a different activation format must be rejected.""" + explicit request for a different activation format must be rejected. + """ config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp6E3M2Dynamic) can_implement, reason = kernel_cls.can_implement(config) assert not can_implement @@ -87,7 +88,8 @@ def test_weight_only_kernels_accept_unquantized_or_mxfp4_activation( """Marlin/Humming never quantize activations, so an unset activation key, or one that already describes MXFP4-shaped data, is tolerated. When an activation key is explicitly set, a warning must be logged noting that it - is ignored, since these kernels are weight-only (A16).""" + is ignored, since these kernels are weight-only (A16). + """ config = MxFp4LinearLayerConfig(activation_quant_key=activation_quant_key) with patch(f"{kernel_cls.__module__}.logger.warning_once") as warning_once: can_implement, reason = kernel_cls.can_implement(config) @@ -116,7 +118,8 @@ def test_weight_only_kernels_reject_non_mxfp4_activation(kernel_cls): ) def test_emulation_kernel_accepts_any_config(activation_quant_key, monkeypatch): """EmulationMxfp4LinearKernel is the universal fallback: it must accept - every supported activation format.""" + every supported activation format. + """ # `EmulationMxfp4LinearKernel.can_implement` gates on `has_quark()`, # which we are not testing here. monkeypatch.setattr( @@ -134,7 +137,8 @@ def test_emulation_kernel_accepts_any_config(activation_quant_key, monkeypatch): def test_emulation_kernel_derives_quant_dequant_func_from_config(monkeypatch): """quant_dequant_func must be derived purely from the config's activation - QuantKey, not set externally.""" + QuantKey, not set externally. + """ # `EmulationMxfp4LinearKernel.can_implement` gates on `has_quark()`, # which we are not testing here. monkeypatch.setattr( @@ -158,7 +162,8 @@ def test_emulation_kernel_derives_quant_dequant_func_from_config(monkeypatch): def test_aiter_kernel_is_supported_requires_native_mx_support(): """AiterMxfp4LinearKernel must not be selected on platforms without - native MX compute, even if AITER itself is importable.""" + native MX compute, even if AITER itself is importable. + """ with patch( "vllm.model_executor.kernels.linear.mxfp4.aiter.current_platform.supports_mx", return_value=False, @@ -194,7 +199,8 @@ class OOTMxFp4LinearKernel(MxFp4LinearKernel): @patch("vllm.model_executor.kernels.linear.current_platform") def test_init_mxfp4_linear_kernel_dispatches_to_registered_kernel(platform_mock): """init_mxfp4_linear_kernel should select a registered kernel that - reports itself as supported, and construct it with a fresh config.""" + reports itself as supported, and construct it with a fresh config. + """ platform_mock._enum = PlatformEnum.OOT register_linear_kernel(OOTMxFp4LinearKernel, PlatformEnum.OOT, "mxfp4") diff --git a/tests/kernels/quantization/test_mxfp4_triton_ep.py b/tests/kernels/quantization/test_mxfp4_triton_ep.py index 93d30f4598..74996d64e3 100644 --- a/tests/kernels/quantization/test_mxfp4_triton_ep.py +++ b/tests/kernels/quantization/test_mxfp4_triton_ep.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests that triton_kernel_moe_forward correctly applies expert_map +"""Tests that triton_kernel_moe_forward correctly applies expert_map remapping when expert parallelism (EP) is enabled. When expert_map is provided, global expert IDs are remapped to local IDs @@ -17,7 +16,8 @@ import torch class TestTritonMoeForwardExpertMap: """Test that triton_kernel_moe_forward applies expert_map remapping - when expert_map is provided (EP active).""" + when expert_map is provided (EP active). + """ def test_expert_map_remap(self): device = "cuda" if torch.cuda.is_available() else "cpu" diff --git a/tests/kernels/quantization/test_mxfp6_kernel_selection.py b/tests/kernels/quantization/test_mxfp6_kernel_selection.py index 2c7b565bea..2c0e94c0f5 100644 --- a/tests/kernels/quantization/test_mxfp6_kernel_selection.py +++ b/tests/kernels/quantization/test_mxfp6_kernel_selection.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for MXFP6 linear kernel selection logic (CPU-only) +"""Tests for MXFP6 linear kernel selection logic (CPU-only). Run `pytest tests/kernels/quantization/test_mxfp6_kernel_selection.py`. """ @@ -41,7 +41,8 @@ def test_can_implement_is_abstract(): def test_emulation_kernel_rejects_non_mxfp6_weights(): """EmulationMxfp6LinearKernel must not implement a non-MXFP6 weight - format.""" + format. + """ config = MxFp6LinearLayerConfig(weight_quant_key=kMxfp4Static) can_implement, reason = EmulationMxfp6LinearKernel.can_implement(config) assert not can_implement @@ -57,7 +58,8 @@ def test_emulation_kernel_accepts_any_supported_config( weight_quant_key, activation_quant_key ): """EmulationMxfp6LinearKernel is the only backend today: it must accept - every supported weight/activation format combination.""" + every supported weight/activation format combination. + """ config = MxFp6LinearLayerConfig( weight_quant_key=weight_quant_key, activation_quant_key=activation_quant_key ) @@ -102,7 +104,8 @@ class OOTMxFp6LinearKernel(MxFp6LinearKernel): def test_init_mxfp6_linear_kernel_dispatches_to_registered_kernel(platform_mock): """init_mxfp6_linear_kernel should select a registered kernel that reports itself as supported/able to implement the given config, and - construct it with that exact config.""" + construct it with that exact config. + """ platform_mock._enum = PlatformEnum.OOT register_linear_kernel(OOTMxFp6LinearKernel, PlatformEnum.OOT, "mxfp6") diff --git a/tests/kernels/quantization/test_nvfp4_emulation.py b/tests/kernels/quantization/test_nvfp4_emulation.py index eb057cebd9..33295b8784 100644 --- a/tests/kernels/quantization/test_nvfp4_emulation.py +++ b/tests/kernels/quantization/test_nvfp4_emulation.py @@ -64,8 +64,7 @@ def loaded_model_files(): class Nvfp4QuantizationEmulationTritonExpertsReference(TritonExperts): - """ - Extension of TritonExperts to support emulated NVFP4 MoE experts. + """Extension of TritonExperts to support emulated NVFP4 MoE experts. It may be used for NVFP4 models when the device does not have native support for this dtype. diff --git a/tests/kernels/quantization/test_nvfp4_quant.py b/tests/kernels/quantization/test_nvfp4_quant.py index 120dc2774a..b8a1cdda46 100644 --- a/tests/kernels/quantization/test_nvfp4_quant.py +++ b/tests/kernels/quantization/test_nvfp4_quant.py @@ -176,8 +176,7 @@ def test_python_util_matches_cpp_allocation( shape: tuple[int, int], is_sf_swizzled_layout: bool, ) -> None: - """ - Verify that the Python utility (create_fp4_output_tensors) allocates + """Verify that the Python utility (create_fp4_output_tensors) allocates tensors with the same shapes and dtypes as the C++ functional variant (scaled_fp4_quant_func). """ diff --git a/tests/kernels/quantization/test_per_token_group_quant.py b/tests/kernels/quantization/test_per_token_group_quant.py index f2326bfd25..c1226efcc8 100644 --- a/tests/kernels/quantization/test_per_token_group_quant.py +++ b/tests/kernels/quantization/test_per_token_group_quant.py @@ -118,8 +118,8 @@ def test_per_token_group_quant_fp8_packed( num_tokens, hidden_dim, group_size, poisoned_scales ): """Test the packed DeepGEMM quantization kernel against the Triton - reference (row-major, UE8M0 scales).""" - + reference (row-major, UE8M0 scales). + """ device = "cuda" torch.manual_seed(42) @@ -206,7 +206,6 @@ def test_per_token_group_quant_fp8_packed_all_zero(): For all-zero input, eps/fp8_max < 1e-10, so the inner fmax clamps back to 1e-10, giving exp2(ceil(log2(1e-10))) = exp2(-33) => UE8M0 byte 0x5E (94). """ - device = "cuda" num_tokens, hidden_dim, group_size = 4, 7168, 128 x = torch.zeros((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) @@ -258,8 +257,8 @@ def test_per_token_group_quant_fp8_packed_all_zero(): def test_per_token_group_quant_fp8_packed_mantissa_rounds_up(): """Inputs whose absmax/max_8bit produces a non-power-of-2 force the mantissa-rounding-up branch (exp_byte += 1). Locks down this behavior - before optimization.""" - + before optimization. + """ device = "cuda" num_tokens, hidden_dim, group_size = 4, 7168, 128 @@ -331,8 +330,8 @@ def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q( ): """When output_q is allocated with shape (tma_aligned_mn, k) instead of (mn, k), the kernel must overwrite the padded mn rows with zeros so - callers can use ``torch.empty`` instead of ``torch.zeros``.""" - + callers can use ``torch.empty`` instead of ``torch.zeros``. + """ device = "cuda" group_size = 128 torch.manual_seed(42) @@ -396,7 +395,6 @@ def test_per_token_group_quant_fp8_packed_large_mn(): This is a differential test that compares fp8 output against Triton output reference when token size sits just above the gridDim.y 2^16 - 1 limit. """ - device = "cuda" group_size = 128 # hidden 2048 -> 2048/128 = 16 groups per row -> kx=16, ry=1: one grid row per mn diff --git a/tests/kernels/quantization/test_quantized_embedding.py b/tests/kernels/quantization/test_quantized_embedding.py index 0e4af0a0c1..35c0e74713 100644 --- a/tests/kernels/quantization/test_quantized_embedding.py +++ b/tests/kernels/quantization/test_quantized_embedding.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for the Triton dequant-gather kernel used by -``CompressedTensorsEmbeddingWNA16Int`` (quantized embedding lookup).""" +``CompressedTensorsEmbeddingWNA16Int`` (quantized embedding lookup). +""" import pytest import torch diff --git a/tests/kernels/quantization/test_rocm_aiter_grouped_quant.py b/tests/kernels/quantization/test_rocm_aiter_grouped_quant.py index c7f0f1eda3..2a5a8081a2 100644 --- a/tests/kernels/quantization/test_rocm_aiter_grouped_quant.py +++ b/tests/kernels/quantization/test_rocm_aiter_grouped_quant.py @@ -29,7 +29,8 @@ pytestmark = pytest.mark.skipif( def test_rocm_aiter_group_fp8_quant_fake_implementation(): """Test that the fake implementation is correctly - defined for torch.ops.vllm.rocm_aiter_group_fp8_quant.""" + defined for torch.ops.vllm.rocm_aiter_group_fp8_quant. + """ # Create test tensors M = 128 N = 4096 @@ -49,7 +50,8 @@ def test_rocm_aiter_group_fp8_quant_fake_implementation(): def test_rocm_aiter_group_fp8_quant_torch_compile_with_cudagraph(): """Test that rocm_aiter_ops.group_fp8_quant with group size 128 can be used with - torch.compile in cudagraph mode.""" + torch.compile in cudagraph mode. + """ # Create test tensors M = 128 N = 4096 diff --git a/tests/kernels/quantization/test_rocm_mxfp4.py b/tests/kernels/quantization/test_rocm_mxfp4.py index 43b5daa45a..7002165af0 100644 --- a/tests/kernels/quantization/test_rocm_mxfp4.py +++ b/tests/kernels/quantization/test_rocm_mxfp4.py @@ -209,7 +209,8 @@ def test_rocm_aiter_fp4_enablement_follows_env_and_arch( @pytest.mark.parametrize("shape", LLAMA_SHAPES) def test_vllm_quant_dequant_mxfp4_matches_reference_on_large_shapes(shape, dtype): """The public vLLM MXFP4 QDQ helper should stay exact against the torch - reference even on larger Llama-like shapes.""" + reference even on larger Llama-like shapes. + """ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( quant_dequant_mxfp4, ) @@ -504,7 +505,7 @@ def test_aiter_fp4_gemm_a4w4_determinism(): ], ) def test_aiter_hardware_fp4_dynamic_quant_format(shape): - """aiter hardware FP4 dynamic quant produces correct output format. + """Aiter hardware FP4 dynamic quant produces correct output format. Tests gfx950 hardware-accelerated FP4 quantization (OCP MXFP4 E2M1). Parity with B200 scaled_fp4_quant: block_size=32, packed uint8 output. diff --git a/tests/kernels/quantization/test_scaled_mm_kernel_selection.py b/tests/kernels/quantization/test_scaled_mm_kernel_selection.py index 8275714744..d682ed05b0 100644 --- a/tests/kernels/quantization/test_scaled_mm_kernel_selection.py +++ b/tests/kernels/quantization/test_scaled_mm_kernel_selection.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for ScaledMM kernel selection logic (CPU-only) +"""Tests for ScaledMM kernel selection logic (CPU-only). Run `pytest tests/kernels/quantization/test_scaled_mm_kernel_selection.py`. """ diff --git a/tests/kernels/quantization/test_triton_scaled_mm.py b/tests/kernels/quantization/test_triton_scaled_mm.py index e829c63a84..55d1f76b2d 100644 --- a/tests/kernels/quantization/test_triton_scaled_mm.py +++ b/tests/kernels/quantization/test_triton_scaled_mm.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the triton_scaled_mm kernel +"""Tests for the triton_scaled_mm kernel. Run `pytest tests/kernels/quantization/test_triton_scaled_mm.py`. """ diff --git a/tests/kernels/test_apply_repetition_penalties.py b/tests/kernels/test_apply_repetition_penalties.py index 8270cf885f..c1551e2e82 100644 --- a/tests/kernels/test_apply_repetition_penalties.py +++ b/tests/kernels/test_apply_repetition_penalties.py @@ -35,8 +35,7 @@ def test_apply_repetition_penalties( dtype: torch.dtype, seed: int, ) -> None: - """ - Test the apply_repetition_penalties custom op + """Test the apply_repetition_penalties custom op against a reference implementation. """ set_random_seed(seed) @@ -86,8 +85,7 @@ def test_apply_repetition_penalties( ) @torch.inference_mode() def test_apply_repetition_penalties_zero_seqs() -> None: - """ - Test the apply_repetition_penalties custom op with num_seqs=0 + """Test the apply_repetition_penalties custom op with num_seqs=0 against a reference implementation. """ num_seqs = 0 diff --git a/tests/kernels/test_awq_int4_to_int8.py b/tests/kernels/test_awq_int4_to_int8.py index f94214785d..dc7124e9ae 100644 --- a/tests/kernels/test_awq_int4_to_int8.py +++ b/tests/kernels/test_awq_int4_to_int8.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for AWQ INT4 W4A8 GEMM pipeline (SGLang kernel migration). +"""Unit tests for AWQ INT4 W4A8 GEMM pipeline (SGLang kernel migration). Part 1: Weight packing tests - convert_weight_packed_scale_zp correctness @@ -44,6 +43,7 @@ def make_awq_checkpoint_data(K, N, group_size, seed=42): float_ref: [K, N] float32, reference dequantized weights weight_int4_orig: [K, N] int32, original int4 values (0-15) zeros_int4_orig: [num_groups, N] int32, original zero points (0-15) + """ rng = np.random.RandomState(seed) num_groups = K // group_size diff --git a/tests/kernels/test_cache_kernels.py b/tests/kernels/test_cache_kernels.py index eff4508b0d..a3cb91242a 100644 --- a/tests/kernels/test_cache_kernels.py +++ b/tests/kernels/test_cache_kernels.py @@ -15,12 +15,10 @@ except ImportError: @pytest.mark.skipif(torch.accelerator.device_count() < 1, reason="Need CUDA device") def test_gather_cache_oob(): - """ - Tests for OOB read in gather_and_maybe_dequant_cache (Issue #27909). + """Tests for OOB read in gather_and_maybe_dequant_cache (Issue #27909). This test constructs a boundary case identified in the issue where seq_starts causes the block_table offset to read out of bounds. """ - block_size = 64 # The kernel only supports the MLA entry sizes. entry_size = 576 diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 73c2027a78..3dc658250f 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Round-trip tests for compressor → FP8 quant + KV cache insert → gather + dequant. +"""Round-trip tests for compressor → FP8 quant + KV cache insert → gather + dequant. These tests cover: A) DeepseekV4 Attention: head_dim=512 (448 FP8 nope + 64 bf16 rope), quant_block=64 @@ -671,7 +670,8 @@ def test_v41_rope_insert_plain_row(compress_ratio: int, store_fp8: bool): def test_v41_compressor_metadata_maps_tokens_to_their_ring(): """The ring group's generic slot mapping is disabled (all PAD), so the builder must map every real token to ``ring_block * capacity + pos % - capacity`` and keep padding tokens at PAD.""" + capacity`` and keep padding tokens at PAD. + """ from unittest.mock import MagicMock from vllm.models.deepseek_v41.compressor import CompressorMetadataBuilder @@ -1096,8 +1096,8 @@ def test_get_c128_boundary(starts, query_start_loc, expected): @pytest.mark.parametrize("block_size", [16, 64]) def test_deepseek_v4_attention_quant_cache_roundtrip(num_tokens: int, block_size: int): """compressed_kv → quantize_and_insert_k_cache → dequantize_and_gather_k_cache - → compare against original.""" - + → compare against original. + """ HEAD_DIM = 512 NOPE_DIM = 448 HEAD_BYTES = 584 # 448 fp8 + 128 bf16 + 8 uint8 scale @@ -1308,9 +1308,9 @@ def test_dequantize_and_gather_k_cache( @pytest.mark.parametrize("num_tokens", [1, 4, 8, 17]) @pytest.mark.parametrize("block_size", [16, 64]) def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int): - """k → indexer_k_quant_and_cache → cp_gather_indexer_k_quant_cache - → manual dequant → compare against original.""" - + """K → indexer_k_quant_and_cache → cp_gather_indexer_k_quant_cache + → manual dequant → compare against original. + """ HEAD_DIM = 128 QUANT_BLOCK_SIZE = 128 # cache_stride = head_dim + (head_dim * 4 / quant_block_size) = 128 + 4 = 132 @@ -1374,7 +1374,6 @@ def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int): def test_indexer_gather_accepts_upper_bound_output(): """Gather only exact cu_seq_lens even when dst is over-allocated.""" - head_dim = 128 quant_block_size = 128 cache_stride = head_dim + head_dim * 4 // quant_block_size @@ -1474,7 +1473,6 @@ def test_indexer_gather_accepts_upper_bound_output(): def test_deepseek_v4_quant_magnitude_range(): """Test that quantization handles a range of magnitudes correctly.""" - HEAD_DIM = 512 NOPE_DIM = 448 HEAD_BYTES = 584 diff --git a/tests/kernels/test_cp_gather_fp8.py b/tests/kernels/test_cp_gather_fp8.py index fe539abb4f..d9157c4110 100644 --- a/tests/kernels/test_cp_gather_fp8.py +++ b/tests/kernels/test_cp_gather_fp8.py @@ -29,6 +29,7 @@ def _build_test_case(seq_lens, block_size, seed=42): Returns: Tuple of (cache, block_table, workspace_starts_t, num_reqs, total_tokens, expected_output). + """ torch.manual_seed(seed) @@ -431,7 +432,8 @@ def test_cp_gather_fp8_with_sequence_starts(gather_seq_lens, seq_starts): ) def test_cp_gather_fp8_large_seqlens(seq_lens, block_size): """Correctness test with large sequence lengths matching benchmark - scenarios (8K-128K prefill).""" + scenarios (8K-128K prefill). + """ ( cache, block_table, diff --git a/tests/kernels/test_deepseek_v4_cpu_kernels.py b/tests/kernels/test_deepseek_v4_cpu_kernels.py index 81c03a0bca..9c3fddb6c1 100644 --- a/tests/kernels/test_deepseek_v4_cpu_kernels.py +++ b/tests/kernels/test_deepseek_v4_cpu_kernels.py @@ -85,7 +85,8 @@ def _decode_token_regions( cache: torch.Tensor, slot_mapping: torch.Tensor, block_size: int ): """Slice out, per valid token, the raw NoPE-fp8 / RoPE-bf16 / scale byte - regions this layout defines.""" + regions this layout defines. + """ valid = slot_mapping >= 0 idx = valid.nonzero(as_tuple=True)[0] slots = slot_mapping[idx] @@ -110,7 +111,8 @@ def _assert_cache_parity(cache_fused, cache_ref, slot_mapping, block_size): """NoPE fp8 + the UE8M0 scale byte are a deterministic quant of an un-rotated bf16 value, so they match tightly; the RoPE region is rotated in fp32 before its bf16 store and eager/kernel can round a tie to - opposite sides, so allow <=1 ULP there.""" + opposite sides, so allow <=1 ULP there. + """ nope_ref, rope_ref, scale_ref = _decode_token_regions( cache_ref, slot_mapping, block_size ) @@ -190,7 +192,8 @@ def _dup_padded_block_cache( num_blocks: int, block_size: int, width: int, pad_elems: int ) -> tuple[torch.Tensor, torch.Tensor]: """Two independently-stored padded caches with identical values, for ops - that mutate the cache in place.""" + that mutate the cache in place. + """ page_stride = block_size * width + pad_elems template = torch.randn(num_blocks, page_stride, dtype=torch.float32) @@ -207,7 +210,8 @@ def _run_save_partial_states_eager( ): """Reference: state_cache[block_idx, pos_in_block, :W] = kv; state_cache[..., W:] = score + ape[position % compress_ratio]; skip - slot < 0. Mutates state_cache in place, matching the op's contract.""" + slot < 0. Mutates state_cache in place, matching the op's contract. + """ state_width = kv.shape[-1] block_size = state_cache.shape[1] valid = slot_mapping >= 0 @@ -267,7 +271,8 @@ def _run_compress_eager( ): """Reference for the per-channel-softmax-pooling + RMSNorm + NoPE-UE8M0 quant + GPT-J RoPE math writing the main (head_dim=512) fp8_ds_mla - layout.""" + layout. + """ num_tokens, window = gather_slots.shape state_width = state_cache.shape[-1] // 2 state_block_size = state_cache.shape[1] @@ -444,7 +449,8 @@ def test_compress_norm_rope_store_cpu_matches_eager( def test_compress_norm_rope_store_cpu_non_contiguous_kv_cache_row_stride(): """The main (output) kv_cache is a separate physical buffer from the - state cache and can independently be a strided view.""" + state cache and can independently be a strided view. + """ eps = 1e-6 compress_ratio, overlap = 128, False num_tokens = 5 @@ -520,7 +526,8 @@ def _run_compress_indexer_eager( """Same softmax-pool + RMSNorm + GPT-J RoPE prologue as the head=512 path, but the indexer's cache quantizes all 128 post-RoPE dims as one FP8 block with a single raw fp32 scale per token (matching - _indexer_k_quant_and_cache_cpu's layout), not 7 separate UE8M0 blocks.""" + _indexer_k_quant_and_cache_cpu's layout), not 7 separate UE8M0 blocks. + """ num_tokens, window = gather_slots.shape state_width = state_cache.shape[-1] // 2 state_block_size = state_cache.shape[1] @@ -648,7 +655,8 @@ def _pack_flash_mla_cache(kv_f32: torch.Tensor, block_size: int): """Quantize kv_f32 ([N, HEAD_DIM]) into the fp8_ds_mla byte layout and return (cache, dequant_ref) -- the exact post-quantization value the kernel's dequant should reproduce. Slot id == row index, packed - contiguously from block 0.""" + contiguously from block 0. + """ n = kv_f32.shape[0] num_blocks = max(1, -(-n // block_size)) cache = _make_cache(num_blocks, block_size) @@ -706,7 +714,8 @@ def _flash_mla_eager( """MQA attention: gather the (already-dequantized) window+compressed KV rows referenced by *_slots (-1 = invalid/masked), append a learned per-head sink logit with zero value contribution, softmax, weighted - sum.""" + sum. + """ num_tokens, num_heads, _ = q.shape q_f32 = q.float() @@ -805,7 +814,8 @@ def test_flash_mla_cpu_kernel_matches_eager( def test_flash_mla_cpu_kernel_all_invalid_falls_back_to_sink(): """When every KV slot is masked out, softmax degenerates to the sink - logit alone and the output must be all-zero (no valid value rows).""" + logit alone and the output must be all-zero (no valid value rows). + """ torch.manual_seed(2) dtype = torch.bfloat16 num_tokens, num_heads, num_window = 3, 4, 5 @@ -845,7 +855,8 @@ def _qnorm_kv_insert_eager( q, kv, positions, cache, slot_mapping, cos_sin_cache, q_head_padded, eps, bs ): """Reference: per-head weight-free RMSNorm + GPT-J RoPE on Q; GPT-J RoPE - + UE8M0 FP8 quant + paged insert on KV.""" + + UE8M0 FP8 quant + paged insert on KV. + """ num_quant_blocks = NOPE_DIM // QUANT_BLOCK num_tokens_full, num_heads_q, _ = q.shape @@ -971,7 +982,8 @@ def test_qnorm_rope_kv_insert_cpu_kernel_matches_eager( @pytest.mark.parametrize("pad", [1, 5]) def test_qnorm_rope_kv_insert_cpu_kernel_dp_padding(pad: int): """slot_mapping shorter than q/kv rows: the KV branch must only touch - the first num_tokens rows while Q-norm+RoPE still runs on all rows.""" + the first num_tokens rows while Q-norm+RoPE still runs on all rows. + """ torch.manual_seed(1) dtype = torch.bfloat16 eps = 1e-6 @@ -1032,7 +1044,8 @@ def _inverse_gptj_rope_o_proj_eager( rope_dim: int, ) -> torch.Tensor: """NoPE dims pass through as fp32; the RoPE segment is de-rotated by - R(-pos) in GPT-J (interleaved even/odd) convention.""" + R(-pos) in GPT-J (interleaved even/odd) convention. + """ nope_dim = o.shape[-1] - rope_dim half_dim = rope_dim // 2 @@ -1098,7 +1111,8 @@ def _write_page( ) -> None: """Write k_fp8 ([block_size, INDEXER_HEAD_DIM]) and scale ([block_size] fp32) into one page row: K-region then scale-region, matching - compress_norm_rope_store_indexer_cpu's layout.""" + compress_norm_rope_store_indexer_cpu's layout. + """ k_region = page_row[: block_size * INDEXER_HEAD_DIM] k_region.copy_(k_fp8.view(torch.uint8).flatten()) scale_region = page_row[block_size * INDEXER_HEAD_DIM :].view(torch.float32) @@ -1134,7 +1148,8 @@ def _build_paged_mqa_logits_batch( ): """Fill pages (a pre-sized [n_pages, buf_width] uint8 tensor, possibly non-contiguous in dim 0) with deterministic-per-seed random K/scale - data.""" + data. + """ torch.manual_seed(seed) batch_size = len(seq_lens) max_seq_len = max(seq_lens) @@ -1200,7 +1215,8 @@ def test_fp8_paged_mqa_logits_cpu_non_contiguous_kv_cache_row_stride(): """The paged K-cache in production is a per-layer view into a shared multi-layer allocation, so its page stride generally exceeds the page byte width -- check the kernel reads a padded/strided buffer identically - to an equivalent tightly-packed one.""" + to an equivalent tightly-packed one. + """ block_size = 64 num_heads = 4 seq_lens = [1, block_size + 7, 3 * block_size] @@ -1326,7 +1342,8 @@ def test_fused_indexer_q_rope_quant_dispatches_to_cpu_kernel(): CPU op on a CPU tensor, not silently fall through to the Triton kernel (which would also "work" under triton-cpu, masking a dispatch regression). The dispatcher does its CPU import inside the CPU branch, - so the patch target is the source module.""" + so the patch target is the source module. + """ torch.manual_seed(0) num_tokens = 1 @@ -1393,7 +1410,8 @@ def _hc_head_torch( ) -> torch.Tensor: """Eager reference for hc_head_fused_cpu: round-trips the RMSNorm output through bf16 before the linear, matching the kernel's bf16 - arithmetic.""" + arithmetic. + """ hc_mult, hidden_size = hidden_states.shape[-2:] outer_shape = hidden_states.shape[:-2] hs_flat = hidden_states.reshape(-1, hc_mult, hidden_size) @@ -1475,7 +1493,8 @@ def test_mhc_post_cpu(num_tokens: int, hidden_size: int): @pytest.mark.parametrize("hidden_size", [4096, 7168]) def test_hc_head_cpu(num_tokens: int, hidden_size: int): """hc_head_fused_cpu vs. the _hc_head_torch eager reference it - replaces.""" + replaces. + """ torch.set_default_device(DEVICE) set_random_seed(0) hc_mult = 4 diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py index c2a7607cba..8deba1241a 100644 --- a/tests/kernels/test_engram.py +++ b/tests/kernels/test_engram.py @@ -383,7 +383,8 @@ def _hash_ids(requests, tables_meta, block_table, cache=None, capture=False): """Run the op on a batch of (token_ids, start, window) requests; `window` is the runner's [depth] lookback for that request (None = all unknown). Token IDs 14 and 22 stand in for image tokens that break n-grams. - Returns one [len(token_ids), ...] hash tensor per request.""" + Returns one [len(token_ids), ...] hash tensor per request. + """ token_map, multipliers, primes, offsets, block_size = tables_meta depth = multipliers.shape[1] - 1 @@ -442,7 +443,8 @@ def test_lookback_window_reproduces_single_instance(runner, capture): V2 supplies every lookback from its device token history and needs no slot cache. V1 supplies prompt positions only and reads generated positions from the slot cache it fills itself. Without any window the - first decode token hashes stale slots (negative control).""" + first decode token hashes stale slots (negative control). + """ block_size = 4 block_table = torch.tensor( [[3, 1, 5, 0], [2, 4, 6, 7]], dtype=torch.int32, device="cuda" @@ -708,7 +710,8 @@ def test_engram_lookup_reuses_jit_across_token_shapes(): def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens): """The fused gather must be bit-exact with the torch dequant path it replaces, from HBM and from pinned host memory alike, and must contribute - zeros for rows another TP rank owns.""" + zeros for rows another TP rank owns. + """ layer = _make_embedding(cpu_offload) cols, rows = 24, layer.part_num_embeddings ids = torch.randint(0, rows, (num_tokens, cols), dtype=torch.int32, device="cuda") diff --git a/tests/kernels/test_fla_layernorm_guard.py b/tests/kernels/test_fla_layernorm_guard.py index 492c099bf1..fad8d90805 100644 --- a/tests/kernels/test_fla_layernorm_guard.py +++ b/tests/kernels/test_fla_layernorm_guard.py @@ -336,7 +336,8 @@ def test_layer_norm_rows_per_block( @torch.inference_mode() def test_strided_input(dtype: torch.dtype) -> None: """Test that the kernel handles non-contiguous (strided) - inputs correctly.""" + inputs correctly. + """ set_random_seed(42) device = torch.device(DEVICE) num_tokens = 128 diff --git a/tests/kernels/test_flex_attention.py b/tests/kernels/test_flex_attention.py index 12d00ec405..6862f8ffb1 100644 --- a/tests/kernels/test_flex_attention.py +++ b/tests/kernels/test_flex_attention.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Integration tests for FlexAttention backend vs default backend""" +"""Integration tests for FlexAttention backend vs default backend.""" import pytest import torch diff --git a/tests/kernels/test_fp32_router_gemm.py b/tests/kernels/test_fp32_router_gemm.py index d330abba53..b6c2a53ed8 100644 --- a/tests/kernels/test_fp32_router_gemm.py +++ b/tests/kernels/test_fp32_router_gemm.py @@ -97,7 +97,8 @@ def test_topk_routing_consistency(num_tokens: int, hidden_dim: int, num_experts: """The gate feeds top-k expert selection: the kernel's top-8 must match an fp64 reference's top-8 per token (ties tolerated). This is the business-level correctness of the router — numeric error only matters - if it flips the argsort.""" + if it flips the argsort. + """ _requires_sm90() top_k = 8 device = torch.device("cuda") diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index 7f74cde05d..6691de7ebf 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -854,7 +854,8 @@ def test_fused_q_no_indexer(num_tokens: int): @pytest.mark.parametrize("has_indexer", [True, False]) def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool): """bf16-query path (FlashMLA sparse, SM90/SM100): only the RoPE'd q_pe is - produced (bf16, unquantized); ql_nope is consumed directly by the caller.""" + produced (bf16, unquantized); ql_nope is consumed directly by the caller. + """ torch.manual_seed(6) dev = "cuda" max_pos = 8192 diff --git a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py index a0dcfc5bb7..3e95913a01 100644 --- a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Standalone unit test for the horizontally-fused DeepseekV4-MLA kernel: +"""Standalone unit test for the horizontally-fused DeepseekV4-MLA kernel: fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert - Q side: optional per-head RMSNorm + GPT-J RoPE on last 64 dims @@ -177,7 +176,8 @@ def _call_fused( def _as_stored_fp8(t: torch.Tensor) -> torch.Tensor: """Reinterpret a float8_e4m3fn-typed kernel output under the real (FNUZ on - gfx942) encoding the kernel actually wrote, without touching the bytes.""" + gfx942) encoding the kernel actually wrote, without touching the bytes. + """ return t.contiguous().view(torch.uint8).view(FP8_STORE_DTYPE) @@ -713,7 +713,8 @@ def test_q_interleaved_is_orthogonal_to_norm_and_rope( @pytest.mark.parametrize("block_size", [16, 64]) def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int): """slot_mapping.size(0) < q.size(0): the kernel must skip padded - tokens in the KV branch while still processing Q on all rows.""" + tokens in the KV branch while still processing Q on all rows. + """ torch.manual_seed(3) device = "cuda" dtype = torch.bfloat16 diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index 7464c90769..20ac8f19f1 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -45,6 +45,7 @@ def quantize_to_mxfp4( Returns: packed: [..., head_dim//2] uint8 2 E2M1 nibbles/byte, low nibble = even index scales: [..., head_dim//32] uint8 1 ue8m0 byte + """ MXFP4_BLOCK_SIZE = 32 orig_shape = x.shape @@ -573,7 +574,8 @@ def test_indexer_k_store_roundtrips_through_rocm_gather(block_size, compress_rat def test_fused_indexer_q_rope_quant_writes_bf16_weights(use_cutedsl): """The MXFP4 path can emit the per-head weights in bf16 (what DeepGEMM's sparse MQA-logits kernels take) instead of fp32; the values are the fp32 - result rounded once, so the scoring kernel needs no cast.""" + result rounded once, so the scoring kernel needs no cast. + """ if use_cutedsl and not has_cutedsl(): pytest.skip("cutedsl (cutlass) not installed") diff --git a/tests/kernels/test_fused_inv_rope_fp8_quant.py b/tests/kernels/test_fused_inv_rope_fp8_quant.py index d27144a38e..97ea8f7e34 100644 --- a/tests/kernels/test_fused_inv_rope_fp8_quant.py +++ b/tests/kernels/test_fused_inv_rope_fp8_quant.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for the fused inverse RoPE + block-scaled FP8 quantization kernel. +"""Unit tests for the fused inverse RoPE + block-scaled FP8 quantization kernel. Tests compare the fused kernel against a reference implementation built from the existing separate operations (inverse RoPE via rotate_neox + FP8 quant @@ -132,6 +131,7 @@ def reference_inv_rope( Returns: o with inverse RoPE applied on the rope portion (bf16). + """ assert cos_sin_cache.dtype == torch.float32 cos_sin = cos_sin_cache[positions] # [T, rope_dim] fp32 @@ -164,6 +164,7 @@ def _ref_ue8m0_quant_block(x_f32: torch.Tensor) -> tuple[torch.Tensor, torch.Ten Returns: x_fp8: same shape, float8_e4m3fn scales: [...] float32, one scale per block + """ absmax = x_f32.abs().amax(dim=-1, keepdim=True).clamp(min=EPS) scale_raw = absmax * (1.0 / FP8_MAX) @@ -192,6 +193,7 @@ def reference_inv_rope_fp8_quant( Returns: o_fp8: [T, G, D] FP8 with strides (D, T*D, 1) o_scale: [T, G, S] FP32 with strides (S, T*S, 1) + """ assert cos_sin_cache.dtype == torch.float32 T, _H, head_dim = o.shape @@ -402,7 +404,6 @@ def test_output_strides(num_tokens, num_heads, n_groups): - FP8: logical [T, G, D] backed by contiguous [G, T, D]. - Scale: MN-major TMA-aligned (column-major: T-stride=1). """ - heads_per_group = num_heads // n_groups max_pos = 4096 device = "cuda" @@ -445,7 +446,8 @@ def test_output_strides(num_tokens, num_heads, n_groups): def test_per_group_contiguity(num_tokens): """FP8 per-group slices must be contiguous. Scale per-group slices are column-major (T-stride=1) — not row-major contiguous, which is - correct for TMA loads.""" + correct for TMA loads. + """ num_heads, n_groups = 64, 8 heads_per_group = num_heads // n_groups max_pos = 4096 @@ -511,7 +513,8 @@ def test_scales_are_power_of_two(): def test_nope_dims_unchanged(): """Nope dimensions (first 448 per head) should only be quantized, not rotated. Verify by dequantizing and comparing against - quantize-only reference (no RoPE).""" + quantize-only reference (no RoPE). + """ num_tokens, num_heads, n_groups = 16, 64, 8 heads_per_group = num_heads // n_groups max_pos = 4096 @@ -902,8 +905,8 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups): @torch.inference_mode() def test_with_real_deepseek_v4_rope(num_tokens, default_vllm_config): """Test with real DeepseekV4ScalingRotaryEmbedding (GPT-J style, - mscale=0, YaRN scaling) matching the production config.""" - + mscale=0, YaRN scaling) matching the production config. + """ num_heads = 64 n_groups = 8 heads_per_group = num_heads // n_groups diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 5fab150d9f..219e42c742 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -703,7 +703,7 @@ def mhc_pre_ref( hc_post_mult_value: float, sinkhorn_repeat: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """mHC pre reference kernel from tilelang repo: https://github.com/tile-ai/tilelang/blob/d135bd1cd2d2eee74fbb41dd0a0831a427194c86/examples/deepseek_mhc/example_mhc_pre.py#L303""" + """MHC pre reference kernel from tilelang repo: https://github.com/tile-ai/tilelang/blob/d135bd1cd2d2eee74fbb41dd0a0831a427194c86/examples/deepseek_mhc/example_mhc_pre.py#L303.""" hc_mult = residual.shape[-2] residual_flat = residual.flatten(-2, -1).float() @@ -742,7 +742,7 @@ def mhc_post_ref( post_layer_mix: torch.Tensor, comb_res_mix: torch.Tensor, ) -> torch.Tensor: - """mHC post reference kernel from tilelang repo: https://github.com/tile-ai/tilelang/blob/d135bd1cd2d2eee74fbb41dd0a0831a427194c86/examples/deepseek_mhc/example_mhc_post.py#L68""" + """MHC post reference kernel from tilelang repo: https://github.com/tile-ai/tilelang/blob/d135bd1cd2d2eee74fbb41dd0a0831a427194c86/examples/deepseek_mhc/example_mhc_post.py#L68.""" term2 = torch.bmm(comb_res_mix.mT, residual.float()) return (x.float().unsqueeze(-2) * post_layer_mix + term2).bfloat16() @@ -1372,7 +1372,8 @@ def _patch_first_rank_pp_group(monkeypatch): def test_deepseek_v4_mhc_broadcast_finalize_sums_hc_streams(monkeypatch): """First finalize (at the end of load_weights) allocates - hc_attn_fn_broadcast as hc_attn_fn summed over hc streams.""" + hc_attn_fn_broadcast as hc_attn_fn summed over hc streams. + """ _patch_first_rank_pp_group(monkeypatch) layer = _make_mhc_decoder_layer(hc_mult=2, hidden_size=8) model = SimpleNamespace(start_layer=0, end_layer=1, layers=[layer]) @@ -1387,7 +1388,8 @@ def test_deepseek_v4_mhc_broadcast_finalize_sums_hc_streams(monkeypatch): def test_deepseek_v4_mhc_broadcast_refit_refreshes_in_place(monkeypatch): """Re-finalizing after a weight refit must copy into the existing broadcast tensor so its address stays stable for captured CUDA graphs, - while picking up the new hc_attn_fn values.""" + while picking up the new hc_attn_fn values. + """ _patch_first_rank_pp_group(monkeypatch) layer = _make_mhc_decoder_layer(hc_mult=2, hidden_size=8) model = SimpleNamespace(start_layer=0, end_layer=1, layers=[layer]) diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py index 997c619809..45fafb0433 100644 --- a/tests/kernels/test_minimax_m3_amd_ops.py +++ b/tests/kernels/test_minimax_m3_amd_ops.py @@ -403,7 +403,8 @@ def test_mxfp8_linear_emulation_bf16_at_load( ): """EmulationMxfp8LinearKernel load-time BF16 dequant (default) and the ``VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD=0`` per-step fallback must produce the - same result; the dtype-match (BF16/FP16 activations) must also hold.""" + same result; the dtype-match (BF16/FP16 activations) must also hold. + """ from vllm.model_executor.kernels.linear.mxfp8.emulation import ( EmulationMxfp8LinearKernel, ) @@ -517,7 +518,8 @@ def test_mxfp8_rocm_native_unaligned_k_dequantizes_at_load(shape): # produced EP garbage). def _capture_expert_mask(expert_mask, *, global_num_experts): """Drive the real ``AiterMxfp8Experts.apply`` mask branch and capture the - ``expert_mask`` it forwards to ``rocm_aiter_ops.fused_moe``.""" + ``expert_mask`` it forwards to ``rocm_aiter_ops.fused_moe``. + """ from types import SimpleNamespace from unittest import mock @@ -568,7 +570,8 @@ def _capture_expert_mask(expert_mask, *, global_num_experts): @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") def test_aiter_mxfp8_apply_forwards_expert_mask_unchanged(): """AiterMxfp8Experts.apply forwards the precomputed 0/1 mask to aiter as-is - (no re-derivation that would collapse an already-0/1 mask to all-ones).""" + (no re-derivation that would collapse an already-0/1 mask to all-ones). + """ # 0/1 local-expert mask over global ids + trailing sentinel (rank owns 0..3). ep_mask = torch.tensor( [1, 1, 1, 1, 0, 0, 0, 0, 0], dtype=torch.int32, device=DEVICE @@ -582,7 +585,8 @@ def test_aiter_mxfp8_apply_forwards_expert_mask_unchanged(): def test_routed_experts_expert_map_delegates_to_kernel(): """RoutedExperts.expert_map returns the 0/1 mask only for kernels that set ``consumes_expert_mask`` (AITER), and the canonical -1 map otherwise -- keyed - on the resolved kernel, not the global aiter switch.""" + on the resolved kernel, not the global aiter switch. + """ from types import SimpleNamespace from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts diff --git a/tests/kernels/test_shuffle_rows.py b/tests/kernels/test_shuffle_rows.py index 08d5883a1e..35c2f26442 100644 --- a/tests/kernels/test_shuffle_rows.py +++ b/tests/kernels/test_shuffle_rows.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the shuffle_rows function +"""Tests for the shuffle_rows function. Run `pytest tests/kernels/test_shuffle_rows.py`. """ @@ -17,7 +17,8 @@ from vllm.platforms import current_platform @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) def test_shuffle_rows_basic(num_tokens: int, hidden_size: int, dtype: torch.dtype): """Test basic functionality of shuffle_rows with various tensor sizes and - dtypes.""" + dtypes. + """ if not current_platform.is_cuda(): pytest.skip("shuffle_rows requires CUDA") @@ -72,7 +73,8 @@ def test_shuffle_rows_permutation( @pytest.mark.parametrize("hidden_size", [256, 512]) def test_shuffle_rows_expansion(num_tokens: int, hidden_size: int): """Test shuffle_rows with expansion (more output tokens than input - tokens).""" + tokens). + """ if not current_platform.is_cuda(): pytest.skip("shuffle_rows requires CUDA") diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index 2895394706..f249adeaff 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -157,8 +157,7 @@ def compare_top_k_results( top_k: int, tolerance: float = 1e-5, ) -> bool: - """ - Compare results from CUDA top_k_per_row with torch.topk. + """Compare results from CUDA top_k_per_row with torch.topk. Both results should be sorted and contain the same top-k elements. """ num_rows = cuda_indices.shape[0] @@ -213,8 +212,7 @@ def validate_topk_against_reference( top_k: int, kernel_name: str, ) -> None: - """ - Validate CUDA top-k results against PyTorch reference implementation. + """Validate CUDA top-k results against PyTorch reference implementation. Args: logits: Input logits tensor @@ -223,6 +221,7 @@ def validate_topk_against_reference( row_ends: Row end positions top_k: Number of top elements to select kernel_name: Name of the kernel being tested (for error messages) + """ num_rows = cuda_indices.shape[0] torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") @@ -248,9 +247,7 @@ def test_top_k_per_row( top_k: int, clean_logits: bool, ) -> None: - """ - Test top_k_per_row. - """ + """Test top_k_per_row.""" set_random_seed(0) torch.set_default_device("cuda:0") @@ -298,9 +295,7 @@ def _run_top_k_per_row_decode_test( clean_logits: bool, data_generation: str, ) -> None: - """ - Helper function to run top_k_per_row_decode test with given parameters. - """ + """Helper function to run top_k_per_row_decode test with given parameters.""" torch.set_default_device("cuda:0") # Create test data @@ -365,9 +360,7 @@ def test_top_k_per_row_decode( clean_logits: bool, data_generation: str, ) -> None: - """ - Test top_k_per_row with seq_lens tensor. - """ + """Test top_k_per_row with seq_lens tensor.""" set_random_seed(0) vocab_size = 20000 _run_top_k_per_row_decode_test( @@ -379,9 +372,7 @@ def test_top_k_per_row_decode( @pytest.mark.parametrize("clean_logits", [True, False]) @torch.inference_mode() def test_top_k_per_row_decode_large_vocab_size(clean_logits: bool) -> None: - """ - Test top_k_per_row_decode with large vocabulary size. - """ + """Test top_k_per_row_decode with large vocabulary size.""" set_random_seed(0) top_k = 2048 batch_size = 2 @@ -886,8 +877,7 @@ def test_deepseek_workspace_topk( next_n: int, backend: str, ) -> None: - """ - Test workspace top-k backends with varying sequence lengths and speculative + """Test workspace top-k backends with varying sequence lengths and speculative decoding. Supports speculative decoding with next_n > 1. """ @@ -1040,8 +1030,7 @@ def run_large_context_topk_test( seed: int = 42, backend: str = "cooperative_topk", ) -> None: - """ - Helper to run a top-k backend test with given parameters. + """Helper to run a top-k backend test with given parameters. Args: batch_size: Number of rows/sequences @@ -1050,6 +1039,7 @@ def run_large_context_topk_test( data_type: Type of test data to generate seed: Random seed for reproducibility backend: Top-k backend to test + """ torch.set_default_device("cuda:0") set_random_seed(seed) @@ -1274,8 +1264,7 @@ def test_cooperative_topk_cs2(top_k: int) -> None: @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) @torch.inference_mode() def test_workspace_topk_correctness(test_config: dict, backend: str) -> None: - """ - Comprehensive correctness tests covering: + """Comprehensive correctness tests covering: - Sequence length edge cases (trivial, boundary, varied) - Very small sequences (< 100 elements) - Mixed sequence lengths in same batch @@ -1339,8 +1328,7 @@ def test_workspace_topk_correctness(test_config: dict, backend: str) -> None: @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) @torch.inference_mode() def test_workspace_topk_algorithm_paths(test_config: dict, backend: str) -> None: - """ - Test different algorithm execution paths (capped at 163840 for DeepSeek V3.2): + """Test different algorithm execution paths (capped at 163840 for DeepSeek V3.2): - Batch size scalability (1, 4, 32, 256) - Single-CTA vs Multi-CTA execution - Extreme configurations (large batch, max context length) @@ -1357,8 +1345,7 @@ def test_workspace_topk_algorithm_paths(test_config: dict, backend: str) -> None @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) @torch.inference_mode() def test_workspace_topk_stress(backend: str) -> None: - """ - Stress test with random configurations to catch edge cases. + """Stress test with random configurations to catch edge cases. Capped at 163840 (DeepSeek V3.2 max context) for realistic testing. """ torch.set_default_device("cuda:0") @@ -1505,8 +1492,7 @@ def test_cooperative_topk_512_tie_workspace_is_per_row() -> None: @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) @torch.inference_mode() def test_workspace_topk(test_config: dict, top_k: int, backend: str) -> None: - """ - Tests specific to workspace top-k backends: + """Tests specific to workspace top-k backends: - Mixed medium/large rows in the same batch (dynamic per-row dispatch) - Boundary around LARGE_THRESHOLD (32K) - Trivial + medium + large rows in a single batch @@ -1567,8 +1553,7 @@ def test_persistent_topk_reused_group_after_short_row() -> None: @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) @torch.inference_mode() def test_workspace_topk_padded_stride(top_k: int, backend: str) -> None: - """ - Test workspace top-k backends with padded logits (large stride, small seq_len) + """Test workspace top-k backends with padded logits (large stride, small seq_len) to simulate the e2e CUDAGraph scenario where fp8_paged_mqa_logits returns [B, max_model_len] with max_model_len=163840. """ @@ -1755,7 +1740,8 @@ def test_sparse_indexer_decode_topk_explicit_backends( ) -> None: """Every explicit sparse-indexer decode top-k backend must restrict each row to its seq_len (dirty data past the end must never be selected) and - match torch.topk on the valid region.""" + match torch.topk on the valid region. + """ from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk @@ -1826,7 +1812,8 @@ def test_sparse_indexer_decode_topk_short_seq_lens( backend: str, workspace_init ) -> None: """seq_len < next_n: derived per-row ends must clamp at 0 (the reference - kernels do max(0, ...)); negative ends are OOB for the ragged kernels.""" + kernels do max(0, ...)); negative ends are OOB for the ragged kernels. + """ from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk diff --git a/tests/kernels/utils.py b/tests/kernels/utils.py index 8bd63872f3..cf68565e9d 100644 --- a/tests/kernels/utils.py +++ b/tests/kernels/utils.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Kernel test utils""" +"""Kernel test utils.""" import itertools import random @@ -94,16 +94,15 @@ def _assert_deterministic( class QKVInputs(NamedTuple): - """ - Data structure for representing unpacked attention inputs, + """Data structure for representing unpacked attention inputs, query/key/values and their sequence lengths. Attributes: - * {query,key,value}: unpacked (batch_size x padded_seq_len x num_heads x head_size) attention inputs * q_seq_lens: query sequence lengths list * kv_seq_lens: shared key/value sequence lengths list + """ query: torch.Tensor @@ -114,16 +113,15 @@ class QKVInputs(NamedTuple): class QKVO(NamedTuple): - """ - Data structure for representing unpacked attention inputs, + """Data structure for representing unpacked attention inputs, alongside unpacked known-correct attention output Attributes: - * qkv: unpacked (batch_size x padded_seq_len x num_heads x head_size) attention inputs * ideal_output: unpacked (batch_size x padded_seq_len x num_heads x head_size) known-correct attention output + """ qkv: QKVInputs @@ -131,11 +129,9 @@ class QKVO(NamedTuple): class PackedQKVInputs(NamedTuple): - """ - Data structure for representing packed attention inputs + """Data structure for representing packed attention inputs. Attributes: - * {query,key,value}: packed (number_of_tokens x num_heads x head_size) attention inputs * q_start_loc_list: list of query start locations within packed tensor @@ -143,6 +139,7 @@ class PackedQKVInputs(NamedTuple): packed tensor * q_seq_lens: query sequence lengths list * kv_seq_lens: shared key/value sequence lengths list + """ query: torch.Tensor @@ -155,16 +152,15 @@ class PackedQKVInputs(NamedTuple): class PackedQKVO(NamedTuple): - """ - Data structure for representing packed attention inputs, + """Data structure for representing packed attention inputs, alongside packed known-correct attention output Attributes: - * packed_qkv: packed (number_of_tokens x num_heads x head_size) attention inputs * ideal_output: packed (number_of_tokens x num_heads x head_size) known-correct attention output + """ packed_qkv: PackedQKVInputs | None @@ -172,13 +168,12 @@ class PackedQKVO(NamedTuple): class KVMemoryMap(NamedTuple): - """ - Data structure for encapsulating KV cache memory mapping. + """Data structure for encapsulating KV cache memory mapping. Attributes: - * block_tables: KV cache block tables * slot_mapping: mapping of sequence offset to physical address + """ block_tables: torch.Tensor @@ -186,18 +181,17 @@ class KVMemoryMap(NamedTuple): class PhaseTestParameters(NamedTuple): - """ - Data structure for encapsulating the test parameters + """Data structure for encapsulating the test parameters for a given test "phase" (prefill or decode phase) and attention scenario (encoder, decoder-self, encoder/decoder-cross) Attributes: - * packed_qkvo: packed (number_of_tokens x num_heads x head_size) attention inputs & known-correct output * kv_mmap: KV cache memory mapping, specific to this test phase & attention scenario + """ packed_qkvo: PackedQKVO @@ -208,13 +202,11 @@ def maybe_make_int_tensor( _list: list[int] | None, device: torch.device | str, ) -> torch.Tensor: - """ - Convert Python int list to a 1D int torch.Tensor on `device` + """Convert Python int list to a 1D int torch.Tensor on `device`. Returns: + If _list is not None: 1D int torch.Tensor on `device`, `None` otherwise. - * If _list is not None: 1D int torch.Tensor on `device` - * None otherwise """ return ( None if _list is None else torch.tensor(_list, dtype=torch.int, device=device) @@ -225,13 +217,11 @@ def maybe_make_long_tensor( _list: list[int] | None, device: torch.device | str, ) -> torch.Tensor: - """ - Convert Python int list to a 1D long torch.Tensor on `device` + """Convert Python int list to a 1D long torch.Tensor on `device`. Returns: + If _list is not None: 1D long torch.Tensor on `device`, `None` otherwise. - * If _list is not None: 1D long torch.Tensor on `device` - * None otherwise """ return ( None if _list is None else torch.tensor(_list, dtype=torch.long, device=device) @@ -239,11 +229,11 @@ def maybe_make_long_tensor( def maybe_max(_list: list | None) -> Number | None: - """ + """Compute the maximum of a list, if it exists. + Returns: + If _list is not None: max(_list), `None` otherwise. - * If _list is not None: max(_list) - * None otherwise """ return None if _list is None else max(_list) @@ -252,19 +242,16 @@ def make_causal_mask( q_max_seq_len: int, kv_max_seq_len: int, ) -> torch.Tensor: - """ - Create a q_max_seq_len x kv_max_seq_len causal mask + """Create a q_max_seq_len x kv_max_seq_len causal mask. Arguments: - - * q_max_seq_len: query max seq len - * kv_max_seq_len: key/value max seq len + q_max_seq_len: query max seq len + kv_max_seq_len: key/value max seq len Returns: + 2D tensor, q_max_seq_len x kv_max_seq_len - * 2D tensor, q_max_seq_len x kv_max_seq_len """ - # Create a matrix where entry (i, j) is True if i >= j mask = torch.triu(torch.ones(q_max_seq_len, kv_max_seq_len), diagonal=1) # Replace True with float('-inf') and False with 0 @@ -281,8 +268,7 @@ def ref_masked_attention( q_seq_lens: list | None = None, kv_seq_lens: list | None = None, ) -> torch.Tensor: - """ - "Golden" masked attention reference. Supports two types of masking: + """Reference ("golden") masked attention. Supports two types of masking: * Basic attention mask, utilizing {q,kv}_seq_lens args to mask out padding elements @@ -290,21 +276,18 @@ def ref_masked_attention( causal Arguments: - - * query: batch_size x q_padded_seq_len x num_heads x head_size - * key: batch_size x kv_padded_seq_len x num_heads x head_size - * value: batch_size x kv_padded_seq_len x num_heads x head_size - * scale: Attention scale factor - * custom_mask: custom attention mask; good place to inject a causal - attention mask - * q_seq_lens: list of unpadded query seq_lens for each batch index - * kv_seq_lens: list of unpadded key/value seq_lens for each batch index + query: batch_size x q_padded_seq_len x num_heads x head_size + key: batch_size x kv_padded_seq_len x num_heads x head_size + value: batch_size x kv_padded_seq_len x num_heads x head_size + scale: Attention scale factor + custom_mask: custom attention mask; good place to inject a causal attention mask + q_seq_lens: list of unpadded query seq_lens for each batch index + kv_seq_lens: list of unpadded key/value seq_lens for each batch index Returns: + Attention result, batch_size x q_padded_seq_len x num_heads x head_size - * Attention result, batch_size x q_padded_seq_len x num_heads x head_size """ - assert q_seq_lens is not None assert kv_seq_lens is not None @@ -346,8 +329,7 @@ def make_qkv( attn_type: AttentionType = AttentionType.ENCODER_DECODER, force_max_len: bool = False, ) -> tuple[QKVInputs, QKVInputs, QKVInputs]: - """ - Construct QKV test tensors for self- and cross-attention. + """Construct QKV test tensors for self- and cross-attention. Generates three query/key/value triplets: @@ -361,30 +343,28 @@ def make_qkv( seqlens Arguments: - - * batch_size - * max_q_seq_len: max query seq len - * max_kv_seq_len: max key/value seq len - * num_heads - * head_size - * is_encoder_decoder_attn: if True, query seqlen may differ from - key/value seqlen (as is often the case for cross-attention); - o/w, query/key/value seqlens match at each batch index - (max_kv_seq_len is unused) - * force_kv_seq_lens: if not None, overrides kv sequence lengths - * attn_type: encoder, decoder self, or enc/dec cross attention - * force_max_len: if True, all query seqlens are max_q_seq_len; o/w query - seqlens are random in [2,max_q_seq_lens]. Same for key/value seqlens - and max_kv_seq_len, unless forced by is_encoder_decoder_attn=False - * device: CPU or CUDA device + batch_size: number of sequences in the batch + max_q_seq_len: max query seq len + max_kv_seq_len: max key/value seq len + num_heads: number of attention heads + head_size: head dimension + is_encoder_decoder_attn: if True, query seqlen may differ from + key/value seqlen (as is often the case for cross-attention); + o/w, query/key/value seqlens match at each batch index + (max_kv_seq_len is unused) + force_kv_seq_lens: if not None, overrides kv sequence lengths + attn_type: encoder, decoder self, or enc/dec cross attention + force_max_len: if True, all query seqlens are max_q_seq_len; o/w query + seqlens are random in [2,max_q_seq_lens]. Same for key/value seqlens + and max_kv_seq_len, unless forced by is_encoder_decoder_attn=False + device: CPU or CUDA device Returns: + Overall QKVInputs structure (containing full unpacked Q/K/V tensors) + Prefill QKVInputs structure (containing all but the last sequence offset) + Decode QKVInputs structure (containing all only the last sequence offset) - * Overall QKVInputs structure (containing full unpacked Q/K/V tensors) - * Prefill QKVInputs structure (containing all but the last sequence offset) - * Decode QKVInputs structure (containing all only the last sequence offset) """ - if force_max_len: q_seq_lens = [max_q_seq_len for _ in range(batch_size)] else: @@ -474,24 +454,21 @@ def make_qkv( def pack_tensor( unpacked_tensor: torch.Tensor, seq_lens: list[int], device: torch.device | str ) -> tuple[torch.Tensor, list[int]]: - """ - Pack a batch_size x padded_seq_len x num_heads x head_size tensor into an + """Pack a batch_size x padded_seq_len x num_heads x head_size tensor into an unpadded number_of_tokens x num_heads x head_size tensor, where number_of_tokens = sum(seq_lens) Arguments: + unpacked_tensor: batch_size x padded_seq_len x num_heads x head_size + seq_lens: list of token counts for each seq + device: CPU or CUDA device - * unpacked_tensor: batch_size x padded_seq_len x num_heads x head_size - * seq_lens: list of token counts for each seq - * device: CPU or CUDA device + Returns: + packed_tensor: number_of_tokens x num_heads x head_size + start_loc_list: start idx of each batch elt in packed_tensor; [0] + + list(itertools.accumulate(seq_lens)) - Returns - - * packed_tensor: number_of_tokens x num_heads x head_size - * start_loc_list: start idx of each batch elt in packed_tensor; [0] + - list(itertools.accumulate(seq_lens)) """ - num_tok = sum(seq_lens) num_heads = unpacked_tensor.shape[-2] head_size = unpacked_tensor.shape[-1] @@ -507,8 +484,7 @@ def pack_tensor( def pack_qkv(qkv: QKVInputs, device: torch.device | str) -> PackedQKVInputs: - """ - Individually pack each of Q, K and V, each with dimensions batch_size x + """Individually pack each of Q, K and V, each with dimensions batch_size x padded_seq_len x num_heads x head_size, into respective number_of_tokens x num_heads x head_size tensors. @@ -517,17 +493,15 @@ def pack_qkv(qkv: QKVInputs, device: torch.device | str) -> PackedQKVInputs: For K and V, number_of_tokens = sum(kv_seq_lens) Arguments: + qkv: Unpacked (batch_size x padded_seq_len x num_heads x head_size) + attention inputs + device: CPU or CUDA device - * qkv: Unpacked (batch_size x padded_seq_len x num_heads x head_size) - attention inputs - * device: CPU or CUDA device + Returns: + Packed (number_of_tokens x num_heads x head_size) QKV inputs + derived from unpacked inputs - Returns - - * Packed (number_of_tokens x num_heads x head_size) QKV inputs - derived from unpacked inputs """ - if qkv.query is None: packed_query = None q_start_loc_list = None @@ -563,26 +537,24 @@ def _make_metadata_tensors( torch.Tensor, int | None, ]: - """ - Build scalar & tensor values required to build attention metadata structure. + """Build scalar & tensor values required to build attention metadata structure. Arguments: - - * seq_lens: list of token-counts for each decoder input seq - * context_lens: list of context length values for each seq - * encoder_seq_lens: list of token-counts for each encoder input seq - * device: CPU or CUDA device + seq_lens: list of token-counts for each decoder input seq + context_lens: list of context length values for each seq + encoder_seq_lens: list of token-counts for each encoder input seq + device: CPU or CUDA device Returns: + seq_lens_tensor: decoder seq_lens list, as tensor + context_lens_tensor: context_lens list, as tensor + max_context_len: max(context_lens) + max_seq_len: max(seq_lens) + seq_start_loc: start idx of each sequence + encoder_seq_lens_tensor: encoder seq_lens list, as tensor + encoder_seq_start_loc: start idx of each encoder sequence + max_encoder_seq_len: encoder seq_lens list, as tensor - * seq_lens_tensor: decoder seq_lens list, as tensor - * context_lens_tensor: context_lens list, as tensor - * max_context_len: max(context_lens) - * max_seq_len: max(seq_lens) - * seq_start_loc: start idx of each sequence - * encoder_seq_lens_tensor: encoder seq_lens list, as tensor - * encoder_seq_start_loc: start idx of each encoder sequence - * max_encoder_seq_len: encoder seq_lens list, as tensor """ seq_lens_tensor = maybe_make_int_tensor(seq_lens, device) context_lens_tensor = maybe_make_int_tensor(context_lens, device) @@ -637,22 +609,21 @@ def make_kv_cache( backend: str, default_val: float = 0.0, ) -> torch.Tensor: - """ - Create a fake KV cache. + """Create a fake KV cache. Arguments: - - * num_blocks: number of blocks in the KV cache - * num_heads: number of attention heads - * head_size: head dimension - * block_size: number of offsets within a block - * device: CPU or CUDA device - * default_val: initialization value for KV cache elements + num_blocks: number of blocks in the KV cache + num_heads: number of attention heads + head_size: head dimension + block_size: number of offsets within a block + device: CPU or CUDA device + backend: attention backend name; only 'FLASH_ATTN' is supported + default_val: initialization value for KV cache elements Returns: + kv_cache: 2 x num_blocks x block_size x num_heads x head_size + for backend 'FLASH_ATTN' - * kv_cache: 2 x num_blocks x block_size x num_heads x head_size - * for backend 'FLASH_ATTN' """ if backend != "FLASH_ATTN": raise ValueError(f"Unknown backend value: '{backend}'. Expected 'FLASH_ATTN'.") @@ -663,8 +634,7 @@ def make_kv_cache( def _num_tokens_to_min_blocks(num_tokens: int, block_size: int) -> int: - """ - Compute the minimum number of blocks required to hold num_tokens tokens, + """Compute the minimum number of blocks required to hold num_tokens tokens, given block_size """ return (num_tokens + block_size - 1) // block_size @@ -683,8 +653,7 @@ def split_slot_mapping( seq_lens: list[int], device: torch.device | str, ): - """ - Split a slot mapping into valid prefill- and decode-phase slot mappings. + """Split a slot mapping into valid prefill- and decode-phase slot mappings. Context: * Your goal is to test (1) prefill of N prompts, with prompt-lengths @@ -708,21 +677,19 @@ def split_slot_mapping( The N excised entries are appended to obtain the decode-phase slot mapping Arguments: - - * slot_mapping_list: Length-P 1D slot mapping (as list) reflecting all N - post-decode sequences - * seq_lens: list of N post-decode sequence lengths (K_i + 1 in the - description above) - * device: cuda, cpu, etc. + slot_mapping_list: Length-P 1D slot mapping (as list) reflecting all N + post-decode sequences + seq_lens: list of N post-decode sequence lengths (K_i + 1 in the + description above) + device: cuda, cpu, etc. Returns: + prefill_slot_mapping: Length-M 1D slot mapping (as Tensor) + reflecting all N prefill prompts + decode_slot_mapping: Length-N 1D slot mapping (as Tensor) reflecting + all N decoded tokens - * prefill_slot_mapping: Length-M 1D slot mapping (as Tensor) - reflecting all N prefill prompts - * decode_slot_mapping: Length-N 1D slot mapping (as Tensor) reflecting - all N decoded tokens """ - prefill_slot_mapping = [] decode_slot_mapping = [] @@ -746,8 +713,7 @@ def make_block_tables_slot_mapping( device: torch.device | str, block_base_addr: int = 0, ) -> tuple[torch.Tensor, list[int], int]: - """ - Construct fake block tables & slot mappings. + """Construct fake block tables & slot mappings. For a sequence with num_tokens tokens the minimum number of required KV cache blocks is @@ -772,19 +738,17 @@ def make_block_tables_slot_mapping( i.e. the total of prefill prompt tokens + decoded tokens. Arguments: - - * block_size: number of offsets per block - * seq_lens: list of token-counts for each sequence - * block_base_addr: the block table base address - * device: CPU or CUDA device + block_size: number of offsets per block + seq_lens: list of token-counts for each sequence + block_base_addr: the block table base address + device: CPU or CUDA device Return: - * block_tables_tensor: block table for sequence * slot_mapping_list: slot mapping for sequence * max_block_idx: the highest block address within this block table - """ + """ # Provision minimum number of KV cache blocks num_blocks_list = [ _num_tokens_to_min_blocks(num_tokens, block_size) for num_tokens in seq_lens @@ -824,14 +788,14 @@ def make_block_tables_slot_mapping( def assert_actual_matches_ideal( test_params: PhaseTestParameters, output_under_test: torch.Tensor, backend: str ) -> None: - """ - Assert that observed output matches the ideal output + """Assert that observed output matches the ideal output contained in the test parameters data structure. Arguments: + test_params: Test parameters including packed ideal output + output_under_test: actually observed output value + backend: attention backend name; only 'FLASH_ATTN' is supported - * test_params: Test parameters including packed ideal output - * output_under_test: actually observed output value """ ideal_output = test_params.packed_qkvo.ideal_output if backend != "FLASH_ATTN": @@ -852,9 +816,7 @@ def fp8_allclose( atol: float = 1e-08, equal_nan: bool = False, ) -> bool: - """ - Reference implementation of torch.allclose - """ + """Reference implementation of torch.allclose.""" torch._refs._check_close_args(name="torch.allclose", a=a, b=b, rtol=rtol, atol=atol) return bool( diff --git a/tests/lora/conftest.py b/tests/lora/conftest.py index 8635779cff..f5ca44f174 100644 --- a/tests/lora/conftest.py +++ b/tests/lora/conftest.py @@ -33,7 +33,6 @@ def should_do_global_cleanup_after_test(request) -> bool: This can provide a ~10x speedup for non-GPU unit tests since they don't need to initialize torch. """ - return not request.node.get_closest_marker("skip_global_cleanup") @@ -367,8 +366,7 @@ def qwen36_moe_3d_lora_files(): @pytest.fixture def reset_default_device(): - """ - Some tests, such as `test_punica_ops.py`, explicitly set the + """Some tests, such as `test_punica_ops.py`, explicitly set the default device, which can affect subsequent tests. Adding this fixture helps avoid this problem. """ diff --git a/tests/lora/test_add_lora.py b/tests/lora/test_add_lora.py index 628f5bc28e..97c37c1895 100644 --- a/tests/lora/test_add_lora.py +++ b/tests/lora/test_add_lora.py @@ -55,8 +55,7 @@ async def requests_processing_time(llm, lora_requests: list[LoRARequest]) -> flo @pytest.mark.asyncio async def test_add_lora(chatglm3_lora_files): - """ - The add_lora function is used to preload some LoRA adapters into the + """The add_lora function is used to preload some LoRA adapters into the engine in anticipation of future requests using these adapters. To test this functionality, we use the async engine to process some requests - We do it twice, once with add_lora() preloading and once without. diff --git a/tests/lora/test_default_mm_loras.py b/tests/lora/test_default_mm_loras.py index 5e3c762f79..3c9505adba 100644 --- a/tests/lora/test_default_mm_loras.py +++ b/tests/lora/test_default_mm_loras.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for applying default registered multimodal loras. -""" +"""Tests for applying default registered multimodal loras.""" import os import unittest.mock as mock @@ -134,7 +132,8 @@ def test_default_mm_lora_fails_with_overridden_lora_request( monkeypatch: pytest.MonkeyPatch, ): """Ensure that if the lora_request conflicts with default_mm_loras, - we use the lora_request.""" + we use the lora_request. + """ run_test( vllm_runner, audio_assets, diff --git a/tests/lora/test_fused_moe_lora_kernel.py b/tests/lora/test_fused_moe_lora_kernel.py index 444b609492..9ee0ed8584 100644 --- a/tests/lora/test_fused_moe_lora_kernel.py +++ b/tests/lora/test_fused_moe_lora_kernel.py @@ -41,8 +41,7 @@ def CEILDIV(x, y): def assign_loras_to_tokens(num_tokens: int, num_sequences: int, max_loras: int): - """ - Split `num_tokens` into `num_sequences` sequences. + """Split `num_tokens` into `num_sequences` sequences. Each sequence randomly selects 1 LoRA index from [0, max_loras), and all tokens in that sequence are assigned this LoRA index. @@ -54,6 +53,7 @@ def assign_loras_to_tokens(num_tokens: int, num_sequences: int, max_loras: int): Returns: torch.Tensor: 1D tensor of shape [num_tokens], where each value is the LoRA index assigned to that token. + """ assert num_sequences > 0 and max_loras > 0 assert num_tokens >= num_sequences, "num_tokens must be >= num_sequences" @@ -81,9 +81,8 @@ def assign_loras_to_tokens(num_tokens: int, num_sequences: int, max_loras: int): def assign_experts_to_tokens(num_tokens: int, num_experts: int, top_k_num: int): - """ - For each token, randomly select `top_k_num` distinct experts out of `num_experts`, - and assign normalized random weights that sum to 1. + """For each token, randomly select `top_k_num` distinct experts out of + `num_experts`, and assign normalized random weights that sum to 1. Args: num_tokens (int): Total number of tokens. @@ -95,6 +94,7 @@ def assign_experts_to_tokens(num_tokens: int, num_experts: int, top_k_num: int): expert index for each token. expert_weights (torch.Tensor): shape [num_tokens, top_k_num], normalized weights (sum = 1 per row). + """ assert top_k_num <= num_experts, "top_k_num must be <= num_experts" @@ -379,8 +379,7 @@ def use_fused_moe_lora_kernel_naive( offset=0, add_inputs=True, ): - """ - Test helper for naive_block_assignment path. + """Test helper for naive_block_assignment path. Skips moe_lora_align_block_size and uses flattened topk_ids as expert_ids. """ config = { @@ -472,8 +471,7 @@ def test_fused_moe_lora_kernel_naive_block_assignment( device, seed, ): - """ - Test the naive_block_assignment path of the fused_moe_lora kernel. + """Test the naive_block_assignment path of the fused_moe_lora kernel. This path is triggered when batch_size * top_k is much smaller than num_experts * max_loras, and skips the moe_lora_align_block_size kernel. """ @@ -1125,7 +1123,8 @@ def test_fused_moe_lora_kernel_one_shot_early_exit(trigger, device): @pytest.mark.parametrize("device", DEVICES) def test_fused_moe_lora_kernel_zero_grid_no_crash(device): """num_active_loras=0 (or num_slices=0) would otherwise launch a grid - with a zero dimension. one-shot wrapper must short-circuit before launch.""" + with a zero dimension. one-shot wrapper must short-circuit before launch. + """ torch.set_default_device(device) set_random_seed(0) num_tokens, top_k, E, max_loras, R, K, N = 8, 2, 8, 4, 16, 1024, 1024 @@ -1187,7 +1186,8 @@ def test_fused_moe_lora_kernel_zero_grid_no_crash(device): def test_fused_moe_lora_kernel_rejects_bad_block_size_m(device): """one-shot must surface a clear assertion when shrink_block_size_m is not a power of 2 / less than 16, instead of the cryptic Triton compile - failure (`arange's range must be a power of 2`).""" + failure (`arange's range must be a power of 2`). + """ torch.set_default_device(device) set_random_seed(0) num_tokens, top_k, E, max_loras, R, K, N = 32, 2, 8, 4, 16, 1024, 1024 diff --git a/tests/lora/test_layers.py b/tests/lora/test_layers.py index 2dab86bc0d..e6f84cf0bc 100644 --- a/tests/lora/test_layers.py +++ b/tests/lora/test_layers.py @@ -126,8 +126,7 @@ def clean_cache_reset_device(reset_default_device): @pytest.fixture(autouse=True) def skip_cuda_with_stage_false(request): - """ - On cuda-like platforms, we use the same kernels for prefill and decode + """On cuda-like platforms, we use the same kernels for prefill and decode stage, and 'stage' is generally ignored, so we only need to test once. """ if current_platform.is_cuda_alike() or current_platform.is_xpu(): @@ -153,8 +152,8 @@ def get_random_id_to_index( num_slots: The number of slots in the mapping. Must be larger than num_loras. log: Whether to log the output. - """ + """ if num_loras > num_slots: raise ValueError( f"num_loras is higher than num_slots: {num_loras} > {num_slots}. " @@ -190,8 +189,8 @@ def populate_loras( repeats: must only be set for column parallel packed layers. Indicates the number of loras to compose together to create a single lora layer. - """ + """ # Dictionary that maps the lora ID to the # corresponding lora weights. lora_dict: dict[int, LoRALayerWeights] = dict() @@ -246,8 +245,8 @@ def create_random_inputs( input_range: the range of values to include in the input. input_range[0] <= possible input values < input_range[1] input_type: the type of values in the input. - """ + """ low, high = input_range inputs: list[torch.Tensor] = [] diff --git a/tests/lora/test_lora_functions.py b/tests/lora/test_lora_functions.py index ecf2ad9ee6..1b841b70ce 100644 --- a/tests/lora/test_lora_functions.py +++ b/tests/lora/test_lora_functions.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Script to test add_lora, remove_lora, pin_lora, list_loras functions. -""" +"""Script to test add_lora, remove_lora, pin_lora, list_loras functions.""" import pytest diff --git a/tests/lora/test_mixtral.py b/tests/lora/test_mixtral.py index 3868bff796..a8f1fdda64 100644 --- a/tests/lora/test_mixtral.py +++ b/tests/lora/test_mixtral.py @@ -32,7 +32,7 @@ def do_sample( @pytest.mark.parametrize("tp_size", [4]) def test_mixtral_lora(mixtral_lora_files, tp_size): - """Original test, the LoRA model has the common target modules, not all""" + """Original test, the LoRA model has the common target modules, not all.""" if ( torch.accelerator.device_count() < tp_size and tp_size > 1 diff --git a/tests/lora/test_moe_lora_align_sum.py b/tests/lora/test_moe_lora_align_sum.py index 80ba0a56ca..df27752d2c 100644 --- a/tests/lora/test_moe_lora_align_sum.py +++ b/tests/lora/test_moe_lora_align_sum.py @@ -249,7 +249,8 @@ def _build_and_run_align( ) def test_moe_lora_align_block_size_mixed_base_and_lora(max_loras): """Regression test for issue #32235: real LoRA slot must not be skipped - when ``active_lora_ids`` has -1 at position 0.""" + when ``active_lora_ids`` has -1 at position 0. + """ out = _build_and_run_align( num_lora_tokens=8, num_base_tokens=8, max_loras=max_loras ) @@ -288,7 +289,8 @@ def test_moe_lora_align_block_size_disabled_adapter_untouched(): kernels. Pins the invariant protected by the ``adapter_enabled`` guard in ``lora_count_and_sort_expert_tokens_kernel``: without it the sort kernel reads uninitialized ``token_mask`` values for disabled slots and - pollutes ``sorted_token_ids`` / ``cumsum_buffer``.""" + pollutes ``sorted_token_ids`` / ``cumsum_buffer``. + """ max_loras = 1 out = _build_and_run_align( num_lora_tokens=16, diff --git a/tests/lora/test_moe_lora_ep_load.py b/tests/lora/test_moe_lora_ep_load.py index ead40b2552..52cccc6b06 100644 --- a/tests/lora/test_moe_lora_ep_load.py +++ b/tests/lora/test_moe_lora_ep_load.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -End-to-end correctness test for 2D MoE LoRA expert-parallel +"""End-to-end correctness test for 2D MoE LoRA expert-parallel load-time slicing """ @@ -20,7 +19,8 @@ NON_EXPERT_MODULES = ("q_proj", "k_proj", "v_proj", "o_proj", "gate") def _expected_lora_modules() -> set[str]: """Replicate the set ``WorkerLoRAManager._load_adapter`` would build - from this model's ``packed_modules_mapping``.""" + from this model's ``packed_modules_mapping``. + """ expected: set[str] = set(NON_EXPERT_MODULES) for expert in range(GLOBAL_NUM_EXPERTS): for proj in EXPERT_PROJECTIONS: @@ -43,7 +43,8 @@ def _load(lora_dir, peft_helper, *, moe_ep_spec, lora_id): def test_moe_lora_ep2_real_qwen3moe(qwen3moe_lora_files, ep_rank): """ep_size=2 against the real Qwen3-MoE adapter: each rank's loaded LoRA has the right size, the right expert membership, and the - right tensor values.""" + right tensor values. + """ peft_helper = PEFTHelper.from_local_dir( qwen3moe_lora_files, max_position_embeddings=4096 ) diff --git a/tests/lora/test_punica_ops.py b/tests/lora/test_punica_ops.py index 9e1f99b3f6..f2bb18aac2 100644 --- a/tests/lora/test_punica_ops.py +++ b/tests/lora/test_punica_ops.py @@ -45,7 +45,8 @@ def _bgmv_shrink( inputs, lora_weight, output, seq_len_tensor, lora_indices, scaling=1.0 ): """Memory-efficient shrink reference: per-LoRA matmul loop. - output[mask] = scaling * inputs[mask] @ weight.T""" + output[mask] = scaling * inputs[mask] @ weight.T + """ exploded = torch.repeat_interleave(lora_indices, seq_len_tensor) for lid in exploded.unique(): if lid < 0: @@ -66,7 +67,8 @@ def _bgmv_expand( add_inputs=False, ): """Memory-efficient expand reference: per-LoRA matmul loop. - output[mask, offset:offset+n] (+)= inputs[mask] @ weight.T""" + output[mask, offset:offset+n] (+)= inputs[mask] @ weight.T + """ exploded = torch.repeat_interleave(lora_indices, seq_len_tensor) for lid in exploded.unique(): if lid < 0: @@ -160,8 +162,7 @@ def check_lora_shrink_kernel( seq_length: int, scaling: float, ): - """ - Compare outputs of torch_ops.sgmv_shrink and triton_ops.lora_shrink + """Compare outputs of torch_ops.sgmv_shrink and triton_ops.lora_shrink kernels. """ data: PunicaTensors = generate_data_for_nslices( @@ -234,8 +235,7 @@ def check_lora_expand_kernel( seq_length: int, add_inputs: bool, ): - """ - Compare outputs of torch_ops.sgmv_expand and triton_ops.lora_expand + """Compare outputs of torch_ops.sgmv_expand and triton_ops.lora_expand kernels. """ data: PunicaTensors = generate_data_for_nslices( @@ -438,9 +438,7 @@ def test_kernels( seed: int, op_type: str, ): - """ - Tests LoRA kernels. - """ + """Tests LoRA kernels.""" torch.set_default_device(device) torch.accelerator.set_device_index(device) set_random_seed(seed) @@ -491,9 +489,7 @@ def test_kernels_hidden_size( seed: int, op_type: str, ): - """ - Tests SGMV and LoRA kernels. - """ + """Tests SGMV and LoRA kernels.""" torch.set_default_device(device) torch.accelerator.set_device_index(device) set_random_seed(seed) @@ -526,8 +522,7 @@ def test_kernels_hidden_size( @pytest.mark.parametrize("device", DEVICES) def test_add_lora_fused_moe_early_exit(device): - """ - Ensures add_lora_fused_moe does not invoke the LoRA kernel or + """Ensures add_lora_fused_moe does not invoke the LoRA kernel or modify the output tensor when no_lora_flag_cpu is True """ from types import SimpleNamespace diff --git a/tests/lora/test_punica_ops_fp8.py b/tests/lora/test_punica_ops_fp8.py index e447cef9d8..18ac6fcd7a 100644 --- a/tests/lora/test_punica_ops_fp8.py +++ b/tests/lora/test_punica_ops_fp8.py @@ -877,7 +877,8 @@ def test_lora_shrink_fp8( quant_mode: str, ): """Test FP8 shrink kernel with per-tensor, per-channel, and block-wise - quantization, comparing against the bf16 baseline.""" + quantization, comparing against the bf16 baseline. + """ torch.set_default_device(device) set_random_seed(seed) @@ -932,7 +933,8 @@ def test_lora_expand_fp8( quant_mode: str, ): """Test FP8 expand kernel with per-tensor, per-channel, and block-wise - quantization, comparing against the bf16 baseline.""" + quantization, comparing against the bf16 baseline. + """ torch.set_default_device(device) set_random_seed(seed) diff --git a/tests/lora/test_punica_xpu_ops.py b/tests/lora/test_punica_xpu_ops.py index 585c97cfa5..36966809eb 100644 --- a/tests/lora/test_punica_xpu_ops.py +++ b/tests/lora/test_punica_xpu_ops.py @@ -86,9 +86,7 @@ def check_bgmv_shrink( device: str, scaling: float, ): - """ - Compare vllm.bgmv_shrink against a reference implementation. - """ + """Compare vllm.bgmv_shrink against a reference implementation.""" seq_length = 1 data: PunicaTensors = generate_data( batches, @@ -130,9 +128,7 @@ def check_bgmv_expand( device: str, add_inputs: bool, ): - """ - Compare vllm.bgmv_expand against a reference implementation. - """ + """Compare vllm.bgmv_expand against a reference implementation.""" seq_length = 1 data: PunicaTensors = generate_data( batches, @@ -172,9 +168,7 @@ def check_bgmv_expand_slice( device: str, add_inputs: bool, ): - """ - Compare vllm.bgmv_expand_slice against a reference implementation. - """ + """Compare vllm.bgmv_expand_slice against a reference implementation.""" seq_length = 1 data: PunicaTensors = generate_data_for_expand_nslices( batches, diff --git a/tests/lora/test_qwen3_unembed.py b/tests/lora/test_qwen3_unembed.py index 4064b443ac..a70df902f7 100644 --- a/tests/lora/test_qwen3_unembed.py +++ b/tests/lora/test_qwen3_unembed.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for Qwen3 unembed LoRA support. +"""Tests for Qwen3 unembed LoRA support. This test creates synthetic LoRA weights that include lm_head (output embedding) to verify that Qwen3 properly supports LoRA on the unembed/lm_head layer. diff --git a/tests/lora/test_qwen3_with_multi_loras.py b/tests/lora/test_qwen3_with_multi_loras.py index 9870da6959..185ab91c7a 100644 --- a/tests/lora/test_qwen3_with_multi_loras.py +++ b/tests/lora/test_qwen3_with_multi_loras.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This script contains: +"""This script contains: 1. test multi loras service with tp >= 2 2. test multi loras request """ @@ -125,8 +124,7 @@ def test_multi_loras_with_tp_sync(): return output_text def reload_lora(name: str): - """ - reload a lora to simulate the case: + """Reload a lora to simulate the case: setting `VLLM_ALLOW_RUNTIME_LORA_UPDATING=true` for dynamic lora loading and unloading """ @@ -205,8 +203,7 @@ def test_multiple_lora_requests(): def test_load_inplace_offline_reload( qwen3_meowing_lora_files: str, qwen3_woofing_lora_files: str ) -> None: - """ - Test that load_inplace=True allows reloading LoRA adapters with the same ID + """Test that load_inplace=True allows reloading LoRA adapters with the same ID in offline mode (using LLM class directly). """ llm = LLM( @@ -255,8 +252,7 @@ def test_load_inplace_offline_reload( def test_load_inplace_false_no_reload( qwen3_meowing_lora_files: str, qwen3_woofing_lora_files: str ) -> None: - """ - Test that load_inplace=False prevents reloading when an adapter + """Test that load_inplace=False prevents reloading when an adapter with the same ID already exists. """ llm = LLM( diff --git a/tests/lora/test_qwenvl.py b/tests/lora/test_qwenvl.py index 3362aa46ed..063de2b587 100644 --- a/tests/lora/test_qwenvl.py +++ b/tests/lora/test_qwenvl.py @@ -46,7 +46,7 @@ class TestConfig: class Qwen2VLTester: - """Test helper for Qwen2 VL models with LoRA""" + """Test helper for Qwen2 VL models with LoRA.""" PROMPT_TEMPLATE = ( "<|im_start|>system\nYou are a helpful assistant.<|im_end|>" @@ -199,7 +199,7 @@ def _enable_deterministic_lora_shrink(monkeypatch: pytest.MonkeyPatch) -> None: def test_qwen2vl_lora(qwen2vl_lora_files): - """Test Qwen 2.0 VL model with LoRA""" + """Test Qwen 2.0 VL model with LoRA.""" config = TestConfig(model_path=QWEN2VL_MODEL_PATH, lora_path=qwen2vl_lora_files) with Qwen2VLTester(config) as tester: # Test with different LoRA IDs @@ -230,7 +230,7 @@ def test_qwen2vl_lora_beam_search(qwen2vl_lora_files): current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests" ) def test_qwen25vl_lora(qwen25vl_lora_files): - """Test Qwen 2.5 VL model with LoRA""" + """Test Qwen 2.5 VL model with LoRA.""" config = TestConfig(model_path=QWEN25VL_MODEL_PATH, lora_path=qwen25vl_lora_files) with Qwen2VLTester(config) as tester: # Test with different LoRA IDs @@ -292,8 +292,7 @@ def test_qwen2vl_multiple_lora_types( qwen2vl_vision_tower_lora_files, monkeypatch: pytest.MonkeyPatch, ): - """ - Test multiple LoRA adapter types (language, vision tower + connector, + """Test multiple LoRA adapter types (language, vision tower + connector, vision tower only) using the same LLM instance to verify mm_encoder_cache behavior with different LoRA requests. diff --git a/tests/lora/test_whisper.py b/tests/lora/test_whisper.py index 6f1a894cf9..73cc550b29 100644 --- a/tests/lora/test_whisper.py +++ b/tests/lora/test_whisper.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Integration tests for Whisper models with LoRA adapters. +"""Integration tests for Whisper models with LoRA adapters. These tests verify that Whisper models can correctly load and use LoRA adapters for speech-to-text transcription tasks. diff --git a/tests/lora/utils.py b/tests/lora/utils.py index e5ce7a8846..8ee2f8af82 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -116,8 +116,7 @@ class PunicaTensors: token_lora_mapping: torch.Tensor def meta(self) -> tuple[int, int]: - """ - Infer max_seq_length and token_nums from the tensors + """Infer max_seq_length and token_nums from the tensors and return them. """ max_seq_length = self.seq_len_tensor.max() diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py index 507fd93e9d..6fdd99deaa 100644 --- a/tests/model_executor/layers/test_fused_shared_expert.py +++ b/tests/model_executor/layers/test_fused_shared_expert.py @@ -808,7 +808,6 @@ def test_models_fse_init( monkeypatch: pytest.MonkeyPatch, ) -> None: """Model construction resolves FSE consistently with Quark quantization.""" - quantization_config: dict[str, Any] = ( get_deepseek_v4_quark_config(["layers.0.ffn.shared_experts"] if exclude else []) if model_type == "deepseek_v4" @@ -1320,7 +1319,8 @@ def test_fp8_shared_expert_fse_allows_symmetric_ignored_layers() -> None: def test_fp8_shared_expert_fse_expands_packed_projections() -> None: """Both shards of a fused projection ignored is symmetric within that - projection, so the mismatch reported is against the routed experts.""" + projection, so the mismatch reported is against the routed experts. + """ quant_config = _fp8_config( ignored_layers=[ "model.layers.0.mlp.shared_experts.gate_proj", @@ -1344,7 +1344,8 @@ def test_fp8_shared_expert_fse_expands_packed_projections() -> None: def test_fp8_shared_expert_fse_propagates_partial_shard_exclusion() -> None: """Half a fused projection excluded is rejected by `is_layer_skipped` - itself, exactly as it is in `Fp8Config.get_quant_method`.""" + itself, exactly as it is in `Fp8Config.get_quant_method`. + """ quant_config = _fp8_config( ignored_layers=["model.layers.0.mlp.shared_experts.gate_proj"] ) diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index 41a06d0546..91fc039569 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -346,7 +346,8 @@ def test_deepseek_v32_dispatches_selected_mha( def test_select_candidate_blocks_tolerates_empty_rows(): """Full-cudagraph decode pads the batch with seq_len-0 rows. The newest block pin must not index -1 for them (device-side assert); they select - no candidate blocks and real rows still pin their newest block.""" + no candidate blocks and real rows still pin their newest block. + """ block_size, topk_blocks = 8, 3 logits = torch.zeros(3, 64, device="cuda") logits[0, 3] = 5.0 # block 0 scores highest for row 0 @@ -457,7 +458,8 @@ def test_candidate_kernels_preserve_packed_bounds_and_padding( def _make_candidates(lens: torch.Tensor, num_candidates: int, cbk: int): """Production-faithful candidates: top min(avail, K) blocks per row — - full coverage for lens <= K*cbk, a fixed-size subset beyond that.""" + full coverage for lens <= K*cbk, a fixed-size subset beyond that. + """ avail = ((lens.cpu() + cbk - 1) // cbk).tolist() rows = [] for a in avail: @@ -493,7 +495,8 @@ def _candidate_token_mask( def _reference_topk(masked: torch.Tensor, topk: int) -> torch.Tensor: """Plain ``torch.topk`` over -inf-masked logits; slots past a row's valid - count are -1 (the indexer's "no token" sentinel).""" + count are -1 (the indexer's "no token" sentinel). + """ values, indices = masked.topk(topk, dim=1) return torch.where(values > float("-inf"), indices.int(), -1) @@ -506,7 +509,8 @@ def _assert_same_selection( ) -> None: """Both top-k index sets must pick the same multiset of logit values per row (-1 padding must match too). Indices themselves may differ: the - top-k kernels break ties among equal values in nondeterministic order.""" + top-k kernels break ties among equal values in nondeterministic order. + """ for r in range(logits.shape[0]): base = int(row_starts[r]) if row_starts is not None else 0 g, e = got[r][got[r] >= 0], expected[r][expected[r] >= 0] @@ -518,7 +522,8 @@ def _assert_same_selection( def _sparse_scratch(rows: int, num_sparse: int, topk: int) -> dict: """Caller-owned buffers the sparse path writes into (the metadata builder - owns these in production).""" + owns these in production). + """ i32 = dict(dtype=torch.int32, device="cuda") return dict( sparse_indices=torch.zeros(rows, num_sparse, **i32), @@ -533,7 +538,8 @@ def _expand_candidates_reference( """PyTorch reference for `candidate_blocks_to_sparse_indices`: expand each candidate block into ``cbk // sbk`` sparse blocks, keep those inside ``[0, ke - ks)``, sort ascending, anchor at ``ks // sbk``, pad by - repeating the last valid block, and count the valid sparse columns.""" + repeating the last valid block, and count the valid sparse columns. + """ ratio = cbk // sbk rows, num_candidates = candidates.shape num_sparse = num_candidates * ratio @@ -559,7 +565,8 @@ def _expand_candidates_reference( @pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel") def test_candidate_blocks_to_sparse_indices_math(): """Candidate->sparse-block expansion: ks anchoring, range filtering, - repeat-last padding, and the per-row valid column count (``end``).""" + repeat-last padding, and the per-row valid column count (``end``). + """ from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( candidate_blocks_to_sparse_indices, ) @@ -611,7 +618,8 @@ def test_candidate_blocks_to_sparse_indices_math(): def test_candidate_blocks_to_sparse_indices_matches_reference(cbk: int, sbk: int): """The Triton expansion equals the PyTorch reference on production-shaped random inputs, including strided (column-sliced) candidate rows and - strided output buffers.""" + strided output buffers. + """ from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( candidate_blocks_to_sparse_indices, ) @@ -681,7 +689,8 @@ def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool): ``torch.topk`` over dense bf16 logits masked to the candidates (the two DeepGEMM kernels agree bitwise on candidate columns), incl. under CUDA graph capture, and a later layer reusing the step's expansion/schedule - selects the same values.""" + selects the same values. + """ deep_gemm = _skip_unless_sm100_sparse_kernels() from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( sparse_mqa_logits_prefill_chunk, @@ -771,7 +780,8 @@ def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool): def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int): """MXFP4 paged cache in DeepGEMM's fused layout: per page, all tokens' - fp4 data first, then their packed UE8M0 scales.""" + fp4 data first, then their packed UE8M0 scales. + """ device = req_lens.device pages_per = (req_lens + page_kv - 1) // page_kv num_pages, max_pages = int(pages_per.sum()), int(pages_per.max()) @@ -807,7 +817,8 @@ def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool """Paged sparse-MQA decode selects exactly the values ``torch.topk`` picks from the dense bf16 masked logits, incl. MTP rows (next_n > 1 flattened to one query row each), CUDA graph capture/replay and schedule reuse by a - later layer.""" + later layer. + """ deep_gemm = _skip_unless_sm100_sparse_kernels() from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( sparse_mqa_logits_paged_decode, @@ -916,7 +927,8 @@ def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool def test_sparse_mqa_logits_accuracy_vs_torch(): """Sparse-kernel numerics vs a PyTorch fp32 reference on the dequantized - MXFP4 inputs.""" + MXFP4 inputs. + """ _skip_unless_sm100_sparse_kernels() from deep_gemm.utils import cast_back_from_fp4 diff --git a/tests/model_executor/model_loader/test_ep_weight_filter.py b/tests/model_executor/model_loader/test_ep_weight_filter.py index 0ff53d24e7..9d177f5991 100644 --- a/tests/model_executor/model_loader/test_ep_weight_filter.py +++ b/tests/model_executor/model_loader/test_ep_weight_filter.py @@ -232,7 +232,8 @@ class TestShouldSkipWeight: class TestSafetensorsWeightsIteratorWithEpFilter: """Verify that EP filtering produces a strict subset of unfiltered loading - and that all expected dense + local expert weights are present.""" + and that all expected dense + local expert weights are present. + """ @pytest.fixture(scope="class") def gpt2_files(self): @@ -262,7 +263,8 @@ class TestSafetensorsWeightsIteratorWithEpFilter: def test_empty_filter_skips_experts_only(self, gpt2_files): """GPT-2 has no expert weights, so even an empty local_expert_ids - set should return all weights (all are dense).""" + set should return all weights (all are dense). + """ all_weights = dict(safetensors_weights_iterator(gpt2_files, False)) filtered_weights = dict( safetensors_weights_iterator(gpt2_files, False, local_expert_ids=set()) @@ -273,7 +275,8 @@ class TestSafetensorsWeightsIteratorWithEpFilter: class TestEpFilterOnSyntheticMoeWeights: """Create synthetic safetensors files with expert-like naming and verify - that the filter correctly skips non-local experts.""" + that the filter correctly skips non-local experts. + """ @pytest.fixture def synthetic_moe_files(self, tmp_path): diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index b4a544d2fc..d123dcb750 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -350,7 +350,8 @@ def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_i @pytest.mark.parametrize("variant", ["fp8", "mxfp8", "nvfp4"]) def test_marlin_prepare_layer_preserves_workspace_address(monkeypatch, variant): """The Marlin fallback prepare_* functions rerun on weight reload and must - reuse the workspace storage whose address captured CUDA graphs hold.""" + reuse the workspace storage whose address captured CUDA graphs hold. + """ from vllm import _custom_ops as ops from vllm.model_executor.layers.quantization.utils import ( marlin_utils, @@ -428,7 +429,8 @@ def test_marlin_prepare_layer_preserves_workspace_address(monkeypatch, variant): def test_marlin_make_workspace_new_rejects_incompatible_existing(monkeypatch): """An incompatible existing workspace means the address captured by CUDA - graphs is already unusable; allocating a replacement would hide that.""" + graphs is already unusable; allocating a replacement would hide that. + """ from vllm.model_executor.layers.quantization.utils import marlin_utils monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4) @@ -609,7 +611,8 @@ class _RecordingQuantMethod(QuantizeMethodBase): class _LateBiasLayer(torch.nn.Module): """Mimics an online-quantized linear: `weight` is created on meta by `create_weights()`, which wraps the loaders, and the linear base registers - `bias` afterwards.""" + `bias` afterwards. + """ def __init__(self, quant_method): super().__init__() diff --git a/tests/model_executor/test_deep_gemm_warmup.py b/tests/model_executor/test_deep_gemm_warmup.py index 6f2a31cf27..ada0737066 100644 --- a/tests/model_executor/test_deep_gemm_warmup.py +++ b/tests/model_executor/test_deep_gemm_warmup.py @@ -42,7 +42,8 @@ def _run_warmup(monkeypatch, layers) -> list[dict]: @pytest.mark.parametrize("scale_name", ["weight_scale", "weight_scale_inv"]) def test_registered_deep_gemm_layer_is_warmed_up(monkeypatch, scale_name) -> None: """Any layer stamped by the DeepGEMM kernel is warmed, regardless of the - quantization method that owns it or the name of its scale parameter.""" + quantization method that owns it or the name of its scale parameter. + """ layer = _block_fp8_layer(256, 128, scale_name) calls = _run_warmup(monkeypatch, [layer]) @@ -54,7 +55,8 @@ def test_registered_deep_gemm_layer_is_warmed_up(monkeypatch, scale_name) -> Non def test_n_multiple_of_64_matches_kernel_selection(monkeypatch) -> None: """DeepGEMM accepts N % 64 == 0 (e.g. DeepSeek kv_a_proj_with_mqa, N=576), - so warmup must not require N % 128 == 0.""" + so warmup must not require N % 128 == 0. + """ assert len(_run_warmup(monkeypatch, [_block_fp8_layer(576, 256)])) == 1 @@ -71,7 +73,7 @@ def test_unstamped_and_mismatched_block_layers_are_skipped(monkeypatch) -> None: @pytest.mark.parametrize("is_bmm", [False, True]) def test_kernel_registers_itself_as_warmup_provider(is_bmm) -> None: - """bmm layers bypass the kernel at runtime, so they must not be warmed.""" + """Bmm layers bypass the kernel at runtime, so they must not be warmed.""" kernel = object.__new__(DeepGemmFp8BlockScaledMMKernel) kernel.is_deep_gemm_supported = False layer = torch.nn.Module() diff --git a/tests/model_executor/test_model_load_with_params.py b/tests/model_executor/test_model_load_with_params.py index 785335e991..7d1bddab0d 100644 --- a/tests/model_executor/test_model_load_with_params.py +++ b/tests/model_executor/test_model_load_with_params.py @@ -23,9 +23,7 @@ REVISION_ROBERTA = os.environ.get("REVISION", "main") current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm." ) def test_model_loading_with_params(vllm_runner, monkeypatch): - """ - Test parameter weight loading with tp>1. - """ + """Test parameter weight loading with tp>1.""" # to use apply_model monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") with vllm_runner( @@ -68,9 +66,7 @@ def test_model_loading_with_params(vllm_runner, monkeypatch): current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm." ) def test_roberta_model_loading_with_params(vllm_runner, monkeypatch): - """ - Test parameter weight loading with tp>1. - """ + """Test parameter weight loading with tp>1.""" # to use apply_model monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") with vllm_runner( @@ -113,9 +109,7 @@ def test_roberta_model_loading_with_params(vllm_runner, monkeypatch): current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm." ) def test_facebook_roberta_model_loading_with_params(vllm_runner, monkeypatch): - """ - Test loading roberta-base model with no lm_head. - """ + """Test loading roberta-base model with no lm_head.""" # to use apply_model monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") model_name = "FacebookAI/roberta-base" diff --git a/tests/model_executor/test_oink_integration.py b/tests/model_executor/test_oink_integration.py index 2f37472b73..b16d716413 100644 --- a/tests/model_executor/test_oink_integration.py +++ b/tests/model_executor/test_oink_integration.py @@ -62,7 +62,6 @@ def test_oink_availability_checks( expected_fused: bool, ): """Test OINK support detection with clean import state for each parameter set.""" - # Use spawn to run function in fresh process with clean imports # TODO migrate to spawn utility: # https://github.com/vllm-project/vllm/issues/41415 diff --git a/tests/model_executor/test_qwen3_omni.py b/tests/model_executor/test_qwen3_omni.py index 30d9909a5a..cf3df85caa 100644 --- a/tests/model_executor/test_qwen3_omni.py +++ b/tests/model_executor/test_qwen3_omni.py @@ -16,8 +16,7 @@ from vllm.multimodal.processing import InputProcessingContext # Helper function to print input IDs with coalesced audio/video tokens. def print_input_ids(input_ids): - """ - Print input IDs, compressing consecutive special tokens. + """Print input IDs, compressing consecutive special tokens. - 151675: <|audio_pad|> - 151656: <|video_pad|> """ @@ -137,7 +136,6 @@ def test_qwen3_omni_get_updates_use_audio_in_video( mock_image_processor, ): """Test the get_updates_use_audio_in_video method directly.""" - from vllm.model_executor.models.qwen3_omni_moe_thinker import ( Qwen3OmniMoeThinkerMultiModalProcessor, Qwen3OmniMoeThinkerProcessingInfo, diff --git a/tests/model_executor/test_qwen3_vl_mrope.py b/tests/model_executor/test_qwen3_vl_mrope.py index 90d9fd6e4f..a50bf319c9 100644 --- a/tests/model_executor/test_qwen3_vl_mrope.py +++ b/tests/model_executor/test_qwen3_vl_mrope.py @@ -51,8 +51,7 @@ class DummyConfig: def make_video_embedding( t, h, w, interleave_text_tokens: tuple[int, int], video_pruning_rate: float = 0.0 ): - """ - Helper function to make a video embedding for a given video size and pruning rate. + """Make a video embedding for a given video size and pruning rate. Args: t: Number of frames. @@ -64,6 +63,7 @@ def make_video_embedding( Returns: Tuple of (unpruned_tokens_sequence, pruned_tokens_sequence, retention_mask) + """ unpruned_tokens_sequence = [] population = list(range(1, 100)) diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index 49986a991e..ed10955f96 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/tests/model_executor/test_routed_experts_capture.py @@ -278,7 +278,7 @@ def test_routed_experts_capturer_single_dp_no_metadata(): def test_routed_experts_capturer_dp_naive_concatenated_all_ranks(): - """n == sum(num_tokens_dp): slice this rank's segment from concatenated topk.""" + """N == sum(num_tokens_dp): slice this rank's segment from concatenated topk.""" capturer = _capturer_with_buffer(dp_rank=1) num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) ctx = SimpleNamespace( @@ -295,7 +295,7 @@ def test_routed_experts_capturer_dp_naive_concatenated_all_ranks(): def test_routed_experts_capturer_dp_modular_local_tokens(): - """n == token_num_per_dp: topk is already local to this DP rank.""" + """N == token_num_per_dp: topk is already local to this DP rank.""" capturer = _capturer_with_buffer(dp_rank=1) num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) ctx = SimpleNamespace( @@ -326,7 +326,8 @@ def test_routed_experts_capturer_dp_unexpected_batch_raises(): def test_routed_experts_attention_group_is_shared_and_fail_closed(): """Both sides key routing data by this gid, so it must skip non-full-attention - groups rather than defaulting to 0, and fail closed when none exists.""" + groups rather than defaulting to 0, and fail closed when none exists. + """ common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32) config = SimpleNamespace( kv_cache_groups=[ diff --git a/tests/model_executor/test_weight_utils.py b/tests/model_executor/test_weight_utils.py index e577690c92..d19ea0ed05 100644 --- a/tests/model_executor/test_weight_utils.py +++ b/tests/model_executor/test_weight_utils.py @@ -57,7 +57,8 @@ class TestMaybeRemapKvScaleName: def test_qkv_proj_k_scale(self): """Qwen3-MoE / llm-compressor format: qkv_proj.k_scale -> attn.k_scale - Regression test for https://github.com/vllm-project/vllm/issues/25047""" + Regression test for https://github.com/vllm-project/vllm/issues/25047 + """ result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.qkv_proj.k_scale", self.PARAMS_DICT ) @@ -65,35 +66,36 @@ class TestMaybeRemapKvScaleName: def test_qkv_proj_v_scale(self): """Qwen3-MoE / llm-compressor format: qkv_proj.v_scale -> attn.v_scale - Regression test for https://github.com/vllm-project/vllm/issues/25047""" + Regression test for https://github.com/vllm-project/vllm/issues/25047 + """ result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.qkv_proj.v_scale", self.PARAMS_DICT ) assert result == "model.layers.0.self_attn.attn.v_scale" def test_modelopt_k_proj_k_scale(self): - """ModelOpt format: k_proj.k_scale -> attn.k_scale""" + """ModelOpt format: k_proj.k_scale -> attn.k_scale.""" result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.k_proj.k_scale", self.PARAMS_DICT ) assert result == "model.layers.0.self_attn.attn.k_scale" def test_modelopt_v_proj_v_scale(self): - """ModelOpt format: v_proj.v_scale -> attn.v_scale""" + """ModelOpt format: v_proj.v_scale -> attn.v_scale.""" result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.v_proj.v_scale", self.PARAMS_DICT ) assert result == "model.layers.0.self_attn.attn.v_scale" def test_deprecated_kv_scale(self): - """Old format: kv_scale -> attn.k_scale (deprecated)""" + """Old format: kv_scale -> attn.k_scale (deprecated).""" result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.kv_scale", self.PARAMS_DICT ) assert result == "model.layers.0.self_attn.attn.k_scale" def test_default_bare_k_scale(self): - """Default format: .k_scale -> .attn.k_scale""" + """Default format: .k_scale -> .attn.k_scale.""" result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.k_scale", self.PARAMS_DICT ) @@ -108,7 +110,8 @@ class TestMaybeRemapKvScaleName: def test_nvfp4_modelopt_k_proj_k_scale(self): """ModelOpt NVFP4 format (e.g. nvidia/Qwen3-30B-A3B-NVFP4): k_proj.k_scale -> attn.k_scale. - Validates that NVFP4 checkpoints are not broken by this change.""" + Validates that NVFP4 checkpoints are not broken by this change. + """ result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.k_proj.k_scale", self.PARAMS_DICT ) @@ -117,7 +120,8 @@ class TestMaybeRemapKvScaleName: def test_nvfp4_modelopt_v_proj_v_scale(self): """ModelOpt NVFP4 format (e.g. nvidia/Qwen3-30B-A3B-NVFP4): v_proj.v_scale -> attn.v_scale. - Validates that NVFP4 checkpoints are not broken by this change.""" + Validates that NVFP4 checkpoints are not broken by this change. + """ result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.v_proj.v_scale", self.PARAMS_DICT ) @@ -125,7 +129,8 @@ class TestMaybeRemapKvScaleName: def test_qwen3_vl_moe_qkv_proj_k_scale(self): """Qwen3-VL-MoE uses the same fused qkv_proj naming as Qwen3-MoE. - Regression test for qwen3_vl_moe.py fix (same bug as #25047).""" + Regression test for qwen3_vl_moe.py fix (same bug as #25047). + """ result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.qkv_proj.k_scale", self.PARAMS_DICT ) @@ -133,7 +138,8 @@ class TestMaybeRemapKvScaleName: def test_qwen3_vl_moe_qkv_proj_v_scale(self): """Qwen3-VL-MoE uses the same fused qkv_proj naming as Qwen3-MoE. - Regression test for qwen3_vl_moe.py fix (same bug as #25047).""" + Regression test for qwen3_vl_moe.py fix (same bug as #25047). + """ result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.qkv_proj.v_scale", self.PARAMS_DICT ) @@ -175,7 +181,8 @@ class TestKvCacheScaleMapper: """The `WeightsMapper` returned by `get_cache_scale_mapper` replaces the per-model `maybe_remap_kv_scale_name` calls. It must remap the same set of checkpoint formats (the non-`params_dict`-dependent ones) and be idempotent - so it composes safely with a model's own qkv/gate_up `hf_to_vllm_mapper`.""" + so it composes safely with a model's own qkv/gate_up `hf_to_vllm_mapper`. + """ def _mapper(self): # `get_cache_scale_mapper` does not use `self`; call it on the base @@ -265,7 +272,8 @@ class TestKvCacheScaleMapper: def test_composes_with_qkv_mapper(self): """Applied together with a model's qkv/gate_up mapper, the regex scale rules run before the substr rename, so scales are normalized to `.attn.` - and regular projections are still fused correctly.""" + and regular projections are still fused correctly. + """ from vllm.model_executor.models.utils import WeightsMapper model_mapper = WeightsMapper( diff --git a/tests/models/inkling/test_mm_towers.py b/tests/models/inkling/test_mm_towers.py index d969c8101a..51daa506c2 100644 --- a/tests/models/inkling/test_mm_towers.py +++ b/tests/models/inkling/test_mm_towers.py @@ -42,7 +42,8 @@ def _assert_close_ulps( compare against the reference's local ulp spacing. ``pre_act`` (the reference pre-activation) adds derivative-propagated slack: near ``gelu(x) ~ 0`` an ulp-level input flip legitimately moves the output by - many of ITS (tiny) ulps, bounded by |gelu'| <= 1.13 times the input ulps.""" + many of ITS (tiny) ulps, bounded by |gelu'| <= 1.13 times the input ulps. + """ g, r = got.float(), ref.float() tol = torch.clamp(max_ulps * _bf16_spacing(ref), min=atol) if pre_act is not None: diff --git a/tests/models/kimi_k3/test_amd_latent_moe_runner.py b/tests/models/kimi_k3/test_amd_latent_moe_runner.py index ec270c04fe..4346a72928 100644 --- a/tests/models/kimi_k3/test_amd_latent_moe_runner.py +++ b/tests/models/kimi_k3/test_amd_latent_moe_runner.py @@ -123,7 +123,8 @@ def _check_writes_only_its_own_shard( device: torch.device, tp_size: int, rank: int ) -> None: """Two ranks that swapped both weight rows and write offsets still sum to - the right total, so inspect each slice before the final collective.""" + the right total, so inspect each slice before the final collective. + """ transform = _build_transform(device) runner = _tail_runner(transform, tp_size) group = get_tp_group().device_group diff --git a/tests/models/kimi_k3/test_aux_attn_res_stream.py b/tests/models/kimi_k3/test_aux_attn_res_stream.py index 26cb76a550..64cec452d2 100644 --- a/tests/models/kimi_k3/test_aux_attn_res_stream.py +++ b/tests/models/kimi_k3/test_aux_attn_res_stream.py @@ -130,7 +130,8 @@ def test_disabled_reproduces_the_plain_residual_sum( def test_taps_the_consumer_layer_when_one_follows(recorder, monkeypatch): """The value the next layer reads is the mixture against *its* weights, - so the tap has to reach forward rather than use the current layer's.""" + so the tap has to reach forward rather than use the current layer's. + """ _set_last_rank(monkeypatch, True) _call(_stub_model(enabled=True), 1, torch.zeros(2), None, torch.zeros(2)) @@ -181,7 +182,8 @@ def test_pending_mlp_output_is_folded_in_rather_than_passed_as_delta( ): """The kernel writes an applied delta back into the prefix in place, which would double-add it into the live residual stream, so the pending output - has to arrive already summed into the prefix with `delta` left None.""" + has to arrive already summed into the prefix with `delta` left None. + """ _set_last_rank(monkeypatch, True) prefix_sum = torch.tensor([1.0, 2.0]) pending = torch.tensor([0.5, 0.25]) diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index fa810217b1..5dad5754b6 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -727,7 +727,8 @@ def test_cudagraph_capture_metadata_avoids_device_to_host_copy( """`build_for_cudagraph_capture` is not capture-only: a DP rank with nothing scheduled re-stages its FULL-graph metadata through it on every dummy step, so it must not synchronize the device. The host-side draft counts have to - come from `query_start_loc_cpu` and match the device-derived values.""" + come from `query_start_loc_cpu` and match the device-derived values. + """ num_speculative_tokens = 2 batch = BatchSpec(seq_lens=[50, 30], query_lens=[3, 3]) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") diff --git a/tests/models/language/generation/test_gemma.py b/tests/models/language/generation/test_gemma.py index c2c8e88e42..307a572afd 100644 --- a/tests/models/language/generation/test_gemma.py +++ b/tests/models/language/generation/test_gemma.py @@ -60,7 +60,8 @@ def test_checkpoint_lm_head_can_override_tied_config(monkeypatch) -> None: def test_gemma4_kv_shared_layer_loads_plain_q_proj() -> None: """KV-shared layers have q_proj instead of a packed qkv_proj; their redundant K/V tensors in original checkpoints have no parameter and are - skipped rather than failing the load.""" + skipped rather than failing the load. + """ model = torch.nn.Module() model.config = SimpleNamespace(num_experts=0) model.start_layer, model.end_layer = 0, 2 @@ -92,7 +93,8 @@ def test_gemma4_kv_shared_layer_loads_plain_q_proj() -> None: @pytest.mark.cpu_test def test_gemma3n_kv_shared_layer_mapper() -> None: """Only non-shared layers pack q/k/v into qkv_proj; KV-shared layers keep - q_proj and drop the redundant K/V tensors original checkpoints ship.""" + q_proj and drop the redundant K/V tensors original checkpoints ship. + """ config = SimpleNamespace(num_hidden_layers=4, num_kv_shared_layers=2) mapper = Gemma3nTextModel.hf_to_vllm_mapper | _kv_sharing_weights_mapper(config) weights = [ diff --git a/tests/models/language/generation/test_hybrid.py b/tests/models/language/generation/test_hybrid.py index c5406fae48..b77670956a 100644 --- a/tests/models/language/generation/test_hybrid.py +++ b/tests/models/language/generation/test_hybrid.py @@ -164,8 +164,7 @@ def test_chunked_prefill_with_parallel_sampling( max_tokens: int, conv_state_layout: str, ) -> None: - """ - Tests chunked prefill in conjunction with n > 1. + """Tests chunked prefill in conjunction with n > 1. In this case, prefill is populated with decoding tokens and we test that it doesn't fail. @@ -199,8 +198,7 @@ def test_mamba_cache_cg_padding( max_tokens: int, conv_state_layout: str, ) -> None: - """ - This test is for verifying that mamba cache is padded to CG captured + """This test is for verifying that mamba cache is padded to CG captured batch size. If it's not, a torch RuntimeError will be raised because tensor dimensions aren't compatible. """ @@ -236,8 +234,7 @@ def test_fail_upon_inc_requests_and_finished_requests_lt_available_blocks( example_prompts, model: str, ) -> None: - """ - This test is for verifying that the hybrid inner state management doesn't + """This test is for verifying that the hybrid inner state management doesn't collapse in case where the number of incoming requests and finished_requests_ids is larger than the maximum mamba block capacity. @@ -263,8 +260,7 @@ def test_state_cleanup( example_prompts, model: str, ) -> None: - """ - This test is for verifying that the Hybrid state is cleaned up between + """This test is for verifying that the Hybrid state is cleaned up between steps. If it's not cleaned, an error would be expected. diff --git a/tests/models/language/generation/test_mistral.py b/tests/models/language/generation/test_mistral.py index fc774bc6e8..a21c5c355a 100644 --- a/tests/models/language/generation/test_mistral.py +++ b/tests/models/language/generation/test_mistral.py @@ -274,7 +274,8 @@ def test_mistral_function_calling(vllm_runner, model: str, dtype: str) -> None: def test_mistral_function_call_nested_json(): """Ensure that the function-name regex captures the entire outermost - JSON block, including nested braces.""" + JSON block, including nested braces. + """ # Create a minimal stub tokenizer that provides the few attributes the # parser accesses (`version` and `get_vocab`). diff --git a/tests/models/language/pooling/test_max_tokens_per_doc.py b/tests/models/language/pooling/test_max_tokens_per_doc.py index d99aa49909..b109b935df 100644 --- a/tests/models/language/pooling/test_max_tokens_per_doc.py +++ b/tests/models/language/pooling/test_max_tokens_per_doc.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for max_tokens_per_doc and max_tokens_per_query. -""" +"""Tests for max_tokens_per_doc and max_tokens_per_query.""" import json import os diff --git a/tests/models/language/pooling/test_splade_sparse_pooler.py b/tests/models/language/pooling/test_splade_sparse_pooler.py index 0b3561f90f..d470637477 100644 --- a/tests/models/language/pooling/test_splade_sparse_pooler.py +++ b/tests/models/language/pooling/test_splade_sparse_pooler.py @@ -29,7 +29,8 @@ from vllm.v1.worker.gpu.states import RequestState @torch.inference_mode def test_splade_pooler_matches_reference_formula(B, T, H, V): """Ensure SPLADESparsePooler forward() matches the mathematical formula: - log1p(relu(logits)) -> max over sequence length (after masking).""" + log1p(relu(logits)) -> max over sequence length (after masking). + """ torch.manual_seed(0) # Prepare [B] sequences of shape [T, H] diff --git a/tests/models/multimodal/generation/test_interleaved.py b/tests/models/multimodal/generation/test_interleaved.py index a773db1982..07ce235853 100644 --- a/tests/models/multimodal/generation/test_interleaved.py +++ b/tests/models/multimodal/generation/test_interleaved.py @@ -23,11 +23,9 @@ NONINTERLEAVED_PROMPT = base_prompt("